[Refactor][GLM-5.3-Flash] Move sparse_attn_indexer_kpool into the model folder and split AMD/NVIDIA (#55358)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
This commit is contained in:
Jiangyun Zhu
@@ -20,10 +20,10 @@ kernel is documented to replicate.
|
||||
|
||||
import torch
|
||||
|
||||
# Bootstrap the glm5next package before entering the indexer module: its
|
||||
# kpool_compress import runs glm5next/__init__, which pulls model ->
|
||||
# attention -> back into sparse_attn_indexer_kpool (attention.py imports the
|
||||
# class at module scope). Production always enters via attention.py first.
|
||||
# Bootstrap the glm5next package before importing kpool_compress directly: its
|
||||
# import runs glm5next/__init__, which pulls model -> attention -> the indexer
|
||||
# dispatcher (sparse_indexer) -> back into the nvidia/amd sparse_indexer
|
||||
# modules. Production always enters via attention.py first.
|
||||
# isort: off
|
||||
import vllm.models.glm5next # noqa: F401
|
||||
from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402
|
||||
@@ -32,13 +32,17 @@ from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402
|
||||
)
|
||||
# isort: on
|
||||
|
||||
import vllm.model_executor.layers.sparse_attn_indexer_kpool as indexer_mod
|
||||
from vllm.model_executor.layers.sparse_attn_indexer_kpool import (
|
||||
from vllm.models.glm5next.common.sparse_indexer import (
|
||||
_decode_topk_seq_lens,
|
||||
_fill_short_decode_causal_indices,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from vllm.models.glm5next.amd import sparse_indexer as indexer_mod
|
||||
else:
|
||||
from vllm.models.glm5next.nvidia import sparse_indexer as indexer_mod
|
||||
|
||||
KPOOL = 4
|
||||
TOPK_TOKENS = 16
|
||||
SELECT_K = TOPK_TOKENS // KPOOL
|
||||
|
||||
@@ -0,0 +1,729 @@
|
||||
# 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_model_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_model_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_model_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_model_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_model_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_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
|
||||
|
||||
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,
|
||||
):
|
||||
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_model_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_model_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,
|
||||
)
|
||||
@@ -23,7 +23,6 @@ from vllm.model_executor.layers.linear import (
|
||||
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
|
||||
from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool
|
||||
from vllm.model_executor.models.deepseek_v2 import (
|
||||
DeepSeekV2FusedQkvAProjLinear,
|
||||
DeepseekV32IndexerCache,
|
||||
@@ -31,6 +30,7 @@ from vllm.model_executor.models.deepseek_v2 import (
|
||||
)
|
||||
from vllm.model_executor.utils import maybe_disable_graph_partition
|
||||
from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
|
||||
from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
|
||||
from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Shared helpers for the glm5next sparse attention indexer (kpool) layers."""
|
||||
|
||||
import torch
|
||||
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
|
||||
|
||||
# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
def _build_decode_scatter_indices(
|
||||
decode_lens: torch.Tensor,
|
||||
num_requests: int,
|
||||
n: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Per-token (request id, intra-request index) for a non-uniform decode
|
||||
batch, with ``n == decode_lens.sum()`` as a host int (avoids a
|
||||
device sync and keeps both repeat_interleaves sync-free).
|
||||
|
||||
Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
|
||||
building it per call would repeat the same repeat_interleave/cumsum chain
|
||||
up to 5x per layer on the eager decode break.
|
||||
"""
|
||||
device = decode_lens.device
|
||||
dl = decode_lens.to(torch.int64)
|
||||
req_id = torch.repeat_interleave(
|
||||
torch.arange(num_requests, device=device, dtype=torch.int64),
|
||||
dl,
|
||||
output_size=n,
|
||||
)
|
||||
req_starts = torch.cumsum(
|
||||
torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
|
||||
dim=0,
|
||||
)
|
||||
# Broadcast the per-request start offsets to per-token (length n ==
|
||||
# dl.sum()) so each token's intra-request index subtracts its own
|
||||
# request's start.
|
||||
starts = torch.repeat_interleave(req_starts, dl, output_size=n)
|
||||
intra = torch.arange(n, device=device, dtype=torch.int64) - starts
|
||||
return req_id, intra
|
||||
|
||||
|
||||
def _scatter_decode_tokens_by_request(
|
||||
tokens: torch.Tensor,
|
||||
pad_value,
|
||||
num_requests: int,
|
||||
lmax: int,
|
||||
scatter_indices: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
|
||||
layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
|
||||
|
||||
Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
|
||||
slot/pos tensors) — it scatters with the shared per-step indices from
|
||||
``_build_decode_scatter_indices``. Used only for the non-uniform
|
||||
(``requires_padding``) decode batch; uniform batches use a zero-copy
|
||||
reshape.
|
||||
"""
|
||||
req_id, intra = scatter_indices
|
||||
out = torch.full(
|
||||
(num_requests, lmax, *tokens.shape[1:]),
|
||||
pad_value,
|
||||
dtype=tokens.dtype,
|
||||
device=tokens.device,
|
||||
)
|
||||
out[req_id, intra] = tokens
|
||||
return out
|
||||
|
||||
|
||||
def _decode_topk_seq_lens(
|
||||
positions: torch.Tensor,
|
||||
decode_lens: torch.Tensor,
|
||||
num_decode_tokens: int,
|
||||
batch_size: int,
|
||||
next_n: int,
|
||||
requires_padding: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
|
||||
|
||||
``pool_topk`` (and the logits it comes from) follow the padded
|
||||
``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
|
||||
``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
|
||||
``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
|
||||
(the uniform-layout shortcut) misaligns every row after the first
|
||||
non-uniform request and, past the decode region, reads prefill tokens'
|
||||
positions; ``expand_pools_and_append_tail`` then anchors the tail at
|
||||
another request's length, dropping the row's real tail tokens or emitting
|
||||
indices past its sequence (out-of-bounds block-table reads). Padded rows
|
||||
get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
|
||||
"""
|
||||
n = batch_size * next_n
|
||||
if not requires_padding:
|
||||
return positions[:n].to(torch.int32) + 1
|
||||
scatter_idx = _build_decode_scatter_indices(
|
||||
decode_lens, batch_size, num_decode_tokens
|
||||
)
|
||||
padded = _scatter_decode_tokens_by_request(
|
||||
positions[:num_decode_tokens].to(torch.int32),
|
||||
-1,
|
||||
batch_size,
|
||||
next_n,
|
||||
scatter_idx,
|
||||
)
|
||||
return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail
|
||||
|
||||
|
||||
def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
|
||||
causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
|
||||
positions = positions.to(torch.int32)
|
||||
rows[:] = causal_range[None, :]
|
||||
rows[causal_range[None, :] > positions[:, None]] = -1
|
||||
|
||||
|
||||
def _fill_short_decode_causal_indices(
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
positions: torch.Tensor | None,
|
||||
num_decode_tokens: int,
|
||||
max_seq_len: int,
|
||||
topk_tokens: int,
|
||||
) -> bool:
|
||||
"""Fill exact causal rows when sparse decode would select every token."""
|
||||
if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
|
||||
return False
|
||||
_fill_causal_indices(
|
||||
topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
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)
|
||||
+75
-390
@@ -2,30 +2,29 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Custom Sparse Attention Indexer layers."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
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.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.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
|
||||
elif current_platform.is_rocm():
|
||||
from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
|
||||
else:
|
||||
from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
|
||||
|
||||
from vllm.utils.deep_gemm import has_deep_gemm
|
||||
from vllm.utils.torch_utils import (
|
||||
LayerNameType,
|
||||
_encode_layer_name,
|
||||
_resolve_layer_name,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
@@ -36,16 +35,9 @@ from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
if current_platform.is_cuda_alike():
|
||||
from vllm import _custom_ops as ops
|
||||
elif current_platform.is_xpu():
|
||||
from vllm._xpu_ops import xpu_ops
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
|
||||
|
||||
# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
|
||||
# kpool write helper: form pools from the current token batch and compress them
|
||||
# into the index K cache via the fused Triton kernel.
|
||||
|
||||
@@ -96,166 +88,6 @@ def _kpool_compress_insert(
|
||||
)
|
||||
|
||||
|
||||
def _build_decode_scatter_indices(
|
||||
decode_lens: torch.Tensor,
|
||||
num_requests: int,
|
||||
n: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Per-token (request id, intra-request index) for a non-uniform decode
|
||||
batch, with ``n == decode_lens.sum()`` as a host int (avoids a
|
||||
device sync and keeps both repeat_interleaves sync-free).
|
||||
|
||||
Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
|
||||
building it per call would repeat the same repeat_interleave/cumsum chain
|
||||
up to 5x per layer on the eager decode break.
|
||||
"""
|
||||
device = decode_lens.device
|
||||
dl = decode_lens.to(torch.int64)
|
||||
req_id = torch.repeat_interleave(
|
||||
torch.arange(num_requests, device=device, dtype=torch.int64),
|
||||
dl,
|
||||
output_size=n,
|
||||
)
|
||||
req_starts = torch.cumsum(
|
||||
torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
|
||||
dim=0,
|
||||
)
|
||||
# Broadcast the per-request start offsets to per-token (length n ==
|
||||
# dl.sum()) so each token's intra-request index subtracts its own
|
||||
# request's start.
|
||||
starts = torch.repeat_interleave(req_starts, dl, output_size=n)
|
||||
intra = torch.arange(n, device=device, dtype=torch.int64) - starts
|
||||
return req_id, intra
|
||||
|
||||
|
||||
def _scatter_decode_tokens_by_request(
|
||||
tokens: torch.Tensor,
|
||||
pad_value,
|
||||
num_requests: int,
|
||||
lmax: int,
|
||||
scatter_indices: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
|
||||
layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
|
||||
|
||||
Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
|
||||
slot/pos tensors) — it scatters with the shared per-step indices from
|
||||
``_build_decode_scatter_indices``. Used only for the non-uniform
|
||||
(``requires_padding``) decode batch; uniform batches use a zero-copy
|
||||
reshape.
|
||||
"""
|
||||
req_id, intra = scatter_indices
|
||||
out = torch.full(
|
||||
(num_requests, lmax, *tokens.shape[1:]),
|
||||
pad_value,
|
||||
dtype=tokens.dtype,
|
||||
device=tokens.device,
|
||||
)
|
||||
out[req_id, intra] = tokens
|
||||
return out
|
||||
|
||||
|
||||
def _decode_topk_seq_lens(
|
||||
positions: torch.Tensor,
|
||||
decode_lens: torch.Tensor,
|
||||
num_decode_tokens: int,
|
||||
batch_size: int,
|
||||
next_n: int,
|
||||
requires_padding: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
|
||||
|
||||
``pool_topk`` (and the logits it comes from) follow the padded
|
||||
``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
|
||||
``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
|
||||
``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
|
||||
(the uniform-layout shortcut) misaligns every row after the first
|
||||
non-uniform request and, past the decode region, reads prefill tokens'
|
||||
positions; ``expand_pools_and_append_tail`` then anchors the tail at
|
||||
another request's length, dropping the row's real tail tokens or emitting
|
||||
indices past its sequence (out-of-bounds block-table reads). Padded rows
|
||||
get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
|
||||
"""
|
||||
n = batch_size * next_n
|
||||
if not requires_padding:
|
||||
return positions[:n].to(torch.int32) + 1
|
||||
scatter_idx = _build_decode_scatter_indices(
|
||||
decode_lens, batch_size, num_decode_tokens
|
||||
)
|
||||
padded = _scatter_decode_tokens_by_request(
|
||||
positions[:num_decode_tokens].to(torch.int32),
|
||||
-1,
|
||||
batch_size,
|
||||
next_n,
|
||||
scatter_idx,
|
||||
)
|
||||
return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail
|
||||
|
||||
|
||||
def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
|
||||
causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
|
||||
positions = positions.to(torch.int32)
|
||||
rows[:] = causal_range[None, :]
|
||||
rows[causal_range[None, :] > positions[:, None]] = -1
|
||||
|
||||
|
||||
def _fill_short_decode_causal_indices(
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
positions: torch.Tensor | None,
|
||||
num_decode_tokens: int,
|
||||
max_seq_len: int,
|
||||
topk_tokens: int,
|
||||
) -> bool:
|
||||
"""Fill exact causal rows when sparse decode would select every token."""
|
||||
if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
|
||||
return False
|
||||
_fill_causal_indices(
|
||||
topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
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_kpool(
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -393,26 +225,13 @@ def sparse_attn_indexer_kpool(
|
||||
else:
|
||||
# standard: per-token fp8 quant + scatter (all tokens).
|
||||
assert scale_fmt is not None
|
||||
if current_platform.is_rocm():
|
||||
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,
|
||||
)
|
||||
else:
|
||||
ops.indexer_k_quant_and_cache(
|
||||
k,
|
||||
kv_cache,
|
||||
slot_mapping,
|
||||
quant_block_size,
|
||||
scale_fmt,
|
||||
)
|
||||
ops.indexer_k_quant_and_cache(
|
||||
k,
|
||||
kv_cache,
|
||||
slot_mapping,
|
||||
quant_block_size,
|
||||
scale_fmt,
|
||||
)
|
||||
|
||||
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
||||
if has_prefill:
|
||||
@@ -466,27 +285,13 @@ def sparse_attn_indexer_kpool(
|
||||
k_scale = k_scale_full[: chunk.total_seq_lens]
|
||||
|
||||
if not chunk.skip_kv_gather:
|
||||
if current_platform.is_rocm():
|
||||
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,
|
||||
)
|
||||
else:
|
||||
ops.cp_gather_indexer_k_quant_cache(
|
||||
kv_cache,
|
||||
k_quant,
|
||||
k_scale,
|
||||
chunk.block_table,
|
||||
chunk.cu_seq_lens,
|
||||
)
|
||||
ops.cp_gather_indexer_k_quant_cache(
|
||||
kv_cache,
|
||||
k_quant,
|
||||
k_scale,
|
||||
chunk.block_table,
|
||||
chunk.cu_seq_lens,
|
||||
)
|
||||
|
||||
q_slice = q_quant[chunk.token_start : chunk.token_end]
|
||||
q_scale_slice = (
|
||||
@@ -504,30 +309,16 @@ def sparse_attn_indexer_kpool(
|
||||
q_slice_cast = q_slice
|
||||
k_quant_cast = k_quant
|
||||
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
|
||||
if current_platform.is_rocm():
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
rocm_fp8_mqa_logits,
|
||||
)
|
||||
from vllm.utils.deep_gemm import fp8_fp4_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,
|
||||
)
|
||||
else:
|
||||
from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
|
||||
|
||||
logits = fp8_fp4_mqa_logits(
|
||||
(q_slice_cast, q_scale_slice),
|
||||
(k_quant_cast, k_scale_cast),
|
||||
weights[chunk.token_start : chunk.token_end],
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
logits = fp8_fp4_mqa_logits(
|
||||
(q_slice_cast, q_scale_slice),
|
||||
(k_quant_cast, k_scale_cast),
|
||||
weights[chunk.token_start : chunk.token_end],
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
num_rows = logits.shape[0]
|
||||
|
||||
# kpool: logits are pool-granular (compress_ratio == index_kpool),
|
||||
@@ -544,28 +335,16 @@ def sparse_attn_indexer_kpool(
|
||||
chunk.token_start : chunk.token_end, :topk_tokens
|
||||
]
|
||||
|
||||
if current_platform.is_xpu():
|
||||
xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined]
|
||||
logits,
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
topk_dst,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
select_k,
|
||||
)
|
||||
else:
|
||||
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,
|
||||
)
|
||||
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)
|
||||
@@ -765,34 +544,18 @@ def sparse_attn_indexer_kpool(
|
||||
if use_fp4_cache
|
||||
else padded_q_quant_decode_tokens
|
||||
)
|
||||
if current_platform.is_rocm():
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
rocm_fp8_paged_mqa_logits,
|
||||
)
|
||||
from vllm.utils.deep_gemm import fp8_fp4_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_model_len,
|
||||
)
|
||||
else:
|
||||
from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
|
||||
|
||||
logits = fp8_fp4_paged_mqa_logits(
|
||||
(padded_q_quant_cast, padded_q_scale),
|
||||
kv_cache,
|
||||
padded_weights[:num_padded_tokens],
|
||||
seq_lens,
|
||||
decode_metadata.block_table,
|
||||
decode_metadata.schedule_metadata,
|
||||
max_model_len=max_model_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
logits = fp8_fp4_paged_mqa_logits(
|
||||
(padded_q_quant_cast, padded_q_scale),
|
||||
kv_cache,
|
||||
padded_weights[:num_padded_tokens],
|
||||
seq_lens,
|
||||
decode_metadata.block_table,
|
||||
decode_metadata.schedule_metadata,
|
||||
max_model_len=max_model_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
num_rows = logits.shape[0]
|
||||
# kpool: logits are pool-granular -> select topk_tokens//kpool pools,
|
||||
# then expand each pool back to its kpool tokens.
|
||||
@@ -819,28 +582,16 @@ def sparse_attn_indexer_kpool(
|
||||
attn_metadata_narrowed.max_seq_len,
|
||||
)
|
||||
else:
|
||||
if current_platform.is_xpu():
|
||||
xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined]
|
||||
logits,
|
||||
next_n,
|
||||
seq_lens,
|
||||
topk_dst,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
select_k,
|
||||
)
|
||||
else:
|
||||
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,
|
||||
)
|
||||
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:
|
||||
@@ -942,33 +693,16 @@ class SparseAttnIndexerKpool(CustomOp):
|
||||
index_kpool: int = 1,
|
||||
positions: torch.Tensor | None = None,
|
||||
):
|
||||
if current_platform.is_cuda() or current_platform.is_xpu():
|
||||
return self.forward_cuda(
|
||||
hidden_states,
|
||||
q_quant,
|
||||
k,
|
||||
weights,
|
||||
gate_score=gate_score,
|
||||
compress_ape=compress_ape,
|
||||
index_kpool=index_kpool,
|
||||
positions=positions,
|
||||
)
|
||||
elif current_platform.is_rocm():
|
||||
return self.forward_hip(
|
||||
hidden_states,
|
||||
q_quant,
|
||||
k,
|
||||
weights,
|
||||
gate_score=gate_score,
|
||||
compress_ape=compress_ape,
|
||||
index_kpool=index_kpool,
|
||||
positions=positions,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"SparseAttnIndexer native forward is only implemented for "
|
||||
"CUDA, ROCm and XPU platforms."
|
||||
)
|
||||
return self.forward_cuda(
|
||||
hidden_states,
|
||||
q_quant,
|
||||
k,
|
||||
weights,
|
||||
gate_score=gate_score,
|
||||
compress_ape=compress_ape,
|
||||
index_kpool=index_kpool,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
@@ -1012,52 +746,3 @@ class SparseAttnIndexerKpool(CustomOp):
|
||||
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_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 rocm_aiter_ops.is_enabled():
|
||||
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_model_len,
|
||||
self.max_total_seq_len,
|
||||
self.topk_indices_buffer,
|
||||
skip_k_cache_insert=self.skip_k_cache_insert,
|
||||
)
|
||||
return self.forward_cuda(
|
||||
hidden_states,
|
||||
q_quant,
|
||||
k,
|
||||
weights,
|
||||
gate_score=gate_score,
|
||||
compress_ape=compress_ape,
|
||||
index_kpool=index_kpool,
|
||||
positions=positions,
|
||||
)
|
||||
raise RuntimeError(
|
||||
"Sparse attention indexer ROCm path is only supported on AITER. "
|
||||
"Please enable aiter with VLLM_ROCM_USE_AITER=1"
|
||||
)
|
||||
@@ -0,0 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Sparse attention indexer (kpool) — hardware-isolated entry point.
|
||||
|
||||
The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers
|
||||
in ``common/``); this module picks the right one for the current platform.
|
||||
"""
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from .amd.sparse_indexer import (
|
||||
SparseAttnIndexerKpool,
|
||||
)
|
||||
else:
|
||||
from .nvidia.sparse_indexer import (
|
||||
SparseAttnIndexerKpool,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SparseAttnIndexerKpool",
|
||||
]
|
||||
Reference in New Issue
Block a user