minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Performance] Optimize Qwen3.5 GDN prefill projection layouts (#36267)

Browse Source
This commit is contained in:
YAMY
2026-09-04 10:32:44 -07:00
committed by GitHub
parent 55509b3f42
commit 07199fa220
7 changed files with 386 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/l2norm.py
+161
View File
@@ -119,6 +119,167 @@ def l2norm_fwd(
return y.view(x_shape_og)
@triton.jit(do_not_specialize=["T"])
def gdn_prefill_qkv_prepare_kernel(
q,
k,
v,
q_out,
k_out,
v_out,
q_stride_t,
q_stride_h,
q_stride_d,
k_stride_t,
k_stride_h,
k_stride_d,
v_stride_t,
v_stride_h,
v_stride_d,
T,
H_QK: tl.constexpr,
H_V: tl.constexpr,
D: tl.constexpr,
BT: tl.constexpr,
BD: tl.constexpr,
):
"""Materialize strided Q/K/V into token-major tensors in one launch."""
token_block = tl.program_id(0)
head_idx = tl.program_id(1)
if head_idx < H_QK:
# Match l2norm_fwd_kernel's block layout so the BF16 reduction tree is
# unchanged for strided inputs.
q_block = tl.make_block_ptr(
q + head_idx * q_stride_h,
(T, D),
(q_stride_t, q_stride_d),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
k_block = tl.make_block_ptr(
k + head_idx * k_stride_h,
(T, D),
(k_stride_t, k_stride_d),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
q_values = tl.load(q_block, boundary_check=(0, 1)).to(tl.float32)
k_values = tl.load(k_block, boundary_check=(0, 1)).to(tl.float32)
q_output_block = tl.make_block_ptr(
q_out + head_idx * D,
(T, D),
(H_QK * D, 1),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
k_output_block = tl.make_block_ptr(
k_out + head_idx * D,
(T, D),
(H_QK * D, 1),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
tl.store(
q_output_block,
q_values.to(q_output_block.dtype.element_ty),
boundary_check=(0, 1),
)
tl.store(
k_output_block,
k_values.to(k_output_block.dtype.element_ty),
boundary_check=(0, 1),
)
else:
value_head = head_idx - H_QK
v_block = tl.make_block_ptr(
v + value_head * v_stride_h,
(T, D),
(v_stride_t, v_stride_d),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
v_values = tl.load(v_block, boundary_check=(0, 1))
v_output_block = tl.make_block_ptr(
v_out + value_head * D,
(T, D),
(H_V * D, 1),
(token_block * BT, 0),
(BT, BD),
(1, 0),
)
tl.store(v_output_block, v_values, boundary_check=(0, 1))
def gdn_prefill_qkv_prepare_fwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Prepare Q/K/V for FlashInfer, materializing only strided inputs."""
if q.ndim != 3 or k.shape != q.shape or v.ndim != 3:
raise ValueError(
"GDN fused prepare requires equal Q/K [T, Hqk, D] shapes and "
"V [T, Hv, D], got "
f"{q.shape=}, {k.shape=}, {v.shape=}"
)
if v.shape[0] != q.shape[0] or v.shape[2] != q.shape[2]:
raise ValueError(
"GDN fused prepare requires common token and head-dim axes, got "
f"{q.shape=}, {v.shape=}"
)
if q.device != k.device or q.device != v.device:
raise ValueError("GDN fused prepare requires Q/K/V on the same device")
if q.dtype != k.dtype or q.dtype != v.dtype:
raise ValueError("GDN fused prepare requires equal Q/K/V dtypes")
T, H_QK, D = q.shape
H_V = v.shape[1]
if D > 512:
raise ValueError(f"GDN fused prepare supports head dim <= 512, got {D}")
if q.is_contiguous() and k.is_contiguous() and v.is_contiguous():
return l2norm_fwd(q, eps), l2norm_fwd(k, eps), v
q_out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
k_out = torch.empty(k.shape, dtype=k.dtype, device=k.device)
v_out = torch.empty(v.shape, dtype=v.dtype, device=v.device)
BT = 16
BD = triton.next_power_of_2(D)
grid = (triton.cdiv(T, BT), H_QK + H_V)
gdn_prefill_qkv_prepare_kernel[grid](
q,
k,
v,
q_out,
k_out,
v_out,
q.stride(0),
q.stride(1),
q.stride(2),
k.stride(0),
k.stride(1),
k.stride(2),
v.stride(0),
v.stride(1),
v.stride(2),
T=T,
H_QK=H_QK,
H_V=H_V,
D=D,
BT=BT,
BD=BD,
num_warps=4,
num_stages=2,
)
return l2norm_fwd(q_out, eps), l2norm_fwd(k_out, eps), v_out
class L2NormFunction(torch.autograd.Function):
@staticmethod
@input_guard
python/sglang/kernels/ops/attention/fla/layernorm_gated.py
+39 -8
View File
@@ -83,6 +83,8 @@ def _layer_norm_fwd_1pass_kernel(
stride_x_row, # how much to increase the pointer when moving by 1 row
stride_y_row,
stride_z_row,
stride_z_token,
stride_z_head,
M, # number of rows in X
N: tl.constexpr, # number of columns in X
eps, # epsilon to avoid division by zero
@@ -90,6 +92,8 @@ def _layer_norm_fwd_1pass_kernel(
ROWS_PER_BLOCK: tl.constexpr,
HAS_BIAS: tl.constexpr,
HAS_Z: tl.constexpr,
Z_IS_3D: tl.constexpr,
Z_HEADS: tl.constexpr,
NORM_BEFORE_GATE: tl.constexpr,
IS_RMS_NORM: tl.constexpr,
ACTIVATION: tl.constexpr,
@@ -123,7 +127,13 @@ def _layer_norm_fwd_1pass_kernel(
x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32)
if HAS_Z and not NORM_BEFORE_GATE:
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
if Z_IS_3D:
z_row_offsets = (rows[:, None] // Z_HEADS) * stride_z_token + (
rows[:, None] % Z_HEADS
) * stride_z_head
else:
z_row_offsets = rows[:, None] * stride_z_row
Z_base = Z + z_row_offsets + col_offsets
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
if ACTIVATION == "swish" or ACTIVATION == "silu":
x *= z * tl.sigmoid(z)
@@ -169,7 +179,13 @@ def _layer_norm_fwd_1pass_kernel(
y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :]
if HAS_Z and NORM_BEFORE_GATE:
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
if Z_IS_3D:
z_row_offsets = (rows[:, None] // Z_HEADS) * stride_z_token + (
rows[:, None] % Z_HEADS
) * stride_z_head
else:
z_row_offsets = rows[:, None] * stride_z_row
Z_base = Z + z_row_offsets + col_offsets
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
if ACTIVATION == "swish" or ACTIVATION == "silu":
y *= z * tl.sigmoid(z)
@@ -223,9 +239,14 @@ def _layer_norm_fwd(
assert N % group_size == 0
ngroups = N // group_size
assert x.stride(-1) == 1
z_is_3d = z is not None and z.ndim == 3
if z is not None:
assert z.stride(-1) == 1
assert z.shape == (M, N)
if z_is_3d:
assert z.shape[0] * z.shape[1] == M
assert z.shape[2] == N
else:
assert z.shape == (M, N)
assert weight.shape == (N,)
assert weight.stride(-1) == 1
if bias is not None:
@@ -280,7 +301,9 @@ def _layer_norm_fwd(
rstd,
x.stride(0),
out.stride(0),
z.stride(0) if z is not None else 0,
z.stride(0) if z is not None and not z_is_3d else 0,
z.stride(0) if z_is_3d else 0,
z.stride(1) if z_is_3d else 0,
M,
group_size,
eps,
@@ -288,6 +311,8 @@ def _layer_norm_fwd(
ROWS_PER_BLOCK=rows_per_block,
HAS_BIAS=bias is not None,
HAS_Z=z is not None,
Z_IS_3D=z_is_3d,
Z_HEADS=z.shape[1] if z_is_3d else 1,
NORM_BEFORE_GATE=norm_before_gate,
IS_RMS_NORM=is_rms_norm,
num_warps=num_warps,
@@ -321,10 +346,16 @@ def rms_norm_gated(
if x.stride(-1) != 1:
x = x.contiguous()
if z is not None:
assert z.shape == x_shape_og
z = z.reshape(-1, z.shape[-1])
if z.stride(-1) != 1:
z = z.contiguous()
if z.shape == x_shape_og:
z = z.reshape(-1, z.shape[-1])
if z.stride(-1) != 1:
z = z.contiguous()
else:
assert len(x_shape_og) == 2
assert z.ndim == 3
assert z.shape[0] * z.shape[1] == x.shape[0]
assert z.shape[2] == x.shape[1]
assert z.stride(-1) == 1
weight = weight.contiguous()
if bias is not None:
bias = bias.contiguous()
python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py
+18
View File
@@ -330,6 +330,24 @@ def fused_qkvzba_split_reshape_cat_contiguous(
return mixed_qkv, z, b, a
def qwen3_5_gdn_prefill_projection_views(
mixed_qkvz,
mixed_ba,
num_heads_qk,
num_heads_v,
head_qk,
head_v,
):
"""Return strided views accepted by the prefill GDN consumers."""
tokens = mixed_qkvz.shape[0]
qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v
mixed_qkv = mixed_qkvz[:, :qkv_dim]
z = mixed_qkvz[:, qkv_dim:].view(tokens, num_heads_v, head_v)
b = mixed_ba[:, :num_heads_v]
a = mixed_ba[:, num_heads_v : 2 * num_heads_v]
return mixed_qkv, z, b, a
# Fusion begins after the quantized GEMMs: qkvz=[q|k|v|z], ba=[b|a].
# This tail unpacks both projections and updates the causal Conv1D state.
python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py
+4 -4
View File
@@ -478,15 +478,15 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
state_checkpoint_every_n_tokens: int = 0,
**kwargs,
) -> tuple:
from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
from sglang.kernels.ops.attention.fla.l2norm import (
gdn_prefill_qkv_prepare_fwd,
)
total_seq_len = q.shape[1]
num_v_heads = v.shape[2]
head_v_dim = v.shape[3]
q_fi = l2norm_fwd(q[0].contiguous())
k_fi = l2norm_fwd(k[0].contiguous())
v_fi = v[0].contiguous()
q_fi, k_fi, v_fi = gdn_prefill_qkv_prepare_fwd(q[0], k[0], v[0])
output = kwargs.get("output")
output_fi = None
python/sglang/srt/models/qwen3_5.py
+18 -4
View File
@@ -27,6 +27,7 @@ import triton
from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
fused_qkvzba_split_reshape_cat_contiguous,
qwen3_5_gdn_prefill_projection_views,
)
from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul
@@ -781,6 +782,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
backend, projected_states_qkvz, projected_states_ba, forward_batch
)
use_strided_prefill_z = False
use_fused_decode_proj_conv = (
_gdn_decode_fused_proj_conv
and forward_batch.forward_mode.is_decode()
@@ -804,7 +806,15 @@ class Qwen3_5GatedDeltaNet(nn.Module):
else:
num_k_heads_tp = triton.cdiv(self.num_k_heads, self.attn_tp_size)
num_v_heads_tp = triton.cdiv(self.num_v_heads, self.attn_tp_size)
mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous(
use_strided_prefill_z = (
_is_cuda and forward_batch.forward_mode.is_extend_without_speculative()
)
split_fn = (
qwen3_5_gdn_prefill_projection_views
if use_strided_prefill_z
else fused_qkvzba_split_reshape_cat_contiguous
)
mixed_qkv, z, b, a = split_fn(
projected_states_qkvz,
projected_states_ba,
num_k_heads_tp,
@@ -844,11 +854,15 @@ class Qwen3_5GatedDeltaNet(nn.Module):
z_shape_og = z.shape
# reshape input data into 2D tensor
core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
z = z.reshape(-1, z.shape[-1])
if use_strided_prefill_z:
z_flat_shape = (z.numel() // z.shape[-1], z.shape[-1])
else:
z = z.reshape(-1, z.shape[-1])
z_flat_shape = z.shape
# Add padding for DP-Attn
if core_attn_out.shape != z.shape:
core_attn_out_pad = torch.zeros_like(z)
if core_attn_out.shape != z_flat_shape:
core_attn_out_pad = z.new_zeros(z_flat_shape)
core_attn_out_pad[: core_attn_out.shape[0], :] = core_attn_out
core_attn_out = core_attn_out_pad
test/registered/attention/test_gdn_prefill_layout.py
+145
View File
@@ -0,0 +1,145 @@
import unittest
import torch
from sglang.kernels.ops.attention.fla.l2norm import (
gdn_prefill_qkv_prepare_fwd,
l2norm_fwd,
)
from sglang.kernels.ops.attention.fla.layernorm_gated import rms_norm_gated
from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
fused_qkv_split_gdn_prefill,
qwen3_5_gdn_prefill_projection_views,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-large")
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestGdnPrefillLayout(unittest.TestCase):
TOKENS = 33
NUM_QK_HEADS = 2
NUM_V_HEADS = 4
HEAD_DIM = 128
def _projection_views(self, dtype):
qkv_dim = (
2 * self.NUM_QK_HEADS * self.HEAD_DIM + self.NUM_V_HEADS * self.HEAD_DIM
)
qkvz = torch.randn(
self.TOKENS,
qkv_dim + self.NUM_V_HEADS * self.HEAD_DIM,
dtype=dtype,
device="cuda",
)
ba = torch.randn(
self.TOKENS,
2 * self.NUM_V_HEADS,
dtype=dtype,
device="cuda",
)
return (qkvz, ba), qwen3_5_gdn_prefill_projection_views(
qkvz,
ba,
self.NUM_QK_HEADS,
self.NUM_V_HEADS,
self.HEAD_DIM,
self.HEAD_DIM,
)
def test_projection_views_preserve_layout(self):
(qkvz, ba), (mixed_qkv, z, b, a) = self._projection_views(torch.bfloat16)
qkv_dim = mixed_qkv.shape[1]
self.assertFalse(mixed_qkv.is_contiguous())
self.assertFalse(z.is_contiguous())
self.assertFalse(b.is_contiguous())
self.assertFalse(a.is_contiguous())
torch.testing.assert_close(mixed_qkv, qkvz[:, :qkv_dim], rtol=0, atol=0)
torch.testing.assert_close(
z.reshape(self.TOKENS, -1), qkvz[:, qkv_dim:], rtol=0, atol=0
)
torch.testing.assert_close(b, ba[:, : self.NUM_V_HEADS], rtol=0, atol=0)
torch.testing.assert_close(a, ba[:, self.NUM_V_HEADS :], rtol=0, atol=0)
def test_qkv_prepare_preserves_dtype_and_matches_materialized_path(self):
for dtype in (torch.bfloat16, torch.float16):
with self.subTest(dtype=dtype):
_, (mixed_qkv, _, _, _) = self._projection_views(dtype)
q_dim = self.NUM_QK_HEADS * self.HEAD_DIM
v_dim = self.NUM_V_HEADS * self.HEAD_DIM
q = mixed_qkv[:, :q_dim].view(
self.TOKENS, self.NUM_QK_HEADS, self.HEAD_DIM
)
k = mixed_qkv[:, q_dim : 2 * q_dim].view(
self.TOKENS, self.NUM_QK_HEADS, self.HEAD_DIM
)
v = mixed_qkv[:, 2 * q_dim : 2 * q_dim + v_dim].view(
self.TOKENS, self.NUM_V_HEADS, self.HEAD_DIM
)
q_out, k_out, v_out = gdn_prefill_qkv_prepare_fwd(q, k, v)
self.assertEqual(q_out.dtype, dtype)
self.assertEqual(k_out.dtype, dtype)
self.assertEqual(v_out.dtype, dtype)
torch.testing.assert_close(
q_out, l2norm_fwd(q.contiguous()), rtol=0, atol=0
)
torch.testing.assert_close(
k_out, l2norm_fwd(k.contiguous()), rtol=0, atol=0
)
torch.testing.assert_close(v_out, v.contiguous(), rtol=0, atol=0)
def test_fused_split_flashinfer_prepare_reuses_contiguous_value(self):
_, (mixed_qkv, _, _, _) = self._projection_views(torch.bfloat16)
q, k, v = fused_qkv_split_gdn_prefill(
mixed_qkv,
self.NUM_QK_HEADS,
self.NUM_QK_HEADS,
self.NUM_V_HEADS,
self.HEAD_DIM,
self.HEAD_DIM,
self.HEAD_DIM,
)
q_out, k_out, v_out = gdn_prefill_qkv_prepare_fwd(q[0], k[0], v[0])
self.assertEqual(v_out.data_ptr(), v.data_ptr())
torch.testing.assert_close(q_out, l2norm_fwd(q[0]), rtol=0, atol=0)
torch.testing.assert_close(k_out, l2norm_fwd(k[0]), rtol=0, atol=0)
def test_strided_gate_matches_contiguous_gate(self):
for dtype in (torch.bfloat16, torch.float16):
for norm_before_gate in (True, False):
with self.subTest(dtype=dtype, norm_before_gate=norm_before_gate):
_, (_, z, _, _) = self._projection_views(dtype)
x = torch.randn(
self.TOKENS * self.NUM_V_HEADS,
self.HEAD_DIM,
dtype=dtype,
device="cuda",
)
weight = torch.randn(self.HEAD_DIM, dtype=dtype, device="cuda")
expected = rms_norm_gated(
x=x,
weight=weight,
bias=None,
z=z.contiguous().view_as(x),
norm_before_gate=norm_before_gate,
is_rms_norm=True,
)
actual = rms_norm_gated(
x=x,
weight=weight,
bias=None,
z=z,
norm_before_gate=norm_before_gate,
is_rms_norm=True,
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py
+1 -1
View File
@@ -62,7 +62,7 @@ class TestFlashInferGDNAlignment(unittest.TestCase):
with mock.patch(
"sglang.kernels.ops.attention.fla.l2norm.l2norm_fwd",
side_effect=lambda tensor: tensor,
side_effect=lambda tensor, eps: tensor,
):
result, _, checkpoints = kernel.extend(
q,