feat(sglang): support checkpoint-engine refit

[tianyi-zhang-02]

What does this PR do?

Adds checkpoint-engine refit for non-colocated SGLang:

• runs one NIXL receiver per SGLang rank
• aligns rank-local streams by weight name when transport bucket boundaries differ
• supports dense and full-weight MoE refit
• registers the SGLang functional test and updates the refit docs

The supported envelope is intentionally narrow: one node per logical engine, sglang_cfg.dp_size=1, and shard_expert_weights=false. Unsupported configurations now fail during setup instead of partway through refit.

The branch is rebased on current main. I also removed a defensive continue that could spin forever if the stream invariant were ever broken; the existing mismatch path now fails loudly.

Closes #3288.

Validation

• 149 unit tests across SGLang generation and weight sync passed in the earlier full run
• focused 2-GPU NIXL probe passed with divergent per-rank bucket boundaries and cross-process CUDA IPC
• ruff and formatting checks pass after the rebase and follow-up fix
• checked the pinned official SGLang 0.5.12.post1 source: the previously suspected begin_weight_update state machine is not present in this version

The remaining acceptance check is the registered GRPO functional test against a real SGLang server in CI; I do not have that environment locally.

#3330 and #3426 overlap with parts of this work. I am happy to keep this standalone or rebase it onto whichever direction you prefer.

cc @RayenTian @yuki-97

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tianyi-zhang-02]

GPU follow-up

I don't have an environment where tests/functional/grpo_sglang_nixl_non_colocated.sh runs end to end, so I validated the transport layer with a focused 2-GPU probe. Ran against nixl 1.3.0, the version this PR pins.

Setup

Single node, 2 GPUs, torch 2.7.1+cu128, nixl 1.3.0 (UCX backend, which bound to real IB HCAs plus cuda_ipc — not a tcp/shm fallback).

Topology mirrors what the PR introduces:

• 2 policy-side senders, one NIXLCheckpointEngine each (cuda:0, cuda:1)
• one actor holding 2 rank-local receivers — the piece this PR adds
• the SGLang HTTP server is stubbed by a separate process that deserializes the flattened bucket from the CUDA IPC handle and compares tensors, so the IPC handoff is real and cross-process

