WIP: add nkigen-lite as a standalone IR-based kernel generation backend

docs/dev_guide/running_on_hardware.md

@@ -0,0 +1,112 @@
+# Running on Neuron Hardware
+
+Practical runtime tricks for running tests, benchmarks, and examples on a
+multi-core Neuron instance (e.g. `trn2.48xlarge`, which exposes 16 NeuronCores
+across 4 devices). These are mostly about **NeuronCore selection** — picking
+which cores a process runs on so that concurrent jobs don't collide.
+
+## Inspecting the device
+
+`neuron-ls` shows every Neuron device, its cores, and which PIDs currently hold
+them:
+
+```bash
+neuron-ls
+```
+
+The `NEURON CORE IDS` and `PID` columns tell you which cores are busy. A core
+with a `PID` is in use; an idle core is free to claim. Use this before starting
+a job to find free cores and avoid contending with someone else's run.
+
+## Selecting specific cores
+
+`NEURON_RT_VISIBLE_CORES` restricts a process to a subset of cores. It accepts a
+single index or an inclusive range:
+
+```bash
+# Run on core 0 only
+NEURON_RT_VISIBLE_CORES=0 python my_script.py
+
+# Run on cores 4,5,6,7 (e.g. a TP=4 job on the second device)
+NEURON_RT_VISIBLE_CORES=4-7 torchrun --nproc-per-node 4 my_script.py
+```
+
+This is the key to running two jobs at once without them fighting over cores —
+for example benchmarking two branches side by side: pin one to `0-3` and the
+other to `4-7`. They land on different physical devices and run concurrently.
+
+`NEURON_RT_NUM_CORES=N` is a related knob that just requests `N` cores without
+pinning specific indices; prefer `NEURON_RT_VISIBLE_CORES` when you need
+determinism about *which* cores are used.
+
+## How the test suite isolates cores
+
+`tests/conftest.py` already handles core isolation for `pytest-xdist` parallel
+runs, so you normally don't set anything by hand:
+
+- `pytest_xdist_auto_num_workers` caps `-n auto` at the number of visible cores
+  (`Spike.get_visible_neuron_core_count()`), since each worker needs its own
+  core. Asking for more workers than cores causes allocation failures.
+- `pytest_configure` maps each xdist worker (`gw0`, `gw1`, …) to one core by
+  setting `NEURON_RT_VISIBLE_CORES=<worker index>` and `NEURON_RT_NUM_CORES=1`.
+
+Implications when running tests:
+
+```bash
+# Let xdist auto-size to the available cores (recommended)
+uv run pytest tests/ -n auto
+
+# To leave some cores free for another job, run on a subset and cap workers
+NEURON_RT_VISIBLE_CORES=0-3 uv run pytest tests/ -n 4
+```
+
+With `-n auto`, if some cores are already busy you should still constrain
+`NEURON_RT_VISIBLE_CORES` to the free ones, otherwise a worker may try to claim
+an occupied core.
+
+## Benchmarking a specific branch (git worktrees)
+
+When comparing branches, a git worktree gives each branch its own checkout. Two
+gotchas:
+
+1. **Make sure you import the code under test.** An editable install
+   (`nkigen_lite`, etc.) may resolve to the *main* checkout, not the worktree.
+   Put the worktree's source first on `PYTHONPATH` so the process picks it up:
+
+   ```bash
+   PYTHONPATH=/path/to/worktree/nkigen-lite/src:$PYTHONPATH python -c \
+     "import nkigen_lite, os; print(os.path.dirname(nkigen_lite.__file__))"
+   ```
+
+   Verify the printed path points into the worktree before trusting the numbers.
+
+2. **Pin the two runs to different cores** so they can run at the same time
+   under the same host/thermal conditions (a fairer comparison than
+   back-to-back runs):
+
+   ```bash
+   # branch A, main checkout
+   NEURON_RT_VISIBLE_CORES=0-3 torchrun --nproc-per-node 4 evaluate.py --benchmark ...
+
+   # branch B, worktree (different cores + worktree on PYTHONPATH)
+   NEURON_RT_VISIBLE_CORES=4-7 \
+   PYTHONPATH=/path/to/worktree/nkigen-lite/src:$PYTHONPATH \
+     torchrun --nproc-per-node 4 evaluate.py --benchmark ...
+   ```
+
+You can inspect the environment of a running process to confirm what it actually
+picked up:
+
+```bash
+cat /proc/<pid>/environ | tr '\0' '\n' | grep -E 'NEURON_RT_VISIBLE_CORES|PYTHONPATH'
+```
+
+## Deprecated / ignored runtime env vars
+
+Neuron Runtime evolves; some env vars are silently ignored on newer runtimes.
+Watch the startup logs for `WARN NRT:nrt_config_parse_init_config` lines. For
+example, `NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS` is no longer supported on
+Neuron Runtime 2.0 ("implicit async exec mode has been removed") — setting it
+has no effect, and async now requires the explicit Neuron Runtime async APIs. If
+a perf knob doesn't move the numbers, check whether the runtime logged that it
+ignored it.

