From cdb8545a91be16f2b80234e0801458ad2a18cf1b Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:17:40 +0800 Subject: [PATCH] [Kernel][Perf] Support Qwen head ratios in fused GDN MTP (#52539) Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> --- .../gdn/fused_gdn_decode_kernel.cu | 72 ++++++++++++------- tests/kernels/mamba/test_gdn_fused_mtp.py | 42 +++++++++++ tests/kernels/test_fused_gdn_post_conv.py | 59 ++++++++++++--- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 3 +- 4 files changed, 142 insertions(+), 34 deletions(-) diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu index 31a16fcf99..1166b41868 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu @@ -146,7 +146,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { return {x, y}; } -template +template __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( const __nv_bfloat16* __restrict__ mixed_qkv, 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( return; } - const int key_head = value_head / 8; + const int key_head = value_head / ValueHeadsPerKeyHead; __shared__ StateT shared_state[2][kChunkV][kDimK]; __shared__ float shared_q[kMaxMtpTokens][kDimK]; __shared__ float shared_k[kMaxMtpTokens][kDimK]; @@ -370,7 +370,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( } } -template +template void launch_gdn_decode_post_conv_mtp( torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log, torch::stable::Tensor const& dt_bias, @@ -395,17 +395,19 @@ void launch_gdn_decode_post_conv_mtp( get_current_cuda_stream(mixed_qkv.get_device_index()); const int num_requests = static_cast(state_indices.size(0)); const dim3 grid(num_requests, num_value_heads); - gdn_decode_post_conv_mtp_kernel<<>>( - static_cast(mixed_qkv.data_ptr()), a, b, - static_cast(a_log.data_ptr()), dt_bias.data_ptr(), - static_cast(state_indices.data_ptr()), - static_cast(cu_seqlens.data_ptr()), - static_cast(num_accepted_tokens.data_ptr()), - static_cast(state.data_ptr()), output_gate, - norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), - num_key_heads, num_value_heads, static_cast(state_indices.size(1)), - dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16, - static_cast(scale), static_cast(norm_eps), strides); + gdn_decode_post_conv_mtp_kernel + <<>>( + static_cast(mixed_qkv.data_ptr()), a, b, + static_cast(a_log.data_ptr()), dt_bias.data_ptr(), + static_cast(state_indices.data_ptr()), + static_cast(cu_seqlens.data_ptr()), + static_cast(num_accepted_tokens.data_ptr()), + static_cast(state.data_ptr()), output_gate, + norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), + num_key_heads, num_value_heads, + static_cast(state_indices.size(1)), dt_bias_type, + norm_weight.scalar_type() == ScalarType::BFloat16, + static_cast(scale), static_cast(norm_eps), strides); const cudaError_t error = cudaGetLastError(); STD_TORCH_CHECK(error == cudaSuccess, "GDN decode MTP post-conv kernel launch failed: ", @@ -479,8 +481,12 @@ void fused_gdn_decode_post_conv_mtp( STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0, "mixed_qkv width is inconsistent with state"); const int num_key_heads = static_cast(key_width / (2 * kDimK)); - STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads, - "GDN decode MTP fusion requires HV/H=8"); + const int value_heads_per_key_head = num_value_heads / num_key_heads; + STD_TORCH_CHECK( + num_value_heads % num_key_heads == 0 && + ((value_heads_per_key_head >= 1 && value_heads_per_key_head <= 4) || + value_heads_per_key_head == 8), + "GDN decode MTP fusion requires HV/H in {1, 2, 3, 4, 8}"); STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 && state_indices.size(1) > 0 && @@ -541,17 +547,35 @@ void fused_gdn_decode_post_conv_mtp( const auto* b_ptr = static_cast(b.data_ptr()); const auto* output_gate_ptr = static_cast(output_gate.data_ptr()); - if (state_scalar_type == ScalarType::Float) { - launch_gdn_decode_post_conv_mtp( - mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, - num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, - output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, - strides); - } else { - launch_gdn_decode_post_conv_mtp<__nv_bfloat16>( + const auto launch = [&]() { + launch_gdn_decode_post_conv_mtp( mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, strides); + }; + const auto dispatch_state_type = [&]() { + if (state_scalar_type == ScalarType::Float) { + launch.template operator()(); + } else { + launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead>(); + } + }; + switch (value_heads_per_key_head) { + case 1: + dispatch_state_type.template operator()<1>(); + break; + case 2: + dispatch_state_type.template operator()<2>(); + break; + case 3: + dispatch_state_type.template operator()<3>(); + break; + case 4: + dispatch_state_type.template operator()<4>(); + break; + default: + dispatch_state_type.template operator()<8>(); + break; } } diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py index 97ed4c48e7..0d78f8ceeb 100644 --- a/tests/kernels/mamba/test_gdn_fused_mtp.py +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -5,6 +5,7 @@ from __future__ import annotations import types +from typing import cast from unittest.mock import patch import pytest @@ -37,6 +38,7 @@ from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( # noq ) from vllm.utils.torch_utils import _encode_layer_name # noqa: E402 from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 + GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 @@ -139,6 +141,46 @@ def _build_layer( return layer +@pytest.mark.parametrize( + "num_v_heads,expected", + [ + pytest.param(2, True, id="ratio1"), + pytest.param(4, True, id="ratio2"), + pytest.param(6, True, id="ratio3"), + pytest.param(8, True, id="ratio4"), + pytest.param(16, True, id="ratio8"), + pytest.param(5, False, id="non-integral-ratio"), + pytest.param(10, False, id="unsupported-ratio5"), + ], +) +def test_fused_mtp_head_ratio_guard(num_v_heads: int, expected: bool) -> None: + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): + pytest.skip("fused GDN decode MTP op is not built") + + layer = types.SimpleNamespace( + num_k_heads=2, + num_v_heads=num_v_heads, + kv_cache=(None, torch.empty(1, dtype=torch.float32, device="cuda")), + gdn_decode_kernel="cuda", + ) + attn_metadata = types.SimpleNamespace( + spec_state_indices_tensor=torch.ones( + 1, SPEC_TOKENS, dtype=torch.int32, device="cuda" + ), + spec_sequence_masks=object(), + num_decodes=0, + num_spec_decodes=1, + ) + + assert ( + QwenGatedDeltaNetAttention._can_use_fused_gdn_mtp_decode( + cast(QwenGatedDeltaNetAttention, layer), + cast(GDNAttentionMetadata, attn_metadata), + ) + is expected + ) + + @torch.inference_mode() def test_fused_forward_uses_packed_entrypoint() -> None: """Fused mode keeps projected QKVZ and BA packed through the model op.""" diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py index b6dacf2e73..7cdeb55b9b 100644 --- a/tests/kernels/test_fused_gdn_post_conv.py +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -217,18 +217,59 @@ def test_fused_post_conv_l0(): @pytest.mark.parametrize( - "tp_size,query_lengths,state_dtype,norm_dtype", + "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype", [ - pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), - pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), - pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), - pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), - pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), - pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param(8, 16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), + pytest.param(8, 4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), + pytest.param(8, 16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), + pytest.param(8, 4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), + pytest.param(8, 16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), + pytest.param(8, 4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param( + 1, + 1, + (4, 4), + torch.float32, + torch.bfloat16, + id="ratio1-tp1-fp32", + ), + pytest.param( + 2, + 1, + (4, 2, 0), + torch.bfloat16, + torch.bfloat16, + id="ratio2-tp1-ragged-bf16", + ), + pytest.param( + 2, + 4, + (8,), + torch.float32, + torch.float32, + id="ratio2-tp4-max-fp32", + ), + pytest.param( + 3, + 2, + (4, 2, 0), + torch.float32, + torch.float32, + id="ratio3-tp2-ragged-fp32", + ), + pytest.param( + 4, + 4, + (8,), + torch.float32, + torch.bfloat16, + id="ratio4-tp4-max-fp32", + ), ], ) @torch.inference_mode() -def test_fused_gdn_decode_post_conv_mtp_ratio8( +def test_fused_gdn_decode_post_conv_mtp_head_ratios( + head_ratio: int, tp_size: int, query_lengths: tuple[int, ...], state_dtype: torch.dtype, @@ -242,7 +283,7 @@ def test_fused_gdn_decode_post_conv_mtp_ratio8( torch.manual_seed(0) device = "cuda" H = 16 // tp_size - HV = 128 // tp_size + HV = head_ratio * H K = V = 128 num_reqs = len(query_lengths) state_width = max(query_lengths) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 27dc4d90d7..d83b076509 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1805,7 +1805,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): and attn_metadata.num_spec_decodes > 0 and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES and self.gdn_decode_kernel == "cuda" - and self.num_v_heads == 8 * self.num_k_heads + and self.num_v_heads % self.num_k_heads == 0 + and self.num_v_heads // self.num_k_heads in (1, 2, 3, 4, 8) and state_indices is not None and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp")