Skip to content

refactor: extract packed capture_rows from the offline SGLang backend - #780

Merged
jiapingW merged 1 commit into
maocheng/colocate-3-rank0-trackerfrom
maocheng/colocate-1-capture-rows
Sep 2, 2026
Merged

refactor: extract packed capture_rows from the offline SGLang backend#780
jiapingW merged 1 commit into
maocheng/colocate-3-rank0-trackerfrom
maocheng/colocate-1-capture-rows

Conversation

@maocheng23

@maocheng23 maocheng23 commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

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

  • Extract capture_rows(input_ids: list[list[int]]) -> (aux_rows, last_rows) from capture_eagle3 on OfflineSGLangCaptureBackend and expose it on OfflineSGLangCapture. One packed prefill captures variable-length rows without building padded tensors. capture_eagle3 delegates 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).
  • DSpark capture-layer setup falls back from the native set_dspark_layers_to_capture hook to the dense DFlash hook and logs the resolved hook. In stock SGLang only deepseek_v4 exposes the DSpark hook; dense targets such as Qwen3 serve the same auxiliary-hidden-state layout through set_dflash_layers_to_capture and 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

  • New 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.py now inspects capture_rows for the request-construction contract.
  • Behavior-preserving refactor for the existing offline path; exercised end to end by the colocated run reported in feat: colocated online training with in-process SGLang capture #783 (to be re-run on 0.5.18 after this rebase, see there).

Checklist

  • Format your code according to the Code Formatting with Pre-Commit (black --check and isort --check-only).
  • Add unit tests.
  • Update documentation as needed (none needed — internal API).

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
maocheng23 force-pushed the maocheng/colocate-3-rank0-tracker branch from fb2bd50 to 9edfa1b Compare September 2, 2026 00:24
@maocheng23
maocheng23 force-pushed the maocheng/colocate-1-capture-rows branch from c4778db to 68cf3fc Compare September 2, 2026 00:24
@maocheng23
maocheng23 marked this pull request as ready for review September 2, 2026 00:25
@jiapingW
jiapingW merged commit 72a2b9e into maocheng/colocate-3-rank0-tracker Sep 2, 2026
1 check passed
@jiapingW
jiapingW deleted the maocheng/colocate-1-capture-rows branch September 2, 2026 04:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants