feat: connect four-model RNGD producers to live ATOM HTTP PD

opt_prime/examples/pd/qwen3_rngd/README.md

@@ -0,0 +1,110 @@
+# 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 two scripts run a finite fixture; the experimental HTTP adapter is described below.
+
+## 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 is exposed by the explicit experimental `serve.py` adapter below;
+it is not automatically registered in npu_pd_bench/serve_batched. Keep model-specific math in this prefiller;
+the portable56-slot import succeeded without communication-library changes.
+
+
+## Experimental HTTP adapter and export improvement (#199)
+
+`serve.py` reuses the existing optimus-rm `SubprocessPrefillWorker`, ATOM batcher,
+`BatchedService` and `BatchedOwner`. `worker_python.sh` redirects only the expected
+prefill worker module launch to `worker.py`; model-specific math stays in
+`prefill.py`. This is an explicit experiment requiring the existing optimus-rm
+runtime on PYTHONPATH, not a new production backend in this checkout.
+
+Place PR195's `furiosa_prefill_python.sh` alongside `worker_python.sh` and use its
+single-card allocation/bootstrap instructions. Keep the task's own compiled
+artifact cache and validated interpreters. The HTTP listener binds only loopback.
+The service is fixed to Qwen3-0.6B, B1/context512/3KV blocks, no prefix caching.
+Prompts above `--width` (default128) are rejected before worker submission. Decode
+requests may queue; only one lane executes. SIGTERM drains queued work before
+closing the worker and slab. Disconnect cancellation and fault recovery have not
+been hardware-validated in this experiment.
+
+```bash
+python3 -m unittest discover -s . -p 'test_*.py'
+# In the writable scratch directory with the validated fixture/artifacts:
+"$ATOM_PYTHON" check_atom_container.py serve.py --cache-dir "$PWD/qwen3-197-atom"
+# Another process in the same allocated Pod:
+"$ATOM_PYTHON" bench_http.py --name pd
+"$ATOM_PYTHON" bench_stream.py --name pd
+# Stop/drain the service, then restart with --local-prefill for the same ATOM-only baseline.
+```
+
+The export fix replaces device-side `torch.stack(...).cpu()` with per-layer
+`.cpu()` followed by CPU stacking. The device graph and width remain unchanged.
+To reproduce the alternating A/B, save the old implementation in the scratch
+working directory (not over the new prefiller):
+
+```bash
+git show 353d15090aa7908e0bb2ea5422f047e59ea9273d:opt_prime/examples/pd/qwen3_rngd/prefill.py > prefill_baseline.py
+./furiosa_prefill_python.sh -m bench_export
+```
+
+Measured warm prefill: 728.435→55.823ms mean (nine pairs); all twelve paired K,V
+and full-logit comparisons exact. Detailed serving results and limitations are
+in `SERVING_EXPERIMENT.md` and `serving_results/`. The old roughly719ms measurement
+above is historical, before this export fix. Width32/64 were also measured but
+changed numerical values, so the verified improvement retains width128.
+
+`safe_open` must be iterated through `.keys()` in this SDK environment; the
+specific lint exemption preserves that behavior.

opt_prime/examples/pd/qwen3_rngd/SERVING_EXPERIMENT.md

