[autoresearch] Goal: fold block-scaled FP8 into DeepGEMM's MegaMoE (the H200-now subset)

deep_gemm/mega/__init__.py

@@ -1,7 +1,7 @@
 import torch
 import types
 from typing import Tuple, Optional
-from ..utils.math import align
+from ..utils.math import align, per_token_cast_to_fp8, unpack_ue8m0_from_int
 
 # noinspection PyBroadException
 try:
@@ -99,6 +99,13 @@ def transform_weights_for_mega_moe(
     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]]:
+    arch_major = torch.cuda.get_device_capability()[0]
+    is_fp8_weights = (l1_weights[0].dtype == torch.float8_e4m3fn)
+
+    if arch_major == 9 and is_fp8_weights:
+        # SM90 FP8 path: no interleaving or UTCCP transpose needed
+        return l1_weights, l2_weights
+
     # L1: interleave gate/up for weight and SF, then transpose SF for UTCCP.
     l1_w = _interleave_weights(l1_weights[0])
     l1_sf = _transpose_sf_for_utccp(_interleave_weights(l1_weights[1]))
@@ -109,6 +116,161 @@ def transform_weights_for_mega_moe(
 
 
 
+def _fp8_mega_moe_composed(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,
+                           activation_clamp: Optional[float] = None,
+                           fast_math: bool = True):
+    """Composed (non-fused) MegaMoE for SM90/Hopper with block-scaled FP8 weights.
+
+    Chains existing m_grouped FP8 GEMMs with a SwiGLU activation step.
+    Single-rank only (multi-rank requires the fused kernel or deep_ep).
+    """
+    num_tokens = y.size(0)
+    hidden = y.size(1)
+    num_experts = sym_buffer.num_experts
+    num_topk = sym_buffer.num_topk
+    intermediate_hidden = sym_buffer.intermediate_hidden
+    num_ranks = sym_buffer.group.size()
+
+    assert num_ranks == 1, "Composed SM90 MegaMoE only supports single-rank"
+    num_experts_per_rank = num_experts
+
+    x_fp8_raw = sym_buffer.x[:num_tokens]
+    x_sf_raw = sym_buffer.x_sf[:num_tokens]
+    topk_idx = sym_buffer.topk_idx[:num_tokens]
+    topk_weights_tensor = sym_buffer.topk_weights[:num_tokens]
+
+    l1_w, l1_w_sf = l1_weights
+    l2_w, l2_w_sf = l2_weights
+
+    # Convert input from UE8M0 per-32 scales to FP32 per-128 scales for SM90 GEMM.
+    # Dequantize from per-32 UE8M0, then re-quantize with per-128 FP32 scales.
+    if x_sf_raw.dtype == torch.int:
+        sf_per32 = unpack_ue8m0_from_int(x_sf_raw)  # [num_tokens, hidden/32] float
+        x_dequant = x_fp8_raw.float()
+        gran_k = 32
+        x_view = x_dequant.view(num_tokens, -1, gran_k)
+        sf_expanded = sf_per32.unsqueeze(2).expand_as(x_view)
+        x_dequant = (x_view * sf_expanded).view(num_tokens, hidden).bfloat16()
+        x_fp8, x_sf = per_token_cast_to_fp8(x_dequant, use_ue8m0=False, gran_k=128)
+    else:
+        x_fp8 = x_fp8_raw
+        x_sf = x_sf_raw
+
+    # --- Step 1: Dispatch (scatter tokens by expert) ---
+    flat_topk_idx = topk_idx.reshape(-1)
+    flat_token_idx = torch.arange(num_tokens, device=y.device).unsqueeze(1).expand(
+        -1, num_topk).reshape(-1)
+    flat_topk_weights = topk_weights_tensor.reshape(-1)
+
+    valid_mask = flat_topk_idx >= 0
+    valid_expert_idx = flat_topk_idx[valid_mask]
+    valid_token_idx = flat_token_idx[valid_mask]
+    valid_topk_weights = flat_topk_weights[valid_mask]
+
+    sorted_order = torch.argsort(valid_expert_idx, stable=True)
+    sorted_expert_idx = valid_expert_idx[sorted_order]
+    sorted_token_idx = valid_token_idx[sorted_order]
+    sorted_topk_weights = valid_topk_weights[sorted_order]
+
+    num_recv_tokens = sorted_token_idx.size(0)
+    if num_recv_tokens == 0:
+        y.zero_()
+        return
+
+    # Build grouped layout: cumulative token counts per expert (psum format)
+    expert_counts = torch.bincount(
+        sorted_expert_idx.int(), minlength=num_experts_per_rank
+    )[:num_experts_per_rank].to(torch.int)
+
+    # Alignment: pad each expert's token count to the alignment boundary
+    mk_alignment = _C.get_theoretical_mk_alignment_for_contiguous_layout(None, None)
+    _C.set_mk_alignment_for_contiguous_layout(mk_alignment)
+
+    # Build padded grouped layout for the GEMM (each expert aligned)
+    padded_expert_counts = ((expert_counts + mk_alignment - 1) // mk_alignment * mk_alignment).to(torch.int)
+    padded_grouped_layout = torch.cumsum(padded_expert_counts, dim=0).to(torch.int)
+    total_padded_tokens = padded_grouped_layout[-1].item()
+
+    # Build padded recv_x: [total_padded_tokens, hidden] FP8, zero-padded per expert
+    recv_x_fp8 = torch.zeros((total_padded_tokens, hidden),
+                             dtype=torch.float8_e4m3fn, device=y.device)
+    recv_x_sf = torch.zeros((total_padded_tokens, hidden // 128),
+                            dtype=x_sf.dtype, device=y.device)
+    recv_topk_weights = torch.zeros(total_padded_tokens, dtype=torch.float32,
+                                    device=y.device)
+
+    # Copy tokens into padded positions
+    src_offset = 0
+    for e in range(num_experts_per_rank):
+        count = expert_counts[e].item()
+        dst_offset = (padded_grouped_layout[e - 1].item() if e > 0 else 0)
+        if count > 0:
+            idx = sorted_token_idx[src_offset:src_offset + count]
+            recv_x_fp8[dst_offset:dst_offset + count] = x_fp8[idx]
+            recv_x_sf[dst_offset:dst_offset + count] = x_sf[idx]
+            recv_topk_weights[dst_offset:dst_offset + count] = sorted_topk_weights[src_offset:src_offset + count]
+            src_offset += count
+
+    # --- Step 2: L1 GEMM ---
+    l1_y = torch.zeros((total_padded_tokens, intermediate_hidden * 2),
+                       dtype=torch.bfloat16, device=y.device)
+    if total_padded_tokens > 0:
+        _C.m_grouped_fp8_fp4_gemm_nt_contiguous(
+            (recv_x_fp8, recv_x_sf), l1_weights, l1_y, padded_grouped_layout,
+            None, None, None, '', False, True, None)
+
+    # --- Step 3: SwiGLU + FP8 requant ---
+    gate = l1_y[:, :intermediate_hidden].float()
+    up = l1_y[:, intermediate_hidden:].float()
+
+    clamp = activation_clamp if activation_clamp is not None else float('inf')
+    if clamp < float('inf'):
+        gate = gate.clamp(-clamp, clamp)
+        up = up.clamp(-clamp, clamp)
+
+    if fast_math:
+        silu_gate = gate * torch.sigmoid(gate)
+    else:
+        silu_gate = torch.nn.functional.silu(gate)
+
+    activation_out = silu_gate * up * recv_topk_weights.unsqueeze(1)
+
+    l2_acts_fp8, l2_acts_sf = per_token_cast_to_fp8(
+        activation_out.bfloat16(), use_ue8m0=False, gran_k=128)
+
+    # --- Step 4: L2 GEMM ---
+    l2_y = torch.zeros((total_padded_tokens, hidden),
+                       dtype=torch.bfloat16, device=y.device)
+    if total_padded_tokens > 0:
+        _C.m_grouped_fp8_fp4_gemm_nt_contiguous(
+            (l2_acts_fp8, l2_acts_sf), l2_weights, l2_y, padded_grouped_layout,
+            None, None, None, '', False, True, None)
+
+    # --- Step 5: Combine (gather + reduce) ---
+    # Collect unpadded results and scatter-add back to output
+    unpadded_results = torch.zeros((num_recv_tokens, hidden),
+                                   dtype=torch.bfloat16, device=y.device)
+    src_offset = 0
+    for e in range(num_experts_per_rank):
+        count = expert_counts[e].item()
+        dst_offset = (padded_grouped_layout[e - 1].item() if e > 0 else 0)
+        if count > 0:
+            unpadded_results[src_offset:src_offset + count] = l2_y[dst_offset:dst_offset + count]
+            src_offset += count
+
+    y.zero_()
+    y.index_add_(0, sorted_token_idx, unpadded_results)
+
+    # Update stats if provided
+    if cumulative_local_expert_recv_stats is not None:
+        cumulative_local_expert_recv_stats.copy_(
+            torch.cumsum(expert_counts, dim=0).to(torch.int))
+
+
 def fp8_fp4_mega_moe(y: torch.Tensor,
                      l1_weights: Tuple[torch.Tensor, torch.Tensor],
                      l2_weights: Tuple[torch.Tensor, torch.Tensor],
@@ -118,6 +280,19 @@ def fp8_fp4_mega_moe(y: torch.Tensor,
                      activation: str = 'swiglu',
                      activation_clamp: Optional[float] = None,
                      fast_math: bool = True):
+    # Detect architecture and weight type to dispatch
+    arch_major = torch.cuda.get_device_capability()[0]
+    l1_w = l1_weights[0] if isinstance(l1_weights, tuple) else l1_weights
+    is_fp8_weights = (l1_w.dtype == torch.float8_e4m3fn)
+
+    # SM90 with FP8 weights: use composed path
+    if arch_major == 9 and is_fp8_weights:
+        _fp8_mega_moe_composed(
+            y, l1_weights, l2_weights, sym_buffer,
+            cumulative_local_expert_recv_stats,
+            activation_clamp, fast_math)
+        return
+
     _C.fp8_fp4_mega_moe(
         y,
         l1_weights, l2_weights,