Skip to content

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

Merged
Tianlei Wu (tianleiwu) merged 76 commits into
mainfrom
copilot/research-onnx-runtime-support
Sep 15, 2026
Merged

Tianlei Wu (tianleiwu) merged 76 commits into
mainfrom
copilot/research-onnx-runtime-support

Conversation

Copilot AI commented Sep 8, 2026 •

Copy link
Copy Markdown
Contributor

Description

This PR extends GatherBlockQuantized (schema + CPU/CUDA/WebGPU kernels) to support FP8 and FP4 block-quantized data in addition to its existing integer formats.

Key semantic additions and fixes:

  • Adds float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz, and float4e2m1 data. For these types, bits is ignored and zero_points is omitted.
  • Allows block_size = 0 for FP8/FP4, meaning the full quantize_axis dimension is one block.
  • Allows FP scales to broadcast on axes other than quantize_axis.
  • Changes existing integer behavior as well as FP behavior: out-of-range indices now consistently produce zero output slices on CPU, CUDA, and WebGPU. This deliberately differs from ONNX Gather; it avoids provider-dependent behavior and preserves asynchronous CUDA execution without stream synchronization. Op-owner sign-off is requested for this contract change.
  • Aligns CUDA defaults with the schema (block_size = 128, quantize_axis = 1) and validates axes, ranks, scale shapes, empty outputs, launch dimensions, and integer bits values.
  • Validates block sizes and integer bits during graph resolution; bits remains intentionally ignored for FP8/FP4.
  • Makes the WebGPU FP8/FP4 path reachable and tests constant FP tables across supported providers.

Tests cover CPU/CUDA/WebGPU behavior, defaults, BFloat16 output, packed FP4 row boundaries, negative axes and indices, invalid indices, empty outputs, and invalid attributes. Generated contrib operator/kernel documentation is included.

Motivation and Context

Qwen3.8-Flash / Qwen4-Exp requires an n-gram embedding gather over a block-quantized FP8/FP4 table. Extending GatherBlockQuantized keeps a single Gather-with-quantization op while adding the dtype and broadcast-scale support required by that table.

Copilot AI and others added 30 commits August 25, 2026 20:18
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…eek-engram' into copilot/add-op-support-for-qwen-3-8-flash-next

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ids, head_offsets

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…utput

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ids, head_offsets

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…utput

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…) support

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
This was referenced Oct 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants