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

Neuron attention: decode fp8 KV records with the hardware fp8 conversion (DEEPINFER_NEURON_FAST_E4M3, default on)

Browse Source 
The record decode was a dozen integer GpSimd ops per element. Viewing the bytes
as float8_e4m3fn (read as trn2 e4m3 under the existing compile flag; DeepInfer
quantizes to 224 on Neuron, where the encodings agree) is one conversion.
Outputs identical (tiny device golden bit-identical; full-model tokens per
step identical), bench68 (vllm-sonnet 512): TPOT 262.0 -> 258.0 ms, TTFT
4.98 -> 4.90 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 05:00:07 +08:00
parent 1dce2c4995
commit 7e032422d2
1 changed files with 14 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
deepinfer/neuron_ops/nki_attention.py
View File

@ -96,8 +96,22 @@ def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
return torch.stack((_e2m1(packed & 15), _e2m1((packed >> 4) & 15)), dim=-1).flatten(-2)
#: DEEPINFER_NEURON_FAST_E4M3: decode fp8 record bytes with the hardware fp8 -> fp32 conversion instead of integer
#: bit arithmetic (a dozen GpSimd ops an element; the attention graphs' largest cost, device profile bench67). The
#: compile flag --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 makes the device read the bytes as trn2 e4m3; DeepInfer
#: quantizes to 224 on Neuron, where the two encodings agree (exponent field <= 14). NaN codes stay NaN (the table
#: path clears them).
FAST_E4M3 = os.environ.get("DEEPINFER_NEURON_FAST_E4M3", "1") == "1"
def decode_v41_fp8(records: torch.Tensor) -> torch.Tensor:
"""``[..., 528]`` uint8 V4.1 fp8 records -> ``[..., 512]`` bf16 (``kv_cache.dequantize_records``)."""
if FAST_E4M3:
lead = records.shape[:-1]
data = records[..., :D_QK].contiguous().view(torch.float8_e4m3fn).to(torch.float32)
sf = records[..., D_QK : D_QK + D_QK // 32].to(torch.int32) - 127
scale = _pow2(torch.clamp(sf, min=-126)).unsqueeze(-1)
return (data.reshape(*lead, D_QK // 32, 32) * scale).reshape(*lead, D_QK).to(torch.bfloat16)
v = records.to(torch.int32)
lead = v.shape[:-1]
data = v[..., :D_QK].reshape(*lead, D_QK // 32, 32)