mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
30c1904d3a4e0b8320fb1807456e30ec3c2a6891
276
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Yongqi ZhuoandClaude Fable 5
|
30c1904d3a |
[CUDA] Unify TMA copy lowering on CuTe algebra (#3106)
* [CUDA] Unify TMA copy lowering on CuTe algebra and support general swizzled layouts Every descriptor-based TMA path (bulk copy, atomic add, gather4/scatter4, im2col) now shares one CuTe-algebra box decomposition. Boxes follow the driver contract (inner box bytes at most the swizzle span, per-dim caps), and the recovered swizzle is widened on demand under an equivalence proof, so general annotated layouts -- including phase-shifted swizzles -- lower correctly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Address review concerns --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
andyluo7
|
ab3202a25c | [ROCm] Preserve FP8 bits in warp shuffles (#3104) | ||
|
Chennes
|
035a2af35b |
[CUDA] Add 16-bit overloads for CUTLASS fast-math functions (#3097)
math.h aliases several half-style names to cutlass::fast_* when a kernel uses selected fast-math intrinsics. CUTLASS provides only float and double overloads for fast_log, fast_sqrt, fast_sin, and fast_cos, and adds half_t overloads only for fast_exp and fast_tanh. Overload resolution can therefore leave a float result that cannot copy-initialize the CUTLASS wrapper. Once a kernel includes math.h, for example through T.__exp, these aliases also affect matching plain math calls in the same kernel. Their compilation could therefore depend on whether another intrinsic pulled math.h in. Add half_t and bfloat16_t overloads for fast_log, fast_sqrt, fast_sin, and fast_cos, plus bfloat16_t for fast_tanh. Each evaluates in float32 and converts back, following the existing fast_exp(bfloat16_t) bridge. Calls that already compiled retain the same numerical path. Add compile-and-run coverage for both wrapper types. |
||
|
Lei WangandClaude Fable 5
|
6279e735f9 |
[Refactor] Give reducers a first-class PartialFragment layout solved by layout inference (#3093)
* Steer reducer destinations during layout inference * Remove extra reducer steering attempt filtering * handle reducer problem3 * bug fix * [Refactor] Give reducers a first-class PartialFragment layout solved by layout inference The reducer's physical plan (narrow vs wide partials) used to be decided after inference by ReducerPlanAndMaterialize, with dst-steering emulating the missing propagation through side channels (wake tables, hint plumbing, ownership gates); the same induced layout was computed twice and kept consistent by convention only, and the reducer split the use-def graph into disconnected components. This makes the reducer's layout a first-class solver value that fully determines the physical plan: * New layout kind `tl.PartialFragment` (PartialFragmentNode : FragmentNode): Fragment algebra whose replication coordinate carries an explicit combine decomposition — under a low-bits convention enforced at construction, `_rep % combine_size` enumerates addend lanes the finalize collective must reduce and `_rep / combine_size` enumerates equal-value copy groups. (storage algebra, combine width) is the complete plan: the collective steps are derivable from the node alone (CombineSteps), the FullParticipant object (combine == replicate == participants, thread == _rep) is the wide fallback, and combine width is part of layout identity (same-storage different-decomposition proposals widen at the commit point instead of unifying silently). A boundary-preserving condensation (CondenseReplicateVarKeepingBoundary) keeps the convention true through projection; kind-aware IsEqual in both dispatch directions prevents silent unification with plain Fragments. * Update nests (ParallelOp) record their reducer_update sites and, once their loop layout solves, propose the induced partial (per-site AnalyzeReducerUpdateSite proofs) or the participant-wide plan. The commit point resolves multi-site disagreement with a monotone widen-on-conflict rule; zero-update epochs are seeded wide before any inference level runs; tl.reducer_force_baseline pins the proposal. * finalize_reducer reads the solved partial and steers an unconstrained destination to its post-collective reading. The wake table, the LayoutInferArgs.reducer_update_sites hint plumbing and the reducer_use_list_ component glue are deleted: reducers travel ordinary use_list_ buffer edges. * ReducerPlanAndMaterialize reads the plan off the node: storage = AsPostCollective(), steps = CombineSteps(), wide lowering iff the node is the FullParticipant object. TryNarrowPlan is retired; what remains (DeriveNarrowLowering) derives only the lowering details that live outside the layout — the packed-lane heuristic and the destination-containment gate (which still demotes gracefully to the wide plan, preserving the unsatisfiable-steering fallback contract). The destination override chains are kept as a verified-no-op fallback with retirement telemetry. * T.annotate_layout accepts a PartialFragment on reducer buffers, pinning plans the automatic solution cannot reach (e.g. a middle plan between narrow and wide). Annotated partials are strict: a disagreeing update site raises with the narrow-rejection reason, and ComputeReducerPinnedCandidate back-propagates the annotation through the update loop into its input fragments. * Layout cost models price register-array spills: a fragment access whose physical slot index depends on the thread variable (detected by the shared FragmentThreadIndexProbe, extracted from HasThreadDependentFragmentIndex) demotes the whole array to local memory, charged as byte-denominated traffic in both models' `mem` channel so it competes with — rather than vetoes — global-bandwidth estimates. ChooseBestCandidate's anti-spill preference now applies regardless of where the loop stores (a fragment->fragment publication copy spills exactly like a fragment->global one). With generation and pricing combined, inference produces the wide publication chain's replicated layouts on its own: the reducer_v2 suite and the mhc_post_bwd production kernels pass with the destination override fully disabled. Testing (all with the kernel cache disabled): reducer_v2 suite 70 passed (new coverage: pinned middle plans, annotation steering input fragments, conflict diagnostics, combine-step roundtrip, same-storage different-combine widening at commit); language suite 1261 passed; layout-inference cost-model tests and maint/layout_inference goldens unchanged; generated CUDA for existing reducer kernels is byte-identical; mhc_post_bwd correctness passes and its six benchmark sizes show no latency change with zero STL/LDL in SASS. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Yongqi ZhuoandClaude Fable 5
|
0ce2a7bacf |
[Bugfix] Make TMA layouts region-aware to keep slices contiguous (#3089)
[Bugfix] Make TMA linear layouts region-aware to keep versioned slices contiguous The inferred "linear" shared layout hoisted every 256-split quotient above all dims, so a pipeline-versioned slice became two gapped chunks: LowerBulk's bijection check fired on warp-specialized wide reductions, and the TMA atomic-add path (which has no such check) silently read the other version. MakeTmaLinearLayout now takes the copied region and orders modes as [region-fixed dims, box quotients, box contents], so a fixed slice owns one contiguous run of whole TMA boxes and the existing rest-loop machinery emits one instruction per box over the true (unaligned) global shape. The atomic-add path shares the helper and computes the region base offset through the layout instead of hand-rolled row-major strides over the physical shape. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Yongqi ZhuoandClaude Fable 5
|
728f2c88a1 |
[Bugfix] Fix wait parity for explicit mbarriers in pipelined loops (#3087)
A synchronous T.gemm(mbar=...) inside T.Pipelined waited on the wrong mbarrier parity. InjectSoftwarePipeline annotated every copy/im2col/gemm call with the ring parity (k / num_stages) % 2, but that law belongs to the compiler-generated pipeline_mbar ring: an explicit mbarrier whose arrive happens inside a tile op is deliberately excluded from ExpandPipelineBarriers, so it remains a single slot that flips once per iteration. The auto-emitted wait asked for a parity that had already completed, returned immediately, and the next iteration's copies overwrote a shared-memory stage the still-running MMA was reading (wrong results observed at num_stages >= 3). Annotate only the ops that consume the phase -- those that auto-emit a paired arrive+wait on an explicit mbarrier: sync gemm with an mbar, and copies carrying an explicit "barrier" that are not is_tma_copy -- and annotate them with the logical iteration parity (k - min) % 2, written against the original loop var. EmitImpl's existing per-block Substitute(loop_var -> normalized_access_index) specializes the expression for every prologue/body/epilogue clone, since the tirx mutators already rewrite PrimExprs inside call annotations. Pipeline-managed TMA copies keep their ring parity in the explicit waits the pass inserts, and the explicit-async T.tcgen05_gemm never auto-waits, so neither is annotated. Drop the matching ring branch from LowerTileOp's fallback: it keyed on the "num_stages" annotation, which InjectSoftwarePipeline strips, so it was dead on pipelined loops and wrong on unpipelined ones -- a compiler-generated per-site barrier flips once per loop iteration regardless of pipeline depth. The fallback is now the zero-based epoch of the nearest enclosing serial loop. Add transform-level tests asserting the annotated phase and the emitted wait parities on sm_90a, so the fix is covered without sm_100 hardware; the loops start at a nonzero min to pin the epoch normalization. Add a tcgen05 correctness test running the failing kernel at num_stages 2 and 3. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Lei Wang
|
cb7538b379 |
[Lang][Reducer] Vectorize contiguous reducer updates (#3079)
* Enable vectorized reducer updates * Mark CUDA-only reducer tests |
||
|
Chennes
|
fe2951323f |
[CUDA] Bridge half-style math intrinsics for 16-bit CUTLASS types (#3077)
TVM lowers 16-bit math operations to CUDA's half-style names, but CUDA provides overloads only for native __half and __nv_bfloat16 values. Calls with CUTLASS half_t and bfloat16_t wrappers therefore failed to compile. Add explicit forwarding overloads for hexp, hexp2, hexp10, hlog, hlog2, hlog10, hsin, hcos, hfloor, hceil, hrint, and htrunc. Each overload converts the wrapper to its native CUDA type, calls the existing intrinsic, and wraps the result back in the original type. Fold the previous standalone bfloat16 hexp bridge into the complete overload set. Codegen and the fast-math path remain unchanged. Add scalar runtime coverage for both wrapper types using materialized temporaries, separate valid input domains, generated-call assertions, and torch references. |
||
|
Sam Sui
|
e8b9da92d2 |
[BugFix][CUDA] Support int4x2 and uint4x2 codegen (#3036)
* fix(cuda): support int4x2 and uint4x2 codegen (#2994, #2995) * fix(cuda): handle packed int4 lane access * fix(cuda): handle packed 4-bit x2 construction * fix(cuda): load packed x2 gathers by logical index * fix(cuda): validate packed int4 vector alignment * fix(cuda): preserve packed vector load fallback |
||
|
KellyFrog
|
baf8a9c467 |
[BugFix] Scalarize Select in automatic vectorization (#3060)
[BugFix] Align Select auto-vectorization with IfThenElse |
||
|
Chennes
|
e54cbabb12 |
[CUDA] Fix warp shuffle for half, bfloat16, and FP8 (#3056)
T.shfl_sync, T.shfl_xor, T.shfl_down, and T.shfl_up lower to CUDA's raw __shfl_*_sync builtins. CUDA has no exact overloads for the narrow floating-point wrappers TileLang emits, so calls either resolved to the float overload and returned a float that could not copy-initialize the wrapper, or were ambiguous. Either way the shuffle failed to compile. Add exact non-template overloads for half_t, bfloat16_t, float_e4m3_t, and float_e5m2_t. The overloads carry the raw bits through CUDA's uint32_t shuffle and preserve mask, lane, delta, and width semantics. A local macro generates the four builtin overloads per type and is undefined immediately after use, leaving codegen and native CUDA overload resolution unchanged. Add exact lane-permutation coverage for all four shuffle operations and all four wrapper types in both materialized-temporary and inline forms. Fixes #2569 |
||
|
Yongqi Zhuo
|
2e43b1395e |
[CUDA] Remove T.{reads,writes} for T.tma_{gather4,scatter4} (#3053)
[CUDA] Remove `T.{reads,writes}` annotations for `T.tma_{gather4,scatter4}`
|
||
|
Srijan Keshri
|
9f336776ee |
[Bugfix] Carry memory_order through vectorized atomic_add (#2924)
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>
|
||
|
Lei Wang
|
23c29fa361 |
[BugFix][Layout] Avoid thread-indexed wide reducer finalize readback (#3049)
Fix wide reducer finalize publication layout |
||
|
Chennes
|
6b19ebd294 |
[CUDA] Fix FP8 min/max codegen (#3047)
[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 |
||
|
Lei Wang
|
9c94f7765e | [Lang][Reducer] Allow conditional reducer finalization (#3044) | ||
|
Lei Wang
|
6328dacd02 |
[Lang][Reducer] Support loop-scoped epochs and legacy default allocations (#3043)
[Reducer] Support epochs in serial control flow |
||
|
Chennes
|
39b3b3e5db |
[CUDA] Add __ldg overloads for 16-bit CUTLASS types (#3028)
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. |
||
|
Lei Wang
|
024fc72896 |
[Lang] Reducer v2: first-class deferred reduction epochs with planned physical lowering (#2940)
* [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 (
|
||
|
Yongqi Zhuo
|
75cf7632b8 | [Testing] Pin the folded-base descriptor form for static ts slices (#2951) | ||
|
Yongqi ZhuoandClaude Fable 5
|
c2d847f497 |
[BugFix][CUDA] Advance tcgen05 ld/st segment pointers in b32 columns (#2952)
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>
|
||
|
Yongqi Zhuo
|
4d1503179f |
[Layout][CUDA] Support reinterpreting (dtype-changing) T.view aliases (#2953)
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. |
||
|
Diao Han
|
6af825e072 |
[Enhancement] Expose cluster_mask on T.tma_copy (#2932)
* [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.
|
||
|
Chennes
|
14901c7722 |
[CUDA] Add pre-SM80 fallback for bf16 atomic add (#2938)
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. |
||
|
Yongqi Zhuo
|
fb1b11f1b8 |
[Feature] Warp specialization schedules and materialization (#2892)
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. |
||
|
Tong WUandYongqi Zhuo
|
13409c1c67 |
[Layout] Support half-subpartition (M=64) TMEM tiles in tcgen05.ld/st (#2880)
* [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> |
||
|
Chaofan Lin
|
cd57a840a1 |
[Refactor][BugFix] Refactor the loop vectorization plan with ConstraintKind (#2935)
* [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 |
||
|
Lei Wang
|
fa2476fabf |
[JIT][FFI] Allocate kernel outputs through the packed API (#2937)
* Support TVM-FFI callee-allocated outputs * Remove redundant TVM-FFI cache fingerprints * Remove non-portable manual out_idx test |
||
|
Chaofan Linandclrs97
|
4a860d9071 |
[Cherry-Pick][BugFix] Fix get_parent_locals frame self-reference leak (#2934)
Fix get_parent_locals frame self-reference leak Co-authored-by: clrs97 <33470839+clrs97@users.noreply.github.com> |
||
|
Chaofan Lin
|
27224884f9 |
[Testing] Run portable regression tests on auto targets (#2914)
[Testing] Remove stale CUDA gates from portable tests |
||
|
KellyFrogandLeiWang1999
|
4e1820c938 |
[BugFix] Deduplicate replicated reducer updates (#2881)
* [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> |
||
|
Chaofan Lin
|
afa3a448fe |
[TIR][Python] Split backend-specific op proxies (#2870)
* [TIR][Python] Split backend-specific op proxies * [TIR][Python] Preserve common op documentation * [TIR][Python] Preserve moved op metadata |
||
|
andyluo7
|
dd9c20dbf7 |
[Testing] Gate CUDA-only tests so non-CUDA backends can run the suite (#2863)
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.
|
||
|
Yongqi Zhuo
|
66e8e47a52 |
[Refactor] Recycle T.unroll(explicit=True) for early explicit unrolling (#2859)
[Refactor] Recycle syntax `T.unroll(explicit=True)` for early explicit unrolling |
||
|
morluto
|
36cae46339 |
[Typo] Correct source spelling errors (#2858)
* fix: correct source spelling errors * style: apply clang-format to reduction helper |
||
|
Yongqi Zhuo
|
caf694446d | [CUDA] Adopt multi-staged buffers in examples (#2836) | ||
|
Lei Wang
|
9e13487941 |
[CUDA][TMA] Separate atomic-add dtype support from layout encoding (#2846)
[CUDA] Refactor TMA atomic add layout validation |
||
|
Tong WU
|
a426ff3657 |
[CUDA] Pack logical TMEM buffers into shared tcgen05.alloc arenas (#2831)
* [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. |
||
|
morluto
|
56a0f7297e |
[BugFix] Reject unsupported TMA atomic add dtypes (#2830)
* [BugFix] Reject unsupported TMA atomic add dtypes * [BugFix] Tighten TMA atomic dtype validation |
||
|
|
8f34abf498 |
[CUDA] Add SM120 NVF4 block-scale MMA support (#2364)
* 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 (
|
||
|
Chaofan Lin
|
50481cce09 | [Refactor] Remove intrinsic compatibility facade (#2812) | ||
|
cjandcaojian5
|
21e8c064ac |
[BugFix] Reject unsupported fast-math input dtypes (#2804)
* Reject unsupported fast-math dtypes * test: cover all unsupported fastmath integer dtypes --------- Co-authored-by: caojian5 <caojian5@huawei.com> |
||
|
KellyFrogandchenkf
|
0bc1913daf |
[BugFix] Resolve partial scalar reduce barrier participation (#2777)
* [CUDA][Reduce] Resolve partial AllReduce participants
* refactor: inline ResolveAllReduceThreadRange into reduce.h and remove allreduce_barriers doc
- Inline ResolveAllReduceThreadRange body from reduce.cc into reduce.h,
add #include <tvm/ir/with_context.h> for With<ConstraintContext> usage.
- Delete src/backend/common/op/reduce.cc (function now in header).
- Remove reduce.cc from CMakeLists.txt TILE_LANG_SRCS.
- Delete docs/compiler_internals/allreduce_barriers.md.
- Remove toctree entry from docs/index.md.
* [CUDA][Reduce] Resolve batched AllReduce participants
* Revert "[CUDA][Reduce] Resolve batched AllReduce participants"
This reverts commit
|
||
|
Chaofan Lin
|
6c3dd971cf |
[BugFix] Skip descriptor TMA for device-bound copy bases (#2803)
* [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 |
||
|
jpwang
|
e41fadbe3d |
[BugFix] Fix warp_reduce truncating int64/uint64 to 32 bits on sm_80+ (#2782)
* [BugFix] Fix warp_reduce truncating int64/uint64 to 32 bits on sm_80+ * Add tests to resolve review comments |
||
|
cjandcaojian5
|
51f88a88ed |
[BugFix] Reject floating-point predicates in vote intrinsics (#2800)
* Reject floating-point predicates in vote intrinsics * Accept boolean vote predicates --------- Co-authored-by: caojian5 <caojian5@huawei.com> |
||
|
cjandcaojian5
|
b1b605dab4 |
[BugFix] Preserve alloc_var initializer dtype (#2801)
* Preserve alloc_var initializer dtype * Test alloc_var float32 initializer range --------- Co-authored-by: caojian5 <caojian5@huawei.com> |
||
|
Yongqi Zhuo
|
1cb4d4f3ac |
[CUDA] Support arbitrary TMEM layouts (#2785)
* [CUDA] Support arbitrary TMEM layouts * Fix bugs |
||
|
Liu YunuoandSiriusNEO
|
7ec5adbeb1 |
[BugFix] Pass buffer row stride to 2D scan kernel to fix silent miscomputation (#2620)
* fix: pass buffer row stride to 2D scan kernel to fix silent miscomputation * update test * lint * fix: pass separate src and dst row strides to 2D scan kernel * test: cover scans with distinct source and destination strides --------- Co-authored-by: SiriusNEO <chaofan@deepseek.com> |
||
|
Lei Wang
|
6171343c06 |
[TIR][Language] Add typed vector lane extraction API (#2789)
* [TIR][Python] Add vector lane extraction API * [TIR][Python] Move Shuffle export to TIR facade * [TIR][Python] Forward Shuffle through op facade |