[zianglih]

Motivation

@HumansAnd

Cover CuTe DSL NVFP4 W4A16 weight reloads used by disjoint Miles training/rollout workers. Raw transfers must update the buffers captured by CUDA graphs and refresh derived scale tensors.

Modifications

• Contract: stage checkpoint-format expert tensors in a pinned CPU replica, copy bytes into the original GPU buffers, then finalize the packed weights.
• Regression: run A → B → A updates; require stable parameter/cached-scale addresses, bit-exact eager and graph outputs against independent fresh loads, changed B output, and exact recovery of A. Poison the captured output before each replay so an unwritten output cannot pass.
• Reuse: extend the existing real FusedMoE NVFP4 backend test and quantization helpers. Adjust its CI time estimate for the added test.
• Scope: this adds regression coverage; the current CuTe DSL reload implementation already passes. The test simulates raw transfer with byte copies; network transport and distributed expert sharding are exercised separately in Miles #3216, which completed two training/rollout rounds and three P2P syncs on four training plus four rollout B300 GPUs. That integration run had zero gradients; changed-weight semantics are checked here by A → B → A. The existing fused-indexer mapping fix in [sglang-miles] Fix p2p weight registry: draft worker overwrite, fused indexer mapping, registration retries #35312 remains a separate sglang-miles dependency.

Accuracy Tests

• Environment: one NVIDIA B300 SXM6 AC GPU from an eight-GPU C2 node, driver 590.48.01; image radixark/miles:dev-202609111228, amd64 manifest sha256:f774dee95a80a213c30489cf82b13295d8a50338ea39f5aaadb0377d3e4a9bba.
• Software: PyTorch 2.13.0+cu130 / CUDA 13.0, FlashInfer 0.6.18, CuTe DSL 4.6.2, TE 2.17.0. SGLang source is this PR on upstream main, selected through PYTHONPATH, rather than the image's sglang-miles source.
• Workload: 8 local experts, hidden/intermediate dimensions 1024, top-k 2, 32 BF16 tokens; serialized NVFP4 weights, W4A16 activations, fused finalize disabled, TP=EP=1. No numerical tolerance was relaxed (rtol=atol=0 against a fresh load of the same packed weights).
• Tested commit: 5c499272f91d08cc5837f770d1f8867b7661bce2.

From this checkout in the image:

CUDA_VISIBLE_DEVICES=1 PYTHONPATH="$PWD/python" \
python test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py \
  TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload

[CI Test Method] TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload
/opt/sglang/lib/python3.12/site-packages/nvidia_cutlass_dsl/dsl_packages/cutlass/_mlir_helpers/op.py:294: DeprecationWarning: make_trivial_tiled_mma with ab_dtype is deprecated, use the overload with separate a_dtype and b_dtype instead
  res_or_list = opFunc(*args, **kwargs, loc=loc)
/opt/sglang/lib/python3.12/site-packages/flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py:2525: DeprecationWarning: Migrated to examples/CuTeDSL/helpers/static_persistent_tile_scheduler.py (BSD-3). The wheel copy will be removed in a future release.
  tile_sched_params = utils.PersistentTileSchedulerParams(
/opt/sglang/lib/python3.12/site-packages/nvidia_cutlass_dsl/dsl_packages/cutlass/_mlir_helpers/op.py:294: DeprecationWarning: FastDivmodDivisor is deprecated in favor of cute.FastDivmodDivisorV2 / cute.fast_divmod_create_divisor_v2. V2 additionally carries the scalar divisor across kernel boundaries (2 MLIR values per object instead of 1), so '.divisor' is readable inside kernels; arithmetic (divmod, //, %) is unchanged.
  res_or_list = opFunc(*args, **kwargs, loc=loc)
/opt/sglang/lib/python3.12/site-packages/nvidia_cutlass_dsl/dsl_packages/cutlass/utils/static_persistent_tile_scheduler.py:292: DeprecationWarning: Migrated to examples/CuTeDSL/helpers/static_persistent_tile_scheduler.py (BSD-3). The wheel copy will be removed in a future release.
  new_params = PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
/opt/sglang/lib/python3.12/site-packages/nvidia_cutlass_dsl/dsl_packages/cutlass/utils/static_persistent_tile_scheduler.py:780: DeprecationWarning: Migrated to examples/CuTeDSL/helpers/static_persistent_tile_scheduler.py (BSD-3). The wheel copy will be removed in a future release.
  return StaticPersistentRuntimeTileScheduler(
/opt/sglang/lib/python3.12/site-packages/nvidia_cutlass_dsl/dsl_packages/cutlass/utils/static_persistent_tile_scheduler.py:725: DeprecationWarning: Migrated to examples/CuTeDSL/helpers/static_persistent_tile_scheduler.py (BSD-3). The wheel copy will be removed in a future release.
  return StaticPersistentRuntimeTileScheduler(
.
----------------------------------------------------------------------
Ran 1 test in 5.169s

OK

Negative control: on the same test body, replace refresh_cutedsl_standard_scales_for_weight_update with a no-op using unittest.mock.patch. The B update must fail against its fresh-load control. Complete output:

test_flashinfer_cutedsl_w4a16_weight_reload (test_nvfp4_moe_backends.TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload) ... [CI Test Method] TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload

  test_flashinfer_cutedsl_w4a16_weight_reload (test_nvfp4_moe_backends.TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload) (state='B') ... FAIL

======================================================================
FAIL: test_flashinfer_cutedsl_w4a16_weight_reload (test_nvfp4_moe_backends.TestNvFp4MoeBackends.test_flashinfer_cutedsl_w4a16_weight_reload) (state='B')
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/hai-workspace/miles-cutedsl-w4a16-p2p/sglang-main/test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py", line 303, in test_flashinfer_cutedsl_w4a16_weight_reload
    torch.testing.assert_close(eager, expected, rtol=0, atol=0)
  File "/opt/sglang/lib/python3.12/site-packages/torch/testing/_comparison.py", line 1631, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not equal!

Mismatched elements: 32206 / 32768 (98.3%)
Greatest absolute difference: 0.14453125 at index (15, 30)
Greatest relative difference: 800.0 at index (22, 512)

----------------------------------------------------------------------
Ran 1 test in 7.419s

FAILED (failures=1)

• Validation boundary: no full-model learning or BF16 parity claim; no RDMA performance claim. Other backend tests were not rerun for this test-only change.

Speed Tests and Profiling

No inference implementation change or speedup claim. Only test execution time was measured.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (Not applicable to this test-only change.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (Exact regression and negative control above; no inference speed change.)
• Follow the SGLang code style guidance.

Full GPU CI matrix remains for maintainer CI.

Review and Merge Process

Normal SGLang review and CI; this PR does not merge or duplicate #35312.

[github-actions]

Thanks @zianglih. Closing this because it is still a draft and has not been updated in 30 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.