mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 22:59:44 +08:00
c24574bfd6
- Layer-major KV / state layout so per-layer views are contiguous graph inputs (the executor copied non-contiguous inputs every call). - One compiled entry per (class, method), rank as a registered buffer, no layer_idx guard: every layer reuses the same graphs and loaded NEFFs. - fp8 weight rescale on the host (numpy LUT, ~2 min instead of 35+ min of eager per-shape compiles); engram sf packing on the host. - int16 ue8m0 packing in fp32 (int32 interleaved slices narrowed to int16 crash neuronx-cc, NCC_IBCG902). - Token embedding and the text model's vision tower stay in host memory on Neuron (2.5 GiB of HBM a rank). - offline example: --max_num_seqs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
"""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
|
|
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]] = {}
|
|
|
|
|
|
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
|
|
)
|
|
|
|
def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
|
|
with hide_cuda():
|
|
return compiled(mod, *args, **kwargs)
|
|
|
|
_ENTRIES[key] = fn
|
|
return functools.partial(fn, module)
|
|
|
|
|
|
@torch._dynamo.disable
|
|
def _boundary() -> None:
|
|
"""A graph boundary and nothing else (vllm-neuron#71's ``_graph_break``)."""
|
|
|
|
|
|
def split() -> None:
|
|
"""Cut the current region graph here (only while compiling; free in eager)."""
|
|
if torch.compiler.is_compiling() and envs.DEEPINFER_NEURON_REGION_SPLIT:
|
|
_boundary()
|