tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py (1)

269-312: Fail fast on inconsistent QKV widths before launching these kernels.

Both wrappers assume the source width is exactly q_dim + k_dim + v_dim and that the head-count arguments multiply back to those dims. If a caller ever passes mismatched values, this becomes raw Triton pointer arithmetic instead of a clean Python exception.

🛡️ Suggested guardrails

 def transpose_and_split_qkv(
     prefill_t: torch.Tensor,
     decode: torch.Tensor,
@@
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@
     num_cols, num_prefill = prefill_t.shape
+    expected_cols = q_dim + k_dim + v_dim
+    if decode.ndim != 2 or num_cols != expected_cols or decode.shape[1] != expected_cols:
+        raise ValueError(
+            f"Expected prefill/decode widths to equal {expected_cols}, got "
+            f"{num_cols} and {decode.shape[1]}"
+        )
+    if q_dim != num_q_heads * head_k_dim or k_dim != num_q_heads * head_k_dim:
+        raise ValueError("Q/K dims must match num_q_heads * head_k_dim")
+    if v_dim != num_v_heads * head_v_dim:
+        raise ValueError("V dim must match num_v_heads * head_v_dim")
     num_decode = decode.shape[0]
     total_seq = num_prefill + num_decode
@@
 def split_qkv_contiguous(
     mixed_qkv: torch.Tensor,
@@
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@
-    seq_len = mixed_qkv.shape[0]
+    seq_len, num_cols = mixed_qkv.shape
+    expected_cols = q_dim + k_dim + v_dim
+    if num_cols != expected_cols:
+        raise ValueError(f"Expected mixed_qkv width {expected_cols}, got {num_cols}")
+    if q_dim != num_q_heads * head_k_dim or k_dim != num_q_heads * head_k_dim:
+        raise ValueError("Q/K dims must match num_q_heads * head_k_dim")
+    if v_dim != num_v_heads * head_v_dim:
+        raise ValueError("V dim must match num_v_heads * head_v_dim")
     src_stride_seq, src_stride_dim = mixed_qkv.stride()

Also applies to: 364-417

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py` around lines 269 -
312, Add preflight checks in transpose_and_split_qkv before launching
_transpose_and_split_qkv_kernel: verify that num_cols == q_dim + k_dim + v_dim
and that q_dim == num_q_heads * head_k_dim and v_dim == num_v_heads * head_v_dim
(and similarly k_dim if applicable), and raise a clear ValueError if any
mismatch occurs. Do the same guard checks in the analogous function at lines
364-417 so callers get a Python exception instead of invalid Triton pointer
arithmetic. Include the unique symbols transpose_and_split_qkv,
_transpose_and_split_qkv_kernel, q_dim, k_dim, v_dim, num_q_heads, head_k_dim,
num_v_heads, head_v_dim when locating where to add these checks.

tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py (1)

91-117: Add a non-contiguous source case for split_qkv_contiguous().

This test only covers a contiguous [T, D] input, but tensorrt_llm/_torch/modules/mamba/gdn_mixer.py Lines 686-688 call the helper with mixed_qkv_t.transpose(0, 1). The stride-sensitive path this kernel was added for is therefore untested right now.

🧪 Suggested test expansion

 `@skip_no_cuda`
 `@pytest.mark.parametrize`("seq_len", [1, 32, 128, 1024])
+@pytest.mark.parametrize("transpose_input", [False, True])
 `@pytest.mark.parametrize`(
     "q_dim,k_dim,v_dim,num_q_heads,head_k_dim,num_v_heads,head_v_dim",
@@
 )
 `@pytest.mark.parametrize`("dtype", [torch.bfloat16])
-def test_split_qkv_contiguous(seq_len, q_dim, k_dim, v_dim,
-                               num_q_heads, head_k_dim, num_v_heads, head_v_dim, dtype):
+def test_split_qkv_contiguous(
+    seq_len,
+    transpose_input,
+    q_dim,
+    k_dim,
+    v_dim,
+    num_q_heads,
+    head_k_dim,
+    num_v_heads,
+    head_v_dim,
+    dtype,
+):
@@
-    mixed_qkv = torch.randn(seq_len, total_dim, dtype=dtype, device=device)
+    if transpose_input:
+        mixed_qkv = torch.randn(total_dim, seq_len, dtype=dtype, device=device).transpose(0, 1)
+        assert not mixed_qkv.is_contiguous()
+    else:
+        mixed_qkv = torch.randn(seq_len, total_dim, dtype=dtype, device=device)

@@
     assert q_out.shape == (1, seq_len, num_q_heads, head_k_dim)
     assert k_out.shape == (1, seq_len, num_q_heads, head_k_dim)
     assert v_out.shape == (1, seq_len, num_v_heads, head_v_dim)
+    assert q_out.is_contiguous()
+    assert k_out.is_contiguous()
+    assert v_out.is_contiguous()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py` around
lines 91 - 117, Update test_split_qkv_contiguous to also exercise the
non-contiguous input path: create a transposed/non-contiguous source (e.g.,
mixed_qkv_t = mixed_qkv.transpose(0,1) or mixed_qkv.t()), call
split_qkv_contiguous with that tensor (matching the same q_dim/k_dim/v_dim/head
args), and then reshape/transpose the kernel outputs back to 2D to compare
against _ref_split_qkv_contiguous(mixed_qkv, ...). Ensure you assert the output
shapes for the transposed case and use torch.testing.assert_close to compare
values (same rtol/atol), so the stride-sensitive branch used by
split_qkv_contiguous is covered.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes