[Tianlei Wu (tianleiwu)]

Description

Adds optional M (row) chunking to the CUDA MatMulNBits fpA_intB path to bound memory that grows with M, targeting large-vocabulary LM heads during chunked prefill. It is off by default and size-gated, so small layers are never split.

Setting	Meaning
ep.cuda.matmul_nbits_m_chunk_size (session config) / ORT_MATMULNBITS_M_CHUNK_SIZE (env)	Max rows of A per fpA_intB launch. 0/unset disables. Session config wins. Also works for the CUDA plugin EP.
ORT_MATMULNBITS_FORCE_CHUNKED=1 (existing)	Now also bypasses the M-chunking size gate (it already bypassed the fallback's size guards).

A node is chunked only when M > chunk and M * (N + K) * sizeof(T) > 256 MiB (the tactic profiler's A and C buffers for an unchunked M; same budget as the existing N-chunked fallback). For Qwen3.8-27B in fp16/bf16:

node	chunked when M >
LM head (N=248320, K=5120)	529
MLP gate/up/down	5957
any node with N + K <= 16384	>= 8192 (effectively never)

Changes

File	Change
matmul_nbits.h / .cc	Parse/validate the chunk size; per-node size gate (FpAIntBRowsPerLaunch); fpA_intB GEMM/GEMV loop over row chunks sharing one chunk-sized workspace, with at most two tactic lookups (full chunks + trailing chunk, which may use the GEMV below 16 rows); cap constructor profiling at max(chunk, 256 MiB / ((N + K) * sizeof(T))); Level-2 DeclareWorkspaceRequirements uses the chunk rows; empty input (M == 0) now returns before runtime g_idx validation, which may sync the stream.
onnxruntime_session_options_config_keys.h	New key kOrtSessionOptionsCudaMatMulNBitsMChunkSize.
docs/contrib_ops/cuda/matmul_nbits.md	New §6.2 (condition, measurements, memory breakdown); env-var table.
tests	C++ chunked-correctness test, Level-2/runtime size-gate test (MatMulNBitsWorkspace.MChunkSizeGate), Python parity / invalid-value / empty-input tests.

The dequantize + cuBLAS fallback and the fused GEMV are unchanged: neither has M-proportional scratch (the fallback's dequant buffer is already bounded by N chunking).

Motivation and Context

For an LM head, the fpA_intB tactic profiler allocates an M x N output buffer per profiled bucket, and a runtime M above the largest profiled bucket (2048 by default) triggers lazy profiling of that larger bucket at run time. Measured on H200 for the Qwen3.8-27B LM head alone (FP16, INT4, block 32, kSameAsRequested arena, NVML sampled from a separate process):

M	chunk	session-creation peak	run peak	session creation	first run	steady run
2048	off	3112 MiB	3080 MiB	17.5 s	584 ms	78.4 ms
2048	256	2234 MiB	3080 MiB	8.2 s	619 ms	79.0 ms
8192	off		10762 MiB	18.5 s	21716 ms	464-480 ms
8192	2048		5990 MiB	18.3 s	2239 ms	316-324 ms

• Above the profiled range, chunking saves 4.7 GiB of peak and makes the first run ~10x faster; the remainder is mostly the Y output.
• Within it, only the session-creation peak drops. The run peak is set by Y, both weight copies, and the CUDA context (breakdown in §6.2).

Limitations / follow-ups

• Default is off pending end-to-end benchmarks. A default that chunks to the largest profiled bucket when the size gate passes looks favorable from the data above.
• The Level-1 (partition-time) estimate ignores chunking and stays a conservative upper bound.
• Not exercised end to end with the plugin EP or CUDA graph capture. For graph capture, warmup must cover the trailing-chunk bucket, the same requirement as today.

Testing

Release build, CUDA 13.0, CMAKE_CUDA_ARCHITECTURES=90-real, onnxruntime_USE_FPA_INTB_GEMM=ON (compact), H200:

• onnxruntime_provider_test --gtest_filter='MatMulNBits*.*:MatMul2Bits*.*': 125 passed.
• onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.*' (with onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON): 291 passed, including MatMulNBitsWorkspace.MChunkSizeGate and the existing workspace agreement / zero-M tests.
• python -m pytest onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py: 7 passed, 9 skipped (offline weight packer not in the build).
• ORT_FPA_INTB_DEBUG=1 trace confirms a small node stays a single launch while the LM head at M=2048, chunk 256 runs 256-row chunks.

[Copilot]

Copilot review overview

🟡 Changes recommended

Profiler bucket rounding defeats the promised memory cap, and configuration documentation and overflow handling need correction.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1 · 1

Open (3)

• Prevent signed overflow when summing N and K · New
• Preserve the exact M cap during profiler allocation · New
• Document the 256 MiB gate for chunked launches · New

What changed in this PR

Adds optional, size-gated row chunking to reduce CUDA MatMulNBits fpA_intB workspace and profiling memory.

Changes:

• Adds chunk-size configuration and chunked GEMM/GEMV execution.
• Updates workspace estimation and empty-input handling.
• Adds documentation and correctness/configuration tests.

File	Description
matmul_nbits.h	Parses configuration and determines chunk sizes.
matmul_nbits.cc	Executes row chunks and adjusts workspace allocation.
onnxruntime_session_options_config_keys.h	Exposes the session configuration key.
matmul_nbits.md	Documents behavior and measurements.
matmul_4bits_test.cc	Tests chunked CUDA correctness.
matmul_nbits_e2e_workspace_test.cc	Tests workspace size gating.
test_op_matmulnbits_prepacked_cuda.py	Tests parity, validation, and empty inputs.

---

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

[Baiju Meswani (baijumeswani)]

Reviewed live head 7407337da3b66187edec0ee7f2b6d324eeb99b93. The existing Copilot threads already cover the actionable overflow, profiler-bucket cap, and public configuration documentation issues, so I am not duplicating those comments.

The row-offset execution, trailing GEMV handling, shared workspace, and empty-input path otherwise look sound. Before using the standalone H200 results to select a default, I recommend adding:

• an end-to-end long-prefill benchmark that reports TTFT and peak memory for the full model, not only the LM-head node;
• correctness/workspace coverage where the natural 256 MiB gate activates, since current chunking parity tests use ORT_MATMULNBITS_FORCE_CHUNKED and the workspace test only exercises a small shape;
• CUDA graph coverage for a shape with a trailing partial chunk, verifying warmup populates both tactic buckets;
• CUDA plugin EP coverage, since the new session configuration is documented as supported there but is not exercised end to end.

The current measurements support keeping this opt-in: the main TTFT benefit is avoiding first-use profiling of an unseen large M bucket, while an M already within the profiled range shows slight first-run overhead.