fast_switch_patch.py

"""Speed patch for PolarQuantizedSwitchLinear at canvas/batch scale.

DiffusionGemma denoising pushes ~2048 (token, expert) routings per switch call
(256-token canvas x top-8). polar_multi_gather_qmv re-reads the input vector
from global memory once per output row, which is ~33 ms/call at that scale.
This patch routes large-N batched calls to: fused Metal dequant (packed tq ->
fp16, one thread per quant group, pure bandwidth) + mx.gather_mm. ~10 ms/call.

Decode-scale calls (n_tokens==1 or small k) keep the existing fused kernels.

Usage:  from fast_switch_patch import patch_switch_fast; patch_switch_fast()
"""

import math

import mlx.core as mx

from turboquant_mlx.layers.polar_switch_linear import PolarQuantizedSwitchLinear
from turboquant_mlx.core.rotation import rotate_input

# Above this many (token, expert) routings, dequant+gather_mm beats the
# per-row gather kernels (measured crossover is well below canvas scale).
LARGE_N_THRESHOLD = 512

_kernel_cache: dict[tuple[int, int], object] = {}


def _get_dequant_kernel(bits: int, group_size: int):
    """One thread per (expert, row, group): unpack + codebook + scale -> fp16."""
    key = (bits, group_size)
    if key in _kernel_cache:
        return _kernel_cache[key]

    n_codes = 1 << bits
    elems_per_u32 = 32 // bits
    mask = (1 << bits) - 1

    source = f"""
    uint gid = thread_position_in_grid.x;
    uint n_groups = scales_shape[2];
    uint out_rows = scales_shape[1];
    uint pw_cols = packed_weight_shape[2];
    uint in_dims = n_groups * {group_size}u;
    uint total = scales_shape[0] * out_rows * n_groups;
    if (gid >= total) return;

    uint g = gid % n_groups;
    uint row = (gid / n_groups) % out_rows;
    uint e = gid / (n_groups * out_rows);

    float cb[{n_codes}];
    for (uint i = 0; i < {n_codes}u; i++) {{ cb[i] = float(codebook[i]); }}

    float scale = float(scales[gid]);
    uint pw_base = (e * out_rows + row) * pw_cols;
    uint out_base = (e * out_rows + row) * in_dims + g * {group_size}u;

    for (uint t = 0; t < {group_size}u; t++) {{
        uint col = g * {group_size}u + t;
        uint word = packed_weight[pw_base + col / {elems_per_u32}u];
        uint code = (word >> ((col % {elems_per_u32}u) * {bits}u)) & {mask}u;
        out[out_base + t] = T(cb[code] * scale);
    }}
    """
    _kernel_cache[key] = mx.fast.metal_kernel(
        name=f"polar_dequant_experts_{bits}bit_gs{group_size}",
        input_names=["packed_weight", "scales", "codebook"],
        output_names=["out"],
        source=source,
        ensure_row_contiguous=True,
    )
    return _kernel_cache[key]


def dequant_experts_fast(layer: PolarQuantizedSwitchLinear) -> mx.array:
    """(num_experts, output_dims, input_dims) fp16 via the fused Metal kernel."""
    kernel = _get_dequant_kernel(layer.bits, layer.group_size)
    e, o = layer.num_experts, layer.output_dims
    n_groups = layer.input_dims // layer.group_size
    total = e * o * n_groups
    return kernel(
        inputs=[layer.weight, layer.scales, layer.codebook],
        template=[("T", layer.scales.dtype)],
        grid=(total, 1, 1),
        threadgroup=(256, 1, 1),
        output_shapes=[(e, o, layer.input_dims)],
        output_dtypes=[layer.scales.dtype],
    )[0]


_orig_call = PolarQuantizedSwitchLinear.__call__


def _patched_call(self, x, indices, sorted_indices=False):
    n_tokens = 1 if x.ndim <= 2 else math.prod(x.shape[:-2])
    k = indices.shape[-1] if indices.ndim >= 1 else 1

    # Canvas/prefill scale: fused dequant + native gather_mm.
    if (n_tokens == k and k >= LARGE_N_THRESHOLD) or (
        n_tokens != 1 and n_tokens != k
    ):
        if self._needs_rotation:
            x = rotate_input(x, self.signs)
        w_deq = dequant_experts_fast(self)
        y = mx.gather_mm(
            x,
            w_deq.swapaxes(-1, -2),
            rhs_indices=indices,
            sorted_indices=sorted_indices,
        )
        if "bias" in self:
            y = y + mx.expand_dims(self["bias"][indices], -2)
        return y

    return _orig_call(self, x, indices, sorted_indices=sorted_indices)


def patch_switch_fast():
    PolarQuantizedSwitchLinear.__call__ = _patched_call
    print(f"[INFO] fast-switch patch active (threshold={LARGE_N_THRESHOLD} routings)")