[ruoqianguo]

Superseded by the upstream follow-up: flashinfer-ai/flashinfer#7542, based on main after flashinfer-ai#6999 merged. This fork PR is kept as the original review record.

---

📌 Description

This PR is stacked on #6999. For the measured N64/K256 configurations, NVFP4 has three AB stages while MXFP4 with compact buffers has four. NVFP4 stores twice as many weight scales (K16 versus K32). Moving scales into a two-slot ring (nvfp4_sf_stages=2) frees space for a fourth AB stage.

Sf2 candidates must satisfy supported geometry, protocol and resource constraints, without GPU-name or token-count checks. The kernel rejects incompatible layouts. Only T512/T1024 defaults change to sf2/R1; the end-to-end test verifies byte-identical output against PR1 sf0.

Performance and reproduction

The table compares sf2 with PR1 sf0 defaults and MX main defaults with manual compact=True. Negative deltas mean sf2 is faster.

T/rank	PR1 sf0 µs	sf2 µs	MX µs	Paired sf2/sf0 Δ	Paired sf2/MX Δ
512	1232.6	1189.6	1192.8	-3.68%	-0.38%
1024	1607.2	1566.6	1615.7	-2.75%	-3.00%

We measured MegaMoE plus top-k reduction on 4×H200 with EP4/H7168/I3072/E384/top-6/clamp10. Each block takes the maximum of four rank medians. Paired ratios are summarized across three rounds and three fresh processes, without dividing independent latency medians. SM clocks were set to 1830 MHz. Software power-cap events were recorded, so a constant clock in every timed block is not guaranteed.

All nine T512 rounds were retained; fresh-process sf2/sf0 medians ranged from -4.83% to -3.16%.

To reproduce PR1, export its defaults first. Setting only nvfp4_sf_stages=0 is insufficient because sf2 entries also change pull-buffer and LUT settings:

python - <<'PYEXPORT'
import json
import runpy

data = runpy.run_path(
    "flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/nvfp4_candidate_data.py"
)
profiles = {
    str(tokens): {
        "tactic": dict(data["NVFP4_SF0_CANDIDATE_SEEDS"][spec[0]]),
        "config": {
            "nvfp4_kperm": 64,
            "nvfp4_fc1_k256_overlap": spec[-1],
        },
    }
    for tokens, spec in data["_DEFAULTS"].items()
}
with open("nvfp4-pr1-sf0-defaults.json", "x") as output:
    json.dump({"per_token": profiles}, output, indent=2)
    output.write("\n")
PYEXPORT

nv_bench_root="$(mktemp -d -p . sm90-pr2-current-XXXXXX)"
for fresh in 1 2 3; do
  FLASHINFER_MOE_EP_KNOB_CACHE="$(pwd)/${nv_bench_root}/knob-${fresh}.json" \
  torchrun --standalone --nproc_per_node=4 benchmarks/bench_moe_ep_sm90_nvfp4.py \
    --tokens 512 1024 --hidden 7168 --intermediate 3072 --experts 384 --top-k 6 \
    --seed 1234 --nv-kperm 64 --nv-overlap auto --mx-compact \
    --nv-reference-knobs-json nvfp4-pr1-sf0-defaults.json \
    --rounds 3 --warmup 10 --iters 50 --padding-us 300 --cooldown-seconds 5 \
    --output "${nv_bench_root}/fresh-${fresh}.json" || exit $?
done

🔍 Related Issues

Base: flashinfer-ai#6999. Separate CUDA implementation: #6343.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

The existing NCCL-EP high-throughput test (tests/moe_ep/test_moe_ep_cudagraph_multirank.py, high_throughput) also fails on unmodified main (eed63cd) in our 4×H200 environment; it is unrelated to this PR.

Default MXFP4/FP8/BF16 cubins match tested main byte-for-byte. The four-format EP4 end-to-end, host/unit and registration tests passed on 4×H200.

EP4 regression reproduction:

for dtype in nvfp4 mxfp4 fp8 bf16; do
  torchrun --standalone --nproc_per_node=4 -m pytest \
    "tests/moe_ep/test_moe_ep_sm90_pull_${dtype}_mega_multirank.py" \
    -vv -ra -rs --tb=short --full -m arch_hopper || exit $?
done

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

Review scale-slot ownership, cross-tile state and memory accounting. This PR is stacked on flashinfer-ai/flashinfer#6999 and will be submitted upstream after PR1 is merged.