[Tianlei Wu (tianleiwu)]

Description

Adds three com.microsoft contrib operators—GatherPacked (CPU and CUDA), MatMulPacked (CUDA), and MoEPacked (CUDA)—that consume 2D and 3D block-packed quantized weight blobs (such as GGUF IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ4_NL, IQ4_XS, Q2_0, Q4_0, Q5_0, Q5_K, Q6_K, Q8_0, F16, BF16, F32) directly in their native block layout without offline weight repacking or separate scale tensors. Direct external initializer buffers (#33877) avoid redundant host-side copies while loading .gguf shards. CUDA execution still copies CPU-backed weights to device memory and requires device storage for the packed weights; it does not execute zero-copy from host-backed shards.

Summary of Changes

Contrib Operator Schemas & Registration

File	Change
onnxruntime/core/graph/contrib_ops/ms_opset.h	Declare GatherPacked, MatMulPacked, and MoEPacked in com.microsoft opset
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Register schemas and shape inference for GatherPacked, MatMulPacked (A, B -> Y with K, N, packed_type), and MoEPacked
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc	Register CPU GatherPacked kernel
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Register CUDA GatherPacked, MatMulPacked, and MoEPacked kernels

Block-Packed Quantization & Kernel Implementations

File	Change
onnxruntime/contrib_ops/cpu/quantization/gguf_quant_tables.h	Shared block-quantization lookup tables (iq2xxs_grid, iq2xs_grid, iq2s_grid, iq3xxs_grid, iq3s_grid), block geometry helper (GgmlRowBytes), and CPU row dequantizer
onnxruntime/contrib_ops/cpu/quantization/gather_packed.cc	Multi-threaded CPU GatherPacked implementation
onnxruntime/contrib_ops/cuda/quantization/gguf_quant_cuda.cuh	CUDA sub-block decoders (DecodeGgufSubBlock8Impl, DecodeGgufSubBlock8SmemLutTpl, DecodeGgufSubBlock16Impl), register-only Iq2xsSignsDev / Iq4NlValueDev helpers, and Tensor Core MmaBf16 wrapper
onnxruntime/contrib_ops/cuda/quantization/gather_packed.{h,cc,cu}	CUDA GatherPacked kernel
onnxruntime/contrib_ops/cuda/quantization/matmul_packed.{h,cc,cu}	CUDA MatMulPacked kernel with warp-K-split Tensor Core (mma.sync.aligned.m16n8k16) and warp GEMV paths for decode (M <= 8), bounded tiled dequantization + cuBLAS GEMM for prefill (M > 8), and FP32 decoded weights for FP16 prefill when the packed format is not F16
onnxruntime/contrib_ops/cuda/moe/moe_packed.{h,cc,cu}	CUDA MoEPacked kernel with fused top-k softmax routing, parallel warp-scan tile bucketing (8 routes/tile in decode, 32 routes/tile in prefill), warp-K-split Tensor Core FC1+SwiGLU in decode, 4-subtile Tensor Core FC1/FC2 in prefill, and 128-bit vectorized finalize
onnxruntime/test/python/transformers/test_gguf_packed_ops.py	Unit tests verifying GatherPacked (CPU vs CUDA bit-exactness across 15 formats), MatMulPacked (M = 1, 7, 8, 32), and MoEPacked (M = 4, 48)

Testing

• Additional regression validation on NVIDIA A100 with CUDA 13: 30 tests and 757 subtests pass both with the normal SM80/SM90a build and with the packed-op kernels compiled as compute_75-only PTX, exercising BF16 SIMT fallback.
• Configured-device capability probing and restoration were verified on both success and failure; focused lintrunner checks pass.
• Unit tests:

  python -m unittest onnxruntime/test/python/transformers/test_gguf_packed_ops.py -v

• End-to-end validation on Qwen3.8-Flash-Next-GGUF (UD-IQ1_S, UD-IQ2_XS, UD-IQ3_XXS) on NVIDIA H200:
  • engine_behaviours.py: 27 / 27 PASS across all three quantization variants.
  • Prefill (TTFT): 0.048–0.052 s at 128 tokens, 0.222–0.244 s at 2K tokens (8,400–9,226 tok/s), and 0.851–0.944 s at 8K tokens (8,682–9,622 tok/s, 17.1x–18.5x speedup over baseline SIMT).
  • Speculative decode (DFlash2 S=8): 20.0–20.9 ms/cycle (183–269 tok/s median, up to 359 tok/s peak).
  • MMLU-Pro (200 items, max_new=8192): 87.50% (UD-IQ1_S), 85.50% (UD-IQ2_XS), 86.00% (UD-IQ3_XXS).

Checklist

• Tests added/updated
• No breaking changes (or documented in description)
• Documentation updated (if applicable)

[Copilot]

🟡 Changes recommended

Unresolved critical defects can cause initialization crashes, invalid memory reads, CUDA launch failures, and incorrect or undefined packed-operator outputs.

11 open issues

• Reject empty vocabularies before gathering rows · New
• Reject top-k values greater than the expert count · New
• Handle shared-memory requests above CUDA's default limit · New
• Gate BF16 Tensor Core dispatch on hardware support · New
• Reject packed formats without a supported decoder · New
• Keep all warp lanes active in partial-tile MMA · New
• Regenerate and commit contrib operator documentation · New
• Stage device initializers into owned CPU tensors · New
• Decode external weights into native-endian float storage · New
• Disable CPU fallback in the CUDA decoder test · New
• Compare MoEPacked outputs against a numerical reference · New

What changed in this PR

Adds block-packed weight support to ONNX Runtime through three contrib operators, enabling embedding lookup, matrix multiplication, and mixture-of-experts execution without offline weight repacking.

Changes:

• Registers GatherPacked for CPU/CUDA and MatMulPacked and MoEPacked for CUDA.
• Adds CUDA dequantization, decode, prefill, and expert-routing kernels with Python tests.
• Extends external-initializer buffer borrowing and device reuse, with loading tests and documentation.

File	Description
onnxruntime/​test/​shared_lib/​test_model_loading.cc	Adds direct-buffer session-loading coverage.
onnxruntime/​test/​python/​transformers/​test_gguf_packed_ops.py	Tests packed formats and decode/prefill paths.
onnxruntime/​test/​framework/​inference_session_test.cc	Tests initializer copying and pointer reuse.
onnxruntime/​core/​session/​inference_session.cc	Passes the direct-buffer option to graph injection.
onnxruntime/​core/​graph/​graph.cc	Extends external-initializer buffer retention.
onnxruntime/​core/​graph/​contrib_ops/​ms_opset.h	Adds the three operators to the Microsoft opset.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Defines schemas and shape inference.
onnxruntime/​core/​framework/​tensorprotoutils.cc	Extends in-memory serialization to small device tensors.
onnxruntime/​core/​framework/​session_state_utils.cc	Reuses graph initializers on their planned device.
onnxruntime/​contrib_ops/​cuda/​quantization/​matmul_packed.h	Declares packed matrix-multiplication interfaces.
onnxruntime/​contrib_ops/​cuda/​quantization/​matmul_packed.cu	Implements decode kernels and cuBLAS prefill.
onnxruntime/​contrib_ops/​cuda/​quantization/​matmul_packed.cc	Validates inputs and dispatches CUDA computation.
onnxruntime/​contrib_ops/​cuda/​quantization/​gguf_quant_cuda.cuh	Adds shared CUDA decoders and BF16 Tensor Core helpers.
onnxruntime/​contrib_ops/​cuda/​quantization/​gather_packed.h	Declares CUDA packed-gather interfaces.
onnxruntime/​contrib_ops/​cuda/​quantization/​gather_packed.cu	Implements gathered-row dequantization.
onnxruntime/​contrib_ops/​cuda/​quantization/​gather_packed.cc	Registers and validates CUDA GatherPacked.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.h	Defines expert-routing parameters and interfaces.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cu	Implements routing, expert computation, and output accumulation.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_packed.cc	Validates expert tensors and chunks execution.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers the three CUDA operators.
onnxruntime/​contrib_ops/​cpu/​quantization/​gather_packed.cc	Implements threaded CPU GatherPacked.
onnxruntime/​contrib_ops/​cpu/​cpu_contrib_kernels.cc	Registers CPU GatherPacked variants.
include/​onnxruntime/​core/​session/​onnxruntime_session_options_config_keys.h	Documents expanded direct-buffer option behavior.
include/​onnxruntime/​core/​graph/​graph.h	Adds the graph-injection borrowing parameter.
docs/​Model_Loading_Performance.md	Describes external-buffer borrowing and device reuse.

_{Review effort: Balanced}

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

🟡 Changes recommended

Unresolved critical CUDA alignment, dimension-range, and BF16 dispatch defects can cause faults or silently incorrect results, while full-matrix dequantization can exhaust device memory.

16 open issues

• Range-check MoE dimensions before narrowing casts · New
• Use scalar finalize for misaligned output pointers · New
• Validate K and grid limits before CUDA launches · New
• Make packed decoder reads safe for unaligned weights · New
• Handle BF16 dispatch when only older PTX was compiled · New
• Reject dimensions that overflow device index types · New
• Tile large prefill dequantization to limit scratch memory · New
• Handle misaligned FP16/BF16 activation loads · New
• Decode GGUF values explicitly as little-endian · New
• Add large-M numerical coverage for chunking and workspace reuse · New
• Preserve ggml copyright and MIT notices for copied tables · New
• Document negative-index behavior and format block requirements · New
• Gate BF16 tests on SM80 capability · New
• Add independent CPU known-answer decoder tests · New
• Cover FP16, prefill, and runtime-M GEMV paths · New
• Add MoE tests for lookup formats and mixed weight pairs · New

11 resolved since last review

• Stage device initializers into owned CPU tensors
• Regenerate and commit contrib operator documentation
• Keep all warp lanes active in partial-tile MMA
• Reject packed formats without a supported decoder
• Gate BF16 Tensor Core dispatch on hardware support
• Handle shared-memory requests above CUDA's default limit
• Reject top-k values greater than the expert count
• Reject empty vocabularies before gathering rows
• Compare MoEPacked outputs against a numerical reference
• Disable CPU fallback in the CUDA decoder test
• Decode external weights into native-endian float storage

_{Review effort: Max}

[Copilot]

🟡 Changes recommended

Unchecked CPU row-size arithmetic can overflow, and CUDA FP16 prefill can produce infinite outputs even when the correct result is finite.

11 open issues

• Check row-size multiplication for int64 overflow · New
• Avoid FP16 overflow when converting large prefill weights · New
• Avoid per-row heap allocations when gathering low-precision output · New
• Clarify GGUF loading claims or implement mapped-buffer access · New
• Probe CUDA function attributes on the configured device · New
• Validate complete quantization blocks during construction · New
• Fall back to SIMT when Tensor Core code is unavailable · New
• Validate output dtype before narrowing to 32 bits · New
• Test BF16 staging-buffer reuse with larger MoE dimensions · New
• Test large-N BF16 decode across quantized formats · New
• Add independent known-answer tests for quantized formats · New

16 resolved since last review

• Handle misaligned FP16/BF16 activation loads
• Tile large prefill dequantization to limit scratch memory
• Reject dimensions that overflow device index types
• Handle BF16 dispatch when only older PTX was compiled
• Make packed decoder reads safe for unaligned weights
• Validate K and grid limits before CUDA launches
• Use scalar finalize for misaligned output pointers
• Range-check MoE dimensions before narrowing casts
• Add MoE tests for lookup formats and mixed weight pairs
• Cover FP16, prefill, and runtime-M GEMV paths
• Add independent CPU known-answer decoder tests
• Gate BF16 tests on SM80 capability
• Document negative-index behavior and format block requirements
• Preserve ggml copyright and MIT notices for copied tables
• Add large-M numerical coverage for chunking and workspace reuse
• Decode GGUF values explicitly as little-endian

_{Review effort: Max}

[Baiju Meswani (baijumeswani)]

What this PR does

This adds GatherPacked on CPU/CUDA and MatMulPacked and MoEPacked on CUDA. They read native block-packed weights without requiring a separate repacking step or separate scale tensors. That is useful for loading existing GGUF models without keeping an expanded copy of their weights. CUDA still needs device memory for the packed weights; host-buffer borrowing and hybrid CPU/GPU offload are separate parts of the wider work.

What looks good

The current version addresses several earlier concerns: dimension and byte-size checks, unaligned buffers, compiled BF16 capability checks and fallback, bounded MatMul prefill scratch, and FP16 prefill overflow. The tests now have independent known answers for all 15 accepted formats, and CUDA sessions disable CPU fallback. These are meaningful improvements, not just comparisons between two new implementations that could share the same mistake.

What remains before approval

Please fix the CPU GatherPacked global-type-reduction failure described inline and add a regression. The public format contract and the decode/prefill precision differences should also be documented clearly before models start depending on these schemas. The inline comments include two further improvement suggestions: reducing empty MoE launches and broadening the independent IQ fixtures.

For the stack as a whole, please restack and test the actual final candidate. Against this head, I reproduced 16 conflicting paths when combining #33962, including validation, decoders, kernels, schemas and tests. The upper branch contains an older copy of the packed implementation, so conflict resolution needs to preserve the newer fixes here as well as the precision/routing additions. #33991 is explicitly test-only and uses an older bottom snapshot; it does not establish that the current merge stack works together.

Before settling the public operator surface, I would also like a clear comparison with #33940's existing-operator approach. Native packed storage is a reasonable requirement; I am not asking for repacking merely to reuse an operator name, or assuming that the alternative covers every GGUF format here. The question is which new public contracts are necessary, while keeping decoding, routing and expert execution reusable.

Reviewed f5f7aee with three independent reviewers. I reproduced the registration filtering failure, the numerical example and tile-bound calculations, and checked the stack conflicts. I did not run a full ORT build, GPU runtime tests, sanitizers or latency benchmarks.

[Tianlei Wu (tianleiwu)]

Addressed the five inline items in 9beea24119: pair-aware GatherPacked reduction, the MoE launch clamp, explicit rounding boundaries, a stable packed-format contract, and independent high-index IQ fixtures.

Validation of this lower candidate

• CUDA 13 Debug Python binding and CUDA provider built successfully; the packaged libraries were verified against the fresh linked outputs.
• All 33 packed-operator tests passed on the freshly built CPU/CUDA runtime: 811 subtests, zero skips, failures, or errors. All five real-parser reduction tests passed.
• Focused Ruff, Python formatting, clang-format, and whitespace checks passed. Generated documentation changes are limited to the three packed schemas.
• For the exact BF16 M=17, k=8, E=512, H=2560, I=7168 workload, an analytic constant-output assertion passed. nsys stats -r cuda_gpu_kern_gb_sum --format csv recorded one FC1 launch with grid (112,136,1) and one FC2 launch with grid (20,136,1): 17,952 combined CTAs. The previous bound calculates to 68,244 CTAs; this is launch-count evidence, not a measured latency improvement.

Comparison with #33940

Contract/capability	#33940's existing-operator approach	This lower PR
Projection storage	Fixed MatMulNBits profiles for offset-UINT8 Q8 and IQ2_XXS, with a separate fixed_scales contract; existing FP4 matmul for NVFP4	Native block blobs with interleaved format metadata and stable IDs for 15 formats
Projection execution	Fixed MatMulNBits profiles are bounded BF16 decode paths, rows 1-8	FP16/BF16 decode and prefill, with documented M/N-dependent weight rounding
Packed embeddings	No packed Gather implementation in that change set	CPU/CUDA GatherPacked
MoE profile	Restricted IQ2 QMoE with ReLU-square and preselected or fused sigmoid/bias-corrected routing	Three native packed expert blobs, SwiGLU, and softmax/top-k routing

These are not strict supersets: this lower PR does not implement NVFP4 or the Bolt-specific routing contract. Offset-UINT8 Q8 also must not be equated with native Q8_0's signed bytes and interleaved scale.

I agree that native storage does not mathematically require new operator names or repacking merely to reuse existing names. The necessary public contracts are the byte codec/layout and dimensions, activation/output types and rounding, and MoE activation/routing semantics. Those need explicit, stable definitions under either operator surface. #33940 demonstrates the alternative of adding guarded fixed profiles to existing schemas; the three names here expose a uniform native-block interface. Private decoding, routing, and expert execution should remain reusable rather than being separated merely by the public names.

Upper-stack qualification remains outstanding

Re-running read-only git merge-tree --write-tree --name-only --no-messages against this committed lower head and #33962 head ca5bf6a4c1308b4374238b37884d07edf4c3bae5 still reports 16 conflicting paths, including validation, decoder tables, kernels, schemas, and tests.

Rui Ren (@rui-ren): #33962 needs replaying/restacking onto 9beea24119aa03ebe05d03e61e3d5d9d57f8aa44, preserving these newer lower fixes alongside its precision/routing additions, followed by validation of that actual combined candidate. I have not rewritten another author's upper branches or claimed combined-stack validation. #33991 remains an older-snapshot, explicitly test-only integration and does not qualify this updated stack.

[Tianlei Wu (tianleiwu)]

CI maintenance update at 69a0a3ca25, preserving the native reviewer fixes in 9beea24119 and the targeted simulator preparation in ccea73c62b.

The preceding exact9beea iOS report contained zero testcase entries: cold simulator boot stalled, the unchanged240-second Detox setup deadline expired, and cleanup retained handles until the90-minute job deadline. The two-file ccea preparation repair passed all nine actual Packed YAML/Node-config regressions and only resets the uniquely selected simulator; it does not restart CoreSimulator globally or weaken the Detox deadline, reinstall, cleanup, or assertions.

The subsequent exactccea React Native42622402833 successfully prepared and launched the device and actually executed three cases. MNIST and Basic Types passed; EPContext reload failed after a real native signal11, with AsyncWorker::toPromise background-thread lambda+56 in the stack. Artifact13679289029 preserves that evidence. This is not another zero-case infrastructure failure and was not blindly retried. The stack identifies the worker thread, not a disassembled source instruction.

The original39a worker has an independently reproduced post-dispatch lifetime bug: resolve/reject dispatch can release its last owner before the background lambda subsequently accesses raw this through onFinished(). The new commit backports the verified f58 completion-before-dispatch repair, preserves exactly-once completion when dispatch throws and on cancellation, and adds settlement-completion assertions with negative component regressions. No Packed operator/provider/kernel or preceding reviewer-fix code was changed.

Fresh validation on the actual Packed69a/f58 header:640 real pinned-Hermes/JSI ASan/LSan lifetime scenarios passed in two independent processes, covering resolve, reject, cancellation, Env invalidation and dispatch exceptions; compiler dependencies confirm the actual checkout header. All five focused EPContext component tests passed, the configured native policy suite passed1/1, and configured C++ formatting, repository TSX ESLint and repository-locked Prettier checks passed. This is local native lifecycle evidence, not a claim that the complete device or optional Packed matrix ran locally.

Actual new-head React Native42634651129 and Windows CUDA42634651253 are IN_PROGRESS. Fresh exact69a Android/iOS three-case/full14 execution and complete current-head CI still need verification before full readiness. Normal human approval remains external.

[Rui Ren (rui-ren)]

LGTM, thanks.