mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Jayce Suanddingsg
779022eaf0
* [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>