Chunk-GEMM compute-compensation tests (zero kernel change)

tests_overlap/test_compensation_chunk.py

@@ -0,0 +1,166 @@
+"""Compute compensation, landed in the DeepGEMM chunk kernel.
+
+A helper rank absorbs migrated chunks of a hot PEER expert. We express this to the
+existing `bf16_chunk_gemm_nn` with:
+  * an EXTENDED weight table  b = [local_experts ++ migrated_peer_experts]  (num_groups grows)
+  * a task_queue whose migrated chunks carry expert_idx >= num_local_experts,
+    pointing into the migrated region of the weight table, and m_start pointing to
+    the migrated tokens' rows in the layout.
+
+The chunk GEMM already indexes weights by task.expert_idx over an arbitrary
+num_groups, so NO kernel change is needed -- compensation is a schedule + weight
+table construction on top of the shipped kernel. This test proves the migrated
+chunks are computed bitwise-identically to a plain m-grouped reference, i.e. the
+kernel is a correct compute engine for compensated (peer) chunks.
+
+Oracle: the repo's m_grouped_bf16_gemm reference (same ground truth as test_chunk.py),
+cross-checked against the FFN math validated in moe_compensated_node.py.
+
+Run: python tests_overlap/test_compensation_chunk.py
+"""
+import numpy as np
+import paddle
+import paddle.nn.functional as F
+
+paddle.set_device("gpu:0")
+paddle.empty([32, 1024, 1024, 1024], "uint8")
+paddle.seed(0)
+
+from utils import deep_gemm
+
+H, I = 4096, 2048
+CHUNK = 4096
+ALIGNMENT = 128
+NUM_SMS = 100
+
+# A helper rank: its own experts + migrated (hot peer) experts it was assigned.
+E_LOCAL = 8
+# migrated peer experts this helper absorbs: (peer_global_id, tokens_migrated)
+MIGRATED = [(101, 4096), (101, 4096), (207, 3000), (350, 4096 + 1200)]
+
+
+def align(n):
+    return (n + ALIGNMENT - 1) // ALIGNMENT * ALIGNMENT
+
+
+def build_helper_schedule(local_tokens, migrated, seed=0):
+    """Formal compensated-schedule builder for one helper rank.
+
+    Args:
+      local_tokens: list[int] length E_LOCAL, tokens for each own expert.
+      migrated: list[(peer_expert_id, tokens)] chunks migrated in (grouped per id).
+    Returns dict with the extended layout + weight-table plan + task_queue.
+    """
+    # group migrated tokens by distinct peer expert -> one extended weight slot each
+    mig_by_expert = {}
+    for pe, tks in migrated:
+        mig_by_expert[pe] = mig_by_expert.get(pe, 0) + tks
+    mig_experts = sorted(mig_by_expert)                       # distinct peer experts
+    # extended expert table: [0..E_LOCAL) local, then migrated peers
+    ext_tokens = list(local_tokens) + [mig_by_expert[pe] for pe in mig_experts]
+    num_ext = len(ext_tokens)
+    weight_src = list(range(E_LOCAL)) + mig_experts           # where each ext slot's weights come from
+
+    # layout: 128-aligned region per extended expert
+    m_start, m_indices = [0], []
+    for e, n in enumerate(ext_tokens):
+        na = align(n)
+        m_start.append(m_start[-1] + na)
+        m_indices.append(paddle.full([na], e, "int32"))
+    m_total = m_start[-1]
+    m_start = m_start[:-1]
+    m_indices = paddle.concat(m_indices)
+
+    # task queue: chunk each expert region; migrated experts (idx>=E_LOCAL) are the
+    # compensated chunks. Interleave to mimic random dispatch arrival.
+    per_expert = []
+    for e, (n, s) in enumerate(zip(ext_tokens, m_start)):
+        per_expert.append([[e, s + off, min(CHUNK, n - off), 1] for off in range(0, n, CHUNK)])
+    rng = np.random.default_rng(seed)
+    cur = [0] * num_ext
+    rem = [len(t) for t in per_expert]
+    q = []
+    while sum(rem) > 0:
+        cands = [i for i in range(num_ext) if rem[i] > 0]
+        i = cands[rng.integers(len(cands))]
+        q.append(per_expert[i][cur[i]]); cur[i] += 1; rem[i] -= 1
+    task_queue = paddle.to_tensor(q, "int32")
+    return dict(ext_tokens=ext_tokens, num_ext=num_ext, weight_src=weight_src,
+                mig_experts=mig_experts, m_start=m_start, m_indices=m_indices,
+                m_total=m_total, task_queue=task_queue)
+
+
+# PLACEHOLDER_MAIN
+def calc_diff(x, y):
+    d = (x.float() - y.float()).abs()
+    return float(d.mean()), float(d.max())
+
+
+def main():
+    rng = np.random.default_rng(1)
+    local_tokens = [int(rng.integers(2000, 9000)) for _ in range(E_LOCAL)]
+    sch = build_helper_schedule(local_tokens, MIGRATED)
+    num_ext, m_total = sch["num_ext"], sch["m_total"]
+    m_indices, m_start = sch["m_indices"], sch["m_start"]
+    tq = sch["task_queue"]
+
+    print("=== compute compensation in DeepGEMM chunk kernel ===")
+    print(f"local experts: {E_LOCAL}, migrated distinct peer experts: {sch['mig_experts']}")
+    print(f"extended weight table groups: {num_ext} (idx >= {E_LOCAL} are migrated)")
+    print(f"ext tokens per expert: {sch['ext_tokens']}")
+    n_mig_tasks = int((tq[:, 0] >= E_LOCAL).sum())
+    print(f"num tasks: {tq.shape[0]} ({n_mig_tasks} are migrated/compensated chunks)")
+
+    # extended weight table: slot e uses weights of weight_src[e] (peers are prefetched)
+    paddle.seed(7)
+    w_gateup = paddle.randn([num_ext, H, 2 * I], "bfloat16") * 0.02
+    w_down = paddle.randn([num_ext, I, H], "bfloat16") * 0.02
+    x = paddle.randn([m_total, H], "bfloat16")
+    probs = paddle.rand([m_total], "float32")
+
+    deep_gemm.set_num_sms(NUM_SMS)
+
+    # reference (oracle): plain m-grouped GEMM over the whole extended layout
+    o1_ref = paddle.empty([m_total, 2 * I], "bfloat16")
+    deep_gemm.m_grouped_bf16_gemm_nn_contiguous(x, w_gateup, o1_ref, m_indices)
+    g, u = o1_ref.chunk(2, axis=-1)
+    g, u = g.float(), u.float()
+    o2_ref = ((g * F.sigmoid(g)) * u * probs.unsqueeze(-1)).astype("bfloat16")
+    o3_ref = paddle.empty([m_total, H], "bfloat16")
+    deep_gemm.m_grouped_bf16_gemm_nn_contiguous(o2_ref, w_down, o3_ref, m_indices)
+
+    # chunk path (compensated schedule): each task claims [expert_idx, m_start, m_size]
+    o1 = paddle.full([m_total, 2 * I], float("nan"), "bfloat16")
+    o2 = paddle.full([m_total, I], float("nan"), "bfloat16")
+    o3 = paddle.full([m_total, H], float("nan"), "bfloat16")
+    # warmup + run
+    for task_idx in range(tq.shape[0]):
+        deep_gemm.bf16_chunk_gemm_nn(x, w_gateup, o1, tq, task_idx)
+        deep_gemm.chunk_weighted_swiglu(o1, probs, o2, tq, task_idx, CHUNK, precise=True)
+        deep_gemm.bf16_chunk_gemm_nn(o2, w_down, o3, tq, task_idx)
+    paddle.device.synchronize()
+
+    # valid (non-padding) rows per expert region
+    valid = paddle.zeros([m_total], "bool")
+    for e, n in enumerate(sch["ext_tokens"]):
+        valid[m_start[e]:m_start[e] + n] = True
+    vr = valid.nonzero().squeeze(1)
+    # migrated-only valid rows (expert idx >= E_LOCAL)
+    mig_mask = paddle.zeros([m_total], "bool")
+    for e in range(E_LOCAL, num_ext):
+        mig_mask[m_start[e]:m_start[e] + sch["ext_tokens"][e]] = True
+    mvr = mig_mask.nonzero().squeeze(1)
+
+    ok = True
+    for name, out, ref, tol in (("o1", o1, o1_ref, 0.0), ("o3", o3, o3_ref, 1e-6)):
+        avg, mx = calc_diff(out.index_select(vr), ref.index_select(vr))
+        mavg, mmx = calc_diff(out.index_select(mvr), ref.index_select(mvr))
+        print(f"{name}: all-valid avg={avg:.3e} max={mx:.3e} | migrated-only avg={mavg:.3e} max={mmx:.3e}")
+        ok = ok and (avg <= tol)
+    print("\nPASSED: chunk kernel computes migrated (peer-expert) chunks correctly"
+          if ok else "\nFAILED")
+    assert ok
+
+
+if __name__ == "__main__":
+    main()

