[zhaoyuanh-nvidia]

Description

Helix MLA context-parallel ranks process the prompt for attention, but each rank owns only part of the persistent KV cache. Persist only the compressed KV records assigned to the rank's cache pages while preserving the temporary K/V needed for context attention.

Build a batch-local map from global token positions to rank-local cache slots and pass it through an optional MLA RoPE/append input. Negative slots suppress persistent KV and scale writes without skipping temporary Q/K preparation. Existing callers that omit this input retain their behavior.

Preprocess and append dense MLA context once, then reuse the prepared K/V for attention. Initialize fresh Helix MLA indptrs even when block reuse and chunked prefill are disabled, clear stale slot metadata when leaving the context path, and reject unsupported cached context before cache mutation.

Tracks TRTLLM-17122, part of TRTLLM-16624. This ports the owner-only write work from #20921 onto main, following the phase-dependent MLA heads merged in #20833. Historical-KV gathering, aggregate request admission/cache allocation, and feature-branch ancestry repair remain separate. This PR does not enable aggregate serving by itself.

Design: DCP for Aggregated Workloads

Scope and compatibility

• Owner-only KV persistence is implemented for fresh dense MLA prefill on the TRTLLM attention backend, including Kimi K3.
• Phase-dependent context heads and construction-time rejection of other dense Helix attention backends come from [TRTLLM-17121][feat] Support phase-dependent MLA heads for Helix prefill #20833 and are unchanged here.
• Sparse MLA paths (DSA and DeepSeek-V4, where sparse_attn_hooks is enabled) retain their existing head layout. This PR does not establish owner-only prefill or context/output-projection compatibility for those sparse Helix paths; these remain follow-up work. The shared metadata initialization change is not sparse Helix enablement.
• Helix context with cached history is rejected before the new owner-write append; distributed historical-KV gathering remains separate.
• Explicit enable_context_mla_with_cached_kv=False is rejected for Helix MLA before metadata allocation: owner-only appends require both cached-token counters and append indptr buffers. The default and explicit True remain supported, and explicit-disable behavior is preserved for non-Helix MLA and non-MLA configurations.

Test Coverage

• tests/unittest/_torch/attention/test_helix_context_slots.py: rank-local slot planning across page boundaries, ragged requests, nonzero prefixes, and CP2/4/8/32; stale-slot reset, generation-metadata preservation, and same-object reuse through decode/non-CP transitions back to changed context lengths and prefixes.
• tests/unittest/_torch/executor/engine/test_metadata.py: factory defaults and explicit overrides; early Helix MLA explicit-disable rejection before backend construction, with ordinary MLA, non-MLA Helix, and non-Helix CP compatibility controls. Default/explicit-enable lifecycle tests run the real metadata constructor and prepare() with CPU storage replacing CUDA allocation/pinning, checking counters, owner slots, and host/device append indptrs across fresh mixed, decode-only, and changed fresh batches on both CP2 ranks.
• tests/unittest/_torch/attention/test_attention_mla.py: update the metadata double for the existing SM90/99/100 chunked-prefill dispatch regression.
• tests/unittest/_torch/modeling/test_kimi_k3_mla.py: distinguish FP8-weight and FP8-cache reference tolerances, including their combined quantization error; retain the weight-only error bound.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py: exact owned writes, untouched non-owned cache, compact allocation, mixed prefill/decode, direct handoff, and a nonzero-prefix writer probe. A global-position/page-sensitive quarter-turn RoPE fixture uses an independent signed-swap oracle; it is not full end-to-end position-varying RoPE lifecycle validation.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py: retain rank-local history in the disaggregated fixture and a separate unsharded FP8-weight/FP8-cache reference control.

Existing attention/executor selections in l0_cpu.yml cover the CPU tests. Four-GPU B300/GB300 post-merge attention selections cover the distributed matrix: BF16/FP8 weights and KV with TP1/CP2, TP2/CP2, and TP1/CP4. This matrix does not cover FP4 CP and is not executed by ordinary single-GPU CI.

