Skip to content

[Feat] FlashInfer: Add NVFP4 W4A16 and SiTU support to MegaMoE - #44036

Merged
Fridge003 merged 6 commits into
sgl-project:mainfrom
zianglih:feat/flashinfer-megamoe-w4a16-followup
Oct 16, 2026
Merged

Fridge003 merged 6 commits into
sgl-project:mainfrom
zianglih:feat/flashinfer-megamoe-w4a16-followup

Conversation

@zianglih

@zianglih zianglih commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

@HumansAnd

Add opt-in NVFP4 W4A16 to FlashInfer MegaMoE for serialized ModelOpt checkpoints
and nvfp4_online, reusing SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1. The default
NVFP4 MegaMoE path remains W4A4.

This is the remaining integration from #39210
on top of merged #43789.
Startup autotuning, cache ownership, native workspace pooling and DP token-count
synchronization are already upstream and are not changed here. The original PR's
branch and history are preserved.

Dependency: the existing FlashInfer 0.7.1 pin includes
FlashInfer #5019.
No Dockerfile or dependency changes.

Modifications

  • Precision and preparation: extend the existing MegaMoE adapter with the
    native W4A16 config/preprocessing. Keep packed NVFP4 weights/block scales,
    per-expert FP32 weight scales, BF16 activations/EP transport/combine and
    post-FC2 routing. Neutralize serialized checkpoint activation scales and
    bypass activation quantization for online W4A16. Retain the existing source
    precision and excluded-layer behavior for online weight conversion.
  • Activation and model wiring: support gated SiLU/SwiGLU and SiTU for NVFP4
    W4A4/W4A16. Kimi K3 uses standard TopK and local EP/shared-expert handling;
    its GEMM1 alpha/clamp fields map to SiTU beta/linear-beta, separately from
    SwiGLU clamps. MXFP8/DeepGEMM reject unwired SiTU. Nemotron ReLU-squared is
    unsupported. Requirements remain SM100/SM103 and attention-DP=TP=EP;
    W4A16 additionally requires BF16 model dtype and combine.
  • Reload: restore canonical loading views over the four prepared
    weight/block-scale allocations, then copy newly prepared data into those
    allocations. Share this loader/preparation lifecycle between W4A4 and W4A16. Keep the live
    MegaMoE layer, tuned profiles and workspace owner so CUDA graph pointers remain
    stable; per-expert runtime scales are staged on each forward. FlashInfer 0.7.1
    captures W4A4 input normalization as a scalar: reloading a different input
    scale is rejected and requires a server restart. Weight, weight-scale and
    per-expert FC1-output scale updates remain supported. The registered
    W4A16 serialized reload test covers both cache-flush modes; the existing
    storage regression covers both precisions.
  • Dummy weights: initialize FP4 global weight_scale_2 tensors to one,
    preserving ordinary random-weight initialization. Retain the simple
    neutral-scale regression test.
  • FC2 reduction: default SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0
    uses external FP32 accumulation. Setting it to 1 permits nondeterministic
    BF16 atomic candidates; tuning can still select external reduction.
  • Current upstream contracts: preserve the default non-MegaMoE draft runner
    and A2A behavior. An unset draft A2A backend now resolves to none for a
    FlashInfer MegaMoE target, and the NVSHMEM guard accounts for that default.
  • Tests: extend the existing adapter/storage and architecture cases, plus
    the registered online NVFP4 GSM8K and serialized weight-reload classes.
    No full Kimi K3 E2E or dedicated NVSHMEM environment tests are added.

Single-node NVSHMEM defaults

For single-node FlashInfer MegaMoE, with no speculative decoding or a default/
explicit none draft A2A backend, set these defaults before initialization:

NVSHMEM_REMOTE_TRANSPORT=none
NVSHMEM_IB_ENABLE_IBGDA=0
NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1

