DeepInfer/deepinfer/tk_compat/quant.py

"""`tile_kernels.quant` -- FP8/FP4 casting, fused SwiGLU, scale-factor layout.

The compiled functions take `num_per_channels` and write into caller-provided `out` buffers; the
references take `block_size` and return fresh tensors. The signatures here are the compiled ones,
because that is what the 4 call sites in the model are written against.

The adapters are thin but not vacuous -- three things genuinely differ:

* `num_per_channels=n` means a per-token block of `(1, n)`, not `(n, n)`.
* `swiglu_forward_and_per_token_cast` fuses activation and cast; the reference has them separate,
  so it is activation followed by cast.
* `cast_back`'s `fmt` is the *output* dtype (`'bf16'`/`'fp32'`), which is easy to misread as the
  quantised format.

Verified against the compiled kernels: `per_token_cast`, `cast_back` and
`swiglu_forward_and_per_token_cast` are bitwise identical on the SM90 recipe
(`scripts/tile_kernels_patch/compare_with_cuda.py`).
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import torch

from deepinfer.tk_compat import use_references

if not use_references():
    from tile_kernels import quant as _compiled

from tile_kernels import torch as _ref


def per_token_cast(
    x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    fmt: str,
    num_per_channels: int,
    x_block_size: tuple[int, int] | None = None,
    use_tma_aligned_col_major_sf: bool = False,
    round_sf: bool = False,
    use_packed_ue8m0: bool = False,
    use_e4m3_sf: bool = False,
    stochastic_cast: bool = False,
    out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Cast to FP8/FP4 with per-token scaling factors."""
    if not use_references():
        return _compiled.per_token_cast(
            x,
            fmt,
            num_per_channels,
            x_block_size,
            use_tma_aligned_col_major_sf,
            round_sf,
            use_packed_ue8m0,
            use_e4m3_sf,
            stochastic_cast,
            out,
        )
    if stochastic_cast:
        raise NotImplementedError(
            "stochastic_cast has no torch reference; it changes rounding, so silently using the "
            "deterministic path would produce plausible but different numbers"
        )
    result = _ref.cast(
        x,
        fmt,
        block_size=(1, num_per_channels),
        x_block_size=x_block_size,
        round_sf=round_sf,
        use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
        use_packed_ue8m0=use_packed_ue8m0,
        use_e4m3_sf=use_e4m3_sf,
    )
    if out is not None:
        # Honour the out-buffer contract: the call sites allocate these and may hold views.
        out[0].copy_(result[0])
        out[1].copy_(result[1])
        return out[0], out[1]
    return result


def cast_back(
    x: tuple[torch.Tensor, torch.Tensor],
    fmt: str,
    x_block_size: tuple[int, int],
    x_special_fmt: str | None = None,
) -> torch.Tensor:
    """Dequantise an FP8/FP4 tensor. `fmt` is the output dtype, not the quantised format."""
    if not use_references():
        return _compiled.cast_back(x, fmt, x_block_size, x_special_fmt)
    if x_special_fmt is not None:
        raise NotImplementedError(f"x_special_fmt={x_special_fmt!r} has no torch reference")
    return _ref.cast_back(x, fmt, block_size=x_block_size)


def swiglu_forward(
    x: torch.Tensor,
    fmt: str,
    psum_num_tokens_per_expert: torch.Tensor | None = None,
    topk_weights: torch.Tensor | None = None,
    routed_scaling_factor: float | None = None,
    clamp_value: float | None = None,
    clamped_count: torch.Tensor | None = None,
) -> torch.Tensor:
    """SwiGLU activation. `fmt` is unused by the compiled path here; kept for signature parity."""
    if not use_references():
        return _compiled.swiglu_forward(
            x,
            fmt,
            psum_num_tokens_per_expert,
            topk_weights,
            routed_scaling_factor,
            clamp_value,
            clamped_count,
        )
    return _ref.swiglu_forward(
        x,
        psum_num_tokens_per_expert,
        topk_weights,
        routed_scaling_factor,
        clamp_value,
        clamped_count,
    )


def swiglu_forward_and_per_token_cast(
    x: torch.Tensor,
    fmt: str,
    num_per_channels: int,
    psum_num_tokens_per_expert: torch.Tensor | None = None,
    topk_weights: torch.Tensor | None = None,
    routed_scaling_factor: float | None = None,
    use_tma_aligned_col_major_sf: bool = False,
    round_sf: bool = False,
    use_packed_ue8m0: bool = False,
    clamp_value: float | None = None,
    clamped_count: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused SwiGLU + per-token cast."""
    if not use_references():
        return _compiled.swiglu_forward_and_per_token_cast(
            x,
            fmt,
            num_per_channels,
            psum_num_tokens_per_expert,
            topk_weights,
            routed_scaling_factor,
            use_tma_aligned_col_major_sf,
            round_sf,
            use_packed_ue8m0,
            clamp_value,
            clamped_count,
        )
    activated = _ref.swiglu_forward(
        x,
        psum_num_tokens_per_expert,
        topk_weights,
        routed_scaling_factor,
        clamp_value,
        clamped_count,
    )
    return _ref.cast(
        activated,
        fmt,
        block_size=(1, num_per_channels),
        round_sf=round_sf,
        use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
        use_packed_ue8m0=use_packed_ue8m0,
    )


def batched_transpose_weight_sf(
    sf: Sequence[torch.Tensor],
    block_size: int,
    transpose_sf1s: Sequence[bool],
    outs: Sequence[tuple[torch.Tensor, torch.Tensor] | None],
) -> list[tuple[torch.Tensor, torch.Tensor]]:
    """Transpose the scale factors of several weights in one launch."""
    if not use_references():
        return _compiled.batched_transpose_weight_sf(sf, block_size, transpose_sf1s, outs)
    # The reference allocates its outputs, so the `outs` slots cannot be filled in place. Check
    # rather than ignore: a caller that pre-allocated would read stale memory.
    if any(o is not None for o in outs):
        raise NotImplementedError("the torch reference allocates its own outputs; pre-allocated `outs` cannot be filled")
    return _ref.batched_transpose_weight_sf(sf, block_size, transpose_sf1s)


def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """Decode packed FP4 pairs. Same on both paths -- it is host-side bit manipulation."""
    return _ref.unpack_from_e2m1fn_x2(x, out_dtype)


def backend() -> str:
    return "ref" if use_references() else "cuda"


def _passthrough(name: str) -> Any:
    """Resolve a name this module does not adapt from whichever backend is active."""
    if not use_references():
        return getattr(_compiled, name)
    return getattr(_ref, name)