tests/example_usage.py

"""Example usage of the kda-neuron-kernels package.

Demonstrates how to invoke each kernel entry point with correct preprocessing.
Runs a small parity check against the fla-core `naive_recurrent_kda` /
`naive_chunk_kda` PyTorch reference to verify the kernel is producing correct
output on your instance.

Requirements to run this script:
    - fla-core installed (`pip install fla-core`); we import it via direct file
      loading to avoid triton dependency issues on CPU-only hosts.
    - PyTorch Native Beta 3+ environment activated (torch-neuronx >= 2.11.3).
    - NKI >= 0.4.0.

Usage:
    source $HOME/workspace/native_venv/bin/activate
    export NEURON_RT_NUM_CORES=4
    python tests/example_usage.py
"""

import importlib.util
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F


def log(msg):
    print(f"[example_usage {time.strftime('%H:%M:%S')}] {msg}", flush=True)


def load_kernels():
    """Load the three kernel entry points from the sibling `build/torch-neuron/` dir."""
    # Assumes this script is at kda-neuron-kernels/tests/example_usage.py
    build_dir = Path(__file__).parent.parent / "build" / "torch-neuron"
    sys.path.insert(0, str(build_dir))
    from nki_kda import kda_recurrent_fwd, kda_recurrent_fwd_state
    from nki_kda_chunked import kda_chunk_step
    return kda_recurrent_fwd, kda_recurrent_fwd_state, kda_chunk_step


def load_fla_naive_reference():
    """Load fla-core's naive_recurrent_kda + naive_chunk_kda without triggering triton import.

    Returns (naive_recurrent_kda, naive_chunk_kda) or (None, None) if fla-core is not installed.
    """
    try:
        import fla  # noqa: F401
    except ImportError:
        return None, None

    fla_naive_path = Path(fla.__file__).parent / "ops" / "kda" / "naive.py"
    if not fla_naive_path.exists():
        return None, None

    spec = importlib.util.spec_from_file_location("fla_kda_naive", str(fla_naive_path))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod.naive_recurrent_kda, mod.naive_chunk_kda


def cos_sim(a, b):
    a = torch.as_tensor(a, dtype=torch.float32).flatten()
    b = torch.as_tensor(b, dtype=torch.float32).flatten()
    return F.cosine_similarity(a, b, dim=0).item()


def max_abs_diff(a, b):
    return float(torch.abs(torch.as_tensor(a, dtype=torch.float32) - torch.as_tensor(b, dtype=torch.float32)).max())


