[sglang] Experimental dense tensor-parallel recipe (NCCL all-reduce over VMM-IPC)

[rchalamala]

Summary

Dense tensor-parallel SGLang SAVE/LOAD recipe — routes the per-layer TP all-reduce through NCCL over Foundry's VMM-IPC bridge. GPU-validated on both Qwen3-8B (dense, TP=2, H200:2) and Qwen3.5-122B-A10B (hybrid Mamba/MoE, TP=4, H200:4), no IMEX/fabric.

The core fix for hybrid models (Qwen3.5): run the pre-capture warmup with the VMM region active on both SAVE and LOAD, so lazy-init persistent buffers (torch.compile codegen, Triton workspace, MoE/mamba state) land at deterministic in-region VMM addresses. The prior cursor-neutral approach (region stopped) put them out-of-region — fine for SAVE (same process), fatal for LOAD (different process → stale addresses → illegal memory access on decode graph replay).

Key changes in hooks.py:

# _run_dense_warmup: region stays ACTIVE (no stop/resume)
_run_warmup_pass(self)         # lazy-init → in-region
fi_backend.decode_cuda_graph_metadata = {}  # drop warmup's FlashInfer wrappers
torch.cuda.empty_cache()       # release transient blocks

# Called on BOTH SAVE and LOAD before prepass/capture/load:
if not _ep_lazy_init_needed() and _dense_save_warmup_enabled():
    _run_dense_warmup(self)

Also: _flashinfer_backend() sees through HybridLinearAttnBackend.full_attn_backend (Qwen3.5), and the SGLang hook is ported to the post-#26735 init_forward_metadata_out_graph 3-method ABC (legacy fallback kept).

LOAD warmup-skip via byte restore (FOUNDRY_SGLANG_WARMUP_RESTORE, default on)

The in-region warmup re-runs real forwards on LOAD (~20 s on 122B) only to reproduce the deterministic persistent buffers it wrote on SAVE. Instead, SAVE dumps the warmup allocation gap's bytes and LOAD restores them:

# SAVE, in _run_dense_warmup's finally, BEFORE empty_cache (gap still fully
# mapped — the caching allocator holds every block it touched):
dump_warmup_region(before_warmup, get_current_alloc_offset())
    # one contiguous cuMemcpyDtoH_v2 if fully mapped, else per-2MiB-page trial
    # copies keeping only cleanly-copied runs -> warmup_region.bin + .json
torch.cuda.empty_cache()   # NOW foundry unmaps freed blocks (holes appear)

# LOAD, after preallocate_for_load_mode() maps the region:
if _warmup_restore_enabled() and warmup_region_available():
    restore_warmup_region()   # cuMemcpyHtoD_v2 each interval back (clamped to
                              # final_alloc_offset), then set the cursor to the
                              # post-empty_cache offset
else:
    _run_dense_warmup(self)   # fallback: archive lacks the dump -> re-run warmup

This is on by default. LOAD only restores when the archive carries the dump; otherwise it falls back to re-running the warmup, so archives produced before this change still load. FOUNDRY_SGLANG_WARMUP_RESTORE=0 forces the re-run path on both.

Robustness:

• Restore is clamped to final_alloc_offset (the LOAD-mapped ceiling). Warmup transients that peak above the final captured offset are dropped (never referenced by the replayed graphs), so a larger-transient model can't drive an H2D copy into unmapped VA.
• Dump runs before empty_cache: afterwards foundry (cuMemFree→cuMemUnmap) releases freed blocks, so a whole-span copy would hit unmapped VA (CUDA_ERROR_INVALID_VALUE). The fragmented fallback (_dump_mapped per-2MiB trial copies over a holed range) is GPU-tested via a dedicated harness entrypoint that maps two VMM blocks around an unmapped hole and asserts the returned intervals/bytes + a healthy post-probe context.

Distinct from FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1, which only reserves the gap without restoring contents — GPU-proven incorrect (coherent but divergent LOAD tokens; the gap holds read-before-write state); that flag stays off by default.

Also fixes a latent race: the shared workspace_root/warmup_state.json was truncated in place by rank 0 while other TP ranks read it at end of capture (JSONDecodeError). Now written atomically (temp + os.replace).

