submission 551966

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  """
- v142: Extreme Python inline optimization.
- - Function pointer swap: warmup function swaps to hot path (no if-check per call)
- - List-indexed cache by M (no dict lookup in hot path)
- - Pre-bound kernel function references (no global lookup)
- - @torch.no_grad() to skip autograd overhead
- - Minimal tuple unpacking in hot path
- - Pre-computed B_q_T strides cached per shape
- - All constants pre-computed during warmup/first-see
- - ASM path for M>64 with same optimizations
+ v178: v120 + Triton compiler globals patches.
+ Patches __oclc_unsafe_math_opt=True, __oclc_finite_only_opt=True,
+ allow_flush_denorm=True, knobs.amd.use_buffer_ops=True.
+ Result: no measurable improvement (MFMA not affected by math flags).
  """
+ import os
+ import sys
+  
+ os.environ['TRITON_DISABLE_LINE_INFO'] = '1'
+  
  from task import input_t, output_t
   
  import torch
  import triton
  import triton.language as tl
+  
+ # Monkey-patch __oclc control constants for unsafe math
+ try:
+     import triton.backends.amd.compiler as _amd_compiler
+     _amd_mod = getattr(_amd_compiler, 'amd', None)
+     if _amd_mod is not None:
+         _orig_set_bool = getattr(_amd_mod, 'set_bool_control_constant', None)
+         if _orig_set_bool is not None:
+             def _patched_set_bool(llvm_mod, name, val):
+                 if name == '__oclc_unsafe_math_opt':
+                     val = True
+                 elif name == '__oclc_finite_only_opt':
+                     val = True
+                 return _orig_set_bool(llvm_mod, name, val)
+             _amd_mod.set_bool_control_constant = _patched_set_bool
+ except Exception:
+     pass
+  
+ try:
+     from triton import knobs as _knobs
+     _knobs.amd.use_buffer_ops = True
+ except Exception:
+     pass
+  
  import aiter
  from aiter import dtypes
  from aiter.ops.triton.quant import _mxfp4_quant_op
  ⋯ 4 unchanged lines
  _reduce_kernel = _gemm_mod._gemm_afp4wfp4_reduce_kernel
  _get_splitk_fn = _gemm_mod.get_splitk
   
- # Pre-bind dtype constants
  _fp4x2 = dtypes.fp4x2
  _fp8_e8m0 = dtypes.fp8_e8m0
  _bf16 = dtypes.bf16
- _uint8 = torch.uint8
- _bfloat16 = torch.bfloat16
- _float32 = torch.float32
   
- # Pre-bind torch functions to avoid global lookups
- _torch_empty = torch.empty
- _torch_full = torch.full
- _triton_cdiv = triton.cdiv
- _triton_np2 = triton.next_power_of_2
   
-  
  @triton.jit
  def _remap_xcd(pid, num_pids, NUM_XCDS: tl.constexpr):
      chunk_size = tl.cdiv(num_pids, NUM_XCDS)
  ⋯ 33 unchanged lines
      K_packed = K_real // 2
      GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
      total_pids = GRID_MN * NUM_KSPLIT
