[bobboli]

Replaces #38282, which GitHub auto-closed when its base branch was deleted and recreated.

Summary

• Capture logical expert IDs from DeepSeek V4 HashTopK before logical-to-physical remapping for R3 replay. Mask CUDA-graph padding, disable target capture from draft modules, and reject the unsupported HashTopK waterfill/capture combination.
• Synchronize CUDA IPC tensor imports in the weight updater before returning to the producer, including flattened buckets and other tensor load formats.
• Support repeated TRT-LLM MXFP4 online loads: permute expert weights in place, retain checkpoint-layout scale parameters, and maintain separate shuffled scale copies with stable kernel addresses. Pass unsigned byte views at the FlashInfer packed-FP4 call boundary without changing load parameters.
• Register the GEMM1 clamp limit as a non-persistent buffer so memory offload preserves it.
• Keep the padded-row Triton fast path off CPU tensors.

Current rebase (2026-10-02)

Rebased the PR-only commits onto sglang-miles at 6c94e57f9ff5b84ba9746f5983ff13ce7e84c68b. Current head: 0f0a367bad0d693ed8af044316111c25d919ed67.

• 091bc49cf3: restore the capture helper's original API, keep device checks at launch sites, validate routed MXFP4 hardware before model construction, and make the flattened-bucket helper's exception-only return contract explicit.
• ec8a0fbf1c: integrate upstream expert padding and input prequantization with hot reload. Recognize signed packed weights when computing hidden size and clear padding before repeat shuffles. Adapt existing tests to the current loader/runtime-context interfaces.

Experiment terminology

• P1: TE MXFP8 training with MXFP8 rollout.
• P2: TE MXFP8 training with MXFP4 weights × MXFP8 activations for rollout routed experts, without trainer-side QAT.
• P3: The P2 setup with routed-expert MXFP4 fake quantization during training; trainer GEMMs still use TE MXFP8.

Review follow-up (2026-10-02)

0f0a367bad adds inverse-shuffle restore for partial MXFP4 updates, live layer-buffer ownership for the clamp, XPU-compatible launch guards, in-place activation constant refresh, and first-build-only allocator cleanup. Empty/weight-only/scale-only sessions failed all six new cases before the restore fix and pass afterward.

The source stack also passed an ARM/GB300 preflight: 69 Miles integration tests, 6 QAT GPU tests and 30 SGLang tests (7 HashTopK, 9 MXFP4 GPU, 5 FP8 ownership, 9 UE8M0). Checkpoint metadata, optimizer/scheduler/RNG state and rollout data-state were audited before the pending P2 resume. This preflight is not an end-to-end training pass.

GitHub CI remains gated on run-ci. The PR author's attempt to add it was denied by GitHub permissions; a maintainer must add the label.

Validation of the preceding rebase

• 320 CPU tests passed, run one test file per process: HashTopK (7), model loader (3), derived-weight cache (2), FP8 runner fallback (4), runner ownership/static buffers (5), MegaMoE (6), routed capture (5), UE8M0 requantization (9), and server arguments (279).
• 3 B300 GPU tests passed: fresh-load versus hot-reload parity and storage-address preservation for aligned/padded expert sizes, plus real MXFP4 x MXFP8 forward/reload with standard and FlashInfer dispatch inputs.
• All applicable changed-file pre-commit checks and git diff --check passed.
• No new end-to-end RL run was performed after this rebase. These local checks do not establish a full GitHub CI pass.

The results below are historical, from earlier heads.

FP8/MXFP8 activation-state follow-up (2026-09-23)

Commit 104716dc91 also fixes the TRT-LLM activation tensors introduced in the new base's FP8/MXFP8 path. These alpha/beta/clamp tensors were ordinary attributes, so the memory-saver static-state backup omitted them. Repeated postprocessing also replaced their storage even though CUDA graphs retained the old addresses. Register them as non-persistent buffers and copy into existing storage on reload.

Evidence: an isolated reproduction using the actual preparation and static-state export/import functions left the old clamp at zero after simulated storage loss; re-preparation created a different pointer holding 10. The same function-level check passes after the fix (restored values, stable storage, no checkpoint entries). The existing runner-ownership test fixture now also covers this lifecycle.

