tilelang/quantize's identity is device-side conversion helpers plus
host-side data-prep utilities (quantization.py's _tir_* converters, the
lop3/mxfp intrin groups); a complete quantizer operator does not belong
in the library. Move the TileLang BF16->NVFP4 kernel factory and its
dispatcher to maint/gemm/tilelang_nvfp4_quantizer.py, next to the maint
benchmark that consumes it.
nvfp4.py now carries exactly the format contract: the _tl_* device
helpers (FP4/UE4M3 encode/decode, scale-source coordinates) any user
kernel can call, the blockscaled_chunk_kmajor packers/swizzles, and the
torch reference quantizer — 524 lines (1334 at series start). The
package re-export list shrinks to the four user-facing names; the alias
wrappers are gone.
Validated: quantize + example CLI suites (35 passed), maint benchmark
and example --verify pass end-to-end.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two structural reductions, no behavior change:
- Merge the three TMA-staged quantize kernel variants (16/32/64x256)
into one _tilelang_nvfp4_blockscaled_quantize_kernel_tiled factory:
their inner loops were line-identical, differing only in rows_per_cta,
stage count and the store path (TMA store vs direct global stores),
all of which are trace-time parameters now. Outputs verified
bitwise-identical to the previous per-variant kernels for every
dispatch configuration (quant bytes and scale source).
- Move the 14 SM120 OMMA scale-selector contract oracles into the
contract test file that is their only consumer; the quantize module
stays focused on quantization and layout packing.
1334 -> 978 lines. Quantize suite: 32 passed.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Production shapes rarely keep M (batch x seq) a multiple of 128. Lift
the divisibility requirements to the validated envelope:
- scale packers zero-pad rows to full 128-row tiles (the same ceil
convention CUTLASS/cuBLAS/CuTeDSL use for the blocked SF tensor), so
the packed source always has ceil(rows/128)*128 rows;
- the example accepts an arbitrary M or an N tail: TMA loads zero-fill
out-of-bounds rows and the C store is predicated. N keeps the
16-byte-aligned-rows rule (N % 8 == 0), the same contiguous-dim
alignment CuTeDSL validates;
- simultaneous M and N tails are rejected with a clear error: the
fragment-to-global copy lowering emits the column guard bound from
the wrong dimension in that case (upstream bug, tracked separately).
Tests: tail-shape runtime correctness (257x384, 130x128, 128x136 with
exact binary-scale references), rejection paths, and packer padding
round-trip.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
blockscaled_chunk_kmajor packed words are byte-identical to the
CUTLASS/CuTeDSL canonical blocked SF layout (tile_atom_to_shape_SF,
sf_vec_size=16, atom ((32,4),(16,4)):((16,4),(0,1)) tiled with order
(2,1,3)); the two views bridge zero-copy via torch view casts. Document
that on the packer and lock it with a cross-layout contract test.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>