Optimize qk norm rope quant FlyDSL launch path

aiter/aot/flydsl/moe.py

@@ -727,6 +727,25 @@ def _make_a_user(a_dtype_user_shape):
             )
             _run_compiled(exe, args)
 
+            # Reduce mode (accumulate=False) runs a separate topk reduction
+            # kernel inside the runtime stage2 wrapper. Precompile it via the
+            # same shared helper the runtime uses so the cache key matches.
+            # Single-GPU path uses use_mask=False (plain); EP/masked reduction
+            # is a multi-GPU path (separately gated) and not covered here.
+            if not accumulate:
+                from aiter.ops.flydsl.moe_kernels import _run_moe_reduction
+
+                _run_moe_reduction(
+                    target,
+                    out,
+                    tokens,
+                    topk,
+                    model_dim,
+                    expert_mask=None,
+                    topk_ids=None,
+                    stream=0,
+                )
+
 
 def compile_one_config(
     kernel_name: str,

aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py

@@ -48,6 +48,7 @@
 from flydsl.expr.typing import T, Vector as Vec
 from flydsl.expr.utils.arith import ArithValue, _to_raw as _raw
 from .kernels_common import dtype_to_elem_type
+from .tensor_shim import _run_compiled
 from flydsl.runtime.device import get_rocm_arch as get_hip_arch
 from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
 from flydsl._mlir import ir
@@ -777,23 +778,24 @@ def _wrap_qkvo(args, kwargs):
                 kwargs[name] = _ptr_arg(kwargs[name])
         return tuple(args), kwargs
 
+    launch_flash_attn_func.compile_hints = dict(_fmha_compile_hints)
+
     def _launch(*args, **kwargs):
         args, kwargs = _wrap_qkvo(args, kwargs)
-        with CompilationContext.compile_hints(_fmha_compile_hints):
-            return launch_flash_attn_func(*args, **kwargs)
+        stream = kwargs.pop("stream", fx.Stream(None))
+        _run_compiled(launch_flash_attn_func, *args, stream)
 
     def _compile(Q, K, V, O, batch_size, seq_len, stream=None):  # noqa: E741
-        with CompilationContext.compile_hints(_fmha_compile_hints):
-            return flyc.compile(
-                launch_flash_attn_func,
-                _ptr_arg(Q),
-                _ptr_arg(K),
-                _ptr_arg(V),
-                _ptr_arg(O),
-                batch_size,
-                seq_len,
-                fx.Stream(stream),
-            )
+        return flyc.compile(
+            launch_flash_attn_func,
+            _ptr_arg(Q),
+            _ptr_arg(K),
+            _ptr_arg(V),
+            _ptr_arg(O),
+            batch_size,
+            seq_len,
+            fx.Stream(stream),
+        )
 
     _launch.compile = _compile
     return _launch

aiter/ops/flydsl/kernels/fmha_gfx1250/fmha_kernel.py

@@ -56,6 +56,7 @@
 from flydsl.expr.rocdl import tdm_ops
 from flydsl.expr.typing import T
 from flydsl.utils.smem_allocator import SmemAllocator
+from ..tensor_shim import _run_compiled
 from flydsl.compiler.kernel_function import (
     CompilationContext,
 )
@@ -3327,15 +3328,6 @@ def _launch(
     _launch_fns[key] = _launch
 
 
-def _run_compiled(exe, args):
-    cf = getattr(exe, "_cf", None)
-    if cf is None:
-        cf = flyc.compile(exe, *args)
-        exe._cf = cf
-    else:
-        cf(*args)
-
-
 def flash_attn_varlen_d192_gfx1250(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -3394,30 +3386,28 @@ def flash_attn_varlen_d192_gfx1250(
 
     _run_compiled(
         _launch_fns[(bool(causal), bool(return_lse))],
-        (
-            out,
-            q,
-            k,
-            v,
-            lse,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            softmax_scale,
-            stride_q_seq,
-            stride_k_seq,
-            stride_v_seq,
-            stride_o_seq,
-            stride_q_head,
-            stride_k_head,
-            stride_v_head,
-            stride_o_head,
-            gqa,
-            max_seqlen_q,
-            max_seqlen_k,
-            nheads_q,
-            batch,
-            torch.cuda.current_stream(),
-        ),
+        out,
+        q,
+        k,
+        v,
+        lse,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        softmax_scale,
+        stride_q_seq,
+        stride_k_seq,
+        stride_v_seq,
+        stride_o_seq,
+        stride_q_head,
+        stride_k_head,
+        stride_v_head,
+        stride_o_head,
+        gqa,
+        max_seqlen_q,
+        max_seqlen_k,
+        nheads_q,
+        batch,
+        torch.cuda.current_stream(),
     )
 
     if return_lse:

aiter/ops/flydsl/kernels/fused_compress_attn.py

@@ -70,7 +70,7 @@
 from flydsl._mlir import ir
 from flydsl._mlir.dialects import llvm, rocdl, scf
 
-from .tensor_shim import _to_raw
+from .tensor_shim import _to_raw, _run_compiled
 
 # --- shape constants --------------------------------------------------------
 BLOCK_THREADS = 64  # 1 wave64; D must be a multiple
@@ -1422,7 +1422,7 @@ def flydsl_fused_compress_attn(
         stream = torch.cuda.current_stream()
     fx_stream = Stream(stream)
 
-    launcher(
+    args = (
         kv_in,
         kv_in.stride(0),
         score_in,
@@ -1447,5 +1447,6 @@ def flydsl_fused_compress_attn(
         bt_arg,
         bt_seq_stride,
         plan_capacity,
-        stream=fx_stream,
+        fx_stream,
     )
+    _run_compiled(launcher, *args)

aiter/ops/flydsl/kernels/fused_compress_attn_hca.py

@@ -59,7 +59,7 @@
 from flydsl.runtime.device import get_rocm_arch
 from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
 
-from .tensor_shim import STensor, _to_raw
+from .tensor_shim import STensor, _to_raw, _run_compiled
 
 # Force-bind LDS-related imports so isort/ruff/format hooks don't drop them
 # (the multi-wave LDS kernel references CompilationContext, STensor,
@@ -1203,7 +1203,7 @@ def flydsl_hca_compress_attn(
         k_split_num_waves=k_split_num_waves,
         slice_size=slice_size,
     )
-    compress_fn(
+    compress_args = (
         kv_in,
         int(kv_in.stride(0)),
         score_in,
@@ -1220,8 +1220,9 @@ def flydsl_hca_compress_attn(
         kv_compressed,
         int(kv_compressed.stride(0)),
         int(plan_capacity),
-        stream=stream_obj,
+        stream_obj,
     )
+    _run_compiled(compress_fn, *compress_args)
 
     rms_weight_is_bf16 = rms_weight.dtype == torch.bfloat16
     norm_fn = compile_hca_norm_rope_scatter(
@@ -1232,7 +1233,7 @@ def flydsl_hca_compress_attn(
         rms_weight_is_bf16=rms_weight_is_bf16,
         rms_eps=rms_eps,
     )
-    norm_fn(
+    norm_args = (
         kv_compressed,
         int(kv_compressed.stride(0)),
         plan_gpu,
@@ -1245,5 +1246,6 @@ def flydsl_hca_compress_attn(
         block_tables,
         int(block_tables.stride(0)),
         int(plan_capacity),
-        stream=stream_obj,
+        stream_obj,
     )
+    _run_compiled(norm_fn, *norm_args)

aiter/ops/flydsl/kernels/moe_gemm_2stage.py

@@ -61,6 +61,7 @@ def bf16_global_atomics_arch_description() -> str:
     crd2idx,
 )
 from .mfma_epilogues import c_shuffle_epilog, default_epilog, mfma_epilog
+from .tensor_shim import _run_compiled
 
 
 @contextmanager
@@ -3972,7 +3973,8 @@ def _ptr_arg(t):
             return flyc.from_c_void_p(fx.Uint8, t.data_ptr())
 
         # Phase 1: GEMM2 (no atomics) -> [tokens*topk, model_dim]
-        self._gemm2_exe(
+        _run_compiled(
+            self._gemm2_exe,
             _ptr_arg(intermediate.view(-1)),
             _ptr_arg(arg_x),
             _ptr_arg(arg_w),
@@ -4002,8 +4004,14 @@ def _ptr_arg(t):
             # Placeholders; kernel ignores them when use_mask=False (compile-time).
             em = torch.empty(0, device=arg_out.device, dtype=torch.int32)
             tk = torch.empty(0, device=arg_out.device, dtype=torch.int32)
-        self._reduce_exe(
-            _ptr_arg(X), _ptr_arg(Y), _ptr_arg(em), _ptr_arg(tk), tokens_in, stream
+        _run_compiled(
+            self._reduce_exe,
+            _ptr_arg(X),
+            _ptr_arg(Y),
+            _ptr_arg(em),
+            _ptr_arg(tk),
+            tokens_in,
+            stream,
         )
 
     @property

aiter/ops/flydsl/kernels/qk_norm_rope_quant.py

@@ -70,7 +70,7 @@
 from flydsl.runtime.device import get_rocm_arch as get_hip_arch
 from flydsl._mlir.dialects import llvm, rocdl
 
-from .tensor_shim import GTensor, _to_raw
+from .tensor_shim import GTensor, _to_raw, _run_compiled
 
 # JIT-free MX-format mode/dtype int mirrors. ``aiter.utility.mx_types``'s
 # pybind11 ``MxScaleRoundMode`` / ``MxDtype`` lazy-load on first attribute
@@ -82,6 +82,29 @@
     MX_DEFAULT_ROUND_MODE as _DEFAULT_MODE,
 )
 
+_STATIC_ADAPTOR_CACHE = {}
+_STATIC_ADAPTOR_CACHE_MAX = 64
+
+
+def _cached_from_dlpack(t: torch.Tensor):
+    key = (
+        int(t.data_ptr()),
+        str(t.device),
+        str(t.dtype),
+        tuple(t.shape),
+        tuple(t.stride()),
+        int(t.storage_offset()),
+    )
+    cached = _STATIC_ADAPTOR_CACHE.get(key)
+    if cached is not None:
+        return cached
+    if len(_STATIC_ADAPTOR_CACHE) >= _STATIC_ADAPTOR_CACHE_MAX:
+        _STATIC_ADAPTOR_CACHE.clear()
+    adaptor = flyc.from_dlpack(t)
+    _STATIC_ADAPTOR_CACHE[key] = adaptor
+    return adaptor
+
+
 # --- shape constants (V4-Pro MVP) -------------------------------------------
 BLOCK_THREADS = 64  # 1 wave64
 
@@ -968,15 +991,24 @@ def flydsl_qk_norm_rope_quant(
 
     if stream is None:
         stream = torch.cuda.current_stream()
-    fx_stream = Stream(stream)
+
+    def _has_direct_state():
+        return getattr(launcher, "_direct_call_state", None) is not None
 
     def _ptr_arg(t):
+        if _has_direct_state():
+            return int(t.data_ptr())
         return flyc.from_c_void_p(fx.Uint8, t.data_ptr())
 
-    q_weight_static = flyc.from_dlpack(q_weight_arg)
-    kv_weight_static = flyc.from_dlpack(kv_weight)
-    cos_static = flyc.from_dlpack(cos_2d)
-    sin_static = flyc.from_dlpack(sin_2d)
+    def _stream_arg():
+        if _has_direct_state():
+            return stream
+        return Stream(stream)
+
+    q_weight_static = _cached_from_dlpack(q_weight_arg)
+    kv_weight_static = _cached_from_dlpack(kv_weight)
+    cos_static = _cached_from_dlpack(cos_2d)
+    sin_static = _cached_from_dlpack(sin_2d)
 
     # HW grid Y is a 16-bit field on AMD HIP → cap 65535 blocks/launch. The
     # kernel uses per-token GTensor base-shift so each chunk's resource span
@@ -992,7 +1024,7 @@ def _ptr_arg(t):
     for start in range(0, T_tok, MAX_GRID_Y):
         n = min(MAX_GRID_Y, T_tok - start)
         end = start + n
-        launcher(
+        args = (
             _ptr_arg(q_view[start:end]),
             _ptr_arg(kv[start:end]),
             q_weight_static,
@@ -1006,7 +1038,8 @@ def _ptr_arg(t):
             _ptr_arg(kv_scale_arg[start:end] if quant else kv_scale_arg),
             kv.stride(0),
             n,
-            stream=fx_stream,
+            _stream_arg(),
         )
+        _run_compiled(launcher, *args)
 
     return q_out, kv_out, (q_scale if quant else None), (kv_scale if quant else None)

aiter/ops/flydsl/linear_attention_prefill_kernels.py

@@ -24,6 +24,7 @@
 import triton
 
 from .kernels.chunk_gated_delta_h import compile_chunk_gated_delta_h
+from .kernels.tensor_shim import _run_compiled
 from ..triton._triton_kernels.gated_delta_rule.utils import (
     prepare_chunk_offsets,
     prepare_num_chunks,
@@ -264,7 +265,8 @@ def _launch_kernel(
 ):
     grid_v = triton.cdiv(V, BV)
     grid_nh = N * H
-    launch_fn(
+    _run_compiled(
+        launch_fn,
         k,
         u,
         w,