The prior 23837cdabf source passed 7 HashTopK tests and the GPU MXFP4 hot-reload test on GB300, but an end-to-end P1 run subsequently exhibited repetitive output and could not collect valid training groups. The activation-state defect is confirmed; attributing the complete rollout failure to it still requires a new RL run. Updated-source integration job completed successfully on GB300: 51 Miles integration tests, 6 GPU MXFP4 QAT tests, and 12 SGLang tests (HashTopK, MXFP4 hot reload, and FP8 runner ownership/activation-state lifecycle). Changed-file pre-commit passed with GITHUB_BASE_REF=sglang-miles. P1 smoke, with all other source/config pins unchanged, collected valid normal-mode samples and completed one valid optimizer step (ppo_kl=0, rollout/train log-prob absolute difference 0.02877). It then failed with host-memory OOM during trainer CPU offload; this is not a three-step end-to-end pass. The first actor_train timer was 703.9 seconds, not a steady-state performance measurement. P1 retry subsequently passed all three optimizer steps and exited successfully (1:03:47), with the same source pins and training settings and additional host-memory diagnostics. actor_train: 699.8 / 127.0 / 105.5 seconds; ppo_kl=0 for all three steps; train/rollout log-prob absolute difference 0.02852 / 0.02860 / 0.02731. No host or CUDA OOM was observed in this run. This does not establish the earlier host-OOM root cause: the physical nodes and Ray-head/PP-stage placement differed. P2 and P3 both failed during the initial SGLang autotune forward, before online weight sync or training: ValueError: Unsupported trtllm-gen input tensor. The MXFP4 method passed signed checkpoint parameters directly to FlashInfer instead of zero-copy unsigned-byte views; the same issue exists in the pre-rebase PR head. The idle outer launchers were cancelled. SGLang commit f9cfbec2d2 fixes that boundary with zero-copy uint8 views, preserving signed load parameters and storage. The new real-kernel case fails with the original exception before the fix; after the fix, all 13 local B300 SGLang tests pass (including real MXFP4 x MXFP8 forward and reload). ARM preflight completed successfully: 70 tests passed (51 Miles, 6 GPU QAT, 13 SGLang), exit 0:0. P2 retry and P3 retry used f9cfbec2d2 and subsequently failed at initial FP32 export quantization, as detailed below. The completed P1 run above used SGLang 104716dc91; it predates this MXFP4-only follow-up.

FP32 export compatibility follow-up (2026-09-23)

P2 retry and P3 retry passed SGLang initialization but failed during initial online weight sync, before any optimizer step: the UE8M0 blockwise-FP8 helper rejected Bridge's FP32 export ([2048, 4096] shared-expert weight). This is separate from the preceding packed-MXFP4 kernel boundary bug.

SGLang commit becddd92ece674066025b3e132768391b658214a accepts BF16 or FP32 without changing the quantization math or rounding through BF16. Both new CPU regression cases fail before the fix. After it: 9 CPU tests and 13 related B300 GPU tests pass; the actual failing shape produces byte-identical FP8 weights and packed scales for BF16 weights versus their lossless FP32 exports. Changed-file pre-commit passes with GITHUB_BASE_REF=sglang-miles.

ARM preflight completed (0:0, 4:35): 79 tests passed (51 Miles, 6 GPU QAT, 22 SGLang).

Completed direct-HF smoke, without Bridge #36

All three arms completed three valid optimizer steps, including repeated online weight updates and normal-mode rollout, and exited successfully. No host/CUDA OOM was observed in these successful runs; all nine logged ppo_kl values are zero, and loss/gradient norms are finite.

Arm	Elapsed	Three actor_train timers (seconds)
P1	1:03:47	699.8 / 127.0 / 105.5
P2	1:04:46	698.6 / 122.6 / 115.9
P3	1:04:32	704.2 / 118.8 / 141.3

Pins: Miles P1 946b00a41c, P2/P3 f7c4529350; unmodified Bridge bridge base 2e09c234a; Megatron-LM f2fc0efcb. P1 used SGLang 104716dc91; P2/P3 used becddd92ec, which adds the packed-MXFP4 boundary fix and FP32 UE8M0 input support. P1 was not rerun at the final SGLang head. All code was committed source; no runtime/site-packages patches or offline weight conversion.

Setup: Preview, TE MXFP8 training, native/cuDNN attention, 8 x 4 GB300, normal mode, R3 enabled, 4K response cap, eval disabled, GPU optimizer with CPU trainer offload; independent experiment records. P3's routed-expert QAT hook installation was confirmed in the runtime log.

This establishes that Bridge #36 is not required for this tested recipe. It does not establish long-run stability, eval quality, QAT benefit, or the cause of the earlier host-OOM failure. First-step timings include warm-up; the subsequent 105.5–141.3 seconds are below the historical 185-second reference-attention figure but are not a controlled comparison with the historical 82.6-second cuDNN result. This is targeted validation, not a full CI pass.

