diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index cad3083..787fa22 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -28,6 +28,7 @@ import functools import os from typing import Callable, List, Optional, Tuple +import numpy as np import torch import torch.distributed as dist @@ -279,14 +280,20 @@ class EngramBuffer(BufferBase): flat = indices.reshape(-1).to(device="cpu", dtype=torch.int64) _rank_timeline(f"eg_d2h{layer_idx}") flat = flat.clamp(0, self._storages[layer_idx].size(0) - 1) - rows = self._storages[layer_idx].index_select(0, flat).view(self.dtype) + _rank_timeline(f"eg_clamp{layer_idx}") + # numpy, not torch.index_select: torch splits these few rows over its intra-op pool, whose threads contend with + # every rank's for the node's CPUs (12-26 ms per call in the 64-rank run vs 0.06 ms here). + idx = flat.numpy() + rows = torch.from_numpy(np.take(self._storages[layer_idx].numpy(), idx, axis=0)) + _rank_timeline(f"eg_isel{layer_idx}") + rows = rows.view(self.dtype) _rank_timeline(f"eg_gather{layer_idx}") fetched = rows.reshape(num_tokens, num_entries_per_token * self.hidden).to(device) _rank_timeline(f"eg_h2d{layer_idx}") if not self._sf_shards: return fetched sf_table, sf_dtype = self._sf_shards[layer_idx] - sf_rows = sf_table.index_select(0, flat).view(sf_dtype) + sf_rows = torch.from_numpy(np.take(sf_table.numpy(), idx, axis=0)).view(sf_dtype) out = fetched, sf_rows.reshape(num_tokens, num_entries_per_token * self.num_sf_packs).to(device) _rank_timeline(f"eg_sf{layer_idx}") return out diff --git a/deepinfer/utils.py b/deepinfer/utils.py index 62facdf..45d37ec 100644 --- a/deepinfer/utils.py +++ b/deepinfer/utils.py @@ -146,10 +146,12 @@ def _neuron_bind_cpus(local_rank: int) -> list[int] | None: The generic split (physical cores // ranks) gave each of 64 ranks one logical CPU, numbered by rank: the step's dispatch thread and the runtime's threads shared it, and devices 0-3 (NUMA 1 on trn2.48xlarge) got CPUs of NUMA 0. DEEPINFER_NEURON_CPU_BIND: ``node`` (default) the device's whole NUMA node, ``core`` one physical core and its - hyperthread on that node, ``legacy`` the generic split. None falls back to the generic split.""" + hyperthread on that node, ``all`` every CPU of the host, ``legacy`` the generic split. None falls back to the generic split.""" mode = os.environ.get("DEEPINFER_NEURON_CPU_BIND", "node") if mode == "legacy": return None + if mode == "all": + return list(range(psutil.cpu_count(logical=True))) import json import subprocess diff --git a/deepinfer/worker/base.py b/deepinfer/worker/base.py index 2cff6a7..734e5e8 100644 --- a/deepinfer/worker/base.py +++ b/deepinfer/worker/base.py @@ -182,8 +182,11 @@ class WorkerBase: torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False torch._C._accelerator_setAllocatorSettings("expandable_segments:True") - torch.set_num_threads(8) - torch.set_num_interop_threads(8) + # Neuron: 64 ranks share the host's CPUs, so 8 intra-op threads each oversubscribe them and every small host op + # (the Engram gather) waits on preempted pool threads. DEEPINFER_NEURON_TORCH_THREADS overrides. + threads = int(os.environ.get("DEEPINFER_NEURON_TORCH_THREADS", "8")) if PLATFORM.is_neuron else 8 + torch.set_num_threads(threads) + torch.set_num_interop_threads(threads) torch.compiler.config.force_cudagraph_gc = False def init_model(self) -> None: