From 1e0cea055c98cfc9bfc9aa3960ea6f5420f8b199 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 14 May 2026 16:57:32 +0000
Subject: [PATCH] cleanup: remove all debug printfs from CUDA kernel and
 weight_transform

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.
---
 .../cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu  | 48 ++-----------------
 src/nvfp4_megamoe_kernel/weight_transform.py  | 18 -------
 2 files changed, 4 insertions(+), 62 deletions(-)

diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
index c4ab503f..fcfc28dd 100644
--- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
+++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
@@ -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) {
diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py
index 2c7ed6c2..893ea0ee 100644
--- a/src/nvfp4_megamoe_kernel/weight_transform.py
+++ b/src/nvfp4_megamoe_kernel/weight_transform.py
@@ -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.