diff --git a/src/flag_blas/runtime/backend/_nvidia/hopper/ops/gemm.py b/src/flag_blas/runtime/backend/_nvidia/hopper/ops/gemm.py
index 93364d3..77ff88d 100644
--- a/src/flag_blas/runtime/backend/_nvidia/hopper/ops/gemm.py
+++ b/src/flag_blas/runtime/backend/_nvidia/hopper/ops/gemm.py
@@ -1904,6 +1904,181 @@ def _hgemm_nn_build_kernel(
 ])
 
 
+# ---------------------------------------------------------------------------
+# Module-level condition predicates for hgemm_tn StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_tn_is_skinny_aligned(m, n, k, aligned, **_kw):
+    return aligned and (
+        (m >= 16384 and max(n, k) <= 2048)
+        or (n >= 16384 and max(m, k) <= 2048)
+    )
+
+
+def _hgemm_tn_is_aligned(m, n, k, aligned, **_kw):
+    return aligned
+
+
+def _hgemm_tn_is_default(**_kw):
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Module-level condition predicates for hgemm_nt StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_nt_is_aligned(m, n, k, aligned, **_kw):
+    return aligned
+
+
+def _hgemm_nt_is_default(**_kw):
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Module-level condition predicates for hgemm_tt StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_tt_is_skinny_aligned(m, n, k, aligned, **_kw):
+    return aligned and m >= 16384 and max(n, k) <= 2048
+
+
+def _hgemm_tt_is_aligned(m, n, k, aligned, **_kw):
+    return aligned
+
+
+def _hgemm_tt_is_default(**_kw):
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Module-level factory functions for hgemm_tn StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_tn_build_kernel3(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    return lambda: _hgemm_tn_kernel3[(
+        triton.cdiv(m, 128) * triton.cdiv(n, 256),
+    )](
+        TensorDescriptor(base=A, shape=[k, m], strides=[lda, 1], block_shape=[64, 128]),
+        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_tn_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_tn_kernel2[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+def _hgemm_tn_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_tn_kernel[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Module-level factory functions for hgemm_nt StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_nt_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_nt_kernel2[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+def _hgemm_nt_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_nt_kernel[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Module-level factory functions for hgemm_tt StaticDispatch
+# ---------------------------------------------------------------------------
+def _hgemm_tt_build_kernel3(
+    A, B, C, m, n, k, lda, ldb, ldc, alpha, beta, beta_is_zero,
+):
+    return lambda: _hgemm_tt_kernel3[(
+        triton.cdiv(m, 128) * triton.cdiv(n, 256),
+    )](
+        TensorDescriptor(base=A, shape=[k, m], strides=[lda, 1], block_shape=[64, 128]),
+        TensorDescriptor(base=B, shape=[n, k], strides=[ldb, 1], block_shape=[256, 64]),
+        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_tt_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_tt_kernel2[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+def _hgemm_tt_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_tt_kernel[grid](
+        A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero,
+    )
+
+
+_HGEMM_TN_DISPATCH = StaticDispatch([
+    # skinny + aligned → kernel3 (TensorDescriptor, hardcoded config)
+    (_hgemm_tn_is_skinny_aligned, _hgemm_tn_build_kernel3),
+    # aligned → kernel2 (block_ptr)
+    (_hgemm_tn_is_aligned, _hgemm_tn_build_kernel2),
+    # default → kernel (original)
+    (_hgemm_tn_is_default, _hgemm_tn_build_kernel),
+])
+
+_HGEMM_NT_DISPATCH = StaticDispatch([
+    # aligned → kernel2 (block_ptr)
+    (_hgemm_nt_is_aligned, _hgemm_nt_build_kernel2),
+    # default → kernel (original)
+    (_hgemm_nt_is_default, _hgemm_nt_build_kernel),
+])
+
+_HGEMM_TT_DISPATCH = StaticDispatch([
+    # skinny + aligned → kernel3 (TensorDescriptor, hardcoded config)
+    (_hgemm_tt_is_skinny_aligned, _hgemm_tt_build_kernel3),
+    # aligned → kernel2 (block_ptr)
+    (_hgemm_tt_is_aligned, _hgemm_tt_build_kernel2),
+    # default → kernel (original)
+    (_hgemm_tt_is_default, _hgemm_tt_build_kernel),
+])
+
+
 def hgemm(
     transa: int,
     transb: int,
@@ -1975,126 +2150,35 @@ def hgemm(
             )
             runner()
         elif transa == CUBLAS_OP_T and transb == CUBLAS_OP_N:
-            is_skinny = (m >= 16384 and max(n, k) <= 2048) or (
-                n >= 16384 and max(m, k) <= 2048
+            runner = _HGEMM_TN_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,
+                ),
             )
-            if aligned and is_skinny:
-                BLOCK_M = 128
-                BLOCK_N = 256
-                BLOCK_K = 64
-                GROUP_M = 8
-                NUM_STAGES = 4
-                NUM_WARPS = 8
-                NUM_CTAS = 1
-                desc_a = TensorDescriptor(
-                    base=A,
-                    shape=[k, m],
-                    strides=[lda, 1],
-                    block_shape=[BLOCK_K, BLOCK_M],
-                )
-                desc_b = TensorDescriptor(
-                    base=B,
-                    shape=[k, n],
-                    strides=[ldb, 1],
-                    block_shape=[BLOCK_K, BLOCK_N],
-                )
-                desc_c = TensorDescriptor(
-                    base=C,
-                    shape=[m, n],
-                    strides=[ldc, 1],
-                    block_shape=[BLOCK_M, BLOCK_N],
-                )
-                grid = (triton.cdiv(m, BLOCK_M) * triton.cdiv(n, BLOCK_N),)
-                _hgemm_tn_kernel3[grid](
-                    desc_a,
-                    desc_b,
-                    desc_c,
-                    alpha,
-                    beta,
-                    m,
-                    n,
-                    k,
-                    beta_is_zero,
-                    BLOCK_M=BLOCK_M,
-                    BLOCK_N=BLOCK_N,
-                    BLOCK_K=BLOCK_K,
-                    GROUP_M=GROUP_M,
-                    num_stages=NUM_STAGES,
-                    num_warps=NUM_WARPS,
-                    num_ctas=NUM_CTAS,
-                )
-            elif aligned:
-                _hgemm_tn_kernel2[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
-            else:
-                _hgemm_tn_kernel[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
+            runner()
         elif transa == CUBLAS_OP_N and transb == CUBLAS_OP_T:
-            if aligned:
-                _hgemm_nt_kernel2[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
-            else:
-                _hgemm_nt_kernel[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
+            runner = _HGEMM_NT_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,
+                ),
+            )
+            runner()
         else:
-            is_skinny = m >= 16384 and max(n, k) <= 2048
-            if is_skinny and aligned:
-                BLOCK_M = 128
-                BLOCK_N = 256
-                BLOCK_K = 64
-                GROUP_M = 8
-                NUM_STAGES = 4
-                NUM_WARPS = 8
-                NUM_CTAS = 1
-                desc_a = TensorDescriptor(
-                    base=A,
-                    shape=[k, m],
-                    strides=[lda, 1],
-                    block_shape=[BLOCK_K, BLOCK_M],
-                )
-                desc_b = TensorDescriptor(
-                    base=B,
-                    shape=[n, k],
-                    strides=[ldb, 1],
-                    block_shape=[BLOCK_N, BLOCK_K],
-                )
-                desc_c = TensorDescriptor(
-                    base=C,
-                    shape=[m, n],
-                    strides=[ldc, 1],
-                    block_shape=[BLOCK_M, BLOCK_N],
-                )
-                grid = (triton.cdiv(m, BLOCK_M) * triton.cdiv(n, BLOCK_N),)
-                _hgemm_tt_kernel3[grid](
-                    desc_a,
-                    desc_b,
-                    desc_c,
-                    alpha,
-                    beta,
-                    m,
-                    n,
-                    k,
-                    beta == 0.0,
-                    BLOCK_M=BLOCK_M,
-                    BLOCK_N=BLOCK_N,
-                    BLOCK_K=BLOCK_K,
-                    GROUP_M=GROUP_M,
-                    num_stages=NUM_STAGES,
-                    num_warps=NUM_WARPS,
-                    num_ctas=NUM_CTAS,
-                )
-            elif aligned:
-                _hgemm_tt_kernel2[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
-            else:
-                _hgemm_tt_kernel[grid](
-                    A, B, C, alpha, beta, m, n, k, lda, ldb, ldc, beta_is_zero
-                )
+            runner = _HGEMM_TT_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,
+                ),
+            )
+            runner()
 
 
 @libentry()
diff --git a/src/flag_blas/utils/code_utils.py b/src/flag_blas/utils/code_utils.py
index 783fff6..91edab0 100644
--- a/src/flag_blas/utils/code_utils.py
+++ b/src/flag_blas/utils/code_utils.py
@@ -1,63 +1,18 @@
-# The code for IndentedBuffer is adapted from
-# https://github.com/pytorch/pytorch/blob/ed48ea9997c2b04736096e4b6669543ab2e627d5/torch/_inductor/utils.py#L742
-# The code for Namespace is adapted from
-# https://github.com/pytorch/pytorch/blob/ed48ea9997c2b04736096e4b6669543ab2e627d5/torch/fx/graph.py#L115
-
-# License from pytorch(https://github.com/pytorch/pytorch)
-
-# From PyTorch:
-
-# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
-# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
-# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
-# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
-# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
-# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
-# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
-# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
-# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
-
-# From Caffe2:
-
-# Copyright (c) 2016-present, Facebook Inc. All rights reserved.
-
-# All contributions by Facebook:
-# Copyright (c) 2016 Facebook Inc.
-
-# All contributions by Google:
-# Copyright (c) 2015 Google Inc.
-# All rights reserved.
-
-# All contributions by Yangqing Jia:
-# Copyright (c) 2015 Yangqing Jia
-# All rights reserved.
-
-# All contributions by Kakao Brain:
-# Copyright 2019-2020 Kakao Brain
-
-# All contributions by Cruise LLC:
-# Copyright (c) 2022 Cruise LLC.
-# All rights reserved.
-
-# All contributions from Caffe:
-# Copyright(c) 2013, 2014, 2015, the respective contributors
-# All rights reserved.
-
-# All other contributions:
-# Copyright(c) 2015, 2016 the respective contributors
-# All rights reserved.
-
-# Caffe2 uses a copyright model similar to Caffe: each contributor holds
-# copyright over their contributions to Caffe2. The project versioning records
-# all such contribution and copyright details. If a contributor wants to further
-# mark their specific copyright on a particular contribution, they should
-# indicate their copyright solely in the commit message of the change when it is
-# committed.
-
-# All rights reserved.
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import builtins
-import contextlib
 import keyword
 import os
 import re
@@ -66,85 +21,115 @@
 from collections import defaultdict
 from io import StringIO
 from pathlib import Path
-from typing import Dict, Set
+from typing import Dict, List, Set
 
 
 class IndentedBuffer:
+    """A buffer for building indented text output (e.g. generated code).
+
+    Maintains a list of lines and a current indentation level. Lines are
+    written with automatic indentation prefix applied.
+    """
+
     tabwidth = 4
 
-    def __init__(self, initial_indent=0):
-        self._lines = []
-        self._indent = initial_indent
+    def __init__(self, initial_indent: int = 0):
+        self._lines: List[str] = []
+        self._level = initial_indent
+
+    # ---- output ----
 
     def getvalue(self) -> str:
         buf = StringIO()
         for line in self._lines:
-            assert isinstance(line, str)
             buf.write(line)
             buf.write("\n")
         return buf.getvalue()
 
-    def clear(self):
+    def clear(self) -> None:
         self._lines.clear()
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return bool(self._lines)
 
-    def prefix(self):
-        return " " * (self._indent * self.tabwidth)
+    # ---- prefix ----
+
+    def _prefix(self) -> str:
+        return " " * (self._level * self.tabwidth)
+
+    def prefix(self) -> str:
+        return self._prefix()
 
-    def newline(self):
-        self.writeline("\n")
+    # ---- write helpers ----
 
-    def writeline(self, line):
+    def newline(self) -> None:
+        """Write a blank line (no indentation)."""
+        self._lines.append("")
+
+    def writeline(self, line: str) -> None:
+        """Write a single line. Non-blank lines are prefixed with current indent."""
         if line.strip():
-            self._lines.append(f"{self.prefix()}{line}")
+            self._lines.append(f"{self._prefix()}{line}")
         else:
             self._lines.append("")
 
-    def tpl(self, format_str, **kwargs):
-        assert isinstance(format_str, str), "format_str must be string of type."
-        format_str = format_str.format(**kwargs)
-        lines = format_str.strip().splitlines()
+    def writelines(self, lines) -> None:
+        """Write a sequence of lines."""
         for line in lines:
-            line = line.replace("\t", " " * self.tabwidth)
             self.writeline(line)
 
-    def writelines(self, lines):
-        for line in lines:
-            self.writeline(line)
-
-    def writemultiline(self, s):
+    def writemultiline(self, s: str) -> None:
+        """Write a multi-line string, splitting on newlines."""
         self.writelines(s.splitlines())
 
-    def indent(self, offset=1):
-        @contextlib.contextmanager
-        def ctx():
-            self._indent += offset
-            try:
-                yield
-            finally:
-                self._indent -= offset
+    def tpl(self, format_str: str, **kwargs) -> None:
+        """Write a template string: format it, then write each line."""
+        formatted = format_str.format(**kwargs)
+        for line in formatted.strip().splitlines():
+            self.writeline(line.replace("\t", " " * self.tabwidth))
+
+    # ---- indent ----
+
+    def indent(self, offset: int = 1):
+        """Context manager that temporarily increases indentation level.
+
+        Usage::
+
+            with buf.indent():
+                buf.writeline("inside")
+        """
+        return self._IndentCtx(self, offset)
+
+    class _IndentCtx:
+        def __init__(self, buf: "IndentedBuffer", offset: int):
+            self._buf = buf
+            self._offset = offset
+
+        def __enter__(self):
+            self._buf._level += self._offset
+            return self
 
-        return ctx()
+        def __exit__(self, *args):
+            self._buf._level -= self._offset
 
 
 class NameSpace:
-    def __init__(self):
-        self._used_names: Set[str] = set()
-        self._base_count: Dict[str, int] = defaultdict(int)
+    """Generates unique, valid Python identifiers from candidate names.
 
-        self._illegal_char_regex = re.compile("[^0-9a-zA-Z_]+")
-        self._name_suffix_regex = re.compile(r"(.*)_(\d+)$")
+    Sanitizes input, avoids Python keywords and builtins, and appends numeric
+    suffixes to prevent collisions.
+    """
 
-    def create_name(self, candidate: str) -> str:
-        """Create a unique name.
+    def __init__(self):
+        self._used: Set[str] = set()
+        self._counters: Dict[str, int] = defaultdict(int)
+        self._sanitize_re = re.compile(r"[^0-9a-zA-Z_]+")
+        self._suffix_re = re.compile(r"(.*)_(\d+)$")
 
-        Arguments:
-            candidate: used as the basis for the unique name, relevant to the user.
-        """
-        # delete all characters that are illegal in a Python identifier
-        candidate = self._illegal_char_regex.sub("_", candidate)
+    def create_name(self, candidate: str) -> str:
+        """Return a unique valid Python identifier based on *candidate*."""
+        # strip illegal characters
+        candidate = self._sanitize_re.sub("_", candidate)
 
         if not candidate:
             candidate = "_unnamed"
@@ -152,35 +137,32 @@ def create_name(self, candidate: str) -> str:
         if candidate[0].isdigit():
             candidate = f"_{candidate}"
 
-        match = self._name_suffix_regex.match(candidate)
-        if match is None:
-            base = candidate
-            num = None
+        # parse base name and optional numeric suffix
+        m = self._suffix_re.match(candidate)
+        if m is None:
+            base, num = candidate, None
         else:
-            base, num_str = match.group(1, 2)
+            base, num_str = m.group(1, 2)
             num = int(num_str)
 
         candidate = base if num is None else f"{base}_{num}"
-        if not num:
-            num = self._base_count[base]
+        if num is None:
+            num = self._counters[base]
 
-        while candidate in self._used_names or self._is_illegal_name(candidate):
+        # advance until unique and legal
+        while candidate in self._used or self._illegal(candidate):
             num += 1
             candidate = f"{base}_{num}"
 
-        self._used_names.add(candidate)
-        self._base_count[base] = num
+        self._used.add(candidate)
+        self._counters[base] = num
         return candidate
 
-    def _is_illegal_name(self, name: str) -> bool:
-        # 1. keywords are never allowed as names.
+    def _illegal(self, name: str) -> bool:
         if name in keyword.kwlist:
             return True
-
-        # 2. Can't shadow a builtin name, unless you *are* that builtin.
         if name in builtins.__dict__:
             return True
-
         return False
 
 
@@ -190,6 +172,7 @@ def write_atomic(
     make_dirs: bool = False,
     encoding: str = "utf-8",
 ) -> None:
+    """Atomically write *content* to *path_* via a temporary file + rename."""
     path = Path(path_)
     if make_dirs:
         path.parent.mkdir(parents=True, exist_ok=True)