sm12x: cap C128A metadata kernel loop at effective_topk (no shape changes)

Cudagraph-safe retry of suggestion #2 from PR vllm-project#41834 comment 4450901180.
Previous attempt (e34daef, reverted) also exposed
``c128a_*_effective_topk`` on the metadata and truncated the buffer
slice inside ``deepseek_v4_attention``; that truncation baked the
shape into the captured forward launch, breaking replay when
``effective_topk`` shifted between capture and replay.

This version only touches the metadata builder (which already runs
*outside* the captured forward), so per-call ``effective_topk``
variation is fine:

1. Pre-fill ``global_decode_buffer[:num_decode_tokens]`` and
   ``prefill_buffer[:num_prefill_tokens]`` with ``-1`` before launch.
2. Compute ``effective_topk_arg = cdiv(max num_compressed across
   in-flight tokens, BLOCK_SIZE) * BLOCK_SIZE``, capped at
   ``max_compressed_tokens``.
3. Kernel inner loop uses ``effective_topk`` (was
   ``max_compressed_tokens``); store mask uses the same.

The buffer entries the kernel skips
(``[effective_topk, max_compressed_tokens)``) stay at ``-1`` from the
pre-fill, so downstream sparse MLA accumulate kernels (which still
iterate the full ``max_compressed_tokens`` width inside the cudagraph)
see only ``-1`` sentinels in the tail and short-circuit them via
``kv_index >= 0`` / ``candidate < valid_len`` checks. No tensor shape
changes inside the captured forward → cudagraph capture/replay
remains correct.

Savings here are limited to the metadata kernel itself; the accumulate
kernels' iteration count is unchanged (their loop bound is the captured
``num_candidates`` shape value, which we deliberately do not narrow).
Bench at long ``max_model_len`` will confirm whether this is enough to
recover a meaningful chunk of the ~27 % TPOT regression observed at
``max_model_len=131072`` vs ``8192``.

Signed-off-by: jasl <jasl9187@hotmail.com>

Author: jasl

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -1082,6 +1082,41 @@ def build_c128a_topk_metadata(
     if num_tokens == 0:
         return global_decode, decode_lens, prefill_local
 
+    BLOCK_SIZE = 1024
+
+    # Compute the smallest BLOCK_SIZE-aligned width that covers every
+    # in-flight token's num_compressed. When ``max_model_len`` is much
+    # larger than the actual prompts (e.g. 1M cap with 2K inputs) the
+    # original ``range(0, max_compressed_tokens, BLOCK_SIZE)`` iterated
+    # over a tail that always wrote ``-1`` for shorter contexts. Capping
+    # the inner loop at ``effective_topk`` cuts those dead iterations.
+    #
+    # This builder runs OUTSIDE the CUDA-graph-captured forward pass, so
+    # ``effective_topk`` may vary freely per call. To keep downstream
+    # (which reads the full ``max_compressed_tokens`` buffer width inside
+    # the captured forward) correct, we pre-fill the active slice with
+    # ``-1`` so the kernel-skipped tail is treated as "invalid" by the
+    # sentinel checks in the sparse MLA accumulate kernels.
+    if num_decode_tokens > 0:
+        global_decode_buffer[:num_decode_tokens].fill_(-1)
+        decode_lens_buffer[:num_decode_tokens].zero_()
+    if num_prefill_tokens > 0:
+        prefill_buffer[:num_prefill_tokens].fill_(-1)
+
+    max_pos = int(positions.max().item())
+    max_num_compressed = min(
+        max((max_pos + 1) // compress_ratio, 0),
+        max_compressed_tokens,
+    )
+    effective_topk_arg = min(
+        max_compressed_tokens,
+        cdiv(max_num_compressed, BLOCK_SIZE) * BLOCK_SIZE,
+    )
+    if effective_topk_arg == 0:
+        # Nothing to write; the fill_(-1) above already produced the
+        # correct "all-invalid" buffer state.
+        return global_decode, decode_lens, prefill_local
+
     _build_c128a_topk_metadata_kernel[(num_tokens,)](
         global_decode_buffer,
         global_decode_buffer.stride(0),
@@ -1090,14 +1125,14 @@ def build_c128a_topk_metadata(
         prefill_buffer.stride(0),
         positions,
         compress_ratio,
-        max_compressed_tokens,
+        effective_topk_arg,
         num_decode_tokens,
         token_to_req_indices,
         block_table,
         block_table.stride(0),
         block_size,
         slot_mapping,
-        BLOCK_SIZE=1024,
+        BLOCK_SIZE=BLOCK_SIZE,
     )
     return global_decode, decode_lens, prefill_local
 
@@ -1114,7 +1149,7 @@ def _build_c128a_topk_metadata_kernel(
     # Inputs
     positions_ptr,
     compress_ratio,
-    max_compressed_tokens,
+    effective_topk,
     num_decode_tokens,
     token_to_req_indices_ptr,
     block_table_ptr,
@@ -1123,20 +1158,26 @@ def _build_c128a_topk_metadata_kernel(
     slot_mapping_ptr,
     BLOCK_SIZE: tl.constexpr,
 ):
+    # ``effective_topk`` is the BLOCK_SIZE-aligned cap that covers every
+    # in-flight token's ``num_compressed`` (computed by the Python
+    # builder). The caller pre-fills the active buffer slice with ``-1``
+    # so entries in ``[effective_topk, max_compressed_tokens)`` remain
+    # ``-1`` (the sentinel the downstream sparse MLA accumulate kernels
+    # use to skip invalid candidates).
     token_idx = tl.program_id(0)
     position = tl.load(positions_ptr + token_idx)
     num_compressed = (position + 1) // compress_ratio
-    num_compressed = tl.minimum(num_compressed, max_compressed_tokens)
+    num_compressed = tl.minimum(num_compressed, effective_topk)
     is_decode = token_idx < num_decode_tokens
 
     if is_decode:
         # --- Decode: block-table lookup → global slot ids + count ---
         is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0
         req_idx = tl.load(token_to_req_indices_ptr + token_idx)
         count = tl.zeros((), dtype=tl.int32)
-        for i in range(0, max_compressed_tokens, BLOCK_SIZE):
+        for i in range(0, effective_topk, BLOCK_SIZE):
             offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < max_compressed_tokens
+            mask = offset < effective_topk
             is_valid = offset < num_compressed
 
             block_indices = offset // block_size
@@ -1161,9 +1202,9 @@ def _build_c128a_topk_metadata_kernel(
     else:
         # --- Prefill: write local indices ---
         pfx_idx = token_idx - num_decode_tokens
-        for i in range(0, max_compressed_tokens, BLOCK_SIZE):
+        for i in range(0, effective_topk, BLOCK_SIZE):
             offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < max_compressed_tokens
+            mask = offset < effective_topk
             tl.store(
                 prefill_local_ptr + pfx_idx * prefill_local_stride + offset,
                 tl.where(offset < num_compressed, offset, -1),