optimize static dispatch

Author: bin913

src/flag_blas/runtime/backend/_nvidia/hopper/ops/gemm.py

@@ -1800,6 +1800,96 @@ def _hgemm_tt_kernel3(
         desc_c.store([offs_m, offs_n], result)
 
 
+# ---------------------------------------------------------------------------
+# Module-level condition predicates for hgemm StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_nn_is_skinny_aligned_large(m, n, k, aligned, **_kw):
+    return (
+        aligned
+        and (m * n > 2048 * 2048)
+        and min(m, n) >= 64
+        and (
+            (m >= 16384 and max(n, k) <= 2048)
+            or (n >= 16384 and max(m, k) <= 2048)
+        )
+    )
+
+
+def _hgemm_nn_is_aligned_large(m, n, k, aligned, **_kw):
+    return aligned and (m * n > 2048 * 2048) and min(m, n) >= 64
+
+
+def _hgemm_nn_is_aligned_small(m, n, k, aligned, **_kw):
+    return aligned and max(m, n) <= 1024
+
+
+def _hgemm_nn_is_default(**_kw):
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Module-level factory functions for hgemm StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_nn_build_kernel4(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    return lambda: _hgemm_nn_kernel4[(
+        triton.cdiv(m, 128) * triton.cdiv(n, 256),
+    )](
+        TensorDescriptor(base=A, shape=[m, k], strides=[lda, 1], block_shape=[128, 64]),
+        TensorDescriptor(base=B, shape=[k, n], strides=[ldb, 1], block_shape=[64, 256]),
+        TensorDescriptor(base=C, shape=[m, n], strides=[ldc, 1], block_shape=[128, 256]),
+        alpha, beta, m, n, k, beta_is_zero,
+        BLOCK_M=128, BLOCK_N=256, BLOCK_K=64, GROUP_M=8,
+        num_stages=4, num_warps=8, num_ctas=1,
+    )
+
+
+def _hgemm_nn_build_kernel3(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    grid = lambda meta: (
+        triton.cdiv(m, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),
+    )
+    return lambda: _hgemm_nn_kernel3[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+def _hgemm_nn_build_kernel2(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    grid = lambda meta: (
+        triton.cdiv(m, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),
+    )
+    return lambda: _hgemm_nn_kernel2[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+def _hgemm_nn_build_kernel(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    grid = lambda meta: (
+        triton.cdiv(m, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),
+    )
+    return lambda: _hgemm_nn_kernel[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+_HGEMM_NN_DISPATCH = StaticDispatch([
+    # skinny + aligned large → kernel4 (TensorDescriptor, hardcoded config)
+    (_hgemm_nn_is_skinny_aligned_large, _hgemm_nn_build_kernel4),
+    # aligned large → kernel3 (TensorDescriptor, autotuned config)
+    (_hgemm_nn_is_aligned_large, _hgemm_nn_build_kernel3),
+    # aligned small → kernel2 (block_ptr)
+    (_hgemm_nn_is_aligned_small, _hgemm_nn_build_kernel2),
+    # default → kernel (original)
+    (_hgemm_nn_is_default, _hgemm_nn_build_kernel),
+])
+
+
 def hgemm(
     transa: int,
     transb: int,
@@ -1861,48 +1951,14 @@ def hgemm(
     aligned = _is_gemm_aligned(A, lda, B, ldb, C, ldc)
     with torch_device_fn.device(A.device):
         if transa == CUBLAS_OP_N and transb == CUBLAS_OP_N:
-            dispatch = StaticDispatch([
-                # skinny + aligned large → kernel4 (TensorDescriptor, hardcoded config)
-                (
-                    lambda m, n, k, aligned, **_kw:
-                        aligned and (m * n > 2048 * 2048) and min(m, n) >= 64
-                        and ((m >= 16384 and max(n, k) <= 2048) or (n >= 16384 and max(m, k) <= 2048)),
-                    lambda: lambda: _hgemm_nn_kernel4[(
-                        triton.cdiv(m, 128) * triton.cdiv(n, 256),
-                    )](
-                        TensorDescriptor(base=A, shape=[m, k], strides=[lda, 1], block_shape=[128, 64]),
-                        TensorDescriptor(base=B, shape=[k, n], strides=[ldb, 1], block_shape=[64, 256]),
-                        TensorDescriptor(base=C, shape=[m, n], strides=[ldc, 1], block_shape=[128, 256]),
-                        alpha, beta, m, n, k, beta_is_zero,
-                        BLOCK_M=128, BLOCK_N=256, BLOCK_K=64, GROUP_M=8,
-                        num_stages=4, num_warps=8, num_ctas=1,
-                    ),
-                ),
-                # aligned large → kernel3 (TensorDescriptor, autotuned config)
-                (
-                    lambda m, n, k, aligned, **_kw:
-                        aligned and (m * n > 2048 * 2048) and min(m, n) >= 64,
-                    lambda: lambda: _hgemm_nn_kernel3[grid](
-                        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
-                    ),
-                ),
-                # aligned small → kernel2 (block_ptr)
-                (
-                    lambda m, n, k, aligned, **_kw:
-                        aligned and max(m, n) <= 1024,
-                    lambda: lambda: _hgemm_nn_kernel2[grid](
-                        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
-                    ),
+            runner = _HGEMM_NN_DISPATCH.lookup_and_build(
+                m, n, k, aligned,
+                context=dict(
+                    A=A, B=B, C=C, m=m, n=n, k=k,
+                    lda=lda, ldb=ldb, ldc=ldc,
+                    alpha=alpha, beta=beta, beta_is_zero=beta_is_zero,
                 ),
-                # default → kernel (original)
-                (
-                    lambda **_kw: True,
-                    lambda: lambda: _hgemm_nn_kernel[grid](
-                        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
-                    ),
-                ),
-            ])
-            runner = dispatch.lookup_and_build(m, n, k, aligned)
+            )
             runner()
         elif transa == CUBLAS_OP_T and transb == CUBLAS_OP_N:
             is_skinny = (m >= 16384 and max(n, k) <= 2048) or (

src/flag_blas/runtime/dispatch.py

@@ -367,10 +367,14 @@ def lookup_and_build(
         n: int,
         k: int,
         aligned: bool,
+        *,
+        context: Optional[dict] = None,
         **extra,
     ) -> Callable[[], None]:
         for condition, factory in self._table:
             if condition(m=m, n=n, k=k, aligned=aligned, **extra):
+                if context is not None:
+                    return factory(**context)
                 return factory()
         raise ValueError(
             f"StaticDispatch: no matching entry for "