Support MXFP4 expert weights in remote prefetch

benchmarks/bench_prefetch.py

@@ -20,6 +20,27 @@
 EXP_H = 3584
 EXP_HP = 3072
 
+K3_BF16_LAYER = [
+    {"transposed": True},                 # gate  -> (I, H)
+    {"transposed": True},                 # up    -> (I, H)
+    {},                                   # down  -> (H, I)
+]
+K3_MXFP4_LAYER = [
+    {"transposed": True, "pack": 2},      # gate packed -> (I, H/2)
+    {"transposed": True, "pack": 2},      # up packed   -> (I, H/2)
+    {"pack": 2},                          # down packed -> (H, I/2)
+    {"scale": True},                      # gate scale  -> (128, 2688)
+    {"scale": True},                      # up scale
+    {"scale": True},                      # down scale
+]
+
+_DT_LABEL = {
+    torch.bfloat16: "bf16",
+    torch.int8: "i8",
+    torch.uint8: "u8",
+}
+_DT_LABEL_TO_TORCH = {v: k for k, v in _DT_LABEL.items()}
+
 # All cases: rank0-initiated, num_sms == 32. Each remote expert is prefetched
 # a *different* number of times in 0..3 (R-1=3 is the max), via `counts[e]`.
 # 0 means that expert is skipped. base_B > epn leaves idle columns so the
@@ -60,9 +81,49 @@
      "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
     {"label": "tiny_512x512", "epn": 8, "H": 512, "Hp": 512, "base_B": 14,
      "counts": [3, 3, 2, 3, 1, 0, 2, 3]},
