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

[GDN/KDA] Fuse SM100 CuteDSL prefill state I/O into the chunk h kernel (#30169)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-07-16 16:40:33 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent db7e6807de
commit a798a2aeea
8 changed files with 299 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/linear/gdn_blackwell/__init__.py
+19 -2
View File
@@ -153,6 +153,7 @@ def chunk_gated_delta_rule_cutedsl(
chunk_indices: torch.Tensor,
chunk_offsets: torch.Tensor,
core_attn_out: torch.Tensor | None = None,
initial_state_indices: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run the GDN chunk CuteDSL prefill kernels.
@@ -162,11 +163,17 @@ def chunk_gated_delta_rule_cutedsl(
v: Value tensor with shape ``[1, T, Hv, V]``.
g: Log-space decay tensor with shape ``[1, T, Hv]``.
beta: Delta-rule beta tensor with shape ``[1, T, Hv]``.
initial_state: Recurrent state with shape ``[N, Hv, V, K]``.
initial_state: Recurrent state with shape ``[N, Hv, V, K]``, or the
state POOL ``[num_slots, Hv, V, K]`` when ``initial_state_indices``
is given.
cu_seqlens: Cumulative sequence lengths with shape ``[N + 1]``.
chunk_indices: Chunk index metadata with shape ``[NT, 2]``.
chunk_offsets: Cumulative chunk offsets with shape ``[N + 1]``.
core_attn_out: Optional output buffer with shape ``[T, Hv, V]``.
initial_state_indices: Optional ``[N]`` int32 pool slots. When given,
the h kernel reads AND writes the pool rows in place (fused state
gather/scatter — no ``[N, Hv, V, K]`` intermediates) and the
returned ``final_state`` is the pool tensor itself.
Returns:
A tuple ``(output, final_state)`` where ``output`` has shape
@@ -213,7 +220,16 @@ def chunk_gated_delta_rule_cutedsl(
head_k_dim,
)
v_new = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
final_state = torch.empty_like(initial_state)
if initial_state_indices is None:
# Dense mode: preserve the return-fresh-final_state contract.
final_state = torch.empty_like(initial_state)
state_indices = torch.arange(
cu_seqlens.numel() - 1, device=q_3d.device, dtype=torch.int32
)
else:
# Pool mode: read and write the pool rows in place.
final_state = initial_state
state_indices = initial_state_indices
h_cutedsl(
k_3d,
u,
@@ -225,6 +241,7 @@ def chunk_gated_delta_rule_cutedsl(
final_state,
cu_seqlens,
chunk_offsets,
state_indices,
)
output = core_attn_out if core_attn_out is not None else torch.empty_like(v_3d)
python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py
+24 -6
View File
@@ -106,6 +106,7 @@ class Sm100ChunkHKernel:
ht: cute.Tensor,
cu_seqlens: cute.Tensor,
chunk_offsets: cute.Tensor,
state_indices: cute.Tensor,
stream: CUstream,
):
tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
@@ -119,7 +120,10 @@ class Sm100ChunkHKernel:
HT_args = self._make_h_tma_args(ht, tma_s2g)
H_args = self._make_h_tma_args(h, tma_s2g)
grid = (self.Hv, h0.shape[0], 1)
# h0/ht may be the full state pool ([num_slots, ...]) rather than a
# per-sequence gather, so the sequence count comes from cu_seqlens and
# each block resolves its state row through state_indices.
grid = (self.Hv, cu_seqlens.shape[0] - 1, 1)
block = (self.num_warps * 32, 1, 1)
self.kernel(
K_args,
@@ -132,6 +136,7 @@ class Sm100ChunkHKernel:
g_cu,
cu_seqlens,
chunk_offsets,
state_indices,
).launch(grid=grid, block=block, stream=stream)
@cute.kernel
@@ -147,6 +152,7 @@ class Sm100ChunkHKernel:
g_cu: cute.Tensor,
cu_seqlens: cute.Tensor,
chunk_offsets: cute.Tensor,
state_indices: cute.Tensor,
):
tid, _, _ = cute.arch.thread_idx()
head_id, seq_id, _ = cute.arch.block_idx()
@@ -220,6 +226,8 @@ class Sm100ChunkHKernel:
eos = cu_seqlens[seq_id + 1]
seqlen = eos - bos
num_chunks = cute.ceil_div(seqlen, BT)
# Row of h0/ht for this sequence (pool slot; fused state gather/scatter).
state_slot = state_indices[seq_id]
if warp_id == 9:
# TMA warp
@@ -234,7 +242,7 @@ class Sm100ChunkHKernel:
H0_size = V_dim * K_dim * self.h_dtype.width // 8
cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
simple_tma_copy(
H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
H0_tma_atom, tmaH0[state_slot, head_id, None, None], sH0, h0_mbar
)
# shape: ((BT, num_BT_tiles), (64, 2))
@@ -531,7 +539,7 @@ class Sm100ChunkHKernel:
cute.arch.barrier(barrier_id=1, number_of_threads=128)
if warp_id_ == 0:
ht_dst = tmaHT[seq_id, head_id, None, None]
ht_dst = tmaHT[state_slot, head_id, None, None]
simple_tma_copy(HT_tma_atom, sH0, ht_dst)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
@@ -676,6 +684,7 @@ class Sm100ChunkHKernel:
total_t = cute.sym_int()
pad_t = cute.sym_int()
total_chunks_n = cute.sym_int()
num_state_slots = cute.sym_int()
num_sequences = cute.sym_int()
cu_entries = cute.sym_int()
@@ -688,13 +697,14 @@ class Sm100ChunkHKernel:
BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
)
h0 = make_fake_tensor(
h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
h_dtype, (num_state_slots, Hv, V_dim, K_dim), divisibility=16
)
ht = make_fake_tensor(
h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
h_dtype, (num_state_slots, Hv, V_dim, K_dim), divisibility=16
)
cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
state_indices = make_fake_tensor(Int32, (num_sequences,), divisibility=1)
kernel = Sm100ChunkHKernel(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
@@ -710,6 +720,7 @@ class Sm100ChunkHKernel:
ht,
cu_seqlens,
chunk_offsets,
state_indices,
stream,
options="--enable-tvm-ffi",
)
@@ -726,10 +737,16 @@ def h_cutedsl(
ht: torch.Tensor,
cu_seqlens: torch.Tensor,
chunk_offsets: torch.Tensor,
state_indices: torch.Tensor,
BT: int = 64,
num_stages: int = 2,
) -> None:
"""Compute H/V_new with the same argument order as the CUDA wrapper."""
"""Compute H/V_new with the same argument order as the CUDA wrapper.
``h0``/``ht`` may be the full state pool; ``state_indices`` [N] int32 maps
each sequence to its row, so state gather/scatter fuses into the kernel's
TMA load/store (no per-call state intermediates).
"""
_, H, K_dim = K.shape
_, Hv, V_dim = V.shape
@@ -748,6 +765,7 @@ def h_cutedsl(
ht,
cu_seqlens,
chunk_offsets,
state_indices,
)
python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py
+35 -4
View File
@@ -112,15 +112,25 @@ def chunk_kda_cutedsl(
v: torch.Tensor, # [T, Hv, V] bf16
g: torch.Tensor, # [T, Hv, K] log-decay. RAW if A_log given, else pre-activated
beta: torch.Tensor, # [T, Hv] fp32, post-sigmoid
h0: torch.Tensor, # [N, Hv, V, K] (initial recurrent state, [V,K] layout)
h0: torch.Tensor, # [N, Hv, V, K] state, or the state POOL with h0_indices
cu_seqlens: torch.Tensor,
scale: float | None = None,
num_sms: int | None = None,
A_log: torch.Tensor | None = None, # [Hv]; if set, activate g internally
dt_bias: torch.Tensor | None = None, # [Hv, K] or [Hv*K]
lower_bound: float | None = None,
h0_indices: torch.Tensor | None = None, # [N] int32 pool slots
):
"""Run the KDA chunk gated-delta-rule prefill. Returns (o [T,Hv,V], ht [N,Hv,V,K])."""
"""Run the KDA chunk gated-delta-rule prefill. Returns (o [T,Hv,V], ht).
Dense mode (``h0_indices is None``): ``h0`` is [N, Hv, V, K]; the final state
is returned in a fresh ``ht`` and ``h0`` is left untouched.
Pool mode: ``h0`` is the state pool [num_slots, Hv, V, K] and ``h0_indices``
maps each sequence to its slot; the h kernel reads AND writes the pool rows
in place (fused state gather/scatter — no [N, Hv, V, K] intermediates), and
the returned ``ht`` is the pool tensor itself.
"""
import torch.nn.functional as F
T, Hv, K = q.shape
@@ -202,8 +212,29 @@ def chunk_kda_cutedsl(
V_new = ws["Vn"][:pad_t]
h_chunks = ws["hc"][:total]
ht = torch.empty_like(h0)
kda_h_cutedsl(KR, U, W, V_new, g_cu, h_chunks, h0, ht, cu_seqlens, chunk_offsets)
if h0_indices is None:
# Dense mode: preserve the return-fresh-ht contract.
ht = torch.empty_like(h0)
state_indices = torch.arange(
cu_seqlens.numel() - 1, device=q.device, dtype=torch.int32
)
else:
# Pool mode: read and write the pool rows in place.
ht = h0
state_indices = h0_indices
kda_h_cutedsl(
KR,
U,
W,
V_new,
g_cu,
h_chunks,
h0,
ht,
cu_seqlens,
chunk_offsets,
state_indices,
)
o = q.new_empty(T, Hv, V, dtype=torch.bfloat16)
kda_o_cutedsl(
python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py
+24 -7
View File
@@ -112,6 +112,7 @@ class Sm100KdaChunkHKernel:
ht: cute.Tensor,
cu_seqlens: cute.Tensor,
chunk_offsets: cute.Tensor,
state_indices: cute.Tensor,
stream: CUstream,
):
tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
@@ -125,7 +126,10 @@ class Sm100KdaChunkHKernel:
HT_args = self._make_h_tma_args(ht, tma_s2g)
H_args = self._make_h_tma_args(h, tma_s2g)
grid = (self.Hv, h0.shape[0], 1)
# h0/ht may be the full state pool ([num_slots, ...]) rather than a
# per-sequence gather, so the sequence count comes from cu_seqlens and
# each block resolves its state row through state_indices.
grid = (self.Hv, cu_seqlens.shape[0] - 1, 1)
block = (self.num_warps * 32, 1, 1)
self.kernel(
K_args,
@@ -138,6 +142,7 @@ class Sm100KdaChunkHKernel:
g_cu,
cu_seqlens,
chunk_offsets,
state_indices,
).launch(grid=grid, block=block, stream=stream)
@cute.kernel
@@ -153,6 +158,7 @@ class Sm100KdaChunkHKernel:
g_cu: cute.Tensor,
cu_seqlens: cute.Tensor,
chunk_offsets: cute.Tensor,
state_indices: cute.Tensor,
):
tid, _, _ = cute.arch.thread_idx()
head_id, seq_id, _ = cute.arch.block_idx()
@@ -226,6 +232,8 @@ class Sm100KdaChunkHKernel:
eos = cu_seqlens[seq_id + 1]
seqlen = eos - bos
num_chunks = cute.ceil_div(seqlen, BT)
# Row of h0/ht for this sequence (pool slot; fused state gather/scatter).
state_slot = state_indices[seq_id]
if warp_id == 9:
# TMA warp
@@ -239,7 +247,7 @@ class Sm100KdaChunkHKernel:
H0_size = V_dim * K_dim * self.h_dtype.width // 8
cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
simple_tma_copy(
H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
H0_tma_atom, tmaH0[state_slot, head_id, None, None], sH0, h0_mbar
)
gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
@@ -514,7 +522,7 @@ class Sm100KdaChunkHKernel:
cute.arch.barrier(barrier_id=1, number_of_threads=128)
if warp_id_ == 0:
ht_dst = tmaHT[seq_id, head_id, None, None]
ht_dst = tmaHT[state_slot, head_id, None, None]
simple_tma_copy(HT_tma_atom, sH0, ht_dst)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
@@ -628,6 +636,7 @@ class Sm100KdaChunkHKernel:
total_t = cute.sym_int()
pad_t = cute.sym_int()
total_chunks_n = cute.sym_int()
num_state_slots = cute.sym_int()
num_sequences = cute.sym_int()
cu_entries = cute.sym_int()
@@ -640,13 +649,14 @@ class Sm100KdaChunkHKernel:
BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
)
h0 = make_fake_tensor(
h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
h_dtype, (num_state_slots, Hv, V_dim, K_dim), divisibility=16
)
ht = make_fake_tensor(
h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
h_dtype, (num_state_slots, Hv, V_dim, K_dim), divisibility=16
)
cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
state_indices = make_fake_tensor(Int32, (num_sequences,), divisibility=1)
kernel = Sm100KdaChunkHKernel(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
@@ -662,6 +672,7 @@ class Sm100KdaChunkHKernel:
ht,
cu_seqlens,
chunk_offsets,
state_indices,
stream,
options="--enable-tvm-ffi",
)
@@ -678,13 +689,19 @@ def kda_h_cutedsl(
ht: torch.Tensor,
cu_seqlens: torch.Tensor,
chunk_offsets: torch.Tensor,
state_indices: torch.Tensor,
BT: int = 64,
num_stages: int = 2,
) -> None:
"""KDA chunk-state kernel. `kg` = per-channel pre-scaled key [T, Hv, K]."""
"""KDA chunk-state kernel. `kg` = per-channel pre-scaled key [T, Hv, K].
``h0``/``ht`` may be the full state pool; ``state_indices`` [N] int32 maps
each sequence to its row, so state gather/scatter fuses into the kernel's
TMA load/store (no per-call state intermediates).
"""
_, Hv, K_dim = kg.shape
_, _, V_dim = V.shape
h_dtype = {torch.bfloat16: BFloat16, torch.float32: Float32}[h0.dtype]
Sm100KdaChunkHKernel.compile(Hv, Hv, K_dim, V_dim, h_dtype, BT, num_stages)(
kg, V, W, V_new, g_cu, h, h0, ht, cu_seqlens, chunk_offsets
kg, V, W, V_new, g_cu, h, h0, ht, cu_seqlens, chunk_offsets, state_indices
)
python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py
+9 -12
View File
@@ -137,38 +137,35 @@ class CuteDSLGDNKernel(LinearAttnKernelBase):
cu_seqlens = query_start_loc.to(torch.int32)
# Pool gather: remap padding (-1) to the last (sentinel) slot.
# Pool state I/O is fused into the h kernel's TMA load/store: pass the
# pool + per-seq slots and the kernel reads h0/writes ht in place at
# those rows (no gather/scatter kernels, no [N, Hv, V, K] intermediates).
# Remap padding (-1) to the last (sentinel) slot.
ssm_cache_indices = torch.where(
cache_indices >= 0,
cache_indices,
ssm_states.shape[0] - 1,
).to(torch.long)
initial_state = ssm_states[ssm_cache_indices].contiguous()
).to(torch.int32)
chunk_indices, chunk_offsets = self._prepare_meta_fn(
cu_seqlens, total_seq_len, chunk_size=64
)
output, final_state = self._extend_fn(
output, _ = self._extend_fn(
q=q_norm,
k=k_norm,
v=v_in,
g=g_in,
beta=beta_in,
initial_state=initial_state,
initial_state=ssm_states,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_offsets=chunk_offsets,
)
ssm_states.index_copy_(
0,
ssm_cache_indices,
final_state.to(ssm_states.dtype),
initial_state_indices=ssm_cache_indices,
)
# Match Triton extend interface: (output, last_recurrent_state, h).
# We've already written state back, so no need to return it.
# The kernel already wrote state back into the pool in place.
return output, None, None
def target_verify(self, *args, **kwargs):
python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py
+8 -6
View File
@@ -120,27 +120,29 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
beta_in = beta[0][:num_tokens].to(torch.float32)
cu_seqlens = query_start_loc.to(torch.int32)
# Pool gather: remap padding (-1) to the last (sentinel) slot. State is
# Pool state I/O is fused into the h kernel's TMA load/store: pass the
# pool + per-seq slots and the kernel reads h0/writes ht in place at
# those rows (no gather/scatter kernels, no [N, HV, V, K] intermediates).
# Remap padding (-1) to the last (sentinel) slot. State is
# [slots, HV, V, K] == cutedsl [V,K] layout, no transpose needed.
ssm_cache_indices = torch.where(
cache_indices >= 0, cache_indices, ssm_states.shape[0] - 1
).to(torch.long)
initial_state = ssm_states[ssm_cache_indices].contiguous()
).to(torch.int32)
o, final_state = self._extend_fn(
o, _ = self._extend_fn(
q_n,
k_n,
v_in,
g_in,
beta_in,
initial_state,
ssm_states,
cu_seqlens,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
h0_indices=ssm_cache_indices,
)
ssm_states.index_copy_(0, ssm_cache_indices, final_state.to(ssm_states.dtype))
# Match chunk_kda's output layout [1, T, HV, V].
return o.unsqueeze(0)