Explicit environment settings are preserved. These defaults apply to all
supported FlashInfer MegaMoE formats and select direct peer memory access,
avoiding the uninitialized proxy global-exit state fixed by
NVIDIA/nvshmem commit af9e2391a849.
The closed NVIDIA/nvshmem#110 has
additional background. Remote IB/UCX transport and proxy-based global-exit/
timeout polling are disabled; bootstrap networking is unaffected. All EP GPUs
must support peer memory access. Multi-node jobs and explicit other draft A2A
backends receive no automatic defaults.

Example server command

Four Blackwell GPUs, FlashInfer 0.7.1, online expert-weight conversion:

CUDA_VISIBLE_DEVICES=0,1,2,3 \
SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1 \
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0 \
SGLANG_FP4_IGNORED_LAYERS='' \
SGLANG_ENABLE_JIT_DEEPGEMM=False \
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 \
  --enable-dp-lm-head \
  --cuda-graph-max-bs-decode 64 \
  --cuda-graph-backend-prefill disabled \
  --max-running-requests 64 \
  --chunked-prefill-size 16384 \
  --mem-fraction-static 0.7

Convert every routed expert for this example; excluded FP8 expert layers cannot
use the NVFP4 MegaMoE runner. With attention DP4, the 16384-token configured
prefill chunk resolves to 4096 tokens per rank. The existing startup autotuner
therefore prepares a 4096-token per-rank prefill profile.

Accuracy Tests

Current source: 7337b0c6aa709132170b4128eb83b8db8c4fb3be, based on
upstream main 12dc62310288926b0f38e33ff2d466dee2bf6bb9. This includes the
CI-required ancestor ff8d964cf5b6c463966bca5016018d2fa0836326.
All six feature commits are patch-identical after rebase; the scope remains 15 files.

Current-head checks: pre-commit and diff checks passed. The focused CPU
adapter harness passed 29 tests; autotune/cache checks passed 15 tests with one
CUDA-only skip; the neutral-scale regression and accessor source check each
passed. These harnesses use real CPU Torch tensors with serving/GPU imports
stubbed or exact-source AST loading; they do not establish CUDA correctness.
The title also passes the exact upstream title-validation script.

29 passed, 1 warning in 1.35s
Ran 16 tests in 3.115s — OK (skipped=1)
Neutral-scale regression: Ran 1 test in 0.001s — OK
Accessor source check: Ran 1 test in 20.960s — OK

C2 evidence scope: the GPU results below predate this rebase. Units, native,
dummy and online cases ran at 7b173e5261ac63ab5a1d6617a1d2ed626ab4aec1;
serialized reload cases ran at 562fe4e70f48de798c9260c71af821b8bad10b5f.
Those two tested revisions have identical production Python trees and retained
tests; only documentation and removal of the extra W4A4 test class differ.
The integration patches are unchanged by the rebase, but surrounding upstream
runtime code changed. GPU validation has not been rerun on the current head.

Environment: eight B300 SXM6 AC GPUs, 275040 MiB each, driver 590.48.01;
four GPUs per model/native case. Image
lmsysorg/sglang:nightly-dev-cu13-20261013-7224fce1, amd64 digest
sha256:009eafc447f47fc595ea5f9820a9d7a07a9d144f3c8b90ab8ef568f4124475a7.
Python 3.12.15, PyTorch 2.14.1+cu130 (CUDA 13.0), CuTe DSL 4.8.0,
NVSHMEM 3.4.5 and nvshmem4py 0.4.0. The image already contained all nine
released FlashInfer 0.7.1 / 0.7.1+cu130 packages. No package install/change or
FlashInfer source build. All 373 Python distributions and six SGLang native
libraries stayed unchanged. Only the source-matched SGLang Rust TreeCore
binding was built into an isolated cache; tests load it with build mode never.

Fixtures: Qwen3 FP8 online checkpoint revision
5a5a776300a41aaa681dd7ff0106608ef2bc90db; NVIDIA Qwen3 NVFP4 serialized checkpoint
revision 2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3. Cached snapshots were checked
against config/index hashes and four shard sizes each; full weight payloads
were not hashed.

Check Result
Local CPU adapter 29 passed; real Torch tensors, native/serving imports stubbed
C2 registered units 65 passed, 35 subtests passed; no failures/errors/skips
Native W4A16 SiLU/SiTU Cold and fresh-process warm passed on all four ranks; 192 graph replays
Dummy W4A16 / W4A4 Each: 4 routed requests, 32 output tokens, finite logprobs
Online W4A16 GSM8K 191/200 correct (95.5%), unchanged threshold >90%
Registered serialized W4A16 reload Both cache-flush modes passed; 1 test, zero failures/errors/skips (port-separated retry)
Local serialized W4A4 reload control Both cache-flush modes passed; 1 test, zero failures/errors/skips
  • Online settings: TP4/attention-DP4/EP4, BF16 compute and KV, online 4over6
    MSE recipe with fast math and E4M3-256, all routed experts converted, memory
    fraction 0.7, no MTP. The unchanged registered test uses decode graph max 128,
    disabled prefill graphs, 200 examples, 128 evaluation threads and a 512-token
    output cap. Effective per-worker prefill is 4096, running cap 3839, KV capacity
    1,965,696 and context limit 262144. The earlier example command is a separate
    configuration, not the measured accuracy setup.
  • Native settings: H=I=256, E=8, top-k=4, EP4; capacities 1/2/4/8/16 and
    DP counts [0,1,1,2], [0,1,2,7], [8,8,8,8], [0,4,9,16]. Two owners/
    startup contexts; five cold AUTO calls per rank, zero warm calls; pooled refs
    2→1→0. Borrowed outputs are cloned immediately. Eager/captured outputs are
    finite and bit-exact against a separate fixed-default native session before
    and after a visible FC1-scale update. This is not an independent BF16 oracle.
  • Dummy settings: two-layer full-geometry fixture, BF16 compute, FP8 E4M3 KV,
    context 2048, per-rank prefill 64, running cap 4, memory fraction 0.4, decode
    graphs 1/2/4 and native profiles 1/2/4/64; no MTP. This checks finite output,
    not model accuracy.
    The measured model configurations resolve as follows (all use TP4/attention-DP4/EP4 and no MTP):
Case Compute / KV dtype Per-rank prefill Running cap per worker Context KV tokens per worker Memory fraction Decode graph max
Online W4A16 BF16 / BF16 4096 3839 262144 1,965,696 0.70 128
Serialized W4A16 reload BF16 / FP8 E4M3 4096 4096 40960 4,930,816 0.87 512
Serialized W4A4 reload BF16 / FP8 E4M3 4096 4096 40960 4,932,224 0.87 512
  • Reload contract: same-checkpoint reload, both cache-flush modes, weights/
    KV/graph release-resume, exact text/token IDs and logprobs within 1e-4. The
    output comparison routes to DP0; it does not prove equality across all DP
    workers. W4A4 is a local control reusing the registered W4A16 class with its
    opt-in set to 0; no additional upstream class.
  • Boundaries: no full Kimi K3/GLM-5.2 E2E, online RL reload, MTP validation,
    or speedup claim. Server cases leave the three NVSHMEM variables unset to
    exercise production defaults; synthetic native cases set them explicitly.
Raw test results
Local CPU adapter:
29 passed, 1 warning in 0.89s

C2 registered units:
65 passed, 22 warnings, 35 subtests passed in 264.17s (0:04:24)

git diff --check: passed
Pre-commit: all applicable checks passed

Warnings were the environment's unknown asyncio_mode pytest setting and
upstream import/API deprecations. No C2 unit cases skipped. These GPU results are from the explicit C2 run
on the source revisions listed above. Its final checks confirmed the clean
GPU-tested source, unchanged package/native-library hashes and Rust cache, and
all eight GPUs idle. GitHub CI evaluates the rebased head separately.

Native result records below have passed=true, bit_exact=true,
borrowed_outputs_cloned=true, fc1_scale_mutation_observed=true, capacities
[1,2,4,8,16], startup contexts 2, and pool refs [2,1,0] in every row.