@@ -0,0 +1,61 @@
+# Qwen3 RNGD → ATOM: host stacking and experimental HTTP serving
+
+- **Date**: 2026-10-23
+- **Task**: Qwen3 serving and export profiling (aicomp #199)
+- **Host / device**: RNGD+ATOM host, dedicated single-card device-plugin Pod; Furiosa 2026.4 fusion4/Mesh2; ATOM rebel-compiler 0.11.0, optimum-rbln 0.11.0.post1, torch 2.11.0+cpu
+- **Commit**: baseline 353d15090aa7908e0bb2ea5422f047e59ea9273d; experiment branch exp/199-qwen3-serving (56d9056 plus working changes before final commit)
+
+## Question
+
+Is the roughly 719 ms Qwen3 research prefill dominated by model computation or KV materialization, and does its improvement survive the existing batched HTTP service?
+
+## Setup
+
+Checkpoint Qwen/Qwen3-0.6B revision c1899de289a04d12100db370d81485cdf75e47ca. Three fixed chat prompts, thinking disabled, 19/19/22 tokens, fixed 16 output tokens for parity diagnostics. RNGD width128, 28 layers, 8 KV heads, dimension128. ATOM own B1/context512/3KV-block artifact from the previous experiment. Prefix caching disabled; B1 serializes three concurrent requests through the existing queue.
+
+Baseline stacks each K/V tensor on RNGD before copying. Improvement copies each layer to CPU and stacks there; model math, widths and sampler unchanged. Existing SubprocessPrefillWorker, AtomRblnEngine, BatchedOwner and HTTP handlers are reused through an explicit experimental launcher. No shared SDK files or other workloads changed. Host-mediated KV, no DMA or producer overlap.
+
+## Command
+
+Use PR195 isolated-container environment and bootstrap; working directory contains the Qwen3 example files and the previous trusted fixture. Set HOME to the SDK owner's home, offline HF cache, OMP_NUM_THREADS=4, MKL_NUM_THREADS=4, RBLN_DEVICES=0, RBLN_NPUS_PER_DEVICE=1, and PYTHONPATH to the existing optimus-rm, rngd_pd and hetero_p2p sources. FURIOSA_PYTHON and ATOM_PYTHON select the previously validated interpreters.
+
+```bash
+# From the checkout, save the exact pre-change implementation into the scratch directory:
+git show 353d15090aa7908e0bb2ea5422f047e59ea9273d:opt_prime/examples/pd/qwen3_rngd/prefill.py > prefill_baseline.py
+./furiosa_prefill_python.sh -m bench_export
+# Service on a dedicated Pod, loopback only. Baseline adds --local-prefill.
+"$ATOM_PYTHON" check_atom_container.py serve.py --cache-dir "$PWD/qwen3-197-atom"
+"$ATOM_PYTHON" bench_http.py --name pd
+"$ATOM_PYTHON" bench_stream.py --name pd
+```
+
+A/B export alternates baseline/optimized order over four rounds and three prompts; round0 excluded from latency means. Whole K,V,logits are compared bit-for-bit each pair. HTTP harness performs twelve sequential requests, three concurrent requests, an over-capacity HTTP400 check and a successful follow-up. Streaming harness performs nine requests plus EOS termination; first round excluded from warm means.
+
+## Results
+
+- Alternating A/B: baseline 728.435 ms mean (688.581–829.348), host stacking 55.823 ms (44.516–67.473), nine warm pairs; 12/12 complete K,V,logit comparisons exact including first round.
+- Phase diagnostic: device stack/export roughly 650 ms versus per-layer CPU export/host stack roughly 11 ms. Layer execution plus final hidden materialization roughly 5 ms; CPU output head roughly 26–40 ms. These are diagnostic wall times with fences, not isolated DMA bandwidth.
+- Width32/64 reduced the original stack cost but changed KV/logits numerically; not adopted. Width128 stays unchanged.
+- Optimized PD HTTP: 16/16 full output sequences equal ATOM-only; 16 remote prefills, zero local fallback. Warm sequential nonstream mean201.439 ms vs ATOM121.831 ms.
+- Warm streaming (six requests): optimized PD TTFT69.567 ms/E2E167.535 ms; ATOM TTFT15.883 ms/E2E128.254 ms.
+- Oversize129-token prompt rejectedHTTP400 before device work; next request succeeds. EOS termination passes both routes. Fixed-length diagnostics intentionally ignoreEOS and are not quality tests.
+
+Failures retained: first profiling script attempted to iterate safe_open directly (TypeError); changed to .keys(), with targeted SIM118 exemption. Initial experimental worker launch used a relative executable path and failed before worker startup; changed to an absolute sibling path. Parent resource_tracker cleaned the abandoned slab. No shared package modified.
+
+Raw logs/results: analysis/local_prefill/native_dma/general/serving199/. Exact local launch scripts retained with this experiment's artifacts; public reproduction uses role/environment variable names.
+
+### Pre-change HTTP baseline and final comparison
+
+| Route | Warm streaming TTFT (ms) | Warm streaming E2E,16 outputs (ms) | Warm nonstream E2E (ms) |
+|---|---:|---:|---:|
+| PD, device-stack baseline |755.042|905.176|887.234|
+| PD, host-stack improvement |69.567|167.535|201.439|
+| ATOM-only |15.883|128.254|121.831|
+
+Streaming means use six requests after excluding the first round; nonstream means use nine requests after excluding the first round. Routes ran sequentially (optimized PD, ATOM, old PD), not a randomized serving trial. No quiet-host guarantee or confidence intervals; the alternating in-process A/B supports the export mechanism independently. All16 nonstream token sequences match across all three routes. All three pass oversized-request rejection/recovery and EOS termination. The pre-change service uses exactly the old prefiller file in the same task-owned scratch directory, restored to the improved file afterwards.
+
+The improved PD first-token gap to ATOM is53.684ms; E2E gap39.281ms in this small streaming sample. Source export and CPU LM head now merit optimization; no producer-side overlap was added. Server drains were successful and the dedicated validation namespace was deleted; other workloads untouched.
+
+## Conclusion
+
+The largest original delay was device-side KV stacking/materialization, not inherent PD disaggregation or transformer-layer compute. Moving the stack to CPU yields about13× prefill speedup without numerical changes and enables the existing HTTP service. ATOM-only remains faster for these short B1 requests; this does not establish native Furiosa parity, throughput at high concurrency, or an advantage for PD. Remaining work: output-head cost, ATOM import/IPC and longer prompts; source-streaming overlap and KONI HTTP integration remain separate validations.

opt_prime/examples/pd/qwen3_rngd/bench_export.py

@@ -0,0 +1,52 @@
+import json
+import time
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from prefill import Prefill
+from safetensors import safe_open
+
+path = Path(snapshot_download("Qwen/Qwen3-0.6B", local_files_only=True))
+c = json.loads((path / "config.json").read_text())
+w = {}
+for shard in path.glob("*.safetensors"):
+    with safe_open(shard, framework="pt") as f:
+        for k in f.keys():  # noqa: SIM118 - safe_open is not iterable
+            w[k] = f.get_tensor(k).float()
+records = torch.load("qwen3_197_kv.pt", weights_only=True)
+from prefill_baseline import Prefill as Baseline
+
+rows = []
+e = Prefill(c, w, width=128, chip=0)
+for repeat in range(4):
+    for i, r in enumerate(records):
+        outputs = {}
+        times = {}
+        for name in (
+            ["baseline", "host_stack"]
+            if repeat % 2 == 0
+            else ["host_stack", "baseline"]
+        ):
+            t = time.perf_counter()
+            outputs[name] = (
+                Baseline.prefill(e, r["ids"])
+                if name == "baseline"
+                else e.prefill(r["ids"])
+            )
+            times[name] = time.perf_counter() - t
+        exact = [
+            torch.equal(a, b)
+            for a, b in zip(outputs["baseline"], outputs["host_stack"])
+        ]
+        assert all(exact), exact
+        row = {
+            "repeat": repeat,
+            "prompt": i,
+            "n": len(r["ids"]),
+            "seconds": times,
+            "kv_logits_exact": exact,
+        }
+        rows.append(row)
+        print(json.dumps(row), flush=True)
+Path("qwen3_199_export_ab.json").write_text(json.dumps(rows, indent=2) + "\n")

opt_prime/examples/pd/qwen3_rngd/bench_http.py

@@ -0,0 +1,68 @@
+import argparse
+import concurrent.futures
+import json
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import torch
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--name", required=True)
+a = ap.parse_args()
+base = "http://127.0.0.1:18197"
+records = torch.load("qwen3_197_kv.pt", weights_only=True)
+rows = []
+
+
+def request(i, repeat, extra=None):
+    body = {
+        "model": "Qwen/Qwen3-0.6B",
+        "prompt": records[i]["ids"],
+        "temperature": 0,
+        "max_tokens": 16,
+        "ignore_eos": True,
+    }
+    body.update(extra or {})
+    req = urllib.request.Request(
+        base + "/v1/completions",
+        json.dumps(body).encode(),
+        headers={
+            "Content-Type": "application/json",
+            "x-request-id": f"{a.name}-{repeat}-{i}",
+        },
+    )
+    t = time.perf_counter()
+    with urllib.request.urlopen(req, timeout=120) as r:
+        out = json.load(r)
+    row = {
+        "repeat": repeat,
+        "prompt": i,
+        "seconds": time.perf_counter() - t,
+        "response": out,
+    }
+    print(json.dumps(row), flush=True)
+    return row
+
+
+for repeat in range(4):
+    for i in range(3):
+        rows.append(request(i, repeat))
+with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
+    rows.extend(pool.map(lambda i: request(i, 4), range(3)))
+# Capacity must be rejected before submitting to the worker, then service must remain usable.
+try:
+    request(0, 5, {"prompt": [1] * 129})
+    raise AssertionError("oversize request unexpectedly accepted")
+except urllib.error.HTTPError as e:
+    assert e.code == 400, e.code
+    rejected = json.loads(e.read())
+rows.append(request(0, 6))
+with urllib.request.urlopen(base + "/metrics") as r:
+    metrics = r.read().decode()
+Path(f"qwen3_199_http_{a.name}.json").write_text(
+    json.dumps({"rows": rows, "oversize": rejected, "metrics": metrics}, indent=2)
+    + "\n"
+)
+print(metrics, flush=True)

opt_prime/examples/pd/qwen3_rngd/bench_stream.py

@@ -0,0 +1,74 @@
+import argparse
+import json
+import time
+import urllib.request
+from pathlib import Path
+
+import torch
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--name", required=True)
+a = ap.parse_args()
+records = torch.load("qwen3_197_kv.pt", weights_only=True)
+rows = []
+for repeat in range(3):
+    for i, r in enumerate(records):
+        payload = {
+            "model": "Qwen/Qwen3-0.6B",
+            "prompt": r["ids"],
+            "max_tokens": 16,
+            "temperature": 0,
+            "ignore_eos": True,
+            "stream": True,
+        }
+        req = urllib.request.Request(
+            "http://127.0.0.1:18197/v1/completions",
+            json.dumps(payload).encode(),
+            headers={
+                "Content-Type": "application/json",
+                "x-request-id": f"{a.name}-stream-{repeat}-{i}",
+            },
+        )
+        t = time.perf_counter()
+        first = None
+        done = False
+        chunks = []
+        with urllib.request.urlopen(req, timeout=120) as f:
+            for line in f:
+                if not line.startswith(b"data: "):
+                    continue
+                if line.strip() == b"data: [DONE]":
+                    done = True
+                    break
+                chunk = json.loads(line[6:])
+                chunks.append(chunk)
+                if (
+                    chunk.get("choices")
+                    and chunk["choices"][0].get("text")
+                    and first is None
+                ):
+                    first = time.perf_counter() - t
+        assert done and first is not None
+        row = {
+            "repeat": repeat,
+            "prompt": i,
+            "ttft_s": first,
+            "e2e_s": time.perf_counter() - t,
+            "text": "".join(
+                c["choices"][0].get("text", "") for c in chunks if c.get("choices")
+            ),
+        }
+        rows.append(row)
+        print(json.dumps(row), flush=True)
+payload.update(stream=False, ignore_eos=False, prompt=records[0]["ids"])
+req = urllib.request.Request(
+    "http://127.0.0.1:18197/v1/completions",
+    json.dumps(payload).encode(),
+    headers={"Content-Type": "application/json", "x-request-id": a.name + "-eos"},
+)
+with urllib.request.urlopen(req, timeout=120) as f:
+    eos = json.load(f)
+assert eos["choices"][0]["finish_reason"] == "stop", eos
+Path(f"qwen3_199_stream_{a.name}.json").write_text(
+    json.dumps({"rows": rows, "eos": eos}, indent=2) + "\n"
+)