From deb6b3231a8b61241a1292dc53cb88d439927faf Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 20:01:47 +0000 Subject: [PATCH] debug: swap M/K in SF remap + add printf diagnostics --- .../cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu | 68 +++++++++++- test_sf_layout_diag.cu | 101 ++++++++++++++++++ 2 files changed, 166 insertions(+), 3 deletions(-) create mode 100644 test_sf_layout_diag.cu 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 b880cfb4..860e41f7 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 @@ -135,9 +135,15 @@ __global__ void remap_sf_to_cutlass_kernel( int m = 0, k_sf = 0; - if constexpr (R == 8) { - 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; + if constexpr (R == 6) { + // K-major SfAtom: ((32, 4, k_tiles), (SFVecSize, 4, m_tiles)) + // Flattened: (inner_k, sub_k, tile_k, inner_m, sub_m, tile_m) + k_sf = cute::get<0>(flat) + cute::get<1>(flat) * 32 + cute::get<2>(flat) * 128; + m = cute::get<3>(flat) + cute::get<4>(flat) * InputSFVectorSize + cute::get<5>(flat) * (InputSFVectorSize * 4); + } else if constexpr (R == 8) { + // With batch dimension: ((32, 4, k_tiles), (SFVecSize, 4, m_tiles), (1,)) + k_sf = cute::get<0>(flat) + cute::get<1>(flat) * 32 + cute::get<2>(flat) * 128; + m = cute::get<3>(flat) + cute::get<4>(flat) * InputSFVectorSize + cute::get<5>(flat) * (InputSFVectorSize * 4); } else { m = 0; k_sf = 0; } @@ -147,6 +153,20 @@ __global__ void remap_sf_to_cutlass_kernel( // SFB: source is (K_sf, MN) row-major → src[k_sf * MN + m] dst[dst_idx] = col_major_src ? src[k_sf * MN + m] : src[m * K_sf + k_sf]; } + + // Diagnostic: print first 10 coordinate mappings (host-side, only dst_idx 0..9) + if (dst_idx < 10) { + printf("[SF-REMAP] dst=%d rank=%d", dst_idx, R); + if constexpr (R >= 1) printf(" f0=%d", (int)cute::get<0>(flat)); + if constexpr (R >= 2) printf(" f1=%d", (int)cute::get<1>(flat)); + if constexpr (R >= 3) printf(" f2=%d", (int)cute::get<2>(flat)); + if constexpr (R >= 4) printf(" f3=%d", (int)cute::get<3>(flat)); + if constexpr (R >= 5) printf(" f4=%d", (int)cute::get<4>(flat)); + if constexpr (R >= 6) printf(" f5=%d", (int)cute::get<5>(flat)); + if constexpr (R >= 7) printf(" f6=%d", (int)cute::get<6>(flat)); + if constexpr (R >= 8) printf(" f7=%d", (int)cute::get<7>(flat)); + printf(" -> m=%d k_sf=%d col_major=%d MN=%d K_sf=%d\n", m, k_sf, col_major_src, MN, K_sf); + } } ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -325,4 +345,46 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( } // extern "C" +// Diagnostic: verify SF coordinate mapping +extern "C" int cutlass_nvfp4_gemm_diag_sf_layout( + int M, int N, int K, + int* out_cosize, int* out_size, int* out_rank, + int* out_shape, int* out_stride, // arrays of 8 ints + int* out_coord_map, // flat (dst_idx -> (m, k_sf)) pairs, max 256 entries + int max_entries +) { + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); + + *out_cosize = cute::cosize(layout_SFA); + *out_size = cute::size(layout_SFA); + + auto shape = layout_SFA.shape(); + auto stride = layout_SFA.stride(); + *out_rank = cute::rank(shape); + + // Write shape and stride + for (int i = 0; i < 8; i++) { + out_shape[i] = (i < *out_rank) ? (int)cute::get<0>(cute::slice_and_offset(shape, stride)) : 0; + out_stride[i] = 0; + } + + // Write coordinate mapping + int total = std::min(*out_cosize, max_entries); + for (int i = 0; i < total; i++) { + auto coord = cute::idx2crd(i, layout_SFA.shape(), layout_SFA.stride()); + auto flat = cute::flatten(coord); + constexpr int R = cute::rank_v; + int m = 0, k_sf = 0; + if constexpr (R == 8) { + 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; + } + out_coord_map[i * 2] = m; + out_coord_map[i * 2 + 1] = k_sf; + } + + return total; +} + #endif diff --git a/test_sf_layout_diag.cu b/test_sf_layout_diag.cu new file mode 100644 index 00000000..b7ec0a57 --- /dev/null +++ b/test_sf_layout_diag.cu @@ -0,0 +1,101 @@ +/** Diagnostic: print the SF layout coordinate mapping for verification. + * Compile and run to see what (m, k_sf) each dst_idx maps to. + */ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace cute; + +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; +using ElementD = cutlass::bfloat16_t; +using ElementC = float; +using LayoutCTag = cutlass::layout::RowMajor; +using LayoutDTag = cutlass::layout::RowMajor; +using ElementAccumulator = float; +using ElementCompute = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using MmaTileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; +constexpr int InputSFVectorSize = 16; + +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutCTag, 4, + ElementD, LayoutDTag, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; + +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, 32, + ElementB, LayoutBTag, 32, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto +>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; + +int main() { + // Test with M=1, N=32, K=32 + int M = 1, N = 32, K = 32; + LayoutSFA layout_SFA = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); + + int total = cute::cosize(layout_SFA); + int logical = cute::size(layout_SFA); + printf("M=%d N=%d K=%d: cosize=%d size=%d\n", M, N, K, total, logical); + + // Print shape and stride + auto shape = layout_SFA.shape(); + auto stride = layout_SFA.stride(); + printf("Layout rank: %d\n", (int)cute::rank(shape)); + + // Print first 20 coordinate mappings + int count = std::min(total, 20); + for (int i = 0; i < count; i++) { + auto coord = cute::idx2crd(i, layout_SFA.shape(), layout_SFA.stride()); + auto flat = cute::flatten(coord); + constexpr int R = cute::rank_v; + printf("dst[%d]: rank=%d", i, R); + if constexpr (R >= 1) printf(" f0=%d", (int)cute::get<0>(flat)); + if constexpr (R >= 2) printf(" f1=%d", (int)cute::get<1>(flat)); + if constexpr (R >= 3) printf(" f2=%d", (int)cute::get<2>(flat)); + if constexpr (R >= 4) printf(" f3=%d", (int)cute::get<3>(flat)); + if constexpr (R >= 5) printf(" f4=%d", (int)cute::get<4>(flat)); + if constexpr (R >= 6) printf(" f5=%d", (int)cute::get<5>(flat)); + if constexpr (R >= 7) printf(" f6=%d", (int)cute::get<6>(flat)); + if constexpr (R >= 8) printf(" f7=%d", (int)cute::get<7>(flat)); + + // Compute m and k_sf using the current formula + int m = 0, k_sf = 0; + if constexpr (R == 8) { + 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; + } + printf(" -> m=%d k_sf=%d\n", m, k_sf); + } + + return 0; +}