Activation Cache reuse Rank AUTO calls Graph replays
silu false 0 5 12
silu false 1 5 12
silu false 2 5 12
silu false 3 5 12
silu true 0 0 12
silu true 1 0 12
silu true 2 0 12
silu true 3 0 12
situ false 0 5 12
situ false 1 5 12
situ false 2 5 12
situ false 3 5 12
situ true 0 0 12
situ true 1 0 12
situ true 2 0 12
situ true 3 0 12
{"w4a16": "1", "passed": true, "requests": 4, "completion_tokens": 32, "all_output_logprobs_finite": true, "layers": 2}
{"w4a16": "0", "passed": true, "requests": 4, "completion_tokens": 32, "all_output_logprobs_finite": true, "layers": 2}

Online:
metrics={'score:std': np.float64(0.20730412441627882), 'score': np.float64(0.955), 'latency': 13.70782539783977, 'output_throughput': 2063.201067943059}
Ran 1 test in 1005.320s
OK
REGISTERED_E2E_RESULT={"case": "online-gsm8k", "passed": true, "tests_run": 1, "failures": 0, "errors": 0, "skipped": 0}

The accuracy harness emits latency and throughput fields above; they are raw
incidental output, not a controlled performance comparison. Disjoint GPU groups
were used concurrently for correctness only. No performance conclusion is drawn.

Initial W4A16 startup on HTTP 31010 collided with the concurrently running
W4A4 server on HTTP 31011: the former's detokenizer and latter's tokenizer
both selected port 31245. No model assertions ran in that attempt. Full logs
were preserved; the retry uses HTTP 34010 with unchanged source, helpers,
checkpoint, settings and assertions. The commands below run sequentially.

Ran 1 test in 170.025s
FAILED (errors=1)
REGISTERED_E2E_RESULT={"case": "serialized-reload-w4a16", "passed": false, "tests_run": 1, "failures": 0, "errors": 1, "skipped": 0}

Completed reload output:

W4A16 (HTTP 34010 retry):
Ran 1 test in 794.303s
OK
REGISTERED_E2E_RESULT={"case": "serialized-reload-w4a16", "passed": true, "tests_run": 1, "failures": 0, "errors": 0, "skipped": 0}

W4A4 (HTTP 31011 local control):
Ran 1 test in 1563.989s
OK
REGISTERED_E2E_RESULT={"case": "serialized-reload-w4a4", "passed": true, "tests_run": 1, "failures": 0, "errors": 0, "skipped": 0}
Public reproduction commands and exact registered settings

These commands reproduce the registered unit selection, online W4A16 accuracy class, and serialized reload checks with the public checkpoints. They do not run a serving performance comparison. Execute the model commands one at a time on four idle Blackwell GPUs; change CUDA_VISIBLE_DEVICES if a different four-GPU group is assigned.

Image and exact source

Use lmsysorg/sglang:nightly-dev-cu13-20261013-7224fce1 (OCI index digest sha256:20d36eeb32df0418a165641b51399fd4dc5239dad57f9b8d1f9974c4c9263ca6). The validation node has eight NVIDIA B300 SXM6 AC GPUs, driver 590.48.01. Each model case uses four GPUs.

The image already includes released FlashInfer Python/cubin 0.7.1, the aggregate and six architecture JIT-cache packages 0.7.1+cu130, Torch 2.14.1+cu130, Python 3.12.15, CuTe DSL 4.8.0, NVSHMEM 3.4.5, and nvshmem4py 0.4.0. No package installation or FlashInfer source build is needed. Runtime CuTe DSL specialization can compile during first use.

Run inside that image:

set -euo pipefail
export TASK_VALIDATION_ROOT="/workspace/sglang-megamoe-validation"
export TASK_SOURCE_SHA="562fe4e70f48de798c9260c71af821b8bad10b5f"
mkdir -p "$TASK_VALIDATION_ROOT"
cd "$TASK_VALIDATION_ROOT"
git clone --filter=blob:none --no-checkout \
  https://github.com/sgl-project/sglang.git sglang
git -C sglang fetch origin "$TASK_SOURCE_SHA"
git -C sglang checkout --detach "$TASK_SOURCE_SHA"
test "$(git -C sglang rev-parse HEAD)" = "$TASK_SOURCE_SHA"
test -z "$(git -C sglang status --porcelain --untracked-files=no)"

