Add experimental Qwen3 RNGD prefill and ATOM handoff validation

opt_prime/examples/pd/qwen3_rngd/README.md

@@ -0,0 +1,63 @@
+# Experimental Qwen3 RNGD prefill → ATOM validation (#197)
+
+Qwen3-0.6B requires Q/K normalization, unbiased projections, and a2048-wide query
+projection despite hidden_size1024. This separate graph implements those semantics;
+it does not loosen the existing Qwen2-only DenseModel contract or register a new
+production serving backend. Quantized weights, scaled RoPE, sliding attention and
+non-Qwen3 configs are rejected. Only Qwen3-0.6B has been hardware-validated here.
+
+## Reproduce
+
+Use the isolated single-card device-plugin Pod setup from PR195's
+`ATOM_CONTAINER_EXPERIMENT.md`, including validated Python environments, HOME,
+offline model cache, device leases, and the fusion4/Mesh2 Furiosa launcher. Requires
+existing `rngd_pd` and optimus-rm modules on PYTHONPATH, plus PR195's
+`atom_stream_upload.py` and `check_atom_container.py`. Do not install into or mutate
+another workload's environment. Work from a writable scratch directory containing
+these files. Use only the allocated devices.
+
+```bash
+python3 test_spec.py
+# Export FURIOSA_PYTHON to the validated Furiosa2026.4 interpreter.
+./furiosa_prefill_python.sh -m run_prefill
+"$ATOM_PYTHON" check_atom_container.py check_atom.py
+```
+
+`run_prefill.py` writes trusted tensor data to qwen3_197_kv.pt. `check_atom.py`
+compiles/loads its own B1/context512/3KV-block artifact under qwen3-197-atom, compares
+CPU FP32 prefill, native ATOM prefill and imported RNGD KV, and writes
+qwen3_197_atom.json. It uses both the existing and portable importers and reads back
+every imported slot. These scripts run a finite fixture, not an HTTP service.
+
+## Hardware evidence (2026-10-23)
+
+RNGD: Furiosa2026.4,fusion4/Mesh2,width128,28layers,8KV heads,head_dim128.
+ATOM: rebel-compiler0.11.0,optimum-rbln0.11.0.post1,torch2.11.0+cpu,
+B1/context512/3KV blocks. Checkpoint Qwen/Qwen3-0.6B; three chat prompts with thinking
+disabled. Exact records in validation.json; compilation provenance in the local
+research experiment2026-10-23-gemma-fix-qwen3-koni.md.
+
+-3/3first tokens match CPU FP32 and native ATOM.
+-3/3fixed16-token PD outputs match native ATOM.
+-3/3full decode-logit sequences match between original and portable importers.
+-336/336slot readbacks exact:3prompts ×2importers ×56slots.
+-RNGD-vs-CPU relative RMS:K0.470–0.491%,V1.757–1.911%,last logits1.164–1.604%.
+-Prefill wall time including host KV/logits:1.321s first,0.723/0.714s subsequent.
+  This is a correctness bring-up, not vendor-native performance parity.
+
+Fixed-length diagnostic generation continues beyondEOS; service termination semantics
+are not tested. The Korean-capital answer is incorrectlyIncheon on both CPU/ATOM in
+this fixture; matching tokens does not establish factual quality. No DMA, zero-copy,
+streamed source export or overlap claim. Host transfer occurs after complete prefill.
+
+The initial compile failed because epsilon/exponent inputs wereBF16 while rms_dsl
+requiresFP32. They now remainFP32; weights/activation carriers remainBF16. Failed
+compile log is retained with the experiment. CPU contract tests reject unsupported
+semantics; hardware records are essential evidence beyond those tests.
+
+## Limits
+
+Other sizes, longer prompts, concurrent requests, prefix reuse, cancellation and
+production recovery/lifetime behavior are unvalidated. The handoff has not been
+wired into npu_pd_bench/serve_batched. Keep model-specific math in this prefiller;
+the portable56-slot import succeeded without communication-library changes.

