From d44cf253c58e4c7cfb8961ae1692f0c6cac7c4c5 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 00:37:23 +0800 Subject: [PATCH] offline_inference_neuron.py --cp Co-Authored-By: Claude Opus 5.5 --- examples/neuron/offline_inference_neuron.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py index 09c2050..01be87e 100644 --- a/examples/neuron/offline_inference_neuron.py +++ b/examples/neuron/offline_inference_neuron.py @@ -36,6 +36,7 @@ PROMPTS = [ parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (basic mode)") 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("--cp", default=1, type=int, help="TP group size: attention heads sharded over cp ranks running one batch (issue #28)") parser.add_argument("--max_tokens", default=16, type=int) parser.add_argument("--num_prompts", default=len(PROMPTS), type=int, help="run only the first N prompts (eager prefill compiles per prompt length)") 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") @@ -44,7 +45,7 @@ parser.add_argument("--max_num_batched_tokens", default=1024, type=int) parser.add_argument("--max_num_seqs", default=1, type=int, help="decode batch; every captured batch size adds one set of region NEFFs per layer (HBM: model code)") args = parser.parse_args() -dp_size = args.ep +dp_size = args.ep // args.cp def _fixed_kv_pool(num_blocks: int) -> None: @@ -69,14 +70,14 @@ def main(dp_size: int, dp_rank: int) -> None: # PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init. from deep_cluster.hardware import get_actual_num_gpus - setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus()) + setup_cuda_visible_devices(dp_rank, args.cp, get_actual_num_gpus()) _fixed_kv_pool(args.num_kv_blocks) llm = create_llm( args.model_name, dp_rank, ep_size=args.ep, - cp_size=1, + cp_size=args.cp, env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"}, data_parallel_size=dp_size, max_num_batched_tokens=args.max_num_batched_tokens,