From 333ce619ef920ddd94a1ada6cb1217791a0044c6 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:22:12 -0700 Subject: [PATCH] fix(cuda): accept fp16/bf16 router_probs, scales and biases in QMoE The native CUDA `com.microsoft::QMoE` kernel hardcoded `router_probs`, fc1/fc2/fc3 `scales`, expert biases and the optional aggregation weights to `Float32`, so an all-fp16 fused QMoE export (e.g. the Mobius fp16 build) was rejected at runtime (`router_probs requires Float32, got Float16`) even though its activations were already handled by the fp16 kernels. router_probs was merely the first check to fail; scales would have failed next. ORT's `com.microsoft::QMoE` binds input, router_probs, scales, biases and output to a single type parameter `T` = {float, float16, bfloat16}, so a valid fp16 graph carries fp16 versions of all these operands. Classify each T-typed float operand independently: f32 operands keep their original pointers (the pure-f32 path and the pre-existing "fp16 activations + f32 scales" graphs stay byte-identical), while fp16/bf16 operands are losslessly widened to f32 scratch via new `qmoe_widen_{f16,bf16}_f32` kernels and fed to the unchanged f32 routing/dequant kernels. Routing (max/expf/sum/normalize/top-k) still accumulates in f32. Adds an all-fp16/bf16 GPU regression (router_probs + all scales + biases + aggregation weights half-typed) asserting conformance to a half-rounded CPU reference across decode and prefill. Measured: `qwen15-moe-qmoe-mobius` (all fp16) now loads and decodes coherent text natively on CUDA with graph capture (captures=1 replays=22 fallbacks=0); `qwen15-moe-qmoe-f32` greedy token ids are byte-identical to the pre-change binary; qmoe_gpu suite 33 passed / 0 failed / 0 ignored under cuda,gpu-tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 --- ...copilot-qmoe-fp16-router-probs-resolved.md | 69 +++++ .../onnx-runtime-ep-cuda/src/kernels/qmoe.rs | 278 ++++++++++++++++-- crates/onnx-runtime-ep-cuda/tests/qmoe_gpu.rs | 95 ++++++ docs/execution/CUDA_COVERAGE.md | 2 +- 4 files changed, 412 insertions(+), 32 deletions(-) create mode 100644 .squad/decisions/inbox/copilot-qmoe-fp16-router-probs-resolved.md diff --git a/.squad/decisions/inbox/copilot-qmoe-fp16-router-probs-resolved.md b/.squad/decisions/inbox/copilot-qmoe-fp16-router-probs-resolved.md new file mode 100644 index 0000000000..c9562bc176 --- /dev/null +++ b/.squad/decisions/inbox/copilot-qmoe-fp16-router-probs-resolved.md @@ -0,0 +1,69 @@ +# Decision: the "fc1 loader blocker" / all-fp16 QMoE rejection is RESOLVED (native CUDA now runs all-fp16 fused QMoE) + +**Author:** Copilot (coding agent) +**Date:** 2026-08-20 +**Scope:** native CUDA `com.microsoft::QMoE`, fp16/bf16 fixtures +**Related:** corrects the phantom in `.squad/identity/now.md:71` and the stale +mixed-precision root-cause note in `.squad/decisions-archive/2026-08.md` (the +2026-08-03 "loader blocker" entry). Those are append-only history and are NOT +rewritten — this decision supersedes them. + +## What was believed (now falsified) + +Two stale framings existed: + +1. A *loader* blocker: the fused fp16 QMoE artifact was thought to be a + **mixed-precision** graph (fp16 activations + **fp32** scales/router_probs) + that ORT rejects at `Session::new` because `com.microsoft::QMoE` binds + `input`, `router_probs`, `scales`, biases and output to a single type + parameter `T`. +2. That the *only* native-kernel blocker for an all-fp16 graph was + `router_probs` hardcoded to fp32. + +## What is actually true (measured, 2026-08-20) + +- Verified against ORT's schema (`ContribOperators.md` / `quantization_defs`): + `com.microsoft::QMoE` has a **single type constraint `T` = {float, float16, + bfloat16}** covering input, router_probs, scales, biases and output. So a + valid fp16 export carries fp16 router_probs **and** fp16 scales **and** fp16 + biases/aggregation weights. +- The current `qwen15-moe-qmoe-mobius` fixture is **internally type-consistent + (all fp16)** — it already satisfies ORT's single-`T` rule. The Mobius export + was fixed; the *loader* blocker is gone. `qwen15-moe-qmoe-f32` is all-f32. +- The remaining blocker was entirely in **our** native CUDA kernel + (`crates/onnx-runtime-ep-cuda/src/kernels/qmoe.rs`), and it was **broader than + router_probs**: `router_probs`, `fc1_scales`, `fc2_scales` (and, latently, + fc3 scales, biases and the optional aggregation weights) were all hardcoded to + require `Float32`. router_probs was merely the first check to fail; scales + would have failed next. + +## Fix + +Per-operand widen of the `T`-typed float operands to f32 device scratch (an +exact fp16/bf16→f32 upcast), then reuse the **unchanged** f32 routing/dequant +kernels (which already accept fp16 *activations*). Design notes: + +- **Accumulation precision unchanged:** routing (max / `expf` / sum / normalize + / top-k) still runs in **f32** in `qmoe_route`; dequant scales are f32. No + softmax/top-k is done in fp16. +- **f32 path is byte-identical:** f32 operands keep their original pointers and + never widen. Measured: `qwen15-moe-qmoe-f32` greedy token ids are identical + between the pre-change control binary (69329c900) and the fixed binary. +- **Backward compatible:** the classification is per-operand, so the existing + synthetic "fp16 activations + f32 scales" graphs keep working unchanged; an + all-fp16 graph now runs too. This is a strict superset of ORT's single-`T` + rule (the native backend is the sole validator — it skips ORT `Session::new`). + +## Evidence + +- `qwen15-moe-qmoe-mobius` (all fp16) now loads + decodes coherent text natively + on CUDA: `": I am trying to find the best way to get the current date in a + format that is easy to read. I"`; CUDA graph capture succeeds + (`captures=1 replays=22 fallbacks=0 invalidations=0`). +- Route-selection dump confirms non-degenerate routing (top-k experts vary + widely across tokens/layers, spanning the full expert range — not collapsed). +- `qmoe_gpu` GPU suite: 33 passed / 0 failed / 0 ignored under + `--features cuda,gpu-tests`, incl. a new all-fp16/bf16 regression + (`qmoe_all_half_router_scales_bias_match_rounded_cpu`). + +**Do not re-investigate an fp16 QMoE "loader blocker" — it no longer exists.** diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/qmoe.rs b/crates/onnx-runtime-ep-cuda/src/kernels/qmoe.rs index 1cfa452fe5..01a1eb6e1b 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/qmoe.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/qmoe.rs @@ -36,6 +36,11 @@ const GATE_UP_ACTIVATE_BF16_OCC_ENTRY: &str = "qmoe_gate_up_activate_bf16_occ"; const COMBINE_F32_ENTRY: &str = "qmoe_combine_f32"; const COMBINE_F16_ENTRY: &str = "qmoe_combine_f16"; const COMBINE_BF16_ENTRY: &str = "qmoe_combine_bf16"; +// Element-wise widen of the T-typed float routing/scale/bias inputs to f32 so +// the f32 routing and dequant kernels can be reused unchanged for fp16/bf16 +// graphs. See `FloatDtype::widen_entry`. +const WIDEN_F16_ENTRY: &str = "qmoe_widen_f16_f32"; +const WIDEN_BF16_ENTRY: &str = "qmoe_widen_bf16_f32"; const LINEAR_ONE_TASK_PER_BLOCK_MAX_ROUTES: usize = 16; const CUDA_SRC: &str = r#" @@ -78,6 +83,35 @@ __device__ __forceinline__ int total_order_key(float value) // memory instead of re-issuing latency-bound global loads. At decode (rows=1) // this replaces a single active thread with the whole block, which was the // dominant decode cost. +#if QMOE_HAS_HALF +// Widen a contiguous fp16/bf16 buffer to f32 (grid-strided). Conversion is +// exact: every fp16/bf16 value is representable in f32, so the reused f32 +// routing/dequant kernels see byte-for-byte the authored values. +extern "C" __global__ void qmoe_widen_f16_f32( + const __half* src, + float* dst, + const unsigned long long count) +{ + for (unsigned long long i = + blockIdx.x * (unsigned long long)blockDim.x + threadIdx.x; + i < count; i += (unsigned long long)blockDim.x * gridDim.x) { + dst[i] = __half2float(src[i]); + } +} + +extern "C" __global__ void qmoe_widen_bf16_f32( + const __nv_bfloat16* src, + float* dst, + const unsigned long long count) +{ + for (unsigned long long i = + blockIdx.x * (unsigned long long)blockDim.x + threadIdx.x; + i < count; i += (unsigned long long)blockDim.x * gridDim.x) { + dst[i] = __bfloat162float(src[i]); + } +} +#endif + extern "C" __global__ void qmoe_route( const float* router_probs, const float* router_weights, @@ -1272,6 +1306,12 @@ struct QuantizedExperts<'a> { scales: &'a TensorView<'a>, zero_points: Option<&'a TensorView<'a>>, bias: Option<&'a TensorView<'a>>, + /// f32 device buffer holding widened scales when the graph is fp16/bf16; + /// `None` for f32 (use `scales` directly). The dequant kernels always read + /// f32 scales, so fp16/bf16 scales are widened once per execute. + scales_override: Option, + /// f32 device buffer holding widened biases when the graph is fp16/bf16. + bias_override: Option, out_features: usize, in_features: usize, packed_in: usize, @@ -1307,7 +1347,7 @@ impl<'a> QuantizedExperts<'a> { packed.dtype, DataType::Uint8, )?; - require_dtype(&format!("{name}_scales"), scales.dtype, DataType::Float32)?; + float_widen_entry(&format!("{name}_scales"), scales.dtype)?; let pack_size = 8 / bits; if !in_features.is_multiple_of(pack_size) { return Err(error(format!( @@ -1349,11 +1389,7 @@ impl<'a> QuantizedExperts<'a> { )?; } if let Some(bias) = bias { - require_dtype( - &format!("{name}_experts_bias"), - bias.dtype, - DataType::Float32, - )?; + float_widen_entry(&format!("{name}_experts_bias"), bias.dtype)?; require_shape( &format!("{name}_experts_bias"), bias.shape, @@ -1380,6 +1416,8 @@ impl<'a> QuantizedExperts<'a> { scales, zero_points, bias, + scales_override: None, + bias_override: None, out_features, in_features, packed_in, @@ -1387,6 +1425,22 @@ impl<'a> QuantizedExperts<'a> { zero_point_bytes, }) } + + /// Device pointer to f32 scales: the widened buffer for fp16/bf16 graphs, + /// or the original f32 initializer. + fn scales_ptr(&self) -> CUdeviceptr { + self.scales_override + .unwrap_or_else(|| tensor_ptr(self.scales)) + } + + /// Device pointer to f32 biases (widened for fp16/bf16), or 0 when absent. + fn bias_ptr(&self) -> CUdeviceptr { + match (self.bias_override, self.bias) { + (Some(ptr), _) => ptr, + (None, Some(bias)) => tensor_ptr(bias), + (None, None) => 0, + } + } } impl Kernel for QMoEKernel { @@ -1430,7 +1484,7 @@ impl Kernel for QMoEKernel { outputs[0].dtype, inputs[0].dtype ))); } - require_dtype("router_probs", inputs[1].dtype, DataType::Float32)?; + float_widen_entry("router_probs", inputs[1].dtype)?; if dtype.needs_half_headers() { self.runtime.require_nvrtc_half_headers("QMoE")?; } @@ -1482,7 +1536,7 @@ impl Kernel for QMoEKernel { } let fc1_size = self.attributes.fc1_size(inter)?; - let fc1 = QuantizedExperts::validate( + let mut fc1 = QuantizedExperts::validate( "fc1", &inputs[2], &inputs[3], @@ -1494,7 +1548,7 @@ impl Kernel for QMoEKernel { self.bits, self.block_size, )?; - let fc2 = QuantizedExperts::validate( + let mut fc2 = QuantizedExperts::validate( "fc2", &inputs[5], &inputs[6], @@ -1509,7 +1563,7 @@ impl Kernel for QMoEKernel { let has_fc3 = optional_input(inputs, 8).is_some(); let uses_separate_gate = self.attributes.uses_separate_gate(has_fc3); - let fc3 = if uses_separate_gate { + let mut fc3 = if uses_separate_gate { Some(QuantizedExperts::validate( "fc3", optional_input(inputs, 8) @@ -1541,7 +1595,7 @@ impl Kernel for QMoEKernel { }; if let Some(router_weights) = optional_input(inputs, 14) { - require_dtype("router_weights", router_weights.dtype, DataType::Float32)?; + float_widen_entry("router_weights", router_weights.dtype)?; require_shape("router_weights", router_weights.shape, &[rows, experts])?; } for (name, tensor) in [("input", &inputs[0]), ("router_probs", &inputs[1])] { @@ -1656,9 +1710,115 @@ impl Kernel for QMoEKernel { ) .transpose()?; + // ORT binds input, router_probs, scales, biases and the optional + // aggregation weights to one type parameter T, so a valid fp16/bf16 + // graph carries fp16/bf16 versions of these operands. The routing and + // dequant kernels read them as f32, so widen any non-f32 operand to f32 + // scratch once per execute — an exact, lossless upcast — and reuse the + // identical f32 kernels. Each operand is classified independently: an + // f32 operand keeps its original pointer (so pure-f32 graphs and the + // pre-existing "fp16 activations + f32 scales" graphs are byte-for-byte + // unchanged), while an fp16/bf16 operand is upcast. + let router_elems = checked_product(&[rows, experts], "router element count")?; + let router_probs_ptr = match float_widen_entry("router_probs", inputs[1].dtype)? { + None => tensor_ptr(&inputs[1]), + Some(entry) => self.widen_to_f32( + &mut scratch, + 11, + capturing, + entry, + tensor_ptr(&inputs[1]), + router_elems, + )?, + }; + let router_weights_ptr = match optional_input(inputs, 14) { + None => 0, + Some(rw) => match float_widen_entry("router_weights", rw.dtype)? { + None => tensor_ptr(rw), + Some(entry) => self.widen_to_f32( + &mut scratch, + 12, + capturing, + entry, + tensor_ptr(rw), + router_elems, + )?, + }, + }; + if let Some(entry) = float_widen_entry("fc1_scales", fc1.scales.dtype)? { + fc1.scales_override = Some(self.widen_to_f32( + &mut scratch, + 13, + capturing, + entry, + tensor_ptr(fc1.scales), + checked_product(fc1.scales.shape, "fc1 scales element count")?, + )?); + } + if let Some(entry) = float_widen_entry("fc2_scales", fc2.scales.dtype)? { + fc2.scales_override = Some(self.widen_to_f32( + &mut scratch, + 14, + capturing, + entry, + tensor_ptr(fc2.scales), + checked_product(fc2.scales.shape, "fc2 scales element count")?, + )?); + } + if let Some(fc3) = fc3.as_mut() { + if let Some(entry) = float_widen_entry("fc3_scales", fc3.scales.dtype)? { + fc3.scales_override = Some(self.widen_to_f32( + &mut scratch, + 15, + capturing, + entry, + tensor_ptr(fc3.scales), + checked_product(fc3.scales.shape, "fc3 scales element count")?, + )?); + } + } + if let Some(bias) = fc1.bias { + if let Some(entry) = float_widen_entry("fc1_experts_bias", bias.dtype)? { + fc1.bias_override = Some(self.widen_to_f32( + &mut scratch, + 16, + capturing, + entry, + tensor_ptr(bias), + checked_product(bias.shape, "fc1 bias element count")?, + )?); + } + } + if let Some(bias) = fc2.bias { + if let Some(entry) = float_widen_entry("fc2_experts_bias", bias.dtype)? { + fc2.bias_override = Some(self.widen_to_f32( + &mut scratch, + 17, + capturing, + entry, + tensor_ptr(bias), + checked_product(bias.shape, "fc2 bias element count")?, + )?); + } + } + if let Some(fc3) = fc3.as_mut() { + if let Some(bias) = fc3.bias { + if let Some(entry) = float_widen_entry("fc3_experts_bias", bias.dtype)? { + fc3.bias_override = Some(self.widen_to_f32( + &mut scratch, + 18, + capturing, + entry, + tensor_ptr(bias), + checked_product(bias.shape, "fc3 bias element count")?, + )?); + } + } + } + self.launch_route( - &inputs[1], - optional_input(inputs, 14), + router_probs_ptr, + router_weights_ptr, route_indices, route_weights, rows, @@ -1669,7 +1829,13 @@ impl Kernel for QMoEKernel { // diffed to distinguish a benign borderline-argmax reassociation from a // real router top-k divergence. Safe only outside graph capture. if !capturing && std::env::var_os("ONNX_GENAI_QMOE_ROUTE_DUMP").is_some() { - self.dump_route_selection(&inputs[1], route_indices, rows, experts, self.attributes.k)?; + self.dump_route_selection( + router_probs_ptr, + route_indices, + rows, + experts, + self.attributes.k, + )?; } if let Some(grouping) = grouping { let fc1_output = fc1_output.expect("grouped QMoE keeps FC1 scratch"); @@ -1821,7 +1987,7 @@ impl QMoEKernel { /// any token divergence is purely downstream (GEMV/argmax), not routing. fn dump_route_selection( &self, - router_probs: &TensorView, + router_probs: CUdeviceptr, route_indices: CUdeviceptr, rows: usize, experts: usize, @@ -1847,7 +2013,7 @@ impl QMoEKernel { rows * experts * std::mem::size_of::(), ) }; - unsafe { self.runtime.dtoh(bytes, tensor_ptr(router_probs))? }; + unsafe { self.runtime.dtoh(bytes, router_probs)? }; } for row in 0..rows { let call = CALL.fetch_add(1, Ordering::Relaxed); @@ -1878,18 +2044,48 @@ impl QMoEKernel { Ok(()) } + /// Widen a contiguous fp16/bf16 device buffer (`entry` selects the kernel) + /// into f32 scratch slot `index`, returning the f32 pointer. Conversion is + /// exact, so the reused f32 routing/dequant kernels are numerically + /// unaffected. Only invoked for fp16/bf16 graphs; f32 never widens. + #[allow(clippy::too_many_arguments)] + fn widen_to_f32( + &self, + scratch: &mut ScratchPool, + index: usize, + capturing: bool, + entry: &str, + src: CUdeviceptr, + elements: usize, + ) -> Result { + let bytes = checked_bytes(elements, std::mem::size_of::(), "widened f32 scratch")?; + let dst = scratch.ensure(&self.runtime, index, bytes, capturing)?; + // The widen kernels use __half/__nv_bfloat16, so ensure NVRTC compiles + // the module with the fp16/bf16 headers available before it is resolved. + self.runtime.require_nvrtc_half_headers("QMoE widen")?; + let function = self.runtime.nvrtc_function(MODULE, CUDA_SRC, entry)?; + let count = as_u64("widen element count", elements)?; + let config = self.pointwise_launch_config(count)?; + let mut builder = self.runtime.stream().launch_builder(&function); + builder.arg(&src).arg(&dst).arg(&count); + // SAFETY: `src` holds `elements` fp16/bf16 values and `dst` was sized for + // the same count of f32; the ABI matches `qmoe_widen_*_f32`. + unsafe { builder.launch(config) } + .map(|_| ()) + .map_err(|err| driver_err("widen QMoE fp16 routing/scale input", err))?; + Ok(dst) + } + fn launch_route( &self, - router_probs: &TensorView, - router_weights: Option<&TensorView>, + router_probs: CUdeviceptr, + router_weights: CUdeviceptr, route_indices: CUdeviceptr, route_weights: CUdeviceptr, rows: usize, experts: usize, ) -> Result<()> { let function = self.runtime.nvrtc_function(MODULE, CUDA_SRC, ROUTE_ENTRY)?; - let router_probs = tensor_ptr(router_probs); - let router_weights = router_weights.map(tensor_ptr).unwrap_or(0); let rows = as_u64("row count", rows)?; let experts = as_i32("expert count", experts)?; let top_k = as_i32("top-k", self.attributes.k)?; @@ -2080,9 +2276,9 @@ impl QMoEKernel { .ok_or_else(|| error("grouped GEMM shared-memory stride overflow"))?, )?; let packed = tensor_ptr(weights.packed); - let scales = tensor_ptr(weights.scales); + let scales = weights.scales_ptr(); let zero_points = weights.zero_points.map(tensor_ptr).unwrap_or(0); - let bias = weights.bias.map(tensor_ptr).unwrap_or(0); + let bias = weights.bias_ptr(); let routes = as_u64("route count", routes)?; let tasks = as_u64("grouped linear task count", tasks)?; let gemm_min_tokens = as_u64( @@ -2149,9 +2345,9 @@ impl QMoEKernel { .nvrtc_function(module, source, dtype.linear_entry())?; let packed = tensor_ptr(weights.packed); let expert_counts = expert_counts.unwrap_or(0); - let scales = tensor_ptr(weights.scales); + let scales = weights.scales_ptr(); let zero_points = weights.zero_points.map(tensor_ptr).unwrap_or(0); - let bias = weights.bias.map(tensor_ptr).unwrap_or(0); + let bias = weights.bias_ptr(); let tasks = checked_product(&[routes, weights.out_features], "linear output task count")?; let grid_x = self.linear_reduction_grid(tasks, routes)?; let config = self.runtime.reduction_launch_config( @@ -2255,17 +2451,15 @@ impl QMoEKernel { std::mem::size_of::() as u32, )?; let fc1_packed = tensor_ptr(fc1.packed); - let fc1_scales = tensor_ptr(fc1.scales); + let fc1_scales = fc1.scales_ptr(); let fc1_zero_points = fc1.zero_points.map(tensor_ptr).unwrap_or(0); - let fc1_bias = fc1.bias.map(tensor_ptr).unwrap_or(0); + let fc1_bias = fc1.bias_ptr(); let fc3_packed = fc3.map(|weights| tensor_ptr(weights.packed)).unwrap_or(0); - let fc3_scales = fc3.map(|weights| tensor_ptr(weights.scales)).unwrap_or(0); + let fc3_scales = fc3.map(|weights| weights.scales_ptr()).unwrap_or(0); let fc3_zero_points = fc3 .and_then(|weights| weights.zero_points.map(tensor_ptr)) .unwrap_or(0); - let fc3_bias = fc3 - .and_then(|weights| weights.bias.map(tensor_ptr)) - .unwrap_or(0); + let fc3_bias = fc3.map(|weights| weights.bias_ptr()).unwrap_or(0); let routes = as_u64("route count", routes)?; let top_k = as_i32("top-k", self.attributes.k)?; let inter = as_i32("intermediate feature count", inter)?; @@ -2499,7 +2693,29 @@ impl QMoEKernel { } } -const SCRATCH_SLOTS: usize = 11; +const SCRATCH_SLOTS: usize = 19; + +/// Classifies a QMoE T-typed float operand (router_probs, scales, biases, +/// aggregation weights) and returns the widen kernel that upcasts it to f32, or +/// `None` when it is already f32 (used directly to keep the f32 path +/// bit-identical). Errors for any non-float dtype. +/// +/// ORT's `com.microsoft::QMoE` binds these operands to a single type parameter +/// `T` = {float, float16, bfloat16}, so a valid fp16 graph carries fp16 scales +/// and router probs. Because the native backend is the sole validator (it skips +/// ORT's `Session::new` type check), we accept the stricter superset where each +/// float operand is independently f32 or a half type: an all-fp16 graph runs, +/// and a mixed graph (fp16 activations, f32 scales) keeps working unchanged. +fn float_widen_entry(name: &str, dtype: DataType) -> Result> { + match dtype { + DataType::Float32 => Ok(None), + DataType::Float16 => Ok(Some(WIDEN_F16_ENTRY)), + DataType::BFloat16 => Ok(Some(WIDEN_BF16_ENTRY)), + other => Err(error(format!( + "{name} requires Float32, Float16, or BFloat16, got {other:?}" + ))), + } +} #[derive(Clone, Copy, Debug, Default)] struct ScratchSlot { diff --git a/crates/onnx-runtime-ep-cuda/tests/qmoe_gpu.rs b/crates/onnx-runtime-ep-cuda/tests/qmoe_gpu.rs index 9f71603b12..3d1cc61c76 100644 --- a/crates/onnx-runtime-ep-cuda/tests/qmoe_gpu.rs +++ b/crates/onnx-runtime-ep-cuda/tests/qmoe_gpu.rs @@ -755,6 +755,101 @@ fn compare(case: Case, dtype: DataType) -> (f32, u32) { error_metrics(&actual, &expected) } +/// Read the f32 values back out of an f32-encoded [`HostTensor`]. +fn f32_values(tensor: &HostTensor) -> Vec { + tensor + .bytes + .chunks_exact(4) + .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) + .collect() +} + +/// Re-encode an f32 [`HostTensor`] as a half (fp16/bf16) tensor, matching how a +/// real all-`T` fp16/bf16 QMoE graph stores its router probs, scales and biases. +fn to_half_tensor(tensor: &HostTensor, dtype: DataType) -> HostTensor { + HostTensor::activation(dtype, &tensor.shape, &f32_values(tensor)) +} + +/// Round an f32 [`HostTensor`] through fp16/bf16 back to f32, so the CPU +/// reference dequantizes with the same values the GPU sees after widening. +fn round_through_half(tensor: &HostTensor, dtype: DataType) -> HostTensor { + let rounded: Vec = f32_values(tensor) + .into_iter() + .map(|value| match dtype { + DataType::Float16 => f16::from_f32(value).to_f32(), + DataType::BFloat16 => bf16::from_f32(value).to_f32(), + _ => value, + }) + .collect(); + HostTensor::f32(&tensor.shape, &rounded) +} + +/// Build an all-half QMoE input set (input, router_probs, all scales, all +/// biases and the aggregation weights are fp16/bf16) plus a matching CPU +/// reference input set whose float operands are the same values rounded through +/// the half type. The uint8 packed weights and zero points stay untouched. +fn all_half_inputs( + case: Case, + dtype: DataType, +) -> (Vec>, Vec>) { + let base = case_inputs(case, dtype); + let mut gpu = base.clone(); + let mut cpu = rounded_cpu_inputs(&base, dtype); + // Indices of the T-typed float operands: router_probs, fc1/fc2/fc3 scales, + // fc1/fc2/fc3 biases and the aggregation weights. + for index in [1usize, 3, 4, 6, 7, 9, 10, 14] { + if let Some(tensor) = base[index].clone() { + gpu[index] = Some(to_half_tensor(&tensor, dtype)); + cpu[index] = Some(round_through_half(&tensor, dtype)); + } + } + (gpu, cpu) +} + +/// The all-`T` fp16/bf16 QMoE path: router_probs, scales, biases and the +/// aggregation weights all carry the input's half type (matching ORT's single +/// type parameter `T`, e.g. the fused Mobius fp16 QMoE export). The GPU widens +/// each to f32 and must match a CPU reference dequantized with the same +/// half-rounded values. +fn compare_all_half(case: Case, dtype: DataType) { + let ep = require_cuda(); + let (gpu_inputs, cpu_inputs) = all_half_inputs(case, dtype); + let expected = run_cpu(case, &cpu_inputs); + let actual = run_gpu(&ep, case, &gpu_inputs, dtype).unwrap(); + assert_conforms(&actual, &expected, case, dtype); +} + +fn all_half_rich_case(rows: usize) -> Case { + Case { + experts: 4, + rows, + hidden: 16, + inter: 16, + bits: 4, + top_k: 2, + activation: "silu", + swiglu_fusion: 0, + affine: true, + fc3: true, + biases: true, + normalize: true, + router_weights: true, + } +} + +#[cfg_attr( + not(feature = "gpu-tests"), + ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner" +)] +#[test] +fn qmoe_all_half_router_scales_bias_match_rounded_cpu() { + // rows=1 exercises the decode/fused path; rows=6 exercises prefill grouping. + for rows in [1, 6] { + compare_all_half(all_half_rich_case(rows), DataType::Float16); + compare_all_half(all_half_rich_case(rows), DataType::BFloat16); + } +} + fn qmoe_64expert_case(rows: usize) -> Case { Case { experts: 64, diff --git a/docs/execution/CUDA_COVERAGE.md b/docs/execution/CUDA_COVERAGE.md index 0d9c5be52c..c731f783bf 100644 --- a/docs/execution/CUDA_COVERAGE.md +++ b/docs/execution/CUDA_COVERAGE.md @@ -102,7 +102,7 @@ not yet wired) · **🔬 custom** (needs a fused NVRTC/CUTLASS kernel). | `FusedMatMulBias` | `com.microsoft` | ✅ | **cuBLASLt** `CUBLASLT_EPILOGUE_BIAS` | Dense rank-2 f32/f16/bf16 with an exact per-N bias vector; bias add is fused into GEMM with no elementwise pass. | | `FusedGemm` | `com.microsoft` | ✅ | **cuBLASLt** `CUBLASLT_EPILOGUE_{BIAS,RELU_BIAS,GELU_BIAS}` | Dense rank-2 f32/f16/bf16; transA/transB and α. CUDA 13's `GELU_BIAS` is the tanh/0.044715 GELU approximation (H200 output differs from exact-erf at the expected ~2.2e-4 for x=1.5); cuBLASLt exposes no exact-erf selector, so this deliberately follows the vendor epilogue rather than adding a separate pass. Bias must be per-N and `beta=1` because `BIAS_POINTER` is unscaled; other β values fail explicitly. Missing `activation` defaults to Relu for the repository optimizer's existing `FusedGemm` contract; empty/Identity selects plain BIAS. | | `MatMulNBits` | `com.microsoft` | ✅ | **NVRTC-custom + cuBLASLt** | Standard packed INT4 `[N, ceil(K/block_size), block_size/2]` weights are block-wise dequantized to f32 on-device, then multiplied by f32 activations with full-f32 accumulation. Supports optional packed zero points, group indices, and fused per-N bias. | -| `QMoE` | `com.microsoft` | ✅ | **NVRTC grouped block-dequant expert GEMV/GEMM** | Single-GPU resident: exact ORT expert-major affine INT1/INT2/INT4/INT8 layouts, optional packed zero points and expert biases, top-k routing/normalization, Relu/Gelu/Silu/Identity and fused or separate-gate SwiGLU. Decode uses per-route GEMV; prefill counts and prefix-groups routes by expert, gathers contiguous rows, runs a tiled dequant-GEMM for experts with at least two tokens, then deterministically combines weighted route outputs. f32/f16/bf16 activations use f32 accumulation; tile width, shared-memory fallback, launch width, and saturation grid derive from detected device limits. Native IQ/MXFP4 blocks remain explicitly unsupported because the ORT QMoE affine inputs cannot encode those layouts; they require a block-quantized MoE operator. Weight paging/prefetch and expert-parallel sharding are deferred. | +| `QMoE` | `com.microsoft` | ✅ | **NVRTC grouped block-dequant expert GEMV/GEMM** | Single-GPU resident: exact ORT expert-major affine INT1/INT2/INT4/INT8 layouts, optional packed zero points and expert biases, top-k routing/normalization, Relu/Gelu/Silu/Identity and fused or separate-gate SwiGLU. Decode uses per-route GEMV; prefill counts and prefix-groups routes by expert, gathers contiguous rows, runs a tiled dequant-GEMM for experts with at least two tokens, then deterministically combines weighted route outputs. f32/f16/bf16 activations use f32 accumulation; the other T-typed float operands (router_probs, fc1/fc2/fc3 scales, expert biases, and the optional aggregation weights) may independently be f32 or the same half type as the input — matching ORT's single type parameter `T` — and any half operand is losslessly widened to f32 before the unchanged f32 routing/dequant kernels run, so an all-fp16 fused QMoE export runs natively and the pure-f32 path stays byte-identical. tile width, shared-memory fallback, launch width, and saturation grid derive from detected device limits. Native IQ/MXFP4 blocks remain explicitly unsupported because the ORT QMoE affine inputs cannot encode those layouts; they require a block-quantized MoE operator. Weight paging/prefetch and expert-parallel sharding are deferred. | ### Convolution