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

Neuron: CP group as a replicated-decode group; MoE all-reduce mode (#28)

Browse Source 
- Followers claim core cp_rank (DEEPINFER_NEURON_LOCAL_RANK at spawn) and get
  the graph's host metadata through shm (no in-graph broadcast on Neuron),
  including the prompt-as-decode bucket and the context limit; they skip the
  sampler. max_num_prefilling_tokens is not divided by cp on Neuron (prompt
  chunks run replicated).
- DEEPINFER_NEURON_NUM_KV_BLOCKS: the fixed KV pool as an env var, so spawned
  CP followers take it too.
- static MoE: when the EP group is the CP group the tokens are already
  replicated: no all-gather, all-reduce instead of reduce-scatter (64x fewer
  expert rows a rank at cp=ep=64).
- run_deepinfer_tiny.py --cp. Host tiny: cp=2 and cp=4 match the cp=1 tokens
  (short 8/9 at a near tie, long 9/9).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 22:46:04 +08:00
parent d188d1b3d1
commit f9cbf6f2e5
9 changed files with 84 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
deepinfer/engine/llm_engine.py
View File

@ -392,7 +392,17 @@ class LLMEngine:
args=args,
daemon=True,
)
worker.start()
# Neuron: the follower claims core cp_rank of CUDA_VISIBLE_DEVICES (deepinfer/platform.py); the spawned
# process copies this environment at start.
saved = os.environ.get("DEEPINFER_NEURON_LOCAL_RANK")
os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = str(cp_rank)
try:
worker.start()
finally:
if saved is None:
os.environ.pop("DEEPINFER_NEURON_LOCAL_RANK", None)
else:
os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = saved
logger.info("_init_workers: non-driver workers started")
driver_kwargs = {

2
deepinfer/environ.py
View File

@ -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: 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_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True)
# Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).

4
deepinfer/model_executor/model_loader.py
View File

@ -128,7 +128,9 @@ def finalize_model_config(model_cfg: AttrDict) -> AttrDict:
ep_size = ep_group.size() if ep_group is not None else 1
model_cfg.gpt.max_num_decoding_tokens = model_cfg.gpt.max_num_seqs * (1 + envs.DEEPINFER_USE_NEXTN_DECODE)
model_cfg.gpt.max_num_prefilling_tokens = model_cfg.gpt.max_num_batched_tokens // PG.context_parallel_group().size()
# Neuron runs prompt chunks replicated across the CP group (through the decode graphs), not sharded.
cp_split = 1 if PLATFORM.is_neuron else PG.context_parallel_group().size()
model_cfg.gpt.max_num_prefilling_tokens = model_cfg.gpt.max_num_batched_tokens // cp_split
min_num_extra_experts = envs.DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS
_configure_moe_redundancy(model_cfg.gpt, min_num_extra_experts, ep_size)

22
deepinfer/neuron_ops/static_moe.py
View File