+     # Pad to multiple of 8 for correct XCD remapping
      total_pids_padded = ((total_pids + 7) // 8) * 8
   
      pid_unified = tl.program_id(axis=0)
  ⋯ 111 unchanged lines
          tl.store(bs_ptr + shuffled_offset, bs_e8m0, mask=bs_store_mask)
   
   
- # ============================================================
- # Cache: dict keyed by (M,K,N) + last-seen fast path
- # Last-seen avoids dict lookup entirely for repeated calls
- # ============================================================
- _fused_cache = {}  # (M,K,N) -> config tuple
- _asm_cache = {}    # (M,K,N) -> config tuple
+ _cache_asm = {}
+ _cache_fused = {}
  _gemm_asm = None
+ _warmup_done = False
   
- # Last-seen fast path: avoids dict lookup for consecutive same-shape calls
- _last_fused_key = None  # (M,K,N) tuple
- _last_fused_cfg = None  # corresponding config
- _last_asm_key = None
- _last_asm_cfg = None
   
+ def custom_kernel(data: input_t) -> output_t:
+     global _gemm_asm, _warmup_done
   
- def _build_fused_config(M, K, N, A_device):
-     """Build and cache fused config for a given (M,K,N). Called once per shape."""
-     K_packed = K >> 1  # K // 2
-     scale_n = (K + 31) >> 5  # (K + 31) // 32
-     SCALE_N_B = ((scale_n + 7) >> 3) << 3  # round up to mult of 8
+     A, B, B_q, B_shuffle, B_scale_sh = data
+     M, K = A.shape
+     N = B_shuffle.shape[0]
   
-     BLOCK_SIZE_M = 16
-     BLOCK_SIZE_N = 64 if M <= 16 else 128
-     BLOCK_SIZE_K = 512
+     # Fused for ALL M<=16 shapes (saves 1 kernel launch)
+     # Fused for M<=32: single launch saves ~2-3us vs quant+ASM
+     # ASM for M>=64: ASM GEMM much faster for larger tiles
+     use_fused = (M <= 64)
   
-     num_pid_m = _triton_cdiv(M, BLOCK_SIZE_M)
-     num_pid_n = _triton_cdiv(N, BLOCK_SIZE_N)
-     base_blocks = num_pid_m * num_pid_n
-     target_ksplit = max(1, 256 // max(1, base_blocks))
+     # Warmup: use ASM path to init aiter module
+     if not _warmup_done:
+         scale_n_valid = (K + 31) // 32
+         SCALE_M = ((M + 255) // 256) * 256
+         SCALE_N = ((scale_n_valid + 7) // 8) * 8
+         BSM = triton.next_power_of_2(M) if M <= 32 else 16
+         grid = (triton.cdiv(M, BSM), triton.cdiv(K, 32))
   
-     NUM_KSPLIT = 1
-     SPLITK_BLOCK_SIZE = K  # 2 * K_packed = K
+         x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=A.device)
+         bs_sh = torch.full((SCALE_M, SCALE_N), 127, dtype=torch.uint8, device=A.device)
   
-     if target_ksplit > 1:
-         sb, bk_adj, nk = _get_splitk_fn(K_packed, BLOCK_SIZE_K, target_ksplit)
-         if bk_adj >= 512:
-             BLOCK_SIZE_K = bk_adj
-             SPLITK_BLOCK_SIZE = sb
-             NUM_KSPLIT = nk
+         _fused_quant_shuffle_kernel[grid](
+             A, x_fp4, bs_sh,
+             A.stride(0), A.stride(1),
+             x_fp4.stride(0), x_fp4.stride(1),
+             M, K, scale_n_valid,
+             SCALE_N=SCALE_N, BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=32,
+             NUM_ITER=1, NUM_STAGES=1, MXFP4_QUANT_BLOCK_SIZE=32,
+             num_warps=1, waves_per_eu=0, num_stages=1,
+         )
   
-     if NUM_KSPLIT > 1:
-         y_pp = _torch_empty((NUM_KSPLIT, M, N), dtype=_float32, device=A_device)
-     else:
-         y_pp = None
-         SPLITK_BLOCK_SIZE = K  # 2 * K_packed
+         result = aiter.gemm_a4w4(
+             x_fp4.view(_fp4x2), B_shuffle,
+             bs_sh.view(_fp8_e8m0), B_scale_sh,
+             dtype=_bf16, bpreshuffle=True,
+         )
+         _warmup_done = True
+         try:
+             _gemm_asm = torch.ops.aiter.gemm_a4w4_asm
+         except Exception:
+             try:
+                 import aiter.jit.core as _jc
+                 _gemm_asm = getattr(_jc, 'gemm_a4w4_asm', None)
+             except Exception:
+                 pass
+         return result
   
-     y = _torch_empty((M, N), dtype=_bfloat16, device=A_device)
+     if use_fused:
+         # --- Fused quant+GEMM: single kernel launch ---
+         key = (M, K, N)
+         c = _cache_fused.get(key)
+         if c is None:
+             K_packed = K // 2
+             scale_n = (K + 31) // 32
+             SCALE_N_B = ((scale_n + 7) // 8) * 8
   
-     total_blocks_raw = NUM_KSPLIT * num_pid_m * num_pid_n
-     total_blocks = ((total_blocks_raw + 7) >> 3) << 3  # pad to mult of 8
+             BLOCK_SIZE_M = 16
+             BLOCK_SIZE_N = 64 if M <= 16 else 128
+             BLOCK_SIZE_K = 512
   
-     bs_stride_n = 32 * SCALE_N_B
-     out_tensor = y if NUM_KSPLIT == 1 else y_pp
+             base_blocks = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
+             target_ksplit = max(1, 256 // max(1, base_blocks))
   
-     # Pre-compute strides for output
-     if NUM_KSPLIT == 1:
-         stride_ck = 0
-         stride_cm = y.stride(0)
-         stride_cn = y.stride(1)
-     else:
-         stride_ck = y_pp.stride(0)
-         stride_cm = y_pp.stride(1)
-         stride_cn = y_pp.stride(2)
+             if target_ksplit > 1:
+                 SPLITK_BLOCK_SIZE, BLOCK_SIZE_K_adj, NUM_KSPLIT = _get_splitk_fn(
+                     K_packed, BLOCK_SIZE_K, target_ksplit
+                 )
+                 if BLOCK_SIZE_K_adj < 512:
+                     # Can't use BSK=512 with this splitk, fallback
+                     BLOCK_SIZE_K_adj = 512
+                     SPLITK_BLOCK_SIZE = 2 * K_packed
+                     NUM_KSPLIT = 1
+                 else:
+                     BLOCK_SIZE_K = BLOCK_SIZE_K_adj
+             else:
+                 NUM_KSPLIT = 1
+                 SPLITK_BLOCK_SIZE = 2 * K_packed
   
-     # Pre-compute reduce params
-     reduce_params = None
-     if NUM_KSPLIT > 1:
-         ACTUAL_KSPLIT = _triton_cdiv(K_packed, SPLITK_BLOCK_SIZE >> 1)
-         grid_reduce = (_triton_cdiv(M, 16), _triton_cdiv(N, 64))
-         np2_ksplit = _triton_np2(NUM_KSPLIT)
-         reduce_params = (grid_reduce, ACTUAL_KSPLIT, np2_ksplit,
-                          y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
-                          y.stride(0), y.stride(1))
+             if NUM_KSPLIT > 1:
+                 y_pp = torch.empty((NUM_KSPLIT, M, N), dtype=torch.float32, device=A.device)
+             else:
+                 y_pp = None
+                 SPLITK_BLOCK_SIZE = 2 * K_packed
   
-     # Pre-compute the grid tuple
-     grid_tuple = (total_blocks,)
+             y = torch.empty((M, N), dtype=torch.bfloat16, device=A.device)
   
-     return (M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
-             NUM_KSPLIT, SPLITK_BLOCK_SIZE,
-             y, y_pp, out_tensor, grid_tuple,
-             stride_ck, stride_cm, stride_cn,
-             bs_stride_n, reduce_params)
+             total_blocks_raw = NUM_KSPLIT * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
+             # Pad to multiple of 8 for correct XCD remapping
+             total_blocks = ((total_blocks_raw + 7) // 8) * 8
   
+             bs_stride_n = 32 * SCALE_N_B
+             bs_stride_k = 1
   
- def _build_asm_config(M, K, N, A, B_shuffle, B_scale_sh):
-     """Build and cache ASM config for a given (M,K,N). Called once per shape."""
-     scale_n_valid = (K + 31) >> 5
-     SCALE_M = ((M + 255) // 256) * 256
-     SCALE_N = ((scale_n_valid + 7) >> 3) << 3
-     padded_m = get_padded_m(M, N, K, 0)
+             c = (K_packed, SCALE_N_B, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
+                  NUM_KSPLIT, SPLITK_BLOCK_SIZE,
+                  y, y_pp, total_blocks, bs_stride_n, bs_stride_k)
+             _cache_fused[key] = c
   
-     BSM = _triton_np2(M) if M <= 32 else 16
-     BSN = 32
-     NUM_ITER_Q = 2
-     grid = (_triton_cdiv(M, BSM), _triton_cdiv(K, BSN * NUM_ITER_Q))
+         (K_packed, SCALE_N_B, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
+          NUM_KSPLIT, SPLITK_BLOCK_SIZE,
+          y, y_pp, total_blocks, bs_stride_n, bs_stride_k) = c
   
-     ck_config = get_GEMM_config(M, N, K)
-     kernel_name = ""
-     split_k = 0
-     if ck_config is not None:
-         split_k = ck_config.get("splitK", 0) or 0
-         kernel_name = ck_config["kernelName"]
+         B_q_u8 = B_q.view(torch.uint8) if B_q.dtype != torch.uint8 else B_q
+         B_q_T = B_q_u8.T
+         B_scale_u8 = B_scale_sh.view(torch.uint8)
   
-     x_fp4 = _torch_empty((M, K >> 1), dtype=_uint8, device=A.device)
-     bs_sh = _torch_full((SCALE_M, SCALE_N), 127, dtype=_uint8, device=A.device)
-     out = _torch_empty((padded_m, N), dtype=_bfloat16, device=A.device)
+         out_tensor = y if NUM_KSPLIT == 1 else y_pp
   
-     x_fp4_view = x_fp4.view(_fp4x2)
-     bs_sh_view = bs_sh.view(_fp8_e8m0)
-     out_view = out[:M] if M < padded_m else out
-  
-     stride_a0 = A.stride(0)
-     stride_a1 = A.stride(1)
-     stride_fp4_0 = x_fp4.stride(0)
-     stride_fp4_1 = x_fp4.stride(1)
-  
-     return (scale_n_valid, SCALE_N, BSM, BSN, NUM_ITER_Q, grid,
-             x_fp4, bs_sh, out, x_fp4_view, bs_sh_view, out_view,
-             kernel_name, split_k,
-             stride_a0, stride_a1, stride_fp4_0, stride_fp4_1)
-  
-  
- # ============================================================
- # The hot path: no warmup check, no dict lookup, minimal Python
- # ============================================================
- @torch.no_grad()
- def _hot_fused(data):
-     """Fused quant+GEMM hot path for M<=64. Absolute minimum Python overhead."""
-     global _last_fused_key, _last_fused_cfg
-     A, B, B_q, B_shuffle, B_scale_sh = data
-     M = A.shape[0]
-     K = A.shape[1]
-     N = B_shuffle.shape[0]
-  
-     # Fast path: check last-seen key (avoids dict lookup for repeated shapes)
-     key = (M, K, N)
-     if key is _last_fused_key or key == _last_fused_key:
-         c = _last_fused_cfg
-     else:
-         c = _fused_cache.get(key)
-         if c is None:
-             c = _build_fused_config(M, K, N, A.device)
-             _fused_cache[key] = c
-         _last_fused_key = key
-         _last_fused_cfg = c
-  
-     # Unpack only what we need
-     (_, _, _, BSM, BSN, BSK,
-      NUM_KSPLIT, SPLITK_BLOCK_SIZE,
-      y, y_pp, out_tensor, grid_tuple,
-      stride_ck, stride_cm, stride_cn,
-      bs_stride_n, reduce_params) = c
-  
-     # B views: must recompute since B changes between phases
-     # .view() and .T are very cheap on contiguous tensors
-     B_q_T = B_q.view(_uint8).T
-     B_scale_u8 = B_scale_sh.view(_uint8)
-  
-     # Launch fused kernel - use pre-computed grid
-     _fused_quant_gemm_kernel[grid_tuple](
-         A, B_q_T, out_tensor, B_scale_u8,
-         M, N, K,
-         A.stride(0), A.stride(1),
-         B_q_T.stride(0), B_q_T.stride(1),
-         stride_ck, stride_cm, stride_cn,
-         bs_stride_n, 1,
-         BLOCK_SIZE_M=BSM,
-         BLOCK_SIZE_N=BSN,
-         BLOCK_SIZE_K=BSK,
-         GROUP_SIZE_M=8,
-         NUM_KSPLIT=NUM_KSPLIT,
-         SPLITK_BLOCK_SIZE=SPLITK_BLOCK_SIZE,
-         QUANT_BLOCK=32,
-         num_warps=8,
-         num_stages=2,
-         waves_per_eu=0,
-     )
-  
-     if reduce_params is not None:
-         grid_reduce, ACTUAL_KSPLIT, np2_ksplit, s0, s1, s2, sy0, sy1 = reduce_params
-         _reduce_kernel[grid_reduce](
-             y_pp, y, M, N,
-             s0, s1, s2,
-             sy0, sy1,
-             16, 64, ACTUAL_KSPLIT, np2_ksplit,
+         _fused_quant_gemm_kernel[(total_blocks,)](
+             A, B_q_T, out_tensor, B_scale_u8,
+             M, N, K,
+             A.stride(0), A.stride(1),
+             B_q_T.stride(0), B_q_T.stride(1),
+             0 if NUM_KSPLIT == 1 else y_pp.stride(0),
+             y.stride(0) if NUM_KSPLIT == 1 else y_pp.stride(1),
+             y.stride(1) if NUM_KSPLIT == 1 else y_pp.stride(2),
+             bs_stride_n, bs_stride_k,
+             BLOCK_SIZE_M=BLOCK_SIZE_M,
+             BLOCK_SIZE_N=BLOCK_SIZE_N,
+             BLOCK_SIZE_K=BLOCK_SIZE_K,
+             GROUP_SIZE_M=8,
+             NUM_KSPLIT=NUM_KSPLIT,
+             SPLITK_BLOCK_SIZE=SPLITK_BLOCK_SIZE,
+             QUANT_BLOCK=32,
+             num_warps=8,
+             num_stages=2,
+             waves_per_eu=0,
+             allow_flush_denorm=True,
          )
   
-     return y
+         if NUM_KSPLIT > 1:
+             ACTUAL_KSPLIT = triton.cdiv(K_packed, (SPLITK_BLOCK_SIZE // 2))
+             grid_reduce = (triton.cdiv(M, 16), triton.cdiv(N, 64))
+             _reduce_kernel[grid_reduce](
+                 y_pp, y, M, N,
+                 y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
+                 y.stride(0), y.stride(1),
+                 16, 64, ACTUAL_KSPLIT,
+                 triton.next_power_of_2(NUM_KSPLIT),
+             )
   
+         return y
   
- @torch.no_grad()
- def _hot_asm(data):
-     """ASM GEMM hot path for M>64. Absolute minimum Python overhead."""
-     global _last_asm_key, _last_asm_cfg
-     A, B, B_q, B_shuffle, B_scale_sh = data
-     M = A.shape[0]
-     K = A.shape[1]
-     N = B_shuffle.shape[0]
-  
-     key = (M, K, N)
-     if key is _last_asm_key or key == _last_asm_key:
-         c = _last_asm_cfg
      else:
-         c = _asm_cache.get(key)
+         # --- ASM GEMM path ---
+         key = (M, K, N)
+         c = _cache_asm.get(key)
          if c is None:
-             c = _build_asm_config(M, K, N, A, B_shuffle, B_scale_sh)
-             _asm_cache[key] = c
-         _last_asm_key = key
-         _last_asm_cfg = c
+             scale_n_valid = (K + 31) // 32
+             SCALE_M = ((M + 255) // 256) * 256
+             SCALE_N = ((scale_n_valid + 7) // 8) * 8
+             padded_m = get_padded_m(M, N, K, 0)
   
-     (scale_n_valid, SCALE_N, BSM, BSN, NUM_ITER_Q, grid,
-      x_fp4, bs_sh, out, x_fp4_view, bs_sh_view, out_view,
-      kernel_name, split_k,
-      stride_a0, stride_a1, stride_fp4_0, stride_fp4_1) = c
+             BSM = triton.next_power_of_2(M) if M <= 32 else 16
+             NW = 1
+             BSN = 32
+             NUM_ITER_Q = 2
+             grid = (triton.cdiv(M, BSM), triton.cdiv(K, BSN * NUM_ITER_Q))
   
-     _fused_quant_shuffle_kernel[grid](
-         A, x_fp4, bs_sh,
-         stride_a0, stride_a1,
-         stride_fp4_0, stride_fp4_1,
-         M, K, scale_n_valid,
-         SCALE_N=SCALE_N, BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=BSN,
-         NUM_ITER=NUM_ITER_Q, NUM_STAGES=NUM_ITER_Q, MXFP4_QUANT_BLOCK_SIZE=32,
-         num_warps=1, waves_per_eu=0, num_stages=NUM_ITER_Q,
-     )
+             ck_config = get_GEMM_config(M, N, K)
+             kernel_name = ""
+             split_k = 0
+             if ck_config is not None:
+                 split_k = ck_config.get("splitK", 0) or 0
+                 kernel_name = ck_config["kernelName"]
   
-     _gemm_asm(x_fp4_view, B_shuffle, bs_sh_view, B_scale_sh,
-               out, kernel_name, None, 1.0, 0.0, True, split_k)
-     return out_view
+             x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=A.device)
+             bs_sh = torch.full((SCALE_M, SCALE_N), 127, dtype=torch.uint8, device=A.device)
+             out = torch.empty((padded_m, N), dtype=torch.bfloat16, device=A.device)
   
+             x_fp4_view = x_fp4.view(_fp4x2)
+             bs_sh_view = bs_sh.view(_fp8_e8m0)
+             out_view = out[:M] if M < padded_m else out
   
- @torch.no_grad()
- def _hot_dispatch(data):
-     """Dispatch to fused or ASM based on M. No warmup check."""
-     M = data[0].shape[0]
-     if M <= 64:
-         return _hot_fused(data)
-     else:
-         return _hot_asm(data)
+             c = (scale_n_valid, SCALE_N, BSM, BSN, NUM_ITER_Q, grid,
+                  x_fp4, bs_sh, out, x_fp4_view, bs_sh_view, out_view,
+                  kernel_name, split_k,
+                  A.stride(0), A.stride(1), x_fp4.stride(0), x_fp4.stride(1))
+             _cache_asm[key] = c
   
+         (scale_n_valid, SCALE_N, BSM, BSN, NUM_ITER_Q, grid,
+          x_fp4, bs_sh, out, x_fp4_view, bs_sh_view, out_view,
+          kernel_name, split_k,
+          stride_a0, stride_a1, stride_fp4_0, stride_fp4_1) = c
   
- def _warmup_kernel(data):
-     """Warmup path: initializes aiter, then swaps custom_kernel to hot path."""
-     global custom_kernel, _gemm_asm
+         _fused_quant_shuffle_kernel[grid](
+             A, x_fp4, bs_sh,
+             stride_a0, stride_a1,
+             stride_fp4_0, stride_fp4_1,
+             M, K, scale_n_valid,
+             SCALE_N=SCALE_N, BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=BSN,
+             NUM_ITER=NUM_ITER_Q, NUM_STAGES=NUM_ITER_Q, MXFP4_QUANT_BLOCK_SIZE=32,
+             num_warps=1, waves_per_eu=0, num_stages=NUM_ITER_Q,
+         )
   
-     A, B, B_q, B_shuffle, B_scale_sh = data
-     M, K = A.shape
-     N = B_shuffle.shape[0]
+         if _gemm_asm is not None:
+             _gemm_asm(x_fp4_view, B_shuffle, bs_sh_view, B_scale_sh,
+                       out, kernel_name, None, 1.0, 0.0, True, split_k)
+             return out_view
   
-     scale_n_valid = (K + 31) >> 5
-     SCALE_M = ((M + 255) // 256) * 256
-     SCALE_N = ((scale_n_valid + 7) >> 3) << 3
-     BSM = _triton_np2(M) if M <= 32 else 16
-     grid = (_triton_cdiv(M, BSM), _triton_cdiv(K, 32))
-  
-     x_fp4 = _torch_empty((M, K >> 1), dtype=_uint8, device=A.device)
-     bs_sh = _torch_full((SCALE_M, SCALE_N), 127, dtype=_uint8, device=A.device)
-  
-     _fused_quant_shuffle_kernel[grid](
-         A, x_fp4, bs_sh,
-         A.stride(0), A.stride(1),
-         x_fp4.stride(0), x_fp4.stride(1),
-         M, K, scale_n_valid,
-         SCALE_N=SCALE_N, BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=32,
-         NUM_ITER=1, NUM_STAGES=1, MXFP4_QUANT_BLOCK_SIZE=32,
-         num_warps=1, waves_per_eu=0, num_stages=1,
-     )
-  
-     result = aiter.gemm_a4w4(
-         x_fp4.view(_fp4x2), B_shuffle,
-         bs_sh.view(_fp8_e8m0), B_scale_sh,
-         dtype=_bf16, bpreshuffle=True,
-     )
-  
-     # Get ASM function
-     try:
-         _gemm_asm = torch.ops.aiter.gemm_a4w4_asm
-     except Exception:
-         try:
-             import aiter.jit.core as _jc
-             _gemm_asm = getattr(_jc, 'gemm_a4w4_asm', None)
-         except Exception:
-             pass
-  
-     # CRITICAL: swap custom_kernel to the hot path
-     # This eliminates the warmup check from ALL future calls
-     custom_kernel = _hot_dispatch
-  
-     return result
-  
-  
- # Start with warmup - gets swapped to _hot_dispatch after first call
- custom_kernel = _warmup_kernel
+         return aiter.gemm_a4w4(
+             x_fp4_view, B_shuffle, bs_sh_view, B_scale_sh,
+             dtype=_bf16, bpreshuffle=True,
+         )