007c645f87
Separate source and destination FP4 scale packing in requant_fp4_to_gran_k so group16-to-group32 conversion always recomputes UE8M0 runtime scales by default. Make prepare_fp4_weights_for_mega_moe accept raw grouped FP4 weights and scales, then perform optional requantization, DeepGEMM scale layout transform, and MegaMoE UTCCP weight transform internally. Update the MegaMoE synthetic benchmark so baseline grouped GEMM uses runtime-layout weights while fused MegaMoE uses transformed weights from the same raw source tensors. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/__init__.py deep_gemm/mega/__init__.py deep_gemm/utils/math.py tests/test_layout.py tests/test_mega_moe.py Tested: git diff --check Not-tested: CUDA build, SM100/B300 runtime, and GLM-5.2 accuracy validation are not available locally.
278 lines
11 KiB
Python
278 lines
11 KiB
Python
import torch
|
|
from typing import Tuple, Optional
|
|
from ..utils.math import align, requant_fp4_to_gran_k, unpack_ue8m0_from_int
|
|
|
|
# noinspection PyBroadException
|
|
try:
|
|
# noinspection PyProtectedMember
|
|
import torch.distributed._symmetric_memory as symm_mem
|
|
import torch.distributed as dist
|
|
except Exception as exception:
|
|
print(f'Failed to load mega kernels, please check your PyTorch version: {exception}')
|
|
|
|
from .. import _C
|
|
|
|
|
|
def _is_sm90() -> bool:
|
|
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
|
|
|
|
|
|
def _is_sm100() -> bool:
|
|
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 10
|
|
|
|
|
|
class SymmBuffer:
|
|
def __init__(self, group: dist.ProcessGroup,
|
|
# MoE arguments
|
|
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'):
|
|
self.group = group
|
|
self.num_experts = num_experts
|
|
self.num_max_tokens_per_rank = num_max_tokens_per_rank
|
|
self.num_topk = num_topk
|
|
self.hidden = hidden
|
|
self.intermediate_hidden = intermediate_hidden
|
|
|
|
# Allocate a symmetric buffer (route by architecture)
|
|
if _is_sm90():
|
|
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe(
|
|
group.size(), num_experts,
|
|
num_max_tokens_per_rank, num_topk,
|
|
hidden, intermediate_hidden,
|
|
use_fp8_dispatch, activation
|
|
)
|
|
elif _is_sm100():
|
|
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe(
|
|
group.size(), num_experts,
|
|
num_max_tokens_per_rank, num_topk,
|
|
hidden, intermediate_hidden,
|
|
use_fp8_dispatch, activation
|
|
)
|
|
else:
|
|
raise RuntimeError('Unsupported architecture for MegaMoE')
|
|
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
|
|
self.handle = symm_mem.rendezvous(self.buffer, group=group)
|
|
self.buffer.zero_()
|
|
self.group.barrier()
|
|
torch.cuda.synchronize()
|
|
|
|
# Create input buffer views
|
|
buffer_views = slice_input_buffers(self.buffer)
|
|
if _is_sm90():
|
|
(self.x, self.x_sf,
|
|
self.topk_idx, self.topk_weights,
|
|
self.l1_acts, self.l1_acts_sf, self.l1_topk_weights,
|
|
self.l2_acts, self.l2_acts_sf,
|
|
self.expert_recv_count_sum,
|
|
self.l1_arrival_count,
|
|
self.l2_arrival_mask,
|
|
self.token_src_metadata,
|
|
self.l1_accum_debug,
|
|
self.combine_acts) = buffer_views
|
|
else:
|
|
(self.x, self.x_sf,
|
|
self.topk_idx, self.topk_weights,
|
|
self.l1_acts, self.l1_acts_sf,
|
|
self.l2_acts, self.l2_acts_sf) = buffer_views
|
|
self.l1_topk_weights = None
|
|
self.expert_recv_count_sum = None
|
|
self.l1_arrival_count = None
|
|
self.l2_arrival_mask = None
|
|
self.token_src_metadata = None
|
|
self.l1_accum_debug = None
|
|
self.combine_acts = None
|
|
|
|
def destroy(self):
|
|
self.handle = None
|
|
self.buffer = None
|
|
self.group = None
|
|
self.x = None
|
|
self.x_sf = None
|
|
self.topk_idx = None
|
|
self.topk_weights = None
|
|
self.l1_acts = None
|
|
self.l1_acts_sf = None
|
|
self.l1_topk_weights = None
|
|
self.l2_acts = None
|
|
self.l2_acts_sf = None
|
|
self.expert_recv_count_sum = None
|
|
self.l1_arrival_count = None
|
|
self.l2_arrival_mask = None
|
|
self.token_src_metadata = None
|
|
self.l1_accum_debug = None
|
|
self.combine_acts = None
|
|
|
|
|
|
def get_symm_buffer_for_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:
|
|
# Token count must be aligned to block sizes
|
|
if _is_sm90():
|
|
alignment = _C.get_token_alignment_for_sm90_mega_moe()
|
|
elif _is_sm100():
|
|
alignment = _C.get_token_alignment_for_mega_moe()
|
|
else:
|
|
raise RuntimeError('Unsupported architecture for MegaMoE')
|
|
num_max_tokens_per_rank = align(num_max_tokens_per_rank, alignment)
|
|
|
|
return SymmBuffer(
|
|
group, num_experts,
|
|
num_max_tokens_per_rank, num_topk,
|
|
hidden, intermediate_hidden,
|
|
use_fp8_dispatch, activation
|
|
)
|
|
|
|
|
|
def _interleave_l1_weight_tensor(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
|
|
# [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up]
|
|
g, n, *rest = t.shape
|
|
half = n // 2
|
|
gate = t[:, :half].reshape(g, half // gran, gran, *rest)
|
|
up = t[:, half:].reshape(g, half // gran, gran, *rest)
|
|
return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))
|
|
|
|
|
|
def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1])
|
|
|
|
|
|
def _transpose_sf_for_utccp(sf: torch.Tensor, gran_k: int = 32) -> torch.Tensor:
|
|
num_groups, mn, packed_sf_k = sf.shape
|
|
assert sf.dtype == torch.int and mn % 128 == 0
|
|
assert 128 % gran_k == 0
|
|
result = (sf.reshape(num_groups, -1, 128 // gran_k, gran_k, packed_sf_k)
|
|
.transpose(2, 3)
|
|
.reshape(num_groups, mn, packed_sf_k))
|
|
return torch.empty_like(sf).copy_(result)
|
|
|
|
|
|
def transform_weights_for_mega_moe_sm90(
|
|
l1_weights: Tuple[torch.Tensor, torch.Tensor],
|
|
l2_weights: Tuple[torch.Tensor, torch.Tensor]
|
|
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
|
|
# L1: interleave FP8 gate/up weights only; SM90 float weight SF stays natural MN-major.
|
|
l1_weights = (_interleave_l1_weight_tensor(l1_weights[0]), l1_weights[1])
|
|
# L2: no transform
|
|
return l1_weights, l2_weights
|
|
|
|
|
|
def transform_weights_for_mega_moe(
|
|
l1_weights: Tuple[torch.Tensor, torch.Tensor],
|
|
l2_weights: Tuple[torch.Tensor, torch.Tensor],
|
|
weight_gran_k: int = 32,
|
|
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
|
|
if _is_sm90():
|
|
return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights)
|
|
# SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose
|
|
l1_interleaved = _interleave_l1_weights(l1_weights)
|
|
l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1], weight_gran_k))
|
|
l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1], weight_gran_k))
|
|
return l1_weights, l2_weights
|
|
|
|
|
|
def _prepare_raw_fp4_weight_for_mega_moe(
|
|
weights: Tuple[torch.Tensor, torch.Tensor],
|
|
source_weight_gran_k: int,
|
|
runtime_weight_gran_k: int,
|
|
source_scale_packed_ue8m0: bool,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
weight, weight_sf = weights
|
|
if source_weight_gran_k != runtime_weight_gran_k:
|
|
if source_weight_gran_k != 16 or runtime_weight_gran_k != 32:
|
|
raise RuntimeError(
|
|
f'Unsupported MegaMoE FP4 weight granularity conversion: '
|
|
f'{source_weight_gran_k} -> {runtime_weight_gran_k}')
|
|
weight, weight_sf = requant_fp4_to_gran_k(
|
|
weight, weight_sf,
|
|
source_weight_gran_k, runtime_weight_gran_k,
|
|
src_scale_packed_ue8m0=source_scale_packed_ue8m0)
|
|
source_scale_packed_ue8m0 = False
|
|
|
|
if source_scale_packed_ue8m0:
|
|
weight_sf = unpack_ue8m0_from_int(weight_sf)
|
|
num_groups, mn, packed_k = weight.shape
|
|
weight_sf = _C.transform_sf_into_required_layout(
|
|
weight_sf, mn, packed_k * 2, (1, runtime_weight_gran_k), num_groups)
|
|
return weight, weight_sf
|
|
|
|
|
|
def prepare_fp4_weights_for_mega_moe(
|
|
l1_weights: Tuple[torch.Tensor, torch.Tensor],
|
|
l2_weights: Tuple[torch.Tensor, torch.Tensor],
|
|
source_weight_gran_k: int = 32,
|
|
runtime_weight_gran_k: int = 32,
|
|
source_scale_packed_ue8m0: bool = False,
|
|
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
|
|
l1_weights = _prepare_raw_fp4_weight_for_mega_moe(
|
|
l1_weights, source_weight_gran_k, runtime_weight_gran_k,
|
|
source_scale_packed_ue8m0)
|
|
l2_weights = _prepare_raw_fp4_weight_for_mega_moe(
|
|
l2_weights, source_weight_gran_k, runtime_weight_gran_k,
|
|
source_scale_packed_ue8m0)
|
|
return transform_weights_for_mega_moe(
|
|
l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
|
|
|
|
|
|
def fp8_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: Optional[Tuple[int, int, int]] = None,
|
|
activation: str = 'swiglu',
|
|
activation_clamp: Optional[float] = None,
|
|
fast_math: bool = True):
|
|
if _is_sm90():
|
|
if recipe is None:
|
|
recipe = (1, 128, 128)
|
|
_C.fp8_mega_moe(
|
|
y,
|
|
l1_weights, l2_weights,
|
|
cumulative_local_expert_recv_stats,
|
|
sym_buffer.buffer,
|
|
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
|
|
sym_buffer.num_max_tokens_per_rank,
|
|
sym_buffer.num_experts, sym_buffer.num_topk,
|
|
recipe,
|
|
activation, activation_clamp,
|
|
fast_math
|
|
)
|
|
elif _is_sm100():
|
|
if recipe is None:
|
|
recipe = (1, 1, 32)
|
|
_C.fp8_fp4_mega_moe(
|
|
y,
|
|
l1_weights, l2_weights,
|
|
cumulative_local_expert_recv_stats,
|
|
sym_buffer.buffer,
|
|
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
|
|
sym_buffer.num_max_tokens_per_rank,
|
|
sym_buffer.num_experts, sym_buffer.num_topk,
|
|
recipe,
|
|
activation, activation_clamp,
|
|
fast_math
|
|
)
|
|
else:
|
|
raise RuntimeError('Unsupported architecture for MegaMoE')
|
|
|
|
|
|
# Backward-compatible alias
|
|
def fp8_fp4_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: Optional[Tuple[int, int, int]] = None,
|
|
activation: str = 'swiglu',
|
|
activation_clamp: Optional[float] = None,
|
|
fast_math: bool = True):
|
|
fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer,
|
|
cumulative_local_expert_recv_stats, recipe,
|
|
activation, activation_clamp, fast_math)
|