fix(gdn): skip the -1 padding sentinel in the chunked extend kernel (#33810)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user