export PYTHONPATH="$TASK_VALIDATION_ROOT/sglang/python"
export PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=1 MAX_JOBS=4
export SGLANG_CACHE_DIR="$TASK_VALIDATION_ROOT/cache/sglang"
export FLASHINFER_WORKSPACE_BASE="$TASK_VALIDATION_ROOT/cache/flashinfer"
export CUTE_DSL_CACHE_DIR="$TASK_VALIDATION_ROOT/cache/cute-dsl"
export TORCH_EXTENSIONS_DIR="$TASK_VALIDATION_ROOT/cache/torch-extensions"
export TRITON_CACHE_DIR="$TASK_VALIDATION_ROOT/cache/triton"
export CUDA_CACHE_PATH="$TASK_VALIDATION_ROOT/cache/cuda"
export TORCHINDUCTOR_CACHE_DIR="$TASK_VALIDATION_ROOT/cache/torchinductor"
export HF_HOME="$TASK_VALIDATION_ROOT/hf-home"
mkdir -p "$TASK_VALIDATION_ROOT/results"

# Exercise the production single-node defaults in server argument handling.
unset NVSHMEM_REMOTE_TRANSPORT NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_DISABLE_LOCAL_ONLY_PROXY

python -B - <<'PY'
import importlib.metadata as md
import os
from pathlib import Path
import flashinfer
import sglang
import torch
assert md.version("flashinfer-python") == "0.7.1"
assert md.version("flashinfer-cubin") == "0.7.1"
for package in (
    "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",
):
    assert md.version(package) == "0.7.1+cu130", package
assert Path(sglang.__file__).resolve().is_relative_to(
    Path(os.environ["TASK_VALIDATION_ROOT"]) / "sglang/python"
)
print({"sglang": sglang.__file__, "flashinfer": flashinfer.__file__,
       "torch": torch.__version__, "cuda": torch.version.cuda})
PY

# Prepare only the matching SGLang Rust TreeCore binding, if not cached.
CUDA_VISIBLE_DEVICES='' CARGO_BUILD_JOBS=4 SGLANG_RUST_BUILD_MODE=auto \
  python -B -c 'from sglang.srt.mem_cache.rust_tree_core.extension import bindings; print(bindings.__file__)'
export SGLANG_RUST_BUILD_MODE=never
CUDA_VISIBLE_DEVICES='' \
  python -B -c 'from sglang.srt.mem_cache.rust_tree_core.extension import bindings; print(bindings.__file__)'

These commands pin the GPU-tested reload revision 562fe4e70f48de798c9260c71af821b8bad10b5f to reproduce the reported C2 evidence. The unit/native/dummy/online runs used 7b173e5261ac63ab5a1d6617a1d2ed626ab4aec1, with identical production code and retained tests. For a new run on the rebased PR, set TASK_SOURCE_SHA=7337b0c6aa709132170b4128eb83b8db8c4fb3be; no GPU result is claimed for that revision here.

Registered unit selection

cd "$TASK_VALIDATION_ROOT/sglang"
CUDA_VISIBLE_DEVICES=0 python -B -m pytest -q -ra \
  test/registered/unit/layers/moe/test_flashinfer_megamoe.py \
  test/registered/unit/layers/test_kimi_mlp_parallel_groups.py \
  test/registered/unit/layer_boundary/test_kimi_k3_stage_kernels.py \
  test/registered/unit/test_runtime_context.py::TestTheAccessorsHaveNoCallersOutsideTheirPackage \
  test/registered/unit/test_runtime_context.py::TestMoeFlagsGroup::test_flashinfer_megamoe_speculative_defaults \
  test/registered/unit/test_runtime_context.py::TestMoeFlagsGroup::test_flashinfer_megamoe_speculative_explicit_wins \
  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_VALIDATION_ROOT/results/units.xml"

Pinned registered-model runner

