| """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.""" |
| |
| 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 |
| 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") |
|
|
| |
| |
| |
| 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() |
|
|
| |
| 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: |
| |
| |
| |
| ref_o, _ = fla_recurrent( |
| q.unsqueeze(0).unsqueeze(2), |
| k.unsqueeze(0).unsqueeze(2), |
| v.unsqueeze(0).unsqueeze(2), |
| g.unsqueeze(0).unsqueeze(2), |
| beta_row.squeeze(-1).unsqueeze(0).unsqueeze(2), |
| scale=1.0, |
| ) |
| 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.") |
|
|
| |
| |
| |
| 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.") |
|
|
| |
| |
| |
| log("=" * 70) |
| log("Wall-clock timing (100 warm iters each)") |
| log("=" * 70) |
|
|
| |
| 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() |
|
|
