clean code

Signed-off-by: Yue Weng <[email protected]>

Author: yweng0828

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -1318,7 +1318,6 @@ MLA_FUNC_DEFINE(__nv_bfloat16)
 template <typename T, typename KVCacheBuffer>
 int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStream_t stream)
 {
-    printf("++++++++++++++++++++++++ in enqueueContext +++++++++++++++++++++++++\n");
     int const headSize = getHeadSize();
 
     int const local_hidden_units_qo = mNumHeads * headSize;
@@ -2162,7 +2161,6 @@ template int AttentionOp::enqueueContext<__nv_bfloat16, KVBlockArray>(
 template <typename T, typename KVCacheBuffer>
 int AttentionOp::enqueueGeneration(EnqueueGenerationParams<T> const& params, cudaStream_t stream)
 {
-    printf("++++++++++++++++++++++++ in enqueueGeneration +++++++++++++++++++++++++\n");
     int const headSize = getHeadSize();
     float const q_scaling = mQScaling;
     float const* logn_scaling_ptr = isLognScaling() ? params.logn_scaling_ptr : nullptr;
@@ -2272,10 +2270,7 @@ int AttentionOp::enqueueGeneration(EnqueueGenerationParams<T> const& params, cud
         // self attn
         XQAParams xqaParams{};
         this->template convertMMHAParamsToXQAParams<T, KVCacheBuffer>(xqaParams, params, /*forConfigurePlugin=*/false);
-        // if (mEnableXQA && mXqaDispatcher->shouldUse(xqaParams))
-        bool shouldUseXQA = mEnableXQA && mXqaDispatcher->shouldUse(xqaParams);
-        printf("++++++++++++++++++++++++ in enqueueGeneration, mEnableXQA: %d, shouldUseXQA: %d +++++++++++++++++++++++++\n", mEnableXQA, shouldUseXQA);
-        if (mEnableXQA && shouldUseXQA)
+        if (mEnableXQA && mXqaDispatcher->shouldUse(xqaParams))
         {
             TLLM_LOG_DEBUG("XQA kernels are selected in the generation phase.");
             xqaParams.stream = stream;

cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp

@@ -98,9 +98,6 @@ bool DecoderXQAImplJIT::mayHavePerfGain(XQAParams const& xqaParams) const
 
 bool DecoderXQAImplJIT::shouldUse(XQAParams const& umbrellaXQAParams, bool forConfigurePlugin)
 {
-    printf("++++++++++++++++++++++++ in decoderXQAImplJIT::shouldUse +++++++++++++++++++++++++\n");
-    // printf("++++++++++++++++++++++++ umbrellaXQAParams: %s +++++++++++++++++++++++++\n", umbrellaXQAParams.toString().c_str());
-
     if (forConfigurePlugin)
     {
         for (int beam_width = 1; beam_width <= umbrellaXQAParams.beam_width; ++beam_width)

cpp/tensorrt_llm/kernels/xqaDispatcher.cpp

@@ -491,119 +491,6 @@ void XqaDispatcher::runImpl(
     }
     else
     {
-        std::cout << "===== debug: in xqaDispatcher.cpp::runImpl" << std::endl;
-        std::cout << "============== debug: print XQA params ==============" << std::endl;
-
-        // batch_size
-        auto batch_size = params.batch_size;
-        auto beam_width = params.beam_width;
-        auto max_draft_len = 12;    // hard code
-        std::cout << "===== debug: batch_size: " << batch_size << ", beam_width: " << beam_width << std::endl;
-
-        // host_past_key_value_lengths
-        if (params.host_past_key_value_lengths)
-        {
-            std::cout << "===== debug: host_past_key_value_lengths: " << *(runtime::ITensor::wrap((void*) params.host_past_key_value_lengths, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size}))) << std::endl;
-        }
-        // host_context_lengths
-        if (params.host_context_lengths)
-        {
-            std::cout << "===== debug: host_context_lengths: " << *(runtime::ITensor::wrap((void*) params.host_context_lengths, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size}))) << std::endl;
-        }
-        // chunked_attention_size
-        auto chunked_attention_size = params.chunked_attention_size;
-        std::cout << "===== debug: chunked_attention_size: " << chunked_attention_size << std::endl;
-        // max_attention_window_size
-        auto max_attention_window_size = params.max_attention_window_size;
-        std::cout << "===== debug: max_attention_window_size: " << max_attention_window_size << std::endl;
-        // cyclic_attention_window_size
-        auto cyclic_attention_window_size = params.cyclic_attention_window_size;
-        std::cout << "===== debug: cyclic_attention_window_size: " << cyclic_attention_window_size << std::endl;
-        // sink_token_length
-        auto sink_token_length = params.sink_token_length;
-        std::cout << "===== debug: sink_token_length: " << sink_token_length << std::endl;
-        // max_past_kv_length
-        auto max_past_kv_length = params.max_past_kv_length;
-        // sequence_lengths
-        if (params.sequence_lengths)
-        {
-            std::cout << "===== debug: sequence_lengths: " << *(runtime::ITensor::wrap((void*) params.sequence_lengths, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size * beam_width}))) << std::endl;
-        }
-        // context_lengths
-        if (params.context_lengths)
-        {
-            std::cout << "===== debug: context_lengths: " << *(runtime::ITensor::wrap((void*) params.context_lengths, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size * beam_width}))) << std::endl;
-        }
-
-        // spec_decoding_packed_mask
-        if (params.spec_decoding_packed_mask)
-        {
-            std::cout << "===== debug: spec_decoding_packed_mask: " << *(runtime::ITensor::wrap((void*) params.spec_decoding_packed_mask, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({max_draft_len+1, batch_size}))) << std::endl;
-        }
-        // spec_decoding_position_offsets
-        if (params.spec_decoding_position_offsets)
-        {
-            std::cout << "===== debug: spec_decoding_position_offsets: " << *(runtime::ITensor::wrap((void*) params.spec_decoding_position_offsets, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size, max_draft_len+1}))) << std::endl;
-        }
-        // spec_decoding_generation_lengths
-        if (params.spec_decoding_generation_lengths)
-        {
-            std::cout << "===== debug: spec_decoding_generation_lengths: " << *(runtime::ITensor::wrap((void*) params.spec_decoding_generation_lengths, nvinfer1::DataType::kINT32,
-                  runtime::ITensor::makeShape({batch_size}))) << std::endl;
-        }
-        // spec_decoding_is_generation_length_variable
-        auto spec_decoding_is_generation_length_variable = params.spec_decoding_is_generation_length_variable;
-        std::cout << "===== debug: spec_decoding_is_generation_length_variable: " << spec_decoding_is_generation_length_variable << std::endl;
-        // spec_decoding_max_generation_length
-        auto spec_decoding_max_generation_length = params.spec_decoding_max_generation_length;
-        std::cout << "===== debug: spec_decoding_max_generation_length: " << spec_decoding_max_generation_length << std::endl;
-        // generation_input_length
-        auto generation_input_length = params.generation_input_length;
-        std::cout << "===== debug: generation_input_length: " << generation_input_length << std::endl;
-        // num_q_heads
-        auto num_q_heads = params.num_q_heads;
-        std::cout << "===== debug: num_q_heads: " << num_q_heads << std::endl;
-        // num_kv_heads
-        auto num_kv_heads = params.num_kv_heads;
-        std::cout << "===== debug: num_kv_heads: " << num_kv_heads << std::endl;
-        // head_size
-        auto head_size = params.head_size;
-        std::cout << "===== debug: head_size: " << head_size << std::endl;
-        // unidirectional
-        auto unidirectional = params.unidirectional;
-        std::cout << "===== debug: unidirectional: " << unidirectional << std::endl;
-        // position_shift_enabled
-        auto position_shift_enabled = params.position_shift_enabled;
-        std::cout << "===== debug: position_shift_enabled: " << position_shift_enabled << std::endl;
-        // remove_padding
-        auto remove_padding = params.remove_padding;
-        std::cout << "===== debug: remove_padding: " << remove_padding << std::endl;
-        // mask_type
-        auto mask_type = params.mask_type;
-        std::cout << "===== debug: mask_type: " << static_cast<int>(mask_type) << std::endl;
-        // paged_kv_cache
-        auto paged_kv_cache = params.paged_kv_cache;
-        std::cout << "===== debug: paged_kv_cache: " << paged_kv_cache << std::endl;
-        // tokens_per_block
-        auto tokens_per_block = params.tokens_per_block;
-        std::cout << "===== debug: tokens_per_block: " << tokens_per_block << std::endl;
-        // max_blocks_per_sequence
-        auto max_blocks_per_sequence = params.max_blocks_per_sequence;
-        std::cout << "===== debug: max_blocks_per_sequence: " << max_blocks_per_sequence << std::endl;
-        // multi_query_tokens
-        auto multi_query_tokens = params.multi_query_tokens;
-        std::cout << "===== debug: multi_query_tokens: " << multi_query_tokens << std::endl;
-        // total_num_input_tokens
-        auto total_num_input_tokens = params.total_num_input_tokens;
-        std::cout << "===== debug: total_num_input_tokens: " << total_num_input_tokens << std::endl;
-        
-        std::cout << "=========================================" << std::endl;
         mDecoderXqaRunner->template dispatch<KVCacheBuffer>(params, kv_cache_buffer, params.stream);
     }
 }

