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 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)
        #: Requests a step can hold (max_num_seqs): the decode rows are max_num_seqs * (1 + nextn).
        self.max_reqs = max(1, max_batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE))
        # 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:
            # A prompt chunk packs at most max_num_seqs requests: bound the table by those and the step's context
            # bucket (without a bound every prefill layer decoded the whole KV pool, issue #29).
            ctx = getattr(input_metadata_device, "neuron_ctx_limit", None)
            nki_attention.set_table_bounds(self.max_reqs if envs.DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS and ctx else None, ctx)
            nki_attention.set_single_request_chunk(self.max_reqs == 1)
        else:
            nki_attention.set_single_request_chunk(False)
            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:
            if regions._HOST_TIMING:
                started = time.perf_counter()
                with hide_cuda():
                    out = self.compiled[batch_size](input_metadata_device)
                self._host_steps = getattr(self, "_host_steps", 0) + 1
                if self._host_steps % 20 == 0 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
                    host = regions.take_host_timings()
                    inside = sum(v[1] for v in host.values())
                    logger.info(
                        f"Neuron host bs={batch_size}: forward {1e3 * (time.perf_counter() - started):.1f} ms, in regions "
                        f"{1e3 * inside:.1f} ms; " + ", ".join(f"{k} {v[0]}x {1e3 * v[1]:.1f} ms" for k, v in sorted(host.items(), key=lambda kv: -kv[1][1]))
                    )
                else:
                    regions.take_host_timings()
                return out
            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