[mickqian]

Motivation

Follow-up to #46659 and #46660. The graph-pool allocation and replay copy logic is repeated across attention wrappers and model-specific breaks. These should share output management without forcing different attention implementations through the same compute path.

Modifications

• Add an optional output_factory to eager_on_graph: allocate outputs before ending the captured segment, retain weak references, and validate layouts before replay writeback.
• Allow direct writes through pass_output=True, preserving H3 embedding's buffer reuse without copying those tensors into themselves.
• Migrate common diffusion attention, H3/VDN attention, H3 embedding, and Qwen Image 2.1 prefix attention. Ordinary eager execution bypasses the factory and extra copies. Existing calls without a factory retain their current behavior.
• Avoid querying the Ulysses group when USP output prediction does not need SP. Update the BCG documentation.

Production code: 115 lines added, 97 removed across five files. Attention algorithms and backend selection are unchanged.

Accuracy Tests

Validated on one H200 with PyTorch 2.14.1+cu130 / CUDA 13.0 at 66de1dcb7a5e7ec9781f66dd751cffbd085182f5:

• SRT BCG basic/copy/break-helper, diffusion attention outputs, and LTX tests: 35 passed, 33 subtests passed.
• Qwen Image 2.1 CPU/CUDA, H3 DiT contract, and VDN suites: 181 passed. Includes exact padded-prefix replay and direct embedding-buffer checks.
• Small real H3 attention modules with random weights: dense SDPA, VDN full-cover, and VDN partial-window branches each captured two segments and matched eager for five changed-input replays at rtol=0, atol=0.
• The standalone LTX USP BCG test fails on the base commit without an initialized SP group and passes with this fix.
• Pre-commit and git diff --check passed.

Full-checkpoint generation, multi-GPU TP/SP, and AMD/XPU have not been validated in this follow-up.

Speed Tests and Profiling

Bounded Qwen prefix-attention microbenchmark, not full-model E2E: one H200, BF16, batch 1, 16 heads x 64, SDPA, 128 padded prefix rows, K/V length Q+384. Measures attention plus one output addition. 20 warmups, then median of 9 batches x 200 calls, synchronized on CUDA. Compared against c43abd6d9e38a14e81bdb77c648a06353e1c10ee; execution order was new, base, new repeat.

Q rows	Mode	Base (us)	New, two runs (us)
256	eager	48.90	38.16 / 39.10
256	BCG	82.40	80.54 / 81.48
2048	eager	75.45	71.83 / 72.00
2048	BCG	97.35	96.27 / 96.93

Output hashes match across commits. Incremental peak PyTorch allocation is unchanged for BCG and lower for eager in both shapes. BCG performance is effectively neutral; eager avoids an unnecessary allocation/copy. No end-to-end speedup is claimed.

Checklist

• Run pre-commit.
• Add focused regression tests.
• Update documentation.
• Report correctness and performance evidence with validation limits.
• Follow the existing code style.

---

CI States

Latest PR Test (Base): ✅ Run #40632784359
Latest PR Test (Extra): ❌ Run #40632784280
Latest PR Test (AMD ROCm 10): ⏳ Run #40632784186

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 28, 2026, 4:36 AM

💡 Tip: Enable Automations to automatically generate PRs for you.

[mickqian]

/tag-and-rerun-ci