Validation

• Current head bfeea235: two DCO-signed commits. Changed-file pre-commit, commit hooks, syntax checks, and git diff --check passed. The full-repository hook run passed the other applicable hooks but could not launch mypy because /bin/bash is unavailable in this Windows environment.
• Supplemental source-isolated host validation: all 20 targeted factory/override/preparation cases pass using the checked-in factory, real metadata dataclasses, and test bodies. The identical command against the old factory fails the three Helix MLA explicit-disable rejection cases; fixed, restored-old, and restored-fixed runs confirm the result. This bypasses package/native imports and MPI setup and is not a native repository pytest or GPU numerical validation pass. Normal repository pytest was attempted and blocked before collection by missing TensorRT-LLM native libraries on Windows.
• Prior head 3d3c463b: native x86/SBSA builds and all 25 intended named CPU cases passed without skips in the targeted CI run. Those results do not validate the current Python guard or added tests and do not establish full required CI success. The explicit-disable fix changes no C++/CUDA sources; current-head native/GPU validation remains pending.

PR Checklist

• The description explains the motivation, implementation, scope, and validation limits.

• Changes include regression coverage and documented existing CI selections.

• The native operator input is optional and backward-compatible; no user-facing serving configuration changes are introduced.

• No new dependencies or ownership changes are introduced.

• The change fits the existing attention architecture; no architecture-diagram update is needed.

• Attention, runtime, disaggregation, and model code owners are assigned as 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.

Dev Engineer Review

• The MLA RoPE/cache append API accepts an optional per-token local-slot map. Negative slots suppress persistent KV writes; a null map retains the previous slot calculation. The operator validates the tensor type, rank, length, and device.
• Helix context metadata maps global token positions to rank-local slots. The MLA path prepares and appends context once, reuses prepared K/V for attention, and rejects cached Helix context before cache mutation.
• Metadata enables cached-KV MLA context for eligible Helix MLA configurations. Verify the default and explicit-override behavior across supported configurations.
• Owner-only writes do not add historical-KV gathering or aggregate request admission and cache allocation. Aggregate serving is not established by this change. FP4 context parallelism is also outside the supplied validation scope.
• Current-head test results and runtime validation are not established by the supplied evidence. Prior targeted results do not establish current-head or full required CI success.

QA Engineer Review

The tests cover owner-slot planning and metadata lifecycle, Helix MLA cache writes, compact allocation, prefill-to-decode handoff, and FP8 tolerance behavior. The Helix integration tests and new planner and metadata tests do not appear in the searched CI or manual-QA lists. test_attention_mla.py appears in three CI lists, but its test_attention_mla case is waived in waives.txt; test_kimi_k3_mla.py appears in l0_b200.yml with a timeout. Coverage verdict: needs follow-up. Confirm CI coverage for the new tests and validate them on the current head.

Per-File QA Perspective

