[Kernel][Perf] Support Qwen head ratios in fused GDN MTP (#52539)
Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com>
kiroxu committed
Aug 18, 2026 at 09:17 UTC
cdb8545a91be16f2b80234e0801458ad2a18cf1b
4 files changed
+142
-34
csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu
+48
-24
@@ -146,7 +146,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) {
146
return {x, y};
147
}
148
149
-template <typename StateT>
149
+template <typename StateT, int ValueHeadsPerKeyHead>
150
__global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
151
const __nv_bfloat16* __restrict__ mixed_qkv,
152
const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b,
@@ -186,7 +186,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
186
return;
187
}
188
189
- const int key_head = value_head / 8;
189
+ const int key_head = value_head / ValueHeadsPerKeyHead;
190
__shared__ StateT shared_state[2][kChunkV][kDimK];
191
__shared__ float shared_q[kMaxMtpTokens][kDimK];
192
__shared__ float shared_k[kMaxMtpTokens][kDimK];
@@ -370,7 +370,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
370
}
371
}
372
373
-template <typename StateT>
373
+template <typename StateT, int ValueHeadsPerKeyHead>
374
void launch_gdn_decode_post_conv_mtp(
375
torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log,
376
torch::stable::Tensor const& dt_bias,
@@ -395,17 +395,19 @@ void launch_gdn_decode_post_conv_mtp(
395
get_current_cuda_stream(mixed_qkv.get_device_index());
396
const int num_requests = static_cast<int>(state_indices.size(0));
397
const dim3 grid(num_requests, num_value_heads);
398
- gdn_decode_post_conv_mtp_kernel<StateT><<<grid, kThreads, 0, stream>>>(
399
- static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
400
- static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
401
- static_cast<const int*>(state_indices.data_ptr()),
402
- static_cast<const int*>(cu_seqlens.data_ptr()),
403
- static_cast<const int*>(num_accepted_tokens.data_ptr()),
404
- static_cast<StateT*>(state.data_ptr()), output_gate,
405
- norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()),
406
- num_key_heads, num_value_heads, static_cast<int>(state_indices.size(1)),
407
- dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16,
408
- static_cast<float>(scale), static_cast<float>(norm_eps), strides);
398
+ gdn_decode_post_conv_mtp_kernel<StateT, ValueHeadsPerKeyHead>
399
+ <<<grid, kThreads, 0, stream>>>(
400
+ static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
401
+ static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
402
+ static_cast<const int*>(state_indices.data_ptr()),
403
+ static_cast<const int*>(cu_seqlens.data_ptr()),
404
+ static_cast<const int*>(num_accepted_tokens.data_ptr()),
405
+ static_cast<StateT*>(state.data_ptr()), output_gate,
406
+ norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()),
407
+ num_key_heads, num_value_heads,
408
+ static_cast<int>(state_indices.size(1)), dt_bias_type,
409
+ norm_weight.scalar_type() == ScalarType::BFloat16,
410
+ static_cast<float>(scale), static_cast<float>(norm_eps), strides);
411
const cudaError_t error = cudaGetLastError();
412
STD_TORCH_CHECK(error == cudaSuccess,
413
"GDN decode MTP post-conv kernel launch failed: ",
@@ -479,8 +481,12 @@ void fused_gdn_decode_post_conv_mtp(
481
STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0,
482
"mixed_qkv width is inconsistent with state");
483
const int num_key_heads = static_cast<int>(key_width / (2 * kDimK));
482
- STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads,
483
- "GDN decode MTP fusion requires HV/H=8");
484
+ const int value_heads_per_key_head = num_value_heads / num_key_heads;
485
+ STD_TORCH_CHECK(
486
+ num_value_heads % num_key_heads == 0 &&
487
+ ((value_heads_per_key_head >= 1 && value_heads_per_key_head <= 4) ||
488
+ value_heads_per_key_head == 8),
489
+ "GDN decode MTP fusion requires HV/H in {1, 2, 3, 4, 8}");
490
491
STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 &&
492
state_indices.size(1) > 0 &&
@@ -541,17 +547,35 @@ void fused_gdn_decode_post_conv_mtp(
547
const auto* b_ptr = static_cast<const __nv_bfloat16*>(b.data_ptr());
548
const auto* output_gate_ptr =
549
static_cast<const __nv_bfloat16*>(output_gate.data_ptr());
544
- if (state_scalar_type == ScalarType::Float) {
545
- launch_gdn_decode_post_conv_mtp<float>(
546
- mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens,
547
- num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr,
548
- output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps,
549
- strides);
550
- } else {
551
- launch_gdn_decode_post_conv_mtp<__nv_bfloat16>(
550
+ const auto launch = [&]<typename StateT, int ValueHeadsPerKeyHead>() {
551
+ launch_gdn_decode_post_conv_mtp<StateT, ValueHeadsPerKeyHead>(
552
mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens,
553
num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr,
554
output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps,
555
strides);
556
+ };
557
+ const auto dispatch_state_type = [&]<int ValueHeadsPerKeyHead>() {
558
+ if (state_scalar_type == ScalarType::Float) {
559
+ launch.template operator()<float, ValueHeadsPerKeyHead>();
560
+ } else {
561
+ launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead>();
562
+ }
563
+ };
564
+ switch (value_heads_per_key_head) {
565
+ case 1:
566
+ dispatch_state_type.template operator()<1>();
567
+ break;
568
+ case 2:
569
+ dispatch_state_type.template operator()<2>();
570
+ break;
571
+ case 3:
572
+ dispatch_state_type.template operator()<3>();
573
+ break;
574
+ case 4:
575
+ dispatch_state_type.template operator()<4>();
576
+ break;
577
+ default:
578
+ dispatch_state_type.template operator()<8>();
579
+ break;
580
}
581
}
tests/kernels/mamba/test_gdn_fused_mtp.py
+42
@@ -5,6 +5,7 @@
5
from __future__ import annotations
6
7
import types
8
+from typing import cast
9
from unittest.mock import patch
10
11
import pytest
@@ -37,6 +38,7 @@ from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( # noq
38
)
39
from vllm.utils.torch_utils import _encode_layer_name # noqa: E402
40
from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402
41
+ GDNAttentionMetadata,
42
GDNAttentionMetadataBuilder,
43
)
44
from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402
@@ -139,6 +141,46 @@ def _build_layer(
141
return layer
142
143
144
+@pytest.mark.parametrize(
145
+ "num_v_heads,expected",
146
+ [
147
+ pytest.param(2, True, id="ratio1"),
148
+ pytest.param(4, True, id="ratio2"),
149
+ pytest.param(6, True, id="ratio3"),
150
+ pytest.param(8, True, id="ratio4"),
151
+ pytest.param(16, True, id="ratio8"),
152
+ pytest.param(5, False, id="non-integral-ratio"),
153
+ pytest.param(10, False, id="unsupported-ratio5"),
154
+ ],
155
+)
156
+def test_fused_mtp_head_ratio_guard(num_v_heads: int, expected: bool) -> None:
157
+ if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"):
158
+ pytest.skip("fused GDN decode MTP op is not built")
159
+
160
+ layer = types.SimpleNamespace(
161
+ num_k_heads=2,
162
+ num_v_heads=num_v_heads,
163
+ kv_cache=(None, torch.empty(1, dtype=torch.float32, device="cuda")),
164
+ gdn_decode_kernel="cuda",
165
+ )
166
+ attn_metadata = types.SimpleNamespace(
167
+ spec_state_indices_tensor=torch.ones(
168
+ 1, SPEC_TOKENS, dtype=torch.int32, device="cuda"
169
+ ),
170
+ spec_sequence_masks=object(),
171
+ num_decodes=0,
172
+ num_spec_decodes=1,
173
+ )
174
+
175
+ assert (
176
+ QwenGatedDeltaNetAttention._can_use_fused_gdn_mtp_decode(
177
+ cast(QwenGatedDeltaNetAttention, layer),
178
+ cast(GDNAttentionMetadata, attn_metadata),
179
+ )
180
+ is expected
181
+ )
182
+
183
+
184
@torch.inference_mode()
185
def test_fused_forward_uses_packed_entrypoint() -> None:
186
"""Fused mode keeps projected QKVZ and BA packed through the model op."""
tests/kernels/test_fused_gdn_post_conv.py
+50
-9
@@ -217,18 +217,59 @@ def test_fused_post_conv_l0():
217
218
219
@pytest.mark.parametrize(
220
- "tp_size,query_lengths,state_dtype,norm_dtype",
220
+ "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype",
221
[
222
- pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"),
223
- pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"),
224
- pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"),
225
- pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"),
226
- pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"),
227
- pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"),
222
+ pytest.param(8, 16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"),
223
+ pytest.param(8, 4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"),
224
+ pytest.param(8, 16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"),
225
+ pytest.param(8, 4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"),
226
+ pytest.param(8, 16, (8,), torch.float32, torch.bfloat16, id="tp16-max"),
227
+ pytest.param(8, 4, (8,), torch.bfloat16, torch.float32, id="tp4-max"),
228
+ pytest.param(
229
+ 1,
230
+ 1,
231
+ (4, 4),
232
+ torch.float32,
233
+ torch.bfloat16,
234
+ id="ratio1-tp1-fp32",
235
+ ),
236
+ pytest.param(
237
+ 2,
238
+ 1,
239
+ (4, 2, 0),
240
+ torch.bfloat16,
241
+ torch.bfloat16,
242
+ id="ratio2-tp1-ragged-bf16",
243
+ ),
244
+ pytest.param(
245
+ 2,
246
+ 4,
247
+ (8,),
248
+ torch.float32,
249
+ torch.float32,
250
+ id="ratio2-tp4-max-fp32",
251
+ ),
252
+ pytest.param(
253
+ 3,
254
+ 2,
255
+ (4, 2, 0),
256
+ torch.float32,
257
+ torch.float32,
258
+ id="ratio3-tp2-ragged-fp32",
259
+ ),
260
+ pytest.param(
261
+ 4,
262
+ 4,
263
+ (8,),
264
+ torch.float32,
265
+ torch.bfloat16,
266
+ id="ratio4-tp4-max-fp32",
267
+ ),
268
],
269
)
270
@torch.inference_mode()
231
-def test_fused_gdn_decode_post_conv_mtp_ratio8(
271
+def test_fused_gdn_decode_post_conv_mtp_head_ratios(
272
+ head_ratio: int,
273
tp_size: int,
274
query_lengths: tuple[int, ...],
275
state_dtype: torch.dtype,
@@ -242,7 +283,7 @@ def test_fused_gdn_decode_post_conv_mtp_ratio8(
283
torch.manual_seed(0)
284
device = "cuda"
285
H = 16 // tp_size
245
- HV = 128 // tp_size
286
+ HV = head_ratio * H
287
K = V = 128
288
num_reqs = len(query_lengths)
289
state_width = max(query_lengths)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+2
-1
@@ -1805,7 +1805,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1805
and attn_metadata.num_spec_decodes > 0
1806
and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES
1807
and self.gdn_decode_kernel == "cuda"
1808
- and self.num_v_heads == 8 * self.num_k_heads
1808
+ and self.num_v_heads % self.num_k_heads == 0
1809
+ and self.num_v_heads // self.num_k_heads in (1, 2, 3, 4, 8)
1810
and state_indices is not None
1811
and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS
1812
and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp")