triton-3bit-kv-cache-kernel — Cross-Byte & Cross-Word 3-Bit KV-Cache Packing, Triton Dequantization & Fused Flash-Decoding

Bit-exact 3-bit (uint3, values in ${0, \dots, 7}$) KV-cache packing/unpacking across byte (3 × uint8 = 24-bit) and GPU word (3 × uint32 = 96-bit) boundaries, fused with on-the-fly register-level affine dequantization inside an online-softmax Flash-Decoding Triton kernel.

---

1. Packing & Unpacking 3-Bit Values Across Byte & Word Boundaries

Because $\gcd(3, 8) = 1$, a tight 3-bit stream crosses a byte boundary every $\text{LCM}(3, 8) = 24$ bits (8 three-bit values = 3 bytes) and crosses a 32-bit GPU word boundary every $\text{LCM}(3, 32) = 96$ bits (32 three-bit values = 3 uint32 words).

A. 24-Bit Octet Layout (8 × 3-bit $\leftrightarrow$ 3 × uint8 bytes b0, b1, b2)

Within every 3-byte group $(b_0, b_1, b_2)$, value $v_2$ (bits 6..8) straddles $b_0 \to b_1$ and value $v_5$ (bits 15..17) straddles $b_1 \to b_2$:

Byte 0 (bits 0..7):   [ v2[1:0] (2b) | v1[2:0] (3b) | v0[2:0] (3b) ]
Byte 1 (bits 8..15):  [ v5[0]   (1b) | v4[2:0] (3b) | v3[2:0] (3b) | v2[2] (1b) ]
Byte 2 (bits 16..23): [ v7[2:0] (3b) | v6[2:0] (3b) | v5[2:1] (2b) ]

Packing (8 × uint3 $\to$ 3 × uint8):

b0 = (v0 & 0x7) | ((v1 & 0x7) << 3) | ((v2 & 0x3) << 6)
b1 = ((v2 >> 2) & 0x1) | ((v3 & 0x7) << 1) | ((v4 & 0x7) << 4) | ((v5 & 0x1) << 7)
b2 = ((v5 >> 1) & 0x3) | ((v6 & 0x7) << 2) | ((v7 & 0x7) << 5)

Unpacking (3 × uint8 $\to$ 8 × uint3):

• Explicit byte-straddle reconstruction:

  v0 = b0 & 0x7
  v1 = (b0 >> 3) & 0x7
  v2 = ((b0 >> 6) & 0x3) | ((b1 & 0x1) << 2)   # 2 LSBs from b0[7:6], 1 MSB from b1[0]
  v3 = (b1 >> 1) & 0x7
  v4 = (b1 >> 4) & 0x7
  v5 = ((b1 >> 7) & 0x1) | ((b2 & 0x3) << 1)   # 1 LSB from b1[7], 2 MSBs from b2[1:0]
  v6 = (b2 >> 2) & 0x7
  v7 = (b2 >> 5) & 0x7

• Branchless 24-bit register assembly inside Triton (triton_fused_3bit_flash_decode.py): By reassembling the 3 bytes into a single 32-bit register word24 = b0 | (b1 << 8) | (b2 << 16), the straddling bits of $v_2$ (bits 6..8) and $v_5$ (bits 15..17) become contiguous inside the register, so all 8 lanes unpack uniformly with zero divergent branches:

  word24 = b0 | (b1 << 8) | (b2 << 16)
  lanes  = tl.arange(0, 8)               # [0, 1, 2, 3, 4, 5, 6, 7]
  q      = (word24[:, None] >> (lanes[None, :] * 3)) & 0x7

B. 96-Bit GPU Coalesced Word Layout (32 × 3-bit $\leftrightarrow$ 3 × uint32 words W0, W1, W2)

For 32-bit aligned GPU loads (32 values = 96 bits = 3 × uint32), lane $i \in {0..31}$ starts at bit offset $3i$, word $w_i = (3i) // 32 \in {0, 1, 2}$, and shift $s_i = (3i) \bmod 32$. Only two lanes straddle 32-bit word boundaries:

• Lane $i = 10$ ($s_{10} = 30$, bits 30..32): 2 LSBs in W0[31:30], 1 MSB in W1[0]
• Lane $i = 21$ ($s_{21} = 31$, bits 63..65): 1 LSB in W1[31], 2 MSBs in W2[1:0]

Both in-word and cross-word lanes unpack branchlessly via:

lo = W[w_i] >> s_i
hi = (W[min(w_i + 1, 2)] << ((32 - s_i) & 31)) * (s_i > 29)
q  = (lo | hi) & 0x7

---

2. How We Prove Bit-Exact Round-Trip (test_roundtrip_exhaustive.py)

1. Exhaustive $2^{24} = 16,777,216$ State-Space Bijection Check (test_1_exhaustive_24bit_bijection & test_5A on CUDA Triton 3.2):
   • Because $\text{LCM}(3, 8) = 24$ bits has exactly $8^8 = 2^{24} = 16,777,216$ states, we test every possible 24-bit state on both CPU and GPU (triton_pack_3bit_bytes_kernel + triton_unpack_3bit_bytes_kernel), verifying both:
     • Left-inverse (Value Round-Trip): $\forall v \in {0..7}^8,; \text{unpack}(\text{pack}(v)) == v$
     • Right-inverse (Byte Bijection): $\forall b \in {0..255}^3,; \text{pack}(\text{unpack}(b)) == b$
2. Walking-Ones / Walking-Zeros Cross-Boundary Lane Isolation (test_2_cross_boundary_lane_isolation):
   • Sweeps all 32 lanes and all 8 values $v \in {0..7}$ against both all-zeros (0b000) and all-ones (0b111) neighbor backgrounds to verify zero bit-bleed on byte-straddling lanes ($v_2, v_5, \dots$) and word-straddling lanes ($v_{10}, v_{21}$), and verifies byte-for-byte little-endian equivalence between uint8[12] and uint32[3].
3. FP16 Affine Dequant $\leftrightarrow$ Re-Quant Round-Trip (test_3_affine_fp16_dequant_roundtrip):
   • Confirms that $\text{clamp}(\text{round}(((q - z) \cdot s) / s + z), 0, 7) == q$ for all $q \in {0..7}$.
4. Live CUDA Triton 3.2 Fused 3-Bit KV-Cache Flash-Decoding Verification (test_5B):
   • Verifies triton_fused_3bit_kv_flash_decode_kernel on GPU against an FP32 dequantized reference (max_abs_err = 0.000000).

Verification Output (NVIDIA GeForce RTX 3060 / PyTorch 2.6.0+cu124 / Triton 3.2.0)

$ python3 test_roundtrip_exhaustive.py
[PASS] Test 1: Exhaustive 2^24 (16,777,216) 3-byte <-> 8x3-bit bijection verified in 0.91s (0 bit errors)
[PASS] Test 2: Cross-byte (v2, v5, ...) & cross-uint32 (v10, v21) walking-pattern isolation verified (0 crosstalk)
[PASS] Test 3: FP16 affine dequant + re-quant round-trip verified on (64, 128) KV tiles (0 mismatches)
[PASS] Test 4: Fused 3-Bit KV-Cache Flash-Decoding reference verified on (B*H=8, S=256, D=64)
[PASS] Test 5A (CUDA Triton 3.2 on NVIDIA GeForce RTX 3060 Laptop GPU): Exhaustive 2^24 (16,777,216) cross-byte pack/unpack states verified on GPU in 1982.2 ms (0 bit errors)
[PASS] Test 5B (CUDA Triton 3.2 on NVIDIA GeForce RTX 3060 Laptop GPU): Fused 3-Bit Cross-Byte KV Flash-Decoding kernel matches FP32 reference (max_abs_err = 0.000000)

ALL 5 CPU + CUDA TRITON GPU VERIFICATION SUITES PASSED 100%.