[WebGPU] Add block-scaled FP8 QMoE

docs/ContribOperators.md

@@ -6121,6 +6121,8 @@ This version of the operator has been available since version 1 of the 'com.micr
         Partial edge blocks are allowed. No zero points or activation scales are used.
         Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead.
         Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization).
+        The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections
+        that exceed 32-bit shader addressing or the device's per-dimension dispatch limit.
 
         Packed byte dimensions are computed as logical_element_count * effective_expert_weight_bits / 8.
         Weight rows must be byte-aligned. Zero-point rows are padded to a whole byte when necessary.
@@ -6158,7 +6160,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>activation_type</tt> : string</dt>
 <dd>Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu</dd>
 <dt><tt>block_size</tt> : int</dt>
-<dd>For integer quantization, size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). Both hidden_size and inter_size must be divisible by the block size. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. Without a positive value, FP8 uses per-expert global scales. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor. </dd>
+<dd>For integer quantization, size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). For integer and FP4 quantization, both hidden_size and inter_size must be divisible by the block size. FP8 with block_size=128 supports partial 128x128 tiles. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. Without a positive value, FP8 uses per-expert global scales. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor. </dd>
 <dt><tt>expert_weight_bits</tt> : int</dt>
 <dd>Number of bits used in quantized weights. Supported values are 2, 4, and 8. Default is 4 bits</dd>
 <dt><tt>fc1_expert_weight_bits</tt> : int</dt>
@@ -8622,5 +8624,3 @@ No versioning maintained for experimental ops.
 <dt><tt>T</tt> : tensor(float)</dt>
 <dd>Constrain input and output types to float32 tensors.</dd>
 </dl>
-
-

onnxruntime/contrib_ops/webgpu/moe/block_fp8_expert_matmul.wgsl.template

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#use .getByOffset .setByOffset
+
+#param has_bias
+#param has_indirect_experts
+#param broadcast_input
+
+$MAIN {
+  let row = global_id.y;
+  let col = global_id.x;
+  if (row >= uniforms.rows || col >= uniforms.cols) {
+    return;
+  }
+
+#if has_indirect_experts
+  let expert_idx = indirect_experts.getByOffset(row);
+  let weight_base = (expert_idx * uniforms.cols + col) * uniforms.inner;
+  let scale_base = expert_idx * uniforms.scale_n_blocks * uniforms.scale_k_blocks +
+                   (col / 128u) * uniforms.scale_k_blocks;
+#else
+  let weight_base = uniforms.weight_byte_offset + col * uniforms.inner;
+  let scale_base = (col / 128u) * uniforms.scale_k_blocks;
+#endif
+#if broadcast_input
+  let input_row = 0u;
+#else
+  let input_row = row;
+#endif
+
+  var value = 0.0;
+  for (var k_block = 0u; k_block < uniforms.scale_k_blocks; k_block++) {
+    let scale = f32(scales.getByOffset(scale_base + k_block));
+    let start = k_block * 128u;
+    let end = min(start + 128u, uniforms.inner);
+    for (var inner = start; inner < end; inner++) {
+      let weight_offset = weight_base + inner;
+      let packed_word = u32(weights.getByOffset(weight_offset / 4u));
+      let weight_code = unpack4xU8(packed_word)[weight_offset % 4u];
+      let weight = bitcast<f32>(kFp8DequantLutBits[weight_code]);
+      value += f32(input.getByOffset(input_row * uniforms.inner + inner)) * weight * scale;
+    }
+  }
+#if has_bias
+#if has_indirect_experts
+  value += f32(bias.getByOffset(expert_idx * uniforms.cols + col));
+#else
+  value += f32(bias.getByOffset(col));
+#endif
+#endif
+  output.setByOffset(row * uniforms.cols + col, output_element_t(value));
+} // MAIN