mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
8a5209ee65
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout, pure-torch replacements for the official reference kernels, golden generation, a teacher-forced DeepInfer runner (LLM offline API) and a comparison tool. CPU mode: 25/27 greedy tokens match, worst logits cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060. - model_loader: check parameters against the platform device, not is_cuda. - deep_gemm paged MQA logits: do not dereference block-table slots past the context. - bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would serialize host compute). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
247 lines
11 KiB
Python
247 lines
11 KiB
Python
"""Pure-torch drop-in for the official DeepSeek-V4.1 reference `inference/kernel.py` (TileLang GPU kernels).
|
|
|
|
Mirrors the public API of kernel.py (act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn,
|
|
hc_split_sinkhorn) with the same numerics, so the unmodified reference `model.py` can run on CPU:
|
|
|
|
import ref_kernels_torch, sys
|
|
sys.modules["kernel"] = ref_kernels_torch # before `import model`
|
|
|
|
Numerics reproduced from kernel.py (line numbers refer to the downloaded kernel.py):
|
|
* fast_round_scale (22-37): scale = 2**ceil(log2(amax * fp32(1/max))) taken from the fp32 bit pattern.
|
|
* act_quant (41-124): per (row, block) amax clamped to >= 1e-4; e4m3 RNE cast after clamp(+-448);
|
|
inplace=True writes back dtype(fp32(e4m3(x/s)) * s) (a quant->dequant round trip).
|
|
* fp4_act_quant (128-204): e8m0 scales: amax >= 6*2**-126, pow2-ceil(amax/6); e4m3 scales: amax >= 6*2**-9,
|
|
s = fp32(e4m3(amax / 6)); values RNE to the e2m1 grid after clamp(+-6); inplace => dtype(fp32(e2m1)*s).
|
|
* fp8_gemm / fp4_gemm (208-307, 478-591): per 32-wide K block, fp32 partial GEMM, then
|
|
acc += partial * scale_a[m, k] * scale_b[n(/32), k] in K-block order; output in torch default dtype.
|
|
FP4 weights: float4_e2m1fn_x2 [N, K/2], element 2j = low nibble, 2j+1 = high nibble.
|
|
* sparse_attn (311-403): FlashAttention-style online softmax over blocks of 64 gathered indices
|
|
(-1 = masked, gathered as zeros with score -inf), running max initialised to -1e30, the probabilities
|
|
are cast to bf16 before the P@V GEMM (acc_s_cast), sink added as exp(sink - max) to the denominator.
|
|
* hc_split_sinkhorn (407-474): sigmoid pre (+eps), 2*sigmoid post, softmax(+eps) then column
|
|
normalisation, then (iters-1) x (row, column) normalisation, each with +eps in the denominator.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
FP8_MAX = 448.0
|
|
FP4_MAX = 6.0
|
|
FP4_TABLE = torch.tensor([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])
|
|
_GEMM_BLOCK_K = 32 # kernel.py: block_K = group_size (32 for this model) / fp4 sub_K = 32
|
|
_ATTN_BLOCK = 64 # kernel.py sparse_attn_kernel: block = 64
|
|
|
|
|
|
def _fast_round_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor:
|
|
"""kernel.py fast_round_scale: 2**ceil(log2(amax * max_inv)), via the IEEE-754 bits."""
|
|
v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous()
|
|
bits = v.view(torch.int32)
|
|
exp = (bits >> 23) & 0xFF
|
|
man = bits & ((1 << 23) - 1)
|
|
e = exp - 127 + (man != 0).to(torch.int32)
|
|
# fast_pow2: (e + 127) << 23 reinterpreted as fp32
|
|
return ((e + 127) << 23).to(torch.int32).view(torch.float32)
|
|
|
|
|
|
def _e2m1_round(v: torch.Tensor) -> torch.Tensor:
|
|
"""fp32 in [-6, 6] -> nearest e2m1 value (round half to even), returned as fp32."""
|
|
a = v.abs()
|
|
code = torch.zeros_like(a, dtype=torch.int64)
|
|
for b in (0.25, 1.25, 2.5, 5.0): # ties go down (to the even code)
|
|
code += a > b
|
|
for b in (0.75, 1.75, 3.5): # ties go up (to the even code)
|
|
code += a >= b
|
|
mag = FP4_TABLE.to(v.device)[code]
|
|
return torch.where(v < 0, -mag, mag)
|
|
|
|
|
|
def _e2m1_codes(v: torch.Tensor) -> torch.Tensor:
|
|
a = v.abs()
|
|
code = torch.zeros_like(a, dtype=torch.int32)
|
|
for b in (0.25, 1.25, 2.5, 5.0):
|
|
code += (a > b).to(torch.int32)
|
|
for b in (0.75, 1.75, 3.5):
|
|
code += (a >= b).to(torch.int32)
|
|
return code | ((v < 0).to(torch.int32) << 3)
|
|
|
|
|
|
def _unpack_fp4(b: torch.Tensor) -> torch.Tensor:
|
|
"""float4_e2m1fn_x2 / int8 [..., K/2] -> fp32 [..., K]; low nibble first."""
|
|
x = b.contiguous().view(torch.uint8)
|
|
table = FP4_TABLE.to(x.device)
|
|
lo, hi = (x & 0x0F).long(), (x >> 4).long()
|
|
return torch.stack([table[lo], table[hi]], dim=-1).flatten(-2)
|
|
|
|
|
|
def _to_scale_dtype(s: torch.Tensor, scale_dtype: torch.dtype) -> torch.Tensor:
|
|
if scale_dtype == torch.float8_e8m0fnu:
|
|
# s is a power of two here: its exponent byte is the e8m0 encoding
|
|
return ((s.contiguous().view(torch.int32) >> 23) & 0xFF).to(torch.uint8).view(torch.float8_e8m0fnu)
|
|
return s.to(scale_dtype)
|
|
|
|
|
|
def act_quant(
|
|
x: torch.Tensor,
|
|
block_size: int = 128,
|
|
scale_fmt: Optional[str] = None,
|
|
scale_dtype: torch.dtype = torch.float32,
|
|
inplace: bool = False,
|
|
):
|
|
"""Block-wise FP8 quantization; inplace=True does the fused quant+dequant back to x.dtype."""
|
|
N = x.size(-1)
|
|
assert N % block_size == 0
|
|
z = x.contiguous()
|
|
xf = z.float().unflatten(-1, (N // block_size, block_size))
|
|
amax = xf.abs().amax(-1).clamp_min(1e-4)
|
|
if scale_fmt is not None:
|
|
s = _fast_round_scale(amax, 1.0 / FP8_MAX)
|
|
else:
|
|
s = amax * torch.tensor(1.0 / FP8_MAX, dtype=torch.float32)
|
|
q = (xf / s.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
|
|
if inplace:
|
|
y = (q.float() * s.unsqueeze(-1)).to(x.dtype).flatten(-2)
|
|
x.copy_(y)
|
|
return x
|
|
return q.flatten(-2), _to_scale_dtype(s, scale_dtype)
|
|
|
|
|
|
def fp4_act_quant(
|
|
x: torch.Tensor,
|
|
block_size: int = 32,
|
|
inplace: bool = False,
|
|
scale_dtype: torch.dtype = torch.float8_e8m0fnu,
|
|
):
|
|
"""FP4 with E8M0 scales (indexer) or E4M3 scales (compressed KV); inplace writes the dequantized values."""
|
|
assert scale_dtype in (torch.float8_e8m0fnu, torch.float8_e4m3fn)
|
|
N = x.size(-1)
|
|
assert N % block_size == 0
|
|
z = x.contiguous()
|
|
xf = z.float().unflatten(-1, (N // block_size, block_size))
|
|
amax = xf.abs().amax(-1)
|
|
if scale_dtype == torch.float8_e4m3fn:
|
|
amax = amax.clamp_min(FP4_MAX * 2.0**-9)
|
|
s = (amax / FP4_MAX).to(torch.float8_e4m3fn).float()
|
|
else:
|
|
amax = amax.clamp_min(FP4_MAX * 2.0**-126)
|
|
s = _fast_round_scale(amax, 1.0 / FP4_MAX)
|
|
v = (xf / s.unsqueeze(-1)).clamp(-FP4_MAX, FP4_MAX)
|
|
if inplace:
|
|
y = (_e2m1_round(v) * s.unsqueeze(-1)).to(x.dtype).flatten(-2)
|
|
x.copy_(y)
|
|
return x
|
|
codes = _e2m1_codes(v).flatten(-2)
|
|
packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).to(torch.uint8).view(torch.float4_e2m1fn_x2)
|
|
s_out = _to_scale_dtype(s, scale_dtype) if scale_dtype == torch.float8_e8m0fnu else s.to(torch.float8_e4m3fn)
|
|
return packed, s_out
|
|
|
|
|
|
def _blocked_gemm(A: torch.Tensor, sa: torch.Tensor, B: torch.Tensor, sb_rows: torch.Tensor, act_group: int) -> torch.Tensor:
|
|
"""A [M,K] fp32 (exact fp8 values), sa [M, K/act_group]; B [N,K] fp32 (exact), sb_rows [N, K/32].
|
|
acc += (A_k @ B_k^T) * sa[:, k // n_sub] * sb[:, k] for each 32-wide K block, in order."""
|
|
M, K = A.shape
|
|
N = B.shape[0]
|
|
n_sub = act_group // _GEMM_BLOCK_K
|
|
acc = torch.zeros(M, N, dtype=torch.float32, device=A.device)
|
|
for k in range(K // _GEMM_BLOCK_K):
|
|
sl = slice(k * _GEMM_BLOCK_K, (k + 1) * _GEMM_BLOCK_K)
|
|
part = A[:, sl] @ B[:, sl].T
|
|
acc += part * sa[:, k // n_sub].unsqueeze(1) * sb_rows[:, k].unsqueeze(0)
|
|
return acc
|
|
|
|
|
|
def fp8_gemm(
|
|
a: torch.Tensor,
|
|
a_s: torch.Tensor,
|
|
b: torch.Tensor,
|
|
b_s: torch.Tensor,
|
|
scale_dtype: torch.dtype = torch.float32,
|
|
block_size: int = 128,
|
|
) -> torch.Tensor:
|
|
"""C[M,N] = A[M,K] @ B[N,K]^T with per-block FP8 scaling (weight scales per (N block, K block))."""
|
|
assert a.is_contiguous() and b.is_contiguous()
|
|
assert a_s.is_contiguous() and b_s.is_contiguous()
|
|
assert block_size == _GEMM_BLOCK_K, "this model only uses 32x32 blocks"
|
|
K = a.size(-1)
|
|
M = a.numel() // K
|
|
N = b.size(0)
|
|
assert K % block_size == 0
|
|
assert a_s.numel() == M * (K // block_size)
|
|
assert b_s.shape == ((N + block_size - 1) // block_size, K // block_size)
|
|
sb_rows = b_s.float().repeat_interleave(block_size, dim=0)[:N]
|
|
acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), b.float(), sb_rows, block_size)
|
|
return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N)
|
|
|
|
|
|
def fp4_gemm(
|
|
a: torch.Tensor,
|
|
a_s: torch.Tensor,
|
|
b: torch.Tensor,
|
|
b_s: torch.Tensor,
|
|
scale_dtype: torch.dtype = torch.float32,
|
|
act_block_size: int = 128,
|
|
) -> torch.Tensor:
|
|
"""C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T; B is [N, K/2] float4_e2m1fn_x2 with per-32 E8M0 scales."""
|
|
assert a.is_contiguous() and b.is_contiguous()
|
|
assert a_s.is_contiguous() and b_s.is_contiguous()
|
|
K = a.size(-1)
|
|
M = a.numel() // K
|
|
N = b.size(0)
|
|
assert act_block_size in (32, 128)
|
|
assert K % act_block_size == 0
|
|
assert a_s.numel() == M * (K // act_block_size)
|
|
assert b_s.shape == (N, K // 32)
|
|
acc = _blocked_gemm(a.view(M, K).float(), a_s.view(M, -1).float(), _unpack_fp4(b), b_s.float(), act_block_size)
|
|
return acc.to(torch.get_default_dtype()).view(*a.size()[:-1], N)
|
|
|
|
|
|
def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
|
|
"""q [b,m,h,d] bf16, kv [b,n,d] bf16, attn_sink [h] fp32, topk_idxs [b,m,topk] int32 (-1 = none)."""
|
|
b, m, h, d = q.shape
|
|
topk = topk_idxs.size(-1)
|
|
qf = q.float()
|
|
sink = attn_sink.float()
|
|
scores_max = torch.full((b, m, h), -1e30, dtype=torch.float32, device=q.device)
|
|
sum_exp = torch.zeros((b, m, h), dtype=torch.float32, device=q.device)
|
|
acc_o = torch.zeros((b, m, h, d), dtype=torch.float32, device=q.device)
|
|
batch_ix = torch.arange(b, device=q.device).view(b, 1, 1)
|
|
n_blocks = (topk + _ATTN_BLOCK - 1) // _ATTN_BLOCK
|
|
for t in range(n_blocks):
|
|
idx = topk_idxs[..., t * _ATTN_BLOCK : (t + 1) * _ATTN_BLOCK].long()
|
|
if idx.size(-1) < _ATTN_BLOCK: # positions past topk read as -1
|
|
idx = torch.cat([idx, idx.new_full((b, m, _ATTN_BLOCK - idx.size(-1)), -1)], dim=-1)
|
|
valid = idx != -1
|
|
kv_blk = kv[batch_ix, idx.clamp_min(0)] # [b, m, 64, d]
|
|
kv_blk = torch.where(valid.unsqueeze(-1), kv_blk, torch.zeros((), dtype=kv.dtype)).float()
|
|
s = torch.einsum("bmhd,bmkd->bmhk", qf, kv_blk)
|
|
s = torch.where(valid.unsqueeze(2), s, torch.tensor(-torch.inf))
|
|
s = s * softmax_scale
|
|
prev = scores_max
|
|
scores_max = torch.maximum(prev, s.amax(-1))
|
|
scale = torch.exp(prev - scores_max)
|
|
p = torch.exp(s - scores_max.unsqueeze(-1))
|
|
sum_exp = sum_exp * scale + p.sum(-1)
|
|
acc_o = acc_o * scale.unsqueeze(-1) + torch.einsum("bmhk,bmkd->bmhd", p.to(torch.bfloat16).float(), kv_blk)
|
|
sum_exp = sum_exp + torch.exp(sink.view(1, 1, h) - scores_max)
|
|
return (acc_o / sum_exp.unsqueeze(-1)).to(q.dtype)
|
|
|
|
|
|
def hc_split_sinkhorn(
|
|
mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6
|
|
):
|
|
hc = hc_mult
|
|
m = mixes.float()
|
|
sc, base = hc_scale.float(), hc_base.float()
|
|
pre = torch.sigmoid(m[..., :hc] * sc[0] + base[:hc]) + eps
|
|
post = 2 * torch.sigmoid(m[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc])
|
|
comb = m[..., 2 * hc :].unflatten(-1, (hc, hc)) * sc[2] + base[2 * hc :].view(hc, hc)
|
|
comb = torch.exp(comb - comb.amax(-1, keepdim=True))
|
|
comb = comb / comb.sum(-1, keepdim=True) + eps
|
|
comb = comb / (comb.sum(-2, keepdim=True) + eps)
|
|
for _ in range(sinkhorn_iters - 1):
|
|
comb = comb / (comb.sum(-1, keepdim=True) + eps)
|
|
comb = comb / (comb.sum(-2, keepdim=True) + eps)
|
|
return pre, post, comb
|