Neuron: head-sharded attention over the CP group (TP decode) (#28)

With a CP group running one replicated batch (DEEPINFER_NEURON_TP_ATTN, on
when cp > 1), each rank keeps its heads' rows of q_b, its groups' slice of wv
(or its heads' columns of one group), the matching out_proj columns and its
attention-sink entries; _wv_and_o_proj all-reduces the partial outputs over
the group (exact: wv and out_proj are linear in the heads; asserts no clamp).
Q-a / KV / indexer / compressor stay replicated.

Host tiny (8 heads, 2 groups): cp=2 (1 group a rank), cp=4 (2 heads of a
group), cp=8 (1 head) all give the golden tokens up to near ties.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/environ.py

@@ -132,6 +132,8 @@ class envs:
     # sequence fits a bucket sizes its dense compressed-KV scopes and indexer logits to the bucket instead of
     # max_model_len (one more set of attention graphs a bucket; issue #29).
     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: 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

@@ -231,7 +231,35 @@ class V4AttentionBase(nn.Module):
         if self.einsum_out_clamp_value is not None:
             wv.clamp_(min=-self.einsum_out_clamp_value, max=self.einsum_out_clamp_value)
         wv = (wv[0].flatten(1, 2), wv[1]) if self.use_fp8_einsum_output else wv.flatten(1, 2)
-        return self.out_proj(wv)
+        out = self.out_proj(wv)
+        tp_group = getattr(self, "tp_group", None)
+        if tp_group is not None:  # head-sharded attention (_neuron_shard_heads): sum the ranks' partial outputs
+            import torch.distributed._functional_collectives as funcol
+
+            out = funcol.wait_tensor(funcol.all_reduce(out.float(), "sum", tp_group)).to(out.dtype)
+        return out
+
+
+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."""
+    w = p.data
+    sf = getattr(p, "sf", None)
+
+    def cut(t: torch.Tensor, full: tuple[int, int], span: tuple[int, int] | None, dim: int) -> torch.Tensor:
+        if span is None:
+            return t
+        a, b = span
+        n = t.shape[dim]
+        if n != full[dim]:  # one scale per 32-block
+            assert n == -(-full[dim] // 32) and a % 32 == 0 and b % 32 == 0, (t.shape, full, span)
+            a, b = a // 32, -(-b // 32)
+        return t.narrow(dim, a, b - a)
+
+    full = tuple(w.shape)
+    p.data = cut(cut(w, full, rows, 0), full, cols, 1).contiguous().clone()
+    if sf is not None:
+        p.sf = cut(cut(sf, full, rows, 0), full, cols, 1).contiguous().clone()
 
 
 class V4Attention(V4AttentionBase):
@@ -268,6 +296,7 @@ class V4Attention(V4AttentionBase):
             self.is_index_source = shared_topology.is_index_source_layer(layer_idx)
             self.is_index_k_source = shared_topology.is_index_k_source_layer(layer_idx)
 
+        self.tp_group = self._neuron_shard_heads(attn_module)
         self._init_q_projections(attn_module, cfg)
 
         # RMSNorm
@@ -278,12 +307,48 @@ class V4Attention(V4AttentionBase):
         # attention sink
         assert attn_module.sink_mode == "swa"
         self.attn_sink_params = runtime_attn_sink(attn_module)
+        if self.tp_group is not None:
+            h0 = self._tp_head0
+            self.attn_sink_params = nn.Parameter(self.attn_sink_params.data[h0 : h0 + self.num_heads].clone(), requires_grad=False)
 
         self._init_compressor_and_indexer(attn_module, layer_idx, cfg, shared_topology=shared_topology)
 
         # o projection
         self._init_wv_o_proj(attn_module, cfg)
 
+    def _neuron_shard_heads(self, attn_module: nn.Module):  # noqa: ANN202
+        """Neuron (issue #28): with a CP group running one replicated decode batch, shard the attention heads over it
+        like vllm-neuron#71's TP. Each rank keeps its heads' rows of q_b, its groups' slice of wv and the matching
+        columns of out_proj; the layer output is all-reduced in _wv_and_o_proj (exact: wv and out_proj are linear in
+        the heads). Q-a / KV / indexer / compressor stay replicated. Returns the group, or None."""
+        if not (PLATFORM.is_neuron and envs.DEEPINFER_NEURON_TP_ATTN):
+            return None
+        group = PG.context_parallel_group()
+        tp = group.size()
+        if tp == 1:
+            return None
+        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
+        assert heads % tp == 0, f"{heads} heads over a TP group of {tp}"
+        h = heads // tp
+        h0 = rank * h
+        assert h % per_group == 0 or per_group % h == 0, (h, per_group)
+        g0, g1 = h0 // per_group, (h0 + h - 1) // per_group + 1
+        o = self.o_head_dim
+        # q_b: this rank's heads' rows.
+        _slice_param(attn_module.Wq.B.weight, rows=(h0 * d, (h0 + h) * d))
+        # wv [groups * o, per_group * d]: this rank's groups, and within a group its heads' columns.
+        cols = (0, per_group * d) if h >= per_group else ((h0 % per_group) * d, (h0 % per_group + h) * d)
+        _slice_param(attn_module.Wv.weight, rows=(g0 * o, g1 * o), cols=cols)
+        # out_proj [hidden, groups * o]: its groups' columns.
+        _slice_param(attn_module.out_proj.weight, cols=(g0 * o, g1 * o))
+        self.num_heads = h
+        self.n_wv_group = g1 - g0
+        self.wv_group_dim = (cols[1] - cols[0])
+        self._tp_head0 = h0
+        return group
+
     def _init_q_projections(self, attn_module: nn.Module, cfg: AttrDict) -> None:
         # With yoco_swa_share_hidden the Q and KV inputs differ, so the projections cannot be fused.
         self.fuse_kv_swa = self.use_swa and not self.yoco_swa_share_hidden