Harry MellorandClaude Opus 5
8c1557a79c
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1686 lines
63 KiB
Python
1686 lines
63 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 (
|
|
WarmupIntRange,
|
|
kernel_launcher,
|
|
zip_inputs,
|
|
)
|
|
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
|
|
LaunchSpec,
|
|
TritonPointerInputVariant,
|
|
TritonWarmupTensor,
|
|
VllmTritonJitKernel,
|
|
triton_scalar_specialization_rep,
|
|
)
|
|
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
|
|
|
|
|
|
@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)
|
|
|
|
|
|
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,
|
|
):
|
|
"""Quantize K tensor and insert into paged K cache.
|
|
|
|
K Cache block layout (block_size=64 tokens):
|
|
- First 64 * 576 = 36864 bytes: Token data
|
|
- Each token: 448 bytes (fp8) + 128 bytes (bf16)
|
|
- Next 64 * 8 = 512 bytes: Scales
|
|
- Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
|
|
- Padded to multiple of 576
|
|
|
|
``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
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
class DequantizeAndGatherKCacheKernel(
|
|
VllmTritonJitKernel["DequantizeAndGatherKCacheKernel.CompileKey"]
|
|
):
|
|
NUM_WORKERS = 128
|
|
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
max_blocks_per_seq: int
|
|
cache_block_size: int
|
|
block_stride: int
|
|
use_fnuz: bool
|
|
has_gather_lens: bool
|
|
offset: int
|
|
|
|
@staticmethod
|
|
@triton.jit
|
|
def 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)
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
max_model_len: int,
|
|
block_table_block_size: int,
|
|
cache_block_size: int,
|
|
offset: int,
|
|
**compile_key_fields: bool,
|
|
) -> CompileKey:
|
|
token_stride = 576
|
|
scale_stride = 8
|
|
unpadded = cache_block_size * (token_stride + scale_stride)
|
|
block_stride = ((unpadded + token_stride - 1) // token_stride) * token_stride
|
|
return self.CompileKey(
|
|
**compile_key_fields,
|
|
max_blocks_per_seq=(max_model_len + block_table_block_size - 1)
|
|
// block_table_block_size,
|
|
cache_block_size=cache_block_size,
|
|
block_stride=block_stride,
|
|
offset=triton_scalar_specialization_rep(offset),
|
|
)
|
|
|
|
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
|
|
max_model_len = vllm_config.model_config.max_model_len
|
|
block_size = vllm_config.cache_config.block_size
|
|
if max_model_len <= 0 or block_size <= 0:
|
|
return []
|
|
|
|
compress_ratios = frozenset(
|
|
max(1, int(compress_ratio))
|
|
for compress_ratio in vllm_config.model_config.hf_config.compress_ratios
|
|
)
|
|
return self._trace_dispatch(self.dispatch)(
|
|
zip_inputs(
|
|
dict(
|
|
cache_block_size=block_size,
|
|
has_gather_lens=True,
|
|
use_fnuz=current_platform.is_fp8_fnuz(),
|
|
offset=1,
|
|
enabled=True,
|
|
),
|
|
dict(
|
|
cache_block_size=block_size,
|
|
has_gather_lens=True,
|
|
use_fnuz=current_platform.is_fp8_fnuz(),
|
|
offset=2,
|
|
enabled=True,
|
|
),
|
|
dict(
|
|
cache_block_size=block_size,
|
|
has_gather_lens=True,
|
|
use_fnuz=current_platform.is_fp8_fnuz(),
|
|
offset=16,
|
|
enabled=True,
|
|
),
|
|
dict(
|
|
cache_block_size=max(block_size // 4, 1),
|
|
has_gather_lens=False,
|
|
use_fnuz=False,
|
|
offset=16,
|
|
enabled=4 in compress_ratios,
|
|
),
|
|
dict(
|
|
cache_block_size=max(block_size // 128, 1),
|
|
has_gather_lens=False,
|
|
use_fnuz=False,
|
|
offset=16,
|
|
enabled=128 in compress_ratios,
|
|
),
|
|
),
|
|
max_model_len=max_model_len,
|
|
block_table_block_size=block_size,
|
|
_when=lambda *, enabled: enabled,
|
|
)
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
return dict(
|
|
out=TritonWarmupTensor(
|
|
torch.bfloat16,
|
|
shape=(1, 1, 512),
|
|
strides=(512, 512, 1),
|
|
),
|
|
k_cache=TritonWarmupTensor(
|
|
torch.uint8,
|
|
shape=(1, 1),
|
|
strides=(compile_key.block_stride, 1),
|
|
),
|
|
seq_lens=int32_ptr,
|
|
gather_lens=(int32_ptr if compile_key.has_gather_lens else None),
|
|
block_table=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.max_blocks_per_seq),
|
|
),
|
|
block_size=compile_key.cache_block_size,
|
|
offset=compile_key.offset,
|
|
use_fnuz=compile_key.use_fnuz,
|
|
)
|
|
|
|
@kernel_launcher
|
|
def __call__(
|
|
self,
|
|
out: torch.Tensor,
|
|
k_cache: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
gather_lens: torch.Tensor | None,
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
offset: int,
|
|
*,
|
|
use_fnuz: bool = False,
|
|
) -> LaunchSpec:
|
|
num_reqs = seq_lens.shape[0]
|
|
return (num_reqs, self.NUM_WORKERS), dict(
|
|
out_stride0=out.stride(0),
|
|
out_stride1=out.stride(1),
|
|
max_blocks_per_seq=block_table.shape[-1],
|
|
fp8_dim=448,
|
|
bf16_dim=64,
|
|
scale_dim=8,
|
|
quant_block=64,
|
|
cache_block_size=block_size,
|
|
token_data_size=576,
|
|
block_stride=k_cache.stride(0),
|
|
output_dim=512,
|
|
fp8_max=448.0,
|
|
n_quant_blocks=7,
|
|
)
|
|
|
|
|
|
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.
|
|
|
|
``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():
|
|
# 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,
|
|
)
|
|
|
|
_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,
|
|
)
|
|
return
|
|
|
|
_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL(
|
|
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(
|
|
global_topk_indices,
|
|
topk_lens,
|
|
topk_indices,
|
|
token_to_req_indices,
|
|
block_table,
|
|
block_size,
|
|
is_valid_token,
|
|
)
|
|
return global_topk_indices, topk_lens
|
|
|
|
|
|
class ComputeGlobalTopkIndicesAndLensKernel(
|
|
VllmTritonJitKernel["ComputeGlobalTopkIndicesAndLensKernel.CompileKey"]
|
|
):
|
|
TRITON_BLOCK_SIZE = 1024
|
|
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
global_topk_indices_stride: int
|
|
topk_indices_stride: int
|
|
topk: int
|
|
block_table_stride: int
|
|
block_size: int
|
|
|
|
@staticmethod
|
|
@triton.jit
|
|
def 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))
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
topk_width: int,
|
|
block_size: int,
|
|
block_table_block_size: int,
|
|
max_model_len: int,
|
|
) -> CompileKey:
|
|
return self.CompileKey(
|
|
global_topk_indices_stride=topk_width,
|
|
topk_indices_stride=topk_width,
|
|
topk=topk_width,
|
|
block_table_stride=(max_model_len + block_table_block_size - 1)
|
|
// block_table_block_size,
|
|
block_size=block_size,
|
|
)
|
|
|
|
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
|
|
index_topk = vllm_config.model_config.hf_config.index_topk
|
|
compress_ratios = vllm_config.model_config.hf_config.compress_ratios
|
|
cache_block_size = vllm_config.cache_config.block_size
|
|
if index_topk <= 0 or cache_block_size <= 0:
|
|
return []
|
|
|
|
max_c128a_topk = (
|
|
((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
|
|
)
|
|
active_c128a_topk_widths = WarmupIntRange(
|
|
128,
|
|
max_c128a_topk + 1,
|
|
advance=lambda width: (
|
|
min(width * 2, max_c128a_topk)
|
|
if width < max_c128a_topk
|
|
else max_c128a_topk + 1
|
|
),
|
|
)
|
|
trace_dispatch = self._trace_dispatch(self.dispatch)
|
|
c4_keys = trace_dispatch(
|
|
topk_width=index_topk if 4 in compress_ratios else (),
|
|
block_size=max(1, cache_block_size // 4),
|
|
block_table_block_size=cache_block_size,
|
|
max_model_len=vllm_config.model_config.max_model_len,
|
|
)
|
|
c128a_keys = trace_dispatch(
|
|
topk_width=active_c128a_topk_widths if 128 in compress_ratios else (),
|
|
block_size=max(1, cache_block_size // 128),
|
|
block_table_block_size=cache_block_size,
|
|
max_model_len=vllm_config.model_config.max_model_len,
|
|
)
|
|
return list(dict.fromkeys((*c4_keys, *c128a_keys)))
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
return dict(
|
|
global_topk_indices=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.topk),
|
|
strides=(compile_key.global_topk_indices_stride, 1),
|
|
),
|
|
topk_lens=int32_ptr,
|
|
topk_indices=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.topk),
|
|
strides=(compile_key.topk_indices_stride, 1),
|
|
),
|
|
token_to_req_indices=int32_ptr,
|
|
block_table=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.block_table_stride),
|
|
),
|
|
block_size=compile_key.block_size,
|
|
is_valid_token=TritonWarmupTensor(torch.bool),
|
|
)
|
|
|
|
@kernel_launcher
|
|
def __call__(
|
|
self,
|
|
global_topk_indices: torch.Tensor,
|
|
topk_lens: torch.Tensor,
|
|
topk_indices: torch.Tensor,
|
|
token_to_req_indices: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
is_valid_token: torch.Tensor,
|
|
) -> LaunchSpec:
|
|
num_tokens = topk_indices.shape[0]
|
|
return (num_tokens,), dict(
|
|
global_topk_indices_stride=global_topk_indices.stride(0),
|
|
topk_indices_stride=topk_indices.stride(0),
|
|
topk=topk_indices.shape[-1],
|
|
block_table_stride=block_table.stride(0),
|
|
TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
|
|
)
|
|
|
|
|
|
_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL = DequantizeAndGatherKCacheKernel()
|
|
_COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL = ComputeGlobalTopkIndicesAndLensKernel()
|
|
|
|
|
|
# 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,
|
|
left_visible: torch.Tensor | None = None,
|
|
right_visible: torch.Tensor | None = None,
|
|
max_image_tokens: int = 0,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
num_tokens = topk_indices.shape[0]
|
|
# max_image_tokens widens the SWA column region for in-image
|
|
# bidirectional visibility; the width is fixed per model so the
|
|
# caller-provided workspace stays valid on image-free batches.
|
|
combined_topk = (
|
|
(topk + window_size + max_image_tokens + _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
|
|
combined_indices.fill_(-1)
|
|
|
|
_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,
|
|
left_visible=left_visible,
|
|
right_visible=right_visible,
|
|
max_image_tokens=max_image_tokens,
|
|
)
|
|
return combined_indices, combined_lens
|
|
|
|
|
|
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 _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(
|
|
VllmTritonJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
|
|
):
|
|
NUM_WORKERS = 256
|
|
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
TOP_K: int
|
|
COMPRESS_RATIO: int
|
|
WINDOW_SIZE: int
|
|
IMAGE_WIDTH: 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,
|
|
left_visible_ptr,
|
|
right_visible_ptr,
|
|
M,
|
|
N,
|
|
TOP_K: tl.constexpr,
|
|
COMPRESS_RATIO: tl.constexpr,
|
|
WINDOW_SIZE: tl.constexpr,
|
|
IMAGE_WIDTH: 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:
|
|
# both the C4A indexer and the C128A metadata builder emit
|
|
# 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
|
|
topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
|
|
if IMAGE_WIDTH > 0:
|
|
# In-image bidirectional visibility: the window starts up to
|
|
# max(left - (window - 1), 0) positions earlier and extends
|
|
# `right` positions past the query token (both 0 outside
|
|
# image spans, reducing to the plain causal window).
|
|
left = tl.load(left_visible_ptr + token_idx)
|
|
right = tl.load(right_visible_ptr + token_idx)
|
|
else:
|
|
left = 0
|
|
right = 0
|
|
left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
|
|
swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
|
|
swa_len = pos + right - 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 + right], the buffer indices are:
|
|
# [N + swa_start - gather_start, N + pos + right - gather_start]
|
|
for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE):
|
|
swa_offset = i + 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,
|
|
compress_ratio: int,
|
|
index_topk: int,
|
|
active_topk_width: int,
|
|
WINDOW_SIZE: int,
|
|
image_width: int = 0,
|
|
) -> CompileKey:
|
|
topk_width = active_topk_width if compress_ratio == 128 else index_topk
|
|
topk = 0 if compress_ratio == 1 else topk_width
|
|
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,
|
|
IMAGE_WIDTH=image_width,
|
|
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)
|
|
hf_config = vllm_config.model_config.hf_config
|
|
index_topk = int(getattr(hf_config, "index_topk", 0) or 0)
|
|
compress_ratios = tuple(
|
|
sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios})
|
|
)
|
|
max_c128a_topk = (
|
|
((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
|
|
)
|
|
active_topk_widths = WarmupIntRange(
|
|
128,
|
|
max_c128a_topk + 1,
|
|
advance=lambda width: (
|
|
min(width * 2, max_c128a_topk)
|
|
if width < max_c128a_topk
|
|
else max_c128a_topk + 1
|
|
),
|
|
)
|
|
image_width = (
|
|
_hf_config_int(vllm_config, "vision_max_n_token", 0)
|
|
if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0
|
|
else 0
|
|
)
|
|
# Warm both the plain-window variant (batches without image spans)
|
|
# and the in-image bidirectional variant.
|
|
image_widths = [0, image_width] if image_width > 0 else [0]
|
|
return self._trace_dispatch(self.dispatch)(
|
|
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,
|
|
),
|
|
),
|
|
compress_ratio=compress_ratios,
|
|
index_topk=index_topk,
|
|
active_topk_width=active_topk_widths,
|
|
WINDOW_SIZE=window_size,
|
|
image_width=image_widths,
|
|
)
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
variant = compile_key.input_variant
|
|
has_image = compile_key.IMAGE_WIDTH > 0
|
|
return dict(
|
|
combined_indices=int32_ptr,
|
|
combined_lens=int32_ptr,
|
|
topk_indices=variant.pointer(
|
|
"topk_indices", torch.int32, shape=(1, compile_key.PADDED_TOP_K)
|
|
),
|
|
query_start_loc=variant.pointer("query_start_loc", torch.int32),
|
|
seq_lens=variant.pointer("seq_lens", torch.int32),
|
|
gather_lens=variant.pointer("gather_lens", torch.int32),
|
|
M=1,
|
|
N=1,
|
|
TOP_K=compile_key.TOP_K,
|
|
COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
|
|
WINDOW_SIZE=compile_key.WINDOW_SIZE,
|
|
# Reproduce the runtime substitution: no-image batches launch with
|
|
# ``topk_indices`` in the visibility pointer slots (see __call__).
|
|
left_visible=(int32_ptr if has_image else None),
|
|
right_visible=(int32_ptr if has_image else None),
|
|
max_image_tokens=compile_key.IMAGE_WIDTH,
|
|
)
|
|
|
|
@kernel_launcher
|
|
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,
|
|
left_visible: torch.Tensor | None = None,
|
|
right_visible: torch.Tensor | None = None,
|
|
max_image_tokens: int = 0,
|
|
) -> LaunchSpec:
|
|
num_reqs = seq_lens.shape[0]
|
|
has_image = left_visible is not None
|
|
image_width = max_image_tokens if has_image else 0
|
|
return (num_reqs, self.NUM_WORKERS), dict(
|
|
combined_indices_stride=combined_indices.stride(0),
|
|
topk_indices_stride=topk_indices.stride(0),
|
|
left_visible_ptr=left_visible if has_image else topk_indices,
|
|
right_visible_ptr=right_visible if has_image else topk_indices,
|
|
TOP_K=TOP_K,
|
|
COMPRESS_RATIO=COMPRESS_RATIO,
|
|
WINDOW_SIZE=WINDOW_SIZE,
|
|
IMAGE_WIDTH=image_width,
|
|
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,
|
|
prefill_left_visible: torch.Tensor | None = None,
|
|
prefill_right_visible: torch.Tensor | None = None,
|
|
max_image_tokens: int = 0,
|
|
) -> 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.
|
|
|
|
When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
|
|
variant), the SWA column region widens by ``max_image_tokens`` and prefill
|
|
tokens inside an image span get a bidirectionally widened window; decode
|
|
rows are padded with -1 across the extra columns.
|
|
"""
|
|
assert decode_swa_indices.dtype == torch.int32
|
|
assert decode_swa_indices.dim() == 2
|
|
swa_index_width = decode_swa_indices.shape[-1]
|
|
assert swa_index_width >= window_size
|
|
has_image = prefill_left_visible is not None
|
|
image_width = max_image_tokens if has_image else 0
|
|
swa_total_width = swa_index_width + image_width
|
|
if has_image:
|
|
assert prefill_right_visible is not None
|
|
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_index_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(
|
|
sparse_indices,
|
|
sparse_topk_lens,
|
|
decode_swa_indices,
|
|
decode_compressed_indices,
|
|
decode_compressed_topk_lens,
|
|
decode_is_valid_token,
|
|
prefill_topk_indices,
|
|
prefill_left_visible if has_image else token_to_req_indices,
|
|
prefill_right_visible if has_image else token_to_req_indices,
|
|
query_start_loc,
|
|
seq_lens,
|
|
token_to_req_indices,
|
|
swa_block_table,
|
|
swa_block_size,
|
|
swa_span,
|
|
compressed_block_table,
|
|
compressed_block_size,
|
|
compressed_span,
|
|
num_decode_tokens=num_decode_tokens,
|
|
window_size=window_size,
|
|
swa_index_width=swa_index_width,
|
|
image_width=image_width,
|
|
compress_ratio=compress_ratio,
|
|
topk=topk,
|
|
padded_topk=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)
|
|
|
|
|
|
class BuildFlashinferMixedSparseIndicesKernel(
|
|
VllmTritonJitKernel["BuildFlashinferMixedSparseIndicesKernel.CompileKey"]
|
|
):
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
window_size: int
|
|
swa_index_width: int
|
|
image_width: int
|
|
compress_ratio: int
|
|
top_k: int
|
|
padded_top_k: int
|
|
decode_compressed_topk: int
|
|
decode_compressed_indices_are_local: bool
|
|
has_decode_compressed_lens: bool
|
|
window_block_size: int
|
|
topk_block_size: int
|
|
|
|
@staticmethod
|
|
@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 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,
|
|
left_visible_ptr,
|
|
right_visible_ptr,
|
|
query_start_loc_ptr,
|
|
seq_lens_ptr,
|
|
token_to_req_indices_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,
|
|
IMAGE_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)
|
|
# Total SWA column count; > SWA_INDEX_WIDTH only for in-image
|
|
# bidirectional visibility (vision variant), where the extra columns are
|
|
# -1-padded for decode rows.
|
|
SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
|
|
|
|
if token_idx < NUM_DECODE_TOKENS:
|
|
for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
|
|
mask = offset < SWA_TOTAL_WIDTH
|
|
values = tl.load(
|
|
decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
|
|
mask=offset < SWA_INDEX_WIDTH,
|
|
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_TOTAL_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_TOTAL_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
|
|
if IMAGE_WIDTH > 0:
|
|
# In-image bidirectional visibility: window starts up to
|
|
# max(left - (window - 1), 0) earlier and extends `right` past pos.
|
|
left = tl.load(left_visible_ptr + token_idx)
|
|
right = tl.load(right_visible_ptr + token_idx)
|
|
else:
|
|
left = 0
|
|
right = 0
|
|
left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
|
|
swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
|
|
swa_len = pos + right - swa_start_pos + 1
|
|
topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
|
|
|
|
for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
|
|
offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
|
|
mask = offset < SWA_TOTAL_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_TOTAL_WIDTH
|
|
+ offset,
|
|
slot_ids,
|
|
mask=mask,
|
|
)
|
|
|
|
tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
swa_index_width: int,
|
|
image_width: int,
|
|
compress_ratio: int,
|
|
index_topk: int,
|
|
active_topk_width: int,
|
|
has_prefill: bool,
|
|
**compile_key_fields: Any,
|
|
) -> CompileKey:
|
|
decode_compressed_topk = (
|
|
active_topk_width
|
|
if compress_ratio == 128
|
|
else index_topk
|
|
if compress_ratio == 4
|
|
else 0
|
|
)
|
|
topk_width = decode_compressed_topk if has_prefill else 0
|
|
padded_top_k = max(topk_width, decode_compressed_topk)
|
|
padded_top_k = (padded_top_k + 3) // 4 * 4
|
|
return self.CompileKey(
|
|
**compile_key_fields,
|
|
compress_ratio=compress_ratio,
|
|
top_k=topk_width,
|
|
decode_compressed_indices_are_local=compress_ratio == 4,
|
|
has_decode_compressed_lens=compress_ratio == 128,
|
|
swa_index_width=swa_index_width,
|
|
image_width=image_width,
|
|
padded_top_k=padded_top_k,
|
|
decode_compressed_topk=decode_compressed_topk,
|
|
window_block_size=next_power_of_2(max(swa_index_width, 1)),
|
|
topk_block_size=next_power_of_2(max(padded_top_k, 1)),
|
|
)
|
|
|
|
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
|
|
hf_config = vllm_config.model_config.hf_config
|
|
window_size = int(getattr(hf_config, "sliding_window", 128) or 128)
|
|
speculative_config = vllm_config.speculative_config
|
|
if speculative_config is not None and speculative_config.use_dspark():
|
|
from vllm.v1.attention.backends.mla.compressor_utils import (
|
|
get_dspark_swa_index_width,
|
|
)
|
|
|
|
swa_index_width: int | tuple[int, int] = (
|
|
window_size,
|
|
get_dspark_swa_index_width(
|
|
window_size,
|
|
speculative_config.num_speculative_tokens,
|
|
),
|
|
)
|
|
else:
|
|
swa_index_width = window_size
|
|
max_image_tokens = (
|
|
int(getattr(hf_config, "vision_max_n_token", 0) or 0)
|
|
if int(getattr(hf_config, "vision_n_layers", 0) or 0) > 0
|
|
else 0
|
|
)
|
|
image_widths = (0, max_image_tokens) if max_image_tokens > 0 else 0
|
|
index_topk = int(getattr(hf_config, "index_topk", 0) or 0)
|
|
compress_ratios = tuple(
|
|
sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios})
|
|
)
|
|
max_c128a_topk = (
|
|
((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
|
|
)
|
|
active_topk_widths = WarmupIntRange(
|
|
128,
|
|
max_c128a_topk + 1,
|
|
advance=lambda width: (
|
|
min(width * 2, max_c128a_topk)
|
|
if width < max_c128a_topk
|
|
else max_c128a_topk + 1
|
|
),
|
|
)
|
|
return self._trace_dispatch(self.dispatch)(
|
|
compress_ratio=compress_ratios,
|
|
index_topk=index_topk,
|
|
active_topk_width=active_topk_widths,
|
|
has_prefill=(False, True),
|
|
window_size=window_size,
|
|
swa_index_width=swa_index_width,
|
|
image_width=image_widths,
|
|
)
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
swa_block_size = 64
|
|
compressed_block_size = (
|
|
swa_block_size
|
|
if compile_key.compress_ratio == 1
|
|
else 256 // compile_key.compress_ratio
|
|
)
|
|
num_warps = (
|
|
4
|
|
if max(compile_key.window_block_size, compile_key.topk_block_size) >= 256
|
|
else 1
|
|
)
|
|
return dict(
|
|
sparse_indices=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(
|
|
1,
|
|
compile_key.swa_index_width
|
|
+ compile_key.image_width
|
|
+ compile_key.padded_top_k,
|
|
),
|
|
),
|
|
sparse_topk_lens=int32_ptr,
|
|
decode_swa_indices=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.swa_index_width),
|
|
),
|
|
decode_compressed_indices=int32_ptr,
|
|
decode_compressed_topk_lens=int32_ptr,
|
|
decode_is_valid_token=TritonWarmupTensor(
|
|
torch.bool
|
|
if compile_key.decode_compressed_indices_are_local
|
|
else torch.int32
|
|
),
|
|
prefill_topk_indices=int32_ptr,
|
|
prefill_left_visible=int32_ptr,
|
|
prefill_right_visible=int32_ptr,
|
|
query_start_loc=int32_ptr,
|
|
seq_lens=int32_ptr,
|
|
token_to_req_indices=int32_ptr,
|
|
swa_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)),
|
|
swa_block_size=swa_block_size,
|
|
swa_block_span=1,
|
|
compressed_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)),
|
|
compressed_block_size=compressed_block_size,
|
|
compressed_block_span=1,
|
|
num_decode_tokens=1,
|
|
window_size=compile_key.window_size,
|
|
swa_index_width=compile_key.swa_index_width,
|
|
image_width=compile_key.image_width,
|
|
compress_ratio=compile_key.compress_ratio,
|
|
topk=compile_key.top_k,
|
|
padded_topk=compile_key.padded_top_k,
|
|
prefill_topk_stride=1,
|
|
decode_compressed_topk=compile_key.decode_compressed_topk,
|
|
decode_compressed_indices_are_local=(
|
|
compile_key.decode_compressed_indices_are_local
|
|
),
|
|
has_decode_compressed_lens=compile_key.has_decode_compressed_lens,
|
|
window_block_size=compile_key.window_block_size,
|
|
topk_block_size=compile_key.topk_block_size,
|
|
num_warps=num_warps,
|
|
)
|
|
|
|
@kernel_launcher
|
|
def __call__(
|
|
self,
|
|
sparse_indices: torch.Tensor,
|
|
sparse_topk_lens: torch.Tensor,
|
|
decode_swa_indices: torch.Tensor,
|
|
decode_compressed_indices: torch.Tensor,
|
|
decode_compressed_topk_lens: torch.Tensor,
|
|
decode_is_valid_token: torch.Tensor,
|
|
prefill_topk_indices: torch.Tensor,
|
|
prefill_left_visible: torch.Tensor,
|
|
prefill_right_visible: 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,
|
|
swa_block_span: int,
|
|
compressed_block_table: torch.Tensor,
|
|
compressed_block_size: int,
|
|
compressed_block_span: int,
|
|
*,
|
|
num_decode_tokens: int,
|
|
window_size: int,
|
|
swa_index_width: int,
|
|
image_width: int,
|
|
compress_ratio: int,
|
|
topk: int,
|
|
padded_topk: int,
|
|
prefill_topk_stride: int,
|
|
decode_compressed_topk: int,
|
|
decode_compressed_indices_are_local: bool,
|
|
has_decode_compressed_lens: bool,
|
|
window_block_size: int,
|
|
topk_block_size: int,
|
|
num_warps: int,
|
|
) -> LaunchSpec:
|
|
return (sparse_indices.shape[0],), dict(
|
|
sparse_indices_stride=sparse_indices.stride(0),
|
|
decode_swa_stride=decode_swa_indices.stride(0),
|
|
decode_compressed_stride=decode_compressed_indices.stride(0),
|
|
prefill_topk_stride=prefill_topk_indices.stride(0),
|
|
left_visible_ptr=prefill_left_visible,
|
|
right_visible_ptr=prefill_right_visible,
|
|
swa_block_table_stride=swa_block_table.stride(0),
|
|
compressed_block_table_stride=compressed_block_table.stride(0),
|
|
NUM_DECODE_TOKENS=num_decode_tokens,
|
|
WINDOW_SIZE=window_size,
|
|
SWA_INDEX_WIDTH=swa_index_width,
|
|
IMAGE_WIDTH=image_width,
|
|
COMPRESS_RATIO=compress_ratio,
|
|
TOP_K=topk,
|
|
PADDED_TOP_K=padded_topk,
|
|
PREFILL_TOPK_STRIDE=prefill_topk_stride,
|
|
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,
|
|
)
|
|
|
|
|
|
_BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL = (
|
|
BuildFlashinferMixedSparseIndicesKernel()
|
|
)
|