diff --git a/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO b/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO deleted file mode 100644 index 59be8beb..00000000 --- a/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO +++ /dev/null @@ -1,7 +0,0 @@ -Metadata-Version: 2.4 -Name: nvfp4-megamoe-kernel -Version: 0.1.0 -Summary: NVFP4 Mega MoE kernel for DeepSeek-V4-Pro on Blackwell (TileLang) -Requires-Python: >=3.10 -Requires-Dist: torch>=2.5 -Requires-Dist: tilelang>=0.1 diff --git a/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt b/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt deleted file mode 100644 index f3641fa0..00000000 --- a/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt +++ /dev/null @@ -1,11 +0,0 @@ -README.md -pyproject.toml -src/nvfp4_megamoe_kernel/__init__.py -src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py -src/nvfp4_megamoe_kernel/symm_buffer.py -src/nvfp4_megamoe_kernel/weight_transform.py -src/nvfp4_megamoe_kernel.egg-info/PKG-INFO -src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt -src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt -src/nvfp4_megamoe_kernel.egg-info/requires.txt -src/nvfp4_megamoe_kernel.egg-info/top_level.txt \ No newline at end of file diff --git a/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt b/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt deleted file mode 100644 index 8b137891..00000000 --- a/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/src/nvfp4_megamoe_kernel.egg-info/requires.txt b/src/nvfp4_megamoe_kernel.egg-info/requires.txt deleted file mode 100644 index 06d97fde..00000000 --- a/src/nvfp4_megamoe_kernel.egg-info/requires.txt +++ /dev/null @@ -1,2 +0,0 @@ -torch>=2.5 -tilelang>=0.1 diff --git a/src/nvfp4_megamoe_kernel.egg-info/top_level.txt b/src/nvfp4_megamoe_kernel.egg-info/top_level.txt deleted file mode 100644 index 0c0c2376..00000000 --- a/src/nvfp4_megamoe_kernel.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -nvfp4_megamoe_kernel diff --git a/src/nvfp4_megamoe_kernel/__init__.py b/src/nvfp4_megamoe_kernel/__init__.py deleted file mode 100644 index 05cdc12d..00000000 --- a/src/nvfp4_megamoe_kernel/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""NVFP4 Mega MoE Kernel — CUTLASS implementation for DeepSeek-V4-Pro on Blackwell.""" - -from nvfp4_megamoe_kernel.nvfp4_mega_moe import ( - nvfp4_mega_moe_full, - nvfp4_mega_moe_l1, - nvfp4_mega_moe_l2, - stage_activation, -) -from nvfp4_megamoe_kernel.weight_transform import ( - transform_nvfp4_weights_for_mega_moe, -) -from nvfp4_megamoe_kernel.symm_buffer import ( - SymmBuffer, - get_symm_buffer_for_nvfp4_mega_moe, -) - -__all__ = [ - "nvfp4_mega_moe_full", - "nvfp4_mega_moe_l1", - "nvfp4_mega_moe_l2", - "stage_activation", - "transform_nvfp4_weights_for_mega_moe", - "SymmBuffer", - "get_symm_buffer_for_nvfp4_mega_moe", -] diff --git a/src/nvfp4_megamoe_kernel/cutedsl/__init__.py b/src/nvfp4_megamoe_kernel/cutedsl/__init__.py deleted file mode 100644 index 0992cc11..00000000 --- a/src/nvfp4_megamoe_kernel/cutedsl/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -""" -NVFP4 MoE kernel using NVIDIA's CuTeDSL ScaledGroupedGemmKernel. - -This replaces the broken C++ CUTLASS kernel. The CuTeDSL kernel handles: -- NVFP4 (Float4E2M1FN + Float8E4M3FN, sf_vec_size=16) natively -- Block-scaled SF layouts (no manual remap needed) -- Full Blackwell pipeline (TMA → MMA → Epilogue overlap) -- Per-expert global scales for NVFP4 - -We just need to: -1. Quantize activations to FP4 (stage_activation) -2. Call the kernel with the right tensor layout -3. Apply MoE routing (gate/up fusion, SiLU, scatter) -""" diff --git a/src/nvfp4_megamoe_kernel/cutedsl/moe.py b/src/nvfp4_megamoe_kernel/cutedsl/moe.py deleted file mode 100644 index 5d82743f..00000000 --- a/src/nvfp4_megamoe_kernel/cutedsl/moe.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -NVFP4 MoE pipeline using CuTeDSL ScaledGroupedGemmKernel. - -Replaces the broken C++ CUTLASS path. Uses NVIDIA's official MoE scaled -grouped GEMM kernel from the CUTLASS CuTeDSL examples. - -Usage: - from nvfp4_megamoe_kernel.cutedsl.moe import nvfp4_mega_moe_full -""" - -import sys -import os -import torch -import cutlass -import cutlass.cute as cute -import cutlass.torch as cutlass_torch -import cutlass.utils as utils -import cutlass.utils.blockscaled_layout as blockscaled_utils - -# Add the CuTeDSL examples to the path so we can import the kernel -_CUTLASS_ROOT = os.environ.get("CUTLASS_ROOT", "/root/cutlass") -_CUTEDSL_EXAMPLES = os.path.join(_CUTLASS_ROOT, "examples/python/CuTeDSL") -if _CUTEDSL_EXAMPLES not in sys.path: - sys.path.insert(0, _CUTEDSL_EXAMPLES) - -from cute.blackwell.kernel.moe.torch_scaled_grouped_mm import ScaledGroupedGemmKernel - -from nvfp4_megamoe_kernel.nvfp4_mega_moe import ( - stage_activation, - _quantize_to_e2m1, -) - -# ── Module-level compiled kernel cache ── -_compiled_l1_kernel = None -_compiled_l2_kernel = None -_l1_kernel_config = None -_l2_kernel_config = None - - -def _get_torch_dtype(cutlass_dtype): - """Convert CUTLASS dtype to PyTorch dtype.""" - mapping = { - cutlass.Float4E2M1FN: torch.float4_e2m1fn_x2, - cutlass.Float8E4M3FN: torch.float8_e4m3fn, - cutlass.Float8E8M0FNU: torch.float8_e8m0fnu, - cutlass.BFloat16: torch.bfloat16, - cutlass.Float16: torch.float16, - cutlass.Float32: torch.float32, - } - return mapping.get(cutlass_dtype) - - -def _torch_tensor_to_cute(torch_tensor: torch.Tensor) -> cute.Tensor: - """Convert a PyTorch GPU tensor to a CuTe tensor with dynamic layout.""" - cute_tensor = cutlass_torch.from_dlpack(torch_tensor) - leading_dim = cutlass_torch.get_leading_dim(torch_tensor) - cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) - return cute_tensor - - -def _compile_kernel_once(kernel, sample_tensors, global_scale_a=None, global_scale_b=None): - """Compile the CuTeDSL kernel on first call, cache the result.""" - import cuda.bindings.driver as cuda - - a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute = sample_tensors - - gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None - gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None - - cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1] - max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - compiled = cute.compile( - kernel, - a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute, - max_active_clusters, stream, - global_scale_a=gsa_cute, - global_scale_b=gsb_cute, - ) - return compiled - - -def run_scaled_grouped_gemm( - mat_a: torch.Tensor, # (tokens_sum, K_packed) float4_e2m1fn_x2 — row-major (K-major for CuTe) - mat_b: torch.Tensor, # (experts, K_packed, N) float4_e2m1fn_x2 — K-major - scale_a: torch.Tensor, # (tokens_sum, K_sf) float8_e4m3fn — row-major - scale_b: torch.Tensor, # (experts, K_sf, N) float8_e4m3fn — K-major after transpose - expert_offsets: torch.Tensor, # (experts,) int32 — cumulative token offsets - global_scale_a: torch.Tensor = None, # (experts,) float32 — NVFP4 per-expert activation scale - global_scale_b: torch.Tensor = None, # (experts,) float32 — NVFP4 per-expert weight scale - mma_tiler_mn: tuple = (128, 128), - cluster_shape_mn: tuple = (1, 1), -) -> torch.Tensor: - """Run the CuTeDSL NVFP4 scaled grouped GEMM. - - 2Dx3D scenario: A(tokens, K) x B(experts, K, N) -> C(tokens, N) - - Args: - mat_a: Activation tensor (tokens_sum, K_packed) in FP4 - mat_b: Weight tensor (experts, K_packed, N) in FP4 - scale_a: Activation block scales (tokens_sum, K_sf) in E4M3 - scale_b: Weight block scales (experts, K_sf, N) in E4M3 - expert_offsets: Cumulative token end offsets per expert - global_scale_a: Per-expert float32 activation global scale (NVFP4) - global_scale_b: Per-expert float32 weight global scale (NVFP4) - - Returns: - Output tensor (tokens_sum, N) in BF16 - """ - global _compiled_l1_kernel, _l1_kernel_config - - tokens_sum = mat_a.shape[0] - k_packed = mat_a.shape[1] - num_experts = mat_b.shape[0] - n_dim = mat_b.shape[2] - k_dim = k_packed * 2 # 2 FP4 values per byte - - # Output tensor - out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) - - # Create kernel config - kernel = ScaledGroupedGemmKernel( - scenario="2Dx3D", - sf_vec_size=16, - accumulate_on_output=False, - separate_tensormap_init=True, - consistent_token_padding=False, - mma_tiler_mnk=(*mma_tiler_mn, 256), - cluster_shape_mnk=(*cluster_shape_mn, 1), - ) - - # Convert to CuTe tensors - a_cute = _torch_tensor_to_cute(mat_a) - b_cute = _torch_tensor_to_cute(mat_b) - sfa_cute = _torch_tensor_to_cute(scale_a) - sfb_cute = _torch_tensor_to_cute(scale_b) - c_cute = _torch_tensor_to_cute(out) - offs_cute = _torch_tensor_to_cute(expert_offsets) - - # Workspace - workspace_size = kernel.get_workspace_size(num_experts) - workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device) - ws_cute = _torch_tensor_to_cute(workspace) - - gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None - gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None - - import cuda.bindings.driver as cuda - cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1] - max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - # Compile and run - compiled = cute.compile( - kernel, - a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute, - max_active_clusters, stream, - global_scale_a=gsa_cute, - global_scale_b=gsb_cute, - ) - - compiled( - a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute, - stream, - global_scale_a=gsa_cute, - global_scale_b=gsb_cute, - ) - torch.cuda.synchronize() - - return out diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md deleted file mode 100644 index 4cd75921..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md +++ /dev/null @@ -1,99 +0,0 @@ -# CUTLASS NVFP4 Block-Scaled GEMM Kernel - -Native Blackwell (SM100) NVFP4 block-scaled GEMM using CUTLASS 3.x. - -## Overview - -This kernel implements the DeepSeek-V4-Pro MoE GEMM operations using CUTLASS's -`MainloopSm100TmaUmmaWarpSpecializedBlockScaled` collective, which invokes the -native `mxf8f6f4.block_scale` tensor core instruction (`tcgen05.mma`) on NVIDIA -Blackwell GPUs. - -### Key Features - -- **Native NVFP4 MMA**: E2M1 × E2M1 with UE4M3 block-16 scaling entirely in hardware -- **No dequantization**: Avoids the costly dequantize-then-BF16-GEMM fallback path -- **TMA + UMMA**: Uses TMA for loading data into shared memory and UMMA for tensor core ops -- **TMEM scale loading**: UE4M3 scale factors loaded into tensor memory via `tcgen05.ld` -- **Grouped expert GEMM**: Per-expert dispatch for MoE with top-k routing - -### Architecture - -``` -E2M1 (int8, 2 vals/byte) + UE4M3 (float8_e4m3fn, group_size=16) - → TMA load to shared memory - → UMMA block-scaled MMA (mxf8f6f4.block_scale) - → float32 accumulator - → BF16 output -``` - -## Data Layout - -| Tensor | Shape | Type | Layout | -|--------|-------|------|--------| -| A (activation) | (M, K//2) | int8 | K-major (ColumnMajor) | -| SFA (activation scales) | (M, K//16) | float8_e4m3fn | K-major (Sm1xxBlockScaledConfig) | -| B (weight) | (N, K//2) | int8 | K-major (ColumnMajor) | -| SFB (weight scales) | (N, K//16) | float8_e4m3fn | K-major (Sm1xxBlockScaledConfig) | -| C (output) | (M, N) | bfloat16 | RowMajor | - -K//2 because E2M1 packs 2 values per byte. -K//16 because UE4M3 block scale has group_size=16. - -## Building on B200 - -```bash -# Inside the Docker container on the B200: -cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm -bash build.sh -``` - -Or manually: - -```bash -export CUTLASS_INCLUDE_DIR=/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include -python3 setup.py build_ext --inplace -``` - -## Testing - -```bash -python3 test_gemm.py -``` - -## Usage in DeepSeek-V4-Pro - -The kernel is automatically used by `nvfp4_mega_moe.py` when: -1. `MEGA_MOE_USE_CUTLASS=1` (default) -2. The CUTLASS extension compiles successfully - -If CUTLASS is unavailable, it falls back to the TileLang or dequantize+BF16 path. - -## CUTLASS Internals - -### Dispatch Policy -`MainloopSm100TmaUmmaWarpSpecializedBlockScaled` - -### TiledMma -UMMA atom: `mxf8f6f4.block_scale` with SFVecSize=16 - -### Scale Factor Layout -Uses `Sm1xxBlockScaledConfig<16>` which defines: -- SfAtom layout for K-major scale factors -- `tile_atom_to_shape_SFA/SFB` for computing the global scale layout -- `deduce_smem_layoutSFA/SFB` for shared memory layout - -### Pipeline -1. TMA loads A, B, SFA, SFB into shared memory -2. UMMA warp-specialized MMA with block scaling -3. Scale factors loaded from shared memory to TMEM via UTCCP -4. Accumulator in float32, converted to BF16 in epilogue - -## Files - -- `cutlass_nvfp4_gemm.cu` — Standalone CUDA kernel (C API) -- `pytorch_binding.cpp` — PyTorch extension binding -- `kernel.py` — Python wrapper with compilation and fallback -- `setup.py` — Build configuration -- `build.sh` — Build script for B200 -- `test_gemm.py` — Test script diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py deleted file mode 100644 index 779f0deb..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro on Blackwell (SM100).""" - -from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import ( - cutlass_nvfp4_blockscaled_gemm, - cutlass_grouped_nvfp4_gemm, -) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh deleted file mode 100644 index 78c233ab..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh +++ /dev/null @@ -1,44 +0,0 @@ -#!/bin/bash -# Build script for CUTLASS NVFP4 block-scaled GEMM on B200 (Blackwell SM100). -# -# Run inside the Docker container: -# docker exec -it deepseek-v4-quant-vllm bash -# cd /path/to/cutlass_nvfp4_gemm && bash build.sh -# -# Or from outside: -# docker exec deepseek-v4-quant-vllm bash -c "cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm && bash build.sh" - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR" - -# CUTLASS include path (inside the Docker container) -export CUTLASS_INCLUDE_DIR="${CUTLASS_INCLUDE_DIR:-/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include}" - -echo "=== CUTLASS NVFP4 GEMM Build ===" -echo "CUTLASS_INCLUDE_DIR: $CUTLASS_INCLUDE_DIR" - -# Verify CUTLASS headers -if [ ! -f "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h" ]; then - echo "ERROR: CUTLASS headers not found at ${CUTLASS_INCLUDE_DIR}" - echo "Set CUTLASS_INCLUDE_DIR to point to the cutlass/include directory." - exit 1 -fi - -# Verify block-scaled MMA header -if [ ! -f "${CUTLASS_INCLUDE_DIR}/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" ]; then - echo "WARNING: Block-scaled MMA header not found. The CollectiveBuilder path will be used." -fi - -echo "Building PyTorch extension..." -python3 setup.py build_ext --inplace 2>&1 | tee build.log - -if [ $? -eq 0 ]; then - echo "=== Build SUCCESS ===" - echo "Extension built. Test with: python3 test_gemm.py" -else - echo "=== Build FAILED ===" - echo "Check build.log for errors." - exit 1 -fi 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 deleted file mode 100644 index 39077d4c..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu +++ /dev/null @@ -1,385 +0,0 @@ -/*************************************************************************************************** - * CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro MoE - **************************************************************************************************/ - -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define CUTLASS_CHECK(status) \ - do { \ - cutlass::Status _s = status; \ - if (_s != cutlass::Status::kSuccess) { \ - fprintf(stderr, "CUTLASS error at %s:%d: %s\n", \ - __FILE__, __LINE__, cutlassGetStatusString(_s)); \ - return -1; \ - } \ - } while (0) - -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - -using namespace cute; - -using ElementA = cutlass::nv_float4_t; -using LayoutATag = cutlass::layout::RowMajor; -constexpr int AlignmentA = 32; - -using ElementB = cutlass::nv_float4_t; -using LayoutBTag = cutlass::layout::ColumnMajor; -constexpr int AlignmentB = 32; - -using ElementD = cutlass::bfloat16_t; -using ElementC = float; -using LayoutCTag = cutlass::layout::RowMajor; -using LayoutDTag = cutlass::layout::RowMajor; -constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; -constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - -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, AlignmentC, - ElementD, LayoutDTag, AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto ->::CollectiveOp; - -using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, OperatorClass, - ElementA, LayoutATag, AlignmentA, - ElementB, LayoutBTag, AlignmentB, - 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 StrideA = typename Gemm::GemmKernel::StrideA; -using StrideB = typename Gemm::GemmKernel::StrideB; -using StrideC = typename Gemm::GemmKernel::StrideC; -using StrideD = typename Gemm::GemmKernel::StrideD; -using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; -using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; - -///////////////////////////////////////////////////////////////////////////////////////////////// -// Scale factor remap: source (row-major or col-major) -> CUTLASS interleaved layout -// -// Scale factor remap: source (row-major or col-major) -> CUTLASS interleaved layout -// -// Iterates over CUTLASS dest indices, uses idx2crd to get the hierarchical coordinate, -// then extracts logical (m, k_sf) from the flattened result. -// -// The flattened coordinate from idx2crd has nested structure. -// For SFA with Step<_2,_1> tiling, the layout shape is: -// ((32, 4, n_m_tiles), (16, 4, n_k_tiles)) -// Flattening gives: (inner_m, sub_m, tile_m, inner_k, sub_k, tile_k) -// where inner_m in [0,32), sub_m in [0,4), tile_m in [0, n_m_tiles) -// inner_k in [0,16) (within one SF group), sub_k in [0,4), tile_k in [0, n_k_tiles) -// -// Logical m = tile_m * 128 + inner_m * 4 + sub_m -// Logical k_sf = tile_k * 4 + sub_k (inner_k is within one SF group — same byte) -// -// NOTE: Allocation must use cute::cosize() (physical size including tile padding), -// not cute::size() (logical size). The dest buffer is zero-initialized so padding -// positions that aren't written are correct zeros. -///////////////////////////////////////////////////////////////////////////////////////////////// - -template -__global__ void remap_sf_to_cutlass_kernel( - const cutlass::float_ue4m3_t* __restrict__ src, - cutlass::float_ue4m3_t* __restrict__ dst, - LayoutSF layout_sf, - int MN, - int K_sf, - int src_stride_mn, - int src_stride_ksf -) { - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int total = MN * K_sf; - if (tid >= total) return; - - int mn = tid / K_sf; - int k_sf = tid % K_sf; - - // Logical K element coordinate, not compact scale-factor coordinate. - int k_elem = k_sf * 16; - - int dst_idx = layout_sf(cute::make_coord(mn, k_elem, 0)); - - dst[dst_idx] = src[mn * src_stride_mn + k_sf * src_stride_ksf]; -} - -// Roundtrip verifier: check that forward remap wrote the correct bytes -template -__global__ void check_sf_forward_kernel( - const cutlass::float_ue4m3_t* src, - const cutlass::float_ue4m3_t* dst, - LayoutSF layout_sf, - int MN, - int K_sf, - int src_stride_mn, - int src_stride_ksf, - int* errors -) { - int tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= MN * K_sf) return; - - int mn = tid / K_sf; - int k_sf = tid % K_sf; - - int src_idx = mn * src_stride_mn + k_sf * src_stride_ksf; - int dst_idx = layout_sf(cute::make_coord(mn, k_sf * 16, 0)); - - auto* src_u8 = reinterpret_cast(src); - auto* dst_u8 = reinterpret_cast(dst); - - if (src_u8[src_idx] != dst_u8[dst_idx]) { - atomicAdd(errors, 1); - } -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// C API -///////////////////////////////////////////////////////////////////////////////////////////////// - -extern "C" { - -int cutlass_nvfp4_gemm_run( - 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 -) { - 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)); - - 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)); - - 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(cute::filter_zeros(layout_SFA)); - int sfb_size = cute::size(cute::filter_zeros(layout_SFB)); - int K_sf = K / InputSFVectorSize; - - cutlass::device_memory::allocation sfa_cutlass(sfa_size); - cutlass::device_memory::allocation sfb_cutlass(sfb_size); - cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); - cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); - - int block = 256; - int sfa_src_total = M * K_sf; - int sfb_src_total = N * K_sf; - remap_sf_to_cutlass_kernel<<<(sfa_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, - M, K_sf, K_sf, 1); // SFA source: row-major (M, K_sf) - remap_sf_to_cutlass_kernel<<<(sfb_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, - N, K_sf, 1, N); // SFB source: row-major (K_sf, N) after transpose - - // One-time roundtrip verification of SF remap - static bool verified = false; - if (!verified) { - verified = true; - cudaStreamSynchronize(stream); - int* d_errors; - cudaMalloc(&d_errors, sizeof(int)); - cudaMemset(d_errors, 0, sizeof(int)); - - check_sf_forward_kernel<<<(sfa_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, - M, K_sf, K_sf, 1, d_errors); - int sfa_errors = 0; - cudaMemcpyAsync(&sfa_errors, d_errors, sizeof(int), cudaMemcpyDeviceToHost, stream); - - cudaMemset(d_errors, 0, sizeof(int)); - check_sf_forward_kernel<<<(sfb_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, - N, K_sf, 1, N, d_errors); - int sfb_errors = 0; - cudaMemcpyAsync(&sfb_errors, d_errors, sizeof(int), cudaMemcpyDeviceToHost, stream); - - cudaStreamSynchronize(stream); - printf("[SF-VERIFY] M=%d N=%d K=%d K_sf=%d sfa_errors=%d sfb_errors=%d " - "sfa_size=%d sfb_size=%d\n", - M, N, K, K_sf, sfa_errors, sfb_errors, sfa_size, sfb_size); - cudaFree(d_errors); - } - - typename Gemm::Arguments arguments { - cutlass::gemm::GemmUniversalMode::kGemm, - {M, N, K, 1}, - { - static_cast(A_ptr), stride_A, - static_cast(B_ptr), stride_B, - sfa_cutlass.get(), layout_SFA, - sfb_cutlass.get(), layout_SFB - }, - { - { alpha, beta }, - nullptr, stride_C, - static_cast(D_ptr), stride_D - } - }; - - Gemm gemm; - CUTLASS_CHECK(gemm.can_implement(arguments)); - - size_t workspace_size = Gemm::get_workspace_size(arguments); - cutlass::device_memory::allocation workspace(workspace_size); - - CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); - CUTLASS_CHECK(gemm.run(stream)); - - return 0; -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// SFB prepack: pre-remap weight scale factors once at load time -///////////////////////////////////////////////////////////////////////////////////////////////// - -extern "C" int cutlass_nvfp4_sfb_size( - int M, int N, int K, - int* out_size -) { - using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); - *out_size = cute::size(cute::filter_zeros(layout_SFB)); - return 0; -} - -extern "C" int cutlass_nvfp4_prepack_sfb_run( - const void* SFB_ptr, - void* SFB_cutlass_ptr, - int M, int N, int K, - cudaStream_t stream -) { - using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); - using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; - - int sfb_size = cute::size(cute::filter_zeros(layout_SFB)); - int K_sf = K / InputSFVectorSize; - - cudaMemsetAsync(static_cast(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream); - - int block = 256; - int sfb_src_total = N * K_sf; - remap_sf_to_cutlass_kernel<<<(sfb_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFB_ptr), - static_cast(SFB_cutlass_ptr), - layout_SFB, - N, K_sf, 1, N // SFB source: row-major (K_sf, N) after transpose - ); - - return 0; -} - -///////////////////////////////////////////////////////////////////////////////////////////////// -// GEMM with prepacked SFB — skips SFB allocation, memset, and remap -///////////////////////////////////////////////////////////////////////////////////////////////// - -extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( - const void* A_ptr, const void* SFA_ptr, - const void* B_ptr, const void* SFB_cutlass_ptr, - void* D_ptr, - int M, int N, int K, - float alpha, float beta, - cudaStream_t stream -) { - 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)); - - 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)); - - 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(cute::filter_zeros(layout_SFA)); - int K_sf = K / InputSFVectorSize; - - // Only remap SFA (activation scales) — SFB is prepacked - cutlass::device_memory::allocation sfa_cutlass(sfa_size); - cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); - - int block = 256; - int sfa_src_total = M * K_sf; - remap_sf_to_cutlass_kernel<<<(sfa_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, - M, K_sf, K_sf, 1); // SFA source: row-major (M, K_sf) - - typename Gemm::Arguments arguments { - cutlass::gemm::GemmUniversalMode::kGemm, - {M, N, K, 1}, - { - static_cast(A_ptr), stride_A, - static_cast(B_ptr), stride_B, - sfa_cutlass.get(), layout_SFA, - static_cast(SFB_cutlass_ptr), layout_SFB - }, - { - { alpha, beta }, - nullptr, stride_C, - static_cast(D_ptr), stride_D - } - }; - - Gemm gemm; - CUTLASS_CHECK(gemm.can_implement(arguments)); - - size_t workspace_size = Gemm::get_workspace_size(arguments); - cutlass::device_memory::allocation workspace(workspace_size); - - CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); - CUTLASS_CHECK(gemm.run(stream)); - - return 0; -} - -} // extern "C" - -#endif diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py deleted file mode 100644 index 76e9bc00..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py +++ /dev/null @@ -1,155 +0,0 @@ -""" -CUTLASS NVFP4 Block-Scaled GEMM — Native Blackwell SM100 kernel. - -Uses the pre-compiled PyTorch CUDA extension (cutlass_nvfp4_gemm._C) -which invokes native mxf8f6f4.block_scale tensor core instructions. -""" - -import os -import torch - -MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) - -try: - from cutlass_nvfp4_gemm import _C - _CUTLASS_AVAILABLE = True -except ImportError: - _CUTLASS_AVAILABLE = False - - -def cutlass_nvfp4_blockscaled_gemm( - A_packed, # (M, K_half) int8 packed E2M1 - SFA, # scale factors for A (float8_e4m3fn) - B_packed, # (K_half, N) int8 packed E2M1, column-major for CUTLASS - SFB, # scale factors for B — either (sf_k, N) float8_e4m3fn row-major, or prepacked CUTLASS layout - M, N, K, # Problem dimensions (K in FP4 elements) - alpha=1.0, # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C - sfb_prepacked=False, # True if SFB is already in CUTLASS layout -): - """Single NVFP4 block-scaled GEMM using CUTLASS. - - If sfb_prepacked=True, SFB is assumed to be in CUTLASS interleaved layout - (from prepack_sfb) and the SFB remap is skipped. - """ - if not _CUTLASS_AVAILABLE: - raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") - if sfb_prepacked: - return _C.forward_prepacked_sfb(A_packed, SFA, B_packed, SFB, M, N, K, alpha) - else: - return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha) - - -def prepack_sfb(SFB, M, N, K): - """Pre-remap SFB weight scales into CUTLASS interleaved layout. - - Call once after weight transform. Returns a tensor that can be passed - to cutlass_nvfp4_blockscaled_gemm with sfb_prepacked=True. - - M is used for layout sizing. Test with different M values to confirm - SFB layout is M-independent; if so, any valid M works (e.g. 128). - """ - if not _CUTLASS_AVAILABLE: - raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") - return _C.prepack_sfb(SFB, M, N, K) - - -def cutlass_grouped_nvfp4_gemm( - x_fp4, # (num_slots_or_tokens, K_half) int8 packed E2M1 - x_sf, # (num_slots_or_tokens, sf_k) float8_e4m3fn block scales - weights, # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS - weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major - slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs - slot_token=None, # (num_slots,) int64 — per-slot token indices (default: arange) - alpha=1.0, # fp32 scalar: D = alpha * A @ B (from stage_activation global scale) - per_expert_alpha=None, # (E_per_rank,) float32 — per-expert alpha overrides scalar alpha -): - """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4. - - Takes 1D per-slot expert IDs and token indices (pre-built by caller). - SFB weight scales are remapped per-expert inside CUTLASS on each call. - NO prepack cache — see nvfp4_mega_moe.py for rationale. - - For L1: x_fp4 has num_tokens rows, slot_token maps slots→rows. - For L2: x_fp4 has num_slots rows, slot_token is just arange(num_slots). - - If per_expert_alpha is provided, each expert uses its own alpha value - (activation_global_scale * weight_global_scale[expert]) instead of the - scalar alpha. This preserves full float32 precision — no lossy float8 - folding of weight global scales. - - Returns: - slot_out: (num_slots, N) bfloat16 — per-slot GEMM results - slot_token: (num_slots,) int64 — token index for each slot - """ - num_slots = slot_expert_ids.shape[0] - K_half = x_fp4.shape[1] - K = K_half * 2 - N = weights.shape[2] - num_experts = weights.shape[0] - - if num_slots == 0: - slot_out = torch.empty(0, N, dtype=torch.bfloat16, device=x_fp4.device) - slot_token_out = torch.empty(0, dtype=torch.int64, device=x_fp4.device) - return slot_out, slot_token_out - - # Use provided slot_token or default to identity mapping - provided_slot_token = slot_token - - if provided_slot_token is None: - slot_token_out = torch.arange(num_slots, device=x_fp4.device) - slot_x = x_fp4 - slot_x_sf = x_sf - else: - slot_token_out = provided_slot_token - slot_x = x_fp4[provided_slot_token].contiguous() - slot_x_sf = x_sf[provided_slot_token].contiguous() - - if MEGA_MOE_DEBUG: - print(f"[cutlass_grouped_gemm] slots={num_slots} K={K} N={N} " - f"experts={num_experts} per_expert_alpha={'yes' if per_expert_alpha is not None else 'no'}") - - slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device) - - for e in range(num_experts): - expert_slots = (slot_expert_ids == e) - if not expert_slots.any(): - continue - - e_idx = expert_slots.nonzero(as_tuple=True)[0] - expert_x = slot_x[e_idx] - expert_x_sf = slot_x_sf[e_idx] - expert_w = weights[e] - expert_w_sf = weight_sf[e] - M_expert = e_idx.shape[0] - - # Per-expert alpha: activation_gs * weight_gs (float32, no precision loss) - expert_alpha = float(per_expert_alpha[e]) if per_expert_alpha is not None else alpha - - if MEGA_MOE_DEBUG and e < 3 and M_expert > 0: - print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} " - f"w shape={expert_w.shape} alpha={expert_alpha:.4e}") - - # Shape/dtype contract asserts — SFB bugs hide in silent shape mismatches - assert expert_x.shape == (M_expert, K // 2), f"expert_x shape {expert_x.shape} != ({M_expert}, {K // 2})" - assert expert_x_sf.shape == (M_expert, K // 16), f"SFA shape {expert_x_sf.shape} != ({M_expert}, {K // 16})" - assert expert_w.shape == (K // 2, N), f"expert_w shape {expert_w.shape} != ({K // 2}, {N})" - assert expert_w_sf.shape == (K // 16, N), f"SFB shape {expert_w_sf.shape} != ({K // 16}, {N})" - assert expert_x_sf.dtype == torch.float8_e4m3fn, f"SFA dtype {expert_x_sf.dtype}" - assert expert_w_sf.dtype == torch.float8_e4m3fn, f"SFB dtype {expert_w_sf.dtype}" - - expert_out = cutlass_nvfp4_blockscaled_gemm( - expert_x, expert_x_sf, - expert_w, expert_w_sf, - M_expert, N, K, - alpha=expert_alpha, - ) - - if MEGA_MOE_DEBUG: - if torch.isnan(expert_out).any() or torch.isinf(expert_out).any(): - raise RuntimeError( - f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. " - f"M={M_expert} N={N} K={K} alpha={expert_alpha:.4e}") - - slot_out[e_idx] = expert_out - - return slot_out, slot_token_out diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp deleted file mode 100644 index 774502df..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp +++ /dev/null @@ -1,131 +0,0 @@ -/** PyTorch binding for CUTLASS NVFP4 block-scaled GEMM */ - -#include -#include -#include - -extern "C" int cutlass_nvfp4_gemm_run( - 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 -); - -extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( - const void* A_ptr, const void* SFA_ptr, - const void* B_ptr, const void* SFB_cutlass_ptr, - void* D_ptr, - int M, int N, int K, - float alpha, float beta, - cudaStream_t stream -); - -extern "C" int cutlass_nvfp4_sfb_size( - int M, int N, int K, - int* out_size -); - -extern "C" int cutlass_nvfp4_prepack_sfb_run( - const void* SFB_ptr, - void* SFB_cutlass_ptr, - int M, int N, int K, - cudaStream_t stream -); - -torch::Tensor cutlass_nvfp4_gemm_forward( - torch::Tensor A_packed, - torch::Tensor SFA, - torch::Tensor B_packed, - torch::Tensor SFB, - int64_t M, int64_t N, int64_t K, - double alpha = 1.0 -) { - auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device())); - - auto stream = c10::cuda::getCurrentCUDAStream(); - cudaStream_t cuda_stream = stream.stream(); - - int rc = cutlass_nvfp4_gemm_run( - A_packed.data_ptr(), SFA.data_ptr(), - B_packed.data_ptr(), SFB.data_ptr(), - D.data_ptr(), - static_cast(M), static_cast(N), static_cast(K), - static_cast(alpha), 0.0f, - cuda_stream - ); - - TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM failed with error code ", rc); - - return D; -} - -torch::Tensor cutlass_nvfp4_gemm_forward_prepacked_sfb( - torch::Tensor A_packed, - torch::Tensor SFA, - torch::Tensor B_packed, - torch::Tensor SFB_cutlass, - int64_t M, int64_t N, int64_t K, - double alpha = 1.0 -) { - auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device())); - - auto stream = c10::cuda::getCurrentCUDAStream(); - cudaStream_t cuda_stream = stream.stream(); - - int rc = cutlass_nvfp4_gemm_run_prepacked_sfb( - A_packed.data_ptr(), SFA.data_ptr(), - B_packed.data_ptr(), SFB_cutlass.data_ptr(), - D.data_ptr(), - static_cast(M), static_cast(N), static_cast(K), - static_cast(alpha), 0.0f, - cuda_stream - ); - - TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM (prepacked SFB) failed with error code ", rc); - - return D; -} - -torch::Tensor prepack_sfb( - torch::Tensor SFB, - int64_t M, - int64_t N, - int64_t K -) { - int size = 0; - int rc = cutlass_nvfp4_sfb_size( - static_cast(M), - static_cast(N), - static_cast(K), - &size - ); - TORCH_CHECK(rc == 0, "sfb_size failed"); - - auto out = torch::empty( - {size}, - torch::dtype(SFB.dtype()).device(SFB.device()) - ); - - auto stream = c10::cuda::getCurrentCUDAStream(); - - rc = cutlass_nvfp4_prepack_sfb_run( - SFB.data_ptr(), - out.data_ptr(), - static_cast(M), - static_cast(N), - static_cast(K), - stream.stream() - ); - - TORCH_CHECK(rc == 0, "prepack_sfb failed"); - - return out; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &cutlass_nvfp4_gemm_forward, "CUTLASS NVFP4 block-scaled GEMM forward"); - m.def("forward_prepacked_sfb", &cutlass_nvfp4_gemm_forward_prepacked_sfb, "CUTLASS NVFP4 GEMM forward with prepacked SFB"); - m.def("prepack_sfb", &prepack_sfb, "Pre-remap SFB weight scales into CUTLASS layout"); -} diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py deleted file mode 100644 index 08de1df0..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Setup script for CUTLASS NVFP4 block-scaled GEMM PyTorch extension.""" - -import os -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -# CUTLASS include directory — prefer the latest from GitHub -CUTLASS_INCLUDE_DIR = os.environ.get( - "CUTLASS_INCLUDE_DIR", - "/root/cutlass/include" -) -if not os.path.exists(os.path.join(CUTLASS_INCLUDE_DIR, "cutlass", "cutlass.h")): - for alt in [ - "/root/cutlass/include", - "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include", - "/usr/local/include/cutlass", - "/opt/cutlass/include", - ]: - if os.path.exists(os.path.join(alt, "cutlass", "cutlass.h")): - CUTLASS_INCLUDE_DIR = alt - break - -CUTLASS_UTIL_INCLUDE = os.path.join(os.path.dirname(CUTLASS_INCLUDE_DIR), "tools", "util", "include") - -include_dirs = [CUTLASS_INCLUDE_DIR] -if os.path.exists(CUTLASS_UTIL_INCLUDE): - include_dirs.append(CUTLASS_UTIL_INCLUDE) - -# CCCL / libcu++ headers (required by CUTLASS 3.x) -CCCL_INCLUDE = "/usr/local/cuda-13.0/targets/x86_64-linux/include/cccl" -if os.path.exists(CCCL_INCLUDE): - include_dirs.append(CCCL_INCLUDE) - -setup( - name="cutlass_nvfp4_gemm", - ext_modules=[ - CUDAExtension( - name="cutlass_nvfp4_gemm._C", - sources=[ - "pytorch_binding.cpp", - "cutlass_nvfp4_gemm.cu", - ], - include_dirs=include_dirs, - extra_compile_args={ - "cxx": [ - "-O3", - "-std=c++17", - "-DCUTLASS_ENABLE_GEMP_OPERATION=1", - "-DCUTLASS_ARCH_SM100_ENABLED=1", - ], - "nvcc": [ - "-gencode=arch=compute_100a,code=sm_100a", - "--expt-relaxed-constexpr", - "-DCUTLASS_ENABLE_GEMP_OPERATION=1", - "-DCUTLASS_ARCH_SM100_ENABLED=1", - "--ptxas-options=-v", - "--ptxas-options=-allow-expensive-optimizations=true", - ], - }, - ), - ], - cmdclass={ - "build_ext": BuildExtension, - }, -) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py deleted file mode 100644 index ec74b4cf..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py +++ /dev/null @@ -1,21 +0,0 @@ -""" -CUTLASS NVFP4 scale factor layout — reference documentation. - -CUTLASS's Sm1xxBlockScaledConfig expects scale factors in a specific -interleaved layout (not simple row-major). The layout is defined by: - - SfAtom = Shape, Shape> - with Stride, Stride<_0, _1>> - (SFVecSize=16 for NVFP4 UE4M3 block-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>) - -The actual remap from row-major → CUTLASS interleaved layout happens -in the CUDA kernel (remap_sf_to_cutlass_kernel in cutlass_nvfp4_gemm.cu), -NOT in Python. This file exists for reference only. - -The CUDA remap uses cute::idx2crd() to invert the CUTLASS layout: -for each linear index in the CUTLASS layout, it computes the logical -(m, k) coordinate and reads from the corresponding row-major position. -""" diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py deleted file mode 100644 index 39d02a64..00000000 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ /dev/null @@ -1,531 +0,0 @@ -""" -NVFP4 Mega MoE Kernel — Full MoE with expert parallelism. - -This is the main kernel that replaces fp8_nvfp4_mega_moe from DeepGEMM. - -Architecture: -- L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- SiLU+Mul activation (per-slot, BEFORE combining expert paths) -- L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- Routing weights applied ONCE at final scatter -- NVLink cross-rank sync handled by caller (not this kernel) -- Expert parallel: each rank handles NUM_EXPERTS/8 experts - -The kernel uses native NVFP4 block-scaled MMA via tcgen05.mma -kind::mxf8f6f4.block_scale on Blackwell (SM100). - -Native NVFP4 path: - E2M1 (int8, 2 vals/byte) × E2M1 + UE4M3 block-16 scales - → native hardware block-scaled MMA in tensor cores - → float32 accumulator - -This replaces the dequantize-then-BF16-GEMM approach. The native path -performs the E2M1 × E2M1 with UE4M3 block scaling entirely in hardware, -avoiding the costly dequantization step. -""" - -import os -import torch - -def unpack_ue4m3_u32(x_u32): - """Unpack uint32 packed UE4M3 scales to float8_e4m3fn. - - Each uint32 contains 4 UE4M3 values packed in bits [0:8], [8:16], [16:24], [24:32]. - Must use bit reinterpret (view), NOT value cast (to) — byte 0x3F is the float8 - whose bits are 0x3F (~0.984), NOT the integer 63. - - CUDA doesn't implement bitwise ops on uint32, so we cast to int32 first. - Supports ND tensors — last dim is the packed dim (N words → N*4 float8 values). - """ - # CUDA uint32 lacks bitwise ops — use int32 - x_i32 = x_u32.to(torch.int32) - *prefix, n_words = x_i32.shape - - # Extract 4 bytes, cast to uint8, then bit-reinterpret to float8_e4m3fn - b0 = (x_i32 & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) - b1 = ((x_i32 >> 8) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) - b2 = ((x_i32 >> 16) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) - b3 = ((x_i32 >> 24) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) - - # Interleave into (*prefix, n_words*4) - out = torch.empty(*prefix, n_words * 4, dtype=torch.float8_e4m3fn, device=x_u32.device) - out[..., 0::4] = b0 - out[..., 1::4] = b1 - out[..., 2::4] = b2 - out[..., 3::4] = b3 - return out - -# CUTLASS native NVFP4 block-scaled GEMM (SM100 Blackwell) -MEGA_MOE_USE_CUTLASS = int(os.environ.get("MEGA_MOE_USE_CUTLASS", "1")) - -try: - from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import ( - cutlass_nvfp4_blockscaled_gemm, - cutlass_grouped_nvfp4_gemm, - ) - _CUTLASS_AVAILABLE = True -except ImportError: - _CUTLASS_AVAILABLE = False - -# DeepSeek-V4-Pro dimensions -HIDDEN = 7168 -INTERMEDIATE = 3072 -NUM_EXPERTS = 256 -NUM_RANKS = 8 -NUM_TOPK = 6 - -# NVFP4 scale parameters -SF_GRANULARITY_K = 16 # UE4M3 group_size -SF_PACK_FACTOR = 4 # 4 UE4M3 values per uint32 - -# Runtime flags -MEGA_MOE_STATIC = int(os.environ.get("MEGA_MOE_STATIC", "0")) -MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) - - - -# --------------------------------------------------------------------------- -# Main kernel entry points -# --------------------------------------------------------------------------- - -def nvfp4_mega_moe_l1( - x_fp4, # (num_tokens, K//2) int8 packed E2M1 - x_sf, # (num_tokens, sf_k_groups) float8_e4m3fn - l1_weights, # (E_per_rank, K//2, 2*INTER) int8, column-major for CUTLASS - l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major - slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs - slot_token, # (num_slots,) int64 — token index per slot - l1_global_sf, # (E_per_rank, 2) or (E_per_rank,) float32 — weight global scales - alpha=1.0, # fp32 scalar from stage_activation global scale -): - """L1 GEMM: gate_up_proj — slot-based, no routing weights. - - Global scale is NOT folded into block scales. Instead, it's applied as a - per-expert multiplier to the GEMM alpha: alpha_expert = alpha * global_sf[expert]. - For L1 with gate+up: gate and up share one GEMM but may have different global scales. - Since the GEMM produces gate|up in one shot, we use a single alpha per expert. - Post-GEMM, we apply the gate/up ratio correction if they differ. - - Actually, for simplicity and correctness: we use the gate global scale as alpha - and correct the up portion after GEMM. But since gate and up global scales - are typically identical in practice, we just use the geometric mean. - - CLEANER APPROACH: use per-expert alpha directly in the grouped GEMM. - The grouped GEMM iterates per expert, so each expert can have its own alpha. - For L1 with separate gate/up global scales, we use the geometric mean - and then apply a correction factor to the up portion. - """ - K_half = x_fp4.shape[1] - K = K_half * 2 - N = l1_weights.shape[2] # 2 * INTERMEDIATE = 6144 - - if MEGA_MOE_DEBUG: - print(f"[nvfp4_moe_l1] tokens={x_fp4.shape[0]} K={K} N={N} slots={slot_expert_ids.shape[0]}") - - x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf - w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales - assert w_sf_fp8.dtype == torch.float8_e4m3fn, f"l1_scales after unpack dtype={w_sf_fp8.dtype}" - - # Compute per-expert alpha: activation_gs * weight_gs - # For L1 with (E, 2) gate/up global scales, use geometric mean per expert - if l1_global_sf.dim() == 2 and l1_global_sf.shape[1] == 2: - # gate_gs and up_gs per expert — use gate_gs for the GEMM alpha, - # then correct the up half post-GEMM - l1_gate_gs = l1_global_sf[:, 0] # (E,) float32 - l1_up_gs = l1_global_sf[:, 1] # (E,) float32 - per_expert_alpha = alpha * l1_gate_gs # (E,) float32 - up_correction = l1_up_gs / l1_gate_gs # (E,) float32 — ratio to apply to up half - else: - per_expert_alpha = alpha * l1_global_sf # (E,) float32 - up_correction = None - - slot_out, slot_token = cutlass_grouped_nvfp4_gemm( - x_fp4, x_sf_fp8, - l1_weights, w_sf_fp8, - slot_expert_ids, - slot_token, - per_expert_alpha=per_expert_alpha, - ) - - # Apply up correction if gate/up global scales differ - if up_correction is not None: - gate_N = N // 2 - # For each slot, apply the correction to the up half - # slot_out is (num_slots, N) — up half is [:, gate_N:] - # Correction factor is per-expert: up_correction[slot_expert_ids] - correction = up_correction[slot_expert_ids].unsqueeze(1) # (num_slots, 1) - slot_out[:, gate_N:] = slot_out[:, gate_N:] * correction.to(slot_out.dtype) - - print(f"[L1-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e}") - return slot_out, slot_token - - -def nvfp4_mega_moe_l2( - x_fp4, # (num_slots, INTER//2) int8 packed E2M1 — already slot-major - x_sf, # (num_slots, sf_k_groups) float8_e4m3fn - l2_weights, # (E_per_rank, INTER//2, HIDDEN) int8, column-major for CUTLASS - l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major - slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs - l2_global_sf, # (E_per_rank,) float32 — weight global scales - alpha=1.0, # fp32 scalar from stage_activation global scale -): - """L2 GEMM: down_proj — slot-based, no routing weights. - - Per-expert alpha = activation_global_scale * weight_global_scale[expert]. - This preserves full float32 precision — no lossy float8 folding. - """ - K_half = x_fp4.shape[1] - K = K_half * 2 - N = l2_weights.shape[2] - - if MEGA_MOE_DEBUG: - print(f"[nvfp4_moe_l2] slots={x_fp4.shape[0]} K={K} N={N} native=1") - - x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf - w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales - assert w_sf_fp8.dtype == torch.float8_e4m3fn, f"l2_scales after unpack dtype={w_sf_fp8.dtype}" - - # Per-expert alpha: activation_gs * weight_gs - per_expert_alpha = alpha * l2_global_sf # (E,) float32 - - slot_out, _ = cutlass_grouped_nvfp4_gemm( - x_fp4, x_sf_fp8, - l2_weights, w_sf_fp8, - slot_expert_ids, - per_expert_alpha=per_expert_alpha, - ) - print(f"[L2-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e} nan={torch.isnan(slot_out).any().item()}") - return slot_out # (num_slots, HIDDEN) bfloat16 - - -# E2M1 (FP4) representable magnitudes: {0, 0.5, 1, 1.5, 2, 3, 4, 6} -# Bit patterns (3-bit, no sign): 000=0, 001=0.5, 010=1, 011=1.5, 100=2, 101=3, 110=4, 111=6 -# Full 4-bit nibble: bit 3 = sign, bits 2:0 = magnitude index -_E2M1_MAGNITUDES = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32) - - -def _quantize_to_e2m1(x_f32): - """Quantize float32 values to E2M1 (FP4) nibble indices. - - Maps each value to the nearest E2M1 representable magnitude, - then packs as 4-bit sign-magnitude nibbles. - - Returns (nibbles, scales) where: - nibbles: (..., N) uint8 with 4-bit sign-magnitude per value - scales: (..., N//16) float8_e4m3fn block scales - """ - *batch, N = x_f32.shape - assert N % 16 == 0, f"Last dim {N} not divisible by 16 (block size)" - - # Reshape into blocks of 16 for block-wise scaling - x_blocks = x_f32.reshape(*batch, N // 16, 16) - - # Per-block absmax determines the scale - block_max = x_blocks.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8) - - # Scale so that the max maps to 6.0 (largest E2M1 magnitude) - scale_f32 = (block_max / 6.0).clamp(min=1e-8, max=448.0) - x_scaled = x_blocks / scale_f32.clamp(min=1e-8) - - # Find nearest E2M1 magnitude for each value - signs = torch.sign(x_scaled) - abs_scaled = x_scaled.abs() - - mags = _E2M1_MAGNITUDES.to(device=abs_scaled.device) - dists = (abs_scaled.unsqueeze(-1) - mags).abs() - idx = dists.argmin(dim=-1) - - idx = idx.clamp(0, 7).to(torch.uint8) - - sign_bit = (signs < 0).to(torch.uint8) - nibbles = (sign_bit << 3) | idx - - nibbles = nibbles.reshape(*batch, N // 2, 2) - packed = (nibbles[..., 1] << 4) | nibbles[..., 0] - - sf = scale_f32.squeeze(-1).to(torch.float8_e4m3fn) - - return packed.to(torch.int8), sf - - -def stage_activation(x_bf16, input_global_scale=None): - """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales. - - Two-level quantization matching the NVFP4 weight format: - 1. Per-tensor global scale: amax / (6.0 * 448.0) [default] or provided - 2. Per-block (16 values) absmax scaling on the normalized values - - Args: - x_bf16: BF16 activation tensor - input_global_scale: If provided, use this as the activation global scale - instead of computing dynamically. WARNING: this is the amax/(6*448) - normalization scale, NOT the checkpoint's input_scale (which is a - different quantity used for alpha computation). Pass None to compute - dynamically from data. - - Returns (x_fp4, x_sf, input_global_scale) where: - x_fp4: packed E2M1 nibbles - x_sf: UE4M3 block scales (NOT folded with global scale) - input_global_scale: fp32 per-tensor scale, applied as GEMM alpha - """ - x_f32 = x_bf16.float() - - if input_global_scale is None: - x_amax = x_f32.abs().amax().to(torch.float32).clamp(min=1e-8) - input_global_scale = x_amax / (6.0 * 448.0) - - x_normalized = x_f32 / input_global_scale - - x_fp4, x_sf = _quantize_to_e2m1(x_normalized) - - return x_fp4, x_sf, input_global_scale - - -def nvfp4_mega_moe_full( - y, # output tensor (num_tokens, HIDDEN) bfloat16 - transformed_l1_weights, # (l1_w, l1_sf, l1_global_sf) from finalize_weights - transformed_l2_weights, # (l2_w, l2_sf, l2_global_sf) from finalize_weights - symm_buffer, # SymmBuffer from get_symm_buffer - activation_clamp=None, # optional clamp value (unused in NVFP4) - fast_math=False, # fast math flag (unused in NVFP4) - l1_input_scale=None, # (num_experts,) float32 — checkpoint input_scale for L1 (w13) - l2_input_scale=None, # (num_experts,) float32 — checkpoint input_scale for L2 (w2) -): - """Full mega_moe forward pass — replaces deep_gemm.mega.fp8_nvfp4_mega_moe. - - Slot-based pipeline (routing weights applied ONCE at final scatter): - 1. Read staged activation from symm_buffer - 2. L1 GEMM → slot output (num_slots, 2*INTER) — per-expert alpha - 3. SiLU + Mul PER SLOT (nonlinearity before combining expert paths) - 4. Quantize activated slots → FP4 - 5. L2 GEMM → slot output (num_slots, HIDDEN) — per-expert alpha - 6. Final scatter: y.index_add_(0, slot_token, slot_weight * l2_slots) - Single routing weight application. - """ - num_tokens = y.shape[0] - device = y.device - dtype = y.dtype - - if MEGA_MOE_STATIC: - if MEGA_MOE_DEBUG: - print(f"[MEGA_MOE_STATIC] Skipping nvfp4_mega_moe, returning zeros " - f"shape=({num_tokens}, {y.shape[1]})") - y.zero_() - return - - # Unpack transformed weights (now includes global_sf) - l1_w, l1_sf, l1_global_sf = transformed_l1_weights - l2_w, l2_sf, l2_global_sf = transformed_l2_weights - - # Expert sanity check — are experts actually distinct? - if not getattr(nvfp4_mega_moe_full, '_expert_sanity', False): - nvfp4_mega_moe_full._expert_sanity = True - for e in range(min(4, l1_w.shape[0])): - w_sample = l1_w[e].view(torch.uint8)[:8, :8] - sf_sample = l1_sf[e].to(torch.float32)[:4, :4] - print(f"[EXPERT-SANITY e={e}] w_bytes[:8,:8]={w_sample.flatten().tolist()[:16]}") - print(f"[EXPERT-SANITY e={e}] sf[:4,:4]={sf_sample.flatten().tolist()[:8]}") - print(f"[EXPERT-SANITY e={e}] l1_global_sf={l1_global_sf[e].tolist()}") - print(f"[EXPERT-SANITY e={e}] l2_global_sf={l2_global_sf[e].tolist()}") - - # Step 1: Read staged activation from symm_buffer - x_fp4 = symm_buffer.x[:num_tokens] - x_sf = symm_buffer.x_sf[:num_tokens] - l1_global_scale = symm_buffer.input_global_scale - - # Diagnostic: check FP4 quantization quality by dequantizing and comparing - if not getattr(nvfp4_mega_moe_full, '_quant_diag', False): - nvfp4_mega_moe_full._quant_diag = True - # Dequantize: FP4 → BF16 round-trip check - x_u8 = x_fp4.view(torch.uint8) - lo = (x_u8 & 0x0F).to(torch.int8) # low nibble - hi = ((x_u8 >> 4) & 0x0F).to(torch.int8) # high nibble - # Interleave back to (num_tokens, K) - x_nibbles = torch.stack([lo, hi], dim=-1).reshape(num_tokens, -1) # (T, K) - signs = (x_nibbles >> 3).float() * -2 + 1 # +1 or -1 - mags = _E2M1_MAGNITUDES.to(device=x_nibbles.device)[(x_nibbles & 0x07).long()] - x_deq = signs * mags # (T, K) in E2M1 magnitudes - # Apply block scales and global scale - sf_expanded = x_sf.to(torch.float32).repeat_interleave(16, dim=-1) # (T, K) - igs = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale - x_reconstructed = x_deq * sf_expanded * igs - print(f"[QUANT-DIAG] x_fp4 amax={x_fp4.view(torch.uint8).float().amax():.0f} " - f"x_sf range=[{x_sf.to(torch.float32).min():.2f}, {x_sf.to(torch.float32).max():.2f}] " - f"igs={igs:.4e}") - print(f"[QUANT-DIAG] reconstructed amax={x_reconstructed.abs().max():.4e} " - f"mean={x_reconstructed.mean():.4e}") - topk_ids = symm_buffer.topk_idx[:num_tokens] - topk_weights = symm_buffer.topk_weights[:num_tokens] - - _x_sf_f32 = x_sf.to(torch.float32) - _igs = l1_global_scale if isinstance(l1_global_scale, float) else l1_global_scale.item() if hasattr(l1_global_scale, 'item') else float(l1_global_scale) - if MEGA_MOE_DEBUG: - print(f"[ALPHA L1] activation_gs={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.max().item():.4e}]") - print(f"[ALPHA L1] l1_global_sf range [{l1_global_sf.min().item():.4e}, {l1_global_sf.max().item():.4e}]") - - # Convert global expert IDs to local expert IDs - num_experts_per_rank = l1_w.shape[0] - experts_start_idx = symm_buffer.experts_start_idx - topk_ids_local = topk_ids - experts_start_idx - - # Build slot mapping for this rank - local_topk = (topk_ids >= experts_start_idx) & (topk_ids < experts_start_idx + num_experts_per_rank) - slot_token, slot_k = local_topk.nonzero(as_tuple=True) - slot_expert_local = topk_ids_local[slot_token, slot_k] - slot_weight = topk_weights[slot_token, slot_k] - num_slots = slot_token.shape[0] - - tokens_routed_locally = local_topk.any(dim=-1).sum().item() - print(f"[ROUTING] tokens_routed_local={tokens_routed_locally}/{num_tokens} " - f"num_slots={num_slots}") - - if MEGA_MOE_DEBUG: - print(f"[nvfp4_mega_moe_full] x_fp4={x_fp4.shape} x_sf={x_sf.shape} " - f"topk_ids range: {topk_ids.min().item()}-{topk_ids.max().item()} " - f"local: {topk_ids_local.min().item()}-{topk_ids_local.max().item()} " - f"slots={num_slots}") - - # Handle no local slots - if num_slots == 0: - y.zero_() - return - - # Ensure alpha is a plain Python float for the base activation global scale - l1_alpha = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale - - # Shape consistency asserts - assert slot_expert_local.ndim == 1 - assert slot_token.ndim == 1 - assert slot_weight.ndim == 1 - assert slot_expert_local.numel() == num_slots - assert slot_token.numel() == num_slots - assert slot_weight.numel() == num_slots - - # BF16 reference: dequantize and run BF16 GEMM for the first slot to compare - if not getattr(nvfp4_mega_moe_full, '_ref_diag', False): - nvfp4_mega_moe_full._ref_diag = True - try: - s0 = slot_token[0].item() - e0 = slot_expert_local[0].item() - # Dump raw GEMM inputs for expert e0 - print(f"[GEMM-DEBUG] expert={e0} s0={s0}") - print(f"[GEMM-DEBUG] x_fp4[s0] first 8 bytes: {x_fp4[s0].view(torch.uint8)[:8].tolist()}") - print(f"[GEMM-DEBUG] x_sf[s0] first 8: {x_sf[s0].to(torch.float32)[:8].tolist()}") - print(f"[GEMM-DEBUG] l1_w[e0] first 8 bytes: {l1_w[e0].view(torch.uint8).flatten()[:8].tolist()}") - print(f"[GEMM-DEBUG] l1_sf[e0] first 8: {l1_sf[e0].to(torch.float32).flatten()[:8].tolist()}") - print(f"[GEMM-DEBUG] l1_global_sf[e0]: {l1_global_sf[e0].tolist()} shape={l1_global_sf[e0].shape}") - print(f"[GEMM-DEBUG] l1_alpha (igs): {l1_alpha:.6e}") - - # Dequantize activation - x_u8 = x_fp4[s0].view(torch.uint8) - lo = (x_u8 & 0x0F).long() - hi = ((x_u8 >> 4) & 0x0F).long() - x_nib = torch.stack([lo, hi], dim=-1).reshape(-1) # (K,) — 1D so simple flatten works - x_signs = (x_nib >> 3).float() * -2 + 1 - x_mags = _E2M1_MAGNITUDES.to(device=x_u8.device)[(x_nib & 0x07)] - x_deq = x_signs * x_mags # (K,) = (7168,) - sf_exp = x_sf[s0].to(torch.float32).repeat_interleave(16, dim=-1) # (K,) - # Dequantize L1 weight for expert e0 - w_u8 = l1_w[e0].view(torch.uint8) - wlo = (w_u8 & 0x0F).long() - whi = ((w_u8 >> 4) & 0x0F).long() - w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0] * 2, w_u8.shape[1]) # (K, N) - w_signs = (w_nib >> 3).float() * -2 + 1 - w_mags = _E2M1_MAGNITUDES.to(device=w_u8.device)[(w_nib & 0x07)] - w_deq = w_signs * w_mags # (K, N) = (7168, 6144) - w_sf_exp = l1_sf[e0].to(torch.float32).repeat_interleave(16, dim=0) # (K, N) - # Full dequant: x = e2m1 * block_sf * igs, w = e2m1 * block_sf * gs - gs = l1_global_sf[e0] # shape (2,) or scalar - igs = l1_alpha # already the input global scale - x_full = (x_deq * sf_exp * igs).to(torch.bfloat16) # (K,) - w_full = (w_deq * w_sf_exp).to(torch.bfloat16) # (K, N) without gs - ref_out = torch.nn.functional.linear(x_full.unsqueeze(0), w_full.T).squeeze(0) # (N,) - # Apply per-half global scale (gate_gs for first half, up_gs for second half) - gn = ref_out.shape[0] // 2 - gs_vals = gs.detach().cpu().tolist() - if isinstance(gs_vals, float) or len(gs_vals) == 1: - ref_out = ref_out * (gs_vals if isinstance(gs_vals, float) else gs_vals[0]) - else: - ref_out[:gn] = ref_out[:gn] * gs_vals[0] - ref_out[gn:] = ref_out[gn:] * gs_vals[1] - nvfp4_mega_moe_full._ref_l1 = (s0, e0, ref_out) - print(f"[BF16-REF-L1] expert={e0} amax={ref_out.abs().max():.4e} mean={ref_out.mean():.4e}") - except Exception as ex: - import traceback - print(f"[BF16-REF-L1] FAILED: {ex}") - traceback.print_exc() - - # Step 2: L1 GEMM — slot-based, per-expert alpha - l1_slots, _ = nvfp4_mega_moe_l1( - x_fp4, x_sf, l1_w, l1_sf, - slot_expert_local, slot_token, - l1_global_sf=l1_global_sf, - alpha=l1_alpha, - ) # (num_slots, 2*INTER) bfloat16 - - # Compare L1 NVFP4 output to BF16 reference - if hasattr(nvfp4_mega_moe_full, '_ref_l1') and not getattr(nvfp4_mega_moe_full, '_ref_comp', False): - nvfp4_mega_moe_full._ref_comp = True - try: - s0, e0, ref = nvfp4_mega_moe_full._ref_l1 - nvfp4_out = l1_slots[0].float() - ref_f = ref.float() - cos = torch.nn.functional.cosine_similarity(nvfp4_out.unsqueeze(0), ref_f.unsqueeze(0)).item() - mse = (nvfp4_out - ref_f).pow(2).mean().item() - print(f"[COSINE-L1] expert={e0} cosine={cos:.6f} mse={mse:.4e} nvfp4_amax={nvfp4_out.abs().max():.4e} ref_amax={ref_f.abs().max():.4e}") - # Dump first 8 output values from each - print(f"[NVFP4-OUT-8] {nvfp4_out[:8].tolist()}") - print(f"[REF-OUT-8] {ref_f[:8].tolist()}") - except Exception as ex: - print(f"[COSINE-L1] FAILED: {ex}") - - # Post-L1 shape asserts - assert l1_slots.shape[0] == num_slots - - if MEGA_MOE_DEBUG: - print(f"[L1-out] nan={torch.isnan(l1_slots).any().item()} " - f"abs_max={l1_slots.abs().max().item():.4e}") - - # Step 3: SiLU + Mul PER SLOT — nonlinearity before combining paths - gate, up = l1_slots.chunk(2, dim=-1) - print(f"[L1-SPLIT] gate amax={gate.abs().max().item():.4e} mean={gate.float().mean().item():.4e} | up amax={up.abs().max().item():.4e} mean={up.float().mean().item():.4e}") - activated = torch.nn.functional.silu(gate) * up - print(f"[SILU-ACT] amax={activated.abs().max().item():.4e} mean={activated.float().mean().item():.4e} nan={torch.isnan(activated).any().item()}") - if activation_clamp is not None: - activated = activated.clamp(max=activation_clamp) - - # Step 4: Quantize activated slots → FP4 - l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated) - - # Pre-L2 shape asserts - assert activated.shape[0] == num_slots - assert l1_fp4.shape[0] == num_slots - assert l1_sf_out.shape[0] == num_slots - l2_alpha = float(l2_global_scale) if not isinstance(l2_global_scale, float) else l2_global_scale - - if MEGA_MOE_DEBUG: - _l1sf_f32 = l1_sf_out.to(torch.float32) - _l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item() - print(f"[ALPHA L2] activation_gs={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}]") - print(f"[ALPHA L2] l2_global_sf range [{l2_global_sf.min().item():.4e}, {l2_global_sf.max().item():.4e}]") - - # Step 5: L2 GEMM — slot-based, per-expert alpha - l2_slots = nvfp4_mega_moe_l2( - l1_fp4, l1_sf_out, l2_w, l2_sf, - slot_expert_local, - l2_global_sf=l2_global_sf, - alpha=l2_alpha, - ) # (num_slots, HIDDEN) bfloat16 - - if MEGA_MOE_DEBUG: - print(f"[L2-out] nan={torch.isnan(l2_slots).any().item()} " - f"abs_max={l2_slots.abs().max().item():.4e}") - - # Step 6: Final scatter — routing weights applied ONCE - y.zero_() - y.index_add_( - 0, - slot_token, - l2_slots * slot_weight.to(l2_slots.dtype).unsqueeze(1), - ) - print(f"[SCATTER] y amax={y.abs().max().item():.4e} mean={y.float().mean().item():.4e} nan={torch.isnan(y).any().item()} slots={num_slots}") diff --git a/src/nvfp4_megamoe_kernel/symm_buffer.py b/src/nvfp4_megamoe_kernel/symm_buffer.py deleted file mode 100644 index d974f1ec..00000000 --- a/src/nvfp4_megamoe_kernel/symm_buffer.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Symmetric buffer for NVLink cross-rank all-reduce in mega_moe. - -Replaces deep_gemm.mega.SymmBuffer and get_symm_buffer_for_nvfp4_mega_moe. -API matches the DeepGEMM signature used in the vLLM deepseek_v4.py patch. -""" - -import os -import torch -import torch.distributed as dist - -MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) - - -class SymmBuffer: - """Symmetric NVLink buffer for expert-parallel cross-rank communication. - - Matches the DeepGEMM SymmBuffer interface expected by the vLLM patch: - - .x: staged activation (FP4 packed) - - .x_sf: staged activation scales (UE4M3 packed) - - .topk_idx: top-k expert indices - - .topk_weights: top-k routing weights - - .buffer: underlying CUDA buffer - - .group: process group - """ - - def __init__(self, group, num_experts, max_num_tokens, top_k, - hidden_size, intermediate_size): - self.group = group - self.num_experts = num_experts - self.max_num_tokens = max_num_tokens - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.experts_start_idx = 0 # set by caller before kernel invocation - - device = torch.cuda.current_device() - - # NVFP4 packed E2M1: 2 FP4 values per byte → K//2 bytes per token. - # Scales are UE4M3 (float8_e4m3fn), one per 16-element group → K//16 - # bytes per token, UNPACKED. This is what `stage_activation` produces - # and what the CUTLASS NVFP4 block-scaled GEMM consumes directly. - # (The DeepGEMM API packed 4 UE4M3 into one uint32 — we don't, because - # our CUTLASS kernel reads scales as float8_e4m3fn.) - sf_k_groups_hidden = hidden_size // 16 - sf_k_groups_inter = intermediate_size // 16 - - # Staging buffers - self.x = torch.empty( - max_num_tokens, hidden_size // 2, - dtype=torch.int8, device=device, - ) - self.x_sf = torch.empty( - max_num_tokens, sf_k_groups_hidden, - dtype=torch.float8_e4m3fn, device=device, - ) - self.topk_idx = torch.empty( - max_num_tokens, top_k, - dtype=torch.int32, device=device, - ) - self.topk_weights = torch.empty( - max_num_tokens, top_k, - dtype=torch.float32, device=device, - ) - - # All-reduce buffer - self.buffer = torch.empty( - max_num_tokens, hidden_size, - dtype=torch.bfloat16, device=device, - ) - - # Per-tensor global scale from stage_activation (fp32 scalar) - # Applied as GEMM alpha: D = global_scale * (A_sf * A_fp4) @ (B_sf * B_fp4) - self.input_global_scale = 1.0 - - if MEGA_MOE_DEBUG: - print(f"[SymmBuffer] x={self.x.shape} x_sf={self.x_sf.shape} " - f"topk_idx={self.topk_idx.shape} topk_weights={self.topk_weights.shape} " - f"buffer={self.buffer.shape}") - - -def get_symm_buffer_for_nvfp4_mega_moe( - group, - num_experts: int, - max_num_tokens: int, - top_k: int, - hidden_size: int, - intermediate_size: int, -) -> SymmBuffer: - """Allocate a symmetric buffer for the NVFP4 mega_moe kernel. - - API matches deep_gemm.mega.get_symm_buffer_for_nvfp4_mega_moe. - """ - return SymmBuffer( - group, num_experts, max_num_tokens, top_k, - hidden_size, intermediate_size, - ) \ No newline at end of file diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py deleted file mode 100644 index 78aa561e..00000000 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ /dev/null @@ -1,108 +0,0 @@ -""" -NVFP4 Weight Transformation for CUTLASS mega_moe kernel. - -Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale) -into the format expected by the CUTLASS block-scaled GEMM kernel: -- Packed FP4 weights (int8, K-major) -- UE4M3 block scales (float8_e4m3fn, row-major — CUTLASS SF remap handles interleaving) -- float32 global scales (NOT folded into block scales — passed separately for per-expert alpha) - -Previous versions folded weight_scale_2 into block scales via float8 round-trip, which caused -25% relative error (product of ~56-448 block_sf × ~4.65e-05 global_sf lands in the low-precision -zone of float8_e4m3fn where step size is 25%). The global scale is now applied as a per-expert -multiplier to the GEMM alpha, preserving full float32 precision. - -Call signature matches the nightly vLLM deepseek_v4.py finalize_weights: - transform_nvfp4_weights_for_mega_moe( - (l1_weight, l1_weight_scale), - (l2_weight, l2_weight_scale), - l1_weight_scale_2=..., - l2_weight_scale_2=..., - ) -""" - -import torch - - -def transform_nvfp4_weights_for_mega_moe( - l1_tuple: tuple[torch.Tensor, torch.Tensor], # (weight, weight_scale) - l2_tuple: tuple[torch.Tensor, torch.Tensor], # (weight, weight_scale) - l1_weight_scale_2: torch.Tensor = None, # float32 global scale for L1 - l2_weight_scale_2: torch.Tensor = None, # float32 global scale for L2 -) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Transform NVFP4 weights for the CUTLASS block-scaled GEMM. - - NO LONGER FOLDS GLOBAL SCALES INTO BLOCK SCALES. - Folding block_sf (float8) × global_sf (float32) → float8 loses ~25% precision - because the product lands in the low-precision zone of float8_e4m3fn. - Instead, global scales are returned separately and applied as per-expert GEMM alpha. - - Args: - l1_tuple: (w13_weight, w13_weight_scale) — gate_up proj - l2_tuple: (w2_weight, w2_weight_scale) — down proj - l1_weight_scale_2: global scale for L1 (float32) - Shape (E, 2) for gate+up, or (E,) per-expert, or scalar - l2_weight_scale_2: global scale for L2 (float32) - Shape (E,) per-expert, or scalar - - Returns: - ((l1_weight, l1_sf, l1_global_sf), (l2_weight, l2_sf, l2_global_sf)) - where global_sf is (E,) float32 — the geometric mean of gate/up for L1, - or the per-expert global scale for L2. - The caller must apply global_sf as a per-expert multiplier to the GEMM alpha. - """ - l1_weight, l1_weight_scale = l1_tuple - l2_weight, l2_weight_scale = l2_tuple - - # Extract global scales as per-expert float32 vectors - # L1: gate/up have separate global scales — store both - # The caller (nvfp4_mega_moe_full) will apply the right one per-expert - if l1_weight_scale_2 is not None: - l1_gs = l1_weight_scale_2.to(torch.float32) - if l1_gs.dim() == 2 and l1_gs.shape[1] == 2: - # (E, 2) — gate_gs and up_gs separate - # For L1 alpha, use the geometric mean (close enough since gate and up - # global scales are typically similar). Actually, we need BOTH because - # the GEMM produces gate and up in one shot. - # Better: just store (E, 2) and let the caller apply post-GEMM scaling. - l1_global_sf = l1_gs # (E, 2) float32 - else: - l1_global_sf = l1_gs # (E,) float32 - else: - l1_global_sf = torch.ones(l1_weight.shape[0], dtype=torch.float32, device=l1_weight.device) - - if l2_weight_scale_2 is not None: - l2_gs = l2_weight_scale_2.to(torch.float32) - l2_global_sf = l2_gs # (E,) or scalar → broadcast to (E,) - if l2_global_sf.dim() == 0: - l2_global_sf = l2_global_sf.expand(l2_weight.shape[0]) - else: - l2_global_sf = torch.ones(l2_weight.shape[0], dtype=torch.float32, device=l2_weight.device) - - # Debug: one-time diagnostic - if not getattr(transform_nvfp4_weights_for_mega_moe, '_diag', False): - transform_nvfp4_weights_for_mega_moe._diag = True - print(f"[WT-XFORM] L1 block_sf range=[{l1_weight_scale.float().min():.4e}, " - f"{l1_weight_scale.float().max():.4e}] unique={torch.unique(l1_weight_scale.view(torch.uint8)).numel()}") - print(f"[WT-XFORM] L1 global_sf: shape={tuple(l1_global_sf.shape)} " - f"range=[{l1_global_sf.min():.4e}, {l1_global_sf.max():.4e}]") - print(f"[WT-XFORM] L2 block_sf range=[{l2_weight_scale.float().min():.4e}, " - f"{l2_weight_scale.float().max():.4e}] unique={torch.unique(l2_weight_scale.view(torch.uint8)).numel()}") - print(f"[WT-XFORM] L2 global_sf: shape={tuple(l2_global_sf.shape)} " - f"range=[{l2_global_sf.min():.4e}, {l2_global_sf.max():.4e}]") - - # Block scales stay as original float8 — NO FOLDING - l1_sf_out = l1_weight_scale.contiguous() - l2_sf_out = l2_weight_scale.contiguous() - - # CUTLASS B is declared ColumnMajor — it expects (K, N) in memory. - # Checkpoint weights are (N, K_half) row-major, so we transpose to (K_half, N) - l1_weight_out = l1_weight.transpose(-2, -1).contiguous() - l2_weight_out = l2_weight.transpose(-2, -1).contiguous() - - # Same for scale factors: (N, sf_k) row-major → (sf_k, N) column-major - l1_sf_out = l1_sf_out.transpose(-2, -1).contiguous() - l2_sf_out = l2_sf_out.transpose(-2, -1).contiguous() - - return (l1_weight_out, l1_sf_out, l1_global_sf), (l2_weight_out, l2_sf_out, l2_global_sf) diff --git a/tests/cudagraph_test.py b/tests/cudagraph_test.py index 5d80f04a..6b949d7d 100644 --- a/tests/cudagraph_test.py +++ b/tests/cudagraph_test.py @@ -111,7 +111,7 @@ def make_dummy_runner(num_experts=32, hidden_size=7168, intermediate_size=3072, return torch.randint(0, 256, shape, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) def rand_sf(*shape, device="cuda"): - return torch.rand(shape, dtype=torch.float8_e4m3fn, device=device) + return torch.rand(shape, dtype=torch.float16, device=device).to(torch.float8_e4m3fn) l1_fp4 = [rand_fp4(3584, intermediate_size, device=device) for _ in range(num_experts)] l1_sf = [rand_sf(3584 // 16, intermediate_size * 2, device=device) for _ in range(num_experts)]