mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
4bc8ed3eaf
- The state buffer's column regions get their own allocations on Neuron: a layer's uint8 SWA cache and fp32 compressor state are both region graph inputs, and aot_autograd compares offsets of differently typed views of one storage in their own element units (false overlap -> 'input mutations on views with different dtypes'). - paged_mqa_logits chunk path: gather the keys the chunk's last token sees (lens.max()), not the first token's. - Debug: DEEPINFER_NEURON_ALIAS_DEBUG prints aliased graph inputs; a cProfile of decode step 8 under DEEPINFER_NEURON_REGION_TIMING. Tiny V4.1 on a NeuronCore, bucket 128: short 8/9 (= eager), long 715 tokens in 6 chunks 9/9 (eager 8/9) in 26 s instead of 978 s of eager prefill. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
226 lines
11 KiB
Python
226 lines
11 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 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.
|
|
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)
|
|
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
|
|
if self._steps == 8 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
|
|
import cProfile
|
|
|
|
profiler = cProfile.Profile()
|
|
profiler.enable()
|
|
started = time.perf_counter()
|
|
with hide_cuda():
|
|
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("tottime").print_stats(40)
|
|
logger.info(f"Neuron step profile (rank 0, step {self._steps}):\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
|