To exercise _aligned_checkpoint_engine_batches the shards are deliberately asymmetric (rank 1 carries 2× rank 0's rows). The probe counts each rank's transport batches and fails if they come out equal, since in that case the alignment path was never exercised at all.

Result

create_backend(UCX): OK
refit returned in 2.91s
transport batches per rank : [5, 6] (diverged=True)
aligned batches consumed   : 9
cross-process IPC loads    : 9
rank 0: 9 weights bit-exact (6912 rows, device=cuda:0)
rank 1: 9 weights bit-exact (13824 rows, device=cuda:1)
kv cache invalidated       : True
engines finalized cleanly
RESULT: PASS

Every weight arrived bit-exact on both ranks, exactly once, per-rank content distinct, across genuinely different transport boundaries (5 vs 6 batches).

Caveats

1. The shard asymmetry is synthetic. Bucket boundaries are a function of each rank's (name, dtype, size) sequence, so ranks with identical shapes produce identical boundaries and the alignment path degenerates to a pass-through. This shows the mechanism is correct when fed divergence; it does not establish how often real workloads diverge. If uneven TP shards cannot actually occur in the supported configurations, the aligner could be replaced with an assert — I'd rather be told that than guess.
2. No GRPO loop, no real SGLang server, no train/token_mult_prob_error check. The functional script still needs a CI run. It is now registered with L1_Functional_Tests_SGLang.sh; previously nothing referenced it, so it had never run.
3. create_backend("UCX") failed intermittently with NIXL_ERR_BACKEND and cleared on re-run. Two hypotheses tested and both refuted: process-local configuration (a matrix over OpenSSL preload × CUDA init × nixl library paths came back clean, including the exact configuration that had failed the run before) and IB availability (several nodes sampled, all showing identical /sys/class/infiniband with the same ports ACTIVE, all passing). Cause unknown. Note this shares an error string with the UCX_TLS item in the environment notes below, but the failing runs had UCX_TLS unset, so I do not think they are the same thing. Flagging only in case it is a known 1.3.0 flake. Happy to give exact run counts if useful.
4. nixl 1.3.0 + Ray SIGSEGVs during interpreter teardown, after finalize() returns cleanly and after the result is known. The probe calls os._exit() once the verdict is in, so a teardown crash cannot masquerade as a transfer failure. That is in the probe only, not in shipped code.

One proposed change: drop an unreachable branch

In _aligned_checkpoint_engine_batches, by the time control reaches this every rank has either just been refilled with a non-empty batch (empty batches raise above) or is finished, and a finished-and-empty rank is already caught by the break/raise immediately above.

         if any(finished[rank] and not queue for rank, queue in enumerate(pending)):
             raise RuntimeError(
                 "Checkpoint-engine streams ended with different weights across "
                 "SGLang ranks."
             )
-        if any(not queue for queue in pending):
-            continue

         aligned = [[] for _queue in pending]

The reason to delete rather than keep it as a defensive guard: if the invariant ever did break, continue returns to the top, the empty rank is skipped by the refill because it is finished, and the same branch is taken again — a silent infinite loop inside a refit, instead of the loud RuntimeError directly above it. It does not guard anything; it would hide something. Happy to keep it if you disagree.

Environment notes, in case they save someone time

• UCX needs a network-capable transport even for an intra-node refit. UCX_TLS=sm,self,cuda_copy,cuda_ipc fails create_backend("UCX") with NIXL_ERR_BACKEND; adding tcp fixes it, as does leaving UCX_TLS unset.
• nixl 1.3.0 is a dispatcher wheel selecting nixl-cu{torch.version.cuda major}; it lists both nixl-cu12 and nixl-cu13 as unconditional Requires-Dist, so an air-gapped install needs --no-deps plus the one backend matching the image.
• Its binding links libssl.so.3 / libcrypto.so.3; UCX itself is vendored in the wheel. On an OpenSSL 1.1 image OpenSSL is the only gap, and it surfaces as a bare ImportError.
• NIXL requires the same bucket_size_bytes on sender and receiver; mismatched sizes fail with createXferReq: length mismatch at index 0. Consistent with _resolve_bucket_size_bytes computing one value for everyone — and it means divergent boundaries can only originate from differing shard shapes.
• A CUDA IPC handle cannot be re-imported by the process that exported it (CUDA error: invalid device context), so any test stubbing update_weights_from_tensor must decode in a separate process.
• Ray clears CUDA_VISIBLE_DEVICES on num_gpus=0 actors; set RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 if the actor manages devices itself.

Also happy to push

A property-test module for the alignment logic: randomized per-rank bucket layouts asserting global weight order is preserved, plus a receiver that poisons its recycled buffer to pin down the double-buffer contract in _receive_weight_chunk_batches. That second one is the valuable half — a fake engine that allocates fresh tensors every batch will pass even if the code holds references into memory NIXL is about to overwrite, so without it the whole class of bug is untestable.

One thing worth confirming, since I could not verify it without a real server: SGLang 0.5.12.post1 asserts self._weight_update_in_progress in update_weights_from_tensor, set only by begin_weight_update, which appears nowhere in this tree. The colocated path posts to the same endpoint and presumably passes CI, so either _apply_sglang_compat_patches rewrites it or the dispatch differs. Does my path land on whatever makes the colocated one work?

[tianyi-zhang-02]

Pushed the property tests I mentioned, plus four more gaps. Each one was verified by mutation — the test fails when the behaviour is broken and passes when it is not — so none of them are assertions that cannot fail:

Property	Mutation it catches
Recycled receive buffers	making the aligner prefetch instead of advancing only empty ranks
Multi-dtype batches	posting only the first dtype group
weight_version across refits	bumping per POST instead of per refit
base_gpu_id remapping	using the physical id instead of the local one
Payload index ↔ SGLang rank	transposing the per-rank payload list

The buffer one is the reason I wanted this in the tree. nixl.py hands out views into a rotating pool, so a fake engine that allocates fresh tensors each batch makes that whole bug class invisible — every test passes while weights are silently corrupted. _RecyclingEngine poisons a buffer once it is recycled. The current aligner survives it because it only calls anext for a rank whose deque is already empty, so at most one batch per rank is outstanding; that is load-bearing and was previously unstated and untested.

The payload-index one is the property with the worst failure mode: SGLang indexes serialized_named_tensors by its own TP rank, so a transposed list loads every shard onto the wrong GPU and still returns success.

149 unit tests passing across tests/unit/models/generation/sglang/ and tests/unit/weight_sync/, ruff and pyrefly clean. Still no GRPO loop and nothing cross-node — the functional script needs a CI run.

[yuki-97]

#3330 and #3426 overlap with parts of this work. I am happy to keep this standalone or rebase it onto whichever direction you prefer.

hi @xiuhu17 @Kh4L, could you help to take a look on this?