79fcfd6abf
Allow SM100 FP4 scale layout transforms to accept group16 and thread weight granularity through the MegaMoE Python wrapper, API checks, and synthetic benchmark entrypoint. Keep fused SM100 MegaMoE compute behind an explicit group16 capability gate until the SFB/TMEM/MMA scale path is updated and validated. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/mega/__init__.py tests/test_mega_moe.py tests/generators.py Tested: git diff --check Not-tested: CUDA build and SM100/B300 runtime validation are not available locally.
235 lines
9.1 KiB
Python
235 lines
9.1 KiB
Python
import torch
|
|
from typing import Tuple, Optional
|
|
from ..utils.math import align
|
|
|
|
# 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 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)
|