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