mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
5e36f90015
The table path decoded the whole KV pool every layer (160 blocks x 128 slots); the fp8 record decode is integer work on GpSimd and took 6-10 ms a layer (device profile of the 6-row attention graph, bench47). A decode step's requests reference at most ceil(ctx / page) + 1 blocks of a scope each: those blocks are chosen with cumsum + scatter (static shapes), gathered whole (one contiguous run each; a per-slot gather was slower than the full decode) and the indices remapped into them. Sizes are rounded to a power of two (neuronx-cc ISA validation 'same_element_count_t2d' otherwise). Prompt-chunk buckets keep the whole pool. The graph runner sets the bounds (requests, context bucket) per replay. bench50 (512 in, 128 out, TP=64, nextn=5): TPOT 234.1 -> 208.4 ms, step 907 -> 805 ms, per-request tokens-per-step identical to bench41 (3.85, 4.38, 1.98). Tiny cp=2 nextn=5 with a 128-block pool: 9/9 golden tokens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
242 lines
12 KiB
Python
242 lines
12 KiB
Python
"""Static decode graphs on Neuron: the CudaGraphRunner contract, backed by torch.compile(backend="neuron").
|
|
|
|
One compiled graph per captured batch size, over persistent device metadata buffers (the same buffers the CUDA
|
|
runner captures over). Dynamo specializes on their shapes (dynamic=False), so each batch size is one NEFF set.
|
|
Replay copies the host metadata into the device buffers and calls the compiled forward; the host->device copy stays
|
|
outside the graph.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import dataclasses
|
|
import math
|
|
import time
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
from deepinfer import envs
|
|
from deepinfer.environ import _descriptors as _env_descriptors
|
|
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
|
|
from deepinfer.worker.cuda_graph_runner import CudaGraphData, CudaGraphRunner
|
|
|
|
if TYPE_CHECKING:
|
|
from deepinfer.model_executor.models.deepseek import DeepseekForCausalLM
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class _GraphKVCaches:
|
|
"""The per-layer cache views, rebuilt from their root allocations inside the graph.
|
|
|
|
The layer caches are views of one KV buffer and one state buffer. Passed as graph inputs, they alias each other,
|
|
and AOT regenerates aliased mutated inputs with as_strided, which cannot be functionalized. Here only the two
|
|
roots are graph inputs; each layer's views are re-sliced from them (slice + view, no as_strided) on access.
|
|
"""
|
|
|
|
def __init__(self, kv_caches: GPUCacheInterface) -> None:
|
|
self.cache_list = kv_caches.cache_list
|
|
self.roots: list[torch.Tensor] = []
|
|
if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions":
|
|
# The cache engine lays every layer view out contiguously on Neuron (_LayerMajorColumns); the region
|
|
# graphs take them as inputs as they are.
|
|
return
|
|
self.recipes: list[dict[str, tuple | None]] = []
|
|
for cache in kv_caches.cache_list:
|
|
recipe = {}
|
|
for f in dataclasses.fields(KVCache):
|
|
t = getattr(cache, f.name)
|
|
recipe[f.name] = None if t is None else self._recipe(t)
|
|
self.recipes.append(recipe)
|
|
|
|
def _recipe(self, t: torch.Tensor) -> tuple:
|
|
storage = t.untyped_storage()
|
|
for i, r in enumerate(self.roots):
|
|
if r.untyped_storage().data_ptr() == storage.data_ptr():
|
|
root = r
|
|
break
|
|
else:
|
|
# One flat uint8 tensor over the whole allocation (dtype views of it have no usable _base).
|
|
root = torch.empty(0, dtype=torch.uint8, device=t.device).set_(storage)
|
|
self.roots.append(root)
|
|
i = len(self.roots) - 1
|
|
rows, row_bytes = t.shape[0], t.stride(0) * t.element_size()
|
|
offset = t.storage_offset() * t.element_size()
|
|
assert offset + math.prod(t.shape[1:]) * t.element_size() <= row_bytes, "cache view must be a column slice of its allocation"
|
|
nbytes = math.prod(t.shape[1:]) * t.element_size()
|
|
recipe = (i, rows, row_bytes, offset, nbytes, tuple(t.shape[1:]), t.dtype)
|
|
rebuilt = self._build(root, recipe)
|
|
assert rebuilt.data_ptr() == t.data_ptr() and rebuilt.shape == t.shape and rebuilt.stride() == t.stride(), (
|
|
f"cache view not expressible as a row slice: {t.shape=} {t.stride()=}"
|
|
)
|
|
return recipe
|
|
|
|
@staticmethod
|
|
def _build(root: torch.Tensor, recipe: tuple) -> torch.Tensor:
|
|
_, rows, row_bytes, offset, nbytes, tail, dtype = recipe
|
|
esize = torch.empty((), dtype=dtype).element_size()
|
|
flat = root[: rows * row_bytes]
|
|
view = flat.view(rows, row_bytes)[:, offset : offset + nbytes].view(rows, *tail[:-1], tail[-1] * esize)
|
|
return view if dtype == torch.uint8 else view.view(dtype)
|
|
|
|
def __getitem__(self, layer_idx: int) -> KVCache | _LazyKVCache:
|
|
if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions":
|
|
# The layer's own contiguous, non-overlapping views as region graph inputs (like vllm-neuron): slicing
|
|
# them from the roots inside the graph baked each layer's byte offsets in, one NEFF per layer.
|
|
return self.cache_list[layer_idx]
|
|
tensors = {}
|
|
for name, recipe in self.recipes[layer_idx].items():
|
|
tensors[name] = None if recipe is None else self._build(self.roots[recipe[0]], recipe)
|
|
return KVCache(**tensors)
|
|
|
|
def is_valid_layeridx(self, layer_idx: int) -> bool:
|
|
return layer_idx < len(self.cache_list)
|
|
|
|
|
|
class _LazyKVCache:
|
|
"""A layer's KVCache whose views are sliced from the roots when read (KVCache's fields, as properties)."""
|
|
|
|
def __init__(self, caches: _GraphKVCaches, layer_idx: int) -> None:
|
|
self._caches = caches
|
|
self._recipes = caches.recipes[layer_idx]
|
|
|
|
def _view(self, name: str) -> torch.Tensor | None:
|
|
recipe = self._recipes[name]
|
|
return None if recipe is None else _GraphKVCaches._build(self._caches.roots[recipe[0]], recipe)
|
|
|
|
ica_cache = property(lambda self: self._view("ica_cache"))
|
|
ica_state = property(lambda self: self._view("ica_state"))
|
|
ca_cache = property(lambda self: self._view("ca_cache"))
|
|
ca_state = property(lambda self: self._view("ca_state"))
|
|
swa_cache = property(lambda self: self._view("swa_cache"))
|
|
|
|
|
|
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
|
|
|
|
def capture(
|
|
self,
|
|
model: DeepseekForCausalLM,
|
|
kv_caches: GPUCacheInterface,
|
|
prepare_input_fn: Callable[[list[SequenceMetadata]], InputMetadata],
|
|
padding_seq_metadata: SequenceMetadata,
|
|
) -> None:
|
|
self.destroy()
|
|
self.model, self.kv_caches = model, kv_caches
|
|
# Resolve every env var now: the lazy descriptor (lock + os.getenv) does not trace.
|
|
for name in _env_descriptors:
|
|
getattr(envs, name)
|
|
backend, fullgraph = envs.DEEPINFER_NEURON_COMPILE_BACKEND, envs.DEEPINFER_NEURON_FULLGRAPH
|
|
logger.info(f"Neuron decode graphs: backend={backend} fullgraph={fullgraph} batch_sizes={self.batch_sizes}")
|
|
|
|
self.padding_seq_metadata_list = [padding_seq_metadata] * self.max_batch_size
|
|
input_metadata_source_full = prepare_input_fn(self.padding_seq_metadata_list)
|
|
input_metadata_device_full = input_metadata_source_full.recursive_fn(lambda x: x.cuda(non_blocking=True))
|
|
|
|
graph_caches = _GraphKVCaches(kv_caches)
|
|
|
|
def forward(input_metadata: InputMetadata) -> CausalLMForwardOutput:
|
|
return model(kv_caches=graph_caches, input_metadata=input_metadata)
|
|
|
|
self.graphs = {}
|
|
for batch_size in reversed(self.batch_sizes):
|
|
|
|
def _cut(x: torch.Tensor, _bs: int = batch_size) -> torch.Tensor:
|
|
assert x.shape[0] == self.max_batch_size
|
|
return x[:_bs]
|
|
|
|
input_metadata_source = input_metadata_source_full.recursive_fn(_cut)
|
|
input_metadata_device = input_metadata_device_full.recursive_fn(_cut)
|
|
if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions":
|
|
compiled = forward # eager step; the layers compile their region graphs (neuron_ops/regions.py)
|
|
else:
|
|
compiled = torch.compile(forward, backend=backend, fullgraph=fullgraph, dynamic=False)
|
|
# First call compiles (and runs) this batch size's graph.
|
|
regions.set_step_rows(batch_size)
|
|
input_metadata_device.recursive_copy_(input_metadata_source)
|
|
with hide_cuda():
|
|
compiled(input_metadata_device)
|
|
for name in input_metadata_source.get_clean_list():
|
|
setattr(input_metadata_source, name, None)
|
|
setattr(input_metadata_device, name, None)
|
|
self.compiled[batch_size] = compiled
|
|
self.graphs[batch_size] = CudaGraphData(None, None, None, input_metadata_source, input_metadata_device)
|
|
|
|
def replay(self, cuda_graph_data: CudaGraphData) -> CausalLMForwardOutput:
|
|
input_metadata_device = cuda_graph_data.input_metadata_device
|
|
# Forward recomputes these (and stores them on the metadata); reset them to the state the graph was traced in.
|
|
for name in input_metadata_device.get_clean_list():
|
|
setattr(input_metadata_device, name, None)
|
|
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)
|
|
regions.set_step_rows(batch_size)
|
|
# Table attention decodes only the blocks the step's requests can reference. A decode request has exactly
|
|
# 1 + nextn rows; a prompt-chunk bucket may pack several prompts, so it keeps the whole pool.
|
|
from deepinfer.neuron_ops import nki_attention
|
|
|
|
if batch_size == nki_attention.PREFILL_BUCKET:
|
|
nki_attention.set_table_bounds(None, None)
|
|
else:
|
|
reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE))
|
|
nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None))
|
|
if not envs.DEEPINFER_NEURON_REGION_TIMING:
|
|
with hide_cuda():
|
|
return self.compiled[batch_size](input_metadata_device)
|
|
torch.neuron.synchronize()
|
|
self._steps = getattr(self, "_steps", 0) + 1
|
|
profiler = None
|
|
self._bucket_steps = getattr(self, "_bucket_steps", 0) + (batch_size > 1)
|
|
profile = self._steps == 8 or (batch_size > 1 and self._bucket_steps == 2) # a decode step, the 2nd prefill chunk
|
|
if profile and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
|
|
import cProfile
|
|
|
|
profiler = cProfile.Profile()
|
|
profiler.enable()
|
|
started = time.perf_counter()
|
|
from deepinfer.neuron_ops.op_census import census
|
|
|
|
with hide_cuda(), census(f"forward bs={batch_size}") if profiler is not None else contextlib.nullcontext():
|
|
out = self.compiled[batch_size](input_metadata_device)
|
|
torch.neuron.synchronize()
|
|
total = time.perf_counter() - started
|
|
if profiler is not None:
|
|
import io
|
|
import pstats
|
|
|
|
profiler.disable()
|
|
buf = io.StringIO()
|
|
pstats.Stats(profiler, stream=buf).sort_stats("cumulative").print_stats(60)
|
|
logger.info(f"Neuron step profile (rank 0, step {self._steps}, bs={batch_size}):\n{buf.getvalue()}")
|
|
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()
|
|
self.compiled = {}
|
|
self.model = self.kv_caches = None
|