coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

neuron: Engram gather via numpy (torch index_select waited 12-26 ms on its contended intra-op pool); DEEPINFER_NEURON_TORCH_THREADS; CPU bind mode 'all'

Browse Source 
Decode step 522 -> 454 ms (vllm-neuron #71 on the same box: 478 ms).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-28 02:10:43 +08:00
parent 81aee64ac2
commit 8374df4b9c
3 changed files with 17 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py vendored
View File

@ -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

4
deepinfer/utils.py
View File

@ -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

7
deepinfer/worker/base.py
View File

@ -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: