Add fused-MoE patches for custom image builds

patches/patch_sparkinfer_projection_tiers.py

@@ -0,0 +1,383 @@
+#!/usr/bin/env python3
+"""Let SparkInfer's fused mixed-Trellis MoE tier per (expert, projection).
+
+Target: sparkinfer from r28 (si200c1db), ABI v6.
+Idempotent, exact anchors, fails closed.
+
+WHY
+    R7 chooses K independently for every (expert, projection). The fused kernel
+    tiers per expert, so an expert that is K3 for gate and K4 for up cannot be
+    expressed and the whole checkpoint falls back to exllamav3's exl3_moe --
+    which measures 27% of HBM bandwidth and is 50% of a decode step.
+
+    Counting the checkpoint per layer, though:
+
+        distinct FC1 K values per layer (|Kg u Ku|):  2 -> 73 layers, 3 -> 2
+        distinct FC2 K values per layer (|Kd|):       2 -> 71 layers, 3 -> 4
+
+    Two bitrates, {3,4}, on both phases in nearly every layer. So tier0 = K3 and
+    tier1 = K4 already serves gate, up AND down; no per-phase trellis_bits is
+    needed. The only inexpressible thing is per-projection MEMBERSHIP, and that
+    is a runtime table:
+
+        descriptor = descriptor_map[combined_expert]   # int32[total_experts]
+        tier       = descriptor >> 8
+        local      = descriptor & 0xFF
+
+WHAT
+    descriptor_map grows from int32[total_experts] to int32[3 * total_experts],
+    rows ordered gate, up, down, same encoding. The dispatch picks a row:
+    FC2 is const_expr; FC1 branches on output_n_tile, which the tile already
+    holds and which trellis3_t256_proj guarantees is projection-aligned (it
+    rejects any geometry where a projection is not an integral number of CTA N
+    tiles). Row 0 is byte-compatible with the old single-row table, so a caller
+    that passes three identical rows reproduces current behaviour exactly.
+
+    No weight-stack surgery. local_expert is now per projection and the gate/up
+    half offset already rides in output_n_tile, so gate at tier-local slot i and
+    up at slot j address as i*expert_stride + gate_half and j*expert_stride +
+    up_half inside the existing w13 stack. Strides, scales, rotations and the
+    route packer are untouched; the stack sizes to max(gate_count, up_count).
+
+CORRECTNESS
+    For any layer where gate and up share a tier assignment, all three rows are
+    identical and every launch resolves exactly as before -- that is the
+    bit-identical gate. Reduction order, fragment layout and accumulation are
+    untouched: this only changes which tier's B pointer a tile reads.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+MARKER = "glm52-r7-projtiers"
+
+MIXED = "sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py"
+
+
+def replace_once(text: str, old: str, new: str, label: str) -> str:
+    count = text.count(old)
+    if count != 1:
+        raise RuntimeError(f"{label}: expected one exact anchor, found {count}")
+    return text.replace(old, new, 1)
+
+
+def patch_mixed(text: str) -> str:
+    # 1. Device tensor: three rows instead of one.
+    text = replace_once(
+        text,
+        """        descriptor_map = cute.make_tensor(
+            descriptor_map_ptr,
+            layout=cute.make_layout((total_experts,), stride=(1,)),
+        )
+""",
+        """        # glm52-r7-projtiers: rows are gate, up, down. Row 0 alone is the
+        # historical layout, so three identical rows reproduce it exactly.
+        descriptor_map = cute.make_tensor(
+            descriptor_map_ptr,
+            layout=cute.make_layout(
+                (cutlass.Int64(3) * total_experts,), stride=(1,)
+            ),
+        )
+""",
+        "descriptor_map tensor",
+    )
+
+    # 2. Dispatch: resolve the row before the lookup. FC2 is compile-time; FC1
+    #    reads the half from output_n_tile, which is already a parameter here.
+    text = replace_once(
+        text,
+        """        combined_expert = block_expert_ids[metadata_block_idx].to(Int32)
+        total_experts = tier0_num_experts + tier1_num_experts
+        if combined_expert >= Int32(0) and combined_expert < total_experts:
+            descriptor = descriptor_map[combined_expert].to(Int32)
+""",
+        """        combined_expert = block_expert_ids[metadata_block_idx].to(Int32)
+        total_experts = tier0_num_experts + tier1_num_experts
+        # glm52-r7-projtiers: gate and up may sit in different tiers, so the
+        # descriptor row is chosen per projection. FC2 resolves at compile time;
+        # FC1 splits on the N half, which trellis3_t256_proj keeps aligned to
+        # whole CTA N tiles.
+        descriptor_row = Int32(2)
+        if cutlass.const_expr(is_fc1):
+            fc1_half_tiles = Int32(self.driver.fc1.n_tiles // 2)
+            descriptor_row = Int32(0)
+            if output_n_tile >= fc1_half_tiles:
+                descriptor_row = Int32(1)
+        if combined_expert >= Int32(0) and combined_expert < total_experts:
+            descriptor = descriptor_map[
+                descriptor_row * total_experts + combined_expert
+            ].to(Int32)
+""",
+        "per-projection descriptor row",
+    )
+
+    # 3. Host validation: the table is three rows now.
+    text = replace_once(
+        text,
+        """    total_experts = launch.tier0_num_experts + launch.tier1_num_experts
+    for name, mapping in (
+        ("global_to_combined", global_to_combined),
+        ("descriptor_map", descriptor_map),
+    ):
+        if (
+            mapping.dtype != torch.int32
+            or mapping.device != x.device
+            or not mapping.is_contiguous()
+            or int(mapping.numel()) != total_experts
+        ):
+""",
+        """    total_experts = launch.tier0_num_experts + launch.tier1_num_experts
+    for name, mapping in (
+        ("global_to_combined", global_to_combined),
+        ("descriptor_map", descriptor_map),
+    ):
+        # glm52-r7-projtiers: descriptor_map carries one row per projection.
+        expected_entries = (
+            3 * total_experts if name == "descriptor_map" else total_experts
+        )
+        if (
+            mapping.dtype != torch.int32
+            or mapping.device != x.device
+            or not mapping.is_contiguous()
+            or int(mapping.numel()) != expected_entries
+        ):
+""",
+        "descriptor_map validation",
+    )
+
+    # 4. FC1 and FC2 hold different expert counts once tiering is per
+    #    projection: on this checkpoint tier0 (K3) carries ~180 gate/up planes
+    #    but ~25 down planes, because down is 87% K4. Equalising by padding
+    #    costs 49.8 GiB per rank on the real bit maps, so the counts must be
+    #    allowed to differ. The FC2 count is read back from the tier's own
+    #    w2_global_scale, which is one fp32 per FC2 expert, so no launch-plan
+    #    or kernel-signature change is needed. The kernel's flat layouts stay
+    #    sized by the FC1 count: they are pointer arithmetic, not bounds checks,
+    #    and descriptor row 2 only ever emits FC2-local indices below the real
+    #    count, so no access lands past the allocation.
+    text = replace_once(
+        text,
+        """    expected_experts = int(expected_experts)
+    bits = int(bits)
+    fc1_cols = 2 * int(intermediate_size)
+""",
+        """    expected_experts = int(expected_experts)
+    # glm52-r7-projtiers
+    fc2_experts = int(tier.w2_global_scale.numel())
+    bits = int(bits)
+    fc1_cols = 2 * int(intermediate_size)
+""",
+        "fc2 expert count",
+    )
+    text = replace_once(
+        text,
+        """        (
+            "w2",
+            tier.w2,
+            torch.int32,
+            expected_experts
+            * (int(intermediate_size) // 16)
+            * (int(hidden_size) // 16)
+            * (8 * bits),
+        ),
+""",
+        """        (
+            "w2",
+            tier.w2,
+            torch.int32,
+            fc2_experts  # glm52-r7-projtiers
+            * (int(intermediate_size) // 16)
+            * (int(hidden_size) // 16)
+            * (8 * bits),
+        ),
+""",
+        "w2 storage extent",
+    )
+    text = replace_once(
+        text,
+        """        (
+            "w2_global_scale",
+            tier.w2_global_scale,
+            torch.float32,
+            expected_experts,
+        ),
+""",
+        """        (
+            "w2_global_scale",
+            tier.w2_global_scale,
+            torch.float32,
+            fc2_experts,  # glm52-r7-projtiers
+        ),
+""",
+        "w2 global scale extent",
+    )
+
+    # 5. The launch count is now an upper bound, not an exact count: with
+    #    per-projection tiering FC1 and FC2 hold different expert counts, and
+    #    the kernel guards dispatch with `local_expert < tierN_num_experts`, so
+    #    the count must cover the LARGER phase. The prepared tier still reports
+    #    its FC1 count, so require only that it fits inside the plan.
+    text = replace_once(
+        text,
+        """        actual_experts = int(tier.num_experts)
+        if actual_experts != int(expected_experts):
+            raise ValueError(
+                f"mixed Trellis {name} has {actual_experts} experts, but the "
+                f"launch plan describes {int(expected_experts)}"
+            )
+""",
+        """        actual_experts = int(tier.num_experts)
+        if actual_experts > int(expected_experts):  # glm52-r7-projtiers
+            raise ValueError(
+                f"mixed Trellis {name} has {actual_experts} experts, which "
+                f"exceeds the {int(expected_experts)} the launch plan allows"
+            )
+""",
+        "tier expert-count bound",
+    )
+
+
+    # 7. The `local_expert < tierN_num_experts` guard uses ONE count per tier,
+    #    which is the FC1 slot count because w13 storage is validated against
+    #    it. Under per-projection tiering FC2 can hold more experts than FC1
+    #    (tier1: 231 down vs 77 gate/up), so that bound wrongly rejects most
+    #    down-projection tiles and silently drops their output -- measured as
+    #    KLD 2.37 against a 0.0625 control. The descriptor map only ever emits
+    #    FC2 locals below the real FC2 count, so the bound is redundant there.
+    text = replace_once(
+        text,
+        """                if tier == Int32(0) and local_expert < tier0_num_experts:""",
+        """                fc1_bound_ok = local_expert < tier0_num_experts
+                if cutlass.const_expr(not is_fc1):
+                    fc1_bound_ok = True  # glm52-r7-projtiers
+                if tier == Int32(0) and fc1_bound_ok:""",
+        "tier0 fc2 bound",
+    )
+    text = replace_once(
+        text,
+        """                elif tier == Int32(1) and local_expert < tier1_num_experts:""",
+        """                fc1_bound_ok1 = local_expert < tier1_num_experts
+                if cutlass.const_expr(not is_fc1):
+                    fc1_bound_ok1 = True  # glm52-r7-projtiers
+                if tier == Int32(1) and fc1_bound_ok1:""",
+        "tier1 fc2 bound",
+    )
+
+    # 6. Host builder for the three-row table.
+    text = replace_once(
+        text,
+        """def combine_trellis_rotations(""",
+        '''def build_projection_tiered_maps(
+    gate_tiers: Sequence[int],
+    up_tiers: Sequence[int],
+    down_tiers: Sequence[int],
+    *,
+    tier_slots: Sequence[int],
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build the route map and the three-row descriptor map for one R7 layer.
+
+    Each argument is one tier id (0 or 1) per global expert, for that
+    projection. Combined expert ids are the global ids -- with per-projection
+    tiering there is no single tier-contiguous ordering, so all tier knowledge
+    lives in the descriptor rows and global_to_combined is the identity.
+
+    Returns (global_to_combined, descriptor_map) where descriptor_map is
+    int32[3 * num_experts] laid out gate, up, down, each entry
+    (tier << 8) | tier_local_index. glm52-r7-projtiers.
+    """
+
+    projections = (
+        ("gate", tuple(int(t) for t in gate_tiers)),
+        ("up", tuple(int(t) for t in up_tiers)),
+        ("down", tuple(int(t) for t in down_tiers)),
+    )
+    num_experts = len(projections[0][1])
+    rows: list[int] = []
+    for name, tiers in projections:
+        if len(tiers) != num_experts:
+            raise ValueError(
+                "mixed Trellis projection tier lists must agree on expert count: "
+                f"{name} has {len(tiers)}, expected {num_experts}"
+            )
+        if any(t not in (0, 1) for t in tiers):
+            raise ValueError(f"mixed Trellis {name} tier ids must be 0 or 1")
+        counters = [0, 0]
+        row = []
+        for tier in tiers:
+            local = counters[tier]
+            counters[tier] += 1
+            if local > 0xFF:
+                raise ValueError(
+                    f"mixed Trellis {name} tier {tier} exceeds 256 experts"
+                )
+            row.append((tier << 8) | local)
+        rows.extend(row)
+    # The launch sizes the combined namespace as the sum of the tier slot
+    # counts, and a tier slots max(gate_count, up_count), so with per-projection
+    # tiering that sum exceeds the real expert count (257 vs 256 on this
+    # checkpoint). Both tables are laid out on the padded stride; the tail
+    # entries are -1 and the kernel skips them, since it guards on
+    # 0 <= combined_expert < total and descriptor >= 0.
+    stride = int(sum(int(v) for v in tier_slots))
+    if stride < num_experts:
+        raise ValueError(
+            f"mixed Trellis tier slots ({stride}) cannot address "
+            f"{num_experts} experts"
+        )
+    global_to_combined = torch.full(
+        (stride,), -1, dtype=torch.int32, device=device
+    )
+    global_to_combined[:num_experts] = torch.arange(
+        num_experts, dtype=torch.int32, device=device
+    )
+    descriptor = torch.full((3 * stride,), -1, dtype=torch.int32, device=device)
+    for row_index in range(3):
+        base = row_index * num_experts
+        descriptor[row_index * stride : row_index * stride + num_experts] = (
+            torch.tensor(
+                rows[base : base + num_experts], dtype=torch.int32, device=device
+            )
+        )
+    return global_to_combined, descriptor
+
+
+def combine_trellis_rotations(''',
+        "projection map builder",
+    )
+    return text
+
+
+PATCHERS = {MIXED: patch_mixed}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--root", default="/opt/venv/lib/python3.12/site-packages", type=Path
+    )
+    parser.add_argument("--check", action="store_true")
+    args = parser.parse_args()
+
+    staged = {}
+    for relative, patcher in PATCHERS.items():
+        target = args.root / relative
+        text = target.read_text()
+        if MARKER in text:
+            print(f"already patched: {relative}")
+            return 0
+        staged[target] = patcher(text)
+        print(f"ready: {relative}")
+    if args.check:
+        print("sparkinfer projection tiers: anchors valid; not applied")
+        return 0
+    for target, text in staged.items():
+        target.write_text(text)
+        compile(text, str(target), "exec")
+    print("sparkinfer projection tiers: applied")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())