[None][feat] Cudagraph updates for helix parallelism

Signed-off-by: Balaram Buddharaju <[email protected]>

Author: brb-nv

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -2,7 +2,7 @@
 import os
 import weakref
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import torch
 
@@ -83,6 +83,8 @@ class TrtllmAttentionWrapper:
     spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor]
     spec_decoding_bl_tree_mask: Optional[torch.Tensor]
     spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor]
+    helix_position_offsets: Optional[torch.Tensor]
+    helix_is_inactive_rank: Optional[torch.Tensor]
     kwargs: dict
 
     def __init__(
@@ -298,10 +300,6 @@ def plan(
         self.sparse_mla_topk = sparse_mla_topk
         self.helix_position_offsets = helix_position_offsets
         self.helix_is_inactive_rank = helix_is_inactive_rank
-        if self.helix_is_inactive_rank is not None and not isinstance(
-                self.helix_is_inactive_rank, torch.Tensor):
-            self.helix_is_inactive_rank = torch.tensor(
-                self.helix_is_inactive_rank, dtype=torch.bool, pin_memory=True)
 
         if max_sequence_length > self.rope_params.max_positions:
             self.rope_params.max_positions = max_sequence_length
@@ -646,6 +644,14 @@ class TrtllmAttentionMetadata(AttentionMetadata):
     spec_decoding_bl_tree_mask: Optional[torch.Tensor] = None
     spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None
 
+    # Flag to enable helix parallelism.
+    enable_helix: bool = False
+
+    # Global position ids of tokens for each sequence in the batch. Given
+    # each helix rank has only a subset of tokens for a sequence, we compute
+    # a global position id for each token here.
+    helix_position_offsets: Optional[torch.Tensor] = None
+
     # Whether the current rank is inactive for helix parallelism.
     # In helix parallelism, only the active rank appends KV cache for the query token
     # and attends to the previously cached tokens as well as the query token. Inactive
@@ -824,11 +830,64 @@ def _post_init_with_buffers(self, buffers) -> None:
                     pin_memory=True,
                 )
 
+        # Allocate static buffers for helix parallelism support
+        if self.enable_helix:
+            self.helix_position_offsets = self.get_empty(
+                buffers,
+                (self.max_num_tokens, ),
+                cache_name="helix_position_offsets",
+                dtype=torch.int,
+                capture_graph=capture_graph,
+            )
+            self.helix_position_offsets_cpu = torch.empty_like(
+                self.helix_position_offsets,
+                device='cpu',
+                pin_memory=True,
+            )
+            self.helix_is_inactive_rank = self.get_empty(
+                buffers,
+                (self.max_num_sequences, ),
+                cache_name="helix_is_inactive_rank",
+                dtype=torch.bool,
+                capture_graph=capture_graph,
+            )
+            self.helix_is_inactive_rank_cpu = torch.empty_like(
+                self.helix_is_inactive_rank,
+                device='cpu',
+                pin_memory=True,
+            )
+
     def on_update_kv_lens(self):
         # After changing the kv_lens/kv_lens_cuda, we may need to update other metadata.
         # Especially for the changes in the _preprocess_inputs() of model_engine.py.
         pass
 
+    def update_helix_param(
+        self,
+        helix_position_offsets: List[int],
+        helix_is_inactive_rank: List[bool],
+    ) -> None:
+        """
+        Update helix parameters by copying into static buffers for CUDA graph compatibility.
+
+        Args:
+            helix_position_offsets: Position offsets for helix parallelism with shape (num_tokens,).
+            helix_is_inactive_rank: Whether the current rank is inactive with shape (batch_size,).
+        """
+        if helix_position_offsets is not None and self.helix_position_offsets is not None:
+            num_tokens = len(helix_position_offsets)
+            self.helix_position_offsets_cpu[:num_tokens].copy_(
+                torch.tensor(helix_position_offsets, dtype=torch.int))
+            self.helix_position_offsets[:num_tokens].copy_(
+                self.helix_position_offsets_cpu[:num_tokens], non_blocking=True)
+
+        if helix_is_inactive_rank is not None and self.helix_is_inactive_rank is not None:
+            batch_size = len(helix_is_inactive_rank)
+            self.helix_is_inactive_rank_cpu[:batch_size].copy_(
+                torch.tensor(helix_is_inactive_rank, dtype=torch.bool))
+            self.helix_is_inactive_rank[:batch_size].copy_(
+                self.helix_is_inactive_rank_cpu[:batch_size], non_blocking=True)
+
     def prepare(self) -> None:
         extra_attrs = get_model_extra_attrs()
         # If model extra attrs is set, attention_metadata is setup in executor.
@@ -868,18 +927,13 @@ def prepare(self) -> None:
 
         if self.enable_flash_mla:
             self.prepare_flash_mla()
-        # number of tokens needed in the kv cache for each sequence after the next pass
-        if self.helix_is_inactive_rank is not None and len(
-                self.helix_is_inactive_rank):
+
+        # number of tokens needed in the kv cache for each sequence after the next pass.
+        if self.enable_helix:
             # If helix is inactive, attend to the previously cached tokens only.
             assert cached_token_lens is not None, "cached_token_lens should be set for helix"
+            active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs]
             kv_lens = cached_token_lens.clone()
