[FlyDSL] Cache compiled FlyDSL launchers on host paths

aiter/ops/flydsl/batched_gemm_mxfp4.py

@@ -15,7 +15,7 @@
 
 from aiter.jit.utils.chip_info import get_gfx
 
-from .kernels.tensor_shim import ptr_arg
+from .kernels.tensor_shim import _run_compiled, ptr_arg
 
 SCALE_GROUP_SIZE = 32
 WMMA_K_GFX1250 = 128
@@ -437,12 +437,13 @@ def flydsl_batched_gemm_a8w4_v2(
     )
 
     a_c = a.contiguous()
-    _launch_v2(
-        out_phys,
+    _run_compiled(
+        _launch_v2,
+        out_phys.view(-1),
         ptr_arg(a_c),
         ptr_arg(w),
-        a_scales.view(torch.int32),
-        w_scales.view(torch.int32),
+        a_scales.view(torch.int32).view(-1),
+        w_scales.view(torch.int32).view(-1),
         M,
         torch.cuda.current_stream(),
         N,
@@ -459,5 +460,19 @@ def flydsl_batched_gemm_a8w4_v2(
         a_is_fp4,
         cluster_m,
         cluster_n,
+        specialization_key=(
+            tile_m,
+            tile_n,
+            tile_k,
+            m_warp,
+            n_warp,
+            out_is_f16,
+            B,
+            layout_mbn,
+            num_buffers,
+            a_is_fp4,
+            cluster_m,
+            cluster_n,
+        ),
     )
     return out_phys.transpose(0, 1) if layout == "mbn" else out_phys

aiter/ops/flydsl/bpreshuffle_gemm_gfx1250.py

@@ -163,8 +163,29 @@ def run_preshuffle_gemm_a8_gfx1250(
     gemm_out = Out if partials is None else partials
     out_is_f16 = 1 if out_dtype == "f16" else 0
 
+    gemm_spec = (
+        tile_m,
+        tile_n,
+        tile_k,
+        m_warp,
+        n_warp,
+        out_is_f16,
+        nb,
+        cluster_m,
+        cluster_n,
+        False,  # is_mxscale
+        SCALE_BLOCK_SIZE,
+        split_k,
+        False,  # batched
+        0,  # preload_ks
+        False,  # a_preshuffle
+        False,  # a_scale_preshuffled
+        False,  # b_scale_32x32
+        False,  # k_tail
+    )
     stream = _fx.Stream(torch.cuda.current_stream(device=XQ.device))
-    _launch_gemm_a8w8(
+    _run_compiled(
+        _launch_gemm_a8w8,
         _ptr_arg(gemm_out),
         _ptr_arg(XQ),
         _ptr_arg(WQ),
@@ -177,18 +198,10 @@ def run_preshuffle_gemm_a8_gfx1250(
         0,
         lda,
         ldc,
-        tile_m,
-        tile_n,
-        tile_k,
-        m_warp,
-        n_warp,
-        out_is_f16,
-        nb,
-        cluster_m,
-        cluster_n,
-        False,
-        SCALE_BLOCK_SIZE,
-        split_k,
+        *gemm_spec[:14],
+        1,  # batch (runtime)
+        *gemm_spec[14:],
+        specialization_key=gemm_spec,
     )
     if partials is not None:
         dense = ldc == N

aiter/ops/flydsl/gather_kv_b_proj_gfx1250.py

@@ -22,9 +22,9 @@ def _projection_runtime():
     from .kernels.gather_kv_b_proj_gfx1250 import (
         launch_gather_kv_b_proj_8wave_gfx1250,
     )
-    from .kernels.tensor_shim import ptr_arg
+    from .kernels.tensor_shim import _run_compiled, ptr_arg
 
-    return launch_gather_kv_b_proj_8wave_gfx1250, ptr_arg, fx
+    return launch_gather_kv_b_proj_8wave_gfx1250, ptr_arg, fx, _run_compiled
 
 
 def _cluster_n(m: int, n_heads: int) -> int:
@@ -45,7 +45,7 @@ def _run_projection(
     cluster_n: int,
 ) -> None:
     """Launch the fixed-width indexed MLA WMMA kernel."""
-    launch_projection, ptr_arg, fx = _projection_runtime()
+    launch_projection, ptr_arg, fx, run_compiled = _projection_runtime()
 
     scale_is_ptr = activation_scale.device.type != "cpu"
     scale_ptr = activation_scale if scale_is_ptr else cache
@@ -62,7 +62,13 @@ def _run_projection(
         0.0 if scale_is_ptr else float(activation_scale),
         fx.Stream(torch.cuda.current_stream(device=cache.device)),
     )
-    launch_projection(*args, n_heads, int(scale_is_ptr), cluster_n)
+    projection_spec = (n_heads, int(scale_is_ptr), cluster_n)
+    run_compiled(
+        launch_projection,
+        *args,
+        *projection_spec,
+        specialization_key=projection_spec,
+    )
 
 
 def unsupported_reason(

aiter/ops/flydsl/grouped_gemm_mxfp4.py

@@ -10,7 +10,7 @@
 import torch
 
 from .kernels.mega_moe_gfx1250.types import Stage2ScatterContext
-from .kernels.tensor_shim import ptr_arg
+from .kernels.tensor_shim import _run_compiled, ptr_arg
 
 _SUPPORTED_CLUSTER_M = (1, 2, 4)
 _SUPPORTED_CLUSTER_N = (1, 2, 3, 4)
@@ -160,7 +160,11 @@ def flydsl_grouped_gemm_a8w4_masked(
     num_buffers = min(num_buffers, max(1, k_tiles))
     has_bias = 1 if bias is not None else 0
     bias_ptr = ptr_arg(bias) if bias is not None else ptr_arg(a)
-    quant_scale_tensor = out if quant_scale is None else quant_scale.view(torch.uint8)
+    # These tensor arguments only supply base addresses, so normalize their
+    # ranks without specializing the compiled callable on buffer shapes.
+    quant_scale_tensor = (
+        out if quant_scale is None else quant_scale.view(torch.uint8)
+    ).view(-1)
     n_tiles = (N + tile_n - 1) // tile_n
     cluster_n = _select_cluster_n(n_tiles, cluster_n)
     cluster_m = _select_cluster_m(cluster_m, cluster_n, stage1_act)
@@ -205,16 +209,8 @@ def flydsl_grouped_gemm_a8w4_masked(
                 f"{stage2_scatter.combine_quant_bits} needs N % {n_align} == 0 to "
                 f"keep the scale plane cache-line aligned, got N={N}"
             )
-    ep_row_map_tensor = ep_row_map if ep_row_map is not None else out
-    launch_gemm_a8w4_tdm(
-        out,
-        ptr_arg(a),
-        ptr_arg(w),
-        a_scales.view(torch.int32),
-        w_scales.view(torch.int32),
-        contiguous_m,
-        stream,
-        N,
+    ep_row_map_tensor = (ep_row_map if ep_row_map is not None else out).view(-1)
+    gemm_spec = (
         K,
         tile_m,
         tile_n,
@@ -224,42 +220,64 @@ def flydsl_grouped_gemm_a8w4_masked(
         out_is_f16,
         num_buffers,
         a_is_fp4,
-        ptr_arg(m_tile_map),
-        n_experts,
-        stage1_act,
-        has_bias,
-        bias_ptr,
-        float(swiglu_limit),
-        stage1_quant_out,
-        quant_wmma_rep,
-        quant_scale_tensor,
+    )
+    route_spec = (n_experts, stage1_act, has_bias)
+    quant_spec = (stage1_quant_out, quant_wmma_rep)
+    tdm_spec = (
         cluster_n,
         cluster_m,
         _select_bool_env("AITER_TDM_NEXT_STAGE_PREFETCH", next_stage_prefetch),
         waves_per_tensor_tdm,
         _select_bool_env("AITER_GROUPED_GEMM_AS_PROLOGUE", tdm_as_in_prologue),
         _select_tdm_b_th(tdm_b_th),
-        enable_ep_scatter=int(enable_ep_scatter),
-        ep_arena_handle=(int(stage2_scatter.arena_handle) if enable_ep_scatter else 0),
-        ep_combine_input_offset=(
-            int(stage2_scatter.combine_input_offset) if enable_ep_scatter else 0
-        ),
-        ep_slot_stride_bytes=(
-            int(stage2_scatter.slot_stride_bytes) if enable_ep_scatter else 0
-        ),
-        ep_destination_stride=int(ep_destination_stride),
-        ep_world_size=int(stage2_scatter.world_size) if enable_ep_scatter else 0,
-        ep_quant_bits=(
-            int(stage2_scatter.combine_quant_bits) if enable_ep_scatter else 0
+    )
+    ep_spec = (
+        int(enable_ep_scatter),
+        int(stage2_scatter.arena_handle) if enable_ep_scatter else 0,
+        int(stage2_scatter.combine_input_offset) if enable_ep_scatter else 0,
+        int(stage2_scatter.slot_stride_bytes) if enable_ep_scatter else 0,
+        int(ep_destination_stride),
+        int(stage2_scatter.world_size) if enable_ep_scatter else 0,
+        int(stage2_scatter.combine_quant_bits) if enable_ep_scatter else 0,
+    )
+    layout_spec = (
+        int(row_major_ascale),
+        int(a_row_stride_bytes),
+        int(a_scale_row_stride_bytes),
+        int(lds_soa_load_interleave),
+        persistent_workers,
+        rolled_drain,
+    )
+    _run_compiled(
+        launch_gemm_a8w4_tdm,
+        out.view(-1),
+        ptr_arg(a),
+        ptr_arg(w),
+        a_scales.view(torch.int32).view(-1),
+        w_scales.view(torch.int32).view(-1),
+        contiguous_m,
+        stream,
+        N,
+        *gemm_spec,
+        ptr_arg(m_tile_map),
+        *route_spec,
+        bias_ptr,
+        float(swiglu_limit),
+        *quant_spec,
+        quant_scale_tensor,
+        *tdm_spec,
+        *ep_spec,
+        ep_row_map_tensor,
+        float(situ_beta),
+        float(situ_linear_beta),
+        *layout_spec,
+        specialization_key=(
+            gemm_spec,
+            route_spec,
+            quant_spec,
+            tdm_spec,
+            ep_spec,
+            layout_spec,
         ),
-        arg_ep_row_map=ep_row_map_tensor,
-        f32_situ_beta=float(situ_beta),
-        f32_situ_linear_beta=float(situ_linear_beta),
-        row_major_ascale=int(row_major_ascale),
-        a_row_stride_bytes=int(a_row_stride_bytes),
-        a_scale_row_stride_bytes=int(a_scale_row_stride_bytes),
-        lds_soa_load_interleave=int(lds_soa_load_interleave),
-        persistent_workers=persistent_workers,
-        rolled_drain=rolled_drain,
     )
     return out