[BBaditi]

Description

CutlassMoeFCRunner reserves an FP32 split-K partials buffer for the fused SwiGLU MoE GEMV (FP16 activations, INT4 weights, swiglu_fusion == 1, see MoeGemvSplitK2SwiGLUSupported) and sized it for every expanded row of the call. Its only reader is the split-K two-pass GEMV. That path runs only when is_moe_gemv_supported() accepts the problem, which needs expanded_num_rows <= moe_gemv::kMaxProfiledExpandedRows (80), and only with FP32 GEMV accumulation (ORT_MOE_GEMV_FP32_ACCUM=1). Above 80 expanded rows the grouped GEMM runs and the buffer is never read, so at prefill sizes the workspace carried a buffer that grows with the prompt and is never used.

This change caps the partials at 80 expanded rows. Calls with up to 80 expanded rows get exactly the same size as before; larger calls keep a fixed 2 * 80 * 2 * inter_size * 4 bytes (0.94 MiB for inter_size 768).

Tests:

• QMoEWorkspaceTest.SplitKSwiGluPartialsDoNotScaleWithPrefillRows fails on main (per-token growth at GEMV sizes 176,640 B, needs at least 225,857 B) and passes with this change.
• QMoEWorkspaceTest.SplitKSwiGluPartialsCoverGemvRows checks that GEMV-sized calls still get the full buffer.

Validation on A100 80GB (SM80), using layer 0 of Qwen3-30B-A3B-Instruct-2507 (INT4, block 128, top-8, 128 experts) as a standalone QMoE model, main vs this branch:

• Outputs are bitwise identical at 8, 80, 88 and 1024 expanded rows, both with ORT_MOE_GEMV_FP32_ACCUM=1 and with the default.
• With ORT_MOE_GEMV_FP32_ACCUM=1, Nsight Systems shows moe_gemv_splitk_partials_kernel and moe_gemv_splitk_reduce_swiglu_kernel launching at 8 and 80 expanded rows, and no GEMV kernels at 88 and 1024.
• compute-sanitizer memcheck reports 0 errors on those runs. memcheck alone would not catch an undersized buffer, because the partials share one workspace allocation with the other buffers. As a control, halving the cap changes the output at 80 expanded rows while memcheck still reports 0 errors, so the parity check is the meaningful one.
• End to end with the full model through ORT GenAI at a 32,640-token prompt and ORT_MOE_GEMV_FP32_ACCUM=1 (one process per build): process peak 47,248 MiB on main vs 43,152 MiB with this change, identical generated tokens, decode 126.2-126.6 vs 126.3-126.6 tokens/s. The buffer is reserved the same way with the default accumulation. An earlier run of the same change on main at e4234e5 measured a peak of 47,254 vs 43,158 MiB in that mode (two processes per build, identical in each).

Motivation and Context

For Qwen3-30B-A3B at a 32,640-token prompt the partials buffer was 2 * 261,120 * 1,536 * 4 B = 3,059 MiB of workspace that was never read.

Limits:

• Only the FP16 + INT4 + fused-SwiGLU runner allocates this buffer, so other dtypes and activations are unaffected.
• The small capped buffer is still reserved when FP32 accumulation is off. Dropping it entirely would tie the workspace layout to that environment variable, so this PR leaves it.
• Tested on SM80 Linux only. CI covers SM86 and Windows.

[Copilot]

🟡 Changes recommended

The workspace tests check only growth deltas, so they can miss underallocation of split-K partials that allows the kernel to overwrite the next workspace region.

1 open issue

• Workspace tests miss missing split-K partial capacity · New

What changed in this PR

This PR reduces CUDA QMoE workspace usage during prefill by capping split-K SwiGLU GEMV partials at the 80-expanded-row dispatch limit, while preserving allocation sizes for GEMV-sized calls.

Changes:

• Caps the partials buffer instead of letting it grow with prompt length.
• Adds regression tests comparing workspace growth at GEMV and prefill sizes.

File	Description
onnxruntime/​test/​contrib_ops/​cuda_kernels/​qmoe_workspace_test.cc	Adds workspace-growth and GEMV-boundary tests.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_kernels.cu	Caps partials allocation using the shared GEMV row limit.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

LGTM — capping moe_gemv_splitk_rows directly at moe_gemv::kMaxProfiledExpandedRows in CutlassMoeFCRunner::getWorkspaceDeviceBufferSizes cleanly eliminates ~3–4 GiB of unused workspace at 32K prefill while keeping getWorkspaceSize() and configureWsPtrs() in lockstep and preserving exact partials capacity for all shapes accepted by is_moe_gemv_supported() (expanded_num_rows <= 80).

+1 to the existing inline suggestion on qmoe_workspace_test.cc to also check absolute partials capacity via swiglu_fusion=1 minus swiglu_fusion=0 so constant-offset under-allocations are caught as well.

[Tianlei Wu (tianleiwu)]

Cap to 80 expanded rows might impact prefill throughput. Please do some testing to avoid regression.

[Copilot]

🟢 Approval recommended

