updates to test - seqlen 64 works

Author: brb-nv

cpp/tensorrt_llm/kernels/mlaKernels.cu

@@ -354,6 +354,14 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe,
     float const* dequant_scale_kv, float host_bmm1_scale, int32_t const* helix_position_offsets,
     bool const* helix_is_inactive_rank)
 {
+    // if (helix_is_inactive_rank != nullptr)
+    // {
+    //     printf("[applyMLARopeAndAssignQKVKernelGeneration] helix_is_inactive_rank: %p\n", helix_is_inactive_rank);
+    // }
+    // else
+    // {
+    //     printf("[applyMLARopeAndAssignQKVKernelGeneration] helix_is_inactive_rank: nullptr\n");
+    // }
 
     // Constants.
     using VecT = typename VecType<T>::Type;

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -231,10 +231,20 @@ class Runner : public RunnerBase
                 if (mla_helix_position_offsets.has_value())
                 {
                     mla_params.helix_position_offsets = mla_helix_position_offsets->data_ptr<int32_t>();
+                    printf("[AttentionOp] helix_position_offsets: %p\n", mla_params.helix_position_offsets);
+                }
+                else
+                {
+                    printf("[AttentionOp] helix_position_offsets: nullptr\n");
                 }
                 if (mla_helix_is_inactive_rank.has_value())
                 {
-                    mla_params.helix_is_inactive_rank = mla_helix_is_inactive_rank->const_data_ptr<bool>();
+                    printf("[AttentionOp] helix_is_inactive_rank: %p\n", mla_helix_is_inactive_rank->data_ptr<bool>());
+                    mla_params.helix_is_inactive_rank = mla_helix_is_inactive_rank->data_ptr<bool>();
+                }
+                else
+                {
+                    printf("[AttentionOp] helix_is_inactive_rank: nullptr\n");
                 }
             }
             else

cpp/tensorrt_llm/thop/dsv3RopeOp.cpp

@@ -108,6 +108,22 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu
     mla_params.helix_position_offsets = args.helix_position_offsets_ptr;
     mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr;
 
+    if (mla_params.helix_position_offsets != nullptr)
+    {
+        printf("[invokeMLARopeGenerationHelper] helix_position_offsets: %p\n", mla_params.helix_position_offsets);
+    }
+    else
+    {
+        printf("[invokeMLARopeGenerationHelper] helix_position_offsets: nullptr\n");
+    }
+    if (mla_params.helix_is_inactive_rank != nullptr)
+    {
+        printf("[invokeMLARopeGenerationHelper] helix_is_inactive_rank: %p\n", mla_params.helix_is_inactive_rank);
+    }
+    else
+    {
+        printf("[invokeMLARopeGenerationHelper] helix_is_inactive_rank: nullptr\n");
+    }
     tk::invokeMLARopeGeneration<T>(mla_params, kv_cache_buffer, stream);
 }
 

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -296,9 +296,10 @@ 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)
+        print(f"[TrtllmAttention] helix_is_inactive_rank: {self.helix_is_inactive_rank}")
 
         if max_sequence_length > self.rope_params.max_positions:
             self.rope_params.max_positions = max_sequence_length

tensorrt_llm/_torch/modules/attention.py

@@ -1702,6 +1702,14 @@ def forward_absorption_generation(
 
         # Compute helix_position_offsets for helix parallelism.
         helix_position_offsets = position_ids if self.mapping.cp_size > 1 else None
+        # Get helix_is_inactive_rank from attn_metadata for helix parallelism.
+        helix_is_inactive_rank = getattr(attn_metadata, 'helix_is_inactive_rank', None)
+
+        if self.mapping.cp_size > 1:
+            assert helix_position_offsets is not None
+            assert helix_is_inactive_rank is not None
+        print(f"[Attention] helix_position_offsets: {helix_position_offsets}")
+        print(f"[Attention] helix_is_inactive_rank: {helix_is_inactive_rank}")
 
         rope_stream = self.aux_stream if not has_fp8_kv_cache else None
         if self.k_b_proj_trans.dtype == torch.bfloat16:
@@ -1727,7 +1735,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 +1766,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: