From 0d2cbcb112748cdb68079d516af1ea28e8366546 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 21 Oct 2026 21:00:54 +0800 Subject: [PATCH] Neuron basic mode: try prefill on every step The DP-wide prefill vote runs every DEEPINFER_PREFILL_ATTEMPT_INTERVAL (70) / ..._ANY_ (100) decode steps. Requests that become schedulable after step 0 (multimodal-processor round trip) then wait 70 placeholder steps; with eager execution a full-model step on the NeuronCore takes ~150 s, so the smoke test sat ~3 hours in placeholder steps. Default both intervals to 1 on Neuron. offline_inference_neuron: force-topk off (#22), 400-block KV pool. Co-Authored-By: Claude Opus 5.5 --- deepinfer/__init__.py | 4 ++++ examples/neuron/offline_inference_neuron.py | 5 +++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/deepinfer/__init__.py b/deepinfer/__init__.py index ddc152d..fc086be 100644 --- a/deepinfer/__init__.py +++ b/deepinfer/__init__.py @@ -47,6 +47,10 @@ if _so_platform_tag().startswith("neuron-"): os.environ.setdefault("DEEPINFER_DFLASH_ENFORCE_EAGER", "1") os.environ.setdefault("DEEPINFER_USE_NEXTN_DECODE", "0") os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0") + # Try prefill on every step: by default the DP-wide prefill vote runs every 70 / 100 decode steps, and an eager + # step on the NeuronCore takes minutes for the full model (a request arriving between votes waited ~3 hours). + os.environ.setdefault("DEEPINFER_PREFILL_ATTEMPT_INTERVAL", "1") + os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", "1") __version__ = "0.1.0" diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py index 3047e5f..16e35df 100644 --- a/examples/neuron/offline_inference_neuron.py +++ b/examples/neuron/offline_inference_neuron.py @@ -37,7 +37,7 @@ parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron ( parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME, help="HF model id or local HF checkpoint directory") parser.add_argument("--ep", default=64, type=int) parser.add_argument("--max_tokens", default=16, type=int) -parser.add_argument("--num_kv_blocks", default=176, type=int, help="fixed KV pool per rank; must hold max_model_len tokens") +parser.add_argument("--num_kv_blocks", default=400, type=int, help="fixed KV pool per rank; the scheduler admits a request only if max_model_len tokens fit, per running sequence") parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)") parser.add_argument("--max_num_batched_tokens", default=1024, type=int) args = parser.parse_args() @@ -59,7 +59,8 @@ def _fixed_kv_pool(num_blocks: int) -> None: def main(dp_size: int, dp_rank: int) -> None: os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0" os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron") - setup_common_env(model_name=args.model_name, sampler_force_topk=4096, set_override_experts=False, enable_eplb=False) + # sampler_force_topk=0: torch.topk with k=4096 over the vocabulary does not compile on the device (issue #22). + setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False) # Pick this rank's NeuronCore before anything initializes deepinfer.PLATFORM (importing deepinfer.worker does): # PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init. from deep_cluster.hardware import get_actual_num_gpus