[pull] main from SandAI-org:main

.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

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

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,

magi_compiler/passes/piecewise_graph/fusion/sm80/cutlass_kernels/swiglu_one_stage.cu

@@ -30,17 +30,19 @@
 // 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 <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
+#include <cstdint>
 #include <memory>
 #include <unordered_map>
 #include <vector>
@@ -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<DualGemmConfig<                                                   \
-          cutlass::gemm::GemmShape<tb_m, tb_n, tb_k>,                           \
-          cutlass::gemm::GemmShape<wa_m, wa_n, wa_k>,                           \
-          stages>>>(label))
+  do {                                                                         \
+    if (MAGI_EVT_MAX_TILES <= 0                                                \
+        || static_cast<int>(configs_.size()) < MAGI_EVT_MAX_TILES) {          \
+      configs_.push_back(std::make_unique<                                     \
+          SwImpl<DualGemmConfig<                                               \
+              cutlass::gemm::GemmShape<tb_m, tb_n, tb_k>,                       \
+              cutlass::gemm::GemmShape<wa_m, wa_n, wa_k>,                       \
+              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<uint64_t>(static_cast<uint32_t>(N_out)) << 32)
+           | static_cast<uint32_t>(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<std::unique_ptr<SwConcept>> configs_;
-  int best_idx_ = -1;     // -1 = not yet autotuned; sticky after first call.
+  std::unordered_map<uint64_t, int> best_idx_map_;  // (N_out, K) → tile index
   at::Tensor ws_;
 };
 

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 <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
+#include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <math.h>
@@ -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<uint64_t>(static_cast<uint32_t>(N)) << 32)
+           | static_cast<uint32_t>(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<std::unique_ptr<EvtConcept>> configs_;
-  int best_idx_ = -1;     // -1 = not yet autotuned; sticky after first call.
+  std::unordered_map<uint64_t, int> best_idx_map_;  // (N, K) → tile index
   at::Tensor ws_;
 }};