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 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-21 21:00:54 +08:00
parent d0f7e4f69c
commit 0d2cbcb112
2 changed files with 7 additions and 2 deletions

View File

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

View File

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