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

[Kernel] Migrate linear-attention, MiniMax-sparse and diffusion kernels to sglang.kernels (RFC #29630, Phase 2.5, 6/7) (#30793)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-15 11:21:36 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent ba5be86d42
commit c00131ebaa
44 changed files with 238 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py
+4 -4
View File
@@ -26,14 +26,14 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
import torch
from sglang.kernels.ops.attention.linear.gdn_blackwell import (
chunk_gated_delta_rule_cutedsl,
prepare_metadata_cutedsl,
)
from sglang.srt.layers.attention.fla.chunk import (
chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
)
from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
from sglang.srt.layers.attention.linear.kernels.gdn_blackwell import (
chunk_gated_delta_rule_cutedsl,
prepare_metadata_cutedsl,
)
# ---------------------------------------------------------------------------
# Helpers (shared shape: pool layout [N, H, K, V] with K-last stride)
benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py
+6 -6
View File
@@ -29,20 +29,20 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
import torch
import torch.nn.functional as F
from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda
from sglang.srt.layers.attention.linear.kernels.kda_blackwell import prepare_metadata
from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_h import (
from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import (
kda_h_cutedsl,
)
from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_kkt_inv_uw import (
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_kkt_inv_uw import (
kkt_inv_uw_cutedsl,
)
from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_o import (
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_o import (
kda_o_cutedsl,
)
from sglang.srt.layers.attention.linear.kernels.kda_blackwell.prologue import (
from sglang.kernels.ops.attention.linear.kda_blackwell.prologue import (
kda_prologue,
)
from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda
BT = 64 # chunk size