biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: pass K_packed/N_packed to _get_compiled_kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-16 19:59:43 +00:00
parent ecc7b83334
commit caf93d6c45
1 changed files with 7 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
cutedsl/bridge.py
View File

@@ -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