coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

neuron: float-only fp8 codec in the compile backend (no GpSimd int ops / remainder)

Browse Source 
Exact for all 256 e4m3fn codes and RNE encode on host; tiny golden 9/9.
vllm-sonnet 512: TPOT 228.6 -> 223.3 ms, TTFT 4349 -> 4117 ms, tokens/step unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 12:41:42 +08:00
parent e7cb740b97
commit d7d2f36ec1
1 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
deepinfer/neuron_ops/neuron_compile.py
View File

@ -43,16 +43,19 @@ def _encode_scalar(value: float) -> int:
def _decode(x_u8: torch.Tensor) -> torch.Tensor:
"""Exact e4m3fn bits -> fp32, elementwise (no table, no bitcast): (8 + m) * 2^(e - 10), subnormals m * 2^-9."""
b = x_u8.to(torch.int32)
e = (b >> 3) & 15
m = (b & 7).to(torch.float32)
p = torch.full_like(m, 2.0**-10)
for bit, factor in ((1, 2.0), (2, 4.0), (4, 16.0), (8, 256.0)):
p = p * torch.where((e & bit) != 0, factor, 1.0)
v = torch.where(e == 0, m * 2.0**-9, (m + 8.0) * p)
v = torch.where(b >= 128, -v, v)
return torch.where((b & 127) == 127, torch.full_like(v, float("nan")), v)
"""Exact e4m3fn bits -> fp32, elementwise (no table, no bitcast): (8 + m) * 2^(e - 10), subnormals m * 2^-9.
Float arithmetic only (Vector / Scalar engines): integer shifts and masks run on GpSimd, and this decode is in
every graph that reads an fp8 activation or cache. ``exp2`` is exact on integers on the device (see
nki_attention._pow2)."""
b = x_u8.to(torch.float32)
neg = b >= 128.0
mag = torch.where(neg, b - 128.0, b)
e = torch.floor(mag * 0.125)
m = mag - 8.0 * e
v = torch.where(e == 0.0, m * 2.0**-9, (m + 8.0) * torch.exp2(e - 10.0))
v = torch.where(neg, -v, v)
return torch.where(mag == 127.0, torch.full_like(v, float("nan")), v)
def _decode_as(dtype: torch.dtype) -> Callable[[torch.Tensor], torch.Tensor]:
@ -65,7 +68,7 @@ def _decode_as(dtype: torch.dtype) -> Callable[[torch.Tensor], torch.Tensor]:
def _rne(m: torch.Tensor) -> torch.Tensor:
f = torch.floor(m)
d = m - f
odd = torch.remainder(f, 2.0) == 1.0
odd = (f - 2.0 * torch.floor(f * 0.5)) == 1.0 # not torch.remainder: it lowers to GpSimd
return f + ((d > 0.5) | ((d == 0.5) & odd)).to(m.dtype)