chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[GDN] Change Attention State Layout from [N, HV, K, V] to [N, HV, V, K] (#20283)

Browse Source
This commit is contained in:
Yuan Luo
2026-03-12 10:53:12 +08:00
committed by GitHub
parent 8787cf4566
commit 649d6f2bc8
7 changed files with 364 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/attention/fla/chunk.py
View File

@@ -141,11 +141,11 @@ def chunk_gated_delta_rule(
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, H, K, V]` for `N` input sequences.
Initial state of shape `[N, H, V, K]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
@@ -157,7 +157,7 @@ def chunk_gated_delta_rule(
o (torch.Tensor):
Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
final_state (torch.Tensor):
Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
Examples::
>>> import torch

68
python/sglang/srt/layers/attention/fla/chunk_delta_h.py
View File

@@ -70,24 +70,24 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
NT = tl.cdiv(T, BT)
boh = i_n * NT
# [BK, BV]
b_h1 = tl.zeros([64, BV], dtype=tl.float32)
# [BV, BK]
b_h1 = tl.zeros([BV, 64], dtype=tl.float32)
if K > 64:
b_h2 = tl.zeros([64, BV], dtype=tl.float32)
b_h2 = tl.zeros([BV, 64], dtype=tl.float32)
if K > 128:
b_h3 = tl.zeros([64, BV], dtype=tl.float32)
b_h3 = tl.zeros([BV, 64], dtype=tl.float32)
if K > 192:
b_h4 = tl.zeros([64, BV], dtype=tl.float32)
b_h4 = tl.zeros([BV, 64], dtype=tl.float32)
# calculate offset
h += ((boh * H + i_h) * K * V).to(tl.int64)
h += ((boh * H + i_h) * V * K).to(tl.int64)
v += ((bos * H + i_h) * V).to(tl.int64)
k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)
w += ((bos * H + i_h) * K).to(tl.int64)
if SAVE_NEW_VALUE:
v_new += ((bos * H + i_h) * V).to(tl.int64)
stride_v = H * V
stride_h = H * K * V
stride_h = H * V * K
stride_k = Hg * K
stride_w = H * K
@@ -95,49 +95,49 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
h0 = initial_state + index * stride_h
ht = initial_state + index * stride_h
if USE_INITIAL_STATE:
h0 = h0 + i_h * K * V
h0 = h0 + i_h * V * K
if INPLACE_UPDATE:
ht = ht + i_h * K * V
ht = ht + i_h * V * K
# load initial state
if USE_INITIAL_STATE:
p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
if K > 64:
p_h0_2 = tl.make_block_ptr(
h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
)
b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)
if K > 128:
p_h0_3 = tl.make_block_ptr(
h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
)
b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)
if K > 192:
p_h0_4 = tl.make_block_ptr(
h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
)
b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)
# main recurrence
for i_t in range(NT):
p_h1 = tl.make_block_ptr(
h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)
h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)
)
tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
if K > 64:
p_h2 = tl.make_block_ptr(
h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
)
tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))
if K > 128:
p_h3 = tl.make_block_ptr(
h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
)
tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1))
if K > 192:
p_h4 = tl.make_block_ptr(
h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
)
tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
@@ -145,25 +145,25 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)
)
b_w = tl.load(p_w, boundary_check=(0, 1))
b_v = tl.dot(b_w, b_h1.to(b_w.dtype))
b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype))
if K > 64:
p_w = tl.make_block_ptr(
w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)
)
b_w = tl.load(p_w, boundary_check=(0, 1))
b_v += tl.dot(b_w, b_h2.to(b_w.dtype))
b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype))
if K > 128:
p_w = tl.make_block_ptr(
w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)
)
b_w = tl.load(p_w, boundary_check=(0, 1))
b_v += tl.dot(b_w, b_h3.to(b_w.dtype))
b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype))
if K > 192:
p_w = tl.make_block_ptr(
w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)
)
b_w = tl.load(p_w, boundary_check=(0, 1))
b_v += tl.dot(b_w, b_h4.to(b_w.dtype))
b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype))
p_v = tl.make_block_ptr(
v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
)
@@ -199,7 +199,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
mask=(o_k1 < K),
other=0.0,
)
b_h1 *= exp(b_gk_last1)[:, None]
b_h1 *= exp(b_gk_last1)[None, :]
if K > 64:
o_k2 = 64 + o_k1
b_gk_last2 = tl.load(
@@ -207,7 +207,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
mask=(o_k2 < K),
other=0.0,
)
b_h2 *= exp(b_gk_last2)[:, None]
b_h2 *= exp(b_gk_last2)[None, :]
if K > 128:
o_k3 = 128 + o_k1
b_gk_last3 = tl.load(
@@ -215,7 +215,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
mask=(o_k3 < K),
other=0.0,
)
b_h3 *= exp(b_gk_last3)[:, None]
b_h3 *= exp(b_gk_last3)[None, :]
if K > 192:
o_k4 = 192 + o_k1
b_gk_last4 = tl.load(
@@ -223,50 +223,50 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
mask=(o_k4 < K),
other=0.0,
)
b_h4 *= exp(b_gk_last4)[:, None]
b_h4 *= exp(b_gk_last4)[None, :]
b_v = b_v.to(k.dtype.element_ty)
p_k = tl.make_block_ptr(
k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
)
b_k = tl.load(p_k, boundary_check=(0, 1))
b_h1 += tl.dot(b_k, b_v)
b_h1 += tl.trans(tl.dot(b_k, b_v))
if K > 64:
p_k = tl.make_block_ptr(
k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
)
b_k = tl.load(p_k, boundary_check=(0, 1))
b_h2 += tl.dot(b_k, b_v)
b_h2 += tl.trans(tl.dot(b_k, b_v))
if K > 128:
p_k = tl.make_block_ptr(
k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)
)
b_k = tl.load(p_k, boundary_check=(0, 1))
b_h3 += tl.dot(b_k, b_v)
b_h3 += tl.trans(tl.dot(b_k, b_v))
if K > 192:
p_k = tl.make_block_ptr(
k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)
)
b_k = tl.load(p_k, boundary_check=(0, 1))
b_h4 += tl.dot(b_k, b_v)
b_h4 += tl.trans(tl.dot(b_k, b_v))
# epilogue
if INPLACE_UPDATE:
p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
if K > 64:
p_ht = tl.make_block_ptr(
ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
)
tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
if K > 128:
p_ht = tl.make_block_ptr(
ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)
ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
)
tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
if K > 192:
p_ht = tl.make_block_ptr(
ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)
ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
)
tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
@@ -302,7 +302,7 @@ def chunk_gated_delta_rule_fwd_h(
)
assert K <= 256, "current kernel does not support head dimension larger than 256."
h = k.new_empty(B, NT, H, K, V)
h = k.new_empty(B, NT, H, V, K)
v_new = torch.empty_like(u) if save_new_value else None

