GPU resident pipeline: batch CUDA attention, head-sharded kv_b, prefill expert groups, W4A16 mixed dispatch (#111)

* Fuse CUDA expert MLP execution

* Group CUDA expert transfers by device

* Instrument grouped CUDA expert execution

* Bound grouped CUDA decode scratch

* Execute expert groups across GPUs in parallel

* Release host backing for multi-GPU experts

* Define quality-preserving memory policies

* Overlap cold expert loading with resident compute

* Adapt expert placement with session LFRU

* Fuse q4 expert gate and up dispatch

* Plan CPU work on physical cores

* Batch grouped expert CUDA kernels

* Separate VRAM and RAM expert placement

* Add ragged multi-sequence decode forward

* feat(runtime): add continuous decode scheduler

* Route concurrent API requests through batch scheduler

* Harden multiplex request lifecycle and framing

* Cancel disconnected multiplex requests

* Bind API port before starting the engine

* fix automatic KV slot allocation

* add native int4 Tensor Core grouped GEMM

* add Tensor Core throughput benchmark

* optimize packed int4 low-row kernels

* add asynchronous CUDA staging streams

* document validated six-GPU dense acceleration

* tune six-GPU expert hot set

* raise validated expert hot-set target

* add CUDA MLA absorption core

* fuse grouped expert gate and up projections

* Warn for explicit lossy routing flags

* Add full-resident expert placement mode

* Adapt VRAM expert slots to live routes

* Accelerate int4 matvec on AVX-512

* Reduce AVX-512 and RoPE decode overhead

* Seed every GPU expert layer after prefill

* Limit live GPU swaps during decode

* CUDA batch MLA attention, kv_b head-sharding, fused o_proj, expert-group dispatch, W4A16 kernels

Lab-qualified on the 6x RTX 5090 machine (914-token request benchmark):
- batch MLA absorption kernel (COLI_CUDA_ATTN=1): whole-batch attention on
  device, 154.8s -> 102.4s
- attention -> o_proj fusion on the layer device: -> 97.4s
- kv_b head-sharding across cards (COLI_CUDA_ATTN_SHARD=1), no weight
  duplication: -> 94.05s
- per-device expert-group dispatch with pinned-buffer async transfers,
  W4A16 tensor-core kernels for the shared expert, OMP hot-thread tuning

Negative results (reverted, kept out): GPU-side weighted scatter-add
(atomics + per-layer D2H lose 43.8%), shared-expert fused small-batch
kernel (-38.8%), W4A4 grouped tensor cores (int4 activations corrupt
output). Details in the lab research log.

* GPU resident pipeline: device-resident prefill attention chain, GPU expert groups in prefill, batched router, W4A16 mixed dispatch

COLI_CUDA_PIPE=1 keeps the prefill data plane on the layer home device;
control flow (routing, cache/pin management) stays on CPU. Any CUDA
failure falls back to the unchanged CPU path.

- Device primitives + unit tests (tests/test_pipe_cuda.cu): rmsnorm
  (strided), interleaved RoPE, silu-mul, residual add, fixed-order row
  merge (no atomics), device-input GEMM, persistent per-device scratch.
  All verified against the engine's CPU math on SM120 (worst 1.2e-5).
- attn_pipe_prefill: q_a -> norm -> q_b -> rope -> kv_a -> norm -> rope ->
  batch attention -> o_proj in one device chain (q_a/q_b/kv_a colocated
  with kv_b); only the final [S,D] and the new KV rows return to host.
  Attention 41.2s -> 30.8s on the 1571-token benchmark.
- Prefill batch-union now uses the GPU expert groups (previously gated to
  S<=64, leaving all VRAM-resident experts idle during prefill - measured
  21ms of GPU expert time in a 148s prefill). Expert phase 78.9s -> 69.0s.
- Router computed as one batched matmul instead of S sequential rows
  (bit-identical math).
- W4A16 tensor-core path for expert groups (COLI_CUDA_TC_W4A16=1) with
  row-count mixed dispatch: >=16 rows per expert use tensor cores, smaller
  batches keep the naive kernel (tensor cores measured negative below
  ~16 rows). Expert phase 69.0s -> 64.3s, decode unaffected.

Net on the 1571-token prefill benchmark: 148.8s -> 114.3-126.8s
(component timings stable across runs; wall drifts +-3-5s because
.coli_usage placement learning shifts the expert tiers between runs).
PROFILO now also prints the prefill-phase breakdown.

* Skip OMP hot-thread tuning when CUDA is enabled

The active-spin worker team measured 66.9s->20.9s on the CPU-only Zen5
build, but on the six-GPU full-residency workload the spinning workers
contend with the CUDA dispatch threads: ~4x slower prefill with the
process stuck near 1.8 cores. Gate the tuning on COLI_CUDA so each
configuration keeps the behavior it was measured to prefer.

* Inc.2a: sparse layers fully resident on the layer device, residual hops cards at layer boundaries

COLI_CUDA_PIPE=2 keeps the residual stream on the layer home device for
consecutive sparse layers (cudaMemcpyPeer at boundaries): in/post norms,
attention chain, both residual adds and the shared-expert MLP run on
device. Per layer only the post-norm activations (router + CPU-tier
experts + group gather), the new KV rows and, on DSA indexer layers, the
pre-attention norm leave the card. Per-layer transfers drop from ~130MB
to ~70MB. A device-side snapshot at layer entry makes any mid-layer CUDA
failure fall back to the unchanged CPU path idempotently.

1571-token prefill: 127.1s (PIPE=1 control) -> 117.6/118.9s, components
attention 30.8->26.1, other 31.8->22.5-24.5; output verified coherent
against the control.

* Head-sharded attention inside the pipe: negative on PCIe star topology, gated opt-in

Slicing q per card from the home device and collecting ctx back
serializes ~95MB/layer through the home card's PCIe link: attention
26.1s -> 41.4/44.4s on the 1571-token benchmark (two repeats), wall
117.6 -> 135-138s. The standalone host-path sharding won because six
cards uploaded from host RAM in parallel; a home-device star has no
such parallelism without NVLink. Kept behind COLI_CUDA_PIPE_SHARD=1
for interconnects where peer bandwidth does not share one root port.

* Inc.3: device-resident KV shadow for decode attention

Decode re-uploaded the whole latent+rope window per layer per token
(~300MB/token at 1571 context). Each layer now keeps a device shadow of
the compressed KV on its kv_b card, bulk-synced when behind and appended
incrementally; the host cache stays canonical. Invalidation on kv_bind
(slot switch), kv_alloc (resize) and on any overwrite of mirrored rows,
with the legacy full-upload path as fallback.

Measured (COLI_CUDA_PIPE gate): short-context decode 5.48 -> 5.59/5.87
tok/s, 1571-context decode 4.14 -> 4.22 tok/s. Decode remains CPU-expert
bound; the shadow removes the transfer tax, not the compute.

* tools: unified user-experience benchmark (bench_ux.sh)

Two fixed scenarios (short chat, long-document QA), TTFT + decode tok/s
+ first-line drift check, TEMP=0 DRAFT=0 enforced, medians over REPS
runs. Encodes the measurement discipline from the lab record: same
binary per comparison, judge medians because .coli_usage placement
learning drifts wall times between runs.

* tools: bench_ux.sh executable bit

* gitignore compiled test binaries

* tools: expert_atlas.py — measure per-expert topic affinity (#175)

Diffs .coli_usage across 10 themed probe batches (code/math/chinese/
prose/science/law/poetry/structured/translation/casual, 3 prompts each)
driven through a running API server — one engine load total. Every
touched expert gets a topic-affinity vector, entropy, and a specialist/
generalist label; output experts.json feeds the Brain page hover.

* serve: persist .coli_usage after every turn in mux mode, not only at exit

run_serve_mux saved the learning cache once at shutdown; a crash lost
the whole session's routing history, and live consumers of the file
(expert_atlas.py diffs it between probe batches) saw a frozen snapshot.
Now saved per turn like the interactive path (165KB write, negligible).

* web: Brain hover shows measured expert atlas when published

If /experts.json (from tools/expert_atlas.py, #175) is served next to
the app, the tooltip upgrades from the depth heuristic to measured
data: specialist/generalist label, entropy, and the top-3 topic
affinities. Row index maps to real layer (row+3, last row = MTP 78).
Falls back to the heuristic when no atlas is published.

---------

Co-authored-by: JustVugg <JustVugg@users.noreply.github.com>

web/src/Brain.tsx

@@ -4,6 +4,7 @@ import { BrainCircuit, Flame, Layers } from "lucide-react"
 import { endpoint } from "@/lib/api"
 
 interface ExpertMap { rows: number; cols: number; map: string; hits: string; seq: number }
+interface AtlasEntry { affinity: Record<string, number>; entropy: number; top: string; label: string }
 
 const TIER_NAME = ["Disk", "RAM", "VRAM"]
 const TIER_RGB: [number, number, number][] = [[58, 71, 80], [90, 155, 216], [78, 214, 165]]
@@ -26,11 +27,19 @@ export function Brain({ baseUrl, apiKey, connected }: { baseUrl: string; apiKey:
   const wrapRef = useRef<HTMLDivElement>(null)
   const [wrapSize, setWrapSize] = useState({ w: 1200, h: 700 })
   const [data, setData] = useState<ExpertMap | null>(null)
+  const [atlas, setAtlas] = useState<Record<string, AtlasEntry> | null>(null)
   const [tip, setTip] = useState<{ x: number; y: number; row: number; col: number; tier: number; heat: number } | null>(null)
   const pulseRef = useRef<Float32Array | null>(null)   // per-expert pulse intensity 0..1
   const lastSeq = useRef(0)
   const rafRef = useRef(0)
 
+  // load the expert atlas if published (measured topic affinity, #175)
+  useEffect(() => {
+    fetch("/experts.json").then(r => r.ok ? r.json() : null).then(d => {
+      if (d?.experts) setAtlas(d.experts)
+    }).catch(() => {})
+  }, [])
+
   // track container size for responsive cell sizing
   useEffect(() => {
     const el = wrapRef.current
@@ -146,14 +155,26 @@ export function Brain({ baseUrl, apiKey, connected }: { baseUrl: string; apiKey:
         <canvas ref={canvasRef} onMouseMove={onMove} onMouseLeave={() => setTip(null)} />
         {!connected && <p className="runtime-unavailable">Connect to the engine to see the cortex.</p>}
       </div>
-      {tip && data && (
+      {tip && data && (() => {
+        const isMtp = tip.row === data.rows - 1
+        const realLayer = isMtp ? 78 : tip.row + 3
+        const entry = atlas?.[`${realLayer}:${tip.col}`]
+        return (
         <div className="brain-tip" style={{ left: tip.x + 14, top: tip.y + 14 }}>
-          <div className="brain-tip-title"><Layers className="size-3" /> Layer row {tip.row}{tip.row === data.rows - 1 ? " (MTP)" : ""} · Expert {tip.col}</div>
+          <div className="brain-tip-title"><Layers className="size-3" /> Layer {realLayer}{isMtp ? " (MTP)" : ""} · Expert {tip.col}</div>
           <div>Tier: <strong style={{ color: ["#8b9aa3", "#5a9bd8", "#4ed6a5"][tip.tier] }}>{TIER_NAME[tip.tier]}</strong></div>
           <div>Heat: <strong>{tip.heat === 0 ? "never routed" : `~2^${tip.heat} selections`}</strong></div>
-          <div className="brain-tip-role">{depthRole(tip.row, data.rows, tip.row === data.rows - 1)}</div>
+          {entry ? <>
+            <div className={entry.label.startsWith("specialist") ? "brain-tip-spec" : undefined}>
+              {entry.label.startsWith("specialist") ? `⭐ Specialist: ${entry.top}` : "Generalist"}
+              <small> (entropy {entry.entropy})</small>
+            </div>
+            <div className="brain-tip-aff">{Object.entries(entry.affinity).sort((a, b) => b[1] - a[1]).slice(0, 3)
+              .map(([c, p]) => `${c} ${Math.round(p * 100)}%`).join(" · ")}</div>
+          </> : <div className="brain-tip-role">{depthRole(tip.row, data.rows, isMtp)}</div>}
         </div>
-      )}
+        )
+      })()}
     </div>
   )
 }

web/src/index.css

@@ -146,3 +146,8 @@ button:focus-visible, input:focus-visible, textarea:focus-visible, select:focus-
   .brain-legend { gap: 8px; font-size: 10px; }
   .brain-tip { max-width: 220px; font-size: 10px; }
 }
+
+/* atlas hover extras */
+.brain-tip-spec { color: var(--primary); font-weight: 700; }
+.brain-tip-spec small, .brain-tip-aff { color: #8b9aa3; font-weight: 400; }
+.brain-tip-aff { font-size: 10px; }