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)