From e5663a0a1d35e55aba3d0bdfbd56a9fb250b5ac7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 09:16:01 +0800 Subject: [PATCH] Neuron: attention TP over lanes of 8 consecutive ranks (DEEPINFER_NEURON_ATTN_TP) (#28) Attention is head-sharded within lanes of 8 consecutive ranks of the CP group and replicated across lanes, as vllm-neuron#71 does: a 64-rank all-reduce is ~1.9 ms of latency against ~0.1 ms for 8 ranks on two adjacent devices. With 64 heads and 8 groups each rank holds exactly one wv group. Host tiny cp=8 with lanes of 2: golden tokens up to near ties. Co-Authored-By: Claude Opus 5.5 --- deepinfer/environ.py | 3 +++ .../model_executor/models/v4_modules/attn.py | 22 +++++++++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index b456230..cbae526 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -134,6 +134,9 @@ class envs: DEEPINFER_NEURON_CTX_BUCKETS = EnvVar.str(default="") # Neuron only: shard the attention heads over the CP group (TP-style decode, issue #28) when cp > 1. DEEPINFER_NEURON_TP_ATTN = EnvVar.bool(default=True) + # Neuron only: attention TP lane width within the CP group (consecutive ranks; replicated across lanes). 0 = the + # whole CP group. + DEEPINFER_NEURON_ATTN_TP = EnvVar.int(default=8) # Neuron only: fixed KV block pool per rank (0 = the engine's memory-profiling prefill). DEEPINFER_NEURON_NUM_KV_BLOCKS = EnvVar.int(default=0) # Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step. diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index ca56c3a..afd9ee9 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/deepinfer/model_executor/models/v4_modules/attn.py @@ -240,6 +240,23 @@ class V4AttentionBase(nn.Module): return out +def _attn_tp_group(cp_group): # noqa: ANN001, ANN202 + """The attention TP group: lanes of DEEPINFER_NEURON_ATTN_TP consecutive ranks of the CP group (the attention is + replicated across lanes). A 64-rank all-reduce costs ~1.9 ms of latency against ~0.1 ms for 8 consecutive ranks + on two adjacent devices (vllm-neuron#71, which shards attention 8-way the same way).""" + import torch.distributed as dist + + from deepinfer.model_executor.parallel_utils.parallel_state import _split_group + + lane = envs.DEEPINFER_NEURON_ATTN_TP + cp = cp_group.size() + if lane <= 0 or lane >= cp: + return cp_group + assert cp % lane == 0, (cp, lane) + world = dist.get_world_size() + return _split_group([list(range(r, r + lane)) for r in range(0, world, lane)], backend=PLATFORM.distributed_backend) + + def _slice_param(p: nn.Parameter, rows: tuple[int, int] | None = None, cols: tuple[int, int] | None = None) -> None: """Keep ``p[rows, cols]`` (and the matching block scales ``p.sf``, per row or per 32-block) in place, as a copy so the full storage is freed.""" @@ -324,9 +341,10 @@ class V4Attention(V4AttentionBase): if not (PLATFORM.is_neuron and envs.DEEPINFER_NEURON_TP_ATTN): return None group = PG.context_parallel_group() - tp = group.size() - if tp == 1: + if group.size() == 1: return None + group = _attn_tp_group(group) + tp = group.size() rank = group.rank() assert getattr(attn_module, "einsum_out_clamp_value", None) is None, "a clamp after wv is not linear in the heads" heads, per_group, d = self.num_heads, self.num_heads // self.n_wv_group, self.qk_head_dim