mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Jayce Suanddingsg
e7da9025c4
* [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>
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
import pytest
|
|
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.cuda.target import normalize_cutedsl_target
|
|
|
|
|
|
def _lower_cutedsl(program):
|
|
if not tvm.runtime.enabled("cuda"):
|
|
pytest.skip("TileLang CuTeDSL codegen requires TVM built with CUDA support.")
|
|
|
|
build_cutedsl = tvm.ffi.get_global_func("target.build.tilelang_cutedsl_without_compile", allow_missing=True)
|
|
if build_cutedsl is None:
|
|
pytest.skip("TileLang CuTeDSL backend is not enabled in this build.")
|
|
|
|
target = normalize_cutedsl_target({"kind": "cutedsl", "arch": "sm_80"})
|
|
assert target is not None
|
|
|
|
with target:
|
|
return lower(program.with_attr("global_symbol", "main"), target=target)
|
|
|
|
|
|
def test_cutedsl_codegen_promotes_narrow_shift_before_wide_cast():
|
|
@T.prim_func
|
|
def prog(A: T.Tensor((1,), "uint16"), B: T.Tensor((1,), "uint32")):
|
|
with T.Kernel(1, threads=1):
|
|
local = T.alloc_local((1,), "uint16")
|
|
local[0] = A[0]
|
|
B[0] = T.cast(local[0] << 20, T.uint32)
|
|
|
|
artifact = _lower_cutedsl(prog)
|
|
|
|
assert "cutlass.Uint32(local[0]) << cutlass.Uint16(20)" in artifact.kernel_source
|
|
assert "cutlass.Uint32(cutlass.Uint16((local[0] << cutlass.Uint16(20))))" not in artifact.kernel_source
|
|
|
|
|
|
def test_cutedsl_codegen_does_not_promote_shift_across_signedness():
|
|
@T.prim_func
|
|
def prog(A: T.Tensor((1,), "int8"), B: T.Tensor((1,), "uint16")):
|
|
with T.Kernel(1, threads=1):
|
|
local = T.alloc_local((1,), "int8")
|
|
local[0] = A[0]
|
|
B[0] = T.cast(local[0] >> 1, T.uint16)
|
|
|
|
artifact = _lower_cutedsl(prog)
|
|
|
|
assert "cutlass.Uint16(local[0]) >>" not in artifact.kernel_source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|