Fix stale track rows corrupting conv checkpoints under the prefill graph (#34184)

2026-08-10 08:57:23 +08:00
parent 56ef810cad
commit 25b7015064
2 changed files with 7 additions and 1 deletions
@@ -1608,6 +1608,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                s["extend_start_loc"][bs:r].fill_(self.raw_num_tokens)
                s["req_pool_indices"][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.
        if (