8
python/sglang/srt/layers/attention/fla/chunk_o.py
View File

@@ -71,7 +71,7 @@ def chunk_fwd_kernel_o(
k += (bos * Hg + i_h // (H // Hg)) * K
v += (bos * H + i_h) * V
o += (bos * H + i_h) * V
h += (i_tg * H + i_h).to(tl.int64) * K * V
h += (i_tg * H + i_h).to(tl.int64) * V * K
b_o = tl.zeros([BT, BV], dtype=tl.float32)
b_A = tl.zeros([BT, BT], dtype=tl.float32)
@@ -84,17 +84,17 @@ def chunk_fwd_kernel_o(
k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)
)
p_h = tl.make_block_ptr(
h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)
h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)
)
# [BT, BK]
b_q = tl.load(p_q, boundary_check=(0, 1))
# [BK, BT]
b_k = tl.load(p_k, boundary_check=(0, 1))
# [BK, BV]
# [BV, BK]
b_h = tl.load(p_h, boundary_check=(0, 1))
# [BT, BK] @ [BK, BV] -> [BT, BV]
b_o += tl.dot(b_q, b_h)
b_o += tl.dot(b_q, tl.trans(b_h))
# [BT, BK] @ [BK, BT] -> [BT, BT]
b_A += tl.dot(b_q, b_k)

