updates to test - seqlen 64 works

Author: brb-nv

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -234,7 +234,7 @@ class Runner : public RunnerBase
                 }
                 if (mla_helix_is_inactive_rank.has_value())
                 {
-                    mla_params.helix_is_inactive_rank = mla_helix_is_inactive_rank->const_data_ptr<bool>();
+                    mla_params.helix_is_inactive_rank = mla_helix_is_inactive_rank->data_ptr<bool>();
                 }
             }
             else

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -296,7 +296,7 @@ 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:
+        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)
 

tensorrt_llm/_torch/modules/attention.py

@@ -1700,8 +1700,14 @@ def forward_absorption_generation(
             device=q.device,
         )
 
-        # Compute helix_position_offsets for helix parallelism.
-        helix_position_offsets = position_ids if self.mapping.cp_size > 1 else None
+        if self.mapping.cp_size > 1:
+            helix_position_offsets = position_ids
+            helix_is_inactive_rank = attn_metadata.helix_is_inactive_rank
+            assert helix_position_offsets is not None, "helix_position_offsets must be provided for helix parallelism."
+            assert helix_is_inactive_rank is not None, "helix_is_inactive_rank must be provided for helix parallelism."
+        else:
+            helix_position_offsets = None
+            helix_is_inactive_rank = None
 
         rope_stream = self.aux_stream if not has_fp8_kv_cache else None
         if self.k_b_proj_trans.dtype == torch.bfloat16:
@@ -1727,7 +1733,9 @@ def forward_absorption_generation(
                                                      mla_bmm2_scale,
                                                      quant_q_buffer,
                                                      helix_position_offsets=
-                                                     helix_position_offsets),
+                                                     helix_position_offsets,
+                                                     helix_is_inactive_rank=
+                                                     helix_is_inactive_rank),
                 self.ln_events[0],
                 self.ln_events[1],
                 rope_stream,
@@ -1756,7 +1764,9 @@ def forward_absorption_generation(
                                                      mla_bmm2_scale,
                                                      quant_q_buffer,
                                                      helix_position_offsets=
-                                                     helix_position_offsets),
+                                                     helix_position_offsets,
+                                                     helix_is_inactive_rank=
+                                                     helix_is_inactive_rank),
                 self.ln_events[0],
                 self.ln_events[1],
                 rope_stream,

tests/unittest/_torch/modules/test_mla_helix.py

@@ -125,18 +125,20 @@ def max_position_embeddings(self) -> int:
     Scenario(batch=16, ctx_len=16384),
     Scenario(batch=16, ctx_len=32768),
     Scenario(batch=16, ctx_len=65536),
+    Scenario(batch=1, ctx_len=64),
 ]
 
 # limit the number of test scenarios to avoid taking too long
 test_scenarios = [
     # note: tests with ctx_len=1024 (or less) are currently failing, most likely due to
     # bad numerics especially with bf16. We ignore those tests for now.
-    all_scenarios[2],
-    all_scenarios[5],
-    all_scenarios[12],
-    all_scenarios[15],
-    all_scenarios[21],
-    all_scenarios[22],
+    # all_scenarios[2],
+    # all_scenarios[5],
+    # all_scenarios[12],
+    # all_scenarios[15],
+    # all_scenarios[21],
+    # all_scenarios[22],
+    all_scenarios[-1],
 ]
 
 
@@ -501,9 +503,16 @@ def _run_mla_distributed(
     start = time.time()
 
     for step in range(gen_steps):
+        helix_is_inactive_rank = []
         for req_id in range(scenario.batch):
             kv_cache_manager.impl.add_token(req_id)
-        cache_add = step if rank == world_size - 1 else 0
+            # Assume last rank is active for all gen steps.
+            if rank == world_size - 1:
+                helix_is_inactive_rank.append(False)
+                cache_add = step
+            else:
+                helix_is_inactive_rank.append(True)
+                cache_add = 0
         cached_tokens_per_seq = [ctx_len_per_gpu + cache_add for _ in range(scenario.batch)]
         if step == 0:
             attn_metadata = get_attention_backend("TRTLLM").Metadata(
@@ -519,12 +528,15 @@ def _run_mla_distributed(
                     num_cached_tokens_per_seq=cached_tokens_per_seq,
                 ),
                 enable_context_mla_with_cached_kv=True,
+                helix_is_inactive_rank=torch.tensor(helix_is_inactive_rank, dtype=torch.bool, device="cuda"),
             )
         else:
             attn_metadata.kv_cache_params = KVCacheParams(
                 use_cache=True,
                 num_cached_tokens_per_seq=cached_tokens_per_seq,
             )
+            attn_metadata.helix_is_inactive_rank = torch.tensor(
+                helix_is_inactive_rank, dtype=torch.bool, device="cuda")
         attn_metadata.prepare()
         extra_attrs["attention_metadata"] = weakref.ref(attn_metadata)
         if not use_cuda_graph: