DeepInfer/tests/neuron/compile/test_nki_hc.py

"""vllm-neuron#71 mHC / RMSNorm / SwiGLU NKI kernels behind DeepInfer's ops (issue #13), on a NeuronCore.

Each case runs DeepInfer's entry point (``mhc_ops`` / ``mhc_modeling`` / ``tilelang_ops.norm`` / ``tile_kernels``
SwiGLU, or ``nki_hc.mixes`` for the shifted ``hc_pre`` mixes) on device tensors -- eagerly and inside
``torch.compile(backend="neuron", fullgraph=True)`` -- and compares with the same entry point on host tensors,
which runs the existing pure-torch shims. Shapes are DeepSeek-V4.1's: the tiny checkpoint (hidden 1024, MoE
intermediate 256) and the full -1017 model (hidden 5120, intermediate 2304), hc_mult 4, 20 Sinkhorn iterations,
hc_eps 1e-6, rms_norm_eps 1e-20. Tolerances are bf16-level: the kernels are not bit-exact emulations.

    NEURON_RT_VISIBLE_CORES=8 DEEPINFER_ACCELERATOR=neuron python -m pytest -q tests/neuron/compile/test_nki_hc.py
"""

from __future__ import annotations

import os

import pytest
import torch

os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
if os.environ.get("DEEPINFER_NEURON_DEVICE") == "cpu":
    pytest.skip("device test (host mode)", allow_module_level=True)
import deepinfer  # noqa: E402,F401  (puts neuron_ops/vendor on sys.path)

try:
    import torch_neuronx  # noqa: F401

    DEV = torch.device("neuron", 0)
except ImportError:  # pragma: no cover
    pytest.skip("needs torch-neuronx and a NeuronCore", allow_module_level=True)

from deepinfer.model_executor.layers.tilelang_ops.norm import norm  # noqa: E402
from deepinfer.model_executor.models.v4_modules import mhc_ops  # noqa: E402
from deepinfer.model_executor.models.v4_modules.mhc_modeling import mhc_pre_big_fuse  # noqa: E402
from deepinfer.neuron_ops import nki_hc, nki_norm  # noqa: E402
from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda  # noqa: E402
from tile_kernels.quant import swiglu_forward_and_per_token_cast  # noqa: E402

HC = 4
ITERS, HC_EPS, NORM_EPS, POST_MULT = 20, 1e-6, 1e-20, 2.0
#: (hidden_size, moe_intermediate_size): tiny checkpoint, DeepSeek-V4.1-Flash-1017.
CONFIGS = {"tiny": (1024, 256), "full": (5120, 2304)}
TOKENS = [1, 130]  # one decode token; more than one 128-row tile plus a tail
MODES = ["eager", "compile"]


USED: list[str] = []


def _counting(original):  # noqa: ANN001, ANN202
    def kernel(*names):  # noqa: ANN002, ANN202
        USED.append(names[0])
        return original(*names)

    return kernel


def run(mode: str, fn, *args):  # noqa: ANN001, ANN201
    """``fn(*args)`` with the args moved to the device, eagerly or as one fullgraph region; results to the host.
    Eagerly, also checks that the call went through an NKI kernel rather than the shim on device tensors."""
    dargs = [a.to(DEV) if isinstance(a, torch.Tensor) else a for a in args]
    if mode == "compile":
        torch._dynamo.reset()
        compiled = torch.compile(fn, backend="neuron", fullgraph=True, dynamic=False)
        with hide_cuda():
            out = compiled(*dargs)
    else:
        USED.clear()
        saved = nki_hc._kernel, nki_norm._kernel
        nki_hc._kernel, nki_norm._kernel = _counting(saved[0]), _counting(saved[1])
        try:
            out = fn(*dargs)
        finally:
            nki_hc._kernel, nki_norm._kernel = saved
        assert USED, "no NKI kernel ran"
    return torch.utils._pytree.tree_map(lambda t: t.cpu() if isinstance(t, torch.Tensor) else t, out)


def close(got: torch.Tensor, ref: torch.Tensor, rtol: float = 1.6e-2, atol: float = 1e-2) -> None:
    assert got.shape == ref.shape and got.dtype == ref.dtype, f"{got.shape} {got.dtype} vs {ref.shape} {ref.dtype}"
    torch.testing.assert_close(got.float(), ref.float(), rtol=rtol, atol=atol)


def _trunk(n: int, h: int) -> torch.Tensor:
    return (torch.randn(1, n, HC, h) * 3).to(torch.bfloat16)


