[Feature Request] Context-parallel (Helix) sharding for MLA in AutoDeploy

[greg-kwasniewski1]

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_process does the online-softmax (LSE) combine. MLA requires CpType.HELIX (Ulysses is rejected for the single KV head).
• Thread cp_size/cp_type through AD's DistConfig -> to_mapping() so the reused base KVCacheManager allocates the per-rank latent shard.
• Wire the Helix combine into the AD trtllm_mla decode op, gated on cp_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=1 vs cp_size=2/4 produce 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

[tanishq0421]

Hi @greg-kwasniewski1, I want to work on this feature request.