Add FP8/FP4 block-quantized support to GatherBlockQuantized contrib op

[Copilot]

Description

This PR extends the existing GatherBlockQuantized contrib op (schema + CPU/CUDA/WebGPU kernels) to support FP8 and FP4 block-quantized data, in addition to its existing int2/int4/uint4/uint8 support. This enables a Gather over a low-precision floating point (FP8 or FP4/Float4E2M1) block-quantized table, with a scales input dequantized per block along quantize_axis.

Key semantic additions to GatherBlockQuantized:

• data may be one of float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, or float4e2m1, in addition to the existing quantized integer types. For these types, bits/zero_points are not applicable (FP8/FP4 quantization is symmetric — dequantization is output = float(data) * scales[block_index]).
• block_size may be 0 for FP8/FP4 data, meaning the entire quantize_axis dimension forms a single block (one scale per row).
• scales supports broadcasting (dim == 1) on any axis other than quantize_axis, enabling a true global per-tensor scale — matching Hugging Face's FP8Embedding (weight_scale is a single scalar for the whole table).

Files changed:

• onnxruntime/core/graph/contrib_ops/contrib_defs.cc — schema (new FP8/FP4 T1 types, block_size == 0 and scale-broadcast semantics) and shape inference.
• onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc — CPU kernel: FP8/FP4 dequantization, per-axis broadcast strides.
• onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc — CPU kernel registrations for the new FP8/FP4 type combinations.
• onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.{cc,cu,cuh} — CUDA kernel support for FP8/FP4 dequantization and scale broadcasting.
• onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.{cc,h} — WebGPU kernel support: host-built WGSL dequantization lookup tables for FP8/FP4 codes (since FP8/FP4 aren't native WGSL shader variable types), and scale-broadcast handling matching the CPU/CUDA kernels.
• onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc — new tests covering FP8/FP4 data, per-row and per-tensor (broadcast) scales, block_size == 0, negative indices, non-leading gather axes, and float16 output, across CPU/CUDA/WebGPU.
• docs/ContribOperators.md / docs/OperatorKernels.md — regenerated docs for the updated schema and kernel registrations.
• onnxruntime/test/contrib_ops/engram_ops_test.cc — ported RunNGramHashMappingEosAcrossChunksTest forward to the base branch's current input-based eos_token_id / reset_on_eos API (filling an EOS-crossing-chunk-boundary coverage gap), since the op's schema had moved on from an attribute-only approach.

Motivation and Context

Support for Qwen3.8-Flash / Qwen4-Exp requires an n-gram embedding gather over a block-quantized FP8/FP4 table. Extending GatherBlockQuantized (rather than adding a separate op) keeps a single Gather-with-quantization op that already has CPU/CUDA/WebGPU coverage and existing graph-optimization/QDQ integration, while adding the FP8/FP4 dtype and broadcast-scale support needed for that table.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Multiple correctness issues were found in the new/updated tests, shape inference, and runtime validation logic (including contract mismatches and a potential divide-by-zero) that must be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR introduces a new com.microsoft contrib operator (GatherFpQuantized) for CPU execution to gather from FP8/FP4 block-scaled tables and dequantize on the fly, and it also updates tests/docs around the Engram/Qwen n-gram hashing functionality.

Changes:

• Add GatherFpQuantized schema + shape inference, CPU kernel implementation, kernel registration, unit tests, and generated documentation entries.
• Update NGramHashMapping reference/test coverage to include an EOS-across-chunks scenario and adjust generated operator documentation text.

File summaries

File	Description
onnxruntime/test/contrib_ops/gather_fp_quantized_op_test.cc	New unit tests for GatherFpQuantized (FP8/FP4, broadcast scales, output types, invalid block size).
onnxruntime/test/contrib_ops/engram_ops_test.cc	Updates NGramHashMappingReference and adds chunked EOS boundary test cases.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Adds GatherFpQuantized operator schema and shape inference to the MS contrib domain.
onnxruntime/contrib_ops/cpu/quantization/gather_fp_quantized.h	Declares the CPU kernel and compute helpers for GatherFpQuantized.
onnxruntime/contrib_ops/cpu/quantization/gather_fp_quantized.cc	Implements the CPU kernel, including scale shape validation and dequantization logic.
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc	Registers GatherFpQuantized CPU kernels for FP8/FP4 data types.
docs/OperatorKernels.md	Adds the GatherFpQuantized kernel listing in the generated kernel table.
docs/ContribOperators.md	Adds generated docs for GatherFpQuantized and updates generated docs text for NGramHashMapping.

Review details

Suppressed comments (3)

onnxruntime/test/contrib_ops/gather_fp_quantized_op_test.cc:123

• Close the DISABLE_FLOAT8_TYPES guard before the Float4-only test section so Float4 builds still compile/run this file.

#if !defined(DISABLE_FLOAT4_TYPES)

onnxruntime/test/contrib_ops/engram_ops_test.cc:328

• The reference implementation masks shifts crossing EOS boundaries to pad_id, but the kernel substitutes eos_token_id (see contrib_ops/cpu/bert/ngram_hash_mapping.cc:150-156). This will make the new EOS-related expectations incorrect.

        if (k > 0 && eos_token_id.has_value()) {
          saw_eos = saw_eos || token == static_cast<T>(*eos_token_id);
          if (saw_eos) {
            token = static_cast<T>(pad_id);
          }

onnxruntime/core/graph/contrib_ops/contrib_defs.cc:4416

• With block_size == 0, effective_block_size becomes data_dim; if data_dim is 0 this skips validation and can lead to runtime divide-by-zero. Instead, require the scales dim on quantize_axis to be exactly 1 whenever block_size == 0, and otherwise require ceil(data_dim/block_size) blocks.

              int64_t effective_block_size = block_size == 0 ? data_shape.dim(i).dim_value() : block_size;
              if (effective_block_size > 0 &&
                  (data_shape.dim(i).dim_value() + effective_block_size - 1) / effective_block_size !=
                      scales_shape.dim(i).dim_value()) {
                fail_shape_inference("data shape and scales shape do not match");

• Files reviewed: 8/8 changed files
• Comments generated: 7
• Review effort level: Lite

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

You can commit the suggested changes from lintrunner.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

🟡 Changes recommended

CUDA support is not registered and has indexing/contract defects, while the registered WebGPU path is known to fail.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 12/12 changed files
• Comments generated: 10
• Review effort level: Balanced

[Copilot]

Fixed: effective_block_size now clamps to std::max<int64_t>(data_shape[quantize_axis_], 1) when block_size_ == 0, avoiding the divide-by-zero for empty quantize-axis dims.

[Copilot]

Fixed: rewrote GatherBlockQuantizedFpKernel to decompose in_idx into per-axis indices via data_dims and compute the quantize-axis block index per-row (instead of a flat block_id = in_idx / block_size), so it now correctly resets at row boundaries when the quantize-axis dimension isn't divisible by block_size.

[Copilot]

Fixed: clamped effective_block_size to std::max<int64_t>(x_shape[quantize_axis], 1) before it's used, avoiding the divide-by-zero for empty quantize-axis dims.

[Copilot]

Fixed: removed FP8/FP4 types from GatherBlockQuantizedT1Constraint() so WebGPU no longer claims these nodes (falls back to CPU) until the shader path is validated on real hardware. Left the existing dequant/shader code in place (now unreachable) with a comment explaining how to re-enable it once verified, and updated the test file's exclusion-list comment accordingly.

[Copilot]

Fixed: added class declarations and BuildKernelCreateInfo registrations in cuda_contrib_kernels.cc for all FP8 (Float8E4M3FN/FNUZ, Float8E5M2/FNUZ) and FP4 (Float4E2M1x2) GatherBlockQuantized variants across {float, MLFloat16, BFloat16} x {int32_t, int64_t}, guarded by the existing DISABLE_FLOAT8_TYPES/DISABLE_FLOAT4_TYPES pattern used elsewhere in the file.

[Copilot]

Fixed: since CUDA forces quantize_axis_ to be the last data axis, in_idx can be decomposed into per-axis indices via a new data_dims TArray. Replaced the old all-or-nothing broadcast rejection with a generic per-axis scale_strides/scale_broadcast_axis construction (0 stride/contribution wherever the scales dim is 1), so arbitrary partial-broadcast shapes on any non-quantize axis are now supported, matching the schema and CPU kernel. Verified by hand against all existing CPU-side FP8/FP4 tests including the 3D non-leading-gather-axis broadcast case.

[Copilot]

Fixed: tightened the condition to block_size == 0 ? is_fp_quantized : (block_size >= 16 && (block_size & (block_size - 1)) == 0), rejecting invalid nonzero values like 8 or 24 at Graph::Resolve(), matching the kernel's ORT_ENFORCE.

[Copilot]

Fixed: effective_block_size is now clamped to std::max<int64_t>(data_shape.dim(i).dim_value(), 1) when block_size == 0, so the scale-shape validation is no longer skipped for an empty quantize-axis dimension (matches the CPU kernel's treatment of the divisor as 1).