clean up

Author: brb-nv

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -562,13 +562,12 @@ def warmup(self, resource_manager: ResourceManager) -> None:
         # Reset the global cuda graph dummy request to None in warmup.
         self.cuda_graph_runner.padding_dummy_request = None
 
+        # TODO: current warmup_request is not suitable for context parallelism.
         cp_type = self.mapping.cp_config.get('cp_type', None)
         if cp_type is not None:
-            if cp_type in [CpType.ULYSSES, CpType.STAR]:
-                assert False, "cp_type must be HELIX for helix benchmarking."
-                print("[ModelEngine::warmup] EARLY RETURN since cp_type ",
-                      cp_type)
-                return
+            logger.info("[ModelEngine::warmup] Skipping warmup for cp_type: ",
+                        cp_type.name)
+            return
 
         self._run_torch_compile_warmup(resource_manager)
         self._run_autotuner_warmup(resource_manager)
@@ -1063,12 +1062,10 @@ def _init_max_seq_len(self):
             # NOTE: py_executor_creator makes sure that the executor uses this
             # smaller value as its max_seq_len too.
             logger.warning(
-                f"\n*******************************************************\n"
-                f"Specified {self.max_seq_len=} is larger than what the model can support\n"
-                f"({inferred_max_seq_len}). NOT Setting max_seq_len to {inferred_max_seq_len}. "
-                f"ARE YOU SURE ABOUT THIS?\n"
-                f"*******************************************************\n")
-            # self.max_seq_len = inferred_max_seq_len
+                f"Specified {self.max_seq_len=} is larger than what the model can support "
+                f"({inferred_max_seq_len}). Setting max_seq_len to {inferred_max_seq_len}. "
+            )
+            self.max_seq_len = inferred_max_seq_len
 
     def _infer_max_seq_len_from_config(self) -> int:
 
@@ -2137,8 +2134,7 @@ def _prepare_tp_inputs_no_cache(
         attn_metadata.padded_num_tokens = padded_num_tokens if padded_num_tokens != num_tokens else None
 
         if self.enable_attention_dp:
-            all_rank_num_tokens = self.dist.allgather(attn_metadata.num_tokens)
-            attn_metadata.all_rank_num_tokens = all_rank_num_tokens
+            attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens
 
         virtual_num_tokens = num_tokens
         if attn_metadata.padded_num_tokens is not None:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1889,9 +1889,6 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
 
         for req in scheduled_batch.generation_requests:
             if req.is_disagg_generation_transmission_complete:
-                print(
-                    "[PyExecutor::_prepare_disagg_gen_transmission_complete]: TRANSMISSION COMPLETE for request ID: ",
-                    req.py_request_id)
                 req.state = LlmRequestState.GENERATION_IN_PROGRESS
                 req.context_current_position = req.prompt_len
                 req.decoding_iter = 1
@@ -1903,9 +1900,6 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
                 beam_width = req.sampling_config.beam_width
 
                 for beam in range(0, beam_width):
-                    print(
-                        f"[PyExecutor::_prepare_disagg_gen_transmission_complete]: Adding new token {torch.tensor(first_gen_tokens[beam])} for beam {beam}."
-                    )
                     req.add_new_token(first_gen_tokens[beam], beam)
 
     @nvtx_range("_recv_disagg_gen_cache")
@@ -2001,24 +1995,12 @@ def _forward_step(self,
         )
         def forward(scheduled_requests, resource_manager, new_tensors_device,
                     gather_context_logits, cache_indirection_buffer):
-            # iter_begin = time.time()
-            result = self.model_engine.forward(
+            return self.model_engine.forward(
                 scheduled_requests,
                 resource_manager,
                 new_tensors_device,
                 gather_context_logits=gather_context_logits,
                 cache_indirection_buffer=cache_indirection_buffer)
-            # torch.cuda.synchronize()
-            # iter_end = time.time()
-            # iter_latency_ms = (iter_end - iter_begin) * 1e3
-            # if self.model_engine.iter_counter > 10 and self.dist.rank == 0:
-            #     logger.info(f"[PyExecutor::_forward_step] CUSTOM LOG: iter={self.model_engine.iter_counter}, "
-            #                 f"rank={self.dist.rank}, "
-            #                 f"active_requests={len(self.active_requests)}, "
-            #                 f"scheduled_generation_requests={len(scheduled_requests.generation_requests)}, "
-            #                 f"scheduled_batch_size={scheduled_requests.batch_size}, "
-            #                 f"iter_latency_ms={iter_latency_ms}ms")
-            return result
 
         try:
             gather_context_logits = any(
@@ -2085,8 +2067,7 @@ def _update_request_states_star_attention(
     @nvtx_range("_update_request_states")
     def _update_request_states(self, scheduled_requests: ScheduledRequests):
         cp_config = self.dist.cp_config
-        # note: helix parallelism uses the same logic as tp parallelism here
-        if 'cp_type' in cp_config and cp_config['cp_type'] != CpType.HELIX:
+        if 'cp_type' in cp_config:
             cp_type = cp_config['cp_type']
             if cp_type == CpType.STAR:
                 self._update_request_states_star_attention(scheduled_requests)

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -177,10 +177,6 @@ def __init__(
         indexer_k_cache_index_head_dim: int = 0,
         **kwargs,
     ) -> None:
-        # Couple of places where we assume tokens_per_block is 32: Let's assert here for now.
-        # 1) block assignment in merge_helix_requests
-        # 2) computation of cache_transceiver_config.max_tokens_in_buffer.
-        assert tokens_per_block == 32, "tokens_per_block must be 32 for helix benchmarking."
         self.mapping = mapping
         self.dtype = dtype
         self.kv_cache_type = kv_cache_type
@@ -443,18 +439,17 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
 
             for req in generation_batch:
                 # Skip allocating KV cache at decode for inactive helix ranks.
-                ##################################################################
-                # TODO: This should be set elsewhere. For now, we hardcode that last rank is active.
-                # Maybe right after pyexecutor._schedule() or in sampler.update_requests() at end of
-                # executor loop for next step.
                 if self.mapping.has_cp_helix():
                     if self.mapping.cp_rank != self.mapping.cp_size - 1:
                         req.py_helix_is_inactive_rank = True
-                ##################################################################
                 if req.py_helix_is_inactive_rank:
-                    # print(f"[ResourceManager::prepare_resources][rank {self.mapping.rank}] Skipping KV allocation for request {req.py_request_id}.")
+                    print(
+                        f"[ResourceManager::prepare_resources][rank {self.mapping.rank}] Skipping KV allocation for request {req.py_request_id}."
+                    )
                     continue
-                # print(f"[ResourceManager::prepare_resources][rank {self.mapping.rank}] Adding KV allocation for request {req.py_request_id}.")
+                print(
+                    f"[ResourceManager::prepare_resources][rank {self.mapping.rank}] Adding KV allocation for request {req.py_request_id}."
+                )
                 self.impl.add_token(req.py_request_id)
                 for _ in range(get_draft_token_length(req)):
                     self.impl.add_token(req.py_request_id)