[Baiju Meswani (baijumeswani)]

Summary

Add optional row tiling to the CUDA grouped-QMoE path to bound temporary workspace usage for larger inputs.

Tiling is disabled by default, preserving the existing execution path. Applications can configure the maximum rows per tile for each ORT session:

"session_options": {
  "ep.cuda.qmoe_row_tile_size": "64"
}

Each tile runs sequentially on the same CUDA stream and reuses its routing metadata and runner workspace.

Changes

• Add row-tile planning and bounded scratch-layout helpers.
• Reuse the grouped-MoE workspace across sequential row tiles.
• Add the ep.cuda.qmoe_row_tile_size session configuration entry.
• Keep ORT_QMOE_ROW_TILE_SIZE as a compatibility fallback.
• Add optional QMoE diagnostics for route, tile size, workspace size, and tactic bucket.

Qualification

• End-to-end generation produced identical token output with tiling enabled.
• In the qualified large MoE workload, a 64-row tile reduced QMoE temporary scratch from approximately 64 MB to 32 MB.

Row tiling remains opt-in.

[Copilot]

Pull request overview

Adds optional row tiling to the CUDA QMoE grouped-MoE execution path to bound temporary workspace usage for large inputs while preserving the existing untiled behavior by default. The change introduces tiling planning/scratch sizing helpers, threads a new session config (ep.cuda.qmoe_row_tile_size) into the QMoE kernel, and adds parity/unit tests for tiling behavior.

Changes:

• Add RowTilePlan / ScratchLayout helpers to plan row tiles and bound routing scratch to tile size.
• Update CUDA QMoE kernel to profile/tune per distinct tile row-count and execute sequential row tiles reusing runner workspace + routing buffers.
• Add unit tests (C++) and a Python parity test validating tiled vs untiled output for an INT4 grouped-MoE shape.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_qmoe_cuda.py	Adds a configurable session option in the test harness and a tiled-vs-untiled parity test.
onnxruntime/test/contrib_ops/qmoe_row_tiling_test.cc	New unit tests for row tile planning/validation and scratch sizing.
onnxruntime/contrib_ops/cuda/moe/qmoe_row_tiling.h	New helper types for row-tiling plan computation and bounded scratch layout calculation.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Adds QMoE row-tiling config/env keys and kernel members to store row tile size + debug flag.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Implements row tiling in QMoE compute, per-tile profiling/tactic selection, bounded scratch allocation, and optional diagnostics.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

The row-tiling math, checked pointer offsets, per-tile tactic/workspace sizing, and same-stream scratch reuse look sound. I found one hot-path performance issue that should be addressed before merge; details are inline. The two existing diagnostics threads remain applicable and are not duplicated here.

[Tianlei Wu (tianleiwu)]

Re-reviewed at bd52958b. All three threads from the previous round are correctly addressed:

• The per-inference host allocations are gone. std::array<..., 2> plus explicit counters is actually stronger than what I suggested: a function-local array cannot reallocate, so the RunnerTileConfig* used later in the tile loop is provably stable where the old iterator was not. I verified the bounds hold on every path (the deterministic and profiling population loops are mutually exclusive, so the counter never exceeds 2), and that dropping #include <vector> is safe since no std::vector is named in the file and auto tactics = active_runner->getTactics() deduces it through moe_kernels.h.
• The range-validation message now names the real knob. Correct in all four combinations of config/env set/unset and parse-vs-range failure.
• The GEMV diagnostic now reports bucketM(expanded), matching the tune key that route actually uses, while the grouped route keeps tile-row buckets.

One non-blocking note that falls outside the diff, so I could not anchor it inline: avg_tokens_per_expert and route_native_fp4 are computed from moe_params.num_rows and active_runner is selected once before the tile loop, but each runMoe only ever sees tile_rows. With tiling enabled on an FP4 prefill, the native-vs-dense crossover heuristic can pick the runner suited to the whole batch rather than the tile. Tactics are profiled per tile M, so this is a heuristic/performance mismatch rather than a correctness issue, and worth revisiting if tiling is ever enabled by default for FP4.

Test coverage worth adding later: CUDA graph capture combined with row_tile_size (the per-tile capture-safe fallback is currently unexercised), tiling on the fp4/nvfp4/fp8 routes that can actually tile, and an exact-divisibility case such as 64 rows with tile 16 to cover the single-entry branch.

Approving: the remaining items are optional follow-ups on a path that is off by default.