Skip to content

[Archive #39210 2026-10-11] NVFP4 W4A16 MegaMoE before rebase - #8

Draft
zianglih wants to merge 21 commits into
archive/pr39210-20261011-basefrom
archive/pr39210-20261011-head
Draft

zianglih wants to merge 21 commits into
archive/pr39210-20261011-basefrom
archive/pr39210-20261011-head

Conversation

@zianglih

Copy link
Copy Markdown
Owner

Motivation

@HumansAnd

Add NVFP4 W4A16 to the existing FlashInfer MegaMoE backend for serialized ModelOpt
checkpoints and nvfp4_online. Reuse SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1
from #35120; the default NVFP4
MegaMoE path remains W4A4.

Merge prerequisite: SGLang must first bump FlashInfer to
0.7.1, which
includes FlashInfer #5019.
The current 0.7.0.post1 pin lacks W4A16 and uses the old W4A4 FC2 keyword.
Dockerfiles and dependency pins are outside this PR's scope.

Modifications

  • Backend and precision: extend the existing MegaMoE preparation/forward path
    and use FlashInfer's native config, preprocessing and pooled workspace APIs.
    W4A16 uses packed NVFP4 weights/block scales, per-expert FP32 weight scales,
    BF16 expert activations/EP transport/combine and post-FC2 routing. Checkpoint
    activation scales are neutralized; online quantization retains its source
    precision and excluded-layer semantics. No new opt-in or quantization interface.
  • Topology and models: SM100/SM103 with attention-DP=TP=EP. GLM DSA/MTP keeps
    synchronized DP token counts and graph buckets while expert inputs remain local.
    Kimi K3 NVFP4 W4A4/W4A16 uses standard TopK and local EP/shared-expert handling;
    SiTU maps GEMM1 alpha/clamp fields to beta/linear-beta, separately from SwiGLU
    clamps. MXFP8/DeepGEMM reject unwired SiTU; Nemotron ReLU-squared is unsupported.
  • Reduction and output: default SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0
    uses external FP32 accumulation. Setting it to 1 permits nondeterministic
    BF16 atomic candidates; tuning may still select external reduction. Use native
    supports_output_view to distinguish W4A4's borrowed output from W4A16's owned output.
  • Reload and dummy weights: restore canonical W4A16 loading views and copy
    prepared weights into existing buffers. Retain the live layer, workspace owner
    and captured addresses; stage runtime scales on every forward. Initialize dummy
    weight_scale_2 to one while preserving ordinary random-weight initialization.
    Serialized W4A16 reload is covered; W4A4 reload and online RL reload are outside scope.

Startup tuning and cache ownership

  • Tune once per geometry and capacity before graph capture and prepare pooled
    workspaces: powers-of-two decode profiles cover captured, eager and adaptive-MTP
    bounds, plus the per-rank prefill ceiling from the shared runtime helper.
    All ranks choose the smallest capacity covering synchronized DP counts,
    including empty ranks; reject overflow.
  • Tune representative capacity-sized inputs, then install fixed workspaces.
    Native first-forward AUTO alone would tune the first live input while caching
    under workspace capacity. FlashInfer owns namespaced persistence; stage each
    profile once across layers and preserve session winners for pooled workspaces.
  • Scope the target rank-cache filename by physical GPU UUID so engines on disjoint
    GPUs do not overwrite each other's records. Same-GPU engines need separate
    SGLANG_CACHE_DIR roots. Cache-disabled runs use unique output filenames.
    Cache naming adds no forward-path work.
  • Use the existing target startup autotune lifecycle; do not enable
    SGLANG_FLASHINFER_AUTOTUNE_EXTEND implicitly. Skipped/disabled tuning retains
    the default native path. Native tuning requires EP=WORLD; draft MegaMoE profiles
    and new deterministic-mode behavior are outside scope.

Single-node NVSHMEM defaults

For supported single-node FlashInfer MegaMoE formats with no speculative decoding
or draft A2A none, set these defaults before NVSHMEM initialization, preserving
explicit environment overrides:

NVSHMEM_REMOTE_TRANSPORT=none
NVSHMEM_IB_ENABLE_IBGDA=0
NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1

These select the direct-P2P path and avoid the uninitialized proxy global-exit
state corrected by NVIDIA/nvshmem commit af9e2391a849.
NVIDIA/nvshmem#110 is closed; its
description and discussion provide additional background.
All EP GPUs must support peer memory access. Proxy-based global exit/timeout
polling and remote IB/UCX transport are disabled; NCCL/UID bootstrap networking
is unaffected. Multi-node jobs and drafts using other A2A backends receive no
automatic defaults. No dedicated environment-variable behavior tests are added.

Example on four Blackwell GPUs with FlashInfer 0.7.1:

SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1 python -m sglang.launch_server \
  --model-path Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \
  --quantization nvfp4_online --dtype bfloat16 \
  --moe-runner-backend flashinfer_megamoe \
  --moe-a2a-backend flashinfer_megamoe \
  --tp-size 4 --attn-dp-size 4 --ep-size 4 \
  --cuda-graph-backend-prefill disabled

Accuracy Tests

SGLang 38114937747fff49d749c3ee6a7e2f240e34530a with official FlashInfer
0.7.1 wheels (release metadata 0b46115d238d29ea2e92c6b9b2cb077d40fbc352).
All nine wheel URLs/hashes match release and install metadata; 363 non-FI package
versions/locations and six SGLang native hashes match the image. Wheel installation
uses no source build; CuTe runtime JIT still occurs at first use.

Check Result
Adapter, autotune/cache, dummy scales, DP metadata, server arguments 49 passed, 21 subtests, 20 warnings, no skips; exit 0
Native W4A16 SwiGLU/SiTU EP4, cold/warm All 16 rank reports passed; four launches exit 0
Online NVFP4 W4A16 Qwen GSM8K 191/200 (0.955); one test passed, exit 0
Serialized NVFP4 W4A16 Qwen reload Both flush modes passed; one test, exit 0
Dummy NVFP4 W4A16/W4A4 startup and generation Both passed; four DP ranks × eight finite-logprob tokens per mode; exit 0

Environment: C2 hu-pdx-111, eight B300 SXM6 AC GPUs (SM103; 275040 MiB each),
four GPUs per native/model run; one visible GPU for units. Image lmsysorg/sglang:nightly-dev-cu13-20261010-961bcf48, digest
sha256:e03dd310dc72e64ac709ceb244a79f20fd5e05448c811ab33cc777d8040061ff.
Driver 590.48.01; Python 3.12.15; Torch 2.14.1+cu130/CUDA 13.0; NVCC 13.0.88;
Triton 3.8.0; CuTe DSL/libraries 4.8.0; Transformers 5.19.0;
sglang-kernel 0.4.9; NVSHMEM 3.4.5/Python 0.4.0; NCCL 2.30.7.
FlashInfer imports installed site-packages; SGLang imports the clean pinned checkout.

Reproduce the released wheel environment
export TASK_ROOT=/data/home/ziangli/sglang-megamoe-validate-20261009
mkdir -p "$TASK_ROOT/setup" "$TASK_ROOT/results"
git clone https://github.com/zianglih/sglang.git "$TASK_ROOT/sglang"
git -C "$TASK_ROOT/sglang" checkout --detach 38114937747fff49d749c3ee6a7e2f240e34530a
python -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache \
  flashinfer-jit-cache-sm80 flashinfer-jit-cache-sm89 flashinfer-jit-cache-sm90a \
  flashinfer-jit-cache-sm100a flashinfer-jit-cache-sm103a flashinfer-jit-cache-sm120f
python - <<'PY'
import os
from pathlib import Path
rows = """flashinfer_python-0.7.1-py3-none-any.whl 69302ed6e91314a85a4392324becc194594a782ba181646d25d5c305e1b7e8c4
flashinfer_cubin-0.7.1-py3-none-any.whl 502eaa5f71e768950ada429618481a886388fcb34a5da6a2258979713cbaa327
flashinfer_jit_cache-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl 75c80222846c11f269ab72c96c1152ee675836deadf84d805e44ac62da411beb
flashinfer_jit_cache_sm80-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl 1788bc850e1b301b06a2c36d67ca333fe3a95e8cec099c0915871db1d72c7626
flashinfer_jit_cache_sm89-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl 5d3673347a7c2b6fb20f758129f215ddc57038df741806b1d07b32aa774bda91
flashinfer_jit_cache_sm90a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl d8bf578c40a236b3f07216302cc7c7ad410ea8c3f12052fcfff8791876b0b1c0
flashinfer_jit_cache_sm100a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl c7dbb9ceede5763448d23ca744ce1ccaaba377cb86e4b15304f9a411fc754808
flashinfer_jit_cache_sm103a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl 79c476784f6641afeecd9289ddd1651b75f868c173f99100d0bb8ba3441d33da
flashinfer_jit_cache_sm120f-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl 596424325c5c34baa9881149e2ef267428a8c05d28659081773500da7b79d263"""
release = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/"
python_wheel = "https://files.pythonhosted.org/packages/32/3f/3c8dbde51499902346f8f38fcdc3f92db3de25b1cdd11361f2f4dd90fd83/"
requirements = []
for line in rows.splitlines():
    filename, sha256 = line.split()
    base = python_wheel if filename.startswith("flashinfer_python-") else release
    requirements.append(base + filename + "#sha256=" + sha256)
Path(os.environ["TASK_ROOT"], "setup/wheels.txt").write_text("\n".join(requirements) + "\n")
PY
python -m pip install --no-deps --only-binary=:all: --require-hashes \
  --report "$TASK_ROOT/setup/wheel-install-report.json" -r "$TASK_ROOT/setup/wheels.txt"
export PYTHONPATH="$TASK_ROOT/sglang/python"
export SGLANG_CACHE_DIR="$TASK_ROOT/cache-wheels/sglang"
export FLASHINFER_WORKSPACE_BASE="$TASK_ROOT/cache-wheels/sglang"
export CUTE_DSL_CACHE_DIR="$TASK_ROOT/cache-wheels/cute-dsl"
export TORCH_EXTENSIONS_DIR="$TASK_ROOT/cache-wheels/torch-extensions"
export HF_HOME="$TASK_ROOT/hf-home"
export PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=1 MAX_JOBS=4

Python/cubin are 0.7.1; the aggregate and six architecture JIT-provider wheels
are 0.7.1+cu130. Setup archive SHA256:
7ed1016101cf70523c93eed2d8f475157854b7b32d9456ba91c37681e598d35f.

Registered units: actual runtime imports; unit cases mock kernel calls.

cd "$TASK_ROOT/sglang"
CUDA_VISIBLE_DEVICES=0 python -B -m pytest -q -ra \
  test/registered/unit/layers/moe/test_flashinfer_megamoe.py \
  test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py \
  test/registered/unit/model_loader/test_weight_utils.py::TestInitializeDummyWeights \
  test/registered/unit/managers/scheduler_components/test_dp_attn.py::TestDPAttnSchedulerMetadata \
  test/registered/unit/server_args/test_server_args.py::TestFlashinferMegaMoeConfig \
  --junitxml="$TASK_ROOT/results/units.xml"
49 passed, 20 warnings, 21 subtests passed in 95.75s (0:01:35)
Validation helper: smoke_autotune_gpu.py
cat > "$TASK_ROOT/smoke_autotune_gpu.py" <<'PY'
"""Native EP4 integration smoke; exercises real AUTO, pooling and graph replay.

Launch with torchrun, pinned SGLang source on PYTHONPATH, and released FI wheels. Run first
with --reuse-cache 0, then --reuse-cache 1 using the same --cache directory.
No candidate or kernel is replaced. The AUTO wrapper only counts invocations.
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import os
import types
from pathlib import Path

import torch
import torch.distributed as dist


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache", type=Path, required=True)
    parser.add_argument("--reuse-cache", type=int, choices=(0, 1), required=True)
    parser.add_argument("--decode-num-tokens", type=int, choices=(2, 16), default=2)
    parser.add_argument("--activation", choices=("silu", "situ"), default="silu")
    args = parser.parse_args()
    os.environ.update(
        SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16="1",
        SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE="0",
        SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE="bf16",
        SGLANG_FLASHINFER_MEGAMOE_MAX_TOKENS_PER_RANK="4",
        FLASHINFER_MOE_EP_KNOB_CACHE="0",
    )
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    dist.init_process_group(
        "nccl", device_id=torch.device("cuda", torch.cuda.current_device())
    )
    rank, world = dist.get_rank(), dist.get_world_size()
    assert world == 4
    torch.manual_seed(100 + rank)

    from flashinfer.autotuner import AutoTuner, autotune
    from flashinfer.moe_ep import MoEEpMegaLayer, MoEEpTensors
    from flashinfer.moe_ep.core.kernel.workspace_pool import pooled_workspace_refcount
    from flashinfer.moe_ep.cute_dsl.megamoe import bf16_nvfp4 as native
    from sglang.srt.layers import dp_attention
    from sglang.srt.layers.moe import flashinfer_megamoe as adapter
    from sglang.srt.layers.moe.flashinfer_megamoe_autotune import (
        megamoe_autotune_context,
    )
    from sglang.srt.runtime_context import get_parallel

    # Synthetic EP metadata uses the current runtime-context ownership API.
    # The real distributed accessor reads this same group during layer construction.
    get_parallel().override_permanently(
        moe_ep_group=types.SimpleNamespace(
            device_group=dist.group.WORLD, world_size=world, rank_in_group=rank
        )
    )
    counts = [2] * world
    dp_attention.get_dp_global_num_tokens = lambda: counts
    auto_calls = 0
    original_auto = native.autotune_bf16_nvfp4_mega_moe

    def count_auto(*a, **kw):
        nonlocal auto_calls
        auto_calls += 1
        return original_auto(*a, **kw)

    native.autotune_bf16_nvfp4_mega_moe = count_auto
    hidden, intermediate, experts, topk = 256, 256, 8, 4
    local_experts = experts // world
    layers = []
    for index in range(2):
        layer = torch.nn.Module()
        layer.layer_id = index
        layer.hidden_size = hidden
        layer.intermediate_size_per_partition = intermediate
        layer.num_experts = experts
        layer.num_local_experts = local_experts
        layer.moe_ep_size, layer.moe_ep_rank = world, rank
        layer.top_k = topk
        layer.quant_config = types.SimpleNamespace(use_per_token_activation=True)
        layer.moe_runner_config = types.SimpleNamespace(
            is_gated=True,
            activation=args.activation,
            gemm1_alpha=4.0,
            gemm1_clamp_limit=25.0,
            swiglu_limit=1.5,
            apply_router_weight_on_input=False,
            routed_scaling_factor=1.0,
        )
        layer.g1_alphas = torch.full((local_experts,), (1 + index) / 16, device="cuda")
        layer.g1_alphas_up = layer.g1_alphas
        layer.g2_alphas = torch.full((local_experts,), (1 + index) / 32, device="cuda")
        for prefix, rows, columns in (
            ("w13", 2 * intermediate, hidden),
            ("w2", hidden, intermediate),
        ):
            layer.register_parameter(
                f"{prefix}_weight",
                torch.nn.Parameter(
                    torch.randint(
                        256,
                        (local_experts, rows, columns // 2),
                        dtype=torch.uint8,
                        device="cuda",
                    ),
                    requires_grad=False,
                ),
            )
            layer.register_parameter(
                f"{prefix}_weight_scale",
                torch.nn.Parameter(
                    torch.ones(local_experts, rows, columns // 16, device="cuda").to(
                        torch.float8_e4m3fn
                    ),
                    requires_grad=False,
                ),
            )
        adapter.prepare_nvfp4_moe_weights_for_flashinfer_megamoe(layer)
        adapter.ensure_nvfp4_moe_layer_for_flashinfer_megamoe(layer)
        layers.append(layer)

    x = torch.randn(16, hidden, dtype=torch.bfloat16, device="cuda")
    ids = torch.stack(
        [torch.randperm(experts, device="cuda")[:topk] for _ in range(16)]
    ).int()
    scores = torch.rand(16, topk, dtype=torch.float32, device="cuda")
    scores /= scores.sum(dim=1, keepdim=True)

    def tensors(layer):
        n = counts[rank]
        return MoEEpTensors(
            x[:n],
            ids[:n],
            scores[:n],
            fc1_alpha=layer.g1_alphas,
            fc2_alpha=layer.g2_alphas,
        )

    def forward(layer):
        t = tensors(layer)
        dispatch = types.SimpleNamespace(
            hidden_states=t.hidden_states,
            topk_output=types.SimpleNamespace(
                topk_ids=t.topk_ids, topk_weights=t.topk_weights
            ),
        )
        info = adapter.FlashInferMegaMoeQuantInfo(
            mega=layer._flashinfer_megamoe_layer,
            mega_forward=layer._flashinfer_megamoe_forward,
            fc1_alpha=t.fc1_alpha,
            fc2_alpha=t.fc2_alpha,
        )
        return adapter.run_flashinfer_megamoe(
            dispatch, info, layer.moe_runner_config
        ).hidden_states

    def sync():
        torch.cuda.synchronize()
        dist.barrier()

    args.cache.mkdir(parents=True, exist_ok=True)
    # Separate startup passes must reuse the process's pooled winner,
    # even when their persistent cache paths differ.
    for index, layer in enumerate(layers):
        cache_path = args.cache / f"layer{index}" / f"rank{rank}.json"
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        if not args.reuse_cache:
            cache_path.unlink(missing_ok=True)
        tuner = AutoTuner.get()
        if cache_path.is_file():
            tuner.load_configs(str(cache_path))
        with (
            autotune(cache=None),
            megamoe_autotune_context(
                decode_num_tokens=args.decode_num_tokens,
                prefill_num_tokens=8,
            ),
        ):
            forward(layer)
        tuner.save_configs(str(cache_path))
        records = tuner.get_namespaced_records("sglang_flashinfer_megamoe")
        assert len(records) == (3 if args.decode_num_tokens == 2 else 5)
        # Generic cache payloads must be identical on all ranks; rank-local
        # timestamps would make the engine discard these caches next time.
        gathered = [None] * world
        dist.all_gather_object(gathered, records)
        assert all(item == records for item in gathered)
        assert cache_path.is_file()
    sync()
    capacities = [1, 2, 8] if args.decode_num_tokens == 2 else [1, 2, 4, 8, 16]
    expected_calls = 0 if args.reuse_cache else len(capacities)
    assert auto_calls == expected_calls, (auto_calls, expected_calls)
    for layer in layers:
        assert sorted(layer._flashinfer_megamoe_forward.workspaces) == capacities
    for capacity in capacities:
        left = (
            layers[0]
            ._flashinfer_megamoe_forward.workspaces[capacity]
            ._backend_workspace
        )
        right = (
            layers[1]
            ._flashinfer_megamoe_forward.workspaces[capacity]
            ._backend_workspace
        )
        assert left is right
        assert pooled_workspace_refcount(left) == 2
    assert os.environ["FLASHINFER_MOE_EP_KNOB_CACHE"] == "0"

    # Independent fixed-default sessions do not reuse the tuned pool entries.
    controls = []
    for layer in layers:
        state = layer._flashinfer_megamoe_forward
        control = MoEEpMegaLayer(
            bootstrap=state.bootstrap,
            fleet_params=dataclasses.replace(state.fleet_params, max_tokens_per_rank=16),
            weights=None,
            backend=dataclasses.replace(
                state.backend,
                megakernel=dataclasses.replace(state.backend.megakernel, knobs={}),
            ),
        )
        control.warmup(tensors(layer))
        controls.append(control)

    checked = []
    shapes = [[0, 1, 1, 2], [0, 1, 2, 7], [8, 8, 8, 8]]
    if args.decode_num_tokens == 16:
        shapes.append([0, 4, 9, 16])
    for shape in shapes:
        counts[:] = shape
        eager = [forward(layer) for layer in layers]
        expected = [
            control.forward(tensors(layer))
            for layer, control in zip(layers, controls, strict=True)
        ]
        sync()
        for actual, ref in zip(eager, expected, strict=True):
            assert torch.isfinite(actual).all()
            assert torch.equal(actual.view(torch.int16), ref.view(torch.int16))
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            captured = [forward(layer) for layer in layers]
        sync()
        for _ in range(2):
            graph.replay()
            sync()
        for actual, ref in zip(captured, eager, strict=True):
            assert torch.equal(actual.view(torch.int16), ref.view(torch.int16))
        for layer in layers:
            layer.g1_alphas.mul_(0.75)
        graph.replay()
        sync()
        expected = [
            control.forward(tensors(layer))
            for layer, control in zip(layers, controls, strict=True)
        ]
        sync()
        for actual, ref in zip(captured, expected, strict=True):
            assert torch.equal(actual.view(torch.int16), ref.view(torch.int16))
        graph.reset()
        checked.append(shape)

    assert auto_calls == expected_calls
    sync()
    survivors = list(layers[1]._flashinfer_megamoe_forward.workspaces.values())
    pooled = [handle._backend_workspace for handle in survivors]
    layers[0]._flashinfer_megamoe_layer.destroy()
    for handle in survivors:
        assert pooled_workspace_refcount(handle._backend_workspace) == 1
    layers[1]._flashinfer_megamoe_layer.destroy()
    assert all(pooled_workspace_refcount(workspace) == 0 for workspace in pooled)
    for control in controls:
        control.destroy()
    print(
        "MEGAMOE_AUTOTUNE_SMOKE="
        + json.dumps(
            dict(
                rank=rank,
                passed=True,
                activation=args.activation,
                precision="w4a16",
                auto_calls=auto_calls,
                reuse_cache=bool(args.reuse_cache),
                capacities=capacities,
                bootstrap_capacity=4,
                decode_num_tokens=args.decode_num_tokens,
                prefill_num_tokens=8,
                startup_contexts=2,
                checked_counts=checked,
                graph_replays=3 * len(shapes),
                bit_exact=True,
                pool_refcounts=[2, 1, 0],
            )
        ),
        flush=True,
    )
    get_parallel().clear_stamp()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()
PY

Native EP4: two layers, hidden/intermediate 256, experts 8/top-k 4; real kernels,
tuning, communication and graph replay with synthetic scheduler counts. Checks:
capacities [1,2,4,8,16], two startup contexts, rank-cache equality, pooled winners,
counts [0,1,1,2]/[0,1,2,7]/[8,8,8,8]/[0,4,9,16], runtime scale updates and
pool refs 2→1→0. silu means SwiGLU; SiTU alpha 4/clamp 25 map to beta/linear-beta.
Observed AUTO calls: 5/rank cold, 0 warm; no candidate was replaced/disabled.
Harness SHA256: 56e343a213fa155ac00a8a6567d20ca213fd61ffac5bc65314ace58ad0e3c6c4.

cd "$TASK_ROOT"
export CUDA_VISIBLE_DEVICES=0,1,2,3
# Synthetic scheduler metadata bypasses ServerArgs; use production defaults explicitly.
export NVSHMEM_REMOTE_TRANSPORT=none NVSHMEM_IB_ENABLE_IBGDA=0
export NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1
for activation in silu situ; do
  for reuse in 0 1; do
    timeout --kill-after=30s 90m torchrun --nnodes=1 --nproc-per-node=4 \
      --master-addr=127.0.0.1 --master-port=29546 \
      "$TASK_ROOT/smoke_autotune_gpu.py" \
      --cache "$TASK_ROOT/results/native-$activation-cache" \
      --reuse-cache "$reuse" --decode-num-tokens 16 --activation "$activation"
  done
done

All four launches exited 0. Complete native fields shared by all 16 reports:

{"passed":true,"precision":"w4a16","capacities":[1,2,4,8,16],"bootstrap_capacity":4,"decode_num_tokens":16,"prefill_num_tokens":8,"startup_contexts":2,"checked_counts":[[0,1,1,2],[0,1,2,7],[8,8,8,8],[0,4,9,16]],"graph_replays":12,"bit_exact":true,"pool_refcounts":[2,1,0]}
activation reuse_cache rank auto_calls
silu false 0 5
silu false 1 5
silu false 2 5
silu false 3 5
silu true 0 0
silu true 1 0
silu true 2 0
silu true 3 0
situ false 0 5
situ false 1 5
situ false 2 5
situ false 3 5
situ true 0 0
situ true 1 0
situ true 2 0
situ true 3 0

Model checks: BF16 expert activations/EP transport, TP=attention-DP=EP=4, MegaMoE runner/A2A, W4A16 and
external FC2 reduction. Online converts all routed experts with the unchanged
4over6 recipe and dense source semantics; GSM8K: 200 examples, 128 threads,
completion API, 512-token cap, score > 0.90. Reload covers both cache-flush modes
and memory release/resume; exact text/token IDs and logprob difference ≤1e-4.
The wrapper changes only snapshot paths/ports, keeps registered flags/assertions,
and writes the dummy fixture manifest.

cd "$TASK_ROOT/sglang"
unset NVSHMEM_REMOTE_TRANSPORT NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_DISABLE_LOCAL_ONLY_PROXY
CUDA_VISIBLE_DEVICES=4,5,6,7 python - <<'PY'
import importlib.util, json, os, socket, unittest
from pathlib import Path
from huggingface_hub import snapshot_download
root = Path(os.environ["TASK_ROOT"])
cases = [
    ("backends/test_flashinfer_nvfp4_online_moe_backend.py",
     "TestFlashinferMegaMoeBackendNvFp4OnlineW4A16",
     "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8",
     "5a5a776300a41aaa681dd7ff0106608ef2bc90db", 31000),
    ("rl/test_update_weights_from_disk_blackwell.py",
     "TestServerUpdateWeightsFromDiskNVFP4W4A16MegaMoE",
     "nvidia/Qwen3-30B-A3B-NVFP4",
     "2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3", 31010),
]
for file, name, repo, revision, port in cases:
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", port))
    snapshot = Path(snapshot_download(repo, revision=revision))
    assert snapshot.name == revision
    index = json.loads((snapshot / "model.safetensors.index.json").read_text())
    assert all((snapshot / f).is_file() for f in set(index["weight_map"].values()))
    record = dict(repo=repo, revision=revision, path=str(snapshot), complete=True)
    directory = root / "results/model-cache"
    directory.mkdir(parents=True, exist_ok=True)
    (directory / (repo.replace("/", "--") + ".json")).write_text(json.dumps(record))
    spec = importlib.util.spec_from_file_location("registered", "test/registered/" + file)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    case = getattr(module, name)
    assert case.model == repo
    case.model = str(snapshot)
    module.DEFAULT_URL_FOR_TEST = case.base_url = f"http://127.0.0.1:{port}"
    result = unittest.TextTestRunner(verbosity=2).run(
        unittest.defaultTestLoader.loadTestsFromTestCase(case))
    print("REGISTERED_E2E_RESULT=" + json.dumps(dict(
        case=name, passed=result.wasSuccessful(), tests_run=result.testsRun,
        failures=len(result.failures), errors=len(result.errors), skipped=len(result.skipped))))
    assert result.testsRun > 0 and result.wasSuccessful() and not result.skipped
PY

Online runtime: BF16 KV on all four ranks (auto in server-info); effective
max-running-requests=3839/DP. Observed 14 tuned capacities:
1,2,4,8,16,32,64,128,256,512,1024,2048,3839,4096.
Raw online footer (exit 0):

metrics={'score:std': np.float64(0.20730412441627882), 'score': np.float64(0.955), 'latency': 11.176999901421368, 'output_throughput': 2535.9220050091903}
ok

----------------------------------------------------------------------
Ran 1 test in 1272.648s

OK
REGISTERED_E2E_RESULT={"case": "online-gsm8k", "passed": true, "tests_run": 1, "failures": 0, "errors": 0, "skipped": 0}

Reload runtime (startup logs): FP8 E4M3 KV, max-running-requests=4096/DP;
13 power-of-two capacities 1–4096, 52 decode graph sizes through 512. Both
flush modes passed release/resume, exact text/token and 1e-4 logprob checks.
Raw reload footer (exit 0):

ok

----------------------------------------------------------------------
Ran 1 test in 1063.293s

OK
REGISTERED_E2E_RESULT={"case": "serialized-reload", "passed": true, "tests_run": 1, "failures": 0, "errors": 0, "skipped": 0}
Validation helper: run_dummy_megamoe.py

Equivalent reproduction with the same server arguments and assertions; the tested helper additionally records fixture and worker metadata.

cat > "$TASK_ROOT/run_dummy_megamoe.py" <<'PY'
import argparse,math,os,socket
from pathlib import Path
import requests
from huggingface_hub import snapshot_download
from sglang.test.test_utils import DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, kill_process_tree, popen_launch_server
p=argparse.ArgumentParser()
p.add_argument("--w4a16",choices=("0","1"),required=True)
p.add_argument("--port",type=int,required=True)
a=p.parse_args()
for k in ("NVSHMEM_REMOTE_TRANSPORT","NVSHMEM_IB_ENABLE_IBGDA","NVSHMEM_DISABLE_LOCAL_ONLY_PROXY"):
    assert k not in os.environ
with socket.socket() as s:s.bind(("127.0.0.1",a.port))
model=snapshot_download("nvidia/Qwen3-30B-A3B-NVFP4",revision="2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3")
assert all((Path(model)/f).is_file() for f in ("config.json","tokenizer.json"))
env={"SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16":a.w4a16,"SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE":"0","SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE":"bf16","SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION":"0"}
args='--dtype bfloat16 --tp-size 4 --attn-dp-size 4 --ep-size 4 --fp4-gemm-backend flashinfer_cutedsl --moe-runner-backend flashinfer_megamoe --moe-a2a-backend flashinfer_megamoe --context-length 2048 --chunked-prefill-size 256 --max-running-requests 16 --cuda-graph-max-bs-decode 8 --cuda-graph-backend-prefill disabled --mem-fraction-static 0.4'.split()
args=["--load-format","dummy","--json-model-override-args",'{"num_hidden_layers":2}',*args]
url=f"http://127.0.0.1:{a.port}"
proc=popen_launch_server(model,url,timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,env=env,other_args=args)
try:
    for rank in range(4):
        r=requests.post(url+"/generate",json={"input_ids":[1,2,3,4],"sampling_params":{"temperature":0,"max_new_tokens":8,"ignore_eos":True},"routed_dp_rank":rank,"return_logprob":True},timeout=180)
        r.raise_for_status()
        m=r.json()["meta_info"];lp=m["output_token_logprobs"]
        assert m["completion_tokens"]==8 and len(lp)==8 and all(math.isfinite(x[0]) for x in lp)
finally:
    kill_process_tree(proc.pid)
PY

Dummy loader: pinned NVIDIA config/tokenizer, only two layers overridden;
real hidden/expert geometry. BF16 combine, external FC2 reduction, per-token
activation 0, TP=attention-DP=EP=4; configured context 2048/prefill 256/requests 16/decode graph 8,
prefill graphs off/memory 0.4. Resolved per-DP prefill 64/requests 4; captured
[1,2,4], FP8 E4M3 KV. Each DP rank generates eight finite-logprob tokens
from IDs [1,2,3,4], temperature 0/ignore-EOS. Helper records responses and worker
metadata, then stops its server. SHA256: b61ff2b30562bec77ff86e8826b5d72d621c8c151e52a3ad5d543afe1ac1a407.

unset NVSHMEM_REMOTE_TRANSPORT NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_DISABLE_LOCAL_ONLY_PROXY
export CUDA_VISIBLE_DEVICES=0,1,2,3
timeout --kill-after=30s 90m python "$TASK_ROOT/run_dummy_megamoe.py" --w4a16 1 --port 31020
timeout --kill-after=30s 90m python "$TASK_ROOT/run_dummy_megamoe.py" --w4a16 0 --port 31030
dummy-w4a16: exit 0
DUMMY_SMOKE_RESULT={"w4a16": "1", "passed": true, "requests": 4, "completion_tokens": 32, "all_output_logprobs_finite": true, "layers": 2}
dummy-w4a4: exit 0
DUMMY_SMOKE_RESULT={"w4a16": "0", "passed": true, "requests": 4, "completion_tokens": 32, "all_output_logprobs_finite": true, "layers": 2}

Scope: native references are independent fixed-default sessions, not a
PyTorch oracle. SiTU uses a small EP4 geometry, not full K3 E2E. Dummy startup
is not accuracy. No W4A4 model-accuracy/reload, online RL reload, Nemotron or
concurrent-engine cache claim. Correctness timings are not serving performance measurements.

Speed Tests and Profiling

End-to-end serving performance has not been measured on the current head.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #38022492097
Latest PR Test (Extra): ❌ Run #38022491936
Latest PR Test (AMD ROCm 10): ❌ Run #38022492254

为单节点 NVFP4 W4A16 设置 P2P-only NVSHMEM 默认值,绕过 proxy 未初始化退出状态;保留显式配置并排除使用其他 A2A 后端的 draft。按要求不增加 behavioral regression tests。
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant