DeepInfer/scripts/flaggems_port/ab_quant.py

"""Is FlagGems' per-token FP8 quant a drop-in for `tile_kernels`'s, on the SM90 recipe?

This is the highest-value candidate in the FlagGems inventory: per-token quantisation runs on every
activation, in attention and in the MLP, so it is the one where replacing a torch reference with a
Triton kernel actually buys something.

Recipe on SM90 (`get_fp8_cast_params`):

    fmt=e4m3, num_per_channels=128, round_sf=True,
    use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True

The two `round_sf` implementations look different but should be the same function. DeepInfer does
`(int_bits + 0x7FFFFF) & 0x7F800000` -- add the maximum mantissa and mask to the exponent field,
which carries into the next exponent unless the mantissa was already zero, i.e. ceiling to a power of
two. FlagGems does `exp2(ceil(log2(s)))`, the same thing spelled arithmetically. Both then divide by
that scale. So this should match bit for bit, not merely closely -- which is why the comparison
insists on exactness and reports any difference rather than tolerating it.

Run: PYTHONPATH=. python scripts/flaggems_port/ab_quant.py
"""

from __future__ import annotations

import sys

import torch

from deepinfer.tk_compat.quant import per_token_cast

DEV = "cuda"
GROUP = 128
FAILED: list[str] = []


def flag_gems_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Call FlagGems' op without importing the package root.

    `import flag_gems` star-imports every op module and loads Triton TLE (which needs libflagcx).
    Pulling the single module in through the package still runs the package `__init__`, so this goes
    through the already-installed package but touches only what it must -- see the note in
    `ab_import_side_effects.py` for why importing the root is safe anyway.
    """
    from flag_gems.ops.per_token_group_quant_fp8 import per_token_group_quant_fp8

    return per_token_group_quant_fp8(x, GROUP, scale_ue8m0=True, column_major_scales=True)


def reference_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    return per_token_cast(
        x,
        "e4m3",
        GROUP,
        round_sf=True,
        use_packed_ue8m0=False,
        use_tma_aligned_col_major_sf=True,
    )


def compare(tag: str, a: torch.Tensor, b: torch.Tensor, exact: bool) -> None:
    if a.shape != b.shape:
        print(f"  DIFF {tag}: shape {tuple(a.shape)} vs {tuple(b.shape)}")
        FAILED.append(f"{tag}:shape")
        return
    if a.dtype != b.dtype:
        print(f"  DIFF {tag}: dtype {a.dtype} vs {b.dtype}")
        FAILED.append(f"{tag}:dtype")
        return
    if exact:
        same = torch.equal(a.view(torch.uint8) if a.dtype.itemsize == 1 else a, b.view(torch.uint8) if b.dtype.itemsize == 1 else b)
        print(f"  {'ok  ' if same else 'DIFF'} {tag}: bitwise identical = {same}")
        if not same:
            FAILED.append(tag)
        return
    diff = (a.float() - b.float()).abs()
    mx = diff.max().item()
    scale = a.float().abs().max().item()
    print(f"  {'ok  ' if mx == 0 else 'DIFF'} {tag}: max abs diff {mx:.3e} (max|a| {scale:.4f})")
    if mx != 0:
        FAILED.append(tag)


def case(tag: str, tokens: int, dim: int, seed: int) -> None:
    torch.manual_seed(seed)
    x = (torch.randn((tokens, dim), device=DEV, dtype=torch.bfloat16) * 0.5).contiguous()
    print(f"\n[{tag}] x {tuple(x.shape)} {x.dtype}")

    q_ref, s_ref = reference_quant(x)
    q_fg, s_fg = flag_gems_quant(x)
    torch.cuda.synchronize()

    compare("quantised values", q_ref, q_fg, exact=True)
    compare("scales", s_ref, s_fg, exact=False)

    # Dequantising both ways is the property that matters downstream: a matching scale with
    # mismatched values would still be wrong after the GEMM.
    deq_ref = q_ref.to(torch.float32) * s_ref.float().repeat_interleave(GROUP, dim=-1)
    deq_fg = q_fg.to(torch.float32) * s_fg.float().repeat_interleave(GROUP, dim=-1)
    compare("dequantised", deq_ref, deq_fg, exact=False)


CASES = [
    ("typical activation", 512, 4096, 0),
    ("single token", 1, 4096, 1),
    ("narrow", 64, 128, 2),
    ("odd token count", 7, 2048, 3),
    ("with small values", 128, 1024, 4),
]

if __name__ == "__main__":
    for c in CASES:
        case(*c)
    print()
    if FAILED:
        print(f"MISMATCH {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_MATCH ({len(CASES)} cases)")