| 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 | KV_cache_u16_ptr, # uint16 view of same storage — Opt#3 SoA loads |
| 50 | # Block table and sequence info |
| 51 | Block_table_ptr, # [B, max_num_blocks] int32 |
| 52 | Seq_lens_ptr, # [B] int32 |
| 53 | # TQ parameters |
| 54 | Centroids_ptr, # [n_centroids] float32 |
| 55 | # Output (intermediate for stage2) |
| 56 | Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32 |
| 57 | # Sink support |
| 58 | Sink_ptr, # [Hq] float32, pre-computed sink attention logits per head |
| 59 | # Strides |
| 60 | stride_qb, |
| 61 | stride_qh, # Q strides: [B, Hq, D] |
| 62 | stride_cache_block: tl.int64, # bytes per block (bs*H*slot_aligned) |
| 63 | stride_bt_b: tl.int64, # block_table stride per batch |
| 64 | stride_mid_b, |
| 65 | stride_mid_h, |
| 66 | stride_mid_s, # mid_o strides |
| 67 | # Constexpr dims |
| 68 | NUM_KV_HEADS: tl.constexpr, |
| 69 | HEAD_DIM: tl.constexpr, |
| 70 | BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages) |
| 71 | NUM_KV_SPLITS: tl.constexpr, |
| 72 | KV_GROUP_SIZE: tl.constexpr, # Hq // Hk |
| 73 | # TQ layout constants |
| 74 | MSE_BITS: tl.constexpr, # 3 or 4 |
| 75 | MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8) |
| 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 | # Opt#3 SoA layout |
| 79 | KEY_DATA_BYTES: tl.constexpr, # MSE_BYTES (MSE) or HEAD_DIM (FP8) |
| 80 | META_REGION_OFFSET: tl.constexpr, # bytes from block start to SoA region |
| 81 | NUM_SOA_FIELDS: tl.constexpr, # 3 (MSE) or 2 (FP8) |
| 82 | SOA_K_NORM: tl.constexpr, # 0 (MSE); unused for FP8 |
| 83 | SOA_V_SCALE: tl.constexpr, # 1 (MSE) / 0 (FP8) |
| 84 | SOA_V_ZERO: tl.constexpr, # 2 (MSE) / 1 (FP8) |
| 85 | # Score constants |
| 86 | ATTN_SCALE: tl.constexpr, # 1/sqrt(D) |
| 87 | # Block tile sizes |
| 88 | BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM) |
| 89 | BLOCK_KV: tl.constexpr, # tokens per tile (16) |
| 90 | KEY_FP8: tl.constexpr, # 1 if K is stored as FP8 |
| 91 | NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids |
| 92 | FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) |
| 93 | USE_SINKS: tl.constexpr = 0, # 1 = use sink tokens for attention anchoring |
| 94 | ): |
| 95 | bid = tl.program_id(0) # batch index |
| 96 | hid = tl.program_id(1) # q_head index |
| 97 | sid = tl.program_id(2) # kv_split index |
| 98 | |
| 99 | kv_head = hid // KV_GROUP_SIZE |
| 100 | |
| 101 | # Sequence length for this batch |
| 102 | seq_len = tl.load(Seq_lens_ptr + bid) |
| 103 | |
| 104 | # KV split range |
| 105 | split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) |
| 106 | split_start = split_len * sid |
| 107 | split_end = tl.minimum(split_start + split_len, seq_len) |
| 108 | |
| 109 | if split_start >= split_end: |
| 110 | return |
| 111 | |
| 112 | # Dimension offsets |
| 113 | d_offs = tl.arange(0, BLOCK_D) |
| 114 | d_mask = d_offs < HEAD_DIM |
| 115 | kv_range = tl.arange(0, BLOCK_KV) |
| 116 | |
| 117 | # Load query vector: q_rot — [BLOCK_D] float32 |
| 118 | q_base = bid * stride_qb + hid * stride_qh |
| 119 | q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32) |
| 120 | |
| 121 | # Precompute byte/bit index vectors for MSE gather loads |
| 122 | if not KEY_FP8: |
| 123 | mse_bit_off = d_offs * MSE_BITS |
| 124 | mse_byte_idx = mse_bit_off // 8 |
| 125 | mse_bit_shift = mse_bit_off % 8 |
| 126 | mse_mask = (1 << MSE_BITS) - 1 |
| 127 | |
| 128 | # Precompute value bit/byte index vectors (loop-invariant) |
| 129 | if VQB == 3: |
| 130 | val_bit_off = d_offs * 3 |
| 131 | val_byte_idx = val_bit_off // 8 |
| 132 | val_bit_shift = val_bit_off % 8 |
| 133 | |
| 134 | # Online softmax accumulators |
| 135 | # Sink tokens provide a pre-computed attention bias that should be |
| 136 | # included in the softmax normalization. For the first split (sid==0), |
| 137 | # we initialize m_prev with the sink logit. L is always 1.0 for all |
| 138 | # splits to match unified attention semantics. |
| 139 | if USE_SINKS: |
| 140 | if sid == 0: |
| 141 | m_prev = tl.load(Sink_ptr + hid).to(tl.float32) |
| 142 | l_prev = 1.0 |
| 143 | else: |
| 144 | m_prev = -float("inf") |
| 145 | l_prev = 0.0 |
| 146 | else: |
| 147 | m_prev = -float("inf") |
| 148 | l_prev = 0.0 |
| 149 | acc = tl.zeros([BLOCK_D], dtype=tl.float32) |
| 150 | |
| 151 | bt_base = bid * stride_bt_b |
| 152 | |
| 153 | # ================================================================ |
| 154 | # TILED LOOP: process BLOCK_KV tokens per iteration |
| 155 | # ================================================================ |
| 156 | for start_n in range(split_start, split_end, BLOCK_KV): |
| 157 | kv_offs = start_n + kv_range |
| 158 | kv_mask = kv_offs < split_end |
| 159 | |
| 160 | page_idx = kv_offs // BLOCK_SIZE |
| 161 | page_off = kv_offs % BLOCK_SIZE |
| 162 | block_nums = tl.load( |
| 163 | Block_table_ptr + bt_base + page_idx, |
| 164 | mask=kv_mask, |
| 165 | other=0, |
| 166 | ).to(tl.int64) |
| 167 | |
| 168 | # Opt#3 SoA addressing: data region then metadata region per block. |
| 169 | slot_within_block = page_off.to(tl.int64) |
| 170 | block_base = block_nums * stride_cache_block |
| 171 | DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES |
| 172 | data_bases = ( |
| 173 | block_base |
| 174 | + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT) |
| 175 | + tl.cast(kv_head, tl.int64) * DATA_BYTES_PER_SLOT |
| 176 | ) |
| 177 | head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast( |
| 178 | kv_head, tl.int64 |
| 179 | ) * (NUM_SOA_FIELDS * BLOCK_SIZE) |
| 180 | knorm_u16_addrs = ( |
| 181 | head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block |
| 182 | ) |
| 183 | vscale_u16_addrs = ( |
| 184 | head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block |
| 185 | ) |
| 186 | vzero_u16_addrs = ( |
| 187 | head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block |
| 188 | ) |
| 189 | |
| 190 | # ============================================================ |
| 191 | # COMPUTE ATTENTION SCORES: [BLOCK_KV] |
| 192 | # ============================================================ |
| 193 | if KEY_FP8: |
| 194 | k_addrs = data_bases[:, None] + d_offs[None, :] |
| 195 | k_raw = tl.load( |
| 196 | KV_cache_ptr + k_addrs, |
| 197 | mask=kv_mask[:, None] & d_mask[None, :], |
| 198 | other=0, |
| 199 | ) |
| 200 | if FP8_E4B15: |
| 201 | k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) |
| 202 | else: |
| 203 | k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) |
| 204 | scores = ( |
| 205 | tl.sum( |
| 206 | tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0), |
| 207 | axis=1, |
| 208 | ) |
| 209 | * ATTN_SCALE |
| 210 | ) |
| 211 | scores = tl.where(kv_mask, scores, -float("inf")) |
| 212 | else: |
| 213 | # MSE unpack + norms |
| 214 | mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :] |
| 215 | mse_raw0 = tl.load( |
| 216 | KV_cache_ptr + mse_addrs0, |
| 217 | mask=kv_mask[:, None] & d_mask[None, :], |
| 218 | other=0, |
| 219 | ).to(tl.int32) |
| 220 | mse_raw1 = tl.load( |
| 221 | KV_cache_ptr + mse_addrs0 + 1, |
| 222 | mask=kv_mask[:, None] & d_mask[None, :], |
| 223 | other=0, |
| 224 | ).to(tl.int32) |
| 225 | raw16 = mse_raw0 | (mse_raw1 << 8) |
| 226 | mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask |
| 227 | |
| 228 | # Centroid gather + dot product |
| 229 | c_vals = tl.load( |
| 230 | Centroids_ptr + mse_idx, |
| 231 | mask=kv_mask[:, None] & d_mask[None, :], |
| 232 | other=0.0, |
| 233 | ) |
| 234 | |
| 235 | # Opt#1: norm-correction (1/||c_vec||) is pre-folded into the |
| 236 | # stored per-token scalar by the MSE store kernel when |
| 237 | # NORM_CORRECTION=1, so we skip the per-tile sum+sqrt+divide. |
| 238 | # NORM_CORRECTION is kept in the signature for API parity; it is |
| 239 | # implicit in the stored norm value now. |
| 240 | |
| 241 | term1 = tl.sum( |
| 242 | tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0), |
| 243 | axis=1, |
| 244 | ) |
| 245 | |
| 246 | # Opt#3: K-norms live in the per-block SoA region. Single u16 |
| 247 | # load per token; contiguous across tokens within a block. |
| 248 | norm_u16 = tl.load( |
| 249 | KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask, other=0 |
| 250 | ) |
| 251 | vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 252 | |
| 253 | scores = vec_norms * term1 * ATTN_SCALE |
| 254 | scores = tl.where(kv_mask, scores, -float("inf")) |
| 255 | |
| 256 | # ============================================================ |
| 257 | # ONLINE SOFTMAX UPDATE (block-level) |
| 258 | # ============================================================ |
| 259 | n_e_max = tl.maximum(tl.max(scores, 0), m_prev) |
| 260 | re_scale = tl.exp(m_prev - n_e_max) |
| 261 | p = tl.exp(scores - n_e_max) |
| 262 | |
| 263 | # ============================================================ |
| 264 | # VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D] |
| 265 | # Opt#3: V data immediately follows K data in the slot (no |
| 266 | # interleaved metadata). Scale/zero come from the SoA metadata |
| 267 | # region as single u16 loads per field per token. |
| 268 | # ============================================================ |
| 269 | val_bases = data_bases + KEY_DATA_BYTES |
| 270 | |
| 271 | if VQB == 3: |
| 272 | val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] |
| 273 | val_raw0 = tl.load( |
| 274 | KV_cache_ptr + val_addrs0, |
| 275 | mask=kv_mask[:, None] & d_mask[None, :], |
| 276 | other=0, |
| 277 | ).to(tl.int32) |
| 278 | val_raw1 = tl.load( |
| 279 | KV_cache_ptr + val_addrs0 + 1, |
| 280 | mask=kv_mask[:, None] & d_mask[None, :], |
| 281 | other=0, |
| 282 | ).to(tl.int32) |
| 283 | raw16 = val_raw0 | (val_raw1 << 8) |
| 284 | v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) |
| 285 | else: # VQB == 4 |
| 286 | vb_idx = d_offs // 2 |
| 287 | vb_shift = (d_offs % 2) * 4 |
| 288 | val_addrs = val_bases[:, None] + vb_idx[None, :] |
| 289 | val_raw = tl.load( |
| 290 | KV_cache_ptr + val_addrs, |
| 291 | mask=kv_mask[:, None] & d_mask[None, :], |
| 292 | other=0, |
| 293 | ).to(tl.int32) |
| 294 | v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32) |
| 295 | |
| 296 | scale_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask, other=0) |
| 297 | zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask, other=0) |
| 298 | v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 299 | v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 300 | values = v_idx * v_scales[:, None] + v_zeros[:, None] |
| 301 | |
| 302 | # ============================================================ |
| 303 | # WEIGHTED VALUE ACCUMULATION |
| 304 | # ============================================================ |
| 305 | acc = acc * re_scale + tl.sum(p[:, None] * values, 0) |
| 306 | l_prev = l_prev * re_scale + tl.sum(p, 0) |
| 307 | m_prev = n_e_max |
| 308 | |
| 309 | # Store partial result |
| 310 | out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s |
| 311 | safe_l = tl.where(l_prev > 0.0, l_prev, 1.0) |
| 312 | tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask) |
| 313 | lse = m_prev + tl.log(safe_l) |
| 314 | tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse) |
| 315 | |
| 316 | |
| 317 | # --------------------------------------------------------------------------- |
| 318 | # Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16 |
| 319 | # --------------------------------------------------------------------------- |
| 320 | |
| 321 | |
| 322 | @triton.jit |
| 323 | def _tq_full_dequant_kv( |
| 324 | KV_cache_ptr, |
| 325 | KV_cache_u16_ptr, |
| 326 | Block_table_ptr, |
| 327 | Centroids_ptr, |
| 328 | K_out_ptr, # [B, Hk, max_seq, D] float16 |
| 329 | V_out_ptr, # [B, Hk, max_seq, D] float16 |
| 330 | stride_ko_b: tl.int64, |
| 331 | stride_ko_h: tl.int64, |
| 332 | stride_ko_s: tl.int64, |
| 333 | stride_vo_b: tl.int64, |
| 334 | stride_vo_h: tl.int64, |
| 335 | stride_vo_s: tl.int64, |
| 336 | stride_cache_block: tl.int64, |
| 337 | stride_bt_b: tl.int64, |
| 338 | HEAD_DIM: tl.constexpr, |
| 339 | BLOCK_SIZE: tl.constexpr, |
| 340 | NUM_KV_HEADS: tl.constexpr, |
| 341 | MSE_BYTES: tl.constexpr, |
| 342 | VQB: tl.constexpr, |
| 343 | VAL_DATA_BYTES: tl.constexpr, |
| 344 | MSE_BITS: tl.constexpr, |
| 345 | KEY_FP8: tl.constexpr, |
| 346 | # Opt#3 SoA layout |
| 347 | KEY_DATA_BYTES: tl.constexpr, |
| 348 | META_REGION_OFFSET: tl.constexpr, |
| 349 | NUM_SOA_FIELDS: tl.constexpr, |
| 350 | SOA_K_NORM: tl.constexpr, |
| 351 | SOA_V_SCALE: tl.constexpr, |
| 352 | SOA_V_ZERO: tl.constexpr, |
| 353 | BLOCK_D: tl.constexpr, |
| 354 | NORM_CORRECTION: tl.constexpr = 0, |
| 355 | FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) |
| 356 | ): |
| 357 | """Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16. |
| 358 | |
| 359 | Opt#3 SoA layout: data region [bs, H, KD+VD] + per-block SoA metadata. |
| 360 | """ |
| 361 | pos = tl.program_id(0) |
| 362 | bh = tl.program_id(1) |
| 363 | bid = bh // NUM_KV_HEADS |
| 364 | hid = bh % NUM_KV_HEADS |
| 365 | |
| 366 | page_idx = pos // BLOCK_SIZE |
| 367 | page_off = pos % BLOCK_SIZE |
| 368 | block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64) |
| 369 | |
| 370 | block_base = block_num * stride_cache_block |
| 371 | slot_within_block = tl.cast(page_off, tl.int64) |
| 372 | DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES |
| 373 | data_base = ( |
| 374 | block_base |
| 375 | + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT) |
| 376 | + tl.cast(hid, tl.int64) * DATA_BYTES_PER_SLOT |
| 377 | ) |
| 378 | head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast( |
| 379 | hid, tl.int64 |
| 380 | ) * (NUM_SOA_FIELDS * BLOCK_SIZE) |
| 381 | knorm_u16_addr = head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block |
| 382 | vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block |
| 383 | vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block |
| 384 | |
| 385 | d_offs = tl.arange(0, BLOCK_D) |
| 386 | d_mask = d_offs < HEAD_DIM |
| 387 | |
| 388 | # === K dequant === |
| 389 | ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s |
| 390 | if KEY_FP8: |
| 391 | k_raw = tl.load(KV_cache_ptr + data_base + d_offs, mask=d_mask, other=0) |
| 392 | if FP8_E4B15: |
| 393 | k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) |
| 394 | else: |
| 395 | k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) |
| 396 | tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) |
| 397 | else: |
| 398 | mse_bit_off = d_offs * MSE_BITS |
| 399 | mse_byte_idx = mse_bit_off // 8 |
| 400 | mse_bit_shift = mse_bit_off % 8 |
| 401 | mse_umask = (1 << MSE_BITS) - 1 |
| 402 | |
| 403 | mse_raw0 = tl.load( |
| 404 | KV_cache_ptr + data_base + mse_byte_idx, mask=d_mask, other=0 |
| 405 | ).to(tl.int32) |
| 406 | mse_raw1 = tl.load( |
| 407 | KV_cache_ptr + data_base + mse_byte_idx + 1, mask=d_mask, other=0 |
| 408 | ).to(tl.int32) |
| 409 | raw16_key = mse_raw0 | (mse_raw1 << 8) |
| 410 | mse_idx = (raw16_key >> mse_bit_shift) & mse_umask |
| 411 | |
| 412 | k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0) |
| 413 | |
| 414 | # Opt#1: norm-correction is pre-folded into the stored scalar. |
| 415 | # Opt#3: K-norm lives in the SoA region — single u16 load. |
| 416 | norm_u16 = tl.load(KV_cache_u16_ptr + knorm_u16_addr) |
| 417 | vec_norm = norm_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 418 | |
| 419 | k_recon = vec_norm * k_mse |
| 420 | tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) |
| 421 | |
| 422 | # === V dequant === |
| 423 | val_base = data_base + KEY_DATA_BYTES |
| 424 | if VQB == 4: |
| 425 | vb_idx = d_offs // 2 |
| 426 | vb_shift = (d_offs % 2) * 4 |
| 427 | val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to( |
| 428 | tl.int32 |
| 429 | ) |
| 430 | v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32) |
| 431 | elif VQB == 3: |
| 432 | val_bit_off = d_offs * 3 |
| 433 | val_byte_idx = val_bit_off // 8 |
| 434 | val_bit_shift = val_bit_off % 8 |
| 435 | val_raw0 = tl.load( |
| 436 | KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0 |
| 437 | ).to(tl.int32) |
| 438 | val_raw1 = tl.load( |
| 439 | KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0 |
| 440 | ).to(tl.int32) |
| 441 | raw16_val = val_raw0 | (val_raw1 << 8) |
| 442 | v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32) |
| 443 | else: |
| 444 | v_idx = tl.zeros([BLOCK_D], dtype=tl.float32) |
| 445 | |
| 446 | # Opt#3: V scale/zero single-u16 loads from SoA region. |
| 447 | sc_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addr) |
| 448 | zr_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addr) |
| 449 | v_scale = sc_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 450 | v_zero = zr_u16.to(tl.float16, bitcast=True).to(tl.float32) |
| 451 | v_vals = v_idx * v_scale + v_zero |
| 452 | |
| 453 | vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s |
| 454 | tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask) |
| 455 | |
| 456 | |
| 457 | # --------------------------------------------------------------------------- |
| 458 | # Stage 2: Reuse from triton_decode_attention.py |
| 459 | # --------------------------------------------------------------------------- |
| 460 | |
| 461 | # --------------------------------------------------------------------------- |
| 462 | # Launcher — cached constants + fused GEMM |
| 463 | # --------------------------------------------------------------------------- |
| 464 | |
| 465 | _layout_cache: dict = {} |
| 466 | |
| 467 | |
| 468 | def _get_layout(D, mse_bits, value_quant_bits, key_packed_size): |
| 469 | """Get cached layout constants.""" |
| 470 | key = (D, mse_bits, value_quant_bits, key_packed_size) |
| 471 | cfg = _layout_cache.get(key) |
| 472 | if cfg is None: |
| 473 | val_data_bytes = math.ceil(D * value_quant_bits / 8) |
| 474 | cfg = { |
| 475 | "mse_bytes": math.ceil(D * mse_bits / 8), |
| 476 | "val_data_bytes": val_data_bytes, |
| 477 | "mse_bits": mse_bits, |
| 478 | "n_centroids": 2**mse_bits, |
| 479 | "BLOCK_D": triton.next_power_of_2(D), |
| 480 | } |
| 481 | _layout_cache[key] = cfg |
| 482 | return cfg |
| 483 | |
| 484 | |
| 485 | def triton_turboquant_decode_attention( |
| 486 | query: torch.Tensor, # [B, Hq, D] — original query |
| 487 | kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 |
| 488 | block_table: torch.Tensor, # [B, max_num_blocks] int32 |
| 489 | seq_lens: torch.Tensor, # [B] int32 |
| 490 | Pi: torch.Tensor, # [D, D] float32 |
| 491 | centroids: torch.Tensor, # [n_centroids] float32 |
| 492 | scale: float, |
| 493 | mse_bits: int, |
| 494 | key_packed_size: int, |
| 495 | value_quant_bits: int, |
| 496 | key_fp8: bool = False, |
| 497 | norm_correction: bool = False, |
| 498 | PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous |
| 499 | # Pre-allocated buffers (optional, avoids per-call allocation) |
| 500 | mid_o_buf: torch.Tensor | None = None, |
| 501 | output_buf: torch.Tensor | None = None, |
| 502 | lse_buf: torch.Tensor | None = None, |
| 503 | buf_holder: Any = None, |
| 504 | max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph) |
| 505 | sinks: torch.Tensor | None = None, # [Hq] float32, pre-computed sink logits |
| 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 | # Opt#3 SoA layout constants (derived locally; match store-side values). |
| 520 | key_data_bytes = D if key_fp8 else cfg["mse_bytes"] |
| 521 | data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"] |
| 522 | meta_region_offset = block_size * Hk * data_bytes_per_slot |
| 523 | num_soa_fields = 2 if key_fp8 else 3 |
| 524 | soa_k_norm = 0 |
| 525 | soa_v_scale = 0 if key_fp8 else 1 |
| 526 | soa_v_zero = 1 if key_fp8 else 2 |
| 527 | kv_cache_u16 = kv_cache.view(torch.uint16) |
| 528 | |
| 529 | # Compute q_rot = q @ Pi.T (rotated query for MSE key scoring) |
| 530 | # FP8 path: pass query directly (float16); kernel casts inline. |
| 531 | # MSE path: still needs external GEMM (cuBLAS), so q_rot is float32. |
| 532 | if key_fp8: |
| 533 | q_rot = query.contiguous() |
| 534 | else: |
| 535 | q_float = query.float() |
| 536 | if PiT is None: |
| 537 | PiT = Pi.T.contiguous() |
| 538 | q_rot = (q_float @ PiT).contiguous() |
| 539 | |
| 540 | NUM_KV_SPLITS = max_num_kv_splits |
| 541 | |
| 542 | if ( |
| 543 | mid_o_buf is not None |
| 544 | and mid_o_buf.shape[0] >= B |
| 545 | and mid_o_buf.shape[2] >= NUM_KV_SPLITS |
| 546 | ): |
| 547 | mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] |
| 548 | else: |
| 549 | mid_o = torch.empty( |
| 550 | B, |
| 551 | Hq, |
| 552 | NUM_KV_SPLITS, |
| 553 | D + 1, |
| 554 | dtype=torch.float32, |
| 555 | device=device, |
| 556 | ) |
| 557 | if buf_holder is not None: |
| 558 | buf_holder._tq_mid_o_buf = mid_o |
| 559 | |
| 560 | # Stage 1: split-KV tiled attention scoring + value accumulation |
| 561 | fp8_e4b15 = _use_fp8_e4b15(device.index or 0) |
| 562 | BLOCK_KV = 4 |
| 563 | grid = (B, Hq, NUM_KV_SPLITS) |
| 564 | |
| 565 | # Prepare sink pointer (convert to float32 on device if provided) |
| 566 | use_sinks = sinks is not None |
| 567 | |
| 568 | _tq_decode_stage1[grid]( |
| 569 | q_rot, |
| 570 | kv_cache, |
| 571 | kv_cache_u16, |
| 572 | block_table, |
| 573 | seq_lens, |
| 574 | centroids, |
| 575 | mid_o, |
| 576 | sinks, |
| 577 | q_rot.stride(0), |
| 578 | q_rot.stride(1), |
| 579 | kv_cache.stride(0), |
| 580 | block_table.stride(0), |
| 581 | mid_o.stride(0), |
| 582 | mid_o.stride(1), |
| 583 | mid_o.stride(2), |
| 584 | NUM_KV_HEADS=Hk, |
| 585 | HEAD_DIM=D, |
| 586 | BLOCK_SIZE=block_size, |
| 587 | NUM_KV_SPLITS=NUM_KV_SPLITS, |
| 588 | KV_GROUP_SIZE=kv_group_size, |
| 589 | MSE_BITS=mse_bits, |
| 590 | MSE_BYTES=cfg["mse_bytes"], |
| 591 | VQB=value_quant_bits, |
| 592 | VAL_DATA_BYTES=cfg["val_data_bytes"], |
| 593 | KEY_DATA_BYTES=key_data_bytes, |
| 594 | META_REGION_OFFSET=meta_region_offset, |
| 595 | NUM_SOA_FIELDS=num_soa_fields, |
| 596 | SOA_K_NORM=soa_k_norm, |
| 597 | SOA_V_SCALE=soa_v_scale, |
| 598 | SOA_V_ZERO=soa_v_zero, |
| 599 | ATTN_SCALE=scale, |
| 600 | BLOCK_D=cfg["BLOCK_D"], |
| 601 | BLOCK_KV=BLOCK_KV, |
| 602 | KEY_FP8=1 if key_fp8 else 0, |
| 603 | NORM_CORRECTION=1 if norm_correction else 0, |
| 604 | FP8_E4B15=fp8_e4b15, |
| 605 | USE_SINKS=1 if use_sinks else 0, |
| 606 | num_warps=1, |
| 607 | num_stages=1, |
| 608 | ) |
| 609 | |
| 610 | # Stage 2: Reduce across KV splits |
| 611 | if output_buf is not None and output_buf.shape[0] >= B: |
| 612 | output = output_buf[:B, :Hq, :D] |
| 613 | else: |
| 614 | output = torch.empty(B, Hq, D, dtype=torch.float32, device=device) |
| 615 | if buf_holder is not None: |
| 616 | buf_holder._tq_output_buf = output |
| 617 | if lse_buf is not None and lse_buf.shape[0] >= B: |
| 618 | lse = lse_buf[:B, :Hq] |
| 619 | else: |
| 620 | lse = torch.empty(B, Hq, dtype=torch.float32, device=device) |
| 621 | if buf_holder is not None: |
| 622 | buf_holder._tq_lse_buf = lse |
| 623 | |
| 624 | grid2 = (B, Hq) |
| 625 | _fwd_kernel_stage2[grid2]( |
| 626 | mid_o, |
| 627 | output, |
| 628 | lse, |
| 629 | seq_lens, |
| 630 | mid_o.stride(0), |
| 631 | mid_o.stride(1), |
| 632 | mid_o.stride(2), |
| 633 | output.stride(0), |
| 634 | output.stride(1), |
| 635 | lse.stride(0), |
| 636 | NUM_KV_SPLITS=NUM_KV_SPLITS, |
| 637 | BLOCK_DV=cfg["BLOCK_D"], |
| 638 | Lv=D, |
| 639 | num_warps=4, |
| 640 | num_stages=2, |
| 641 | ) |
| 642 | |
| 643 | return output.to(query.dtype) |