From 0a5cfe0433673073a9a782d54026340d529956be Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 16 May 2026 20:45:46 +0000
Subject: [PATCH] =?UTF-8?q?add=20kernel=20compile=20caching=20=E2=80=94=20?=
 =?UTF-8?q?compile=20once,=20invoke=20on=20subsequent=20calls?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

First call: cute.compile() with real tensors (warmup).
Subsequent calls: just invoke compiled() with new CuTe views.
No cute.compile() in the forward path = cudagraph-safe.
---
 cutedsl/bridge.py | 44 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py
index 4d9ba7a0..cc05fbaa 100644
--- a/cutedsl/bridge.py
+++ b/cutedsl/bridge.py
@@ -29,6 +29,9 @@ from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
 
 E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
 
+# Cache compiled kernels by (num_experts, device, K, N)
+_compiled_kernel_cache = {}
+
 # Cached LUT for E2M1 quantization (created once per device, cudagraph-safe)
 _NVFP4_STEP_LUT_CACHE = {}
 _NVFP4_STEP_LUT_LOCK = threading.Lock()
@@ -292,6 +295,10 @@ def run_nvfp4_grouped_gemm(
     """Run the CuTeDSL NVFP4 scaled grouped GEMM.
     
     2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
+    
+    CUDAGraph-compatible: kernel is compiled once on first call (warmup),
+    then only the compiled kernel is invoked on subsequent calls.
+    No torch.cuda.synchronize() or .item() in the forward path.
     """
     num_experts = mat_b.shape[0]
     n_dim = mat_b.shape[2]  # packed N (in float4 elements)
@@ -333,16 +340,31 @@ def run_nvfp4_grouped_gemm(
     max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
-    compiled = cute.compile(
-        kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
-        max_active_clusters, stream,
-        global_scale_a=gsa_c, global_scale_b=gsb_c,
-    )
-
-    compiled(
-        a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
-        stream,
-        global_scale_a=gsa_c, global_scale_b=gsb_c,
-    )
+    # Cache compiled kernel by (num_experts, K, N)
+    cache_key = (num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn,
+                 mat_a.shape[1], mat_b.shape[2])
+    
+    if cache_key not in _compiled_kernel_cache:
+        # First call: compile with real tensors
+        compiled = cute.compile(
+            kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
+            max_active_clusters, stream,
+            global_scale_a=gsa_c, global_scale_b=gsb_c,
+        )
+        # Warmup run
+        compiled(
+            a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
+            stream,
+            global_scale_a=gsa_c, global_scale_b=gsb_c,
+        )
+        _compiled_kernel_cache[cache_key] = compiled
+    else:
+        # Subsequent calls: just invoke the compiled kernel
+        compiled = _compiled_kernel_cache[cache_key]
+        compiled(
+            a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
+            stream,
+            global_scale_a=gsa_c, global_scale_b=gsb_c,
+        )
 
     return out