---

CI States

Latest PR Test (Base): ❌ Run #37047037453
Latest PR Test (Extra): ❌ Run #37047036885
Latest PR Test (AMD ROCm 10): ❌ Run #37047037268

[Zhichenzzz]

🤖 Posted autonomously by Claude Code acting on the user's behalf. Single-pass diff review of this PR; one correctness concern found. Please push back if the conclusion is off.

update_weights_from_tensor now ignores the flattened_bucket helper's return value.

Before this PR, the flattened_bucket branch returned the result of _update_weights_from_flattened_bucket(...) directly. After the refactor, that call's return value is discarded and the function falls through to the unconditional return True, "Success":

https://github.com/sgl-project/sglang/blob/d6d709637bf12c95377c200b3a3fb9fe40736102/python/sglang/srt/model_executor/model_runner_components/weight_updater.py#L571-L588

Severity note (verified against the head source): today _update_weights_from_flattened_bucket can only fail by raising — its sole return is (True, "Success") (weight_updater.py#L423-L453) — so there is no live behavior change. It is a latent trap: if that helper ever starts reporting failure via its return tuple, the caller will silently report success. Suggest either propagating the status or changing the helper to return None so the convention is explicit.

[HJSang]

LGTM — nice set of fixes, and the comments explaining why each invariant exists make this much easier to review than it could have been.

Things I checked independently and that hold up:

• The shuffle refactor is equivalent against FlashInfer 0.6.18: shuffle_matrix_a(x, t) == x[get_shuffle_matrix_a_row_indices(x, t)] and shuffle_matrix_sf_a(x, t) == block_scale_interleave(x[get_shuffle_matrix_sf_a_row_indices(x, t)]), with num_elts_per_sf defaulting to 16. Folding [w1, w3] -> [w3, w1] into the gather as w3_w1_rows[perm[r]] reproduces the old shuffle(reorder(w)) ordering, and half is valid for the scales too since their row count matches w13.shape[1].
• supports_output_view is a real per-backend attribute in flashinfer/moe_ep (modes/mega_layer.py plus the cutedsl backends), so gating on it is a genuine fix rather than a guess — the DeepGEMM mega path doesn't advertise it.
• get_platform().is_sm100 resolves to device_capability_majors=[10], so sm_103 (B300/GB300) passes. Good, given where this was validated.
• Every call site of the new capture_routed_experts_if_allowed(allow_capture, ...) signature is updated, and HashTopK is a plain nn.Module, so the isinstance chain in disable_routed_experts_capture_for_draft is safe.
• Merges cleanly onto current sglang-miles (only topk.py overlaps the 12 commits since the base).

A few comments, none blocking.

1. Mxfp4FlashinferTrtllmMoEMethod has no restore_weights_before_loading

This is the one I'd most like addressed. process_weights_after_loading now permutes w13_weight / w2_weight in place, which makes it non-idempotent. restore_weight in model_loader/loader.py is the hook for exactly this ("Undo in-place quant packing so fresh weights can be loaded"), and Mxfp4MoEMethod.restore_weights_before_loading implements it for the Marlin path.

Without it, begin_weight_update -> a session that updates only non-MoE tensors -> end_weight_update runs postprocess_weight over already-permuted weights and permutes them a second time, silently. A full-model sync is unaffected, so this doesn't touch the validated recipe — but /begin_weight_update and /update_weights_from_tensor are public endpoints and the name list is caller-controlled. Applying the inverse permutation (argsort(w13_rows) / argsort(w2_rows)) in a restore hook would make the postprocess correct either way.

Related, in test_mxfp4_flashinfer_trtllm_hot_reload.py: the restore_weight(layer, torch.device("cuda")) call and the assertions after it are currently a no-op, since _apply_quant_method_hook does a hasattr check and skips the method. So the "loader hooks bracket every update session; neither may disturb the checkpoint-layout parameters" half of that test isn't actually exercised yet.

2. The *_shuffled split has two consequences worth a comment

• Memory: the fp32 checkpoint-layout scales used to be freed when they were replaced by the shuffled version. Now both layouts stay resident; the fp32 copy alone is roughly a quarter the size of w13_weight. Necessary for reloadability, just worth stating.
• EPLB: DeepseekV2MoE.get_moe_weights picks up w13_weight_scale_inv_shuffled only because copy_or_rebind_param registers it as a Parameter and filter_moe_weight_param_global_expert admits anything with shape[0] == num_local_experts. That's load-bearing — if these ever became buffers, a rebalance would update the checkpoint-layout scales and leave the kernel reading stale shuffled ones. A one-line comment pinning that would help.

3. apply() reads the method attribute, not the buffer

gemm1_clamp_limit=self._gemm1_clamp_limit_tensor — but the reason for register_buffer is that the layer's copy is what offload maintains. They agree today because _import_static_state restores in place, but nn.Module._apply (.to(), .cuda()) replaces _buffers[key] outright and would strand the method's reference. Reading layer._gemm1_clamp_limit_tensor matches what the FP8 side already does.

4. _can_fuse_padded_region: x.is_cuda also excludes XPU

On XPU both _is_cuda and _is_npu are False, so a tensor reaches this predicate and now falls back to eager even though Triton works there. The comment says "fall back to eager for host tensors"; not x.is_cpu matches that intent exactly and keeps ROCm and XPU on the fused path.

5. Smaller things

• get_quant_method: after the SM100 raise, the is_sm90 or is_sm120 branch is unreachable for flashinfer_trtllm_routed, so the comment above it now describes only the flashinfer_mxfp4 arm. One check placed just before the TRT-LLM return would cover both backends and remove the dead branch — at the cost of changing flashinfer_mxfp4 behaviour on unsupported arches, which may be a separate decision.
• apply(): w13 / w2 are now unconditionally uint8 views, so ... if w2.dtype == torch.uint8 else ... can't take the else.
• _update_weights_from_flattened_bucket's return True, "Success" is now discarded by the caller.
• _prepare_flashinfer_trtllm_activation_params: existing.fill_(float(value)) would skip building tensor on the reload path. Also existing is not None and tensor is None falls into the register_buffer branch, which is the pointer-orphaning case the in-place branch exists to prevent — unreachable while moe_runner_config is fixed for the layer's lifetime, but the condition reads as though it were handled.
• torch.cuda.empty_cache() at the end of process_weights_after_loading now runs per MoE layer on every weight-update session. With the in-place design there's far less to reclaim than before, so it may be worth restricting to the first build.
• disable_routed_experts_capture_for_draft: elif hasattr(module, "allow_routed_experts_capture") would cover HashTopK and any future router class without the lazy import and without another elif next time.
• HashTopK.allow_routed_experts_capture could be a class attribute next to the docstring with the same "draft-side blocks set this False" comment that TopKConfig carries; right now the HashTopK side has no explanation of what the flag is for.

6. CI hasn't run on this head

All the check failures are the same gate — Missing required label 'run-ci' — and every suite is skipping. Worth adding the label so the six new test files actually execute before this lands.

[bobboli]

@Zhichenzzz Addressed the flattened-bucket return-value finding in 091bc49: the internal helper now explicitly returns None and raises on reconstruction/load failure. The public method still synchronizes the device before returning (True, "Success"), so no helper status is silently discarded and CUDA IPC lifetime protection is preserved.

[bobboli]

@HJSang Follow-up to your review: addressed in 0f0a367.

1. The TRT-LLM MXFP4 method now implements the QuantizeMethodBase restore contract. It inverse-permutes weights in place before updates, preserving parameter/storage identity; repeated restore/postprocess calls are safe. Empty, weight-only and scale-only update sessions failed all 6 regression cases before the fix and pass afterward, with both padded and aligned dimensions.
2. Documented the retained FP32 scales' memory cost and why shuffled scales must remain Parameters for EPLB.
3. apply() reads the layer's live clamp buffer. The real-kernel test now replaces buffers through Module._apply and invalidates the old clamp reference.
4. Launch sites use not tensor.is_cpu, preserving the XPU/ROCm fused path. XPU execution was not tested locally.
5. Removed the dead uint8 conditional, made activation-buffer reload use fill_ and reject enabled/disabled transitions, and limited empty_cache() to the first layout build. The flattened helper's None return was fixed in 091bc49. HashTopK now shares TopKConfig, so its flag documentation and draft opt-out follow the existing contract rather than a new duck-typed attribute. The earlier model-load validation keeps the existing flashinfer_mxfp4 architecture support unchanged.
6. I tried adding run-ci, but GitHub rejected it with AddLabelsToLabelable permission denied. Could a maintainer please add the label?

Validation: 9 B300 GPU cases, 7 HashTopK CPU tests and 5 FP8 ownership/lifecycle CPU tests passed; changed-file pre-commit passed. This is not a full GitHub CI pass.

[Zhichenzzz]

LGTM