[FlashInfer] Add startup autotuning for MegaMoE

[zianglih]

Motivation

@HumansAnd

Add startup autotuning for the existing FlashInfer MegaMoE NVFP4 W4A4 and MXFP8
backends. The generic FlashInfer autotune context does not trigger MegaMoE's
native tuner. Tune representative decode/prefill capacities before CUDA graph
capture, then select a fixed workspace using synchronized DP token counts.

This is the standalone autotune prerequisite for
#39210. W4A16 integration stays
in that separate PR. Uses the FlashInfer 0.7.1 dependency already on main.

Modifications

• Startup: participate in the existing target-model autotune lifecycle by
  default. Prepare powers-of-two decode capacities plus the exact decode ceiling
  and one prefill ceiling. Decode bounds cover captured graphs, eager request
  limits, and maximum speculative width. Preparing the native prefill profile
  does not enable SGLANG_FLASHINFER_AUTOTUNE_EXTEND.
• Representative inputs: repeat or slice the startup batch to each capacity;
  synthesize inputs for idle ranks. Native knobs="auto" otherwise measures its
  first live input while caching under workspace capacity. Tune temporary owners,
  then destroy them and warm fixed public workspace handles before capture.
• Runtime selection: use the largest synchronized DP count to choose the
  smallest covering profile on every EP rank. Retain local expert inputs and
  existing DP bookkeeping; reject capacity overflow. No tuning or cache-file
  access occurs in ordinary forwards.
• Workspace ownership: use FlashInfer's reference-counted pool instead of
  assigning private workspaces across layers. Stage existing per-layer NVFP4
  scales on every forward; use the public output-view capability and released
  enable_in_kernel_fc2_reduce keyword. The FC2 flag permits atomic candidates;
  the selected tactic can still use external reduction.
• Persistence: store native tactic payloads in the existing FlashInfer
  namespaced cache. Stage each geometry/capacity once per startup context and
  retain process-local winners for already-pooled workspaces. Target MegaMoE
  cache filenames include physical GPU UUIDs; engines sharing physical GPUs
  need separate SGLANG_CACHE_DIR roots. Cache-disabled runs use unique filenames.
• Boundaries: native tuning requires EP to span WORLD because FlashInfer
  0.7.1's tuner uses WORLD collectives. Draft profiles are not added. Existing
  disable, deterministic-mode, and skip-op controls preserve normal forwarding;
  DeepGEMM keeps its existing forward path. No W4A16/SiTU, quantization/loader,
  model-specific, NVSHMEM, Dockerfile, or dependency-pin changes.

Use --flashinfer-autotune-skip-ops flashinfer_megamoe to skip just this tuner,
or --disable-flashinfer-autotune to disable startup FlashInfer tuning.

Accuracy Tests

Environment and source

• Tested source: 0b0ff32ee67e9218fe4d5867a1732436adcef6b8, based on upstream
  208df292e541246c4b88a11131474c804119047b. The pinned checkout, clean tracked
  files, and source import paths were verified before each unit/native run.
• Image: lmsysorg/sglang:nightly-dev-cu13-20261012-e6072935; index digest
  sha256:b5c488694d74c13587af6057b90220498c4869571c8e1e3c43dd6f43ae721306,
  amd64 manifest digest
  sha256:e20a28a4badee7d3de71881e7c2c160c7079465bf1ac21edec31a99c2602507d.
• Hardware: four NVIDIA B300 SXM6 AC GPUs per native/serving case, with
  275040 MiB reported per GPU; driver 590.48.01. Native runs used physical
  GPUs 0–3 while independent correctness validation used GPUs 4–7. Timing from
  these runs is not performance evidence.
• Runtime: Python 3.12.15, PyTorch 2.14.1+cu130 / CUDA 13.0,
  flashinfer-python==0.7.1, flashinfer-cubin==0.7.1,
  flashinfer-jit-cache==0.7.1+cu130, CUTLASS DSL and its CUDA 13 libs 4.8.0,
  nvidia-nvshmem-cu13==3.4.5, nvshmem4py-cu13==0.4.0.
• FlashInfer installation: installed the official released binary wheels,
  including the CUDA 13 architecture providers. Build metadata identifies
  release commit 0b46115d238d29ea2e92c6b9b2cb077d40fbc352. No FlashInfer source
  build or editable installation. All non-FlashInfer Python distributions and the image's
  sgl_kernel shared-library hashes were unchanged. CuTe DSL kernels specialize
  normally at runtime.