• cpp/tensorrt_llm/kernels/mlaKernels.cu: Verify owner-local slot routing and negative-slot write suppression in both KV-copy paths. Confirm null-map fallback behavior.
• cpp/tensorrt_llm/kernels/mlaKernels.h: The append launcher API gains a local-slot pointer. Verify declaration and implementation signatures match callers.
• cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp: The operator gains an optional slot tensor with type, rank, length, and device checks. Verify omitted-tensor behavior and dtype dispatch.
• tensorrt_llm/_torch/attention/backends/trtllm.py: Helix context preparation adds a per-token local-slot tensor. Verify ownership mapping and clearing or recomputation when metadata is reused.
• tensorrt_llm/_torch/attention/mla.py: The Helix path reuses prepared K/V and rejects cached context. Verify rejection leaves cache state unchanged and ordinary MLA behavior remains intact.
• tensorrt_llm/_torch/pyexecutor/engine/metadata.py: Metadata defaults change for eligible Helix MLA configurations. Verify explicit overrides and behavior outside those configurations.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py: Clears owner-slot metadata in the disaggregated decode fixture. No matching CI or manual-QA list entry was found.
• tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py: Adds tests for owned prefill, compact caches, FP8 configurations, and context-head layouts. No matching CI or manual-QA list entry was found.
• tests/unittest/_torch/attention/test_attention_mla.py: Updates the metadata test double. The file appears in three CI lists; the test_attention_mla case is waived in waives.txt.
• tests/unittest/_torch/attention/test_helix_context_slots.py: Adds CPU tests for slot ownership, reset behavior, and metadata reuse. No matching CI or manual-QA list entry was found.
• tests/unittest/_torch/executor/engine/test_metadata.py: Adds tests for Helix MLA defaults, explicit disable behavior, and preparation lifecycle. No matching CI or manual-QA list entry was found.
• tests/unittest/_torch/modeling/test_kimi_k3_mla.py: Adjusts reference tolerances for FP8 weights and KV cache. The file appears in l0_b200.yml with a timeout.

[zhaoyuanh-nvidia]

/bot run --disable-fail-fast --extra-stage "GB300-4_GPUs-PyTorch-Post-Merge-1,GB300-4_GPUs-PyTorch-Post-Merge-2,GB300-4_GPUs-PyTorch-Post-Merge-3"

[tensorrt-cicd]

PR_Github #81372 [ run ] triggered by Bot. Commit: 347e182 Link to invocation

[zhaoyuanh-nvidia]

/bot run --disable-fail-fast --stage-list "CPU-Generic-x86-1,CPU-Generic-x86-2"

[tensorrt-cicd]

PR_Github #81414 [ run ] triggered by Bot. Commit: 3d3c463 Link to invocation

[tensorrt-cicd]

PR_Github #81372 [ run ] completed with state ABORTED. Commit: 347e182

Link to invocation

[tensorrt-cicd]

PR_Github #81414 [ run ] completed with state SUCCESS. Commit: 3d3c463
/LLM/main/L0_MergeRequest_PR pipeline #67090 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[zhaoyuanh-nvidia]

/bot run --disable-fail-fast --stage-list "CPU-Generic-x86-1,CPU-Generic-x86-2"

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
py-filelock	CVE-2026-112906	filelock is a platform-independent file lock for Python. From 4.1.0 until 4.2.1, SoftFileLock versions from 4.1.0 can reuse a stale observation when an attacker pre-dates and removes its own marker, at least two lock-user processes are waiting, coarse or matching modification times preserve the observation, and an honest contender creates a fresh marker on the reused inode, causing that live marker to be removed as stale. This issue is fixed in version 4.2.1.	MEDIUM

[coderabbitai]

Walkthrough

This change adds rank-local Helix context slots to MLA cache appends. Metadata plans the slots, and the append kernel uses them to select KV-cache destinations. MLA context handling and tests cover ownership, cache allocation, and prefill-to-decode behavior.

Changes

Helix MLA Context KV Ownership

Layer / File(s)	Summary
Plan Helix context-local slots tensorrt_llm/_torch/attention/backends/trtllm.py, tests/unittest/_torch/attention/test_helix_context_slots.py	Attention metadata computes rank-local slots for context tokens and clears them when context or cache-mapping conditions are not met. CPU tests cover ownership, reset behavior, and metadata reuse.
Forward local slots to KV-cache writes cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp, cpp/tensorrt_llm/kernels/mlaKernels.h, cpp/tensorrt_llm/kernels/mlaKernels.cu	The append operator accepts and validates optional local slots, then forwards them to the kernel. The kernel uses supplied slots or the previous cache index and skips KV writes for negative slots.
Use owned slots in MLA context tensorrt_llm/_torch/attention/mla.py, tensorrt_llm/_torch/pyexecutor/engine/metadata.py, tests/unittest/_torch/attention/test_attention_mla.py, tests/unittest/_torch/executor/engine/test_metadata.py	MLA context handling routes owned Helix KV through the TRTLLM append path and uses cache RoPE data. The metadata default also enables context MLA for Helix parallelism.
Validate ownership and prefill-to-decode handoff tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py, tests/unittest/_torch/attention/multi_gpu/test_mla_helix_context.py, tests/unittest/_torch/modeling/test_kimi_k3_mla.py	Tests cover owned KV writes, compact cache allocation, prefill-to-decode handoff, FP8 cache cases, and reference outputs across Helix configurations.

