diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/README.md b/examples/jit_cpp/fused_hadamard_quant_a5/README.md new file mode 100644 index 00000000..29ee6228 --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/README.md @@ -0,0 +1,155 @@ +# fused_hadamard_quant_a5 - an order-K Hadamard and MXFP4 in one launch + +`x -> order-K Hadamard -> E2M1 nibbles + one E8M0 scale per 32`, as a +single kernel on the Ascend 950 / A5 (`dav-c310-vec`) vector core, JIT-compiled +with `bisheng` and loaded through `ctypes`. `K` is a template parameter over 26 +widths; one `.so` holds an instantiation per width and the launcher dispatches on +it, so there is no rebuild per size. + +`fused_hadamard_quant_b32_a5` is the companion that rotates independent +32-element blocks instead. Prefer that one for a `K` that is not a power of two; +prefer this one when the method calls for a rotation across the whole row. + +The rotation is row wide: every output element depends on all `K` inputs, which +is why `K` must be a power of two. + +The transform runs in two phases, because Sylvester factors as +`H_K = H_(K/256) (x) H_256`. Phase 1 does the order-256 transform inside every +256-element window, each window an independent deinterleave-load, add/sub, +concat-halves-store repeated eight times. Phase 2 pairs windows `(a, a|t)` +elementwise for `log2(K/256)` further stages. Windows are independent in phase 1 +and window pairs are independent in phase 2, so no register holds more than one +window and the row width is not capped -- a single-phase butterfly keeping a +whole row in registers stops at 4096, where a row is already 16 chunks against +16 register slots. The stage count is the same either way: +`8 + log2(K/256) = log2(K)`. + +The MXFP4 group stays 32 and no longer lines up with the rotation, which costs +nothing: the quantizer takes the rotated tile in 32-element blocks whatever +produced it. + +## Fusing the pair is 2.45x the two separate launches + +Unfused, this is two passes over HBM: the butterfly writes the rotated tile out +and the quantizer reads it straight back. Fused, that tile never leaves UB and +only the nibbles and scales are written. Bytes per element tell the whole story: +6.53 unfused against 2.53 fused. + +| K | 2 launches | fused | vs 2 | rel err | spread | +|---|--:|--:|--:|--:|--:| +| 1024 | 38.8 | 28.4 | 1.37x | 0.0 | 17.7% | +| 4096 | 297.7 | 121.7 | **2.45x** | 0.0 | 2.8% | +| 8192 | 612.3 | 247.4 | **2.47x** | 0.0 | 2.2% | +| 16384 | 1209.9 | 493.8 | **2.45x** | 0.0 | 1.0% | + +M = 16384, microseconds per launch, what `benchmark.py` prints. Both arms agree +to a relative error of 0.0, checked before either is timed. Byte traffic +predicts 6.53 / 2.53 = 2.58x, and the three clean widths measure 2.45-2.47x. + +K=1024 is lower for a reason worth knowing rather than hiding. The unfused +intermediate is `2*M*k` bytes, so at K=1024 it is 32 MB against a 128 MiB L2 and +the unfused arm reads much of it from cache rather than HBM -- which flatters the +arm fusing is measured against and understates the result. Its 17.7% bracket +spread is the same thing showing up as noise. +The 19.4% bracket spread on that row is the same thing showing up as noise. +K=4096 and above is where the comparison is clean, and that is where the +prediction and the measurement meet. + +## It runs at about the speed of a copy of its input + +| K | fused | d2d copy | vs copy | fused GB/s | copy GB/s | +|---|--:|--:|--:|--:|--:| +| 1024 | 121.5 | 191.5 | **1.58x** | 1398 | 1402 | +| 4096 | 122.9 | 191.7 | **1.56x** | 1382 | 1400 | +| 8192 | 122.5 | 192.3 | **1.57x** | 1386 | 1396 | +| 16384 | 122.5 | 189.3 | **1.55x** | 1387 | 1418 | + +64Mi elements per launch. The fused column is flat -- 121.5 to 122.9 us across a +16x range of row width -- because the transform is entirely hidden under the DMA +at every width. Both arms reach much the same bandwidth, the kernel 1382-1398 +GB/s against the copy's 1396-1418, so the kernel is not moving bytes faster than +a copy; it is moving 1.58x fewer of them, 2.53 B/element against 4.0. + +Getting there took two changes, and their sizes are worth recording. The +cross-window stages were originally addressed by a shift-and-OR index computed +per register slot inside the unrolled fold, which cost 266 us of a 388 us kernel +at K=16384; walking nested loops over `base + m*step` instead, with the same +memory pattern and the same number of passes, cut that phase to about 23 us. +Fusing the passes on top (`FUSED_CROSS_FUSE`) added a further 1.16x. Set it to 1 +to get one stage per pass and measure the difference. + +Measured on an `Ascend950PR_9589`: 64 vector cores, 128 MiB L2, 1.65 GHz, HBM +peak 1.6 TB/s, so the kernel reaches 86-87% of peak. The copy is a +reference for what moving the bytes costs, not a proven lower bound -- +it is a vendor kernel doing a simpler job. HBM peak is the closer thing +to a real ceiling, and that is the number above. Other A5 parts have +different HBM, and absolute GB/s from one part should not be compared against +another -- the ratios above are the portable numbers. + +## Correctness + +The kernel cannot be bit-exact against a torch expression: it rotates in bf16 +with a specific operand order and no torch formulation reproduces that tree. So +`test_fused_hadamard_quant_a5.py` establishes it three ways, strongest first. + +1. **Scale bytes** must match a reference that rotates in fp32 and quantizes with + `torch_npu`. A scale is a power of two derived from a block maximum, so bf16 + rounding inside the butterfly almost never moves it -- disagreeing scales mean + a wrong rotation, not different rounding. Threshold 98%; measured 99.8%. +2. **Dequantized values** must track that reference to within MXFP4's own + resolution. This catches a correct-looking permutation, which a check on the + packed bytes would not. Threshold 5%; measured 0.36%. +3. **The output must be non-trivial.** A kernel that writes nothing, or writes + its input back, is the characteristic silent failure on this hardware and + would pass a loose tolerance. Separate tests assert the nibbles are neither + all-zero nor degenerate, and that the result differs from quantizing without + the rotation. + +Two structural cases are covered because both have hidden real bugs here. The +width list spans both **unroll classes** -- the butterfly unrolls by 8 or by 4 +depending on `rows_for(k) * k / 256`, they are different code paths, and class +membership is derived from the built `.so` rather than hardcoded, because raising +the tile size once moved four widths between classes and left the matrix +single-class. And one test uses a batch deep enough that every core walks +several tiles, since a shallow batch leaves the buffer rotation, the prefetch and +the drain unexercised. + +```bash +python3 -m pytest -q test_fused_hadamard_quant_a5.py +``` + +## Running the benchmark + +```bash +./run_benchmark.sh # or: python3 benchmark.py --device 0 +``` + +Needs a CANN whose PTO carries MXFP4 (`Exp2DStrided` in `pto/npu/a5/TQuant.hpp`). +9.1.0 and 9.2.0 both do; 9.0.0 does not. + +## Tunables + +All compile-time, with the shipped defaults. Every combination is checked by +`static_assert`, so a tile that will not fit UB or a prefetch depth that would +deadlock fails to compile rather than misbehaving. + +| flag | default | what it is | +|---|---|---| +| `FUSED_TILE_ELEMS` | 24576 | elements per UB tile, 48 KB in bf16 | +| `FUSED_BUFFERS` | 3 | UB pipeline buffers | +| `FUSED_PREFETCH` | 2 | tiles in flight ahead | + +The same source builds the reduced kernels the ladder needs: +`FUSED_ROTATE_ONLY` drops the quantizer and `FUSED_NO_ROTATE` drops the +butterfly, so the two arms differ in what they fuse and in nothing else. +`FUSED_BUFFERS` above 4 does not build at K=4096 -- five slots need 311,040 +bytes of UB against 253,952 available. + +The butterfly is the unnormalised Sylvester matrix, so its output is `sqrt(K)` +larger than an orthogonal Hadamard's, and that factor is left to the caller. +Scale `x` by `1/sqrt(K)` going in if orthogonal semantics are wanted. Whether +`E8M0` could have absorbed it instead depends on the width: at `K` = 64, 256, +1024 or 4096 the factor is 8, 16, 32 or 64 and is itself a power of two, but at +32, 128, 512 or 2048 it is not, so the scale cannot take it and the nibbles +would genuinely differ. Leaving it out is the one behaviour that holds at every +supported width. diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/benchmark.py b/examples/jit_cpp/fused_hadamard_quant_a5/benchmark.py new file mode 100644 index 00000000..2964c643 --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/benchmark.py @@ -0,0 +1,214 @@ +"""Fused block-32 Hadamard + MXFP4 quantize on Ascend A5, against the unfused +pair and against a device-to-device copy. + +Two comparisons: + + A. the fusion ladder. Unfused is two launches -- the Hadamard, then the + quantizer -- at 6.53 B/element, against 2.53 fused. Both arms are this + kernel, built from one source with the unwanted half compiled out, so they + differ in what they fuse and in nothing else. + + B. the fused kernel against a d2d copy of its input, as a reference for what + moving the bytes costs. Not a proven lower bound: the copy is a vendor + kernel doing a simpler job, and nothing here shows it is optimal. + +Method: wall clock on a saturated queue, medians over TRIALS brackets of +LAUNCHES launches, inputs drawn from a rotating pool so a bracket cannot be +served from cache. Every arm is checked against the unfused arm before any of +them is timed. +""" + +import argparse +import statistics +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch_npu # noqa + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from jit_util_fused_a5 import ( # noqa + MX_BLOCK, + build_and_load, +) + +# At M=16384 the unfused intermediates are 2*M*k bytes. Below K=4096 that fits +# the 128 MiB L2, so the unfused arms partly read from cache and the ladder +# understates fusing -- 2.1x at K=1024 against 4.1x at K=4096. Kept in the sweep +# because the effect is worth seeing, not because those rows are the headline. +SHAPES = (1024, 4096, 8192, 16384) +COPY_ELEMS = 1 << 26 +M = 16384 +TRIALS = 15 +LAUNCHES = 20 +WARMUP = 5 +POOL_BYTES = 256 * 1024 * 1024 +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) + + +def trials(call, depth, launches=LAUNCHES): + for _ in range(WARMUP): + call(0) + torch.npu.synchronize() + out = [] + for t in range(TRIALS): + torch.npu.synchronize() + t0 = time.perf_counter() + for i in range(launches): + call((t * launches + i) % depth) + torch.npu.synchronize() + out.append((time.perf_counter() - t0) * 1e6 / launches) + med = statistics.median(out) + return med, 100 * (max(out) - min(out)) / med + + +def dequant(q, s, k): + q = q.cpu() + lo, hi = q & 0x0F, (q >> 4) & 0x0F + codes = torch.stack([lo, hi], dim=-1).reshape(q.shape[0], -1) + mag = E2M1[(codes & 0x07).long()] + sign = torch.where(codes & 0x08 != 0, -1.0, 1.0) + scale = torch.exp2(s.cpu().float() - 127.0).repeat_interleave(MX_BLOCK, dim=-1) + return (mag * sign * scale).reshape(-1, k) + + +def bench_ladder(k): + """A: two launches, then one.""" + depth = max(2, min(16, POOL_BYTES // max(M * k * 4, 1))) + x = [torch.randn(M, k, dtype=torch.bfloat16, device="npu") for _ in range(depth)] + + # one source, three builds: both halves, the rotation alone, the quantizer + # alone. Same tiling, same UB layout, same buffer count in each, so the + # difference between the arms is the fusion and nothing else. + fused = build_and_load(k=k, verbose=False) + rotate_only = build_and_load( + k=k, verbose=False, extra_defs=("-DFUSED_ROTATE_ONLY",) + ) + quant = build_and_load(k=k, verbose=False, extra_defs=("-DFUSED_NO_ROTATE",)) + + q = torch.empty((M, k // 2), dtype=torch.uint8, device="npu") + s = torch.empty((M, k // MX_BLOCK), dtype=torch.uint8, device="npu") + rot = torch.empty((M, k), dtype=torch.bfloat16, device="npu") + torch.npu.synchronize() + + def two(i): # Hadamard, then quantize + rotate_only(x[i % depth], out=(rot.view(torch.uint8), s)) + quant(rot, out=(q, s)) + + def one(i): # both in one launch + fused(x[i % depth], out=(q, s)) + + # correctness gate: the arms must agree before either is timed + two(0) + torch.npu.synchronize() + ref = dequant(q.clone(), s.clone(), k) + one(0) + torch.npu.synchronize() + got = dequant(q, s, k) + rel = ((got - ref).abs().mean() / ref.abs().mean().clamp_min(1e-6)).item() + + # A disagreeing arm is a bug, not a datum: stop rather than print a table + # whose rows measure different computations. bf16 rounding differences + # between a fused and an unfused rotation land near 1e-3, not near 1. + if rel > 0.05: + raise SystemExit( + f"K={k}: arms disagree, rel={rel:.4f} -- the ladder is not measuring " + "the same computation in every arm" + ) + + t2, s2 = trials(two, depth) + t1, s1 = trials(one, depth) + x.clear() + torch.npu.empty_cache() + return { + "k": k, + "two_us": round(t2, 1), + "fused_us": round(t1, 1), + "vs_two": round(t2 / t1, 2), + "rel": round(rel, 5), + "spread_pct": round(max(s2, s1), 1), + } + + +def bench_copy(k): + """B: the fused kernel against a copy of its input.""" + batch = max(128, (COPY_ELEMS // k) // 128 * 128) + depth = max(2, min(16, POOL_BYTES // max(batch * k * 4, 1))) + x = [ + torch.randn(batch, k, dtype=torch.bfloat16, device="npu") for _ in range(depth) + ] + dst = [torch.empty_like(x[0]) for _ in range(depth)] + fused = build_and_load(k=k, verbose=False) + q = torch.empty((batch, k // 2), dtype=torch.uint8, device="npu") + s = torch.empty((batch, k // MX_BLOCK), dtype=torch.uint8, device="npu") + torch.npu.synchronize() + + tf, sf = trials(lambda i: fused(x[i % depth], out=(q, s)), depth) + tc, sc = trials(lambda i: dst[i % depth].copy_(x[i % depth]), depth) + # the kernel reads 2 B and writes 0.5 + 1/32 per element; the copy 2 and 2 + kernel_gbs = batch * k * (2 + 0.5 + 1 / MX_BLOCK) / (tf * 1e-6) / 1e9 + copy_gbs = batch * k * 4.0 / (tc * 1e-6) / 1e9 + x.clear() + dst.clear() + torch.npu.empty_cache() + return { + "k": k, + "batch": batch, + "fused_us": round(tf, 1), + "copy_us": round(tc, 1), + "vs_copy": round(tc / tf, 2), + "fused_gbs": round(kernel_gbs), + "copy_gbs": round(copy_gbs), + "spread_pct": round(max(sf, sc), 1), + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--device", type=int, default=0) + args = ap.parse_args() + torch.npu.set_device(args.device) + torch.manual_seed(20260826) + np.random.seed(0) + + print(f"=== A. fusion ladder, M={M} (microseconds per launch) ===") + print( + f"{'K':>7} {'2 launches':>11} {'fused':>8} {'vs 2':>6} {'rel':>8} " + f"{'spread':>7}" + ) + for k in SHAPES: + try: + r = bench_ladder(k) + except (RuntimeError, SystemExit) as exc: + print(f"{k:>7} skipped: {str(exc)[:56]}") + continue + print( + f"{r['k']:>7} {r['two_us']:>11.1f} {r['fused_us']:>8.1f} " + f"{r['vs_two']:>5.2f}x {r['rel']:>8.4f} {r['spread_pct']:>6.1f}%" + ) + + print(f"\n=== B. fused vs a d2d copy ({COPY_ELEMS // 1024}Ki elements) ===") + print( + f"{'K':>7} {'batch':>8} {'fused':>8} {'copy':>8} {'vs copy':>8} " + f"{'fused GB/s':>11} {'copy GB/s':>10} {'spread':>7}" + ) + for k in SHAPES: + try: + r = bench_copy(k) + except RuntimeError as exc: + print(f"{k:>7} skipped: {str(exc)[:56]}") + continue + print( + f"{r['k']:>7} {r['batch']:>8} {r['fused_us']:>8.1f} {r['copy_us']:>8.1f} " + f"{r['vs_copy']:>7.2f}x {r['fused_gbs']:>11} {r['copy_gbs']:>10} " + f"{r['spread_pct']:>6.1f}%" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/fused_hadamard_quant_a5.cpp b/examples/jit_cpp/fused_hadamard_quant_a5/fused_hadamard_quant_a5.cpp new file mode 100644 index 00000000..c3571d4a --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/fused_hadamard_quant_a5.cpp @@ -0,0 +1,1052 @@ +// Block-32 Hadamard fused with MXFP4 quantization, one launch. +// +// x -> (x @ H) -> E2M1 nibbles + one E8M0 scale per 32 +// +// Unfused this is two passes over HBM: read x / write rotated, then read +// rotated / write nibbles+scales. Fused it is read x / write nibbles+scales, so +// on a DMA-bound op the saving is close to the whole second pass. +// +// Built from two kernels that are already measured and merged upstream: +// fast_hadamard_a5 supplies the butterfly, mxfp4_quant_a5 the four quant +// passes, the tiling and the outputs. Both are left doing what they already do; +// what is new here is that the rotated tile never leaves UB. +// +// The butterfly was fp16 upstream and is bf16 here, which costs nothing +// structurally: vlds/vsts are bit-width ops on vector_u16 (DINTLV_B16 / +// NORM_B16), so only the arithmetic type changes, by reference cast. That is +// the same idiom mxfp4_quant_a5 already uses for its max reduction. +// +// The difference from v1 is not just a pinned width. v1 rotates a whole row, so +// K must be a power of two and at most 2048. Here the Hadamard is always 32 +// wide and a row is a sequence of independent 32-blocks, which decouples the +// rotation from the row width: K goes back to any multiple of 32, so 4096 and +// 11008 work, which a row-wide rotation rejects. +// +// It falls out of the tile being a flat run of Rows*K elements. Blocks are +// contiguous and 32 long, the butterfly window is 256 = eight blocks, and +// blocks are independent -- so one window covers eight of them and never has to +// care whether they came from the same row. +// +// Also: the MXFP4 group is 32 and the Hadamard block is 32, so a scale covers +// exactly one rotated block. No reshaping, and no group straddling a rotation. +#include +#include +#include +using namespace pto; + +// Row widths with an instantiation. The full set the quantizer supports: a +constexpr unsigned SUPPORTED_K[] = {32, 64, 128, 256, 512, + 1024, 2048, 4096, 8192, 16384}; +constexpr unsigned SUPPORTED_COUNT = + sizeof(SUPPORTED_K) / sizeof(SUPPORTED_K[0]); +// --- butterfly geometry, from fast_hadamard_a5 ------------------------------- +// WINDOW is two registers: the deinterleave load splits a 2*lanes run into +// even/odd halves, and the concat-halves store puts them back. +constexpr unsigned SLOTS = 8; // unroll width: register sets per sweep +constexpr unsigned HAD_ALIGN = 512; + +constexpr unsigned MX_BLOCK = 32; // MXFP4 block: 32 elements, one E8M0 scale + +// The three pipeline parameters, overridable for tuning. The defaults are the +// tuned point: 24576 is the largest tile that fits UB at all, and 3 is the only +// buffer count it fits at. It beats the 16384/4/2 this kernel shipped with by +// 1.063-1.066x on large launches, bit-exact, with no regression at any shape +// measured. Numbers and the full grid are in the README. +// +// Overriding is safe in the way that matters: every combination is checked by +// the static_asserts at the end of QuantShape, so a tile that will not fit UB, +// or a prefetch depth that would deadlock, fails to COMPILE rather than +// misbehaving. And the host reads rows-per-tile back from the .so +// (hadamard_mxfp4_full_rows_for), so a changed TILE_ELEMS cannot desynchronise +// from the harness. +#ifndef FUSED_BUFFERS +#define FUSED_BUFFERS 3 +#endif +#ifndef FUSED_PREFETCH +#define FUSED_PREFETCH 2 +#endif +#ifndef FUSED_TILE_ELEMS +#define FUSED_TILE_ELEMS 24576 // 48 KB bf16 +#endif + +// Store the butterfly halves with vscatter instead of vsts. +// +// 0 vsts NORM_B16 pair -- what ships, and the default. +// 1 vscatter with an IDENTITY index. Same instruction count, same registers, +// same dependency chain, bit-identical output; the ONLY difference is the +// opcode. This exists to price vscatter against vsts, because vscatter's +// cost on A5 is not documented and the vendor's bf16 TTRANS is built from +// it -- and two TTRANS calls are 96% of the pure-PTO kernel's cost, so it +// could be far dearer than a plain store. +// +// Mode 2 was a ROL5 index meant to absorb the rotation fixup into the store at +// no extra instruction -- the fixup is 13.8% of the kernel (82.74 -> 71.31 us +// with -DFUSED_NO_ROTFIX, paired 1.164x, resolved). It is GONE, refuted by +// mode 1: vscatter costs +28.96 us per call against vsts (111.57 against 82.61, +// paired 0.741x, resolved), so the opcode swap alone costs 2.5x what the fixup +// it would remove is worth. Break-even needed vscatter under ~5.5x a vsts. +// +// This also explains the pure-PTO kernel: its two TTRANS calls are 96% of its +// 2052 us, and the vendor builds bf16 TTRANS out of vgather2/vscatter. + +#ifndef FUSED_SCATTER +#define FUSED_SCATTER 0 +#endif +#if FUSED_SCATTER > 1 +#error "FUSED_SCATTER=2 (ROL5 store) was measured and refuted; see above" +#endif + +constexpr unsigned DEF_BUFFERS = FUSED_BUFFERS; // UB pipeline buffers +constexpr unsigned DEF_PREFETCH = FUSED_PREFETCH; // tiles in flight ahead +constexpr unsigned TILE_ELEMS = FUSED_TILE_ELEMS; +// RULE: every GM move_tile is one row and a Tile refuses a row under 32 bytes. +// The scale row is the smallest, at tile_elems/32 bytes, so a tile must be a +// whole multiple of 32*MX_BLOCK elements. DMA sets this grain, not the compute. +constexpr unsigned TILE_GRAIN = 1024; +template +struct Gcd { + static constexpr unsigned value = Gcd::value; +}; +template +struct Gcd { + static constexpr unsigned value = A; +}; + +// ROWS_PER_TILE: the largest row count whose tile is a whole number of grains. +// Not TILE_ELEMS / K -- for a large odd factor (768 = 32*24) the quotient is +// not a multiple of the grain. Zero means inadmissible; Rows asserts on it. +// +// Rows*K is a multiple of TILE_GRAIN exactly when Rows is a multiple of +// TILE_GRAIN / gcd(K, TILE_GRAIN), so the answer is the largest such multiple +// within cap. Counting down from cap one at a time instead costs a template +// instantiation per step: at TILE_ELEMS 32768 and K=32 that is 1024 of them, +// which is the compiler's default depth limit, and the tile could not be raised +// without hitting it. This form has no such ceiling. +template +struct RowsFor { + static constexpr unsigned cap = TILE_ELEMS / K > 1u ? TILE_ELEMS / K : 1u; + static constexpr unsigned step = TILE_GRAIN / Gcd::value; + static constexpr unsigned value = (cap / step) * step; +}; + +#if defined(FUSED_ROTATE_ONLY) && defined(MXFP4_TQUANT) +#error "FUSED_ROTATE_ONLY does nothing in a TQuant build: TQuant owns them" +#endif + +#ifdef __CCE_AICORE__ +constexpr unsigned B16_LANES = 128; // bf16 lanes in one vector register +// vcgmax on b16 groups 16 lanes, 8 results in lanes 0..7. +constexpr unsigned VCGMAX_B16_GROUP = 16; +constexpr unsigned VCGMAX_B16_RESULTS = B16_LANES / VCGMAX_B16_GROUP; +static_assert(VCGMAX_B16_RESULTS == 8, "block_abs_max stores with PAT_VL8"); +// RULE: vsts needs a 32-byte-aligned UB address, else 507035. Tile refuses a +// sub-32-byte DMA, so the padding is squeezed out in UB, not on the way to GM. +constexpr unsigned VSTS_ALIGN = 32; +constexpr unsigned GROUP_PITCH_B16 = VSTS_ALIGN / 2u; // in b16 elements +// RULE: vselr indices reach only the low 128 source bytes: 4 groups per gather. +constexpr unsigned GROUPS_PER_COMPACT = 4; +constexpr unsigned EVENT_SLOTS = 8; +static_assert(EVENT_SLOTS == 8, "extend buffer_free's initialiser first"); +constexpr unsigned UB_ALIGN = 512; +// A5 has 256 KB. +constexpr unsigned UB_BYTES = PTO_UBUF_SIZE_BYTES; + +// bf16 bit-field constants. bf16 is 1-8-7, so a magnitude's biased exponent is +// simply bits >> 7 once the sign is cleared. +constexpr uint16_t BF16_ABS = 0x7FFFu; // clears the sign bit +constexpr int16_t BF16_MANT_BITS = 7; +constexpr int16_t E8M0_BIAS_ADJ = -2; // byte = b - 2 (Algorithm 1, FLOOR) +constexpr int16_t RECIP_OFFSET = 256; // 1/X exponent field = 256 - b +// b must stay in a window where 1/X is finite, non-subnormal bf16: field 256-b +// must land in [2, 254]. Clamp b, then derive BOTH outputs from the clamped b. +constexpr int16_t B_MIN = 2; +constexpr int16_t B_MAX = 254; + +// RULE: a constexpr function cannot be called from [aicore] code. +template +struct Log2 { + static constexpr unsigned value = 1 + Log2<(Value >> 1)>::value; +}; +template <> +struct Log2<1> { + static constexpr unsigned value = 0; +}; + +// Largest unroll width <= Limit that divides Windows exactly, so the sweep +// covers the tile with no tail. A value template rather than a constexpr +// function: [aicore] code may not call one, even to initialise a constexpr. +template +struct UnrollFor { + static constexpr unsigned value = + (Windows % Limit == 0) ? Limit : UnrollFor::value; +}; +template +struct UnrollFor { + static constexpr unsigned value = 1u; +}; + +template +struct RoundUp { + static constexpr unsigned value = (Bytes + UB_ALIGN - 1) & ~(UB_ALIGN - 1); +}; + +// Every derived size for one instantiation. +template +struct QuantShape { + static constexpr unsigned tile_elems = Rows * K; + static constexpr unsigned row_elems = K; + static constexpr unsigned rows_in_tile = Rows; + static constexpr unsigned blocks = tile_elems / MX_BLOCK; + static constexpr unsigned in_bytes = tile_elems * 2u; + static constexpr unsigned q_bytes = tile_elems / 2u; + static constexpr unsigned scale_bytes = blocks; + + // One "group" is one vcgmax: 8 blocks == 256 elements. + static constexpr unsigned groups = blocks / VCGMAX_B16_RESULTS; + // Round UP: the last bite may be partial. Safe only because the buffers + // below are sized from these counts, and nothing reads what a partial bite + // writes past `blocks`. The asserts at the end of this struct pin both. + static constexpr unsigned compact_iters = + (groups + GROUPS_PER_COMPACT - 1u) / GROUPS_PER_COMPACT; + static constexpr unsigned b_iters = (blocks + B16_LANES - 1u) / B16_LANES; + static constexpr unsigned c_iters = tile_elems / (2u * B16_LANES); + + // maxima_bytes carries a register of read-ahead: the gather reads one + // register past its last input offset. The rest are sized from what the loops + // WRITE, since the two rounded counts can exceed their data by one bite. + static constexpr unsigned maxima_bytes = + compact_iters * GROUPS_PER_COMPACT * GROUP_PITCH_B16 * 2u + B16_LANES; + static constexpr unsigned packed_bytes = + compact_iters * GROUPS_PER_COMPACT * VCGMAX_B16_RESULTS * 2u; + static constexpr unsigned aligned_in = RoundUp::value; + static constexpr unsigned aligned_q = RoundUp::value; + static constexpr unsigned aligned_s = RoundUp::value; + static constexpr unsigned aligned_max = RoundUp::value; + static constexpr unsigned aligned_packed = RoundUp::value; + static constexpr unsigned aligned_mult = + RoundUp::value; + + static constexpr unsigned slot_stride = aligned_in + aligned_q + aligned_s; + static constexpr unsigned scratch_base = NBuffers * slot_stride; + // every scratch region SlotOffset hands out, in the same order, so the two + // cannot drift: omitting one here silently shrinks the UB-overflow guard + static constexpr unsigned ub_needed = + scratch_base + aligned_max + aligned_packed + aligned_mult; + +#ifdef MXFP4_TQUANT + // TQuant reads its per-block maxima a whole register at a time, so the span + // rounds up to b_iters registers, and its reducer flushes one 32-byte block + // past the last group. + static constexpr unsigned tquant_max_elems = b_iters * B16_LANES; + static constexpr unsigned tquant_max_bytes = + tquant_max_elems * 2u + VSTS_ALIGN; + static constexpr unsigned tquant_scaling_bytes = blocks * 2u; + static_assert(tquant_max_bytes <= aligned_max, + "TQuant maxima do not fit the maxima region"); + static_assert(tquant_scaling_bytes <= aligned_mult, + "TQuant scaling does not fit the reciprocal region"); + // numGroups truncates inside TQuant, and a partial 8-group window takes a + // different store path. + static_assert(tile_elems % MX_BLOCK == 0, "TQuant would drop a group"); + static_assert(blocks % 8u == 0, "TQuant would take its vstus tail path"); +#endif + + // --- butterfly geometry, in two phases + // ------------------------------ The rotation is order K, and Sylvester + // factors as H_K = H_(K/256) (x) H_256. So a row is treated as K/256 windows + // of 256: phase 1 runs the order-256 transform inside every window, phase 2 + // pairs windows (a, a|t) elementwise for log2(K/256) stages. Both phases are + // made of independent pieces -- windows in phase 1, window pairs in phase 2 + // -- so neither has a width limit. Holding a whole row in registers instead + // stops at K=4096, where the row is already 16 chunks against 16 slots. + // + // The stage count is the same either way: 8 + log2(K/256) = log2(K). + static constexpr unsigned had_lanes_b16 = B16_LANES; + static constexpr unsigned had_window = 2u * had_lanes_b16; + static constexpr unsigned log2_k = Log2::value; + static constexpr unsigned log2_window = Log2::value; + // A window is always 256 elements: one row segment when the row is wider, + // several whole rows packed together when it is narrower. + static constexpr unsigned rows_per_window = + K < had_window ? had_window / K : 1u; + static constexpr unsigned had_group = had_window; + // Only a row narrower than a window leaves a packing tail to undo. + static constexpr unsigned rotations = + log2_window - (log2_k < log2_window ? log2_k : log2_window); + static constexpr unsigned phase1_stages = + log2_k < log2_window ? log2_k : log2_window; + static constexpr unsigned windows_per_row = + K < had_window ? 1u : K / had_window; + static constexpr unsigned phase2_stages = log2_k - phase1_stages; + static constexpr unsigned upper = had_group / 2u; + static constexpr unsigned lanes = + upper < had_lanes_b16 ? upper : had_lanes_b16; + static constexpr unsigned chunks = upper / lanes; + static constexpr unsigned windows_per_tile = tile_elems / had_group; + static constexpr unsigned groups_per_iter = + UnrollFor::value; + static constexpr unsigned had_iters = + tile_elems / had_group / groups_per_iter; + static constexpr unsigned sweep_stride = groups_per_iter * had_group; + // How many register slots one sweep call must use. This has to be derived + // from groups_per_iter, NOT from SLOTS. + // + // A slot addresses window `Slot / chunks` at chunk `Slot % chunks`, so a call + // with N slots covers N/chunks windows, while the loop advances + // groups_per_iter windows per iteration. Instantiating the sweep with SLOTS + // when groups_per_iter is smaller makes consecutive iterations overlap and + // runs the last one off the end of the tile. + // + // groups_per_iter = SLOTS / chunks makes these agree by construction. + static constexpr unsigned sweep_slots = groups_per_iter * chunks; + + static_assert(K >= MX_BLOCK && !(K & (K - 1u)), + "an order-K butterfly needs K a power of two, at least one " + "MXFP4 block wide"); + static_assert(chunks == 1u, "a 256-wide window is exactly two registers"); + static_assert(phase1_stages + phase2_stages == log2_k, + "the two phases must add up to the full transform"); + static_assert(windows_per_row * had_window == K || K < had_window, + "a row must be a whole number of windows"); + static_assert(tile_elems % had_group == 0, + "tile must be a whole number of butterfly windows"); + static_assert(windows_per_tile % groups_per_iter == 0, + "UnrollFor must divide the window count exactly"); + static_assert(sweep_slots <= SLOTS, + "a sweep would need more register slots than SLOTS declares"); + // The tiling condition the overlap bug violated: one iteration's slots must + // cover exactly the windows the stride advances, no more and no fewer. + static_assert(sweep_slots / chunks * had_group == sweep_stride, + "sweep slots and sweep_stride disagree: iterations overlap"); + + static_assert(Rows > 0, "no Rows makes Rows*K a whole TILE_GRAIN: bad K"); + static_assert(K % MX_BLOCK == 0, "a block may not straddle a row boundary"); + static_assert(TILE_GRAIN == VSTS_ALIGN * MX_BLOCK, "grain != scale DMA row"); + static_assert(tile_elems % TILE_GRAIN == 0, "tile is not a whole grain"); + static_assert(tile_elems % (2u * B16_LANES) == 0, "pack_nibbles wants 256"); + static_assert(scale_bytes % VSTS_ALIGN == 0, "scale row is not a legal DMA"); + static_assert(q_bytes % VSTS_ALIGN == 0, "nibble row is not a legal DMA"); + static_assert(in_bytes % VSTS_ALIGN == 0, "input row is not a legal DMA"); + static_assert(blocks % VCGMAX_B16_RESULTS == 0, + "blocks != whole vcgmax groups"); + // the rounded-up passes run one bite past the data; prove they stay inside + static_assert(b_iters * B16_LANES <= aligned_s, "scale tail overruns"); + static_assert(b_iters * B16_LANES * 2u <= aligned_mult, + "recips tail overruns"); + static_assert(packed_bytes <= aligned_packed, "compaction tail overruns"); + static_assert(groups * VSTS_ALIGN <= aligned_max, "padded maxima overrun"); + static_assert( + c_iters * VCGMAX_B16_RESULTS <= b_iters * B16_LANES, + "pack_nibbles would index recips past what derive_scales wrote"); + static_assert(sizeof(bfloat16_t) == 2, "RowsFor assumes 2-byte elements"); + static_assert(ub_needed <= UB_BYTES, "UB overflow"); + // Strictly less, not <=. The once-per-launch D load signals on EVENT_ID7 over + // MTE2 -> V, and buffer_free[7] is also EVENT_ID7 on that same pipe pair, so + // at NBuffers == EVENT_SLOTS buffer 7 and the D preamble would share a + // channel. Unreachable at the shipped NBuffers of 3; this stops it becoming + // reachable. + static_assert(NBuffers < EVENT_SLOTS, + "NBUF must leave EVENT_ID7 free for the D preamble"); + static_assert(NPrefetch < NBuffers, + "PREFETCH == NBUF deadlocks the pipeline"); +}; + +// Byte offsets within a pipeline slot, plus shared scratch. Constexpr +// *variables* for the reason above, so the slot base is multiplied in at +// the use site. +template +struct SlotOffset { + static constexpr unsigned input = 0; + static constexpr unsigned nibbles = Shape::aligned_in; + static constexpr unsigned scales = Shape::aligned_in + Shape::aligned_q; + static constexpr unsigned maxima = Shape::scratch_base; + static constexpr unsigned packed = Shape::scratch_base + Shape::aligned_max; + static constexpr unsigned reciprocal = packed + Shape::aligned_packed; +}; + +#ifdef __DAV_VEC__ +// A flat run of Elems values in GM, and the matching UB tile, for one dtype. +template +using GmShape = pto::Shape<1, 1, 1, 1, Elems>; +template +using GmStride = pto::Stride<1, 1, 1, Elems, 1>; +template +using UbTile = Tile; +// Same tile with a RUNTIME valid column count, zero-filling the rest in UB, for +// the one partial tile a batch can end on. +template +using UbTilePart = + Tile; + +// ------------------------------------------------------- block_abs_max +// Per-32-element magnitude max. A 2:1 fold makes 16 lanes == one block, +// which is what vcgmax's group size requires. A 4:1 fold silently reports +// max(block 2j, block 2j+1) instead. +// --- the rotation ------------------------------------------------------------ +// One sweep: deinterleave-load a window, add/sub, and store the halves back +// concatenated. Registers are vector_u16 because vlds/vsts are bit-width ops; +// the arithmetic type is chosen by reference cast, which is how bf16 costs +// nothing here. All loads precede all stores, which the comma fold guarantees +// by evaluating left to right -- required, not stylistic, since a store would +// otherwise clobber a window a later load still needs. +using HadRegs = vector_u16[SLOTS]; + +template +inline AICORE void sweep(__ubuf__ uint16_t *tile, uint32_t base, MaskReg all, + HadRegs &even, HadRegs &odd, HadRegs &sum, + HadRegs &diff, vector_u16 &idx_lo, vector_u16 &idx_hi, + std::index_sequence) { + constexpr unsigned g = Shape::had_group, up = Shape::upper; + constexpr unsigned ln = Shape::lanes, ch = Shape::chunks; + (vlds(even[Slot], odd[Slot], + tile + base + Slot / ch * g + Slot % ch * 2u * ln, 0, DINTLV_B16), + ...); + (vadd((vector_bf16 &)sum[Slot], (vector_bf16 &)even[Slot], + (vector_bf16 &)odd[Slot], all), + ...); + (vsub((vector_bf16 &)diff[Slot], (vector_bf16 &)even[Slot], + (vector_bf16 &)odd[Slot], all), + ...); + // A 256-element window packs eight independent 32-blocks, which leaves the + // result rotated right by log2(window/block) = 3; these register-only + // deinterleaves undo it, fused into the final stage using the pair that is + // dead by then. + // + // ATTRIBUTION SWITCH. vdintlv was measured at ~20x a vadd, and there are + // Rotations per slot here against five arithmetic ops, so this fixup may be + // most of the butterfly's cost. -DFUSED_NO_ROTFIX drops the deinterleaves and + // keeps everything else, including both stores, so the difference is the + // fixup. It PRODUCES WRONG OUTPUT -- the registers selected below then hold + // loaded values rather than rotated ones -- so it is for timing only and the + // benchmark's correctness gate will reject it. +#ifndef FUSED_NO_ROTFIX + if constexpr (Rotations >= 1) { + (vdintlv(even[Slot], odd[Slot], sum[Slot], diff[Slot]), ...); + } + if constexpr (Rotations >= 2) { + (vdintlv(sum[Slot], diff[Slot], even[Slot], odd[Slot]), ...); + } + if constexpr (Rotations >= 3) { + (vdintlv(even[Slot], odd[Slot], sum[Slot], diff[Slot]), ...); + } +#endif + HadRegs &lo = (Rotations % 2 == 1) ? even : sum; + HadRegs &hi = (Rotations % 2 == 1) ? odd : diff; +#if FUSED_SCATTER == 0 + (vsts(lo[Slot], tile + base + Slot / ch * g + Slot % ch * ln, 0, NORM_B16, + all), + ...); + (vsts(hi[Slot], tile + base + Slot / ch * g + up + Slot % ch * ln, 0, + NORM_B16, all), + ...); +#else + // vscatter instead of vsts, one for one. The index decides which experiment + // this is; see the FUSED_SCATTER comment at the top of the file. Both halves + // address the SAME window base, because with a permuting index the upper half + // is no longer a contiguous run at +upper. + (vscatter(lo[Slot], tile + base + Slot / ch * g, idx_lo, all), ...); + (vscatter(hi[Slot], tile + base + Slot / ch * g, idx_hi, all), ...); + (void)up; +#endif +} + +// --- phase 2: the cross-window stages --------------------------------------- +// +// After phase 1 every 256-element window holds its own order-256 transform, and +// H_K = H_(K/256) (x) H_256 leaves log2(K/256) stages that pair windows +// elementwise. Measured, these stages were 69-73% of the kernel at K=8192 and +// 16384, against phase 1 costing nothing at all -- phase 1 walks windows in +// order while a phase-2 stage reads windows 256*t apart, up to 16 KB at the +// last stage, and that stride is the cost. Op count is not the explanation: +// phase 1 does more loads and stores per element and is entirely hidden. +// +// So the stages are FUSED. A group of 2^R windows is loaded once, R stages run +// register to register, and the group is stored once -- R strided passes become +// one. FUSED_CROSS_FUSE=1 reproduces one stage per pass, which is what the A/B +// against this is. +#ifndef FUSED_CROSS_FUSE +#define FUSED_CROSS_FUSE 3 +#endif +constexpr unsigned CROSS_FUSE = FUSED_CROSS_FUSE; +static_assert( + (1u << CROSS_FUSE) <= SLOTS, + "a fused group holds 2^FUSED_CROSS_FUSE windows, and each needs a " + "register slot in both ping-pong arrays"); + +// One stage across the window axis, register to register. Reading from `src` +// and writing to `dst` is what removes the copy an in-place butterfly needs: +// window m's new value combines it with its partner m^bit, added when the bit +// is clear and subtracted the other way when it is set, so every window is one +// instruction and no temporary survives the stage. +template +inline AICORE void one_window(MaskReg all, HadRegs &s_lo, HadRegs &s_hi, + HadRegs &d_lo, HadRegs &d_hi) { + constexpr std::size_t P = M ^ (1u << Bit); + if constexpr ((M & (1u << Bit)) == 0) { + vadd((vector_bf16 &)d_lo[M], (vector_bf16 &)s_lo[M], (vector_bf16 &)s_lo[P], + all); + vadd((vector_bf16 &)d_hi[M], (vector_bf16 &)s_hi[M], (vector_bf16 &)s_hi[P], + all); + } else { + vsub((vector_bf16 &)d_lo[M], (vector_bf16 &)s_lo[P], (vector_bf16 &)s_lo[M], + all); + vsub((vector_bf16 &)d_hi[M], (vector_bf16 &)s_hi[P], (vector_bf16 &)s_hi[M], + all); + } +} + +template +inline AICORE void window_stage(MaskReg all, HadRegs &s_lo, HadRegs &s_hi, + HadRegs &d_lo, HadRegs &d_hi, + std::index_sequence) { + (one_window(all, s_lo, s_hi, d_lo, d_hi), ...); +} + +// R stages over the same registers, alternating direction so neither array is +// read and written in the same stage. +template +inline AICORE void window_stages(MaskReg all, HadRegs &a_lo, HadRegs &a_hi, + HadRegs &b_lo, HadRegs &b_hi, + std::index_sequence ms) { + if constexpr (I < R) { + if constexpr (I % 2u == 0u) + window_stage(all, a_lo, a_hi, b_lo, b_hi, ms); + else + window_stage(all, b_lo, b_hi, a_lo, a_hi, ms); + window_stages(all, a_lo, a_hi, b_lo, b_hi, ms); + } +} + +// One group: 2^R windows in, R stages, 2^R windows out. Window m of the group +// sits at m * 2^S0 windows from the base, because m's bits map onto window bits +// S0..S0+R-1 and those are contiguous. +template +inline AICORE void cross_group(__ubuf__ uint16_t *tile, uint32_t base, + MaskReg all, HadRegs &a_lo, HadRegs &a_hi, + HadRegs &b_lo, HadRegs &b_hi, + std::index_sequence ms) { + constexpr unsigned ln = B16_LANES, win = Shape::had_window; + constexpr unsigned step = (1u << S0) * win; + (vlds(a_lo[M], tile + base + M * step, 0, NORM), ...); + (vlds(a_hi[M], tile + base + M * step + ln, 0, NORM), ...); + window_stages(all, a_lo, a_hi, b_lo, b_hi, ms); + // R stages land back in a_* when R is even and in b_* when it is odd + if constexpr (R % 2u == 0u) { + (vsts(a_lo[M], tile + base + M * step, 0, NORM_B16, all), ...); + (vsts(a_hi[M], tile + base + M * step + ln, 0, NORM_B16, all), ...); + } else { + (vsts(b_lo[M], tile + base + M * step, 0, NORM_B16, all), ...); + (vsts(b_hi[M], tile + base + M * step + ln, 0, NORM_B16, all), ...); + } +} + +// Every group for one fused pass. A group's base has window bits S0..S0+R-1 +// clear, so the bases are (hi << (S0+R)) | lo over the two remaining ranges, +// and rows never mix. +template +inline AICORE void cross_pass(__ubuf__ uint16_t *tile, MaskReg all, + HadRegs &a_lo, HadRegs &a_hi, HadRegs &b_lo, + HadRegs &b_hi) { + constexpr unsigned win = Shape::had_window; + constexpr unsigned nwin = Shape::windows_per_row; + constexpr unsigned lo_span = 1u << S0; + constexpr unsigned hi_step = 1u << (S0 + R); + constexpr auto ms = std::make_index_sequence<(1u << R)>{}; + static_assert(nwin % hi_step == 0, "a fused group must fit the row"); + for (uint16_t row = 0; row < (uint16_t)Shape::rows_in_tile; ++row) + for (uint16_t hi = 0; hi < (uint16_t)(nwin / hi_step); ++hi) + for (uint16_t lo = 0; lo < (uint16_t)lo_span; ++lo) + cross_group( + tile, + (uint32_t)row * Shape::row_elems + + ((uint32_t)hi * hi_step + (uint32_t)lo) * win, + all, a_lo, a_hi, b_lo, b_hi, ms); +} + +// Walk the stages in fused groups, largest first, with whatever remains taking +// a narrower final group. +template +inline AICORE void cross_from(__ubuf__ uint16_t *tile, MaskReg all, + HadRegs &a_lo, HadRegs &a_hi, HadRegs &b_lo, + HadRegs &b_hi) { + if constexpr (S0 < Shape::phase2_stages) { + constexpr unsigned left = Shape::phase2_stages - S0; + constexpr unsigned R = left < CROSS_FUSE ? left : CROSS_FUSE; + cross_pass(tile, all, a_lo, a_hi, b_lo, b_hi); + // the next group reads what this one wrote + mem_bar(VST_VLD); + cross_from(tile, all, a_lo, a_hi, b_lo, b_hi); + } +} + +template +__tf__ static AICORE void cross_windows(__ubuf__ uint16_t *tile) { + __VEC_SCOPE__ { + uint32_t lane_count = B16_LANES; + MaskReg all = CreatePredicate(lane_count); + vector_u16 a_lo[SLOTS], a_hi[SLOTS], b_lo[SLOTS], b_hi[SLOTS]; + cross_from(tile, all, a_lo, a_hi, b_lo, b_hi); + } +} + +// log2(K) stages over the tile already in UB, in place. The quant passes read +// the same buffer straight afterwards, which is the point of the fusion. +template +__tf__ static AICORE void rotate(__ubuf__ uint16_t *tile) { + // sweep_slots, not SLOTS: see the derivation in Shape. The register arrays + // are sized SLOTS and a shorter pack leaves the top ones unused. + constexpr auto slots = std::make_index_sequence{}; + constexpr unsigned plain = + Shape::phase1_stages - (Shape::rotations ? 1u : 0u); + __VEC_SCOPE__ { + uint32_t lane_count = Shape::lanes; + MaskReg all = CreatePredicate(lane_count); + vector_u16 even[SLOTS], odd[SLOTS], sum[SLOTS], diff[SLOTS]; + // Scatter indices, built once per tile and unused when FUSED_SCATTER == 0. + // There is no vands, so the low three bits of the lane come out as + // l - (l >> 3) * 8. vshrs/vmuls/vadds are vector-scalar; vadd/vsub are + // vector-vector. + vector_u16 idx_lo, idx_hi; + vci((vector_s16 &)idx_lo, (int16_t)0, INC_ORDER); +#if FUSED_SCATTER == 1 + // IDENTITY index: byte-for-byte what the vsts pair does, so the output must + // be bit-identical and the only difference measured is the opcode price. + vdup(idx_hi, (uint16_t)Shape::upper, all, MODE_ZEROING); + vadd(idx_hi, idx_lo, idx_hi, all); +#endif + // Step by a literal 1 with the stride folded into base: the loop analyser + // only verifies a tripcount for a literal step, and 1 divides any bound, so + // had_iters may be template-dependent. + for (uint16_t stage = 0; stage < (uint16_t)plain; ++stage) { + for (uint16_t iter = 0; iter < (uint16_t)Shape::had_iters; ++iter) + sweep(tile, (uint32_t)iter * Shape::sweep_stride, all, even, + odd, sum, diff, idx_lo, idx_hi, slots); + mem_bar(VST_VLD); + } + if constexpr (Shape::rotations > 0) { + for (uint16_t iter = 0; iter < (uint16_t)Shape::had_iters; ++iter) + sweep( + tile, (uint32_t)iter * Shape::sweep_stride, all, even, odd, sum, + diff, idx_lo, idx_hi, slots); + mem_bar(VST_VLD); + } + } +} + +template +__tf__ static AICORE void block_abs_max(__ubuf__ uint16_t *input, + __ubuf__ uint16_t *maxima) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + // PAT_VL8 matches VCGMAX_B16_RESULTS + MaskReg low_eight = pset_b16(PAT_VL8); + vector_u16 abs_mask; + vdup(abs_mask, BF16_ABS, all_lanes, MODE_ZEROING); + + for (uint16_t group = 0; group < (uint16_t)Shape::groups; ++group) { + const uint32_t base = (uint32_t)group * 256u; + vector_u16 even, odd, folded, grouped; + vlds(even, odd, input + base, 0, + DINTLV_B16); // lane i: elements 2i, 2i+1 + vand(even, even, abs_mask, all_lanes); + vand(odd, odd, abs_mask, all_lanes); + // sign cleared, so a signed max over the bit patterns IS a magnitude max + vmax((vector_s16 &)folded, (vector_s16 &)even, (vector_s16 &)odd, + all_lanes); + vcgmax((vector_s16 &)grouped, (vector_s16 &)folded, all_lanes); + // 32-byte pitch, not 16: see VSTS_ALIGN + vsts(grouped, maxima + (uint32_t)group * GROUP_PITCH_B16, 0, NORM_B16, + low_eight); + } + mem_bar(VST_VLD); + } +} + +// ------------------------------------------------------ compact_maxima +// Squeeze out the padding VSTS_ALIGN forces: output byte i takes input byte +// 2*(i & 0xF0) + (i & 0x0F). +template +__tf__ static AICORE void compact_maxima(__ubuf__ uint16_t *padded, + __ubuf__ uint16_t *packed) { + __VEC_SCOPE__ { + MaskReg all_byte_lanes = pset_b8(PAT_ALL); + MaskReg low_32 = pset_b16(PAT_VL32); // 32 b16 == 64 bytes + vector_u8 byte_index, high_half, low_half, high_mask, low_mask; + vci((vector_s8 &)byte_index, (int8_t)0, INC_ORDER); + vdup(high_mask, (uint8_t)0xF0, all_byte_lanes, MODE_ZEROING); + vdup(low_mask, (uint8_t)0x0F, all_byte_lanes, MODE_ZEROING); + vand(high_half, byte_index, high_mask, all_byte_lanes); + vand(low_half, byte_index, low_mask, all_byte_lanes); + vadd((vector_s8 &)high_half, (vector_s8 &)high_half, (vector_s8 &)high_half, + all_byte_lanes); // 2*high_half + vadd((vector_s8 &)byte_index, (vector_s8 &)high_half, (vector_s8 &)low_half, + all_byte_lanes); + + for (uint16_t gather = 0; gather < (uint16_t)Shape::compact_iters; + ++gather) { + vector_u16 padded_chunk, packed_chunk; + const uint32_t src_offset = + (uint32_t)gather * GROUPS_PER_COMPACT * GROUP_PITCH_B16; + vlds(padded_chunk, padded + src_offset, 0, NORM); + vselr((vector_u8 &)packed_chunk, (vector_u8 &)padded_chunk, byte_index); + vsts(packed_chunk, + packed + (uint32_t)gather * GROUPS_PER_COMPACT * VCGMAX_B16_RESULTS, + 0, NORM_B16, low_32); + } + mem_bar(VST_VLD); + } +} + +// -------------------------------------------------------- derive_scales +// maxima -> E8M0 scale byte + one bf16 reciprocal per block. pack_nibbles +// reads this array with E2B_B16, whose x16 replication matches its +// pair-granular deinterleave exactly, so no duplication is needed here. +template +__tf__ static AICORE void derive_scales(__ubuf__ uint16_t *maxima, + __ubuf__ uint16_t *recips_out, + __ubuf__ uint16_t *scale_out) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + vector_u16 bias; + vdup(bias, (uint16_t)RECIP_OFFSET, all_lanes, MODE_ZEROING); + + for (uint16_t chunk = 0; chunk < (uint16_t)Shape::b_iters; ++chunk) { + vector_u16 block_max, exponent, scale_byte, reciprocal; + vlds(block_max, maxima + (uint32_t)chunk * B16_LANES, 0, NORM); + // bit 15 is already clear, so this shift alone yields the biased exponent + vshrs(exponent, block_max, BF16_MANT_BITS, all_lanes, MODE_ZEROING); + vmaxs(exponent, exponent, B_MIN, all_lanes); + vmins(exponent, exponent, B_MAX, all_lanes); + vadds(scale_byte, exponent, E8M0_BIAS_ADJ, all_lanes); + vsts(scale_byte, scale_out + (uint32_t)chunk * 64u, 0, PK_B16, all_lanes); + vsub(reciprocal, bias, exponent, all_lanes); + vshls(reciprocal, reciprocal, BF16_MANT_BITS, all_lanes, MODE_ZEROING); + vsts(reciprocal, recips_out + (uint32_t)chunk * B16_LANES, 0, NORM_B16, + all_lanes); + } + mem_bar(VST_VLD); + } +} + +// --------------------------------------------------------- pack_nibbles +// Scale, cast, pack -- 256 elements per iteration, no gather. +// One vcvt puts 64 bytes at byte STRIDE 4, offset chosen by +// PART_P0..P3, so converting two halves into offsets 0 and 1, OR-ing, and +// storing with PK_B32 (keeps the low 2 bytes of each 4-byte group) writes 128 +// CONTIGUOUS bytes. RULE: fp4 packs two elements per byte, so DINTLV_B16 would +// pair element 4k with 4k+2 -- deinterleave at b32 (pairs) to keep (4k, 4k+1) +// together. That also puts both b16 lanes of a half in block j/8, so E2B_B16's +// x16 replication is exact and one multiplier register serves both halves. +template +__tf__ static AICORE void pack_nibbles(__ubuf__ uint16_t *input, + __ubuf__ uint16_t *reciprocal, + __ubuf__ uint8_t *nibble_out) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + MaskReg all_byte_lanes = pset_b8(PAT_ALL); + MaskReg all_b32_lanes = pset_b32(PAT_ALL); + + for (uint16_t chunk = 0; chunk < (uint16_t)Shape::c_iters; ++chunk) { + vector_u16 recips; + vector_u32 even, odd; + vector_bf16 scaled_even, scaled_odd; + vector_f4e2m1x2 packed_even, packed_odd, packed; + vlds(recips, reciprocal + (uint32_t)chunk * VCGMAX_B16_RESULTS, 0, + E2B_B16); + vlds(even, odd, (__ubuf__ uint32_t *)input + (uint32_t)chunk * B16_LANES, + 0, DINTLV_B32); + vmul(scaled_even, (vector_bf16 &)even, (vector_bf16 &)recips, all_lanes); + vmul(scaled_odd, (vector_bf16 &)odd, (vector_bf16 &)recips, all_lanes); + vcvt(packed_even, scaled_even, all_lanes, ROUND_R, PART_P0); + vcvt(packed_odd, scaled_odd, all_lanes, ROUND_R, PART_P1); + vor((vector_u8 &)packed, (vector_u8 &)packed_even, + (vector_u8 &)packed_odd, all_byte_lanes); + // 256 elements in, but PK_B32 keeps 2 of every 4 bytes: 128 bytes out + vsts((vector_u16 &)packed, + (__ubuf__ uint16_t *)(nibble_out + (uint32_t)chunk * B16_LANES), 0, + PK_B32, all_b32_lanes); + } + mem_bar(VST_VLD); + } +} + +#ifdef MXFP4_TQUANT +// Requires PTO 9.1.0: 9.0.0 has no MXFP4 quantizer. Included here, not at file +// scope, because this region is inside the device-pass guard. +#include + +// ------------------------------------------------------- tquant_passes +// One vendor tile op in place of block_abs_max, compact_maxima, derive_scales +// and pack_nibbles. validCols is tile_elems even on the partial tile: the load +// already zero-fills the pad, and a short validCols would send TQuant's own +// ZeroPadSourceTile over the input slot. Offsets::packed is left allocated and +// unused, since reclaiming it would move slot_stride. +template +inline AICORE void tquant_passes(uint32_t input_offset, uint32_t nibble_offset, + uint32_t scale_offset) { + static_assert(sizeof(float4_e2m1x2_t) == 1, + "the nibble tile assumes one byte per float4_e2m1x2_t"); + static_assert(REPEAT_BYTE / sizeof(bfloat16_t) == B16_LANES, + "tquant_max_elems assumes a 128-lane b16 vector"); + UbTile source; + UbTile nibbles; + UbTile scales; + UbTile block_max; + UbTile reciprocal; + TASSIGN(source, input_offset); + TASSIGN(nibbles, nibble_offset); + TASSIGN(scales, scale_offset); + TASSIGN(block_max, SlotOffset::maxima); + TASSIGN(reciprocal, SlotOffset::reciprocal); + // TEMPLATE order is Out, Src, Exp, Max, Scaling; ARGUMENT order is dst, exp, + // max, scaling, src. PTO 9.1.0 release inserted a `bool Exp2DStrided` second + // template parameter that 9.1.0-beta.3 does not have; the tile types are in a + // non-deduced position, so neither spelling can be dropped. benchmark.py + // compiles both and keeps whichever the local headers accept. +#ifdef MXFP4_TQUANT_EXP2D + TQuant_MXFP4_E2M1_Impl( + nibbles.data(), scales.data(), block_max.data(), reciprocal.data(), + source.data(), 1u, Shape::tile_elems); +#else + TQuant_MXFP4_E2M1_Impl( + nibbles.data(), scales.data(), block_max.data(), reciprocal.data(), + source.data(), 1u, Shape::tile_elems); +#endif +} +#endif // MXFP4_TQUANT + +// Move one tile of `T` between GM and UB. Partial carries only `valid` +// elements: the load zero-fills the rest of the UB tile so the compute passes +// still see whole registers, and the store truncates so padding never reaches +// GM. +template +inline AICORE void move_tile(uint32_t tile_index, uint32_t ub_offset, + __gm__ void *gm_base, uint32_t valid = 0) { + std::conditional_t, UbTile> ub; + TASSIGN(ub, ub_offset); + if constexpr (Partial) ub.ColMaskInternal = (int)valid; + GlobalTensor, GmStride> gm( + (__gm__ T *)gm_base + (uint64_t)tile_index * Elems, GmShape()); + if constexpr (ToUb) { + TLOAD(ub, gm); + } else { + TSTORE(gm, ub); + } +} + +// Start the async load of this core's nth tile, if it has one. A function, +// not a lambda: set_flag/wait_flag do not resolve inside a lambda. +template +inline AICORE void issue_tile_load(uint32_t nth_tile, uint32_t core_id, + uint32_t core_count, uint32_t tiles, + uint32_t full_tiles, uint32_t tail_elems, + const event_t *buffer_free, + __gm__ void *input_gm) { + const uint32_t tile_index = core_id + nth_tile * core_count; + if (tile_index >= tiles) return; + const uint32_t buffer = nth_tile % Buffers; + const uint32_t off = buffer * Shape::slot_stride + SlotOffset::input; + wait_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[buffer]); + // at most one tile is partial, and only when batch does not fill it + if (tile_index == full_tiles) { + move_tile(tile_index, off, + input_gm, tail_elems); + } else { + move_tile(tile_index, off, input_gm); + } + set_flag(PIPE_MTE2, PIPE_V, buffer_free[buffer]); +} +#endif // __DAV_VEC__ +#endif // __CCE_AICORE__ + +// The pipeline: each core walks a strided subset of the tiles, keeping Prefetch +// loads in flight so DMA and the vector pipe overlap. +#if defined(__CCE_AICORE__) && defined(__DAV_VEC__) +// A device function rather than the kernel body, so a caller that wants the +// pipeline over a sub-range can reach it directly. mxfp4_quant below is the +// entry point and the only caller here. +template +inline AICORE void quant_tiles(__gm__ void *input_gm, __gm__ void *nibble_gm, + __gm__ void *scale_gm, uint32_t batch) { + using Shape = QuantShape; + using Offsets = SlotOffset; + set_mask_norm(); + set_vector_mask(-1, -1); + const event_t buffer_free[EVENT_SLOTS] = {EVENT_ID0, EVENT_ID1, EVENT_ID2, + EVENT_ID3, EVENT_ID4, EVENT_ID5, + EVENT_ID6, EVENT_ID7}; + const uint32_t core_id = get_block_idx(), core_count = get_block_num(); + // the remainder, if any, rides along as one extra partial tile + const uint32_t full_tiles = batch / Rows; + const uint32_t tail_elems = (batch % Rows) * K; + const uint32_t tiles = full_tiles + (tail_elems ? 1u : 0u); + + for (unsigned i = 0; i < NBuffers; ++i) // every buffer starts free + set_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[i]); + for (unsigned i = 0; i < NPrefetch; ++i) + issue_tile_load(i, core_id, core_count, tiles, full_tiles, + tail_elems, buffer_free, input_gm); + + uint32_t issued = 0; + for (uint32_t tile_index = core_id; tile_index < tiles; + tile_index += core_count, ++issued) { + const uint32_t buffer = issued % NBuffers; + // issued ahead of the wait below, so this load overlaps this tile's compute + issue_tile_load(issued + NPrefetch, core_id, core_count, + tiles, full_tiles, tail_elems, buffer_free, + input_gm); + wait_flag(PIPE_MTE2, PIPE_V, buffer_free[buffer]); + const uint32_t slot_base = buffer * Shape::slot_stride; +#ifdef MXFP4_TQUANT + tquant_passes(slot_base + Offsets::input, + slot_base + Offsets::nibbles, + slot_base + Offsets::scales); +#else + // name the UB regions once; inline casts are noise at every call site + using B16 = __ubuf__ uint16_t *; + B16 input_ub = (B16)(uintptr_t)(slot_base + Offsets::input); + B16 scale_ub = (B16)(uintptr_t)(slot_base + Offsets::scales); + B16 maxima_ub = (B16)(uintptr_t)Offsets::maxima; + B16 packed_ub = (B16)(uintptr_t)Offsets::packed; + B16 recips_ub = (B16)(uintptr_t)Offsets::reciprocal; + __ubuf__ uint8_t *nibble_ub = + (__ubuf__ uint8_t *)(uintptr_t)(slot_base + Offsets::nibbles); + // rotate in place, then quantize the rotated tile without it ever leaving + // UB +#ifndef FUSED_NO_ROTATE + rotate(input_ub); + // phase 2 finishes the transform when a row is wider than one window. A + // separate call because __tf__ may not call __tf__, and a separate + // __VEC_SCOPE__ because it needs its own register set. + // + // ATTRIBUTION SWITCH. -DFUSED_NO_CROSS drops phase 2 and keeps phase 1 and + // the quantizer, so the difference is what the cross-window stages cost. + // It PRODUCES WRONG OUTPUT for K > 256 -- the transform is unfinished -- so + // it is for timing only and every correctness gate will reject it. +#ifndef FUSED_NO_CROSS + if constexpr (Shape::phase2_stages > 0) { + cross_windows(input_ub); + } +#endif +#else + // Diagnostic build: same kernel, same tiling, same UB layout and buffer + // count -- only the butterfly removed. Comparing this against the quantizer + // alone separates the butterfly's vector cost from the cost of fusing at + // all (extra UB regions, so fewer buffers, so less overlap). + (void)0; +#endif +#ifndef FUSED_ROTATE_ONLY + block_abs_max(input_ub, maxima_ub); + compact_maxima(maxima_ub, packed_ub); + derive_scales(packed_ub, recips_ub, scale_ub); + pack_nibbles(input_ub, recips_ub, nibble_ub); +#else + // The other half of the fusion question. FUSED_NO_ROTATE keeps the + // quantizer and drops the butterfly; this keeps the butterfly and drops the + // quantizer, storing the rotated bf16 tile instead. Chained with the + // standalone quantizer it is the UNFUSED reference: two launches, two + // passes over HBM, 4 + 2.53 B/elem against the fused kernel's 2.53. + // + // Same tiling, UB layout and buffer count as the fused build, so the only + // differences against it are the arithmetic skipped and the bytes stored. + (void)scale_ub; + (void)maxima_ub; + (void)packed_ub; + (void)recips_ub; + (void)nibble_ub; +#endif +#endif + set_flag(PIPE_V, PIPE_MTE3, buffer_free[buffer]); + wait_flag(PIPE_V, PIPE_MTE3, buffer_free[buffer]); +#ifdef FUSED_ROTATE_ONLY + // `nibble_gm` carries the rotated bf16 tile here and `scale_gm` is + // untouched, so the launcher signature does not change. The harness + // allocates 2K bytes per row for it, not K/2. + if (tile_index == full_tiles) { + move_tile( + tile_index, slot_base + Offsets::input, nibble_gm, tail_elems); + } else { + move_tile( + tile_index, slot_base + Offsets::input, nibble_gm); + } + (void)scale_gm; +#else + if (tile_index == full_tiles) { + move_tile( + tile_index, slot_base + Offsets::nibbles, nibble_gm, tail_elems / 2u); + move_tile( + tile_index, slot_base + Offsets::scales, scale_gm, + tail_elems / MX_BLOCK); + } else { + move_tile( + tile_index, slot_base + Offsets::nibbles, nibble_gm); + move_tile( + tile_index, slot_base + Offsets::scales, scale_gm); + } +#endif + set_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[buffer]); + } + for (unsigned i = 0; i < NBuffers; ++i) // drain + wait_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[i]); +} +#endif // __CCE_AICORE__ && __DAV_VEC__ + +template +__global__ AICORE void mxfp4_quant(__gm__ void *input_gm, + __gm__ void *nibble_gm, + __gm__ void *scale_gm, uint32_t batch) { +#ifdef __DAV_VEC__ + quant_tiles(input_gm, nibble_gm, scale_gm, + batch); +#else + (void)input_gm; + (void)nibble_gm; + (void)scale_gm; + (void)batch; +#endif +} + +#ifndef FUSED_INCLUDE_ONLY // define to take the device code without hosts +// ---------------------------------------------------------------- entry points +// One .so serves every K: fold over SUPPORTED_K for the instantiation. +template +inline void launch_for_k(uint32_t block_dim, void *stream, uint8_t *input, + uint8_t *nibbles, uint8_t *scales, uint32_t batch, + uint32_t k, std::index_sequence) { + ((k == SUPPORTED_K[Idx] + ? (void)(mxfp4_quant::value, + DEF_BUFFERS, DEF_PREFETCH> + <<>>(input, nibbles, scales, + batch)) + : (void)0), + ...); +} + +// An unsupported k is a silent no-op; the host validates +// (check_row_width). +extern "C" void call_hadamard_mxfp4_full(uint32_t block_dim, void *stream, + uint8_t *input, uint8_t *nibbles, + uint8_t *scales, uint32_t batch, + uint32_t k) { + launch_for_k(block_dim, stream, input, nibbles, scales, batch, k, + std::make_index_sequence{}); +} + +template +inline uint32_t rows_for_k(uint32_t k, std::index_sequence) { + uint32_t rows = 0; + ((k == SUPPORTED_K[Idx] ? (void)(rows = RowsFor::value) + : (void)0), + ...); + return rows; // 0 for an unsupported k +} + +extern "C" uint32_t hadamard_mxfp4_full_rows_for(uint32_t k) { + return rows_for_k(k, std::make_index_sequence{}); +} +#endif // FUSED_INCLUDE_ONLY diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/jit_util_fused_a5.py b/examples/jit_cpp/fused_hadamard_quant_a5/jit_util_fused_a5.py new file mode 100644 index 00000000..1c0bbcee --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/jit_util_fused_a5.py @@ -0,0 +1,161 @@ +"""Build and load the block-32 fused Hadamard + MXFP4 quantize kernel. + +Deliberately thin: the kernel's own launcher dispatches on K, so this only has to +compile one .so and hand back a callable. Modelled on mxfp4_quant's jit helper, +with the entry points renamed and the width list narrowed to those the rotation +supports. +""" + +import ctypes +import os +import subprocess +from pathlib import Path + +import torch +import torch_npu # noqa + +HERE = Path(__file__).resolve().parent +BUILDDIR = HERE / "build" +_LIB_NAME = "fused_full.so" +SOURCE = HERE / "fused_hadamard_quant_a5.cpp" + +MX_BLOCK = 32 +VECTOR_CORES = 64 # vector cores on an A5 +# The rotation is order K, so K must be a power of two. The width is not +# otherwise capped: the kernel does the transform as window-local work plus +# cross-window stages, and neither phase holds more than one 256-element window +# in registers. Must match SUPPORTED_K in the kernel. +SUPPORTED_K = (32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384) + + +def _flags(home): + return ( + f"-xcce --cce-aicore-arch=dav-c310-vec -DREGISTER_BASE " + f"-std=c++17 -O2 -fPIC -Wno-ignored-attributes -Wno-macro-redefined " + f"-mllvm -cce-aicore-stack-size=0x8000 " + f"-mllvm -cce-aicore-function-stack-size=0x8000 " + f"-mllvm -cce-aicore-addr-transform " + f"-mllvm -cce-aicore-dcci-insert-for-scalar=false -Xhost-start -Xhost-end " + f"-I{home}/aarch64-linux/include -I{home}/include" + ).split() + + +def compile_kernel(verbose=True, extra_defs=()): + """Compile the fused kernel to a .so. One .so serves every supported K. + + extra_defs are extra -D tokens for a tuning or A/B variant. They go into the + .so NAME as well as the command line, so a variant can never be served from + the default build's cache -- silently timing the wrong binary is the failure + this guards. + """ + home = os.environ.get("ASCEND_HOME_PATH") or os.environ.get("ASCEND_TOOLKIT_HOME") + if not home: + raise RuntimeError("source a CANN set_env.sh first: ASCEND_HOME_PATH is unset") + BUILDDIR.mkdir(parents=True, exist_ok=True) + tag = "".join("_" + d.lstrip("-D").replace("=", "") for d in sorted(extra_defs)) + # Reuse an .so newer than its source. These kernels unroll to hundreds of + # tile instructions and a rebuild can outlast the task queue's 600 s cap, so + # recompiling per call is not merely wasteful. + cached = BUILDDIR / _LIB_NAME.replace(".so", f"{tag}.so") + if cached.exists() and cached.stat().st_mtime > SOURCE.stat().st_mtime: + if verbose: + print("reusing", cached) + return cached + obj = BUILDDIR / f"fused_b32{tag}.o" + lib = cached + for step in ( + [ + f"{home}/bin/bisheng", + *_flags(home), + *extra_defs, + "-c", + str(SOURCE), + "-o", + str(obj), + ], + [ + f"{home}/bin/bisheng", + "-fPIC", + "-shared", + "--cce-fatobj-link", + f"-Wl,-soname,{lib.name}", + str(obj), + "-o", + str(lib), + ], + ): + if verbose: + print("compile:", " ".join(step[:3]), "...") + subprocess.run(step, check=True) + return lib + + +def current_stream_ptr(): + return ctypes.c_void_p(torch.npu.current_stream().npu_stream) + + +# The butterfly is the UNNORMALISED Sylvester matrix, so its output is sqrt(32) +# larger than an orthogonal block Hadamard's. That factor is deliberate and left +# to the caller: MXFP4's E8M0 scale is a power of two and sqrt(32) is not, so the +# scale cannot absorb it and the nibbles genuinely differ. Scale x by +# 1/sqrt(32) on the way in if orthogonal semantics are wanted. + + +def build_and_load(k=256, verbose=True, extra_defs=()): + """Return `fused(x) -> (nibbles, scales)` for row width `k`. + + Allocates its outputs, mirroring `torch_npu.npu_dynamic_mx_quant`, so the two + are comparable on the same call path. + + extra_defs reaches the compiler, so the reduced builds the benchmark's ladder + needs come from this one source: FUSED_ROTATE_ONLY leaves the butterfly + alone, FUSED_NO_ROTATE leaves the quantizer alone. + """ + if k not in SUPPORTED_K: + raise ValueError( + f"K={k} has no instantiation; supported: {sorted(SUPPORTED_K)}. " + "Widths must be a multiple of 32 with an instantiation." + ) + lib = ctypes.CDLL(str(compile_kernel(verbose=verbose, extra_defs=extra_defs))) + launch = lib.call_hadamard_mxfp4_full + launch.argtypes = [ + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_uint32, + ] + launch.restype = None + rows_for = lib.hadamard_mxfp4_full_rows_for + rows_for.argtypes = [ctypes.c_uint32] + rows_for.restype = ctypes.c_uint32 + + def fused(x, out=None): + if x.dtype != torch.bfloat16: + raise TypeError(f"expected bfloat16, got {x.dtype}") + if x.shape[-1] != k: + raise ValueError(f"expected last dim {k}, got {tuple(x.shape)}") + if not x.is_contiguous(): + raise ValueError("expected a contiguous tensor; call .contiguous()") + batch = x.numel() // k + if out is None: + q = torch.empty((batch, k // 2), dtype=torch.uint8, device=x.device) + s = torch.empty((batch, k // MX_BLOCK), dtype=torch.uint8, device=x.device) + else: + q, s = out + launch( + VECTOR_CORES, + current_stream_ptr(), + ctypes.c_void_p(x.data_ptr()), + ctypes.c_void_p(q.data_ptr()), + ctypes.c_void_p(s.data_ptr()), + batch, + k, + ) + return q, s + + fused.rows_for = lambda: rows_for(k) + fused.k = k + return fused diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/run_benchmark.sh b/examples/jit_cpp/fused_hadamard_quant_a5/run_benchmark.sh new file mode 100755 index 00000000..d575da23 --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/run_benchmark.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# One-command on-device benchmark for fused_hadamard_quant_a5 on an Ascend 950 (A5). +# Requires a real A5 device, torch + torch_npu, and bisheng (CANN toolkit). +# +# Needs a CANN whose PTO carries MXFP4 (Exp2DStrided in pto/npu/a5/TQuant.hpp): +# 9.1.0 and 9.2.0 both do, 9.0.0 does not. +if [[ -z "${ASCEND_TOOLKIT_HOME:-}" && -z "${ASCEND_HOME_PATH:-}" ]]; then + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +: "${ASCEND_HOME_PATH:=${ASCEND_TOOLKIT_HOME:-/usr/local/Ascend/ascend-toolkit/latest}}" +export ASCEND_HOME_PATH +cd "${SCRIPT_DIR}" +exec python3 benchmark.py "$@" diff --git a/examples/jit_cpp/fused_hadamard_quant_a5/test_fused_hadamard_quant_a5.py b/examples/jit_cpp/fused_hadamard_quant_a5/test_fused_hadamard_quant_a5.py new file mode 100644 index 00000000..ac73b42d --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_a5/test_fused_hadamard_quant_a5.py @@ -0,0 +1,332 @@ +"""Does the fused kernel rotate and quantize correctly? + +The fused kernel cannot be bit-exact against a torch reference: it rotates in +bf16 with a specific operand order, and no torch expression reproduces that +tree. So correctness is established in three ways instead, from strongest to +weakest: + +1. **The scale bytes** must match a reference that rotates in fp32 and quantizes + with `torch_npu`. A scale is a power of two derived from a block maximum, so + bf16 rounding inside the butterfly almost never moves it -- if scales + disagree, the rotation is wrong, not merely rounded differently. +2. **The dequantized values** must track the fp32-rotated reference to within + MXFP4's own resolution. This catches a correct-looking permutation, which a + relative-error check on the packed bytes would not. +3. **The output must be non-trivial.** A kernel that writes nothing, or writes + the input back, is the characteristic silent failure on this hardware, and it + would otherwise pass a loose tolerance. +""" + +import numpy as np +import pytest +import torch +import torch_npu # noqa + +from jit_util_fused_a5 import ( + MX_BLOCK, + SUPPORTED_K, + VECTOR_CORES, + build_and_load, +) + +VENDOR_DST_TYPE = 296 # E2M1, matching mxfp4_quant_a5's tests + + +def hadamard_matrix(n): + """Natural-order Sylvester +/-1 matrix, unnormalised -- the convention + fast_hadamard_a5 and its tests use.""" + m = np.array([[1.0]], dtype=np.float64) + while m.shape[0] < n: + m = np.block([[m, m], [m, -m]]) + return m + + +def fwht(x): + """Textbook iterative Walsh-Hadamard transform, natural order. + + A second implementation on purpose. The kernel does the transform as + window-local work plus cross-window stages; this does explicit strided + butterflies, so a mistake in the kernel's decomposition cannot cancel out of + both sides. `test_fwht_matches_the_explicit_matrix` pins this against the + Sylvester matrix at the widths where building that matrix is cheap -- at + K=16384 the matrix alone is a gigabyte, which is why the tests reference this + instead of the matrix directly. + """ + y = np.array(x, dtype=np.float64, copy=True) + k = y.shape[-1] + h = 1 + while h < k: + for i in range(0, k, 2 * h): + a = y[..., i : i + h].copy() + b = y[..., i + h : i + 2 * h].copy() + y[..., i : i + h] = a + b + y[..., i + h : i + 2 * h] = a - b + h *= 2 + return y + + +@pytest.mark.parametrize("k", (32, 64, 256, 1024)) +def test_fwht_matches_the_explicit_matrix(k): + """The reference implementation must equal x @ H_k. + + Every other test references `fwht`, so this is what makes those tests mean + "matches the Hadamard matrix" rather than "matches my other loop". + """ + rng = np.random.default_rng(k) + x = rng.standard_normal((4, k)) + want = x @ hadamard_matrix(k) + got = fwht(x) + assert np.abs(got - want).max() < 1e-9 * max(np.abs(want).max(), 1.0) + + +def reference(x, k): + """Rotate each row by the order-K transform, then quantize. + + One rotation across the whole row, matching the kernel: every output element + depends on all k inputs. Computed by `fwht`, a different decomposition from + the kernel's and itself pinned against the explicit Sylvester matrix, so a + mistake in the kernel cannot cancel out of both sides. fp64 on the host, so + the reference does not depend on the kernel's bf16 arithmetic. + """ + rot = torch.from_numpy(fwht(x.float().cpu().numpy())).to(torch.bfloat16).npu() + q, s = torch_npu.npu_dynamic_mx_quant(rot, dst_type=VENDOR_DST_TYPE) + return rot, q, s.reshape(s.shape[0], -1)[:, : k // MX_BLOCK] + + +E2M1_LEVELS = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 +) + + +def dequant(q, s, k): + """Unpack E2M1 nibbles and apply the E8M0 scale, on the host in fp32.""" + q = q.cpu() + lo, hi = q & 0x0F, (q >> 4) & 0x0F + codes = torch.stack([lo, hi], dim=-1).reshape(q.shape[0], -1) + mag = E2M1_LEVELS[(codes & 0x07).long()] + vals = torch.where(codes >= 8, -mag, mag) + exp = s.cpu().to(torch.int32) - 127 + scale = torch.ldexp(torch.ones_like(exp, dtype=torch.float32), exp) + return vals.reshape(-1, k // MX_BLOCK, MX_BLOCK) * scale.unsqueeze(-1) + + +@pytest.fixture(scope="module", autouse=True) +def seeded(): + torch.manual_seed(20260818) + torch.npu.set_device(0) + + +# A spread rather than all eight -- each width is a separate .so compile, so the +# full set costs minutes. The spread has to cover both SLOT-ADDRESSING CLASSES. +# +# A slot addresses group `Slot / chunks` at chunk `Slot % chunks`. For K <= 256 a +# group is a whole window holding several rows and chunks is 1; for K >= 512 a +# group is one row spread over 2, 4, 8 or 16 chunks. Those are different index +# arithmetic, and a matrix covering only one would leave the other live in +# production and dead in CI. That has happened on the predecessor of this +# kernel: every width in a five-width matrix landed in one class, the uncovered +# path was broken, the sweep ran past the end of the tile, and nothing noticed. +# +# Class membership is DERIVED from the width below rather than listed here, so it +# cannot go stale against a tuning change. +WIDTHS = (32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384) + + +@pytest.mark.parametrize("k", WIDTHS) +def test_matches_reference(k): + """Scales exact, dequantized values within MXFP4 resolution.""" + batch = 64 + fused = build_and_load(k=k, verbose=False) + x = torch.randn(batch, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + _, q_ref, s_ref = reference(x, k) + + scale_match = (s.cpu() == s_ref.cpu()).float().mean().item() + assert scale_match > 0.98, ( + f"K={k}: only {scale_match:.3f} of scale bytes match the fp32-rotated " + "reference -- that is a wrong rotation, not bf16 rounding" + ) + + got = dequant(q, s, k).reshape(batch, k) + want = dequant(q_ref, s_ref, k).reshape(batch, k) + denom = want.abs().mean().clamp_min(1e-6) + rel = (got - want).abs().mean() / denom + assert rel < 0.05, f"K={k}: dequantized mean rel error {rel:.4f} vs reference" + + +# Deep enough that the pipeline machinery runs, at both slot classes and at both +# ends of the rows range: 256 packs 96 rows per tile and never chunks, 4096 is +# 16 chunks with 6 rows. +DEEP_WIDTHS = (256, 1024, 8192) +TILES_PER_CORE = 4 + + +@pytest.mark.parametrize("k", DEEP_WIDTHS) +def test_matches_reference_many_tiles_per_core(k): + """The same check as test_matches_reference, but with a full pipeline. + + test_matches_reference uses batch=64, which at K=4096 is 11 tiles spread over + 64 cores: one tile for eleven cores and none for the rest. So the buffer + rotation (issued % NBuffers), the prefetch and the drain never run there, and + a fault in any of them cannot fail that test. This sizes the batch so every + core walks several tiles, and leaves a remainder so the partial tail tile is + taken too -- except at K=14336, where a tile is one row and no batch can + leave a remainder. + """ + fused = build_and_load(k=k, verbose=False) + rows = fused.rows_for() + batch = rows * VECTOR_CORES * TILES_PER_CORE + rows // 2 + 1 + tiles = -(-batch // rows) + + # The point of the test is the depth, so assert it rather than trusting that + # a TILE_ELEMS change left it intact. + assert tiles / VECTOR_CORES >= 3, ( + f"K={k}: {tiles} tiles over {VECTOR_CORES} cores is too shallow to " + "exercise the buffer rotation" + ) + # A tile is `rows` rows, so where rows == 1 every batch is a whole number of + # tiles and the kernel's partial branch is unreachable by construction. Only + # claim the tail where one can exist. + assert batch % rows or rows == 1, f"K={k}: batch {batch} is whole tiles" + + x = torch.randn(batch, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + _, q_ref, s_ref = reference(x, k) + + scale_match = (s.cpu() == s_ref.cpu()).float().mean().item() + assert scale_match > 0.98, ( + f"K={k}, {tiles} tiles: only {scale_match:.3f} of scale bytes match the " + "fp32-rotated reference" + ) + + got = dequant(q, s, k).reshape(batch, k) + want = dequant(q_ref, s_ref, k).reshape(batch, k) + denom = want.abs().mean().clamp_min(1e-6) + rel = (got - want).abs().mean() / denom + assert rel < 0.05, f"K={k}, {tiles} tiles: mean rel error {rel:.4f} vs reference" + + +@pytest.mark.parametrize("k", WIDTHS) +def test_output_is_nontrivial(k): + """A kernel that writes nothing, or echoes its input, must fail here.""" + fused = build_and_load(k=k, verbose=False) + x = torch.randn(32, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + assert q.any().item(), f"K={k}: nibbles are all zero" + assert s.any().item(), f"K={k}: scale bytes are all zero" + assert len(torch.unique(q.cpu())) > 4, f"K={k}: nibbles are degenerate" + + +@pytest.mark.parametrize("k", WIDTHS) +def test_rotation_actually_happened(k): + """The rotation must change the answer. + + Quantizing x directly and quantizing (x @ H) should differ; if the fused + output matches the unrotated quantization, `rotate` is a no-op -- which is + precisely the failure a tolerance-based check would wave through. + """ + fused = build_and_load(k=k, verbose=False) + x = torch.randn(64, k, dtype=torch.bfloat16, device="npu") + q_fused, _ = fused(x) + q_plain, _ = torch_npu.npu_dynamic_mx_quant(x, dst_type=VENDOR_DST_TYPE) + torch.npu.synchronize() + same = (q_fused.cpu() == q_plain.cpu()).float().mean().item() + assert same < 0.6, ( + f"K={k}: fused output matches the UNROTATED quantization at {same:.3f} " + "-- the rotation is not happening" + ) + + +def chunk_count(k): + """Shape::chunks, recomputed on the host. + + A group is one row once the row is at least a window wide, so chunks is + (k/2)/128 there and 1 below it, where several rows share a window instead. + The two cases take different slot addressing, so a matrix that covered only + one would leave the other live in production and dead in CI. + """ + upper = (k if k >= 256 else 256) // 2 + return upper // min(upper, 128) + + +def test_width_matrix_covers_both_chunk_classes(): + """The matrix must exercise chunks == 1 AND chunks > 1. + + Guards the gap itself rather than one instance of it: the original five widths + were in one chunk class, so the other path was dead code in CI while live + in production. A width added later -- or a change to TILE_ELEMS, which is what + actually happened -- must not quietly return the matrix to one class. + + The classes are derived here rather than asserted against a stored list. A + stored list is a snapshot of TILE_ELEMS, so it goes stale on a tuning change + and then reports a tile change as a width bug. When this fails it names the + supported widths that would restore coverage, since that is the fix. + """ + seen = {} + for k in WIDTHS: + rows = build_and_load(k=k, verbose=False).rows_for() + assert rows > 0, f"K={k}: rows_for returned 0" + seen.setdefault(chunk_count(k) > 1, []).append(k) + if len(seen) < 2: + missing = {} + for k in sorted(SUPPORTED_K): + cls = chunk_count(k) > 1 + if cls not in seen: + missing.setdefault(cls, []).append(k) + pytest.fail( + f"the width matrix only exercises chunked={sorted(seen)}: {seen}. " + f"Both paths are live in production. Add one of: " + f"{ {c: v[:6] for c, v in missing.items()} }" + ) + + +@pytest.mark.parametrize("k", (256, 512, 4096)) +def test_constant_row_is_a_delta(k): + """A constant row becomes ONE delta for the whole row. + + H's first column is all ones, so a constant row sums into element 0 and + cancels across all k-1 others. This is the check that separates this kernel + from a block-wise rotation: there, a constant row gives k/32 deltas, one per + block. Here anything past element 0 means the butterfly stopped short of + spanning the row. + + 256 packs rows into a window, 512 chunks by 2 and 4096 by 16, so the three + cover both slot-addressing classes. Every 32-block after the first is all + zeros, so only block 0 carries signal. + """ + fused = build_and_load(k=k, verbose=False) + x = torch.ones(8, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + # dequant returns (batch, k/32, 32); flatten so element 0 is the row's, not + # block 0's -- indexing the blocked shape takes all 32 of block 0, and 31 of + # those are legitimately zero + vals = dequant(q, s, k).reshape(8, k) + lead, rest = vals[:, 0].abs(), vals[:, 1:].abs() + assert (lead > 0).all(), "element 0 should carry the whole row's sum" + assert rest.max() <= lead.min() * 0.05, ( + f"a constant row should rotate to a single delta; leaked " + f"{rest.max():.3f} against a lead of {lead.min():.3f}. One delta per " + f"32-block instead means the rotation is still block-wide." + ) + + +def test_unsupported_k_is_rejected(): + """Widths without an instantiation must raise on the host. + + The dispatch would otherwise fall through silently and hand back the caller's + buffers untouched. + """ + for bad in (31, 33, 100, 0, 4095): + with pytest.raises((ValueError, TypeError)): + build_and_load(k=bad, verbose=False) + + +def test_wrong_dtype_is_rejected(): + fused = build_and_load(k=256, verbose=False) + for dtype in (torch.float16, torch.float32): + with pytest.raises(TypeError): + fused(torch.randn(16, 256, dtype=dtype, device="npu")) diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/README.md b/examples/jit_cpp/fused_hadamard_quant_b32_a5/README.md new file mode 100644 index 00000000..a4a5fcdb --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/README.md @@ -0,0 +1,139 @@ +# fused_hadamard_quant_b32_a5 - a block-32 Hadamard and MXFP4 in one launch + +`x -> block-32 Hadamard -> E2M1 nibbles + one E8M0 scale per 32`, as a +single kernel on the Ascend 950 / A5 (`dav-c310-vec`) vector core, JIT-compiled +with `bisheng` and loaded through `ctypes`. `K` is a template parameter over 28 +widths from 32 to 16384; one `.so` holds an instantiation per width and the +launcher dispatches on it, so there is no rebuild per size. The rotation is 32 +wide however long the row is, so the width costs nothing in registers and is +bounded only by the tile and DMA arithmetic. + +`fused_hadamard_quant_a5` is the companion that rotates the whole row instead. +Pick this one when `K` is not a power of two, or when the widest rotation is not +wanted: MXFP4's scale covers 32 elements, so a 32-wide rotation already matches +the quantizer's granularity, and on heavy-tailed data it measured 4-5% lower +quantization error than a full-row rotation at K=4096, because spreading an +outlier across the whole row lifts every block's shared scale instead of just +one block's. + +The rotation is 32 wide rather than row wide, which is what lets `K` be any +multiple of 32 instead of a power of two: a row is a run of independent 32-blocks, +the butterfly window is 256 = eight blocks, and a block never straddles a row. +The MXFP4 group is also 32, so one scale covers exactly one rotated block. + +## Fusing the pair is 2.45-2.54x the two separate launches + +Unfused, this is two passes over HBM: the butterfly writes the rotated tile out +and the quantizer reads it straight back. Fused, that tile never leaves UB and +only the nibbles and scales are written. Bytes per element tell the whole story: +6.53 unfused against 2.53 fused. + +| K | 2 launches | fused | vs 2 | rel err | spread | +|---|--:|--:|--:|--:|--:| +| 32 | 29.1 | 13.6 | 2.14x | 0.0 | 9.2% | +| 1024 | 37.5 | 27.2 | 1.38x | 0.0 | 14.5% | +| 4096 | 293.4 | 119.7 | **2.45x** | 0.0 | 3.7% | +| 16384 | 1206.1 | 474.6 | **2.54x** | 0.0 | 1.3% | + +M = 16384, microseconds per launch, what `benchmark.py` prints. Both arms agree +to a relative error of 0.0, checked before either is timed. Byte traffic +predicts 6.53 / 2.53 = 2.58x, and the two clean widths measure 2.45x and 2.54x. + +The other two rows are not traffic results, for different reasons. At K=32 a row +is 0.5M elements at M=16384 and the fused arm's 13.6 us is the dispatch floor, so +its 2.14x is two launches against one rather than anything about bytes. At K=1024 +the unfused intermediate is `2*M*k` = 32 MB against a 128 MiB L2, so the unfused +arm reads much of it from cache rather than HBM, which flatters the arm fusing is +measured against; the 14.5% bracket spread on that row is the same thing showing +up as noise. The copy section below has neither problem, since it runs 64Mi +elements whatever `K` is. + +## It runs at about the speed of a copy of its input + +| K | fused | d2d copy | vs copy | fused GB/s | copy GB/s | +|---|--:|--:|--:|--:|--:| +| 32 | 122.1 | 191.1 | **1.57x** | 1391 | 1404 | +| 1024 | 122.4 | 192.4 | **1.57x** | 1388 | 1395 | +| 4096 | 117.1 | 192.8 | **1.65x** | 1450 | 1392 | +| 16384 | 117.3 | 192.4 | **1.64x** | 1448 | 1395 | + +64Mi elements per launch. Both arms reach much the same bandwidth -- the kernel +1398-1477 GB/s against the copy's 1394-1420 -- so the kernel is not moving bytes +faster than a copy, it is moving 1.58x fewer of them: 2.53 B/element against 4.0. +The butterfly and the quantizer are both hidden under the DMA. That +also means there is nothing to win from instruction selection here; the tile size +is what mattered, and a `vsts` to `vscatter` swap with bit-identical output cost ++29 us. + +Measured on an `Ascend950PR_9589`: 64 vector cores, 128 MiB L2, 1.65 GHz, HBM +peak 1.6 TB/s, so the kernel reaches 88-91% of peak. The copy is a +reference for what moving the bytes costs, not a proven lower bound -- +it is a vendor kernel doing a simpler job. HBM peak is the closer thing +to a real ceiling, and that is the number above. Other A5 parts have +different HBM, and absolute GB/s from one part should not be compared against +another -- the ratios above are the portable numbers. + +## Correctness + +The kernel cannot be bit-exact against a torch expression: it rotates in bf16 +with a specific operand order and no torch formulation reproduces that tree. So +`test_fused_hadamard_quant_b32_a5.py` establishes it three ways, strongest first. + +1. **Scale bytes** must match a reference that rotates in fp32 and quantizes with + `torch_npu`. A scale is a power of two derived from a block maximum, so bf16 + rounding inside the butterfly almost never moves it -- disagreeing scales mean + a wrong rotation, not different rounding. Threshold 98%; measured 99.8%. +2. **Dequantized values** must track that reference to within MXFP4's own + resolution. This catches a correct-looking permutation, which a check on the + packed bytes would not. Threshold 5%; measured 0.36%. +3. **The output must be non-trivial.** A kernel that writes nothing, or writes + its input back, is the characteristic silent failure on this hardware and + would pass a loose tolerance. Separate tests assert the nibbles are neither + all-zero nor degenerate, and that the result differs from quantizing without + the rotation. + +Two structural cases are covered because both have hidden real bugs here. The +width list spans both **unroll classes** -- the butterfly unrolls by 8 or by 4 +depending on `rows_for(k) * k / 256`, they are different code paths, and class +membership is derived from the built `.so` rather than hardcoded, because raising +the tile size once moved four widths between classes and left the matrix +single-class. And one test uses a batch deep enough that every core walks +several tiles, since a shallow batch leaves the buffer rotation, the prefetch and +the drain unexercised. + +```bash +python3 -m pytest -q test_fused_hadamard_quant_b32_a5.py +``` + +## Running the benchmark + +```bash +./run_benchmark.sh # or: python3 benchmark.py --device 0 +``` + +Needs a CANN whose PTO carries MXFP4 (`Exp2DStrided` in `pto/npu/a5/TQuant.hpp`). +9.1.0 and 9.2.0 both do; 9.0.0 does not. + +## Tunables + +All compile-time, with the shipped defaults. Every combination is checked by +`static_assert`, so a tile that will not fit UB or a prefetch depth that would +deadlock fails to compile rather than misbehaving. + +| flag | default | what it is | +|---|---|---| +| `FUSED_TILE_ELEMS` | 24576 | elements per UB tile, 48 KB in bf16 | +| `FUSED_BUFFERS` | 3 | UB pipeline buffers | +| `FUSED_PREFETCH` | 2 | tiles in flight ahead | + +The same source builds the reduced kernels the ladder needs: +`FUSED_ROTATE_ONLY` drops the quantizer and `FUSED_NO_ROTATE` drops the +butterfly, so the two arms differ in what they fuse and in nothing else. +`FUSED_BUFFERS` above 4 does not build at K=4096 -- five slots need 311,040 +bytes of UB against 253,952 available. + +The butterfly is the unnormalised Sylvester matrix, so its output is `sqrt(32)` +larger than an orthogonal block Hadamard's. That is left to the caller rather +than absorbed: MXFP4's `E8M0` scale is a power of two and `sqrt(32)` is not, so +the scale cannot take it up and the nibbles would genuinely differ. Scale `x` by +`1/sqrt(32)` going in if orthogonal semantics are wanted. diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/benchmark.py b/examples/jit_cpp/fused_hadamard_quant_b32_a5/benchmark.py new file mode 100644 index 00000000..a712d9fc --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/benchmark.py @@ -0,0 +1,214 @@ +"""Fused block-32 Hadamard + MXFP4 quantize on Ascend A5, against the unfused +pair and against a device-to-device copy. + +Two comparisons: + + A. the fusion ladder. Unfused is two launches -- the Hadamard, then the + quantizer -- at 6.53 B/element, against 2.53 fused. Both arms are this + kernel, built from one source with the unwanted half compiled out, so they + differ in what they fuse and in nothing else. + + B. the fused kernel against a d2d copy of its input, as a reference for what + moving the bytes costs. Not a proven lower bound: the copy is a vendor + kernel doing a simpler job, and nothing here shows it is optimal. + +Method: wall clock on a saturated queue, medians over TRIALS brackets of +LAUNCHES launches, inputs drawn from a rotating pool so a bracket cannot be +served from cache. Every arm is checked against the unfused arm before any of +them is timed. +""" + +import argparse +import statistics +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch_npu # noqa + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from jit_util_fused_b32_a5 import ( # noqa + MX_BLOCK, + build_and_load, +) + +# At M=16384 the unfused intermediates are 2*M*k bytes. Below K=4096 that fits +# the 128 MiB L2, so the unfused arms partly read from cache and the ladder +# understates fusing -- 2.1x at K=1024 against 4.1x at K=4096. Kept in the sweep +# because the effect is worth seeing, not because those rows are the headline. +SHAPES = (32, 1024, 4096, 16384) +COPY_ELEMS = 1 << 26 +M = 16384 +TRIALS = 15 +LAUNCHES = 20 +WARMUP = 5 +POOL_BYTES = 256 * 1024 * 1024 +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) + + +def trials(call, depth, launches=LAUNCHES): + for _ in range(WARMUP): + call(0) + torch.npu.synchronize() + out = [] + for t in range(TRIALS): + torch.npu.synchronize() + t0 = time.perf_counter() + for i in range(launches): + call((t * launches + i) % depth) + torch.npu.synchronize() + out.append((time.perf_counter() - t0) * 1e6 / launches) + med = statistics.median(out) + return med, 100 * (max(out) - min(out)) / med + + +def dequant(q, s, k): + q = q.cpu() + lo, hi = q & 0x0F, (q >> 4) & 0x0F + codes = torch.stack([lo, hi], dim=-1).reshape(q.shape[0], -1) + mag = E2M1[(codes & 0x07).long()] + sign = torch.where(codes & 0x08 != 0, -1.0, 1.0) + scale = torch.exp2(s.cpu().float() - 127.0).repeat_interleave(MX_BLOCK, dim=-1) + return (mag * sign * scale).reshape(-1, k) + + +def bench_ladder(k): + """A: two launches, then one.""" + depth = max(2, min(16, POOL_BYTES // max(M * k * 4, 1))) + x = [torch.randn(M, k, dtype=torch.bfloat16, device="npu") for _ in range(depth)] + + # one source, three builds: both halves, the rotation alone, the quantizer + # alone. Same tiling, same UB layout, same buffer count in each, so the + # difference between the arms is the fusion and nothing else. + fused = build_and_load(k=k, verbose=False) + rotate_only = build_and_load( + k=k, verbose=False, extra_defs=("-DFUSED_ROTATE_ONLY",) + ) + quant = build_and_load(k=k, verbose=False, extra_defs=("-DFUSED_NO_ROTATE",)) + + q = torch.empty((M, k // 2), dtype=torch.uint8, device="npu") + s = torch.empty((M, k // MX_BLOCK), dtype=torch.uint8, device="npu") + rot = torch.empty((M, k), dtype=torch.bfloat16, device="npu") + torch.npu.synchronize() + + def two(i): # Hadamard, then quantize + rotate_only(x[i % depth], out=(rot.view(torch.uint8), s)) + quant(rot, out=(q, s)) + + def one(i): # both in one launch + fused(x[i % depth], out=(q, s)) + + # correctness gate: the arms must agree before either is timed + two(0) + torch.npu.synchronize() + ref = dequant(q.clone(), s.clone(), k) + one(0) + torch.npu.synchronize() + got = dequant(q, s, k) + rel = ((got - ref).abs().mean() / ref.abs().mean().clamp_min(1e-6)).item() + + # A disagreeing arm is a bug, not a datum: stop rather than print a table + # whose rows measure different computations. bf16 rounding differences + # between a fused and an unfused rotation land near 1e-3, not near 1. + if rel > 0.05: + raise SystemExit( + f"K={k}: arms disagree, rel={rel:.4f} -- the ladder is not measuring " + "the same computation in every arm" + ) + + t2, s2 = trials(two, depth) + t1, s1 = trials(one, depth) + x.clear() + torch.npu.empty_cache() + return { + "k": k, + "two_us": round(t2, 1), + "fused_us": round(t1, 1), + "vs_two": round(t2 / t1, 2), + "rel": round(rel, 5), + "spread_pct": round(max(s2, s1), 1), + } + + +def bench_copy(k): + """B: the fused kernel against a copy of its input.""" + batch = max(128, (COPY_ELEMS // k) // 128 * 128) + depth = max(2, min(16, POOL_BYTES // max(batch * k * 4, 1))) + x = [ + torch.randn(batch, k, dtype=torch.bfloat16, device="npu") for _ in range(depth) + ] + dst = [torch.empty_like(x[0]) for _ in range(depth)] + fused = build_and_load(k=k, verbose=False) + q = torch.empty((batch, k // 2), dtype=torch.uint8, device="npu") + s = torch.empty((batch, k // MX_BLOCK), dtype=torch.uint8, device="npu") + torch.npu.synchronize() + + tf, sf = trials(lambda i: fused(x[i % depth], out=(q, s)), depth) + tc, sc = trials(lambda i: dst[i % depth].copy_(x[i % depth]), depth) + # the kernel reads 2 B and writes 0.5 + 1/32 per element; the copy 2 and 2 + kernel_gbs = batch * k * (2 + 0.5 + 1 / MX_BLOCK) / (tf * 1e-6) / 1e9 + copy_gbs = batch * k * 4.0 / (tc * 1e-6) / 1e9 + x.clear() + dst.clear() + torch.npu.empty_cache() + return { + "k": k, + "batch": batch, + "fused_us": round(tf, 1), + "copy_us": round(tc, 1), + "vs_copy": round(tc / tf, 2), + "fused_gbs": round(kernel_gbs), + "copy_gbs": round(copy_gbs), + "spread_pct": round(max(sf, sc), 1), + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--device", type=int, default=0) + args = ap.parse_args() + torch.npu.set_device(args.device) + torch.manual_seed(20260826) + np.random.seed(0) + + print(f"=== A. fusion ladder, M={M} (microseconds per launch) ===") + print( + f"{'K':>7} {'2 launches':>11} {'fused':>8} {'vs 2':>6} {'rel':>8} " + f"{'spread':>7}" + ) + for k in SHAPES: + try: + r = bench_ladder(k) + except (RuntimeError, SystemExit) as exc: + print(f"{k:>7} skipped: {str(exc)[:56]}") + continue + print( + f"{r['k']:>7} {r['two_us']:>11.1f} {r['fused_us']:>8.1f} " + f"{r['vs_two']:>5.2f}x {r['rel']:>8.4f} {r['spread_pct']:>6.1f}%" + ) + + print(f"\n=== B. fused vs a d2d copy ({COPY_ELEMS // 1024}Ki elements) ===") + print( + f"{'K':>7} {'batch':>8} {'fused':>8} {'copy':>8} {'vs copy':>8} " + f"{'fused GB/s':>11} {'copy GB/s':>10} {'spread':>7}" + ) + for k in SHAPES: + try: + r = bench_copy(k) + except RuntimeError as exc: + print(f"{k:>7} skipped: {str(exc)[:56]}") + continue + print( + f"{r['k']:>7} {r['batch']:>8} {r['fused_us']:>8.1f} {r['copy_us']:>8.1f} " + f"{r['vs_copy']:>7.2f}x {r['fused_gbs']:>11} {r['copy_gbs']:>10} " + f"{r['spread_pct']:>6.1f}%" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/fused_hadamard_quant_b32_a5.cpp b/examples/jit_cpp/fused_hadamard_quant_b32_a5/fused_hadamard_quant_b32_a5.cpp new file mode 100644 index 00000000..d334295e --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/fused_hadamard_quant_b32_a5.cpp @@ -0,0 +1,886 @@ +// Block-32 Hadamard fused with MXFP4 quantization, one launch. +// +// x -> (x @ H) -> E2M1 nibbles + one E8M0 scale per 32 +// +// Unfused this is two passes over HBM: read x / write rotated, then read +// rotated / write nibbles+scales. Fused it is read x / write nibbles+scales, so +// on a DMA-bound op the saving is close to the whole second pass. +// +// Built from two kernels that are already measured and merged upstream: +// fast_hadamard_a5 supplies the butterfly, mxfp4_quant_a5 the four quant +// passes, the tiling and the outputs. Both are left doing what they already do; +// what is new here is that the rotated tile never leaves UB. +// +// The butterfly was fp16 upstream and is bf16 here, which costs nothing +// structurally: vlds/vsts are bit-width ops on vector_u16 (DINTLV_B16 / +// NORM_B16), so only the arithmetic type changes, by reference cast. That is +// the same idiom mxfp4_quant_a5 already uses for its max reduction. +// +// The difference from v1 is not just a pinned width. v1 rotates a whole row, so +// K must be a power of two and at most 2048. Here the Hadamard is always 32 +// wide and a row is a sequence of independent 32-blocks, which decouples the +// rotation from the row width: K goes back to any multiple of 32, so 4096 and +// 11008 work, which a row-wide rotation rejects. +// +// It falls out of the tile being a flat run of Rows*K elements. Blocks are +// contiguous and 32 long, the butterfly window is 256 = eight blocks, and +// blocks are independent -- so one window covers eight of them and never has to +// care whether they came from the same row. +// +// Also: the MXFP4 group is 32 and the Hadamard block is 32, so a scale covers +// exactly one rotated block. No reshaping, and no group straddling a rotation. +#include +#include +#include +using namespace pto; + +// Row widths with an instantiation. The full set the quantizer supports: a +// 32-wide rotation puts no power-of-two constraint on the row, so 4096 and +// 11008-style widths are back. Also add any new width to the jit helper. +constexpr unsigned SUPPORTED_K[] = {32, 64, 96, 128, 192, 256, 512, + 768, 896, 1024, 1152, 1280, 1408, 1536, + 1664, 1792, 2048, 2560, 2816, 3072, 3584, + 4096, 5120, 6144, 7168, 8192, 14336, 16384}; +constexpr unsigned SUPPORTED_COUNT = + sizeof(SUPPORTED_K) / sizeof(SUPPORTED_K[0]); +// --- butterfly geometry, from fast_hadamard_a5 ------------------------------- +// WINDOW is two registers: the deinterleave load splits a 2*lanes run into +// even/odd halves, and the concat-halves store puts them back. +constexpr unsigned SLOTS = 8; // unroll width: register sets per sweep +constexpr unsigned HAD_ALIGN = 512; +constexpr unsigned HAD_BLOCK = 32; // the Hadamard block, == MX_BLOCK + +constexpr unsigned MX_BLOCK = 32; // MXFP4 block: 32 elements, one E8M0 scale + +// The three pipeline parameters, overridable for tuning. The defaults are the +// tuned point: 24576 is the largest tile that fits UB at all, and 3 is the only +// buffer count it fits at. It beats the 16384/4/2 this kernel shipped with by +// 1.063-1.066x on large launches, bit-exact, with no regression at any shape +// measured. Numbers and the full grid are in the README. +// +// Overriding is safe in the way that matters: every combination is checked by +// the static_asserts at the end of QuantShape, so a tile that will not fit UB, +// or a prefetch depth that would deadlock, fails to COMPILE rather than +// misbehaving. And the host reads rows-per-tile back from the .so +// (hadamard_mxfp4_b32_rows_for), so a changed TILE_ELEMS cannot desynchronise +// from the harness. +#ifndef FUSED_BUFFERS +#define FUSED_BUFFERS 3 +#endif +#ifndef FUSED_PREFETCH +#define FUSED_PREFETCH 2 +#endif +#ifndef FUSED_TILE_ELEMS +#define FUSED_TILE_ELEMS 24576 // 48 KB bf16 +#endif + +// Store the butterfly halves with vscatter instead of vsts. +// +// 0 vsts NORM_B16 pair -- what ships, and the default. +// 1 vscatter with an IDENTITY index. Same instruction count, same registers, +// same dependency chain, bit-identical output; the ONLY difference is the +// opcode. This exists to price vscatter against vsts, because vscatter's +// cost on A5 is not documented and the vendor's bf16 TTRANS is built from +// it -- and two TTRANS calls are 96% of the pure-PTO kernel's cost, so it +// could be far dearer than a plain store. +// +// Mode 2 was a ROL5 index meant to absorb the rotation fixup into the store at +// no extra instruction -- the fixup is 13.8% of the kernel (82.74 -> 71.31 us +// with -DFUSED_NO_ROTFIX, paired 1.164x, resolved). It is GONE, refuted by +// mode 1: vscatter costs +28.96 us per call against vsts (111.57 against 82.61, +// paired 0.741x, resolved), so the opcode swap alone costs 2.5x what the fixup +// it would remove is worth. Break-even needed vscatter under ~5.5x a vsts. +// +// This also explains the pure-PTO kernel: its two TTRANS calls are 96% of its +// 2052 us, and the vendor builds bf16 TTRANS out of vgather2/vscatter. + +#ifndef FUSED_SCATTER +#define FUSED_SCATTER 0 +#endif +#if FUSED_SCATTER > 1 +#error "FUSED_SCATTER=2 (ROL5 store) was measured and refuted; see above" +#endif + +constexpr unsigned DEF_BUFFERS = FUSED_BUFFERS; // UB pipeline buffers +constexpr unsigned DEF_PREFETCH = FUSED_PREFETCH; // tiles in flight ahead +constexpr unsigned TILE_ELEMS = FUSED_TILE_ELEMS; +// RULE: every GM move_tile is one row and a Tile refuses a row under 32 bytes. +// The scale row is the smallest, at tile_elems/32 bytes, so a tile must be a +// whole multiple of 32*MX_BLOCK elements. DMA sets this grain, not the compute. +constexpr unsigned TILE_GRAIN = 1024; +// ROWS_PER_TILE: the largest row count whose tile is a whole number of grains. +// Not TILE_ELEMS / K -- for a large odd factor (768 = 32*24) the quotient is +// not a multiple of the grain. Zero means inadmissible; Rows asserts on it. +template +struct Gcd { + static constexpr unsigned value = Gcd::value; +}; +template +struct Gcd { + static constexpr unsigned value = A; +}; + +// Rows*K is a multiple of TILE_GRAIN exactly when Rows is a multiple of +// TILE_GRAIN / gcd(K, TILE_GRAIN), so the answer is the largest such multiple +// within cap. Counting down from cap one step at a time costs a template +// instantiation per step, which is 768 of them at K=32 and would grow past the +// compiler's depth limit if the tile were raised. +template +struct RowsFor { + static constexpr unsigned cap = TILE_ELEMS / K > 1u ? TILE_ELEMS / K : 1u; + static constexpr unsigned step = TILE_GRAIN / Gcd::value; + static constexpr unsigned value = (cap / step) * step; +}; + +#if defined(FUSED_ROTATE_ONLY) && defined(MXFP4_TQUANT) +#error "FUSED_ROTATE_ONLY does nothing in a TQuant build: TQuant owns them" +#endif + +#ifdef __CCE_AICORE__ +constexpr unsigned B16_LANES = 128; // bf16 lanes in one vector register +// vcgmax on b16 groups 16 lanes, 8 results in lanes 0..7. +constexpr unsigned VCGMAX_B16_GROUP = 16; +constexpr unsigned VCGMAX_B16_RESULTS = B16_LANES / VCGMAX_B16_GROUP; +static_assert(VCGMAX_B16_RESULTS == 8, "block_abs_max stores with PAT_VL8"); +// RULE: vsts needs a 32-byte-aligned UB address, else 507035. Tile refuses a +// sub-32-byte DMA, so the padding is squeezed out in UB, not on the way to GM. +constexpr unsigned VSTS_ALIGN = 32; +constexpr unsigned GROUP_PITCH_B16 = VSTS_ALIGN / 2u; // in b16 elements +// RULE: vselr indices reach only the low 128 source bytes: 4 groups per gather. +constexpr unsigned GROUPS_PER_COMPACT = 4; +constexpr unsigned EVENT_SLOTS = 8; +static_assert(EVENT_SLOTS == 8, "extend buffer_free's initialiser first"); +constexpr unsigned UB_ALIGN = 512; +// A5 has 256 KB. +constexpr unsigned UB_BYTES = PTO_UBUF_SIZE_BYTES; + +// bf16 bit-field constants. bf16 is 1-8-7, so a magnitude's biased exponent is +// simply bits >> 7 once the sign is cleared. +constexpr uint16_t BF16_ABS = 0x7FFFu; // clears the sign bit +constexpr int16_t BF16_MANT_BITS = 7; +constexpr int16_t E8M0_BIAS_ADJ = -2; // byte = b - 2 (Algorithm 1, FLOOR) +constexpr int16_t RECIP_OFFSET = 256; // 1/X exponent field = 256 - b +// b must stay in a window where 1/X is finite, non-subnormal bf16: field 256-b +// must land in [2, 254]. Clamp b, then derive BOTH outputs from the clamped b. +constexpr int16_t B_MIN = 2; +constexpr int16_t B_MAX = 254; + +// RULE: a constexpr function cannot be called from [aicore] code. +template +struct Log2 { + static constexpr unsigned value = 1 + Log2<(Value >> 1)>::value; +}; +template <> +struct Log2<1> { + static constexpr unsigned value = 0; +}; + +// Largest unroll width <= Limit that divides Windows exactly, so the sweep +// covers the tile with no tail. A value template rather than a constexpr +// function: [aicore] code may not call one, even to initialise a constexpr. +template +struct UnrollFor { + static constexpr unsigned value = + (Windows % Limit == 0) ? Limit : UnrollFor::value; +}; +template +struct UnrollFor { + static constexpr unsigned value = 1u; +}; + +template +struct RoundUp { + static constexpr unsigned value = (Bytes + UB_ALIGN - 1) & ~(UB_ALIGN - 1); +}; + +// Every derived size for one instantiation. +template +struct QuantShape { + static constexpr unsigned tile_elems = Rows * K; + static constexpr unsigned blocks = tile_elems / MX_BLOCK; + static constexpr unsigned in_bytes = tile_elems * 2u; + static constexpr unsigned q_bytes = tile_elems / 2u; + static constexpr unsigned scale_bytes = blocks; + + // One "group" is one vcgmax: 8 blocks == 256 elements. + static constexpr unsigned groups = blocks / VCGMAX_B16_RESULTS; + // Round UP: the last bite may be partial. Safe only because the buffers + // below are sized from these counts, and nothing reads what a partial bite + // writes past `blocks`. The asserts at the end of this struct pin both. + static constexpr unsigned compact_iters = + (groups + GROUPS_PER_COMPACT - 1u) / GROUPS_PER_COMPACT; + static constexpr unsigned b_iters = (blocks + B16_LANES - 1u) / B16_LANES; + static constexpr unsigned c_iters = tile_elems / (2u * B16_LANES); + + // maxima_bytes carries a register of read-ahead: the gather reads one + // register past its last input offset. The rest are sized from what the loops + // WRITE, since the two rounded counts can exceed their data by one bite. + static constexpr unsigned maxima_bytes = + compact_iters * GROUPS_PER_COMPACT * GROUP_PITCH_B16 * 2u + B16_LANES; + static constexpr unsigned packed_bytes = + compact_iters * GROUPS_PER_COMPACT * VCGMAX_B16_RESULTS * 2u; + static constexpr unsigned aligned_in = RoundUp::value; + static constexpr unsigned aligned_q = RoundUp::value; + static constexpr unsigned aligned_s = RoundUp::value; + static constexpr unsigned aligned_max = RoundUp::value; + static constexpr unsigned aligned_packed = RoundUp::value; + static constexpr unsigned aligned_mult = + RoundUp::value; + + static constexpr unsigned slot_stride = aligned_in + aligned_q + aligned_s; + static constexpr unsigned scratch_base = NBuffers * slot_stride; + // every scratch region SlotOffset hands out, in the same order, so the two + // cannot drift: omitting one here silently shrinks the UB-overflow guard + static constexpr unsigned ub_needed = + scratch_base + aligned_max + aligned_packed + aligned_mult; + +#ifdef MXFP4_TQUANT + // TQuant reads its per-block maxima a whole register at a time, so the span + // rounds up to b_iters registers, and its reducer flushes one 32-byte block + // past the last group. + static constexpr unsigned tquant_max_elems = b_iters * B16_LANES; + static constexpr unsigned tquant_max_bytes = + tquant_max_elems * 2u + VSTS_ALIGN; + static constexpr unsigned tquant_scaling_bytes = blocks * 2u; + static_assert(tquant_max_bytes <= aligned_max, + "TQuant maxima do not fit the maxima region"); + static_assert(tquant_scaling_bytes <= aligned_mult, + "TQuant scaling does not fit the reciprocal region"); + // numGroups truncates inside TQuant, and a partial 8-group window takes a + // different store path. + static_assert(tile_elems % MX_BLOCK == 0, "TQuant would drop a group"); + static_assert(blocks % 8u == 0, "TQuant would take its vstus tail path"); +#endif + + // --- butterfly geometry over 32-element blocks + // ------------------------------ Fixed at HAD_BLOCK, independent of K: the + // tile is a flat run of Rows*K elements, so it is (Rows*K)/32 independent + // blocks and the window packs eight of them. + static constexpr unsigned had_lanes_b16 = B16_LANES; + static constexpr unsigned had_window = 2u * had_lanes_b16; + static constexpr unsigned log2_block = Log2::value; + static constexpr unsigned log2_window = Log2::value; + static constexpr unsigned rows_per_window = had_window / HAD_BLOCK; + static constexpr unsigned had_group = HAD_BLOCK * rows_per_window; + static constexpr unsigned rotations = log2_window - log2_block; + static constexpr unsigned upper = had_group / 2u; + static constexpr unsigned lanes = + upper < had_lanes_b16 ? upper : had_lanes_b16; + static constexpr unsigned chunks = upper / lanes; + static constexpr unsigned windows_per_tile = tile_elems / had_group; + static constexpr unsigned groups_per_iter = + UnrollFor::value; + static constexpr unsigned had_blocks_per_tile = tile_elems / HAD_BLOCK; + static constexpr unsigned had_iters = + tile_elems / had_group / groups_per_iter; + static constexpr unsigned sweep_stride = groups_per_iter * had_group; + // How many register slots one sweep call must use. This has to be derived + // from groups_per_iter, NOT from SLOTS. + // + // A slot addresses window `Slot / chunks` at chunk `Slot % chunks`, so a call + // with N slots covers N/chunks windows, while the loop advances + // groups_per_iter windows per iteration. Instantiating the sweep with SLOTS + // when groups_per_iter is smaller makes consecutive iterations overlap and + // runs the last one off the end of the tile. + // + // The row-wide variant cannot hit this because it defines + // groups_per_iter = SLOTS / chunks, which makes the two agree by + // construction. This kernel picks groups_per_iter with UnrollFor instead -- + // deliberately, because a 32-wide rotation's window count need not be a + // multiple of 8 -- and that is exactly what decoupled them. + static constexpr unsigned sweep_slots = groups_per_iter * chunks; + + static_assert(HAD_BLOCK == MX_BLOCK, + "a scale must cover exactly one rotated block"); + static_assert(had_window % HAD_BLOCK == 0, + "the butterfly window must be whole blocks"); + static_assert(tile_elems % had_group == 0, + "tile must be a whole number of butterfly windows"); + static_assert(windows_per_tile % groups_per_iter == 0, + "UnrollFor must divide the window count exactly"); + static_assert(chunks == 1u, "a 32-wide rotation packs; it never chunks"); + static_assert(sweep_slots <= SLOTS, + "a sweep would need more register slots than SLOTS declares"); + // The tiling condition the overlap bug violated: one iteration's slots must + // cover exactly the windows the stride advances, no more and no fewer. + static_assert(sweep_slots / chunks * had_group == sweep_stride, + "sweep slots and sweep_stride disagree: iterations overlap"); + + static_assert(Rows > 0, "no Rows makes Rows*K a whole TILE_GRAIN: bad K"); + static_assert(K % MX_BLOCK == 0, "a block may not straddle a row boundary"); + static_assert(TILE_GRAIN == VSTS_ALIGN * MX_BLOCK, "grain != scale DMA row"); + static_assert(tile_elems % TILE_GRAIN == 0, "tile is not a whole grain"); + static_assert(tile_elems % (2u * B16_LANES) == 0, "pack_nibbles wants 256"); + static_assert(scale_bytes % VSTS_ALIGN == 0, "scale row is not a legal DMA"); + static_assert(q_bytes % VSTS_ALIGN == 0, "nibble row is not a legal DMA"); + static_assert(in_bytes % VSTS_ALIGN == 0, "input row is not a legal DMA"); + static_assert(blocks % VCGMAX_B16_RESULTS == 0, + "blocks != whole vcgmax groups"); + // the rounded-up passes run one bite past the data; prove they stay inside + static_assert(b_iters * B16_LANES <= aligned_s, "scale tail overruns"); + static_assert(b_iters * B16_LANES * 2u <= aligned_mult, + "recips tail overruns"); + static_assert(packed_bytes <= aligned_packed, "compaction tail overruns"); + static_assert(groups * VSTS_ALIGN <= aligned_max, "padded maxima overrun"); + static_assert( + c_iters * VCGMAX_B16_RESULTS <= b_iters * B16_LANES, + "pack_nibbles would index recips past what derive_scales wrote"); + static_assert(sizeof(bfloat16_t) == 2, "RowsFor assumes 2-byte elements"); + static_assert(ub_needed <= UB_BYTES, "UB overflow"); + // Strictly less, not <=. The once-per-launch D load signals on EVENT_ID7 over + // MTE2 -> V, and buffer_free[7] is also EVENT_ID7 on that same pipe pair, so + // at NBuffers == EVENT_SLOTS buffer 7 and the D preamble would share a + // channel. Unreachable at the shipped NBuffers of 3; this stops it becoming + // reachable. + static_assert(NBuffers < EVENT_SLOTS, + "NBUF must leave EVENT_ID7 free for the D preamble"); + static_assert(NPrefetch < NBuffers, + "PREFETCH == NBUF deadlocks the pipeline"); +}; + +// Byte offsets within a pipeline slot, plus shared scratch. Constexpr +// *variables* for the reason above, so the slot base is multiplied in at +// the use site. +template +struct SlotOffset { + static constexpr unsigned input = 0; + static constexpr unsigned nibbles = Shape::aligned_in; + static constexpr unsigned scales = Shape::aligned_in + Shape::aligned_q; + static constexpr unsigned maxima = Shape::scratch_base; + static constexpr unsigned packed = Shape::scratch_base + Shape::aligned_max; + static constexpr unsigned reciprocal = packed + Shape::aligned_packed; +}; + +#ifdef __DAV_VEC__ +// A flat run of Elems values in GM, and the matching UB tile, for one dtype. +template +using GmShape = pto::Shape<1, 1, 1, 1, Elems>; +template +using GmStride = pto::Stride<1, 1, 1, Elems, 1>; +template +using UbTile = Tile; +// Same tile with a RUNTIME valid column count, zero-filling the rest in UB, for +// the one partial tile a batch can end on. +template +using UbTilePart = + Tile; + +// ------------------------------------------------------- block_abs_max +// Per-32-element magnitude max. A 2:1 fold makes 16 lanes == one block, +// which is what vcgmax's group size requires. A 4:1 fold silently reports +// max(block 2j, block 2j+1) instead. +// --- the rotation ------------------------------------------------------------ +// One sweep: deinterleave-load a window, add/sub, and store the halves back +// concatenated. Registers are vector_u16 because vlds/vsts are bit-width ops; +// the arithmetic type is chosen by reference cast, which is how bf16 costs +// nothing here. All loads precede all stores, which the comma fold guarantees +// by evaluating left to right -- required, not stylistic, since a store would +// otherwise clobber a window a later load still needs. +using HadRegs = vector_u16[SLOTS]; + +template +inline AICORE void sweep(__ubuf__ uint16_t *tile, uint32_t base, MaskReg all, + HadRegs &even, HadRegs &odd, HadRegs &sum, + HadRegs &diff, vector_u16 &idx_lo, vector_u16 &idx_hi, + std::index_sequence) { + constexpr unsigned g = Shape::had_group, up = Shape::upper; + constexpr unsigned ln = Shape::lanes, ch = Shape::chunks; + (vlds(even[Slot], odd[Slot], + tile + base + Slot / ch * g + Slot % ch * 2u * ln, 0, DINTLV_B16), + ...); + (vadd((vector_bf16 &)sum[Slot], (vector_bf16 &)even[Slot], + (vector_bf16 &)odd[Slot], all), + ...); + (vsub((vector_bf16 &)diff[Slot], (vector_bf16 &)even[Slot], + (vector_bf16 &)odd[Slot], all), + ...); + // A 256-element window packs eight independent 32-blocks, which leaves the + // result rotated right by log2(window/block) = 3; these register-only + // deinterleaves undo it, fused into the final stage using the pair that is + // dead by then. + // + // ATTRIBUTION SWITCH. vdintlv was measured at ~20x a vadd, and there are + // Rotations per slot here against five arithmetic ops, so this fixup may be + // most of the butterfly's cost. -DFUSED_NO_ROTFIX drops the deinterleaves and + // keeps everything else, including both stores, so the difference is the + // fixup. It PRODUCES WRONG OUTPUT -- the registers selected below then hold + // loaded values rather than rotated ones -- so it is for timing only and the + // benchmark's correctness gate will reject it. +#ifndef FUSED_NO_ROTFIX + if constexpr (Rotations >= 1) { + (vdintlv(even[Slot], odd[Slot], sum[Slot], diff[Slot]), ...); + } + if constexpr (Rotations >= 2) { + (vdintlv(sum[Slot], diff[Slot], even[Slot], odd[Slot]), ...); + } + if constexpr (Rotations >= 3) { + (vdintlv(even[Slot], odd[Slot], sum[Slot], diff[Slot]), ...); + } +#endif + HadRegs &lo = (Rotations % 2 == 1) ? even : sum; + HadRegs &hi = (Rotations % 2 == 1) ? odd : diff; +#if FUSED_SCATTER == 0 + (vsts(lo[Slot], tile + base + Slot / ch * g + Slot % ch * ln, 0, NORM_B16, + all), + ...); + (vsts(hi[Slot], tile + base + Slot / ch * g + up + Slot % ch * ln, 0, + NORM_B16, all), + ...); +#else + // vscatter instead of vsts, one for one. The index decides which experiment + // this is; see the FUSED_SCATTER comment at the top of the file. Both halves + // address the SAME window base, because with a permuting index the upper half + // is no longer a contiguous run at +upper. + (vscatter(lo[Slot], tile + base + Slot / ch * g, idx_lo, all), ...); + (vscatter(hi[Slot], tile + base + Slot / ch * g, idx_hi, all), ...); + (void)up; +#endif +} + +// log2(K) stages over the tile already in UB, in place. The quant passes read +// the same buffer straight afterwards, which is the point of the fusion. +template +__tf__ static AICORE void rotate(__ubuf__ uint16_t *tile) { + // sweep_slots, not SLOTS: see the derivation in Shape. The register arrays + // are sized SLOTS and a shorter pack leaves the top ones unused. + constexpr auto slots = std::make_index_sequence{}; + constexpr unsigned plain = Shape::log2_block - (Shape::rotations ? 1u : 0u); + __VEC_SCOPE__ { + uint32_t lane_count = Shape::lanes; + MaskReg all = CreatePredicate(lane_count); + vector_u16 even[SLOTS], odd[SLOTS], sum[SLOTS], diff[SLOTS]; + // Scatter indices, built once per tile and unused when FUSED_SCATTER == 0. + // There is no vands, so the low three bits of the lane come out as + // l - (l >> 3) * 8. vshrs/vmuls/vadds are vector-scalar; vadd/vsub are + // vector-vector. + vector_u16 idx_lo, idx_hi; + vci((vector_s16 &)idx_lo, (int16_t)0, INC_ORDER); +#if FUSED_SCATTER == 1 + // IDENTITY index: byte-for-byte what the vsts pair does, so the output must + // be bit-identical and the only difference measured is the opcode price. + vdup(idx_hi, (uint16_t)Shape::upper, all, MODE_ZEROING); + vadd(idx_hi, idx_lo, idx_hi, all); +#endif + // Step by a literal 1 with the stride folded into base: the loop analyser + // only verifies a tripcount for a literal step, and 1 divides any bound, so + // had_iters may be template-dependent. + for (uint16_t stage = 0; stage < (uint16_t)plain; ++stage) { + for (uint16_t iter = 0; iter < (uint16_t)Shape::had_iters; ++iter) + sweep(tile, (uint32_t)iter * Shape::sweep_stride, all, even, + odd, sum, diff, idx_lo, idx_hi, slots); + mem_bar(VST_VLD); + } + if constexpr (Shape::rotations > 0) { + for (uint16_t iter = 0; iter < (uint16_t)Shape::had_iters; ++iter) + sweep( + tile, (uint32_t)iter * Shape::sweep_stride, all, even, odd, sum, + diff, idx_lo, idx_hi, slots); + mem_bar(VST_VLD); + } + } +} + +template +__tf__ static AICORE void block_abs_max(__ubuf__ uint16_t *input, + __ubuf__ uint16_t *maxima) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + // PAT_VL8 matches VCGMAX_B16_RESULTS + MaskReg low_eight = pset_b16(PAT_VL8); + vector_u16 abs_mask; + vdup(abs_mask, BF16_ABS, all_lanes, MODE_ZEROING); + + for (uint16_t group = 0; group < (uint16_t)Shape::groups; ++group) { + const uint32_t base = (uint32_t)group * 256u; + vector_u16 even, odd, folded, grouped; + vlds(even, odd, input + base, 0, + DINTLV_B16); // lane i: elements 2i, 2i+1 + vand(even, even, abs_mask, all_lanes); + vand(odd, odd, abs_mask, all_lanes); + // sign cleared, so a signed max over the bit patterns IS a magnitude max + vmax((vector_s16 &)folded, (vector_s16 &)even, (vector_s16 &)odd, + all_lanes); + vcgmax((vector_s16 &)grouped, (vector_s16 &)folded, all_lanes); + // 32-byte pitch, not 16: see VSTS_ALIGN + vsts(grouped, maxima + (uint32_t)group * GROUP_PITCH_B16, 0, NORM_B16, + low_eight); + } + mem_bar(VST_VLD); + } +} + +// ------------------------------------------------------ compact_maxima +// Squeeze out the padding VSTS_ALIGN forces: output byte i takes input byte +// 2*(i & 0xF0) + (i & 0x0F). +template +__tf__ static AICORE void compact_maxima(__ubuf__ uint16_t *padded, + __ubuf__ uint16_t *packed) { + __VEC_SCOPE__ { + MaskReg all_byte_lanes = pset_b8(PAT_ALL); + MaskReg low_32 = pset_b16(PAT_VL32); // 32 b16 == 64 bytes + vector_u8 byte_index, high_half, low_half, high_mask, low_mask; + vci((vector_s8 &)byte_index, (int8_t)0, INC_ORDER); + vdup(high_mask, (uint8_t)0xF0, all_byte_lanes, MODE_ZEROING); + vdup(low_mask, (uint8_t)0x0F, all_byte_lanes, MODE_ZEROING); + vand(high_half, byte_index, high_mask, all_byte_lanes); + vand(low_half, byte_index, low_mask, all_byte_lanes); + vadd((vector_s8 &)high_half, (vector_s8 &)high_half, (vector_s8 &)high_half, + all_byte_lanes); // 2*high_half + vadd((vector_s8 &)byte_index, (vector_s8 &)high_half, (vector_s8 &)low_half, + all_byte_lanes); + + for (uint16_t gather = 0; gather < (uint16_t)Shape::compact_iters; + ++gather) { + vector_u16 padded_chunk, packed_chunk; + const uint32_t src_offset = + (uint32_t)gather * GROUPS_PER_COMPACT * GROUP_PITCH_B16; + vlds(padded_chunk, padded + src_offset, 0, NORM); + vselr((vector_u8 &)packed_chunk, (vector_u8 &)padded_chunk, byte_index); + vsts(packed_chunk, + packed + (uint32_t)gather * GROUPS_PER_COMPACT * VCGMAX_B16_RESULTS, + 0, NORM_B16, low_32); + } + mem_bar(VST_VLD); + } +} + +// -------------------------------------------------------- derive_scales +// maxima -> E8M0 scale byte + one bf16 reciprocal per block. pack_nibbles +// reads this array with E2B_B16, whose x16 replication matches its +// pair-granular deinterleave exactly, so no duplication is needed here. +template +__tf__ static AICORE void derive_scales(__ubuf__ uint16_t *maxima, + __ubuf__ uint16_t *recips_out, + __ubuf__ uint16_t *scale_out) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + vector_u16 bias; + vdup(bias, (uint16_t)RECIP_OFFSET, all_lanes, MODE_ZEROING); + + for (uint16_t chunk = 0; chunk < (uint16_t)Shape::b_iters; ++chunk) { + vector_u16 block_max, exponent, scale_byte, reciprocal; + vlds(block_max, maxima + (uint32_t)chunk * B16_LANES, 0, NORM); + // bit 15 is already clear, so this shift alone yields the biased exponent + vshrs(exponent, block_max, BF16_MANT_BITS, all_lanes, MODE_ZEROING); + vmaxs(exponent, exponent, B_MIN, all_lanes); + vmins(exponent, exponent, B_MAX, all_lanes); + vadds(scale_byte, exponent, E8M0_BIAS_ADJ, all_lanes); + vsts(scale_byte, scale_out + (uint32_t)chunk * 64u, 0, PK_B16, all_lanes); + vsub(reciprocal, bias, exponent, all_lanes); + vshls(reciprocal, reciprocal, BF16_MANT_BITS, all_lanes, MODE_ZEROING); + vsts(reciprocal, recips_out + (uint32_t)chunk * B16_LANES, 0, NORM_B16, + all_lanes); + } + mem_bar(VST_VLD); + } +} + +// --------------------------------------------------------- pack_nibbles +// Scale, cast, pack -- 256 elements per iteration, no gather. +// One vcvt puts 64 bytes at byte STRIDE 4, offset chosen by +// PART_P0..P3, so converting two halves into offsets 0 and 1, OR-ing, and +// storing with PK_B32 (keeps the low 2 bytes of each 4-byte group) writes 128 +// CONTIGUOUS bytes. RULE: fp4 packs two elements per byte, so DINTLV_B16 would +// pair element 4k with 4k+2 -- deinterleave at b32 (pairs) to keep (4k, 4k+1) +// together. That also puts both b16 lanes of a half in block j/8, so E2B_B16's +// x16 replication is exact and one multiplier register serves both halves. +template +__tf__ static AICORE void pack_nibbles(__ubuf__ uint16_t *input, + __ubuf__ uint16_t *reciprocal, + __ubuf__ uint8_t *nibble_out) { + __VEC_SCOPE__ { + MaskReg all_lanes = pset_b16(PAT_ALL); + MaskReg all_byte_lanes = pset_b8(PAT_ALL); + MaskReg all_b32_lanes = pset_b32(PAT_ALL); + + for (uint16_t chunk = 0; chunk < (uint16_t)Shape::c_iters; ++chunk) { + vector_u16 recips; + vector_u32 even, odd; + vector_bf16 scaled_even, scaled_odd; + vector_f4e2m1x2 packed_even, packed_odd, packed; + vlds(recips, reciprocal + (uint32_t)chunk * VCGMAX_B16_RESULTS, 0, + E2B_B16); + vlds(even, odd, (__ubuf__ uint32_t *)input + (uint32_t)chunk * B16_LANES, + 0, DINTLV_B32); + vmul(scaled_even, (vector_bf16 &)even, (vector_bf16 &)recips, all_lanes); + vmul(scaled_odd, (vector_bf16 &)odd, (vector_bf16 &)recips, all_lanes); + vcvt(packed_even, scaled_even, all_lanes, ROUND_R, PART_P0); + vcvt(packed_odd, scaled_odd, all_lanes, ROUND_R, PART_P1); + vor((vector_u8 &)packed, (vector_u8 &)packed_even, + (vector_u8 &)packed_odd, all_byte_lanes); + // 256 elements in, but PK_B32 keeps 2 of every 4 bytes: 128 bytes out + vsts((vector_u16 &)packed, + (__ubuf__ uint16_t *)(nibble_out + (uint32_t)chunk * B16_LANES), 0, + PK_B32, all_b32_lanes); + } + mem_bar(VST_VLD); + } +} + +#ifdef MXFP4_TQUANT +// Requires PTO 9.1.0: 9.0.0 has no MXFP4 quantizer. Included here, not at file +// scope, because this region is inside the device-pass guard. +#include + +// ------------------------------------------------------- tquant_passes +// One vendor tile op in place of block_abs_max, compact_maxima, derive_scales +// and pack_nibbles. validCols is tile_elems even on the partial tile: the load +// already zero-fills the pad, and a short validCols would send TQuant's own +// ZeroPadSourceTile over the input slot. Offsets::packed is left allocated and +// unused, since reclaiming it would move slot_stride. +template +inline AICORE void tquant_passes(uint32_t input_offset, uint32_t nibble_offset, + uint32_t scale_offset) { + static_assert(sizeof(float4_e2m1x2_t) == 1, + "the nibble tile assumes one byte per float4_e2m1x2_t"); + static_assert(REPEAT_BYTE / sizeof(bfloat16_t) == B16_LANES, + "tquant_max_elems assumes a 128-lane b16 vector"); + UbTile source; + UbTile nibbles; + UbTile scales; + UbTile block_max; + UbTile reciprocal; + TASSIGN(source, input_offset); + TASSIGN(nibbles, nibble_offset); + TASSIGN(scales, scale_offset); + TASSIGN(block_max, SlotOffset::maxima); + TASSIGN(reciprocal, SlotOffset::reciprocal); + // TEMPLATE order is Out, Src, Exp, Max, Scaling; ARGUMENT order is dst, exp, + // max, scaling, src. PTO 9.1.0 release inserted a `bool Exp2DStrided` second + // template parameter that 9.1.0-beta.3 does not have; the tile types are in a + // non-deduced position, so neither spelling can be dropped. benchmark.py + // compiles both and keeps whichever the local headers accept. +#ifdef MXFP4_TQUANT_EXP2D + TQuant_MXFP4_E2M1_Impl( + nibbles.data(), scales.data(), block_max.data(), reciprocal.data(), + source.data(), 1u, Shape::tile_elems); +#else + TQuant_MXFP4_E2M1_Impl( + nibbles.data(), scales.data(), block_max.data(), reciprocal.data(), + source.data(), 1u, Shape::tile_elems); +#endif +} +#endif // MXFP4_TQUANT + +// Move one tile of `T` between GM and UB. Partial carries only `valid` +// elements: the load zero-fills the rest of the UB tile so the compute passes +// still see whole registers, and the store truncates so padding never reaches +// GM. +template +inline AICORE void move_tile(uint32_t tile_index, uint32_t ub_offset, + __gm__ void *gm_base, uint32_t valid = 0) { + std::conditional_t, UbTile> ub; + TASSIGN(ub, ub_offset); + if constexpr (Partial) ub.ColMaskInternal = (int)valid; + GlobalTensor, GmStride> gm( + (__gm__ T *)gm_base + (uint64_t)tile_index * Elems, GmShape()); + if constexpr (ToUb) { + TLOAD(ub, gm); + } else { + TSTORE(gm, ub); + } +} + +// Start the async load of this core's nth tile, if it has one. A function, +// not a lambda: set_flag/wait_flag do not resolve inside a lambda. +template +inline AICORE void issue_tile_load(uint32_t nth_tile, uint32_t core_id, + uint32_t core_count, uint32_t tiles, + uint32_t full_tiles, uint32_t tail_elems, + const event_t *buffer_free, + __gm__ void *input_gm) { + const uint32_t tile_index = core_id + nth_tile * core_count; + if (tile_index >= tiles) return; + const uint32_t buffer = nth_tile % Buffers; + const uint32_t off = buffer * Shape::slot_stride + SlotOffset::input; + wait_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[buffer]); + // at most one tile is partial, and only when batch does not fill it + if (tile_index == full_tiles) { + move_tile(tile_index, off, + input_gm, tail_elems); + } else { + move_tile(tile_index, off, input_gm); + } + set_flag(PIPE_MTE2, PIPE_V, buffer_free[buffer]); +} +#endif // __DAV_VEC__ +#endif // __CCE_AICORE__ + +// The pipeline: each core walks a strided subset of the tiles, keeping Prefetch +// loads in flight so DMA and the vector pipe overlap. +#if defined(__CCE_AICORE__) && defined(__DAV_VEC__) +// A device function rather than the kernel body, so a caller that wants the +// pipeline over a sub-range can reach it directly. mxfp4_quant below is the +// entry point and the only caller here. +template +inline AICORE void quant_tiles(__gm__ void *input_gm, __gm__ void *nibble_gm, + __gm__ void *scale_gm, uint32_t batch) { + using Shape = QuantShape; + using Offsets = SlotOffset; + set_mask_norm(); + set_vector_mask(-1, -1); + const event_t buffer_free[EVENT_SLOTS] = {EVENT_ID0, EVENT_ID1, EVENT_ID2, + EVENT_ID3, EVENT_ID4, EVENT_ID5, + EVENT_ID6, EVENT_ID7}; + const uint32_t core_id = get_block_idx(), core_count = get_block_num(); + // the remainder, if any, rides along as one extra partial tile + const uint32_t full_tiles = batch / Rows; + const uint32_t tail_elems = (batch % Rows) * K; + const uint32_t tiles = full_tiles + (tail_elems ? 1u : 0u); + + for (unsigned i = 0; i < NBuffers; ++i) // every buffer starts free + set_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[i]); + for (unsigned i = 0; i < NPrefetch; ++i) + issue_tile_load(i, core_id, core_count, tiles, full_tiles, + tail_elems, buffer_free, input_gm); + + uint32_t issued = 0; + for (uint32_t tile_index = core_id; tile_index < tiles; + tile_index += core_count, ++issued) { + const uint32_t buffer = issued % NBuffers; + // issued ahead of the wait below, so this load overlaps this tile's compute + issue_tile_load(issued + NPrefetch, core_id, core_count, + tiles, full_tiles, tail_elems, buffer_free, + input_gm); + wait_flag(PIPE_MTE2, PIPE_V, buffer_free[buffer]); + const uint32_t slot_base = buffer * Shape::slot_stride; +#ifdef MXFP4_TQUANT + tquant_passes(slot_base + Offsets::input, + slot_base + Offsets::nibbles, + slot_base + Offsets::scales); +#else + // name the UB regions once; inline casts are noise at every call site + using B16 = __ubuf__ uint16_t *; + B16 input_ub = (B16)(uintptr_t)(slot_base + Offsets::input); + B16 scale_ub = (B16)(uintptr_t)(slot_base + Offsets::scales); + B16 maxima_ub = (B16)(uintptr_t)Offsets::maxima; + B16 packed_ub = (B16)(uintptr_t)Offsets::packed; + B16 recips_ub = (B16)(uintptr_t)Offsets::reciprocal; + __ubuf__ uint8_t *nibble_ub = + (__ubuf__ uint8_t *)(uintptr_t)(slot_base + Offsets::nibbles); + // rotate in place, then quantize the rotated tile without it ever leaving + // UB +#ifndef FUSED_NO_ROTATE + rotate(input_ub); +#else + // Diagnostic build: same kernel, same tiling, same UB layout and buffer + // count -- only the butterfly removed. Comparing this against the quantizer + // alone separates the butterfly's vector cost from the cost of fusing at + // all (extra UB regions, so fewer buffers, so less overlap). + (void)0; +#endif +#ifndef FUSED_ROTATE_ONLY + block_abs_max(input_ub, maxima_ub); + compact_maxima(maxima_ub, packed_ub); + derive_scales(packed_ub, recips_ub, scale_ub); + pack_nibbles(input_ub, recips_ub, nibble_ub); +#else + // The other half of the fusion question. FUSED_NO_ROTATE keeps the + // quantizer and drops the butterfly; this keeps the butterfly and drops the + // quantizer, storing the rotated bf16 tile instead. Chained with the + // standalone quantizer it is the UNFUSED reference: two launches, two + // passes over HBM, 4 + 2.53 B/elem against the fused kernel's 2.53. + // + // Same tiling, UB layout and buffer count as the fused build, so the only + // differences against it are the arithmetic skipped and the bytes stored. + (void)scale_ub; + (void)maxima_ub; + (void)packed_ub; + (void)recips_ub; + (void)nibble_ub; +#endif +#endif + set_flag(PIPE_V, PIPE_MTE3, buffer_free[buffer]); + wait_flag(PIPE_V, PIPE_MTE3, buffer_free[buffer]); +#ifdef FUSED_ROTATE_ONLY + // `nibble_gm` carries the rotated bf16 tile here and `scale_gm` is + // untouched, so the launcher signature does not change. The harness + // allocates 2K bytes per row for it, not K/2. + if (tile_index == full_tiles) { + move_tile( + tile_index, slot_base + Offsets::input, nibble_gm, tail_elems); + } else { + move_tile( + tile_index, slot_base + Offsets::input, nibble_gm); + } + (void)scale_gm; +#else + if (tile_index == full_tiles) { + move_tile( + tile_index, slot_base + Offsets::nibbles, nibble_gm, tail_elems / 2u); + move_tile( + tile_index, slot_base + Offsets::scales, scale_gm, + tail_elems / MX_BLOCK); + } else { + move_tile( + tile_index, slot_base + Offsets::nibbles, nibble_gm); + move_tile( + tile_index, slot_base + Offsets::scales, scale_gm); + } +#endif + set_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[buffer]); + } + for (unsigned i = 0; i < NBuffers; ++i) // drain + wait_flag(PIPE_MTE3, PIPE_MTE2, buffer_free[i]); +} +#endif // __CCE_AICORE__ && __DAV_VEC__ + +template +__global__ AICORE void mxfp4_quant(__gm__ void *input_gm, + __gm__ void *nibble_gm, + __gm__ void *scale_gm, uint32_t batch) { +#ifdef __DAV_VEC__ + quant_tiles(input_gm, nibble_gm, scale_gm, + batch); +#else + (void)input_gm; + (void)nibble_gm; + (void)scale_gm; + (void)batch; +#endif +} + +#ifndef FUSED_INCLUDE_ONLY // define to take the device code without hosts +// ---------------------------------------------------------------- entry points +// One .so serves every K: fold over SUPPORTED_K for the instantiation. +template +inline void launch_for_k(uint32_t block_dim, void *stream, uint8_t *input, + uint8_t *nibbles, uint8_t *scales, uint32_t batch, + uint32_t k, std::index_sequence) { + ((k == SUPPORTED_K[Idx] + ? (void)(mxfp4_quant::value, + DEF_BUFFERS, DEF_PREFETCH> + <<>>(input, nibbles, scales, + batch)) + : (void)0), + ...); +} + +// An unsupported k is a silent no-op; the host validates +// (check_row_width). +extern "C" void call_hadamard_mxfp4_b32(uint32_t block_dim, void *stream, + uint8_t *input, uint8_t *nibbles, + uint8_t *scales, uint32_t batch, + uint32_t k) { + launch_for_k(block_dim, stream, input, nibbles, scales, batch, k, + std::make_index_sequence{}); +} + +template +inline uint32_t rows_for_k(uint32_t k, std::index_sequence) { + uint32_t rows = 0; + ((k == SUPPORTED_K[Idx] ? (void)(rows = RowsFor::value) + : (void)0), + ...); + return rows; // 0 for an unsupported k +} + +extern "C" uint32_t hadamard_mxfp4_b32_rows_for(uint32_t k) { + return rows_for_k(k, std::make_index_sequence{}); +} +#endif // FUSED_INCLUDE_ONLY diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/jit_util_fused_b32_a5.py b/examples/jit_cpp/fused_hadamard_quant_b32_a5/jit_util_fused_b32_a5.py new file mode 100644 index 00000000..9c21e31d --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/jit_util_fused_b32_a5.py @@ -0,0 +1,189 @@ +"""Build and load the block-32 fused Hadamard + MXFP4 quantize kernel. + +Deliberately thin: the kernel's own launcher dispatches on K, so this only has to +compile one .so and hand back a callable. Modelled on mxfp4_quant's jit helper, +with the entry points renamed and the width list narrowed to those the rotation +supports. +""" + +import ctypes +import os +import subprocess +from pathlib import Path + +import torch +import torch_npu # noqa + +HERE = Path(__file__).resolve().parent +BUILDDIR = HERE / "build" +_LIB_NAME = "fused_b32.so" +SOURCE = HERE / "fused_hadamard_quant_b32_a5.cpp" + +MX_BLOCK = 32 +VECTOR_CORES = 64 # vector cores on an A5 +# The rotation is always 32 wide, so it puts no power-of-two constraint on the +# row: any width the quantizer supports works, 4096 included. Must match +# SUPPORTED_K in the kernel. +SUPPORTED_K = ( + 32, + 64, + 96, + 128, + 192, + 256, + 512, + 768, + 896, + 1024, + 1152, + 1280, + 1408, + 1536, + 1664, + 1792, + 2048, + 2560, + 2816, + 3072, + 3584, + 4096, + 5120, + 6144, + 7168, + 8192, + 14336, + 16384, +) + + +def _flags(home): + return ( + f"-xcce --cce-aicore-arch=dav-c310-vec -DREGISTER_BASE " + f"-std=c++17 -O2 -fPIC -Wno-ignored-attributes -Wno-macro-redefined " + f"-mllvm -cce-aicore-stack-size=0x8000 " + f"-mllvm -cce-aicore-function-stack-size=0x8000 " + f"-mllvm -cce-aicore-addr-transform " + f"-mllvm -cce-aicore-dcci-insert-for-scalar=false -Xhost-start -Xhost-end " + f"-I{home}/aarch64-linux/include -I{home}/include" + ).split() + + +def compile_kernel(verbose=True, extra_defs=()): + """Compile the fused kernel to a .so. One .so serves every supported K. + + extra_defs are extra -D tokens for a tuning or A/B variant. They go into the + .so NAME as well as the command line, so a variant can never be served from + the default build's cache -- silently timing the wrong binary is the failure + this guards. + """ + home = os.environ.get("ASCEND_HOME_PATH") or os.environ.get("ASCEND_TOOLKIT_HOME") + if not home: + raise RuntimeError("source a CANN set_env.sh first: ASCEND_HOME_PATH is unset") + BUILDDIR.mkdir(parents=True, exist_ok=True) + tag = "".join("_" + d.lstrip("-D").replace("=", "") for d in sorted(extra_defs)) + # Reuse an .so newer than its source. These kernels unroll to hundreds of + # tile instructions and a rebuild can outlast the task queue's 600 s cap, so + # recompiling per call is not merely wasteful. + cached = BUILDDIR / _LIB_NAME.replace(".so", f"{tag}.so") + if cached.exists() and cached.stat().st_mtime > SOURCE.stat().st_mtime: + if verbose: + print("reusing", cached) + return cached + obj = BUILDDIR / f"fused_b32{tag}.o" + lib = cached + for step in ( + [ + f"{home}/bin/bisheng", + *_flags(home), + *extra_defs, + "-c", + str(SOURCE), + "-o", + str(obj), + ], + [ + f"{home}/bin/bisheng", + "-fPIC", + "-shared", + "--cce-fatobj-link", + f"-Wl,-soname,{lib.name}", + str(obj), + "-o", + str(lib), + ], + ): + if verbose: + print("compile:", " ".join(step[:3]), "...") + subprocess.run(step, check=True) + return lib + + +def current_stream_ptr(): + return ctypes.c_void_p(torch.npu.current_stream().npu_stream) + + +# The butterfly is the UNNORMALISED Sylvester matrix, so its output is sqrt(32) +# larger than an orthogonal block Hadamard's. That factor is deliberate and left +# to the caller: MXFP4's E8M0 scale is a power of two and sqrt(32) is not, so the +# scale cannot absorb it and the nibbles genuinely differ. Scale x by +# 1/sqrt(32) on the way in if orthogonal semantics are wanted. + + +def build_and_load(k=256, verbose=True, extra_defs=()): + """Return `fused(x) -> (nibbles, scales)` for row width `k`. + + Allocates its outputs, mirroring `torch_npu.npu_dynamic_mx_quant`, so the two + are comparable on the same call path. + + extra_defs reaches the compiler, so the reduced builds the benchmark's ladder + needs come from this one source: FUSED_ROTATE_ONLY leaves the butterfly + alone, FUSED_NO_ROTATE leaves the quantizer alone. + """ + if k not in SUPPORTED_K: + raise ValueError( + f"K={k} has no instantiation; supported: {sorted(SUPPORTED_K)}. " + "Widths must be a multiple of 32 with an instantiation." + ) + lib = ctypes.CDLL(str(compile_kernel(verbose=verbose, extra_defs=extra_defs))) + launch = lib.call_hadamard_mxfp4_b32 + launch.argtypes = [ + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_uint32, + ] + launch.restype = None + rows_for = lib.hadamard_mxfp4_b32_rows_for + rows_for.argtypes = [ctypes.c_uint32] + rows_for.restype = ctypes.c_uint32 + + def fused(x, out=None): + if x.dtype != torch.bfloat16: + raise TypeError(f"expected bfloat16, got {x.dtype}") + if x.shape[-1] != k: + raise ValueError(f"expected last dim {k}, got {tuple(x.shape)}") + if not x.is_contiguous(): + raise ValueError("expected a contiguous tensor; call .contiguous()") + batch = x.numel() // k + if out is None: + q = torch.empty((batch, k // 2), dtype=torch.uint8, device=x.device) + s = torch.empty((batch, k // MX_BLOCK), dtype=torch.uint8, device=x.device) + else: + q, s = out + launch( + VECTOR_CORES, + current_stream_ptr(), + ctypes.c_void_p(x.data_ptr()), + ctypes.c_void_p(q.data_ptr()), + ctypes.c_void_p(s.data_ptr()), + batch, + k, + ) + return q, s + + fused.rows_for = lambda: rows_for(k) + fused.k = k + return fused diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/run_benchmark.sh b/examples/jit_cpp/fused_hadamard_quant_b32_a5/run_benchmark.sh new file mode 100755 index 00000000..7f9a50c0 --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/run_benchmark.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# One-command on-device benchmark for fused_hadamard_quant_b32_a5 on an Ascend 950 (A5). +# Requires a real A5 device, torch + torch_npu, and bisheng (CANN toolkit). +# +# Needs a CANN whose PTO carries MXFP4 (Exp2DStrided in pto/npu/a5/TQuant.hpp): +# 9.1.0 and 9.2.0 both do, 9.0.0 does not. +if [[ -z "${ASCEND_TOOLKIT_HOME:-}" && -z "${ASCEND_HOME_PATH:-}" ]]; then + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +: "${ASCEND_HOME_PATH:=${ASCEND_TOOLKIT_HOME:-/usr/local/Ascend/ascend-toolkit/latest}}" +export ASCEND_HOME_PATH +cd "${SCRIPT_DIR}" +exec python3 benchmark.py "$@" diff --git a/examples/jit_cpp/fused_hadamard_quant_b32_a5/test_fused_hadamard_quant_b32_a5.py b/examples/jit_cpp/fused_hadamard_quant_b32_a5/test_fused_hadamard_quant_b32_a5.py new file mode 100644 index 00000000..45bc2f59 --- /dev/null +++ b/examples/jit_cpp/fused_hadamard_quant_b32_a5/test_fused_hadamard_quant_b32_a5.py @@ -0,0 +1,301 @@ +"""Does the fused kernel rotate and quantize correctly? + +The fused kernel cannot be bit-exact against a torch reference: it rotates in +bf16 with a specific operand order, and no torch expression reproduces that +tree. So correctness is established in three ways instead, from strongest to +weakest: + +1. **The scale bytes** must match a reference that rotates in fp32 and quantizes + with `torch_npu`. A scale is a power of two derived from a block maximum, so + bf16 rounding inside the butterfly almost never moves it -- if scales + disagree, the rotation is wrong, not merely rounded differently. +2. **The dequantized values** must track the fp32-rotated reference to within + MXFP4's own resolution. This catches a correct-looking permutation, which a + relative-error check on the packed bytes would not. +3. **The output must be non-trivial.** A kernel that writes nothing, or writes + the input back, is the characteristic silent failure on this hardware, and it + would otherwise pass a loose tolerance. +""" + +import numpy as np +import pytest +import torch +import torch_npu # noqa + +from jit_util_fused_b32_a5 import ( + MX_BLOCK, + SUPPORTED_K, + VECTOR_CORES, + build_and_load, +) + +VENDOR_DST_TYPE = 296 # E2M1, matching mxfp4_quant_a5's tests + + +def hadamard_matrix(n): + """Natural-order Sylvester +/-1 matrix, unnormalised -- the convention + fast_hadamard_a5 and its tests use.""" + m = np.array([[1.0]], dtype=np.float64) + while m.shape[0] < n: + m = np.block([[m, m], [m, -m]]) + return m + + +def reference(x, k): + """Rotate each 32-block in fp32 on the host, then quantize with the vendor op. + + Block-diagonal, matching the kernel: a row of k is k/32 independent + rotations. fp32 deliberately, so the reference does not depend on the + kernel's bf16 arithmetic. + """ + h = torch.from_numpy(hadamard_matrix(MX_BLOCK)).to(torch.float32) + flat = x.float().cpu().reshape(-1, MX_BLOCK) @ h + rot = flat.reshape(x.shape[0], k).to(torch.bfloat16).npu() + q, s = torch_npu.npu_dynamic_mx_quant(rot, dst_type=VENDOR_DST_TYPE) + return rot, q, s.reshape(s.shape[0], -1)[:, : k // MX_BLOCK] + + +E2M1_LEVELS = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 +) + + +def dequant(q, s, k): + """Unpack E2M1 nibbles and apply the E8M0 scale, on the host in fp32.""" + q = q.cpu() + lo, hi = q & 0x0F, (q >> 4) & 0x0F + codes = torch.stack([lo, hi], dim=-1).reshape(q.shape[0], -1) + mag = E2M1_LEVELS[(codes & 0x07).long()] + vals = torch.where(codes >= 8, -mag, mag) + exp = s.cpu().to(torch.int32) - 127 + scale = torch.ldexp(torch.ones_like(exp, dtype=torch.float32), exp) + return vals.reshape(-1, k // MX_BLOCK, MX_BLOCK) * scale.unsqueeze(-1) + + +@pytest.fixture(scope="module", autouse=True) +def seeded(): + torch.manual_seed(20260818) + torch.npu.set_device(0) + + +# A spread rather than all 26 -- each width is a separate .so compile, so the +# full set costs minutes. But the spread has to cover both UNROLL CLASSES, which +# the original five did not. +# +# The butterfly's unroll width is UnrollFor, so a width whose +# window count is not a multiple of 8 unrolls by 4 instead. Those two paths are +# different code. Every one of the original five (64, 256, 1024, 4096, 14336) +# unrolls by 8, so the unroll-by-4 path was never executed -- and it was broken: +# the sweep was instantiated with 8 slots regardless, so iterations overlapped by +# four windows and the last ran 1024 elements past the tile. K=96 failed this +# file's own thresholds and nothing noticed. +# +# Which class a width lands in is NOT a property of the width: it is +# rows_for(k) * k / 256, and rows_for depends on TILE_ELEMS. Raising the tile from +# 16384 to 24576 moved 96, 192, 768 and 2816 from unroll-4 to unroll-8 and left +# the matrix single-class -- caught by the test below, which is why the class +# membership is now DERIVED from the .so instead of being listed here. +# +# What the list has to guarantee is coverage: at the shipped tile some width here +# must land in each class. 896 and 3584 are the unroll-4 members at 24576 (both +# 84 windows), chosen at opposite ends of the range; the rest are unroll-8. +WIDTHS = (32, 64, 96, 192, 256, 768, 896, 1024, 2816, 3584, 4096, 14336, 16384) + + +@pytest.mark.parametrize("k", WIDTHS) +def test_matches_reference(k): + """Scales exact, dequantized values within MXFP4 resolution.""" + batch = 64 + fused = build_and_load(k=k, verbose=False) + x = torch.randn(batch, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + _, q_ref, s_ref = reference(x, k) + + scale_match = (s.cpu() == s_ref.cpu()).float().mean().item() + assert scale_match > 0.98, ( + f"K={k}: only {scale_match:.3f} of scale bytes match the fp32-rotated " + "reference -- that is a wrong rotation, not bf16 rounding" + ) + + got = dequant(q, s, k).reshape(batch, k) + want = dequant(q_ref, s_ref, k).reshape(batch, k) + denom = want.abs().mean().clamp_min(1e-6) + rel = (got - want).abs().mean() / denom + assert rel < 0.05, f"K={k}: dequantized mean rel error {rel:.4f} vs reference" + + +# Deep enough that the pipeline machinery runs, at both unroll classes and at +# both ends of the rows range: 896 is unroll-4 with 24 rows per tile, 14336 is +# unroll-8 with a single row. +DEEP_WIDTHS = (896, 4096, 14336) +TILES_PER_CORE = 4 + + +@pytest.mark.parametrize("k", DEEP_WIDTHS) +def test_matches_reference_many_tiles_per_core(k): + """The same check as test_matches_reference, but with a full pipeline. + + test_matches_reference uses batch=64, which at K=4096 is 11 tiles spread over + 64 cores: one tile for eleven cores and none for the rest. So the buffer + rotation (issued % NBuffers), the prefetch and the drain never run there, and + a fault in any of them cannot fail that test. This sizes the batch so every + core walks several tiles, and leaves a remainder so the partial tail tile is + taken too -- except at K=14336, where a tile is one row and no batch can + leave a remainder. + """ + fused = build_and_load(k=k, verbose=False) + rows = fused.rows_for() + batch = rows * VECTOR_CORES * TILES_PER_CORE + rows // 2 + 1 + tiles = -(-batch // rows) + + # The point of the test is the depth, so assert it rather than trusting that + # a TILE_ELEMS change left it intact. + assert tiles / VECTOR_CORES >= 3, ( + f"K={k}: {tiles} tiles over {VECTOR_CORES} cores is too shallow to " + "exercise the buffer rotation" + ) + # A tile is `rows` rows, so where rows == 1 every batch is a whole number of + # tiles and the kernel's partial branch is unreachable by construction. Only + # claim the tail where one can exist. + assert batch % rows or rows == 1, f"K={k}: batch {batch} is whole tiles" + + x = torch.randn(batch, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + _, q_ref, s_ref = reference(x, k) + + scale_match = (s.cpu() == s_ref.cpu()).float().mean().item() + assert scale_match > 0.98, ( + f"K={k}, {tiles} tiles: only {scale_match:.3f} of scale bytes match the " + "fp32-rotated reference" + ) + + got = dequant(q, s, k).reshape(batch, k) + want = dequant(q_ref, s_ref, k).reshape(batch, k) + denom = want.abs().mean().clamp_min(1e-6) + rel = (got - want).abs().mean() / denom + assert rel < 0.05, f"K={k}, {tiles} tiles: mean rel error {rel:.4f} vs reference" + + +@pytest.mark.parametrize("k", WIDTHS) +def test_output_is_nontrivial(k): + """A kernel that writes nothing, or echoes its input, must fail here.""" + fused = build_and_load(k=k, verbose=False) + x = torch.randn(32, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + assert q.any().item(), f"K={k}: nibbles are all zero" + assert s.any().item(), f"K={k}: scale bytes are all zero" + assert len(torch.unique(q.cpu())) > 4, f"K={k}: nibbles are degenerate" + + +@pytest.mark.parametrize("k", WIDTHS) +def test_rotation_actually_happened(k): + """The rotation must change the answer. + + Quantizing x directly and quantizing (x @ H) should differ; if the fused + output matches the unrotated quantization, `rotate` is a no-op -- which is + precisely the failure a tolerance-based check would wave through. + """ + fused = build_and_load(k=k, verbose=False) + x = torch.randn(64, k, dtype=torch.bfloat16, device="npu") + q_fused, _ = fused(x) + q_plain, _ = torch_npu.npu_dynamic_mx_quant(x, dst_type=VENDOR_DST_TYPE) + torch.npu.synchronize() + same = (q_fused.cpu() == q_plain.cpu()).float().mean().item() + assert same < 0.6, ( + f"K={k}: fused output matches the UNROTATED quantization at {same:.3f} " + "-- the rotation is not happening" + ) + + +def unroll_width(k, rows): + """UnrollFor, recomputed on the host. + + windows_per_tile is tile_elems/had_group with had_group = 256, i.e. + rows*k/256 -- NOT rows*k/32, which is the block count. Getting that wrong is + what made a first pass at this conclude the unroll was 8 everywhere. + """ + windows = rows * k // 256 + limit = 8 + while limit > 1 and windows % limit: + limit //= 2 + return limit + + +def test_width_matrix_covers_both_unroll_classes(): + """The matrix must exercise unroll-by-8 AND unroll-by-4. + + Guards the gap itself rather than one instance of it: the original five widths + were all unroll-by-8, so the other path was dead code in CI while being live + in production. A width added later -- or a change to TILE_ELEMS, which is what + actually happened -- must not quietly return the matrix to one class. + + The classes are derived here rather than asserted against a stored list. A + stored list is a snapshot of TILE_ELEMS, so it goes stale on a tuning change + and then reports a tile change as a width bug. When this fails it names the + supported widths that would restore coverage, since that is the fix. + """ + seen = {} + for k in WIDTHS: + rows = build_and_load(k=k, verbose=False).rows_for() + assert rows > 0, f"K={k}: rows_for returned 0" + seen.setdefault(unroll_width(k, rows), []).append(k) + if len(seen) < 2: + missing = {} + for k in sorted(SUPPORTED_K): + rows = build_and_load(k=k, verbose=False).rows_for() + cls = unroll_width(k, rows) + if cls not in seen: + missing.setdefault(cls, []).append(k) + pytest.fail( + f"the width matrix only exercises unroll {sorted(seen)}: {seen}. " + f"Both paths are live in production. Add one of: " + f"{ {c: v[:6] for c, v in missing.items()} }" + ) + + +@pytest.mark.parametrize("k", (256, 96, 768)) +def test_constant_row_is_a_delta(k): + """A constant row becomes one delta per 32-block: a sharp structural check. + + H's first column is all ones, so each block sums into its own element 0 and + cancels across the other 31. Smearing means the butterfly is pairing wrongly; + a single delta per *row* would mean it rotated the whole row instead of each + block, which is the specific bug this variant exists to avoid. + + 96 and 768 are unroll-by-4 widths. This check on one of those would have + caught the sweep overlap directly: an overlapped window gets rotated twice, + and a twice-rotated constant block is 32x the input in element 0 rather than + a clean delta. + """ + fused = build_and_load(k=k, verbose=False) + x = torch.ones(8, k, dtype=torch.bfloat16, device="npu") + q, s = fused(x) + torch.npu.synchronize() + vals = dequant(q, s, k).reshape(8, k // MX_BLOCK, MX_BLOCK) + lead, rest = vals[:, :, 0].abs(), vals[:, :, 1:].abs() + assert (lead > 0).all(), "each block's element 0 should carry its block sum" + assert rest.max() <= lead.min() * 0.05, ( + f"each 32-block should rotate to a delta; leaked {rest.max():.3f} " + f"against a lead of {lead.min():.3f}" + ) + + +def test_unsupported_k_is_rejected(): + """Widths without an instantiation must raise on the host. + + The dispatch would otherwise fall through silently and hand back the caller's + buffers untouched. + """ + for bad in (31, 33, 100, 0, 4095): + with pytest.raises((ValueError, TypeError)): + build_and_load(k=bad, verbose=False) + + +def test_wrong_dtype_is_rejected(): + fused = build_and_load(k=256, verbose=False) + for dtype in (torch.float16, torch.float32): + with pytest.raises(TypeError): + fused(torch.randn(16, 256, dtype=dtype, device="npu"))