[Feature Request] Bandwidth-optimal tensor-parallel sharding of MLA (latent-cache sharding / TPLA)

Summary

Under tensor parallelism (TP), Multi-head Latent Attention (MLA) replicates its single compressed latent KV cache on every rank: the latent is shared across all query heads, so each rank must store and read the full latent regardless of TP size. The per-rank latent KV read is therefore a non-scaling term in decode (it does not shrink as TP grows), capping decode throughput/latency for long-context, high-concurrency MLA serving. This issue tracks an investigation into TP sharding strategies that reduce the per-rank latent KV footprint and read bandwidth.

Background / problem

• AutoDeploy already TP-shards MLA weights and attention FLOPs (q_b_proj colwise, kv_b_proj per-head, o_proj rowwise + all_reduce), but the cached latent (kv_lora_rank + qk_rope, e.g. 576/token for DeepSeek-V3) is produced by a replicated kv_a_proj and is duplicated across TP ranks.
• Per-rank decode read volume V(P) ~= weights/P + KV(1): every term scales with TP size P except the latent KV read, which is constant. That single non-scaling term is the bandwidth wall for MLA decode.

Directions under investigation

• TPLA (approximate, drop-in on MLA checkpoints): shard the latent feature dimension across ranks; each rank runs local attention over its latent slice and outputs are combined with the existing o_proj all_reduce. Requires an orthogonal (Hadamard/PCA) reparameterization of kv_b_proj/kv_a_layernorm to recover accuracy (independent per-shard softmax is not exact). ~1.8-1.9x decode speedups reported in the literature at 32K context.
• Exact alternatives (no approximation, no retraining) that shard the latent KV on a different axis: context/sequence parallelism (sequence split + flash-attention log-sum-exp merge) and attention data parallelism (request split). The feature/head axis cannot be sharded exactly under TP because softmax couples the full latent.

Methodology (first step)

Before implementing a real path, measure the upper bound: an env-gated hack (TPLA_LATENT_DIV=G) shrinks the per-rank cached/read latent by a factor G (numerically garbage, perf-representative) and we benchmark decode inter-token latency / throughput vs G across context lengths. If the ceiling is worthwhile, implement the accuracy-preserving path.

Scope

• Surface: tensorrt_llm/_torch/auto_deploy/custom_ops/mla/, tensorrt_llm/_torch/auto_deploy/transform/library/sharding*.py, AutoDeploy model-registry MLA configs.
• Affected components: MLA custom ops + KV-cache initializers, TP sharding transforms, (later) a load-time weight reparameterization.

Tests / validation

• Upper-bound benchmark on an MLA model (DeepSeek-V2-Lite / V3) on 8xH100 / 8xB200.
• For a real implementation: multi-GPU equivalence (G=1 matches head-TP MLA exactly), accuracy eval after reparam, decode throughput.

Risk

library-visible (MLA custom ops, sharding IR, KV-cache handlers; multi-GPU CI).

Links

• Investigation notes / hand-off doc: mla_sharding_investigation.md (branch gk/mla-sharding-investigation).
• JIRA: None