From 7e032422d2210f5609a8fd7f795ec70e717ab65b Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 05:00:07 +0800 Subject: [PATCH] Neuron attention: decode fp8 KV records with the hardware fp8 conversion (DEEPINFER_NEURON_FAST_E4M3, default on) 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 --- deepinfer/neuron_ops/nki_attention.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index a94d213..5ba7ef2 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -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)