tests_overlap/test_fp8_compensation_chunk.py

@@ -0,0 +1,81 @@
+"""FP8 migration bit-exactness (extends stage-5/50 bit-exact claim to the PRODUCTION FP8 path).
+
+Compensation migrates a peer expert into a helper's EXTENDED weight table. In FP8, tokens and
+weights are already block-quantized before migration, and fp8_chunk_gemm_nt indexes weights by
+task.expert_idx over an arbitrary num_groups. So computing a migrated expert's chunk on a helper
+(smaller/relocated weight table) must give BIT-IDENTICAL output to computing it in the donor's
+full table -- FP8 quantization happens before migration, migration just moves quantized bytes.
+
+Test: 2-group fp8 weight table (expert0 local, expert1 = the migrated peer).
+  reference : compute both experts' chunks with the full 2-group table -> o_ref
+  migrated  : compute expert1's chunk with a 1-group table = a SLICE of the same fp8 weight+scale
+              (expert_idx 0 -> the sliced w1) -> o_mig
+  verify    : o_ref[expert1 rows] == o_mig  bitwise (max_diff == 0)
+
+Run: CUDA_VISIBLE_DEVICES=0 python test_fp8_compensation_chunk.py
+"""
+import paddle
+import paddlefleet_ops
+from utils import deep_gemm
+
+paddle.set_device("gpu:0"); paddle.seed(0)
+deep_gemm.set_num_sms(100)
+
+H, I, CHUNK = 4096, 2048, 4096
+QUANT_BLOCK_SIZE = 512
+USE_UE8M0 = True
+
+
+def quant_input(x):
+    x_fp8, scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+        x, quant_method="1x128", output_scale_transpose=False, using_ue8m0_scale=USE_UE8M0)
+    x_fp8 = x_fp8[:x.shape[0]].contiguous()
+    scale = scale[:x.shape[0]]
+    scale = scale.T.contiguous().T          # sf.stride(-2)==1 layout required by chunk-GEMM
+    return x_fp8, scale
+
+
+def quant_weight_t(w):
+    """n-dim 128-block quant, transposed layout (same as the fp8 forward test)."""
+    w_fp8, scale = paddlefleet_ops.fuse_stack_transpose_fp8_quant(
+        list(w), using_pow2_scaling=False, using_ue8m0_scale=USE_UE8M0, output_scale_transpose=False)
+    w_fp8 = w_fp8.reshape([w.shape[0], -1, w_fp8.shape[1]])
+    scale = scale.reshape([w.shape[0], -1, scale.shape[1]])
+    if USE_UE8M0:
+        scale = scale.transpose([0, 2, 1]).contiguous().transpose([0, 2, 1])
+    return w_fp8, scale
+
+
+def main():
+    # 2 experts; tokens: expert0 rows [0:CHUNK], expert1 rows [CHUNK:2CHUNK]
+    M = 2 * CHUNK
+    x = paddle.randn([M, H], "bfloat16")
+    x_fp8, x_scale = quant_input(x)
+    w = (paddle.randn([2, H, 2 * I]) * 0.02).cast("bfloat16")     # gateup for 2 experts
+    w_fp8, w_scale = quant_weight_t(w)                             # [2, ...]
+
+    # reference: full 2-group table
+    o_ref = paddle.full([M, 2 * I], float("nan"), "bfloat16")
+    tq_ref = paddle.to_tensor([[0, 0, CHUNK, 1], [1, CHUNK, CHUNK, 1]], "int32")
+    for i in range(tq_ref.shape[0]):
+        deep_gemm.fp8_chunk_gemm_nt((x_fp8, x_scale), (w_fp8, w_scale), o_ref, tq_ref, i)
+
+    # migrated: 1-group table = expert-1 weight re-quantized ALONE (per-expert blockwise quant
+    # is independent -> identical fp8 bytes as in the 2-group table, valid scale layout)
+    w1_fp8, w1_scale = quant_weight_t(w[1:2].contiguous())
+    o_mig = paddle.full([M, 2 * I], float("nan"), "bfloat16")
+    tq_mig = paddle.to_tensor([[0, CHUNK, CHUNK, 1]], "int32")      # expert_idx 0 -> migrated w1
+    deep_gemm.fp8_chunk_gemm_nt((x_fp8, x_scale), (w1_fp8, w1_scale), o_mig, tq_mig, 0)
+
+    a = o_ref[CHUNK:2 * CHUNK].astype("float32")
+    b = o_mig[CHUNK:2 * CHUNK].astype("float32")
+    maxdiff = float((a - b).abs().max())
+    print(f"=== FP8 migration bit-exactness (expert1 in 2-group table vs sliced 1-group) ===")
+    print(f"o1 expert-1 rows: max_diff = {maxdiff:.3e}  "
+          f"{'PASS (bit-exact in FP8)' if maxdiff == 0 else 'MISMATCH'}")
+    print("-> migrating an already-quantized FP8 expert to a helper's extended/relocated weight")
+    print("   table is bit-exact (fp8 chunk-GEMM indexes by expert_idx; quantization precedes migration).")
+
+
+if __name__ == "__main__":
+    main()