Support >INT_MAX CUDA 4-bit quantization with block-aligned chunked launches #17

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

• No global int64 conversion of common kernels or existing ABI.
• No 8-bit blockwise, GEMV/GEMM, LLM.int8, optimizer, HIP/XPU, serialization-format, or public-schema expansion.
• Packed storage pointer arithmetic must be in bytes even when quant_storage has itemsize greater than one.
• Preserve existing stream behavior and error checking for every chunk; do not add host synchronization.
• Recheck upstream Support quantizing tensors when numel() > INT_MAX bitsandbytes-foundation/bitsandbytes#1785 and open work immediately before implementation. Abort rather than duplicate a newly active implementation, and rebase cleanly if nearby Add k-bit blockwise quantization (K=2-5) with warp-level CUDA kernels bitsandbytes-foundation/bitsandbytes#1858 or fix: Stop quantize_blockwise and quantize_4bit from mutating user-provided absmax bitsandbytes-foundation/bitsandbytes#1863 lands.

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.