• Source-checkout setup: the image sets SGLANG_RUST_BUILD_MODE=never, but
  its bundled Rust TreeCore binary targets newer source than this PR. Build the
  pinned SGLang Rust extension with Rust/Cargo 1.92.0 into the
  source-fingerprinted task cache before
  launching the server, then keep normal server mode never. This does not
  rebuild FlashInfer or alter installed packages.
• NVSHMEM setup: the single-node runs explicitly use
  NVSHMEM_REMOTE_TRANSPORT=none, NVSHMEM_IB_ENABLE_IBGDA=0, and
  NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1. The native harness bypasses ServerArgs;
  this PR does not change NVSHMEM defaults. The local-only proxy issue is
  addressed by NVIDIA/nvshmem@af9e239,
  with context in NVIDIA/nvshmem#110.

Released-wheel and source setup

Inside the image, with the pinned SGLang checkout under $TASK_ROOT/sglang
and the retained standalone validation helpers under $TASK_ROOT/c2-tests:

export TASK_ROOT="$PWD"
export PYTHONPATH="$TASK_ROOT/sglang/python"
export SGLANG_CACHE_DIR="$TASK_ROOT/cache/sglang"
export FLASHINFER_WORKSPACE_BASE="$TASK_ROOT/cache/flashinfer"
export CUTE_DSL_CACHE_DIR="$TASK_ROOT/cache/cute-dsl"
export TORCH_EXTENSIONS_DIR="$TASK_ROOT/cache/torch-extensions"
export PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 OMP_NUM_THREADS=1 MAX_JOBS=4
mkdir -p "$TASK_ROOT/results" "$TASK_ROOT/setup"
test "$(git -C "$TASK_ROOT/sglang" rev-parse HEAD)" = \
  0b0ff32ee67e9218fe4d5867a1732436adcef6b8

cat > "$TASK_ROOT/setup/wheels.txt" <<'WHEELS'
https://files.pythonhosted.org/packages/32/3f/3c8dbde51499902346f8f38fcdc3f92db3de25b1cdd11361f2f4dd90fd83/flashinfer_python-0.7.1-py3-none-any.whl#sha256=69302ed6e91314a85a4392324becc194594a782ba181646d25d5c305e1b7e8c4
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_cubin-0.7.1-py3-none-any.whl#sha256=502eaa5f71e768950ada429618481a886388fcb34a5da6a2258979713cbaa327
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=75c80222846c11f269ab72c96c1152ee675836deadf84d805e44ac62da411beb
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm80-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=1788bc850e1b301b06a2c36d67ca333fe3a95e8cec099c0915871db1d72c7626
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm89-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=5d3673347a7c2b6fb20f758129f215ddc57038df741806b1d07b32aa774bda91
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm90a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=d8bf578c40a236b3f07216302cc7c7ad410ea8c3f12052fcfff8791876b0b1c0
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm100a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=c7dbb9ceede5763448d23ca744ce1ccaaba377cb86e4b15304f9a411fc754808
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm103a-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=79c476784f6641afeecd9289ddd1651b75f868c173f99100d0bb8ba3441d33da
https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.1/flashinfer_jit_cache_sm120f-0.7.1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl#sha256=596424325c5c34baa9881149e2ef267428a8c05d28659081773500da7b79d263
WHEELS
python -m pip download --no-deps --only-binary=:all: --require-hashes \
  --dest "$TASK_ROOT/setup/downloaded-wheels" -r "$TASK_ROOT/setup/wheels.txt"
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 -m pip install --no-deps --only-binary=:all: --require-hashes \
  -r "$TASK_ROOT/setup/wheels.txt"

Before serving from the pinned checkout, prepare its Rust TreeCore cache:

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__)'
CUDA_VISIBLE_DEVICES="" SGLANG_RUST_BUILD_MODE=never \
python -B -c 'from sglang.srt.mem_cache.rust_tree_core.extension import bindings; print(bindings.__file__)'

The loader uses locked Cargo dependencies and a source/toolchain/PyTorch
fingerprint; server launches reuse this artifact with mode never. The initial
serving attempt reached the end of MegaMoE tuning but failed while loading this
unrelated missing extension. Its complete log is retained separately from the
repaired-environment retry.

Source imports use PYTHONPATH; remaining Python/CUDA packages come from the image.
Separate source, FlashInfer, CuTe DSL, PyTorch extension, and tuning cache roots
were used for this PR. Only the verified model download cache was shared.

Registered unit tests