def main():
    log("Loading kernels...")
    kda_recurrent_fwd, kda_recurrent_fwd_state, kda_chunk_step = load_kernels()

    log("Loading fla-core reference (optional)...")
    fla_recurrent, fla_chunk = load_fla_naive_reference()
    have_fla = fla_recurrent is not None and fla_chunk is not None
    if not have_fla:
        log("  fla-core not available; will run kernels but skip parity check")

    # =========================================================================
    # Example 1: Recurrent kernel on a small sequence
    # =========================================================================
    log("=" * 70)
    log("Example 1: kda_recurrent_fwd on S=16 tokens, single (batch, head)")
    log("=" * 70)

    torch.manual_seed(42)
    S = 16
    D = 128

    q_raw = torch.randn(S, D, dtype=torch.float32)
    k_raw = torch.randn(S, D, dtype=torch.float32)
    v = torch.randn(S, D, dtype=torch.float32) * 0.3
    g = -torch.rand(S, D, dtype=torch.float32) * 0.01
    beta_row = torch.rand(S, 1) - 0.5 + 1.0
    beta = beta_row.expand(S, D).contiguous()

    # Wrapper preprocessing: L2-norm q, k and scale q by 1/sqrt(D)
    q = F.normalize(q_raw, p=2, dim=-1) * (D ** -0.5)
    k = F.normalize(k_raw, p=2, dim=-1)

    q_dev = q.to("neuron")
    k_dev = k.to("neuron")
    v_dev = v.to("neuron")
    g_dev = g.to("neuron")
    beta_dev = beta.to("neuron")

    log(f"  Input shapes: q, k, v, g, beta each ({S}, {D})")

    t0 = time.time()
    out = kda_recurrent_fwd(q_dev, k_dev, v_dev, g_dev, beta_dev)
    out_cpu = torch.as_tensor(out).to("cpu") if isinstance(out, torch.Tensor) else torch.from_numpy(np.asarray(out))
    log(f"  kda_recurrent_fwd returned shape={tuple(out_cpu.shape)} in {time.time()-t0:.2f}s")

    if have_fla:
        # fla reference contract: takes unnormalized q, k, computes scale=K^-0.5 internally.
        # Pass q_raw (unnormalized) but with L2-norm; fla applies K^-0.5 with default scale.
        # Actually to match our NKI q_double = l2norm(q) * K^-0.5 exactly, pass scale=1.0:
        ref_o, _ = fla_recurrent(
            q.unsqueeze(0).unsqueeze(2),  # (B=1, S, H=1, K)
            k.unsqueeze(0).unsqueeze(2),
            v.unsqueeze(0).unsqueeze(2),
            g.unsqueeze(0).unsqueeze(2),
            beta_row.squeeze(-1).unsqueeze(0).unsqueeze(2),  # (B=1, S, H=1)
            scale=1.0,  # already scaled by K^-0.5
        )
        ref_slice = ref_o[0, :, 0]
        cs = cos_sim(out_cpu, ref_slice)
        md = max_abs_diff(out_cpu, ref_slice)
        log(f"  Parity vs fla naive_recurrent_kda: cos_sim={cs:.6f}  max_diff={md:.6e}")
        if cs < 0.9995:
            log(f"  WARNING: cos_sim below 0.9995. Expected ~0.99997 at S=16 on random inputs.")
        else:
            log(f"  OK: parity within expected floor for BF16 tensor engine accumulation.")

    # =========================================================================
    # Example 2: Chunked kernel on one 128-token chunk
    # =========================================================================
    log("=" * 70)
    log("Example 2: kda_chunk_step on C=128 tokens, single (batch, head)")
    log("=" * 70)

    C = 128
    torch.manual_seed(42)
    q_raw = torch.randn(C, D, dtype=torch.float32)
    k_raw = torch.randn(C, D, dtype=torch.float32)
    v = torch.randn(C, D, dtype=torch.float32) * 0.3
    g_step = -torch.rand(C, D, dtype=torch.float32) * 0.01
    beta_row = torch.rand(C, 1) - 0.5 + 1.0
    beta = beta_row.expand(C, D).contiguous()

    q = F.normalize(q_raw, p=2, dim=-1) * (D ** -0.5)
    k = F.normalize(k_raw, p=2, dim=-1)

    gc = torch.cumsum(g_step, dim=0)
    g_last = gc[-1:, :].expand(C, D).contiguous()

    state = torch.zeros(C, D, dtype=torch.float32)

    q_dev = q.to("neuron")
    k_dev = k.to("neuron")
    v_dev = v.to("neuron")
    beta_dev = beta.to("neuron")
    gc_dev = gc.to("neuron")
    gl_dev = g_last.to("neuron")
    state_dev = state.to("neuron")

    t0 = time.time()
    chunk_out, state_new = kda_chunk_step(q_dev, k_dev, v_dev, beta_dev, gc_dev, gl_dev, state_dev)
    chunk_out_cpu = torch.as_tensor(chunk_out).to("cpu") if isinstance(chunk_out, torch.Tensor) else torch.from_numpy(np.asarray(chunk_out))
    state_new_cpu = torch.as_tensor(state_new).to("cpu") if isinstance(state_new, torch.Tensor) else torch.from_numpy(np.asarray(state_new))
    log(f"  kda_chunk_step returned chunk_out shape={tuple(chunk_out_cpu.shape)} "
        f"state shape={tuple(state_new_cpu.shape)} in {time.time()-t0:.2f}s")

    if have_fla:
        ref_out, _ = fla_chunk(
            q.unsqueeze(0).unsqueeze(2),
            k.unsqueeze(0).unsqueeze(2),
            v.unsqueeze(0).unsqueeze(2),
            g_step.unsqueeze(0).unsqueeze(2),
            beta_row.squeeze(-1).unsqueeze(0).unsqueeze(2),
            scale=1.0,
            chunk_size=128,
        )
        ref_slice = ref_out[0, :, 0]
        cs = cos_sim(chunk_out_cpu, ref_slice)
        md = max_abs_diff(chunk_out_cpu, ref_slice)
        log(f"  Parity vs fla naive_chunk_kda: cos_sim={cs:.6f}  max_diff={md:.6e}")
        if cs < 0.999:
            log(f"  WARNING: cos_sim below 0.999. Expected ~0.99988 at C=128 on random inputs.")
        else:
            log(f"  OK: parity within expected floor for scalar-mean approximation.")

    # =========================================================================
    # Wall-clock timing (post-warmup)
    # =========================================================================
    log("=" * 70)
    log("Wall-clock timing (100 warm iters each)")
    log("=" * 70)

    # Warmup
    for _ in range(3):
        _ = kda_recurrent_fwd(q_dev[:S], k_dev[:S], v_dev[:S], g_step[:S].to("neuron"), beta_dev[:S])
    try:
        torch.neuron.synchronize()
    except Exception:
        pass

    N = 100
    t0 = time.time()
    for _ in range(N):
        _ = kda_chunk_step(q_dev, k_dev, v_dev, beta_dev, gc_dev, gl_dev, state_dev)
    try:
        torch.neuron.synchronize()
    except Exception:
        pass
    chunk_us = (time.time() - t0) / N * 1e6
    log(f"  kda_chunk_step (C=128): {chunk_us:.1f} μs per call")
    log(f"    Per-token effective: {chunk_us / 128:.3f} μs/token")

    log("Done.")


if __name__ == "__main__":
    main()