@ -66,6 +66,18 @@ def _all_gather(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: in
return funcol.wait_tensor(funcol.all_gather_tensor(t, 0, group))
def tokens_replicated(moe: DeepseekMoE) -> bool:
"""Every EP rank holds the same tokens: the EP group is the CP group, whose members run one replicated batch
(Neuron, issue #28). The MoE then needs no all-gather, and an all-reduce instead of a reduce-scatter."""
from deepinfer.model_executor.parallel_utils.parallel_state import PG
return moe.ep_size > 1 and PG.context_parallel_group().size() == moe.ep_size
def _all_reduce(t: torch.Tensor, group: torch.distributed.ProcessGroup) -> torch.Tensor:
return funcol.wait_tensor(funcol.all_reduce(t, "sum", group))
def _reduce_scatter(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: int, rank: int) -> torch.Tensor:
if size == 1:
return t
@ -99,7 +111,8 @@ def static_moe_forward(
cap = max(moe.max_num_decoding_tokens, num_tokens)
parts = x_parts + [topk_idx.to(torch.int32), topk_weights.to(torch.float32)]
if size == 1:
replicated = tokens_replicated(moe)
if size == 1 or replicated:
# Nothing to gather. Skipping the byte packing also keeps dtype bitcasts of graph inputs out of the region
# graph, which neuronx-cc mishandles (InferNonlocalTensors, issue #25).
g_parts = parts
@ -130,7 +143,7 @@ def static_moe_forward(
_expert_gemm(down_in, down_w, down, e, moe.use_fp8_act, recipe_a, recipe_b)
acc = acc + down.float()
out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
out = _all_reduce(acc, group) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
shared = moe.shared_forward(hidden_states)
if shared is not None:
out = out + shared.float()
@ -193,7 +206,8 @@ def _nki_moe_forward(
# Local expert ids relative to this rank, from a rank *tensor*: the graph stays identical across ranks.
idx = topk_idx.to(torch.int32)
w = topk_weights.to(torch.float32)
if size > 1:
replicated = tokens_replicated(moe)
if size > 1 and not replicated:
pad = cap - num_tokens
x = torch.nn.functional.pad(x, (0, 0, 0, pad))
idx = torch.nn.functional.pad(idx, (0, 0, 0, pad), value=-1)
@ -221,7 +235,7 @@ def _nki_moe_forward(
)
acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float()
regions.split()
out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
out = _all_reduce(acc, group) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
regions.split()
if moe.has_shared:
shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)

6
deepinfer/worker/base.py
View File

@ -216,6 +216,12 @@ class WorkerBase:
Profiles the peak memory usage of the model and returns the maximum
number of GPU and CPU cache blocks that can be allocated.
"""
if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_NUM_KV_BLOCKS > 0:
# A fixed pool instead of the profile prefill (it does not fit next to the Neuron graphs). An env var, not a
# patch of this method, so CP followers (spawned processes) take the same pool as their driver.
self.model_runner.set_block_size(self.cache_config.block_size)
self.cache_config.num_gpu_blocks = envs.DEEPINFER_NEURON_NUM_KV_BLOCKS
return envs.DEEPINFER_NEURON_NUM_KV_BLOCKS
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()

35
deepinfer/worker/model_runner.py
View File

@ -249,6 +249,19 @@ class ModelRunner:
if not run_cuda_graph:
assert hasattr(input_metadata, "is_prompt")
input_metadata_dict.update(input_metadata.to_flatten_dict())
elif PLATFORM.is_neuron:
# No in-graph broadcast on Neuron (the GPU graphs broadcast their device metadata themselves):
# send the graph's host metadata through shm; followers copy it into their own source buffers.
input_metadata_dict["neuron_graph"] = {
"graph_size": graph_size,
"prompt_as_decode": prompt_as_decode,
"ctx_limit": getattr(input_metadata, "neuron_ctx_limit", None),
"source": {
k: v.clone()
for k, v in cuda_graph_used.input_metadata_source.to_flatten_dict().items()
if isinstance(v, torch.Tensor)
},
}
broadcast_tensor_dict(input_metadata_dict, src=0)
return input_metadata, total_q, run_cuda_graph, cuda_graph_used, graph_size
@ -260,7 +273,21 @@ class ModelRunner:
run_cuda_graph: bool = input_metadata_dict.pop("run_cuda_graph")
total_q = input_metadata_dict["total_q"]
num_reqs = input_metadata_dict["num_reqs"]
if run_cuda_graph:
neuron_graph = input_metadata_dict.pop("neuron_graph", None)
if run_cuda_graph and neuron_graph is not None:
graph_size = neuron_graph["graph_size"]
cuda_graph_used = self.cuda_graph_runner.graphs[graph_size]
source = cuda_graph_used.input_metadata_source
for key, value in neuron_graph["source"].items():
name, _, sub = key.partition(".")
target = getattr(source, name)
(target[sub] if sub else target).copy_(value)
input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
input_metadata.neuron_ctx_limit = neuron_graph["ctx_limit"]
if neuron_graph["prompt_as_decode"]:
input_metadata = copy.copy(input_metadata)
input_metadata.is_prompt = True
elif run_cuda_graph:
graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q)
input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
else:
@ -287,7 +314,7 @@ class ModelRunner:
yoco_prefill_truncation_mode,
envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS,
)
if PG.context_parallel_group().size() > 1:
if PG.context_parallel_group().size() > 1 and not PLATFORM.is_neuron:
decode_inputs = sampler_inputs.decode
metadata_dict = {
"is_prompt": input_metadata.is_prompt,
@ -385,6 +412,10 @@ class ModelRunner:
IMPORTANT: This function must NOT contain any cuda synchronize.
"""
if PLATFORM.is_neuron and not self.is_driver_worker:
# Only the driver samples; a follower's sampler work (lm_head, prompt logprobs) has no collectives on
# Neuron, so it is skipped together with the sampler-input broadcast.
return SamplerOutput(post_funcs=[])
sampler_inputs = self.prepare_sampler_inputs(seq_metadata_list, input_metadata)
kwargs = {
"input_metadata": input_metadata,

1
examples/neuron/bench_neuron.py
View File

@ -50,6 +50,7 @@ dp_size = args.ep
def _fixed_kv_pool(num_blocks: int) -> None:
os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_blocks) # CP followers read it too
from deepinfer.worker.base import WorkerBase
def profile_num_available_blocks(self): # noqa: ANN001, ANN202

1
examples/neuron/offline_inference_neuron.py
View File

@ -48,6 +48,7 @@ dp_size = args.ep
def _fixed_kv_pool(num_blocks: int) -> None:
os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_blocks) # CP followers read it too
"""Skip WorkerBase.profile_num_available_blocks (a max_num_batched_tokens dummy prefill)."""
from deepinfer.worker.base import WorkerBase

14
examples/neuron/run_deepinfer_tiny.py
View File

@ -46,6 +46,9 @@ def parse_args():
ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)")
ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards")
ap.add_argument("--attn-p-bf16", action="store_true", help="experiment: patch the flash_mla shim with the proposed bf16-P fix")
ap.add_argument("--cp", type=int, default=1,
help="context-parallel group of one engine (decode replicated on every member, EP over the group): "
"CUDA_VISIBLE_DEVICES lists its cores (default 0..cp-1)")
ap.add_argument("--ep", type=int, default=1,
help="expert parallel size (= data parallel size): one engine process per rank, every rank runs the "
"same teacher-forced prompts. Rank r uses entry r of CUDA_VISIBLE_DEVICES (default 0..ep-1), "
@ -54,14 +57,14 @@ def parse_args():
return ap.parse_args()
def setup_env(device: str) -> None:
def setup_env(device: str, cp: int = 1) -> None:
# Must happen before deepinfer is imported (see examples/_config.py).
os.environ["DEEPINFER_ACCELERATOR"] = "neuron"
os.environ["DEEPINFER_NEURON_DEVICE"] = device
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1")
os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") # one engine rank; mapped to the target device
os.environ.setdefault("CUDA_VISIBLE_DEVICES", ",".join(str(i) for i in range(cp))) # mapped to the target
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
@ -85,6 +88,7 @@ def install_workarounds(num_kv_blocks: int, attn_p_bf16: bool = False) -> None:
# max_num_batched_tokens tokens (903 s on the host for 2048 tokens) and then sizes the KV cache from
# mem_get_info(), which on cpu is host RAM -> 0.5 * 2 TB = 559 GiB of KV blocks. Use a fixed small pool instead.
if num_kv_blocks > 0:
os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_kv_blocks) # CP followers read it too
import torch.distributed as _dist # noqa: F401
from deepinfer.worker.base import WorkerBase
@ -247,7 +251,7 @@ def _ep_rank_main(args, rank: int, core: str) -> None:
def run(args, rank: int = 0):
setup_env(args.device)
setup_env(args.device, args.cp)
os.makedirs(args.out, exist_ok=True)
import torch
@ -271,8 +275,8 @@ def run(args, rank: int = 0):
llm = LLM(
model=args.hf_ckpt,
trust_remote_code=True,
expert_parallel_size=args.ep,
context_parallel_size=1,
expert_parallel_size=args.ep * args.cp,
context_parallel_size=args.cp,
data_parallel_size=args.ep,
engine_rank=rank,
chunked_prefill=True,