mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +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>
111 lines
4.7 KiB
Python
111 lines
4.7 KiB
Python
"""Block-quantized GEMMs through the vllm-neuron#71 NKI kernels (issue #10, #25).
|
|
|
|
DeepInfer keeps fp8 weights as ``[n, k]`` e4m3 with per-row, per-32-column packed ue8m0 scales (expanded from the
|
|
checkpoint's 32x32 blocks) and fp4 experts as ``[e, n, k/2]`` e2m1 pairs with per-row scales. The kernels take the
|
|
checkpoint form: e4m3 bytes with a ``[ceil(n/32), ceil(k/32)]`` e8m0 scale, e2m1 bytes with ``[n, k/32]``. The
|
|
kernel scales are prepared once after loading (``prepare_model``) and kept on the weight as ``_nki_scale``.
|
|
|
|
Activations arrive fp8-quantized (1x32 blocks); they are dequantized to bf16 exactly (an e4m3 value times a power of
|
|
two) and fed to the kernel. ``fp4_linear`` re-quantizes them through the same 1x32 fp8 round trip, which is
|
|
idempotent on already-quantized values.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import os
|
|
|
|
import torch
|
|
|
|
ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_GEMM", "1") == "1"
|
|
|
|
|
|
def _bytes(sf: torch.Tensor) -> torch.Tensor:
|
|
"""Packed ue8m0 scales (int16 pairs, uint8 or float8_e8m0fnu) -> uint8 exponent bytes, row-major."""
|
|
return sf.contiguous().view(torch.uint8)
|
|
|
|
|
|
def _kernel_scale(weight: torch.Tensor, sf: torch.Tensor) -> torch.Tensor | None:
|
|
"""The kernel's scale for one weight, or None if its scales do not fit the kernel's layout."""
|
|
k_blocks = -(-(weight.shape[-1] * (2 if weight.dtype == torch.int8 else 1)) // 32)
|
|
b = _bytes(sf)[..., :k_blocks]
|
|
if weight.dtype == torch.int8: # fp4 experts: per-row scales, as the kernel takes them
|
|
return b.contiguous()
|
|
n = weight.shape[-2]
|
|
blocks = b[..., ::32, :]
|
|
if not torch.equal(blocks.repeat_interleave(32, dim=-2)[..., :n, :], b): # rows of a block must share a scale
|
|
return None
|
|
return blocks.contiguous()
|
|
|
|
|
|
def prepare_model(model: torch.nn.Module) -> int:
|
|
"""Attach ``_nki_scale`` to every fp8 / fp4 weight the kernels can take. Returns the count."""
|
|
if not ENABLED:
|
|
return 0
|
|
count = 0
|
|
for p in model.parameters():
|
|
sf = getattr(p, "sf", None)
|
|
if sf is None or p.dtype not in (torch.float8_e4m3fn, torch.int8) or p.dim() not in (2, 3):
|
|
continue
|
|
# On the host: the scale tensors are small, and eager device ops would compile once per weight shape.
|
|
sf_host = (sf.data if isinstance(sf, torch.nn.Parameter) else sf).cpu()
|
|
scale = _kernel_scale(p.data, sf_host)
|
|
if scale is not None:
|
|
scale = scale.to(p.device)
|
|
p._nki_scale = scale
|
|
# The weight bytes as their own uint8 tensor (same storage): a dtype bitcast of a graph input feeding the
|
|
# kernel crashes neuronx-cc (InferNonlocalTensors), a uint8 input does not.
|
|
p._nki_weight = p.data.view(torch.uint8)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
@functools.cache
|
|
def _kernel(name: str): # noqa: ANN202
|
|
from torch_neuronx.nki_hop import wrap_nki
|
|
|
|
from deepinfer.neuron_ops.nki import blockquant
|
|
|
|
return wrap_nki(getattr(blockquant, name))
|
|
|
|
|
|
def _activation_bf16(a: torch.Tensor, sfa: torch.Tensor, k: int) -> torch.Tensor:
|
|
from deep_gemm._quant import decode_data, sf_to_scales
|
|
|
|
return (decode_data(a) * sf_to_scales(sfa, k)).to(torch.bfloat16)
|
|
|
|
|
|
def supported(a: torch.Tensor, b: torch.Tensor, k: int, n: int) -> bool:
|
|
return (
|
|
ENABLED
|
|
and a.device.type == "neuron"
|
|
and getattr(b, "_nki_scale", None) is not None
|
|
and k % 128 == 0
|
|
and n % 32 == 0
|
|
)
|
|
|
|
|
|
def gemm(a: torch.Tensor, sfa: torch.Tensor, b: torch.Tensor, d: torch.Tensor, group: int | None = None) -> None:
|
|
"""``d = dequant(a) @ dequant(b[group])^T`` through the NKI kernels."""
|
|
m = a.shape[0]
|
|
x = _activation_bf16(a, sfa, a.shape[-1])
|
|
scale = b._nki_scale if group is None else b._nki_scale[group]
|
|
w = b._nki_weight if group is None else b._nki_weight[group]
|
|
name = "fp4_linear" if b.dtype == torch.int8 else "fp8_linear_raw"
|
|
out = _kernel(name)(x, w, scale)
|
|
d.copy_(out[:m].to(d.dtype))
|
|
|
|
|
|
def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> torch.Tensor:
|
|
"""bf16 ``x [m, k]`` against a prepared fp8 / fp4 weight, the kernel quantizing ``x`` to fp8 (1x32, power-of-two
|
|
scales, max 224) on the way in -- the same quantization the shims apply, so callers can skip the fp8 activation
|
|
and its int16 scales entirely. Returns bf16 ``[m, n]``."""
|
|
scale = w._nki_scale if group is None else w._nki_scale[group]
|
|
weight = w._nki_weight if group is None else w._nki_weight[group]
|
|
name = "fp4_linear" if w.dtype == torch.int8 else "fp8_linear"
|
|
return _kernel(name)(x, weight, scale)[: x.shape[0]]
|
|
|
|
|
|
def has_kernel(w: torch.Tensor | None) -> bool:
|
|
return ENABLED and w is not None and getattr(w, "_nki_scale", None) is not None and w.device.type == "neuron"
|