Resolve conflicts in merge_shared_memory_allocations.cc and codegen_cuda.cc:
upstream's new shared common/storage_size.h GetBufferStorageSizeBytes helper
covers the packed-FP4 exact-bits sizing this branch introduced (and extends
it to 4-bit integer types), so the sizing conflicts take the upstream side
and the branch-local sizing helpers are dropped. The NVFP4
byte-offset -> logical-index conversion in the non-alias rewrite stays, as
upstream has no equivalent. codegen_cuda.cc's 4-bit shared-allocation
handling was equivalently refactored upstream (is_byte_packed_4bit) and
already includes the scalar float4_e2m1fn case.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
T.copy_ue4m3_scale_tile and T.ue4m3_scale_tile_source_coords were
NVFP4-specific helpers on the general T.* surface. Remove both and move the
abstraction down a level:
- The example now declares the packed scale source with its honest "tile rows"
shape ([n_mn_blocks * n_k_blocks * block_MN, block_K // 64] uint32): the
blockscaled_chunk_kmajor layout stores each (mn_block, k_block) tile
contiguously, so with this view every tile is a plain rectangular slice and
staging is a two-line T.Parallel loop with row-offset addressing only. The
host obtains the view zero-copy via
swizzle_blockscaled_chunk_kmajor_scale_words(...).reshape(-1, 4).
- tilelang.quantize.nvfp4 gains blockscaled_chunk_kmajor_tile_source_coords,
the hand-written offset formula for staging schedules that keep the 2-D
[MN, K/64] view (the maint benchmark's producer warp group uses it).
Compiler-internal staging was prototyped first (plain T.copy of the slice)
and rejected for now: the raw-pointer bulk-TMA path this lowers to on some
shapes never arrives on its mbarrier (deadlock, tracked upstream), and the
synchronous fallback stalls the producer warps (-8.5% at 2048^3). The example
comment documents this so the loop can become a plain T.copy once fixed.
Validation on RTX 5090: 123 tests pass; example outputs bitwise-identical to
the previous macro implementation on 5 shapes including M/N tail tiles
(257x384x512, 130x128x256, 128x136x256); example perf at parity
(100.1-100.4% over 2048^3/4096^3/8192^3/8192x8192x1536, interleaved A/B,
generated CUDA identical modulo a loop variable name); maint WS benchmark
kernel source byte-identical with --verify passing and perf unchanged
(641.9 / 1310.5 / 1149.5 TF at 2048^3 / 8192^3 / 8192x8192x1536).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Address maintainer review follow-ups:
- Move the four loose SM120 NVFP4 files under maint/gemm/ into a dedicated
maint/gemm/gemm_sm120/ folder to distinguish them from other gemm types,
updating the repo-root path resolution, the benchmark usage string, the
cutlass_nvf4_ref.cu source path, and the quantizer loader in the scale-layout
test.
- Guard the SM120-hardware test points with
tilelang.testing.requires_cuda_compute_version_eq(12, 0): replace the manual
capability skip in the example CLI test and tighten the block-scale kernel
tests from _ge(12, 0) to _eq(12, 0).
- Reuse the native tilelang.cuda.target.target_is_sm120 helper in
_is_explicit_non_sm120_cuda instead of hand-parsing the arch string, keeping
the explicit-arch guard so targets without an arch attr never trip the C++
ICHECK on hosts that have no SM120 GPU.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
T.atomic_load and T.atomic_store previously accepted every order in the shared memory-order map, including operation-illegal combinations that trigger device-side assertions in libcu++.
Validate memory orders against operation-specific allow lists before lowering, report invalid values with a frontend ValueError, and update the documentation and regression coverage accordingly.
Fixes#2574
* [BugFix] Return previous value for atomic min/max/add (scalar + vectorized)
`atomic_min`/`atomic_max` (scalar) and `atomic_addx2`/`atomic_addx4`
(vectorized) typed their call with a non-void return under
`return_prev=True` but always selected the void `*_elem_op`, so codegen
emitted an empty binding (`prev = ;`) and nvcc failed.
Select the `*_ret_elem_op` variant under `return_prev`, as `atomic_add`
already does:
- min/max: the ret ops, codegen handlers, and device templates
(AtomicMaxRet/AtomicMinRet) already existed.
- addx2/addx4 (float32): register the ret ops (builtin.cc/.h) and add
codegen handlers emitting AtomicAddx2Ret/AtomicAddx4Ret (already in
the device headers) as expressions; the return is typed as the packed
2-/4-lane vector so the store matches.
- addx2 (fp16/bf16): the ret store LHS is `uint1`; codegen wraps the
half2/bf162 result with `tl::to_uint1`.
- addx4 (fp16/bf16): add AtomicAddx4Ret(half_t*)/(bfloat16_t*) device
overloads that run two AtomicAddx2Ret halves and return them packed as
uint2 (a fp16x4 is stored as uint2).
The 16-bit x4 path is per-pair atomic (two half2 atomics), not
whole-quad single-atomic -- hardware has no vector atomic wider than
`.f16x2` for 16-bit types. This matches the existing fp32-x4 scalar
fallback, and the docstring only promises "quad-width atomic addition".
Adds tests across scalar/x2/x4 and float32/fp16/bf16.
Fixes#2667
* minor fix for missing conditiions
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
`T.infinity` aborted with an internal `LOG(FATAL)` ("Cannot decide
infinity for type float8_e5m2") because `infinity_op` had no fp8 branch.
e5m2 is the only fp8 format with a representable inf, so add a branch
returning +inf, and emit it in codegen via its bit pattern
(`::bitcast(0x7c)`, mirroring the tfloat32 path) since inf/NaN cannot be
spelled as a float literal for the fp8 constructor. Other fp8 formats
(e4m3*, e3m4, e8m0, finite fnuz incl. e5m2fnuz) have no inf and keep the
fatal. Adds e5m2 to the infinity test.
Follow-up to the package export slimming: blockscaled_chunk_kmajor_word_offset
is no longer re-exported at the tilelang.quantize package level.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
nvfp4.py: remove the two remaining bring-up era quantize kernels (the
legacy no-TMA-load 16x256 schedule and the k64_per_cta=1 generic path,
reachable only through tuning knobs nothing uses) and the use_tma_load /
k64_per_cta parameters. The dispatcher now maps straight onto the single
tiled kernel factory; outputs verified bitwise-identical to the previous
code for every remaining configuration. 978 -> 710 lines (1334 before
the slimming series).
Example: drop the ones / constant-scale debug input modes and their CLI
flags; the random-FP4 + random-binary-scale path is the only
verification mode and checks exact equality.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Production shapes rarely keep M (batch x seq) a multiple of 128. Lift
the divisibility requirements to the validated envelope:
- scale packers zero-pad rows to full 128-row tiles (the same ceil
convention CUTLASS/cuBLAS/CuTeDSL use for the blocked SF tensor), so
the packed source always has ceil(rows/128)*128 rows;
- the example accepts an arbitrary M or an N tail: TMA loads zero-fill
out-of-bounds rows and the C store is predicated. N keeps the
16-byte-aligned-rows rule (N % 8 == 0), the same contiguous-dim
alignment CuTeDSL validates;
- simultaneous M and N tails are rejected with a clear error: the
fragment-to-global copy lowering emits the column guard bound from
the wrong dimension in that case (upstream bug, tracked separately).
Tests: tail-shape runtime correctness (257x384, 130x128, 128x136 with
exact binary-scale references), rejection paths, and packer padding
round-trip.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The example and the maint warp-specialized benchmark each carried a
private macro with the blockscaled_chunk_kmajor tile addressing spelled
out by hand. Move that knowledge into the library next to
T.mma_gemm_blockscaled:
- T.ue4m3_scale_tile_source_coords: the single source of the packed
scale tile addressing (tiles are contiguous in this layout);
words_per_kblock stays a parameter so future contracts
(block_K != 256, MX scale packings) reuse it.
- T.copy_ue4m3_scale_tile: cooperative whole-block staging macro with
trace-time dtype/shape validation, used by the example. The scale
operand stays visible in user kernels; only the layout math is hidden.
The maint kernel keeps its producer-warp-group staging schedule and
reuses the coords helper. Validated: nvf4 language + example CLI tests
(91 passed incl. a new validation test), example and WS benchmark
--verify pass.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merging the SM120 block-scale mode into TensorCoreIntrinEmitter changed
mma()'s 4th/5th positional parameters to SFA_buf/SFB_buf, so the plain
gemm lowering's positional mma(A_local, B_local, C_buf, ki) call bound
ki to SFA_buf and raised 'Scale buffers require TensorCoreIntrinEmitter
block-scale mode' for every non-blockscaled MMA GEMM (CI examples:
attention_sink, deepseek_nsa/mhc/v32).
Keep the base (A, B, C, k_inner) positional signature, make all
scale-related parameters keyword-only, forward k_inner to the base
implementation, and drop the unused blockscaled mma_atom override that
shadowed the base method the same way. Add a signature-contract
regression test.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Move the persistent CUTLASS comparison harness to maint/gemm and keep the public SM120 NVFP4 example focused on a non-persistent TileLang kernel using T.mma_gemm_blockscaled and normal T.copy output.
Add a source-level guard so the public example does not grow CUTLASS build or subprocess benchmark plumbing again.
Remove the public micro_pipeline strategy knob from T.mma_gemm_blockscaled and route the packed-scale SM120 path internally through the package-pingpong lowering. Hide SM120 strategy helpers from the public TIR API and remove stale emitter experiments that were only used by the strategy sweep.
Restore the optimized persistent SM120 NVFP4 example as the benchmark path, but keep the language surface clean by using sf_layout to select the packed scale contract and a regular T.copy epilogue instead of private C-fragment store APIs.
Validation: pre-commit run --files on changed files; pytest testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py -q => 89 passed; 512^3 correctness passed; 8192^3 128x128x256 stages=2 timing => 0.7198 ms, 1527.57 TFLOPS.
* [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
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>
* 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
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
* 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>
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
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
* [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>
* 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.
* 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>
* Refactor kernel launch materialization
* [Refactor] Improve git commit ID retrieval and version labeling logic
* Enhance the get_git_commit_id function to check for the existence of the .git directory before attempting to retrieve the commit hash, ensuring more robust behavior.
* Update the dynamic_metadata function to conditionally set the NO_VERSION_LABEL environment variable based on the presence of the .git directory.
* Modify the GitHub Actions workflow to unset the NO_VERSION_LABEL environment variable during the installation and version check of the package, allowing for accurate version labeling in SDist builds.
* [Refactor] Update WebGPU backend registration
* Import the WebGPU package to automatically register its pass pipeline.
* Remove the explicit registration of the WebGPU pass pipeline, streamlining the backend setup process.
* [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>
* Add TMA copy for packed global FP4 to unpacked SMEM.
Wire IsValidTMACopyDtypePair and ALIGN16B tensor map descriptors for
float4_e2m1fn global to float4_e2m1_unpacked shared transfers, and
extend language tests for TMA and SIMT copy restrictions.
* Extend tcgen05 MMA metadata for FP4 f8f6f4 and dtype helpers.
* Refactor tcgen05_gemm_blockscaled API and lowering.
* Use separate A/B dtypes in GEMM tileop emitters across backends.
* Add maint/gemm correctness evaluation scripts and update gemm_sm100 examples.
* Fix profiler assert_allclose to compare outputs only.
* lint
* canonicalize data types
* add dpgm fp8-fp4 kernel
* resolve comments
* Refine FP4 unpacked SMEM dtype handling
Treat float4_e2m1_unpacked as a SMEM/TMA-only storage tag instead of a general vector/register dtype. Remove vectorized Python aliases and keep codegen uses constrained to shared-memory storage/reference paths that lower to uint8_t.
Allow packed FP4 global tensors to TMA-load into unpacked shared memory with the ALIGN16B TensorMap format, while keeping unpacked storage out of TMA store paths. Use payload-bit accounting for global strides and mbarrier transactions.
Add focused dtype and TMA descriptor tests, and validate the real tcgen05 fp4-unpacked GEMM paths.
* Restrict FP4 unpack TMA to loads
* Hide vectorized unpacked FP4 language aliases
* Update TVM for unpacked FP4 symbol exports
* Rebind TVM submodule to merged unpacked FP4 export fix
* Enable cache for persistent tcgen5 GEMM example
* Add float4_e2m1_unpacked dtype
* Fix TIR cast for float4_e2m1_unpacked dtype.
Route custom[float4_e2m1_unpacked] through Float4E2M1Unpacked FFI
so T.float4_e2m1_unpacked(0.0) works in prim_func builders.
* Limit float4_e2m1_unpacked vector lanes and fix TMA transaction bytes.
Cap unpacked FP4 vector dtypes at x32 and count TMA/mbarrier
transactions using 4-bit payload size (0.5 byte per element).
* lint
* [Pipeline] Refactor software pipeline transforms
* enhance
* Refine pipeline handling for guarded binds
* Refactor IfStmtBindingRewriter and related classes for improved access analysis
- Removed unused Target parameter from IfStmtBindingRewriter.
- Introduced IfStmtAccessCollector to handle buffer access collection.
- Updated access analysis methods to utilize the new collector.
- Cleaned up includes and removed obsolete functions in access_analysis.
- Enhanced pipeline utility headers with necessary includes for consistency.
