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