mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
"""fp32 / bf16 -> e4m3 bytes through the hardware conversion (DeepInfer, issue #29).
|
|
|
|
``neuron_compile.encode_codes`` builds the e4m3 code of every element with ~30 elementwise passes (an exponent search
|
|
by comparisons, then RNE), which is most of the KV-cache store in a 512-row prompt chunk's projection graph. Here the
|
|
tensor engine's own fp32 -> fp8e4m3fn conversion does it in one pass, and the fp8 tile is stored as its bytes. Inputs
|
|
are clamped to ``+-limit`` first (the trn2 e4m3 range, ``neuron_ops.fp8.FP8_MAX``), inside which the hardware format
|
|
and e4m3fn share every bit pattern.
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
#: Partition width.
|
|
PT = 128
|
|
|
|
|
|
@nki.jit
|
|
def kernel(x, limit):
|
|
"""``x [rows, d]`` (fp32 or bf16), ``limit [1, 1]`` fp32 -> ``[rows, d]`` uint8 e4m3 bytes of ``clamp(x, +-limit)``."""
|
|
rows, d = x.shape
|
|
out = nl.ndarray((rows, d), dtype=nl.uint8, buffer=nl.shared_hbm)
|
|
lim = nl.load(limit[0:1, 0:1], dtype=nl.float32)
|
|
|
|
def body(r0, nr):
|
|
t = nl.load(x[r0 : r0 + nr, 0:d], dtype=nl.float32)
|
|
hi = nl.broadcast_to(lim, (nr, 1))
|
|
c = nl.minimum(nl.maximum(t, nl.multiply(hi, -1.0)), hi)
|
|
q = nl.ndarray((nr, d), dtype=nl.float8_e4m3fn, buffer=nl.sbuf)
|
|
nisa.tensor_copy(dst=q, src=c)
|
|
nl.store(out[r0 : r0 + nr, 0:d], value=q.view(nl.uint8))
|
|
|
|
full, tail = rows // PT, rows % PT
|
|
for i in nl.affine_range(full):
|
|
body(i * PT, PT)
|
|
if tail:
|
|
body(full * PT, tail)
|
|
return out
|