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

bug fixes

Browse Source
This commit is contained in:
biondizzle
2026-05-15 09:25:33 +00:00
parent a780bb5fde
commit f29b96de09
3 changed files with 77 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files

344
NEXT_OPTIMIZATION.md
View File

@@ -1,344 +0,0 @@
pre-remap the weight scale factors (SFB) once, and stop remapping/allocating them inside every GEMM.
This is lower hanging than rewriting stage_activation in Triton, because weights are static after load. Activation scales SFA still need per-call remap, but SFB does not.
Current hot path in cutlass_nvfp4_gemm.cu:
```
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);
remap_sf_to_cutlass_kernel<<<...>>>(SFA_ptr, sfa_cutlass.get(), ...);
remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...);
```
Target shape
Keep this per GEMM:
```
SFA row-major activation scales
→ remap dynamically
```
Change this:
```
SFB weight scales
→ already CUTLASS-remapped
→ pass directly to GEMM
```
So the GEMM run becomes:
```
// still dynamic
remap SFA
// no remap
use prepacked_SFB_cutlass directly
```
Step 1: add a prepack_sfb CUDA entrypoint
Add a new exported function next to cutlass_nvfp4_gemm_run.
Something like:
```
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(layout_SFB);
int K_sf = K / InputSFVectorSize;
cudaMemsetAsync(
static_cast<ElementSF*>(SFB_cutlass_ptr),
0,
sfb_size * sizeof(ElementSF),
stream);
int block = 256;
remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>(
static_cast<const ElementSF*>(SFB_ptr),
static_cast<ElementSF*>(SFB_cutlass_ptr),
layout_SFB,
N,
K_sf,
true // SFB source is (K_sf, N)
);
return 0;
}
```
You’ll also want a size query helper:
```
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(layout_SFB);
return 0;
}
```
Important: pass the same M, N, K geometry you’ll use for the GEMM at first. Later you can test whether SFB layout is independent of M; I suspect it is effectively N/K-driven, but don’t assume that until you print sizes for several M.
Step 2: expose it in pytorch_binding.cpp
Add declarations:
```
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
);
```
Then add a Python-visible wrapper:
```
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;
}
```
Register it:
```
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &cutlass_nvfp4_gemm_forward,
"CUTLASS NVFP4 block-scaled GEMM forward");
m.def("prepack_sfb", &prepack_sfb,
"Pre-remap SFB weight scales into CUTLASS layout");
}
```
Step 3: add a GEMM path that accepts prepacked SFB
Add a second C API entrypoint, or a boolean flag, for:
```
cutlass_nvfp4_gemm_run_prepacked_sfb(...)
```
Inside it, delete this:
```
cutlass::device_memory::allocation<ElementSF> sfb_cutlass(sfb_size);
cudaMemsetAsync(sfb_cutlass.get(), 0, ...);
remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...);
```
and use the passed pointer directly:
```
static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
```
The relevant section becomes:
```
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream);
remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>(
static_cast<const ElementSF*>(SFA_ptr),
sfa_cutlass.get(),
layout_SFA,
M,
K_sf,
false);
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
}
};
```
Step 4: prepack in Python after weight transform
In weight_transform.py, you currently return:
```
return (l1_weight_out, l1_sf_out), (l2_weight_out, l2_sf_out)
```
Do not do the prepack directly there unless extension import order is clean. Safer first pass: do it lazily in nvfp4_mega_moe_full() once.
Add helper:
```
def _maybe_prepack_weight_sf(weights, weight_sf, N, K, tag):
cache_attr = f"_prepacked_{tag}"
cached = getattr(_maybe_prepack_weight_sf, cache_attr, None)
if cached is not None:
return cached
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm import _C
E = weight_sf.shape[0]
packed = []
# Use M=128 initially to match the MMA tile boundary.
# Later test if SFB size/layout is stable across M.
M_for_layout = 128
for e in range(E):
packed.append(
_C.prepack_sfb(
weight_sf[e],
M_for_layout,
N,
K,
)
)
packed = torch.stack(packed, dim=0).contiguous()
setattr(_maybe_prepack_weight_sf, cache_attr, packed)
return packed
```
Then before L1/L2 calls:
```
l1_N = l1_w.shape[2]
l1_K = l1_w.shape[1] * 2
l1_sf_prepacked = _maybe_prepack_weight_sf(l1_w, l1_sf, l1_N, l1_K, "l1")
l2_N = l2_w.shape[2]
l2_K = l2_w.shape[1] * 2
l2_sf_prepacked = _maybe_prepack_weight_sf(l2_w, l2_sf, l2_N, l2_K, "l2")
```
Then pass l1_sf_prepacked / l2_sf_prepacked into the new prepacked-SFB GEMM path.
Validation test
Before trusting it, compare old vs new on one expert:
```
old = cutlass_nvfp4_blockscaled_gemm(
expert_x,
expert_x_sf,
expert_w,
expert_w_sf,
M_expert,
N,
K,
alpha=alpha,
)
sfb_pre = _C.prepack_sfb(expert_w_sf, 128, N, K)
new = cutlass_nvfp4_blockscaled_gemm_prepacked_sfb(
expert_x,
expert_x_sf,
expert_w,
sfb_pre,
M_expert,
N,
K,
alpha=alpha,
)
print("max diff", (old - new).abs().max())
print("cos", torch.nn.functional.cosine_similarity(
old.flatten().float(),
new.flatten().float(),
dim=0,
))
```
Expected:
```
max diff = 0 or tiny BF16-level difference
cos ≈ 1.0
```
If it fails only when M_expert != 128, then LayoutSFB is M-dependent. In that case, cache by an M bucket:
```
bucket_m = ((M_expert + 127) // 128) * 128
cache_key = (tag, bucket_m, N, K)
```
But test first. If SFB layout size and values are stable across M, keep one prepack per expert per layer.
Why this is worth doing now
This removes, per active expert GEMM:
```
1 cudaMalloc/free-ish allocation for SFB
1 cudaMemsetAsync for SFB padding
1 remap kernel launch for SFB
```
Across L1 + L2 and many routed experts, that’s a very clean win without touching routing semantics, quantization, or the actual MMA tile.

