| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton fused TurboQuant decode attention. |
| 4 | |
| 5 | Decode path: Triton stage1 (split-KV tiled attention scoring + value |
| 6 | accumulation) + stage2 (log-sum-exp reduction across splits). |
| 7 | |
| 8 | Supports FP8 (E4M3) keys, 3-bit and 4-bit uniform quantized values. |
| 9 | """ |
| 10 | |
| 11 | import math |
| 12 | from typing import Any |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.triton_utils import tl, triton |
| 18 | from vllm.v1.attention.ops.triton_decode_attention import ( |
| 19 | _fwd_kernel_stage2, |
| 20 | ) |
| 21 | |
| 22 | _FP8_E4B15: dict[int, int] = {} |
| 23 | |
| 24 | |
| 25 | def _use_fp8_e4b15(device: int = 0) -> int: |
| 26 | """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0. |
| 27 | On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format). |
| 28 | """ |
| 29 | if device not in _FP8_E4B15: |
| 30 | if current_platform.is_cuda_alike(): |
| 31 | cap = torch.cuda.get_device_capability(device) |
| 32 | _FP8_E4B15[device] = 1 if cap < (8, 9) else 0 |
| 33 | else: |
| 34 | _FP8_E4B15[device] = 0 |
| 35 | return _FP8_E4B15[device] |
| 36 | |
| 37 | |
| 38 | # --------------------------------------------------------------------------- |
| 39 | # Stage 1: Fused TQ score + value accumulation (BLOCK_KV tiled) |
| 40 | # --------------------------------------------------------------------------- |
| 41 | |
| 42 | |
| 43 | @triton.jit |
| 44 | def _tq_decode_stage1( |
| 45 | # Precomputed query projection |
| 46 | Q_rot_ptr, # [B, Hq, D] float32 |
| 47 | # Compressed KV cache (combined K+V) |
| 48 | KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8 |
| 49 | # Block table and sequence info |
| 50 | Block_table_ptr, # [B, max_num_blocks] int32 |
| 51 | Seq_lens_ptr, # [B] int32 |
| 52 | # TQ parameters |
| 53 | Centroids_ptr, # [n_centroids] float32 |
| 54 | # Output (intermediate for stage2) |
| 55 | Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32 |
| 56 | # Strides |
| 57 | stride_qb, |
| 58 | stride_qh, # Q strides: [B, Hq, D] |
| 59 | stride_cache_block, |
| 60 | stride_cache_pos, |
| 61 | stride_cache_head, # KV cache |
| 62 | stride_bt_b, # block_table stride per batch |
| 63 | stride_mid_b, |
| 64 | stride_mid_h, |
| 65 | stride_mid_s, # mid_o strides |
| 66 | # Constexpr dims |
| 67 | NUM_KV_HEADS: tl.constexpr, |
| 68 | HEAD_DIM: tl.constexpr, |
| 69 | BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages) |
| 70 | NUM_KV_SPLITS: tl.constexpr, |
| 71 | KV_GROUP_SIZE: tl.constexpr, # Hq // Hk |
| 72 | # TQ layout constants |
| 73 | MSE_BITS: tl.constexpr, # 3 or 4 |
| 74 | MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8) |
| 75 | KPS: tl.constexpr, # key_packed_size |
| 76 | VQB: tl.constexpr, # value_quant_bits (4 or 8=FP8) |
| 77 | VAL_DATA_BYTES: tl.constexpr, # ceil(D * vqb / 8) or D for FP8 |
| 78 | # Score constants |
| 79 | ATTN_SCALE: tl.constexpr, # 1/sqrt(D) |
| 80 | # Block tile sizes |
| 81 | BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM) |
| 82 | BLOCK_KV: tl.constexpr, # tokens per tile (16) |
| 83 | KEY_FP8: tl.constexpr, # 1 if K is stored as FP8 |
| 84 | NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids |
| 85 | FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) |
| 86 | ): |
| 87 | bid = tl.program_id(0) # batch index |
| 88 | hid = tl.program_id(1) # q_head index |
| 89 | sid = tl.program_id(2) # kv_split index |
| 90 | |
| 91 | kv_head = hid // KV_GROUP_SIZE |
| 92 | |
| 93 | # Sequence length for this batch |
| 94 | seq_len = tl.load(Seq_lens_ptr + bid) |
| 95 | |
| 96 | # KV split range |
| 97 | split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) |
| 98 | split_start = split_len * sid |
| 99 | split_end = tl.minimum(split_start + split_len, seq_len) |
| 100 | |
| 101 | if split_start >= split_end: |
| 102 | return |
| 103 | |
| 104 | # Dimension offsets |
| 105 | d_offs = tl.arange(0, BLOCK_D) |
| 106 | d_mask = d_offs < HEAD_DIM |
| 107 | kv_range = tl.arange(0, BLOCK_KV) |
| 108 | |
| 109 | # Load query vector: q_rot — [BLOCK_D] float32 |
| 110 | q_base = bid * stride_qb + hid * stride_qh |
| 111 | q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32) |
| 112 | |
| 113 | # Precompute byte/bit index vectors for MSE gather loads |
| 114 | if not KEY_FP8: |
| 115 | mse_bit_off = d_offs * MSE_BITS |
| 116 | mse_byte_idx = mse_bit_off // 8 |
| 117 | mse_bit_shift = mse_bit_off % 8 |
| 118 | mse_mask = (1 << MSE_BITS) - 1 |
| 119 | |
| 120 | # Precompute value bit/byte index vectors (loop-invariant) |
| 121 | if VQB == 3: |
| 122 | val_bit_off = d_offs * 3 |
| 123 | val_byte_idx = val_bit_off // 8 |
| 124 | val_bit_shift = val_bit_off % 8 |
| 125 | |
| 126 | # Online softmax accumulators |
| 127 | m_prev = -float("inf") |
| 128 | l_prev = 0.0 |
| 129 | acc = tl.zeros([BLOCK_D], dtype=tl.float32) |
| 130 | |
| 131 | bt_base = bid * stride_bt_b |
| 132 | |
| 133 | # ================================================================ |
| 134 | # TILED LOOP: process BLOCK_KV tokens per iteration |
| 135 | # ================================================================ |
| 136 | for start_n in range(split_start, split_end, BLOCK_KV): |
| 137 | kv_offs = start_n + kv_range |
| 138 | kv_mask = kv_offs < split_end |
| 139 | |
| 140 | page_idx = kv_offs // BLOCK_SIZE |
| 141 | page_off = kv_offs % BLOCK_SIZE |
| 142 | block_nums = tl.load( |
| 143 | Block_table_ptr + bt_base + page_idx, |
| 144 | mask=kv_mask, |
| 145 | other=0, |
| 146 | ).to(tl.int64) |
| 147 | |
| 148 | slot_bases = ( |
| 149 | block_nums * stride_cache_block |
| 150 | + page_off.to(tl.int64) * stride_cache_pos |
| 151 | + tl.cast(kv_head, tl.int64) * stride_cache_head |
| 152 | ) |
| 153 | |
| 154 | # ============================================================ |
| 155 | # COMPUTE ATTENTION SCORES: [BLOCK_KV] |
| 156 | # ============================================================ |
| 157 | if KEY_FP8: |
| 158 | k_addrs = slot_bases[:, None] + d_offs[None, :] |
| 159 | k_raw = tl.load( |
| 160 | KV_cache_ptr + k_addrs, |
| 161 | mask=kv_mask[:, None] & d_mask[None, :], |
| 162 | other=0, |
| 163 | ) |
| 164 | if FP8_E4B15: |
| 165 | k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) |
| 166 | else: |
| 167 | k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) |
| 168 | scores = ( |
| 169 | tl.sum( |
| 170 | tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0), |
| 171 | axis=1, |
| 172 | ) |
| 173 | * ATTN_SCALE |
| 174 | ) |
| 175 | scores = tl.where(kv_mask, scores, -float("inf")) |
| 176 | else: |
| 177 | # MSE unpack + norms |
| 178 | mse_addrs0 = slot_bases[:, None] + mse_byte_idx[None, :] |
| 179 | mse_raw0 = tl.load( |
| 180 | KV_cache_ptr + mse_addrs0, |
| 181 | mask=kv_mask[:, None] & d_mask[None, :], |
| 182 | other=0, |
| 183 | ).to(tl.int32) |
| 184 | mse_raw1 = tl.load( |
| 185 | KV_cache_ptr + mse_addrs0 + 1, |
| 186 | mask=kv_mask[:, None] & d_mask[None, :], |
| 187 | other=0, |
| 188 | ).to(tl.int32) |
| 189 | raw16 = mse_raw0 | (mse_raw1 << 8) |
| 190 | mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask |
| 191 | |
| 192 | # Centroid gather + dot product |
| 193 | c_vals = tl.load( |
| 194 | Centroids_ptr + mse_idx, |
| 195 | mask=kv_mask[:, None] & d_mask[None, :], |
| 196 | other=0.0, |
| 197 | ) |
| 198 | |
| 199 | # Norm correction: re-normalize centroid vector to unit norm |
| 200 | if NORM_CORRECTION: |
| 201 | c_norm_sq = tl.sum( |
| 202 | tl.where(d_mask[None, :], c_vals * c_vals, 0.0), |
| 203 | axis=1, |
| 204 | ) |
| 205 | c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16) |
| 206 | c_vals = c_vals * c_inv_norm[:, None] |
| 207 | |
| 208 | term1 = tl.sum( |
| 209 | tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0), |
| 210 | axis=1, |
| 211 | ) |
| 212 | |
| 213 | # Load norms (fp16 -> fp32): norms are at MSE_BYTES offset |
| 214 | norm_bases = slot_bases + MSE_BYTES |
| 215 | n_lo = tl.load(KV_cache_ptr + norm_bases, mask=kv_mask, other=0).to( |
| 216 | tl.uint16 |
| 217 | ) |
| 218 | n_hi = tl.load(KV_cache_ptr + norm_bases + 1, mask=kv_mask, other=0).to( |
| 219 | tl.uint16 |
| 220 | ) |
| 221 | vec_norms = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 222 | |
| 223 | scores = vec_norms * term1 * ATTN_SCALE |
| 224 | scores = tl.where(kv_mask, scores, -float("inf")) |
| 225 | |
| 226 | # ============================================================ |
| 227 | # ONLINE SOFTMAX UPDATE (block-level) |
| 228 | # ============================================================ |
| 229 | n_e_max = tl.maximum(tl.max(scores, 0), m_prev) |
| 230 | re_scale = tl.exp(m_prev - n_e_max) |
| 231 | p = tl.exp(scores - n_e_max) |
| 232 | |
| 233 | # ============================================================ |
| 234 | # VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D] |
| 235 | # ============================================================ |
| 236 | val_bases = slot_bases + KPS |
| 237 | |
| 238 | if VQB == 3: |
| 239 | val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] |
| 240 | val_raw0 = tl.load( |
| 241 | KV_cache_ptr + val_addrs0, |
| 242 | mask=kv_mask[:, None] & d_mask[None, :], |
| 243 | other=0, |
| 244 | ).to(tl.int32) |
| 245 | val_raw1 = tl.load( |
| 246 | KV_cache_ptr + val_addrs0 + 1, |
| 247 | mask=kv_mask[:, None] & d_mask[None, :], |
| 248 | other=0, |
| 249 | ).to(tl.int32) |
| 250 | raw16 = val_raw0 | (val_raw1 << 8) |
| 251 | v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) |
| 252 | |
| 253 | sc_bases = val_bases + VAL_DATA_BYTES |
| 254 | sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to( |
| 255 | tl.uint16 |
| 256 | ) |
| 257 | sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to( |
| 258 | tl.uint16 |
| 259 | ) |
| 260 | v_scales = ( |
| 261 | (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 262 | ) |
| 263 | zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to( |
| 264 | tl.uint16 |
| 265 | ) |
| 266 | zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to( |
| 267 | tl.uint16 |
| 268 | ) |
| 269 | v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 270 | values = v_idx * v_scales[:, None] + v_zeros[:, None] |
| 271 | else: # VQB == 4 |
| 272 | vb_idx = d_offs // 2 |
| 273 | vb_shift = (d_offs % 2) * 4 |
| 274 | val_addrs = val_bases[:, None] + vb_idx[None, :] |
| 275 | val_raw = tl.load( |
| 276 | KV_cache_ptr + val_addrs, |
| 277 | mask=kv_mask[:, None] & d_mask[None, :], |
| 278 | other=0, |
| 279 | ).to(tl.int32) |
| 280 | v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32) |
| 281 | |
| 282 | sc_bases = val_bases + VAL_DATA_BYTES |
| 283 | sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to( |
| 284 | tl.uint16 |
| 285 | ) |
| 286 | sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to( |
| 287 | tl.uint16 |
| 288 | ) |
| 289 | v_scales = ( |
| 290 | (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 291 | ) |
| 292 | zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to( |
| 293 | tl.uint16 |
| 294 | ) |
| 295 | zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to( |
| 296 | tl.uint16 |
| 297 | ) |
| 298 | v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 299 | values = v_idx * v_scales[:, None] + v_zeros[:, None] |
| 300 | |
| 301 | # ============================================================ |
| 302 | # WEIGHTED VALUE ACCUMULATION |
| 303 | # ============================================================ |
| 304 | acc = acc * re_scale + tl.sum(p[:, None] * values, 0) |
| 305 | l_prev = l_prev * re_scale + tl.sum(p, 0) |
| 306 | m_prev = n_e_max |
| 307 | |
| 308 | # Store partial result |
| 309 | out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s |
| 310 | safe_l = tl.where(l_prev > 0.0, l_prev, 1.0) |
| 311 | tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask) |
| 312 | lse = m_prev + tl.log(safe_l) |
| 313 | tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse) |
| 314 | |
| 315 | |
| 316 | # --------------------------------------------------------------------------- |
| 317 | # Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16 |
| 318 | # --------------------------------------------------------------------------- |
| 319 | |
| 320 | |
| 321 | @triton.jit |
| 322 | def _tq_full_dequant_kv( |
| 323 | KV_cache_ptr, |
| 324 | Block_table_ptr, |
| 325 | Centroids_ptr, |
| 326 | K_out_ptr, # [B, Hk, max_seq, D] float16 |
| 327 | V_out_ptr, # [B, Hk, max_seq, D] float16 |
| 328 | stride_ko_b, |
| 329 | stride_ko_h, |
| 330 | stride_ko_s, |
| 331 | stride_vo_b, |
| 332 | stride_vo_h, |
| 333 | stride_vo_s, |
| 334 | stride_cache_block, |
| 335 | stride_cache_pos, |
| 336 | stride_cache_head, |
| 337 | stride_bt_b, |
| 338 | HEAD_DIM: tl.constexpr, |
| 339 | BLOCK_SIZE: tl.constexpr, |
| 340 | NUM_KV_HEADS: tl.constexpr, |
| 341 | MSE_BYTES: tl.constexpr, |
| 342 | KPS: tl.constexpr, |
| 343 | VQB: tl.constexpr, |
| 344 | VAL_DATA_BYTES: tl.constexpr, |
| 345 | MSE_BITS: tl.constexpr, |
| 346 | KEY_FP8: tl.constexpr, |
| 347 | BLOCK_D: tl.constexpr, |
| 348 | NORM_CORRECTION: tl.constexpr = 0, |
| 349 | FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) |
| 350 | ): |
| 351 | """Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16.""" |
| 352 | pos = tl.program_id(0) |
| 353 | bh = tl.program_id(1) |
| 354 | bid = bh // NUM_KV_HEADS |
| 355 | hid = bh % NUM_KV_HEADS |
| 356 | |
| 357 | page_idx = pos // BLOCK_SIZE |
| 358 | page_off = pos % BLOCK_SIZE |
| 359 | block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64) |
| 360 | slot_base = ( |
| 361 | block_num * stride_cache_block |
| 362 | + tl.cast(page_off, tl.int64) * stride_cache_pos |
| 363 | + tl.cast(hid, tl.int64) * stride_cache_head |
| 364 | ) |
| 365 | |
| 366 | d_offs = tl.arange(0, BLOCK_D) |
| 367 | d_mask = d_offs < HEAD_DIM |
| 368 | |
| 369 | # === K dequant === |
| 370 | ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s |
| 371 | if KEY_FP8: |
| 372 | k_raw = tl.load(KV_cache_ptr + slot_base + d_offs, mask=d_mask, other=0) |
| 373 | if FP8_E4B15: |
| 374 | k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) |
| 375 | else: |
| 376 | k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) |
| 377 | tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) |
| 378 | else: |
| 379 | # MSE unpack (3-bit or 4-bit) + norms |
| 380 | mse_bit_off = d_offs * MSE_BITS |
| 381 | mse_byte_idx = mse_bit_off // 8 |
| 382 | mse_bit_shift = mse_bit_off % 8 |
| 383 | mse_umask = (1 << MSE_BITS) - 1 |
| 384 | |
| 385 | mse_raw0 = tl.load( |
| 386 | KV_cache_ptr + slot_base + mse_byte_idx, mask=d_mask, other=0 |
| 387 | ).to(tl.int32) |
| 388 | mse_raw1 = tl.load( |
| 389 | KV_cache_ptr + slot_base + mse_byte_idx + 1, mask=d_mask, other=0 |
| 390 | ).to(tl.int32) |
| 391 | raw16_key = mse_raw0 | (mse_raw1 << 8) |
| 392 | mse_idx = (raw16_key >> mse_bit_shift) & mse_umask |
| 393 | |
| 394 | k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0) |
| 395 | |
| 396 | # Norm correction: re-normalize centroid vector to unit norm |
| 397 | if NORM_CORRECTION: |
| 398 | c_norm_sq = tl.sum(tl.where(d_mask, k_mse * k_mse, 0.0), axis=0) |
| 399 | c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16) |
| 400 | k_mse = k_mse * c_inv_norm |
| 401 | |
| 402 | # Norms at MSE_BYTES offset (no QJL bytes) |
| 403 | norm_base = slot_base + MSE_BYTES |
| 404 | n_lo = tl.load(KV_cache_ptr + norm_base).to(tl.uint16) |
| 405 | n_hi = tl.load(KV_cache_ptr + norm_base + 1).to(tl.uint16) |
| 406 | vec_norm = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 407 | |
| 408 | k_recon = vec_norm * k_mse |
| 409 | tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) |
| 410 | |
| 411 | # === V dequant === |
| 412 | val_base = slot_base + KPS |
| 413 | if VQB == 4: |
| 414 | vb_idx = d_offs // 2 |
| 415 | vb_shift = (d_offs % 2) * 4 |
| 416 | val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to( |
| 417 | tl.int32 |
| 418 | ) |
| 419 | v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32) |
| 420 | |
| 421 | sc_base = val_base + VAL_DATA_BYTES |
| 422 | sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16) |
| 423 | sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16) |
| 424 | v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 425 | zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16) |
| 426 | zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16) |
| 427 | v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 428 | v_vals = v_idx * v_scale + v_zero |
| 429 | elif VQB == 3: |
| 430 | # 3-bit value unpack: 8 values per 3 bytes |
| 431 | val_bit_off = d_offs * 3 |
| 432 | val_byte_idx = val_bit_off // 8 |
| 433 | val_bit_shift = val_bit_off % 8 |
| 434 | val_raw0 = tl.load( |
| 435 | KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0 |
| 436 | ).to(tl.int32) |
| 437 | val_raw1 = tl.load( |
| 438 | KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0 |
| 439 | ).to(tl.int32) |
| 440 | raw16_val = val_raw0 | (val_raw1 << 8) |
| 441 | v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32) |
| 442 | |
| 443 | sc_base = val_base + VAL_DATA_BYTES |
| 444 | sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16) |
| 445 | sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16) |
| 446 | v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 447 | zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16) |
| 448 | zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16) |
| 449 | v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) |
| 450 | v_vals = v_idx * v_scale + v_zero |
| 451 | else: |
| 452 | v_vals = tl.zeros([BLOCK_D], dtype=tl.float32) |
| 453 | |
| 454 | vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s |
| 455 | tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask) |
| 456 | |
| 457 | |
| 458 | # --------------------------------------------------------------------------- |
| 459 | # Stage 2: Reuse from triton_decode_attention.py |
| 460 | # --------------------------------------------------------------------------- |
| 461 | |
| 462 | # --------------------------------------------------------------------------- |
| 463 | # Launcher — cached constants + fused GEMM |
| 464 | # --------------------------------------------------------------------------- |
| 465 | |
| 466 | _layout_cache: dict = {} |
| 467 | |
| 468 | |
| 469 | def _get_layout(D, mse_bits, value_quant_bits, key_packed_size): |
| 470 | """Get cached layout constants.""" |
| 471 | key = (D, mse_bits, value_quant_bits, key_packed_size) |
| 472 | cfg = _layout_cache.get(key) |
| 473 | if cfg is None: |
| 474 | val_data_bytes = math.ceil(D * value_quant_bits / 8) |
| 475 | cfg = { |
| 476 | "mse_bytes": math.ceil(D * mse_bits / 8), |
| 477 | "val_data_bytes": val_data_bytes, |
| 478 | "mse_bits": mse_bits, |
| 479 | "n_centroids": 2**mse_bits, |
| 480 | "BLOCK_D": triton.next_power_of_2(D), |
| 481 | } |
| 482 | _layout_cache[key] = cfg |
| 483 | return cfg |
| 484 | |
| 485 | |
| 486 | def triton_turboquant_decode_attention( |
| 487 | query: torch.Tensor, # [B, Hq, D] — original query |
| 488 | kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 |
| 489 | block_table: torch.Tensor, # [B, max_num_blocks] int32 |
| 490 | seq_lens: torch.Tensor, # [B] int32 |
| 491 | Pi: torch.Tensor, # [D, D] float32 |
| 492 | centroids: torch.Tensor, # [n_centroids] float32 |
| 493 | scale: float, |
| 494 | mse_bits: int, |
| 495 | key_packed_size: int, |
| 496 | value_quant_bits: int, |
| 497 | key_fp8: bool = False, |
| 498 | norm_correction: bool = False, |
| 499 | PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous |
| 500 | # Pre-allocated buffers (optional, avoids per-call allocation) |
| 501 | mid_o_buf: torch.Tensor | None = None, |
| 502 | output_buf: torch.Tensor | None = None, |
| 503 | lse_buf: torch.Tensor | None = None, |
| 504 | buf_holder: Any = None, |
| 505 | max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph) |
| 506 | ) -> torch.Tensor: |
| 507 | """Launch fused TQ decode attention (Triton stage1 + stage2). |
| 508 | |
| 509 | Returns: output tensor [B, Hq, D] in query's dtype. |
| 510 | """ |
| 511 | B, Hq, D = query.shape |
| 512 | Hk = kv_cache.shape[2] |
| 513 | block_size = kv_cache.shape[1] |
| 514 | kv_group_size = Hq // Hk |
| 515 | device = query.device |
| 516 | |
| 517 | cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size) |
| 518 | |
| 519 | # Compute q_rot = q @ Pi.T (rotated query for MSE key scoring) |
| 520 | # FP8 path: pass query directly (float16); kernel casts inline. |
| 521 | # MSE path: still needs external GEMM (cuBLAS), so q_rot is float32. |
| 522 | if key_fp8: |
| 523 | q_rot = query.contiguous() |
| 524 | else: |
| 525 | q_float = query.float() |
| 526 | if PiT is None: |
| 527 | PiT = Pi.T.contiguous() |
| 528 | q_rot = (q_float @ PiT).contiguous() |
| 529 | |
| 530 | NUM_KV_SPLITS = max_num_kv_splits |
| 531 | |
| 532 | if ( |
| 533 | mid_o_buf is not None |
| 534 | and mid_o_buf.shape[0] >= B |
| 535 | and mid_o_buf.shape[2] >= NUM_KV_SPLITS |
| 536 | ): |
| 537 | mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] |
| 538 | else: |
| 539 | mid_o = torch.empty( |
| 540 | B, |
| 541 | Hq, |
| 542 | NUM_KV_SPLITS, |
| 543 | D + 1, |
| 544 | dtype=torch.float32, |
| 545 | device=device, |
| 546 | ) |
| 547 | if buf_holder is not None: |
| 548 | buf_holder._tq_mid_o_buf = mid_o |
| 549 | |
| 550 | # Stage 1: split-KV tiled attention scoring + value accumulation |
| 551 | fp8_e4b15 = _use_fp8_e4b15(device.index or 0) |
| 552 | BLOCK_KV = 4 |
| 553 | grid = (B, Hq, NUM_KV_SPLITS) |
| 554 | _tq_decode_stage1[grid]( |
| 555 | q_rot, |
| 556 | kv_cache, |
| 557 | block_table, |
| 558 | seq_lens, |
| 559 | centroids, |
| 560 | mid_o, |
| 561 | q_rot.stride(0), |
| 562 | q_rot.stride(1), |
| 563 | kv_cache.stride(0), |
| 564 | kv_cache.stride(1), |
| 565 | kv_cache.stride(2), |
| 566 | block_table.stride(0), |
| 567 | mid_o.stride(0), |
| 568 | mid_o.stride(1), |
| 569 | mid_o.stride(2), |
| 570 | NUM_KV_HEADS=Hk, |
| 571 | HEAD_DIM=D, |
| 572 | BLOCK_SIZE=block_size, |
| 573 | NUM_KV_SPLITS=NUM_KV_SPLITS, |
| 574 | KV_GROUP_SIZE=kv_group_size, |
| 575 | MSE_BITS=mse_bits, |
| 576 | MSE_BYTES=cfg["mse_bytes"], |
| 577 | KPS=key_packed_size, |
| 578 | VQB=value_quant_bits, |
| 579 | VAL_DATA_BYTES=cfg["val_data_bytes"], |
| 580 | ATTN_SCALE=scale, |
| 581 | BLOCK_D=cfg["BLOCK_D"], |
| 582 | BLOCK_KV=BLOCK_KV, |
| 583 | KEY_FP8=1 if key_fp8 else 0, |
| 584 | NORM_CORRECTION=1 if norm_correction else 0, |
| 585 | FP8_E4B15=fp8_e4b15, |
| 586 | num_warps=1, |
| 587 | num_stages=1, |
| 588 | ) |
| 589 | |
| 590 | # Stage 2: Reduce across KV splits |
| 591 | # Output in query dtype — eliminates float16_copy kernel after stage2 |
| 592 | out_dtype = query.dtype |
| 593 | if ( |
| 594 | output_buf is not None |
| 595 | and output_buf.shape[0] >= B |
| 596 | and output_buf.dtype == out_dtype |
| 597 | ): |
| 598 | output = output_buf[:B, :Hq, :D] |
| 599 | else: |
| 600 | output = torch.empty(B, Hq, D, dtype=out_dtype, device=device) |
| 601 | if buf_holder is not None: |
| 602 | buf_holder._tq_output_buf = output |
| 603 | if lse_buf is not None and lse_buf.shape[0] >= B: |
| 604 | lse = lse_buf[:B, :Hq] |
| 605 | else: |
| 606 | lse = torch.empty(B, Hq, dtype=torch.float32, device=device) |
| 607 | if buf_holder is not None: |
| 608 | buf_holder._tq_lse_buf = lse |
| 609 | |
| 610 | grid2 = (B, Hq) |
| 611 | _fwd_kernel_stage2[grid2]( |
| 612 | mid_o, |
| 613 | output, |
| 614 | lse, |
| 615 | seq_lens, |
| 616 | mid_o.stride(0), |
| 617 | mid_o.stride(1), |
| 618 | mid_o.stride(2), |
| 619 | output.stride(0), |
| 620 | output.stride(1), |
| 621 | lse.stride(0), |
| 622 | NUM_KV_SPLITS=NUM_KV_SPLITS, |
| 623 | BLOCK_DV=cfg["BLOCK_D"], |
| 624 | Lv=D, |
| 625 | OUTPUT_FP16=1 if out_dtype == torch.float16 else 0, |
| 626 | num_warps=4, |
| 627 | num_stages=2, |
| 628 | ) |
| 629 | |
| 630 | return output # already in query dtype |