[Bugfix] Allocate packed outputs in fused_q_kv_rmsnorm so q_b_proj keeps its low-latency GEMM path at decode (#53109)
Signed-off-by: yasong.wang <yasong.wang@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Roy Wang committed
Aug 28, 2026 at 10:08 UTC
b74ac0bc61d8b8f74848664b1a3b4e878f047861
2 files changed
+40
-2
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+31
@@ -51,6 +51,37 @@ def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
51
torch.testing.assert_close(kv_out, kv_ref, **tol)
52
53
54
+@pytest.mark.parametrize("num_tokens", [1, 16])
55
+def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int):
56
+ """Regression guard: the outputs must be packed row-major even when the
57
+ inputs are column slices of a wider fused-projection buffer.
58
+
59
+ empty_like preserves the strides of size-1 dims, so with num_tokens == 1
60
+ a [1, q_size] slice of a [1, q_size + kv_size] buffer used to produce a
61
+ qr_out with row stride q_size + kv_size. Downstream dispatchers that
62
+ require packed row-major inputs then reject the tensor and silently fall
63
+ back to a slower GEMM path on every decode step."""
64
+ device = "cuda"
65
+ dtype = torch.bfloat16
66
+ q_size, kv_size = 192, 576
67
+ fused = torch.randn(num_tokens, q_size + kv_size, dtype=dtype, device=device)
68
+ qr, kv = fused.split([q_size, kv_size], dim=-1)
69
+ qw = torch.randn(q_size, dtype=dtype, device=device)
70
+ kvw = torch.randn(kv_size, dtype=dtype, device=device)
71
+ eps = 1e-6
72
+
73
+ qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
74
+
75
+ assert qr_out.stride() == (q_size, 1)
76
+ assert kv_out.stride() == (kv_size, 1)
77
+
78
+ qr_ref = _ref_rmsnorm(qr, qw, eps)
79
+ kv_ref = _ref_rmsnorm(kv, kvw, eps)
80
+ tol = dict(rtol=1e-2, atol=1e-2)
81
+ torch.testing.assert_close(qr_out, qr_ref, **tol)
82
+ torch.testing.assert_close(kv_out, kv_ref, **tol)
83
+
84
+
85
@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
86
def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
87
"""Regression guard: grid used to be (2, num_tokens), hitting CUDA's
vllm/models/common/ops/fused_qk_rmsnorm.py
+9
-2
@@ -77,8 +77,15 @@ def fused_q_kv_rmsnorm(
77
q_size = qr.shape[1]
78
kv_size = kv.shape[1]
79
num_tokens = qr.shape[0]
80
- qr_out = torch.empty_like(qr)
81
- kv_out = torch.empty_like(kv)
80
+ # Allocate packed outputs rather than empty_like: qr/kv are typically
81
+ # column slices of a fused qkv_a projection, and for num_tokens == 1
82
+ # empty_like preserves the parent's row stride (torch keeps strides of
83
+ # size-1 dims), e.g. (2112, 1) for a [1, 1536] q slice. Downstream
84
+ # dispatchers that require packed row-major inputs (e.g. the Kimi-K3
85
+ # low-latency GEMM's _is_packed_row_major check) then reject the tensor
86
+ # and silently fall back to cuBLAS on every decode step.
87
+ qr_out = torch.empty(qr.shape, dtype=qr.dtype, device=qr.device)
88
+ kv_out = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
89
if num_tokens == 0:
90
return qr_out, kv_out
91