97
README.md
View File

@@ -33,22 +33,54 @@ Input hidden states (BF16)
Writes topk_ids / topk_weights to SymmBuffer
│
▼
┌─────────────────────────────────────────┐
│ nvfp4_mega_moe_full │ ← nvfp4_mega_moe.py
│ │
│ 1. Read staged activation from buffer │
│ 2. L1 GEMM: gate_up_proj │ ← CUTLASS NVFP4 block-scaled
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
│ → BF16 output (6144-wide) │
│ 3. SiLU(gate) * up (activation) │
│ 4. stage_activation: BF16 → FP4 │ ← proper E2M1 quantization
│ 5. L2 GEMM: down_proj │ ← CUTLASS NVFP4 block-scaled
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
│ → BF16 output (7168-wide) │
│ 6. Write to output tensor │ ← caller handles cross-rank all-reduce
└─────────────────────────────────────────┘
┌─────────────────────────────────────────────────┐
│ nvfp4_mega_moe_full │ ← nvfp4_mega_moe.py
│ │
│ 1. Read staged activation from buffer │
│ 2. Build slot mapping (token, topk) → local │
│ expert, routing weight │
│ 3. L1 GEMM: gate_up_proj (slot-based) │ ← CUTLASS NVFP4 block-scaled
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
│ → BF16 per-slot output (6144-wide) │
│ 4. SiLU(gate) * up PER SLOT │
│ (nonlinearity before combining paths) │
│ 5. stage_activation: BF16 → FP4 │ ← proper E2M1 quantization
│ 6. L2 GEMM: down_proj (slot-based) │ ← CUTLASS NVFP4 block-scaled
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
│ → BF16 per-slot output (7168-wide) │
│ 7. Final scatter: │
│ y.index_add_(slot_token, │
│ slot_weight * l2_slots) │
│ Routing weight applied ONCE at scatter │
└─────────────────────────────────────────────────┘
│
▼
Cross-rank all-reduce (vLLM built-in)
```
### Slot-Based Dispatch
The kernel uses a **slot representation** instead of collapsing expert outputs early. A slot is one `(token, topk_expert)` pair. For a batch of T tokens with top-6 routing, there are up to 6T slots (fewer if some experts are out of the local rank's range).
**Why slots?** Two bugs in the previous approach:
1. **SiLU after summing is mathematically wrong.** `silu(Σ wᵢ·gateᵢ) * (Σ wᵢ·upᵢ) ≠ Σ wᵢ·silu(gateᵢ)·upᵢ`. The nonlinearity must happen per-expert-path before combining.
2. **Routing weights applied twice.** The old grouped GEMM applied `topk_weights` in its scatter loop, and was called for both L1 and L2 — squaring the weights.
The slot approach fixes both: SiLU+Mul happens per-slot, and routing weights are applied exactly once at the final `index_add_` scatter.
### Prepacked SFB (Weight Scale Factors)
Weight scale factors (SFB) are pre-remapped into CUTLASS interleaved layout once at first forward pass (lazily cached). This eliminates per-GEMM:
- 1 `cudaMalloc`-ish allocation for SFB
- 1 `cudaMemsetAsync` for SFB padding
- 1 remap kernel launch for SFB
The `cutlass_nvfp4_gemm_run_prepacked_sfb` C entry point accepts the prepacked SFB pointer directly. Only SFA (activation scales) is remapped dynamically — those change every forward pass.
---
### vLLM Startup Sequence (how our code plugs in)
```
@@ -85,6 +117,7 @@ Input hidden states (BF16)
6. Profile run (warmup)
└─ First forward pass to allocate KV cache, etc.
└─ This is where the CUTLASS GEMM first executes
└─ SFB weight scales are prepacked into CUTLASS layout (lazy, cached)
7. Ready to serve
```
@@ -96,14 +129,14 @@ Input hidden states (BF16)
```
nvfp4_megamoe_kernel/
├── __init__.py # Public API exports
├── nvfp4_mega_moe.py # Main kernel: nvfp4_mega_moe_full, nvfp4_mega_moe_l1/l2, stage_activation
├── nvfp4_mega_moe.py # Main kernel: nvfp4_mega_moe_full, L1/L2, stage_activation, prepack
├── weight_transform.py # Weight prep: fold global scale, pack UE4M3
├── symm_buffer.py # GPU buffer allocation for MoE dispatch
│
└── cutlass_nvfp4_gemm/ # CUTLASS CUDA extension (the actual hardware kernel)
├── cutlass_nvfp4_gemm.cu # CUDA: CUTLASS GEMM + SF remap kernel
├── pytorch_binding.cpp # PyTorch C++ binding (_C.forward)
├── kernel.py # Python: cutlass_grouped_nvfp4_gemm (per-expert loop)
├── cutlass_nvfp4_gemm.cu # CUDA: CUTLASS GEMM + SF remap + prepack SFB + prepacked-SFB GEMM path
├── pytorch_binding.cpp # PyTorch C++ binding (forward, forward_prepacked_sfb, prepack_sfb)
├── kernel.py # Python: cutlass_grouped_nvfp4_gemm (slot-based, per-expert loop)
├── sf_layout.py # CUTLASS SF layout reference docs
├── setup.py # Build config (nvcc, CUTLASS include paths)
├── build.sh # Build script
@@ -117,9 +150,9 @@ nvfp4_megamoe_kernel/
|------|-------------|--------------|
| `weight_transform.py` | Once at startup (weight loading) | Takes raw NVFP4 checkpoint weights, folds global scales into block scales. Returns scales as `float8_e4m3fn` (not packed uint32). Output: `((l1_w, l1_sf), (l2_w, l2_sf))` |
| `symm_buffer.py` | Once at startup (buffer alloc) | Pre-allocates GPU tensors for activations, scales, routing data, and all-reduce. These persist across forward passes. |
| `nvfp4_mega_moe.py` | Every forward pass | Orchestrates the MoE: reads from symm buffer → L1 GEMM → activation → re-quantize → L2 GEMM → output. Contains `stage_activation` (BF16→FP4 quantize for L1→L2) and `unpack_ue4m3_u32` (uint32 packed scales → float8). |
| `cutlass_nvfp4_gemm/kernel.py` | Every forward pass (called by nvfp4_mega_moe) | Per-expert loop: gather tokens for each expert, call CUTLASS GEMM, scatter results with routing weights. |
| `cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu` | Every forward pass (CUDA kernel) | The actual CUTLASS kernel: native NVFP4 block-scaled GEMM + GPU-side scale factor remap (row-major → CUTLASS interleaved layout). |
| `nvfp4_mega_moe.py` | Every forward pass | Orchestrates the MoE: reads from symm buffer → build slot mapping → L1 GEMM → SiLU+Mul per-slot → re-quantize → L2 GEMM → final index_add_ scatter with routing weights. Contains `stage_activation` (BF16→FP4), `unpack_ue4m3_u32`, and `_prepack_weight_sf` (lazy SFB prepack). |
| `cutlass_nvfp4_gemm/kernel.py` | Every forward pass (called by nvfp4_mega_moe) | Slot-based per-expert loop: gather slots for each expert, call CUTLASS GEMM (with prepacked SFB), write results to slot buffer. No routing weights — caller handles scatter. |
| `cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu` | Every forward pass (CUDA kernel) | The actual CUTLASS kernel: native NVFP4 block-scaled GEMM + GPU-side SFA remap. SFB remap done once at prepack time. Two GEMM entry points: standard (both remap) and prepacked-sfb (SFA remap only). |
| `cutlass_nvfp4_gemm/sf_layout.py` | Reference only | Documents the CUTLASS SfAtom layout. Not used at runtime (remap is in CUDA). |
---
@@ -179,7 +212,7 @@ f7 = 0 — unused
| 0 | 0 | 0 | 0 | T | 0 | 0 | 0 | 0 |
| 1 | 0 | 0 | 0 | T | 1 | 0 | 0 | 0 |
| 4 | 0 | 1 | 0 | T | 0 | 0 | 0 | 0 |
| 16 | 1 | 0 | 0 | T | 0 | 0 | 0 | 0 |
| 16 | 1 | 0 | 0 | T | 0 | 0 |0 | 0 |
| 511 | 31 | 3 | 0 | T | 3 | 0 | 0 | 0 |
| 512 | 0 | 0 | 0 | T | 0 | 1 | 0 | 0 |
| 1024 | 0 | 0 | 0 | T | 0 | 2 | 0 | 0 |
@@ -250,6 +283,16 @@ Additionally, the dest buffer must be zero-initialized before remap because CUTL
**Fix:** Correct extraction: `m = f0 + f1*32 + f2*128`, `k_sf = f4 + f5*4`. Zero-init dest buffer before remap.
**Diagnostic trail:** Constant-scale test (all SF=1.0) → cosine 1.0 proved FP4 path was correct. Real scales → cosine 0.83 proved SF remap was broken. Single-element probes (SFA[0,0] vs SFA[0,3]) proved only k_group=0 worked. Printf dump of flat coordinates at specific indices revealed flat_rank=8 and the correct extraction formula.
### 6. SiLU after summing expert paths (math error)
**File:** `nvfp4_mega_moe.py`
**Bug:** The old grouped GEMM collapsed expert outputs into a weighted sum, then applied SiLU+Mul on the sum. `silu(Σ wᵢ·gateᵢ) * (Σ wᵢ·upᵢ) ≠ Σ wᵢ·silu(gateᵢ)·upᵢ`. The nonlinearity must happen per-expert-path.
**Fix:** Slot-based dispatch — L1 GEMM returns per-slot output, SiLU+Mul applied per-slot, L2 GEMM per-slot, routing weights applied once at final `index_add_` scatter.
### 7. Routing weights applied twice
**File:** `cutlass_nvfp4_gemm/kernel.py`
**Bug:** `cutlass_grouped_nvfp4_gemm` applied `topk_weights` in its scatter loop. Called for both L1 and L2, each expert's contribution was scaled by `topk_weight²`.
**Fix:** GEMM returns per-slot results with no routing weights. Single `y.index_add_(0, slot_token, slot_weight * l2_slots)` at the end.
### Diagnostic: constant-scale test (smoking gun for SF bugs)
When all scale factors are set to UE4M3(1.0):
@@ -282,13 +325,15 @@ The CUTLASS extension builds inside the container during `pip install` of the nv
---
## Known Issues
## Known Issues / TODO
1. ~~**MoE dispatch is slow**~~ — Fixed. Slot-based `index_add_` replaces the Python double loop over tokens×topk. Routing weights applied once at final scatter. Per-expert loop still exists (gather+GEMM) but scatter is vectorized.
1. ~~**MoE dispatch is slow**~~ — Fixed. Slot-based `index_add_` replaces the Python double loop over tokens×topk. Routing weights applied once at final scatter.
2. **stage_activation is Python** — Re-quantization from L1 BF16 output to FP4 for L2 input runs in PyTorch. Should use the Triton staging kernel for speed and consistency with vLLM's built-in staging.
3. **SF remap allocates every call** — `cudaMemset` + remap kernel runs per GEMM invocation. Could pre-compute the CUTLASS-layout buffer once during weight transform.
3. ~~**SF remap allocates every call**~~ — Fixed. SFB weight scales are prepacked into CUTLASS layout once (lazy, cached per layer). Only SFA (activation scales) remapped dynamically.
4. **Per-expert GEMM dispatch is serial Python loop** — The `cutlass_grouped_nvfp4_gemm` iterates over 48 experts in a Python `for` loop. Each iteration launches one CUTLASS GEMM. Could benefit from a true grouped GEMM kernel or CUDA-side expert dispatch.
---
@@ -306,4 +351,4 @@ The CUTLASS extension builds inside the container during `pip install` of the nv
- **Kernel:** `sweetapi.com/biondizzle/nvfp4-megamoe-kernel` (branch: master)
- **Deployment:** `sweetapi.com/biondizzle/deepseek-v4-quant` (branch: modelopt-nvfp4)
- **Local:** `~/dev/nvfp4-mojo/`, `~/dev/deepseek-v4-quant/`
- **Local:** `~/dev/nvfp4-megamoe-kernel/`, `~/dev/deepseek-v4-quant/`