-            helix_is_inactive_rank_cpu = torch.tensor(
-                self.helix_is_inactive_rank,
-                dtype=torch.bool,
-                device='cpu',
-            )
-            active_rank = ~helix_is_inactive_rank_cpu
             kv_lens[active_rank] += self.seq_lens_kv[active_rank]
         else:
             kv_lens = cached_token_lens + self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv
@@ -1485,7 +1539,6 @@ def forward(
         mrope_config: Optional[dict] = None,
         attention_window_size: Optional[int] = None,
         softmax_stats_tensor: Optional[torch.Tensor] = None,
-        helix_position_offsets: Optional[torch.Tensor] = None,
         enable_attn_nvfp4_output: bool = True,
         output: Optional[torch.Tensor] = None,
         output_sf: Optional[torch.Tensor] = None,
@@ -1596,7 +1649,7 @@ def forward(
             sparse_attn_indices_block_size=sparse_attn_indices_block_size,
             sparse_mla_topk=metadata.sparse_mla_topk if hasattr(
                 metadata, 'sparse_mla_topk') else 0,
-            helix_position_offsets=helix_position_offsets,
+            helix_position_offsets=metadata.helix_position_offsets,
             helix_is_inactive_rank=metadata.helix_is_inactive_rank,
         )
         out_dtype = None
@@ -1856,8 +1909,6 @@ def mla_rope_generation(
         mla_bmm1_scale: torch.Tensor,
         mla_bmm2_scale: torch.Tensor,
         quant_q_buffer: torch.Tensor,
-        helix_position_offsets: Optional[torch.Tensor] = None,
-        helix_is_inactive_rank: Optional[torch.Tensor] = None,
         out_scale: Optional[torch.Tensor] = None,
     ) -> None:
         """
@@ -1878,13 +1929,9 @@ def mla_rope_generation(
         assert metadata.kv_cache_manager is not None
         sink_token_length = 0
 
-        # Ensure helix_is_inactive_rank and position_ids are on the same device.
-        if helix_is_inactive_rank is not None:
-            assert helix_is_inactive_rank.device == helix_position_offsets.device, \
-                f"helix_is_inactive_rank must be on the same device as helix_position_offsets, " \
-                f"got {helix_is_inactive_rank.device} vs {helix_position_offsets.device}"
-
-        mla_tensor_params = [helix_position_offsets, helix_is_inactive_rank]
+        mla_tensor_params = [
+            metadata.helix_position_offsets, metadata.helix_is_inactive_rank
+        ]
 
         torch.ops.trtllm.mla_rope_generation(
             fused_q,

tensorrt_llm/_torch/modules/attention.py

@@ -708,7 +708,7 @@ def __init__(
         dtype: torch.dtype = None,
         dense_bias: Optional[bool] = None,
         config: Optional[ModelConfig] = None,
-        enable_unit_test: bool = False,
+        enable_helix_test: bool = False,
         mapping_with_cp: Optional[Mapping] = None,
         reduce_output: bool = True,
     ):
@@ -733,7 +733,7 @@ def __init__(
             dtype (torch.dtype): The data type.
             dense_bias (bool): Whether to use bias in the output projection layer.
             config (ModelConfig): The model configuration.
-            enable_unit_test (bool): Whether to enable unit test.
+            enable_helix_test (bool): Whether to enable helix unit test.
         """
         super().__init__()
         self.layer_idx = layer_idx
@@ -754,7 +754,7 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
         self.pos_embd_params = pos_embd_params
         self.dense_bias = dense_bias
-        self.enable_unit_test = enable_unit_test
+        self.enable_helix_test = enable_helix_test
         if dense_bias is None:
             self.dense_bias = bias
 
@@ -816,7 +816,7 @@ def __init__(
         self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size -
                                        1) // tp_size
 
-        if self.enable_unit_test:
+        if self.enable_helix_test:
             rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps",
                                    1e-6)
         else:
