coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/neuron_ops/nki/blockquant.py
whn09 9c8c85d83d Neuron: per-layer region graphs + vllm-neuron#71 NKI kernels (issues #10, #11, #13, #25) 
- Decode runs eagerly and compiles per-layer regions (attention half, MoE) with stage boundaries
  (neuron_ops/regions.py, DEEPINFER_NEURON_GRAPH_MODE=regions); layer KV views are sliced from the
  roots inside the graph (_LazyKVCache).
- vllm-neuron#71 NKI kernels (Apache-2.0, vendored under neuron_ops/nki/ with attribution), called via
  torch_neuronx.nki_hop.wrap_nki: blockquant (fp8/fp4 GEMMs incl. grouped einsum, MoE and shared
  expert on bf16 activations), sparse_attn / index_score / compress_pool (nki_attention.py),
  hc_mixes / hc_expand / hc_collapse / hc_attn_out / rmsnorm / swiglu (nki_hc.py, nki_norm.py).
- Compile-safe formulations: row-major scales, no dtype bitcasts of graph inputs (uint8 weight
  tensors), whole-row KV stores, byte assembly in int32 narrowed once at the store, arithmetic int16
  scale packing, plain fma/rsqrt and bf16 F.linear under compile; default backend deepinfer_neuron
  (torch fp8 carried as uint8, avoiding NCC_EOCP001 with the kernels' OCP fp8).
- Host: tiny model 9/9 vs golden on all prompts; device: kernel tests pass (13 attention, 88 hc/norm);
  region decode graph compiles except the attention projection stage (being reduced).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-22 17:36:54 +08:00

1236 lines
47 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/blockquant.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""Shape-general block-quantized linear kernels — the arithmetic every other
module in this model is mostly made of.
fp8 `[out, in]` e4m3 + `[ceil(out/32), ceil(in/32)]` e8m0 attention, shared experts, lm head path
fp4 `[out, in//2]` two e2m1 per byte + `[out, in/32]` e8m0 routed experts
bf16 `[out, in]`, no scale compressor, indexer wk, router
* The activation round trip divides by `FP8_PEAK = 224`, not by the
reference's 448, because Trn2's `nl.float8_e4m3` saturates at 240. It
costs one mantissa bit against `model.act_quant` and is what the
compiler's `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` flag is about.
* The weight decoders leave the exponent bias out of the mantissa decode and
fold it into the scale, so a decoded value is a *very* small bf16 (an e4m3
`eeee = 0` lands on a bf16 denormal) and only the product with the scale is
in range. Do not "simplify" either half on its own.
"""
import nki
import nki.isa as nisa
import nki.language as nl
#: Partition width of the tensor engine; every tile is 128 rows.
PT = 128
#: Widest free-axis tile. A psum bank holds 2 KiB per partition = 512 fp32.
NT = 512
#: Quantization block along the contracted axis, both for the activation and
#: for the stored scales.
BLK = 32
FP8_MAX = 448.0
#: Half of `FP8_MAX`: see the module docstring.
FP8_PEAK = FP8_MAX / 2.0
FP8_AMAX_FLOOR = 1e-4
MGROUP_BYTES = 64 * 1024
WLOOK = 8
XLOOK = 2
BYTE_HI = 256.0
E4M3_MAG2 = 0x7F007F00
E4M3_SGN2 = 0x80008000 - (1 << 32)
#: fp8: a decoded e4m3 comes out 2**-120 too small, and the stored scale means
#: 2**(code-127), so one multiply by 2**(code-7) closes both.
SC_EXP_ADD = 120
#: fp4: 0x1040. Multiplying a 4-bit code by this lands the sign in bit 15, the
#: two exponent bits in 7..8 and the mantissa bit in 6, which is a bf16 whose
#: exponent field *is* the e2m1 exponent.
FP4_NIB_MUL = 4160.0
#: sign | exponent(2) | mantissa(1) of the bf16 built above; everything else
#: the multiply produced is spill.
FP4_FIELD = 0x81C0
#: fp4: a decoded e2m1 is 2**-126 too small and the stored scale is
#: 2**(code-127), so the closing factor is 2**(code-1). Split as
#: 2**(code-1-FP4_SPLIT) on the bf16 weight and 2**FP4_SPLIT on the fp32
FP4_SPLIT = 64
FP4_POST = float(2.0**FP4_SPLIT)
FP4_EXP_BIAS = 127 - 1 - FP4_SPLIT
def _mgroup(mtiles, ktiles):
"""How many m-tiles to stage at once: the largest power of two that both
divides `mtiles` and keeps the staged activation under `MGROUP_BYTES`.
"""
mg = 1
while (
mg * 2 <= mtiles
and mtiles % (mg * 2) == 0
and (mg * 2) * ktiles * PT * 2 <= MGROUP_BYTES
):
mg *= 2
return mg
#: Free-axis tile widths to try, widest first. Any divisor of `n` not larger
#: than `NT` is correct; wider is fewer tiles, so fewer instructions.
_NTILES = (NT, 256, PT, 64, 32, 16, 8, 4, 2)
def _ntile(n):
"""Free-axis tile width: as wide as a psum bank allows, and a divisor."""
if n <= NT:
return n
for cand in _NTILES:
if n % cand == 0:
return cand
return 1
def _pow2_ceil_log2(x):
"""`2 ** ceil(log2 x)` and its reciprocal, by exponent arithmetic."""
bits = nl.copy(x, dtype=nl.float32).view(nl.int32)
exp = nl.bitwise_and(nl.right_shift(bits, 23), 0xFF)
man = nl.bitwise_and(bits, 0x7FFFFF)
bump = nl.copy(nl.greater(man, 0), dtype=nl.int32)
e = nl.add(exp, bump)
scale = nl.copy(nl.left_shift(e, 23)).view(nl.float32)
inv = nl.copy(nl.left_shift(nl.subtract(254, e), 23)).view(nl.float32)
return scale, inv
def _qdq_fp8(tile, rows, n):
"""Round `tile` through blockwise fp8 and return it in bf16."""
nb = n // BLK
blocked = tile.reshape((rows, nb, BLK))
ab = nl.ndarray((rows, nb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.activation(dst=ab, op=nl.abs, data=blocked, bias=0.0, scale=1.0)
amax = nl.maximum(
nl.copy(nl.max(ab, axis=2, keepdims=True), dtype=nl.float32), FP8_AMAX_FLOOR
)
scale, rcp = _pow2_ceil_log2(nl.multiply(amax, 1.0 / FP8_PEAK))
rcp_b = nl.copy(rcp, dtype=nl.bfloat16)
scale_b = nl.copy(scale, dtype=nl.bfloat16)
q8 = nl.ndarray((rows, nb, BLK), dtype=nl.float8_e4m3fn, buffer=nl.sbuf)
nisa.tensor_tensor(
dst=q8, data1=blocked, op=nl.multiply, data2=nl.broadcast_to(rcp_b, (rows, nb, BLK))
)
deq = nl.ndarray((rows, nb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_tensor(
dst=deq, data1=q8, op=nl.multiply, data2=nl.broadcast_to(scale_b, (rows, nb, BLK))
)
return deq.reshape((rows, n))
def _decode_e4m3(byte_tile, rows, k):
"""e4m3 bytes to bf16, 2**-120 too small (see `SC_EXP_ADD`)."""
t = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=t, data=byte_tile, op0=nl.multiply, operand0=BYTE_HI, engine=nisa.engine.scalar
)
hi = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=hi.view(nl.uint32),
data=t.view(nl.uint32),
op0=nl.bitwise_and,
operand0=E4M3_MAG2,
op1=nl.right_shift,
operand1=4,
engine=nisa.engine.vector,
)
bits = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
nisa.scalar_tensor_tensor(
dst=bits.view(nl.uint32),
data=t.view(nl.uint32),
op0=nl.bitwise_and,
operand0=E4M3_SGN2,
op1=nl.bitwise_or,
operand1=hi.view(nl.uint32),
)
return bits.view(nl.bfloat16)
def _dequant_fp8(w_u8, s_u8, dst_hbm, n, k):
"""fp8 `[n, k]` with a `32 x 32`-tile e8m0 scale, into a bf16 HBM scratch."""
kb = k // BLK
rows = PT if n % PT == 0 else BLK
nsb = rows // BLK
for nb in nl.affine_range(n // rows):
wq = _decode_e4m3(
nl.load(w_u8[nb * rows : (nb + 1) * rows, 0:k], dtype=nl.uint8), rows, k
)
sf = nl.ndarray((rows, kb), dtype=nl.bfloat16, buffer=nl.sbuf)
for i in nl.static_range(nsb):
sb = nl.load(s_u8[nb * nsb + i : nb * nsb + i + 1, 0:kb], dtype=nl.int32)
one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(nl.float32)
nisa.tensor_copy(
dst=sf[i * BLK : (i + 1) * BLK, 0:kb],
src=nl.broadcast_to(one, (BLK, kb)),
)
rep = nl.broadcast_to(sf.reshape((rows, kb, 1)), (rows, kb, BLK))
prod = nl.ndarray((rows, k), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_tensor(
dst=prod, data1=wq, data2=rep.reshape((rows, k)), op=nl.multiply
)
nl.store(dst_hbm[nb * rows : (nb + 1) * rows, 0:k], value=prod)
#: Widest n-slab one `_dequant_fp8_k` iteration decodes at once, in elements.
_DQ_SLAB = 8192
def _nslab(n):
"""Widest divisor of `n` at most `_DQ_SLAB` and a multiple of `BLK`."""
if n <= _DQ_SLAB:
return n
best = BLK
for cand in range(BLK, _DQ_SLAB + 1, BLK):
if n % cand == 0:
best = cand
return best
def _dequant_fp8_k(w_u8, s_u8, dst_hbm, k, n):
"""fp8 `[k, n]` **k-major** with a `32 x 32`-tile e8m0 scale, into bf16 HBM."""
assert n % BLK == 0, f"n={n} must be a multiple of {BLK} for a k-major scale"
slab = _nslab(n)
nb = slab // BLK
nsb = PT // BLK
for kt in nl.affine_range(k // PT):
for ns in nl.affine_range(n // slab):
wq = _decode_e4m3(
nl.load(
w_u8[kt * PT : (kt + 1) * PT, ns * slab : (ns + 1) * slab],
dtype=nl.uint8,
),
PT,
slab,
)
sf = nl.ndarray((PT, nb), dtype=nl.bfloat16, buffer=nl.sbuf)
for i in nl.static_range(nsb):
sb = nl.load(
s_u8[
kt * nsb + i : kt * nsb + i + 1,
ns * nb : (ns + 1) * nb,
],
dtype=nl.int32,
)
one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(
nl.float32
)
nisa.tensor_copy(
dst=sf[i * BLK : (i + 1) * BLK, 0:nb],
src=nl.broadcast_to(one, (BLK, nb)),
)
rep = nl.broadcast_to(sf.reshape((PT, nb, 1)), (PT, nb, BLK))
prod = nl.ndarray((PT, slab), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_tensor(
dst=prod, data1=wq, data2=rep.reshape((PT, slab)), op=nl.multiply
)
nl.store(
dst_hbm[kt * PT : (kt + 1) * PT, ns * slab : (ns + 1) * slab],
value=prod,
)
def _dequant_fp4(w_u8, s_u8, dst_hbm, n, k):
"""fp4 `[n, k//2]` with a per-row, per-32-column e8m0 scale, into bf16 HBM."""
kb = k // BLK
half = BLK // 2
rows = PT if n % PT == 0 else BLK
for nb in nl.affine_range(n // rows):
pk = nl.load(
w_u8[nb * rows : (nb + 1) * rows, 0 : k // 2], dtype=nl.uint8
).reshape((rows, kb, half))
pk16 = nl.ndarray((rows, kb, half), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=pk16, data=pk, op0=nl.multiply, operand0=1.0, engine=nisa.engine.scalar
)
nib = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=nib[:, :, 0:BLK:2], data=pk16, op0=nl.bitwise_and, operand0=0x0F
)
nisa.tensor_scalar(
dst=nib[:, :, 1:BLK:2],
data=pk16,
op0=nl.right_shift,
operand0=4,
op1=nl.bitwise_and,
operand1=0x0F,
)
spread = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=spread,
data=nib,
op0=nl.multiply,
operand0=FP4_NIB_MUL,
engine=nisa.engine.scalar,
)
bits = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=bits, data=spread, op0=nl.bitwise_and, operand0=FP4_FIELD
)
code = nl.load(s_u8[nb * rows : (nb + 1) * rows, 0:kb], dtype=nl.uint8)
sf = nl.ndarray((rows, kb), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=sf,
data=code,
op0=nl.multiply,
operand0=128.0,
op1=nl.add,
operand1=float(FP4_EXP_BIAS * 128),
engine=nisa.engine.scalar,
)
rep = nl.broadcast_to(
sf.view(nl.bfloat16).reshape((rows, kb, 1)), (rows, kb, BLK)
)
prod = nl.ndarray((rows, kb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_tensor(
dst=prod, data1=bits.view(nl.bfloat16), data2=rep, op=nl.multiply
)
nl.store(
dst_hbm[nb * rows : (nb + 1) * rows, 0:k], value=prod.reshape((rows, k))
)
def _gemm(x_hbm, w_hbm, dst_hbm, m, k, n, out_scale=1.0, qdq=False, w_kmajor=False):
"""`dst[m, n] = x[m, k] @ w[n, k].T`, both operands bf16 in HBM."""
ktiles = k // PT
nt = _ntile(n)
mg = _mgroup(m // PT, ktiles)
ntiles = n // nt
for mgb in nl.affine_range((m // PT) // mg):
xt = nl.ndarray((PT, mg, ktiles, PT), dtype=nl.bfloat16, buffer=nl.sbuf)
def stage_x(mi):
for kb in nl.static_range(ktiles):
nisa.dma_transpose(
dst=xt[0:PT, mi, kb, 0:PT],
src=x_hbm[
(mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
kb * PT : (kb + 1) * PT,
],
)
def ntile(nb, jit):
ntp = -(-nt // 16) * 16
wt = nl.ndarray((PT, ktiles, ntp), dtype=nl.bfloat16, buffer=nl.sbuf)
def stage_w(kb):
if w_kmajor:
nisa.dma_copy(
dst=wt[0:PT, kb, 0:nt],
src=w_hbm[kb * PT : (kb + 1) * PT, nb * nt : (nb + 1) * nt],
)
else:
nisa.dma_transpose(
dst=wt[0:PT, kb, 0:nt],
src=w_hbm[nb * nt : (nb + 1) * nt, kb * PT : (kb + 1) * PT],
)
for kb in nl.static_range(min(WLOOK, ktiles)):
stage_w(kb)
for mi in nl.static_range(mg):
if jit and mi + XLOOK < mg:
stage_x(mi + XLOOK)
acc = nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum)
for kb in nl.static_range(ktiles):
if mi == 0 and kb + WLOOK < ktiles:
stage_w(kb + WLOOK)
nisa.nc_matmul(
dst=acc,
stationary=xt[0:PT, mi, kb, 0:PT],
moving=wt[0:PT, kb, 0:nt],
accumulate=(kb > 0),
)
if out_scale == 1.0:
row = nl.copy(acc, dtype=nl.bfloat16)
else:
row = nl.ndarray((PT, nt), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=row, data=acc, op0=nl.multiply, operand0=out_scale
)
nl.store(
dst_hbm[
(mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
nb * nt : (nb + 1) * nt,
],
value=_qdq_fp8(row, PT, nt) if qdq else row,
)
for mi0 in nl.static_range(min(XLOOK, mg)):
stage_x(mi0)
ntile(0, mg > XLOOK)
if ntiles > 1:
for nb in nl.affine_range(ntiles - 1):
ntile(nb + 1, False)
def _check(k, n):
assert k % PT == 0, f"k={k} must be a multiple of {PT}"
assert k % BLK == 0, f"k={k} must be a multiple of {BLK}"
assert n % _ntile(n) == 0, f"n={n} has no usable free-axis tile"
def _stage_rows(x, m, k, dst_hbm, quantize):
"""Copy `x[m, k]` into a bf16 `[ceil(m/128)*128, k]` HBM scratch."""
full = m // PT
tail = m - full * PT
for mb in nl.affine_range(full):
tile = nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16)
nl.store(
dst_hbm[mb * PT : (mb + 1) * PT, 0:k],
value=_qdq_fp8(tile, PT, k) if quantize else tile,
)
if tail:
padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.memset(padded, 0.0)
nisa.tensor_copy(
dst=padded[0:tail, 0:k],
src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
)
nl.store(
dst_hbm[full * PT : (full + 1) * PT, 0:k],
value=_qdq_fp8(padded, PT, k) if quantize else padded,
)
@nki.jit
def fp8_linear(x, w, scale):
"""`x[m, k]` (+ fp8 activation round trip) against fp8 `w[n, k]`."""
m, k = x.shape
n = w.shape[0]
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_dequant_fp8(w, scale, wd, n, k)
_gemm(xq, wd, out, mpad, k, n)
return out
@nki.jit
def fp8_linear_raw(x, w, scale):
"""`fp8_linear` without the activation round trip."""
m, k = x.shape
n = w.shape[0]
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xs, False)
_dequant_fp8(w, scale, wd, n, k)
_gemm(xs, wd, out, mpad, k, n)
return out
def _fp8_2core_span(w, scale, programs, core):
"""`(span, lo, w_rows, scale_rows)` for this core's share of an fp8 GEMM.
The column cut is the same free axis `fp4_linear_2core` divides, with one
extra thing to get right: **fp8 carries its scale at one row per `BLK`
weight rows**, not one per weight row. `_dequant_fp8` reads
`s_u8[nb * nsb + i]` while loading `w_u8[nb * rows ...]` with
`nsb = rows // BLK`, so a weight slice starting at row `lo` has to be paired
with the scale slice starting at `lo // BLK`, and pairing it with `lo` would
dequantise every column of this core's half against the wrong exponent --
silently, with plausible-looking magnitudes.
It divides exactly. The caller asserts `n % (programs * PT) == 0`, so `span`
is a multiple of `PT = 128` and `lo = core * span` is a multiple of 128,
which is four whole `BLK = 32` scale rows. `_dequant_fp8` then takes
`rows = PT` on the slice, exactly as it does on the whole weight, so the
per-tile arithmetic inside it is unchanged rather than merely equivalent.
"""
n = w.shape[0]
k = w.shape[1]
span = n // programs
lo = core * span
return (
span,
lo,
w[lo : lo + span, 0:k],
scale[lo // BLK : (lo + span) // BLK],
)
def _fp8_2core_check(programs, n, name):
assert programs == 2, (
f"launch grid is {programs}; {name} divides its output columns between "
"the two physical cores of a logical NeuronCore and computes only its "
f"own half, so at grid 1 it would return a half-empty result. Launch it "
f"as `wrap_nki({name})[2]`, or call the undivided kernel instead."
)
assert n % (programs * PT) == 0, (
f"n={n} does not divide into {programs} shares of a multiple of {PT}"
)
@nki.jit
def fp8_linear_2core(x, w, scale):
"""`fp8_linear` with the output columns divided across both physical cores.
Launch at `grid=(2,)`. See `fp4_linear_2core` for why the split exists and
what it costs (a duplicated `_stage_rows` into each core's private HBM), and
`_fp8_2core_span` for the scale-row arithmetic this storage adds.
Measured by iteration 5 at `in512_bs1` on 64 ranks, as the second half of
`DEEPSEEK_V41_BQ_2CORE=fp4,fp8`: TTFT 888.2 -> 841.8 ms against the fp4-only
arm, TPOT 111.2 ms against a 143.3 ms baseline, per-stream decode 8.99
tok/s. That iteration was rejected for an unrelated verifier failure and the
kernel was lost with its worktree; `.autohelix/notes/iter-5.md` §8-§9 is the
implementation record it was rewritten from.
**`nki.simulate` cannot check this kernel.** It raises `OverflowError` on
`E4M3_SGN2` inside `_decode_e4m3` in the numpy backend -- in the undivided
kernel too, so it is not the split's doing (iter-5.md §9). Validation is a
device run against the undivided arm plus the index argument above.
"""
m, k = x.shape
n = w.shape[0]
programs = nl.num_programs(0)
core = nl.program_id(0)
_fp8_2core_check(programs, n, "fp8_linear_2core")
span, lo, w_share, s_share = _fp8_2core_span(w, scale, programs, core)
_check(k, span)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_dequant_fp8(w_share, s_share, wd, span, k)
_gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span)
return out
@nki.jit
def fp8_linear_raw_2core(x, w, scale):
"""`fp8_linear_raw` split across both physical cores; see `fp8_linear_2core`.
The only difference from the quantising form is `_stage_rows(..., False)`,
exactly as `fp8_linear_raw` differs from `fp8_linear`.
"""
m, k = x.shape
n = w.shape[0]
programs = nl.num_programs(0)
core = nl.program_id(0)
_fp8_2core_check(programs, n, "fp8_linear_raw_2core")
span, lo, w_share, s_share = _fp8_2core_span(w, scale, programs, core)
_check(k, span)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xs, False)
_dequant_fp8(w_share, s_share, wd, span, k)
_gemm(xs, wd, out[0:mpad, lo : lo + span], mpad, k, span)
return out
#: SBUF bytes a partition may spend on the staged weight in the fused GEMM.
_WT_BYTES = 96 * 1024
def _nchunk(n, nt, ktiles):
"""Widest multiple of `nt` that divides `n` and fits `_WT_BYTES`."""
cap = max(nt, _WT_BYTES // (ktiles * 2))
best, cand = nt, nt
while cand <= cap:
if n % cand == 0:
best = cand
cand += nt
return best
def _decode_fp8_tile(w, s, dst, kb, c0, nc):
"""Decode the k-major fp8 tile `w[kb*128:(kb+1)*128, c0:c0+nc]` into `dst`."""
nbk = nc // BLK
nsb = PT // BLK
byte = nl.load(w[kb * PT : (kb + 1) * PT, c0 : c0 + nc], dtype=nl.uint8)
wq = _decode_e4m3(byte, PT, nc)
sf = nl.ndarray((PT, nbk), dtype=nl.bfloat16, buffer=nl.sbuf)
for i in nl.static_range(nsb):
sb = nl.load(
s[kb * nsb + i : kb * nsb + i + 1, c0 // BLK : c0 // BLK + nbk],
dtype=nl.int32,
)
one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(nl.float32)
nisa.tensor_copy(
dst=sf[i * BLK : (i + 1) * BLK, 0:nbk],
src=nl.broadcast_to(one, (BLK, nbk)),
)
rep = nl.broadcast_to(sf.reshape((PT, nbk, 1)), (PT, nbk, BLK))
nisa.tensor_tensor(
dst=dst, data1=wq, data2=rep.reshape((PT, nc)), op=nl.multiply
)
def _stage_x_tile(x_hbm, xt, mgb, mg, ktiles, mi):
"""One m-tile of the activation, transposed into SBUF."""
for kb in nl.static_range(ktiles):
nisa.dma_transpose(
dst=xt[0:PT, mi, kb, 0:PT],
src=x_hbm[
(mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
kb * PT : (kb + 1) * PT,
],
)
def _gemm_fp8_fused(x_hbm, w, s, dst_hbm, m, k, n):
"""`dst[m, n] = x[m, k] @ w[k, n]` with the weight decoded in SBUF."""
ktiles = k // PT
nt = _ntile(n)
mg = _mgroup(m // PT, ktiles)
nc = _nchunk(n, nt, ktiles)
for mgb in nl.affine_range((m // PT) // mg):
xt = nl.ndarray((PT, mg, ktiles, PT), dtype=nl.bfloat16, buffer=nl.sbuf)
for mi in nl.static_range(mg):
_stage_x_tile(x_hbm, xt, mgb, mg, ktiles, mi)
for cb in nl.affine_range(n // nc):
wt = nl.ndarray((PT, ktiles, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
for kb in nl.static_range(ktiles):
_decode_fp8_tile(w, s, wt[0:PT, kb, 0:nc], kb, cb * nc, nc)
for nb in nl.static_range(nc // nt):
for mi in nl.static_range(mg):
acc = nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum)
for kb in nl.static_range(ktiles):
nisa.nc_matmul(
dst=acc,
stationary=xt[0:PT, mi, kb, 0:PT],
moving=wt[0:PT, kb, nb * nt : (nb + 1) * nt],
accumulate=(kb > 0),
)
nl.store(
dst_hbm[
(mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
],
value=nl.copy(acc, dtype=nl.bfloat16),
)
@nki.jit
def fp8_linear_fused(x, w, scale):
"""`fp8_linear` with the weight decoded in SBUF — see `_gemm_fp8_fused`."""
m, k = x.shape
n = w.shape[1]
assert w.shape[0] == k, f"fp8 weight {w.shape} is not k-major for k={k}"
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_gemm_fp8_fused(xq, w, scale, out, mpad, k, n)
return out
@nki.jit
def fp8_linear_kmajor(x, w, scale):
"""**Unreachable**, and kept as the measured k-major arm — see
`_dequant_fp8_k`. Bit-identical to `fp8_linear` on the same weight
transposed, and slower end to end.
"""
m, k = x.shape
n = w.shape[1]
assert w.shape[0] == k, f"fp8 weight {w.shape} is not k-major for k={k}"
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((k, n), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_dequant_fp8_k(w, scale, wd, k, n)
_gemm(xq, wd, out, mpad, k, n, w_kmajor=True)
return out
@nki.jit
def fp4_linear(x, w, scale):
"""`x[m, k]` (+ fp8 activation round trip) against fp4 `w[n, k//2]`."""
m, k = x.shape
n = w.shape[0]
assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_dequant_fp4(w, scale, wd, n, k)
_gemm(xq, wd, out, mpad, k, n, out_scale=FP4_POST)
return out
@nki.jit
def fp4_linear_2core(x, w, scale):
"""`fp4_linear` with the output columns divided across the two physical
cores of a logical NeuronCore. Launch it at `grid=(2,)`.
**Why this exists.** `wrap_nki` hardcodes `grid=[]`
(`libtorch_neuronx_lite/nki/nki_hop.py:657`), which `compile_nki` turns
into `lnc = 1`, while importing lite sets `NEURON_LOGICAL_NC_CONFIG=2`.
Iteration 4 read the lowering off `module.mlir` for `fp4_linear` at both
grids, 128x5120x2304 — the exact shape a routed expert step runs:
grid 1 one function @blockquant.fp4_linear 360 nc_matmul
grid 2 two functions @...fp4_linear_lnc{0,1} 720 nc_matmul
One function and no `lnc_id` at grid 1, so there is no code for the second
physical core; and at grid 2 *every* op count is exactly doubled, because
`fp4_linear` divides nothing by `nl.program_id`, so both cores would
compute the same GEMM over the same data. Hence this: the same work, cut in
two, so the two per-core functions' op counts **sum** to the grid-1 total
instead of doubling it.
The cut is the free axis. `_gemm` writes `dst_hbm[rows, nb*nt:(nb+1)*nt]`
and `_dequant_fp4` reads a row range of `w`, so core `p` owns columns
`[p*span, (p+1)*span)` of the output and the matching rows of the weight:
disjoint reads, disjoint writes, no barrier and no exchange. What it does
duplicate is `_stage_rows` — `m x k` bf16 into each core's *private* HBM
scratch, 1.3 MB at m=128, k=5120 — because both cores need every row of the
stationary operand and a barrier to share one copy would cost more than the
copy does.
Reproduce the lowering check with
`/mnt/nvme/tmp/iter4lnc/compile_bq_lnc.py`, which compiles through
`TracerFrontend` (what `wrap_nki` defaults to) and opens no device.
Nothing calls this yet: `model._nki_entry` wraps at `wrap_nki(...)`, i.e.
grid 1, and this kernel asserts on a grid of 1 rather than silently running
half a GEMM. Wiring it means taking grid 2 **for this entry point only** —
`wrap_nki(fp4_linear_2core)[2]` — and leaving every unsplit kernel at
grid 1, because a kernel written for one core runs twice over the same data
at grid 2.
"""
m, k = x.shape
n = w.shape[0]
assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
programs = nl.num_programs(0)
core = nl.program_id(0)
assert programs == 2, (
f"launch grid is {programs}; fp4_linear_2core divides its output "
"columns between the two physical cores of a logical NeuronCore and "
"computes only its own half, so at grid 1 it would return a half-"
"empty result. Launch it as `wrap_nki(fp4_linear_2core)[2]`, or call "
"`fp4_linear` instead."
)
assert n % (programs * PT) == 0, (
f"n={n} does not divide into {programs} shares of a multiple of {PT}"
)
span = n // programs
lo = core * span
_check(k, span)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xq, True)
_dequant_fp4(w[lo : lo + span, 0 : k // 2], scale[lo : lo + span], wd, span, k)
_gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span, out_scale=FP4_POST)
return out
#: The low nibble of each uint16 lane of a uint32, i.e. two nibbles a lane-op.
FP4_NIB_MASK2 = 0x000F000F
#: `FP4_FIELD` doubled and signed: NKI takes a 32-bit bitwise operand as int32,
#: the same way `E4M3_SGN2` is written.
FP4_FIELD2 = 0x81C081C0 - (1 << 32)
#: PSUM holds 8 banks of 2 KiB a partition. An accumulation chain has to stay
#: in one place for the whole k loop, so the chunking below is sized to keep
#: `mtiles * (nc // nt)` live accumulators inside both of these.
_PSUM_TILES = 6
_PSUM_FP32 = 6 * 2048 // 4
def _fp4_plan(n, mtiles):
"""`(nt, nc)`: the PSUM tile width and the n-chunk the k loop runs over."""
nt = 1
for cand in range(1, NT + 1):
if n % cand == 0:
nt = cand
room = min(max(1, _PSUM_TILES // mtiles), max(1, _PSUM_FP32 // (mtiles * nt)))
best = 1
for cand in range(1, room + 1):
if n % (nt * cand) == 0:
best = cand
return nt, nt * best
def _stage_rows_fp4(x, m, k, dst_hbm):
"""`_stage_rows(..., quantize=True)` with the k axis de-interleaved."""
kq = k // 4
full = m // PT
tail = m - full * PT
for mb in nl.affine_range(full):
tile = _qdq_fp8(nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16), PT, k)
blocked = tile.reshape((PT, kq, 4))
half = nl.ndarray((PT, kq), dtype=nl.bfloat16, buffer=nl.sbuf)
for c in range(4):
nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kq, c])
nl.store(dst_hbm[mb * PT : (mb + 1) * PT, c, 0:kq], value=half)
if tail:
padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.memset(padded, 0.0)
nisa.tensor_copy(
dst=padded[0:tail, 0:k],
src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
)
tile = _qdq_fp8(padded, PT, k)
blocked = tile.reshape((PT, kq, 4))
half = nl.ndarray((PT, kq), dtype=nl.bfloat16, buffer=nl.sbuf)
for c in range(4):
nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kq, c])
nl.store(dst_hbm[full * PT : (full + 1) * PT, c, 0:kq], value=half)
def _fp4_decode_tile(wb, scb, dst, nc, c):
"""Nibble `c` of every packed pair in `wb[*, nc]`, dequantized into `dst`."""
nib = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
if c == 0:
nisa.tensor_scalar(
dst=nib.view(nl.uint32),
data=wb.view(nl.uint32),
op0=nl.bitwise_and,
operand0=FP4_NIB_MASK2,
engine=nisa.engine.vector,
)
else:
nisa.tensor_scalar(
dst=nib.view(nl.uint32),
data=wb.view(nl.uint32),
op0=nl.right_shift,
operand0=4 * c,
op1=nl.bitwise_and,
operand1=FP4_NIB_MASK2,
engine=nisa.engine.vector,
)
pre = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=pre, data=nib, op0=nl.multiply, operand0=FP4_NIB_MUL,
engine=nisa.engine.scalar,
)
bits = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=bits.view(nl.uint32),
data=pre.view(nl.uint32),
op0=nl.bitwise_and,
operand0=FP4_FIELD2,
engine=nisa.engine.vector,
)
nisa.tensor_tensor(
dst=dst,
data1=bits.view(nl.bfloat16),
data2=scb.view(nl.bfloat16),
op=nl.multiply,
)
def _fp4_scale_rows(s16, srep, n, k):
"""`srep[r, j] = s[j, r // 8]` — the e8m0 code of the 32-value block that
k-pair `r` belongs to, k-major, one byte a pair.
Invalid access of 8 partitions starting at partition 8
"""
npair = k // 64
tw = _ntile(n)
for tb in range(-(-npair // 64)):
p0 = tb * 64
pn = min(64, npair - p0)
sp = nl.ndarray((64, n), dtype=nl.uint16, buffer=nl.sbuf)
for q in range(n // tw):
nisa.dma_transpose(
dst=sp[0:pn, q * tw : (q + 1) * tw],
src=s16[q * tw : (q + 1) * tw, p0 : p0 + pn],
)
spu = sp.view(nl.uint8)
for h in range(2):
half = nl.ndarray((64, n), dtype=nl.uint8, buffer=nl.sbuf)
nisa.tensor_copy(dst=half[0:pn, 0:n], src=spu[0:pn, h : 2 * n : 2])
for i in range(8):
nl.store(srep[p0 : p0 + pn, h, i, 0:n], value=half[0:pn, 0:n])
def _gemm_fp4_fused(x4_hbm, w16, srep, dst_hbm, m, k, n):
"""`dst[m, n] = x[m, k] @ w[n, k].T` with the fp4 weight decoded in SBUF."""
kq = k // 4
rfull = kq // PT
rrem = kq - rfull * PT
rtiles = rfull
if rrem:
rtiles = rfull + 1
mtiles = m // PT
nt, nc = _fp4_plan(n, mtiles)
ntiles = nc // nt
at = nl.ndarray((PT, 4, rtiles, m), dtype=nl.bfloat16, buffer=nl.sbuf)
for mi in range(mtiles):
for c in range(4):
for rb in range(rtiles):
part = PT
if rb >= rfull:
part = rrem
nisa.dma_transpose(
dst=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
src=x4_hbm[
mi * PT : (mi + 1) * PT, c, rb * PT : rb * PT + part
],
)
for cb in nl.affine_range(n // nc):
accs = []
for mi in range(mtiles):
row = []
for nb in range(ntiles):
row.append(nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum))
accs.append(row)
for rb in range(rtiles):
part = PT
if rb >= rfull:
part = rrem
wb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
for q in range(ntiles):
lo = cb * nc + q * nt
nisa.dma_transpose(
dst=wb[0:part, q * nt : (q + 1) * nt],
src=w16[lo : lo + nt, rb * PT : rb * PT + part],
)
cod = nl.load(
srep[rb * PT : rb * PT + part, cb * nc : (cb + 1) * nc],
dtype=nl.uint8,
)
scb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=scb[0:part, 0:nc],
data=cod,
op0=nl.multiply,
operand0=128.0,
op1=nl.add,
operand1=float(FP4_EXP_BIAS * 128),
engine=nisa.engine.scalar,
)
for c in range(4):
dqf = nl.ndarray((PT, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
_fp4_decode_tile(wb, scb, dqf, nc, c)
for nb in range(ntiles):
for mi in range(mtiles):
nisa.nc_matmul(
dst=accs[mi][nb],
stationary=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
moving=dqf[0:part, nb * nt : (nb + 1) * nt],
accumulate=(rb > 0 or c > 0),
)
for mi in range(mtiles):
for nb in range(ntiles):
row = nl.ndarray((PT, nt), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=row, data=accs[mi][nb], op0=nl.multiply, operand0=FP4_POST
)
nl.store(
dst_hbm[
mi * PT : (mi + 1) * PT,
cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
],
value=row,
)
@nki.jit
def fp4_linear_fused(x, w, scale):
"""`fp4_linear` with the weight transposed as bytes and decoded in SBUF.
* `tpot = step / tokens_per_step`, and the divisor fell 3.6 → 2.818
because the needle case emitted 32 tokens where the record-holding run
emitted 19. The first five decode steps accepted 12 of 25 drafts — the
baseline's own rate — so this is EOS placement, not a worse drafter.
**To land this kernel the decode pass has to reach 324 ms**
(= 115.108 × 2.818), from 355.
* `ttft_short` is half a *fully traced* prefill pass, and this kernel
emits ~1 261 activity rows an execution where `fp4_linear` emits 2 230.
Freeing that much NRT trace budget makes **untouched** graphs get priced
more completely: in the same window `69d801c7`'s worst-traced execution
went 5 044 → 40 533 rows at an unchanged 15 executions and its measured
engine time rose 28%, while `544192c5`, whose coverage did not move,
did not move in time either (0.08%).
"""
m, k = x.shape
n = w.shape[0]
assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
assert k % 64 == 0, f"k={k} must be a multiple of 64 for the uint16 view"
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
x4 = nl.ndarray((mpad, 4, k // 4), dtype=nl.bfloat16, buffer=nl.private_hbm)
srep = nl.ndarray((k // 64, 2, 8, n), dtype=nl.uint8, buffer=nl.private_hbm)
_stage_rows_fp4(x, m, k, x4)
_fp4_scale_rows(scale.view(nl.uint16), srep, n, k)
_gemm_fp4_fused(
x4, w.view(nl.uint16), srep.reshape((k // 4, n)), out, mpad, k, n
)
return out
#: Select byte 1 of each uint16 lane, in place (signed, as `E4M3_SGN2` is).
FP8_HI_MASK2 = 0xFF00FF00 - (1 << 32)
#: Select byte 0 of each uint16 lane; a following `<< 8` moves it to byte 1,
#: which is where `_decode_e4m3`'s masks expect the e4m3 byte to sit.
FP8_LO_MASK2 = 0x00FF00FF
def _fp8_scale_rows(s16, srep, n, k, ns):
"""`srep[r, j] = scale[j // 32, r // 16]`, the e8m0 code of the `32 x 32`
block that k-pair `r` and output `j` belong to, k-major, one byte a pair.
"""
npair = (k // 32) // 2
tw = _ntile(ns)
for tb in range(-(-npair // 64)):
p0 = tb * 64
pn = min(64, npair - p0)
sp = nl.ndarray((64, ns), dtype=nl.uint16, buffer=nl.sbuf)
for q in range(ns // tw):
nisa.dma_transpose(
dst=sp[0:pn, q * tw : (q + 1) * tw],
src=s16[q * tw : (q + 1) * tw, p0 : p0 + pn],
)
spu = sp.view(nl.uint8)
dei = nl.ndarray((PT, ns), dtype=nl.uint8, buffer=nl.sbuf)
nisa.tensor_copy(dst=dei[0:pn, 0:ns], src=spu[0:pn, 0 : 2 * ns : 2])
nisa.tensor_copy(dst=dei[64 : 64 + pn, 0:ns], src=spu[0:pn, 1 : 2 * ns : 2])
for h in range(2):
rep = nl.ndarray((PT, ns, 32), dtype=nl.uint8, buffer=nl.sbuf)
nisa.tensor_copy(
dst=rep[0:pn, 0:ns, 0:32],
src=nl.broadcast_to(
dei[h * 64 : h * 64 + pn, 0:ns].reshape((pn, ns, 1)),
(pn, ns, 32),
),
)
flat = rep.reshape((PT, ns * 32))
for i in range(16):
nl.store(srep[p0 : p0 + pn, h, i, 0:n], value=flat[0:pn, 0:n])
def _stage_rows_fp8k(x, m, k, dst_hbm, quantize):
"""`_stage_rows` with the k axis split by `k % 2`."""
kh = k // 2
full = m // PT
tail = m - full * PT
for mb in nl.affine_range(full):
tile = nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16)
if quantize:
tile = _qdq_fp8(tile, PT, k)
blocked = tile.reshape((PT, kh, 2))
half = nl.ndarray((PT, kh), dtype=nl.bfloat16, buffer=nl.sbuf)
for c in range(2):
nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kh, c])
nl.store(dst_hbm[mb * PT : (mb + 1) * PT, c, 0:kh], value=half)
if tail:
padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.memset(padded, 0.0)
nisa.tensor_copy(
dst=padded[0:tail, 0:k],
src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
)
tile = _qdq_fp8(padded, PT, k) if quantize else padded
blocked = tile.reshape((PT, kh, 2))
half = nl.ndarray((PT, kh), dtype=nl.bfloat16, buffer=nl.sbuf)
for c in range(2):
nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kh, c])
nl.store(dst_hbm[full * PT : (full + 1) * PT, c, 0:kh], value=half)
def _fp8_decode_tile(wb, scb, dst, nc, c):
"""Byte `c` of every packed pair in `wb[*, nc]`, dequantized into `dst`."""
t = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
if c:
nisa.tensor_scalar(
dst=t.view(nl.uint32), data=wb.view(nl.uint32),
op0=nl.bitwise_and, operand0=FP8_HI_MASK2, engine=nisa.engine.vector,
)
else:
nisa.tensor_scalar(
dst=t.view(nl.uint32), data=wb.view(nl.uint32),
op0=nl.bitwise_and, operand0=FP8_LO_MASK2,
op1=nl.left_shift, operand1=8, engine=nisa.engine.vector,
)
hi = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=hi.view(nl.uint32), data=t.view(nl.uint32),
op0=nl.bitwise_and, operand0=E4M3_MAG2,
op1=nl.right_shift, operand1=4, engine=nisa.engine.vector,
)
bits = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.scalar_tensor_tensor(
dst=bits.view(nl.uint32), data=t.view(nl.uint32),
op0=nl.bitwise_and, operand0=E4M3_SGN2,
op1=nl.bitwise_or, operand1=hi.view(nl.uint32),
)
nisa.tensor_tensor(
dst=dst, data1=bits.view(nl.bfloat16), data2=scb.view(nl.bfloat16),
op=nl.multiply,
)
#: SBUF a partition may spend on the resident activation tiles of the fused fp8
#: GEMM (`2 * rtiles * m * 2` bytes). The MTP input projection, `k = 15360`, is
#: 60 k-pair tiles and so needs 30 KiB at `m = 128` and 120 KiB at `m = 512` —
_FP8_AT_BYTES = 64 * 1024
def _gemm_fp8_tfused(x2_hbm, w16, srep, dst_hbm, m, k, n):
"""`dst[m, n] = x[m, k] @ w[n, k].T` with the fp8 weight decoded in SBUF."""
kh = k // 2
rfull = kh // PT
rrem = kh - rfull * PT
rtiles = rfull
if rrem:
rtiles = rfull + 1
mtiles = m // PT
nt, nc = _fp4_plan(n, mtiles)
ntiles = nc // nt
assert 2 * rtiles * m * 2 <= _FP8_AT_BYTES, (
f"activation tiles are {2 * rtiles * m * 2} bytes a partition for "
f"k={k} m={m}; see _FP8_AT_BYTES"
)
at = nl.ndarray((PT, 2, rtiles, m), dtype=nl.bfloat16, buffer=nl.sbuf)
for mi in range(mtiles):
for c in range(2):
for rb in range(rtiles):
part = PT
if rb >= rfull:
part = rrem
nisa.dma_transpose(
dst=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
src=x2_hbm[mi * PT : (mi + 1) * PT, c, rb * PT : rb * PT + part],
)
for cb in nl.affine_range(n // nc):
accs = []
for mi in range(mtiles):
row = []
for nb in range(ntiles):
row.append(nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum))
accs.append(row)
for rb in range(rtiles):
part = PT
if rb >= rfull:
part = rrem
wb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
for q in range(ntiles):
lo = cb * nc + q * nt
nisa.dma_transpose(
dst=wb[0:part, q * nt : (q + 1) * nt],
src=w16[lo : lo + nt, rb * PT : rb * PT + part],
)
cod = nl.load(
srep[rb * PT : rb * PT + part, cb * nc : (cb + 1) * nc],
dtype=nl.uint8,
)
scb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=scb[0:part, 0:nc], data=cod,
op0=nl.multiply, operand0=128.0,
op1=nl.add, operand1=float(SC_EXP_ADD * 128),
engine=nisa.engine.scalar,
)
for c in range(2):
dqf = nl.ndarray((PT, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
_fp8_decode_tile(wb, scb, dqf, nc, c)
for nb in range(ntiles):
for mi in range(mtiles):
nisa.nc_matmul(
dst=accs[mi][nb],
stationary=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
moving=dqf[0:part, nb * nt : (nb + 1) * nt],
accumulate=(rb > 0 or c > 0),
)
for mi in range(mtiles):
for nb in range(ntiles):
nl.store(
dst_hbm[
mi * PT : (mi + 1) * PT,
cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
],
value=nl.copy(accs[mi][nb], dtype=nl.bfloat16),
)
@nki.jit
def fp8_linear_tfused(x, w, scale, quantize=True):
"""`fp8_linear` with the weight transposed as bytes and decoded in SBUF."""
m, k = x.shape
n = w.shape[0]
assert w.shape[1] == k, f"fp8 weight {w.shape} is not [n, k] for k={k}"
assert k % 64 == 0, f"k={k} must be a multiple of 64"
assert n % BLK == 0, f"n={n} must be a multiple of {BLK}"
_check(k, n)
mpad = -(-m // PT) * PT
ns = -(-n // BLK)
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
x2 = nl.ndarray((mpad, 2, k // 2), dtype=nl.bfloat16, buffer=nl.private_hbm)
srep = nl.ndarray((k // 64, 2, 16, n), dtype=nl.uint8, buffer=nl.private_hbm)
_stage_rows_fp8k(x, m, k, x2, quantize)
_fp8_scale_rows(scale.view(nl.uint16), srep, n, k, ns)
_gemm_fp8_tfused(
x2, w.view(nl.uint16), srep.reshape((k // 2, n)), out, mpad, k, n
)
return out
@nki.jit
def fp32w_linear(x, w):
"""`x[m, k] @ w[n, k].T` where `w` is stored **fp32**, narrowed in-kernel."""
m, k = x.shape
n = w.shape[0]
_check(k, n)
mpad = -(-m // PT) * PT
npad = -(-n // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
ws = nl.ndarray((npad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xs, False)
_stage_rows(w, n, k, ws, False)
_gemm(xs, ws, out, mpad, k, n)
return out
@nki.jit
def bf16_linear(x, w):
"""`x[m, k] @ w[n, k].T`, both bf16, no quantization anywhere."""
m, k = x.shape
n = w.shape[0]
_check(k, n)
mpad = -(-m // PT) * PT
out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
_stage_rows(x, m, k, xs, False)
_gemm(xs, w, out, mpad, k, n)
return out
Reference in New Issue View Git Blame Copy Permalink