From 3629d113fb49d84247bba3e485730fb7d063a758 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 03:19:56 +0800 Subject: [PATCH] Neuron: prompt chunks through the decode region graphs (issue #26) - PrepareArgs.prompt_as_decode: decode-layout metadata for a prompt chunk; input_ids are the prompt's own tokens instead of last token + MTP drafts. - DEEPINFER_NEURON_PREFILL_BUCKET adds one graph size; the model runner pads a prompt chunk to it with the decode placeholders (every EP rank takes the same bucket, so the static collectives match). Sampling still sees a prompt step. - nki_attention.paged_mqa_logits: one gather + one kernel call for a chunk (one sequence) instead of a T-way loop. - bf16 linear for the LM head eagerly on a NeuronCore (b.float() was a 2.6 GiB copy a call); DEEPINFER_NEURON_REGION_TIMING; --num_prompts. Host mode, tiny V4.1, chunk 128: same tokens as eager prefill (short 9/9, long 715 tokens in 6 chunks 8/9 against golden). Co-Authored-By: Claude Opus 5.5 --- .../input_metadata/prepare_tensor_input.hpp | 3 ++ .../prepare_tensor_input_base.cpp | 5 ++-- csrc/engine/pybind_output.cpp | 1 + deepinfer/environ.py | 6 ++++ .../model_executor/model_metadata_utils.py | 5 +++- deepinfer/neuron_ops/nki_attention.py | 26 +++++++++++++++-- deepinfer/neuron_ops/regions.py | 22 ++++++++++++++- deepinfer/neuron_ops/vendor/deep_gemm/gemm.py | 3 +- deepinfer/worker/model_runner.py | 28 +++++++++++++++++-- deepinfer/worker/neuron_graph_runner.py | 24 +++++++++++++++- examples/neuron/offline_inference_neuron.py | 6 ++-- 11 files changed, 116 insertions(+), 13 deletions(-) diff --git a/csrc/engine/input_metadata/prepare_tensor_input.hpp b/csrc/engine/input_metadata/prepare_tensor_input.hpp index a2eaf99..9656bf6 100644 --- a/csrc/engine/input_metadata/prepare_tensor_input.hpp +++ b/csrc/engine/input_metadata/prepare_tensor_input.hpp @@ -25,6 +25,9 @@ struct PrepareArgs { int32_t image_sentinel = -1; bool is_prompt = true; bool force_load_balance = false; + // decode-shaped metadata for a prompt chunk (Neuron prefill through the decode graphs): input_ids are the prompt's + // own tokens rather than the last token + MTP drafts; everything else follows the decode layout. + bool prompt_as_decode = false; // prompt-only int cp_size = 1; diff --git a/csrc/engine/input_metadata/prepare_tensor_input_base.cpp b/csrc/engine/input_metadata/prepare_tensor_input_base.cpp index 04be7f5..2c050e1 100644 --- a/csrc/engine/input_metadata/prepare_tensor_input_base.cpp +++ b/csrc/engine/input_metadata/prepare_tensor_input_base.cpp @@ -97,12 +97,13 @@ static py::dict prepare_common_base(const PrepareArgs &args) { auto input_ids = get_or_create_tensor(result, has_out, "input_ids", {total_q}, torch::kInt32); auto p_input_ids = input_ids.mutable_data_ptr(); + const bool prompt_ids = is_prompt || args.prompt_as_decode; i_q = 0; for (auto p_seq : seq_metadata_list) { auto &token_ids = p_seq->seq_data->token_ids; static_assert(std::is_same_v); int seqlen_q_golden; - if (is_prompt) { + if (prompt_ids && p_seq->is_prompt) { TORCH_CHECK(int(token_ids.size()) >= p_seq->seqlen_k); seqlen_q_golden = p_seq->seqlen_q; } else { @@ -110,7 +111,7 @@ static py::dict prepare_common_base(const PrepareArgs &args) { seqlen_q_golden = 1; } memcpy(p_input_ids + i_q, token_ids.data() + p_seq->seqlen_k - p_seq->seqlen_q, seqlen_q_golden * sizeof(int32_t)); - if (!is_prompt) { + if (seqlen_q_golden < p_seq->seqlen_q) { int seqlen_q_draft = p_seq->seqlen_q - 1; TORCH_CHECK(int(p_seq->mtp.draft_tokens.size()) >= seqlen_q_draft); for (int j = 0; j < seqlen_q_draft; ++j) { diff --git a/csrc/engine/pybind_output.cpp b/csrc/engine/pybind_output.cpp index 995fd42..27e05a7 100644 --- a/csrc/engine/pybind_output.cpp +++ b/csrc/engine/pybind_output.cpp @@ -67,6 +67,7 @@ void pybind_engine_output(pybind11::module &engine) { .def_readwrite("yoco_prefill_keep_tokens", &PrepareArgs::yoco_prefill_keep_tokens) .def_readwrite("max_backtracking_tokens", &PrepareArgs::max_backtracking_tokens) .def_readwrite("force_load_balance", &PrepareArgs::force_load_balance) + .def_readwrite("prompt_as_decode", &PrepareArgs::prompt_as_decode) .def_readwrite("dense_compress_ratios", &PrepareArgs::dense_compress_ratios) .def_readwrite("sparse_compress_ratios", &PrepareArgs::sparse_compress_ratios); diff --git a/deepinfer/environ.py b/deepinfer/environ.py index b09f076..9e98be6 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -114,8 +114,14 @@ class envs: DEEPINFER_NEURON_GRAPH_MODE = EnvVar.str(default="regions") # Neuron only: cut the region graphs further at the stage boundaries marked by neuron_ops.regions.split(). DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=True) + # Neuron only: synchronize around every region graph call and log per-kind times each decode step (profiling). + DEEPINFER_NEURON_REGION_TIMING = EnvVar.bool() # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron. DEEPINFER_NEURON_HOST_EMBED = EnvVar.bool(default=True) + # Neuron only, regions mode: run prompt chunks through the decode region graphs padded to this many tokens + # (issue #26, the vllm-neuron#71 layout). 0 = eager prefill. Needs max_num_batched_tokens <= bucket and + # DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket (the SWA / compressor rings must hold a whole chunk). + DEEPINFER_NEURON_PREFILL_BUCKET = EnvVar.int(default=0) # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break). DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool() diff --git a/deepinfer/model_executor/model_metadata_utils.py b/deepinfer/model_executor/model_metadata_utils.py index 25c8367..fd839a5 100644 --- a/deepinfer/model_executor/model_metadata_utils.py +++ b/deepinfer/model_executor/model_metadata_utils.py @@ -152,8 +152,10 @@ class ModelMetadata: window_size: int, swa_block_size: int, metadata_zos: MetadataZOS | None = None, + prompt_as_decode: bool = False, ) -> dict: - is_prompt = seq_metadata_list[0].is_prompt + # prompt_as_decode: decode-layout metadata for a prompt chunk (Neuron prefill through the decode graphs). + is_prompt = seq_metadata_list[0].is_prompt and not prompt_as_decode prepare_func = self.prepare_prompt if is_prompt else self.prepare_decode @@ -168,6 +170,7 @@ class ModelMetadata: args.cp_size = PG.context_parallel_group().size() else: args.is_prompt = False + args.prompt_as_decode = prompt_as_decode args.max_num_blocks = max_num_blocks args.out = out_dict args.window_size = window_size diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 75f196f..43d1da1 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -29,6 +29,8 @@ import os import torch ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_ATTN", "1") == "1" +#: DEEPINFER_NEURON_PREFILL_BUCKET (issue #26): a step of exactly this many query tokens is one prompt chunk. +PREFILL_BUCKET = int(os.environ.get("DEEPINFER_NEURON_PREFILL_BUCKET", "0")) D_QK = 512 INDEX_DIM = 128 @@ -320,6 +322,23 @@ def paged_mqa_logits( slot = torch.clamp(pos // page_kv, max=width - 1) lens = context_lens.reshape(num_q, 1).to(torch.int32) valid = (pos.view(1, n) < lens) & (pos.view(1, n) < width * page_kv) + w = weights.to(torch.float32) + kernel = _kernel("index_score", "weighted_scores") + + if num_q == PREFILL_BUCKET: + # A prompt chunk run as a decode step (issue #26): its tokens are one sequence's, so they share token 0's block + # table and differ only in their context length. Gather the keys once and score every token in one kernel call + # (padding rows past the chunk score garbage on the same keys; their results are never read). + page = block_table[0].to(torch.int32).index_select(0, slot) + page = torch.where(pos < lens[0], page, torch.zeros_like(page)).view(1, n) + offset = (pos % page_kv).view(1, n) + data = gather_rows(kv_cache, page, offset, data_bytes, data_bytes).to(torch.int32) + sf = gather_rows(kv_cache, page, offset, SF_BYTES, SF_BYTES, base=page_kv * data_bytes).to(torch.int32) + values = _unpack_e2m1(data) if is_fp4 else _e4m3(data) # [1, n, 128] fp32 + keys = (values.view(n, SF_BYTES, 32) * _pow2(sf.view(n, SF_BYTES) - 127).unsqueeze(-1)).view(n, INDEX_DIM) + scores = kernel(q.reshape(num_q * heads, -1).contiguous(), keys.to(torch.bfloat16), w.transpose(0, 1).contiguous()) + return _finish_logits(scores, valid, max_context_len, num_q, device) + rows = torch.arange(num_q, device=device, dtype=torch.int32).view(num_q, 1) * width + slot.view(1, n) page = block_table.to(torch.int32).reshape(-1).index_select(0, rows.reshape(-1)).view(num_q, n) page = torch.where(valid, page, torch.zeros_like(page)) # stale table slots past the context: read page 0 @@ -331,12 +350,13 @@ def paged_mqa_logits( keys = (values.view(num_q, n, SF_BYTES, 32) * _pow2(sf - 127).unsqueeze(-1)).view(num_q, n, INDEX_DIM) keys = keys.to(torch.bfloat16) - w = weights.to(torch.float32) - kernel = _kernel("index_score", "weighted_scores") # Each query token has its own context: one kernel call per token (the kernel shares k across its tokens). scores = torch.cat([kernel(q[t].contiguous(), keys[t], w[t].reshape(heads, 1)) for t in range(num_q)], dim=0) - scores = torch.where(valid, scores, torch.full_like(scores, float("-inf")))[:, :max_context_len] + return _finish_logits(scores, valid, max_context_len, num_q, device) + +def _finish_logits(scores: torch.Tensor, valid: torch.Tensor, max_context_len: int, num_q: int, device) -> torch.Tensor: # noqa: ANN001 + scores = torch.where(valid, scores, torch.full_like(scores, float("-inf")))[:, :max_context_len] stride = _ceil(max_context_len, 256) # deep_gemm's 512 B row alignment logits = torch.empty((num_q, stride), dtype=torch.bfloat16, device=device)[:, :max_context_len] logits.copy_(scores.to(torch.bfloat16)) diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index d5795b0..aa71490 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -8,6 +8,8 @@ fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF from __future__ import annotations import functools +import time +from collections import defaultdict from collections.abc import Callable from typing import Any @@ -35,6 +37,15 @@ def raise_dynamo_limits() -> None: _ENTRIES: dict[tuple[type, str], Callable[..., Any]] = {} +#: DEEPINFER_NEURON_REGION_TIMING: seconds per region kind in the current step (each call synchronized). +_TIMES: dict[str, list[float]] = defaultdict(list) + + +def take_timings() -> dict[str, list[float]]: + """The region timings recorded since the last call, then reset.""" + out = dict(_TIMES) + _TIMES.clear() + return out def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]: @@ -55,9 +66,18 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]: getattr(type(module), name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False ) + label = f"{type(module).__name__}.{name}" + def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any: with hide_cuda(): - return compiled(mod, *args, **kwargs) + if not envs.DEEPINFER_NEURON_REGION_TIMING: + return compiled(mod, *args, **kwargs) + torch.neuron.synchronize() + started = time.perf_counter() + out = compiled(mod, *args, **kwargs) + torch.neuron.synchronize() + _TIMES[label].append(time.perf_counter() - started) + return out _ENTRIES[key] = fn return functools.partial(fn, module) diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py index a009c40..da5e23e 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py @@ -106,9 +106,10 @@ def _bf16_gemm( if k == 0: _write_output(torch.zeros((m, n), dtype=torch.float32, device=d.device), d, c, alpha) return - if torch.compiler.is_compiling(): + if torch.compiler.is_compiling() or b.device.type == "neuron": # NEURON graphs: one bf16 dot with fp32 accumulation and no materialized fp32 transpose (vllm-neuron#71 runs # these as bf16 linears; an fp32 transpose feeding a dot is a neuronx-cc InferNonlocalTensors trigger). + # Eagerly on a NeuronCore too: ``b.float()`` of the LM head is a 2.6 GiB copy a call. _write_output(torch.nn.functional.linear(a, b).float(), d, c, alpha) return _write_output(a.float() @ b.float().transpose(0, 1), d, c, alpha) diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index c0cf54d..ea44665 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import copy import time from dataclasses import dataclass from typing import TYPE_CHECKING @@ -79,6 +80,13 @@ class ModelRunner: max_batch_size=self.scheduler_config.max_num_seqs * (1 + self.nextn), input_on_device=envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING, ) + bucket = getattr(self.cuda_graph_runner, "prefill_bucket", 0) + if bucket: + assert self.scheduler_config.max_num_batched_tokens <= bucket, ( + f"max_num_batched_tokens ({self.scheduler_config.max_num_batched_tokens}) > DEEPINFER_NEURON_PREFILL_BUCKET" + ) + # The chunk's tokens are all stored before any of them attends: the rings must hold a whole chunk. + assert envs.DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket, "set DEEPINFER_MAX_BACKTRACKING_TOKENS >= the prefill bucket" self.dflash_use_cuda_graph = ( self.use_cuda_graph and not envs.DEEPINFER_DFLASH_ENFORCE_EAGER @@ -124,6 +132,7 @@ class ModelRunner: seq_metadata_list: list[SequenceMetadata], out_dict: dict | None = None, metadata_zos: MetadataZOS | None = None, + prompt_as_decode: bool = False, ) -> InputMetadata | None: """Build input metadata on CPU, or on device for ZOS decode. @@ -138,6 +147,7 @@ class ModelRunner: window_size=self.window_size, swa_block_size=self.swa_block_size, metadata_zos=metadata_zos, + prompt_as_decode=prompt_as_decode, ) if out_dict is not None: return None @@ -181,16 +191,30 @@ class ModelRunner: if total_q is None: total_q = sum(seq.seqlen_q for seq in seq_metadata_list) is_prompt = seq_metadata_list[0].is_prompt - run_cuda_graph = not is_prompt and self.use_cuda_graph and bool(self.cuda_graph_runner.graphs) + # Neuron (issue #26): a prompt chunk runs through the decode graphs as a decode-shaped step padded to the + # prefill bucket. Every EP rank takes the same bucket (idle ranks run a placeholder prompt), so the static + # collectives inside the graphs match. + prefill_bucket = getattr(self.cuda_graph_runner, "prefill_bucket", 0) if self.use_cuda_graph else 0 + prompt_as_decode = is_prompt and prefill_bucket > 0 + if prompt_as_decode: + assert total_q <= prefill_bucket, f"prompt chunk of {total_q} tokens exceeds DEEPINFER_NEURON_PREFILL_BUCKET" + # The indexer scores the chunk against one block table (nki_attention.paged_mqa_logits). + assert num_reqs == 1, "a Neuron prefill step through the decode graphs takes one sequence" + run_cuda_graph = (not is_prompt or prompt_as_decode) and self.use_cuda_graph and bool(self.cuda_graph_runner.graphs) if run_cuda_graph: - graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q) + graph_size, cuda_graph_used = self.cuda_graph_runner.select(prefill_bucket if prompt_as_decode else total_q) self._prepare_input_tensors( seq_metadata_list + self.cuda_graph_runner.padding_seq_metadata_list[: graph_size - total_q], out_dict=cuda_graph_used.input_metadata_source.to_dict(), metadata_zos=metadata_zos, + prompt_as_decode=prompt_as_decode, ) input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs) + if prompt_as_decode: + # The graph runs on its decode-mode metadata; sampling and post-processing see a prompt step. + input_metadata = copy.copy(input_metadata) + input_metadata.is_prompt = True else: graph_size, cuda_graph_used = 0, None input_metadata = self._prepare_input_tensors( diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 8b68edc..effb38c 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -10,6 +10,7 @@ from __future__ import annotations import dataclasses import math +import time from collections.abc import Callable from typing import TYPE_CHECKING @@ -21,6 +22,7 @@ from deepinfer.logger import init_logger from deepinfer.model_executor.input_metadata import InputMetadata from deepinfer.model_executor.model_output import CausalLMForwardOutput from deepinfer.neuron_ops import neuron_compile # noqa: F401 (registers the "deepinfer_neuron" backend) +from deepinfer.neuron_ops import regions from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda from deepinfer.sequence import SequenceMetadata from deepinfer.worker.cache_engine import GPUCacheInterface, KVCache @@ -120,6 +122,11 @@ class _LazyKVCache: class NeuronGraphRunner(CudaGraphRunner): def __init__(self, max_batch_size: int, input_on_device: bool = False) -> None: super().__init__(max_batch_size, input_on_device) + # Prompt chunks run as decode-shaped steps of this many tokens (issue #26): one more graph size. + self.prefill_bucket = envs.DEEPINFER_NEURON_PREFILL_BUCKET + if self.prefill_bucket > self.max_batch_size: + self.max_batch_size = self.prefill_bucket + self.batch_sizes = [*self.batch_sizes, self.prefill_bucket] self.compiled: dict[int, Callable[..., CausalLMForwardOutput]] = {} self.model: DeepseekForCausalLM | None = None self.kv_caches: GPUCacheInterface | None = None @@ -179,8 +186,23 @@ class NeuronGraphRunner(CudaGraphRunner): if not self.input_on_device: input_metadata_device.recursive_copy_(cuda_graph_data.input_metadata_source) batch_size = next(bs for bs, data in self.graphs.items() if data is cuda_graph_data) + if not envs.DEEPINFER_NEURON_REGION_TIMING: + with hide_cuda(): + return self.compiled[batch_size](input_metadata_device) + torch.neuron.synchronize() + started = time.perf_counter() with hide_cuda(): - return self.compiled[batch_size](input_metadata_device) + out = self.compiled[batch_size](input_metadata_device) + torch.neuron.synchronize() + total = time.perf_counter() - started + timings = regions.take_timings() + in_regions = sum(sum(v) for v in timings.values()) + if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: + return out + parts = ", ".join(f"{k} {len(v)}x {1e3 * sum(v) / len(v):.2f} ms" for k, v in sorted(timings.items())) + logger.info(f"Neuron step bs={batch_size}: {1e3 * total:.1f} ms, regions {1e3 * in_regions:.1f} ms ({parts}), " + f"outside regions {1e3 * (total - in_regions):.1f} ms") + return out def destroy(self) -> None: super().destroy() diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py index ef375fb..c462aea 100644 --- a/examples/neuron/offline_inference_neuron.py +++ b/examples/neuron/offline_inference_neuron.py @@ -37,6 +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_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") 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) @@ -86,8 +87,9 @@ def main(dp_size: int, dp_rank: int) -> None: from deepinfer.sampling_params import SamplingParams - params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in PROMPTS] - outputs = llm.generate(prompts=PROMPTS, sampling_params=params, use_tqdm=False) + prompts = PROMPTS[: args.num_prompts] + params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in prompts] + outputs = llm.generate(prompts=prompts, sampling_params=params, use_tqdm=False) for i, output in enumerate(outputs): text = llm.llm_engine.tokenizer.decode([x.token_id for x in output]) logps = [x.token_logprobs for x in output]