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

fix(gdn): skip the -1 padding sentinel in the chunked extend kernel (#33810)

Browse Source
This commit is contained in:
Yuwei An
2026-08-06 18:48:47 -07:00
committed by GitHub
parent 1e08b865f9
commit db8f3cdd11
2 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/chunk_delta_h.py
+5 -2
View File
@@ -119,6 +119,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
# per-slot pitch spans ALL layers' state, not H*V*K. int64: envelope pitches # per-slot pitch spans ALL layers' state, not H*V*K. int64: envelope pitches
# overflow an int32 index product. # overflow an int32 index product.
index = tl.load(initial_state_indices + i_n).to(tl.int64) index = tl.load(initial_state_indices + i_n).to(tl.int64)
# Padded rows carry the -1 sentinel; the decode kernel guards on it
# (fused_recurrent.py), the chunked extend path did not.
valid_state = index >= 0
h0 = initial_state + index * stride_init_state h0 = initial_state + index * stride_init_state
ht = initial_state + index * stride_init_state ht = initial_state + index * stride_init_state
if USE_INITIAL_STATE: if USE_INITIAL_STATE:
@@ -127,7 +130,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
ht = ht + i_h * V * K ht = ht + i_h * V * K
# load initial state # load initial state
if USE_INITIAL_STATE: if USE_INITIAL_STATE and valid_state:
p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (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) b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
if K > 64: if K > 64:
@@ -290,7 +293,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
b_h4 += tl.trans(tl.dot(b_k, b_v)) b_h4 += tl.trans(tl.dot(b_k, b_v))
# epilogue # epilogue
if INPLACE_UPDATE: if INPLACE_UPDATE and valid_state:
p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (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)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
if K > 64: if K > 64:
test/registered/8-gpu-models/test_qwen35.py
+1 -1
View File
@@ -30,7 +30,7 @@ class TestQwen35(unittest.TestCase):
"--tool-call-parser=qwen3_coder", "--tool-call-parser=qwen3_coder",
"--mem-fraction-static=0.8", "--mem-fraction-static=0.8",
] ]
dp_args = ["--dp=8", "--enable-dp-attention", "--disable-prefill-cuda-graph"] dp_args = ["--dp=8", "--enable-dp-attention"]
mtp_args = [ mtp_args = [
"--speculative-algorithm=EAGLE", "--speculative-algorithm=EAGLE",
"--speculative-num-steps=3", "--speculative-num-steps=3",