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 <noreply@anthropic.com>

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.

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