onnxruntime
5b249f65 - [CUDA] Speed up 8-bit MatMulNBits dequantization with byte permutes (#31350)

Commit
3 days ago
[CUDA] Speed up 8-bit MatMulNBits dequantization with byte permutes (#31350) ### Description The uint8 → half/float dequantization in the 8-bit `MatMulNBits` GEMV kernels is **issue-bound, not bandwidth-bound**. Replacing the per-element integer-to-float converts with the "magic half" byte-permute decode makes the M=1 kernel **18% faster**, with **bit-identical** output. `ncu` on an H200 (SM90), 8-bit `MatMulNBits` with K=2048, N=248320 (an int8 `lm_head`): | metric | value | |---|---| | Compute (SM) Throughput | **86.46 %** | | DRAM Throughput | 35.02 % | | L1/TEX Cache Throughput | 66.00 % | | Achieved Occupancy | 94.51 % | | Executed Ipc Active | **3.49** inst/cycle (of 4) | Occupancy is already ~95% and DRAM is only a third utilized, so the kernel is limited by the ALU/convert instructions in the dequantization rather than by weight traffic. The fix is to issue fewer instructions per unpacked weight. ### Approach The half bit pattern `0x6400 | q` encodes exactly `1024 + q` for any `q ∈ [0, 255]`: the `2^10` exponent makes the low 10 mantissa bits integer-valued with `ulp == 1`. A single `__byte_perm` against the constant `0x64646464` therefore materializes **two halves at once**, so 8 `u16→half` converts plus 4 pack operations collapse into 4 permutes. Subtracting the matching biased zero point `0x6400 | zp` (i.e. `1024 + zp`) cancels the `1024` offset exactly. This is the same technique already vendored in this repository — see `start_byte_for_fp16 = 0x64646464` in [`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h`](onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h#L54) (CUTLASS's `FastInterleavedAndBiasedNumericArrayConverter`) — applied here to the non-interleaved MatMulNBits layout. **This is bit-identical to the previous code, not an approximation.** Both operands are exact halves and the result `q - zp` lies in `[-255, 255]`, which half represents exactly, so no rounding is introduced anywhere. ### Key Changes | Function | Used by | Change | |---|---|---| | `AccumulateEightElements8b` (half overload) | `MatMulFloat8bKernelM1` | 8 converts + 4 packs → 4 `__byte_perm` | | `DequantizeEight8b<T>` | `MatMulFloat8bKernelBatched` | 8 int→float converts + 8 subtracts → 4 permutes + 4 `__hsub2` | The pre-existing scalar loop is retained under `#else` for `__CUDA_ARCH__ < 530`, matching the guard already used by the half accumulate path. No API, kernel-launch, or dispatch changes. ### Results Measured on H200 (SM90), nsys median over ~60 launches, N=248320, K=2048: | kernel | before | after | | |---|---:|---:|---:| | `MatMulFloat8bKernelM1` | 248.6 µs | **203.0 µs** | **−18%** | | `MatMulFloat8bKernelBatched` | 301.6 µs | 299.6 µs | −1% | `ncu` after the change shows Compute SM 87.2% / DRAM 42.9% / L1-TEX 82.5% — the kernel has moved off the instruction-issue limit and toward the L1 limit. The batched kernel gains little because it is **not** issue-bound: L1/TEX throughput is 98.5% and occupancy is capped at 34.9% (3 blocks, 80 registers/thread). Its bottleneck is activation re-reads — the `[M,K]` activation tile is re-read by every one of the N/16 blocks — which needs shared-memory staging and a register reduction. That is left to a separate change. ### Testing Notes Existing coverage exercises both modified paths; because the transform is bit-identical, no test changes are required. ```bash ninja -j32 onnxruntime_provider_test ./onnxruntime_provider_test --gtest_filter='*MatMul8Bits*:*MatMulNBits*:*MatMul8bits*' ``` Result: **74 tests, 71 PASSED, 0 FAILED**, 3 skipped (`DynamicZeroPoints_AsymmetricCompInt8`, `SharedPrepackedWeights_DynamicZeroPoints_AsymmetricCompInt8`, `Float16_Comprehensive` — all skipped on `main` as well). Reviewers may want to sanity-check the `__byte_perm` selector nibbles: bytes 0-3 come from `x`, bytes 4-7 from `y`, and nibble 0 (LSB) selects result byte 0. So `__byte_perm(lo32, 0x64646464, 0x4140)` yields `{q0, 0x64, q1, 0x64}`, i.e. the two halves `1024+q0` and `1024+q1` (`half2` stores `.x` in the low 16 bits).
Author
Parents
Loading