2c7543130b
Add a top-level MegaMoE helper that handles source/runtime FP4 weight granularity before applying the existing MegaMoE weight layout transform. Use the helper from the synthetic MegaMoE benchmark so SGLang can later follow the same contract for GLM-5.2 NVFP4 group16 checkpoints. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/__init__.py deep_gemm/mega/__init__.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.
254 lines
10 KiB
Python
254 lines
10 KiB
Python
import torch
|
|
from typing import Tuple, Optional
|
|
from ..utils.math import align, requant_fp4_to_gran_k
|
|
|
|
# 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_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,
|
|
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
|
|
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}')
|
|
l1_weights = requant_fp4_to_gran_k(
|
|
l1_weights[0], l1_weights[1], source_weight_gran_k, runtime_weight_gran_k)
|
|
l2_weights = requant_fp4_to_gran_k(
|
|
l2_weights[0], l2_weights[1], source_weight_gran_k, runtime_weight_gran_k)
|
|
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)
|