revert: restore mxf4nvf4/block16 code (correct path for sm_100a)
Reverted to commit 36b439e's NVFP4 kernel code: - kGranK=16, mxf4nvf4.block_scale.scale_vec::4X - float_ue4m3_t instruction descriptor - Block16 SF layout (4X TMEM) - UE4M3 L1 epilogue - No UE4M3→UE8M0 conversion, no block16→block32 merge The mxf4nvf4.scale_vec::4X PTX instruction compiles successfully on both sm_100 and sm_100f with CUDA 13.0. The previous build 17 error was likely from a different cause, not the arch flag. Python: reverted transform_nvfp4_weights_for_mega_moe to use pack_ue4m3_to_int32 with gran_k=16, no UE8M0 conversion.
This commit is contained in:
@@ -138,93 +138,22 @@ def _pack_nvfp4_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
def transform_nvfp4_weights_for_mega_moe(
|
||||
l1_weights: Tuple[torch.Tensor, torch.Tensor],
|
||||
l2_weights: Tuple[torch.Tensor, torch.Tensor],
|
||||
l1_weight_scale_2: Optional[torch.Tensor] = None,
|
||||
l2_weight_scale_2: Optional[torch.Tensor] = None
|
||||
l2_weights: Tuple[torch.Tensor, torch.Tensor]
|
||||
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Transform NVFP4 expert weights for the mega_moe kernel.
|
||||
|
||||
Uses deep_gemm.transform_sf_into_required_layout for proper TMA-aligned
|
||||
UTCCP layout with recipe (1, 1, 16) for NVFP4 group_size=16.
|
||||
NVFP4 weights come as (weight, scale) where:
|
||||
- weight: uint8 E2M1 packed, shape (num_experts, N, K//2)
|
||||
- scale: float8_e4m3fn UE4M3 block scales, shape (num_experts, N, K//16)
|
||||
|
||||
The kernel expects (weight, packed_sf) where packed_sf is int32 UTCCP layout.
|
||||
"""
|
||||
from deep_gemm import transform_sf_into_required_layout
|
||||
|
||||
def fold_global_scale(sf: torch.Tensor, scale_2: Optional[torch.Tensor]) -> torch.Tensor:
|
||||
if scale_2 is None:
|
||||
return sf
|
||||
sf_f32 = sf.to(torch.float32)
|
||||
if scale_2.dim() == 1:
|
||||
scale_2 = scale_2.view(-1, 1, 1)
|
||||
sf_f32 = sf_f32 * scale_2
|
||||
sf_f32 = sf_f32.clamp(0.0, 448.0)
|
||||
return sf_f32.to(torch.float8_e4m3fn)
|
||||
|
||||
l1_sf = fold_global_scale(l1_weights[1], l1_weight_scale_2)
|
||||
l2_sf = fold_global_scale(l2_weights[1], l2_weight_scale_2)
|
||||
|
||||
# Merge NVFP4 block16 scales → block32 for SM100 (scale_vec::2X)
|
||||
# B200 (SM100) doesn't support scale_vec::4X (block16) — requires SM103/SM120
|
||||
# Take the max of each pair of adjacent block16 scales for block32
|
||||
def merge_block16_to_block32(sf):
|
||||
# sf: (experts, mn, K//16) float8_e4m3fn
|
||||
# output: (experts, mn, K//32) uint8 (UE8M0)
|
||||
# SM100 (B200) doesn't support mxf4nvf4 — must use mxf8f6f4 with UE8M0 scales
|
||||
# Convert UE4M3 → float32 → UE8M0 (power-of-2)
|
||||
sf_f32 = sf.to(torch.float32)
|
||||
# Take max of adjacent pairs
|
||||
sf_merged = torch.maximum(sf_f32[..., 0::2], sf_f32[..., 1::2])
|
||||
# Convert to UE8M0: extract exponent byte from float32 bit pattern
|
||||
# UE8M0: uint8 = float32_bits[31:23] (8 exponent bits)
|
||||
# Note: PyTorch doesn't support >> on uint32, cast to int32 first
|
||||
sf_bits = sf_merged.view(torch.int32) # reinterpret float32 bits as int32
|
||||
sf_ue8m0 = ((sf_bits >> 23) & 0xFF).to(torch.uint8)
|
||||
return sf_ue8m0
|
||||
|
||||
l1_sf_32 = merge_block16_to_block32(l1_sf)
|
||||
l2_sf_32 = merge_block16_to_block32(l2_sf)
|
||||
|
||||
num_experts = l1_weights[0].shape[0]
|
||||
l1_n = l1_weights[0].shape[1]
|
||||
l1_k = l1_weights[0].shape[2] * 2
|
||||
l2_n = l2_weights[0].shape[1]
|
||||
l2_k = l2_weights[0].shape[2] * 2
|
||||
|
||||
# Pack UE8M0 (uint8) block scales into int32 for DeepGEMM TMA consumption
|
||||
# Same packing as MXFP4: 4 uint8 → 1 int32
|
||||
def pack_uint8_to_int32(sf):
|
||||
assert sf.dtype == torch.uint8
|
||||
assert sf.shape[-1] % 4 == 0
|
||||
packed = (sf[..., 0::4].to(torch.int32) |
|
||||
(sf[..., 1::4].to(torch.int32) << 8) |
|
||||
(sf[..., 2::4].to(torch.int32) << 16) |
|
||||
(sf[..., 3::4].to(torch.int32) << 24))
|
||||
return packed.contiguous()
|
||||
|
||||
l1_sf_packed = pack_uint8_to_int32(l1_sf_32)
|
||||
l2_sf_packed = pack_uint8_to_int32(l2_sf_32)
|
||||
|
||||
print(f"[NVFP4-MoE] l1_sf_32: shape={l1_sf_32.shape}, l1_sf_packed: shape={l1_sf_packed.shape}")
|
||||
print(f"[NVFP4-MoE] l2_sf_32: shape={l2_sf_32.shape}, l2_sf_packed: shape={l2_sf_packed.shape}")
|
||||
print(f"[NVFP4-MoE] l1_n={l1_n} l1_k={l1_k} l2_n={l2_n} l2_k={l2_k}")
|
||||
|
||||
# Transpose to MN-major layout (stride(-2)=1) and make contiguous
|
||||
# transform_sf_into_required_layout expects MN-major input for TMA stride checks
|
||||
l1_sf_mn = l1_sf_packed.transpose(-2, -1).contiguous().transpose(-2, -1)
|
||||
l2_sf_mn = l2_sf_packed.transpose(-2, -1).contiguous().transpose(-2, -1)
|
||||
|
||||
# Transform SF into TMA-aligned UTCCP layout using DeepGEMM's C++ function
|
||||
# recipe (1, 32): gran_mn=1, gran_k=16
|
||||
l1_sf_transformed = transform_sf_into_required_layout(
|
||||
l1_sf_mn, l1_n, l1_k, (1, 32), num_experts)
|
||||
l2_sf_transformed = transform_sf_into_required_layout(
|
||||
l2_sf_mn, l2_n, l2_k, (1, 32), num_experts)
|
||||
|
||||
# L1: interleave gate/up
|
||||
l1_interleaved = _interleave_l1_weights((l1_weights[0], l1_sf_packed))
|
||||
# DeepGEMM expects int8 (kPackedFP4 = torch.kInt8)
|
||||
l1_out = (l1_interleaved[0].view(torch.int8), l1_sf_transformed)
|
||||
l2_out = (l2_weights[0].view(torch.int8), l2_sf_transformed)
|
||||
return l1_out, l2_out
|
||||
# L1: interleave gate/up, then pack + transpose SF for UTCCP
|
||||
l1_interleaved = _interleave_l1_weights(l1_weights)
|
||||
l1_weights = (l1_interleaved[0], _pack_nvfp4_sf_for_utccp(l1_interleaved[1]))
|
||||
# L2: only pack + transpose SF for UTCCP
|
||||
l2_weights = (l2_weights[0], _pack_nvfp4_sf_for_utccp(l2_weights[1]))
|
||||
return l1_weights, l2_weights
|
||||
|
||||
|
||||
def fp8_fp4_mega_moe(y: torch.Tensor,
|
||||
@@ -250,49 +179,12 @@ def fp8_fp4_mega_moe(y: torch.Tensor,
|
||||
)
|
||||
|
||||
|
||||
def get_symm_buffer_for_nvfp4_mega_moe(
|
||||
group: "dist.ProcessGroup",
|
||||
num_experts: int,
|
||||
num_max_tokens_per_rank: int, num_topk: int,
|
||||
hidden: int, intermediate_hidden: int,
|
||||
use_fp8_dispatch: bool = True,
|
||||
activation: str = 'swiglu') -> SymmBuffer:
|
||||
"""Allocate a SymmBuffer sized for NVFP4 mega_moe (group_size=16)."""
|
||||
from .. import _C
|
||||
num_max_tokens_per_rank = align(num_max_tokens_per_rank,
|
||||
_C.get_token_alignment_for_nvfp4_mega_moe())
|
||||
buf = SymmBuffer.__new__(SymmBuffer)
|
||||
buf.group = group
|
||||
buf.num_experts = num_experts
|
||||
buf.num_max_tokens_per_rank = num_max_tokens_per_rank
|
||||
buf.num_topk = num_topk
|
||||
buf.hidden = hidden
|
||||
buf.intermediate_hidden = intermediate_hidden
|
||||
# Use NVFP4-specific buffer size (2x SF due to group_size=16)
|
||||
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_nvfp4_mega_moe(
|
||||
group.size(), num_experts,
|
||||
num_max_tokens_per_rank, num_topk,
|
||||
hidden, intermediate_hidden,
|
||||
use_fp8_dispatch, activation)
|
||||
import torch.distributed._symmetric_memory as symm_mem
|
||||
import torch.distributed as dist
|
||||
buf.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
|
||||
buf.handle = symm_mem.rendezvous(buf.buffer, group=group)
|
||||
buf.buffer.zero_()
|
||||
buf.group.barrier()
|
||||
torch.cuda.synchronize()
|
||||
buf.x, buf.x_sf, buf.topk_idx, buf.topk_weights, \
|
||||
buf.l1_acts, buf.l1_acts_sf, buf.l2_acts, buf.l2_acts_sf = \
|
||||
slice_input_buffers(buf.buffer)
|
||||
return buf
|
||||
|
||||
|
||||
def fp8_nvfp4_mega_moe(y: torch.Tensor,
|
||||
l1_weights: Tuple[torch.Tensor, torch.Tensor],
|
||||
l2_weights: Tuple[torch.Tensor, torch.Tensor],
|
||||
sym_buffer: SymmBuffer,
|
||||
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
|
||||
recipe: Tuple[int, int, int] = (1, 1, 32),
|
||||
recipe: Tuple[int, int, int] = (1, 1, 16),
|
||||
activation: str = 'swiglu',
|
||||
activation_clamp: Optional[float] = None,
|
||||
fast_math: bool = True):
|
||||
|
||||
Reference in New Issue
Block a user