[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)

* Merge with private repo

* Update README

* Update README

* Update README

* Add PyTorch requirements

* Fix sync scopes for MQA logits (#256)

* Update README

deep_gemm/utils/__init__.py

@@ -1,3 +1,4 @@
 from . import math, layout
 from .layout import *
 from .math import *
+from .dist import init_dist, uneven_all_gather

deep_gemm/utils/dist.py

@@ -0,0 +1,74 @@
+import inspect
+import os
+import torch
+import torch.distributed as dist
+from typing import Tuple
+
+_local_rank = None
+
+
+def init_dist(local_rank: int, num_local_ranks: int) -> Tuple[int, int, dist.ProcessGroup]:
+    # NOTES: you may rewrite this function with your own cluster settings
+    ip = os.getenv('MASTER_ADDR', '127.0.0.1')
+    port = int(os.getenv('MASTER_PORT', '8361'))
+    num_nodes = int(os.getenv('WORLD_SIZE', 1))
+    node_rank = int(os.getenv('RANK', 0))
+
+    # Set local rank
+    global _local_rank
+    _local_rank = local_rank
+
+    sig = inspect.signature(dist.init_process_group)
+    params = {
+        'backend': 'nccl',
+        'init_method': f'tcp://{ip}:{port}',
+        'world_size': num_nodes * num_local_ranks,
+        'rank': node_rank * num_local_ranks + local_rank,
+    }
+    if 'device_id' in sig.parameters:
+        # noinspection PyTypeChecker
+        params['device_id'] = torch.device(f'cuda:{local_rank}')
+    dist.init_process_group(**params)
+    torch.set_default_device('cuda')
+    torch.cuda.set_device(local_rank)
+
+    return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes)))
+
+
+def uneven_all_gather(tensor: torch.Tensor, dim: int = 0, group: dist.ProcessGroup = None) -> torch.Tensor:
+    world_size = dist.get_world_size(group)
+
+    # Exchange sizes
+    local_dim_size = torch.tensor([tensor.shape[dim]], device=tensor.device, dtype=torch.long)
+    all_dim_sizes = [torch.zeros_like(local_dim_size) for _ in range(world_size)]
+    dist.all_gather(all_dim_sizes, local_dim_size, group=group)
+    all_dim_sizes = [s.item() for s in all_dim_sizes]
+    max_dim_size = max(all_dim_sizes)
+
+    # Pad
+    if tensor.shape[dim] < max_dim_size:
+        pad_shape = list(tensor.shape)
+        pad_shape[dim] = max_dim_size - tensor.shape[dim]
+        padding = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device)
+        tensor_padded = torch.cat([tensor, padding], dim=dim)
+    else:
+        tensor_padded = tensor.contiguous()
+
+    # All-gather
+    gathered = [torch.zeros_like(tensor_padded) for _ in range(world_size)]
+    dist.all_gather(gathered, tensor_padded, group=group)
+
+    # Remove padding
+    trimmed = [
+        torch.narrow(gathered[i], dim, 0, all_dim_sizes[i])
+        for i in range(world_size)
+    ]
+    return torch.cat(trimmed, dim=dim)
+
+
+def dist_print(s: str = '', once_in_node: bool = False) -> None:
+    global _local_rank
+    assert _local_rank is not None
+    if not once_in_node or _local_rank == 0:
+        print(s, flush=True)
+    dist.barrier()

deep_gemm/utils/layout.py

@@ -10,7 +10,11 @@ except ImportError:
     pass
 
 # Valid for all CUDA versions
-from .._C import get_mk_alignment_for_contiguous_layout
+from .._C import (
+    set_mk_alignment_for_contiguous_layout,
+    get_mk_alignment_for_contiguous_layout,
+    get_theoretical_mk_alignment_for_contiguous_layout,
+)
 
 # Some alias
 get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout

deep_gemm/utils/math.py

@@ -11,21 +11,30 @@ def align(x: int, y: int) -> int:
 
 
 def ceil_to_ue8m0(x: torch.Tensor):
