133b71e0be
Signed-off-by: pavelzak <pavel.zakharov@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
1257 lines
49 KiB
Python
1257 lines
49 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Custom Sparse Attention Indexer layers."""
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
import vllm.envs as envs
|
|
from vllm import _custom_ops as ops
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
|
from vllm.config import CUDAGraphMode, get_current_vllm_config
|
|
from vllm.distributed import get_dcp_group, get_pcp_group
|
|
from vllm.forward_context import get_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.custom_op import CustomOp
|
|
from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
|
|
apply_candidate_mask as _apply_candidate_mask,
|
|
)
|
|
from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
|
|
select_candidate_blocks as _select_candidate_blocks,
|
|
)
|
|
from vllm.model_executor.layers.indexer_topk import (
|
|
RADIX_TOPK_WORKSPACE_SIZE,
|
|
get_indexer_topk,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
get_fp8_min_max,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.utils.deep_gemm import (
|
|
fp8_fp4_mqa_logits,
|
|
fp8_fp4_paged_mqa_logits,
|
|
has_deep_gemm,
|
|
)
|
|
from vllm.utils.import_utils import has_cutedsl
|
|
from vllm.utils.torch_utils import (
|
|
LayerNameType,
|
|
_encode_layer_name,
|
|
_resolve_layer_name,
|
|
direct_register_custom_op,
|
|
)
|
|
from vllm.v1.attention.backends.mla.indexer import (
|
|
DeepseekV32IndexerMetadata,
|
|
)
|
|
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
|
|
from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k
|
|
from vllm.v1.worker.workspace import current_workspace_manager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
|
|
MXFP4_BLOCK_SIZE = 32
|
|
|
|
|
|
def _assert_cutedsl_dcp_merge_supported(
|
|
logits: torch.Tensor,
|
|
topk_indices: torch.Tensor,
|
|
k: int,
|
|
) -> None:
|
|
# The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL
|
|
# stable-topk selector); there is no PyTorch fallback. The first cut targets
|
|
# Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix
|
|
# sizing); the Triton pack itself has no shape/topk constraints.
|
|
if not has_cutedsl():
|
|
raise RuntimeError(
|
|
"DCP sparse-indexer merge requires CuteDSL; install it or disable DCP."
|
|
)
|
|
if logits.device.type != "cuda":
|
|
raise RuntimeError("DCP sparse-indexer merge requires CUDA tensors.")
|
|
if logits.dtype != torch.float32 or topk_indices.dtype != torch.int32:
|
|
raise RuntimeError(
|
|
"DCP sparse-indexer merge requires fp32 logits and int32 indices."
|
|
)
|
|
if k not in (512, 1024, 2048):
|
|
raise RuntimeError(
|
|
f"DCP sparse-indexer merge requires index_topk in (512, 1024, 2048); "
|
|
f"got {k}."
|
|
)
|
|
|
|
|
|
def _merge_dcp_topk_global(
|
|
logits: torch.Tensor,
|
|
topk_indices: torch.Tensor,
|
|
topk_tokens: int,
|
|
dcp_rank: int,
|
|
dcp_world_size: int,
|
|
cp_interleave: int,
|
|
row_starts: torch.Tensor | None = None,
|
|
) -> None:
|
|
"""Merge each DCP rank's local top-K into the global top-K.
|
|
|
|
``topk_indices`` are this rank's local top-K positions into its 1/N KV
|
|
shard. A token in the global top-K must also be in its owning rank's local
|
|
top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at
|
|
most that many on its own rank), so exchanging only the per-rank local
|
|
candidates is exact -- equivalent to all-gathering the full logit matrix,
|
|
but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole
|
|
score row. Overwrites ``topk_indices`` with global token ids (``-1`` for
|
|
padding); the attention backend localizes them back to physical slots per
|
|
rank.
|
|
"""
|
|
if dcp_world_size <= 1:
|
|
return
|
|
|
|
# CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's
|
|
# (score, global_id) candidates on-device, all-gather, then the CuteDSL
|
|
# stable-topk selector.
|
|
_assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens)
|
|
from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
|
|
pack_dcp_topk_candidates_cutedsl,
|
|
stable_topk_from_gathered_candidates_cutedsl,
|
|
)
|
|
|
|
packed = torch.empty(
|
|
(*topk_indices.shape, 2),
|
|
dtype=torch.float32,
|
|
device=topk_indices.device,
|
|
)
|
|
pack_dcp_topk_candidates_cutedsl(
|
|
logits,
|
|
topk_indices,
|
|
packed,
|
|
dcp_rank,
|
|
dcp_world_size,
|
|
cp_interleave,
|
|
row_starts,
|
|
)
|
|
gathered = get_dcp_group().all_gather(packed, dim=1)
|
|
stable_topk_from_gathered_candidates_cutedsl(
|
|
gathered, topk_tokens, out=topk_indices
|
|
)
|
|
|
|
|
|
def dcp_gather_kv_rows(
|
|
local_padded: torch.Tensor,
|
|
deinterleave_idx: torch.Tensor,
|
|
gathered_buf: torch.Tensor,
|
|
out_buf: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""All-gather this rank's KV shard and return it in global token order."""
|
|
gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]]
|
|
dist.all_gather_into_tensor(
|
|
gathered, local_padded.contiguous(), group=get_dcp_group().device_group
|
|
)
|
|
return torch.index_select(
|
|
gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]]
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _fused_indexer_q_rope_quant_kernel(
|
|
positions,
|
|
q,
|
|
q_s0,
|
|
q_s1,
|
|
cos_sin_cache,
|
|
cos_sin_s0,
|
|
q_fp8,
|
|
q_fp8_s0,
|
|
q_fp8_s1,
|
|
weights,
|
|
weights_s0,
|
|
weights_s1,
|
|
weights_out,
|
|
weights_out_s0,
|
|
weights_out_s1,
|
|
softmax_scale,
|
|
head_scale,
|
|
fp8_min: tl.constexpr,
|
|
fp8_max: tl.constexpr,
|
|
is_neox: tl.constexpr,
|
|
):
|
|
token = tl.program_id(0)
|
|
head = tl.program_id(1)
|
|
offs32 = tl.arange(0, 32)
|
|
offs64 = tl.arange(0, 64)
|
|
|
|
pos = tl.load(positions + token)
|
|
cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32)
|
|
sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32)
|
|
q_base = q + token * q_s0 + head * q_s1
|
|
out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1
|
|
|
|
if is_neox:
|
|
# NeoX layout, x0 = q[0:32], x1 = q[32:64]
|
|
x0 = tl.load(q_base + offs32).to(tl.float32)
|
|
x1 = tl.load(q_base + 32 + offs32).to(tl.float32)
|
|
else:
|
|
# interleaved layout
|
|
# x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...]
|
|
x0 = tl.load(q_base + offs32 * 2).to(tl.float32)
|
|
x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32)
|
|
r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
|
|
r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32)
|
|
amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1)))
|
|
|
|
q_nope = tl.load(q_base + 64 + offs64).to(tl.float32)
|
|
amax = tl.maximum(amax, tl.max(tl.abs(q_nope)))
|
|
scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max)
|
|
# e8m0 format
|
|
q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
|
|
|
|
if is_neox:
|
|
tl.store(
|
|
out_base + offs32,
|
|
tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
|
|
)
|
|
tl.store(
|
|
out_base + 32 + offs32,
|
|
tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
|
|
)
|
|
else:
|
|
tl.store(
|
|
out_base + offs32 * 2,
|
|
tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
|
|
)
|
|
tl.store(
|
|
out_base + offs32 * 2 + 1,
|
|
tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
|
|
)
|
|
tl.store(
|
|
out_base + 64 + offs64,
|
|
tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
|
|
)
|
|
|
|
weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32)
|
|
tl.store(
|
|
weights_out + token * weights_out_s0 + head * weights_out_s1,
|
|
weight * q_scale * softmax_scale * head_scale,
|
|
)
|
|
|
|
|
|
def fused_indexer_q_rope_quant(
|
|
positions: torch.Tensor,
|
|
q: torch.Tensor,
|
|
cos_sin_cache: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
softmax_scale: float,
|
|
head_scale: float,
|
|
is_neox: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
assert current_platform.is_cuda()
|
|
assert q.dtype == torch.bfloat16
|
|
assert q.shape[-1] == 128
|
|
assert cos_sin_cache.shape[-1] == 64
|
|
assert weights.shape == q.shape[:2]
|
|
|
|
q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype())
|
|
weights_out = torch.empty_like(weights, dtype=torch.float32)
|
|
fp8_min, fp8_max = get_fp8_min_max()
|
|
_fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])](
|
|
positions,
|
|
q,
|
|
q.stride(0),
|
|
q.stride(1),
|
|
cos_sin_cache,
|
|
cos_sin_cache.stride(0),
|
|
q_fp8,
|
|
q_fp8.stride(0),
|
|
q_fp8.stride(1),
|
|
weights,
|
|
weights.stride(0),
|
|
weights.stride(1),
|
|
weights_out,
|
|
weights_out.stride(0),
|
|
weights_out.stride(1),
|
|
softmax_scale,
|
|
head_scale,
|
|
fp8_min=fp8_min,
|
|
fp8_max=fp8_max,
|
|
is_neox=is_neox,
|
|
num_warps=1,
|
|
)
|
|
return q_fp8, weights_out
|
|
|
|
|
|
def _gather_workspace_shapes(
|
|
total_seq_lens: int,
|
|
head_dim: int,
|
|
fp8_dtype: torch.dtype,
|
|
use_fp4_cache: bool,
|
|
) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
|
|
"""Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
|
|
the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
|
|
scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
|
|
(T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
|
|
if use_fp4_cache:
|
|
return (
|
|
((total_seq_lens, head_dim // 2), torch.uint8),
|
|
((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
|
|
)
|
|
return (
|
|
((total_seq_lens, head_dim), fp8_dtype),
|
|
((total_seq_lens, 4), torch.uint8),
|
|
)
|
|
|
|
|
|
def kv_cache_as_quant_view(
|
|
kv_cache: torch.Tensor,
|
|
head_dim: int,
|
|
use_fp4_cache: bool,
|
|
) -> torch.Tensor:
|
|
"""4D ``[num_blocks, block_size, 1, head_width]`` view expected by
|
|
DeepGEMM, from the 3D indexer kv-cache allocation."""
|
|
if use_fp4_cache:
|
|
assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
|
|
num_blocks, block_size, _ = kv_cache.shape
|
|
page_bytes = int(kv_cache.stride(0))
|
|
fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
|
|
return torch.as_strided(
|
|
kv_cache,
|
|
size=(num_blocks, block_size, 1, fp4_bytes),
|
|
stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
|
|
)
|
|
return kv_cache.unsqueeze(-2)
|
|
|
|
|
|
@eager_break_during_capture
|
|
def sparse_attn_indexer(
|
|
hidden_states: torch.Tensor,
|
|
k_cache_prefix: LayerNameType,
|
|
kv_cache: torch.Tensor,
|
|
q_quant: torch.Tensor,
|
|
q_scale: torch.Tensor | None,
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
quant_block_size: int,
|
|
scale_fmt: str | None,
|
|
topk_tokens: int,
|
|
head_dim: int,
|
|
max_model_len: int,
|
|
total_seq_lens: int,
|
|
topk_indices_buffer: torch.Tensor,
|
|
skip_k_cache_insert: bool,
|
|
use_pcp: bool,
|
|
dense_mha_metadata_layer_name: LayerNameType,
|
|
use_fp4_cache: bool = False,
|
|
dcp_rank: int = 0,
|
|
dcp_world_size: int = 1,
|
|
cp_kv_cache_interleave_size: int = 1,
|
|
skip_topk_buffer_clear: bool = False,
|
|
candidate_blocks: torch.Tensor | None = None,
|
|
candidate_block_size: int = 0,
|
|
candidate_write: bool = False,
|
|
topk_backend: str = "auto",
|
|
) -> torch.Tensor:
|
|
# careful! this will be None in dummy run
|
|
forward_context = get_forward_context()
|
|
attn_metadata = forward_context.attn_metadata
|
|
fp8_dtype = current_platform.fp8_dtype()
|
|
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
|
|
|
|
if candidate_blocks is not None:
|
|
# Candidate blocks are request-local; the DCP-sharded logits layout
|
|
# would need per-rank translation that is not implemented.
|
|
assert dcp_world_size == 1, (
|
|
"v4.1 two-level candidate filtering is not supported with DCP."
|
|
)
|
|
assert candidate_block_size > 0
|
|
|
|
# assert isinstance(attn_metadata, dict)
|
|
if not isinstance(attn_metadata, dict):
|
|
# Reserve workspace for indexer during profiling run
|
|
values_spec, scales_spec = _gather_workspace_shapes(
|
|
total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
|
|
)
|
|
profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [
|
|
values_spec,
|
|
scales_spec,
|
|
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
|
|
]
|
|
if use_pcp and dcp_world_size > 1:
|
|
# The PCP+DCP path takes an all-gather destination and a
|
|
# de-interleaved result.
|
|
gather_spec = _gather_workspace_shapes(
|
|
total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
|
|
)
|
|
profile_specs.extend(gather_spec * 2)
|
|
current_workspace_manager().get_simultaneous(*profile_specs)
|
|
|
|
# Dummy allocation to simulate for peak logits tensor memory during inference.
|
|
# FP8 elements so elements == bytes
|
|
max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
|
|
_ = torch.empty(
|
|
max_logits_elems, dtype=torch.uint8, device=hidden_states.device
|
|
)
|
|
|
|
return sparse_attn_indexer_fake(
|
|
hidden_states,
|
|
k_cache_prefix,
|
|
kv_cache,
|
|
q_quant,
|
|
q_scale,
|
|
k,
|
|
weights,
|
|
quant_block_size,
|
|
scale_fmt,
|
|
topk_tokens,
|
|
head_dim,
|
|
max_model_len,
|
|
total_seq_lens,
|
|
topk_indices_buffer,
|
|
skip_k_cache_insert,
|
|
use_pcp,
|
|
dense_mha_metadata_layer_name,
|
|
use_fp4_cache,
|
|
candidate_blocks=candidate_blocks,
|
|
candidate_block_size=candidate_block_size,
|
|
candidate_write=candidate_write,
|
|
)
|
|
attn_metadata_narrowed = attn_metadata[k_cache_prefix]
|
|
assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
|
|
slot_mapping = attn_metadata_narrowed.slot_mapping
|
|
has_decode = attn_metadata_narrowed.num_decodes > 0
|
|
has_prefill = attn_metadata_narrowed.num_prefills > 0
|
|
num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
|
|
|
|
# q_scale is required iff the FP4 cache path is enabled; the FP8 path
|
|
# folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
|
|
if use_fp4_cache:
|
|
assert q_scale is not None, "use_fp4_cache=True requires q_scale"
|
|
else:
|
|
assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
|
|
|
|
# During speculative decoding, k may be padded to the CUDA graph batch
|
|
# size while slot_mapping only covers actual tokens. Truncate k to avoid
|
|
# out-of-bounds reads in the kernel.
|
|
# Keep PCP padding so every rank contributes the same all-gather shape.
|
|
num_tokens = slot_mapping.shape[0]
|
|
if k is not None:
|
|
# MTP draft decodes are replicated and use an unexpanded slot mapping.
|
|
if use_pcp and num_tokens > k.shape[0]:
|
|
num_tokens //= get_pcp_group().world_size
|
|
k = k[:num_tokens]
|
|
|
|
if not skip_k_cache_insert:
|
|
assert k is not None
|
|
k, slot_mapping_for_cache = maybe_gather_indexer_k(
|
|
k,
|
|
slot_mapping,
|
|
num_decode_tokens,
|
|
use_pcp,
|
|
)
|
|
# scale_fmt can be None, but the function expects str
|
|
assert scale_fmt is not None
|
|
assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
|
|
ops.indexer_k_quant_and_cache(
|
|
k,
|
|
kv_cache,
|
|
slot_mapping_for_cache,
|
|
quant_block_size,
|
|
scale_fmt,
|
|
)
|
|
# The indexer and main MLA may classify the same short extend differently
|
|
# because they use independent decode thresholds. Only the main MLA route
|
|
# can determine whether the top-k indices will be consumed.
|
|
if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL:
|
|
dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name)
|
|
if dense_mha_layer:
|
|
mla_metadata = attn_metadata.get(dense_mha_layer)
|
|
prefill_metadata = getattr(mla_metadata, "prefill", None)
|
|
if (
|
|
getattr(prefill_metadata, "use_dense_mha", False)
|
|
and getattr(mla_metadata, "num_decode_tokens", -1) == 0
|
|
and not torch.cuda.is_current_stream_capturing()
|
|
):
|
|
# Deliberately leave the buffer untouched. Dense MHA does not
|
|
# consume top-k indices for this batch; clearing it would be
|
|
# unnecessary work.
|
|
return topk_indices_buffer
|
|
|
|
# The buffer must be pre-filled with -1 (the "no token" sentinel) before the
|
|
# top-k kernels scatter valid indices into it. On the fused deepseek_v32
|
|
# nvidia path, _fused_norm_rope_kernel already cleared the same
|
|
# [:num_tokens, :topk] region earlier in this forward, so skip the redundant
|
|
# fill.
|
|
if not skip_topk_buffer_clear:
|
|
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
|
if has_prefill:
|
|
prefill_metadata = attn_metadata_narrowed.prefill
|
|
assert prefill_metadata is not None
|
|
|
|
# Get the full shared workspace buffers once (will allocate on first use).
|
|
# Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
|
|
# MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
|
|
# scales) based on use_fp4_cache.
|
|
workspace_manager = current_workspace_manager()
|
|
values_spec, scales_spec = _gather_workspace_shapes(
|
|
total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
|
|
)
|
|
# PCP + DCP needs two more pairs: the rank-major all-gather destination
|
|
# and the de-interleaved result.
|
|
pcp_chunks = [
|
|
(c, c.pcp_deinterleave_idx)
|
|
for c in prefill_metadata.chunks
|
|
if c.pcp_deinterleave_idx is not None
|
|
]
|
|
gather_specs: list[tuple[tuple[int, int], torch.dtype]] = []
|
|
if pcp_chunks:
|
|
gathered_rows = max(
|
|
dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks
|
|
)
|
|
deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks)
|
|
for rows in (gathered_rows, deinterleaved_rows):
|
|
gather_specs.extend(
|
|
_gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache)
|
|
)
|
|
k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous(
|
|
values_spec,
|
|
scales_spec,
|
|
*gather_specs,
|
|
)
|
|
for chunk in prefill_metadata.chunks:
|
|
cu_seqlen_ks = chunk.cu_seqlen_ks
|
|
cu_seqlen_ke = chunk.cu_seqlen_ke
|
|
assert chunk.local_cu_seq_lens is not None
|
|
k_quant = k_quant_full[: chunk.max_local_total_seq_lens]
|
|
k_scale = k_scale_full[: chunk.max_local_total_seq_lens]
|
|
if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0:
|
|
ops.cp_gather_indexer_k_quant_cache(
|
|
kv_cache,
|
|
k_quant,
|
|
k_scale,
|
|
chunk.block_table,
|
|
chunk.local_cu_seq_lens,
|
|
)
|
|
|
|
# PCP + DCP KV all-gather.
|
|
deinterleave_idx = chunk.pcp_deinterleave_idx
|
|
if deinterleave_idx is not None:
|
|
# local_total_seq_lens is the PCP-padded extent here, so every
|
|
# rank contributes an identically shaped shard.
|
|
gathered_values, gathered_scales, out_values, out_scales = gather_bufs
|
|
shard = k_quant[: chunk.local_total_seq_lens]
|
|
k_quant = dcp_gather_kv_rows(
|
|
shard, deinterleave_idx, gathered_values, out_values
|
|
)
|
|
k_scale = dcp_gather_kv_rows(
|
|
k_scale[: chunk.local_total_seq_lens],
|
|
deinterleave_idx,
|
|
gathered_scales,
|
|
out_scales,
|
|
)
|
|
|
|
q_slice = q_quant[chunk.token_start : chunk.token_end]
|
|
q_scale_slice = (
|
|
q_scale[chunk.token_start : chunk.token_end]
|
|
if q_scale is not None
|
|
else None
|
|
)
|
|
topk_indices = topk_indices_buffer[
|
|
chunk.token_start : chunk.token_end, :topk_tokens
|
|
]
|
|
|
|
if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0:
|
|
logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
|
|
topk_indices.fill_(-1)
|
|
else:
|
|
# DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
|
|
# (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
|
|
if use_fp4_cache:
|
|
q_slice_cast = q_slice.view(torch.int8)
|
|
k_quant_cast = k_quant.view(torch.int8)
|
|
k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
|
|
else:
|
|
q_slice_cast = q_slice
|
|
k_quant_cast = k_quant
|
|
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
|
|
if current_platform.is_xpu():
|
|
if q_scale_slice is not None:
|
|
raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
|
|
logits = torch.ops.vllm.xpu_fp8_mqa_logits(
|
|
q_slice_cast,
|
|
k_quant_cast,
|
|
k_scale_cast,
|
|
weights[chunk.token_start : chunk.token_end],
|
|
cu_seqlen_ks,
|
|
cu_seqlen_ke,
|
|
)
|
|
else:
|
|
logits = fp8_fp4_mqa_logits(
|
|
(q_slice_cast, q_scale_slice),
|
|
(k_quant_cast, k_scale_cast),
|
|
weights[chunk.token_start : chunk.token_end],
|
|
cu_seqlen_ks,
|
|
cu_seqlen_ke,
|
|
clean_logits=False,
|
|
)
|
|
num_rows = logits.shape[0]
|
|
if candidate_blocks is not None:
|
|
# Two-level selection (v4.1): the candidate source
|
|
# publishes its top blocks; later indexers mask their
|
|
# scores to them. Both before the row top-k.
|
|
chunk_candidates = candidate_blocks[
|
|
chunk.token_start : chunk.token_end
|
|
]
|
|
if candidate_write:
|
|
_select_candidate_blocks(
|
|
logits,
|
|
cu_seqlen_ks,
|
|
cu_seqlen_ke,
|
|
chunk_candidates.shape[1],
|
|
candidate_block_size,
|
|
chunk_candidates,
|
|
)
|
|
else:
|
|
_apply_candidate_mask(
|
|
logits,
|
|
cu_seqlen_ks,
|
|
cu_seqlen_ke,
|
|
chunk_candidates,
|
|
candidate_block_size,
|
|
)
|
|
ops.top_k_per_row_prefill(
|
|
logits,
|
|
cu_seqlen_ks,
|
|
cu_seqlen_ke,
|
|
topk_indices,
|
|
num_rows,
|
|
logits.stride(0),
|
|
logits.stride(1),
|
|
topk_tokens,
|
|
)
|
|
|
|
if deinterleave_idx is None:
|
|
# Under the PCP path the top-k already ran over the whole
|
|
# context.
|
|
_merge_dcp_topk_global(
|
|
logits,
|
|
topk_indices,
|
|
topk_tokens,
|
|
dcp_rank,
|
|
dcp_world_size,
|
|
cp_kv_cache_interleave_size,
|
|
row_starts=chunk.cu_seqlen_ks,
|
|
)
|
|
|
|
if has_decode:
|
|
decode_metadata = attn_metadata_narrowed.decode
|
|
assert decode_metadata is not None
|
|
kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
|
|
decode_lens = decode_metadata.decode_lens
|
|
# requires_padding can be computed False for ragged warmup/mixed
|
|
# batches (seen on SM120 TP=2: 8 tokens over 6 seqs -> uniform
|
|
# reshape below would crash). The token-count/batch-size divisibility
|
|
# check is a necessary-but-not-sufficient proxy for "uniform batch"
|
|
# (e.g. decode_lens=[1, 2, 3] divides evenly by 3 but isn't uniform);
|
|
# it catches the specific case seen in practice without requiring a
|
|
# device sync to inspect decode_lens directly, which would break
|
|
# CUDA-graph-capture safety on this path. Shared by both the pack
|
|
# branch below and the matching unpack branch further down — they
|
|
# must agree on which path was taken.
|
|
needs_padded_path = decode_metadata.requires_padding or (
|
|
num_decode_tokens % decode_lens.shape[0] != 0
|
|
)
|
|
if num_decode_tokens == 0:
|
|
padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])
|
|
padded_q_scale = (
|
|
q_scale[:1].reshape(1, 1, *q_scale.shape[1:])
|
|
if q_scale is not None
|
|
else None
|
|
)
|
|
elif needs_padded_path:
|
|
# pad in edge case where we have short chunked prefill length <
|
|
# decode_threshold since we unstrictly split
|
|
# prefill and decode by decode_threshold
|
|
# (currently set to 1 + speculative tokens).
|
|
# FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK —
|
|
# downstream context_lens masks stale slots). MXFP4 Q is two
|
|
# uint8 tensors (values + ue8m0 scales) — use the dedicated uint8
|
|
# packer with pad_byte=0 so padded slots dequantize to 0 and
|
|
# can't produce NaN/Inf in the logits kernel.
|
|
if q_scale is not None:
|
|
padded_q_quant_decode_tokens = pack_seq_triton(
|
|
q_quant[:num_decode_tokens], decode_lens, pad_value=0
|
|
)
|
|
padded_q_scale = pack_seq_triton(
|
|
q_scale[:num_decode_tokens], decode_lens, pad_value=0
|
|
)
|
|
else:
|
|
padded_q_quant_decode_tokens = pack_seq_triton(
|
|
q_quant[:num_decode_tokens], decode_lens
|
|
)
|
|
padded_q_scale = None
|
|
else:
|
|
padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
|
|
decode_lens.shape[0], -1, *q_quant.shape[1:]
|
|
)
|
|
if q_scale is not None:
|
|
padded_q_scale = q_scale[:num_decode_tokens].reshape(
|
|
decode_lens.shape[0], -1, *q_scale.shape[1:]
|
|
)
|
|
else:
|
|
padded_q_scale = None
|
|
# TODO: move and optimize below logic with triton kernels
|
|
batch_size = padded_q_quant_decode_tokens.shape[0]
|
|
next_n = padded_q_quant_decode_tokens.shape[1]
|
|
num_padded_tokens = batch_size * next_n
|
|
seq_lens = decode_metadata.seq_lens[:batch_size]
|
|
# seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
|
|
# otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
|
|
# the downstream topk kernels accept both 1D and 2D.
|
|
padded_q_quant_cast = (
|
|
padded_q_quant_decode_tokens.view(torch.int8)
|
|
if use_fp4_cache
|
|
else padded_q_quant_decode_tokens
|
|
)
|
|
if current_platform.is_xpu():
|
|
if padded_q_scale is not None:
|
|
raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q")
|
|
seq_lens_xpu = (
|
|
seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
|
|
)
|
|
logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits(
|
|
padded_q_quant_cast,
|
|
kv_cache,
|
|
weights[:num_padded_tokens],
|
|
seq_lens_xpu,
|
|
decode_metadata.block_table,
|
|
decode_metadata.schedule_metadata,
|
|
max_model_len,
|
|
)
|
|
else:
|
|
logits = fp8_fp4_paged_mqa_logits(
|
|
(padded_q_quant_cast, padded_q_scale),
|
|
kv_cache,
|
|
weights[:num_padded_tokens],
|
|
seq_lens,
|
|
decode_metadata.block_table,
|
|
decode_metadata.schedule_metadata,
|
|
max_model_len=max_model_len,
|
|
clean_logits=False,
|
|
indices=decode_metadata.indices,
|
|
)
|
|
num_rows = logits.shape[0]
|
|
if candidate_blocks is not None:
|
|
# Two-level selection (v4.1) on the decode logits; columns are
|
|
# request-local compressed positions. seq_lens is (B, next_n)
|
|
# for native spec decode (per-row effective lens) and (B, 1)
|
|
# otherwise.
|
|
vis = seq_lens.reshape(-1)
|
|
row_repeat = next_n if vis.numel() != num_rows else 1
|
|
vis = vis[:num_rows]
|
|
decode_candidates = candidate_blocks[:num_rows]
|
|
if candidate_write:
|
|
_select_candidate_blocks(
|
|
logits,
|
|
None,
|
|
vis,
|
|
decode_candidates.shape[1],
|
|
candidate_block_size,
|
|
decode_candidates,
|
|
row_repeat,
|
|
)
|
|
else:
|
|
_apply_candidate_mask(
|
|
logits,
|
|
None,
|
|
vis,
|
|
decode_candidates,
|
|
candidate_block_size,
|
|
row_repeat,
|
|
)
|
|
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
|
|
|
|
# The backend comes from the layer (config is only readable at model
|
|
# construction); dispatchers are cached per backend.
|
|
get_indexer_topk(topk_backend)(
|
|
logits,
|
|
seq_lens,
|
|
next_n,
|
|
topk_indices,
|
|
topk_tokens,
|
|
attn_metadata_narrowed.max_seq_len,
|
|
)
|
|
|
|
if decode_metadata.global_seq_lens is not None:
|
|
_merge_dcp_topk_global(
|
|
logits,
|
|
topk_indices,
|
|
topk_tokens,
|
|
dcp_rank,
|
|
dcp_world_size,
|
|
cp_kv_cache_interleave_size,
|
|
)
|
|
|
|
if needs_padded_path:
|
|
# if padded, we need to unpack
|
|
# the topk indices removing padded tokens
|
|
topk_indices = unpack_seq_triton(
|
|
topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
|
|
decode_lens,
|
|
)
|
|
topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = (
|
|
topk_indices
|
|
)
|
|
|
|
return topk_indices_buffer
|
|
|
|
|
|
def sparse_attn_indexer_fake(
|
|
hidden_states: torch.Tensor,
|
|
k_cache_prefix: LayerNameType,
|
|
kv_cache: torch.Tensor,
|
|
q_quant: torch.Tensor,
|
|
q_scale: torch.Tensor | None,
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
quant_block_size: int,
|
|
scale_fmt: str | None,
|
|
topk_tokens: int,
|
|
head_dim: int,
|
|
max_model_len: int,
|
|
total_seq_lens: int,
|
|
topk_indices_buffer: torch.Tensor | None,
|
|
skip_k_cache_insert: bool,
|
|
use_pcp: bool,
|
|
dense_mha_metadata_layer_name: LayerNameType,
|
|
use_fp4_cache: bool = False,
|
|
dcp_rank: int = 0,
|
|
dcp_world_size: int = 1,
|
|
cp_kv_cache_interleave_size: int = 1,
|
|
skip_topk_buffer_clear: bool = False,
|
|
candidate_blocks: torch.Tensor | None = None,
|
|
candidate_block_size: int = 0,
|
|
candidate_write: bool = False,
|
|
topk_backend: str = "auto",
|
|
) -> torch.Tensor:
|
|
return topk_indices_buffer
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="sparse_attn_indexer",
|
|
op_func=sparse_attn_indexer,
|
|
mutates_args=["topk_indices_buffer", "candidate_blocks"],
|
|
fake_impl=sparse_attn_indexer_fake,
|
|
dispatch_key=current_platform.dispatch_key,
|
|
)
|
|
|
|
|
|
@CustomOp.register("sparse_attn_indexer")
|
|
class SparseAttnIndexer(CustomOp):
|
|
"""Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
|
|
separate custom op since it involves heavy custom kernels like `mqa_logits`,
|
|
`paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
|
|
specific memory layout or implementation for different hardware backends to
|
|
achieve optimal performance.
|
|
|
|
For now, the default native path will use CUDA backend path. Other platform
|
|
may requires add the corresponding Custom Op name `sparse_attn_indexer` to
|
|
`custom_ops` in `CompilationConfig` to enable the platform specific path.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
k_cache,
|
|
quant_block_size: int,
|
|
scale_fmt: str,
|
|
topk_tokens: int,
|
|
head_dim: int,
|
|
max_model_len: int,
|
|
max_total_seq_len: int,
|
|
topk_indices_buffer: torch.Tensor,
|
|
skip_k_cache_insert: bool = False,
|
|
use_fp4_cache: bool = False,
|
|
compress_ratio: int = 1,
|
|
candidate_blocks: torch.Tensor | None = None,
|
|
candidate_block_size: int = 0,
|
|
candidate_write: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.k_cache = k_cache
|
|
self.quant_block_size = quant_block_size
|
|
self.scale_fmt = scale_fmt
|
|
self.topk_tokens = topk_tokens
|
|
self.head_dim = head_dim
|
|
self.max_model_len = max_model_len
|
|
self.max_total_seq_len = max_total_seq_len
|
|
self.topk_indices_buffer = topk_indices_buffer
|
|
self.skip_k_cache_insert = skip_k_cache_insert
|
|
self.use_fp4_cache = use_fp4_cache
|
|
self.compress_ratio = compress_ratio
|
|
# v4.1 two-level selection: the candidate source indexer writes the
|
|
# top candidate blocks here; later indexers mask their scores with it.
|
|
self.candidate_blocks = candidate_blocks
|
|
self.candidate_block_size = candidate_block_size
|
|
self.candidate_write = candidate_write
|
|
self.dense_mha_metadata_layer_name = ""
|
|
# DCP scalars are constant for the run; resolve them here (config is set
|
|
# during model construction) and pass them into the custom op, rather
|
|
# than threading them through per-step metadata.
|
|
vllm_config = get_current_vllm_config()
|
|
parallel_config = vllm_config.parallel_config
|
|
self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
|
|
self._parallel_config = parallel_config
|
|
self.dcp_world_size = parallel_config.decode_context_parallel_size
|
|
self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
|
|
self.use_pcp = parallel_config.prefill_context_parallel_size > 1
|
|
self._cp_kv_cache_interleave_size: int | None = None
|
|
if current_platform.is_cuda() and not has_deep_gemm():
|
|
raise RuntimeError(
|
|
"Sparse Attention Indexer CUDA op requires DeepGEMM support in "
|
|
"the current vLLM environment."
|
|
)
|
|
|
|
if vllm_config.kernel_config.enable_jit_warmup:
|
|
from vllm.v1.attention.ops.common import (
|
|
_PACK_SEQ_TRITON_KERNEL,
|
|
_UNPACK_SEQ_TRITON_KERNEL,
|
|
)
|
|
|
|
pack_dtype = torch.uint8 if use_fp4_cache else current_platform.fp8_dtype()
|
|
_PACK_SEQ_TRITON_KERNEL.register_warmup(
|
|
dtype=pack_dtype,
|
|
pad_value=0 if use_fp4_cache else -float("inf"),
|
|
)
|
|
_UNPACK_SEQ_TRITON_KERNEL.register_warmup()
|
|
|
|
if self.dcp_world_size > 1 and current_platform.is_cuda() and has_cutedsl():
|
|
from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( # noqa: E501
|
|
_PACK_DCP_TOPK_CANDIDATES_KERNEL,
|
|
_STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL,
|
|
)
|
|
|
|
_PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup()
|
|
_STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup()
|
|
|
|
@property
|
|
def cp_kv_cache_interleave_size(self) -> int:
|
|
"""With PD+DCP, the real value isn't known until block_size is finalized,
|
|
which happens after this layer is built. Safe to cache after the first access,
|
|
as long as the adjustment always runs before any forward pass
|
|
(it's set up in Worker.initialize_from_config, ahead of warmup/serving).
|
|
"""
|
|
if self._cp_kv_cache_interleave_size is None:
|
|
value = self._parallel_config.cp_kv_cache_interleave_size
|
|
if isinstance(get_forward_context().attn_metadata, dict):
|
|
self._cp_kv_cache_interleave_size = value
|
|
return value
|
|
return self._cp_kv_cache_interleave_size
|
|
|
|
def forward_native(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
):
|
|
if current_platform.is_cuda() or current_platform.is_xpu():
|
|
return self.forward_cuda(hidden_states, q_quant, k, weights)
|
|
elif current_platform.is_rocm():
|
|
if self.use_pcp and self.dcp_world_size > 1:
|
|
raise NotImplementedError(
|
|
"The ROCm sparse-indexer path does not support PCP+DCP."
|
|
)
|
|
return self.forward_hip(hidden_states, q_quant, k, weights)
|
|
elif current_platform.is_cpu():
|
|
return self.forward_cpu(hidden_states, q_quant, k, weights)
|
|
else:
|
|
raise NotImplementedError(
|
|
"SparseAttnIndexer native forward is only implemented for "
|
|
"CUDA, ROCm, XPU and CPU platforms."
|
|
)
|
|
|
|
def forward_cuda(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
):
|
|
# FP8 path: single tensor (per-token scale is folded into `weights`).
|
|
# FP4 path: (values, scales) tuple with scales required by the kernel.
|
|
if isinstance(q_quant, tuple):
|
|
q_values, q_scale = q_quant
|
|
else:
|
|
q_values, q_scale = q_quant, None
|
|
return torch.ops.vllm.sparse_attn_indexer(
|
|
hidden_states,
|
|
_encode_layer_name(self.k_cache.prefix),
|
|
self.k_cache.kv_cache,
|
|
q_values,
|
|
q_scale,
|
|
k,
|
|
weights,
|
|
self.quant_block_size,
|
|
self.scale_fmt,
|
|
self.topk_tokens,
|
|
self.head_dim,
|
|
self.max_model_len,
|
|
self.max_total_seq_len,
|
|
self.topk_indices_buffer,
|
|
self.skip_k_cache_insert,
|
|
self.use_pcp,
|
|
_encode_layer_name(self.dense_mha_metadata_layer_name),
|
|
self.use_fp4_cache,
|
|
self.dcp_rank,
|
|
self.dcp_world_size,
|
|
self.cp_kv_cache_interleave_size,
|
|
False,
|
|
candidate_blocks=self.candidate_blocks,
|
|
candidate_block_size=self.candidate_block_size,
|
|
candidate_write=self.candidate_write,
|
|
topk_backend=self.topk_backend,
|
|
)
|
|
|
|
def forward_xpu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_fp8: torch.Tensor,
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
):
|
|
return self.forward_cuda(hidden_states, q_fp8, k, weights)
|
|
|
|
def forward_hip(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
):
|
|
assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
|
|
assert isinstance(q_quant, torch.Tensor), (
|
|
"AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
|
|
)
|
|
from vllm.platforms.rocm import on_gfx11, on_gfx950
|
|
|
|
if (
|
|
rocm_aiter_ops.is_enabled()
|
|
or rocm_aiter_ops.is_rdna_aiter_enabled()
|
|
or on_gfx11()
|
|
# The so-called AITER sparse indexer op has a native gfx950 path:
|
|
# its cache insert, MQA logits, and top-k fallbacks are implemented
|
|
# by local Triton/C++ kernels and do not require the aiter package.
|
|
or on_gfx950()
|
|
):
|
|
return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
|
|
hidden_states,
|
|
_encode_layer_name(self.k_cache.prefix),
|
|
self.k_cache.kv_cache,
|
|
q_quant,
|
|
k,
|
|
weights,
|
|
self.quant_block_size,
|
|
self.scale_fmt,
|
|
self.topk_tokens,
|
|
self.head_dim,
|
|
self.max_model_len,
|
|
self.max_total_seq_len,
|
|
self.topk_indices_buffer,
|
|
skip_k_cache_insert=self.skip_k_cache_insert,
|
|
compress_ratio=self.compress_ratio,
|
|
candidate_blocks=self.candidate_blocks,
|
|
candidate_block_size=self.candidate_block_size,
|
|
candidate_write=self.candidate_write,
|
|
)
|
|
raise RuntimeError(
|
|
"Sparse attention indexer ROCm path requires AITER or a supported "
|
|
"native architecture (gfx950/gfx11)."
|
|
)
|
|
|
|
def forward_cpu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor | None,
|
|
weights: torch.Tensor,
|
|
):
|
|
"""CPU sparse attention indexer: cache write stays eager Python glue
|
|
(own K-cache layout, not shared with the main attention cache
|
|
write). PREFILL and DECODE both call the ported
|
|
``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels,
|
|
which read the paged K-cache directly via ``page_table`` -- no
|
|
eager gather step, no per-request Python loop.
|
|
|
|
``prefill_metadata.chunks`` always has exactly one entry here:
|
|
``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk
|
|
split to always return the whole step's prefill batch as one
|
|
chunk, since the base chunking only bounds CUDA/XPU's dense M*N
|
|
logits tensor and flat K-gather workspace, neither of which this
|
|
paged kernel allocates.
|
|
"""
|
|
assert not self.use_fp4_cache, (
|
|
"CPU sparse indexer doesn't support fp4 cache yet"
|
|
)
|
|
assert isinstance(q_quant, torch.Tensor), (
|
|
"CPU sparse_attn_indexer expects a single FP8 q_quant tensor"
|
|
)
|
|
assert self.dcp_world_size <= 1 and not self.use_pcp, (
|
|
"CPU sparse indexer doesn't support decode/prefill context parallelism yet."
|
|
)
|
|
|
|
forward_context = get_forward_context()
|
|
attn_metadata = forward_context.attn_metadata
|
|
attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None
|
|
if isinstance(attn_metadata, dict):
|
|
metadata = attn_metadata[self.k_cache.prefix]
|
|
assert isinstance(metadata, DeepseekV32IndexerMetadata)
|
|
attn_metadata_narrowed = metadata
|
|
if attn_metadata_narrowed is None:
|
|
# Profiling/dummy run: no real metadata to act on.
|
|
return self.topk_indices_buffer
|
|
|
|
kv_cache = self.k_cache.kv_cache
|
|
topk_tokens = self.topk_tokens
|
|
topk_indices_buffer = self.topk_indices_buffer
|
|
slot_mapping = attn_metadata_narrowed.slot_mapping
|
|
has_decode = attn_metadata_narrowed.num_decodes > 0
|
|
has_prefill = attn_metadata_narrowed.num_prefills > 0
|
|
num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
|
|
|
|
num_tokens = slot_mapping.shape[0]
|
|
if k is not None:
|
|
k = k[:num_tokens]
|
|
|
|
if not self.skip_k_cache_insert:
|
|
# Only reachable via DeepseekV32Attention with
|
|
# prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu
|
|
# (the kernels this used to call) have been removed as unused/
|
|
# untested (csrc/cpu/sgl-kernels/store_cache.cpp).
|
|
raise NotImplementedError(
|
|
"SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False "
|
|
"(prefill context parallel on CPU) is not supported."
|
|
)
|
|
|
|
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
|
|
|
if has_prefill:
|
|
assert topk_tokens == 512, (
|
|
"topk_transform_512_cpu only supports index_topk == 512."
|
|
)
|
|
prefill_metadata = attn_metadata_narrowed.prefill
|
|
assert prefill_metadata is not None
|
|
assert len(prefill_metadata.chunks) == 1, (
|
|
"forward_cpu expects the prefill metadata builder to always "
|
|
"produce a single chunk -- see "
|
|
"DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks."
|
|
)
|
|
chunk = prefill_metadata.chunks[0]
|
|
# kv_cache is a per-layer view into vLLM's shared multi-layer
|
|
# cache allocation, so its block stride generally exceeds
|
|
# block_size * page_width; fp8_paged_mqa_logits_cpu reads
|
|
# kv_view.stride(0) explicitly, so no copy is needed here.
|
|
kv_view = kv_cache.view(kv_cache.shape[0], -1)
|
|
block_size = kv_cache.shape[1]
|
|
q_slice = q_quant[chunk.token_start : chunk.token_end]
|
|
topk_indices = topk_indices_buffer[
|
|
chunk.token_start : chunk.token_end, :topk_tokens
|
|
]
|
|
if chunk.local_total_seq_lens == 0:
|
|
topk_indices.fill_(-1)
|
|
else:
|
|
assert chunk.local_cu_seq_lens is not None
|
|
# Each token's own (per-request, DCP-local) causal length.
|
|
local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
|
|
# Recover each token's owning request from its row-start
|
|
# tag. Ties (from zero-length requests) are harmless: those
|
|
# tokens have local length 0 and never dereference
|
|
# page_table.
|
|
req_idx = (
|
|
torch.searchsorted(
|
|
chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True
|
|
)
|
|
- 1
|
|
)
|
|
page_table = chunk.block_table[req_idx]
|
|
# The true max over this chunk's own rows, NOT
|
|
# chunk.max_local_total_seq_lens (that field sums every
|
|
# request's length in the chunk, bounding the old
|
|
# flat-gather buffer this paged path no longer allocates).
|
|
max_seq_len = int(local_seq_lens.max().item())
|
|
|
|
logits = ops.fp8_paged_mqa_logits_cpu(
|
|
q_slice,
|
|
kv_view,
|
|
weights[chunk.token_start : chunk.token_end],
|
|
local_seq_lens,
|
|
page_table,
|
|
block_size,
|
|
max_seq_len,
|
|
)
|
|
out_page_scratch = torch.empty(
|
|
(q_slice.shape[0], topk_tokens),
|
|
dtype=torch.int32,
|
|
device=kv_cache.device,
|
|
)
|
|
ops.topk_transform_512_cpu(
|
|
logits,
|
|
local_seq_lens,
|
|
page_table,
|
|
out_page_scratch,
|
|
block_size,
|
|
topk_indices,
|
|
)
|
|
|
|
if has_decode:
|
|
decode_metadata = attn_metadata_narrowed.decode
|
|
assert decode_metadata is not None
|
|
assert not decode_metadata.requires_padding, (
|
|
"CPU sparse indexer decode path does not support speculative "
|
|
"decoding (native MTP) yet."
|
|
)
|
|
batch_size = decode_metadata.decode_lens.shape[0]
|
|
if batch_size > 0:
|
|
# No native MTP on CPU (asserted above) => exactly one
|
|
# query token per decode request, so the flat slice below
|
|
# is already the batch-major layout
|
|
# fp8_paged_mqa_logits_cpu wants.
|
|
assert num_decode_tokens == batch_size, (
|
|
"CPU sparse indexer decode path expects exactly one query "
|
|
"token per decode request."
|
|
)
|
|
assert topk_tokens == 512, (
|
|
"topk_transform_512_cpu only supports index_topk == 512."
|
|
)
|
|
seq_lens = decode_metadata.seq_lens
|
|
seq_lens = (
|
|
seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
|
|
)
|
|
block_table = decode_metadata.block_table[:batch_size]
|
|
block_size = kv_cache.shape[1]
|
|
kv_view = kv_cache.view(kv_cache.shape[0], -1)
|
|
|
|
logits = ops.fp8_paged_mqa_logits_cpu(
|
|
q_quant[:num_decode_tokens],
|
|
kv_view,
|
|
weights[:num_decode_tokens],
|
|
seq_lens,
|
|
block_table,
|
|
block_size,
|
|
attn_metadata_narrowed.max_seq_len,
|
|
)
|
|
|
|
# out_page_indices is a required kernel output but unused:
|
|
# the indexer's topk output must stay local/compressed-
|
|
# context positions (resolved later by
|
|
# DeepseekV4CPUAttention.forward_mqa via
|
|
# map_local_to_global_slots_cpu).
|
|
out_page_scratch = torch.empty(
|
|
(batch_size, topk_tokens),
|
|
dtype=torch.int32,
|
|
device=kv_cache.device,
|
|
)
|
|
ops.topk_transform_512_cpu(
|
|
logits,
|
|
seq_lens,
|
|
block_table,
|
|
out_page_scratch,
|
|
block_size,
|
|
topk_indices_buffer[:num_decode_tokens, :topk_tokens],
|
|
)
|
|
|
|
return topk_indices_buffer
|