| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Optimized Triton TurboQuant decode attention (v2). |
| 4 | |
| 5 | FLUTE-paper optimizations applied: |
| 6 | 1. Grouped Q heads: grid over (B, Hk, splits) instead of (B, Hq, splits). |
| 7 | Each program loads BLOCK_M Q heads sharing a KV head into a 2D tile, |
| 8 | enabling tl.dot on tensor cores (MFMA/WMMA) for both Q·K and P·V. |
| 9 | 2. Vectorized pair LUT: precompute `pair_table[i][j] = (T[i], T[j])` |
| 10 | offline. At runtime, extract adjacent index pairs and fetch two |
| 11 | dequantized centroids with a single gather, halving LUT lookups. |
| 12 | 3. exp2 instead of exp: scores pre-scaled by log2(e) so the hardware- |
| 13 | native exp2 instruction replaces the more expensive exp. |
| 14 | 4. Wider index extraction: for 4-bit MSE, two adjacent 4-bit indices share |
| 15 | a byte. One byte load yields both, eliminating redundant loads. |
| 16 | 5. Centroids pre-warmed in L1 at kernel start. |
| 17 | 6. BLOCK_KV = TILE_SIZE raised to 16-32 (from 4), reducing loop iterations |
| 18 | and softmax rescaling overhead. |
| 19 | |
| 20 | Stage 2 is reused unchanged from triton_decode_attention.py. |
| 21 | """ |
| 22 | |
| 23 | import math |
| 24 | |
| 25 | import torch |
| 26 | |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.triton_utils import tl, triton |
| 29 | |
| 30 | # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory |
| 31 | # pressure (mirrors the pattern used in triton_decode_attention.py and |
| 32 | # triton_turboquant_decode.py). |
| 33 | _is_hip = current_platform.is_rocm() |
| 34 | |
| 35 | # On ROCm, bf16 has the same MFMA throughput as fp16 but wider dynamic range |
| 36 | # (8-bit exponent vs 5-bit), which is safer for attention scores. On CUDA, |
| 37 | # fp16 tensor cores may be faster than bf16 for some shapes. |
| 38 | _DOT_DTYPE = tl.bfloat16 if _is_hip else tl.float16 |
| 39 | |
| 40 | |
| 41 | # --------------------------------------------------------------------------- |
| 42 | # Pair LUT construction (FLUTE) — called once at launcher time |
| 43 | # --------------------------------------------------------------------------- |
| 44 | |
| 45 | |
| 46 | def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor: |
| 47 | """Build vectorized pair lookup table. |
| 48 | |
| 49 | For N centroids, returns a [N, N, 2] float32 tensor where |
| 50 | pair_lut[i, j] = (centroids[i], centroids[j]). |
| 51 | Flattened to [N*N, 2] for kernel indexing: pair_lut[i*N + j]. |
| 52 | |
| 53 | For 4-bit MSE (N=16): 16*16*2*4 = 2048 bytes — fits in L1/smem. |
| 54 | For 3-bit MSE (N=8): 8*8*2*4 = 512 bytes. |
| 55 | """ |
| 56 | N = centroids.shape[0] |
| 57 | # pair_lut[i,j,0] = centroids[i], pair_lut[i,j,1] = centroids[j] |
| 58 | c = centroids.float() |
| 59 | lut = torch.empty(N, N, 2, dtype=torch.float32, device=centroids.device) |
| 60 | lut[:, :, 0] = c[:, None] |
| 61 | lut[:, :, 1] = c[None, :] |
| 62 | return lut.reshape(N * N, 2).contiguous() |
| 63 | |
| 64 | |
| 65 | # --------------------------------------------------------------------------- |
| 66 | # Launcher v2 |
| 67 | # --------------------------------------------------------------------------- |
| 68 | |
| 69 | _layout_cache: dict = {} |
| 70 | |
| 71 | |
| 72 | def _get_layout(D, mse_bits, value_quant_bits, key_packed_size): |
| 73 | key = (D, mse_bits, value_quant_bits, key_packed_size) |
| 74 | cfg = _layout_cache.get(key) |
| 75 | if cfg is None: |
| 76 | val_data_bytes = math.ceil(D * value_quant_bits / 8) |
| 77 | cfg = { |
| 78 | "mse_bytes": math.ceil(D * mse_bits / 8), |
| 79 | "val_data_bytes": val_data_bytes, |
| 80 | "mse_bits": mse_bits, |
| 81 | "n_centroids": 2**mse_bits, |
| 82 | "BLOCK_D": triton.next_power_of_2(D), |
| 83 | } |
| 84 | _layout_cache[key] = cfg |
| 85 | return cfg |
| 86 | |
| 87 | |
| 88 | def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: |
| 89 | """Return a fresh pair-LUT for ``centroids`` on each call. |
| 90 | |
| 91 | The LUT is tiny (N*N*2 fp32, e.g. 2KB for 4-bit MSE) so the build cost |
| 92 | is negligible compared to attention. We avoid caching by data_ptr() |
| 93 | because CUDA allocator memory reuse across different centroid tensors |
| 94 | can silently return a stale LUT (subtle correctness bug). If this ever |
| 95 | shows up on a profile, cache by a hash-of-values fingerprint instead. |
| 96 | """ |
| 97 | return build_pair_lut(centroids) |