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,