DeepInfer/deepinfer/worker/neuron_graph_runner.py

"""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
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.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.roots: list[torch.Tensor] = []
        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:
        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.recipes)


class NeuronGraphRunner(CudaGraphRunner):
    def __init__(self, max_batch_size: int, input_on_device: bool = False) -> None:
        super().__init__(max_batch_size, input_on_device)
        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)
            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)
        with hide_cuda():
            return self.compiled[batch_size](input_metadata_device)

    def destroy(self) -> None:
        super().destroy()
        self.compiled = {}
        self.model = self.kv_caches = None