DeepInfer/deepinfer/neuron_ops/regions.py

"""Region graphs for the Neuron decode step (issue #25), the vllm-neuron#71 layout.

A whole decode step in one graph crashes neuronx-cc (issue #9) and takes long to compile. As in vllm-neuron#71, the
step instead runs eagerly and compiles a few entry points per layer (the attention half and the MoE), each its own
fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF cache compiles each kind once.
"""

from __future__ import annotations

import functools
import os
import time
from collections import defaultdict
from collections.abc import Callable
from typing import Any

import torch

from deepinfer import PLATFORM, envs


def enabled(input_metadata: Any) -> bool:
    """Region graphs: Neuron, decode, DEEPINFER_ENFORCE_EAGER=0 and DEEPINFER_NEURON_GRAPH_MODE=regions."""
    return (
        PLATFORM.is_neuron
        and not envs.DEEPINFER_ENFORCE_EAGER
        and envs.DEEPINFER_NEURON_GRAPH_MODE == "regions"
        and not input_metadata.is_prompt
    )


def raise_dynamo_limits() -> None:
    """Every layer recompiles each entry (its parameters and attributes differ); keep dynamo compiling them. Past
    the accumulated limit dynamo silently runs frames eagerly instead (see vllm-neuron#71 model.py)."""
    for name in ("recompile_limit", "cache_size_limit", "accumulated_recompile_limit", "accumulated_cache_size_limit"):
        if hasattr(torch._dynamo.config, name):
            setattr(torch._dynamo.config, name, 1 << 30)


_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)
#: Experiment (issue #27): synchronize after every N region calls (0 = never).
_SYNC_EVERY = int(os.environ.get("DEEPINFER_NEURON_SYNC_EVERY", "0"))
_CALLS = [0]


def count(label: str) -> None:
    """Count an event in the per-step timings (DEEPINFER_NEURON_REGION_TIMING)."""
    _TIMES[label].append(0.0)


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]:
    """``type(module).<name>`` as a region, compiled once per (class, method) and called with the module.

    One compiled callable for all layers of a class: Dynamo then reuses the same graph (and the same loaded NEFF)
    for every layer whose guards match (same layer kind and shapes). A per-module compile gave every layer its own
    compiled artifact and NEFF load -- the same code loaded 43 times filled ~6 GB of HBM with model code.
    """
    key = (type(module), name)
    fn = _ENTRIES.get(key)
    if fn is None:
        from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda

        raise_dynamo_limits()
        # fullgraph=False: split() marks sub-region boundaries inside the entry (smaller graphs for neuronx-cc).
        compiled = torch.compile(
            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():
                if not envs.DEEPINFER_NEURON_REGION_TIMING:
                    out = compiled(mod, *args, **kwargs)
                    if _SYNC_EVERY:
                        _CALLS[0] += 1
                        if _CALLS[0] % _SYNC_EVERY == 0:
                            torch.neuron.synchronize()
                    return out
                torch.neuron.synchronize()
                started = time.perf_counter()
                _MARK[0], _MARK[1], _MARK[2] = started, 0, label
                out = compiled(mod, *args, **kwargs)
                torch.neuron.synchronize()
                done = time.perf_counter()
                if _MARK[1]:
                    _TIMES[f"{label}.split{_MARK[1]}"].append(done - _MARK[0])
                _MARK[0] = 0.0
                _TIMES[label].append(done - started)
                return out

        _ENTRIES[key] = fn
    return functools.partial(fn, module)


_MARK: list = [0.0, 0, ""]


@torch._dynamo.disable
def _boundary() -> None:
    """A graph boundary and nothing else (vllm-neuron#71's ``_graph_break``). Under DEEPINFER_NEURON_REGION_TIMING
    also records the time since the previous boundary (or the region's start) as ``split<i>``."""
    if envs.DEEPINFER_NEURON_REGION_TIMING:
        torch.neuron.synchronize()
        now = time.perf_counter()
        if _MARK[0]:
            _TIMES[f"{_MARK[2]}.split{_MARK[1]}"].append(now - _MARK[0])
        _MARK[0], _MARK[1] = now, _MARK[1] + 1


#: Rows of the step being run (set by NeuronGraphRunner.replay): a prefill chunk has more than one.
_STEP_ROWS = [1]
#: Multi-row steps up to this many rows (DSpark verification: 1 + drafts) keep only the coarse splits
#: (DEEPINFER_NEURON_SPLIT_MIN_ROWS). Whole 6-row layers crash neuronx-cc on the full model, and the 12 fine
#: splits a layer cost ~1 ms each in launches.
_SPLIT_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_SPLIT_MIN_ROWS", "16"))


def set_step_rows(rows: int) -> None:
    _STEP_ROWS[0] = rows


def split(coarse: bool = False) -> None:
    """Cut the current region graph here (only while compiling; free in eager). Always for prefill chunks under
    DEEPINFER_NEURON_REGION_SPLIT_PREFILL: neuronx-cc runs the split 128-row attention ~2x faster (62 vs 122 ms a
    layer, issue #26), while decode's 1-row graphs are faster whole. Multi-row steps up to _SPLIT_MIN_ROWS rows cut
    only at the ``coarse`` boundaries."""
    if not torch.compiler.is_compiling():
        return
    rows = _STEP_ROWS[0]
    multi_row = envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and rows > 1
    if envs.DEEPINFER_NEURON_REGION_SPLIT or (multi_row and (coarse or rows > _SPLIT_MIN_ROWS)):
        _boundary()