docs/index.md

@@ -44,6 +44,7 @@ user_guide/distributed_execution
 
 dev_guide/extending_language
 dev_guide/testing
+dev_guide/running_on_hardware
 dev_guide/building_docs
 api/index
 ```

examples/models/qwen3/dump_layer_ir.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Dump the full generated tensor_ir and nki_ir for one fused Qwen3-30B-A3B
+MoE transformer layer (prefill), TP=4 per-rank shapes.
+
+Companion to dump_nki_ir.py (per building-block pattern) — this one is the WHOLE
+fused layer, so it captures the fully-unrolled MoE expert loop
+(`for b: for t: for e in top_k`), which the per-pattern dumps miss and which the
+performance ledger flags as the dominant 30B bottleneck.
+
+The MoE loop unrolls to B*L*top_k feed-forward invocations at trace time, so
+this is run at a SHORT context length (default L=8) to stay tractable; op counts
+scale ~linearly in L. Set QWEN3_DUMP_SEQ to change it.
+
+    QWEN3_BACKEND=nkigen-lite uv run python dump_layer_ir.py
+"""
+
+import os
+from collections import Counter
+from unittest import mock
+
+import numpy as np
+import torch.distributed as dist
+
+WORLD_SIZE = 4
+mock.patch.object(dist, "is_initialized", lambda: True).start()
+mock.patch.object(dist, "get_world_size", lambda *a, **k: WORLD_SIZE).start()
+mock.patch.object(dist, "get_rank", lambda *a, **k: 0).start()
+
+from config import Config  # noqa: E402
+from kernels.transformer_layer import transformer_layer  # noqa: E402
+
+from nkipy.core.trace import NKIPyKernel  # noqa: E402
+from nkigen_lite.tensor_ir.passes import lower_to_nki  # noqa: E402
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "nki_ir_dumps")
+DT = np.dtype("float32")
+
+# Qwen3-30B-A3B per-rank (TP=4) dims, matching test_oom_repro.py.
+HIDDEN = 2048
+HEAD_DIM = 128
+NUM_HEADS = 32
+NUM_KV_HEADS = 4
+N_EXPERTS = 128
+TOP_K = 8
+INTERMEDIATE = 192
+QKV_OUT = 1280
+O_IN = 1024
+SEQ = int(os.environ.get("QWEN3_DUMP_SEQ", "8"))
+
+
+def _z(shape):
+    return np.zeros(shape, dtype=DT)
+
+
+def main():
+    os.makedirs(OUT_DIR, exist_ok=True)
+    cfg = Config(
+        hidden_size=HIDDEN,
+        num_heads=NUM_HEADS,
+        head_dim=HEAD_DIM,
+        num_kv_heads=NUM_KV_HEADS,
+        num_layers=1,
+        num_experts_per_tok=TOP_K,
+        num_experts=N_EXPERTS,
+        context_len=SEQ,
+        max_new_tokens=4,
+        intermediate_size=INTERMEDIATE,
+    )
+    n_local_kv = max(1, NUM_KV_HEADS // WORLD_SIZE)
+    cache = _z((1, cfg.max_seq_len, n_local_kv, HEAD_DIM))
+
+    arrays = dict(
+        x=_z((1, SEQ, HIDDEN)),
+        start_pos=None,
+        qkv_weight=_z((HIDDEN, QKV_OUT)),
+        o_weight=_z((O_IN, HIDDEN)),
+        input_weight=_z((HIDDEN,)),
+        q_norm_weight=_z((HEAD_DIM,)),
+        k_norm_weight=_z((HEAD_DIM,)),
+        post_attention_weight=_z((HIDDEN,)),
+        router_weight=_z((HIDDEN, N_EXPERTS)),
+        gate_up_weight=_z((N_EXPERTS, HIDDEN, 2 * INTERMEDIATE)),
+        down_weight=_z((N_EXPERTS, INTERMEDIATE, HIDDEN)),
+        cache_k=cache,
+        cache_v=cache.copy(),
+        configs=cfg,
+    )
+
+    print(f"Tracing fused MoE transformer layer (L={SEQ}, top_k={TOP_K}) ...")
+    k_snap = NKIPyKernel.trace(transformer_layer, backend="nkigen-lite")
+    tg = k_snap.specialize(**arrays)._graph
+    tensor_dump = tg.dump()
+    tensor_hist = Counter(op.opcode for op in tg.ops)
+    print(f"  tensor_ir ops: {len(tg.ops)}")
+
+    kernel = NKIPyKernel.trace(transformer_layer, backend="nkigen-lite")
+    tensor_graph = kernel.specialize(**arrays)._graph
+    print("Lowering tensor_ir -> nki_ir ...")
+    nki = lower_to_nki(tensor_graph)
+    nki_hist = Counter(o.opcode for o in nki.ops)
+    print(f"  nki_ir ops: {len(nki.ops)}  (expansion {len(nki.ops)/len(tg.ops):.1f}x)")
+
+    ti_path = os.path.join(OUT_DIR, "transformer_layer.tensor_ir")
+    with open(ti_path, "w") as f:
+        f.write(f"# Fused Qwen3-30B-A3B MoE transformer layer (TP=4, L={SEQ}) "
+                f"— tensor_ir\n")
+        f.write(f"# total tensor_ir ops: {len(tg.ops)}\n\n")
+        f.write("=" * 70 + "\n## opcode histogram\n" + "=" * 70 + "\n")
+        for op, n in tensor_hist.most_common():
+            f.write(f"  {op:24s} {n}\n")
+        f.write("\n" + "=" * 70 + "\n## tensor_ir (full)\n" + "=" * 70 + "\n")
+        f.write(tensor_dump + "\n")
+    print(f"  wrote {ti_path}")
+
+    nki_path = os.path.join(OUT_DIR, "transformer_layer.nki")
+    with open(nki_path, "w") as f:
+        f.write(f"# Fused Qwen3-30B-A3B MoE transformer layer (TP=4, L={SEQ}) "
+                f"— nki_ir\n")
+        f.write(f"# tensor_ir ops: {len(tg.ops)}   nki_ir ops: {len(nki.ops)}"
+                f"   (expansion {len(nki.ops)/len(tg.ops):.1f}x)\n")
+        f.write(f"# NOTE: MoE loop = B*L*top_k = 1*{SEQ}*{TOP_K} feed-forward "
+                f"invocations, fully unrolled; scales ~linearly in L.\n\n")
+        f.write("=" * 70 + "\n## nki_ir opcode histogram\n" + "=" * 70 + "\n")
+        for op, n in nki_hist.most_common():
+            f.write(f"  {op:24s} {n}\n")
+        f.write("\n" + "=" * 70 + "\n## nki_ir (full)\n" + "=" * 70 + "\n")
+        f.write(nki.dump() + "\n")
+    print(f"  wrote {nki_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/models/qwen3/dump_nki_ir.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Dump the lowered nki_ir for each Qwen3 building-block pattern, for inspection.
+
+Traces each kernel, lowers tensor_ir -> nki_ir, and writes per pattern into
+nki_ir_dumps/<name>.nki: the tensor_ir graph, an nki opcode histogram, and the
+full nki_ir text dump. Lets us examine exactly what instructions each pattern
+expands into (and where the HBM round-trips / scaffolding are).
+
+This mirrors qwen3_embedding/dump_nki_ir.py but uses the Qwen3-30B-A3B MoE
+building blocks and per-rank (TP=4) shapes. torch.distributed is mocked to a
+TP=4 world (single process) so the head splits and intermediate sizes match a
+real shard, exactly like test_oom_repro.py.
+
+    uv run python dump_nki_ir.py
+"""
+
+import os
+from collections import Counter
+from unittest import mock
+
+import numpy as np
+import torch.distributed as dist
+
+# The kernels divide head counts / hidden by world_size, so the dump must see
+# the same TP=4 world the real shards are built for (see test_oom_repro.py).
+WORLD_SIZE = 4
+mock.patch.object(dist, "is_initialized", lambda: True).start()
+mock.patch.object(dist, "get_world_size", lambda *a, **k: WORLD_SIZE).start()
+mock.patch.object(dist, "get_rank", lambda *a, **k: 0).start()
+
+from kernels.rmsnorm import rmsnorm_kernel  # noqa: E402
+from kernels.softmax import softmax_kernel  # noqa: E402
+from kernels.feedforward import feedforward_kernel, silu_kernel_  # noqa: E402
+from kernels.rope import apply_rotary_emb_kernel  # noqa: E402
+
+from nkipy.core.trace import NKIPyKernel  # noqa: E402
+from nkigen_lite.tensor_ir.passes import lower_to_nki  # noqa: E402
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "nki_ir_dumps")
+DT = np.dtype("float32")  # f32 so both backends share one numerics path
+NORM_EPS = 1e-6
+
+# Qwen3-30B-A3B dims (global; kernels divide head counts / intermediate by TP).
+HIDDEN = 2048
+HEAD_DIM = 128
+NUM_HEADS = 32          # -> 8 local q heads at TP=4
+NUM_KV_HEADS = 4        # -> 1 local kv head at TP=4
+INTER = 768 // WORLD_SIZE  # per-rank expert intermediate = 192
+SEQ = 128
+
+N_LOCAL_HEADS = NUM_HEADS // WORLD_SIZE        # 8
+N_LOCAL_KV_HEADS = max(1, NUM_KV_HEADS // WORLD_SIZE)  # 1
+
+# Shapes for standalone building blocks.
+HID = (1, SEQ, HIDDEN)                         # hidden states for norm / ffn
+SCORE = (1, N_LOCAL_HEADS, SEQ, SEQ)           # attention scores for softmax
+QSHAPE = (1, SEQ, N_LOCAL_HEADS, HEAD_DIM)     # query post-reshape (BSHD)
+KSHAPE = (1, SEQ, N_LOCAL_KV_HEADS, HEAD_DIM)  # key post-reshape (BSHD)
+
+
+def matmul_gup(x, w):
+    return np.matmul(x, w)
+
+
+def dump(name, fn, arrays):
+    k = NKIPyKernel.trace(fn, backend="nkigen-lite")
+    ir = k.specialize(**arrays)
+    tensor_graph = ir._graph
+    # lower_to_nki mutates the tensor graph in place (canonicalize/decompose),
+    # so snapshot its dump from a fresh trace for the "tensor_ir" section.
+    k2 = NKIPyKernel.trace(fn, backend="nkigen-lite")
+    tg = k2.specialize(**arrays)._graph
+    tensor_dump = tg.dump()
+
+    nki = lower_to_nki(tensor_graph)
+    hist = Counter(o.opcode for o in nki.ops)
+
+    path = os.path.join(OUT_DIR, f"{name}.nki")
+    with open(path, "w") as f:
+        f.write(f"# Pattern: {name}\n")
+        f.write(f"# inputs: " + ", ".join(
+            f"{key}={getattr(v, 'shape', v)}" for key, v in arrays.items()) + "\n")
+        f.write(f"# nki_ir total ops: {len(nki.ops)}\n\n")
+        f.write("=" * 70 + "\n## tensor_ir (after trace)\n" + "=" * 70 + "\n")
+        f.write(tensor_dump + "\n\n")
+        f.write("=" * 70 + "\n## nki_ir opcode histogram\n" + "=" * 70 + "\n")
+        for op, n in hist.most_common():
+            f.write(f"  {op:24s} {n}\n")
+        f.write("\n" + "=" * 70 + "\n## nki_ir (full)\n" + "=" * 70 + "\n")
+        f.write(nki.dump() + "\n")
+    print(f"  wrote {path}  ({len(nki.ops)} nki ops)")
+
+
+def main():
+    rng = np.random.default_rng(0)
+    os.makedirs(OUT_DIR, exist_ok=True)
+
+    dump("rmsnorm", rmsnorm_kernel,
+         {"x": rng.standard_normal(HID).astype(DT),
+          "weight": rng.standard_normal(HIDDEN).astype(DT), "eps": NORM_EPS})
+
+    dump("softmax", softmax_kernel,
+         {"x": rng.standard_normal(SCORE).astype(DT)})
+
+    dump("silu", silu_kernel_,
+         {"x": rng.standard_normal((1, SEQ, INTER)).astype(DT)})
+
+    # Single-expert feed-forward (one expert's gate_up / down shard).
+    dump("feedforward", feedforward_kernel,
+         {"x": rng.standard_normal(HID).astype(DT),
+          "gate_up_weight": rng.standard_normal((HIDDEN, 2 * INTER)).astype(DT),
+          "down_weight": rng.standard_normal((INTER, HIDDEN)).astype(DT)})
+
+    # QKV projection matmul (hidden -> per-rank qkv width).
+    qkv_out = (N_LOCAL_HEADS + 2 * N_LOCAL_KV_HEADS) * HEAD_DIM
+    dump("matmul_qkv", matmul_gup,
+         {"x": rng.standard_normal(HID).astype(DT),
+          "w": rng.standard_normal((HIDDEN, qkv_out)).astype(DT)})
+
+    # RoPE applied to q/k (post-reshape BSHD).
+    half = HEAD_DIM // 2
+    dump("rope", apply_rotary_emb_kernel,
+         {"xq": rng.standard_normal(QSHAPE).astype(DT),
+          "xk": rng.standard_normal(KSHAPE).astype(DT),
+          "freqs_cos": rng.standard_normal((SEQ, half)).astype(DT),
+          "freqs_sin": rng.standard_normal((SEQ, half)).astype(DT)})
+
+    print(f"\nDumps in {OUT_DIR}/")
+
+
+if __name__ == "__main__":
+    main()