diff --git a/sglang-research/python/sglang/QuantKernel/__init__.py b/sglang-research/python/sglang/QuantKernel/__init__.py index 10002cd25..c5a05045e 100644 --- a/sglang-research/python/sglang/QuantKernel/__init__.py +++ b/sglang-research/python/sglang/QuantKernel/__init__.py @@ -1,4 +1,4 @@ -"""Quantization-related kernels (Hadamard + int2 KV fusion).""" +"""Quantization-related kernels (Hadamard + int2/int1/PQ KV fusion).""" from sglang.QuantKernel.fused_hadamard_int2_kv import ( MAX_HADAMARD_ORDER, @@ -7,6 +7,12 @@ validate_hadamard_order_for_kv_fuse_int2, ) from sglang.QuantKernel.gpu_flush_int2 import gpu_flush_int2 +from sglang.QuantKernel.gpu_flush_int1 import gpu_flush_int1 +from sglang.QuantKernel.oscar_rotation_clip_int1_kv import ( + quantized_set_kv_int1_oscar_rotate_k_clip_triton, + quantized_set_kv_int1_pretransformed_clip_triton, + quantized_set_kv_int1_pretransformed_triton, +) __all__ = [ "MAX_HADAMARD_ORDER", @@ -14,4 +20,8 @@ "quantized_set_kv_int2_pretransformed_triton", "validate_hadamard_order_for_kv_fuse_int2", "gpu_flush_int2", + "gpu_flush_int1", + "quantized_set_kv_int1_oscar_rotate_k_clip_triton", + "quantized_set_kv_int1_pretransformed_clip_triton", + "quantized_set_kv_int1_pretransformed_triton", ] diff --git a/sglang-research/python/sglang/QuantKernel/gpu_flush_int1.py b/sglang-research/python/sglang/QuantKernel/gpu_flush_int1.py new file mode 100644 index 000000000..fb511083b --- /dev/null +++ b/sglang-research/python/sglang/QuantKernel/gpu_flush_int1.py @@ -0,0 +1,823 @@ +""" +GPU-only HP -> int1 flush for the unified mixed KV pool. + +INT1 variant of :mod:`sglang.QuantKernel.gpu_flush_int2`. Same plan/apply +structure and same FlushPlan dataclass shape. The only differences are: + + * The quant cache holds ``head_dim // 8`` bytes per (token, head) + instead of ``head_dim // 4``. + * The quant kernel packs 8 quant slots per byte (bit ``i`` ↔ position + ``i * BLOCK_OCTANT + bo`` of the row). + * ``max_q == 1`` so ``scale = range`` (no divisor) and the quant slot + holds {0, 1}. + +Scale/zero layout (interleaved ``(scale, zero)`` per group) is unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Tuple + +import torch +import triton +import triton.language as tl + + +@dataclass +class FlushPlanInt1: + """Output of ``gpu_flush_int1_plan``; consumed by ``gpu_flush_int1_apply``. + + Same shape as the INT2 ``FlushPlan``; duplicated as its own dataclass so + the int1 / int2 plan objects don't accidentally cross paths. + """ + + returned_slot_ids: torch.Tensor + src_hp_slot: torch.Tensor + flush_pos: torch.Tensor + valid_mask: torch.Tensor + dst_quant_slots: torch.Tensor + bs: int + flush_interval: int + + +# --------------------------------------------------------------------------- +# Plan kernel — identical to int2 (it does not touch the quant arena). +# --------------------------------------------------------------------------- + + +@triton.jit +def _flush_plan_kernel_int1( + seq_lens_ptr, + prefix_lens_ptr, + req_pool_indices_ptr, + dst_quant_slot_ptr, + req_to_token_ptr, + flush_mask_ptr, + src_hp_slot_out_ptr, + returned_slot_ids_ptr, + flush_pos_out_ptr, + valid_mask_out_ptr, + max_ctx, + rtt_stride_row, + HP_PREFIX_TOKENS: tl.constexpr, + HP_RECENT_TOKENS: tl.constexpr, + HP_OFFSET: tl.constexpr, + FLUSH_INTERVAL: tl.constexpr, +): + i = tl.program_id(0) + do_flush = tl.load(flush_mask_ptr + i).to(tl.int32) + seq_len = tl.load(seq_lens_ptr + i).to(tl.int32) + prefix_len = tl.load(prefix_lens_ptr + i).to(tl.int32) + req_pool_idx = tl.load(req_pool_indices_ptr + i).to(tl.int64) + + for j in tl.static_range(FLUSH_INTERVAL): + out_idx = i * FLUSH_INTERVAL + j + dst_q = tl.load(dst_quant_slot_ptr + out_idx).to(tl.int64) + + valid = 0 + src_hp = tl.full((), -1, tl.int64) + flush_pos = tl.full((), -1, tl.int32) + + if do_flush == 1 and HP_RECENT_TOKENS > 0: + fp = seq_len - HP_RECENT_TOKENS - (FLUSH_INTERVAL - 1) + j + if fp >= prefix_len and fp >= 0: + loc = tl.load( + req_to_token_ptr + req_pool_idx * rtt_stride_row + fp.to(tl.int64) + ).to(tl.int64) + if loc >= HP_OFFSET: + src_hp = loc - HP_OFFSET + valid = 1 + flush_pos = fp + + tl.store(valid_mask_out_ptr + out_idx, tl.full((), valid, tl.int8)) + tl.store(flush_pos_out_ptr + out_idx, flush_pos) + + if valid == 1: + tl.store(returned_slot_ids_ptr + out_idx, src_hp + HP_OFFSET) + tl.store(src_hp_slot_out_ptr + out_idx, src_hp) + else: + tl.store(returned_slot_ids_ptr + out_idx, dst_q) + tl.store(src_hp_slot_out_ptr + out_idx, -1) + + +# --------------------------------------------------------------------------- +# Fused INT1 quant kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _fused_flush_quant_body_int1( + hp_base, + quant_base, + sz_base, + src_hp_slot, + dst_quant_slot, + active, + head_idx, + HP_STRIDE_LOC: tl.constexpr, + HP_STRIDE_HEAD: tl.constexpr, + HP_STRIDE_DIM: tl.constexpr, + Q_STRIDE_LOC: tl.constexpr, + Q_STRIDE_HEAD: tl.constexpr, + Q_STRIDE_DIM: tl.constexpr, + SZ_STRIDE_LOC: tl.constexpr, + SZ_STRIDE_HEAD: tl.constexpr, + SZ_STRIDE_DIM: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_OCTANT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK_TOK: tl.constexpr, + CLIP_INDEX: tl.constexpr, + BSEARCH_ITERS: tl.constexpr, + LLOYD_MAX: tl.constexpr, +): + """Quantize ``BLOCK_TOK`` (src_hp_slot, head) HP rows into int1 at the + matching ``dst_quant_slot``s. 8-way pack mirrors INT2's quartered split. + + When ``LLOYD_MAX`` is set the per-group binary quantizer matches the + prefill grouped set_kv kernel: boundary at the group mean, centroids at + ``mean ± std*sqrt(2/pi)`` (scale = ``2*sqrt(2/pi)*std``). Without it the + legacy uniform binary (centroids at the clipped min/max) is used, which is + catastrophic for 1-bit since both levels land at the distribution extremes. + """ + full_offs = tl.arange(0, HEAD_DIM) + base = ( + src_hp_slot[:, None] * HP_STRIDE_LOC + + head_idx * HP_STRIDE_HEAD + + full_offs[None, :] * HP_STRIDE_DIM + ) + acc = tl.load( + hp_base + base, + mask=active[:, None], + other=0.0, + ).to(tl.float32) + + if CLIP_INDEX >= 0: + abs_acc = tl.abs(acc) + if BSEARCH_ITERS > 0: + target_above = HEAD_DIM - CLIP_INDEX + thr_lo = tl.zeros([BLOCK_TOK], dtype=tl.float32) + thr_hi = tl.max(abs_acc, axis=1) + for _ in tl.static_range(BSEARCH_ITERS): + thr_mid = (thr_lo + thr_hi) * 0.5 + cnt_above = tl.sum((abs_acc > thr_mid[:, None]).to(tl.int32), axis=1) + too_many = cnt_above > target_above + thr_lo = tl.where(too_many, thr_mid, thr_lo) + thr_hi = tl.where(too_many, thr_hi, thr_mid) + thr = thr_hi + else: + sorted_acc = tl.sort(abs_acc) + pick = (full_offs == CLIP_INDEX)[None, :] + thr = tl.sum(tl.where(pick, sorted_acc, 0.0), axis=1) + acc = tl.minimum( + tl.maximum(acc, -thr[:, None]), + thr[:, None], + ) + + grouped = tl.reshape(acc, (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) + if LLOYD_MAX: + LM_C_1BIT: tl.constexpr = 0.79788456 # sqrt(2/pi) + group_mean = tl.sum(grouped, axis=2) / GROUP_SIZE + group_diff = grouped - group_mean[:, :, None] + group_var = tl.sum(group_diff * group_diff, axis=2) / GROUP_SIZE + group_std = tl.sqrt(group_var + 1e-8) + scale = tl.maximum(2.0 * LM_C_1BIT * group_std, 1e-8) + zero = 0.5 - group_mean / scale + else: + val_min = tl.min(grouped, axis=2) + val_max = tl.max(grouped, axis=2) + scale = tl.maximum(val_max - val_min, 1e-8) # max_q = 1 + zero = tl.math.div_rn(-val_min, scale) + + # Octant split via reshape + permute + split×3 on fp32 acc, broadcasting + # per-group scale/zero through the same pipeline. Mirrors INT2's quartered + # split (split×2); just one more level deep. + acc_r = tl.reshape(acc, (BLOCK_TOK, 8, BLOCK_OCTANT)) + acc_p = tl.permute(acc_r, (0, 2, 1)) + acc_s = tl.reshape(acc_p, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + a_lo, a_hi = tl.split(acc_s) + a_e0, a_e2 = tl.split(a_lo) + a_e1, a_e3 = tl.split(a_hi) + vals0, vals4 = tl.split(a_e0) + vals2, vals6 = tl.split(a_e2) + vals1, vals5 = tl.split(a_e1) + vals3, vals7 = tl.split(a_e3) + + scale_3d = tl.broadcast_to(scale[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) + zero_3d = tl.broadcast_to(zero[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) + scale_flat = tl.reshape(scale_3d, (BLOCK_TOK, HEAD_DIM)) + zero_flat = tl.reshape(zero_3d, (BLOCK_TOK, HEAD_DIM)) + + sr = tl.reshape(scale_flat, (BLOCK_TOK, 8, BLOCK_OCTANT)) + sp = tl.permute(sr, (0, 2, 1)) + ss = tl.reshape(sp, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + s_lo, s_hi = tl.split(ss) + s_e0, s_e2 = tl.split(s_lo) + s_e1, s_e3 = tl.split(s_hi) + s0, s4 = tl.split(s_e0) + s2, s6 = tl.split(s_e2) + s1, s5 = tl.split(s_e1) + s3, s7 = tl.split(s_e3) + + zr = tl.reshape(zero_flat, (BLOCK_TOK, 8, BLOCK_OCTANT)) + zp = tl.permute(zr, (0, 2, 1)) + zs = tl.reshape(zp, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + z_lo, z_hi = tl.split(zs) + z_e0, z_e2 = tl.split(z_lo) + z_e1, z_e3 = tl.split(z_hi) + z0, z4 = tl.split(z_e0) + z2, z6 = tl.split(z_e2) + z1, z5 = tl.split(z_e1) + z3, z7 = tl.split(z_e3) + + # Clamp the rounding result to {0, 1} before packing. For the uniform path + # this is a no-op (qf in [0.5, 1.5] -> trunc gives {0, 1}); for the LM path + # it is required, since the LM scale = 2*c*std makes qf = (val-mean)/scale + + # 1.0 land in [-1, 2] for tail values, which would otherwise pack as 2 or + # wrap to 255 and corrupt the byte. Clamped, the boundary sits exactly at + # val == mean (q = 1 iff val >= mean), matching the prefill grouped kernel. + q0 = tl.minimum(tl.maximum(tl.math.div_rn(vals0, s0) + z0 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q1 = tl.minimum(tl.maximum(tl.math.div_rn(vals1, s1) + z1 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q2 = tl.minimum(tl.maximum(tl.math.div_rn(vals2, s2) + z2 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q3 = tl.minimum(tl.maximum(tl.math.div_rn(vals3, s3) + z3 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q4 = tl.minimum(tl.maximum(tl.math.div_rn(vals4, s4) + z4 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q5 = tl.minimum(tl.maximum(tl.math.div_rn(vals5, s5) + z5 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q6 = tl.minimum(tl.maximum(tl.math.div_rn(vals6, s6) + z6 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + q7 = tl.minimum(tl.maximum(tl.math.div_rn(vals7, s7) + z7 + 0.5, 0.0), 1.0).to( + tl.uint8 + ) + packed = ( + q0 + | (q1 << 1) + | (q2 << 2) + | (q3 << 3) + | (q4 << 4) + | (q5 << 5) + | (q6 << 6) + | (q7 << 7) + ) + + dim_offs_o = tl.arange(0, BLOCK_OCTANT) + cache_offset = ( + dst_quant_slot[:, None] * Q_STRIDE_LOC + + head_idx * Q_STRIDE_HEAD + + dim_offs_o[None, :] * Q_STRIDE_DIM + ) + tl.store(quant_base + cache_offset, packed, mask=active[:, None]) + + group_ids = tl.arange(0, NUM_GROUPS) + sz_offset_base = dst_quant_slot[:, None] * SZ_STRIDE_LOC + head_idx * SZ_STRIDE_HEAD + tl.store( + sz_base + sz_offset_base + (group_ids[None, :] * 2) * SZ_STRIDE_DIM, + scale, + mask=active[:, None], + ) + tl.store( + sz_base + sz_offset_base + (group_ids[None, :] * 2 + 1) * SZ_STRIDE_DIM, + zero, + mask=active[:, None], + ) + + +@triton.jit +def _fused_flush_quant_kernel_int1( + hp_k_ptrs_ptr, + hp_v_ptrs_ptr, + quant_k_ptrs_ptr, + quant_v_ptrs_ptr, + k_sz_ptrs_ptr, + v_sz_ptrs_ptr, + hp_k_sample_ptr, + hp_v_sample_ptr, + quant_k_sample_ptr, + quant_v_sample_ptr, + k_sz_sample_ptr, + v_sz_sample_ptr, + src_hp_slot_ptr, + dst_quant_slot_ptr, + valid_mask_ptr, + num_flush_tokens, + num_heads, + num_layers, + HP_K_STRIDE_LOC: tl.constexpr, + HP_K_STRIDE_HEAD: tl.constexpr, + HP_K_STRIDE_DIM: tl.constexpr, + HP_V_STRIDE_LOC: tl.constexpr, + HP_V_STRIDE_HEAD: tl.constexpr, + HP_V_STRIDE_DIM: tl.constexpr, + Q_K_STRIDE_LOC: tl.constexpr, + Q_K_STRIDE_HEAD: tl.constexpr, + Q_K_STRIDE_DIM: tl.constexpr, + Q_V_STRIDE_LOC: tl.constexpr, + Q_V_STRIDE_HEAD: tl.constexpr, + Q_V_STRIDE_DIM: tl.constexpr, + K_SZ_STRIDE_LOC: tl.constexpr, + K_SZ_STRIDE_HEAD: tl.constexpr, + K_SZ_STRIDE_DIM: tl.constexpr, + V_SZ_STRIDE_LOC: tl.constexpr, + V_SZ_STRIDE_HEAD: tl.constexpr, + V_SZ_STRIDE_DIM: tl.constexpr, + K_HEAD_DIM: tl.constexpr, + K_BLOCK_OCTANT: tl.constexpr, + K_NUM_GROUPS: tl.constexpr, + K_GROUP_SIZE: tl.constexpr, + V_HEAD_DIM: tl.constexpr, + V_BLOCK_OCTANT: tl.constexpr, + V_NUM_GROUPS: tl.constexpr, + V_GROUP_SIZE: tl.constexpr, + BLOCK_TOK: tl.constexpr, + K_CLIP_INDEX: tl.constexpr, + V_CLIP_INDEX: tl.constexpr, + K_BSEARCH_ITERS: tl.constexpr, + V_BSEARCH_ITERS: tl.constexpr, + LLOYD_MAX: tl.constexpr, +): + pid_tok = tl.program_id(0) + head = tl.program_id(1) + layer = tl.program_id(2) + if head >= num_heads or layer >= num_layers: + return + + tok_offs = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + tok_mask = tok_offs < num_flush_tokens + + valid = tl.load(valid_mask_ptr + tok_offs, mask=tok_mask, other=0).to(tl.int32) + if tl.max(valid, axis=0) == 0: + return + + active = tok_mask & (valid != 0) + src = tl.load(src_hp_slot_ptr + tok_offs, mask=tok_mask, other=0).to(tl.int64) + dst = tl.load(dst_quant_slot_ptr + tok_offs, mask=tok_mask, other=0).to(tl.int64) + head64 = head.to(tl.int64) + + hp_k_base = tl.load(hp_k_ptrs_ptr + layer).to( + tl.pointer_type(hp_k_sample_ptr.dtype.element_ty) + ) + q_k_base = tl.load(quant_k_ptrs_ptr + layer).to( + tl.pointer_type(quant_k_sample_ptr.dtype.element_ty) + ) + sz_k_base = tl.load(k_sz_ptrs_ptr + layer).to( + tl.pointer_type(k_sz_sample_ptr.dtype.element_ty) + ) + _fused_flush_quant_body_int1( + hp_k_base, + q_k_base, + sz_k_base, + src, + dst, + active, + head64, + HP_K_STRIDE_LOC, + HP_K_STRIDE_HEAD, + HP_K_STRIDE_DIM, + Q_K_STRIDE_LOC, + Q_K_STRIDE_HEAD, + Q_K_STRIDE_DIM, + K_SZ_STRIDE_LOC, + K_SZ_STRIDE_HEAD, + K_SZ_STRIDE_DIM, + K_HEAD_DIM, + K_BLOCK_OCTANT, + K_NUM_GROUPS, + K_GROUP_SIZE, + BLOCK_TOK, + K_CLIP_INDEX, + K_BSEARCH_ITERS, + LLOYD_MAX, + ) + hp_v_base = tl.load(hp_v_ptrs_ptr + layer).to( + tl.pointer_type(hp_v_sample_ptr.dtype.element_ty) + ) + q_v_base = tl.load(quant_v_ptrs_ptr + layer).to( + tl.pointer_type(quant_v_sample_ptr.dtype.element_ty) + ) + sz_v_base = tl.load(v_sz_ptrs_ptr + layer).to( + tl.pointer_type(v_sz_sample_ptr.dtype.element_ty) + ) + _fused_flush_quant_body_int1( + hp_v_base, + q_v_base, + sz_v_base, + src, + dst, + active, + head64, + HP_V_STRIDE_LOC, + HP_V_STRIDE_HEAD, + HP_V_STRIDE_DIM, + Q_V_STRIDE_LOC, + Q_V_STRIDE_HEAD, + Q_V_STRIDE_DIM, + V_SZ_STRIDE_LOC, + V_SZ_STRIDE_HEAD, + V_SZ_STRIDE_DIM, + V_HEAD_DIM, + V_BLOCK_OCTANT, + V_NUM_GROUPS, + V_GROUP_SIZE, + BLOCK_TOK, + V_CLIP_INDEX, + V_BSEARCH_ITERS, + LLOYD_MAX, + ) + + +# --------------------------------------------------------------------------- +# Remap kernel — identical to INT2. +# --------------------------------------------------------------------------- + + +@triton.jit +def _flush_remap_kernel_int1( + req_pool_indices_ptr, + flush_pos_ptr, + dst_quant_slot_ptr, + valid_mask_ptr, + req_to_token_ptr, + rtt_stride_row, + FLUSH_INTERVAL: tl.constexpr, +): + i = tl.program_id(0) + j = tl.program_id(1) + out_idx = i * FLUSH_INTERVAL + j + valid = tl.load(valid_mask_ptr + out_idx).to(tl.int32) + if valid == 0: + return + req = tl.load(req_pool_indices_ptr + i).to(tl.int64) + fp = tl.load(flush_pos_ptr + out_idx).to(tl.int64) + dst = tl.load(dst_quant_slot_ptr + out_idx).to(tl.int32) + tl.store(req_to_token_ptr + req * rtt_stride_row + fp, dst) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def _resolve_kv_quant_config_int1( + head_dim: int, num_scale_groups: int +) -> Tuple[int, int, int]: + """Return (BLOCK_OCTANT, NUM_GROUPS, GROUP_SIZE) for the int1 quant kernel.""" + if head_dim % num_scale_groups != 0: + raise ValueError( + f"head_dim ({head_dim}) must be divisible by num_scale_groups " + f"({num_scale_groups})" + ) + if (head_dim & (head_dim - 1)) != 0: + raise ValueError( + f"head_dim ({head_dim}) must be a power of two for the int1 " + f"flush quant kernel" + ) + if head_dim % 8 != 0: + raise ValueError( + f"head_dim ({head_dim}) must be a multiple of 8 for int1 packing" + ) + block_octant = head_dim // 8 + group_size = head_dim // num_scale_groups + return block_octant, num_scale_groups, group_size + + +def _flush_clip_index_int1(clip_ratio: float, head_dim: int) -> int: + if clip_ratio <= 0.0: + return -1 + idx = int(clip_ratio * head_dim) + if idx >= head_dim: + idx = head_dim - 1 + if idx < 0: + idx = 0 + return idx + + +def _flush_elements_per_thread_int1(dtype: torch.dtype) -> int: + if dtype == torch.bfloat16: + return 8 + if dtype.is_floating_point and dtype.itemsize == 1: + return 16 + raise AssertionError( + f"flush quant kernel requires bf16 or fp8 HP dtype, got {dtype}" + ) + + +def _flush_block_tok_and_num_warps_int1( + flush_interval: int, head_dim: int, elements_per_thread: int +) -> Tuple[int, int]: + fi_pow2 = triton.next_power_of_2(max(1, int(flush_interval))) + block_tok = 2 + while block_tok * head_dim < 32 * elements_per_thread: + block_tok *= 2 + if block_tok > fi_pow2: + block_tok = fi_pow2 + total_elems = block_tok * head_dim + vectors_per_warp = 32 * elements_per_thread + num_warps = triton.next_power_of_2( + max(1, triton.cdiv(total_elems, vectors_per_warp)) + ) + return block_tok, num_warps + + +def gpu_flush_int1_plan( + *, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + dst_quant_slots: torch.Tensor, + req_to_token: torch.Tensor, + flush_mask: torch.Tensor, + hp_prefix_tokens: int, + hp_recent_tokens: int, + hp_global_offset: int, + flush_interval: int, +): + bs = int(seq_lens.shape[0]) + if bs == 0 or flush_interval <= 0: + return None + + assert req_to_token.dtype == torch.int32 + assert seq_lens.dtype == torch.int32 + assert prefix_lens.dtype == torch.int32 + assert req_pool_indices.dtype == torch.int64 + assert dst_quant_slots.dtype == torch.int64 + assert dst_quant_slots.numel() == bs * flush_interval + assert flush_mask.shape == (bs,), ( + f"flush_mask shape {tuple(flush_mask.shape)} != ({bs},)" + ) + flush_mask_i8 = flush_mask.to(torch.int8) + + device = seq_lens.device + total_flush_slots = bs * flush_interval + returned_slot_ids = torch.empty( + (total_flush_slots,), dtype=torch.int64, device=device + ) + src_hp_slot = torch.empty((total_flush_slots,), dtype=torch.int64, device=device) + flush_pos = torch.empty((total_flush_slots,), dtype=torch.int32, device=device) + valid_mask = torch.empty((total_flush_slots,), dtype=torch.int8, device=device) + + rtt_stride_row = int(req_to_token.stride(0)) + + _flush_plan_kernel_int1[(bs,)]( + seq_lens, + prefix_lens, + req_pool_indices, + dst_quant_slots, + req_to_token, + flush_mask_i8, + src_hp_slot, + returned_slot_ids, + flush_pos, + valid_mask, + int(req_to_token.shape[1]), + rtt_stride_row, + HP_PREFIX_TOKENS=int(hp_prefix_tokens), + HP_RECENT_TOKENS=int(hp_recent_tokens), + HP_OFFSET=int(hp_global_offset), + FLUSH_INTERVAL=int(flush_interval), + num_warps=1, + num_stages=1, + ) + + return FlushPlanInt1( + returned_slot_ids=returned_slot_ids, + src_hp_slot=src_hp_slot, + flush_pos=flush_pos, + valid_mask=valid_mask, + dst_quant_slots=dst_quant_slots, + bs=bs, + flush_interval=int(flush_interval), + ) + + +def gpu_flush_int1_apply( + plan: FlushPlanInt1, + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + hp_k_ptrs: torch.Tensor, + hp_v_ptrs: torch.Tensor, + quant_k_ptrs: torch.Tensor, + quant_v_ptrs: torch.Tensor, + k_sz_ptrs: torch.Tensor, + v_sz_ptrs: torch.Tensor, + hp_k_sample: torch.Tensor, + hp_v_sample: torch.Tensor, + quant_k_sample: torch.Tensor, + quant_v_sample: torch.Tensor, + k_sz_sample: torch.Tensor, + v_sz_sample: torch.Tensor, + hp_k_strides: Tuple[int, int, int], + hp_v_strides: Tuple[int, int, int], + quant_k_strides: Tuple[int, int, int], + quant_v_strides: Tuple[int, int, int], + k_sz_strides: Tuple[int, int, int], + v_sz_strides: Tuple[int, int, int], + num_heads: int, + head_dim: int, + v_head_dim: int, + k_num_scale_groups: int, + v_num_scale_groups: int, + num_layers: int, + k_clip_ratio: float = 0.0, + v_clip_ratio: float = 0.0, + lloyd_max: bool = False, +) -> None: + bs = plan.bs + flush_interval = plan.flush_interval + total_flush_slots = bs * flush_interval + + safe_src_hp_slot = plan.src_hp_slot.clamp(min=0) + + k_block_octant, k_num_groups, k_group_size = _resolve_kv_quant_config_int1( + head_dim, k_num_scale_groups + ) + v_block_octant, v_num_groups, v_group_size = _resolve_kv_quant_config_int1( + v_head_dim, v_num_scale_groups + ) + + k_clip_index = _flush_clip_index_int1(k_clip_ratio, head_dim) + v_clip_index = _flush_clip_index_int1(v_clip_ratio, v_head_dim) + + elements_per_thread = _flush_elements_per_thread_int1(hp_k_sample.dtype) + block_tok, num_warps = _flush_block_tok_and_num_warps_int1( + flush_interval, head_dim, elements_per_thread + ) + grid = ( + triton.cdiv(total_flush_slots, block_tok), + num_heads, + int(num_layers), + ) + _fused_flush_quant_kernel_int1[grid]( + hp_k_ptrs, + hp_v_ptrs, + quant_k_ptrs, + quant_v_ptrs, + k_sz_ptrs, + v_sz_ptrs, + hp_k_sample, + hp_v_sample, + quant_k_sample, + quant_v_sample, + k_sz_sample, + v_sz_sample, + safe_src_hp_slot, + plan.dst_quant_slots, + plan.valid_mask, + total_flush_slots, + num_heads, + int(num_layers), + HP_K_STRIDE_LOC=hp_k_strides[0], + HP_K_STRIDE_HEAD=hp_k_strides[1], + HP_K_STRIDE_DIM=hp_k_strides[2], + HP_V_STRIDE_LOC=hp_v_strides[0], + HP_V_STRIDE_HEAD=hp_v_strides[1], + HP_V_STRIDE_DIM=hp_v_strides[2], + Q_K_STRIDE_LOC=quant_k_strides[0], + Q_K_STRIDE_HEAD=quant_k_strides[1], + Q_K_STRIDE_DIM=quant_k_strides[2], + Q_V_STRIDE_LOC=quant_v_strides[0], + Q_V_STRIDE_HEAD=quant_v_strides[1], + Q_V_STRIDE_DIM=quant_v_strides[2], + K_SZ_STRIDE_LOC=k_sz_strides[0], + K_SZ_STRIDE_HEAD=k_sz_strides[1], + K_SZ_STRIDE_DIM=k_sz_strides[2], + V_SZ_STRIDE_LOC=v_sz_strides[0], + V_SZ_STRIDE_HEAD=v_sz_strides[1], + V_SZ_STRIDE_DIM=v_sz_strides[2], + K_HEAD_DIM=int(head_dim), + K_BLOCK_OCTANT=k_block_octant, + K_NUM_GROUPS=k_num_groups, + K_GROUP_SIZE=k_group_size, + V_HEAD_DIM=int(v_head_dim), + V_BLOCK_OCTANT=v_block_octant, + V_NUM_GROUPS=v_num_groups, + V_GROUP_SIZE=v_group_size, + BLOCK_TOK=block_tok, + K_CLIP_INDEX=k_clip_index, + V_CLIP_INDEX=v_clip_index, + K_BSEARCH_ITERS=(int(head_dim).bit_length() - 1) if head_dim >= 64 else 0, + V_BSEARCH_ITERS=(int(v_head_dim).bit_length() - 1) if v_head_dim >= 64 else 0, + LLOYD_MAX=bool(lloyd_max), + num_warps=num_warps, + num_stages=1, + ) + + rtt_stride_row = int(req_to_token.stride(0)) + _flush_remap_kernel_int1[(bs, flush_interval)]( + req_pool_indices, + plan.flush_pos, + plan.dst_quant_slots, + plan.valid_mask, + req_to_token, + rtt_stride_row, + FLUSH_INTERVAL=int(flush_interval), + num_warps=1, + num_stages=1, + ) + + +def gpu_flush_int1( + *, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + dst_quant_slots: torch.Tensor, + req_to_token: torch.Tensor, + flush_mask: torch.Tensor, + hp_k_ptrs: torch.Tensor, + hp_v_ptrs: torch.Tensor, + quant_k_ptrs: torch.Tensor, + quant_v_ptrs: torch.Tensor, + k_sz_ptrs: torch.Tensor, + v_sz_ptrs: torch.Tensor, + hp_k_sample: torch.Tensor, + hp_v_sample: torch.Tensor, + quant_k_sample: torch.Tensor, + quant_v_sample: torch.Tensor, + k_sz_sample: torch.Tensor, + v_sz_sample: torch.Tensor, + hp_k_strides: Tuple[int, int, int], + hp_v_strides: Tuple[int, int, int], + quant_k_strides: Tuple[int, int, int], + quant_v_strides: Tuple[int, int, int], + k_sz_strides: Tuple[int, int, int], + v_sz_strides: Tuple[int, int, int], + hp_prefix_tokens: int, + hp_recent_tokens: int, + hp_global_offset: int, + num_heads: int, + head_dim: int, + v_head_dim: int, + k_num_scale_groups: int, + v_num_scale_groups: int, + num_layers: int, + flush_interval: int, + k_clip_ratio: float = 0.0, + v_clip_ratio: float = 0.0, + lloyd_max: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + plan = gpu_flush_int1_plan( + seq_lens=seq_lens, + prefix_lens=prefix_lens, + req_pool_indices=req_pool_indices, + dst_quant_slots=dst_quant_slots, + req_to_token=req_to_token, + flush_mask=flush_mask, + hp_prefix_tokens=hp_prefix_tokens, + hp_recent_tokens=hp_recent_tokens, + hp_global_offset=hp_global_offset, + flush_interval=flush_interval, + ) + if plan is None: + device = seq_lens.device + empty = torch.empty((0,), dtype=torch.int64, device=device) + mask = torch.empty((0,), dtype=torch.int8, device=device) + return empty, mask + + gpu_flush_int1_apply( + plan, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + hp_k_ptrs=hp_k_ptrs, + hp_v_ptrs=hp_v_ptrs, + quant_k_ptrs=quant_k_ptrs, + quant_v_ptrs=quant_v_ptrs, + k_sz_ptrs=k_sz_ptrs, + v_sz_ptrs=v_sz_ptrs, + hp_k_sample=hp_k_sample, + hp_v_sample=hp_v_sample, + quant_k_sample=quant_k_sample, + quant_v_sample=quant_v_sample, + k_sz_sample=k_sz_sample, + v_sz_sample=v_sz_sample, + hp_k_strides=hp_k_strides, + hp_v_strides=hp_v_strides, + quant_k_strides=quant_k_strides, + quant_v_strides=quant_v_strides, + k_sz_strides=k_sz_strides, + v_sz_strides=v_sz_strides, + num_heads=num_heads, + head_dim=head_dim, + v_head_dim=v_head_dim, + k_num_scale_groups=k_num_scale_groups, + v_num_scale_groups=v_num_scale_groups, + num_layers=num_layers, + k_clip_ratio=k_clip_ratio, + v_clip_ratio=v_clip_ratio, + lloyd_max=lloyd_max, + ) + + return plan.returned_slot_ids, plan.valid_mask diff --git a/sglang-research/python/sglang/QuantKernel/gpu_flush_int2.py b/sglang-research/python/sglang/QuantKernel/gpu_flush_int2.py index 7b238966c..f3b3af111 100644 --- a/sglang-research/python/sglang/QuantKernel/gpu_flush_int2.py +++ b/sglang-research/python/sglang/QuantKernel/gpu_flush_int2.py @@ -44,10 +44,10 @@ class FlushPlan: """ returned_slot_ids: torch.Tensor # int64 [bs * flush_interval] - src_hp_slot: torch.Tensor # int64 [bs * flush_interval] - flush_pos: torch.Tensor # int32 [bs * flush_interval] - valid_mask: torch.Tensor # int8 [bs * flush_interval] - dst_quant_slots: torch.Tensor # int64 [bs * flush_interval] (carried through) + src_hp_slot: torch.Tensor # int64 [bs * flush_interval] + flush_pos: torch.Tensor # int32 [bs * flush_interval] + valid_mask: torch.Tensor # int8 [bs * flush_interval] + dst_quant_slots: torch.Tensor # int64 [bs * flush_interval] (carried through) bs: int flush_interval: int @@ -59,16 +59,16 @@ class FlushPlan: @triton.jit def _flush_plan_kernel( - seq_lens_ptr, # int32 [bs] - prefix_lens_ptr, # int32 [bs] - req_pool_indices_ptr, # int64 [bs] - dst_quant_slot_ptr, # int64 [bs * FLUSH_INTERVAL] pre-allocated - req_to_token_ptr, # int32 [num_req_slots, max_ctx] - flush_mask_ptr, # int8 [bs] -- 1 if request flushes this step - src_hp_slot_out_ptr, # int64 [bs, FLUSH_INTERVAL] src hp slot or -1 - returned_slot_ids_ptr, # int64 [bs, FLUSH_INTERVAL] hp slot or dst_quant_slot - flush_pos_out_ptr, # int32 [bs, FLUSH_INTERVAL] flush_pos or -1 - valid_mask_out_ptr, # int8 [bs, FLUSH_INTERVAL] 1 if flushed else 0 + seq_lens_ptr, # int32 [bs] + prefix_lens_ptr, # int32 [bs] + req_pool_indices_ptr, # int64 [bs] + dst_quant_slot_ptr, # int64 [bs * FLUSH_INTERVAL] pre-allocated + req_to_token_ptr, # int32 [num_req_slots, max_ctx] + flush_mask_ptr, # int8 [bs] -- 1 if request flushes this step + src_hp_slot_out_ptr, # int64 [bs, FLUSH_INTERVAL] src hp slot or -1 + returned_slot_ids_ptr, # int64 [bs, FLUSH_INTERVAL] hp slot or dst_quant_slot + flush_pos_out_ptr, # int32 [bs, FLUSH_INTERVAL] flush_pos or -1 + valid_mask_out_ptr, # int8 [bs, FLUSH_INTERVAL] 1 if flushed else 0 max_ctx, rtt_stride_row, HP_PREFIX_TOKENS: tl.constexpr, @@ -146,13 +146,13 @@ def _flush_plan_kernel( @triton.jit def _fused_flush_quant_body( - hp_base, # pointer to hp_dtype arena for one (layer, K|V) - quant_base, # pointer to uint8 arena (int2 packed view) - sz_base, # pointer to scale_dtype arena - src_hp_slot, # int64 [BLOCK_TOK] - dst_quant_slot, # int64 [BLOCK_TOK] - active, # int1 [BLOCK_TOK] per-row valid mask - head_idx, # int64 scalar + hp_base, # pointer to hp_dtype arena for one (layer, K|V) + quant_base, # pointer to uint8 arena (int2 packed view) + sz_base, # pointer to scale_dtype arena + src_hp_slot, # int64 [BLOCK_TOK] + dst_quant_slot, # int64 [BLOCK_TOK] + active, # int1 [BLOCK_TOK] per-row valid mask + head_idx, # int64 scalar HP_STRIDE_LOC: tl.constexpr, HP_STRIDE_HEAD: tl.constexpr, HP_STRIDE_DIM: tl.constexpr, @@ -209,9 +209,7 @@ def _fused_flush_quant_body( thr_hi = tl.max(abs_acc, axis=1) for _ in tl.static_range(BSEARCH_ITERS): thr_mid = (thr_lo + thr_hi) * 0.5 - cnt_above = tl.sum( - (abs_acc > thr_mid[:, None]).to(tl.int32), axis=1 - ) + cnt_above = tl.sum((abs_acc > thr_mid[:, None]).to(tl.int32), axis=1) too_many = cnt_above > target_above thr_lo = tl.where(too_many, thr_mid, thr_lo) thr_hi = tl.where(too_many, thr_hi, thr_mid) @@ -249,12 +247,8 @@ def _fused_flush_quant_body( vals0, vals2 = tl.split(a_even) vals1, vals3 = tl.split(a_odd) - scale_3d = tl.broadcast_to( - scale[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE) - ) - zero_3d = tl.broadcast_to( - zero[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE) - ) + scale_3d = tl.broadcast_to(scale[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) + zero_3d = tl.broadcast_to(zero[:, :, None], (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) scale_flat = tl.reshape(scale_3d, (BLOCK_TOK, HEAD_DIM)) zero_flat = tl.reshape(zero_3d, (BLOCK_TOK, HEAD_DIM)) @@ -287,9 +281,7 @@ def _fused_flush_quant_body( tl.store(quant_base + cache_offset, packed, mask=active[:, None]) group_ids = tl.arange(0, NUM_GROUPS) - sz_offset_base = ( - dst_quant_slot[:, None] * SZ_STRIDE_LOC + head_idx * SZ_STRIDE_HEAD - ) + sz_offset_base = dst_quant_slot[:, None] * SZ_STRIDE_LOC + head_idx * SZ_STRIDE_HEAD tl.store( sz_base + sz_offset_base + (group_ids[None, :] * 2) * SZ_STRIDE_DIM, scale, @@ -321,9 +313,9 @@ def _fused_flush_quant_kernel( k_sz_sample_ptr, v_sz_sample_ptr, # Flush plan (flat view of [bs, FLUSH_INTERVAL]) - src_hp_slot_ptr, # int64 [num_flush_tokens] clamped >= 0 - dst_quant_slot_ptr, # int64 [num_flush_tokens] - valid_mask_ptr, # int8 [num_flush_tokens] + src_hp_slot_ptr, # int64 [num_flush_tokens] clamped >= 0 + dst_quant_slot_ptr, # int64 [num_flush_tokens] + valid_mask_ptr, # int8 [num_flush_tokens] num_flush_tokens, num_heads, num_layers, @@ -480,11 +472,11 @@ def _fused_flush_quant_kernel( @triton.jit def _flush_remap_kernel( - req_pool_indices_ptr, # int64 [bs] - flush_pos_ptr, # int32 [bs * FLUSH_INTERVAL] - dst_quant_slot_ptr, # int64 [bs * FLUSH_INTERVAL] - valid_mask_ptr, # int8 [bs * FLUSH_INTERVAL] - req_to_token_ptr, # int32 [num_req_slots, max_ctx] + req_pool_indices_ptr, # int64 [bs] + flush_pos_ptr, # int32 [bs * FLUSH_INTERVAL] + dst_quant_slot_ptr, # int64 [bs * FLUSH_INTERVAL] + valid_mask_ptr, # int8 [bs * FLUSH_INTERVAL] + req_to_token_ptr, # int32 [num_req_slots, max_ctx] rtt_stride_row, FLUSH_INTERVAL: tl.constexpr, ): @@ -574,22 +566,25 @@ def _flush_block_tok_and_num_warps( if block_tok > fi_pow2: block_tok = fi_pow2 total_elems = block_tok * head_dim - assert total_elems % (32 * elements_per_thread) == 0, ( - f"BLOCK_TOK={block_tok} head_dim={head_dim} " - f"epp={elements_per_thread}: tile doesn't divide cleanly into " - "128-bit/thread loads" + # Small test/debug intervals (notably K=1 with head_dim=64) contain fewer + # than one warp's worth of 128-bit vectors. They are still valid Triton + # tiles; launch one warp and let inactive lanes idle instead of rejecting + # the configuration. Production shapes remain exactly vector-balanced. + vectors_per_warp = 32 * elements_per_thread + num_warps = triton.next_power_of_2( + max(1, triton.cdiv(total_elems, vectors_per_warp)) ) - num_warps = total_elems // (32 * elements_per_thread) return block_tok, num_warps + def gpu_flush_int2_plan( *, - seq_lens: torch.Tensor, # int32 [bs] - prefix_lens: torch.Tensor, # int32 [bs] - req_pool_indices: torch.Tensor, # int64 [bs] - dst_quant_slots: torch.Tensor, # int64 [bs * flush_interval] - req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] - flush_mask: torch.Tensor, # bool [bs] -- per-request gate + seq_lens: torch.Tensor, # int32 [bs] + prefix_lens: torch.Tensor, # int32 [bs] + req_pool_indices: torch.Tensor, # int64 [bs] + dst_quant_slots: torch.Tensor, # int64 [bs * flush_interval] + req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] + flush_mask: torch.Tensor, # bool [bs] -- per-request gate hp_prefix_tokens: int, hp_recent_tokens: int, hp_global_offset: int, @@ -623,7 +618,9 @@ def gpu_flush_int2_plan( device = seq_lens.device total_flush_slots = bs * flush_interval - returned_slot_ids = torch.empty((total_flush_slots,), dtype=torch.int64, device=device) + returned_slot_ids = torch.empty( + (total_flush_slots,), dtype=torch.int64, device=device + ) src_hp_slot = torch.empty((total_flush_slots,), dtype=torch.int64, device=device) flush_pos = torch.empty((total_flush_slots,), dtype=torch.int32, device=device) valid_mask = torch.empty((total_flush_slots,), dtype=torch.int8, device=device) @@ -665,8 +662,8 @@ def gpu_flush_int2_plan( def gpu_flush_int2_apply( plan: FlushPlan, *, - req_pool_indices: torch.Tensor, # int64 [bs] - req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] + req_pool_indices: torch.Tensor, # int64 [bs] + req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] # Pool-held metadata (built once at pool construction): hp_k_ptrs: torch.Tensor, hp_v_ptrs: torch.Tensor, @@ -800,19 +797,19 @@ def gpu_flush_int2_apply( def gpu_flush_int2( *, - seq_lens: torch.Tensor, # int32 [bs] - prefix_lens: torch.Tensor, # int32 [bs] - req_pool_indices: torch.Tensor, # int64 [bs] - dst_quant_slots: torch.Tensor, # int64 [bs * flush_interval] - req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] - flush_mask: torch.Tensor, # bool [bs] -- per-request gate + seq_lens: torch.Tensor, # int32 [bs] + prefix_lens: torch.Tensor, # int32 [bs] + req_pool_indices: torch.Tensor, # int64 [bs] + dst_quant_slots: torch.Tensor, # int64 [bs * flush_interval] + req_to_token: torch.Tensor, # int32 [num_req_slots, max_ctx] + flush_mask: torch.Tensor, # bool [bs] -- per-request gate # Pool-held metadata (built once at pool construction): - hp_k_ptrs: torch.Tensor, # int64 [num_layers] - hp_v_ptrs: torch.Tensor, # int64 [num_layers] - quant_k_ptrs: torch.Tensor, # int64 [num_layers] - quant_v_ptrs: torch.Tensor, # int64 [num_layers] - k_sz_ptrs: torch.Tensor, # int64 [num_layers] - v_sz_ptrs: torch.Tensor, # int64 [num_layers] + hp_k_ptrs: torch.Tensor, # int64 [num_layers] + hp_v_ptrs: torch.Tensor, # int64 [num_layers] + quant_k_ptrs: torch.Tensor, # int64 [num_layers] + quant_v_ptrs: torch.Tensor, # int64 [num_layers] + k_sz_ptrs: torch.Tensor, # int64 [num_layers] + v_sz_ptrs: torch.Tensor, # int64 [num_layers] hp_k_sample: torch.Tensor, hp_v_sample: torch.Tensor, quant_k_sample: torch.Tensor, @@ -904,3 +901,122 @@ def gpu_flush_int2( ) return plan.returned_slot_ids, plan.valid_mask + + +def gpu_flush_pqk_int2v_apply( + plan: "FlushPlan", + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + kv_pool, # UnifiedInt2HPKVPool with pq_k_int2v dtype +) -> None: + """Apply phase for PQ K + INT2 LM V flush. + + K: gather HP K, PQ-encode → k_buffer codes. + V: gather HP V, INT2 LM-encode → v_buffer + v_scales_zeros. + Remap: same _flush_remap_kernel as INT2. + """ + from sglang.QuantKernel.oscar_rotation_clip_int2_kv import ( + _launch_grouped_clip_int2, + _launch_single_clip_int2, + ) + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_encode_k + from sglang.srt.mem_cache.kv_quant_kernels import _get_num_scale_groups + + bs = plan.bs + flush_interval = plan.flush_interval + + safe_src = plan.src_hp_slot.clamp(min=0) # [bs * flush_interval] + valid = plan.valid_mask.bool() # [bs * flush_interval] + dst = plan.dst_quant_slots # [bs * flush_interval] int32 + + valid_dst = dst[valid] # [n_valid] + + for layer_idx in range(kv_pool.layer_num): + layer_id = kv_pool.start_layer + layer_idx + # --- K: PQ encode --- + hp_k = kv_pool.hp_k_buffer[layer_idx] # [hp_slots, heads, head_dim] + gathered_k = hp_k[safe_src][valid] # [n_valid, heads, head_dim] + k_fp16 = ( + gathered_k + if gathered_k.dtype == torch.float16 + else gathered_k.to(torch.float16) + ) + pq_encode_k( + k_fp16, + valid_dst, + kv_pool.k_buffer[layer_idx], + kv_pool.get_pq_codebook(layer_id), + kv_pool.get_pq_cb_norms(layer_id), + ) + # RVQ stage-2: encode the residual (k - PQ(k)) into k_buffer2 (reuses pq kernels). + if getattr(kv_pool, "_is_rvq", False): + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_decode_k + + n_valid = k_fp16.shape[0] + head_dim_k = k_fp16.shape[-1] + recon1 = pq_decode_k( + kv_pool.k_buffer[layer_idx][valid_dst], + kv_pool.get_pq_codebook(layer_id), + n_valid, + head_dim_k, + ).to(k_fp16.dtype) + resid = (k_fp16 - recon1).contiguous() + pq_encode_k( + resid, + valid_dst, + kv_pool.k_buffer2[layer_idx], + kv_pool.get_rvq_cb2(layer_id), + kv_pool.get_rvq_cb2_norms(layer_id), + ) + + # --- V: PQ encode (if SGLANG_PQ_V_CODEBOOK set) else INT2 LM encode --- + hp_v = kv_pool.hp_v_buffer[layer_idx] # [hp_slots, heads, v_head_dim] + gathered_v = hp_v[safe_src][valid] # [n_valid, heads, v_head_dim] + v_float = gathered_v.to(kv_pool.hp_dtype) + if getattr(kv_pool, "_pq_v", False): + vcb = kv_pool.get_pq_v_codebook(layer_id) + n_sub_v = vcb.shape[0] + pq_encode_k( + v_float.to(torch.float16), + valid_dst, + kv_pool.v_buffer[layer_idx][ + ..., :n_sub_v + ], # store codes in first n_sub bytes + vcb, + kv_pool.get_pq_v_cb_norms(layer_id), + ) + continue + n_groups_v = _get_num_scale_groups(kv_pool.v_scales_zeros[layer_idx]) + if n_groups_v == 1: + _launch_single_clip_int2( + v_float, + valid_dst, + kv_pool.v_buffer[layer_idx], + kv_pool.v_scales_zeros[layer_idx], + kv_pool._v_clip_ratio, + hp_global_offset=None, + lloyd_max=kv_pool._lloyd_max, + ) + else: + _launch_grouped_clip_int2( + v_float, + valid_dst, + kv_pool.v_buffer[layer_idx], + kv_pool.v_scales_zeros[layer_idx], + kv_pool._v_clip_ratio, + hp_global_offset=None, + ) + + rtt_stride_row = int(req_to_token.stride(0)) + _flush_remap_kernel[(bs, flush_interval)]( + req_pool_indices, + plan.flush_pos, + plan.dst_quant_slots, + plan.valid_mask, + req_to_token, + rtt_stride_row, + FLUSH_INTERVAL=int(flush_interval), + num_warps=1, + num_stages=1, + ) diff --git a/sglang-research/python/sglang/QuantKernel/oscar_rotation_clip_int1_kv.py b/sglang-research/python/sglang/QuantKernel/oscar_rotation_clip_int1_kv.py new file mode 100644 index 000000000..bd068f900 --- /dev/null +++ b/sglang-research/python/sglang/QuantKernel/oscar_rotation_clip_int1_kv.py @@ -0,0 +1,980 @@ +""" +Oscar-style per-row quantile clip + int1 KV cache pack kernels. + +INT1 mirrors :mod:`sglang.QuantKernel.oscar_rotation_clip_int2_kv` but with +``max_q == 1`` (two quant levels {0, 1}) and an 8-way byte pack. Storage is +``head_dim // 8`` packed uint8 bytes (vs ``head_dim // 4`` for INT2). The +scale/zero layout is identical to INT2 (interleaved ``(scale, zero)`` pairs +per group), so the dequantize formula ``(q - zero) * scale`` is unchanged +except that ``q`` now comes from a 1-bit slot. + +The octanted split inside the kernel uses three rounds of ``tl.split``: +``(BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)`` → 2 → 4 → 8 leaves of shape +``(BLOCK_TOK, BLOCK_OCTANT)``. Mirrors the INT2 "quartered split via +``tl.split×2``" idiom one level deeper. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.mem_cache.kv_quant_kernels import ( + _get_num_scale_groups, + _is_power_of_two, +) + + +# --------------------------------------------------------------------------- +# Fused threshold + clip + int1 pack kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _pretransformed_int1_set_kv_clip_single_kernel( + input_ptr, + loc_ptr, + cache_ptr, + scales_zeros_ptr, + num_tokens, + num_heads, + input_stride_token, + input_stride_head, + input_stride_dim, + cache_stride_loc, + cache_stride_head, + cache_stride_dim, + sz_stride_loc, + sz_stride_head, + sz_stride_dim, + HP_OFFSET: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_OCTANT: tl.constexpr, + BLOCK_TOK: tl.constexpr, + CLIP_INDEX: tl.constexpr, + LLOYD_MAX: tl.constexpr, +): + """Multi-row fused threshold + single-scale clip + int1 pack. + + Mirrors the INT2 single-scale kernel but packs 8 quant slots per byte. + ``BLOCK_OCTANT == HEAD_DIM // 8``. + """ + pid_tok = tl.program_id(0) + head_idx = tl.program_id(1) + if head_idx >= num_heads: + return + + tok_offs = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + tok_mask = tok_offs < num_tokens + + cache_loc = tl.load(loc_ptr + tok_offs, mask=tok_mask, other=0) + if HP_OFFSET >= 0: + active = tok_mask & (cache_loc < HP_OFFSET) + else: + active = tok_mask + + full_offs = tl.arange(0, HEAD_DIM) + base = ( + tok_offs[:, None] * input_stride_token + + head_idx * input_stride_head + + full_offs[None, :] * input_stride_dim + ) + rows = tl.load( + input_ptr + base, + mask=tok_mask[:, None], + other=0.0, + ).to(tl.float32) + + if CLIP_INDEX >= 0: + abs_rows = tl.abs(rows) + sorted_rows = tl.sort(abs_rows) + pick = (full_offs == CLIP_INDEX)[None, :] + thr = tl.sum(tl.where(pick, sorted_rows, 0.0), axis=1) # [BLOCK_TOK] + rows = tl.minimum( + tl.maximum(rows, -thr[:, None]), + thr[:, None], + ) + + if LLOYD_MAX: + LM_C_1BIT: tl.constexpr = 0.79788456 # sqrt(2/pi) + mean = tl.sum(rows, axis=1) / HEAD_DIM + diff = rows - mean[:, None] + var = tl.sum(diff * diff, axis=1) / HEAD_DIM + std = tl.sqrt(var + 1e-8) + scale = tl.maximum(2.0 * LM_C_1BIT * std, 1e-8) + zero = 0.5 - mean / scale + else: + val_min = tl.min(rows, axis=1) + val_max = tl.max(rows, axis=1) + val_range = tl.maximum(val_max - val_min, 1e-8) + scale = val_range # max_q = 1 → scale = range / 1 + zero = -val_min / scale + + # Octant split: reshape (BLOCK_TOK, 8*BLOCK_OCTANT) → + # (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2) via reshape + permute, then split×3. + rows_r = tl.reshape(rows, (BLOCK_TOK, 8, BLOCK_OCTANT)) + rows_p = tl.permute(rows_r, (0, 2, 1)) + rows_s = tl.reshape(rows_p, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + h_lo, h_hi = tl.split(rows_s) # each [BLOCK_TOK, BLOCK_OCTANT, 2, 2] + e0, e2 = tl.split(h_lo) # each [BLOCK_TOK, BLOCK_OCTANT, 2] + e1, e3 = tl.split(h_hi) + vals0, vals4 = tl.split(e0) # each [BLOCK_TOK, BLOCK_OCTANT] + vals2, vals6 = tl.split(e2) + vals1, vals5 = tl.split(e1) + vals3, vals7 = tl.split(e3) + + # Lloyd-Max centroids do not bound tail values, so the rounded expression + # can fall outside {0, 1}. Clamp before the uint8 cast; otherwise 2 or a + # wrapped negative value spills into neighboring packed bits. This also + # makes the decision boundary exactly ``value >= mean``, matching the + # grouped prefill and decode-flush kernels. + q0 = tl.minimum( + tl.maximum(vals0 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q1 = tl.minimum( + tl.maximum(vals1 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q2 = tl.minimum( + tl.maximum(vals2 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q3 = tl.minimum( + tl.maximum(vals3 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q4 = tl.minimum( + tl.maximum(vals4 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q5 = tl.minimum( + tl.maximum(vals5 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q6 = tl.minimum( + tl.maximum(vals6 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + q7 = tl.minimum( + tl.maximum(vals7 / scale[:, None] + zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + + packed = ( + q0 + | (q1 << 1) + | (q2 << 2) + | (q3 << 3) + | (q4 << 4) + | (q5 << 5) + | (q6 << 6) + | (q7 << 7) + ) + + dim_offs_o = tl.arange(0, BLOCK_OCTANT) + cache_offset = ( + cache_loc[:, None] * cache_stride_loc + + head_idx * cache_stride_head + + dim_offs_o[None, :] * cache_stride_dim + ) + tl.store(cache_ptr + cache_offset, packed, mask=active[:, None]) + + sz_offset_base = cache_loc * sz_stride_loc + head_idx * sz_stride_head + tl.store( + scales_zeros_ptr + sz_offset_base + 0 * sz_stride_dim, + scale, + mask=active, + ) + tl.store( + scales_zeros_ptr + sz_offset_base + 1 * sz_stride_dim, + zero, + mask=active, + ) + + +@triton.jit +def _pretransformed_int1_set_kv_clip_grouped_kernel( + input_ptr, + loc_ptr, + cache_ptr, + scales_zeros_ptr, + num_tokens, + num_heads, + input_stride_token, + input_stride_head, + input_stride_dim, + cache_stride_loc, + cache_stride_head, + cache_stride_dim, + sz_stride_loc, + sz_stride_head, + sz_stride_dim, + HEAD_DIM: tl.constexpr, + BLOCK_OCTANT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + GROUP_SIZE: tl.constexpr, + HP_OFFSET: tl.constexpr, + BLOCK_TOK: tl.constexpr, + CLIP_INDEX: tl.constexpr, + LLOYD_MAX: tl.constexpr, +): + """Multi-row fused threshold + groupwise clip + int1 pack (8 per byte). + + When LLOYD_MAX=True: per-group LM binary quantizer (boundary=group mean, + centroid=group_mean ± group_std*sqrt(2/pi)). + """ + pid_tok = tl.program_id(0) + head_idx = tl.program_id(1) + if head_idx >= num_heads: + return + + tok_offs = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + tok_mask = tok_offs < num_tokens + + cache_loc = tl.load(loc_ptr + tok_offs, mask=tok_mask, other=0) + if HP_OFFSET >= 0: + active = tok_mask & (cache_loc < HP_OFFSET) + else: + active = tok_mask + + full_offs = tl.arange(0, HEAD_DIM) + base = ( + tok_offs[:, None] * input_stride_token + + head_idx * input_stride_head + + full_offs[None, :] * input_stride_dim + ) + acc = tl.load( + input_ptr + base, + mask=tok_mask[:, None], + other=0.0, + ).to(tl.float32) + + if CLIP_INDEX >= 0: + abs_acc = tl.abs(acc) + sorted_acc = tl.sort(abs_acc) + pick = (full_offs == CLIP_INDEX)[None, :] + thr = tl.sum(tl.where(pick, sorted_acc, 0.0), axis=1) + acc = tl.minimum( + tl.maximum(acc, -thr[:, None]), + thr[:, None], + ) + + grouped = tl.reshape(acc, (BLOCK_TOK, NUM_GROUPS, GROUP_SIZE)) + if LLOYD_MAX: + LM_C_1BIT: tl.constexpr = 0.79788456 # sqrt(2/pi) + group_mean = tl.sum(grouped, axis=2) / GROUP_SIZE + group_diff = grouped - group_mean[:, :, None] + group_var = tl.sum(group_diff * group_diff, axis=2) / GROUP_SIZE + group_std = tl.sqrt(group_var + 1e-8) + scale = tl.maximum(2.0 * LM_C_1BIT * group_std, 1e-8) + zero = 0.5 - group_mean / scale + quant = (group_diff >= 0.0).to(tl.uint8) + else: + val_min = tl.min(grouped, axis=2) + val_max = tl.max(grouped, axis=2) + scale = tl.maximum(val_max - val_min, 1e-8) + zero = tl.math.div_rn(-val_min, scale) + quant = ( + tl.math.div_rn(grouped, scale[:, :, None]) + zero[:, :, None] + 0.5 + ).to(tl.uint8) + quant_flat = tl.reshape(quant, (BLOCK_TOK, HEAD_DIM)) + quant_r = tl.reshape(quant_flat, (BLOCK_TOK, 8, BLOCK_OCTANT)) + quant_p = tl.permute(quant_r, (0, 2, 1)) + quant_s = tl.reshape(quant_p, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + h_lo, h_hi = tl.split(quant_s) + e0, e2 = tl.split(h_lo) + e1, e3 = tl.split(h_hi) + q0, q4 = tl.split(e0) + q2, q6 = tl.split(e2) + q1, q5 = tl.split(e1) + q3, q7 = tl.split(e3) + + packed = ( + q0 + | (q1 << 1) + | (q2 << 2) + | (q3 << 3) + | (q4 << 4) + | (q5 << 5) + | (q6 << 6) + | (q7 << 7) + ) + + dim_offs_o = tl.arange(0, BLOCK_OCTANT) + cache_offset = ( + cache_loc[:, None] * cache_stride_loc + + head_idx * cache_stride_head + + dim_offs_o[None, :] * cache_stride_dim + ) + tl.store(cache_ptr + cache_offset, packed, mask=active[:, None]) + + group_ids = tl.arange(0, NUM_GROUPS) + sz_offset_base = cache_loc[:, None] * sz_stride_loc + head_idx * sz_stride_head + tl.store( + scales_zeros_ptr + sz_offset_base + (group_ids[None, :] * 2) * sz_stride_dim, + scale, + mask=active[:, None], + ) + tl.store( + scales_zeros_ptr + + sz_offset_base + + (group_ids[None, :] * 2 + 1) * sz_stride_dim, + zero, + mask=active[:, None], + ) + + +def _can_use_grouped_clip_kernel( + head_dim: int, scales_zeros_buffer: torch.Tensor +) -> bool: + num_groups = _get_num_scale_groups(scales_zeros_buffer) + if num_groups == 1: + return True + if head_dim % num_groups != 0: + return False + group_size = head_dim // num_groups + return _is_power_of_two(num_groups) and _is_power_of_two(group_size) + + +def _clip_index(clip_ratio: float, head_dim: int) -> int: + if clip_ratio <= 0.0: + return -1 + idx = int(clip_ratio * head_dim) + if idx >= head_dim: + idx = head_dim - 1 + if idx < 0: + idx = 0 + return idx + + +def _vectorized_elems_per_thread(dtype: torch.dtype) -> int: + if dtype == torch.bfloat16: + return 8 + if dtype.is_floating_point and dtype.itemsize == 1: + return 16 + raise AssertionError( + f"clip int1 kernel requires bf16 or fp8 input dtype, got {dtype}" + ) + + +def _pick_block_tok_and_num_warps( + head_dim: int, elements_per_thread: int +) -> Tuple[int, int]: + block_tok = 4 + while block_tok * head_dim < 32 * elements_per_thread: + block_tok *= 2 + total_elems = block_tok * head_dim + assert total_elems % (32 * elements_per_thread) == 0, ( + f"BLOCK_TOK={block_tok} head_dim={head_dim} epp={elements_per_thread}: " + "tile size doesn't divide cleanly into 128-bit/thread loads" + ) + num_warps = total_elems // (32 * elements_per_thread) + return block_tok, num_warps + + +def _launch_single_clip_int1( + data: torch.Tensor, + loc: torch.Tensor, + buf: torch.Tensor, + sz_buf: torch.Tensor, + clip_ratio: float, + hp_global_offset=None, + lloyd_max: bool = False, +) -> None: + num_tokens, num_heads, head_dim = data.shape + if num_tokens == 0: + return + assert _is_power_of_two(head_dim), ( + f"clip int1 kernel requires power-of-two head_dim, got {head_dim}" + ) + assert head_dim % 8 == 0, ( + f"head_dim must be divisible by 8 for INT1, got {head_dim}" + ) + elements_per_thread = _vectorized_elems_per_thread(data.dtype) + block_tok, num_warps = _pick_block_tok_and_num_warps(head_dim, elements_per_thread) + grid = (triton.cdiv(num_tokens, block_tok), num_heads) + _pretransformed_int1_set_kv_clip_single_kernel[grid]( + data, + loc, + buf, + sz_buf, + num_tokens, + num_heads, + data.stride(0), + data.stride(1), + data.stride(2), + buf.stride(0), + buf.stride(1), + buf.stride(2), + sz_buf.stride(0), + sz_buf.stride(1), + sz_buf.stride(2), + HP_OFFSET=-1 if hp_global_offset is None else int(hp_global_offset), + HEAD_DIM=head_dim, + BLOCK_OCTANT=head_dim // 8, + BLOCK_TOK=block_tok, + CLIP_INDEX=_clip_index(clip_ratio, head_dim), + LLOYD_MAX=lloyd_max, + num_warps=num_warps, + num_stages=1, + ) + + +def _launch_grouped_clip_int1( + data: torch.Tensor, + loc: torch.Tensor, + buf: torch.Tensor, + sz_buf: torch.Tensor, + clip_ratio: float, + hp_global_offset=None, + lloyd_max: bool = False, +) -> None: + num_tokens, num_heads, head_dim = data.shape + if num_tokens == 0: + return + num_groups = _get_num_scale_groups(sz_buf) + + group_size = head_dim // num_groups + block_octant = triton.next_power_of_2(head_dim // 8) + elements_per_thread = _vectorized_elems_per_thread(data.dtype) + + block_tok, num_warps = _pick_block_tok_and_num_warps(head_dim, elements_per_thread) + + assert _is_power_of_two(head_dim), ( + f"clip int1 kernel requires power-of-two head_dim, got {head_dim}" + ) + assert head_dim % 8 == 0, ( + f"head_dim must be divisible by 8 for INT1, got {head_dim}" + ) + assert _is_power_of_two(num_groups) and head_dim % num_groups == 0 + + grid = (triton.cdiv(num_tokens, block_tok), num_heads) + _pretransformed_int1_set_kv_clip_grouped_kernel[grid]( + data, + loc, + buf, + sz_buf, + num_tokens, + num_heads, + data.stride(0), + data.stride(1), + data.stride(2), + buf.stride(0), + buf.stride(1), + buf.stride(2), + sz_buf.stride(0), + sz_buf.stride(1), + sz_buf.stride(2), + HEAD_DIM=head_dim, + BLOCK_OCTANT=block_octant, + NUM_GROUPS=num_groups, + GROUP_SIZE=group_size, + HP_OFFSET=-1 if hp_global_offset is None else int(hp_global_offset), + BLOCK_TOK=block_tok, + CLIP_INDEX=_clip_index(clip_ratio, head_dim), + LLOYD_MAX=lloyd_max, + num_warps=num_warps, + num_stages=1, + ) + + +def quantized_set_kv_int1_pretransformed_clip_triton( + cache_k: torch.Tensor, + cache_v: torch.Tensor, + loc: torch.Tensor, + k_cache_buffer: torch.Tensor, + v_cache_buffer: torch.Tensor, + k_scales_zeros_buffer: torch.Tensor, + v_scales_zeros_buffer: torch.Tensor, + clip_ratio_k: float, + clip_ratio_v: float, + hp_global_offset=None, + lloyd_max: bool = False, +) -> None: + """Fused threshold + clip + quantize + int1-pack for already-rotated K/V.""" + assert cache_k.shape[:2] == cache_v.shape[:2], ( + f"K/V shape mismatch in pretransformed_clip: {cache_k.shape} vs {cache_v.shape}" + ) + num_tokens, _num_heads, k_head_dim = cache_k.shape + v_head_dim = cache_v.shape[-1] + assert k_head_dim % 8 == 0 and v_head_dim % 8 == 0, ( + "K/V head dims must be divisible by 8 for INT1, got " + f"K={k_head_dim}, V={v_head_dim}" + ) + + if num_tokens == 0: + return + + k_grouped_ok = _can_use_grouped_clip_kernel(k_head_dim, k_scales_zeros_buffer) + v_grouped_ok = _can_use_grouped_clip_kernel(v_head_dim, v_scales_zeros_buffer) + if not (k_grouped_ok and v_grouped_ok): + raise NotImplementedError( + f"pretransformed_clip int1 kernel requires power-of-two group configs " + f"(k_head_dim={k_head_dim}, v_head_dim={v_head_dim}, " + f"k_num_groups={_get_num_scale_groups(k_scales_zeros_buffer)}, " + f"v_num_groups={_get_num_scale_groups(v_scales_zeros_buffer)})" + ) + + if _get_num_scale_groups(k_scales_zeros_buffer) == 1: + _launch_single_clip_int1( + cache_k, + loc, + k_cache_buffer, + k_scales_zeros_buffer, + clip_ratio_k, + hp_global_offset, + lloyd_max=lloyd_max, + ) + else: + _launch_grouped_clip_int1( + cache_k, + loc, + k_cache_buffer, + k_scales_zeros_buffer, + clip_ratio_k, + hp_global_offset, + lloyd_max=lloyd_max, + ) + + if _get_num_scale_groups(v_scales_zeros_buffer) == 1: + _launch_single_clip_int1( + cache_v, + loc, + v_cache_buffer, + v_scales_zeros_buffer, + clip_ratio_v, + hp_global_offset, + lloyd_max=lloyd_max, + ) + else: + _launch_grouped_clip_int1( + cache_v, + loc, + v_cache_buffer, + v_scales_zeros_buffer, + clip_ratio_v, + hp_global_offset, + lloyd_max=lloyd_max, + ) + + +# --------------------------------------------------------------------------- +# Fused rotate (K) + clip + quantize + int1 pack kernel for K and V +# --------------------------------------------------------------------------- + + +@triton.jit +def _kv_oscar_rotate_k_clip_single_kernel_int1( + k_input_ptr, + v_input_ptr, + R_ptr, + loc_ptr, + k_cache_ptr, + v_cache_ptr, + k_sz_ptr, + v_sz_ptr, + num_tokens, + num_heads, + k_input_stride_token, + k_input_stride_head, + k_input_stride_dim, + v_input_stride_token, + v_input_stride_head, + v_input_stride_dim, + R_stride_in, + R_stride_out, + k_cache_stride_loc, + k_cache_stride_head, + k_cache_stride_dim, + v_cache_stride_loc, + v_cache_stride_head, + v_cache_stride_dim, + k_sz_stride_loc, + k_sz_stride_head, + k_sz_stride_dim, + v_sz_stride_loc, + v_sz_stride_head, + v_sz_stride_dim, + HP_OFFSET: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_OCTANT: tl.constexpr, + BLOCK_TOK: tl.constexpr, + K_CLIP_INDEX: tl.constexpr, + V_CLIP_INDEX: tl.constexpr, + BSEARCH_ITERS: tl.constexpr, + LLOYD_MAX: tl.constexpr, +): + """Single fused kernel: rotate(K) + clip(K) + quant(K) + pack(K), then + clip(V) + quant(V) + pack(V) sharing the same token tile. INT1 variant. + """ + pid_tok = tl.program_id(0) + head_idx = tl.program_id(1) + if head_idx >= num_heads: + return + + tok_offs = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + tok_mask = tok_offs < num_tokens + + cache_loc = tl.load(loc_ptr + tok_offs, mask=tok_mask, other=0) + if HP_OFFSET >= 0: + active = tok_mask & (cache_loc < HP_OFFSET) + else: + active = tok_mask + + full_offs = tl.arange(0, HEAD_DIM) + dim_offs_o = tl.arange(0, BLOCK_OCTANT) + + # ---------- K: load + rotate + clip + scale/zero + pack + write ---------- + k_base = ( + tok_offs[:, None] * k_input_stride_token + + head_idx * k_input_stride_head + + full_offs[None, :] * k_input_stride_dim + ) + k_tile = tl.load( + k_input_ptr + k_base, + mask=tok_mask[:, None], + other=0.0, + ) + + r_in = tl.arange(0, HEAD_DIM) + r_out = tl.arange(0, HEAD_DIM) + R_offs = r_in[:, None] * R_stride_in + r_out[None, :] * R_stride_out + R_tile = tl.load(R_ptr + R_offs) + + k_rows = tl.dot(k_tile, R_tile, out_dtype=tl.float32) + + if K_CLIP_INDEX >= 0: + abs_rows = tl.abs(k_rows) + if BSEARCH_ITERS > 0: + target_above = HEAD_DIM - K_CLIP_INDEX + thr_lo = tl.zeros([BLOCK_TOK], dtype=tl.float32) + thr_hi = tl.max(abs_rows, axis=1) + for _ in tl.static_range(BSEARCH_ITERS): + thr_mid = (thr_lo + thr_hi) * 0.5 + cnt_above = tl.sum((abs_rows > thr_mid[:, None]).to(tl.int32), axis=1) + too_many = cnt_above > target_above + thr_lo = tl.where(too_many, thr_mid, thr_lo) + thr_hi = tl.where(too_many, thr_hi, thr_mid) + thr = thr_hi + else: + sorted_rows = tl.sort(abs_rows) + pick = (full_offs == K_CLIP_INDEX)[None, :] + thr = tl.sum(tl.where(pick, sorted_rows, 0.0), axis=1) + k_rows = tl.minimum( + tl.maximum(k_rows, -thr[:, None]), + thr[:, None], + ) + + if LLOYD_MAX: + K_LM_C_1BIT: tl.constexpr = 0.79788456 + k_mean = tl.sum(k_rows, axis=1) / HEAD_DIM + k_diff = k_rows - k_mean[:, None] + k_std = tl.sqrt(tl.sum(k_diff * k_diff, axis=1) / HEAD_DIM + 1e-8) + k_scale = tl.maximum(2.0 * K_LM_C_1BIT * k_std, 1e-8) + k_zero = 0.5 - k_mean / k_scale + else: + k_min = tl.min(k_rows, axis=1) + k_max = tl.max(k_rows, axis=1) + k_scale = tl.maximum(k_max - k_min, 1e-8) + k_zero = -k_min / k_scale + + k_r = tl.reshape(k_rows, (BLOCK_TOK, 8, BLOCK_OCTANT)) + k_p = tl.permute(k_r, (0, 2, 1)) + k_s = tl.reshape(k_p, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + k_lo, k_hi = tl.split(k_s) + k_e0, k_e2 = tl.split(k_lo) + k_e1, k_e3 = tl.split(k_hi) + k_v0, k_v4 = tl.split(k_e0) + k_v2, k_v6 = tl.split(k_e2) + k_v1, k_v5 = tl.split(k_e1) + k_v3, k_v7 = tl.split(k_e3) + k_q0 = tl.minimum( + tl.maximum(k_v0 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q1 = tl.minimum( + tl.maximum(k_v1 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q2 = tl.minimum( + tl.maximum(k_v2 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q3 = tl.minimum( + tl.maximum(k_v3 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q4 = tl.minimum( + tl.maximum(k_v4 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q5 = tl.minimum( + tl.maximum(k_v5 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q6 = tl.minimum( + tl.maximum(k_v6 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_q7 = tl.minimum( + tl.maximum(k_v7 / k_scale[:, None] + k_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + k_packed = ( + k_q0 + | (k_q1 << 1) + | (k_q2 << 2) + | (k_q3 << 3) + | (k_q4 << 4) + | (k_q5 << 5) + | (k_q6 << 6) + | (k_q7 << 7) + ) + + k_cache_offset = ( + cache_loc[:, None] * k_cache_stride_loc + + head_idx * k_cache_stride_head + + dim_offs_o[None, :] * k_cache_stride_dim + ) + tl.store(k_cache_ptr + k_cache_offset, k_packed, mask=active[:, None]) + + k_sz_base = cache_loc * k_sz_stride_loc + head_idx * k_sz_stride_head + tl.store(k_sz_ptr + k_sz_base + 0 * k_sz_stride_dim, k_scale, mask=active) + tl.store(k_sz_ptr + k_sz_base + 1 * k_sz_stride_dim, k_zero, mask=active) + + # ---------- V: load + clip + scale/zero + pack + write (no rotate) ------ + v_base = ( + tok_offs[:, None] * v_input_stride_token + + head_idx * v_input_stride_head + + full_offs[None, :] * v_input_stride_dim + ) + v_rows = tl.load( + v_input_ptr + v_base, + mask=tok_mask[:, None], + other=0.0, + ).to(tl.float32) + + if V_CLIP_INDEX >= 0: + abs_rows = tl.abs(v_rows) + if BSEARCH_ITERS > 0: + target_above = HEAD_DIM - V_CLIP_INDEX + thr_lo = tl.zeros([BLOCK_TOK], dtype=tl.float32) + thr_hi = tl.max(abs_rows, axis=1) + for _ in tl.static_range(BSEARCH_ITERS): + thr_mid = (thr_lo + thr_hi) * 0.5 + cnt_above = tl.sum((abs_rows > thr_mid[:, None]).to(tl.int32), axis=1) + too_many = cnt_above > target_above + thr_lo = tl.where(too_many, thr_mid, thr_lo) + thr_hi = tl.where(too_many, thr_hi, thr_mid) + thr = thr_hi + else: + sorted_rows = tl.sort(abs_rows) + pick = (full_offs == V_CLIP_INDEX)[None, :] + thr = tl.sum(tl.where(pick, sorted_rows, 0.0), axis=1) + v_rows = tl.minimum( + tl.maximum(v_rows, -thr[:, None]), + thr[:, None], + ) + + if LLOYD_MAX: + V_LM_C_1BIT: tl.constexpr = 0.79788456 + v_mean = tl.sum(v_rows, axis=1) / HEAD_DIM + v_diff = v_rows - v_mean[:, None] + v_std = tl.sqrt(tl.sum(v_diff * v_diff, axis=1) / HEAD_DIM + 1e-8) + v_scale = tl.maximum(2.0 * V_LM_C_1BIT * v_std, 1e-8) + v_zero = 0.5 - v_mean / v_scale + else: + v_min = tl.min(v_rows, axis=1) + v_max = tl.max(v_rows, axis=1) + v_scale = tl.maximum(v_max - v_min, 1e-8) + v_zero = -v_min / v_scale + + v_r = tl.reshape(v_rows, (BLOCK_TOK, 8, BLOCK_OCTANT)) + v_p = tl.permute(v_r, (0, 2, 1)) + v_s = tl.reshape(v_p, (BLOCK_TOK, BLOCK_OCTANT, 2, 2, 2)) + v_lo, v_hi = tl.split(v_s) + v_e0, v_e2 = tl.split(v_lo) + v_e1, v_e3 = tl.split(v_hi) + v_v0, v_v4 = tl.split(v_e0) + v_v2, v_v6 = tl.split(v_e2) + v_v1, v_v5 = tl.split(v_e1) + v_v3, v_v7 = tl.split(v_e3) + v_q0 = tl.minimum( + tl.maximum(v_v0 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q1 = tl.minimum( + tl.maximum(v_v1 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q2 = tl.minimum( + tl.maximum(v_v2 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q3 = tl.minimum( + tl.maximum(v_v3 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q4 = tl.minimum( + tl.maximum(v_v4 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q5 = tl.minimum( + tl.maximum(v_v5 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q6 = tl.minimum( + tl.maximum(v_v6 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_q7 = tl.minimum( + tl.maximum(v_v7 / v_scale[:, None] + v_zero[:, None] + 0.5, 0.0), 1.0 + ).to(tl.uint8) + v_packed = ( + v_q0 + | (v_q1 << 1) + | (v_q2 << 2) + | (v_q3 << 3) + | (v_q4 << 4) + | (v_q5 << 5) + | (v_q6 << 6) + | (v_q7 << 7) + ) + + v_cache_offset = ( + cache_loc[:, None] * v_cache_stride_loc + + head_idx * v_cache_stride_head + + dim_offs_o[None, :] * v_cache_stride_dim + ) + tl.store(v_cache_ptr + v_cache_offset, v_packed, mask=active[:, None]) + + v_sz_base = cache_loc * v_sz_stride_loc + head_idx * v_sz_stride_head + tl.store(v_sz_ptr + v_sz_base + 0 * v_sz_stride_dim, v_scale, mask=active) + tl.store(v_sz_ptr + v_sz_base + 1 * v_sz_stride_dim, v_zero, mask=active) + + +def _pick_block_tok_and_num_warps_for_dot( + head_dim: int, elements_per_thread: int +) -> Tuple[int, int]: + block_tok, num_warps = _pick_block_tok_and_num_warps(head_dim, elements_per_thread) + if block_tok < 16: + block_tok = 16 + total_elems = block_tok * head_dim + assert total_elems % (32 * elements_per_thread) == 0, ( + f"BLOCK_TOK={block_tok} head_dim={head_dim} epp={elements_per_thread}: " + "tile size doesn't divide cleanly into 128-bit/thread loads" + ) + num_warps = total_elems // (32 * elements_per_thread) + return block_tok, num_warps + + +def quantized_set_kv_int1_oscar_rotate_k_clip_triton( + cache_k_unrotated: torch.Tensor, + cache_v_rotated: torch.Tensor, + R_k: torch.Tensor, + loc: torch.Tensor, + k_cache_buffer: torch.Tensor, + v_cache_buffer: torch.Tensor, + k_scales_zeros_buffer: torch.Tensor, + v_scales_zeros_buffer: torch.Tensor, + clip_ratio_k: float, + clip_ratio_v: float, + hp_global_offset=None, + lloyd_max: bool = False, +) -> None: + """Single-launch fused oscar K-rotation + clip + quantize + int1 pack.""" + assert cache_k_unrotated.shape == cache_v_rotated.shape, ( + "K/V shape mismatch in oscar_rotate_k_clip (int1 kernel requires identical " + f"shapes incl. head_dim): {cache_k_unrotated.shape} vs {cache_v_rotated.shape}" + ) + num_tokens, num_heads, head_dim = cache_k_unrotated.shape + if num_tokens == 0: + return + + assert head_dim % 8 == 0, ( + f"head_dim must be divisible by 8 for INT1, got {head_dim}" + ) + assert _is_power_of_two(head_dim), ( + f"oscar rotate+clip int1 kernel requires power-of-two head_dim, got {head_dim}" + ) + assert R_k.shape == (head_dim, head_dim), ( + f"R_k must be [head_dim, head_dim]={head_dim}x{head_dim}, " + f"got {tuple(R_k.shape)}" + ) + assert R_k.dtype == cache_k_unrotated.dtype, ( + f"R_k dtype ({R_k.dtype}) must match input dtype ({cache_k_unrotated.dtype})" + ) + + if _get_num_scale_groups(k_scales_zeros_buffer) != 1: + raise NotImplementedError( + "oscar rotate+clip+quant fused int1 kernel requires single-scale K layout " + f"(got num_groups={_get_num_scale_groups(k_scales_zeros_buffer)})" + ) + if _get_num_scale_groups(v_scales_zeros_buffer) != 1: + raise NotImplementedError( + "oscar rotate+clip+quant fused int1 kernel requires single-scale V layout " + f"(got num_groups={_get_num_scale_groups(v_scales_zeros_buffer)})" + ) + + elements_per_thread = _vectorized_elems_per_thread(cache_k_unrotated.dtype) + block_tok, num_warps = _pick_block_tok_and_num_warps_for_dot( + head_dim, elements_per_thread + ) + grid = (triton.cdiv(num_tokens, block_tok), num_heads) + _kv_oscar_rotate_k_clip_single_kernel_int1[grid]( + cache_k_unrotated, + cache_v_rotated, + R_k, + loc, + k_cache_buffer, + v_cache_buffer, + k_scales_zeros_buffer, + v_scales_zeros_buffer, + num_tokens, + num_heads, + cache_k_unrotated.stride(0), + cache_k_unrotated.stride(1), + cache_k_unrotated.stride(2), + cache_v_rotated.stride(0), + cache_v_rotated.stride(1), + cache_v_rotated.stride(2), + R_k.stride(0), + R_k.stride(1), + k_cache_buffer.stride(0), + k_cache_buffer.stride(1), + k_cache_buffer.stride(2), + v_cache_buffer.stride(0), + v_cache_buffer.stride(1), + v_cache_buffer.stride(2), + k_scales_zeros_buffer.stride(0), + k_scales_zeros_buffer.stride(1), + k_scales_zeros_buffer.stride(2), + v_scales_zeros_buffer.stride(0), + v_scales_zeros_buffer.stride(1), + v_scales_zeros_buffer.stride(2), + HP_OFFSET=-1 if hp_global_offset is None else int(hp_global_offset), + HEAD_DIM=head_dim, + BLOCK_OCTANT=head_dim // 8, + BLOCK_TOK=block_tok, + K_CLIP_INDEX=_clip_index(clip_ratio_k, head_dim), + V_CLIP_INDEX=_clip_index(clip_ratio_v, head_dim), + BSEARCH_ITERS=(head_dim.bit_length() - 1) if head_dim >= 64 else 0, + LLOYD_MAX=bool(lloyd_max), + num_warps=num_warps, + num_stages=1, + ) + + +# --------------------------------------------------------------------------- +# Non-clip pretransformed pack (used when clip ratios are 0). Mirrors the +# `quantized_set_kv_int2_pretransformed_triton` entry point from the +# `fused_hadamard_int2_kv` module so the unified pool can dispatch the same +# way when no clip is requested. The implementation just reuses the clip +# kernel with CLIP_INDEX = -1 (disable). +# --------------------------------------------------------------------------- + + +def quantized_set_kv_int1_pretransformed_triton( + cache_k: torch.Tensor, + cache_v: torch.Tensor, + loc: torch.Tensor, + k_cache_buffer: torch.Tensor, + v_cache_buffer: torch.Tensor, + k_scales_zeros_buffer: torch.Tensor, + v_scales_zeros_buffer: torch.Tensor, + hp_global_offset=None, +) -> None: + """No-clip int1 pretransformed pack (clip kernel with CLIP_INDEX=-1).""" + quantized_set_kv_int1_pretransformed_clip_triton( + cache_k, + cache_v, + loc, + k_cache_buffer, + v_cache_buffer, + k_scales_zeros_buffer, + v_scales_zeros_buffer, + clip_ratio_k=0.0, + clip_ratio_v=0.0, + hp_global_offset=hp_global_offset, + ) diff --git a/sglang-research/python/sglang/QuantKernel/oscar_rotation_pq_k_kv.py b/sglang-research/python/sglang/QuantKernel/oscar_rotation_pq_k_kv.py new file mode 100644 index 000000000..eee7151cf --- /dev/null +++ b/sglang-research/python/sglang/QuantKernel/oscar_rotation_pq_k_kv.py @@ -0,0 +1,585 @@ +""" +Product Quantization (PQ) for K cache with INT2 LM fallback for V. + +PQ K scheme: + - K split into N_SUB=16 sub-vectors of SUB_DIM=8 channels each + - Each sub-vector encoded to 1 byte (8-bit index into 256 centroids) + - True 1.0 bpe for K (shared codebook, no per-token scale/zero) + - SQNR on real Qwen3-4B-Thinking data: +9.09 dB K (vs INT2 LM +8.33 dB at 2.5 bpe) + +Codebook format (fp16, stored on GPU): + - K codebook: [N_SUB, N_CENTROIDS, SUB_DIM] = [16, 256, 8] fp16 = 64KB + - Precomputed from calibration data via K-means + - Saved to: /k_pq_codebook_n16_c256_d8.pt + +V is handled by existing INT2 LM kernel (SGLANG_LLOYD_MAX=1) at 2.5 bpe. + +Combined BPE: K=1.0 bpe, V=2.5 bpe → avg 1.75 bpe +Combined SQNR: K=+9.09 dB, V=+8.34 dB → avg +8.7 dB (exceeds INT2 uniform) + +Encode kernel: + - grid: (cdiv(num_tokens, BLOCK_TOK), num_heads) + - For each (tok_block, head): 16 sub-vectors × 8D + - Per sub-vector: element-wise dot product loop over SUB_DIM dimensions + (avoids tensor core padding complications for K=8) + - L2² = ||k||² + ||c||² - 2·k·c, argmin over 256 centroids + +Decode kernel: + - grid: (cdiv(num_tokens, BLOCK_TOK), num_heads) + - For each (tok_block, head): 16 sub-vector lookups + - Gather centroid vectors from codebook and write to output +""" + +from __future__ import annotations +import math, os +from typing import Optional + +import torch +import triton +import triton.language as tl + +N_SUB = 16 # number of sub-vectors per head +SUB_DIM = 8 # channels per sub-vector +N_CENTROIDS = 256 # 2^8, 1 byte per code + +# ─── codebook helpers ───────────────────────────────────────────────────────── + + +def build_pq_codebook( + data: torch.Tensor, # [N_samples, head_dim] + n_sub: int = N_SUB, + n_centroids: int = N_CENTROIDS, + sub_dim: int = SUB_DIM, + n_iter: int = 30, + seed: int = 42, +) -> torch.Tensor: + """ + Train per-position PQ codebook using K-means on float32 data. + Returns codebook [n_sub, n_centroids, sub_dim] as float16. + """ + from scipy.cluster.vq import kmeans2 + import numpy as np + + N, D = data.shape + assert D == n_sub * sub_dim + sub_data = data.float().numpy().reshape(N, n_sub, sub_dim) + + books = [] + for s in range(n_sub): + x = sub_data[:, s, :] # [N, sub_dim] + print(f" sub {s}/{n_sub}: kmeans {n_centroids} centroids on {N} samples...") + cb, _ = kmeans2(x, n_centroids, minit="points", iter=n_iter, seed=seed) + books.append(cb) + + codebook = torch.from_numpy(np.stack(books)).to(torch.float16) + return codebook # [n_sub, n_centroids, sub_dim] + + +def save_pq_codebook(codebook: torch.Tensor, path: str) -> None: + torch.save(codebook, path) + + +def load_pq_codebook(path: str, device="cuda") -> torch.Tensor: + cb = torch.load(path, map_location="cpu").to(torch.float16) + return cb.to(device).contiguous() # [n_sub, n_centroids, sub_dim] + + +def pq_codebook_norms(codebook: torch.Tensor) -> torch.Tensor: + """Precompute ||c||² for each centroid: [n_sub, n_centroids] fp32.""" + return (codebook.float() ** 2).sum(dim=-1) # float32 for accumulation accuracy + + +# ─── Triton encode kernel ───────────────────────────────────────────────────── + + +@triton.jit +def _pq_encode_k_kernel( + k_ptr, # [num_tok, num_heads, HEAD_DIM] float16 input (row-major) + loc_ptr, # [num_tok] int32 — position in KV cache + codes_ptr, # [max_tok, num_heads, N_SUB] uint8 output (KV cache buffer) + cb_ptr, # [N_SUB, N_CENTROIDS, SUB_DIM] float16 codebook (contiguous) + cb_norm2_ptr, # [N_SUB, N_CENTROIDS] float32 ||centroid||² + num_tokens, + num_heads, + k_stride_tok, + k_stride_head, + # k_stride_dim assumed = 1 + codes_stride_loc, + codes_stride_head, + codes_stride_sub, + HEAD_DIM: tl.constexpr, # 128 + N_SUB: tl.constexpr, # 16 + SUB_DIM: tl.constexpr, # 8 + N_CENTROIDS: tl.constexpr, # 256 + BLOCK_TOK: tl.constexpr, # e.g. 16 + HP_OFFSET: tl.constexpr, +): + pid_tok = tl.program_id(0) + pid_head = tl.program_id(1) + + tok_range = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + active = tok_range < num_tokens + + # cache location for each token in this block + cache_loc = tl.load(loc_ptr + tok_range, mask=active, other=0) + if HP_OFFSET >= 0: + active &= cache_loc < HP_OFFSET + + cent_range = tl.arange(0, N_CENTROIDS) # [N_CENTROIDS] + + for s in tl.static_range(N_SUB): + # Accumulate: ||k_sub||² [BLOCK_TOK] and k·c [BLOCK_TOK, N_CENTROIDS] + k_norm2 = tl.zeros([BLOCK_TOK], dtype=tl.float32) + dot = tl.zeros([BLOCK_TOK, N_CENTROIDS], dtype=tl.float32) + + for d in tl.static_range(SUB_DIM): + # Load k_d for this dimension: [BLOCK_TOK] + k_d = tl.load( + k_ptr + + tok_range * k_stride_tok + + pid_head * k_stride_head + + (s * SUB_DIM + d), + mask=active, + other=0.0, + ).to(tl.float32) + + k_norm2 += k_d * k_d + + # Load codebook column d for sub-vector s: [N_CENTROIDS] + cb_d = tl.load(cb_ptr + (s * N_CENTROIDS + cent_range) * SUB_DIM + d).to( + tl.float32 + ) + + # Outer product accumulation: [BLOCK_TOK, N_CENTROIDS] + dot += k_d[:, None] * cb_d[None, :] + + # Load precomputed ||c||²: [N_CENTROIDS] + cb_n2 = tl.load(cb_norm2_ptr + s * N_CENTROIDS + cent_range).to(tl.float32) + + # L2² = ||k||² + ||c||² - 2·k·c + d2 = k_norm2[:, None] + cb_n2[None, :] - 2.0 * dot # [BLOCK_TOK, N_CENTROIDS] + + # Nearest centroid + code = tl.argmin(d2, axis=1).to(tl.uint8) # [BLOCK_TOK] + + # Store code to KV cache buffer + out_off = ( + cache_loc * codes_stride_loc + + pid_head * codes_stride_head + + s * codes_stride_sub + ) + tl.store(codes_ptr + out_off, code, mask=active) + + +# ─── Triton decode kernel ───────────────────────────────────────────────────── + + +@triton.jit +def _pq_decode_k_kernel( + codes_ptr, # [max_tok, num_heads, N_SUB] uint8 + out_ptr, # [num_out_tok, num_heads, HEAD_DIM] float16 output + cb_ptr, # [N_SUB, N_CENTROIDS, SUB_DIM] float16 + num_tokens, + num_heads, + codes_stride_loc, + codes_stride_head, + codes_stride_sub, + out_stride_tok, + out_stride_head, + out_stride_dim, + HEAD_DIM: tl.constexpr, + N_SUB: tl.constexpr, + SUB_DIM: tl.constexpr, + N_CENTROIDS: tl.constexpr, + BLOCK_TOK: tl.constexpr, +): + pid_tok = tl.program_id(0) + pid_head = tl.program_id(1) + + tok_range = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + active = tok_range < num_tokens + + sub_range = tl.arange(0, SUB_DIM) + + for s in tl.static_range(N_SUB): + code = tl.load( + codes_ptr + + tok_range * codes_stride_loc + + pid_head * codes_stride_head + + s * codes_stride_sub, + mask=active, + other=0, + ).to(tl.int32) # [BLOCK_TOK] + + # Gather centroid vectors: [BLOCK_TOK, SUB_DIM] + # layout: cb[s, code, :] → cb_ptr + s*N_CENTROIDS*SUB_DIM + code*SUB_DIM + d + cb_idx = (s * N_CENTROIDS + code[:, None]) * SUB_DIM + sub_range[None, :] + recon = tl.load( + cb_ptr + cb_idx, mask=active[:, None], other=0.0 + ) # [BLOCK_TOK, SUB_DIM] fp16 + + # Write reconstructed sub-vector to output + out_off = ( + tok_range[:, None] * out_stride_tok + + pid_head * out_stride_head + + (s * SUB_DIM + sub_range[None, :]) * out_stride_dim + ) + tl.store(out_ptr + out_off, recon, mask=active[:, None]) + + +@triton.jit +def _pq_decode_k_at_locs_kernel( + codes_ptr, + loc_ptr, + out_ptr, + cb_ptr, + num_tokens, + num_heads, + codes_stride_loc, + codes_stride_head, + codes_stride_sub, + out_stride_tok, + out_stride_head, + out_stride_dim, + N_SUB: tl.constexpr, + SUB_DIM: tl.constexpr, + N_CENTROIDS: tl.constexpr, + BLOCK_TOK: tl.constexpr, + HP_OFFSET: tl.constexpr, +): + """Decode cache rows selected by loc; HP locs produce zero rows.""" + pid_tok = tl.program_id(0) + pid_head = tl.program_id(1) + tok_range = pid_tok * BLOCK_TOK + tl.arange(0, BLOCK_TOK) + token_mask = tok_range < num_tokens + cache_loc = tl.load(loc_ptr + tok_range, mask=token_mask, other=0).to(tl.int64) + active = token_mask + if HP_OFFSET >= 0: + active &= cache_loc < HP_OFFSET + safe_loc = tl.where(active, cache_loc, 0) + sub_range = tl.arange(0, SUB_DIM) + + for sub in tl.static_range(N_SUB): + code = tl.load( + codes_ptr + + safe_loc * codes_stride_loc + + pid_head * codes_stride_head + + sub * codes_stride_sub, + mask=active, + other=0, + ).to(tl.int32) + cb_idx = (sub * N_CENTROIDS + code[:, None]) * SUB_DIM + sub_range[None, :] + recon = tl.load(cb_ptr + cb_idx, mask=active[:, None], other=0.0) + out_off = ( + tok_range[:, None] * out_stride_tok + + pid_head * out_stride_head + + (sub * SUB_DIM + sub_range[None, :]) * out_stride_dim + ) + tl.store(out_ptr + out_off, recon, mask=token_mask[:, None]) + + +# ─── Python launch wrappers ─────────────────────────────────────────────────── + + +def pq_encode_k( + cache_k: torch.Tensor, # [num_tokens, num_heads, head_dim] bf16/fp16 + loc: torch.Tensor, # [num_tokens] int32 — positions in KV cache buffer + k_codes_buffer: torch.Tensor, # [max_tokens, num_heads, N_SUB] uint8 output + codebook: torch.Tensor, # [N_SUB, N_CENTROIDS, SUB_DIM] fp16 + codebook_norms: torch.Tensor, # [N_SUB, N_CENTROIDS] fp32 + block_tok: int = 16, + hp_global_offset: Optional[int] = None, +) -> None: + """Encode K tensor using product quantization, writing codes into the KV cache buffer.""" + num_tokens, num_heads, head_dim = cache_k.shape + if num_tokens == 0: + return + + n_sub, n_centroids, sub_dim = codebook.shape + assert head_dim == n_sub * sub_dim, ( + f"head_dim {head_dim} != n_sub*sub_dim {n_sub * sub_dim}" + ) + # Codes are stored as uint8, so any n_centroids <= 256 is valid (RVQ stage2 uses 16). + assert n_centroids <= 256, f"n_centroids {n_centroids} must be <= 256 (uint8 codes)" + + k_fp16 = cache_k.to(torch.float16).contiguous() + cb = codebook.to(torch.float16).contiguous() + cb_n2 = codebook_norms.float().contiguous() + + grid = (triton.cdiv(num_tokens, block_tok), num_heads) + _pq_encode_k_kernel[grid]( + k_fp16, + loc, + k_codes_buffer, + cb, + cb_n2, + num_tokens, + num_heads, + k_fp16.stride(0), + k_fp16.stride(1), + k_codes_buffer.stride(0), + k_codes_buffer.stride(1), + k_codes_buffer.stride(2), + HEAD_DIM=head_dim, + N_SUB=n_sub, + SUB_DIM=sub_dim, + N_CENTROIDS=n_centroids, + BLOCK_TOK=block_tok, + HP_OFFSET=-1 if hp_global_offset is None else int(hp_global_offset), + num_warps=4, + num_stages=2, + ) + + +def pq_decode_k( + k_codes_buffer: torch.Tensor, # [max_tokens, num_heads, N_SUB] uint8 + codebook: torch.Tensor, # [N_SUB, N_CENTROIDS, SUB_DIM] fp16 + num_decode_tokens: int, + head_dim: int, + block_tok: int = 16, +) -> torch.Tensor: + """Decode PQ codes to reconstructed K tensor [num_decode_tokens, num_heads, head_dim] fp16.""" + _, num_heads, n_sub = k_codes_buffer.shape + n_centroids, sub_dim = codebook.shape[1], codebook.shape[2] + + out = torch.empty( + num_decode_tokens, + num_heads, + head_dim, + dtype=torch.float16, + device=k_codes_buffer.device, + ) + cb = codebook.to(torch.float16).contiguous() + + grid = (triton.cdiv(num_decode_tokens, block_tok), num_heads) + _pq_decode_k_kernel[grid]( + k_codes_buffer, + out, + cb, + num_decode_tokens, + num_heads, + k_codes_buffer.stride(0), + k_codes_buffer.stride(1), + k_codes_buffer.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + HEAD_DIM=head_dim, + N_SUB=n_sub, + SUB_DIM=sub_dim, + N_CENTROIDS=n_centroids, + BLOCK_TOK=block_tok, + num_warps=4, + num_stages=2, + ) + return out + + +def pq_decode_k_at_locs( + k_codes_buffer: torch.Tensor, + loc: torch.Tensor, + codebook: torch.Tensor, + head_dim: int, + block_tok: int = 16, + hp_global_offset: Optional[int] = None, +) -> torch.Tensor: + """Decode cache rows selected by loc without data-dependent indexing.""" + num_tokens = int(loc.shape[0]) + _, num_heads, n_sub = k_codes_buffer.shape + n_centroids, sub_dim = codebook.shape[1], codebook.shape[2] + assert head_dim == n_sub * sub_dim + out = torch.empty( + num_tokens, + num_heads, + head_dim, + dtype=torch.float16, + device=k_codes_buffer.device, + ) + if num_tokens == 0: + return out + cb = codebook.to(torch.float16).contiguous() + grid = (triton.cdiv(num_tokens, block_tok), num_heads) + _pq_decode_k_at_locs_kernel[grid]( + k_codes_buffer, + loc, + out, + cb, + num_tokens, + num_heads, + k_codes_buffer.stride(0), + k_codes_buffer.stride(1), + k_codes_buffer.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + N_SUB=n_sub, + SUB_DIM=sub_dim, + N_CENTROIDS=n_centroids, + BLOCK_TOK=block_tok, + HP_OFFSET=-1 if hp_global_offset is None else int(hp_global_offset), + num_warps=4, + num_stages=1, + ) + return out + + +# ─── Python reference (for SQNR validation) ────────────────────────────────── + + +def pq_encode_decode_python( + k: torch.Tensor, # [N, head_dim] fp32 + codebook: torch.Tensor, # [N_SUB, N_CENTROIDS, SUB_DIM] fp32 +) -> torch.Tensor: + """Reference PQ encode+decode on CPU, returns reconstructed K [N, head_dim] fp32.""" + N, D = k.shape + n_sub, n_centroids, sub_dim = codebook.shape + k_np = k.float().numpy() + cb_np = codebook.float().numpy() + sub_data = k_np.reshape(N, n_sub, sub_dim) + + recon_subs = [] + for s in range(n_sub): + x = sub_data[:, s, :] # [N, sub_dim] + cb = cb_np[s] # [n_centroids, sub_dim] + # L2 distance + d2 = ((x[:, None, :] - cb[None, :, :]) ** 2).sum(axis=2) # [N, n_centroids] + codes = d2.argmin(axis=1) # [N] + recon_subs.append(cb[codes]) # [N, sub_dim] + + import numpy as np + + recon = np.concatenate(recon_subs, axis=1) # [N, head_dim] + return torch.from_numpy(recon) + + +# ─── codebook training ──────────────────────────────────────────────────────── + + +def train_and_save_codebook_from_dumps( + dump_base: str, + rot_dir: str, + save_path: str, + n_sub: int = N_SUB, + n_centroids: int = N_CENTROIDS, + sub_dim: int = SUB_DIM, + max_samples: int = 300_000, +): + """Train PQ codebook from actual model K dumps after OSCAR rotation.""" + import glob + import numpy as np + + rk_data = torch.load(f"{rot_dir}/k_rotation_qqt_r_h_pbr.pt", map_location="cpu") + R_k = {v["layer_id"]: v["rotation"] for v in rk_data["layers"].values()} + + qkv_base = f"{dump_base}/qkv_dumps/gpqa" + layers = sorted( + int(d.split("_")[1]) + for d in os.listdir(qkv_base) + if d.startswith("layer_") and os.path.isdir(f"{qkv_base}/{d}") + ) + + k_chunks = [] + for lid in layers: + for kf in sorted(glob.glob(f"{qkv_base}/layer_{lid}/k/*.pt")): + k = torch.load(kf, map_location="cpu").float() @ R_k[lid] + k_chunks.append(k.reshape(-1, n_sub * sub_dim)) + + k_all = torch.cat(k_chunks, dim=0) # [N, head_dim] + print(f"Total samples: {k_all.shape[0]}") + + rng = np.random.default_rng(42) + idx = rng.choice(k_all.shape[0], min(max_samples, k_all.shape[0]), replace=False) + k_train = k_all[idx] + print(f"Training on {k_train.shape[0]} samples...") + + codebook = build_pq_codebook(k_train, n_sub, n_centroids, sub_dim) + save_pq_codebook(codebook, save_path) + print(f"Saved: {save_path} shape={codebook.shape}") + return codebook + + +# ─── standalone test ───────────────────────────────────────────────────────── + + +def _test_sqnr(n_tok=1024, n_heads=8, seed=0): + """ + Quick correctness + SQNR test using synthetic Gaussian K data. + + NOTE: synthetic i.i.d. Gaussian data yields ~5 dB SQNR (scalar quantization limit). + Real Qwen3 K data after OSCAR rotation has inter-channel correlations that PQ + exploits, yielding ~9 dB (measured via sqnr_1bit_exploration.py). + """ + import numpy as np + + torch.manual_seed(seed) + head_dim = N_SUB * SUB_DIM # 128 + device = "cuda" + + # Synthetic: N(0,1) with random per-token-head scale (mimics real K) + k = torch.randn(n_tok, n_heads, head_dim, device=device) + scale = torch.rand(n_tok, n_heads, 1, device=device) * 2 + 0.1 + k = k * scale + + # Train codebook from same data (upper bound; real use: separate calibration set) + k_flat = k.reshape(-1, head_dim).cpu() + print("Training codebook...") + codebook_cpu = build_pq_codebook(k_flat, N_SUB, N_CENTROIDS, SUB_DIM, n_iter=20) + codebook = codebook_cpu.to(device) + cb_norms = pq_codebook_norms(codebook) + + # Encode via Triton + loc = torch.arange(n_tok, device=device, dtype=torch.int32) + codes_buf = torch.zeros(n_tok, n_heads, N_SUB, device=device, dtype=torch.uint8) + pq_encode_k(k, loc, codes_buf, codebook, cb_norms) + + # Decode via Triton + k_recon = pq_decode_k(codes_buf, codebook, n_tok, head_dim) + + # Triton SQNR + k_f = k.float() + r_f = k_recon.float() + sqnr_db = 10 * torch.log10((k_f**2).mean() / ((k_f - r_f) ** 2).mean()).item() + + # Python reference on first 64 tokens (correctness check) + n_check = 64 + k_py = k[:n_check].reshape(-1, head_dim) + k_py_recon = pq_encode_decode_python(k_py.cpu(), codebook_cpu) + py_sqnr = 10 * math.log10( + (k_py.float() ** 2).mean().item() + / ((k_py.float().cpu() - k_py_recon.float()) ** 2).mean().item() + ) + + # Code match: Triton codes vs Python codes + k_sub = k[:n_check].cpu().float().numpy().reshape(-1, N_SUB, SUB_DIM) + cb_np = codebook_cpu.float().numpy() + py_codes = np.zeros((n_check * n_heads, N_SUB), dtype=np.uint8) + for s in range(N_SUB): + x = k_sub[:, s, :] + d2 = ((x[:, None, :] - cb_np[s][None, :, :]) ** 2).sum(axis=2) + py_codes[:, s] = d2.argmin(axis=1) + triton_codes = codes_buf[:n_check].cpu().numpy().reshape(-1, N_SUB) + code_match = (py_codes == triton_codes).mean() * 100 + + print(f"Triton SQNR: {sqnr_db:.2f} dB (synthetic i.i.d. Gaussian)") + print(f"Python ref SQNR: {py_sqnr:.2f} dB (same data, fp32 codebook)") + print( + f"Code match: {code_match:.2f}% (Triton vs Python, first {n_check} tokens)" + ) + print(f" [Real Qwen3 K data yields ~+9 dB from inter-channel correlations]") + return sqnr_db, code_match + + +if __name__ == "__main__": + import sys + + if len(sys.argv) > 1 and sys.argv[1] == "train": + if len(sys.argv) != 5: + raise SystemExit( + "usage: oscar_rotation_pq_k_kv.py train " + " " + ) + train_and_save_codebook_from_dumps(*sys.argv[2:5]) + else: + sqnr, code_match = _test_sqnr() + ok = sqnr > 4.5 and code_match > 95.0 + print("PASS" if ok else "FAIL") diff --git a/sglang-research/python/sglang/srt/environ.py b/sglang-research/python/sglang/srt/environ.py index e064c4ca5..c27b6250c 100644 --- a/sglang-research/python/sglang/srt/environ.py +++ b/sglang-research/python/sglang/srt/environ.py @@ -236,6 +236,19 @@ class Envs: # the default uniform min-max. Applies only to single-scale pretransformed # clip kernels (num_groups == 1). Requires oscar rotation + clip enabled. SGLANG_LLOYD_MAX = EnvBool(False) + # Path to a PQ K codebook .pt file for the pq_k_int2v KV dtype. + # Expected keys: codebooks ([N_SUB, N_CENTS, SUB_DIM] list), n_sub, sub_dim, n_centroids. + SGLANG_PQ_K_CODEBOOK = EnvStr("") + # Optional per-layer PQ V codebook. When set with pq_k_int2v, V is stored + # as one uint8 centroid index per sub-vector instead of INT2 values. + SGLANG_PQ_V_CODEBOOK = EnvStr("") + # PQ decode kernel tuning. Zero selects the shape-dependent default. + # -1 = auto (ADC for graph batch <4), 0 = reconstruct K, 1 = ADC. + SGLANG_PQ_USE_ADC = EnvInt(-1) + SGLANG_PQ_BLOCK_N = EnvInt(0) + SGLANG_PQ_BLOCK_H = EnvInt(0) + SGLANG_PQ_NUM_WARPS = EnvInt(0) + SGLANG_PQ_NUM_STAGES = EnvInt(0) SGLANG_MIXED_KV_HP_MAX_SPLITS = EnvInt(8) HADAMARD_ORDER = EnvInt(16) diff --git a/sglang-research/python/sglang/srt/layers/attention/flashattention_backend.py b/sglang-research/python/sglang/srt/layers/attention/flashattention_backend.py index 148941727..a6a1331bb 100644 --- a/sglang-research/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/sglang-research/python/sglang/srt/layers/attention/flashattention_backend.py @@ -552,13 +552,14 @@ def forward_extend( k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, ): - # Int2 takes a dedicated rotation-aware path below: Q/K/V are + # Int2/int1/PQ take a dedicated rotation-aware path below: Q/K/V are # pre-rotated (Hadamard or Oscar) once, and the rotated K/V is then # written to the pool with ``already_hadamard_transformed=True`` so the # pool does not rotate a second time. We therefore skip the eager save - # here for int2 and defer to the quantized prefill branch. + # here and defer to the quantized prefill branch. int2_deferred_save = ( - self.kv_cache_dtype_str == "int2" and not self.use_mla + self.kv_cache_dtype_str in ("int2", "int1", "pq_k_int2v") + and not self.use_mla ) if k is not None: assert v is not None @@ -608,7 +609,7 @@ def forward_extend( ) window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1) k_descale, v_descale = None, None - if self.kv_cache_dtype_str == "int2": + if self.kv_cache_dtype_str in ("int2", "int1", "pq_k_int2v"): pass elif ( self.kv_cache_dtype_str != "auto" @@ -684,7 +685,7 @@ def forward_extend( # Use Flash Attention for prefill if not self.use_mla: # Do multi-head attention - if self.kv_cache_dtype_str == "int2": + if self.kv_cache_dtype_str in ("int2", "int1", "pq_k_int2v"): # Rotation-aware path that shares logic with TritonAttnBackend # (see python/sglang/srt/layers/attention/quantized_kv_prefill.py). # Handles the pure-int2 ``MHATokenToKVPool`` and the mixed @@ -1192,10 +1193,11 @@ def forward_decode( # kv_cache_dtype=int2`` at startup, so this branch should never be # reached with int2; the explicit guard here is a belt-and-braces # assertion rather than a silent ``.to("int2")`` TypeError. - assert self.kv_cache_dtype_str != "int2", ( - "FA3 forward_decode does not support int2 KV cache; use " - "--prefill-attention-backend fa3 --decode-attention-backend " - "triton instead (HybridAttnBackend routes int2 decode to triton)." + assert self.kv_cache_dtype_str not in ("int2", "int1", "pq_k_int2v"), ( + f"FA3 forward_decode does not support {self.kv_cache_dtype_str} " + "KV cache; use --prefill-attention-backend fa3 " + "--decode-attention-backend triton instead (HybridAttnBackend " + f"routes {self.kv_cache_dtype_str} decode to triton)." ) if self.kv_cache_dtype_str != "auto" and layer.head_dim <= 256: if layer.k_scale is not None: diff --git a/sglang-research/python/sglang/srt/layers/attention/quantized_kv_prefill.py b/sglang-research/python/sglang/srt/layers/attention/quantized_kv_prefill.py index ef813d521..43b54e070 100644 --- a/sglang-research/python/sglang/srt/layers/attention/quantized_kv_prefill.py +++ b/sglang-research/python/sglang/srt/layers/attention/quantized_kv_prefill.py @@ -121,16 +121,14 @@ def prepare_quantized_extend_qkv( """ need_v_inverse = False kv_dtype = kv_pool.dtype - if kv_dtype != "int2": + if kv_dtype not in ("int2", "int1", "pq_k_int2v"): return q, k, v, need_v_inverse if _pool_uses_oscar_rotation(kv_pool): layer_idx = layer.layer_id - kv_pool.start_layer R_k = kv_pool._R_k[layer_idx] R_v = kv_pool._R_v[layer_idx] - v_rotation_absorbed = bool( - getattr(layer, "oscar_v_rotation_absorbed", False) - ) + v_rotation_absorbed = bool(getattr(layer, "oscar_v_rotation_absorbed", False)) if not q_already_hadamard_transformed: q = _apply_oscar_rotation(q, R_k) if not kv_already_hadamard_transformed: @@ -176,6 +174,7 @@ def _mixed_prefix_dequant_kernel( HP_OFFSET: tl.constexpr, GROUP_SIZE: tl.constexpr, BLOCK_DIM: tl.constexpr, + INT1: tl.constexpr, ): token_idx = tl.program_id(0) head_idx = tl.program_id(1) @@ -184,9 +183,13 @@ def _mixed_prefix_dequant_kernel( slot = tl.load(prefix_indices_ptr + token_idx) is_hp = slot >= HP_OFFSET - quarter_dim = head_dim // 4 - byte_offsets = offs % quarter_dim + pack_bits: tl.constexpr = 1 if INT1 else 2 + pack_factor: tl.constexpr = 8 if INT1 else 4 + byte_dim: tl.constexpr = head_dim // pack_factor + byte_mask: tl.constexpr = (1 << pack_bits) - 1 + + byte_offsets = offs % byte_dim packed = tl.load( quant_ptr + slot * quant_stride_token @@ -195,8 +198,8 @@ def _mixed_prefix_dequant_kernel( mask=(~is_hp) & dim_mask, other=0, ) - shift = (offs // quarter_dim) * 2 - q = ((packed >> shift) & 0x03).to(tl.float32) + shift = (offs // byte_dim) * pack_bits + q = ((packed >> shift) & byte_mask).to(tl.float32) group_ids = offs // GROUP_SIZE scale = tl.load( @@ -237,6 +240,171 @@ def _mixed_prefix_dequant_kernel( ) +@triton.jit +def _mixed_prefix_pq_dequant_kernel( + prefix_indices_ptr, + codes_ptr, + codebook_ptr, + codes2_ptr, + codebook2_ptr, + hp_ptr, + out_ptr, + num_tokens, + num_heads, + codes_stride_token: tl.constexpr, + codes_stride_head: tl.constexpr, + codes_stride_sub: tl.constexpr, + codes2_stride_token: tl.constexpr, + codes2_stride_head: tl.constexpr, + codes2_stride_sub: tl.constexpr, + hp_stride_token: tl.constexpr, + hp_stride_head: tl.constexpr, + hp_stride_dim: tl.constexpr, + out_stride_token: tl.constexpr, + out_stride_head: tl.constexpr, + out_stride_dim: tl.constexpr, + HP_OFFSET: tl.constexpr, + HEAD_DIM: tl.constexpr, + SUB_DIM: tl.constexpr, + N_CENTROIDS: tl.constexpr, + N_CENTROIDS2: tl.constexpr, + BLOCK_DIM: tl.constexpr, + HAS_STAGE2: tl.constexpr, +): + """Decode a mixed HP/PQ prefix without data-dependent Python masks.""" + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + offs = tl.arange(0, BLOCK_DIM) + dim_mask = offs < HEAD_DIM + + slot = tl.load(prefix_indices_ptr + token_idx).to(tl.int64) + is_hp = slot >= HP_OFFSET + quant_slot = tl.where(is_hp, 0, slot) + hp_slot = tl.where(is_hp, slot - HP_OFFSET, 0) + + sub_idx = offs // SUB_DIM + sub_off = offs % SUB_DIM + code = tl.load( + codes_ptr + + quant_slot * codes_stride_token + + head_idx * codes_stride_head + + sub_idx * codes_stride_sub, + mask=(~is_hp) & dim_mask, + other=0, + ).to(tl.int64) + quant_val = tl.load( + codebook_ptr + (sub_idx * N_CENTROIDS + code) * SUB_DIM + sub_off, + mask=(~is_hp) & dim_mask, + other=0.0, + ).to(tl.float32) + + if HAS_STAGE2: + code2 = tl.load( + codes2_ptr + + quant_slot * codes2_stride_token + + head_idx * codes2_stride_head + + sub_idx * codes2_stride_sub, + mask=(~is_hp) & dim_mask, + other=0, + ).to(tl.int64) + quant_val += tl.load( + codebook2_ptr + (sub_idx * N_CENTROIDS2 + code2) * SUB_DIM + sub_off, + mask=(~is_hp) & dim_mask, + other=0.0, + ).to(tl.float32) + + hp_val = tl.load( + hp_ptr + + hp_slot * hp_stride_token + + head_idx * hp_stride_head + + offs * hp_stride_dim, + mask=is_hp & dim_mask, + other=0.0, + ) + out_val = tl.where(is_hp, hp_val, quant_val) + tl.store( + out_ptr + + token_idx * out_stride_token + + head_idx * out_stride_head + + offs * out_stride_dim, + out_val, + mask=(token_idx < num_tokens) & (head_idx < num_heads) & dim_mask, + ) + + +def _mixed_prefix_dequantize_pq_tensor( + prefix_indices: torch.Tensor, + codes: torch.Tensor, + codebook: torch.Tensor, + hp: torch.Tensor, + hp_offset: int, + model_dtype: torch.dtype, + codes2: Optional[torch.Tensor] = None, + codebook2: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens = prefix_indices.shape[0] + num_heads = codes.shape[1] + n_sub, n_centroids, sub_dim = codebook.shape + head_dim = n_sub * sub_dim + out = torch.empty( + (num_tokens, num_heads, head_dim), + dtype=model_dtype, + device=prefix_indices.device, + ) + if num_tokens == 0: + return out + + has_stage2 = codes2 is not None + if has_stage2: + assert codebook2 is not None + assert codebook2.shape[0] == n_sub + assert codebook2.shape[2] == sub_dim + codes2_arg = codes2 + codebook2_arg = codebook2 + n_centroids2 = int(codebook2.shape[1]) + else: + # Triton still requires pointer arguments even when the constexpr + # branch is disabled. + codes2_arg = codes + codebook2_arg = codebook + n_centroids2 = n_centroids + + grid = (num_tokens, num_heads) + _mixed_prefix_pq_dequant_kernel[grid]( + prefix_indices, + codes, + codebook, + codes2_arg, + codebook2_arg, + hp, + out, + num_tokens, + num_heads, + codes.stride(0), + codes.stride(1), + codes.stride(2), + codes2_arg.stride(0), + codes2_arg.stride(1), + codes2_arg.stride(2), + hp.stride(0), + hp.stride(1), + hp.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + HP_OFFSET=int(hp_offset), + HEAD_DIM=head_dim, + SUB_DIM=sub_dim, + N_CENTROIDS=n_centroids, + N_CENTROIDS2=n_centroids2, + BLOCK_DIM=triton.next_power_of_2(head_dim), + HAS_STAGE2=has_stage2, + num_warps=4, + num_stages=1, + ) + return out + + def _mixed_prefix_dequantize_tensor( prefix_indices: torch.Tensor, quantized: torch.Tensor, @@ -282,6 +450,7 @@ def _mixed_prefix_dequantize_tensor( HP_OFFSET=int(hp_offset), GROUP_SIZE=group_size, BLOCK_DIM=triton.next_power_of_2(head_dim), + INT1=(quantized.shape[-1] == head_dim // 8), num_warps=4, num_stages=1, ) @@ -325,9 +494,44 @@ def dequantize_prefix_kv( getattr(kv_pool, "mixed_kv_enabled", None) is not None and kv_pool.mixed_kv_enabled() ): - assert kv_pool.dtype == "int2", ( - f"Unsupported quantized KV dtype: {kv_pool.dtype}" - ) + assert kv_pool.dtype in ( + "int2", + "int1", + "pq_k_int2v", + ), f"Unsupported quantized KV dtype: {kv_pool.dtype}" + if kv_pool.dtype == "pq_k_int2v": + hp_off = kv_pool.hp_global_offset + out_k = _mixed_prefix_dequantize_pq_tensor( + prefix_indices, + kv_pool.get_raw_key_buffer(layer_id), + kv_pool.get_pq_codebook(layer_id), + kv_pool.get_hp_key_buffer(layer_id), + hp_off, + model_dtype, + codes2=kv_pool.get_raw_key_buffer2(layer_id), + codebook2=kv_pool.get_rvq_cb2(layer_id), + ) + vcb = kv_pool.get_pq_v_codebook(layer_id) + if vcb is not None: + out_v = _mixed_prefix_dequantize_pq_tensor( + prefix_indices, + kv_pool.get_raw_value_buffer(layer_id), + vcb, + kv_pool.get_hp_value_buffer(layer_id), + hp_off, + model_dtype, + ) + else: + out_v = _mixed_prefix_dequantize_tensor( + prefix_indices, + kv_pool.get_raw_value_buffer(layer_id), + kv_pool.get_value_scales_zeros(layer_id), + kv_pool.get_hp_value_buffer(layer_id), + hp_off, + kv_pool.v_head_dim, + model_dtype, + ) + return out_k, out_v return ( _mixed_prefix_dequantize_tensor( prefix_indices, @@ -351,11 +555,37 @@ def dequantize_prefix_kv( raw_k = kv_pool.get_raw_key_buffer(layer_id)[prefix_indices] raw_v = kv_pool.get_raw_value_buffer(layer_id)[prefix_indices] - scales_k = kv_pool.get_key_scales_zeros(layer_id)[prefix_indices] scales_v = kv_pool.get_value_scales_zeros(layer_id)[prefix_indices] - assert kv_pool.dtype == "int2", ( - f"Unsupported quantized KV dtype: {kv_pool.dtype}" - ) + assert kv_pool.dtype in ( + "int2", + "int1", + "pq_k_int2v", + ), f"Unsupported quantized KV dtype: {kv_pool.dtype}" + if kv_pool.dtype == "pq_k_int2v": + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_decode_k + + n = int(raw_k.shape[0]) + _cb = kv_pool.get_pq_codebook(layer_id) + head_dim_k = int(_cb.shape[0]) * int(_cb.shape[2]) + k_out = pq_decode_k(raw_k, _cb, n, head_dim_k) + # RVQ: add stage-2 residual reconstruction (None for plain PQ K). + _cb2 = kv_pool.get_rvq_cb2(layer_id) + if _cb2 is not None: + raw_k2 = kv_pool.get_raw_key_buffer2(layer_id)[prefix_indices] + k_out = k_out + pq_decode_k(raw_k2, _cb2, n, head_dim_k) + k_out = k_out.to(model_dtype) + _vcb = kv_pool.get_pq_v_codebook(layer_id) + if _vcb is not None: + n_sub_v = int(_vcb.shape[0]) + v_out = pq_decode_k(raw_v[..., :n_sub_v], _vcb, n, kv_pool.v_head_dim).to( + model_dtype + ) + else: + v_out = dequantize_kv_int2_triton( + raw_v, scales_v, kv_pool.v_head_dim, model_dtype + ) + return k_out, v_out + scales_k = kv_pool.get_key_scales_zeros(layer_id)[prefix_indices] return ( dequantize_kv_int2_triton(raw_k, scales_k, kv_pool.head_dim, model_dtype), dequantize_kv_int2_triton(raw_v, scales_v, kv_pool.v_head_dim, model_dtype), @@ -375,12 +605,18 @@ def apply_inverse_v_rotation( ``result`` must have shape ``[..., v_head_dim]``; callers should reshape beforehand if their output is stored flattened. """ - if not need_v_inverse or kv_pool.dtype != "int2": + if not need_v_inverse: return result + # Oscar rotation: V is stored in R_v-rotated space regardless of whether + # the rotation was absorbed into qkv_proj or applied explicitly at store + # time. Always apply R_v^T to recover the original space. if _pool_uses_oscar_rotation(kv_pool): layer_idx = layer.layer_id - kv_pool.start_layer R_v = kv_pool._R_v[layer_idx] return (result.to(R_v.dtype) @ R_v.T).contiguous() + # Hadamard rotation (non-Oscar pools) only applies to scalar int2/int1. + if kv_pool.dtype not in ("int2", "int1"): + return result return _apply_segmented_hadamard_transform(result) @@ -417,6 +653,7 @@ def _cpu_int_list(values) -> Optional[list[int]]: return [int(v) for v in values.tolist()] return [int(v) for v in values] + def build_prefix_indices_from_req_to_token( req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, diff --git a/sglang-research/python/sglang/srt/layers/attention/triton_backend.py b/sglang-research/python/sglang/srt/layers/attention/triton_backend.py index 54c4e67b3..0df9d404f 100644 --- a/sglang-research/python/sglang/srt/layers/attention/triton_backend.py +++ b/sglang-research/python/sglang/srt/layers/attention/triton_backend.py @@ -202,6 +202,9 @@ def __init__( from sglang.srt.layers.attention.triton_ops.decode_attention import ( decode_attention_fwd, decode_attention_fwd_int2_unified, + decode_attention_fwd_int1_unified, + decode_attention_fwd_int1_via_dequant, + decode_attention_fwd_pqk_int2v_unified, decode_attention_fwd_quantized, ) from sglang.srt.layers.attention.triton_ops.extend_attention import ( @@ -219,6 +222,15 @@ def __init__( self.decode_attention_fwd_int2_unified = torch.compiler.disable( decode_attention_fwd_int2_unified ) + self.decode_attention_fwd_int1_unified = torch.compiler.disable( + decode_attention_fwd_int1_unified + ) + self.decode_attention_fwd_int1_via_dequant = torch.compiler.disable( + decode_attention_fwd_int1_via_dequant + ) + self.decode_attention_fwd_pqk_int2v_unified = torch.compiler.disable( + decode_attention_fwd_pqk_int2v_unified + ) self.extend_attention_fwd = torch.compiler.disable(extend_attention_fwd) self.extend_attention_fwd_unified = torch.compiler.disable( extend_attention_fwd_unified @@ -1430,6 +1442,11 @@ def forward_extend( window_kv_offsets = None kv_pool = forward_batch.token_to_kv_pool + # int1 cannot use the quantized-dense prefill fast path (which lives + # in quantized_kv_prefill.py and is int2-specific); int1 prefill + # writes go through the unified pool's oscar rotate+pack int1 kernel + # and prefill attention falls back to standard fa3/triton on the HP + # tier (no quant read at prefill). use_quantized_dense_prefill = ( hasattr(kv_pool, "dtype") and kv_pool.dtype == "int2" @@ -1745,9 +1762,18 @@ def forward_decode( ): attn_logits = self.forward_metadata.swa_attn_logits - # Int2 quantized KV cache path (the only supported quant tier). + # Int2/Int1/Ternary quantized KV cache path. Ternary shares the int2 + # 2-bit storage and dequant kernel (zero=1.0, scale=max|x| stored in + # sz_buf → packed {0,1,2} dequants to {-1,0,+1}*scale). kv_pool = forward_batch.token_to_kv_pool - if hasattr(kv_pool, "dtype") and kv_pool.dtype == "int2": + if hasattr(kv_pool, "dtype") and kv_pool.dtype in ( + "int2", + "int1", + "pq_k_int2v", + ): + kv_dtype = kv_pool.dtype + is_int1 = kv_dtype == "int1" + is_pqk_int2v = kv_pool.dtype == "pq_k_int2v" uses_oscar = _pool_uses_oscar_rotation(kv_pool) q_for_decode = q.contiguous().view(-1, layer.tp_q_head_num, layer.qk_head_dim) @@ -1756,13 +1782,11 @@ def forward_decode( ) mixed_decode_enabled = ( self.enable_mixed_kv - and kv_pool.dtype == "int2" and sinks is None and mixed_decode_metadata_available ) if ( self.enable_mixed_kv - and kv_pool.dtype == "int2" and mixed_decode_metadata_available and sinks is not None ): @@ -1770,18 +1794,7 @@ def forward_decode( "Mixed KV windows do not support sink tokens in Triton decode." ) - # Hard guarantee that the upstream gating actually held: if mixed - # KV is enabled with an int2 pool, ``init_forward_metadata`` must - # have built the per-tier indices. Falling through to the - # non-mixed ``decode_attention_fwd_quantized`` path would treat - # HP slot ids (>= HP_OFFSET) as quant slot ids and read OOB - # garbage from the quant buffer. The known offenders are the - # ``spec_info != None`` decode-or-idle paths (currently gated out - # at server-args / model-runner level); this assertion makes the - # gating load-bearing at the kernel boundary so any future - # widening of those upstream gates surfaces here loudly instead - # of silently corrupting attention output. - if self.enable_mixed_kv and kv_pool.dtype == "int2": + if self.enable_mixed_kv: assert mixed_decode_metadata_available, ( "Mixed-KV pool active but mixed decode metadata not built. " "spec_info / non-decode-or-idle paths must not reach the " @@ -1800,54 +1813,110 @@ def forward_decode( q_for_decode = apply_segmented_hadamard_transform(q_for_decode) if mixed_decode_enabled: bs = q_for_decode.shape[0] - self.decode_attention_fwd_int2_unified( - q_for_decode, - kv_pool.get_hp_key_buffer(layer.layer_id), - kv_pool.get_hp_value_buffer(layer.layer_id), - kv_pool.get_raw_key_buffer(layer.layer_id), - kv_pool.get_raw_value_buffer(layer.layer_id), - kv_pool.get_key_scales_zeros(layer.layer_id), - kv_pool.get_value_scales_zeros(layer.layer_id), - o.view(-1, layer.tp_q_head_num, layer.v_head_dim), - self.forward_metadata.mixed_hp_kv_indptr, - self.forward_metadata.mixed_hp_kv_indices, - self.forward_metadata.mixed_quant_kv_indptr, - self.forward_metadata.mixed_quant_kv_indices, - self.forward_metadata.mixed_attn_logits[:bs], - self.forward_metadata.mixed_attn_lse[:bs], - self.forward_metadata.mixed_hp_num_kv_splits[:bs], - self.forward_metadata.mixed_quant_num_kv_splits[:bs], - self.max_hp_kv_splits, - self.max_kv_splits, - layer.scaling, - logit_cap=logits_soft_cap, - sinks=sinks, - xai_temperature_len=layer.xai_temperature_len, - ) + if is_pqk_int2v: + self.decode_attention_fwd_pqk_int2v_unified( + q_for_decode, + kv_pool.get_hp_key_buffer(layer.layer_id), + kv_pool.get_hp_value_buffer(layer.layer_id), + kv_pool.get_raw_key_buffer(layer.layer_id), + kv_pool.get_raw_value_buffer(layer.layer_id), + kv_pool.get_key_scales_zeros(layer.layer_id), + kv_pool.get_value_scales_zeros(layer.layer_id), + kv_pool.get_pq_codebook(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + self.forward_metadata.mixed_hp_kv_indptr, + self.forward_metadata.mixed_hp_kv_indices, + self.forward_metadata.mixed_quant_kv_indptr, + self.forward_metadata.mixed_quant_kv_indices, + self.forward_metadata.mixed_attn_logits[:bs], + self.forward_metadata.mixed_attn_lse[:bs], + self.forward_metadata.mixed_hp_num_kv_splits[:bs], + self.forward_metadata.mixed_quant_num_kv_splits[:bs], + self.max_hp_kv_splits, + self.max_kv_splits, + layer.scaling, + logit_cap=logits_soft_cap, + sinks=sinks, + xai_temperature_len=layer.xai_temperature_len, + quant_k_buffer2=kv_pool.get_raw_key_buffer2(layer.layer_id), + pq_codebook2=kv_pool.get_rvq_cb2(layer.layer_id), + pq_v_codebook=kv_pool.get_pq_v_codebook(layer.layer_id), + ) + else: + unified_kernel = ( + self.decode_attention_fwd_int1_unified + if is_int1 + else self.decode_attention_fwd_int2_unified + ) + unified_kernel( + q_for_decode, + kv_pool.get_hp_key_buffer(layer.layer_id), + kv_pool.get_hp_value_buffer(layer.layer_id), + kv_pool.get_raw_key_buffer(layer.layer_id), + kv_pool.get_raw_value_buffer(layer.layer_id), + kv_pool.get_key_scales_zeros(layer.layer_id), + kv_pool.get_value_scales_zeros(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + self.forward_metadata.mixed_hp_kv_indptr, + self.forward_metadata.mixed_hp_kv_indices, + self.forward_metadata.mixed_quant_kv_indptr, + self.forward_metadata.mixed_quant_kv_indices, + self.forward_metadata.mixed_attn_logits[:bs], + self.forward_metadata.mixed_attn_lse[:bs], + self.forward_metadata.mixed_hp_num_kv_splits[:bs], + self.forward_metadata.mixed_quant_num_kv_splits[:bs], + self.max_hp_kv_splits, + self.max_kv_splits, + layer.scaling, + logit_cap=logits_soft_cap, + sinks=sinks, + xai_temperature_len=layer.xai_temperature_len, + ) else: - # Use optimized quantized attention kernel - self.decode_attention_fwd_quantized( - q_for_decode, - kv_pool.get_raw_key_buffer(layer.layer_id), - kv_pool.get_raw_value_buffer(layer.layer_id), - kv_pool.get_key_scales_zeros(layer.layer_id), - kv_pool.get_value_scales_zeros(layer.layer_id), - o.view(-1, layer.tp_q_head_num, layer.v_head_dim), - kv_indptr, - kv_indices, - self.forward_metadata.attn_logits, - self.forward_metadata.attn_lse, - self.forward_metadata.num_kv_splits, - self.max_kv_splits, - layer.scaling, - kv_pool.dtype, - logit_cap=logits_soft_cap, - sinks=sinks, - xai_temperature_len=layer.xai_temperature_len, - ) - # int2: V is always rotated, so apply the inverse rotation to the - # output. Oscar mode uses ``o @ R_v.T``; Hadamard mode re-applies - # the segmented FWHT (self-inverse with 1/sqrt(N)). + if is_int1: + self.decode_attention_fwd_int1_via_dequant( + q_for_decode, + kv_pool.get_raw_key_buffer(layer.layer_id), + kv_pool.get_raw_value_buffer(layer.layer_id), + kv_pool.get_key_scales_zeros(layer.layer_id), + kv_pool.get_value_scales_zeros(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + kv_indptr, + kv_indices, + self.forward_metadata.attn_logits, + self.forward_metadata.attn_lse, + self.forward_metadata.num_kv_splits, + self.max_kv_splits, + layer.scaling, + logit_cap=logits_soft_cap, + sinks=sinks, + xai_temperature_len=layer.xai_temperature_len, + ) + else: + self.decode_attention_fwd_quantized( + q_for_decode, + kv_pool.get_raw_key_buffer(layer.layer_id), + kv_pool.get_raw_value_buffer(layer.layer_id), + kv_pool.get_key_scales_zeros(layer.layer_id), + kv_pool.get_value_scales_zeros(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + kv_indptr, + kv_indices, + self.forward_metadata.attn_logits, + self.forward_metadata.attn_lse, + self.forward_metadata.num_kv_splits, + self.max_kv_splits, + layer.scaling, + kv_dtype, + logit_cap=logits_soft_cap, + sinks=sinks, + xai_temperature_len=layer.xai_temperature_len, + ) + # int2/int1: V is stored in R_v-rotated space (either because it + # was explicitly rotated at store time, or because qkv_proj was + # modified to produce V_r = V @ R_v directly when V rotation is + # absorbed). Either way, apply the inverse rotation to the output. + # Oscar mode: o @ R_v.T. Hadamard mode: segmented FWHT (self-inv). if uses_oscar: R_v = kv_pool._R_v[oscar_layer_idx] o3 = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) diff --git a/sglang-research/python/sglang/srt/layers/attention/triton_ops/decode_attention.py b/sglang-research/python/sglang/srt/layers/attention/triton_ops/decode_attention.py index 16d231cd1..63974f355 100644 --- a/sglang-research/python/sglang/srt/layers/attention/triton_ops/decode_attention.py +++ b/sglang-research/python/sglang/srt/layers/attention/triton_ops/decode_attention.py @@ -500,7 +500,13 @@ def _decode_grouped_att_m_fwd( batch, head_num = q.shape[0], q.shape[1] kv_group_num = q.shape[1] // k_buffer.shape[1] - BLOCK_H = 16 + requested_block_h = 16 + if requested_block_h >= kv_group_num: + BLOCK_H = triton.next_power_of_2(kv_group_num) + else: + BLOCK_H = requested_block_h + while BLOCK_H > 1 and kv_group_num % BLOCK_H != 0: + BLOCK_H //= 2 MAX_KV_SPLITS = max_kv_splits grid = ( batch, @@ -627,7 +633,9 @@ def _fwd_kernel_stage2( if WRITE_LSE: # Per-seq log-sum-exp = e_max + log(e_sum). O_lse has shape # [bs, num_heads]; batch stride = stride_obs // Lv = num_heads. - tl.store(O_lse + cur_batch * (stride_obs // Lv) + cur_head, e_max + tl.log(e_sum)) + tl.store( + O_lse + cur_batch * (stride_obs // Lv) + cur_head, e_max + tl.log(e_sum) + ) def _decode_softmax_reducev_fwd( @@ -962,6 +970,7 @@ def _fwd_kernel_stage1_quant_int2( Lv: tl.constexpr, xai_temperature_len: tl.constexpr, GROUP_SIZE: tl.constexpr, + INT1: tl.constexpr, ): cur_batch = tl.program_id(0) cur_head = tl.program_id(1) @@ -1024,14 +1033,22 @@ def _fwd_kernel_stage1_quant_int2( other=0, ) - # Load packed INT2 K (uint8, 4 values per byte) + # Load packed K into four logical quarters. INT2 has one crumb + # from each quarter per byte. INT1 has two octants per quarter, + # so each physical byte is loaded twice with a different bit. offs_d_packed = tl.arange(0, BLOCK_DMODEL // 4) mask_d_packed = offs_d_packed < (Lk // 4) + if INT1: + k_byte_offsets = offs_d_packed % (Lk // 8) + k_bit_in_quarter = offs_d_packed // (Lk // 8) + else: + k_byte_offsets = offs_d_packed + k_bit_in_quarter = tl.zeros([BLOCK_DMODEL // 4], dtype=tl.int32) offs_buf_k_packed = ( kv_loc[:, None] * stride_buf_kbs + cur_kv_head * stride_buf_kh - + offs_d_packed[None, :] + + k_byte_offsets[None, :] ) k_quant_packed = tl.load( K_Buffer + offs_buf_k_packed, @@ -1049,9 +1066,7 @@ def _fwd_kernel_stage1_quant_int2( safe_group_k_q1 = tl.where(mask_d_packed, offs_group_k_q1, 0) safe_group_k_q2 = tl.where(mask_d_packed, offs_group_k_q2, 0) safe_group_k_q3 = tl.where(mask_d_packed, offs_group_k_q3, 0) - offs_sz_k = ( - kv_loc[:, None] * stride_sz_kbs + cur_kv_head * stride_sz_kh - ) + offs_sz_k = kv_loc[:, None] * stride_sz_kbs + cur_kv_head * stride_sz_kh k_scale_q0 = tl.load( K_Scales_Zeros + offs_sz_k + 2 * safe_group_k_q0[None, :], mask=(offs_n[:, None] < split_kv_end) & (mask_d_packed[None, :]), @@ -1092,23 +1107,34 @@ def _fwd_kernel_stage1_quant_int2( mask=(offs_n[:, None] < split_kv_end) & (mask_d_packed[None, :]), other=0.0, ) - # Dequantize INT2 K inline: unpack 4 crumbs and dequantize per-group. - k_q0 = ( - ((k_quant_packed & 0x03).to(tl.float32) - k_zero_q0) - * k_scale_q0 - ).to(q_q0.dtype) - k_q1 = ( - (((k_quant_packed >> 2) & 0x03).to(tl.float32) - k_zero_q1) - * k_scale_q1 - ).to(q_q0.dtype) - k_q2 = ( - (((k_quant_packed >> 4) & 0x03).to(tl.float32) - k_zero_q2) - * k_scale_q2 - ).to(q_q0.dtype) - k_q3 = ( - (((k_quant_packed >> 6) & 0x03).to(tl.float32) - k_zero_q3) - * k_scale_q3 - ).to(q_q0.dtype) + if INT1: + k_q0_raw = (k_quant_packed >> k_bit_in_quarter[None, :]) & 0x01 + k_q1_raw = ( + k_quant_packed >> (2 + k_bit_in_quarter[None, :]) + ) & 0x01 + k_q2_raw = ( + k_quant_packed >> (4 + k_bit_in_quarter[None, :]) + ) & 0x01 + k_q3_raw = ( + k_quant_packed >> (6 + k_bit_in_quarter[None, :]) + ) & 0x01 + else: + k_q0_raw = k_quant_packed & 0x03 + k_q1_raw = (k_quant_packed >> 2) & 0x03 + k_q2_raw = (k_quant_packed >> 4) & 0x03 + k_q3_raw = (k_quant_packed >> 6) & 0x03 + k_q0 = ((k_q0_raw.to(tl.float32) - k_zero_q0) * k_scale_q0).to( + q_q0.dtype + ) + k_q1 = ((k_q1_raw.to(tl.float32) - k_zero_q1) * k_scale_q1).to( + q_q0.dtype + ) + k_q2 = ((k_q2_raw.to(tl.float32) - k_zero_q2) * k_scale_q2).to( + q_q0.dtype + ) + k_q3 = ((k_q3_raw.to(tl.float32) - k_zero_q3) * k_scale_q3).to( + q_q0.dtype + ) else: offs_sz_k_1d = kv_loc * stride_sz_kbs + cur_kv_head * stride_sz_kh k_scale_1d = tl.load( @@ -1121,21 +1147,33 @@ def _fwd_kernel_stage1_quant_int2( mask=offs_n < split_kv_end, other=0.0, ) + if INT1: + k_q0_raw = (k_quant_packed >> k_bit_in_quarter[None, :]) & 0x01 + k_q1_raw = ( + k_quant_packed >> (2 + k_bit_in_quarter[None, :]) + ) & 0x01 + k_q2_raw = ( + k_quant_packed >> (4 + k_bit_in_quarter[None, :]) + ) & 0x01 + k_q3_raw = ( + k_quant_packed >> (6 + k_bit_in_quarter[None, :]) + ) & 0x01 + else: + k_q0_raw = k_quant_packed & 0x03 + k_q1_raw = (k_quant_packed >> 2) & 0x03 + k_q2_raw = (k_quant_packed >> 4) & 0x03 + k_q3_raw = (k_quant_packed >> 6) & 0x03 k_q0 = ( - ((k_quant_packed & 0x03).to(tl.float32) - k_zero_1d[:, None]) - * k_scale_1d[:, None] + (k_q0_raw.to(tl.float32) - k_zero_1d[:, None]) * k_scale_1d[:, None] ).to(q_q0.dtype) k_q1 = ( - (((k_quant_packed >> 2) & 0x03).to(tl.float32) - k_zero_1d[:, None]) - * k_scale_1d[:, None] + (k_q1_raw.to(tl.float32) - k_zero_1d[:, None]) * k_scale_1d[:, None] ).to(q_q0.dtype) k_q2 = ( - (((k_quant_packed >> 4) & 0x03).to(tl.float32) - k_zero_1d[:, None]) - * k_scale_1d[:, None] + (k_q2_raw.to(tl.float32) - k_zero_1d[:, None]) * k_scale_1d[:, None] ).to(q_q0.dtype) k_q3 = ( - (((k_quant_packed >> 6) & 0x03).to(tl.float32) - k_zero_1d[:, None]) - * k_scale_1d[:, None] + (k_q3_raw.to(tl.float32) - k_zero_1d[:, None]) * k_scale_1d[:, None] ).to(q_q0.dtype) # Compute QK from 4 partial dot products @@ -1155,14 +1193,20 @@ def _fwd_kernel_stage1_quant_int2( qk = tl.where(offs_n < split_kv_end, qk, float("-inf")) - # Load packed INT2 V + # Load packed V into four logical quarters (same layout as K). offs_dv_packed = tl.arange(0, BLOCK_DV // 4) mask_dv_packed = offs_dv_packed < (Lv // 4) + if INT1: + v_byte_offsets = offs_dv_packed % (Lv // 8) + v_bit_in_quarter = offs_dv_packed // (Lv // 8) + else: + v_byte_offsets = offs_dv_packed + v_bit_in_quarter = tl.zeros([BLOCK_DV // 4], dtype=tl.int32) offs_buf_v_packed = ( kv_loc[:, None] * stride_buf_vbs + cur_kv_head * stride_buf_vh - + offs_dv_packed[None, :] + + v_byte_offsets[None, :] ) v_quant_packed = tl.load( V_Buffer + offs_buf_v_packed, @@ -1180,9 +1224,7 @@ def _fwd_kernel_stage1_quant_int2( safe_group_v_q1 = tl.where(mask_dv_packed, offs_group_v_q1, 0) safe_group_v_q2 = tl.where(mask_dv_packed, offs_group_v_q2, 0) safe_group_v_q3 = tl.where(mask_dv_packed, offs_group_v_q3, 0) - offs_sz_v = ( - kv_loc[:, None] * stride_sz_vbs + cur_kv_head * stride_sz_vh - ) + offs_sz_v = kv_loc[:, None] * stride_sz_vbs + cur_kv_head * stride_sz_vh v_scale_q0 = tl.load( V_Scales_Zeros + offs_sz_v + 2 * safe_group_v_q0[None, :], mask=(offs_n[:, None] < split_kv_end) & (mask_dv_packed[None, :]), @@ -1223,23 +1265,34 @@ def _fwd_kernel_stage1_quant_int2( mask=(offs_n[:, None] < split_kv_end) & (mask_dv_packed[None, :]), other=0.0, ) - # Dequantize INT2 V inline: unpack 4 crumbs per-group. - v_q0 = ( - ((v_quant_packed & 0x03).to(tl.float32) - v_zero_q0) - * v_scale_q0 - ).to(q_q0.dtype) - v_q1 = ( - (((v_quant_packed >> 2) & 0x03).to(tl.float32) - v_zero_q1) - * v_scale_q1 - ).to(q_q0.dtype) - v_q2 = ( - (((v_quant_packed >> 4) & 0x03).to(tl.float32) - v_zero_q2) - * v_scale_q2 - ).to(q_q0.dtype) - v_q3 = ( - (((v_quant_packed >> 6) & 0x03).to(tl.float32) - v_zero_q3) - * v_scale_q3 - ).to(q_q0.dtype) + if INT1: + v_q0_raw = (v_quant_packed >> v_bit_in_quarter[None, :]) & 0x01 + v_q1_raw = ( + v_quant_packed >> (2 + v_bit_in_quarter[None, :]) + ) & 0x01 + v_q2_raw = ( + v_quant_packed >> (4 + v_bit_in_quarter[None, :]) + ) & 0x01 + v_q3_raw = ( + v_quant_packed >> (6 + v_bit_in_quarter[None, :]) + ) & 0x01 + else: + v_q0_raw = v_quant_packed & 0x03 + v_q1_raw = (v_quant_packed >> 2) & 0x03 + v_q2_raw = (v_quant_packed >> 4) & 0x03 + v_q3_raw = (v_quant_packed >> 6) & 0x03 + v_q0 = ((v_q0_raw.to(tl.float32) - v_zero_q0) * v_scale_q0).to( + q_q0.dtype + ) + v_q1 = ((v_q1_raw.to(tl.float32) - v_zero_q1) * v_scale_q1).to( + q_q0.dtype + ) + v_q2 = ((v_q2_raw.to(tl.float32) - v_zero_q2) * v_scale_q2).to( + q_q0.dtype + ) + v_q3 = ((v_q3_raw.to(tl.float32) - v_zero_q3) * v_scale_q3).to( + q_q0.dtype + ) else: offs_sz_v_1d = kv_loc * stride_sz_vbs + cur_kv_head * stride_sz_vh v_scale_1d = tl.load( @@ -1252,21 +1305,33 @@ def _fwd_kernel_stage1_quant_int2( mask=offs_n < split_kv_end, other=0.0, ) + if INT1: + v_q0_raw = (v_quant_packed >> v_bit_in_quarter[None, :]) & 0x01 + v_q1_raw = ( + v_quant_packed >> (2 + v_bit_in_quarter[None, :]) + ) & 0x01 + v_q2_raw = ( + v_quant_packed >> (4 + v_bit_in_quarter[None, :]) + ) & 0x01 + v_q3_raw = ( + v_quant_packed >> (6 + v_bit_in_quarter[None, :]) + ) & 0x01 + else: + v_q0_raw = v_quant_packed & 0x03 + v_q1_raw = (v_quant_packed >> 2) & 0x03 + v_q2_raw = (v_quant_packed >> 4) & 0x03 + v_q3_raw = (v_quant_packed >> 6) & 0x03 v_q0 = ( - ((v_quant_packed & 0x03).to(tl.float32) - v_zero_1d[:, None]) - * v_scale_1d[:, None] + (v_q0_raw.to(tl.float32) - v_zero_1d[:, None]) * v_scale_1d[:, None] ).to(q_q0.dtype) v_q1 = ( - (((v_quant_packed >> 2) & 0x03).to(tl.float32) - v_zero_1d[:, None]) - * v_scale_1d[:, None] + (v_q1_raw.to(tl.float32) - v_zero_1d[:, None]) * v_scale_1d[:, None] ).to(q_q0.dtype) v_q2 = ( - (((v_quant_packed >> 4) & 0x03).to(tl.float32) - v_zero_1d[:, None]) - * v_scale_1d[:, None] + (v_q2_raw.to(tl.float32) - v_zero_1d[:, None]) * v_scale_1d[:, None] ).to(q_q0.dtype) v_q3 = ( - (((v_quant_packed >> 6) & 0x03).to(tl.float32) - v_zero_1d[:, None]) - * v_scale_1d[:, None] + (v_q3_raw.to(tl.float32) - v_zero_1d[:, None]) * v_scale_1d[:, None] ).to(q_q0.dtype) n_e_max = tl.maximum(tl.max(qk, 0), e_max) @@ -1395,6 +1460,7 @@ def _fwd_grouped_kernel_stage1_quant_int2( xai_temperature_len: tl.constexpr, L: tl.constexpr, GROUP_SIZE: tl.constexpr, + INT1: tl.constexpr, ): cur_batch = tl.program_id(0) cur_head_id = tl.program_id(1) @@ -1498,14 +1564,20 @@ def _fwd_grouped_kernel_stage1_quant_int2( other=0, ) - # Load packed INT2 K in transposed format for efficient dot product + # Load packed K in transposed logical-quarter format. offs_d_packed = tl.arange(0, BLOCK_D // 4) mask_d_packed = offs_d_packed < (L // 4) + if INT1: + k_byte_offsets = offs_d_packed % (L // 8) + k_bit_in_quarter = offs_d_packed // (L // 8) + else: + k_byte_offsets = offs_d_packed + k_bit_in_quarter = tl.zeros([BLOCK_D // 4], dtype=tl.int32) offs_buf_k_packed = ( kv_loc[None, :] * stride_buf_kbs + cur_kv_head * stride_buf_kh - + offs_d_packed[:, None] + + k_byte_offsets[:, None] ) k_packed = tl.load( K_Buffer + offs_buf_k_packed, @@ -1530,75 +1602,99 @@ def _fwd_grouped_kernel_stage1_quant_int2( ) k_scale_q0_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k[:, None], - mask=offs_n[None, :] < split_kv_end, other=1.0, + mask=offs_n[None, :] < split_kv_end, + other=1.0, ) k_zero_q0_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k[:, None] + 1, - mask=offs_n[None, :] < split_kv_end, other=0.0, + mask=offs_n[None, :] < split_kv_end, + other=0.0, ) k_scale_q1_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q1[:, None], - mask=offs_n[None, :] < split_kv_end, other=1.0, + mask=offs_n[None, :] < split_kv_end, + other=1.0, ) k_zero_q1_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q1[:, None] + 1, - mask=offs_n[None, :] < split_kv_end, other=0.0, + mask=offs_n[None, :] < split_kv_end, + other=0.0, ) k_scale_q2_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q2[:, None], - mask=offs_n[None, :] < split_kv_end, other=1.0, + mask=offs_n[None, :] < split_kv_end, + other=1.0, ) k_zero_q2_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q2[:, None] + 1, - mask=offs_n[None, :] < split_kv_end, other=0.0, + mask=offs_n[None, :] < split_kv_end, + other=0.0, ) k_scale_q3_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q3[:, None], - mask=offs_n[None, :] < split_kv_end, other=1.0, + mask=offs_n[None, :] < split_kv_end, + other=1.0, ) k_zero_q3_grp = tl.load( K_Scales_Zeros + offs_sz_k + 2 * offs_grp_k_q3[:, None] + 1, - mask=offs_n[None, :] < split_kv_end, other=0.0, + mask=offs_n[None, :] < split_kv_end, + other=0.0, ) # Broadcast per-group across GROUP_SIZE dims via reshape. k_scale_q0 = tl.reshape( - tl.broadcast_to(k_scale_q0_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_scale_q0_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_zero_q0 = tl.reshape( - tl.broadcast_to(k_zero_q0_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_zero_q0_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_scale_q1 = tl.reshape( - tl.broadcast_to(k_scale_q1_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_scale_q1_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_zero_q1 = tl.reshape( - tl.broadcast_to(k_zero_q1_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_zero_q1_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_scale_q2 = tl.reshape( - tl.broadcast_to(k_scale_q2_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_scale_q2_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_zero_q2 = tl.reshape( - tl.broadcast_to(k_zero_q2_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_zero_q2_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_scale_q3 = tl.reshape( - tl.broadcast_to(k_scale_q3_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_scale_q3_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) k_zero_q3 = tl.reshape( - tl.broadcast_to(k_zero_q3_grp[:, None, :], - (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N)), + tl.broadcast_to( + k_zero_q3_grp[:, None, :], + (NUM_GROUPS_QUARTER, GROUP_SIZE, BLOCK_N), + ), (BLOCK_D // 4, BLOCK_N), ) else: @@ -1610,46 +1706,94 @@ def _fwd_grouped_kernel_stage1_quant_int2( grp_q1: tl.constexpr = (1 * (BLOCK_D // 4)) // GROUP_SIZE grp_q2: tl.constexpr = (2 * (BLOCK_D // 4)) // GROUP_SIZE grp_q3: tl.constexpr = (3 * (BLOCK_D // 4)) // GROUP_SIZE - k_scale_q0_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q0, - mask=offs_n < split_kv_end, other=1.0) - k_zero_q0_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q0 + 1, - mask=offs_n < split_kv_end, other=0.0) - k_scale_q1_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q1, - mask=offs_n < split_kv_end, other=1.0) - k_zero_q1_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q1 + 1, - mask=offs_n < split_kv_end, other=0.0) - k_scale_q2_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q2, - mask=offs_n < split_kv_end, other=1.0) - k_zero_q2_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q2 + 1, - mask=offs_n < split_kv_end, other=0.0) - k_scale_q3_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q3, - mask=offs_n < split_kv_end, other=1.0) - k_zero_q3_t = tl.load(K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q3 + 1, - mask=offs_n < split_kv_end, other=0.0) - k_scale_q0 = tl.broadcast_to(k_scale_q0_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_zero_q0 = tl.broadcast_to(k_zero_q0_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_scale_q1 = tl.broadcast_to(k_scale_q1_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_zero_q1 = tl.broadcast_to(k_zero_q1_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_scale_q2 = tl.broadcast_to(k_scale_q2_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_zero_q2 = tl.broadcast_to(k_zero_q2_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_scale_q3 = tl.broadcast_to(k_scale_q3_t[None, :], (BLOCK_D // 4, BLOCK_N)) - k_zero_q3 = tl.broadcast_to(k_zero_q3_t[None, :], (BLOCK_D // 4, BLOCK_N)) + k_scale_q0_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q0, + mask=offs_n < split_kv_end, + other=1.0, + ) + k_zero_q0_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q0 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + k_scale_q1_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q1, + mask=offs_n < split_kv_end, + other=1.0, + ) + k_zero_q1_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q1 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + k_scale_q2_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q2, + mask=offs_n < split_kv_end, + other=1.0, + ) + k_zero_q2_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q2 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + k_scale_q3_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q3, + mask=offs_n < split_kv_end, + other=1.0, + ) + k_zero_q3_t = tl.load( + K_Scales_Zeros + offs_sz_k_1d + 2 * grp_q3 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + k_scale_q0 = tl.broadcast_to( + k_scale_q0_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_zero_q0 = tl.broadcast_to( + k_zero_q0_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_scale_q1 = tl.broadcast_to( + k_scale_q1_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_zero_q1 = tl.broadcast_to( + k_zero_q1_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_scale_q2 = tl.broadcast_to( + k_scale_q2_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_zero_q2 = tl.broadcast_to( + k_zero_q2_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_scale_q3 = tl.broadcast_to( + k_scale_q3_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) + k_zero_q3 = tl.broadcast_to( + k_zero_q3_t[None, :], (BLOCK_D // 4, BLOCK_N) + ) # Cast scales/zeros to q's dtype ONCE so the per-element dequant # below stays entirely in bf16 (saves 2 fp32↔bf16 casts per crumb). k_scale_q0 = k_scale_q0.to(q_q0.dtype) - k_zero_q0 = k_zero_q0.to(q_q0.dtype) + k_zero_q0 = k_zero_q0.to(q_q0.dtype) k_scale_q1 = k_scale_q1.to(q_q0.dtype) - k_zero_q1 = k_zero_q1.to(q_q0.dtype) + k_zero_q1 = k_zero_q1.to(q_q0.dtype) k_scale_q2 = k_scale_q2.to(q_q0.dtype) - k_zero_q2 = k_zero_q2.to(q_q0.dtype) + k_zero_q2 = k_zero_q2.to(q_q0.dtype) k_scale_q3 = k_scale_q3.to(q_q0.dtype) - k_zero_q3 = k_zero_q3.to(q_q0.dtype) - # Dequantize INT2 K inline: unpack 4 crumbs per-group. - # k_packed shape: [BLOCK_D//4, BLOCK_N] (transposed) - k_q0 = ((k_packed & 0x03).to(q_q0.dtype) - k_zero_q0) * k_scale_q0 - k_q1 = (((k_packed >> 2) & 0x03).to(q_q0.dtype) - k_zero_q1) * k_scale_q1 - k_q2 = (((k_packed >> 4) & 0x03).to(q_q0.dtype) - k_zero_q2) * k_scale_q2 - k_q3 = (((k_packed >> 6) & 0x03).to(q_q0.dtype) - k_zero_q3) * k_scale_q3 + k_zero_q3 = k_zero_q3.to(q_q0.dtype) + if INT1: + k_q0_raw = (k_packed >> k_bit_in_quarter[:, None]) & 0x01 + k_q1_raw = (k_packed >> (2 + k_bit_in_quarter[:, None])) & 0x01 + k_q2_raw = (k_packed >> (4 + k_bit_in_quarter[:, None])) & 0x01 + k_q3_raw = (k_packed >> (6 + k_bit_in_quarter[:, None])) & 0x01 + else: + k_q0_raw = k_packed & 0x03 + k_q1_raw = (k_packed >> 2) & 0x03 + k_q2_raw = (k_packed >> 4) & 0x03 + k_q3_raw = (k_packed >> 6) & 0x03 + k_q0 = (k_q0_raw.to(q_q0.dtype) - k_zero_q0) * k_scale_q0 + k_q1 = (k_q1_raw.to(q_q0.dtype) - k_zero_q1) * k_scale_q1 + k_q2 = (k_q2_raw.to(q_q0.dtype) - k_zero_q2) * k_scale_q2 + k_q3 = (k_q3_raw.to(q_q0.dtype) - k_zero_q3) * k_scale_q3 else: offs_sz_k_1d = kv_loc * stride_sz_kbs + cur_kv_head * stride_sz_kh k_scale_1d = tl.load( @@ -1662,18 +1806,28 @@ def _fwd_grouped_kernel_stage1_quant_int2( mask=offs_n < split_kv_end, other=0.0, ).to(q_q0.dtype) - k_q0 = ( - (k_packed & 0x03).to(q_q0.dtype) - k_zero_1d[None, :] - ) * k_scale_1d[None, :] - k_q1 = ( - ((k_packed >> 2) & 0x03).to(q_q0.dtype) - k_zero_1d[None, :] - ) * k_scale_1d[None, :] - k_q2 = ( - ((k_packed >> 4) & 0x03).to(q_q0.dtype) - k_zero_1d[None, :] - ) * k_scale_1d[None, :] - k_q3 = ( - ((k_packed >> 6) & 0x03).to(q_q0.dtype) - k_zero_1d[None, :] - ) * k_scale_1d[None, :] + if INT1: + k_q0_raw = (k_packed >> k_bit_in_quarter[:, None]) & 0x01 + k_q1_raw = (k_packed >> (2 + k_bit_in_quarter[:, None])) & 0x01 + k_q2_raw = (k_packed >> (4 + k_bit_in_quarter[:, None])) & 0x01 + k_q3_raw = (k_packed >> (6 + k_bit_in_quarter[:, None])) & 0x01 + else: + k_q0_raw = k_packed & 0x03 + k_q1_raw = (k_packed >> 2) & 0x03 + k_q2_raw = (k_packed >> 4) & 0x03 + k_q3_raw = (k_packed >> 6) & 0x03 + k_q0 = (k_q0_raw.to(q_q0.dtype) - k_zero_1d[None, :]) * k_scale_1d[ + None, : + ] + k_q1 = (k_q1_raw.to(q_q0.dtype) - k_zero_1d[None, :]) * k_scale_1d[ + None, : + ] + k_q2 = (k_q2_raw.to(q_q0.dtype) - k_zero_1d[None, :]) * k_scale_1d[ + None, : + ] + k_q3 = (k_q3_raw.to(q_q0.dtype) - k_zero_1d[None, :]) * k_scale_1d[ + None, : + ] # Compute QK as ONE fused MMA instead of 4 small ones by stacking # the 4 dequantized quarters into a contiguous D axis. @@ -1684,18 +1838,18 @@ def _fwd_grouped_kernel_stage1_quant_int2( # We use tl.join (which adds a new last axis) + tl.reshape to # interleave: [BLOCK_D//4, BLOCK_N] -> [4, BLOCK_D//4, BLOCK_N] # via two binary joins -> permute -> reshape to [BLOCK_D, BLOCK_N]. - k_01 = tl.join(k_q0, k_q1) # [BLOCK_D//4, BLOCK_N, 2] - k_23 = tl.join(k_q2, k_q3) # [BLOCK_D//4, BLOCK_N, 2] - k_full = tl.join(k_01, k_23) # [BLOCK_D//4, BLOCK_N, 2, 2] + k_01 = tl.join(k_q0, k_q1) # [BLOCK_D//4, BLOCK_N, 2] + k_23 = tl.join(k_q2, k_q3) # [BLOCK_D//4, BLOCK_N, 2] + k_full = tl.join(k_01, k_23) # [BLOCK_D//4, BLOCK_N, 2, 2] k_full = tl.reshape(k_full, (BLOCK_D // 4, BLOCK_N, 4)) - k_full = tl.permute(k_full, (2, 0, 1)) # [4, BLOCK_D//4, BLOCK_N] + k_full = tl.permute(k_full, (2, 0, 1)) # [4, BLOCK_D//4, BLOCK_N] k_full = tl.reshape(k_full, (BLOCK_D, BLOCK_N)) - q_01 = tl.join(q_q0, q_q1) # [BLOCK_H, BLOCK_D//4, 2] + q_01 = tl.join(q_q0, q_q1) # [BLOCK_H, BLOCK_D//4, 2] q_23 = tl.join(q_q2, q_q3) - q_full = tl.join(q_01, q_23) # [BLOCK_H, BLOCK_D//4, 2, 2] + q_full = tl.join(q_01, q_23) # [BLOCK_H, BLOCK_D//4, 2, 2] q_full = tl.reshape(q_full, (BLOCK_H, BLOCK_D // 4, 4)) - q_full = tl.permute(q_full, (0, 2, 1)) # [BLOCK_H, 4, BLOCK_D//4] + q_full = tl.permute(q_full, (0, 2, 1)) # [BLOCK_H, 4, BLOCK_D//4] q_full = tl.reshape(q_full, (BLOCK_H, BLOCK_D)) qk = tl.dot(q_full, k_full) @@ -1712,24 +1866,31 @@ def _fwd_grouped_kernel_stage1_quant_int2( mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") ) - # Load packed INT2 V and dequantize. V layout: [BLOCK_N, BLOCK_D//4] + # Load packed V into four logical quarters. offs_d_packed_v = tl.arange(0, BLOCK_D // 4) + if INT1: + v_byte_offsets = offs_d_packed_v % (L // 8) + v_bit_in_quarter = offs_d_packed_v // (L // 8) + else: + v_byte_offsets = offs_d_packed_v + v_bit_in_quarter = tl.zeros([BLOCK_D // 4], dtype=tl.int32) offs_buf_v_packed = ( kv_loc[:, None] * stride_buf_vbs + cur_kv_head * stride_buf_vh - + offs_d_packed_v[None, :] + + v_byte_offsets[None, :] ) v_packed = tl.load( V_Buffer + offs_buf_v_packed, - mask=(offs_n[:, None] < split_kv_end) & (offs_d_packed_v[None, :] < (L // 4)), + mask=(offs_n[:, None] < split_kv_end) + & (offs_d_packed_v[None, :] < (L // 4)), other=0, ) # Load V scales and zeros for dequantization if GROUPED: if FAST: - NUM_GROUPS_QUARTER: tl.constexpr = (BLOCK_D // 4) // GROUP_SIZE - offs_grp_v = tl.arange(0, NUM_GROUPS_QUARTER) + NUM_GROUPS_QUARTER_V: tl.constexpr = (BLOCK_D // 4) // GROUP_SIZE + offs_grp_v = tl.arange(0, NUM_GROUPS_QUARTER_V) offs_grp_v_q1 = (BLOCK_D // 4) // GROUP_SIZE + offs_grp_v offs_grp_v_q2 = 2 * (BLOCK_D // 4) // GROUP_SIZE + offs_grp_v offs_grp_v_q3 = 3 * (BLOCK_D // 4) // GROUP_SIZE + offs_grp_v @@ -1738,74 +1899,98 @@ def _fwd_grouped_kernel_stage1_quant_int2( ) v_scale_q0_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v[None, :], - mask=offs_n[:, None] < split_kv_end, other=1.0, + mask=offs_n[:, None] < split_kv_end, + other=1.0, ) v_zero_q0_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v[None, :] + 1, - mask=offs_n[:, None] < split_kv_end, other=0.0, + mask=offs_n[:, None] < split_kv_end, + other=0.0, ) v_scale_q1_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q1[None, :], - mask=offs_n[:, None] < split_kv_end, other=1.0, + mask=offs_n[:, None] < split_kv_end, + other=1.0, ) v_zero_q1_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q1[None, :] + 1, - mask=offs_n[:, None] < split_kv_end, other=0.0, + mask=offs_n[:, None] < split_kv_end, + other=0.0, ) v_scale_q2_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q2[None, :], - mask=offs_n[:, None] < split_kv_end, other=1.0, + mask=offs_n[:, None] < split_kv_end, + other=1.0, ) v_zero_q2_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q2[None, :] + 1, - mask=offs_n[:, None] < split_kv_end, other=0.0, + mask=offs_n[:, None] < split_kv_end, + other=0.0, ) v_scale_q3_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q3[None, :], - mask=offs_n[:, None] < split_kv_end, other=1.0, + mask=offs_n[:, None] < split_kv_end, + other=1.0, ) v_zero_q3_grp = tl.load( V_Scales_Zeros + offs_sz_v + 2 * offs_grp_v_q3[None, :] + 1, - mask=offs_n[:, None] < split_kv_end, other=0.0, + mask=offs_n[:, None] < split_kv_end, + other=0.0, ) v_scale_q0 = tl.reshape( - tl.broadcast_to(v_scale_q0_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_scale_q0_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_zero_q0 = tl.reshape( - tl.broadcast_to(v_zero_q0_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_zero_q0_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_scale_q1 = tl.reshape( - tl.broadcast_to(v_scale_q1_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_scale_q1_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_zero_q1 = tl.reshape( - tl.broadcast_to(v_zero_q1_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_zero_q1_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_scale_q2 = tl.reshape( - tl.broadcast_to(v_scale_q2_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_scale_q2_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_zero_q2 = tl.reshape( - tl.broadcast_to(v_zero_q2_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_zero_q2_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_scale_q3 = tl.reshape( - tl.broadcast_to(v_scale_q3_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_scale_q3_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) v_zero_q3 = tl.reshape( - tl.broadcast_to(v_zero_q3_grp[:, :, None], - (BLOCK_N, NUM_GROUPS_QUARTER, GROUP_SIZE)), + tl.broadcast_to( + v_zero_q3_grp[:, :, None], + (BLOCK_N, NUM_GROUPS_QUARTER_V, GROUP_SIZE), + ), (BLOCK_N, BLOCK_D // 4), ) else: @@ -1815,45 +2000,94 @@ def _fwd_grouped_kernel_stage1_quant_int2( v_grp_q1: tl.constexpr = (1 * (BLOCK_D // 4)) // GROUP_SIZE v_grp_q2: tl.constexpr = (2 * (BLOCK_D // 4)) // GROUP_SIZE v_grp_q3: tl.constexpr = (3 * (BLOCK_D // 4)) // GROUP_SIZE - v_scale_q0_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q0, - mask=offs_n < split_kv_end, other=1.0) - v_zero_q0_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q0 + 1, - mask=offs_n < split_kv_end, other=0.0) - v_scale_q1_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q1, - mask=offs_n < split_kv_end, other=1.0) - v_zero_q1_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q1 + 1, - mask=offs_n < split_kv_end, other=0.0) - v_scale_q2_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q2, - mask=offs_n < split_kv_end, other=1.0) - v_zero_q2_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q2 + 1, - mask=offs_n < split_kv_end, other=0.0) - v_scale_q3_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q3, - mask=offs_n < split_kv_end, other=1.0) - v_zero_q3_t = tl.load(V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q3 + 1, - mask=offs_n < split_kv_end, other=0.0) - v_scale_q0 = tl.broadcast_to(v_scale_q0_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_zero_q0 = tl.broadcast_to(v_zero_q0_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_scale_q1 = tl.broadcast_to(v_scale_q1_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_zero_q1 = tl.broadcast_to(v_zero_q1_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_scale_q2 = tl.broadcast_to(v_scale_q2_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_zero_q2 = tl.broadcast_to(v_zero_q2_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_scale_q3 = tl.broadcast_to(v_scale_q3_t[:, None], (BLOCK_N, BLOCK_D // 4)) - v_zero_q3 = tl.broadcast_to(v_zero_q3_t[:, None], (BLOCK_N, BLOCK_D // 4)) + v_scale_q0_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q0, + mask=offs_n < split_kv_end, + other=1.0, + ) + v_zero_q0_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q0 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + v_scale_q1_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q1, + mask=offs_n < split_kv_end, + other=1.0, + ) + v_zero_q1_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q1 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + v_scale_q2_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q2, + mask=offs_n < split_kv_end, + other=1.0, + ) + v_zero_q2_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q2 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + v_scale_q3_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q3, + mask=offs_n < split_kv_end, + other=1.0, + ) + v_zero_q3_t = tl.load( + V_Scales_Zeros + offs_sz_v_1d + 2 * v_grp_q3 + 1, + mask=offs_n < split_kv_end, + other=0.0, + ) + v_scale_q0 = tl.broadcast_to( + v_scale_q0_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_zero_q0 = tl.broadcast_to( + v_zero_q0_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_scale_q1 = tl.broadcast_to( + v_scale_q1_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_zero_q1 = tl.broadcast_to( + v_zero_q1_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_scale_q2 = tl.broadcast_to( + v_scale_q2_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_zero_q2 = tl.broadcast_to( + v_zero_q2_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_scale_q3 = tl.broadcast_to( + v_scale_q3_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) + v_zero_q3 = tl.broadcast_to( + v_zero_q3_t[:, None], (BLOCK_N, BLOCK_D // 4) + ) # Cast V scales/zeros to q's dtype ONCE so per-element dequant # below stays in bf16 (saves 2 fp32↔bf16 casts per crumb). v_scale_q0 = v_scale_q0.to(q_q0.dtype) - v_zero_q0 = v_zero_q0.to(q_q0.dtype) + v_zero_q0 = v_zero_q0.to(q_q0.dtype) v_scale_q1 = v_scale_q1.to(q_q0.dtype) - v_zero_q1 = v_zero_q1.to(q_q0.dtype) + v_zero_q1 = v_zero_q1.to(q_q0.dtype) v_scale_q2 = v_scale_q2.to(q_q0.dtype) - v_zero_q2 = v_zero_q2.to(q_q0.dtype) + v_zero_q2 = v_zero_q2.to(q_q0.dtype) v_scale_q3 = v_scale_q3.to(q_q0.dtype) - v_zero_q3 = v_zero_q3.to(q_q0.dtype) - # Dequantize INT2 V inline: unpack 4 crumbs per-group. - v_q0 = ((v_packed & 0x03).to(q_q0.dtype) - v_zero_q0) * v_scale_q0 - v_q1 = (((v_packed >> 2) & 0x03).to(q_q0.dtype) - v_zero_q1) * v_scale_q1 - v_q2 = (((v_packed >> 4) & 0x03).to(q_q0.dtype) - v_zero_q2) * v_scale_q2 - v_q3 = (((v_packed >> 6) & 0x03).to(q_q0.dtype) - v_zero_q3) * v_scale_q3 + v_zero_q3 = v_zero_q3.to(q_q0.dtype) + if INT1: + v_q0_raw = (v_packed >> v_bit_in_quarter[None, :]) & 0x01 + v_q1_raw = (v_packed >> (2 + v_bit_in_quarter[None, :])) & 0x01 + v_q2_raw = (v_packed >> (4 + v_bit_in_quarter[None, :])) & 0x01 + v_q3_raw = (v_packed >> (6 + v_bit_in_quarter[None, :])) & 0x01 + else: + v_q0_raw = v_packed & 0x03 + v_q1_raw = (v_packed >> 2) & 0x03 + v_q2_raw = (v_packed >> 4) & 0x03 + v_q3_raw = (v_packed >> 6) & 0x03 + v_q0 = (v_q0_raw.to(q_q0.dtype) - v_zero_q0) * v_scale_q0 + v_q1 = (v_q1_raw.to(q_q0.dtype) - v_zero_q1) * v_scale_q1 + v_q2 = (v_q2_raw.to(q_q0.dtype) - v_zero_q2) * v_scale_q2 + v_q3 = (v_q3_raw.to(q_q0.dtype) - v_zero_q3) * v_scale_q3 else: offs_sz_v_1d = kv_loc * stride_sz_vbs + cur_kv_head * stride_sz_vh v_scale_1d = tl.load( @@ -1866,18 +2100,28 @@ def _fwd_grouped_kernel_stage1_quant_int2( mask=offs_n < split_kv_end, other=0.0, ).to(q_q0.dtype) - v_q0 = ( - (v_packed & 0x03).to(q_q0.dtype) - v_zero_1d[:, None] - ) * v_scale_1d[:, None] - v_q1 = ( - ((v_packed >> 2) & 0x03).to(q_q0.dtype) - v_zero_1d[:, None] - ) * v_scale_1d[:, None] - v_q2 = ( - ((v_packed >> 4) & 0x03).to(q_q0.dtype) - v_zero_1d[:, None] - ) * v_scale_1d[:, None] - v_q3 = ( - ((v_packed >> 6) & 0x03).to(q_q0.dtype) - v_zero_1d[:, None] - ) * v_scale_1d[:, None] + if INT1: + v_q0_raw = (v_packed >> v_bit_in_quarter[None, :]) & 0x01 + v_q1_raw = (v_packed >> (2 + v_bit_in_quarter[None, :])) & 0x01 + v_q2_raw = (v_packed >> (4 + v_bit_in_quarter[None, :])) & 0x01 + v_q3_raw = (v_packed >> (6 + v_bit_in_quarter[None, :])) & 0x01 + else: + v_q0_raw = v_packed & 0x03 + v_q1_raw = (v_packed >> 2) & 0x03 + v_q2_raw = (v_packed >> 4) & 0x03 + v_q3_raw = (v_packed >> 6) & 0x03 + v_q0 = (v_q0_raw.to(q_q0.dtype) - v_zero_1d[:, None]) * v_scale_1d[ + :, None + ] + v_q1 = (v_q1_raw.to(q_q0.dtype) - v_zero_1d[:, None]) * v_scale_1d[ + :, None + ] + v_q2 = (v_q2_raw.to(q_q0.dtype) - v_zero_1d[:, None]) * v_scale_1d[ + :, None + ] + v_q3 = (v_q3_raw.to(q_q0.dtype) - v_zero_1d[:, None]) * v_scale_1d[ + :, None + ] n_e_max = tl.maximum(tl.max(qk, 1), e_max) re_scale = tl.exp(e_max - n_e_max) @@ -1955,6 +2199,7 @@ def _decode_att_m_fwd_quant_int2( sm_scale, logit_cap, xai_temperature_len=-1, + int1=False, ): """ INT2 quantized KV cache attention wrapper (MHA). @@ -1965,10 +2210,9 @@ def _decode_att_m_fwd_quant_int2( if _is_hip: BLOCK = 8 MAX_KV_SPLITS = max_kv_splits - # For INT2, the buffer stores packed values (head_dim//4) - # But we need to work with the actual head_dim - Lk = k_buffer.shape[-1] * 4 # Unpack to get real dimension - Lv = v_buffer.shape[-1] * 4 + pack_factor = 8 if int1 else 4 + Lk = k_buffer.shape[-1] * pack_factor + Lv = v_buffer.shape[-1] * pack_factor batch, head_num = q.shape[0], q.shape[1] @@ -1984,9 +2228,7 @@ def _decode_att_m_fwd_quant_int2( BLOCK_DMODEL = triton.next_power_of_2(Lk) BLOCK_DV = triton.next_power_of_2(Lv) - group_size = _get_shared_kv_scale_group_size( - Lk, Lv, k_scales_zeros, v_scales_zeros - ) + group_size = _get_shared_kv_scale_group_size(Lk, Lv, k_scales_zeros, v_scales_zeros) _fwd_kernel_stage1_quant_int2[grid]( q, @@ -2025,6 +2267,7 @@ def _decode_att_m_fwd_quant_int2( Lk=Lk, Lv=Lv, GROUP_SIZE=group_size, + INT1=int1, ) @@ -2043,6 +2286,7 @@ def _decode_grouped_att_m_fwd_quant_int2( sm_scale, logit_cap, xai_temperature_len=-1, + int1=False, ): """ INT2 quantized KV cache attention wrapper (GQA/MQA). @@ -2058,21 +2302,18 @@ def _decode_grouped_att_m_fwd_quant_int2( crumb → mask/shift → cast → sub zero → mul scale → tl.dot) over more KV tokens; smaller BLOCK_H lowers register pressure so more blocks fit per SM. """ - # For INT2, k_buffer is packed, so actual head dim is 4x the last dimension. - # K and V share the same head dim in this path (no MLA/DPE split). - L = k_buffer.shape[-1] * 4 - assert v_buffer.shape[-1] * 4 == L, "INT2 KV cache requires Lk == Lv" + pack_factor = 8 if int1 else 4 + L = k_buffer.shape[-1] * pack_factor + assert v_buffer.shape[-1] * pack_factor == L, "Quantized KV cache requires Lk == Lv" BLOCK_D = triton.next_power_of_2(L) - group_size = _get_shared_kv_scale_group_size( - L, L, k_scales_zeros, v_scales_zeros - ) + group_size = _get_shared_kv_scale_group_size(L, L, k_scales_zeros, v_scales_zeros) batch, head_num = q.shape[0], q.shape[1] kv_group_num = q.shape[1] // k_buffer.shape[1] MAX_KV_SPLITS = max_kv_splits - # Tile heuristic + # Tile heuristic if kv_group_num <= 8: if batch >= 16: _bn_default, _bh_default, _nw_default = 32, 4, 1 @@ -2085,7 +2326,15 @@ def _decode_grouped_att_m_fwd_quant_int2( _bh_default = 16 if batch >= 16 else 8 _nw_default = 4 BLOCK = int(os.environ.get("SGL_INT2_BLOCK_N", _bn_default)) - BLOCK_H = int(os.environ.get("SGL_INT2_BLOCK_H", _bh_default)) + requested_block_h = max(1, int(os.environ.get("SGL_INT2_BLOCK_H", _bh_default))) + if requested_block_h >= kv_group_num: + BLOCK_H = triton.next_power_of_2(kv_group_num) + else: + BLOCK_H = triton.next_power_of_2(requested_block_h) + if BLOCK_H > requested_block_h: + BLOCK_H //= 2 + while BLOCK_H > 1 and kv_group_num % BLOCK_H != 0: + BLOCK_H //= 2 num_warps = int(os.environ.get("SGL_INT2_NUM_WARPS", _nw_default)) num_stages = int(os.environ.get("SGL_INT2_NUM_STAGES", 3)) @@ -2137,6 +2386,7 @@ def _decode_grouped_att_m_fwd_quant_int2( num_stages=num_stages, L=L, GROUP_SIZE=group_size, + INT1=int1, **extra_kargs, ) @@ -2463,7 +2713,921 @@ def decode_attention_fwd_int2_unified( ) if quant_kv_indices.numel() > 0: + # The grouped kernel uses one shared D axis. Fall back to the + # per-query-head kernel for legal attention layouts with Lk != Lv. + if kv_group_num == 1 or quant_k_buffer.shape[-1] != quant_v_buffer.shape[-1]: + _decode_att_m_fwd_quant_int2( + q, + quant_k_buffer, + quant_v_buffer, + quant_k_scales_zeros, + quant_v_scales_zeros, + quant_logits, + quant_lse, + quant_kv_indptr, + quant_kv_indices, + quant_num_kv_splits, + quant_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + else: + _decode_grouped_att_m_fwd_quant_int2( + q, + quant_k_buffer, + quant_v_buffer, + quant_k_scales_zeros, + quant_v_scales_zeros, + quant_logits, + quant_lse, + quant_kv_indptr, + quant_kv_indices, + quant_num_kv_splits, + quant_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + + _unified_stage2( + attn_logits, + attn_lse, + o, + total_splits=total_splits, + ) + return o + + +@triton.jit +def _pq_build_lut_kernel( + Q, + Codebook, + Lut, + HEAD_DIM: tl.constexpr, + N_SUB: tl.constexpr, + SUB_DIM: tl.constexpr, + N_CENTROIDS: tl.constexpr, +): + batch_q_head = tl.program_id(0) + sub = tl.program_id(1) + centroids = tl.arange(0, N_CENTROIDS) + acc = tl.zeros([N_CENTROIDS], dtype=tl.float32) + for dim in tl.static_range(SUB_DIM): + q_value = tl.load(Q + batch_q_head * HEAD_DIM + sub * SUB_DIM + dim).to( + tl.float32 + ) + centroid = tl.load( + Codebook + (sub * N_CENTROIDS + centroids) * SUB_DIM + dim + ).to(tl.float32) + acc += q_value * centroid + tl.store( + Lut + (batch_q_head * N_SUB + sub) * N_CENTROIDS + centroids, + acc, + ) + + +def _build_pq_lut(q: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor: + batch, q_heads, head_dim = q.shape + n_sub, n_centroids, sub_dim = codebook.shape + assert q.is_contiguous() and codebook.is_contiguous() + assert head_dim == n_sub * sub_dim + lut = torch.empty( + (batch, q_heads, n_sub, n_centroids), + dtype=torch.float32, + device=q.device, + ) + _pq_build_lut_kernel[(batch * q_heads, n_sub)]( + q, + codebook, + lut, + HEAD_DIM=head_dim, + N_SUB=int(n_sub), + SUB_DIM=int(sub_dim), + N_CENTROIDS=int(n_centroids), + num_warps=8, + num_stages=2, + ) + return lut + + +@triton.jit +def _fwd_grouped_kernel_stage1_pq( + Q, + K_Codes, + V_Buffer, + V_Scales_Zeros, + K_Codebook, + K_Lut, + K_Codes2, + K_Codebook2, + K_Lut2, + V_Codebook, + sm_scale, + kv_indptr, + kv_indices, + Att_Out, + Att_Lse, + num_kv_splits, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_ks, + stride_lut_b, + stride_lut_h, + stride_lut_sub, + stride_lut_centroid, + stride_k2bs, + stride_k2h, + stride_k2s, + stride_lut2_b, + stride_lut2_h, + stride_lut2_sub, + stride_lut2_centroid, + stride_vbs, + stride_vh, + stride_vs, + stride_vszbs, + stride_vszh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_DK: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + logit_cap: tl.constexpr, + xai_temperature_len: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, + K_SUB_DIM: tl.constexpr, + K_N_CENTROIDS: tl.constexpr, + K2_N_CENTROIDS: tl.constexpr, + V_SUB_DIM: tl.constexpr, + V_N_CENTROIDS: tl.constexpr, + V_GROUP_SIZE: tl.constexpr, + HAS_K_STAGE2: tl.constexpr, + V_IS_PQ: tl.constexpr, + USE_ADC: tl.constexpr, +): + """Graph-safe PQ-K attention with optional RVQ-K and PQ-V. + + The kernel follows the tier-local indptr directly. No gathered tensor has + a data-dependent Python shape, so padded CUDA Graph index buffers are safe. + """ + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + split_kv_id = tl.program_id(2) + + if BLOCK_H < kv_group_num: + VALID_BLOCK_H: tl.constexpr = BLOCK_H + else: + VALID_BLOCK_H: tl.constexpr = kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = (cur_head < (cur_head_id + 1) * VALID_BLOCK_H) & (cur_head < q_head_num) + cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + + offs_dk = tl.arange(0, BLOCK_DK) + offs_dv = tl.arange(0, BLOCK_DV) + mask_dk = offs_dk < Lk + mask_dv = offs_dv < Lv + q = tl.load( + Q + cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dk[None, :], + mask=mask_h[:, None] & mask_dk[None, :], + other=0.0, + ) + + batch_kv_start = tl.load(kv_indptr + cur_batch) + seq_len = tl.load(kv_indptr + cur_batch + 1) - batch_kv_start + kv_splits = tl.load(num_kv_splits + cur_batch) + kv_len_per_split = tl.cdiv(tl.cdiv(seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + split_start = kv_len_per_split * split_kv_id + split_end = tl.minimum(split_start + kv_len_per_split, seq_len) + + if xai_temperature_len > 0: + offs_qidx = seq_len - 1 + xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len)) + qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale + xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, qtemp, 1.0) + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + if split_end > split_start: + k_sub = offs_dk // K_SUB_DIM + k_sub_off = offs_dk % K_SUB_DIM + if not V_IS_PQ: + v_byte_dim: tl.constexpr = Lv // 4 + v_byte_off = offs_dv % v_byte_dim + v_shift = (offs_dv // v_byte_dim) * 2 + v_group = offs_dv // V_GROUP_SIZE + + for start_n in range(split_start, split_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + valid_n = offs_n < split_end + kv_loc = tl.load( + kv_indices + batch_kv_start + offs_n, + mask=valid_n, + other=0, + ).to(tl.int64) + + if USE_ADC: + qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32) + for sub in tl.static_range(Lk // K_SUB_DIM): + code = tl.load( + K_Codes + + kv_loc * stride_kbs + + cur_kv_head * stride_kh + + sub * stride_ks, + mask=valid_n, + other=0, + ).to(tl.int32) + qk += tl.load( + K_Lut + + cur_batch * stride_lut_b + + cur_head[:, None] * stride_lut_h + + sub * stride_lut_sub + + code[None, :] * stride_lut_centroid, + mask=mask_h[:, None] & valid_n[None, :], + other=0.0, + ) + if HAS_K_STAGE2: + code2 = tl.load( + K_Codes2 + + kv_loc * stride_k2bs + + cur_kv_head * stride_k2h + + sub * stride_k2s, + mask=valid_n, + other=0, + ).to(tl.int32) + qk += tl.load( + K_Lut2 + + cur_batch * stride_lut2_b + + cur_head[:, None] * stride_lut2_h + + sub * stride_lut2_sub + + code2[None, :] * stride_lut2_centroid, + mask=mask_h[:, None] & valid_n[None, :], + other=0.0, + ) + qk *= sm_scale + else: + k_code = tl.load( + K_Codes + + kv_loc[None, :] * stride_kbs + + cur_kv_head * stride_kh + + k_sub[:, None] * stride_ks, + mask=mask_dk[:, None] & valid_n[None, :], + other=0, + ).to(tl.int64) + k = tl.load( + K_Codebook + + (k_sub[:, None] * K_N_CENTROIDS + k_code) * K_SUB_DIM + + k_sub_off[:, None], + mask=mask_dk[:, None] & valid_n[None, :], + other=0.0, + ).to(q.dtype) + if HAS_K_STAGE2: + k_code2 = tl.load( + K_Codes2 + + kv_loc[None, :] * stride_k2bs + + cur_kv_head * stride_k2h + + k_sub[:, None] * stride_k2s, + mask=mask_dk[:, None] & valid_n[None, :], + other=0, + ).to(tl.int64) + k += tl.load( + K_Codebook2 + + (k_sub[:, None] * K2_N_CENTROIDS + k_code2) * K_SUB_DIM + + k_sub_off[:, None], + mask=mask_dk[:, None] & valid_n[None, :], + other=0.0, + ).to(q.dtype) + qk = tl.dot(q, k) * sm_scale + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + if xai_temperature_len > 0: + qk *= xai_temperature_reg + qk = tl.where(mask_h[:, None] & valid_n[None, :], qk, float("-inf")) + + next_e_max = tl.maximum(tl.max(qk, axis=1), e_max) + rescale = tl.exp(e_max - next_e_max) + p = tl.exp(qk - next_e_max[:, None]) + acc *= rescale[:, None] + + if V_IS_PQ and BLOCK_DV == Lv: + # Decode and accumulate one 8-D product-quantization subspace + # at a time. This avoids materializing [BLOCK_N, 128] V in a + # single program and loads each code byte only once. + v_dim = tl.arange(0, V_SUB_DIM) + v_sub_ids = tl.arange(0, Lv // V_SUB_DIM) + for sub in tl.static_range(Lv // V_SUB_DIM): + code = tl.load( + V_Buffer + + kv_loc * stride_vbs + + cur_kv_head * stride_vh + + sub * stride_vs, + mask=valid_n, + other=0, + ).to(tl.int64) + values = tl.load( + V_Codebook + + (sub * V_N_CENTROIDS + code[:, None]) * V_SUB_DIM + + v_dim[None, :], + mask=valid_n[:, None], + other=0.0, + ).to(q.dtype) + partial = tl.dot(p.to(values.dtype), values) + partial_full = tl.broadcast_to( + partial[:, None, :], + ( + BLOCK_H, + Lv // V_SUB_DIM, + V_SUB_DIM, + ), + ) + sub_mask = v_sub_ids[None, :, None] == sub + acc += tl.reshape( + tl.where(sub_mask, partial_full, 0.0), + (BLOCK_H, BLOCK_DV), + ) + elif V_IS_PQ: + v_sub = offs_dv // V_SUB_DIM + v_sub_off = offs_dv % V_SUB_DIM + v_code = tl.load( + V_Buffer + + kv_loc[:, None] * stride_vbs + + cur_kv_head * stride_vh + + v_sub[None, :] * stride_vs, + mask=valid_n[:, None] & mask_dv[None, :], + other=0, + ).to(tl.int64) + v = tl.load( + V_Codebook + + (v_sub[None, :] * V_N_CENTROIDS + v_code) * V_SUB_DIM + + v_sub_off[None, :], + mask=valid_n[:, None] & mask_dv[None, :], + other=0.0, + ).to(q.dtype) + acc += tl.dot(p.to(v.dtype), v) + else: + v_packed = tl.load( + V_Buffer + + kv_loc[:, None] * stride_vbs + + cur_kv_head * stride_vh + + v_byte_off[None, :] * stride_vs, + mask=valid_n[:, None] & mask_dv[None, :], + other=0, + ) + v_q = ((v_packed >> v_shift[None, :]) & 0x03).to(tl.float32) + v_sz_base = kv_loc[:, None] * stride_vszbs + cur_kv_head * stride_vszh + v_scale = tl.load( + V_Scales_Zeros + v_sz_base + 2 * v_group[None, :], + mask=valid_n[:, None] & mask_dv[None, :], + other=1.0, + ).to(tl.float32) + v_zero = tl.load( + V_Scales_Zeros + v_sz_base + 2 * v_group[None, :] + 1, + mask=valid_n[:, None] & mask_dv[None, :], + other=0.0, + ).to(tl.float32) + v = ((v_q - v_zero) * v_scale).to(q.dtype) + acc += tl.dot(p.to(v.dtype), v) + e_sum = e_sum * rescale + tl.sum(p, axis=1) + e_max = next_e_max + + out_base = ( + cur_batch * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + ) + tl.store( + Att_Out + out_base + offs_dv[None, :], + acc / e_sum[:, None], + mask=mask_h[:, None] & mask_dv[None, :], + ) + lse_off = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + ) // Lv + tl.store( + Att_Lse + lse_off, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +def _decode_grouped_att_m_fwd_pq( + q, + k_codes, + v_buffer, + v_scales_zeros, + k_codebook, + att_out, + att_lse, + kv_indptr, + kv_indices, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len=-1, + k_codes2=None, + k_codebook2=None, + v_codebook=None, +): + """Launch graph-safe inline PQ/RVQ decode for MHA, GQA, or MQA.""" + Lk = int(q.shape[-1]) + Lv = int(att_out.shape[-1]) + k_n_sub, k_n_centroids, k_sub_dim = k_codebook.shape + assert int(k_n_sub) * int(k_sub_dim) == Lk + assert k_codes.shape[-1] == k_n_sub + + has_k_stage2 = k_codes2 is not None + if has_k_stage2: + assert k_codebook2 is not None + assert k_codebook2.shape[0] == k_n_sub + assert k_codebook2.shape[2] == k_sub_dim + k_codes2_arg = k_codes2 + k_codebook2_arg = k_codebook2 + k2_n_centroids = int(k_codebook2.shape[1]) + else: + k_codes2_arg = k_codes + k_codebook2_arg = k_codebook + k2_n_centroids = int(k_n_centroids) + + batch, q_head_num = q.shape[0], q.shape[1] + adc_mode = envs.SGLANG_PQ_USE_ADC.get() + use_adc = batch < 4 if adc_mode < 0 else adc_mode != 0 + if use_adc: + k_lut = _build_pq_lut(q, k_codebook) + k_lut2 = _build_pq_lut(q, k_codebook2_arg) if has_k_stage2 else k_lut + lut_strides = k_lut.stride() + lut2_strides = k_lut2.stride() + else: + # Pointer/stride placeholders for the compile-time disabled ADC branch. + k_lut = k_codebook + k_lut2 = k_codebook2_arg + lut_strides = (0, 0, k_codebook.stride(0), k_codebook.stride(1)) + lut2_strides = ( + 0, + 0, + k_codebook2_arg.stride(0), + k_codebook2_arg.stride(1), + ) + + v_is_pq = v_codebook is not None + if v_is_pq: + v_n_sub, v_n_centroids, v_sub_dim = v_codebook.shape + assert int(v_n_sub) * int(v_sub_dim) == Lv + assert v_buffer.shape[-1] == v_n_sub + v_codebook_arg = v_codebook + v_group_size = Lv + else: + assert v_buffer.shape[-1] * 4 == Lv + v_n_centroids = 1 + v_sub_dim = 1 + v_codebook_arg = k_codebook + v_num_groups = int(v_scales_zeros.shape[-1]) // 2 + assert Lv % v_num_groups == 0 + v_group_size = Lv // v_num_groups + + kv_group_num = q_head_num // k_codes.shape[1] + configured_block_h = envs.SGLANG_PQ_BLOCK_H.get() + requested_block_h = ( + configured_block_h + if configured_block_h > 0 + else (1 if use_adc and batch < 4 else 4) + ) + if requested_block_h >= kv_group_num: + block_h = triton.next_power_of_2(kv_group_num) + else: + block_h = 1 << (requested_block_h.bit_length() - 1) + while block_h > 1 and kv_group_num % block_h != 0: + block_h //= 2 + # H100 tuning for Qwen/Llama D=128. Low batch benefits from a larger + # sequence tile (fewer Python-range loop iterations); high batch already + # exposes enough blocks and needs the lower-register BN=32 kernel. + large_tile_safe = max(Lk, Lv) <= 128 and v_is_pq and not has_k_stage2 + if large_tile_safe: + if batch < 16: + default_block_n, default_warps, default_stages = 256, 8, 2 + else: + default_block_n, default_warps, default_stages = 64, 4, 2 + else: + default_block_n, default_warps, default_stages = 32, 4, 2 + configured_block_n = envs.SGLANG_PQ_BLOCK_N.get() + configured_warps = envs.SGLANG_PQ_NUM_WARPS.get() + configured_stages = envs.SGLANG_PQ_NUM_STAGES.get() + requested_block_n = triton.next_power_of_2( + max( + 16, + configured_block_n if configured_block_n > 0 else default_block_n, + ) + ) + max_block_n = 256 if large_tile_safe else 64 + block_n = min(requested_block_n, max_block_n) + requested_warps = configured_warps if configured_warps > 0 else default_warps + num_warps = min(8, triton.next_power_of_2(requested_warps)) + requested_stages = configured_stages if configured_stages > 0 else default_stages + max_stages = 2 if block_n >= 256 else 3 + num_stages = min(max(1, requested_stages), max_stages) + block_dk = triton.next_power_of_2(Lk) + block_dv = triton.next_power_of_2(Lv) + grid = ( + batch, + triton.cdiv(q_head_num, min(block_h, kv_group_num)), + max_kv_splits, + ) + _fwd_grouped_kernel_stage1_pq[grid]( + q, + k_codes, + v_buffer, + v_scales_zeros, + k_codebook, + k_lut, + k_codes2_arg, + k_codebook2_arg, + k_lut2, + v_codebook_arg, + sm_scale, + kv_indptr, + kv_indices, + att_out, + att_lse, + num_kv_splits, + q.stride(0), + q.stride(1), + k_codes.stride(0), + k_codes.stride(1), + k_codes.stride(2), + *lut_strides, + k_codes2_arg.stride(0), + k_codes2_arg.stride(1), + k_codes2_arg.stride(2), + *lut2_strides, + v_buffer.stride(0), + v_buffer.stride(1), + v_buffer.stride(2), + v_scales_zeros.stride(0), + v_scales_zeros.stride(1), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + kv_group_num=kv_group_num, + q_head_num=q_head_num, + BLOCK_DK=block_dk, + BLOCK_DV=block_dv, + BLOCK_N=block_n, + BLOCK_H=block_h, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + logit_cap=logit_cap, + xai_temperature_len=xai_temperature_len, + Lk=Lk, + Lv=Lv, + K_SUB_DIM=int(k_sub_dim), + K_N_CENTROIDS=int(k_n_centroids), + K2_N_CENTROIDS=k2_n_centroids, + V_SUB_DIM=int(v_sub_dim), + V_N_CENTROIDS=int(v_n_centroids), + V_GROUP_SIZE=v_group_size, + HAS_K_STAGE2=has_k_stage2, + V_IS_PQ=v_is_pq, + USE_ADC=use_adc, + num_warps=num_warps, + num_stages=num_stages, + ) + + +def decode_attention_fwd_pqk_int2v_unified( + q, + hp_k_buffer, + hp_v_buffer, + quant_k_buffer, # PQ codes [cache, heads, N_SUB] uint8 + quant_v_buffer, # INT2 V [cache, heads, v_head_dim//4] uint8 + quant_k_scales_zeros, # unused for PQ K (dummy) + quant_v_scales_zeros, + pq_codebook, # [N_SUB, N_CENTS, SUB_DIM] fp16 + o, + hp_kv_indptr, + hp_kv_indices, + quant_kv_indptr, + quant_kv_indices, + attn_logits, + attn_lse, + hp_num_kv_splits, + quant_num_kv_splits, + hp_max_kv_splits, + quant_max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, + xai_temperature_len=-1, + quant_k_buffer2=None, # RVQ stage-2 codes (None = plain PQ K) + pq_codebook2=None, # RVQ stage-2 codebook + pq_v_codebook=None, # PQ V codebook (None = INT2 V) +): + """Unified HP + PQ-K + INT2-V decode attention. + + HP stage: FP16 K/V attention (unchanged from INT2 unified). + Quant stage: pre-decode PQ K codes + INT2 V to FP16, then FP16 attention. + Stage-2: shared LSE reduce (same as INT2 unified). + """ + if sinks is not None: + raise NotImplementedError( + "Mixed KV windows do not support sink tokens in pqk_int2v decode." + ) + + total_splits = hp_max_kv_splits + quant_max_kv_splits + assert attn_logits.shape[2] == total_splits + + attn_lse.fill_(float("-inf")) + + hp_logits = attn_logits[:, :, :hp_max_kv_splits, :] + hp_lse = attn_lse[:, :, :hp_max_kv_splits] + quant_logits = attn_logits[:, :, hp_max_kv_splits:, :] + quant_lse = attn_lse[:, :, hp_max_kv_splits:] + + kv_group_num = q.shape[1] // hp_k_buffer.shape[1] + + # HP stage: FP16 K+V attention + if hp_kv_indices.numel() > 0: if kv_group_num == 1: + _decode_att_m_fwd( + q, + hp_k_buffer, + hp_v_buffer, + hp_logits, + hp_lse, + hp_kv_indptr, + hp_kv_indices, + hp_num_kv_splits, + hp_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + else: + _decode_grouped_att_m_fwd( + q, + hp_k_buffer, + hp_v_buffer, + hp_logits, + hp_lse, + hp_kv_indptr, + hp_kv_indices, + hp_num_kv_splits, + hp_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + + # Quant stage: decode PQ/RVQ centroids inline while following indptr. + # The index buffer may be a padded CUDA Graph allocation; only entries + # selected by quant_kv_indptr are ever loaded. + if quant_kv_indices.numel() > 0: + _decode_grouped_att_m_fwd_pq( + q, + quant_k_buffer, + quant_v_buffer, + quant_v_scales_zeros, + pq_codebook, + quant_logits, + quant_lse, + quant_kv_indptr, + quant_kv_indices, + quant_num_kv_splits, + quant_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + k_codes2=quant_k_buffer2, + k_codebook2=pq_codebook2, + v_codebook=pq_v_codebook, + ) + + _unified_stage2(attn_logits, attn_lse, o, total_splits=total_splits) + return o + + +# --------------------------------------------------------------------------- +# INT1 decode attention: gather+dequant the referenced int1 rows to bf16, then +# route through the standard non-quantized attention kernels. Slower than an +# inline-dequant kernel but avoids duplicating the int2 stage-1 kernels (which +# would need 8 octants instead of 4 quarters and ~2x the body size). +# --------------------------------------------------------------------------- + + +def _dequant_int1_for_attention( + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + k_scales_zeros: torch.Tensor, + v_scales_zeros: torch.Tensor, + kv_indices: torch.Tensor, + head_dim_k: int, + head_dim_v: int, + out_dtype: torch.dtype, +): + """Gather + dequant K and V int1 rows at ``kv_indices``. + + Returns ``(k_bf16, v_bf16, remapped_kv_indices)`` where ``k_bf16`` / + ``v_bf16`` have shape ``(n_indices, num_heads, head_dim_{k,v})`` and + ``remapped_kv_indices`` is ``arange(n_indices)`` so the caller can drop + the dequant buffer in as a stand-in for an HP K/V buffer. + """ + from sglang.srt.mem_cache.kv_quant_kernels import ( + gather_dequantize_kv_int1_triton, + ) + + n_indices = int(kv_indices.shape[0]) + k_dequant = gather_dequantize_kv_int1_triton( + k_buffer, k_scales_zeros, kv_indices, head_dim_k, out_dtype + ) + v_dequant = gather_dequantize_kv_int1_triton( + v_buffer, v_scales_zeros, kv_indices, head_dim_v, out_dtype + ) + remapped = torch.arange(n_indices, dtype=kv_indices.dtype, device=kv_indices.device) + return k_dequant, v_dequant, remapped + + +def decode_attention_fwd_int1_via_dequant( + q, + k_buffer, + v_buffer, + k_scales_zeros, + v_scales_zeros, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, + xai_temperature_len=-1, + output_lse=None, +): + """INT1 decode attention via gather+dequant. Dispatches to the standard + non-quantized attention path after materializing the referenced int1 + rows as bf16. + """ + if kv_indices.numel() == 0: + return o + + head_dim_k = int(k_buffer.shape[-1]) * 8 + head_dim_v = int(v_buffer.shape[-1]) * 8 + out_dtype = q.dtype + + k_dequant, v_dequant, remapped = _dequant_int1_for_attention( + k_buffer, + v_buffer, + k_scales_zeros, + v_scales_zeros, + kv_indices, + head_dim_k, + head_dim_v, + out_dtype, + ) + + kv_group_num = q.shape[1] // v_buffer.shape[1] + if kv_group_num == 1: + _decode_att_m_fwd( + q, + k_dequant, + v_dequant, + attn_logits, + attn_lse, + kv_indptr, + remapped, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + else: + _decode_grouped_att_m_fwd( + q, + k_dequant, + v_dequant, + attn_logits, + attn_lse, + kv_indptr, + remapped, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + + _decode_softmax_reducev_fwd( + attn_logits, + attn_lse, + q, + o, + 1.0, # v_scale (already dequantized) + v_dequant, + kv_indptr, + num_kv_splits, + max_kv_splits, + sinks, + output_lse=output_lse, + ) + return o + + +def decode_attention_fwd_int1_unified( + q, + hp_k_buffer, + hp_v_buffer, + quant_k_buffer, + quant_v_buffer, + quant_k_scales_zeros, + quant_v_scales_zeros, + o, + hp_kv_indptr, + hp_kv_indices, + quant_kv_indptr, + quant_kv_indices, + attn_logits, + attn_lse, + hp_num_kv_splits, + quant_num_kv_splits, + hp_max_kv_splits, + quant_max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, + xai_temperature_len=-1, +): + """Unified HP + int1 decode attention via gather+dequant of the quant + portion. Mirrors :func:`decode_attention_fwd_int2_unified` but routes + int1 through the non-quantized stage-1 kernel after materializing the + referenced rows as bf16. + """ + if sinks is not None: + raise NotImplementedError( + "Mixed KV windows do not support sink tokens in Triton decode yet." + ) + + total_splits = hp_max_kv_splits + quant_max_kv_splits + assert attn_logits.shape[2] == total_splits, ( + f"attn_logits split dim ({attn_logits.shape[2]}) must equal hp_max_kv_splits " + f"({hp_max_kv_splits}) + quant_max_kv_splits ({quant_max_kv_splits})" + ) + + attn_lse.fill_(float("-inf")) + + hp_logits = attn_logits[:, :, :hp_max_kv_splits, :] + hp_lse = attn_lse[:, :, :hp_max_kv_splits] + quant_logits = attn_logits[:, :, hp_max_kv_splits:, :] + quant_lse = attn_lse[:, :, hp_max_kv_splits:] + + kv_group_num = q.shape[1] // hp_k_buffer.shape[1] + + if hp_kv_indices.numel() > 0: + if kv_group_num == 1: + _decode_att_m_fwd( + q, + hp_k_buffer, + hp_v_buffer, + hp_logits, + hp_lse, + hp_kv_indptr, + hp_kv_indices, + hp_num_kv_splits, + hp_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + else: + _decode_grouped_att_m_fwd( + q, + hp_k_buffer, + hp_v_buffer, + hp_logits, + hp_lse, + hp_kv_indptr, + hp_kv_indices, + hp_num_kv_splits, + hp_max_kv_splits, + sm_scale, + logit_cap, + xai_temperature_len, + ) + + if quant_kv_indices.numel() > 0: + if kv_group_num == 1 or quant_k_buffer.shape[-1] != quant_v_buffer.shape[-1]: _decode_att_m_fwd_quant_int2( q, quant_k_buffer, @@ -2479,6 +3643,7 @@ def decode_attention_fwd_int2_unified( sm_scale, logit_cap, xai_temperature_len, + int1=True, ) else: _decode_grouped_att_m_fwd_quant_int2( @@ -2496,6 +3661,7 @@ def decode_attention_fwd_int2_unified( sm_scale, logit_cap, xai_temperature_len, + int1=True, ) _unified_stage2( diff --git a/sglang-research/python/sglang/srt/mem_cache/common.py b/sglang-research/python/sglang/srt/mem_cache/common.py index 5fab44718..dcbc673cf 100644 --- a/sglang-research/python/sglang/srt/mem_cache/common.py +++ b/sglang-research/python/sglang/srt/mem_cache/common.py @@ -17,6 +17,11 @@ gpu_flush_int2, gpu_flush_int2_apply, gpu_flush_int2_plan, + gpu_flush_pqk_int2v_apply, +) +from sglang.QuantKernel.gpu_flush_int1 import ( + gpu_flush_int1_apply, + gpu_flush_int1_plan, ) if TYPE_CHECKING: @@ -732,7 +737,7 @@ def _alloc_for_decode_mixed(batch: ScheduleBatch, token_per_req: int) -> torch.T # ``allocator.free``, whose ``torch.unique`` host-syncs only against this # short pre-wait prefix instead of the previous forward. See # plan-for-a-fix-starry-russell.md. - plan = gpu_flush_int2_plan( + plan_kwargs = dict( seq_lens=seq_lens_int32, prefix_lens=prefix_lens_gpu, req_pool_indices=req_pool_indices_int64, @@ -744,6 +749,12 @@ def _alloc_for_decode_mixed(batch: ScheduleBatch, token_per_req: int) -> torch.T hp_global_offset=kv_pool.hp_global_offset, flush_interval=flush_interval, ) + use_int1_flush = getattr(kv_pool, "dtype", None) == "int1" + use_pqk_flush = getattr(kv_pool, "dtype", None) == "pq_k_int2v" + if use_int1_flush: + plan = gpu_flush_int1_plan(**plan_kwargs) + else: + plan = gpu_flush_int2_plan(**plan_kwargs) if plan is not None: # Free everything returned by the kernel in one call: flushed HP @@ -763,8 +774,7 @@ def _alloc_for_decode_mixed(batch: ScheduleBatch, token_per_req: int) -> torch.T wait_pending_forward() if plan is not None: - gpu_flush_int2_apply( - plan, + apply_kwargs = dict( req_pool_indices=req_pool_indices_int64, req_to_token=batch.req_to_token_pool.req_to_token, hp_k_ptrs=kv_pool._flush_hp_k_ptrs, @@ -794,6 +804,19 @@ def _alloc_for_decode_mixed(batch: ScheduleBatch, token_per_req: int) -> torch.T k_clip_ratio=kv_pool._k_clip_ratio, v_clip_ratio=kv_pool._v_clip_ratio, ) + if use_pqk_flush: + gpu_flush_pqk_int2v_apply( + plan, + req_pool_indices=req_pool_indices_int64, + req_to_token=batch.req_to_token_pool.req_to_token, + kv_pool=kv_pool, + ) + elif use_int1_flush: + gpu_flush_int1_apply( + plan, lloyd_max=getattr(kv_pool, "_lloyd_max", False), **apply_kwargs + ) + else: + gpu_flush_int2_apply(plan, **apply_kwargs) if batch.model_config.is_encoder_decoder: locs = batch.encoder_lens + batch.seq_lens diff --git a/sglang-research/python/sglang/srt/mem_cache/kv_quant_kernels.py b/sglang-research/python/sglang/srt/mem_cache/kv_quant_kernels.py index 71f4c2fab..2950ba434 100644 --- a/sglang-research/python/sglang/srt/mem_cache/kv_quant_kernels.py +++ b/sglang-research/python/sglang/srt/mem_cache/kv_quant_kernels.py @@ -1,7 +1,9 @@ """ -Triton kernels for efficient INT2 KV cache quantization. +Triton kernels for efficient INT2/INT1 KV cache quantization. """ +from typing import Optional + import torch import triton import triton.language as tl @@ -579,3 +581,374 @@ def dequantize_kv_int2_triton( num_stages=1, ) return output + + +# --------------------------------------------------------------------------- +# INT1 helpers +# --------------------------------------------------------------------------- +# INT1 packs 8 quant values per byte (1-bit slots). Slot i sits in bits [i, i+1). +# Storage shape is ``[cache_size, num_heads, head_dim // 8]`` uint8. Per-group +# scale/zero layout is identical to INT2 (interleaved scale/zero pairs). +# Bit ``i`` of byte ``b`` corresponds to original head_dim position +# ``i * (head_dim // 8) + b`` (so the 8 lanes in one byte are stride-1 sub- +# samples of the head_dim row, the same convention as INT2's "quartered split" +# but one level deeper). + + +@triton.jit +def _dequantize_kv_int1_kernel( + quantized_ptr, scales_zeros_ptr, output_ptr, + cache_size, num_heads, head_dim, + quant_stride_cache, quant_stride_head, quant_stride_dim, + sz_stride_cache, sz_stride_head, sz_stride_dim, + out_stride_cache, out_stride_head, out_stride_dim, + BLOCK_SIZE_DIM: tl.constexpr, +): + cache_idx = tl.program_id(0) + head_idx = tl.program_id(1) + if cache_idx >= cache_size or head_idx >= num_heads: + return + + sz_base = cache_idx * sz_stride_cache + head_idx * sz_stride_head + scale = tl.load(scales_zeros_ptr + sz_base + 0 * sz_stride_dim).to(tl.float32) + zero = tl.load(scales_zeros_ptr + sz_base + 1 * sz_stride_dim).to(tl.float32) + + octant_dim = head_dim // 8 + dim_offsets = tl.arange(0, BLOCK_SIZE_DIM) + dim_mask = dim_offsets < octant_dim + + quant_offset = ( + cache_idx * quant_stride_cache + head_idx * quant_stride_head + + dim_offsets * quant_stride_dim + ) + packed = tl.load(quantized_ptr + quant_offset, mask=dim_mask, other=0) + + out_base = cache_idx * out_stride_cache + head_idx * out_stride_head + for i in tl.static_range(8): + d_i = (((packed >> i) & 0x01).to(tl.float32) - zero) * scale + tl.store( + output_ptr + out_base + (dim_offsets + i * octant_dim) * out_stride_dim, + d_i, + mask=dim_mask, + ) + + +@triton.jit +def _dequantize_kv_int1_grouped_kernel( + quantized_ptr, + scales_zeros_ptr, + output_ptr, + cache_size, + num_heads, + quant_stride_cache, + quant_stride_head, + quant_stride_dim, + sz_stride_cache, + sz_stride_head, + sz_stride_dim, + out_stride_cache, + out_stride_head, + out_stride_dim, + GROUP_SIZE: tl.constexpr, + NUM_GROUPS_OCTANT: tl.constexpr, +): + """Groupwise INT1 dequantize. The 2D tile is shaped + ``[NUM_GROUPS_OCTANT, GROUP_SIZE]`` so that each of the 8 1-bit slots + inside a packed byte at ``(g, e)`` consistently belongs to a single group: + slot k uses group ``g + k * NUM_GROUPS_OCTANT``. + Requires ``num_groups % 8 == 0``. + """ + cache_idx = tl.program_id(0) + head_idx = tl.program_id(1) + if cache_idx >= cache_size or head_idx >= num_heads: + return + + octant_dim = NUM_GROUPS_OCTANT * GROUP_SIZE + + g_ids = tl.arange(0, NUM_GROUPS_OCTANT) + e_ids = tl.arange(0, GROUP_SIZE) + dim_offsets_2d = g_ids[:, None] * GROUP_SIZE + e_ids[None, :] + + quant_offset = ( + cache_idx * quant_stride_cache + + head_idx * quant_stride_head + + dim_offsets_2d * quant_stride_dim + ) + packed = tl.load(quantized_ptr + quant_offset) + + sz_base = cache_idx * sz_stride_cache + head_idx * sz_stride_head + out_base = cache_idx * out_stride_cache + head_idx * out_stride_head + + for i in tl.static_range(8): + gi = g_ids + i * NUM_GROUPS_OCTANT + s_i = tl.load(scales_zeros_ptr + sz_base + (gi * 2) * sz_stride_dim).to(tl.float32) + z_i = tl.load(scales_zeros_ptr + sz_base + (gi * 2 + 1) * sz_stride_dim).to(tl.float32) + q_i = ((packed >> i) & 0x01).to(tl.float32) + d_i = (q_i - z_i[:, None]) * s_i[:, None] + tl.store( + output_ptr + out_base + (dim_offsets_2d + i * octant_dim) * out_stride_dim, + d_i, + ) + + +@triton.jit +def _gather_dequantize_kv_int1_kernel( + quantized_ptr, + scales_zeros_ptr, + indices_ptr, + output_ptr, + n_indices, + num_heads, + head_dim, + quant_stride_cache, + quant_stride_head, + quant_stride_dim, + sz_stride_cache, + sz_stride_head, + sz_stride_dim, + out_stride_token, + out_stride_head, + out_stride_dim, + BLOCK_SIZE_DIM: tl.constexpr, +): + """Gather + dequant int1 single-scale kernel. Reads slot id from + ``indices_ptr[token_idx]`` and writes dequantized row to + ``output[token_idx]``. + """ + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + if token_idx >= n_indices or head_idx >= num_heads: + return + + slot = tl.load(indices_ptr + token_idx).to(tl.int64) + + sz_base = slot * sz_stride_cache + head_idx * sz_stride_head + scale = tl.load(scales_zeros_ptr + sz_base + 0 * sz_stride_dim).to(tl.float32) + zero = tl.load(scales_zeros_ptr + sz_base + 1 * sz_stride_dim).to(tl.float32) + + octant_dim = head_dim // 8 + dim_offsets = tl.arange(0, BLOCK_SIZE_DIM) + dim_mask = dim_offsets < octant_dim + + quant_offset = ( + slot * quant_stride_cache + head_idx * quant_stride_head + + dim_offsets * quant_stride_dim + ) + packed = tl.load(quantized_ptr + quant_offset, mask=dim_mask, other=0) + + out_base = token_idx * out_stride_token + head_idx * out_stride_head + for i in tl.static_range(8): + d_i = (((packed >> i) & 0x01).to(tl.float32) - zero) * scale + tl.store( + output_ptr + out_base + (dim_offsets + i * octant_dim) * out_stride_dim, + d_i, + mask=dim_mask, + ) + + +@triton.jit +def _gather_dequantize_kv_int1_grouped_kernel( + quantized_ptr, + scales_zeros_ptr, + indices_ptr, + output_ptr, + n_indices, + num_heads, + quant_stride_cache, + quant_stride_head, + quant_stride_dim, + sz_stride_cache, + sz_stride_head, + sz_stride_dim, + out_stride_token, + out_stride_head, + out_stride_dim, + GROUP_SIZE: tl.constexpr, + NUM_GROUPS_OCTANT: tl.constexpr, +): + """Gather + dequant int1 grouped kernel. ``num_groups % 8 == 0`` required.""" + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + if token_idx >= n_indices or head_idx >= num_heads: + return + + slot = tl.load(indices_ptr + token_idx).to(tl.int64) + + octant_dim = NUM_GROUPS_OCTANT * GROUP_SIZE + + g_ids = tl.arange(0, NUM_GROUPS_OCTANT) + e_ids = tl.arange(0, GROUP_SIZE) + dim_offsets_2d = g_ids[:, None] * GROUP_SIZE + e_ids[None, :] + + quant_offset = ( + slot * quant_stride_cache + + head_idx * quant_stride_head + + dim_offsets_2d * quant_stride_dim + ) + packed = tl.load(quantized_ptr + quant_offset) + + sz_base = slot * sz_stride_cache + head_idx * sz_stride_head + out_base = token_idx * out_stride_token + head_idx * out_stride_head + + for i in tl.static_range(8): + gi = g_ids + i * NUM_GROUPS_OCTANT + s_i = tl.load(scales_zeros_ptr + sz_base + (gi * 2) * sz_stride_dim).to(tl.float32) + z_i = tl.load(scales_zeros_ptr + sz_base + (gi * 2 + 1) * sz_stride_dim).to(tl.float32) + q_i = ((packed >> i) & 0x01).to(tl.float32) + d_i = (q_i - z_i[:, None]) * s_i[:, None] + tl.store( + output_ptr + out_base + (dim_offsets_2d + i * octant_dim) * out_stride_dim, + d_i, + ) + + +def gather_dequantize_kv_int1_triton( + quantized: torch.Tensor, + scales_zeros: torch.Tensor, + indices: torch.Tensor, + head_dim: int, + model_dtype: torch.dtype, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Gather rows of an INT1 KV buffer at ``indices`` and dequantize to + ``model_dtype``. Result has shape ``(len(indices), num_heads, head_dim)``. + + Used by the int1 decode path: rather than implementing a separate + int1-aware attention kernel, the rows referenced by ``kv_indices`` are + materialized as bf16 once per layer and fed to the standard non-quantized + attention kernel with a remapped index range ``[0, len(indices))``. + """ + assert head_dim % 8 == 0, ( + f"head_dim ({head_dim}) must be divisible by 8 for INT1" + ) + n_indices = int(indices.shape[0]) + num_heads = quantized.shape[1] + if out is None: + out = torch.empty( + (n_indices, num_heads, head_dim), + dtype=model_dtype, + device=quantized.device, + ) + else: + assert out.shape == (n_indices, num_heads, head_dim), ( + f"out shape {tuple(out.shape)} != " + f"({n_indices}, {num_heads}, {head_dim})" + ) + assert out.dtype == model_dtype + if n_indices == 0: + return out + + num_groups = _get_num_scale_groups(scales_zeros) + grid = (n_indices, num_heads) + + if num_groups == 1: + BLOCK_SIZE_DIM = triton.next_power_of_2(head_dim // 8) + _gather_dequantize_kv_int1_kernel[grid]( + quantized, + scales_zeros, + indices, + out, + n_indices, + num_heads, + head_dim, + quantized.stride(0), + quantized.stride(1), + quantized.stride(2), + scales_zeros.stride(0), + scales_zeros.stride(1), + scales_zeros.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + BLOCK_SIZE_DIM=BLOCK_SIZE_DIM, + ) + return out + + group_size = head_dim // num_groups + if not _can_use_triton_groupwise(num_groups, group_size, packing=8): + raise NotImplementedError( + f"int1 KV gather+dequant: unsupported quant grouping " + f"(num_groups={num_groups}, group_size={group_size})." + ) + + _gather_dequantize_kv_int1_grouped_kernel[grid]( + quantized, + scales_zeros, + indices, + out, + n_indices, + num_heads, + quantized.stride(0), + quantized.stride(1), + quantized.stride(2), + scales_zeros.stride(0), + scales_zeros.stride(1), + scales_zeros.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + GROUP_SIZE=group_size, + NUM_GROUPS_OCTANT=num_groups // 8, + num_warps=1, + num_stages=1, + ) + return out + + +def dequantize_kv_int1_triton( + quantized: torch.Tensor, scales_zeros: torch.Tensor, + head_dim: int, model_dtype: torch.dtype, +) -> torch.Tensor: + """Dequantize INT1 KV cache to ``model_dtype``.""" + assert head_dim % 8 == 0, ( + f"head_dim ({head_dim}) must be divisible by 8 for INT1" + ) + cache_size, num_heads, _ = quantized.shape + output = torch.empty( + (cache_size, num_heads, head_dim), dtype=model_dtype, device=quantized.device + ) + grid = (cache_size, num_heads) + num_groups = _get_num_scale_groups(scales_zeros) + + if num_groups == 1: + BLOCK_SIZE_DIM = triton.next_power_of_2(head_dim // 8) + _dequantize_kv_int1_kernel[grid]( + quantized, scales_zeros, output, + cache_size, num_heads, head_dim, + quantized.stride(0), quantized.stride(1), quantized.stride(2), + scales_zeros.stride(0), scales_zeros.stride(1), scales_zeros.stride(2), + output.stride(0), output.stride(1), output.stride(2), + BLOCK_SIZE_DIM=BLOCK_SIZE_DIM, + ) + return output + + group_size = head_dim // num_groups + if not _can_use_triton_groupwise(num_groups, group_size, packing=8): + raise NotImplementedError( + f"int1 KV dequantize: unsupported quant grouping " + f"(num_groups={num_groups}, group_size={group_size}). " + f"The int1 Triton kernel requires num_groups and group_size to be " + f"powers of two with num_groups % 8 == 0." + ) + + _dequantize_kv_int1_grouped_kernel[grid]( + quantized, + scales_zeros, + output, + cache_size, + num_heads, + quantized.stride(0), + quantized.stride(1), + quantized.stride(2), + scales_zeros.stride(0), + scales_zeros.stride(1), + scales_zeros.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + GROUP_SIZE=group_size, + NUM_GROUPS_OCTANT=num_groups // 8, + num_warps=1, + num_stages=1, + ) + return output diff --git a/sglang-research/python/sglang/srt/mem_cache/memory_pool.py b/sglang-research/python/sglang/srt/mem_cache/memory_pool.py index 017a3ab7e..91de12678 100644 --- a/sglang-research/python/sglang/srt/mem_cache/memory_pool.py +++ b/sglang-research/python/sglang/srt/mem_cache/memory_pool.py @@ -768,10 +768,16 @@ def __init__( self.device = device if model_dtype is not None: self.model_dtype = model_dtype - elif dtype == "int2": + elif dtype in ("int2", "int1", "pq_k_int2v"): raise ValueError(f"model_dtype is required for {dtype} kv cache") - if dtype in (torch.float8_e5m2, torch.float8_e4m3fn, "int2"): + if dtype in ( + torch.float8_e5m2, + torch.float8_e4m3fn, + "int2", + "int1", + "pq_k_int2v", + ): # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2 self.store_dtype = torch.uint8 else: @@ -893,10 +899,11 @@ def __init__( else v_head_dim if v_head_dim is not None else head_dim ) self.kv_cache_quant_group_size = kv_cache_quant_group_size - # Scale/zero dtype for int2-packed scales. fp32 preserves the historical - # default; bf16/fp16 are opt-in via env. Unused for non-int2 dtypes. + # Scale/zero dtype for int2/int1-packed scales. fp32 preserves the + # historical default; bf16/fp16 are opt-in via env. Unused for other + # dtypes. self.scale_dtype = scale_dtype if scale_dtype is not None else torch.float32 - if self.dtype == "int2": + if self.dtype in ("int2", "int1", "pq_k_int2v"): self.k_quant_group_size, self.k_num_scale_groups = ( self._resolve_quant_grouping(self.head_dim, "K") ) @@ -1013,10 +1020,10 @@ def _create_buffers(self): if self.dtype == "int2": assert ( self.head_dim % 4 == 0 - ), f"head_dim: {self.head_dim}, kv cache dtype: int2" + ), f"head_dim: {self.head_dim}, kv cache dtype: {self.dtype}" assert ( self.v_head_dim % 4 == 0 - ), f"v_head_dim: {self.v_head_dim}, kv cache dtype: int2" + ), f"v_head_dim: {self.v_head_dim}, kv cache dtype: {self.dtype}" self.k_buffer = [ torch.zeros( (self.size + self.page_size, self.head_num, self.head_dim // 4), diff --git a/sglang-research/python/sglang/srt/mem_cache/unified_kv_pool.py b/sglang-research/python/sglang/srt/mem_cache/unified_kv_pool.py index 893f88b5b..fc3f4dd7e 100644 --- a/sglang-research/python/sglang/srt/mem_cache/unified_kv_pool.py +++ b/sglang-research/python/sglang/srt/mem_cache/unified_kv_pool.py @@ -1,10 +1,14 @@ """ -Unified HP + int2 KV cache pool. +Unified HP + int2/int1 KV cache pool. Quant arena: paged with ``N_Q`` slots per page. HP arena: shared HP-prefix pool (paged) followed by per-request HP-recent ring slabs. Slot id namespace is flat (``[0, num_quant_pages*N_Q)`` quant, ``[HP_OFFSET, ...)`` HP), and kernels dispatch by ``slot >= HP_OFFSET``. + +INT2 packs 4 quant slots per byte (head_dim // 4 bytes per row); INT1 packs +8 quant slots per byte (head_dim // 8 bytes per row). The pool is parameterized +by ``pack_factor`` so a single class serves both dtypes. """ from __future__ import annotations @@ -21,9 +25,20 @@ quantized_set_kv_int2_pretransformed_triton, ) from sglang.QuantKernel.oscar_rotation_clip_int2_kv import ( + _launch_grouped_clip_int2, + _launch_single_clip_int2, quantized_set_kv_int2_oscar_rotate_k_clip_triton, quantized_set_kv_int2_pretransformed_clip_triton, ) +from sglang.QuantKernel.oscar_rotation_clip_int1_kv import ( + quantized_set_kv_int1_oscar_rotate_k_clip_triton, + quantized_set_kv_int1_pretransformed_clip_triton, + quantized_set_kv_int1_pretransformed_triton, +) +from sglang.QuantKernel.oscar_rotation_pq_k_kv import ( + pq_decode_k_at_locs, + pq_encode_k, +) from sglang.srt.mem_cache.kv_quant_kernels import _get_num_scale_groups from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs @@ -42,6 +57,109 @@ GB = 1024 * 1024 * 1024 +def _validate_pq_codebook_file( + data: dict, expected_head_dim: int, label: str +) -> Tuple[int, int, int]: + """Return (n_sub, n_centroids, sub_dim) after strict shape validation.""" + if "codebooks_stage1" in data: + stage1 = data["codebooks_stage1"] + stage2 = data.get("codebooks_stage2") + if stage1.ndim != 4 or stage2 is None or stage2.ndim != 4: + raise ValueError( + f"{label} RVQ codebooks must be rank-4 stage1/stage2 tensors" + ) + if ( + stage1.shape[0] != stage2.shape[0] + or stage1.shape[1] != stage2.shape[1] + or stage1.shape[3] != stage2.shape[3] + ): + raise ValueError( + f"{label} RVQ stage shapes are incompatible: " + f"{tuple(stage1.shape)} vs {tuple(stage2.shape)}" + ) + n_sub, n_centroids, sub_dim = map(int, stage1.shape[1:]) + n_layers = int(stage1.shape[0]) + if int(stage2.shape[2]) > 256: + raise ValueError( + f"{label} RVQ stage2 has {stage2.shape[2]} centroids; " + "uint8 codes support at most 256" + ) + if int(stage2.shape[2]) & (int(stage2.shape[2]) - 1): + raise ValueError( + f"{label} RVQ stage2 centroid count must be a power of two, " + f"got {stage2.shape[2]}" + ) + elif "codebooks_per_layer" in data: + codebooks = data["codebooks_per_layer"] + if codebooks.ndim != 4: + raise ValueError( + f"{label} codebooks_per_layer must be rank 4, got " + f"shape={tuple(codebooks.shape)}" + ) + n_sub, n_centroids, sub_dim = map(int, codebooks.shape[1:]) + n_layers = int(codebooks.shape[0]) + elif "codebooks" in data: + books = data["codebooks"] + if not books: + raise ValueError(f"{label} shared codebook list is empty") + stacked = torch.stack(books) + if stacked.ndim != 3: + raise ValueError( + f"{label} shared codebooks must stack to rank 3, got " + f"shape={tuple(stacked.shape)}" + ) + n_sub, n_centroids, sub_dim = map(int, stacked.shape) + n_layers = None + else: + raise ValueError(f"{label} file contains no supported codebook tensor") + + metadata = { + "n_sub": n_sub, + "sub_dim": sub_dim, + } + if "n_centroids" in data: + metadata["n_centroids"] = n_centroids + for key, actual in metadata.items(): + if key in data and int(data[key]) != actual: + raise ValueError( + f"{label} metadata {key}={data[key]} does not match tensor " + f"shape ({actual})" + ) + if n_centroids > 256: + raise ValueError( + f"{label} has {n_centroids} centroids; uint8 codes support at most 256" + ) + if n_centroids & (n_centroids - 1): + raise ValueError( + f"{label} centroid count must be a power of two, got {n_centroids}" + ) + if n_sub * sub_dim != expected_head_dim: + raise ValueError( + f"{label} codebook reconstructs {n_sub}*{sub_dim}=" + f"{n_sub * sub_dim} dims, expected {expected_head_dim}" + ) + if n_layers is not None: + raw_layer_ids = data.get("layer_ids", list(range(n_layers))) + if len(raw_layer_ids) != n_layers: + raise ValueError( + f"{label} has {n_layers} layer codebooks but " + f"{len(raw_layer_ids)} layer_ids" + ) + layer_ids = [] + for raw_layer_id in raw_layer_ids: + try: + layer_id = int(raw_layer_id) + exact = float(raw_layer_id) == layer_id + except (TypeError, ValueError): + exact = False + if not exact: + raise ValueError(f"{label} layer_id {raw_layer_id!r} is not an integer") + layer_ids.append(layer_id) + if len(set(layer_ids)) != len(layer_ids): + raise ValueError(f"{label} layer_ids must be unique: {layer_ids}") + return n_sub, n_centroids, sub_dim + + @triton.jit def _set_mixed_hp_buffer_kernel( src_ptr, @@ -88,9 +206,7 @@ def _resolve_torch_dtype(name: str, *, kind: str) -> torch.dtype: return torch.float16 if n in ("fp32", "float32"): return torch.float32 - raise ValueError( - f"Unsupported {kind} dtype: {name}. Expected bf16/fp16/fp32." - ) + raise ValueError(f"Unsupported {kind} dtype: {name}. Expected bf16/fp16/fp32.") def resolve_scale_dtype(name: str) -> torch.dtype: @@ -155,11 +271,61 @@ def __init__( scale_dtype: torch.dtype = torch.bfloat16, num_hp_prefix_slots: int = 0, ): - assert dtype == "int2", ( - "UnifiedInt2HPKVPool supports only int2 quant tier; got %s" % dtype + assert dtype in ("int2", "int1", "pq_k_int2v"), ( + "UnifiedInt2HPKVPool supports int2/int1/pq_k_int2v quant tiers; got %s" + % dtype ) + # Bit-width-specific packing: int2 → 4 vals/byte, + # scalar int1 → 8 vals/byte for both K and V, and PQ → one byte per + # sub-vector code (16 bytes for the default 16x8 codebook). + # For pq_k_int2v: K uses N_SUB=16 codes/token-head = 16 bytes (same as int1 @ 128-dim), + # V uses int2 packing (4 vals/byte = 32 bytes). Split into k/v pack factors. + self._k_pack_factor = 8 if dtype in ("int1", "pq_k_int2v") else 4 + self._v_pack_factor = 8 if dtype == "int1" else 4 + self._pack_factor = ( + self._k_pack_factor + ) # back-compat alias (used in base asserts) + # RVQ stage-2 state — MUST be defined before _create_arenas() (which allocates + # k_buffer2 when _is_rvq). _is_rvq is decided in the early peek below; the cb2 + # codebooks themselves are loaded later (only needed at flush/decode, not alloc). + self._is_rvq: bool = False + self._rvq_cb2_per_layer: Optional[list] = None + self._rvq_cb2_norms_per_layer: Optional[list] = None + self.k_buffer2: Optional[list] = None + # Early codebook peek for pq_k_int2v: n_sub may differ from the n_sub=16 + # default (e.g. CQ-16c8b uses n_sub=8 → k_pack_factor=16). We must know + # the correct k_pack_factor BEFORE _create_arenas() allocates k_buffer, or + # the buffer gets the wrong last dimension and pq_decode_k infers the wrong + # N_SUB, causing OOB reads into the codebook. Only read the scalar field + # here; full codebook tensors are loaded to GPU later (lines 303+). + if dtype == "pq_k_int2v": + _pq_path_early = envs.SGLANG_PQ_K_CODEBOOK.get() + if _pq_path_early: + _hdr = torch.load( + _pq_path_early, map_location="cpu", weights_only=False + ) + _n_sub_early, _, _ = _validate_pq_codebook_file(_hdr, head_dim, "PQ K") + _early_pack = head_dim // _n_sub_early + if _early_pack != self._k_pack_factor: + self._k_pack_factor = _early_pack + self._pack_factor = _early_pack + # RVQ if the codebook carries a stage-2 (residual) book. + self._is_rvq = "codebooks_stage1" in _hdr + # Early peek for PQ V: set v_pack_factor so v_buffer is sized to the PQ + # codes (v_head_dim // n_sub_v = 16B for n_sub=16), not the INT2 32B — + # this is the true 1.0-bpe-stored V (clean packing). Must precede _create_arenas. + _pqv_path_early = envs.SGLANG_PQ_V_CODEBOOK.get() + if _pqv_path_early: + _vhdr = torch.load( + _pqv_path_early, map_location="cpu", weights_only=False + ) + _vhd_early = v_head_dim if v_head_dim is not None else head_dim + _v_n_sub_early, _, _ = _validate_pq_codebook_file( + _vhdr, _vhd_early, "PQ V" + ) + self._v_pack_factor = _vhd_early // _v_n_sub_early # Work around KVCache.__init__ dtype validation: it stores ``dtype`` as - # a string and sets ``store_dtype=torch.uint8`` for int2. + # a string and sets ``store_dtype=torch.uint8`` for int2/int1. super().__init__( size=num_quant_pages, # used by base class for sizing heuristics only page_size=1, @@ -230,11 +396,13 @@ def __init__( self.v_quant_group_size, self.v_num_scale_groups = self._resolve_quant_grouping( self.v_head_dim, "V" ) - assert self.head_dim % 4 == 0, ( - f"head_dim={self.head_dim} must be divisible by 4 for int2 packing" + assert self.head_dim % self._k_pack_factor == 0, ( + f"head_dim={self.head_dim} must be divisible by {self._k_pack_factor} " + f"for {dtype} K packing" ) - assert self.v_head_dim % 4 == 0, ( - f"v_head_dim={self.v_head_dim} must be divisible by 4 for int2 packing" + assert self.v_head_dim % self._v_pack_factor == 0, ( + f"v_head_dim={self.v_head_dim} must be divisible by " + f"{self._v_pack_factor} for {dtype} V packing" ) self._create_arenas() @@ -276,9 +444,169 @@ def __init__( self._lloyd_max, ) + # PQ K codebook (only for pq_k_int2v dtype). + # Supports two file formats: + # (a) shared codebook: {"codebooks": list[tensor], "n_sub", "sub_dim", "n_centroids", ...} + # (b) per-layer codebooks: {"codebooks_per_layer": Tensor[L, N_SUB, N_CENTS, SUB_DIM], ...} + # Per-layer is strongly preferred — shared codebooks fail when layers have + # very different K scales (e.g. Qwen3-8B layer 0 has K ~7x larger than other layers). + self._pq_codebook: Optional[torch.Tensor] = ( + None # [N_SUB, N_CENTS, SUB_DIM] (shared) + ) + self._pq_codebooks_per_layer: Optional[list] = ( + None # list of [N_SUB, N_CENTS, SUB_DIM] + ) + self._pq_cb_norms: Optional[torch.Tensor] = ( + None # [N_SUB, N_CENTS] (shared, or None) + ) + self._pq_cb_norms_per_layer: Optional[list] = None # list of [N_SUB, N_CENTS] + # RVQ stage-2 (residual VQ): _is_rvq / k_buffer2 / cb2 lists are initialized + # earlier (before _create_arenas, which allocates k_buffer2). The stage-2 + # codebooks are populated in the load block below when codebooks_stage1 present. + # PQ V (optional): when SGLANG_PQ_V_CODEBOOK is set, V uses PQ instead of INT2. + # The early codebook peek above changes _v_pack_factor before arena + # construction, so v_buffer is exactly n_sub bytes wide. + self._pq_v: bool = False + self._pq_v_codebooks_per_layer: Optional[list] = None + self._pq_v_cb_norms_per_layer: Optional[list] = None + if dtype == "pq_k_int2v": + _pqv_path = envs.SGLANG_PQ_V_CODEBOOK.get() + if _pqv_path: + _vd = torch.load(_pqv_path, map_location="cpu", weights_only=False) + _validate_pq_codebook_file(_vd, self.v_head_dim, "PQ V") + _vcbs = _vd["codebooks_per_layer"].to( + torch.float16 + ) # [L, N_SUB, N_CENTS, SUB_DIM] + _vids = _vd.get("layer_ids", list(range(len(_vcbs)))) + _v_l2i = {int(l): i for i, l in enumerate(_vids)} + _vdev = torch.device(self.device) + self._pq_v_codebooks_per_layer = [] + self._pq_v_cb_norms_per_layer = [] + for lid in range(self.start_layer, self.start_layer + self.layer_num): + vc = _vcbs[_v_l2i[lid]].to(_vdev).contiguous() + self._pq_v_codebooks_per_layer.append(vc) + self._pq_v_cb_norms_per_layer.append( + (vc.float() ** 2).sum(-1).to(dtype=torch.float32, device=_vdev) + ) + self._pq_v = True + logger.info( + "UnifiedInt2HPKVPool: PQ V codebooks loaded from %s (n_layers=%d n_sub=%d " + "n_cents=%d sqnr_avg=%.2f dB) — V uses PQ, stored in v_buffer[:n_sub]", + _pqv_path, + self.layer_num, + _vd["n_sub"], + _vd["n_centroids"], + _vd.get("sqnr_avg", float("nan")), + ) + pq_path = envs.SGLANG_PQ_K_CODEBOOK.get() + if not pq_path: + raise ValueError( + "SGLANG_PQ_K_CODEBOOK must point to a .pt codebook file when dtype=pq_k_int2v" + ) + _pq_data = torch.load(pq_path, map_location="cpu", weights_only=False) + _n_sub, _n_cents, _sub_dim = _validate_pq_codebook_file( + _pq_data, self.head_dim, "PQ K" + ) + dev = torch.device(self.device) + if "codebooks_stage1" in _pq_data: + # RVQ format: stage1 = PQ codebook, stage2 = residual codebook. + _s1 = _pq_data["codebooks_stage1"].to( + torch.float16 + ) # [L, N_SUB, N_CENTS1, SUB_DIM] + _s2 = _pq_data["codebooks_stage2"].to( + torch.float16 + ) # [L, N_SUB, N_CENTS2, SUB_DIM] + _layer_ids = _pq_data.get("layer_ids", list(range(len(_s1)))) + _lid_to_cb_idx = {int(lid): i for i, lid in enumerate(_layer_ids)} + self._pq_codebooks_per_layer = [] + self._pq_cb_norms_per_layer = [] + self._rvq_cb2_per_layer = [] + self._rvq_cb2_norms_per_layer = [] + for lid in range(self.start_layer, self.start_layer + self.layer_num): + c1 = _s1[_lid_to_cb_idx[lid]].to(dev).contiguous() + c2 = _s2[_lid_to_cb_idx[lid]].to(dev).contiguous() + self._pq_codebooks_per_layer.append(c1) + self._pq_cb_norms_per_layer.append( + (c1.float() ** 2).sum(-1).to(dtype=torch.float32, device=dev) + ) + self._rvq_cb2_per_layer.append(c2) + self._rvq_cb2_norms_per_layer.append( + (c2.float() ** 2).sum(-1).to(dtype=torch.float32, device=dev) + ) + self._is_rvq = True + logger.info( + "UnifiedInt2HPKVPool: RVQ K codebooks loaded from %s (n_layers=%d n_sub=%d " + "stage1_cents=%d stage2_cents=%d sub_dim=%d sqnr_stage1=%.2f sqnr_rvq=%.2f dB)", + pq_path, + self.layer_num, + _n_sub, + _s1.shape[2], + _s2.shape[2], + _sub_dim, + _pq_data.get("sqnr_stage1_avg", float("nan")), + _pq_data.get("sqnr_rvq_avg", float("nan")), + ) + elif "codebooks_per_layer" in _pq_data: + # Per-layer codebooks: [n_layers, N_SUB, N_CENTS, SUB_DIM] fp32 + _all_cbs = _pq_data["codebooks_per_layer"].to( + torch.float16 + ) # [L, N_SUB, N_CENTS, SUB_DIM] + _layer_ids = _pq_data.get("layer_ids", list(range(len(_all_cbs)))) + assert len(_layer_ids) == _all_cbs.shape[0] + # Build layer_index → codebook mapping (same ordering as layer_ids) + _lid_to_cb_idx = {int(lid): i for i, lid in enumerate(_layer_ids)} + self._pq_codebooks_per_layer = [] + self._pq_cb_norms_per_layer = [] + for lid in range(self.start_layer, self.start_layer + self.layer_num): + cb = _all_cbs[_lid_to_cb_idx[lid]].to(dev).contiguous() + self._pq_codebooks_per_layer.append(cb) + self._pq_cb_norms_per_layer.append( + (cb.float() ** 2).sum(-1).to(dtype=torch.float32, device=dev) + ) + sqnr_info = _pq_data.get("sqnr_avg", float("nan")) + logger.info( + "UnifiedInt2HPKVPool: PQ K per-layer codebooks loaded from %s " + "(n_layers=%d n_sub=%d n_cents=%d sub_dim=%d avg_sqnr=%.2f dB)", + pq_path, + self.layer_num, + _n_sub, + _n_cents, + _sub_dim, + sqnr_info, + ) + else: + # Legacy shared codebook format + _books = _pq_data["codebooks"] # list of [N_CENTS, SUB_DIM] tensors + self._pq_codebook = torch.stack(_books).to( + dtype=torch.float16, device=dev + ) + self._pq_cb_norms = ( + (self._pq_codebook.float() ** 2) + .sum(-1) + .to(dtype=torch.float32, device=dev) + ) + logger.info( + "UnifiedInt2HPKVPool: PQ K shared codebook loaded from %s " + "(n_sub=%d n_cents=%d sub_dim=%d sqnr_train=%.2f dB) " + "[WARNING: shared codebook degrades on layers with different K scales]", + pq_path, + _n_sub, + _n_cents, + _sub_dim, + _pq_data.get("sqnr_train", float("nan")), + ) + # Log the effective pack factor (already set via the early peek above). + new_pack = self.head_dim // _n_sub + if new_pack != 8: # non-default → worth calling out + logger.info( + "UnifiedInt2HPKVPool: PQ K n_sub=%d → k_pack_factor=%d (%.2f bpe K)", + _n_sub, + new_pack, + 8.0 / (self.head_dim / _n_sub), + ) + hp_total_slots = ( - self.num_hp_prefix_slots - + self.max_req_slots * self.hp_recent_ring_size + self.num_hp_prefix_slots + self.max_req_slots * self.hp_recent_ring_size ) self._finalize_allocation_log(hp_total_slots) hp_itemsize = torch.empty(0, dtype=self.hp_dtype).element_size() @@ -342,10 +670,7 @@ def hp_global_offset(self) -> int: @property def hp_size(self) -> int: - return ( - self.num_hp_prefix_slots - + self.max_req_slots * self.hp_recent_ring_size - ) + return self.num_hp_prefix_slots + self.max_req_slots * self.hp_recent_ring_size @property def quant_size(self) -> int: @@ -374,7 +699,9 @@ def release_req_slab(self, req_pool_idx) -> None: self._next_slab_offset[i] = 0 self._flush_counter[i] = 0 - def _resolve_quant_grouping(self, head_dim: int, tensor_name: str) -> tuple[int, int]: + def _resolve_quant_grouping( + self, head_dim: int, tensor_name: str + ) -> tuple[int, int]: group_size = ( head_dim if self.kv_cache_quant_group_size is None @@ -398,8 +725,7 @@ def _create_arenas(self): # [per-req recent slab 1] ... Quant arena is paged with N_Q slots # per page; scales/zeros are quant-only. hp_total_slots = ( - self.num_hp_prefix_slots - + self.max_req_slots * self.hp_recent_ring_size + self.num_hp_prefix_slots + self.max_req_slots * self.hp_recent_ring_size ) with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( @@ -409,15 +735,29 @@ def _create_arenas(self): ): self.k_buffer = [ torch.zeros( - (self.num_quant_pages * self.N_Q, self.head_num, self.head_dim // 4), + ( + self.num_quant_pages * self.N_Q, + self.head_num, + self.head_dim // self._k_pack_factor, + ), dtype=torch.uint8, device=self.device, ) for _ in range(self.layer_num) ] + if self._is_rvq: + # RVQ stage-2 codes: identical shape to k_buffer (same n_sub). + self.k_buffer2 = [ + torch.zeros_like(self.k_buffer[i]) + for i in range(self.layer_num) + ] self.v_buffer = [ torch.zeros( - (self.num_quant_pages * self.N_Q, self.head_num, self.v_head_dim // 4), + ( + self.num_quant_pages * self.N_Q, + self.head_num, + self.v_head_dim // self._v_pack_factor, + ), dtype=torch.uint8, device=self.device, ) @@ -425,7 +765,11 @@ def _create_arenas(self): ] self.k_scales_zeros = [ torch.zeros( - (self.num_quant_pages * self.N_Q, self.head_num, 2 * self.k_num_scale_groups), + ( + self.num_quant_pages * self.N_Q, + self.head_num, + 2 * self.k_num_scale_groups, + ), dtype=self.scale_dtype, device=self.device, ) @@ -433,7 +777,11 @@ def _create_arenas(self): ] self.v_scales_zeros = [ torch.zeros( - (self.num_quant_pages * self.N_Q, self.head_num, 2 * self.v_num_scale_groups), + ( + self.num_quant_pages * self.N_Q, + self.head_num, + 2 * self.v_num_scale_groups, + ), dtype=self.scale_dtype, device=self.device, ) @@ -506,6 +854,8 @@ def _strides(t: torch.Tensor) -> tuple: def get_kv_size_bytes(self): k = sum(get_tensor_size_bytes(t) for t in self.k_buffer) + if self.k_buffer2 is not None: + k += sum(get_tensor_size_bytes(t) for t in self.k_buffer2) k += sum(get_tensor_size_bytes(s) for s in self.k_scales_zeros) k += sum(get_tensor_size_bytes(t) for t in self.hp_k_buffer) v = sum(get_tensor_size_bytes(t) for t in self.v_buffer) @@ -516,6 +866,46 @@ def get_kv_size_bytes(self): def _layer_index(self, layer_id: int) -> int: return layer_id - self.start_layer + def get_pq_codebook(self, layer_id: int) -> torch.Tensor: + """Return the PQ K codebook for this layer (per-layer or shared).""" + if self._pq_codebooks_per_layer is not None: + return self._pq_codebooks_per_layer[self._layer_index(layer_id)] + return self._pq_codebook + + def get_pq_cb_norms(self, layer_id: int) -> torch.Tensor: + """Return the PQ K centroid norms for this layer (per-layer or shared).""" + if self._pq_cb_norms_per_layer is not None: + return self._pq_cb_norms_per_layer[self._layer_index(layer_id)] + return self._pq_cb_norms + + def get_pq_v_codebook(self, layer_id: int) -> Optional[torch.Tensor]: + """PQ V codebook for this layer, or None if V is not PQ (uses INT2).""" + if self._pq_v_codebooks_per_layer is None: + return None + return self._pq_v_codebooks_per_layer[self._layer_index(layer_id)] + + def get_pq_v_cb_norms(self, layer_id: int) -> Optional[torch.Tensor]: + if self._pq_v_cb_norms_per_layer is None: + return None + return self._pq_v_cb_norms_per_layer[self._layer_index(layer_id)] + + def get_rvq_cb2(self, layer_id: int) -> Optional[torch.Tensor]: + """RVQ stage-2 (residual) codebook for this layer, or None if not RVQ.""" + if self._rvq_cb2_per_layer is None: + return None + return self._rvq_cb2_per_layer[self._layer_index(layer_id)] + + def get_rvq_cb2_norms(self, layer_id: int) -> Optional[torch.Tensor]: + if self._rvq_cb2_norms_per_layer is None: + return None + return self._rvq_cb2_norms_per_layer[self._layer_index(layer_id)] + + def get_raw_key_buffer2(self, layer_id: int) -> Optional[torch.Tensor]: + """RVQ stage-2 codes buffer for this layer, or None if not RVQ.""" + if self.k_buffer2 is None: + return None + return self.k_buffer2[self._layer_index(layer_id)] + def get_key_buffer(self, layer_id: int) -> torch.Tensor: # Triton backend asks for the quant view in the mixed path; HP view is # accessed via ``get_hp_key_buffer``. @@ -547,13 +937,16 @@ def get_hp_value_buffer(self, layer_id: int) -> torch.Tensor: def get_raw_kv_buffer(self, layer_id: int): idx = self._layer_index(layer_id) - return { + buffers = { "k_buffer": self.k_buffer[idx], "v_buffer": self.v_buffer[idx], "k_scales_zeros": self.k_scales_zeros[idx], "v_scales_zeros": self.v_scales_zeros[idx], - "dtype": "int2", + "dtype": self.dtype, } + if self.k_buffer2 is not None: + buffers["k_buffer2"] = self.k_buffer2[idx] + return buffers def _split_global_locs(self, loc: torch.Tensor): loc64 = loc.to(torch.int64) @@ -596,9 +989,7 @@ def _prepare_hp_kv_tensors( """ if already_rotated: return cache_k.to(self.hp_dtype), cache_v.to(self.hp_dtype) - return self._rotate_kv_inplace( - layer_id, cache_k, cache_v, v_rotation_absorbed - ) + return self._rotate_kv_inplace(layer_id, cache_k, cache_v, v_rotation_absorbed) def _set_hp_kv_buffer( self, @@ -646,33 +1037,49 @@ def _set_quant_kv_buffer_extend( # already be in R_v space (rotation absorbed) — the kernel does # not rotate V. Requires single-scale layout (num_groups == 1) for # both K and V scales/zeros. - if envs.SGLANG_OSCAR_FUSED_ROTATE_CLIP_QUANT.get(): - assert v_rotation_absorbed, ( - "V rotation must be absorbed for fused oscar K-rotation + clip + quant + set" - ) - use_fused_rotate = ( envs.SGLANG_OSCAR_FUSED_ROTATE_CLIP_QUANT.get() and not already_hadamard_transformed and v_rotation_absorbed and clip_on + and self.head_dim == self.v_head_dim and _get_num_scale_groups(self.k_scales_zeros[idx]) == 1 and _get_num_scale_groups(self.v_scales_zeros[idx]) == 1 + and self.dtype != "pq_k_int2v" ) if use_fused_rotate: - quantized_set_kv_int2_oscar_rotate_k_clip_triton( - cache_k.to(self.hp_dtype), - cache_v.to(self.hp_dtype), - self._R_k[idx], - quant_loc, - self.k_buffer[idx], - self.v_buffer[idx], - self.k_scales_zeros[idx], - self.v_scales_zeros[idx], - self._k_clip_ratio, - self._v_clip_ratio, - hp_global_offset=mixed_hp_offset, + assert v_rotation_absorbed, ( + "V rotation must be absorbed for fused oscar K-rotation + clip + quant + set" ) + if self.dtype == "int1": + quantized_set_kv_int1_oscar_rotate_k_clip_triton( + cache_k.to(self.hp_dtype), + cache_v.to(self.hp_dtype), + self._R_k[idx], + quant_loc, + self.k_buffer[idx], + self.v_buffer[idx], + self.k_scales_zeros[idx], + self.v_scales_zeros[idx], + self._k_clip_ratio, + self._v_clip_ratio, + hp_global_offset=mixed_hp_offset, + lloyd_max=self._lloyd_max, + ) + else: + quantized_set_kv_int2_oscar_rotate_k_clip_triton( + cache_k.to(self.hp_dtype), + cache_v.to(self.hp_dtype), + self._R_k[idx], + quant_loc, + self.k_buffer[idx], + self.v_buffer[idx], + self.k_scales_zeros[idx], + self.v_scales_zeros[idx], + self._k_clip_ratio, + self._v_clip_ratio, + hp_global_offset=mixed_hp_offset, + ) return if not already_hadamard_transformed: @@ -683,8 +1090,113 @@ def _set_quant_kv_buffer_extend( cache_k = cache_k.to(self.hp_dtype) cache_v = cache_v.to(self.hp_dtype) - if not clip_on: - quantized_set_kv_int2_pretransformed_triton( + if not clip_on and self.dtype != "pq_k_int2v": + if self.dtype == "int1": + quantized_set_kv_int1_pretransformed_triton( + cache_k, + cache_v, + quant_loc, + self.k_buffer[idx], + self.v_buffer[idx], + self.k_scales_zeros[idx], + self.v_scales_zeros[idx], + hp_global_offset=mixed_hp_offset, + ) + else: + quantized_set_kv_int2_pretransformed_triton( + cache_k, + cache_v, + quant_loc, + self.k_buffer[idx], + self.v_buffer[idx], + self.k_scales_zeros[idx], + self.v_scales_zeros[idx], + hp_global_offset=mixed_hp_offset, + ) + return + + if self.dtype == "pq_k_int2v": + # K: PQ encode (OSCAR rotation already applied via cache_k = rotated). + k_for_pq = ( + cache_k if cache_k.dtype == torch.float16 else cache_k.to(torch.float16) + ) + if k_for_pq.shape[0] > 0: + pq_encode_k( + k_for_pq, + quant_loc, + self.k_buffer[idx], + self.get_pq_codebook(layer_id), + self.get_pq_cb_norms(layer_id), + hp_global_offset=mixed_hp_offset, + ) + # RVQ stage 2 encodes the stage-1 reconstruction residual. + if self._is_rvq: + cb1 = self.get_pq_codebook(layer_id) + recon1 = pq_decode_k_at_locs( + self.k_buffer[idx], + quant_loc, + cb1, + self.head_dim, + hp_global_offset=mixed_hp_offset, + ).to(k_for_pq.dtype) + pq_encode_k( + (k_for_pq - recon1).contiguous(), + quant_loc, + self.k_buffer2[idx], + self.get_rvq_cb2(layer_id), + self.get_rvq_cb2_norms(layer_id), + hp_global_offset=mixed_hp_offset, + ) + + # V: PQ encode into the compact n_sub-byte buffer when configured; + # otherwise retain the INT2 Lloyd-Max path. + if self._pq_v: + if cache_v.shape[0] > 0: + pq_encode_k( + cache_v.to(torch.float16), + quant_loc, + self.v_buffer[idx], + self.get_pq_v_codebook(layer_id), + self.get_pq_v_cb_norms(layer_id), + hp_global_offset=mixed_hp_offset, + ) + else: + v_grouped_ok = _get_num_scale_groups(self.v_scales_zeros[idx]) == 1 + if v_grouped_ok: + _launch_single_clip_int2( + cache_v, + quant_loc, + self.v_buffer[idx], + self.v_scales_zeros[idx], + self._v_clip_ratio, + hp_global_offset=mixed_hp_offset, + lloyd_max=self._lloyd_max, + ) + else: + _launch_grouped_clip_int2( + cache_v, + quant_loc, + self.v_buffer[idx], + self.v_scales_zeros[idx], + self._v_clip_ratio, + hp_global_offset=mixed_hp_offset, + ) + elif self.dtype == "int1": + quantized_set_kv_int1_pretransformed_clip_triton( + cache_k, + cache_v, + quant_loc, + self.k_buffer[idx], + self.v_buffer[idx], + self.k_scales_zeros[idx], + self.v_scales_zeros[idx], + self._k_clip_ratio, + self._v_clip_ratio, + hp_global_offset=mixed_hp_offset, + lloyd_max=self._lloyd_max, + ) + else: + quantized_set_kv_int2_pretransformed_clip_triton( cache_k, cache_v, quant_loc, @@ -692,23 +1204,11 @@ def _set_quant_kv_buffer_extend( self.v_buffer[idx], self.k_scales_zeros[idx], self.v_scales_zeros[idx], + self._k_clip_ratio, + self._v_clip_ratio, hp_global_offset=mixed_hp_offset, + lloyd_max=self._lloyd_max, ) - return - - quantized_set_kv_int2_pretransformed_clip_triton( - cache_k, - cache_v, - quant_loc, - self.k_buffer[idx], - self.v_buffer[idx], - self.k_scales_zeros[idx], - self.v_scales_zeros[idx], - self._k_clip_ratio, - self._v_clip_ratio, - hp_global_offset=mixed_hp_offset, - lloyd_max=self._lloyd_max, - ) def _set_mixed_hp_kv_buffer( self, @@ -777,7 +1277,9 @@ def set_kv_buffer( if loc.numel() == 0: return - layer_id = layer_id_override if layer_id_override is not None else layer.layer_id + layer_id = ( + layer_id_override if layer_id_override is not None else layer.layer_id + ) v_rotation_absorbed = bool(getattr(layer, "oscar_v_rotation_absorbed", False)) if is_decode: @@ -823,6 +1325,8 @@ def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): for l in range(self.layer_num): if tgt_q.numel() > 0: self.k_buffer[l][tgt_q] = self.k_buffer[l][src_q] + if self.k_buffer2 is not None: + self.k_buffer2[l][tgt_q] = self.k_buffer2[l][src_q] self.v_buffer[l][tgt_q] = self.v_buffer[l][src_q] self.k_scales_zeros[l][tgt_q] = self.k_scales_zeros[l][src_q] self.v_scales_zeros[l][tgt_q] = self.v_scales_zeros[l][src_q] diff --git a/sglang-research/python/sglang/srt/model_executor/model_runner.py b/sglang-research/python/sglang/srt/model_executor/model_runner.py index cb3829629..111ff202e 100644 --- a/sglang-research/python/sglang/srt/model_executor/model_runner.py +++ b/sglang-research/python/sglang/srt/model_executor/model_runner.py @@ -2014,7 +2014,7 @@ def configure_kv_cache_dtype(self): f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2" ) self.kv_cache_dtype = self.dtype - elif self.server_args.kv_cache_dtype == "int2": + elif self.server_args.kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): self.kv_cache_dtype = self.server_args.kv_cache_dtype else: raise ValueError( diff --git a/sglang-research/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/sglang-research/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 69dad16b7..6e717240f 100644 --- a/sglang-research/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/sglang-research/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -497,7 +497,7 @@ def _init_pools(self: ModelRunner): enable_mixed_kv = ( envs.SGLANG_ENABLE_MIXED_KV_WINDOWS.get() and _attention_supports_mixed_kv(self.server_args) - and self.kv_cache_dtype == "int2" + and self.kv_cache_dtype in ("int2", "int1", "pq_k_int2v") and not self.is_hybrid_swa and self.server_args.disaggregation_mode in (None, "null") and self.server_args.speculative_algorithm is None @@ -562,10 +562,11 @@ def _init_pools(self: ModelRunner): (hp_prefix_pool + n_q - 1) // n_q * n_q ) logger.info( - "Enable unified mixed KV (int2): prefix=%s recent=%s " + "Enable unified mixed KV (%s): prefix=%s recent=%s " "flush_interval=%s num_quant_pages=%s N_Q=%s " "hp_dtype=%s scale_dtype=%s max_total_num_tokens=%s " "max_req_slots=%s hp_prefix_pool_tokens=%s", + self.kv_cache_dtype, p_tokens, envs.SGLANG_MIXED_KV_RECENT_TOKENS.get(), n_q, @@ -602,17 +603,26 @@ def _init_pools(self: ModelRunner): scale_dtype=scale_dtype, ) else: - # For int2 KV cache, scale/zero dtype is configurable via - # SGLANG_MIXED_KV_SCALE_DTYPE (defaults to float32 to match - # the historical behavior). For non-int2 dtypes the pool - # ignores this kwarg. + if self.kv_cache_dtype in ("int1", "pq_k_int2v"): + raise ValueError( + f"{self.kv_cache_dtype} KV cache requires the unified " + "mixed HP+quant pool. Set " + "SGLANG_ENABLE_MIXED_KV_WINDOWS=1 and use a supported " + "Triton decode configuration without speculative " + "decoding, disaggregation, or hybrid SWA." + ) + # For int2/int1 KV cache, scale/zero dtype is configurable + # via SGLANG_MIXED_KV_SCALE_DTYPE (defaults to float32 to + # match the historical behavior). For non-int2/int1 dtypes + # the pool ignores this kwarg. int2_scale_dtype = None - if self.kv_cache_dtype == "int2": + if self.kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): int2_scale_dtype = resolve_scale_dtype( envs.SGLANG_MIXED_KV_SCALE_DTYPE.get() ) logger.info( - "int2 KV cache: scale_dtype=%s", + "%s KV cache: scale_dtype=%s", + self.kv_cache_dtype, envs.SGLANG_MIXED_KV_SCALE_DTYPE.get(), ) self.token_to_kv_pool = MHATokenToKVPool( diff --git a/sglang-research/python/sglang/srt/model_executor/pool_configurator.py b/sglang-research/python/sglang/srt/model_executor/pool_configurator.py index 1a2c3ed41..2771f94ca 100644 --- a/sglang-research/python/sglang/srt/model_executor/pool_configurator.py +++ b/sglang-research/python/sglang/srt/model_executor/pool_configurator.py @@ -105,22 +105,31 @@ def _get_int_kv_bytes_per_head_pair( group_size: Optional[int], scale_dtype_bytes: int = 4, ) -> int: - """Bytes per *quant-token* per head-pair for the int2 KV cache. + """Bytes per *quant-token* per head-pair for the int2/int1 KV cache. Used by both the non-mixed path and the mixed path: the scheduler's ``max_total_num_tokens`` is denominated in quant tokens (= slot ids on the - int2 tier), and the mixed allocator's ``size`` = ``(num_pages - 1) * N_Q`` - is also in quant tokens, so the leak check + int2/int1 tier), and the mixed allocator's ``size`` = + ``(num_pages - 1) * N_Q`` is also in quant tokens, so the leak check ``size - available - evictable - protected`` closes in a single unit. """ - assert kv_cache_dtype == "int2", ( - f"Only int2 quant KV is supported, got {kv_cache_dtype}" - ) - pack_factor = 4 + if kv_cache_dtype == "int2": + k_pack_factor = 4 + v_pack_factor = 4 + elif kv_cache_dtype == "int1": + k_pack_factor = 8 + v_pack_factor = 8 + elif kv_cache_dtype == "pq_k_int2v": + k_pack_factor = 8 # N_SUB=16 codes = 16 bytes for head_dim=128 + v_pack_factor = 4 # INT2 V + else: + raise AssertionError( + f"Only int2/int1/pq_k_int2v quant KV is supported, got {kv_cache_dtype}" + ) k_groups = _resolve_quant_group_count(k_head_dim, group_size) v_groups = _resolve_quant_group_count(v_head_dim, group_size) - packed_k_bytes = k_head_dim // pack_factor - packed_v_bytes = v_head_dim // pack_factor + packed_k_bytes = k_head_dim // k_pack_factor + packed_v_bytes = v_head_dim // v_pack_factor # Interleaved (scale, zero) per group in ``scale_dtype``. scales_zeros_bytes = 2 * scale_dtype_bytes * (k_groups + v_groups) return packed_k_bytes + packed_v_bytes + scales_zeros_bytes @@ -227,7 +236,7 @@ def _compute_cell_size(self, mr: ModelRunner, num_layers: int) -> int: # ``(head_dim + v_head_dim) * hp_dtype_bytes`` (== 1 HP token worth). # Plus the scales/zeros arena sized for every page's quant view. enable_mixed_kv = ( - kv_cache_dtype == "int2" + kv_cache_dtype in ("int2", "int1", "pq_k_int2v") and envs.SGLANG_ENABLE_MIXED_KV_WINDOWS.get() and _attention_supports_mixed_kv(mr.server_args) and not mr.is_hybrid_swa @@ -235,7 +244,7 @@ def _compute_cell_size(self, mr: ModelRunner, num_layers: int) -> int: and mr.server_args.speculative_algorithm is None ) - if kv_cache_dtype == "int2": + if kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): scale_dtype = resolve_scale_dtype(envs.SGLANG_MIXED_KV_SCALE_DTYPE.get()) scale_bytes = torch.empty(0, dtype=scale_dtype).element_size() else: @@ -252,7 +261,7 @@ def _compute_cell_size(self, mr: ModelRunner, num_layers: int) -> int: model_config.v_head_dim, kv_quant_group_size ) # max_total_num_tokens is denominated in *quant tokens* (slot ids - # on the int2 tier). This matches the unified allocator's + # on the int2/int1 tier). This matches the unified allocator's # scheduler-facing ``size = (num_pages - 1) * N_Q``. bytes_per_head = _get_unified_mixed_kv_bytes_per_quant_token( model_config.head_dim, @@ -264,7 +273,7 @@ def _compute_cell_size(self, mr: ModelRunner, num_layers: int) -> int: n_q, ) kv_size = None - elif kv_cache_dtype == "int2": + elif kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): bytes_per_head = _get_int_kv_bytes_per_head_pair( model_config.head_dim, model_config.v_head_dim, @@ -360,16 +369,16 @@ def __init__(self, mr: ModelRunner): tp_size = get_attention_tp_size() if ( - kv_cache_dtype == "int2" + kv_cache_dtype in ("int2", "int1", "pq_k_int2v") and kv_quant_group_size is not None ): raise ValueError( "--kv-cache-quant-group-size is only supported for the " - "full-attention Triton int2 KV cache path and is not supported " - "with hybrid SWA models" + "full-attention Triton int2/int1/PQ KV cache path and is not " + "supported with hybrid SWA models" ) - if kv_cache_dtype == "int2": + if kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): full_per_token = model_config.get_num_kv_heads(tp_size) * ( _get_int_kv_bytes_per_head_pair( model_config.head_dim, diff --git a/sglang-research/python/sglang/srt/models/llama.py b/sglang-research/python/sglang/srt/models/llama.py index b8ad74015..1b6d7cd8a 100644 --- a/sglang-research/python/sglang/srt/models/llama.py +++ b/sglang-research/python/sglang/srt/models/llama.py @@ -51,6 +51,7 @@ kv_cache_scales_loader, maybe_remap_kv_scale_name, ) +from sglang.srt.models.utils import maybe_absorb_oscar_v_rotation_into_qkv from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_npu, make_layers from sglang.utils import get_exception_traceback @@ -194,6 +195,9 @@ def __init__( quant_config=quant_config, prefix=add_prefix("attn", prefix), ) + # OSCAR: K QQT-rotation is applied pool-side by layer_id; V SST-rotation is + # folded into qkv_proj at load (maybe_absorb_oscar_v_rotation_into_qkv). + self.attn.oscar_v_rotation_absorbed = False def forward_prepare_native(self, positions, hidden_states): qkv, _ = self.qkv_proj(hidden_states) @@ -678,6 +682,10 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): else: logger.warning(f"Parameter {name} not found in params_dict") + maybe_absorb_oscar_v_rotation_into_qkv( + self.model, quant_config=self.quant_config, model_label="Llama" + ) + def get_weights_by_name( self, name: str, truncate_size: int = 100, tp_size: int = 1 ) -> Optional[torch.Tensor]: diff --git a/sglang-research/python/sglang/srt/server_args.py b/sglang-research/python/sglang/srt/server_args.py index 51d20bd7d..7b4d198b9 100644 --- a/sglang-research/python/sglang/srt/server_args.py +++ b/sglang-research/python/sglang/srt/server_args.py @@ -2608,9 +2608,9 @@ def _handle_page_size(self): self.page_size = 1 def _unified_mixed_kv_active(self) -> bool: - """Return True when --kv-cache-dtype int2 + the SGLANG_ENABLE_MIXED_KV_* + """Return True when --kv-cache-dtype int2/int1 + the SGLANG_ENABLE_MIXED_KV_* environment variables would route ModelRunner through the unified - HP+int2 KV pool. Mirrors the gates in + HP+quant KV pool. Mirrors the gates in ``model_runner_kv_cache_mixin._init_pools`` and ``pool_configurator._attention_supports_mixed_kv``. """ @@ -2618,7 +2618,7 @@ def _unified_mixed_kv_active(self) -> bool: if not envs.SGLANG_ENABLE_MIXED_KV_WINDOWS.get(): return False - if self.kv_cache_dtype != "int2": + if self.kv_cache_dtype not in ("int2", "int1", "pq_k_int2v"): return False ab = self.attention_backend pab = self.prefill_attention_backend or ab @@ -4178,8 +4178,10 @@ def add_cli_args(parser: argparse.ArgumentParser): "bfloat16", "fp4_e2m1", "int2", + "int1", + "pq_k_int2v", ], - help='Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+. "int2" uses the Triton quantized KV cache path.', + help='Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+. "int2" uses the Triton quantized KV cache path. "int1" extends int2 to 1-bit packing (8 values/byte). "pq_k_int2v" uses Product-Quantized K (1.0 bpe) with INT2 V, or PQ V when SGLANG_PQ_V_CODEBOOK is configured.', ) parser.add_argument( "--kv-cache-quant-group-size", @@ -6491,29 +6493,29 @@ def check_server_args(self): if self.kv_cache_quant_group_size is not None: if self.kv_cache_quant_group_size <= 0: raise ValueError("--kv-cache-quant-group-size must be positive") - if self.kv_cache_dtype != "int2": + if self.kv_cache_dtype not in ("int2", "int1", "pq_k_int2v"): raise ValueError( "--kv-cache-quant-group-size is only supported with " - "--kv-cache-dtype int2" + "--kv-cache-dtype int2, int1, or pq_k_int2v" ) if self.model_path.lower() not in ["none", "dummy"] and getattr( self.get_model_config(), "is_hybrid_swa", False ): raise ValueError( "--kv-cache-quant-group-size is only supported for the " - "full-attention int2 KV cache path (Triton backend, or " + "full-attention int2/int1 KV cache path (Triton backend, or " "hybrid FA3-prefill + Triton-decode) and is not " "supported with hybrid SWA models" ) - # int2 has no FA3 decode reader path; it lives only in the Triton - # backend. FA3 prefill is supported via HybridAttnBackend + # int2/int1 have no FA3 decode reader path; they live only in the + # Triton backend. FA3 prefill is supported via HybridAttnBackend # (--prefill-attention-backend fa3 --decode-attention-backend - # triton), but sending decode to FA3 with int2 would crash in + # triton), but sending decode to FA3 with int2/int1 would crash in # forward_decode (``q.to("int2")``). Reject the unsupported # combinations up front so we fail at startup with a clear message # rather than a cryptic TypeError mid-forward. - if self.kv_cache_dtype == "int2": + if self.kv_cache_dtype in ("int2", "int1", "pq_k_int2v"): bad_backend = None if ( self.attention_backend not in (None, "triton") @@ -6530,9 +6532,10 @@ def check_server_args(self): ) if bad_backend is not None: raise ValueError( - "--kv-cache-dtype int2 requires the Triton decode path. " - f"Got {bad_backend}. Use either `--attention-backend " - "triton` or `--prefill-attention-backend fa3 " + f"--kv-cache-dtype {self.kv_cache_dtype} requires the " + f"Triton decode path. Got {bad_backend}. Use either " + "`--attention-backend triton` or " + "`--prefill-attention-backend fa3 " "--decode-attention-backend triton`." ) diff --git a/sglang-research/test/registered/kernels/test_unified_int2_kv.py b/sglang-research/test/registered/kernels/test_unified_int2_kv.py index 695dc960d..13c440ec9 100644 --- a/sglang-research/test/registered/kernels/test_unified_int2_kv.py +++ b/sglang-research/test/registered/kernels/test_unified_int2_kv.py @@ -56,6 +56,64 @@ def _identity_rotation_paths(head_dim: int, layer_num: int) -> tuple[str, str]: return k_path, v_path +def _pq_codebook_path( + name: str, + *, + head_dim: int, + layer_num: int, + n_sub: int = 8, + n_centroids: int = 16, + rvq: bool = False, +) -> str: + """Create a deterministic, small per-layer PQ/RVQ fixture.""" + assert head_dim % n_sub == 0 + sub_dim = head_dim // n_sub + generator = torch.Generator().manual_seed(2026) + stage1 = torch.randn( + layer_num, + n_sub, + n_centroids, + sub_dim, + generator=generator, + dtype=torch.float32, + ) + path = os.path.join(_IDENTITY_ROT_DIR, f"{name}.pt") + common = { + "layer_ids": list(range(layer_num)), + "n_sub": n_sub, + "sub_dim": sub_dim, + } + if rvq: + stage2 = 0.25 * torch.randn( + layer_num, + n_sub, + n_centroids, + sub_dim, + generator=generator, + dtype=torch.float32, + ) + torch.save( + { + **common, + "n_cents1": n_centroids, + "n_cents2": n_centroids, + "codebooks_stage1": stage1, + "codebooks_stage2": stage2, + }, + path, + ) + else: + torch.save( + { + **common, + "n_centroids": n_centroids, + "codebooks_per_layer": stage1, + }, + path, + ) + return path + + def _make_pool( num_quant_pages: int = 64, layer_num: int = 2, @@ -67,6 +125,12 @@ def _make_pool( hp_recent_tokens: int = 128, max_req_slots: int = 64, num_hp_prefix_slots: int = 256, + dtype: str = "int2", + k_clip_ratio: float = 0.0, + v_clip_ratio: float = 0.0, + lloyd_max: bool = False, + pq_k_codebook: str = "", + pq_v_codebook: str = "", # Backward-compat for tests that still pass ``num_pages``. num_pages: int = None, ): @@ -77,22 +141,23 @@ def _make_pool( num_quant_pages = num_pages os.environ.setdefault("HADAMARD_ORDER", "16") hp_dtype, scale_dtype = _resolve_dtypes() - assert head_dim == v_head_dim, ( - "identity rotation fixture assumes head_dim == v_head_dim" - ) - k_path, v_path = _identity_rotation_paths(head_dim, layer_num) + k_path, _ = _identity_rotation_paths(head_dim, layer_num) + _, v_path = _identity_rotation_paths(v_head_dim, layer_num) with ( envs.SGLANG_OSCAR_K_ROTATION_PATH.override(k_path), envs.SGLANG_OSCAR_V_ROTATION_PATH.override(v_path), - envs.SGLANG_OSCAR_K_CLIP_RATIO.override(0.0), - envs.SGLANG_OSCAR_V_CLIP_RATIO.override(0.0), + envs.SGLANG_OSCAR_K_CLIP_RATIO.override(k_clip_ratio), + envs.SGLANG_OSCAR_V_CLIP_RATIO.override(v_clip_ratio), + envs.SGLANG_LLOYD_MAX.override(lloyd_max), + envs.SGLANG_PQ_K_CODEBOOK.override(pq_k_codebook), + envs.SGLANG_PQ_V_CODEBOOK.override(pq_v_codebook), ): return UnifiedInt2HPKVPool( num_quant_pages=num_quant_pages, hp_dtype=hp_dtype, hp_prefix_tokens=hp_prefix_tokens, hp_recent_tokens=hp_recent_tokens, - dtype="int2", + dtype=dtype, head_num=head_num, head_dim=head_dim, layer_num=layer_num, @@ -126,7 +191,7 @@ def _make_allocator(pool=None, **overrides): hp_recent_ring_size=pool.hp_recent_ring_size, max_req_slots=pool.max_req_slots, num_hp_prefix_slots=pool.num_hp_prefix_slots, - dtype="int2", + dtype=pool.dtype, hp_dtype=pool.hp_dtype, device="cuda", kvcache=pool, @@ -355,11 +420,9 @@ def leak(evictable: int, protected: int, session: int) -> int: def test_available_size_pools_quant_and_hp_prefix(self): """``available_size`` covers both pooled tiers (quant + HP-prefix).""" a = self.allocator - expected = ( - (a.free_pages.numel() + a.release_pages.numel()) * a.N_Q - + (a.hp_prefix_free_pages.numel() + a.hp_prefix_release_pages.numel()) - * a.N_Q - ) + expected = (a.free_pages.numel() + a.release_pages.numel()) * a.N_Q + ( + a.hp_prefix_free_pages.numel() + a.hp_prefix_release_pages.numel() + ) * a.N_Q self.assertEqual(a.available_size(), expected) @@ -380,7 +443,9 @@ def test_decoupled_arena_shapes(self): ) self.assertEqual(pool.k_buffer[0].shape[0], 32 * pool.N_Q) # HP arena: prefix pool + per-req recent slabs. - expected_hp = pool.num_hp_prefix_slots + pool.max_req_slots * pool.hp_recent_ring_size + expected_hp = ( + pool.num_hp_prefix_slots + pool.max_req_slots * pool.hp_recent_ring_size + ) self.assertEqual(pool.hp_k_buffer[0].shape[0], expected_hp) # HP-recent base aligns with the prefix region size. self.assertEqual(pool.hp_recent_base, pool.num_hp_prefix_slots) @@ -550,6 +615,1091 @@ def test_mixed_hot_paths_avoid_dynamic_cuda_masks(self): self.assertNotIn("cu_seqlens_k_cpu", fa_src) +class OneBitStorageCorrectnessTest(unittest.TestCase): + class _Layer: + layer_id = 0 + + @staticmethod + def _lm_int1_reference(x: torch.Tensor): + x_f = x.float() + mean = x_f.mean(dim=-1) + diff = x_f - mean[..., None] + std = torch.sqrt((diff * diff).mean(dim=-1) + 1e-8) + scale = torch.clamp(2.0 * 0.79788456 * std, min=1e-8) + zero = 0.5 - mean / scale + quant = (x_f >= mean[..., None]).to(torch.uint8) + block_octant = x.shape[-1] // 8 + quant = quant.reshape(*quant.shape[:-1], 8, block_octant) + shifts = torch.arange(8, device=x.device, dtype=torch.uint8).view( + *((1,) * (quant.ndim - 2)), 8, 1 + ) + packed = torch.sum(quant << shifts, dim=-2).to(torch.uint8) + return packed, torch.stack((scale, zero), dim=-1) + + def test_pq_codebook_metadata_mismatch_fails_before_allocation(self): + path = _pq_codebook_path("pq_bad_metadata", head_dim=64, layer_num=1, n_sub=8) + data = torch.load(path, map_location="cpu", weights_only=False) + data["n_sub"] = 7 + torch.save(data, path) + with self.assertRaisesRegex(ValueError, "metadata n_sub"): + _make_pool( + num_pages=2, + layer_num=1, + head_num=1, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=path, + ) + data["n_sub"] = 8 + data["layer_ids"] = [0, 0] + data["codebooks_per_layer"] = data["codebooks_per_layer"].repeat(2, 1, 1, 1) + torch.save(data, path) + with self.assertRaisesRegex(ValueError, "layer_ids must be unique"): + _make_pool( + num_pages=2, + layer_num=1, + head_num=1, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=path, + ) + non_power_path = _pq_codebook_path( + "pq_bad_centroids", + head_dim=64, + layer_num=1, + n_sub=8, + n_centroids=16, + ) + non_power = torch.load(non_power_path, map_location="cpu", weights_only=False) + non_power["codebooks_per_layer"] = non_power["codebooks_per_layer"][:, :, :15] + non_power["n_centroids"] = 15 + torch.save(non_power, non_power_path) + with self.assertRaisesRegex(ValueError, "power of two"): + _make_pool( + num_pages=2, + layer_num=1, + head_num=1, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=non_power_path, + ) + + def test_int1_prefill_single_group_is_byte_exact_for_k_and_v(self): + _ensure_cuda() + from sglang.srt.mem_cache.kv_quant_kernels import ( + gather_dequantize_kv_int1_triton, + ) + + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="int1", + k_clip_ratio=1.0, + v_clip_ratio=1.0, + lloyd_max=True, + ) + self.assertEqual(pool.k_buffer[0].shape[-1], 8) + self.assertEqual(pool.v_buffer[0].shape[-1], 8) + + pool.k_buffer[0].fill_(0xA5) + pool.v_buffer[0].fill_(0xA5) + torch.manual_seed(7) + cache_k = 3.0 * torch.randn( + 3, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + cache_v = 3.0 * torch.randn( + 3, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + loc = torch.tensor([2, 5, 9], dtype=torch.int64, device="cuda") + + pool.set_kv_buffer( + self._Layer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + torch.cuda.synchronize() + + ref_k, ref_k_sz = self._lm_int1_reference(cache_k) + ref_v, ref_v_sz = self._lm_int1_reference(cache_v) + self.assertTrue(torch.equal(pool.k_buffer[0][loc], ref_k)) + self.assertTrue(torch.equal(pool.v_buffer[0][loc], ref_v)) + self.assertTrue((pool.k_buffer[0][3] == 0xA5).all()) + self.assertTrue((pool.v_buffer[0][3] == 0xA5).all()) + torch.testing.assert_close( + pool.k_scales_zeros[0][loc].float(), ref_k_sz, atol=2e-2, rtol=2e-2 + ) + torch.testing.assert_close( + pool.v_scales_zeros[0][loc].float(), ref_v_sz, atol=2e-2, rtol=2e-2 + ) + + v_roundtrip = gather_dequantize_kv_int1_triton( + pool.v_buffer[0], + pool.v_scales_zeros[0], + loc, + pool.v_head_dim, + pool.hp_dtype, + ) + self.assertEqual(v_roundtrip.shape, cache_v.shape) + self.assertTrue(torch.isfinite(v_roundtrip).all()) + + def test_fused_int1_oscar_prefill_matches_lloyd_max_reference(self): + _ensure_cuda() + from sglang.srt.environ import envs + + class _AbsorbedLayer: + layer_id = 0 + oscar_v_rotation_absorbed = True + + kwargs = dict( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="int1", + k_clip_ratio=1.0, + v_clip_ratio=1.0, + lloyd_max=True, + ) + fused_pool = _make_pool(**kwargs) + reference_pool = _make_pool(**kwargs) + torch.manual_seed(9) + cache_k = torch.randn( + 3, + fused_pool.head_num, + fused_pool.head_dim, + dtype=fused_pool.hp_dtype, + device="cuda", + ) + cache_v = torch.randn_like(cache_k) + loc = torch.tensor([2, 5, 9], dtype=torch.int64, device="cuda") + + reference_pool.set_kv_buffer( + self._Layer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + with envs.SGLANG_OSCAR_FUSED_ROTATE_CLIP_QUANT.override(True): + fused_pool.set_kv_buffer( + _AbsorbedLayer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=False, + ) + torch.cuda.synchronize() + + self.assertTrue( + torch.equal(fused_pool.k_buffer[0][loc], reference_pool.k_buffer[0][loc]) + ) + self.assertTrue( + torch.equal(fused_pool.v_buffer[0][loc], reference_pool.v_buffer[0][loc]) + ) + torch.testing.assert_close( + fused_pool.k_scales_zeros[0][loc], + reference_pool.k_scales_zeros[0][loc], + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + fused_pool.v_scales_zeros[0][loc], + reference_pool.v_scales_zeros[0][loc], + atol=2e-2, + rtol=2e-2, + ) + + def test_pq_v_prefill_writes_only_compact_code_rows(self): + _ensure_cuda() + from sglang.QuantKernel.oscar_rotation_pq_k_kv import ( + pq_decode_k, + pq_encode_k, + ) + from sglang.srt.layers.attention.quantized_kv_prefill import ( + dequantize_prefix_kv, + ) + + k_path = _pq_codebook_path("pq_k_prefill", head_dim=64, layer_num=1, n_sub=8) + v_path = _pq_codebook_path("pq_v_prefill", head_dim=64, layer_num=1, n_sub=8) + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=k_path, + pq_v_codebook=v_path, + ) + self.assertEqual(pool.k_buffer[0].shape[-1], 8) + self.assertEqual(pool.v_buffer[0].shape[-1], 8) + + pool.k_buffer[0].fill_(0xA5) + pool.v_buffer[0].fill_(0xA5) + torch.manual_seed(11) + cache_k = torch.randn( + 2, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + cache_v = torch.randn( + 2, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + loc = torch.tensor([4, 8], dtype=torch.int64, device="cuda") + + ref_k = torch.full_like(pool.k_buffer[0], 0xA5) + ref_v = torch.full_like(pool.v_buffer[0], 0xA5) + pq_encode_k( + cache_k, + loc, + ref_k, + pool.get_pq_codebook(0), + pool.get_pq_cb_norms(0), + ) + pq_encode_k( + cache_v, + loc, + ref_v, + pool.get_pq_v_codebook(0), + pool.get_pq_v_cb_norms(0), + ) + pool.set_kv_buffer( + self._Layer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + torch.cuda.synchronize() + + self.assertTrue(torch.equal(pool.k_buffer[0][loc], ref_k[loc])) + self.assertTrue(torch.equal(pool.v_buffer[0][loc], ref_v[loc])) + self.assertTrue((pool.k_buffer[0][5] == 0xA5).all()) + self.assertTrue((pool.v_buffer[0][5] == 0xA5).all()) + + got_k, got_v = dequantize_prefix_kv(pool, 0, loc, pool.hp_dtype) + expected_k = pq_decode_k( + ref_k[loc], pool.get_pq_codebook(0), len(loc), pool.head_dim + ).to(pool.hp_dtype) + expected_v = pq_decode_k( + ref_v[loc], pool.get_pq_v_codebook(0), len(loc), pool.v_head_dim + ).to(pool.hp_dtype) + self.assertTrue(torch.equal(got_k, expected_k)) + self.assertTrue(torch.equal(got_v, expected_v)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pool.set_kv_buffer( + self._Layer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + graph.replay() + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool.k_buffer[0][loc], ref_k[loc])) + self.assertTrue(torch.equal(pool.v_buffer[0][loc], ref_v[loc])) + + def test_rvq_prefill_and_cache_move_preserve_stage2_codes(self): + _ensure_cuda() + from sglang.QuantKernel.oscar_rotation_pq_k_kv import ( + pq_decode_k, + pq_encode_k, + ) + + rvq_path = _pq_codebook_path( + "rvq_k_prefill", head_dim=64, layer_num=1, n_sub=8, rvq=True + ) + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=rvq_path, + ) + self.assertIsNotNone(pool.k_buffer2) + + torch.manual_seed(19) + cache_k = torch.randn( + 2, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + cache_v = torch.randn( + 2, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + src = torch.tensor([3, 7], dtype=torch.int64, device="cuda") + + ref_stage1 = torch.zeros_like(pool.k_buffer[0]) + ref_stage2 = torch.zeros_like(pool.k_buffer2[0]) + pq_encode_k( + cache_k, + src, + ref_stage1, + pool.get_pq_codebook(0), + pool.get_pq_cb_norms(0), + ) + recon1 = pq_decode_k( + ref_stage1[src], + pool.get_pq_codebook(0), + len(src), + pool.head_dim, + ).to(cache_k.dtype) + pq_encode_k( + (cache_k - recon1).contiguous(), + src, + ref_stage2, + pool.get_rvq_cb2(0), + pool.get_rvq_cb2_norms(0), + ) + + pool.set_kv_buffer( + self._Layer(), + src, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool.k_buffer[0][src], ref_stage1[src])) + self.assertTrue(torch.equal(pool.k_buffer2[0][src], ref_stage2[src])) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pool.set_kv_buffer( + self._Layer(), + src, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + graph.replay() + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool.k_buffer[0][src], ref_stage1[src])) + self.assertTrue(torch.equal(pool.k_buffer2[0][src], ref_stage2[src])) + + dst = torch.tensor([12, 13], dtype=torch.int64, device="cuda") + pool.move_kv_cache(dst, src) + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool.k_buffer[0][dst], ref_stage1[src])) + self.assertTrue(torch.equal(pool.k_buffer2[0][dst], ref_stage2[src])) + + def test_int1_unified_decode_is_graph_safe_and_matches_dense_reference(self): + _ensure_cuda() + from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd_int1_unified, + ) + from sglang.srt.mem_cache.kv_quant_kernels import ( + gather_dequantize_kv_int1_triton, + ) + + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="int1", + k_clip_ratio=1.0, + v_clip_ratio=1.0, + lloyd_max=True, + ) + torch.manual_seed(29) + quant_k = torch.randn( + 4, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + quant_v = torch.randn( + 4, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + quant_locs = torch.tensor([4, 5, 6, 7], dtype=torch.int64, device="cuda") + pool.set_kv_buffer( + self._Layer(), + quant_locs, + quant_k, + quant_v, + already_hadamard_transformed=True, + ) + + pool.hp_k_buffer[0][0:2] = torch.randn( + 2, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + pool.hp_v_buffer[0][0:2] = torch.randn( + 2, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + + # GQA=18 exercises a non-power-of-two group larger than BLOCK_H=16. + bs, q_heads = 2, 36 + q = torch.randn(bs, q_heads, pool.head_dim, dtype=pool.hp_dtype, device="cuda") + hp_indptr = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") + hp_indices = torch.tensor([0, 1], dtype=torch.int64, device="cuda") + quant_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + # Deliberately leave an invalid padded tail. A graph-safe kernel must + # obey indptr and never derive a Python slice length from it. + quant_indices = torch.full( + (16,), + pool.k_buffer[0].shape[0] + 123, + dtype=torch.int64, + device="cuda", + ) + quant_indices[:4] = quant_locs + hp_splits = torch.ones((bs,), dtype=torch.int32, device="cuda") + quant_splits = torch.ones((bs,), dtype=torch.int32, device="cuda") + attn_logits = torch.empty( + (bs, q_heads, 2, pool.v_head_dim), + dtype=torch.float32, + device="cuda", + ) + attn_lse = torch.empty((bs, q_heads, 2), dtype=torch.float32, device="cuda") + sm_scale = pool.head_dim**-0.5 + + def launch(q_arg, out_arg): + return decode_attention_fwd_int1_unified( + q_arg, + pool.hp_k_buffer[0], + pool.hp_v_buffer[0], + pool.k_buffer[0], + pool.v_buffer[0], + pool.k_scales_zeros[0], + pool.v_scales_zeros[0], + out_arg, + hp_indptr, + hp_indices, + quant_indptr, + quant_indices, + attn_logits, + attn_lse, + hp_splits, + quant_splits, + 1, + 1, + sm_scale, + ) + + dequant_k = gather_dequantize_kv_int1_triton( + pool.k_buffer[0], + pool.k_scales_zeros[0], + quant_locs, + pool.head_dim, + pool.hp_dtype, + ) + dequant_v = gather_dequantize_kv_int1_triton( + pool.v_buffer[0], + pool.v_scales_zeros[0], + quant_locs, + pool.v_head_dim, + pool.hp_dtype, + ) + + def dense_reference(q_arg): + ref = torch.empty_like(q_arg) + kv_group = q_heads // pool.head_num + for batch_idx in range(bs): + keys = torch.cat( + ( + pool.hp_k_buffer[0][batch_idx : batch_idx + 1], + dequant_k[batch_idx * 2 : batch_idx * 2 + 2], + ), + dim=0, + ).float() + values = torch.cat( + ( + pool.hp_v_buffer[0][batch_idx : batch_idx + 1], + dequant_v[batch_idx * 2 : batch_idx * 2 + 2], + ), + dim=0, + ).float() + for q_head in range(q_heads): + kv_head = q_head // kv_group + scores = ( + keys[:, kv_head] @ q_arg[batch_idx, q_head].float() + ) * sm_scale + ref[batch_idx, q_head] = ( + (torch.softmax(scores, dim=0)[:, None] * values[:, kv_head]) + .sum(dim=0) + .to(ref.dtype) + ) + return ref + + eager_out = torch.empty_like(q) + launch(q, eager_out) + torch.cuda.synchronize() + torch.testing.assert_close(eager_out, dense_reference(q), atol=4e-2, rtol=4e-2) + + static_q = q.clone() + static_out = torch.empty_like(q) + # Warm every Triton specialization before capture. + launch(static_q, static_out) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(static_q, static_out) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, dense_reference(static_q), atol=4e-2, rtol=4e-2 + ) + + replay_q = torch.randn_like(static_q) + static_q.copy_(replay_q) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, dense_reference(replay_q), atol=4e-2, rtol=4e-2 + ) + + def test_pq_unified_decode_is_graph_safe_and_matches_dense_reference(self): + _ensure_cuda() + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_decode_k + from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd_pqk_int2v_unified, + ) + + k_path = _pq_codebook_path( + "pq_k_decode_graph", head_dim=64, layer_num=1, n_sub=8 + ) + v_path = _pq_codebook_path( + "pq_v_decode_graph", head_dim=64, layer_num=1, n_sub=8 + ) + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=64, + dtype="pq_k_int2v", + pq_k_codebook=k_path, + pq_v_codebook=v_path, + ) + torch.manual_seed(31) + quant_k = torch.randn( + 4, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + quant_v = torch.randn( + 4, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + quant_locs = torch.tensor([4, 5, 6, 7], dtype=torch.int64, device="cuda") + pool.set_kv_buffer( + self._Layer(), + quant_locs, + quant_k, + quant_v, + already_hadamard_transformed=True, + ) + pool.hp_k_buffer[0][0:2] = torch.randn( + 2, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + pool.hp_v_buffer[0][0:2] = torch.randn( + 2, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + + # GQA=6 deliberately exercises a non-power-of-two group. + bs, q_heads = 2, 12 + q = torch.randn(bs, q_heads, pool.head_dim, dtype=pool.hp_dtype, device="cuda") + hp_indptr = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") + hp_indices = torch.tensor([0, 1], dtype=torch.int64, device="cuda") + quant_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + quant_indices = torch.full( + (16,), + pool.k_buffer[0].shape[0] + 321, + dtype=torch.int64, + device="cuda", + ) + quant_indices[:4] = quant_locs + hp_splits = torch.ones((bs,), dtype=torch.int32, device="cuda") + quant_splits = torch.ones((bs,), dtype=torch.int32, device="cuda") + attn_logits = torch.empty( + (bs, q_heads, 2, pool.v_head_dim), + dtype=torch.float32, + device="cuda", + ) + attn_lse = torch.empty((bs, q_heads, 2), dtype=torch.float32, device="cuda") + sm_scale = pool.head_dim**-0.5 + + def launch(q_arg, out_arg): + return decode_attention_fwd_pqk_int2v_unified( + q_arg, + pool.hp_k_buffer[0], + pool.hp_v_buffer[0], + pool.k_buffer[0], + pool.v_buffer[0], + pool.k_scales_zeros[0], + pool.v_scales_zeros[0], + pool.get_pq_codebook(0), + out_arg, + hp_indptr, + hp_indices, + quant_indptr, + quant_indices, + attn_logits, + attn_lse, + hp_splits, + quant_splits, + 1, + 1, + sm_scale, + pq_v_codebook=pool.get_pq_v_codebook(0), + ) + + dequant_k = pq_decode_k( + pool.k_buffer[0][quant_locs], + pool.get_pq_codebook(0), + len(quant_locs), + pool.head_dim, + ).to(pool.hp_dtype) + dequant_v = pq_decode_k( + pool.v_buffer[0][quant_locs], + pool.get_pq_v_codebook(0), + len(quant_locs), + pool.v_head_dim, + ).to(pool.hp_dtype) + + def dense_reference(q_arg): + ref = torch.empty_like(q_arg) + kv_group = q_heads // pool.head_num + for batch_idx in range(bs): + keys = torch.cat( + ( + pool.hp_k_buffer[0][batch_idx : batch_idx + 1], + dequant_k[batch_idx * 2 : batch_idx * 2 + 2], + ), + dim=0, + ).float() + values = torch.cat( + ( + pool.hp_v_buffer[0][batch_idx : batch_idx + 1], + dequant_v[batch_idx * 2 : batch_idx * 2 + 2], + ), + dim=0, + ).float() + for q_head in range(q_heads): + kv_head = q_head // kv_group + scores = ( + keys[:, kv_head] @ q_arg[batch_idx, q_head].float() + ) * sm_scale + ref[batch_idx, q_head] = ( + (torch.softmax(scores, dim=0)[:, None] * values[:, kv_head]) + .sum(dim=0) + .to(ref.dtype) + ) + return ref + + eager_out = torch.empty_like(q) + launch(q, eager_out) + torch.cuda.synchronize() + torch.testing.assert_close(eager_out, dense_reference(q), atol=4e-2, rtol=4e-2) + + static_q = q.clone() + static_out = torch.empty_like(q) + launch(static_q, static_out) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(static_q, static_out) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, dense_reference(static_q), atol=4e-2, rtol=4e-2 + ) + + replay_q = torch.randn_like(static_q) + static_q.copy_(replay_q) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, dense_reference(replay_q), atol=4e-2, rtol=4e-2 + ) + + from sglang.srt.environ import envs + + with envs.SGLANG_PQ_USE_ADC.override(1): + adc_q = q.clone() + adc_out = torch.empty_like(q) + launch(adc_q, adc_out) + torch.cuda.synchronize() + torch.testing.assert_close( + adc_out, dense_reference(adc_q), atol=4e-2, rtol=4e-2 + ) + adc_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(adc_graph): + launch(adc_q, adc_out) + adc_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + adc_out, dense_reference(adc_q), atol=4e-2, rtol=4e-2 + ) + + def test_graph_decode_supports_distinct_k_and_v_head_dims(self): + _ensure_cuda() + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_decode_k + from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd_int1_unified, + decode_attention_fwd_pqk_int2v_unified, + ) + from sglang.srt.mem_cache.kv_quant_kernels import ( + gather_dequantize_kv_int1_triton, + ) + + for dtype in ("int1", "pq_k_int2v"): + with self.subTest(dtype=dtype): + kwargs = dict( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=64, + v_head_dim=32, + dtype=dtype, + ) + if dtype == "int1": + kwargs.update( + k_clip_ratio=1.0, + v_clip_ratio=1.0, + lloyd_max=True, + ) + else: + kwargs.update( + pq_k_codebook=_pq_codebook_path( + "pq_k_distinct_dims", + head_dim=64, + layer_num=1, + n_sub=8, + ), + pq_v_codebook=_pq_codebook_path( + "pq_v_distinct_dims", + head_dim=32, + layer_num=1, + n_sub=4, + ), + ) + pool = _make_pool(**kwargs) + + torch.manual_seed(43) + cache_k = torch.randn( + 3, + pool.head_num, + pool.head_dim, + dtype=pool.hp_dtype, + device="cuda", + ) + cache_v = torch.randn( + 3, + pool.head_num, + pool.v_head_dim, + dtype=pool.hp_dtype, + device="cuda", + ) + loc = torch.tensor([4, 5, 6], dtype=torch.int64, device="cuda") + pool.set_kv_buffer( + self._Layer(), + loc, + cache_k, + cache_v, + already_hadamard_transformed=True, + ) + + q_heads = 6 # GQA=3, also non-power-of-two. + q = torch.randn( + 1, + q_heads, + pool.head_dim, + dtype=pool.hp_dtype, + device="cuda", + ) + hp_indptr = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + hp_indices = torch.empty((0,), dtype=torch.int64, device="cuda") + quant_indptr = torch.tensor([0, 3], dtype=torch.int32, device="cuda") + quant_indices = torch.full( + (8,), + pool.k_buffer[0].shape[0] + 55, + dtype=torch.int64, + device="cuda", + ) + quant_indices[:3] = loc + splits = torch.ones((1,), dtype=torch.int32, device="cuda") + attn_logits = torch.empty( + (1, q_heads, 2, pool.v_head_dim), + dtype=torch.float32, + device="cuda", + ) + attn_lse = torch.empty( + (1, q_heads, 2), dtype=torch.float32, device="cuda" + ) + sm_scale = pool.head_dim**-0.5 + + if dtype == "int1": + + def launch(q_arg, out_arg): + return decode_attention_fwd_int1_unified( + q_arg, + pool.hp_k_buffer[0], + pool.hp_v_buffer[0], + pool.k_buffer[0], + pool.v_buffer[0], + pool.k_scales_zeros[0], + pool.v_scales_zeros[0], + out_arg, + hp_indptr, + hp_indices, + quant_indptr, + quant_indices, + attn_logits, + attn_lse, + splits, + splits, + 1, + 1, + sm_scale, + ) + + reconstructed_k = gather_dequantize_kv_int1_triton( + pool.k_buffer[0], + pool.k_scales_zeros[0], + loc, + pool.head_dim, + pool.hp_dtype, + ) + reconstructed_v = gather_dequantize_kv_int1_triton( + pool.v_buffer[0], + pool.v_scales_zeros[0], + loc, + pool.v_head_dim, + pool.hp_dtype, + ) + else: + + def launch(q_arg, out_arg): + return decode_attention_fwd_pqk_int2v_unified( + q_arg, + pool.hp_k_buffer[0], + pool.hp_v_buffer[0], + pool.k_buffer[0], + pool.v_buffer[0], + pool.k_scales_zeros[0], + pool.v_scales_zeros[0], + pool.get_pq_codebook(0), + out_arg, + hp_indptr, + hp_indices, + quant_indptr, + quant_indices, + attn_logits, + attn_lse, + splits, + splits, + 1, + 1, + sm_scale, + pq_v_codebook=pool.get_pq_v_codebook(0), + ) + + reconstructed_k = pq_decode_k( + pool.k_buffer[0][loc], + pool.get_pq_codebook(0), + len(loc), + pool.head_dim, + ).to(pool.hp_dtype) + reconstructed_v = pq_decode_k( + pool.v_buffer[0][loc], + pool.get_pq_v_codebook(0), + len(loc), + pool.v_head_dim, + ).to(pool.hp_dtype) + + def reference(q_arg): + ref = torch.empty( + (1, q_heads, pool.v_head_dim), + dtype=q_arg.dtype, + device=q_arg.device, + ) + kv_group = q_heads // pool.head_num + for q_head in range(q_heads): + kv_head = q_head // kv_group + scores = ( + reconstructed_k[:, kv_head].float() + @ q_arg[0, q_head].float() + ) * sm_scale + ref[0, q_head] = ( + ( + torch.softmax(scores, dim=0)[:, None] + * reconstructed_v[:, kv_head].float() + ) + .sum(dim=0) + .to(ref.dtype) + ) + return ref + + static_q = q.clone() + static_out = torch.empty( + (1, q_heads, pool.v_head_dim), + dtype=q.dtype, + device=q.device, + ) + launch(static_q, static_out) + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, reference(static_q), atol=4e-2, rtol=4e-2 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(static_q, static_out) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, reference(static_q), atol=4e-2, rtol=4e-2 + ) + + def test_rvq_k_int2_v_decode_graph_path_matches_reconstruction(self): + _ensure_cuda() + from sglang.QuantKernel.oscar_rotation_pq_k_kv import pq_decode_k + from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd_pqk_int2v_unified, + ) + from sglang.srt.mem_cache.kv_quant_kernels import dequantize_kv_int2_triton + + rvq_path = _pq_codebook_path( + "rvq_k_decode_graph", + head_dim=128, + layer_num=1, + n_sub=16, + rvq=True, + ) + pool = _make_pool( + num_pages=16, + layer_num=1, + head_num=2, + head_dim=128, + v_head_dim=128, + dtype="pq_k_int2v", + pq_k_codebook=rvq_path, + ) + torch.manual_seed(37) + quant_k = torch.randn( + 3, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + quant_v = torch.randn( + 3, pool.head_num, pool.v_head_dim, dtype=pool.hp_dtype, device="cuda" + ) + loc = torch.tensor([4, 5, 6], dtype=torch.int64, device="cuda") + pool.set_kv_buffer( + self._Layer(), + loc, + quant_k, + quant_v, + already_hadamard_transformed=True, + ) + + q = torch.randn( + 1, pool.head_num, pool.head_dim, dtype=pool.hp_dtype, device="cuda" + ) + hp_indptr = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + hp_indices = torch.zeros((1,), dtype=torch.int64, device="cuda") + quant_indptr = torch.tensor([0, 3], dtype=torch.int32, device="cuda") + quant_indices = torch.full( + (8,), + pool.k_buffer[0].shape[0] + 99, + dtype=torch.int64, + device="cuda", + ) + quant_indices[:3] = loc + hp_splits = torch.ones((1,), dtype=torch.int32, device="cuda") + quant_splits = torch.ones((1,), dtype=torch.int32, device="cuda") + attn_logits = torch.empty( + (1, pool.head_num, 2, pool.head_dim), + dtype=torch.float32, + device="cuda", + ) + attn_lse = torch.empty( + (1, pool.head_num, 2), dtype=torch.float32, device="cuda" + ) + sm_scale = pool.head_dim**-0.5 + + def launch(q_arg, out_arg): + return decode_attention_fwd_pqk_int2v_unified( + q_arg, + pool.hp_k_buffer[0], + pool.hp_v_buffer[0], + pool.k_buffer[0], + pool.v_buffer[0], + pool.k_scales_zeros[0], + pool.v_scales_zeros[0], + pool.get_pq_codebook(0), + out_arg, + hp_indptr, + hp_indices, + quant_indptr, + quant_indices, + attn_logits, + attn_lse, + hp_splits, + quant_splits, + 1, + 1, + sm_scale, + quant_k_buffer2=pool.get_raw_key_buffer2(0), + pq_codebook2=pool.get_rvq_cb2(0), + ) + + reconstructed_k = ( + pq_decode_k( + pool.k_buffer[0][loc], + pool.get_pq_codebook(0), + len(loc), + pool.head_dim, + ) + + pq_decode_k( + pool.k_buffer2[0][loc], + pool.get_rvq_cb2(0), + len(loc), + pool.head_dim, + ) + ).to(pool.hp_dtype) + reconstructed_v = dequantize_kv_int2_triton( + pool.v_buffer[0][loc], + pool.v_scales_zeros[0][loc], + pool.v_head_dim, + pool.hp_dtype, + ) + + def reference(q_arg): + ref = torch.empty_like(q_arg) + for head in range(pool.head_num): + scores = ( + reconstructed_k[:, head].float() @ q_arg[0, head].float() + ) * sm_scale + ref[0, head] = ( + ( + torch.softmax(scores, dim=0)[:, None] + * reconstructed_v[:, head].float() + ) + .sum(dim=0) + .to(ref.dtype) + ) + return ref + + static_q = q.clone() + static_out = torch.empty_like(q) + launch(static_q, static_out) + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, reference(static_q), atol=4e-2, rtol=4e-2 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(static_q, static_out) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + static_out, reference(static_q), atol=4e-2, rtol=4e-2 + ) + + class GpuFlushInt2Test(unittest.TestCase): """Round-trip tests for ``gpu_flush_int2``. @@ -980,9 +2130,7 @@ class RadixCacheMixedTrimTest(unittest.TestCase): function reads. This keeps the test pure-Python and CUDA-free. """ - def _trim( - self, committed_len, hp_prefix, hp_recent, flush_interval, page_size - ): + def _trim(self, committed_len, hp_prefix, hp_recent, flush_interval, page_size): """Re-derive trim via the same formula path as ``_mixed_kv_tail_to_drop``. We replicate the logic here rather than instantiate a full @@ -1075,9 +2223,7 @@ def _reference_build( quant_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=device) hp_indptr[1:] = torch.cumsum(hp_lens, dim=0) quant_indptr[1:] = torch.cumsum(quant_lens, dim=0) - hp_parts = [ - rows[i][hp_mask[i]] - hp_offset for i in range(bs) - ] + hp_parts = [rows[i][hp_mask[i]] - hp_offset for i in range(bs)] quant_parts = [rows[i][quant_mask[i]] for i in range(bs)] hp_flat = ( torch.cat(hp_parts) @@ -1155,9 +2301,7 @@ def test_matches_python_reference(self): # Randomly classify each (req, pos) token as HP or quant by choosing a # slot id above or below ``hp_offset``. Positions >= seq_len are 0 # (the test expects them to be ignored via the valid_mask). - req_to_token = torch.zeros( - (bs, max_ctx), dtype=torch.int32, device="cuda" - ) + req_to_token = torch.zeros((bs, max_ctx), dtype=torch.int32, device="cuda") for i in range(bs): n = int(seq_lens[i].item()) tier = torch.randint(0, 2, (n,), device="cuda") # 0=quant, 1=hp @@ -1193,9 +2337,7 @@ def test_matches_python_reference(self): torch.cuda.synchronize() self.assertTrue(torch.equal(hp_lens, ref_hp_lens), "hp_lens mismatch") - self.assertTrue( - torch.equal(quant_lens, ref_quant_lens), "quant_lens mismatch" - ) + self.assertTrue(torch.equal(quant_lens, ref_quant_lens), "quant_lens mismatch") self.assertTrue(torch.equal(hp_indptr, ref_hp_indptr), "hp_indptr mismatch") self.assertTrue( torch.equal(quant_indptr, ref_quant_indptr), "quant_indptr mismatch" @@ -1226,9 +2368,7 @@ def test_no_masked_select_sync(self): max_ctx = 1024 seq_lens = torch.full((bs,), 800, dtype=torch.int32, device="cuda") req_pool_indices = torch.arange(bs, dtype=torch.int64, device="cuda") - req_to_token = torch.zeros( - (bs, max_ctx), dtype=torch.int32, device="cuda" - ) + req_to_token = torch.zeros((bs, max_ctx), dtype=torch.int32, device="cuda") req_to_token[:, :800] = hp_offset + 1 # all HP hp_lens, quant_lens, *_ = self._build_via_triton(