fix(sglang): port the DataFlow capture backend to sglang 0.5.14 (up-16)

[maocheng23]

Summary

Makes up-16's DataFlow target-capture backend work on sglang 0.5.14, and bumps the pin from 0.5.9 → 0.5.14 (matching main).

On main, the 0.5.14 move landed as #605 (eagle3 / custom backends / shared infra) plus the follow-up branch patch/fix_sgl0.5.14_dflash (the DFlash-specific _extend fixes #605 didn't include). Both were written against the old layout (specforge/modeling/target/dflash_target_model.py). up-16 consolidated eagle3 + dflash extend into a single SGLangCaptureBackend (specforge/inference/target_engine/sglang_backend/capture.py), so the same sglang-API adaptations are applied there once and cover both draft algorithms (and DSpark, which builds on the DFlash path).

Changes (sglang_backend/capture.py)

• Import move: prepare_mlp_sync_batch_raw now comes from sglang.srt.managers.scheduler_components.dp_attn (moved from …scheduler_dp_attn_mixin in 0.5.14). This is the import that was hard-crashing the whole backend on 0.5.14.
• build(): 0.5.14 split the post-load setup out of ModelRunner.initialize() (which now only loads weights). Since we drive the ModelRunner directly (no scheduler/TpModelWorker), replicate it: alloc_memory_pool() / init_attention_backends() / init_cuda_graphs(). Also set chunked_prefill_size=-1 — we prefill the whole batch in one _forward_extend, and 0.5.14's eager runner otherwise pre-allocates buffers sized to the default chunked_prefill_size (8192) and errors when our token count exceeds it.
• _forward_extend(): get_model_worker_batch() was removed — ForwardBatch.init_new now consumes the ScheduleBatch directly and reads capture_hidden_mode off it. Also materialize prefill_input_ids_cpu → device (the scheduler normally does this H2D copy in resolve_forward_inputs; we bypass the scheduler).
• Req construction (dflash + eagle3 + eagle3-vlm): Req.fill_ids was removed in favor of full_untruncated_fill_ids (an array) + integer fill_len.
• input_lens captured before the forward — origin_input_ids is released during prepare_for_extend/forward on 0.5.13+.
• pyproject.toml: sglang==0.5.9 → sglang==0.5.14.

sglang_backend/{utils,patch,model_runner}.py needed no changes.

Testing

• Against a real sglang 0.5.14 env (venv on sci-h200, 8×H200): the backend imports and the DFlash sglang capture path runs — full GPU smoke in progress; results to follow in a comment.
• Reference parity: the diffs mirror patch/fix_sgl0.5.14_dflash (DFlash) and upgraded sglang to v0.5.14 #605 (shared/eagle3), re-expressed for the consolidated backend.

Notes

• Scoped to the sglang-version coupling in the capture backend; no algorithm/behavior changes.
• Base is dataflow-up-16-zerocopy so DSpark (feat: DSpark trainer (DFlash + Markov/confidence heads + L1 distillation) #613) can rebase on top and exercise the DFlash sglang last_hidden_states path on 0.5.14.

🤖 Generated with Claude Code

[maocheng23]

✅ Validated on real sglang 0.5.14 (venv on sci-h200, 8×H200, torch 2.11): the DFlash sglang capture path builds the target and runs a real forward — generate_dflash_data on Qwen2.5-0.5B returns hidden_states=(2, 16, 2688) (3 capture layers × 896), finite. Backend imports cleanly and _forward_extend completes end-to-end.

Note: needed two extra patch.py fixes beyond the initial commit (parallel-state init) — init_model_parallel_group dropped pynccl_use_current_stream, and compute_dp_attention_world_info now returns a 4-tuple. Both mirror main's post-#605 patch.py. Now folded in.

[gemini-code-assist]

Code Review

This pull request updates the sglang dependency from version 0.5.9 to 0.5.14 and adapts the sglang_backend/capture.py module to accommodate several breaking changes introduced in sglang versions 0.5.13 and 0.5.14. These adaptations include updating import paths, disabling chunked prefill, replicating post-load setup steps for the ModelRunner, handling host-to-device input ID copies, adjusting ForwardBatch initialization, capturing input lengths prior to forwarding, and replacing deprecated request attributes like fill_ids with full_untruncated_fill_ids and fill_len. There are no review comments, so I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[maocheng23]

🔁 Re-validated on sglang 0.5.14 after the eagle3 fix (f59f014) + comment trim:

• eagle3 sglang capture now runs end-to-end (wrap_eagle3_logits=True): hidden_states=(2,16,2688), target=(2,16,151936) → EAGLE3_SGLANG_OK. Confirms the replaced_logits_processor_forward_for_eagle3 / multi_item_delimiter_indices fix.
• dflash sglang capture still OK ((2,16,2688)) — no regression from the update.

Both eagle3 and dflash capture paths now work on 0.5.14.