* [CUDA] Pack logical TMEM buffers into shared tcgen05.alloc arenas
LowerSharedTmem gave every T.alloc_tmem buffer its own tcgen05.alloc,
rounded up to a power of two. Small buffers therefore each burned a
32-column allocation: a 384-column accumulator next to three 4-column
block-scale buffers needs 396 columns but asked the hardware for
512 + 32 + 32 + 32 = 608 of the 512 a CTA has, so that kernel could not
be expressed at all. Kernels worked around it by declaring one wide
buffer and slicing scale factors out of it by hand.
Plan the allocations instead. PlanTmemArenas packs logical buffers into
shared allocations, and each buffer's column offset inside the one it
shares is folded into every address it forms -- both the encoded
coordinate of a BufferLoad, which tcgen05.ld/st consume, and the
(base, offset) argument pairs of the tl.ptx_tcgen05_* intrinsics.
Buffers at least 32 columns wide keep the 32-column alignment
tcgen05.alloc would have given them alone; narrower ones use the
4-column granularity of tcgen05.cp.32x128b.warpx4, the same boundary
CUTLASS and DeepGEMM place these operands on.
A buffer joins an existing allocation only when that strictly reduces the
total column count, so packing can never grow a kernel's TMEM footprint
and a kernel whose buffers already fit lowers exactly as before. The
mxfp8 block-scaled GEMM drops from 320 to 288 columns; flash attention
and single-accumulator GEMMs are untouched.
Buffers with an explicit T.deallocate_tmem keep their own allocation,
since releasing one early would release whatever shares it. A base
address that reaches an operation this pass cannot offset is rejected
rather than silently addressing the start of the allocation. A block
that allocates more than 512 columns is now reported at compile time
instead of failing inside tcgen05.alloc on device.
* [Test] Cover TMEM arena planning
Pins the plan the packer produces, not just that it produces one:
- 396 columns across four buffers become one 512-column allocation at
offsets 0 / 384 / 388 / 392, matching DeepGEMM's hand-written layout
for this kernel column for column
- block-scale factors share an allocation while the accumulator keeps
its own, because packing all three would cost more than it saves
- equal-cost cases keep separate allocations, so a kernel that already
fits is untouched
- allocations are issued widest first, as PTX requires
- packing never allocates more than separate allocations would, over six
shape combinations
- an explicit T.deallocate_tmem keeps its own allocation
- a dynamic pipeline-stage coordinate still carries the column offset
- a block over the 512-column budget is reported
* [Test] Roundtrip two TMEM buffers that share one allocation
Covers the tcgen05.st/ld address path for a buffer packed at a nonzero
column offset, which no existing kernel exercises: the block-scale
buffers that share an allocation today are never read back into
registers.
The wide buffer's roundtrip is the assertion that matters. The narrow
one would still round-trip if its store and its load dropped the column
offset together -- but its store would land on the wide buffer's first
columns, so only the wide readback catches it. Verified by mutation:
dropping the offset in LowerSharedTmem fails this test.
* [Test] Guard the TMEM transform tests on Blackwell
Gate every TMEM test the same way, matching the decorators the tcgen05
language and kernel tests already carry.
These three lower IR for a hardcoded sm_100 target and need no device, so
this does trade running them on any machine for consistency with the rest
of the TMEM suite.
* 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>
* [BugFix] Skip descriptor TMA for device-bound copy bases
* [BugFix] Run TMA base annotation independently of WS
* [Refactor] Expose TMA base annotation as a standalone pass
* [Docs] Restore warp-specialization pipeline comment
* [BugFix] Honor nan_propagate in reduce max/min/absmax clear=False write-back
MakeUpdate (the clear=False write-back merge) used plain Max/Min for
max/min/absmax reducers, silently dropping NaN when nan_propagate=True.
Now it emits max_nan/min_nan on fp16/bf16 when nan_propagate is set,
mirroring the existing MakeReduce path.
GHI: #2697
* [BugFix] Honor nan_propagate in reduce max/min/absmax clear=False write-back
MakeUpdate (the clear=False write-back merge) used plain Max/Min for
max/min/absmax reducers, silently dropping NaN when nan_propagate=True.
Now it emits max_nan/min_nan on fp16/bf16 when nan_propagate is set,
mirroring the existing MakeReduce path.
GHI: #2697
* [BugFix] Cast to the destination dtype in the scalar T.copy path
The scalar fast path in T.copy lowers a single-element copy directly to a
BufferStore. Unlike the general region path, which casts the loaded value to
the destination dtype in src/op/copy.cc, it stored the source value unchanged.
A cross-dtype scalar copy therefore failed the strict BufferStore dtype check.
This affects the common gather pattern of copying an int64 index tensor
(torch's default for argsort and topk) into an int32 on-device scratch buffer.
Mirror the general path by casting the source value to the destination
BufferLoad dtype when they differ, preserving any index lanes.
Add regression coverage for int64-to-int32 narrowing and int16-to-int32
widening.
Fixes#2689
* [BugFix] Convert the test index dtype with as_torch()
The TileLang dtype subclasses str, so getattr(torch, index_dtype) works, but
as_torch() states the conversion directly instead of relying on that detail.
* Add language dialect overlays
* Refactor language dialects into static backend layers
* Make language facade a static re-export and extract sparse layout leaf
Replace the runtime-activated language facade with a static
`from tilelang.cuda.language import *` re-export: drop the .pyi stubs +
generator + py.typed, the globals()-based __all__ scraping, and the
_activate_cuda_facade() setattr. Backend dialects (cuda/metal/rocm) now
build on tilelang.language.common with per-backend tir overlays.
Import-time dtype defaults in the mma/wgmma/mfma macro generators are
pinned to the dtypes leaf so they no longer touch the half-initialized
facade during bootstrap. Metal generators import tilelang.metal.language
(not the cuda facade) so the eager builder's real getattr resolves the
simdgroup ops; a lazy __getattr__ breaks the metal.language <-> intrinsics
cycle.
Extract the 2:4 sparsity layout metadata into
tilelang/cuda/intrinsics/sparse_layout.py (a dtypes-only leaf) so mma_sp
imports it directly and tilelang.utils.sparse drops out of the bootstrap
closure, letting its @tilelang.jit compressor go back to a module-level
decorator.
Exports unchanged: common=558, cuda/facade=683.
* Remove unwired BackendModule scaffolding
The BackendModule manifest and per-backend backend.py registrations were
included by mistake: nothing imports the per-backend backend.py files, the
ensure_*_loaded lazy-loaders have no call sites, and the mechanism
duplicates the register_lazy_execution_backends wiring already in
tilelang/backend/__init__.py. Drop the dead code and its unit test.
* Import tvm.arith in metal macro generator
The simdgroup stride check uses arith.Analyzer but the module was missing
the tvm.arith import.
* Use string dtype defaults in CUDA macro generators
Default dtype arguments are evaluated at import time, when the language
facade is only partially initialized, so they cannot reference T.*. They
previously used the tilelang.language.dtypes leaf; since every consumer
treats these as strings (annotated str, wrapped via DataType before use),
plain string literals are equivalent and drop the _dtypes dependency.
* Drop MPSIntrinEmitter from metal language facade
The emitter is only consumed directly from tilelang.metal.intrinsics, never
via the language surface, and exposing it here required a lazy __getattr__ to
break the metal.language <-> metal.intrinsics import cycle. Removing it from
__all__ eliminates the cycle and the deferred import.
* Restore dtype forwarding for ptx_ldmatrix
T.ptx_ldmatrix previously accepted a dtype= keyword via _dtype_forward (from
the old language/ast/ir.py path). The CUDA dialect alias dropped the wrapper
while every neighboring ptx_* operator kept it, changing the public call
contract. Wrap it to match.
* Re-export tir Select/min/max wrappers and dedup metal __all__
Add Select/min/max to tilelang.language.tir.ir __all__ so the common
facade re-exports the documented wrappers instead of the raw TVM
symbols. Build metal __all__ via dict.fromkeys to match the sibling
dialects and satisfy PLE0604.
---------
Co-authored-by: SiriusNEO <chaofan@deepseek.com>
* [Transform] Never classify side-effecting binds as replayable
IsReplayableScalarBind (introduced in #2245) decides whether a scalar
bind may be dissolved and its value re-evaluated at every use site
(IfStmtBinding, default-on) or replayed across pipeline stages
(PipelinePlanning, InjectSoftwarePipeline). The predicate checked the
dtype and read/write disjointness, but never the value's effect kind,
so a bind like
pos = T.atomic_add(counter[0], 1, return_prev=True)
used at two sites (bound check + store index) was inlined into both use
sites, executing the atomic once per use site instead of exactly once
and silently corrupting compaction/histogram/top-k kernels since
v0.1.11. The atomic intrinsic is correctly registered kOpaque; its
pointer argument enters the expression through tvm_access_ptr, which
the access collector records as a read, so the write-disjointness check
passed and the bind was declared replayable.
Reject any bind whose value has SideEffect > kReadState, mirroring the
effect guard StmtSimplifier already applies to let inlining. Read-only
values remain replayable; read consistency is enforced by the existing
write-set disjointness check.
The shipped examples/deepseek_v32/topk_selector.py uses exactly this
pattern and is miscompiled on v0.1.11/v0.1.12; its test only printed
recall and could not fail, so it also gains a recall assertion.
Regression tests:
- transform: IfStmtBinding must keep a bind whose value is an atomic
RMW returning the previous value.
- language (CUDA): compaction kernel using atomic_add(return_prev=True)
must contain exactly one AtomicAddRet and produce a permutation with
counter == N.
* [Transform] Validate scalar bind replay safety
Track access-pointer write masks during bind and pipeline analysis, and reject unsafe precomputed replay annotations during pipeline injection.
Add regressions for atomic target reads and side-effecting replay masks.
* Remove recall calculation from top-k selector
Removed recall calculation and assertion for top-k selection.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
T.atomic_load and T.atomic_store previously accepted every order in the shared memory-order map, including operation-illegal combinations that trigger device-side assertions in libcu++.
Validate memory orders against operation-specific allow lists before lowering, report invalid values with a frontend ValueError, and update the documentation and regression coverage accordingly.
Fixes#2574
* [BugFix] Return previous value for atomic min/max/add (scalar + vectorized)
`atomic_min`/`atomic_max` (scalar) and `atomic_addx2`/`atomic_addx4`
(vectorized) typed their call with a non-void return under
`return_prev=True` but always selected the void `*_elem_op`, so codegen
emitted an empty binding (`prev = ;`) and nvcc failed.
Select the `*_ret_elem_op` variant under `return_prev`, as `atomic_add`
already does:
- min/max: the ret ops, codegen handlers, and device templates
(AtomicMaxRet/AtomicMinRet) already existed.
- addx2/addx4 (float32): register the ret ops (builtin.cc/.h) and add
codegen handlers emitting AtomicAddx2Ret/AtomicAddx4Ret (already in
the device headers) as expressions; the return is typed as the packed
2-/4-lane vector so the store matches.
- addx2 (fp16/bf16): the ret store LHS is `uint1`; codegen wraps the
half2/bf162 result with `tl::to_uint1`.
- addx4 (fp16/bf16): add AtomicAddx4Ret(half_t*)/(bfloat16_t*) device
overloads that run two AtomicAddx2Ret halves and return them packed as
uint2 (a fp16x4 is stored as uint2).
The 16-bit x4 path is per-pair atomic (two half2 atomics), not
whole-quad single-atomic -- hardware has no vector atomic wider than
`.f16x2` for 16-bit types. This matches the existing fp32-x4 scalar
fallback, and the docstring only promises "quad-width atomic addition".
Adds tests across scalar/x2/x4 and float32/fp16/bf16.
Fixes#2667
* minor fix for missing conditiions
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
`T.infinity` aborted with an internal `LOG(FATAL)` ("Cannot decide
infinity for type float8_e5m2") because `infinity_op` had no fp8 branch.
e5m2 is the only fp8 format with a representable inf, so add a branch
returning +inf, and emit it in codegen via its bit pattern
(`::bitcast(0x7c)`, mirroring the tfloat32 path) since inf/NaN cannot be
spelled as a float literal for the fp8 constructor. Other fp8 formats
(e4m3*, e3m4, e8m0, finite fnuz incl. e5m2fnuz) have no inf and keep the
fatal. Adds e5m2 to the infinity test.
* [CUDA] Add T.__fns intrinsic for find-nth-set bit
Expose CUDA __fns in TileLang for direct lookup of the k-th active lane in a 32-bit bitmask, complementing the existing T.__ffs helper.
* Add builtin intrinsics for shared ldst and atomics
Use PTX v4 atomics for 16-bit CUDA atomic_addx4
The initial fp16/bf16 atomic_addx4 fix avoided the generic float4 path by splitting each operation into two AtomicAddx2 calls. That fixed correctness, but it still used the simulated implementation on Hopper where PTX provides native vector atomics.
Add sm90+ half_t and bfloat16_t AtomicAddx4 overloads backed by the PTX vector forms atom.global.v4.f16.add.noftz and atom.global.v4.bf16.add.noftz, including the existing memory-order variants. Keep the x2 fallback for pre-sm90 targets because PTX vector f16/bf16 atomics require sm90 or newer.
Keep the CUDA language regression focused on fp16/bf16 atomic_addx4 codegen and offset coverage.
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
* Add split-k for example_warp_specialize_flashmla
* Rotate QK issue to end of main loop to allow more latency hiding with TMA issue
* During rescale, read max score from REG instead of SMEM
* Move K_pe_shared_1 load
* Use stmatrix for SP0 write to reduce SMEM bank conflict
* Support arbitrary SMEM layout in MMA lowering
* Remove runtime guard for increase_descriptor_offset
T.fill of an int8 buffer with a negative scalar (other than -1) silently
wrote only 1 of every 4 elements correctly and clobbered the other 3 to
-1. The vectorized fill packs four int8 lanes into one 32-bit store via
make_int(v,v,v,v); a negative signed char is integer-promoted to int and
sign-extended before the shifts/OR, so its high 1-bits flood the
neighbouring lanes (e.g. make_int(-7,-7,-7,-7) returned 0xFFFFFFF9 =
[-7,-1,-1,-1]).
Build the packed value from explicit unsigned bytes so the sign-extension
bits cannot leak across lanes. Add an int8 negative-fill regression test.
Fixes#2427
* Add stmatrix m16n8 support
* lint
* Refactor stmatrix code generation to handle optional shape argument
Updated the handling of the shape argument in the stmatrix function across CUDA, CuTeDSL, and HIP code generators. The shape is now checked as the last argument, and the payload start index is adjusted accordingly. This change improves consistency and clarity in the code generation process for tile language operations.
* Refactor stmatrix function naming for consistency
Updated the naming convention of stmatrix functions in CUDA and related files to improve clarity and consistency. The function names now follow a uniform pattern, placing the shape descriptor before the axis identifier. This change enhances readability and aligns with the recent updates in code generation logic.
* Enhance extern call statement printing in CUDA codegen
Updated the `VisitExpr_` method in `codegen_cuda.cc` to improve the clarity of extern call statements. The function now includes comments indicating that the transformation, number, and shape are encoded in the helper name, and only the shared memory pointer and value operands are printed. This change aims to enhance code readability and maintainability.
* Refactor extern call expression printing in CUDA codegen
Introduced a helper function to streamline the printing of extern call expressions in the `VisitExpr_` method of `codegen_cuda.cc`. This change enhances code readability by reducing redundancy and improving the clarity of the generated output for specific function calls, such as `tl::pack_b8x4`.
* Refactor shape argument handling in stmatrix code generation
Updated the handling of the shape argument in the `VisitExpr_` method for both CUDA and CuTeDSL code generators. The shape is now checked as the fourth argument instead of the last, improving clarity and consistency in extern call statement printing. This change enhances the robustness of the code generation process for tile language operations.
* Refactor shape argument handling in stmatrix code generation
Updated the handling of the shape argument in the `VisitExpr_` method for both CUDA and CuTeDSL code generators. The shape is now checked as the last argument and validated accordingly, improving clarity and consistency in the code generation process for tile language operations.
* Refactor tilelang lower function context management
Updated the context management in the `test_tilelang_copy_uses_stmatrix_m16n8_for_sm100_int8_shared_store` test by removing the unnecessary `PassContext` wrapper. This change simplifies the code and enhances clarity in the test setup for tile language operations.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Vectorized atomic_add into an fp16/bf16 destination with an fp32 source
bit-reinterpreted the source float2 as half2/__nv_bfloat162 instead of
numerically converting it, silently producing wrong results. This hit
both the auto-vectorized T.atomic_add path (AtomicAddx2 value form) and
the explicit T.atomic_addx2 API (pointer form).
Add float2 / float* converting overloads (round-to-nearest) to ToHalf2
and ToBfloat162 so the vectorized helper converts instead of
reinterpreting the source bytes. Add mixed-dtype regression tests for
both entry points.
Fixes#2382
morluto