From 29c4317d086433dc3836d38740024340f9fc2cc5 Mon Sep 17 00:00:00 2001
From: qescccczmr <bryentonxemmrj@gmail.com>
Date: Fri, 28 Aug 2026 08:34:57 +0000
Subject: [PATCH 1/6] feat: add symmetric-memory LM-head all-gather

---
 .../backends/cuda/comm/symm_mem_allgather.py  | 521 ++++++++++++++++++
 lmdeploy/pytorch/envs.py                      |   2 +
 lmdeploy/pytorch/nn/embedding.py              |  25 +
 3 files changed, 548 insertions(+)
 create mode 100644 lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py

diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
new file mode 100644
index 0000000000..261086a1fe
--- /dev/null
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -0,0 +1,521 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Symmetric-memory ``multimem.st`` all-gather along the hidden (last) dim.
+
+Each rank stores its ``[T, H/TP]`` shard into a multicast buffer in one NVLink
+pass instead of an NCCL ring; ``create_state`` rendezvous once so launches are
+CUDA-graph capturable.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+# Each thread moves _NUMEL_PER_THREAD bf16 via one 128-bit multimem op; the
+# grid-strided block count is tunable in [_MIN_BLOCKS, _MAX_BLOCKS].
+_BLOCK_THREADS = 1024
+_NUMEL_PER_THREAD = 8
+_MIN_BLOCKS = 4
+_MAX_BLOCKS = 32
+_SUPPORTED_WORLD_SIZES = {2, 4, 8}
+
+
+# ------------------------------------------------------------------------------
+# Low-level PTX helpers
+# ------------------------------------------------------------------------------
+
+
+@triton.jit
+def _multimem_st_128(multicast_ptrs, x, y, z, w, mask):
+    return tl.inline_asm_elementwise(
+        """
+        {
+            .reg .pred %p0;
+            setp.eq.s32 %p0, $6, 1;
+            @!%p0 bra end;
+            multimem.st.relaxed.sys.global.v4.f32 [$1], {$2, $3, $4, $5};
+            end:
+        }
+        """,
+        '=r,l,r,r,r,r,r',
+        args=[multicast_ptrs, x, y, z, w, mask.to(tl.int32)],
+        dtype=(tl.uint32),
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _local_ld_128(in_ptr, mask):
+    return tl.inline_asm_elementwise(
+        """
+        {
+            .reg .pred %p0;
+            setp.eq.s32 %p0, $5, 1;
+            @!%p0 bra end;
+            ld.relaxed.sys.global.v4.b32 {$0, $1, $2, $3}, [$4];
+            end:
+        }
+        """,
+        '=r,=r,=r,=r,l,r',
+        args=[in_ptr, mask.to(tl.int32)],
+        dtype=(tl.uint32, tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_tid():
+    return tl.inline_asm_elementwise(
+        """
+        mov.u32 $0, %tid.x;
+        mov.u32 $1, %tid.y;
+        mov.u32 $2, %tid.z;
+        """,
+        '=r,=r,=r',
+        [],
+        dtype=(tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_ntid():
+    return tl.inline_asm_elementwise(
+        """
+        mov.u32 $0, %ntid.x;
+        mov.u32 $1, %ntid.y;
+        mov.u32 $2, %ntid.z;
+        """,
+        '=r,=r,=r',
+        [],
+        dtype=(tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_flat_tid():
+    tid_x, tid_y, tid_z = _get_tid()
+    ntid_x, ntid_y, _ = _get_ntid()
+    return tid_z * ntid_y * ntid_x + tid_y * ntid_x + tid_x
+
+
+@triton.jit
+def _sync_threads():
+    tl.inline_asm_elementwise(
+        'bar.sync 0;', '=r', [], dtype=tl.int32, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _fence_proxy_alias():
+    """Order multicast writes before observing the unicast buffer alias."""
+    tl.inline_asm_elementwise(
+        'fence.proxy.alias;', '=r', [], dtype=tl.int32, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _send_signal(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            send_signal:
+                atom.global.relaxed.sys.cas.b32 %tmp32_0, [$1], 0, 1;
+                setp.eq.u32 %p0, %tmp32_0, 0;
+                @!%p0 bra send_signal;
+        }
+        """,
+        '=r, l',
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _send_signal_release(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            send_signal:
+                atom.global.release.sys.cas.b32 %tmp32_0, [$1], 0, 1;
+                setp.eq.u32 %p0, %tmp32_0, 0;
+                @!%p0 bra send_signal;
+        }
+        """,
+        '=r, l',
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wait_signal(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            wait_signal:
+                atom.global.sys.relaxed.cas.b32 %tmp32_0, [$1], 1, 0;
+                setp.eq.u32 %p0, %tmp32_0, 1;
+                @!%p0 bra wait_signal;
+        }
+        """,
+        '=r, l',
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wait_signal_acquire(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            wait_signal:
+                atom.global.sys.acquire.cas.b32 %tmp32_0, [$1], 1, 0;
+                setp.eq.u32 %p0, %tmp32_0, 1;
+                @!%p0 bra wait_signal;
+        }
+        """,
+        '=r, l',
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _blockwise_barrier(
+    signal_pad_ptrs,
+    rank: tl.constexpr,
+    world_size: tl.constexpr,
+    sem: tl.constexpr,
+):
+    block_id = (
+        tl.program_id(2) * tl.num_programs(1) * tl.num_programs(0)
+        + tl.program_id(1) * tl.num_programs(0)
+        + tl.program_id(0)
+    )
+    flat_tid = _get_flat_tid()
+
+    remote_ranks = tl.arange(0, world_size)
+    signal_pad_ptrs = signal_pad_ptrs.to(tl.pointer_type(tl.uint64))
+    remote_signal_pad_addrs = tl.load(signal_pad_ptrs + remote_ranks).to(
+        tl.pointer_type(tl.uint32)
+    )
+    send_addrs = remote_signal_pad_addrs + block_id * world_size + rank
+
+    local_signal_pad_addr = tl.load(signal_pad_ptrs + rank).to(
+        tl.pointer_type(tl.uint32)
+    )
+    wait_addrs = local_signal_pad_addr + block_id * world_size + remote_ranks
+
+    if flat_tid < world_size:
+        if sem == 'relaxed':
+            _send_signal(send_addrs)
+            _wait_signal(wait_addrs)
+        else:
+            _send_signal_release(send_addrs)
+            _wait_signal_acquire(wait_addrs)
+
+
+@triton.jit
+def _all_gather_kernel_inner(
+    input_ptr,
+    multicast_ptr,
+    signal_pad_ptr,
+    total_tokens,
+    hidden_offset,
+    LOCAL_HIDDEN: tl.constexpr,
+    TOTAL_HIDDEN: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUMEL_PER_THREAD: tl.constexpr,
+    RANK: tl.constexpr,
+    WORLD_SIZE: tl.constexpr,
+    SKIP_ENTRY_SYNC: tl.constexpr,
+) -> None:
+    if SKIP_ENTRY_SYNC == 0:
+        _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem='relaxed')
+        _sync_threads()
+
+    chunks_per_row: tl.constexpr = LOCAL_HIDDEN // NUMEL_PER_THREAD
+    total_hidden_chunks: tl.constexpr = TOTAL_HIDDEN // NUMEL_PER_THREAD
+    hidden_offset_chunks = hidden_offset // NUMEL_PER_THREAD
+    total_chunks = total_tokens * chunks_per_row
+
+    pid = tl.program_id(axis=0)
+    tid = _get_flat_tid()
+    block_start = pid * BLOCK_SIZE
+
+    while block_start < total_chunks:
+        chunk = block_start + tid
+        mask = chunk < total_chunks
+        row = chunk // chunks_per_row
+        col_chunk = chunk % chunks_per_row
+
+        in_ptr = input_ptr.to(tl.pointer_type(tl.uint64)) + chunk * 2
+        out_chunk = row * total_hidden_chunks + hidden_offset_chunks + col_chunk
+        out_ptr = (
+            multicast_ptr.to(tl.int64).to(tl.pointer_type(tl.uint64)) + out_chunk * 2
+        )
+        x, y, z, w = _local_ld_128(in_ptr, mask)
+        _multimem_st_128(out_ptr, x, y, z, w, mask)
+        block_start += tl.num_programs(axis=0) * BLOCK_SIZE
+
+    # The producer writes through the multicast VA and callers consume through
+    # the ordinary symmetric-buffer VA.  Hopper requires an alias-proxy fence
+    # before the release/acquire completion handshake.
+    _fence_proxy_alias()
+    _sync_threads()
+    _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem='acq_rel')
+
+
+# ------------------------------------------------------------------------------
+# Public API
+# ------------------------------------------------------------------------------
+
+
+@dataclass
+class MultimemAllGatherState:
+    group: dist.ProcessGroup
+    rank_in_group: int
+    world_size: int
+    device: torch.device
+    max_token_num: int
+    hidden_dim: int
+    comm_buff: torch.Tensor
+    # Rendezvous handle; stable for the buffer's lifetime, resolved once.
+    symm_mem_hdl: Any
+
+
+def create_state(
+    group: dist.ProcessGroup,
+    rank_in_group: int,
+    max_tokens: int,
+    hidden_size: int,
+    device: torch.device | None = None,
+) -> MultimemAllGatherState:
+    """Allocate and rendezvous the symmetric-memory buffer.
+
+    Collective: call
+    once outside CUDA-graph capture with identical args on every rank.
+    """
+    assert type(group) is dist.ProcessGroup, f"Expected ProcessGroup, got {type(group)}"
+    assert hidden_size % _NUMEL_PER_THREAD == 0, (
+        f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
+        f"bf16 for 16-byte multimem.st row alignment"
+    )
+    device = device or torch.device(f"cuda:{torch.cuda.current_device()}")
+
+    # Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
+    pad_bytes = _MAX_BLOCKS * group.size() * 4
+    symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
+    with torch.inference_mode(False), torch.no_grad():
+        comm_buff = symm_mem.empty(
+            (max_tokens, hidden_size), dtype=torch.bfloat16, device=device
+        )
+    hdl = symm_mem.rendezvous(comm_buff, group=group)
+    assert hdl.rank == rank_in_group, (
+        f"symm_mem handle rank {hdl.rank} != rank_in_group {rank_in_group}; the "
+        f"hidden-shard offset would be wrong"
+    )
+    return MultimemAllGatherState(
+        group=group,
+        rank_in_group=rank_in_group,
+        world_size=group.size(),
+        device=device,
+        max_token_num=max_tokens,
+        hidden_dim=hidden_size,
+        comm_buff=comm_buff,
+        symm_mem_hdl=hdl,
+    )
+
+
+def _launch_config(local_numel: int):
+    assert local_numel % _NUMEL_PER_THREAD == 0
+    return _MIN_BLOCKS, _BLOCK_THREADS, _BLOCK_THREADS // 32, _NUMEL_PER_THREAD
+
+
+def all_gather_inner(
+    state: MultimemAllGatherState,
+    hidden_states: torch.Tensor,
+    tp_hidden_dim: int,
+    skip_entry_sync: bool = False,
+    safe: bool = True,
+) -> torch.Tensor:
+    """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim.
+
+    ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``,
+    else a view into the symmetric buffer (valid until the next collective).
+    """
+    world_size = state.world_size
+    assert hidden_states.dtype == torch.bfloat16, 'Only bfloat16 is supported'
+    assert hidden_states.is_contiguous(), 'hidden_states must be contiguous'
+    assert hidden_states.data_ptr() % 16 == 0, (
+        f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
+        f"16-byte aligned for 128-bit multimem.st"
+    )
+    assert (
+        tp_hidden_dim % world_size == 0
+    ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
+    local_hidden = tp_hidden_dim // world_size
+    assert local_hidden % _NUMEL_PER_THREAD == 0, (
+        f"per-rank hidden shard ({local_hidden}) must be a multiple of "
+        f"{_NUMEL_PER_THREAD} bf16"
+    )
+    assert tp_hidden_dim <= state.hidden_dim, (
+        f"comm buffer too narrow: tp_hidden_dim={tp_hidden_dim} > "
+        f"state.hidden_dim={state.hidden_dim}"
+    )
+    total_tokens, in_hidden = hidden_states.shape
+    assert (
+        in_hidden == local_hidden
+    ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
+    assert (
+        total_tokens <= state.max_token_num
+    ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
+
+    hidden_offset = local_hidden * state.rank_in_group
+    symm_mem_hdl = state.symm_mem_hdl
+    num_blocks, block_size, num_warps, numel_per_thread = _launch_config(
+        total_tokens * local_hidden
+    )
+    grid = (num_blocks, 1, 1)
+    _all_gather_kernel_inner[grid](
+        input_ptr=hidden_states,
+        multicast_ptr=symm_mem_hdl.multicast_ptr,
+        signal_pad_ptr=symm_mem_hdl.signal_pad_ptrs_dev,
+        total_tokens=total_tokens,
+        hidden_offset=hidden_offset,
+        LOCAL_HIDDEN=local_hidden,
+        TOTAL_HIDDEN=state.hidden_dim,
+        BLOCK_SIZE=block_size,
+        NUMEL_PER_THREAD=numel_per_thread,
+        RANK=symm_mem_hdl.rank,
+        WORLD_SIZE=symm_mem_hdl.world_size,
+        SKIP_ENTRY_SYNC=1 if skip_entry_sync else 0,
+        num_warps=num_warps,
+    )
+    output = state.comm_buff[:total_tokens, :tp_hidden_dim]
+    return output.clone() if safe else output
+
+
+# ------------------------------------------------------------------------------
+# Guarded wrapper
+# ------------------------------------------------------------------------------
+
+
+class MultimemAllGatherer:
+    """Guarded last-dim multimem all-gather with NCCL fallback.
+
+    Owns one symmetric buffer built lazily on the first eager call, and uses the kernel only when the input fits its
+    dtype/shape/alignment contract. The returned tensor owns its storage so a later collective cannot overwrite logits
+    that are still being consumed on another stream.
+    """
+
+    _UNINIT = object()
+
+    def __init__(
+        self,
+        group: dist.ProcessGroup,
+        rank: int,
+        gathered_width: int,
+        max_tokens: int,
+        *,
+        enabled: bool = True,
+    ):
+        self._group = group
+        self._rank = rank
+        self._gathered_width = gathered_width
+        self._max_tokens = int(max_tokens)
+        # None => always NCCL; _UNINIT => build on first eager call.
+        self._state = self._UNINIT if enabled else None
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor | None:
+        state = self._state
+        if state is self._UNINIT:
+            state = self._build(x)
+            if state is not self._UNINIT:
+                self._state = state
+        if (
+            state is not None
+            and state is not self._UNINIT
+            and x.dtype == torch.bfloat16
+            and x.dim() == 2
+            and x.is_contiguous()
+            and 0 < x.shape[0] <= state.max_token_num
+            and x.data_ptr() % 16 == 0
+            and x.shape[-1] * state.world_size == state.hidden_dim
+        ):
+            return all_gather_inner(
+                state,
+                x,
+                tp_hidden_dim=self._gathered_width,
+                skip_entry_sync=False,
+                safe=True,
+            )
+        return None
+
+    def _build(self, x: torch.Tensor):
+        if x.dim() != 2 or x.dtype != torch.bfloat16:
+            return None
+        if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
+            # Can't allocate under capture; retry later.
+            return self._UNINIT
+        if x.shape[-1] % _NUMEL_PER_THREAD != 0:
+            return None
+        try:
+            world_size = dist.get_world_size(self._group)
+            # tl.arange requires a power-of-two extent.  Keep unsupported TP
+            # layouts on the existing NCCL path rather than failing at JIT.
+            if world_size not in _SUPPORTED_WORLD_SIZES:
+                return None
+            state = create_state(
+                group=self._group,
+                rank_in_group=self._rank,
+                max_tokens=self._max_tokens,
+                hidden_size=self._gathered_width,
+            )
+            if state.symm_mem_hdl.multicast_ptr == 0:
+                # No multicast for this world size / arch; multimem.st would
+                # write nowhere. Fall back to NCCL.
+                logger.warning(
+                    'multimem all-gather disabled (no multicast for world_size=%d)',
+                    state.world_size,
+                )
+                return None
+            return state
+        except Exception as e:
+            logger.warning('multimem all-gather disabled (%s)', e)
+            return None
diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py
index 7212a9df3f..de7eddca16 100644
--- a/lmdeploy/pytorch/envs.py
+++ b/lmdeploy/pytorch/envs.py
@@ -238,6 +238,8 @@ def _patched_get_env(
     # cuda communicator
     enable_flashinfer_allreduce = env_to_bool('LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE', False)
     enable_symm_mem_allreduce = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE', False)
+    enable_symm_mem_lmhead = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD', False)
+    symm_mem_lmhead_max_mb = max(1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB', 64))
 
     # opt-ttft
     opt_ttft_policy = env_to_choice('LMDEPLOY_PT_TTFT_POLICY', 'size', {'fifo', 'size'})
diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py
index c25abaf611..f965b2a843 100644
--- a/lmdeploy/pytorch/nn/embedding.py
+++ b/lmdeploy/pytorch/nn/embedding.py
@@ -3,6 +3,7 @@
 import torch.distributed as dist
 from torch import nn
 
+from lmdeploy.pytorch import envs as _envs
 from lmdeploy.pytorch.backends import OpType, get_backend
 from lmdeploy.pytorch.distributed import get_dist_group, get_dist_manager, get_tp_world_rank
 from lmdeploy.pytorch.weight_loader.model_weight_loader import default_weight_loader
@@ -47,6 +48,7 @@ def __init__(
 
         dist_group = get_dist_group(layer_type=layer_type)
         self.tp_group = dist_group.gpu_group
+        self.tp_rank = dist_group.rank
 
         if is_tp and self.tp > 1:
             self.vocab_size_padded = pad_vocab_size(self.vocab_size, self.padding_size)
@@ -143,6 +145,22 @@ def __init__(
         builder = get_backend().get_layer_impl_builder(OpType.Linear)
         self.impl = builder.build(hidden_size, self.vocab_size_padded, bias, dtype=dtype)
 
+        self._symm_mem_gatherer = None
+        if self.all_reduce and _envs.enable_symm_mem_lmhead:
+            try:
+                from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer
+
+                gathered_width = self.tp * self.vocab_size_padded
+                capacity = _envs.symm_mem_lmhead_max_mb * 1024 * 1024
+                max_tokens = capacity // (gathered_width * torch.bfloat16.itemsize)
+                if max_tokens > 0:
+                    self._symm_mem_gatherer = MultimemAllGatherer(group=self.tp_group,
+                                                                  rank=self.tp_rank,
+                                                                  gathered_width=gathered_width,
+                                                                  max_tokens=max_tokens)
+            except ImportError:
+                pass
+
     def tie_weights(self, embedding: ParallelEmbedding):
         """Tie the local LM-head shard to a parallel embedding shard."""
         self.weight = embedding.weight
@@ -158,6 +176,13 @@ def all_gather_logits(self, local_logits: torch.Tensor) -> torch.Tensor:
         if not self.all_reduce:
             return local_logits[..., :self.vocab_size]
 
+        if self._symm_mem_gatherer is not None:
+            local_logits_2d = local_logits.reshape(-1, local_logits.shape[-1])
+            gathered = self._symm_mem_gatherer(local_logits_2d)
+            if gathered is not None:
+                output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], )
+                return gathered.reshape(output_shape)[..., :self.vocab_size]
+
         input_size = local_logits.size()
         output_size = (input_size[0] * self.tp, ) + input_size[1:]
         logits = local_logits.new_empty(output_size)

From be132c84560873248fd946ecbf962b22fccc7804 Mon Sep 17 00:00:00 2001
From: qescccczmr <bryentonxemmrj@gmail.com>
Date: Thu, 3 Sep 2026 04:20:24 +0000
Subject: [PATCH 2/6] test: harden symmetric-memory LM-head integration

---
 benchmark/profile_lmhead_allgather.py         | 227 ++++++++++++++++++
 .../backends/cuda/comm/symm_mem_allgather.py  | 199 ++++++++++-----
 .../pytorch/engine/executor/ray_executor.py   |  40 ++-
 lmdeploy/pytorch/nn/embedding.py              |  83 ++++++-
 .../engine/test_ray_executor_symm_mem.py      | 127 ++++++++++
 tests/pytorch/nn/test_lm_head_symm_mem.py     |  25 ++
 tests/pytorch/test_symm_mem_allgather.py      | 132 ++++++++++
 7 files changed, 756 insertions(+), 77 deletions(-)
 create mode 100644 benchmark/profile_lmhead_allgather.py
 create mode 100644 tests/pytorch/engine/test_ray_executor_symm_mem.py
 create mode 100644 tests/pytorch/nn/test_lm_head_symm_mem.py
 create mode 100644 tests/pytorch/test_symm_mem_allgather.py

diff --git a/benchmark/profile_lmhead_allgather.py b/benchmark/profile_lmhead_allgather.py
new file mode 100644
index 0000000000..7b7a212faf
--- /dev/null
+++ b/benchmark/profile_lmhead_allgather.py
@@ -0,0 +1,227 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""Compare the production NCCL LM-head gather with symmetric memory.
+
+Example:
+
+.. code-block:: bash
+
+   torchrun --standalone --nproc-per-node=8 \
+     benchmark/profile_lmhead_allgather.py \
+     --tokens 1 8 32 128 --warmup 20 --repeat 200
+
+The benchmark reports the maximum CUDA-event latency across TP ranks. V1
+includes the owning clone performed by ``MultimemAllGatherer``.
+"""
+
+import argparse
+import math
+import os
+import statistics
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer
+
+
+def _nccl_gather(local: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+    """Match ``ParallelLMHead`` allocation, gather and vocabulary layout."""
+    world_size = dist.get_world_size(group)
+    output = local.new_empty((world_size, ) + tuple(local.shape))
+    dist.all_gather_into_tensor(output, local, group=group)
+    return output.movedim(0, 1).reshape(local.shape[0], -1)
+
+
+def _percentile(samples: list[float], q: float) -> float:
+    index = max(0, min(len(samples) - 1, math.ceil(q * len(samples)) - 1))
+    return sorted(samples)[index]
+
+
+def _measure_paired(
+    baseline: Callable[[], torch.Tensor],
+    target: Callable[[], torch.Tensor],
+    *,
+    warmup: int,
+    repeat: int,
+    group: dist.ProcessGroup,
+    local_rank: int,
+) -> dict[str, float]:
+    """Interleave both paths and reduce every sample to the slowest rank."""
+    for _ in range(warmup):
+        baseline()
+        target()
+    torch.cuda.synchronize()
+    dist.barrier(group=group, device_ids=[local_rank])
+
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
+    rank_times = torch.empty(2, dtype=torch.float64, device='cuda')
+    baseline_samples = []
+    target_samples = []
+    for iteration in range(repeat):
+        dist.barrier(group=group, device_ids=[local_rank])
+        functions = (baseline, target) if iteration % 2 == 0 else (target, baseline)
+        outputs = []
+        for index, function in enumerate(functions):
+            starts[index].record()
+            outputs.append(function())
+            ends[index].record()
+        ends[-1].synchronize()
+        elapsed = [starts[index].elapsed_time(ends[index]) * 1000.0 for index in range(2)]
+        if iteration % 2:
+            elapsed.reverse()
+        rank_times[0] = elapsed[0]
+        rank_times[1] = elapsed[1]
+        dist.all_reduce(rank_times, op=dist.ReduceOp.MAX, group=group)
+        baseline_us, target_us = rank_times.tolist()
+        baseline_samples.append(baseline_us)
+        target_samples.append(target_us)
+        del outputs
+
+    baseline_median = statistics.median(baseline_samples)
+    target_median = statistics.median(target_samples)
+    return {
+        'baseline_median_us': baseline_median,
+        'target_median_us': target_median,
+        'baseline_p95_us': _percentile(baseline_samples, 0.95),
+        'target_p95_us': _percentile(target_samples, 0.95),
+        'speedup': baseline_median / target_median,
+    }
+
+
+def _capture(function: Callable[[], torch.Tensor], group: dist.ProcessGroup, local_rank: int):
+    for _ in range(3):
+        function()
+    torch.cuda.synchronize()
+    dist.barrier(group=group, device_ids=[local_rank])
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = function()
+    torch.cuda.synchronize()
+    dist.barrier(group=group, device_ids=[local_rank])
+    return graph.replay, output, graph
+
+
+def _report(rank: int, tokens: int, phase: str, metric: dict[str, float]) -> None:
+    if rank != 0:
+        return
+    print(
+        f'RESULT tokens={tokens} phase={phase} '
+        f'base_median_us={metric["baseline_median_us"]:.3f} '
+        f'v1_median_us={metric["target_median_us"]:.3f} '
+        f'base_p95_us={metric["baseline_p95_us"]:.3f} '
+        f'v1_p95_us={metric["target_p95_us"]:.3f} '
+        f'speedup={metric["speedup"]:.4f}x',
+        flush=True,
+    )
+
+
+@torch.inference_mode()
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--vocab-size', type=int, default=154880)
+    parser.add_argument('--hidden-size', type=int, default=4096)
+    parser.add_argument('--tokens', type=int, nargs='+', default=[1, 8, 32, 128])
+    parser.add_argument('--warmup', type=int, default=20)
+    parser.add_argument('--repeat', type=int, default=200)
+    parser.add_argument('--skip-graph', action='store_true')
+    args = parser.parse_args()
+
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group('nccl')
+    group = dist.group.WORLD
+    rank = dist.get_rank(group)
+    world_size = dist.get_world_size(group)
+    if args.vocab_size % world_size:
+        raise ValueError(f'vocab-size {args.vocab_size} must be divisible by TP={world_size}')
+
+    local_vocab = args.vocab_size // world_size
+    gatherer = MultimemAllGatherer(group, rank, args.vocab_size, max(args.tokens))
+    if not gatherer.prepare(torch.device('cuda', local_rank)):
+        raise RuntimeError('symmetric-memory multicast is unavailable for this TP topology')
+
+    torch.manual_seed(1701 + rank)
+    weight = torch.randn((local_vocab, args.hidden_size), dtype=torch.bfloat16, device='cuda')
+    if rank == 0:
+        print(
+            f'CONFIG gpu={torch.cuda.get_device_name(local_rank)!r} tp={world_size} dtype=bf16 '
+            f'hidden={args.hidden_size} global_vocab={args.vocab_size} local_vocab={local_vocab} '
+            f'warmup={args.warmup} repeat={args.repeat}',
+            flush=True,
+        )
+
+    for tokens in args.tokens:
+        torch.manual_seed(2026 + rank * 17 + tokens)
+        local_logits = torch.randn((tokens, local_vocab), dtype=torch.bfloat16, device='cuda')
+        baseline_output = _nccl_gather(local_logits, group)
+        target_output = gatherer(local_logits)
+        if target_output is None:
+            raise RuntimeError(f'symmetric-memory path rejected tokens={tokens}')
+        torch.testing.assert_close(target_output, baseline_output, rtol=0, atol=0)
+
+        def baseline():
+            return _nccl_gather(local_logits, group)
+
+        def target():
+            return gatherer(local_logits)
+
+        metric = _measure_paired(
+            baseline,
+            target,
+            warmup=args.warmup,
+            repeat=args.repeat,
+            group=group,
+            local_rank=local_rank,
+        )
+        _report(rank, tokens, 'gather', metric)
+
+        torch.manual_seed(4096 + tokens)
+        hidden = torch.randn((tokens, args.hidden_size), dtype=torch.bfloat16, device='cuda')
+
+        def baseline_e2e():
+            return _nccl_gather(F.linear(hidden, weight), group)
+
+        def target_e2e():
+            return gatherer(F.linear(hidden, weight))
+
+        torch.testing.assert_close(target_e2e(), baseline_e2e(), rtol=0, atol=0)
+        metric = _measure_paired(
+            baseline_e2e,
+            target_e2e,
+            warmup=args.warmup,
+            repeat=args.repeat,
+            group=group,
+            local_rank=local_rank,
+        )
+        _report(rank, tokens, 'gemm_gather', metric)
+
+        if not args.skip_graph:
+            baseline_replay, baseline_graph_output, baseline_graph = _capture(baseline_e2e, group, local_rank)
+            target_replay, target_graph_output, target_graph = _capture(target_e2e, group, local_rank)
+            baseline_replay()
+            target_replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(target_graph_output, baseline_graph_output, rtol=0, atol=0)
+            metric = _measure_paired(
+                baseline_replay,
+                target_replay,
+                warmup=args.warmup,
+                repeat=args.repeat,
+                group=group,
+                local_rank=local_rank,
+            )
+            _report(rank, tokens, 'graph_gemm_gather', metric)
+            del baseline_replay, baseline_graph_output, baseline_graph
+            del target_replay, target_graph_output, target_graph
+            torch.cuda.synchronize()
+
+    if rank == 0:
+        print('PASS correctness=bitwise', flush=True)
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()
diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
index 261086a1fe..ae97bc9082 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -325,6 +325,7 @@ def create_state(
     max_tokens: int,
     hidden_size: int,
     device: torch.device | None = None,
+    comm_buff: torch.Tensor | None = None,
 ) -> MultimemAllGatherState:
     """Allocate and rendezvous the symmetric-memory buffer.
 
@@ -336,20 +337,16 @@ def create_state(
         f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
         f"bf16 for 16-byte multimem.st row alignment"
     )
-    device = device or torch.device(f"cuda:{torch.cuda.current_device()}")
-
-    # Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
-    pad_bytes = _MAX_BLOCKS * group.size() * 4
-    symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
-    with torch.inference_mode(False), torch.no_grad():
-        comm_buff = symm_mem.empty(
-            (max_tokens, hidden_size), dtype=torch.bfloat16, device=device
-        )
+    device = torch.device(device or torch.device(f"cuda:{torch.cuda.current_device()}"))
+    if device.type == 'cuda' and device.index is None:
+        device = torch.device('cuda', torch.cuda.current_device())
+
+    if comm_buff is None:
+        comm_buff = _allocate_symmetric_buffer(group, max_tokens, hidden_size, device)
+    elif (comm_buff.shape != (max_tokens, hidden_size) or comm_buff.dtype != torch.bfloat16
+          or comm_buff.device != device or not comm_buff.is_contiguous() or comm_buff.storage_offset() != 0):
+        raise ValueError('preallocated symmetric buffer does not match state')
     hdl = symm_mem.rendezvous(comm_buff, group=group)
-    assert hdl.rank == rank_in_group, (
-        f"symm_mem handle rank {hdl.rank} != rank_in_group {rank_in_group}; the "
-        f"hidden-shard offset would be wrong"
-    )
     return MultimemAllGatherState(
         group=group,
         rank_in_group=rank_in_group,
@@ -362,6 +359,20 @@ def create_state(
     )
 
 
+def _allocate_symmetric_buffer(
+    group: dist.ProcessGroup,
+    max_tokens: int,
+    hidden_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Perform only the local allocation half of state construction."""
+    # Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
+    pad_bytes = _MAX_BLOCKS * group.size() * 4
+    symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
+    with torch.inference_mode(False), torch.no_grad():
+        return symm_mem.empty((max_tokens, hidden_size), dtype=torch.bfloat16, device=device)
+
+
 def _launch_config(local_numel: int):
     assert local_numel % _NUMEL_PER_THREAD == 0
     return _MIN_BLOCKS, _BLOCK_THREADS, _BLOCK_THREADS // 32, _NUMEL_PER_THREAD
@@ -439,9 +450,11 @@ def all_gather_inner(
 class MultimemAllGatherer:
     """Guarded last-dim multimem all-gather with NCCL fallback.
 
-    Owns one symmetric buffer built lazily on the first eager call, and uses the kernel only when the input fits its
-    dtype/shape/alignment contract. The returned tensor owns its storage so a later collective cannot overwrite logits
-    that are still being consumed on another stream.
+    Owns one symmetric buffer and admits the input contract collectively on the
+    first eager call. Subsequent calls require the same dtype, layout and device
+    on every rank; token count may vary but must be TP-invariant. The returned
+    tensor owns its storage so a later collective cannot overwrite logits that
+    are still being consumed on another stream.
     """
 
     _UNINIT = object()
@@ -459,63 +472,131 @@ def __init__(
         self._rank = rank
         self._gathered_width = gathered_width
         self._max_tokens = int(max_tokens)
+        self._enabled = enabled
+        self._graph_ready = False
+        self._runtime_admitted = False
         # None => always NCCL; _UNINIT => build on first eager call.
         self._state = self._UNINIT if enabled else None
 
     def __call__(self, x: torch.Tensor) -> torch.Tensor | None:
         state = self._state
         if state is self._UNINIT:
-            state = self._build(x)
+            state = self._build(x.device)
             if state is not self._UNINIT:
                 self._state = state
-        if (
-            state is not None
-            and state is not self._UNINIT
-            and x.dtype == torch.bfloat16
-            and x.dim() == 2
-            and x.is_contiguous()
-            and 0 < x.shape[0] <= state.max_token_num
-            and x.data_ptr() % 16 == 0
-            and x.shape[-1] * state.world_size == state.hidden_dim
-        ):
-            return all_gather_inner(
-                state,
-                x,
-                tp_hidden_dim=self._gathered_width,
-                skip_entry_sync=False,
-                safe=True,
-            )
-        return None
-
-    def _build(self, x: torch.Tensor):
-        if x.dim() != 2 or x.dtype != torch.bfloat16:
+        if state is None or state is self._UNINIT:
+            return None
+
+        eligible = self._is_static_input_eligible(state, x)
+        capturing = torch.cuda.is_current_stream_capturing()
+        if not self._runtime_admitted:
+            # Admission uses host-visible consensus and cannot run under graph
+            # capture. In that uncommon first-call case every rank retains the
+            # existing NCCL path captured by ParallelLMHead.
+            if capturing:
+                return None
+            if not self.agree(eligible, state.device):
+                self._state = None
+                return None
+            self._runtime_admitted = True
+        elif not eligible:
+            raise RuntimeError('multimem all-gather input contract changed after TP-wide admission')
+
+        if not 0 < x.shape[0] <= state.max_token_num:
+            # ParallelLMHead presents identical token counts to all TP ranks,
+            # so this dynamic capacity fallback cannot split the collective.
             return None
+        # State allocation may be prepared before capture while Triton has not
+        # compiled this shape yet. Keep that first call on NCCL rather than
+        # trying to JIT-compile inside a CUDA Graph.
+        if capturing and not self._graph_ready:
+            return None
+        output = all_gather_inner(
+            state,
+            x,
+            tp_hidden_dim=self._gathered_width,
+            skip_entry_sync=False,
+            safe=True,
+        )
+        self._graph_ready = True
+        return output
+
+    @staticmethod
+    def _is_static_input_eligible(state: MultimemAllGatherState, x: torch.Tensor) -> bool:
+        """Check properties that remain stable for one LM-head instance."""
+        return (x.dtype == torch.bfloat16 and x.dim() == 2 and x.device == state.device and x.is_contiguous()
+                and x.data_ptr() % 16 == 0 and x.shape[-1] % _NUMEL_PER_THREAD == 0
+                and x.shape[-1] * state.world_size == state.hidden_dim)
+
+    def prepare(self, device: torch.device | str) -> bool:
+        """Collectively allocate and rendezvous the arena before graph capture."""
+        if self._state is self._UNINIT:
+            state = self._build(torch.device(device))
+            if state is not self._UNINIT:
+                self._state = state
+        return self._state is not None and self._state is not self._UNINIT
+
+    def agree(self, local_ready: bool, device: torch.device | str) -> bool:
+        """Return a TP-wide setup decision so ranks never split paths."""
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError('TP readiness consensus is not graph capturable')
+        ready = torch.tensor(int(local_ready), dtype=torch.int32, device=device)
+        dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=self._group)
+        return bool(ready.item())
+
+    def release(self) -> None:
+        """Drop the device arena so model offload can reclaim its storage."""
+        self._state = self._UNINIT if self._enabled else None
+        self._graph_ready = False
+        self._runtime_admitted = False
+
+    def _build(self, device: torch.device):
+        device = torch.device(device)
+        if device.type == 'cuda' and device.index is None:
+            device = torch.device('cuda', torch.cuda.current_device())
         if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
             # Can't allocate under capture; retry later.
             return self._UNINIT
-        if x.shape[-1] % _NUMEL_PER_THREAD != 0:
+        if self._gathered_width % _NUMEL_PER_THREAD != 0:
             return None
+        world_size = dist.get_world_size(self._group)
+        # tl.arange requires a power-of-two extent. Group size is identical on
+        # all ranks, so this fallback decision cannot split the protocol.
+        if world_size not in _SUPPORTED_WORLD_SIZES:
+            return None
+
+        # Allocate locally first, then make one TP-wide admission decision. No
+        # rank may rendezvous while a peer can still take a local fallback.
+        comm_buff = None
+        allocation_error = None
         try:
-            world_size = dist.get_world_size(self._group)
-            # tl.arange requires a power-of-two extent.  Keep unsupported TP
-            # layouts on the existing NCCL path rather than failing at JIT.
-            if world_size not in _SUPPORTED_WORLD_SIZES:
-                return None
-            state = create_state(
-                group=self._group,
-                rank_in_group=self._rank,
-                max_tokens=self._max_tokens,
-                hidden_size=self._gathered_width,
-            )
-            if state.symm_mem_hdl.multicast_ptr == 0:
-                # No multicast for this world size / arch; multimem.st would
-                # write nowhere. Fall back to NCCL.
+            comm_buff = _allocate_symmetric_buffer(self._group, self._max_tokens, self._gathered_width, device)
+        except Exception as exc:
+            allocation_error = exc
+        if not self.agree(comm_buff is not None, device):
+            if self._rank == 0:
                 logger.warning(
-                    'multimem all-gather disabled (no multicast for world_size=%d)',
-                    state.world_size,
+                    'multimem all-gather disabled because a TP rank could not allocate its symmetric arena%s',
+                    f': {allocation_error}' if allocation_error else '',
                 )
-                return None
-            return state
-        except Exception as e:
-            logger.warning('multimem all-gather disabled (%s)', e)
             return None
+        if comm_buff is None:
+            raise RuntimeError('TP admitted a missing symmetric-memory arena')
+
+        # Every rank is committed from this point. A rendezvous error must
+        # propagate; turning it into a rank-local fallback can deadlock peers.
+        state = create_state(
+            group=self._group,
+            rank_in_group=self._rank,
+            max_tokens=self._max_tokens,
+            hidden_size=self._gathered_width,
+            device=device,
+            comm_buff=comm_buff,
+        )
+        multicast_ready = state.symm_mem_hdl.rank == self._rank and state.symm_mem_hdl.multicast_ptr != 0
+        if not self.agree(multicast_ready, device):
+            if self._rank == 0:
+                logger.warning('multimem all-gather disabled (invalid TP-wide multicast handle for world_size=%d)',
+                               state.world_size)
+            return None
+        return state
diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py
index 5f749a7fa6..35d7f9c520 100644
--- a/lmdeploy/pytorch/engine/executor/ray_executor.py
+++ b/lmdeploy/pytorch/engine/executor/ray_executor.py
@@ -97,6 +97,19 @@ def _update_runtime_env_nsys(runtime_env: dict):
     return runtime_env
 
 
+def _needs_symm_mem_device_setup(dist_config: DistConfig) -> bool:
+    """Whether Ray workers must preserve peer CUDA-device visibility.
+
+    Ray normally narrows every actor to one visible GPU, which remaps that
+    device to local ordinal zero. CUDA symmetric-memory rendezvous identifies
+    allocations by the process-local device ordinal, so TP actors on one host
+    must instead inherit the full visibility and select their assigned GPU.
+    """
+    from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem
+    return (should_try_symm_mem(dist_config)
+            or (_envs.enable_symm_mem_lmhead and dist_config.attn_tp > 1))
+
+
 class RemoteLogger:
     """Remote logger."""
 
@@ -172,6 +185,12 @@ def set_assigned_cuda_device(self):
         local_rank = (visible_devices.split(',').index(physical_device_id)
                       if visible_devices else int(physical_device_id))
         self.set_device(local_rank)
+        return {
+            'node_ip': self.node_ip,
+            'ray_gpu_ids': [str(gpu_id) for gpu_id in gpu_ids],
+            'cuda_visible_devices': visible_devices,
+            'current_device': torch.cuda.current_device(),
+        }
 
     def set_env(self, envs: dict[str, str]):
         for key, value in envs.items():
@@ -253,10 +272,9 @@ def __init__(
 
         device_ctx = DeviceContext(device_type)
         with get_device_manager().context(device_ctx):
-            self._try_symm_mem = False
+            self._needs_symm_mem_device_setup = False
             if device_type == 'cuda':
-                from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem
-                self._try_symm_mem = should_try_symm_mem(dist_config)
+                self._needs_symm_mem_device_setup = _needs_symm_mem_device_setup(dist_config)
             logger.info('Init ray cluster.')
             attn_tp = dist_config.attn_tp
             self.ray_ctx = RayContext(attn_tp, dp=dist_config.dp, device_type=device_type)
@@ -698,7 +716,7 @@ def _init_workers_ray(self, placement_group: PlacementGroup, worker_kwargs: dict
             if device_str == 'GPU':
                 runtime_env = dict()
                 runtime_env = _update_runtime_envs(runtime_env)
-                if self._try_symm_mem:
+                if self._needs_symm_mem_device_setup:
                     # Symmetric-memory IPC needs peer TP GPUs to stay visible.
                     # Keep the inherited visibility and bind each actor below.
                     runtime_env['env_vars']['RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES'] = '1'
@@ -728,10 +746,20 @@ def _init_distributed_environment_by_device(self, device_str: str):
         driver_ip = _get_master_addr()
         if device_str == 'cuda':
             self.workers = self._sort_workers(driver_ip, self.workers)
-            if self._try_symm_mem:
+            if self._needs_symm_mem_device_setup:
                 # Ray did not narrow CUDA visibility, so select each actor's
                 # placement-group assignment before distributed initialization.
-                ray.get([worker.set_assigned_cuda_device.remote() for worker in self.workers])
+                bindings = ray.get([
+                    worker.set_assigned_cuda_device.remote()
+                    for worker in self.workers
+                ])
+                device_keys = {
+                    (binding['node_ip'], binding['current_device'])
+                    for binding in bindings
+                }
+                if len(device_keys) != len(bindings):
+                    raise RuntimeError('Ray symmetric-memory workers must bind unique CUDA '
+                                       f'devices on each node, got: {bindings}')
 
         elif device_str == 'ascend':
             self._init_ascend_distributed_environment(driver_ip)
diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py
index f965b2a843..68161eafdd 100644
--- a/lmdeploy/pytorch/nn/embedding.py
+++ b/lmdeploy/pytorch/nn/embedding.py
@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import logging
+
 import torch
 import torch.distributed as dist
 from torch import nn
@@ -9,6 +11,23 @@
 from lmdeploy.pytorch.weight_loader.model_weight_loader import default_weight_loader
 
 DEFAULT_VOCAB_PADDING_SIZE = 64
+logger = logging.getLogger(__name__)
+
+
+def _tp_agree(local_ready: bool, device: torch.device, group: dist.ProcessGroup) -> bool:
+    """Resolve an optional LM-head provider decision on every TP rank."""
+    ready = torch.tensor(int(local_ready), dtype=torch.int32, device=device)
+    dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=group)
+    return bool(ready.item())
+
+
+def _tp_same_config(values: tuple[int, ...], device: torch.device, group: dist.ProcessGroup) -> bool:
+    """Return whether every TP rank supplied the same integer config."""
+    lower = torch.tensor(values, dtype=torch.int64, device=device)
+    upper = lower.clone()
+    dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=group)
+    dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=group)
+    return bool(torch.equal(lower, upper))
 
 
 def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
@@ -146,25 +165,65 @@ def __init__(
         self.impl = builder.build(hidden_size, self.vocab_size_padded, bias, dtype=dtype)
 
         self._symm_mem_gatherer = None
-        if self.all_reduce and _envs.enable_symm_mem_lmhead:
+        self._symm_mem_device = self.weight.device
+        self._symm_mem_dtype = self.weight.dtype
+        if self.all_reduce and self.weight.device.type == 'cuda':
+            device = self.weight.device
+            if device.index is None:
+                device = torch.device('cuda', torch.cuda.current_device())
+            requested = _envs.enable_symm_mem_lmhead and self.weight.dtype == torch.bfloat16
+            if not _tp_agree(requested, device, self.tp_group):
+                return
+
+            gathered_width = self.tp * self.vocab_size_padded
+            capacity = _envs.symm_mem_lmhead_max_mb * 1024 * 1024
+            max_tokens = capacity // (gathered_width * torch.bfloat16.itemsize)
+            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
+
+            gatherer_cls = None
             try:
                 from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer
-
-                gathered_width = self.tp * self.vocab_size_padded
-                capacity = _envs.symm_mem_lmhead_max_mb * 1024 * 1024
-                max_tokens = capacity // (gathered_width * torch.bfloat16.itemsize)
-                if max_tokens > 0:
-                    self._symm_mem_gatherer = MultimemAllGatherer(group=self.tp_group,
-                                                                  rank=self.tp_rank,
-                                                                  gathered_width=gathered_width,
-                                                                  max_tokens=max_tokens)
-            except ImportError:
-                pass
+                gatherer_cls = MultimemAllGatherer
+            except ImportError as exc:
+                if self.tp_rank == 0:
+                    logger.warning('symmetric-memory LM-head unavailable: %s', exc)
+            if not _tp_agree(gatherer_cls is not None, device, self.tp_group):
+                return
+
+            gatherer = gatherer_cls(group=self.tp_group,
+                                    rank=self.tp_rank,
+                                    gathered_width=gathered_width,
+                                    max_tokens=max_tokens)
+            if gatherer.prepare(device):
+                self._symm_mem_gatherer = gatherer
 
     def tie_weights(self, embedding: ParallelEmbedding):
         """Tie the local LM-head shard to a parallel embedding shard."""
         self.weight = embedding.weight
 
+    def _apply(self, fn, recurse=True):
+        """Keep the symmetric arena aligned with model device moves."""
+        previous_device = self._symm_mem_device
+        previous_dtype = self._symm_mem_dtype
+        result = super()._apply(fn, recurse=recurse)
+        current_device = self.weight.device
+        current_dtype = self.weight.dtype
+        gatherer = self._symm_mem_gatherer
+        if gatherer is not None and (current_device != previous_device or current_dtype != previous_dtype):
+            if previous_device.type == 'cuda':
+                gatherer.release()
+            self._symm_mem_device = current_device
+            self._symm_mem_dtype = current_dtype
+            if current_dtype != torch.bfloat16:
+                self._symm_mem_gatherer = None
+            elif current_device.type == 'cuda' and not gatherer.prepare(current_device):
+                self._symm_mem_gatherer = None
+        return result
+
     def get_local_logits(self, hidden_states: torch.Tensor):
         """Compute logits for the vocabulary shard owned by this rank."""
         if hidden_states.dtype != self.weight.dtype:
diff --git a/tests/pytorch/engine/test_ray_executor_symm_mem.py b/tests/pytorch/engine/test_ray_executor_symm_mem.py
new file mode 100644
index 0000000000..429cca6caf
--- /dev/null
+++ b/tests/pytorch/engine/test_ray_executor_symm_mem.py
@@ -0,0 +1,127 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+@pytest.fixture(scope='module')
+def ray_executor_module():
+    pytest.importorskip('ray')
+    from lmdeploy.pytorch.engine.executor import ray_executor
+    return ray_executor
+
+
+def _dist_config(*, attn_tp=8):
+    return SimpleNamespace(dp=1, ep=1, attn_tp=attn_tp, enable_microbatch=False)
+
+
+@pytest.mark.parametrize(
+    ('enable_allreduce', 'enable_lmhead', 'attn_tp', 'expected'),
+    [
+        (False, False, 8, False),
+        (True, False, 8, True),
+        (False, True, 8, True),
+        (False, True, 1, False),
+    ],
+)
+def test_needs_symm_mem_device_setup(monkeypatch, ray_executor_module, enable_allreduce, enable_lmhead, attn_tp,
+                                     expected):
+    monkeypatch.setattr(ray_executor_module._envs, 'enable_symm_mem_allreduce', enable_allreduce)
+    monkeypatch.setattr(ray_executor_module._envs, 'enable_symm_mem_lmhead', enable_lmhead)
+
+    actual = ray_executor_module._needs_symm_mem_device_setup(_dist_config(attn_tp=attn_tp))
+
+    assert actual is expected
+
+
+class _RemoteMethod:
+
+    def __init__(self, result):
+        self.remote = Mock(return_value=result)
+
+
+@pytest.mark.parametrize('required', [False, True])
+def test_ray_worker_runtime_env_tracks_device_setup(monkeypatch, ray_executor_module, required):
+    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
+    executor._needs_symm_mem_device_setup = required
+    executor.dist_config = _dist_config(attn_tp=2)
+
+    placement_group = SimpleNamespace(bundle_specs=[{'GPU': 1}, {'GPU': 1}])
+    remote_options = []
+
+    class _RemoteActor:
+
+        def remote(self, **kwargs):
+            return object()
+
+    def fake_remote(**options):
+        remote_options.append(options)
+        return lambda actor_cls: _RemoteActor()
+
+    monkeypatch.setattr(ray_executor_module, 'get_device_str', lambda: 'GPU')
+    monkeypatch.setattr(ray_executor_module, '_update_runtime_envs', lambda _: {'env_vars': {}})
+    monkeypatch.setattr(ray_executor_module, 'PlacementGroupSchedulingStrategy', lambda **kwargs: kwargs)
+    monkeypatch.setattr(ray_executor_module.ray, 'remote', fake_remote)
+    monkeypatch.setattr(ray_executor_module._envs, 'ray_external_pg_bundles', [])
+    monkeypatch.setattr(ray_executor_module._envs, 'ray_nsys_enable', False)
+
+    workers = executor._init_workers_ray(placement_group, worker_kwargs={})
+
+    assert len(workers) == 2
+    assert len(remote_options) == 2
+    for options in remote_options:
+        actual = options['runtime_env']['env_vars'].get('RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES')
+        assert actual == ('1' if required else None)
+
+
+@pytest.mark.parametrize('required', [False, True])
+def test_ray_device_binding_tracks_device_setup(monkeypatch, ray_executor_module, required):
+    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
+    executor._needs_symm_mem_device_setup = required
+    reports = [{'node_ip': '127.0.0.1', 'current_device': rank} for rank in range(2)]
+    workers = [SimpleNamespace(set_assigned_cuda_device=_RemoteMethod(report)) for report in reports]
+    executor.workers = workers
+    executor._sort_workers = Mock(return_value=workers)
+    ray_get = Mock(return_value=reports)
+    monkeypatch.setattr(ray_executor_module, '_get_master_addr', lambda: '127.0.0.1')
+    monkeypatch.setattr(ray_executor_module.ray, 'get', ray_get)
+
+    executor._init_distributed_environment_by_device('cuda')
+
+    for worker in workers:
+        assert worker.set_assigned_cuda_device.remote.call_count == int(required)
+    assert ray_get.call_count == int(required)
+
+
+def test_ray_device_binding_rejects_duplicate_local_ordinals(monkeypatch, ray_executor_module):
+    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
+    executor._needs_symm_mem_device_setup = True
+    reports = [
+        {'node_ip': '127.0.0.1', 'current_device': 0},
+        {'node_ip': '127.0.0.1', 'current_device': 0},
+    ]
+    workers = [SimpleNamespace(set_assigned_cuda_device=_RemoteMethod(report)) for report in reports]
+    executor.workers = workers
+    executor._sort_workers = Mock(return_value=workers)
+    monkeypatch.setattr(ray_executor_module, '_get_master_addr', lambda: '127.0.0.1')
+    monkeypatch.setattr(ray_executor_module.ray, 'get', Mock(return_value=reports))
+
+    with pytest.raises(RuntimeError, match='must bind unique CUDA devices'):
+        executor._init_distributed_environment_by_device('cuda')
+
+
+def test_lmhead_only_does_not_enable_symm_mem_allreduce(monkeypatch, ray_executor_module):
+    from lmdeploy.pytorch.backends.cuda.comm import communicator
+
+    monkeypatch.setattr(communicator._envs, 'enable_flashinfer_allreduce', False)
+    monkeypatch.setattr(communicator._envs, 'enable_symm_mem_allreduce', False)
+    monkeypatch.setattr(communicator._envs, 'enable_symm_mem_lmhead', True)
+
+    actual = communicator.build_cuda_communicator(
+        cpu_group=object(),
+        device_group=object(),
+        dist_config=_dist_config(attn_tp=8),
+    )
+
+    assert actual is None
diff --git a/tests/pytorch/nn/test_lm_head_symm_mem.py b/tests/pytorch/nn/test_lm_head_symm_mem.py
new file mode 100644
index 0000000000..757becef4c
--- /dev/null
+++ b/tests/pytorch/nn/test_lm_head_symm_mem.py
@@ -0,0 +1,25 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest.mock import Mock
+
+import torch
+from torch import nn
+
+from lmdeploy.pytorch.nn.embedding import ParallelLMHead
+
+
+def test_lm_head_apply_releases_arena_when_tied_weight_already_moved():
+    head = ParallelLMHead.__new__(ParallelLMHead)
+    nn.Module.__init__(head)
+    head.register_parameter('weight', nn.Parameter(torch.empty(1), requires_grad=False))
+    gatherer = Mock()
+    head._symm_mem_gatherer = gatherer
+    # A tied embedding can move the shared Parameter before LM-head._apply.
+    head._symm_mem_device = torch.device('cuda:0')
+    head._symm_mem_dtype = head.weight.dtype
+
+    result = head._apply(lambda tensor: tensor)
+
+    assert result is head
+    gatherer.release.assert_called_once_with()
+    gatherer.prepare.assert_not_called()
+    assert head._symm_mem_device == torch.device('cpu')
diff --git a/tests/pytorch/test_symm_mem_allgather.py b/tests/pytorch/test_symm_mem_allgather.py
new file mode 100644
index 0000000000..5933584044
--- /dev/null
+++ b/tests/pytorch/test_symm_mem_allgather.py
@@ -0,0 +1,132 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+comm = pytest.importorskip('lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather')
+
+
+def _gatherer():
+    return comm.MultimemAllGatherer(group=object(), rank=0, gathered_width=16, max_tokens=4)
+
+
+def _disable_cuda_capture_probe(monkeypatch):
+    monkeypatch.setattr(comm.torch.cuda, 'is_available', lambda: False)
+    monkeypatch.setattr(comm.dist, 'get_world_size', lambda group: 2)
+
+
+def test_prepare_rejects_before_rendezvous_when_any_rank_cannot_allocate(monkeypatch):
+    _disable_cuda_capture_probe(monkeypatch)
+    gatherer = _gatherer()
+    events = []
+
+    def _allocate(*args, **kwargs):
+        events.append('allocate')
+        raise RuntimeError('local allocation failed')
+
+    def _agree(local_ready, device):
+        events.append(('agree', local_ready, device))
+        return False
+
+    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', _allocate)
+    monkeypatch.setattr(comm, 'create_state', lambda *args, **kwargs: pytest.fail('rendezvous must not be entered'))
+    monkeypatch.setattr(gatherer, 'agree', _agree)
+
+    assert gatherer.prepare(torch.device('cuda:3')) is False
+    assert events == ['allocate', ('agree', False, torch.device('cuda:3'))]
+    assert gatherer._state is None
+
+
+def test_prepare_agrees_before_rendezvous_and_preserves_device(monkeypatch):
+    _disable_cuda_capture_probe(monkeypatch)
+    gatherer = _gatherer()
+    events = []
+    arena = torch.empty(0)
+
+    def _allocate(group, max_tokens, hidden_size, device):
+        events.append(('allocate', device))
+        return arena
+
+    def _agree(local_ready, device):
+        events.append(('agree', local_ready, device))
+        return True
+
+    def _create_state(**kwargs):
+        events.append(('rendezvous', kwargs['device'], kwargs['comm_buff']))
+        return SimpleNamespace(
+            symm_mem_hdl=SimpleNamespace(multicast_ptr=1, rank=0),
+            world_size=2,
+            max_token_num=4,
+        )
+
+    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', _allocate)
+    monkeypatch.setattr(comm, 'create_state', _create_state)
+    monkeypatch.setattr(gatherer, 'agree', _agree)
+
+    assert gatherer.prepare(torch.device('cuda:3')) is True
+    assert events == [
+        ('allocate', torch.device('cuda:3')),
+        ('agree', True, torch.device('cuda:3')),
+        ('rendezvous', torch.device('cuda:3'), arena),
+        ('agree', True, torch.device('cuda:3')),
+    ]
+
+
+def test_prepare_does_not_turn_rendezvous_failure_into_local_fallback(monkeypatch):
+    _disable_cuda_capture_probe(monkeypatch)
+    gatherer = _gatherer()
+    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', lambda *args, **kwargs: torch.empty(0))
+    monkeypatch.setattr(gatherer, 'agree', lambda local_ready, device: True)
+
+    def _rendezvous_failure(**kwargs):
+        raise RuntimeError('collective rendezvous failed')
+
+    monkeypatch.setattr(comm, 'create_state', _rendezvous_failure)
+
+    with pytest.raises(RuntimeError, match='collective rendezvous failed'):
+        gatherer.prepare(torch.device('cuda:3'))
+    assert gatherer._state is gatherer._UNINIT
+
+
+def test_release_drops_arena_and_resets_graph_admission():
+    gatherer = _gatherer()
+    gatherer._state = object()
+    gatherer._graph_ready = True
+    gatherer._runtime_admitted = True
+
+    gatherer.release()
+
+    assert gatherer._state is gatherer._UNINIT
+    assert gatherer._graph_ready is False
+    assert gatherer._runtime_admitted is False
+
+
+def test_first_call_collectively_admits_static_input_contract(monkeypatch):
+    gatherer = _gatherer()
+    gatherer._state = SimpleNamespace(
+        max_token_num=4,
+        world_size=2,
+        hidden_dim=16,
+        device=torch.device('cpu'),
+    )
+    admission = Mock(return_value=True)
+    monkeypatch.setattr(gatherer, 'agree', admission)
+    monkeypatch.setattr(comm.torch.cuda, 'is_current_stream_capturing', lambda: False)
+    monkeypatch.setattr(comm, 'all_gather_inner', lambda state, x, **kwargs: x)
+    value = torch.empty((1, 8), dtype=torch.bfloat16)
+
+    assert gatherer(value) is value
+    assert gatherer(value) is value
+    admission.assert_called_once_with(True, torch.device('cpu'))
+    assert gatherer._runtime_admitted is True
+
+
+def test_direct_path_rejects_unaligned_local_width_before_kernel():
+    gatherer = _gatherer()
+    gatherer._state = SimpleNamespace(max_token_num=4, world_size=2, hidden_dim=8, device=torch.device('cpu'))
+    gatherer.agree = lambda local_ready, device: local_ready
+
+    assert gatherer(torch.empty((1, 4), dtype=torch.bfloat16)) is None
+    assert gatherer._state is None

From ab66c3131dc73554b0fc2651b77f0fc27a699fab Mon Sep 17 00:00:00 2001
From: qescccczmr <bryentonxemmrj@gmail.com>
Date: Thu, 3 Sep 2026 04:23:37 +0000
Subject: [PATCH 3/6] style: apply docformatter

---
 .../pytorch/backends/cuda/comm/symm_mem_allgather.py | 12 ++++++------
 lmdeploy/pytorch/engine/executor/ray_executor.py     |  7 +++----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
index ae97bc9082..fa4660de1c 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -450,11 +450,10 @@ def all_gather_inner(
 class MultimemAllGatherer:
     """Guarded last-dim multimem all-gather with NCCL fallback.
 
-    Owns one symmetric buffer and admits the input contract collectively on the
-    first eager call. Subsequent calls require the same dtype, layout and device
-    on every rank; token count may vary but must be TP-invariant. The returned
-    tensor owns its storage so a later collective cannot overwrite logits that
-    are still being consumed on another stream.
+    Owns one symmetric buffer and admits the input contract collectively on the first eager call. Subsequent calls
+    require the same dtype, layout and device on every rank; token count may vary but must be TP-invariant. The returned
+    tensor owns its storage so a later collective cannot overwrite logits that are still being consumed on another
+    stream.
     """
 
     _UNINIT = object()
@@ -529,7 +528,8 @@ def _is_static_input_eligible(state: MultimemAllGatherState, x: torch.Tensor) ->
                 and x.shape[-1] * state.world_size == state.hidden_dim)
 
     def prepare(self, device: torch.device | str) -> bool:
-        """Collectively allocate and rendezvous the arena before graph capture."""
+        """Collectively allocate and rendezvous the arena before graph
+        capture."""
         if self._state is self._UNINIT:
             state = self._build(torch.device(device))
             if state is not self._UNINIT:
diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py
index 35d7f9c520..fd4443abb5 100644
--- a/lmdeploy/pytorch/engine/executor/ray_executor.py
+++ b/lmdeploy/pytorch/engine/executor/ray_executor.py
@@ -100,10 +100,9 @@ def _update_runtime_env_nsys(runtime_env: dict):
 def _needs_symm_mem_device_setup(dist_config: DistConfig) -> bool:
     """Whether Ray workers must preserve peer CUDA-device visibility.
 
-    Ray normally narrows every actor to one visible GPU, which remaps that
-    device to local ordinal zero. CUDA symmetric-memory rendezvous identifies
-    allocations by the process-local device ordinal, so TP actors on one host
-    must instead inherit the full visibility and select their assigned GPU.
+    Ray normally narrows every actor to one visible GPU, which remaps that device to local ordinal zero. CUDA symmetric-
+    memory rendezvous identifies allocations by the process-local device ordinal, so TP actors on one host must instead
+    inherit the full visibility and select their assigned GPU.
     """
     from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem
     return (should_try_symm_mem(dist_config)

From 6712d63fc84805467ed95713a76e1c4d04fb98d8 Mon Sep 17 00:00:00 2001
From: qescccczmr <bryentonxemmrj@gmail.com>
Date: Thu, 3 Sep 2026 05:02:30 +0000
Subject: [PATCH 4/6] chore: remove benchmark and test files

---
 benchmark/profile_lmhead_allgather.py         | 227 ------------------
 .../engine/test_ray_executor_symm_mem.py      | 127 ----------
 tests/pytorch/nn/test_lm_head_symm_mem.py     |  25 --
 tests/pytorch/test_symm_mem_allgather.py      | 132 ----------
 4 files changed, 511 deletions(-)
 delete mode 100644 benchmark/profile_lmhead_allgather.py
 delete mode 100644 tests/pytorch/engine/test_ray_executor_symm_mem.py
 delete mode 100644 tests/pytorch/nn/test_lm_head_symm_mem.py
 delete mode 100644 tests/pytorch/test_symm_mem_allgather.py

diff --git a/benchmark/profile_lmhead_allgather.py b/benchmark/profile_lmhead_allgather.py
deleted file mode 100644
index 7b7a212faf..0000000000
--- a/benchmark/profile_lmhead_allgather.py
+++ /dev/null
@@ -1,227 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-"""Compare the production NCCL LM-head gather with symmetric memory.
-
-Example:
-
-.. code-block:: bash
-
-   torchrun --standalone --nproc-per-node=8 \
-     benchmark/profile_lmhead_allgather.py \
-     --tokens 1 8 32 128 --warmup 20 --repeat 200
-
-The benchmark reports the maximum CUDA-event latency across TP ranks. V1
-includes the owning clone performed by ``MultimemAllGatherer``.
-"""
-
-import argparse
-import math
-import os
-import statistics
-from collections.abc import Callable
-
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-
-from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer
-
-
-def _nccl_gather(local: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
-    """Match ``ParallelLMHead`` allocation, gather and vocabulary layout."""
-    world_size = dist.get_world_size(group)
-    output = local.new_empty((world_size, ) + tuple(local.shape))
-    dist.all_gather_into_tensor(output, local, group=group)
-    return output.movedim(0, 1).reshape(local.shape[0], -1)
-
-
-def _percentile(samples: list[float], q: float) -> float:
-    index = max(0, min(len(samples) - 1, math.ceil(q * len(samples)) - 1))
-    return sorted(samples)[index]
-
-
-def _measure_paired(
-    baseline: Callable[[], torch.Tensor],
-    target: Callable[[], torch.Tensor],
-    *,
-    warmup: int,
-    repeat: int,
-    group: dist.ProcessGroup,
-    local_rank: int,
-) -> dict[str, float]:
-    """Interleave both paths and reduce every sample to the slowest rank."""
-    for _ in range(warmup):
-        baseline()
-        target()
-    torch.cuda.synchronize()
-    dist.barrier(group=group, device_ids=[local_rank])
-
-    starts = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
-    ends = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
-    rank_times = torch.empty(2, dtype=torch.float64, device='cuda')
-    baseline_samples = []
-    target_samples = []
-    for iteration in range(repeat):
-        dist.barrier(group=group, device_ids=[local_rank])
-        functions = (baseline, target) if iteration % 2 == 0 else (target, baseline)
-        outputs = []
-        for index, function in enumerate(functions):
-            starts[index].record()
-            outputs.append(function())
-            ends[index].record()
-        ends[-1].synchronize()
-        elapsed = [starts[index].elapsed_time(ends[index]) * 1000.0 for index in range(2)]
-        if iteration % 2:
-            elapsed.reverse()
-        rank_times[0] = elapsed[0]
-        rank_times[1] = elapsed[1]
-        dist.all_reduce(rank_times, op=dist.ReduceOp.MAX, group=group)
-        baseline_us, target_us = rank_times.tolist()
-        baseline_samples.append(baseline_us)
-        target_samples.append(target_us)
-        del outputs
-
-    baseline_median = statistics.median(baseline_samples)
-    target_median = statistics.median(target_samples)
-    return {
-        'baseline_median_us': baseline_median,
-        'target_median_us': target_median,
-        'baseline_p95_us': _percentile(baseline_samples, 0.95),
-        'target_p95_us': _percentile(target_samples, 0.95),
-        'speedup': baseline_median / target_median,
-    }
-
-
-def _capture(function: Callable[[], torch.Tensor], group: dist.ProcessGroup, local_rank: int):
-    for _ in range(3):
-        function()
-    torch.cuda.synchronize()
-    dist.barrier(group=group, device_ids=[local_rank])
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        output = function()
-    torch.cuda.synchronize()
-    dist.barrier(group=group, device_ids=[local_rank])
-    return graph.replay, output, graph
-
-
-def _report(rank: int, tokens: int, phase: str, metric: dict[str, float]) -> None:
-    if rank != 0:
-        return
-    print(
-        f'RESULT tokens={tokens} phase={phase} '
-        f'base_median_us={metric["baseline_median_us"]:.3f} '
-        f'v1_median_us={metric["target_median_us"]:.3f} '
-        f'base_p95_us={metric["baseline_p95_us"]:.3f} '
-        f'v1_p95_us={metric["target_p95_us"]:.3f} '
-        f'speedup={metric["speedup"]:.4f}x',
-        flush=True,
-    )
-
-
-@torch.inference_mode()
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('--vocab-size', type=int, default=154880)
-    parser.add_argument('--hidden-size', type=int, default=4096)
-    parser.add_argument('--tokens', type=int, nargs='+', default=[1, 8, 32, 128])
-    parser.add_argument('--warmup', type=int, default=20)
-    parser.add_argument('--repeat', type=int, default=200)
-    parser.add_argument('--skip-graph', action='store_true')
-    args = parser.parse_args()
-
-    local_rank = int(os.environ['LOCAL_RANK'])
-    torch.cuda.set_device(local_rank)
-    dist.init_process_group('nccl')
-    group = dist.group.WORLD
-    rank = dist.get_rank(group)
-    world_size = dist.get_world_size(group)
-    if args.vocab_size % world_size:
-        raise ValueError(f'vocab-size {args.vocab_size} must be divisible by TP={world_size}')
-
-    local_vocab = args.vocab_size // world_size
-    gatherer = MultimemAllGatherer(group, rank, args.vocab_size, max(args.tokens))
-    if not gatherer.prepare(torch.device('cuda', local_rank)):
-        raise RuntimeError('symmetric-memory multicast is unavailable for this TP topology')
-
-    torch.manual_seed(1701 + rank)
-    weight = torch.randn((local_vocab, args.hidden_size), dtype=torch.bfloat16, device='cuda')
-    if rank == 0:
-        print(
-            f'CONFIG gpu={torch.cuda.get_device_name(local_rank)!r} tp={world_size} dtype=bf16 '
-            f'hidden={args.hidden_size} global_vocab={args.vocab_size} local_vocab={local_vocab} '
-            f'warmup={args.warmup} repeat={args.repeat}',
-            flush=True,
-        )
-
-    for tokens in args.tokens:
-        torch.manual_seed(2026 + rank * 17 + tokens)
-        local_logits = torch.randn((tokens, local_vocab), dtype=torch.bfloat16, device='cuda')
-        baseline_output = _nccl_gather(local_logits, group)
-        target_output = gatherer(local_logits)
-        if target_output is None:
-            raise RuntimeError(f'symmetric-memory path rejected tokens={tokens}')
-        torch.testing.assert_close(target_output, baseline_output, rtol=0, atol=0)
-
-        def baseline():
-            return _nccl_gather(local_logits, group)
-
-        def target():
-            return gatherer(local_logits)
-
-        metric = _measure_paired(
-            baseline,
-            target,
-            warmup=args.warmup,
-            repeat=args.repeat,
-            group=group,
-            local_rank=local_rank,
-        )
-        _report(rank, tokens, 'gather', metric)
-
-        torch.manual_seed(4096 + tokens)
-        hidden = torch.randn((tokens, args.hidden_size), dtype=torch.bfloat16, device='cuda')
-
-        def baseline_e2e():
-            return _nccl_gather(F.linear(hidden, weight), group)
-
-        def target_e2e():
-            return gatherer(F.linear(hidden, weight))
-
-        torch.testing.assert_close(target_e2e(), baseline_e2e(), rtol=0, atol=0)
-        metric = _measure_paired(
-            baseline_e2e,
-            target_e2e,
-            warmup=args.warmup,
-            repeat=args.repeat,
-            group=group,
-            local_rank=local_rank,
-        )
-        _report(rank, tokens, 'gemm_gather', metric)
-
-        if not args.skip_graph:
-            baseline_replay, baseline_graph_output, baseline_graph = _capture(baseline_e2e, group, local_rank)
-            target_replay, target_graph_output, target_graph = _capture(target_e2e, group, local_rank)
-            baseline_replay()
-            target_replay()
-            torch.cuda.synchronize()
-            torch.testing.assert_close(target_graph_output, baseline_graph_output, rtol=0, atol=0)
-            metric = _measure_paired(
-                baseline_replay,
-                target_replay,
-                warmup=args.warmup,
-                repeat=args.repeat,
-                group=group,
-                local_rank=local_rank,
-            )
-            _report(rank, tokens, 'graph_gemm_gather', metric)
-            del baseline_replay, baseline_graph_output, baseline_graph
-            del target_replay, target_graph_output, target_graph
-            torch.cuda.synchronize()
-
-    if rank == 0:
-        print('PASS correctness=bitwise', flush=True)
-    dist.destroy_process_group()
-
-
-if __name__ == '__main__':
-    main()
diff --git a/tests/pytorch/engine/test_ray_executor_symm_mem.py b/tests/pytorch/engine/test_ray_executor_symm_mem.py
deleted file mode 100644
index 429cca6caf..0000000000
--- a/tests/pytorch/engine/test_ray_executor_symm_mem.py
+++ /dev/null
@@ -1,127 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-import pytest
-
-
-@pytest.fixture(scope='module')
-def ray_executor_module():
-    pytest.importorskip('ray')
-    from lmdeploy.pytorch.engine.executor import ray_executor
-    return ray_executor
-
-
-def _dist_config(*, attn_tp=8):
-    return SimpleNamespace(dp=1, ep=1, attn_tp=attn_tp, enable_microbatch=False)
-
-
-@pytest.mark.parametrize(
-    ('enable_allreduce', 'enable_lmhead', 'attn_tp', 'expected'),
-    [
-        (False, False, 8, False),
-        (True, False, 8, True),
-        (False, True, 8, True),
-        (False, True, 1, False),
-    ],
-)
-def test_needs_symm_mem_device_setup(monkeypatch, ray_executor_module, enable_allreduce, enable_lmhead, attn_tp,
-                                     expected):
-    monkeypatch.setattr(ray_executor_module._envs, 'enable_symm_mem_allreduce', enable_allreduce)
-    monkeypatch.setattr(ray_executor_module._envs, 'enable_symm_mem_lmhead', enable_lmhead)
-
-    actual = ray_executor_module._needs_symm_mem_device_setup(_dist_config(attn_tp=attn_tp))
-
-    assert actual is expected
-
-
-class _RemoteMethod:
-
-    def __init__(self, result):
-        self.remote = Mock(return_value=result)
-
-
-@pytest.mark.parametrize('required', [False, True])
-def test_ray_worker_runtime_env_tracks_device_setup(monkeypatch, ray_executor_module, required):
-    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
-    executor._needs_symm_mem_device_setup = required
-    executor.dist_config = _dist_config(attn_tp=2)
-
-    placement_group = SimpleNamespace(bundle_specs=[{'GPU': 1}, {'GPU': 1}])
-    remote_options = []
-
-    class _RemoteActor:
-
-        def remote(self, **kwargs):
-            return object()
-
-    def fake_remote(**options):
-        remote_options.append(options)
-        return lambda actor_cls: _RemoteActor()
-
-    monkeypatch.setattr(ray_executor_module, 'get_device_str', lambda: 'GPU')
-    monkeypatch.setattr(ray_executor_module, '_update_runtime_envs', lambda _: {'env_vars': {}})
-    monkeypatch.setattr(ray_executor_module, 'PlacementGroupSchedulingStrategy', lambda **kwargs: kwargs)
-    monkeypatch.setattr(ray_executor_module.ray, 'remote', fake_remote)
-    monkeypatch.setattr(ray_executor_module._envs, 'ray_external_pg_bundles', [])
-    monkeypatch.setattr(ray_executor_module._envs, 'ray_nsys_enable', False)
-
-    workers = executor._init_workers_ray(placement_group, worker_kwargs={})
-
-    assert len(workers) == 2
-    assert len(remote_options) == 2
-    for options in remote_options:
-        actual = options['runtime_env']['env_vars'].get('RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES')
-        assert actual == ('1' if required else None)
-
-
-@pytest.mark.parametrize('required', [False, True])
-def test_ray_device_binding_tracks_device_setup(monkeypatch, ray_executor_module, required):
-    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
-    executor._needs_symm_mem_device_setup = required
-    reports = [{'node_ip': '127.0.0.1', 'current_device': rank} for rank in range(2)]
-    workers = [SimpleNamespace(set_assigned_cuda_device=_RemoteMethod(report)) for report in reports]
-    executor.workers = workers
-    executor._sort_workers = Mock(return_value=workers)
-    ray_get = Mock(return_value=reports)
-    monkeypatch.setattr(ray_executor_module, '_get_master_addr', lambda: '127.0.0.1')
-    monkeypatch.setattr(ray_executor_module.ray, 'get', ray_get)
-
-    executor._init_distributed_environment_by_device('cuda')
-
-    for worker in workers:
-        assert worker.set_assigned_cuda_device.remote.call_count == int(required)
-    assert ray_get.call_count == int(required)
-
-
-def test_ray_device_binding_rejects_duplicate_local_ordinals(monkeypatch, ray_executor_module):
-    executor = ray_executor_module.RayExecutor.__new__(ray_executor_module.RayExecutor)
-    executor._needs_symm_mem_device_setup = True
-    reports = [
-        {'node_ip': '127.0.0.1', 'current_device': 0},
-        {'node_ip': '127.0.0.1', 'current_device': 0},
-    ]
-    workers = [SimpleNamespace(set_assigned_cuda_device=_RemoteMethod(report)) for report in reports]
-    executor.workers = workers
-    executor._sort_workers = Mock(return_value=workers)
-    monkeypatch.setattr(ray_executor_module, '_get_master_addr', lambda: '127.0.0.1')
-    monkeypatch.setattr(ray_executor_module.ray, 'get', Mock(return_value=reports))
-
-    with pytest.raises(RuntimeError, match='must bind unique CUDA devices'):
-        executor._init_distributed_environment_by_device('cuda')
-
-
-def test_lmhead_only_does_not_enable_symm_mem_allreduce(monkeypatch, ray_executor_module):
-    from lmdeploy.pytorch.backends.cuda.comm import communicator
-
-    monkeypatch.setattr(communicator._envs, 'enable_flashinfer_allreduce', False)
-    monkeypatch.setattr(communicator._envs, 'enable_symm_mem_allreduce', False)
-    monkeypatch.setattr(communicator._envs, 'enable_symm_mem_lmhead', True)
-
-    actual = communicator.build_cuda_communicator(
-        cpu_group=object(),
-        device_group=object(),
-        dist_config=_dist_config(attn_tp=8),
-    )
-
-    assert actual is None
diff --git a/tests/pytorch/nn/test_lm_head_symm_mem.py b/tests/pytorch/nn/test_lm_head_symm_mem.py
deleted file mode 100644
index 757becef4c..0000000000
--- a/tests/pytorch/nn/test_lm_head_symm_mem.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from unittest.mock import Mock
-
-import torch
-from torch import nn
-
-from lmdeploy.pytorch.nn.embedding import ParallelLMHead
-
-
-def test_lm_head_apply_releases_arena_when_tied_weight_already_moved():
-    head = ParallelLMHead.__new__(ParallelLMHead)
-    nn.Module.__init__(head)
-    head.register_parameter('weight', nn.Parameter(torch.empty(1), requires_grad=False))
-    gatherer = Mock()
-    head._symm_mem_gatherer = gatherer
-    # A tied embedding can move the shared Parameter before LM-head._apply.
-    head._symm_mem_device = torch.device('cuda:0')
-    head._symm_mem_dtype = head.weight.dtype
-
-    result = head._apply(lambda tensor: tensor)
-
-    assert result is head
-    gatherer.release.assert_called_once_with()
-    gatherer.prepare.assert_not_called()
-    assert head._symm_mem_device == torch.device('cpu')
diff --git a/tests/pytorch/test_symm_mem_allgather.py b/tests/pytorch/test_symm_mem_allgather.py
deleted file mode 100644
index 5933584044..0000000000
--- a/tests/pytorch/test_symm_mem_allgather.py
+++ /dev/null
@@ -1,132 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-comm = pytest.importorskip('lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather')
-
-
-def _gatherer():
-    return comm.MultimemAllGatherer(group=object(), rank=0, gathered_width=16, max_tokens=4)
-
-
-def _disable_cuda_capture_probe(monkeypatch):
-    monkeypatch.setattr(comm.torch.cuda, 'is_available', lambda: False)
-    monkeypatch.setattr(comm.dist, 'get_world_size', lambda group: 2)
-
-
-def test_prepare_rejects_before_rendezvous_when_any_rank_cannot_allocate(monkeypatch):
-    _disable_cuda_capture_probe(monkeypatch)
-    gatherer = _gatherer()
-    events = []
-
-    def _allocate(*args, **kwargs):
-        events.append('allocate')
-        raise RuntimeError('local allocation failed')
-
-    def _agree(local_ready, device):
-        events.append(('agree', local_ready, device))
-        return False
-
-    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', _allocate)
-    monkeypatch.setattr(comm, 'create_state', lambda *args, **kwargs: pytest.fail('rendezvous must not be entered'))
-    monkeypatch.setattr(gatherer, 'agree', _agree)
-
-    assert gatherer.prepare(torch.device('cuda:3')) is False
-    assert events == ['allocate', ('agree', False, torch.device('cuda:3'))]
-    assert gatherer._state is None
-
-
-def test_prepare_agrees_before_rendezvous_and_preserves_device(monkeypatch):
-    _disable_cuda_capture_probe(monkeypatch)
-    gatherer = _gatherer()
-    events = []
-    arena = torch.empty(0)
-
-    def _allocate(group, max_tokens, hidden_size, device):
-        events.append(('allocate', device))
-        return arena
-
-    def _agree(local_ready, device):
-        events.append(('agree', local_ready, device))
-        return True
-
-    def _create_state(**kwargs):
-        events.append(('rendezvous', kwargs['device'], kwargs['comm_buff']))
-        return SimpleNamespace(
-            symm_mem_hdl=SimpleNamespace(multicast_ptr=1, rank=0),
-            world_size=2,
-            max_token_num=4,
-        )
-
-    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', _allocate)
-    monkeypatch.setattr(comm, 'create_state', _create_state)
-    monkeypatch.setattr(gatherer, 'agree', _agree)
-
-    assert gatherer.prepare(torch.device('cuda:3')) is True
-    assert events == [
-        ('allocate', torch.device('cuda:3')),
-        ('agree', True, torch.device('cuda:3')),
-        ('rendezvous', torch.device('cuda:3'), arena),
-        ('agree', True, torch.device('cuda:3')),
-    ]
-
-
-def test_prepare_does_not_turn_rendezvous_failure_into_local_fallback(monkeypatch):
-    _disable_cuda_capture_probe(monkeypatch)
-    gatherer = _gatherer()
-    monkeypatch.setattr(comm, '_allocate_symmetric_buffer', lambda *args, **kwargs: torch.empty(0))
-    monkeypatch.setattr(gatherer, 'agree', lambda local_ready, device: True)
-
-    def _rendezvous_failure(**kwargs):
-        raise RuntimeError('collective rendezvous failed')
-
-    monkeypatch.setattr(comm, 'create_state', _rendezvous_failure)
-
-    with pytest.raises(RuntimeError, match='collective rendezvous failed'):
-        gatherer.prepare(torch.device('cuda:3'))
-    assert gatherer._state is gatherer._UNINIT
-
-
-def test_release_drops_arena_and_resets_graph_admission():
-    gatherer = _gatherer()
-    gatherer._state = object()
-    gatherer._graph_ready = True
-    gatherer._runtime_admitted = True
-
-    gatherer.release()
-
-    assert gatherer._state is gatherer._UNINIT
-    assert gatherer._graph_ready is False
-    assert gatherer._runtime_admitted is False
-
-
-def test_first_call_collectively_admits_static_input_contract(monkeypatch):
-    gatherer = _gatherer()
-    gatherer._state = SimpleNamespace(
-        max_token_num=4,
-        world_size=2,
-        hidden_dim=16,
-        device=torch.device('cpu'),
-    )
-    admission = Mock(return_value=True)
-    monkeypatch.setattr(gatherer, 'agree', admission)
-    monkeypatch.setattr(comm.torch.cuda, 'is_current_stream_capturing', lambda: False)
-    monkeypatch.setattr(comm, 'all_gather_inner', lambda state, x, **kwargs: x)
-    value = torch.empty((1, 8), dtype=torch.bfloat16)
-
-    assert gatherer(value) is value
-    assert gatherer(value) is value
-    admission.assert_called_once_with(True, torch.device('cpu'))
-    assert gatherer._runtime_admitted is True
-
-
-def test_direct_path_rejects_unaligned_local_width_before_kernel():
-    gatherer = _gatherer()
-    gatherer._state = SimpleNamespace(max_token_num=4, world_size=2, hidden_dim=8, device=torch.device('cpu'))
-    gatherer.agree = lambda local_ready, device: local_ready
-
-    assert gatherer(torch.empty((1, 4), dtype=torch.bfloat16)) is None
-    assert gatherer._state is None

From 1a4e06e6815e7f6fa52ca13394f14c704a94aae9 Mon Sep 17 00:00:00 2001
From: qescccczmr <105876751+qescccczmr@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:57:49 +0800
Subject: [PATCH 5/6] Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
index fa4660de1c..85c3839b8d 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -332,7 +332,8 @@ def create_state(
     Collective: call
     once outside CUDA-graph capture with identical args on every rank.
     """
-    assert type(group) is dist.ProcessGroup, f"Expected ProcessGroup, got {type(group)}"
+    if not isinstance(group, dist.ProcessGroup):
+        raise TypeError(f"Expected dist.ProcessGroup, got {type(group)}")
     assert hidden_size % _NUMEL_PER_THREAD == 0, (
         f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
         f"bf16 for 16-byte multimem.st row alignment"

From 669b1e0074e0ef56bc92937761400f6a1eb5a450 Mon Sep 17 00:00:00 2001
From: Codex <codex@example.invalid>
Date: Fri, 4 Sep 2026 03:18:56 +0000
Subject: [PATCH 6/6] perf: tune opt-in symmetric-memory LM-head gather

---
 .../backends/cuda/comm/symm_mem_allgather.py  | 913 +++++++++++++++---
 lmdeploy/pytorch/envs.py                      |  23 +
 lmdeploy/pytorch/nn/embedding.py              |  50 +-
 3 files changed, 832 insertions(+), 154 deletions(-)

diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
index 85c3839b8d..851527b966 100644
--- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
+++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py
@@ -9,6 +9,7 @@
 
 import logging
 from dataclasses import dataclass
+from functools import lru_cache
 from typing import Any
 
 import torch
@@ -17,17 +18,43 @@
 import triton
 import triton.language as tl
 
+from lmdeploy.pytorch import envs as _envs
+
 logger = logging.getLogger(__name__)
 
 # Each thread moves _NUMEL_PER_THREAD bf16 via one 128-bit multimem op; the
 # grid-strided block count is tunable in [_MIN_BLOCKS, _MAX_BLOCKS].
 _BLOCK_THREADS = 1024
+_BLOCK_THREAD_CANDIDATES = (256, 512, 1024)
 _NUMEL_PER_THREAD = 8
 _MIN_BLOCKS = 4
 _MAX_BLOCKS = 32
+_TARGET_GRID_STRIDE_ITERS = 4
+# A separate one-CTA barrier is cheaper once payload parallelism would make
+# the per-CTA signal protocol issue many duplicate cross-rank CAS operations.
+# The threshold is deliberately conservative; ``B<16`` retains the original
+# single-kernel path and avoids paying two extra launches for tiny decode M.
+_SINGLE_BARRIER_MIN_BLOCKS = 16
 _SUPPORTED_WORLD_SIZES = {2, 4, 8}
 
 
+def _is_cuda_graph_capturing() -> bool:
+    """Probe capture state without touching the CUDA driver on CPU paths."""
+    try:
+        available = torch.cuda.is_available()
+    except RuntimeError:
+        available = False
+    if not available:
+        return False
+    try:
+        return bool(torch.cuda.is_current_stream_capturing())
+    except RuntimeError:
+        # Some CPU-only test runners expose the CUDA module but no active
+        # driver. Treat those calls as eager; normal device checks disable the
+        # provider when no CUDA driver is available.
+        return False
+
+
 # ------------------------------------------------------------------------------
 # Low-level PTX helpers
 # ------------------------------------------------------------------------------
@@ -74,44 +101,18 @@ def _local_ld_128(in_ptr, mask):
 
 
 @triton.jit
-def _get_tid():
+def _get_tid_x():
+    """Return the lane's linear thread id for the x-only launch contract."""
     return tl.inline_asm_elementwise(
-        """
-        mov.u32 $0, %tid.x;
-        mov.u32 $1, %tid.y;
-        mov.u32 $2, %tid.z;
-        """,
-        '=r,=r,=r',
-        [],
-        dtype=(tl.uint32, tl.uint32, tl.uint32),
-        is_pure=True,
-        pack=1,
-    )
-
-
-@triton.jit
-def _get_ntid():
-    return tl.inline_asm_elementwise(
-        """
-        mov.u32 $0, %ntid.x;
-        mov.u32 $1, %ntid.y;
-        mov.u32 $2, %ntid.z;
-        """,
-        '=r,=r,=r',
+        'mov.u32 $0, %tid.x;',
+        '=r',
         [],
-        dtype=(tl.uint32, tl.uint32, tl.uint32),
+        dtype=tl.uint32,
         is_pure=True,
         pack=1,
     )
 
 
-@triton.jit
-def _get_flat_tid():
-    tid_x, tid_y, tid_z = _get_tid()
-    ntid_x, ntid_y, _ = _get_ntid()
-    return tid_z * ntid_y * ntid_x + tid_y * ntid_x + tid_x
-
-
 @triton.jit
 def _sync_threads():
     tl.inline_asm_elementwise(
@@ -221,33 +222,39 @@ def _blockwise_barrier(
     rank: tl.constexpr,
     world_size: tl.constexpr,
     sem: tl.constexpr,
+    slot_offset: tl.constexpr = 0,
 ):
-    block_id = (
-        tl.program_id(2) * tl.num_programs(1) * tl.num_programs(0)
-        + tl.program_id(1) * tl.num_programs(0)
-        + tl.program_id(0)
-    )
-    flat_tid = _get_flat_tid()
+    # Every caller launches an x-only grid and Triton maps ``num_warps`` to an
+    # x-only CUDA thread block.  Specializing that invariant removes the
+    # generic y/z CTA and thread-index arithmetic from both barriers.
+    block_id = tl.program_id(0) + slot_offset
+    flat_tid = _get_tid_x()
 
-    remote_ranks = tl.arange(0, world_size)
-    signal_pad_ptrs = signal_pad_ptrs.to(tl.pointer_type(tl.uint64))
-    remote_signal_pad_addrs = tl.load(signal_pad_ptrs + remote_ranks).to(
-        tl.pointer_type(tl.uint32)
-    )
-    send_addrs = remote_signal_pad_addrs + block_id * world_size + rank
-
-    local_signal_pad_addr = tl.load(signal_pad_ptrs + rank).to(
-        tl.pointer_type(tl.uint32)
-    )
-    wait_addrs = local_signal_pad_addr + block_id * world_size + remote_ranks
+    # Keep the cast in a distinct SSA value. Triton cannot merge the
+    # pointer-typed branch value with the original int64 tensor.
+    signal_pad_ptrs_u64 = signal_pad_ptrs.to(tl.pointer_type(tl.uint64))
 
     if flat_tid < world_size:
+        # One lane is assigned to one peer.  Keeping the peer index scalar
+        # avoids materializing a rank-wide pointer vector in the Triton IR and
+        # shortens the address live range (the NVIDIA backend may scalarize
+        # either form, but this spelling also keeps the protocol explicit).
+        # Self-send/self-wait remains intentional: it keeps the epoch complete
+        # even for a single local rank in test/fake providers.
+        peer = flat_tid
+        remote_signal_pad_addr = tl.load(signal_pad_ptrs_u64 + peer).to(
+            tl.pointer_type(tl.uint32))
+        local_signal_pad_addr = tl.load(signal_pad_ptrs_u64 + rank).to(
+            tl.pointer_type(tl.uint32))
+        send_addr = (remote_signal_pad_addr + block_id * world_size + rank)
+        wait_addr = (local_signal_pad_addr + block_id * world_size + peer)
+
         if sem == 'relaxed':
-            _send_signal(send_addrs)
-            _wait_signal(wait_addrs)
+            _send_signal(send_addr)
+            _wait_signal(wait_addr)
         else:
-            _send_signal_release(send_addrs)
-            _wait_signal_acquire(wait_addrs)
+            _send_signal_release(send_addr)
+            _wait_signal_acquire(wait_addr)
 
 
 @triton.jit
@@ -264,6 +271,7 @@ def _all_gather_kernel_inner(
     RANK: tl.constexpr,
     WORLD_SIZE: tl.constexpr,
     SKIP_ENTRY_SYNC: tl.constexpr,
+    SKIP_EXIT_SYNC: tl.constexpr,
 ) -> None:
     if SKIP_ENTRY_SYNC == 0:
         _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem='relaxed')
@@ -275,7 +283,7 @@ def _all_gather_kernel_inner(
     total_chunks = total_tokens * chunks_per_row
 
     pid = tl.program_id(axis=0)
-    tid = _get_flat_tid()
+    tid = _get_tid_x()
     block_start = pid * BLOCK_SIZE
 
     while block_start < total_chunks:
@@ -293,12 +301,39 @@ def _all_gather_kernel_inner(
         _multimem_st_128(out_ptr, x, y, z, w, mask)
         block_start += tl.num_programs(axis=0) * BLOCK_SIZE
 
-    # The producer writes through the multicast VA and callers consume through
+    # The payload writes through the multicast VA and callers consume through
     # the ordinary symmetric-buffer VA.  Hopper requires an alias-proxy fence
-    # before the release/acquire completion handshake.
+    # before the release/acquire completion handshake.  In split-barrier mode
+    # the payload kernel is followed (on the same stream) by a one-CTA
+    # completion barrier, so a grid-wide CTA barrier and per-CTA remote CAS are
+    # unnecessary here; kernel completion is the grid synchronization point.
     _fence_proxy_alias()
-    _sync_threads()
-    _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem='acq_rel')
+    if SKIP_EXIT_SYNC == 0:
+        _sync_threads()
+        _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem='acq_rel')
+
+
+@triton.jit
+def _one_block_barrier_kernel(
+    signal_pad_ptr,
+    RANK: tl.constexpr,
+    WORLD_SIZE: tl.constexpr,
+    SLOT: tl.constexpr,
+    RELEASE: tl.constexpr,
+):
+    """One-CTA cross-rank rendezvous for payload launches."""
+    if RELEASE:
+        _blockwise_barrier(signal_pad_ptr,
+                           RANK,
+                           WORLD_SIZE,
+                           sem='acq_rel',
+                           slot_offset=SLOT)
+    else:
+        _blockwise_barrier(signal_pad_ptr,
+                           RANK,
+                           WORLD_SIZE,
+                           sem='relaxed',
+                           slot_offset=SLOT)
 
 
 # ------------------------------------------------------------------------------
@@ -333,21 +368,36 @@ def create_state(
     once outside CUDA-graph capture with identical args on every rank.
     """
     if not isinstance(group, dist.ProcessGroup):
-        raise TypeError(f"Expected dist.ProcessGroup, got {type(group)}")
+        raise TypeError(f'Expected ProcessGroup, got {type(group)}')
+    if max_tokens <= 0:
+        raise ValueError(f'max_tokens must be positive, got {max_tokens}')
+    if rank_in_group < 0 or rank_in_group >= group.size():
+        raise ValueError(
+            f'rank_in_group={rank_in_group} is outside group size={group.size()}')
     assert hidden_size % _NUMEL_PER_THREAD == 0, (
         f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
         f"bf16 for 16-byte multimem.st row alignment"
     )
-    device = torch.device(device or torch.device(f"cuda:{torch.cuda.current_device()}"))
+    device = torch.device(
+        device or torch.device(f"cuda:{torch.cuda.current_device()}"))
     if device.type == 'cuda' and device.index is None:
         device = torch.device('cuda', torch.cuda.current_device())
 
     if comm_buff is None:
-        comm_buff = _allocate_symmetric_buffer(group, max_tokens, hidden_size, device)
-    elif (comm_buff.shape != (max_tokens, hidden_size) or comm_buff.dtype != torch.bfloat16
-          or comm_buff.device != device or not comm_buff.is_contiguous() or comm_buff.storage_offset() != 0):
+        comm_buff = _allocate_symmetric_buffer(group, max_tokens, hidden_size,
+                                                device)
+    elif (comm_buff.shape != (max_tokens, hidden_size)
+          or comm_buff.dtype != torch.bfloat16
+          or comm_buff.device != device or not comm_buff.is_contiguous()
+          or comm_buff.storage_offset() != 0):
         raise ValueError('preallocated symmetric buffer does not match state')
     hdl = symm_mem.rendezvous(comm_buff, group=group)
+    # Do not raise on a rank-local handle mismatch here.  ``rendezvous`` is a
+    # TP collective and the caller must be able to run one more TP-wide
+    # validity reduction before deciding whether to disable the provider;
+    # raising on only the mismatching rank would leave its peers in that
+    # reduction (or in the next launch) indefinitely.  ``_build`` validates
+    # rank/world-size/pointers collectively after this function returns.
     return MultimemAllGatherState(
         group=group,
         rank_in_group=rank_in_group,
@@ -367,16 +417,129 @@ def _allocate_symmetric_buffer(
     device: torch.device,
 ) -> torch.Tensor:
     """Perform only the local allocation half of state construction."""
-    # Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
+    # Pad holds the per-CTA slots for the payload protocol.  Split barriers
+    # use slot 0 for entry and slot 1 for a one-CTA completion rendezvous;
+    # consequently the latter consumes only the next single block slot (not
+    # another ``_MAX_BLOCKS`` range).  The max() never shrinks the pad.
     pad_bytes = _MAX_BLOCKS * group.size() * 4
-    symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
+    current_pad = symm_mem.get_signal_pad_size()
+    if current_pad < pad_bytes:
+        try:
+            # PyTorch requires this setting before the first symmetric
+            # allocation in a process.  If another symmetric-memory user has
+            # already allocated with a smaller pad, do not turn the situation
+            # into a rank-local exception: the caller catches this error and
+            # TP-wide disables the optional provider.
+            symm_mem.set_signal_pad_size(pad_bytes)
+        except RuntimeError:
+            if symm_mem.get_signal_pad_size() < pad_bytes:
+                raise
     with torch.inference_mode(False), torch.no_grad():
-        return symm_mem.empty((max_tokens, hidden_size), dtype=torch.bfloat16, device=device)
+        return symm_mem.empty((max_tokens, hidden_size),
+                              dtype=torch.bfloat16,
+                              device=device)
+
+
+@lru_cache(maxsize=256)
+def _launch_config_cached(local_numel: int, configured_threads: int,
+                          configured_blocks: int, autotune: bool,
+                          total_tokens: int, world_size: int):
+    """Resolve a launch shape once for each static input/configuration pair.
+
+    Decode reuses a small set of flattened token counts.  Keeping the policy
+    resolution out of the hot path is useful in its own right, and including
+    the environment values in the cache key keeps unit tests and long-lived
+    processes that construct more than one tuning profile deterministic.
+    """
+    if local_numel <= 0 or local_numel % _NUMEL_PER_THREAD != 0:
+        raise ValueError(
+            f'local_numel must be a positive multiple of '
+            f'{_NUMEL_PER_THREAD}, got {local_numel}')
+    chunks = local_numel // _NUMEL_PER_THREAD
+
+    if configured_threads and configured_threads not in _BLOCK_THREAD_CANDIDATES:
+        raise ValueError(
+            'LMDEPLOY_SYMM_MEM_LMHEAD_BLOCK_THREADS must be one of '
+            f'{_BLOCK_THREAD_CANDIDATES}, got {configured_threads}')
+    if configured_blocks and not _MIN_BLOCKS <= configured_blocks <= _MAX_BLOCKS:
+        raise ValueError(
+            'LMDEPLOY_SYMM_MEM_LMHEAD_BLOCKS must be in '
+            f'[{_MIN_BLOCKS}, {_MAX_BLOCKS}], got {configured_blocks}')
+
+    if configured_threads:
+        block_threads = configured_threads
+    elif not autotune:
+        block_threads = _BLOCK_THREADS
+    elif chunks <= 2048:
+        block_threads = 256
+    elif chunks <= 8192:
+        block_threads = 512
+    else:
+        block_threads = 1024
+
+    if configured_blocks:
+        num_blocks = configured_blocks
+    elif not autotune:
+        num_blocks = _MIN_BLOCKS
+    else:
+        chunks_per_block = block_threads * _TARGET_GRID_STRIDE_ITERS
+        num_blocks = (chunks + chunks_per_block - 1) // chunks_per_block
+        max_blocks = _MAX_BLOCKS
+        if world_size == 2 and 0 < total_tokens <= 32:
+            # On TP2 the signal traffic is already small, so M=6..32 is
+            # dominated by the two per-CTA barriers before enough payload is
+            # available to profit from a wide grid.  H200 paired sweeps pick
+            # four CTAs here; cap at four while retaining the full
+            # 32-CTA range for larger payloads and TP4/TP8.
+            max_blocks = _MIN_BLOCKS
+        elif world_size == 2 and 32 < total_tokens <= 64:
+            # A small middle band benefits from a little more payload
+            # parallelism, but a full 32-CTA grid would still multiply the
+            # entry/exit signal traffic before the copy is saturated.
+            max_blocks = 8
+        elif world_size >= 4 and 0 < total_tokens <= 8:
+            max_blocks = 8
+        elif world_size >= 4 and 8 < total_tokens <= 32:
+            max_blocks = 16
+        num_blocks = min(max_blocks, max(_MIN_BLOCKS, num_blocks))
+
+    return (num_blocks, block_threads, block_threads // 32,
+            _NUMEL_PER_THREAD)
+
+
+def _launch_config(local_numel: int, *, total_tokens: int = 0,
+                   world_size: int = 0):
+    """Return the cached launch shape for the current process policy."""
+    return _launch_config_cached(
+        int(local_numel),
+        int(_envs.symm_mem_lmhead_block_threads),
+        int(_envs.symm_mem_lmhead_blocks),
+        bool(_envs.symm_mem_lmhead_autotune),
+        int(total_tokens),
+        int(world_size),
+    )
+
 
+def _use_single_barrier(num_blocks: int, world_size: int) -> bool:
+    """Select the barrier protocol for one payload launch.
 
-def _launch_config(local_numel: int):
-    assert local_numel % _NUMEL_PER_THREAD == 0
-    return _MIN_BLOCKS, _BLOCK_THREADS, _BLOCK_THREADS // 32, _NUMEL_PER_THREAD
+    In ``single`` mode the entry and completion rendezvous are separate
+    one-CTA kernels around a payload kernel.  Stream ordering then provides a
+    grid-wide completion point, while only one CTA performs cross-rank CAS.
+    This is safe for arbitrary grid sizes (unlike a per-CTA remote barrier,
+    which scales its signal traffic with the number of resident CTAs).
+    """
+    mode = getattr(_envs, 'symm_mem_lmhead_barrier_mode', 'auto')
+    if mode == 'single':
+        return True
+    if mode == 'per_block':
+        return False
+    # For TP2/TP4 the two extra launches usually cost more than the small
+    # signal reduction at decode sizes.  TP8 issues 4*W signal atomics per CTA
+    # across the entry/exit barriers, where the split protocol starts paying
+    # back for genuinely large payload grids.
+    return (world_size >= 8
+            and num_blocks >= _SINGLE_BARRIER_MIN_BLOCKS)
 
 
 def all_gather_inner(
@@ -385,44 +548,59 @@ def all_gather_inner(
     tp_hidden_dim: int,
     skip_entry_sync: bool = False,
     safe: bool = True,
+    *,
+    _validated: bool = False,
 ) -> torch.Tensor:
     """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim.
 
     ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``,
     else a view into the symmetric buffer (valid until the next collective).
+    ``_validated`` is reserved for :meth:`MultimemAllGatherer.fast_call`,
+    whose admission check already enforces the immutable dtype/layout/width
+    contract.  Public/direct callers retain all defensive assertions.
     """
     world_size = state.world_size
-    assert hidden_states.dtype == torch.bfloat16, 'Only bfloat16 is supported'
-    assert hidden_states.is_contiguous(), 'hidden_states must be contiguous'
-    assert hidden_states.data_ptr() % 16 == 0, (
-        f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
-        f"16-byte aligned for 128-bit multimem.st"
-    )
-    assert (
-        tp_hidden_dim % world_size == 0
-    ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
+    if not _validated:
+        assert hidden_states.dtype == torch.bfloat16, 'Only bfloat16 is supported'
+        assert hidden_states.is_contiguous(), 'hidden_states must be contiguous'
+        assert hidden_states.data_ptr() % 16 == 0, (
+            f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
+            f"16-byte aligned for 128-bit multimem.st"
+        )
+        assert (
+            tp_hidden_dim % world_size == 0
+        ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
     local_hidden = tp_hidden_dim // world_size
-    assert local_hidden % _NUMEL_PER_THREAD == 0, (
-        f"per-rank hidden shard ({local_hidden}) must be a multiple of "
-        f"{_NUMEL_PER_THREAD} bf16"
-    )
-    assert tp_hidden_dim <= state.hidden_dim, (
-        f"comm buffer too narrow: tp_hidden_dim={tp_hidden_dim} > "
-        f"state.hidden_dim={state.hidden_dim}"
-    )
     total_tokens, in_hidden = hidden_states.shape
-    assert (
-        in_hidden == local_hidden
-    ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
-    assert (
-        total_tokens <= state.max_token_num
-    ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
+    if not _validated:
+        assert local_hidden % _NUMEL_PER_THREAD == 0, (
+            f"per-rank hidden shard ({local_hidden}) must be a multiple of "
+            f"{_NUMEL_PER_THREAD} bf16"
+        )
+        assert tp_hidden_dim <= state.hidden_dim, (
+            f"comm buffer too narrow: tp_hidden_dim={tp_hidden_dim} > "
+            f"state.hidden_dim={state.hidden_dim}"
+        )
+        assert (
+            in_hidden == local_hidden
+        ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
+        assert (
+            total_tokens <= state.max_token_num
+        ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
 
     hidden_offset = local_hidden * state.rank_in_group
     symm_mem_hdl = state.symm_mem_hdl
     num_blocks, block_size, num_warps, numel_per_thread = _launch_config(
-        total_tokens * local_hidden
+        total_tokens * local_hidden,
+        total_tokens=total_tokens,
+        world_size=world_size,
     )
+    split_barrier = _use_single_barrier(num_blocks, world_size)
+    if split_barrier and not skip_entry_sync:
+        # The one-CTA entry barrier is launched before the payload kernel on
+        # the same stream.  It retains the reuse protection of the original
+        # protocol while avoiding one barrier per payload CTA.
+        barrier_inner(state, slot=0, release=False)
     grid = (num_blocks, 1, 1)
     _all_gather_kernel_inner[grid](
         input_ptr=hidden_states,
@@ -436,13 +614,39 @@ def all_gather_inner(
         NUMEL_PER_THREAD=numel_per_thread,
         RANK=symm_mem_hdl.rank,
         WORLD_SIZE=symm_mem_hdl.world_size,
-        SKIP_ENTRY_SYNC=1 if skip_entry_sync else 0,
+        SKIP_ENTRY_SYNC=1 if (skip_entry_sync or split_barrier) else 0,
+        SKIP_EXIT_SYNC=1 if split_barrier else 0,
         num_warps=num_warps,
     )
+    if split_barrier:
+        # A kernel boundary on one CUDA stream waits for every payload CTA.
+        # The completion barrier can therefore be a single CTA and still
+        # publish the whole grid before the output view is returned.
+        barrier_inner(state, slot=1, release=True)
     output = state.comm_buff[:total_tokens, :tp_hidden_dim]
     return output.clone() if safe else output
 
 
+def barrier_inner(state: MultimemAllGatherState, *, slot: int,
+                  release: bool) -> None:
+    """Synchronize ranks without moving payload data.
+
+    ``slot`` selects independent signal-pad storage. Slot zero is the entry
+    epoch and slot one is the split-payload completion epoch.
+    """
+    if slot not in (0, 1):
+        raise ValueError(f'barrier slot must be 0 or 1, got {slot}')
+    handle = state.symm_mem_hdl
+    _one_block_barrier_kernel[(1, )](
+        signal_pad_ptr=handle.signal_pad_ptrs_dev,
+        RANK=state.rank_in_group,
+        WORLD_SIZE=state.world_size,
+        SLOT=slot,
+        RELEASE=1 if release else 0,
+        num_warps=1,
+    )
+
+
 # ------------------------------------------------------------------------------
 # Guarded wrapper
 # ------------------------------------------------------------------------------
@@ -451,10 +655,10 @@ def all_gather_inner(
 class MultimemAllGatherer:
     """Guarded last-dim multimem all-gather with NCCL fallback.
 
-    Owns one symmetric buffer and admits the input contract collectively on the first eager call. Subsequent calls
-    require the same dtype, layout and device on every rank; token count may vary but must be TP-invariant. The returned
-    tensor owns its storage so a later collective cannot overwrite logits that are still being consumed on another
-    stream.
+    Owns one symmetric buffer built lazily on the first eager call, and uses
+    the kernel only when the input fits its dtype/shape/alignment contract.
+    The returned tensor owns its storage so a later collective cannot overwrite
+    logits that are still being consumed on another stream.
     """
 
     _UNINIT = object()
@@ -473,119 +677,466 @@ def __init__(
         self._gathered_width = gathered_width
         self._max_tokens = int(max_tokens)
         self._enabled = enabled
+        self._logged_dispatch = False
         self._graph_ready = False
+        # CUDA Graph capture is shape-specialized. A single bool is not
+        # sufficient when warmup first sees M=1 and capture later replays M=8;
+        # an unseen shape could trigger Triton compilation inside capture.
+        self._graph_ready_shapes = set()
+        # Static input admission is a setup-time collective.  It prevents a
+        # rank-local alignment/device check from sending one TP rank into NCCL
+        # while its peers enter the multimem signal protocol.
         self._runtime_admitted = False
+
         # None => always NCCL; _UNINIT => build on first eager call.
         self._state = self._UNINIT if enabled else None
 
-    def __call__(self, x: torch.Tensor) -> torch.Tensor | None:
-        state = self._state
-        if state is self._UNINIT:
-            state = self._build(x.device)
-            if state is not self._UNINIT:
-                self._state = state
+    def __call__(self, x: torch.Tensor, *,
+                 safe: bool = True) -> torch.Tensor | None:
+        """Gather logits, optionally returning the reusable arena view.
+
+        ``safe=True`` keeps the public owning-output contract. ``safe=False``
+        is reserved for a serialized consumer that finishes using the result
+        before the next collective on this instance.
+        """
+        state = self.get_state(x)
         if state is None or state is self._UNINIT:
             return None
 
         eligible = self._is_static_input_eligible(state, x)
-        capturing = torch.cuda.is_current_stream_capturing()
+        state_device = getattr(state, 'device', x.device)
+        # Admission and CUDA-graph capture are mutually exclusive.  The
+        # engine warms each decode shape eagerly before capture; an accidental
+        # first call inside capture therefore stays on the portable path.
+        if x.dim() >= 1:
+            shape_key = (int(x.shape[0]), int(x.shape[-1]))
+        else:
+            shape_key = (0, 0)
+        # Once a shape has been launched eagerly, its Triton specialization is
+        # ready and the same shape is safe to replay under a CUDA Graph.  Do
+        # not query the CUDA driver on every steady-state decode call; on
+        # Hopper that host query is measurable at M=1.
+        capture_probe_needed = (not self._runtime_admitted
+                                or shape_key not in self._graph_ready_shapes)
+        capturing = (capture_probe_needed and _is_cuda_graph_capturing())
         if not self._runtime_admitted:
-            # Admission uses host-visible consensus and cannot run under graph
-            # capture. In that uncommon first-call case every rank retains the
-            # existing NCCL path captured by ParallelLMHead.
             if capturing:
                 return None
-            if not self.agree(eligible, state.device):
-                self._state = None
+            if not self.agree(eligible, state_device):
+                # Every rank receives the same result from ``agree``.  Drop
+                # this attempted arena and leave lazy admission open: a
+                # transient prefill/layout probe must not permanently disable
+                # the provider needed by a later valid decode call.
+                self.release()
+                return None
+            if not self._agree_first_shape(x, state_device):
+                # The portable all-gather API also requires equal tensor
+                # shapes. Converge before either rank enters the multimem
+                # protocol when the first TP call violates that invariant;
+                # keep the provider retryable after this rejected probe.
+                self.release()
                 return None
             self._runtime_admitted = True
         elif not eligible:
-            raise RuntimeError('multimem all-gather input contract changed after TP-wide admission')
+            raise RuntimeError(
+                'multimem all-gather input contract changed after TP-wide '
+                'admission')
+
+        # Token count is intentionally the one dynamic dimension.  The model
+        # scheduler presents the same flattened row count to every TP rank;
+        # a capacity miss consequently follows the same NCCL fallback on all
+        # ranks.  Known decode/MTP shapes are pre-reserved at construction.
+        if not 0 < x.shape[0] <= state.max_token_num:
+            return None
+        if capturing:
+            return None
+        if not self._logged_dispatch:
+            self._logged_dispatch = True
+            if self._rank == 0:
+                logger.warning(
+                    'multimem all-gather direct dispatch active '
+                    '(tokens=%d, local_width=%d)', x.shape[0], x.shape[-1])
+        output = all_gather_inner(
+            state,
+            x,
+            tp_hidden_dim=self._gathered_width,
+            skip_entry_sync=False,
+            safe=safe,
+            _validated=True,
+        )
+        self._graph_ready = True
+        self._graph_ready_shapes.add(shape_key)
+        return output
+
+    def fast_call(self, x: torch.Tensor, *, safe: bool = True
+                  ) -> torch.Tensor | None:
+        """Run an admitted gather with a reduced dispatch check.
 
+        ``ParallelLMHead`` owns a stable BF16 contiguous logits tensor after
+        the first eager admission.  The fast path keeps the inexpensive static
+        contract guard for fail-closed behavior, while avoiding setup
+        collectives and repeated state construction; arbitrary users retain
+        the fully validating ``__call__`` API.
+        """
+        if not self._runtime_admitted:
+            return self(x, safe=safe)
+        state = self._state
+        if state is None or state is self._UNINIT:
+            return None
+        if x.dim() != 2 or x.device != state.device:
+            return self(x, safe=safe)
+        if not self._is_static_input_eligible(state, x):
+            # After TP-wide admission this is a programming/weight-contract
+            # violation, not a rank-local reason to switch to NCCL. Raising
+            # keeps all ranks from silently taking different collective paths.
+            raise RuntimeError(
+                'multimem all-gather input contract changed after TP-wide '
+                'admission')
         if not 0 < x.shape[0] <= state.max_token_num:
-            # ParallelLMHead presents identical token counts to all TP ranks,
-            # so this dynamic capacity fallback cannot split the collective.
             return None
-        # State allocation may be prepared before capture while Triton has not
-        # compiled this shape yet. Keep that first call on NCCL rather than
-        # trying to JIT-compile inside a CUDA Graph.
-        if capturing and not self._graph_ready:
+        shape_key = (int(x.shape[0]), int(x.shape[-1]))
+        if ((not self._runtime_admitted
+             or shape_key not in self._graph_ready_shapes)
+                and _is_cuda_graph_capturing()):
             return None
         output = all_gather_inner(
             state,
             x,
             tp_hidden_dim=self._gathered_width,
             skip_entry_sync=False,
-            safe=True,
+            safe=safe,
+            _validated=True,
         )
         self._graph_ready = True
+        self._graph_ready_shapes.add(shape_key)
         return output
 
     @staticmethod
-    def _is_static_input_eligible(state: MultimemAllGatherState, x: torch.Tensor) -> bool:
-        """Check properties that remain stable for one LM-head instance."""
-        return (x.dtype == torch.bfloat16 and x.dim() == 2 and x.device == state.device and x.is_contiguous()
-                and x.data_ptr() % 16 == 0 and x.shape[-1] % _NUMEL_PER_THREAD == 0
+    def _is_static_input_eligible(state: MultimemAllGatherState,
+                                  x: torch.Tensor) -> bool:
+        """Check the properties shared by all calls of one LM-head."""
+        state_device = getattr(state, 'device', x.device)
+        return (x.dtype == torch.bfloat16 and x.dim() == 2
+                and x.device == state_device and x.is_contiguous()
+                and x.data_ptr() % 16 == 0
+                and x.shape[-1] % _NUMEL_PER_THREAD == 0
                 and x.shape[-1] * state.world_size == state.hidden_dim)
 
+    def _agree_first_shape(self, x: torch.Tensor,
+                           device: torch.device | str) -> bool:
+        """Check the first dynamic row/width shape across the TP group."""
+        if not self._uses_real_process_group():
+            return True
+        rows = int(x.shape[0]) if x.dim() >= 1 else -1
+        width = int(x.shape[-1]) if x.dim() >= 1 else -1
+        shape = torch.tensor((rows, width), dtype=torch.int64, device=device)
+        lower = shape.clone()
+        upper = shape.clone()
+        dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=self._group)
+        dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=self._group)
+        return bool(torch.equal(lower, upper))
+
+    def get_state(self, x: torch.Tensor):
+        """Return the lazily rendezvoused symmetric-memory state."""
+        state = self._state
+        if state is self._UNINIT:
+            # A CUDA tensor of any shape still participates in lazy build.  The
+            # subsequent TP-wide static admission in ``__call__`` decides
+            # whether this particular input is eligible.  This matters when a
+            # malformed tensor reaches only one rank: all ranks must enter the
+            # same build/admission sequence instead of one rank returning
+            # before its peers reach ``rendezvous``. CPU callers cannot
+            # rendezvous a CUDA arena and keep the state open for a later CUDA
+            # call.
+            if x.device.type != 'cuda':
+                # Keep the old lazy-test/portable behavior for an explicitly
+                # valid CPU-shaped probe.  Real CUDA callers take the branch
+                # below; a CPU input is never sent to a symmetric-memory
+                # kernel by the LM-head backend.
+                if (x.dim() != 2 or x.dtype != torch.bfloat16
+                        or x.shape[-1] % _NUMEL_PER_THREAD != 0):
+                    return self._UNINIT
+                state = self._build(x.device)
+                if state is self._UNINIT or state is None:
+                    # A CPU probe must not permanently disable a provider that
+                    # may be moved to CUDA later.
+                    return self._UNINIT
+                self._state = state
+                return state
+            if not torch.cuda.is_available():
+                return self._UNINIT
+            if _is_cuda_graph_capturing():
+                return self._UNINIT
+            state = self._build(x.device)
+            if state is not self._UNINIT:
+                self._state = state
+        return state
+
     def prepare(self, device: torch.device | str) -> bool:
-        """Collectively allocate and rendezvous the arena before graph
-        capture."""
+        """Collectively materialize the arena before graph capture.
+
+        Call this on every rank after the TP process group is ready.  Forward
+        retains the lazy path for compatibility, but production warmup uses
+        this method so allocation/rendezvous never occurs in a captured region.
+        """
+        device = torch.device(device)
+        if device.type == 'cuda' and device.index is None:
+            device = torch.device('cuda', torch.cuda.current_device())
+        if _is_cuda_graph_capturing():
+            raise RuntimeError(
+                'symmetric-memory prepare must run before CUDA Graph capture')
+        if self._uses_real_process_group():
+            # ``prepare`` is normally called with the same phase on every TP
+            # rank.  Still resolve a mixed READY/UNINIT (or DISABLED) phase
+            # before entering ``_build``: otherwise one rank could skip the
+            # symmetric rendezvous while a peer enters it after a device
+            # transition.  The two reductions are setup-only and never occur
+            # on the decode hot path.
+            phase = (2 if self._state is None else
+                     1 if self._state is not self._UNINIT else 0)
+            phase_min = torch.tensor(phase, dtype=torch.int32, device=device)
+            phase_max = phase_min.clone()
+            dist.all_reduce(phase_min, op=dist.ReduceOp.MIN, group=self._group)
+            dist.all_reduce(phase_max, op=dist.ReduceOp.MAX, group=self._group)
+            if int(phase_min.item()) != int(phase_max.item()):
+                if int(phase_max.item()) == 2:
+                    # A disabled rank cannot safely be rebuilt by its peers;
+                    # make the provider uniformly disabled instead.
+                    if self._state is not None and self._state is not self._UNINIT:
+                        self.release()
+                    self._state = None
+                    return False
+                # READY/UNINIT mismatch: all enabled ranks rebuild the same
+                # arena and rendezvous in the same order.
+                if self._state is not self._UNINIT:
+                    self.release()
+                self._state = self._UNINIT if self._enabled else None
+            elif int(phase_min.item()) == 1:
+                state = self._state
+                local_state_ok = (
+                    getattr(state, 'device', None) == device
+                    and getattr(state, 'rank_in_group', -1) == self._rank
+                    and getattr(state, 'group', None) is self._group)
+                if not self.agree(local_state_ok, device):
+                    # A stale arena (for example after an offload cycle) must
+                    # not remain ready on only a subset of TP ranks.
+                    self.release()
+                    self._state = None
+                    return False
+                config = torch.tensor(
+                    (state.world_size, state.max_token_num,
+                     state.hidden_dim, int(self._runtime_admitted)),
+                    dtype=torch.int64,
+                    device=device,
+                )
+                config_min = config.clone()
+                config_max = config.clone()
+                dist.all_reduce(config_min,
+                                op=dist.ReduceOp.MIN,
+                                group=self._group)
+                dist.all_reduce(config_max,
+                                op=dist.ReduceOp.MAX,
+                                group=self._group)
+                if not torch.equal(config_min, config_max):
+                    # A stale/mismatched ready arena cannot safely be reused;
+                    # converge to the same disabled state rather than letting
+                    # one rank launch with a different layout.
+                    self.release()
+                    self._state = None
+                    return False
+                return True
+            elif int(phase_min.item()) == 2:
+                return False
+
         if self._state is self._UNINIT:
-            state = self._build(torch.device(device))
+            state = self._build(device)
             if state is not self._UNINIT:
                 self._state = state
         return self._state is not None and self._state is not self._UNINIT
 
-    def agree(self, local_ready: bool, device: torch.device | str) -> bool:
-        """Return a TP-wide setup decision so ranks never split paths."""
-        if torch.cuda.is_current_stream_capturing():
+    def admit_static(self, device: torch.device | str) -> bool:
+        """Mark a prepared, fixed-layout provider ready for hot dispatch.
+
+        ``ParallelLMHead`` calls this only after its TP-wide BF16/shape
+        contract and :meth:`prepare` have succeeded.  The dynamic row count
+        remains checked by :meth:`fast_call`; skipping the first-call shape
+        reductions removes setup collectives from the first decode request.
+        Generic users should keep using :meth:`__call__`, which performs the
+        defensive admission checks itself.
+        """
+        if _is_cuda_graph_capturing():
+            raise RuntimeError(
+                'symmetric-memory admission must run before CUDA Graph '
+                'capture')
+        state = self._state
+        if state is None or state is self._UNINIT:
+            return False
+        device = torch.device(device)
+        if device.type == 'cuda' and device.index is None:
+            device = torch.device('cuda', torch.cuda.current_device())
+        local_ready = (device == getattr(state, 'device', None)
+                       and getattr(state, 'rank_in_group', -1) == self._rank
+                       and getattr(state, 'world_size', -1)
+                       in _SUPPORTED_WORLD_SIZES
+                       and getattr(state, 'hidden_dim', -1)
+                       == self._gathered_width)
+        if not self.agree(local_ready, device):
+            self.release()
+            self._state = None
+            return False
+        self._runtime_admitted = True
+        return True
+
+    def _uses_real_process_group(self) -> bool:
+        """Whether setup calls can safely use TP collectives."""
+        return (dist.is_initialized()
+                and isinstance(self._group, dist.ProcessGroup))
+
+    def agree(self, local_ready: bool,
+              device: torch.device | str) -> bool:
+        """Return TP-wide readiness so ranks never split collective paths.
+
+        This setup-only collective is intentionally forbidden during CUDA
+        Graph capture.  Runtime kernels use the immutable decision recorded by
+        their owner after this method returns.
+        """
+        # ``is_current_stream_capturing`` itself can query the CUDA driver and
+        # raise on CPU-only/unit-test processes.  Consensus is only used as a
+        # setup operation, so avoid that probe when CUDA is unavailable.
+        if _is_cuda_graph_capturing():
             raise RuntimeError('TP readiness consensus is not graph capturable')
+        # Keep the helper usable by CPU/unit-test callers that construct a
+        # provider without a real process group.  Production CUDA paths always
+        # pass an initialized ``ProcessGroup`` and therefore take the
+        # collective branch below.
+        if not self._uses_real_process_group():
+            return bool(local_ready)
+        device = torch.device(device)
+        # A CPU/Gloo test group may call this helper without CUDA.  Production
+        # symmetric-memory groups remain CUDA/NCCL and keep the fast device
+        # reduction below.
+        if device.type == 'cuda' and not torch.cuda.is_available():
+            device = torch.device('cpu')
         ready = torch.tensor(int(local_ready), dtype=torch.int32, device=device)
         dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=self._group)
         return bool(ready.item())
 
     def release(self) -> None:
-        """Drop the device arena so model offload can reclaim its storage."""
+        """Drop device arenas so model offload can reclaim their storage."""
+        state = self._state
+        state_device = getattr(state, 'device', None)
+        if (state is not self._UNINIT and state is not None
+                and state_device is not None
+                and torch.device(state_device).type == 'cuda'
+                and _is_cuda_graph_capturing()):
+            # A captured graph may still hold the arena address.  Clearing the
+            # Python reference here would make a later replay use freed
+            # storage; release only after capture has ended and the stream is
+            # quiescent.
+            raise RuntimeError(
+                'cannot release symmetric-memory arena during CUDA Graph '
+                'capture')
+        if (state is not self._UNINIT and state is not None
+                and state_device is not None
+                and torch.device(state_device).type == 'cuda'
+                and torch.cuda.is_available()
+                and not _is_cuda_graph_capturing()):
+            # Device transitions are quiescent in the model-agent lifecycle,
+            # but a prior launch can still be queued on this stream. Synchronize
+            # before dropping the symmetric arena.
+            # This synchronization is off the steady-state forward path.
+            torch.cuda.current_stream(state_device).synchronize()
         self._state = self._UNINIT if self._enabled else None
         self._graph_ready = False
+        self._graph_ready_shapes.clear()
         self._runtime_admitted = False
 
+        self._logged_dispatch = False
+
     def _build(self, device: torch.device):
         device = torch.device(device)
         if device.type == 'cuda' and device.index is None:
             device = torch.device('cuda', torch.cuda.current_device())
-        if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
+        if device.type != 'cuda':
+            return None
+        if _is_cuda_graph_capturing():
             # Can't allocate under capture; retry later.
             return self._UNINIT
-        if self._gathered_width % _NUMEL_PER_THREAD != 0:
-            return None
         world_size = dist.get_world_size(self._group)
-        # tl.arange requires a power-of-two extent. Group size is identical on
-        # all ranks, so this fallback decision cannot split the protocol.
-        if world_size not in _SUPPORTED_WORLD_SIZES:
+        # Validate all rank-local construction inputs before any rank enters
+        # symmetric-memory rendezvous.  A malformed rank must join the same
+        # reduction as its peers and cause a uniform NCCL fallback instead of
+        # making the remaining ranks wait forever in ``rendezvous``.  The
+        # non-ProcessGroup branch is retained for lightweight fake providers;
+        # production builders always use the collective branch.
+        if self._uses_real_process_group():
+            group_rank = dist.get_rank(self._group)
+            try:
+                capability_ok = (torch.cuda.get_device_capability(device)
+                                 >= (9, 0))
+            except (RuntimeError, AssertionError):
+                capability_ok = False
+            local_contract = (
+                self._rank == group_rank
+                and self._max_tokens > 0
+                and self._gathered_width > 0
+                and self._gathered_width % _NUMEL_PER_THREAD == 0
+                and world_size in _SUPPORTED_WORLD_SIZES
+                and capability_ok
+            )
+            if not self.agree(local_contract, device):
+                return None
+        elif (self._max_tokens <= 0 or self._gathered_width <= 0
+              or self._gathered_width % _NUMEL_PER_THREAD != 0
+              or world_size not in _SUPPORTED_WORLD_SIZES):
             return None
+        # tl.arange requires a power-of-two extent.  Group size is identical
+        # on all ranks, so this fallback decision cannot split the protocol.
 
-        # Allocate locally first, then make one TP-wide admission decision. No
-        # rank may rendezvous while a peer can still take a local fallback.
+        # Allocate locally first, then make one TP-wide admission decision.
+        # No rank may enter symmetric-memory rendezvous while a peer can still
+        # take a local allocation fallback, otherwise the peer remains stuck
+        # in this collective forever.
         comm_buff = None
         allocation_error = None
         try:
-            comm_buff = _allocate_symmetric_buffer(self._group, self._max_tokens, self._gathered_width, device)
+            comm_buff = _allocate_symmetric_buffer(
+                self._group,
+                self._max_tokens,
+                self._gathered_width,
+                device,
+            )
         except Exception as exc:
             allocation_error = exc
-        if not self.agree(comm_buff is not None, device):
+        allocation_ok = comm_buff is not None
+        if self._uses_real_process_group():
+            allocation_ok = False
+            try:
+                allocation_ok = (isinstance(comm_buff, torch.Tensor)
+                                 and comm_buff.shape ==
+                                 (self._max_tokens, self._gathered_width)
+                                 and comm_buff.dtype == torch.bfloat16
+                                 and comm_buff.device == device
+                                 and comm_buff.is_contiguous()
+                                 and comm_buff.storage_offset() == 0
+                                 and comm_buff.data_ptr() % 16 == 0)
+            except Exception:
+                allocation_ok = False
+        if not self.agree(allocation_ok, device):
             if self._rank == 0:
                 logger.warning(
-                    'multimem all-gather disabled because a TP rank could not allocate its symmetric arena%s',
-                    f': {allocation_error}' if allocation_error else '',
-                )
+                    'multimem all-gather disabled because a TP rank could not '
+                    'allocate its symmetric arena%s',
+                    f': {allocation_error}' if allocation_error else '')
+            # Successful ranks may have allocated a private arena while a
+            # peer failed.  Drop that reference before taking NCCL fallback.
+            comm_buff = None
             return None
         if comm_buff is None:
             raise RuntimeError('TP admitted a missing symmetric-memory arena')
 
-        # Every rank is committed from this point. A rendezvous error must
-        # propagate; turning it into a rank-local fallback can deadlock peers.
+        # From this point every rank is committed to the collective.  Do not
+        # catch rendezvous errors and attempt a local fallback.
         state = create_state(
             group=self._group,
             rank_in_group=self._rank,
@@ -594,10 +1145,80 @@ def _build(self, device: torch.device):
             device=device,
             comm_buff=comm_buff,
         )
-        multicast_ready = state.symm_mem_hdl.rank == self._rank and state.symm_mem_hdl.multicast_ptr != 0
+
+        handle = state.symm_mem_hdl
+        # Handle fields are read after a collective rendezvous. Normalize all
+        # of them under one guard so a malformed/downstream handle still
+        # reaches the final TP-wide ``agree`` instead of throwing on one rank
+        # while its peers wait there.
+        multicast_ready = False
+        try:
+            handle_rank_raw = getattr(handle, 'rank', None)
+            handle_world_raw = getattr(handle, 'world_size', None)
+            handle_rank = (int(handle_rank_raw)
+                           if handle_rank_raw is not None else -1)
+            handle_world = (int(handle_world_raw)
+                            if handle_world_raw is not None else -1)
+            multicast_ptr = int(getattr(handle, 'multicast_ptr', 0) or 0)
+            signal_ptrs = getattr(handle, 'signal_pad_ptrs_dev', None)
+            # Torch 2.13 exposes this field as a raw device address; a few
+            # downstream builds wrap it in an address-bearing Tensor.
+            if hasattr(signal_ptrs, 'data_ptr'):
+                signal_addr = int(signal_ptrs.data_ptr())
+            else:
+                signal_addr = int(signal_ptrs or 0)
+            # Device addresses are unsigned in the real handle.  Requiring a
+            # strictly positive aligned value also rejects malformed/mock
+            # handles that happen to satisfy ``(-16) % 8 == 0``.
+            signal_ready = signal_addr > 0 and signal_addr % 8 == 0
+            if signal_ready and hasattr(signal_ptrs, 'numel'):
+                signal_ready = int(signal_ptrs.numel()) >= world_size
+            signal_pad_size = getattr(handle, 'signal_pad_size', None)
+            if signal_pad_size is None:
+                # ``signal_pad_size`` is not exposed on every PyTorch handle;
+                # in those builds the process-wide symmetric-memory setting is
+                # the authoritative bound.  Query it only during admission
+                # (never from the decode path), and fail closed when an
+                # explicitly available API reports a short/invalid pad.
+                get_pad_size = getattr(symm_mem, 'get_signal_pad_size', None)
+                if callable(get_pad_size):
+                    try:
+                        signal_pad_size = int(get_pad_size())
+                    except Exception:
+                        signal_pad_size = -1
+            if signal_pad_size is not None:
+                # Per-CTA slots occupy the complete configured range.  A
+                # downstream Torch build may expose a smaller pad in its
+                # handle (or globally) even when the pointer itself is
+                # non-null; reject it before any kernel can index past the
+                # allocation.
+                try:
+                    signal_ready = (signal_ready
+                                    and int(signal_pad_size) >=
+                                    _MAX_BLOCKS * world_size * 4)
+                except (TypeError, ValueError, OverflowError):
+                    signal_ready = False
+            multicast_ready = (handle_rank == self._rank
+                               and handle_world == world_size
+                               and multicast_ptr > 0
+                               and multicast_ptr % 16 == 0
+                               and signal_ready)
+        except Exception:
+            multicast_ready = False
         if not self.agree(multicast_ready, device):
             if self._rank == 0:
-                logger.warning('multimem all-gather disabled (invalid TP-wide multicast handle for world_size=%d)',
-                               state.world_size)
+                logger.warning(
+                    'multimem all-gather disabled (invalid TP-wide symmetric '
+                    'handle for world_size=%d)', state.world_size)
+            state.comm_buff = None
             return None
+        if self._rank == 0:
+            # Ray workers normally inherit the server's WARNING log level.
+            logger.warning(
+                'multimem all-gather enabled (world_size=%d, '
+                'gathered_width=%d, max_tokens=%d)',
+                state.world_size,
+                self._gathered_width,
+                state.max_token_num,
+            )
         return state
diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py
index de7eddca16..eefeeb8eac 100644
--- a/lmdeploy/pytorch/envs.py
+++ b/lmdeploy/pytorch/envs.py
@@ -238,8 +238,31 @@ def _patched_get_env(
     # cuda communicator
     enable_flashinfer_allreduce = env_to_bool('LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE', False)
     enable_symm_mem_allreduce = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE', False)
+    # Keep NCCL as the production default.  Symmetric-memory LM-head support
+    # is selected only when this process-start environment flag is explicitly
+    # set to ``1``; all tuning flags below are inert otherwise.
     enable_symm_mem_lmhead = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD', False)
     symm_mem_lmhead_max_mb = max(1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB', 64))
+    # ``0`` means use the shape-aware launch policy. Non-zero values are
+    # deterministic overrides useful for paired tuning on a fixed GPU.
+    symm_mem_lmhead_blocks = max(
+        0, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_BLOCKS', 0))
+    symm_mem_lmhead_block_threads = max(
+        0, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_BLOCK_THREADS', 0))
+    symm_mem_lmhead_autotune = env_to_bool(
+        'LMDEPLOY_SYMM_MEM_LMHEAD_AUTOTUNE', True)
+    # ``auto`` uses one cross-rank barrier CTA for larger payload grids and
+    # keeps the original per-CTA protocol for tiny grids.  The split form
+    # avoids multiplying signal CAS operations by the payload CTA count.
+    symm_mem_lmhead_barrier_mode = env_to_choice(
+        'LMDEPLOY_SYMM_MEM_LMHEAD_BARRIER_MODE', 'auto',
+        {'auto', 'per_block', 'single'})
+    # Minimum flattened token rows for the symmetric-memory LM-head gather.
+    # Smaller calls use the portable NCCL ``all_gather_into_tensor`` path;
+    # this is useful for decode/M=1 where launch overhead can exceed the
+    # communication savings.  Set this to ``1`` to force V1 for every shape.
+    symm_mem_lmhead_min_tokens = max(
+        1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MIN_TOKENS', 2))
 
     # opt-ttft
     opt_ttft_policy = env_to_choice('LMDEPLOY_PT_TTFT_POLICY', 'size', {'fifo', 'size'})
diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py
index 68161eafdd..9d47ff9200 100644
--- a/lmdeploy/pytorch/nn/embedding.py
+++ b/lmdeploy/pytorch/nn/embedding.py
@@ -167,6 +167,7 @@ def __init__(
         self._symm_mem_gatherer = None
         self._symm_mem_device = self.weight.device
         self._symm_mem_dtype = self.weight.dtype
+        self._logged_nccl_small_token_fallback = False
         if self.all_reduce and self.weight.device.type == 'cuda':
             device = self.weight.device
             if device.index is None:
@@ -178,7 +179,14 @@ def __init__(
             gathered_width = self.tp * self.vocab_size_padded
             capacity = _envs.symm_mem_lmhead_max_mb * 1024 * 1024
             max_tokens = capacity // (gathered_width * torch.bfloat16.itemsize)
-            same_config = _tp_same_config((capacity, gathered_width, max_tokens), device, self.tp_group)
+            barrier_mode = {'auto': 0, 'per_block': 1, 'single': 2}.get(
+                _envs.symm_mem_lmhead_barrier_mode, -1)
+            same_config = _tp_same_config(
+                (capacity, gathered_width, max_tokens,
+                 _envs.symm_mem_lmhead_blocks,
+                 _envs.symm_mem_lmhead_block_threads,
+                 int(_envs.symm_mem_lmhead_autotune), barrier_mode,
+                 _envs.symm_mem_lmhead_min_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')
@@ -198,7 +206,14 @@ def __init__(
                                     rank=self.tp_rank,
                                     gathered_width=gathered_width,
                                     max_tokens=max_tokens)
-            if gatherer.prepare(device):
+            prepared = gatherer.prepare(device)
+            admit_static = getattr(gatherer, 'admit_static', None)
+            if prepared and callable(admit_static):
+                # Freeze the static BF16/layout contract during model setup;
+                # the first decode call can then dispatch directly without
+                # another TP-wide admission reduction.
+                prepared = admit_static(device)
+            if prepared:
                 self._symm_mem_gatherer = gatherer
 
     def tie_weights(self, embedding: ParallelEmbedding):
@@ -220,8 +235,13 @@ def _apply(self, fn, recurse=True):
             self._symm_mem_dtype = current_dtype
             if current_dtype != torch.bfloat16:
                 self._symm_mem_gatherer = None
-            elif current_device.type == 'cuda' and not gatherer.prepare(current_device):
-                self._symm_mem_gatherer = None
+            elif current_device.type == 'cuda':
+                prepared = gatherer.prepare(current_device)
+                admit_static = getattr(gatherer, 'admit_static', None)
+                if prepared and callable(admit_static):
+                    prepared = admit_static(current_device)
+                if not prepared:
+                    self._symm_mem_gatherer = None
         return result
 
     def get_local_logits(self, hidden_states: torch.Tensor):
@@ -237,10 +257,24 @@ def all_gather_logits(self, local_logits: torch.Tensor) -> torch.Tensor:
 
         if self._symm_mem_gatherer is not None:
             local_logits_2d = local_logits.reshape(-1, local_logits.shape[-1])
-            gathered = self._symm_mem_gatherer(local_logits_2d)
-            if gathered is not None:
-                output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], )
-                return gathered.reshape(output_shape)[..., :self.vocab_size]
+            if local_logits_2d.shape[0] >= _envs.symm_mem_lmhead_min_tokens:
+                # ParallelLMHead has a stable BF16/layout contract after the
+                # first collective admission.  Use the admitted fast path so
+                # steady-state decode avoids repeated host-side validation;
+                # arbitrary gatherer users retain the defensive __call__ API.
+                fast_call = getattr(self._symm_mem_gatherer, 'fast_call', None)
+                gathered = (fast_call(local_logits_2d)
+                            if callable(fast_call) else
+                            self._symm_mem_gatherer(local_logits_2d))
+                if gathered is not None:
+                    output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], )
+                    return gathered.reshape(output_shape)[..., :self.vocab_size]
+            elif not self._logged_nccl_small_token_fallback:
+                logger.info(
+                    'symmetric-memory LM-head fallback to NCCL for token_rows=%s '
+                    '(minimum=%s)', local_logits_2d.shape[0],
+                    _envs.symm_mem_lmhead_min_tokens)
+                self._logged_nccl_small_token_fallback = True
 
         input_size = local_logits.size()
         output_size = (input_size[0] * self.tp, ) + input_size[1:]