[BugFix] Fix H100 CuTeDSL full-run edge cases

.gitignore

@@ -18,3 +18,7 @@ _version.py
 
 # workspace
 workspace
+
+.venv-*
+.test-logs
+analysis

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}"

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)

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
-
-

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

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
 

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()