1dc2d854c1
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
1470 lines
52 KiB
Python
1470 lines
52 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
|
|
preparation.
|
|
|
|
- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into
|
|
the paged cache.
|
|
- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged
|
|
cache for sparse/SWA prefill.
|
|
- compute_global_topk_indices_and_lens: map local topk indices to global KV
|
|
cache slots and count valid entries.
|
|
- combine_topk_swa_indices: concatenate topk compressed indices with SWA
|
|
window indices for sparse prefill.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
get_fp8_min_max,
|
|
)
|
|
from vllm.model_executor.warmup.jit_warmup import (
|
|
VllmJitKernel,
|
|
_WarmupInputRows,
|
|
zip_inputs,
|
|
)
|
|
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
|
|
TritonPointerInputVariant,
|
|
TritonWarmupTensor,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.utils.import_utils import has_cutedsl
|
|
from vllm.utils.math_utils import next_power_of_2
|
|
|
|
# Per-token byte width of the two paged fp8 records. Both put a page's whole
|
|
# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the
|
|
# width the KV-cache spec publishes -- identifies the record:
|
|
# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64
|
|
# dims each plus a pad byte.
|
|
# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0
|
|
# scales of 32 dims each. FlashMLA's ``ModelType::V41``.
|
|
# V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in
|
|
# the low nibble), then 32 e4m3 scales of 16 dims each.
|
|
# FlashMLA's ``ModelType::V41_FP4``, compressed cache only.
|
|
V4_BYTES_PER_TOKEN = 584
|
|
V41_BYTES_PER_TOKEN = 528
|
|
V41_QUANT_BLOCK = 32
|
|
V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16
|
|
V41_NVFP4_BYTES_PER_TOKEN = 288
|
|
V41_NVFP4_QUANT_BLOCK = 16
|
|
V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK # 32
|
|
|
|
|
|
@triton.jit
|
|
def quantize_and_insert_k_kernel(
|
|
# Input tensors
|
|
k_ptr, # [num_tokens, 512] bf16
|
|
slot_mapping_ptr, # [num_tokens] int64
|
|
# Output tensor
|
|
k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view)
|
|
# Dimensions
|
|
num_tokens,
|
|
input_dim: tl.constexpr, # 512
|
|
fp8_dim: tl.constexpr, # 448
|
|
bf16_dim: tl.constexpr, # 64
|
|
scale_dim: tl.constexpr, # 8
|
|
quant_block: tl.constexpr, # 64 (quantization block size)
|
|
cache_block_size: tl.constexpr, # 64 (paged cache block size)
|
|
token_data_size: tl.constexpr, # 576 bytes per token data
|
|
block_stride: tl.constexpr, # total bytes per block (padded)
|
|
fp8_max: tl.constexpr,
|
|
n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding)
|
|
use_fnuz: tl.constexpr = False,
|
|
):
|
|
"""Quantize K tensor and insert into paged K cache.
|
|
|
|
K Cache block layout (block_size=64 tokens):
|
|
- [0, 64*576): Token data, each token has 448 fp8 + 128 bf16
|
|
- [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales
|
|
- [64*576 + 64*8, block_stride): Padding
|
|
|
|
One program per token.
|
|
|
|
``use_fnuz=True`` selects FNUZ (``tl.float8e4b8``); default OCP
|
|
(``tl.float8e4nv``) matches every production caller.
|
|
"""
|
|
pid = tl.program_id(0)
|
|
|
|
if pid >= num_tokens:
|
|
return
|
|
|
|
# Get slot mapping
|
|
slot_idx = tl.load(slot_mapping_ptr + pid)
|
|
if slot_idx == -1:
|
|
return
|
|
|
|
block_idx = slot_idx // cache_block_size
|
|
pos_in_block = slot_idx % cache_block_size
|
|
|
|
# Input pointer for this token
|
|
input_row_ptr = k_ptr + pid * input_dim
|
|
|
|
# int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks
|
|
# (e.g. >= 57K at block_stride ~37K). Matches gather path below.
|
|
cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
|
|
|
|
# Token data pointer: token data is stored contiguously at start of block
|
|
# Each token's data is at offset pos_in_block * token_data_size
|
|
token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
|
|
|
|
# Scale pointer: scales are stored after ALL token data in the block
|
|
# Scale for this token is at offset (64 * 576) + pos_in_block * 8
|
|
token_scale_ptr = (
|
|
cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim
|
|
)
|
|
|
|
# Token data layout: [0:448] fp8, [448:576] bf16
|
|
token_fp8_ptr = token_data_ptr
|
|
token_bf16_ptr = token_data_ptr + fp8_dim
|
|
|
|
# ========== Quantize and store FP8 portion (first 448 elements) ==========
|
|
# Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent)
|
|
for qblock_idx in tl.static_range(n_quant_blocks):
|
|
qblock_start = qblock_idx * quant_block
|
|
|
|
if qblock_start < fp8_dim:
|
|
offsets = qblock_start + tl.arange(0, quant_block)
|
|
mask = offsets < fp8_dim
|
|
|
|
# Load bf16 input
|
|
x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0)
|
|
|
|
# Compute absmax scale (same as CUDA kernel)
|
|
abs_x = tl.abs(x)
|
|
block_max = tl.max(abs_x, axis=0)
|
|
block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4)
|
|
|
|
# UE8M0: Round scale UP to next power of 2
|
|
# scale = 2^ceil(log2(block_max / fp8_max))
|
|
raw_scale = block_max / fp8_max
|
|
log_scale = tl.log2(raw_scale)
|
|
exponent = tl.ceil(log_scale) # Round UP to next integer exponent
|
|
scale = tl.exp2(exponent) # scale = 2^exponent (power of 2)
|
|
|
|
# Quantize to fp8: fp8_value = bf16_value / scale
|
|
x_scaled = x / scale
|
|
x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max)
|
|
|
|
# Convert to fp8 (FNUZ on gfx942, OCP elsewhere), then bitcast to uint8.
|
|
if use_fnuz:
|
|
x_fp8 = x_clamped.to(tl.float8e4b8)
|
|
else:
|
|
x_fp8 = x_clamped.to(tl.float8e4nv)
|
|
x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
|
|
|
|
# Store as uint8 (1 byte each)
|
|
tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask)
|
|
|
|
# UE8M0 scale encoding: stored_value = exponent + 127 (bias)
|
|
# During dequant: scale = 2^(stored_value - 127)
|
|
encoded_scale = exponent + 127.0
|
|
encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0)
|
|
tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8))
|
|
|
|
# Padding scale at index 7
|
|
tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8))
|
|
|
|
# ========== Store BF16 portion (last 64 elements, no quantization) ==========
|
|
bf16_input_offset = fp8_dim
|
|
|
|
# Process bf16 in chunks of 16
|
|
bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
|
|
for i in tl.static_range(bf16_dim // 16):
|
|
chunk_offsets = i * 16 + tl.arange(0, 16)
|
|
bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets)
|
|
tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
|
|
|
|
|
|
@triton.jit
|
|
def _quantize_and_insert_k_mxfp8_kernel(
|
|
k_ptr, # [num_tokens, 512] bf16
|
|
slot_mapping_ptr, # [num_tokens] int64
|
|
k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view)
|
|
num_tokens,
|
|
head_dim: tl.constexpr, # 512, RoPE dims included
|
|
scale_dim: tl.constexpr, # 16
|
|
quant_block: tl.constexpr, # 32 dims per UE8M0 scale
|
|
cache_block_size: tl.constexpr, # paged cache block size, in tokens
|
|
block_stride: tl.constexpr, # total bytes per block (padded)
|
|
fp8_max: tl.constexpr,
|
|
use_fnuz: tl.constexpr = False,
|
|
):
|
|
"""Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache.
|
|
|
|
One program per token. Unlike the V4 record there is no bf16 tail, so every
|
|
dim is scaled.
|
|
"""
|
|
pid = tl.program_id(0)
|
|
if pid >= num_tokens:
|
|
return
|
|
|
|
slot_idx = tl.load(slot_mapping_ptr + pid)
|
|
if slot_idx == -1:
|
|
return
|
|
|
|
block_idx = slot_idx // cache_block_size
|
|
pos_in_block = slot_idx % cache_block_size
|
|
|
|
# int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks.
|
|
cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
|
|
token_data_ptr = cache_block_ptr + pos_in_block * head_dim
|
|
token_scale_ptr = (
|
|
cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
|
|
)
|
|
|
|
d = tl.arange(0, head_dim)
|
|
x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32)
|
|
|
|
tiles = tl.reshape(x, (scale_dim, quant_block))
|
|
# Floor the absmax like the CUDA encoder so an all-zero tile still lands on
|
|
# a finite scale.
|
|
amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4)
|
|
exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max)))
|
|
scaled = tl.clamp(
|
|
tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max
|
|
)
|
|
if use_fnuz: # noqa: SIM108
|
|
fp8 = scaled.to(tl.float8e4b8)
|
|
else:
|
|
fp8 = scaled.to(tl.float8e4nv)
|
|
tl.store(
|
|
token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,))
|
|
)
|
|
|
|
# UE8M0 encoding: stored_value = exponent + 127 (bias).
|
|
encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0)
|
|
tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8))
|
|
|
|
|
|
def quantize_and_insert_k_cache(
|
|
k: torch.Tensor, # [num_tokens, 512] bf16
|
|
k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8
|
|
slot_mapping: torch.Tensor, # [num_tokens] int64
|
|
block_size: int = 64,
|
|
is_ue8m0: bool = True,
|
|
use_fnuz: bool = False,
|
|
bytes_per_token: int = V4_BYTES_PER_TOKEN,
|
|
):
|
|
"""Quantize K tensor and insert into paged K cache.
|
|
|
|
``bytes_per_token`` picks the record (see the module header): the V4 one,
|
|
or V4.1's all-dims MXFP8 one.
|
|
|
|
``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on
|
|
platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,
|
|
which is used by OCP-encoded caches even on gfx942.
|
|
"""
|
|
assert k.dim() == 2 and k.shape[1] == 512, (
|
|
f"K must be [num_tokens, 512], got {k.shape}"
|
|
)
|
|
assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}"
|
|
assert is_ue8m0, "Only support ue8m0 quantization."
|
|
|
|
# NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to
|
|
# padding. Always use slot_mapping.shape[0] as the token count.
|
|
num_tokens = slot_mapping.shape[0]
|
|
block_stride = k_cache.stride(0) # bytes per block
|
|
|
|
if bytes_per_token == V41_BYTES_PER_TOKEN:
|
|
if use_fnuz and not current_platform.is_fp8_fnuz():
|
|
raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
|
|
_quantize_and_insert_k_mxfp8_kernel[(num_tokens,)](
|
|
k,
|
|
slot_mapping,
|
|
k_cache,
|
|
num_tokens,
|
|
head_dim=512,
|
|
scale_dim=V41_NUM_SCALES,
|
|
quant_block=V41_QUANT_BLOCK,
|
|
cache_block_size=block_size,
|
|
block_stride=block_stride,
|
|
fp8_max=get_fp8_min_max()[1]
|
|
if use_fnuz
|
|
else torch.finfo(torch.float8_e4m3fn).max,
|
|
use_fnuz=use_fnuz,
|
|
)
|
|
return
|
|
|
|
TOKEN_FP8_DIM = 448
|
|
TOKEN_BF16_DIM = 64
|
|
TOKEN_SCALE_DIM = 8
|
|
QUANT_BLOCK_SIZE = 64
|
|
if use_fnuz:
|
|
if not current_platform.is_fp8_fnuz():
|
|
raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
|
|
_, FP8_MAX = get_fp8_min_max()
|
|
else:
|
|
FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
|
|
TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
|
|
|
|
grid = (num_tokens,)
|
|
|
|
quantize_and_insert_k_kernel[grid](
|
|
k,
|
|
slot_mapping,
|
|
k_cache,
|
|
num_tokens,
|
|
input_dim=512,
|
|
fp8_dim=TOKEN_FP8_DIM,
|
|
bf16_dim=TOKEN_BF16_DIM,
|
|
scale_dim=TOKEN_SCALE_DIM,
|
|
quant_block=QUANT_BLOCK_SIZE,
|
|
cache_block_size=block_size,
|
|
token_data_size=TOKEN_DATA_SIZE,
|
|
block_stride=block_stride,
|
|
fp8_max=FP8_MAX,
|
|
n_quant_blocks=8,
|
|
use_fnuz=use_fnuz,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _dequantize_and_gather_k_kernel(
|
|
out_ptr,
|
|
out_stride0,
|
|
out_stride1,
|
|
k_cache_ptr,
|
|
seq_lens_ptr,
|
|
block_table_ptr,
|
|
offset,
|
|
gather_lens_ptr,
|
|
# Constants
|
|
max_blocks_per_seq: tl.constexpr,
|
|
fp8_dim: tl.constexpr, # 448
|
|
bf16_dim: tl.constexpr, # 64
|
|
scale_dim: tl.constexpr, # 8
|
|
quant_block: tl.constexpr, # 64 (quantization block size)
|
|
cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size)
|
|
token_data_size: tl.constexpr, # 576 bytes per token data
|
|
block_stride: tl.constexpr, # total bytes per block (padded) int32
|
|
output_dim: tl.constexpr, # 512
|
|
fp8_max: tl.constexpr,
|
|
n_quant_blocks: tl.constexpr, # 7 real blocks
|
|
use_fnuz: tl.constexpr = False,
|
|
):
|
|
batch_idx = tl.program_id(0)
|
|
worker_id = tl.program_id(1)
|
|
num_workers = tl.num_programs(1)
|
|
|
|
seq_len = tl.load(seq_lens_ptr + batch_idx)
|
|
if gather_lens_ptr is not None: # noqa: SIM108
|
|
gather_len = tl.load(gather_lens_ptr + batch_idx)
|
|
else:
|
|
# Gather all tokens
|
|
gather_len = seq_len
|
|
start_pos = seq_len - gather_len
|
|
|
|
for i in range(worker_id, gather_len, num_workers):
|
|
# Calculate the actual token index in the sequence
|
|
pos = start_pos + i
|
|
|
|
# Calculate which block and position within block
|
|
block_in_seq = pos // cache_block_size
|
|
pos_in_block = pos % cache_block_size
|
|
|
|
# Get physical block index from block table
|
|
block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
|
|
physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
|
|
|
|
# int64: physical_block_idx * block_stride can exceed 2^31 with many
|
|
# KV-cache blocks (e.g. >= 57K at block_stride ~37K).
|
|
cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
|
|
|
|
# Token data pointer
|
|
token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
|
|
|
|
# Scale pointer: after all token data
|
|
token_scale_ptr = (
|
|
cache_block_ptr
|
|
+ cache_block_size * token_data_size
|
|
+ pos_in_block * scale_dim
|
|
)
|
|
|
|
# Token data layout: [0:448] fp8, [448:576] bf16
|
|
token_fp8_ptr = token_data_ptr
|
|
token_bf16_ptr = token_data_ptr + fp8_dim
|
|
|
|
# Output pointer for this token (flattened)
|
|
output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
|
|
|
|
# ========== Dequantize FP8 portion using UE8M0 ==========
|
|
for qblock_idx in tl.static_range(n_quant_blocks):
|
|
qblock_start = qblock_idx * quant_block
|
|
|
|
if qblock_start < fp8_dim:
|
|
offsets = qblock_start + tl.arange(0, quant_block)
|
|
mask = offsets < fp8_dim
|
|
|
|
# Load quantized fp8 values (stored as uint8)
|
|
x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)
|
|
|
|
# Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere).
|
|
if use_fnuz:
|
|
x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
|
|
else:
|
|
x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
|
|
|
|
# Convert fp8 to float32 for computation
|
|
x_float = x_fp8.to(tl.float32)
|
|
|
|
# Load and decode UE8M0 scale
|
|
# UE8M0: scale = 2^(stored_value - 127)
|
|
encoded_scale = tl.load(token_scale_ptr + qblock_idx)
|
|
exponent = encoded_scale.to(tl.float32) - 127.0
|
|
scale = tl.exp2(exponent)
|
|
|
|
# Dequantize: bf16_value = fp8_value * scale
|
|
x_dequant = x_float * scale
|
|
|
|
# Store as bf16
|
|
tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
|
|
|
|
# ========== Copy BF16 portion directly ==========
|
|
bf16_output_offset = fp8_dim # After 448 elements in output
|
|
|
|
# Read bf16 from cache
|
|
bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
|
|
|
|
# Process in chunks of 16
|
|
for j in tl.static_range(bf16_dim // 16):
|
|
chunk_offsets = j * 16 + tl.arange(0, 16)
|
|
bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets)
|
|
tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
|
|
|
|
|
|
@triton.jit
|
|
def _dequantize_and_gather_k_mxfp8_kernel(
|
|
out_ptr,
|
|
out_stride0,
|
|
out_stride1,
|
|
k_cache_ptr,
|
|
seq_lens_ptr,
|
|
block_table_ptr,
|
|
offset,
|
|
gather_lens_ptr,
|
|
max_blocks_per_seq: tl.constexpr,
|
|
head_dim: tl.constexpr, # 512
|
|
scale_dim: tl.constexpr, # 16
|
|
quant_block: tl.constexpr, # 32
|
|
cache_block_size: tl.constexpr,
|
|
block_stride: tl.constexpr,
|
|
use_fnuz: tl.constexpr = False,
|
|
):
|
|
"""Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place."""
|
|
batch_idx = tl.program_id(0)
|
|
worker_id = tl.program_id(1)
|
|
num_workers = tl.num_programs(1)
|
|
|
|
seq_len = tl.load(seq_lens_ptr + batch_idx)
|
|
if gather_lens_ptr is not None: # noqa: SIM108
|
|
gather_len = tl.load(gather_lens_ptr + batch_idx)
|
|
else:
|
|
gather_len = seq_len
|
|
start_pos = seq_len - gather_len
|
|
|
|
d = tl.arange(0, head_dim)
|
|
s = tl.arange(0, scale_dim)
|
|
for i in range(worker_id, gather_len, num_workers):
|
|
pos = start_pos + i
|
|
block_in_seq = pos // cache_block_size
|
|
pos_in_block = pos % cache_block_size
|
|
|
|
block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
|
|
physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
|
|
|
|
# int64: physical_block_idx * block_stride can exceed 2^31.
|
|
cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
|
|
token_data_ptr = cache_block_ptr + pos_in_block * head_dim
|
|
token_scale_ptr = (
|
|
cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
|
|
)
|
|
|
|
x_uint8 = tl.load(token_data_ptr + d)
|
|
if use_fnuz:
|
|
x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
|
|
else:
|
|
x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
|
|
tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block))
|
|
|
|
# UE8M0: scale = 2^(stored_value - 127).
|
|
encoded = tl.load(token_scale_ptr + s)
|
|
scale = tl.exp2(encoded.to(tl.float32) - 127.0)
|
|
dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
|
|
|
|
output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
|
|
tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
|
|
|
|
|
|
@triton.jit
|
|
def _dequantize_and_gather_k_nvfp4_kernel(
|
|
out_ptr,
|
|
out_stride0,
|
|
out_stride1,
|
|
k_cache_ptr,
|
|
seq_lens_ptr,
|
|
block_table_ptr,
|
|
offset,
|
|
gather_lens_ptr,
|
|
max_blocks_per_seq: tl.constexpr,
|
|
head_dim: tl.constexpr, # 512
|
|
scale_dim: tl.constexpr, # 32
|
|
quant_block: tl.constexpr, # 16
|
|
cache_block_size: tl.constexpr,
|
|
block_stride: tl.constexpr,
|
|
):
|
|
"""Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace."""
|
|
batch_idx = tl.program_id(0)
|
|
worker_id = tl.program_id(1)
|
|
num_workers = tl.num_programs(1)
|
|
|
|
seq_len = tl.load(seq_lens_ptr + batch_idx)
|
|
if gather_lens_ptr is not None: # noqa: SIM108
|
|
gather_len = tl.load(gather_lens_ptr + batch_idx)
|
|
else:
|
|
gather_len = seq_len
|
|
start_pos = seq_len - gather_len
|
|
|
|
packed_bytes: tl.constexpr = head_dim // 2
|
|
d = tl.arange(0, head_dim)
|
|
for i in range(worker_id, gather_len, num_workers):
|
|
pos = start_pos + i
|
|
block_in_seq = pos // cache_block_size
|
|
pos_in_block = pos % cache_block_size
|
|
|
|
physical_block_idx = tl.load(
|
|
block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq
|
|
)
|
|
page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
|
|
|
|
packed = tl.load(
|
|
page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes)
|
|
)
|
|
# Even element in the low nibble, odd in the high nibble.
|
|
codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32))
|
|
# e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6.
|
|
mag_code = codes & 7
|
|
e = (mag_code >> 1).to(tl.float32)
|
|
m = (mag_code & 1).to(tl.float32)
|
|
mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0))
|
|
vals = tl.where(codes >= 8, -mag, mag)
|
|
|
|
sf = tl.load(
|
|
page
|
|
+ cache_block_size * packed_bytes
|
|
+ pos_in_block * scale_dim
|
|
+ tl.arange(0, scale_dim)
|
|
)
|
|
scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32)
|
|
tiles = tl.reshape(vals, (scale_dim, quant_block))
|
|
dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
|
|
|
|
output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
|
|
tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
|
|
|
|
|
|
def dequantize_and_gather_k_cache_triton(
|
|
# [num_reqs, max_num_tokens, head_size]
|
|
out: torch.Tensor,
|
|
# [num_blocks, block_size, head_bytes]
|
|
k_cache: torch.Tensor,
|
|
# [num_reqs]
|
|
seq_lens: torch.Tensor,
|
|
# [num_reqs]
|
|
gather_lens: torch.Tensor | None,
|
|
# [num_reqs, max_blocks_per_seq]
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
offset: int,
|
|
use_fnuz: bool = False,
|
|
) -> None:
|
|
num_reqs = seq_lens.shape[0]
|
|
NUM_WORKERS = 128
|
|
if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN:
|
|
# Keep short gathers unchanged; cap total workers for multi-request chunks.
|
|
num_workers = min(
|
|
max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset)
|
|
)
|
|
_dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)](
|
|
out,
|
|
out.stride(0),
|
|
out.stride(1),
|
|
k_cache,
|
|
seq_lens,
|
|
block_table,
|
|
offset,
|
|
gather_lens,
|
|
max_blocks_per_seq=block_table.shape[-1],
|
|
head_dim=512,
|
|
scale_dim=V41_NVFP4_NUM_SCALES,
|
|
quant_block=V41_NVFP4_QUANT_BLOCK,
|
|
cache_block_size=block_size,
|
|
block_stride=k_cache.stride(0),
|
|
)
|
|
return
|
|
|
|
if k_cache.shape[-1] == V41_BYTES_PER_TOKEN:
|
|
_dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)](
|
|
out,
|
|
out.stride(0),
|
|
out.stride(1),
|
|
k_cache,
|
|
seq_lens,
|
|
block_table,
|
|
offset,
|
|
gather_lens,
|
|
max_blocks_per_seq=block_table.shape[-1],
|
|
head_dim=512,
|
|
scale_dim=V41_NUM_SCALES,
|
|
quant_block=V41_QUANT_BLOCK,
|
|
cache_block_size=block_size,
|
|
block_stride=k_cache.stride(0),
|
|
use_fnuz=use_fnuz,
|
|
)
|
|
return
|
|
|
|
TOKEN_FP8_DIM = 448
|
|
TOKEN_BF16_DIM = 64
|
|
TOKEN_SCALE_DIM = 8
|
|
QUANT_BLOCK_SIZE = 64
|
|
FP8_MAX = 448.0
|
|
TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
|
|
|
|
_dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
|
|
out,
|
|
out.stride(0),
|
|
out.stride(1),
|
|
k_cache,
|
|
seq_lens,
|
|
block_table,
|
|
offset,
|
|
gather_lens,
|
|
max_blocks_per_seq=block_table.shape[-1],
|
|
fp8_dim=TOKEN_FP8_DIM,
|
|
bf16_dim=TOKEN_BF16_DIM,
|
|
scale_dim=TOKEN_SCALE_DIM,
|
|
quant_block=QUANT_BLOCK_SIZE,
|
|
cache_block_size=block_size,
|
|
token_data_size=TOKEN_DATA_SIZE,
|
|
block_stride=k_cache.stride(0),
|
|
output_dim=512,
|
|
fp8_max=FP8_MAX,
|
|
n_quant_blocks=7,
|
|
use_fnuz=use_fnuz,
|
|
)
|
|
|
|
|
|
def dequantize_and_gather_k_cache(
|
|
# [num_reqs, max_num_tokens, head_size]
|
|
out: torch.Tensor,
|
|
# [num_blocks, block_size, head_bytes]
|
|
k_cache: torch.Tensor,
|
|
# [num_reqs]
|
|
seq_lens: torch.Tensor,
|
|
# [num_reqs]
|
|
gather_lens: torch.Tensor | None,
|
|
# [num_reqs, max_blocks_per_seq]
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
offset: int,
|
|
use_fnuz: bool = False,
|
|
) -> None:
|
|
"""Dequantize and gather a paged DSv4 K cache.
|
|
|
|
The record is read off ``k_cache.shape[-1]``; see the module header. Only
|
|
the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton
|
|
path.
|
|
|
|
``use_fnuz`` MUST match the encoder of the specific cache being read:
|
|
``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
|
|
``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
|
|
writes FNUZ on gfx942 and OCP on gfx950).
|
|
"""
|
|
if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN:
|
|
# lazily import, otherwise some tests fail due to CUDA driver init failure.
|
|
from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
|
|
_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
|
|
)
|
|
|
|
mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN
|
|
_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
|
|
out=out,
|
|
k_cache=k_cache,
|
|
seq_lens=seq_lens,
|
|
gather_lens=gather_lens,
|
|
block_table=block_table,
|
|
block_size=block_size,
|
|
offset=offset,
|
|
fp8_dim=512 if mxfp8 else 448,
|
|
group_size=V41_QUANT_BLOCK if mxfp8 else 64,
|
|
)
|
|
return
|
|
|
|
dequantize_and_gather_k_cache_triton(
|
|
out,
|
|
k_cache,
|
|
seq_lens,
|
|
gather_lens,
|
|
block_table,
|
|
block_size,
|
|
offset,
|
|
use_fnuz=use_fnuz,
|
|
)
|
|
|
|
|
|
def compute_global_topk_indices_and_lens(
|
|
topk_indices: torch.Tensor,
|
|
token_to_req_indices: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
is_valid_token: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Map local topk indices to global KV cache slots and count valid entries.
|
|
|
|
Fuses three operations into a single kernel:
|
|
1. Block-table lookup (local index → global slot id)
|
|
2. Valid-entry counting (topk_lens per token)
|
|
3. Masking padding tokens to length 0
|
|
"""
|
|
num_tokens = topk_indices.shape[0]
|
|
global_topk_indices = torch.empty_like(topk_indices)
|
|
topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
|
|
_compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
|
|
global_topk_indices,
|
|
global_topk_indices.stride(0),
|
|
topk_lens,
|
|
topk_indices,
|
|
topk_indices.stride(0),
|
|
topk_indices.shape[-1],
|
|
token_to_req_indices,
|
|
block_table,
|
|
block_table.stride(0),
|
|
block_size,
|
|
is_valid_token,
|
|
TRITON_BLOCK_SIZE=1024,
|
|
)
|
|
return global_topk_indices, topk_lens
|
|
|
|
|
|
@triton.jit
|
|
def _compute_global_topk_indices_and_lens_kernel(
|
|
global_topk_indices_ptr,
|
|
global_topk_indices_stride: tl.constexpr,
|
|
topk_lens_ptr,
|
|
topk_indices_ptr,
|
|
topk_indices_stride: tl.constexpr,
|
|
topk: tl.constexpr,
|
|
token_to_req_indices_ptr,
|
|
block_table_ptr,
|
|
block_table_stride: tl.constexpr,
|
|
block_size: tl.constexpr,
|
|
is_valid_token_ptr,
|
|
TRITON_BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
token_idx = tl.program_id(0)
|
|
is_valid_token = tl.load(is_valid_token_ptr + token_idx)
|
|
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
|
|
|
|
count = tl.zeros((), dtype=tl.int32)
|
|
for i in range(0, topk, TRITON_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
|
|
mask = offset < topk
|
|
|
|
local_idx = tl.load(
|
|
topk_indices_ptr + token_idx * topk_indices_stride + offset,
|
|
mask=mask,
|
|
other=-1,
|
|
)
|
|
is_valid = local_idx >= 0
|
|
|
|
block_indices = local_idx // block_size
|
|
block_numbers = tl.load(
|
|
block_table_ptr + req_idx * block_table_stride + block_indices,
|
|
mask=mask & is_valid,
|
|
)
|
|
block_offsets = local_idx % block_size
|
|
|
|
slot_ids = block_numbers * block_size + block_offsets
|
|
slot_ids = tl.where(is_valid, slot_ids, -1)
|
|
tl.store(
|
|
global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset,
|
|
slot_ids,
|
|
mask=mask,
|
|
)
|
|
count += tl.sum(is_valid.to(tl.int32), axis=0)
|
|
|
|
# Zero out length for padding tokens.
|
|
tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
|
|
|
|
|
|
# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see
|
|
# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is
|
|
# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both.
|
|
# The extra slots stay as -1 sentinels and `combined_lens` caps the valid
|
|
# range via `topk_length`, so padding is a no-op at kernel level.
|
|
_SPARSE_PREFILL_TOPK_ALIGNMENT = 128
|
|
|
|
|
|
def combine_topk_swa_indices(
|
|
topk_indices: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
gather_lens: torch.Tensor,
|
|
window_size: int,
|
|
compress_ratio: int,
|
|
topk: int,
|
|
M: int,
|
|
N: int,
|
|
out: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
num_tokens = topk_indices.shape[0]
|
|
# The width is fixed per model so the caller-provided workspace stays
|
|
# valid on every batch.
|
|
combined_topk = (
|
|
(topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
|
|
// _SPARSE_PREFILL_TOPK_ALIGNMENT
|
|
* _SPARSE_PREFILL_TOPK_ALIGNMENT
|
|
)
|
|
if out is None:
|
|
combined_indices = torch.full(
|
|
(num_tokens, combined_topk),
|
|
fill_value=-1,
|
|
dtype=torch.int32,
|
|
device=topk_indices.device,
|
|
)
|
|
combined_lens = torch.empty(
|
|
num_tokens, dtype=torch.int32, device=topk_indices.device
|
|
)
|
|
else:
|
|
combined_indices, combined_lens = out
|
|
|
|
_COMBINE_TOPK_SWA_INDICES_KERNEL(
|
|
combined_indices,
|
|
combined_lens,
|
|
topk_indices,
|
|
query_start_loc,
|
|
seq_lens,
|
|
gather_lens,
|
|
M,
|
|
N,
|
|
TOP_K=topk,
|
|
COMPRESS_RATIO=compress_ratio,
|
|
WINDOW_SIZE=window_size,
|
|
)
|
|
return combined_indices, combined_lens
|
|
|
|
|
|
_COMBINE_TOPK_SWA_NUM_WORKERS = 256
|
|
|
|
|
|
# Representative pointer alignment variants for Triton pointer specialization.
|
|
_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs(
|
|
dict(
|
|
topk_indices=True,
|
|
query_start_loc=True,
|
|
seq_lens=True,
|
|
gather_lens=True,
|
|
),
|
|
dict(
|
|
topk_indices=True,
|
|
query_start_loc=False,
|
|
seq_lens=False,
|
|
gather_lens=True,
|
|
),
|
|
dict(
|
|
topk_indices=False,
|
|
query_start_loc=False,
|
|
seq_lens=False,
|
|
gather_lens=False,
|
|
),
|
|
)
|
|
|
|
|
|
def _hf_config_int(vllm_config: Any, name: str, default: int) -> int:
|
|
model_config = getattr(vllm_config, "model_config", None)
|
|
hf_config = getattr(model_config, "hf_config", None)
|
|
return int(getattr(hf_config, name, default) or default)
|
|
|
|
|
|
def _combine_topk_swa_warmup_inputs(vllm_config: Any) -> _WarmupInputRows:
|
|
"""One warmup row per v4.1 layer type (compress ratio) in the model.
|
|
|
|
Every layer slices the shared ``index_topk``-wide indices buffer, so the
|
|
padded width is the same for all rows; SWA-only layers (ratio 0) pass
|
|
``TOP_K=0``.
|
|
"""
|
|
model_config = getattr(vllm_config, "model_config", None)
|
|
hf_config = getattr(model_config, "hf_config", None)
|
|
compress_ratios = getattr(hf_config, "compress_ratios", None) or [0]
|
|
topk = _hf_config_int(vllm_config, "index_topk", 0)
|
|
return zip_inputs(
|
|
*(
|
|
dict(compress_ratio=ratio, topk=topk if ratio > 0 else 0, topk_width=topk)
|
|
for ratio in sorted({int(ratio) for ratio in compress_ratios})
|
|
)
|
|
)
|
|
|
|
|
|
def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int:
|
|
scheduler_config = getattr(vllm_config, "scheduler_config", None)
|
|
return int(getattr(scheduler_config, name, default) or default)
|
|
|
|
|
|
class CombineTopkSwaIndicesKernel(
|
|
VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
|
|
):
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
TOP_K: int
|
|
COMPRESS_RATIO: int
|
|
WINDOW_SIZE: int
|
|
PADDED_TOP_K: int
|
|
input_variant: TritonPointerInputVariant
|
|
|
|
@staticmethod
|
|
@triton.jit(
|
|
do_not_specialize=[
|
|
"combined_indices_stride",
|
|
"topk_indices_stride",
|
|
"M",
|
|
"N",
|
|
]
|
|
)
|
|
def kernel(
|
|
combined_indices_ptr,
|
|
combined_indices_stride,
|
|
combined_lens_ptr,
|
|
topk_indices_ptr,
|
|
topk_indices_stride,
|
|
query_start_loc_ptr,
|
|
seq_lens_ptr,
|
|
gather_lens_ptr,
|
|
M,
|
|
N,
|
|
TOP_K: tl.constexpr,
|
|
COMPRESS_RATIO: tl.constexpr,
|
|
WINDOW_SIZE: tl.constexpr,
|
|
PADDED_TOP_K: tl.constexpr,
|
|
):
|
|
batch_idx = tl.program_id(0)
|
|
worker_id = tl.program_id(1)
|
|
num_workers = tl.num_programs(1)
|
|
|
|
# query_start_loc is a global tensor; rebase to chunk-local offsets
|
|
# by subtracting the chunk's starting value.
|
|
base = tl.load(query_start_loc_ptr)
|
|
query_start = tl.load(query_start_loc_ptr + batch_idx) - base
|
|
query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
|
|
query_len = query_end - query_start
|
|
seq_len = tl.load(seq_lens_ptr + batch_idx)
|
|
gather_len = tl.load(gather_lens_ptr + batch_idx)
|
|
start_pos = seq_len - query_len
|
|
# The SWA portion of the gathered buffer starts from position
|
|
# (seq_len - gather_len), not position 0. We need this offset
|
|
# to correctly index into the gathered buffer.
|
|
gather_start = seq_len - gather_len
|
|
|
|
for token_idx in range(query_start + worker_id, query_end, num_workers):
|
|
# topk_len is fully determined by the query token's absolute position:
|
|
# the indexer emits min((pos + 1) // compress_ratio, topk_tokens)
|
|
# valid entries. Caller passes TOP_K=0 for SWA-only layers to zero
|
|
# this out.
|
|
token_idx_in_query = token_idx - query_start
|
|
pos = start_pos + token_idx_in_query
|
|
if COMPRESS_RATIO > 0:
|
|
topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
|
|
else:
|
|
# SWA-only layers pass TOP_K=0; skip the division entirely
|
|
# (integer div by constexpr 0 is UB and yields garbage, which
|
|
# turns the store index below negative).
|
|
topk_len = 0
|
|
swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
|
|
# gather_len already excludes context below the request's replay
|
|
# start (SWA bounded replay), so the window cannot start before
|
|
# the gathered buffer does.
|
|
swa_start = tl.maximum(swa_start, gather_start)
|
|
swa_len = pos - swa_start + 1
|
|
|
|
offset = tl.arange(0, PADDED_TOP_K)
|
|
mask = offset < topk_len
|
|
topk_indices = tl.load(
|
|
topk_indices_ptr + token_idx * topk_indices_stride + offset,
|
|
mask=mask,
|
|
)
|
|
tl.store(
|
|
combined_indices_ptr + token_idx * combined_indices_stride + offset,
|
|
topk_indices + M * batch_idx,
|
|
mask=mask,
|
|
)
|
|
# Index into gathered buffer: N + (position - gather_start)
|
|
# For positions [swa_start, pos], the buffer indices are:
|
|
# [N + swa_start - gather_start, N + pos - gather_start]
|
|
swa_offset = tl.arange(0, WINDOW_SIZE)
|
|
tl.store(
|
|
combined_indices_ptr
|
|
+ token_idx * combined_indices_stride
|
|
+ topk_len
|
|
+ swa_offset,
|
|
M * batch_idx + N + swa_start + swa_offset - gather_start,
|
|
mask=swa_offset < swa_len,
|
|
)
|
|
|
|
combined_len = topk_len + swa_len
|
|
tl.store(combined_lens_ptr + token_idx, combined_len)
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
topk_width: int,
|
|
topk_indices: bool,
|
|
query_start_loc: bool,
|
|
seq_lens: bool,
|
|
gather_lens: bool,
|
|
topk: int,
|
|
compress_ratio: int,
|
|
WINDOW_SIZE: int,
|
|
) -> CompileKey:
|
|
padded_topk = next_power_of_2(topk_width)
|
|
input_variant = TritonPointerInputVariant.from_alignment(
|
|
topk_indices=topk_indices,
|
|
query_start_loc=query_start_loc,
|
|
seq_lens=seq_lens,
|
|
gather_lens=gather_lens,
|
|
)
|
|
return self.CompileKey(
|
|
TOP_K=topk,
|
|
COMPRESS_RATIO=compress_ratio,
|
|
WINDOW_SIZE=WINDOW_SIZE,
|
|
PADDED_TOP_K=padded_topk,
|
|
input_variant=input_variant,
|
|
)
|
|
|
|
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
|
|
if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0:
|
|
return []
|
|
|
|
window_size = _hf_config_int(vllm_config, "sliding_window", 128)
|
|
return self._trace_dispatch(self.dispatch)(
|
|
_combine_topk_swa_warmup_inputs(vllm_config),
|
|
_COMBINE_TOPK_SWA_POINTER_INPUTS,
|
|
WINDOW_SIZE=window_size,
|
|
)
|
|
|
|
def compile(self, compile_key: CompileKey) -> None:
|
|
warmup = getattr(self.kernel, "warmup", None)
|
|
assert warmup is not None
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
input_variant = compile_key.input_variant
|
|
warmup(
|
|
int32_ptr,
|
|
1, # do not specialize combined_indices_stride
|
|
int32_ptr,
|
|
input_variant.pointer("topk_indices", torch.int32),
|
|
1, # do not specialize topk_indices_stride
|
|
input_variant.pointer("query_start_loc", torch.int32),
|
|
input_variant.pointer("seq_lens", torch.int32),
|
|
input_variant.pointer("gather_lens", torch.int32),
|
|
1, # do not specialize M
|
|
1, # do not specialize N
|
|
TOP_K=compile_key.TOP_K,
|
|
COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
|
|
WINDOW_SIZE=compile_key.WINDOW_SIZE,
|
|
PADDED_TOP_K=compile_key.PADDED_TOP_K,
|
|
grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
|
|
)
|
|
|
|
def __call__(
|
|
self,
|
|
combined_indices: torch.Tensor,
|
|
combined_lens: torch.Tensor,
|
|
topk_indices: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
gather_lens: torch.Tensor,
|
|
M: int,
|
|
N: int,
|
|
*,
|
|
TOP_K: int,
|
|
COMPRESS_RATIO: int,
|
|
WINDOW_SIZE: int,
|
|
) -> None:
|
|
num_reqs = seq_lens.shape[0]
|
|
self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
|
|
combined_indices,
|
|
combined_indices.stride(0),
|
|
combined_lens,
|
|
topk_indices,
|
|
topk_indices.stride(0),
|
|
query_start_loc,
|
|
seq_lens,
|
|
gather_lens,
|
|
M,
|
|
N,
|
|
TOP_K=TOP_K,
|
|
COMPRESS_RATIO=COMPRESS_RATIO,
|
|
WINDOW_SIZE=WINDOW_SIZE,
|
|
PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
|
|
)
|
|
|
|
|
|
_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel()
|
|
|
|
|
|
def build_flashinfer_mixed_sparse_indices(
|
|
decode_swa_indices: torch.Tensor,
|
|
decode_compressed_indices: torch.Tensor | None,
|
|
decode_compressed_topk_lens: torch.Tensor | None,
|
|
prefill_topk_indices: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
token_to_req_indices: torch.Tensor,
|
|
swa_block_table: torch.Tensor,
|
|
swa_block_size: int,
|
|
compressed_block_table: torch.Tensor | None,
|
|
compressed_block_size: int,
|
|
window_size: int,
|
|
compress_ratio: int,
|
|
topk: int,
|
|
decode_compressed_indices_are_local: bool = False,
|
|
decode_is_valid_token: torch.Tensor | None = None,
|
|
swa_block_span: int | None = None,
|
|
compressed_block_span: int | None = None,
|
|
*,
|
|
replay_start: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
|
|
|
|
Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width +
|
|
padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the
|
|
rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length
|
|
per token). Decode tokens read precomputed SWA/compressed indices; prefill
|
|
tokens derive their SWA window from the position and translate local
|
|
compressed indices to global slots via the block tables.
|
|
|
|
``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every
|
|
prefill token's window: positions below it hold no window KV.
|
|
"""
|
|
assert decode_swa_indices.dtype == torch.int32
|
|
assert decode_swa_indices.dim() == 2
|
|
swa_total_width = decode_swa_indices.shape[-1]
|
|
assert swa_total_width >= window_size
|
|
if decode_compressed_topk_lens is not None:
|
|
assert decode_compressed_topk_lens.dtype == torch.int32
|
|
assert prefill_topk_indices.dtype == torch.int32
|
|
assert prefill_topk_indices.dim() == 2
|
|
assert query_start_loc.dtype == torch.int32
|
|
assert seq_lens.dtype == torch.int32
|
|
assert token_to_req_indices.dtype == torch.int32
|
|
assert swa_block_table.dtype == torch.int32
|
|
|
|
num_decode_tokens = decode_swa_indices.shape[0]
|
|
num_prefill_tokens = prefill_topk_indices.shape[0]
|
|
num_tokens = num_decode_tokens + num_prefill_tokens
|
|
assert token_to_req_indices.shape[0] >= num_tokens
|
|
if decode_compressed_topk_lens is not None:
|
|
assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens
|
|
|
|
decode_compressed_topk = 0
|
|
if decode_compressed_indices is None:
|
|
decode_compressed_indices = prefill_topk_indices
|
|
else:
|
|
assert decode_compressed_indices.dtype == torch.int32
|
|
assert decode_compressed_indices.dim() == 2
|
|
assert decode_compressed_indices.shape[0] == num_decode_tokens
|
|
decode_compressed_topk = decode_compressed_indices.shape[-1]
|
|
if decode_compressed_topk > 0 and decode_compressed_indices_are_local:
|
|
assert decode_is_valid_token is not None
|
|
assert decode_is_valid_token.dtype == torch.bool
|
|
assert decode_is_valid_token.shape[0] >= num_decode_tokens
|
|
else:
|
|
decode_is_valid_token = token_to_req_indices
|
|
|
|
if compressed_block_table is None:
|
|
compressed_block_table = swa_block_table
|
|
assert compressed_block_table.dtype == torch.int32
|
|
has_decode_compressed_lens = decode_compressed_topk_lens is not None
|
|
if decode_compressed_topk_lens is None:
|
|
decode_compressed_topk_lens = token_to_req_indices
|
|
|
|
# The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk
|
|
# index row to start on a 16-byte boundary: the kernel loads the compressed
|
|
# indices with 128-bit (16-byte) vectorized loads, so a misaligned row would
|
|
# fault or read across rows. 16 bytes = 4 int32 indices, so round the topk
|
|
# width (and hence the row stride, since the SWA columns are fixed-width) up
|
|
# to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded
|
|
# by ``sparse_topk_lens``, so padding never changes the attention result.
|
|
padded_topk = max(topk, decode_compressed_topk)
|
|
padded_topk = (padded_topk + 3) // 4 * 4
|
|
sparse_indices = torch.empty(
|
|
(num_tokens, swa_total_width + padded_topk),
|
|
dtype=torch.int32,
|
|
device=decode_swa_indices.device,
|
|
)
|
|
sparse_topk_lens = torch.empty(
|
|
num_tokens, dtype=torch.int32, device=decode_swa_indices.device
|
|
)
|
|
if num_tokens == 0:
|
|
return sparse_indices, sparse_topk_lens
|
|
|
|
window_block_size = triton.next_power_of_2(max(swa_total_width, 1))
|
|
topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
|
|
max_block_size = max(window_block_size, topk_block_size)
|
|
num_warps = 4 if max_block_size >= 256 else 1
|
|
|
|
# block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV.
|
|
swa_span = swa_block_size if swa_block_span is None else swa_block_span
|
|
compressed_span = (
|
|
compressed_block_size
|
|
if compressed_block_span is None
|
|
else compressed_block_span
|
|
)
|
|
_build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
|
|
sparse_indices,
|
|
sparse_indices.stride(0),
|
|
sparse_topk_lens,
|
|
decode_swa_indices,
|
|
decode_swa_indices.stride(0),
|
|
decode_compressed_indices,
|
|
decode_compressed_indices.stride(0),
|
|
decode_compressed_topk_lens,
|
|
decode_is_valid_token,
|
|
prefill_topk_indices,
|
|
prefill_topk_indices.stride(0),
|
|
query_start_loc,
|
|
seq_lens,
|
|
token_to_req_indices,
|
|
replay_start,
|
|
swa_block_table,
|
|
swa_block_table.stride(0),
|
|
swa_block_size,
|
|
swa_span,
|
|
compressed_block_table,
|
|
compressed_block_table.stride(0),
|
|
compressed_block_size,
|
|
compressed_span,
|
|
NUM_DECODE_TOKENS=num_decode_tokens,
|
|
WINDOW_SIZE=window_size,
|
|
SWA_INDEX_WIDTH=swa_total_width,
|
|
COMPRESS_RATIO=compress_ratio,
|
|
TOP_K=topk,
|
|
PADDED_TOP_K=padded_topk,
|
|
PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1],
|
|
DECODE_COMPRESSED_TOPK=decode_compressed_topk,
|
|
DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
|
|
HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
|
|
WINDOW_BLOCK_SIZE=window_block_size,
|
|
TOPK_BLOCK_SIZE=topk_block_size,
|
|
num_warps=num_warps,
|
|
)
|
|
return sparse_indices, sparse_topk_lens
|
|
|
|
|
|
@triton.jit
|
|
def _remap_flashinfer_index(values, block_size, block_span):
|
|
# FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
|
|
# packed pages (#44577) need block*block_size+off -> block*block_span+off.
|
|
# TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
|
|
is_valid = values >= 0
|
|
safe_values = tl.where(is_valid, values, 0)
|
|
values = (safe_values // block_size) * block_span
|
|
values += safe_values % block_size
|
|
return tl.where(is_valid, values, -1)
|
|
|
|
|
|
@triton.jit(
|
|
do_not_specialize=[
|
|
"sparse_indices_stride",
|
|
"decode_swa_stride",
|
|
"decode_compressed_stride",
|
|
"prefill_topk_stride",
|
|
"swa_block_table_stride",
|
|
"swa_block_size",
|
|
"swa_block_span",
|
|
"compressed_block_table_stride",
|
|
"compressed_block_size",
|
|
"compressed_block_span",
|
|
"NUM_DECODE_TOKENS",
|
|
"PREFILL_TOPK_STRIDE",
|
|
]
|
|
)
|
|
def _build_flashinfer_mixed_sparse_indices_kernel(
|
|
sparse_indices_ptr,
|
|
sparse_indices_stride,
|
|
sparse_topk_lens_ptr,
|
|
decode_swa_indices_ptr,
|
|
decode_swa_stride,
|
|
decode_compressed_indices_ptr,
|
|
decode_compressed_stride,
|
|
decode_compressed_topk_lens_ptr,
|
|
decode_is_valid_token_ptr,
|
|
prefill_topk_indices_ptr,
|
|
prefill_topk_stride,
|
|
query_start_loc_ptr,
|
|
seq_lens_ptr,
|
|
token_to_req_indices_ptr,
|
|
replay_start_ptr,
|
|
swa_block_table_ptr,
|
|
swa_block_table_stride,
|
|
swa_block_size,
|
|
swa_block_span,
|
|
compressed_block_table_ptr,
|
|
compressed_block_table_stride,
|
|
compressed_block_size,
|
|
compressed_block_span,
|
|
NUM_DECODE_TOKENS,
|
|
WINDOW_SIZE: tl.constexpr,
|
|
SWA_INDEX_WIDTH: tl.constexpr,
|
|
COMPRESS_RATIO: tl.constexpr,
|
|
TOP_K: tl.constexpr,
|
|
PADDED_TOP_K: tl.constexpr,
|
|
PREFILL_TOPK_STRIDE,
|
|
DECODE_COMPRESSED_TOPK: tl.constexpr,
|
|
DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr,
|
|
HAS_DECODE_COMPRESSED_LENS: tl.constexpr,
|
|
WINDOW_BLOCK_SIZE: tl.constexpr,
|
|
TOPK_BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
token_idx = tl.program_id(0)
|
|
|
|
if token_idx < NUM_DECODE_TOKENS:
|
|
for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
|
|
mask = offset < SWA_INDEX_WIDTH
|
|
values = tl.load(
|
|
decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
|
|
mask=mask,
|
|
other=-1,
|
|
)
|
|
values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
|
|
tl.store(
|
|
sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
|
|
values,
|
|
mask=mask,
|
|
)
|
|
|
|
compressed_len = tl.zeros((), dtype=tl.int32)
|
|
for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
|
|
mask = offset < PADDED_TOP_K
|
|
values = tl.load(
|
|
decode_compressed_indices_ptr
|
|
+ token_idx * decode_compressed_stride
|
|
+ offset,
|
|
mask=offset < DECODE_COMPRESSED_TOPK,
|
|
other=-1,
|
|
)
|
|
if DECODE_COMPRESSED_INDICES_ARE_LOCAL:
|
|
token_valid = tl.load(decode_is_valid_token_ptr + token_idx)
|
|
is_valid = values >= 0
|
|
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
|
|
block_indices = values // compressed_block_size
|
|
block_numbers = tl.load(
|
|
compressed_block_table_ptr
|
|
+ req_idx * compressed_block_table_stride
|
|
+ block_indices,
|
|
mask=mask & is_valid,
|
|
other=-1,
|
|
)
|
|
block_offsets = values % compressed_block_size
|
|
values = block_numbers * compressed_block_size + block_offsets
|
|
values = tl.where(is_valid, values, -1)
|
|
compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)
|
|
values = _remap_flashinfer_index(
|
|
values, compressed_block_size, compressed_block_span
|
|
)
|
|
tl.store(
|
|
sparse_indices_ptr
|
|
+ token_idx * sparse_indices_stride
|
|
+ SWA_INDEX_WIDTH
|
|
+ offset,
|
|
values,
|
|
mask=mask,
|
|
)
|
|
|
|
if DECODE_COMPRESSED_TOPK == 0:
|
|
compressed_len = tl.zeros((), dtype=tl.int32)
|
|
elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
|
|
if HAS_DECODE_COMPRESSED_LENS:
|
|
compressed_len = tl.load(decode_compressed_topk_lens_ptr + token_idx)
|
|
else:
|
|
compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
|
|
|
|
tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len)
|
|
return
|
|
|
|
prefill_idx = token_idx - NUM_DECODE_TOKENS
|
|
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
|
|
query_start = tl.load(query_start_loc_ptr + req_idx)
|
|
query_end = tl.load(query_start_loc_ptr + req_idx + 1)
|
|
query_len = query_end - query_start
|
|
seq_len = tl.load(seq_lens_ptr + req_idx)
|
|
start_pos = seq_len - query_len
|
|
token_idx_in_query = token_idx - query_start
|
|
pos = start_pos + token_idx_in_query
|
|
swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
|
|
# SWA bounded replay: no window KV exists below the request's replay start.
|
|
swa_start_pos = tl.maximum(swa_start_pos, tl.load(replay_start_ptr + req_idx))
|
|
swa_len = pos - swa_start_pos + 1
|
|
if COMPRESS_RATIO > 0:
|
|
topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
|
|
else:
|
|
# SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB.
|
|
topk_len = 0
|
|
|
|
for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
|
|
mask = offset < SWA_INDEX_WIDTH
|
|
pos_offset = swa_start_pos + offset
|
|
block_indices = pos_offset // swa_block_size
|
|
block_numbers = tl.load(
|
|
swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices,
|
|
mask=mask & (offset < swa_len),
|
|
other=-1,
|
|
)
|
|
block_offsets = pos_offset % swa_block_size
|
|
slot_ids = block_numbers * swa_block_size + block_offsets
|
|
slot_ids = tl.where(offset < swa_len, slot_ids, -1)
|
|
slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)
|
|
tl.store(
|
|
sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
|
|
slot_ids,
|
|
mask=mask,
|
|
)
|
|
|
|
for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
|
|
mask = offset < PADDED_TOP_K
|
|
local_idx = tl.load(
|
|
prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset,
|
|
mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len),
|
|
other=-1,
|
|
)
|
|
is_valid = local_idx >= 0
|
|
block_indices = local_idx // compressed_block_size
|
|
block_numbers = tl.load(
|
|
compressed_block_table_ptr
|
|
+ req_idx * compressed_block_table_stride
|
|
+ block_indices,
|
|
mask=mask & is_valid,
|
|
other=-1,
|
|
)
|
|
block_offsets = local_idx % compressed_block_size
|
|
slot_ids = block_numbers * compressed_block_size + block_offsets
|
|
slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
|
|
slot_ids = _remap_flashinfer_index(
|
|
slot_ids, compressed_block_size, compressed_block_span
|
|
)
|
|
tl.store(
|
|
sparse_indices_ptr
|
|
+ token_idx * sparse_indices_stride
|
|
+ SWA_INDEX_WIDTH
|
|
+ offset,
|
|
slot_ids,
|
|
mask=mask,
|
|
)
|
|
|
|
tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len)
|