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

cleanup: remove all debug printfs from CUDA kernel and weight_transform

Browse Source 
Removed printf from remap kernel (flat_rank dump, coordinate probes,
first-coord log). Removed weight_scale_2 debug prints from
weight_transform.py. Production-ready now.
This commit is contained in:
biondizzle
2026-05-14 16:57:32 +00:00
parent 839835cba4
commit 1e0cea055c
2 changed files with 4 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
View File

@@ -126,61 +126,21 @@ __global__ void remap_sf_to_cutlass_kernel(
constexpr int R = cute::rank_v<decltype(flat)>;
// Debug: print rank once (only thread 0)
if (dst_idx == 0) {
printf("[remap] flat_rank=%d MN=%d K_sf=%d total=%d\n", R, MN, K_sf, total);
}
// Debug: single-thread print at specific indices
// Only thread 0 prints, and only at specific dst_idx values
// This avoids garbled output from concurrent threads
if (threadIdx.x == 0 && blockIdx.x == 0) {
// Print a series of indices to understand the decomposition
int debug_indices[] = {0, 1, 4, 16, 64, 128, 256, 511, 512, 513, 516, 1024, 2048, 4096, 8192, 65536, 131072};
int n_debug = sizeof(debug_indices) / sizeof(debug_indices[0]);
for (int di = 0; di < n_debug; di++) {
int didx = debug_indices[di];
if (didx < total) {
auto dcoord = cute::idx2crd(didx, layout_sf.shape(), layout_sf.stride());
auto dflat = cute::flatten(dcoord);
printf("[remap] idx=%d f=", didx);
if constexpr (R >= 8) {
printf("%d,%d,%d,%d,%d,%d,%d,%d",
int(cute::get<0>(dflat)), int(cute::get<1>(dflat)),
int(cute::get<2>(dflat)), int(cute::get<3>(dflat)),
int(cute::get<4>(dflat)), int(cute::get<5>(dflat)),
int(cute::get<6>(dflat)), int(cute::get<7>(dflat)));
}
printf("\n");
}
}
}
int m = 0, k_sf = 0;
if constexpr (R == 8) {
// 8 flattened coordinates from idx2crd:
// f0 = inner_m (0..31) — varies fastest within M atom
// f1 = sub_m (0..3) — second M sub-coordinate
// f2 = tile_m (0..) — M tile index
// f3 = step_m stride (degenerate — always equals total) — NOT a coordinate
// f4 = sub_k (0..3) — K sub-coordinate within atom
// f5 = tile_k (0..) — K tile index
// f6, f7 = 0 (unused higher tiling)
// f0 = inner_m (0..31), f1 = sub_m (0..3), f2 = tile_m (0..)
// f3 = step_m stride (degenerate — always equals total, not a coordinate)
// f4 = sub_k (0..3), f5 = tile_k (0..), f6 = 0, f7 = 0
//
// CuTe "first sub varies fastest" for Shape<32, 4>:
// m = f0 + f1 * 32 + f2 * 128
// k_sf = f4 + f5 * 4
// k_sf = f4 + f5 * 4
m = cute::get<0>(flat) + cute::get<1>(flat) * 32 + cute::get<2>(flat) * 128;
k_sf = cute::get<4>(flat) + cute::get<5>(flat) * 4;
} else {
// Fallback: index 0 and 1
m = 0; k_sf = 0;
if (dst_idx == 0) printf("[remap] UNEXPECTED flat_rank=%d\n", R);
}
if (dst_idx == 0) {
printf("[remap] first coord: m=%d k_sf=%d (src_idx would be %d)\n", m, k_sf, m * K_sf + k_sf);
}
if (m < MN && k_sf < K_sf) {

18
src/nvfp4_megamoe_kernel/weight_transform.py
View File

@@ -75,25 +75,7 @@ def transform_nvfp4_weights_for_mega_moe(
l1_weight, l1_weight_scale = l1_tuple
l2_weight, l2_weight_scale = l2_tuple
# Debug: verify weight_scale_2 shape and values
if l1_weight_scale_2 is not None:
l1_gs = l1_weight_scale_2
print(f"[weight_transform] L1 weight_scale_2: shape={l1_gs.shape} dtype={l1_gs.dtype} "
f"min={l1_gs.min().item():.6f} max={l1_gs.max().item():.6f} "
f"numel={l1_gs.numel()}")
# Check if all experts have the same global scale (expected for ModelOpt NVFP4)
if l1_gs.numel() > 1:
unique = l1_gs.flatten().unique()
print(f"[weight_transform] L1 weight_scale_2 unique values: {len(unique)} "
f"samples: {l1_gs.flatten()[:8].tolist()}")
if l2_weight_scale_2 is not None:
l2_gs = l2_weight_scale_2
print(f"[weight_transform] L2 weight_scale_2: shape={l2_gs.shape} dtype={l2_gs.dtype} "
f"min={l2_gs.min().item():.6f} max={l2_gs.max().item():.6f} "
f"numel={l2_gs.numel()}")
# Fold global scales into block scales
# Both L1 and L2 use per-expert global scales (shape (E,1) or (E,)).
# The logical_widths branch was wrong: it treated gs as per-projection
# scalars and only used experts 0 and 1's scales for ALL experts.
# The else branch correctly broadcasts each expert's own global scale.