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>
688 lines
27 KiB
Python
688 lines
27 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""Utility functions for sparse MLA backends."""
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
from vllm.config import VllmConfig
|
||
from vllm.distributed import get_dcp_group
|
||
from vllm.model_executor.warmup.jit_warmup import (
|
||
kernel_launcher,
|
||
zip_inputs,
|
||
)
|
||
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
|
||
LaunchSpec,
|
||
TritonWarmupTensor,
|
||
VllmTritonJitKernel,
|
||
triton_scalar_specialization_rep,
|
||
)
|
||
from vllm.triton_utils import tl, triton
|
||
from vllm.utils.math_utils import cdiv
|
||
from vllm.v1.worker.block_table import get_block_table_width
|
||
|
||
|
||
def request_row_bounds(req_idx: np.ndarray) -> np.ndarray:
|
||
"""Bounds of the runs of adjacent rows that belong to one request: run
|
||
``r`` is rows ``[bounds[r], bounds[r + 1])``.
|
||
|
||
Under PCP a rank holds two adjacent chunk rows of a split prefill; the
|
||
sparse backends give such a run one KV region.
|
||
"""
|
||
assert req_idx.size > 0
|
||
bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True])
|
||
assert bounds.size - 1 == np.unique(req_idx).size, (
|
||
"rows of one request must be adjacent"
|
||
)
|
||
return bounds
|
||
|
||
|
||
def flat_kv_row_view(
|
||
kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim]
|
||
block_size: int,
|
||
) -> tuple[torch.Tensor, int]:
|
||
"""Flat [row, head_dim] view of a paged cache and its physical rows per block.
|
||
|
||
Token offset is ``block_idx * block_stride_rows + offset_in_block``.
|
||
When other layers' pages sit between consecutive blocks of this cache,
|
||
``block_stride_rows`` exceeds ``block_size``; those in-between rows are never
|
||
indexed (`triton_convert_req_index_to_global_index` ensures this).
|
||
"""
|
||
num_blocks, _, head_dim = kv_cache.shape
|
||
assert kv_cache.stride(0) % head_dim == 0, (
|
||
"block stride is not a whole number of rows; flat row indexing would "
|
||
"silently misaddress"
|
||
)
|
||
block_stride_rows = kv_cache.stride(0) // head_dim
|
||
num_rows = (num_blocks - 1) * block_stride_rows + block_size
|
||
rows = kv_cache.as_strided((num_rows, head_dim), (head_dim, 1))
|
||
return rows, block_stride_rows
|
||
|
||
|
||
# Kernel with prefill workspace support and valid count tracking
|
||
|
||
|
||
class ConvertReqIndexToGlobalIndexKernel(
|
||
VllmTritonJitKernel["ConvertReqIndexToGlobalIndexKernel.CompileKey"]
|
||
):
|
||
@dataclass(frozen=True)
|
||
class CompileKey:
|
||
block_size: int
|
||
block_stride_rows: int
|
||
block_n: int
|
||
has_prefill_workspace: bool
|
||
count_valid: bool
|
||
single_tile: bool
|
||
compact_to_front: bool
|
||
dcp_size: int
|
||
dcp_rank: int
|
||
dcp_interleave: int
|
||
num_warps: int
|
||
num_topk_tokens: int
|
||
block_table_stride: int
|
||
|
||
@staticmethod
|
||
@triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"])
|
||
def kernel(
|
||
req_id_ptr, # int32 [num_tokens]
|
||
block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
|
||
token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
|
||
out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
|
||
valid_count_ptr, # int32 [num_tokens] - output valid count per row
|
||
prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill
|
||
workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr
|
||
workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace
|
||
# shapes (compile-time where possible)
|
||
max_num_blocks_per_req,
|
||
BLOCK_SIZE: tl.constexpr,
|
||
BLOCK_STRIDE_ROWS: tl.constexpr,
|
||
BLOCK_N: tl.constexpr, # tile width along columns
|
||
HAS_PREFILL: tl.constexpr,
|
||
COUNT_VALID: tl.constexpr, # whether to count valid indices
|
||
# BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
|
||
# is an in-register reduction and needs no atomic.
|
||
SINGLE_TILE: tl.constexpr,
|
||
# When set, scatter valid slots to a contiguous prefix [0, valid_count) using
|
||
# valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
|
||
# -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
|
||
# Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
|
||
# prefix is unspecified (only the selected set matters).
|
||
COMPACT_TO_FRONT: tl.constexpr,
|
||
# DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
|
||
DCP_SIZE: tl.constexpr,
|
||
DCP_RANK: tl.constexpr,
|
||
DCP_INTERLEAVE: tl.constexpr,
|
||
# strides (in elements)
|
||
bt_stride0,
|
||
bt_stride1,
|
||
ti_stride0,
|
||
ti_stride1,
|
||
out_stride0,
|
||
out_stride1,
|
||
):
|
||
# program_id(0) -> token_id (row)
|
||
# program_id(1) -> tile index along columns
|
||
token_id = tl.program_id(0)
|
||
tile_id = tl.program_id(1)
|
||
|
||
# Each program covers BLOCK_N consecutive columns
|
||
indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
|
||
|
||
# Load request id for this token (no mask: grid is exact)
|
||
req = tl.load(req_id_ptr + token_id)
|
||
|
||
# Load token indices for this tile
|
||
ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
|
||
tok = tl.load(ti_ptr) # int32
|
||
|
||
# Only token == -1 should propagate as -1
|
||
is_invalid_tok = tok < 0
|
||
is_prefill = False
|
||
if HAS_PREFILL:
|
||
prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
|
||
is_prefill = prefill_req_id >= 0
|
||
|
||
# DCP de-interleave the global token id into this rank's local slot.
|
||
# Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
|
||
# DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
|
||
# remote) and local_idx == tok, so this reduces to the non-DCP path; with
|
||
# DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
|
||
owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
|
||
is_remote = owning_rank != DCP_RANK
|
||
local_idx = (
|
||
tok // (DCP_SIZE * DCP_INTERLEAVE)
|
||
) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
|
||
|
||
# Compute block id and in-block offset
|
||
block_id = local_idx // BLOCK_SIZE
|
||
inblock_off = local_idx % BLOCK_SIZE
|
||
|
||
# Guard block_table access
|
||
valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
|
||
bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
|
||
is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill)
|
||
base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
|
||
out_val = base * BLOCK_STRIDE_ROWS + inblock_off
|
||
|
||
# Override with prefill output if prefill is enabled
|
||
if HAS_PREFILL:
|
||
workspace_start = tl.load(
|
||
workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
|
||
)
|
||
# Under DCP the prefill workspace is the all-gather of every rank's
|
||
# shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok.
|
||
prefill_out = (
|
||
owning_rank * workspace_rank_stride + workspace_start + local_idx
|
||
)
|
||
out_val = tl.where(is_prefill, prefill_out, out_val)
|
||
out_val = tl.where(is_invalid_tok, -1, out_val)
|
||
|
||
if COMPACT_TO_FRONT:
|
||
# Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
|
||
# sum gives each valid lane a distinct local offset; one atomic add of the
|
||
# tile's valid count reserves a contiguous base across racing tiles. The
|
||
# out buffer is pre-filled with -1, so unwritten tail slots stay -1.
|
||
# With no racing tiles the base is 0 and the allocator becomes a store.
|
||
is_valid = (~is_invalid_tok).to(tl.int32)
|
||
local_offset = tl.cumsum(is_valid) - is_valid
|
||
tile_valid_count = tl.sum(is_valid)
|
||
if SINGLE_TILE:
|
||
base = 0
|
||
tl.store(valid_count_ptr + token_id, tile_valid_count)
|
||
else:
|
||
base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
|
||
dest = base + local_offset
|
||
out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
|
||
tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
|
||
else:
|
||
# Store results in place (input column == output column).
|
||
out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
|
||
tl.store(out_ptr_ij, out_val)
|
||
|
||
# Accumulate the tile's valid count into the row total; a single tile's
|
||
# reduction *is* the total.
|
||
if COUNT_VALID:
|
||
tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
|
||
if SINGLE_TILE:
|
||
tl.store(valid_count_ptr + token_id, tile_valid_count)
|
||
else:
|
||
tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
|
||
|
||
def dispatch( # type: ignore[override]
|
||
self,
|
||
*,
|
||
BLOCK_SIZE: int,
|
||
BLOCK_STRIDE_ROWS: int,
|
||
BLOCK_N: int,
|
||
NUM_TOPK_TOKENS: int,
|
||
HAS_PREFILL_WORKSPACE: bool,
|
||
COUNT_VALID: bool,
|
||
COMPACT_TO_FRONT: bool,
|
||
DCP_SIZE: int,
|
||
DCP_RANK: int,
|
||
DCP_INTERLEAVE: int,
|
||
max_num_blocks_per_req: int,
|
||
) -> CompileKey:
|
||
tiling = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID)
|
||
return self.CompileKey(
|
||
block_size=BLOCK_SIZE,
|
||
block_stride_rows=BLOCK_STRIDE_ROWS,
|
||
block_n=tiling[1],
|
||
has_prefill_workspace=HAS_PREFILL_WORKSPACE,
|
||
count_valid=COUNT_VALID,
|
||
single_tile=tiling[0],
|
||
compact_to_front=COMPACT_TO_FRONT,
|
||
dcp_size=DCP_SIZE,
|
||
dcp_rank=DCP_RANK,
|
||
dcp_interleave=DCP_INTERLEAVE,
|
||
num_warps=tiling[3],
|
||
num_topk_tokens=NUM_TOPK_TOKENS,
|
||
block_table_stride=triton_scalar_specialization_rep(max_num_blocks_per_req),
|
||
)
|
||
|
||
def get_warmup_keys(
|
||
self,
|
||
vllm_config: VllmConfig,
|
||
*,
|
||
block_stride_rows: int,
|
||
) -> list[CompileKey]:
|
||
block_size = vllm_config.cache_config.block_size
|
||
dcp_size = vllm_config.parallel_config.decode_context_parallel_size
|
||
dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
|
||
dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
|
||
num_topk_tokens = vllm_config.model_config.hf_text_config.index_topk
|
||
max_num_blocks = cdiv(
|
||
vllm_config.model_config.max_model_len,
|
||
block_size * dcp_size,
|
||
)
|
||
max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
|
||
pcp_dcp_combos = []
|
||
if (
|
||
vllm_config.parallel_config.prefill_context_parallel_size > 1
|
||
and dcp_size > 1
|
||
):
|
||
# PCP + DCP attends prefill rows over a DCP-gathered workspace.
|
||
pcp_dcp_combos.append(
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=True,
|
||
COUNT_VALID=True,
|
||
COMPACT_TO_FRONT=True,
|
||
DCP_SIZE=dcp_size,
|
||
DCP_RANK=dcp_rank,
|
||
DCP_INTERLEAVE=dcp_interleave,
|
||
)
|
||
)
|
||
return self._trace_dispatch(self.dispatch)(
|
||
zip_inputs(
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=False,
|
||
COUNT_VALID=False,
|
||
COMPACT_TO_FRONT=False,
|
||
DCP_SIZE=1,
|
||
DCP_RANK=0,
|
||
DCP_INTERLEAVE=1,
|
||
),
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=False,
|
||
COUNT_VALID=True,
|
||
COMPACT_TO_FRONT=True,
|
||
DCP_SIZE=1,
|
||
DCP_RANK=0,
|
||
DCP_INTERLEAVE=1,
|
||
),
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=True,
|
||
COUNT_VALID=True,
|
||
COMPACT_TO_FRONT=True,
|
||
DCP_SIZE=1,
|
||
DCP_RANK=0,
|
||
DCP_INTERLEAVE=1,
|
||
),
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=False,
|
||
COUNT_VALID=True,
|
||
COMPACT_TO_FRONT=False,
|
||
DCP_SIZE=dcp_size,
|
||
DCP_RANK=dcp_rank,
|
||
DCP_INTERLEAVE=dcp_interleave,
|
||
),
|
||
dict(
|
||
HAS_PREFILL_WORKSPACE=False,
|
||
COUNT_VALID=True,
|
||
COMPACT_TO_FRONT=dcp_size > 1,
|
||
DCP_SIZE=dcp_size,
|
||
DCP_RANK=dcp_rank,
|
||
DCP_INTERLEAVE=dcp_interleave,
|
||
),
|
||
*pcp_dcp_combos,
|
||
),
|
||
BLOCK_SIZE=block_size,
|
||
BLOCK_STRIDE_ROWS=block_stride_rows,
|
||
BLOCK_N=128,
|
||
NUM_TOPK_TOKENS=num_topk_tokens,
|
||
max_num_blocks_per_req=max_num_blocks_per_req,
|
||
)
|
||
|
||
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
||
int32_ptr = TritonWarmupTensor(torch.int32)
|
||
token_shape = (1, compile_key.num_topk_tokens)
|
||
token_strides = (16, 1)
|
||
return dict(
|
||
req_id=int32_ptr,
|
||
block_table=TritonWarmupTensor(
|
||
torch.int32,
|
||
shape=(1, compile_key.block_table_stride),
|
||
),
|
||
token_indices=TritonWarmupTensor(
|
||
torch.int32,
|
||
shape=token_shape,
|
||
strides=token_strides,
|
||
),
|
||
out=TritonWarmupTensor(
|
||
torch.int32,
|
||
shape=token_shape,
|
||
strides=token_strides,
|
||
),
|
||
valid_counts=int32_ptr if compile_key.count_valid else None,
|
||
prefill_workspace_request_ids=(
|
||
int32_ptr if compile_key.has_prefill_workspace else None
|
||
),
|
||
prefill_workspace_starts=(
|
||
int32_ptr if compile_key.has_prefill_workspace else None
|
||
),
|
||
max_num_blocks_per_req=1,
|
||
workspace_rank_stride=1,
|
||
BLOCK_SIZE=compile_key.block_size,
|
||
BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,
|
||
BLOCK_N=compile_key.block_n,
|
||
NUM_TOPK_TOKENS=compile_key.num_topk_tokens,
|
||
HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,
|
||
COUNT_VALID=compile_key.count_valid,
|
||
COMPACT_TO_FRONT=compile_key.compact_to_front,
|
||
DCP_SIZE=compile_key.dcp_size,
|
||
DCP_RANK=compile_key.dcp_rank,
|
||
DCP_INTERLEAVE=compile_key.dcp_interleave,
|
||
)
|
||
|
||
@kernel_launcher
|
||
def __call__(
|
||
self,
|
||
req_id: torch.Tensor,
|
||
block_table: torch.Tensor,
|
||
token_indices: torch.Tensor,
|
||
out: torch.Tensor,
|
||
valid_counts: torch.Tensor | None,
|
||
prefill_workspace_request_ids: torch.Tensor | None,
|
||
prefill_workspace_starts: torch.Tensor | None,
|
||
*,
|
||
workspace_rank_stride: int,
|
||
max_num_blocks_per_req: int,
|
||
BLOCK_SIZE: int,
|
||
BLOCK_STRIDE_ROWS: int,
|
||
BLOCK_N: int,
|
||
NUM_TOPK_TOKENS: int,
|
||
HAS_PREFILL_WORKSPACE: bool,
|
||
COUNT_VALID: bool,
|
||
COMPACT_TO_FRONT: bool,
|
||
DCP_SIZE: int,
|
||
DCP_RANK: int,
|
||
DCP_INTERLEAVE: int,
|
||
) -> LaunchSpec:
|
||
single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
|
||
NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID
|
||
)
|
||
# Exact 2D grid: tokens × column tiles
|
||
return (req_id.shape[0], tiles_per_row), dict(
|
||
valid_count_ptr=valid_counts,
|
||
prefill_request_id_ptr=prefill_workspace_request_ids,
|
||
workspace_starts_ptr=prefill_workspace_starts,
|
||
BLOCK_N=block_n,
|
||
HAS_PREFILL=HAS_PREFILL_WORKSPACE,
|
||
SINGLE_TILE=single_tile,
|
||
# Strides in elements
|
||
bt_stride0=block_table.stride(0),
|
||
bt_stride1=block_table.stride(1),
|
||
ti_stride0=token_indices.stride(0),
|
||
ti_stride1=token_indices.stride(1),
|
||
out_stride0=out.stride(0),
|
||
out_stride1=out.stride(1),
|
||
num_warps=num_warps,
|
||
)
|
||
|
||
|
||
def _remap_tiling(
|
||
NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool
|
||
) -> tuple[bool, int, int, int]:
|
||
"""Pick the column tiling for the index remap kernel.
|
||
|
||
Counting the valid slots per row is the only reason the column tiles have to
|
||
talk to each other, so when counting give one program the whole row: the
|
||
count becomes an in-register reduction plus a plain store, needing neither
|
||
atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this
|
||
needs a power-of-two width; other top-k sizes stay tiled and atomic.
|
||
|
||
Returns:
|
||
(single_tile, block_n, tiles_per_row, num_warps)
|
||
|
||
"""
|
||
single_tile = (
|
||
count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS
|
||
)
|
||
if single_tile:
|
||
return True, NUM_TOPK_TOKENS, 1, 8
|
||
return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4
|
||
|
||
|
||
def triton_convert_req_index_to_global_index(
|
||
req_id: torch.Tensor, # int32 [num_tokens]
|
||
block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req]
|
||
token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]
|
||
BLOCK_SIZE: int = 64,
|
||
BLOCK_STRIDE_ROWS: int | None = None,
|
||
NUM_TOPK_TOKENS: int = 2048,
|
||
BLOCK_N: int = 128, # tile width along columns
|
||
HAS_PREFILL_WORKSPACE: bool = False,
|
||
prefill_workspace_request_ids: torch.Tensor | None = None,
|
||
prefill_workspace_starts: torch.Tensor | None = None,
|
||
prefill_workspace_rank_stride: int | None = None,
|
||
dcp_size: int = 1,
|
||
dcp_rank: int = 0,
|
||
cp_kv_cache_interleave_size: int = 1,
|
||
return_valid_counts: bool = False,
|
||
out: torch.Tensor | None = None,
|
||
valid_counts_out: torch.Tensor | None = None,
|
||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||
"""out[token_id, indice_id] =
|
||
block_table[req_id[token_id],
|
||
token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
|
||
+ token_indices[token_id, indice_id] % BLOCK_SIZE
|
||
|
||
Only when token_indices[token_id, indice_id] == -1 do we output -1.
|
||
For safety, we also output -1 if the derived block_id would be
|
||
out-of-bounds.
|
||
|
||
When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets
|
||
instead of global cache slots. prefill_workspace_request_ids and
|
||
prefill_workspace_starts must be provided.
|
||
|
||
prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else
|
||
prefill request index (maps to prefill_workspace_starts)
|
||
prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace
|
||
starts for each prefill request
|
||
|
||
When return_valid_counts is True, also returns the count of valid (non -1)
|
||
indices per row, computed during the same kernel pass (no extra overhead).
|
||
"""
|
||
assert req_id.dtype == torch.int32
|
||
assert block_table.dtype == torch.int32
|
||
assert token_indices.dtype == torch.int32
|
||
assert req_id.shape[0] == token_indices.shape[0], (
|
||
f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "
|
||
"must cover the same tokens; the grid is sized by req_id but the output "
|
||
"is allocated like token_indices, so a longer req_id writes out of bounds"
|
||
)
|
||
assert token_indices.shape[1] == NUM_TOPK_TOKENS
|
||
assert NUM_TOPK_TOKENS % BLOCK_N == 0, (
|
||
f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"
|
||
)
|
||
|
||
if HAS_PREFILL_WORKSPACE:
|
||
assert prefill_workspace_request_ids is not None
|
||
assert prefill_workspace_starts is not None
|
||
assert prefill_workspace_request_ids.dtype == torch.int32
|
||
assert prefill_workspace_starts.dtype == torch.int32
|
||
|
||
num_tokens = req_id.shape[0]
|
||
max_num_blocks_per_req = block_table.shape[1]
|
||
|
||
single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts)
|
||
|
||
# Ensure contiguous tensors on the same device
|
||
req_id_c = req_id.contiguous()
|
||
block_table_c = block_table.contiguous()
|
||
token_indices_c = token_indices.contiguous()
|
||
# When return_valid_counts, the kernel scatters valid entries to a
|
||
# contiguous prefix [0, valid_count) and leaves the tail unwritten, so
|
||
# pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
|
||
# [:topk_length] == exactly the valid set (no dropped tokens).
|
||
if out is None:
|
||
out = (
|
||
torch.full_like(token_indices_c, -1)
|
||
if return_valid_counts
|
||
else torch.empty_like(token_indices_c)
|
||
)
|
||
else:
|
||
assert out.dtype == token_indices_c.dtype
|
||
assert out.device == token_indices_c.device
|
||
assert out.shape == token_indices_c.shape
|
||
assert out.is_contiguous()
|
||
if return_valid_counts:
|
||
out.fill_(-1)
|
||
|
||
valid_counts: torch.Tensor | None = None
|
||
if return_valid_counts:
|
||
if valid_counts_out is None:
|
||
# Zero-init only matters for the atomic accumulation path.
|
||
alloc = torch.empty if single_tile else torch.zeros
|
||
valid_counts = alloc(
|
||
num_tokens, dtype=torch.int32, device=token_indices.device
|
||
)
|
||
else:
|
||
assert valid_counts_out.dtype == torch.int32
|
||
assert valid_counts_out.device == token_indices.device
|
||
assert valid_counts_out.shape == (num_tokens,)
|
||
assert valid_counts_out.is_contiguous()
|
||
valid_counts = valid_counts_out
|
||
if not single_tile:
|
||
valid_counts.zero_()
|
||
|
||
# Prepare prefill pointers
|
||
if HAS_PREFILL_WORKSPACE:
|
||
assert prefill_workspace_request_ids is not None # for mypy
|
||
assert prefill_workspace_starts is not None # for mypy
|
||
assert prefill_workspace_request_ids.is_contiguous()
|
||
assert prefill_workspace_starts.is_contiguous()
|
||
|
||
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
|
||
req_id_c,
|
||
block_table_c,
|
||
token_indices_c,
|
||
out,
|
||
valid_counts,
|
||
prefill_workspace_request_ids,
|
||
prefill_workspace_starts,
|
||
# shapes / constexprs
|
||
workspace_rank_stride=prefill_workspace_rank_stride or 0,
|
||
max_num_blocks_per_req=max_num_blocks_per_req,
|
||
BLOCK_SIZE=BLOCK_SIZE,
|
||
BLOCK_STRIDE_ROWS=(
|
||
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
|
||
),
|
||
BLOCK_N=BLOCK_N,
|
||
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
|
||
HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
|
||
COUNT_VALID=return_valid_counts,
|
||
COMPACT_TO_FRONT=return_valid_counts,
|
||
DCP_SIZE=dcp_size,
|
||
DCP_RANK=dcp_rank,
|
||
DCP_INTERLEAVE=cp_kv_cache_interleave_size,
|
||
)
|
||
|
||
if return_valid_counts:
|
||
assert valid_counts is not None
|
||
return out, valid_counts
|
||
return out
|
||
|
||
|
||
def triton_filter_and_convert_dcp_index(
|
||
req_id: torch.Tensor,
|
||
block_table: torch.Tensor,
|
||
token_indices: torch.Tensor,
|
||
dcp_size: int,
|
||
dcp_rank: int,
|
||
cp_kv_cache_interleave_size: int = 1,
|
||
BLOCK_SIZE: int = 64,
|
||
BLOCK_STRIDE_ROWS: int | None = None,
|
||
NUM_TOPK_TOKENS: int = 2048,
|
||
BLOCK_N: int = 128,
|
||
return_valid_counts: bool = False,
|
||
compact_valid_to_front: bool = True,
|
||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||
"""Filter global per-request indices to this DCP rank's local slots.
|
||
|
||
With ``compact_valid_to_front`` (default), the conversion kernel scatters
|
||
this rank's owned slots to a contiguous prefix ``[0, valid_count)`` and
|
||
leaves the rest ``-1``. DCP filtering marks non-owned slots ``-1`` and so
|
||
creates interior gaps; the trtllm-gen sparse kernel reads the first
|
||
``valid_count`` entries of each row, so they must be a contiguous prefix.
|
||
Compaction is fused into the kernel (atomic slot allocator) rather than a
|
||
separate sort/gather pass. Prefix order is unspecified (only the set matters).
|
||
"""
|
||
assert dcp_size >= 1
|
||
assert 0 <= dcp_rank < dcp_size
|
||
# Interleave groups must align to KV blocks (globally enforced by
|
||
# VllmConfig: block_size % cp_kv_cache_interleave_size == 0); assert the
|
||
# local invariant so local_idx // BLOCK_SIZE never straddles a group.
|
||
assert BLOCK_SIZE % cp_kv_cache_interleave_size == 0, (
|
||
f"BLOCK_SIZE ({BLOCK_SIZE}) must be divisible by "
|
||
f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})."
|
||
)
|
||
assert req_id.dtype == torch.int32
|
||
assert block_table.dtype == torch.int32
|
||
assert token_indices.dtype == torch.int32
|
||
assert token_indices.shape[1] == NUM_TOPK_TOKENS
|
||
assert NUM_TOPK_TOKENS % BLOCK_N == 0
|
||
|
||
if dcp_size == 1:
|
||
return triton_convert_req_index_to_global_index(
|
||
req_id,
|
||
block_table,
|
||
token_indices,
|
||
BLOCK_SIZE=BLOCK_SIZE,
|
||
BLOCK_STRIDE_ROWS=BLOCK_STRIDE_ROWS,
|
||
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
|
||
BLOCK_N=BLOCK_N,
|
||
return_valid_counts=return_valid_counts,
|
||
)
|
||
|
||
num_tokens = req_id.shape[0]
|
||
max_num_blocks_per_req = block_table.shape[1]
|
||
|
||
req_id_c = req_id.contiguous()
|
||
block_table_c = block_table.contiguous()
|
||
token_indices_c = token_indices.contiguous()
|
||
|
||
# The compaction uses the valid-count buffer as a slot allocator, so it
|
||
# requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
|
||
count_valid = return_valid_counts or compact_valid_to_front
|
||
|
||
# The compaction builds on the counting, so it shares the tiling.
|
||
single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)
|
||
|
||
if compact_valid_to_front:
|
||
out = torch.full_like(token_indices_c, -1)
|
||
else:
|
||
out = torch.empty_like(token_indices_c)
|
||
|
||
valid_counts: torch.Tensor | None = None
|
||
if count_valid:
|
||
# Zero-init only matters for the atomic accumulation path.
|
||
alloc = torch.empty if single_tile else torch.zeros
|
||
valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
|
||
|
||
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
|
||
req_id_c,
|
||
block_table_c,
|
||
token_indices_c,
|
||
out,
|
||
valid_counts,
|
||
# No prefill workspace on the DCP decode path.
|
||
None,
|
||
None,
|
||
workspace_rank_stride=0,
|
||
max_num_blocks_per_req=max_num_blocks_per_req,
|
||
BLOCK_SIZE=BLOCK_SIZE,
|
||
BLOCK_STRIDE_ROWS=(
|
||
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
|
||
),
|
||
BLOCK_N=BLOCK_N,
|
||
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
|
||
HAS_PREFILL_WORKSPACE=False,
|
||
COUNT_VALID=count_valid,
|
||
COMPACT_TO_FRONT=compact_valid_to_front,
|
||
DCP_SIZE=dcp_size,
|
||
DCP_RANK=dcp_rank,
|
||
DCP_INTERLEAVE=cp_kv_cache_interleave_size,
|
||
)
|
||
|
||
if return_valid_counts:
|
||
assert valid_counts is not None
|
||
return out, valid_counts
|
||
return out
|
||
|
||
|
||
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()
|