The existing registered suites run in the full image without the local CPU
bootstrap harness. Their own test fixtures/mocks remain in place. This covers
15 adapter/profile cases, runner/cache startup, the legacy-accessor source rule,
and DP scheduler metadata; native CUDA execution is covered separately below.

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/test_runtime_context.py::TestTheAccessorsHaveNoCallersOutsideTheirPackage \
  test/registered/unit/managers/scheduler_components/test_dp_attn.py::TestDPAttnSchedulerMetadata \
  --junitxml="$TASK_ROOT/results/units.xml"

Raw footer:

33 passed, 21 warnings, 7 subtests passed in 181.16s (0:03:01)

JUnit records 40 entries including the seven subtests, with zero failures,
errors, or skips. Full warning output and JUnit are retained with the raw logs.

Native W4A4 and MXFP8 autotune/capture validation

• Fixture: EP4/WORLD, synthetic SwiGLU, hidden/intermediate size 256, eight
  experts, top-k 4, two layer owners and separate startup contexts. Both paths
  use their actual SGLang weight preprocessing and the released native kernels.
• Profiles: bootstrap 4, decode ceiling 16, prefill ceiling 8; capacities
  [1, 2, 4, 8, 16]. Synchronized DP counts are supplied as
  [0, 1, 1, 2], [0, 1, 2, 7], [8, 8, 8, 8], and [0, 4, 9, 16].
  These include unequal and idle ranks; full scheduler behavior is outside
  this fixture.
• Oracle: compare against independent fixed-default native MegaMoE sessions
  (knobs={}, capacity 16), with immediate clones of pooled output views for
  both owners. This is the same released FlashInfer kernel family, not a BF16
  or PyTorch mathematical oracle. Both results must be finite and bitwise
  equal as BF16; no tolerance is applied.
• Capture: two unchanged CUDA graph replays and one replay after mutating
  activations, plus FC1 scales for W4A4, for each count fixture. Assert shared
  workspace identity and reference counts 2 → 1 → 0 during owner teardown.
• Persistence: the native AUTO wrapper only counts calls. All candidates,
  kernels, tuning, and dispatch remain unchanged. Cold startup invokes AUTO
  five times per rank; warm restart invokes it zero times. Both owners' eight
  cache files per precision contain identical namespaced records for all five
  capacities.

The retained smoke_autotune_gpu.py is a standalone validation helper, outside
this PR's source diff. The full native commands are:

export CUDA_VISIBLE_DEVICES=0,1,2,3
export NVSHMEM_REMOTE_TRANSPORT=none
export NVSHMEM_IB_ENABLE_IBGDA=0
export NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1
export SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=0
export SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0
export SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE=bf16
export SGLANG_FLASHINFER_MEGAMOE_MAX_TOKENS_PER_RANK=4
export FLASHINFER_MOE_EP_KNOB_CACHE=0

for precision in w4a4 mxfp8; 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/c2-tests/smoke_autotune_gpu.py" \
      --cache "$TASK_ROOT/results/native-$precision-cache" --reuse-cache "$reuse" \
      --decode-num-tokens 16 --precision "$precision"
  done
done

Complete per-rank verdicts (cold and warm processes all exited 0):

Precision	Cache	Rank	AUTO calls	Comparisons	Graph replays	Max abs error	Bit exact	Pool refs
w4a4	cold	0	5	24	12	0.0	True	2→1→0
w4a4	cold	1	5	24	12	0.0	True	2→1→0
w4a4	cold	2	5	24	12	0.0	True	2→1→0
w4a4	cold	3	5	24	12	0.0	True	2→1→0
w4a4	warm	0	0	24	12	0.0	True	2→1→0
w4a4	warm	1	0	24	12	0.0	True	2→1→0
w4a4	warm	2	0	24	12	0.0	True	2→1→0
w4a4	warm	3	0	24	12	0.0	True	2→1→0
mxfp8	cold	0	5	24	12	0.0	True	2→1→0
mxfp8	cold	1	5	24	12	0.0	True	2→1→0
mxfp8	cold	2	5	24	12	0.0	True	2→1→0
mxfp8	cold	3	5	24	12	0.0	True	2→1→0
mxfp8	warm	0	0	24	12	0.0	True	2→1→0
mxfp8	warm	1	0	24	12	0.0	True	2→1→0
mxfp8	warm	2	0	24	12	0.0	True	2→1→0
mxfp8	warm	3	0	24	12	0.0	True	2→1→0

All 384 comparisons were finite and bitwise exact, with maximum absolute error
0.0: 312 compared nonempty tensors and 72 exercised idle-rank empty outputs. The 16 rank/run combinations each performed 12 graph replays. Full stdout,
per-comparison records, selected native knobs, cache payloads, command receipts,
source/wheel provenance, exit codes, and UTC timestamps are retained locally
in c2-20261012/remote-results/ and NATIVE_VALIDATION.json.

W4A4 serving accuracy

Model: nvidia/Qwen3-30B-A3B-NVFP4, pinned Hugging Face revision
2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3. TP4 / attention DP4 / EP4;
max running requests 64; decode CUDA graphs up to 64; prefill CUDA graphs
disabled; chunked prefill 4096; static memory fraction 0.7; BF16 combine;
per-token activation quantization and in-kernel FC2 reduction disabled.
Startup FlashInfer tuning remains enabled by default. No speculative decoding.

Full server command (the test helper resolves this public model to the pinned
local snapshot):

export MODEL_PATH="$(hf download nvidia/Qwen3-30B-A3B-NVFP4 \
  --revision 2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3)"
unset SGLANG_FLASHINFER_MEGAMOE_MAX_TOKENS_PER_RANK FLASHINFER_MOE_EP_KNOB_CACHE
CUDA_VISIBLE_DEVICES=0,1,2,3 \
SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=0 \
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=0 \
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0 \
SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE=bf16 \
SGLANG_ENABLE_JIT_DEEPGEMM=False \
NVSHMEM_REMOTE_TRANSPORT=none \
NVSHMEM_IB_ENABLE_IBGDA=0 \
NVSHMEM_DISABLE_LOCAL_ONLY_PROXY=1 \
sglang serve \
  --model-path "$MODEL_PATH" \
  --dtype bfloat16 \
  --tp-size 4 --attn-dp-size 4 --ep-size 4 \
  --moe-runner-backend flashinfer_megamoe \
  --moe-a2a-backend flashinfer_megamoe \
  --cuda-graph-backend-prefill disabled \
  --cuda-graph-max-bs-decode 64 \
  --max-running-requests 64 \
  --chunked-prefill-size 4096 \
  --mem-fraction-static 0.7 \
  --device cuda --host 127.0.0.1 --port 31100

After server readiness, run 200 GSM8K examples with 128 client threads,
completion API, a 512-token output cap, and the script defaults
(temperature 0, top-p 1). The test requires score strictly greater than 0.90.

python -m sglang.test.run_eval \
  --base-url http://127.0.0.1:31100 \
  --model "$MODEL_PATH" \
  --eval-name gsm8k --num-examples 200 --num-threads 128 \
  --api completion --max-tokens 512

Result: passed, 92.5% (185/200), with 200 logged HTTP 200 completion
responses and no traceback. Both the test and its supervisor exited 0. Raw
result (latency/throughput are incidental evaluator fields, not benchmark evidence):

{
  "score:std": 0.26339134382131846,
  "score": 0.925,
  "latency": 24.68711337633431,
  "output_throughput": 1023.2869114708564
}

The effective per-rank settings were 16 running requests, 1024-token chunked
prefill, and captured decode batches [1, 2, 4, 8, 12, 16]; the requested limits
above produced these DP-adjusted runtime limits. The checkpoint selected FP8
E4M3 KV cache; context
length was 40960. Startup reused native profiles [1, 2, 4, 8, 16, 1024] from
the first attempt's completed tuning: six cache-hit entries and no new native
tuning entries in the successful retry. No source or package edits were needed.

Supplemental local checks and remaining scope

• CPU checks on macOS 27.0.1 arm64 / Python 3.12.12 / PyTorch 2.14.0: adapter
  harness 15 passed; cache harness Ran 16 tests ... OK (skipped=1) including
  four local cache-isolation cases; AST-scoped unchanged legacy-accessor check
  Ran 1 test ... OK. The CPU harnesses stub CUDA/native interfaces; Gloo,
  process spawning, runtime-context overrides, and cache-file operations are
  real. Their CUDA-only skip is separate from the zero-skip image result above.
• All eight changed Python files parse; git diff --check and applicable
  pre-commit checks pass. Mintlify site checks were unavailable locally because
  mint is not installed; the MDX changes update existing table text.
• MXFP8 full-model serving, larger native shapes, other activations, other
  combine/reduction modes, and BF16-oracle numerics are not validated by these
  native fixtures. W4A16 remains in its separate PR.

git diff --check 208df292e541246c4b88a11131474c804119047b...HEAD
git diff --name-only 208df292e541246c4b88a11131474c804119047b...HEAD | xargs pre-commit run --files

Speed Tests and Profiling

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

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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

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 #38217773506
Latest PR Test (Extra): ❌ Run #38217773092
Latest PR Test (AMD ROCm 10): ❌ Run #38217773465

[zianglih]

/tag-and-rerun-ci

[nvpohanh]

cc @leejnau @wenscarl