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

Initial: TileLang NVFP4 mega_moe kernel package

Browse Source 
- nvfp4_mega_moe_full: drop-in replacement for deep_gemm.mega.fp8_nvfp4_mega_moe
- transform_nvfp4_weights_for_mega_moe: weight transformation (tested)
- SymmBuffer + get_symm_buffer_for_nvfp4_mega_moe: API-matching stubs
- MEGA_MOE_STATIC=1 support for pipeline testing
- pyproject.toml for pip install
This commit is contained in:
biondizzle
2026-05-13 15:44:51 +00:00
commit c2b752c2fe
13 changed files with 1227 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
README.md Normal file
View File

@@ -0,0 +1,40 @@
# NVFP4 Mega MoE Kernel — Mojo Rewrite
Rewrite of the DeepGEMM `fp8_nvfp4_mega_moe` kernel in Mojo.
## 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
## Architecture
The kernel performs NVFP4 (E2M1 + UE4M3 block16 scales) matrix multiply
for MoE (Mixture of Experts) with expert parallelism across NVLink.
### 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
### 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
```

16
pyproject.toml Normal file
View File

@@ -0,0 +1,16 @@
[build-system]
requires = ["setuptools>=68.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "nvfp4-megamoe-kernel"
version = "0.1.0"
description = "NVFP4 Mega MoE kernel for DeepSeek-V4-Pro on Blackwell (TileLang)"
requires-python = ">=3.10"
dependencies = [
"torch>=2.5",
"tilelang>=0.1",
]
[tool.setuptools.packages.find]
where = ["src"]

32
src/layout.mojo Normal file
View File

@@ -0,0 +1,32 @@
"""
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
...

24
src/mega_moe.mojo Normal file
View File

@@ -0,0 +1,24 @@
"""
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

256
src/nvfp4_blockscaled_gemm.py Normal file
View File

@@ -0,0 +1,256 @@
"""
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

163
src/nvfp4_mega_moe.py Normal file
View File

@@ -0,0 +1,163 @@
"""
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)

31
src/nvfp4_megamoe_kernel/__init__.py Normal file
View File

@@ -0,0 +1,31 @@
"""NVFP4 Mega MoE Kernel — TileLang implementation for DeepSeek-V4-Pro on Blackwell."""
from nvfp4_megamoe_kernel.nvfp4_mega_moe import (
nvfp4_mega_moe_full,
nvfp4_mega_moe_l1,
nvfp4_mega_moe_l2,
stage_activation,
)
from nvfp4_megamoe_kernel.weight_transform import (
fold_global_scale,
pack_ue4m3_to_uint32,
interleave_l1_weights,
transform_nvfp4_weights_for_tilelang as transform_nvfp4_weights_for_mega_moe,
)
from nvfp4_megamoe_kernel.symm_buffer import (
SymmBuffer,
get_symm_buffer_for_nvfp4_mega_moe,
)
__all__ = [
"nvfp4_mega_moe_full",
"nvfp4_mega_moe_l1",
"nvfp4_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",
]

154
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py Normal file
View File

@@ -0,0 +1,154 @@
"""
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 os
import torch
# DeepSeek-V4-Pro dimensions
HIDDEN = 7168
INTERMEDIATE = 3072
NUM_EXPERTS = 256
NUM_RANKS = 8
NUM_TOPK = 6
# NVFP4 scale parameters
SF_GRANULARITY_K = 16 # UE4M3 group_size
SF_PACK_FACTOR = 4 # 4 UE4M3 values per uint32
# Runtime flags
MEGA_MOE_STATIC = int(os.environ.get("MEGA_MOE_STATIC", "0"))
MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0"))
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,
):
"""L1 GEMM: gate_up_proj — FP4 x FP4 → BF16 with block scaling.
TODO: TileLang JIT kernel (nvfp4_blockscaled_gemm_2cta_persistent pattern).
"""
raise NotImplementedError("nvfp4_mega_moe_l1 TileLang kernel not yet implemented")
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,
):
"""L2 GEMM: down_proj — FP4 x FP4 → BF16 with block scaling.
TODO: TileLang JIT kernel (same pattern as L1).
"""
raise NotImplementedError("nvfp4_mega_moe_l2 TileLang kernel not yet implemented")
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.
"""
from vllm.model_executor.layers.quantization.utils.fp4_utils import (
per_token_cast_to_fp4,
)
return per_token_cast_to_fp4(x_bf16)
def nvfp4_mega_moe_full(
y, # output tensor (num_tokens, HIDDEN) bfloat16
transformed_l1_weights, # (l1_w, l1_sf) tuple from finalize_weights
transformed_l2_weights, # (l2_w, l2_sf) tuple from finalize_weights
symm_buffer, # SymmBuffer from get_symm_buffer
activation_clamp=None, # optional clamp value (unused in NVFP4)
fast_math=False, # fast math flag (unused in NVFP4)
):
"""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.
"""
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
# 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]
topk_ids = symm_buffer.topk_idx[:num_tokens]
topk_weights = symm_buffer.topk_weights[:num_tokens]
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(
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)
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(
l1_fp4, l1_sf_out, l2_w, l2_sf,
topk_ids, topk_weights, num_experts_per_rank,
)
# Step 6: Write to output
y.copy_(l2_output)

86
src/nvfp4_megamoe_kernel/symm_buffer.py Normal file
View File

@@ -0,0 +1,86 @@
"""Symmetric buffer for NVLink cross-rank all-reduce in mega_moe.
Replaces deep_gemm.mega.SymmBuffer and get_symm_buffer_for_nvfp4_mega_moe.
API matches the DeepGEMM signature used in the vLLM deepseek_v4.py patch.
"""
import os
import torch
import torch.distributed as dist
MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0"))
class SymmBuffer:
"""Symmetric NVLink buffer for expert-parallel cross-rank communication.
Matches the DeepGEMM SymmBuffer interface expected by the vLLM patch:
- .x: staged activation (FP4 packed)
- .x_sf: staged activation scales (UE4M3 packed)
- .topk_idx: top-k expert indices
- .topk_weights: top-k routing weights
- .buffer: underlying CUDA buffer
- .group: process group
"""
def __init__(self, group, num_experts, max_num_tokens, top_k,
hidden_size, intermediate_size):
self.group = group
self.num_experts = num_experts
self.max_num_tokens = max_num_tokens
self.top_k = top_k
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
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)
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,
)
self.topk_idx = torch.empty(
max_num_tokens, top_k,
dtype=torch.int32, device=device,
)
self.topk_weights = torch.empty(
max_num_tokens, top_k,
dtype=torch.float32, device=device,
)
# All-reduce buffer
self.buffer = torch.empty(
max_num_tokens, hidden_size,
dtype=torch.bfloat16, device=device,
)
if MEGA_MOE_DEBUG:
print(f"[SymmBuffer] x={self.x.shape} x_sf={self.x_sf.shape} "
f"topk_idx={self.topk_idx.shape} topk_weights={self.topk_weights.shape} "
f"buffer={self.buffer.shape}")
def get_symm_buffer_for_nvfp4_mega_moe(
group,
num_experts: int,
max_num_tokens: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
) -> SymmBuffer:
"""Allocate a symmetric buffer for the NVFP4 mega_moe kernel.
API matches deep_gemm.mega.get_symm_buffer_for_nvfp4_mega_moe.
"""
return SymmBuffer(
group, num_experts, max_num_tokens, top_k,
hidden_size, intermediate_size,
)

118
src/nvfp4_megamoe_kernel/weight_transform.py Normal file
View File

@@ -0,0 +1,118 @@
"""
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)
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
"""
# 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()
return transformed_weight, sf_packed
# 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

95
src/quantize.mojo Normal file
View File

@@ -0,0 +1,95 @@
"""
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))

65
src/staging.py Normal file
View File

@@ -0,0 +1,65 @@
"""
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

147
src/weight_transform.py Normal file
View File

@@ -0,0 +1,147 @@
"""
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()