mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-07 09:32:37 +08:00
0ce2a7bacf01ab4622d4745ca48b004f1caa14aa
22
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 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 |
||
|
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 (
|
||
|
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> |
||
|
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.
|
||
|
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
|
||
|
Yuanyuan Zhaoandzyy3077
|
a69708c1d5 |
[BugFix] Reject non-power-of-two AllReduce widths (#2611)
Co-authored-by: zyy3077 <zyy3077@users.noreply.github.com> |
||
|
Xiangwen Wang
|
4433981c99 |
[Enhancement] Add local buffer reduction lowering (#2693)
* Add local buffer reduction lowering * Minor fix * Simplify |
||
|
Zihao WangandLeiWang1999
|
5dfa0f5d6e |
[BugFix] Fix grouped reduce_sum over-counts on straddle layout (#2424)
* Fix reduce ownership for straddled layouts * Guard reduce thread ownership projection --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Jayce Suanddingsg
|
779022eaf0 |
[BugFix][CuTeDSL] Complete TileKernels benchmark support (#2319)
* [BugFix] CuTeDSL backend support for TileKernels benchmarks TileKernels benchmark coverage with TILELANG_TARGET=cutedsl exercises backend paths that were either missing or not legal for CuTeDSL-generated Python/MLIR code. The failures were not limited to one kernel: packed FP4/FP8 quantization flows need sub-byte vector materialization, MoE mapping uses warp vote/shuffle and grid synchronization, and benchmark launch wrappers need dtype-aware host plumbing for generated CuTeDSL modules. Before this change the H100 benchmark run exposed several classes of failures: missing shuffle/vote intrinsics such as shfl_xor/match_any/popc in CuTeDSL lowering, invalid scalar/vector casts for FP4 E2M1 and packed UE8M0 paths, dynamic if expressions whose branches narrowed or changed integer types, unsupported sub-byte load/store/copy patterns, and a fragile one-counter software grid barrier that could reset before all waiting blocks observed release. Lower the missing TileLang operations to CuTeDSL runtime helpers, add Float4E2M1FN software pack/unpack and zero-fill support, make sub-byte vector loads/stores use packed Uint8 storage where required, cast conditional branches to the expected result dtype, and teach the CuTeDSL wrapper about float4/float8 dtype plumbing and C++ min/max expression rendering. Also replace the sync_grid single-counter barrier with a counter plus generation sense barrier so a waiting block cannot miss the release when the arrival counter is reset. This keeps the fixes in the CuTeDSL backend rather than adding TileKernels-specific workarounds. For example, TileKernels expand_to_fused_with_sf relies on TileLang safe-memory legalization for aligned SF loops; the CuTeDSL generated code now preserves the hidden_sf guard instead of requiring the benchmark source to special-case packed UE8M0 layouts. Validation on H100 (sm_90, 114 SMs) with CUDA visible device 0, Python 3.12, TILELANG_TARGET=cutedsl, TK_FULL_TEST=1, TILELANG_DISABLE_CACHE=1, and PYTHONPATH=/home/jayce.su/new/tilelang: - cmake --build build -j 16 - ruff check tilelang/contrib/cutedsl tilelang/jit/adapter/cutedsl tilelang/jit/adapter/utils.py - git diff --check - Original failing node: tests/moe/test_expand_to_fused.py::test_expand_to_fused_with_sf_benchmark[num_send_tokens=4001-num_topk=2-num_experts=4-num_ep_ranks=64-hidden=576-num_per_channels=32-use_tma_aligned_col_major_sf=True-round_sf=True-use_packed_ue8m0=True], 1 passed with -n 2 - Full-run local pressure slice around the failing MoE benchmark region, 80 passed with -n 2 in 379.53s Author-email: jayce.su@enflame-tech.com Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Complete CuTeDSL support for TileKernels benchmarks TileKernels benchmark coverage with TILELANG_TARGET=cutedsl exposes several general CuTeDSL backend gaps. These are not TileKernels-specific workarounds, H100-only specializations, or changes required only because of a particular CUDA/CUTLASS DSL version pairing. The H100 CUDA 13.2 + nvidia-cutlass-dsl-libs-base 4.5.0 environment is the reproducer that exercises the missing backend coverage; the underlying issues are generic to the emitted CuTeDSL code shapes. Root causes: - Predicated T.copy and safe memory legalization lower BufferLoad-or-zero values into conditional expressions. Emitting these values through tl.where evaluates both branches eagerly, which can issue unsafe packed loads and can force CuTeDSL/MLIR through unsupported vector select and unrealized_conversion_cast paths for FP8/FP4 storage. - Safe memory legalization can bind the zero branch through a local zero-like variable, let binding, cast, or broadcast-zero tensor. The CuTeDSL codegen must recognize those zero-like values when deciding whether a conditional load can be lowered lazily and when materializing Float4E2M1FN broadcasts. - CuTeDSL generated code has to distinguish rmem Tensor objects from TensorSSA values. Appending .load() blindly is invalid once a branch or cast already produced TensorSSA, and this showed up in the TileLang CI cast example as AttributeError: 'TensorSSA' object has no attribute 'load'. - Scalar FP8 widening and narrowing require CuTeDSL tensor casts on aligned vector values. Plain scalar construction falls back to unsupported scalar FP8 conversion paths, so the backend has to bridge scalar casts through a packed vector form. - Sub-byte vectors such as Float4E2M1FN cannot be copied lane by lane because CuTeDSL does not support sub-byte scalar dereference. The backend must preserve TensorSSA/packed-vector store form for these values, and scalarized local/shared sub-byte vector buffers must use the sub-byte packing factor as the make_tensor_at_offset div_by assumption. - Generated CuTeDSL kernels can call CUDA-template math helpers such as __habs and __float2half_rz, and math wrappers such as log1p/isfinite. Those helper names must be provided by tilelang.contrib.cutedsl.math when the Python CuTeDSL backend is selected. - sync_grid publishes block arrival through thread 0, but the memory ordering has to cover all threads in the block. A CTA sync alone does not provide device-scope ordering for global writes performed by non-zero threads. Fixes: - Recognize zero-like vars in conditional store lowering, including let bindings and bind values that are casted or broadcast-zero tensors. This keeps predicated loads lazy and avoids eager tl.where evaluation for packed and vectorized T.copy paths. - Add CuTeDSL helper bridges for Tensor/TensorSSA normalization. tl.as_tensor_ssa avoids calling .load() on TensorSSA values, and tl.as_rmem_tensor spills TensorSSA values into rmem tensors before byte recasts for sub-byte stores. - Add scalar FP8 bridge lowering that packs Float16/Float32/BFloat16 or FP8 scalar values into a four-lane rmem tensor, uses tl.cast_tensor, and reads lane zero. This matches CuTeDSL/NVGPU alignment requirements without forcing front-end kernels to perturb tile shapes just to avoid scalar casts. - Keep sub-byte vector stores in TensorSSA form instead of copying lane by lane, keep non-load vector store RHS materialized as rmem tensors so CuTeDSL .store() receives a TensorSSA value, and use 8 / element_bits as the div_by packing factor for scalarized local/shared sub-byte vector buffer views. - Gate the Float16 nvgpu.cvt_fpext fast path to floating-point source TensorSSA types only, so integer TensorSSA casts do not enter an FP-only conversion op. - Strengthen sync_grid by fencing all block threads at device scope before thread 0 publishes arrival, then waiting at block scope and fencing again before user code resumes. - Expose the missing CuTeDSL math helpers. __habs clears the sign bit for Float16/BFloat16 scalars and returns the original dtype. __float2half_rz uses inline PTX cvt.rz.f16.f32 so generated E5M6 cast-back code can preserve CUDA round-toward-zero semantics. log1p and isfinite are wired through the CuTeDSL math module and fastmath canonicalization. - Keep the defensive dynamic-symbolic None fallback visible by logging a warning before substituting 0. Version/specialization assessment: - The patch is backend-generic: it fixes how TileLang emits CuTeDSL for predicated loads, FP8 scalar casts, Tensor/TensorSSA boundaries, sub-byte vector stores, grid synchronization, and generated math helpers. It does not specialize TileKernels cases, tensor shapes, H100 SM counts, or benchmark parameters. - A future CuTeDSL version might accept some currently rejected IR forms, but relying on eager branch evaluation, missing helper symbols, incorrect TensorSSA .load() calls, or thread-0-only grid-sync ordering would still be fragile. The corrected lowering is semantically closer to TIR: predicated loads are lazy, sub-byte storage stays packed, TensorSSA values are not treated as tensors, and scalar FP8 conversion follows the vector cast path required by CuTeDSL/NVGPU. - CUDA 13.2 and CUTLASS DSL libs 4.5.0 are recorded as the verified reproducer, not as the reason for the patch. The same TileLang backend gaps can be triggered by any workload that emits these CuTeDSL patterns. Verified on H100 PCIe (sm_90, CUDA_VISIBLE_DEVICES=0 maps to NVIDIA H100 PCIe), CUDA 13.2, Python 3.12.3, PyTorch 2.12.0+cu132, nvidia-cutlass-dsl-libs-base 4.5.0, and local TileLang built from this branch: - cmake --build build -j$(nproc) - python3.12 -m py_compile tilelang/contrib/cutedsl/math.py tilelang/contrib/cutedsl/utils.py tilelang/contrib/cutedsl/grid_sync.py tilelang/jit/adapter/cutedsl/adapter.py - pre-commit run --files $(git diff --name-only origin/main...HEAD) - examples/cast/test_example_cast.py::test_example_group_per_split_token_cast_to_fp8: 1 passed - representative per_token_cast, cast_back, cast_back_e5m6, per_block_cast_lossless, per_block_cast, and per_channel_cast benchmark failures: 6 passed - TileKernels full benchmark with TILELANG_TARGET=cutedsl, TK_BENCHMARK_BACKEND=event, and pytest -n 2: 10165 passed, 160 warnings in 2:44:58 Jayce: JayceSu98 <jayce.su@enflame-tech.com> Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Avoid zero-sized fake tensors in CuTeDSL cubin generation CuTeDSL cubin generation builds fake compact tensors to compile wrapper code before the real runtime tensors are launched. TileKernels exercises benchmark and correctness cases with empty dynamic dimensions, such as zero-token inputs. Passing those zero extents directly into make_fake_compact_tensor can make the CuTeDSL fake layout invalid even though the shape is only a compilation placeholder and the real runtime tensor metadata is passed separately. Clamp only the fake tensor shape to positive extents before constructing the CuTeDSL placeholder. This keeps the cubin generation path valid for empty-input cases without changing runtime tensor shapes, strides, launch parameters, or kernel behavior. This is not a TileKernels-specific workaround or a CUDA toolkit version workaround. It fixes the CuTeDSL adapter boundary so compile-time fake tensors satisfy CuTeDSL layout construction requirements while preserving the actual runtime contract. Verified with: - python -m py_compile tilelang/jit/adapter/cutedsl/wrapper.py - git diff --check -- tilelang/jit/adapter/cutedsl/wrapper.py Jayce: JayceSu98 <jayce.su@enflame-tech.com> Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Stabilize batched BF16 reduce test input CUDA CI exposed a flaky failure in python/language/test_tilelang_language_reduce.py::test_reduce[sum-bfloat16-64x128-f2f-t256-b4]. The failing run had only one mismatched output element, with an absolute difference slightly above the 0.1 BF16 tolerance and a large relative error because the reference value was small. The runtime code on this branch only changes CuTeDSL backend paths, and the same test passed locally on the ordinary CUDA backend when run in isolation. The failure is caused by the test input depending on the worker-local torch RNG state: the module seeds once at import time, but xdist load scheduling lets each worker consume a different number of random tensors before this parameterized case runs. That makes the BF16 sum case occasionally sample an input whose rounded reduction error crosses the tolerance boundary. Make the parameterized reduce input schedule-independent by deriving a deterministic seed from the case parameters and using a local CPU torch.Generator for both integer and floating inputs. This preserves randomized coverage per case while decoupling correctness from xdist worker assignment, prior cases in the same worker, and full-suite ordering. Verified on H100 PCIe with CUDA visible device 1, Python 3.12.3, PyTorch 2.12.0+cu132, and the local TileLang dev build: - ruff check testing/python/language/test_tilelang_language_reduce.py - python -m pytest -q -s --tb=short testing/python/language/test_tilelang_language_reduce.py::test_reduce[sum-bfloat16-64x128-f2f-t256-b4] - python -m pytest --verbose --color=yes --durations=0 --showlocals --cache-clear --maxfail=3 --numprocesses=8 testing/python/language/test_tilelang_language_reduce.py - 10 consecutive reruns of the failed BF16 reduce case Jayce: JayceSu98 <jayce.su@enflame-tech.com> Co-authored-by: dingsg <shengge.ding@enflame-tech.com> --------- Co-authored-by: dingsg <shengge.ding@enflame-tech.com> |
||
|
Lei Wang
|
4e9b0543d5 |
[Reduce][Codegen] Guard packed local reduce ramp loads (#2298)
Fix packed reduce ramp contiguity check |
||
|
Lei Wang
|
124ba6affe |
[CUDA][Reduce] Fix packed mixed-dtype reduce casts (#2276)
Fix packed reduce casts for mixed dtypes |
||
|
KurisandLeiWang1999
|
719d23dce5 |
feat: auto-vectorize bf16/fp16 reduce with packed add2 intrinsics (#2112)
* fix: add missing TvmLogDebugSettings::ParseSpec and VerboseEnabledImpl for TVM_LOG_CUSTOMIZE builds When TVM_LOG_CUSTOMIZE=1, TVM's logging.cc skips compiling TvmLogDebugSettings::ParseSpec and VerboseEnabledImpl (guarded by #if TVM_LOG_CUSTOMIZE == 0). However libtilelang.so calls these functions via the inline FromFlag(), causing a runtime symbol lookup error. Add the missing implementations to tilelang's own logging.cc. * feat: auto-vectorize bf16/fp16 reduce with packed add2 intrinsics * style: clang-format * fix: fp16 batch pack — fix Shuffle codegen taking address of temporary - Add tl::pack_half2 helper in common.h that returns uint1, avoiding taking address of a temporary __pack_half2 return value - Update ShuffleNode codegen for fp16x2 to use pack_half2 instead of &(__pack_half2(...)) - Re-enable fp16 batch pack (remove bf16-only restriction in can_batch_pack) - Add SumOp_fp16x2/MaxOp_fp16x2/MinOp_fp16x2 packed reducer structs * refactor: simplify reduce packing API - Replace ReducePackConfig with simple GetPreferedVectorizedSize(DataType, Target) -> int - Merge getBatchReducerSuffix into MakeCodegenReducer(vsize) - MakeReduce(vsize, acc, b) -> optional<PrimExpr> - MakeInitValue(vsize) handles Broadcast internally - Guard nan_propagate in MakeCodegenReducer vsize>1 path (no packed nan ops exist) * fix: guard nan_propagate in MakeReduce packed path, mirror scalar structure * feat: support nan_propagate in packed reduce (max2_nan/min2_nan) * fix: add bitwise reduce ops to MakeReduce vsize=1 path * fix lint error --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
|
6548c05f0c |
[Feature] Batched AllReduce for better T.reduce performance (#1976)
* [Feature] Batched AllReduce for better T.reduce performance (#1928)
When each thread holds multiple values to reduce, the previous code
emitted one AllReduce call per element inside a loop, causing redundant
synchronization barriers (e.g. 8 values × 6 syncs = 48 syncs).
This change adds batch support to AllReduce so all per-thread values are
reduced in a single butterfly pass sharing barriers (6 syncs total).
- Add `batch_size` and `workspace_stride` template parameters to
`tl::AllReduce` (defaults preserve existing scalar behaviour).
- In ReduceOpNode::Lower, detect when batching is beneficial
(batch_size > 1 and cross-warp reduction) and restructure into three
phases: loop(init + local reduce), batched AllReduce, loop(copy-back).
- Apply same optimisation to FinalizeReducerOpNode::Lower.
- Mirror template changes in HIP AllReduce.
Closes #1928
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Test] Add codegen tests for batched AllReduce
Verify that the batched AllReduce path (batch_size > 1, workspace_stride)
is emitted in the generated CUDA source for reduce_max, reduce_sum, and
reduce_min with cross-warp reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Fix batched AllReduce codegen for HIP/ROCm targets
HIP AllReduce template has no Barrier type parameter, so the batched
path must omit tl::SyncThreadsBarrier for ROCm targets. Adds
TargetIsRocm() checks in both reduce.cc and finalize_reducer.cc.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Fix] Use target-aware warp size for cross-warp reduction threshold
Address CodeRabbit review: ROCm wavefronts are 64-wide, so the cross-warp
threshold should be 64 (not 32) on HIP targets. This avoids unnecessary
workspace allocation and batched codegen for single-wavefront reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Style] Fix ruff format for test assertion
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Support batched AllReduce in CuteDSL Python codegen
The CuteDSL codegen translates C++ AllReduce template calls to Python.
With the new batch_size and workspace_stride template params, the
generated Python calls AllReduce() with 7 args instead of 5. Add
batch_size and workspace_stride params to the Python AllReduce function
and implement batched reduction logic for both standard and Hopper paths.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Feature] Add explicit batch parameter to T.reduce for user-controlled batched AllReduce
Expose a `batch` parameter on T.reduce (and all T.reduce_* wrappers,
default 1 = scalar / current behaviour). When batch > 1 the compiler:
- validates batch divides the per-thread output element count N;
- emits ceil(N/batch) calls to the new AllReduce::run_batch interface,
each sharing a single pair of barriers across all batch elements;
- allocates workspace of size reducing_threads * batch (only for cross-
warp reduce, i.e. reducing_threads > 32).
Key differences from the existing auto-detection approach:
* The user controls when batching is enabled — no silent smem increase.
* workspace_stride = reducing_threads (not block_size), so workspace is
minimal and does not blow past the GPU shared-memory limit.
* run_batch (not run) avoids C++ overload-resolution ambiguity when a
pointer is passed as the first argument.
Files changed:
- src/op/reduce.h: add `int batch{1}` field to ReduceOpNode
- src/op/reduce.cc: read "batch" annotation; validate; emit N/batch
batched AllReduce calls with address_of per-chunk pointers
- src/tl_templates/cuda/reduce.h: rename batch run() -> run_batch(),
fix recursive call
- src/tl_templates/hip/reduce.h: same
- tilelang/language/reduce_op.py: add batch param to reduce() and all
reduce_max/min/sum/absmax/abssum/bitand/bitor/bitxor wrappers
- testing/python/language/test_tilelang_language_reduce.py: update
test_batched_allreduce_codegen to use explicit batch=2 and match
run_batch in generated source
* [Fix] Use run_batch in finalize_reducer batched AllReduce codegen
finalize_reducer.cc also generates batched AllReduce calls (used by
T.alloc_reducer / T.finalize_reducer), but was still emitting ::run
instead of ::run_batch, causing a compiler error because the batched
pointer-based interface is now named run_batch to avoid overload
ambiguity.
* rocm fix
* [Feature] Add explicit batch parameter to T.finalize_reducer
Mirrors the batch= API already available on T.reduce:
- batch=1 (default) keeps the existing scalar AllReduce::run path
- batch>1 emits AllReduce::run_batch, sharing barriers across batch
output elements and reducing barrier count by batch×
C++ changes:
- FinalizeReducerOpNode gains an int batch{1} field (registered via
reflection) and reads it from the 'batch' annotation in the constructor
- Lower() validates batch against the layout output-element count and
uses effective_batch to drive use_batch / codegen
Python changes:
- finalize_reducer(reducer, batch=1) – new optional argument
- batch < 1 raises ValueError (consistent with T.reduce)
Tests: add test_batched_finalize_reducer_codegen covering sum/max/min
with explicit batch values, verifying run_batch is emitted with the
correct template argument.
* [Test] Add comprehensive tests for T.finalize_reducer batch parameter
Covers all cases introduced by the explicit batch= API:
Codegen tests:
- batch>1 emits AllReduce::run_batch with correct template argument
- batch=1 (default) does NOT emit run_batch (scalar path)
Correctness tests:
- batch=1 scalar path: sum/max/min across float32/float16 (PASS)
- batch==block_M full-batch: correct by coincidence (PASS)
- batch>1 partial-batch: xfail(strict=True) – known bug where
run_batch receives buffer->data directly but fragment layout is
non-contiguous per thread; root cause documented in
work/bug-finalize-reducer-batch.md
Error-input tests:
- batch=0 / batch<0 raise ValueError at Python layer
- batch > total output elements raises at compile/lower time
- batch that does not evenly divide output elements raises
* [Test] Rewrite finalize_reducer tests to follow clean outline
Replace the previous over-engineered test suite (4 case lists, 2 split
correctness functions, per-case xfail) with a minimal set organised by
the three categories that matter:
1. Codegen – single parametrised test covers both batch=1 (no run_batch)
and batch>1 (run_batch with correct template arg), reusing one case
list for both codegen and correctness.
2. Correctness – only batch=1 (scalar path), which is currently correct.
batch>1 correctness tests are omitted until the underlying fragment
layout bug is fixed (tracked in work/bug-finalize-reducer-batch.md).
3. Error inputs – batch<1 collapsed into one parametrised test; compile-
time errors (batch>output, batch not divisible) each get one test
using a shared _make_error_kernel helper.
* [Test] Rewrite reduce tests following clean outline
Replace fragmented per-function tests with four parametrized suites:
- test_reduce: op × dtype × src_scope × dst_scope × threads, covers
fragment→fragment, shared→shared, fragment→shared for sum/max/min/
abssum/absmax with float32/float16/int32/int64
- test_reduce_clear: op × src_scope × dst_scope with clear=False,
verifies pre-filled dst is accumulated correctly
- test_reduce_batch_codegen: verify batch>1 emits AllReduce::run_batch
- test_finalize_reducer_codegen/correctness/invalid_batch: batch
parameter feature tests, sharing one FINALIZE_REDUCER_CASES list
Use tilelang.jit lazy style and torch.testing.assert_close directly,
no profiler wrapper.
* [Test] Merge batch codegen+correctness into test_reduce
Add batch column to REDUCE_CASES; batch>1 cases verify both that
run_batch appears in generated source and that numerical output matches
reference, within the same test_reduce parametrization.
Remove the now-redundant test_reduce_batch_codegen.
* [Test] Add more batch>1 reduce cases covering all ops and dtypes
Add batch=2/4/8 cases for sum/max/min/abssum/absmax across float32/
float16/bfloat16, with both fragment and shared src/dst scopes.
Relax tolerance to 1e-1 for float16/bfloat16 to account for larger
rounding error in multi-element reductions.
* [Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug
* Revert "[Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug"
This reverts commit
|
||
|
Lei WangandZhiwen Mo
|
f8dc61c02c |
test: reduce CI runtime for slow Python suites (#1932)
* test: trim CI runtime for slow Python suites * remove useless test * Update TVM submodule to latest commit and remove manual device assert test from CI. Replace with a call to tilelang.testing.main(). * lint fix * fix * revert test for eager jit --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
|
48a8f4a6ce |
[Feature] Hierarchical reduction and warp reduction intrinsics support (#1762)
* [Feature] Add new hierarchical reduce with less workspace size and new ptx * Lint * Add intrinsic functions and more types * Fix sm100a redux instruction * Fix data type * [Test] Add tests for reduce operations * [Refactor] Unify AllReduce barrier mechanism via policy template Replace the duplicated run()/run_hopper() methods in AllReduce with a single run() that accepts a Barrier policy template parameter (SyncThreadsBarrier or NamedBarrier<N>). Extract the shared inter-warp reduction logic into a warp_inter_reduce helper, and split the dispatch into private hierarchical_reduce/butterfly_reduce methods. Update codegen in reduce.cc and finalize_reducer.cc to emit NamedBarrier<all_threads> for SM >= 90 targets instead of the old all_threads + run_hopper pattern. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add handling for non-float types in warp_reduce function This update introduces a new case in the warp_reduce function to cast non-float types to float using static_cast. This enhancement improves type flexibility and ensures compatibility with a broader range of input types during reduction operations. * [Refactor] Remove hierarchical reduce because of performance regression * [BugFix] Buffer for reduction --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
ShaobinChenandSiriusNEO
|
215f436b05 |
[BugFix] Fix reduce_sum with clear=False not accumulating correctly (#1778)
* Fix type annotations for T.reshape and T.view * Fix issue #1666: reduce_sum with clear=False not accumulating correctly * address comments and add testcases * add more tests --------- Co-authored-by: SiriusNEO <chaofan@deepseek.com> |
||
|
Lei Wang
|
c750fb8a98 |
[Enhancement] Update examples and tests for improved type handling functionality (#1448)
* [Enhancement] Update examples and tests for improved type handling and functionality
- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.
* [Refactor] Update accumulation data type to float32 across examples
- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.
* [Refactor] Standardize data type usage across benchmark scripts
- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.
* [Refactor] Standardize data type usage in templates and scripts
- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.
* [Refactor] Standardize data type usage in examples and benchmarks
- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.
* [Refactor] Import dtypes from language.v2 module
- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.
* fix
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.
* [Refactor] Update data type handling for consistency and clarity
- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.
* [Enhancement] Improve data type handling and error messaging
- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.
* [Fix] Correct boolean flag in GEMM SP test case
- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.
* [Refactor] Standardize data type usage in various modules
- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.
* [Refactor] Update argument parsing for data types in benchmarks
- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.
* [Refactor] Update data type handling in benchmark and example scripts
- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.
* [Refactor] Fix data type conversion in multiple scripts
- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.
* [Refactor] Update float8 data type usage across multiple scripts
- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.
* [Refactor] Enhance float8 data type handling in CUDA code generation
- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Streamline float8 data type handling in CUDA and related modules
- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Remove unnecessary cache disabling in float8 example script
- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.
* [Refactor] Update data type usage in debug print tests
- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.
* lint fix
* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples
* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
|
||
|
Lei Wang
|
29051439db | [Lint] Phaseout Yapf format and embrace ruff format (#1417) | ||
|
Lei Wang
|
67cc861136 |
[Enhancement] Add thread count validation for ReduceOp fragment layout inference (#1225)
* [Enhancement] Add thread count validation for ReduceOp fragment layout inference * Introduced a check to ensure that the thread count is divisible by the replicate extent during layout inference in ReduceOpNode. This validation prevents layout inference failures and provides detailed error messages to guide users in resolving issues related to thread block sizes and fragment layouts. * Updated tests to remove unsupported configurations that could lead to layout inference errors, ensuring more robust testing scenarios. * lint fix |
||
|
Lei Wang
|
bc37ea69d5 |
[Language] Efficient T.reduce_ with shared memory input/output (#1080)
* Support reduce ss * lint fix * test fix * lint fix |