NVIDIA / NVIDIA/TensorRT-LLM

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

Open
#15,367 1 comment 0 reactions 1 assignee View on GitHub

@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_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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.