[zhaoyuanh-nvidia]

Description

Helix aggregate chunked prefill must attend to earlier MLA KV records that are stored across context-parallel (CP) ranks. Before this change, a cached context could not reconstruct its full token history from its rank-local persistent cache.

This PR derives each rank's owned read intervals from the existing global history-chunk plan, loads the corresponding compressed KV records, and uses the variable-size CP all-gather to assemble each bounded history chunk. Every rank participates in the same collective sequence, including iterations where it contributes zero records. FP8 cache records are dequantized locally to the query dtype before communication. A CUDA reorder operator restores request and token order, then separates compressed KV from the shared RoPE key segment (k_pe) for the existing MLA expansion, attention, and softmax-merge path. Persistent KV remains sharded.

Cached Helix contexts require chunked prefill on SM100 or newer; the unsupported nonchunked path fails explicitly. This PR supplies historical-KV gathering for dense MLA. Request admission, owner-only writes, and initial KV allocation are provided by #20923, #20921, and #20920 on the feature branch. Full-checkpoint aggregate serving, prefix reuse, and performance remain integration work.

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

Test Coverage

• tests/unittest/_torch/attention/kernels/test_mla_helix_reorder.py checks exact record ordering and K/k_pe splitting for FP16, BF16, and FP32, including empty requests and page/cycle boundaries.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py checks multiple real prefill chunks, unequal and empty CP contributions, BF16/FP8 weights and KV storage, TP/CP layouts, exact cache ownership, mixed context/decode batches, metadata restoration, and direct decode handoff.
• The new attention tests are included through the GB300 multi-GPU CI list's unittest/_torch/attention selection.

Validation

• This GitHub port is one DCO-signed commit on the current feat/helix-agg head and has the same stable patch ID as the validated seven-file implementation. Python syntax, git diff --check, changed-file formatting/lint hooks, and the waiver and AST test-list checks pass. The repository-wide legacy-config hook reports pre-existing stale entries; the two test-list hooks were rerun successfully with Python 3.12 because the Windows hook uses Python 3.9.
• Earlier GB300/SM103 validation of the implementation passed a native build, 48 exact CUDA reorder cases, 20 distributed-history cases, 85 existing GPU regressions, 17 CPU metadata cases, and 43 cache/runtime cases. A parent/fix control reproduced the missing-history failure on the parent and passed with the feature. These results were obtained on an earlier base; they do not establish GPU correctness on this GitHub head.
• Native rebuild, GPU regression, full-checkpoint Kimi K3 aggregate serving, and performance validation on the combined GitHub feature branch remain pending.

PR Checklist

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

• Regression coverage exercises the reorder operator and distributed cached-history handoff.

• No user-facing configuration, 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.