Save this local runner outside the checkout. It downloads or reuses the exact public snapshot, changes only its model path and port, and invokes the actual registered test. For the W4A4 control it changes the same reload class's environment flag and display label in this process; it adds no upstream test class and keeps the inherited assertions intact.

cat > "$TASK_VALIDATION_ROOT/run_registered.py" <<'PY'
import argparse
import importlib.util
import json
import os
from pathlib import Path
import socket
import unittest
from huggingface_hub import snapshot_download

parser = argparse.ArgumentParser()
parser.add_argument("case", choices=("online", "reload-w4a16", "reload-w4a4"))
parser.add_argument("--port", type=int, required=True)
args = parser.parse_args()
root = Path(os.environ["TASK_VALIDATION_ROOT"])
with socket.socket() as probe:
    probe.bind(("127.0.0.1", args.port))
for key in ("NVSHMEM_REMOTE_TRANSPORT", "NVSHMEM_IB_ENABLE_IBGDA", "NVSHMEM_DISABLE_LOCAL_ONLY_PROXY"):
    assert key not in os.environ, f"Unset {key} to exercise production defaults"

if args.case == "online":
    repo = "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8"
    revision = "5a5a776300a41aaa681dd7ff0106608ef2bc90db"
    filename = "backends/test_flashinfer_nvfp4_online_moe_backend.py"
    class_name = "TestFlashinferMegaMoeBackendNvFp4OnlineW4A16"
else:
    repo = "nvidia/Qwen3-30B-A3B-NVFP4"
    revision = "2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3"
    filename = "rl/test_update_weights_from_disk_blackwell.py"
    class_name = "TestServerUpdateWeightsFromDiskNVFP4W4A16MegaMoE"

snapshot = Path(snapshot_download(repo_id=repo, revision=revision))
assert snapshot.name == revision
index = json.loads((snapshot / "model.safetensors.index.json").read_text())
assert all((snapshot / name).is_file() for name in set(index["weight_map"].values()))
source = root / "sglang/test/registered" / filename
spec = importlib.util.spec_from_file_location("registered_megamoe_reproduction", source)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
module.DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{args.port}"
case = getattr(module, class_name)
assert case.model == repo
case.model = str(snapshot)
case.base_url = module.DEFAULT_URL_FOR_TEST
if args.case == "reload-w4a4":
    case.launch_env = {
        **case.launch_env,
        "SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16": "0",
    }
    case.backend_test_suites = (
        {**case.backend_test_suites[0], "name": "flashinfer_megamoe_nvfp4_w4a4_local_control"},
    )
print(json.dumps({
    "case": args.case, "registered_class": class_name,
    "public_model": repo, "revision": revision, "source": str(source),
    "model_path": str(snapshot), "url": case.base_url,
    "launch_env": getattr(case, "launch_env", getattr(case, "extra_env", {})),
    "extra_args": getattr(case, "extra_args", None),
    "backend_test_suites": getattr(case, "backend_test_suites", None),
    "update_test_suites": getattr(case, "update_test_suites", None),
}, indent=2), flush=True)
result = unittest.TextTestRunner(verbosity=2).run(
    unittest.defaultTestLoader.loadTestsFromTestCase(case)
)
print(json.dumps({
    "case": args.case, "passed": result.wasSuccessful(),
    "tests_run": result.testsRun, "failures": len(result.failures),
    "errors": len(result.errors), "skipped": len(result.skipped),
}), flush=True)
raise SystemExit(not (result.wasSuccessful() and result.testsRun == 1 and not result.skipped))
PY

Online W4A16 GSM8K

cd "$TASK_VALIDATION_ROOT/sglang"
CUDA_VISIBLE_DEVICES=0,1,2,3 timeout --kill-after=30s 120m \
  python -B "$TASK_VALIDATION_ROOT/run_registered.py" online --port 31000

The unchanged class uses TP4/attention-DP4/EP4, BF16 model dtype, DP LM head, MegaMoE runner/A2A, nvfp4_online, decode graph maximum 128, disabled prefill graphs, and static memory fraction 0.7. It evaluates 200 GSM8K examples through the completion API with 128 threads, a 512-token limit, and score > 0.90.