def _prepared(fn: torch.Tensor) -> torch.Tensor:
    """``fn`` on the device with its bf16 hi / lo split attached, as ``nki_hc.prepare_model`` does at load."""
    d = fn.to(DEV)
    nki_hc.split_fn(d)
    return d


def _hc_params(h: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """fn ``[(2 + hc) * hc, hc * h]`` fp32 (norm weight folded in), scale ``[3]``, base ``[24]``."""
    width = (2 + HC) * HC
    fn = torch.randn(width, HC * h) * (HC * h) ** -0.5
    scale = torch.tensor([0.7, 0.4, 1.5])
    base = torch.cat([torch.full((HC,), -1.0), torch.zeros(HC), (torch.eye(HC) * 6.0 - 3.0).flatten()])
    base = base + 0.1 * torch.randn(width)
    return fn, scale, base


@pytest.fixture(autouse=True)
def _seed() -> None:
    torch.manual_seed(0)


# ---------------------------------------------------------------------------------------------------------- mHC


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_expand(cfg: str, n: int, mode: str) -> None:
    h = CONFIGS[cfg][0]
    x = torch.randn(n, h).to(torch.bfloat16)
    got = run(mode, mhc_ops.expand_to_mhc, x, HC)
    assert torch.equal(got, mhc_ops.expand_to_mhc(x, HC))  # pure data movement: exact


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_pre_apply_mix(cfg: str, n: int, mode: str) -> None:
    h = CONFIGS[cfg][0]
    x = _trunk(n, h)
    mix = torch.rand(1, n, HC, 1) + HC_EPS
    close(run(mode, mhc_ops.mhc_pre_apply_mix, x, mix), mhc_ops.mhc_pre_apply_mix(x, mix))


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_post(cfg: str, n: int, mode: str) -> None:
    h = CONFIGS[cfg][0]
    x = torch.randn(1, n, h).to(torch.bfloat16)
    residual = _trunk(n, h)
    post = 2 * torch.rand(1, n, HC, 1)
    comb = mhc_ops.sinkhorn_normalize(torch.randn(1, n, HC, HC), ITERS, HC_EPS)
    close(run(mode, mhc_ops.mhc_post, x, residual, post, comb), mhc_ops.mhc_post(x, residual, post, comb))


def _shim_mixes(residual, fn, scale, base):  # noqa: ANN001, ANN202
    """hc.py's shifted ``hc_pre`` mixes on the shim path."""
    mixes = mhc_ops.mhc_pre_norm_fn(residual, fn, None, NORM_EPS)
    pre, post, comb = mhc_ops.mhc_pre_split_mixes(mixes, scale, base, HC, POST_MULT, HC_EPS)
    return pre, post, mhc_ops.sinkhorn_normalize(comb, repeat=ITERS, eps=HC_EPS)


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_mixes_shifted(cfg: str, n: int, mode: str) -> None:
    h = CONFIGS[cfg][0]
    residual = _trunk(n, h)
    fn, scale, base = _hc_params(h)
    assert nki_hc.mixes_supported(residual.to(DEV), NORM_EPS, HC_EPS, HC_EPS, POST_MULT, ITERS)
    got = run(mode, nki_hc.mixes, residual, _prepared(fn), scale, base)
    ref = _shim_mixes(residual, fn, scale, base)
    for g, r, name in zip(got, ref, ("pre", "post", "comb"), strict=True):
        assert g.shape == r.shape, name
        torch.testing.assert_close(g, r, rtol=1e-4, atol=2e-5, msg=lambda m, name=name: f"{name}: {m}")
    # doubly stochastic up to Sinkhorn's residual
    torch.testing.assert_close(got[2].sum(-2), torch.ones(1, n, HC), rtol=0, atol=1e-3)


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_pre_big_fuse(cfg: str, n: int, mode: str) -> None:
    h = CONFIGS[cfg][0]
    residual = _trunk(n, h)
    fn, scale, base = _hc_params(h)

    def f(r, fn, s, b):  # noqa: ANN001, ANN202
        return mhc_pre_big_fuse(r, fn, s, b, NORM_EPS, HC_EPS, HC_EPS, POST_MULT, ITERS)

    got_post, got_comb, got_in = run(mode, f, residual, _prepared(fn), scale, base)
    ref_post, ref_comb, ref_in = f(residual, fn, scale, base)
    torch.testing.assert_close(got_post, ref_post, rtol=1e-4, atol=2e-5)
    torch.testing.assert_close(got_comb, ref_comb, rtol=1e-4, atol=2e-5)
    close(got_in, ref_in)


@pytest.mark.parametrize("cfg", CONFIGS)
def test_fn_gemm_fp32_in_graph(cfg: str) -> None:
    """neuronx-cc narrows an fp32 matmul's operands to bf16 inside a graph; the prepared hi / lo split does not.
    (If the plain product ever becomes fp32-accurate in a graph, the split can go.)"""
    h = CONFIGS[cfg][0]
    flat = _trunk(130, h).reshape(130, HC * h)
    fn = _hc_params(h)[0]
    ref = flat.double() @ fn.double().T
    rng = ref.abs().max().item()
    split = run("compile", nki_hc._fn_gemm, flat, _prepared(fn)).double()
    plain = run("compile", lambda x, f: x.float() @ f.T, flat, fn).double()
    assert (split - ref).abs().max().item() < 1e-5 * rng
    assert (plain - ref).abs().max().item() > 1e-4 * rng  # known: bf16-level (~2e-3 of the range)


def test_mixes_fall_back_on_other_constants() -> None:
    r = _trunk(1, 128).to(DEV)
    assert not nki_hc.mixes_supported(r, NORM_EPS, HC_EPS, HC_EPS, POST_MULT, 10)
    assert not nki_hc.mixes_supported(r, 1e-6, HC_EPS, HC_EPS, POST_MULT, ITERS)
    assert not nki_hc.mixes_supported(r.cpu(), NORM_EPS, HC_EPS, HC_EPS, POST_MULT, ITERS)


# ------------------------------------------------------------------------------------------------------ RMSNorm

#: V4.1 norm widths: hidden, q_lora_rank (full), head_dim (kv_norm), index_head_dim (k_norm).
NORM_DIMS = {"tiny": [1024, 256, 512, 128], "full": [5120, 1280, 512, 128]}


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_norm_bf16(cfg: str, n: int, mode: str) -> None:
    for d in NORM_DIMS[cfg]:
        x = (torch.randn(n, d) * 2).to(torch.bfloat16)
        w = (1 + 0.1 * torch.randn(d)).to(torch.bfloat16)
        close(run(mode, lambda x, w: norm(x, w, eps=NORM_EPS), x, w), norm(x, w, eps=NORM_EPS))


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
def test_norm_residual_out_scale(n: int, mode: str) -> None:
    d = 5120
    x, res = (torch.randn(n, d).to(torch.bfloat16) for _ in range(2))
    w = (1 + 0.1 * torch.randn(d)).to(torch.bfloat16)

    def f(x, w, r):  # noqa: ANN001, ANN202
        return norm(x, w, residual=r, eps=1e-6, out_scale=1.5, return_bf16=True)

    (got_y, got_bf16), got_res = run(mode, f, x, w, res)
    (ref_y, ref_bf16), ref_res = f(x, w, res)
    close(got_y, ref_y)
    close(got_bf16, ref_bf16)
    assert torch.equal(got_res, ref_res)


def _dequant(q: torch.Tensor, sf: torch.Tensor, npc: int) -> torch.Tensor:
    """e4m3 values times their ue8m0 scales (int16 words, low byte first)."""
    e = sf.contiguous().view(torch.uint8).reshape(*q.shape[:-1], -1)[..., : q.shape[-1] // npc]
    return q.float() * torch.exp2(e.float() - 127).repeat_interleave(npc, -1)


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_norm_fp8(cfg: str, n: int, mode: str) -> None:
    """The production cast: e4m3, 1x32, ue8m0 packed into int16 (get_fp8_cast_params), plus the bf16 copy."""
    from deepinfer.model_executor.layers.linear import get_fp8_cast_params

    params = get_fp8_cast_params()
    npc = params["num_per_channels"]
    d = CONFIGS[cfg][0]
    x = (torch.randn(n, d) * 2).to(torch.bfloat16)
    w = (1 + 0.1 * torch.randn(d)).to(torch.bfloat16)

    def f(x, w):  # noqa: ANN001, ANN202
        return norm(x, w, eps=NORM_EPS, return_bf16=True, **params)

    (got_q, got_sf), got_bf16 = run(mode, f, x, w)
    (ref_q, ref_sf), ref_bf16 = f(x, w)
    close(got_bf16, ref_bf16)
    assert got_q.dtype == torch.float8_e4m3fn and got_q.shape == ref_q.shape
    assert got_sf.dtype == ref_sf.dtype == torch.int16 and got_sf.shape == ref_sf.shape
    # Scales are exponents of each group's amax: equal unless the bf16 result moved an amax across a power of two.
    same_sf = (got_sf.view(torch.uint8) == ref_sf.view(torch.uint8)).float().mean().item()
    assert same_sf >= 0.99, same_sf
    got_d, ref_d = _dequant(got_q, got_sf, npc), _dequant(ref_q, ref_sf, npc)
    # one e4m3 step (3 mantissa bits) of slack on top of the bf16 difference
    torch.testing.assert_close(got_d, ref_d, rtol=0.13, atol=ref_d.abs().max().item() * 2**-9)
    # and the cast itself is exact on the kernel's own bf16 result
    q2, sf2 = nki_norm.fp8_cast(got_bf16, npc, True, True, torch.int16)
    assert torch.equal(q2.view(torch.uint8), got_q.view(torch.uint8)) and torch.equal(sf2, got_sf)


def test_ue8m0_exponent_matches_bits() -> None:
    """The bitcast-free exponent against the shim's ``((bits - 1) >> 23) + 1`` (device, compiled)."""
    p = torch.exp2(torch.arange(-30, 30).float())
    v = torch.cat([torch.rand(8192) * 100 + 1e-6, p, p * (1 + 2**-23), p * (1 - 2**-24)])
    want = ((v.view(torch.int32) - 1) >> 23) + 1
    assert torch.equal(run("compile", nki_norm.ue8m0_exponent, v), want)


# ------------------------------------------------------------------------------------------------------- SwiGLU


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
@pytest.mark.parametrize("routed", [False, True])
def test_swiglu(cfg: str, n: int, mode: str, routed: bool) -> None:
    """static_moe's call: bf16 out, clamp 10 (swiglu_limit), routed experts with their per-row weight."""
    inter = CONFIGS[cfg][1]
    x = (torch.randn(n, 2 * inter) * 5).to(torch.bfloat16)  # crosses the clamp
    w = torch.rand(n) if routed else None

    def f(x, w):  # noqa: ANN001, ANN202
        return swiglu_forward_and_per_token_cast(
            x, fmt="bf16", num_per_channels=None, topk_weights=w, routed_scaling_factor=1.5 if routed else None,
            clamp_value=10.0,
        )

    close(run(mode, f, x, w), f(x, w), atol=2e-2)


# ------------------------------------------------------------------------------------------- one hc_pre region


@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("n", TOKENS)
@pytest.mark.parametrize("cfg", CONFIGS)
def test_shifted_post_pre_chain(cfg: str, n: int, mode: str) -> None:
    """HCLayerMixin._hc_fused_post_pre on the shifted (V4.1) path as one graph: the previous sub-layer's post, the
    collapse with the carried pre mix, the next mixes and the fp8 RMSNorm."""
    from deepinfer.model_executor.layers.linear import get_fp8_cast_params

    h = CONFIGS[cfg][0]
    prev_out = torch.randn(n, h).to(torch.bfloat16)
    residual = _trunk(n, h)
    post = 2 * torch.rand(1, n, HC, 1)
    comb = mhc_ops.sinkhorn_normalize(torch.randn(1, n, HC, HC), ITERS, HC_EPS)
    prev_pre = torch.rand(1, n, HC, 1) + HC_EPS
    fn, scale, base = _hc_params(h)
    w = (1 + 0.1 * torch.randn(h)).to(torch.bfloat16)

    def f(prev_out, residual, post, comb, prev_pre, fn, scale, base, w):  # noqa: ANN001, ANN202
        trunk = mhc_ops.mhc_post(prev_out.unsqueeze(0), residual, post, comb)
        layer_input = mhc_ops.mhc_pre_apply_mix(trunk, prev_pre)
        if trunk.device.type == "neuron":
            pre, new_post, new_comb = nki_hc.mixes(trunk, fn, scale, base)
        else:
            pre, new_post, new_comb = _shim_mixes(trunk, fn, scale, base)
        (q, sf), y = norm(layer_input.squeeze(0), w, eps=NORM_EPS, return_bf16=True, **get_fp8_cast_params())
        return trunk, pre, new_post, new_comb, y, q, sf

    args = (prev_out, residual, post, comb, prev_pre, fn, scale, base, w)
    got, ref = run(mode, f, *args[:5], _prepared(fn), *args[6:]), f(*args)
    close(got[0], ref[0])
    for g, r in zip(got[1:4], ref[1:4], strict=True):
        torch.testing.assert_close(g, r, rtol=1e-4, atol=2e-5)
    close(got[4], ref[4])
    npc = get_fp8_cast_params()["num_per_channels"]
    torch.testing.assert_close(_dequant(got[5], got[6], npc), _dequant(ref[5], ref[6], npc), rtol=0.13, atol=1e-2)