[NO REVIEW][atom/atom-vllm] fix MTP accept-rate drop at con>=256 on sparse-MLA (DSA) decode

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)

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,

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