From caf93d6c45b3055a0c6f8eec55219b660346e599 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 19:59:43 +0000 Subject: [PATCH] fix: pass K_packed/N_packed to _get_compiled_kernel --- cutedsl/bridge.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index 09232f7d..08792612 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -292,11 +292,12 @@ def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"): _compiled_kernel_cache = {} -def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn): +def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn, K_packed, N_packed): """Get or compile the CuTeDSL grouped GEMM kernel (cached by shape config). - The kernel compilation is deterministic for a given (num_experts, device, tiler, cluster) - config, so we cache it to avoid recompiling on every forward call. + The kernel compilation is deterministic for a given config, so we cache it + to avoid recompiling on every forward call. K_packed and N_packed are needed + because the compiled kernel's TMA descriptors are sized from compilation shapes. """ cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn, K_packed, N_packed) if cache_key in _compiled_kernel_cache: @@ -317,11 +318,8 @@ def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn): max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - # We need dummy tensors to compile against — shapes must match runtime tensors - # The compiled kernel uses mark_layout_dynamic but TMA descriptors - # are sized based on the compilation shapes - K_packed = mat_a.shape[1] # actual K packed dimension - N_packed = mat_b.shape[2] # actual N dimension + # Dummy tensors for compilation — use actual K/N dimensions + # (TMA descriptors are sized from compilation shapes) tokens = 1 dummy_a = torch.zeros(tokens, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) @@ -396,7 +394,7 @@ def run_nvfp4_grouped_gemm( out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device) compiled, kernel, max_active_clusters = _get_compiled_kernel( - num_experts, device, mma_tiler_mn, cluster_shape_mn + num_experts, device, mma_tiler_mn, cluster_shape_mn, K_packed, N_packed ) # Convert to CuTe tensors with dynamic layout