diff --git a/NEXT_OPTIMIZATION.md b/NEXT_OPTIMIZATION.md deleted file mode 100644 index 107150d1..00000000 --- a/NEXT_OPTIMIZATION.md +++ /dev/null @@ -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 sfa_cutlass(sfa_size); -cutlass::device_memory::allocation sfb_cutlass(sfb_size); - -cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); -cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); - -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(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(SFB_ptr), - static_cast(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(M), - static_cast(N), - static_cast(K), - &size - ); - TORCH_CHECK(rc == 0, "sfb_size failed"); - - auto out = torch::empty( - {size}, - torch::dtype(SFB.dtype()).device(SFB.device()) - ); - - auto stream = c10::cuda::getCurrentCUDAStream(); - - rc = cutlass_nvfp4_prepack_sfb_run( - SFB.data_ptr(), - out.data_ptr(), - static_cast(M), - static_cast(N), - static_cast(K), - stream.stream() - ); - - TORCH_CHECK(rc == 0, "prepack_sfb failed"); - - return out; -} -``` - -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 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(SFB_cutlass_ptr), layout_SFB -``` - -The relevant section becomes: - -``` -cutlass::device_memory::allocation 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(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(A_ptr), stride_A, - static_cast(B_ptr), stride_B, - sfa_cutlass.get(), layout_SFA, - static_cast(SFB_cutlass_ptr), layout_SFB - }, - { - { alpha, beta }, - nullptr, stride_C, - static_cast(D_ptr), stride_D - } -}; -``` - -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. \ No newline at end of file diff --git a/README.md b/README.md index 1da1713a..dab90e23 100644 --- a/README.md +++ b/README.md @@ -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/` diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 2b1dab0c..d4b2fd67 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -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