From 58dc36e21c6643e4481a5ccf835b22d2ac9885ad Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 20:28:15 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20compile=20fresh=20each=20call=20?= =?UTF-8?q?=E2=80=94=20cached=20compile=20produces=20wrong=20TMA=20descrip?= =?UTF-8?q?tors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CuTeDSL kernel's TMA descriptors are bound to the compilation-time tensor addresses. Caching the compiled kernel and reusing it with different tensor allocations produces wrong memory access patterns (cosine 0.5 instead of 0.99). Fresh compilation is proven correct (cosine 0.989). We can optimize later with proper TMA descriptor reinitialization. --- cutedsl/bridge.py | 77 +++++++++++++++-------------------------------- 1 file changed, 25 insertions(+), 52 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index 7d46e8e1..3ea58ebb 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -318,19 +318,29 @@ def run_nvfp4_grouped_gemm( 2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N) - Compiles with real tensors (not dummy shapes) because the CuTeDSL - kernel's TMA descriptors are sized from compilation-time tensor shapes. - Compilation is cached per (num_experts, K, N) and happens once. + Compiles with real tensors each call. The CuTeDSL kernel's TMA + descriptors are bound to the compilation-time tensor addresses, + so caching across different tensor allocations produces wrong results. The forward call (after compilation) is cudagraph-safe. """ num_experts = mat_b.shape[0] - n_dim = mat_b.shape[2] # N dimension + n_dim = mat_b.shape[2] tokens_sum = mat_a.shape[0] device = mat_a.device out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device) + kernel = ScaledGroupedGemmKernel( + scenario="2Dx3D", + sf_vec_size=SF_VEC_SIZE, + accumulate_on_output=False, + separate_tensormap_init=True, + consistent_token_padding=False, + mma_tiler_mnk=(*mma_tiler_mn, 256), + cluster_shape_mnk=(*cluster_shape_mn, 1), + ) + # Convert to CuTe tensors with dynamic layout def to_cute(t): ct = cutlass_torch.from_dlpack(t) @@ -343,60 +353,23 @@ def run_nvfp4_grouped_gemm( c_c = to_cute(out) offs_c = to_cute(expert_offsets) - workspace_size = 0 + workspace_size = kernel.get_workspace_size(num_experts) + workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device) + ws_c = to_cute(workspace) + gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None import cuda.bindings.driver as cuda + cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - # Check cache — compile with real tensors on first use per (experts, K, N) - K_packed = mat_a.shape[1] - N_packed = mat_b.shape[2] - cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn, K_packed, N_packed) - use_cache = True # TEMP: set False to always recompile (debug) - - if not (use_cache and cache_key in _compiled_kernel_cache): - kernel = ScaledGroupedGemmKernel( - scenario="2Dx3D", - sf_vec_size=SF_VEC_SIZE, - accumulate_on_output=False, - separate_tensormap_init=True, - consistent_token_padding=False, - mma_tiler_mnk=(*mma_tiler_mn, 256), - cluster_shape_mnk=(*cluster_shape_mn, 1), - ) - - cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] - max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) - - workspace_size = kernel.get_workspace_size(num_experts) - workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device) - ws_c = to_cute(workspace) - - 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, - ) - - # Warm up with real data - 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, - ) - torch.cuda.synchronize() - - _compiled_kernel_cache[cache_key] = (compiled, kernel, max_active_clusters) - - compiled, kernel, max_active_clusters = _compiled_kernel_cache[cache_key] - - # Allocate workspace if not already done during compilation - if workspace_size == 0: - workspace_size = kernel.get_workspace_size(num_experts) - workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device) - ws_c = to_cute(workspace) + 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,