diff --git a/benchmark/bench_linear_attention/bench_gdn_qkv_split.py b/benchmark/bench_linear_attention/bench_gdn_qkv_split.py new file mode 100644 index 000000000..30aaf9e26 --- /dev/null +++ b/benchmark/bench_linear_attention/bench_gdn_qkv_split.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import argparse + +import torch + +from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill + +DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Benchmark GDN prefill QKV split fallback vs fused Triton path." + ) + parser.add_argument("--seq-len", type=int, default=8192) + parser.add_argument("--num-q-heads", type=int, default=16) + parser.add_argument("--num-k-heads", type=int, default=16) + parser.add_argument("--num-v-heads", type=int, default=16) + parser.add_argument("--head-q", type=int, default=128) + parser.add_argument("--head-k", type=int, default=128) + parser.add_argument("--head-v", type=int, default=128) + parser.add_argument("--dtype", choices=DTYPES.keys(), default="bf16") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + return parser.parse_args() + + +def make_non_contiguous_view(src: torch.Tensor) -> torch.Tensor: + backing = torch.empty( + src.shape[1], + src.shape[0], + dtype=src.dtype, + device=src.device, + ) + view = backing.transpose(0, 1) + view.copy_(src) + return view + + +def split_reference( + mixed_qkv: torch.Tensor, + num_q_heads: int, + num_k_heads: int, + num_v_heads: int, + head_q: int, + head_k: int, + head_v: int, +): + q_dim = num_q_heads * head_q + k_dim = num_k_heads * head_k + v_dim = num_v_heads * head_v + actual_seq_len = mixed_qkv.shape[0] + query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + query = query.reshape(1, actual_seq_len, num_q_heads, head_q).contiguous() + key = key.reshape(1, actual_seq_len, num_k_heads, head_k).contiguous() + value = value.reshape(1, actual_seq_len, num_v_heads, head_v).contiguous() + return query, key, value + + +@torch.inference_mode() +def benchmark(fn, warmup: int, iters: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def check_close(actual, expected): + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor, rtol=0, atol=0) + + +def run_case(name: str, mixed_qkv: torch.Tensor, args): + shape_args = ( + args.num_q_heads, + args.num_k_heads, + args.num_v_heads, + args.head_q, + args.head_k, + args.head_v, + ) + expected = split_reference(mixed_qkv, *shape_args) + actual = fused_qkv_split_gdn_prefill(mixed_qkv, *shape_args) + check_close(actual, expected) + + baseline_us = benchmark( + lambda: split_reference(mixed_qkv, *shape_args), + args.warmup, + args.iters, + ) + fused_us = benchmark( + lambda: fused_qkv_split_gdn_prefill(mixed_qkv, *shape_args), + args.warmup, + args.iters, + ) + speedup = baseline_us / fused_us + print(f"{name:>12} {baseline_us:12.2f} {fused_us:12.2f} {speedup:10.2f}x") + + +def main(): + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark") + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = DTYPES[args.dtype] + qkv_dim = ( + args.num_q_heads * args.head_q + + args.num_k_heads * args.head_k + + args.num_v_heads * args.head_v + ) + mixed_qkv = torch.randn(args.seq_len, qkv_dim, dtype=dtype, device=device) + mixed_qkv_strided = make_non_contiguous_view(mixed_qkv) + + print( + f"seq_len={args.seq_len} qkv_dim={qkv_dim} dtype={args.dtype} " + f"warmup={args.warmup} iters={args.iters}" + ) + print(f"{'layout':>12} {'baseline_us':>12} {'fused_us':>12} {'speedup':>11}") + run_case("contiguous", mixed_qkv, args) + run_case("strided", mixed_qkv_strided, args) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/jit_kernel/triton/gdn_fused_proj.py b/python/sglang/jit_kernel/triton/gdn_fused_proj.py index d7d07da73..076a1cc3f 100644 --- a/python/sglang/jit_kernel/triton/gdn_fused_proj.py +++ b/python/sglang/jit_kernel/triton/gdn_fused_proj.py @@ -308,3 +308,100 @@ def fused_qkvzba_split_reshape_cat_contiguous( num_stages=3, ) return mixed_qkv, z, b, a + + +@triton.jit +def fused_qkv_split_gdn_prefill_kernel( + q, + k, + v, + mixed_qkv, + MIXED_QKV_STRIDE_T: tl.constexpr, + MIXED_QKV_STRIDE_D: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_K_HEADS: tl.constexpr, + NUM_V_HEADS: tl.constexpr, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + HEAD_V: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + i_t = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + + q_dim: tl.constexpr = NUM_Q_HEADS * HEAD_Q + k_dim: tl.constexpr = NUM_K_HEADS * HEAD_K + v_dim: tl.constexpr = NUM_V_HEADS * HEAD_V + qk_dim: tl.constexpr = q_dim + k_dim + qkv_dim: tl.constexpr = qk_dim + v_dim + + mask = offsets < qkv_dim + values = tl.load( + mixed_qkv + i_t * MIXED_QKV_STRIDE_T + offsets * MIXED_QKV_STRIDE_D, + mask=mask, + ) + + q_mask = offsets < q_dim + tl.store(q + i_t * q_dim + offsets, values, mask=q_mask) + + k_offsets = offsets - q_dim + k_mask = (offsets >= q_dim) & (offsets < qk_dim) + tl.store(k + i_t * k_dim + k_offsets, values, mask=k_mask) + + v_offsets = offsets - qk_dim + v_mask = (offsets >= qk_dim) & (offsets < qkv_dim) + tl.store(v + i_t * v_dim + v_offsets, values, mask=v_mask) + + +def fused_qkv_split_gdn_prefill( + mixed_qkv: torch.Tensor, + num_q_heads: int, + num_k_heads: int, + num_v_heads: int, + head_q: int, + head_k: int, + head_v: int, +): + """Split packed post-conv GDN QKV into contiguous FLA prefill tensors. + + `mixed_qkv` is laid out per token as `[all_q | all_k | all_v]`. The FLA + chunk kernels consume separate contiguous `[1, T, H, D]` tensors, so this + fused split replaces three independent `aten::copy_` kernels from the + generic FLA input guard. `mixed_qkv` may be a strided `[T, qkv_dim]` view. + """ + seq_len = mixed_qkv.shape[0] + q = torch.empty( + (1, seq_len, num_q_heads, head_q), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + k = torch.empty( + (1, seq_len, num_k_heads, head_k), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + v = torch.empty( + (1, seq_len, num_v_heads, head_v), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + + qkv_dim = num_q_heads * head_q + num_k_heads * head_k + num_v_heads * head_v + fused_qkv_split_gdn_prefill_kernel[(seq_len,)]( + q, + k, + v, + mixed_qkv, + mixed_qkv.stride(0), + mixed_qkv.stride(1), + num_q_heads, + num_k_heads, + num_v_heads, + head_q, + head_k, + head_v, + BLOCK_SIZE=triton.next_power_of_2(qkv_dim), + num_warps=8, + num_stages=3, + ) + return q, k, v diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index 5feb33906..fec725d01 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +import os from typing import Optional, Tuple import torch @@ -20,6 +21,9 @@ from sglang.srt.layers.attention.fla.utils import ( NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] CHUNK_SIZE = 64 +GDN_CHUNK_H_BV = int(os.getenv("SGLANG_GDN_CHUNK_H_BV", "32")) +GDN_CHUNK_H_NUM_WARPS = int(os.getenv("SGLANG_GDN_CHUNK_H_NUM_WARPS", "4")) +GDN_CHUNK_H_NUM_STAGES = int(os.getenv("SGLANG_GDN_CHUNK_H_NUM_STAGES", "2")) @triton.autotune( @@ -32,8 +36,16 @@ CHUNK_SIZE = 64 # because cloning the cache pool for each benchmark exceeds available memory. # NT_BUCKET is kept in the autotune key for forward-compatibility (allows # future per-bucket configs once the kernel is refactored to write final - # state to a separate output buffer). - configs=[triton.Config({"BV": 32}, num_warps=4, num_stages=2)], + # state to a separate output buffer). The env knobs keep this single-config + # property while allowing model/hardware-local validation of the selected + # tile without corrupting the state pool through multi-config autotune. + configs=[ + triton.Config( + {"BV": GDN_CHUNK_H_BV}, + num_warps=GDN_CHUNK_H_NUM_WARPS, + num_stages=GDN_CHUNK_H_NUM_STAGES, + ) + ], key=["H", "K", "V", "BT", "USE_GK", "NT_BUCKET"], **autotune_cache_kwargs, ) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index a10fcfa8e..ade156e98 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -26,6 +26,11 @@ if not is_cpu(): CHUNK_SIZE as FLA_CHUNK_SIZE, ) +if is_cuda(): + from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill + +MAX_FUSED_QKV_SPLIT_DIM = 8192 + if is_cuda(): from sglang.srt.layers.attention.mamba.causal_conv1d import ( causal_conv1d_fn as causal_conv1d_fn_cuda, @@ -444,16 +449,27 @@ class GDNAttnBackend(MambaAttnBackendBase): seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ).transpose(0, 1)[:seq_len] - query, key, value = torch.split( - mixed_qkv, - [layer.q_dim, layer.k_dim, layer.v_dim], - dim=-1, - ) - - actual_seq_len = query.shape[0] - query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim) - key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim) - value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) + actual_seq_len = mixed_qkv.shape[0] + qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim + if is_cuda() and qkv_dim <= MAX_FUSED_QKV_SPLIT_DIM: + query, key, value = fused_qkv_split_gdn_prefill( + mixed_qkv, + layer.num_q_heads, + layer.num_k_heads, + layer.num_v_heads, + layer.head_q_dim, + layer.head_k_dim, + layer.head_v_dim, + ) + else: + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim) + key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim) + value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) if is_target_verify: core_attn_out = self.kernel_dispatcher.target_verify(