@dataclass
class ROCMAiterMLASparseMetadataBuilder(
AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
):
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
def __init__(
self,
kv_cache_spec: AttentionSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
):
self.kv_cache_spec = kv_cache_spec
self.model_config = vllm_config.model_config
self.model_dtype = vllm_config.model_config.dtype
parallel_config = vllm_config.parallel_config
self.device = device
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.vllm_config = vllm_config
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
self.mla_dims = get_mla_dims(self.model_config)
self.topk_tokens = vllm_config.model_config.hf_config.index_topk
# Bounds the KV-split heuristic (see `_sparse_decode_max_split`).
self._num_compute_units = current_platform.num_compute_units()
self.max_model_len_tensor = torch.tensor(
[self.model_config.max_model_len], device=device, dtype=torch.int32
)
# this is ignored by `flash_mla_with_kvcache` if indices not None
self.dummy_block_table = torch.empty(
(1, 1), dtype=torch.int32, device=self.device
)
self.req_id_per_token_buffer = torch.zeros(
(vllm_config.scheduler_config.max_num_batched_tokens,),
dtype=torch.int32,
device=device,
)
self.qo_indptr = torch.arange(
0, max_num_batched_tokens + 1, dtype=torch.int32, device=device
)
self.paged_kv_last_page_len = torch.ones(
max_num_batched_tokens, dtype=torch.int32, device=device
)
# These two needs to be calculated in runtime,
# but we still needs to prepare the buffer
self.paged_kv_indices = torch.zeros(
[max_num_batched_tokens * self.topk_tokens],
dtype=torch.int32,
device=device,
)
self.paged_kv_indptr = torch.zeros(
[max_num_batched_tokens + 1], dtype=torch.int32, device=device
)
# ----- Persistent MLA metadata buffers -----
# The aiter sparse decode kernel supports a "persistent" path that
# uses precomputed work-splitting metadata for better load balancing
# across CUs. Mirrors the approach used in rocm_aiter_mla.py.
#
# In the sparse case each query token is its own "batch" entry in the
# qo_indptr (qo_indptr = [0, 1, 2, ..., num_tokens]) and max_qo_len=1.
# We pad get_mla_metadata_info_v1's batch_size to max_num_batched_tokens
# so the buffers are large enough for any decode shape we might see.
from aiter import dtypes, get_mla_metadata_info_v1
# Aiter sparse MLA also requires num_heads >= 16 (will be padded by
# AiterMLAHelper.get_mla_padded_q in forward).
self._num_attention_heads = max(16, self.num_heads)
q_dtype = self.model_dtype
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
kv_cache_dtype_str = "fp8"
else:
kv_cache_dtype_str = "bf16"
kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
(
(work_meta_data_size, work_meta_data_type),
(work_indptr_size, work_indptr_type),
(work_info_set_size, work_info_set_type),
(reduce_indptr_size, reduce_indptr_type),
(reduce_final_map_size, reduce_final_map_type),
(reduce_partial_map_size, reduce_partial_map_type),
) = get_mla_metadata_info_v1(
max_num_batched_tokens,
1,
self._num_attention_heads,
q_dtype,
kv_dtype,
is_sparse=True,
fast_mode=True,
)
self._mla_work_meta_data = torch.empty(
work_meta_data_size, dtype=work_meta_data_type, device=device
)
self._mla_work_indptr = torch.empty(
work_indptr_size, dtype=work_indptr_type, device=device
)
self._mla_work_info_set = torch.empty(
work_info_set_size, dtype=work_info_set_type, device=device
)
self._mla_reduce_indptr = torch.empty(
reduce_indptr_size, dtype=reduce_indptr_type, device=device
)
self._mla_reduce_final_map = torch.empty(
reduce_final_map_size, dtype=reduce_final_map_type, device=device
)
self._mla_reduce_partial_map = torch.empty(
reduce_partial_map_size,
dtype=reduce_partial_map_type,
device=device,
)
self._prev_req_extent: int = 0
self._prev_indices_extent: int = 0
self._prev_metadata_key: tuple | None = None
def _sparse_decode_max_split(self, max_seq_len: int) -> int:
"""Cap ``max_split_per_batch`` for the aiter sparse-MLA decode reduce.
The reduce only covers the selected tokens per row (``<= topk_tokens``),
so aiter's default (``-1`` => split across every CU) over-fragments it.
Mirror ``triton_mla.py``: aim for a minimum work per split, round to a
power of two, and cap by the CU count. Numerics are unchanged.
"""
effective_len = min(max_seq_len, self.topk_tokens)
min_work_per_split = 128
ideal_splits = triton.next_power_of_2(
max(1, effective_len // min_work_per_split)
)
return min(ideal_splits, self._num_compute_units)
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> ROCMAiterMLASparseMetadata:
num_tokens = common_attn_metadata.num_actual_tokens
starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
seg_lengths = np.diff(starts)
req_id_per_token = np.repeat(
np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
)
# Only re-zero the shrink-tail. paged_kv_indptr is fully rewritten
# by the cumsum below. paged_kv_indices entries past new_indices_extent
# are never read (the attention kernel only touches the ranges
# defined by paged_kv_indptr).
new_req_extent = int(req_id_per_token.shape[0])
new_indices_extent = num_tokens * self.topk_tokens
if self._prev_req_extent > new_req_extent:
self.req_id_per_token_buffer[new_req_extent : self._prev_req_extent].fill_(
0
)
if self._prev_indices_extent > new_indices_extent:
self.paged_kv_indices[new_indices_extent : self._prev_indices_extent].fill_(
0
)
self._prev_req_extent = new_req_extent
self._prev_indices_extent = new_indices_extent
self.req_id_per_token_buffer[:new_req_extent].copy_(
torch.from_numpy(req_id_per_token), non_blocking=True
)
query_lens = (
common_attn_metadata.query_start_loc[1:]
- common_attn_metadata.query_start_loc[:-1]
)
seq_lens = common_attn_metadata.seq_lens
sparse_seqlen = generate_sparse_seqlen_triton(
query_lens,
seq_lens,
common_attn_metadata.query_start_loc,
self.topk_tokens,
num_tokens,
common_attn_metadata.max_query_len,
)
torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1])
self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens])
req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
qo_indptr = self.qo_indptr[: num_tokens + 1]
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens]
paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1]
paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
# ----- Compute persistent MLA metadata -----
# The aiter sparse decode kernel uses qseqlen=1 (each query token is
# treated as its own batch entry), so persistent metadata can always
# be precomputed here. The kernel switches to the persistent
# work-stealing path automatically when work_meta_data is non-None.
# The output is a deterministic function of the per-request query and
# context lengths (both clamped to topk_tokens, past which per-token KV
# length saturates) and num_heads; fingerprint those CPU-side and skip
# the launch when nothing changed.
num_reqs = common_attn_metadata.num_reqs
clamped_seq_lens = np.minimum(
common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
self.topk_tokens,
)
clamped_context_lens = np.minimum(
common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths,
self.topk_tokens,
)
metadata_key = (
num_tokens,
int(common_attn_metadata.max_query_len),
self._num_attention_heads,
clamped_seq_lens.tobytes(),
clamped_context_lens.tobytes(),
seg_lengths.tobytes(),
)
if metadata_key != self._prev_metadata_key:
from aiter import get_mla_metadata_v1
max_split_per_batch = self._sparse_decode_max_split(
int(common_attn_metadata.max_seq_len)
)
get_mla_metadata_v1(
qo_indptr,
paged_kv_indptr,
paged_kv_last_page_len,
self._num_attention_heads,
1,
True,
self._mla_work_meta_data,
self._mla_work_info_set,
self._mla_work_indptr,
self._mla_reduce_indptr,
self._mla_reduce_final_map,
self._mla_reduce_partial_map,
page_size=1,
kv_granularity=16,
max_seqlen_qo=1,
uni_seqlen_qo=1,
fast_mode=True,
max_split_per_batch=max_split_per_batch,
)
# The persistent metadata buffers are read by graph replay. Order
# the async metadata write before the graph-captured decode kernel.
torch.cuda.current_stream(self.device).synchronize()
self._prev_metadata_key = metadata_key
metadata = ROCMAiterMLASparseMetadata(
num_reqs=common_attn_metadata.num_reqs,
max_query_len=common_attn_metadata.max_query_len,
max_seq_len=common_attn_metadata.max_seq_len,
num_actual_tokens=common_attn_metadata.num_actual_tokens,
query_start_loc=common_attn_metadata.query_start_loc,
slot_mapping=common_attn_metadata.slot_mapping,
block_table=common_attn_metadata.block_table_tensor,
req_id_per_token=req_id_per_token,
block_size=self.kv_cache_spec.block_size,
attn_out_dtype=self.model_dtype,
topk_tokens=self.topk_tokens,
qo_indptr=qo_indptr,
paged_kv_last_page_len=paged_kv_last_page_len,
paged_kv_indices=paged_kv_indices,
paged_kv_indptr=paged_kv_indptr,
work_meta_data=self._mla_work_meta_data,
work_indptr=self._mla_work_indptr,
work_info_set=self._mla_work_info_set,
reduce_indptr=self._mla_reduce_indptr,
reduce_final_map=self._mla_reduce_final_map,
reduce_partial_map=self._mla_reduce_partial_map,
)
return metadata