Commit Graph
552 Commits
Author SHA1 Message Date
Lei Wang 3af87b7fb6 [TIR][Transform] Fix flat Bind modeling in parallel race checks (#2665)
[TIR] Fix flat Bind modeling in parallel race checks
2026-07-15 10:55:23 +08:00
Lei Wang 70548a1707 [BugFix][CUDA][HIP] Correct packed shared memory allocation sizes (#2660)
Fix packed shared memory allocation sizes
2026-07-14 16:11:22 +08:00
penguin_wwy effebb6bf0 [BugFix] Fix segfault when tunable params default to None (#2657) 2026-07-14 14:11:20 +08:00
Lei Wang 1ac5a01a4e [TIR][Runtime] Enforce host-evaluable assumptions at runtime (#2655)
* Add runtime checks for host-evaluable assumptions

* Preserve nullable buffer semantics for assume checks

* Rename assume runtime-check marker
2026-07-14 12:24:26 +08:00
Da1L8-Xanddaniu.ld 6ed02aee47 [Autotune] Support per-config pass_configs in autotuning (#2496)
* feat(autotuner): support per-config pass_configs with code deduplication

- Extract _merge_pass_configs_into_compile_args helper to deduplicate
  pass_configs merge logic across grouped_compile.py and tuner.py
- Replace json.dumps uniformity check with direct equality comparison
- Extract _make_jit_compile_func factory to unify lazy/eager jit_compile
- Fix decorator-mode cache isolation: use compile() directly when
  per_config_pass_configs is present to avoid JIT cache key collision
- Add example and test for per-config pass_configs tuning

* docs: add docstrings to meet coverage threshold

* refactor(autotuner): move pass_configs resolution to tuner so grouped compile handles uniform configs

* fix(autotuner): define _PASS_CONFIGS_KEY constant and fix jit_compile race condition

- Replace all magic string "__pass_configs__" with _PASS_CONFIGS_KEY constant
- Move _pass_configs_lock to cover both branches in jit_compile to prevent
  concurrent reads of temporarily mutated pass_configs

* refactor(autotuner): introduce type aliases for compile-unit bookkeeping

- Add ConfigArg, UnitItem, UnitResult, CompileUnit type aliases
- Replace all long tuple type annotations with aliases
- Expand example configs to include num_stages=[1,2] and enable_grouped_compile

---------

Co-authored-by: daniu.ld <daniu.ld@alibaba-inc.com>
2026-07-13 22:24:19 +08:00
禾可 207f3a75f2 [BugFix] Reject non-positive thread extents in T.Kernel (#2653)
* [BugFix] Reject non-positive thread extents in T.Kernel

`_normalize_threads` passed its argument straight through, so a non-positive
extent such as `T.Kernel(1, threads=-1)` compiled and launched cleanly while
writing nothing. Validate the normalized extents and raise instead.

Concrete integer extents only, so symbolic extents keep working.

Fixes #2633.

Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>

* [BugFix] Document _normalize_threads contract

Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>

---------

Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>
2026-07-13 17:29:01 +08:00
Chennes c8c49d5025 [BugFix] Reject invalid dtypes in T.dp4a (#2652)
T.dp4a previously accepted arbitrary buffer dtypes and passed them to backend DP4A implementations, which could silently reinterpret the operand bytes and produce incorrect results.

Require both inputs to use int8 and the accumulator to use int32. Add regression coverage for valid Buffer and BufferLoad forms and for invalid input and accumulator dtypes.

Fixes #2600
2026-07-13 17:13:24 +08:00
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>
2026-07-13 11:50:03 +08:00
Yongqi Zhuo 1e07514959 [CUDA] Support fp32x2 ops as reducers (#2637) 2026-07-12 21:47:32 +08:00
Tong WUandLeiWang1999 228c7c0486 [Feature] Support iket profiler for CUDA backend (#2515)
* Add experimental IKET CUDA backend hooks

* Keep IKET APIs under namespace

* Use real Perfetto screenshot for IKET example

* Show IKET events in Perfetto timeline screenshot

* Add IKET validation overhead example

* Remove IKET validation overhead example

* Fix IKET lint issues

* Move IKET integration to CUDA tools

* Add tools documentation and fix workflow regressions

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-07-12 20:03:30 +08:00
Yuanyuan Zhao bd76473830 [BugFix] Preserve guard identity in LoopUnswitching (#2585) 2026-07-11 14:29:20 +08:00
Lei Wang 250c1fc9ed [TIR][Codegen] Preserve decoupled cast buffer scope (#2545)
Preserve lexical scope for decoupled casts
2026-07-08 14:36:12 +08:00
Lei Wang 8533d2a987 [Language][Scheduler] Expose scalar tile scheduler state (#2553)
Expose scalar tile scheduler state
2026-07-08 14:34:55 +08:00
Lei Wang efce11ed25 [CUDA][Codegen] Keep RNG state in kernel scope (#2540)
Fix CUDA RNG state scope
2026-07-07 14:41:03 +08:00
Lei Wang 0d97af896e [CUDA][Cache] Include compile options in binary cache key (#2532)
* Include CUDA compile options in binary cache key

* enhance
2026-07-07 00:09:22 +08:00
Shuyi LinandClaude 16af6f805e [Feature][Tool] Add pass_visualizer: structure-tree pass browser (#2449)
* [Feature] Add pass_visualizer: structure-tree pass browser

Add an interactive, pass-by-pass IR structure-tree visualizer under
tilelang/utils/pass_visualizer. It complements the existing text-level
pass_diff tool by rendering the SBlock structure tree and expanding tile
ops by field name, with per-class operator highlighting (tile op / sync /
lowered hardware intrinsic) across the CUDA lowering prologue.

Includes a CLI entry point, an example gemm_relu kernel, a pytest suite,
and a docs section contrasting it with pass_diff.

Co-Authored-By: Claude <noreply@anthropic.com>

* Move pass_visualizer to tilelang/tools and fix prologue + review issues

- Move package from tilelang/utils to tilelang/tools per review
- Add missing MaterializeKernelLaunch pass (fixes CUDA CI modulo-by-zero)
- Escape "</" in embedded JSON to prevent <script> injection in HTML report
- Use explicit utf-8 encoding for source/output file I/O
- Guard _fmt_shape against symbolic shapes and validate module spec

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
2026-07-06 15:39:54 +08:00
penguin_wwy af2483f904 [BugFix] Fix llvm auto backend resolution (#2519) 2026-07-06 10:55:13 +08:00
Lei Wang 14fc963089 [Transform][Layout] Avoid thread-indexed replicated fragment readback (#2514)
Avoid thread-indexed replicated fragment readback
2026-07-03 14:52:24 +08:00
Chennes b8dc37e66f [BugFix] Fix LayoutInference divide-by-zero on non-power-of-two broadcast (#2469)
* [BugFix] Avoid zero-extent layout leftovers

* [Test] Cover 24/40 widths in #2394 CUDA numerical regression
2026-07-03 14:30:26 +08:00
Tong WU aca581263a [Feature] Expose multiple CUDA intrinsics (#2473)
* [CUDA] Add T.__fns intrinsic for find-nth-set bit

Expose CUDA __fns in TileLang for direct lookup of the k-th active lane in a 32-bit bitmask, complementing the existing T.__ffs helper.

* Add builtin intrinsics for shared ldst and atomics
2026-07-02 17:29:54 +08:00
penguin_wwyandLeiWang1999 47d1b48c21 [BugFix] Skip source-compilation options when exporting LLVM module (#2467)
* [BugFix] Skip source-compilation options when exporting LLVM module cache

* Add regression test

* Refactor kernel cache export handling and improve test coverage

- Renamed `_get_compile_args` to `_get_source_compile_args` for clarity.
- Introduced `_get_export_link_args` to manage export link arguments.
- Updated `_safe_write_executable` to accept export kwargs directly.
- Enhanced `TVMFFIKernelCache` to determine export kwargs based on target host.
- Modified tests to validate new export behavior and ensure source options are correctly handled for different targets.

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-07-02 16:41:07 +08:00
Jayce Suanddingsg 39ade83ff2 [BugFix][CUDA] Use PTX v4 atomics for fp16/bf16 atomic_addx4 (#2492)
Use PTX v4 atomics for 16-bit CUDA atomic_addx4

The initial fp16/bf16 atomic_addx4 fix avoided the generic float4 path by splitting each operation into two AtomicAddx2 calls. That fixed correctness, but it still used the simulated implementation on Hopper where PTX provides native vector atomics.

Add sm90+ half_t and bfloat16_t AtomicAddx4 overloads backed by the PTX vector forms atom.global.v4.f16.add.noftz and atom.global.v4.bf16.add.noftz, including the existing memory-order variants. Keep the x2 fallback for pre-sm90 targets because PTX vector f16/bf16 atomics require sm90 or newer.

Keep the CUDA language regression focused on fp16/bf16 atomic_addx4 codegen and offset coverage.

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
2026-07-02 16:21:02 +08:00
Lei Wang 3ebd73de4f [Pipeline] Fix physical async wait counts (#2505)
* Fix physical wait counts for async pipeline groups

* Add commit group tracking for async pipeline in inject_pipeline.cc

- Introduced buffer_to_commit_group mapping and commit_group_count in AsyncStateGlobal.
- Enhanced wait handling for tail consumers in async pipeline to ensure correct physical wait counts.
- Updated tests to validate fine-grained physical waits and tail drain waits for async pipeline.
2026-07-02 15:08:39 +08:00
Chennes e2c2c207e6 [BugFix] Sign-extend packed uint32 signed decode (#2500)
_tir_u32_to_int_to_float decoded signed sub-word fields from uint32 storage as unsigned values because it only masked the selected field before casting. As a result, negative encodings such as signed int4 `0xF` decoded as `15.0` instead of `-1.0`.

Sign-extend the extracted field through int32 before casting to the requested float dtype, matching the existing packed signed-int decode helper.

Add CUDA regression coverage for signed int2, int4, and int8 values decoded from uint32 storage.

Fixes #2482
2026-07-02 13:12:56 +08:00
Yongqi Zhuo b1e1e250a2 [Backend] [CUDA] Support GMMA/UMMA lowering for sliced SMEM layout (actually arbitrary layout) (#2452)
* Add split-k for example_warp_specialize_flashmla

* Rotate QK issue to end of main loop to allow more latency hiding with TMA issue

* During rescale, read max score from REG instead of SMEM

* Move K_pe_shared_1 load

* Use stmatrix for SP0 write to reduce SMEM bank conflict

* Support arbitrary SMEM layout in MMA lowering

* Remove runtime guard for increase_descriptor_offset
2026-07-01 18:21:53 +08:00
Xiangwen Wang be2360a5d0 [Enhancement] Fix T.assume and loop bounds to eliminate redundant boundary checks (#2502)
Fix T.assume and loop bounds to eliminate redundant boundary checks
2026-07-01 18:14:01 +08:00
Lei Wang 465c6c5e8c [CUDA] Reduce template include overhead (#2474)
* [Refactor] Remove unsupported tl_gemm and tl_gemm_sp operations

* Removed references to tl_gemm and tl_gemm_sp from various files, including code generation and built-in definitions, as these operations are currently unsupported.
* Updated related comments and diagnostics to reflect the removal of these operations, ensuring clarity in the codebase.
* This change simplifies the handling of unsupported operations and improves overall code maintainability.

* Reduce CUDA template include overhead

* Fix CUDA reduce fast min max helpers

* lint fix

* Fix CUDA swizzle ceil div helper

* Restore CUDA wait_wgmma intrinsic helper

* Fix CUDA cython TMA host adapter include

* Remove unnecessary check for Cutlass fast math header in CUDA code generation

* Update LibraryGenerator to correctly handle CUDA gencode flags for shared-library compilation, ensuring compatibility with target architecture and avoiding Hopper-only instruction issues.

* Refactor smem_ptr_to_uint function to use __cvta_generic_to_shared for improved clarity and performance.

* Refactor example scripts and CUDA math intrinsics for improved performance and clarity

- Updated `example_dequant_gemm_fp4_hopper.py` to directly call `run_regression_perf()` and print latency, removing unused argument parsing code.
- Modified `example_topk.py` to specify the backend as "cupti" in the benchmarking function.
- Enhanced CUDA code generation by introducing `RequiresTileLangMathHeader` to manage math header dependencies for specific functions.
- Updated barrier operations in `barrier.h` to use the correct syntax for shared memory barriers.
- Improved testing for bfloat16 math intrinsics by replacing pytest skips with decorators for CUDA availability checks.

* Refactor CUDA code for improved readability and consistency

- Reformatted the `RequiresTileLangMathHeader` function for better readability.
- Updated assembly syntax in `barrier.h` for clarity and consistency.
- Cleaned up whitespace in the bfloat16 math intrinsics test file to enhance code cleanliness.

* Remove unnecessary whitespace in CUDA code for improved readability

* Add TileLang intrinsic for prefetching TMA descriptor

- Introduced `prefetch_tma_descriptor` intrinsic in C++ and updated relevant files to use `call_intrin` instead of `call_extern`.
- Updated CUDA code generation to handle the new intrinsic correctly.
- Modified tests to ensure the new intrinsic is called appropriately and that no `call_extern` is used for internal TileLang operations.
- Enhanced documentation in SKILL.md regarding the addition of new TileLang intrinsics.

* Add bfloat16 hexp overload and update common.h for compatibility

- Introduced `hexp` overload for `bfloat16_t` in `common.h` to bridge TileLang and CUDA's handling of exponential functions.
- Updated tests to verify the presence of the new `hexp` definition in `common.h`.
- Enhanced CUDA code generation to support the new intrinsic for better performance and compatibility.

* Update common.h for CUDA compatibility with aligned double4 types

- Adjusted preprocessor condition in `common.h` to maintain compatibility with CUDA versions prior to 13 by ensuring proper definition of aligned double4 types.
- This change enhances the compatibility of the code with older NVRTC built-in vector headers.
2026-07-01 16:53:07 +08:00
Lei Wang 0d0e96861a [Env] Require JSON for default target config (#2491)
Require JSON for default target config
2026-06-30 13:27:34 +08:00
Xiangwen Wang 03407d5c33 [Enhancement] Add vectorized fp8x2 <-> fp16/bf16 cast codegen (#2475)
* Add vectorized fp8x2 <-> fp16/bf16 cast codegen

* Minor fix

* Fix lint

* Add unit test
2026-06-29 16:00:57 +08:00
Chennes ed00dfcd7f [BugFix] Improve diagnostic for T.serial fragment access (#2462)
* [BugFix] Reject unsupported T.serial access to thread-distributed fragments with clear diagnostics

A T.serial loop runs sequentially inside each CUDA thread, but a fragment's
elements are distributed across threads, so indexing a fragment by a serial
loop variable has no valid thread-ownership mapping. Three fuzzer reports share
this root cause:

- #2393: serial reduce of a GEMM-output fragment -> unhelpful internal error
  ("contains inner var j").
- #2395: serial write into a fragment then T.copy out -> nvcc "identifier
  undefined" (an owner-thread guard built from serial loop vars was emitted
  outside the loop that binds them).
- #2396: serial write of a fragment to global -> silently wrong (threads
  race-write each cell).

AddWrapperForSingleBufStore wraps bare fragment stores in a degenerate
T.parallel(1) and only rejected constant non-zero fragment indices, so a
dynamic (e.g. serial-loop) index slipped through and lowered to broken code
(#2395/#2396). Reject any fragment index that is not the constant all-zero
index in this fallback path, with a message pointing to T.Parallel /
T.reduce_sum. Rephrase the existing inner-var layout check (#2393) to explain
the cause instead of dumping "contains inner var".

Also collect every access to a buffer in a statement (not just the last) so the
diagnostic is not bypassed when one fragment is accessed at both a constant and
a variable index in the same statement; without it such a case falls through to
a cryptic downstream StructuralEqual internal error. This is a completeness
hardening only, not required by the three issues above.

Fixes #2393
Fixes #2395
Fixes #2396

* Generalize fragment diagnostic to non-parallel loops

inner_vars_ is populated for every non-parallel For (serial, unroll,
vectorized), so the message must not hard-code T.serial. Reword to
"non-parallel loop variable" per review feedback.

* Narrow fix scope to #2393 only

Remove fragment index validation from AddWrapperForSingleBufStore and
delete tests for #2395/#2396. These issues cannot be caught at this pass
stage because T.grid and T.serial are indistinguishable (both appear as
ForKind.SERIAL before LowerTileOp).

Keep only the ParallelOpNode diagnostic improvement which successfully
catches #2393.
2026-06-28 13:05:51 +08:00
Lei Wang 46f3d1c924 [CUDA][JIT][Cache] Add cross-host CUDA binary cache (#2459)
* Add CUDA target code attr support

* Add CUDA target normalization support

Introduced a new function to normalize CUDA targets, ensuring that the target is correctly identified and transformed based on the detected architecture. Added a test to verify that the bare CUDA target uses the detected architecture accurately.

* Add cross-host CUDA binary cache

* Nest host kernel cache under version namespace
2026-06-26 15:56:04 +08:00
Ftoandcoderabbitai[bot] ab083cb3c6 [BugFix] Fix T.Persistent dropping tiles when last dim is not a multiple of group_size (#2455)
* [BugFix] Fix T.Persistent dropping tiles when last dim is not a multiple of group_size (#2433)

* Update src/ir.cc

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2026-06-26 14:12:53 +08:00
Lei Wang 23d2f25673 [CUDA] Add target code attribute support (#2454)
* Add CUDA target code attr support

* Add CUDA target normalization support

Introduced a new function to normalize CUDA targets, ensuring that the target is correctly identified and transformed based on the detected architecture. Added a test to verify that the bare CUDA target uses the detected architecture accurately.
2026-06-25 16:49:20 +08:00
Lei Wang dfc7e953ed [Pipeline] Simplify async copy lowering (#2444)
* Clean up async copy pipeline lowering

* FIX

* enhance

* refactor pipeline
2026-06-24 15:11:05 +08:00
Chennes 74d4f0ed66 [BugFix] Fix make_int sign-extending negative int8 lanes (#2438)
T.fill of an int8 buffer with a negative scalar (other than -1) silently
wrote only 1 of every 4 elements correctly and clobbered the other 3 to
-1. The vectorized fill packs four int8 lanes into one 32-bit store via
make_int(v,v,v,v); a negative signed char is integer-promoted to int and
sign-extended before the shifts/OR, so its high 1-bits flood the
neighbouring lanes (e.g. make_int(-7,-7,-7,-7) returned 0xFFFFFFF9 =
[-7,-1,-1,-1]).

Build the packed value from explicit unsigned bytes so the sign-extension
bits cannot leak across lanes. Add an int8 negative-fill regression test.

Fixes #2427
2026-06-23 14:40:11 +08:00
Tong WUandLeiWang1999 52995c4b7f [Feature] Support stmatrix m16n8 on Blackwell (#2417)
* Add stmatrix m16n8 support

* lint

* Refactor stmatrix code generation to handle optional shape argument

Updated the handling of the shape argument in the stmatrix function across CUDA, CuTeDSL, and HIP code generators. The shape is now checked as the last argument, and the payload start index is adjusted accordingly. This change improves consistency and clarity in the code generation process for tile language operations.

* Refactor stmatrix function naming for consistency

Updated the naming convention of stmatrix functions in CUDA and related files to improve clarity and consistency. The function names now follow a uniform pattern, placing the shape descriptor before the axis identifier. This change enhances readability and aligns with the recent updates in code generation logic.

* Enhance extern call statement printing in CUDA codegen

Updated the `VisitExpr_` method in `codegen_cuda.cc` to improve the clarity of extern call statements. The function now includes comments indicating that the transformation, number, and shape are encoded in the helper name, and only the shared memory pointer and value operands are printed. This change aims to enhance code readability and maintainability.

* Refactor extern call expression printing in CUDA codegen

Introduced a helper function to streamline the printing of extern call expressions in the `VisitExpr_` method of `codegen_cuda.cc`. This change enhances code readability by reducing redundancy and improving the clarity of the generated output for specific function calls, such as `tl::pack_b8x4`.

* Refactor shape argument handling in stmatrix code generation

Updated the handling of the shape argument in the `VisitExpr_` method for both CUDA and CuTeDSL code generators. The shape is now checked as the fourth argument instead of the last, improving clarity and consistency in extern call statement printing. This change enhances the robustness of the code generation process for tile language operations.

* Refactor shape argument handling in stmatrix code generation

Updated the handling of the shape argument in the `VisitExpr_` method for both CUDA and CuTeDSL code generators. The shape is now checked as the last argument and validated accordingly, improving clarity and consistency in the code generation process for tile language operations.

* Refactor tilelang lower function context management

Updated the context management in the `test_tilelang_copy_uses_stmatrix_m16n8_for_sm100_int8_shared_store` test by removing the unnecessary `PassContext` wrapper. This change simplifies the code and enhances clarity in the test setup for tile language operations.

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-06-23 14:21:07 +08:00
jpwang 088f3814ef [BugFix] Fix packed vector type printing for bfloat16 with metal codegen (#2437)
* [BugFix] Fix packed vector type printing for bfloat16 with metal backend

* Resolve review comments

* Resolve review comments
2026-06-23 13:32:11 +08:00
Chennes 9f76d522ed [BugFix] Fix vectorized atomic_add dtype mismatch reinterpret (#2414)
Vectorized atomic_add into an fp16/bf16 destination with an fp32 source
bit-reinterpreted the source float2 as half2/__nv_bfloat162 instead of
numerically converting it, silently producing wrong results. This hit
both the auto-vectorized T.atomic_add path (AtomicAddx2 value form) and
the explicit T.atomic_addx2 API (pointer form).

Add float2 / float* converting overloads (round-to-nearest) to ToHalf2
and ToBfloat162 so the vectorized helper converts instead of
reinterpreting the source bytes. Add mixed-dtype regression tests for
both entry points.

Fixes #2382
2026-06-19 23:21:35 +08:00
Yongqi Zhuo 65dbc9837b [Backend] Support TMA lowering for arbitrary (swizzled) SMEM layout (#2380) 2026-06-18 18:42:01 +08:00
Witherstrike a4ab99e466 [Feature] Add LLVM backend support (#2409)
feat: add end-to-end LLVM backend execution support

Enable LLVM as a TileLang execution backend in CMake and route llvm targets through TVM LLVM codegen during lowering.

Fix tvm_ffi output tensor device inference by deriving the output device from tensor inputs, and add LLVM execution regression tests for matmul, T.gemm, T.copy, and T.While.
2026-06-18 15:25:03 +08:00
Chennes d34109aeb6 [BugFix] Fix vectorized fp16<->bf16 cast compilation (#2407)
A vectorized T.copy between fp16 and bf16 failed to compile. The CUDA
cast codegen has dedicated vectorized branches for fp16<->fp32,
bf16<->fp32, fp8 and fp4, but none for fp16<->bf16, so it fell into the
elementwise fallback. There the lanes load as native __half/__nv_bfloat16,
and a direct (half_t)(__nv_bfloat16) (or reverse) is an ambiguous
user-defined conversion, breaking the build.

Route cross-half elementwise casts through float to disambiguate, and
add a cross-dtype copy regression test for both directions.

Fixes #2385
2026-06-17 13:44:36 +08:00
Jayce Suanddingsg e7da9025c4 [BugFix][CuTeDSL] Fix TileKernels scan, optional-shape, and e5m6 paths (#2369)
* [BugFix] Complete CuTeDSL scan lowering for TileKernels

TileKernels exercises CuTeDSL scan and routing patterns that were not fully covered by the existing Python backend helpers.

The MoE get_fused_mapping kernel computes expert base offsets with T.cumsum over a 256-entry shared buffer. CuTeDSL CumSum1D only performed a single warp-local prefix scan, while the CUDA scan template uses the first warp to scan all 32-element segments and carry the result between segments. As a result, expert prefixes reset at later warp-sized chunks, positions for experts such as 32 and 64 overlapped earlier expert ranges, and pos_to_expert validation failed on H100.

Bring the CuTeDSL scan helpers in line with the CUDA InclusiveScanLine behavior: add reusable one-warp line scans with carry across 32-element segments, route CumSum1D and CumMax1D through them, and use the same line scan for CumSum2D/CumMax2D on both row-wise and column-wise axes. This removes the previous dim=0 H<=32 limitation and fixes multi-segment cummax as well as cumsum.

The same TileKernels path also relies on CuTeDSL codegen accepting non-flat buffer accesses after TIR lowering. Add a shared index linearization helper for loads, stores, address_of/reinterpret, and atomic pointer generation, and lower tl.match_any_sync through the existing CuTeDSL warp intrinsic. Thread binding now falls back to the IterVar name when thread_tag has already been normalized away.

Verified on H100 CUDA/CuTeDSL with CUDA_VISIBLE_DEVICES=0 and TILELANG_TARGET=cutedsl:

- cmake --build build -j$(nproc)

- ruff format --check tilelang/contrib/cutedsl/reduce.py

- ruff check tilelang/contrib/cutedsl/reduce.py

- git diff --check

- pytest testing/python/target/test_tilelang_codegen_cutedsl_scan.py testing/python/language/test_tilelang_language_scan.py -q -s --tb=short: 13 passed

- pytest testing/python/language/test_tilelang_language_warp_vote.py::test_match_any_sync -q -s --tb=short: 1 passed

- pytest tests/moe/test_get_fused_mapping.py::test_get_fused_mapping[num_send_tokens=4001-num_topk=2-num_experts=72-num_ep_ranks=1-alignment=64] -q -s --tb=short: 1 passed

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>

* [BugFix] Resolve CuTeDSL dynamic shapes from live tensors

CuTeDSL records dynamic symbolic dimensions in the same ordering as the CUDA wrapper, but it previously kept only the first buffer shape or stride location for each symbol. That is not sufficient for kernels with optional tensor arguments. TileKernels reduce_fused declares num_tokens on topk_weights, token_topk_to_pos, and out; when with_weights is false the topk_weights argument is None, so the first recorded source is not a live tensor.

The previous adapter fallback substituted 0 when that first source was None. On H100 with TILELANG_TARGET=cutedsl this passed num_tokens=0 to the generated launcher even though token_topk_to_pos had shape (4001, 2), producing grid=[0, 1, 1] and a CUDA error 1 launch failure in reduce_fused.

Keep all shape/stride candidates for each dynamic symbol and resolve the runtime value from the first candidate backed by a real torch.Tensor. Reuse the same resolution helper for allocated output shapes and for dynamic arguments passed to the generated CuTeDSL module, and fail clearly if no live tensor source exists.

Verified with a CuTeDSL adapter regression test for repeated dynamic shape symbols with a None first candidate, and with the TileKernels reduce_fused H100 failure case that previously launched with num_tokens=0.

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>

* [BugFix] Permit dead optional CuTeDSL stride symbols

TileKernels no-scale MoE expand paths can keep a dynamic stride symbol in the generated CuTeDSL host wrapper ABI even when the optional scale tensor is absent from the lowered device kernel variant. In that case the runtime argument is intentionally None, so the stride-only symbol has no live tensor from which the adapter can read torch stride metadata.

The previous dynamic-symbol fix made missing live tensor sources strict to avoid unsafe shape fallbacks such as launching with num_tokens=0 when the first candidate tensor was optional. That strict behavior is still required for shape symbols because they control output allocation and launch dimensions.

Limit the fallback to symbols whose candidates are stride-only. Such values are ABI placeholders for optional strided tensor parameters in pruned variants; returning 0 preserves the generated call signature without weakening dynamic shape resolution. Add adapter regressions covering both the optional stride fallback and the strict missing-shape error.

Verified on H100 with TILELANG_TARGET=cutedsl:

- ruff format/check on the touched files

- pytest testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::test_cutedsl_adapter_resolves_dynamic_symbol_from_live_tensor_candidate testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::test_cutedsl_adapter_allows_optional_stride_symbol_without_live_tensor -q

- pytest tests/moe/test_expand_to_fused.py::test_expand_to_fused[num_send_tokens=4001-num_topk=2-num_experts=9-num_ep_ranks=8-hidden=576] -vv -s --tb=short

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>

* [BugFix] Fix CuTeDSL optional ABI and early-return codegen

TileKernels full functional coverage with TILELANG_TARGET=cutedsl exposed additional CuTeDSL gaps in top2_sum_gate after the broader TileKernels fixes.

First, the host adapter still appended dynamic shape symbols for optional tensors that had been compiled out by static flags. In top2_sum_gate, to_physical_map/logical_count can be None when logical expert mapping is disabled, but their shape-only symbol remains in the wrapper ABI. Keep output allocation strict, but allow dead host ABI dynamic shape arguments to use a zero placeholder when no live tensor source exists.

Second, CuTeDSL cannot emit dynamic thread returns directly, so codegen rewrites if cond: thread_return(); rest into if not cond: rest. The existing recognizer only handled a bare thread_return then body, and missed the common if cond: stores; thread_return() pattern. Preserve the pre-return stores and guard the following statements so masked tokens cannot be overwritten by the fallthrough path.

Third, vector binary codegen reused SSA aliases for operands containing BufferLoad nodes. That is not valid for mutable local tensors after an intervening store. The softmax top2_sum_gate path restored raw logits to scores_local, but the subsequent scores_local + bias expression reused an older softmax load and ranked by softmax(logits)+bias instead of logits+bias. Avoid SSA alias reuse for vector operands containing BufferLoad so the generated code reloads the current tensor value.

Verified on H100 with the local TileLang build and TileKernels CuTeDSL backend:

- cmake --build build -j 32

- ruff format/check on touched Python files

- pytest testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::{dynamic symbol candidate tests} -q: 3 passed

- tests/moe/test_top2_sum_gate.py::test_top2_sum_gate[num_groups=0-num_topk_groups=0-num_routed_experts=72-num_shared_experts=1-num_topk=6] with TILELANG_TARGET=cutedsl TK_FULL_TEST=1: 1 passed

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>

* [BugFix] Promote narrow CuTeDSL shifts before wide casts

TileKernels e5m6 quantization kernels pack 8 truncated fp16 values into three uint32 words. The TileLang source writes expressions such as T.cast(half_u16[i] << 20, T.uint32), which the CUDA C++ backend evaluates with normal integer promotion before the final uint32 cast.

CuTeDSL uses strongly typed cutlass integer wrappers instead of CUDA C integer promotion. The previous CuTeDSL lowering cast the shift result back to the narrow TIR result type first, producing code like cutlass.Uint32(cutlass.Uint16((half_u16[0] << cutlass.Uint16(20)))). For shifts by more than the uint16 width this drops all high bits before the packed uint32 store, corrupting e5m6 payload bytes. Cast-back then faithfully decoded the corrupted payload into nan/inf values, so TileKernels full correctness failed in cast_back_e5m6 and per_token_cast_to_e5m6 cases.

When a narrow integer shift is immediately cast to a wider integer type, emit the CuTeDSL shift with the left operand promoted to the target integer type before shifting. This preserves the intended packed-bit construction while leaving narrow shift expressions that are not widened on the existing path.

Add a CuTeDSL codegen regression test for uint16 << 20 widened to uint32 so future changes do not reintroduce a narrow in-place shift.

Verified on H100 CUDA/CuTeDSL with TILELANG_TARGET=cutedsl and TILELANG_DISABLE_CACHE=1:

- cmake --build build -j 32

- python -m pytest testing/python/target/test_tilelang_codegen_cutedsl_integer_shift.py -q

- TileKernels focused e5m6 quant tests: tests/quant/test_per_token_cast_to_e5m6.py tests/quant/test_cast_back_e5m6.py -n 2: 168 passed, 112 skipped

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>

---------

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
2026-06-17 13:41:26 +08:00
cjandcaojian5 7b7df67acf [feat] add pass diff show for debugging (#2375)
Co-authored-by: caojian5 <caojian5@huawei.com>
2026-06-16 19:38:22 +08:00
Tong WU 41b95a5dfa [Feature] Support st.bulk for shared zero fill on SM100+ (#2403)
feat(cuda): use st.bulk for shared zero fill
2026-06-15 21:40:01 +08:00
Chennes c9e5112332 [BugFix] Fix bf16 CUDA exp self-recursion (#2402)
`hexp` is #define'd to `cutlass::fast_exp`, so the hand-written
`cutlass::fast_exp(bfloat16_t)` overload's `return ::hexp(x)` expanded into a
call to itself: bf16 `T.exp` recursed and returned garbage.

Route through float, mirroring cutlass's own `half_t fast_exp`
(cutlass/fast_math.h), so the call resolves to `cutlass::fast_exp(float)`
(= ::expf). Verified against torch.exp using the #2383 repro.

Fixes #2383.
2026-06-15 21:21:32 +08:00
Lei Wang 9371a66721 [Language] Fix TMA 1D load test lowering (#2401)
* Fix TMA 1D load test lowering

* Refactor test execution in tilelang TMA 1D test script

Updated the main execution flow to directly call the test function instead of running an elementwise add operation, improving clarity and focus on testing functionality.
2026-06-15 15:00:06 +08:00
Lei Wang b35a9aaba2 [CUDA] Align swizzled TMA shared buffers (#2391)
Fix TMA swizzled smem alignment
2026-06-15 01:42:06 +08:00
Chennes b7489e2608 [CUDA] Add SM75 MMA dispatchers for FP16 accumulation and UINT8 (#2392)
The SM75 GEMM emitter derives MMA shapes from input dtype bits alone, so it
already emits the (16,8,8) shape for fp16->fp16 accumulation (identical to the
registered fp16->fp32 case) and the (8,8,16) shape for uint8 inputs (identical
to signed int8). But the C++ dispatcher table only covered the three configs
tested by #2198 (fp16->fp32, int8, int4), so requesting fp16->fp16-accum or
uint8 GEMM on Turing failed at nvcc with an opaque
`tl::mma_sync: unsupported configuration` static_assert.

Add the two missing tl::detail SM75 atoms and register their dispatchers,
following the precedent set by #2198's hand-written int4/uint4 structs (CuTe
ships only fp16->fp32 and signed-int8 SM75 atoms). The uint8 atom uses the
non-saturating PTX form to match the existing signed-int8 path.

Verified on a TITAN RTX (SM75): emitted PTX is
mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 and
mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32, and end-to-end GEMM numerics
match torch.
2026-06-14 02:41:36 +08:00
赖泉laiquanandMAC 17c4b384df Fix bf16 CUDA rsqrt overload (#2386)
fix: add bf16 CUDA rsqrt overload

Co-authored-by: MAC <mac@MacBook-Air.local>
2026-06-12 23:30:11 +08:00
Qiping PanandLeiWang1999 200190346a [Fix] Bind statement missing role marker (#2362)
* fix: name

* [Transform] Enhance role management and scalar dependency propagation

- Introduced a new Role enumeration value `kNeutral` to better represent variable states.
- Added `CombineRoles` function to manage role combinations effectively.
- Implemented `ScalarAccess` structure and `ScalarAccessCollector` class for tracking variable definitions and usages.
- Enhanced `PropagateScalarDependencyRoles` function to propagate roles based on scalar dependencies across statements.
- Updated role assignments in various statement visitors to utilize the new role management logic, improving clarity and correctness in role handling.

* Update main execution flow in test_tilelang_language_let.py to enable tilelang testing framework

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-06-12 17:35:42 +08:00