+    # --- Quantized experts -------------------------------------------------
+    {"label": "mxfp4_gate_up", "epn": 8, "H": EXP_H, "Hp": EXP_HP,
+     "transposed": True, "pack": 2,
+     "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
+    {"label": "mxfp4_down", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "pack": 2,
+     "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
+    {"label": "mxfp4_scale", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "scale": True,
+     "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
+    {"label": "k3_layer_bf16", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "base_B": 14,
+     "parts": K3_BF16_LAYER, "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
+    {"label": "k3_layer_mxfp4", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "base_B": 14,
+     "dtype": torch.uint8,
+     "parts": K3_MXFP4_LAYER, "counts": [3, 0, 2, 1, 3, 0, 2, 1]},
 ]
 
 
+def derive_extents(H, Hp, part):
+    if part.get("scale"):
+        # One ue8m0 byte per 32 values, re-cut into whole 128x128 tiles.
+        nbytes = H * Hp // 32
+        return 128, nbytes // 128
+    out, contracted = (Hp, H) if part.get("transposed") else (H, Hp)
+    return out, contracted // int(part.get("pack", 1))
+
+
+def resolve_parts(case):
+    H, Hp = int(case["H"]), int(case["Hp"])
+    default_dt = case.get("dtype", torch.bfloat16)
+    resolved = []
+    for part in case.get("parts") or [case]:
+        th, thp = derive_extents(H, Hp, part)
+        resolved.append((th, thp, part.get("dtype", default_dt)))
+    return resolved
+
+
+def fill_random(shape, dtype, gen, dev="cuda"):
+    if dtype.is_floating_point:
+        return torch.randn(shape, dtype=dtype, device=dev, generator=gen)
+    info = torch.iinfo(dtype)
+    return torch.randint(info.min, info.max, shape, dtype=dtype, device=dev,
+                         generator=gen)
+
+
 def setup():
     dist.init_process_group(backend="nccl")
     rank = dist.get_rank()
@@ -90,43 +151,50 @@ def bench_case(case, args, rank, R):
     dev = "cuda"
     epn = int(case["epn"])
     E = R * epn
-    H = int(case["H"])
-    Hp = int(case["Hp"])
-    B = pad_dim0_for_alignment([int(case["base_B"]), H, Hp], torch.bfloat16)
+    parts = resolve_parts(case)
+    th0, thp0, dt0 = parts[0]
+    B = pad_dim0_for_alignment([int(case["base_B"]), th0, thp0], dt0)
     counts = case.get("counts") or [int(case.get("nremote", 1))] * epn
     num_sms = NUM_SMS
 
-    assert H % 128 == 0 and Hp % 128 == 0, \
-        f"{case['label']}: H and Hp must be multiples of 128"
+    for th, thp, _dt in parts:
+        assert th % 128 == 0 and thp % 128 == 0, \
+            f"{case['label']}: derived extents must be multiples of 128, " \
+            f"got ({th}, {thp})"
     assert len(counts) == epn, f"{case['label']}: counts must have epn={epn} entries"
     assert all(0 <= c < R for c in counts), f"{case['label']}: each count must be in [0, R)"
 
-    # The remote expert table lives on rank1; rank0 reads it over NVLink.
-    padded_E = pad_dim0_for_alignment([E, H, Hp], torch.bfloat16)
-    mapped = create_nvl_single_owner_tensor(
-        [padded_E, H, Hp], torch.bfloat16, owner_rank=1, local_rank=rank
-    )
-    remote_expert = mapped[:E]
-    prefetch_buffers = torch.empty(R * B, H, Hp, dtype=torch.bfloat16, device=dev)
-
-    if rank == 1:
-        gen = torch.Generator(device=dev).manual_seed(321 + rank)
-        remote_expert.copy_(torch.randn(
-            E, H, Hp, dtype=torch.bfloat16, device=dev, generator=gen
-        ))
+    # The remote expert table lives on one owner; rank0 reads it over NVLink
+    owner_rank = int(args.owner_rank)
+    remote_experts, prefetch_buffers = [], []
+    for i, (th, thp, dt) in enumerate(parts):
+        padded_E = pad_dim0_for_alignment([E, th, thp], dt)
+        mapped = create_nvl_single_owner_tensor(
+            [padded_E, th, thp], dt, owner_rank=owner_rank, local_rank=rank
+        )
+        remote = mapped[:E]
+        if rank == owner_rank:
+            gen = torch.Generator(device=dev).manual_seed(321 + rank + i)
+            remote.copy_(fill_random((E, th, thp), dt, gen, dev))
+        remote_experts.append(remote)
+        prefetch_buffers.append(
+            torch.empty(R * B, th, thp, dtype=dt, device=dev)
+        )
+
     plan = expert_plan(R, B, epn, counts, dev)
     experts_to_copy = plan.flatten() if rank == 0 else \
         torch.full((R * B,), -1, dtype=torch.int32, device=dev)
     torch.cuda.synchronize()
     dist.barrier(device_ids=[torch.cuda.current_device()])
 
     def prefetch_once():
-        launch_prefetch(
-            remote_expert,
-            prefetch_buffers,
-            experts_to_copy,
-            num_sms=num_sms,
-        )
+        for remote, buf in zip(remote_experts, prefetch_buffers):
+            launch_prefetch(
+                remote,
+                buf,
+                experts_to_copy,
+                num_sms=num_sms,
+            )
 
     # Warmup (also JIT-compiles the kernel) then capture the iters loop into a
     # single CUDA graph to strip launch/python overhead.
@@ -162,22 +230,25 @@ def prefetch_once():
     worst_us = start.elapsed_time(end) * 1e3 / args.iters if rank == 0 else 0.0
 
     # Critical-path traffic for the prefetching rank: read every consumed slot
-    # from the remote table over NVLink (2B/elem), then write it to the local
-    # prefetch buffer (2B/elem).
+    # from the remote table over NVLink, then write it to the local prefetch
+    # buffer.
     slots = int(sum(counts))
-    tile = H * Hp * 2
-    bytes_per_rank = slots * tile * 2
+    per_slot = sum(th * thp * dt.itemsize for th, thp, dt in parts)
+    bytes_per_rank = slots * per_slot * 2
     bw_gbs = bytes_per_rank / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0
     # pure NVLink read traffic: only the remote expert-table reads (buffer
     # writes are local HBM, off the NVLink path).
-    comm_gbs = slots * tile / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0
+    comm_gbs = slots * per_slot / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0
 
     dist.barrier(device_ids=[torch.cuda.current_device()])
-    return worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots
+    dt_label = (_DT_LABEL.get(dt0, str(dt0))
+                if len({dt for _, _, dt in parts}) == 1 else "mix")
+    return (worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots,
+            len(parts), dt_label, int(case["H"]), int(case["Hp"]))
 
 
 def explicit_single_config_requested(argv):
-    shape_flags = ("--epn", "--H", "--Hp", "--B", "--nremote")
+    shape_flags = ("--epn", "--H", "--Hp", "--B", "--nremote", "--dtype")
     for arg in argv:
         for flag in shape_flags:
             if arg == flag or arg.startswith(flag + "="):
@@ -196,8 +267,17 @@ def parse_args():
     parser.add_argument("--Hp", type=int, default=3072)
     parser.add_argument("--B", type=int, default=14)
     parser.add_argument("--nremote", type=int, default=2)
+    parser.add_argument("--dtype", choices=sorted(_DT_LABEL_TO_TORCH),
+                        default="bf16",
+                        help="Element type of the prefetched tensor "
+                             "(u8 = packed MXFP4 or its ue8m0 scales)")
     parser.add_argument("--warmup", type=int, default=5)
     parser.add_argument("--iters", type=int, default=20)
+    parser.add_argument("--owner-rank", type=int, default=1,
+                        help="Rank whose GPU physically holds the remote expert "
+                             "rank0 is always the reader. An owner on rank0's own node "
+                             "measures intra-node NVLink, one on another node "
+                             "measures cross-node MNNVL fabric.")
     parser.add_argument("--no-graph", action="store_true",
                         help="Time plain launches instead of a CUDA graph (for NCU)")
     return parser.parse_args()
@@ -210,6 +290,10 @@ def main():
 
     if args.single and args.suite:
         raise ValueError("Use only one of --single or --suite")
+    if not 0 <= args.owner_rank < R:
+        raise ValueError(
+            f"--owner-rank must be in [0, {R}), got {args.owner_rank}"
+        )
 
     run_single = args.single or (
         not args.suite and explicit_single_config_requested(sys.argv[1:])
@@ -220,6 +304,7 @@ def main():
             "epn": args.epn,
             "H": args.H,
             "Hp": args.Hp,
+            "dtype": _DT_LABEL_TO_TORCH[args.dtype],
             "base_B": args.B,
             "nremote": args.nremote,
         }]
@@ -229,20 +314,22 @@ def main():
     if rank == 0:
         print(
             f"MoonEP Prefetch Benchmark (R={R}, warmup={args.warmup}, "
-            f"iters={args.iters})"
+            f"iters={args.iters}, reader=rank0, owner=rank{args.owner_rank}, "
+            f"num_sms={NUM_SMS})"
         )
         print(
-            f"{'Config':<20} {'E':>5} {'B':>4} {'H':>7} {'Hp':>6} "
+            f"{'Config':<20} {'E':>5} {'B':>4} {'dt':>5} {'N':>3} {'H':>7} {'Hp':>6} "
             f"{'SMs':>5} {'Slots':>6} {'Data(MB)':>10} {'Worst(us)':>10} {'BW(GB/s)':>9} {'CommBW':>8}"
         )
-        print("-" * 101)
+        print("-" * 111)
 
     for case in cases:
-        worst_us, mb, bw_gbs, comm_gbs, E, B, slots = bench_case(case, args, rank, R)
+        (worst_us, mb, bw_gbs, comm_gbs, E, B, slots,
+         ntensor, dt_label, H, Hp) = bench_case(case, args, rank, R)
         if rank == 0:
             print(
-                f"{case['label']:<20} {E:>5} {B:>4} "
-                f"{case['H']:>7} {case['Hp']:>6} {NUM_SMS:>5} {slots:>6} "
+                f"{case['label']:<20} {E:>5} {B:>4} {dt_label:>5} {ntensor:>3} "
+                f"{H:>7} {Hp:>6} {NUM_SMS:>5} {slots:>6} "
                 f"{mb:>10.2f} {worst_us:>10.2f} {bw_gbs:>9.2f} {comm_gbs:>8.2f}"
             )
 

moonep/api.py

@@ -68,7 +68,7 @@
 from .dispatch_epilogue import launch_dispatch_epilogue
 from .combine import launch_combine
 from .combine_prologue import launch_combine_prologue
-from .prefetch import launch_prefetch
+from .prefetch import _ELEM_TYPES, launch_prefetch, retile_for_prefetch
 from .grad_reduce import launch_grad_reduce
 
 logger = logging.getLogger(__name__)
@@ -161,6 +161,7 @@ def _launch_full_weight_prefetches(
     full_up_weight: torch.Tensor,
     full_down_weight: torch.Tensor,
     experts_to_copy: torch.Tensor,
+    scales: tuple[torch.Tensor, ...] | None = None,
 ) -> None:
     E = int(ctx['E'])
     num_sms = int(ctx['num_sms'])
@@ -171,6 +172,14 @@ def _launch_full_weight_prefetches(
             experts_to_copy,
             num_sms=num_sms,
         )
+    for full_scale in scales or ():
+        tiled = retile_for_prefetch(full_scale)
+        launch_prefetch(
+            tiled[:E],
+            tiled[E:],
+            experts_to_copy,
+            num_sms=num_sms,
+        )
 
 
 def _launch_full_grad_reduces(
@@ -675,11 +684,13 @@ def _run_prefetch_weight_on_current_stream(
         ctx: dict,
         experts_to_copy: torch.Tensor,
         weight_prefetch_args,
+        scale_prefetch_args=None,
     ) -> None:
         _launch_full_weight_prefetches(
             ctx,
             *weight_prefetch_args,
             experts_to_copy[int(ctx['rank'])],
+            scales=scale_prefetch_args,
         )
 
     def dispatch(
@@ -821,6 +832,9 @@ def prefetch_weight(
         full_gate_weight: torch.Tensor | None = None,
         full_up_weight: torch.Tensor | None = None,
         full_down_weight: torch.Tensor | None = None,
+        full_gate_scale: torch.Tensor | None = None,
+        full_up_scale: torch.Tensor | None = None,
+        full_down_scale: torch.Tensor | None = None,
     ):
         """Prefetch the remote expert weights selected by ``plan`` into the
         local prefetch slots (dispatch fwd, weight side).
@@ -829,9 +843,14 @@ def prefetch_weight(
             plan: MoonEPCommPlan returned by ``dispatch``.
             async_finish: run on the comm stream and return a CUDA event.
             full_gate_weight / full_up_weight / full_down_weight:
-                [E+B, H, H'] bf16 contiguous weight tensors; rows [0, E) are
-                source expert weights, rows [E, E+B) are the prefetch slots
-                filled by this call.
+                [E+B, H, H'] contiguous weight tensors; rows [0, E) are source
+                expert weights, rows [E, E+B) are the prefetch slots filled by
+                this call. bf16 for unquantized experts, uint8 for MXFP4 (e2m1
+                packs two values per byte, so H' is K/2).
+            full_gate_scale / full_up_scale / full_down_scale:
+                optional [E+B, ...] contiguous block-scale tensors, same row
+                convention. Required for quantized experts and omitted for bf16
+                ones.
 
         Returns:
             None in synchronous mode, or the comm-stream CUDA event when
@@ -849,22 +868,38 @@ def prefetch_weight(
         assert all(w is not None for w in weight_prefetch_args), \
             "prefetch_weight tensors must be provided together"
         for w in weight_prefetch_args:
-            assert w.dtype == torch.bfloat16 and w.is_contiguous()
+            assert w.dtype in _ELEM_TYPES, \
+                f"prefetch_weight: unsupported weight dtype {w.dtype}"
+            assert w.is_contiguous()
             assert w.ndim == 3 and int(w.shape[0]) == int(ctx['E']) + int(ctx['B'])
 
+        scale_prefetch_args = (full_gate_scale, full_up_scale, full_down_scale)
+        if any(s is not None for s in scale_prefetch_args):
+            assert all(s is not None for s in scale_prefetch_args), \
+                "prefetch_weight scales must be provided together"
+            for s in scale_prefetch_args:
+                assert s.is_contiguous()
+                assert s.ndim >= 2 and int(s.shape[0]) == int(ctx['E']) + int(ctx['B'])
+        else:
+            scale_prefetch_args = None
+
         if not async_finish:
             self._run_prefetch_weight_on_current_stream(
                 ctx,
                 plan.experts_to_copy,
                 weight_prefetch_args,
+                scale_prefetch_args,
             )
             return None
 
         main_stream = torch.cuda.current_stream()
         comm = self._comm_stream
         assert comm is not None, "MoonEP Buffer communication stream is not initialized"
 
-        self._record_streams((plan.experts_to_copy, *weight_prefetch_args), comm)
+        self._record_streams(
+            (plan.experts_to_copy, *weight_prefetch_args, *(scale_prefetch_args or ())),
+            comm,
+        )
         input_ready = main_stream.record_event()
         comm.wait_event(input_ready)
 
@@ -873,6 +908,7 @@ def prefetch_weight(
                 ctx,
                 plan.experts_to_copy,
                 weight_prefetch_args,
+                scale_prefetch_args,
             )
             done = comm.record_event()
 

moonep/buffer.py

@@ -22,12 +22,6 @@
     nvl_multicast_bind_map,
 )
 
-_ELEM_SIZE = {
-    torch.float32: 4,
-    torch.bfloat16: 2,
-    torch.int32: 4,
-}
-
 # How VMM allocations are shared between the ranks of an EP group. "auto" (the
 # default) picks fabric handles when the group spans more than one node and the
 # device supports them, and POSIX fds otherwise; "fabric" / "fd" force one.
@@ -103,8 +97,7 @@ def pad_dim0_for_alignment(chunk_shape: list[int], dtype: torch.dtype) -> int:
 
     Returns the padded dim0 value (>= chunk_shape[0]).
     """
-    elem_size = _ELEM_SIZE[dtype]
-    inner_size = elem_size
+    inner_size = dtype.itemsize
     for d in chunk_shape[1:]:
         inner_size *= d  # bytes per row