64
python/sglang/srt/layers/attention/fla/fused_recurrent.py
View File

@@ -64,17 +64,17 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
if not IS_KDA:
p_g = g + bos * HV + i_hv
else:
p_gk = g + (bos * HV + i_hv) * K + o_k
p_gk = g + (bos * H + i_h) * K + o_k
p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
mask_k = o_k < K
mask_v = o_v < V
mask_h = mask_k[:, None] & mask_v[None, :]
mask_h = mask_v[:, None] & mask_k[None, :]
b_h = tl.zeros([BK, BV], dtype=tl.float32)
b_h = tl.zeros([BV, BK], dtype=tl.float32)
if USE_INITIAL_STATE:
p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
p_h0 = h0 + i_nh * V * K + o_v[:, None] * K + o_k[None, :]
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
for _ in range(0, T):
@@ -86,24 +86,24 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
b_q = b_q * scale
# [BK, BV]
# [BV, BK]
if not IS_KDA:
b_g = tl.load(p_g).to(tl.float32)
b_h *= exp(b_g)
else:
b_gk = tl.load(p_gk).to(tl.float32)
b_h *= exp(b_gk[:, None])
b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32)
b_h *= exp(b_gk[None, :])
# [BV]
b_v -= tl.sum(b_h * b_k[:, None], 0)
b_v -= tl.sum(b_h * b_k[None, :], 1)
if IS_BETA_HEADWISE:
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
else:
b_beta = tl.load(p_beta).to(tl.float32)
b_v *= b_beta
# [BK, BV]
b_h += b_k[:, None] * b_v[None, :]
# [BV, BK]
b_h += b_v[:, None] * b_k[None, :]
# [BV]
b_o = tl.sum(b_h * b_q[:, None], 0)
b_o = tl.sum(b_h * b_q[None, :], 1)
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
p_q += H * K
@@ -113,11 +113,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
if not IS_KDA:
p_g += HV
else:
p_gk += HV * K
p_gk += H * K
p_beta += HV * (V if IS_BETA_HEADWISE else 1)
if STORE_FINAL_STATE:
p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
p_ht = ht + i_nh * V * K + o_v[:, None] * K + o_k[None, :]
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
@@ -144,7 +144,7 @@ def fused_recurrent_gated_delta_rule_fwd(
o = q.new_empty(NK, *v.shape)
if output_final_state:
final_state = q.new_empty(N, HV, K, V, dtype=torch.float32)
final_state = q.new_empty(N, HV, V, K, dtype=torch.float32)
else:
final_state = None
@@ -252,11 +252,11 @@ def fused_recurrent_gated_delta_rule(
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, HV, K, V]` for `N` input sequences.
Initial state of shape `[N, HV, V, K]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
Whether to output the final state of shape `[N, HV, V, K]`. Default: `False`.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
@@ -264,7 +264,7 @@ def fused_recurrent_gated_delta_rule(
o (torch.Tensor):
Outputs of shape `[B, T, HV, V]`.
final_state (torch.Tensor):
Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
Final state of shape `[N, HV, V, K]` if `output_final_state=True` else `None`.
Examples::
>>> import torch
>>> import torch.nn.functional as F
@@ -277,7 +277,7 @@ def fused_recurrent_gated_delta_rule(
>>> v = torch.randn(B, T, HV, V, device='cuda')
>>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
>>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
>>> h0 = torch.randn(B, HV, K, V, device='cuda')
>>> h0 = torch.randn(B, HV, V, K, device='cuda')
>>> o, ht = fused_gated_recurrent_delta_rule(
q, k, v, g, beta,
initial_state=h0,
@@ -413,9 +413,9 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
mask_k = o_k < K
mask_v = o_v < V
mask_h = mask_k[:, None] & mask_v[None, :]
mask_h = mask_v[:, None] & mask_k[None, :]
b_h = tl.zeros([BK, BV], dtype=tl.float32)
b_h = tl.zeros([BV, BK], dtype=tl.float32)
if USE_INITIAL_STATE:
idx = tl.load(h0_indices + i_n)
# Add bounds checking for idx
@@ -424,8 +424,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
h0_source
+ idx * HV * K * V
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[:, None] * K
+ o_k[None, :]
)
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
@@ -449,8 +449,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
+ cache_idx * cache_steps * HV * K * V
+ step_offset
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[:, None] * K
+ o_k[None, :]
)
b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32)
@@ -466,17 +466,17 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
# [BK, BV]
b_h *= exp(b_g)
# [BV]
b_v -= tl.sum(b_h * b_k[:, None], 0)
b_v -= tl.sum(b_h * b_k[None, :], 1)
if IS_BETA_HEADWISE:
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
else:
b_beta = tl.load(p_beta).to(tl.float32)
b_v *= b_beta
# [BK, BV]
b_h += b_k[:, None] * b_v[None, :]
# [BV, BK]
b_h += b_v[:, None] * b_k[None, :]
# [BV]
if not DISABLE_OUTPUT_CALCULATION:
b_o = tl.sum(b_h * b_q[:, None], 0)
b_o = tl.sum(b_h * b_q[None, :], 1)
# core attn output
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
@@ -490,8 +490,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
+ cache_idx * cache_steps * HV * K * V
+ step_offset
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[:, None] * K
+ o_k[None, :]
)
tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
@@ -513,8 +513,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
h0_source
+ idx * HV * K * V
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[:, None] * K
+ o_k[None, :]
)
tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)