The class sets these actual recipe/precision variables; the runner does not replace them:

FLASHINFER_NVFP4_4OVER6=1
FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE
FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=1
FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1
SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=0
SGLANG_FP4_IGNORED_LAYERS=
SGLANG_ENABLE_JIT_DEEPGEMM=False

The example server configuration in the PR body is separate: its graph/running caps are 64 and configured chunk is 16384. Those values are not substituted into this registered accuracy test.

Serialized W4A16 reload and local W4A4 control

Run sequentially:

cd "$TASK_VALIDATION_ROOT/sglang"
CUDA_VISIBLE_DEVICES=0,1,2,3 timeout --kill-after=30s 120m \
  python -B "$TASK_VALIDATION_ROOT/run_registered.py" reload-w4a16 --port 31010

CUDA_VISIBLE_DEVICES=0,1,2,3 timeout --kill-after=30s 120m \
  python -B "$TASK_VALIDATION_ROOT/run_registered.py" reload-w4a4 --port 31011

Both use serialized NVIDIA NVFP4 weights, BF16 model dtype, TP4/attention-DP4/EP4, --fp4-gemm-backend flashinfer_cutedsl, MegaMoE runner/A2A, --enable-memory-saver, and disabled prefill graphs. SGLANG_MEMORY_SAVER_CUDA_GRAPH=1 is set by the inherited launcher. In-kernel FC2 reduction and per-token activation quantization are both disabled. The W4A16 opt-in is 1 for the registered test and 0 for the local W4A4 control.

Both retain flush_cache=True and False with abort_all_requests=False, weight/KV/graph release-resume, exact text/token-ID comparisons, and log-probability tolerance 1e-4. Generation uses "The capital of France is", temperature 0, and at most 16 new tokens, explicitly routed to DP rank 0. These are same-checkpoint reload tests; they do not establish changed-checkpoint accuracy or all-rank output equality. A W4A4 checkpoint changing the captured input activation scale requires an engine restart.

No full Kimi K3/GLM-5.2 E2E, online RL reload, or speedup claim is included.

Speed Tests and Profiling

No controlled serving benchmark or speed comparison was run for this PR.
No speedup is claimed.

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 #38620506291
Latest PR Test (Extra): ❌ Run #38620505630
Latest PR Test (AMD ROCm 10): ❌ Run #38620505417

@github-actions github-actions Bot added documentation Improvements or additions to documentation quant LLM Quantization blackwell SM100/SM120 labels Oct 13, 2026
@zianglih
zianglih marked this pull request as ready for review October 13, 2026 22:02
@zianglih

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 13, 2026
@zianglih zianglih changed the title [FlashInfer] Add NVFP4 W4A16 and SiTU support to MegaMoE [Feat] FlashInfer: Add NVFP4 W4A16 and SiTU support to MegaMoE Oct 14, 2026
@ziang-and
ziang-and force-pushed the feat/flashinfer-megamoe-w4a16-followup branch 2 times, most recently from 571dab4 to e81c6bd Compare October 14, 2026 05:45
@ziang-and
ziang-and force-pushed the feat/flashinfer-megamoe-w4a16-followup branch from e81c6bd to 7337b0c Compare October 14, 2026 19:03
@zianglih

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

1 similar comment
@zianglih

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@Fridge003
Fridge003 merged commit 7098ec3 into sgl-project:main Oct 16, 2026
437 of 483 checks passed
@Fridge003 Fridge003 added the release-highlight Candidate PR for release note highlight label Oct 16, 2026
@zianglih
zianglih deleted the feat/flashinfer-megamoe-w4a16-followup branch October 16, 2026 07:45
ByronHsu added a commit to ByronHsu/sglang that referenced this pull request Oct 16, 2026
Resolve the SiTU backend allow-list in modelopt_quant.py: keep upstream's
FlashInfer MegaMoE (sgl-project#44036) and this PR's CuTe DSL v2 standard path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

blackwell SM100/SM120 documentation Improvements or additions to documentation quant LLM Quantization release-highlight Candidate PR for release note highlight run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants