CollectiveX: FP8 precision + low-latency (decode) mode, unified sweep runtime (#2278)

Author: Oseltamivir

.github/workflows/collectivex-sweep.yml

@@ -23,6 +23,10 @@ on:
         description: Keep only shards whose expert-parallel degree is in this comma-list (8 keeps EP8 only, dropping EP16 so GB SKUs co-schedule with 8-GPU SKUs; blank = all)
         type: string
         default: ''
+      modes:
+        description: "Keep only shards in this comma-list of modes (normal, low-latency); low-latency is a decode-only EP8 addition; blank = all"
+        type: string
+        default: ''
 concurrency:
   group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }}
   cancel-in-progress: false
@@ -44,12 +48,14 @@ jobs:
           INPUT_ONLY_SKU: ${{ inputs.only_sku }}
           INPUT_EXCLUDE_SKUS: ${{ inputs.exclude_skus }}
           INPUT_EP_SIZES: ${{ inputs.ep_sizes }}
+          INPUT_MODES: ${{ inputs.modes }}
         run: |
           set -euo pipefail
           args=(--backend "$INPUT_BACKEND")
           [ -n "$INPUT_ONLY_SKU" ] && args+=(--only-sku "$INPUT_ONLY_SKU")
           [ -n "$INPUT_EXCLUDE_SKUS" ] && args+=(--exclude-skus "$INPUT_EXCLUDE_SKUS")
           [ -n "$INPUT_EP_SIZES" ] && args+=(--ep-sizes "$INPUT_EP_SIZES")
+          [ -n "$INPUT_MODES" ] && args+=(--modes "$INPUT_MODES")
           python3 sweep_matrix.py "${args[@]}" --out matrix_full.json >/dev/null
           python3 - "$GITHUB_OUTPUT" <<'PY'
           import json
@@ -94,10 +100,10 @@ jobs:
       COLLX_SHARD_SKU: ${{ matrix.sku }}
       COLLECTIVEX_CANONICAL_GHA: '1'
       COLLECTIVEX_SOURCE_SHA: ${{ github.sha }}
-      # Consolidated shards run one bounded build-group in one Slurm allocation.
-      # MI300X's compute partition has a 180-minute ceiling; the other production
-      # pools accept 300 minutes. Allocations release as soon as the shard finishes.
-      COLLX_TIME: ${{ matrix.sku == 'mi300x' && '180' || '300' }}
+      # Consolidated shards run one bounded build-group in one Slurm allocation;
+      # every production pool accepts 300 minutes. Allocations release as soon as
+      # the shard finishes.
+      COLLX_TIME: '300'
       COLLECTIVEX_EXECUTION_ID: ${{ github.run_id }}_${{ github.run_attempt }}_${{ matrix.id }}
       COLLX_JOB_PARENT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || '/tmp' }}
       COLLX_JOB_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) }}

experimental/CollectiveX/README.md

@@ -12,11 +12,21 @@ responsibility. The full measurement methodology is in [docs/methodology.md](doc
 ## Execution Profile
 
 The workload uses packed placement and one pinned `fixed-profile` resource configuration per
-backend/topology; there is no tuning sweep. Dispatch and combine are fixed BF16 on every backend;
-precision is not a swept dimension. Every case runs the single normal-mode contract:
-
-- Normal mode uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only
-  combine. Coverage is uniform routing only.
+backend/topology; there is no tuning sweep. Combine is always BF16; dispatch precision is a swept
+dimension — a BF16 control plus a caller-prequantized FP8 dispatch on every backend. Coverage is
+uniform routing only. Cases run in one of two modes:
+
+- `normal` uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only,
+  unweighted rank-sum combine. It runs the full decode and prefill ladders.
+- `low-latency` uses each backend's decode-optimized kernel family: on DeepEP the legacy
+  `deep_ep.Buffer` IBGDA `low_latency_dispatch`/`low_latency_combine` (a per-expert padded receive
+  and a source-side gate-weighted combine); on MoRI the `IntraNodeLL` kernel (single-call,
+  pure-intranode, same compact layout and unweighted rank-sum combine as `IntraNode`). It is a
+  decode-phase-only, per-SKU-capability-gated addition whose runnable set differs from `normal`'s, so
+  it is enabled from each SKU's `ll_backends` registry entry (currently DeepEP V2 EP8 on H100/H200/B200
+  and MoRI EP8 on MI300X/MI325X/MI355X). Scoped single-node EP8 runs over the intra-node NVLink/XGMI
+  low-latency path (no `/dev/gdrdrv` needed — validated on H200 with it absent); NVSHMEM/IBGDA on the
+  wire is only a multi-node scale-out (EP16) concern.
 
 Cases use a fixed timing profile from `configs/sweep.json`: 256 trials x 8 timed iterations (2048
 samples per component) with 32 synchronized full roundtrip warmups before each measured component at
@@ -29,13 +39,16 @@ Correctness is checked against an implementation-independent oracle that reprodu
 two-level reduction — intra-scale-up-domain FP32, then a BF16 cast of each domain's partial for the
 scale-out send. The combine gate is a tight max elementwise relative error below `8 * 2^-8`
 (denominator clamped at 0.02), which holds across scale-up and multi-node scale-out topologies
-alike. Any failed rank or point makes the case ineligible in the result it writes.
+alike. Under FP8 dispatch the oracle applies the same per-token cast round-trip to its semantic
+payload, so the dispatched-payload compare stays bit-exact and the combine gate is unchanged — the
+quantization is modeled, not tolerated. Any failed rank or point makes the case ineligible in the
+result it writes.
 
 The matrix covers H100, H200, B200, B300, GB200, GB300, MI300X, MI325X, and MI355X. `sweep_matrix.py` materializes
 the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls
-and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`, and
-`--ep-sizes` select a subset; the matrix is generated per dispatch, with no frozen digest or locked
-case count.
+and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`,
+`--ep-sizes`, and `--precisions` select a subset; the matrix is generated per dispatch, with no
+frozen digest or locked case count.
 
 | Systems | EP8 | EP16 |
 |---|---|---|
@@ -48,8 +61,8 @@ scale-up domain.
 
 | Backend | Current scope |
 |---|---|
-| DeepEP V2 | PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out |
-| MoRI | AMD EP8 uses IntraNode-family kernels (MI355X IntraNode, MI300X/MI325X asyncLL); EP16 pins InterNodeV1 over 2x8 XGMI + RDMA |
+| DeepEP V2 | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode/EP8 only |
+| MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950); combine stays BF16 (`quant_type=none`) alongside BF16 dispatch |
 
 DeepEP V2 means the `ElasticBuffer` implementation introduced by
 [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build.

experimental/CollectiveX/bench/ep_backend.py

@@ -61,18 +61,31 @@ class EPBackend(abc.ABC):
     Subclasses implement the transport (create_buffer, dispatch, stage,
     combine, recv_tokens, inspect_dispatch, combine_transformed);
     everything the driver and the oracles need beyond that is provided here.
-    Dispatch and combine are fixed BF16, so no adapter selects a precision codec.
+    Combine is always BF16; an adapter that supports FP8 dispatch overrides
+    SUPPORTED_PRECISIONS and the semantic_payload/_encode_dispatch hooks.
     """
 
     name: str = ""
     SUPPORTED_MODES: tuple = ("normal",)
+    # Dispatch precisions the adapter realizes. BF16 is the universal control; an
+    # adapter that also sends an FP8-quantized dispatch payload widens this.
+    SUPPORTED_PRECISIONS: tuple = ("bf16",)
     stage_device_work = False
     combine_needs_redispatch = False
     dispatch_needs_combine_cleanup = False
     # Adapters that reduce activations and top-k weights independently must carry
     # the complete local weighted expert sum in the activation tensor.
     combine_weight_semantics = "unweighted-rank-sum"
     roundtrip_only = False
+    # Realized wire formats recorded in the artifact. Combine is always BF16;
+    # dispatch_dtype is overridden per-run by an FP8 adapter (e.g. "fp8-e4m3fn").
+    dispatch_dtype = "bf16"
+    combine_dtype = "bf16"
+    # Logical byte model for one dispatched copy: bytes per activation value and
+    # per-copy scale bytes. BF16 moves 2 bytes/value with no scale payload; an FP8
+    # adapter sends 1 byte/value plus (for a blockwise codec) per-block FP32 scales.
+    dispatch_value_bytes = 2
+    dispatch_scale_bytes_per_copy = 0
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
@@ -90,6 +103,11 @@ def __init__(self, args, rank, world_size, local_rank, device):
         self.mode = args.mode
         if self.mode not in self.SUPPORTED_MODES:
             raise ValueError(f"{self.name} does not support mode {self.mode!r}")
+        self.precision = args.precision
+        if self.precision not in self.SUPPORTED_PRECISIONS:
+            raise ValueError(
+                f"{self.name} does not support precision {self.precision!r}"
+            )
 
     # ---- Abstract transport contract -------------------------------------------------
 
@@ -179,17 +197,44 @@ def _build_rank_inputs(self, args, tokens_per_rank) -> RankInputs:
             global_weights=w_g,
         )
 
+    def semantic_payload(self, x):
+        """The BF16 values the oracle should expect for a dispatched payload.
+
+        Identity for a backend that sends x unchanged. An FP8 backend overrides this
+        to apply the exact quant->dequant round-trip the kernel transports, so the
+        dispatched-payload compare stays bit-exact and the combine gate stays tight.
+        """
+        return x
+
+    def _encode_dispatch(self, x):
+        """Return (dispatch_payload, oracle_semantic) for the source activations x.
+
+        Base identity: send x, no separate oracle payload (BF16). An FP8 adapter
+        returns the caller-prequantized dispatch payload and the dequantized BF16 the
+        oracle must expect after the backend's own dequant.
+        """
+        return x, None
+
     def make_problem(self, T, idx, weights, x):
-        """Assemble the per-shape problem namespace (BF16 dispatch sends x directly)."""
+        """Assemble the per-shape problem namespace.
+
+        dispatch_x is the payload actually sent (x itself in BF16; the caller-
+        prequantized encoding under FP8). oracle_x, when set, is the dequantized BF16
+        the combine oracle must expect, so the tight gate needs no tolerance change.
+        """
         import torch
 
-        return types.SimpleNamespace(
+        dispatch_x, oracle_semantic = self._encode_dispatch(x)
+        problem = types.SimpleNamespace(
             T=T,
             x=x,
-            dispatch_x=x,
+            dispatch_x=dispatch_x,
             topk_idx=idx.to(self._topk_idx_dtype()),
             topk_weights=weights.to(torch.float32),
         )
+        if oracle_semantic is not None:
+            problem.oracle_x = oracle_semantic
+        return problem
 
     def _topk_idx_dtype(self):
         """Integer dtype the backend's kernels expect for top-k routing indices."""