[zhaoyuanh-nvidia]

Description

Aggregated Helix prefill runs the full prompt on every context-parallel (CP) rank for attention, but each rank must persist only the compressed MLA KV records in the cache pages it owns. Previously, prefill could write non-owned records into rank-local cache storage, so the cache was not a valid handoff to Helix decode.

This PR maps each new context token's global position to a rank-local KV slot using the Helix page ownership rule. A slot of -1 suppresses the persistent KV and scale-cache writes on non-owning ranks. The existing MLA RoPE/append operator accepts these optional slots and continues to prepare temporary Q/K for every prompt token, which dense prefill still needs to compute attention. Dense MLA performs the append/preprocessing before attention and consumes the prepared temporary K/V without appending a second time. Callers that omit the slots retain the existing behavior.

Cached-context attention still requires distributed historical-KV gathering and fails explicitly until TRTLLM-17123 provides it. This PR covers owner-only persistence for dense MLA on the TRTLLM attention backend; it does not complete aggregated serving by itself. The companion #20759 provides aggregate full-prompt admission and cache allocation.

Tracks TRTLLM-17122, part of TRTLLM-16624. Design: DCP for Aggregated Workloads.

Test Coverage

• tests/unittest/_torch/attention/test_helix_context_slots.py checks the host-side global-to-local slot mapping across requests, page boundaries, and CP ranks; it is selected by l0_cpu.yml.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py checks exact cache ownership, untouched non-owned storage, compact cache allocation, and direct prefill-to-decode handoff with BF16/FP8 weights and KV caches. Its distributed cases include TP1/CP2, TP2/CP2, and TP1/CP4.
• The existing disaggregated decode fixture keeps its rank-local input history, and the Kimi MLA reference test adjusts the combined FP8-weight/FP8-cache numerical bound using an unsharded control.
• A cached-context case verifies that the unsupported distributed-history path fails before modifying the cache; a separate operator check covers nonzero global prefixes and RoPE.

Validation

• On this GitHub port, Python syntax, git diff --check, changed-file formatting/lint hooks, and DCO pass. The waiver and test-list checks pass manually with Python 3.12. The legacy-config hook reports pre-existing stale entries.
• The earlier implementation passed a native GB300 build, 85 GPU regression cases, and 17 CPU metadata cases, including a parent/fix control that reproduced the incorrect cache writes. Those runs used an earlier base and do not establish GPU correctness on this GitHub branch.
• Native rebuild and GPU regression on the combined GitHub feature branch remain pending. Full-checkpoint Kimi K3 aggregate serving and performance require the companion integration work.

PR Checklist

• The description covers the motivation, behavior, scope, and validation limits.

• Regression coverage exercises rank-local ownership and the prefill-to-decode cache handoff.

• The new operator argument is optional for existing callers; no dependency or ownership changes are introduced.

• Attention, runtime, and kernel owners are appropriate reviewers.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see the available CI bot commands, comment /bot help.