A tile-region `T.atomic_add(..., memory_order=...)` lost its ordering
whenever the add auto-vectorized: the vectorizer rebuilt the wide op from
only `{dst, src}`, so `AtomicAddx2`/`AtomicAddx4` were emitted without the
order argument and fell back to their `relaxed` default. The identical call
honored the order when it lowered scalar-wise, so whether the requested
ordering survived depended only on whether the per-thread extent was
divisible by the atomic vector width.
Carry the trailing operands through the wide reconstruction. Codegen and
the device helpers already accept the order (`args.size() > 2` guard, and
`int memory_order` parameters that branch to ordered PTX), and
`atomic_addx2/x4_elem_op` are declared with `num_inputs = 3` -- the
vectorizer was simply under-filling the op. The order is copied as-is
rather than visited, since it is a lane-invariant scalar that must not be
broadcast across the vector lanes.
Carrying the order also wakes up inline PTX that was previously
unreachable: with the order always defaulted to relaxed, the non-relaxed
branches were constant-folded away before ptxas ever saw them. Three
problems surface once they are live, all verified with ptxas 13.3:
1. `atom` with vector operands requires sm_90, so the ordered fp16 pair
failed to assemble on sm_80/sm_86/sm_89 -- and fp16 vectorizes to x2 on
every target, so that is every ordered fp16 region atomic_add. Use the
packed `f16x2` form instead of `.v2.f16`: it encodes the same pairwise
add but is available from sm_60, so the wide op stays a single atomic
instruction on every supported target rather than needing an arch split.
sm_89 SASS confirms a native ATOMG.E.ADD.F16x2.RN.STRONG.GPU. This also
drops the per-half pack/unpack, since a half2 reinterprets directly to
the 32-bit operand.
2. The bf16 v2 asm was missing the `.noftz` modifier that ptxas requires
for `.bf16`, so `atom.{release,acquire,acq_rel}.gpu.global.add.v2.bf16`
failed to assemble on every arch, sm_90 included. The scalar bf16
helper next to it already spells it correctly.
3. bf16 has no packed equivalent below sm_90 -- `bf16x2` and `.v2.bf16`
both require sm_90, as does scalar `atom.add.bf16` -- so unlike fp16 it
genuinely cannot express the ordering in one instruction there. Guard
the sm_90 path and emulate the ordering below it with device-scope
fences around a relaxed bf16x2 atomicAdd. `__threadfence()` lowers to
`fence.sc.gpu`, the same device scope as the `.gpu` qualifier on the
ordered PTX, and is two-way where release and acquire need only one-way,
so this is conservative rather than lossy. `AtomicAddx2Ret` takes the
same guard, since it shares the helper. Nothing reachable instantiates
the bf16 Ret overloads with a non-relaxed order today, so that path has
no test; it is guarded to keep the file consistent and to stay correct
once region-form `return_prev=True` is implemented.
Add a regression test in fp16 and bf16. It disables the kernel cache,
since the key does not cover the native library by default and would
otherwise serve stale pre-fix source, and it runs the kernel so the
ordered wide path is checked for correctness and not just for the emitted
argument. The bf16 case covers the pre-sm_90 fenced branch on Ampere/Ada
and the sm_90 instruction on Hopper.
Also correct the stale memory_order sentence in atomic_max and
atomic_min, which forward the order through the shared region lowering
just like atomic_add does.
Fixes#2688
Signed-off-by: Srijan Keshri <srijankeshri007@gmail.com>
[CUDA] Add FP8 min/max overloads
CUDA codegen emits unqualified min/max calls for TIR MinNode and MaxNode
expressions. FP8 operands promote to float, but the FP8 wrappers have no
implicit conversion back from the result, so T.max, T.min, and T.clamp
failed during nvcc compilation.
Add tl::min/max overloads for float_e4m3_t and float_e5m2_t. The overloads
compute through fminf/fmaxf and convert the result back to FP8. ADL lets the
existing scalar and per-lane vector codegen resolve these overloads without
FP8-specific codegen branches.
Add runtime coverage for both FP8 formats and compile coverage for vectorized
expressions and the e4m3fn spelling.
Fixes#2985
T.__ldg failed to compile for float16 and bfloat16 buffers because CUDA
provides __ldg overloads for __half and __nv_bfloat16, while TileLang uses
CUTLASS half_t and bfloat16_t for these scalar types.
Add forwarding overloads in common.h that reinterpret the bit-compatible
pointers as the native CUDA types and wrap the loaded values back in the
CUTLASS types. The native 16-bit __ldg is available on SM60+, so no
architecture-specific fallback is needed.
Update the stale 32-bit-only API wording and add exact round-trip coverage
for both dtypes.
* [Lang] Reducer v2: first-class deferred reduction epochs with planned physical lowering
Fixes#2408 at its root. v1 implemented reducers as fragment buffers with
a special layout, so no pass owned the contribution-multiplicity
semantics: loop-layout replication silently multiplied contributions
(128 threads summing 8 elements returned 576 instead of 36), the
finalize width had to be guessed from the storage layout's
ReplicateExtent, and every compiler stage carried buffer-identity
special cases to compensate (see #2897, #2881).
Reducer v2 makes the reducer a first-class deferred reduction epoch.
Core invariant: `T.reducer_update` is the single source of contribution
multiplicity; physical layout can never change how many times a logical
contribution is combined.
New IR contract (`local.reducer` virtual scope + three first-class ops):
acc = T.alloc_reducer(shape, dtype, op="sum", seed=None)
T.reducer_init(acc)
T.reducer_update(acc[indices], contribution) # inside T.Parallel
T.finalize_reducer(acc, dst) # out-of-place
Pipeline:
* VerifyReducerEpoch (early): lifecycle state machine + opaque-access
rules; ordinary reads/writes, clear/fill, aliasing, updates outside
T.Parallel, double init/finalize are compile errors.
* ReducerPlanAndMaterialize (after LayoutInference; loop layouts are
read-only inputs): decides storage, communication, and execution
multiplicity per epoch, then materializes the first-class ops into
ordinary IR. Two physical plans:
- Narrow plan: project the reduction axes (parallel vars absent from
the update indices, plus loop replication) out of the update site's
loop layout to obtain compact per-thread partial storage; updates
run unguarded on every replica; the collective covers only the
thread-expression splits sourced from reduction axes (reusing the
ComputeReducerLayout / CollectThreadReduceSteps machinery T.reduce
is built on). Zero splits means zero collectives (LocalComplete).
Proof obligations: direct var-to-dim indices, replica-safe
contribution values, full-block coverage, a single power-of-two
step, structurally equal plans across sites, destination-layout
containment (with a controlled override for provably unconstrained
single-slot destinations).
- Wide plan (FullParticipant baseline, always available): one full
logical-shape partial per participant, canonical-replica-guarded
updates, participant-wide AllReduce, seed combined exactly once
post-collective. Any narrow-plan proof failure falls back here
deterministically - never to a compile error or wrong code.
* Generic multiplicity marker `tl.parallel_multiplicity`: PartitionLoop
lowers it to a REP == 0 guard knowing nothing about reducers; the
fully_replicated_reducer_buffers compatibility hook is retired.
* FinalizeReducerOp takes explicit plan operands (reducing_threads,
scale), closing the derive-width-from-ReplicateExtent channel.
* VerifyReducerConsumed (after LowerTileOp): surviving v2 constructs
are compile errors, never silently-wrong code.
* Pass config `tl.reducer_force_baseline` forces the wide plan for
differential testing.
Legacy compatibility and v1 removal: CanonicalizeLegacyReducer rewrites
v1 syntax (clear + read-modify-write stores + in-place finalize) into
v2 ops on a strict whitelist; non-identity fills on idempotent combines
are forwarded as one-time seeds, non-zero sum fills are rejected, batch
annotations are forwarded, and post-finalize reads are redirected to a
synthesized destination. With the shim in place, the entire v1 lowering
path is deleted: LayoutReducer, the PartitionLoop/LowerParallelLoop
reducer-buffer plumbing, the ParallelOp rep=ALL special cases, and the
vectorizer's reducer parameter chain.
Testing (H100, kernel cache disabled): 128 passed / 2 skipped across
the legacy reduce suite (93, unchanged, through the shim), the new
reducer v2 suite (22, incl. the 36-vs-576 regression, narrow-plan
codegen assertions, lifecycle diagnostics, forced-baseline
differentials), and the copy suite; the pipelined GEMV example (legacy
syntax) remains numerically correct.
* [Refactor] Consolidate reducer ops into src/op/reducer.{h,cc}
FinalizeReducerOp is no longer a standalone legacy op — it is the
collective that ReducerPlanAndMaterialize emits — so it moves next to
the first-class reducer ops. Its combine-op field now uses the same
ReducerV2OpType enum as the planner, removing the implicit
static_cast coupling between two parallel enums.
The former LayoutReducer pass files are gone entirely: the
ReducerInfo/ReducerRepType annotation types had no consumers left, and
the legacy `reducer_info` annotation key (still read by the
canonicalization shim and the data-race verifier) moves to
op/reducer.h. Stale includes in layout_inference / lower_tile_op are
dropped.
No functional change; 115 reducer tests pass.
* [Doc] Clarify the intent of tl.reducer_force_baseline
The switch is not a workaround for expected narrow-plan bugs; it exposes
the reducer design's canonical reference lowering for differential
testing, field escape-hatch use, and plan-choice A/B measurement.
* [Layout] Opaque tile-op loops must not self-plan their layout
A parallel loop whose body contains an opaque tile-op call (e.g. a
reducer update) has no visible write to anchor layout inference. In
free mode such a loop used to self-plan a partition, and the
register-count-based attempt selection then preferred it: the plan
minimizes this loop's own footprint (e.g. one x-vector element per
thread) while silently forcing a column-major layout onto the fragment
operands it reads, degrading their producing copies from vectorized to
scalar accesses. On a 4096x4096 fp16 GEMV this cost 40% end-to-end
versus the legacy reducer lowering.
Fix (generic, not reducer-aware): when a loop carries an opaque
tile-op call, prefer propagating its layout from a fragment operand
that is addressed purely by the parallel loop vars, deferring while
such an operand is still unsolved so the producing ops fix the
operands first. Self-planning remains the fallback when no operand can
drive the loop (e.g. operands indexed by inner serial reduction vars,
which require the loop to own the operand layout instead) or when
operand propagation cannot express the loop's iteration space.
With the fix the GEMV benchmark matches the legacy lowering (0.075 ms
vs 0.074 ms) and narrow-plan row reductions stay ahead (AllReduce<16>
with one partial slot per thread, ~6% faster than legacy); 129 tests
pass across the reducer v1/v2 and copy suites.
* [Lang] Reducer v2 stage 4: narrow-plan coverage, packed accumulation, bitwise ops
Narrow-plan coverage and observability:
- tl.enable_reducer_plan_verbose pass config: INFO-level per-epoch plan
decisions (narrow/wide, collective widths, rejection reasons).
- Multi-step collectives: interleaved thread splits lower to one AllReduce
per step instead of falling back to the wide plan.
- Destination override generalized: multi-slot destinations and transitive
copy chains (e.g. fp32->fp16 staging hops) adopt the induced layout when
the whole chain is unconstrained; ParallelOp candidate selection falls
back to the operand-compatible layout when the register-preferred
candidate conflicts with an overridden buffer.
- Update-index generalization: permuted loop-var indices and constant unit
dims (acc[j, i], acc[0, i]) rebuild the induced fragment over the
reducer's own dim order; all-constant scalar indices keep the synthetic
unit input ComputeReducerLayout leaves behind.
- Seed and batch now work under narrow plans: the finalize lowering
combines the seed once per slot after the collective, and explicit-width
steps compose with batched AllReduce.
Packed partial accumulation (16-bit floats):
- Single-site narrow epochs on fp16/bf16 split each per-thread partial
into two lanes keyed by the parity of an on-thread reduction source (the
innermost enclosing serial loop, or the innermost parallel reduction dim
when its low bit stays on one thread), halving the serial combine
dependence chain.
- A per-thread fold recombines the lanes into the plain induced storage
before the (unchanged) collective, so the plan's communication and
destination proofs are untouched. Lane parity never changes which thread
a contribution lands on, so the plan stays correct under any layout.
Bitwise combine ops:
- alloc_reducer(op="bitand"/"bitor"/"bitxor") for integer dtypes under
both plans (identities: all-ones / 0 / 0). Float dtypes are rejected at
allocation time; the legacy v1 path stays sum/max/min only.
Tests: reducer v2 suite grows to 38 (packed structural + numeric +
differential, bitwise under both plans, permuted/unit-dim/scalar indices,
narrow seed); the v1 codegen batch tests pin tl.reducer_force_baseline
since narrow plans legitimately skip run_batch. v1 suite stays green (93).
* [Lang] Move the reducer seed to T.reducer_init(acc, init)
The logical starting value used to be declared at allocation time
(alloc_reducer(seed=...)) and smuggled to the planner through the
reducer_info_v2 block annotation. Declaring it where the epoch opens
reads naturally ("the reduction starts from this value") and makes it a
first-class operand of the init op:
acc = T.alloc_reducer((M,), dtype, op="sum")
T.reducer_init(acc, 100.0) # optional; default = combine identity
- ReducerInitOp accepts an optional second argument carrying the value;
the annotation now carries only the combine op.
- VerifyReducerEpoch checks the value's dtype at the init site (Python
numbers are auto-converted by the frontend).
- CanonicalizeLegacyReducer forwards v1 non-identity idempotent fills as
the synthesized reducer_init's argument instead of via the annotation.
- Semantics unchanged: physical partials always initialize to the
identity; the value is combined exactly once per logical output at
finalize time, so physical replication can never multiply it.
Tests: seed tests migrated to the new API plus a dtype-mismatch
diagnostic (v2 suite now 39); v1 suite stays green (93).
* [Layout] Name the operand-driven-loop contract: kTLPerIterationOp
The opaque-loop rule (6749631b) worked but read poorly: detection
sniffed the tile-op registry (minus a tl.region special case), the
policy lived in comments spread over three code blocks, and "can this
operand drive the loop?" was answered twice with different machinery
(a try/catch around the generic source-buffer choice, plus a separate
drivability scan for the defer decision).
Restructure without changing the architecture:
- New op attribute kTLPerIterationOp declares the contract at the op
definition site: the op executes once per iteration inside T.Parallel
and its region-mediated effects are invisible to statement-level
scans. tl.reducer_update is its only member today. Both consumers now
read the attribute instead of keeping their own lists: ParallelOp's
detection (previously registry sniffing) and the nested-loop
checker's _PARALLEL_SAFE_TILE_OPS whitelist (deleted).
- Free-mode inference for such loops is now a single decision:
FindOperandDriver returns the best solved drivable operand (indices
pure in the loop vars, most index dims first) or reports a pending
one; the loop adopts, defers, or falls through to self-planning. The
exception-as-control-flow try/catch is gone from the policy path (a
small guard remains for non-affine driver indices, which drivability
does not test). One behavioral refinement: when the generic
most-indices read buffer is not drivable but another solved drivable
operand exists, the loop now adopts that operand instead of falling
back to self-planning.
Verified: reducer v2 (39) + v1 (93) + copy suites pass; the reducer
benchmark (rowsum fp32/fp16, GEMV fp16) reproduces identical plans and
timings to the previous commit.
* [Lang] reducer_update is a per-iteration builtin intrinsic, not a tile op
reducer_update never was a tile op in substance: it executes once per
iteration INSIDE T.Parallel (no other tile op does), owns no layout (the
enclosing loop and the planner decide physics), and never lowers on its
own. Registering it as one forced two categories of friction: consumers
that had to exempt it (the nested-loop checker whitelist, ParallelOp's
registry sniffing before kTLPerIterationOp), and an argument convention
(tl.region wrapping) whose only real payload was the multi-dim indices.
Re-register it as the builtin `tl.reducer_update`:
- args are (acc[indices], value) with a plain BufferLoad target — an
update descriptor keeping structured indices for the planner. Every
access scan now sees the acc read for free; the write stays
undeclared, which is sound for reducers specifically: updates commute
(combine ops are associative/commutative) and VerifyReducerEpoch pins
init/finalize to straight-line code, so no cross-statement write
ordering exists to lose inside pipelined/parallel regions.
- ParseReducerUpdate replaces the TileOperator ctor; verifier, planner
and the legacy shim consume it. The verifier visits the target's
indices but not the target load itself (a reducer load anywhere else
is still the usual compile error).
- Pipeline statement classification counts kTLPerIterationOp calls as
compute (BufferRegionCollector, IsPureCopyStmt/GetSinglePureCopyTileOp
would otherwise classify an update-only loop as copy-like now that the
call no longer parses as a tile op); the shared IsPerIterationOpCall
helper also replaces ParallelOp's inline attribute lookup.
- The nested-loop checker exemption is gone entirely: a builtin has no
TLOpBuilder, so it passes the tile-op ban by construction.
Verified: reducer v2 (39) + v1 (93) + copy suites pass; the reducer
benchmark (all cases T.Pipelined, exercising the pipeline dependency and
classification paths) reproduces identical plans and timings.
* [Cleanup] Drop the TLPerIterationOp attribute for a direct predicate
The attribute named a category that does not carve reality at a joint:
nearly every builtin intrinsic "executes per iteration" inside
T.Parallel, so the name suggested a structural property while actually
being an opt-in marker with exactly one member. What its two consumers
really key on is reducer_update itself — the hidden write to a
deferred-layout buffer (ParallelOp inference) and compute-vs-copy
statement classification (pipeline planning).
Replace it with IsReducerUpdateCall next to the op declaration; both
consumers and ParallelOpNode::has_reducer_update_ now say plainly what
they check. A future sibling op (e.g. a rescale intrinsic) extends the
one predicate.
* [Doc] Record the operand-driven defer protocol's known limitation
Deferring assumes the pending drivable operand will be solved by its
producer. An operand with no producer at all (an update reading a
fragment nothing ever writes) defers on every free-level round and
surfaces as a compile-time "layout can not be inferred" ICHECK — never
wrong code, and only on degenerate input. A sound fix needs the
inference engine to order solving (or guarantee a final round); that
redesign is deliberately deferred to a follow-up rather than patched
into the protocol here.
* [Layout] Drop the operand-driven defer protocol from ParallelOp
The three-piece workaround (has_reducer_update_ detection, the
FindOperandDriver three-state ladder, free-mode defer) patched solve
ordering for anchor-less loops from inside a single op. The real defect
is the engine's register-count attempt scoring: the bias row-broadcast
counterexample reproduces the same mis-selection with no anchor-less
loop involved, so the protocol treated a symptom. Direction for the
proper fix (IO-aware cost model + declared solve ordering) is recorded
in the layout RFC.
Reducer-update loops return to the generic two-candidate free
inference. The candidate-validation fallback still bounds the damage:
fp16 GEMV 0.0615 ms vs 0.0736 ms v1 baseline (down from 0.0346 ms with
the protocol), rowsum unchanged across all sizes. reducer v2 (39),
v1 reduce (93), atomic, and copy suites pass.
* [Cleanup] Unwrap the free-inference branch's vacuous layout guard
The inner !loop_layout_.defined() check was the seam where the deleted
operand-driven block could adopt a driver layout before the generic
two-candidate path ran. The enclosing branch already requires an
undefined layout and nothing assigns one in between, so the guard is
always true now.
* [Cleanup] Drop the unreachable post-selection candidate fallback
ChooseBestCandidate already validates both candidates with the same
flags and returns a valid one whenever either validates (single-valid
returns it directly, both-valid picks among valid, both-invalid leaves
nothing to fall back to). The post-hoc re-validation could therefore
never flip the choice; an instrumented run across the reducer suites
and the rowsum/GEMV benchmarks confirms it never fires. It was added
alongside the destination-chain override work and mistakenly credited
for making multi-slot overrides safe — the chain registration is what
did that. Comments referencing the fallback updated.
* [Cleanup] Tighten loop_partition after the multiplicity-marker rework
- MultiplicityMarkerLowerer only rewrites statement-level AttrStmt
markers; derive from StmtMutator instead of StmtExprMutator so
expression subtrees are not traversed.
- Drop a duplicated example paragraph in the partition guard comment.
- Discard InverseWithLevel's IterMapLevel explicitly instead of
keeping an unused pair binding.
* [Test] Cover reducer accumulation inside a T.Pipelined body
Pipeline planning must classify the per-iteration update intrinsic as
compute while the shared staging copy forms the only copy stage; the
epoch spans every pipeline iteration with a single finalize. This path
previously had no coverage at all.
Warp specialization is disabled in the test: a WS-split epoch is out of
scope for reducer v2 and currently surfaces as the finalize lowering's
thread-bounds ICHECK rather than a user-facing diagnostic.
* [Cleanup] Revert the pipeline-planning reducer classification marks
The three IsReducerUpdateCall sites marked update intrinsics as compute
during pipeline statement classification. Empirically they change
nothing: with the marks compiled out, both a plain pipelined
accumulation and an adversarial mixed statement (copy-shaped shared
store plus an update in one parallel loop) produce byte-identical
kernels — a bare intrinsic never classifies as a copy on any consumer
path that matters while warp-specialized reducer epochs are rejected
outright. pipeline_planning.cc returns to its upstream state, and the
now-unused IsReducerUpdateCall predicate is dropped with it.
* [Lint] Apply pre-commit clang-format
The multi-segment recursion in tcgen05_{ld,st}_core advanced the
register pointer in dst_t/src_t elements while segment lengths count
b32 columns, so a 16-bit copy whose column count is not a power of two
(e.g. 48 columns = one x32 + one x16 segment) accessed the tail
segments' registers at half the intended offset -- a silent
miscompilation. Advance the pointer in b32 units instead.
Add asymmetric whole-store/whole-load roundtrips at (128, 96) bf16: a
symmetric roundtrip cancels the unit error, so exactly one side is
split into power-of-two (single-segment) slices, catching each
direction independently.
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Aliasing buffers must induce the same storage-bit map,
bit(k) = flat_output(k) * elem_bits: layout outputs are measured in the
buffer's own elements, so a reinterpreting view needs unit conversion on
the last (stride-1) output mode, which the generic flat-index reshape
cannot express (it drops the position inside a storage slot).
TryReinterpretReshape implements the rule uniformly for every element
width, replacing the packed-sub-byte special case:
- Narrowing by f: old element q becomes new elements f*q .. f*q+f-1,
always representable as last_output * f + lane. The sub-byte
lane-append/collapse forms are deleted: the trailing lane mode was
only ever a legality certificate for the widening direction.
- Widening by f: representable only when every aligned group of f old
elements sits contiguously in one aligned slot, now proven per
reshape (plus thread-mapping invariance for fragments); layouts that
interleave storage below the new element width are a hard error
instead of silently miscompiling.
* [Enhancement] Expose cluster_mask on T.tma_copy
TMA multicast already composes with tma_copy's split-phase contract: in
LowerBulkCopy, use_multicast is orthogonal to the flag that suppresses the
wait, so annotations={"cluster_mask": ...} produces an issue-only multicast
today. Only the named parameter was missing.
Promote it to a documented keyword argument, matching the parameter that
copy_cluster already exposes under the same name, and write down the
semantics: the lowest-ranked CTA in the mask issues the multicast, other
in-mask CTAs issue nothing, a CTA outside the mask falls back to a unicast
load, and every CTA runs its own expect_tx against its local barrier.
No lowering change.
* [Testing] Drop the duplicated TMA store arrive/wait
T.copy already lowers shared -> global to tma_store + tma_store_arrive +
tma_store_wait, so the explicit pair after it emitted a second, redundant
arrive/wait.
On SM70-79, normalize_atomic_type leaves bfloat16_t unchanged, but CUDA
provides no atomicAdd overload for it. Scalar, return-prev, and vectorized
bf16 atomic adds therefore failed to compile; AtomicAddRet also fell through
to unsupported cuda::atomic_ref<bfloat16_t>::fetch_add.
Implement atomicAdd(bfloat16_t*) with a 16-bit CAS loop, route AtomicAddRet
through it, and add a per-element scalar fallback for x2 that the existing
x4 path composes below SM80. Guard the overload to SM70+ so unsupported
16-bit atomicCAS does not affect older targets. Add contended numerical
coverage and an explicit sm_75 compile-only regression for the scalar,
vectorized, and return-prev helpers.
This PR adds a new set of annotations for describing arbitrary warp specialization scheme, with the abstraction of pipelines. A pipeline supports 4 kinds of operations: producer_acquire, producer_commit, consumer_wait, and consumer_release. Behind the scenes, mbarriers are created to implement the synchronization between warps.
This PR introduces a new pass, MaterializeWSSchedule, that transforms an annotated kernel to a warp-specialized kernel.
The set of infrastructure enables further development of fully automatic warp specialization in TileLang. I am working on this.
* [Layout] Support half-subpartition (M=64) TMEM tiles in tcgen05.ld/st
A PTX Layout F fragment -- the 1SM M=64 shape make_tmem_frg_a/c selects -- occupies only the LOW 16 datapaths of each 32-datapath sub-partition, so its four warp groups sit 32 datapaths apart with nothing in between. ExpandTcgen05Layout assumed every tile covers 128 contiguous datapaths, so no instruction could be planned for one: an M=64 TS GEMM could not get its A operand into TMEM at all, dying with "Failed to find a suitable instruction for tcgen05.st". Tracks #2877.
No new PTX and no new atoms. The 32dp wrappers already issue the 16x{64,128,256}b shapes twice, once per half, so this adds wrappers that stop after the first issue -- and whether a tile wants one form or the other is a property of the TILE, so the planner derives it rather than the metas carrying a flag: right_inverse's maximal contiguous run is 128 when the fragment fills every datapath and exactly the atom's own extent when it is Layout F. One meta per width still covers both, and the plan reports the datapaths one warp ends up covering -- the number the wrapper is named after, which codegen prints straight into tcgen05_{ld,st}_NNdp*Nx.
Replication is one explicit stamp -- what a single wrapper call covers -- rather than a blocked product whose complement happens to enumerate the gaps in the order the wrapper wants them. That coincidence holds only for a dense tile: for Layout F the complement's fastest gap is the unoccupied datapaths above the atom, so it would pack the warps in at stride 16 instead of one sub-partition apart. Spelling the stamp out costs a few lines and serves both.
Mapping it back to logical indices takes two inverses, because two kinds of tile arrive. inv_prefix spans the maximal contiguous chunk -- one whole issue for a full-datapath tile, and valid even where the tile is not injective, as a pack::16b tile is not. A Layout F issue reaches four datapath groups 32 apart so no contiguous prefix covers it, but such a tile IS injective, and left_inverse is then a genuine physical -> logical map with no contiguity requirement. Either way the placement is derived from the tile, not assumed of it.
Which datapaths of a sub-partition a warp covers is read off the tile too -- the run along the datapath axis alone, clamped to the 32 a warp owns -- so an atom too wide for it (32dp32b against Layout F) falls out as a non-division rather than needing a case of its own, and the answer does not depend on whether the tile's columns count values or b32 slots.
Two things a 16-bit Layout F operand gets wrong if taken from the old code:
pack::16b was judged by "does the fragment cover its codomain". Layout F leaves gaps on the datapath axis, not inside its columns, and was misread as half-filled -- emitting the modifier for a tile whose columns are full. TmemFragmentNeedsPack16b now asks the intended question directly: is every column of every occupied datapath full?
a partial sub-partition cannot be planned in value columns. Only the width-1 32dp32b atom is insensitive to the unit; the wider atoms pick their b32 column with a LANE bit, so stepping their columns by a value column splits each register's pair across two columns, transposing K with the column's low bit and silently permuting the A operand as [0,2,1,3] within every group of four. The stamp therefore steps the atom by a whole b32 column and carries the packed partner as its own fastest value axis. Only how many tile columns share a b32 column has to be passed in: a value-column layout cannot say that about itself.
Verified on B200: an M=64 TS GEMM with A staged through registers into TMEM now matches torch for K in {64,128,256,512} (and M=128 is unchanged); recovering the operand the MMA actually read is the identity, where it was the permutation above. Layout F register->TMEM->register round trips are bit-exact for fp32 and packed bf16 over one and two warpgroups; the fragments are proven against the CUTLASS Copy_Traits DstLayout formulas symbolically, for all lanes/registers/repetitions/warps at once; full-datapath tiles still lower to the identical 32dp instruction.
* lint
* [CUDA] Make the tcgen05_ld/st datapath argument optional
A six-argument call written before the datapath count existed still
means the sub-partition-filling 32dp wrappers, so hand-written
intrinsic calls keep working.
* [Layout] Derive the tcgen05 issue footprint from per-axis projections
Columns per issue now come from right_inverse of the column projection,
as review suggested: a Layout F tile's gaps live only on the datapath
axis, so the column projection is gap-free and the derivation needs no
partial-subpartition special case. The datapath footprint is collected
mode-wise and both sides join into one chunk that doubles as the
rest-domain tiler and the dense physical->logical map.
A batched Layout F tile -- per-batch column gaps on top of the datapath
gaps -- now plans one issue per contiguous column run instead of
finding no plan.
* [CuTeDSL] Support tcgen05 TMEM copies incl. half-subpartition tiles
The cutedsl runtime had only the four 32dp loads and none of the other
pieces a TMEM kernel reaches: tcgen05.st was missing entirely, the
before/after_thread_sync fences were emitted but never implemented, and
ptx_tcgen05_mma_ss/ts had drifted an enable_2cta argument behind the
lowering. Add the full store set (32dp and 16dp), the 16dp loads and
the fences, reject 2-CTA MMA explicitly, and fix the chained-segment
column advance to step seg_x * columns-per-x (2/4/8 for the
64/128/256b shapes). Roundtrip and TS-GEMM tests cover both wrapper
families on sm100.
* Refactor
---------
Co-authored-by: Yongqi Zhuo <yzhuo@deepseek.com>
* [BugFix] Preserve invariant store scalarization
Classify vectorization constraints explicitly so broadcast loads, memory accesses, local accesses, and semantic scalarization requirements cannot be conflated by scope-based bucketing.
* [Refactor] Preserve vectorization planner diagnostics
Move the original strategy rationale into the new classification helpers and retain the established verbose summary and strategy log messages while adding separate constraint details.
* [Docs] Restore vectorization strategy overview
* [BugFix] Guard replicated reducer updates
Use only the canonical T.Parallel replica when updating fully replicated reducer buffers, preventing over-counting across differently replicated fragment layouts. Add short-extent and mixed-extent numerical regressions.
* [Refactor] Use ordered reducer buffer lookup
Keep the small reducer buffer collection as an ffi::Array and perform identity membership checks with Buffer::same_as, avoiding unordered STL in loop partitioning.
* [Refactor] Clarify replicated reducer guard naming
* [Docs] Mark reducer partition hook as temporary
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
A large number of tests import CUDA-only internals directly (for example
`tilelang.cuda.pipeline.CUDAPassPipelineBodyPrologue`) or exercise
NVIDIA-specific features such as TMA, LDGSTG, mbarrier, warp specialization,
tcgen05 and tfloat32. `src/cuda/CMakeLists.txt` only compiles a small
"ALWAYS" subset when `USE_CUDA=OFF`, so those FFI symbols are never
registered and the tests fail with:
AttributeError: module 'tilelang.cuda._ffi_api' has no attribute
'AnnotateDeviceBoundTmaCopies'
This is not ROCm-specific -- the same tests fail on any build without the
CUDA backend, including Metal and CPU-only.
Add `@tilelang.testing.requires_cuda` to those tests, plus three
portability fixes that keep coverage on non-CUDA backends instead of
skipping it:
* `test_tilelang_transform_layout_inference.py`: CUDA emits
`__launch_bounds__(N, 1)` while HIP emits `__launch_bounds__(N)`.
Match both.
* `test_tilelang_transform_pipeline_planning.py`: PipelinePlanning
deliberately strips pipeline annotations on ROCm targets other than
gfx950 (see the guard in src/transform/pipeline_planning.cc), so assert
the documented sequential fallback there.
* `test_tilelang_kernel_gemm.py`: the two f32f32f32 cases use `T.tfloat32`,
which is NVIDIA-only, so narrow `requires_cuda_or_cdna` to
`requires_cuda`.
Also refresh three stale RDNA tests: gfx12 is now a supported generation
(`get_arch` accepts 11 and 12), so they can no longer assert that gfx1200
is rejected. Force an unsupported generation instead, which is the only
way to reach the guard now that `target_is_rdna` reports only gfx11/gfx12.
Verified on 8x MI300X (gfx942, ROCm 7.2 container, torch nightly rocm7.2):
the ROCm suite goes from 105 failures to 0 with no loss of passing tests.
* [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>
Yongqi Zhuo