examples/llm-api/quickstart_advanced.py

@@ -8,10 +8,9 @@
                                  TorchCompileConfig)
 
 example_prompts = [
-    # "Hello, my name is",
-    # "The capital of France is",
-    # "The future of AI is",
-    "You are a good assistant. Please tell me the capital of France is",
+    "Hello, my name is",
+    "The capital of France is",
+    "The future of AI is",
 ]
 
 

tensorrt_llm/_torch/attention_backend/interface.py

@@ -330,8 +330,9 @@ def restore_from_spec_dec(self) -> None:
             setattr(self, f, v)
         self._saved_tensors.clear()
 
-    def update_spec_dec_param(self, is_spec_decoding_enabled, spec_metadata, spec_tree_manager,
-                              max_draft_len, max_total_draft_tokens):
+    def update_spec_dec_param(self, is_spec_decoding_enabled, spec_metadata,
+                              spec_tree_manager, max_draft_len,
+                              max_total_draft_tokens):
         """
         Hook to be called when using TRTLLM attention backend in spec-dec mode.
         """

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1044,9 +1044,10 @@ def prepare_context_mla_with_cached_kv(self,
                      out=self.host_ctx_kv_indptr[1:self.num_contexts + 1])
         self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
             self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
-        
-    def update_spec_dec_param(self, is_spec_decoding_enabled, spec_metadata, spec_tree_manager, 
-                              max_draft_len, max_total_draft_tokens):
+
+    def update_spec_dec_param(self, is_spec_decoding_enabled, spec_metadata,
+                              spec_tree_manager, max_draft_len,
+                              max_total_draft_tokens):
         # spec_dec mode should only be enabled for pre-Blackwell machines and when there's a spec-dec tree.
         self.is_spec_decoding_enabled = is_spec_decoding_enabled and get_sm_version(
         ) < 100
