feat: CUTLASS NVFP4 mega_moe kernel — slot-based L1/L2, source-first SF remap
Major changes from initial TileLang prototype: Kernel: - CUTLASS NVFP4 block-scaled GEMM (SM100 Blackwell, OpClassBlockScaledTensorOp) - Slot-based dispatch: L1 GEMM → SiLU+Mul per-slot → L2 GEMM → index_add scatter - 1D slot_expert_ids passed to both L1 and L2 (no 2D topk_ids rebuild) - slot_token gathered in cutlass_grouped_nvfp4_gemm when provided SF Remap (source-first): - Iterates logical (m, k_sf) source grid, uses layout_sf(make_coord(m, k_sf)) for CUTLASS dest index — no idx2crd/flatten coordinate extraction - 2D kernel launch: dim3 block(32,8), grid over (K_sf, MN) - Uses cute::cosize() for physical allocation size (not cute::size) - SFA: (MN, K_sf) row-major; SFB: (K_sf, MN) row-major (col-major) Weight transform: - UE4M3 unpack with bit reinterpret (not value cast) - Global scale folding (weight_scale_2) for gate/up split - clamp(0,448) → float8_e4m3fn, transpose (N,K)→(K,N) for CUTLASS No prepack cache: - SFB remapped per-call inside CUTLASS (~µs, not the bottleneck) - See README for why prepack cache must never return (OOM, CUDA graphs, M-dependent layout, cross-layer collisions) Stage activation: - Nearest-neighbor E2M1 quantization (no clamp, no uniform steps) - Per-tensor global scale → alpha for L2 GEMM Bug fixes: - _fold_global_scale: removed broken logical_widths branch - unpack_ue4m3_u32: int32 for CUDA bitwise, view not to, ND support - Correct expert param mapping for NVFP4 checkpoint - SiLU applied per-slot (not after summing expert paths)
This commit is contained in:
6
.dockerignore
Normal file
6
.dockerignore
Normal file
@@ -0,0 +1,6 @@
|
||||
docker/
|
||||
scripts/
|
||||
*.egg-info/
|
||||
.git/
|
||||
.gitignore
|
||||
README.md
|
||||
3
.gitignore
vendored
Normal file
3
.gitignore
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.egg-info/
|
||||
50
Dockerfile
Normal file
50
Dockerfile
Normal file
@@ -0,0 +1,50 @@
|
||||
# DeepSeek V4 NVFP4 vLLM + CUTLASS NVFP4 Mega MoE Kernel
|
||||
FROM vllm/vllm-openai:nightly-x86_64
|
||||
|
||||
# Remove broken nixl_ep (built against CUDA 12, image is CUDA 13)
|
||||
RUN pip uninstall -y nixl-ep; rm -rf /usr/local/lib/python3.12/dist-packages/nixl_ep
|
||||
|
||||
RUN apt-get update && apt-get install -y git screen cmake libcusolver-dev-13-0 libcusparse-dev-13-0 libcublas-dev-13-0 libcurand-dev-13-0 libcufft-dev-13-0 libnvjitlink-dev-13-0 && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Remove the broken symlink if it exists
|
||||
RUN rm -f /usr/local/cuda/lib64/libcudart.so.12
|
||||
|
||||
ENV CUDA_HOME=/usr/local/cuda
|
||||
ENV TORCH_CUDA_ARCH_LIST="10.0"
|
||||
|
||||
# Clone latest CUTLASS (has NVFP4 block-scaled MMA support)
|
||||
RUN git clone --depth 1 https://github.com/NVIDIA/cutlass.git /root/cutlass
|
||||
|
||||
ARG CACHE_BUSTER=${TIMESTAMP}
|
||||
|
||||
# Copy and install the NVFP4 mega_moe kernel (from this repo)
|
||||
COPY src/ /root/nvfp4-megamoe-kernel/src/
|
||||
COPY pyproject.toml /root/nvfp4-megamoe-kernel/pyproject.toml
|
||||
RUN cd /root/nvfp4-megamoe-kernel && pip install -e .
|
||||
|
||||
# Build the CUTLASS NVFP4 block-scaled GEMM extension
|
||||
RUN cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm && \
|
||||
mkdir -p cutlass_nvfp4_gemm && \
|
||||
CUTLASS_INCLUDE_DIR=/root/cutlass/include \
|
||||
TORCH_CUDA_ARCH_LIST=10.0 \
|
||||
python3 setup.py build_ext --inplace
|
||||
|
||||
# Install TileLang (for potential future use)
|
||||
RUN pip install tilelang
|
||||
|
||||
ENV PYTHONPATH="/root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm:/root/nvfp4-megamoe-kernel:${PYTHONPATH}"
|
||||
|
||||
# Patch vLLM — overwrite model files and register architecture
|
||||
ARG VLLM_MODELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models
|
||||
ARG VLLM_LAYERS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers
|
||||
|
||||
COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py
|
||||
COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py
|
||||
|
||||
RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \
|
||||
${VLLM_MODELS_DIR}/registry.py
|
||||
|
||||
# Verify
|
||||
RUN python3 -c "import torch; import cutlass_nvfp4_gemm._C; print('CUTLASS NVFP4 OK')" && \
|
||||
python3 -c "import vllm; print('vLLM OK')" && \
|
||||
python3 -c "import nvfp4_megamoe_kernel; print('NVFP4 kernel OK')"
|
||||
383
README.md
383
README.md
@@ -1,40 +1,361 @@
|
||||
# NVFP4 Mega MoE Kernel — Mojo Rewrite
|
||||
# nvfp4-megamoe-kernel
|
||||
|
||||
Rewrite of the DeepGEMM `fp8_nvfp4_mega_moe` kernel in Mojo.
|
||||
Native NVFP4 block-scaled MoE kernel for DeepSeek-V4-Pro on NVIDIA Blackwell (SM100).
|
||||
|
||||
## Why Mojo?
|
||||
- Python-like syntax, C-level performance
|
||||
- Direct GPU programming without PTX inline asm
|
||||
- Safer than CUDA C++ (ownership, borrowing)
|
||||
- Better ergonomics for complex kernel development
|
||||
Replaces the broken `fp8_nvfp4_mega_moe` kernel from DeepGEMM with a working CUTLASS-based implementation that emits real `SM100_MMA_MXF4_SS` tensor core instructions.
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
The kernel performs NVFP4 (E2M1 + UE4M3 block16 scales) matrix multiply
|
||||
for MoE (Mixture of Experts) with expert parallelism across NVLink.
|
||||
DeepSeek-V4-Pro is a 384-expert MoE model with expert parallelism across 8 ranks (B200 GPUs). Each rank handles 48 experts. For each token, the router picks the top-6 experts.
|
||||
|
||||
### Key operations:
|
||||
1. **Staging** — quantize BF16 activation to FP4 (E2M1) with UE8M0 scales
|
||||
2. **TMA load** — load packed FP4 weights and UE4M3 scales from global memory
|
||||
3. **UMMA** — `mxf4nvf4` matrix multiply with block scaling
|
||||
4. **Epilogue** — quantize L1 output (BF16 → FP4 + UE4M3 scales for L2)
|
||||
5. **NVLink sync** — cross-rank barrier and buffer management
|
||||
### The MoE Forward Pass
|
||||
|
||||
### NVFP4 specifics (vs MXFP4):
|
||||
- group_size=16 (UE4M3 block scales), not group_size=32 (UE8M0)
|
||||
- 2 SF K-columns per BLOCK_K (128/16/4=2), not 1
|
||||
- Weights are E2M1 packed int8 (2 values per byte)
|
||||
- `mxf4nvf4` UMMA instruction with `scale_vec::4X`
|
||||
|
||||
## Structure
|
||||
```
|
||||
src/
|
||||
mega_moe.mojo — main kernel entry point
|
||||
staging.mojo — activation quantization (BF16 → FP4)
|
||||
tma.mojo — TMA descriptor creation and copy
|
||||
umma.mojo — UMMA descriptor and MMA operations
|
||||
epilogue.mojo — output quantization and TMA store
|
||||
barrier.mojo — NVLink cluster sync and symm buffer
|
||||
layout.mojo — weight transformation and SF layout
|
||||
utils.mojo — math helpers, UE4M3 packing
|
||||
Input hidden states (BF16)
|
||||
│
|
||||
▼
|
||||
┌─────────────────┐
|
||||
│ Shared Experts │ ← vLLM native FlashInfer CUTLASS NVFP4 path
|
||||
│ (gate + up → │ (not our kernel)
|
||||
│ SiLU * up → │
|
||||
│ down) │
|
||||
└─────────────────┘
|
||||
│
|
||||
▼
|
||||
Staging Kernel (vLLM built-in)
|
||||
BF16 → packed E2M1 (int8) + UE4M3 block-16 scales (uint32)
|
||||
Writes to SymmBuffer.x / SymmBuffer.x_sf
|
||||
│
|
||||
▼
|
||||
Router (vLLM built-in)
|
||||
Writes topk_ids / topk_weights to SymmBuffer
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────┐
|
||||
│ 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.
|
||||
|
||||
### SFB (Weight Scale Factors) — Remapped Per-Call, NOT Cached
|
||||
|
||||
Weight scale factors (SFB) are remapped from row-major to CUTLASS interleaved layout on every GEMM call. This is a lightweight scatter kernel (~µs) and is NOT the bottleneck compared to the GEMM itself.
|
||||
|
||||
⚠️ **DO NOT ADD A PREPACK CACHE FOR SFB.** Previous attempts caused critical issues:
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **OOM** | ~1.75 GiB per prepacked tensor × 61 MoE layers × 2 (L1+L2) = ~214 GiB — exceeds B200 capacity |
|
||||
| **Peak memory 2×** | `torch.stack` held all expert tensors + final stack simultaneously before LRU eviction |
|
||||
| **CUDA graph trap** | LRU eviction frees tensors that CUDA graphs still reference → use-after-free → silent corruption or crash |
|
||||
| **M-dependent layout** | `prepack_sfb(M=128)` assumed SFB layout size is M-independent (never verified). If wrong, entire prepack is invalid |
|
||||
| **Cross-layer cache collision** | Tag-based cache (`"l1"`/`"l2"`) returned layer N-1's data for layer N. Fixed with data_ptr key, but the cache itself was the root problem |
|
||||
|
||||
The per-call remap costs microseconds. The cache cost was hours of debugging. Don't repeat this mistake.
|
||||
|
||||
---
|
||||
|
||||
### vLLM Startup Sequence (how our code plugs in)
|
||||
|
||||
```
|
||||
1. vLLM engine init
|
||||
└─ ModelOptNvFp4Config selected (NVFP4 quantization scheme)
|
||||
└─ FlashInferCutlassNvFp4LinearKernel for linear layers
|
||||
|
||||
2. Model construction
|
||||
└─ DeepseekV4ForCausalLM → DeepseekV4MoE → DeepseekV4DecoderLayer
|
||||
Each layer has: attention + MoE block
|
||||
MoE block has: shared experts + 384 routed experts
|
||||
|
||||
3. Weight loading
|
||||
└─ 95 safetensor shards loaded
|
||||
└─ weight, weight_scale, weight_scale_2 loaded per linear
|
||||
|
||||
4. process_weights_after_loading ← THIS IS WHERE WE HOOK IN
|
||||
└─ ModelOptNvFp4LinearMethod swizzles/pads weights for CUTLASS
|
||||
└─ finalize_mega_moe_weights()
|
||||
└─ weight_transform.py: transform_nvfp4_weights_for_mega_moe()
|
||||
• Folds weight_scale_2 (global scale) into weight_scale (block scale)
|
||||
• UE4M3 block-16 scales: 4 values packed per uint32
|
||||
• Returns ((l1_w, l1_sf), (l2_w, l2_sf)) per rank
|
||||
|
||||
5. SymmBuffer allocation
|
||||
└─ symm_buffer.py: get_symm_buffer_for_nvfp4_mega_moe()
|
||||
• Pre-allocates GPU buffers for:
|
||||
- x: int8 packed E2M1 activations
|
||||
- x_sf: uint32 packed UE4M3 activation scales
|
||||
- topk_idx: int32 expert indices
|
||||
- topk_weights: float32 routing weights
|
||||
- buffer: BF16 all-reduce buffer
|
||||
|
||||
6. Profile run (warmup)
|
||||
└─ First forward pass to allocate KV cache, etc.
|
||||
└─ This is where the CUTLASS GEMM first executes
|
||||
└─ SFB weight scales remapped per-expert inside CUTLASS (no cache)
|
||||
|
||||
7. Ready to serve
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## File Map
|
||||
|
||||
```
|
||||
nvfp4_megamoe_kernel/
|
||||
├── __init__.py # Public API exports
|
||||
├── nvfp4_mega_moe.py # Main kernel: nvfp4_mega_moe_full, L1/L2, stage_activation
|
||||
├── 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 + 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
|
||||
├── test_gemm.py # Standalone test
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### What each file does (in call order)
|
||||
|
||||
| File | When it runs | What it does |
|
||||
|------|-------------|--------------|
|
||||
| `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 → 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) and `unpack_ue4m3_u32`. NO prepack cache — SFB remapped per-call inside CUTLASS. |
|
||||
| `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 (SFB remapped inside C extension), 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 and SFB remap. |
|
||||
| `cutlass_nvfp4_gemm/sf_layout.py` | Reference only | Documents the CUTLASS SfAtom layout. Not used at runtime (remap is in CUDA). |
|
||||
|
||||
---
|
||||
|
||||
## Data Formats
|
||||
|
||||
### Weights
|
||||
- **Packed E2M1** (`int8`): 2 FP4 values per byte. Shape: `(E_per_rank, N, K//2)`, K-major layout.
|
||||
- **UE4M3 block scales** (`float8_e4m3fn`): 1 scale per 16 FP4 values (group_size=16). Shape: `(E_per_rank, N, K//16)`. Returned as `float8_e4m3fn` from `weight_transform.py` — NOT packed uint32. The CUTLASS GEMM consumes float8 directly.
|
||||
|
||||
### Activations (after staging kernel)
|
||||
- **Packed E2M1** (`int8`): Shape: `(num_tokens, K//2)`.
|
||||
- **UE4M3 scales** (`uint32`): 4 UE4M3 values packed per uint32. Shape: `(num_tokens, K//64)`. Unpacked to `float8_e4m3fn` via `unpack_ue4m3_u32` before reaching the CUTLASS GEMM.
|
||||
|
||||
### GEMM dimensions (DeepSeek-V4-Pro)
|
||||
- **L1 (gate_up_proj):** M×6144×7168 (per expert)
|
||||
- **L2 (down_proj):** M×7168×3072 (per expert)
|
||||
- 48 experts per rank (384 total / 8 ranks), top-6 routing
|
||||
|
||||
---
|
||||
|
||||
## CUTLASS Scale Factor Remap
|
||||
|
||||
CUTLASS's `Sm1xxBlockScaledConfig` expects scale factors in a specific interleaved layout, not simple row-major. The SfAtom is:
|
||||
|
||||
```
|
||||
Atom Shape: Shape<Shape<32, 4>, Shape<16, 4>>
|
||||
Atom Stride: Stride<Stride<16, 4>, Stride<0, 1>>
|
||||
Tiling: Step<_2, _1> (M tiled with step 2, K with step 1)
|
||||
```
|
||||
|
||||
Our source data is row-major `(M, K_sf)` where `K_sf = K / 16`. The remap kernel (`remap_sf_to_cutlass_kernel` in `cutlass_nvfp4_gemm.cu`) converts from row-major to CUTLASS's interleaved layout.
|
||||
|
||||
### How the remap works
|
||||
|
||||
The kernel iterates over CUTLASS destination indices, uses `cute::idx2crd` to get the hierarchical coordinate, then `cute::flatten` to get a flat tuple of 8 sub-indices. From those, we extract logical `(m, k_sf)` and read from the row-major source.
|
||||
|
||||
### Flattened coordinate decomposition (flat_rank=8)
|
||||
|
||||
From the SfAtom layout with Step<_2, _1> tiling, `flatten(idx2crd(idx, ...))` produces 8 values:
|
||||
|
||||
```
|
||||
f0 = inner_m (0..31) — varies fastest within M atom
|
||||
f1 = sub_m (0..3) — second M sub-coordinate
|
||||
f2 = tile_m (0..) — M tile index
|
||||
f3 = step_m stride — degenerate (always = sfa_size, not a coordinate)
|
||||
f4 = sub_k (0..3) — K sub-coordinate within atom
|
||||
f5 = tile_k (0..) — K tile index
|
||||
f6 = 0 — unused
|
||||
f7 = 0 — unused
|
||||
```
|
||||
|
||||
#### Empirical coordinate dump (MN=8192, K_sf=448, T = sfa_size = 58720256)
|
||||
|
||||
| idx | f0 | f1 | f2 | f3 | f4 | f5 | f6 | f7 |
|
||||
| ----- | --- | --- | --- | --- | --- | --- | --- | --- |
|
||||
| 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 |
|
||||
| 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 |
|
||||
| 2048 | 0 | 0 | 0 | T | 0 | 4 | 0 | 0 |
|
||||
| 4096 | 0 | 0 | 0 | T | 0 | 8 | 0 | 0 |
|
||||
| 8192 | 0 | 0 | 0 | T | 0 | 16 | 0 | 0 |
|
||||
| 65536 | 0 | 0 | 1 | T | 0 | 16 | 0 | 0 |
|
||||
| 131072 | 0 | 0 | 2 | T | 0 | 32 | 0 | 0 |
|
||||
|
||||
#### Extraction formula
|
||||
|
||||
CuTe uses "first sub varies fastest" for `Shape<32, 4>`:
|
||||
|
||||
```cpp
|
||||
m = f0 + f1 * 32 + f2 * 128;
|
||||
k_sf = f4 + f5 * 4;
|
||||
```
|
||||
|
||||
This was verified with 6 independent probes:
|
||||
|
||||
| Probe | Source | Expected | Result |
|
||||
|-------|--------|----------|--------|
|
||||
| SFA[1, 0] = 2.0 | row 1 changes | ✅ only row 1 | Confirms f0 term |
|
||||
| SFA[32, 0] = 2.0 | row 32 changes | ✅ only row 32 | Confirms f1*32, rules out f0*4+f1 |
|
||||
| SFA[128, 0] = 2.0 | row 128 changes | ✅ only row 128 | Confirms f2*128 |
|
||||
| SFA[0, 1] = 2.0 | row 0 changes (k=1) | ✅ only row 0 | Confirms f4 term |
|
||||
| SFA[0, 4] = 2.0 | row 0 changes (k=4) | ✅ only row 0 | Confirms f5*4 term |
|
||||
| SFA[0, 100] = 2.0 | row 0 changes (k=100) | ✅ only row 0 | Confirms tile-overflow range |
|
||||
|
||||
#### Why the previous remap was broken
|
||||
|
||||
The previous code used `cute::get<0>(flat)` and `cute::get<1>(flat)` to extract (m, k). Since flatten produces `(inner_m, sub_m, tile_m, ...)` in order, `get<0>` and `get<1>` are both **M sub-indices** — they carry no K information. This caused only `k_group=0` to work; all other K-groups were silently mapped to the wrong source offset.
|
||||
|
||||
Additionally, the dest buffer must be zero-initialized before remap because CUTLASS pads to tile boundaries (128 × 64), making the dest buffer larger than `M * K_sf`. Unmapped padding slots reading garbage caused sporadic wrong results.
|
||||
|
||||
---
|
||||
|
||||
## Bugs Found & Fixed
|
||||
|
||||
### 1. unpack_ue4m3_u32: value cast vs bit reinterpret
|
||||
**File:** `nvfp4_mega_moe.py`
|
||||
**Bug:** `(x_u32 & 0xFF).to(torch.int32).to(torch.float8_e4m3fn)` converts integer 63 → float8(63.0).
|
||||
**Fix:** `(x_u32 & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn)` reinterprets bit pattern 0x3F → float8(~0.984).
|
||||
**Also:** `uint32` lacks CUDA bitwise ops — cast to `int32` first.
|
||||
**Impact:** Corrupted every activation scale fed to the L1 GEMM. Weight scales were fine (already float8 from weight_transform). "Structured garbage" recipe.
|
||||
|
||||
### 2. stage_activation: three independent bugs
|
||||
**File:** `nvfp4_moe.py`
|
||||
**Bug A:** `clamp(0, 15)` zeroed every negative value. E2M1 is sign-magnitude 4-bit (bit3=sign, bits2:0=mag).
|
||||
**Bug B:** Stored `block_max` but divided by `block_max/6.0` → stored scale was 6× too large.
|
||||
**Bug C:** Uniform 0.5 step doesn't match E2M1 values {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6} — non-uniform above ±2.
|
||||
**Fix:** Rewrote with proper nearest-neighbor E2M1 quantization.
|
||||
**Impact:** Half the L1→L2 activation was zeroed, 6× scale mismatch, quantization noise on top.
|
||||
|
||||
### 3. _fold_global_scale: logical_widths branch
|
||||
**File:** `weight_transform.py`
|
||||
**Bug:** `logical_widths=[3072, 3072]` caused the function to apply expert 0's scale to gate half and expert 1's scale to up half of ALL experts. All other experts' global scales were discarded.
|
||||
**Fix:** Removed the `logical_widths` branch entirely. The `else` branch correctly broadcasts each expert's own `(E, 1)` global scale across `(E, N, K//16)`.
|
||||
|
||||
### 4. L1 weight interleave removed (transpose still needed)
|
||||
**File:** `weight_transform.py`
|
||||
**Bug:** `_interleave_l1_weights` assumed gate/up were pre-interleaved in groups of 16 and that the kernel used 2CTA UMMA layout. vLLM uses plain concat `[gate; up]` along the output dim, and our CUTLASS kernel uses `ClusterShape<1, 1, 1>`.
|
||||
**Fix:** Removed the interleave function. Weights still need a transpose from checkpoint layout `(N, K_half)` row-major to CUTLASS layout `(K_half, N)` column-major — this is standard row→column conversion, not interleaving. Both L1 and L2 weights and scales are transposed.
|
||||
|
||||
### 5. SF remap: idx2crd+flatten coordinate extraction
|
||||
**File:** `cutlass_nvfp4_gemm.cu`
|
||||
**Bug:** `cute::flatten(coord)` produces 8 sub-indices (flat_rank=8). `get<0>` and `get<1>` are both M sub-indices (inner_m, sub_m), carrying zero K information. Only k_group=0 worked; all other K-groups were silently wrong.
|
||||
**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):
|
||||
- **Cosine = 1.0000, MSE = 0.19** (expected FP4 quantization noise)
|
||||
|
||||
With real (variable) scale factors and the broken remap:
|
||||
- **Cosine = 0.83** → scales are misaligned, not fundamentally broken
|
||||
|
||||
After the fix with correct coordinate extraction:
|
||||
- **Cosine = 1.0000, MSE = 0.0** → perfect match with dequantized reference
|
||||
|
||||
---
|
||||
|
||||
## Build & Deploy (B200)
|
||||
|
||||
```bash
|
||||
# On B200 host — CUTLASS must be cloned and mounted
|
||||
cd /root/nvidia-meeting/deepseek-v4-quant/
|
||||
|
||||
# Rebuild container (CUTLASS is host-mounted at /root/cutlass)
|
||||
KERNEL_CACHE_BUSTER=$(date +%s) docker compose build --no-cache
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
The CUTLASS extension builds inside the container during `pip install` of the nvfp4-megamoe-kernel package. It needs:
|
||||
- CUDA 13.0 toolkit (in the vllm/vllm-openai:nightly image)
|
||||
- CUTLASS headers at `/root/cutlass/include/`
|
||||
- CCCL headers at `/usr/local/cuda-13.0/targets/x86_64-linux/include/cccl/`
|
||||
- Device with SM100 compute capability (B200)
|
||||
|
||||
---
|
||||
|
||||
## 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.
|
||||
|
||||
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**~~ — 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.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `MEGA_MOE_STATIC` | 0 | Set to 1 to skip MoE kernel entirely (return zeros) |
|
||||
| `MEGA_MOE_DEBUG` | 0 | Set to 1 for verbose logging |
|
||||
| `SKIP_ATTENTION` | 0 | Skip attention layers (debug) |
|
||||
|
||||
---
|
||||
|
||||
## Repos
|
||||
|
||||
- **Kernel:** `sweetapi.com/biondizzle/nvfp4-megamoe-kernel` (branch: master)
|
||||
- **Deployment:** `sweetapi.com/biondizzle/deepseek-v4-quant` (branch: modelopt-nvfp4)
|
||||
- **Local:** `~/dev/nvfp4-megamoe-kernel/`, `~/dev/deepseek-v4-quant/`
|
||||
|
||||
22
build_and_run.sh
Executable file
22
build_and_run.sh
Executable file
@@ -0,0 +1,22 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
# Bust the Docker build cache by injecting a timestamp
|
||||
TIMESTAMP=$(date +%s)
|
||||
sed -i -E "s/ARG CACHE_BUSTER=.*/ARG CACHE_BUSTER=${TIMESTAMP}/" Dockerfile
|
||||
|
||||
echo "=== Stopping existing container ==="
|
||||
docker compose down 2>/dev/null || true
|
||||
|
||||
echo "=== Building (no cache) ==="
|
||||
docker compose build --no-cache
|
||||
|
||||
echo "=== Starting ==="
|
||||
docker compose up -d
|
||||
|
||||
# Restore Dockerfile so git diff stays clean
|
||||
sed -i -E "s/ARG CACHE_BUSTER=.*/ARG CACHE_BUSTER=\${TIMESTAMP}/" Dockerfile
|
||||
|
||||
echo "=== Done. Container: $(docker compose ps -q) ==="
|
||||
37
docker-compose.yml
Normal file
37
docker-compose.yml
Normal file
@@ -0,0 +1,37 @@
|
||||
services:
|
||||
vllm:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
ports:
|
||||
- "8000:8000"
|
||||
environment:
|
||||
- OMP_NUM_THREADS=128
|
||||
- CUDA_LAUNCH_BLOCKING=0
|
||||
- TORCH_SHOW_CPP_STACKTRACES=0
|
||||
- MEGA_MOE_DEBUG=0
|
||||
- MEGA_MOE_STATIC=0
|
||||
- NVFP4_DEBUG=0
|
||||
- NVFP4_DEBUG_SYNC=0
|
||||
- SKIP_ATTENTION=0
|
||||
- MEGA_MOE_USE_CUTLASS=1
|
||||
- DG_JIT_DEBUG=0
|
||||
- DEEP_GEMM_JIT_DEBUG=0
|
||||
command:
|
||||
- /model
|
||||
- --trust-remote-code
|
||||
- --enable-expert-parallel
|
||||
- --tensor-parallel-size=8
|
||||
- --enforce-eager
|
||||
- --tokenizer-mode=deepseek_v4
|
||||
- --host=0.0.0.0
|
||||
- --port=8000
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
volumes:
|
||||
- /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4:/model:ro
|
||||
@@ -1,32 +0,0 @@
|
||||
"""
|
||||
NVFP4 weight transformation and SF layout utilities.
|
||||
|
||||
Port of deep_gemm.mega.transform_nvfp4_weights_for_mega_moe
|
||||
"""
|
||||
|
||||
from math import ceil_div
|
||||
|
||||
fn fold_global_scale_into_block_scales(
|
||||
weight_scale: Tensor[float8_e4m3fn], # (N, K//16) UE4M3 block scales
|
||||
weight_scale_2: Tensor[float32], # (num_logical,) or scalar global scale
|
||||
logical_widths: List[int], # per-logical-weight row counts
|
||||
) -> Tensor[float32]:
|
||||
"""Fold global scale into block scales: UE4M3 * FP32 -> FP32"""
|
||||
# Convert UE4M3 to float32, multiply by global scale
|
||||
# For MergedColumnParallelLinear, expand per-logical global scale
|
||||
...
|
||||
|
||||
fn pack_ue4m3_to_int32(sf: Tensor[float8_e4m3fn]) -> Tensor[int32]:
|
||||
"""Pack 4 UE4M3 values (4 bytes) into one int32 for DeepGEMM TMA"""
|
||||
# View as uint8, pack 4 consecutive bytes into int32
|
||||
...
|
||||
|
||||
fn transform_sf_into_required_layout(
|
||||
sf_mn: Tensor[int32], # MN-major packed SF
|
||||
N: int, K: int,
|
||||
recipe: Tuple[int, int], # (gran_mn, gran_k)
|
||||
num_groups: int,
|
||||
) -> Tensor[int32]:
|
||||
"""Transform SF into TMA-aligned UTCCP layout for DeepGEMM"""
|
||||
# Call into DeepGEMM's C++ layout transform
|
||||
...
|
||||
@@ -1,24 +0,0 @@
|
||||
"""
|
||||
NVFP4 Mega MoE Kernel — Mojo Rewrite
|
||||
|
||||
This is a from-scratch rewrite of the DeepGEMM fp8_nvfp4_mega_moe kernel.
|
||||
The CUDA C++ version crashes on B200 with CUDA_ERROR_LAUNCH_FAILED in the
|
||||
SM100_MMA_MXF4NVF4 instruction. This Mojo rewrite aims to:
|
||||
1. Produce a correct, working NVFP4 mega_moe kernel
|
||||
2. Be more maintainable than 1200+ lines of CUDA template metaprogramming
|
||||
3. Leverage Mojo's GPU programming model for cleaner TMA/UMMA integration
|
||||
|
||||
NVFP4 format:
|
||||
- Weights: E2M1 packed (2 x 4-bit values per byte), int8 container
|
||||
- Block scales: UE4M3 (float8_e4m3fn), group_size=16
|
||||
- Global scale: float32 scalar
|
||||
- Activation: FP8 e4m3fn with UE8M0 per-token scales
|
||||
|
||||
Key differences from MXFP4 (group_size=32, UE8M0):
|
||||
- 2 SF K-columns per BLOCK_K (128/16/4=2) instead of 1
|
||||
- mxf4nvf4 UMMA instruction with scale_vec::4X
|
||||
- Different weight transformation (gran_k=16 vs gran_k=32)
|
||||
"""
|
||||
|
||||
# TODO: Implement as Mojo GPU kernel
|
||||
# Waiting for Mojo GPU programming docs / SDK setup
|
||||
@@ -1,256 +0,0 @@
|
||||
"""
|
||||
NVFP4 Block-Scaled GEMM — TileLang
|
||||
2CTA persistent kernel for Blackwell SM100
|
||||
|
||||
Adapted from tilelang/examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.py
|
||||
|
||||
Key NVFP4 differences from MXFP8:
|
||||
- sf_granularity_k=16 (UE4M3 block16) instead of 128 (UE8M0)
|
||||
- 2 SF uint32 words per BLOCK_K instead of 1 (sf_load_period=1, not 4)
|
||||
- sf_a_id/sf_b_id cycle through 0,1,2,3 within each SF word
|
||||
- Must call tcgen05_gemm_blockscaled twice per K-block with different SF words
|
||||
- in_dtype="float4_e2m1fnx2" (packed FP4) instead of "float8_e4m3fn" (FP8)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang.carver.arch import driver
|
||||
from tilelang.profiler import do_bench
|
||||
|
||||
# NVFP4 constants
|
||||
NVFP4_SF_GRAN_K = 16 # UE4M3 group_size
|
||||
NVFP4_SF_PACK = 4 # 4 UE4M3 values per uint32
|
||||
NVFP4_SF_K_PER_WORD = NVFP4_SF_GRAN_K * NVFP4_SF_PACK # 64 K-elements per uint32 word
|
||||
# For BLOCK_K=128: 128/64 = 2 SF words per K-block
|
||||
# Each word covers 4 UMMA sub-atoms (sf_id 0-3)
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def nvfp4_blockscaled_gemm_2cta_persistent(
|
||||
A, # (M, K) packed FP4
|
||||
B, # (N, K) packed FP4 (K-major for NN)
|
||||
SFA, # (sf_k_groups * M) uint32 packed UE4M3
|
||||
SFB, # (sf_k_groups * N) uint32 packed UE4M3
|
||||
block_M=128,
|
||||
block_N=256,
|
||||
block_K=128,
|
||||
in_dtype="float4_e2m1fnx2",
|
||||
out_dtype="bfloat16",
|
||||
accum_dtype="float32",
|
||||
num_stages=2,
|
||||
sf_granularity_k=NVFP4_SF_GRAN_K,
|
||||
use_tma_store=True,
|
||||
store_block_N=64,
|
||||
):
|
||||
M, N, K = T.const("M, N, K")
|
||||
|
||||
assert block_M == 128
|
||||
assert block_N == 256
|
||||
assert block_K == 128
|
||||
|
||||
half_N = block_N // 2
|
||||
k_iters = T.ceildiv(K, block_K)
|
||||
|
||||
# NVFP4 SF layout:
|
||||
# sf_granularity_k=16, pack_factor=4 → 64 K-elements per uint32
|
||||
# BLOCK_K=128 → 2 SF words per K-block
|
||||
# sf_load_period=1 (load every K-block, unlike MXFP8 which loads every 4)
|
||||
sf_words_per_k_block = block_K // (sf_granularity_k * NVFP4_SF_PACK) # 2
|
||||
sf_k_groups = T.ceildiv(K, sf_granularity_k * NVFP4_SF_PACK) # total SF groups across K
|
||||
|
||||
A: T.Tensor[[M, K], in_dtype]
|
||||
B: T.Tensor[[N, K], in_dtype]
|
||||
SFA: T.Tensor[[sf_k_groups * M], T.uint32]
|
||||
SFB: T.Tensor[[sf_k_groups * N], T.uint32]
|
||||
C = T.empty((M, N), out_dtype)
|
||||
|
||||
sm_num = driver.get_num_sms()
|
||||
num_clusters = sm_num // 2
|
||||
m_blocks = T.ceildiv(M, block_M)
|
||||
m_clusters = m_blocks // 2
|
||||
n_blocks = T.ceildiv(N, block_N)
|
||||
waves = T.ceildiv(m_blocks * n_blocks, sm_num)
|
||||
group_size = 16
|
||||
assert n_blocks % (2 * group_size) == 0
|
||||
|
||||
with T.Kernel(sm_num, threads=256, cluster_dims=2) as (block_id):
|
||||
cta_id = T.block_rank_in_cluster()
|
||||
T.assume(cta_id < 2)
|
||||
|
||||
# Shared memory — FP4 packed (K is halved because 2 values per byte)
|
||||
A_shared = T.alloc_shared((num_stages, block_M, block_K // 2), "uint8")
|
||||
B_shared = T.alloc_shared((num_stages, block_K // 2, half_N), "uint8")
|
||||
|
||||
# Shared memory for SF — 2 uint32 words per K-block for NVFP4
|
||||
SFA_shared = T.alloc_shared((num_stages, block_M, sf_words_per_k_block), T.uint32)
|
||||
SFB_shared = T.alloc_shared((num_stages, block_N, sf_words_per_k_block), T.uint32)
|
||||
|
||||
# Tensor memory
|
||||
C_tmem = T.alloc_tmem([block_M, block_N], accum_dtype)
|
||||
SFA_tmem = T.alloc_tmem([block_M, block_M // 128 * 4], T.uint32)
|
||||
SFB_tmem = T.alloc_tmem([block_M, block_N // 128 * 4], T.uint32)
|
||||
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
C_local_cast = T.alloc_fragment((block_M, block_N), out_dtype)
|
||||
C_shared = T.alloc_shared((block_M, store_block_N), out_dtype)
|
||||
|
||||
loaded = T.alloc_barrier([32] * num_stages)
|
||||
with_sf_full = T.alloc_cluster_barrier([32 * 2] * num_stages)
|
||||
consumed = T.alloc_cluster_barrier([1] * num_stages)
|
||||
tmem_full = T.alloc_cluster_barrier([1])
|
||||
tmem_empty = T.alloc_cluster_barrier([128 * 2])
|
||||
|
||||
tx = T.get_thread_binding()
|
||||
warp_idx = tx // 32
|
||||
|
||||
if warp_idx == 0:
|
||||
# Warp 0: TMA load
|
||||
for w in T.serial(waves):
|
||||
cluster_id = block_id // 2
|
||||
tile_id = num_clusters * w + cluster_id
|
||||
bx_cluster = (tile_id // group_size) % m_clusters
|
||||
bx = bx_cluster * 2 + cta_id
|
||||
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size
|
||||
|
||||
if bx * block_M < M and by * block_N < N:
|
||||
for k in T.serial(k_iters):
|
||||
phase = w * k_iters + k
|
||||
stage = phase % num_stages
|
||||
parity = (phase // num_stages) & 1
|
||||
|
||||
T.mbarrier_wait_parity(consumed[stage], parity ^ 1)
|
||||
|
||||
# TMA load A (packed FP4)
|
||||
T.tma_copy(
|
||||
A[bx * block_M, k * block_K],
|
||||
A_shared[stage, :, :],
|
||||
barrier=loaded[stage],
|
||||
)
|
||||
# TMA load B (packed FP4)
|
||||
T.tma_copy(
|
||||
B[k * block_K, by * block_N + cta_id * half_N],
|
||||
B_shared[stage, :, :],
|
||||
barrier=loaded[stage],
|
||||
)
|
||||
# TMA load SFA — 2 words per K-block for NVFP4
|
||||
# Word 0 covers K[k*128 : k*128+64]
|
||||
# Word 1 covers K[k*128+64 : k*128+128]
|
||||
for sf_word in T.unroll(sf_words_per_k_block):
|
||||
sf_group = k * sf_words_per_k_block + sf_word
|
||||
T.tma_copy(
|
||||
SFA[sf_group * M + bx * block_M :
|
||||
sf_group * M + (bx + 1) * block_M],
|
||||
SFA_shared[stage, :, sf_word],
|
||||
barrier=loaded[stage],
|
||||
)
|
||||
# TMA load SFB — 2 words per K-block
|
||||
for sf_word in T.unroll(sf_words_per_k_block):
|
||||
sf_group = k * sf_words_per_k_block + sf_word
|
||||
T.tma_copy(
|
||||
SFB[sf_group * N + by * block_N :
|
||||
sf_group * N + (by + 1) * block_N],
|
||||
SFB_shared[stage, :, sf_word],
|
||||
barrier=loaded[stage],
|
||||
)
|
||||
T.mbarrier_arrive(loaded[stage])
|
||||
|
||||
elif warp_idx == 1 and cta_id == 0:
|
||||
# Warp 1: MMA issue + UTCCP
|
||||
for w in T.serial(waves):
|
||||
cluster_id = block_id // 2
|
||||
tile_id = num_clusters * w + cluster_id
|
||||
bx_cluster = (tile_id // group_size) % m_clusters
|
||||
bx = bx_cluster * 2 + cta_id
|
||||
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size
|
||||
|
||||
if bx * block_M < M and by * block_N < N:
|
||||
for k in T.serial(k_iters):
|
||||
phase = w * k_iters + k
|
||||
stage = phase % num_stages
|
||||
parity = (phase // num_stages) & 1
|
||||
|
||||
T.mbarrier_wait_parity(with_sf_full[stage], parity)
|
||||
|
||||
# NVFP4: 2 SF words per K-block → 2 sub-GEMM calls
|
||||
# Each sub-GEMM covers 64 K-elements (BLOCK_K/2)
|
||||
for sf_word in T.unroll(sf_words_per_k_block):
|
||||
# UTCCP: copy SF word from shared to tensor memory
|
||||
T.tcgen05_cp_warpx4(
|
||||
SFA_shared[stage, :, sf_word],
|
||||
SFA_tmem,
|
||||
use_2cta=True,
|
||||
)
|
||||
T.tcgen05_cp_warpx4(
|
||||
SFB_shared[stage, :, sf_word],
|
||||
SFB_tmem,
|
||||
use_2cta=True,
|
||||
)
|
||||
|
||||
# Block-scaled GEMM
|
||||
# For NVFP4: sf_a_id selects which of 4 packed UE4M3
|
||||
# values in the uint32 word to use.
|
||||
# With sf_granularity_k=16 and UMMA_K=64:
|
||||
# 4 SF values per UMMA atom, sf_id=0 covers all 4
|
||||
# (the hardware auto-cycles through 0,1,2,3 internally)
|
||||
T.tcgen05_gemm_blockscaled(
|
||||
A_shared[stage, :, :],
|
||||
B_shared[stage, :, :],
|
||||
C_tmem,
|
||||
SFA_tmem,
|
||||
SFB_tmem,
|
||||
transpose_B=True,
|
||||
mbar=consumed[stage],
|
||||
clear_accum=(k == 0 and sf_word == 0),
|
||||
sf_a_id=0,
|
||||
sf_b_id=0,
|
||||
use_2cta=True,
|
||||
)
|
||||
|
||||
T.tcgen05_mma_arrive(tmem_full, arrive_2cta=True)
|
||||
|
||||
elif warp_idx == 2:
|
||||
# Warp 2: SF transpose
|
||||
for w in T.serial(waves):
|
||||
cluster_id = block_id // 2
|
||||
tile_id = num_clusters * w + cluster_id
|
||||
bx_cluster = (tile_id // group_size) % m_clusters
|
||||
bx = bx_cluster * 2 + cta_id
|
||||
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size
|
||||
|
||||
if bx * block_M < M and by * block_N < N:
|
||||
for k in T.serial(k_iters):
|
||||
phase = w * k_iters + k
|
||||
stage = phase % num_stages
|
||||
parity = (phase // num_stages) & 1
|
||||
|
||||
T.mbarrier_wait_parity(loaded[stage], parity)
|
||||
# Transpose all SF words for this K-block
|
||||
for sf_word in T.unroll(sf_words_per_k_block):
|
||||
T.tcgen05_sf_warp_transpose(SFA_shared[stage, :, sf_word])
|
||||
T.tcgen05_sf_warp_transpose(SFB_shared[stage, :, sf_word])
|
||||
T.fence_proxy_async()
|
||||
T.mbarrier_arrive(with_sf_full[stage], 0)
|
||||
|
||||
# Epilogue: write C from tmem to global memory
|
||||
for w in T.serial(waves):
|
||||
cluster_id = block_id // 2
|
||||
tile_id = num_clusters * w + cluster_id
|
||||
bx_cluster = (tile_id // group_size) % m_clusters
|
||||
bx = bx_cluster * 2 + cta_id
|
||||
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size
|
||||
|
||||
if bx * block_M < M and by * block_N < N:
|
||||
T.mbarrier_wait_parity(tmem_full, w & 1)
|
||||
T.copy(C_tmem, C_local)
|
||||
T.copy(C_local, C_local_cast)
|
||||
T.copy(C_local_cast, C_shared)
|
||||
|
||||
if use_tma_store:
|
||||
T.copy(C_shared, C[bx * block_M, by * block_N])
|
||||
else:
|
||||
T.copy(C_shared, C[bx * block_M, by * block_N], disable_tma=True)
|
||||
|
||||
T.mbarrier_arrive(tmem_empty, 0)
|
||||
|
||||
return C
|
||||
@@ -1,163 +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
|
||||
- L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales)
|
||||
- NVLink cross-rank sync via symm buffer
|
||||
- Expert parallel: each rank handles NUM_EXPERTS/8 experts
|
||||
|
||||
The kernel is written in TileLang, compiled to SM100 (Blackwell) CUBIN.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang.carver.arch import driver
|
||||
|
||||
# DeepSeek-V4-Pro dimensions
|
||||
HIDDEN = 7168
|
||||
INTERMEDIATE = 3072
|
||||
NUM_EXPERTS = 256
|
||||
NUM_RANKS = 8
|
||||
NUM_TOPK = 6
|
||||
|
||||
# Block sizes for the block-scaled GEMM
|
||||
BLOCK_M = 192 # tokens per tile
|
||||
BLOCK_K = 128
|
||||
BLOCK_N = 128
|
||||
|
||||
# NVFP4 scale parameters
|
||||
SF_GRANULARITY_K = 16 # UE4M3 group_size
|
||||
SF_PACK_FACTOR = 4 # 4 UE4M3 values per uint32
|
||||
SF_WORDS_PER_K_BLOCK = BLOCK_K // (SF_GRANULARITY_K * SF_PACK_FACTOR) # 2
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def nvfp4_mega_moe_l1(
|
||||
# Activation (packed FP4, from staging kernel)
|
||||
X, # (num_tokens, K//2) int8 packed E2M1
|
||||
X_SF, # (num_tokens, sf_k_groups) uint32 packed UE4M3
|
||||
# L1 weights (pre-transformed for mega_moe)
|
||||
L1_W, # (num_experts_per_rank, 2*INTERMEDIATE, K//2) int8 K-major
|
||||
L1_SF, # (num_experts_per_rank, 2*INTERMEDIATE, sf_k_groups) uint32 TMA-aligned
|
||||
# Routing
|
||||
TopkIdx, # (num_tokens, NUM_TOPK) int32
|
||||
TopkW, # (num_tokens, NUM_TOPK) float32
|
||||
# Output
|
||||
Y, # (num_tokens, 2*INTERMEDIATE) bfloat16
|
||||
num_experts_per_rank,
|
||||
):
|
||||
"""
|
||||
L1 GEMM for mega_moe: gate_up_proj
|
||||
X(FP4) @ L1_W(FP4)^T → Y(BF16) with UE4M3 block scaling
|
||||
|
||||
This handles the gate_up_proj for all experts in the rank.
|
||||
Each token is routed to NUM_TOPK experts, and the GEMM is computed
|
||||
per expert group using persistent scheduling.
|
||||
"""
|
||||
num_tokens = T.const("num_tokens")
|
||||
K = T.const("K")
|
||||
INTER = T.const("INTERMEDIATE")
|
||||
|
||||
k_iters = T.ceildiv(K, BLOCK_K)
|
||||
sf_k_groups = T.ceildiv(K, SF_GRANULARITY_K * SF_PACK_FACTOR)
|
||||
|
||||
with T.Kernel(
|
||||
T.ceildiv(num_tokens, BLOCK_M),
|
||||
T.ceildiv(2 * INTER, BLOCK_N),
|
||||
threads=256,
|
||||
cluster_dims=2,
|
||||
) as (bx, by):
|
||||
cta_id = T.block_rank_in_cluster()
|
||||
T.assume(cta_id < 2)
|
||||
|
||||
# ... (persistent scheduling, expert routing, L1 GEMM with block scaling)
|
||||
# This follows the same pattern as nvfp4_blockscaled_gemm_2cta_persistent
|
||||
# but adds expert scheduling and topk weight scaling
|
||||
|
||||
pass # Full implementation in progress
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def nvfp4_mega_moe_l2(
|
||||
# L1 output (quantized to FP4 for L2 input)
|
||||
X, # (num_tokens, INTER//2) int8 packed E2M1
|
||||
X_SF, # (num_tokens, sf_k_groups) uint32 packed UE4M3
|
||||
# L2 weights
|
||||
L2_W, # (num_experts_per_rank, HIDDEN, INTER//2) int8 K-major
|
||||
L2_SF, # (num_experts_per_rank, HIDDEN, sf_k_groups) uint32 TMA-aligned
|
||||
# Routing
|
||||
TopkIdx, # (num_tokens, NUM_TOPK) int32
|
||||
TopkW, # (num_tokens, NUM_TOPK) float32
|
||||
# Output
|
||||
Y, # (num_tokens, HIDDEN) bfloat16
|
||||
num_experts_per_rank,
|
||||
):
|
||||
"""
|
||||
L2 GEMM for mega_moe: down_proj
|
||||
X(FP4) @ L2_W(FP4)^T → Y(BF16) with UE4M3 block scaling
|
||||
|
||||
After SiLU+Mul on the L1 output, the result is quantized to FP4
|
||||
and fed into L2.
|
||||
"""
|
||||
pass # Symmetric to L1
|
||||
|
||||
|
||||
def nvfp4_mega_moe_full(
|
||||
hidden_states, # (num_tokens, HIDDEN) bfloat16
|
||||
topk_weights, # (num_tokens, NUM_TOPK) float32
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32
|
||||
l1_weights, # L1 weights (transformed for mega_moe)
|
||||
l1_scales, # L1 UE4M3 scales (transformed)
|
||||
l2_weights, # L2 weights (transformed for mega_moe)
|
||||
l2_scales, # L2 UE4M3 scales (transformed)
|
||||
symm_buffer, # NVLink symm buffer for cross-rank sync
|
||||
):
|
||||
"""
|
||||
Full mega_moe forward pass:
|
||||
1. Stage: quantize BF16 hidden_states → FP4 + UE4M3 scales
|
||||
2. L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with block scaling)
|
||||
3. SiLU + Mul (activation)
|
||||
4. Quantize L1 output → FP4 + UE4M3 scales
|
||||
5. L2 GEMM: down_proj (FP4 x FP4 → BF16 with block scaling)
|
||||
6. NVLink sync + reduce across ranks
|
||||
"""
|
||||
# Step 1: Stage activation (BF16 → FP4 quantization)
|
||||
# This is the staging kernel from patches/staging_kernel.py
|
||||
x_fp4, x_sf = stage_activation(hidden_states)
|
||||
|
||||
# Step 2: L1 GEMM
|
||||
l1_output = nvfp4_mega_moe_l1(
|
||||
x_fp4, x_sf, l1_weights, l1_scales, topk_ids, topk_weights)
|
||||
|
||||
# Step 3: SiLU + Mul
|
||||
gate, up = l1_output.chunk(2, dim=-1)
|
||||
activated = torch.nn.functional.silu(gate) * up
|
||||
|
||||
# Step 4: Quantize L1 output → FP4
|
||||
l1_fp4, l1_sf = stage_activation(activated)
|
||||
|
||||
# Step 5: L2 GEMM
|
||||
l2_output = nvfp4_mega_moe_l2(
|
||||
l1_fp4, l1_sf, l2_weights, l2_scales, topk_ids, topk_weights)
|
||||
|
||||
# Step 6: NVLink reduce
|
||||
output = nvlink_reduce(l2_output, topk_weights, symm_buffer)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def stage_activation(x_bf16):
|
||||
"""Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales.
|
||||
|
||||
This replaces the Triton staging kernel from patches/staging_kernel.py.
|
||||
"""
|
||||
# E2M1 quantization with UE4M3 block scaling
|
||||
# For now, use PyTorch reference implementation
|
||||
# TODO: Write as TileLang kernel for full pipeline integration
|
||||
from deep_gemm import per_token_cast_to_fp4
|
||||
return per_token_cast_to_fp4(x_bf16)
|
||||
7
src/nvfp4_megamoe_kernel.egg-info/PKG-INFO
Normal file
7
src/nvfp4_megamoe_kernel.egg-info/PKG-INFO
Normal file
@@ -0,0 +1,7 @@
|
||||
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
|
||||
11
src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt
Normal file
11
src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt
Normal file
@@ -0,0 +1,11 @@
|
||||
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
src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt
Normal file
1
src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
2
src/nvfp4_megamoe_kernel.egg-info/requires.txt
Normal file
2
src/nvfp4_megamoe_kernel.egg-info/requires.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
torch>=2.5
|
||||
tilelang>=0.1
|
||||
1
src/nvfp4_megamoe_kernel.egg-info/top_level.txt
Normal file
1
src/nvfp4_megamoe_kernel.egg-info/top_level.txt
Normal file
@@ -0,0 +1 @@
|
||||
nvfp4_megamoe_kernel
|
||||
@@ -1,4 +1,4 @@
|
||||
"""NVFP4 Mega MoE Kernel — TileLang implementation for DeepSeek-V4-Pro on Blackwell."""
|
||||
"""NVFP4 Mega MoE Kernel — CUTLASS implementation for DeepSeek-V4-Pro on Blackwell."""
|
||||
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import (
|
||||
nvfp4_mega_moe_full,
|
||||
@@ -7,10 +7,7 @@ from nvfp4_megamoe_kernel.nvfp4_mega_moe import (
|
||||
stage_activation,
|
||||
)
|
||||
from nvfp4_megamoe_kernel.weight_transform import (
|
||||
fold_global_scale,
|
||||
pack_ue4m3_to_uint32,
|
||||
interleave_l1_weights,
|
||||
transform_nvfp4_weights_for_tilelang as transform_nvfp4_weights_for_mega_moe,
|
||||
transform_nvfp4_weights_for_mega_moe,
|
||||
)
|
||||
from nvfp4_megamoe_kernel.symm_buffer import (
|
||||
SymmBuffer,
|
||||
@@ -20,12 +17,9 @@ from nvfp4_megamoe_kernel.symm_buffer import (
|
||||
__all__ = [
|
||||
"nvfp4_mega_moe_full",
|
||||
"nvfp4_mega_moe_l1",
|
||||
"nvfp4_moe_l2",
|
||||
"nvfp4_mega_moe_l2",
|
||||
"stage_activation",
|
||||
"transform_nvfp4_weights_for_mega_moe",
|
||||
"fold_global_scale",
|
||||
"pack_ue4m3_to_uint32",
|
||||
"interleave_l1_weights",
|
||||
"SymmBuffer",
|
||||
"get_symm_buffer_for_nvfp4_mega_moe",
|
||||
]
|
||||
|
||||
99
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md
Normal file
99
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md
Normal file
@@ -0,0 +1,99 @@
|
||||
# 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
|
||||
6
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py
Normal file
6
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
"""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,
|
||||
)
|
||||
44
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh
Normal file
44
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh
Normal file
@@ -0,0 +1,44 @@
|
||||
#!/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
|
||||
@@ -0,0 +1,312 @@
|
||||
/***************************************************************************************************
|
||||
* 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
|
||||
//
|
||||
// Source-first remap: iterates over the logical (m, k_sf) source grid (which is
|
||||
// a simple 2D row-major tensor), and for each valid element, computes the
|
||||
// physical CUTLASS dest index using the layout's forward mapping.
|
||||
//
|
||||
// This is the inverse of the old idx2crd approach (which iterated dest → source
|
||||
// and relied on fragile coordinate extraction from flattened hierarchical coords).
|
||||
// Source-first is simpler, correct by construction, and avoids the idx2crd bugs.
|
||||
//
|
||||
// 2D kernel launch: blockDim = (32, 8), grid over (K_sf, MN).
|
||||
// SFA: source is (MN, K_sf) row-major → src[m * K_sf + k_sf]
|
||||
// SFB: source is (K_sf, MN) row-major → src[k_sf * MN + m]
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template<typename LayoutSF>
|
||||
__global__ void remap_sf_to_cutlass_kernel(
|
||||
const cutlass::float_ue4m3_t* __restrict__ src,
|
||||
cutlass::float_ue4m3_t* __restrict__ dst, // CUTLASS interleaved layout (zero-initialized)
|
||||
LayoutSF layout_sf, // CuTe layout for dst
|
||||
int MN, int K_sf, // Source dimensions (in SF groups)
|
||||
bool col_major_src = false // true if source is (K_sf, MN) row-major
|
||||
) {
|
||||
int k_sf = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int m = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (m >= MN || k_sf >= K_sf) return;
|
||||
|
||||
// Compute the CUTLASS physical index from logical (m, k_sf) via the layout
|
||||
int dst_idx = layout_sf(cute::make_coord(m, k_sf));
|
||||
|
||||
// Read from source (row-major or col-major) and write to CUTLASS position
|
||||
dst[dst_idx] = col_major_src ? src[k_sf * MN + m] : src[m * K_sf + k_sf];
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// 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::cosize(layout_SFA);
|
||||
int sfb_size = cute::cosize(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);
|
||||
|
||||
dim3 block(32, 8);
|
||||
dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y);
|
||||
dim3 grid_sfb((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y);
|
||||
|
||||
remap_sf_to_cutlass_kernel<<<grid_sfa, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false);
|
||||
remap_sf_to_cutlass_kernel<<<grid_sfb, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, true);
|
||||
|
||||
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::cosize(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::cosize(layout_SFB);
|
||||
int K_sf = K / InputSFVectorSize;
|
||||
|
||||
cudaMemsetAsync(static_cast<ElementSF*>(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream);
|
||||
|
||||
dim3 block(32, 8);
|
||||
dim3 grid((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y);
|
||||
remap_sf_to_cutlass_kernel<<<grid, 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;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// 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::cosize(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);
|
||||
|
||||
dim3 block(32, 8);
|
||||
dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y);
|
||||
remap_sf_to_cutlass_kernel<<<grid_sfa, 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
|
||||
}
|
||||
};
|
||||
|
||||
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
|
||||
146
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
Normal file
146
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
Normal file
@@ -0,0 +1,146 @@
|
||||
"""
|
||||
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 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).
|
||||
|
||||
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}")
|
||||
|
||||
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]
|
||||
|
||||
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}")
|
||||
|
||||
# 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=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}")
|
||||
|
||||
slot_out[e_idx] = expert_out
|
||||
|
||||
return slot_out, slot_token_out
|
||||
131
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp
Normal file
131
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp
Normal file
@@ -0,0 +1,131 @@
|
||||
/** 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");
|
||||
}
|
||||
65
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py
Normal file
65
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py
Normal file
@@ -0,0 +1,65 @@
|
||||
"""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,
|
||||
},
|
||||
)
|
||||
21
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py
Normal file
21
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py
Normal file
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
109
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py
Normal file
109
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py
Normal file
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
Test script for CUTLASS NVFP4 block-scaled GEMM.
|
||||
|
||||
Verifies the kernel against the dequantize-then-BF16 reference implementation.
|
||||
Run on the B200 server after building the extension.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add parent directory to path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from nvfp4_megamoe_kernel.nvfp4_dequant import unpack_e2m1_to_bf16
|
||||
|
||||
|
||||
def test_cutlass_nvfp4_gemm():
|
||||
"""Test single GEMM: C = A @ B^T with native NVFP4 block-scaled MMA."""
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available, skipping test")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
M, N, K = 128, 128, 256 # Small test dimensions
|
||||
|
||||
# Create random E2M1-packed inputs
|
||||
A_packed = torch.randint(-128, 127, (M, K // 2), dtype=torch.int8, device=device)
|
||||
B_packed = torch.randint(-128, 127, (N, K // 2), dtype=torch.int8, device=device)
|
||||
|
||||
# Create random UE4M3 block16 scales
|
||||
A_scales = torch.randn(M, K // 16, dtype=torch.float8_e4m3fn, device=device)
|
||||
B_scales = torch.randn(N, K // 16, dtype=torch.float8_e4m3fn, device=device)
|
||||
# Make positive (UE4M3 is unsigned)
|
||||
A_scales = A_scales.abs().clamp(min=0.0625, max=448.0)
|
||||
B_scales = B_scales.abs().clamp(min=0.0625, max=448.0)
|
||||
|
||||
# Reference: dequantize then BF16 GEMM
|
||||
A_bf16 = unpack_e2m1_to_bf16(A_packed, A_scales)
|
||||
B_bf16 = unpack_e2m1_to_bf16(B_packed, B_scales)
|
||||
C_ref = torch.matmul(A_bf16, B_bf16.t())
|
||||
|
||||
# CUTLASS expects B in column-major: (K_half, N) for weights, (sf_k, N) for scales
|
||||
B_packed_cm = B_packed.t().contiguous()
|
||||
B_scales_cm = B_scales.t().contiguous()
|
||||
|
||||
# CUTLASS native NVFP4 GEMM
|
||||
try:
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
C_cutlass = cutlass_nvfp4_blockscaled_gemm(A_packed, A_scales, B_packed_cm, B_scales_cm)
|
||||
|
||||
# Compare (NVFP4 has low precision, so use loose tolerance)
|
||||
diff = (C_cutlass.float() - C_ref.float()).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
print(f"CUTLASS vs Reference: max_diff={max_diff:.4f}, mean_diff={mean_diff:.4f}")
|
||||
print(f"CUTLASS output shape: {C_cutlass.shape}, dtype: {C_cutlass.dtype}")
|
||||
print(f"Reference output shape: {C_ref.shape}, dtype: {C_ref.dtype}")
|
||||
|
||||
if max_diff < 1.0: # Loose tolerance for 4-bit arithmetic
|
||||
print("✓ CUTLASS NVFP4 GEMM test PASSED")
|
||||
else:
|
||||
print("✗ CUTLASS NVFP4 GEMM test FAILED (max_diff too high)")
|
||||
except Exception as e:
|
||||
print(f"✗ CUTLASS NVFP4 GEMM test FAILED with exception: {e}")
|
||||
print("Falling back to reference implementation (dequantize+BF16)")
|
||||
|
||||
|
||||
def test_grouped_gemm():
|
||||
"""Test grouped expert GEMM for MoE dispatch."""
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available, skipping test")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
num_tokens = 64
|
||||
K = 256
|
||||
N = 128
|
||||
E = 4 # Small number of experts for testing
|
||||
top_k = 2
|
||||
|
||||
# Create inputs
|
||||
x_packed = torch.randint(-128, 127, (num_tokens, K // 2), dtype=torch.int8, device=device)
|
||||
x_scales = torch.randn(num_tokens, K // 16, dtype=torch.float8_e4m3fn, device=device).abs().clamp(min=0.0625, max=448.0)
|
||||
# Weights in column-major for CUTLASS: (E, K_half, N) and (E, sf_k, N)
|
||||
weights = torch.randint(-128, 127, (E, K // 2, N), dtype=torch.int8, device=device)
|
||||
weight_scales = torch.randn(E, K // 16, N, dtype=torch.float8_e4m3fn, device=device).abs().clamp(min=0.0625, max=448.0)
|
||||
topk_ids = torch.randint(0, E, (num_tokens, top_k), dtype=torch.int32, device=device)
|
||||
topk_weights = torch.rand(num_tokens, top_k, dtype=torch.float32, device=device)
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
|
||||
try:
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_grouped_nvfp4_gemm
|
||||
output = cutlass_grouped_nvfp4_gemm(
|
||||
x_packed, x_scales, weights, weight_scales, topk_ids, topk_weights
|
||||
)
|
||||
print(f"Grouped GEMM output shape: {output.shape}, dtype: {output.dtype}")
|
||||
print("✓ Grouped NVFP4 GEMM test PASSED")
|
||||
except Exception as e:
|
||||
print(f"✗ Grouped NVFP4 GEMM test FAILED: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("CUTLASS NVFP4 Block-Scaled GEMM Tests")
|
||||
print("=" * 60)
|
||||
test_cutlass_nvfp4_gemm()
|
||||
print()
|
||||
test_grouped_gemm()
|
||||
@@ -5,17 +5,70 @@ 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
|
||||
- SiLU+Mul activation (per-slot, BEFORE combining expert paths)
|
||||
- L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales)
|
||||
- NVLink cross-rank sync via symm buffer
|
||||
- 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 is written in TileLang, compiled to SM100 (Blackwell) CUBIN.
|
||||
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)
|
||||
# Primary path: uses CUTLASS MainloopSm100TmaUmmaWarpSpecializedBlockScaled
|
||||
# which invokes mxf8f6f4.block_scale tensor core instructions directly.
|
||||
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
|
||||
@@ -32,47 +85,153 @@ 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) uint32 packed UE4M3
|
||||
l1_weights, # (E_per_rank, 2*INTER, K//2) int8 K-major
|
||||
l1_scales, # (E_per_rank, 2*INTER, sf_k_groups) uint32 packed UE4M3
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32
|
||||
topk_weights, # (num_tokens, NUM_TOPK) float32
|
||||
num_experts_per_rank,
|
||||
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
|
||||
alpha=1.0, # fp32 scalar from stage_activation global scale
|
||||
):
|
||||
"""L1 GEMM: gate_up_proj — FP4 x FP4 → BF16 with block scaling.
|
||||
|
||||
TODO: TileLang JIT kernel (nvfp4_blockscaled_gemm_2cta_persistent pattern).
|
||||
"""L1 GEMM: gate_up_proj — slot-based, no routing weights.
|
||||
|
||||
Takes pre-built slot mapping (slot_expert_ids, slot_token) from the outer
|
||||
routing logic. Returns (slot_out, slot_token) where each slot is one
|
||||
(token, topk) pair.
|
||||
"""
|
||||
raise NotImplementedError("nvfp4_mega_moe_l1 TileLang kernel not yet implemented")
|
||||
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}"
|
||||
|
||||
slot_out, slot_token = cutlass_grouped_nvfp4_gemm(
|
||||
x_fp4, x_sf_fp8,
|
||||
l1_weights, w_sf_fp8,
|
||||
slot_expert_ids, # 1D per-slot expert IDs
|
||||
slot_token, # 1D per-slot token indices
|
||||
alpha=alpha,
|
||||
)
|
||||
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_tokens, INTER//2) int8 packed E2M1
|
||||
x_sf, # (num_tokens, sf_k_groups) uint32 packed UE4M3
|
||||
l2_weights, # (E_per_rank, HIDDEN, INTER//2) int8 K-major
|
||||
l2_scales, # (E_per_rank, HIDDEN, sf_k_groups) uint32 packed UE4M3
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32
|
||||
topk_weights, # (num_tokens, NUM_TOPK) float32
|
||||
num_experts_per_rank,
|
||||
x_fp4, # (num_slots, INTER//2) int8 packed E2M1
|
||||
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 (from L1 routing)
|
||||
slot_token, # (num_slots,) int64 — token index per slot (from L1)
|
||||
alpha=1.0, # fp32 scalar from stage_activation global scale
|
||||
):
|
||||
"""L2 GEMM: down_proj — FP4 x FP4 → BF16 with block scaling.
|
||||
|
||||
TODO: TileLang JIT kernel (same pattern as L1).
|
||||
"""L2 GEMM: down_proj — slot-based, no routing weights.
|
||||
|
||||
Reuses the same slot mapping from L1 (same slot_token and slot_expert_ids).
|
||||
"""
|
||||
raise NotImplementedError("nvfp4_mega_moe_l2 TileLang kernel not yet implemented")
|
||||
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}"
|
||||
|
||||
slot_out, _ = cutlass_grouped_nvfp4_gemm(
|
||||
x_fp4, x_sf_fp8,
|
||||
l2_weights, w_sf_fp8,
|
||||
slot_expert_ids, # 1D per-slot expert IDs — GEMM handles directly
|
||||
alpha=alpha,
|
||||
)
|
||||
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):
|
||||
"""Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales.
|
||||
|
||||
This replaces the Triton staging kernel from patches/staging_kernel.py.
|
||||
Two-level quantization matching the NVFP4 weight format:
|
||||
1. Per-tensor global scale: amax / (6.0 * 448.0)
|
||||
2. Per-block (16 values) absmax scaling on the normalized values
|
||||
|
||||
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
|
||||
"""
|
||||
from vllm.model_executor.layers.quantization.utils.fp4_utils import (
|
||||
per_token_cast_to_fp4,
|
||||
)
|
||||
return per_token_cast_to_fp4(x_bf16)
|
||||
x_f32 = x_bf16.float()
|
||||
|
||||
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(
|
||||
@@ -84,71 +243,154 @@ def nvfp4_mega_moe_full(
|
||||
fast_math=False, # fast math flag (unused in NVFP4)
|
||||
):
|
||||
"""Full mega_moe forward pass — replaces deep_gemm.mega.fp8_nvfp4_mega_moe.
|
||||
|
||||
API matches the DeepGEMM fp8_nvfp4_mega_moe call signature used in
|
||||
the vLLM deepseek_v4.py patch:
|
||||
|
||||
fp8_nvfp4_mega_moe(y, l1_weights, l2_weights, symm_buffer,
|
||||
activation_clamp=..., fast_math=...)
|
||||
|
||||
Pipeline:
|
||||
1. Read staged activation from symm_buffer (already quantized by staging kernel)
|
||||
2. L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with block scaling)
|
||||
3. SiLU + Mul (activation)
|
||||
4. Quantize L1 output → FP4 + UE4M3 scales
|
||||
5. L2 GEMM: down_proj (FP4 x FP4 → BF16 with block scaling)
|
||||
6. NVLink sync + reduce across ranks → write to y
|
||||
|
||||
When MEGA_MOE_STATIC=1, returns zeros (bypass) for pipeline testing.
|
||||
|
||||
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) — NO routing weights
|
||||
3. SiLU + Mul PER SLOT (nonlinearity before combining expert paths)
|
||||
4. Quantize activated slots → FP4
|
||||
5. L2 GEMM → slot output (num_slots, HIDDEN) — NO routing weights
|
||||
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
|
||||
l1_w, l1_sf = transformed_l1_weights
|
||||
l2_w, l2_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]}")
|
||||
|
||||
# Step 1: Read staged activation from symm_buffer
|
||||
# The staging has already been done by _stage_deepseek_v4_mega_moe_inputs
|
||||
# and stored in symm_buffer.x, symm_buffer.x_sf
|
||||
x_fp4 = symm_buffer.x[:num_tokens]
|
||||
x_sf = symm_buffer.x_sf[:num_tokens]
|
||||
l1_global_scale = symm_buffer.input_global_scale
|
||||
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] alpha={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.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={topk_ids.shape} l1_w={l1_w.shape} l2_w={l2_w.shape}")
|
||||
|
||||
# Step 2: L1 GEMM
|
||||
num_experts_per_rank = l1_w.shape[0]
|
||||
l1_output = nvfp4_mega_moe_l1(
|
||||
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 (C extension can't handle torch scalars)
|
||||
l1_alpha = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale
|
||||
|
||||
# Shape consistency asserts — catch mismatched slot mappings early
|
||||
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
|
||||
|
||||
# SFB weight scales are remapped per-expert inside CUTLASS on each call.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# NO PREPACK CACHE — see README for rationale.
|
||||
# DO NOT add a prepack cache. Previous attempts caused:
|
||||
# - OOM: ~1.75 GiB per prepacked tensor × 61 layers = 214 GiB
|
||||
# - Peak memory 2× during torch.stack before eviction
|
||||
# - CUDA graph use-after-free on evicted entries
|
||||
# - M_for_layout=128 assumption (unverified M-independence)
|
||||
# The SFB remap is a small scatter kernel (~µs) — not the bottleneck.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Step 2: L1 GEMM — slot-based, no routing weights
|
||||
l1_slots, _ = nvfp4_mega_moe_l1(
|
||||
x_fp4, x_sf, l1_w, l1_sf,
|
||||
topk_ids, topk_weights, num_experts_per_rank,
|
||||
)
|
||||
|
||||
# Step 3: SiLU + Mul
|
||||
gate, up = l1_output.chunk(2, dim=-1)
|
||||
slot_expert_local, slot_token,
|
||||
alpha=l1_alpha,
|
||||
) # (num_slots, 2*INTER) bfloat16
|
||||
|
||||
# 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)
|
||||
activated = torch.nn.functional.silu(gate) * up
|
||||
if activation_clamp is not None:
|
||||
activated = activated.clamp(max=activation_clamp)
|
||||
|
||||
# Step 4: Quantize L1 output → FP4
|
||||
l1_fp4, l1_sf_out = stage_activation(activated)
|
||||
|
||||
# Step 5: L2 GEMM
|
||||
l2_output = nvfp4_mega_moe_l2(
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
print(f"[silu] nan={torch.isnan(activated).any().item()} "
|
||||
f"abs_max={activated.abs().max().item():.4e}")
|
||||
|
||||
# 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] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}]")
|
||||
|
||||
# Step 5: L2 GEMM — slot-based, no routing weights
|
||||
l2_slots = nvfp4_mega_moe_l2(
|
||||
l1_fp4, l1_sf_out, l2_w, l2_sf,
|
||||
topk_ids, topk_weights, num_experts_per_rank,
|
||||
slot_expert_local, slot_token,
|
||||
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),
|
||||
)
|
||||
|
||||
# Step 6: Write to output
|
||||
y.copy_(l2_output)
|
||||
|
||||
@@ -31,21 +31,27 @@ class SymmBuffer:
|
||||
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 values per byte), so K//2
|
||||
sf_k_groups_hidden = hidden_size // (16 * 4) # UE4M3 block16, 4 packed per uint32
|
||||
sf_k_groups_inter = intermediate_size // (16 * 4)
|
||||
|
||||
# Staging buffers (matching DeepGEMM layout)
|
||||
# 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.uint32, device=device,
|
||||
dtype=torch.float8_e4m3fn, device=device,
|
||||
)
|
||||
self.topk_idx = torch.empty(
|
||||
max_num_tokens, top_k,
|
||||
@@ -62,6 +68,10 @@ class SymmBuffer:
|
||||
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} "
|
||||
@@ -83,4 +93,4 @@ def get_symm_buffer_for_nvfp4_mega_moe(
|
||||
return SymmBuffer(
|
||||
group, num_experts, max_num_tokens, top_k,
|
||||
hidden_size, intermediate_size,
|
||||
)
|
||||
)
|
||||
@@ -1,118 +1,162 @@
|
||||
"""
|
||||
NVFP4 Weight Transformation for TileLang mega_moe kernel.
|
||||
NVFP4 Weight Transformation for CUTLASS mega_moe kernel.
|
||||
|
||||
Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale)
|
||||
into the TMA-aligned format expected by the block-scaled GEMM kernel:
|
||||
- Packed FP4 weights (uint8, K-major)
|
||||
- Packed UE4M3 scales (uint32, TMA-aligned UTCCP layout)
|
||||
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)
|
||||
|
||||
This replaces deep_gemm.mega.transform_nvfp4_weights_for_mega_moe.
|
||||
Weight scales are returned as float8_e4m3fn (NOT packed uint32). The CUTLASS GEMM
|
||||
consumes float8 scales directly; only activation scales from the staging kernel come
|
||||
as uint32 and need unpack_ue4m3_u32.
|
||||
|
||||
This replaces deep_gemma.mega.transform_nvfp4_weights_for_mega_moe.
|
||||
|
||||
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
|
||||
import math
|
||||
|
||||
|
||||
def fold_global_scale(
|
||||
weight_scale: torch.Tensor, # (N, K//16) float8_e4m3fn
|
||||
weight_scale_2: torch.Tensor, # (num_logical,) or scalar float32
|
||||
logical_widths: list[int] = None, # per-logical-weight row counts
|
||||
def _fold_global_scale(
|
||||
weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn
|
||||
weight_scale_2: torch.Tensor, # (E,) or (E, 2) or scalar float32
|
||||
) -> torch.Tensor:
|
||||
"""Fold global scale into block scales: UE4M3 * FP32 → UE4M3 → float32.
|
||||
"""Fold global scale into block scales: UE4M3 * FP32 → float32.
|
||||
|
||||
Returns: (N, K//16) float32 folded block scales.
|
||||
For fused projections (w13 = gate+up), weight_scale_2 is (E, 2):
|
||||
scale_2[e, 0] applies to gate_proj rows, scale_2[e, 1] applies to up_proj rows.
|
||||
N is split: gate = weight_scale[:, :N//2, :], up = weight_scale[:, N//2:, :]
|
||||
For single projections (w2), weight_scale_2 is (E,) or scalar.
|
||||
"""
|
||||
sf_f32 = weight_scale.to(torch.float32)
|
||||
|
||||
if weight_scale_2.numel() == 1:
|
||||
sf_f32 = sf_f32 * weight_scale_2.to(torch.float32)
|
||||
elif weight_scale_2.numel() > 1 and logical_widths is not None:
|
||||
# Per-logical-weight global scale (e.g., gate_up_proj has 2)
|
||||
expanded = []
|
||||
for i, w in enumerate(logical_widths):
|
||||
if i < len(weight_scale_2):
|
||||
expanded.append(weight_scale_2[i].flatten()[0].expand(w))
|
||||
global_scale = torch.cat(expanded).to(torch.float32).unsqueeze(1)
|
||||
sf_f32 = sf_f32 * global_scale
|
||||
gs = weight_scale_2.to(torch.float32)
|
||||
|
||||
if gs.numel() == 1:
|
||||
sf_f32 = sf_f32 * gs
|
||||
elif gs.dim() == 2 and gs.shape[1] == 2:
|
||||
# Fused projection: (E, 2) — gate and up have separate global scales
|
||||
# weight_scale is (E, N, K//16), N = gate_N + up_N
|
||||
gate_N = sf_f32.shape[1] // 2
|
||||
gs_gate = gs[:, 0].unsqueeze(-1) # (E, 1)
|
||||
gs_up = gs[:, 1].unsqueeze(-1) # (E, 1)
|
||||
sf_f32[:, :gate_N, :] = sf_f32[:, :gate_N, :] * gs_gate.unsqueeze(-1)
|
||||
sf_f32[:, gate_N:, :] = sf_f32[:, gate_N:, :] * gs_up.unsqueeze(-1)
|
||||
else:
|
||||
sf_f32 = sf_f32 * weight_scale_2.max().to(torch.float32)
|
||||
|
||||
# Per-expert global scale — broadcast multiply
|
||||
while gs.dim() < sf_f32.dim():
|
||||
gs = gs.unsqueeze(-1)
|
||||
sf_f32 = sf_f32 * gs.expand_as(sf_f32)
|
||||
|
||||
return sf_f32
|
||||
|
||||
|
||||
def pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack 4 UE4M3 (float8_e4m3fn) values into one uint32.
|
||||
|
||||
Input: (..., K//16) float8_e4m3fn
|
||||
Output: (..., K//64) uint32 (4 values packed per word)
|
||||
"""
|
||||
# View as raw bytes
|
||||
sf_u8 = sf.view(torch.uint8) # (..., K//16) uint8
|
||||
shape = sf_u8.shape
|
||||
assert shape[-1] % 4 == 0, f"Last dim {shape[-1]} not divisible by 4"
|
||||
|
||||
# Pack 4 consecutive uint8 values into one uint32
|
||||
def _pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack 4 UE4M3 (float8_e4m3fn) values into one uint32."""
|
||||
sf_u8 = sf.view(torch.uint8)
|
||||
assert sf_u8.shape[-1] % 4 == 0, f"Last dim {sf_u8.shape[-1]} not divisible by 4"
|
||||
packed = (sf_u8[..., 0::4].to(torch.int32) |
|
||||
(sf_u8[..., 1::4].to(torch.int32) << 8) |
|
||||
(sf_u8[..., 2::4].to(torch.int32) << 16) |
|
||||
(sf_u8[..., 3::4].to(torch.int32) << 24))
|
||||
|
||||
return packed.contiguous()
|
||||
|
||||
|
||||
def interleave_l1_weights(
|
||||
weight: torch.Tensor, # (E, 2*INTER, K//2) int8, K-major
|
||||
) -> torch.Tensor:
|
||||
"""Interleave L1 (gate_up) weights for 2CTA UMMA.
|
||||
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], tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Transform NVFP4 weights for the CUTLASS block-scaled GEMM.
|
||||
|
||||
The gate and up projections are interleaved in groups of 8 rows
|
||||
to match the UMMA 2CTA schedule.
|
||||
"""
|
||||
E, N, K_half = weight.shape
|
||||
assert N % 16 == 0, f"N={N} not divisible by 16"
|
||||
Matches the call signature from nightly vLLM deepseek_v4.py finalize_weights.
|
||||
|
||||
# Reshape to (E, N//16, 16, K_half) → interleave pairs of 8
|
||||
w = weight.view(E, N // 16, 16, K_half)
|
||||
w_gate = w[:, :, :8, :]
|
||||
w_up = w[:, :, 8:16, :]
|
||||
interleaved = torch.stack([w_gate, w_up], dim=3).reshape(E, N, K_half)
|
||||
|
||||
return interleaved.contiguous()
|
||||
|
||||
|
||||
def transform_nvfp4_weights_for_tilelang(
|
||||
weight: torch.Tensor, # (E, N, K//2) int8, K-major packed FP4
|
||||
weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn UE4M3
|
||||
weight_scale_2: torch.Tensor, # (E, 1) or (E, 2) float32 global scale
|
||||
N: int, # output dimension (2*INTER for L1, HIDDEN for L2)
|
||||
K: int, # input dimension (HIDDEN for L1, INTER for L2)
|
||||
gran_k: int = 16, # NVFP4 group_size
|
||||
is_l1: bool = False, # True for gate_up_proj (needs interleaving)
|
||||
logical_widths: list[int] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Transform NVFP4 weights for the TileLang block-scaled GEMM.
|
||||
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)
|
||||
l2_weight_scale_2: global scale for L2 (float32)
|
||||
|
||||
Returns:
|
||||
transformed_weight: (E, N, K//2) int8, K-major, contiguously interleaved
|
||||
transformed_sf: (E, N, K//64) uint32, packed UE4M3 with folded global scale
|
||||
((l1_weight, l1_sf_packed), (l2_weight, l2_sf_packed))
|
||||
"""
|
||||
# Step 1: Fold global scale into block scales
|
||||
sf_folded = fold_global_scale(weight_scale, weight_scale_2, logical_widths)
|
||||
|
||||
# Step 2: Clamp and convert back to UE4M3
|
||||
sf_clamped = sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
||||
|
||||
# Step 3: Pack UE4M3 into uint32
|
||||
sf_packed = pack_ue4m3_to_uint32(sf_clamped) # (E, N, K//64)
|
||||
|
||||
# Step 4: Interleave L1 weights if needed
|
||||
if is_l1:
|
||||
transformed_weight = interleave_l1_weights(weight)
|
||||
l1_weight, l1_weight_scale = l1_tuple
|
||||
l2_weight, l2_weight_scale = l2_tuple
|
||||
|
||||
# DEBUG: check raw scales before folding
|
||||
l1_sf_f32_raw = l1_weight_scale.to(torch.float32)
|
||||
l1_gs_raw = l1_weight_scale_2.to(torch.float32) if l1_weight_scale_2 is not None else None
|
||||
if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug', False):
|
||||
transform_nvfp4_weights_for_mega_moe._sf_debug = True
|
||||
print(f"[SF-DEBUG] raw l1_sf dtype={l1_weight_scale.dtype} range=[{l1_sf_f32_raw.min().item():.4e}, {l1_sf_f32_raw.max().item():.4e}] "
|
||||
f"unique_raw={torch.unique(l1_weight_scale.view(torch.uint8)).numel()}")
|
||||
if l1_gs_raw is not None:
|
||||
print(f"[SF-DEBUG] l1_gs dtype={l1_weight_scale_2.dtype} shape={tuple(l1_weight_scale_2.shape)} "
|
||||
f"range=[{l1_gs_raw.min().item():.4e}, {l1_gs_raw.max().item():.4e}] "
|
||||
f"unique_gs={torch.unique(l1_gs_raw).numel()}")
|
||||
if l1_gs_raw.dim() == 2 and l1_gs_raw.shape[1] == 2:
|
||||
print(f"[SF-DEBUG] gate gs unique={torch.unique(l1_gs_raw[:, 0]).numel()} "
|
||||
f"up gs unique={torch.unique(l1_gs_raw[:, 1]).numel()}")
|
||||
|
||||
# DEBUG: check L2 scales
|
||||
l2_sf_f32_raw = l2_weight_scale.to(torch.float32)
|
||||
l2_gs_raw = l2_weight_scale_2.to(torch.float32) if l2_weight_scale_2 is not None else None
|
||||
if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug_l2', False):
|
||||
transform_nvfp4_weights_for_mega_moe._sf_debug_l2 = True
|
||||
print(f"[SF-DEBUG-L2] raw l2_sf dtype={l2_weight_scale.dtype} range=[{l2_sf_f32_raw.min().item():.4e}, {l2_sf_f32_raw.max().item():.4e}] "
|
||||
f"unique_raw={torch.unique(l2_weight_scale.view(torch.uint8)).numel()}")
|
||||
if l2_gs_raw is not None:
|
||||
print(f"[SF-DEBUG-L2] l2_gs dtype={l2_weight_scale_2.dtype} shape={tuple(l2_weight_scale_2.shape)} "
|
||||
f"range=[{l2_gs_raw.min().item():.4e}, {l2_gs_raw.max().item():.4e}] "
|
||||
f"unique_gs={torch.unique(l2_gs_raw).numel()}")
|
||||
|
||||
# Post-fold diagnostics — one-time
|
||||
if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug_fold', False):
|
||||
transform_nvfp4_weights_for_mega_moe._sf_debug_fold = True
|
||||
l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2) if l1_weight_scale_2 is not None else l1_weight_scale.to(torch.float32)
|
||||
l1_sf_out_check = l1_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
||||
l2_sf_folded = _fold_global_scale(l2_weight_scale, l2_weight_scale_2) if l2_weight_scale_2 is not None else l2_weight_scale.to(torch.float32)
|
||||
l2_sf_out_check = l2_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
||||
print(f"[SF-FOLD] l1 pre-fold unique_u8={torch.unique(l1_weight_scale.view(torch.uint8)).numel()} "
|
||||
f"post-fold unique_u8={torch.unique(l1_sf_out_check.view(torch.uint8)).numel()} "
|
||||
f"range=[{l1_sf_folded.min().item():.4e}, {l1_sf_folded.max().item():.4e}]")
|
||||
print(f"[SF-FOLD] l2 pre-fold unique_u8={torch.unique(l2_weight_scale.view(torch.uint8)).numel()} "
|
||||
f"post-fold unique_u8={torch.unique(l2_sf_out_check.view(torch.uint8)).numel()} "
|
||||
f"range=[{l2_sf_folded.min().item():.4e}, {l2_sf_folded.max().item():.4e}]")
|
||||
|
||||
# Fold global scales into block scales
|
||||
# The logical_widths branch was wrong: it treated gs as per-projection
|
||||
# scalars and only used experts 0 and 1's scales for ALL experts.
|
||||
# The else branch correctly broadcasts each expert's own global scale.
|
||||
if l1_weight_scale_2 is not None:
|
||||
l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2)
|
||||
else:
|
||||
transformed_weight = weight.contiguous()
|
||||
|
||||
return transformed_weight, sf_packed
|
||||
l1_sf_folded = l1_weight_scale.to(torch.float32)
|
||||
|
||||
if l2_weight_scale_2 is not None:
|
||||
l2_sf_folded = _fold_global_scale(l2_weight_scale, l2_weight_scale_2)
|
||||
else:
|
||||
l2_sf_folded = l2_weight_scale.to(torch.float32)
|
||||
|
||||
# Alias for drop-in replacement of deep_gemm.mega.transform_nvfp4_weights_for_mega_moe
|
||||
transform_nvfp4_weights_for_mega_moe = transform_nvfp4_weights_for_tilelang
|
||||
# Clamp and convert back to UE4M3
|
||||
l1_sf_out = l1_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn).contiguous()
|
||||
l2_sf_out = l2_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn).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)
|
||||
# which is column-major (N, K_half). This is a one-time cost at load time.
|
||||
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), (l2_weight_out, l2_sf_out)
|
||||
|
||||
@@ -1,95 +0,0 @@
|
||||
"""
|
||||
NVFP4 quantization utilities — E2M1 packing and UE4M3 scale handling.
|
||||
Core math layer for the NVFP4 mega_moe kernel rewrite.
|
||||
"""
|
||||
|
||||
# E2M1 magnitude lookup table (positive values only)
|
||||
# Index 0-7 maps to: 0, 0.5, 1, 1.5, 2, 3, 4, 6
|
||||
def e2m1_magnitude(index: Int) -> Float64:
|
||||
if index == 0: return 0.0
|
||||
if index == 1: return 0.5
|
||||
if index == 2: return 1.0
|
||||
if index == 3: return 1.5
|
||||
if index == 4: return 2.0
|
||||
if index == 5: return 3.0
|
||||
if index == 6: return 4.0
|
||||
if index == 7: return 6.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def quantize_e2m1(value: Float64) -> UInt8:
|
||||
"""Quantize a float64 value to E2M1 (4-bit), returning the 4-bit nibble with sign."""
|
||||
var sign = 0
|
||||
var abs_val = value
|
||||
if value < 0.0:
|
||||
sign = 1
|
||||
abs_val = -value
|
||||
|
||||
# Find best E2M1 match
|
||||
var best_idx = 0
|
||||
var best_err = abs_val # error for idx=0
|
||||
|
||||
for i in range(1, 8):
|
||||
mag = e2m1_magnitude(i)
|
||||
err = abs(abs_val - mag)
|
||||
if err < best_err:
|
||||
best_err = err
|
||||
best_idx = i
|
||||
|
||||
return (sign << 3) | best_idx
|
||||
|
||||
|
||||
def unpack_e2m1(packed: UInt8, idx: Int) -> Float64:
|
||||
"""Unpack one E2M1 value from a packed byte.
|
||||
idx=0 -> low nibble, idx=1 -> high nibble.
|
||||
"""
|
||||
nibble: UInt8
|
||||
if idx == 0:
|
||||
nibble = packed & 0x0F
|
||||
else:
|
||||
nibble = (packed >> 4) & 0x0F # keep sign bit
|
||||
|
||||
sign = (nibble >> 3) & 1
|
||||
mag_idx = nibble & 0x07
|
||||
magnitude = e2m1_magnitude(Int(mag_idx))
|
||||
|
||||
if sign:
|
||||
return -magnitude
|
||||
return magnitude
|
||||
|
||||
|
||||
def dequantize_nvfp4_weight(
|
||||
packed_weight: UInt8,
|
||||
block_scale: Float64,
|
||||
group_offset: Int,
|
||||
) -> Float64:
|
||||
"""Dequantize a single NVFP4 weight element.
|
||||
weight = E2M1_magnitude * block_scale
|
||||
(global_scale is already folded into block_scale)
|
||||
"""
|
||||
e2m1_value = unpack_e2m1(packed_weight, group_offset)
|
||||
return e2m1_value * block_scale
|
||||
|
||||
|
||||
def main() raises:
|
||||
# Test E2M1 quantization round-trip
|
||||
print("E2M1 quantization test:")
|
||||
for val in [-6.0, -4.0, -3.0, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]:
|
||||
packed = quantize_e2m1(val)
|
||||
unpacked = unpack_e2m1(packed, 0)
|
||||
print(" ", val, " -> E2M1 -> ", unpacked)
|
||||
|
||||
# Test packed byte (two E2M1 values)
|
||||
print("\nPacked byte test:")
|
||||
lo = 1.5
|
||||
hi = -3.0
|
||||
packed = (quantize_e2m1(hi) << 4) | quantize_e2m1(lo)
|
||||
print(" lo=", lo, " hi=", hi, " packed=", packed)
|
||||
print(" unpack lo=", unpack_e2m1(packed, 0), " unpack hi=", unpack_e2m1(packed, 1))
|
||||
|
||||
# Test NVFP4 dequantization
|
||||
print("\nNVFP4 dequantization test:")
|
||||
packed_w = UInt8(0x36) # low=6.0, high=3.0
|
||||
scale = 0.5
|
||||
print(" packed=0x36, scale=0.5, lo=", dequantize_nvfp4_weight(packed_w, scale, 0))
|
||||
print(" packed=0x36, scale=0.5, hi=", dequantize_nvfp4_weight(packed_w, scale, 1))
|
||||
@@ -1,65 +0,0 @@
|
||||
"""
|
||||
NVFP4 Activation Quantization — BF16 → packed FP4 + UE4M3 block16 scales
|
||||
|
||||
Port of patches/staging_kernel.py to TileLang.
|
||||
|
||||
This quantizes BF16 hidden states to:
|
||||
- Packed FP4 (E2M1, 2 values per uint8 byte)
|
||||
- UE4M3 block16 scales (4 values per uint32 word, group_size=16)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def fp4_quant_kernel(
|
||||
X, # (M, K) bfloat16 input
|
||||
block_M=128,
|
||||
block_K=128,
|
||||
group_size=16, # NVFP4 UE4M3 group_size
|
||||
):
|
||||
"""Quantize BF16 → packed FP4 (E2M1) + UE4M3 block16 scales.
|
||||
|
||||
Output:
|
||||
- X_FP4: (M, K//2) uint8 packed E2M1
|
||||
- X_SF: (M, K//group_size//4) uint32 packed UE4M3 scales
|
||||
"""
|
||||
M, K = T.const("M, K")
|
||||
X: T.Tensor[[M, K], "bfloat16"]
|
||||
|
||||
X_FP4 = T.empty((M, K // 2), "uint8")
|
||||
X_SF = T.empty((M, K // (group_size * 4)), "uint32")
|
||||
|
||||
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(K, block_K), threads=128) as (bx, by):
|
||||
# Each block quantizes a (block_M, block_K) tile
|
||||
tx = T.get_thread_binding()
|
||||
|
||||
row = bx * block_M + tx // (block_K // 2)
|
||||
col = tx % (block_K // 2)
|
||||
|
||||
if row < M and col < K // 2:
|
||||
# Load 2 BF16 values
|
||||
val_lo = X[row, col * 2]
|
||||
val_hi = X[row, col * 2 + 1]
|
||||
|
||||
# Quantize to E2M1 (4-bit each, pack 2 per byte)
|
||||
# ... (E2M1 quantization logic)
|
||||
packed = T.uint8(0) # placeholder
|
||||
X_FP4[row, col] = packed
|
||||
|
||||
# Compute block scale for group of 16 elements
|
||||
# ... (UE4M3 scale computation)
|
||||
|
||||
return X_FP4, X_SF
|
||||
|
||||
|
||||
def fp4_quant_reference(x_bf16, group_size=16):
|
||||
"""Reference FP4 quantization using PyTorch + DeepGEMM.
|
||||
|
||||
Used for correctness verification of the TileLang kernel.
|
||||
"""
|
||||
from deep_gemm import per_token_cast_to_fp4
|
||||
x_fp4, x_sf = per_token_cast_to_fp4(x_bf16)
|
||||
return x_fp4, x_sf
|
||||
@@ -1,147 +0,0 @@
|
||||
"""
|
||||
NVFP4 Weight Transformation for TileLang mega_moe kernel.
|
||||
|
||||
Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale)
|
||||
into the TMA-aligned format expected by the block-scaled GEMM kernel:
|
||||
- Packed FP4 weights (uint8, K-major)
|
||||
- Packed UE4M3 scales (uint32, TMA-aligned UTCCP layout)
|
||||
|
||||
This replaces deep_gemm.mega.transform_nvfp4_weights_for_mega_moe.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
|
||||
|
||||
def fold_global_scale(
|
||||
weight_scale: torch.Tensor, # (N, K//16) float8_e4m3fn
|
||||
weight_scale_2: torch.Tensor, # (num_logical,) or scalar float32
|
||||
logical_widths: list[int] = None, # per-logical-weight row counts
|
||||
) -> torch.Tensor:
|
||||
"""Fold global scale into block scales: UE4M3 * FP32 → UE4M3 → float32.
|
||||
|
||||
Returns: (N, K//16) float32 folded block scales.
|
||||
"""
|
||||
sf_f32 = weight_scale.to(torch.float32)
|
||||
|
||||
if weight_scale_2.numel() == 1:
|
||||
sf_f32 = sf_f32 * weight_scale_2.to(torch.float32)
|
||||
elif weight_scale_2.numel() > 1 and logical_widths is not None:
|
||||
# Per-logical-weight global scale (e.g., gate_up_proj has 2)
|
||||
expanded = []
|
||||
for i, w in enumerate(logical_widths):
|
||||
if i < len(weight_scale_2):
|
||||
expanded.append(weight_scale_2[i].flatten()[0].expand(w))
|
||||
global_scale = torch.cat(expanded).to(torch.float32).unsqueeze(1)
|
||||
sf_f32 = sf_f32 * global_scale
|
||||
else:
|
||||
sf_f32 = sf_f32 * weight_scale_2.max().to(torch.float32)
|
||||
|
||||
return sf_f32
|
||||
|
||||
|
||||
def pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack 4 UE4M3 (float8_e4m3fn) values into one uint32.
|
||||
|
||||
Input: (..., K//16) float8_e4m3fn
|
||||
Output: (..., K//64) uint32 (4 values packed per word)
|
||||
"""
|
||||
# View as raw bytes
|
||||
sf_u8 = sf.view(torch.uint8) # (..., K//16) uint8
|
||||
shape = sf_u8.shape
|
||||
assert shape[-1] % 4 == 0, f"Last dim {shape[-1]} not divisible by 4"
|
||||
|
||||
# Pack 4 consecutive uint8 values into one uint32
|
||||
packed = (sf_u8[..., 0::4].to(torch.int32) |
|
||||
(sf_u8[..., 1::4].to(torch.int32) << 8) |
|
||||
(sf_u8[..., 2::4].to(torch.int32) << 16) |
|
||||
(sf_u8[..., 3::4].to(torch.int32) << 24))
|
||||
|
||||
return packed.contiguous()
|
||||
|
||||
|
||||
def interleave_l1_weights(
|
||||
weight: torch.Tensor, # (E, 2*INTER, K//2) int8, K-major
|
||||
) -> torch.Tensor:
|
||||
"""Interleave L1 (gate_up) weights for 2CTA UMMA.
|
||||
|
||||
The gate and up projections are interleaved in groups of 8 rows
|
||||
to match the UMMA 2CTA schedule.
|
||||
"""
|
||||
E, N, K_half = weight.shape
|
||||
assert N % 16 == 0, f"N={N} not divisible by 16"
|
||||
|
||||
# Reshape to (E, N//16, 16, K_half) → interleave pairs of 8
|
||||
w = weight.view(E, N // 16, 16, K_half)
|
||||
# Interleave: [gate_0..7, up_0..7, gate_8..15, up_8..15, ...]
|
||||
# Actually: pairs of 8 rows from gate and up
|
||||
w_gate = w[:, :, :8, :]
|
||||
w_up = w[:, :, 8:16, :]
|
||||
interleaved = torch.stack([w_gate, w_up], dim=3).reshape(E, N, K_half)
|
||||
|
||||
return interleaved.contiguous()
|
||||
|
||||
|
||||
def transform_nvfp4_weights_for_tilelang(
|
||||
weight: torch.Tensor, # (E, N, K//2) int8, K-major packed FP4
|
||||
weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn UE4M3
|
||||
weight_scale_2: torch.Tensor, # (E, 1) or (E, 2) float32 global scale
|
||||
N: int, # output dimension (2*INTER for L1, HIDDEN for L2)
|
||||
K: int, # input dimension (HIDDEN for L1, INTER for L2)
|
||||
gran_k: int = 16, # NVFP4 group_size
|
||||
is_l1: bool = False, # True for gate_up_proj (needs interleaving)
|
||||
logical_widths: list[int] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Transform NVFP4 weights for the TileLang block-scaled GEMM.
|
||||
|
||||
Returns:
|
||||
transformed_weight: (E, N, K//2) int8, K-major, contiguously interleaved
|
||||
transformed_sf: (E, N, K//64) uint32, packed UE4M3 with folded global scale
|
||||
"""
|
||||
device = weight.device
|
||||
|
||||
# Step 1: Fold global scale into block scales
|
||||
sf_folded = fold_global_scale(weight_scale, weight_scale_2, logical_widths)
|
||||
|
||||
# Step 2: Clamp and convert back to UE4M3
|
||||
sf_clamped = sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
||||
|
||||
# Step 3: Pack UE4M3 into uint32
|
||||
sf_packed = pack_ue4m3_to_uint32(sf_clamped) # (E, N, K//64)
|
||||
|
||||
# Step 4: Interleave L1 weights if needed
|
||||
if is_l1:
|
||||
transformed_weight = interleave_l1_weights(weight)
|
||||
else:
|
||||
transformed_weight = weight.contiguous()
|
||||
|
||||
# Step 5: Ensure K-major layout (already K-major from checkpoint)
|
||||
# The weight should be (E, N, K//2) with stride (N*K//2, K//2, 1) for K-major
|
||||
|
||||
return transformed_weight, sf_packed
|
||||
|
||||
|
||||
def main():
|
||||
"""Test the weight transformation with dummy data."""
|
||||
E = 4 # small test
|
||||
N = 64
|
||||
K = 128
|
||||
|
||||
device = "cpu" # No GPU on sandbox
|
||||
weight = torch.randint(-128, 127, (E, N, K // 2), dtype=torch.int8, device=device)
|
||||
weight_scale = torch.randn(E, N, K // 16, device=device).abs().to(torch.float8_e4m3fn)
|
||||
weight_scale_2 = torch.ones(E, 2, dtype=torch.float32, device=device)
|
||||
|
||||
transformed_w, transformed_sf = transform_nvfp4_weights_for_tilelang(
|
||||
weight, weight_scale, weight_scale_2, N, K, is_l1=True,
|
||||
logical_widths=[32, 32],
|
||||
)
|
||||
|
||||
print(f"Weight: {transformed_w.shape} {transformed_w.dtype}")
|
||||
print(f"SF: {transformed_sf.shape} {transformed_sf.dtype}")
|
||||
print(f"Expected SF: (E={E}, N={N}, K//64={K//64})")
|
||||
print("Weight transformation test passed!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
2304
vllm/patches/deepseek_v4.py
Normal file
2304
vllm/patches/deepseek_v4.py
Normal file
File diff suppressed because it is too large
Load Diff
1155
vllm/patches/deepseek_v4_attention.py
Normal file
1155
vllm/patches/deepseek_v4_attention.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user