opt_prime/examples/pd/qwen3_rngd/check_atom.py

@@ -0,0 +1,111 @@
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from atom_stream_upload import import_atom_span
+from opt_prime.accel.platforms.atom.engine import AtomRblnEngine
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+mid = "Qwen/Qwen3-0.6B"
+physical = int(next(Path("/dev").glob("rbln[0-9]*")).name[4:])
+e = AtomRblnEngine(
+    devices=[physical],
+    max_seq_len=512,
+    batch_size=1,
+    kv_blocks=3,
+    cache_dir=str(Path("qwen3-197-atom").resolve()),
+)
+e.start(mid)
+records = torch.load("qwen3_197_kv.pt", weights_only=True)
+ref = AutoModelForCausalLM.from_pretrained(mid, dtype=torch.float32).eval()
+tok = AutoTokenizer.from_pretrained(mid)
+b = e.batcher(prefix_cache=False)
+rows = []
+
+
+def decode(first, n):
+    out = [first]
+    vectors = []
+    for j in range(15):
+        z = (
+            b.decoder(
+                input_ids=torch.tensor([[out[-1]]]),
+                cache_position=torch.tensor([[n + j]], dtype=torch.int32),
+                block_tables=b.blocks.table([0]),
+            )
+            .logits[0, -1]
+            .float()
+        )
+        assert torch.isfinite(z).all()
+        vectors.append(z.clone())
+        out.append(int(z.argmax()))
+    return out, torch.stack(vectors)
+
+
+with torch.inference_mode():
+    for r in records:
+        ids = r["ids"]
+        n = len(ids)
+        b.blocks.admit(0, ids)
+        block = b.blocks.lane_block[0]
+        nl = b._prefill(0, ids).logits[0, -1].float()
+        native, _ = decode(int(nl.argmax()), n)
+        hf = ref(torch.tensor([ids]), use_cache=True)
+        hl = hf.logits[0, -1].float()
+        ll = r["logits"].float()
+        ck = torch.stack([l.keys[0] for l in hf.past_key_values.layers])
+        cv = torch.stack([l.values[0] for l in hf.past_key_values.layers])
+        kr = float(
+            (r["keys"].float() - ck).square().mean().sqrt() / ck.square().mean().sqrt()
+        )
+        vr = float(
+            (r["values"].float() - cv).square().mean().sqrt()
+            / cv.square().mean().sqrt()
+        )
+        variants = {}
+        saved = []
+        for name in ["original", "portable"]:
+            if name == "original":
+                b.blocks.import_span(0, 0, r["keys"], r["values"])
+            else:
+                import_atom_span(b.blocks, 0, 0, r["keys"], r["values"])
+            exact = 0
+            for meta, buf in zip(b.blocks.metas, b.blocks._staging()):
+                expected = buf[block, :, :n].copy()
+                buf.fill(0)
+                b.blocks.io_rt.get_kv_cache_for_layer(meta["name"], block, 0, n, buf)
+                assert np.array_equal(expected, buf[block, :, :n])
+                exact += 1
+            tokens, vectors = decode(int(ll.argmax()), n)
+            saved.append(vectors)
+            variants[name] = {
+                "tokens": tokens,
+                "text": tok.decode(tokens),
+                "readback_exact_slots": exact,
+                "native_equal": tokens == native,
+            }
+        row = {
+            "prompt": r["prompt"],
+            "prompt_tokens": n,
+            "first_cpu": int(hl.argmax()),
+            "first_rngd": int(ll.argmax()),
+            "first_atom": int(nl.argmax()),
+            "key_relative_rms": kr,
+            "value_relative_rms": vr,
+            "logits_relative_rms": float(
+                (ll - hl).square().mean().sqrt() / hl.square().mean().sqrt()
+            ),
+            "prefill_s": r["prefill_s"],
+            "native_tokens": native,
+            "native_text": tok.decode(native),
+            "variants": variants,
+            "importer_logits_exact": bool(torch.equal(*saved)),
+        }
+        rows.append(row)
+        print(json.dumps(row, ensure_ascii=False), flush=True)
+        b.blocks.release(0, None)
+Path("qwen3_197_atom.json").write_text(
+    json.dumps(rows, ensure_ascii=False, indent=2) + "\n"
+)
+e.close()