8
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
View File

@@ -359,6 +359,9 @@ def nvfp4_mega_moe_full(
y.zero_()
return
# Ensure alpha is a plain Python float (C extension can't handle torch scalars)
l1_alpha = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale
# Prepack SFB weight scales into CUTLASS layout (lazy, once per layer)
l1_N = l1_w.shape[2]
l1_K = l1_w.shape[1] * 2
@@ -372,7 +375,7 @@ def nvfp4_mega_moe_full(
l1_slots, _ = nvfp4_mega_moe_l1(
x_fp4, x_sf, l1_w, l1_sf_prepacked,
topk_ids_local,
alpha=l1_global_scale,
alpha=l1_alpha,
sfb_prepacked=True,
) # (num_slots, 2*INTER) bfloat16
@@ -392,6 +395,7 @@ def nvfp4_mega_moe_full(
# Step 4: Quantize activated slots → FP4
l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated)
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)
@@ -402,7 +406,7 @@ def nvfp4_mega_moe_full(
l2_slots = nvfp4_mega_moe_l2(
l1_fp4, l1_sf_out, l2_w, l2_sf_prepacked,
topk_ids_local, slot_token,
alpha=l2_global_scale,
alpha=l2_alpha,
sfb_prepacked=True,
) # (num_slots, HIDDEN) bfloat16