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>
This commit is contained in:
ZacharyZcR
Executable
+51
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
# Unified user-experience benchmark: fixed scenarios, fixed decoding, medians.
|
||||
#
|
||||
# Reports, per scenario: TTFT (prefill wall), decode tok/s, and the first line
|
||||
# of the generated text (drift check). Runs each scenario REPS times and prints
|
||||
# every sample; judge by the median, not the best.
|
||||
#
|
||||
# Usage: SNAP=/path/to/model [REPS=3] [PIPE=2] tools/bench_ux.sh [engine-binary]
|
||||
# The engine env (COLI_CUDA, PIN, ...) is inherited from the caller so the same
|
||||
# script exercises CPU-only, CUDA and pipeline configurations.
|
||||
#
|
||||
# Discipline (docs/experiments/glm52-6x5090-2026-07-12.md):
|
||||
# - TEMP=0 DRAFT=0 always: greedy, no speculation, one variable at a time.
|
||||
# - Same binary for every configuration under comparison.
|
||||
# - .coli_usage drifts placement between runs: compare medians of >=3 reps,
|
||||
# or snapshot/restore the usage file around the battery.
|
||||
set -u
|
||||
GLM="${1:-./glm}"
|
||||
REPS="${REPS:-3}"
|
||||
export TEMP=0 DRAFT=0
|
||||
[ -n "${PIPE:-}" ] && export COLI_CUDA_PIPE="$PIPE"
|
||||
|
||||
SHORT_PROMPT="Explain why the sky is blue in simple terms."
|
||||
LONGDOC_PROMPT=$(python3 - <<'PY'
|
||||
base=("蜂鸟是世界上最小的鸟类之一，主要分布在美洲大陆。它们的翅膀每秒可以扇动五十到八十次，"
|
||||
"使它们能够在空中悬停、倒飞，甚至垂直起降。蜂鸟的新陈代谢极快，心跳每分钟可达一千两百次，"
|
||||
"因此它们必须不断进食来维持能量。它们的主要食物是花蜜，同时也会捕食小型昆虫和蜘蛛来补充蛋白质。"
|
||||
"蜂鸟的喙细长，适合深入花朵内部吸食花蜜，舌头呈管状，可以快速伸缩。在授粉过程中，"
|
||||
"蜂鸟扮演着重要的角色，许多美洲植物依赖蜂鸟传粉。")
|
||||
print((base*12)+"请根据上文回答：蜂鸟的主要食物是什么？")
|
||||
PY
|
||||
)
|
||||
|
||||
scenario(){ # name prompt ngen
|
||||
local name=$1 prompt=$2 ngen=$3
|
||||
for r in $(seq 1 "$REPS"); do
|
||||
local log; log=$(mktemp)
|
||||
NGEN=$ngen PROMPT="$prompt" "$GLM" 64 4 4 >"$log" 2>&1
|
||||
local line; line=$(grep -aE "prefill .* decode .*tok/s" "$log" | tail -1)
|
||||
local head; head=$(grep -av "^\[" "$log" | grep -avE "PROFILE|PROFILO|ATTENTION|expert|CUDA|prefill|specul|^---|TOPP|stop|Motore|caricato|prompt:|token" \
|
||||
| tail -1 | tr -d '\n' | cut -c1-48)
|
||||
printf "%-10s r%d %s\n" "$name" "$r" "$line"
|
||||
printf "%-10s r%d text: %s\n" "$name" "$r" "$head"
|
||||
rm -f "$log"
|
||||
done
|
||||
}
|
||||
|
||||
echo "== bench_ux: reps=$REPS pipe=${COLI_CUDA_PIPE:-unset} $(date +%F' '%T) =="
|
||||
scenario chat "$SHORT_PROMPT" 96
|
||||
scenario longdoc "$LONGDOC_PROMPT" 96
|
||||
echo "== bench_ux done $(date +%T) =="
|
||||
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Expert Atlas (#175): measure per-expert topic affinity by diffing .coli_usage
|
||||
across themed probe batches, served through a running colibri API server.
|
||||
|
||||
Protocol per category: snapshot .coli_usage -> send probes -> snapshot again;
|
||||
the delta is that category's expert-activation spectrum. One engine load total.
|
||||
|
||||
Output: experts.json — for every (layer, expert): counts per category,
|
||||
normalized affinity, entropy, and a "specialist" label when one topic dominates.
|
||||
|
||||
Usage (server already running with the model):
|
||||
python3 tools/expert_atlas.py --api http://127.0.0.1:8000 \
|
||||
--usage /path/to/model/.coli_usage --out experts.json --ngen 64
|
||||
"""
|
||||
import argparse, json, math, time, urllib.request
|
||||
|
||||
PROBES = {
|
||||
"code": [
|
||||
"Write a Python function that parses a CSV file and returns a dict keyed by the first column.",
|
||||
"Explain the difference between a mutex and a semaphore, with a C example.",
|
||||
"Refactor this into idiomatic Rust: for i in range(len(xs)): total += xs[i] * 2",
|
||||
],
|
||||
"math": [
|
||||
"Prove that the square root of 2 is irrational.",
|
||||
"Compute the derivative of x^3 * ln(x) and explain each step.",
|
||||
"A fair die is rolled 4 times. What is the probability of at least one six?",
|
||||
],
|
||||
"chinese": [
|
||||
"请用中文解释一下什么是光合作用,以及它对地球生态系统的重要性。",
|
||||
"把这句话翻译成中文并解释语法:The early bird catches the worm.",
|
||||
"写一段关于秋天的短文,一百字左右。",
|
||||
],
|
||||
"english_prose": [
|
||||
"Write a vivid paragraph describing an old lighthouse keeper watching a storm arrive.",
|
||||
"Summarize the plot of Romeo and Juliet in three sentences.",
|
||||
"Continue this story: The last train left the station, and Maria realized her mistake.",
|
||||
],
|
||||
"science": [
|
||||
"Explain how mRNA vaccines work at the cellular level.",
|
||||
"Why is the sky blue during the day but red at sunset?",
|
||||
"Describe the life cycle of a massive star, from formation to supernova.",
|
||||
],
|
||||
"law": [
|
||||
"Explain the difference between a patent, a trademark, and a copyright.",
|
||||
"What are the key elements required to form a legally binding contract?",
|
||||
"Summarize what 'due process' means in constitutional law.",
|
||||
],
|
||||
"poetry": [
|
||||
"Write a short poem about a hummingbird in the style of Emily Dickinson.",
|
||||
"Compose a haiku about winter rain, then explain its imagery.",
|
||||
"Write four rhyming lines about the sea at night.",
|
||||
],
|
||||
"structured": [
|
||||
'Convert to JSON: name Alice, age 30, hobbies reading and chess, address 5 Oak St.',
|
||||
"Write a SQL query returning the top 5 customers by total order value, with the schema you assume.",
|
||||
"Write a regex that matches ISO-8601 dates and explain each part.",
|
||||
],
|
||||
"translation": [
|
||||
"Translate into French, German and Spanish: 'Knowledge is the only treasure that grows when shared.'",
|
||||
"Translate this Italian sentence to English and comment on nuance: 'In bocca al lupo per domani.'",
|
||||
"Translate into Japanese: 'The meeting has been moved to next Tuesday afternoon.'",
|
||||
],
|
||||
"casual": [
|
||||
"Hey! Any tips for staying awake during boring afternoon meetings?",
|
||||
"What should I cook tonight? I have eggs, rice, tomatoes and some cheese.",
|
||||
"My friend is always late. How do I tell them it bothers me without being rude?",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def read_usage(path):
|
||||
counts = {}
|
||||
try:
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
p = line.split()
|
||||
if len(p) == 3:
|
||||
counts[(int(p[0]), int(p[1]))] = int(p[2])
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return counts
|
||||
|
||||
|
||||
def diff(after, before):
|
||||
return {k: v - before.get(k, 0) for k, v in after.items() if v - before.get(k, 0) > 0}
|
||||
|
||||
|
||||
def chat(api, prompt, ngen):
|
||||
body = json.dumps({"model": "glm-5.2-colibri", "stream": False, "max_tokens": ngen,
|
||||
"messages": [{"role": "user", "content": prompt}]}).encode()
|
||||
req = urllib.request.Request(f"{api}/v1/chat/completions", data=body,
|
||||
headers={"Content-Type": "application/json"})
|
||||
with urllib.request.urlopen(req, timeout=600) as r:
|
||||
json.load(r)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--api", default="http://127.0.0.1:8000")
|
||||
ap.add_argument("--usage", required=True, help="path to the model's .coli_usage")
|
||||
ap.add_argument("--out", default="experts.json")
|
||||
ap.add_argument("--ngen", type=int, default=64)
|
||||
a = ap.parse_args()
|
||||
|
||||
spectra = {}
|
||||
for cat, prompts in PROBES.items():
|
||||
before = read_usage(a.usage)
|
||||
t0 = time.time()
|
||||
for p in prompts:
|
||||
chat(a.api, p, a.ngen)
|
||||
time.sleep(2) # let the engine flush .coli_usage
|
||||
spectra[cat] = diff(read_usage(a.usage), before)
|
||||
total = sum(spectra[cat].values())
|
||||
print(f"[{cat}] {len(spectra[cat])} experts touched, {total} selections, {time.time()-t0:.0f}s", flush=True)
|
||||
|
||||
cats = list(PROBES.keys())
|
||||
experts = {}
|
||||
for cat, spec in spectra.items():
|
||||
for k, v in spec.items():
|
||||
experts.setdefault(k, {c: 0 for c in cats})[cat] = v
|
||||
|
||||
atlas = {}
|
||||
for (layer, eid), counts in experts.items():
|
||||
total = sum(counts.values())
|
||||
if total < 8:
|
||||
continue # too few observations to characterise
|
||||
aff = {c: v / total for c, v in counts.items() if v}
|
||||
ent = -sum(p * math.log2(p) for p in aff.values())
|
||||
top = max(aff, key=aff.get)
|
||||
label = f"specialist: {top}" if aff[top] >= 0.45 and ent < 2.2 else "generalist"
|
||||
atlas[f"{layer}:{eid}"] = {"counts": counts, "affinity": {c: round(p, 3) for c, p in aff.items()},
|
||||
"entropy": round(ent, 2), "top": top, "label": label}
|
||||
|
||||
spec_n = sum(1 for v in atlas.values() if v["label"].startswith("specialist"))
|
||||
with open(a.out, "w") as f:
|
||||
json.dump({"categories": cats, "ngen": a.ngen, "experts": atlas}, f)
|
||||
print(f"\natlas: {len(atlas)} experts characterised, {spec_n} specialists -> {a.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user