[None][fix] DSv4 fused FP8 Q-quant: q_pe must be 3D for thop.attention (NVIDIA#14433)

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>

Author: mingyangHao

tensorrt_llm/_torch/modules/attention.py

@@ -1760,21 +1760,21 @@ def _deepseek_v4_q_b_layernorm_fused_fp8(self, q_proj: torch.Tensor):
             self._quant_scale_qkv = torch.tensor([1.0],
                                                  dtype=torch.float32,
                                                  device=q_proj.device)
-        # Allocate the interleaved [N, H*head_dim] FP8 buffer (nope filled by
-        # this op, rope slot left for applyMLARopeAndAssignQKVKernelOptContext) and the bf16 rope buffer.
+        # q_pe is 3D so thop.attention's sparse-MLA context branch passes its
+        # q_pe->dim() == 3 check; the kernel op consumes the flat 2D view.
         num_tokens = q_proj.shape[0]
         rope_dim = self.qk_head_dim - self.kv_lora_rank
         quant_q_buffer = q_proj.new_empty(
             (num_tokens, self.num_heads_tp * self.qk_head_dim),
             dtype=torch.float8_e4m3fn)
-        q_pe = q_proj.new_empty((num_tokens, self.num_heads_tp * rope_dim))
+        q_pe = q_proj.new_empty((num_tokens, self.num_heads_tp, rope_dim))
         torch.ops.trtllm.deepseek_v4_q_norm_fused_fp8(
             q_proj,
             quant_q_buffer,
-            q_pe,
+            q_pe.view(num_tokens, self.num_heads_tp * rope_dim),
             self.num_heads_tp,
             self.qk_head_dim,
-            self.kv_lora_rank,  # nope_dim
+            self.kv_lora_rank,
             float(self.q_b_layernorm.variance_epsilon),
             self._quant_scale_qkv,
         )

tests/integration/test_lists/test-db/l0_b200_ds.yml

@@ -25,6 +25,7 @@ l0_b200_ds:
   - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py TIMEOUT (90)
   - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py TIMEOUT (90)
   - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py TIMEOUT (15)
+  - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15)
   # ------------- Disaggregated transfer component tests (single GPU) ---------------
   - unittest/disaggregated/test_agent.py
   - unittest/disaggregated/test_agent_multi_backends.py

tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py

@@ -168,6 +168,42 @@ def test_deepseek_v4_q_norm_fused_fp8_zero_rows():
     assert q_pe.shape == (0, num_heads * rope_dim)
 
 
+@pytest.mark.parametrize("num_tokens", [1, 7, 129])
+def test_deepseek_v4_q_b_layernorm_fused_fp8_returns_3d_q_pe(num_tokens):
+    """Lock down the q_pe dim==3 contract expected by thop.attention's
+    sparse-MLA context branch (TORCH_CHECK(q_pe->dim() == 3))."""
+    import types
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch.modules.attention import MLA
+
+    num_heads = 16
+    qk_head_dim = 512
+    kv_lora_rank = 448
+    rope_dim = qk_head_dim - kv_lora_rank
+    stub = SimpleNamespace(
+        num_heads_tp=num_heads,
+        qk_head_dim=qk_head_dim,
+        kv_lora_rank=kv_lora_rank,
+        q_b_layernorm=SimpleNamespace(variance_epsilon=1e-6),
+    )
+    fused = types.MethodType(MLA._deepseek_v4_q_b_layernorm_fused_fp8, stub)
+    q_proj = torch.randn(
+        num_tokens, num_heads * qk_head_dim, dtype=torch.bfloat16, device="cuda"
+    ).contiguous()
+
+    placeholder_q, quant_q_buffer, q_pe, scale = fused(q_proj)
+
+    assert q_pe.shape == (num_tokens, num_heads, rope_dim)
+    assert q_pe.stride(2) == 1
+    assert q_pe.is_contiguous()
+    assert quant_q_buffer.shape == (num_tokens, num_heads * qk_head_dim)
+    assert quant_q_buffer.dtype == torch.float8_e4m3fn
+    assert placeholder_q.data_ptr() == q_proj.data_ptr()
+    assert scale.shape == (1,) and scale.dtype == torch.float32
+    assert float(scale.item()) == 1.0
+
+
 @pytest.mark.parametrize("num_tokens", [1, 7, 129])
 @pytest.mark.parametrize("num_heads", [1, 16, 128])
 def test_deepseek_v4_q_norm_fused_fp8_interleaved_layout(num_tokens, num_heads):