diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 36b68889ae582..9ae17df7e5114 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -24,6 +24,7 @@ and have been significantly modified for ONNX Runtime — see 10. [FP8 (W8A16) Details](#10-fp8-w8a16-details) 11. [WFP4AFP8 Details](#11-wfp4afp8-details) 12. [Future / Deferred Modes](#12-future--deferred-modes) + - [12.1 MoE GEMV Optimization Summary](#121-moe-gemv-optimization-summary) 13. [Testing](#13-testing) 14. [Build Configuration](#14-build-configuration) 15. [Limitations & Known Issues](#15-limitations--known-issues) @@ -175,16 +176,42 @@ values that require them are rejected at construction time: ## 4. Architecture Dispatch & Kernel Paths -The runner selects between three CUTLASS kernel families at runtime. The choice is +The runner selects between three CUTLASS kernel families and one small-row GEMV +fast path at runtime. The choice is made by `CutlassMoeFCRunner::supportsTmaWarpSpecialized()` and the dispatch headers under [onnxruntime/contrib_ops/cuda/llm/moe_gemm/](onnxruntime/contrib_ops/cuda/llm/moe_gemm/). | Path | CUTLASS class | Used for | SM range | |------|---------------|----------|----------| +| **MoE GEMV fast path** | `fpA_intB_gemv`-based custom kernel | INT4/INT8 per-column W*A16 and symmetric INT4/INT8 block-wise W*A16 with FP16 or BF16 activations and true decode row counts | SM80+ | | **Ampere GemmGrouped** | `cutlass::gemm::kernel::GemmGrouped` | INT4/INT8 W*A16, FP8 W8A16 dequant fallback, FP32 | SM75–SM89, plus all mixed-input on SM90/SM120 | | **TMA Warp-Specialized (mixed-input)** | `CollectiveBuilderMixedInput` | Same-type FP16×FP16 / BF16×BF16, native MXFP4 W4A16 | SM90 (same-type), SM120 (FP4 W4A16) | | **Block-Scaled Tensor Op** | `OpClassBlockScaledTensorOp` | Native FP8×MXFP4 (`wfp4afp8`) | SM100+ (Blackwell) | +The MoE GEMV fast path is selected before the Ampere grouped GEMM for integer +QMoE when all of the following are true: + +- activation/output dtype is FP16 or BF16; +- scales and biases use the same dtype as the activation; +- weights are INT4 or INT8 for per-column scales, or INT4/INT8 for symmetric block-wise scales; +- `block_size <= 0` for per-column INT4/INT8, or `block_size` is 32, 64, or 128 for block-wise INT4/INT8; +- block-wise GEMV is symmetric only; if zero-point compensation is present, dispatch falls back to grouped GEMM; +- `expanded_num_rows = num_tokens * top_k` is in `(0, 8]`; +- `N >= 512` and `K >= 512`; +- if `expanded_num_rows > 4`, the logical MoE intermediate size is at least 512; +- `N` is divisible by the column-interleaved tile width (32 for INT4, 16 for INT8), + and `K` satisfies the kernel step and, for block-wise scales, complete-block alignment. + +Asymmetric block-wise quantization, broader row counts, and dimensions +outside the profiled gate stay on grouped GEMM until profile data shows an +end-to-end GEMV win. FP16 and BF16 share the same dispatch gate and custom +kernels; for a given shape, BF16 routes to GEMV exactly where FP16 does and +shows comparable latency. + +Set `ORT_DISABLE_MOE_GEMV=1` before process start to force the grouped GEMM +fallback for debugging, benchmarking, or bisecting numerical differences. The +switch is cached on first use. + ### 4.1 Per-mode dispatch matrix | Mode | SM75-89 (Ampere/Ada) | SM90 (Hopper) | SM100 (Blackwell) | SM120 (RTX 5090) | @@ -338,6 +365,10 @@ This is the layout transform applied either offline by `weights_prepacked=0` (see [§5.1](#51-weights-input-2--5--8)). +INT MoE/QMoE CUDA kernels, including the small-row MoE GEMV path, consume the +SM80 `ColumnMajorTileInterleave<64, 4>` layout. Pack INT4/INT8 MoE weights with +the SM80 target layout even when the runtime GPU is Hopper or newer. + 1. **Input layout**: `[N, K]` per expert (Out × In), 2 elements per byte for INT4. 2. **Transpose & signed conversion**: - Unpack `uint4 [0, 15]` → subtract 8 → `int8 [-8, 7]`. @@ -424,9 +455,9 @@ Dequantization (symmetric): `W = (W_stored - 128) * scale`. | Architecture | Activation | Supported `block_size` | |--------------|-----------|------------------------| -| SM75–89 (Turing/Ampere/Ada) | FP16/BF16 | 64, 128 | -| SM90 (Hopper) | FP16/BF16 | any multiple of 64 | -| SM100/120 (Blackwell) | FP16/BF16 | falls back to Ampere — 64 or 128 | +| SM75–89 (Turing/Ampere/Ada) | FP16/BF16 | 32, 64, 128 | +| SM90 (Hopper) | FP16/BF16 | falls back to Ampere — 32, 64, 128 | +| SM100/120 (Blackwell) | FP16/BF16 | falls back to Ampere — 32, 64, 128 | For MXFP4, the block size is fixed at **32** by the format. @@ -882,6 +913,110 @@ software dequant of the mixed-input path. The schema reserves the necessary input slots (18–21) so adding these modes will not change the operator interface. +### 12.1 MoE GEMV Optimization Summary + +The MoE GEMV work targets decode-sized integer QMoE workloads where grouped GEMM +launch overhead, prologue overhead, and intermediate traffic dominate the FC +compute. Detailed measurements are recorded in +[qmoe_gemv_experiments.md](qmoe_gemv_experiments.md). This section is the +implementation summary and current backlog. + +#### Completed per-column INT4 work + +Per-column here means the original INT4 W4A16 path with one scale per output +column: scale tensors are `[E, N]`, and `block_size <= 0`. + +| Area | What is implemented | Result | +|------|---------------------|--------| +| Benchmarking | Added `profile_qmoe_gemv.py` and `profile_qmoe_gemv.sh` to run GEMV-enabled and `ORT_DISABLE_MOE_GEMV=1` grouped-GEMM profiles in separate processes. | Stable A/B profiles with the `benchmark` NVTX range; use `parse_nsys.py --pattern '%'` so fallback CUTLASS kernels are visible. | +| Route policy | Dispatch is data-gated to FP16/BF16 integer QMoE decode shapes, `expanded_num_rows <= 8`, `N >= 512`, `K >= 512`, and profiled alignment constraints. | Keeps tiny shapes and unprofiled row counts on grouped GEMM. | +| Row-to-expert lookup | The prologue materializes the local expert id for each permuted row and passes it to the GEMV kernels. | Removes repeated prefix-offset scans inside each N-tile CTA; GPT-OSS and Gemma model-shape kernels improved. | +| FC1 interleaved SwiGLU | The FC1 GEMV path can apply interleaved SwiGLU in the GEMV epilogue for the profiled FP16/BF16 INT4 path. | Removes the separate activation launch and FC1 intermediate traffic; GPT-OSS and Gemma improved end-to-end. | +| One-row finalize | `num_rows == 1`, `top_k <= 4` has a static top-k finalize specialization. | Modest GPT-OSS finalize improvement while preserving the existing FC2 GEMV parallelism. | +| End-to-end GPT-OSS | The final per-column GEMV path was validated in ORT GenAI on GPT-OSS-20B INT4. | About 15% faster than the grouped-GEMM baseline and about 8% faster than the FasterTransformer reference at batch 1. | + +#### Completed per-column INT8 work + +Per-column INT8 means W8A16 with one symmetric scale per output column: scale +tensors are `[E, N]` and `block_size <= 0`. The per-column path previously +required `block_size == 0` exactly in `is_moe_gemv_supported`, but the QMoE +runtime carries per-column scales as `group_size = -1` (the `QuantParams::Int` +default), so per-column INT4 *and* INT8 silently fell back to grouped GEMM. The +gate now treats any `group_size <= 0` as the per-column case, matching the GEMV +launcher dispatch that already mapped `group_size <= 0` to the `GroupSize == 0` +kernel. + +| Area | What is implemented | Result | +|------|---------------------|--------| +| Gate fix | `is_moe_gemv_supported` accepts `group_size <= 0` (per-column) alongside 32/64/128 block-wise group sizes. | Per-column INT4 and INT8 now reach the GEMV kernels, and block-wise INT4/INT8 includes `block_size=32`. | +| INT8 details | The existing `(half, uint8_t)` and `(__nv_bfloat16, uint8_t)` GEMV kernel details cover per-column INT8 with no new instantiation. | FC1 interleaved-SwiGLU and FC2 per-column INT8 GEMV run for FP16 and BF16. | +| Profiling | `int8_per_column_*_1024x4096_e8` and `gpt_oss_20b_*_int8_2880x2880_e32` cases profiled in GEMV and `ORT_DISABLE_MOE_GEMV=1` modes. | Real `moe_gemv_kernel` / `moe_gemv_interleaved_swiglu_kernel` confirmed; about 1.2x–1.4x lower benchmark latency than grouped GEMM with valid output. | + +#### Completed block-wise INT4/INT8 work + + +Block-wise here means `quant_type="int"` with `block_size` 32, 64, or 128 and scales +provided as `[E, N, K / block_size]`. QMoE prepack/runtime transposes those +scales to `[E, K_blocks, N]`, and the GEMV kernels consume that same layout. + +| Area | What is implemented | Result | +|------|---------------------|--------| +| INT4 and INT8 details | The GEMV kernel details support `(half, cutlass::uint4b_t)`, `(half, uint8_t)`, and the matching `__nv_bfloat16` weight-type pairs. | Both weight types use the SM80 column-interleaved `fpA_intB` layout on all GPUs, matching the grouped-GEMM path, for FP16 and BF16 activations. | +| Group-size dispatch | GEMV templates now cover `GroupSize == 0`, 32, 64, and 128. | Per-column and block-wise paths share the same kernel structure while preserving complete-block checks. | +| Scale indexing | Block-wise scale loads use `real_offset_k / GroupSize * n + real_offset_n` with a K-loop scale step. | Reuses QMoE's existing `[E, K_blocks, N]` runtime scale layout; no new scale pack format is needed. | +| Symmetric-only gate | Block-wise GEMV runs only when zero-point compensation is absent. | Asymmetric block-wise models stay on grouped GEMM until a zero-point GEMV path is implemented and profiled. | +| Model-shape b64 profile | GPT-OSS-20B and Qwen3.6-35B-A3B were profiled with `block_size=64`. | Both use real GEMV kernels under the current 512 threshold and show about 1.4x lower benchmark latency than grouped GEMM fallback. | + +#### Completed BF16 enablement + +BF16 activations now share the exact dispatch gate and custom GEMV kernels with +FP16. The runtime gate relaxes from `T == half` to `T == half || T == +__nv_bfloat16`, and `__nv_bfloat16` template instantiations were added for the +per-column INT4, block-wise INT4/INT8, and interleaved-SwiGLU GEMV kernels. + +| Area | What is implemented | Result | +|------|---------------------|--------| +| Gate relaxation | `tryLaunchMoeGemvIntSymmetric` and the interleaved-SwiGLU variant accept `__nv_bfloat16` activations with `ScaleBiasType == T`. | For a given shape, BF16 routes to GEMV exactly where FP16 does. | +| Kernel instantiation | `moe_gemv.cu` adds `__nv_bfloat16` details/instantiations (group sizes 0/32/64/128, INT4/INT8, bias on/off) under `ENABLE_BF16`. | The custom FC1/FC2 GEMV kernels run for BF16; no grouped-GEMM fallback when the FP16 gate would route. | +| Profiling | GPT-OSS-20B, Qwen3.6-35B-A3B, and Gemma model shapes profiled with `block_size=64` for both dtypes. | BF16 matches FP16 routing and latency within noise (about 1.3x–1.5x faster than grouped GEMM); SwiGLU BF16 parity tests pass. | + +#### Experiments rejected after profiling + +| Experiment | Why it was rejected | +|------------|---------------------| +| Broad GEMV enablement for tiny 128x256 cases | Raw GEMV compute kernels were faster, but end-to-end ORT loop latency was worse than grouped GEMM. | +| Expanded rows 8/16 for 1024x4096 before model-specific tuning | GEMV compute stayed competitive, but total latency regressed for larger expanded-row counts. | +| `CtaN=16, Threads=128` | Fewer N-tile CTAs did not offset lower per-CTA efficiency; GPT and Gemma kernels slowed down. | +| `CtaN=8, Threads=64` | Lower thread count slowed both profiled model-size cases. | +| Map-specialized launch | Avoiding the runtime map/prefix branch was neutral or slightly worse and added code size. | +| Naive FC2 GEMV + finalize fusion | Serialized top-k expert GEMVs inside one CTA; GPT-OSS FC2/finalize kernel time regressed sharply. | +| One-row finalize with 128 threads | Underutilized the 2880-wide GPT-OSS output; the 256-thread static top-k variant was better. | +| Asymmetric block-size-128 fallback stress cases | Existing grouped-GEMM parity tolerance was exceeded in this environment; they were not kept in the GEMV-focused matrix. | + +#### Remaining ideas not tried + +- A better FC2/finalize design that preserves parallelism across top-k experts + and N tiles, then reduces partial results with bounded numerical change. +- Architecture-specific dispatch thresholds or a small autotuner for SM80, SM89, + SM90, SM100, and SM120 rather than one broad hand-tuned gate. +- Asymmetric block-wise GEMV with zero-point compensation in the kernel, if model + demand and grouped-GEMM baseline data justify the maintenance cost. +- More model-shape block-wise profiling, especially `block_size=128`, INT8, and + end-to-end GenAI runs for models that ship block-wise QMoE weights. +- Native validation on SM100/SM120 for interactions between GEMV routing and the + FP4 / WFP4AFP8 paths. + +Keep `ORT_DISABLE_MOE_GEMV=1` available. It is useful for A/B testing, fallback +validation, and bisecting numerical or performance regressions. For quick +profiling, use +[profile_qmoe_gemv.sh](../../../onnxruntime/test/python/transformers/profile_qmoe_gemv.sh): + +```bash +onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --case gpt_oss_20b_m1_top4_fp16_2880x2880_e32 \ + --block-size 64 --warmup 5 --repeat 100 +``` + --- ## 13. Testing diff --git a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md new file mode 100644 index 0000000000000..5659ff4f2a091 --- /dev/null +++ b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md @@ -0,0 +1,975 @@ +# QMoE GEMV Profiling Experiments + +This file records QMoE INT4/INT8 GEMV profiling results so future kernel and dispatch changes can be compared against a stable baseline. + +## 2026-06-12 Baseline: SM90, Warmup 5, Repeat 100 + +### Setup + +- Machine/GPU: local CUDA machine, SM90 reported by `torch.cuda.get_device_capability()`. +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release`. +- Python: `~/onnxruntime/.venv/bin/python`. +- Nsight Systems: `~/cuda13.0/bin/nsys`. +- Benchmark script: `onnxruntime/test/python/transformers/profile_qmoe_gemv.sh`. +- Warmup: 5 ORT runs before the measured `benchmark` NVTX range. +- Repeat: 100 measured ORT runs inside the `benchmark` NVTX range. +- Common shape settings: hidden size 128, intermediate size 256, 4 experts, top-k 2, INT4 symmetric per-channel weights, SwiGLU interleaved activation. +- Profile log: `/tmp/qmoe_gemv_warmup5_r100.log`. +- Nsight artifacts: `/tmp/qmoe_gemv__warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`. + +Command template: + +```bash +pushd /tmp >/dev/null +PATH=~/cuda13.0/bin:$PATH \ +PYTHONPATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/onnxruntime/test/python/transformers \ +~/onnxruntime/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --python ~/onnxruntime/.venv/bin/python \ + --case --warmup 5 --repeat 100 \ + -o /tmp/qmoe_gemv__warmup5_r100 +popd >/dev/null +``` + +Parse command used for kernel-level comparisons: + +```bash +python onnxruntime/test/python/transformers/parse_nsys.py \ + /tmp/qmoe_gemv__warmup5_r100_.sqlite \ + --nvtx-range benchmark --pattern '%' +``` + +`--pattern '%'` is important for fallback runs because the grouped-GEMM compute kernels are named as CUTLASS kernels and are not shown by the default parser pattern. + +### End-To-End ORT Loop Timing + +Lower is better. `GEMV/GEMM` compares the host-observed average latency printed by `profile_qmoe_gemv.py`. + +| Case | Tokens | Expanded rows | DType | GEMV ms | GEMM fallback ms | GEMV/GEMM | Result | +|------|--------|---------------|-------|---------|------------------|-----------|--------| +| `m1_top2_fp16_128x256` | 1 | 2 | FP16 | 0.0967 | 0.0706 | 1.37x | GEMV slower | +| `m4_top2_fp16_128x256` | 4 | 8 | FP16 | 0.1372 | 0.0705 | 1.94x | GEMV slower | +| `m8_top2_fp16_128x256` | 8 | 16 | FP16 | 0.1381 | 0.0702 | 1.97x | GEMV slower | +| `m1_top2_bf16_128x256` | 1 | 2 | BF16 | 0.0972 | 0.0706 | 1.38x | GEMV slower | + +### Primary Compute Kernel Timing + +Values are average kernel duration in microseconds inside the measured NVTX range. GEMV has two compute kernels per ORT run in these cases, corresponding to the two MoE FC stages. GEMM fallback has two grouped-GEMM kernels per ORT run, except `m8_top2_fp16_128x256` appears as one aggregated CUTLASS kernel row with 200 calls. + +| Case | GEMV compute avg us | GEMM compute avg us | Notes | +|------|---------------------|---------------------|-------| +| `m1_top2_fp16_128x256` | 2.82, 2.56 | 3.98, 3.44 | GEMV compute kernels are faster, but total ORT loop is slower. | +| `m4_top2_fp16_128x256` | 3.01, 2.74 | 3.72, 3.42 | GEMV compute kernels are faster, but total ORT loop is slower. | +| `m8_top2_fp16_128x256` | 3.09, 2.75 | 3.55 aggregated over 200 calls | GEMV compute kernels are faster per call, but total ORT loop is slower. | +| `m1_top2_bf16_128x256` | 2.84, 2.61 | 4.22, 3.49 | GEMV compute kernels are faster, but total ORT loop is slower. | + +### Observations + +- With explicit warmup and `nsys` NVTX filtering, the host-observed ORT loop latency is not aligned with the earlier low-repeat smoke numbers that suggested a GEMV win. This warmed baseline shows grouped GEMM fallback faster for these small 128x256 cases. +- Kernel-only compute timing still shows the custom GEMV compute kernels faster than grouped GEMM compute kernels. The end-to-end loss likely comes from non-compute overhead around the GEMV path, dispatch/prologue behavior, extra launches, or measurement sensitivity at very small latencies. +- The default `parse_nsys.py` kernel filter can hide fallback CUTLASS grouped-GEMM kernels. Use `--pattern '%'` for GEMV-vs-GEMM profile comparisons. +- This baseline strengthens the P1/P2 work items: avoid repeated row-to-expert scans, tune tile/threshold decisions with data, and consider fusing FC1 activation or FC2 finalize/scatter before expanding GEMV coverage. + +### Next Experiments + +- Sweep larger realistic FC dimensions, especially model-like hidden/intermediate sizes, because this baseline only covers 128x256 test-sized shapes. +- Sweep expanded rows `{1, 2, 4, 8, 16, 32, 64}` while keeping shape and dtype fixed. +- Capture CUDA API timing inside the `benchmark` range to identify host-side synchronization or launch overhead differences. +- Record per-architecture results for SM80/SM89/SM90/SM100/SM120 when available. + +## 2026-06-12 Larger Shape Sweep: SM90, 1024x4096, Warmup 5, Repeat 100 + +### Setup + +- Same environment as the baseline section. +- Custom profiler arguments were used instead of named cases. +- Common shape settings: hidden size 1024, intermediate size 4096, 8 experts, top-k 2, INT4 symmetric per-channel weights, SwiGLU interleaved activation. +- Profile log: `/tmp/qmoe_gemv_1024x4096_e8_warmup5_r100.log`. +- Nsight artifacts: `/tmp/qmoe_gemv_custom__1024x4096_e8_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`. + +Command template: + +```bash +pushd /tmp >/dev/null +PATH=~/cuda13.0/bin:$PATH \ +PYTHONPATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/onnxruntime/test/python/transformers \ +~/onnxruntime/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --python ~/onnxruntime/.venv/bin/python \ + --batch-size 1 --sequence-length \ + --hidden-size 1024 --intermediate-size 4096 \ + --num-experts 8 --top-k 2 --dtype FLOAT16 \ + --warmup 5 --repeat 100 \ + -o /tmp/qmoe_gemv_custom_m_top2_float16_1024x4096_e8_warmup5_r100 +popd >/dev/null +``` + +### End-To-End ORT Loop Timing + +Lower is better. `GEMV/GEMM` compares the host-observed average latency printed by `profile_qmoe_gemv.py`. + +| Case | Tokens | Expanded rows | DType | GEMV ms | GEMM fallback ms | GEMV/GEMM | Result | +|------|--------|---------------|-------|---------|------------------|-----------|--------| +| `custom_m1_top2_float16_1024x4096_e8` | 1 | 2 | FP16 | 0.0622 | 0.0811 | 0.77x | GEMV faster | +| `custom_m4_top2_float16_1024x4096_e8` | 4 | 8 | FP16 | 0.1542 | 0.0873 | 1.77x | GEMV slower | +| `custom_m8_top2_float16_1024x4096_e8` | 8 | 16 | FP16 | 0.2083 | 0.0976 | 2.13x | GEMV slower | +| `custom_m1_top2_bfloat16_1024x4096_e8` | 1 | 2 | BF16 | 0.2144 | 0.0844 | 2.54x | GEMV slower | + +### Primary Compute Kernel Timing + +Values are average kernel duration in microseconds inside the measured NVTX range. + +| Case | GEMV compute avg us | GEMM compute avg us | Notes | +|------|---------------------|---------------------|-------| +| `custom_m1_top2_float16_1024x4096_e8` | 6.78, 4.70 | 14.61, 8.55 | GEMV wins both compute and end-to-end. | +| `custom_m4_top2_float16_1024x4096_e8` | 12.05, 10.00 | 16.56, 12.71 | GEMV compute is faster, but total ORT loop is slower. | +| `custom_m8_top2_float16_1024x4096_e8` | 23.19, 13.50 | 26.34, 14.29 | GEMV compute advantage narrows and end-to-end latency is worse. | +| `custom_m1_top2_bfloat16_1024x4096_e8` | 6.99, 4.69 | 12.86 aggregated over 200 calls | BF16 GEMV is much slower end-to-end despite comparable compute-kernel timing. | + +### Observations + +- FP16 GEMV shows a real end-to-end win for single-token decode at 1024x4096, unlike the tiny 128x256 test shape. +- The previous broad `expanded_num_rows <= 64` dispatch threshold was too optimistic on SM90 for this shape. At expanded rows 8 and 16, grouped GEMM fallback is faster end-to-end. +- The initial P1 dispatch cutoff is therefore conservative: FP16 only, `expanded_num_rows <= 2`, and `N/K >= 1024`. Collect more model-size points around expanded rows 2, 4, 8, and 16 before enabling GEMV beyond true single-token decode. + +## 2026-06-12 Post-Threshold Dispatch Smoke: SM90, Warmup 1, Repeat 2 + +### Setup + +- Same build and Python environment as above, after syncing the rebuilt CUDA provider into the Python package `capi` directory. +- Nsight Systems profiles used the measured `benchmark` NVTX range and counted kernels matching `%moe_gemv%`. +- These runs validate routing only; repeat count is too small for performance comparison. + +| Case | Expanded rows | DType | Expected route | `moe_gemv` calls in benchmark | Result | +|------|---------------|-------|----------------|--------------------------------|--------| +| `m1_top2_fp16_128x256` | 2 | FP16 | grouped GEMM fallback (`N/K < 1024`) | 0 | Passed | +| `custom_m1_top2_float16_1024x4096_e8` | 2 | FP16 | GEMV | 4 | Passed | +| `custom_m4_top2_float16_1024x4096_e8` | 8 | FP16 | grouped GEMM fallback (`expanded_num_rows > 2`) | 0 | Passed | +| `custom_m1_top2_bfloat16_1024x4096_e8` | 2 | BF16 | grouped GEMM fallback (FP16-only GEMV gate) | 0 | Passed | + +## 2026-06-12 Actual Model Dimensions: SM90, FP16, Warmup 5, Repeat 20 + +### Setup + +- Same build and Python environment as above. +- Nsight Systems profiles used the measured `benchmark` NVTX range. +- Enabled mode leaves `ORT_DISABLE_MOE_GEMV` unset. Fallback mode sets `ORT_DISABLE_MOE_GEMV=1`. +- Nsight artifacts: `/tmp/qmoe_actual__warmup5_r20_{enabled,gemm}.{nsys-rep,sqlite}`. +- Summary log: `/tmp/qmoe_actual_model_dims_fp16_warmup5_r20_summary.txt`. + +Model dimensions came from the Hugging Face configs below. Qwen's shared expert +is ignored in this benchmark. + +| Model case | Source dimensions | Tokens | Expanded rows | Enabled ms | Fallback ms | `moe_gemv` calls in enabled profile | Route | +|------------|-------------------|--------|---------------|------------|-------------|--------------------------------------|-------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | `hidden_size=2880`, `intermediate_size=2880`, `num_local_experts=32`, `top_k=4` | 1 | 4 | 0.0909 | 0.0944 | 0 | grouped GEMM fallback | +| `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256` | `hidden_size=2048`, `moe_intermediate_size=512`, `num_experts=256`, `top_k=8` | 1 | 8 | 0.0744 | 0.0743 | 0 | grouped GEMM fallback | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | `hidden_size=2816`, `moe_intermediate_size=704`, `num_experts=128`, `top_k=8` | 1 | 8 | 0.0819 | 0.0820 | 0 | grouped GEMM fallback | + +### Observations + +- The actual single-token top-k values for these models produce expanded rows 4 + or 8, so they intentionally stay on grouped GEMM under the current + `expanded_num_rows <= 2` GEMV gate. +- Enabled and forced-fallback timings are nearly identical, and the enabled + profiles contain zero custom `moe_gemv` kernels, confirming the dispatch route. +- These profiles establish the current grouped-GEMM baseline for actual model + dimensions. Future GEMV work for model-realistic top-k needs targeted row-count + improvements before expanding the dispatch gate. + +## 2026-06-12 Actual Model Dimensions With Relaxed GEMV Gate: SM90, FP16, Warmup 5, Repeat 100 + +### Setup + +- Same build and Python environment as above. +- Dispatch was temporarily relaxed to `expanded_num_rows <= 8` and `N/K >= 512` + to force GEMV coverage for GPT-OSS-20B, Qwen3.6-35B-A3B, and Gemma-4-26B-A4B + FP16 model-size cases. +- Nsight artifacts: `/tmp/qmoe_actual__relaxed_warmup5_r100_{enabled,gemm}.{nsys-rep,sqlite}`. +- Summary log: `/tmp/qmoe_actual_model_dims_fp16_relaxed_warmup5_r100_summary.txt`. + +| Model case | Expanded rows | Enabled route | Enabled ms | Fallback ms | GEMV/GEMM | GEMV calls | Primary enabled compute avg us | Primary fallback compute avg us | Decision | +|------------|---------------|---------------|------------|-------------|-----------|------------|--------------------------------|---------------------------------|----------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 4 | GEMV | 0.0760 | 0.0905 | 0.84x | 200 | 15.35, 11.94 | 20.11, 18.48 | Keep GEMV candidate | +| `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256` | 8 | GEMV | 0.0865 | 0.0707 | 1.22x | 200 | 20.37, 18.77 | 10.33 aggregated over 200 calls | Reject current GEMV gate | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 8 | GEMV | 0.0745 | 0.0801 | 0.93x | 200 | 14.61, 11.85 | 13.58 aggregated over 200 calls | Keep GEMV candidate | + +### Observations + +- GPT-OSS-20B benefits from GEMV at expanded rows 4 with square 2880x2880 FC + dimensions. +- Gemma-4-26B-A4B shows a smaller but positive end-to-end win at expanded rows 8 + with 2816x704 / 704x2816 FC dimensions. +- Qwen3.6-35B-A3B regresses with the current GEMV kernel at expanded rows 8 and + 2048x512 / 512x2048 FC dimensions. The primary GEMV kernels are slower than + grouped GEMM for this 512-wide MoE hidden size. +- Based on this pass, the dispatch gate should keep `expanded_num_rows <= 8` and + `N/K >= 512`, but require the logical MoE intermediate size and each GEMV call + dimension to be at least 704 when `expanded_num_rows > 4`. The logical + intermediate-size check is needed because Qwen's gated FC1 has `N=1024` even + though the MoE hidden size is 512. This keeps GPT-OSS and Gemma enabled while + routing the measured Qwen regression to grouped GEMM. A future autotuner could + replace this hand cutoff. + +### Final Route Check After Logical Intermediate-Size Guard + +After adding the logical intermediate-size guard, short Nsight profiles with +warmup 1 and repeat 2 confirmed the intended final routing: + +| Model case | Expected route | `moe_gemv` calls in benchmark | Result | +|------------|----------------|--------------------------------|--------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | GEMV | 4 | Passed | +| `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256` | grouped GEMM fallback | 0 | Passed | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | GEMV | 4 | Passed | + +## 2026-06-13 P1 Row-To-Expert Map: SM90, FP16, Warmup 5, Repeat 100 + +### Setup + +- Same build and Python environment as above. +- Change under test: the prologue now writes the local expert id for each + permuted row into `permuted_token_selected_experts`, and the INT4 + per-channel MoE GEMV kernel uses that direct row-to-expert map instead of + scanning `expert_first_token_offset` in every N-tile CTA. The prefix-offset + scan remains as a fallback when no map is passed. +- Nsight artifacts: + `/tmp/qmoe_actual_gpt_oss_20b_p1_row_expert_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}` + and + `/tmp/qmoe_actual_gemma4_26b_p1_row_expert_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`. +- Summary logs: + `/tmp/qmoe_actual_gpt_oss_20b_p1_row_expert_warmup5_r100.log` and + `/tmp/qmoe_actual_gemma4_26b_p1_row_expert_warmup5_r100.log`. + +### End-To-End ORT Loop Timing + +| Model case | Expanded rows | Enabled route | Enabled ms | Fallback ms | GEMV/GEMM | GEMV calls | Result | +|------------|---------------|---------------|------------|-------------|-----------|------------|--------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 4 | GEMV | 0.0721 | 0.0941 | 0.77x | 200 | GEMV faster | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 8 | GEMV | 0.0610 | 0.0795 | 0.77x | 200 | GEMV faster | + +### Primary Compute Kernel Timing + +Values are average kernel duration in microseconds inside the measured NVTX +range. The two GEMV compute rows correspond to FC1 and FC2. + +| Model case | GEMV compute avg us | Fallback compute avg us | Previous GEMV compute avg us | Notes | +|------------|---------------------|-------------------------|------------------------------|-------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 13.69, 10.22 | 20.32, 18.55 | 15.35, 11.94 | Direct map reduces GEMV kernel time by about 11% and 14% versus the relaxed-gate baseline. | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 7.17, 4.56 | 18.82, 7.99 | 14.61, 11.85 | Direct map removes a large per-tile scan cost at top-k 8 for this 704-wide case. | + +### Observations + +- The P1 row-to-expert map improves both actual-model GEMV candidates without + changing the dispatch policy. +- The improvement is largest for Gemma, where expanded rows 8 and many N tiles + made the repeated prefix scan particularly visible. +- Qwen remains routed to grouped GEMM by the logical-intermediate-size guard. A + quick enabled-mode smoke run reported valid output for + `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256`. + +## 2026-06-12 P1 Tile-Shape Probe: SM90, FP16, Warmup 5, Repeat 100 + +### Setup + +- Same build and Python environment as above. +- Goal: test the next P1 tile-shape knobs against the actual model-size cases + that route to GEMV after the row-to-expert map optimization. +- Variants tested: + - `CtaN=16, Threads=128`: halves the number of N-tile CTAs. + - `CtaN=8, Threads=64`: reduces threads per CTA while preserving the N tile. + - `CtaN=8, Threads=128` with a map-specialized launch: splits direct-map and + prefix-scan kernels so the common direct-map path avoids the runtime branch. +- Nsight artifacts: + `/tmp/qmoe_actual_*_ctan16_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`, + `/tmp/qmoe_actual_*_threads64_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`, + and `/tmp/qmoe_actual_*_specialized_map_warmup5_r100_{gemv,gemm}.{nsys-rep,sqlite}`. + +### End-To-End ORT Loop Timing + +| Variant | Model case | Enabled ms | Fallback ms | Result | +|---------|------------|------------|-------------|--------| +| `CtaN=16, Threads=128` | `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 0.0852 | 0.0872 | Reject | +| `CtaN=16, Threads=128` | `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 0.0644 | 0.0791 | Reject | +| `CtaN=8, Threads=64` | `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 0.0832 | 0.0993 | Reject | +| `CtaN=8, Threads=64` | `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 0.0645 | 0.0799 | Reject | +| map-specialized launch | `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 0.0726 | 0.0909 | Reject, neutral/slightly worse kernels | +| map-specialized launch | `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 0.0600 | 0.0790 | Reject, neutral/slightly worse kernels | + +### Primary GEMV Compute Kernel Timing + +Values are average kernel duration in microseconds inside the measured NVTX +range. Previous best is the row-to-expert map build with `CtaN=8, Threads=128`: +GPT `13.69, 10.22` us and Gemma `7.17, 4.56` us. + +| Variant | GPT GEMV avg us | Gemma GEMV avg us | Decision | +|---------|-----------------|-------------------|----------| +| `CtaN=16, Threads=128` | 19.48, 17.52 | 11.07, 5.65 | Wider N tile is slower for both models. | +| `CtaN=8, Threads=64` | 18.77, 16.78 | 11.09, 5.09 | Fewer threads are slower for both models. | +| map-specialized launch | 14.00, 10.10 | 7.25, 4.62 | No clear kernel win; keep the simpler launch. | + +### Observations + +- The existing `CtaN=8, Threads=128` tile remains the best measured choice for + the current SM90 actual-model cases. +- Halving N tiles with `CtaN=16` reduced CTA count but lost too much per-CTA + efficiency. +- `Threads=64` did not help the 704-wide Gemma case and significantly hurt the + larger GPT-OSS case. +- The direct-map branch in the current kernel is not a visible bottleneck after + the row-to-expert map optimization, so the extra specialized launch variants + are not worth the additional code size. + +## 2026-06-13 FC1 Interleaved SwiGLU Fusion: SM90, FP16, Warmup 5, Repeat 100 + +### Setup + +- Same build and Python environment as above. +- Change under test: FC1 INT4 per-channel GEMV has an interleaved SwiGLU epilogue + for the profiled FP16 path. It computes adjacent gate/linear FC1 columns, + applies the existing `alpha=1.702`, `beta=1.0`, `limit=7.0` SwiGLU formula, + and writes post-activation `[expanded_rows, inter_size]` directly to the FC2 + input buffer. +- The unfused GEMV plus `doGatedActivationKernel` path remains the fallback for + non-FP16, non-INT4-per-channel, groupwise quantization, non-interleaved + activations, and shapes rejected by the existing GEMV dispatch policy. +- Nsight artifacts: + `/tmp/ort_qmoe_profile/qmoe_swiglu_fused_gpt_{gemv,gemm}.{nsys-rep,sqlite}`, + `/tmp/ort_qmoe_profile/qmoe_swiglu_fused_gemma_{gemv,gemm}.{nsys-rep,sqlite}`, + and `/tmp/ort_qmoe_profile/qmoe_swiglu_fused_qwen_{gemv,gemm}.{nsys-rep,sqlite}`. + +### Correctness + +- Focused QMoE GEMV smoke: + `TestQMoEGemvBenchmark::test_decode_latency` passed. +- Large enough FP16 INT4 SwiGLU parity case to exercise the fused path + (`hidden_size=1024`, `intermediate_size=512`, `num_experts=4`, `top_k=2`) had + max absolute difference `0.0009766` before the helper's parity check and + `0.000893` in the built-in parity check. + +### End-To-End ORT Loop Timing + +Lower is better. Previous best is the P1 row-to-expert map build with the same +`CtaN=8, Threads=128` tile shape. + +| Model case | Expanded rows | Enabled route | Fused GEMV ms | Previous best GEMV ms | Fallback ms | Result | +|------------|---------------|---------------|---------------|-----------------------|-------------|--------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 4 | GEMV | 0.0681 | 0.0721 | 0.1013 | GEMV faster; fusion improves enabled path by about 5.5%. | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 8 | GEMV | 0.0580 | 0.0610 | 0.0811 | GEMV faster; fusion improves enabled path by about 4.9%. | +| `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256` | 8 | grouped GEMM fallback | 0.0706 | N/A | 0.0710 | No custom GEMV kernels; route unchanged by fusion. | + +### Primary Kernel Timing + +Values are average kernel duration in microseconds inside the measured NVTX +range. The fused GEMV compute rows correspond to fused FC1 and unfused FC2. + +| Model case | Fused GEMV compute avg us | Previous best GEMV compute avg us | Removed activation avg us from fallback profile | Notes | +|------------|---------------------------|-----------------------------------|-----------------------------------------------|-------| +| `gpt_oss_20b_m1_top4_fp16_2880x2880_e32` | 13.93, 10.11 | 13.69, 10.22 | 3.23 | FC1 GEMV is slightly slower due to the fused epilogue, but removing the activation launch gives a net end-to-end win. | +| `gemma4_26b_a4b_m1_top8_fp16_2816x704_e128` | 8.30, 4.57 | 7.17, 4.56 | 1.82 | The fused epilogue costs more for the small FC1 tile, but still beats the separate activation launch end-to-end. | +| `qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256` | N/A | N/A | 1.76 | The 512-wide intermediate-size guard keeps Qwen on grouped GEMM. | + +### Observations + +- FC1 interleaved SwiGLU fusion is a modest but real win for the two actual-model + cases that already route to GEMV. +- The win comes from launch and memory-traffic removal, not from faster FC1 + compute. The fused FC1 kernel is slightly slower than the prior FC1 GEMV + kernel, so this optimization should stay narrowly gated to the measured path. +- The FC2 finalize/scatter fusion remains a larger but riskier opportunity. It + can remove another launch, but it needs a design that preserves the current + float accumulation behavior across top-k experts. + +## 2026-06-13 GPT-OSS FC2 / Finalize Follow-Up: SM90, FP16, Warmup 5, Repeat 100 + +### Setup + +- Same build and Python environment as above. +- Target case: `gpt_oss_20b_m1_top4_fp16_2880x2880_e32`. +- Goal: find remaining GPT-OSS-specific opportunity after FC1 interleaved SwiGLU + fusion. The prior profile had FC1 GEMV around `13.93` us, FC2 GEMV around + `10.11` us, and finalize around `5.65` us. +- Nsight artifacts: + `/tmp/ort_qmoe_profile/qmoe_fc2_finalize_fused_gpt_{gemv,gemm}.{nsys-rep,sqlite}` + for the rejected fused-FC2 prototype and + `/tmp/ort_qmoe_profile/qmoe_finalize_onerow_gpt_{gemv,gemm}.{nsys-rep,sqlite}` + for the retained one-row finalize specialization. + +### Rejected Prototype: FC2 GEMV + Finalize Fusion + +The prototype assigned one CTA to each original token and N tile, then looped +over the token's top-k experts inside that CTA. It avoided atomics and preserved +the top-k accumulation order, but it also serialized the four GPT-OSS FC2 GEMVs +that previously ran as separate expanded-row CTAs. + +| Variant | Enabled ms | Fused FC2/finalize avg us | Decision | +|---------|------------|---------------------------|----------| +| FC2 GEMV + finalize fusion | 0.0848 | 33.21 | Reject; removed from code. | + +### Retained Prototype: One-Row Finalize Specialization + +The retained change keeps the existing FC2 GEMV kernel and only specializes +`finalizeMoeRoutingKernelLauncher` for `num_rows == 1` and `top_k <= 4`. The +specialized kernel caches `unpermuted_row_to_permuted_row`, expert ids, and +routing scales once in shared memory instead of reloading them for every output +vector element. + +| Variant | Enabled ms | FC1 GEMV avg us | FC2 GEMV avg us | Finalize avg us | Result | +|---------|------------|-----------------|-----------------|-----------------|--------| +| FC1 fused baseline | 0.0681 | 13.93 | 10.11 | ~5.65 | Baseline for comparison. | +| one-row finalize specialization | 0.0681 | 13.98 | 10.09 | 5.50 | Kept as first step; small finalize-kernel win, end-to-end neutral within noise. | +| static top-k one-row specialization | 0.0671 | 13.93 | 10.01 | 5.00 | Keep; top-k 4 compile-time unroll improves finalize and gives the best GPT-OSS enabled latency so far. | +| static top-k with 128 threads | 0.0684 | 13.82 | 10.24 | 6.48 | Reject; smaller block underutilizes the 2880-wide output. | + +### Correctness + +- Large enough FP16 INT4 SwiGLU top-k 4 parity case + (`hidden_size=1024`, `intermediate_size=1024`, `num_experts=8`, `top_k=4`) had + max absolute difference `0.0006104` before the helper's parity check and + `0.000732` in the built-in parity check. + +### Observations + +- GPT-OSS still has room around FC2 and finalize, but a simple single-CTA + FC2/finalize fusion is the wrong shape because it gives up expanded-row + parallelism. +- A viable future FC2/finalize design likely needs to preserve parallelism across + top-k experts and N tiles, then reduce partial results without atomics or with + carefully bounded numerical change. +- The current one-row finalize specialization is intentionally small: it is useful + for GPT-style single-token decode and does not alter routing for larger batches + or top-k 8 models. Dispatching exact top-k specializations is worthwhile for + GPT-OSS top-k 4; reducing the one-row block from 256 to 128 threads is not. + +## 2026-06-13 GenAI End-To-End Throughput: GPT-OSS-20B INT4, SM90 (H200), Batch 1 + +### Setup + +- Measures full ONNX Runtime GenAI token-generation throughput, not isolated + kernel time, so it captures the real end-to-end impact of the MoE GEMV path. +- GPU: single H200 (SM90), `CUDA_VISIBLE_DEVICES=1` on an otherwise idle GPU. + GPU 0 was busy with another job during early runs and produced corrupted + numbers, so all results below use an idle GPU. +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release` (branch + `tlw/rel-1.27.0_qmoe_update`), CUDA 13.0, `CMAKE_CUDA_ARCHITECTURES=89;90`. +- ONNX Runtime GenAI: `0.14.0-dev`, venv `~/.venv_src_8f0278c` (Python 3.14). +- Model: `gpt-oss-20b` INT4 per-channel, + `~/models/gpt-oss-20b/cuda/cuda-int4-kquant-block-32-mixed/` + (`hidden=inter=2880`, 32 experts, top-k 4, SwiGLU interleaved). Single-token + decode expands to 4 rows per step, which routes to the INT4 per-channel GEMV. +- Three configurations compared: + - **Cutlass baseline (grouped GEMM)**: GEMV disabled via + `ORT_DISABLE_MOE_GEMV=1`, so FC1/FC2 use the cutlass grouped-GEMM path. + - **GEMV (final)**: default build with the INT4 per-channel MoE GEMV fast path, + including the row-to-expert map, FC1 interleaved SwiGLU fusion, and static + top-k one-row finalize specialization. + - **FT baseline**: FasterTransformer MoE kernel used in ORT 1.26 (or ORT GenAI 0.14.1) reference token-generation throughput for the same model and prompt lengths. +- Correctness: both GEMV-enabled and `ORT_DISABLE_MOE_GEMV=1` produce identical + correct output ("Paris is the capital of France.") for the sanity prompt, and + the Nsight trace confirms `moe_gemv_kernel` executes for FC1 and FC2. + +Command template (run once per configuration): + +```bash +cd ~/onnxruntime-genai/benchmark/python +source ~/.venv_src_8f0278c/bin/activate +export LD_LIBRARY_PATH=~/cuda13.0/lib64:~/cudnn9.19_cuda13/lib:$LD_LIBRARY_PATH +# Add ORT_DISABLE_MOE_GEMV=1 for the cutlass baseline. +CUDA_VISIBLE_DEVICES=1 python benchmark_e2e.py \ + -i ~/models/gpt-oss-20b/cuda/cuda-int4-kquant-block-32-mixed/ \ + -b 1 -l 128,1024,2048 -g 256 -r 10 -w 2 \ + --use_random_tokens --chat_template '{input}' \ + -pm 1 -e cuda -mn gpt-oss-20b -pr int4 -o results_gptoss/.csv +``` + +### Token-Generation Throughput + +Higher is better. Values are average token-generation throughput in tokens per +second (tps) for batch size 1, 256 generated tokens, 10 repeats, 2 warmups. + +| Prompt length | Cutlass baseline (gemm) tps | GEMV (final) tps | FT baseline tps | GEMV vs cutlass | GEMV vs FT | +|---------------|-----------------------------|------------------|-----------------|-----------------|------------| +| 128 | 248.9 | 288.0 | 265.2 | +15.7% | +8.6% | +| 1024 | 237.8 | 272.2 | 252.9 | +14.5% | +7.6% | +| 2048 | 231.3 | 265.0 | 245.6 | +14.6% | +7.9% | + +### Observations + +- The final MoE GEMV path beats both the cutlass grouped-GEMM baseline and the + FasterTransformer 0.14.1 reference at every measured prompt length. +- Versus the cutlass grouped-GEMM baseline, GEMV improves end-to-end + token-generation throughput by roughly 15% across prompt lengths. +- Versus the FT 0.14.1 baseline, GEMV is about 8% faster, meeting the goal of + outperforming FasterTransformer for GPT-OSS-20B single-token decode. +- These end-to-end gains are consistent with the kernel-level wins recorded in + the row-to-expert map, FC1 interleaved SwiGLU fusion, and static top-k finalize + sections above. +- Always confirm the benchmark GPU is idle (`nvidia-smi`) before recording + numbers; contention on a shared GPU inflates sampling latency and corrupts the + throughput measurement. + +## 2026-06-13 Block-Wise INT4/INT8 GEMV: SM90, `block_size=64`, Warmup 5, Repeat 100 + +### Setup + +- Goal: extend the CUDA QMoE GEMV path from INT4 per-column quantization to + symmetric block-wise integer quantization. +- GPU: H200, SM90. +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release`. +- Python: `~/onnxruntime/.venv/bin/python`. +- Nsight Systems: `~/cuda13.0/bin/nsys`. +- All model-shape runs used `--warmup 5 --repeat 100 --block-size 64` and parsed + the `benchmark` NVTX range. +- Kernel parsing used `parse_nsys.py --pattern '%'` so the CUTLASS fallback + kernels and custom GEMV kernels both appear. + +Implementation summary: + +- The GEMV kernels now support symmetric INT4 and INT8 block-wise scales for + group sizes 64 and 128. +- Block-wise scale inputs are `[E, N, K_blocks]`; QMoE prepack/runtime transposes + them to `[E, K_blocks, N]`, and GEMV consumes that same layout. +- Block-wise GEMV is symmetric only. When zero-point compensation is present, + dispatch rejects GEMV and falls back to grouped GEMM. +- The per-column INT4 path is unchanged and still uses `GroupSize == 0`. + +Command template: + +```bash +cd /tmp +source ~/onnxruntime/.venv/bin/activate +export CUDA_HOME=~/cuda13.0 +export CUDNN_HOME=~/cudnn9.19_cuda13 +export PATH=$CUDA_HOME/bin:$PATH +export LD_LIBRARY_PATH=/usr/lib64/openmpi/lib:$CUDA_HOME/lib64:$CUDNN_HOME/lib64:$CUDNN_HOME/lib:${LD_LIBRARY_PATH:-} +export PYTHONPATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/onnxruntime/test/python/transformers + +~/onnxruntime/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --python ~/onnxruntime/.venv/bin/python \ + --case \ + --block-size 64 --warmup 5 --repeat 100 \ + -o /tmp/qmoe_profile_20260613/ +``` + +### Correctness and Smoke Coverage + +- Provider build passed after threading `MOEParallelismConfig` into the static + `gemm1` helper. +- Python syntax passed for `test_qmoe_cuda.py` and `profile_qmoe_gemv.py`. +- Focused parity passed: `python -m pytest .../test_qmoe_cuda.py -k "blockwise"` + reported `40 passed, 35 deselected`. +- Short benchmark smokes on SM90 produced finite output for INT4/INT8, + `block_size` 64/128, with `expanded_num_rows=2`. + +| Case | Quant Bits | Block Size | Latency (ms) | Status | +|------|------------|------------|--------------|--------| +| `blockwise_int4_b64_m1_top2_fp16_1024x4096_e8` | 4 | 64 | 0.074851 | passed | +| `blockwise_int4_b128_m1_top2_fp16_1024x4096_e8` | 4 | 128 | 0.129937 | passed | +| `blockwise_int8_b64_m1_top2_fp16_1024x4096_e8` | 8 | 64 | 0.078567 | passed | +| `blockwise_int8_b128_m1_top2_fp16_1024x4096_e8` | 8 | 128 | 0.081153 | passed | + +### Route Investigation + +- The symmetric INT4 block-wise ONNX graph has empty zero-point inputs 11/12 and + no zero-point initializers. +- Temporary diagnostics showed early profiler/tactic invocations could reject + GEMV, but the measured benchmark loop had null zero-point pointers and launched + the custom kernels. +- Temporary diagnostics were removed before final profiling; the provider no + longer contains `MOE_GEMV_DEBUG`, `QMOE_PREPACK_DEBUG`, or `QMOE_COMPUTE_DEBUG` + strings. +- Route verification must come from Nsight kernels inside the `benchmark` NVTX + range, not from the benchmark JSON alone. + +### GPT-OSS-20B Shape, INT4, `block_size=64` + +Shape: `M=1`, `top_k=4`, `expanded_num_rows=4`, `hidden_size=2880`, +`intermediate_size=2880`, `num_experts=32`, FP16. + +Artifacts: + +- GEMV enabled: `/tmp/qmoe_profile_20260613/gpt_oss_20b_b64_gemv.sqlite` +- GEMV disabled: `/tmp/qmoe_profile_20260613/gpt_oss_20b_b64_gemm.sqlite` + +| Mode | Benchmark latency (ms) | Key FC kernels in benchmark range | +|------|------------------------|-----------------------------------| +| GEMV enabled | 0.068951 | `moe_gemv_interleaved_swiglu_kernel`: 100 calls, 14.37 us avg; `moe_gemv_kernel`: 100 calls, 10.91 us avg | +| GEMV disabled | 0.096503 | `MoeFCGemm`: 200 calls, 22.36 us avg | + +Result: real GEMV route, about 1.40x faster than grouped GEMM fallback by +end-to-end benchmark latency. + +### Qwen3-6-35B-A3B Shape, INT4, `block_size=64` + +Shape: `M=1`, `top_k=8`, `expanded_num_rows=8`, `hidden_size=2048`, +`intermediate_size=512`, `num_experts=256`, FP16. + +The old `kMinProfiledProblemDimForExpandedRowsAbove4 = 704` policy produced a +fallback-vs-fallback comparison: + +| Mode | Benchmark latency (ms) | Route | +|------|------------------------|-------| +| GEMV enabled | 0.072524 | grouped GEMM fallback, `MoeFCGemm`: 200 calls, 10.90 us avg | +| GEMV disabled | 0.074525 | grouped GEMM fallback, `MoeFCGemm`: 200 calls, 10.93 us avg | + +After lowering `kMinProfiledProblemDimForExpandedRowsAbove4` to 512: + +Artifacts: + +- GEMV enabled: `/tmp/qmoe_profile_20260613/qwen3_6_35b_a3b_b64_gate512_gemv.sqlite` +- GEMV disabled: `/tmp/qmoe_profile_20260613/qwen3_6_35b_a3b_b64_gate512_gemm.sqlite` + +| Mode | Benchmark latency (ms) | Key FC kernels in benchmark range | +|------|------------------------|-----------------------------------| +| GEMV enabled | 0.052437 | `moe_gemv_interleaved_swiglu_kernel`: 100 calls, 5.07 us avg; `moe_gemv_kernel`: 100 calls, 3.28 us avg | +| GEMV disabled | 0.073088 | `MoeFCGemm`: 200 calls, 10.81 us avg | + +Result: the 512 gate makes Qwen a true GEMV-vs-GEMM comparison and improves +end-to-end latency by about 1.39x. FC kernel time drops from about 2.162 ms +total per 100 iterations for grouped GEMM to about 0.835 ms total per 100 +iterations for GEMV. + +### Decision + +- Keep `block_size=64` as the immediate model-shape profiling target. +- Keep symmetric block-wise INT4/INT8 support for both 64 and 128 in the + implementation, but only report model-shape profiling for 64 in this round. +- Keep `kMinProfiledProblemDimForExpandedRowsAbove4 = 512` so Qwen-style + `top_k=8`, `intermediate_size=512` decode runs can use the custom GEMV path. + +## 2026-06-14 Block Size 32 vs 64: SM90, FP16 INT4, Warmup 5, Repeat 100 + +### Setup + +- Goal: compare the newly enabled INT4 `block_size=32` path against the existing + `block_size=64` path on real model-shaped QMoE decode workloads. +- GPU: single H200 (SM90). +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release`, CUDA 13.0, + after the `block_size=32` CUTLASS/GEMV changes were built and installed. +- Python: `~/onnxruntime/.venv/bin/python`. +- Nsight Systems: `nsys 2025.3.2.367`. The first `nsys` run hit a + `/tmp/nvidia/nsight_systems` permission issue; rerunning with `TMPDIR` under + the artifact directory fixed it. +- Artifacts: `/tmp/qmoe_profile_block32_vs64_20260614_024925/`. + +The profiling CLI was updated to accept `--block-size 32`; prior to this run it +only allowed `0/64/128` even though the benchmark helper could construct custom +cases. + +Command template: + +```bash +cd ~/onnxruntime +source .venv/bin/activate +export CUDA_VERSION=13.0 +export CUDA_HOME=~/cuda13.0 +export CUDNN_HOME=~/cudnn_9.19_cuda13 +export PATH=$CUDA_HOME/bin:$PATH +export LD_LIBRARY_PATH=/usr/lib64/openmpi/lib:$CUDA_HOME/lib64:$CUDNN_HOME/lib64:$CUDNN_HOME/lib:${LD_LIBRARY_PATH:-} +export PYTHONPATH=$PWD/build/cu130/Release:$PWD/onnxruntime/test/python/transformers:${PYTHONPATH:-} + +python onnxruntime/test/python/transformers/profile_qmoe_gemv.py \ + --case --block-size <32|64> --warmup 5 --repeat 100 + +TMPDIR=/tmp/qmoe_profile_block32_vs64_20260614_024925/tmp \ +bash onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --case --block-size <32|64> --warmup 5 --repeat 100 \ + -o /tmp/qmoe_profile_block32_vs64_20260614_024925/ +``` + +### Standalone End-To-End Benchmark Latency + +Lower is better. These numbers are from the plain benchmark loop, outside nsys, +with the default GEMV-enabled route. Every case reported +`has_invalid_output=false`. + +| Model case | Expanded rows | Block size | Latency (ms) | Relative to b64 | +|------------|---------------|------------|--------------|-----------------| +| `gpt_oss_20b` (`2880x2880`, e32, top4) | 4 | 32 | 0.068118 | 1.006x | +| `gpt_oss_20b` (`2880x2880`, e32, top4) | 4 | 64 | 0.067712 | 1.000x | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | 8 | 32 | 0.050836 | 0.991x | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | 8 | 64 | 0.051312 | 1.000x | + +Result: `block_size=32` and `block_size=64` are effectively tied on the default +GEMV path for these model-shaped decode cases. GPT-OSS is about 0.6% faster with +`block_size=64`; Qwen3.6-35B-A3B is about 0.9% faster with `block_size=32`. Both +differences are small enough to treat as measurement noise unless repeated +end-to-end model runs show the same direction. + +### Nsight Route Confirmation And FC Kernel Timing + +All rows below are restricted to the `benchmark` NVTX range. The GEMV-enabled +route uses `moe_gemv_interleaved_swiglu_kernel` for FC1 and `moe_gemv_kernel` for +FC2. The disabled route (`ORT_DISABLE_MOE_GEMV=1`) uses CUTLASS `MoeFCGemm`. + +| Model case | Block size | Mode | Route kernel | Calls | Avg us | +|------------|------------|------|--------------|-------|--------| +| `gpt_oss_20b` | 32 | GEMV enabled | `moe_gemv_interleaved_swiglu_kernel` | 100 | 17.878 | +| `gpt_oss_20b` | 32 | GEMV enabled | `moe_gemv_kernel` | 100 | 12.978 | +| `gpt_oss_20b` | 32 | GEMV disabled | `MoeFCGemm` | 200 | 43.283 | +| `gpt_oss_20b` | 64 | GEMV enabled | `moe_gemv_interleaved_swiglu_kernel` | 100 | 17.655 | +| `gpt_oss_20b` | 64 | GEMV enabled | `moe_gemv_kernel` | 100 | 12.364 | +| `gpt_oss_20b` | 64 | GEMV disabled | `MoeFCGemm` | 200 | 41.113 | +| `qwen3_6_35b_a3b` | 32 | GEMV enabled | `moe_gemv_interleaved_swiglu_kernel` | 100 | 6.266 | +| `qwen3_6_35b_a3b` | 32 | GEMV enabled | `moe_gemv_kernel` | 100 | 4.000 | +| `qwen3_6_35b_a3b` | 32 | GEMV disabled | `MoeFCGemm` | 200 | 20.719 | +| `qwen3_6_35b_a3b` | 64 | GEMV enabled | `moe_gemv_interleaved_swiglu_kernel` | 100 | 6.201 | +| `qwen3_6_35b_a3b` | 64 | GEMV enabled | `moe_gemv_kernel` | 100 | 3.930 | +| `qwen3_6_35b_a3b` | 64 | GEMV disabled | `MoeFCGemm` | 200 | 20.000 | + +The custom GEMV kernels are present for both `block_size=32` and 64. Kernel-level +timing is slightly lower for `block_size=64` in this profile, but the standalone +benchmark latency is essentially flat between the two block sizes. + +### nsys Wrapper Benchmark Latency + +These numbers are useful for route comparisons but include profiling overhead; +prefer the standalone table above for clean latency. Lower is better. + +| Model case | Mode | Block size 32 (ms) | Block size 64 (ms) | +|------------|------|--------------------|--------------------| +| `gpt_oss_20b` | GEMV enabled | 0.075520 | 0.072995 | +| `gpt_oss_20b` | GEMV disabled | 0.140200 | 0.134682 | +| `qwen3_6_35b_a3b` | GEMV enabled | 0.055514 | 0.064243 | +| `qwen3_6_35b_a3b` | GEMV disabled | 0.093866 | 0.096064 | + +The profiled fallback runs confirm GEMV remains substantially faster than the +grouped-GEMM fallback for both block sizes and both model shapes. + +### Artifacts + +Primary summaries: + +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/benchmark_results.jsonl` +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/route_kernel_summary.tsv` +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/artifacts.txt` + +Nsight Systems produced one `.nsys-rep` and one `.sqlite` for each +`{model, block_size, mode}` tuple. Examples: + +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/nsys_gpt_oss_20b_m1_top4_fp16_2880x2880_e32_b32_gemv.sqlite` +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/nsys_gpt_oss_20b_m1_top4_fp16_2880x2880_e32_b64_gemv.sqlite` +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/nsys_qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256_b32_gemv.sqlite` +- `/tmp/qmoe_profile_block32_vs64_20260614_024925/nsys_qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256_b64_gemv.sqlite` + +### Decision + +- Keep `block_size=32` enabled for INT4 QMoE. It reaches the same custom GEMV + route as `block_size=64` for the GPT-OSS-20B and Qwen3.6-35B-A3B decode + shapes, with no meaningful end-to-end latency regression in the standalone + benchmark loop. +- Do not tune the GEMV gate differently for 32 vs 64 based on this data. The + current model-shape route is valid for both block sizes. +- Continue using nsys NVTX-range kernel evidence, not benchmark JSON alone, when + verifying whether a block-wise QMoE case actually reached GEMV or fell back to + grouped GEMM. + +## 2026-06-13 BF16 GEMV Enablement: SM90, Warmup 5, Repeat 100 + +### Setup + +- Goal: extend the CUDA QMoE GEMV fast path from FP16-only activations to BF16, + reaching dispatch and kernel parity with FP16. +- GPU: single H200 (SM90), `CUDA_VISIBLE_DEVICES=1` on an otherwise idle GPU + (`nvidia-smi` confirmed 0% before recording). +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release`, CUDA 13.0. +- Python: `~/onnxruntime/.venv_cu130/bin/python` (Python 3.14). +- Nsight Systems: `~/cuda13.0/bin/nsys`. Kernel parsing used + `parse_nsys.py --nvtx-range benchmark --pattern '%'` so both the CUTLASS + grouped-GEMM fallback and the custom GEMV kernels appear. +- All model-shape runs used `--warmup 5 --repeat 100 --block-size 64`. + +Implementation summary: + +- The runtime gate in `moe_kernels.cu` relaxes from `std::is_same_v` to + `std::is_same_v || std::is_same_v` for both + `tryLaunchMoeGemvIntSymmetric` and the interleaved-SwiGLU variant. +- `moe_gemv.cu` adds `__nv_bfloat16` `DetailsForTAndWeight` specializations for + `cutlass::uint4b_t` and `uint8_t`, plus the matching `__nv_bfloat16` template + instantiations (group sizes 0/64/128, INT4/INT8, bias on/off) under + `ENABLE_BF16`. +- BF16 and FP16 now share one dispatch gate and one set of custom kernels. + +Command template (run once per dtype): + +```bash +cd /tmp +export CUDA_HOME=~/cuda13.0 +export CUDNN_HOME=~/cudnn9.19_cuda13 +export PATH=$CUDA_HOME/bin:$PATH +export LD_LIBRARY_PATH=/usr/lib64/openmpi/lib:$CUDA_HOME/lib64:$CUDNN_HOME/lib64:$CUDNN_HOME/lib:${LD_LIBRARY_PATH:-} +export PYTHONPATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/onnxruntime/test/python/transformers + +CUDA_VISIBLE_DEVICES=1 \ +~/onnxruntime/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \ + --python ~/onnxruntime/.venv_cu130/bin/python \ + --case --dtype \ + --block-size 64 --warmup 5 --repeat 100 \ + -o /tmp/qmoe_gemv_bf16_20260613/ +``` + +### Routing Parity + +Both dtypes route to the same custom kernels for the same shape: +`moe_gemv_interleaved_swiglu_kernel` for FC1 and `moe_gemv_kernel` for FC2. + +| Case | Expanded rows | DType | FC1 kernel | FC2 kernel | Route | +|------|---------------|-------|------------|------------|-------| +| `gpt_oss_20b` (`2880x2880`, e32, top4) | 4 | FP16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | +| `gpt_oss_20b` (`2880x2880`, e32, top4) | 4 | BF16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | 8 | FP16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | 8 | BF16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | +| `gemma4_26b_a4b` (`2816x704`, e128, top8) | 8 | FP16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | +| `gemma4_26b_a4b` (`2816x704`, e128, top8) | 8 | BF16 | `moe_gemv_interleaved_swiglu_kernel` | `moe_gemv_kernel` | GEMV | + +### End-To-End Benchmark Latency (block_size=64, INT4) + +Lower is better. `Enabled` is the default GEMV build; `Fallback` sets +`ORT_DISABLE_MOE_GEMV=1`. Values are the benchmark-loop latency in milliseconds. + +| Model case | DType | Enabled ms | Fallback ms | Speedup | +|------------|-------|------------|-------------|---------| +| `gpt_oss_20b` (`2880x2880`, e32, top4) | FP16 | 0.0683 | 0.0969 | 1.42x | +| `gpt_oss_20b` (`2880x2880`, e32, top4) | BF16 | 0.0673 | 0.0995 | 1.48x | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | FP16 | 0.0512 | 0.0730 | 1.43x | +| `qwen3_6_35b_a3b` (`2048x512`, e256, top8) | BF16 | 0.0526 | 0.0724 | 1.38x | +| `gemma4_26b_a4b` (`2816x704`, e128, top8) | FP16 | 0.0595 | 0.0827 | 1.39x | +| `gemma4_26b_a4b` (`2816x704`, e128, top8) | BF16 | 0.0604 | 0.0771 | 1.28x | + +### Primary FC Compute Kernel Timing + +Average kernel duration in microseconds inside the measured `benchmark` NVTX +range. GEMV columns are FC1 (`moe_gemv_interleaved_swiglu_kernel`) and FC2 +(`moe_gemv_kernel`); the fallback column is the CUTLASS `MoeFCGemm` average over +its 200 calls. + +| Model case | DType | FC1 GEMV us | FC2 GEMV us | Fallback `MoeFCGemm` us | +|------------|-------|-------------|-------------|--------------------------| +| `gpt_oss_20b` | FP16 | 14.17 | 10.97 | 22.14 | +| `gpt_oss_20b` | BF16 | 14.58 | 11.00 | 22.71 | +| `qwen3_6_35b_a3b` | FP16 | 5.11 | 3.24 | 10.67 | +| `qwen3_6_35b_a3b` | BF16 | 5.29 | 3.35 | 10.94 | +| `gemma4_26b_a4b` | FP16 | 7.52 | 4.59 | 14.96 | +| `gemma4_26b_a4b` | BF16 | 9.45 | 5.29 | 12.28 | + +### Standalone INT4/INT8 Synthetic Shape (`1024x4096`, e8, top2) + +| Case | DType | Route | FC1 GEMV us | FC2 GEMV us | +|------|-------|-------|-------------|-------------| +| `blockwise_int4_b64` | FP16 | GEMV | 4.53 | 6.95 | +| `blockwise_int4_b64` | BF16 | GEMV | 4.62 | 7.01 | +| `blockwise_int8_b64` | FP16 | grouped GEMM fallback | — | — | +| `blockwise_int8_b64` | BF16 | grouped GEMM fallback | — | — | + +The INT4 GEMV kernel times are within noise across dtypes. End-to-end latency on +this tiny `e8` shape is high-variance (the disabled path swung between 0.17 and +0.27 ms across repeats), so only the stable in-NVTX kernel times are reported +here. INT8 block-wise stays on grouped GEMM for this shape, and that fallback is +dtype-independent (both FP16 and BF16 fall back identically). + +### Correctness + +- Focused SwiGLU BF16 parity tests passed: + `pytest test_qmoe_cuda.py -k "swiglu and bf16"` reported `16 passed, 59 + deselected`. +- Every benchmark case above reported `has_invalid_output=false`, so the + GEMV-enabled BF16 output matched the reference within tolerance. + +### Per-Channel Note + +The default per-channel (`block_size=0`) `gpt_oss` benchmark case stays on +grouped GEMM in this build for both FP16 and BF16 (the trace shows two +`MoeFCGemm` calls plus a separate `doGatedActivationKernel`). The fallback is +dtype-independent, so BF16 still matches FP16 behavior. The current branch's +confirmed model-shape GEMV coverage is the block-wise (`block_size=64`) path used +in the tables above. + +> Update (2026-06-13): the per-channel fallback observed here was a gate bug +> (`group_size = -1` rejected by `is_moe_gemv_supported`), not a fundamental +> limitation. It is fixed in the "INT8 Per-Column GEMV Enablement" section below, +> after which per-column INT4/INT8 route to GEMV for FP16 and BF16. + + +### Decision + +- Keep the relaxed gate: BF16 and FP16 share one dispatch path and one set of + custom GEMV kernels. +- For every profiled model shape, BF16 routes to GEMV exactly where FP16 does and + matches FP16 latency within measurement noise. + +## 2026-06-13 INT8 Per-Column GEMV Enablement: SM90, `block_size=0`, Warmup 5, Repeat 100 + +### Setup + +- Goal: route per-column (`block_size <= 0`) symmetric INT8 W8A16 QMoE decode + shapes to the custom GEMV fast path for FP16 and BF16 activations. +- GPU: single H200 (SM90), `CUDA_VISIBLE_DEVICES=1` on an otherwise idle GPU + (`nvidia-smi` confirmed 0% before recording). +- ONNX Runtime build: `~/onnxruntime/build/cu130/Release`, CUDA 13.0, + `CMAKE_CUDA_ARCHITECTURES=89;90`. Rebuilt with `--clean_moe`. +- Python: `~/onnxruntime/.venv_cu130/bin/python` (Python 3.14). +- Nsight Systems: `~/cuda13.0/bin/nsys`. Kernel parsing used + `parse_nsys.py --nvtx-range benchmark --pattern '%'`. +- Artifacts: `/tmp/qmoe_gemv_int8pc_20260613/_{gemv,gemm}.{nsys-rep,sqlite}`. + +### Root Cause Of The Prior Per-Column Fallback + +The earlier BF16 section's "Per-Channel Note" observed that per-channel +(`block_size=0`) cases stayed on grouped GEMM. The cause was in the GEMV gate: + +- The QMoE runtime carries per-column scales through `QuantParams::Int`, which + leaves `groupwise.group_size` at its struct default of `-1` + (`moe_kernels.h`). +- `is_moe_gemv_supported` previously required `group_size == 0` exactly for the + per-column case, so `group_size = -1` was rejected (`sup=0`) and dispatch fell + back to grouped GEMM — for per-column INT4 *and* INT8. +- The GEMV launcher dispatch (`dispatch_moe_gemv_group_size` and the + interleaved-SwiGLU variant) already maps any `group_size <= 0` to the + `GroupSize == 0` per-column kernel. + +The fix relaxes `is_moe_gemv_supported` to accept any `group_size <= 0` as the +per-column case (alongside block-wise 64/128). No new kernel instantiation is +needed: `(half, uint8_t)` and `(__nv_bfloat16, uint8_t)` GEMV details already +exist from the block-wise INT8 work. + +### Routing Confirmation + +Both FC stages now route to the custom kernels for per-column INT8. Values are +average kernel duration in microseconds inside the measured `benchmark` NVTX +range; FC1 is `moe_gemv_interleaved_swiglu_kernel`, FC2 is `moe_gemv_kernel`. + +| Case | Expanded rows | DType | FC1 GEMV us | FC2 GEMV us | Route | +|------|---------------|-------|-------------|-------------|-------| +| `int8_per_column_m1_top2_*_1024x4096_e8` | 2 | FP16 | 5.15 | 5.26 | GEMV | +| `int8_per_column_m1_top2_*_1024x4096_e8` | 2 | BF16 | 6.24 | 5.75 | GEMV | +| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | 4 | FP16 | 22.57 | 11.90 | GEMV | +| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | 4 | BF16 | 22.71 | 12.13 | GEMV | + +### End-To-End Benchmark Latency + +Lower is better. `Enabled` is the default GEMV build; `Fallback` sets +`ORT_DISABLE_MOE_GEMV=1`. Values are the benchmark-loop latency in milliseconds. +Every case reported `has_invalid_output=false`. + +| Model case | DType | Enabled ms | Fallback ms | Speedup | +|------------|-------|------------|-------------|---------| +| `int8_per_column_m1_top2_1024x4096_e8` | FP16 | 0.0566 | 0.0816 | 1.44x | +| `int8_per_column_m1_top2_1024x4096_e8` | BF16 | 0.0578 | 0.0862 | 1.49x | +| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | FP16 | 0.0785 | 0.0947 | 1.21x | +| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | BF16 | 0.0785 | 0.0989 | 1.26x | + +### Correctness + +- Focused INT8 per-column SwiGLU parity tests passed: + `pytest test_qmoe_cuda.py -k "test_swiglu_qmoe_parity_1 or + test_swiglu_qmoe_parity_3 or test_swiglu_qmoe_parity_bf16_1 or + test_swiglu_qmoe_parity_bf16_3"` reported `4 passed`. +- Regression check: `pytest -k "TestSwigluQMoE or TestQMoEIntPrePackSmoke"` + reported `34 passed, 4 skipped`. +- The per-column INT4 and block-wise INT4/INT8 routes are unchanged: block-wise + cases still pass `group_size = 64/128` and per-column INT4 now also reaches + GEMV via the same `group_size <= 0` relaxation. + +### Decision + +- Keep the relaxed gate: `is_moe_gemv_supported` accepts `group_size <= 0` as the + per-column case for INT4 and INT8. +- Per-column INT8 W8A16 decode shapes route to GEMV for both FP16 and BF16 and + beat the grouped-GEMM fallback at every profiled shape. diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h index a888ea3e71487..91f648760a965 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h @@ -323,7 +323,7 @@ struct GemmFpAIntB { {problem_size_k * kInterleave, params.problem_size.n() / kInterleave}, thread_idx, tb_offset_B, params.gather_B_indices); - typename MatrixCoord::Index scale_row_extent = isFinegrained(Mma::QuantOp) ? problem_size_k / 64 : 1; + typename MatrixCoord::Index scale_row_extent = isFinegrained(Mma::QuantOp) ? problem_size_k / params.group_size : 1; typename Mma::IteratorScale iterator_scale = initialize_scale( params.params_scale, params.ref_scale.data(), params.ref_zero.data(), {scale_row_extent, params.problem_size.n()}, thread_idx, tb_offset_scale, params.group_size); diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h index ab8ae054db048..ecfebd7a258d9 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h @@ -498,7 +498,7 @@ struct MoeFCGemm { __syncthreads(); if constexpr (use_dq_gemm::value) { - typename MatrixCoord::Index scale_row_extent = isFinegrained(Mma::QuantOp) ? gemm_k / 64 : 1; + typename MatrixCoord::Index scale_row_extent = isFinegrained(Mma::QuantOp) ? gemm_k / params.group_size : 1; typename Mma::IteratorScale iterator_scale = initialize_scale(LayoutScaleZero(ldm_Scale), reinterpret_cast(ptr_Scale), reinterpret_cast(ptr_Zero), diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_base.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_base.h index cad280febbe76..667250d53f9bd 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_base.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_base.h @@ -68,8 +68,10 @@ class DqMmaBase { static_assert(DequantOp != WeightOnlyQuantOp::UNDEFINED, ""); - // Finegrained scales get streamed in via cp.async - static constexpr int ScalebiasStages = isFinegrained(DequantOp) ? Stages : 1; + static constexpr int kMinFinegrainedGroupSize = 32; + static constexpr int kFinegrainedScaleRowsPerStage = Shape::kK / kMinFinegrainedGroupSize; + // Finegrained scales get streamed in via cp.async. + static constexpr int ScalebiasStages = isFinegrained(DequantOp) ? Stages * kFinegrainedScaleRowsPerStage : 1; // We always have scales. static constexpr int ScaleElementsPerStage = Shape::kN; // We sometimes have a bias diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h index 5db74039469c4..ebefb25e8aa78 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h @@ -138,7 +138,7 @@ class DqMmaMultistage::value * IteratorScale::kAlignment / 8; - cutlass::arch::cp_async(smem_scale_ptr, gmem_scale_ptr, iterator_scale.valid()); + int const current_group = iterator_scale.row_groupsize64_ * 64 / iterator_scale.group_size_; + int const next_group = (iterator_scale.row_groupsize64_ * 64 + Shape::kK) / iterator_scale.group_size_; + int const scale_rows_to_advance = next_group - current_group; + int const scale_rows_to_copy = scale_rows_to_advance > 0 ? scale_rows_to_advance : 1; - if (gmem_zero_ptr != nullptr) { - cutlass::arch::cp_async(smem_zero_ptr, gmem_zero_ptr, iterator_scale.valid()); - } + CUTLASS_PRAGMA_UNROLL + for (int scale_row = 0; scale_row < Base::kFinegrainedScaleRowsPerStage; ++scale_row) { + if (scale_row < scale_rows_to_copy) { + gmem_scale_ptr = iterator_scale.get_scale(); + gmem_zero_ptr = iterator_scale.get_zero(); - if (iterator_scale.group_size_ == 64) { - iterator_scale.add_tile_offset({1, 0}); - } else if (iterator_scale.group_size_ == 128) { - if constexpr (Shape::kK == 128) { - iterator_scale.add_tile_offset({1, 0}); - } else if constexpr (Shape::kK == 64) { - if (iterator_scale.row_groupsize64_ & 0x1) { + smem_scale_ptr = reinterpret_cast(this->smem_iterator_scale_.get_scale()); + smem_zero_ptr = reinterpret_cast(this->smem_iterator_scale_.get_zero()); + + cutlass::arch::cp_async(smem_scale_ptr, gmem_scale_ptr, iterator_scale.valid()); + + if (gmem_zero_ptr != nullptr) { + cutlass::arch::cp_async(smem_zero_ptr, gmem_zero_ptr, iterator_scale.valid()); + } + + if (scale_row < scale_rows_to_advance) { iterator_scale.add_tile_offset({1, 0}); } - } else { - static_assert(Shape::kK == 0, "Unsupported k tile shape, can only be 64 or 128"); } - } - iterator_scale.row_groupsize64_++; + this->smem_iterator_scale_.add_tile_offset({1, 0}); + } - this->smem_iterator_scale_.add_tile_offset({1, 0}); + iterator_scale.row_groupsize64_ += Shape::kK / 64; } CUTLASS_DEVICE @@ -469,7 +499,7 @@ class DqMmaMultistagewarp_tile_iterator_A_; ++this->warp_tile_iterator_B_; - warp_dequantizer_.add_pointer_offset(Shape::kN); + advance_dequantizer_after_load(0); iterator_A.clear_mask(gemm_k_iterations == 0); iterator_B.clear_mask(gemm_k_iterations == 0); @@ -508,6 +538,12 @@ class DqMmaMultistagewarp_tile_iterator_B_; } + int const scale_row = scale_row_for_warp_mma(warp_mma_k); + if (warp_mma_k > 0 && scale_row != scale_row_for_warp_mma(warp_mma_k - 1)) { + warp_dequantizer_.load(warp_frag_scales, warp_frag_zeros); + advance_dequantizer_after_load(scale_row); + } + typename TransformBAfterLDS::result_type converted_frag_B = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]); warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales, warp_frag_zeros); @@ -564,7 +600,7 @@ class DqMmaMultistagesmem_iterator_A_.add_tile_offset({0, -Base::kStages}); this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); - this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0}); + this->smem_iterator_scale_.add_tile_offset({-Base::kStages * Base::kFinegrainedScaleRowsPerStage, 0}); smem_write_stage_idx = 0; } else { ++smem_write_stage_idx; @@ -575,7 +611,7 @@ class DqMmaMultistagewarp_tile_iterator_B_.add_tile_offset( {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0}); - warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN); + warp_dequantizer_.add_pointer_offset(-Base::kStages * Base::kFinegrainedScaleRowsPerStage * Shape::kN); smem_read_stage_idx = 0; } else { ++smem_read_stage_idx; @@ -591,7 +627,7 @@ class DqMmaMultistage 0 ? scale_rows_to_advance : 1; - arch::global_load(tb_frag_scales, gmem_scale_ptr, iterator_scale.valid()); + CUTLASS_PRAGMA_UNROLL + for (int scale_row = 0; scale_row < Base::kFinegrainedScaleRowsPerStage; ++scale_row) { + if (scale_row < scale_rows_to_copy) { + auto gmem_scale_ptr = iterator_scale.get_scale(); + auto gmem_zero_ptr = iterator_scale.get_zero(); - if (gmem_zero_ptr != nullptr) { - arch::global_load( - tb_frag_zeros, gmem_zero_ptr, iterator_scale.valid()); - } + arch::global_load(tb_frag_scales, gmem_scale_ptr, iterator_scale.valid()); - typename TransformScale::result_type tb_frag_scales_fp16 = transformScale(tb_frag_scales); - typename TransformScale::result_type tb_frag_zeros_fp16; - if (gmem_zero_ptr != nullptr) - tb_frag_zeros_fp16 = transformScale(tb_frag_zeros); + if (gmem_zero_ptr != nullptr) { + arch::global_load( + tb_frag_zeros, gmem_zero_ptr, iterator_scale.valid()); + } - auto frag_scale_ptr_fp16 = reinterpret_cast(&tb_frag_scales_fp16); - auto frag_zero_ptr_fp16 = reinterpret_cast(&tb_frag_zeros_fp16); - auto smem_scale_ptr = this->smem_iterator_scale_.get_scale(); - auto smem_zero_ptr = this->smem_iterator_scale_.get_zero(); + typename TransformScale::result_type tb_frag_scales_fp16 = transformScale(tb_frag_scales); + typename TransformScale::result_type tb_frag_zeros_fp16; + if (gmem_zero_ptr != nullptr) + tb_frag_zeros_fp16 = transformScale(tb_frag_zeros); - if (iterator_scale.valid()) { - auto smem_offset = cast_smem_ptr_to_uint(smem_scale_ptr); - arch::shared_store(smem_offset, frag_scale_ptr_fp16); + auto frag_scale_ptr_fp16 = reinterpret_cast(&tb_frag_scales_fp16); + auto frag_zero_ptr_fp16 = reinterpret_cast(&tb_frag_zeros_fp16); + auto smem_scale_ptr = this->smem_iterator_scale_.get_scale(); + auto smem_zero_ptr = this->smem_iterator_scale_.get_zero(); - if (gmem_zero_ptr != nullptr) { - smem_offset = cast_smem_ptr_to_uint(smem_zero_ptr); - arch::shared_store(smem_offset, frag_zero_ptr_fp16); - } - } + if (iterator_scale.valid()) { + auto smem_offset = cast_smem_ptr_to_uint(smem_scale_ptr); + arch::shared_store(smem_offset, frag_scale_ptr_fp16); - if (iterator_scale.group_size_ == 64) { - iterator_scale.add_tile_offset({1, 0}); - } else if (iterator_scale.group_size_ == 128) { - if constexpr (Shape::kK == 128) { - iterator_scale.add_tile_offset({1, 0}); - } else if constexpr (Shape::kK == 64) { - if (iterator_scale.row_groupsize64_ & 0x1) { + if (gmem_zero_ptr != nullptr) { + smem_offset = cast_smem_ptr_to_uint(smem_zero_ptr); + arch::shared_store(smem_offset, frag_zero_ptr_fp16); + } + } + + if (scale_row < scale_rows_to_advance) { iterator_scale.add_tile_offset({1, 0}); } - } else { - static_assert(Shape::kK == 0, "Unsupported k tile shape, can only be 64 or 128"); } - } - iterator_scale.row_groupsize64_++; + this->smem_iterator_scale_.add_tile_offset({1, 0}); + } - this->smem_iterator_scale_.add_tile_offset({1, 0}); + iterator_scale.row_groupsize64_ += Shape::kK / 64; } /// Perform a threadblock-scoped matrix multiply-accumulate @@ -323,7 +347,7 @@ class DqMmaPipelinedwarp_tile_iterator_A_; ++this->warp_tile_iterator_B_; - warp_dequantizer_.add_pointer_offset(Shape::kN); + advance_dequantizer_after_load(0); Operator warp_mma; @@ -368,13 +392,13 @@ class DqMmaPipelinedsmem_iterator_A_.add_tile_offset({0, -Base::kStages}); this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); - this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0}); + this->smem_iterator_scale_.add_tile_offset({-Base::kStages * Base::kFinegrainedScaleRowsPerStage, 0}); } else { this->warp_tile_iterator_A_.add_tile_offset( {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); this->warp_tile_iterator_B_.add_tile_offset( {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0}); - warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN); + warp_dequantizer_.add_pointer_offset(-Base::kStages * Base::kFinegrainedScaleRowsPerStage * Shape::kN); } smem_write_stage_idx ^= 1; @@ -409,6 +433,12 @@ class DqMmaPipelined 0 && scale_row != scale_row_for_warp_mma(warp_mma_k - 1)) { + warp_dequantizer_.load(warp_frag_scales, warp_frag_zero); + advance_dequantizer_after_load(scale_row); + } + typename TransformBAfterLDS::result_type converted_frag_B = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]); warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales, warp_frag_zero); warp_mma(accum, warp_frag_A[warp_mma_k % 2], converted_frag_B, accum, warp_tileB_k_compute_offset); @@ -417,7 +447,7 @@ class DqMmaPipelined::value / 8; + const LongIndex tb_scale_row = threadblock_offset.row() * 64 / group_size; + const LongIndex tb_row_byte_offset = tb_scale_row * params_.stride_ * sizeof_bits::value / 8; const LongIndex tb_col_byte_offset = threadblock_offset.column() * sizeof_bits::value / 8; pointer_scale_ += (tb_row_byte_offset + tb_col_byte_offset); @@ -159,7 +160,7 @@ class FineGrainedScaleZeroIterator + namespace onnxruntime::llm { namespace kernels { namespace fpA_intB_gemv { @@ -41,6 +43,10 @@ struct MathWrapper { return __half2half2(v); } + __device__ __forceinline__ static float2 to_float2(Type2 const& v) { + return __half22float2(v); + } + __device__ __forceinline__ static Type2 fma2(Type2 const& a, Type2 const& b, Type2 const& c) { return __hfma2(a, b, c); } @@ -65,6 +71,14 @@ struct MathWrapper { #endif } + __device__ __forceinline__ static float2 to_float2(Type2 const& v) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __bfloat1622float2(v); +#else + return float2{0.f, 0.f}; +#endif + } + __device__ __forceinline__ static Type2 fma2(Type2 const& a, Type2 const& b, Type2 const& c) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) return __hfma2(a, b, c); @@ -147,22 +161,44 @@ __device__ __forceinline__ void pack_to_vec2(void* dst, void* src, int n) { } } -template +template ::Type> __device__ __forceinline__ void mma(void* acc, void* w_pack2, void* act) { using Type = typename MathWrapper::Type; using Type2 = typename MathWrapper::Type2; static_assert(N % 2 == 0); static constexpr int VecN = N / 2; + if constexpr (std::is_same_v) { + // fp32 accumulation: keep the per-thread K-dimension dot product in float to avoid + // the precision loss of accumulating a long chain in 16-bit (especially bf16). #pragma unroll - for (int m = 0; m < M; ++m) { + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int n = 0; n < VecN; ++n) { + float2 acc2 = reinterpret_cast(acc)[m * VecN + n]; +#pragma unroll + for (int k = 0; k < K; ++k) { + float const a = static_cast(reinterpret_cast(act)[m * K + k]); + float2 const w = MathWrapper::to_float2( + reinterpret_cast(w_pack2)[n * K + k]); + acc2.x += w.x * a; + acc2.y += w.y * a; + } + reinterpret_cast(acc)[m * VecN + n] = acc2; + } + } + } else { +#pragma unroll + for (int m = 0; m < M; ++m) { #pragma unroll - for (int n = 0; n < VecN; ++n) { + for (int n = 0; n < VecN; ++n) { #pragma unroll - for (int k = 0; k < K; ++k) { - reinterpret_cast(acc)[m * VecN + n] = MathWrapper::fma2( - reinterpret_cast(w_pack2)[n * K + k], - MathWrapper::to_vec2(reinterpret_cast(act)[m * K + k]), - reinterpret_cast(acc)[m * VecN + n]); + for (int k = 0; k < K; ++k) { + reinterpret_cast(acc)[m * VecN + n] = MathWrapper::fma2( + reinterpret_cast(w_pack2)[n * K + k], + MathWrapper::to_vec2(reinterpret_cast(act)[m * K + k]), + reinterpret_cast(acc)[m * VecN + n]); + } } } } @@ -180,7 +216,8 @@ __device__ __forceinline__ T warp_reduce_sum(T& val) { return val; } -template +template ::Type> __device__ __forceinline__ void epilogue(void* out, int stride, void* tile_acc, void* bias, float alpha) { using Type = typename MathWrapper::Type; static constexpr int Interleave = Details::kInterleave; @@ -195,7 +232,7 @@ __device__ __forceinline__ void epilogue(void* out, int stride, void* tile_acc, for (int m = 0; m < CtaM; ++m) { #pragma unroll for (int n = 0; n < CtaN; ++n) { - float v = static_cast(reinterpret_cast(tile_acc)[m * CtaN + n]); + float v = static_cast(reinterpret_cast(tile_acc)[m * CtaN + n]); v = warp_reduce_sum(v); if (lane_id < Interleave * ThreadsPerInterleavedTile && lane_id % ThreadsPerInterleavedTile == 0) { shmem[warp_id * CtaM * CtaN * Interleave + m * CtaN * Interleave + n * Interleave + lane_id / ThreadsPerInterleavedTile] = v; diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc index de02de3090671..1aeb99332fa77 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc @@ -7,6 +7,7 @@ #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h" #include +#include #include namespace onnxruntime::llm::kernels::cutlass_kernels { @@ -127,33 +128,86 @@ std::optional MoeGemmProfiler::runProfiling(int maxM, M return best_config; } -void MoeGemmProfiler::profileTactics(CutlassMoeFCRunnerInterface* runner, nvinfer::DataType dtype, +void MoeGemmProfiler::profileTactics(CutlassMoeFCRunnerInterface* runner, weight_only::GemmDims const& dims, MoeGemmId const& gemmId) { ORT_LLM_LOG_ENTRY(); - // Check if already cached - (void)dtype; - auto it = config_cache_.find(gemmId); - if (it != config_cache_.end()) { - return; // Already profiled + + // Profile per M bucket: decode (small M) and prefill (large M) prefer different tile shapes, + // so cache a separate best config for each bucket instead of a single shape-only config. + int const bucket = bucketM(dims.maxM); + auto& bucket_map = config_cache_[gemmId]; + if (bucket_map.find(bucket) != bucket_map.end()) { + return; // Already profiled for this (GemmId, M bucket). } // Initialize backend with correct types initBackend(runner, gemmId); - // Run profiling - int maxM = static_cast(dims.maxM); - auto result = runProfiling(maxM, gemmId); + // Run profiling at the bucket's representative M. + auto result = runProfiling(bucket, gemmId); + + // Cache result for this bucket + bucket_map[bucket] = result; +} - // Cache result - config_cache_[gemmId] = result; +int MoeGemmProfiler::bucketM(int64_t m) { + // Snap M up to the next power of two so a handful of buckets cover the full range. M=1 (the + // common batch-1 decode case) gets its own bucket and therefore its own decode-tuned tactic. + if (m <= 1) { + return 1; + } + // Saturate large M values into one bucket to keep the cache bounded. + constexpr int64_t kMaxBucket = 8192; + if (m >= kMaxBucket) { + return static_cast(kMaxBucket); + } + int64_t bucket = 1; + while (bucket < m) { + bucket <<= 1; + } + return static_cast(bucket); } std::optional MoeGemmProfiler::getBestConfig(int m, MoeGemmId const& id) const { ORT_LLM_LOG_ENTRY(); - (void)m; // M is already factored into profiling auto it = config_cache_.find(id); - if (it != config_cache_.end()) { - return it->second; + if (it == config_cache_.end()) { + return std::nullopt; + } + auto const& bucket_map = it->second; + int const bucket = bucketM(m); + + // Exact bucket profiled: use it. + auto exact = bucket_map.find(bucket); + if (exact != bucket_map.end()) { + return exact->second; + } + + // Not profiled for this exact bucket. Fall back to the nearest profiled bucket: prefer the + // smallest profiled bucket >= requested (tuned for at least this much work), otherwise the + // largest profiled bucket below it. + std::optional best_ge; + int best_ge_bucket = std::numeric_limits::max(); + std::optional best_lt; + int best_lt_bucket = -1; + for (auto const& kv : bucket_map) { + if (kv.first >= bucket) { + if (kv.first < best_ge_bucket) { + best_ge_bucket = kv.first; + best_ge = kv.second; + } + } else { + if (kv.first > best_lt_bucket) { + best_lt_bucket = kv.first; + best_lt = kv.second; + } + } + } + if (best_ge_bucket != std::numeric_limits::max()) { + return best_ge; + } + if (best_lt_bucket >= 0) { + return best_lt; } return std::nullopt; } diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h index 0d73a6e1dfacd..705bae181dca4 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h @@ -90,13 +90,22 @@ class MoeGemmProfiler { sm_ = sm; } - // Profile tactics for a GEMM problem using GemmProfilerBackend - void profileTactics(CutlassMoeFCRunnerInterface* runner, onnxruntime::llm::nvinfer::DataType dtype, + // Profile tactics for a GEMM problem using GemmProfilerBackend. + // Profiles (and caches) the best config for the M bucket that contains dims.maxM. The first + // call for a new (GemmId, M-bucket) pair runs the profiler; subsequent calls return immediately. + // The data/weight types are taken from gemmId, so no separate dtype argument is needed. + void profileTactics(CutlassMoeFCRunnerInterface* runner, weight_only::GemmDims const& dims, MoeGemmId const& gemmId); - // Get best config for a given M and GemmId + // Get best config for a given M and GemmId. Selects the config profiled for the M bucket that + // contains m, so small-M (decode) GEMMs use a decode-tuned tile instead of a prefill-tuned one. std::optional getBestConfig(int m, MoeGemmId const& id) const; + // Snap a row count M to a representative profiling bucket. Decode (small M) and prefill + // (large M) favor very different CUTLASS tile shapes, so we keep a separate best config per + // bucket rather than reusing a single shape-only config. Buckets are powers of two. + static int bucketM(int64_t m); + private: // Initialize backend for profiling void initBackend(CutlassMoeFCRunnerInterface* runner, MoeGemmId const& gemmId); @@ -108,8 +117,8 @@ class MoeGemmProfiler { GemmProfilerBackend backend_; CutlassMoeFCRunnerInterface* runner_{nullptr}; - // Cached results: (M, GemmId) -> best config - mutable std::unordered_map, MoeGemmIdHash> config_cache_; + // Cached results: GemmId -> (M bucket -> best config) + mutable std::unordered_map>, MoeGemmIdHash> config_cache_; // Profiler parameters int num_experts_{0}; diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu new file mode 100644 index 0000000000000..6c198e0765989 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu @@ -0,0 +1,610 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv.h" + +#include +#include +#include + +#include "core/common/common.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h" + +namespace onnxruntime::llm { +namespace kernels { +namespace fpA_intB_gemv { + +// MoE batched GEMV kernel. One thread-block (CtaM = 1 row) processes a single +// expanded row; the row's expert determines the weight/scale/bias base pointers. +// Mirrors the dense fpA_intB_gemv `kernel<>` body for GroupSize=0 (per-channel), +// EnableActScale=false, EnableZero=false, ApplyAlphaInAdvance=false. +template +__global__ void moe_gemv_kernel(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, + int num_experts, + int64_t weight_expert_stride, int64_t scale_expert_stride, int n, int k) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)) + using AccessTypeA = typename Details::AccessTypeA; + using AccessTypeW = typename Details::AccessTypeW; + + static constexpr bool Mandatory = true; + static constexpr int CtaM = 1; + static constexpr int StepK = Details::kStepK; + static constexpr int CtaK = StepK * Threads; + static_assert(CtaN % 2 == 0); + if constexpr (GroupSize != 0) { + static_assert((CtaK / Details::kInterleave) % GroupSize == 0); + } + + int const row = blockIdx.x; + + int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; + // Fallback path for prologues that have not materialized the row-to-expert map. +#pragma unroll 1 + for (int e = 0; e < num_experts && permuted_row_to_expert == nullptr; ++e) { + if (row >= static_cast(expert_first_token_offset[e + 1])) { + expert = e + 1; + continue; + } + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + weight += expert * weight_expert_stride; + scales += static_cast(expert) * scale_expert_stride; + if constexpr (EnableBias) { + bias += static_cast(expert) * n; + } + + int const origin_k = k, interleaved_k = k * Details::kInterleave; + + int const tile_id_m = row, tile_id_n = blockIdx.y, tid = threadIdx.x; + int const offset_m = tile_id_m * CtaM, interleaved_offset_n = tile_id_n * CtaN; + int const real_offset_n = interleaved_offset_n * Details::kInterleave + + ((tid * StepK / Details::LayoutDetails::kTileSize) % Details::kInterleave); + int const real_offset_k = + (tid * StepK / (Details::kInterleave * Details::LayoutDetails::kTileSize)) * Details::LayoutDetails::kTileSize + + ((tid * StepK) % Details::LayoutDetails::kTileSize); + + GMemIterator act_iterator( + act, offset_m * origin_k + real_offset_k, CtaK / Details::kInterleave, origin_k); + GMemIterator weight_iterator( + weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, + CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); + GMemIterator scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + + out += offset_m * n + tile_id_n * CtaN * Details::kInterleave; + if constexpr (EnableBias) { + bias += tile_id_n * CtaN * Details::kInterleave; + } + + AccT tile_acc[CtaM * CtaN]; + fill(tile_acc, static_cast(0.f)); + + TypeA vec_scale[CtaN]; + if constexpr (GroupSize == 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, 0, i); + } + } + + for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { + TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; + uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + if constexpr (GroupSize != 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, iter, i); + } + } +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + weight_iterator.load(tile_w_quantized, iter, i); + dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); + pack_to_vec2(tile_w_pack2, tile_w, i); + } +#pragma unroll + for (int i = 0; i < CtaM; ++i) { + act_iterator.load(tile_a, iter, i); + mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + } + } + epilogue(out, n, tile_acc, bias, 1.0f); +#endif +} + +template +__device__ __forceinline__ void swiglu_epilogue(void* out, void* tile_acc, void* bias, + cutlass_kernels::ActivationParams activation_params) { + static constexpr int Interleave = Details::kInterleave; + static constexpr int ThreadsPerInterleavedTile = Details::kThreadsPerInterleavedTile; + static constexpr int WarpSize = Details::kWarpSize; + static constexpr int WarpNum = Threads / WarpSize; + static constexpr int RawCols = CtaN * Interleave; + static_assert(CtaM == 1); + static_assert(RawCols % 2 == 0); + static_assert(Threads % WarpSize == 0); + + __shared__ float shmem[CtaM * CtaN * Interleave * WarpNum]; + int tid = threadIdx.x; + int warp_id = tid / WarpSize, lane_id = tid % WarpSize; +#pragma unroll + for (int n = 0; n < CtaN; ++n) { + float v = static_cast(reinterpret_cast(tile_acc)[n]); + v = warp_reduce_sum(v); + if (lane_id < Interleave * ThreadsPerInterleavedTile && lane_id % ThreadsPerInterleavedTile == 0) { + shmem[warp_id * RawCols + n * Interleave + lane_id / ThreadsPerInterleavedTile] = v; + } + } + __syncthreads(); + +#pragma unroll + for (int pair = tid; pair < RawCols / 2; pair += Threads) { + int const gate_idx = pair * 2; + int const linear_idx = gate_idx + 1; + float gate = 0.f; + float linear = 0.f; +#pragma unroll + for (int warp = 0; warp < WarpNum; ++warp) { + gate += shmem[warp * RawCols + gate_idx]; + linear += shmem[warp * RawCols + linear_idx]; + } + if constexpr (EnableBias) { + gate += static_cast(reinterpret_cast(bias)[gate_idx]); + linear += static_cast(reinterpret_cast(bias)[linear_idx]); + } + if (isfinite(activation_params.limit)) { + gate = fminf(gate, activation_params.limit); + linear = fminf(fmaxf(linear, -activation_params.limit), activation_params.limit); + } + linear += activation_params.beta; + float const sigmoid = 1.0f / (1.0f + expf(-activation_params.alpha * gate)); + reinterpret_cast(out)[pair] = static_cast(gate * sigmoid * linear); + } +} + +template +__global__ void moe_gemv_interleaved_swiglu_kernel( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t weight_expert_stride, int64_t scale_expert_stride, int inter_size, int k, + cutlass_kernels::ActivationParams activation_params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)) + using AccessTypeA = typename Details::AccessTypeA; + using AccessTypeW = typename Details::AccessTypeW; + + static constexpr bool Mandatory = true; + static constexpr int CtaM = 1; + static constexpr int StepK = Details::kStepK; + static constexpr int CtaK = StepK * Threads; + static_assert(CtaN % 2 == 0); + if constexpr (GroupSize != 0) { + static_assert((CtaK / Details::kInterleave) % GroupSize == 0); + } + + int const row = blockIdx.x; + + int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; +#pragma unroll 1 + for (int e = 0; e < num_experts && permuted_row_to_expert == nullptr; ++e) { + if (row >= static_cast(expert_first_token_offset[e + 1])) { + expert = e + 1; + continue; + } + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + float const* alpha = activation_params.swiglu_alpha; + float const* beta = activation_params.swiglu_beta; + float const* limit = activation_params.swiglu_limit; + activation_params.alpha = alpha ? alpha[expert] : activation_params.alpha; + activation_params.beta = beta ? beta[expert] : activation_params.beta; + activation_params.limit = limit ? limit[expert] : activation_params.limit; + + int const n = inter_size * 2; + weight += expert * weight_expert_stride; + scales += static_cast(expert) * scale_expert_stride; + if constexpr (EnableBias) { + bias += static_cast(expert) * n; + } + + int const origin_k = k, interleaved_k = k * Details::kInterleave; + + int const tile_id_m = row, tile_id_n = blockIdx.y, tid = threadIdx.x; + int const offset_m = tile_id_m * CtaM, interleaved_offset_n = tile_id_n * CtaN; + int const real_offset_n = interleaved_offset_n * Details::kInterleave + + ((tid * StepK / Details::LayoutDetails::kTileSize) % Details::kInterleave); + int const real_offset_k = + (tid * StepK / (Details::kInterleave * Details::LayoutDetails::kTileSize)) * Details::LayoutDetails::kTileSize + + ((tid * StepK) % Details::LayoutDetails::kTileSize); + + GMemIterator act_iterator( + act, offset_m * origin_k + real_offset_k, CtaK / Details::kInterleave, origin_k); + GMemIterator weight_iterator( + weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, + CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); + GMemIterator scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + + out += offset_m * inter_size + tile_id_n * CtaN * Details::kInterleave / 2; + if constexpr (EnableBias) { + bias += tile_id_n * CtaN * Details::kInterleave; + } + + AccT tile_acc[CtaM * CtaN]; + fill(tile_acc, static_cast(0.f)); + + TypeA vec_scale[CtaN]; + if constexpr (GroupSize == 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, 0, i); + } + } + + for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { + TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; + uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + if constexpr (GroupSize != 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, iter, i); + } + } +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + weight_iterator.load(tile_w_quantized, iter, i); + dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); + pack_to_vec2(tile_w_pack2, tile_w, i); + } +#pragma unroll + for (int i = 0; i < CtaM; ++i) { + act_iterator.load(tile_a, iter, i); + mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + } + } + swiglu_epilogue(out, tile_acc, bias, activation_params); +#endif +} + +template +static void launch_moe_gemv(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, + int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, + cudaStream_t stream) { + int64_t const weight_expert_stride = n * k / Details::kElemsPerByteW; + int64_t const scale_expert_stride = GroupSize == 0 ? n : ((k + GroupSize - 1) / GroupSize) * n; + dim3 grid(static_cast(expanded_num_rows), static_cast(n / (CtaN * Details::kInterleave))); + dim3 block(Threads); + if (bias != nullptr) { + moe_gemv_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k)); + } else { + moe_gemv_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k)); + } +} + +template +static void launch_moe_gemv_interleaved_swiglu( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + int64_t const n = inter_size * 2; + int64_t const weight_expert_stride = n * k / Details::kElemsPerByteW; + int64_t const scale_expert_stride = GroupSize == 0 ? n : ((k + GroupSize - 1) / GroupSize) * n; + dim3 grid(static_cast(expanded_num_rows), static_cast(n / (CtaN * Details::kInterleave))); + dim3 block(Threads); + if (bias != nullptr) { + moe_gemv_interleaved_swiglu_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(inter_size), static_cast(k), activation_params); + } else { + moe_gemv_interleaved_swiglu_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(inter_size), static_cast(k), activation_params); + } +} + +template +static void dispatch_moe_gemv_group_size(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, + int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t n, int64_t k, + int group_size, cudaStream_t stream) { + if (group_size <= 0) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 32) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 64) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 128) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else { + ORT_THROW("unsupported MoE GEMV group_size: ", group_size); + } +} + +template +static void dispatch_moe_gemv_interleaved_swiglu_group_size( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int group_size, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + if (group_size <= 0) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 32) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 64) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 128) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else { + ORT_THROW("unsupported MoE GEMV group_size: ", group_size); + } +} + +} // namespace fpA_intB_gemv + +namespace moe_gemv { + +namespace fiv = onnxruntime::llm::kernels::fpA_intB_gemv; + +// CtaN/Threads match the dense per-channel (EnableZero=false) m-tile config. +static constexpr int kCtaN = 8; +static constexpr int kThreads = 128; +// int4 ColumnMajorInterleave (Sm80) tile width along N. +static constexpr int kTileSizeK = 64; +static constexpr int kInt4Interleave = 128 * 8 / (kTileSizeK * 4); // = 4 +static constexpr int kInt8Interleave = 128 * 8 / (kTileSizeK * 8); // = 2 + +// Maps a runtime accumulator-precision choice to a compile-time type tag so the +// host launcher can select the fp32- or 16-bit-accumulation kernel instantiation. +template +struct TypeTag { + using type = T; +}; + +// Opt-in: accumulate the GEMV inner product in 16-bit (fp16) instead of the default +// fp32. Honored only for fp16 activations (bf16 always accumulates in fp32). Set +// ORT_MOE_GEMV_FP16_ACCUM=1 to measure the perf/accuracy tradeoff of 16-bit accumulation. +inline bool MoeGemvUseFp16Accum() { + static bool const enabled = []() { + char const* v = std::getenv("ORT_MOE_GEMV_FP16_ACCUM"); + return v != nullptr && v[0] == '1'; + }(); + return enabled; +} + +bool is_moe_gemv_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, + int weight_bits, int group_size) { + if (sm < 80) { + return false; + } + if (weight_bits != 4 && weight_bits != 8) { + return false; + } + // group_size <= 0 selects the per-column (per-channel) path; block-wise scales must be 32, 64, or 128. + if (group_size > 0 && group_size != 32 && group_size != 64 && group_size != 128) { + return false; + } + // Keep the first block-wise GEMV implementation on complete K blocks. + if (group_size > 0 && k % group_size != 0) { + return false; + } + if (expanded_num_rows <= 0 || expanded_num_rows > kMaxProfiledExpandedRows) { + return false; + } + if (n < kMinProfiledProblemDim || k < kMinProfiledProblemDim) { + return false; + } + if (expanded_num_rows > kMaxProfiledExpandedRowsForSmallProblemDim && + (n < kMinProfiledProblemDimForExpandedRowsAbove4 || k < kMinProfiledProblemDimForExpandedRowsAbove4)) { + return false; + } + // n must tile evenly; k must tile evenly into StepK along interleaved-K. + int const interleave = weight_bits == 4 ? kInt4Interleave : kInt8Interleave; + if (n % (kCtaN * interleave) != 0) { + return false; + } + int64_t const interleaved_k = k * interleave; + int const step_k = 128 / weight_bits; + if (interleaved_k % step_k != 0) { + return false; + } + return true; +} + +bool is_moe_gemv_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k) { + return is_moe_gemv_supported(sm, expanded_num_rows, n, k, 4, 0); +} + +template +struct DetailsForTAndWeight; + +template <> +struct DetailsForTAndWeight { + using Details = fiv::KernelDetails; + using TypeA = half; + static constexpr int kWeightBits = 4; +}; + +template <> +struct DetailsForTAndWeight { + using Details = fiv::KernelDetails; + using TypeA = half; + static constexpr int kWeightBits = 8; +}; + +#ifdef ENABLE_BF16 +template <> +struct DetailsForTAndWeight<__nv_bfloat16, cutlass::uint4b_t> { + using Details = fiv::KernelDetails; + using TypeA = __nv_bfloat16; + static constexpr int kWeightBits = 4; +}; + +template <> +struct DetailsForTAndWeight<__nv_bfloat16, uint8_t> { + using Details = fiv::KernelDetails; + using TypeA = __nv_bfloat16; + static constexpr int kWeightBits = 8; +}; +#endif + +template +void launch_moe_gemv_int_symmetric(T const* act, WeightType const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, + int num_experts, + int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, int sm, + cudaStream_t stream) { + ORT_UNUSED_PARAMETER(sm); + using Details = typename DetailsForTAndWeight::Details; + using TypeA = typename DetailsForTAndWeight::TypeA; + // Accumulate in fp32 by default. fp16 activations may opt back into 16-bit accumulation + // via ORT_MOE_GEMV_FP16_ACCUM=1; bf16 always accumulates in fp32 (16-bit bf16 accumulation + // is too lossy). use_fp32_accum selects the kernel's AccT at runtime. + bool const use_fp32_accum = !std::is_same_v || !MoeGemvUseFp16Accum(); + auto launch = [&](auto acc_tag) { + using AccT = typename decltype(acc_tag)::type; + fiv::dispatch_moe_gemv_group_size( + const_cast(reinterpret_cast(act)), + const_cast(reinterpret_cast(weight)), + const_cast(reinterpret_cast(scales)), + const_cast(reinterpret_cast(bias)), + reinterpret_cast(out), + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, group_size, stream); + }; + if (use_fp32_accum) { + launch(TypeTag{}); + } else { + launch(TypeTag{}); + } +} + +template +void launch_moe_gemv_int_symmetric_interleaved_swiglu( + T const* act, WeightType const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int group_size, int sm, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + ORT_UNUSED_PARAMETER(sm); + using Details = typename DetailsForTAndWeight::Details; + using TypeA = typename DetailsForTAndWeight::TypeA; + // Accumulate in fp32 by default (see launch_moe_gemv_int_symmetric for the policy). + bool const use_fp32_accum = !std::is_same_v || !MoeGemvUseFp16Accum(); + auto launch = [&](auto acc_tag) { + using AccT = typename decltype(acc_tag)::type; + fiv::dispatch_moe_gemv_interleaved_swiglu_group_size( + const_cast(reinterpret_cast(act)), + const_cast(reinterpret_cast(weight)), + const_cast(reinterpret_cast(scales)), + const_cast(reinterpret_cast(bias)), + reinterpret_cast(out), + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size, k, group_size, + activation_params, stream); + }; + if (use_fp32_accum) { + launch(TypeTag{}); + } else { + launch(TypeTag{}); + } +} + +template +void launch_moe_gemv_int4_per_channel(T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, + int num_experts, + int64_t expanded_num_rows, int64_t n, int64_t k, int sm, cudaStream_t stream) { + launch_moe_gemv_int_symmetric( + act, reinterpret_cast(weight), scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, 0, sm, stream); +} + +template +void launch_moe_gemv_int4_per_channel_interleaved_swiglu( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int sm, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + launch_moe_gemv_int_symmetric_interleaved_swiglu( + act, reinterpret_cast(weight), scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, inter_size, k, 0, sm, activation_params, stream); +} + +template void launch_moe_gemv_int_symmetric( + half const*, cutlass::uint4b_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric( + half const*, uint8_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric_interleaved_swiglu( + half const*, cutlass::uint4b_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, cudaStream_t); +template void launch_moe_gemv_int_symmetric_interleaved_swiglu( + half const*, uint8_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, cudaStream_t); + +template void launch_moe_gemv_int4_per_channel(half const*, uint8_t const*, half const*, half const*, half*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, + int, cudaStream_t); +template void launch_moe_gemv_int4_per_channel_interleaved_swiglu( + half const*, uint8_t const*, half const*, half const*, half*, int64_t const*, int const*, int, int64_t, + int64_t, int64_t, int, cutlass_kernels::ActivationParams, cudaStream_t); + +#ifdef ENABLE_BF16 +template void launch_moe_gemv_int_symmetric<__nv_bfloat16, cutlass::uint4b_t>( + __nv_bfloat16 const*, cutlass::uint4b_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric<__nv_bfloat16, uint8_t>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric_interleaved_swiglu<__nv_bfloat16, cutlass::uint4b_t>( + __nv_bfloat16 const*, cutlass::uint4b_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, + cudaStream_t); +template void launch_moe_gemv_int_symmetric_interleaved_swiglu<__nv_bfloat16, uint8_t>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, + cudaStream_t); + +template void launch_moe_gemv_int4_per_channel<__nv_bfloat16>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, cudaStream_t); +template void launch_moe_gemv_int4_per_channel_interleaved_swiglu<__nv_bfloat16>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, cutlass_kernels::ActivationParams, + cudaStream_t); +#endif +} // namespace moe_gemv +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h new file mode 100644 index 0000000000000..b4dfe1c59f02a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Batched GEMV fast path for symmetric int weight-only MoE at small expanded row +// counts (e.g. batch-1 decode with top_k experts). Each expanded row is a single +// token-expert pair; one thread-block handles one row and offsets the +// weight/scale/bias pointers by that row's expert. Reuses the device-side math +// (layout, dequantize, mma, epilogue) from the dense fpA_intB_gemv kernel. + +#pragma once + +#include +#include + +#include "contrib_ops/cuda/llm/moe_gemm/common.h" + +namespace onnxruntime::llm { +namespace kernels { +namespace moe_gemv { + +inline constexpr int64_t kMaxProfiledExpandedRows = 8; +inline constexpr int64_t kMaxProfiledExpandedRowsForSmallProblemDim = 4; +inline constexpr int64_t kMinProfiledProblemDim = 512; +// Lowered from 704 to 512 so block-wise decode shapes (e.g. Qwen top_k=8, +// inter_size=512) take the GEMV path. This also covers per-column INT4 shapes +// with inter_size in [512, 704); both bands are gated by ORT_DISABLE_MOE_GEMV. +inline constexpr int64_t kMinProfiledProblemDimForExpandedRowsAbove4 = 512; + +// Returns true if the batched MoE GEMV fast path supports this problem shape. +// Requirements: FP16/BF16 activations, sm >= 80, small expanded_num_rows, supported +// INT weight type, supported group size, and n divisible by the kernel tile width. +bool is_moe_gemv_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, + int weight_bits, int group_size); + +// Backward-compatible per-channel INT4 shape check. +bool is_moe_gemv_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k); + +// Launches symmetric INT MoE GEMV. group_size <= 0 means per-channel scales; +// group_size 32/64/128 means block-wise scales laid out as [num_experts, k_blocks, n]. +// T is half or __nv_bfloat16. WeightType is cutlass::uint4b_t or uint8_t. +template +void launch_moe_gemv_int_symmetric( + T const* act, WeightType const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t n, int64_t k, int group_size, int sm, cudaStream_t stream); + +// Launches symmetric INT MoE GEMV and fuses interleaved SwiGLU activation. +// weight/bias use raw FC1 output width n = 2 * inter_size. Scales are +// [num_experts, n] for group_size <= 0 and [num_experts, k_blocks, n] for +// block-wise group_size 32/64/128. +template +void launch_moe_gemv_int_symmetric_interleaved_swiglu( + T const* act, WeightType const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t inter_size, int64_t k, int group_size, int sm, cutlass_kernels::ActivationParams activation_params, + cudaStream_t stream); + +// Launches the int4 per-channel MoE GEMV. +// act: [expanded_num_rows, k] permuted activations (row-major) +// weight: [num_experts, k, n] packed int4 in Sm80 ColumnMajorInterleave layout (uint8) +// scales: [num_experts, n] per-channel scales (T) +// bias: [num_experts, n] per-expert bias (T) or nullptr +// out: [expanded_num_rows, n] (row-major) +// expert_first_token_offset: [num_experts + 1] prefix offsets of permuted rows +// permuted_row_to_expert: [expanded_num_rows] local expert id for each permuted row, or nullptr to scan offsets +// T is half or __nv_bfloat16. +template +void launch_moe_gemv_int4_per_channel( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t n, int64_t k, int sm, cudaStream_t stream); + +// Launches the int4 per-channel MoE GEMV and fuses interleaved SwiGLU activation. +// weight/scales/bias use raw FC1 output width [num_experts, k, 2 * inter_size] +// out is post-activation [expanded_num_rows, inter_size] +// Only interleaved SwiGLU layout (`swiglu_fusion == 1`) is supported. +template +void launch_moe_gemv_int4_per_channel_interleaved_swiglu( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t inter_size, int64_t k, int sm, cutlass_kernels::ActivationParams activation_params, + cudaStream_t stream); + +} // namespace moe_gemv +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index ab0e2d9e01901..b0b6d66eca4c5 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu @@ -15,6 +15,7 @@ */ #include +#include #include #include #include @@ -60,6 +61,7 @@ #include "contrib_ops/cuda/llm/kernels/quantization.cuh" #include "contrib_ops/cuda/llm/moe_gemm/common.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_kernels.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh" #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_utils.cuh" @@ -71,6 +73,134 @@ using namespace onnxruntime::llm::kernels; using namespace onnxruntime::llm::common; namespace onnxruntime::llm::kernels::cutlass_kernels { + +// Master switch for the symmetric INT MoE GEMV fast path. Enabled by default; +// set ORT_DISABLE_MOE_GEMV=1 to fall back to the CUTLASS grouped GEMM. +inline bool MoeGemvDisabledByEnv() { + static bool const disabled = []() { + char const* v = std::getenv("ORT_DISABLE_MOE_GEMV"); + return v != nullptr && v[0] == '1'; + }(); + return disabled; +} + +inline bool MoeGemvRejectedByProfiledInterSize(int64_t expanded_num_rows, int64_t inter_size) { + return expanded_num_rows > onnxruntime::llm::kernels::moe_gemv::kMaxProfiledExpandedRowsForSmallProblemDim && + inter_size < onnxruntime::llm::kernels::moe_gemv::kMinProfiledProblemDimForExpandedRowsAbove4; +} + +template +constexpr int MoeGemvWeightBits() { + if constexpr (std::is_same_v) { + return 4; + } else if constexpr (std::is_same_v) { + return 8; + } else { + return 0; + } +} + +// Attempts the batched symmetric INT MoE GEMV. Returns true if it ran (output +// written), false if the configuration is unsupported and the caller must fall +// back to the grouped GEMM. Compiles to a no-op (returns false) for type +// combinations other than (half activations, int4/int8 weights, ScaleBias==T). +template +bool tryLaunchMoeGemvIntSymmetric(T const* input, WeightType const* weights, ScaleBiasType const* scales, + ScaleBiasType const* weight_zeros, ScaleBiasType const* biases, T* output, + int64_t const* expert_first_token_offset, int num_experts_per_node, + int const* permuted_row_to_expert, int64_t expanded_num_rows, + int64_t n, int64_t k, int sm, int group_size, + bool disabled, cudaStream_t stream) { + if constexpr ((std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && std::is_same_v) { + bool const env_disabled = MoeGemvDisabledByEnv(); + bool const has_block_zeros = group_size > 0 && weight_zeros != nullptr; + constexpr int weight_bits = MoeGemvWeightBits(); + if (disabled || env_disabled || has_block_zeros) { + return false; + } + if (!onnxruntime::llm::kernels::moe_gemv::is_moe_gemv_supported( + sm, expanded_num_rows, n, k, weight_bits, group_size)) { + return false; + } + onnxruntime::llm::kernels::moe_gemv::launch_moe_gemv_int_symmetric( + input, weights, scales, biases, output, expert_first_token_offset, permuted_row_to_expert, + num_experts_per_node, expanded_num_rows, n, k, group_size, sm, stream); + return true; + } else { + (void)input; + (void)weights; + (void)scales; + (void)weight_zeros; + (void)biases; + (void)output; + (void)expert_first_token_offset; + (void)num_experts_per_node; + (void)permuted_row_to_expert; + (void)expanded_num_rows; + (void)n; + (void)k; + (void)sm; + (void)group_size; + (void)disabled; + (void)stream; + return false; + } +} + +template +bool tryLaunchMoeGemvIntSymmetricInterleavedSwiGLU( + T const* input, WeightType const* weights, ScaleBiasType const* scales, ScaleBiasType const* weight_zeros, + ScaleBiasType const* biases, T* output, + int64_t const* expert_first_token_offset, int num_experts_per_node, int const* permuted_row_to_expert, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int sm, int group_size, + bool disabled, cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + if constexpr ((std::is_same_v || std::is_same_v) && + (std::is_same_v || std::is_same_v) && std::is_same_v) { + bool const env_disabled = MoeGemvDisabledByEnv(); + bool const has_block_zeros = group_size > 0 && weight_zeros != nullptr; + if (disabled || env_disabled || has_block_zeros) { + return false; + } + int64_t const n = inter_size * 2; + constexpr int weight_bits = MoeGemvWeightBits(); + if (activation_params.swiglu_fusion != 1) { + return false; + } + if (activation_params.activation_type != ActivationType::Swiglu && + activation_params.activation_type != ActivationType::SwigluBias) { + return false; + } + if (!onnxruntime::llm::kernels::moe_gemv::is_moe_gemv_supported( + sm, expanded_num_rows, n, k, weight_bits, group_size)) { + return false; + } + onnxruntime::llm::kernels::moe_gemv::launch_moe_gemv_int_symmetric_interleaved_swiglu( + input, weights, scales, biases, output, expert_first_token_offset, permuted_row_to_expert, + num_experts_per_node, expanded_num_rows, inter_size, k, group_size, sm, activation_params, stream); + return true; + } else { + (void)input; + (void)weights; + (void)scales; + (void)weight_zeros; + (void)biases; + (void)output; + (void)expert_first_token_offset; + (void)num_experts_per_node; + (void)permuted_row_to_expert; + (void)expanded_num_rows; + (void)inter_size; + (void)k; + (void)sm; + (void)group_size; + (void)disabled; + (void)activation_params; + (void)stream; + return false; + } +} + /** * Takes the input maps and prepares the expanded maps for min latency * @param num_active_experts_per_node: Number of active experts on current node @@ -272,6 +402,7 @@ void buildMinLatencyActiveExpertMaps(int* num_active_experts_per_node, float* ex template __global__ void fusedBuildExpertMapsSortFirstTokenKernel(int const* const token_selected_experts, int* const permuted_row_to_unpermuted_row, int* const unpermuted_row_to_permuted_row, + int* const permuted_token_selected_experts, int64_t* const expert_first_token_offset, int64_t const num_tokens, int const experts_per_token, int const start_expert, int const end_expert, int const num_experts_per_node) { // Only using block wise collective so we can only have one block @@ -342,6 +473,7 @@ __global__ void fusedBuildExpertMapsSortFirstTokenKernel(int const* const token_ int const permuted_row = local_token_permuted_indices[i]; permuted_row_to_unpermuted_row[permuted_row] = unpermuted_row; unpermuted_row_to_permuted_row[unpermuted_row] = permuted_row; + permuted_token_selected_experts[permuted_row] = local_token_selected_experts[i]; } } @@ -356,7 +488,8 @@ __global__ void fusedBuildExpertMapsSortFirstTokenKernel(int const* const token_ template bool fusedBuildExpertMapsSortFirstTokenDispatch(int const* token_selected_experts, int* permuted_row_to_unpermuted_row, - int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, int64_t const num_tokens, + int* unpermuted_row_to_permuted_row, int* permuted_token_selected_experts, + int64_t* expert_first_token_offset, int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream) { ORT_ENFORCE(num_experts_per_node == (end_expert - start_expert), @@ -394,15 +527,17 @@ bool fusedBuildExpertMapsSortFirstTokenDispatch(int const* token_selected_expert CUDA_CALL_THROW(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_size)); CUDA_CALL_THROW(cudaLaunchKernelEx(&config, kernel, token_selected_experts, permuted_row_to_unpermuted_row, - unpermuted_row_to_permuted_row, expert_first_token_offset, num_tokens, experts_per_token, start_expert, - end_expert, num_experts_per_node)); + unpermuted_row_to_permuted_row, permuted_token_selected_experts, + expert_first_token_offset, num_tokens, experts_per_token, start_expert, end_expert, + num_experts_per_node)); return true; } template bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_experts, int* permuted_row_to_unpermuted_row, - int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, int64_t const num_tokens, + int* unpermuted_row_to_permuted_row, int* permuted_token_selected_experts, + int64_t* expert_first_token_offset, int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream) { int const block_size = num_tokens; @@ -421,13 +556,14 @@ bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_exper } return func(token_selected_experts, permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, - expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, start_expert, end_expert, - stream); + permuted_token_selected_experts, expert_first_token_offset, num_tokens, num_experts_per_node, + experts_per_token, start_expert, end_expert, stream); } template bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_experts, int* permuted_row_to_unpermuted_row, - int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, int64_t const num_tokens, + int* unpermuted_row_to_permuted_row, int* permuted_token_selected_experts, + int64_t* expert_first_token_offset, int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream) { auto func = &fusedBuildExpertMapsSortFirstTokenBlockSize<1, LOG2_NUM_EXPERTS>; @@ -458,12 +594,13 @@ bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_exper } } return func(token_selected_experts, permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, - expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, start_expert, end_expert, - stream); + permuted_token_selected_experts, expert_first_token_offset, num_tokens, num_experts_per_node, + experts_per_token, start_expert, end_expert, stream); } bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* permuted_row_to_unpermuted_row, - int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, int64_t const num_tokens, + int* unpermuted_row_to_permuted_row, int* permuted_token_selected_experts, + int64_t* expert_first_token_offset, int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream) { // We need enough bits to represent [0, num_experts_per_node+1] (inclusive) i.e. num_experts_per_node + 2 values @@ -477,8 +614,9 @@ bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* &fusedBuildExpertMapsSortFirstTokenBlockSize<8>, &fusedBuildExpertMapsSortFirstTokenBlockSize<9>}; return funcs[expert_log - 1](token_selected_experts, permuted_row_to_unpermuted_row, - unpermuted_row_to_permuted_row, expert_first_token_offset, num_tokens, num_experts_per_node, - experts_per_token, start_expert, end_expert, stream); + unpermuted_row_to_permuted_row, permuted_token_selected_experts, + expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, + start_expert, end_expert, stream); } ORT_LLM_LOG_DEBUG(onnxruntime::MakeString("Experts per node ", num_experts_per_node, " does not have supported fused moe prologues")); return false; @@ -1455,7 +1593,6 @@ __global__ void finalizeMoeRoutingKernel(GemmOutputType const* expanded_permuted OutputType* reduced_unpermuted_output, ScaleBiasType const* bias, float const* scales, int const* unpermuted_row_to_permuted_row, int const* token_selected_experts, int64_t const orig_cols, int64_t const experts_per_token, int const num_experts_per_node, int const start_expert_id) { - assert(orig_cols % 4 == 0); int64_t const original_row = blockIdx.x; int64_t const num_rows = gridDim.x; auto const offset = original_row * orig_cols; @@ -1463,6 +1600,7 @@ __global__ void finalizeMoeRoutingKernel(GemmOutputType const* expanded_permuted // Load 128-bits per thread, according to the smallest data type we read/write constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / std::min(sizeof_bits::value, sizeof_bits::value); + assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0); int64_t const start_offset = threadIdx.x; int64_t const stride = FINALIZE_THREADS_PER_BLOCK; @@ -1515,6 +1653,77 @@ __global__ void finalizeMoeRoutingKernel(GemmOutputType const* expanded_permuted #endif } +template +__global__ void finalizeMoeRoutingOneRowKernel(GemmOutputType const* expanded_permuted_rows, + OutputType* reduced_unpermuted_output, ScaleBiasType const* bias, float const* scales, + int const* unpermuted_row_to_permuted_row, int const* token_selected_experts, + int64_t const orig_cols, int const num_experts_per_node, + int const start_expert_id) { + static_assert(ExpertsPerToken > 0 && ExpertsPerToken <= 4); + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / std::min(sizeof_bits::value, sizeof_bits::value); + assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0); + + __shared__ int expanded_permuted_rows_for_topk[ExpertsPerToken]; + __shared__ int expert_ids_for_topk[ExpertsPerToken]; + __shared__ float row_scales_for_topk[ExpertsPerToken]; + + int const tid = threadIdx.x; + if (tid < ExpertsPerToken) { + int const expert_id = token_selected_experts[tid] - start_expert_id; + expert_ids_for_topk[tid] = expert_id; + expanded_permuted_rows_for_topk[tid] = + (expert_id >= 0 && expert_id < num_experts_per_node) ? unpermuted_row_to_permuted_row[tid] : -1; + row_scales_for_topk[tid] = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[tid]; + } + __syncthreads(); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + int64_t const start_offset = tid; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD; + + using BiasElem = cutlass::Array; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + auto const* bias_v = reinterpret_cast(bias); + auto const* expanded_permuted_rows_v = reinterpret_cast(expanded_permuted_rows); + auto* reduced_row_ptr_v = reinterpret_cast(reduced_unpermuted_output); + + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { + ComputeElem thread_output; + thread_output.fill(0); +#pragma unroll + for (int k_idx = 0; k_idx < ExpertsPerToken; ++k_idx) { + int const expanded_permuted_row = expanded_permuted_rows_for_topk[k_idx]; + if (expanded_permuted_row < 0) { + continue; + } + + auto const* expanded_permuted_rows_row_ptr = expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col; + + ComputeElem expert_result = arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); + if (bias) { + auto const* bias_ptr = bias_v + expert_ids_for_topk[k_idx] * num_elems_in_col; + expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); + } + + thread_output = thread_output + row_scales_for_topk[k_idx] * expert_result; + } + + OutputElem output_elem = arrayConvert(thread_output); + reduced_row_ptr_v[elem_index] = output_elem; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + // Final kernel to unpermute and scale // This kernel unpermutes the original data, does the k-way reduction and performs the final skip connection. template @@ -1523,8 +1732,6 @@ __global__ void finalizeMoeRoutingNoFillingKernel(GemmOutputType const* expanded int const* const unpermuted_row_to_permuted_row, int const* permuted_row_to_unpermuted_row, int const* token_selected_experts, int64_t const* expert_first_token_offset, int64_t const num_rows, int64_t const orig_cols, int64_t const experts_per_token, int const num_experts_per_node, int const start_expert_id) { - assert(orig_cols % 4 == 0); - #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) asm volatile("griddepcontrol.wait;"); #endif @@ -1556,6 +1763,7 @@ __global__ void finalizeMoeRoutingNoFillingKernel(GemmOutputType const* expanded // Load 128-bits per thread, according to the smallest data type we read/write constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / std::min(sizeof_bits::value, sizeof_bits::value); + assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0); int64_t const start_offset = threadIdx.x; int64_t const stride = FINALIZE_THREADS_PER_BLOCK; @@ -1613,6 +1821,10 @@ void finalizeMoeRoutingKernelLauncher(GemmOutputType const* expanded_permuted_ro // Only add bias on rank 0 for tensor parallelism bool const is_rank_0 = parallelism_config.tp_rank == 0; ScaleBiasType const* bias_ptr = is_rank_0 ? bias : nullptr; + constexpr int64_t kFinalizeElemPerThread = 128 / std::min(sizeof_bits::value, sizeof_bits::value); + ORT_ENFORCE(cols % kFinalizeElemPerThread == 0, + "MoE finalize requires cols to be divisible by ", kFinalizeElemPerThread, + " for vectorized 128-bit loads, got ", cols, "."); int num_experts_per_node_int = SafeInt(num_experts_per_node); int const start_expert_id = num_experts_per_node_int * parallelism_config.ep_rank; @@ -1645,12 +1857,39 @@ void finalizeMoeRoutingKernelLauncher(GemmOutputType const* expanded_permuted_ro int64_t const threads = FINALIZE_THREADS_PER_BLOCK; config.gridDim = blocks; config.blockDim = threads; - auto func = final_scales - ? &finalizeMoeRoutingKernel - : &finalizeMoeRoutingKernel; - cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, - unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, - start_expert_id); + if (num_rows == 1 && experts_per_token > 0 && experts_per_token <= 4) { +#define LAUNCH_FINALIZE_ONE_ROW(EXPERTS_PER_TOKEN) \ + do { \ + auto func = final_scales \ + ? &finalizeMoeRoutingOneRowKernel \ + : &finalizeMoeRoutingOneRowKernel; \ + cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, \ + unpermuted_row_to_permuted_row, token_selected_experts, cols, num_experts_per_node_int, start_expert_id); \ + } while (0) + + switch (experts_per_token) { + case 1: + LAUNCH_FINALIZE_ONE_ROW(1); + break; + case 2: + LAUNCH_FINALIZE_ONE_ROW(2); + break; + case 3: + LAUNCH_FINALIZE_ONE_ROW(3); + break; + default: + LAUNCH_FINALIZE_ONE_ROW(4); + break; + } +#undef LAUNCH_FINALIZE_ONE_ROW + } else { + auto func = final_scales + ? &finalizeMoeRoutingKernel + : &finalizeMoeRoutingKernel; + cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, + start_expert_id); + } } } @@ -1914,10 +2153,10 @@ CutlassMoeFCRunner: // in the case of unfused activation we overlap permuted_data and fc1_result // we need to calculate the max possible size, so use the max of all three size_t overlapped_gemm1_gemm2_inputs_size = std::max(permuted_data_size, fc2_result_size); - // When glu_inter_elems is 0 we are always fused, otherwise we may need the un-fused case - if (glu_inter_elems > 0) { - overlapped_gemm1_gemm2_inputs_size = std::max(overlapped_gemm1_gemm2_inputs_size, fc1_result_size); - } + // In the gated (glu) case fc1_result_ used to alias overlapped_gemm1_gemm2_inputs. That made the + // fused-SwiGLU GEMV read and write the same buffer with mismatched row strides (input hidden_size, + // output inter_size), corrupting any launch that spanned more than one wave. fc1_result_ now uses + // its own dedicated buffer (fc1_result_dedicated below), so inputs no longer needs to fit fc1_result_size. size_t const alpha_scale_ptr_array_size = num_experts_per_node * sizeof(float*); @@ -1927,6 +2166,10 @@ CutlassMoeFCRunner: overlapped_gemm1_gemm2_outputs_size = std::max(std::max(glu_inter_size, fc2_result_size), overlapped_gemm1_gemm2_outputs_size); } + // Dedicated buffer for the gated FC1 result so the fused-SwiGLU GEMV never aliases its own input + // (see comment above). Only the gated / has-glu path uses it; zero otherwise. + size_t const fc1_result_dedicated_size = (glu_inter_elems > 0) ? fc1_result_size : 0; + size_t smoothed_act_size = use_awq ? std::max(permuted_elems, interbuf_elems) * sizeof(T) * 2 : 0; // Extra workspace required by AWQ for smoothing activations @@ -1950,6 +2193,7 @@ CutlassMoeFCRunner: ADD(permuted_token_final_scales); ADD(overlapped_gemm1_gemm2_inputs); ADD(overlapped_gemm1_gemm2_outputs); + ADD(fc1_result_dedicated); ADD_NAME(alpha_scale_ptr_array_fc1, alpha_scale_ptr_array_size); ADD_NAME(alpha_scale_ptr_array_fc2, alpha_scale_ptr_array_size); ADD(fp4_act_scale); @@ -2018,10 +2262,12 @@ void CutlassMoeFCRunner{input, - total_tokens_including_expert, fc1_expert_weights, - /*scales*/ quant_params.groupwise.group_size > 0 - ? static_cast(quant_params.groupwise.fc1.weight_scales) - : fc1_int_scales, - /*zeros*/ quant_params.groupwise.group_size > 0 - ? static_cast(quant_params.groupwise.fc1.weight_zeros) - : nullptr, - fc1_expert_biases, static_cast(use_ampere_activation_fusion ? output : intermediate_result), - alpha_scale_ptr_array, /*occupancy*/ nullptr, - use_ampere_activation_fusion ? fc1_activation_type : ActivationType::Identity, expanded_num_rows, - /*N*/ int64_t(fc1_out_size), - /*K*/ hidden_size, num_experts_per_node, quant_params.groupwise.group_size, bias_is_broadcast, - use_ampere_activation_fusion, stream, activation_params, config}; - gemm_runner.moeGemmBiasAct(universal_input, TmaWarpSpecializedGroupedGemmInput{}); + bool const fc1_did_fused_gemv = tryLaunchMoeGemvIntSymmetricInterleavedSwiGLU( + input, fc1_expert_weights, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_scales) + : fc1_int_scales, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_zeros) + : nullptr, + fc1_expert_biases, output, expert_first_token_offset, num_experts_per_node, + permuted_row_to_expert, expanded_num_rows, inter_size, hidden_size, + onnxruntime::llm::common::getSMVersion(), quant_params.groupwise.group_size, + /*disabled=*/parallelism_config.ep_size > 1 || use_ampere_activation_fusion || !bias_is_broadcast || + MoeGemvRejectedByProfiledInterSize(expanded_num_rows, inter_size), + activation_params, stream); + + // Run the GEMM with activation function overridden with `Identity`, we do the activation separately. + // Fast path: int4 per-channel MoE GEMV for small expanded-row counts (e.g. batch-1 decode). + bool const fc1_did_gemv = fc1_did_fused_gemv || tryLaunchMoeGemvIntSymmetric( + input, fc1_expert_weights, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_scales) + : fc1_int_scales, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_zeros) + : nullptr, + fc1_expert_biases, static_cast(intermediate_result), expert_first_token_offset, num_experts_per_node, + permuted_row_to_expert, expanded_num_rows, /*n=*/static_cast(fc1_out_size), /*k=*/hidden_size, + onnxruntime::llm::common::getSMVersion(), + quant_params.groupwise.group_size, + /*disabled=*/parallelism_config.ep_size > 1 || use_ampere_activation_fusion || !bias_is_broadcast || + MoeGemvRejectedByProfiledInterSize(expanded_num_rows, inter_size), + stream); + if (!fc1_did_gemv) { + auto universal_input = GroupedGemmInput{input, + total_tokens_including_expert, fc1_expert_weights, + /*scales*/ quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_scales) + : fc1_int_scales, + /*zeros*/ quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc1.weight_zeros) + : nullptr, + fc1_expert_biases, static_cast(use_ampere_activation_fusion ? output : intermediate_result), + alpha_scale_ptr_array, /*occupancy*/ nullptr, + use_ampere_activation_fusion ? fc1_activation_type : ActivationType::Identity, expanded_num_rows, + /*N*/ int64_t(fc1_out_size), + /*K*/ hidden_size, num_experts_per_node, quant_params.groupwise.group_size, bias_is_broadcast, + use_ampere_activation_fusion, stream, activation_params, config}; + gemm_runner.moeGemmBiasAct(universal_input, TmaWarpSpecializedGroupedGemmInput{}); + } sync_check_cuda_error(stream); - if (!use_ampere_activation_fusion) { + if (!use_ampere_activation_fusion && !fc1_did_fused_gemv) { using GatedActOutputType = std::conditional_t; if (is_gated_activation) { doGatedActivation( @@ -2259,6 +2540,7 @@ void CutlassMoeFCRunner{input, total_tokens_including_expert, - fc2_expert_weights, - quant_params.groupwise.group_size > 0 - ? static_cast(quant_params.groupwise.fc2.weight_scales) - : fc2_int_scales, - quant_params.groupwise.group_size > 0 - ? static_cast(quant_params.groupwise.fc2.weight_zeros) - : nullptr, - nullptr, static_cast(gemm_output), - alpha_scale_ptr_array, /*occupancy*/ nullptr, ActivationType::Identity, expanded_num_rows, - /*N*/ hidden_size, - /*K*/ inter_size, - num_experts_per_node, - quant_params.groupwise.group_size, - /*bias_is_broadcast*/ false, - /*use_fused_moe*/ false, - stream, - activation_params, - config}; - gemm_runner.moeGemmBiasAct(universal_input, tma_ws_input); + // Note: expanded_num_rows, to check this value, it's greater than num_rows * num_experts_per_node + // Fast path: int4 per-channel MoE GEMV (no bias here; fc2 bias applied in finalizeMoeRouting). + // Keep the decode GEMV path single-EP until token-drop/all-to-all + // cases are profiled and validated. + bool const fc2_did_gemv = tryLaunchMoeGemvIntSymmetric( + input, fc2_expert_weights, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc2.weight_scales) + : fc2_int_scales, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc2.weight_zeros) + : nullptr, + /*biases*/ nullptr, static_cast(gemm_output), expert_first_token_offset, num_experts_per_node, + permuted_row_to_expert, expanded_num_rows, /*n=*/hidden_size, /*k=*/inter_size, onnxruntime::llm::common::getSMVersion(), + quant_params.groupwise.group_size, + /*disabled=*/parallelism_config.ep_size > 1 || using_tma_ws_gemm2 || + MoeGemvRejectedByProfiledInterSize(expanded_num_rows, inter_size), + stream); + if (!fc2_did_gemv) { + auto universal_input = GroupedGemmInput{input, total_tokens_including_expert, + fc2_expert_weights, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc2.weight_scales) + : fc2_int_scales, + quant_params.groupwise.group_size > 0 + ? static_cast(quant_params.groupwise.fc2.weight_zeros) + : nullptr, + nullptr, static_cast(gemm_output), + alpha_scale_ptr_array, /*occupancy*/ nullptr, ActivationType::Identity, expanded_num_rows, + /*N*/ hidden_size, + /*K*/ inter_size, + num_experts_per_node, + quant_params.groupwise.group_size, + /*bias_is_broadcast*/ false, + /*use_fused_moe*/ false, + stream, + activation_params, + config}; + gemm_runner.moeGemmBiasAct(universal_input, tma_ws_input); + } sync_check_cuda_error(stream); bool has_different_output_type_ampere = (use_w4afp8 || use_fp8) && !using_tma_ws_gemm2; @@ -2459,8 +2760,10 @@ void CutlassMoeFCRunner& gemm_runner, @@ -449,6 +451,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, int64_t const expanded_num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts_per_node, int64_t const experts_per_token, float const** alpha_scale_ptr_array, + int const* permuted_row_to_expert, cudaStream_t stream, MOEParallelismConfig parallelism_config, cutlass_extensions::CutlassGemmConfig config); @@ -469,8 +472,8 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { static_cast(fc1_expert_weights), static_cast(fc1_expert_biases), num_valid_tokens_ptr, static_cast(fc1_int_scales), fc1_fp8_dequant, fc2_fp8_quant, fc1_fp4_act_flat, fc2_fp4_act_flat, quant_params, num_rows, expanded_num_rows, hidden_size, inter_size, - num_experts_per_node, fc1_activation_type, alpha_scale_ptr_array, bias_is_broadcast, stream, config, - activation_params); + num_experts_per_node, fc1_activation_type, alpha_scale_ptr_array, nullptr, bias_is_broadcast, stream, + MOEParallelismConfig{}, config, activation_params); } void gemm2(void const* const input, void* const gemm_output, void* const final_output, @@ -491,7 +494,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { static_cast(fc2_int_scales), fc2_fp8_dequant, fc2_fp4_act_flat, quant_params, token_topk_unpermuted_scales, token_topk_permuted_scales, unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row, token_selected_experts, num_valid_tokens_ptr, num_rows, expanded_num_rows, - hidden_size, inter_size, num_experts_per_node, experts_per_token, alpha_scale_ptr_array, + hidden_size, inter_size, num_experts_per_node, experts_per_token, alpha_scale_ptr_array, nullptr, stream, parallelism_config, config); } diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h index 62f69831d3236..bdae2159e9a3c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h @@ -37,11 +37,12 @@ namespace onnxruntime::llm::kernels { namespace cutlass_kernels { -// These kernels are used in moeUtilOp.cpp +// Utility kernels used by the MoE runner to build expert maps, expand rows, and finalize routing. int64_t computeNumTokensPerBlock(int64_t const num_tokens, int64_t const num_experts_per_node); -bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* unpermuted_token_selected_experts, - int* permuted_source_token_ids, int64_t* expert_first_token_offset, int64_t const num_tokens, +bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int* permuted_token_selected_experts, + int64_t* expert_first_token_offset, int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe.cc b/onnxruntime/contrib_ops/cuda/moe/moe.cc index 603ba2b96b81b..4ab84fa082ccf 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe.cc @@ -9,6 +9,8 @@ #include "contrib_ops/cuda/llm/moe_gemm/moe_kernels.h" #include "contrib_ops/cuda/llm/common/env_utils.h" +#include + using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; @@ -17,6 +19,16 @@ namespace onnxruntime { namespace contrib { namespace cuda { +namespace { +void LogSwigluFusionRemapOnce() { + static std::once_flag log_warning; + std::call_once(log_warning, []() { + LOGS_DEFAULT(WARNING) << "MoE swiglu_fusion is 0 with no fc3_experts_weights; assuming interleaved " + "SwiGLU layout for backward compatibility."; + }); +} +} // namespace + #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ MoE, kMSDomain, 1, T, kCudaExecutionProvider, \ @@ -42,8 +54,21 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { const Tensor* fc3_experts_bias_optional = context->Input(7); using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; + + // Backward compatibility: the published gpt-oss-20b model (and any model exported by ORT < 1.27) + // hard-coded the interleaved SwiGLU fusion layout and did not emit a swiglu_fusion attribute, so it + // falls back to the default of 0 ("not fused"). When the activation is SwiGLU, swiglu_fusion is 0, + // and there is no separate FC3 weight, the gate and value projections are actually pre-fused into FC1 + // (interleaved layout). Treat this as swiglu_fusion == 1 so those legacy models keep working. + int swiglu_fusion = swiglu_fusion_; + if (activation_type_ == ActivationType::Swiglu && swiglu_fusion == 0 && + fc3_experts_weights_optional == nullptr) { + swiglu_fusion = 1; + LogSwigluFusionRemapOnce(); + } + bool is_fused_swiglu = (activation_type_ == ActivationType::Swiglu) && - (swiglu_fusion_ != 0) && + (swiglu_fusion != 0) && (fc3_experts_weights_optional == nullptr); MoEParameters moe_params; @@ -133,23 +158,23 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { // GEMM 1 MoeGemmId id1(static_cast(moe_params.inter_size), static_cast(moe_params.hidden_size), dtype, MoeGemmId::GemmType::Gemm1); - if (mGemmId1 != id1) { - mGemmId1 = id1; + { + // profileTactics caches per (GemmId, M bucket); calling it every forward lets decode + // (small M) and prefill (large M) each profile and select their own best tile shape. GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.inter_size), static_cast(moe_params.hidden_size)); - mGemmProfiler.profileTactics(&moe_runner, dtype, dims, id1); + mGemmProfiler.profileTactics(&moe_runner, dims, id1); } - auto config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), mGemmId1); + auto config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id1); // GEMM 2 MoeGemmId id2(static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), dtype, MoeGemmId::GemmType::Gemm2); - if (mGemmId2 != id2) { - mGemmId2 = id2; + { GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size)); - mGemmProfiler.profileTactics(&moe_runner, dtype, dims, id2); + mGemmProfiler.profileTactics(&moe_runner, dims, id2); } - auto config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), mGemmId2); + auto config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id2); moe_runner.setTactic(config1, config2); } @@ -301,7 +326,7 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { onnxruntime::llm::kernels::cutlass_kernels::ActivationParams params(kernel_activation_type); params.alpha = activation_alpha_; params.beta = activation_beta_; - params.swiglu_fusion = swiglu_fusion_; + params.swiglu_fusion = swiglu_fusion; params.limit = swiglu_limit_; return params; }(), diff --git a/onnxruntime/contrib_ops/cuda/moe/moe.h b/onnxruntime/contrib_ops/cuda/moe/moe.h index 58dbb2c70e3d0..79d97cfcfd754 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe.h @@ -24,8 +24,6 @@ class MoE final : public CudaKernel, public MoEBase { private: mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler mGemmProfiler; - mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmId mGemmId1; - mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmId mGemmId2; mutable std::mutex mGemmProfilerMutex; }; diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index 8e78076288015..de97ce71323f0 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -23,12 +23,23 @@ #include #include +#include #include using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; +namespace { +void LogQMoESwigluFusionRemapOnce() { + static std::once_flag log_warning; + std::call_once(log_warning, []() { + LOGS_DEFAULT(WARNING) << "QMoE swiglu_fusion is 0; assuming interleaved SwiGLU layout " + "for backward compatibility."; + }); +} +} // namespace + namespace onnxruntime { namespace contrib { namespace cuda { @@ -258,6 +269,18 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { "QMoE in CUDA execution provider does not support separate fc3_experts_weights. " "Gate and up projection weights must be pre-concatenated into fc1."); + // Backward compatibility: the published gpt-oss-20b model (and any model exported by ORT < 1.27) + // hard-coded the interleaved SwiGLU fusion layout and did not emit a swiglu_fusion attribute, so it + // falls back to the default of 0 ("not fused"). QMoE never has a separate FC3 (enforced above), so a + // SwiGLU activation with swiglu_fusion == 0 means the gate and value projections are actually pre-fused + // into FC1 (interleaved layout). Treat this as swiglu_fusion == 1 so those legacy models keep working. + int swiglu_fusion = swiglu_fusion_; + if (activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Swiglu && + swiglu_fusion == 0) { + swiglu_fusion = 1; + LogQMoESwigluFusionRemapOnce(); + } + const Tensor* fc1_zeros = packed_fc1_bias_ ? nullptr : context->Input(11); const Tensor* fc2_zeros = packed_fc2_bias_ ? nullptr : context->Input(12); @@ -299,10 +322,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { "QMoE row-wise quantization (block_size <= 0) does not support zero_points. " "Remove fc*_zero_points or use block-wise quantization."); } - if (block_size_ > 0 && block_size_ < 64 && has_any_zero_point) { + if (block_size_ > 0 && block_size_ < 32 && has_any_zero_point) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "QMoE asymmetric zero_points are currently supported only when block_size >= 64. " - "Use block_size >= 64 or remove fc*_zero_points."); + "QMoE asymmetric zero_points are currently supported only when block_size >= 32. " + "Use block_size >= 32 or remove fc*_zero_points."); } int64_t pack_size = expert_weight_bits_ == 4 ? 2 : 1; @@ -468,23 +491,23 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // GEMM 1: N=fc1_out_size (doubled for gated), K=hidden_size MoeGemmId id1(static_cast(fc1_out_size), static_cast(moe_params.hidden_size), dtype, wtype, MoeGemmId::GemmType::Gemm1); - if (mGemmId1 != id1) { - mGemmId1 = id1; + { + // profileTactics caches per (GemmId, M bucket); calling it every forward lets decode + // (small M) and prefill (large M) each profile and select their own best tile shape. GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), fc1_out_size, static_cast(moe_params.hidden_size)); - mGemmProfiler.profileTactics(m_moe_runner.get(), dtype, dims, id1); + mGemmProfiler.profileTactics(m_moe_runner.get(), dims, id1); } - config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), mGemmId1); + config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id1); // GEMM 2 MoeGemmId id2(static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), dtype, wtype, MoeGemmId::GemmType::Gemm2); - if (mGemmId2 != id2) { - mGemmId2 = id2; + { GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size)); - mGemmProfiler.profileTactics(m_moe_runner.get(), dtype, dims, id2); + mGemmProfiler.profileTactics(m_moe_runner.get(), dims, id2); } - config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), mGemmId2); + config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id2); m_moe_runner->setTactic(config1, config2); } @@ -1008,7 +1031,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { onnxruntime::llm::kernels::cutlass_kernels::ActivationParams params(activation_type_); params.alpha = activation_alpha_; params.beta = activation_beta_; - params.swiglu_fusion = swiglu_fusion_; + params.swiglu_fusion = swiglu_fusion; params.limit = swiglu_limit_; return params; }(), diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index 2bbadc205b5d8..91c84f919bc82 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -117,8 +117,6 @@ class QMoE final : public CudaKernel, public MoEBase { IAllocatorUniquePtr packed_fc2_act_scale_; mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler mGemmProfiler; - mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmId mGemmId1; - mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmId mGemmId2; mutable std::mutex mGemmProfilerMutex; }; diff --git a/onnxruntime/test/python/transformers/profile_qmoe_gemv.py b/onnxruntime/test/python/transformers/profile_qmoe_gemv.py new file mode 100644 index 0000000000000..0f71409fa2206 --- /dev/null +++ b/onnxruntime/test/python/transformers/profile_qmoe_gemv.py @@ -0,0 +1,112 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +""" +Profiling script for the CUDA QMoE GEMV decode path. + +Usage: + python profile_qmoe_gemv.py --case m1_top2_fp16_128x256 --warmup 5 --repeat 100 + + nsys profile -t cuda,nvtx -o qmoe_gemv --export=sqlite \ + python profile_qmoe_gemv.py --case m1_top2_fp16_128x256 --warmup 5 --repeat 100 + python parse_nsys.py qmoe_gemv.sqlite --nvtx-range benchmark +""" + +import argparse +import json +import os + +import torch +from test_qmoe_cuda import ( + _QMOE_GEMV_BENCHMARK_RESULT_PREFIX, + _qmoe_gemv_benchmark_case, + _qmoe_gemv_benchmark_cases, + run_qmoe_gemv_benchmark, +) + + +def _custom_case_from_args(args): + case = dict(_qmoe_gemv_benchmark_case(args.case)) + custom_fields = { + "batch_size": args.batch_size, + "sequence_length": args.sequence_length, + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "num_experts": args.num_experts, + "top_k": args.top_k, + "onnx_dtype": args.dtype, + "quant_bits": args.quant_bits, + "block_size": args.block_size, + } + case.update({key: value for key, value in custom_fields.items() if value is not None}) + + if any(value is not None for value in custom_fields.values()): + case["name"] = ( + f"custom_m{case['batch_size'] * case['sequence_length']}_top{case['top_k']}_" + f"{case['onnx_dtype'].lower()}_{case['hidden_size']}x{case['intermediate_size']}_" + f"e{case['num_experts']}_int{case.get('quant_bits', 4)}_b{case.get('block_size', 0)}" + ) + + return case + + +def main(): + parser = argparse.ArgumentParser(description="Profile CUDA QMoE GEMV decode") + parser.add_argument("--case", default="m1_top2_fp16_128x256", help="Benchmark case name") + parser.add_argument("--list-cases", action="store_true", help="List available benchmark case names and exit") + parser.add_argument("--batch-size", type=int, help="Override batch size") + parser.add_argument("--sequence-length", type=int, help="Override sequence length") + parser.add_argument("--hidden-size", type=int, help="Override hidden size") + parser.add_argument("--intermediate-size", type=int, help="Override intermediate size") + parser.add_argument("--num-experts", type=int, help="Override number of experts") + parser.add_argument("--top-k", type=int, help="Override top-k experts per token") + parser.add_argument("--dtype", choices=["FLOAT16", "BFLOAT16"], help="Override ONNX dtype") + parser.add_argument("--quant-bits", type=int, choices=[4, 8], help="Override QMoE integer weight bits") + parser.add_argument("--block-size", type=int, choices=[0, 32, 64, 128], help="Override QMoE INT block size") + parser.add_argument("--warmup", type=int, default=5, help="Warmup iterations before the benchmark NVTX range") + parser.add_argument("--repeat", type=int, default=100, help="Benchmark iterations") + parser.add_argument( + "--disable-gemv", + action="store_true", + help="Run the grouped GEMM fallback by setting ORT_DISABLE_MOE_GEMV=1 before session creation", + ) + parser.add_argument( + "--nvtx", + action="store_true", + help="Wrap the measured loop in an NVTX range named 'benchmark'", + ) + args = parser.parse_args() + + if args.list_cases: + for case in _qmoe_gemv_benchmark_cases(): + print(case["name"]) + return + + case = _custom_case_from_args(args) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for QMoE GEMV profiling") + + os.environ["ORT_QMOE_GEMV_BENCHMARK_REPEATS"] = str(max(1, args.repeat)) + os.environ["ORT_QMOE_GEMV_BENCHMARK_WARMUP"] = str(max(0, args.warmup)) + if args.nvtx: + os.environ["ORT_QMOE_GEMV_BENCHMARK_NVTX"] = "1" + else: + os.environ.pop("ORT_QMOE_GEMV_BENCHMARK_NVTX", None) + + if args.disable_gemv: + os.environ["ORT_DISABLE_MOE_GEMV"] = "1" + else: + os.environ.pop("ORT_DISABLE_MOE_GEMV", None) + + result = run_qmoe_gemv_benchmark(case) + if result["has_invalid_output"]: + raise RuntimeError("QMoE GEMV profiling produced NaN or Inf output") + + print(_QMOE_GEMV_BENCHMARK_RESULT_PREFIX + json.dumps(result, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh b/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh new file mode 100755 index 0000000000000..c452ee17fec00 --- /dev/null +++ b/onnxruntime/test/python/transformers/profile_qmoe_gemv.sh @@ -0,0 +1,130 @@ +#!/bin/bash +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# +# Profile the CUDA QMoE GEMV decode path with nsys. +# +# Usage: +# ./profile_qmoe_gemv.sh +# ./profile_qmoe_gemv.sh --list-cases +# ./profile_qmoe_gemv.sh --case m8_top2_fp16_128x256 --warmup 5 --repeat 200 +# ./profile_qmoe_gemv.sh --case gpt_oss_20b_m1_top4_fp16_2880x2880_e32 --warmup 5 --repeat 100 +# ./profile_qmoe_gemv.sh --batch-size 1 --sequence-length 1 --hidden-size 1024 --intermediate-size 4096 --num-experts 8 --top-k 2 --quant-bits 8 --block-size 128 +# CUDA_VISIBLE_DEVICES=1 ./profile_qmoe_gemv.sh -o /tmp/qmoe_gemv +# + +set -e +set -o pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CASE="m1_top2_fp16_128x256" +WARMUP=5 +REPEAT=100 +OUTPUT_NAME="qmoe_gemv_profile" +PY="${PYTHON:-python}" +EXTRA_ARGS=() +LIST_CASES=0 + +while [[ "$#" -gt 0 ]]; do + case $1 in + --case) + CASE="$2" + shift + ;; + --list-cases) + LIST_CASES=1 + ;; + --batch-size|--sequence-length|--hidden-size|--intermediate-size|--num-experts|--top-k|--dtype|--quant-bits|--block-size) + EXTRA_ARGS+=("$1" "$2") + shift + ;; + --repeat) + REPEAT="$2" + shift + ;; + --warmup) + WARMUP="$2" + shift + ;; + --python) + PY="$2" + shift + ;; + -o|--output) + OUTPUT_NAME="$2" + shift + ;; + *) + echo "Unknown option: $1" + echo "Usage: $0 [--list-cases] [--case NAME] [--batch-size N] [--sequence-length N] [--hidden-size N] [--intermediate-size N] [--num-experts N] [--top-k N] [--dtype FLOAT16|BFLOAT16] [--quant-bits 4|8] [--block-size 0|32|64|128] [--warmup N] [--repeat N] [--python PYTHON] [-o NAME]" + exit 1 + ;; + esac + shift +done + +if [[ "${LIST_CASES}" -eq 1 ]]; then + "${PY}" "${SCRIPT_DIR}/profile_qmoe_gemv.py" --list-cases + exit 0 +fi + +if ! command -v nsys >/dev/null; then + echo "Error: nsys not found. Install NVIDIA Nsight Systems or add it to PATH." + exit 1 +fi + +HAVE_NVTX=0 +if "${PY}" -c "import nvtx" 2>/dev/null; then + HAVE_NVTX=1 +else + echo "Note: 'nvtx' package not installed. NVTX range markers will be disabled." + echo " Install with: pip install nvtx" + echo " Falling back to --skip-first to exclude warmup-like first calls." +fi + +echo "" +echo "========================================" +echo " Profiling: CUDA QMoE GEMV" +echo "========================================" +echo "Case: ${CASE}" +echo "Warmup: ${WARMUP}" +echo "Repeat: ${REPEAT}" +echo "Python: ${PY}" +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then + echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}" +fi +if [[ "${#EXTRA_ARGS[@]}" -gt 0 ]]; then + echo "Custom args: ${EXTRA_ARGS[*]}" +fi + +profile_one() { + local mode="$1" + local disable_arg="" + local base="${OUTPUT_NAME}_${mode}" + if [[ "${mode}" == "gemm" ]]; then + disable_arg="--disable-gemv" + fi + + echo "" + echo "---- Profiling ${mode} ----" + rm -f "${base}.nsys-rep" "${base}.sqlite" + nsys profile -t cuda,nvtx --force-overwrite true -o "${base}" --export=sqlite \ + "${PY}" "${SCRIPT_DIR}/profile_qmoe_gemv.py" \ + --case "${CASE}" "${EXTRA_ARGS[@]}" --warmup "${WARMUP}" --repeat "${REPEAT}" --nvtx ${disable_arg} + + echo "" + echo "---- Kernel results (${mode}) ----" + if [[ "${HAVE_NVTX}" -eq 1 ]]; then + "${PY}" "${SCRIPT_DIR}/parse_nsys.py" "${base}.sqlite" --nvtx-range benchmark + else + "${PY}" "${SCRIPT_DIR}/parse_nsys.py" "${base}.sqlite" --skip-first 1 + fi +} + +profile_one gemv +profile_one gemm + +echo "" +echo "Done." diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index 9677542270a53..f7aa3cf48b6bf 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -152,9 +152,7 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True): q_weight_reshaped = q_weight.reshape(n, -1) # Pack weights for CUDA mixed-gemm kernel (FpA_IntB format), and qMoE kernel uses the same format. - # Pin arch=80: the QMoE grouped MoE GEMM always runs the Ampere (SM80) kernel -- even on SM90 -- - # so it consumes the SM80 (column-interleaved) layout on every GPU. Auto-detect (force_arch=-1) - # would emit the non-interleaved SM90 layout on Hopper and produce wrong results. + # INT MoE/QMoE kernels consume the SM80 column-interleaved layout, including on newer GPUs. processed_q_weight = _quantize.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 4, 80) # So we need to DEQUANTIZE back to get `result`. @@ -235,10 +233,8 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True): ) q_weight_reshaped = q_weight.reshape(n, -1) - # Pack weights for CUDA mixed-gemm kernel (FpA_IntB format). - # Pin arch=80: the QMoE grouped MoE GEMM always runs the Ampere (SM80) kernel -- even on SM90 -- - # so it consumes the SM80 (column-interleaved) layout on every GPU. Auto-detect (force_arch=-1) - # would emit the non-interleaved SM90 layout on Hopper and produce wrong results. + # Pack weights for CUDA mixed-gemm kernel (FpA_IntB format) + # INT MoE/QMoE kernels consume the SM80 column-interleaved layout, including on newer GPUs. processed_q_weight = _quantize.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 8, 80) # Dequantize for reference @@ -1526,7 +1522,7 @@ def test_phi3_moe_parity(self, batch_size, sequence_length, quant_bits, onnx_typ phi3_qmoe_test_cases = list( itertools.product( [1, 4], # batch_size - [1, 8], # sequence_length + [1, 8], # sequence_length; keeps expanded rows <= 64 to exercise the INT4 MoE GEMV path [TensorProto.FLOAT16], # onnx type, None mean fp32 for bits = 0, fp16 for bits > 0 [True], # normalize_routing_weights ) diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index c56383d2851d3..fa741bbadc272 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -18,9 +18,12 @@ # while maintaining computational efficiency. # -------------------------------------------------------------------------- import copy +import json +import os import time import unittest from collections import OrderedDict +from contextlib import nullcontext import numpy import torch @@ -33,6 +36,14 @@ import onnxruntime from onnxruntime.capi import _pybind_state as _pybind +try: + import nvtx + + has_nvtx = True +except ImportError: + has_nvtx = False + nvtx = None + try: from onnx import TensorProto @@ -81,6 +92,17 @@ class TensorProtoPlaceholder: else: ort_provider = ["CPUExecutionProvider"] + +def _qmoe_benchmark_nvtx_range(name="benchmark", color="green"): + if os.getenv("ORT_QMOE_GEMV_BENCHMARK_NVTX") != "1": + return nullcontext() + + if not has_nvtx: + return nullcontext() + + return nvtx.annotate(name, color=color) + + torch.manual_seed(42) numpy.random.seed(42) @@ -287,6 +309,39 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True, asymmetric: bool scale, quantized_storage, dequantized, zero_point_storage """ block_size = weights.shape[1] + if not asymmetric and block_size > 256: + n, k = weights.shape + weights_float = weights.detach().float() + if is_4_bit_quantization: + scale = weights_float.abs().amax(dim=1, keepdim=True) / 7.0 + scale = torch.clamp(scale, min=torch.finfo(torch.float32).eps) + q_weight = torch.clamp(torch.round(weights_float / scale), -8, 7).to(torch.int16) + 8 + q_weight = q_weight.to(torch.uint8).contiguous() + q_low = q_weight[:, 0::2] + q_high = q_weight[:, 1::2] + if q_high.shape[1] < q_low.shape[1]: + q_high = F.pad(q_high, (0, 1)) + q_packed = q_low | (q_high << 4) + processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_packed.cpu().numpy(), n, k, 4, 80) + processed_q_weight_torch = ( + torch.from_numpy(processed_q_weight).reshape(k, n // 2).to(weights.device).view(torch.uint8) + ) + dequantized = (q_weight.to(weights.dtype) - 8.0) * scale.to(weights.device).to(weights.dtype) + return scale.to(weights.device).to(torch.float16), processed_q_weight_torch, dequantized, None + else: + # 8-bit per-column (per-channel) symmetric quantization. Weights are biased + # uint8 values centered at 128, matching the kernel's (q - 128) * scale dequant. + scale = weights_float.abs().amax(dim=1, keepdim=True) / 127.0 + scale = torch.clamp(scale, min=torch.finfo(torch.float32).eps) + q_weight = torch.clamp(torch.round(weights_float / scale), -127, 127).to(torch.int16) + 128 + q_weight = q_weight.to(torch.uint8).contiguous() + processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight.cpu().numpy(), n, k, 8, 80) + processed_q_weight_torch = ( + torch.from_numpy(processed_q_weight).reshape(k, n).to(weights.device).view(torch.uint8) + ) + dequantized = (q_weight.to(weights.dtype) - 128.0) * scale.to(weights.device).to(weights.dtype) + return scale.to(weights.device).to(torch.float16), processed_q_weight_torch, dequantized, None + return quant_dequant_blockwise(weights, block_size, is_4_bit_quantization, asymmetric) @@ -855,14 +910,25 @@ def ort_forward( print("DEBUG: ORT inference completed successfully") if enable_performance_test: - repeat = 100 - s = time.time() - for _ in range(repeat): - iobinding.synchronize_inputs() - self.ort_sess.run_with_iobinding(iobinding) - iobinding.synchronize_outputs() - e = time.time() + warmup = max(0, int(os.getenv("ORT_QMOE_GEMV_BENCHMARK_WARMUP", "5"))) + repeat = max(1, int(os.getenv("ORT_QMOE_GEMV_BENCHMARK_REPEATS", "100"))) + with _qmoe_benchmark_nvtx_range("warmup", "yellow"): + for _ in range(warmup): + iobinding.synchronize_inputs() + self.ort_sess.run_with_iobinding(iobinding) + iobinding.synchronize_outputs() + + with _qmoe_benchmark_nvtx_range("benchmark", "green"): + torch.cuda.synchronize() + s = time.perf_counter() + for _ in range(repeat): + iobinding.synchronize_inputs() + self.ort_sess.run_with_iobinding(iobinding) + iobinding.synchronize_outputs() + torch.cuda.synchronize() + e = time.perf_counter() time_ms = (e - s) / repeat * 1000 + self.last_ort_latency_ms = time_ms is_swiglu = hasattr(self, "use_swiglu") and self.use_swiglu is_interleaved = getattr(self, "swiglu_fusion", 0) == 1 act_type = f"SwiGLU(interleaved={is_interleaved})" if is_swiglu else "SiLU" @@ -947,12 +1013,18 @@ def recreate_onnx_model(self): w2_bias_list.append(w2_bias.detach().cpu()) torch_dtype = onnx_to_torch_type_map[self.onnx_dtype] if self.onnx_dtype else torch.float32 - # For BF16 quantized: keep expert weights in float32 so the PyTorch reference + # For quantized MoE: keep expert weights in float32 so the PyTorch reference # computes in float32 (PhiMoESwiGLUMLP.forward casts input to weight dtype). - # ORT's CUTLASS kernel accumulates int8 products in float32 before applying the - # BF16 scale, matching float32 precision. Storing weights as BF16 causes - # catastrophic cancellation for near-zero outputs due to the 7-bit mantissa. - ref_weight_dtype = torch.float32 if (torch_dtype == torch.bfloat16 and self.quant_bits > 0) else torch_dtype + # ORT's CUTLASS grouped GEMM and the decode GEMV kernel both accumulate the + # weight*activation products in float32 before applying the FP16/BF16 scale, so a + # float32 reference matches the kernel's accumulation precision. Storing weights in + # the low-precision dtype causes catastrophic cancellation for near-zero outputs + # (BF16's 7-bit / FP16's 10-bit mantissa) and makes the reference itself lossy. + ref_weight_dtype = ( + torch.float32 + if (torch_dtype in (torch.bfloat16, torch.float16) and self.quant_bits > 0) + else torch_dtype + ) if self.use_swiglu: if getattr(self, "swiglu_fusion", 0) == 1: @@ -1149,7 +1221,7 @@ def parity_check(self): dtype_str = ort_dtype_name_map[self.onnx_dtype] tolerance_key = f"{dtype_str}:{self.quant_bits}" if tolerance_key in ort_dtype_quant_bits_tolerance_map: - base_atol, rtol = ort_dtype_quant_bits_tolerance_map[tolerance_key] + base_atol, _rtol = ort_dtype_quant_bits_tolerance_map[tolerance_key] # Increase tolerance for asymmetric quantization due to different computation path if self.use_asymmetric_quant: @@ -1318,11 +1390,15 @@ def __init__( expert.w2.weight, is_4_bit, asymmetric=use_effective_asymmetric_quant ) - # For BF16 quantized: keep weights in float32 so the PyTorch reference - # computes in float32, matching ORT's CUTLASS kernel that accumulates int8 - # products in float32 before applying the BF16 scale. + # For quantized MoE: keep weights in float32 so the PyTorch reference computes + # in float32, matching ORT's CUTLASS grouped GEMM and decode GEMV kernel that + # both accumulate weight*activation products in float32 before applying the + # FP16/BF16 scale. A low-precision reference is itself lossy and would mask the + # kernel's accumulation precision. ref_weight_dtype = ( - torch.float32 if (torch_dtype == torch.bfloat16 and self.quant_bits > 0) else torch_dtype + torch.float32 + if (torch_dtype in (torch.bfloat16, torch.float16) and self.quant_bits > 0) + else torch_dtype ) expert.w1.weight.data = w1_qdq.to(ref_weight_dtype) expert.w2.weight.data = w2_qdq.to(ref_weight_dtype) @@ -1423,6 +1499,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Define test cases for different MoE types phi3_test_cases = [ + (1, 1, 4), # decode-sized INT4 per-channel path exercises the MoE GEMV fast path (1, 32, 4), (1, 32, 8), (2, 16, 4), @@ -1433,16 +1510,112 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: phi3_blockwise_test_cases = [ (1, 1, 4, 32), # tiny debug case for asymmetric ZP compensation (1, 32, 4, 32), # batch_size, sequence_length, quant_bits, block_size + (1, 32, 4, 64), + (1, 32, 4, 128), + (1, 32, 8, 32), (1, 32, 8, 64), + (1, 32, 8, 128), (2, 16, 4, 32), + (2, 16, 8, 32), (2, 16, 8, 64), ] phi3_blockwise_asymmetric_test_cases = [ + (1, 1, 4, 32), + (1, 1, 8, 32), (1, 32, 4, 64), (1, 32, 8, 64), + (1, 32, 8, 128), (2, 16, 8, 64), ] +# These cases use expanded rows > 4 with K < 512, which is outside the profiled +# GEMV range and therefore exercises the CUTLASS grouped GEMM path. +qmoe_cutlass_gemm_blockwise_test_cases = [ + (1, 3, 4, 32), + (1, 3, 8, 32), +] + +qmoe_cutlass_gemm_second_scale_row_test_cases = [ + (4, False), + (4, True), + (8, False), + (8, True), +] + + +def _run_qmoe_cutlass_gemm_second_scale_row_regression(test_case, quant_bits, use_asymmetric_quant): + hidden_size = 128 + intermediate_size = 128 + sequence_length = 8 + num_experts = 1 + top_k = 1 + block_size = 32 + onnx_dtype = TensorProto.FLOAT16 + torch_dtype = onnx_to_torch_type_map[onnx_dtype] + + is_4_bit = quant_bits == 4 + + fc1 = torch.zeros((intermediate_size, hidden_size), device=device, dtype=torch_dtype) + fc2 = torch.zeros((hidden_size, intermediate_size), device=device, dtype=torch_dtype) + + fc1[0, :block_size] = 1.0 / 1024.0 + fc1[0, block_size : 2 * block_size] = 1.0 + fc2[0, 0] = 1.0 + + fc1_scale, fc1_weight, fc1_qdq, fc1_zp = quant_dequant_blockwise( + fc1, block_size, is_4_bit, asymmetric=use_asymmetric_quant + ) + fc2_scale, fc2_weight, fc2_qdq, fc2_zp = quant_dequant_blockwise( + fc2, block_size, is_4_bit, asymmetric=use_asymmetric_quant + ) + + model = create_moe_onnx_graph( + hidden_size=hidden_size, + sequence_length=sequence_length, + num_experts=num_experts, + top_k=top_k, + intermediate_size=intermediate_size, + torch_dtype=torch.float32, + onnx_dtype=onnx_dtype, + fc1_experts_weights=fc1_weight.unsqueeze(0), + fc2_experts_weights=fc2_weight.unsqueeze(0), + fc1_scales=fc1_scale.unsqueeze(0), + fc2_scales=fc2_scale.unsqueeze(0), + fc1_zero_points=fc1_zp.unsqueeze(0) if fc1_zp is not None else None, + fc2_zero_points=fc2_zp.unsqueeze(0) if fc2_zp is not None else None, + use_swiglu=False, + use_quant=True, + quant_bits=quant_bits, + block_size=block_size, + ) + + previous_disable_gemv = os.environ.get("ORT_DISABLE_MOE_GEMV") + os.environ["ORT_DISABLE_MOE_GEMV"] = "1" + try: + sess_options = onnxruntime.SessionOptions() + sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + sess = onnxruntime.InferenceSession( + model, + sess_options, + providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")], + ) + finally: + if previous_disable_gemv is None: + os.environ.pop("ORT_DISABLE_MOE_GEMV", None) + else: + os.environ["ORT_DISABLE_MOE_GEMV"] = previous_disable_gemv + + x = torch.zeros((sequence_length, hidden_size), device=device, dtype=torch_dtype) + x[:, block_size : 2 * block_size] = 1.0 if use_asymmetric_quant else -1.0 + router = torch.zeros((sequence_length, num_experts), device=device, dtype=torch_dtype) + + ort_output = sess.run(None, {"input": x.cpu().numpy(), "router_probs": router.cpu().numpy()})[0] + fc1_output = torch.matmul(x.float(), fc1_qdq.float().T) + expected = torch.matmul(F.silu(fc1_output), fc2_qdq.float().T).cpu().numpy().astype(numpy.float16) + + test_case.assertGreater(abs(expected[0, 0]), 20.0) + numpy.testing.assert_allclose(ort_output, expected, rtol=2e-2, atol=2.5e-1) + @unittest.skipIf(not torch.cuda.is_available(), "skipping QMoE test since it requires CUDA.") class TestPhiQMoE(unittest.TestCase): @@ -1539,8 +1712,6 @@ def test_phi3_qmoe_asymmetric_parity(self, batch_size, sequence_length, quant_bi @parameterized.expand(phi3_blockwise_test_cases) def test_phi3_qmoe_blockwise_parity(self, batch_size, sequence_length, quant_bits, block_size): - if quant_bits == 8: - self.skipTest("8-bit blockwise quantization is not supported on CUDA") torch.manual_seed(42) numpy.random.seed(42) @@ -1573,8 +1744,6 @@ def test_phi3_qmoe_blockwise_parity(self, batch_size, sequence_length, quant_bit @parameterized.expand(phi3_blockwise_test_cases) def test_phi3_qmoe_blockwise_parity_bf16(self, batch_size, sequence_length, quant_bits, block_size): - if quant_bits == 8: - self.skipTest("8-bit blockwise quantization is not supported on CUDA") torch.manual_seed(142) numpy.random.seed(142) @@ -1619,6 +1788,52 @@ def test_phi3_qmoe_blockwise_asymmetric_parity(self, batch_size, sequence_length ) phi3_moe.parity_check() + @parameterized.expand(qmoe_cutlass_gemm_blockwise_test_cases) + def test_phi3_qmoe_blockwise_cutlass_gemm_parity(self, batch_size, sequence_length, quant_bits, block_size): + torch.manual_seed(44) + numpy.random.seed(44) + + test_config = f"batch_size={batch_size}, sequence_length={sequence_length}, quant_bits={quant_bits}, block_size={block_size}" + print(f"Running Phi3 QMoE block-wise CUTLASS GEMM test: {test_config}") + + config = PhiMoEConfig(hidden_size=128, intermediate_size=256, num_local_experts=4, num_experts_per_tok=2) + + phi3_moe = PhiMoESparseMoeBlock( + config, + batch_size=batch_size, + sequence_length=sequence_length, + quant_bits=quant_bits, + onnx_dtype=TensorProto.FLOAT16, + block_size=block_size, + use_asymmetric_quant=False, + ) + phi3_moe.parity_check() + + @parameterized.expand(qmoe_cutlass_gemm_blockwise_test_cases) + def test_phi3_qmoe_blockwise_cutlass_gemm_parity_bf16(self, batch_size, sequence_length, quant_bits, block_size): + torch.manual_seed(144) + numpy.random.seed(144) + + test_config = f"batch_size={batch_size}, sequence_length={sequence_length}, quant_bits={quant_bits}, block_size={block_size} (BF16)" + print(f"Running Phi3 QMoE block-wise CUTLASS GEMM test (BF16): {test_config}") + + config = PhiMoEConfig(hidden_size=128, intermediate_size=256, num_local_experts=4, num_experts_per_tok=2) + + phi3_moe = PhiMoESparseMoeBlock( + config, + batch_size=batch_size, + sequence_length=sequence_length, + quant_bits=quant_bits, + onnx_dtype=TensorProto.BFLOAT16, + block_size=block_size, + use_asymmetric_quant=False, + ) + phi3_moe.parity_check() + + @parameterized.expand(qmoe_cutlass_gemm_second_scale_row_test_cases) + def test_phi3_qmoe_blockwise_cutlass_gemm_second_scale_row(self, quant_bits, use_asymmetric_quant): + _run_qmoe_cutlass_gemm_second_scale_row_regression(self, quant_bits, use_asymmetric_quant) + swiglu_test_cases = [ (1, 32, 4), @@ -1632,13 +1847,20 @@ def test_phi3_qmoe_blockwise_asymmetric_parity(self, batch_size, sequence_length (1, 1, 4, 32), # tiny debug case for asymmetric ZP compensation (1, 32, 4, 32), # batch_size, sequence_length, quant_bits, block_size (1, 32, 4, 64), # New case for group_size=64 + (1, 32, 4, 128), + (1, 32, 8, 32), (1, 32, 8, 64), + (1, 32, 8, 128), (2, 16, 4, 32), + (2, 16, 8, 32), (2, 16, 8, 64), ] swiglu_blockwise_asymmetric_test_cases = [ + (1, 1, 4, 32), + (1, 1, 8, 32), (1, 32, 4, 64), (1, 32, 8, 64), + (1, 32, 8, 128), (2, 16, 8, 64), ] @@ -1812,6 +2034,52 @@ def test_swiglu_qmoe_blockwise_asymmetric_parity(self, batch_size, sequence_leng ) swiglu_moe.parity_check() + @parameterized.expand(qmoe_cutlass_gemm_blockwise_test_cases) + def test_swiglu_qmoe_blockwise_cutlass_gemm_parity(self, batch_size, sequence_length, quant_bits, block_size): + torch.manual_seed(44) + numpy.random.seed(44) + + test_config = f"batch_size={batch_size}, sequence_length={sequence_length}, quant_bits={quant_bits}, block_size={block_size}" + print(f"Running SwiGLU block-wise CUTLASS GEMM test: {test_config}") + + config = SwigluMoeConfig(hidden_size=128, intermediate_size=256, num_local_experts=4, num_experts_per_token=2) + + swiglu_moe = SwigluMoEBlock( + config, + batch_size=batch_size, + sequence_length=sequence_length, + quant_bits=quant_bits, + onnx_dtype=TensorProto.FLOAT16, + block_size=block_size, + use_asymmetric_quant=False, + ) + swiglu_moe.parity_check() + + @parameterized.expand(qmoe_cutlass_gemm_blockwise_test_cases) + def test_swiglu_qmoe_blockwise_cutlass_gemm_parity_bf16(self, batch_size, sequence_length, quant_bits, block_size): + torch.manual_seed(144) + numpy.random.seed(144) + + test_config = f"batch_size={batch_size}, sequence_length={sequence_length}, quant_bits={quant_bits}, block_size={block_size} (BF16)" + print(f"Running SwiGLU block-wise CUTLASS GEMM test (BF16): {test_config}") + + config = SwigluMoeConfig(hidden_size=128, intermediate_size=256, num_local_experts=4, num_experts_per_token=2) + + swiglu_moe = SwigluMoEBlock( + config, + batch_size=batch_size, + sequence_length=sequence_length, + quant_bits=quant_bits, + onnx_dtype=TensorProto.BFLOAT16, + block_size=block_size, + use_asymmetric_quant=False, + ) + swiglu_moe.parity_check() + + @parameterized.expand(qmoe_cutlass_gemm_second_scale_row_test_cases) + def test_swiglu_qmoe_blockwise_cutlass_gemm_second_scale_row(self, quant_bits, use_asymmetric_quant): + _run_qmoe_cutlass_gemm_second_scale_row_regression(self, quant_bits, use_asymmetric_quant) + def has_bf16_qmoe(): """Check if BF16 QMoE is supported (requires Ampere or newer GPU).""" @@ -1866,6 +2134,321 @@ def test_swiglu_qmoe_bf16_parity(self, batch_size, sequence_length, quant_bits): swiglu_moe.parity_check() +_QMOE_GEMV_BENCHMARK_RESULT_PREFIX = "QMOE_GEMV_BENCHMARK_RESULT " + + +def _qmoe_gemv_benchmark_cases(): + return [ + { + "name": "m1_top2_fp16_128x256", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 128, + "intermediate_size": 256, + "num_experts": 4, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "m4_top2_fp16_128x256", + "batch_size": 1, + "sequence_length": 4, + "hidden_size": 128, + "intermediate_size": 256, + "num_experts": 4, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "m8_top2_fp16_128x256", + "batch_size": 1, + "sequence_length": 8, + "hidden_size": 128, + "intermediate_size": 256, + "num_experts": 4, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "m1_top2_bf16_128x256", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 128, + "intermediate_size": 256, + "num_experts": 4, + "top_k": 2, + "onnx_dtype": "BFLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "gpt_oss_20b_m1_top4_fp16_2880x2880_e32", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_experts": 32, + "top_k": 4, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "qwen3_6_35b_a3b_m1_top8_fp16_2048x512_e256", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2048, + "intermediate_size": 512, + "num_experts": 256, + "top_k": 8, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "gemma4_26b_a4b_m1_top8_fp16_2816x704_e128", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2816, + "intermediate_size": 704, + "num_experts": 128, + "top_k": 8, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "blockwise_int4_b64_m1_top2_fp16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 64, + }, + { + "name": "blockwise_int4_b128_m1_top2_fp16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 4, + "block_size": 128, + }, + { + "name": "blockwise_int8_b64_m1_top2_fp16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 8, + "block_size": 64, + }, + { + "name": "blockwise_int8_b128_m1_top2_fp16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 8, + "block_size": 128, + }, + { + "name": "gpt_oss_20b_m1_top4_bf16_2880x2880_e32", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_experts": 32, + "top_k": 4, + "onnx_dtype": "BFLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "qwen3_6_35b_a3b_m1_top8_bf16_2048x512_e256", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2048, + "intermediate_size": 512, + "num_experts": 256, + "top_k": 8, + "onnx_dtype": "BFLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "gemma4_26b_a4b_m1_top8_bf16_2816x704_e128", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2816, + "intermediate_size": 704, + "num_experts": 128, + "top_k": 8, + "onnx_dtype": "BFLOAT16", + "quant_bits": 4, + "block_size": 0, + }, + { + "name": "blockwise_int4_b64_m1_top2_bf16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "BFLOAT16", + "quant_bits": 4, + "block_size": 64, + }, + { + "name": "blockwise_int8_b64_m1_top2_bf16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "BFLOAT16", + "quant_bits": 8, + "block_size": 64, + }, + { + "name": "gpt_oss_20b_m1_top4_int8_fp16_2880x2880_e32", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_experts": 32, + "top_k": 4, + "onnx_dtype": "FLOAT16", + "quant_bits": 8, + "block_size": 0, + }, + { + "name": "gpt_oss_20b_m1_top4_int8_bf16_2880x2880_e32", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_experts": 32, + "top_k": 4, + "onnx_dtype": "BFLOAT16", + "quant_bits": 8, + "block_size": 0, + }, + { + "name": "int8_per_column_m1_top2_fp16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "FLOAT16", + "quant_bits": 8, + "block_size": 0, + }, + { + "name": "int8_per_column_m1_top2_bf16_1024x4096_e8", + "batch_size": 1, + "sequence_length": 1, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_experts": 8, + "top_k": 2, + "onnx_dtype": "BFLOAT16", + "quant_bits": 8, + "block_size": 0, + }, + ] + + +def _qmoe_gemv_benchmark_case(case_name): + for case in _qmoe_gemv_benchmark_cases(): + if case["name"] == case_name: + return case + + case_names = ", ".join(case["name"] for case in _qmoe_gemv_benchmark_cases()) + raise ValueError(f"Unknown QMoE GEMV benchmark case '{case_name}'. Available cases: {case_names}") + + +def run_qmoe_gemv_benchmark(case): + seed = 4242 + torch.manual_seed(seed) + numpy.random.seed(seed) + + onnx_dtype = getattr(TensorProto, case["onnx_dtype"]) + torch_dtype = onnx_to_torch_type_map[onnx_dtype] + config = PhiMoEConfig( + hidden_size=case["hidden_size"], + intermediate_size=case["intermediate_size"], + num_local_experts=case["num_experts"], + num_experts_per_tok=case["top_k"], + ) + qmoe = PhiMoESparseMoeBlock( + config, + batch_size=case["batch_size"], + sequence_length=case["sequence_length"], + quant_bits=case.get("quant_bits", 4), + onnx_dtype=onnx_dtype, + block_size=case.get("block_size", 0), + use_asymmetric_quant=False, + ) + hidden_states = torch.randn( + case["batch_size"], case["sequence_length"], case["hidden_size"], device=device, dtype=torch_dtype + ) + output = qmoe.ort_forward(hidden_states, enable_performance_test=True) + + return { + "case": case["name"], + "block_size": case.get("block_size", 0), + "disable_gemv": os.getenv("ORT_DISABLE_MOE_GEMV") == "1", + "expanded_num_rows": case["batch_size"] * case["sequence_length"] * case["top_k"], + "has_invalid_output": bool(torch.isnan(output).any() or torch.isinf(output).any()), + "latency_ms": qmoe.last_ort_latency_ms, + "quant_bits": case.get("quant_bits", 4), + "sm": torch.cuda.get_device_capability()[0] * 10 + torch.cuda.get_device_capability()[1], + } + + +def run_qmoe_gemv_benchmark_case(case_name=None): + case = _qmoe_gemv_benchmark_case( + case_name or os.getenv("ORT_QMOE_GEMV_BENCHMARK_CASE", _qmoe_gemv_benchmark_cases()[0]["name"]) + ) + return run_qmoe_gemv_benchmark(case) + + +@unittest.skipIf(not torch.cuda.is_available(), "skipping QMoE GEMV benchmark since it requires CUDA.") +@unittest.skipIf( + os.getenv("ORT_QMOE_GEMV_BENCHMARK") != "1", + "Set ORT_QMOE_GEMV_BENCHMARK=1 to run the opt-in QMoE GEMV benchmark.", +) +class TestQMoEGemvBenchmark(unittest.TestCase): + def test_decode_latency(self): + result = run_qmoe_gemv_benchmark_case() + self.assertFalse(result["has_invalid_output"]) + print(_QMOE_GEMV_BENCHMARK_RESULT_PREFIX + json.dumps(result, sort_keys=True)) + + @unittest.skipIf(True, "Skipping QMoE benchmark tests") class TestQMoESwiGLUBenchmark(unittest.TestCase): """Benchmark tests for QMoE SwiGLU performance measurement."""