File size: 8,707 Bytes
dab7980 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | """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()
|