[Justin Chu (justinchuby)]

Description

Add opt-in activation_quantization="dynamic_fp8" to the existing com.microsoft::QMoE v1 operator, supporting FLOAT, FLOAT16 and BFLOAT16 input/router/output tensors. The default "none" preserves weight-only execution.

• Keep FLOAT8E4M3FN expert weights and FP32 square-block scale grids unchanged in storage; reuse sparse routing and separate gate/up or fused SwiGLU layouts.
• Dynamically quantize consecutive K128 groups before gate/up projection and directly after FP32 SiLU/SwiGLU. Use saturating nearest-even E4M3 conversion, max(amax,1e-10) * float32(1/448) scales, and FP32 div.full normalization. Both arithmetic choices matter at FP8 midpoints; deterministic regressions cover the one-ULP differences from precise division.
• Execute native FP8 mma.sync.m16n8k32 instructions and apply each K128 group's scales in FP32. Gate/up outputs use the input type, with no extra FP16/BF16 round before the second quantizer. FLOAT gate/up stays FP32, including values above 65504. Down projection, routed products and top-k reduction remain FP32 through the final output cast. Routing coefficients round to the input type before multiplication.
• Restrict the new FLOAT CUDA registration to FP8 weights/FP32 scales, preserving CPU placement of existing FLOAT/integer models. FLOAT weight-only execution is not enabled. Dynamic execution bypasses unused A16 runner construction/prepacking rather than interpreting FLOAT as BF16.
• Add schema/type/shape checks, explicit unsupported-profile/EP errors, numerical-boundary and CUDA-graph tests, and public documentation. No dense expert-weight buffer is allocated.

Initial profile: FP32/FP16/BF16 inputs; FP32 weight scales; quant_type="fp8", expert_weight_bits=8, block_size=128; 128-aligned hidden/intermediate dimensions; standard SiLU/SwiGLU; native SM89+ kernel; at most 65,535 token/expert pairs per invocation. Biases, supplied activation/global scales, zero points and router_weights are rejected. CPU/WebGPU reject the mode explicitly. Requires the existing onnxruntime_USE_FP8_QMOE=ON build option; disabling fused execution is an error, not a W8A16 fallback.

Validation

Fresh isolated Linux Release build, CUDA 13.3 and H200 SM90. No performance, non-H200 hardware, or full-checkpoint qualification is claimed here.

Build command (from the source directory, with ROOT=/tmp/ort-drama3-w8a8-c6396368 and PY=$ROOT/venv/bin/python) — passed; the final FLOAT/normalization changes were incrementally rebuilt and the wheel reinstalled:

$PY tools/ci_build/build.py --config Release --update --build --build_wheel --skip_tests --parallel 16 --build_dir $ROOT/build --cmake_generator Ninja --use_cuda --cuda_home /usr/local/cuda --cudnn_home /usr/lib/python3.12/site-packages/nvidia/cudnn --skip_submodule_sync --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=90 onnxruntime_USE_FP8_QMOE=ON onnxruntime_USE_FP4_QMOE=OFF onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON onnxruntime_QUICK_BUILD=ON onnxruntime_USE_FLASH_ATTENTION=OFF onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION=OFF onnxruntime_USE_TRT_FUSED_ATTENTION=OFF onnxruntime_USE_XQA=OFF

Using TEST=$ROOT/src/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py, the following commands ran from $ROOT on the final plain-QMoE build:

$PY -m pytest -q "$TEST" -k 'not integer_scales_still_require_activation_type and not fused_silu_accumulator_precision'
# 85 passed, 6 skipped, 5 deselected (baseline failures below)
ORT_RUN_LARGE_FP8_QMOE_TEST=1 $PY -m pytest -q "$TEST" -k dynamic_fp8_drama3_projection_shapes
# 3 passed: FLOAT/FLOAT16/BFLOAT16, E=256, H=2048, I=512, top-k=8

The skipped cases include opt-in large shapes (run separately above) and intentionally unsupported FLOAT/weight-only CUDA graph combinations. The suite also verifies that FLOAT/integer QMoE remains on CPU with CUDA first in the provider list.

From the build output directory:

./onnxruntime_provider_test --gtest_filter='QMoEShapeInferenceTest.*'
# 13 passed, including FLOAT shape/type cases
GTEST_FILTER='CUDA_EP_Unittest.QMoEDynamicFp8*' ./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All
# 4 internal CUDA tests passed, inside the outer loader test

python -m ruff check onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py, python -m ruff format --check onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py, and git diff --check — passed. Changed C++ lines were formatted with clang-format 18. All 17 public QMoE attribute descriptions match the freshly built schema.

A separate six-case comparison against the pinned Fish Speech FP8 implementation (541566932abbaebcb0a40e562a1f846738383f38) passed decode (1 token) and prefill (33 tokens) for all three dtypes at the full expert dimensions above with seeded synthetic weights. FLOAT used strict atol=2e-6, rtol=2e-5; maximum absolute error was 8.94e-8 for decode and 2.384e-7 for prefill. FLOAT16 maximum error was 4.883e-4; BFLOAT16 was bit-identical for decode and had maximum error 1.221e-4 for prefill. This verifies operator/source semantics, not full-model or checkpoint fidelity. No assertions were loosened to qualify FLOAT; its normalization midpoint discrepancy was fixed and regression-tested.

Baseline qualification: the broader initial run had five failures. Recompiling baseline HEAD's QMoE host implementation and relinking against the identical dense kernels reproduced all five: four existing dense-fallback SiLU precision/overflow comparisons and one compact-build BF16 INT8 runner availability limitation. These unchanged paths remain excluded from the scoped regression command above, rather than hiding their failures behind a passing total.

Motivation and Context

Existing block-FP8 QMoE preserves weight storage but uses FP16/BF16 activations. It does not reproduce models that dynamically quantize activations before both expert projections. Extend the general operator rather than introduce a model-specific op, dequantize complete expert banks, approximate W8A8 using A16 GEMM, or change the existing W8A16 default. FLOAT input/output support additionally permits FP32 residual/tower streams without rounding gate/up intermediates to a low dtype.

[Copilot]

🟡 Changes recommended

The quantizer has unresolved launch-dimension and NaN-preservation defects.

2 open issues

• Quantizer converts NaN activations to finite FP8 values · New
• Quantizer grid.y overflows for large inter_size · New

What changed in this PR

Adds opt-in dynamic FP8 activation quantization to com.microsoft::QMoE v1, enabling native CUDA W8A8 execution while preserving the default weight-only path.

Changes:

• Adds dynamic quantization, native FP8 projections, and FP32 routed reduction without dense expert-weight buffers.
• Validates the supported profile and explicitly rejects unsupported CPU/WebGPU execution.
• Adds numerical, shape-inference, and CUDA-graph tests, plus public documentation.

File	Description
onnxruntime/​test/​python/​transformers/​test_qmoe_fp8_cuda.py	Adds dynamic FP8 reference and integration tests.
onnxruntime/​test/​contrib_ops/​qmoe_shape_inference_test.cc	Tests supported shapes and invalid metadata.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​qmoe_fp8_compaction_test.cc	Tests quantization boundaries and routing precision.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Defines the new attribute and profile validation.
onnxruntime/​contrib_ops/​webgpu/​moe/​qmoe.h	Rejects dynamic FP8 on WebGPU.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Declares dynamic FP8 kernel interfaces.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Implements quantization, native projections, and reduction.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.h	Tracks the dynamic FP8 execution mode.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Integrates validation, scratch allocation, and dispatch.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_quantization_cpu.cc	Rejects dynamic FP8 on CPU.
docs/​ContribOperators.md	Documents the new QMoE attribute and semantics.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Explains numerical behavior and CUDA requirements.

_{Review effort: Balanced}