save changes to get seqlen 64 working

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

@@ -528,14 +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=helix_is_inactive_rank,
+                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 = helix_is_inactive_rank
+            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: