fix: use float16->float8 cast for rand_sf (torch.rand doesn't support float8)
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
torch>=2.5
|
||||
tilelang>=0.1
|
||||
@@ -1 +0,0 @@
|
||||
nvfp4_megamoe_kernel
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)
|
||||
"""
|
||||
@@ -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
|
||||
@@ -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<Stages, SchedPipe, AccPipe, ClusterShape, ArchTag>`
|
||||
|
||||
### 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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -1,385 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro MoE
|
||||
**************************************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cute/tensor.hpp>
|
||||
#include <cutlass/tensor_ref.h>
|
||||
#include <cutlass/gemm/dispatch_policy.hpp>
|
||||
#include <cutlass/gemm/collective/collective_builder.hpp>
|
||||
#include <cutlass/epilogue/collective/collective_builder.hpp>
|
||||
#include <cutlass/detail/sm100_blockscaled_layout.hpp>
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/gemm_universal.hpp>
|
||||
#include <cutlass/gemm/kernel/tile_scheduler_params.h>
|
||||
#include <cutlass/float_subbyte.h>
|
||||
#include <cutlass/util/packed_stride.hpp>
|
||||
#include <cutlass/util/device_memory.h>
|
||||
|
||||
#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<cutlass::float_e2m1_t>;
|
||||
using LayoutATag = cutlass::layout::RowMajor;
|
||||
constexpr int AlignmentA = 32;
|
||||
|
||||
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<ElementD>::value;
|
||||
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::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<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
using 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<typename LayoutSF>
|
||||
__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<class LayoutSF>
|
||||
__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<const uint8_t*>(src);
|
||||
auto* dst_u8 = reinterpret_cast<const uint8_t*>(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<ElementSF> sfa_cutlass(sfa_size);
|
||||
cutlass::device_memory::allocation<ElementSF> 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<const ElementSF*>(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<const ElementSF*>(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<const ElementSF*>(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<const ElementSF*>(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<const ArrayElementA*>(A_ptr), stride_A,
|
||||
static_cast<const ArrayElementB*>(B_ptr), stride_B,
|
||||
sfa_cutlass.get(), layout_SFA,
|
||||
sfb_cutlass.get(), layout_SFB
|
||||
},
|
||||
{
|
||||
{ alpha, beta },
|
||||
nullptr, stride_C,
|
||||
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(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<uint8_t> 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<ElementSF*>(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<const ElementSF*>(SFB_ptr),
|
||||
static_cast<ElementSF*>(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<ElementSF> 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<const ElementSF*>(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<const ArrayElementA*>(A_ptr), stride_A,
|
||||
static_cast<const ArrayElementB*>(B_ptr), stride_B,
|
||||
sfa_cutlass.get(), layout_SFA,
|
||||
static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
|
||||
},
|
||||
{
|
||||
{ alpha, beta },
|
||||
nullptr, stride_C,
|
||||
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(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<uint8_t> workspace(workspace_size);
|
||||
|
||||
CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream));
|
||||
CUTLASS_CHECK(gemm.run(stream));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
|
||||
#endif
|
||||
@@ -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
|
||||
@@ -1,131 +0,0 @@
|
||||
/** PyTorch binding for CUTLASS NVFP4 block-scaled GEMM */
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
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<int>(M), static_cast<int>(N), static_cast<int>(K),
|
||||
static_cast<float>(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<int>(M), static_cast<int>(N), static_cast<int>(K),
|
||||
static_cast<float>(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<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(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<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(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");
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
@@ -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<_32, _4>, Shape<SFVecSize, _4>>
|
||||
with Stride<Stride<_16, _4>, 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.
|
||||
"""
|
||||
@@ -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}")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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)]
|
||||
|
||||
Reference in New Issue
Block a user