@@ -1108,8 +1108,8 @@ def _attn_forward_gen(self, attn_backend: AttentionBackend, q: torch.Tensor,
                 v,
                 attn_metadata,
                 softmax_stats_tensor=softmax_stats,
-                helix_position_offsets=position_ids,
-                **kwargs)
+                **kwargs,
+            )
             # this is the post-processing of helix parallel attention,
             # similar to the post-processing of ring attention
             kv_lora_rank = partial_o.shape[-1] // self.num_heads_tp
@@ -1135,7 +1135,7 @@ def _attn_forward_gen(self, attn_backend: AttentionBackend, q: torch.Tensor,
     def create_output(self, hidden_states: torch.Tensor, num_contexts: int):
         num_tokens = hidden_states.shape[0]
         hidden_size = self.o_proj.in_features
-        if self.enable_unit_test and num_contexts > 0:
+        if self.enable_helix_test and num_contexts > 0:
             # note: for testing Helix parallelism, we ensure that the output is
             # large enough for the context phase, but we then cut it again in
             # `forward_context`
@@ -1379,6 +1379,12 @@ def forward_context_default(
             -1,
         )
 
+        if self.enable_helix_test:
+            # While helix parallelism is mainly meant for generation, we set the
+            # helix position offsets for the context phase to get the math right
+            # in test_mla_helix.py.
+            attn_metadata.helix_position_offsets = position_ids
+
         k = torch.empty_like(q).view(-1, self.num_heads_tp, self.qk_head_dim)
         maybe_compiled_copy_(
             k[..., :self.qk_nope_head_dim],
@@ -1388,17 +1394,13 @@ def forward_context_default(
                                                        self.qk_rope_head_dim)
         k = k.view(-1, self.num_heads_tp * self.qk_head_dim)
 
-        helix_position_offsets = position_ids if self.mapping.has_cp_helix(
-        ) else None
-
         attn_output = self.mha.forward(
             q,
             k,
             v,
             attn_metadata,
             attention_input_type=AttentionInputType.context_only,
             latent_cache=latent_cache,
-            helix_position_offsets=helix_position_offsets,
             out_scale=self.out_scale,
             output=output,
         )
@@ -1769,12 +1771,6 @@ def forward_absorption_generation(
             device=q.device,
         )
 
-        helix_position_offsets, helix_is_inactive_rank = None, None
-        if self.mapping.has_cp_helix():
-            helix_position_offsets = position_ids
-            helix_is_inactive_rank = attn_metadata.helix_is_inactive_rank
-            assert helix_position_offsets is not None and helix_is_inactive_rank is not None, "helix_position_offsets and helix_is_inactive_rank must be provided for helix parallelism."
-
         rope_stream = self.aux_stream if not has_fp8_kv_cache else None
         if self.k_b_proj_trans.dtype == torch.bfloat16:
             # [num_heads, num_tokens, self.qk_nope_head_dim]
@@ -1799,8 +1795,7 @@ def forward_absorption_generation(
                     mla_bmm1_scale,
                     mla_bmm2_scale,
                     quant_q_buffer,
-                    helix_position_offsets=helix_position_offsets,
-                    helix_is_inactive_rank=helix_is_inactive_rank),
+                ),
                 self.ln_events[0],
                 self.ln_events[1],
                 rope_stream,
@@ -1829,8 +1824,7 @@ def forward_absorption_generation(
                     mla_bmm1_scale,
                     mla_bmm2_scale,
                     quant_q_buffer,
-                    helix_position_offsets=helix_position_offsets,
-                    helix_is_inactive_rank=helix_is_inactive_rank),
+                ),
                 self.ln_events[0],
                 self.ln_events[1],
                 rope_stream,
@@ -2182,7 +2176,7 @@ def forward(
                               output=attn_output,
                               latent_cache_gen=latent_cache_gen)
 
-        if self.enable_unit_test and self.mapping.has_cp_helix():
+        if self.enable_helix_test and self.mapping.has_cp_helix():
             # note: for allowing testing Helix parallelism, we ensure that
             # the output is compatible with o_proj even in the context phase,
             # thus we cut it to num_heads_tp_cp * v_head_dim

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -488,9 +488,10 @@ def __init__(
         self.py_orig_prompt_len = self.orig_prompt_len
         self.py_max_new_tokens = self.max_new_tokens
         self.py_min_length = self.sampling_config.min_length
+        # `seqlen_this_rank_cp`, `total_input_len_cp`, and `py_helix_is_inactive_rank` are relevant to helix parallelism.
+        self.seqlen_this_rank_cp = self.prompt_len
+        self.total_input_len_cp = self.prompt_len
         self.py_helix_is_inactive_rank = False
-        self.seqlen_this_rank_cp = 0
-        self.total_input_len_cp = 0
         self.py_batch_idx = None
         self.py_draft_pages_allocated = 0
         self.py_rewind_len = 0