opt_prime/examples/pd/qwen3_rngd/prefill.py

@@ -0,0 +1,196 @@
+"""Experimental Qwen3 full-prefix graph; requires validated Furiosa environment."""
+
+import math
+
+import torch
+from rngd_pd.engines.fsdk import (
+    CompileModule,
+    FCConfig,
+    TacticKernelModule,
+    dev_str,
+    hint,
+    sync,
+)
+from rngd_pd.kernels.prefill_chunk import rms_dsl, softmax_dsl, swiglu_dsl, zipadd_dsl
+from spec import validate
+
+BF = torch.bfloat16
+
+
+class Layer(torch.nn.Module):
+    def __init__(self, s, width):
+        super().__init__()
+        self.s = s
+        self.width = width
+        h = s["hidden_size"]
+        d = s["head_dim"]
+        nh = s["num_attention_heads"]
+        nk = s["num_key_value_heads"]
+        self.norm = TacticKernelModule(rms_dsl(width, h))
+        self.add = TacticKernelModule(zipadd_dsl(width, h))
+        self.qnorm = TacticKernelModule(rms_dsl(width * nh, d))
+        self.knorm = TacticKernelModule(rms_dsl(width * nk, d))
+        self.act = TacticKernelModule(swiglu_dsl(width, s["intermediate_size"]))
+        self.softmax = TacticKernelModule(
+            softmax_dsl(nk, (nh // nk) * width, width).replace("/bf16", "/f32")
+        )
+
+    def forward(
+        self,
+        x,
+        eps,
+        exponent,
+        qeps,
+        qexp,
+        keps,
+        kexp,
+        n1,
+        n2,
+        qn,
+        kn,
+        wq,
+        wk,
+        wv,
+        wo,
+        wg,
+        wu,
+        wd,
+        rotation,
+        mask,
+    ):
+        s = self.s
+        c = self.width
+        d = s["head_dim"]
+        nh = s["num_attention_heads"]
+        nk = s["num_key_value_heads"]
+        g = nh // nk
+        h = self.norm(x, n1, eps, exponent)[0]
+        q = self.qnorm((h @ wq).reshape(c * nh, d), qn, qeps, qexp)[0].reshape(c, nh, d)
+        k = self.knorm((h @ wk).reshape(c * nk, d), kn, keps, kexp)[0].reshape(c, nk, d)
+        q = torch.matmul(q, rotation).to(BF)
+        k = torch.matmul(k, rotation).to(BF)
+        keys = k.permute(1, 2, 0)
+        values = (h @ wv).reshape(c, nk, d).permute(1, 0, 2)
+        query = q.reshape(c, nk, g, d).permute(1, 2, 0, 3).reshape(nk, g * c, d)
+        probs = self.softmax((query.float() @ keys.float()) * (d**-0.5) + mask.float())[
+            0
+        ].to(BF)
+        att = (
+            (probs @ values).reshape(nk, g, c, d).permute(2, 0, 1, 3).reshape(c, nh * d)
+        )
+        res = self.add(att @ wo, x)[0]
+        h = self.norm(res, n2, eps, exponent)[0]
+        out = self.add(self.act(h @ wg, h @ wu)[0] @ wd, res)[0]
+        return out, torch.stack((keys, values.transpose(-1, -2)))
+
+
+class Prefill:
+    def __init__(self, config, weights, width=128, chip=0):
+        s = self.s = validate(config)
+        self.width = width
+        self.chip = chip
+        if width <= 0 or width > 512 or width % 32:
+            raise ValueError("width must be a multiple of32 within512")
+        device = dev_str(chip)
+        self.device = device
+
+        def put(x):
+            return x.contiguous().to(BF).to(device)
+
+        self.embedding = weights["model.embed_tokens.weight"]
+        h = s["hidden_size"]
+        d = s["head_dim"]
+        nh = s["num_attention_heads"]
+        nk = s["num_key_value_heads"]
+        eps = s["rms_norm_eps"]
+        c = width
+        self.fixed = [
+            torch.full((r, col), v, dtype=torch.float32).to(device)
+            for r, col, v in [
+                (c, h, eps * h),
+                (c, h, -0.5),
+                (c * nh, d, eps * d),
+                (c * nh, d, -0.5),
+                (c * nk, d, eps * d),
+                (c * nk, d, -0.5),
+            ]
+        ]
+        self.layers = []
+        for i in range(s["num_hidden_layers"]):
+
+            def w(k, i=i):
+                return weights[f"model.layers.{i}." + k]
+
+            norms = [
+                put(w(k).repeat(n, 1) * math.sqrt(dim))
+                for k, n, dim in [
+                    ("input_layernorm.weight", c, h),
+                    ("post_attention_layernorm.weight", c, h),
+                    ("self_attn.q_norm.weight", c * nh, d),
+                    ("self_attn.k_norm.weight", c * nk, d),
+                ]
+            ]
+            mats = [
+                put(w(k).T)
+                for k in [
+                    "self_attn.q_proj.weight",
+                    "self_attn.k_proj.weight",
+                    "self_attn.v_proj.weight",
+                    "self_attn.o_proj.weight",
+                    "mlp.gate_proj.weight",
+                    "mlp.up_proj.weight",
+                    "mlp.down_proj.weight",
+                ]
+            ]
+            self.layers.append(norms + mats)
+        inv = 1 / (s["rope_theta"] ** (torch.arange(0, d, 2).float() / d))
+        angle = torch.arange(c)[:, None] * inv[None]
+        rot = torch.zeros(c, d, d)
+        idx = torch.arange(d // 2)
+        rot[:, idx, idx] = angle.cos()
+        rot[:, idx + d // 2, idx + d // 2] = angle.cos()
+        rot[:, idx, idx + d // 2] = angle.sin()
+        rot[:, idx + d // 2, idx] = -angle.sin()
+        self.rotation = put(rot)
+        self.mask = put(
+            torch.triu(torch.full((c, c), -1e4), 1).repeat(nh // nk, 1)[None]
+        )
+        tensors = [
+            torch.zeros(c, h, dtype=BF),
+            *self.fixed,
+            *self.layers[0],
+            self.rotation,
+            self.mask,
+        ]
+        sample = tuple(torch.zeros(tuple(t.shape), dtype=t.dtype) for t in tensors)
+        self.module = CompileModule.from_module(
+            Layer(s, c),
+            sample,
+            compiler_config=FCConfig(tactic_hint=hint("ForLlmModelPrefill")),
+        )
+        self.module.to(device)
+        sync(chip)
+        self.norm = weights["model.norm.weight"].float()
+        self.head = weights.get("lm_head.weight", self.embedding).float()
+
+    def prefill(self, tokens):
+        if not 0 < len(tokens) <= self.width:
+            raise ValueError("prompt exceeds compiled capacity")
+        x = (
+            self.embedding[tokens + [tokens[-1]] * (self.width - len(tokens))]
+            .to(BF)
+            .to(self.device)
+        )
+        keys = []
+        values = []
+        for w in self.layers:
+            x, kv = self.module(x, *self.fixed, *w, self.rotation, self.mask)
+            keys.append(kv[0])
+            values.append(kv[1].transpose(-1, -2))
+        sync(self.chip)
+        n = len(tokens)
+        k = torch.stack(keys).cpu()[:, :, :, :n].transpose(-1, -2)
+        v = torch.stack(values).cpu()[:, :, :n]
+        z = x.cpu()[n - 1].float()
+        z = z * torch.rsqrt(z.square().mean() + self.s["rms_norm_eps"]) * self.norm
+        return k, v, z @ self.head.T