@@ -1086,16 +1087,18 @@ def update_spec_dec_param(self, is_spec_decoding_enabled, spec_metadata, spec_tr
             is_target_model = not spec_metadata.is_draft_model
             is_using_tree = self.is_spec_dec_tree or self.is_spec_dec_dynamic_tree
             if is_target_model and is_using_tree:
-                assert spec_metadata.spec_dec_mode.is_eagle3(), "Tree decoding is only supported for Eagle3 now"
+                assert spec_metadata.spec_dec_mode.is_eagle3(
+                ), "Tree decoding is only supported for Eagle3 now"
                 # If is the dynamic tree
                 if self.is_spec_dec_dynamic_tree:
                     # TODO: add dynamic tree logic
                     assert False, "Dynamic tree is not supported yet"
                 # If is the static tree
                 else:
-                    self.spec_decoding_position_offsets[:,].copy_(
-                        spec_tree_manager.spec_dec_position_offsets[0, :],
-                        non_blocking=True)
+                    self.spec_decoding_position_offsets[
+                        :,
+                    ].copy_(spec_tree_manager.spec_dec_position_offsets[0, :],
+                            non_blocking=True)
                     self.spec_decoding_packed_mask[:, :, :].copy_(
                         spec_tree_manager.spec_dec_packed_mask[0, :, :],
                         non_blocking=True)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1222,7 +1222,7 @@ def _prepare_tp_inputs(
         num_accepted_draft_tokens = []  # per request
         # if using tree decoding, we need to store the request type and accepted path for each request,
         # which will be used to update the hidden_states_read_indices.
-        request_accepted_path = {} # per request
+        request_accepted_path = {}  # per request
 
         for request in scheduled_requests.context_requests:
             request_ids.append(request.py_request_id)
@@ -1237,7 +1237,9 @@ def _prepare_tp_inputs(
             gather_ids.append(len(input_ids) - 1)
             sequence_lengths.append(len(prompt_tokens))
             num_accepted_draft_tokens.append(len(prompt_tokens) - 1)
-            request_accepted_path[request.py_request_id] = request.py_num_accepted_draft_tokens_indices
+            request_accepted_path[
+                request.
+                py_request_id] = request.py_num_accepted_draft_tokens_indices
             prompt_lengths.append(len(prompt_tokens))
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
@@ -1323,7 +1325,9 @@ def _prepare_tp_inputs(
         previous_pos_indices = []
         for request in extend_requests:
             request_ids.append(request.py_request_id)
-            request_accepted_path[request.py_request_id] = request.py_num_accepted_draft_tokens_indices
+            request_accepted_path[
+                request.
+                py_request_id] = request.py_num_accepted_draft_tokens_indices
             # the request has no previous tensor:
             # (1) next_draft_tokens_device is None, which means overlap scheduler is disabled; or
             # (2) a dummy request; or
@@ -1359,13 +1363,15 @@ def _prepare_tp_inputs(
                     assert spec_tree_manager is not None
                     assert num_draft_tokens == spec_tree_manager.max_total_draft_tokens
                     position_ids.extend(
-                        past_seen_token_num + spec_tree_manager.spec_dec_position_offsets[0] # [max_total_draft_tokens + 1]
+                        past_seen_token_num +
+                        spec_tree_manager.spec_dec_position_offsets[
+                            0]  # [max_total_draft_tokens + 1]
                     )
                 else:
                     position_ids.extend(
                         list(
-                            range(past_seen_token_num, past_seen_token_num + 1 +
-                                num_draft_tokens)))
+                            range(past_seen_token_num,
+                                  past_seen_token_num + 1 + num_draft_tokens)))
                 num_cached_tokens_per_seq.append(past_seen_token_num)
                 request.cached_tokens = num_cached_tokens_per_seq[-1]
                 # update batch index
@@ -1390,12 +1396,15 @@ def _prepare_tp_inputs(
                     assert spec_tree_manager is not None
                     assert num_draft_tokens == spec_tree_manager.max_total_draft_tokens
                     position_ids.extend(
-                        past_seen_token_num + spec_tree_manager.spec_dec_position_offsets[0] # [max_total_draft_tokens + 1]
+                        past_seen_token_num +
+                        spec_tree_manager.spec_dec_position_offsets[
+                            0]  # [max_total_draft_tokens + 1]
                     )
                 else:
                     position_ids.extend(
                         list(
-                            range(past_seen_token_num, past_seen_token_num + 1 +
+                            range(
+                                past_seen_token_num, past_seen_token_num + 1 +
                                 self.runtime_draft_len)))
                 # previous tensor
                 previous_batch_indices.append(previous_batch_idx)
@@ -1433,7 +1442,9 @@ def _prepare_tp_inputs(
             sequence_lengths.append(1 + self.original_max_draft_len)
             num_accepted_draft_tokens.append(
                 request.py_num_accepted_draft_tokens)
-            request_accepted_path[request.py_request_id] = request.py_num_accepted_draft_tokens_indices
+            request_accepted_path[
+                request.
+                py_request_id] = request.py_num_accepted_draft_tokens_indices
             prompt_lengths.append(request.py_prompt_len)
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
@@ -2241,15 +2252,14 @@ def _get_lora_params_from_requests(self,
         return lora_params
 
     @nvtx_range("_prepare_inputs")
-    def _prepare_inputs(
-            self,
-            scheduled_requests: ScheduledRequests,
-            kv_cache_manager: KVCacheManager,
-            attn_metadata: AttentionMetadata,
-            spec_metadata: Optional[SpecMetadata] = None,
-            new_tensors_device: Optional[SampleStateTensors] = None,
-            cache_indirection_buffer: Optional[torch.Tensor] = None,
-            resource_manager: Optional[ResourceManager] = None):
+    def _prepare_inputs(self,
+                        scheduled_requests: ScheduledRequests,
+                        kv_cache_manager: KVCacheManager,
+                        attn_metadata: AttentionMetadata,
+                        spec_metadata: Optional[SpecMetadata] = None,
+                        new_tensors_device: Optional[SampleStateTensors] = None,
+                        cache_indirection_buffer: Optional[torch.Tensor] = None,
+                        resource_manager: Optional[ResourceManager] = None):
         if self.mapping is not None and 'cp_type' in self.mapping.cp_config:
             cp_type = self.mapping.cp_config['cp_type']
             if CpType.STAR == cp_type:
@@ -2297,7 +2307,8 @@ def forward(
                 self.model_is_wrapped, spec_metadata.is_spec_dec_tree)
             attn_metadata.update_spec_dec_param(
                 is_spec_dec_mode, spec_metadata, spec_tree_manager,
-                self.original_max_draft_len, self.original_max_total_draft_tokens)
+                self.original_max_draft_len,
+                self.original_max_total_draft_tokens)
         else:
             spec_resource_manager = None
             spec_metadata = None