refactor: extract packed capture_rows from the offline SGLang backend - #780
Merged
jiapingW merged 1 commit intoSep 2, 2026
Merged
Conversation
5 tasks
This was referenced Sep 1, 2026
maocheng23
force-pushed
the
maocheng/colocate-1-capture-rows
branch
from
September 1, 2026 23:30
7998faa to
c4778db
Compare
maocheng23
changed the base branch from
main
to
maocheng/colocate-3-rank0-tracker
September 1, 2026 23:31
Expose capture_rows(input_ids) on OfflineSGLangCaptureBackend and OfflineSGLangCapture so callers can capture variable-length rows in one packed prefill without building padded tensors. capture_eagle3 delegates to it and keeps its exact request construction and outputs. DSpark capture-layer setup now falls back from the native set_dspark_layers_to_capture hook to the dense DFlash hook (with a log line naming the resolved hook) so dense targets such as Qwen3 can serve DSpark capture on stock SGLang 0.5.14 models. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
maocheng23
force-pushed
the
maocheng/colocate-3-rank0-tracker
branch
from
September 2, 2026 00:24
fb2bd50 to
9edfa1b
Compare
maocheng23
force-pushed
the
maocheng/colocate-1-capture-rows
branch
from
September 2, 2026 00:24
c4778db to
68cf3fc
Compare
maocheng23
marked this pull request as ready for review
September 2, 2026 00:25
jiapingW
approved these changes
Sep 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Stack 3/5 replacing draft #766. This PR is the capture-engine primitive the colocated runtime builds on: a packed variable-length capture API on the offline SGLang backend, plus DSpark capture-hook compatibility for dense targets.
Modifications
capture_rows(input_ids: list[list[int]]) -> (aux_rows, last_rows)fromcapture_eagle3onOfflineSGLangCaptureBackendand expose it onOfflineSGLangCapture. One packed prefill captures variable-length rows without building padded tensors.capture_eagle3delegates to it and keeps its exact request construction (set_extend_range, sglang 0.5.18), outputs, and pool clearing, so offline feature preparation stays byte-identical (unlike feat: add scalable colocated online training #766, no attention-mask trimming is bundled in).set_dspark_layers_to_capturehook to the dense DFlash hook and logs the resolved hook. In stock SGLang onlydeepseek_v4exposes the DSpark hook; dense targets such as Qwen3 serve the same auxiliary-hidden-state layout throughset_dflash_layers_to_captureand previously failed outright. The hook table is one module constant (_CAPTURE_LAYER_HOOKS).Related Issues
Splits #766. Stack: #781 teardown-abort ← #782 rank0-tracker ← #3 (this) ← #783 colocated-core ← #784 hybrid-shard.
Accuracy Test
tests/test_runtime/test_sglang_capture_hooks.py: hook preference, DFlash fallback, missing hooks, unknown method. Runs where sglang is installed (CI), like the existing sglang compat tests.test_sglang_0518_compat.pynow inspectscapture_rowsfor the request-construction contract.Checklist
black --checkandisort --check-only).