Fix stale track rows corrupting conv checkpoints under the prefill graph (#34184)
This commit is contained in:
@@ -1608,6 +1608,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
s["extend_start_loc"][bs:r].fill_(self.raw_num_tokens)
|
s["extend_start_loc"][bs:r].fill_(self.raw_num_tokens)
|
||||||
s["req_pool_indices"][bs:r].zero_()
|
s["req_pool_indices"][bs:r].zero_()
|
||||||
s["orig_seq_lens"][bs:r].zero_()
|
s["orig_seq_lens"][bs:r].zero_()
|
||||||
|
# The captured track scatter reads these rows too, and a stale
|
||||||
|
# mask row still carries a live destination slot: it would land
|
||||||
|
# this replay's window in an earlier request's checkpoint.
|
||||||
|
registry = self.buffer_registry
|
||||||
|
for name in ("mamba_track_mask", "mamba_track_indices"):
|
||||||
|
if registry.has_slot(name):
|
||||||
|
registry.get_slot(name).buffer[bs:r].zero_()
|
||||||
|
|
||||||
# Refresh the static buffer the captured graph reads from.
|
# Refresh the static buffer the captured graph reads from.
|
||||||
if (
|
if (
|
||||||
|
|||||||
@@ -163,7 +163,6 @@ class TestInklingSmallNvfp4Deterministic(CustomTestCase):
|
|||||||
"--mem-fraction-static",
|
"--mem-fraction-static",
|
||||||
"0.85",
|
"0.85",
|
||||||
"--enable-deterministic-inference",
|
"--enable-deterministic-inference",
|
||||||
"--disable-prefill-cuda-graph",
|
|
||||||
],
|
],
|
||||||
env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
|
env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user