Problem
CUDA NF4/FP4 quantize_4bit and dequantize_4bit currently pass the logical element count through C entry points and kernel launchers whose n is int. The existing large-tensor test therefore stops at 2**31 - 1 elements. This is reachable well below the memory capacity of B300-class GPUs: the boundary is about 4 GiB for FP16/BF16 input and 8 GiB for FP32 input.
This issue tracks a bounded investigation of ABI-preserving support for numel() > INT_MAX. It follows the limitation documented in upstream #1782, the <= INT_MAX overflow fix in #1784, and the high-priority 4-bit scope in #1785. Performance and correctness are measurement gates, not assumed outcomes.
Proposed scope
- Keep every existing int32 C symbol and the complete
n <= INT_MAX Python path unchanged.
- Add parallel large-count C entry points for CUDA NF4/FP4 quantization and dequantization. Their
n argument is int64_t; the Python ctypes declarations must use ctypes.c_int64 explicitly.
- Select the new entry points only when
numel() > INT_MAX.
- Inside each new host helper, use
chunk_elems = floor(INT_MAX / blocksize) * blocksize
and invoke the existing typed int32 launchers once per chunk on the same stream/launch ordering.
- For a chunk beginning at logical element
base, use raw-byte packed offsets:
- quantize:
A + base, out_bytes + base / 2, absmax + base / blocksize
- dequantize:
A_bytes + base / 2, absmax + base / blocksize, out + base
- Pass each launcher an explicit int32 local chunk count. All supported 4-bit blocksizes are even, so every nonfinal boundary owns complete quantization blocks and complete packed bytes; only the global final chunk may contain an odd nibble tail.
- Preserve existing output allocation,
QuantState, .out, caller-provided buffers, schemas, fake/meta behavior, and serialization.
compress_statistics=True remains in validation scope, but does not expand this issue into general >INT_MAX 8-bit blockwise support: for B300-feasible 4-bit inputs, the nested absmax tensor remains far below INT_MAX and should continue through the existing path.
Correctness gates
- Construct a reference by applying the unchanged upstream int32 ops to the same block-aligned chunks and concatenating packed bytes/absmax. Require bitwise equality for candidate quantized bytes and absmax.
- Dequantize the reference chunks with unchanged upstream calls and require bitwise equality to the candidate large dequantization output.
- Cover
INT_MAX, INT_MAX + 1, INT_MAX + blocksize - 1, INT_MAX + blocksize + 1, and more than two large chunks with an odd final tail where memory permits.
- Cover FP16, BF16, and FP32; NF4 and FP4; blocksizes 32/64/128 plus representative 256/4096;
quant_storage wider than uint8 via raw-byte views; nested statistics on/off; allocating and .out APIs; 1-D/multidimensional shapes; zero, finite, and extreme values.
- For all
n <= INT_MAX, prove that the old symbols are selected and outputs/state remain bitwise identical. Verify that old C symbols retain their original signatures and remain resolvable.
- Run relevant 4-bit op/functional/Linear4bit/Params4bit/serialization tests, opcheck/fake/compile coverage, the official CUDA multi-architecture build, and repository lint/pre-commit checks.
B300 validation and acceptance
- Reproduce the current direct failure at
INT_MAX + 1 in a child process so a native failure cannot terminate the validation driver.
- On one B300 allocation, compare isolated baseline and candidate builds from the same upstream commit using the official CUDA target list and assert the loaded library plus old/new symbols.
- Validate large cases for one, two, and more than two launches as memory permits, including a public
functional.quantize_4bit / Params4bit._quantize path.
- Compare performance against an explicit sequence of unchanged int32 chunk calls over the same input, not against the invalid direct baseline call. Report CUDA-event and host-wall medians, p10/p90, effective bandwidth, launch count, and peak allocation.
- Accept only if every large boundary completes with bitwise chunk-oracle parity, normal sizes regress by no more than 1%, the large path reaches at least 95% of the manual chunk reference throughput, and it adds no material allocation beyond the public outputs/state.
Non-goals and risks
If exact packed parity or bounded throughput cannot be achieved with this additive chunking design, record a no-go rather than broadening the work into an all-kernel int64 rewrite.
Problem
CUDA NF4/FP4
quantize_4bitanddequantize_4bitcurrently pass the logical element count through C entry points and kernel launchers whosenisint. The existing large-tensor test therefore stops at2**31 - 1elements. This is reachable well below the memory capacity of B300-class GPUs: the boundary is about 4 GiB for FP16/BF16 input and 8 GiB for FP32 input.This issue tracks a bounded investigation of ABI-preserving support for
numel() > INT_MAX. It follows the limitation documented in upstream #1782, the<= INT_MAXoverflow fix in #1784, and the high-priority 4-bit scope in #1785. Performance and correctness are measurement gates, not assumed outcomes.Proposed scope
n <= INT_MAXPython path unchanged.nargument isint64_t; the Python ctypes declarations must usectypes.c_int64explicitly.numel() > INT_MAX.chunk_elems = floor(INT_MAX / blocksize) * blocksizeand invoke the existing typed int32 launchers once per chunk on the same stream/launch ordering.
base, use raw-byte packed offsets:A + base,out_bytes + base / 2,absmax + base / blocksizeA_bytes + base / 2,absmax + base / blocksize,out + baseQuantState,.out, caller-provided buffers, schemas, fake/meta behavior, and serialization.compress_statistics=Trueremains in validation scope, but does not expand this issue into general >INT_MAX 8-bit blockwise support: for B300-feasible 4-bit inputs, the nested absmax tensor remains far below INT_MAX and should continue through the existing path.Correctness gates
INT_MAX,INT_MAX + 1,INT_MAX + blocksize - 1,INT_MAX + blocksize + 1, and more than two large chunks with an odd final tail where memory permits.quant_storagewider than uint8 via raw-byte views; nested statistics on/off; allocating and.outAPIs; 1-D/multidimensional shapes; zero, finite, and extreme values.n <= INT_MAX, prove that the old symbols are selected and outputs/state remain bitwise identical. Verify that old C symbols retain their original signatures and remain resolvable.B300 validation and acceptance
INT_MAX + 1in a child process so a native failure cannot terminate the validation driver.functional.quantize_4bit/Params4bit._quantizepath.Non-goals and risks
quant_storagehas itemsize greater than one.If exact packed parity or bounded throughput cannot be achieved with this additive chunking design, record a no-go rather than broadening the work into an all-kernel int64 rewrite.