Repository navigation
Conversation
| self._refresh_cache_indices() | ||
|
|
||
| def _prepare_track_indices(self, forward_batch: ForwardBatch): | ||
| indices = forward_batch.mamba_track_indices |
There was a problem hiding this comment.
[bug] Unified _prepare_slot_indices now always calls _prepare_track_indices, which does a bare forward_batch.mamba_track_indices load. Decode/prefill replay views include that field (None when tracking is off). EAGLE draft-extend replay does not: EAGLEDraftExtendCudaGraphRunner.replay and MultiLayerEagleDraftExtendCudaGraphRunner pass a SimpleNamespace with req_pool_indices / seq_lens / spec_info but no mamba_track_indices (python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py:609 and python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py:483). Inkling's hybrid wrapper still forwards init_forward_metadata_out_graph to the sidecar on DRAFT_EXTEND_V2 (the draft runs its own convs). The new load therefore raises AttributeError on unified-memory draft-extend graph replay and aborts the rest of _prepare_slot_indices, so active-slot translation is skipped too. Decode capture already documents the intended contract: the registry slot is the virtual source and the backend copies into its own static buffer.
Suggestion: Treat a missing field as None so cache-index prep still runs, and add mamba_track_indices (virtual registry slice, or None) to both draft-extend replay views the same way decode_cuda_graph_runner._replay_forward_batch does. If draft-extend graphs are supposed to refresh checkpoint destinations, the capture batch needs that slot as well; today capture builds a ForwardBatch with the default None, so the track-save is specialized out of the graph independently of this PR.
There was a problem hiding this comment.
Modifications.
1. Explicitly provide checkpoint inputs in both draft-extend runners : Added mamba_track_indices to capture, replay, and multi-layer staging. Disabled tracking uses explicit None; replay refreshes live IDs and padding, or clears the buffer to sink IDs while preserving its address.
2. Keep checkpoint destinations correct across translation and graph replay : Physical destinations now live in separate, fixed-address backend metadata, while ForwardBatch retains virtual IDs to prevent double translation. The draft accepted-window checkpoint path is preserved.
3. Additional fix found during validation: Propagate the lazy-allocation policy consistently.
Fixed a scheduler/request-pool policy mismatch that could send an unallocated -1 slot to cleanup. Added regressions for both two- and three-pool configurations.
Validation summary
- 74 tests and 125 subtests passed, covering CUDA capture/replay, checkpoint readback, compaction, padding, disabled tracking, staging, and production buffer/pool construction.
- Earlier serving validation: 41 requests across five reduced-Inkling/Qwen configurations, with no observed negative-slot frees or accounting failures.
3dbf5db to
273a48a
Compare
273a48a to
9c90ccf
Compare
ch-wan
left a comment
ch-wan
left a comment
There was a problem hiding this comment.
The post-verify translate, the draft-extend runner wiring, and the lazy-policy propagation each fill a real gap. The core fix, though, now overlaps #41144, which landed on main with the opposite convention for forward_batch.mamba_track_indices.
Suggested path forward. #41144 already fixes the target prefill/decode checkpoint writes, the corruption every Inkling + unified-memory prefix-cache user hits, and adds an end-to-end bit-exact regression test. So this PR shouldn't land as-is. Three gaps remain on main that this PR addresses. I'd rebase and split it into two focused PRs:
-
A: spec-path translation. Replace #41144's in-place rebind with this PR's convention:
ForwardBatchkeeps virtual ids, and physical ids live in backend metadata. Going by the code on main (not run on GPU), the rebind is fragile in two places:- The draft-extend replay views carry no
mamba_track_indices, but #41144 reads the field without a fallback in_prepare_slot_indices(Mamba2AttnBackendusesgetattrfor the same field). Inkling MTP + unified memory + draft-extend CUDA graph should therefore raiseAttributeError. - Its
data_ptrguard only recognizes its own backend's buffer. The per-step re-plan inmulti_layer_eagle_worker_v2.pytherefore translates a second time for every step after the first.
On top of that, add the
commit_conv_state_after_mtp_verifytranslate and the draft-extend view field. #41144 deferred the translate as DSPARK-only, but the EAGLE commit inspec_utils.pyalso callsupdate_mamba_state_after_mtp_verify, and I found nothing that restricts it to DSPARK. Drop the[:bs]slicing change unless it's needed. - The draft-extend replay views carry no
-
B: propagate the lazy policy to the unified factories, together with a
-1filter inStreamingSession._free_slot_mamba/session_held_mamba_slots. Today the scheduler runs lazy while the unified pool allocates both ping-pong slots up front.mamba_lazy_prealloc_at_boundarythen skips every boundary, and the second slot sits idle.
Please attach a full-model Inkling + unified-memory run to each PR. The numbers in the current description come from older bases (81b81664a6 / 5bebe7a033).
Inline comments: 1 blocker (double translation after the rebase), 1 medium (lazy mode + streaming session), 1 question, 4 nits/perf.
| def _prepare_track_indices(self, forward_batch: ForwardBatch): | ||
| # Every metadata view must explicitly provide virtual IDs or None. | ||
| indices = forward_batch.mamba_track_indices | ||
| if indices is None or self._slot_gather_recordable: | ||
| self.sconv_metadata.track_cache_indices = indices | ||
| return | ||
| n = indices.shape[0] | ||
| assert n <= self._graph_track_indices.shape[0], ( | ||
| "checkpoint-index buffer too small for the forward batch" | ||
| ) | ||
| # Keep the batch virtual across repeated eager metadata initialization | ||
| # and across draft backends. Captured consumers read this fixed buffer. | ||
| out = self._graph_track_indices[:n] | ||
| out.copy_(self._translate_mamba_indices(indices)) | ||
| self.sconv_metadata.track_cache_indices = out |
There was a problem hiding this comment.
[blocker] Rebase onto #41144: the merged tree translates twice.
#41144 (on main as da3eb6d) fixes the same bug differently. At the end of _prepare_slot_indices it rebinds forward_batch.mamba_track_indices to a physical buffer (_translate_track_indices / _track_indices_buf). This branch merges onto main with no conflict, but the result keeps both paths, and _prepare_track_indices assumes the field is still virtual:
- Eager multi-layer EAGLE draft extend:
multi_layer_eagle_worker_v2.pypre-plans runner[0]'s backend, then callsinit_forward_metadataagain on the sameforward_batchfor every step. From the second prep on,_prepare_track_indicesreads [Unified Memory] Fix Inkling conv-checkpoint track ids written to virtual slot numbers #41144's physical buffer and producesv2p[v2p[v]]. [Unified Memory] Fix Inkling conv-checkpoint track ids written to virtual slot numbers #41144'sdata_ptrguard only recognizes its own backend's buffer, so the other steps' backends translate again on that side too. - Graph path:
MultiLayerEagleMultiStepDraftExtendCudaGraphRunnerhands oneSimpleNamespaceto every backend in itsfor backend in backendsloop, so every backend after the first gets the same double translation.
After any compaction, checkpoints then land in another request's slot or in a free one.
This PR's convention (ForwardBatch stays virtual, physical ids live in backend metadata) is the one that holds up across multiple backends. So I'd delete #41144's _translate_track_indices / _track_indices_buf during the rebase rather than keep two translations. Keep the translate in commit_conv_state_after_mtp_verify: #41144 explicitly left it as a follow-up, so it isn't covered on main.
There was a problem hiding this comment.
Removed the in-place rebind and its translation buffer. ForwardBatch retains virtual IDs, and each backend prepares its own physical destinations. A focused multi-backend test reproduces the upstream failure and passes with this fix.
| max_num_reqs=max_num_reqs, | ||
| enable_memory_saver=get_exec().features.enable_memory_saver, | ||
| enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer, | ||
| enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy, |
There was a problem hiding this comment.
[medium] Lazy ping-pong on the unified pool reaches an unfiltered free.
Propagating the flag is right. But with it, the unified pool holds lazy ping-pong buffers with a -1 second entry for the first time. free_mamba_cache filters != -1 under lazy mode, but StreamingSession._free_slot_mamba (session/streaming_session.py:499) frees slot.kv.mamba_ping_pong_track_buffer whole.
On the unified allocator that free indexes virtual_to_physical[-1]. Under eager compaction it trips _raise_stale_slot_assertion when a session slot closes; under lazy compaction it frees a bogus page. session_held_mamba_slots also counts numel() (2) for one allocated slot.
The same gap exists on the static pool in lazy mode, but this change is what makes it reachable on the unified pool. Filtering -1 in the session free and count (or sharing a helper with free_mamba_cache) should land together with this change.
There was a problem hiding this comment.
Addressed in seperated PR in #42072 : session frees and held-slot counts now exclude -1.
| cache_params, | ||
| mamba_layer_ids: List[int], | ||
| enable_mamba_extra_buffer: bool, | ||
| enable_mamba_extra_buffer_lazy: bool = False, |
There was a problem hiding this comment.
nit: the bug this fixes is a caller silently dropping this flag, and a False default keeps that failure mode open for the next factory or fixture. The parameters are keyword-only and enable_mamba_extra_buffer just above has no default. Consider dropping the default here and in init_unified_mamba_pools / init_unified_mamba_swa_pools, so a missed flag becomes a TypeError.
| req_pool_indices=buffers.req_pool_indices[:bs], | ||
| seq_lens=buffers.seq_lens[:bs], |
There was a problem hiding this comment.
Question: this changes the view for every draft-extend backend, not just Inkling. req_pool_indices / seq_lens go from the full max_bs buffers to [:bs], while seq_lens_cpu a few lines below stays full length (the multi-layer replay view does the same).
flashattention_backend slices [:bs] itself, so I don't see a crash there. Still, it's an unrelated behavior change inside a checkpoint fix, and it leaves the GPU and CPU seq-len views with different lengths. Was it needed so req_pool_indices and mamba_track_indices have the same row count in _prepare_slot_indices? If so, please slice seq_lens_cpu the same way and add a comment saying why. Otherwise, drop the change.
| if buffers.mamba_track_indices is not None: | ||
| buffers.mamba_track_indices[:bs].zero_() | ||
| if forward_batch.mamba_track_indices is not None: | ||
| buffers.mamba_track_indices[:raw_bs].copy_( | ||
| forward_batch.mamba_track_indices | ||
| ) |
There was a problem hiding this comment.
perf: this adds two launches to every draft-extend replay (zero_ over [:bs], then copy_ over [:raw_bs]). They run outside the grouped foreach copy just below, which exists to cut launch count, and every hybrid model with enable_mamba_extra_buffer pays for them, not only Inkling.
Appending (buffers.mamba_track_indices[:raw_bs], forward_batch.mamba_track_indices) to copy_dsts / copy_srcs, and zeroing only [raw_bs:bs] when there is padding, keeps this within the existing launch. The multi-layer prepare() has the same pattern.
There was a problem hiding this comment.
Ordinary replay now uses the existing grouped copy and zeros only padding.
The multi-layer runner uses fused input preparation, so it retains one separate checkpoint copy while avoiding redundant zeroing.
| if forward_batch.forward_mode.is_draft_extend_v2(): | ||
| # Draft extend snapshots the accepted window in its own cache update. | ||
| # Inert prefill tracking would overwrite its checkpoint destinations. | ||
| self.sconv_metadata.track_conv_indices = None | ||
| return |
There was a problem hiding this comment.
nit: this early return deliberately skips the inert track scatter for draft extend. That contradicts the __init__ comment at L152-155, which says a capture batch without tracking metadata must still launch it. Please scope that comment to exclude draft extend, so nobody later "fixes" this return and brings back the overwrite you're avoiding.
| mamba_track_indices=( | ||
| None | ||
| if buffers.mamba_track_indices is None | ||
| else buffers.mamba_track_indices[:bs] | ||
| ), |
There was a problem hiding this comment.
nit: None if buffers.mamba_track_indices is None else buffers.mamba_track_indices[:bs] now appears six times across the two runners, and the zero/copy refresh block twice. _prepare_track_indices reads the field with a bare attribute access, so a view that omits it is an AttributeError by design. A small helper on the buffers (e.g. track_indices(bs) and refresh_track_indices(src, raw_bs, bs)) keeps the next view from forgetting it.
9c90ccf to
0c15e77
Compare
Rebased and pushed the checkpoint/spec-path changes here |
ch-wan
added
run-ci
0c15e77 to
fcce553
Compare
|
/rerun-group unified-memory |
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 |
mechanical_provable
…e and replay Handle optional replay tracking fields, preserve virtual source buffers, and keep draft accepted-window snapshots separate from prefill tracking. Cover both EAGLE runners with real compaction and graph replay, including disabled tracking and padding. Register GPU regressions in the kernel suite. non_mechanical_provable
…zation Keep kernel-facing checkpoint slots in backend metadata so repeated eager planning and distinct draft backends always translate the original virtual IDs. Route fused and unfused consumers through that metadata and cover repeated planning, staging, and inert prefill capture. non_mechanical_provable
Require mamba_track_indices on metadata views, with None representing disabled tracking. Wire the multi-layer staging view to the shared virtual input buffer and cover missing fields, explicit None, and staging checkpoint writes. non_mechanical_provable
prepare_for_draft_extend now reads draft_model_runner.model_config.model_is_mrope when it lays out the draft-extend lengths on the GPU.
fcce553 to
b5253d8
Compare
ch-wan
left a comment
ch-wan
left a comment
There was a problem hiding this comment.
Thanks for the rebase and the split. Everything from the last round is addressed: #41144's in-place rebind is gone and ForwardBatch.mamba_track_indices stays virtual across backends, the verify commit translates its destinations, both draft-extend runners carry the checkpoint buffer through capture, staging and replay, the [:bs] slicing is reverted, and the lazy-policy and session fixes moved to #42072. LGTM with two small items.
- [minor] The new draft-extend checkpoint write is the one device-visible change on the static pool, and none of the posted runs exercise it. Details inline. Related description fix: on main the captured draft-extend graph doesn't omit the draft conv checkpoint write. It writes it into reserved slot 0.
- nit: the history still contains
fix(mem-cache): propagate lazy Mamba checkpoint policy to unified pools, whichreconcile checkpoint metadata with upstream and split lazy policyreverts, so the net diff carries none of it. Please drop the add/revert pair. Otherwise the default squash message advertises a change that now lives in #42072.
perf: the multi-layer prepare() still does one separate checkpoint copy because its other inputs go through the fused prepare kernel. That's fine as is.
Inline comments: 1 minor.
| if forward_batch.forward_mode.is_draft_extend_v2(): | ||
| # Draft extend snapshots the accepted window in its own cache update. | ||
| # Inert prefill tracking would overwrite its checkpoint destinations. | ||
| self.sconv_metadata.track_conv_indices = None | ||
| return |
There was a problem hiding this comment.
[minor] No end-to-end run exercises the new draft-extend checkpoint write.
On the static pool, this early return plus the capture batch now carrying mamba_track_indices in both draft-extend runners is the only change with a visible effect on the device. On main, the draft-extend capture batch has no track indices, so the inert branch below rebinds forward_batch.mamba_track_indices to the zero buffer and _update_sconv_cache_for_draft_extend bakes those zeros into the graph. Every graph replay then writes the draft's boundary conv window into reserved slot 0, while eager draft extend writes the real ping-pong slot. With this PR the graph writes the real slot too.
Where it shows up: static pool + Inkling MTP (EAGLE or multi-layer EAGLE) + extra_buffer + draft-extend CUDA graph. A boundary crossed during decode is cached, and a later request that hits that prefix restores the draft conv state from the cached slot (multi-layer EAGLE copies the draft pools in _apply_deferred_mamba_init_to_draft_pools). On main that request restores a stale window, so accept length on cache hits should go up with this PR. The posted runs can't show it: the KL checks only read target logprobs, and accept length is 2.736 on both sides.
Could you add a shared-prefix accept-length comparison, main vs this PR? For example, a multi-turn run where turn 2 reuses turn 1's generated tokens past a mamba_track_interval boundary. If that's too costly, a line in the description naming test_draft_runners_capture_and_replay_checkpoint_destinations as the coverage for this path works too.
There was a problem hiding this comment.
Thanks for pointing out the coverage gap. I added a shared-prefix comparison between the merge parent (1e49077) and the merged fix (1d02a36), exercising draft-extend graph replay and restoration into both draft pools.
- Across 12 varied inputs, aggregate draft acceptance changed from 72.67% to 73.06%. One proof case consistently required 29 rather than 30 verify calls for the same 64 output tokens, reproduced in five additional flush/reseed repetitions.
- Results were not uniformly positive: a separate 8-token comparison had unchanged verify counts and one fewer accepted draft overall. I added the setup and limitations to the description and explicitly named test_draft_runners_capture_and_replay_checkpoint_destinations as the focused regression coverage.
…current state A per-depth conv-chain MTP head (Inkling) owns one recurrent-state block per depth. `StateHostKind` registers as the host for `LAYER_STATE`: the mamba sub-pool's entries take the draft's state block after the host's streams, laid out stream-major like the host's own block, so the slot that is the request's carries both and the host's whole-entry clear and copy restore the draft's state with the target's on every radix checkpoint. A host without a mamba sub-pool declines the draft, which its factory then refuses for an EAGLE draft as before. `MambaSubPoolSpec` gains the fused entry math (`host_entry_bytes`, `draft_offset_in_entry`, an aligned `entry_bytes` when a state region is set; byte-identical unfused), `build_mamba_entry_views` takes an entry stride and a block offset, `UnifiedKVPool.build_draft_state_views` views the draft block, and `UnifiedMambaPool.clear_slots`, `copy_from` and `move_kv_cache` work on whole entries: compaction relocates live and cached state slots, and a per-tensor move would leave the draft block behind. The draft worker binds `UnifiedDraftMambaPool` through the "state" binder: a request table cloned from the target's (same request->slot mappings, same allocator) whose `mamba_pool` is a view of the draft block keyed by the runner's own state layer ids. `UnifiedDraftMambaPool.clear_slots` / `copy_from` act on the draft's own block: a fused draft shares the host's slot, so the host's whole-entry ops never run for it and nothing else would reset the draft's conv and temporal streams. The multi-layer EAGLE worker drives its state pools with the request pool's translated indices, so the block is cleared and copied on decode and draft-extend, not only on plain extend. The boot solve prices what the factory carves out: with a fused state region the per-slot state bytes are the fused entry (host + draft + pad), not the host's raw entry, so `max_mamba_cache_size` is solved against the buffer that is actually allocated. `fused_entry_bytes` gains the "mamba" builder that answers that price; it prices the state layers in this runner's span, as the mamba pool factory slices them (`mamba2_cache_params.layers` is the whole model's). The draft worker's check of the built model now refuses recurrent layers only when the placement gives the runner no state lanes, and its KV dtype check reads the dense regions alone. Tests: state placement rows (per-depth and replicated heads, decline without a state host), the mixed boot-log line, the mamba builder pricing this runner's state layers, the fused state entry math, byte-disjoint host and draft blocks within a slot, the host's whole-entry clear, copy and move carrying the draft block, the draft pool's own slot ops clearing its block and sparing the host's, and the fused request-pool clone on the tri-pool factory. This makes Inkling's MTP draft placeable under the unified pool. The two track-id paths that reaches, draft-extend replay views and the verify commit, hand the conv kernels physical slots since sgl-project#38229. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2 participants
Motivation
On the unified memory pool the scheduler hands out virtual Mamba slot ids, while the Inkling short-conv kernels write checkpoints to physical slots. #41144 fixed the main path by rewriting
forward_batch.mamba_track_indicesin place to a physical buffer. Two gaps remain:commit_conv_state_after_mtp_verifypasses the scheduler's virtualmamba_track_indicesto the verify scatter untranslated. DSPARK on the unified pool reaches it.The in-place rewrite also assumes one backend prepares each batch. When several do, as multi-layer EAGLE does with one draft backend per step, every backend after the first translates the already-physical ids again (
v2p[v2p[v]]). Main currently rejects--enable-unified-memorywith speculative algorithms other than DSPARK, so serving does not reach this today; keeping the batch virtual removes the hazard instead of relying on that restriction.Modifications
ForwardBatch.mamba_track_indicesstays virtual.InklingShortConvAttnBackendtranslates it into a fixed-address int64 buffer it owns and stores the physical ids inInklingShortConvMetadata.track_cache_indices. All Inkling checkpoint consumers (fused and unfused, decode, extend and draft extend) read that field. The static pool keeps the identity path. [Unified Memory] Fix Inkling conv-checkpoint track ids written to virtual slot numbers #41144's in-place rebind and its buffer are removed.commit_conv_state_after_mtp_verifytranslates the checkpoint destinations, as the GDN backend already does.enable_mamba_extra_bufferis set, pass it at capture, staging and replay, refresh live ids on replay and zero only the padded rows. The ordinary runner folds the refresh into its grouped copy. Helpers live inspeculative/draft_checkpoint.py.Accuracy Tests
Validated on an NVIDIA B200 against main
048b2a07f3, with this PR rebased onto it; the end-to-end runs also include #39982 and #42072.test_inkling_checkpoint_indices.py: 17 tests and 29 subtests pass. The 33 related existing test files give the same results as main.test_draft_runners_capture_and_replay_checkpoint_destinations, which checks checkpoint destinations during CUDA graph capture and replay.thinkingmachines/Inkling, revisiontest),--enable-deterministic-inference, 32 LongBench prompts: logprob match, prefill cache hit and decode cache hit are atavg_kl_div=0.0on the static, lazy and unified pools, and multi-layer EAGLE MTP (two heads) on the static pool is at 0.0 as well.test_inkling.pypasses 5/5 andtest_inkling_unified.py9/9.avg_kl_div=0.116with no cache hit, the same on main and with this PR), so the end-to-end KL check cannot isolate the verify-commit translation. The unit test above covers it.Speed Tests and Profiling
Speculative accept length over the same workloads, read from the decode log (main / this PR):
Per step, the unified pool adds one translate and an int64 copy of at most
bsentries; the ordinary draft-extend runner folds its refresh into the existing grouped copy, and the multi-layer runner adds one copy.Follow-up: shared-prefix checkpoint restoration
Compared the merge parent (
1e490772e5) with the merged fix(
1d02a36bb7) on 12 varied inputs. Eight baseline cases had draftacceptance below 100%.
The proof-case change reproduced in five additional cache-flush/reseed
repetitions per revision. Its cold control required 29 verify calls on
both revisions. All paired output token IDs matched between revisions.
Setup and limitations
RTX 5090 / WSL, reduced 8-layer Inkling test checkpoint, FA4, static
pools, extra_buffer, page/track interval 128, and two-head multi-layer
EAGLE with draft-extend CUDA graphs. Both revisions used the same
SM120 MoE compatibility override and diagnostic instrumentation.
Each seed generated 640 tokens. Follow-up requests reused a generated
768- or 896-token checkpoint with exactly one uncached input token.
Logs confirmed draft graph replay and checkpoint restore calls for
both draft pools.
The five proof-case repetitions were selected after the initial
comparison and were not independent server replicates.
Results were not uniformly positive: a separate 8-token continuation
experiment used 49 verify calls on both revisions, with accepted drafts
changing from 48/98 to 47/98. Two inputs had warm/cold output differences
identically on both revisions; their causes remain unestablished.
These results demonstrate a small, workload-specific acceptance
improvement, not a general throughput gain.
Checklist
CI States
Latest PR Test (Base): ⏳ Run #37193176706
Latest PR Test (Extra): ⏳ Run #37193176482
Latest PR Test (AMD ROCm 10): ⏳ Run #37193176705