From 1f3f3f0ce4411f1683f393fede60bf02b54ed32d Mon Sep 17 00:00:00 2001 From: JayceSu98 Date: Fri, 12 Jun 2026 06:29:40 +0000 Subject: [PATCH 1/4] [BugFix][Engram] Support non-divisible SM90 block counts The README declares SM90 GPUs as supported, but the Engram grad-w reduce kernel hard-coded four persistent batches and asserted that the persistent block count was divisible by four. That assumption is not true for all SM90 parts. For example, H100 PCIe exposes 114 SMs, so the CUDA/CuTeDSL full test path can fail before code generation reaches the actual kernel logic. Choose the largest batch count up to four that divides the persistent block count. This keeps the original four-batch layout when it is valid, while allowing other SM90 block counts to use the same persistent kernel shape without tripping the launch-time assertion. Co-authored-by: dingsg --- tile_kernels/engram/engram_grad_w_reduce_kernel.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tile_kernels/engram/engram_grad_w_reduce_kernel.py b/tile_kernels/engram/engram_grad_w_reduce_kernel.py index 5d13a29..ad0f1f3 100644 --- a/tile_kernels/engram/engram_grad_w_reduce_kernel.py +++ b/tile_kernels/engram/engram_grad_w_reduce_kernel.py @@ -1,5 +1,4 @@ import os - import torch import tilelang from tilelang import language as T @@ -21,8 +20,7 @@ def get_engram_grad_w_reduce_kernel( blk_d = 512 assert hidden_size % blk_d == 0 num_tiles = hidden_size // blk_d - num_batches = 4 - assert num_persistent_blocks % num_batches == 0 + num_batches = max(batch for batch in range(4, 0, -1) if num_persistent_blocks % batch == 0) num_rows = num_persistent_blocks // num_batches @T.prim_func From 07739b7d22a80bb4daf2f94cbf9a38bccd9cef57 Mon Sep 17 00:00:00 2001 From: JayceSu98 Date: Fri, 12 Jun 2026 06:30:01 +0000 Subject: [PATCH 2/4] [BugFix][MoE] Guard fused mapping edge tiles Full H100/CuTeDSL coverage exercises MoE routing shapes where the generated work loops are rounded up to alignment or warp-size boundaries. The logical tensors remain sized by the real hidden dimension, scaling-factor dimension, and token-topk count. The expand-to-fused kernel previously iterated over aligned hidden and scaling-factor extents and wrote every lane into the destination tensors. For tail dimensions this can address columns outside the real tensor shape, especially with TMA-aligned column-major scaling-factor layouts. The fused mapping kernel also used a Select expression whose inactive arm could still expose an out-of-range topk_idx load to code generation, and its fallback allocation estimate rounded the per-expert alignment upper bound down to an alignment multiple. Guard aligned expand stores with the real hidden/scaling-factor extents, initialize the shared cumsum buffer for all threads, avoid the speculative out-of-range topk load, and reserve a rounded-up per-expert alignment upper bound before trimming with the synchronized expert counts. Co-authored-by: dingsg --- tile_kernels/moe/expand_to_fused_kernel.py | 24 ++++++++++++-------- tile_kernels/moe/get_fused_mapping_kernel.py | 9 ++++++-- 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/tile_kernels/moe/expand_to_fused_kernel.py b/tile_kernels/moe/expand_to_fused_kernel.py index 92a18c4..87bee96 100644 --- a/tile_kernels/moe/expand_to_fused_kernel.py +++ b/tile_kernels/moe/expand_to_fused_kernel.py @@ -56,13 +56,15 @@ def expand_to_fused_kernel( if pid_token < num_expanded_tokens: if pos_to_expert[pid_token] < 0: for i in T.Parallel(hidden_aligned): - expanded_x[pid_token, i] = 0 + if i < hidden: + expanded_x[pid_token, i] = 0 if num_per_channels is not None: for i in T.Parallel(hidden_sf_aligned): - if use_tma_aligned_col_major_sf: - expanded_x_sf[i, pid_token] = 0 - else: - expanded_x_sf[pid_token, i] = 0 + if i < hidden_sf: + if use_tma_aligned_col_major_sf: + expanded_x_sf[i, pid_token] = 0 + else: + expanded_x_sf[pid_token, i] = 0 if pid_token >= num_tokens: T.thread_return() @@ -80,13 +82,15 @@ def expand_to_fused_kernel( T.assume(pos_local[k] < num_expanded_tokens) if pos_local[k] >= 0: for i in T.Parallel(hidden_aligned): - expanded_x[pos_local[k], i] = x_fragment[i] + if i < hidden: + expanded_x[pos_local[k], i] = x_fragment[i] if num_per_channels is not None: for i in T.Parallel(hidden_sf_aligned): - if use_tma_aligned_col_major_sf: - expanded_x_sf[i, pos_local[k]] = x_sf_fragment[i] - else: - expanded_x_sf[pos_local[k], i] = x_sf_fragment[i] + if i < hidden_sf: + if use_tma_aligned_col_major_sf: + expanded_x_sf[i, pos_local[k]] = x_sf_fragment[i] + else: + expanded_x_sf[pos_local[k], i] = x_sf_fragment[i] return expand_to_fused_kernel diff --git a/tile_kernels/moe/get_fused_mapping_kernel.py b/tile_kernels/moe/get_fused_mapping_kernel.py index 998189e..a4457cb 100644 --- a/tile_kernels/moe/get_fused_mapping_kernel.py +++ b/tile_kernels/moe/get_fused_mapping_kernel.py @@ -104,6 +104,7 @@ def get_fused_mapping_kernel( T.sync_grid() cumsum_shared = T.alloc_shared((num_threads,), T.int32) + cumsum_shared[thread_idx] = 0 expert_num_elements = T.alloc_var(T.int32, init=0) expert_num_elements_aligned = T.alloc_var(T.int32, init=0) prefix_expert_num_elements = T.alloc_var(T.int32, init=0) @@ -143,7 +144,9 @@ def get_fused_mapping_kernel( lane_mask_rev = ~lane_mask for i in T.serial(start + lane_idx, aligned_end, warp_size): T.assume(0 <= i) - expert_idx = T.Select(i < numel, T.int32(topk_idx_1d[i]), -1) + expert_idx = T.alloc_var(T.int32, init=-1) + if i < numel: + expert_idx = T.int32(topk_idx_1d[i]) mask = T.call_extern(T.uint32, '__match_any_sync', 0xFFFFFFFF, expert_idx) count = T.popcount(mask & lane_mask) @@ -201,7 +204,9 @@ def get_fused_mapping( should_sync = False if num_expanded_tokens == 0 and not force_no_sync: should_sync = True - num_expanded_tokens = (num_tokens * num_topk + (alignment - 1) * num_experts) // alignment * alignment + # Each expert range is aligned independently, so reserve the rounded-up + # upper bound before trimming with num_tokens_per_expert. + num_expanded_tokens = align(num_tokens * num_topk + (alignment - 1) * num_experts, alignment) # Allocate output num_sms = get_num_sms() From d01cf19dd81d7a839b80a9b8bf8fda316df46048 Mon Sep 17 00:00:00 2001 From: JayceSu98 Date: Fri, 12 Jun 2026 06:30:17 +0000 Subject: [PATCH 3/4] [BugFix][Quant] Guard scale tails and FP32 absmax The full H100/CuTeDSL correctness run covers token and hidden sizes that do not land exactly on the kernel tile, scaling-factor block, or E5M6 packing granularity. Those edge tiles exposed two independent quantization issues. First, cast-back and lossless/per-token cast kernels could load or store scaling factors for padded token/channel blocks. The same edge condition also applied to packed E5M6 output rows and columns, where the final tile may not contain a full logical token row or hidden group. Second, per-token scale generation used in_config.dtype as the absmax fragment dtype. That dtype describes the input tensor storage, not the reduction accumulator. For BF16 and narrow quantization paths, keeping absmax in the input dtype can round or overflow the scale path before the final output cast, producing byte mismatches against the PyTorch reference. Guard scale loads/stores and packed E5M6 stores with the real token/channel extents, keep invalid scale lanes zero-initialized, and use FP32 fragments for per-token absmax reductions before computing the reciprocal scale. Co-authored-by: dingsg --- tile_kernels/quant/cast_back_kernel.py | 9 +++++++-- .../quant/per_block_cast_lossless_kernel.py | 18 ++++++++++++------ tile_kernels/quant/per_token_cast_kernel.py | 8 +++++--- .../quant/per_token_cast_to_e5m6_kernel.py | 12 ++++++++---- 4 files changed, 32 insertions(+), 15 deletions(-) diff --git a/tile_kernels/quant/cast_back_kernel.py b/tile_kernels/quant/cast_back_kernel.py index 69d6a6f..36c47af 100644 --- a/tile_kernels/quant/cast_back_kernel.py +++ b/tile_kernels/quant/cast_back_kernel.py @@ -60,11 +60,16 @@ def cast_back_kernel( out_fragment = T.alloc_fragment((TILE_M, TILE_K), out_dtype) T.copy(x[pid_token * TILE_M, pid_hidden * TILE_K], x_shared, disable_tma=True) + num_sf_token_blocks = T.ceildiv(num_tokens, num_per_tokens) + num_sf_channel_blocks = T.ceildiv(hidden, num_per_channels) for i, j in T.Parallel(T.ceildiv(TILE_M, num_per_tokens), T.ceildiv(TILE_K, num_per_channels)): token_index = pid_token * TILE_M // num_per_tokens + i channel_index = pid_hidden * TILE_K // num_per_channels + j - sf = load_sf(x_sf, token_index, channel_index, in_config) - sf_shared[i, j] = transform_sf(sf, in_config) + if token_index < num_sf_token_blocks and channel_index < num_sf_channel_blocks: + sf = load_sf(x_sf, token_index, channel_index, in_config) + sf_shared[i, j] = transform_sf(sf, in_config) + else: + sf_shared[i, j] = 0 for i, j in T.Parallel(TILE_M, TILE_K): out_fragment[i, j] = x_shared[i, j] * sf_shared[i // num_per_tokens, j // num_per_channels] diff --git a/tile_kernels/quant/per_block_cast_lossless_kernel.py b/tile_kernels/quant/per_block_cast_lossless_kernel.py index 9645aef..84f9c70 100644 --- a/tile_kernels/quant/per_block_cast_lossless_kernel.py +++ b/tile_kernels/quant/per_block_cast_lossless_kernel.py @@ -88,10 +88,13 @@ def per_block_cast_lossless_kernel( # Load scaling factor of x to fragment T.fill(x_sf_fragment, 0) + num_in_sf_blocks_m = T.ceildiv(num_tokens, in_config.sf_block[0]) + num_in_sf_blocks_k = T.ceildiv(hidden, in_config.sf_block[1]) for i, j in T.Parallel(num_in_sf_per_block_m, num_in_sf_per_block_k): m_idx = pid_token * block_m // in_config.sf_block[0] + i k_idx = pid_hidden * block_k // in_config.sf_block[1] + j - x_sf_fragment[i, j] = load_sf(x_sf, m_idx, k_idx, in_config) + if m_idx < num_in_sf_blocks_m and k_idx < num_in_sf_blocks_k: + x_sf_fragment[i, j] = load_sf(x_sf, m_idx, k_idx, in_config) # Alloc fragments x_sf_uint32_fragment = T.alloc_fragment((num_in_sf_per_block_m, num_in_sf_per_block_k), T.uint32) @@ -137,14 +140,17 @@ def per_block_cast_lossless_kernel( x_out_fragment[i, j] = T.cast(T.float32(x_in_shared[i, j]) * sf, out_config.dtype) # Store scaling factor back to global memory + num_out_sf_blocks_m = T.ceildiv(num_tokens, out_config.sf_block[0]) + num_out_sf_blocks_k = T.ceildiv(hidden, out_config.sf_block[1]) for i, j in T.Parallel(num_out_sf_per_block_m, num_out_sf_per_block_k): sf_m_idx = pid_token * num_out_sf_per_block_m + i sf_k_idx = pid_hidden * num_out_sf_per_block_k + j - if out_config.use_packed_ue8m0: - sf = T.uint8(out_sf_uint32_fragment[i, j]) - else: - sf = transform_sf_to_fp32(out_sf_uint32_fragment[i, j]) - store_sf(out_sf, sf, sf_m_idx, sf_k_idx, out_config) + if sf_m_idx < num_out_sf_blocks_m and sf_k_idx < num_out_sf_blocks_k: + if out_config.use_packed_ue8m0: + sf = T.uint8(out_sf_uint32_fragment[i, j]) + else: + sf = transform_sf_to_fp32(out_sf_uint32_fragment[i, j]) + store_sf(out_sf, sf, sf_m_idx, sf_k_idx, out_config) T.copy(x_out_fragment, out[pid_token * block_m: (pid_token + 1) * block_m, pid_hidden * block_k: (pid_hidden + 1) * block_k]) diff --git a/tile_kernels/quant/per_token_cast_kernel.py b/tile_kernels/quant/per_token_cast_kernel.py index 1678648..7234b7f 100644 --- a/tile_kernels/quant/per_token_cast_kernel.py +++ b/tile_kernels/quant/per_token_cast_kernel.py @@ -115,7 +115,8 @@ def per_token_cast_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv # Store casted values @@ -131,7 +132,7 @@ def per_token_cast_kernel( sf = load_sf(out_sf, pid_token * block_m + i, pid_hidden * num_groups + j, out_config) sf_inv_fragment[i, j] = 1 / sf else: - amax_fragment = T.alloc_fragment((block_m, num_groups), in_config.dtype) + amax_fragment = T.alloc_fragment((block_m, num_groups), T.float32) x_fragment_reshaped = T.reshape(x_fragment, [block_m, num_groups, num_per_channels]) # Reduce SF T.reduce_absmax(x_fragment_reshaped, amax_fragment, dim=2) @@ -142,7 +143,8 @@ def per_token_cast_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv # Store casted values diff --git a/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py b/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py index 0956c13..01be08a 100644 --- a/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py +++ b/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py @@ -119,7 +119,7 @@ def per_token_cast_to_e5m6_kernel( # Copy input into registers T.copy(x[pid_token * block_m, pid_hidden * block_k], x_fragment) - amax_fragment = T.alloc_fragment((block_m, num_groups), in_config.dtype) + amax_fragment = T.alloc_fragment((block_m, num_groups), T.float32) x_fragment_reshaped = T.reshape(x_fragment, [block_m, num_groups, num_per_channels]) # Reduce SF T.reduce_absmax(x_fragment_reshaped, amax_fragment, dim=2) @@ -130,7 +130,8 @@ def per_token_cast_to_e5m6_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv T.annotate_layout({ @@ -150,8 +151,11 @@ def per_token_cast_to_e5m6_kernel( for j in T.serial(8): in_local[j] = out_fragment[x, y * 8 + j] float_to_e5m6(in_local, out_local) - for j in T.serial(3): - out[pid_token * block_m + x, pid_hidden * (block_k // 8 * 3) + y * 3 + j] = out_local[j] + m_idx = pid_token * block_m + x + k_idx = pid_hidden * block_k + y * 8 + if m_idx < num_tokens and k_idx < hidden: + for j in T.serial(3): + out[m_idx, (k_idx // 8) * 3 + j] = out_local[j] return per_token_cast_to_e5m6_kernel From b39520e32c613871d9c19735b61b523ffb9a2618 Mon Sep 17 00:00:00 2001 From: JayceSu98 Date: Fri, 12 Jun 2026 06:30:42 +0000 Subject: [PATCH 4/4] [Test] Add reusable CuTeDSL full-run controls Full H100/CuTeDSL validation needs to run correctness and benchmark coverage under xdist while preserving host GPU pinning, separating TileLang compile caches, and collecting enough per-worker information to debug late failures without re-running the whole suite blindly. Keep externally provided CUDA_VISIBLE_DEVICES instead of overwriting it in every xdist worker, allow per-worker TileLang cache directories, add optional per-test CUDA synchronization for benchmark/debug runs, and write optional trace/failure reports controlled by environment variables. Benchmark-only sweeps also need a way to record new results before a baseline file exists. Add an opt-in missing-baseline override and make the timer backend configurable so local H100 runs can use CUDA events when CUPTI is not the desired backend. Add a small run_cutedsl_tests.sh wrapper that rebuilds a local TileLang checkout, selects correctness/benchmark/all modes, stores logs and JSONL output under .test-logs, and leaves local virtualenv/log directories ignored. Co-authored-by: dingsg --- .gitignore | 4 + run_cutedsl_tests.sh | 181 ++++++++++++++++++++++++++++++ tests/moe/test_expand_to_fused.py | 2 + tests/pytest_benchmark_plugin.py | 55 +++++++-- 4 files changed, 235 insertions(+), 7 deletions(-) create mode 100755 run_cutedsl_tests.sh diff --git a/.gitignore b/.gitignore index 705d765..13410ba 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,7 @@ _version.py # workspace workspace + +.venv-* +.test-logs +analysis diff --git a/run_cutedsl_tests.sh b/run_cutedsl_tests.sh new file mode 100755 index 0000000..c7e8fff --- /dev/null +++ b/run_cutedsl_tests.sh @@ -0,0 +1,181 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: + ./run_cutedsl_tests.sh [all|correctness|benchmark] [options] [-- extra pytest args] + +Modes: + all Run correctness + benchmark together. This is the default. + correctness Run correctness only. + benchmark Run benchmark cases only. + +Options: + -n, --workers N Pytest xdist worker count. Default: 2 + --tilelang-dir PATH Local TileLang checkout. Default: $TILELANG_DIR or ../tilelang + --python PATH Python executable. Default: ./.venv-tk-test/bin/python + -h, --help Show this help + +Environment overrides: + TK_CUDA_VISIBLE_DEVICES Visible GPU list for this host. Default: $CUDA_VISIBLE_DEVICES or 0 + TK_BENCHMARK_BACKEND Benchmark timer backend. Default: event + TK_FULL_TEST Full parameter coverage flag. Default: 1 + TILELANG_DISABLE_CACHE TileLang cache toggle. Default: 1 + TK_BUILD_JOBS TileLang build jobs. Default: 16 + TK_SKIP_BUILD Skip `cmake --build` when set to 1 + +Examples: + ./run_cutedsl_tests.sh + ./run_cutedsl_tests.sh correctness + ./run_cutedsl_tests.sh benchmark -n 2 + ./run_cutedsl_tests.sh all -- tests/engram/test_engram_fused_weight.py +EOF +} + +MODE="all" +WORKERS=2 +TILELANG_DIR="${TILELANG_DIR:-}" +PYTHON_BIN="" +EXTRA_PYTEST_ARGS=() + +while (($#)); do + case "$1" in + all|correctness|benchmark) + MODE="$1" + shift + ;; + -n|--workers) + WORKERS="$2" + shift 2 + ;; + --tilelang-dir) + TILELANG_DIR="$2" + shift 2 + ;; + --python) + PYTHON_BIN="$2" + shift 2 + ;; + --) + shift + EXTRA_PYTEST_ARGS+=("$@") + break + ;; + -h|--help) + usage + exit 0 + ;; + *) + EXTRA_PYTEST_ARGS+=("$1") + shift + ;; + esac +done + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -z "${PYTHON_BIN}" ]]; then + PYTHON_BIN="${ROOT_DIR}/.venv-tk-test/bin/python" +fi +if [[ -z "${TILELANG_DIR}" ]]; then + TILELANG_DIR="${ROOT_DIR}/../tilelang" +fi + +if [[ ! -x "${PYTHON_BIN}" ]]; then + echo "[error] Python executable not found: ${PYTHON_BIN}" >&2 + exit 1 +fi +if [[ ! -d "${TILELANG_DIR}" ]]; then + echo "[error] TileLang directory not found: ${TILELANG_DIR}" >&2 + exit 1 +fi +TILELANG_DIR="$(cd "${TILELANG_DIR}" && pwd)" + +CUDA_VISIBLE_DEVICES_VALUE="${TK_CUDA_VISIBLE_DEVICES:-${CUDA_VISIBLE_DEVICES:-0}}" +BENCHMARK_BACKEND_VALUE="${TK_BENCHMARK_BACKEND:-event}" +FULL_TEST_VALUE="${TK_FULL_TEST:-1}" +DISABLE_CACHE_VALUE="${TILELANG_DISABLE_CACHE:-1}" +BUILD_JOBS_VALUE="${TK_BUILD_JOBS:-16}" +SKIP_BUILD_VALUE="${TK_SKIP_BUILD:-0}" +STAMP="$(date -u +%Y%m%d_%H%M%S)" + +mkdir -p "${ROOT_DIR}/.test-logs" + +RUN_LABEL="${MODE}" +LOG_PREFIX="${ROOT_DIR}/.test-logs/${RUN_LABEL}_cutedsl_n${WORKERS}_${BENCHMARK_BACKEND_VALUE}_${STAMP}" +LOG_PATH="${LOG_PREFIX}.log" +JSONL_PATH="${LOG_PREFIX}.jsonl" +FAILURE_PREFIX="${LOG_PREFIX}.failure" +TRACE_PREFIX="${LOG_PREFIX}.trace" +CACHE_DIR="${LOG_PREFIX}.cache" + +export PYTHONPATH="${TILELANG_DIR}${PYTHONPATH:+:${PYTHONPATH}}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES_VALUE}" +export TILELANG_TARGET="cutedsl" +export TILELANG_DISABLE_CACHE="${DISABLE_CACHE_VALUE}" +export TK_FULL_TEST="${FULL_TEST_VALUE}" +export TILELANG_CACHE_DIR="${CACHE_DIR}" +export TK_TILELANG_CACHE_PER_WORKER="1" +export TK_BENCHMARK_FAILURE_REPORT="${FAILURE_PREFIX}" +export TK_BENCHMARK_TRACE_REPORT="${TRACE_PREFIX}" + +PYTEST_ARGS=( + tests + -n "${WORKERS}" + --tb=short + -ra +) + +case "${MODE}" in + all) + export TK_BENCHMARK_BACKEND="${BENCHMARK_BACKEND_VALUE}" + export TK_BENCHMARK_ALLOW_MISSING_BASELINES="1" + PYTEST_ARGS+=( + --run-benchmark + "--benchmark-output=${JSONL_PATH}" + ) + ;; + correctness) + ;; + benchmark) + export TK_BENCHMARK_BACKEND="${BENCHMARK_BACKEND_VALUE}" + export TK_BENCHMARK_ALLOW_MISSING_BASELINES="1" + PYTEST_ARGS+=( + --run-benchmark + -m benchmark + "--benchmark-output=${JSONL_PATH}" + ) + ;; +esac + +if ((${#EXTRA_PYTEST_ARGS[@]})); then + PYTEST_ARGS+=("${EXTRA_PYTEST_ARGS[@]}") +fi + +echo "[run] mode=${MODE}" +echo "[run] workers=${WORKERS}" +echo "[run] tilelang_dir=${TILELANG_DIR}" +echo "[run] python=${PYTHON_BIN}" +echo "[run] CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "[run] log=${LOG_PATH}" +if [[ "${MODE}" != "correctness" ]]; then + echo "[run] benchmark_jsonl=${JSONL_PATH}" +fi + +"${PYTHON_BIN}" - <<'PY' +import os +import torch + +print(f"[env] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}") +print(f"[env] torch.cuda.device_count()={torch.cuda.device_count()}") +for idx in range(torch.cuda.device_count()): + print(f"[env] cuda:{idx} -> {torch.cuda.get_device_name(idx)}") +PY + +if [[ "${SKIP_BUILD_VALUE}" != "1" ]]; then + echo "[build] cmake --build ${TILELANG_DIR}/build -j ${BUILD_JOBS_VALUE}" + cmake --build "${TILELANG_DIR}/build" -j "${BUILD_JOBS_VALUE}" +fi + +echo "[pytest] ${PYTHON_BIN} -m pytest ${PYTEST_ARGS[*]}" +"${PYTHON_BIN}" -m pytest "${PYTEST_ARGS[@]}" 2>&1 | tee "${LOG_PATH}" diff --git a/tests/moe/test_expand_to_fused.py b/tests/moe/test_expand_to_fused.py index d8e5b08..f5d31aa 100644 --- a/tests/moe/test_expand_to_fused.py +++ b/tests/moe/test_expand_to_fused.py @@ -85,6 +85,8 @@ def generate_test_data_expand_to_fused_with_sf(params): round_sf=round_sf, use_packed_ue8m0=use_packed_ue8m0, ) + if os.environ.get('TK_SYNC_AFTER_EXPAND_WITH_SF_CAST') == '1': + torch.cuda.synchronize() pos_to_expert, _, _, token_topk_to_pos, _, _, _, _ = tile_kernels.moe.get_fused_mapping(topk_idx, num_experts, 0, 16) return (x_fp8, x_sf, pos_to_expert, token_topk_to_pos, num_tokens) diff --git a/tests/pytest_benchmark_plugin.py b/tests/pytest_benchmark_plugin.py index d4cdc91..b0b7c54 100644 --- a/tests/pytest_benchmark_plugin.py +++ b/tests/pytest_benchmark_plugin.py @@ -68,14 +68,22 @@ def pytest_configure(config): if worker_id is not None: gpu_id = int(worker_id.replace('gw', '')) num_gpus = torch.cuda.device_count() - os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id % num_gpus) + pinned_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES') + if not pinned_visible_devices: + os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id % num_gpus) + if os.environ.get('TK_TILELANG_CACHE_PER_WORKER') == '1': + cache_dir = os.environ.get('TILELANG_CACHE_DIR') + if cache_dir: + worker_cache_dir = os.path.join(cache_dir, worker_id) + os.makedirs(worker_cache_dir, exist_ok=True) + os.environ['TILELANG_CACHE_DIR'] = worker_cache_dir # Restrict each worker's GPU memory to (total - 10 GB) / workers_per_gpu. # PYTEST_XDIST_WORKER_COUNT is set by pytest-xdist automatically. total_workers = int(os.environ.get('PYTEST_XDIST_WORKER_COUNT', '1')) workers_per_gpu = math.ceil(total_workers / num_gpus) _reserve_bytes = 10 * (1024 ** 3) # 10 GB reserved for system / frameworks - total_mem = torch.cuda.mem_get_info(0)[1] + total_mem = torch.cuda.get_device_properties(0).total_memory usable_mem = max(total_mem - _reserve_bytes, 0) mem_per_worker = usable_mem / workers_per_gpu fraction = mem_per_worker / total_mem @@ -103,6 +111,40 @@ def pytest_collection_modifyitems(config, items): # Use `-m benchmark` explicitly if you want ONLY benchmarks. +@pytest.fixture(autouse=True) +def _sync_after_benchmark(request): + yield + if ( + os.environ.get('TK_BENCHMARK_SYNC_AFTER_TEST') == '1' + and 'benchmark' in request.node.keywords + and torch.cuda.is_available() + ): + torch.cuda.synchronize() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + report = outcome.get_result() + worker_id = os.environ.get('PYTEST_XDIST_WORKER', 'main') + trace_prefix = os.environ.get('TK_BENCHMARK_TRACE_REPORT') + if trace_prefix and report.when == 'call': + trace_path = f'{trace_prefix}.{worker_id}.log' + with open(trace_path, 'a') as f: + f.write(f'{report.outcome}: {report.nodeid}\n') + + output_prefix = os.environ.get('TK_BENCHMARK_FAILURE_REPORT') + if not output_prefix or report.when != 'call' or not report.failed: + return + + output_path = f'{output_prefix}.{worker_id}.log' + with open(output_path, 'a') as f: + f.write(f'nodeid: {report.nodeid}\n') + f.write(f'worker: {worker_id}\n') + f.write(str(report.longrepr)) + f.write('\n\n') + + # --------------------------------------------------------------------------- # Regression detection & exit code @@ -159,7 +201,8 @@ def pytest_sessionfinish(session, exitstatus): results, baselines, regressions, improvements, missing = result # Stash for pytest_terminal_summary session.config._benchmark_detection = result - if (regressions or missing) and exitstatus == 0: + allow_missing = os.environ.get('TK_BENCHMARK_ALLOW_MISSING_BASELINES') == '1' + if (regressions or (missing and not allow_missing)) and exitstatus == 0: session.exitstatus = 1 @@ -421,7 +464,6 @@ def _record(*, kernel, operation, params, time_us, bandwidth_gbs=None, extras=No with request.config._benchmark_results_lock: request.config._benchmark_results.append(record) - return _record @@ -440,7 +482,8 @@ def benchmark_timer(): from tilelang.profiler.bench import do_bench def _timer(fn, **overrides): - kwargs = dict(backend='cupti', warmup=0, rep=30) + backend = os.environ.get('TK_BENCHMARK_BACKEND', 'cupti') + kwargs = dict(backend=backend, warmup=0, rep=30) kwargs.update(overrides) return do_bench(fn, **kwargs) * 1e3 # ms → us @@ -473,5 +516,3 @@ def _load_baselines(): rec = json.loads(line) baselines[_make_key(rec)] = rec return baselines - -