Upgrade note: the dense/single-GPU/DP paths now run an in-region warmup (they previously did not), which shifts their VMM cursor trajectory. An archive captured before this change is not LOAD-compatible with a build after it — re-capture (re-run SAVE) after upgrading.

Validation

Model	TP	GPUs	SAVE→SAVE offset	SAVE==LOAD tokens
Qwen3-8B	2	H200:2	120517033984 (reproducible)	16/16 exact
Qwen3.5-122B-A10B	4	H200:4	120376524800 (all 4 ranks)	16/16 exact

Zero faults (MMU/Xid/CUDA error) on both. Byte-identical SAVE/LOAD VMM offset trajectory at every checkpoint. For 122B with warmup-restore on by default, LOAD cold-start ~94 s (all 4 ranks restore 144 MiB as one contiguous interval) vs ~127 s re-running warmup vs ~249 s cold SAVE, still 16/16 exact SAVE==LOAD and base==LOAD. The fragmented _dump_mapped fallback passes its dedicated single-GPU test.

122B-specific flags: --enforce-disable-flashinfer-allreduce-fusion, --max-mamba-cache-size 32, --disable-radix-cache.

Link to Devin session: https://modal.devinenterprise.com/sessions/625ef9aabbff4b318f291f968d259cfc
Requested by: @rchalamala

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR that start with 'DevinAI' or '@devin'.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[devin-ai-integration]

Devin Review found 2 new potential issues.

[devin-ai-integration]

🔍 build_capture_fb_view accesses buffers.seq_lens_cpu not used by the legacy path

build_capture_fb_view (python/foundry/integration/sglang/graph_ops.py:271) reads buffers.seq_lens_cpu[:bs], an attribute the legacy pre-pass never accessed (the legacy path derived CPU seq lens via seq_lens.cpu() inside reuse_pre_pass_init at hooks.py:565). If cuda_graph_runner.buffers in the targeted sglang build does not expose seq_lens_cpu as an attribute, this raises AttributeError during the pre-pass on both SAVE and LOAD. sglang is not vendored in this repo so I could not verify the attribute exists on the buffers namespace; the other accessed fields (seq_lens, req_pool_indices, positions, encoder_lens) are confirmed used by existing legacy code, but seq_lens_cpu is new. Worth confirming against the pinned sglang HEAD build.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

buffers.seq_lens_cpu does exist on post-#26735 sglang. In cuda_graph_runner.py, DecodeInputBuffers declares seq_lens_cpu and the runner allocates + fills it (self.seq_lens_cpu.fill_(seq_len_fill_value) / self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)), and sglang's own build_replay_fb_view reads buffers.seq_lens_cpu[:bs] — this port uses the same field the upstream replay view does.

Also now GPU-validated: I re-ran the dense TP recipe (Qwen3-8B, TP=2, 2×H200) against the sglang foundry HEAD (the new init_forward_metadata_out_graph API, which is what exercises build_capture_fb_view). All 4 phases healthy, errors=[], no AttributeError in the pre-pass. So this path is exercised, not just the pre-rename base.

[devin-ai-integration]

📝 Info: reuse_out_graph re-planning semantics differ from the legacy reuse_pre_pass_init

The new-API shim reuse_out_graph (python/foundry/integration/sglang/hooks.py:635-672) sets fi_backend.forward_metadata from the pre-allocated wrappers and then calls real_out_graph(forward_batch, in_capture=False) to re-run the planner without re-allocating, whereas the legacy reuse_pre_pass_init (hooks.py:477-574) explicitly calls indices_updater_decode.update(...) / indices_updater_prefill.update(...). The correctness of the new path hinges on the assumption (stated in comments) that init_forward_metadata_out_graph(..., in_capture=False) re-plans against the existing wrappers, does NOT allocate a second _int_workspace_buffer, and does NOT overwrite forward_metadata. This is an assumption about sglang internals that can't be verified from this repo; it is reportedly GPU-validated. If a future sglang changes in_capture=False to re-set forward_metadata or re-allocate, the VMM cursor determinism (or the manually-set forward_metadata) would silently break.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

Both assumptions check out against the real post-#26735 sglang, and the path is now GPU-validated.

