[ROCm][Perf] Reuse compiled FlyDSL host launchers

vllm/models/deepseek_v4/amd/mono/runner.py

@@ -21,6 +21,8 @@
 
 import torch
 
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
+
 from .common.plan import BLOCKS
 from .dims import (
     EXPERTS,
@@ -359,7 +361,8 @@ def forward(
                 comp_block_table.stride(0),
                 topk_indices,
             )
-        k1(
+        _run_compiled(
+            k1,
             residual.data_ptr(),
             x.data_ptr(),
             post_mix.data_ptr(),
@@ -400,10 +403,11 @@ def forward(
             scratch.data_ptr(),
             self.epoch.data_ptr(),
             self.tl[0].data_ptr() if self.tl is not None else 0,
-            stream=st,
+            st,
         )
         select(self.normed[:M], self.qr[:M])
-        k2(
+        _run_compiled(
+            k2,
             self.q.data_ptr(),
             swa_cache.data_ptr(),
             swa_cache.stride(0),
@@ -458,6 +462,6 @@ def forward(
             self.rank,
             self.epoch.data_ptr(),
             self.tl[1].data_ptr() if self.tl is not None else 0,
-            stream=st,
+            st,
         )
         return outs

vllm/models/deepseek_v41/amd/mono/runner.py

@@ -22,6 +22,8 @@
 
 import torch
 
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
+
 from .attention.front import FRONT_POINTS
 from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, WINDOW, Dims
 from .common.arch import GFX942
@@ -386,7 +388,8 @@ def forward(
         else:
             comp, comp_stride, comp_block = self._dummy, 0, 1
             bt, bt_stride, topk = self._dummy, 0, self._dummy
-        k1(
+        _run_compiled(
+            k1,
             residual.data_ptr(),
             x.data_ptr(),
             post_mix.data_ptr(),
@@ -427,7 +430,7 @@ def forward(
             self.scratch(M).data_ptr(),
             self.epoch.data_ptr(),
             self.tl1.data_ptr() if self.timeline else 0,
-            stream=st,
+            st,
         )
         if select is not None:
             select(self.normed[:M], self.qr[:M])
@@ -450,7 +453,8 @@ def forward(
                     comp_block_table=bt,
                     comp_block=comp_block,
                 )
-        k2(
+        _run_compiled(
+            k2,
             self.q.data_ptr(),
             swa_cache.data_ptr(),
             swa_cache.stride(0),
@@ -505,7 +509,7 @@ def forward(
             self.rank,
             self.epoch.data_ptr(),
             self.tl2.data_ptr() if self.timeline else 0,
-            stream=st,
+            st,
         )
         return outs
 
@@ -533,7 +537,8 @@ def ffn(
         if outs is None:
             outs = self._outs(M, residual)
         out, res_out, post_out, comb_out, pre_out = outs
-        self.ffn_kernel(M, w.gate_w.shape[0], topk)(
+        _run_compiled(
+            self.ffn_kernel(M, w.gate_w.shape[0], topk),
             part.data_ptr(),
             residual.data_ptr(),
             post_mix.data_ptr(),
@@ -564,6 +569,6 @@ def ffn(
             self.rank,
             self.epoch.data_ptr(),
             self.tl3.data_ptr() if self.timeline else 0,
-            stream=torch.cuda.current_stream(),
+            torch.cuda.current_stream(),
         )
         return outs

vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py

@@ -73,6 +73,7 @@
     publish,
     sreg,
 )
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 HIDDEN = 7168
 NA = 56  # attn CTAs
@@ -987,7 +988,8 @@ def kda_pre(
     layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and
     ``layer`` this launch's slot below ``TAG_SLOTS``."""
     f = build(key)
-    f(
+    _run_compiled(
+        f,
         prefix.data_ptr(),
         ptr(delta),
         blocks.data_ptr(),
@@ -1017,5 +1019,5 @@ def kda_pre(
         epoch.data_ptr(),
         health.data_ptr(),
         alarm,
-        stream=torch.cuda.current_stream(),
+        torch.cuda.current_stream(),
     )

vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py

@@ -64,6 +64,7 @@
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv
 from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe as k2b
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 O0, O1 = NA, BLOCKS  # o_proj CTAs: every CTA past AttnRes's
 O_N = O1 - O0  # 200
@@ -540,7 +541,8 @@ def k2_launch(
     """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in
     place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``."""
     f = build(key)
-    f(
+    _run_compiled(
+        f,
         core.data_ptr(),
         w_o.data_ptr(),
         prefix.data_ptr(),
@@ -569,5 +571,5 @@ def k2_launch(
         layer,
         health.data_ptr(),
         alarm,
-        stream=torch.cuda.current_stream(),
+        torch.cuda.current_stream(),
     )

vllm/v1/attention/ops/flydsl_kernels/__init__.py

@@ -1,2 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def _run_compiled(launcher, *args) -> None:
+    """Reuse the compiled callable of a factory-style host launcher."""
+    import flydsl.compiler as flyc
+    from flydsl.utils import env
+
+    if env.compile.compile_only:
+        flyc.compile(launcher, *args)
+        return
+
+    compiled = getattr(launcher, "_cf", None)
+    if compiled is not None:
+        compiled(*args)
+        return
+
+    # flyc.compile executes once; do not launch its result again on this call.
+    launcher._cf = flyc.compile(launcher, *args)

vllm/v1/attention/ops/flydsl_turboquant_decode.py

@@ -29,6 +29,7 @@
 
 from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 logger = init_logger(__name__)
 
@@ -698,7 +699,8 @@ def flydsl_turboquant_decode_attention(
             num_partitions * tile_groups_per_partition * kv_compute_block,
             worst_case_max_seq_len,
         )
-    launch(
+    _run_compiled(
+        launch,
         segm_out,
         segm_sum,
         segm_max,

vllm/v1/attention/ops/flydsl_ultraquant_decode.py

@@ -15,6 +15,7 @@
 
 from vllm.logger import init_logger
 from vllm.triton_utils import triton
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 logger = init_logger(__name__)
 
@@ -773,7 +774,8 @@ def flydsl_ultraquant_decode_attention(
             num_partitions * tile_groups_per_partition * kv_compute_block,
             worst_case_max_seq_len,
         )
-    launch(
+    _run_compiled(
+        launch,
         segm_out,
         segm_sum,
         segm_max,