-    assert x.view(-1).amax().item() > 0
-    return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
+    bits = x.abs().float().view(torch.int)
+    exp = ((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int()
+    return (exp.clamp(1, 254) << 23).view(torch.float)
 
 
-def per_token_cast_to_fp8(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:
+def pack_ue8m0_to_int(x: torch.Tensor):
+    assert x.dtype == torch.float and x.size(-1) % 4 == 0
+    assert (x.view(torch.int) & ((1 << 23) - 1) == 0).all()
+    return (x.view(torch.int) >> 23).to(torch.uint8).view(torch.int)
+
+
+def per_token_cast_to_fp8(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128,
+                          use_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
     assert x.dim() == 2
     m, n = x.shape
     padded_n = align(n, gran_k)
     x_padded = torch.empty((m, padded_n), dtype=x.dtype, device=x.device).fill_(0)
     x_padded[:, :n] = x
-    x_view = x_padded.view(m, -1, gran_k)
-    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
+    x_view = x_padded.view(m, padded_n // gran_k, gran_k)
+    x_amax = x_view.abs().float().amax(dim=2).view(m, padded_n // gran_k).clamp(1e-4)
     sf = x_amax / 448.0
     sf = ceil_to_ue8m0(sf) if use_ue8m0 else sf
-    return (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, padded_n)[:, :n].contiguous(), sf
+    x_fp8 = (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, padded_n)[:, :n].contiguous()
+    return x_fp8, pack_ue8m0_to_int(sf) if use_packed_ue8m0 else sf
 
 
 def per_channel_cast_to_fp8(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -70,13 +79,14 @@ def _quantize_to_fp4_e2m1(x: torch.Tensor) -> torch.Tensor:
     code = idx.to(torch.uint8)
     sign = (x < 0) & (idx != 0)
     code = code | (sign.to(torch.uint8) << 3)
-    return code  # uint8, 0..15
+    return code.view(torch.int8)
 
 
-def per_token_cast_to_fp4(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:
-    assert x.dim() == 2
+def per_token_cast_to_fp4(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128,
+                          use_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
     m, n = x.shape
     assert n % 2 == 0
+    assert not use_packed_ue8m0 or use_ue8m0
     padded_n = align(n, gran_k)
     x_padded = torch.zeros((m, padded_n), dtype=x.dtype, device=x.device)
     x_padded[:, :n] = x
@@ -85,23 +95,49 @@ def per_token_cast_to_fp4(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128) -
     sf = x_amax / 6.0
     sf = ceil_to_ue8m0(sf) if use_ue8m0 else sf
     x_scaled = x_view * (1.0 / sf.unsqueeze(2))
-    codes = _quantize_to_fp4_e2m1(x_scaled).view(m, padded_n)  # uint8, (m, padded_n)
+    codes = _quantize_to_fp4_e2m1(x_scaled).view(m, padded_n)  # int8, (m, padded_n)
     codes2 = codes.view(m, padded_n // 2, 2)
-    packed = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4)  # uint8
-    return packed[:, :n // 2].contiguous(), sf
+    packed = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4)  # int8
+    return packed[:, :n // 2].contiguous(), pack_ue8m0_to_int(sf) if use_packed_ue8m0 else sf
 
 
 def transpose_packed_fp4(a: torch.Tensor) -> torch.Tensor:
-    assert a.dtype == torch.uint8
+    assert a.dtype == torch.int8
     assert a.dim() == 2
     m, n2 = a.shape
     n = n2 * 2
     assert (m % 2) == 0
     lo = a & 0x0F
     hi = (a >> 4) & 0x0F
-    codes = torch.empty((m, n), device=a.device, dtype=torch.uint8)
+    codes = torch.empty((m, n), device=a.device, dtype=torch.int8)
     codes[:, 0::2], codes[:, 1::2] = lo, hi
     codes_t = codes.transpose(0, 1).contiguous()
     codes2 = codes_t.view(n, m // 2, 2)
     out = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4)
-    return out.contiguous()
+    return out.contiguous()
+
+
+def _dequantize_from_fp4_e2m1(x: torch.Tensor) -> torch.Tensor:
+    fp4_values = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=x.device, dtype=torch.float)
+    sign, value_idx = (x & 0x08) != 0, (x & 0x07).to(torch.int)
+    value = fp4_values[value_idx]
+    return torch.where(sign & (value_idx != 0), -value, value)
+
+
+def unpack_ue8m0_from_int(packed_sf: torch.Tensor) -> torch.Tensor:
+    return (packed_sf.view(torch.uint8).to(torch.int) << 23).view(torch.float)
+
+
+def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128,
+                       use_packed_ue8m0: bool = False) -> torch.Tensor:
+    m, n2 = packed.shape
+    n = n2 * 2
+    if use_packed_ue8m0:
+        sf = unpack_ue8m0_from_int(sf)
+    unpacked = torch.zeros((m, n), dtype=torch.int8, device=packed.device)
+    unpacked[:, ::2] = packed & 0x0F
+    unpacked[:, 1::2] = (packed >> 4) & 0x0F
+    x_dequantized = _dequantize_from_fp4_e2m1(unpacked)
+    group_idx = torch.arange(n, device=packed.device) // gran_k
+    x_restored = x_dequantized * sf[:, group_idx]
+    return x_restored