From 07199fa220cc353a2b5230b8bf47a0dead7ea96c Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:32:44 -0700 Subject: [PATCH] [Performance] Optimize Qwen3.5 GDN prefill projection layouts (#36267) --- .../kernels/ops/attention/fla/l2norm.py | 161 ++++++++++++++++++ .../ops/attention/fla/layernorm_gated.py | 47 ++++- .../ops/attention/triton_gdn_fused_proj.py | 18 ++ .../linear/kernels/gdn_flashinfer.py | 8 +- python/sglang/srt/models/qwen3_5.py | 22 ++- .../attention/test_gdn_prefill_layout.py | 145 ++++++++++++++++ .../test_gdn_flashinfer_alignment.py | 2 +- 7 files changed, 386 insertions(+), 17 deletions(-) create mode 100644 test/registered/attention/test_gdn_prefill_layout.py diff --git a/python/sglang/kernels/ops/attention/fla/l2norm.py b/python/sglang/kernels/ops/attention/fla/l2norm.py index a55e322c3..30b445d3c 100644 --- a/python/sglang/kernels/ops/attention/fla/l2norm.py +++ b/python/sglang/kernels/ops/attention/fla/l2norm.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py index 05631f006..c0de08c45 100644 --- a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py +++ b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py @@ -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() diff --git a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py index 96f741de2..2c6096531 100644 --- a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py +++ b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py @@ -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. diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index 27b4c948b..e66cee0ab 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 924b70ce8..2de7eeee9 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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 diff --git a/test/registered/attention/test_gdn_prefill_layout.py b/test/registered/attention/test_gdn_prefill_layout.py new file mode 100644 index 000000000..f6c9955f6 --- /dev/null +++ b/test/registered/attention/test_gdn_prefill_layout.py @@ -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() diff --git a/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py b/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py index 6f0a44437..a05ddb63e 100644 --- a/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py +++ b/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py @@ -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,