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

debug: swap M/K in SF remap + add printf diagnostics

Browse Source
This commit is contained in:
biondizzle
2026-05-15 20:01:47 +00:00
parent 22f0457ccf
commit deb6b3231a
2 changed files with 166 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<i>(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<decltype(flat)>;
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

101
test_sf_layout_diag.cu Normal file
View File

@@ -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 <cstdio>
#include <cutlass/cutlass.h>
#include <cute/tensor.hpp>
#include <cutlass/float_subbyte.h>
#include <cutlass/gemm/dispatch_policy.hpp>
#include <cutlass/gemm/collective/collective_builder.hpp>
#include <cutlass/epilogue/collective/collective_builder.hpp>
#include <cutlass/gemm/device/gemm_universal_adapter.h>
#include <cutlass/gemm/kernel/gemm_universal.hpp>
using namespace cute;
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using LayoutATag = cutlass::layout::RowMajor;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>,
CollectiveMainloop,
CollectiveEpilogue,
void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
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<decltype(flat)>;
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;
}