From e7cb740b97054e0fdf70ea970edf35f844313ae9 Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 08:18:39 +0800 Subject: [PATCH] Neuron attention: decode fp4 KV records in float arithmetic (Vector engine), not integer bit ops (GpSimd) The compressed-KV scopes are fp4 records; their decode was nibble masks and shifts plus the integer e2m1 / e4m3 decoders on GpSimd (three such chains were the largest cost of the ratio-2 attention graph, device profile bench72). Now: nibbles from floor(b / 16), e2m1 from a piecewise-linear map of the magnitude code, scales through the hardware fp8 conversion. Bit-identical to the integer decode for every byte (host check). bench73 (vllm-sonnet 512, TP=64, nextn=5): TPOT 243.5 -> 228.6 ms, TTFT 4.84 -> 4.35 s, step 699 -> 666 ms, tokens per step identical. Co-Authored-By: Claude Opus 5.5 --- deepinfer/neuron_ops/nki_attention.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index d1ada54..b1c2479 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -119,8 +119,27 @@ def decode_v41_fp8(records: torch.Tensor) -> torch.Tensor: return _e4m3(data, sf).reshape(*lead, D_QK).to(torch.bfloat16) +def _decode_v41_fp4_fast(records: torch.Tensor) -> torch.Tensor: + """``decode_v41_fp4`` in float arithmetic (Vector engine) instead of integer bit ops (GpSimd): the nibbles come + from floor(b / 16), e2m1 from a piecewise-linear map of the 3-bit magnitude code, the e4m3 scales from the + hardware fp8 conversion (see FAST_E4M3). Same values code by code.""" + lead = records.shape[:-1] + b = records[..., : D_QK // 2].to(torch.float32) + hi = torch.floor(b * 0.0625) + lo = b - 16.0 * hi + codes = torch.stack((lo, hi), dim=-1).flatten(-2) # even element in the low nibble + neg = codes >= 8.0 + c = torch.where(neg, codes - 8.0, codes) + mag = torch.where(c < 4.0, 0.5 * c, torch.where(c < 6.0, c - 2.0, 2.0 * c - 8.0)) + values = torch.where(neg, -mag, mag).reshape(*lead, D_QK // 16, 16) + scales = records[..., D_QK // 2 : D_QK // 2 + D_QK // 16].contiguous().view(torch.float8_e4m3fn).to(torch.float32) + return (values * scales.unsqueeze(-1)).reshape(*lead, D_QK).to(torch.bfloat16) + + def decode_v41_fp4(records: torch.Tensor) -> torch.Tensor: """``[..., 288]`` uint8 V4.1 fp4 records -> ``[..., 512]`` bf16.""" + if FAST_E4M3: + return _decode_v41_fp4_fast(records) v = records.to(torch.int32) lead = v.shape[:-1] values = _unpack_e2m1(v[..., : D_QK // 2]).reshape(*lead, D_QK // 16, 16)