class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
"""ROCm sparse MLA attention layer for DeepSeek V4.1."""
backend_cls = DeepseekV41ROCMAiterMLASparseBackend
swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
def _indexer_cls(
self, k_cache: DeepseekV4IndexerCache | None
) -> type[DeepseekV4Indexer]:
if k_cache is not None and k_cache.rocm_mxfp4:
return DeepseekV41RocmMxfp4Indexer
if _gfx942_monokernel():
return DeepseekV41Gfx942Indexer
return super()._indexer_cls(k_cache)
def __init__(self, *args, **kwargs):
vllm_config = args[0] if args else kwargs["vllm_config"]
super().__init__(*args, **kwargs)
# CUDA executes WO_A with a quantized grouped-BMM kernel. ROCm's
# correctness path below dequantizes WO_A once and uses torch.einsum,
# so retain the ordinary MXFP8 linear kernel for post-load processing
# instead of asking for the CUDA-only BMM kernel.
self.wo_a.is_bmm = False
self._has_kv_transfer = vllm_config.kv_transfer_config is not None
# Block scale for the preshuffled weight; None = not preshuffled.
self._wqa_wkv_scale: torch.Tensor | None = None
self._wo_b_scale: torch.Tensor | None = None
self._fused_compressor_weight: torch.Tensor | None
self.register_buffer("_fused_compressor_weight", None, persistent=False)
self._fused_compressor_split_sizes: tuple[int, int] | None = None
# Decode ragged topk metadata is a pure function of the indices its
# index source published, so every consumer of a source rebuilds the
# same thing. The source owns one cache per compress ratio, refreshed
# when it runs; consumers below it read through.
self._topk_ragged_cache: dict[int, _TopkRagged] = {}
self._prefill_topk_ragged_cache: dict[
str, tuple[torch.Tensor, torch.Tensor]
] = {}
self._index_source_prefix: str | None = None
if self.compress_ratio > 0:
assert self.index_source_layer_id is not None
self._index_source_prefix = _replace_layer_index(
self.prefix, self.index_source_layer_id
)
if self._index_source_prefix not in self._static_forward_context:
raise NotImplementedError(
f"Index source {self._index_source_prefix} not found on "
"this rank; PP splits inside a v4.1 index-sharing group "
"are not supported."
)
self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)
if self.compressor is None and self.indexer is None:
# Dense layers have nothing to overlap; keep the base serial path.
self.aux_stream_list = None
@classmethod
def get_padded_num_q_heads(cls, num_heads: int) -> int:
return num_heads
def prepare_attn_preshuffle(self) -> None:
from vllm._aiter_ops import rocm_aiter_ops
if not rocm_aiter_ops.is_enabled():
return
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
_upcast_e8m0_to_fp32,
)
from vllm.model_executor.utils import replace_parameter
def _prep(linear) -> torch.Tensor | None:
w = getattr(linear, "weight", None)
if w is None or w.dim() != 2:
return None
# K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.
if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:
return None
ws = getattr(linear, "weight_scale_inv", None) # per-block scale
if ws is None:
return None
if ws.dtype == torch.float8_e8m0fnu:
ws = _upcast_e8m0_to_fp32(ws).contiguous()
# Shuffle the weight in place (single weight, no unshuffled copy).
replace_parameter(
linear,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
return ws
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
self._wo_b_scale = _prep(self.wo_b)
def prepare_compressor_gemm_fusion(self) -> bool:
# V4.1 derives index keys from the source compressor's emitted latent
# and has no nested ``indexer.compressor``. Keep the projections
# separate and use the shared linear/PyTorch correctness path.
return False
def _bpre_attn_gemm(
self,
weight: torch.Tensor,
scale: torch.Tensor,
x: torch.Tensor,
reduce_tp: bool,
) -> torch.Tensor:
from vllm._aiter_ops import rocm_aiter_ops
x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True)
out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
x_fp8, weight, x_scale, scale, output_dtype=x.dtype
)
if reduce_tp and get_tensor_model_parallel_world_size() > 1:
out = tensor_model_parallel_all_reduce(out)
return out
def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
if self._wqa_wkv_scale is not None and hidden_states.dim() == 2:
return self._bpre_attn_gemm(
self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False
)
return super()._fused_wqa_wkv_gemm(hidden_states)
def _run_parallel_input_projections(
self, hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
if self.indexer is not None and self.compressor is None:
# Serial: qr must sit on the default stream before the q-side
# fork below can consume it (ROCm).
aux_streams = self.aux_stream_list
self.aux_stream_list = None
try:
return super()._run_parallel_input_projections(hidden_states)
finally:
self.aux_stream_list = aux_streams
return super()._run_parallel_input_projections(hidden_states)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
llama_4_scaling: torch.Tensor | None = None,
) -> torch.Tensor:
if (
self.indexer is None
or self._prepare_and_attn_fn == self._prepare_and_attn_eager
or not current_platform.enable_multi_stream_overlap(
self.aux_stream_list, get_forward_context().attn_metadata
)
):
# Sequential fallback: no forks outside the capture region, where
# HIP event sync is unreliable; MRV1 keeps the input prep in its
# wide eager region (#51430).
aux_streams = self.aux_stream_list
self.aux_stream_list = None
try:
return super().forward(positions, hidden_states, llama_4_scaling)
finally:
self.aux_stream_list = aux_streams
attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
if self.compressor is not None:
q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_full(
hidden_states, positions
)
else:
q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_reindex(
hidden_states, positions
)
self._sparse_indexer_and_attn(
hidden_states,
index_q,
index_q_scale,
index_weights,
q,
kv,
positions,
attn_out,
)
return self._o_proj(attn_out, positions)
def _forward_csa2_full(
self, hidden_states: torch.Tensor, positions: torch.Tensor
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
"""CSA layers: one fork replaces the base pipeline's three.
Each fork/join is a HIP event pair per side stream on ROCm, so the
base three-fork pipeline is expensive there; the compressor chain and
the indexer weights projection run fully in parallel with the default
chain. The indexer K write and q-side stay serial after the join:
``wk`` consumes the aux-0 latent and the q-side's wq_b consumes the
default-stream qr.
"""
attn_metadata = get_forward_context().attn_metadata
compressor = self.compressor
indexer = self.indexer
assert compressor is not None and indexer is not None
aux_streams = self.aux_stream_list
assert aux_streams is not None and len(aux_streams) >= 2
def default_chain():
qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
q = self._wq_b_proj(qr, qr_scale).view(
-1, self.n_local_heads, self.head_dim
)
return (
self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata),
qr,
qr_scale,
kv,
)
def compressor_chain():
kv_score = torch.mm(
hidden_states,
compressor.fused_wkv_wgate.weight.T,
out_dtype=torch.float32,
)
latent = compressor(kv_score, positions)
compressor.insert_cache(latent, positions, self.rotary_emb)
return latent
def indexer_weights_chain():
# ReplicatedLinear returns (output, bias); bias is None.
weights, _ = indexer.weights_proj(hidden_states)
return weights
(q, qr, qr_scale, kv), (latent, indexer_weights) = execute_in_parallel(
default_chain,
[compressor_chain, indexer_weights_chain],
self.ln_events[0],
self.ln_events[1:3],
aux_streams[:2],
enable=True,
default_first=True,
)
indexer._produce_k(latent, positions, self.indexer_rotary_emb)
index_q, index_q_scale, index_weights_out = indexer.forward_q(
qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb
)
return q, kv, index_q, index_q_scale, index_weights_out
def _forward_csa2_reindex(
self, hidden_states: torch.Tensor, positions: torch.Tensor
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
"""Indexer-only layers: the q-side overlaps the SWA q path.
The input projections run serial first: the q-side consumes qr, so
it can only overlap the SWA q projection and KV insert.
"""
indexer = self.indexer
assert indexer is not None
attn_metadata = get_forward_context().attn_metadata
# Serial on these layers via the _run_parallel_input_projections
# override: qr must sit on the default stream before the fork.
qr_kv, _, indexer_weights = self._run_parallel_input_projections(hidden_states)
assert indexer_weights is not None
qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
def swa_q_chain():
q = self._wq_b_proj(qr, qr_scale).view(
-1, self.n_local_heads, self.head_dim
)
return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
def indexer_q_chain():
return indexer.forward_q(
qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb
)
aux_streams = self.aux_stream_list
assert aux_streams is not None
q, aux_results = execute_in_parallel(
swa_q_chain,
[indexer_q_chain],
self.ln_events[0],
self.ln_events[1:2],
aux_streams[:1],
enable=True,
default_first=True,
)
index_q, index_q_scale, index_weights_out = aux_results[0]
return q, kv, index_q, index_q_scale, index_weights_out
@functools.cached_property
def _wq_b_uses_aiter_block_scaled(self) -> bool:
"""True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.
Cached: the linear kernels and the aiter env gates are fixed once
the model is built, so this is evaluated at the first forward
only.
The fused norm+quant path is only valid if the quant and GEMM it
replaces are exactly the aiter ones; otherwise fall back to the
shared path.
"""
from vllm._aiter_ops import rocm_aiter_ops
from vllm.model_executor.kernels.linear.scaled_mm import (
Fp8BlockScaledMMLinearKernel,
)
if not rocm_aiter_ops.is_linear_fp8_enabled():
return False
linears = [self.wq_b]
if self.indexer is not None:
linears.append(self.indexer.wq_b)
for linear in linears:
kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
return False
return True
def _split_qkv_and_norm(
self, qr_kv: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
"""Fuse q/kv RMSNorm + per-1x128 fp8 q quant into one aiter kernel.
The shared path norms q and kv in one triton kernel and the wq_b
linears then re-read the bf16 qr to quantize it. The aiter kernel
computes both RMSNorms (fp32 accumulate) and the fp8 group quant
in a single pass, writing fp8 qr + group scales directly; both
wq_b GEMMs (attention and indexer) then consume that pair and
skip their own input quant. kv stays bf16: the fused insert
kernel RoPE/quantizes it itself. Falls back to the shared path
when the aiter linear path is not active.
"""
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
if not (
qr.dim() == 2
and qr.shape[0] > 0
and self.q_lora_rank % 128 == 0
and self._wq_b_uses_aiter_block_scaled
):
return super()._split_qkv_and_norm(qr_kv)
from vllm._aiter_ops import rocm_aiter_ops
return rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
q=qr,
q_weight=self.q_norm.weight.data,
q_epsilon=self.eps,
kv=kv,
kv_weight=self.kv_norm.weight.data,
kv_epsilon=self.eps,
group_size=128,
transpose_scale=False,
)
def _alloc_attn_out(
self, num_tokens: int, hidden_states: torch.Tensor
) -> torch.Tensor | QuantizedActivation:
if not _ON_GFX950:
return super()._alloc_attn_out(num_tokens, hidden_states)
# wo_a's MXFP8 input: the decode reduce writes it directly, prefill
# rows are rotated and quantized after their bf16 attention.
width = self.n_local_heads * self.head_dim
data = torch.empty(
(num_tokens, width), dtype=torch.float8_e4m3fn, device=hidden_states.device
)
scale = torch.empty(
(num_tokens, width // 32), dtype=torch.uint8, device=hidden_states.device
)
return QuantizedActivation(
data=data,
scale=scale,
orig_dtype=hidden_states.dtype,
orig_shape=data.shape,
quant_key=kMxfp8Dynamic,
)
def _o_proj(
self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor
) -> torch.Tensor:
if isinstance(attn_out, QuantizedActivation):
z = rocm_mxfp8_wo_a_bmm(
attn_out.data,
attn_out.scale,
self.wo_a,
self.n_local_groups,
self.o_lora_rank,
)
return self._wo_b_after_wo_a(z)
o = attn_out[:, : self.n_local_heads, :]
# ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
z = rocm_inv_rope_einsum(
self.rotary_emb,
o,
positions,
self.rope_head_dim,
self.n_local_groups,
self.o_lora_rank,
self.wo_a,
inverse_rope=False,
)
return self._wo_b_after_wo_a(z.flatten(1))
def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor:
if self._wo_b_scale is not None and zf.dim() == 2:
return self._bpre_attn_gemm(
self.wo_b.weight, self._wo_b_scale, zf, self.wo_b.reduce_results
)
return self.wo_b(zf)
def forward_mqa(
self,
q: torch.Tensor,
kv: torch.Tensor,
positions: torch.Tensor,
output: torch.Tensor | QuantizedActivation,
) -> None:
mxfp8_out = output if isinstance(output, QuantizedActivation) else None
if mxfp8_out is None:
assert isinstance(output, torch.Tensor)
assert output.shape == q.shape, (
f"output buffer shape {output.shape} must match q shape {q.shape}"
)
assert output.dtype == q.dtype, (
f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
)
forward_context = get_forward_context()
attn_metadata = forward_context.attn_metadata
if attn_metadata is None:
# Warmup dummy run: no real metadata. Reserve the workspace
# _forward_prefill would; the dequantize / topk / sparse_fwd kernels
# are skipped this step. The aiter prefill reads the caches in place
# and needs none.
swa_only = self.compress_ratio == 0
N = (
0
if swa_only
else (self.max_model_len + self.compress_ratio - 1)
// self.compress_ratio
)
M = N + self.window_size + self.max_num_batched_tokens
if not self._use_aiter_sparse_mla:
shapes = self._prefill_workspace_shapes(
M, self.max_num_batched_tokens, q
)
current_workspace_manager().get_simultaneous(*shapes)
if mxfp8_out is None:
output.zero_()
else:
mxfp8_out.data.zero_()
mxfp8_out.scale.zero_()
return
assert isinstance(attn_metadata, dict)
rocm_metadata = cast(
DeepseekV4FlashMLAMetadata | None,
attn_metadata.get(self.compressed_cache_prefix)
if self.compressed_cache_prefix is not None
else None,
)
swa_metadata = cast(
DeepseekV41ROCMAiterSparseSWAMetadata | None,
attn_metadata.get(self.swa_cache_layer.prefix),
)
assert swa_metadata is not None
swa_only = self.compress_ratio == 0
self_kv_cache = None if swa_only else self._compressed_kv_cache()
swa_kv_cache = self.swa_cache_layer.kv_cache
num_decodes = swa_metadata.num_decodes
num_prefills = swa_metadata.num_prefills
num_decode_tokens = swa_metadata.num_decode_tokens
if num_prefills > 0:
self._forward_prefill(
q=q[num_decode_tokens:],
positions=positions[num_decode_tokens:],
compressed_k_cache=self_kv_cache,
swa_k_cache=swa_kv_cache,
output=(
output[num_decode_tokens:]
if isinstance(output, torch.Tensor)
else None
),
attn_metadata=rocm_metadata,
swa_metadata=swa_metadata,
mxfp8_out=(
None
if mxfp8_out is None
else (
mxfp8_out.data[num_decode_tokens:],
mxfp8_out.scale[num_decode_tokens:],
)
),
)
if mxfp8_out is not None:
# The reduce epilogue rotates and quantizes every decode row and
# the prefill path already did its own, so nothing is owed here.
if num_decodes > 0:
self._forward_decode(
q=q[:num_decode_tokens],
positions=positions[:num_decode_tokens],
kv_cache=self_kv_cache,
swa_metadata=swa_metadata,
attn_metadata=rocm_metadata,
swa_only=swa_only,
output=None,
output_mxfp8=(
mxfp8_out.data[:num_decode_tokens],
mxfp8_out.scale[:num_decode_tokens],
),
)
return
assert isinstance(output, torch.Tensor)
rotated = 0
if num_decodes > 0:
rotated = self._forward_decode(
q=q[:num_decode_tokens],
positions=positions[:num_decode_tokens],
kv_cache=self_kv_cache,
swa_metadata=swa_metadata,
attn_metadata=rocm_metadata,
swa_only=swa_only,
output=output[:num_decode_tokens],
)
# Rows [rotated, end) were not rotated in a kernel's store and owe the
# standalone pass. Settle that here rather than in _o_proj: the split
# is batch-dependent and _o_proj runs compiled, where such a value
# freezes at its trace-time value.
end = num_decode_tokens if self._use_aiter_sparse_mla else output.shape[0]
rocm_inverse_rope_rows_(
output[rotated:end, : self.n_local_heads, :],
positions[rotated:end],
self.rotary_emb.cos_sin_cache,
self.rope_head_dim,
)
def _decode_topk_ragged(
self,
swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
num_decodes: int,
num_decode_tokens: int,
) -> _TopkRagged:
"""Ragged form of the topk indices this layer's index source published.
The packing depends only on the shared ``topk_indices_buffer`` and
per-step metadata, apart from the layer's own compress ratio, so the
source memoizes one result per ratio for all the layers below it.
"""
assert attn_metadata is not None
assert swa_metadata.is_valid_token is not None
assert self.topk_indices_buffer is not None
assert self._index_source_prefix is not None
source = self._static_forward_context[self._index_source_prefix]
if source is self:
# Fresh indices as of this layer; drop what the last step cached.
source._topk_ragged_cache = {}
cached = source._topk_ragged_cache.get(self.compress_ratio)
if cached is not None:
return cached
built = compute_global_topk_ragged_indices_and_indptr(
self.topk_indices_buffer[:num_decode_tokens],
swa_metadata.token_to_req_indices,
attn_metadata.block_table[:num_decodes],
attn_metadata.block_size // self.compress_ratio,
swa_metadata.is_valid_token[:num_decode_tokens],
)
source._topk_ragged_cache[self.compress_ratio] = built
return built
def _forward_decode(
self,
q: torch.Tensor,
positions: torch.Tensor,
kv_cache: torch.Tensor | None,
swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_only: bool,
output: torch.Tensor | None,
output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> int:
"""Returns how many leading rows the decode epilogue inverse-RoPE'd."""
num_decodes = swa_metadata.num_decodes
num_decode_tokens = swa_metadata.num_decode_tokens
topk_lens = None
topk_ragged_indices = None
topk_ragged_indptr = None
if not swa_only:
(
topk_ragged_indices,
topk_ragged_indptr,
topk_lens,
) = self._decode_topk_ragged(
swa_metadata=swa_metadata,
attn_metadata=attn_metadata,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
)
if self._use_aiter_sparse_mla:
assert swa_metadata.decode_swa_ragged_indices is not None
assert swa_metadata.decode_swa_ragged_indptr is not None
# The kernel's store rotates (and for MXFP8, quantizes) every row.
self._aiter_sparse_mla(
q,
output,
self.swa_cache_layer.kv_cache,
swa_metadata.decode_swa_ragged_indices,
swa_metadata.decode_swa_ragged_indptr,
kv_cache,
topk_ragged_indices,
topk_ragged_indptr,
positions=positions,
output_mxfp8=output_mxfp8,
)
return q.shape[0]
return rocm_sparse_attn_decode(
q=q,
kv_cache=kv_cache,
swa_k_cache=self.swa_cache_layer.kv_cache,
swa_only=swa_only,
topk_indices=None,
topk_lens=topk_lens,
swa_indices=swa_metadata.decode_swa_indices,
swa_lens=swa_metadata.decode_swa_lens,
swa_ragged_indices=swa_metadata.decode_swa_ragged_indices,
swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr,
topk_ragged_indices=topk_ragged_indices,
topk_ragged_indptr=topk_ragged_indptr,
attn_sink=self.attn_sink,
scale=self.scale,
head_dim=self.head_dim,
nope_head_dim=self.nope_head_dim,
rope_head_dim=self.rope_head_dim,
output=output,
output_mxfp8=output_mxfp8,
inv_rope_positions=positions,
inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
self.kv_cache_dtype,
self._has_kv_transfer,
not swa_only and kv_cache is not None,
),
)
def _forward_prefill(
self,
q: torch.Tensor,
positions: torch.Tensor,
compressed_k_cache: torch.Tensor | None,
swa_k_cache: torch.Tensor,
output: torch.Tensor | None,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata,
mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> None:
"""Sparse prefill into bf16 ``output``, or into ``mxfp8_out``.
With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
inverse-RoPE'd and quantized into it in one pass.
"""
if self._use_aiter_sparse_mla:
self._forward_prefill_aiter(
q=q,
positions=positions,
compressed_k_cache=compressed_k_cache,
swa_k_cache=swa_k_cache,
output=output,
attn_metadata=attn_metadata,
swa_metadata=swa_metadata,
mxfp8_out=mxfp8_out,
)
return
swa_only = attn_metadata is None
num_prefill_tokens = swa_metadata.num_prefill_tokens
num_decodes = swa_metadata.num_decodes
num_decode_tokens = swa_metadata.num_decode_tokens
seq_lens = swa_metadata.prefill_seq_lens
gather_lens = swa_metadata.prefill_gather_lens
assert seq_lens is not None
assert gather_lens is not None
# CPU copy, so sizing the gather grid from it does not sync.
seq_lens_cpu = swa_metadata.prefill_seq_lens_cpu
assert seq_lens_cpu is not None
query_start_loc_cpu = swa_metadata.query_start_loc_cpu
query_start_loc = swa_metadata.query_start_loc
assert query_start_loc_cpu is not None
assert query_start_loc is not None
prefill_token_base = query_start_loc_cpu[num_decodes]
# Local indices filled by the index source; SWA-only layers pass
# top_k=0 and never read them.
assert self.topk_indices_buffer is not None
topk_indices = self.topk_indices_buffer[num_decode_tokens:]
topk_indices = topk_indices[:num_prefill_tokens]
if not swa_only:
top_k = topk_indices.shape[-1]
N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
else:
top_k = 0
N = 0
M = N + self.window_size + self.max_num_batched_tokens
chunk_plan = swa_metadata.get_prefill_chunk_plan(
compress_ratio=self.compress_ratio,
prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
has_compressed=not swa_only,
)
assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
workspace = current_workspace_manager().get_simultaneous(
*self._prefill_workspace_shapes(M, num_prefill_tokens, q)
)
kv_rows = workspace[0].view(-1, q.shape[-1])
if mxfp8_out is not None:
output = workspace[1]
assert output is not None
for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
chunk_size = chunk_end - chunk_start
kv = kv_rows[: chunk_size * chunk_M].view(chunk_size, chunk_M, -1)
if not swa_only:
assert attn_metadata is not None
assert compressed_k_cache is not None
block_table = attn_metadata.block_table[num_decodes:]
# compressed_k_cache is OCP on every platform (Triton encoder).
dequantize_and_gather_k_cache(
kv[:chunk_size],
compressed_k_cache,
seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
gather_lens=None,
block_table=block_table[chunk_start:chunk_end],
block_size=attn_metadata.block_size // self.compress_ratio,
offset=0,
use_fnuz=False,
max_gather_len=int(seq_lens_cpu[chunk_start:chunk_end].max())
// self.compress_ratio,
)
swa_block_table = swa_metadata.block_table[num_decodes:]
dequantize_and_gather_k_cache(
kv[:chunk_size],
swa_k_cache,
seq_lens=seq_lens[chunk_start:chunk_end],
gather_lens=gather_lens[chunk_start:chunk_end],
block_table=swa_block_table[chunk_start:chunk_end],
block_size=swa_metadata.block_size,
offset=chunk_N,
use_fnuz=current_platform.is_fp8_fnuz(),
)
query_start = (
query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
)
query_end = (
query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
)
combined_indices, combined_lens = combine_topk_swa_indices(
topk_indices[query_start:query_end],
query_start_loc[
num_decodes + chunk_start : num_decodes + chunk_end + 1
],
seq_lens[chunk_start:chunk_end],
gather_lens[chunk_start:chunk_end],
self.window_size,
self.compress_ratio,
top_k,
chunk_M,
chunk_N,
)
rocm_sparse_attn_prefill(
q=q[query_start:query_end],
kv=kv.view(-1, 1, q.shape[-1]),
indices=combined_indices,
topk_length=combined_lens,
scale=self.scale,
head_dim=self.head_dim,
nope_head_dim=self.nope_head_dim,
rope_head_dim=self.rope_head_dim,
attn_sink=self.attn_sink,
output=output[query_start:query_end],
)
if mxfp8_out is not None:
rocm_inverse_rope_mxfp8_rows(
output,
positions[:num_prefill_tokens],
self.rotary_emb.cos_sin_cache,
self.rope_head_dim,
mxfp8_out[0][:num_prefill_tokens],
mxfp8_out[1][:num_prefill_tokens],
)
def _prefill_workspace_shapes(
self, M: int, num_prefill_tokens: int, q: torch.Tensor
) -> list[tuple[tuple[int, ...], torch.dtype]]:
"""The prefill gather buffer, plus bf16 rows for the MXFP8 output.
One ``get_simultaneous`` call: separate calls alias the same memory.
"""
shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
]
if _ON_GFX950:
shapes.append(
(
(num_prefill_tokens, self.n_local_heads, self.head_dim),
torch.bfloat16,
)
)
return shapes
def _prefill_topk_ragged(
self,
swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata,
attn_metadata: DeepseekV4FlashMLAMetadata,
compressed_k_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prefill counterpart of ``_decode_topk_ragged``: layers that share an
index source and a compressed cache reuse one build per step."""
assert swa_metadata.token_to_req_indices is not None
assert swa_metadata.query_start_loc is not None
assert swa_metadata.seq_lens is not None
assert swa_metadata.is_valid_token is not None
assert self.topk_indices_buffer is not None
assert self._index_source_prefix is not None
assert self.compressed_cache_prefix is not None
source = self._static_forward_context[self._index_source_prefix]
if source is self:
source._prefill_topk_ragged_cache = {}
cached = source._prefill_topk_ragged_cache.get(self.compressed_cache_prefix)
if cached is not None:
return cached
num_decode_tokens = swa_metadata.num_decode_tokens
num_prefill_tokens = swa_metadata.num_prefill_tokens
built = build_prefill_topk_ragged_indices(
self.topk_indices_buffer[
num_decode_tokens : num_decode_tokens + num_prefill_tokens
],
swa_metadata.token_to_req_indices,
swa_metadata.query_start_loc,
swa_metadata.seq_lens,
swa_metadata.is_valid_token,
attn_metadata.block_table,
block_size=attn_metadata.block_size // self.compress_ratio,
compress_ratio=self.compress_ratio,
num_compressed=(self.max_model_len + self.compress_ratio - 1)
// self.compress_ratio,
token_offset=num_decode_tokens,
num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
)
source._prefill_topk_ragged_cache[self.compressed_cache_prefix] = built
return built
def _forward_prefill_aiter(
self,
q: torch.Tensor,
positions: torch.Tensor,
compressed_k_cache: torch.Tensor | None,
swa_k_cache: torch.Tensor,
output: torch.Tensor | None,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata,
mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> None:
"""Prefill straight off the paged caches, with no bf16 gather.
The store applies the inverse RoPE, and quantizes into ``mxfp8_out``.
"""
num_prefill_tokens = swa_metadata.num_prefill_tokens
assert swa_metadata.prefill_swa_ragged_indices is not None
assert swa_metadata.prefill_swa_ragged_indptr is not None
topk_indices = None
topk_indptr = None
if attn_metadata is None:
compressed_k_cache = None
else:
assert compressed_k_cache is not None
topk_indices, topk_indptr = self._prefill_topk_ragged(
swa_metadata, attn_metadata, compressed_k_cache
)
self._aiter_sparse_mla(
q[:num_prefill_tokens],
None if output is None else output[:num_prefill_tokens],
swa_k_cache,
swa_metadata.prefill_swa_ragged_indices,
swa_metadata.prefill_swa_ragged_indptr,
compressed_k_cache,
topk_indices,
topk_indptr,
positions=positions[:num_prefill_tokens],
output_mxfp8=(
None
if mxfp8_out is None
else (
mxfp8_out[0][:num_prefill_tokens],
mxfp8_out[1][:num_prefill_tokens],
)
),
)
def _aiter_sparse_mla(
self,
q: torch.Tensor,
output: torch.Tensor | None,
swa_k_cache: torch.Tensor,
swa_indices: torch.Tensor,
swa_indptr: torch.Tensor,
compressed_k_cache: torch.Tensor | None,
topk_indices: torch.Tensor | None,
topk_indptr: torch.Tensor | None,
positions: torch.Tensor | None = None,
output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> None:
from vllm._aiter_ops import rocm_aiter_ops
assert (output is None) != (output_mxfp8 is None)
# Both fp8 caches are read in place: the SWA window as the main segment,
# the top-k compressed tokens as the extra one. With positions the
# store applies the inverse RoPE; output_mxfp8 replaces output.
rocm_aiter_ops.triton_sparse_mla_fwd(
q,
swa_k_cache,
output,
self.scale,
swa_indptr,
swa_indices,
attn_sink=self.attn_sink[: q.shape[1]],
extra_kv_buffer=compressed_k_cache,
extra_kv_indptr=topk_indptr,
extra_kv_indices=topk_indices,
inv_rope_positions=positions,
inv_rope_cos_sin_cache=(
None if positions is None else self.rotary_emb.cos_sin_cache
),
out_mxfp8=output_mxfp8,
)