[sglang-miles] Kimi K3 colocated RL: LoRA fixes and in-place MXFP4 Marlin reload

[yueming-yuan]

The Kimi K3 RL serving work reached sglang-miles through main, and the miles K3 PR (radixark/miles#1825) now runs against sglang-miles instead of the old sglang-miles-k3 branch. Running it on radixark/miles:dev found five places where the merge lost or never had what colocated K3 LoRA RL needs. Each commit is one of them; all were found by reading the traceback back to the branch diff, and the last two by --check-weight-update-equal / --check-lora-weight-equal.

• num_shared_experts (K3's config key) was not recognized when sizing the shared-expert LoRA buffers, so init_state died with KeyError('gate_up_proj') on the first dense layer.
• The output gate K3 installs by patching o_proj.forward never ran under LoRA: RowParallelLinearWithLoRA calls quant_method.apply on the base layer, so the MLA output went out ungated and the gate GEMM on the alt stream was never joined (capturing stream has unjoined work). Restores the lora_input_transform hook from the K3 branch (3768a87).
• Routed-expert LoRA buffers were sized by hidden_size; K3's latent MoE feeds the experts routed_expert_hidden_size (3584).
• post_load_weights read o_proj.params_dtype through the LoRA wrapper.
• _merge_bfa_weights stored f_b_proj.weight as a second Parameter on the attention module, which hides f_b_proj.weight from named_parameters(); every weight update silently skipped it (fb03504). Alias the storage instead.

Running the full model on 16 GB300 nodes (TP16, Marlin, --no-lora-use-virtual-experts) then showed the train/rollout KL growing with the adapter norm (0.0023 at step 0, 0.0047 by step 60) while the July stack stayed flat. A line-by-line diff of the LoRA path against that stack found four more losses, each its own commit:

• Mxfp4MoEMethod never exposed get_marlin_quant_info, so FusedMoEWithLoRA refused the Marlin backend outright; the LoRA Marlin runner also only ran silu_and_mul, so K3's gated SiTU experts failed at CUDA-graph capture.
• The Marlin tile padding of the expert intermediate (K3 at TP16: 192 -> 256) was ignored by the MoE LoRA hooks: the up delta was written at the logical stride into the gate padding and the down LoRA-A consumed the padded activation. Restores output_slice_stride in fused_moe_lora and the marlin_intermediate_size plumbing.
• The interleaved gate/up LoRA-B layout keyed on gemm1_alpha alone; K3 sets it for SiTU but loads w1/w3 contiguously, so every routed-expert gate/up delta was scrambled.
• The fused MoE front read the shared-expert base weights straight through the LoRA wrappers, so the shared_experts adapters the trainer streams were never applied.
• The LoRA Marlin runner used atomic-add for MXFP4 where the base runner does not.
• --lora-no-cpu-backup: the streamed adapter no longer keeps a per-TP-rank host copy (~45 GB/rank for the 896-expert adapter); the IPC views are installed into the pool at apply time.

Full-parameter RL re-pushes the MXFP4 base on every sync, which needs one more piece:

• Mxfp4MoEMethod.process_weights_after_loading repacked the experts into fresh Parameters, so a later load_weights found no weight_loader on them and the runner's begin_weight_update had nothing to restore. The method now records the checkpoint and Marlin layouts per layer, restore_weights_before_loading views the storage back to the checkpoint shape, and the repack runs expert by expert into the same storage, so the Parameter objects (and their weight_loader) survive. Unit test: test/registered/e2e/rl/test_mxfp4_marlin_inplace_reload.py.
• The last commit (Mxfp4MarlinComparable accepting the checker's is_shuffled) is a temporary cherry-pick of a standalone fix for sglang-miles; it drops out of this PR once that lands.

Validation

miles CI (radixark/miles#1825, run-ci-model-scripts on 8xH200, radixark/miles:dev + this branch synced at test time) on Pinaster/Kimi-K3-4layer-64experts:

• test_kimi_k3_4layer_lora_ci.py: two rollouts and syncs; --check-lora-weight-equal matches on every rank and the base weight checker compares equal after each sync.
• test_kimi_k3_4layer_ci.py (full-parameter): the MXFP4 base is re-quantized and re-pushed on every sync; --check-weight-update-equal compares the Marlin experts equal in dequantized space ([sglang-miles] Compare MXFP4 Marlin experts in dequantized space in the weight checker #40028) after the in-place reload.

The Marlin-padding, interleave and shared-expert fixes were validated on the 4-layer TP8 layout with a trained adapter re-sharded from the 16-node run (rollout-0 logprob abs diff back to the zero-adapter floor); the padded-intermediate path itself only triggers at TP16 and was validated on the 16-node GB300 run.

sglang's own CI does not run on sglang-miles PRs (the check-changes maintenance gate), so the registered tests that exercise the touched files were run by hand on a fresh 8xH200 devbox (radixark/miles:dev, this branch at e5dbafcada installed in place of the baked sglang):

test	result
e2e/rl/test_mxfp4_marlin_inplace_reload.py, e2e/rl/test_weight_checker_mxfp4_marlin.py	1 + 9 passed
unit/utils/test_weight_checker_comparator.py	11 passed
lora/test_fused_moe_lora_kernel.py	108 passed
quant/test_marlin_moe.py, kernels/ops/moe/test_moe_wna16_marlin.py	2 passed / 678 subtests; 7 passed, 24 skipped
lora/test_virtual_experts_kernels.py, unit/lora/test_experimental_sgl_marlin_{multi_prefill,shared_outer_reduce}.py	13 + 2 + 4 passed
unit/models/test_kimi_k3_bfa_overlap.py, kernels/ops/test_kimi_k3_prerequisite_ops.py, unit/layers/quantization/test_mxfp4_sm90_cutlass.py	4 + 12 + 19 passed
rl/test_lora_load_from_tensor.py, lora/test_multi_lora_backend.py, lora/test_lora_overlap_loading.py, lora/test_lora_eviction.py	6 + 3 + 9 + 2 passed
lora/test_lora_qwen3_8b_logprob_diff.py, lora/test_lora_moe_tp_logprob_diff.py (TP2)	2 passed; 2 passed, 1 skipped
GPT-OSS-20B MXFP4, Marlin runner vs triton_kernel, greedy 32 tokens x 4 prompts	3/4 identical text, max abs logprob diff <= 0.056; the 4th diverges after 20 identical tokens

Five files fail identically on plain sglang-miles (reproduced on the same devbox with the base branch installed) and are not touched by this PR: unit/utils/test_weight_checker.py (a helper returns 2 values, the test unpacks 3), rl/test_weight_checker_e2e.py and rl/test_update_weights_from_tensor.py (they push weights without begin_weight_update, which the scheduler has asserted on since the RL session protocol), lora/test_moe_lora_info.py (its fake forward_batch lacks global_dp_buffer_len), lora/test_lora_spec_decoding.py (EAGLE3 draft context-length check). Not runnable on H200: the B200/B300-only registrations.

Follow-ups not in scope

• The AR-fusion (_symm_o_proj_forward) and SP-MoE (_sp_o_proj_forward) patches on o_proj.forward are bypassed by the LoRA wrapper the same way the output gate was; neither is active in the smoke run (auto-probe off, no symm mem). They need an output-side hook before K3 LoRA runs with AR fusion.

---

CI States

Latest PR Test (Base): ❌ Run #35702051199
Latest PR Test (Extra): ❌ Run #35702050805
Latest PR Test (AMD ROCm 10): ❌ Run #35702051045