Skip to content

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

Draft
maocheng23 wants to merge 1 commit into
mainfrom
maocheng/colocate-1-capture-rows
Draft

refactor: extract packed capture_rows from the offline SGLang backend#780
maocheng23 wants to merge 1 commit into
mainfrom
maocheng/colocate-1-capture-rows

Conversation

@maocheng23

Copy link
Copy Markdown
Collaborator

Motivation

Stack 1/5 replacing draft #766 (colocated online training), split into reviewable steps. This PR is the capture-engine primitive the colocated runtime builds on: a packed variable-length capture API, plus DSpark capture-hook compatibility for dense targets.

Modifications

  • Extract capture_rows(input_ids: list[list[int]]) 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, outputs, and pool-clearing behavior — offline feature preparation is 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, logging the resolved hook. Dense targets such as Qwen3 expose the same auxiliary-hidden-state layout only through set_dflash_layers_to_capture on stock SGLang 0.5.14, and previously failed outright. (feat: add scalable colocated online training #766's additional EAGLE3-hook fallback tier is dropped: no supported build needs it and its layout equivalence was unverified.)

Related Issues

Splits #766. Stack: #1 (this)#2 teardown-abort ← #3 rank0-tracker ← #4 colocated-core ← #5 hybrid-shard.

Accuracy Test

  • tests/test_runtime/test_sglang_0514_compat.py gains hook-preference, fallback, and missing-hook coverage (runs in CI with sglang installed).
  • Behavior-preserving refactor for the existing offline path; exercised end to end by the 90-step colocated run reported in fix bug #4 of this stack.

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>
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.

1 participant