mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
28fb291db1
No behaviour change. `CONTRIBUTING.md` requires `ruff check --no-cache` and `ruff format --check --no-cache` to pass, and the SM90 branch predates that requirement: 72 findings, all in the files added here. Most of it is mechanical (`ruff format`), plus: - annotations for the helper signatures (`torch.Tensor` / `int` / `Generator` / return types); hooks are typed as `Callable[[nn.Module, tuple, object], None]` rather than left bare; - two `if`/`else` assignments become ternaries (SIM108), and the index read in `h20_real_requant_check.py` uses a context manager (SIM115). Also corrects the `sparse_mla_fwd` docstring, which still described the old `0 <= idx <= q_position` masking contract. Passing `is_causal=True` there is what produced the all-NaN layer 0: indices address the assembled `[swa cache | current swa | compressed]` buffer, so the row number is not a valid causal bound and the engine must pass False.
96 lines
4.4 KiB
Python
96 lines
4.4 KiB
Python
"""SM90 (Hopper) requantization helpers.
|
||
|
||
The V4.1 checkpoint is exported for Blackwell: FP4 experts and FP8 dense weights, both
|
||
with packed UE8M0 scales on a per-32 block grid. DeepGEMM on SM90 requires FP32 scales
|
||
on a per-128 grid, so the weights are dequantized once, while loading, and re-quantized
|
||
into the SM90 recipe. Everything here runs on the GPU, per tensor, so peak memory stays
|
||
at a couple of blocks.
|
||
|
||
Layout conventions used by the DeepInfer loaders:
|
||
|
||
* dense 2D weights: ``(n_out, k)`` FP8 with scales ``(n_out / block, k / block)``;
|
||
* routed experts: ``(n_local * k, n_out)`` FP8 with scales ``(n_local * k / block, n_out / block)``
|
||
(contraction-dim major, because ``process_moe_weight`` transposes into expert-major
|
||
``(n_local, n_out, k)`` right after loading).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
|
||
FP8_MAX = 448.0
|
||
FP4_BLOCK = 32
|
||
|
||
# E2M1 values indexed by the 4-bit code (sign bit + 3-bit magnitude).
|
||
_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0)
|
||
|
||
|
||
def ue8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
|
||
"""F8_E8M0 (exponent only) → fp32 scale, i.e. ``2 ** (bits - 127)``."""
|
||
exponents = scale.view(torch.uint8).to(torch.int32) - 127
|
||
return torch.ldexp(torch.ones_like(exponents, dtype=torch.float32), exponents)
|
||
|
||
|
||
def _unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
|
||
"""``(..., k // 2)`` int8 packed E2M1 → ``(..., k)`` float32 values."""
|
||
table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=packed.device)
|
||
bytes_view = packed.view(torch.uint8)
|
||
low = (bytes_view & 0x0F).to(torch.int64)
|
||
high = (bytes_view >> 4).to(torch.int64)
|
||
# nibble order: low nibble is the even k element (DeepSeek export packs MLX-style pairs)
|
||
out = torch.empty((*packed.shape[:-1], packed.shape[-1] * 2), dtype=torch.float32, device=packed.device)
|
||
out[..., 0::2] = table[low]
|
||
out[..., 1::2] = table[high]
|
||
return out
|
||
|
||
|
||
def _expand_blocks(block_scale: torch.Tensor, rows: int, cols: int, block: int) -> torch.Tensor:
|
||
"""``(rows / block, cols / block)`` → ``(rows, cols)`` scale map."""
|
||
return block_scale.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1)[:rows, :cols]
|
||
|
||
|
||
def _requant(x: torch.Tensor, dst_block: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""``(rows, cols)`` float32 → FP8 E4M3 with ``dst_block × dst_block`` FP32 scales."""
|
||
rows, cols = x.shape
|
||
assert rows % dst_block == 0 and cols % dst_block == 0, f"{x.shape=} is not a multiple of {dst_block}"
|
||
amax = x.abs().view(rows // dst_block, dst_block, cols // dst_block, dst_block).amax(dim=(1, 3))
|
||
scale = (amax / FP8_MAX).clamp_min(torch.finfo(torch.float32).tiny)
|
||
q = (x.view(rows // dst_block, dst_block, cols // dst_block, dst_block) / scale[:, None, :, None]).to(torch.float8_e4m3fn)
|
||
return q.view(rows, cols), scale
|
||
|
||
|
||
@torch.no_grad()
|
||
def requant_dense_fp8(
|
||
weight: torch.Tensor, sf: torch.Tensor, src_block: int = FP4_BLOCK, dst_block: int = 128
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""FP8 dense weight with per-``src_block`` scales → FP8 with per-``dst_block`` FP32 scales."""
|
||
rows, cols = weight.shape
|
||
x = weight.to(torch.float32) * _expand_blocks(sf.to(torch.float32), rows, cols, src_block)
|
||
return _requant(x, dst_block)
|
||
|
||
|
||
@torch.no_grad()
|
||
def requant_expert_fp4(
|
||
packed: torch.Tensor,
|
||
sf: torch.Tensor,
|
||
out_block: int = 128,
|
||
contraction_major_sf: bool = True,
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""Packed FP4 expert slice ``(n_out, k // 2)`` + ``(n_out, k / 32)`` UE8M0 scales → FP8 + FP32 scales.
|
||
|
||
Returns ``(fp8 (n_out, k), sf (k / block, n_out / block) fp32)`` when ``contraction_major_sf``,
|
||
which is the layout ``process_moe_weight`` expects for the grouped GEMM.
|
||
"""
|
||
n_out, half_k = packed.shape
|
||
k = half_k * 2
|
||
# Expert scales are per-output-channel already: shape `(n_out, k / 32)`. Only the contraction
|
||
# axis is blocked, so only that axis gets expanded -- running this through `_expand_blocks`
|
||
# (which repeats *both* axes) would hand every output channel the scales of the first
|
||
# `n_out / 32` rows, scrambling the whole expert.
|
||
assert tuple(sf.shape) == (n_out, k // FP4_BLOCK), f"unexpected expert scale layout: {tuple(sf.shape)}"
|
||
x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=1)
|
||
q, scale = _requant(x, out_block)
|
||
if contraction_major_sf:
|
||
scale = scale.t().contiguous()
|
||
return q, scale
|