[issue9] KV Cache 通信组件多表 destination-major 字节打包参考实现

src/code/issue9/README.md

@@ -0,0 +1,22 @@
+# Issue 9 单机可完成部分：多表 AlltoAllv 聚合参考实现
+
+`fused_alltoallv_reference.py` 把各表数据按“destination-major、table-minor”顺序
+打成一个字节 buffer，使每个目标 rank 只需要一个连续 segment。接收端根据每表
+dtype、尾部 shape 和 split 元数据恢复结果。
+
+本机模拟覆盖：
+
+- 多张表 first-dim 不同；
+- 表间 hidden shape 和 dtype 不同；
+- split 中存在零长度消息；
+- 接收顺序与逐表 AlltoAllv 一致，并逐位相等。
+
+```bash
+python -m pytest src/code/issue9/test_fused_alltoallv_reference.py
+python src/code/issue9/fused_alltoallv_reference.py
+```
+
+默认演示把 8 次逐表 launch 合成 1 次，模型控制面开销下降 87.5%。这只是 CPU
+参考实现与启动开销模型；下一阶段需要把同一布局接到 NCCL/DeepEP 数据面，在多机多卡
+实测 launch 数、协商开销与端到端时间。
+

src/code/issue9/fused_alltoallv_reference.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""Byte-exact reference for fusing multiple table-wise AlltoAllv calls."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class TableDescriptor:
+    dtype: str
+    trailing_shape: tuple[int, ...]
+    splits: tuple[int, ...]
+    row_bytes: int
+
+
+@dataclass
+class PackedRank:
+    buffer: np.ndarray
+    destination_offsets: tuple[int, ...]
+    tables: tuple[TableDescriptor, ...]
+
+    @property
+    def world_size(self) -> int:
+        return len(self.destination_offsets) - 1
+
+    @property
+    def destination_sizes(self) -> tuple[int, ...]:
+        return tuple(
+            self.destination_offsets[index + 1] - self.destination_offsets[index]
+            for index in range(self.world_size)
+        )
+
+
+def _descriptor(table: np.ndarray, splits: Sequence[int], world_size: int) -> TableDescriptor:
+    if table.ndim < 1:
+        raise ValueError("each table must have at least one dimension")
+    if len(splits) != world_size or any(split < 0 for split in splits):
+        raise ValueError("each split vector must contain world_size non-negative entries")
+    if sum(splits) != table.shape[0]:
+        raise ValueError("split sizes must sum to the table's first dimension")
+    row_elements = int(np.prod(table.shape[1:], dtype=np.int64)) if table.ndim > 1 else 1
+    return TableDescriptor(
+        dtype=table.dtype.str,
+        trailing_shape=tuple(table.shape[1:]),
+        splits=tuple(int(split) for split in splits),
+        row_bytes=row_elements * table.dtype.itemsize,
+    )
+
+
+def pack_rank(
+    tables: Sequence[np.ndarray], splits_per_table: Sequence[Sequence[int]], world_size: int
+) -> PackedRank:
+    """Pack one rank's tables in destination-major, table-minor order."""
+
+    if world_size <= 0 or not tables or len(tables) != len(splits_per_table):
+        raise ValueError("world_size must be positive and every table needs a split vector")
+    contiguous = [np.ascontiguousarray(table) for table in tables]
+    descriptors = [
+        _descriptor(table, splits, world_size)
+        for table, splits in zip(contiguous, splits_per_table, strict=True)
+    ]
+    table_starts = [np.cumsum((0,) + descriptor.splits) for descriptor in descriptors]
+
+    chunks: list[np.ndarray] = []
+    destination_offsets = [0]
+    for destination in range(world_size):
+        for table, descriptor, starts in zip(
+            contiguous, descriptors, table_starts, strict=True
+        ):
+            row_start = int(starts[destination])
+            row_end = int(starts[destination + 1])
+            chunk = table[row_start:row_end].reshape(-1).view(np.uint8)
+            chunks.append(chunk)
+        destination_offsets.append(sum(chunk.nbytes for chunk in chunks))
+
+    buffer = np.concatenate(chunks) if chunks else np.empty(0, dtype=np.uint8)
+    return PackedRank(buffer, tuple(destination_offsets), tuple(descriptors))
+
+
+def _validate_compatible(packed_ranks: Sequence[PackedRank]) -> tuple[int, int]:
+    if not packed_ranks:
+        raise ValueError("at least one packed rank is required")
+    world_size = len(packed_ranks)
+    if any(packed.world_size != world_size for packed in packed_ranks):
+        raise ValueError("packed rank count must equal every split vector's world size")
+    num_tables = len(packed_ranks[0].tables)
+    for packed in packed_ranks[1:]:
+        if len(packed.tables) != num_tables:
+            raise ValueError("all ranks must contain the same number of tables")
+        for reference, candidate in zip(packed_ranks[0].tables, packed.tables, strict=True):
+            if (reference.dtype, reference.trailing_shape) != (
+                candidate.dtype,
+                candidate.trailing_shape,
+            ):
+                raise ValueError("matching tables must use the same dtype and trailing shape")
+    return world_size, num_tables
+
+
+def simulate_fused_alltoallv(packed_ranks: Sequence[PackedRank]) -> list[list[np.ndarray]]:
+    """Simulate one fused byte-buffer AlltoAllv and unpack its table results."""
+
+    world_size, num_tables = _validate_compatible(packed_ranks)
+    outputs: list[list[np.ndarray]] = []
+    for destination in range(world_size):
+        destination_tables: list[np.ndarray] = []
+        for table_index in range(num_tables):
+            chunks = []
+            for source in packed_ranks:
+                descriptor = source.tables[table_index]
+                segment_start = source.destination_offsets[destination]
+                table_offset = sum(
+                    source.tables[index].splits[destination]
+                    * source.tables[index].row_bytes
+                    for index in range(table_index)
+                )
+                rows = descriptor.splits[destination]
+                byte_count = rows * descriptor.row_bytes
+                raw = source.buffer[
+                    segment_start + table_offset : segment_start + table_offset + byte_count
+                ]
+                shape = (rows,) + descriptor.trailing_shape
+                chunks.append(np.frombuffer(raw, dtype=np.dtype(descriptor.dtype)).reshape(shape).copy())
+            descriptor = packed_ranks[0].tables[table_index]
+            if chunks:
+                destination_tables.append(np.concatenate(chunks, axis=0))
+            else:
+                destination_tables.append(
+                    np.empty((0,) + descriptor.trailing_shape, dtype=np.dtype(descriptor.dtype))
+                )
+        outputs.append(destination_tables)
+    return outputs
+
+
+def simulate_separate_alltoallv(
+    rank_tables: Sequence[Sequence[np.ndarray]],
+    rank_splits: Sequence[Sequence[Sequence[int]]],
+) -> list[list[np.ndarray]]:
+    """Reference semantics for launching one AlltoAllv per table."""
+
+    world_size = len(rank_tables)
+    packed = [
+        pack_rank(tables, splits, world_size)
+        for tables, splits in zip(rank_tables, rank_splits, strict=True)
+    ]
+    _, num_tables = _validate_compatible(packed)
+    outputs: list[list[np.ndarray]] = [[] for _ in range(world_size)]
+    for destination in range(world_size):
+        for table_index in range(num_tables):
+            chunks = []
+            for source_rank in range(world_size):
+                table = np.ascontiguousarray(rank_tables[source_rank][table_index])
+                splits = rank_splits[source_rank][table_index]
+                start = sum(splits[:destination])
+                end = start + splits[destination]
+                chunks.append(table[start:end])
+            outputs[destination].append(np.concatenate(chunks, axis=0))
+    return outputs
+
+
+def demo_report(world_size: int, num_tables: int, launch_overhead_us: float) -> dict[str, Any]:
+    if world_size <= 0 or num_tables <= 0 or launch_overhead_us < 0:
+        raise ValueError("invalid demo parameters")
+    rng = np.random.default_rng(20260726)
+    rank_tables = []
+    rank_splits = []
+    for rank in range(world_size):
+        tables = []
+        splits_for_rank = []
+        for table_index in range(num_tables):
+            splits = rng.integers(0, 6, size=world_size, dtype=np.int64)
+            if splits.sum() == 0:
+                splits[rank] = 1
+            rows = int(splits.sum())
+            hidden = 4 * (table_index + 1)
+            dtype = np.float16 if table_index % 2 == 0 else np.float32
+            values = np.arange(rows * hidden, dtype=dtype).reshape(rows, hidden) + rank * 100
+            tables.append(values)
+            splits_for_rank.append(splits.tolist())
+        rank_tables.append(tables)
+        rank_splits.append(splits_for_rank)
+
+    packed = [
+        pack_rank(tables, splits, world_size)
+        for tables, splits in zip(rank_tables, rank_splits, strict=True)
+    ]
+    fused = simulate_fused_alltoallv(packed)
+    separate = simulate_separate_alltoallv(rank_tables, rank_splits)
+    bitwise_equal = all(
+        np.array_equal(fused[destination][table], separate[destination][table])
+        for destination in range(world_size)
+        for table in range(num_tables)
+    )
+    payload_bytes = sum(item.buffer.nbytes for item in packed)
+    separate_control_us = num_tables * launch_overhead_us
+    fused_control_us = launch_overhead_us
+    return {
+        "schema_version": 1,
+        "world_size": world_size,
+        "num_tables": num_tables,
+        "bitwise_equal_to_separate": bitwise_equal,
+        "payload_bytes": payload_bytes,
+        "per_rank_destination_bytes": [list(item.destination_sizes) for item in packed],
+        "launches": {"separate": num_tables, "fused": 1},
+        "modeled_control_us": {"separate": separate_control_us, "fused": fused_control_us},
+        "modeled_control_saving_ratio": (
+            0.0 if separate_control_us == 0 else 1.0 - fused_control_us / separate_control_us
+        ),
+        "warning": "control-time saving is a launch-overhead model, not a multi-rank measurement",
+    }
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Validate multi-table AlltoAllv packing")
+    parser.add_argument("--world-size", type=int, default=4)
+    parser.add_argument("--num-tables", type=int, default=8)
+    parser.add_argument("--launch-overhead-us", type=float, default=8.0)
+    parser.add_argument(
+        "--output", type=Path, default=Path("src/code/issue9/results/fusion_reference.json")
+    )
+    return parser
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+    report = demo_report(args.world_size, args.num_tables, args.launch_overhead_us)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    print(json.dumps(report, ensure_ascii=False, indent=2))
+    print(f"Wrote {args.output}")
+    return 0 if report["bitwise_equal_to_separate"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/code/issue9/results/fusion_reference.json

@@ -0,0 +1,43 @@
+{
+  "schema_version": 1,
+  "world_size": 4,
+  "num_tables": 8,
+  "bitwise_equal_to_separate": true,
+  "payload_bytes": 19160,
+  "per_rank_destination_bytes": [
+    [
+      808,
+      1200,
+      1264,
+      1368
+    ],
+    [
+      1072,
+      1296,
+      1392,
+      912
+    ],
+    [
+      1232,
+      1304,
+      1032,
+      1192
+    ],
+    [
+      1040,
+      1560,
+      1248,
+      1240
+    ]
+  ],
+  "launches": {
+    "separate": 8,
+    "fused": 1
+  },
+  "modeled_control_us": {
+    "separate": 64.0,
+    "fused": 8.0
+  },
+  "modeled_control_saving_ratio": 0.875,
+  "warning": "control-time saving is a launch-overhead model, not a multi-rank measurement"
+}

src/code/issue9/test_fused_alltoallv_reference.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from fused_alltoallv_reference import (
+    demo_report,
+    pack_rank,
+    simulate_fused_alltoallv,
+    simulate_separate_alltoallv,
+)
+
+
+def test_fused_matches_separate_with_mixed_shapes_dtypes_and_zero_splits():
+    rank_tables = [
+        [
+            np.arange(12, dtype=np.float32).reshape(3, 4),
+            np.arange(10, dtype=np.int16).reshape(5, 2),
+        ],
+        [
+            (np.arange(16, dtype=np.float32) + 100).reshape(4, 4),
+            (np.arange(4, dtype=np.int16) + 100).reshape(2, 2),
+        ],
+    ]
+    rank_splits = [
+        [[1, 2], [0, 5]],
+        [[3, 1], [2, 0]],
+    ]
+    packed = [pack_rank(rank_tables[index], rank_splits[index], 2) for index in range(2)]
+
+    fused = simulate_fused_alltoallv(packed)
+    separate = simulate_separate_alltoallv(rank_tables, rank_splits)
+
+    for destination in range(2):
+        for table in range(2):
+            np.testing.assert_array_equal(fused[destination][table], separate[destination][table])
+
+
+def test_demo_reduces_launch_count_and_preserves_bits():
+    report = demo_report(world_size=4, num_tables=8, launch_overhead_us=8)
+    assert report["bitwise_equal_to_separate"] is True
+    assert report["launches"] == {"separate": 8, "fused": 1}
+    assert report["modeled_control_saving_ratio"] == 0.875
+