feat: add symmetric-memory LM-head all-gather

[qescccczmr]

Motivation

Tensor-parallel ParallelLMHead currently materializes each vocabulary shard and uses dist.all_gather_into_tensor to reconstruct full logits. For the small token counts common in decode/speculative decode, collective launch, allocation, and rank-layout overhead are a meaningful fraction of LM-head latency.

This PR adds an opt-in Hopper/NVLink symmetric-memory V1 path. Each rank writes its BF16 logits shard through multimem.st into a symmetric full-logits arena, then returns an owning output tensor. The existing NCCL implementation remains the default and fallback.

Modification

• Add a Triton multimem.st last-dimension all-gather for TP2/TP4/TP8 BF16 logits.
• Integrate it into the existing ParallelLMHead.forward path after the native local GEMM.
• Add TP-wide setup and first-use admission so ranks cannot split between symmetric-memory and NCCL paths during rendezvous.
• Preserve the existing NCCL path for disabled/unsupported configurations and requests larger than the configured arena.
• Release/rebuild the non-module arena across model device/dtype transitions, including engine sleep/wakeup.
• Keep all peer GPUs visible in Ray workers, explicitly bind the assigned CUDA device, and reject duplicate per-node bindings.

The feature is disabled by default:

export LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD=1
export LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB=128  # optional; default 64

Token-level performance

The following is a synthetic kernel-level benchmark, not an end-to-end serving
benchmark. It measures the TP-local BF16 logits all-gather at the GLM5.2
shape (V=154880, H=6144) on 8x NVIDIA H200 with TP8. M is the number of
logits rows submitted to one collective. Each value is the rank-wise maximum
latency across the eight ranks; the reported value is the median of 150
iterations after 30 warmup iterations.

• Base: dist.all_gather_into_tensor plus the existing output layout.
• V1: MultimemAllGatherer (multimem.st), with the production-safe owning
  output (safe=True, including the output clone).
• Both paths use BF16 and identical inputs. Correctness passed exact
  bitwise comparison (rtol=0, atol=0).

Rows M	Base AG (us)	V1 AG (us)	AG speedup	Base GEMM+AG (us)	V1 GEMM+AG (us)	GEMM+AG speedup
1	52.624	58.128	0.905x	107.376	103.808	1.034x
6	65.312	61.488	1.062x	116.336	107.568	1.082x
8	64.912	57.952	1.120x	121.968	110.288	1.106x
32	102.528	90.880	1.128x	171.520	151.984	1.129x
128	214.624	174.608	1.229x	269.264	222.048	1.213x

M=1 is representative of a single-token decode step;
M=6 is representative of one target token plus five speculative draft tokens.
LargerM values represent batched decode or chunked/prefill-sized work. The table
does not include model scheduling, KV-cache work, sampling, or network/server
overhead, so it must not be interpreted as TTFT/TPOT or serving throughput.
The static input contract (BF16, contiguous/aligned local logits, device and shard width) is collectively admitted once and must remain unchanged for that LM-head instance. ParallelLMHead provides this invariant; token count can vary and uses the NCCL fallback when it exceeds arena capacity.

BC-breaking

None. The optimization is opt-in and the existing NCCL path remains unchanged when it is disabled or unavailable.

Use cases

Enable LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD=1 for tensor-parallel CUDA deployments on supported Hopper/NVLink systems. Unsupported environments continue through the existing NCCL implementation.

Checklist

1. Lint checks are run for the changed files.
2. No new downstream dependency is introduced.

CuTeDSL/TileLang fused-epilogue experiments are intentionally out of scope for this V1 PR.

[Copilot]

🟡 Changes recommended

The new symmetric-memory all-gather implementation contains a correctness bug (type(group) is dist.ProcessGroup) that will fail for valid process groups and prevent the feature from functioning when enabled.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR introduces an opt-in symmetric-memory (“multimem.st”) all-gather fast path for tensor-parallel ParallelLMHead logits to reduce per-step collective overhead in small-token decode, while preserving the existing NCCL all_gather_into_tensor path as the default/fallback. It also adjusts Ray worker GPU visibility/device binding to support symmetric-memory rendezvous requirements.

Changes:

• Add a Triton-based symmetric-memory last-dimension all-gather implementation (MultimemAllGatherer) and integrate it into ParallelLMHead.all_gather_logits when enabled and eligible.
• Add environment flags to control the feature (LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD, LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB).
• Update Ray executor initialization to keep peer GPUs visible and bind each worker to its assigned CUDA device when symmetric-memory features are in use.

File summaries

File	Description
lmdeploy/pytorch/nn/embedding.py	Adds TP-wide opt-in admission and a symmetric-memory all-gather fast path for ParallelLMHead, with NCCL fallback.
lmdeploy/pytorch/envs.py	Introduces env toggles and sizing for symmetric-memory LM-head arena.
lmdeploy/pytorch/engine/executor/ray_executor.py	Preserves full CUDA visibility and enforces unique per-node device binding for Ray workers when symmetric-memory is needed.
lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py	New Triton + symmetric-memory implementation for last-dim all-gather, plus guarded wrapper for safe fallback/capture behavior.

Review details

Suppressed comments (1)

lmdeploy/pytorch/nn/embedding.py:185

• This warning message is emitted both when TP ranks have inconsistent arena config and when max_tokens <= 0 (arena too small for even one token row), which makes the log misleading and harder to debug.

            same_config = _tp_same_config((capacity, gathered_width, max_tokens), device, self.tp_group)
            if max_tokens <= 0 or not same_config:
                if self.tp_rank == 0:
                    logger.warning('symmetric-memory LM-head disabled because TP ranks have inconsistent arena config')
                return

• Files reviewed: 4/4 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.