tensorrt_llm/_torch/pyexecutor/model_engine.py (2)

1526-1554: HELIX: avoid per‑beam allgather and fix position_id to be global.

• Calling cp_allgather inside the beam loop is redundant and can cause unnecessary collectives.
• More importantly, generation position_ids for HELIX should use the global position (sum across CP ranks). Compute once per request and reuse across beams.

This also sets up a single global position_id you can reuse for MRoPE generation below (see next comment).

Apply:

             request_ids.append(request.py_request_id)
             beam_width = request.sampling_config.beam_width
-            for beam in range(beam_width):
+            # Compute local and global position_id once per request (avoid per-beam collectives)
+            past_seen_token_num = (request.max_beam_num_tokens - 1
+                                   if (new_tokens_device is None or request.is_dummy or request.py_batch_idx is None)
+                                   else request.max_beam_num_tokens)
+            helix_position_id = past_seen_token_num
+            if self.mapping.has_cp_helix():
+                helix_position_id = sum(self.dist.cp_allgather(past_seen_token_num))
+
+            for beam in range(beam_width):
                 # the request has no previous tensor:
                 # (1) new_tokens_device is None, which means overlap scheduler is disabled; or
                 # (2) a dummy request; or
                 # (3) the first step in the generation server of disaggregated serving
                 if new_tokens_device is None or request.is_dummy or request.py_batch_idx is None:
                     # skip adding input_ids of CUDA graph dummy requests so that new_tokens_device
                     # can be aligned to the correct positions.
                     if not request.is_cuda_graph_dummy:
                         input_ids.append(request.get_last_tokens(beam))
-                    past_seen_token_num = request.max_beam_num_tokens - 1
                 else:
                     # the request has previous tensor
                     # previous_batch_indices is used per request, not per beam
                     # Only append it once for the first beam of each request
                     first_beam = 0
                     if beam == first_beam:
                         previous_batch_indices.append(request.py_batch_idx)
-                    past_seen_token_num = request.max_beam_num_tokens
-                position_id = past_seen_token_num
-                if self.mapping.has_cp_helix():
-                    # Do an allgather among CP ranks to get the complete sequence length seen by all CP ranks.
-                    past_seen_token_nums = self.dist.cp_allgather(
-                        past_seen_token_num)
-                    position_id = sum(past_seen_token_nums)
-                position_ids.append(position_id)
+                position_ids.append(helix_position_id)
                 num_cached_tokens_per_seq.append(past_seen_token_num)

---

1566-1572: HELIX + MRoPE: position deltas must add to global position, not local.

Generation MRoPE position ids currently add deltas to the local past_seen_token_num, diverging from the global position_ids you just fixed. Use the same global position.

Apply:

-                        gen_mrope_position_ids = (past_seen_token_num +
+                        gen_mrope_position_ids = (helix_position_id +
                                                   mrope_position_deltas).expand(
                                                       3, 1, 1)

tensorrt_llm/_torch/distributed/communicator.py (1)

452-460: Restore broadcast return value (and cover Helix CP)

The refactor drops the return statements, so broadcast() now always yields None, breaking every call site that relied on the object being returned. It also forgets to handle Helix CP, so a Helix-only mapping falls through the else branch and returns None. Please reinstate the returns and cover both CP flavors:

-        if self.mapping.has_cp_ulysses():
-            self.broadcast_cp(obj, root)
-        elif self.mapping.has_tp():
-            self.broadcast_tp(obj, root)
-        else:
-            pass
+        if self.mapping.has_cp_ulysses() or self.mapping.has_cp_helix():
+            return self.broadcast_cp(obj, root)
+        if self.mapping.has_tp():
+            return self.broadcast_tp(obj, root)
+        return obj

Without this, Helix ranks lose their config broadcast entirely.