• Return value ignored by the capture loop. init_forward_metadata_out_graph is a void method upstream (base_attn_backend.py), and the CUDA graph capture path calls it purely for side effects — attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) — without consuming a return. So reuse_out_graph returning None on the reuse-success path is fine; I only forward the real method's return on fallback for symmetry.
• in_capture=False does not allocate. In flashinfer_backend.py::init_forward_metadata_out_graph, the only allocator (_prepare_cuda_graph_metadata, which does the torch.empty for _int_workspace_buffer and creates the wrappers) is guarded by if in_capture:. With in_capture=False it just re-runs indices_updater_*.update(...) against the already-allocated *_cuda_graph_metadata[bs] wrappers. That's exactly why the reuse shim delegates to in_capture=False — no fresh graph memory. It also skips the forward_metadata assignment (only _prepare_* sets it), which is why the shim sets forward_metadata explicitly before delegating.

GPU evidence (new-API path): Qwen3-8B, TP=2, 2×H200 on sglang foundry HEAD, baseline→SAVE→SAVE→LOAD: per-rank final_alloc_offset=121831948288 identical across both SAVE passes and identical to the pre-rename-base run, LOAD healthy in 79s vs 250s cold SAVE, and temperature-0 SAVE-engine==LOAD-engine exact-token match on all 16/16 prompts, errors=[]. So the VMM cursor does not diverge SAVE-vs-LOAD on the ported path.

[devin-ai-integration]

Following up now that the new-API path has been exercised much harder than the 8B check:

