minor changes

Author: brb-nv

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -693,6 +693,7 @@ def _merge_helix_requests(self, new_requests: list[RequestQueueItem],
                 input_token_ids=input_ids_this_rank,
                 position_ids=position_ids_this_rank,
             )
+            req.total_input_len_cp = input_len
             req_with_children.append(req)
             if req.child_requests:
                 req_with_children.extend(req.child_requests)

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1148,6 +1148,7 @@ def _executor_loop(self):
                     break
 
                 self._pause_requests(scheduled_batch.paused_requests)
+
                 finished_requests = []
 
                 can_queue = self._can_queue(scheduled_batch)
@@ -1898,7 +1899,6 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
                 ctx_draft_tokens = req.context_phase_params.draft_tokens
                 req.py_draft_tokens = [] if ctx_draft_tokens is None else ctx_draft_tokens
                 beam_width = req.sampling_config.beam_width
-
                 for beam in range(0, beam_width):
                     req.add_new_token(first_gen_tokens[beam], beam)
 

tests/unittest/_torch/modules/test_mla_helix.py

@@ -80,7 +80,7 @@ class Scenario:
     rope_original_max_position_embeddings: int = 4096
     rope_type: str = "yarn"
     model_type: str = "deepseek_v3"
-    kv_cache_tokens_per_block: int = 64
+    kv_cache_tokens_per_block: int = 32
     # TODO only 1 is supported for now here
     predicted_tokens_per_seq: int = 1
     bias: bool = False