No active review issues are recorded for the change reducing CUDA QMoE workspace usage on large prefill calls.

No open issues 🎉

1 resolved since last review

• Workspace tests miss missing split-K partial capacity

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Thanks for the careful write-up. The cap is correct. I traced every consumer:

• The partials are only written by moe_gemv_splitk_partials_kernel and read by the reduce kernel. Their layout is [SplitK, expanded_num_rows, 2*inter] FP32, so the new formula is exactly the required size.
• That path is only reachable through gemm1 → tryLaunchMoeGemvIntSymmetricInterleavedSwiGLU → is_moe_gemv_shape_supported, which rejects expanded_num_rows > 80. The runtime expanded_num_rows (num_rows * experts_per_token) equals num_moe_inputs, so min(num_moe_inputs, 80) always covers the launch.
• Sizes for 80 rows or fewer are unchanged (Swiglu/SwigluBias are gated, so the old glu_inter_elems was exactly rows*inter*2). getWorkspaceSize() and configureWsPtrs() stay in sync, and callers that size once for a larger shape still work because the size only grows with rows.
• The exact-size test (swiglu_fusion=1 minus swiglu_fusion=0) catches the constant shortfall that the growth-rate test couldn't.

Follow-up on my earlier prefill-throughput comment

• Standard QMoE op: I agree prefill shouldn't change. Kernel selection is gated on rows only, and the row-tile size comes from the row_tile_size option (or the INT2 packed-prefill budget), not from the runner workspace.
• QMoE offload, FP16 + INT4 + swiglu_fusion=1: prefill planning does change here.
  • The INT4 prefill planner (qmoe_offload.cc, capacity loop around L8124) starts at the maximum capacity (≥256 routes) and halves it until getWorkspaceSize(capacity, …) fits effective_gpu_limit (session.moe.offload.device_working_bytes). FP16 uses CutlassMoeFCRunner<half, uint4b_t, half>, and ORT_QMOE_INT4_PREFILL_GEMM is on by default.
  • Each candidate used to include 16 * inter_size bytes of partials per route (48 MiB at 4096 routes for inter 768). That is now capped at 80 rows (0.94 MiB). When the working budget is the limit, the planner can pick a larger wave_route_slots, or use the INT4 CUTLASS runner where it used to fall back. Both change prefill chunking and dispatch.
  • This is probably a speedup, and correctness holds (runMoe there uses 64 ≤ num_r ≤ capacity, and the 80-row cap covers GEMV-sized calls). But the validation above doesn't cover it. test_v4_int4_cutlass_* are skipped in CI (QUICK_BUILD), and the only FP16 case sets no working budget and asserts the maximum capacity. BF16 offload is not affected because its runner never reserves partials.
  • Could you mention this in the description and, if possible, A/B one FP16 offload prefill with a tight device_working_bytes? ORT_QMOE_SHARED_RESIDENT_LOG=1 prints the chosen capacity=.

CI: none of the 3 failures come from this PR.

• Windows GPU CUDA CI Pipeline Test Job fails only test_qmoe_shared_direct.py::test_cuda_int3_tensor_core_compatibility_shadow, which fails the same way on main (https://github.com/microsoft/onnxruntime/actions/runs/43508257837).
• The two Kernel Documentation jobs fail because docs/OperatorKernels.md is out of date on main (int8 Abs/Neg/Relu/Sign).
• The new QMoEWorkspaceTest.* cases run and pass on Linux CUDA and Windows CUDA.

One optional inline suggestion below.

[BBaditi]

I ran the FP16 offload A/B on A100 with one Qwen3-30B-A3B QMoE layer at a 4,096-token prefill (60 warm runs in A-B-B-A order, main vs this PR):

• 256 MiB (default, all 128 experts on GPU): main capacity 2048 (18.10 ms) vs this PR capacity 2048 (18.13 ms)
• 64 MiB (all 128 experts on GPU): main capacity 1024 (19.42 ms) vs this PR capacity 2048 (18.15 ms)
• 48 MiB (all 128 experts on GPU): main MoE GEMV (44.62 ms) vs this PR capacity 1024 (19.39 ms)
• 32 MiB (all 128 experts on GPU): main MoE GEMV (46.13 ms) vs this PR capacity 1024 (27.70 ms)
• 128 MiB (31 on GPU, rest streamed): main capacity 512 (208.7 ms) vs this PR capacity 1024 (208.1 ms)
• 112 MiB (31 on GPU, rest streamed): main MoE GEMV (222.5 ms) vs this PR capacity 256 (213.2 ms)

At 48 MiB, prefill runs as 16 tiles of 256 tokens (2,048 routes) and needs at least 1,024-route capacity when all 128 experts stay on GPU. On main that capacity needs 22 MiB of workspace (12 MiB of it split-K partials), which does not fit, so every tile falls back to MoE GEMV (38.9 ms GPU time). With this PR the workspace drops to 11 MiB, so each tile runs two 1,024-route INT4 CUTLASS grouped GEMM calls instead (11.0 ms GPU time). Outputs are bit-identical whenever both use CUTLASS, and BF16 is identical at every budget.