Add scale factor remap kernel: remap simple row-major SFs to CUTLASS interleaved layout
This commit is contained in:
@@ -22,7 +22,6 @@
|
||||
#include <cutlass/util/packed_stride.hpp>
|
||||
#include <cutlass/util/device_memory.h>
|
||||
|
||||
// CUTLASS_CHECK macro (not in headers, define our own)
|
||||
#define CUTLASS_CHECK(status) \
|
||||
do { \
|
||||
cutlass::Status _s = status; \
|
||||
@@ -38,40 +37,34 @@
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// NVFP4 × NVFP4 → BF16 GEMM (for L1 and L2 MoE layers)
|
||||
// NVFP4 × NVFP4 → BF16 GEMM
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// A matrix configuration (activation, FP4 + UE4M3 block scales)
|
||||
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
|
||||
using LayoutATag = cutlass::layout::RowMajor;
|
||||
constexpr int AlignmentA = 32;
|
||||
|
||||
// B matrix configuration (weight, FP4 + UE4M3 block scales)
|
||||
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor; // K-major
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignmentB = 32;
|
||||
|
||||
// C/D matrix configuration (BF16 output)
|
||||
using ElementD = cutlass::bfloat16_t;
|
||||
using ElementC = float; // Epilogue compute type
|
||||
using ElementC = float;
|
||||
using LayoutCTag = cutlass::layout::RowMajor;
|
||||
using LayoutDTag = cutlass::layout::RowMajor;
|
||||
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
|
||||
|
||||
// Kernel functional config
|
||||
using ElementAccumulator = float;
|
||||
using ElementCompute = float;
|
||||
using ArchTag = cutlass::arch::Sm100;
|
||||
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
|
||||
|
||||
// Kernel perf config
|
||||
using MmaTileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
|
||||
constexpr int InputSFVectorSize = 16; // UE4M3 group_size = 16
|
||||
constexpr int InputSFVectorSize = 16;
|
||||
|
||||
// Epilogue: standard LinComb (alpha * AB + beta * C) → BF16 output
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
MmaTileShape, ClusterShape,
|
||||
@@ -100,7 +93,6 @@ using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
// Layout and stride types from the kernel
|
||||
using StrideA = typename Gemm::GemmKernel::StrideA;
|
||||
using StrideB = typename Gemm::GemmKernel::StrideB;
|
||||
using StrideC = typename Gemm::GemmKernel::StrideC;
|
||||
@@ -109,50 +101,113 @@ using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
|
||||
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// C API for PyTorch binding
|
||||
// Scale factor remap kernel
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Remap simple row-major (MN, K//16) scales to CUTLASS interleaved layout
|
||||
// The CUTLASS layout is computed by Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA
|
||||
// which tiles SfAtom to the problem shape.
|
||||
//
|
||||
// For SFVecSize=16, K-major:
|
||||
// SfAtom = Shape<Shape<32,4>, Shape<16,4>>
|
||||
// Stride<Stride<16,4>, Stride<0,1>>
|
||||
//
|
||||
// tile_to_shape(SfAtom{}, make_shape(M,K), Step<_2,_1>{}) produces a 2D layout
|
||||
// where the first mode (row/M) is tiled with step 2 (interleaved with the second mode)
|
||||
// and the second mode (col/K) is tiled with step 1.
|
||||
//
|
||||
// For our remap, we iterate over the CUTLASS layout indices and for each,
|
||||
// compute which (row, k_group) from the source it corresponds to, then copy.
|
||||
|
||||
template<typename LayoutSF>
|
||||
__global__ void remap_sf_to_cutlass_layout_kernel(
|
||||
const cutlass::float_ue4m3_t* __restrict__ src, // (MN, K//16) row-major
|
||||
cutlass::float_ue4m3_t* __restrict__ dst, // CUTLASS layout
|
||||
LayoutSF layout_sf, // CuTe layout for dst
|
||||
int MN, int K_sf // Source dimensions
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = cute::size(layout_sf);
|
||||
if (idx >= total) return;
|
||||
|
||||
// The CUTLASS layout maps idx -> (row_coord, k_coord)
|
||||
// We need to figure out which (row, k_group) this corresponds to
|
||||
// in the simple row-major source layout.
|
||||
//
|
||||
// For the SfAtom with SFVecSize=16, K-major:
|
||||
// The layout interleaves scale factors in groups of 128 rows
|
||||
// and 64 K-groups (4 "tiles" of 16).
|
||||
//
|
||||
// Simple approach: use the layout's coordinate function.
|
||||
// The layout maps a linear index to a logical coordinate (m, k).
|
||||
// We can extract (m, k) and then index into our source as m * K_sf + k.
|
||||
|
||||
// Get the 2D coordinate from the layout
|
||||
auto coord = layout_sf.get_flat_coord(idx);
|
||||
int m = cute::get<0>(coord); // Row index in the layout's coordinate system
|
||||
int k = cute::get<1>(coord); // K-group index
|
||||
|
||||
// Map to source index
|
||||
if (m < MN && k < K_sf) {
|
||||
dst[idx] = src[m * K_sf + k];
|
||||
} else {
|
||||
dst[idx] = cutlass::float_ue4m3_t(0); // Zero padding
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// C API
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
extern "C" {
|
||||
|
||||
// Run a single NVFP4 × NVFP4 → BF16 GEMM
|
||||
// A: (M, K/2) int8 packed E2M1, SFA: UE4M3 block scales, B: (N, K/2) int8 packed E2M1, SFB: UE4M3 block scales
|
||||
// Output: (M, N) bfloat16
|
||||
int cutlass_nvfp4_gemm_run(
|
||||
const void* A_ptr, const void* SFA_ptr, // A matrix + scale factors
|
||||
const void* B_ptr, const void* SFB_ptr, // B matrix + scale factors
|
||||
void* D_ptr, // Output D matrix (bfloat16)
|
||||
int M, int N, int K, // Problem dimensions (K in E2M1 elements, so actual K = K_arg * 2)
|
||||
const void* A_ptr, const void* SFA_ptr,
|
||||
const void* B_ptr, const void* SFB_ptr,
|
||||
void* D_ptr,
|
||||
int M, int N, int K,
|
||||
float alpha, float beta,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
// Compute strides
|
||||
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
|
||||
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
|
||||
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
|
||||
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
|
||||
|
||||
// Compute scale factor layouts
|
||||
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
|
||||
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
|
||||
|
||||
// Cast void* to proper element types
|
||||
using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
|
||||
using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
|
||||
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
|
||||
|
||||
int sfa_size = cute::size(layout_SFA);
|
||||
int sfb_size = cute::size(layout_SFB);
|
||||
int K_sf = K / InputSFVectorSize;
|
||||
|
||||
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
|
||||
cutlass::device_memory::allocation<ElementSF> sfb_cutlass(sfb_size);
|
||||
|
||||
// Remap scales from simple row-major to CUTLASS interleaved layout
|
||||
int block = 256;
|
||||
remap_sf_to_cutlass_layout_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf);
|
||||
remap_sf_to_cutlass_layout_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf);
|
||||
|
||||
typename Gemm::Arguments arguments {
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{M, N, K, 1},
|
||||
{ // Mainloop arguments
|
||||
{
|
||||
static_cast<const ArrayElementA*>(A_ptr), stride_A,
|
||||
static_cast<const ArrayElementB*>(B_ptr), stride_B,
|
||||
static_cast<const ElementSF*>(SFA_ptr), layout_SFA,
|
||||
static_cast<const ElementSF*>(SFB_ptr), layout_SFB
|
||||
sfa_cutlass.get(), layout_SFA,
|
||||
sfb_cutlass.get(), layout_SFB
|
||||
},
|
||||
{ // Epilogue arguments
|
||||
{
|
||||
{ alpha, beta },
|
||||
nullptr, stride_C, // C matrix (not used, beta=0)
|
||||
nullptr, stride_C,
|
||||
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(D_ptr), stride_D
|
||||
}
|
||||
};
|
||||
|
||||
156
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py
Normal file
156
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py
Normal file
@@ -0,0 +1,156 @@
|
||||
"""
|
||||
CUTLASS NVFP4 scale factor layout transformation.
|
||||
|
||||
CUTLASS's Sm1xxBlockScaledConfig expects scale factors in a specific
|
||||
interleaved layout (not simple row-major). This module transforms
|
||||
our simple (M, K//16) float8_e4m3fn scales into CUTLASS's expected layout.
|
||||
|
||||
The layout is defined by:
|
||||
SfAtom = Shape<Shape<_32, _4>, Shape<SFVecSize, _4>>
|
||||
with Stride<Stride<_16, _4>, Stride<_0, _1>>
|
||||
(for K-major, SFVecSize=16)
|
||||
|
||||
layout_SFA = tile_to_shape(SfAtom{}, make_shape(M, K), Step<_2, _1>)
|
||||
layout_SFB = tile_to_shape(SfAtom{}, make_shape(N, K), Step<_2, _1>)
|
||||
|
||||
This creates a 2D layout where scale factors are interleaved in blocks
|
||||
of 32 rows × (16*4) = 64 columns, with specific stride patterns.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
|
||||
|
||||
def compute_cutlass_sf_size(M: int, K: int, sf_vec_size: int = 16) -> int:
|
||||
"""Compute the number of scale factor elements needed for CUTLASS layout.
|
||||
|
||||
The CUTLASS layout tile_to_shape(SfAtom{}, make_shape(M, K), Step<_2, _1>)
|
||||
produces a layout with a specific size that may differ from M * (K // sf_vec_size).
|
||||
"""
|
||||
# SfAtom shape: (32, 4, sf_vec_size, 4) = (32*4, sf_vec_size*4) = (128, 64) for sf_vec_size=16
|
||||
# tile_to_shape tiles this atom to cover (M, K) with Step<_2, _1>
|
||||
# The resulting size = ceil(M / 128) * 128 * ceil(K / 64) * 64 / sf_vec_size...
|
||||
# Actually it's simpler: the layout size = the total number of elements in the tiled layout
|
||||
|
||||
# For SfAtom with SFVecSize=16:
|
||||
# SfAtom has shape (32, 4, 16, 4) which is a 2D layout of shape (128, 64)
|
||||
# (after composition). The actual element count per tile = 128 * 4 = 512
|
||||
# (the K dimension is sf_vec_size * 4 = 64, but only 4 "columns" in the SF atom)
|
||||
|
||||
# Actually, let's just compute it empirically. The size of the layout for (M, K) is:
|
||||
# For SfAtom: Shape<(32,4), (16,4)> = (128, 64) in logical space
|
||||
# Each atom has 128 * 4 = 512 elements (the (32,4) × (4) in the first dim groups)
|
||||
# Wait, that's wrong. Let me think about this differently.
|
||||
|
||||
# The SfAtom for K-major, SFVecSize=16:
|
||||
# Shape<Shape<32,4>, Shape<16,4>> with Stride<Stride<16,4>, Stride<0,1>>
|
||||
# This is a 2D layout. The shape is (32*4, 16*4) = (128, 64) in terms of
|
||||
# logical coordinates, but the actual number of elements = 32*4*16*4 = 8192? No.
|
||||
#
|
||||
# Actually: Shape<Shape<32,4>, Shape<16,4>> is a nested shape.
|
||||
# Flattened: 32 * 4 * 16 * 4 = 8192 elements per atom? That can't be right.
|
||||
#
|
||||
# No - this is a 2D layout. The shape is (32*4, 16*4) = (128, 64).
|
||||
# Number of elements = 128 * 64 = 8192 per atom? Still too many.
|
||||
#
|
||||
# Actually, the CuTe layout system works differently.
|
||||
# Shape<Shape<32,4>, Shape<16,4>> has 32*4 = 128 elements in the first mode
|
||||
# and 16*4 = 64 elements in the second mode. Total = 128 * 64 = 8192? No.
|
||||
#
|
||||
# A CuTe Layout<Shape, Stride> maps logical indices to memory offsets.
|
||||
# For SfAtom: the total number of elements = product of all shape components
|
||||
# = 32 * 4 * 16 * 4 = 8192.
|
||||
# But that's per "atom". The tile_to_shape function tiles this.
|
||||
#
|
||||
# For the full layout with (M, K):
|
||||
# size = ceil(M/128) * 128 * ceil(K/64) * 64 ... no.
|
||||
# Actually the number of SF elements should be roughly M * (K / SFVecSize)
|
||||
# = M * K / 16. The interleaving just changes the stride pattern, not the count.
|
||||
#
|
||||
# The total number of elements in the layout is the product of the shape.
|
||||
# tile_to_shape preserves the element count (it's just a reshaping/tiling).
|
||||
#
|
||||
# For simple (M, K) with SFVecSize=16:
|
||||
# Total SF elements = M * (K // SFVecSize) = M * (K // 16)
|
||||
# But CUTLASS may pad to multiples of 128 (rows) and 64 (cols)
|
||||
M_pad = math.ceil(M / 128) * 128
|
||||
K_pad = math.ceil(K / 64) * 64 # 64 = SFVecSize * 4
|
||||
return M_pad * (K_pad // sf_vec_size)
|
||||
|
||||
|
||||
def transform_scales_to_cutlass_layout(
|
||||
sf: torch.Tensor, # (..., M, K//16) float8_e4m3fn, simple row-major
|
||||
M: int,
|
||||
K: int,
|
||||
sf_vec_size: int = 16,
|
||||
) -> torch.Tensor:
|
||||
"""Transform scale factors from simple row-major to CUTLASS interleaved layout.
|
||||
|
||||
CUTLASS expects scale factors in the Sm1xxBlockScaledConfig layout:
|
||||
- SfAtom = Shape<Shape<32,4>, Shape<16,4>> with Stride<Stride<16,4>, Stride<0,1>>
|
||||
- tile_to_shape(SfAtom{}, make_shape(M, K), Step<_2, _1>)
|
||||
|
||||
This produces an interleaved layout where:
|
||||
- Scales are grouped in blocks of 128 rows × 4 "scale columns"
|
||||
- Within each block, 32 groups of 4 interleaved with stride pattern
|
||||
- The K dimension maps to groups of sf_vec_size consecutive scales
|
||||
|
||||
For our purposes, the simplest approach is to allocate a buffer of the
|
||||
right size and let CUTLASS's own layout computation handle the indexing.
|
||||
|
||||
Since the total number of elements is the same (just reordered), we can
|
||||
compute the CUTLASS layout on the Python side and remap.
|
||||
|
||||
For now, we use a simpler approach: pass the scale factors in their
|
||||
natural layout and adjust the CUTLASS kernel to use a compatible stride.
|
||||
"""
|
||||
# The CUTLASS block-scaled GEMM reads scales using TMA descriptors
|
||||
# that are initialized with the layout_SFA/layout_SFB computed by
|
||||
# Sm1xxBlockScaledConfig::tile_atom_to_shape_SFA.
|
||||
#
|
||||
# The key insight: the CUTLASS 72b example fills scales using
|
||||
# a cute::Tensor with this layout. If we provide data in a different
|
||||
# layout, the TMA loads will read wrong addresses.
|
||||
#
|
||||
# However, for the _single_ expert GEMM (not grouped), CUTLASS handles
|
||||
# the TMA descriptor setup internally based on the stride/layout we pass.
|
||||
# If we pass our scales with a simple row-major stride, CUTLASS will
|
||||
# still try to read them using the interleaved layout — which is wrong.
|
||||
#
|
||||
# The fix: we need to either:
|
||||
# 1. Remap our data to match the CUTLASS layout (complex)
|
||||
# 2. Or use a different CUTLASS API that accepts custom strides
|
||||
#
|
||||
# For now, let's compute the CUTLASS layout size, allocate, and
|
||||
# do a simple remap. The remap is: for each (row, k_group) in the
|
||||
# CUTLASS layout, write the scale factor from our simple layout.
|
||||
|
||||
# For SFVecSize=16, the SfAtom has:
|
||||
# Shape<(32,4), (16,4)> with Stride<(16,4), (0,1)>
|
||||
#
|
||||
# This means for the K-major atom:
|
||||
# - Row index (first mode): 32 groups of 4, stride (16, 4)
|
||||
# So element (i, j) in (32, 4) maps to offset i*16 + j*4
|
||||
# - Col index (second mode): 16 groups of 4, stride (0, 1)
|
||||
# So element (k, l) in (16, 4) maps to offset k*0 + l*1 = l
|
||||
#
|
||||
# Combined: offset = (i*16 + j*4) + l * 128
|
||||
# (because the second mode stride (0, 1) means col dimension
|
||||
# doesn't contribute to the first mode and vice versa)
|
||||
#
|
||||
# Actually, for a 2D CuTe layout:
|
||||
# Shape<(32,4), (16,4)> = (128, 64)
|
||||
# Stride<(16,4), (0,1)> means:
|
||||
# First mode (row): stride 16, 4 -> 32*4=128 rows, row stride = (16,4)
|
||||
# Second mode (col): stride 0, 1 -> 16*4=64 cols, col stride = (0,1)
|
||||
#
|
||||
# For element at (r, c) where r = i*4+j (i in 0..31, j in 0..3)
|
||||
# c = k*4+l (k in 0..15, l in 0..3):
|
||||
# offset = i*16 + j*4 + l*128
|
||||
#
|
||||
# Wait, that's not right either. Let me just compute the layout properly.
|
||||
# The total number of SF elements should be M * (K // 16) for our data.
|
||||
|
||||
# For now, let's just return the scales as-is and see if CUTLASS can handle
|
||||
# a simple stride. The real fix would be to use CUTLASS's own layout helpers.
|
||||
return sf
|
||||
Reference in New Issue
Block a user