Priority: Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Metadata as TrtllmAttentionMetadata
  participant MLA as MLA.forward_context_default
  participant Attention as TRTLLM attention
  participant Kernel as applyMLARopeAppendPagedKVAssignQKernel
  Metadata->>Attention: Pass computed Helix local slots
  MLA->>Attention: Append latent cache and assign Q
  Attention->>Kernel: Forward local slots
  Kernel->>Kernel: Select destination and write owned KV

Loading

Suggested reviewers: litaotju

Merge Risk: Minimal · up to bfeea

Dense Helix prefill is intended to persist KV only on its owning rank. No concrete merge-blocking regression was established in the reviewed paths.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title follows the required ticket/type format and clearly describes the main change: owner-only MLA KV writes for Helix prefill.
Description check		The description is complete and relevant. It explains the motivation, implementation, scope, compatibility limits, test coverage, validation status, and checklist items.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 27 functions across 12 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[yuxianq]

@farazkh80, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~

Assignment explanation

Components	kernels/auxiliary (25%); native-ops/qkv-rope (10%); backends/trtllm (20%); modules/mla (20%); runtime/metadata-workspaces (5%); integration/attention-tests (20%)
Tags	algorithm:dense (100%); representation:mla (100%); phase:prefill (100%); features:context-parallelism (40%); features:paged-kv (30%); features:rope (20%); features:quantization (10%); context:llm (100%)
PR effort	8 points — Requires tracing global-to-local ownership through metadata, native validation, CUDA writes and MLA dispatch while preserving temporary attention inputs and avoiding duplicate preprocessing. Quantized storage and metadata reuse add substantial reasoning; truncated implementation and test sections limit verification.
Importance	Important — Rank-local addressing and negative-slot suppression change persistent CUDA cache writes. MLA prefill also changes preprocessing order and latent-cache forwarding, creating sensitive distributed-cache and RoPE correctness contracts.
Affinity (0–1)	0.492 component/tag match; 0.369 after PIC preference
Outstanding points at selection	0 → 8 (+8)
Points already reserved on this PR	0
Effective weight	1
Completed review points (preceding 168 hours)	48
Combined points after assignment	56
Demand-derived target	48.00 points (soft share)
Additional fair share before assignment	0.00 points
Share explanation	Whole PR / qualification or coverage constraint exceeds a soft share; assignment is allowed.
Selection rule	affinity with 0.3-share guard
Assignment scope (farazkh80)	all

Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.

[chienchunhung]

Overall LGTM; 1 minor finding (please see inline comment). Thanks!

[chienchunhung]

With Helix MLA and enable_context_mla_with_cached_kv=False, preparation still sets helix_context_local_slots. This guard then reads num_ctx_cached_tokens, which is initialized only when the flag is enabled, causing AttributeError. The new explicit-disable test checks only the forwarded flag. Consider rejecting this combination early or initializing owner-append metadata independently, and test the preparation lifecycle.

PS: The normal engine default enables the flag; I found no production Helix caller using this override.

[zhaoyuanh-nvidia]

Fixed in bfeea23. Helix MLA now rejects explicit enable_context_mla_with_cached_kv=False before metadata allocation, while other configurations retain the override. Added regression coverage for rejection and the default/explicit-enable preparation lifecycle, including reuse across mixed and decode-only batches.