fix: use scale_vec::2X (block32) for SM100 B200 compatibility
scale_vec::4X (block16) requires SM103/SM120 (B300/GB300), not SM100 (B200). Revert to block32 with UE4M3 scales. Same TMEM layout as MXFP4 but with UE4M3 scale format instead of UE8M0. Changes: - kGranK: 16 → 32 - PTX: scale_vec::4X → scale_vec::2X - SF layout: same as MXFP4 (K/32, K/128 for int32 packed) - UTCCP: i*8 → i*4 (2X layout, same as MXFP4) - TMEM columns: same as MXFP4 (SF_BLOCK_M/32, SF_BLOCK_N/32) - Python: merge NVFP4 block16→block32 scales (max of adjacent pairs) - recipe: (1,1,16) → (1,1,32)
This commit is contained in:
@@ -162,6 +162,20 @@ def transform_nvfp4_weights_for_mega_moe(
|
||||
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) float8_e4m3fn
|
||||
sf_f32 = sf.to(torch.float32)
|
||||
# Take max of adjacent pairs (preserves magnitude, avoids underflow)
|
||||
sf_merged = torch.maximum(sf_f32[..., 0::2], sf_f32[..., 1::2])
|
||||
return sf_merged.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
||||
|
||||
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
|
||||
@@ -179,8 +193,8 @@ def transform_nvfp4_weights_for_mega_moe(
|
||||
(sf_u8[..., 3::4].to(torch.int32) << 24))
|
||||
return packed.contiguous()
|
||||
|
||||
l1_sf_packed = pack_ue4m3_to_int32(l1_sf)
|
||||
l2_sf_packed = pack_ue4m3_to_int32(l2_sf)
|
||||
l1_sf_packed = pack_ue4m3_to_int32(l1_sf_32)
|
||||
l2_sf_packed = pack_ue4m3_to_int32(l2_sf_32)
|
||||
|
||||
# Transpose to MN-major layout (stride(-2)=1) and make contiguous
|
||||
# transform_sf_into_required_layout expects MN-major input for TMA stride checks
|
||||
@@ -188,11 +202,11 @@ def transform_nvfp4_weights_for_mega_moe(
|
||||
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, 16): gran_mn=1, gran_k=16
|
||||
# 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, 16), num_experts)
|
||||
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, 16), num_experts)
|
||||
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))
|
||||
@@ -267,7 +281,7 @@ def fp8_nvfp4_mega_moe(y: 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, 16),
|
||||
recipe: Tuple[int, int, int] = (1, 1, 32),
|
||||
activation: str = 'swiglu',
|
||||
activation_clamp: Optional[float] = None,
|
||||
fast_math: bool = True):
|
||||
|
||||
Reference in New Issue
Block a user