From f9cbf6f2e5e4b4d3616409ea8682f3cc0ec0c47a Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 22:46:04 +0800 Subject: [PATCH] Neuron: CP group as a replicated-decode group; MoE all-reduce mode (#28) - 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 --- deepinfer/engine/llm_engine.py | 12 ++++++- deepinfer/environ.py | 2 ++ deepinfer/model_executor/model_loader.py | 4 ++- deepinfer/neuron_ops/static_moe.py | 22 ++++++++++--- deepinfer/worker/base.py | 6 ++++ deepinfer/worker/model_runner.py | 35 +++++++++++++++++++-- examples/neuron/bench_neuron.py | 1 + examples/neuron/offline_inference_neuron.py | 1 + examples/neuron/run_deepinfer_tiny.py | 14 ++++++--- 9 files changed, 84 insertions(+), 13 deletions(-) diff --git a/deepinfer/engine/llm_engine.py b/deepinfer/engine/llm_engine.py index cb87941..271b35c 100644 --- a/deepinfer/engine/llm_engine.py +++ b/deepinfer/engine/llm_engine.py @@ -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 = { diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 6df4fd9..8cb6971 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: 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). diff --git a/deepinfer/model_executor/model_loader.py b/deepinfer/model_executor/model_loader.py index 4af10db..81a4997 100644 --- a/deepinfer/model_executor/model_loader.py +++ b/deepinfer/model_executor/model_loader.py @@ -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) diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index ba86d74..393f4a0 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -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) diff --git a/deepinfer/worker/base.py b/deepinfer/worker/base.py index 31fc4e4..d13e3eb 100644 --- a/deepinfer/worker/base.py +++ b/deepinfer/worker/base.py @@ -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() diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 742998a..b93c4c1 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -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, diff --git a/examples/neuron/bench_neuron.py b/examples/neuron/bench_neuron.py index 96ef1d5..56899e5 100644 --- a/examples/neuron/bench_neuron.py +++ b/examples/neuron/bench_neuron.py @@ -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 diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py index c462aea..09c2050 100644 --- a/examples/neuron/offline_inference_neuron.py +++ b/examples/neuron/offline_inference_neuron.py @@ -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 diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py index 5b1dad8..7cb0228 100644 --- a/examples/neuron/run_deepinfer_tiny.py +++ b/examples/neuron/run_deepinfer_tiny.py @@ -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,