mirrors / lmcache--lmcache

项目文件夹

lmcache--lmcache
文件
main
T
wehub-resource-sync f213ec8976 chore: import upstream snapshot with attribution
2026-07-13 12:24:33 +08:00

446 行
16 KiB
Python
原始文件 永久链接 Blame 文件历史

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Fused Triton kernels for TurboQuant KV store.
Two kernels:
1. _tq_fused_store_fp8: FP8 key scatter + value uniform quantization.
2. _tq_fused_store_mse: Fused binary-search bucketize + MSE index
packing + value quantization.
The launcher `triton_turboquant_store` selects the appropriate kernel.
"""
# Standard
import math
# Third Party
import torch
import triton
import triton.language as tl
# First Party
from lmcache.v1.distributed.serde.turboquant.decode_kernel import _use_fp8_e4b15
# ═══════════════════════════════════════════════════════════════════════
# Shared: value uniform quantization + pack + scale/zero store
# ═══════════════════════════════════════════════════════════════════════
@triton.jit
def _store_quantized_value(
Value_ptr,
KV_cache_ptr,
base, # pid * D offset into Value_ptr
slot_base, # byte offset into KV_cache_ptr for this slot+head
d_offs, # tl.arange(0, BLOCK_D)
d_mask, # d_offs < D
D: tl.constexpr,
KPS: tl.constexpr,
VQB: tl.constexpr,
VAL_DATA_BYTES: tl.constexpr,
BLOCK_D: tl.constexpr,
BLOCK_VAL: tl.constexpr,
BLOCK_GRP: tl.constexpr,
):
"""Uniform quantization of values to VQB bits, pack, and store with scale/zero."""
val_cache_offset = KPS
if VQB == 3:
val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
tl.float32
)
val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
v_scale = (val_max - val_min) / 7.0
v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
q_vals = tl.minimum(
tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7
)
grp_offs = tl.arange(0, BLOCK_GRP)
grp_mask = grp_offs < (D // 8)
q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8])
shifts_3bit = tl.arange(0, 8) * 3
packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1)
b0 = (packed_24 & 0xFF).to(tl.uint8)
b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
tl.store(
KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3,
b0,
mask=grp_mask,
)
tl.store(
KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 1,
b1,
mask=grp_mask,
)
tl.store(
KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 2,
b2,
mask=grp_mask,
)
sc_offset = val_cache_offset + VAL_DATA_BYTES
sc_f16 = v_scale.to(tl.float16)
sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
tl.store(
KV_cache_ptr + slot_base + sc_offset + 1,
((sc_u16 >> 8) & 0xFF).to(tl.uint8),
)
zr_f16 = val_min.to(tl.float16)
zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
tl.store(
KV_cache_ptr + slot_base + sc_offset + 3,
((zr_u16 >> 8) & 0xFF).to(tl.uint8),
)
else: # VQB == 4
val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
tl.float32
)
val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
v_scale = (val_max - val_min) / 15.0
v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
# Quantize all D elements from register (no re-load)
q_all = tl.minimum(
tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15
)
# Reshape to pairs and pack two 4-bit values per byte
q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2])
shifts_4 = tl.arange(0, 2) * 4
packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
val_offs = tl.arange(0, BLOCK_D // 2)
val_mask = val_offs < VAL_DATA_BYTES
tl.store(
KV_cache_ptr + slot_base + val_cache_offset + val_offs,
packed_val,
mask=val_mask,
)
sc_offset = val_cache_offset + VAL_DATA_BYTES
sc_f16 = v_scale.to(tl.float16)
sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
tl.store(
KV_cache_ptr + slot_base + sc_offset + 1,
((sc_u16 >> 8) & 0xFF).to(tl.uint8),
)
zr_f16 = val_min.to(tl.float16)
zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
tl.store(
KV_cache_ptr + slot_base + sc_offset + 3,
((zr_u16 >> 8) & 0xFF).to(tl.uint8),
)
# ═══════════════════════════════════════════════════════════════════════
# FP8 key store + value uniform quantization
# ═══════════════════════════════════════════════════════════════════════
@triton.jit
def _tq_fused_store_fp8(
Key_ptr, # [NH, D] float16/bfloat16 — raw keys
Value_ptr, # [NH, D] float16/bfloat16 — raw values
KV_cache_ptr, # [total_bytes] uint8 (flattened view)
Slot_mapping_ptr, # [N] int32 — per-token slot indices
# Cache strides (for computing byte offsets)
stride_cache_block: tl.constexpr,
stride_cache_pos: tl.constexpr,
stride_cache_head: tl.constexpr,
# Dimensions
D: tl.constexpr,
H: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
BLOCK_D: tl.constexpr,
# TQ layout
KPS: tl.constexpr,
# Value quantization
VQB: tl.constexpr,
VAL_DATA_BYTES: tl.constexpr,
# Packing block sizes
BLOCK_VAL: tl.constexpr,
BLOCK_GRP: tl.constexpr = 16,
FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
):
"""FP8 key cast+scatter + value uniform quantization."""
pid = tl.program_id(0)
token_idx = pid // H
head_idx = pid % H
slot = tl.load(Slot_mapping_ptr + token_idx)
if slot < 0:
return
blk = slot // BLOCK_SIZE
off = slot % BLOCK_SIZE
slot_base = (
blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
)
base = pid * D
# ── FP8 KEY: cast to FP8 in-kernel and store ─────────────────
d_offs = tl.arange(0, BLOCK_D)
d_mask = d_offs < D
k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0)
k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
k_bytes = k_fp8.to(tl.uint8, bitcast=True)
tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)
# ── VALUE QUANTIZE + PACK ───────────────────────────────────────
_store_quantized_value(
Value_ptr,
KV_cache_ptr,
base,
slot_base,
d_offs,
d_mask,
D=D,
KPS=KPS,
VQB=VQB,
VAL_DATA_BYTES=VAL_DATA_BYTES,
BLOCK_D=BLOCK_D,
BLOCK_VAL=BLOCK_VAL,
BLOCK_GRP=BLOCK_GRP,
)
# ═══════════════════════════════════════════════════════════════════════
# Fused MSE store: bucketize + MSE index pack + norm store + value pack
# (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel)
# ═══════════════════════════════════════════════════════════════════════
@triton.jit
def _tq_fused_store_mse(
# Post-rotation inputs
Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT)
Norms_ptr, # [NH] float32 — key vector norms (||k||)
Value_ptr, # [NH, D] float32 — raw values
# Quantization tables
Midpoints_ptr, # [n_centroids-1] float32
# Cache and indexing
KV_cache_ptr, # [total_bytes] uint8 (flattened view)
Slot_mapping_ptr, # [N] int32 — per-token slot indices
# Cache strides
stride_cache_block: tl.constexpr,
stride_cache_pos: tl.constexpr,
stride_cache_head: tl.constexpr,
# Dimensions
D: tl.constexpr,
H: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
BLOCK_D: tl.constexpr,
# TQ layout
MSE_BYTES: tl.constexpr,
KPS: tl.constexpr,
# Value quantization
VQB: tl.constexpr,
VAL_DATA_BYTES: tl.constexpr,
# Packing block sizes
BLOCK_VAL: tl.constexpr,
# MSE params
MSE_BITS: tl.constexpr,
N_CENTROIDS: tl.constexpr,
BLOCK_GRP: tl.constexpr = 16,
):
"""Fused MSE quantize + pack + store.
Performs binary-search bucketize, MSE index packing, norm storage,
and value quantization in one kernel.
"""
pid = tl.program_id(0)
token_idx = pid // H
head_idx = pid % H
slot = tl.load(Slot_mapping_ptr + token_idx)
if slot < 0:
return
blk = slot // BLOCK_SIZE
off = slot % BLOCK_SIZE
slot_base = (
blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
)
base = pid * D
d_offs = tl.arange(0, BLOCK_D)
d_mask = d_offs < D
# ── 1. BINARY SEARCH BUCKETIZE ───────────────────────────────────
# Midpoints are sorted (N_CENTROIDS-1 values); binary search finds
# insertion point in MSE_BITS iterations vs N_CENTROIDS-1 for linear.
y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)
lo = tl.zeros([BLOCK_D], dtype=tl.int32)
hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)
for _ in range(MSE_BITS):
mid = (lo + hi) >> 1
# Clamp to valid midpoint index [0, N_CENTROIDS-2] for load safety;
# the search result (lo) is still correct since converged lanes
# don't change.
safe_mid = tl.minimum(mid, N_CENTROIDS - 2)
mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)
lo = tl.where(y_vec >= mid_val, mid + 1, lo)
hi = tl.where(y_vec >= mid_val, hi, mid)
idx = tl.minimum(lo, N_CENTROIDS - 1)
# ── 2. PACK MSE INDICES from register idx ─────────────────────────
if MSE_BITS == 4:
idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])
shifts_4 = tl.arange(0, 2) * 4
packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
mse_offs = tl.arange(0, BLOCK_D // 2)
mse_mask = mse_offs < MSE_BYTES
tl.store(KV_cache_ptr + slot_base + mse_offs, packed, mask=mse_mask)
elif MSE_BITS == 3:
grp_offs = tl.arange(0, BLOCK_GRP)
grp_mask = grp_offs < (D // 8)
idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])
shifts_3 = tl.arange(0, 8) * 3
packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)
b0 = (packed_24 & 0xFF).to(tl.uint8)
b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
tl.store(KV_cache_ptr + slot_base + grp_offs * 3, b0, mask=grp_mask)
tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 1, b1, mask=grp_mask)
tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 2, b2, mask=grp_mask)
# ── 3. STORE vec_norm (fp16, 2 bytes) ─────────────────────────────
norm_offset = MSE_BYTES
vn_f16 = tl.load(Norms_ptr + pid).to(tl.float16)
vn_u16 = vn_f16.to(tl.uint16, bitcast=True)
tl.store(KV_cache_ptr + slot_base + norm_offset, (vn_u16 & 0xFF).to(tl.uint8))
tl.store(
KV_cache_ptr + slot_base + norm_offset + 1, ((vn_u16 >> 8) & 0xFF).to(tl.uint8)
)
# ── 4. VALUE QUANTIZE + PACK ──────────────────────────────────────
_store_quantized_value(
Value_ptr,
KV_cache_ptr,
base,
slot_base,
d_offs,
d_mask,
D=D,
KPS=KPS,
VQB=VQB,
VAL_DATA_BYTES=VAL_DATA_BYTES,
BLOCK_D=BLOCK_D,
BLOCK_VAL=BLOCK_VAL,
BLOCK_GRP=BLOCK_GRP,
)
# ═══════════════════════════════════════════════════════════════════════
# Launcher
# ═══════════════════════════════════════════════════════════════════════
def triton_turboquant_store(
key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE)
value: torch.Tensor, # [N, H, D] — raw values
kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
slot_mapping: torch.Tensor, # [N] int32
PiT: torch.Tensor, # [D, D] float32
midpoints: torch.Tensor, # [n_centroids-1] float32
mse_bits: int,
key_packed_size: int,
value_quant_bits: int,
key_fp8: bool = False,
):
"""Launch TQ store kernel (FP8 or MSE path)."""
N, H, D = key.shape
NH = N * H
block_size = kv_cache.shape[1]
BLOCK_D = triton.next_power_of_2(D)
mse_bytes = math.ceil(D * mse_bits / 8)
n_centroids = 2**mse_bits
val_data_bytes = math.ceil(D * value_quant_bits / 8)
BLOCK_VAL = triton.next_power_of_2(val_data_bytes)
# Cache strides (element_size=1 for uint8, so stride in bytes = stride())
stride_block = kv_cache.stride(0)
stride_pos = kv_cache.stride(1)
stride_head = kv_cache.stride(2)
block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1
# ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ──
if key_fp8:
k_flat = key.reshape(NH, D).contiguous()
v_flat = value.reshape(NH, D).contiguous()
fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0)
grid = (NH,)
_tq_fused_store_fp8[grid](
k_flat,
v_flat,
kv_cache.view(-1),
slot_mapping,
stride_cache_block=stride_block,
stride_cache_pos=stride_pos,
stride_cache_head=stride_head,
D=D,
H=H,
BLOCK_SIZE=block_size,
BLOCK_D=BLOCK_D,
KPS=key_packed_size,
VQB=value_quant_bits,
VAL_DATA_BYTES=val_data_bytes,
BLOCK_VAL=BLOCK_VAL,
BLOCK_GRP=block_grp,
FP8_E4B15=fp8_e4b15,
num_warps=4,
num_stages=1,
)
return
# ── MSE PATH: external GEMM + fused bucketize/pack kernel ──
# Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)
k_flat = key.float().reshape(NH, D)
norms = k_flat.norm(dim=1, keepdim=True)
x_hat = k_flat / (norms + 1e-8)
y = x_hat @ PiT
v_flat = value.float().reshape(NH, D)
# Fused kernel: bucketize + MSE index pack + norm store + value pack
grid = (NH,)
_tq_fused_store_mse[grid](
y,
norms.squeeze(1),
v_flat,
midpoints,
kv_cache.view(-1),
slot_mapping,
stride_cache_block=stride_block,
stride_cache_pos=stride_pos,
stride_cache_head=stride_head,
D=D,
H=H,
BLOCK_SIZE=block_size,
BLOCK_D=BLOCK_D,
MSE_BYTES=mse_bytes,
KPS=key_packed_size,
VQB=value_quant_bits,
VAL_DATA_BYTES=val_data_bytes,
BLOCK_VAL=BLOCK_VAL,
MSE_BITS=mse_bits,
N_CENTROIDS=n_centroids,
BLOCK_GRP=block_grp,
num_warps=4,
num_stages=1,
)
在新工单中引用 查看 Git Blame 复制永久链接