iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][Perf] Support Qwen head ratios in fused GDN MTP (#52539)

Browse Source 
Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com>
This commit is contained in:
kiroxu
2026-08-17 18:17:40 -07:00
committed by GitHub
parent c296851a7d
commit cdb8545a91
4 changed files with 142 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu
+48 -24
View File
@@ -146,7 +146,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) {
return {x, y};
}
template <typename StateT>
template <typename StateT, int ValueHeadsPerKeyHead>
__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 <typename StateT>
template <typename StateT, int ValueHeadsPerKeyHead>
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<int>(state_indices.size(0));
const dim3 grid(num_requests, num_value_heads);
gdn_decode_post_conv_mtp_kernel<StateT><<<grid, kThreads, 0, stream>>>(
static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
static_cast<const int*>(state_indices.data_ptr()),
static_cast<const int*>(cu_seqlens.data_ptr()),
static_cast<const int*>(num_accepted_tokens.data_ptr()),
static_cast<StateT*>(state.data_ptr()), output_gate,
norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()),
num_key_heads, num_value_heads, static_cast<int>(state_indices.size(1)),
dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16,
static_cast<float>(scale), static_cast<float>(norm_eps), strides);
gdn_decode_post_conv_mtp_kernel<StateT, ValueHeadsPerKeyHead>
<<<grid, kThreads, 0, stream>>>(
static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
static_cast<const int*>(state_indices.data_ptr()),
static_cast<const int*>(cu_seqlens.data_ptr()),
static_cast<const int*>(num_accepted_tokens.data_ptr()),
static_cast<StateT*>(state.data_ptr()), output_gate,
norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()),
num_key_heads, num_value_heads,
static_cast<int>(state_indices.size(1)), dt_bias_type,
norm_weight.scalar_type() == ScalarType::BFloat16,
static_cast<float>(scale), static_cast<float>(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<int>(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<const __nv_bfloat16*>(b.data_ptr());
const auto* output_gate_ptr =
static_cast<const __nv_bfloat16*>(output_gate.data_ptr());
if (state_scalar_type == ScalarType::Float) {
launch_gdn_decode_post_conv_mtp<float>(
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 = [&]<typename StateT, int ValueHeadsPerKeyHead>() {
launch_gdn_decode_post_conv_mtp<StateT, ValueHeadsPerKeyHead>(
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 = [&]<int ValueHeadsPerKeyHead>() {
if (state_scalar_type == ScalarType::Float) {
launch.template operator()<float, ValueHeadsPerKeyHead>();
} 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;
}
}
tests/kernels/mamba/test_gdn_fused_mtp.py
+42
View File
@@ -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."""
tests/kernels/test_fused_gdn_post_conv.py
+50 -9
View File
@@ -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)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+2 -1
View File
@@ -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")