[GDN] Optimize prefill QKV split dispatch (#26206)

Co-authored-by: BBuf <xiaoyu.zhang@radixark.net>

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()