[Feature Request] Context-parallel (Helix) sharding for MLA in AutoDeploy
Open
@greg-kwasniewski1 is already working on this.
Since Jun 15, 2026.
AutoDeploy
feature request
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Summary
AutoDeploy (AD) supports TP/EP/attention-DP sharding for MLA, but not context parallelism (CP). Under TP, the MLA latent KV cache is replicated on every rank (single latent KV head), making the per-rank latent read/footprint a non-scaling term in decode. This tracks adding exact, no-retrain context-parallel (Helix) sharding of the MLA latent KV cache in AD, so each rank stores/reads only 1/cp_size of the latent.
Proposed approach
- Reuse base TRT-LLM's Helix CP machinery (no new kernels): each CP rank runs MLA attention over its KV shard producing a partial output + softmax stats; an all-to-all exchanges them and
torch.ops.trtllm.helix_post_processdoes the online-softmax (LSE) combine. MLA requiresCpType.HELIX(Ulysses is rejected for the single KV head). - Thread
cp_size/cp_typethrough AD'sDistConfig->to_mapping()so the reused baseKVCacheManagerallocates the per-rank latent shard. - Wire the Helix combine into the AD
trtllm_mladecode op, gated oncp_size>1, mirroring the existing MoE all-to-all pattern (_check_moe_alltoall/_run_moe_with_alltoall). - Decode-only split-KV MVP first (tokens replicated, MLP token reduce-scatter deferred); prefill ring-CP is a follow-up.
Scope
- Surface:
tensorrt_llm/_torch/auto_deploy/utils/dist_config.py,.../custom_ops/mla/trtllm_mla.py,.../shim/ad_executor.py,.../custom_ops/attention_interface.py,.../transform/library/sharding.py,.../llm_args.py. - Affected components: AD DistConfig, MLA custom op + KV-cache wiring, sharding transforms, executor dist/process-group setup.
Tests / validation
- Multi-GPU equivalence test (CP is exact):
cp_size=1vscp_size=2/4produce identical logits on the same input. - Per-rank latent cache footprint scales
1/cp_size. - Runs on 8xH100 (CP keeps head_dim=576; NCCL all-to-all path).
Risk
library-visible (MLA custom op, KV-cache wiring, sharding transforms; multi-GPU CI).
Links
- JIRA: None
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.