From db8f3cdd11eba31258157ce30912a30200361cae Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Thu, 6 Aug 2026 18:48:47 -0700 Subject: [PATCH] fix(gdn): skip the -1 padding sentinel in the chunked extend kernel (#33810) --- python/sglang/kernels/ops/attention/fla/chunk_delta_h.py | 7 +++++-- test/registered/8-gpu-models/test_qwen35.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py index 7f79c6be6..2fe5e623d 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py @@ -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 # overflow an int32 index product. 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 ht = initial_state + index * stride_init_state if USE_INITIAL_STATE: @@ -127,7 +130,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht = ht + i_h * V * K # 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)) b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) 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)) # 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)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 64: diff --git a/test/registered/8-gpu-models/test_qwen35.py b/test/registered/8-gpu-models/test_qwen35.py index 32ca62bd0..6d9f091f9 100644 --- a/test/registered/8-gpu-models/test_qwen35.py +++ b/test/registered/8-gpu-models/test_qwen35.py @@ -30,7 +30,7 @@ class TestQwen35(unittest.TestCase): "--tool-call-parser=qwen3_coder", "--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 = [ "--speculative-algorithm=EAGLE", "--speculative-num-steps=3",