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