diff --git a/.github/workflows/evt_test.yml b/.github/workflows/evt_test.yml index 2ad4183..2794c38 100644 --- a/.github/workflows/evt_test.yml +++ b/.github/workflows/evt_test.yml @@ -1,6 +1,6 @@ name: EVT CUTLASS Test -# CUTLASS EVT epilogue-fusion tests need ~35 min of nvcc compilation +# CUTLASS EVT epilogue-fusion tests need ~17 min of nvcc compilation # per PyTorch version. Run manually from the Actions tab when needed. # # Usage: Actions -> EVT CUTLASS Test -> Run workflow diff --git a/magi_compiler/passes/piecewise_graph/fusion/common/codegen_shared.py b/magi_compiler/passes/piecewise_graph/fusion/common/codegen_shared.py index b572fb6..c3d9c3d 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/common/codegen_shared.py +++ b/magi_compiler/passes/piecewise_graph/fusion/common/codegen_shared.py @@ -16,7 +16,38 @@ from __future__ import annotations +import os import textwrap +from typing import Sequence, TypeVar + +_T = TypeVar("_T") + + +def max_tile_candidates() -> int | None: + """Return ``MAGI_EVT_MAX_TILES`` when set to a positive int, else ``None``. + + Tests set this to 1 so nvcc instantiates a single tile instead of the + full autotune list. Unset in production (unlimited). + """ + raw = os.environ.get("MAGI_EVT_MAX_TILES") + if raw is None or not str(raw).strip(): + return None + try: + n = int(raw) + except ValueError as exc: + raise ValueError(f"MAGI_EVT_MAX_TILES must be a positive int, got {raw!r}") from exc + if n < 1: + raise ValueError(f"MAGI_EVT_MAX_TILES must be >= 1, got {n}") + return n + + +def limit_tile_candidates(candidates: Sequence[_T]) -> list[_T]: + """Slice ``candidates`` to ``MAGI_EVT_MAX_TILES`` when that env is set.""" + n = max_tile_candidates() + if n is None: + return list(candidates) + return list(candidates)[:n] + _DTYPE_TO_CUTLASS = {"bfloat16": "cutlass::bfloat16_t", "float16": "cutlass::half_t", "float32": "float"} diff --git a/magi_compiler/passes/piecewise_graph/fusion/evt_runtime.py b/magi_compiler/passes/piecewise_graph/fusion/evt_runtime.py index 7ecf6b1..f6eb096 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/evt_runtime.py +++ b/magi_compiler/passes/piecewise_graph/fusion/evt_runtime.py @@ -43,6 +43,7 @@ from magi_compiler.config import get_compile_config +from .common.codegen_shared import max_tile_candidates from .evt_ir import Accum, AuxLoad, ColBroadcast, Compute, RowBroadcast, Store from .sm80.evt_codegen import render_evt_cu as _render_evt_cu_sm80 from .sm90.evt_codegen import render_evt_cu as _render_evt_cu_sm90 @@ -288,29 +289,29 @@ def _compile_evt_module( b_dtype: torch.dtype, b_layout: str = "row", m_bucket: str = "medium", - N: int = 0, - K: int = 0, alignment_a_bits: int = 128, alignment_b_bits: int = 128, alignment_c_bits: int = 128, ): """Render + JIT-compile the EVT kernel for ``ir_json``. Process-level cached. - Each distinct (N, K) gets its own module so autotune state is isolated. + (N, K) are launch-time problem sizes and do not change the generated .cu; + the C++ runner caches autotune winners in a per-(N, K) map. Alignment + bits stay in the key because they are baked into the CUTLASS typedefs. """ arch = _device_arch_tag() + tiles = max_tile_candidates() fast_key = ( ir_json, a_dtype, b_dtype, b_layout, m_bucket, - N, - K, alignment_a_bits, alignment_b_bits, alignment_c_bits, arch, + tiles, ) cached = _MODULE_FAST_CACHE.get(fast_key) if cached is not None: @@ -320,23 +321,21 @@ def _compile_evt_module( raise ValueError(f"b_layout must be 'row' or 'col', got {b_layout!r}") a_str = _DTYPE_TO_STR[a_dtype] b_str = _DTYPE_TO_STR[b_dtype] - extended = json.dumps( - { - "ir": ir_json, - "a": a_str, - "b": b_str, - "b_layout": b_layout, - "m_bucket": m_bucket, - "N": int(N), - "K": int(K), - "alignA_bits": int(alignment_a_bits), - "alignB_bits": int(alignment_b_bits), - "alignC_bits": int(alignment_c_bits), - "arch": arch, - "version": 10, - }, - sort_keys=True, - ).encode("utf-8") + payload = { + "ir": ir_json, + "a": a_str, + "b": b_str, + "b_layout": b_layout, + "m_bucket": m_bucket, + "alignA_bits": int(alignment_a_bits), + "alignB_bits": int(alignment_b_bits), + "alignC_bits": int(alignment_c_bits), + "arch": arch, + "version": 11, + } + if tiles is not None: + payload["max_tiles"] = tiles + extended = json.dumps(payload, sort_keys=True).encode("utf-8") key = hashlib.sha256(extended).hexdigest() cached = _MODULE_CACHE.get(key) @@ -476,17 +475,16 @@ def _node_from_dict(d): raise ValueError(f"Unknown IR kind {kind!r}") -# Per-(m_bucket, N, K, align) cache — separate modules so each runner has its -# own autotune state (best_idx_). +# Per-(m_bucket, align) cache. (N, K) are launch-time sizes; the C++ runner +# stores autotune winners in a per-(N_out, K) map inside one shared module. _SWIGLU_FAST_CACHE: dict = {} _SWIGLU_BUILD_LOCKS: dict = {} -def _compile_swiglu_dual( - m_bucket: str, N: int, K: int, alignment_a_bits: int = 128, alignment_b_bits: int = 128, alignment_c_bits: int = 128 -): - """Lazy-load a per-(bucket, N, K, align) DualGemm kernel module.""" - fast_key = (m_bucket, int(N), int(K), int(alignment_a_bits), int(alignment_b_bits), int(alignment_c_bits)) +def _compile_swiglu_dual(m_bucket: str, alignment_a_bits: int = 128, alignment_b_bits: int = 128, alignment_c_bits: int = 128): + """Lazy-load a per-(bucket, align) DualGemm kernel module.""" + tiles = max_tile_candidates() + fast_key = (m_bucket, int(alignment_a_bits), int(alignment_b_bits), int(alignment_c_bits), tiles) cached = _SWIGLU_FAST_CACHE.get(fast_key) if cached is not None: return cached @@ -510,9 +508,10 @@ def _compile_swiglu_dual( if not os.path.exists(src): raise FileNotFoundError(f"vendored swiglu source not found: {src}") cache_root = get_compile_config().cache_root_dir - # Build dir embeds (arch, bucket, N, K, align) — stale cross-arch - # binaries cause cudaErrorInvalidDeviceFunction. - build_tag = f"{m_bucket}_N{N}_K{K}" f"_aA{alignment_a_bits}_aB{alignment_b_bits}_aC{alignment_c_bits}" + # Build dir embeds (arch, bucket, align, optional tile cap). + # Stale cross-arch binaries cause cudaErrorInvalidDeviceFunction. + tiles_tag = f"_tiles{tiles}" if tiles is not None else "" + build_tag = f"{m_bucket}_aA{alignment_a_bits}_aB{alignment_b_bits}_aC{alignment_c_bits}{tiles_tag}" build_dir = os.path.join(cache_root, "evt_kernels", arch_tag, f"swiglu_dual_{build_tag}") os.makedirs(build_dir, exist_ok=True) mod_name = f"magi_swiglu_dual_{build_tag}" @@ -557,6 +556,7 @@ def _compile_swiglu_dual( f"-DMAGI_SWIGLU_ALIGN_A_BITS={int(alignment_a_bits)}", f"-DMAGI_SWIGLU_ALIGN_B_BITS={int(alignment_b_bits)}", f"-DMAGI_SWIGLU_ALIGN_C_BITS={int(alignment_c_bits)}", + *([f"-DMAGI_EVT_MAX_TILES={int(tiles)}"] if tiles is not None else []), ], build_directory=build_dir, verbose=False, @@ -569,8 +569,9 @@ def _compile_swiglu_dual( # ── Dispatch fast-cache ────────────────────────────────────────────────────── # Collapses out_dtype_from_id → _m_bucket → _compile_* → mod.attr-lookup -# into a single dict.get(). Keyed by (kind, ir_json, dtypes, N, K, m_bucket, -# out_dtype); reaches steady state after the first call per (site, bucket). +# into a single dict.get(). Keyed by (kind, ir_json, dtypes, B sizes, +# m_bucket, out_dtype) so alignment is still derived per shape; the +# compiled .so is shared across (N, K) that hash to the same align. class _DispatchEntry: __slots__ = ("kernel_call", "is_evt", "out_dtype") @@ -593,7 +594,7 @@ def _resolve_dispatch(kind, ir_json, a_dtype, b_dtype, N_w, K_w, m_bucket, out_d # K alignment also covers ldB=2K. align_bits = _runtime_align_bits(K_w, a_dtype) mod = _compile_swiglu_dual( - m_bucket, N_w, K_w, alignment_a_bits=align_bits, alignment_b_bits=align_bits, alignment_c_bits=alignment_c_bits + m_bucket, alignment_a_bits=align_bits, alignment_b_bits=align_bits, alignment_c_bits=alignment_c_bits ) sw7 = json.loads(ir_json) if ir_json else {} sw7_alpha = float(sw7.get("alpha", 1.702)) @@ -620,8 +621,6 @@ def _sw7_call(A, B, D, _fn=kernel_fn, _a=sw7_alpha, _l=sw7_limit, _o=sw7_one): b_dtype, b_layout=b_layout, m_bucket=m_bucket, - N=N_w, - K=K_w, alignment_a_bits=alignment_a_bits, alignment_b_bits=alignment_b_bits, alignment_c_bits=alignment_c_bits, diff --git a/magi_compiler/passes/piecewise_graph/fusion/sm80/cutlass_kernels/swiglu_one_stage.cu b/magi_compiler/passes/piecewise_graph/fusion/sm80/cutlass_kernels/swiglu_one_stage.cu index a4f01d4..59bbad8 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/sm80/cutlass_kernels/swiglu_one_stage.cu +++ b/magi_compiler/passes/piecewise_graph/fusion/sm80/cutlass_kernels/swiglu_one_stage.cu @@ -30,10 +30,11 @@ // smem stages; their accumulators stay in registers and a custom // SwigluCombine epilogue functor combines them and writes only D. // -// AUTOTUNE: at first call per (M, N, K) tuple the runner times every -// registered (TileShape, WarpShape, Stages) candidate and caches the -// fastest one. Candidate set is sized to the sm_120 / Ada SMEM budget -// (~96 KB per CTA); see SwAutoTuneRunner for SMEM math. +// AUTOTUNE: first call per (N_out, K) times every registered +// (TileShape, WarpShape, Stages) candidate and caches the winner in +// best_idx_map_. Different M inside one bucket reuse that choice. +// Candidate set is sized to the sm_120 / Ada SMEM budget (~96 KB per +// CTA); see SwAutoTuneRunner for SMEM math. #include #include @@ -41,6 +42,7 @@ #include #include +#include #include #include #include @@ -240,12 +242,21 @@ class SwImpl : public SwConcept { // AutoTune runner — first call per (M, N_out, K) shape times all candidates. //////////////////////////////////////////////////////////////////////////////// +#ifndef MAGI_EVT_MAX_TILES +#define MAGI_EVT_MAX_TILES 0 +#endif + #define SW_TILE(tb_m, tb_n, tb_k, wa_m, wa_n, wa_k, stages, label) \ - configs_.push_back(std::make_unique< \ - SwImpl, \ - cutlass::gemm::GemmShape, \ - stages>>>(label)) + do { \ + if (MAGI_EVT_MAX_TILES <= 0 \ + || static_cast(configs_.size()) < MAGI_EVT_MAX_TILES) { \ + configs_.push_back(std::make_unique< \ + SwImpl, \ + cutlass::gemm::GemmShape, \ + stages>>>(label)); \ + } \ + } while (0) class SwAutoTuneRunner { public: @@ -322,15 +333,9 @@ class SwAutoTuneRunner { cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.device().index()).stream(); - // Single autotune per module. The .cu is compiled per (M-bucket, N, K) - // on the Python side — every distinct weight (N, K) gets its own .cu, - // so this runner instance hosts exactly one (N, K) and one bucket. The - // first call autotunes; all subsequent calls (any M in the bucket) - // reuse `best_idx_`. - if (best_idx_ < 0) { - best_idx_ = autotune(ea, stream); - } - int idx = best_idx_; + // One .so is shared across all (N, K) for this (bucket, align). Autotune + // winners live in best_idx_map_ so a new shape only re-times candidates. + int idx = lookup_or_autotune(ea, stream); auto& gemm = configs_[idx]; size_t ws_sz = gemm->get_workspace_size(ea); @@ -351,6 +356,22 @@ class SwAutoTuneRunner { int num_configs() const { return (int)configs_.size(); } private: + static uint64_t nk_key(int N_out, int K) { + return (static_cast(static_cast(N_out)) << 32) + | static_cast(K); + } + + int lookup_or_autotune(const SwArgs& ea, cudaStream_t stream) { + const uint64_t key = nk_key(ea.N_out, ea.K); + auto it = best_idx_map_.find(key); + if (it != best_idx_map_.end()) { + return it->second; + } + int idx = autotune(ea, stream); + best_idx_map_.emplace(key, idx); + return idx; + } + int autotune(const SwArgs& ea, cudaStream_t stream) { int best_idx = -1; float best_time = 1e30f; @@ -397,7 +418,7 @@ class SwAutoTuneRunner { } std::vector> configs_; - int best_idx_ = -1; // -1 = not yet autotuned; sticky after first call. + std::unordered_map best_idx_map_; // (N_out, K) → tile index at::Tensor ws_; }; diff --git a/magi_compiler/passes/piecewise_graph/fusion/sm80/evt_codegen.py b/magi_compiler/passes/piecewise_graph/fusion/sm80/evt_codegen.py index a4e8471..08329c4 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/sm80/evt_codegen.py +++ b/magi_compiler/passes/piecewise_graph/fusion/sm80/evt_codegen.py @@ -29,6 +29,7 @@ _DTYPE_TO_CUTLASS, _VALID_ALIGN_BITS, _emit_custom_functor, + limit_tile_candidates, ) from ..evt_ir import Accum, AuxLoad, ColBroadcast, Compute, RowBroadcast, Store, walk_leaves @@ -68,7 +69,7 @@ def _emit_tile_candidates(m_bucket: str) -> str: """Emit C++ EVT_TILE_CANDIDATE(...) statements for the given M bucket.""" - candidates = _TILE_CANDIDATES_SM120.get(m_bucket, _TILE_CANDIDATES_SM120["medium"]) + candidates = limit_tile_candidates(_TILE_CANDIDATES_SM120.get(m_bucket, _TILE_CANDIDATES_SM120["medium"])) lines = [] for bm, bn, bk, wm, wn, wk, stages, label in candidates: lines.append(f' EVT_TILE_CANDIDATE({bm}, {bn}, {bk}, {wm}, {wn}, {wk}, ' f'{stages}, "{label}");') @@ -208,6 +209,7 @@ def _emit_args_tree(node, leaf_args: Dict[int, str], indent: int = 4) -> str: #include #include #include +#include #include #include #include @@ -312,8 +314,9 @@ def _emit_args_tree(node, leaf_args: Dict[int, str], indent: int = 4) -> str: }}; //////////////////////////////////////////////////////////////////////////////// -// Autotune runner — one candidate per tile/warp/stages combination; first call -// at a new (M, N, K) tuple times every candidate and caches the winner. +// Autotune runner — one candidate per tile/warp/stages combination. First call +// at each (N, K) times every candidate; later calls with the same (N, K) reuse +// the cached winner. Different M inside one bucket share that choice. //////////////////////////////////////////////////////////////////////////////// struct EvtArgs {{ @@ -484,15 +487,10 @@ class EvtAutoTuneRunner {{ cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.device().index()).stream(); - // Single autotune per module. The .cu is compiled per (IR, M-bucket, - // b_layout, N, K) on the Python side — every distinct weight (N, K) - // gets its own .cu, so this runner instance hosts exactly one (N, K) - // and one bucket of M values. Autotune once on the first call; all - // subsequent calls (any M inside the bucket) reuse `best_idx_`. - if (best_idx_ < 0) {{ - best_idx_ = autotune(ea, stream); - }} - int idx = best_idx_; + // One .so is shared across all (N, K) for this (IR, bucket, layout, + // align). Autotune winners live in best_idx_map_ so a new shape only + // re-times candidates — it does not trigger another nvcc. + int idx = lookup_or_autotune(ea, stream); auto& gemm = configs_[idx]; size_t ws_sz = gemm->get_workspace_size(ea); @@ -511,6 +509,22 @@ class EvtAutoTuneRunner {{ int num_configs() const {{ return (int)configs_.size(); }} private: + static uint64_t nk_key(int N, int K) {{ + return (static_cast(static_cast(N)) << 32) + | static_cast(K); + }} + + int lookup_or_autotune(const EvtArgs& ea, cudaStream_t stream) {{ + const uint64_t key = nk_key(ea.N, ea.K); + auto it = best_idx_map_.find(key); + if (it != best_idx_map_.end()) {{ + return it->second; + }} + int idx = autotune(ea, stream); + best_idx_map_.emplace(key, idx); + return idx; + }} + int autotune(const EvtArgs& ea, cudaStream_t stream) {{ int best_idx = -1; float best_time = 1e30f; @@ -582,7 +596,7 @@ class EvtAutoTuneRunner {{ }} std::vector> configs_; - int best_idx_ = -1; // -1 = not yet autotuned; sticky after first call. + std::unordered_map best_idx_map_; // (N, K) → tile index at::Tensor ws_; }}; diff --git a/magi_compiler/passes/piecewise_graph/fusion/sm90/cutlass_kernels/swiglu_one_stage.cu b/magi_compiler/passes/piecewise_graph/fusion/sm90/cutlass_kernels/swiglu_one_stage.cu index cb2b907..9578893 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/sm90/cutlass_kernels/swiglu_one_stage.cu +++ b/magi_compiler/passes/piecewise_graph/fusion/sm90/cutlass_kernels/swiglu_one_stage.cu @@ -25,10 +25,11 @@ // B : (N, K) bf16 row-major (torch.nn.Linear weight convention; N even) // D : (M, N/2) bf16 row-major (strided view of (M, ldd) host-padded buffer) // -// AUTOTUNE: at first call per (M, N, K) tuple the runner times every -// registered (TileShape, Stages) candidate and caches the fastest one. The -// candidate set targets H100's ~228 KiB dynamic-smem budget; per-stage smem -// for Sm90DualGemm = (BM + 2*BN) * BK * 2 (bf16) * stages. +// AUTOTUNE: first call per (N_out, K) times every registered +// (TileShape, Stages) candidate and caches the winner in best_idx_map_. +// Different M inside one bucket reuse that choice. The candidate set +// targets H100's ~228 KiB dynamic-smem budget; per-stage smem for +// Sm90DualGemm = (BM + 2*BN) * BK * 2 (bf16) * stages. // // Built by magi_compiler/passes/piecewise_graph/fusion/cutlass_fusion/ // evt_runtime.py::_compile_swiglu_dual when the live device's compute @@ -41,7 +42,9 @@ #include #include +#include #include +#include #include #include "cutlass/cutlass.h" @@ -227,11 +230,20 @@ class SwSm90Impl : public SwSm90Concept { // AutoTune runner — first call per (M, N_out, K) shape times all candidates. //////////////////////////////////////////////////////////////////////////////// +#ifndef MAGI_EVT_MAX_TILES +#define MAGI_EVT_MAX_TILES 0 +#endif + #define SW_SM90_TILE(bm, bn, bk, stages, label) \ - configs_.push_back(std::make_unique< \ - SwSm90Impl, cute::Int, cute::Int>, \ - stages>>>(label)) + do { \ + if (MAGI_EVT_MAX_TILES <= 0 \ + || static_cast(configs_.size()) < MAGI_EVT_MAX_TILES) { \ + configs_.push_back(std::make_unique< \ + SwSm90Impl, cute::Int, cute::Int>, \ + stages>>>(label)); \ + } \ + } while (0) class SwSm90AutoTuneRunner { public: @@ -317,13 +329,9 @@ class SwSm90AutoTuneRunner { cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.device().index()).stream(); - // Single autotune per module. The .cu is compiled per (m_bucket, N, K, - // alignA, alignB, alignC) on the Python side — every distinct shape - // bucket gets its own runner instance with isolated `best_idx_`. - if (best_idx_ < 0) { - best_idx_ = autotune(ea, stream); - } - int idx = best_idx_; + // One .so is shared across all (N, K) for this (bucket, align). Autotune + // winners live in best_idx_map_ so a new shape only re-times candidates. + int idx = lookup_or_autotune(ea, stream); auto& gemm = configs_[idx]; size_t ws_sz = gemm->get_workspace_size(ea); @@ -344,6 +352,22 @@ class SwSm90AutoTuneRunner { int num_configs() const { return (int)configs_.size(); } private: + static uint64_t nk_key(int N_out, int K) { + return (static_cast(static_cast(N_out)) << 32) + | static_cast(K); + } + + int lookup_or_autotune(const SwArgs& ea, cudaStream_t stream) { + const uint64_t key = nk_key(ea.N_out, ea.K); + auto it = best_idx_map_.find(key); + if (it != best_idx_map_.end()) { + return it->second; + } + int idx = autotune(ea, stream); + best_idx_map_.emplace(key, idx); + return idx; + } + int autotune(const SwArgs& ea, cudaStream_t stream) { int best_idx = -1; float best_time = 1e30f; @@ -387,7 +411,7 @@ class SwSm90AutoTuneRunner { } std::vector> configs_; - int best_idx_ = -1; // -1 = not yet autotuned; sticky after first call. + std::unordered_map best_idx_map_; // (N_out, K) → tile index at::Tensor ws_; }; diff --git a/magi_compiler/passes/piecewise_graph/fusion/sm90/evt_codegen.py b/magi_compiler/passes/piecewise_graph/fusion/sm90/evt_codegen.py index f14a7c8..83d6438 100644 --- a/magi_compiler/passes/piecewise_graph/fusion/sm90/evt_codegen.py +++ b/magi_compiler/passes/piecewise_graph/fusion/sm90/evt_codegen.py @@ -35,6 +35,7 @@ _DTYPE_TO_CUTLASS, _VALID_ALIGN_BITS, _emit_custom_functor, + limit_tile_candidates, ) from ..evt_ir import Accum, AuxLoad, ColBroadcast, Compute, RowBroadcast, Store, walk_leaves @@ -79,7 +80,7 @@ def _emit_tile_candidates(m_bucket: str) -> str: """Emit C++ EVT_TILE_CANDIDATE(...) statements for the given M bucket.""" - candidates = _TILE_CANDIDATES_SM90.get(m_bucket, _TILE_CANDIDATES_SM90["medium"]) + candidates = limit_tile_candidates(_TILE_CANDIDATES_SM90.get(m_bucket, _TILE_CANDIDATES_SM90["medium"])) lines = [] for tm, tn, tk, cm, cn, ck, schedule, label in candidates: kernel_sched, epi_sched = _SCHEDULE_TYPES[schedule] @@ -209,8 +210,10 @@ def _emit_args_tree(node, leaf_args: Dict[int, str], indent: int = 8) -> str: #include #include #include +#include #include #include +#include #include #include "cutlass/cutlass.h" @@ -517,15 +520,10 @@ class EvtAutoTuneRunner {{ ea.ptr_extras.reserve({n_extras}); {extras_ptrs} - // Single autotune per module. The .cu is compiled per (IR, M-bucket, - // b_layout, N, K) on the Python side — every distinct weight (N, K) - // gets its own .cu, so this runner instance hosts exactly one (N, K) - // and one bucket of M values. Autotune once on the first call; all - // subsequent calls (any M inside the bucket) reuse `best_idx_`. - if (best_idx_ < 0) {{ - best_idx_ = autotune(ea, stream); - }} - int idx = best_idx_; + // One .so is shared across all (N, K) for this (IR, bucket, layout, + // align). Autotune winners live in best_idx_map_ so a new shape only + // re-times candidates — it does not trigger another nvcc. + int idx = lookup_or_autotune(ea, stream); auto& gemm = configs_[idx]; size_t ws_sz = gemm->get_workspace_size(ea); @@ -544,6 +542,22 @@ class EvtAutoTuneRunner {{ int num_configs() const {{ return (int)configs_.size(); }} private: + static uint64_t nk_key(int N, int K) {{ + return (static_cast(static_cast(N)) << 32) + | static_cast(K); + }} + + int lookup_or_autotune(const EvtArgs& ea, cudaStream_t stream) {{ + const uint64_t key = nk_key(ea.N, ea.K); + auto it = best_idx_map_.find(key); + if (it != best_idx_map_.end()) {{ + return it->second; + }} + int idx = autotune(ea, stream); + best_idx_map_.emplace(key, idx); + return idx; + }} + int autotune(const EvtArgs& ea, cudaStream_t stream) {{ int best_idx = -1; float best_time = 1e30f; @@ -619,7 +633,7 @@ class EvtAutoTuneRunner {{ }} std::vector> configs_; - int best_idx_ = -1; // -1 = not yet autotuned; sticky after first call. + std::unordered_map best_idx_map_; // (N, K) → tile index at::Tensor ws_; }}; diff --git a/tests/feature_tests/pass/test_matmul_epilogue_fusion.py b/tests/feature_tests/pass/test_matmul_epilogue_fusion.py index a184038..77cdc58 100644 --- a/tests/feature_tests/pass/test_matmul_epilogue_fusion.py +++ b/tests/feature_tests/pass/test_matmul_epilogue_fusion.py @@ -32,8 +32,13 @@ ``magi_epilogue.matmul_fused_epilogue`` node. 3. Negative fallback: shapes / dtypes / chains the EVT pass does NOT support must keep the original ``aten.mm`` and run through cuBLAS. + +This file sets ``MAGI_EVT_MAX_TILES=1`` so JIT compiles a single tile +instead of the full autotune list. Numerical checks still run; they do +not assert that the fastest tile was chosen. """ +import os from typing import Optional import pytest @@ -45,6 +50,11 @@ from magi_compiler.api import magi_compile from magi_compiler.config import get_compile_config +# Correctness tests do not need the full autotune tile set. One candidate +# cuts CUTLASS template instantiations ~6-8x. Production is unchanged +# unless MAGI_EVT_MAX_TILES is set in the environment. +os.environ.setdefault("MAGI_EVT_MAX_TILES", "1") + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") _HAS_CUTLASS = get_compile_config().has_cutlass @@ -332,20 +342,22 @@ def test_evt_unary_activations_fuse(epi_name, epi_fn, atol, rtol): @_EVT_CAPABLE -def test_evt_relu_native(): - """Plain ``aten.relu`` variants must fuse and preserve emitted output dtype.""" +def test_evt_relu_mm_plus_1d_bias(): + """``relu((mm + bias_N).float())`` — 1-D bias as RowBroadcast, fp32 relu out.""" - class Fp32Relu(nn.Module): + class M(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(_N, _K)) + self.bias = nn.Parameter(torch.randn(_N)) def forward(self, a): - return torch.relu(torch.mm(a, self.weight.permute(1, 0)).float()) + return torch.relu((torch.mm(a, self.weight.permute(1, 0)) + self.bias).float()) _compile_and_check( - Fp32Relu(), + M(), (_input_a(),), + atol=1.5, expect_fused=1, expect_kinds=["evt_col"], expect_out_dtype=torch.float32, @@ -444,68 +456,27 @@ def forward(self, a): @_EVT_CAPABLE -def test_evt_mm_plus_1d_bias(): - """``silu(mm + bias_N)`` — 1-D bias as RowBroadcast extras.""" - - class M(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(_N, _K)) - self.bias = nn.Parameter(torch.randn(_N)) - - def forward(self, a): - y = torch.mm(a, self.weight.permute(1, 0)) + self.bias - return high_precision_silu(y, out_dtype=torch.bfloat16) +def test_evt_aux_loads_padded_and_repeated_fuse(): + """Padded-stride AuxLoad + repeated / multiple extras in one fused chain. - _compile_and_check( - M(), - (_input_a(),), - atol=1.5, - expect_fused=1, - expect_kinds=["evt_col"], - expect_out_dtype=torch.bfloat16, - expect_actual_dtype=torch.bfloat16, - ) - - -@_EVT_CAPABLE -def test_evt_aux_load_padded_stride(): - """AuxLoad with padded row stride (stride(0) > N) must fuse and read correctly.""" + ``gate`` is a (M, N) view of a wider buffer (stride(0) > N). It is also + reused at two EVT positions, alongside two distinct AuxLoad tensors. + """ class M(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(_N, _K)) - def forward(self, a, gate): - y = torch.mm(a, self.weight.permute(1, 0)) * gate - return y.to(torch.bfloat16) + def forward(self, a, gate, r1, r2): + y = torch.mm(a, self.weight.permute(1, 0)) + return (y * gate + gate + r1 + r2).to(torch.bfloat16) a = _input_a() N_padded = _N + 64 gate_buf = torch.randn(_M, N_padded, device="cuda", dtype=torch.bfloat16) gate = gate_buf[:, :_N] # shape (_M, _N), stride (N_padded, 1) assert gate.stride() == (N_padded, 1), f"Expected padded stride, got {gate.stride()}" - _compile_and_check( - M(), (a, gate), atol=0.0, rtol=0.1, expect_fused=1, expect_kinds=["evt_col"], dynamic_arg_dims={"a": 0, "gate": 0} - ) - - -@_EVT_CAPABLE -def test_evt_multiple_and_repeated_aux_loads_fuse(): - """Multiple AuxLoad extras, with one tensor reused at multiple EVT positions.""" - - class M(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(_N, _K)) - - def forward(self, a, gate, r1, r2): - y = torch.mm(a, self.weight.permute(1, 0)) - return (y * gate + gate + r1 + r2).to(torch.bfloat16) - - a = _input_a() - gate = torch.randn(_M, _N, device="cuda", dtype=torch.bfloat16) r1 = torch.randn(_M, _N, device="cuda", dtype=torch.bfloat16) r2 = torch.randn(_M, _N, device="cuda", dtype=torch.bfloat16) _compile_and_check( @@ -1050,6 +1021,104 @@ def test_evt_codegen_sm80_per_node_compute_dtype(): assert "float, float" in src +def _evt_tile_invocations(src: str) -> int: + return len([line for line in src.splitlines() if line.startswith(" EVT_TILE_CANDIDATE(")]) + + +def test_evt_max_tiles_env_limits_codegen(monkeypatch): + """MAGI_EVT_MAX_TILES=1 must emit exactly one tile candidate per arch.""" + monkeypatch.setenv("MAGI_EVT_MAX_TILES", "1") + from magi_compiler.passes.piecewise_graph.fusion.evt_ir import Accum, Compute, Store + from magi_compiler.passes.piecewise_graph.fusion.sm80.evt_codegen import render_evt_cu as render_sm80 + from magi_compiler.passes.piecewise_graph.fusion.sm90.evt_codegen import render_evt_cu as render_sm90 + + ir = Store(Compute("silu", (Accum(),)), "bfloat16") + for src in (render_sm80(ir, "bfloat16", "bfloat16"), render_sm90(ir, "bfloat16", "bfloat16")): + assert _evt_tile_invocations(src) == 1 + + +def test_evt_max_tiles_unset_keeps_full_bucket(monkeypatch): + """Unset MAGI_EVT_MAX_TILES must keep the full medium-bucket candidate list.""" + monkeypatch.delenv("MAGI_EVT_MAX_TILES", raising=False) + from magi_compiler.passes.piecewise_graph.fusion.evt_ir import Accum, Compute, Store + from magi_compiler.passes.piecewise_graph.fusion.sm80.evt_codegen import _TILE_CANDIDATES_SM120 + from magi_compiler.passes.piecewise_graph.fusion.sm80.evt_codegen import render_evt_cu as render_sm80 + from magi_compiler.passes.piecewise_graph.fusion.sm90.evt_codegen import _TILE_CANDIDATES_SM90 + from magi_compiler.passes.piecewise_graph.fusion.sm90.evt_codegen import render_evt_cu as render_sm90 + + ir = Store(Compute("silu", (Accum(),)), "bfloat16") + assert _evt_tile_invocations(render_sm80(ir, "bfloat16", "bfloat16")) == len(_TILE_CANDIDATES_SM120["medium"]) + assert _evt_tile_invocations(render_sm90(ir, "bfloat16", "bfloat16")) == len(_TILE_CANDIDATES_SM90["medium"]) + + +@_EVT_CAPABLE +def test_evt_nvcc_compile_cache_is_hit(monkeypatch): + """Same IR must hit the nvcc cache: memory reuse, (N, K) sharing, and on-disk .so. + + ``cpp_extension.load`` is the nvcc entry point. After a cold JIT it must + not be called again for the same IR/align, including a different (N, K) + and a reload after dropping the process-level module dicts. + """ + from pathlib import Path + + import torch.utils.cpp_extension as cpp_extension + + from magi_compiler.passes.piecewise_graph.fusion import evt_runtime as rt + from magi_compiler.passes.piecewise_graph.fusion.evt_ir import Accum, Compute, Store, to_canonical_json + + load_calls: list = [] + real_load = cpp_extension.load + + def counting_load(*args, **kwargs): + load_calls.append(kwargs.get("name", args[0] if args else "?")) + return real_load(*args, **kwargs) + + monkeypatch.setattr(cpp_extension, "load", counting_load) + + ir_json = to_canonical_json(Store(Compute("silu", (Accum(),)), "bfloat16")) + dt = torch.bfloat16 + compile_kw = dict(ir_json=ir_json, a_dtype=dt, b_dtype=dt, b_layout="col", m_bucket="medium") + + rt._MODULE_CACHE.clear() + rt._MODULE_FAST_CACHE.clear() + rt._DISPATCH_CACHE.clear() + + mod1 = rt._compile_evt_module(**compile_kw) + assert len(load_calls) == 1, f"cold JIT should nvcc once, got {load_calls}" + assert hasattr(mod1, "evt_matmul_out") + sos = list(Path(get_compile_config().cache_root_dir).rglob("magi_evt_*.so")) + assert sos, "cold JIT must write an on-disk EVT .so" + + mod2 = rt._compile_evt_module(**compile_kw) + assert mod2 is mod1 + assert len(load_calls) == 1, f"in-process cache missed; extra nvcc {load_calls}" + + # Same IR, different problem N — align is still 128-bit, must share the .so. + entry = rt._resolve_dispatch("evt_col", ir_json, dt, dt, 1032, 1024, "medium", dt) + assert entry.is_evt + assert len(load_calls) == 1, f"different (N, K) recompiled; extra nvcc {load_calls}" + + rt._MODULE_CACHE.clear() + rt._MODULE_FAST_CACHE.clear() + rt._DISPATCH_CACHE.clear() + mod3 = rt._compile_evt_module(**compile_kw) + assert len(load_calls) == 1, f"on-disk .so miss fell back to nvcc; extra {load_calls}" + assert hasattr(mod3, "evt_matmul_out") + + +def test_evt_codegen_autotune_is_per_nk_map(): + """Generated runners must cache autotune winners by (N, K), not one sticky idx.""" + from magi_compiler.passes.piecewise_graph.fusion.evt_ir import Accum, Compute, Store + from magi_compiler.passes.piecewise_graph.fusion.sm80.evt_codegen import render_evt_cu as render_sm80 + from magi_compiler.passes.piecewise_graph.fusion.sm90.evt_codegen import render_evt_cu as render_sm90 + + ir = Store(Compute("silu", (Accum(),)), "bfloat16") + for src in (render_sm80(ir, "bfloat16", "bfloat16"), render_sm90(ir, "bfloat16", "bfloat16")): + assert "best_idx_map_" in src + assert "lookup_or_autotune" in src + assert "int best_idx_ = -1" not in src + + def test_evt_codegen_sm90_per_node_compute_dtype(): """SM90 codegen emits per-node element types in Sm90Compute.""" from magi_compiler.passes.piecewise_graph.fusion.evt_ir import Accum, Compute, Store