753 lines
31 KiB
Python
753 lines
31 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Custom Sparse Attention Indexer layers."""
|
|
|
|
import torch
|
|
|
|
import vllm.envs as envs
|
|
from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
|
from vllm.config import get_current_vllm_config_or_none
|
|
from vllm.forward_context import get_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.custom_op import CustomOp
|
|
from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
|
|
from vllm.models.glm5next.common.sparse_indexer import (
|
|
RADIX_TOPK_WORKSPACE_SIZE,
|
|
_build_decode_scatter_indices,
|
|
_decode_topk_seq_lens,
|
|
_fill_causal_indices,
|
|
_fill_short_decode_causal_indices,
|
|
_gather_workspace_shapes,
|
|
_scatter_decode_tokens_by_request,
|
|
kv_cache_as_quant_view,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.torch_utils import (
|
|
LayerNameType,
|
|
_encode_layer_name,
|
|
_resolve_layer_name,
|
|
)
|
|
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.worker.workspace import current_workspace_manager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# kpool write helper: form pools from the current token batch and compress them
|
|
# into the index K cache via the fused Triton kernel.
|
|
|
|
|
|
def _kpool_compress_insert(
|
|
k: torch.Tensor,
|
|
gate_score: torch.Tensor,
|
|
ape: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
kpool: int,
|
|
head_dim: int,
|
|
round_scale: bool,
|
|
) -> None:
|
|
"""Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots.
|
|
|
|
``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec):
|
|
only the *last* token of each complete pool carries a valid (>=0) slot;
|
|
intra-pool tokens are -1. Every position is treated as a pool-completion
|
|
candidate and non-completions are masked off inside the kernel. Compacting
|
|
the valid rows first costs two device syncs on the eager prefill path and
|
|
buys nothing numerically. Assumes pool-aligned chunk starts.
|
|
"""
|
|
n = slot_mapping.shape[0]
|
|
# No pool can complete in a batch smaller than one pool; also keeps the
|
|
# clamped gather indices below in bounds.
|
|
if n < kpool:
|
|
return
|
|
pos = torch.arange(n, device=k.device)
|
|
valid = slot_mapping >= 0
|
|
# Drop pools whose start falls before the batch (leading padding); their
|
|
# gate/k data is undefined anyway.
|
|
write_mask = valid & (pos >= kpool - 1)
|
|
offs = torch.arange(kpool, device=k.device)
|
|
idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :]
|
|
kpool_ops.kpool_compress_and_write_cache(
|
|
kv_cache,
|
|
k[idx], # [n, kpool, head_dim]
|
|
gate_score[idx],
|
|
ape,
|
|
slot_mapping.to(torch.int64),
|
|
pool_size=kpool,
|
|
head_dim=head_dim,
|
|
write_mask=write_mask,
|
|
round_scale=round_scale,
|
|
write_cache=True,
|
|
return_compressed=False,
|
|
)
|
|
|
|
|
|
@eager_break_during_capture
|
|
def sparse_attn_indexer_kpool(
|
|
hidden_states: torch.Tensor,
|
|
k_cache_prefix: LayerNameType,
|
|
kv_cache: torch.Tensor,
|
|
q_quant: torch.Tensor,
|
|
q_scale: torch.Tensor | None,
|
|
k: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
quant_block_size: int,
|
|
scale_fmt: str | None,
|
|
topk_tokens: int,
|
|
head_dim: int,
|
|
max_pool_len: int,
|
|
total_seq_lens: int,
|
|
topk_indices_buffer: torch.Tensor,
|
|
skip_k_cache_insert: bool,
|
|
use_fp4_cache: bool = False,
|
|
# kpool params (Plan-A: gate is consumed at write time and read back at
|
|
# topk time to softmax-weight the pool).
|
|
gate_score: torch.Tensor | None = None,
|
|
compress_ape: torch.Tensor | None = None,
|
|
index_kpool: int = 1,
|
|
positions: torch.Tensor | None = None,
|
|
# Paged tail cache (in-progress pool's raw K + gate score), replacing the
|
|
# transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix]
|
|
# for the tail group's token-granular slot_mapping. None on the dummy/profiling
|
|
# path and when the tail cache is disabled.
|
|
tail_kv_cache: torch.Tensor | None = None,
|
|
tail_prefix: str | None = None,
|
|
) -> torch.Tensor:
|
|
# careful! this will be None in dummy run
|
|
attn_metadata = get_forward_context().attn_metadata
|
|
fp8_dtype = current_platform.fp8_dtype()
|
|
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
|
|
|
|
# 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
|
|
)
|
|
current_workspace_manager().get_simultaneous(
|
|
values_spec,
|
|
scales_spec,
|
|
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
|
|
)
|
|
|
|
# Reserve profiler-visible memory for the worst-case decode logits,
|
|
# whose shape is [B * next_n, max_pool_len]. This profiling branch
|
|
# returns before invoking the logits kernel itself.
|
|
cfg = get_current_vllm_config_or_none()
|
|
worst_decode_tokens = 0
|
|
if cfg is not None:
|
|
sched = cfg.scheduler_config
|
|
num_spec = (
|
|
cfg.speculative_config.num_speculative_tokens
|
|
if cfg.speculative_config is not None
|
|
else 0
|
|
)
|
|
worst_decode_tokens = min(
|
|
sched.max_num_seqs * (num_spec + 1),
|
|
sched.max_num_batched_tokens,
|
|
)
|
|
# float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
|
|
decode_logits_elems = worst_decode_tokens * max_pool_len * 4
|
|
prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
|
|
max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
|
|
_ = torch.empty(
|
|
max_logits_elems, dtype=torch.uint8, device=hidden_states.device
|
|
)
|
|
|
|
return topk_indices_buffer
|
|
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.
|
|
num_tokens = slot_mapping.shape[0]
|
|
if k is not None:
|
|
k = k[:num_tokens]
|
|
|
|
if not skip_k_cache_insert:
|
|
assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
|
|
if index_kpool > 1 and gate_score is not None and compress_ape is not None:
|
|
# kpool prefill write: pool kpool consecutive prefill tokens via
|
|
# softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots.
|
|
# Decode tokens (the first num_decode_tokens in the batch) cannot be
|
|
# pooled here — their pool's earlier tokens are not in this batch —
|
|
# so they are deferred to the tail-buffer kernel in has_decode.
|
|
# compress_ratio == index_kpool makes slot_mapping pool-granular.
|
|
n_prefill = num_tokens - num_decode_tokens
|
|
if n_prefill > 0:
|
|
# decode tokens are batched first; prefill tokens follow.
|
|
prefill_slice = slice(num_decode_tokens, num_tokens)
|
|
_kpool_compress_insert(
|
|
k[prefill_slice],
|
|
gate_score[prefill_slice],
|
|
compress_ape,
|
|
kv_cache,
|
|
slot_mapping[prefill_slice],
|
|
index_kpool,
|
|
head_dim,
|
|
round_scale=(scale_fmt is not None),
|
|
)
|
|
# Persist each request's incomplete prefill pool so decode can
|
|
# finish it, including after PD transfer. Tail slots use
|
|
# ``pos % kpool`` within the request's tail block. Processing
|
|
# only the batch's trailing tokens would miss all but the last
|
|
# request in a multi-request prefill.
|
|
if tail_kv_cache is not None and tail_prefix is not None:
|
|
tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix))
|
|
if tail_meta is not None:
|
|
assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
|
|
kpool_ops.kpool_seed_tail_cache(
|
|
tail_kv_cache,
|
|
k[prefill_slice],
|
|
gate_score[prefill_slice],
|
|
tail_meta.slot_mapping[prefill_slice],
|
|
index_kpool,
|
|
head_dim,
|
|
)
|
|
else:
|
|
# standard: per-token fp8 quant + scatter (all tokens).
|
|
assert scale_fmt is not None
|
|
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
|
indexer_k_quant_and_cache_triton,
|
|
)
|
|
|
|
indexer_k_quant_and_cache_triton(
|
|
k,
|
|
kv_cache,
|
|
slot_mapping,
|
|
quant_block_size,
|
|
scale_fmt,
|
|
)
|
|
|
|
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
|
if has_prefill:
|
|
prefill_metadata = attn_metadata_narrowed.prefill
|
|
assert prefill_metadata is not None
|
|
|
|
# Short sequences select every pool, so skip sparse scoring and fill
|
|
# the top-k buffer with all causal token indices. The index-K cache was
|
|
# already written above.
|
|
n_prefill_sf = num_tokens - num_decode_tokens
|
|
# Host-side short-prefill predicate: max_prefill_seq_len is computed
|
|
# in the metadata builder (exact for prefill rows) and equals
|
|
# positions[prefill_slice].max() + 1, so this replaces a
|
|
# positions.max().item() device sync per layer. -1 (unknown metadata)
|
|
# falls back to the device-side check.
|
|
if prefill_metadata.max_prefill_seq_len >= 0:
|
|
short_prefill = (
|
|
n_prefill_sf > 0
|
|
and positions is not None
|
|
and prefill_metadata.max_prefill_seq_len <= topk_tokens
|
|
)
|
|
else:
|
|
short_prefill = (
|
|
n_prefill_sf > 0
|
|
and positions is not None
|
|
and int(positions[num_decode_tokens:num_tokens].max().item()) + 1
|
|
<= topk_tokens
|
|
)
|
|
if short_prefill:
|
|
# short_prefill is only True when positions is not None (above),
|
|
# but narrow explicitly for the indexer below.
|
|
assert positions is not None
|
|
_pos = positions[num_decode_tokens:num_tokens].to(torch.int32)
|
|
_buf = topk_indices_buffer[num_decode_tokens:num_tokens]
|
|
_fill_causal_indices(_buf, _pos)
|
|
|
|
# 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
|
|
)
|
|
k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
|
|
values_spec,
|
|
scales_spec,
|
|
)
|
|
for chunk in prefill_metadata.chunks if not short_prefill else ():
|
|
k_quant = k_quant_full[: chunk.total_seq_lens]
|
|
k_scale = k_scale_full[: chunk.total_seq_lens]
|
|
|
|
if not chunk.skip_kv_gather:
|
|
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
|
cp_gather_indexer_k_quant_cache_triton,
|
|
)
|
|
|
|
cp_gather_indexer_k_quant_cache_triton(
|
|
kv_cache,
|
|
k_quant,
|
|
k_scale,
|
|
chunk.block_table,
|
|
chunk.cu_seq_lens,
|
|
token_to_seq=chunk.token_to_seq,
|
|
)
|
|
|
|
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
|
|
)
|
|
# 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)
|
|
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
|
rocm_fp8_mqa_logits,
|
|
)
|
|
|
|
assert q_scale_slice is None
|
|
logits = rocm_fp8_mqa_logits(
|
|
q_slice_cast,
|
|
(k_quant_cast, k_scale_cast),
|
|
weights[chunk.token_start : chunk.token_end],
|
|
chunk.cu_seqlen_ks,
|
|
chunk.cu_seqlen_ke,
|
|
)
|
|
num_rows = logits.shape[0]
|
|
|
|
# kpool: logits are pool-granular (compress_ratio == index_kpool),
|
|
# so topk selects pools. We pick topk_tokens // kpool pools then
|
|
# expand each pool back to its kpool constituent tokens.
|
|
select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
|
|
if index_kpool > 1:
|
|
pool_topk = torch.empty(
|
|
(num_rows, select_k), dtype=torch.int32, device=logits.device
|
|
)
|
|
topk_dst = pool_topk
|
|
else:
|
|
topk_dst = topk_indices_buffer[
|
|
chunk.token_start : chunk.token_end, :topk_tokens
|
|
]
|
|
|
|
torch.ops._C.top_k_per_row_prefill(
|
|
logits,
|
|
chunk.cu_seqlen_ks,
|
|
chunk.cu_seqlen_ke,
|
|
topk_dst,
|
|
num_rows,
|
|
logits.stride(0),
|
|
logits.stride(1),
|
|
select_k,
|
|
)
|
|
|
|
if index_kpool > 1:
|
|
pool_ids = pool_topk.to(torch.int64)
|
|
if positions is not None:
|
|
# Fused expand-pools + append-tail into one Triton kernel
|
|
# (replaces ~25 elementwise ops). seq_len is token-granular
|
|
# (pos+1); the kernel derives pool_len internally.
|
|
q_seq = (
|
|
positions[chunk.token_start : chunk.token_end].to(torch.int32)
|
|
+ 1
|
|
)
|
|
expanded = kpool_ops.expand_pools_and_append_tail(
|
|
pool_ids, q_seq, index_kpool
|
|
)
|
|
else:
|
|
valid = pool_ids >= 0
|
|
expanded = kpool_ops.expand_pools_to_tokens(
|
|
pool_ids, valid, topk_tokens, index_kpool
|
|
)
|
|
topk_indices_buffer[
|
|
chunk.token_start : chunk.token_end, : expanded.shape[-1]
|
|
] = expanded
|
|
|
|
if has_decode:
|
|
decode_metadata = attn_metadata_narrowed.decode
|
|
assert decode_metadata is not None
|
|
kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes
|
|
kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
|
|
|
|
# Update the tail before reading logits; completed pools are compressed
|
|
# into the slot supplied by slot_mapping.
|
|
# Spec verification groups tokens by request and preserves position
|
|
# order so each token is stashed before the next completes its pool.
|
|
# Positions must remain token-granular because the kernel derives the
|
|
# pool phase and tail index from ``pos % kpool``.
|
|
if (
|
|
index_kpool > 1
|
|
and gate_score is not None
|
|
and compress_ape is not None
|
|
and positions is not None
|
|
and not skip_k_cache_insert
|
|
):
|
|
num_requests = attn_metadata_narrowed.num_decodes
|
|
# Kpool writes must recover the original request grouping after the
|
|
# indexer's flattened decode path. Host metadata avoids a CUDA graph
|
|
# sync when choosing the uniform or padded layout.
|
|
per_req_lens = decode_metadata.per_req_decode_lens
|
|
if per_req_lens is not None:
|
|
use_uniform = (
|
|
decode_metadata.decode_is_uniform
|
|
and num_decode_tokens
|
|
== num_requests * decode_metadata.write_max_decode_len
|
|
)
|
|
group_lens = per_req_lens
|
|
lmax = decode_metadata.write_max_decode_len
|
|
else:
|
|
# Legacy metadata without per-request lens: fall back to the
|
|
# host-side requires_padding flag. Unreached now (per-request
|
|
# lens is always populated for decode), kept defensive.
|
|
use_uniform = not decode_metadata.requires_padding
|
|
group_lens = decode_metadata.decode_lens
|
|
lmax = int(decode_metadata.decode_lens.max().item())
|
|
if not use_uniform:
|
|
# Non-uniform decode_lens (mixed plain-decode + spec-verify, or
|
|
# a variable MTP-verify batch): scatter actual tokens into a
|
|
# padded [B, lmax] layout. int32 tensors can't go through
|
|
# pack_seq_triton (float/uint8 only). The scatter indices are
|
|
# shared by all five scatters below (and the tail slot one).
|
|
scatter_idx = _build_decode_scatter_indices(
|
|
group_lens, num_requests, num_decode_tokens
|
|
)
|
|
dec_k = _scatter_decode_tokens_by_request(
|
|
k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx
|
|
)
|
|
dec_gate = _scatter_decode_tokens_by_request(
|
|
gate_score[:num_decode_tokens],
|
|
0,
|
|
num_requests,
|
|
lmax,
|
|
scatter_idx,
|
|
)
|
|
dec_slot = _scatter_decode_tokens_by_request(
|
|
slot_mapping[:num_decode_tokens],
|
|
-1,
|
|
num_requests,
|
|
lmax,
|
|
scatter_idx,
|
|
)
|
|
dec_pos = _scatter_decode_tokens_by_request(
|
|
positions[:num_decode_tokens].to(torch.int32),
|
|
-1,
|
|
num_requests,
|
|
lmax,
|
|
scatter_idx,
|
|
)
|
|
else:
|
|
next_n = num_decode_tokens // num_requests
|
|
shape2 = (num_requests, next_n)
|
|
dec_k = k[:num_decode_tokens].view(*shape2, head_dim)
|
|
dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim)
|
|
dec_slot = slot_mapping[:num_decode_tokens].view(shape2)
|
|
dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2)
|
|
tail_meta = (
|
|
attn_metadata.get(_resolve_layer_name(tail_prefix))
|
|
if tail_prefix is not None
|
|
else None
|
|
)
|
|
# Paged tail cache replaces the transient _DECODE_TAIL ring. Group
|
|
# the tail group's token-granular slot_mapping per-request, mirroring
|
|
# dec_slot / dec_pos, so the kernel gets each request's current-token
|
|
# tail slot (block * kpool + pos % kpool).
|
|
if tail_meta is not None:
|
|
assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
|
|
if tail_meta is None or tail_kv_cache is None:
|
|
dec_tail_slot = None
|
|
elif not use_uniform:
|
|
dec_tail_slot = _scatter_decode_tokens_by_request(
|
|
tail_meta.slot_mapping[:num_decode_tokens],
|
|
-1,
|
|
num_requests,
|
|
lmax,
|
|
scatter_idx,
|
|
)
|
|
else:
|
|
dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2)
|
|
# The compress kernel writes the raw fp8 cache (not the quant view);
|
|
# pass the underlying kv_cache, not kv_cache_quant_view.
|
|
if dec_tail_slot is not None:
|
|
# Single batched launch over [num_requests, next_n] replaces the
|
|
# per-token sequential loop. The kernel iterates each request's
|
|
# tokens in position order internally, preserving the
|
|
# pool-completion read-after-stash dependency that the loop
|
|
# provided. Inputs are already grouped per request (uniform:
|
|
# view; non-uniform: _scatter_decode_tokens_by_request padded to
|
|
# [B, lmax]) — no per-token .contiguous() copies needed.
|
|
kpool_ops.kpool_decode_update_and_maybe_write_cache_batched(
|
|
kv_cache_raw,
|
|
tail_kv_cache,
|
|
dec_tail_slot,
|
|
dec_k,
|
|
dec_gate,
|
|
compress_ape,
|
|
dec_slot,
|
|
dec_pos,
|
|
index_kpool,
|
|
head_dim,
|
|
round_scale=(scale_fmt is not None),
|
|
)
|
|
if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices(
|
|
topk_indices_buffer,
|
|
positions,
|
|
num_decode_tokens,
|
|
attn_metadata_narrowed.max_seq_len,
|
|
topk_tokens,
|
|
):
|
|
return topk_indices_buffer
|
|
decode_lens = decode_metadata.decode_lens
|
|
if decode_metadata.requires_padding:
|
|
# Padding also covers short chunked prefills classified as decode.
|
|
# MXFP4 uses zero-byte padding so padded slots dequantize to zero.
|
|
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
|
|
padded_weights = pack_seq_triton(
|
|
weights[:num_decode_tokens], decode_lens, pad_value=0
|
|
).reshape(-1, *weights.shape[1:])
|
|
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
|
|
padded_weights = weights[:num_decode_tokens]
|
|
# 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
|
|
)
|
|
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
|
rocm_fp8_paged_mqa_logits,
|
|
)
|
|
|
|
assert padded_q_scale is None
|
|
logits = rocm_fp8_paged_mqa_logits(
|
|
padded_q_quant_cast,
|
|
kv_cache,
|
|
padded_weights[:num_padded_tokens],
|
|
seq_lens,
|
|
decode_metadata.block_table,
|
|
decode_metadata.schedule_metadata,
|
|
max_model_len=max_pool_len,
|
|
)
|
|
num_rows = logits.shape[0]
|
|
# kpool: logits are pool-granular -> select topk_tokens//kpool pools,
|
|
# then expand each pool back to its kpool tokens.
|
|
select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
|
|
if index_kpool > 1:
|
|
pool_topk = torch.empty(
|
|
(num_rows, select_k), dtype=torch.int32, device=logits.device
|
|
)
|
|
topk_dst = pool_topk
|
|
else:
|
|
topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
|
|
|
|
torch.ops._C.top_k_per_row_decode(
|
|
logits,
|
|
next_n,
|
|
seq_lens,
|
|
topk_dst,
|
|
num_rows,
|
|
logits.stride(0),
|
|
logits.stride(1),
|
|
select_k,
|
|
)
|
|
|
|
# Resolve to token-level indices in the output buffer.
|
|
if index_kpool > 1:
|
|
pool_ids = pool_topk.to(torch.int64)
|
|
n = pool_topk.shape[0]
|
|
# Decode seq_lens are pool-granular; recover token lengths from
|
|
# positions using the padded [B, next_n] row layout when needed.
|
|
if positions is not None:
|
|
dec_seq = _decode_topk_seq_lens(
|
|
positions,
|
|
decode_lens,
|
|
num_decode_tokens,
|
|
batch_size,
|
|
next_n,
|
|
decode_metadata.requires_padding,
|
|
)
|
|
else:
|
|
dec_seq = decode_metadata.seq_lens[:n]
|
|
if dec_seq.ndim == 2:
|
|
dec_seq = dec_seq[:, -1]
|
|
dec_seq = dec_seq.to(torch.int32)
|
|
out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool)
|
|
else:
|
|
out = topk_dst
|
|
|
|
if decode_metadata.requires_padding:
|
|
# Drop padded query rows introduced by the next_n padding above.
|
|
out = unpack_seq_triton(
|
|
out.reshape(batch_size, -1, out.shape[-1]), decode_lens
|
|
)
|
|
topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out
|
|
|
|
return topk_indices_buffer
|
|
|
|
|
|
@CustomOp.register("sparse_attn_indexer_kpool")
|
|
class SparseAttnIndexerKpool(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_pool_len: int,
|
|
max_total_seq_len: int,
|
|
topk_indices_buffer: torch.Tensor,
|
|
skip_k_cache_insert: bool = False,
|
|
use_fp4_cache: bool = False,
|
|
tail_cache=None,
|
|
):
|
|
super().__init__()
|
|
self.k_cache = k_cache
|
|
self.tail_cache = tail_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_pool_len = max_pool_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
|
|
|
|
def forward_hip(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
*,
|
|
gate_score: torch.Tensor | None = None,
|
|
compress_ape: torch.Tensor | None = None,
|
|
index_kpool: int = 1,
|
|
positions: torch.Tensor | None = None,
|
|
):
|
|
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"
|
|
)
|
|
if not rocm_aiter_ops.is_enabled():
|
|
raise RuntimeError(
|
|
"Sparse attention indexer ROCm path is only supported on AITER. "
|
|
"Please enable aiter with VLLM_ROCM_USE_AITER=1"
|
|
)
|
|
if index_kpool <= 1:
|
|
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_pool_len,
|
|
self.max_total_seq_len,
|
|
self.topk_indices_buffer,
|
|
skip_k_cache_insert=self.skip_k_cache_insert,
|
|
)
|
|
return sparse_attn_indexer_kpool(
|
|
hidden_states,
|
|
self.k_cache.prefix,
|
|
self.k_cache.kv_cache,
|
|
q_quant,
|
|
None,
|
|
k,
|
|
weights,
|
|
self.quant_block_size,
|
|
self.scale_fmt,
|
|
self.topk_tokens,
|
|
self.head_dim,
|
|
self.max_pool_len,
|
|
self.max_total_seq_len,
|
|
self.topk_indices_buffer,
|
|
self.skip_k_cache_insert,
|
|
self.use_fp4_cache,
|
|
gate_score,
|
|
compress_ape,
|
|
index_kpool,
|
|
positions,
|
|
self.tail_cache.kv_cache if self.tail_cache is not None else None,
|
|
self.tail_cache.prefix if self.tail_cache is not None else None,
|
|
)
|
|
|
|
def forward_native(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
k: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
*,
|
|
gate_score: torch.Tensor | None = None,
|
|
compress_ape: torch.Tensor | None = None,
|
|
index_kpool: int = 1,
|
|
positions: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
return self.forward_hip(
|
|
hidden_states,
|
|
q_quant,
|
|
k,
|
|
weights,
|
|
gate_score=gate_score,
|
|
compress_ape=compress_ape,
|
|
index_kpool=index_kpool,
|
|
positions=positions,
|
|
)
|