From ad6b47a141105db5f7c037b95bdffea981f93586 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 22:53:32 +0800 Subject: [PATCH] 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 --- deepinfer/environ.py | 2 + .../model_executor/models/v4_modules/attn.py | 67 ++++++++++++++++++- 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 8cb6971..b456230 100644 --- a/deepinfer/environ.py +++ b/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. diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index 7564ed0..ca56c3a 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/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