Repository navigation
fix(inkling): translate unified-memory checkpoint destinations to physical slots #38229
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -85,6 +85,8 @@ class InklingShortConvMetadata(msgspec.Struct): | |
| # [B, conv_kernel - 1] input positions whose conv window feeds the prefix | ||
| # cache. Extend only, and only when tracking is on. | ||
| track_conv_indices: Optional[torch.Tensor] = None | ||
| # Kernel-facing checkpoint slots. ForwardBatch retains the virtual source. | ||
| track_cache_indices: Optional[torch.Tensor] = None | ||
|
|
||
|
|
||
| class InklingShortConvAttnBackend(ShortConvAttnBackend): | ||
|
|
@@ -144,7 +146,10 @@ def _alloc_graph_buffers(self): | |
| self._graph_track_conv_indices = torch.zeros( | ||
| (max_bs, self.conv_state_len), dtype=torch.int64, device=dev | ||
| ) | ||
| # Prefill can capture before init_cuda_graph_state. Keep the translated | ||
| # checkpoint destinations at one address across both graph phases. | ||
| self._graph_track_indices = torch.empty(max_bs, dtype=torch.int64, device=dev) | ||
| # Inert prefill track fields (excluding draft extend): a warmup batch that | ||
| # carries no tracking metadata must still LAUNCH the track scatter | ||
| # (all rows masked off), or the python-level `if` specializes the | ||
| # scatter out of the captured graph. | ||
|
|
@@ -155,8 +160,6 @@ def _alloc_graph_buffers(self): | |
| self._graph_track_inert_seqlens = torch.zeros( | ||
| max_bs, dtype=torch.int64, device=dev | ||
| ) | ||
| # Same address-stability requirement; the base only sizes this from | ||
| # init_cuda_graph_state, which the prefill graph never calls. | ||
| self._alloc_cache_indices_buf(max_bs) | ||
|
|
@@ -250,6 +253,7 @@ def init_forward_metadata_capture_cpu_graph(self, *args, **kwargs): | |
|
|
||
| def _prepare_slot_indices(self, forward_batch: ForwardBatch): | ||
| self._reset_step_state() | ||
| self._prepare_track_indices(forward_batch) | ||
| req_pool_indices = forward_batch.req_pool_indices | ||
| n = req_pool_indices.shape[0] | ||
| buf = self._cache_indices_buf | ||
|
|
@@ -271,22 +275,22 @@ def _prepare_slot_indices(self, forward_batch: ForwardBatch): | |
| return | ||
| self.forward_metadata = self._forward_metadata(forward_batch) | ||
| self._refresh_cache_indices() | ||
|
|
||
| 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 | ||
|
Comment on lines
+279
to
+293
Collaborator
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
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
Contributor
Author
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. |
||
|
|
||
| def _refresh_sconv_metadata( | ||
| self, forward_batch: ForwardBatch, *, on_graph_path: bool | ||
|
|
@@ -451,6 +455,11 @@ def _refresh_track_conv_indices( | |
| ``batch_size`` rows while the track lengths cover only live requests, and | ||
| every row it may read must index inside *this* replay's token buffer. | ||
| """ | ||
| 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 | ||
|
Comment on lines
+458
to
+462
Collaborator
There was a problem hiding this comment. nit: this early return deliberately skips the inert track scatter for draft extend. That contradicts the
Comment on lines
+458
to
+462
Collaborator
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 Where it shows up: static pool + Inkling MTP (EAGLE or multi-layer EAGLE) + 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
Contributor
Author
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.
|
||
| if forward_batch.mamba_track_mask is None: | ||
| if not on_graph_path: | ||
| self.sconv_metadata.track_conv_indices = None | ||
|
|
@@ -460,6 +469,9 @@ def _refresh_track_conv_indices( | |
| rows = forward_batch.batch_size | ||
| forward_batch.mamba_track_mask = self._graph_track_inert_mask[:rows] | ||
| forward_batch.mamba_track_indices = self._graph_track_inert_indices[:rows] | ||
| self.sconv_metadata.track_cache_indices = self._graph_track_inert_indices[ | ||
| :rows | ||
| ] | ||
| forward_batch.mamba_track_seqlens = self._graph_track_inert_seqlens[:rows] | ||
| rows = forward_batch.batch_size | ||
| query_start_loc = self.sconv_metadata.query_start_loc | ||
|
|
@@ -532,6 +544,8 @@ def commit_conv_state_after_mtp_verify( | |
| slot_ids = self._translate_mamba_indices( | ||
| pool.get_mamba_indices(req_pool_indices) | ||
| ) | ||
| if mamba_track_indices is not None: | ||
| mamba_track_indices = self._translate_mamba_indices(mamba_track_indices) | ||
| scatter_mamba_states_after_mtp_verify( | ||
| pool.get_speculative_mamba2_params_all_layers(), | ||
| slot_ids, | ||
|
|
@@ -603,6 +617,7 @@ def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): | |
| and self.supports_draft_extend_metadata_staging | ||
| ): | ||
| self.short_conv_backend._reset_step_state() | ||
| self.short_conv_backend._prepare_track_indices(forward_batch) | ||
| self.short_conv_backend._refresh_sconv_metadata( | ||
| forward_batch, on_graph_path=True | ||
| ) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| from typing import Optional | ||
|
|
||
| import torch | ||
|
|
||
|
|
||
| def track_indices(buffer: Optional[torch.Tensor], bs: int): | ||
| return None if buffer is None else buffer[:bs] | ||
|
|
||
|
|
||
| def refresh_track_indices(buffer, src, *, raw_bs, bs, copy_dsts=None, copy_srcs=None): | ||
| """Refresh virtual destinations without clearing rows that will be copied. | ||
|
|
||
| Ordinary EAGLE joins its existing grouped copy. The multi-layer runner | ||
| uses a fused preparation kernel for other inputs and copies here instead. | ||
| """ | ||
| if buffer is None: | ||
| return | ||
| if src is None: | ||
| buffer[:bs].zero_() | ||
| return | ||
| if raw_bs < bs: | ||
| buffer[raw_bs:bs].zero_() | ||
| if copy_dsts is None: | ||
| buffer[:raw_bs].copy_(src) | ||
| else: | ||
| copy_dsts.append(buffer[:raw_bs]) | ||
| copy_srcs.append(src) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[bug] Unified
_prepare_slot_indicesnow always calls_prepare_track_indices, which does a bareforward_batch.mamba_track_indicesload. Decode/prefill replay views include that field (None when tracking is off). EAGLE draft-extend replay does not:EAGLEDraftExtendCudaGraphRunner.replayandMultiLayerEagleDraftExtendCudaGraphRunnerpass aSimpleNamespacewithreq_pool_indices/seq_lens/spec_infobut nomamba_track_indices(python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py:609andpython/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py:483). Inkling's hybrid wrapper still forwardsinit_forward_metadata_out_graphto the sidecar onDRAFT_EXTEND_V2(the draft runs its own convs). The new load therefore raisesAttributeErroron 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 waydecode_cuda_graph_runner._replay_forward_batchdoes. If draft-extend graphs are supposed to refresh checkpoint destinations, the capture batch needs that slot as well; today capture builds aForwardBatchwith the defaultNone, so the track-save is specialized out of the graph independently of this PR.Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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