perf(KERNEL-ATTN-MLA-SPARSE): the MLA context gather still launches one block per token and carries the deleted token_to_seq back-map

Row: KERNEL-ATTN-MLA-SPARSE

Found by wave PORTQ-4 (#2680)
re-deriving PORT-NOW entry 132 of 5559679229..e126687a9a, upstream
c76a425278 vllm#51739
"[Kernel] Optimize long-context MLA cache gathers".

This entry was first labelled NOT_APPLICABLE by this wave and a fresh review
found that false. The correction and the reasoning that produced the error are
recorded in .agents/sync/2026-09-03-portq4.md §6.10, because the mistake has a
shape worth publishing, not only a fix.

What is wrong

The commit has two halves and this tree carries the surface for one of them.

The mla_attention.py half passes a token-bounded destination slice to
cp_gather_cache — the fp8 and DCP sibling, which this tree does not implement
and records as out of scope at include/vt/ops.h:5233-5235.

The csrc/libtorch_stable/cache_kernels.cu half rewrites a different kernel:
gather_and_maybe_dequant_cache. That is the one this tree ports, and its header
says so in as many words:

include/vt/ops.h:5225-5227
  // Ported 1:1 from vllm/csrc/libtorch_stable/cache_kernels.cu:992-1064
  // (`vllm::gather_and_maybe_dequant_cache`) + its host wrapper `:1099-1157`

Upstream's hunks (@@ -1120,18 +1158,17 @@ and @@ -1189,14 +1226,21 @@) rename
the kernel to gather_and_maybe_dequant_cache_page, replace the per-token launch
with a page-task grid, and delete the token_to_seq back-map entirely
((void)token_to_seq; in the wrapper).

This tree still carries the pre-fix shape verbatim:

• src/vt/cuda/cuda_mla_prefill.cu:310 — GatherMlaCacheKernel still takes
  const int32_t* __restrict__ token_to_seq, used at :316
  (const int batch_id = token_to_seq[token_id];).
• src/vt/cuda/cuda_mla_prefill.cu:343-344 — const int threads = 64; and
  const unsigned grid = static_cast<unsigned>(num_tokens);, which is exactly
  upstream's deleted constexpr int32_t thread_block_size = 64; dim3 grid(num_tokens);.

Upstream replaces that launch with
ceil_div(num_tokens, block_size) + 2 * num_reqs blocks at 128 or 256 threads,
so one CUDA block handles one page rather than one token. On a long context
the old launch spawns one block per token, most of which stride the same pages.

It is reachable in production, not test-only

vt::GatherMlaCache is called at
include/vllm/model_executor/layers/attention/mla_chunked_context.h:274, inside
BuildMlaChunkedContext, whose non-test callers are:

• src/vllm/model_executor/models/deepseek_v2.cpp:228
• src/vllm/model_executor/models/kimi_linear_device.cpp:2028
• src/vllm/model_executor/models/minicpm3.cpp:119

The header calls it "the FIRST step of every chunked-context iteration".

Why REAL_GAP rather than inert

The semantics are identical either way — this is a pure launch-geometry rewrite,
so no token changes and no gate would go from red to green. It is a REAL_GAP
because the surface is here, reached on the default configuration, and behind
upstream; nothing gates it. It is filed with the same standing as the tranche's
other two performance gaps
(#2684,
#2696).

Size

Port map_gather_page_task, rewrite GatherMlaCacheKernel as a page-task kernel,
change the launch, and drop token_to_seq from the signature and from
MlaChunkedContextMetadata if nothing else reads it. Roughly 120-200 lines
across src/vt/cuda/cuda_mla_prefill.cu, include/vt/ops.h and
include/vllm/model_executor/layers/attention/mla_chunked_context.h, plus the
ported upstream kernel test.

KERNEL-KV-CACHE (.agents/kernel-matrix.md:133) co-owns it as the row that
mirrors upstream's cache_kernels.cu; the anchor lives under the MLA campaign,
which is why KERNEL-ATTN-MLA-SPARSE (.agents/kernel-matrix.md:140) is named
first.

What was not determined

Nothing was executed — no build, no test run, no GPU, no lease. The size of
the win is upstream's claim plus the launch arithmetic, not a measurement here;
token_to_seq is a host-built [MAX_TOKEN_ACROSS_CHUNK] tensor, so dropping it
also removes an allocation, but that was not measured either. Whoever ports this
should benchmark it rather than quote upstream.

Start with the port comments in include/vt/ops.h:5225-5235 and compare the upstream cache_kernels.cu changes with src/vt/cuda/cuda_mla_prefill.cu:310-344. Trace GatherMlaCache through mla_chunked_context.h:274 and its listed callers, then port the page-task launch, remove token_to_seq where unused, and add the upstream kernel test. Done means the relevant tests pass and a benchmark measures the launch and allocation changes.