cpp/tensorrt_llm/kernels/mlaKernels.cu (1)

1-3: Update copyright year to 2025

The copyright header should include the current year (2025) according to coding guidelines.

 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
  *

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

681-681: Remove commented debug print statement.

This appears to be leftover debugging code.

-        # print("prepare TrtllmAttentionMetadata")
         extra_attrs = get_model_extra_attrs()

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T04:09:12.904Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx} : All class templates, function templates, class template member functions and class template static members shall be instantiated at least once.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T04:09:12.904Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx,cu,py} : All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. The block should be prepended to the top of all files, including .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

cpp/tensorrt_llm/kernels/mlaKernels.cu (1)

1-15: Update copyright year to 2025.

According to the coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header that includes the current year.

 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
  *

tensorrt_llm/_torch/attention_backend/trtllm.py (2)

681-681: Remove commented-out debug print statement.

-        # print("prepare TrtllmAttentionMetadata")

---

201-201: Line exceeds maximum length of 120 characters.

-            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU.
+            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and 
+                                         generation requests, with shape (2) on CPU.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T07:34:42.734Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx} : All class templates, function templates, class template member functions and class template static members shall be instantiated at least once.

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

201-201: Line exceeds maximum length limit.

Line 201 is 148 characters long, exceeding the 120 character limit.

Consider breaking this line for better readability:

-            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU.
+            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and 
+                                         generation requests, with shape (2) on CPU.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T07:34:42.734Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx,cu,py} : All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T07:34:42.734Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx} : All class templates, function templates, class template member functions and class template static members shall be instantiated at least once.

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

201-201: Consider breaking long documentation line for better readability.

The documentation line exceeds the 120-character limit (148 characters). Consider reformatting:

-            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU.
+            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests 
+                and generation requests, with shape (2) on CPU.

cpp/tensorrt_llm/thop/attentionOp.cpp (1)

443-444: Clarify the validation logic for better readability.

The validation condition is_mla_enable || is_fused_qkv with the error message "Only fused QKV is supported for non-MLA attention now" could be confusing. Consider making the intent clearer:

-TLLM_CHECK_WITH_INFO(is_mla_enable || is_fused_qkv, "Only fused QKV is supported for non-MLA attention now");
+TLLM_CHECK_WITH_INFO(is_mla_enable || is_fused_qkv, "Either MLA must be enabled or fused QKV must be used");

Or restructure the logic to be more explicit about the requirements for each case.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T07:34:42.734Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx,cu,py} : All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-01T07:34:42.734Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx} : All class templates, function templates, class template member functions and class template static members shall be instantiated at least once.

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

tensorrt_llm/_torch/attention_backend/trtllm.py (2)

70-70: Consider using Optional typing for consistency.

The softmax_stats_tensor attribute should use Optional[torch.Tensor] typing to be consistent with other optional tensor attributes in the class.

-    softmax_stats_tensor: Optional[torch.Tensor]
+    softmax_stats_tensor: Optional[torch.Tensor]

---

201-201: Fix line length violation.

This line exceeds the 120-character limit specified in the coding guidelines.

The line should be wrapped or shortened to comply with the coding standards. Consider breaking the parameter documentation across multiple lines or shortening the description.

cpp/tensorrt_llm/thop/attentionOp.cpp (1)

191-191: Remove unnecessary blank line.

There's an extra blank line that should be removed for consistency with the codebase formatting.

                mla_params.latent_cache = static_cast<T const*>(latent_cache->data_ptr());