16
python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py
View File

@@ -99,8 +99,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
h0_source
+ idx * HV * K * V
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[None, :] * K
+ o_k[:, None]
)
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
@@ -137,8 +137,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
+ cache_idx * cache_steps * HV * K * V
+ step_offset
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[None, :] * K
+ o_k[:, None]
)
b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32)
@@ -207,8 +207,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
+ cache_idx * cache_steps * HV * K * V
+ step_offset
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[None, :] * K
+ o_k[:, None]
)
tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
@@ -234,8 +234,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
h0_source
+ idx * HV * K * V
+ i_hv * K * V
+ o_k[:, None] * V
+ o_v[None, :]
+ o_v[None, :] * K
+ o_k[:, None]
)
tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)

17
python/sglang/srt/layers/attention/fla/kda.py
View File

@@ -59,11 +59,11 @@ def fused_recurrent_kda_fwd(
num_stages = 3
num_warps = 1
o = torch.empty_like(k)
o = q.new_empty(NK, *v.shape)
if inplace_final_state:
final_state = initial_state
else:
final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)
final_state = q.new_empty(N, HV, V, K, dtype=initial_state.dtype)
stride_init_state_token = initial_state.stride(0)
stride_final_state_token = final_state.stride(0)
@@ -113,6 +113,7 @@ def fused_recurrent_kda_fwd(
num_stages=num_stages,
)
o = o.squeeze(0)
return o, final_state
@@ -756,11 +757,11 @@ def chunk_gla_fwd_kernel_o(
(1, 0),
)
p_h = tl.make_block_ptr(
h + (i_tg * H + i_h) * K * V,
(K, V),
(V, 1),
(i_k * BK, i_v * BV),
(BK, BV),
h + (i_tg * H + i_h) * V * K,
(V, K),
(K, 1),
(i_v * BV, i_k * BK),
(BV, BK),
(1, 0),
)
@@ -776,7 +777,7 @@ def chunk_gla_fwd_kernel_o(
# works but dkw, owing to divine benevolence
# [BT, BV]
if i_k >= 0:
b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))
b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
p_v = tl.make_block_ptr(
v + (bos * H + i_h) * V,
(T, V),

274
test/registered/attention/test_chunk_gated_delta_rule.py Normal file
View File

@@ -0,0 +1,274 @@
import unittest
import torch
from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
from sglang.srt.layers.attention.fla.fused_recurrent import (
fused_recurrent_gated_delta_rule,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu")
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestChunkGatedDeltaRule(unittest.TestCase):
"""Test chunk_gated_delta_rule against token-by-token fused_recurrent reference."""
ATOL = 2e-2
RTOL = 1e-2
def _run_reference(self, pool_init, cache_indices, q, k, v, g, beta):
"""Per-batch token-by-token reference using fused_recurrent_gated_delta_rule.
initial_state shape: [N, H, V, K] (native layout on this branch).
"""
B = cache_indices.shape[0]
T_per_seq = q.shape[1] // B
pool = pool_init.clone()
h_cur = pool[cache_indices].contiguous().clone()
o_list = []
for b in range(B):
sl = slice(b * T_per_seq, (b + 1) * T_per_seq)
o_b, h_b = fused_recurrent_gated_delta_rule(
q=q[0, sl].unsqueeze(0),
k=k[0, sl].unsqueeze(0),
v=v[0, sl].unsqueeze(0),
g=g[0, sl].unsqueeze(0),
beta=beta[0, sl].unsqueeze(0),
initial_state=h_cur[b : b + 1],
output_final_state=True,
use_qk_l2norm_in_kernel=True,
)
o_list.append(o_b)
h_cur[b] = h_b[0]
pool[cache_indices] = h_cur
return torch.cat(o_list, dim=1), pool
def _run_chunk(self, pool_init, cache_indices, q, k, v, g, beta, cu_seqlens):
"""Run chunk_gated_delta_rule with native [V, K] pool."""
pool = pool_init.clone()
o, _, _ = chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=pool,
initial_state_indices=cache_indices,
cu_seqlens=cu_seqlens,
head_first=False,
use_qk_l2norm_in_kernel=True,
)
return o, pool
def _check_shape(
self, B, T_per_seq, H, K, V, pool_size, sequential_indices=False, seed=42
):
"""Run correctness check for one (B, T_per_seq, H, K, V, pool_size) config."""
device = "cuda"
dtype = torch.bfloat16
T = B * T_per_seq
torch.manual_seed(seed)
if sequential_indices:
cache_indices = torch.arange(B, dtype=torch.int32, device=device)
else:
perm = torch.randperm(pool_size, device=device)[:B]
cache_indices = perm.to(torch.int32)
pool_init = (
torch.randn(pool_size, H, V, K, dtype=torch.float32, device=device) * 0.1
)
cu_seqlens = torch.zeros(B + 1, dtype=torch.long, device=device)
cu_seqlens[1:] = (
torch.arange(1, B + 1, dtype=torch.long, device=device) * T_per_seq
)
q = torch.randn(1, T, H, K, dtype=dtype, device=device)
k = torch.randn(1, T, H, K, dtype=dtype, device=device)
v = torch.randn(1, T, H, V, dtype=dtype, device=device)
g = torch.nn.functional.logsigmoid(
torch.randn(1, T, H, dtype=dtype, device=device)
)
beta = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device))
o_ref, pool_ref = self._run_reference(
pool_init, cache_indices, q, k, v, g, beta
)
o_new, pool_new = self._run_chunk(
pool_init, cache_indices, q, k, v, g, beta, cu_seqlens
)
self.assertTrue(
torch.allclose(
o_ref.float(), o_new.float(), atol=self.ATOL, rtol=self.RTOL
),
f"Output mismatch: max_diff="
f"{(o_ref.float() - o_new.float()).abs().max().item():.2e}",
)
ref_slots = pool_ref[cache_indices].contiguous()
new_slots = pool_new[cache_indices].contiguous()
self.assertTrue(
torch.allclose(
ref_slots.float(), new_slots.float(), atol=self.ATOL, rtol=self.RTOL
),
f"State mismatch: max_diff="
f"{(ref_slots.float() - new_slots.float()).abs().max().item():.2e}",
)
# ------------------------------------------------------------------
# Production-style configs (Qwen3-Next)
# ------------------------------------------------------------------
def test_production_nt1(self):
self._check_shape(B=4, T_per_seq=64, H=16, K=128, V=128, pool_size=32)
def test_production_nt2(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
def test_production_nt4(self):
self._check_shape(B=4, T_per_seq=256, H=16, K=128, V=128, pool_size=32)
# ------------------------------------------------------------------
# Batch size sweep
# ------------------------------------------------------------------
def test_batch_1(self):
self._check_shape(B=1, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
def test_batch_2(self):
self._check_shape(B=2, T_per_seq=128, H=16, K=128, V=128, pool_size=32)
def test_batch_8(self):
self._check_shape(B=8, T_per_seq=128, H=16, K=128, V=128, pool_size=64)
def test_batch_16(self):
self._check_shape(B=16, T_per_seq=64, H=16, K=128, V=128, pool_size=128)
def test_batch_32(self):
self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256)
# ------------------------------------------------------------------
# Head count sweep
# ------------------------------------------------------------------
def test_heads_4(self):
self._check_shape(B=4, T_per_seq=128, H=4, K=128, V=128, pool_size=32)
def test_heads_8(self):
self._check_shape(B=4, T_per_seq=128, H=8, K=128, V=128, pool_size=32)
def test_heads_32(self):
self._check_shape(B=4, T_per_seq=128, H=32, K=128, V=128, pool_size=32)
def test_heads_64(self):
self._check_shape(B=4, T_per_seq=128, H=64, K=128, V=128, pool_size=32)
# ------------------------------------------------------------------
# K != V (exercises that [V,K] != [K,V] byte-order matters)
# ------------------------------------------------------------------
def test_dim_64x64(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=64, pool_size=32)
def test_dim_k_lt_v(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=128, pool_size=32)
def test_dim_k_gt_v(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=64, pool_size=32)
def test_dim_256x256(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=256, V=256, pool_size=32)
# ------------------------------------------------------------------
# Short sequences (T < chunk_size=64)
# ------------------------------------------------------------------
def test_seqlen_1(self):
self._check_shape(B=4, T_per_seq=1, H=16, K=128, V=128, pool_size=32)
def test_seqlen_7(self):
self._check_shape(B=4, T_per_seq=7, H=16, K=128, V=128, pool_size=32)
def test_seqlen_16(self):
self._check_shape(B=4, T_per_seq=16, H=16, K=128, V=128, pool_size=32)
def test_seqlen_32(self):
self._check_shape(B=4, T_per_seq=32, H=16, K=128, V=128, pool_size=32)
# ------------------------------------------------------------------
# Multi-chunk and large pool
# ------------------------------------------------------------------
def test_multi_chunk_nt8(self):
self._check_shape(B=4, T_per_seq=512, H=16, K=128, V=128, pool_size=32)
def test_large_pool(self):
self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=512)
# ------------------------------------------------------------------
# Combined stress
# ------------------------------------------------------------------
def test_stress(self):
self._check_shape(B=32, T_per_seq=128, H=32, K=128, V=128, pool_size=256)
# ------------------------------------------------------------------
# Sequential-index variants (pool_size == B, indices = 0..B-1)
# ------------------------------------------------------------------
def test_seq_idx_b4(self):
self._check_shape(
B=4,
T_per_seq=128,
H=16,
K=128,
V=128,
pool_size=4,
sequential_indices=True,
)
def test_seq_idx_b8(self):
self._check_shape(
B=8,
T_per_seq=128,
H=16,
K=128,
V=128,
pool_size=8,
sequential_indices=True,
)
def test_seq_idx_h32(self):
self._check_shape(
B=4,
T_per_seq=128,
H=32,
K=128,
V=128,
pool_size=4,
sequential_indices=True,
)
def test_seq_idx_h64(self):
self._check_shape(
B=4,
T_per_seq=128,
H=64,
K=128,
V=128,
pool_size=4,
sequential_indices=True,
)
def test_seq_idx_stress(self):
self._check_shape(
B=32,
T_per_seq=128,
H=32,
K=128,
V=128,
pool_size=32,
sequential_indices=True,
)
if __name__ == "__main__":
unittest.main()