mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
8f34abf498
* Add SM120 NVFP4 blockscaled GEMM support
* Address SM120 NVFP4 review cleanup
* Remove SM120 private C-fragment store helpers
* Unify SM120 blockscaled MMA TIR helper
* Clean SM120 NVFP4 blockscaled fast path
Remove the public micro_pipeline strategy knob from T.mma_gemm_blockscaled and route the packed-scale SM120 path internally through the package-pingpong lowering. Hide SM120 strategy helpers from the public TIR API and remove stale emitter experiments that were only used by the strategy sweep.
Restore the optimized persistent SM120 NVFP4 example as the benchmark path, but keep the language surface clean by using sf_layout to select the packed scale contract and a regular T.copy epilogue instead of private C-fragment store APIs.
Validation: pre-commit run --files on changed files; pytest testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py -q => 89 passed; 512^3 correctness passed; 8192^3 128x128x256 stages=2 timing => 0.7198 ms, 1527.57 TFLOPS.
* Prune SM120 NVFP4 debug lowering paths
* Simplify SM120 NVFP4 blockscaled example
Move the persistent CUTLASS comparison harness to maint/gemm and keep the public SM120 NVFP4 example focused on a non-persistent TileLang kernel using T.mma_gemm_blockscaled and normal T.copy output.
Add a source-level guard so the public example does not grow CUTLASS build or subprocess benchmark plumbing again.
* Remove SM120 fulltile debug macros and compile flags
The SM120 fulltile package path now has a single production behavior:
row-major shared A/B views and the packed-U4 ldmatrix.b16 operand fetch.
Drop the TL_SM120_FULLTILE_* configuration macros (ROWMAJOR_VIEW,
FP4_LDSM, FP4_REG_SHIFT, CUTE_ACCUM_*) together with the dead compact
view helpers, and stop injecting -D flags from the example and the
maint benchmark.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Fix TensorCoreIntrinEmitter.mma base signature regression
Merging the SM120 block-scale mode into TensorCoreIntrinEmitter changed
mma()'s 4th/5th positional parameters to SFA_buf/SFB_buf, so the plain
gemm lowering's positional mma(A_local, B_local, C_buf, ki) call bound
ki to SFA_buf and raised 'Scale buffers require TensorCoreIntrinEmitter
block-scale mode' for every non-blockscaled MMA GEMM (CI examples:
attention_sink, deepseek_nsa/mhc/v32).
Keep the base (A, B, C, k_inner) positional signature, make all
scale-related parameters keyword-only, forward k_inner to the base
implementation, and drop the unused blockscaled mma_atom override that
shadowed the base method the same way. Add a signature-contract
regression test.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Pin scale layout byte-compat with CuTeDSL blocked SF layout
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>
* Remove unreachable SM120 blockscale exploration code
Back-traced every symbol the example and the maint WS benchmark pull in,
then walked transitive reachability from the production lowering
(gemm_mma.py) and tests:
- TensorCoreIntrinEmitter: drop 16 unreachable micro-pipeline strategy
variants left over from bring-up (prefetched-scale orderings, owner /
owner-wide / kpack scale copies, serpentine issue orders, selector
probe) — the production path uses ldmatrix_a/b, mma, ldscale(_fragment),
mma_full_b_atom_with_scale_fragments and the fulltile package backend.
- gemm_sm120.h: drop two device helpers nothing emits
(dual-issue same-B mma wrapper and the accum-inplace variant).
- quantize/nvfp4.py: drop the one quantize kernel variant with no
dispatch site (32x256 tma_io).
No public surface change; language API stays T.mma_gemm_blockscaled +
ptx_mma_block_scale. Validated with a fresh JIT cache: nvf4 language +
example CLI + quantize layout + atom mma tests (122 passed), example and
WS benchmark --verify both pass.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Scope shared-memory bit-exact sizing to packed scalar NVFP4
The exact-bits shared storage sizing previously applied to every dtype
with element_bits % 8 != 0 (int4/uint4/uint2/...), halving e.g. int4
shared allocations while SharedByteOffsetToLogicalIndexOffset kept the
legacy bytes-per-element conversion for those dtypes — an inconsistent
size/index pair for any non-FP4 sub-byte buffer in a merged arena.
Restrict the new sizing to float4_e2m1fn scalar so it mirrors the
byte-offset -> logical-index special case exactly: NVFP4 buffers get
packed two-per-byte sizing and indexing, every other dtype keeps
bit-exact upstream behavior. This pass now changes nothing outside the
NVFP4 path.
Validated after rebuilding libtilelang.so: nvf4 language + example CLI +
quantize layout + access_ptr codegen + atom mma tests (129 passed),
example and WS benchmark --verify pass with a fresh JIT cache.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Add T.copy_ue4m3_scale_tile scale staging helper
The example and the maint warp-specialized benchmark each carried a
private macro with the blockscaled_chunk_kmajor tile addressing spelled
out by hand. Move that knowledge into the library next to
T.mma_gemm_blockscaled:
- T.ue4m3_scale_tile_source_coords: the single source of the packed
scale tile addressing (tiles are contiguous in this layout);
words_per_kblock stays a parameter so future contracts
(block_K != 256, MX scale packings) reuse it.
- T.copy_ue4m3_scale_tile: cooperative whole-block staging macro with
trace-time dtype/shape validation, used by the example. The scale
operand stays visible in user kernels; only the layout math is hidden.
The maint kernel keeps its producer-warp-group staging schedule and
reuses the coords helper. Validated: nvf4 language + example CLI tests
(91 passed incl. a new validation test), example and WS benchmark
--verify pass.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Support M or N tail tiles in the SM120 NVFP4 example
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>
* Slim tilelang/quantize/nvfp4.py
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>
* Converge quantizer kernels and drop example debug input modes
nvfp4.py: remove the two remaining bring-up era quantize kernels (the
legacy no-TMA-load 16x256 schedule and the k64_per_cta=1 generic path,
reachable only through tuning knobs nothing uses) and the use_tma_load /
k64_per_cta parameters. The dispatcher now maps straight onto the single
tiled kernel factory; outputs verified bitwise-identical to the previous
code for every remaining configuration. 978 -> 710 lines (1334 before
the slimming series).
Example: drop the ones / constant-scale debug input modes and their CLI
flags; the random-FP4 + random-binary-scale path is the only
verification mode and checks exact equality.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Keep tilelang.quantize to device helpers and format contract
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>
* Import layout oracle from the nvfp4 submodule in the language test
Follow-up to the package export slimming: blockscaled_chunk_kmajor_word_offset
is no longer re-exported at the tilelang.quantize package level.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Group SM120 NVFP4 maint files and use native target API
Address maintainer review follow-ups:
- Move the four loose SM120 NVFP4 files under maint/gemm/ into a dedicated
maint/gemm/gemm_sm120/ folder to distinguish them from other gemm types,
updating the repo-root path resolution, the benchmark usage string, the
cutlass_nvf4_ref.cu source path, and the quantizer loader in the scale-layout
test.
- Guard the SM120-hardware test points with
tilelang.testing.requires_cuda_compute_version_eq(12, 0): replace the manual
capability skip in the example CLI test and tighten the block-scale kernel
tests from _ge(12, 0) to _eq(12, 0).
- Reuse the native tilelang.cuda.target.target_is_sm120 helper in
_is_explicit_non_sm120_cuda instead of hand-parsing the arch string, keeping
the explicit-arch guard so targets without an arch attr never trip the C++
ICHECK on hosts that have no SM120 GPU.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Keep NVFP4 scale staging out of the tilelang.language surface
T.copy_ue4m3_scale_tile and T.ue4m3_scale_tile_source_coords were
NVFP4-specific helpers on the general T.* surface. Remove both and move the
abstraction down a level:
- The example now declares the packed scale source with its honest "tile rows"
shape ([n_mn_blocks * n_k_blocks * block_MN, block_K // 64] uint32): the
blockscaled_chunk_kmajor layout stores each (mn_block, k_block) tile
contiguously, so with this view every tile is a plain rectangular slice and
staging is a two-line T.Parallel loop with row-offset addressing only. The
host obtains the view zero-copy via
swizzle_blockscaled_chunk_kmajor_scale_words(...).reshape(-1, 4).
- tilelang.quantize.nvfp4 gains blockscaled_chunk_kmajor_tile_source_coords,
the hand-written offset formula for staging schedules that keep the 2-D
[MN, K/64] view (the maint benchmark's producer warp group uses it).
Compiler-internal staging was prototyped first (plain T.copy of the slice)
and rejected for now: the raw-pointer bulk-TMA path this lowers to on some
shapes never arrives on its mbarrier (deadlock, tracked upstream), and the
synchronous fallback stalls the producer warps (-8.5% at 2048^3). The example
comment documents this so the loop can become a plain T.copy once fixed.
Validation on RTX 5090: 123 tests pass; example outputs bitwise-identical to
the previous macro implementation on 5 shapes including M/N tail tiles
(257x384x512, 130x128x256, 128x136x256); example perf at parity
(100.1-100.4% over 2048^3/4096^3/8192^3/8192x8192x1536, interleaved A/B,
generated CUDA identical modulo a loop variable name); maint WS benchmark
kernel source byte-identical with --verify passing and perf unchanged
(641.9 / 1310.5 / 1149.5 TF at 2048^3 / 8192^3 / 8192x8192x1536).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Drop stale scale-load metadata from the SM120 WS benchmark
_SM120_SCALE_LOAD = "tma" and _SM120_AB_SHARED_STORAGE = "packed" had no
consumers (leftovers from removed strategy knobs), and the block_K guard
message claimed the scales are TMA-loaded when the producer warp group
actually stages them with a plain lane-strided copy. Delete the dead
constants and state the guard's real reason (the packed-scale MMA contract
consumes four K64 scale words per tile).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Strip trailing blank lines left in gemm_op.py
The removal of the NVFP4 staging helpers left two blank lines at EOF, which
end-of-file-fixer flags in Quick Lint / pre-commit.ci.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Trim the SF staging comment to the load-bearing caveat
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Consolidate NVFP4 tests into a single file
Review follow-up: fold the three NVFP4 test files into
test_tilelang_language_nvf4_mma_block_scale.py and drop the deep
contract-oracle suites.
- The example CLI file is gone; its tail-tile e2e run and the
unsupported-tail rejection test move here, the source-string checks are
dropped.
- The quantize scale-layout file is gone; the load-bearing packer contracts
move here (device word-offset fixed cases, packer vs independent byte
oracle, CuTeDSL canonical SF byte equivalence, packer/encoder known values
and rejections, torch-reference quantizer contracts, the GPU quantizer
cross-check, and swizzle tail padding).
- The OMMA selector / issue-schedule / package-lifecycle oracle suites are
removed from both files; the production contract stays enforced by the
lowering asserts and the e2e correctness runs.
3 files / 1836 lines -> 1 file / 923 lines; 39 tests pass locally on
RTX 5090 (full-suite run, fresh JIT cache).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Evaluate emitter dtype defaults lazily under the dialect facade
The language dialect split (#2734) makes tilelang.language import
tilelang.cuda.language, which transitively imports this module while
tilelang.language is still initializing. Default argument values and
dataclass field defaults run at import time, so spelling them T.float16 /
T.float32 now dereferences a partially initialized module and raises.
Use the plain dtype string literals, matching the convention the rest of
this file (and upstream's copy) already follows; T.* stays confined to
function bodies.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Restore the scale-layout contract tests; keep the example CLI file folded in
Revise the earlier consolidation (c34e7502) after reconsideration: the
quantize scale-layout suite comes back unchanged - it is the CPU-executable
contract between the Python packers and the CUDA-side addressing (and the
only part of the NVFP4 coverage that actually runs on CI hosts without an
SM120 GPU). The net test-file change against the pre-consolidation state is
now just: the example CLI file is gone, with its tail-tile e2e run and
unsupported-tail rejection test living in the language test file, and the
two brittle source-string checks dropped.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* refactpr example
* [SM120] Simplify NVFP4 benchmark
* [SM120] Simplify NVFP4 correctness comparison
* [SM120] Inline NVFP4 quantizer pass configs
* [SM120] Move block-scaled MMA helper to instruction headers
* [CUDA] Fix packed FP4 address codegen
* [SM120] Simplify NVFP4 lowering internals
* [SM120] Generalize NVFP4 block-scaled MMA lowering
* [SM120] Fix package macro source test
* [SM120] Remove redundant local FP4 access pointer test
* [SM120] Split block-scaled MMA emitter
* [SM120] Split block-scaled GEMM lowering
* [SM120] Isolate block-scaled GEMM plumbing
---------
Co-authored-by: qutao <qutao@qutaodeMacBook-Pro.local>
Co-authored-by: qutao <qutao@aishi.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Rachmanino <109033598+Rachmanino@users.noreply.github.com>