-

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-04T02:12:17.550Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx,cu,py} : All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-04T02:12:17.550Z
Learning: Applies to **/*.{cpp,h,hpp,cc,cxx} : All class templates, function templates, class template member functions and class template static members shall be instantiated at least once.

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu (1)

2-2: Update copyright header to include current year.

The copyright header shows 2022-2024, but the coding guidelines require including the current year (2025).

Apply this diff to fix the copyright header:

- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

157-157: Fix line length violation in docstring.

The addition of total_kv_lens parameter is correct, but the docstring on line 201 exceeds the 120 character limit.

Apply this diff to fix the line length:

-            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU.
+            total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context 
+                requests and generation requests, with shape (2) on CPU.

Also applies to: 201-201

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.

cpp/tensorrt_llm/common/attentionOp.cpp (2)

770-772: Under-allocation: include mFP8ContextMLA in BMM scale workspace sizes

fmha_bmm1_scale_size and fmha_bmm2_scale_size only account for mFP8ContextFMHA. FP8 context MLA also consumes these buffers (you pass bmm1/bmm2 scales into params.mla_param and use them in invokeMLAContextFp8Quantize). This can under-allocate the context workspace and corrupt subsequent allocations.

Apply this fix:

-    size_t const fmha_bmm1_scale_size = mFP8ContextFMHA ? sizeof(float) * 2 : 0;
-    size_t const fmha_bmm2_scale_size = mFP8ContextFMHA ? sizeof(float) : 0;
+    size_t const fmha_bmm1_scale_size = (mFP8ContextFMHA || mFP8ContextMLA) ? sizeof(float) * 2 : 0;
+    size_t const fmha_bmm2_scale_size = (mFP8ContextFMHA || mFP8ContextMLA) ? sizeof(float) : 0;

---

1655-1683: Null q/k/v pointers when latent_cache is absent break FP8 quantization

invokeMLAContextFp8Quantize() requires params.mla_param->q_buf/k_buf/v_buf to be set. In the default MLA context path (latent_cache == nullptr), these pointers are never initialized before the quant call, causing TLLM_CHECK failures or undefined behavior.

Set q/k/v pointers for the default path prior to quantization:

             params.mla_param->host_bmm1_scale
                 = 1 / (mQScaling * sqrt((float) (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim)));
+            // Ensure q/k/v are set when latent_cache is absent (default separate Q/K/V input).
+            if (params.mla_param->q_buf == nullptr)
+            {
+                // attention_input already points to Q in context FMHA path.
+                params.mla_param->q_buf = const_cast<T*>(attention_input);
+            }
+            if (params.mla_param->k_buf == nullptr)
+            {
+                // k_ptr comes from THOP attention binding for separate K input.
+                params.mla_param->k_buf = const_cast<T*>(params.k_ptr);
+            }
+            if (params.mla_param->v_buf == nullptr)
+            {
+                // v_buf is const in MlaParams.
+                params.mla_param->v_buf = params.v_ptr;
+            }
             if (params.mla_param->latent_cache != nullptr)
             {
                 // compute RoPE and set compressed_kv + k_pe by invokeMLARopeContext if latent_cache is not nullptr
                 invokeMLARopeContext<T, KVCacheBuffer>(*params.mla_param, kv_cache_buffer, stream);
             }
             if (mFP8ContextMLA)
             {
                 invokeMLAContextFp8Quantize(*params.mla_param, params.total_kv_len, stream);
             }

tensorrt_llm/_torch/modules/attention.py (1)

1164-1172: Syntax error in type annotation for latent_cache

The argument annotation is split across lines, producing invalid Python syntax.

Apply this fix:

-        latent_cache: torch.
-        Tensor,  # compressed_kv + k_pe [context_tokens, 1, lora_size + rope_size]
+        latent_cache: torch.Tensor,  # compressed_kv + k_pe [context_tokens, 1, lora_size + rope_size]

cpp/tensorrt_llm/thop/attentionOp.cpp (3)

140-147: Guard against non-contiguous/invalid Q layout before taking raw pointer

The enqueue path assumes the inner-most dimension is contiguous. Add light shape/stride checks to fail-fast when Q is not laid out as expected.

Apply this diff:

-        auto stream = at::cuda::getCurrentCUDAStream(qkv_or_q.get_device());
-        T* attention_input = static_cast<T*>(qkv_or_q.slice(0, token_offset).data_ptr());
+        auto stream = at::cuda::getCurrentCUDAStream(qkv_or_q.get_device());
+        TORCH_CHECK(qkv_or_q.dim() >= 2, "qkv_or_q must be at least 2D: [tokens, hidden]");
+        TORCH_CHECK(qkv_or_q.strides()[qkv_or_q.dim() - 1] == 1, "qkv_or_q inner-most dimension must be contiguous");
+        T* attention_input = static_cast<T*>(qkv_or_q.slice(0, token_offset).data_ptr());

---

652-655: Fix potential null deref when beam_width > 1 but cache_indirection is absent

This uses cache_indirection.value() without has_value() guard. Prefer the same safe pattern you already use inside Runner::run.

Apply this diff:

-    int32_t const max_attention_window_size
-        = beam_width == 1 ? attention_window_size : cache_indirection.value().size(2);
+    int32_t const max_attention_window_size = beam_width == 1
+        ? attention_window_size
+        : (cache_indirection.has_value() ? cache_indirection.value().size(2) : attention_window_size);

---

576-601: Missing has_value() checks for MLA params (q_lora_rank/kv_lora_rank/...): will crash if unset

You call .value() on several optionals without verifying has_value(). Add checks and provide a clear error message.

Apply this diff:

     if (is_mla_enable)
     {
         // MLA does not support NVFP4 output yet.
         TLLM_CHECK(!is_fp4_out);
 
         TLLM_CHECK(host_kv_cache_pool_mapping.has_value());
         int32_t const layer_num = host_kv_cache_pool_mapping.value().size(0);
 
+        TLLM_CHECK(q_lora_rank.has_value() && kv_lora_rank.has_value() && qk_nope_head_dim.has_value()
+                && qk_rope_head_dim.has_value() && v_head_dim.has_value(),
+            "MLA requires q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, and v_head_dim");
+
         op->mIsMLAEnabled = true;
         op->mMLAParams = {static_cast<int>(q_lora_rank.value()), static_cast<int>(kv_lora_rank.value()),
             static_cast<int>(qk_nope_head_dim.value()), static_cast<int>(qk_rope_head_dim.value()),
             static_cast<int>(v_head_dim.value()), static_cast<int>(predicted_tokens_per_seq),
             static_cast<int>(layer_num)};

cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h (8)

1850-1851: Document mask/layout fields for readability (avoid magic numbers)

The 0/1 (mask) and 3 (layout) columns correspond to AttentionMaskType and AttentionInputLayout respectively. Inline annotations improve clarity and reduce regressions.

Apply:

-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90_kernel", 213248, 384, 64, 0, 3, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90_kernel", 213248, 384, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90},
-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 3, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90},

If this header is generated, consider emitting enum names (e.g., static_cast(AttentionInputLayout::SEPARATE_Q_K_V)) in the generator.

---

4333-4334: SM100 NL-tiled: annotate mask/layout and confirm shared wrapper + implementation

• Add inline comments for mask/layout.
• Both rows share run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled. Confirm this is intended for both mask kinds.
• Ensure the wrapper’s definition exists (see extern at Line 1083).

-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_100, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_kernel_nl_tiled", 81920, 128, 64, 0, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_100, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_kernel_nl_tiled", 81920, 128, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled},
-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_100, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_sm100_kernel_nl_tiled", 81920, 128, 64, 1, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_100, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_sm100_kernel_nl_tiled", 81920, 128, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled},

---

4743-4744: SM120 NL-tiled BF16: annotate mask/layout and ensure wrapper definition exists

• Add inline mask/layout comments.
• Confirm the wrapper run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled (extern at Line 1224) is implemented and linked.
• Validate that sharing the same wrapper for causal and non-causal is intentional.

-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_kernel_nl_tiled", 81920, 128, 64, 0, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_kernel_nl_tiled", 81920, 128, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled},
-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_sm120_kernel_nl_tiled", 81920, 128, 64, 1, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_sm120_kernel_nl_tiled", 81920, 128, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled},

---

212-216: SM90 separate Q-K-V wrappers: ensure implementations exist to avoid link errors

The new externs for SM90 separate Q-K-V context paths look correct, but prior verification found no corresponding definitions for these symbols, which will lead to linker failures if still missing:

• run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90 (Line 212)
• run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90 (Line 216)

Please add their definitions in the appropriate .cu/.cpp unit (consistent with other run_fmha_v2_* wrappers) and ensure runtime dispatch reaches them.

Run this to confirm definitions exist (expects at least one non-header definition per symbol):

#!/bin/bash
set -euo pipefail
syms=(
  run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90
  run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90
)
for s in "${syms[@]}"; do
  echo ">>> Checking definition for $s"
  if ! rg -nP -C2 -g '!**/*.h' -g '*.{cu,cpp,cc}' "^[[:space:]]*(?:__host__\\s+|__global__\\s+)?void[[:space:]]+${s}[[:space:]]*\\("; then
    echo "!! Missing definition for: $s"
    exit 1
  fi
