diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 3340e43a0..7528158c1 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -17,6 +17,7 @@ fused_qk_rope_concat_and_cache_mla, get_hip_quant, ) +from aiter.ops.triton.utils.device_info import get_num_sms # The segmented (page_size>1) MLA cache kernels only exist in newer aiter # builds. Import them lazily so that the default page_size=1 path keeps working @@ -1627,10 +1628,15 @@ def triton_gather_kv_indices_sparse( assert topk_indices.shape[1] == NUM_TOPK_TOKENS assert NUM_TOPK_TOKENS % BLOCK_N == 0 - # MTP decode can carry metadata tensors padded to a larger query layout - # than the number of rows produced by the current indexer call. Keep all - # per-token inputs aligned to the actual valid intersection before launch; - # otherwise the kernel may read past topk_indices. + # MTP decode can carry metadata tensors padded to a larger query layout than + # the number of rows the indexer actually produced (e.g. real tokens < + # batch_size * max_seqlen_q under cudagraph padding, so token_to_seq_idxs is + # longer than topk_indices). Align every per-token input to the valid + # intersection before launch, otherwise the kernel reads past topk_indices. + # This truncation only drops padding rows, never real requests: the chunked + # indexer launches in sparse_attn_indexer already guarantee every real row + # has fresh top-k content (that is what fixes the con>=256 stale-index drop), + # so no stale rows survive within the valid range. num_tokens = min( token_to_seq_idxs.shape[0], topk_indices.shape[0], @@ -1663,3 +1669,10 @@ def triton_gather_kv_indices_sparse( ti_stride1, ) return out_buf + + +def sparse_indexer_decode_rows_per_launch(wave_per_eu: int = 2) -> int: + # Keep sparse-indexer decode launches within the CU wave capacity. Large + # single launches can leave stale sparse indices in the persistent buffer on + # the GLM MTP path, even though standalone kernel probes may cover all rows. + return max(1, get_num_sms() * wave_per_eu) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 6a363ac8f..9b8478975 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -58,6 +58,7 @@ triton_convert_req_index_to_global_index, triton_convert_req_index_to_global_index_dsa_prefill, triton_gather_kv_indices_sparse, + sparse_indexer_decode_rows_per_launch, ) from atom.model_ops.base_attention import Attention from atom.model_ops.embed_head import ( @@ -1350,32 +1351,69 @@ def sparse_attn_indexer( assert batch_size == context.batch_size num_padded_tokens = batch_size * next_n batch_size, next_n, heads, _ = padded_q_fp8_decode_tokens.shape - logits = torch.empty( - [batch_size * next_n, max_model_len], dtype=torch.float32, device="cuda" - ) - deepgemm_fp8_paged_mqa_logits( - padded_q_fp8_decode_tokens, - kv_cache, - weights[:num_padded_tokens], - logits, - decode_metadata.context_lens, - attn_metadata.block_tables, - max_model_len, - KVBlockSize=runner_block_size, - Preshuffle=True, - ) - num_rows = logits.shape[0] assert topk_tokens == 2048, "top_k_per_row assumes size 2048" topk_indices_decode = topk_indices[:num_decode_tokens, :topk_tokens] - top_k_per_row_decode( - logits, - next_n, - decode_metadata.context_lens, - topk_indices_decode, - num_rows, - logits.stride(0), - logits.stride(1), - ) + # Keep each sparse-indexer decode launch under the row cliff observed at + # large MTP batches; otherwise later rows can retain stale sparse KV + # indices in the persistent buffer. + max_indexer_rows = sparse_indexer_decode_rows_per_launch() + if num_padded_tokens > max_indexer_rows: + chunk_reqs = max(1, max_indexer_rows // next_n) + for req_start in range(0, batch_size, chunk_reqs): + req_end = min(req_start + chunk_reqs, batch_size) + row_start = req_start * next_n + row_end = req_end * next_n + logits = torch.empty( + [row_end - row_start, max_model_len], + dtype=torch.float32, + device="cuda", + ) + deepgemm_fp8_paged_mqa_logits( + padded_q_fp8_decode_tokens[req_start:req_end], + kv_cache, + weights[row_start:row_end], + logits, + decode_metadata.context_lens[req_start:req_end], + attn_metadata.block_tables[req_start:req_end], + max_model_len, + KVBlockSize=runner_block_size, + Preshuffle=True, + ) + top_k_per_row_decode( + logits, + next_n, + decode_metadata.context_lens[req_start:req_end], + topk_indices_decode[row_start:row_end], + row_end - row_start, + logits.stride(0), + logits.stride(1), + ) + else: + logits = torch.empty( + [num_padded_tokens, max_model_len], + dtype=torch.float32, + device="cuda", + ) + deepgemm_fp8_paged_mqa_logits( + padded_q_fp8_decode_tokens, + kv_cache, + weights[:num_padded_tokens], + logits, + decode_metadata.context_lens, + attn_metadata.block_tables, + max_model_len, + KVBlockSize=runner_block_size, + Preshuffle=True, + ) + top_k_per_row_decode( + logits, + next_n, + decode_metadata.context_lens, + topk_indices_decode, + num_padded_tokens, + logits.stride(0), + logits.stride(1), + ) if attn_metadata.max_seqlen_q > 1: triton_gather_kv_indices_sparse( attn_metadata.sparse_kv_indptr, diff --git a/atom/plugin/vllm/attention/layer_sparse_mla.py b/atom/plugin/vllm/attention/layer_sparse_mla.py index 704c27e4b..4912756d1 100644 --- a/atom/plugin/vllm/attention/layer_sparse_mla.py +++ b/atom/plugin/vllm/attention/layer_sparse_mla.py @@ -24,6 +24,7 @@ from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits +from atom.model_ops.attention_mla import sparse_indexer_decode_rows_per_launch from atom.plugin.prepare import is_vllm from atom.utils import envs from atom.utils.custom_register import direct_register_custom_op @@ -458,35 +459,81 @@ def sparse_attn_indexer_plugin_mode( next_n = padded_q_fp8_decode_tokens.shape[1] assert batch_size == decode_metadata.seq_lens.shape[0] num_padded_tokens = batch_size * next_n - logits = torch.empty( - [batch_size * next_n, max_model_len], dtype=torch.float32, device="cuda" - ) - deepgemm_fp8_paged_mqa_logits( - padded_q_fp8_decode_tokens, - kv_cache, - weights[:num_padded_tokens], - logits, - decode_metadata.seq_lens, - decode_metadata.block_table, - max_model_len, - ChunkK=256, - KVBlockSize=kv_block_size, - Preshuffle=preshuffle_cache, - WavePerEU=2, - ) - - num_rows = logits.shape[0] assert topk_tokens == 2048, "top_k_per_row assumes size 2048" - topk_indices_decode = topk_indices[:num_decode_tokens, :topk_tokens] - top_k_per_row_decode( - logits, - next_n, - decode_metadata.seq_lens, - topk_indices_decode, - num_rows, - logits.stride(0), - logits.stride(1), - ) + if decode_metadata.requires_padding: + padded_weights_decode = pack_seq_triton( + weights[:num_decode_tokens], decode_lens + ).reshape(num_padded_tokens, *weights.shape[1:]) + topk_indices_decode = torch.empty( + [num_padded_tokens, topk_tokens], + dtype=torch.int32, + device=topk_indices.device, + ) + else: + padded_weights_decode = weights[:num_padded_tokens] + topk_indices_decode = topk_indices[:num_decode_tokens, :topk_tokens] + max_indexer_rows = sparse_indexer_decode_rows_per_launch() + if num_padded_tokens > max_indexer_rows: + chunk_reqs = max(1, max_indexer_rows // next_n) + for req_start in range(0, batch_size, chunk_reqs): + req_end = min(req_start + chunk_reqs, batch_size) + row_start = req_start * next_n + row_end = req_end * next_n + logits = torch.empty( + [row_end - row_start, max_model_len], + dtype=torch.float32, + device="cuda", + ) + deepgemm_fp8_paged_mqa_logits( + padded_q_fp8_decode_tokens[req_start:req_end], + kv_cache, + padded_weights_decode[row_start:row_end], + logits, + decode_metadata.seq_lens[req_start:req_end], + decode_metadata.block_table[req_start:req_end], + max_model_len, + ChunkK=256, + KVBlockSize=kv_block_size, + Preshuffle=preshuffle_cache, + WavePerEU=2, + ) + top_k_per_row_decode( + logits, + next_n, + decode_metadata.seq_lens[req_start:req_end], + topk_indices_decode[row_start:row_end], + row_end - row_start, + logits.stride(0), + logits.stride(1), + ) + else: + logits = torch.empty( + [num_padded_tokens, max_model_len], + dtype=torch.float32, + device="cuda", + ) + deepgemm_fp8_paged_mqa_logits( + padded_q_fp8_decode_tokens, + kv_cache, + padded_weights_decode, + logits, + decode_metadata.seq_lens, + decode_metadata.block_table, + max_model_len, + ChunkK=256, + KVBlockSize=kv_block_size, + Preshuffle=preshuffle_cache, + WavePerEU=2, + ) + top_k_per_row_decode( + logits, + next_n, + decode_metadata.seq_lens, + topk_indices_decode, + num_padded_tokens, + logits.stride(0), + logits.stride(1), + ) if decode_metadata.requires_padding: # if padded, we need to unpack