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

Fix gate stride for 4D decode layouts (#31231)

Browse Source 
Co-authored-by: lmzheng <lmzheng@fb.com>
Co-authored-by: michael604work <michael604@meta.com>
This commit is contained in:
Lianmin Zheng
2026-07-14 20:06:50 -07:00
committed by GitHub
co-authored by lmzheng michael604work
parent b4fdce3b63
commit 4ae9cc3c81
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py
+3 -1
View File
@@ -281,7 +281,9 @@ def fused_sigmoid_gating_delta_rule_update(
# Both paths (KDA/GDN) advance p_a once per token, so use the token-axis stride.
# For 2D a ([T, ...]) this is stride(0); for 3D a ([B, T, ...]) this is stride(1).
# Using stride()[-2] covers GDN [T, HV] and KDA layouts ([T, HV*K] / [B, T, HV*K]).
stride_a = a.stride()[-2]
# KDA decode also passes 4-D [B, T, H, K], where [-2] is the head stride, not the
# token stride; take dim 1 explicitly for that layout.
stride_a = a.stride()[1] if a.ndim == 4 else a.stride()[-2]
HV = v.shape[2]
N = B if cu_seqlens is None else len(cu_seqlens) - 1
BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)