done
echo "All SM90 S_q_k_v wrappers found."

---

1083-1083: SM100 NL-tiled separate Q-K-V wrapper: definition likely missing

Extern added:

• run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled

A previous review found no implementation. Add the definition and ensure both non-causal and causal registry entries route correctly to this wrapper.

#!/bin/bash
set -euo pipefail
sym='run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm100_nl_tiled'
rg -nP -C2 -g '!**/*.h' -g '*.{cu,cpp,cc}' "^[[:space:]]*(?:__host__\\s+|__global__\\s+)?void[[:space:]]+${sym}[[:space:]]*\\(" || {
  echo "Missing definition for: $sym"
  exit 1
}

---

1224-1224: SM120 NL-tiled BF16 separate Q-K-V wrapper: add implementation

Extern added:

• run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled

This was previously flagged as missing. Provide the .cu/.cpp implementation and include it in the build.

#!/bin/bash
set -euo pipefail
sym='run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_sm120_nl_tiled'
rg -nP -C2 -g '!**/*.h' -g '*.{cu,cpp,cc}' "^[[:space:]]*(?:__host__\\s+|__global__\\s+)?void[[:space:]]+${sym}[[:space:]]*\\(" || {
  echo "Missing definition for: $sym"
  exit 1
}

---

1254-1258: SM120 E4M3 64×64 separate Q-K-V wrappers: add missing implementations

Externs added:

• run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled (Line 1254)
• run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled (Line 1258)

Previous checks found no definitions. Implement both wrappers and ensure they are linked.

#!/bin/bash
set -euo pipefail
syms=(
  run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled
  run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled
)
for s in "${syms[@]}"; do
  echo ">>> Checking definition for $s"
  rg -nP -C2 -g '!**/*.h' -g '*.{cu,cpp,cc}' "^[[:space:]]*(?:__host__\\s+|__global__\\s+)?void[[:space:]]+${s}[[:space:]]*\\(" || {
    echo "Missing definition for: $s"
    exit 1
  }
done
echo "All SM120 E4M3 S_q_k_v wrappers found."

---

4833-4834: SM120 E4M3 separate Q-K-V rows: annotate mask/layout and implement wrappers