• 8B / TP=2 on renamed sglang HEAD (foundry-org/sglang foundry, post-#26735, the init_forward_metadata_out_graph ABC): baseline→SAVE→SAVE→LOAD, errors=[], per-rank final_alloc_offset reproducible across both SAVE passes, 16/16 exact temperature-0 SAVE-engine==LOAD-engine token match.
• The in_capture=False "no second _int_workspace_buffer alloc" assumption is directly falsifiable and holds: if the replan re-allocated the workspace, the SAVE final_alloc_offset would drift and LOAD would land at a different cursor. It doesn't — SAVE final_alloc_offset is byte-identical across passes and LOAD's post-load cursor matches it exactly.
• Now also exercised on a hybrid MoE (Qwen3.5-122B-A10B, TP=4) where FlashInfer is nested under HybridLinearAttnBackend.full_attn_backend: the pre-pass runs cleanly and SAVE/LOAD reach byte-identical VMM offsets (save_after_pre_init == LOAD after_pre_init; SAVE final_alloc_offset == LOAD after_load_all_graphs).

So the new-API branch is validated on the renamed build, not just the pre-rename base — the README claim is backed by GPU evidence. (122B still has a separate, non-FlashInfer decode-graph replay issue unrelated to this shim; tracked separately.)

[devin-ai-integration]

Correct that this hinges on sglang's in_capture=False semantics (re-plan against existing wrappers, no second _int_workspace_buffer alloc, no forward_metadata overwrite). That's why the shim sets forward_metadata explicitly before the re-plan call and why we assert byte-identical VMM offsets between SAVE and LOAD as the guardrail: save_after_pre_init == LOAD after_pre_init and SAVE final_alloc_offset == LOAD after_load_all_graphs (validated equal on 122B). If a future sglang changed in_capture=False to re-allocate or re-set metadata, those offsets would diverge and the run would fault/mismatch rather than silently drift — so the determinism check is the tripwire. The legacy indices_updater_*.update(...) path remains for pre-#26735 sglang. Not making a code change here; flagging that the offset-equality assertion is the intended protection against exactly this drift.

[devin-ai-integration]

Devin Review found 1 new potential issue.

[devin-ai-integration]

📝 Info: Dense warmup now also affects the previously-untouched single-GPU / DP dense paths

Before this PR, only the EP path ran a pre-capture warmup; the comment at python/foundry/integration/sglang/hooks.py:474-475 still claims 'the validated dense / single-GPU / DP paths are untouched'. That comment refers only to the EP-gated block, but _dense_save_warmup_enabled() defaults on and gates the new _run_dense_warmup for ALL non-EP paths (SAVE hooks.py:599, LOAD hooks.py:525-528), including plain single-GPU and DP. So those paths now DO run a warmup + torch.cuda.empty_cache() and shift their VMM cursor trajectory relative to prior behavior. This appears intentional (the feature targets lazy-init hybrid models) and remains symmetric SAVE↔LOAD, but any previously-saved single-GPU/DP archive produced before this change would be incompatible with a LOAD after it (different cursor trajectory). Re-capture is required after upgrading.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

Good call-out — the neutrality invariant is exactly what needs to hold, so I re-validated the dense path with the warmup default-ON (H200:2, Qwen3-8B, --tp-size 2, baseline → SAVE → SAVE → LOAD):

[save    ] errors=[] offsets=['120439439360', '120439439360']
[save2   ] errors=[] offsets=['120439439360', '120439439360']
[load    ] errors=[]
PRIMARY  temp-0 tokens SAVE-engine == LOAD-engine : True   (16/16 prompts)
SAVE final_alloc_offset reproducible              : True

So with the warmup running, SAVE still produces a per-rank final_alloc_offset identical across both passes, and the restored LOAD engine reproduces the warmed SAVE engine token-for-token on all 16 prompts. The neutrality holds because the warmup runs entirely under stop_allocation_region() (activations go through the ordinary CUDA allocator, VMM cursor frozen), the caching allocator is emptied afterwards, and the per-bs FlashInfer wrappers the warmup populates are dropped so the real pre-pass re-allocates them in-region — the cursor entering that pre-pass is byte-for-byte what LOAD sees.

Scope is intentional: I kept it default-on for all non-EP paths (with FOUNDRY_SGLANG_SAVE_WARMUP=0 to opt out) rather than special-casing dense-TP, since the same lazy-init-inside-capture failure affects any model that defers torch.compile/JIT to first forward. final_alloc_offset here (120439439360) differs from the earlier run only because this harness lowered --cuda-graph-max-bs 256→8; determinism and SAVE==LOAD are independent of that.

[devin-ai-integration]

Agreed the old comment was misleading — fixed in a3fe80f. It now reads that the EP block is EP-only and that dense/single-GPU/DP get their own in-region warmup via _run_dense_warmup on both SAVE and LOAD. Your compatibility note is correct and intended: this shifts the dense/DP VMM cursor trajectory, so an archive captured before this change is not LOAD-compatible after it (re-capture required after upgrading). I'll add that to the README's warmup-restore section alongside the default-on note.

[devin-ai-integration]

Devin Review found 2 new potential issues.

[devin-ai-integration]

🔍 Fragmented-gap fallback assumes failed cuMemcpyDtoH leaves context healthy

_dump_mapped (python/foundry/integration/sglang/runtime.py:262-302) first tries one contiguous cuMemcpyDtoH_v2 and, on failure, probes per-2MiB page keeping only cleanly-copied runs. The correctness of the probe relies on a failed copy over a partially-unmapped span being a synchronous argument-validation error that leaves the CUDA context usable (as the docstring claims). If instead an unmapped page in the middle produces an asynchronous fault, it could leave a sticky error on the context. This fallback path is not exercised by the validated runs (which report a single contiguous interval), so it is effectively untested. Worth validating before relying on the fragmented path.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

Correct that the validated runs only exercised the single-contiguous path (one 144 MiB interval), so the per-page fallback is defensive. On the context-health assumption: cuMemcpyDtoH_v2 validates the CUdeviceptr range against the driver's VA mappings synchronously and returns CUDA_ERROR_INVALID_VALUE (1) before issuing any copy for an unmapped/partially-unmapped span — it is an argument-validation error, not an asynchronous kernel fault, so it does not leave a sticky error on the context. We actually observed exactly this earlier in development: the pre-fix whole-span copy over an unmapped gap failed with cuda driver error 1 (invalid argument) and the process stayed healthy (it raised cleanly through _cu_check), which is why the fix moved the dump before empty_cache. The primary guarantee is that the dump runs while the gap is fully mapped (caching allocator still holds every block), so the contiguous copy succeeds and the fallback is belt-and-suspenders for a hypothetical mid-warmup unmap. I'd rather keep it than silently produce a short dump; happy to add a fragmented-path GPU test if you want it exercised explicitly.