onnxruntime

43989a71 - QMoE CUDA: input validation, prepack cleanups, and packaging pipeline fix (#28607)

QMoE CUDA: input validation, prepack cleanups, and packaging pipeline fix (#28607) ## Description Follow-up to #28583. Addresses review feedback that landed after merge (input validation, redundant memset, dead branches in `PrePackComputeBias`) and fixes a pre-existing latent CUTLASS issue that surfaced as a packaging pipeline failure once MoE GEMM kernels were built with a multi-arch `CMAKE_CUDA_ARCHITECTURES` list spanning pre-Ampere and Ampere+ targets. ## Summary of Changes ### Packaging pipeline build fix | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h` | Replace the unconditional `static_assert(false, ...)` in the pre-Ampere `#else` branch of `MoeFCGemm::operator()` with `CUTLASS_NOT_IMPLEMENTED()` plus a comment explaining why this is safe. | Background: `moe_gemm_kernels_*.cu` instantiate `MoeFCGemm` through `MoeGemmRunner<...>::dispatchToArch`, which contains *runtime* (not `constexpr`) `if (sm_ >= 80 && sm_ < 90)` branches. NVCC therefore instantiates the kernel for every requested device target, including pre-Sm80 device compile passes. The old `static_assert(false, ...)` fired on those passes whenever `CMAKE_CUDA_ARCHITECTURES` contained any arch below 80 (e.g. the packaging pipeline list `52-real;61-real;75-real;86-real;89-real;90-virtual`). Replacing it with `CUTLASS_NOT_IMPLEMENTED()` lets NVCC emit a runtime trap stub for pre-Sm80, while runtime dispatch in `MoeGemmRunner::dispatchToArch()` already guarantees `sm_ >= 80` before the kernel is ever launched, so the stub is unreachable in practice. ### Address PR #28583 post-merge review | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu` | Add `ValidateScaledZP4BitBatchedArgs` (positive `experts`/`n`/`k_blocks`, `experts ≤ 65535` for the `gridDim.z` limit) and call it from both `LaunchQMoEScaledZP4BitBatched` overloads. Matches the validation style of `LaunchQMoERepackFP4ColToRow`. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` (`PrePackSwizzleBlockScales`) | Remove the redundant `cudaMemsetAsync` of the destination buffer. `QMoEBlockScaleInterleaveKernel`'s `(batch, row, col) -> offset` map is a bijection over the padded output extent and writes 0 for padded source positions, so every output byte is already written. Comment explains the invariant. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` (`PrePackComputeBias`, 4-bit block-wise) | Add `ORT_ENFORCE` checks for positive shape dims and an `INT_MAX/2` bound on `packed_k_blocks` (parity with `PrePackSwizzleBlockScales` / `PrePackRepackFP4Weights`). Drop the shadowed `bool is_fp16 = is_fp16_; bool is_bf16 = !is_fp16_;` locals in favour of `is_fp16_`. Replace the dead-branch ternary `(is_fp16 \|\| is_bf16 ? 2 : 4)` with `sizeof(uint16_t)` and a clarifying comment, and remove the unreachable `else ORT_THROW(...)` (the QMoE type path is strictly FP16/BF16). | ## Testing - Built locally with CUDA 12.8 against the failing CI arch list (`-DCMAKE_CUDA_ARCHITECTURES="52-real;61-real;75-real;86-real;89-real;90-virtual"`) and confirmed `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_bf16_bf16.cu.o` compiles cleanly (only an `sm_<75` deprecation warning, no `static_assert` failure). - Existing QMoE Python tests (`onnxruntime/test/python/transformers/test_qmoe_cuda.py`, `test_qmoe_cpu.py`) exercise the affected `PrePackSwizzleBlockScales` / `PrePackComputeBias` paths under `--config Debug` builds and continue to pass; the added `ORT_ENFORCE` checks only trigger on invalid shapes that are not produced by the supported QMoE input contract. - No behaviour change on supported devices: `dispatchToArch` already gates `MoeFCGemm` behind `sm_ >= 80`, so the new `CUTLASS_NOT_IMPLEMENTED()` stub is unreachable at runtime. ## Motivation and Context Once #28583 enabled the MoE GEMM kernels as part of the contrib CUDA build, packaging pipelines (which target a wide arch range to maximise GPU coverage) started failing on the pre-Ampere device compile passes. The kernel-side fix in this PR resolves the immediate breakage while keeping the cmake-level binary-size optimisation (per-kernel arch pinning, TensorRT-LLM style) as a follow-up — CMake's `CUDA_ARCHITECTURES` is target/directory-scoped only, so the proper way to restrict per-kernel archs is an OBJECT-library refactor, which is intentionally not in scope here. ## Checklist - [x] Tests added/updated (input validation covered by existing QMoE tests; the new `ORT_ENFORCE` checks fail loudly on out-of-contract shapes) - [x] No documentation changes needed - [x] No breaking changes - [x] Local packaging-pipeline arch list verified to compile