• Add inline comments for mask/layout to avoid magic numbers.
• Ensure both E4M3 wrappers are implemented:
  • run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled (Line 4833/4834)
  • run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled (Line 4842/4843)

-{ DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_kernel_nl_tiled", 32768, 128, 64, 0, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled},
+{ DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_kernel_nl_tiled", 32768, 128, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled},
-{ DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_causal_sm120_kernel_nl_tiled", 32768, 128, 64, 1, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled},
+{ DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_causal_sm120_kernel_nl_tiled", 32768, 128, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled},
-{ DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_kernel_nl_tiled", 32768, 128, 64, 0, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled},
+{ DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_kernel_nl_tiled", 32768, 128, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled},
-{ DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_causal_output_bf16_sm120_kernel_nl_tiled", 32768, 128, 64, 1, 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled},
+{ DATA_TYPE_E4M3, DATA_TYPE_BF16, 0, 64, 64, 192, 128, 0, 0, 0, kSM_120, nullptr, 0, "fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_causal_output_bf16_sm120_kernel_nl_tiled", 32768, 128, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, false, true, true, true, false, true, run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled},

Quick definition check:

#!/bin/bash
set -euo pipefail
syms=(
  run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_sm120_nl_tiled
  run_fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_k_v_192x128_output_bf16_sm120_nl_tiled
)
missing=0
for s in "${syms[@]}"; do
  if ! rg -nP -C2 -g '!**/*.h' -g '*.{cu,cpp,cc}' "^[[:space:]]*(?:__host__\\s+|__global__\\s+)?void[[:space:]]+${s}[[:space:]]*\\(" ; then
    echo "Missing definition for: $s"
    missing=1
  fi
done
exit $missing

Also applies to: 4842-4843

cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h (1)

569-576: Fix KV head-indexing for GQA/MQA in Gmem_tile_q_k_v

The ctor currently sets idx = binfo.bidh for all tensors. That is incorrect when numKvHeads < numHeads (GQA/MQA) for K/V: the K/V head index must be kv_head_id = bidh / headsPerKv. Mirror the mapping used elsewhere.

Apply this patch:

-        // The row offset in the batched GEMM, including the sequence offset.
-        int64_t row_offset = (int64_t) (row + cta_row_offset + seq_offset) * params_q_k_v_stride_in_bytes_;
-        // Add the head index.
-        int64_t idx = binfo.bidh;
+        // The row offset in the batched GEMM, including the sequence offset.
+        int64_t row_offset = (int64_t) (row + cta_row_offset + seq_offset) * params_q_k_v_stride_in_bytes_;
+        // Add the head index (map Q head -> KV head for GQA/MQA).
+        int64_t idx;
+        if (qkv_offset == 0)
+        {
+            // Q tensor: one head per Q head.
+            idx = binfo.bidh;
+        }
+        else
+        {
+            // K or V tensor: map Q head id to KV head id.
+            int const num_heads = params.h;
+            int const num_kv_heads = max(1, params.h_kv);
+            int const heads_per_kv = max(1, num_heads / num_kv_heads);
+            int const kv_head_id = binfo.bidh / heads_per_kv;
+            idx = kv_head_id;
+        }

cpp/tensorrt_llm/kernels/mlaKernels.cu (1)

944-962: Quantize kernel launch: guard dims and grid before launch

The quantize kernel is instantiated for QK_NOPE=128, QK_ROPE=64, V=128. Add runtime checks to fail clearly if params meta dims differ, and validate grid dims against device limits (prior suggestion).

Apply this at the top of invokeMLAContextFp8Quantize():

     TLLM_LOG_DEBUG("MLA RoPE Context: Quantizing separate qkv to FP8");

     if (params.acc_q_len > 0)
     {
         constexpr int threads_per_block = 384;
         dim3 grid(int(tensorrt_llm::common::divUp(total_kv_len, 48)), 1, params.head_num);
+        // Sanity check the expected head dims for this specialization.
+        TLLM_CHECK_WITH_INFO(params.meta.qk_nope_head_dim == 128
+                && params.meta.qk_rope_head_dim == 64
+                && params.meta.v_head_dim == 128,
+            "MLA FP8 quantization currently specialized for (qk_nope, qk_rope, v) = (128, 64, 128); got (%d, %d, %d).",
+            params.meta.qk_nope_head_dim, params.meta.qk_rope_head_dim, params.meta.v_head_dim);
+        // Validate grid dimensions to avoid exceeding device limits.
+        int maxGridDimX = 0, maxGridDimZ = 0;
+        cudaDeviceGetAttribute(&maxGridDimX, cudaDevAttrMaxGridDimX, 0);
+        cudaDeviceGetAttribute(&maxGridDimZ, cudaDevAttrMaxGridDimZ, 0);
+        TLLM_CHECK_WITH_INFO(grid.x <= maxGridDimX && grid.z <= maxGridDimZ,
+            "Quantize kernel grid dims exceed device limits: (%d,%d,%d) vs max (%d,*,%d).",
+            grid.x, grid.y, grid.z, maxGridDimX, maxGridDimZ);

cpp/kernels/fmha_v2/fmha_test.py (2)

168-169: Include SM121 in FP8 context MLA gating

Allow FP8 MLA tests to run on SM121 as well. This avoids false skips on 12.1 parts.

Apply this diff:

-    if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 120:
-        pytest.skip("FP8 MLAs are only supported on sm120 currently.")
+    if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version not in (120, 121):
+        pytest.skip("FP8 MLAs are only supported on sm120/sm121 currently.")

---

213-214: Include SM121 in FP8 generation MLA gating

Mirror the context change to avoid false skips on 12.1.

Apply this diff:

-    if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 120:
-        pytest.skip("FP8 MLAs are only supported on sm120 currently.")
+    if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version not in (120, 121):
+        pytest.skip("FP8 MLAs are only supported on sm120/sm121 currently.")

cpp/tensorrt_llm/thop/attentionOp.cpp (1)

168-185: Also validate K/V dtype and device to match Q (carry-over from prior review)

The context-MLA path still lacks dtype/device checks for k and v. This can cause UB when kernels reinterpret memory.

Apply this diff to add the missing validations:

                 TORCH_CHECK(k.has_value());
                 TORCH_CHECK(v.has_value());
                 TORCH_CHECK(k->dim() == 2);
                 TORCH_CHECK(v->dim() == 2);
                 TORCH_CHECK(k->strides()[1] == 1);
                 TORCH_CHECK(v->strides()[1] == 1);
+                TORCH_CHECK(k->scalar_type() == qkv_or_q.scalar_type(), "K dtype must match Q dtype");
+                TORCH_CHECK(v->scalar_type() == qkv_or_q.scalar_type(), "V dtype must match Q dtype");
+                TORCH_CHECK(k->scalar_type() == v->scalar_type(), "K and V dtypes must match");
+                TORCH_CHECK(k->device() == qkv_or_q.device() && v->device() == qkv_or_q.device(),
+                    "K/V must be on the same device as Q");

cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h (1)

1860-1861: SM90 softmax rows: annotate mask/layout and confirm shared wrapper is intended

• Add inline comments for mask/layout to avoid magic numbers.
• Both non-causal and causal rows reference the same run wrapper. Verify the wrapper branches on mAttentionMaskType correctly for CAUSAL vs PADDING.

-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90_kernel", 213248, 384, 64, 0, 3, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90_kernel", 213248, 384, 64, /* mask=PADDING */ 0, /* layout=SEPARATE_Q_K_V */ 3, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90},
-{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_softmax_tma_ws_sm90_kernel", 213248, 384, 64, 1, 3, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90},
+{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_softmax_tma_ws_sm90_kernel", 213248, 384, 64, /* mask=CAUSAL */ 1, /* layout=SEPARATE_Q_K_V */ 3, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_softmax_tma_ws_sm90},

cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu (5)

149-161: Remove duplicate early-return check for merge_op.

merge_op_val == 0 is checked twice back-to-back. The second check is redundant and can be removed to reduce instruction count and warp divergence.

Apply this diff:

     int64_t merge_op_val = merge_op[batch_idx];
     if (merge_op_val == 0)
     {
         return; // skip this batch
     }

-    size_t const head_dim_vec_idx = (threadIdx.x % KT::kVecPerHead);
-    size_t const head_dim_idx = head_dim_vec_idx * KT::kElemPerThread;
-
-    if (merge_op_val == 0)
-    {
-        return; // skip this batch
-    }
+    size_t const head_dim_vec_idx = (threadIdx.x % KT::kVecPerHead);
+    size_t const head_dim_idx = head_dim_vec_idx * KT::kElemPerThread;

---

162-185: Use 64-bit/size_t for global offsets to avoid potential overflow on long sequences.

Offsets derived from cu_q_seq_len (int64_t) are narrowed to int and used in pointer indexing. For large concatenated sequences, this can overflow 32-bit. Prefer int64_t/size_t for offsets.

If your workloads guarantee <2^31 elements, this is low risk; otherwise, consider the change below.

-    int const curr_q_len = static_cast<int>(cu_q_seq_len[batch_idx + 1] - cu_q_seq_len[batch_idx]);
-    int const global_q_offset = cu_q_seq_len[batch_idx];
+    int64_t const curr_q_len = cu_q_seq_len[batch_idx + 1] - cu_q_seq_len[batch_idx];
+    int64_t const global_q_offset = cu_q_seq_len[batch_idx];

     for (int local_token_idx = (threadIdx.x / KT::kVecPerHead) + blockIdx.x * KT::kTokenPerBlock;
          local_token_idx < curr_q_len; local_token_idx += gridDim.x * KT::kTokenPerBlock)
     {
         // load softmax stat
-        int const global_softmax_stats_offset = (global_q_offset + local_token_idx) * num_heads + head_idx;
+        size_t const global_softmax_stats_offset
+            = static_cast<size_t>(global_q_offset + local_token_idx) * static_cast<size_t>(num_heads)
+            + static_cast<size_t>(head_idx);
         float2 curr_stats = curr_softmax_stats[global_softmax_stats_offset];
         float2 pre_stats = pre_softmax_stats[global_softmax_stats_offset];

         // load attn
         typename KT::VecReader pre_attn_reader{};
         typename KT::VecReader curr_attn_reader{};
         typename KT::VecReader merged_attn_reader{};

-        int const global_attn_offset
-            = (global_q_offset + local_token_idx) * num_heads * head_size + head_idx * head_size;
+        size_t const global_attn_offset = static_cast<size_t>(global_q_offset + local_token_idx)
+            * static_cast<size_t>(num_heads) * static_cast<size_t>(head_size)
+            + static_cast<size_t>(head_idx) * static_cast<size_t>(head_size);

---

201-203: Use device fast-math for exp to avoid double promotion and improve throughput.

In device code, std::exp may promote to double and is slower. Use __expf (or expf) with float inputs for better performance and consistency with fmaxf.

-            float pre_shift = std::exp(pre_stats.x - merged_stats.x);
-            float curr_shift = std::exp(curr_stats.x - merged_stats.x);
+            float pre_shift = __expf(pre_stats.x - merged_stats.x);
+            float curr_shift = __expf(curr_stats.x - merged_stats.x);

Note: If you rely on precise IEEE semantics, prefer expf; __expf trades a bit of precision for speed.

---

246-261: Type kvSrc as TCache and make it const for correctness and clarity.

The cache storage element type is TCache, which can differ from T (e.g., FP8). Casting to T* is confusing and may mislead future refactors. Use TCache const* and keep the vectorized load const-correct.

-            auto* kvSrc = reinterpret_cast<T*>(kv_cache.getKBlockPtr(batch_idx, token_idx_in_kv_cache));
+            auto const* kvSrc
+                = reinterpret_cast<TCache const*>(kv_cache.getKBlockPtr(batch_idx, token_idx_in_kv_cache));
             // head_idx === 0
             auto kvBlockIdx
                 = kv_cache.getKVLocalIdx(token_idx_in_kv_cache, 0, KT::kVecPerHead, static_cast<int>(head_dim_vec_idx));
-            auto ld_data = (reinterpret_cast<typename KT::VecT*>(kvSrc))[kvBlockIdx];
+            auto ld_data = (reinterpret_cast<typename KT::VecT const*>(kvSrc))[kvBlockIdx];

---

138-221: Minor: consider aligning math intrinsics and documenting head_size assumptions.

• You use fmaxf and std::exp together; after switching to __expf/expf, math intrinsics will be consistent.
• Head size is hardcoded in traits (128) and checked at runtime; a brief comment in the kernel noting that KT::kHeadSize must be 128 would help future maintainers.

cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_40_sm89.cubin.cpp (1)

1-4: Header-check exception for LFS pointer files

This .cpp is a Git LFS pointer, so adding a copyright header would break the pointer format. Ensure the header/lint rules exclude cubin pointer files under cubin/ to avoid false positives in CI.

If needed, I can propose an exclude pattern update for your header-check tooling (e.g., regex to skip paths matching **/cubin/*.cubin.cpp).

cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp (1)

155-158: SEPARATE_Q_K_V runtime path: ensure q/k/v pointers are present

When using the separate Q/K/V layout, qPtr/kPtr/vPtr must be non-null. Add a defensive check to fail fast if callers forget to set them.

Apply this diff locally within FmhaDispatcher::run:

-        else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_Q_K_V)
+        else if (mFixedParams.attentionInputLayout == AttentionInputLayout::SEPARATE_Q_K_V)
         {
             qkvLayout = kernels::QkvLayout::SeparateQkv;
+            TLLM_CHECK_WITH_INFO(runnerParams.qPtr && runnerParams.kPtr && runnerParams.vPtr,
+                "SEPARATE_Q_K_V layout requires non-null qPtr/kPtr/vPtr.");
         }

cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h (1)

634-640: Preserve move_col(steps) in all Q/K/V tile variants

To keep the API symmetric with rewind_col(int steps) and consistent across sibling tiles, add a default steps parameter to every move_col that currently lacks it.

Affected locations in
cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h:

• Gmem_tile_q_k_v::move_col (around line 635)
• Gmem_tile_contiguous_kv::move_col (around line 1190)

Proposed patch:

--- a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h
+++ b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h
@@ -633,7 +633,9 @@ struct Gmem_tile_q_k_v {
-    inline __device__ void move_col()
+    inline __device__ void move_col(int const steps = 1)
     {
-        q_k_v_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8);
-        // Update col_in_bytes_ to ensure load predicates work
-        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG;
+        q_k_v_ptr_    += (int64_t) COLS * (BITS_PER_ELEMENT / 8) * steps;
+        // Update col_in_bytes_ to ensure load predicates work
+        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG       * steps;
     }
@@ -1188,7 +1188,9 @@ struct Gmem_tile_contiguous_kv {
-    inline __device__ void move_col()
+    inline __device__ void move_col(int const steps = 1)
     {
-        kv_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8);
-        // Update col_in_bytes_ to ensure load predicates work
-        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG;
+        kv_ptr_       += (int64_t) COLS * (BITS_PER_ELEMENT / 8) * steps;
+        // Update col_in_bytes_ to ensure load predicates work
+        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG       * steps;
     }

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

205-206: Docstring line exceeds 120 chars

Break the host_total_kv_lens docstring line to satisfy the style limit flagged by static analysis.

-            host_total_kv_lens (torch.Tensor): The tensor to store the total KV lens for context requests and generation requests, with shape (2) on CPU.
+            host_total_kv_lens (torch.Tensor): The tensor to store the total KV lens for
+                context requests and generation requests, with shape (2) on CPU.

cpp/tensorrt_llm/kernels/mlaKernels.cu (1)

820-921: QuantizeCopyInputToFp8Kernel: consider generalizing grid.x token grouping

grid.x uses divUp(total_kv_len, 48). 48 is the LCM of tokens-per-block for QK and V for BF16/FP16, but not for FP32; it still works but may underutilize SMs on some shapes. Consider computing grid.x based on the larger of QK_TOKENS_PER_BLOCK and V_TOKENS_PER_BLOCK (or their LCM) per specialization to balance occupancy and iteration count.

tests/unittest/_torch/test_attention_mla.py (2)

671-673: Broaden enable_flash_mla gating beyond (9, 0)

Flash MLA appears to be available on newer arch as well. Consider enabling on SM100/SM120 too to avoid false negatives on capable hardware.

-                enable_flash_mla=torch.cuda.get_device_capability() == (9, 0),
+                enable_flash_mla=torch.cuda.get_device_capability() in {(9, 0), (10, 0), (12, 0)},

---

750-757: Line length in diagnostic prints exceeds 120

Split long f-strings for readability and to satisfy linters.

-            print(
-                f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}"
-            )
-            print(
-                f"Reference output mean: {ref_result.abs().mean().item()}, max: {ref_result.abs().max().item()}"
-            )
-            print(
-                f"Difference mean: {(result - ref_result).abs().mean().item()}, max: {(result - ref_result).abs().max().item()}"
-            )
+            print(f"{backend_name} output mean: {result.abs().mean().item()}, "
+                  f"max: {result.abs().max().item()}")
+            print(f"Reference output mean: {ref_result.abs().mean().item()}, "
+                  f"max: {ref_result.abs().max().item()}")
+            print(f"Difference mean: {(result - ref_result).abs().mean().item()}, "
+                  f"max: {(result - ref_result).abs().max().item()}")

cpp/kernels/fmha_v2/setup.py (2)

3649-3660: Clarify gating intent for MLA vs normal return-softmax cases

The new skip_mla_combination fixes earlier over-enumeration. Minor readability tweak: make MLA layout intent explicit and avoid double-negatives.

Apply this diff:

-        # for normal attention, we only need contiguous kv as input layout when returning softmax.
-        skip_combination = return_softmax and input_layout != InputLayout.CONTIGUOUS_Q_KV
-        # for context mla, we need separate qkv as input layout when returning softmax.
-        skip_mla_combination = return_softmax and input_layout != InputLayout.SEPARATE_Q_K_V
+        # For normal attention w/ return_softmax, only CONTIGUOUS_Q_KV is needed.
+        skip_combination = return_softmax and input_layout != InputLayout.CONTIGUOUS_Q_KV
+        # For context MLA w/ return_softmax, only SEPARATE_Q_K_V is meaningful.
+        is_mla_layout = input_layout == InputLayout.SEPARATE_Q_K_V
+        skip_mla_combination = return_softmax and not is_mla_layout

---

6367-6391: Fix E712 comparisons to True/False in kernel filter (ruff)

Pythonic conditionals improve readability and satisfy lints. Replace explicit True/False comparisons.

Apply this diff:

-                  # Deepseek MLA (generation 576/512 paged)
-                  or (kspec.sm            in [90, 100, 120]
+                  # Deepseek MLA (generation 576/512 paged)
+                  or (kspec.sm            in [90, 100, 120]
                   and kspec.dtype         in ['bf16', 'e4m3_fp32']
                   and kspec.head_size     == 576
                   and kspec.head_size_v   == 512
                   and kspec.input_layout == InputLayout.Q_PAGED_KV
                   and kspec.sage_block_sizes is None
                   and kspec.version       == 2
-                  and kspec.cross_mha     == False
-                  and kspec.flash_attention == True
-                  and kspec.warp_specialization == False
-                  and kspec.tiled == True)
+                  and not kspec.cross_mha
+                  and kspec.flash_attention
+                  and not kspec.warp_specialization
+                  and kspec.tiled)
@@
-                  # Deepseek MLA (context 192/128 separate-q-k-v)
-                  or (kspec.sm            in [90, 100, 120]
-                  and kspec.dtype         in ['bf16', 'e4m3_fp32']
+                  # Deepseek MLA (context 192/128 separate-q-k-v)
+                  or (kspec.sm            in [90, 100, 120]
+                  and kspec.dtype         in ['bf16', 'e4m3_fp32']
                   and kspec.head_size     == 192
                   and kspec.head_size_v   == 128
                   and kspec.input_layout == InputLayout.SEPARATE_Q_K_V
                   and kspec.sage_block_sizes is None
                   and kspec.version       == 2
-                  and kspec.cross_mha     == False
-                  and kspec.flash_attention == True
-                  and ((kspec.warp_specialization == True and kspec.alibi == False)   # sm90
-                    or (kspec.warp_specialization == False and kspec.tiled == True))  # non-sm90
-                  and kspec.enable_attn_logit_softcapping == False)
+                  and not kspec.cross_mha
+                  and kspec.flash_attention
+                  and ((kspec.warp_specialization and not kspec.alibi)      # sm90
+                    or (not kspec.warp_specialization and kspec.tiled))     # non-sm90
+                  and not kspec.enable_attn_logit_softcapping)

cpp/tensorrt_llm/common/attentionOp.cpp (1)

2628-2639: Using SEPARATE_Q_K_V and headSizeQkNope for context MLA is appropriate

Forcing AttentionInputLayout::SEPARATE_Q_K_V for context MLA and plumbed headSizeQkNope matches the kernel-param needs introduced in fmha runners. The headSizeV note is documented; consider an assertion or static check in debug builds if future models diverge.

cpp/tensorrt_llm/kernels/mlaKernels.h (1)

54-104: q/k mutability justified; consider typing FP8 quant buffers

You made q_buf and k_buf mutable because RoPE modifies them in-place; v_buf is const, which is correct. Consider typing the FP8 quant buffers as __nv_fp8_e4m3* instead of void* to avoid repeated casts and improve readability across declarations and call sites.

If changing the header type is too intrusive now, at least add brief comments in callers where casts happen to reduce future confusion.

cpp/tensorrt_llm/thop/attentionOp.h (1)

37-61: Minor consistency: prefer one optional type

This header mixes std::optionaltorch::Tensor and torch::optionaltorch::Tensor. Both work, but pick one style for consistency (project convention leans towards std::optional in C++17 code).

cpp/tensorrt_llm/thop/attentionOp.cpp (6)

67-85: Run() signature explosion: consider bundling args to improve safety and readability

The expanded run signature is getting unwieldy and error-prone. Consider introducing a small POD/struct (e.g., RunArgs) to group the many tensors and scalars. This reduces callsite mistakes, eases evolution, and matches the pattern you already use for enqueue params.

---

668-681: Workspace device selection is correct, but add capacity check when external buffer is provided

If the provided workspace tensor’s dtype/device mismatches q’s device or is not byte, implicit assumptions could break. Consider validating dtype and device.

Apply this diff:

     if (workspace_.has_value())
     {
         if (workspace_.value().numel() < workspace_size)
         {
             TLLM_LOG_WARNING("Attention workspace size is not enough, increase the size from %ld bytes to %ld bytes",
                 workspace_.value().numel(), workspace_size);
             workspace_.value().resize_({workspace_size});
         }
-        workspace = workspace_.value();
+        TORCH_CHECK(workspace_.value().dtype() == torch::kByte,
+            "Workspace tensor must have dtype=Byte");
+        TORCH_CHECK(workspace_.value().device() == qkv_or_q.device(),
+            "Workspace tensor must be on same device as inputs");
+        workspace = workspace_.value();
     }
     else
     {
-        workspace = torch::empty({workspace_size}, torch::dtype(torch::kByte).device(qkv_or_q.device()));
+        workspace = torch::empty({workspace_size}, torch::dtype(torch::kByte).device(qkv_or_q.device()));
     }

---

469-487: Top-level guards on fused vs separate inputs make sense, but add explicit K/V dtype/device checks early

You validate presence of K/V when needed. For earlier/faster failure, mirror the dtype/device checks here (in addition to Runner::run), so errors are caught before dispatch/templating.

Would you like me to draft a minimal guard block here?

---

285-292: Validate softmax_stats_tensor dtype/shape before casting to float2*

You correctly enforce float32 dtype, but since it’s reinterpreted as float2, also assert that numel is even and inner stride aligns to sizeof(float).

Proposed snippet (contextual):

TORCH_CHECK(softmax_stats_tensor->numel() % 2 == 0, "softmax_stats_tensor must have even number of floats for float2");

---

1-16: Update copyright year

Per coding guidelines, use the current year. Update 1993-2024 to 1993-2025.

Apply this diff:

- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &

---

683-712: Request: add targeted tests for separate-QKV context MLA across all paths

Given the behavioral changes (default/KV-reuse/chunked), please add tests that:

• validate K/V dtype/device mismatch raises,
• verify total_kv_len is honored in enqueue,
• cover both SM90 and SM120 config gates (feature flags permitting).

I can scaffold unit tests for test_attention_mla.py covering these scenarios. Want me to draft them?