mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Gabriel Wuandyuxic
391cf5b6b9
* fix(cutedsl): add rsqrt/rsqrtf support in codegen
- Add rsqrt/rsqrtf to fastmath function mapping
- Implement ieee_frsqrt using tl.rsqrt
- Export rsqrt from contrib/cutedsl/math.py
Fixes: norm/test_rms_norm.py
* fix(cutedsl): add atomic_add support and fix shared.dyn access
- Add atomic_add_elem_op and atomic_add_ret_elem_op handling
- Exclude shared.dyn from direct indexing to avoid out-of-bounds access
(shared.dyn shape is set to (1,) but accessed with arbitrary indices)
Fixes: gemv/test_example_gemv.py
* fix(cutedsl): add builtin::ptx_ldmatrix support
Implement builtin::ptx_ldmatrix by reusing tl.ptx_ldmatrix_x* functions.
* feat(cutedsl): add ptx_mma support with all MMA types
Add complete PTX MMA intrinsic support for CuTeDSL backend:
- Implement ptx_mma codegen in codegen_cutedsl.cc
- Add GetVarPtr_ helper for local buffer pointer extraction
- Create ptx_mma.py with inline PTX assembly wrappers for:
- FP16: m16n8k16 -> f16/f32
- BF16: m16n8k16 -> f32
- INT8/UINT8: m16n8k32 -> i32
- INT4/UINT4: m16n8k32 -> i32
- TF32: m16n8k4, m16n8k8 -> f32
- FP64: m8n8k4 -> f64
- FP8 e4m3/e5m2: m16n8k32 -> f16/f32
- Fix BroadcastNode to emit typed values for make_filled_tensor
- Update ptx_ldmatrix to use GetVarPtr_ for local buffers
* feat(cutedsl): comprehensive atomic operations support
- Add atomic.py module with full atomic operations:
- AtomicAdd: float32 (nvvm.atomicrmw), float16 (inline PTX with i16 bitcast),
int32/int64 (nvvm.atomicrmw)
- AtomicAddx2/x4: vectorized float32 via PTX atom.global.add.v2/v4.f32
- AtomicMax/Min: int32/int64 (nvvm.atomicrmw), float32 (CAS loop via inline PTX)
- AtomicLoad/Store: all memory orderings via inline PTX (relaxed/acquire/release/seq_cst)
- Add codegen support for atomic_load/store/max/min/addx2/addx4 in codegen_cutedsl.cc
- Move AtomicAdd from __init__.py to atomic.py, extend with float16/int64 support
* feat(cutedsl): add fabsf/fabs support in codegen and math module
* feat(cutedsl): add wait_wgmma support via nvvm.wgmma_wait_group_sync_aligned
* feat(cutedsl): support i32 min/max
* fix(cutedsl): handle FP8 narrow precision alignment in codegen
CuTeDSL requires all FP8 vector operations to be 32-bit aligned, but
TIR generates float8x2 (16-bit) vectors. Additionally, MLIR vector
extractelement on FP8 types fails with unrealized_conversion_cast.
Fix CastNode, BufferLoadNode, and BufferStoreNode to:
- Pad unaligned FP8 casts to aligned width (e.g. f32x4 -> fp8x4)
- Return rmem tensor names (not .load()) so downstream uses rmem
element access instead of MLIR vector extractelement
- Use padded aligned stores for local rmem, scalar element stores
for global/shared memory (to handle non-aligned offsets)
* fix(cutedsl): use pointer arithmetic for shared.barrier buffer access
_Pointer type from tl.alloc_smem() does not support subscript [i],
but supports pointer arithmetic (ptr + i). This fixes warp_specialize
tests that use mbarrier.
- GetBufferRef_: use (vid + index) for shared.barrier instead of vid[index]
- BufferLoadNode: skip .load() suffix for shared.barrier (returns pointer, not value)
Fixes 3/4 warp_specialize tests.
* fix(cutedsl): handle single coordinate in tma_load
When tma_load receives a single Int32 coordinate (for 1D tensors) instead
of a tuple, wrap it in a tuple before iterating. This fixes the error:
'Int32' object is not iterable
Fixes flash_decoding/test_example_flash_decoding.py
* feat(cutedsl): add decode_i4u/s_to_f16 and warp primitives
- Add quantize.py with decode_i4u_to_f16 and decode_i4s_to_f16
using inline PTX (lop3.b32 + sub.f16x2)
- Add warp.py with __activemask, __shfl_down_sync, __shfl_sync
using CuTeDSL's shuffle_sync functions
- Update codegen to map decode_* and __* functions to tl.* namespace
Fixes dequant_gemv_fp16xint4 test.
* feat(cutedsl): add decode_fp4_to_bf16_twiddling
Implement FP4 to BF16 conversion using bit twiddling technique.
- Use prmt.b32 for endianness handling
- Use mul.bf16x2 for scaling
- Pack output with correct high/low BF16 layout
Fixes test_example_dequant_gemm_bf16_mxfp4_hopper.
* fix(cutedsl): support reinterpret for arbitrary expressions and fix shift type promotion
- reinterpret: handle non-BufferLoad args with tl.bitcast()
- shift_right/shift_left: wrap result in dtype() to preserve type
- CuTeDSL Int8>>4 returns Int32 (Python promotion), CUDA keeps int8_t
- Add bitcast() function using llvm.bitcast MLIR op
* feat: disable v2->v1 fallback
* fix(cutedsl): TensorSSA for store (Select/scalar) and TF32 MMA dispatch
- ArithValue: emit tl.where() for Select and wrap scalar store RHS in
make_filled_tensor().load() so CuTeDSL Tensor.store() gets TensorSSA.
- TF32: accept fp32/f32/float32 for m16n8k4 and m16n8k8 in ptx_mma dispatch
(e.g. deepseek_mhc).
* feat(cutedsl): implement CumSum1D and CumSum2D for cumulative sum operations
- Add CumSum1D and CumSum2D classes in reduce.py based on CUDA templates
- Add __shfl_up_sync wrapper in warp.py using CuTeDSL's shuffle_sync_up
- Implement warp-level prefix sum (forward/reverse) using shuffle operations
- Use @cute.jit decorator and nested if for DSL-compatible control flow
- Fixes gdn/test_example_cumsum_compilation test (7/8 gdn tests now pass)
* feat(cutedsl): implement tma_reduce for TMA atomic add reduction
* fix(cutedsl): handle if_then_else as tl.where and unify Select/where arm dtypes
- Override builtin::if_then_else in CallNode to emit tl.where() instead of
Python ternary (which yields ArithValue, not TensorSSA)
- Use make_rmem_tensor + element assignment instead of make_filled_tensor
(which creates unsupported vector<1xT> types for int)
- Unify arm dtypes in both SelectNode and if_then_else: when true/false arms
differ (e.g. BFloat16 vs Float32), cast both to common type via .to()
- Fix BufferStore: skip make_filled_tensor wrapping for Select/if_then_else
values (already TensorSSA); extract [0] for scalar element assignment
- Fix GetBufferRef_: extract Ramp base for make_tensor_at_offset offset
parameter to avoid emitting tuple offsets
Fixes: cast/group_per_split_token, topk, deepseek_mhc, flash varlen
* fix(cutedsl): convert non-integer offsets to Int64 in make_tensor_at_offset
Complex arithmetic on cutlass.Int64 values (involving tl.min, //, %) can
produce Float64 ArithValue that CuTeDSL Pointer.__add__ does not accept.
Detect non-integer offsets via is_integer() and wrap with cutlass.Int64().
Fixes blocksparse_attention varlen_indice and varlen_mask (6/6 passed).
* fix(cutedsl): fix narrow unaligned load/store for sub-32-bit types
Three issues in the narrow unaligned (e.g. uint8x2, fp8x2) code paths:
1. Load from shared/global: using aligned_lanes (4) elements caused OOB
near buffer boundaries and MisalignedAddress for non-aligned offsets.
Fix: load exactly value_lanes (2) elements with no cute.assume.
2. Load from rmem: cute.assume(offset, divby=aligned_lanes) silently
truncates offsets not divisible by aligned_lanes (e.g. offset=6 with
divby=4 → 4). Fix: drop cute.assume for rmem (no hardware constraint).
3. Store to rmem: padded (aligned_lanes,) .store() writes beyond
value_lanes elements, causing overlapping writes and OOB at tensor
boundaries. Also triggers CuTeDSL's 32-bit alignment check for FP8.
Fix: use element-by-element assignment (vid[i]=val).
Fixes dequantize_gemm w4a8 (cudaErrorMisalignedAddress → correct results).
* fix(cutedsl): handle signless int types in tl.min/tl.max fallback
* fix(cutedsl): handle _Pointer src_values in AtomicAddx2/x4
* fix(cutedsl): remove double PrintIndent in cp_async_gs_conditional codegen
* fix(example): support SM100 in attention_sink bwd configs
Change get_bwd_configs() to use >= 90 instead of == 90 so that SM100
(Blackwell) uses the same config as SM90 (Hopper): block_M=128,
block_N=32, num_stages=2, threads=256.
* feat: implement WGMMA V2 (ss + rs) for CuTeDSL backend
Implement explicit WGMMA descriptor management ops in the CuTeDSL backend,
replacing LOG(FATAL) stubs with working codegen for all 8 WGMMA-related ops.
New Python files:
- typing.py: WGMMA type map (string dtype -> nvvm.WGMMATypes)
- gemm_V2.py: GmmaDescriptor class, initialize/increase descriptor,
warpgroup sync ops, wgmma_ss (via nvvm.wgmma_mma_async),
wgmma_rs (via PTX inline asm)
Key findings and fixes:
- nvvm.mma_smem_desc does not correctly pack the layout_type/swizzle
field in the WGMMA descriptor; replaced with manual bit packing
- wgmma_rs implemented using llvm.inline_asm with dynamically generated
PTX and constraint strings based on M/N/K Constexpr parameters
C++ codegen changes (codegen_cutedsl.cc):
- warpgroup_arrive/commit_batch/wait/fence_operand
- ptx_wgmma_ss (15 args) and ptx_wgmma_rs (14 args)
- initialize_wgmma_descriptor (arg reorder: TIR->Python)
- increase_descriptor_offset
All GEMM tests (4/4) and flash attention tests (14/14) pass.
* fix(cutedsl): replace break with guard variable pattern for CuTeDSL compatibility
* feat(cutedsl): implement tcgen05 (SM100/Blackwell) MMA support
Add CuTeDSL codegen and Python runtime support for tcgen05 MMA operations:
- New gemm_tcgen05.py: Tcgen05SmemDescriptor, descriptor initialization,
MMA variants (ss/ws_ss/ts), mbarrier arrive, TMEM alloc/dealloc/load
- All MMA and commit instructions guarded by elect_one_sync via inline PTX
(@q predicate on individual instructions, not on block scopes)
- TMEM load uses recursive power-of-2 splitting with direct MLIR emission
- C++ codegen: replace 6 LOG(FATAL) stubs with proper code generation for
tcgen05 MMA/descriptor/TMEM ops, handle descriptor scope allocation
- Fix shared.barrier pointer access (skip .iterator for alloc_smem results)
- Fix mbarrier_wait: use blocking PTX loop instead of single try_wait call
* fix(cutedsl): correct tcgen05_ld 64b/128b/256b PTX variants and cleanup
- Fix tcgen05_ld_32dp64bNx/128bNx/256bNx: use correct PTX instructions
(16x64b/16x128b/16x256b) with two-half decomposition instead of
hardcoded 32x32b
- Generalize _emit_tmem_ld_segment to accept ptx_type and regs_per_x
- Remove dead expression in ptx_mma.py (c_ptr + a_offset)
- Remove unused GemmCuTeDSL and is_cutedsl_target imports
* fix(cutedsl): review cleanup — warp shuffle, mbar timeout, atomic addr space, dead code
- Fix warp.py: correct mask_and_clamp formulas to match CUDA semantics
for __shfl_down_sync, __shfl_up_sync, __shfl_sync (width param was ignored)
- Fix mbar.py: use timeout_ns parameter instead of hardcoded 0x989680
- Fix atomic.py: use generic address space (ld.f32, atom.cas.b32) instead of
.global for AtomicMax/AtomicMin CAS loops, supporting shared memory pointers
- Add ICHECK for non-unit Ramp stride in codegen_cutedsl.cc GetBufferRef_
- Remove unreachable duplicate elif gemm_inst.is_tcgen5mma() branch
- Remove unused _to_ir_value helper from ptx_mma.py
* fix(cutedsl): skip BulkLoad1D for CuTeDSL + fix uint broadcast type mismatch
Two fixes for dequant/groupedgemm_bf16_mxfp4_hopper on CuTeDSL:
1. Skip BulkLoad1D/BulkStore1D for CuTeDSL target in copy.cc:
cp_async_bulk_shared_cluster_global (raw 1D TMA) combined with WGMMA
in the same kernel triggers a ptxas ICE in the NVPTX backend.
Falls through to descriptor-based BulkLoad/BulkStore instead.
2. Use signed int type for BroadcastNode fill values in codegen_cutedsl.cc:
CuTeDSL/MLIR normalizes unsigned integer tensor loads to signed types
(Uint8 pointer -> i8 elements). BroadcastNode fill values now use
matching signed types to avoid Int8/Uint8 mismatch in tl.where().
* fix(cutedsl): support fp16/bf16 in AtomicAddx2 via atom.add.noftz.v2.f16/bf16
* ci(cutedsl): add CuTeDSL examples to CI pipeline
Add a new GitHub Actions step that runs all examples with
TILELANG_TARGET=cutedsl. Known failures (sparse ops) are marked as
xfail and unsupported tests (convolution, stream-K, flash_decoding FA3)
are auto-skipped via pytest_collection_modifyitems in conftest.py.
The step uses continue-on-error: true initially until stabilized.
* chore: fix pre-commit
* fix(adapter): use name-based fallback for tir.Var dynamic symbolic lookup
tir.Var uses object identity for hash/eq, so when param_shapes (from
the original PrimFunc) and dynamic_symbolic_map (from the transformed
PrimFunc in ir_module) contain different tir.Var objects with the same
name, direct dict lookup fails with KeyError.
Add a secondary name-based index (_dynamic_symbolic_name_map) and a
_lookup_dynamic_symbolic() helper that falls back to name matching.
This fixes KeyError: batch under pytest-xdist parallel execution.
Affected adapters: cutedsl, nvrtc. (cython and tvm_ffi already had
a str()-based workaround.)
* chore: remove dead GemmCuTeDSL class and update stale V1 comments
- Delete `gemm_cutedsl.py` (no references anywhere in codebase)
- Remove outdated comment in `gemm_op.py`
- Update `run_local_ci_test.sh` comment: CuTeDSL no longer requires V1
- Keep `TILELANG_USE_GEMM_V1` env var as a debugging escape hatch
* refactor(cutedsl): reorganize contrib/cutedsl module structure
- Rename gemm_V1.py/gemm_V2.py to snake_case (gemm_v1.py/gemm_v2.py)
- Merge typing.py (19 lines) and mbar.py (53 lines) into utils.py and cpasync.py
- Extract utility functions from __init__.py into new utils.py module
- Slim __init__.py to pure re-exports
- Add __all__ to all submodules for explicit public API documentation
* fix(ci): use .empty() instead of .size() >= 1 in pipeline_planning.cc
* fix: coderabbit comments
* fix(cutedsl): correct WGMMA result register count for f16/bf16 output
wgmma_ss/wgmma_rs used M*N//128 for result registers, which is only
correct for 32-bit output types. For f16/bf16, each i32 register holds
2 packed elements, halving the count. Derive from cutlass type .width
instead of hardcoding.
Also update barrier test assertion to match codegen's pointer-arithmetic
syntax: (barriers + i) instead of barriers[i].
* refactor(cutedsl): templatize ptx_mma.py with factory function
Replace 13 near-identical @dsl_user_op PTX MMA functions (~840 lines)
with a single _make_ptx_mma factory and a config table (~120 lines).
Also correctly interpolates a_layout/b_layout into the PTX mnemonic
instead of hardcoding row.col.
* fix(cutedsl): fix NaN CAS loop, wgmma_rs int accum, and stale pipeline state
- atomic.py: Use integer-domain comparison (setp.ne.b32) in AtomicMax/Min
CAS loops to prevent infinite spin when values are NaN
- gemm_v2.py: Use "=r"/i32 constraints for wgmma_rs when C_dtype is s32
(int8 GEMM) instead of hardcoded "=f"/f32
- reduce.py: Add assert H <= 32 guard for CumSum2D dim=0 column-wise path
- pipeline_planning.cc: Clear pending tcgen05 state unconditionally after
arrive to prevent stale entries from leaking to future arrives
* fix(cutedsl): fix ptx_mma brace escaping and wgmma_rs DSL compatibility
- ptx_mma.py: Double-brace register lists in PTX template so
.format(a_layout=, b_layout=) doesn't interpret {$0, ...} as
format fields (KeyError regression from templatization)
- gemm_v2.py: Replace if/else branches inside DSL-traced loops with
ternary expressions to avoid CuTeDSL "operands type change inside
dynamic if" error
- ci.yml: Remove continue-on-error from cutedsl CI, reduce maxfail
to 3, increase parallelism to 4
* feat(cutedsl): implement ptx_mma_sp (sparse MMA) codegen and runtime
* fix(cutedsl): guard saturate param, move CuTeDSL CI step after core tests
- ptx_mma/ptx_mma_sp: raise NotImplementedError if saturate=True
- warpgroup_fence_operand: add comment explaining no-op rationale
- ci.yml: move CuTeDSL examples after CUDA/ROCm/Metal test steps
so a CuTeDSL failure no longer blocks core tests
* feat(cutedsl): implement 16 previously unsupported ops
Add support for thread index queries (get_lane_idx, get_warp_idx,
get_warp_idx_sync, get_warp_group_idx), exp10 math function,
warp-level reductions (sum, max, min, bitand, bitor), and
IEEE-754 compliant arithmetic with explicit rounding modes
(add, sub, mul, fma, rcp, sqrt, div).
Reduces unsupported op count from 26 to 10. The remaining ops
are either high-complexity (tma_load_im2col, tl_gemm_sp),
require wrapper changes (sync_grid), or are legacy/unused.
* fix(cutedsl): wrap non-int32 integer literals with CuTeDSL type constructors
CuTeDSL's as_numeric() converts bare Python ints to Int32, causing
"Type mismatch" errors when storing to narrower tensors (e.g. Int16).
Override IntImmNode to emit cutlass.Int16(0) instead of bare 0 for
non-int32 literals. Also add a BufferStore narrowing guard as safety net.
Fixes compress_kernel Int32→Int16 error in gemm_sp/custom_compress.
* feat(cutedsl): implement sync_grid with cooperative kernel launch
Add grid-level synchronization via a software barrier using
llvm.mlir.global counter + inline PTX atomic spin-wait. The CuTeDSL
wrapper detects sync_grid usage and automatically switches to
cuLaunchCooperativeKernel to guarantee all blocks are resident.
- grid_sync.py: module-global counter + PTX barrier protocol
- wrapper.py: CPP_COOPERATIVE_KERNEL_LAUNCH_TEMPLATE
- codegen_cutedsl.cc: emit tl.sync_grid() instead of LOG(FATAL)
---------
Co-authored-by: yuxic <yuxic@nvidia.com>