diff --git a/docs/plans/exl3-graphs.md b/docs/plans/exl3-graphs.md new file mode 100644 index 0000000000..1972a782d8 --- /dev/null +++ b/docs/plans/exl3-graphs.md @@ -0,0 +1,250 @@ +# CUDA graphs for EXL3 Flash Next with the expert offload pool + +Branch `exl3-c-graphs` (base `exl3-c-main` c5bfe87b). Opt-in: `TF_FLASHNEXT_EXL3_GRAPH=1`. + +## What blocked graphs: the pool's host cut + +Every offloaded MoE layer ended its segment with a D2H of the picks and a stream synchronize (`cuda_exl3_pass.pooled`), +the host planned hits/misses/victims (`exl3_round.plan`), staged the misses on the pool's copy stream and only then +launched the experts. 48 host waits a token, and a step that cannot be captured. + +## Design chosen: the pool planned on the device + +Measured first, then chosen over "graph A / host plan / graph B per layer": + +- `flashnext_exl3_dplan.cu` `fn_exl3_dplan` (one block a layer) runs the host planner's rules for a one-pass round + without prefetch or double-buffer protection on device copies of each decoder layer's state (line -> expert, stamps, + expert -> line, lifetime counts, clock): observe, touch held, misses hottest first, first empty line else the + coldest resident not in the round. It rewrites the pool's own pointer/width tables in place (fixed addresses) and + writes a copy list. `fn_exl3_dcopy` copies the misses from the pinned host bank through its mapped (unified) + address into the lines: 12.6 GB/s vs 13.6 GB/s for the copy engine (`flashnext-exl3-graph plan`). +- Equality: device planner == host planner over 400/600/600 random rounds (10/47/63 lines, 1/4/6 rows): lines, + stamps, counts, hit/copy/eviction counts, all 1536 table entries, every resident line's bytes. +- Ownership (`cuda_exl3_dplan.zig`): the host (prefill, the MTP layer's pool, eager rounds) and the device take turns; + `toDevice` uploads the host state after draining the copy streams, `toHost` downloads and rebuilds the Vram, counts + and statistics before any host plan of a decoder pool. A decode step then has no host wait between its first launch + and its logits. +- `cuda_exl3_graph.zig` captures that step: `Engine.round` (one row a stream) and `Engine.forward` windows with every + row's logits (MTP verify windows) of at most lines/10 rows. At tp=2 the graph holds the step's 48 NCCL sums and the + head's gathers (both ranks capture the same calls). Pinned uploads (ids, positions, n-gram rows) become copy nodes + read at replay; `prepare` rewrites them before each launch. Keys: rows, the sequences' buffers, attention chunk and + score-tile seams, scratch address and capacity, scores, logits, keep/streams, planner state. A key is captured at + its 6th sighting (`TF_FLASHNEXT_EXL3_GRAPH_AFTER`, prefill layers at the 2nd: `..._PREFILL_AFTER`); a new key of a + cached shape updates that exec in place (`cuGraphExecUpdate`, 2-3 ms instead of a 20-35 ms instantiate), at tp=2 + too: the captured template is kept alive with the exec because NCCL's captured resources belong to it + (`TF_FLASHNEXT_EXL3_GRAPH_TP_UPDATE=0` instantiates instead). +- A round with more misses than lines free of its own experts reads the extra experts in place from the mapped bank + (same bytes; slower) and clears their entries afterwards (`fn_exl3_drestore`), so any round up to + `TF_FLASHNEXT_EXL3_GRAPH_ROWS` (16, the MTP window) plans on the device. +- Fallbacks: more than 16 rows, media, or no pools take the host cut (eager), with a residency pull; + a failed capture turns graphs off for the process and runs the step eagerly on the device plan + (`TF_FLASHNEXT_EXL3_GRAPH_STRICT=1` makes it an error instead). +- `TF_FLASHNEXT_EXL3_DPLAN=1` alone: the device plan launched eagerly (no capture). +- `TF_FLASHNEXT_EXL3_OVERLAP=1` (default off, not validated on the pack): held experts compute while the misses' copy + runs on a side stream. + +## Prefill + +The expert-major deep prefill's dense half of each layer (sub-blocks: HC, DeltaNet/attention, write-backs, router, +top-k) is captured per layer (`prefillBegin/End` hook in `cuda_exl3_prefill.chunk`): a DeltaNet layer's graph depends +only on the chunk's shape and replays for later chunks; an attention layer's launch geometry follows the position, so +it is re-captured each chunk and updated in place. For TP=2, the graph captures each sub-block's wait on its incoming +NCCL sum as an external event-wait node. This is intended to preserve the eager schedule where ready sub-blocks begin +while later sub-block reductions are still running. Graphs are cleared before their reduction events are destroyed. +The experts' half stays eager: its passes, ring copies and member lists are planned on the host from each layer's picks +(data-dependent shapes). + +## Measurements (par1-cs13, 2x 5060 Ti, idle cards under gpurun) + +Engine test `flashnext-exl3-graph eq` (3 golden prompts, 48 greedy rounds + the same continuation as windows 3,1,4,2; +every step's bf16 logits bit-equal to the host cut, both ranks), decode ms/token: + +| tp | prompt | host cut | device plan | graph | +|----|--------|----------|-------------|-------| +| 2 | ids16 | 26.2 | 21.6 | 21.7 | +| 2 | France | 27.9 | 23.7 | 23.7 | +| 2 | Once upon | 28.9 | 25.9 | 25.1 | +| 1 | ids16 | 31.5 | 31.1 | 30.7 | +| 1 | France | 35.5 | 35.0 | 34.3 | +| 1 | Once upon | 36.8 | 36.0 | 35.3 | + +Two streams in one round (tp=2): 32.7 ms a round both ways, bit-equal. + +Why graphs add little over the device plan: a step is GPU-bound. nsys (tp=1): ~20.6 ms of kernels a token (EXL3 +grouped experts 5.5 ms, rows=1 EXL3 linears ~7 ms, native HC lanes ~4.4 ms) plus ~10 ms of expert copies at 86% +hits; host launch gaps were already hidden. The host cut's cost was the waits (both ranks stall at tp=2). + +Serve gate at tp=2 (tpgate template, two BENCH runs of 128 tokens, ms/token; DRAFTED==UNDRAFTED and BENCH sha +2c5d814cfa80, the eager one, in every run): + +| config | drafts off | drafts on | +|--------|-----------|-----------| +| eager (host cut) | 29.1 / 27.4 | 16.1 / 15.2 | +| device plan, no capture (`DPLAN=1`) | 26.9 / 24.5 | 16.3 / 15.6 | +| graphs, capture at 2nd sighting, tp=2 instantiate | 26.6 / 24.1 | 20.3 / 19.8 | +| graphs, 2nd sighting, in-place updates | 26.6 / 24.1 | 18.8 / 17.6 | +| graphs, 6th sighting, in-place updates (default) | 26.7 / 24.2 | 17.4 / 16.3 | + +Replays: 301 of 417 graph-eligible steps in the gate (72%), 11 captures (9 of them in-place updates), 105 steps eager +on the device plan (keys not yet seen 6 times: every request brings new sequence/draft buffers, so new keys). With +drafts on, a capture (15-20 ms of host work the GPU waits for) costs more than the replays save on a GPU-bound step. + +Prefill graph overlap validation (TP=2, 8192-row prompt, last-row logits plus 8 decode steps bit-equal): graphs +capture and replay the dense half of each layer, with external event-wait nodes for each expert-parallel reduction. +The graphs are cleared before those events are destroyed. With 1024-row chunks, eager measured 770.5 rows/s and +steady graph runs 771.0-772.0; with 2048-row chunks, eager measured 1116.3 and graphs 1112.1-1117.1; with +4096-row chunks, eager measured 1468.0 and graph replays 1469.5-1471.1. These small differences are within run +variation; a graph speedup is not demonstrated. An 8192-row chunk OOMed and fell back to the non-deep path, so its +output is not a valid graph measurement. The larger fitting chunk improves prefill throughput for both modes; it +does not establish a graph-specific win. + +Decode fork experiment (TP=2, same kernels and bit-equal logits): toggling multipart projection stream forking on +helped DPLAN eager in these runs, while graph changes were mixed and near measurement noise. This does not explain +or fix the graph-vs-eager decode gap. The earlier controlled serve gate remains the relevant result: graphs did not +reliably beat DPLAN eager. + +Local, uncommitted MTP verify-tail A/B (TP=2, blocking NCCL, graph mode): moving the existing greedy `Head.best` +selection and its small result copy onto the main forward stream before the graph's completion boundary preserved +the serving outputs and bench hash, but the paired 128-token gate showed no observable decode-time change at its +reported precision. This is one A/B pair, not a repeatability proof; the optimization remains opt-in and is not +claimed as a graph speedup. + +## Combined TP=2 paging/cache gate (2026-10-20) + +The rebuilt server now reports prefill graph replay/capture counters, decode coverage, decode preparation time and +cumulative graph capture time. On par1-cs13, a two-request 12,040-token TP=2 run used BF16 KV paging, MTP, prompt +snapshots and prefill/decode graph mode. With capture-after=1, each rank reported 72 prefill graph replays and 216 +captures across cold prefill. With the normal capture-after=2 threshold, parity remained exact, but prefill reported +0 replays, 192 eager layer executions and 96 captures; cumulative prefill capture time fell from about 0.5 s to +about 0.3 s per rank. Warm and rewind restored 12,033 prompt tokens; divergent branches reused 12,023. All responses +matched cold-cache token hashes. + +Two alternating DPLAN-versus-graph repetitions at 4,139 tokens preserved exact output hashes and warm/branch/rewind +cache hits in both modes. Request-level timings were within run/lane variation; the paired result does not establish +a speedup. The matched 12k DPLAN run also preserved exact hashes, while its per-lane prefill times varied under +concurrent scheduling. Treat these as correctness and instrumentation evidence, not a performance claim. + +Follow-up at 16,040 prompt tokens (TP=2, two concurrent lanes, BF16 paging, MTP, default prefill capture threshold +2) confirmed the threshold behavior and exercised actual prefill replay. With the original candidate matcher, both +ranks reported 72 prefill graph replays, 216 eager layer executions, 96 captures, and 314 ms cumulative prefill +capture time. The two lanes' cold/warm/branch/rewind hashes matched between DPLAN and graph mode; warm reused 16,033 +tokens and branch reused 16,023. This resolves the earlier 12k zero-replay observation: each lane did not revisit a +captured same-shape chunk, so the default second-sighting threshold captured but had no later chunk to replay. It was +expected warm-up, not a failed replay lookup. + +The repeated graph runs' cold prefill times were 27.92/16.39 s and 28.00/16.51 s for the two lanes; the DPLAN run +was 27.85/16.47 s. This is no demonstrated graph-specific speedup (differences are below or comparable to run/lane +variation). A follow-up matcher keeps the full exact replay key but normalizes only the paging-window-origin hash for +capture sightings and driver-checked update candidates; sequence/scratch/event ownership remains strict. The updated +gate reported 72 replays, 192 eager executions, 120 captures and 24 successful updates per rank, with exact cache +hashes. Capture time was 325 ms; DPLAN/graph cold-prefill times were 16.44/27.80 s and 16.51/27.99 s. The change +reduces eager fallback and reuses executable objects through `Exec.update`, but does not demonstrate a speedup; +incompatible graph topology still falls back to instantiation. Next, profile the expert/pool host boundary and dense +graph replay versus eager time before further topology changes. + +The decode profile with the normal threshold showed about 50 ms of host preparation over 96 steps, versus roughly +0.45–0.49 s cumulative graph capture/instantiation per rank. A capture-after=1 run raised decode capture time to +roughly 0.68–0.69 s per rank without a repeatable end-to-end gain. Thus delaying capture reduces one-time overhead, +but neither threshold has demonstrated faster output. The expert pass, pool transfers and ring scheduling remain +eager; these are the high-value prefill boundaries to profile next. MTP verification and result readback also remain +outside the graph. Keep graphs opt-in until an alternating, repeated TP=2 A/B demonstrates a repeatable improvement +for each phase. 200k × 2 remains deferred until graph speed and cache reuse are optimized. + +## Not proven / open + +- Prefill graph capture/replay is now validated with TP=2 deep prefill at chunk sizes 1024, 2048 and 4096, including + bit-identical logits and post-prefill decode. It restores the lost reduction overlap but has not shown a repeatable + throughput win over eager. +- The experts' half of prefill is not graphed (host-planned passes). +- In-place reads are proven bit-equal (tp=2, 24 lines, 2936 reads) but slow; windows at few lines lose to the host + cut's DMA passes. +- The overlap split (`TF_FLASHNEXT_EXL3_OVERLAP=1`) is unvalidated on the pack. +- Graphs stay opt-in (`TF_FLASHNEXT_EXL3_GRAPH=1`); a repeatable decode speed win over `TF_FLASHNEXT_EXL3_DPLAN=1` + remains unproven. +- The greedy verify-tail experiment has no demonstrated benefit and stays default-off pending a repeated profile or + a measured GPU/host critical-path reduction. +- 200k contexts and int8 KV are out of scope for this gate. Two concurrent BF16-paged server streams under graphs + now pass at 12k tokens; the 200k × 2 capacity target remains deferred. + +## Making tp=2 decode faster (branch exl3-c-fast, base exl3-c-main 7988c6fa) + +Profile (nsys, `flashnext-exl3-graph prof`: cuProfilerStart/Stop over 16 graphed rounds of 1 row and 16 verify +windows of 3 rows, both ranks, idle cards, ~99% hits), per round: rank 0 wall 20.2 ms (GPU busy 19.6), rank 1 19.6; +EXL3 dense projections 6.1 (GDN in_proj q|k|v 2.2, z / attention 1.4, GDN out_proj 1.1, lm_head half 0.56), +native bf16 lanes (HC down/up, mixer) 4.5, NCCL 3.5 on rank 0 vs 2.0 on rank 1 (60 all-reduces; rank 0 waits ~1.5 ms +for rank 1's shared expert and extra misses), experts 2.9 / 3.3, HC glue 0.9, router/top-k 0.6, GDN 0.5, device +plan 0.25, miss copies 0.24 / 0.74, attention 0.1. Serve profile with drafts on: the GPU idled about half the time; +every kernel ran eagerly (verify windows rarely repeat a graph key) and each eager all-reduce cost ~294 us of host +time (6801 cudaStreamWaitEvent calls = 2.0 s on rank 0 over one 128-token request). + +What helped (all bit-identical: same kernels on the same inputs): + +1. The n-gram rows a step gathers on the host (`Ngram.rows`, EXL3 table): inline instead of items/4 threads, and a + one-pass row unpack: 0.6-0.8 ms a token -> 0.05-0.13 ms (3-row window 2.35 -> 0.15-0.36 ms). It sat on the + critical path in both modes (graph: before launch; eager: joined at layer 1). +2. `TF_FLASHNEXT_NCCL_BLOCKING=1`: a blocking NCCL communicator. Eager all-reduce 109 us -> 20.6 us (the captured + one is 20 us); `Collective.wait` also spins before its 50 us sleeps. Opt-in: a blocking init has no deadline if + the peer never joins. +3. Multi-part EXL3 projections' parts on parallel streams (`cuda_exl3_fork.zig`, default on): GPU busy 19.6 -> 18.8 + ms a round, within noise end to end. + +Serve gate at tp=2 (BENCH 128 tokens twice, ms/token; DRAFTED==UNDRAFTED and sha 2c5d814cfa80 in every run): + +| tree / config | drafts off | drafts on | +|---|---|---| +| 7988c6fa eager | 30.1 / 29.7 | 14.9 / 14.6 | +| 7988c6fa `GRAPH=1` | 25.7 / 23.5 | 15.4 / 14.8 | +| exl3-c-fast `GRAPH=1 NCCL_BLOCKING=1` | 23.7 / 21.4 | 12.9 / 12.0 | +| exl3-c-fast `NCCL_BLOCKING=1` (eager) | 23.4 / 21.3 | 12.2 / 11.6 | + +With cheap eager collectives the step is GPU-bound and graphs no longer add speed (drafts on they lose a little to +captures of keys that do not repeat). + +Tried, no gain or rejected: +- upstream's EXL3 linear body re-instantiated at 2-4 programs an SM (128 registers, small spills): bit-equal, + 19.67 -> 19.71 ms/token (the 1-row projection is bound by trellis decode and its serial k walk, not occupancy); +- the device plan's held/miss overlap split (`TF_FLASHNEXT_EXL3_OVERLAP=1`): 22.3 -> 22.1 ms rounds, 29.4 -> 29.9 ms + windows; stays off; +- resident MTP experts (`TF_FLASHNEXT_EXL3_MTP_RESIDENT=1`): costs rank 0 17 pool lines, no change; +- `NCCL_PROTO=LL`: captured all-reduce 19.8 -> 11.4 us, serve drafts off 24.3/22.0 -> 23.8/21.1 ms, but tp=2 deep + prefill 1473 -> 1288 rows/s (-12%): not set by default (it is process-wide); +- prefill layer graphs at tp=2 before external event waits: bit-equal, 1474 -> 1324 rows/s. The measured external + event-wait change restores the overlap and bit parity; it does not produce a repeatable graph-specific speedup. + +What remains is kernel time: ~6 ms of 1-3-row EXL3 projections (decode-bound at ~190-290 GB/s), ~4.5 ms of bf16 HC +lanes (already 290-390 GB/s of 448), ~3 ms of experts, ~2 ms of collectives, ~1.5 ms of rank-0 wait. Faster +projection kernels with a different split or fewer HC bytes change the fp32 order (numerics) and would have to be +default-off with logit diffs; the MTP head runs on rank 0 alone (rank 1 idle ~28% with drafts on). + +## TP=2 BF16-paged prefill profile (2026-10-20) + +Added `flashnext-exl3-graph ... prefill-prof N eager|graph` as a bounded Nsight Systems harness. It requires +`TF_FLASHNEXT_KVPAGE=1`, warms the same 8,192-token input twice, measures one forward with 4,096-row chunks, then +checks the final logits hash across both ranks. The eager and graph runs both produced hash `ffcf463b229c2aa6`; +graph mode replayed 96 dense-layer graphs. The paired non-Nsight run measured 5,597.1 ms eager and 5,639.5 ms +graph (one pair; not a speedup claim). In the Nsight-captured pair, rank 0 measured 5,693.6 ms eager and 5,696.4 ms +graph, with the same hash. Differences are below what this single pair can establish. + +The traces show that graphs are doing real work but are not currently shortening the prefill critical path: + +- Eager submitted 20,559 `cuLaunchKernel` calls; graph mode submitted 2,103 plus 96 `cuGraphLaunch` calls. Dense graph + replay therefore removes much of the host launch stream, but not the expert boundary. +- Each rank moved about 28.8 GB host-to-device in roughly 48,000 copy operations in either mode; measured copy-engine + time summed to 2.79 s eager / 2.77 s graph. These are trace totals across operations, not additive wall-time slices. +- The measured forward scheduled about 14.4k expert copies on rank 0 (14.0k counted as prefetch), 570 expert passes, + and about 35 ms of CPU planning. Rank 1's copy/pass counts were similarly stable. +- `host_wait_ns` is the sum of intervals from enqueueing the picks D2H through stream synchronization; it includes + preceding dense kernels and reduction dependencies, not just the D2H or removable CPU stall. It was about 5.19 s + eager and 5.28 s graph in the untraced pair. Nsight also reports about 4.86/4.95 s total across 201 + `cuStreamSynchronize` API calls; do not interpret either sum as a pure synchronization penalty. +- All-reduce and the largest EXL3 prompt kernels had nearly equal aggregate time in eager and graph traces. The + roughly 2.8 s of H2D work overlaps other GPU work, so the operation totals cannot be summed into a critical path. + +The evidence says not to keep tuning capture thresholds or dense graph launch mechanics as the next prefill lever. +The next design target is the expert boundary: preserve exact GPU routing while supplying expert work and transfers +without a per-layer host decision on the critical path. The implementation should begin with a bounded two-layer +pipeline: GPU-built stable member lists, fixed-address double-buffered staging, and the existing expert arithmetic +and TP=2 reduction order. Before committing to full-layer streaming, measure each rank's largest layer-bank size and +available memory after BF16 KV paging, pools, scratch, and snapshots. Keep a selective-transfer fallback for cases +where streaming unused experts costs more than the host cut. Gate skewed routing, empty experts, ownership masks, +page-window rebases, cache branch/rewind, snapshots, and eager-versus-graph logit parity before optimizing timings. diff --git a/docs/plans/exl3-kv-paging.md b/docs/plans/exl3-kv-paging.md new file mode 100644 index 0000000000..b50518940f --- /dev/null +++ b/docs/plans/exl3-kv-paging.md @@ -0,0 +1,460 @@ +# Dynamic KV paging for long contexts (KVP) + +## Experimental rank-shared BF16 block arena + +The default remains a private pager per sequence. `TF_FLASHNEXT_KVPAGE=1` with +`TF_FLASHNEXT_KVPAGE_SHARED=1` opts into a fixed rank-owned allocation. Set +`TF_FLASHNEXT_KVPAGE_SHARED_LANES` to the admitted lane count (default 8, maximum +32). The existing window, rows and slots settings retain their meanings; slots +now form one pool across every lane and attention layer. Ordinary-context TP2 +hardware gates pass within the scope below; a repeatable performance improvement +and 200k-context capacity remain unqualified. + +Runtime source `ccb32fc3` passes 153 EXL3 host tests and real CUDA sm_120 builds. +Matched private/shared GPU fixtures force sparse paging with 1,024 slots and +compare 54 per-step logit hashes and six logical-state records per rank. They +also check exact restored state, owned-entry payload/ages, recycled and +asymmetric TP generations, and a foreign request's windows/history/continuation. +Each restored graph phase replays 201 prefill and seven decode graphs per rank. +Two concurrent 16,424-token HTTP histories with MTP and graphs pass cold, warm, +branch and rewind against the private-pager reference, reusing 16,407–16,417 +tokens. This is not an independent golden-model check or a 200k-capacity gate. + +A separate same-binary eight-lane allocation diagnostic on source `0de41d37` +measures 920 MiB lower sampled physical GPU peaks per rank. It uses 1,024 input +and 128 output tokens per request; do not relabel it as the later source or +long-context performance evidence. Full settings, hashes and limitations are in +the [qualification receipt](https://github.com/ashhart/TensorFold/pull/1113#issuecomment-6325790752). + +The slot count is global in shared mode, not per sequence: eight lanes with +16,384 slots share 16,384 entries rather than reserving eight private pools. +This shares allocation capacity, not identical-prefix payloads. Reduced total +residency can increase misses; memory savings are not a throughput claim. Fixed +windows are prepaid for every configured lane, so low-concurrency operation can +use more resident memory than private, lazily created pagers. Compare physical +peaks with the same active concurrency, window geometry and expert pools. + +Each K/V side is `[all lane/layer windows | shared block slots]`. A pager uses +its own window base and a checked relative offset to the shared slots, retaining +the upstream BF16 attention load path. Canonical RAM, window origins, logical +history, indexer keys and recurrent state remain sequence-owned. Every sequence +must use the engine's common admitted capacity (`context` plus MTP slack, bounded +by the model window). A mismatched capacity or arena geometry is rejected. + +Startup allocates and charges all fixed windows, shared payload, global mapping +tables and paging work once, before expert-pool admission. Per-sequence admission +charges the remaining sequence buffers. The lane limit also caps startup room. +Request teardown releases only its generation's tags; it does not free the arena. +With graph recycling enabled, compatible BF16 sequence/pager allocations can now +remain in a bounded four-entry pool. Reuse drains previous consumers, resets all +request state and canonical storage, and reattaches a fresh paging owner. Old +request graphs retire: retaining addresses does not authorize replay under an +old generation. Incompatible retained allocations are evicted before new +allocation, and disabling recycling purges that engine's retained sequences. +These request-boundary operations include context completion; they do not add +waits to each normal token. On runtime `4afc6169`, both sm_120 cards pass the real +private/shared lifecycle fixture with device and mapped-host indexers, stable +allocation addresses, complete reset, fresh owners, snapshot restoration and +foreign-lane preservation. The host c8 regression also remains within its +original combined device/pinned allocation limit across incompatible capacities. +This does not establish a c8 hardware memory saving or cross-request graph reuse. + +`Engine.forward` and `Engine.round` acquire a rank-wide batch lease containing all +participating sequence generations. Pager calls validate membership, and the +lease covers paging plus every attention consumer, including captured/replayed +rounds. Normal release acknowledges the existing `finishPaged`/collective +completion boundary and introduces no additional device-wide wait. A submitted +failure poisons touched targets and retains the lease until draining completes. + +Private snapshots retain their exact version 1 contract. Shared snapshots use +explicit version 2: canonical BF16 prefixes, complete owned windows, ordinary +state, local marks/counters and only the sequence's resident blocks. Restore +assigns destination slots and generations, invalidates evicted inverse maps, and +preserves owned block ages against a monotonic shared clock. Global physical slot +numbers and another lane's payload are not restored. Unrelated windows and +canonical history remain untouched. If residency is full, ordinary eviction can +remove a foreign cached block while preserving its canonical history. Snapshot +admission includes the checked maximum owned-entry payload and metadata. + +## Graph replay and separate cache diagnostics + +The follow-on native BF16 prefill path now retains validated attention graph +templates and rebinds origin-dependent pointers, planner/remap scalars and +score/chunk grids without recapturing compatible chunks. Paging identity, +allocation, stream and event keys remain strict. Dense/sparse and paging-split +topology changes still recapture, retained templates have separate budget +tickets, and a partially failed setter poisons replay. Runtime `4afc6169` passes +the private and shared TP2 device fixtures: 24 steps of 256 rows per storage +mode, exact eager/rebound logits and recurrent/conv/PLE/indexer/canonical BF16 +state, 241 rebinds per rank in each mode, and zero-rebind recapture at the +dense/sparse and paging-split transitions. Scratch is poisoned between runs. + +The same runtime's native binary passes two 16,424-token HTTP histories with +32-token MTP output through cold/warm/branch/rewind prefix reuse. All eight +response hashes match the separate `ca2b13c8` cold reference; warm/branch/rewind +reuse covers 16,407–16,417 tokens. Both ranks sample 144 prefill-origin rebinds +and six sequence-buffer reuses. That correctness run overlapped debug-runner +compilation, so its timing is not performance evidence. Independent golden +fidelity, repeatable graph speedup and 200k-context capacity remain open. + +Reference EXL3 projections also reset split-K counters inside their existing +input rotation, removing one separate memset node per 128-row reference block. +This leaves projection arithmetic unchanged. Fresh NVCC 13.4.92 images for +sm_80/86/89/120/121 retain the existing rotation resource records; this is compile +resource evidence, not GPU arithmetic or throughput evidence. See the +[incremental checkpoint](https://github.com/ashhart/TensorFold/pull/1113#issuecomment-6328093792) +for the separately tested baseline and current qualification boundaries. +The isolated `52ed4ac3` device build separately passes the poisoned-counter, +canary, repeated eager/graph reset regression and existing EXL3 FP64 checks on +sm_120. Pack-dependent Python oracles were skipped by that synthetic run; it is +not a qualification of the later prefill/lifecycle changes or a speedup result. + +`TF_FLASHNEXT_EXL3_GRAPH_STATE_REBIND=1` now permits exact-pair GDN state +rebinding for verification widths 2–16, not just width 2. In-place width 1 and +multi-lane layouts remain outside this rebinding optimization. Allocation, +generation, paging, nonstate-parameter and TP agreement checks remain strict; +partial update failures poison the affected binding. This does not make every +operation graphed or eliminate prefill-origin and request-allocation churn. + +Runtime `713e379b` retains all four prompt/token/text hashes against `ccb32fc3` +in two TP2 waves of two 16,384-input / 256-output requests with shared BF16 KV, +MTP and graphs. Source `b3a7b565` adds a real GPU fixture for widths 3, 4 and 16: +each width on each rank passes exact eager/graph state and logits, one capture +and three actual rebind replays. These are replay correctness/reuse gates, not +evidence of a repeatable throughput improvement. + +`TF_FLASHNEXT_CACHE_TELEMETRY=1` enables diagnostic reports on both TP ranks. +It is off by default; enabled sampling waits at completed model-call boundaries +and therefore must not be used to establish normal-path performance or overlap. +`EXL3-CACHE v=1` records distinguish rank, phase, model rows and operation count. +Explicit phase metadata crosses the TP control path; common MTP verification +enters and restores its own phase. Unclassified work remains `unknown`. + +Expert reports separate host-planner picks, executed device-planner picks, +mandatory/prefetch pool copies, mapped in-place reads, and expert-major ring +demand/copies. Opt-in ring attribution separates demand, initial-chunk prefetch +and cross-layer prefetch copies, requested bytes and unused residency. A held +retry keeps the original copy's provenance; unused residency can be observed in +a later interval than its copy. Disabled attribution allocates no slot records. +Direct DMA coalesces adjacent projections only inside a proven common pinned +bank; expert-copy and logical-byte counters retain their original meaning. +KV reports count unique requested and missing blocks before the +active window, sparse calls, rebases and written rows. Their denominators differ: +routed picks are not distinct experts, and KV block requests are not token positions. Requested +or logical payload bytes are not measured PCIe traffic. Automatic decoder +hooks exclude standalone MTP-head work and transfers outside model calls, such +as snapshot restoration. Intervals crossing unaccounted resets, restores, +recycled owners or ambiguous wraps are rejected rather than reported as zero misses; +shared paging uses completed-boundary arena clocks, not stale per-layer epochs. + +On runtime `975c1df9`, the diagnostics-enabled TP2 shared-BF16 cache-reuse gate +passes all eight cold/warm/branch/rewind response-token hashes against the +cache-disabled reference. Separate expert and KV counters pass validation on +both ranks, including matching phase denominators and actual block misses. +A fresh diagnostics-disabled c8 run passes two waves of 1,024-input / 128-output +requests with identical full prompt/token/text hashes and concurrent graph +replay. These are bounded diagnostics, not a repeatable speedup or 200k result; +settings, binary hashes and measured counters are in the +[replay/cache qualification](https://github.com/ashhart/TensorFold/pull/1113#issuecomment-6327563987). + +The serving prefill default is 4,096 rows. Benchmark harnesses may explicitly +override it to 2,048; comparing those settings is not a new-code speedup. + +## Snapshot and cache reuse validation (2026-10-20) + +Version 1 snapshots preserve canonical BF16 RAM rows after flushing every dirty window, the full compact device +windows and slot payloads, both mapping directions, marks, clocks, epochs and LRU stamps. They also retain recurrent +and convolution state, partial raw index keys, complete pooled blocks, PLE tail/history, next logits, and an optional +settled MTP head. Restore copies into destination allocations; source lifetime does not constrain snapshot lifetime. +Transient MTP verification generations and captured sequence graphs are invalidated before restored state is used. + +The pager configuration is frozen on the loaded engine. Paged serving lanes allocate the configured context capacity +already included in stream admission, so different prompt lengths use compatible table strides. Snapshot admission +includes its fixed windows, tables and header in addition to logical history; paged snapshots are charged to host RAM. +Capacity or geometry changes are rejected before mutation. This format preserves the warm pager cache, rather than +reconstructing a cold cache. + +TP2 restore validates both ranks, copies and waits on both ranks, then commits. The owning scheduler remains excluded +until acknowledgements complete. A copy or peer failure poisons the session; a partially restored destination cannot +execute. Failed asynchronous submissions drain before their memory is freed, following the existing CUDA completion +policy (terminate the rank if neither its stream nor context can drain). + +Current scope: snapshot/cache reuse correctness and prefill/decode graph replay with two ordinary-length lanes. +The two concurrent 200,000-token gate is deferred. Its later sizing candidate is `--context 200256 --parallel 2`, +`TF_FLASHNEXT_EXL3_STREAMS=2`, BF16, host IKC and snapshots, paging window 4096 / slots 49152 / rows 2048, engine chunk +512 and deep prefill rows 2048. The serving client must assert the exact prompt-token count and return nonzero on any +failed check. Configuration is a sizing candidate, not a validated capacity claim. + +Validation: focused host/lifecycle coverage passes, including restore-copy submission failures, preflight rejection, +source drop/repeated restore, MTP stale-publication rejection, TP2 prepare/copy/commit ordering, and failed-completion +draining. CUDA lifecycle tests pass at boundaries 0, 1, 3, 4, 2051, 2052, 4096, 4099 and 6147, checking windows, +slots and mapping tables. The TP2 two-GPU graph test passes with two paged BF16 lanes and exact state/logit parity; +each rank replayed prefill graphs 72 times and decode graphs 5 times after restore. A real TP2 HTTP chat-history gate +with BF16 paging, graphs, MTP drafting and two concurrent requests passed warm, branch and rewind reuse at 12,040 +prompt tokens: warm/rewind reused 12,033 tokens, branch reused 12,023, and all response hashes matched the cold +baseline. With capture-after=1, each rank observed 72 prefill graph replays and 216 captures. With the default +capture-after=2, each rank observed 0 prefill replays, 192 eager layer executions and 96 captures, with lower +one-time capture cost. Decode logs expose host preparation and capture time; capture-after=2 reduced cumulative +capture time versus capture-after=1. Paired DPLAN-vs-graph gates preserved exact hashes, but timings show no +repeatable graph speedup. This proves ordinary-length cache reuse, not 200k capacity. +Live-rank overflow/failure injection remains open. The user explicitly defers 200k × 2 until graphs and cache reuse +are optimized. Int8 KV is out of scope; BF16 is the required primary format. + +Status: active implementation, 2026-10-20. Branch `exl3-c-main`. Task: prove a dynamic KV +paging cache at long context on upstream's attention kernels: a sequence's full bf16 K/V in pinned host memory, the +card holding a recent dense window plus an LRU cache of the 4-key blocks the sparse indexer selects. + +Every number is tagged: **measured** (par1-cs13, one idle RTX 5060 Ti under `~/bin/gpurun`, sm_120), **replay** +(offline simulation over measured selection traces), or *estimate* (arithmetic, not a run). + +## 1. What the attention reads (measured traces, replayed) + +Tool: `tf-cuda-test exl3-kvpage DIR trace IDS N,N.. DECODE OUT prompt` runs the family's Engine (EXL3 offload pool, +expert-major deep prefill, as serve attaches it) over the first N ids of `~/pf/long_ids.bin` (306k ids of real text: +the repo's docs tokenized), then DECODE greedy rounds, and records after every `fn_qsa_select` each row's 512 +selected blocks (`cuda_exl3_kvpage.Trace`, a one-line hook at the top of `cuda_batch.zig` attnRows). Traces: +32,768 / 65,536 / 200,000 prompt tokens, every prompt row of every attention layer plus 256 real decode rows each +(2.40 M rows at 200k, 2.3 GiB). `tools/exl3_kvpage_replay.py` replays them. + +Selection churn (replay over the decode rows, mean over the 12 layers): + +| prompt tokens | blocks in context a layer | overlap with the previous row | new blocks a row a layer | union of a 4-row MTP window | share of the 512 in the last 4,096 tokens | +|---:|---:|---:|---:|---:|---:| +| 32,768 | 8,192 | 0.62 (3 layers) | 197 | 1,002 | 0.21 | +| 65,536 | 16,384 | 0.55 | 230 | 1,093 | 0.27 | +| 200,000 | 50,000 | 0.49 (0.39-0.60 by layer) | 263 | 1,175 | 0.11 | + +K10's 0.80-0.88 (6.5k tokens) does not hold at long context: at 200k consecutive decode rows share half their blocks. + +Prompt rows: a 256-row attention call's selections cover (union, blocks) 4,068 mean / 7,202 max at <32k, 9,538 / +12,136 at 32-64k, 11,692 / 16,699 at 64-128k and 14,175 / 18,852 at 128-200k (layer 0; layer 1 is the widest: max +24,564). Outside a 4,096-token window the max is 23,595 blocks (184 MiB of K/V) at 200k. + +### Cache replay + +One slot pool shared by the 12 layers (keys layer × blocks + block; exact LRU at call granularity; calls in +execution order: the prompt layer-major inside each deep chunk, then decode 12 layers a token), window W tokens a +layer always on the card. Hit rate = selected blocks found on the card (window or slots) / selected; bytes = misses × +8 KiB (4 keys × K and V × 2 heads × 256 × bf16). MTP: 4 consecutive decode rows as one call a layer. + +| prompt | W | slots (MiB) | decode hit | misses a token (p95) | MiB a token | MTP-4 hit | MiB a 4-row window | prefill copies | per prompt row | +|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 32,768 | 4,096 | 12,288 (96) | 0.819 | 1,111 (1,990) | 8.7 | 0.640 | 34.7 | 2.6 GiB | 83 KiB | +| | | 24,576 (192) | 0.907 | 572 (1,223) | 4.5 | 0.813 | 18.0 | 2.5 GiB | 80 KiB | +| | | 49,152 (384) | 0.963 | 228 (569) | 1.8 | 0.925 | 7.2 | 2.0 GiB | 63 KiB | +| 65,536 | 4,096 | 12,288 (96) | 0.779 | 1,358 (2,115) | 10.6 | 0.596 | 41.3 | 12.4 GiB (overflow) | 199 KiB | +| | | 24,576 (192) | 0.863 | 839 (1,598) | 6.6 | 0.747 | 25.9 | 11.0 GiB | 176 KiB | +| | | 49,152 (384) | 0.913 | 532 (1,118) | 4.2 | 0.837 | 16.7 | 10.5 GiB | 168 KiB | +| | | 98,304 (768) | 0.954 | 280 (734) | 2.2 | 0.914 | 8.8 | 8.2 GiB | 131 KiB | +| 200,000 | 4,096 | 24,576 (192) | 0.764 | 1,447 (2,478) | 11.3 | 0.585 | 45.7 | 104.7 GiB | 549 KiB | +| | | 49,152 (384) | 0.831 | 1,039 (1,930) | 8.1 | 0.704 | 32.6 | 90.9 GiB | 477 KiB | +| | | 98,304 (768) | 0.884 | 712 (1,598) | 5.6 | 0.798 | 22.3 | 88.6 GiB | 465 KiB | +| | 8,192 | 98,304 (768) | 0.890 | 679 (1,504) | 5.3 | 0.807 | 21.2 | 84.2 GiB | 442 KiB | +| | 2,560 | 49,152 (384) | 0.827 | 1,062 (1,980) | 8.3 | 0.698 | 33.3 | 92.7 GiB | 486 KiB | + +(Per-layer pools of the same total do about the same: 8,192 a layer at 200k: decode hit 0.897.) 12,288 slots +overflow during the 200k prompt (a call's union does not fit); 24,576 is the floor at 200k with 256-row calls. + +Reading: the cache is not a "hit rate 99%" cache at 200k; it is a bandwidth reducer. At 0.38-0.77 GiB of slots a +decode token moves 5.6-8.1 MiB (all 12 layers) instead of the 48 MiB its 12 × 512 selected blocks hold, and the card +does not hold the 4.58 GiB of a 200k stream's full K/V. At PCIe ~12-13 GB/s (K10: one pinned H2D 13.9 +GB/s, kernel gather from mapped memory 12.5 GB/s) that is 0.45-0.7 ms a token (*estimate*), next to a decode step of +~57-70 ms on this card (the expert pool's misses dominate). A 4-row MTP verify window moves 22-33 MiB (~2-2.7 ms, +*estimate*); the prompt moves ~0.45-0.55 MiB a row (90 GiB for 200k: ~7.5 s next to a ~170-250 s prefill, +*estimate*). + +## 2. Design (built) + +Files: `zig/src/families/flashnext/cuda_exl3_kvpage.zig` (trace, `Paged`, the attention rows over it), +`zig/kernels/cuda/flashnext_exl3_kvpage.cu` (`kvp_plan`, `kvp_fetch`, `kvp_remap`), `zig/tests/cuda/exl3_kvpage_tests.zig` +(`tf-cuda-test exl3-kvpage`), `tools/exl3_kvpage_replay.py`. Hooks (one line each): `cuda_batch.zig` attnRows +(`if (exl3_kvpage.takes(self.seq)) return exl3_kvpage.attnRows(...)`), `cuda_batch_state.zig` SeqState `.paged`, +`cuda_prompt.zig` Seq `.paged` / view / deinit / newSeq (`kvpage.forSeq`; kc/vc shrink to one row), `cuda_snap.zig` +(refuses a paged sequence), `cuda_lanes_load.zig` (stream budget), kernel list (2 lines), barrel, main.zig. +**No upstream `.cu` changed.** + +Switch: `TF_FLASHNEXT_KVPAGE=1` (default off; all-on-device bf16 stays the default), sizes +`TF_FLASHNEXT_KVPAGE_WINDOW` (4096), `_SLOTS` (49152, shared by the 12 layers), `_ROWS` (max(1536, engine chunk)). + +- **Host store**: one mapped pinned buffer a sequence, `[layer][K|V][cap][2][256]` bf16 (4.6 GiB at 200k). +- **Card**: one K buffer and one V buffer `[12 windows of wr rows | 4 × slots rows][2][256]`, wr = window + rows + 4; + tables: blk2slot `[12][blocks]`, mark `[12][blocks]`, slot2blk / stamp `[slots]`, counters `[12][8]`. +- **Writes**: upstream's `fn_attn_prep` writes the rows into the layer's window (its K/V base pointer offset by + −win0 rows: prep indexes by absolute position). When a call would run past the window, the window moves: the rows + that leave are written to host memory (`cuMemcpyDtoDAsync` to the mapped pointer) and the last `window` tokens move + down in non-overlapping pieces. So the prompt writes its K/V to host as it goes (one window step behind) and decode + writes nothing per token. +- **Reads**: after upstream's `fn_qsa_select` lists a call's keys (positions), `kvp_plan` (one 1024-thread block) + collects the call's unique blocks before the window, stamps the hits with the call's epoch, takes for the misses the + least recently used slots (exact LRU: a 4-pass radix selection of the smallest stamps, as `fn_qsa_select` selects the + largest scores), and updates the tables; `kvp_fetch` copies each miss's 8 KiB from mapped host memory into its slot + (a kernel, 128 blocks); `kvp_remap` rewrites each sparse row's key list from positions to rows of the compact + buffer (window rows, or slot base + 4 × slot + offset). Upstream's `fn_attn_chunks` / `fn_attn_merge` then run + unchanged on the compact buffer. Dense rows (<= 2,051 keys) read rows 0..n−1, which is the window while it has not + moved (it moves only past position window + rows > 2,051). +- **Shared slots**: one pool for all layers, because a prompt runs layer-major (one layer's calls take the whole + pool: a 256-row call at 200k needs up to 23.6k blocks) and decode spreads evenly (replay: same hit rate as per-layer + pools of the same total). +- Everything runs on the compute stream (plan → fetch → remap → attention): the misses are only known after the + layer's own selection, and the next layer depends on this one, so there is nothing to overlap the fetch with inside + a token (see §5). + +## 3. Proof (measured, `tf-cuda-test exl3-kvpage`) + +**Kernels** (`exl3-kvpage none kernel 6144|49152`, synthetic key lists, 60 calls of 1 or 64 rows with 15% drift a +row, window start at 40,000): 898,338 keys remapped, **every one reads its own host row** (K and V), 0 bad; misses +GPU 13,657 vs host LRU model 13,666 (6,144 slots) and 9,642 vs 9,642 (49,152 slots); no overflow. + +**Whole model, one card** (`exl3-kvpage DIR cmp IDS N.. DECODE 4096 49152 [NEEDLE]`; the family's Engine with the EXL3 +offload pool and the deep prefill as serve attaches them, `TF_FLASHNEXT_CHUNK=512 TF_FLASHNEXT_EXL3_PREFILL_ROWS=4096`; +each prompt through an all-on-device sequence, then the same Engine with a paged sequence; real text from +`~/pf/long_ids.bin`): + +| prompt | last-row logits | greedy tokens | K/V rows equal (12 layers × K,V) | prefill ms all-on-device / paged | decode ms/token all-on-device / paged | +|---:|---|---|---|---:|---:| +| 32,768 | bit-equal | 128/128 equal | 24/24 (32,896 rows) | 34,605 / 34,828 (+0.6%) | 63.48 / 63.95 (+0.47) | +| 65,536 | bit-equal | 128/128 equal | 24/24 | 71,075 / 72,594 (+2.1%) | 64.33 / 65.20 (+0.87) | +| 200,000 | bit-equal | 128/128 equal | 24/24 (200,128 rows) | 244,232 / 254,237 (+4.1%) | 69.44 / 70.61 (+1.17) | +| 65,536 needle at 0.2 | bit-equal | 24/24 equal | 24/24 | 75,385 / 76,342 | 51.93 / 50.76 | +| 131,072 needle at 0.2 | bit-equal | 24/24 equal | 24/24 | 159,563 / 164,551 | 49.17 / 49.63 | +| 65,536 needle at 0.35 | bit-equal | 24/24 equal | 24/24 | 75,384 / 76,314 | 59.65 / 58.89 | +| 98,304 needle at 0.35 | bit-equal | 24/24 equal | 24/24 | 115,648 / 118,685 | 50.21 / 49.83 | + +Copies and memory (paged runs): + +| prompt | prefill hit / fetched | decode hit (outside the window) | fetched a token | sequence on the card all-on-device / paged | +|---:|---|---:|---:|---| +| 32,768 | 93.7% / 1.9 GiB | 94.1% | 2.30 MiB | 5.484 / 1.391 GiB (cap 200k: freed 4.09 GiB) | +| 65,536 | 90.8% / 10.3 GiB | 86.9% | 4.40 MiB | same allocation | +| 200,000 | 84.4% / 90.4 GiB | 80.5% | 8.26 MiB | same allocation | + +The paged side itself is 0.486 GiB (windows 0.10, slots 0.375, tables); the rest of the 1.39 GiB is what stays as +upstream has it (indexer keys 0.57 and pooled keys 0.14 at 200k, GDN states, PLE tail). Pinned host: 4.6 GiB a 200k +stream. The measured copies match the replay (200k: 8.26 vs 8.1 MiB a token, 80.5% vs 80.9% outside-window hit, +prefill 90.4 vs 90.9 GiB). + +**Needle**: a sentence with a 5-digit code inserted at 20% of the text and a question appended. At 131,072 tokens +both caches answer "70259" (the code; position ~26k, far outside the window and long since evicted from and fetched +back into the slots). At 65,536 both caches emit `<|im_end|>` and start a chat turn instead of answering (identical +tokens: the model's behaviour on this prompt, not the cache's). With the code at 35% of the text both caches answer "48613" at 65,536 tokens and +"70259" at 98,304 tokens. + +**tp=2 serve** (`tensorfold-native serve --tp 2 --context 65536 --prompt-cache-gib 0`, both cards, upstream's +expert-parallel EXL3 path; `~/tensorfold-c-kvpage/tpkv.sh`): two real-text prompts of 27,632 and 57,546 tokens, +greedy 32 tokens, `TF_FLASHNEXT_KVPAGE` off then on: **the completions are identical** on both prompts. With paging +the stream budget is smaller, so the expert pools grow from 52 / 59 to 63 / 71 lines a layer (rank 0 / rank 1) and +the requests take 123 s / 260 s instead of 132 s / 278 s (prefill plus 32 tokens, one run each). Not instrumented: +per-rank paging counters at tp=2 (the hit/miss numbers above are tp=1); the tp=2 equality is over 2 × 32 tokens. + +## 4. Cost (why it is cheap, and what it is not) + +Decode: +0.5 / +0.9 / +1.2 ms a token at 32k / 65k / 200k on 63-69 ms tokens (+0.7-1.7%). 8.3 MiB at ~12.5 GB/s is +~0.7 ms; the plan/remap launches the rest. Prefill: +0.6% / +2.1% / +4.1% (90 GiB of fetches at 200k ≈ 7.5 s at PCIe +rate, 10 s measured, of 244 s). Memory: −4.09 GiB a 200k stream on the card for 4.6 GiB of pinned host memory. +Not measured: MTP verify windows (the tp=2 serve run had drafts at the server default; only its completions were compared) through the paged cache in the engine (the replay gives 22-33 MiB a 4-row window at +200k, ≈ 2-2.7 ms, *estimate*); decode rounds of several lanes (the hook runs per lane, not exercised here); expert-pool +contention for PCIe beyond what the decode numbers above contain. + +## 5. Limits and next steps + +1. A call whose union of blocks outside the window exceeds the free slots must fail closed. Device-side overflow + signalling and a host-side + check before consuming the result are wired into the forward path; synthetic and 16k TP=2 deep-prefill tests + pass. The planned 200k × 2 stress/overflow-injection gate remains open. Slots must be >= ~24k at 200k with + 256-row prompt calls (replay max 23.6k). +2. Paged prompt-cache snapshots now flush and preserve BF16 host K/V, device windows, table/LRU metadata, recurrent, + indexer, PLE and MTP state. TP2 restore validates and commits both ranks before resuming. Ordinary-context two-lane + warm/branch/rewind cache reuse passes with exact cold-output parity; deterministic failure injection on live ranks + remains open. +3. Paging captures and replays both deep-prefill layer graphs and decode graphs at TP=2. The 16,384-row CUDA test + verifies exact graph-vs-eager state/logit parity; the 12,040-token two-request HTTP test records actual prefill + graph replay and paging snapshots. These establish correctness, not a speed benefit. Paired DPLAN-vs-graph runs + remain within variation; optimize/profile before claiming graphs improve performance. +4. The fetch is on the compute stream. Overlap needs prediction (prefetch for row t+1 from row t's selection), but at + 200k consecutive rows share only half their blocks, so most misses are new blocks; not worth it at +1 ms a token. +5. int8 K/V (K10) is out of scope here (bf16 only, as asked); the paging and int8 would compose (the host store and + slots would halve). + +## 6. Acceptance gates (current requirement) + +BF16 remains the production KV format. RAM paging is a capacity tier for the same BF16 rows, not a reason to switch +to int8. The single-card 200k parity result and the short TP=2 parity run above are prototype evidence only; they do +not satisfy the following acceptance gates: + +1. **Graph performance (immediate):** keep the DPLAN-vs-graph A/B on the same TP2, two-lane BF16-paged workload; + profile capture/replay, host prepare time and eager fallback by phase. Optimize prefill and decode separately, + preserving exact outputs and NCCL overlap. Do not enable graphs by default until repeated runs show a benefit. +2. **Safety:** deterministically inject slot overflow and restore/copy failures on TP2; fail closed before attention + consumes invalid rows, drain both ranks, and prove the session is discarded cleanly. +3. **Cache lifecycle:** retain exact warm/branch/rewind parity and exercise eviction/reload and sequence reuse with + MTP draft rejection/rollback. Existing ordinary-context two-lane tests pass; failure cleanup remains open. +4. **Long context (later):** only after graph and cache reuse are optimized, prove the 200k × 2 TP2 capacity target + with BF16 paging and graphs enabled. Keep a resident-BF16 comparison where it fits; otherwise compare every + fetched/evicted K/V row against the resident oracle in bounded fixtures plus full-request logits/tokens. +5. **Observability:** expose per-rank overflow, hit/miss, pinned-RAM and device-byte counters in the acceptance logs. + +Until these pass, paging stays opt-in and must not be described as production-ready. CPU expert execution and int8 KV +remain out of scope for this BF16 paging milestone. + +## 7. Integration audit (updated 2026-10-20) + +The allocator/admission mismatch is fixed in commit `8528ca2c`: when host indexer keys are enabled, the +placement-aware `newSeq` path honors the BF16 pager instead of allocating resident K/V, and the admission budget +uses the same sequence layout, configured chunk and one-row placeholders. Quantized KV and split-head placement are +rejected when paging is enabled. Later work adds pager-aware snapshot state and TP2 graph replay counters. + +Latest hardware gates (par1-cs13, both 5060 Ti cards, TP=2): graph snapshot/restore parity passes on both ranks. +The 12,040-token two-request HTTP gate passes with BF16 paging, MTP, graph mode and exact warm/branch/rewind cache +parity; each rank logs prefill replays and captures. Four alternating 4,139-token DPLAN/graph runs also preserve exact +cross-mode hashes and cache hits, but do not show a repeatable graph speedup. The 16,384-row deep-prefill graph gate +continues to pass eager parity. This is ordinary-context correctness, not 200k × 2 capacity. + +Remaining blockers: + +- **Graph performance:** repeated same-workload TP2 A/B and phase-level profiling are still needed; graphs are correct + but have not demonstrated a speed win. +- **Overflow stress:** deterministic TP2 overflow/failure cleanup and clean collective drain remain open. +- **Mixed serving:** two-lane cache reuse passes at 12k tokens with MTP and graphs; long-context paging and mixed + failure/rollback under cache eviction remain unproven. + +The next priority is graph profiling/optimization and deterministic TP2 overflow/failure cleanup. The 200k × 2 +capacity run is deferred until those are complete. For resident parity, use replicated attention +(`TF_FLASHNEXT_ATTN_SPLIT=0`) so a different attention reduction order does not confound the cache comparison. + +## 8. Persistent expert-placement priors (2026-10-22) + +Set `TF_FLASHNEXT_EXL3_PROFILE_DIR` to a writable POSIX cache directory to opt in. Each model/TP-rank file stores +decoder and optional MTP lifetime expert counts, not weights, KV, prompts, graph executables or VRAM addresses. +The advisory identity hashes loaded routed trellis bytes and config/index contents; it is not a full-checkpoint +checksum because router/dense weights are excluded. Pool capacity is deliberately excluded, so a profile can seed +a differently sized cache. This does not redistribute experts between TP ranks or implement elastic VRAM budgeting. + +After pool attachment, valid priors decay once by two (nonzero minimum one). Only positive-count local routed +experts warm the existing empty lines, hottest first; copies and pointer-table writes finish before residency is +published. Warming does not invent routing observations. Actual host/device planners keep learning. The separate +deep-prefill ring histogram is not persisted, and a subsequent prefill may discard startup residency: this is an +advisory prior, not an established startup or throughput improvement. + +Graceful engine release drains and pulls the owning device planner before snapshotting decoder/MTP counts and +before freeing graphs or expert pools. Saving replaces the snapshot rather than adding imported counts again. +A nonblocking lifetime lock prevents competing writers. Private directories/files, a bounded versioned format, +SHA-256 and synced atomic replacement protect the last good profile. Missing/corrupt/mismatched files and storage +errors fall back cold or warn; GPU warm-copy failures fail startup. There is no per-token disk I/O or synchronization. +An abrupt kill can lose updates since the previous save; power-loss durability of the directory rename is not promised. + +Validation: 211 EXL3 host tests and native compilation pass; 15 focused persistence/lifecycle tests pass in Debug +and ReleaseSafe on Linux. The source `53ef5037` CUDA fixture passes on both cards with real copied bytes/tables, +device-learned counts, rank isolation and pool-size changes (10 to 23 lines). + +Native source `babc9762e2dbc999c9b68df252d619830203d693`, binary SHA-256 +`86a7b20b6bac2bc5b2b78f9c89654a21cd1c0bd2ad9e6806ee1dc70abe110907`, passes two fresh-process TP2 runs with +MTP, graphs and shared BF16 paging: two waves of two concurrent 1,024-input/64-output requests per process, +prefix retention disabled. Both ranks save, reload, warm and update their profiles; all response hashes match +across restart. Startup warms 1,568 rank-0 and 2,304 rank-1 lines inside the fixed 32/48-line pools. Saved decoder +totals after restart equal the decayed prior plus a fresh observation total on each rank, not a doubled prior. +The harness signals the HTTP leader and lets its TP stop protocol release the follower; killing both processes +directly is not graceful follower shutdown. Existing shutdown allocator warnings are not a clean-teardown claim. + +Routed-bank fingerprinting costs about 16.6 seconds per rank at startup in this gate. No profile throughput gain, +deep-prefill histogram persistence, long-context capacity or independent model fidelity is established by this test. +The profile does not relax any BF16 paging, graph or golden-fidelity gate. diff --git a/docs/plans/exl3-on-upstream-final-report.md b/docs/plans/exl3-on-upstream-final-report.md new file mode 100644 index 0000000000..078dd8bb97 --- /dev/null +++ b/docs/plans/exl3-on-upstream-final-report.md @@ -0,0 +1,1373 @@ +# EXL3 Qwen3.8-Flash-Next on upstream TensorFold: final report + +Branch `exl3-on-upstream-wip` (staging branch `exl3-c-main`), originally based on upstream tag `v1.1.0` and now merged with current upstream `main`. +Machine: par1-cs13, two RTX 5060 Ti 16 GiB (sm_120), 251 GB RAM, pack +`Qwen3.8-Flash-Next-exl3-3.05bpw`. All measurements are test measurements on idle cards under a lock +wrapper, not benchmarks. + +## Current handover (2026-10-23) + +Upstream draft PR [#1113](https://github.com/ashhart/TensorFold/pull/1113) is open from the fork branch +`exl3-on-upstream-wip` into upstream `main`. The branch integrates current upstream at `d1018b318`; exact upstream +commit ancestry is retained. The PR description contains the proposed dependency-ordered split plan. This remains +one large draft and is not merge-ready. + +The user's acceptance target is TP=2 with CUDA graphs on both decode and prefill, and graphs must improve +performance. BF16 is the main KV format; dynamic KV paging to RAM is also required. CPU expert execution is +deferred. Decode graphs match eager for eligible captured shapes, but input preparation, the MTP chain's absorb +step, and host readback remain outside capture. A new serving ABBA shows a narrow no-draft decode gain but a +material MTP-drafted regression; graphs therefore do not yet meet the performance acceptance gate. Prefill graphs +capture each layer's dense sub-blocks; expert scheduling and ring +copies remain eager. A new event-wait-node path preserves overlap with TP=2 sub-block reductions and passes the +bit-equality gate. Across 8192-row prompts with 1024-, 2048- and 4096-row chunks, graph throughput tracks eager +within run variation; a graph-specific prefill win is not demonstrated. An 8192-row chunk OOMed and fell back to +the non-deep path, so it is excluded from graph results. Graph performance remains an open acceptance gate. + +A TP=2 real-text HTTP ABBA on the same native binary, with graphs and shared BF16 paging enabled, measured a small +prefill-only improvement from setting expert-ring lookahead to one layer: 5.1851 s versus 5.2499 s for an 8,192-token +prompt (+1.25% throughput). All 12 measured one-token outputs matched. Lookahead 1 becomes the default, but this is +not a decode gain, concurrent-throughput result, or meaningful progress toward the 2× target by itself. The isolated +synthetic-ID profiler shows a similar ~1% result at 8k and 16k rows, with about 2.8% more expert copies. + +The near-200k × 2 TP=2 BF16 needle test (195,938 and 195,957 prompt tokens) passed on this tree before the final +agent changes; it was not rerun on the exact PR head. Treat that result as demonstrated on a near-final tree, not +as a final-head gate or an exact-200k test. The follow-up implementation now combines BF16 paging, snapshots, TP=2, +graphs, MTP and two concurrent HTTP chat histories at ordinary lengths. Restored graph state/logits match exactly; +warm, branch and rewind cache phases each reuse 4,122/4,139 prompt tokens with exact response hashes versus cold +cache. The 200k × 2 gate remains explicitly deferred. Live-rank overflow/failure injection and graph speedup across +the intended drafted workload remain open. BF16 is required; int8 KV is out of scope. + +A follow-up on `e22c4552a` passes this gate on both RTX 5060 Ti ranks: shared BF16 paging, recycled request allocation, +snapshot restore, exact state/logit continuation, and graph-mode decode. The earlier zero-capture result came from a +128-entry FIFO warmup history evicting the 192-shape working set (two lanes × two chunk banks × 48 layers). Raising +history capacity to 224 allowed each rank to capture 112 graph entries; the bounded graph cache then replayed all 112 +on the next run, including 28 paged-attention replays and four safe generation/origin rebinds. The other 80 shapes ran +eagerly as intended by the 112-entry cache bound. This proves cache reuse and correctness, not a latency gain; graph +speedup remains an open measured-performance gate. + +## 2026-10-23 score-grid graph-key ABBA follow-up + +An existing four-leg TP=2 HTTP ABBA (bucket off/on/on/off) used both RTX 5060 Ti cards, `--parallel 2`, 8,192 real-text +prompt IDs, 512 generated tokens, graph-enabled decode and prefill, shared BF16 KV paging, 38 expert lines, prompt +cache disabled, and drafts both off and on. Each leg discarded a warm-up request, then measured two requests per +draft mode. All requests returned the same full token hash +`600815b14b64d917c86eebf8279154089d014500a6d213246692cf86d918680f`; drafted/undrafted parity held. + +Mean prefill was 5.2437 s with bucketing off and 5.2473 s on (0.07% slower with the candidate). Mean undrafted decode +was 26.911 versus 27.007 tok/s (+0.36%); drafted decode was 25.214 versus 25.381 tok/s (+0.66%). These deltas are +within run variation and do not qualify as an improvement. The candidate reduced main-round graph instantiations +from 8 to 6 in the server summary, while prefill graph counters were unchanged; the additional capture reuse did not +translate into measured end-to-end speed. The score-grid bucket implementation is therefore rejected and removed +from the source tree. This result rules out this bucketing approach as a meaningful performance step; it does not +close the separate shared-snapshot/prefill-replay correctness gate. + +## 2026-10-23 MTP verify-best ABBA + +An exclusive TP=2 HTTP ABBA (`TF_FLASHNEXT_EXL3_VERIFY_BEST` unset/on/on/unset) tested the existing default-off +path that moves each verify row's top-1 selection and asynchronous record copy into the verify tail, avoiding a +second selection in `draws`. Both RTX 5060 Ti cards ran the same binary +(`e3bf1015f9e0d1847d83f5b651bdabefa48c5ecfc0c45d665e7c5197da6827e3`), 8,192 real-text prompt IDs, 512-token +outputs, TP=2, `--parallel 2`, CUDA graphs, shared BF16 KV paging, 38 expert lines, and no prompt cache. Each leg +discarded warm-up requests and measured two plain and two drafted requests. Drafted requests actually proposed 522 +tokens and accepted 403 on every leg; undrafted requests proposed/accepted none. Every output had the same full token +hash `600815b14b64d917c86eebf8279154089d014500a6d213246692cf86d918680f`. + +Mean prefill was 5.1738 s with the option off and 5.1771 s on (0.06% slower). Undrafted decode was 26.960 versus +26.950 tok/s (-0.04%); drafted decode was 25.261 versus 25.259 tok/s (-0.01%). This is no measurable improvement; +leave the option off by default. Graph logs show stable main-round replay, but reported zero MTP-chain graph captures +in that workload despite the accepted drafts, so that run did not establish that the draft chain itself was graphed. +The following short-decode ABBA separately confirms MTP-chain graph replay for its own workload; the difference +between those run configurations remains unexplained. + +## 2026-10-23 clean TP2 short-decode graph/eager ABBA + +A follow-up clean-binary ABBA (`graph/eager/eager/graph`) repeated the short HTTP decode gate on source head +`7f8b7a194`, without temporary chain-tracing instrumentation. Binary SHA-256 was +`a20e5165c138dbbf0192c3fb62aa091acf277aae40eddbfc66c2247efba36385`. Both RTX 5060 Ti cards ran TP=2 with +`--parallel 2`, a 32,768-token window, shared BF16 KV paging, 38 expert lines/layer, blocking NCCL, and prompt cache +disabled. Each leg ran the same golden checks and two 128-token measurements per draft mode. Output hash prefix +`2c5d814cfa80` matched across all runs; drafted runs accepted 87 of 115 proposals in every leg. + +Undrafted decode averaged 30.03 tok/s eager versus 33.72 tok/s with graphs (**+12.3%**). Drafted decode averaged +41.96 tok/s eager versus 43.91 tok/s with graphs (**+4.6%**). Rank 0 reported 188 main-round graph replays and 54 +MTP-chain replays after four captures; rank 1 has no MTP head. This confirms the MTP chain graph captures and +replays on this short greedy workload. It does not explain why the separate 8,192-ID/512-output verify-best ABBA +reported zero chain replays, and it does not supersede that longer drafted regression (eager was 15.6% faster). +Treat this as a workload-specific short-decode improvement, not a general graph win or progress toward 2×. The +prefill timer was about 0.25 s in graph legs versus 1.39 s eager, but prefill graph counters were zero and repeated +fixture requests may reuse warmed request state; exclude this from any prefill claim. TP=2 prefill graph speedup +remains unproven. + +## 2026-10-20 upstream-main integration and release checks + +The merge resolves the shared Flash Next hooks against current upstream `main`. On par1-cs13, Zig 0.17 +`zig build test-cuda-host` passes after fixing a test fixture that left the new paging option undefined; +`zig build native -Dnvcc=nvcc -Dsm=120 -j16` also passes. `zig fmt --check zig/src zig/tests` passes after formatting +the three reported Zig files. These checks validate compilation and host tests only; the final-head TP=2 serving, +graph-performance, BF16 paging/snapshot/cache-reuse, and 200k × 2 gates remain to be rerun. + +The lean checker does not pass: current upstream `main` has 444 findings and this branch has 1,761, leaving +1,317 new findings. Treat this as an explicit CONTRIBUTING blocker; the PR stays draft. On 2026-10-20, removed +118 tool co-author/session trailer lines from the 76 feature commits using a commit-object rewrite scoped to the +feature range. Original author identities, subjects, complete tree, and exact upstream `main` commit IDs are +preserved. The branch history is now grouped into five implementation commits, one upstream integration merge, +and one focused documentation commit. The draft omits the obsolete v1.1.0 and int8-first planning memos. + +## 2026-10-20 prefill graph profile addendum + +The new `prefill-prof` gate ran TP=2 on both GPUs with BF16 KV paging enabled, 8,192 rows, 4,096-row chunks, two +warm-ups, and one measured forward per mode. Both eager and graph runs returned the same cross-rank logits hash +(`ffcf463b229c2aa6`); graph mode replayed 96 dense-layer graphs. One matched untraced pair measured 5,597.1 ms +eager and 5,639.5 ms graph. In the matched Nsight pair, rank 0 measured 5,693.6 ms eager and 5,696.4 ms graph. +Neither pair establishes a speedup. + +Nsight shows the graph replacing eager kernel submissions (20,559 `cuLaunchKernel` calls versus 2,103 plus 96 +`cuGraphLaunch` calls) while expert-copy volume and kernel/collective totals remain nearly unchanged. Each rank +transferred about 28.8 GB H2D in approximately 48k operations; rank 0 scheduled 14,426 expert copies (13,993 counted +as prefetch) and 570 passes, with about 35 ms CPU planning. The measured picks-D2H-to-synchronize interval is about +5.2 s cumulatively, but includes the preceding dense kernels and NCCL dependencies; it is not pure CPU stall time. +The next prefill performance work is therefore a proper expert-routing/transfer pipeline, not another capture-threshold +change. The design investigation should first establish staging-arena memory fit after BF16 paging, pools, scratch, +and snapshots; then prototype GPU-built stable member lists and fixed-buffer double-buffered transfers while keeping +existing expert arithmetic and reduction order. + +## 2026-10-22 TP2 prefill profile and reduction-overlap A/B + +At head `5a94359c`, Nsight Systems profiled a 16,384-row TP2 prefill with graphs, BF16 KV paging, 4,096-row deep +chunks, two warm-ups and one measured forward. Both ranks replayed 168 dense-layer graphs, captured 24 shapes, and +reported the same logits hash (`662339a95f0de2d7`). The measured test-harness wall time was about 11.3 s per rank; +this is not a production HTTP benchmark. In the per-rank summed kernel durations, NCCL ring all-reduce accounts for +24–27%, `tf_exl3_prompt_2` about 22%, GDN about 8%, and attention about 6%. The CUDA memory-operation report shows +about 5.0–5.1 s of cumulative host-to-device memcpy durations per rank. Those intervals can overlap other GPU work; +they are not isolated expert-miss stall time. The kernel-duration sums are likewise not additive wall time. + +A follow-up, non-Nsight ABBA test compared only `TF_FLASHNEXT_EXL3_PREFILL_OVERLAP` (1, 0, 0, 1) using the same +binary and graph/paging configuration. Each run used 16,384 rows, 4,096-row chunks, two graph warm-ups, and the +same output hash. Graph replay counts, expert-copy counts and planning time were unchanged. Overlap-on wall times +were 11,215.0 and 11,212.1 ms; overlap-off times were 12,475.7 and 12,472.0 ms. The on-pair mean is 11,213.6 ms +versus 12,473.9 ms off, a 10.1% wall-time reduction (11.2% higher rows/s). This is a repeatable targeted prefill +optimization, not a 2× result and not proof of eager-vs-graph speedup. It supports keeping sub-block reductions +overlapped; next compare collective grouping and remove the post-reduction placement/writeback work only behind +exact rank and graph parity gates. + +An opt-in reduced-branch follow-up (commit `99265db1`, `TF_FLASHNEXT_EXL3_PREFILL_REDUCED_BRANCH=1`) skips the +post-all-reduce `tpPlace` and feeds the contiguous FP32 sum directly to HC writeback with one unit-weight slot. +The default remains off. The targeted `flashnext-hc` GPU regression passed the new exact hidden/FP32-sum, +signed-zero, BF16-tie and short-tail checks. A separate graph-on, BF16-paged 16k TP2 ABBA run (flag 0, 1, 1, 0) +measured 11,207.6 ms off and 11,033.4 ms on, about 1.6% higher rows/s. All four runs replayed 168 graphs per rank +and returned hash `662339a95f0de2d7`; expert-copy counts and planner time were unchanged. This is a small measured +gain, not a 2× result. Snapshot/cache continuation under the opt-in mode and eager-vs-graph parity remain to be +qualified before enabling it by default. + +A four-run 4,096/6,144/6,144/4,096 ABBA comparison used 16,384 rows, TP=2, graphs, BF16 paging, overlap on, and +reduced writeback off. Mean wall time was 11,210.7 ms at 4,096 rows and 10,757.7 ms at 6,144 rows (4.0% lower wall +time, 4.2% higher rows/s). All eight rank-runs returned hash `662339a95f0de2d7`; replay counts were 168 versus 120. +Rank 0 expert copies fell from 25,588 to 21,648 (-15.4%) and passes from 1,144 to 845 (-26.1%). This is a stable, +small chunk-size gain, not a 2× result. Larger chunks consume more scratch/memory and 8,192 already OOMed on this +configuration. + +## 2026-10-22 prompt chunk and decode-miss profile + +On fork head `06d58b2e`, increasing the EXL3 prompt chunk from 256 to 2,048 rows reduced the matched 8,192-row +TP=2 graph + BF16-paged prefill from the prior ~5.48 s baseline to 5.206 s (about 5% lower wall time). Two current +repeats measured 5,206.4/5,206.1 ms and 5,205.9/5,205.5 ms on ranks 0/1, with the same hash +`ffcf463b229c2aa6`, 72 graph replays, and no eager layers. A 32,768-row run completed in 21,853.8/21,853.6 ms +with matching hash `5894636da2406e08` and 360 graph replays. These results demonstrate the chunk-size setting works +at both shapes; only the 8,192-row result has a matched prior-chunk timing. This is not a graph-on/off comparison. + +A separate 16,384-context, 64-step TP=2 decode profile with BF16 paging compared eager DPLAN and graph replay in the +same harness. It passed output-hash parity (`fba5a0bab565af03`) on both ranks. Graph measured 24.30 ms/token versus +24.26 ms/token eager (+0.16% latency, effectively no change); it reused an existing graph for all 64 steps with no +captures or executable updates. The measured graph segment recorded 14,017 expert hits, 1,135 misses/copies across +3,072 passes, and 4,575 KV-page misses (300,405 required page slots; 768 calls; no rebases or overflow). These +counters identify expert and KV miss frequency, not their exposed stall time. Graph replay correctness is confirmed; +a decode graph speedup is not. + +## 2026-10-23 TP2 prefill reduction-overlap and chunk ABBA + +On fork source head `85df2b43d`, an exclusive four-leg ABBA compared the existing overlapped per-sub-block NCCL +reduction path (default) against `TF_FLASHNEXT_EXL3_PREFILL_OVERLAP=0`. Both RTX 5060 Ti cards ran the same 8,192-row +TP2 prefill with BF16 KV paging and prefill graphs enabled. Each leg replayed 72 dense-layer graphs, had zero eager +fallbacks, and returned cross-rank logits hash `ffcf463b229c2aa6`. Rank wall times were 5,056.2/5,056.1 ms and +5,059.2/5,059.2 ms with overlap on, versus 5,886.9/5,886.6 ms and 5,884.8/5,884.6 ms with overlap off. The paired +means are 5,057.7 ms and 5,885.7 ms: overlap reduced latency 14.1% (16.4% more rows/s). Expert copies, passes and +planning time were unchanged. This confirms the existing overlap implementation is valuable; it is not a new code +change, a graph-on/off comparison, or a 2x result. + +A second four-leg ABBA on the same source and workload tested `TF_FLASHNEXT_CHUNK=2048` against the default 1,024. +All legs passed graph/paging and output-hash checks. The 1,024-row legs averaged 5,055.6 ms; 2,048-row legs +averaged 5,300.7 ms (4.85% higher latency). Reject this larger chunk setting for this 8k workload. This is +consistent with the earlier 256-to-2,048 improvement being workload/starting-point specific; it does not justify +globally increasing the current default. + +An additional counter-only telemetry run attributed rank 0/1 cross-layer prefetches to 13,098/13,719 copies, of +which 4,712/4,778 (36.0%/34.8%) were evicted unused; their logical payloads were 8.69/8.81 GB. The telemetry +sampling mode adds a synchronization boundary, so its wall time is not used. A separate uninstrumented four-leg +ABBA compared current `ring_ahead=2` with `TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD=0`: mean wall was 5,060.0 ms with +lookahead and 6,409.9 ms without it (+26.7% latency when disabled), despite 6.5–6.9% fewer expert copies. Every +run replayed 72 graphs and returned the same logits hash. Thus prefetch traffic is not safely removable wholesale; +the next opportunity is to improve which experts are prefetched, while retaining the asynchronous lookahead. + +## 2026-10-23 TP2 expert-prefetch lookahead ABBA + +An exclusive TP2 `prefill-prof` ABBA on both RTX 5060 Ti cards compared `ring_ahead=1` against `ring_ahead=2` at +8,192 and 16,384 rows. Both tests enabled BF16 KV paging and CUDA prefill graphs; warm-up/capture ran before the +measured forward. Inputs are deterministic PRNG IDs (seed `0x9e3779b9`), so this is an isolated schedule profile, not +a real-text performance receipt. Each rank returned the same logits hash for both settings (`ffcf463b229c2aa6` at 8k, +`662339a95f0de2d7` at 16k), with 72 and 168 graph replays respectively and zero eager fallbacks. + +At 8k, mean slower-rank wall time was 5,005.6 ms with lookahead 1 and 5,058.0 ms with lookahead 2 (+1.05% rows/s for 1). +At 16k, it was 10,257.1 ms and 10,367.4 ms (+1.07% rows/s for 1). The result repeats across these two tested +shapes. Lookahead 1 made about 2.8% more expert copies, so the profile suggests improved overlap/wait time rather than +less traffic or fewer misses. + +A real-text HTTP follow-up used source head `60c92357b` plus the one-line default change, the same native binary +throughout (`e3bf1015f9e0d1847d83f5b651bdabefa48c5ecfc0c45d665e7c5197da6827e3`), Qwen3.8-Flash-Next EXL3 pack +fingerprint `e98f5301e6eb58ab48a117e892f71aecd19b449cec5513a620c53cdcb95ddce4`, two RTX 5060 Ti 16 GiB (SM 12.0), +Linux 6.8.0-147-generic, NVIDIA driver 615.71.09, CUDA 13.4, and Zig 0.17.0. Each server was TP=2, context 32k, +parallel capacity 2, 38 expert lines, shared BF16 KV paging, CUDA decode/prefill graphs, drafts off, and prompt cache +off. The request was serial (not concurrent): first 8,192 IDs from the same locally held real-text fixture +(SHA-256 `b74a9216bb35cf88763989f4fd30ce93dfd726d033403b5c8c268b2b96c2d9b8`), one generated token. The first +request per server was discarded as graph warm-up; three fresh measured requests followed. Four exclusive server +legs ran lookahead 1/2/2/1. Mean server-reported prefill was 5.1851 s at 1 and 5.2499 s at 2: **1.25% higher +prefill throughput** at lookahead 1. All 12 measured outputs had token hash +`c6f3ac57944a531490cd39902d0f777715fd005efac9a30622d5f5205e7f6894`; both graph flags were enabled and server logs +confirm shared BF16 paging. The HTTP server does not emit per-request replay counts. The same binary's synthetic +`prefill-prof` gate separately logged 72/168 graph replays at 8k/16k. This confirms a small, repeatable prefill gain on one real-text prompt, +not a decode win, concurrency throughput result, 2× result, or snapshot/prefix-reuse gate. The runtime default is +reduced to 1; `TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD` remains the override. Snapshot/cache behavior was not changed +and is not re-claimed by this benchmark. + +## 2026-10-23 TP2 reduced HC writeback experiment + +An exclusive four-leg ABBA compared the existing TP2 HC writeback (`TF_FLASHNEXT_EXL3_PREFILL_REDUCED_BRANCH=0`) +with the reduced-writeback path (`=1`) on both RTX 5060 Ti cards. The workload was an 8,192-row graph-on prefill +with BF16 KV paging, 1,024-row chunks, overlapped NCCL reductions, and identical rank/output hashes +(`ffcf463b229c2aa6`). Default legs averaged 5,062.2 ms; reduced-writeback legs averaged 4,958.6 ms: **2.05% lower +latency / 2.09% higher rows/s**. All legs replayed 72 prefill graphs with zero eager layer fallbacks; expert copies, +pass counts, and planner behavior matched. This is a promising measured opt-in prefill improvement, not a 2x gain, +not a serving/decode result, and not yet the default. The reduced path changes HC writeback arithmetic; snapshot/ +restore on private BF16-paged lanes also passed exact continuation and graph replay. The matching shared-paging +cache-reuse gate failed with `GraphBudgetExceeded` during decode graph replay after restoring shared snapshots; +both ranks matched the paired BF16 logits through step 2 before the failure. The option remains opt-in, and shared +snapshot/cache reuse with graphs is not yet validated. The failure is being investigated as graph-budget/cache +admission behavior, not treated as evidence of a numerical mismatch. + +## 2026-10-23 TP2 HTTP graph/eager decode ABBA + +An exclusive off/on/on/off serving comparison used the same rebuilt binary on both RTX 5060 Ti cards, TP=2, +`--parallel 2`, BF16 KV paging with shared pages, 38 expert-pool lines per layer, prompt cache disabled, and a +16-token completion prompt with a 512-token output. Each server was restarted per leg; the gate first ran its three +golden fixtures, then measured two undrafted and two drafted completions. Thus this is a two-stream-capacity server +but serial HTTP timing (not concurrent-request aggregate throughput). All four measured outputs used hash +`b170e8433207`, and drafted/undrafted results matched. Graph-on runs replayed the stable round graph 477 times, +the MTP chain 21 times, and captured four shapes; prefill graph counters stayed zero for this short prompt. + +Mean decode throughput was 32.36 tok/s eager versus 34.79 tok/s graph-on without drafts (+7.5%), but 49.17 tok/s +eager versus 41.51 tok/s graph-on with drafts (−15.6%). The four ABBA legs consistently show that the current +graph-enabled drafted path is slower, despite replay. The server reported first-completion prefill/TTFT around +1.39–1.43 s in eager legs and 0.26–0.32 s in graph legs; because this short prompt triggered no prefill graph and +these values varied with server warm state, do not treat them as a prefill-graph result. This test does not establish +a prefill graph speedup or overall multi-request improvement. Retain graphs as opt-in and profile the MTP +verify/absorb and host boundaries before changing defaults; the TP=2 graph performance requirement is not met yet. + +## Cumulative performance context + +The earliest recorded one-row kernel/decode-step optimization reduced a test-harness decode step from 2.4–3.0 s to +0.41–0.45 s (roughly 5.3–7.3×); that work is already part of the pushed baseline and is not a new gain from the +current chunk-size change. Historical prefill notes report 100–210 rows/s for the old batched pass, versus 1,225– +1,555 rows/s in the later TP=2 deep-prefill runs (nominally 5.8–15.6×). The old and new numbers were not collected +as a controlled same-binary ABBA experiment, so this range is directional, not a causal speedup claim. The named +`5a94359c`→`03d4dbf3` controlled TP=2 prefill ABBA measured **4.90% more rows/s** (16,384 rows, 4,096-row chunks, +BF16 paging, graphs on) with identical output hashes. A later lookahead ABBA measured +1.25% 8k prefill throughput; +reduced HC writeback measured +2.09% 8k prefill rows/s but remains opt-in because shared-snapshot graph replay +failed. The clean short-decode ABBA measured graph-on gains of 12.3% undrafted and 4.6% drafted at 128 output +tokens; a separate 512-token ABBA measured +7.5% undrafted but −15.6% drafted. These are workload-specific, not +cumulative gains and do not establish 2×. Historical prefill records are roughly 5.8–15.6× directional and +one-row harness decode 5.3–7.3×, but those are not controlled same-workload comparisons. There is no defensible +single cumulative end-to-end decode speedup from project start. + +## 2026-10-23 batched EXL3 dispatch and TP=2 two-stream gate + +The batched decode keep-buffer path was calling the affine expert path for EXL3 layers, whose affine pointers are not +initialized by the EXL3 loader. The focused host test also requires an unattached expert pool to fail before any +keep-buffer copy. The fix now preflights `needsPool()`, copies the requested keep buffers, and dispatches through +`exl3_pass.moe`. On the current fork head plus this change, `zig build test-flashnext-exl3-host` passes and +`zig build native -Dnvcc=nvcc -Dsm=120 -j16` passes (one existing unused-variable nvcc warning). + +The rebuilt binary passed an HTTP TP=2 `--parallel 2` decode-graph test on both RTX 5060 Ti cards at 16,384 context +with BF16 KV paging, 68 expert lines/rank, two planned streams, a 0.5 GiB memory reserve, drafts off, and prompt +cache disabled. Two 128-token concurrent requests for each of two prompts, repeated twice, matched their solo token +hashes (4/4 per prompt); the serial controls also matched (1/1 per prompt), with no failed or unmeasured requests. +Both ranks logged decode graph replay. This test had no prefill graph replays and is not a performance comparison; +the line-count/reserve settings are admission accommodations, not defaults. Attempts at 76 lines and at the 76-line +single-stream graph budget were rejected at startup for insufficient memory; reducing the pool to 68 lines admitted +the two-stream run. + +A separate 1 GiB prompt-cache run through `tools/bench_cache.py` used the same TP=2/two-stream BF16-paged graph setup, +a 2,571-token completion prompt, 64 output tokens, two repeated requests and two concurrent requests. All outputs +matched the first response hashes, but every request reported `cache_state=cold`, `cached_tokens=0`, zero stored +states and zero hits; the strict cache gate therefore failed qualification. This completion fixture does not +invalidate the previously passing chat warm/branch/rewind cache gate, but cache reuse for this fixture remains +unproven and the difference in request path needs investigation. No speedup is claimed for this correctness fix. + +## Matched 5a94359c TP2 prefill comparison (2026-10-22) + +An exclusive four-run ABBA comparison (baseline/current/current/baseline) tested the pushed `5a94359c` source and +current `03d4dbf3` source on both RTX 5060 Ti cards. Each sample measured one 16,384-row deterministic prefill +forward with TP=2, CUDA prefill graphs, BF16 KV paging, blocking NCCL, 4,096-row deep chunks and overlap enabled. +The baseline averaged 13,165.4 ms; current averaged 12,550.1 ms: **4.67% lower latency / 4.90% higher rows/s**. +All four rank-pairs replayed 168 graphs and produced the same cross-revision logits hash +`662339a95f0de2d7`; expert copies and pass counts were unchanged. This is a measured prefill gain over the named +baseline, not a serving or multi-context result, not attributable to one isolated commit, and far below 2×. It +does not establish decode improvement. The matched graph-decode ABBA completed: the steady phase measured +22.50/22.47 ms per round on baseline and 22.45/22.50 ms on current (means 22.485 vs 22.475 ms); four-row windows +measured 36.66/36.59 ms baseline and 36.65/36.64 ms current. Each run captured two graphs and replayed 252 times +in the steady phase. The average round-time difference is 0.04%, noise-level rather than an improvement. This +profile used a fixed 16-token prompt and 64 rounds; BF16 paging was enabled, but the short context caused no KV page +misses. An earlier exploratory run selected eager mode by mistake and is excluded from this graph-on comparison. + +### Quantified gap to the 2× objective + +On these matched workloads, doubling prefill throughput against `5a94359c` requires reducing 16k prefill latency from +13.1654 s to at most 6.5827 s (about 2,489 rows/s); the measured `03d4dbf3` result is 12.5501 s, so another 47.5% +latency reduction is required. Doubling graph decode throughput requires reducing the measured 22.485 ms/round to at +most 11.2425 ms; `03d4dbf3` measured 22.475 ms, essentially unchanged. The decode fixture is short-context with no +KV page misses, and the prefill fixture is a warmed synthetic forward; these thresholds quantify those fixtures only, +not long-context HTTP serving. No current result establishes the 2× objective. + +## 2026-10-22 16k paged graph/eager miss diagnostic + +A current-head diagnostic on `d467247b` used TP=2, BF16 paging, a 16,384-token repeated-QBench-ID fixture, 64 +decode steps, and an eager/graph/graph/eager sequence. This is a synthetic repeated-token stress profile, not a +real-text or serving benchmark. Mean measured latency was 34.735 ms/token for eager and 34.650 ms/token for graph +(0.25% lower, within noise). All four runs returned hash `4c2d8973e43477cd`; graph runs replayed an already captured +graph 64 times with no captures or updates. Rank 0 recorded 3,616 expert misses/copies, 230,795 KV-page misses in +768 page calls, and no rebases or overflow in every run. The controlled ordering confirms separate miss telemetry +and graph/eager output parity under paging, but no graph speedup. + +## 2026-10-22 persisted expert-placement A/B + +An orderly-shutdown TP=2 HTTP gate verified that both ranks save and restore their lifetime expert-placement profiles +across process restarts. Rank 0 warmed 2,942 expert lines and rank 1 warmed 3,140. A short-prompt ABBA comparison +then used the same 128-token input and 64-token greedy output, TP=2 on both RTX 5060 Ti cards, CUDA graphs enabled, +BF16 KV paging enabled, and identical full output hashes (`c538c2913c77b8713a3f73bf289d20e253144aa010ae46452329fe19d5fdb74a`). +The warm-profile runs averaged 1,767 ms TTFT and 27.72 ms/token decode; the two independent cold-profile runs averaged +1,943 ms TTFT and 29.63 ms/token. This is 9.0% lower TTFT and 6.5% lower decode latency for this short-prompt, +no-drafts workload. It is a measured cache-placement gain, not a 2x result and not evidence for long-context or +drafted decode. Each restart also spent about 16.6 s per rank fingerprinting routed weights before serving; that +startup cost is outside request latency and remains an optimization target. No production default or source change +was made for this experiment; the existing opt-in persisted placement path was measured as implemented. + +Two follow-up optimization screens do not qualify for production integration. First, an experimental GPU-built +worklist on an 8,192-row TP2 graph prefill compared host/GPU/GPU/host planning. Mean wall time was 5,159.7 ms for +the host planner and 5,114.0 ms for the GPU worklist (0.89% lower latency); all four runs and both ranks matched +hash `ffcf463b229c2aa6`, with 72 graph replays, 12,513 copies and 570 passes per sample. CPU planning fell from +18.35 to 2.30 ms, but the path still downloads counts and synchronizes every layer; measured pick-wait intervals +remained about 4.77–4.81 s and include GPU dependencies. KV paging was enabled but had zero page misses at this +8k shape. The small end-to-end delta does not justify integrating the prototype as the requested staging pipeline. + +Second, an isolated SM120 asynchronous activation double-buffer kernel candidate was tested in TP2 graph prefill +at 16,384 rows and 4,096-row deep chunks, with `prompt_chunk=2048`, BF16 paging enabled and greedy decode disabled. +The baseline/candidate/candidate/baseline means were 10,711.4 ms and 10,677.4 ms (0.32% lower latency). All runs +reported hash `4b9f5c87b87ad84a`, with 96 graph replays and 48 captures per rank. The candidate dispatched only for +the 2,560×10,240 mul1 5-bit projection; other shapes fell back to the existing kernel. The delta is too small to +integrate, and the isolated prototype remains outside the pushed branch. + +Third, an experimental exact-order TP2 prefill collective-submission grouping was tested on the same 16,384-row, +graph-enabled workload with BF16 paging and four exclusive ABBA samples (group size 1, 2, 2, 1). Both ranks +preserved hash `662339a95f0de2d7`; each run replayed 168 graphs and retained identical expert-copy/pass counts. +Group size 1 averaged 12,538.7 ms, while group size 2 averaged 13,391.1 ms (6.8% higher latency). The profiled +pick-wait interval also rose from about 11.87 s to 12.71 s. Host tests and CUDA build passed, but GPU results do +not support the change, so the prototype remains unintegrated and is not pushed. + +Fourth, an exact-order 32-column-fragment EXL3 linear candidate for K=2560, N=10240, 5-bit mul1 was validated on +both RTX 5060 Ti cards. Four ABBA trials per card, with 128 iterations per sample, covered M=1 and M=4, graph and +eager replay, and resident and rotating weights. Exact output parity passed for F16/BF16/FP32 inputs and outputs, +bias, and both supported trellis layouts. In graph replay, resident-weight latency fell from 44.82 to 26.74 us at +M=1 (40.3%) and 45.35 to 27.46 us at M=4 (39.4%); rotating-weight latency fell from 61.23 to 42.89 us at M=1 +(30.0%) and 61.58 to 43.16 us at M=4 (29.9%). This isolated kernel gain did not translate to serving: a TP2 +HTTP ABBA (baseline/candidate/candidate/baseline), with graphs and BF16 K/V RAM paging enabled, generated 128 +greedy tokens from the 16-id prompt. All runs had hash `2c5d814cfa80`, and drafted/undrafted outputs matched. +Undrafted decode averaged 42.81 tok/s baseline versus 43.41 candidate (+1.4%); drafted decode averaged 56.33 +versus 56.28 tok/s (-0.1%). Prefill graphs were not exercised by this short-prompt gate. The end-to-end differences +are within run variation. A follow-up used the same four-server ABBA design and 16-id prompt but generated 512 +tokens per measured request. With the same decode-graph/BF16-paging configuration, all output hashes were +`b170e8433207` and drafted/undrafted outputs still matched. Undrafted decode averaged 43.08 tok/s baseline versus +43.53 candidate (+1.0%); drafted decode averaged 56.85 versus 57.01 tok/s (+0.3%). Decode graphs replayed in every +leg; prefill graphs did not run. The active BF16 page window was 4096 tokens, so this was not a test of KV page +misses. The longer output does not change the conclusion: the end-to-end differences remain within run variation, +so the runtime dispatch stays out of the PR; the isolated candidate is not promoted. + +The CPU AVX2 investigation is not ready for integration. A single-core, pinned feasibility harness executed a full +synthetic 3-bit expert and matched the established F64 oracle bit-for-bit, but took a median 2.35 s/expert; most time +was materializing rotated F64 weights. That is about 17,400× the 135 us estimated H2D payload floor, so this +materializing implementation is rejected. It is not a CUDA-parity result or an impossibility result for a fused +packed AVX2 kernel. A TP2 decode-gate PCIe telemetry sample on PCIe 4.0 x8 links measured peak Rx at 3.4 GB/s on +rank 0 and 2.5 GB/s on rank 1; during the active decode segment the averages were 1.8 and 0.85 GB/s respectively. +These are well below nominal x8 link bandwidth, so the sample does not support PCIe saturation as the current decode +bottleneck. Nsight Compute counters required admin access (`ERR_NVGPUCTRPERM` for the unprivileged run); collected +on card 0 for the synthetic 2560x10240 dense linear, baseline and no-prefetch each read about 16.44 MB from DRAM. +The reported elapsed DRAM-throughput metric was 56.69% of peak for baseline and 66.35% for the candidate. This +does not show HBM saturation for that kernel, but it is not a whole-model or KV-paging profile. A routed expert is +about 1.758 MiB packed and 9.83 MFLOPs per decode row. Any further CPU candidate must reproduce CUDA's +FP32/MMA/activation rounding, include activation round-trip, scheduling and synchronization, and beat the current +GPU miss route end-to-end. Start with sparse decode +misses, not prefill (where an uploaded expert is reused across many rows); never skip or approximate a routed expert. +If RAM bandwidth is already occupied by KV paging, CPU execution may make things worse. + +## Decode critical-path profile and next kernel experiment + +An Nsight Systems trace from the current tree (`a939413c`, with the inspected EXL3 linear, expert and TP2 collective +sources unchanged from `5a94359c`) measured 64 warmed single-token rounds and 64 four-row verification windows. +The short-context single-token phase averaged 23.10 ms/round, with about 22.01 ms between first and last GPU kernels; +the verification windows averaged 38.73 ms. These test-harness intervals exclude MTP draft generation/acceptance. +Although KV paging was configured, the trace contained no `kvp_*` fetch kernels, so it does not characterize long- +context paging misses. + +The largest active decode costs were approximately 5.7 ms/round in dense EXL3 projections and 4.3 ms in BF16 HC +matrix products. The repeated `tf_exl3_linear_10_2_8` launch (K=2560, N=10240, 5-bit mul1, one row) used 80 CTAs, +256 threads and 147 registers/thread, at roughly 69 us per call across 36 calls/round. A kernel-only no-prefetch +candidate keeps launch geometry, MMA, split/reduction and rounding unchanged. On both RTX 5060 Ti cards it passed +byte-exact output checks at 1 and 4 rows. Seven ABBA trials measured graph replay at 63.5 us baseline versus +58.0-58.2 us candidate for one row (-8.4 to -8.7%), and 63.7-63.8 versus 58.2-58.3 us for four rows +(-8.5 to -8.7%). Rotating weights improved 9.2-9.4% and 9.6-9.8%, respectively; warm-weight latency improved +about 28%. The default-off runtime dispatch was then exercised through the TP2 serving gate with graphs enabled. +Baseline and candidate preserved the same `2c5d814cfa80` output hash, and drafted/undrafted outputs matched on every +prompt. Across the two 128-token samples, plain decode was 41.93/46.41 tok/s baseline and 42.05/46.81 candidate; +drafted decode was 63.40/66.47 baseline and 63.02/66.39 candidate. This is within run variation and provides no +reliable end-to-end improvement. Keep the candidate default-off and do not promote it into the PR absent a larger, +repeatable serving win; the isolated kernel gain does not move the 2x objective materially. + +Per-rank mapped-copy kernels occupied 2.11/2.58 ms per round, grouped expert kernels 2.25/2.67 ms, while all-reduce +kernel sums were 3.60 ms on rank 0 and 2.66 ms on rank 1. This suggests rank waiting/imbalance, but does not isolate +NCCL transport. Recorded H2D DMA was only about 349 KB / 0.19 ms during the single-token phase because mapped +expert loads occur inside kernels; those DMA counters do not establish PCIe saturation. Even halving both dense +EXL3 and HC active time would indicate only roughly 1.3× before secondary effects. The 2× objective therefore +requires a combined system-level improvement; no single bottleneck or current optimization proves it. + +A 128-row dense-prompt kernel prototype was also tested directly on the RTX 5060 Ti. Its GPU fixture passed exact +output, FP32/BF16, tail and memory-guard checks across all four traced projection shapes; the candidate used 128 +registers and 34,816 bytes of shared memory with no local spill. However, paired timings (seven samples, 32 launches +per window, eager and graph, warm and rotating weights) were slower than the existing 64-row tile on every shape: +about 1.01× to 1.31× its latency. It is rejected and is not integrated or pushed. Reusing more decoded weights per +CTA increased pressure enough to erase the expected benefit. A 64-row activation-pipeline alternative remains a +prototype investigation; no dense-kernel speedup is claimed. + +An opt-in BF16 HC up-projection/mix fusion was screened with a standalone CUDA exactness test on sm_120. It passed +byte-exact projection, rounded-product, mixed-output and FP32 reduction checks, including changed-input graph replay. +Four eager/captured ABBA pairs on synthetic nonzero BF16 weights measured baseline versus fused per-call latency at +rows 1/4/8/16: eager 7.45/7.64/8.13/10.09 us versus 9.15/9.43/10.82/16.39 us; graph 8.22/8.22/10.25/12.28 us +versus 10.27/10.41/12.31/17.55 us. The candidate lost at every width and is rejected; these are isolated synthetic +kernel timings, not model-weight or serving measurements. The one-launch design is not sufficient to offset its +resource/occupancy cost on the RTX 5060 Ti. + +## 2026-10-23 expert-copy and overlap follow-up + +An isolated `fn_exl3_dcopy` grid-size screen used the existing CUDA test fixture on both RTX 5060 Ti cards. For 10 +expert lines (17.58 MiB), 36–576 blocks measured 1.458–1.466 ms (12.6 GB/s); the equivalent pinned H2D DMA fixture +measured 1.352 ms (13.6 GB/s). Results were effectively identical across cards and block counts, so changing copy +kernel grid size is rejected. The transfer path is near the measured link ceiling; reducing bytes/misses is more +promising than adding blocks. + +The existing side-stream split that computes resident experts while mapped misses copy was tested graph-on in a +four-leg TP2 HTTP ABBA (off/on/on/off), with shared BF16 paging, 38 lines/layer, `--parallel 2`, blocking NCCL, +prompt cache disabled and two 128-token samples per draft mode per leg. All outputs matched hash `2c5d814cfa80`, +golden responses matched and MTP acceptance stayed 87/115. Undrafted means were 35.07 tok/s baseline versus 34.96 +overlap; drafted means were 44.75 versus 45.13 tok/s. The -0.3% / +0.8% differences are within run variation, so +overlap remains opt-in and is not promoted. The short prompts produced no prefill graph replays. + +An opt-in `TF_FLASHNEXT_CACHE_TELEMETRY=1` run on the same short-context configuration collected per-step cache +counters (the boundary synchronizations add overhead, so its tok/s values are not performance measurements). Across +323 decode and 107 verify samples per rank, device expert-pick miss fractions were 15.4% on rank 0 and 17.5% on +rank 1; BF16 KV block misses were zero. This confirms meaningful expert-cache misses but does not exercise long-context +KV paging. Next, measure per-layer miss distributions and compare fixed versus frequency-informed line budgets under +the same graph-on TP2 ABBA; retain only a repeatable serving gain with exact outputs and unchanged BF16 paging/cache +semantics. CPU AVX2 is not the next step: transfer utilization and the prior exact-but-slow CPU prototype do not +support it. + +A separate opt-in telemetry diagnostic then submitted two concurrent 30,000-ID prompts at TP2 with shared BF16 KV +paging, context 32,768 and 38 expert lines/layer. The prompt IDs were formed by repeating two golden-fixture ID lists; +this tests paging mechanics, not natural-text quality. Both requests completed and both ranks reported page fetches and +rebases. Across 60,000 prompt rows, each rank counted 4,284,422 BF16 block requests and about 302,000 misses (7.05%, +about 2.47 GB logical fetch payload per rank); the two-token verify phase counted 678/684 misses and about 5.55/5.60 +MB per rank. This is direct evidence that BF16 KV paging reaches the RAM-backed sparse-fetch path at long context. +Telemetry adds synchronization, so no throughput claim is made. Decode graph mode was enabled, but this short output +did not produce a replay counter; prefill graphs were not enabled, so this run is not graph-performance evidence. + +An opt-in per-layer extension to that telemetry reused the existing planner counters and reported each rank's +routed hit/miss picks by decoder layer. On a TP2 graph-on HTTP run with `--parallel 2`, 32,768-token configured +context, the 128-token HTTP gate, and prompt cache disabled, aggregate pick-miss rates were rank 0: prefill +20.05%, decode 9.72%, verify 7.56%; rank 1: prefill 18.79%, decode 8.47%, verify 5.80%. Highest decode miss counts +in this short fixture were rank 0 layer 1 (339) and layer 6 (248), and rank 1 layer 0 (328) and layer 21 (240). +These are aggregate diagnostic counts, not a natural long-context distribution; opt-in telemetry synchronization +means the run is not a performance benchmark, and it did not measure KV paging misses. Drafted and undrafted outputs +matched on every prompt; two of six prompts did not match the strict golden check. Next use these counts to screen a +fixed-total, minimum-floor per-layer line-budget allocator, then validate under an uninstrumented TP2 ABBA with exact +output and BF16 paging/cache-reuse gates before keeping any policy change. + +## 2026-10-23 persisted expert placement under BF16 paging + +A fresh per-rank placement profile was collected from the same six-prompt fixture and then compared against cold +placement in an uninstrumented TP2 ABBA. Both runs used the same server binary, `--parallel 2`, 32,768-token context, +graphs enabled, BF16 KV paging/shared allocation enabled, prompt cache disabled, and fixed pool budgets of 38 lines +on rank 0 and 49 on rank 1. Each gate requested 128 output tokens; the client benchmark requests were sequential, +so this is not a concurrent-request benchmark. Full serving output hashes were represented by gate hash +`2c5d814cfa80` in all four runs; drafted output equaled undrafted output, and two of six prompts failed the strict +golden check in every configuration. + +| Decode mode | Cold map | Fresh map | Throughput change | +|---|---:|---:|---:| +| Drafts off | 35.19 tok/s | 38.04 tok/s | +8.1% | +| Drafts on | 47.07 tok/s | 48.92 tok/s | +3.9% | + +Separate opt-in telemetry runs with the same prompt fixture measured combined decode/verify device-pick misses at +15.43% → 14.22% on rank 0 and 11.89% → 10.44% on rank 1. These telemetry runs include boundary synchronization +and are used only for miss counts, not throughput. Decode graphs replayed in the throughput runs; this short fixture +did not establish long-context KV-page misses, prefill-graph benefit, or two simultaneous live requests. + +The next miss-reduction candidate is a fixed-total per-layer line budget informed by these measurements, not a larger +rank-wide allocation. The device planner and prefill-ring layout require variable-sized planner-state offsets and +graph-safe launch/state handling. A local ragged-pool prototype was screened below, but is not committed or promoted. +The exact startup budget must preserve BF16 KV, graph scratch, and the minimum pool floor. + +A profile from an older workload was also tested and regressed this fixture: misses rose to 19.50% on rank 0 and +17.13% on rank 1; its uninstrumented ABBA averaged 1.7% lower undrafted and 4.0% lower drafted throughput than cold +placement. Profile reuse is therefore workload-sensitive. This experiment validates the existing opt-in profile +mechanism for a matched prompt mix; it does not justify enabling historical profiles by default or claim a general +serving speedup. Next work: make profile freshness/recency safe, then repeat with a longer representative workload +and simultaneous requests before promoting a placement policy. + +## 2026-10-24 profile-sized line-budget screen + +The ragged-pool prototype fixes the device picks-buffer capacity from the minimum line floor, so profile-weighted line +counts preserve the modeled variable CUDA allocation exactly. On the two TP2 ranks, startup logged 4,806,791,424 and +5,780,005,248 bytes allocated, equal to each rank's uniform baseline; pool sizes ranged 30–73 lines (rank 0, uniform +54) and 44–84 (rank 1, uniform 65). This is an allocator-budget equality, not a full independent VRAM audit. + +An uninstrumented graph-on cold/profile/profile/cold ABBA used two concurrent mixed requests, 256 generated tokens, +32,768 configured context, prompt cache off, and shared BF16 KV. Aggregate results were 34.5/34.8/34.7/34.5 tok/s; +each leg had zero failures and exact alone/serial equality, with the same two output hashes. The profile arms averaged +34.75 tok/s versus 34.50 for the controls (+0.7%), which is not a reliable serving gain. + +A separate matched telemetry profile/control pair used the same workload and output hashes. Verify device-pick misses +changed from 110,253 to 101,117 on rank 0 (−8.3%) and from 88,309 to 91,727 on rank 1 (+3.9%): combined verify misses +fell 2.9%. Combined decode misses increased 3.5%. Telemetry synchronizes at boundaries and is only miss evidence; +these counters are rank-asymmetric and do not establish a robust total-miss or throughput improvement. Keep the +allocator experimental. Next, score candidate budgets by replaying ordered demand through the production admission and +eviction policy, then repeat with held-out concurrent prompts before considering promotion. + +The held-out check used two different concurrent code/chat prompts with the profile trained on the original pair. +Under telemetry, profile versus uniform measured verify misses of 121,127 versus 113,923 on rank 0 (+6.3%) and +117,102 versus 98,834 on rank 1 (+18.5%); combined verify misses rose 11.9%, and combined decode misses rose 26.8%. +The uninstrumented aggregate was 33.1 versus 33.9 tok/s (−2.4%). All legs had zero failures and identical per-request +hashes. This falsifies generalization for this profile-sizing heuristic; do not promote it. Review also found two +contradictory/stale expectations in the new allocator tests and that profile loading would activate sizing without an +explicit opt-in. The prototype remains an experiment; the next candidate is ordered-demand replay, not another +histogram tweak. + +The separately committed recent-demand eviction policy (`7d0970d35`, default-off) was then tested on the same held-out +pair in a graph-on TP2/c=2 ABBA. Uniform legs averaged 34.45 aggregate tok/s; recent-demand legs averaged 33.50 tok/s +(−2.8%). A matched telemetry pair measured 270,350 combined decode/verify misses with recent-demand versus 259,479 +with uniform placement (+4.2%); verify misses rose 3.6% and decode misses 6.8%. All four throughput legs and both +telemetry legs had zero failures and identical per-request hashes. Keep recent-demand default-off; it helped the original +pair but did not generalize to this held-out mix. + +Review also identified a startup-ledger undercount after adding the recent-demand state: the old estimate omitted 4,610 +u32 words per planner pool (18,440 bytes), or about 0.84–0.86 MiB/rank for the current pool counts. The local startup +ledger now derives state bytes from the planner layout. This is a correctness/accounting fix, not a cache-speed claim. + +### Next expert-miss experiment + +On 2026-10-24, an opt-in DPLAN v3 diagnostic captured 1,024 decode rounds per pool in a 1 MiB trace and exactly replayed them. The TP2 graph-on run used BF16 paged KV, a 10-line expert pool, and one golden prompt on the direct Engine path (not the shared-arena server path). All 96 decoder-pool traces (48 per rank) had zero drops and 1,024/1,024 verified transitions; both ranks produced output hash `da8b74115bc2f3bc`. Decode graphs replayed 1,018/1,024 rounds per rank. TP2's rank masks split 491,520 routed picks across both ranks (236,927 on rank 0; 254,593 on rank 1). The production lifetime policy recorded 122,984 misses/copies on rank 0 and 137,392 on rank 1: 260,376/491,520 combined picks (53.0%), with 231,144 hits. The offline evaluator's lifetime copy count exactly matched production. Combined LRU replayed 293,031 copies (+12.5%); decayed frequency replayed 258,797 (−0.61%, 1,579 fewer). The previous interim note mistakenly doubled the denominator; the rank-specific counts above are verified from all 96 trace files. + +A second exact TP2 graph-on trace used 56 expert lines/layer, matching the recent serving configuration, with the same 1,024-token single-prompt fixture. All 96 traces replayed 1,024/1,024 transitions with matching rank hashes. Production misses/copies were 17,629 on rank 0 and 21,477 on rank 1 (39,106/491,520 combined picks, 7.95%). Offline decayed-frequency replay reduced copies to 37,753 (−3.46%); guarded admission reduced copies to 33,720 (−13.8%) while mapping 4,023 cold picks rather than copying them. These are trace counterfactuals, not measured throughput wins; mapped-read cost is not represented. + +The v4 trace format adds strict checkpoints for device-planner ownership transfers. A sequential two-prompt run (512 decode rounds each, 10 lines/layer) passed exact replay on all 96 pools with two checkpoints and 1,024/1,024 transitions; the per-segment TP2 output hashes matched across ranks, with combined hash `1f63a7addea9e585`. Production recorded 81,671 misses on rank 0 and 89,854 on rank 1 (171,525/491,520 combined picks, 34.9%). This demonstrates reuse across sequential prompt segments only—not concurrent serving, throughput, or 200k context. No policy is promoted until checkpoint-aware counterfactuals and uninstrumented serving tests pass. + +Checkpoint-aware offline replay now cross-checks each trace's lifetime-policy rounds, hit/miss picks, copies, evictions, mapped experts, and mapped picks against the exact device-planner replay; the 96-pool run passed this guard. Across its recorded device segments, LRU caused 182,921 copies (+6.64%); decayed-frequency caused 165,699 (−3.40%, rank 0 −3.55%, rank 1 −3.25%). Guarded admission cut simulated copies to 17,478 (−89.8%) only by mapping 167,186 cold expert picks. These are segment-conditioned counterfactuals: at each checkpoint candidates restart from the production state, and host-prefill intervals/costs are excluded. The mapped-read byte traffic and latency are not modeled, so guarded admission is not a speedup claim. Next, evaluate a default-off decayed-frequency runtime candidate and quantify actual mapped traffic on held-out concurrent prompts; retain the current policy unless both-rank misses and uninstrumented TP2 graph-on serving improve. + +The held-out regressions reject both histogram-sized pools and recent-only eviction as general policies. Current DPLAN +telemetry records aggregate/per-layer hit, miss-pick, copy, eviction, and in-place totals, but those totals cannot +reconstruct planner order or reuse distance. The ring contains only a bounded recent window and is not a complete +trace. The planner also admits every miss when a line is available; it does not compare a newcomer against a resident +before evicting, so cache-pollution/bypass is not measured yet. + +The bounded, opt-in DPLAN trace now records ordered picks and exact planner transitions, including strict checkpoints +when host prefill transfers pool ownership back to the device planner. Replay verifies every device segment and the +recorded checkpoint states; it deliberately excludes the host-prefill interval. Candidate policy replay is therefore +segment-conditioned, not continuous across requests. Keep tracing out of throughput runs. Next evaluate a default-off +decayed-frequency runtime candidate, count actual mapped-read traffic, and test on held-out concurrent prompts. Reject +dropped or discontinuous traces; preserve TP2 budgets, the 10-line floor, graph capture/replay, snapshots, cache isolation, +and BF16 KV paging. Promotion requires miss improvement on both ranks plus an uninstrumented graph-on serving gain +outside baseline variation. Experimental policies remain opt-in/default-off until then. + +## What this adds + +The EXL3 pack now runs through upstream's own Flash Next CUDA family and server (`tensorfold-native serve`), +on one card and at tp=2. Integration touches existing upstream files and adds new files under +`zig/src/families/flashnext/cuda_exl3_*.zig`, `zig/kernels/cuda/flashnext_exl3_*.cu` and `zig/tests/cuda/`. + +| Area | What exists | Main files | +|---|---|---| +| Weights | EXL3 dense, shared, routed experts, 5-bit lm_head, PLE n-gram table read from the pack | `cuda_exl3_load.zig`, `cuda_exl3_store.zig`, `cuda_exl3_ple.zig` | +| Expert offload | Per-layer device line pool fed from a pinned host bank; LRU plus use count; misses copied on a side stream | `cuda_exl3_pool.zig`, `exl3_offload.zig`, `exl3_round.zig` | +| tp=2 | Expert parallel: rank 0 routed 0-255 (+MTP head), rank 1 routed 256-511 + shared expert; one fp32 all-reduce per layer; per-rank pools and half-banks; vocabulary-split lm_head | `cuda_exl3_tp.zig`, `flashnext_exl3_tp.cu` | +| Prefill | Expert-major deep chunks (thousands of rows), one ring of lines across all layers, 64-member tile kernel for the mul1 3-bit experts, tp=2 version with overlapped per-sub-block reduction, expert-major MTP head | `cuda_exl3_prefill*.zig`, `flashnext_exl3_prefill.cu` | +| MTP | Draft rows, drafts with the prompt cache at tp=2, rank 1 keeps no head cache in its windows | `cuda_exl3_mtp.zig`, `cuda_rank_engine.zig` hooks | +| Long context | Attention head split across ranks (12 query heads + 1 KV head each), indexer keys in mapped pinned host memory, 1024-row chunk, prompt states parked in host memory | `cuda_exl3_heads.zig`, `cuda_exl3_kv.zig`, `cuda_snap.zig` | +| KV paging | Full BF16 K/V of the 12 attention layers in pinned host memory with a device window and block cache; allocation/budget geometry now unified, default off | `cuda_exl3_kvpage.zig`, `cuda_exl3_kv.zig`, `flashnext_exl3_kvpage.cu` | +| CUDA graphs | Eligible decode rounds/windows capture the 48-layer pass, NCCL sums and head gathers; prefill captures dense layer sub-blocks; greedy MTP proposals have a conditional graph; graph instantiation runs on a worker | `cuda_exl3_dplan.zig`, `cuda_exl3_graph.zig`, `cuda_exl3_chain.zig`, `cuda_exl3_recycle.zig` | + +## Verification + +References: a float64 oracle (`exl3_*_cpu.zig`, kept as test code) and an independent Python run of the +upstream-quality reference (`tools/exl3_ref_golden.py`). + +- **Single card, served over HTTP:** drafted output equals undrafted on every prompt; lm_head matches float64 to + 2.95e-3 of peak; teacher-forced top-1 equals the reference at all 16 tested prompt positions. This is a narrow + sample, not an end-to-end strict golden pass. +- **tp=2, served over HTTP** (`--tp 2 --rank R`): drafted equals undrafted, bench hash `2c5d814cfa80` in every + configuration; layers 0, 3 and 17 are bit-identical to tp=1 when run all-resident. The serving script's hash + checks are truncated and it can exit successfully after a golden mismatch; interpret its status as evidence, + not a strict pass/fail gate. +- **Near-200k context, 2 concurrent requests, tp=2, bf16 KV:** 195,938- and 195,957-token needle requests sent together + were both answered correctly, with drafts on and off, drafted equal to undrafted. +- **KV paging:** at 200k on one card, last-row logits, 128 greedy tokens and every K/V row are bit-equal to keeping + K/V on the card; sequence VRAM fell from 5.484 GiB to 1.391 GiB. TP=2 paging has only been tried sequentially + at shorter contexts, without token hashes or paging counters. +- **Graphs:** logits of every graphed step are bit-equal to eager on both ranks; the device planner equals the host + planner on 1,600 random rounds. + +### Test measurements (ms per token, tp=2 serve gate, two 128-token runs each) + +| Configuration | drafts off | drafts on | +|---|---|---| +| Before the stall fixes, eager | 29.7 - 30.1 | 14.6 - 14.9 | +| Defaults (n-gram gather inline) | 24.7 - 26.2 | 11.9 - 12.6 | +| Blocking NCCL communicator | 21.2 - 23.3 | 11.5 - 12.2 | +| Graphs + blocking NCCL (final branch) | 21.1 - 23.4 | 11.5 - 13.7 | + +The three lower rows were reproduced on the final branch head (`44d65597`) by a separate gate run: bench hash +`2c5d814cfa80` and drafted equal to undrafted in every configuration; two of three golden prompts match. + +Prefill, real text, rows per second: tp=2 deep 1,555 at 8k, 1,486 at 32k, 1,421 at 65k, 1,225 for a real 200k +prompt (tp=1 deep: 1,531 / 1,461 / 1,395 / 818). The old batched pass ran at 100 - 210 rows/s. A prior note's +168 s “on one card” duration is near the tp=2 rate (about 163 s by the reported rate), not the tp=1 rate (about +244 s); its configuration is unverified, so it is excluded. Saved concurrent-run logs show 312.1 / 562.5 s without drafts +and 380.1 / 693.2 s with drafts; the earlier 304 / 554 s figures are not reconciled to those logs. + +## What graphs do and do not give + +For eligible shapes, decode graphs match eager bit for bit, but a repeatable performance win over the optimized +eager path has **not been demonstrated** on this hardware. The decode speedup came from removing host stalls: the +pool's per-layer host wait (replaced by the device planner), thread spawns in the n-gram gather, and the non-blocking +NCCL communicator's host cost. Once those are gone a step is bound by GPU kernels (profile: about 20 ms of kernels +per token; dense EXL3 projections 6 ms, bf16 hyper-connection lanes 4.5 ms, all-reduces 2 - 3.5 ms, experts 3 ms). +With drafts on, the first drafted request after start is slower in the recorded gate; graph instantiation now runs +on a worker, but graph execution still synchronizes for results and timings do not isolate first-use capture from +replay. A TP=2 projection-fork experiment preserved logits but did not produce a repeatable graph-specific gain. +Graphs stay opt-in (`TF_FLASHNEXT_EXL3_GRAPH=1`). Prefill graph captures now preserve sub-block reduction overlap +with external event-wait nodes and pass bit parity. At chunk sizes 1024, 2048 and 4096, graph throughput remained +within run variation of eager; the former 10% regression is removed, but a speedup is not proven. An 8192-row +chunk request OOMed and used fallback, so it is not counted. Prefill graphs remain opt-in +(`TF_FLASHNEXT_EXL3_GRAPH_PREFILL=1`). + +## Prefill graph replay follow-up (2026-10-20) + +A TP=2, two-lane HTTP run at 16,040 prompt tokens used BF16 paging, MTP, snapshots and the default +`TF_FLASHNEXT_EXL3_GRAPH_PREFILL_AFTER=2`. Both ranks reported 72 prefill graph replays, 216 eager layer executions, +96 captures and about 314 ms cumulative capture time. Warm/rewind reused 16,033 tokens and branch reused 16,023; +all response hashes matched the cold baselines. The earlier 12k run's zero replays are explained by the threshold: +the available same-shape chunks only reached capture, with no later same-shape chunk to replay. + +The DPLAN run and two graph runs measured cold prefill at 27.85/16.47 s, 27.92/16.39 s and 28.00/16.51 s across +the two concurrent lanes. This is within run/lane variation and is not a demonstrated graph speedup. A follow-up +keeps the complete exact replay key, but normalizes only the position and paging-window-origin hash for capture +sightings and driver-checked executable-update candidates; sequence/scratch/event ownership remains strict. The +updated gate retained exact output/cache hashes and reported 72 replays, 192 eager layer executions, 120 captures and +24 successful `Exec.update` calls per rank. Capture time was 325 ms; DPLAN/graph cold prefill measured 16.44/27.80 s +and 16.51/27.99 s. This exercises executable reuse across page-window changes, but shows no speedup; incompatible +topology still uses the instantiate fallback. Next profile the expert/pool host boundary and dense graph replay +versus eager time before making more topology changes. + +## Switches + +| Variable | Effect | Default | +|---|---|---| +| `TF_FLASHNEXT_EXL3_GRAPH=1` | Decode graphs (and prefill dense-half graphs where enabled) | off | +| `TF_FLASHNEXT_EXL3_GRAPH_PREFILL=1` | Prefill layer graphs at tp=2 | off | +| `TF_FLASHNEXT_NCCL_BLOCKING=1` | Blocking NCCL communicator (init has no deadline if the peer never joins) | off | +| `TF_FLASHNEXT_ATTN_SPLIT=0` | Replicated attention at tp=2 (the split is on by default) | split on | +| `TF_FLASHNEXT_IKC_HOST=1` | Indexer keys in mapped pinned host memory (needed for 200k x 2) | off | +| `TF_FLASHNEXT_KVPAGE=1` | K/V paging to host memory; takes precedence over the head split | off | +| `TF_FLASHNEXT_EXL3_PREFILL=0` | Turn the deep prefill off | on | +| `TF_FLASHNEXT_SHARD_HEAD=0` | Whole lm_head on both ranks | split on | +| `TF_FLASHNEXT_DRAFT_HEAD=affine4` | Cheaper draft head; changes which drafts are proposed (88/111 accepted instead of 90/113), output stays bit-equal | off | + +The 200k x 2 tp=2 run used resident BF16 KV (not dynamic KV paging) and was run with +`TF_FLASHNEXT_EXL3_STREAMS=2 TF_FLASHNEXT_IKC_HOST=1 +TF_FLASHNEXT_EXL3_PREFILL_ROWS=2048 TENSORFOLD_MEMORY_RESERVE_GIB=0.75 TF_FLASHNEXT_NCCL_BLOCKING=1 --parallel 2`; +the pool had 11 lines per layer on rank 0 and 29 on rank 1 (floor 10). + +## Known gaps and risks + +1. **Golden gate fails on one prompt.** The 16-token prompt leaves the Python reference at step 3 (ours 9564, + reference 6124, reference margin 0.125, one bf16 step) on one card and at tp=2. Two of three golden prompts + match exactly. Treated as a bf16 near-tie; not proven. A teacher-forced margin study would settle it. +2. **The head split is not bit-identical to replicated attention** (different fp32 order in the o_proj sum). + Greedy continuations after 8k and 32k real tokens differ after 3 - 4 tokens; not shown to be near-ties because + the server returns no logprobs. +3. **tp=2 parity is not established.** A saved 39-position teacher-forced run matched tp=1 at 37/39 positions + and failed at positions 18 and 33; logits also differ due to fp32 reduction order. At 65k and 200k the + deep-prefill last row swaps a near-tie pair between tp=1 and tp=2. The single-layer 16-position top-1 check + does not cover this case. +4. **200k x 2 remains deferred.** Rank-0 pool headroom is tight and simultaneous prefill can hurt decode. The prior + near-200k resident-BF16 retrieval gate is not a claim for paged 200k capacity. Ordinary-length prompt-cache reuse + now works with TP=2, BF16 paging, MTP and graph mode: two concurrent chat histories passed warm, branch and rewind + phases with 4,122/4,139 prompt tokens reused and exact response-hash parity against cold-cache baselines. +5. **BF16 KV paging and snapshots now pass ordinary-length TP=2 correctness gates.** A two-GPU graph test restored + two paged lanes and preserved exact state/logits; each rank replayed prefill graphs 72 times and decode graphs 5 + times. The HTTP cache gate proves actual snapshot reuse, MTP continuation and two concurrent requests. Still open: + repeatable graph speedup/profiling, live-rank overflow and failure injection, and long-context paged capacity. + 200k × 2 is deliberately not the next gate. The user’s priority is to make graph execution and cache reuse robust + and fast first. Int8 KV is out of scope; BF16 remains the primary and required format. +6. **200k reference parity is not established:** a saved 200k prefill comparison exited 1 with max logit + difference 1.1406 and top-1 6009 versus 19251 (reference margin 0.25). The successful 200k × 2 needle answers + establish retrieval/serving behavior for those prompts, not strict full-logit parity. +7. **Graph path:** graphs execute and replay for both prefill and decode, including TP2 BF16 paged snapshot restore. + However, neither decode nor prefill graphs have a repeatable measured speedup yet. MTP verification also has a redundant GPU best-selection/readback + boundary after full-logit readback. An opt-in verify-tail experiment preserved output/hash but showed no + observable speed change in one TP=2 A/B pair; no graph speedup is claimed. +8. **Prefill at tp=2 is only about 2% faster than tp=1 at equal chunk sizes:** the replicated dense work + (projections, GDN, attention, hyper-connections) is the bottleneck, not the experts. +9. **Blocking NCCL is opt-in** until its init has a timeout. +10. Existing automated gates need stricter pass/fail behavior before serving as acceptance gates: the HTTP gate can + print golden mismatches without failing, and compares only truncated drafted/undrafted hash prefixes; the + prefill comparison allows a close-margin top-1 and can skip missing reference dumps. The passing summaries + above must be read with those limits. +11. **MTP draft-token parity is not directly asserted for graph replay.** The graph test checks main-model logits, + not the draft proposals or replay count. Final drafted/undrafted output equality does not prove the proposed + tokens matched. The draft-row test covers limited inputs and does not assert argmax equality. +12. `zig build test-cuda-host` fails on clean upstream v1.1.0 on macOS as well; `test-flashnext-exl3-host` + passes. Upstream PR #1046, which touches our hook points, could not be read. + +## Next steps + +1. Meet the graph acceptance gate: reduce decode host boundaries (first target is MTP verify's redundant best- + selection/readback), make decode graphs measurably faster than optimized eager decode, and extend prefill capture + beyond dense sub-blocks without losing reduction overlap. Recheck with repeated TP=2 gates before changing defaults. +2. Exercise fail-closed paging overflow and collective cleanup on live TP2 ranks; snapshots/MTP rollback and + ordinary-context cache reuse are implemented and pass their current correctness gates. Defer the 200k × 2 + acceptance run until graph execution and cache reuse are optimized. +3. Add an init timeout to the blocking NCCL communicator before making it the tp=2 default. +4. Make the HTTP and prefill gates fail on strict mismatches, compare full hashes, and require reference artifacts + for long-context parity runs; investigate the 200k reference-logit failure. +5. Assert actual graph replay and compare full MTP draft-token sequences against eager; strengthen the draft-row + test to assert the decision-critical top-1 output. +6. Margin study for the golden near-tie and for the head split (needs logprobs from the server). +7. Spread the MTP head across both ranks (rank 1 idles about 28% with drafts on); split the replicated dense + work to speed up prefill at tp=2. +8. The single-core materializing F64/AVX2 prototype is rejected (2.35 s/expert). A fused packed CPU expert path + remains unbuilt; only reopen it after the CUDA kernel experiments and only if a bounded CUDA-semantics-correct + cold-decode test beats the GPU miss route end-to-end without harming BF16 paging. Never omit routed experts. +9. Split this branch into reviewable PRs after the TP=2 graph gates pass. + +## Branch map + +`exl3-c-main` merges: core adapter, pool, PLE, draft rows, expert-parallel TP=2, serving, split `lm_head`, deep +prefill, BF16 KV paging, snapshots/cache reuse, and CUDA graphs. Current design and validation notes: +`docs/plans/exl3-kv-paging.md` and `docs/plans/exl3-graphs.md`. Superseded v1.1.0 and int8-first planning memos +were removed from the draft because they describe obsolete decisions and conflict with the BF16 paging scope. + +Not done on purpose: no change to `docs/recipes`, no family or model registration (held until the strict golden +gate passes). + +## 2026-10-24 recent-demand device eviction + +The graph/DPLAN path now has an opt-in 256-row, per-layer expert-demand ring (`TF_FLASHNEXT_EXL3_RECENT_EVICTION=1`). +Lifetime counts still seed initial placement and remain available to profile persistence; the recent score only affects +decode miss ordering and device-pool victim selection. The policy does not skip routed experts, alter logits, or enter +request snapshots. Host prefill allocation is unchanged. It is off by default. + +Validation on par1-cs13 used the `Qwen3.8-Flash-Next-exl3-3.05bpw` pack, both RTX 5060 Ti cards (sm_120), TP=2, +32,768 configured context, CUDA graphs on, `--parallel 2`, shared BF16 KV paging on, and prompt cache off. Two concurrent +HTTP prompts generated 512 output IDs each with drafting disabled. With synchronized per-layer telemetry, baseline decode +misses were 31,844 on rank 0 and 25,446 on rank 1 (57,290 total); recent-demand measured 29,012 and 22,060 (51,072, +−10.9%). Both per-request output hashes matched between policies. This count comparison is diagnostic and is not used +for throughput. + +Two uninstrumented base/recent pairs in opposite order used the same two concurrent prompts and output length. Mean +per-request decode was 27.97 tok/s baseline versus 28.53 tok/s recent-demand (+2.0%); all four output hashes matched. +This is a small result from one prompt pair, not a broad speedup claim. A separate graph-on, shared-BF16-paged, concurrent +512-output MTP-on run also matched both no-draft output hashes; telemetry misses were 729→726, so the larger no-draft +miss reduction does not apply to the MTP path in this fixture. + +Validation also passed `zig build test-flashnext-exl3-host`, `zig fmt --check` on the changed Zig files, `git diff --check`, +and the native sm_120 build. The option remains experimental and default-off pending held-out prompts, longer sessions, +and broader MTP/profile interactions. + +## 2026-10-24 prefill-graph churn attribution and two-chunk window ABBA + +Setup for every leg below: par1-cs13, two RTX 5060 Ti, TP=2, `--parallel 2`, shared BF16 KV paging, decode and prefill +graphs on, MTP drafts on, prompt cache off, two concurrent 16,384-token prompts, 128 generated tokens, two warm-up and +two measured batches per leg. All 16 measured requests across the legs below produced the same token hash +`44f75f563088…`. + +**Prefill graph churn is not the bottleneck.** With 38 expert lines and prefill graphs on, the server counters read +`shared lanes reassigned 0`, `admission misses 0`, and prefill graphs replay/eager/instantiated/updated 1200/96/240/0 with +`capture 440 ms` (rank 0) and `401 ms` (rank 1) for the whole run, against about 14 s of prefill per batch. The 96 eager +runs and 240 captures are a one-time fill (240 sighted, 96 run eagerly before instantiation); they did not recur in the +measured batches. The lane-reassignment hypothesis is refuted; even a perfect graph policy bounds the saving at well under +1 s. + +**Prefill is bound by the expert stream over PCIe.** Opt-in cache telemetry (`TF_FLASHNEXT_CACHE_TELEMETRY=1`) showed rank 0 +moving 427 GB through the expert ring for four batches (about 107 GB per batch, about 13 GB per 4096-row chunk), with +13% of ring copies unused. At 14.4 s per batch that is about 7.4 GB/s per rank over a PCIe 4.0 x8 link. + +**The two-chunk window (`TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK=1`) shortens prefill when it is admitted.** At 38 lines per +layer it is declined (`extra_device=588070912`, free 6.36 GB against 6.60 GB protected). At 26 lines per layer it is +admitted. ABBA on one binary (two-chunk, one-chunk, one-chunk, two-chunk), 26 lines per layer: + +| leg | mean prefill, measured batches | aggregate decode tok/s | +|---|---|---| +| two-chunk (A) | 12.503, 12.511 | 15.75, 15.74 | +| one-chunk (B) | 14.428, 14.422 | 14.92, 14.94 | +| one-chunk (B) | 15.732, 14.420 | 14.92, 14.97 | +| two-chunk (A) | 12.537, 12.543 | 15.71, 15.76 | + +Two-chunk prefill is about 13% shorter (12.52 s against 14.42 s, ignoring the one 15.73 s outlier in a one-chunk leg), with the +same direction in both orders, about 15% more prefill throughput. Reference points at 38 lines with one chunk (separate +runs, not part of the ABBA): prefill 14.41 s, decode 15.87 tok/s. Two-chunk at 26 lines therefore costs about 0.8% decode +against the 38-line configuration. The higher decode throughput of the two-chunk legs over the one-chunk legs at 26 lines +(about +5.5%) is not a decode speedup: see the correction below. + +This is a measured, correctness-preserving prefill gain for this workload, not the 2× target. It depends on trading 12 of 38 +expert lines per layer for the window's extra device memory, and it was measured on one prompt shape and one machine. The +option stays default-off until held-out prompt shapes, longer prompts, and the MTP and snapshot paths are re-checked. + +### Prompt-length and MTP sweep of the two-chunk window (same day) + +Same binary, machine, 26 lines per layer and flags as the ABBA above; one server start per leg and order two-chunk, +one-chunk, one-chunk, two-chunk (A B B A). Per leg one warm-up batch (discarded) and two measured batches of two +concurrent prompts, 128 generated tokens. Every measured request returned the same token hash across all four legs in each +row (8,192 and 30,000 tokens without drafts, 16,384 tokens with request-level MTP drafts). + +| prompt tokens | drafts | mean prefill s, two-chunk (4 batches) | mean prefill s, one-chunk (4 batches) | change | +|---|---|---|---|---| +| 8,192 | off | 4.925 | 6.47 (one 7.76 outlier excluded) | −24% | +| 30,000 | off | 25.92 | 29.50 | −12% | +| 16,384 | on | 15.37 | 15.23 (one 16.58 outlier excluded) | +1% (no gain) | + +The window shortens prefill for every no-draft shape tried, with the same direction in both orders. With MTP drafts on it +did not at first: drafted prompts also ran as 8,192-row paired passes, but `head()` in `cuda_exl3_prefill.zig` declined any +call with more rows than the deep chunk (`dp.rows < n`), so the MTP head silently took the slow per-chunk path for every +paired pass. Aggregate decode for drafted 16,384-token requests was 11.5 tok/s in both modes, below the undrafted 15.7 tok/s, +which matches the earlier finding that MTP-drafted decode regresses on this setup. + +**Head fix.** `head()` now splits a paired window's rows into two deep-sized slices (`headSplit`: the first aligned up to +`att_rows` and at most one deep chunk, the second no larger and not under `min_rows`, so the second slice cannot decline after +the first has run). The log line `the MTP head runs expert-major over deep chunks (4096 rows first)` confirms the deep head path +engages. Re-measured with the same ABBA procedure (two-chunk, one-chunk, one-chunk, two-chunk), 16,384 prompt tokens, +drafts on, 26 lines per layer, all 16 measured requests with token hash `44f75f56…`: + +| leg | mean prefill s, two measured batches | aggregate decode tok/s | +|---|---|---| +| two-chunk (A1) | 13.177, 13.231 | 12.80, 12.78 | +| one-chunk (B1) | 15.214, 15.215 | 11.65, 11.66 | +| one-chunk (B2) | 15.211, 15.232 | 11.67, 11.67 | +| two-chunk (A2) | 13.179, 13.211 | 12.77, 12.81 | + +Drafted prefill is 13.2% shorter (13.20 s against 15.22 s) and drafted aggregate decode is 9.6% higher (12.79 against 11.66 +tok/s); the decode difference is again not isolated to a cause. One-chunk drafted prefill is unchanged by the fix (15.2 s). +Host validation: `zig fmt --check` passes and `zig build test-flashnext-exl3-host` passes on par1-cs25 including a new +`headSplit` test. These sweeps do not cover the snapshot path or prompts above 30,000 tokens; the window stays default-off. + +**Snapshot gate.** `run-snapshot-gate.sh` (`tf-cuda-test flashnext-exl3-graph … snapshot-shared`, TP=2, shared BF16 paging, graphs +on, `TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK=1`) was rebuilt from the tree with the head change and passed on both ranks (exit +codes 0/0): shared-slot snapshot remapping, recycled generation, foreign windows/history/continuation, exact eager/graph +state and logits, prefill replays 112 and decode replays 3. This gate's prompts are short, so it shows the window-on build +keeps snapshot behavior intact; it does not by itself prove the snapshot path on 8,192-row paired passes. + +### Admission thresholds and a 60,000-token prompt + +The window needs `extra_device=588070912` bytes and is admitted only when `free >= protected + extra`. Free memory stayed at about +6.35–6.37 GB in every probe; what moved was `protected`, which falls about 69 MB per expert line per layer. Probes +(TP=2, 32,768-token context unless noted, window on): 38 lines declined (protected 6.575 GB), 38 lines with a 17,408-token +context declined (6.575 GB; a smaller window does not help), 29 lines declined (5.953 GB, still about 0.18 GB short), 26 lines +admitted. At a 65,536-token context 26 lines is declined (protected 6.089 GB) and 21 lines is admitted. The practical rule on +this card is therefore a loss of 12 lines per layer at a 32k context and 17 at a 64k context. + +A 60,000-token prompt at a 65,536-token context and 21 lines, ABBA on one binary (two concurrent prompts, 128 generated tokens, +drafts off, one discarded warm-up batch and two measured batches per leg), all measured requests with token hash +`566ecfab…`: + +| leg | mean prefill s, two measured batches | aggregate decode tok/s | +|---|---|---| +| two-chunk (A1) | 60.44, 60.59 | 5.30, 5.29 | +| one-chunk (B1) | 65.37, 65.45 | 5.30, 5.30 | +| one-chunk (B2) | 65.37, 64.05 | 5.30, 5.30 | +| two-chunk (A2) | 60.62, 60.67 | 5.29, 5.29 | + +Prefill is 7.3% shorter with the window (60.58 s against 65.0 s; 65.0 is the mean of all four one-chunk batches). The gain +shrinks with prompt length (24% at 8k, 13% at 16k, 12% at 30k, 7% at 60k), consistent with the fixed dense and attention work +growing relative to the expert stream. At this length the prefill graph counters show 1330 replays, 1438 eager runs and 1141 +admission misses against a 112-entry graph cache, so very long prompts mostly run eagerly. Decode is unchanged by the window. + +### Correction: the window gives no decode gain + +The decode figures above that favored the window (about +5% to +10%) are most likely a prefill/decode interleaving effect, not a +faster decode. With two concurrent prompts, the first request decodes while the second is still prefilling, and the +reported per-request and aggregate decode rates include that overlap. A shorter prefill ends the overlap sooner, so both +requests report a higher rate (7.86/20.5 tok/s with the window against 7.47/19.5 without, TTFT 9.9/20.0 s against 10.4/21.0 s). +Supporting evidence: decode-phase expert cache counters were identical between window-on and window-off server logs +(25,152 plans, 176,210 hit picks, 68,222 miss picks, 66,693 device copies, a 72.1% device hit rate at 26 lines), and decode +graph coverage was the same (468 round-2 replays, 8 eager). This was not confirmed with a decode-only run. Treat the window as +a prefill optimization only; no decode improvement from it is claimed. + +### How much of decode is exposed expert-copy time (estimate) + +Two decode-phase telemetry runs on the same binary and workload (TP=2, graphs, shared BF16 paging, two concurrent 16,384-token +prompts, 128 generated tokens, window off) differ only in expert lines per layer: + +| lines per layer | device copies per decode step (rank 0) | device miss rate | sum of per-request decode tok/s | +|---|---|---|---| +| 38 | 95.3 | 20.8% | 29.7 | +| 26 | 127.3 | 27.9% | 26.95 | + +Treating step time as proportional to 1/tok/s, 32 more copies per step cost about 3.4 ms, or about 0.106 ms per copy. A copy is +one 1.76 MiB line, so that is roughly 17 GB/s: close to the PCIe 4.0 x8 link rate, which means these copies are almost entirely +exposed (not hidden behind compute). At 38 lines that puts about 10 ms of a roughly 34 ms step (about 30%) on expert misses. + +Consequences (estimates from two operating points, not a controlled sweep): removing every miss would cut step time by at most +about 30% (about 1.4× decode), so cache policy alone cannot reach 2×, and a 6% reduction in copies, the size of the best offline +replay result, is worth about 2% of step time, which is inside run variation. Large decode gains need either hiding the copies +(a router-independent prefetch of the next layer's likely experts) or shrinking the 23–24 ms non-miss part of the step. Neither +has been attempted. + +### Next-layer expert prefetch: predictability from a DPLAN trace (negative result) + +To size a prefetch that hides decode expert copies, a default-off dump was added to `cuda_exl3_dplan.zig`: with +`TF_FLASHNEXT_EXL3_DPLAN_TRACE=1` and `TF_FLASHNEXT_EXL3_DPLAN_TRACE_DIR=`, each pool's raw v4 trace buffer is written to +`dplan--pool.bin` when the planner is destroyed (only a gracefully stopped rank dumps; a rank killed with SIGKILL does not). +`tools/exl3_dplan_predict.py ` parses the files and replays them. + +Capture: rank 0 of a TP=2 graph-on server (26 lines per layer, shared BF16 paging, window off), two concurrent 4,096-token +prompts, 700 generated tokens each, drafts off. 48 decoder pools retained 1,074 aligned rounds (rows per round equal across all +layers in every round); the trace has a discontinuity from request ownership transfers, and 1,026 later rounds were dropped by the +bounded buffer. + +Method: train a co-occurrence table on the first half of the rounds (picks of layer L, same round, to picks of layer L+1), then on +the second half predict the top-N experts of layer L+1 from layer L's picks, excluding experts already resident in a simulated +26-line LRU cache, and compare against that layer's real misses: + +| prefetched experts per layer | share of misses covered | share of issued copies that were needed | +|---|---|---| +| 2 | 30.9% | 13.0% | +| 4 | 46.8% | 9.8% | +| 8 | 63.6% | 6.7% | +| 16 | 78.7% | 4.1% | +| 32 | 88.9% | 2.3% | + +Even the best precision wastes 87% of the copies, and at the measured 0.106 ms per copy (near PCIe line rate) the wasted +traffic costs more than the saved misses. For example top-2 issues 2 copies to save about 0.26 misses per layer-round: about +0.21 ms of link time against about 0.03 ms saved. This predictor does not support a prefetch. The next-layer prefetch is +rejected as a decode optimization on this evidence. + +Limits: the prompt is a repeated golden-fixture ID list, so generation is repetitive and the simulated LRU missed only 10.9% +of needed experts against 20.8–27.9% in the serving telemetry; a diverse workload would lower predictability further, not raise +it. One rank, one capture, a plain LRU model rather than the real planner, and a single co-occurrence predictor were tried. + +### Two-chunk window on real, different prompts per request + +All earlier window measurements sent the same repeated-ID prompt to both requests. This ABBA (two-chunk, one-chunk, one-chunk, +two-chunk; same binary, 26 lines per layer, 32,768 context, TP=2, graphs, shared BF16 paging, drafts off, 128 generated tokens) +sent two different text prompts built from the checkout's own files: request 1 is Zig source (16,008 tokens), request 2 is +markdown docs and test sources (12,758 tokens). One discarded warm-up batch and two measured batches per leg. Token hashes +were identical across all four legs for both requests (`1c2777f5…`, `a65bd936…`). + +| leg | request 1 prefill s | request 2 prefill s | batch wall s | +|---|---|---|---| +| two-chunk (A1) | 17.95, 17.95 | 8.14, 8.14 | 26.52, 26.51 | +| one-chunk (B1) | 21.19, 21.25 | 10.66, 10.70 | 30.35, 30.35 | +| one-chunk (B2) | 21.20, 21.27 | 10.68, 10.71 | 30.33, 30.33 | +| two-chunk (A2) | 13.16*, 17.99 | 8.16*, 8.17 | 26.62, 26.62 | + +(*one measured A2 batch reported a lower prefill for each request, but the same batch wall time; the lower figures come from the +two requests' prefills overlapping differently in that batch and are not used.) Two-chunk shortens the whole batch by 12.5% +(26.5 s against 30.3 s), request 1 prefill by 15% and request 2 prefill by 24%, with identical outputs. The gain therefore holds +for different, non-repeating prompt content and unequal prompt lengths, not only for the repeated-ID fixture. + +### Trimming the graph reservation to fit the window at more lines (dead end) + +The window's admission gap at 38 lines is about 0.8 GB. The graph budget is a heuristic (reservation 660 MiB plus headroom 128 MiB; +`exec/template/transient` 16/16/32 MiB, "not measured exec bytes"), so two env overrides were tried on the 38/32-line server: +`TF_FLASHNEXT_EXL3_GRAPH_RESERVATION_MIB=128` with `..._HEADROOM_MIB=32` fails startup with `InvalidGraphPolicy`, and +`..._RESERVATION_MIB=400` with `..._HEADROOM_MIB=64` (32 lines) starts but the window is declined with +`incremental graph reservation unavailable (required=0 supplied=419430400)`: an explicit reservation override disables the +window's own graph reservation. Closing the gap would need code that applies `withWindowReservation` to manual values and a measured +(not heuristic) graph footprint; neither was attempted. The window therefore still fits only after trimming 12 lines per layer at a +32k context (17 at 64k). Making it the default is a decode-versus-prefill trade that depends on output length: at 26 lines decode is +about 6% slower (14.9 against 15.87 tok/s aggregate) while a 16k prefill is 1.9 s shorter, so the window pays back only when generated +tokens are fewer than about 300 per 16k-token prompt. It is left opt-in. + +### Decode profile under nsys and a copy-kernel sweep (no change found) + +nsys 2026.3 on each rank in turn (`--trace=cuda --cuda-graph-trace=node`), TP=2, 38 lines per layer, graphs on, shared BF16 paging, +window off, two concurrent 2,048-token prompts, 256 generated tokens, three batches. Rank 0 and rank 1 kernel totals: + +| | rank 0 | rank 1 | +|---|---|---| +| total kernel time | 31.79 s | 31.12 s | +| NCCL all-reduce | 5.68 s | 5.05 s | +| everything else | 26.11 s | 26.07 s | + +Rank 0's biggest kernels: all-reduce 17.9%, `fn_exl3_dcopy` (expert miss copy, on a side stream) 12.6%, `tf_exl3_linear_10_2_8` 9.9%, +`tf_exl3_linear_10_2_4` 7.8%, `tf_exl3_grouped_2` 7.6%, `tf_qmmf` kernels about 10.5% together. The GPU is busy about 82% of the +decode window. The all-reduce median is 28 µs but its mean is 142 µs, so most of its time is waiting for the other rank. The +non-all-reduce work is equal on both ranks and both wait about the same, so the skew is symmetric (per-layer variation in misses on +whichever rank is unlucky), not one structurally heavier rank. A per-instance arrival comparison was attempted and discarded: the two +ranks' all-reduce counts differ (37,446 against 36,192, prefill all-reduces included), so the pairing is invalid. + +Sweep of `TF_FLASHNEXT_EXL3_DCOPY_BLOCKS` (default 144 blocks of 256 threads on a 36-SM card) at 12, 24, 48 and 144, ABBA-style +repeats, same workload: summed per-request decode 58.3/58.9 tok/s for every value, identical to the digit, with identical hashes. The +copy kernel's block count does not matter; the copies are link-bound. No default change. + +Consequence: about 16% of kernel time is all-reduce wait caused by miss variance, and eliminating that variance would be worth at most a +few percent per step. The decode step has no single large inefficiency to remove with a tuning knob. + +### Hit/miss overlap on by default (+4% decode without drafts) + +`TF_FLASHNEXT_EXL3_OVERLAP` (the held experts compute while the misses' copies run on a side stream, then the misses' experts +follow) was compiled but default-off and unmeasured. It is active in graph mode (same graph coverage). With the stock 144-block copy +kernel it gave no gain (summed per-request decode 58.1/58.6 against 58.3-58.5/59.0-59.1 tok/s at 2,048-token prompts), because the +copy kernel's blocks fill the 36 SMs and the held experts' kernels cannot run beside them. With few copy blocks the gain appears and +plateaus (2,048-token prompts, 256 generated tokens, 38 lines, summed per-request decode, two measured batches each): + +| overlap | copy blocks | decode tok/s | +|---|---|---| +| off | 144 | 58.3-58.5, 59.0 | +| on | 144 | 58.0-58.1, 58.6 | +| on | 32 | 59.2, 59.9 | +| on | 16 | 60.4, 61.0 | +| on | 2, 4, 6, 8 | 60.5-60.6, 61.0-61.1 | + +The copy is link-bound (a block-count sweep without overlap was identical to the digit), so 2 blocks already saturate it. Default +change: the split is on unless `TF_FLASHNEXT_EXL3_OVERLAP=0`, and with it on the copy kernel uses a quarter of the SMs, at least 8 +blocks (9 on the 36-SM RTX 5060 Ti); `TF_FLASHNEXT_EXL3_DCOPY_BLOCKS` still overrides. + +Validation on the exact new defaults, ABBA against `OVERLAP=0` with 144 blocks, 16,384-token prompts, 128 generated tokens, 38 lines, +window off, TP=2, graphs, shared BF16 paging; summed per-request decode tok/s (aggregate in parentheses), all hashes `44f75f…`: + +| drafts | new default (two legs, two batches each) | old behavior (two legs, two batches each) | change | +|---|---|---|---| +| off | 31.99, 32.34, 32.01, 32.16 (16.07-16.12) | 30.79, 30.90, 30.68, 30.93 (15.79-15.83) | +4.4% (+1.6% aggregate) | +| on | 19.87, 19.88, 19.79, 19.83 (12.37-12.40) | 20.12, 20.04, 20.00, 20.04 (12.44-12.48) | -1.0% | + +The MTP chain does not use the split and loses about 1% to the smaller copy-block count. The snapshot gate +(`flashnext-exl3-graph … snapshot-shared`) passes on both ranks with the new defaults, `zig fmt --check` and +`zig build test-flashnext-exl3-host` pass. + +Configuration note: on this setup drafted decode (about 20 tok/s summed, 71% draft acceptance, 62 rounds for 128 tokens) is far slower +than undrafted decode (about 32 tok/s), so serving with `--no-drafts` is the faster setting for these two cards today. + +### Drafts versus no drafts on real text (new defaults) + +Same server (38 lines, overlap default, window off, TP=2, graphs, shared BF16 paging), two different text prompts (16,008 and 12,758 +tokens, built as in the real-text window test), 192 generated tokens, one discarded warm-up batch and two measured batches per +mode. Summed per-request decode: drafts off 25.45 and 26.80 tok/s, drafts on 20.06 and 21.13 tok/s (about 21% lower with drafts). +Draft acceptance was 77% (161/208) on the first prompt and 50% (130/260) on the second. Token hashes were identical between the +two modes for both requests, so drafting is exact but slower here: each drafted round verifies more rows, needs more distinct +experts, and the extra expert misses cost more than the accepted extra tokens save. `--no-drafts` is the faster setting on this +setup. + +### Requests decode plain by default when experts are offloaded + +Drafted rounds lose to plain decode on offloaded experts (previous section). The lane engine now ignores a request's draft flag when +the engine has offloaded expert layers (`plainByDefault()` in `cuda_exl3_pass.zig`, used where `cuda_lanes.zig` decides `drafting`); +`TF_FLASHNEXT_EXL3_DRAFTS=1` restores drafting. Engines without offloaded experts are unchanged. The MTP head and its draft stream +memory are still loaded; only the per-request drafting is skipped (a later change could also skip loading them). + +Validation on the new binary, 16,384-token prompts, 128 generated tokens, 38 lines, overlap default, window off, TP=2, graphs, shared +BF16 paging, two measured batches per row, summed per-request decode tok/s, all hashes `44f75f…`: + +| server setting | request draft=false | request draft=true | +|---|---|---| +| new default | 32.01, 32.14 | 32.30, 32.32 (127 rounds, no drafts accepted) | +| `TF_FLASHNEXT_EXL3_DRAFTS=1` | 32.02, 32.22 | 19.83, 19.81 (62 rounds, 65/92 accepted) | + +A request that asks for drafts therefore decodes at the plain rate by default (about 63% faster than before on this prompt and about +21% faster on real text), with identical output. The snapshot gate passes on both ranks on a freshly built runner, and +`zig fmt --check`, `zig build test-flashnext-exl3-host` and `zig build compile-cuda-tests` pass. + +### Second, independent ABBA of the default-path changes (real text, drafts requested) + +Different content from the repeated-ID runs: two different text prompts (16,008 and 12,758 tokens), 192 generated tokens, 38 lines, +TP=2, graphs, shared BF16 paging, window off, one discarded warm-up batch and two measured batches per leg and request mode. New = +the branch defaults (overlap on with few copy blocks, plain decode on offloaded experts); old = `TF_FLASHNEXT_EXL3_OVERLAP=0`, +`TF_FLASHNEXT_EXL3_DCOPY_BLOCKS=144`, `TF_FLASHNEXT_EXL3_DRAFTS=1`. Order new, old, old, new. Token hashes were identical for both +requests in every leg and mode (`b143a9…`, `2b3c44…`). + +| requests ask for | new: summed decode tok/s, batch wall s | old: summed decode tok/s, batch wall s | +|---|---|---| +| drafts (a normal client) | 25.80, 26.96, 26.95, 27.02; 32.3, 32.2, 32.2, 32.2 | 20.50, 20.78, 20.52, 19.75; 39.2, 39.0, 39.2, 39.1 | +| no drafts | 26.65, 26.81, 25.47, 26.83; 32.4, 32.4, 32.4, 32.3 | 25.90, 26.07, 25.92, 26.07; 32.8, 32.8, 32.8, 32.8 | + +For a client that asks for drafts the batch finishes in 32.2 s instead of 39.1 s (17.6% shorter, 21% more throughput) and summed +decode is about 30% higher. For a client that already disabled drafts the overlap alone gives about 2-3% on summed decode and 1.3% +on batch wall time. Together with the repeated-ID ABBA (+4.4% undrafted, +61% for drafted requests) this is two independent +workloads in both orders. + +### Comparison with Strata on the same two RTX 5060 Ti (single request) + +Strata 0.2.4.2 (`7dd5cf1`, engine 0.1.48 for CUDA 13) was installed on cs13 with `./setup.sh --model IQ3_XXS --context 32768 --gpus all +--yes`: the Qwen3.8-Flash-Next IQ3_XXS GGUF (47.0 GB, about 3.06 bits per weight), automatic layer split across both cards, its MTP draft +layer with `--spec 4 --spec-min-p 0.5`, int8 KV, expert cache `auto` (about 15.7 GiB of VRAM used on each card), one request at a time +(its default). Neither Strata nor project-maya has an EXL3 path for this model, so the weights differ in format (a GGUF IQ3_XXS against our +EXL3 3.05 bpw) and quality was not compared. Ours was the branch defaults (overlap on, plain decode, window off, 38 lines per layer, TP=2, +graphs, shared BF16 paging, 32,768 context). Both were driven by `tools/exl3_single_request_bench.py` (chat endpoint, thinking off, +temperature 0, `ignore_eos`, 192 requested tokens, a unique prefix per run so no prefix cache applies), three requests per prompt. The +prompts are a 70-token story request and the two real-text prompts used earlier (Zig source and markdown docs). Strata stopped some code +answers early despite `ignore_eos` (41-96 tokens), so its decode figures for those rows cover fewer tokens. + +| prompt | metric | ours (EXL3, TP=2) | Strata (IQ3_XXS) | +|---|---|---|---| +| 70 tokens | time to first token | 0.73-0.76 s | 0.57-0.78 s | +| 70 tokens | decode | 33.7-36.5 tok/s | 75.1-87.1 tok/s | +| about 16,050 tokens (Zig) | time to first token | 10.38-10.50 s | 7.47-7.62 s | +| about 16,050 tokens (Zig) | decode | 26.6-29.1 tok/s | 83.1-103.6 tok/s | +| about 12,800 tokens (docs) | time to first token | 9.45-9.65 s | 6.74-6.85 s | +| about 12,800 tokens (docs) | decode | 28.7-29.3 tok/s | 101.1-102.1 tok/s | + +On this hardware Strata decodes 2.3 to 3.5 times faster and prefills long prompts about 1.4 times faster (about 2,100 against 1,550 tokens/s +at 16k); short-prompt first-token time is the same. Our two-chunk window (26 lines) closes part of the prefill gap (about 12.5 s per 16k +prompt in the two-request runs) but not the decode gap. From the design documents and the profile, the likely causes of the decode gap are +that Strata computes cold experts on the CPU instead of copying them over PCIe (our exposed miss copies are about 30% of a step), splits layers +across the cards (pipeline) instead of expert-parallel with an all-reduce every layer (about 16% of our kernel time is all-reduce wait), and +makes its MTP drafting pay off (ours loses to plain decode). These causes were not isolated by experiment. + +### Why Strata decodes faster: the expert cache hit rate (measured in its log; causes inferred) + +Strata prints per request: `expert cache 94.1% hit (+1.3% of the routed experts over PCIe)`, 95.2%, 97.5%, 97.2%, 97.9% and 98.9% (+0.1%) on the +same prompts used above. Ours at 38 lines per layer measured a 20.8% device miss rate (a 79% hit rate) in the decode telemetry, and 27.9% at 26 +lines (72% hit). So Strata moves roughly 5 to 20 times fewer experts over PCIe per token, and the exposed miss copies that take about 30% of our +decode step (previous sections) are close to absent for it. This alone accounts for a large part of the decode gap, before any CPU lane or MTP effect. + +What is visible in Strata's setup log and files explains how it gets there: +- It fills each GPU's expert cache with 7,507 experts in 10.95 GiB of VRAM (about 1.5 MB per IQ3_XXS expert), roughly 15,000 experts across the two + cards, about 61% of the model's 24,576. Ours holds 38 lines x 49 layers per rank, 3.23 GiB per rank, about 3,700 experts across both ranks (15%). +- It splits layers across the two cards, so each card holds the dense weights of only its own layers. Ours is expert-parallel: each rank's ledger + shows `persistent=7.53 GB` (dense weights, replicated or mostly replicated, and other fixed state), `scratch`, `deep_prefill`, `pool_scratch` + and `graph_reservation`, leaving 3.47 GB for expert pools out of a 16 GB card. Why persistent is 7.5 GB was not broken down in this work. +- It ships a pretrained hot-expert profile (`data/expert-profile.bin`, passed as `--expert-profile`) that seeds the cache. Our profile-sized pools + failed to generalize on held-out prompts earlier (-2.4% throughput), but those were sizing experiments, not seeding the same cache. +- Its CPU lane runs cold experts from host RAM. The vendored kernel `src/dsv41/vendor/exl3_cpu.cpp` (MIT, from project-maya) implements the EXL3 mul1 + codebook for K bits 1-8 with AVX-512 or AVX2+FMA. cs13's EPYC 7452 has AVX2+FMA only (no AVX-512), so the AVX2 tier would run there. With 94-99% + hits the CPU lane serves only the remaining 1-6%. + +Consequences for the plan (hypotheses to test, not results): raising our expert cache from 3.2 GiB to near 10 GiB per rank is worth more than a CPU +lane or MTP tuning. At 26 to 38 lines, every 12 added lines measured about +10% decode (summed per-request 26.95 to 29.7 tok/s), and the hit rate +rises faster than linearly as the cache approaches the hot set. The first steps are to find what the 7.5 GB `persistent` holds and what can move to +host memory or be sharded across the ranks, then to seed the pools from a hot-expert profile at startup. + +### Where the 16 GB goes, and how many expert lines fit + +Per-rank memory ledger at 38 lines (default features, `--parallel 2`, 32,768 context, shared BF16 KV paging), from the two ranks' startup logs: + +| item | rank 0 | rank 1 | +|---|---|---| +| weights resident on the device | about 5.0 GB | about 5.0 GB | +| shared BF16 KV window (`common allocation`) | 2.02 GB | 2.02 GB | +| MTP head (loaded unless `--no-drafts`) | 0.50 GB | 0 | +| prompt scratch | 0.76 GB | 0.76 GB | +| deep prefill rows (4096) | 0.73 GB | 0.73 GB | +| pool scratch | 0.52 GB | 0.31 GB | +| two streams | 1.09 GB | 0.81 GB | +| graph reservation | 0.39 GB | 0.39 GB | +| startup reserve floor | 1.07 GB (`TF_FLASHNEXT_EXL3_RESERVE_MB`, default 1024) | same | +| expert pools (38 lines) | 3.47 GB | 3.54 GB | + +The ledger line `free_after_pools` (2.8 GB and 3.4 GB) is not spare room: admission also holds back the reserve and the future bytes (streams, graphs, +planner), so 38 lines is already near the limit; 44 lines at the default settings and 56 lines were refused with `InsufficientStartupMemory`. +Real headroom comes from not loading the head (`--no-drafts` frees 0.5 GB on rank 0) and from a smaller reserve floor. + +Sweep with `--no-drafts` and `TF_FLASHNEXT_EXL3_RESERVE_MB=512`, branch defaults otherwise (overlap on, window off, TP=2, graphs), single request, +`tools/exl3_single_request_bench.py`, 192 forced tokens, three requests per prompt (decode tok/s): + +| lines per layer | 70-token prompt | 16k Zig prompt | 12.8k docs prompt | result | +|---|---|---|---|---| +| 38 (default, drafts loaded) | 33.7-36.5 | 26.6-29.1 | 28.7-29.3 | baseline | +| 44 | 34.5-36.4 | 28.6-30.1 | 30.0-33.9 | fits | +| 48 | 34.1-36.7 | 29.3-34.7 | 29.9-31.4 | fits | +| 52 | 35.8-38.2 | 30.9-31.5 | 30.6-32.7 | fits | +| 56 | | | | starts loading, never serves | + +Going from 38 to 52 lines (+37% cache) is worth about +6% to +12% decode, and time to first token is unchanged (10.4 s at 16k). That is a +configuration change that needs no code, but it is far short of the gap: Strata's cache holds about four times as many experts and reaches a +94-99% hit rate. To get there the resident state per rank has to shrink by several GB (weights, the 2 GB KV window, prefill scratch), which means a layer +split or other sharding of the dense weights rather than another tuning knob. + +### What the pack holds, and a sized plan for a larger expert cache + +Tensor sizes by role in `Qwen3.8-Flash-Next-exl3-3.05bpw` (84.95 GB total, from the safetensors headers, `tools/exl3_pack_sizes.py`): + +| role | size | where it lives today | +|---|---|---| +| routed experts | 45.77 GB | pinned host banks (about 22 GB per rank), expert pools on the device | +| per-layer n-gram embedding shards | 32.64 GB | host memory | +| DeltaNet attention | 1.32 GB | device, both ranks | +| hyper-connection mixes | 1.28 GB | device, both ranks | +| embed_tokens | 1.27 GB | device | +| MTP head | 1.04 GB (0.5 GB on the device, rank 0) | device unless `--no-drafts` | +| lm_head | 0.40 GB | device (split by vocabulary) | +| full attention | 0.38 GB | device (heads split) | +| shared expert, router, projections, other | about 0.6 GB | device | + +The on-device dense set is therefore about 5.8 GB, and the measured 5.0 GB resident on each rank says most of it is replicated, not sharded. +A layer split would roughly halve it (-2.9 GB per rank). Against Strata's 10.95 GiB of expert cache per card, our 3.5 GB of pools is paying for: +12.5 GB of other state per rank at 38 lines (weights 5.0, KV window 2.0, prompt scratch 0.76, deep prefill 0.73, pool scratch 0.3-0.5, streams 0.8-1.1, +graph reservation 0.39, reserve floor 1.07, head 0.5 on rank 0). + +Sized estimate of what could be recovered per rank (arithmetic from the ledger above, not tested): layer-split dense weights 2.9 GB, a smaller KV +device window about 1.0 GB, the reserve floor 0.5 GB, prefill rows and scratch about 0.4 GB, no MTP head 0.5 GB on rank 0. Together about 5 GB, which would take +the expert pools from 3.5 GB to roughly 8.5 GB (about 2.4 times the cache; at 69-84 MB per line that is on the order of 100 lines). From the 38 to 52 +line sweep, +37% cache gave +6% to +12% decode, so even a 2.4-fold cache is more likely to land decode in the 40s of tokens per second than at +Strata's 75-104, which also benefits from its CPU lane and from MTP drafting that pays off at a 95% hit rate. The expert-cache enlargement is the first +and largest step; the CPU lane (the vendored `exl3_cpu.cpp` runs on cs13's AVX2) and revisiting MTP follow it. + +### 68 lines per layer: the shared KV window was sized for 8 lanes + +With shared BF16 KV paging on, `TF_FLASHNEXT_KVPAGE_SHARED_LANES` defaults to 8 lanes (`cuda_exl3_kvpage.zig`), while the server runs +`--parallel 2`. Setting it to 2 shrinks the `shared BF16 KV common allocation` from 2,023,234,372 to 809,961,172 bytes per rank (1.21 GB), and the pools can +take it. Sweep with `--no-drafts`, `TF_FLASHNEXT_EXL3_RESERVE_MB=512` and `TF_FLASHNEXT_KVPAGE_SHARED_LANES=2`, branch defaults otherwise, single +request, `tools/exl3_single_request_bench.py`, 192 forced tokens, three requests per prompt (decode tok/s; the 70-token row excludes the first, +cold request): + +| lines per layer | 70-token prompt | 16k Zig prompt | 12.8k docs prompt | result | +|---|---|---|---|---| +| 38 (defaults, drafts loaded, 8 lanes) | 34.2-36.5 | 26.6-29.1 | 28.7-29.3 | baseline | +| 52 (no drafts, 512 MB reserve) | 36.2-38.2 | 30.9-31.5 | 30.6-32.7 | fits | +| 60 (plus 2 lanes) | 38.6-39.6 | 32.6-34.4 (first 27.9) | 31.9-34.1 | fits | +| 68 (plus 2 lanes) | 40.4-40.9 | 35.1-35.4 (first 32.2) | 32.9-34.0 | fits | +| 76 (plus 2 lanes) | | | | refused | + +`tools/exl3_serve_tp2.sh` now defaults to this configuration (68 lines, no drafts, 512 MB reserve, KV lanes equal to `PARALLEL`) and was run end to end on cs13: +38.8-40.7 tok/s on the 70-token prompt, 33.3-34.7 on the 16k prompt, 32.9-35.4 on the 12.8k prompt, time to first token unchanged (0.57 s, +10.4 s, 9.5 s). Against the defaults that is about +13% to +20% decode on short prompts and +16% to +26% on long ones, for no code change. Limits: the +KV window sized for two lanes means a third concurrent request would not fit, 72 lines was not tried, the 512 MB reserve was tested only in these +runs, and the gap to Strata (75-104 tok/s) is still 2.1 to 3 times. + +### One user: 80 lines, and where the curve flattens + +For a single-user server (`--parallel 1`, `TF_FLASHNEXT_KVPAGE_SHARED_LANES=1`, `TF_FLASHNEXT_EXL3_RESERVE_MB=256`) the shared KV allocation is 0.61 GB per rank. At 76 +lines the stream check then refused to start ("the CUDA memory budget fits no stream"), because it applies a separate 1.0 GiB device reserve. That reserve is +`TENSORFOLD_MEMORY_RESERVE_GIB` (default the larger of 1 GiB and 1/32 of the card; the code allows 0.5 GiB at the least). With it set to 0.5, 72, 76 and 80 lines serve. +Single request, `tools/exl3_single_request_bench.py`, 192 forced tokens, three requests per prompt, decode tok/s (the 70-token row excludes the cold first request): + +| setup | lines | 70-token | 16k Zig | 12.8k docs | +|---|---|---|---|---| +| defaults (drafts loaded, 8 KV lanes, `--parallel 2`) | 38 | 34.2-36.5 | 26.6-29.1 | 28.7-29.3 | +| no drafts, 2 KV lanes, reserve 512 MB, `--parallel 2` | 68 | 40.4-40.9 | 35.1-35.4 | 32.9-34.0 | +| plus `--parallel 1`, 1 KV lane, reserves 256 MB and 0.5 GiB | 72 | 40.7-40.7 | 34.6-35.7 | 33.3-35.1 | +| same | 76 | 42.1-42.5 | 34.3-35.6 | 34.9-36.2 | +| same | 80 | 41.7-42.8 | 34.7-36.1 | 34.8-36.8 | + +Beyond about 68 lines each further four lines add only 1% to 3%, and the long-prompt rows barely move (35 tok/s, about 28.5 ms per token). That matches the +nsys profile: the roughly 23 ms of each step that is not expert-miss copying (all-reduce waiting, the dense and expert kernels) is now most of the step. Strata's 75-104 tok/s means 10 to 13 ms per +token, so its per-step cost is also less than half of our non-miss time even before drafting. The cache enlargement is therefore done (about +20% to +30% over the defaults); the next gains have to come from the +step itself (the per-layer all-reduce of expert-parallel TP, kernel efficiency) and from making drafting pay, not from more lines. +`tools/exl3_serve_tp2.sh` uses 80 lines when `PARALLEL=1` and 68 otherwise. The 80-line single-user setup was run through the sweep script, not through the launch script, and the 0.5 GiB reserve is the +supported minimum with no headroom for other GPU users. + +### The decode step at 80 lines, from nsys (rank 0) + +nsys trace of rank 0 with the single-user setup (80 lines, `--parallel 1`, no drafts, graphs; a short story prompt, 256 generated tokens, three runs at +40.3, 44.1 and 44.7 tok/s under the profiler), counted per decode step (765 steps, `tools/exl3_nsys_step_breakdown.py`). Wall time was 24.2 ms per step under the profiler +and the non-copy kernels kept the GPU 81% busy. + +| category | ms per step | launches per step | +|---|---|---| +| EXL3 matmuls (`tf_exl3_linear_*`, `tf_exl3_grouped_*`, gate-up and down epilogues) | 10.85 | 819 | +| `qmmf` matmuls | 4.56 | 234 | +| NCCL all-reduce and gather | 3.07 | 49 | +| expert miss copy (side stream, overlapped) | 2.07 | 47 | +| hyper-connection kernels | 0.95 | 390 | +| router | 0.65 | 96 | +| DeltaNet kernels | 0.42 | 108 | +| planner, attention and KV paging, other | 0.85 | about 400 | + +Largest single kernels: `tf_exl3_linear_10_2_8` 3.87 ms (49 launches, 79 us each, 80-96 blocks), `tf_exl3_linear_10_2_4` 3.41 ms (108 launches, 31 us), NCCL all-reduce 2.94 ms +(48 launches, 61 us each, mostly waiting for the other rank), `qmmf` 2.68 and 1.87 ms, `tf_exl3_grouped_2` 2.57 ms (192 launches, 13 us). + +What this says: the expert misses are no longer the cost (2.07 ms of copy, hidden on a side stream); the step is about 15.4 ms of matmul kernels, 3 ms of all-reduce wait and the rest small kernels. +The dense matmuls run identically on both ranks. Pack inspection (`tools/exl3_pack_sizes.py` and the trellis shapes): routed experts are 3-bit, the DeltaNet projections (1.3 GB) and `lm_head` are 5-bit, +and the on-device dense set is about 5.8 GB, so a rank streams a few GB of dense weights per token (roughly 3.6 GB per step including experts and the head, which is about 8 ms at the card's 448 GB/s). +That mapping from kernel names to weights is an inference from sizes and launch counts, not confirmed by a per-kernel trace of the weights each reads. Implications that follow (hypotheses): sharding the dense +projections across the ranks would cut that per-rank streaming roughly in half but adds collectives per layer, and each NCCL call already costs 61 us here, so it only pays if the +gathers are folded into the existing per-layer all-reduce; a layer split removes the all-reduce (about 3 ms) and the replication but leaves one card idle per token for a single stream. +Neither was prototyped. + +### Which weights the decode kernels read, and how close they run to the memory limit + +Kernel names decode as `tf_exl3_linear_{2 x bits}_{codebook}_{warps}` (`exl3_linear.cu`: `K2` is twice the width in bits, `CB` 2 is the mul1 codebook, `WK` the warps per block), so +`tf_exl3_linear_10_2_8` is 5-bit mul1 with 8 warps and `tf_exl3_linear_10_2_4` is 5-bit with 4 warps. The pack's 5-bit tensors (trellis shapes, `tools/exl3_pack_sizes.py` +and the bit listing) are the DeltaNet projections (`in_proj_qkv` 0.59 GB, `in_proj_z` 0.35 GB, `out_proj` 0.35 GB over 36 layers), full attention (`q_proj` 0.24 GB, `o_proj` 0.12 GB, +`k_proj` and `v_proj` 0.01 GB each over 12 layers), the shared expert (0.15 GB over 48 layers; rank 1 only) and `lm_head` (0.40 GB, vocabulary split). Routed experts are 3-bit. +Matching launches per step and grid sizes (grid x is N / 128) from the 80-line nsys trace of rank 0: + +| kernel | launches per step | grid | likely weights (inferred) | bytes read per step | ms per step | achieved bandwidth | +|---|---|---|---|---|---|---| +| `tf_exl3_linear_10_2_8` | 49 | 80 or 96 blocks | DeltaNet `in_proj_qkv` (36, N 10240) + full-attention `q_proj` (12, N 12288) + one more | about 0.83 GB | 3.87 | about 215 GB/s | +| `tf_exl3_linear_10_2_4` | 108 | 16 to 160 blocks | DeltaNet `in_proj_z` and `out_proj` (72) + full-attention `o_proj`, `k_proj`, `v_proj` (36) | about 0.85 GB | 3.41 | about 250 GB/s | +| `qmmf_kernel` (two variants) | 234 | small grids | 4 hyper-connection mix matrices per layer (192) plus others; 1.28 GB in the pack | about 1.3 GB | 4.56 | about 280 GB/s | +| `tf_exl3_grouped_2` | 192 | 11 x 5 x 8 etc. | routed experts (3-bit), the rank's share | small | 2.57 | not computed | + +The card's memory bandwidth is 448 GB/s, so these dense kernels run at roughly 50% to 62% of the limit, and about 11.8 ms of the 24 ms step is spent streaming dense weights that both ranks read in full. The kernel-to-weight match +rests on launch counts, grid sizes and the pack's tensor sizes; it was not confirmed by tracing which tensor each launch reads, and the bandwidths assume each weight is read once. + +Two ways to cut that 11.8 ms while keeping tensor parallelism (neither prototyped): (1) make the kernels faster, by moving the 5-bit EXL3 dense matmuls and the `qmmf` matrices from about 55% to about 85% +of the memory limit, worth roughly 4 ms of the step, with no structural change; (2) shard the dense projections across the two ranks (column-parallel input projections, row-parallel +`out_proj`, and the hyper-connection mixes folded into the existing per-layer reduction), which would halve each rank's streaming (up to about 5 to 6 ms) but adds a second +collective per layer, which costs 25 to 60 us per call here, so it pays only if the extra reductions are fused into the one we already do. + +### Why the dense kernels run at about half the memory limit, and what stops a quick fix + +The launch shape of every EXL3 linear comes from `plan()` in `zig/src/cuda/exl3.zig`, called as `plan(k, n, 192, 8)`: when the output has at least 64 blocks of 128 columns (N >= 8192) it uses 8 warps and no split of K, so the grid is just N / 128 blocks +(80 or 96 for the DeltaNet and full-attention input projections), about 2.2 blocks of 256 threads per SM on a 36-SM card; only narrower outputs (under 64 blocks) are split along K until they reach 192 blocks. That under-occupies the memory system +for the widest matrices, which fits the achieved 215 GB/s of the 448 GB/s limit. The obvious change is to split K for those shapes too (and so use about 192 blocks), but the function's comment says the plan is chosen from the shape alone "so a row's +bits never depend on the row count (linear.py plan)": the split changes the order of the fp32 accumulation, so the outputs would no longer be bit-identical to the current kernels or to the reference, and the branch's earlier verification +(layers 0, 3 and 17 bit-identical to TP=1, graphs bit-equal to eager) depends on that. Any speedup of the dense kernels therefore has to either keep the accumulation order (deeper prefetch or more loads in flight per +thread inside the existing grid) or be accepted as near-exact and checked by token hashes instead of bit equality. The same applies to sharding the dense projections across the ranks, which splits sums between cards. Neither was built. + +### Result: splitting K for the wide 5-bit matmuls does not speed decode (rejected) + +To test the under-occupancy explanation in the previous section, `plan()` was given an opt-in (`TF_EXL3_WIDE_SPLIT=1`) that splits K for outputs of 64 to 191 column blocks, which moves the 80- and 96-block +projections to 160-320 blocks. It built and passed the host tests. A/B on cs13 at 72 lines per layer, single-user (`PARALLEL=1`, no drafts), the launch script, three requests per prompt +(decode tok/s; legs in the order off, on, on; the second "off" leg did not finish and is not used): + +| prompt | split off | split on (two legs) | +|---|---|---| +| 70-token | 39.2, 40.7, 40.6 | 38.8, 40.5, 41.6 and 38.8, 41.8, 41.8 | +| 16k Zig | 33.9, 34.4, 35.3 | 33.9, 33.7, 36.4 and 34.8, 33.3, 36.2 | +| 12.8k docs | 34.4, 34.4, 35.2 | 33.2, 35.6, 35.2 and 33.4, 34.2, 34.6 | + +No gain outside run variation (about 3% on these rows), so more blocks in flight for the widest matrices is not what limits them; the extra partial-sum scratch also made +80 lines no longer fit (rank 1 was 0.02 GiB short of its stream budget). The change was removed. The output-hash comparison planned for this test was not usable (the bench prefixes each request with a unique id). +The 215-280 GB/s the dense kernels reach is therefore not explained by launch shape; the trellis decode work per byte, instruction issue, or the replicated small kernels around them (390 hyper-connection +launches and 819 matmul launches per step) are the remaining suspects, none profiled at instruction level (ncu was not run). + +### The cards are not throttled; the dense kernels are not bandwidth-bound + +`nvidia-smi` sampled every 0.5 s on both 5060 Ti while the single-user setup (72 lines, `PARALLEL=1`, no drafts) decoded 512-token replies (567 busy samples at 50% utilization or more, per card): + +| | GPU 0 | GPU 1 | +|---|---|---| +| SM clock | 2,898 MHz average (max 3,210) | 2,879 MHz (max 3,210) | +| memory clock | 13,801 MHz average (max 14,001) | 13,801 MHz | +| power | 102 W average (limit 180) | 103 W | +| GPU utilization / memory utilization | 95% / 45% | 95% / 47% | +| performance state | P1 | P1 | + +There is no power cap, no memory down-clock (no compute-state P2 cap) and no thermal limit (59-63 C). The GPU is busy 95% of the time but its memory system only 45%, which agrees with the nsys estimate that the dense 5-bit +and `qmmf` matmuls reach about 215-280 GB/s of the 448 GB/s available: their time goes to latency or instruction issue inside the kernels, not to waiting on DRAM. Decode in this run was 40.6-43.0 tok/s on the short prompt and +29.9-38.3 tok/s on the 16k prompt (512 tokens, three requests each; the long-prompt spread is wide). Finding out which of latency, instruction issue or occupancy is responsible needs the performance counters (`ncu`), which this account cannot read. diff --git a/tools/exl3_dplan_predict.py b/tools/exl3_dplan_predict.py new file mode 100644 index 0000000000..6aaa2724db --- /dev/null +++ b/tools/exl3_dplan_predict.py @@ -0,0 +1,87 @@ +"""Next-layer expert predictability from DPLAN v4 trace dumps (TF_FLASHNEXT_EXL3_DPLAN_TRACE_DIR). Usage: dir.""" +import struct, glob, os, re, collections, sys +D = sys.argv[1] +pools = {} +for f in glob.glob(D + "/*.bin"): + idx = int(re.search(r"pool(\d+)", f).group(1)) + w = struct.unpack("<%dI" % (os.path.getsize(f) // 4), open(f, "rb").read()) + layer, lines, maxr, sw, stride, cap, rec = w[3], w[4], w[5], w[6], w[7], w[8], w[9] + off = 16 + 2 * sw + rounds = [] + for k in range(rec): + r = w[off + k * stride: off + (k + 1) * stride] + rows = r[0] + picks = [struct.unpack(" 0] +K = min(len(pools[i][2]) for i in L) +rows_ok = sum(1 for k in range(K) if len({pools[i][2][k][0] for i in L}) == 1) +print("layers with data", len(L), "rounds", K, "rounds with equal rows across layers", rows_ok) +sets = {i: [pools[i][2][k][1] for k in range(K)] for i in L} +LINES = pools[L[0]][1] +# per-layer LRU simulation -> miss sets +miss = {} +for i in L: + cache = collections.OrderedDict() + ms = [] + for k in range(K): + need = sets[i][k] + m = set(e for e in need if e not in cache) + for e in need: + if e in cache: cache.move_to_end(e) + for e in m: + while len(cache) >= LINES: + for old in cache: + if old not in need: + del cache[old]; break + else: + break + cache[e] = 1 + ms.append(m) + miss[i] = ms +tot_miss = sum(len(m) for i in L for m in miss[i]); tot_need = sum(len(s) for i in L for s in sets[i]) +print("LRU(%d lines) misses per layer-round %.2f of %.2f needed (%.1f%%)" % (LINES, tot_miss / (K * len(L)), tot_need / (K * len(L)), 100 * tot_miss / tot_need)) +T = K // 2 +def evaluate(Ns, use_prev_round=False): + # co-occurrence: layer i -> layer i+1 within the same round, trained on rounds < T + out = {N: [0, 0, 0, 0] for N in Ns} # covered misses, total misses, predicted-not-cached, useful(pred & need & notcached) + for a, b in zip(L, L[1:]): + co = collections.defaultdict(collections.Counter) + freq = collections.Counter() + for k in range(T): + freq.update(sets[b][k]) + for e in sets[a][k]: + co[e].update(sets[b][k]) + # replay cache state for layer b through test rounds again to know what's cached + cache = collections.OrderedDict() + for k in range(K): + need = sets[b][k] + if k >= T: + score = collections.Counter() + for e in sets[a][k]: + tot = sum(co[e].values()) or 1 + for x, c in co[e].items(): score[x] += c / tot + for x, c in freq.items(): score[x] += 0.01 * c / T + ranked = [x for x, _ in score.most_common() if x not in cache] + m = set(x for x in need if x not in cache) + for N in Ns: + pred = set(ranked[:N]) + out[N][0] += len(m & pred); out[N][1] += len(m); out[N][2] += len(pred); out[N][3] += len(m & pred) + m = set(e for e in need if e not in cache) + for e in need: + if e in cache: cache.move_to_end(e) + for e in m: + while len(cache) >= LINES: + for old in cache: + if old not in need: + del cache[old]; break + else: break + cache[e] = 1 + return out +res = evaluate([2, 4, 8, 16, 32]) +for N, (cov, tm, pred, useful) in res.items(): + print("prefetch top-%d from previous layer's picks: covers %.1f%% of misses; %.2f copies issued per layer-round, %.1f%% of issued copies useful" % (N, 100 * cov / tm, pred / ((K - T) * (len(L) - 1)), 100 * useful / pred)) diff --git a/tools/exl3_kv_budget.py b/tools/exl3_kv_budget.py new file mode 100644 index 0000000000..58e9b96cba --- /dev/null +++ b/tools/exl3_kv_budget.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +"""Per-card memory budget for Qwen3.8-Flash-Next EXL3 at 200k context on two 16 GiB cards (docs/plans/exl3-kv-memory.md). + + exl3_kv_budget.py [CONTEXT] + +Every byte formula is upstream's (cuda_prompt.zig seqBytes, cuda_batch_state.zig Scratch.bytes, cuda_mtp_draft.zig +Draft.bytes) with the K10 levers applied (cuda_exl3_kv.zig seqBytes gives the same numbers; its host test checks the +default policy against upstream's). Inputs that are not formulas are marked where they come from. +""" +import sys + +G = 2 ** 30 +CTX = int(sys.argv[1]) if len(sys.argv) > 1 else 200_000 +CAP = CTX + 16 # the window a stream holds: context + mtp.max_rows (cuda_lanes_load.zig) + +# -- model shapes (cuda_prompt_wt.zig, cuda_gdn.zig) ------------------------------------------------------------------ +dims, streams, low, vocab, experts_n, slots, moe_w, tile = 2560, 4, 320, 248320, 512, 11, 640, 16 +wide = streams * dims +q_heads, kv_heads, head_dim, index_heads, index_dim, half, proj_n, out_k, nch = 24, 2, 256, 4, 128, 32, 13952, 6144, 5 +idw, att_rows, hc_sk, out_sk, plan_small = 2052, 256, 32, 8, 1024 +nk, nv, dk, dv, taps = 16, 48, 128, 128, 4 +conv_dim = 2 * nk * dk + nv * dv +value_dim = nv * dv +proj_row = max(conv_dim + nv * dv + 2 * nv, proj_n) +pss_n = (dims // 256) * streams +state_n = nv * dv * dk +linear_n, attn_n, ple_tail_n = 36, 12, 9 +cdiv = lambda a, b: (a + b - 1) // b + + +def scratch(cap, whole_head=False): + """cuda_batch_state.zig Scratch.bytes(cap), term for term; whole_head drops head_local (cap * vocab bytes), which + only a vocabulary-split head reads (cuda_batch.zig projectHead).""" + pairs = cap * slots + nblk = cdiv(pairs, plan_small) + items = cdiv(pairs, tile) + experts_n + 1 + n = cap * 4 + (cap + 1) * 16 + n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4 + n += cap * wide * 2 + 8 * cap * dims * 4 + n += cap * low * 2 + cap * (low // 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims // 32) * 4 + n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * value_dim * (2 + 4 + 4) + cap * (value_dim // 32) * 4 + n += cap * (2 * nk * dk + 2 * nv) * 4 + n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2 + n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k // 32) * 4 + n += att_rows * (idw + 2) * 4 + n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4 + n += cap * moe_w * 6 + cap * dims * 4 + n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2 + n += cap * dims * 2 + cap * (dims // 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4 + n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + (0 if whole_head else cap * vocab) + half * 4 + 4096 + return n + + +def kvrow(kv, heads): + per = {"bf16": head_dim * 2, "int8": head_dim + head_dim // 32 * 2, "int4": head_dim // 2 + head_dim // 32 * 2}[kv] + return 2 * heads * per + + +def pooled(cap): + return cdiv(cap, 4) * index_dim * 2 + + +def seq(cap, kv="bf16", heads=2, kv_host=False, ikc_host=False): + """(device, host) bytes of one stream: cuda_prompt.zig seqBytes with the levers.""" + gdn = linear_n * ((taps - 1) * conv_dim * 2 + state_n * 4) + k = attn_n * cap * kvrow(kv, heads) + i = attn_n * cap * index_dim * 2 + dev = gdn + (0 if kv_host else k) + (0 if ikc_host else i) + attn_n * pooled(cap) + ple_tail_n * wide * 2 + return dev, (k if kv_host else 0) + (i if ikc_host else 0) + + +def draft(cap, kv="bf16", kv_host=False, ikc_host=False, rows=1536, draft_n=79591): + """cuda_mtp_draft.zig Draft.bytes (rank 0, MTP head's one attention layer under the same policy).""" + fixed = max(rows, 16) * wide * 2 + wide * 2 + 16 * vocab * 2 + draft_n * 2 + 16 * 16 + 16 * wide * 2 + linear_n * state_n * 4 + return fixed + (0 if kv_host else cap * kvrow(kv, kv_heads)) + (0 if ikc_host else cap * index_dim * 2) + pooled(cap) + + +# -- inputs that are not formulas ----------------------------------------------------------------------------------------- +CARD = 15.568 # GiB, cuMemGetInfo total on a 5060 Ti (measured, exl3-kv bench) +CONTEXT_MODULES = 0.40 # GiB, CUDA context + every module the engine loads (estimate; our test's modules alone: 0.134 measured) +WEIGHTS = {0: 5.107 + 0.475, 1: 5.245} # GiB, T10's expert-parallel placement (read from the pack's headers); rank 0 + MTP head +OTHER = 0.05 + 0.015 + 0.05 # GiB: block scores at 200k (scoreBytes, read), lanes logits (read), EXL3 Dense scratch (estimate) +LINE = 0.136 # GiB a pool line in every one of 48 layers (T10: 10 lines/layer = 1.36 GiB) +FLOOR = 10 # lines a layer, the pool's floor +# decode hit rate at tp=1 by lines a layer, measured by E2 on real routing (512 experts in one pool) +HIT = [(10, 0.33), (24, 0.505), (48, 0.64), (96, 0.78), (160, 0.88)] +MISS_MS = 0.147 # one 3-bit expert staged from a pinned bank (measured, E2) + + +def hit_tp1(lines): + import math + if lines <= HIT[0][0]: + return HIT[0][1] * lines / HIT[0][0] + for (a, ha), (b, hb) in zip(HIT, HIT[1:]): + if lines <= b: + t = (math.log(lines) - math.log(a)) / (math.log(b) - math.log(a)) + return ha + t * (hb - ha) + return HIT[-1][1] + + +def miss_ms(lines_rank): + """Expert-miss staging a token on one rank (estimate): a rank serves its 256 experts, ~5 picks a row a layer; its + L lines hit like a tp=1 pool of 2L lines over all 512.""" + h = hit_tp1(2 * lines_rank) + return 48 * 5 * (1 - h) * MISS_MS, h + + +LEVERS = [ + # name, kv dtype, heads a rank, kv_host, ikc_host, chunk, whole-head scratch + ("A upstream: bf16 KV, chunk 1536", "bf16", 2, False, False, 1536, False), + ("B int8 KV", "int8", 2, False, False, 1536, False), + ("C = B + index keys in host memory", "int8", 2, False, True, 1536, False), + ("D = C + scratch: chunk 1024, no head_local", "int8", 2, False, True, 1024, True), + ("E = D + attention head split", "int8", 1, False, True, 1024, True), + ("F bf16 KV + head split + index keys on host + D's scratch", "bf16", 1, False, True, 1024, True), + ("G host KV (int8, mapped) + index keys on host + D's scratch", "int8", 2, True, True, 1024, True), +] + + +def table(drafts=True): + print(f"drafts {'on (rank 0 holds every stream' + chr(39) + 's draft cache)' if drafts else 'off'}") + print(f"| lever (context {CTX:,}) | seqs | rank 0 seq+draft GiB | rank 0 left for pool GiB | lines/layer r0 | r1 lines | decode miss ms/token (est.) | host GiB |") + print("|---|---:|---:|---:|---:|---:|---:|---:|") + for name, kv, heads, kvh, ikh, chunk, whole in LEVERS: + for n in (1, 2, 4): + d, h = seq(CAP, kv, heads, kvh, ikh) + dr = draft(CAP, kv, kvh, ikh, rows=chunk) if drafts else 0 + fixed = CONTEXT_MODULES + scratch(chunk, whole) / G + OTHER + left = {r: CARD - WEIGHTS[r] - fixed - n * d / G - (n * dr / G if r == 0 else 0) for r in (0, 1)} + lines = {r: max(0, int(left[r] / LINE)) for r in (0, 1)} + ok = min(lines.values()) >= FLOOR + ms, hit = miss_ms(min(lines.values())) if ok else (float("nan"), 0) + host = n * h / G + print(f"| {name} | {n} | {n * (d + dr) / G:.2f} | {left[0]:.2f} | {lines[0] if left[0] > 0 else 'no fit'} | {lines[1] if left[1] > 0 else 'no fit'} | {'%.1f (hit %.2f)' % (ms, hit) if ok else 'below floor'} | {host:.1f} |") + + +if __name__ == "__main__": + print(f"scratch GiB by chunk: " + ", ".join(f"{c}: {scratch(c) / G:.3f} ({scratch(c, True) / G:.3f} whole head)" for c in (1536, 1024, 512, 256))) + d, _ = seq(CAP) + print(f"one stream at {CTX:,} (bf16, upstream): {d / G:.3f} GiB; draft cache {draft(CAP) / G:.3f} GiB") + print() + table(True) + print() + table(False) diff --git a/tools/exl3_kvpage_replay.py b/tools/exl3_kvpage_replay.py new file mode 100644 index 0000000000..07f717a032 --- /dev/null +++ b/tools/exl3_kvpage_replay.py @@ -0,0 +1,339 @@ +#!/usr/bin/env python3 +"""Offline replay of EXL3 Flash Next sparse-attention selection traces (tf-cuda-test exl3-kvpage DIR trace ...). + +A trace holds, per attention layer (0..11) and per attention call (a 256-row prompt block or one decode row), each +row's 512 selected 4-key blocks (or "dense": the row reads every key, <= 2051). This replays a device block cache: +a recent dense window of W tokens (always on the card) plus C block slots per layer with exact LRU at call +granularity (the GPU's policy in cuda_exl3_kvpage.zig: a call's union is stamped, then the oldest slots are evicted). + +Usage: exl3_kvpage_replay.py TRACE.bin [--windows 2048,4096,8192] [--caches 256,...] [--json OUT] +""" +import argparse +import json +import sys + +import numpy as np + +RATIO = 4 +TOP = 512 +BLOCK_BYTES = 4 * 2 * 2 * 256 * 2 # 4 keys x (K and V) x 2 KV heads x 256 dims x bf16 = 8 KiB +LAYERS = 12 + + +def parse(path, ordered=None): + """Per layer: [(pos, kind, rows)]; with `ordered` (a list) also every call in file (= execution) order.""" + raw = np.fromfile(path, dtype=np.uint8) + calls = [[] for _ in range(LAYERS)] # (pos, kind, [rows: array or None]) + at = 0 + n = raw.size + while at < n: + ai, pos, m, kind = np.frombuffer(raw[at:at + 16].tobytes(), dtype=np.uint32) + at += 16 + rows = [] + for _ in range(int(m)): + cnt = int(np.frombuffer(raw[at:at + 2].tobytes(), dtype=np.uint16)[0]) + at += 2 + if cnt == 0: + rows.append(None) + else: + rows.append(np.frombuffer(raw[at:at + 2 * cnt].tobytes(), dtype=np.uint16).astype(np.int64)) + at += 2 * cnt + calls[int(ai)].append((int(pos), int(kind), rows)) + if ordered is not None: + ordered.append((int(ai), int(pos), int(kind), rows)) + return calls + + +class Lru: + """C slots, exact LRU over calls; blocks are ints < nblocks.""" + + def __init__(self, cap, nblocks): + self.cap = cap + self.stamp = np.full(nblocks, -1, dtype=np.int64) # last call that used the block while resident + self.res = np.zeros(nblocks, dtype=bool) + self.count = 0 + self.t = 0 + + def use(self, need): + """`need`: unique blocks of one call (not in the window). Returns misses (overflow counted as misses).""" + self.t += 1 + hits = self.res[need] + miss = need[~hits] + self.stamp[need] = self.t + over = self.count + miss.size - self.cap + if over > 0: + resident = np.flatnonzero(self.res) + old = resident[self.stamp[resident] < self.t] + k = min(over, old.size) + if k > 0: + ev = old[np.argpartition(self.stamp[old], k - 1)[:k]] if k < old.size else old + self.res[ev] = False + self.count -= ev.size + room = self.cap - self.count + take = miss[:max(room, 0)] + self.res[take] = True + self.count += take.size + return int(miss.size), int(max(miss.size - take.size, 0)) + + +def window_start(end, w): + """First token of the dense window for a call whose last row is at end - 1 (rounded down to a block).""" + return max(0, (end - w) // RATIO * RATIO) + + +def outside(blocks, w0): + return blocks[blocks * RATIO < w0] + + +def decode_stats(calls, windows, caches, nblocks, warm): + """Per decode row: overlap with the previous row, new blocks, hit rates and bytes for each (W, C).""" + out = {} + rows = [(p, r[0]) for (p, k, r) in calls if k == 1 and r[0] is not None] + if len(rows) < 2: + return out + ov = [np.intersect1d(a, b, assume_unique=True).size / TOP for (_, a), (_, b) in zip(rows, rows[1:])] + new = [np.setdiff1d(b, a, assume_unique=True).size for (_, a), (_, b) in zip(rows, rows[1:])] + out["rows"] = len(rows) + out["pos"] = [rows[0][0], rows[-1][0]] + out["overlap_prev_mean"] = float(np.mean(ov)) + out["new_blocks_per_row_mean"] = float(np.mean(new)) + u4 = [] + for i in range(0, len(rows) - 3, 4): + u4.append(np.unique(np.concatenate([rows[i + j][1] for j in range(4)])).size) + out["mtp4_union_mean"] = float(np.mean(u4)) if u4 else None + for w in windows: + frac_win = np.mean([1 - outside(b, window_start(p + 1, w)).size / TOP for p, b in rows]) + out[f"W{w}_frac_in_window"] = float(frac_win) + for c in caches: + for group in (1, 4): + lru = Lru(c, nblocks) + if warm is not None: + for need in warm(w): + lru.use(need) + miss = 0 + over = 0 + sel = 0 + needed = 0 + per = [] + for i in range(0, len(rows) - group + 1, group): + end = rows[i + group - 1][0] + 1 + w0 = window_start(end, w) + allb = np.unique(np.concatenate([rows[i + j][1] for j in range(group)])) + need = outside(allb, w0) + mi, ov_ = lru.use(need) + miss += mi + over += ov_ + sel += allb.size + needed += need.size + per.append(mi) + calls_n = len(per) + key = f"W{w}_C{c}" + ("" if group == 1 else "_mtp4") + out[key] = { + "hit_rate_selected": 1 - miss / max(sel, 1), + "hit_rate_outside_window": 1 - miss / max(needed, 1), + "misses_per_call_mean": miss / max(calls_n, 1), + "misses_per_call_p95": float(np.percentile(per, 95)) if per else 0, + "misses_per_call_max": int(max(per)) if per else 0, + "bytes_per_call_per_layer": miss / max(calls_n, 1) * BLOCK_BYTES, + "overflow": over, + } + return out + + +def prompt_calls(calls, w): + """Each prompt call's unique needed blocks outside the window, in order (for warming and for prompt replay).""" + res = [] + for p, k, rows in calls: + if k != 0: + continue + sp = [r for r in rows if r is not None] + if not sp: + res.append(np.zeros(0, dtype=np.int64)) + continue + end = p + len(rows) + w0 = window_start(end, w) + res.append(np.unique(outside(np.concatenate(sp), w0))) + return res + + +def prompt_stats(calls, windows, caches, nblocks): + out = {} + pc = [(p, rows) for p, k, rows in calls if k == 0] + if not pc: + return out + out["calls"] = len(pc) + out["rows"] = sum(len(r) for _, r in pc) + # union of a 256-row call's selections (all blocks, and in buckets of position) + un = [] + for p, rows in pc: + sp = [r for r in rows if r is not None] + if sp: + un.append((p, np.unique(np.concatenate(sp)).size)) + if un: + arr = np.array(un) + out["union_per_call_mean"] = float(arr[:, 1].mean()) + out["union_per_call_max"] = int(arr[:, 1].max()) + buckets = {} + for lo, hi in ((0, 32768), (32768, 65536), (65536, 131072), (131072, 1 << 30)): + sel = arr[(arr[:, 0] >= lo) & (arr[:, 0] < hi)] + if sel.size: + buckets[f"{lo}-{hi}"] = {"mean": float(sel[:, 1].mean()), "max": int(sel[:, 1].max())} + out["union_by_pos"] = buckets + for w in windows: + needs = prompt_calls(calls, w) + out[f"W{w}_outside_union_mean"] = float(np.mean([x.size for x in needs])) + out[f"W{w}_outside_union_max"] = int(max(x.size for x in needs)) + for c in caches: + lru = Lru(c, nblocks) + miss = 0 + over = 0 + for need in needs: + mi, ov_ = lru.use(need) + miss += mi + over += ov_ + out[f"W{w}_C{c}"] = { + "miss_blocks_total": miss, + "miss_bytes_per_row_per_layer": miss * BLOCK_BYTES / out["rows"], + "miss_gib_total_per_layer": miss * BLOCK_BYTES / 2**30, + "overflow_calls_blocks": over, + } + return out + + +def shared_replay(ordered, windows, totals, nblocks): + """One slot pool for all 12 layers of the sequence (keys layer * nblocks + block), calls in execution order: + the prompt (layer-major inside a deep chunk), then decode rounds (12 layers a token). Also MTP 4-row verify + windows: 4 consecutive decode tokens as one call a layer.""" + out = {} + decode = [(ai, p, rows[0]) for ai, p, k, rows in ordered if k == 1 and rows[0] is not None] + prompt = [(ai, p, rows) for ai, p, k, rows in ordered if k == 0] + prompt_rows = sum(len(r) for _, _, r in prompt) // LAYERS + tokens = len(decode) // LAYERS + for w in windows: + pneeds = [] + for ai, p, rows in prompt: + sp = [r for r in rows if r is not None] + if not sp: + pneeds.append((ai, np.zeros(0, dtype=np.int64))) + continue + w0 = window_start(p + len(rows), w) + pneeds.append((ai, np.unique(outside(np.concatenate(sp), w0)) + ai * nblocks)) + dneeds = [(ai, outside(b, window_start(p + 1, w)) + ai * nblocks, b.size) for ai, p, b in decode] + # MTP: tokens t..t+3 of one layer as one call (the layers interleaved per window) + by_layer = [[x for x in decode if x[0] == li] for li in range(LAYERS)] + mneeds = [] + for t0 in range(0, tokens - 3, 4): + for li in range(LAYERS): + rows = by_layer[li][t0:t0 + 4] + if len(rows) < 4: + continue + w0 = window_start(rows[-1][1] + 1, w) + allb = np.unique(np.concatenate([r[2] for r in rows])) + mneeds.append((li, outside(allb, w0) + li * nblocks, allb.size)) + out[f"W{w}_prompt_union_max"] = int(max((x.size for _, x in pneeds), default=0)) + for tot in totals: + for kind in ("decode", "mtp4"): + lru = Lru(tot, LAYERS * nblocks) + pm = 0 + pov = 0 + for _, need in pneeds: + mi, ov_ = lru.use(need) + pm += mi + pov += ov_ + seq = dneeds if kind == "decode" else mneeds + miss = 0 + sel = 0 + needed = 0 + per_tok = [] + acc = 0 + for i, (_, need, nsel) in enumerate(seq): + mi, _ = lru.use(need) + miss += mi + sel += nsel + needed += need.size + acc += mi + if i % LAYERS == LAYERS - 1: + per_tok.append(acc) + acc = 0 + key = f"W{w}_S{tot}" + if kind == "decode": + out[key + "_prompt"] = {"miss_blocks": pm, "miss_bytes_per_row": pm * BLOCK_BYTES / max(prompt_rows, 1), "miss_gib": pm * BLOCK_BYTES / 2**30, "overflow_blocks": pov} + unit = "token" if kind == "decode" else "window" + out[key + "_" + kind] = { + "hit_rate_selected": 1 - miss / max(sel, 1), + "hit_rate_outside_window": 1 - miss / max(needed, 1), + f"misses_per_{unit}": float(np.mean(per_tok)) if per_tok else 0, + f"misses_per_{unit}_p95": float(np.percentile(per_tok, 95)) if per_tok else 0, + f"misses_per_{unit}_max": int(max(per_tok)) if per_tok else 0, + f"bytes_per_{unit}": float(np.mean(per_tok)) * BLOCK_BYTES if per_tok else 0, + } + print(key, kind, json.dumps(out[key + "_" + kind]), flush=True) + print(key, "prompt", json.dumps(out[key + "_prompt"]), flush=True) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("trace") + ap.add_argument("--windows", default="2048,4096,8192") + ap.add_argument("--caches", default="256,512,1024,2048,4096,8192") + ap.add_argument("--json") + ap.add_argument("--layers", default="all") + ap.add_argument("--shared", help="slot totals for one pool shared by the 12 layers, e.g. 12288,24576,49152") + a = ap.parse_args() + windows = [int(x) for x in a.windows.split(",")] + caches = [int(x) for x in a.caches.split(",")] + ordered = [] if a.shared else None + calls = parse(a.trace, ordered) + maxpos = max((p + len(r)) for c in calls for (p, k, r) in c) + nblocks = maxpos // RATIO + 16 + if a.shared: + rep = {"trace": a.trace, "tokens": maxpos, "shared": shared_replay(ordered, windows, [int(x) for x in a.shared.split(",")], nblocks)} + if a.json: + with open(a.json, "w") as f: + json.dump(rep, f, indent=1) + return 0 + layers = range(LAYERS) if a.layers == "all" else [int(x) for x in a.layers.split(",")] + report = {"trace": a.trace, "tokens": maxpos, "layers": {}} + for li in layers: + c = calls[li] + has_prompt = any(k == 0 for _, k, _ in c) + warm = (lambda w, c=c: prompt_calls(c, w)) if has_prompt else None + report["layers"][li] = {"decode": decode_stats(c, windows, caches, nblocks, warm), "prompt": prompt_stats(c, windows, caches, nblocks)} + d = report["layers"][li]["decode"] + print(f"layer {li}: decode rows {d.get('rows')}, overlap {d.get('overlap_prev_mean', 0):.3f}, new/row {d.get('new_blocks_per_row_mean', 0):.1f}, mtp4 union {d.get('mtp4_union_mean')}", flush=True) + # summary over layers + summ = {} + for w in windows: + for c in caches: + for kind in ("", "_mtp4"): + key = f"W{w}_C{c}{kind}" + vals = [report["layers"][li]["decode"].get(key) for li in layers] + vals = [v for v in vals if v] + if vals: + summ["decode_" + key] = {k: float(np.mean([v[k] for v in vals])) for k in vals[0]} + vals = [report["layers"][li]["prompt"].get(f"W{w}_C{c}") for li in layers] + vals = [v for v in vals if v] + if vals: + summ["prompt_" + f"W{w}_C{c}"] = {k: float(np.mean([v[k] for v in vals])) for k in vals[0]} + for k in ("overlap_prev_mean", "new_blocks_per_row_mean", "mtp4_union_mean"): + vals = [report["layers"][li]["decode"].get(k) for li in layers] + vals = [v for v in vals if v is not None] + if vals: + summ["decode_" + k] = float(np.mean(vals)) + summ["decode_" + k + "_by_layer"] = [round(float(v), 4) for v in vals] + for w in windows: + vals = [report["layers"][li]["decode"].get(f"W{w}_frac_in_window") for li in layers] + vals = [v for v in vals if v is not None] + if vals: + summ[f"decode_W{w}_frac_in_window"] = float(np.mean(vals)) + report["summary"] = summ + for k, v in summ.items(): + print(k, json.dumps(v) if isinstance(v, dict) else v) + if a.json: + with open(a.json, "w") as f: + json.dump(report, f, indent=1) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/exl3_len_abba.py b/tools/exl3_len_abba.py new file mode 100644 index 0000000000..fe4cafce50 --- /dev/null +++ b/tools/exl3_len_abba.py @@ -0,0 +1,97 @@ +"""Concurrent two-request HTTP bench (repeated-ID prompts) used for the EXL3 ABBA runs: PORT IDS.json OUT.json LENGTH DRAFT.""" +import concurrent.futures +import json +import statistics +import sys +import time +import urllib.request + + +def one(port, prompt, tokens=128, draft=False): + body = { + "prompt": prompt, + "max_tokens": tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + "draft": draft, + } + request = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", + json.dumps(body).encode(), + {"Content-Type": "application/json"}, + ) + started = time.perf_counter() + first = last = None + usage = {} + runtime = {} + text_events = 0 + complete = False + with urllib.request.urlopen(request, timeout=1800) as response: + for raw in response: + line = raw.decode().strip() + if line == "data: [DONE]": + complete = True + break + if not line.startswith("data:"): + continue + event = json.loads(line[5:]) + usage = event.get("usage") or usage + runtime = event.get("tensorfold") or runtime + for choice in event.get("choices", []): + piece = choice.get("text") or (choice.get("delta") or {}).get("content") or "" + if piece: + now = time.perf_counter() + first = first or now + last = now + text_events += 1 + generated = int(usage.get("completion_tokens", 0)) + span = (last - first) if first is not None and last is not None else 0.0 + return { + "complete": complete, + "prompt_tokens": usage.get("prompt_tokens"), + "tokens": generated, + "ttft_s": (first - started) if first is not None else None, + "first_token_at": first, + "last_token_at": last, + "prefill_s": runtime.get("prefill_seconds"), + "decode_tps_server": runtime.get("tokens_per_second"), + "decode_tps_client": ((generated - 1) / span) if generated > 1 and span > 0 else None, + "itl_ms_client": (1000 * span / (generated - 1)) if generated > 1 and span > 0 else None, + "token_sha256": runtime.get("token_sha256"), + "text_events": text_events, + } + + +def batch(port, prompt, tokens, draft=False): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(lambda _: one(port, prompt, tokens, draft), range(2))) + end = time.perf_counter() + total = sum(row["tokens"] for row in results) + first_times = [row["first_token_at"] for row in results if row["first_token_at"] is not None] + last_times = [row["last_token_at"] for row in results if row["last_token_at"] is not None] + decode_span = max(last_times) - min(first_times) if first_times and last_times else 0.0 + return { + "requests": results, + "batch_wall_s": end - start, + "aggregate_decode_tps": total / decode_span if decode_span > 0 else None, + "end_to_end_aggregate_tps": total / (end - start) if end > start else None, + "mean_prefill_s": statistics.mean(row["prefill_s"] for row in results if row["prefill_s"] is not None), + } + + +if __name__ == "__main__": + port = int(sys.argv[1]) + with open(sys.argv[2], encoding="utf-8") as source: + base = json.load(source) + length = int(sys.argv[4]); draft = sys.argv[5] == "1" + ids = (base * (length // len(base) + 1))[:length] + report = {"prompt_tokens": len(ids), "tokens": 128, "drafts": draft, "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + result = batch(port, ids, int(__import__("os").environ.get("TOK", "128")), draft) + report["batches"].append({"phase": phase, **result}) + print(json.dumps(report["batches"][-1]), flush=True) + with open(sys.argv[3], "w", encoding="utf-8") as target: + json.dump(report, target, indent=2) diff --git a/tools/exl3_longctx_client.py b/tools/exl3_longctx_client.py new file mode 100644 index 0000000000..e420a4c0c2 --- /dev/null +++ b/tools/exl3_longctx_client.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Long-context checks against tensorfold-native serve (EXL3 Flash Next, docs/plans/exl3-kv-memory.md). + +usage: + exl3_longctx_client.py PORT needle CTX[,CTX..] [CONC] [TEXT] + A haystack of real text (TEXT, default ~/pf/long_text.txt) cut to about CTX prompt tokens (counted by the + server's /tokenize on the rendered chat), one needle sentence at a depth, a question at the end; greedy chat + reply (thinking off). CONC requests at once, each with its own needle, depth and text offset. Every request + runs drafts off, then all again drafts on: the answer must hold the needle's value and the drafted reply must + equal the undrafted one (token_sha256). + exl3_longctx_client.py PORT cont N[,N..] GEN [IDS] + The first N real-text ids (IDS, default ~/pf/long_ids.bin), GEN greedy tokens, drafts off: prints the reply's + sha and ids (for comparing two servers' continuations).""" +import hashlib, json, os, struct, sys, threading, time, urllib.request + +port, mode = int(sys.argv[1]), sys.argv[2] +base = f"http://127.0.0.1:{port}" + + +def post(path, body, timeout=7200): + req = urllib.request.Request(base + path, json.dumps(body).encode(), {"Content-Type": "application/json"}) + return json.load(urllib.request.urlopen(req, timeout=timeout)) + + +NEEDLES = [ + ("The secret code for the blue vault is 7391-4826.", "What is the secret code for the blue vault?", "7391-4826"), + ("The password for the red door is amber-heron-55.", "What is the password for the red door?", "amber-heron-55"), + ("The launch date written on the green folder is 14 March 2031.", "What launch date is written on the green folder?", "14 March 2031"), + ("The name of the night watchman at the north gate is Oskar Lindqvist.", "What is the name of the night watchman at the north gate?", "Lindqvist"), +] +DEPTHS = [0.5, 0.25, 0.75, 0.1] + + +def messages(hay, needle, question): + content = ("Read the following document carefully. Somewhere in it is one sentence that does not belong to it.\n\n" + + hay + "\n\nQuestion: " + question + " Answer with the exact value only.") + return [{"role": "user", "content": content}] + + +def build(text, ctx, i): + needle, question, answer = NEEDLES[i % len(NEEDLES)] + depth = DEPTHS[i % len(DEPTHS)] + off = (i * 97_000) % max(1, len(text) // 2) + chars = int(ctx * 3.3) + for _ in range(4): + hay = text[off:off + chars] + if len(hay) < chars: + hay = (text[off:] + "\n" + text)[:chars] + cut = int(len(hay) * depth) + cut = hay.rfind("\n", 0, cut) + 1 or cut + doc = hay[:cut] + needle + "\n" + hay[cut:] + msgs = messages(doc, needle, question) + n = post("/v1/tokenize", {"messages": msgs, "chat_template_kwargs": {"enable_thinking": False}})["count"] + if abs(n - ctx) <= max(64, ctx // 200) and n <= ctx: + break + chars = int(chars * (ctx - 48) / n) + return msgs, answer, n + + +def chat(msgs, draft, max_tokens=32): + body = {"messages": msgs, "max_tokens": max_tokens, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}} + if not draft: + body["draft"] = False + t0 = time.time() + r = post("/v1/chat/completions", body) + r["_wall"] = time.time() - t0 + return r + + +def run_all(reqs, draft): + out = [None] * len(reqs) + + def one(k): + try: + out[k] = chat(reqs[k][0], draft) + except Exception as e: # noqa: BLE001 - a failed request is reported, not raised + out[k] = {"error": repr(e)} + + ts = [threading.Thread(target=one, args=(k,)) for k in range(len(reqs))] + t0 = time.time() + for t in ts: + t.start() + for t in ts: + t.join() + return out, time.time() - t0 + + +if mode == "needle": + ctxs = [int(x) for x in sys.argv[3].split(",")] + conc = int(sys.argv[4]) if len(sys.argv) > 4 else 1 + text = open(os.path.expanduser(sys.argv[5] if len(sys.argv) > 5 else "~/pf/long_text.txt"), encoding="utf-8", errors="replace").read() + all_ok = True + for ctx in ctxs: + reqs = [build(text, ctx, i) for i in range(conc)] + res = {} + for draft in (False, True): + out, wall = run_all(reqs, draft) + res[draft] = out + for k, r in enumerate(out): + if "error" in r: + print(f"NEEDLE ctx={ctx} req={k} drafts={'on' if draft else 'off'}: ERROR {r['error']}", flush=True) + all_ok = False + continue + tf, sp = r.get("tensorfold", {}), r.get("speculative", {}) + ans = r["choices"][0]["message"]["content"] + ok = reqs[k][1].lower() in ans.lower() + all_ok &= ok + print(f"NEEDLE ctx={ctx} req={k} prompt_tokens={r['usage']['prompt_tokens']} drafts={'on' if draft else 'off'}: " + f"{'CORRECT' if ok else 'WRONG'} answer={ans.strip()[:80]!r} want={reqs[k][1]!r} prefill {tf.get('prefill_seconds', 0):.1f} s " + f"decode {tf.get('tokens_per_second', 0):.1f} tok/s accepted {sp.get('accepted')}/{sp.get('drafted')} " + f"sha {tf.get('token_sha256', '')[:12]} wall {r['_wall']:.1f} s", flush=True) + print(f"NEEDLE ctx={ctx} conc={conc} drafts={'on' if draft else 'off'}: all requests done in {wall:.1f} s", flush=True) + for k in range(conc): + a, b = res[False][k], res[True][k] + same = "error" not in a and "error" not in b and a["tensorfold"]["token_sha256"] == b["tensorfold"]["token_sha256"] + all_ok &= same + print(f"NEEDLE ctx={ctx} req={k}: drafted {'==' if same else '!='} undrafted", flush=True) + print("NEEDLE ALL OK" if all_ok else "NEEDLE NOT ALL OK", flush=True) +elif mode == "cont": + ns = [int(x) for x in sys.argv[3].split(",")] + gen = int(sys.argv[4]) + raw = open(os.path.expanduser(sys.argv[5] if len(sys.argv) > 5 else "~/pf/long_ids.bin"), "rb").read() + ids = list(struct.unpack("<%dI" % (len(raw) // 4), raw)) + for n in ns: + body = {"prompt": ids[:n], "max_tokens": gen, "temperature": 0, "ignore_eos": True, "draft": False} + r = post("/v1/completions", body) + tf = r["tensorfold"] + print(f"CONT n={n} gen={gen}: sha {tf['token_sha256'][:16]} prefill {tf['prefill_seconds']:.1f} s decode {tf['tokens_per_second']:.1f} tok/s " + f"text {r['choices'][0]['text'][:100]!r}", flush=True) +else: + sys.exit(__doc__) diff --git a/tools/exl3_nsys_step_breakdown.py b/tools/exl3_nsys_step_breakdown.py new file mode 100644 index 0000000000..97827342e3 --- /dev/null +++ b/tools/exl3_nsys_step_breakdown.py @@ -0,0 +1,41 @@ +"""Per-decode-step kernel time by category from an nsys SQLite export of a rank (counts fn_exl3_dplan launches / 48 as steps). Edit f.""" +import sqlite3, collections, re, subprocess, os +f = "/home/chenco_adm/nsys80-r0" +if not os.path.exists(f + ".sqlite"): + subprocess.run(["nsys", "stats", "--report", "cuda_gpu_kern_sum", "-o", "/dev/null", f + ".nsys-rep"], capture_output=True) +c = sqlite3.connect(f + ".sqlite") +names = dict(c.execute("select id,value from StringIds")) +rows = c.execute("select start,end,demangledName from CUPTI_ACTIVITY_KIND_KERNEL order by start").fetchall() +dp = [r for r in rows if "fn_exl3_dplan" in names[r[2]]] +steps = len(dp) // 48 +t0, t1 = dp[0][0], dp[-1][1] +print("dplan launches", len(dp), "decode steps", steps, "window s", round((t1 - t0) / 1e9, 2), "ms/step wall", round((t1 - t0) / 1e6 / steps, 2)) +def cat(n): + if "AllReduce" in n or "AllGather" in n: return "nccl all-reduce/gather" + if "fn_exl3_dcopy" in n: return "expert miss copy (side stream)" + if "tf_exl3_linear" in n or "tf_exl3_grouped" in n or "tf_exl3_group" in n or "tf_exl3_gateup" in n or "tf_exl3_down" in n or "tf_exl3_rot" in n: return "expert/dense exl3 matmuls" + if "qmmf" in n: return "qmmf matmuls" + if "fn_gdn" in n: return "DeltaNet (gdn)" + if "fn_hc_" in n: return "hyper-connection" + if "attn" in n or "kvp_" in n or "fn_qsa" in n: return "attention / kv paging" + if "tf_route" in n or "router" in n: return "router" + if "dplan" in n: return "device planner" + return "other" +tot = collections.Counter(); cnt = collections.Counter(); other = collections.Counter() +for s, e, n in rows: + if s < t0 or e > t1: continue + nm = names[n]; k = cat(nm); tot[k] += e - s; cnt[k] += 1 + if k == "other": other[nm[:70]] += e - s +for k, v in tot.most_common(): + print("%-34s %7.2f ms/step %5d launches/step" % (k, v / 1e6 / steps, cnt[k] // steps)) +print("sum of kernel time (all streams) ms/step", round(sum(tot.values()) / 1e6 / steps, 2)) +print("top 'other':") +for k, v in other.most_common(8): print(" %6.2f ms/step %s" % (v / 1e6 / steps, k)) +# main stream busy fraction: union over all non-dcopy kernels +iv = sorted((s, e) for s, e, n in rows if t0 <= s and e <= t1 and "dcopy" not in names[n]) +busy = 0; cs, ce = iv[0] +for s, e in iv[1:]: + if s > ce: busy += ce - cs; cs, ce = s, e + else: ce = max(ce, e) +busy += ce - cs +print("busy fraction (non-copy kernels)", round(busy / (t1 - t0), 3)) diff --git a/tools/exl3_pack_sizes.py b/tools/exl3_pack_sizes.py new file mode 100644 index 0000000000..6525fdc3b8 --- /dev/null +++ b/tools/exl3_pack_sizes.py @@ -0,0 +1,26 @@ +"""Group an EXL3 pack's safetensors tensors by role and print their sizes (reads headers only). Edit P for the pack directory.""" +import json, struct, glob, collections +P = "/data/models/Qwen3.8-Flash-Next-exl3-3.05bpw" +g = collections.Counter() +def cat(n): + for k in ("ngram_embedding", "ple_", ".ple."): + if k in n and "ngram" in n: return "ple n-gram embedding shards" + if ".ple." in n: return "ple other (projections)" + if "mtp." in n: return "mtp (head)" + if "experts." in n and "shared" not in n: return "routed experts" + if "shared_expert" in n: return "shared expert" + if "hyper_connection" in n: return "hyper-connection mixes" + if "linear_attn" in n: return "linear (DeltaNet) attention" + if "self_attn" in n: return "full attention" + if "embed_tokens" in n: return "embed_tokens" + if "lm_head" in n: return "lm_head" + if "router" in n or "gate." in n: return "router" + return "other: " + n.split(".")[-2] if "." in n else "other" +for f in sorted(glob.glob(P + "/*.safetensors")): + with open(f, "rb") as h: + n = struct.unpack("","rank":0,"phase":"prefill", + "enabled":true,"replays_before":0,"replays_after":3,"source":"rank0.log:42"} +The counters must cover that request on that rank and phase. Only a positive +replay delta qualifies; enabled flags, captures, old totals, global counters +without request attribution and unrelated IDs do not. This is a collector input +schema, NOT an existing server API. The harness cannot authenticate its source; +retain the underlying logs. Missing instrumentation/evidence fails the gate. +The collector must finish writing before the HTTP phase ends. --graph-evidence +also implies --require-graph-replay. With neither option graphs stay unverified. + +Exit 0 qualifies the requested checks only; exit 1 fails them, exit 2 is bad input. +""" + +import argparse +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +import hashlib +import json +import math +from pathlib import Path +import re +import struct +import sys +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +import uuid + + +LIMITATION = "HTTP overlap alone does not prove two admitted, live server lanes or TP2 residency." +CORRELATION = [ + "Match client_request_id (X-Request-ID), server_response_id, response request ID and UTC intervals to logs.", + "Record binary/commit identity and BF16 paging, TP2, parallel=2 and sequence-memory reservation settings.", + "On BOTH ranks, correlate distinct sequence/lane IDs, prompt lengths, positions and overlapping live lifetimes.", + "Correlate /metrics running-request samples with per-rank lane traces; queued HTTP requests are insufficient.", + "Record pager misses/writebacks/rebases/overflow, per-rank VRAM peaks, host/pinned bytes and expert-pool lines.", + "For cache reuse, record snapshot save/restore IDs, prefix lengths, admission/eviction bytes and source lanes.", + "Verify snapshot branching against independent cold continuations; record graph replay and MTP rollback if enabled.", + "For graph qualification, retain request-attributed prefill/decode replay counter deltas from both ranks; enabled/capture counts alone do not qualify.", +] + + +def count(value): + return value if type(value) is int and value >= 0 else None + + +def reject_constant(value): + raise ValueError(f"invalid JSON constant: {value}") + + +def stamp(): + return {"utc": datetime.now(timezone.utc).isoformat(), "monotonic_ns": time.monotonic_ns()} + + +def load_fixture(path, length): + with Path(path).open("rb") as fixture: + size = fixture.seek(0, 2) + if size % 4: + raise ValueError(f"{path}: fixture byte length must be divisible by four") + if size // 4 < length: + raise ValueError(f"{path}: {size // 4} tokens available, {length} required") + fixture.seek(0) + raw = fixture.read(length * 4) + if len(raw) != length * 4: + raise ValueError(f"{path}: fixture changed while reading") + return list(struct.unpack(f"<{length}I", raw)), { + "path": str(path), "source_tokens": size // 4, "submitted_tokens": length, + "selected_prefix": size // 4 > length, "prompt_u32le_sha256": hashlib.sha256(raw).hexdigest(), + } + + +def inspect_response(row, payload, expected_prompt, expected_completion, require_reuse): + if not isinstance(payload, dict): + row["errors"].append("response must be a JSON object") + return + row["server_response_id"] = payload.get("id") + runtime = payload.get("tensorfold") + row["tensorfold"] = runtime + runtime = runtime if isinstance(runtime, dict) else {} + if "error" in payload: + row["errors"].append(f"server error: {json.dumps(payload['error'])}") + usage = payload.get("usage") + usage = usage if isinstance(usage, dict) else {} + row["usage"] = usage + row["prompt_tokens"] = count(usage.get("prompt_tokens")) + row["completion_tokens"] = count(usage.get("completion_tokens")) + details = usage.get("prompt_tokens_details") + cached = count(details.get("cached_tokens")) if isinstance(details, dict) else None + if cached is not None and (row["prompt_tokens"] is None or cached > row["prompt_tokens"]): + cached = None + row["cached_tokens"] = cached + row["cache_reuse_observed"] = cached is not None and cached > 0 + choices = payload.get("choices") + choice = choices[0] if isinstance(choices, list) and len(choices) == 1 else None + reason = choice.get("finish_reason") if isinstance(choice, dict) else None + row["finish_reason"] = reason + row["complete"] = ( + reason in ("stop", "length") and "error" not in payload + and payload.get("complete", True) is True and runtime.get("complete", True) is True + and payload.get("status", "completed") in ("completed", "complete") + ) + if not row["complete"]: + row["errors"].append("incomplete or invalid completion state") + if row["prompt_tokens"] != expected_prompt: + row["errors"].append(f"prompt count must be exactly {expected_prompt}, got {row['prompt_tokens']}") + if row["completion_tokens"] != expected_completion: + row["errors"].append(f"completion count must be exactly {expected_completion}, got {row['completion_tokens']}") + digest = runtime.get("token_sha256") + short = runtime.get("token_sha") + valid_hash = isinstance(digest, str) and re.fullmatch(r"[0-9a-f]{64}", digest) is not None + if short is not None: + valid_hash = (valid_hash and isinstance(short, str) + and re.fullmatch(r"(?:[0-9a-f]{12}|[0-9a-f]{64})", short) is not None + and digest.startswith(short)) + row["token_sha256"] = digest if valid_hash else None + if not valid_hash: + row["errors"].append("missing, invalid or inconsistent TensorFold token_sha256 evidence") + if require_reuse and not row["cache_reuse_observed"]: + row["errors"].append("required positive cache reuse evidence is absent") + + +def request_completion(endpoint, model, prompt, completion_tokens, timeout, draft, + run_id, label, require_reuse=False, barrier=None): + row = { + "label": label, "client_request_id": f"{run_id}-{label}", + "submitted_prompt_tokens": len(prompt), "requested_completion_tokens": completion_tokens, + "start": None, "end": None, "http_status": None, "server_response_id": None, + "response_request_id": None, "prompt_tokens": None, "completion_tokens": None, + "cached_tokens": None, "cache_reuse_observed": False, "complete": False, + "token_sha256": None, "tensorfold": None, "errors": [], + } + try: + body = json.dumps({"model": model, "prompt": prompt, "max_tokens": completion_tokens, + "temperature": 0, "seed": 1234, "ignore_eos": True, + "draft": draft, "stream": False, "n": 1}).encode("utf-8") + request = urllib.request.Request(endpoint, body, { + "Content-Type": "application/json", "Accept": "application/json", + "X-Request-ID": row["client_request_id"], + }) + if barrier is not None: + barrier.wait(timeout=min(timeout, 30)) + row["start"] = stamp() + with urllib.request.urlopen(request, timeout=timeout) as response: + row["http_status"] = response.status + row["response_request_id"] = response.headers.get("X-Request-ID") + raw = response.read() + row["end"] = stamp() + if row["http_status"] != 200: + row["errors"].append(f"unexpected HTTP status {row['http_status']}") + inspect_response(row, json.loads(raw, parse_constant=reject_constant), + len(prompt), completion_tokens, require_reuse) + except urllib.error.HTTPError as error: + row["http_status"] = error.code + row["errors"].append(f"HTTPError: {error.code} {error.reason}") + error.close() + except Exception as error: + row["errors"].append(f"{type(error).__name__}: {error}") + finally: + if row["end"] is None: + row["end"] = stamp() + row["qualified"] = not row["errors"] + return row + + +def overlap(rows): + if any(row["start"] is None or row["end"] is None for row in rows): + return {"seconds": 0, "start_utc": None, "end_utc": None, "observed": False} + latest = max(rows, key=lambda row: row["start"]["monotonic_ns"])["start"] + earliest = min(rows, key=lambda row: row["end"]["monotonic_ns"])["end"] + seconds = max(0, earliest["monotonic_ns"] - latest["monotonic_ns"]) / 1e9 + return {"seconds": seconds, "start_utc": latest["utc"] if seconds > 0 else None, + "end_utc": earliest["utc"] if seconds > 0 else None, "observed": seconds > 0} + + +def graph_evidence(path, requests): + evidence = { + "source_file": str(path) if path else None, "observations": [], "errors": [], + "scope": "externally collected per-request replay deltas; retain source logs for verification", + "qualified": False, + } + if path is None: + evidence["errors"].append("graph replay requires an external --graph-evidence JSONL file") + return evidence + request_ids = {row["client_request_id"] for row in requests} + expected = {(request_id, rank, phase) + for request_id in request_ids for rank in (0, 1) for phase in ("prefill", "decode")} + seen = set() + try: + with path.open(encoding="utf-8") as source: + for line_number, line in enumerate(source, 1): + if not line.strip(): + continue + item = json.loads(line, parse_constant=reject_constant) + if not isinstance(item, dict): + raise ValueError(f"line {line_number}: expected an observation object") + if not isinstance(item.get("client_request_id"), str): + raise ValueError(f"line {line_number}: client_request_id must be a string") + if item["client_request_id"] not in request_ids: + continue + rank, phase = item.get("rank"), item.get("phase") + if type(rank) is not int or rank not in (0, 1) or phase not in ("prefill", "decode"): + raise ValueError(f"line {line_number}: invalid TP2 rank or graph phase") + key = (item["client_request_id"], rank, phase) + if key in seen: + raise ValueError(f"line {line_number}: duplicate observation for {key}") + before, after = count(item.get("replays_before")), count(item.get("replays_after")) + origin = item.get("source") + if (item.get("enabled") is not True or before is None or after is None or after <= before + or not isinstance(origin, str) or not origin.strip()): + raise ValueError(f"line {line_number}: enabled graphs, a positive replay delta and source are required") + seen.add(key) + evidence["observations"].append({**item, "replay_delta": after - before}) + except (OSError, ValueError) as error: + evidence["errors"].append(f"{type(error).__name__}: {error}") + for request_id, rank, phase in sorted(expected - seen): + evidence["errors"].append(f"missing graph replay evidence for {request_id} rank={rank} phase={phase}") + if not requests: + evidence["errors"].append("no concurrent requests to correlate") + evidence["qualified"] = not evidence["errors"] + return evidence + + +def run(args): + prompts, fixtures = [], [] + for path in (args.fixture_a, args.fixture_b): + prompt, evidence = load_fixture(path, args.context_length) + prompts.append(prompt) + fixtures.append(evidence) + if prompts[0] == prompts[1]: + raise ValueError("fixtures must provide distinct token histories for the two requests") + prefix_length = args.warm_prefix_tokens + if prefix_length and prompts[0][:prefix_length] != prompts[1][:prefix_length]: + raise ValueError("fixtures do not share the requested warm prefix") + run_id = uuid.uuid4().hex + endpoint = args.base_url.rstrip("/") + "/v1/completions" + report = { + "schema_version": 1, "run_id": run_id, "started": stamp(), + "scope": "HTTP acceptance only; server configuration is externally verified", + "limitation": LIMITATION, "server_lane_residency_proven": False, + "snapshot_restore_proven": False, "correlation_evidence_needed": CORRELATION, + "configuration": {"context_length": args.context_length, "completion_tokens": args.completion_tokens, + "warm_prefix_tokens": prefix_length, "draft": args.draft, + "require_cache_reuse": args.require_cache_reuse or bool(prefix_length), + "require_graph_replay": args.require_graph_replay or args.graph_evidence is not None}, + "graph_replay_evidence": None, + "fixtures": fixtures, "warmup": None, "requests": [], "errors": [], + } + if prefix_length: + report["warmup"] = request_completion( + endpoint, args.model, prompts[0][:prefix_length], 1, args.timeout, args.draft, run_id, "warmup") + if not report["warmup"]["qualified"]: + report["errors"].append("shared-prefix warmup failed; concurrent phase was not submitted") + if not report["errors"]: + barrier = threading.Barrier(2) + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit( + request_completion, endpoint, args.model, prompt, args.completion_tokens, + args.timeout, args.draft, run_id, f"branch-{index}", + args.require_cache_reuse or bool(prefix_length), barrier, + ) for index, prompt in enumerate(prompts)] + report["requests"] = [future.result() for future in futures] + if any(not row["qualified"] for row in report["requests"]): + report["errors"].append("one or more concurrent requests failed qualification") + report["http_overlap"] = overlap(report["requests"]) if report["requests"] else { + "seconds": 0, "start_utc": None, "end_utc": None, "observed": False, + } + if not report["http_overlap"]["observed"]: + report["errors"].append("two HTTP request intervals did not overlap") + if report["configuration"]["require_graph_replay"]: + report["graph_replay_evidence"] = graph_evidence(args.graph_evidence, report["requests"]) + if not report["graph_replay_evidence"]["qualified"]: + report["errors"].append("required graph replay evidence did not qualify") + report["finished"] = stamp() + report["qualified"] = not report["errors"] + return report + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("fixture_a", type=Path) + parser.add_argument("fixture_b", type=Path) + parser.add_argument("--base-url", required=True, help="origin, without /v1/completions") + parser.add_argument("--model", required=True) + parser.add_argument("--context-length", required=True, type=int, help="exact prompt tokens per request") + parser.add_argument("--completion-tokens", type=int, default=256, help="exact reply tokens; ignore_eos is set") + parser.add_argument("--warm-prefix-tokens", type=int, default=0) + parser.add_argument("--require-cache-reuse", action="store_true") + parser.add_argument("--require-graph-replay", action="store_true", + help="require attributed prefill/decode replay deltas for both requests on both ranks") + parser.add_argument("--graph-evidence", type=Path, + help="external JSONL observations read after requests finish; implies --require-graph-replay") + parser.add_argument("--draft", action="store_true") + parser.add_argument("--timeout", type=float, default=7200, help="HTTP socket timeout in seconds") + parser.add_argument("--output", type=Path, help="also write the JSON receipt to this path") + args = parser.parse_args(argv) + if args.context_length <= 0 or args.completion_tokens <= 0: + parser.error("context length and completion tokens must be positive") + if not 0 <= args.warm_prefix_tokens < args.context_length: + parser.error("warm prefix must be zero or shorter than the exact context length") + if not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("timeout must be positive and finite") + url = urllib.parse.urlsplit(args.base_url) + if url.scheme not in ("http", "https") or not url.netloc or url.query or url.fragment or url.path not in ("", "/"): + parser.error("base URL must be an http(s) origin with no path, query or fragment") + print(LIMITATION, file=sys.stderr) + try: + report = run(args) + code = 0 if report["qualified"] else 1 + except (OSError, ValueError) as error: + report = {"schema_version": 1, "qualified": False, "errors": [f"{type(error).__name__}: {error}"], + "limitation": LIMITATION, "correlation_evidence_needed": CORRELATION} + code = 2 + encoded = json.dumps(report, indent=2, allow_nan=False) + "\n" + print(encoded, end="") + if args.output: + try: + args.output.write_text(encoded, encoding="utf-8") + except OSError as error: + print(f"cannot save receipt: {error}", file=sys.stderr) + return 2 + return code + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/exl3_ref_compare.py b/tools/exl3_ref_compare.py new file mode 100644 index 0000000000..dc3c35b33b --- /dev/null +++ b/tools/exl3_ref_compare.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Compare the independent reference golden (tools/exl3_ref_golden.py) with our engine's outputs. + + exl3_ref_compare.py golden.json --gen NAME=gen.log ... --tf NAME=tf.log ... + +``gen.log``: ``tf-cuda-test exl3-forward gen 48 `` (line ``GEN ids=...``). +``tf.log``: ``tf-cuda-test exl3-forward tf 48 hc`` (the per-position table: our device +top-1 and logit, the f64 oracle's top-1 and logit, our top-5 ids). Positions past the golden prompt are matched +against the golden greedy steps when the tf ids are the prompt followed by the golden continuation. +""" + +from __future__ import annotations + +import argparse +import json +import re + + +def parse_gen(path): + for line in open(path): + if line.startswith("GEN ids="): + return [int(x) for x in line.split("=", 1)[1].strip().split(",") if x] + raise ValueError(f"{path}: no GEN ids line") + + +ROW = re.compile(r"^\s*(\d+)\s+(\d+)\s+(\d+) \(\s*([-\d.]+)\)\s+(\d+) \(\s*([-\d.]+)\)\s+(\S+)\s+([-\d.]+)\s+([\d,]+)\s*$") + + +def parse_tf(path): + rows = [] + for line in open(path): + m = ROW.match(line) + if m: + p, inp, g, gl, o, ol, match, gap, five = m.groups() + rows.append({"pos": int(p), "input": int(inp), "gpu": int(g), "gpu_logit": float(gl), "oracle": int(o), + "oracle_logit": float(ol), "match": match, "oracle_gap_of_gpu_pick": float(gap), + "gpu_top5": [int(x) for x in five.split(",")]}) + return rows + + +def ref_positions(p): + """Reference rows by absolute position along prompt + greedy continuation: (input id, top5 ids, top5 logits).""" + ids = p["ids"] + p["greedy_ids"] + out = {} + for r in p["teacher_forced"]: + out[r["pos"]] = (ids[r["pos"]], r["top5_ids"], r["top5_logits"]) + n = len(p["ids"]) + for k, r in enumerate(p["greedy_steps"][1:], start=1): + out[n - 1 + k] = (ids[n - 1 + k], r["top5_ids"], r["top5_logits"]) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("golden") + ap.add_argument("--gen", action="append", default=[]) + ap.add_argument("--tf", action="append", default=[]) + ap.add_argument("--json", help="write the comparison here") + a = ap.parse_args() + g = json.load(open(a.golden)) + by = {p["name"]: p for p in g["prompts"]} + report = {"gen": {}, "tf": {}} + for spec in a.gen: + name, path = spec.split("=", 1) + p, ours = by[name], parse_gen(path) + ref = p["greedy_ids"] + n = min(len(ref), len(ours)) + first = next((i for i in range(n) if ref[i] != ours[i]), None) + agree = sum(ref[i] == ours[i] for i in range(n)) + rec = {"ref": ref, "ours": ours, "agree_prefix": first if first is not None else n, "same_positions": agree, + "n": n, "first_divergence": first} + if first is not None: + st = p["greedy_steps"][first] + rec["ref_top5_at_divergence"] = list(zip(st["top5_ids"], st["top5_logits"])) + ours_tok = ours[first] + rec["ref_logit_of_our_pick"] = (st["top5_logits"][st["top5_ids"].index(ours_tok)] + if ours_tok in st["top5_ids"] else None) + rec["ref_margin_top1_minus_our_pick"] = (None if rec["ref_logit_of_our_pick"] is None + else st["top5_logits"][0] - rec["ref_logit_of_our_pick"]) + report["gen"][name] = rec + print(f"GEN {name!r}: {agree}/{n} positions equal, identical prefix {rec['agree_prefix']}" + + ("" if first is None else f"; first divergence at step {first}: ref {ref[first]} ours {ours[first]}, " + f"ref top-5 {rec['ref_top5_at_divergence']}, ref margin over our pick " + f"{rec['ref_margin_top1_minus_our_pick']}")) + for spec in a.tf: + name, path = spec.split("=", 1) + p, rows = by[name], parse_tf(path) + refpos = ref_positions(p) + out = [] + print(f"TF {name!r} ({len(rows)} positions)") + print(" pos input ref_top1(logit) gpu_top1(logit) f64_top1(logit) ref=gpu ref=f64 gpu=f64 " + "|ref-gpu| |ref-f64| top5(ref)==top5(gpu)") + for r in rows: + if r["pos"] not in refpos: + continue + inp, ti, tv = refpos[r["pos"]] + if inp != r["input"]: + raise ValueError(f"pos {r['pos']}: input {r['input']} vs golden {inp}") + d_g = abs(tv[0] - r["gpu_logit"]) if ti[0] == r["gpu"] else None + d_o = abs(tv[0] - r["oracle_logit"]) if ti[0] == r["oracle"] else None + row = {**r, "ref_top1": ti[0], "ref_logit": tv[0], "ref_top5": ti, "ref_top5_logits": tv, + "ref_eq_gpu": ti[0] == r["gpu"], "ref_eq_oracle": ti[0] == r["oracle"], + "abs_ref_gpu": d_g, "abs_ref_oracle": d_o, "top5_set_eq": set(ti) == set(r["gpu_top5"])} + out.append(row) + f = lambda x: " - " if x is None else f"{x:.3f}" # noqa: E731 + print(f" {r['pos']:>3} {r['input']:>6} {ti[0]:>6}({tv[0]:7.3f}) {r['gpu']:>6}({r['gpu_logit']:7.3f})" + f" {r['oracle']:>6}({r['oracle_logit']:7.3f}) {'yes' if row['ref_eq_gpu'] else 'NO ':>7}" + f" {'yes' if row['ref_eq_oracle'] else 'NO ':>7} {'yes' if r['gpu'] == r['oracle'] else 'NO ':>7}" + f" {f(d_g):>8} {f(d_o):>8} {row['top5_set_eq']}") + n = len(out) + dg = [x["abs_ref_gpu"] for x in out if x["abs_ref_gpu"] is not None] + do = [x["abs_ref_oracle"] for x in out if x["abs_ref_oracle"] is not None] + summ = {"positions": n, "ref_eq_gpu": sum(x["ref_eq_gpu"] for x in out), + "ref_eq_oracle": sum(x["ref_eq_oracle"] for x in out), + "gpu_eq_oracle": sum(x["gpu"] == x["oracle"] for x in out), + "top5_set_eq": sum(x["top5_set_eq"] for x in out), + "max_abs_ref_gpu_top1_logit": max(dg) if dg else None, + "mean_abs_ref_gpu_top1_logit": sum(dg) / len(dg) if dg else None, + "max_abs_ref_oracle_top1_logit": max(do) if do else None, + "mean_abs_ref_oracle_top1_logit": sum(do) / len(do) if do else None} + print(" " + json.dumps(summ)) + report["tf"][name + "@" + path] = {"summary": summ, "rows": out} + if a.json: + json.dump(report, open(a.json, "w"), indent=1) + + +if __name__ == "__main__": + main() diff --git a/tools/exl3_ref_golden.py b/tools/exl3_ref_golden.py new file mode 100644 index 0000000000..a0a5d88843 --- /dev/null +++ b/tools/exl3_ref_golden.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +"""Independent reference run for the Qwen3.8-Flash-Next EXL3 pack: golden token ids and logits from the +Python TensorFold CUDA engine (tensorfold 0.6.x, ``families/qwen4_exp/cuda``), which no oracle of ours wrote. + +The Python engine runs an EXL3 pack on ONE GPU with every routed expert resident in VRAM (43.5 GiB here); a +16 GiB card cannot hold that and the engine has no expert offload on CUDA (``--ssd-experts`` is MLX only, +``--tp 2`` is refused for EXL3 packs). This harness changes WHERE the routed-expert trellis buffer lives and +nothing else: ``exl3.expert_table`` is re-executed from its own source with the one line that allocates +``big`` on the device replaced by page-locked host memory registered with CUDA (UVA: the kernels read the +same bytes over PCIe through the same device pointer). Every kernel, the maths and the sampling are the +reference engine's own. The capacity admission is bypassed (the weights do not "fit" by its rule), and the +engine is built directly (``decode.Engine``) without MTP, graphs or the HTTP server. + +Usage (par1-cs13, one card): + ~/bin/gpurun 1 env TORCH_EXTENSIONS_DIR=$HOME/tensorfold-c-ref/ext \ + /data/src/kt-venv/bin/python tools/exl3_ref_golden.py --tf ~/tf-final/src \ + --pack /data/models/Qwen3.8-Flash-Next-exl3-3.05bpw --out golden.json +""" + +from __future__ import annotations + +import argparse +import inspect +import json +import platform +import subprocess +import sys +import time +from pathlib import Path + +PROMPT_IDS = [760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13] +TEXTS = ["The capital of France is", "Once upon a time, in a small village by the sea"] + + +def host_mapped_factory(torch, keep): + """A CUDA-visible uint8 tensor of n bytes in page-locked host memory (cudaHostRegister, portable|mapped).""" + + import numpy as np + + cudart = torch.cuda.cudart() + + def make(n: int): + arr = np.empty((n,), dtype=np.uint8) + err = cudart.cudaHostRegister(arr.ctypes.data, n, 3) # cudaHostRegisterPortable | cudaHostRegisterMapped + if int(err) != 0: + raise RuntimeError(f"cudaHostRegister({n}) failed: {err}") + keep.append(arr) + + class _Iface: + __cuda_array_interface__ = {"shape": (n,), "typestr": "|u1", "data": (arr.ctypes.data, False), + "version": 2, "strides": None} + + t = torch.as_tensor(_Iface(), device="cuda") + if t.data_ptr() != arr.ctypes.data: + raise RuntimeError("device pointer of the registered host buffer differs from its host pointer (no UVA)") + return t + + return make + + +def patch_expert_table(exl3_mod, make_host): + src = inspect.getsource(exl3_mod.expert_table) + old = "big = torch.empty((total,), dtype=torch.uint8, device=device)" + if src.count(old) != 1: + raise RuntimeError("expert_table source changed; the one-line host-memory patch no longer applies") + src = src.replace(old, "big = _HOST_MAPPED(total)") + ns = exl3_mod.__dict__ + ns["_HOST_MAPPED"] = make_host + exec(compile(src, exl3_mod.__file__ + " [host-mapped experts]", "exec"), ns) + return src + + +def top(torch, logits_row, k=5): + row = logits_row.float() + lse = float(torch.logsumexp(row, dim=-1)) + v, i = torch.topk(row, k) + return [int(x) for x in i.cpu()], [float(x) for x in v.cpu()], lse + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--tf", required=True, help="tf-final/src (the reference tensorfold package, read only)") + ap.add_argument("--pack", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--gen", type=int, default=24) + ap.add_argument("--layers-only", action="store_true", help="load and stop (smoke test)") + ap.add_argument("--build-only", action="store_true", + help="JIT-build the reference's CUDA extensions for sm_120 without touching a GPU " + "(run with CUDA_VISIBLE_DEVICES= outside the GPU lock)") + args = ap.parse_args() + sys.path.insert(0, args.tf) + t_start = time.time() + + import torch + + if args.build_only: + torch.cuda.get_device_capability = lambda *a, **k: (12, 0) + torch.cuda.is_available = lambda: True + from tensorfold.families.qwen4_exp.cuda.engine import build_kernels as _bk + _bk(exl3=True) + print("[ref] extensions built", flush=True) + return 0 + import tensorfold + from tokenizers import Tokenizer + + from tensorfold.families.qwen4_exp.cuda import exl3 as x3mod + from tensorfold.families.qwen4_exp.cuda import decode as dec + from tensorfold.families.qwen4_exp.cuda.engine import build_kernels + from tensorfold.families.qwen4_exp.cuda.forward import commit + + torch.cuda.set_device(0) + keep: list = [] + patched = patch_expert_table(x3mod, host_mapped_factory(torch, keep)) + x3mod.PREFILL_ROWS = 64 # the n-gram staging / expert scratch rows (prompts here are <= 40 tokens) + + tok = Tokenizer.from_file(str(Path(args.pack) / "tokenizer.json")) + build_kernels(exl3=True) + t0 = time.time() + w = x3mod.load(Path(args.pack), "cuda", mtp=False, draft_vocab=None) + load_s = time.time() - t0 + mem = torch.cuda.max_memory_allocated() / 2**30 + print(f"[ref] loaded in {load_s:.0f} s, device peak {mem:.2f} GiB, host-mapped expert buffers " + f"{sum(a.nbytes for a in keep) / 2**30:.2f} GiB, centred offset {w.meta.get('centred_offset')}", flush=True) + if args.layers_only: + return 0 + e = dec.Engine(w, capacity=128, max_rows=1, prefill_rows=64, graphs=False) + + @torch.no_grad() + def teacher_forced(ids): + """Serial path: prefill ids[0], then one decode step per token; logits after every position.""" + rows = [] + dec.prefill_begin(e, ids, mtp=False) + last = dec.prefill_chunk(e, ids[:1], 0, mtp=False) + rows.append(top(torch, last[0])) + for t in ids[1:]: + lg = e.forward([t]) + rows.append(top(torch, lg[0])) + commit(w, e.st, e.buf, 1, 1) + return rows + + @torch.no_grad() + def greedy(ids, n): + """The engine's own prompt path (one prefill chunk), then n-1 greedy serial decode steps.""" + dec.prefill_begin(e, ids, mtp=False) + last = dec.prefill_chunk(e, ids, 0, mtp=False) + steps = [top(torch, last[0])] + out = [steps[0][0][0]] + while len(out) < n: + lg = e.forward([out[-1]]) + steps.append(top(torch, lg[0])) + commit(w, e.st, e.buf, 1, 1) + out.append(steps[-1][0][0]) + return out, steps + + def pack_rows(rows, ids=None): + res = [] + for p, (ti, tv, lse) in enumerate(rows): + r = {"pos": p, "top5_ids": ti, "top5_logits": tv, "logsumexp": lse, "margin12": tv[0] - tv[1]} + if ids is not None and p + 1 < len(ids): + r["next_actual"] = ids[p + 1] + res.append(r) + return res + + prompts = [{"name": "ids16", "ids": PROMPT_IDS, "text": tok.decode(PROMPT_IDS)}] + for text in TEXTS: + prompts.append({"name": text, "ids": tok.encode(text, add_special_tokens=False).ids, "text": text}) + results = [] + for p in prompts: + t1 = time.time() + tf_rows = teacher_forced(p["ids"]) + gen, steps = greedy(p["ids"], args.gen) + # the prefill path's last-row logits against the serial path's at the same position + d = abs(tf_rows[-1][1][0] - steps[0][1][0]) + results.append({**p, "teacher_forced": pack_rows(tf_rows, p["ids"]), + "teacher_forced_top1": [r[0][0] for r in tf_rows], + "greedy_ids": gen, "greedy_text": tok.decode(gen), "greedy_steps": pack_rows(steps), + "prefill_vs_serial_last_top1_logit_absdiff": d, "seconds": time.time() - t1}) + print(f"[ref] {p['name']!r}: {tok.decode(gen)!r} ({time.time() - t1:.1f} s)", flush=True) + + def sh(cmd): + try: + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=20).stdout.strip() + except Exception as exc: # noqa: BLE001 + return f"error: {exc}" + + golden = { + "what": "Independent reference: Python TensorFold CUDA engine (qwen4_exp/cuda) on the EXL3 pack; routed " + "expert trellises in host-mapped pinned memory (only allocation site changed), serial decode path, " + "no MTP, no graphs, bf16 KV, greedy argmax", + "pack": args.pack, + "command": " ".join(sys.argv), + "started_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(t_start)), + "finished_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "versions": {"tensorfold": tensorfold.__version__, "torch": torch.__version__, "torch_cuda": torch.version.cuda, + "python": platform.python_version(), "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "driver": sh("nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1"), + "nvcc": sh("nvcc --version | tail -1")}, + "load_seconds": load_s, "device_peak_gib": mem, "centred_offset": w.meta.get("centred_offset"), + "patched_expert_table_line": "big = _HOST_MAPPED(total)", + "harness_changes": ["exl3.expert_table: routed-expert trellis buffer in cudaHostRegister'd host memory (UVA)", + "exl3.PREFILL_ROWS=64; Engine(capacity=128, max_rows=1, prefill_rows=64, graphs=False); " + "load(mtp=False, draft_vocab=None); capacity admission not run"], + "notes": "teacher_forced[p] are the logits after position p (predicting ids[p+1]); greedy_steps[0] is " + "the prefill chunk's last row, later steps serial decode", + "prompts": results, + } + Path(args.out).write_text(json.dumps(golden, indent=1)) + print(f"[ref] wrote {args.out}", flush=True) + del patched + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/exl3_serve_gate.py b/tools/exl3_serve_gate.py new file mode 100755 index 0000000000..94bcfd93a8 --- /dev/null +++ b/tools/exl3_serve_gate.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +"""EXL3 serve gate: greedy 24 tokens for each golden prompt through tensorfold-native serve, drafts off +("draft": false) and on (the server's MTP), compared by the reply's token_sha256 (sha256 of the comma-joined +decimal ids, server/tokens.zig tokenSha256) against the golden's greedy ids and against each other. +usage: s9_client.py PORT GOLDEN.json [extra known continuations as name=ids,...]""" +import hashlib, json, sys, time, urllib.request + +port, golden = int(sys.argv[1]), json.load(open(sys.argv[2])) +extra = {} +for a in sys.argv[3:]: + k, v = a.split("=", 1) + extra[k] = [int(x) for x in v.split(",") if x] + +def sha(ids): + return hashlib.sha256(",".join(str(i) for i in ids).encode()).hexdigest() + +def ask(ids, draft, n=24): + body = {"prompt": ids, "max_tokens": n, "temperature": 0, "ignore_eos": True} + if not draft: + body["draft"] = False + req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions", json.dumps(body).encode(), {"Content-Type": "application/json"}) + t0 = time.time() + r = json.load(urllib.request.urlopen(req, timeout=600)) + r["_wall"] = time.time() - t0 + return r + +ok = True +report = [] +for p in golden["prompts"]: + want = p["greedy_ids"] + row = {"name": p["name"]} + for draft in (False, True): + r = ask(p["ids"], draft) + tf, sp = r["tensorfold"], r.get("speculative", {}) + got = tf["token_sha256"] + known = [k for k, v in extra.items() if sha(v) == got and k.startswith(p["name"][:4])] + key = "drafts" if draft else "no_drafts" + row[key] = {"sha": got[:12], "equals_golden": got == sha(want), "equals": known, "text": r["choices"][0]["text"], + "decode_tok_s": round(tf["tokens_per_second"], 2), "prefill_s": round(tf["prefill_seconds"], 3), + "drafted": sp.get("drafted"), "accepted": sp.get("accepted"), "acceptance_rate": sp.get("acceptance_rate"), + "tokens_per_round": sp.get("tokens_per_round"), "drafts_flag": tf["drafts"]} + row["drafted_equals_undrafted"] = row["drafts"]["sha"] == row["no_drafts"]["sha"] + row["golden_sha"] = sha(want)[:12] + ok &= row["drafted_equals_undrafted"] + report.append(row) + print(json.dumps(row, indent=1)) +print("DRAFTED==UNDRAFTED on every prompt:", ok) +# EXL3_GATE_BENCH=N: one longer greedy reply (the 16-id prompt, N tokens), drafts off then on, twice each +n = int(__import__("os").environ.get("EXL3_GATE_BENCH", "0")) +if n: + ids = golden["prompts"][0]["ids"] + for draft in (False, True, False, True): + r = ask(ids, draft, n) + tf, sp = r["tensorfold"], r.get("speculative", {}) + print(f"BENCH {n} tokens drafts={draft}: decode {tf['tokens_per_second']:.2f} tok/s ({1000 / tf['tokens_per_second']:.1f} ms/token), " + f"prefill {tf['prefill_seconds'] * 1000:.0f} ms, accepted {sp.get('accepted')}/{sp.get('drafted')}, sha {tf['token_sha256'][:12]}") diff --git a/tools/exl3_serve_tp2.sh b/tools/exl3_serve_tp2.sh new file mode 100755 index 0000000000..1c9db8bc9f --- /dev/null +++ b/tools/exl3_serve_tp2.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Launch one TP=2 rank of the EXL3 Qwen3.8-Flash-Next server with the measured-fastest settings on two RTX 5060 Ti (16 GB). +# Usage: tools/exl3_serve_tp2.sh MODEL_DIR RANK PORT MASTER_PORT [extra serve args]; start rank 1 first, then rank 0. +# FAST=0 uses the plain defaults (38 lines, drafts loaded). FAST=1 (default) drops the MTP head, trims the reserve to 512 MB and sizes the +# shared BF16 KV window for PARALLEL lanes instead of the default 8, which together let 68 lines per layer fit on a 16 GB card. +# Plain decode is faster than drafting on offloaded experts here; see docs/plans/exl3-on-upstream-final-report.md. +set -euo pipefail +model=$1 rank=$2 port=$3 master=$4 +shift 4 +bin=${TENSORFOLD_BIN:-./zig-out/native/bin/tensorfold-native} +export TENSORFOLD_CHECKPOINT_SHA256=${TENSORFOLD_CHECKPOINT_SHA256:-$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64)} +export TF_FLASHNEXT_KVPAGE=1 TF_FLASHNEXT_KVPAGE_SHARED=1 TF_FLASHNEXT_EXL3_GRAPH=1 TF_FLASHNEXT_EXL3_GRAPH_PREFILL=1 +export TF_FLASHNEXT_EXL3_GRAPH_PREFILL_AFTER=2 TF_FLASHNEXT_EXL3_GRAPH_AFTER=2 TF_FLASHNEXT_EXL3_GRAPH_ASYNC=0 +if [ "${FAST:-1}" = 1 ]; then + export TF_FLASHNEXT_KVPAGE_SHARED_LANES=${PARALLEL:-2} + if [ "${PARALLEL:-2}" = 1 ]; then + # one user: a 0.5 GiB device reserve (the supported minimum) and a 256 MB planner reserve let 80 lines fit + export TF_FLASHNEXT_EXL3_LINES=${LINES:-80} TF_FLASHNEXT_EXL3_RESERVE_MB=${RESERVE_MB:-256} TENSORFOLD_MEMORY_RESERVE_GIB=${MEMORY_RESERVE_GIB:-0.5} + else + export TF_FLASHNEXT_EXL3_LINES=${LINES:-68} TF_FLASHNEXT_EXL3_RESERVE_MB=${RESERVE_MB:-512} + fi + extra=(--no-drafts) +else + export TF_FLASHNEXT_EXL3_LINES=${LINES:-38} + extra=() +fi +CUDA_VISIBLE_DEVICES=$rank exec "$bin" serve "$model" --backend cuda --context "${CONTEXT:-32768}" --tp 2 --rank "$rank" \ + --master 127.0.0.1 --master-port "$master" --port "$port" --no-update-check --parallel "${PARALLEL:-2}" --prompt-cache-gib 0 "${extra[@]}" "$@" diff --git a/tools/exl3_single_request_bench.py b/tools/exl3_single_request_bench.py new file mode 100644 index 0000000000..20cf458464 --- /dev/null +++ b/tools/exl3_single_request_bench.py @@ -0,0 +1,47 @@ +"""Single-request prefill/decode bench over the OpenAI chat API, shared by our server and Strata: PORT.""" +import json, os, sys, time, urllib.request, uuid +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import exl3_text_abba as t + +PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8080 +TOK = int(os.environ.get("TOK", "192")) +root = "~/tensorfold-twochunk-dbg" +a = t.corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], 50000) +b = t.corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], 50000) + + +def one(prompt): + body = { + "model": os.environ.get("MODEL", "qwen3.8-flash-next-iq3_xxs"), + "messages": [{"role": "user", "content": "Run %s.\n%s" % (uuid.uuid4().hex, prompt)}], + "max_tokens": TOK, "temperature": 0, "ignore_eos": True, "stream": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": False}, + } + req = urllib.request.Request("http://127.0.0.1:%d/v1/chat/completions" % PORT, json.dumps(body).encode(), {"Content-Type": "application/json"}) + start = time.perf_counter(); first = last = None; usage = {}; pieces = 0 + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode().strip() + if line == "data: [DONE]": + break + if not line.startswith("data:"): + continue + ev = json.loads(line[5:]) + usage = ev.get("usage") or usage + for c in ev.get("choices", []): + d = c.get("delta") or {} + if d.get("content") or d.get("reasoning_content"): + now = time.perf_counter(); first = first or now; last = now; pieces += 1 + n = int(usage.get("completion_tokens", 0)) + span = (last - first) if first and last else 0 + return {"prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": n, + "ttft_s": round(first - start, 2) if first else None, + "decode_tps": round((n - 1) / span, 2) if n > 1 and span > 0 else None} + + +short = "Write a long, detailed story about a lighthouse keeper who finds a message in a bottle. Keep writing until told to stop." +for name, prompt in (("short", short), ("zig", a), ("docs", b)): + for rep in range(3): + r = one(prompt); r["prompt"] = name; r["rep"] = rep + print(json.dumps(r), flush=True) diff --git a/tools/exl3_snapshot_gate.py b/tools/exl3_snapshot_gate.py new file mode 100644 index 0000000000..eb1ef83481 --- /dev/null +++ b/tools/exl3_snapshot_gate.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Strict two-lane serving gate for BF16 paged prompt-cache restore.""" + +import argparse +import concurrent.futures +import json +import threading +import time +import urllib.request +from pathlib import Path + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("url") + parser.add_argument("--tokens", type=int, default=4099) + parser.add_argument("--decode", type=int, default=32) + parser.add_argument("--drafts", action="store_true") + parser.add_argument("--reference", type=Path) + parser.add_argument("--record-reference", action="store_true", help="record a cold baseline; use a fresh server with prompt cache disabled") + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + if args.tokens < 32 or args.decode < 1: + parser.error("require at least 32 prompt tokens and one decode token") + if args.record_reference and not args.output: + parser.error("--record-reference requires --output") + if args.record_reference and args.reference: + parser.error("--record-reference and --reference are mutually exclusive") + shared = " cache" * max(args.tokens, 1) + history = [{"role": "user", "content": shared}, {"role": "assistant", "content": "I have reviewed the shared document."}] + prompts = [history + [{"role": "user", "content": f"Original ending for lane {lane}."}] for lane in range(2)] + reference = json.loads(args.reference.read_text()) if args.reference else None + report = {"prompt_tokens": args.tokens, "decode_tokens": args.decode, "drafts": args.drafts, "phases": {}} + if reference: + for field in ("prompt_tokens", "decode_tokens", "drafts"): + if reference[field] != report[field]: + raise AssertionError(f"reference configuration mismatch: {field}") + + def pair(request_prompts): + barrier = threading.Barrier(2) + + def ask(ids): + started = time.perf_counter() + body = {"messages": ids, "max_tokens": args.decode, "temperature": 0, "ignore_eos": True, "draft": args.drafts} + request = urllib.request.Request(args.url.rstrip("/") + "/v1/chat/completions", json.dumps(body).encode(), {"Content-Type": "application/json"}) + barrier.wait(timeout=10) + with urllib.request.urlopen(request, timeout=900) as reply: + response = json.load(reply) + usage = response["usage"] + if usage["prompt_tokens"] < args.tokens or usage["completion_tokens"] != args.decode: + raise AssertionError(f"inexact token counts: {usage}") + runtime = response["tensorfold"] + digest = runtime.get("token_sha256") or runtime["token_sha"] + if len(digest) < 12: + raise AssertionError("missing output token digest") + return { + "sha": digest, + "cached_tokens": usage["prompt_tokens_details"]["cached_tokens"], + "prompt_tokens": usage["prompt_tokens"], + "elapsed_seconds": time.perf_counter() - started, + "prefill_seconds": runtime.get("prefill_seconds"), + "tokens_per_second": runtime.get("tokens_per_second"), + } + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + return list(pool.map(ask, request_prompts)) + + branches = [history + [{"role": "user", "content": f"Edited ending for branch {lane}."}] for lane in range(2)] + for phase, ids in (("cold", prompts), ("warm", prompts), ("branch", branches), ("rewind", prompts)): + results = pair(ids) + for lane, result in enumerate(results): + if args.record_reference and result["cached_tokens"] != 0: + raise AssertionError(f"reference phase {phase} lane {lane} was not cold: {result}") + if not args.record_reference and phase != "cold" and not 0 < result["cached_tokens"] < result["prompt_tokens"]: + raise AssertionError(f"{phase} lane {lane} did not restore a cached prefix: {result}") + if phase in ("warm", "rewind") and result["sha"] != report["phases"]["cold"][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from its cold continuation") + if reference and result["sha"] != reference["phases"][phase][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from reference continuation") + report["phases"][phase] = results + print(json.dumps({"phase": phase, "lanes": results}), flush=True) + if args.output: + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.record_reference: + print("PASS: recorded no-cache cold continuations for cold, warm, branch and rewind inputs") + else: + print("PASS: exact prompt/reply lengths, two submitted lanes, warm/branch/rewind cache hits and cold-baseline parity") + + +if __name__ == "__main__": + main() diff --git a/tools/exl3_text_abba.py b/tools/exl3_text_abba.py new file mode 100644 index 0000000000..5e10cf1af4 --- /dev/null +++ b/tools/exl3_text_abba.py @@ -0,0 +1,45 @@ +"""Concurrent two-request bench with two different real-text prompts built from a checkout: PORT OUT.json CHARS.""" +import concurrent.futures, glob, json, os, statistics, sys, time, urllib.request +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import exl3_len_abba as base + + +def corpus(patterns, chars): + out, total = [], 0 + for pattern in patterns: + for path in sorted(glob.glob(os.path.expanduser(pattern), recursive=True)): + try: + text = open(path, encoding="utf-8").read() + except Exception: + continue + out.append("// file: %s\n%s\n" % (os.path.basename(path), text)) + total += len(out[-1]) + if total >= chars: + return "".join(out)[:chars] + return "".join(out)[:chars] + + +def post(port, prompt, tokens): + return base.one(port, prompt, tokens, False) + + +if __name__ == "__main__": + port, out_path, chars, tokens = int(sys.argv[1]), sys.argv[2], int(sys.argv[3]), int(os.environ.get("TOK", "128")) + root = "~/tensorfold-twochunk-dbg" + a = corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], chars) + b = corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], chars) + prompts = [a, b] + report = {"chars": [len(a), len(b)], "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + rows = list(pool.map(lambda p: post(port, p, tokens), prompts)) + batch = { + "phase": phase, + "requests": rows, + "mean_prefill_s": statistics.mean(r["prefill_s"] for r in rows), + "batch_wall_s": time.perf_counter() - start, + } + report["batches"].append(batch) + print(json.dumps(batch), flush=True) + json.dump(report, open(out_path, "w"), indent=2) diff --git a/tools/tests/test_exl3_paged_acceptance.py b/tools/tests/test_exl3_paged_acceptance.py new file mode 100644 index 0000000000..f94ec94450 --- /dev/null +++ b/tools/tests/test_exl3_paged_acceptance.py @@ -0,0 +1,388 @@ +"""Model-free acceptance harness checks using a local HTTP server.""" + +from contextlib import contextmanager, redirect_stderr, redirect_stdout +import copy +import hashlib +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import importlib.util +import io +import json +from pathlib import Path +import struct +import tempfile +import threading +import unittest +from unittest.mock import patch + + +SCRIPT = Path(__file__).resolve().parents[1] / "exl3_paged_acceptance.py" +SPEC = importlib.util.spec_from_file_location("exl3_paged_acceptance", SCRIPT) +harness = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(harness) + + +def reply(body): + digest = hashlib.sha256(json.dumps(body["prompt"]).encode()).hexdigest() + return { + "id": "fake-response", "choices": [{"text": "fake output", "finish_reason": "length"}], + "usage": {"prompt_tokens": len(body["prompt"]), "completion_tokens": body["max_tokens"], + "prompt_tokens_details": {"cached_tokens": 0}}, + "tensorfold": {"token_sha256": digest, "token_sha": digest[:12]}, + } + + +@contextmanager +def server(transform=lambda body, payload: payload, status=200, raw=None, observe=None): + requests = [] + lock = threading.Lock() + simultaneous = threading.Barrier(2) + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + with lock: + requests.append({"body": body, "request_id": self.headers.get("X-Request-ID"), + "path": self.path}) + if "branch-" in self.headers.get("X-Request-ID", ""): + try: + simultaneous.wait(timeout=3) + except threading.BrokenBarrierError: + pass + payload = transform(body, reply(body)) + if observe is not None: + observe(self.headers.get("X-Request-ID"), body) + encoded = json.dumps(payload).encode() if raw is None else raw + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.send_header("X-Request-ID", "fake-http-request") + self.end_headers() + try: + self.wfile.write(encoded) + except (BrokenPipeError, ConnectionResetError): + pass + + def log_message(self, *args): + pass + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=httpd.serve_forever, kwargs={"poll_interval": 0.01}) + thread.start() + try: + yield f"http://127.0.0.1:{httpd.server_port}", requests + finally: + httpd.shutdown() + thread.join(timeout=5) + httpd.server_close() + + +class AcceptanceTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.fixture_a = Path(self.directory.name) / "a.bin" + self.fixture_b = Path(self.directory.name) / "b.bin" + self.fixture_a.write_bytes(struct.pack("<5I", 11, 12, 13, 14, 15)) + self.fixture_b.write_bytes(struct.pack("<4I", 11, 12, 90, 91)) + + def invoke(self, base, *extra): + output, errors = io.StringIO(), io.StringIO() + with redirect_stdout(output), redirect_stderr(errors): + code = harness.main([ + str(self.fixture_a), str(self.fixture_b), "--base-url", base, "--model", "fake", + "--context-length", "4", "--completion-tokens", "3", *extra, + ]) + return code, json.loads(output.getvalue()), errors.getvalue() + + def graph_collector(self, mutate=lambda rows: rows): + path = Path(self.directory.name) / "graph-observations.jsonl" + path.write_text("", encoding="utf-8") + lock = threading.Lock() + + def observe(request_id, body): + if "branch-" not in request_id: + return + rows = [{"client_request_id": request_id, "rank": rank, "phase": phase, "enabled": True, + "replays_before": 10, "replays_after": 12, "source": f"fake-rank-{rank}.log:42"} + for rank in (0, 1) for phase in ("prefill", "decode")] + with lock, path.open("a", encoding="utf-8") as output: + for row in mutate(rows): + output.write(json.dumps(row) + "\n") + + return path, observe + + def test_exact_concurrent_requests_and_receipt(self): + output = Path(self.directory.name) / "receipt.json" + with server() as (base, requests): + code, report, errors = self.invoke(base, "--output", str(output)) + self.assertEqual(code, 0) + self.assertEqual(json.loads(output.read_text()), report) + self.assertEqual(len(requests), 2) + self.assertEqual({tuple(item["body"]["prompt"]) for item in requests}, {(11, 12, 13, 14), (11, 12, 90, 91)}) + self.assertTrue(all(item["path"] == "/v1/completions" for item in requests)) + self.assertTrue(all(item["body"]["stream"] is False for item in requests)) + self.assertEqual(report["fixtures"][0]["source_tokens"], 5) + self.assertTrue(report["fixtures"][0]["selected_prefix"]) + self.assertEqual(report["fixtures"][0]["prompt_u32le_sha256"], + hashlib.sha256(struct.pack("<4I", 11, 12, 13, 14)).hexdigest()) + self.assertTrue(report["http_overlap"]["observed"]) + self.assertGreater(report["http_overlap"]["seconds"], 0) + self.assertFalse(report["server_lane_residency_proven"]) + self.assertFalse(report["snapshot_restore_proven"]) + self.assertIsNone(report["graph_replay_evidence"]) + self.assertIn("HTTP overlap alone", errors) + self.assertGreater(len(report["correlation_evidence_needed"]), 0) + for row in report["requests"]: + self.assertEqual(row["prompt_tokens"], 4) + self.assertEqual(row["completion_tokens"], 3) + self.assertTrue(row["complete"]) + self.assertEqual(row["server_response_id"], "fake-response") + self.assertEqual(row["response_request_id"], "fake-http-request") + self.assertEqual(len(row["token_sha256"]), 64) + self.assertLess(row["start"]["monotonic_ns"], row["end"]["monotonic_ns"]) + self.assertTrue(row["start"]["utc"].endswith("+00:00")) + self.assertEqual(row["tensorfold"]["token_sha256"], row["token_sha256"]) + + def test_warm_shared_prefix_precedes_concurrent_branches(self): + warmed = threading.Event() + + def reuse(body, payload): + if len(body["prompt"]) == 2: + warmed.set() + elif warmed.is_set(): + payload["usage"]["prompt_tokens_details"]["cached_tokens"] = 2 + return payload + + with server(reuse) as (base, requests): + code, report, _ = self.invoke(base, "--warm-prefix-tokens", "2") + self.assertEqual(code, 0) + self.assertEqual(len(requests), 3) + self.assertEqual(requests[0]["body"]["prompt"], [11, 12]) + self.assertEqual(requests[0]["body"]["max_tokens"], 1) + self.assertTrue(all(row["cached_tokens"] == 2 for row in report["requests"])) + self.assertLessEqual(report["warmup"]["end"]["monotonic_ns"], + min(row["start"]["monotonic_ns"] for row in report["requests"])) + + def test_warmup_failure_stops_dependent_requests(self): + def incomplete(body, payload): + payload["choices"][0]["finish_reason"] = None + return payload + + with server(incomplete) as (base, requests): + code, report, _ = self.invoke(base, "--warm-prefix-tokens", "2") + self.assertEqual(code, 1) + self.assertEqual(len(requests), 1) + self.assertFalse(report["warmup"]["complete"]) + self.assertEqual(report["requests"], []) + + def test_required_reuse_fails_when_missing_zero_or_invalid(self): + for cached in (None, 0, True, "2", 2.0, -1, 5): + with self.subTest(cached=cached): + def change(body, payload): + payload["usage"]["prompt_tokens_details"] = {"cached_tokens": cached} + return payload + + with server(change) as (base, _): + code, report, _ = self.invoke(base, "--require-cache-reuse") + self.assertEqual(code, 1) + self.assertTrue(all(not row["cache_reuse_observed"] for row in report["requests"])) + + def test_warm_phase_implicitly_requires_reuse(self): + with server() as (base, _): + code, report, _ = self.invoke(base, "--warm-prefix-tokens", "2") + self.assertEqual(code, 1) + self.assertTrue(report["configuration"]["require_cache_reuse"]) + + def test_short_cached_branches_require_prefill_and_decode_replays_on_both_ranks(self): + path, observe = self.graph_collector() + + def reuse(body, payload): + payload["usage"]["prompt_tokens_details"]["cached_tokens"] = 2 + return payload + + with server(reuse, observe=observe) as (base, _): + code, report, _ = self.invoke(base, "--warm-prefix-tokens", "2", + "--require-graph-replay", "--graph-evidence", str(path)) + self.assertEqual(code, 0) + self.assertEqual(report["configuration"]["context_length"], 4) + self.assertTrue(report["configuration"]["require_graph_replay"]) + self.assertTrue(report["graph_replay_evidence"]["qualified"]) + self.assertEqual(len(report["graph_replay_evidence"]["observations"]), 8) + self.assertTrue(all(row["replay_delta"] == 2 for row in report["graph_replay_evidence"]["observations"])) + self.assertFalse(report["snapshot_restore_proven"]) + + def test_graph_requirement_fails_without_collector_file(self): + for extra in (("--require-graph-replay",), ("--graph-evidence", str(Path(self.directory.name) / "missing"))): + with self.subTest(extra=extra): + with server() as (base, _): + code, report, _ = self.invoke(base, *extra) + self.assertEqual(code, 1) + self.assertFalse(report["graph_replay_evidence"]["qualified"]) + + def test_graph_evidence_option_implies_gate_and_unrelated_ids_do_not_satisfy_it(self): + def stale(rows): + return [{**row, "client_request_id": "previous-run-branch-0"} for row in rows] + + path, observe = self.graph_collector(stale) + with server(observe=observe) as (base, _): + code, report, _ = self.invoke(base, "--graph-evidence", str(path)) + self.assertEqual(code, 1) + self.assertTrue(report["configuration"]["require_graph_replay"]) + self.assertEqual(report["graph_replay_evidence"]["observations"], []) + + def test_graphs_enabled_captures_and_total_counts_do_not_establish_replay(self): + changes = [ + {"replays_after": 10, "captures": 100}, {"replays_after": 9}, + {"replays_before": None, "replays_after": 100}, + {"replays_before": False}, {"replays_after": "12"}, + {"enabled": False}, {"source": ""}, {"client_request_id": []}, + ] + for changed in changes: + with self.subTest(changed=changed): + path, observe = self.graph_collector(lambda rows: [{**row, **changed} for row in rows]) + with server(observe=observe) as (base, _): + code, report, _ = self.invoke(base, "--graph-evidence", str(path)) + self.assertEqual(code, 1) + self.assertFalse(report["graph_replay_evidence"]["qualified"]) + + def test_graph_evidence_missing_rank_phase_or_branch_fails(self): + filters = [lambda rows: [row for row in rows if row["rank"] == 0], + lambda rows: [row for row in rows if row["phase"] == "decode"], + lambda rows: [row for row in rows if row["client_request_id"].endswith("branch-0")]] + for select in filters: + with self.subTest(select=select): + path, observe = self.graph_collector(select) + with server(observe=observe) as (base, _): + code, report, _ = self.invoke(base, "--graph-evidence", str(path)) + self.assertEqual(code, 1) + self.assertTrue(any("missing graph replay" in error + for error in report["graph_replay_evidence"]["errors"])) + + def test_duplicate_and_malformed_graph_observations_fail(self): + path, observe = self.graph_collector(lambda rows: rows + [rows[0]]) + with server(observe=observe) as (base, _): + code, report, _ = self.invoke(base, "--graph-evidence", str(path)) + self.assertEqual(code, 1) + self.assertTrue(any("duplicate" in error for error in report["graph_replay_evidence"]["errors"])) + path.write_text('{"client_request_id":', encoding="utf-8") + with server() as (base, _): + code, report, _ = self.invoke(base, "--graph-evidence", str(path)) + self.assertEqual(code, 1) + self.assertTrue(report["graph_replay_evidence"]["errors"]) + + def test_usage_counts_must_be_exact_non_boolean_integers(self): + for field, value in (("prompt_tokens", 3), ("prompt_tokens", None), ("prompt_tokens", True), + ("prompt_tokens", 4.0), ("completion_tokens", 2), ("completion_tokens", None)): + with self.subTest(field=field, value=value): + def change(body, payload): + payload["usage"][field] = value + return payload + + with server(change) as (base, _): + code, report, _ = self.invoke(base) + self.assertEqual(code, 1) + self.assertFalse(report["qualified"]) + + def test_one_bad_branch_fails_pair(self): + def change(body, payload): + if body["prompt"][-1] == 91: + payload["usage"]["prompt_tokens"] = 3 + return payload + + with server(change) as (base, _): + code, report, _ = self.invoke(base) + self.assertEqual(code, 1) + self.assertTrue(report["requests"][0]["qualified"]) + self.assertFalse(report["requests"][1]["qualified"]) + + def test_server_errors_incomplete_and_bad_hashes_fail(self): + changes = [ + {"error": {"message": "pager overflow"}}, + {"choices": [{"finish_reason": None}]}, + {"choices": [{"finish_reason": "error"}]}, + {"choices": []}, {"choices": [None]}, {"complete": False}, {"status": "incomplete"}, + {"tensorfold": {"complete": False, "token_sha256": "a" * 64}}, + {"tensorfold": {}}, {"tensorfold": {"token_sha256": "not-a-hash"}}, + {"tensorfold": {"token_sha256": "a" * 64, "token_sha": "b" * 12}}, + ] + for changed in changes: + with self.subTest(changed=changed): + def change(body, payload): + payload.update(copy.deepcopy(changed)) + return payload + + with server(change) as (base, _): + code, report, _ = self.invoke(base) + self.assertEqual(code, 1) + self.assertTrue(all(row["errors"] for row in report["requests"])) + + def test_http_error_and_invalid_or_truncated_json_fail(self): + for status, raw in ((500, None), (200, b'{"choices":'), (200, b'[]'), + (200, b'{"tensorfold":{"seconds":NaN}}')): + with self.subTest(status=status, raw=raw): + with server(status=status, raw=raw) as (base, _): + code, report, _ = self.invoke(base) + self.assertEqual(code, 1) + self.assertTrue(all(not row["complete"] for row in report["requests"])) + self.assertTrue(all(row["end"] is not None for row in report["requests"])) + + def test_timeout_is_a_failed_receipt_with_both_request_errors(self): + release = threading.Event() + + def stalled(body, payload): + release.wait(timeout=0.1) + return payload + + with server(stalled) as (base, _): + code, report, _ = self.invoke(base, "--timeout", "0.03") + release.set() + self.assertEqual(code, 1) + self.assertTrue(all(not row["complete"] and row["errors"] for row in report["requests"])) + + def test_nonoverlapping_intervals_fail_cli_even_with_valid_responses(self): + with server() as (base, _): + with patch.object(harness, "overlap", return_value={ + "seconds": 0, "start_utc": None, "end_utc": None, "observed": False, + }): + code, report, _ = self.invoke(base) + self.assertEqual(code, 1) + self.assertTrue(all(row["qualified"] for row in report["requests"])) + self.assertIn("two HTTP request intervals did not overlap", report["errors"]) + + def test_short_misaligned_and_missing_fixtures_fail_before_http(self): + for raw in (struct.pack("<3I", 1, 2, 3), b'bad', None): + with self.subTest(raw=raw): + if raw is None: + self.fixture_b.unlink() + else: + self.fixture_b.write_bytes(raw) + with server() as (base, requests): + code, report, _ = self.invoke(base) + self.assertEqual(code, 2) + self.assertEqual(requests, []) + self.assertFalse(report["qualified"]) + + def test_different_shared_prefix_fails_before_http(self): + with server() as (base, requests): + code, report, _ = self.invoke(base, "--warm-prefix-tokens", "3") + self.assertEqual(code, 2) + self.assertEqual(requests, []) + self.assertIn("do not share", report["errors"][0]) + + def test_overlap_uses_monotonic_intervals(self): + rows = [ + {"start": {"monotonic_ns": 1, "utc": "first-start"}, + "end": {"monotonic_ns": 4, "utc": "first-end"}}, + {"start": {"monotonic_ns": 3, "utc": "second-start"}, + "end": {"monotonic_ns": 8, "utc": "second-end"}}, + ] + self.assertEqual(harness.overlap(rows), {"seconds": 1e-9, "observed": True, + "start_utc": "second-start", "end_utc": "first-end"}) + rows[1]["start"]["monotonic_ns"] = 4 + self.assertFalse(harness.overlap(rows)["observed"]) + rows[1]["start"] = None + self.assertFalse(harness.overlap(rows)["observed"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/tests/test_exl3_snapshot_gate.py b/tools/tests/test_exl3_snapshot_gate.py new file mode 100644 index 0000000000..7363c7c345 --- /dev/null +++ b/tools/tests/test_exl3_snapshot_gate.py @@ -0,0 +1,86 @@ +"""Model-free checks for the strict chat-history snapshot gate.""" + +from contextlib import redirect_stderr, redirect_stdout +import hashlib +import importlib.util +import io +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + + +SCRIPT = Path(__file__).resolve().parents[1] / "exl3_snapshot_gate.py" +SPEC = importlib.util.spec_from_file_location("exl3_snapshot_gate", SCRIPT) +gate = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(gate) + + +class Response: + def __init__(self, payload): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self): + return json.dumps(self.payload).encode() + + +class SnapshotGateTests(unittest.TestCase): + def invoke(self, args, cached): + requests = [] + + def open_request(request, timeout): + body = json.loads(request.data) + requests.append((request.full_url, body)) + digest = hashlib.sha256(json.dumps(body["messages"]).encode()).hexdigest() + return Response({ + "usage": {"prompt_tokens": 80, "completion_tokens": 2, + "prompt_tokens_details": {"cached_tokens": cached}}, + "tensorfold": {"token_sha256": digest}, + }) + + output, errors = io.StringIO(), io.StringIO() + with patch.object(gate.urllib.request, "urlopen", open_request), redirect_stdout(output), redirect_stderr(errors): + try: + gate.main(["http://example.invalid", "--tokens", "40", "--decode", "2", *args]) + code = 0 + except AssertionError as error: + errors.write(str(error)) + code = 1 + return code, output.getvalue(), errors.getvalue(), requests + + def test_records_cold_reference_then_requires_exact_chat_cache_reuse(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + receipt = Path(directory) / "receipt.json" + code, _, _, cold_requests = self.invoke(["--record-reference", "--output", str(reference)], 0) + self.assertEqual(code, 0) + self.assertEqual(len(cold_requests), 8) + self.assertTrue(all(url.endswith("/v1/chat/completions") for url, _ in cold_requests)) + self.assertTrue(all(len(body["messages"]) == 3 for _, body in cold_requests)) + code, output, _, warm_requests = self.invoke(["--reference", str(reference), "--output", str(receipt)], 40) + self.assertEqual(code, 0) + self.assertEqual(len(warm_requests), 8) + self.assertIn("PASS:", output) + saved = json.loads(receipt.read_text()) + for phase in ("warm", "branch", "rewind"): + self.assertTrue(all(0 < lane["cached_tokens"] < lane["prompt_tokens"] + for lane in saved["phases"][phase])) + + def test_missing_warm_cache_hit_fails(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + self.invoke(["--record-reference", "--output", str(reference)], 0) + code, output, errors, _ = self.invoke(["--reference", str(reference)], 0) + self.assertEqual(code, 1) + self.assertIn("did not restore a cached prefix", errors) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/zig/cuda_host/exl3_linear_reset.cpp b/tools/zig/cuda_host/exl3_linear_reset.cpp new file mode 100644 index 0000000000..5321541e5c --- /dev/null +++ b/tools/zig/cuda_host/exl3_linear_reset.cpp @@ -0,0 +1,27 @@ +#include +#include +#include + +#define __device__ +#define __forceinline__ inline +#include "../../../zig/kernels/cuda/exl3_linear_reset.cuh" + +int main() { + constexpr int guard = 16; + constexpr int capacity = 8192; + std::array storage; + for (int blocks : {1, 2, 5, 16, 64}) { + for (int count : {0, 8, 512, 7760, 1, capacity, 8}) { + storage.fill(0x5a5a5a5a); + for (int block = blocks - 1; block >= 0; --block) { + for (int thread = 127; thread >= 0; --thread) { + reset_linear_counters(storage.data() + guard, count, block, thread, blocks); + } + } + assert(std::all_of(storage.begin(), storage.begin() + guard, [](int value) { return value == 0x5a5a5a5a; })); + assert(std::all_of(storage.begin() + guard, storage.begin() + guard + count, [](int value) { return value == 0; })); + assert(std::all_of(storage.begin() + guard + count, storage.end(), [](int value) { return value == 0x5a5a5a5a; })); + } + } + reset_linear_counters(nullptr, 0, 0, 0, 1); +} diff --git a/tools/zig/exl3_prompt_ldmatrix_fixture.py b/tools/zig/exl3_prompt_ldmatrix_fixture.py new file mode 100644 index 0000000000..38925b1fe7 --- /dev/null +++ b/tools/zig/exl3_prompt_ldmatrix_fixture.py @@ -0,0 +1,550 @@ +import argparse +import array +import ctypes as ct +import hashlib +import json +import math +from pathlib import Path +import random +import statistics +import struct +import sys + + +SHAPES = ((2560, 2560), (2560, 6144), (2560, 10240), (6144, 2560)) +GUARD = 256 +CHUNK = 256 +HANDLE = ct.c_void_p +PTR = ct.c_uint64 +UINT = ct.c_uint +INT = ct.c_int + + +class Matrix(ct.Structure): + _fields_ = [(name, UINT) for name in ("k", "n", "bits", "tiles_n", "split")] + + +def emit(**values): + print(json.dumps(values, sort_keys=True), flush=True) + + +def digest(data): + return hashlib.sha256(data).hexdigest() + + +def mapping_test(): + assert ct.sizeof(Matrix) == 20 + words = 0 + negative_control = False + for tile_row in range(4): + for tile_k in range(8): + supplied = [(16 * tile_row + (lane & 15)) * 136 + 16 * tile_k + (lane >> 4) * 8 for lane in range(32)] + assert all(address % 8 == 0 for address in supplied) + for lane in range(32): + base = (16 * tile_row + lane // 4) * 136 + 16 * tile_k + 2 * (lane % 4) + scalar = [base, base + 8 * 136, base + 8, base + 8 * 136 + 8] + matrix = [supplied[8 * quadrant + lane // 4] + 2 * (lane % 4) for quadrant in range(4)] + assert matrix == scalar + assert all(0 <= address < 64 * 136 - 1 for address in matrix) + negative_control |= [matrix[0], matrix[2], matrix[1], matrix[3]] != scalar + words += 4 + assert negative_control + assert exact_values()[0] != exact_values()[1] + exact_fixture_test() + return words + + +class Driver: + def __init__(self, device): + self.lib = ct.CDLL("libcuda.so.1") + self.buffers = [] + self.modules = [] + self.events = [] + self.context = HANDLE() + self.stream = HANDLE() + signatures = { + "cuInit": [UINT], "cuDriverGetVersion": [ct.POINTER(INT)], + "cuDeviceGet": [ct.POINTER(INT), INT], "cuDeviceGetName": [HANDLE, INT, INT], + "cuDeviceGetAttribute": [ct.POINTER(INT), INT, INT], + "cuCtxCreate_v2": [ct.POINTER(HANDLE), UINT, INT], "cuCtxDestroy_v2": [HANDLE], + "cuStreamCreate": [ct.POINTER(HANDLE), UINT], "cuStreamDestroy_v2": [HANDLE], + "cuStreamSynchronize": [HANDLE], "cuModuleLoad": [ct.POINTER(HANDLE), ct.c_char_p], + "cuModuleUnload": [HANDLE], "cuModuleGetFunction": [ct.POINTER(HANDLE), HANDLE, ct.c_char_p], + "cuFuncGetAttribute": [ct.POINTER(INT), INT, HANDLE], + "cuMemAlloc_v2": [ct.POINTER(PTR), ct.c_size_t], "cuMemFree_v2": [PTR], + "cuMemcpyHtoD_v2": [PTR, HANDLE, ct.c_size_t], "cuMemcpyDtoH_v2": [HANDLE, PTR, ct.c_size_t], + "cuMemsetD8Async": [PTR, ct.c_ubyte, ct.c_size_t, HANDLE], + "cuLaunchKernel": [HANDLE, UINT, UINT, UINT, UINT, UINT, UINT, UINT, HANDLE, ct.POINTER(HANDLE), HANDLE], + "cuEventCreate": [ct.POINTER(HANDLE), UINT], "cuEventDestroy_v2": [HANDLE], + "cuEventRecord": [HANDLE, HANDLE], "cuEventSynchronize": [HANDLE], + "cuEventElapsedTime": [ct.POINTER(ct.c_float), HANDLE, HANDLE], + "cuStreamBeginCapture_v2": [HANDLE, INT], "cuStreamEndCapture": [HANDLE, ct.POINTER(HANDLE)], + "cuGraphInstantiateWithFlags": [ct.POINTER(HANDLE), HANDLE, ct.c_uint64], + "cuGraphLaunch": [HANDLE, HANDLE], "cuGraphExecDestroy": [HANDLE], "cuGraphDestroy": [HANDLE], + "cuGetErrorName": [INT, ct.POINTER(ct.c_char_p)], + } + for name, arguments in signatures.items(): + function = getattr(self.lib, name) + function.restype = INT + function.argtypes = arguments + self.call("cuInit", 0) + ordinal = INT() + self.call("cuDeviceGet", ct.byref(ordinal), device) + self.call("cuCtxCreate_v2", ct.byref(self.context), 0, ordinal) + self.call("cuStreamCreate", ct.byref(self.stream), 1) + self.device = ordinal.value + + def call(self, name, *arguments): + result = getattr(self.lib, name)(*arguments) + if result: + message = ct.c_char_p() + self.lib.cuGetErrorName(result, ct.byref(message)) + raise RuntimeError(f"{name}: {result} {message.value!r}") + + def synchronize(self): + self.call("cuStreamSynchronize", self.stream) + + def attribute(self, number): + value = INT() + self.call("cuDeviceGetAttribute", ct.byref(value), number, self.device) + return value.value + + def launch(self, function, grid, threads, arguments): + pointers = (HANDLE * len(arguments))(*(ct.cast(ct.byref(value), HANDLE) for value in arguments)) + self.call("cuLaunchKernel", function, *grid, threads, 1, 1, 0, self.stream, pointers, None) + + def event(self): + event = HANDLE() + self.call("cuEventCreate", ct.byref(event), 0) + self.events.append(event) + return event + + def close(self): + try: + if self.stream: + self.synchronize() + finally: + for event in reversed(self.events): + self.lib.cuEventDestroy_v2(event) + for module in reversed(self.modules): + self.lib.cuModuleUnload(module) + for pointer in reversed(self.buffers): + self.lib.cuMemFree_v2(pointer) + if self.stream: + self.lib.cuStreamDestroy_v2(self.stream) + if self.context: + self.lib.cuCtxDestroy_v2(self.context) + + +class Buffer: + def __init__(self, driver, size): + self.driver, self.size, self.poison = driver, size, 0xA5 + self.base = PTR() + driver.call("cuMemAlloc_v2", ct.byref(self.base), size + 2 * GUARD) + driver.buffers.append(self.base) + self.ptr = self.base.value + GUARD + self.reset() + + def reset(self, poison=0xA5): + self.poison = poison + self.driver.call("cuMemsetD8Async", self.base, poison, self.size + 2 * GUARD, self.driver.stream) + + def upload(self, data): + assert len(data) == self.size + self.driver.synchronize() + host = ct.create_string_buffer(data) + self.driver.call("cuMemcpyHtoD_v2", self.ptr, host, self.size) + self.expected_hash = digest(data) + + def read(self, live=None): + host = ct.create_string_buffer(self.size + 2 * GUARD) + self.driver.call("cuMemcpyDtoH_v2", host, self.base, len(host)) + data = host.raw + live = self.size if live is None else live + assert data[:GUARD] == bytes([self.poison]) * GUARD, "prefix guard" + assert data[GUARD + live:] == bytes([self.poison]) * (self.size - live + GUARD), "suffix guard" + return data[GUARD:GUARD + live] + + +class Image: + def __init__(self, driver, path, label): + self.driver, self.label = driver, label + module = HANDLE() + driver.call("cuModuleLoad", ct.byref(module), str(path.resolve()).encode()) + driver.modules.append(module) + self.functions = {} + resources = {} + for name in ("tf_exl3_rotate_h", "tf_exl3_prompt_2"): + function = HANDLE() + driver.call("cuModuleGetFunction", ct.byref(function), module, name.encode()) + self.functions[name] = function + resource = {} + for key, number in (("max_threads", 0), ("shared_bytes_driver", 1), ("local_bytes", 3), ("registers", 4), ("binary_version", 6)): + value = INT() + driver.call("cuFuncGetAttribute", ct.byref(value), number, function) + resource[key] = value.value + assert resource["max_threads"] >= (256 if name.endswith("_2") else 32) + resources[name] = resource + emit(image=label, sha256=digest(path.read_bytes()), resources=resources) + + def rotate(self, case, count, offset=0): + self.driver.launch(self.functions["tf_exl3_rotate_h"], (case.k // 128, count, 1), 32, + [PTR(case.input.ptr + offset * case.stride * 2), UINT(case.stride), PTR(case.suh.ptr), PTR(case.rotated[self.label].ptr), UINT(case.k)]) + + def prompt(self, case, count, f32, bank=0, offset=0): + self.driver.launch(self.functions["tf_exl3_prompt_2"], (case.n // 128, (count + 63) // 64, 1), 256, + [PTR(case.rotated[self.label].ptr), UINT(count), PTR(case.banks[bank].ptr), PTR(case.svh.ptr), + Matrix(case.k, case.n, case.bits, case.n // 16, 1), PTR(case.output[self.label].ptr + offset * case.n * (4 if f32 else 2)), UINT(f32)]) + + +def half_bytes(values): + return b"".join(struct.pack("> 16) & 1)) >> 16) & 0xFFFF for word in words)).tobytes() + + +def same(actual, expected, label): + assert len(actual) == len(expected), f"{label}: length mismatch" + if actual != expected: + first = next(index for index, pair in enumerate(zip(actual, expected)) if pair[0] != pair[1]) + raise AssertionError(f"{label}: first differing byte {first}") + + +def exact_values(): + scale = struct.unpack("> shift) & 255 for shift in (0, 8, 16, 24)) + values.append(struct.unpack(" CHUNK and prefix is not None: + same(data[:len(prefix)], prefix, "256-row seam") + for bank in range(1, self.bank_count): + for f32 in (True, False): + self.compare(CHUNK, f32, bank, bank % 2) + random_hashes = dict(self.bank_hash) + self.exact_check() + self.nonfinite_check() + self.bank_hash = random_hashes + self.restore() + emit(shape=[self.k, self.n], byte_guard_repeat_checks="PASS", banks_checked=self.bank_count, fp64_exact_fixture="PASS", nonfinite_classes="PASS") + + def nonfinite_check(self): + def transform(values): + values = list(values) + for step in (1, 2, 4, 8, 16, 32, 64): + for base in range(0, 128, 2 * step): + for offset in range(step): + lower, upper = base + offset, base + offset + step + left, right = values[lower], values[upper] + values[lower], values[upper] = left + right, left - right + return values + def category(value): + return "nan" if math.isnan(value) else "positive_inf" if value == math.inf else "negative_inf" if value == -math.inf else "finite" + tile_bytes = 32 * self.bits + self.banks[0].upload((bytes(tile_bytes) + bytes([255]) * tile_bytes) * (self.n // 32) * (self.k // 16)) + self.suh.upload(half_bytes([1.0] * self.k)) + self.svh.upload(half_bytes([1.0] * self.n)) + source = array.array("H", [0]) * (self.rows * self.stride) + families = ((0x7FC0,), (0x7F80,), (0xFF80,), (0x7F80, 0xFF80)) + expected = [] + weights = exact_values() + for row, words in enumerate(families): + for column, word in enumerate(words): + source[row * self.stride + column] = word + values = [struct.unpack("= 2 and options.samples >= 1 and options.iterations >= 1 and options.warmup >= 0 + shapes = [tuple(map(int, value.split(","))) for value in options.shape] if options.shape else SHAPES + assert not options.profile_only or options.shape and len(shapes) == 1, "profile-only requires exactly one --shape" + assert all(len(shape) in (2, 3) and shape[:2] in SHAPES and (len(shape) == 2 or 1 <= shape[2] <= 8) for shape in shapes) + driver = Driver(options.device) + try: + name = ct.create_string_buffer(256) + version = INT() + driver.call("cuDeviceGetName", name, len(name), driver.device) + driver.call("cuDriverGetVersion", ct.byref(version)) + l2 = driver.attribute(38) + assert l2 > 0 + emit(device=options.device, gpu=name.value.decode(), cuda_driver=version.value, l2_bytes=l2, + python=sys.version, seed=20261022, submission=options.submission, profile_only=options.profile_only, + runtime_flag_match="must be established by external build receipts") + options.seed = 20261022 + images = [Image(driver, options.baseline, "baseline"), Image(driver, options.candidate, "candidate")] + for shape in shapes: + first_buffer = len(driver.buffers) + case = Case(driver, images, shape, options, l2) + if options.profile_only: + case.prepare() + image = images[0 if options.profile_only == "baseline" else 1] + image.prompt(case, CHUNK, options.profile_f32) + driver.synchronize() + case.output[image.label].read(CHUNK * case.n * (4 if options.profile_f32 else 2)) + emit(profile_only=image.label, rows=CHUNK, shape=[case.k, case.n], bits=case.bits, + cache="uncontrolled_after_upload", byte_comparison_run=False, prompt_launches=1) + else: + case.checks(options) + if not options.checks_only: + case.timing(options) + case.immutable() + driver.synchronize() + for pointer in reversed(driver.buffers[first_buffer:]): + driver.call("cuMemFree_v2", pointer) + del driver.buffers[first_buffer:] + emit(result="PROFILE_ONLY" if options.profile_only else "PASS", shapes=len(shapes), + qualification="microfixture only; serving performance unqualified") + finally: + driver.close() + + +if __name__ == "__main__": + main() diff --git a/tools/zig/test_cuda_sources.py b/tools/zig/test_cuda_sources.py index 53e4adcab8..562c396ce4 100644 --- a/tools/zig/test_cuda_sources.py +++ b/tools/zig/test_cuda_sources.py @@ -77,6 +77,15 @@ def test_real_exl3_layout(self): result = self.parse(check.HEADERS / "exl3_layout.cpp") self.assertEqual(result.returncode, 0, result.stderr) + def test_exl3_linear_reset_covers_dirty_reused_counters_without_touching_guards(self): + executable = self.root / "linear_reset" + result = subprocess.run([self.compiler, "-std=c++20", "-fsanitize=address,undefined", + str(check.HEADERS / "exl3_linear_reset.cpp"), "-o", str(executable)], + capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + result = subprocess.run([str(executable)], capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + def test_layout_rejects_wrong_coordinate(self): header = (check.ROOT / "zig/kernels/cuda/exl3_decode.cuh").read_text() bad = header.replace("return (lane >> 2) + 8 * (j >> 2);", "return (lane >> 2);", 1) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6dfec448ca..34c38b7852 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -398,7 +398,17 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st b.step("test-api-boundaries", "Runtime thread initialization failure regression").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"Failure boundaries:"} })).step); b.step("test-review-api", "Ownership rollback: ownership and final batching regressions").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"Ownership rollback:"} })).step); const check_native = nativeServer(b, cuda_target, .debug, check_cuda, check_mods.core, check_mods.lanes, check_mods.nemotron, check_mods.heat, check_mods.qwen35, check_mods.qwen_vision, check_mods.tokenizer, check_mods.living, build_options, false, null); - b.step("check-native-cuda-host", "Compile the real native CUDA server without kernels or GPU execution").dependOn(&check_native.step); + const native_check = b.step("check-native-cuda-host", "Compile the real native CUDA server without kernels or GPU execution"); + native_check.dependOn(&check_native.step); + const prefill_body = b.addObject(.{ .name = "exl3-prefill-body", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/exl3_prefill_body_compile.zig"), + .target = cuda_target, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "flashnext", .module = flashnextModule(b, cuda_target, .debug, check_cuda, check_mods.core, check_mods.lanes, check_mods.qwen_vision) }}, + }) }); + _ = prefill_body.getEmittedBin(); + native_check.dependOn(&prefill_body.step); const offer_tests = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_vision_offer.zig"), .target = host, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "nemotron", .module = mods.nemotron }, .{ .name = "qwen35", .module = mods.qwen35 } } }), .filters = if (video_filter) |f| &.{f} else &.{} })); offer_tests.has_side_effects = true; @@ -463,6 +473,21 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them + // the Flash Next family's EXL3 hooks alone (every test named "exl3: ..."), without the rest of the host suite + b.step("test-flashnext-exl3-host", "Flash Next EXL3 admission, store, sizing and loader host tests").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision), .filters = &.{"exl3: "} })).step); + const trace_writer_build = b.addSystemCommand(&.{ "c++", "-std=c++17", "-O1" }); + trace_writer_build.addFileArg(b.path("zig/src/families/flashnext/exl3_dplan_trace_kernel_test.cpp")); + trace_writer_build.addArg("-o"); + const trace_writer_binary = trace_writer_build.addOutputFileArg("dplan-trace-writer"); + trace_writer_build.addFileInput(b.path("zig/kernels/cuda/flashnext_exl3_dplan.cu")); + const trace_writer_run = std.Build.Step.Run.create(b, "dplan-trace-writer"); + trace_writer_run.addFileArg(trace_writer_binary); + const trace_fixture = trace_writer_run.addOutputFileArg("demand-v4.exdpt"); + const checkpoint_fixture = trace_writer_run.addOutputFileArg("demand-v4-checkpoints.exdpt"); + const trace_writer_tests = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/exl3_dplan_trace_writer_test.zig"), .target = host, .optimize = .debug }); + trace_writer_tests.addAnonymousImport("dplan_trace_fixture", .{ .root_source_file = trace_fixture }); + trace_writer_tests.addAnonymousImport("dplan_checkpoint_fixture", .{ .root_source_file = checkpoint_fixture }); + b.step("test-flashnext-exl3-trace-writer-host", "CPU-only CUDA planner writer parity and exact Zig replay; requires c++").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = trace_writer_tests })).step); for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.qwen_vision, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const selected = b.option([]const u8, "modelopt-host-filter", "One portable ModelOpt regression") orelse "modelopt host:"; const run = b.addRunArtifact(b.addTest(.{ .root_module = mods.nemotron, .filters = &.{selected} })); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 25315ca223..e8aaf28705 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -106,6 +106,12 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3", .flags = &.{"-O3"} }, .{ .name = "drama3_fp8", .flags = &.{"-O3"} }, .{ .name = "drama3_codec", .flags = &.{"-O3"} }, + .{ .name = "flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) + .{ .name = "flashnext_exl3_tp", .flags = glue }, // ours: Flash Next's EXL3 expert-parallel masks and placement + .{ .name = "flashnext_exl3_dplan", .flags = &.{"-O3"} }, // ours: Flash Next's EXL3 pool planned on the device (decode graphs) + .{ .name = "flashnext_exl3_kv", .flags = glue }, // ours: Flash Next attention over an int8 KV cache (K10 prototype) + .{ .name = "flashnext_exl3_kvpage", .flags = glue }, // ours: Flash Next K/V paged to host memory, a device block cache + .{ .name = "flashnext_exl3_prefill", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // ours: Flash Next EXL3 prefill experts, 64-member tiles any width }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/exl3_candidate.cu b/zig/kernels/cuda/exl3_candidate.cu index a288cb5b45..fc04dbb10e 100644 --- a/zig/kernels/cuda/exl3_candidate.cu +++ b/zig/kernels/cuda/exl3_candidate.cu @@ -123,9 +123,22 @@ __device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned ro const uint32_t b[2][2] = {{tf_exl3_pair(w[0], w[1]), tf_exl3_pair(w[2], w[3])}, {tf_exl3_pair(w[4], w[5]), tf_exl3_pair(w[6], w[7])}}; #pragma unroll for (int i = 0; i < 4; ++i) { - const __half* x = xs + (16 * i + g8) * TF_EXL3_XS + 16 * r + 2 * t4; - const uint32_t a[4] = {*reinterpret_cast(x), *reinterpret_cast(x + 8 * TF_EXL3_XS), - *reinterpret_cast(x + 8), *reinterpret_cast(x + 8 * TF_EXL3_XS + 8)}; + uint32_t a[4]; +#if __CUDA_ARCH__ >= 1200 + if constexpr (C == 2) { + const __half* address = xs + (16 * i + (lane & 15)) * TF_EXL3_XS + 16 * r + (lane >> 4) * 8; + const unsigned shared_address = (unsigned)__cvta_generic_to_shared(address); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(shared_address)); + } else +#endif + { + const __half* x = xs + (16 * i + g8) * TF_EXL3_XS + 16 * r + 2 * t4; + a[0] = *reinterpret_cast(x); + a[1] = *reinterpret_cast(x + 8 * TF_EXL3_XS); + a[2] = *reinterpret_cast(x + 8); + a[3] = *reinterpret_cast(x + 8 * TF_EXL3_XS + 8); + } tf_exl3_mma(acc[i][0], a, b[0][0], b[0][1]); tf_exl3_mma(acc[i][1], a, b[1][0], b[1][1]); } diff --git a/zig/kernels/cuda/exl3_linear.cu b/zig/kernels/cuda/exl3_linear.cu index 6b805e4390..ef045b213c 100644 --- a/zig/kernels/cuda/exl3_linear.cu +++ b/zig/kernels/cuda/exl3_linear.cu @@ -4,6 +4,7 @@ #include #include "exl3_decode.cuh" +#include "exl3_linear_reset.cuh" namespace tf_exl3 { @@ -31,8 +32,14 @@ struct RotInArgs { half* xh; // fp16 [m, k] out int x_dtype; int m, k; + int* counters; + int counter_count; }; +static_assert(sizeof(RotInArgs) == 56); +static_assert(__builtin_offsetof(RotInArgs, counters) == 40); +static_assert(__builtin_offsetof(RotInArgs, counter_count) == 48); + __device__ __forceinline__ void load4(const void* p, int dtype, size_t i, float (&v)[4]) { if (dtype == F32) { const float4 u = *reinterpret_cast(static_cast(p) + i); @@ -102,6 +109,7 @@ __device__ __forceinline__ void step_lane_words(const uint32_t (&raw)[8 * lane_w // xh = fp16(((x * suh) @ H) / sqrt(128)): one warp per 128 columns, one block of 4 warps per 512. extern "C" __global__ void __launch_bounds__(128) tf_exl3_rot_in(RotInArgs a) { + if (blockIdx.y == 0) reset_linear_counters(a.counters, a.counter_count, blockIdx.x, threadIdx.x, gridDim.x); const int blk = blockIdx.x * 4 + (threadIdx.x >> 5), row = blockIdx.y, lane = threadIdx.x & 31; if (blk * 128 >= a.k) return; const int k = blk * 128 + 4 * lane; diff --git a/zig/kernels/cuda/exl3_linear_reset.cuh b/zig/kernels/cuda/exl3_linear_reset.cuh new file mode 100644 index 0000000000..fc17103300 --- /dev/null +++ b/zig/kernels/cuda/exl3_linear_reset.cuh @@ -0,0 +1,5 @@ +#pragma once + +__device__ __forceinline__ void reset_linear_counters(int* counters, int count, int block, int thread, int blocks) { + for (int index = block * 128 + thread; index < count; index += blocks * 128) counters[index] = 0; +} diff --git a/zig/kernels/cuda/flashnext_exl3.cu b/zig/kernels/cuda/flashnext_exl3.cu new file mode 100644 index 0000000000..1fef7e4643 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3.cu @@ -0,0 +1,17 @@ +// The Flash Next family's EXL3 glue (cuda_exl3_pass.zig): an fp16 row as two bf16 rows whose sum is it exactly. + +#include +#include + +// hi = bf16(x), lo = bf16(x - hi): an fp16 has 11 significant bits, hi keeps 8 and the remainder has at most 3, so +// lo holds it exactly (bf16 has fp32's exponent range) and hi + lo == x for every finite fp16. Two bf16 products +// then give x . w with exact products and fp32 sums, where one bf16 row would round x first. +extern "C" __global__ void fn_exl3_split_f16(const __half* __restrict__ x, __nv_bfloat16* __restrict__ hi, + __nv_bfloat16* __restrict__ lo, long long n) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + const float v = __half2float(x[i]); + const __nv_bfloat16 h = __float2bfloat16_rn(v); + hi[i] = h; + lo[i] = __float2bfloat16_rn(v - __bfloat162float(h)); +} diff --git a/zig/kernels/cuda/flashnext_exl3_dplan.cu b/zig/kernels/cuda/flashnext_exl3_dplan.cu new file mode 100644 index 0000000000..ff506d0124 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_dplan.cu @@ -0,0 +1,451 @@ +// The Flash Next family's EXL3 device-side pool plan (cuda_exl3_dplan.zig): one layer's offload round planned on the +// device, so a decode step needs no host wait between the router and the experts and runs as ONE captured graph. +// +// The rules are the host planner's (exl3_round.zig `plan` over exl3_offload.zig's Vram, one pass, no prefetch, no +// double-buffer protection): the round's routed picks raise the layer's lifetime counts; held experts are touched +// (stamp = now); misses, hottest first (count desc, appearances desc, id asc), each take the first empty line, else the +// coldest resident line not used by this round (count asc, stamp asc, higher line). A miss's three table entries point +// into its line, the evicted expert's are zeroed, and its copy (host bank -> line) joins the work list that +// fn_exl3_dcopy runs next. Every line must fit: the caller only plans rounds of at most `lines` distinct experts. +// +// State, per layer (u32 words at `st`): line_id[L] (expert or -1), stamp[L], line_of[512] (line or -1), counts[512], +// then now, rounds, hit_slots, miss_slots, copies, evictions, overflow, in-place reads (more misses than lines +// free of the round's experts: read from the mapped bank in place for that round). + +#include + +#define DP_EXPERTS 512 +#define DP_SHARED 512 +#define DP_MAXU 256 +#define DP_RECENT_WINDOW 256 +#define DP_RECENT_WORDS_PER_ROW (DP_EXPERTS / 32) +#define DP_RECENT_WORDS (DP_RECENT_WINDOW * DP_RECENT_WORDS_PER_ROW) + +struct DpMisc { + unsigned now, rounds, hit_slots, miss_slots, copies, evictions, overflow, inplace, recent_cursor, recent_filled; +}; + +struct DpTraceHeader { + unsigned magic, version, pool_index, layer, lines, max_rows, seed_words, record_words; + unsigned capacity, recorded, dropped_lo, dropped_hi, discontinuity, invalid_launch, checkpoints, dropped_checkpoints; +}; + +static_assert(sizeof(DpTraceHeader) == 64, "trace header ABI"); +static_assert(sizeof(DpMisc) == 40, "planner misc ABI"); + +static constexpr unsigned DP_TRACE_WORDS = 1024 * 1024 / 4; + +__device__ __forceinline__ bool dp_trace_valid(DpTraceHeader* header, int lines) { + const unsigned long long state_words = 2ull * static_cast(lines) + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10; + const unsigned long long record_words = 15ull + 11ull * header->max_rows; + const unsigned long long base = 16 + 2 * state_words; + const unsigned long long checkpoint_words = 4 + 2 * state_words; + if (lines <= 0 || header->magic != 0x44505431 || header->version != 4 || + header->lines != static_cast(lines) || header->max_rows == 0 || + base >= DP_TRACE_WORDS || record_words > DP_TRACE_WORDS - base || + header->seed_words != state_words || header->record_words != record_words || + header->capacity != (DP_TRACE_WORDS - base) / record_words || header->recorded > header->capacity || + header->discontinuity > 1 || header->invalid_launch > 1 || + header->checkpoints > (DP_TRACE_WORDS - base) / checkpoint_words || + base + static_cast(header->recorded) * record_words > + DP_TRACE_WORDS - static_cast(header->checkpoints) * checkpoint_words) { + header->invalid_launch = 1; + return false; + } + return true; +} + +extern "C" __global__ void fn_exl3_dcheckpoint(const unsigned* st, int lines, unsigned* trace, int transfer) { + if (trace == nullptr || threadIdx.x != 0 || blockIdx.x != 0) return; + auto* header = reinterpret_cast(trace); + if (!dp_trace_valid(header, lines)) return; + if (transfer != 0 && transfer != 1) { + header->invalid_launch = 1; + return; + } + unsigned* seed = trace + 16; + unsigned* terminal = seed + header->seed_words; + const bool initialized = header->recorded != 0 || header->checkpoints != 0; + if (transfer == 0) { + if (initialized) + for (unsigned index = 0; index < header->seed_words; ++index) + if (terminal[index] != st[index]) header->discontinuity = 1; + return; + } + const unsigned stride = 4 + 2 * header->seed_words; + const unsigned end = 16 + 2 * header->seed_words + header->recorded * header->record_words; + const unsigned start = DP_TRACE_WORDS - header->checkpoints * stride; + if (start - end < stride || header->dropped_lo != 0 || header->dropped_hi != 0 || header->dropped_checkpoints != 0) { + if (header->dropped_checkpoints != 0xffffffffu) ++header->dropped_checkpoints; + return; + } + unsigned* checkpoint = trace + start - stride; + checkpoint[0] = header->recorded; + checkpoint[1] = 1; + checkpoint[2] = header->seed_words; + checkpoint[3] = header->checkpoints; + for (unsigned index = 0; index < header->seed_words; ++index) { + if (!initialized) seed[index] = terminal[index] = st[index]; + checkpoint[4 + index] = terminal[index]; + checkpoint[4 + header->seed_words + index] = st[index]; + terminal[index] = st[index]; + } + ++header->checkpoints; +} + +__device__ __forceinline__ unsigned* dp_trace_begin(const int* picks, int n, unsigned* trace, + const unsigned* st, int recent_enabled, int lines) { + if (trace == nullptr) return nullptr; + DpTraceHeader* header = reinterpret_cast(trace); + if (!dp_trace_valid(header, lines)) return nullptr; + if (n <= 0 || n % 11 != 0 || static_cast(n / 11) > header->max_rows) { + header->invalid_launch = 1; + return nullptr; + } + const unsigned end = 16 + 2 * header->seed_words + header->recorded * header->record_words; + const unsigned start = DP_TRACE_WORDS - header->checkpoints * (4 + 2 * header->seed_words); + if (start - end < header->record_words || header->dropped_checkpoints != 0) { + unsigned long long dropped = static_cast(header->dropped_hi) << 32 | header->dropped_lo; + if (dropped != ~0ull) ++dropped; + header->dropped_lo = static_cast(dropped); + header->dropped_hi = static_cast(dropped >> 32); + return nullptr; + } + unsigned* seed = trace + 16; + unsigned* terminal = seed + header->seed_words; + for (unsigned index = 0; index < header->seed_words; ++index) { + if (header->recorded == 0 && header->checkpoints == 0) seed[index] = st[index]; + else if (terminal[index] != st[index]) header->discontinuity = 1; + } + unsigned* round = terminal + header->seed_words + header->recorded * header->record_words; + round[0] = static_cast(n / 11); + round[1] = 0; + round[2] = 0; + round[3] = static_cast(recent_enabled); + unsigned seen[DP_EXPERTS / 32] = {}; + unsigned* raw = round + 15; + for (int pick = 0; pick < n; ++pick) { + const int expert = picks[pick]; + raw[pick] = static_cast(expert); + if (expert < 0 || expert >= DP_SHARED) continue; + round[1] += 1; + const unsigned bit = 1u << (expert & 31); + if ((seen[expert >> 5] & bit) == 0) round[2] += 1; + seen[expert >> 5] |= bit; + } + return round; +} + +__device__ __forceinline__ void dp_trace_end(unsigned* trace, unsigned* round, const unsigned* st, + const DpMisc* misc, const int* picks, int n, + const int* ovf, int n_ovf) { + if (round == nullptr) return; + DpTraceHeader* header = reinterpret_cast(trace); + const unsigned* counters = reinterpret_cast(misc); + for (unsigned index = 0; index < 10; ++index) round[4 + index] = counters[index]; + unsigned mapped_experts[DP_EXPERTS / 32] = {}; + for (int mapped = 0; mapped < n_ovf; ++mapped) { + const int expert = ovf[mapped]; + mapped_experts[expert >> 5] |= 1u << (expert & 31); + } + unsigned mapped_picks = 0; + for (int pick = 0; pick < n; ++pick) { + const int expert = picks[pick]; + if (expert >= 0 && expert < DP_SHARED && (mapped_experts[expert >> 5] & (1u << (expert & 31)))) + ++mapped_picks; + } + round[14] = mapped_picks; + unsigned* terminal = trace + 16 + header->seed_words; + for (unsigned index = 0; index < header->seed_words; ++index) terminal[index] = st[index]; + header->recorded += 1; +} + +__device__ __forceinline__ bool dp_colder(unsigned ca, unsigned sa, int ia, unsigned cb, unsigned sb, int ib) { + if (ca != cb) return ca < cb; + if (sa != sb) return sa < sb; + return ia > ib; +} + +// One block of 256 threads. picks [n] int32 (-1 skipped, 512 the shared expert, masked to the rank's ids at tp=2). +// tab: three int64 [E] pointer tables then (k2tab) three int32 [E] width tables (the pool's, exl3_pool.Tables). +// src: per routed expert three host-bank device addresses, lens: their byte counts, k2: their half-bit widths. +// work: per copy three {src, dst, len} triples (int64), work_n: the copy count. +extern "C" __global__ void __launch_bounds__(256) fn_exl3_dplan( + const int* __restrict__ picks, int n, unsigned* __restrict__ st, int lines, int recent_enabled, long long* __restrict__ tab0, + long long* __restrict__ tab1, long long* __restrict__ tab2, int* __restrict__ k2t0, int* __restrict__ k2t1, + int* __restrict__ k2t2, const long long* __restrict__ src, const long long* __restrict__ lens, + const int* __restrict__ k2, unsigned long long line_base, unsigned long long line_bytes, unsigned long long off0, + unsigned long long off1, unsigned long long off2, long long* __restrict__ work, int* __restrict__ work_n, + int* __restrict__ hit_picks, int* __restrict__ miss_picks, int* __restrict__ ovf, int* __restrict__ ovf_n, + unsigned* __restrict__ trace) { + int* line_id = reinterpret_cast(st); + unsigned* stamp = st + lines; + int* line_of = reinterpret_cast(st + 2 * lines); + unsigned* counts = st + 2 * lines + DP_EXPERTS; + unsigned* recent_counts = counts + DP_EXPERTS; + unsigned* recent_rows = recent_counts + DP_EXPERTS; + DpMisc* misc = reinterpret_cast(recent_rows + DP_RECENT_WORDS); + const unsigned* score = recent_enabled ? recent_counts : counts; + + __shared__ unsigned short uniq[DP_MAXU]; + __shared__ unsigned short app[DP_MAXU]; + __shared__ unsigned short miss[DP_MAXU]; + __shared__ unsigned keep[DP_EXPERTS / 32]; + __shared__ int n_uniq, n_miss, n_copy, n_ovf; + __shared__ unsigned now; + __shared__ unsigned rkey_c[256], rkey_s[256]; + __shared__ int rkey_i[256]; + const int t = threadIdx.x; + unsigned* trace_round = nullptr; + + for (int word = t; word < DP_EXPERTS / 32; word += blockDim.x) keep[word] = 0; + __syncthreads(); + if (t == 0) { + trace_round = dp_trace_begin(picks, n, trace, st, recent_enabled, lines); + now = ++misc->now; + misc->rounds += 1; + if (recent_enabled) { + const int rows = n / 11; + for (int row = 0; row < rows; ++row) { + unsigned* expired = recent_rows + misc->recent_cursor * DP_RECENT_WORDS_PER_ROW; + for (int word = 0; word < DP_RECENT_WORDS_PER_ROW; ++word) { + unsigned removed = expired[word]; + while (removed != 0) { + const int bit = __ffs(removed) - 1; + const int expert = word * 32 + bit; + recent_counts[expert] -= 1; + removed &= removed - 1; + } + expired[word] = 0; + } + for (int slot = 0; slot < 11; ++slot) { + const int expert = picks[row * 11 + slot]; + if (expert < 0 || expert >= DP_SHARED) continue; + const unsigned bit = 1u << (expert & 31); + unsigned& word = expired[expert >> 5]; + if ((word & bit) == 0) { + word |= bit; + recent_counts[expert] += 1; + } + } + misc->recent_cursor = (misc->recent_cursor + 1) % DP_RECENT_WINDOW; + if (misc->recent_filled < DP_RECENT_WINDOW) misc->recent_filled += 1; + } + } + int nu = 0; + bool over = false; + for (int i = 0; i < n; ++i) { + const int p = picks[i]; + if (p < 0 || p >= DP_SHARED) continue; + // observe: the lifetime count (saturating) and this round's appearances + if (counts[p] != 0xffffffffu) counts[p] += 1; + if (line_of[p] >= 0) misc->hit_slots += 1; else misc->miss_slots += 1; + int j = 0; + while (j < nu && uniq[j] != p) ++j; + if (j == nu) { + if (nu == DP_MAXU) { over = true; continue; } + uniq[nu] = static_cast(p); + app[nu] = 0; + keep[p >> 5] |= 1u << (p & 31); + ++nu; + } + app[j] += 1; + } + // held experts are touched; misses keep first-appearance order, then sort hottest first (insertion sort) + int nm = 0; + for (int j = 0; j < nu; ++j) { + const int id = uniq[j]; + const int ln = line_of[id]; + if (ln >= 0) { + stamp[ln] = now; + } else { + miss[nm] = static_cast(j); // index into uniq/app + ++nm; + } + } + for (int a = 1; a < nm; ++a) { + const unsigned short x = miss[a]; + int b = a - 1; + while (b >= 0) { + const unsigned short y = miss[b]; + const unsigned cx = score[uniq[x]], cy = score[uniq[y]]; + const bool hotter = cx != cy ? cx > cy : (app[x] != app[y] ? app[x] > app[y] : uniq[x] < uniq[y]); + if (!hotter) break; + miss[b + 1] = y; + --b; + } + miss[b + 1] = x; + } + if (over) misc->overflow += 1; // more distinct experts than DP_MAXU: the caller's row bound prevents it + n_uniq = nu; + n_miss = over ? 0 : nm; + n_copy = 0; + n_ovf = 0; + } + __syncthreads(); + // the split (optional): picks held when the round started (and the shared expert) vs the misses this round copies; + // the experts of the first set run while the copies of the second land + if (hit_picks != nullptr) { + for (int i = t; i < n; i += blockDim.x) { + const int p = picks[i]; + const bool held = p >= DP_SHARED || (p >= 0 && line_of[p] >= 0); + hit_picks[i] = p < 0 ? -1 : (held ? p : -1); + miss_picks[i] = p < 0 ? -1 : (held ? -1 : p); + } + __syncthreads(); + } + const unsigned long long offs[3] = {off0, off1, off2}; + long long* tabs[3] = {tab0, tab1, tab2}; + int* k2s[3] = {k2t0, k2t1, k2t2}; + for (int m = 0; m < n_miss; ++m) { + // victim: the first empty line, else the coldest resident not in this round (count, stamp, higher line) + unsigned bc = 0xffffffffu, bs = 0xffffffffu; + int bi = -1; + bool bempty = false; + for (int i = t; i < lines; i += blockDim.x) { + const int id = line_id[i]; + if (id < 0) { + if (!bempty || i < bi) { bempty = true; bi = i; } + continue; + } + if (bempty) continue; + if (keep[id >> 5] & (1u << (id & 31))) continue; + const unsigned c = score[id], s = stamp[i]; + if (bi < 0 || dp_colder(c, s, i, bc, bs, bi)) { bc = c; bs = s; bi = i; } + } + // empty lines win (encoded as count 0 with a flag in the stamp key's top bit), lowest index among them + rkey_c[t] = bempty ? 0u : bc; + rkey_s[t] = bempty ? 0u : bs; + rkey_i[t] = bempty ? -2 - bi : bi; // -2 - i < -1 marks an empty line i + __syncthreads(); + for (int w = blockDim.x / 2; w > 0; w >>= 1) { + if (t < w) { + const int ia = rkey_i[t], ib = rkey_i[t + w]; + bool take_b; + if (ib == -1) take_b = false; + else if (ia == -1) take_b = true; + else if (ia <= -2 || ib <= -2) { + // an empty line beats a resident one; between empties the lower line + if (ia <= -2 && ib <= -2) take_b = (-2 - ib) < (-2 - ia); + else take_b = ib <= -2; + } else take_b = dp_colder(rkey_c[t + w], rkey_s[t + w], ib, rkey_c[t], rkey_s[t], ia); + if (take_b) { + rkey_c[t] = rkey_c[t + w]; + rkey_s[t] = rkey_s[t + w]; + rkey_i[t] = ib; + } + } + __syncthreads(); + } + if (t == 0) { + const int r = rkey_i[0]; + const int v = r <= -2 ? -2 - r : r; + const int id = uniq[miss[m]]; + if (v >= 0) { + const int old = line_id[v]; + if (old >= 0) { + line_of[old] = -1; + for (int w = 0; w < 3; ++w) { + tabs[w][old] = 0; + k2s[w][old] = 0; + } + misc->evictions += 1; + } + line_id[v] = id; + line_of[id] = v; + stamp[v] = now; + const unsigned long long base = line_base + static_cast(v) * line_bytes; + long long* wk = work + static_cast(n_copy) * 9; + for (int w = 0; w < 3; ++w) { + tabs[w][id] = static_cast(base + offs[w]); + k2s[w][id] = k2[id * 3 + w]; + wk[w * 3 + 0] = src[id * 3 + w]; + wk[w * 3 + 1] = static_cast(base + offs[w]); + wk[w * 3 + 2] = lens[id * 3 + w]; + } + n_copy += 1; + misc->copies += 1; + } else { + // no line is free of this round's experts: the expert is read in place from the mapped host bank + // this round (same bytes, so the same results; slower), and its entries are cleared after the round + for (int w = 0; w < 3; ++w) { + tabs[w][id] = src[id * 3 + w]; + k2s[w][id] = k2[id * 3 + w]; + } + ovf[n_ovf] = id; + n_ovf += 1; + misc->inplace += 1; + } + } + __syncthreads(); + } + if (t == 0) { + *work_n = n_copy; + *ovf_n = n_ovf; + dp_trace_end(trace, trace_round, st, misc, picks, n, ovf, n_ovf); + } +} + +// After the round's experts: the in-place (host-bank) experts' table entries cleared again (they hold no line). +extern "C" __global__ void fn_exl3_drestore(const int* __restrict__ ovf, const int* __restrict__ ovf_n, + long long* __restrict__ tab0, long long* __restrict__ tab1, + long long* __restrict__ tab2, int* __restrict__ k2t0, int* __restrict__ k2t1, + int* __restrict__ k2t2) { + const int n = *ovf_n; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const int id = ovf[i]; + tab0[id] = 0; + tab1[id] = 0; + tab2[id] = 0; + k2t0[id] = 0; + k2t1[id] = 0; + k2t2[id] = 0; + } +} + +// The planned copies: each segment {src (host bank, mapped), dst (a line), len (bytes, a multiple of 16)} read as +// 16-byte words by the whole grid. Reads of pinned host memory over PCIe: many independent loads in flight. +extern "C" __global__ void __launch_bounds__(256) fn_exl3_dcopy(const long long* __restrict__ work, + const int* __restrict__ work_n) { + const int segs = *work_n * 3; + const long long tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const long long stride = static_cast(gridDim.x) * blockDim.x; + for (int g = 0; g < segs; ++g) { + const uint4* __restrict__ s = reinterpret_cast(work[g * 3 + 0]); + uint4* __restrict__ d = reinterpret_cast(work[g * 3 + 1]); + const long long words = work[g * 3 + 2] / 16; + long long i = tid; + for (; i + 3 * stride < words; i += 4 * stride) { + const uint4 a = s[i], b = s[i + stride], c = s[i + 2 * stride], e = s[i + 3 * stride]; + d[i] = a; + d[i + stride] = b; + d[i + 2 * stride] = c; + d[i + 3 * stride] = e; + } + for (; i < words; i += stride) d[i] = s[i]; + } +} + +// The draft chain on the device (cuda_exl3_chain.zig): fn_row_best's record [best, log-sum-exp, column bits, -] of a +// draft row picks the draft vocabulary's id at that column as the next chained step's input token. +extern "C" __global__ void fn_exl3_draft_token(const float* __restrict__ best, const unsigned* __restrict__ ids, + int* __restrict__ token) { + if (threadIdx.x == 0 && blockIdx.x == 0) token[0] = static_cast(ids[__float_as_int(best[2])]); +} + +// The chain's step decision on the device: the next token as fn_exl3_draft_token, then upstream's stop rule +// (Ops.draft: a pick below `confidence` ends the chain) folded into `alive`, which also drives the conditional +// node around the next head step (cudaGraphSetConditional) and is kept per step for the host. +extern "C" __global__ void fn_exl3_draft_next(const float* __restrict__ best, const unsigned* __restrict__ ids, + int* __restrict__ token, int* __restrict__ alive, + int* __restrict__ alive_out, cudaGraphConditionalHandle handle, + double confidence) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + token[0] = static_cast(ids[__float_as_int(best[2])]); + const double prob = exp(static_cast(best[0]) - static_cast(best[1])); + const int low = confidence > 0.0 && prob < confidence; + const int a = alive[0] && !low; + alive[0] = a; + alive_out[0] = a; + cudaGraphSetConditional(handle, static_cast(a)); +} diff --git a/zig/kernels/cuda/flashnext_exl3_kv.cu b/zig/kernels/cuda/flashnext_exl3_kv.cu new file mode 100644 index 0000000000..86b86642c5 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_kv.cu @@ -0,0 +1,402 @@ +// Flash Next attention over an int8 KV cache (ours, K10 prototype): upstream's flashnext_attn.cu prep / chunks / merge +// with the cache stored as int8 codes plus one fp16 scale a 32-group, dequantised into the same bf16 smem tiles. +// +// Storage (per attention layer, per sequence): K8/V8 [pos][HK][256] int8 and KS/VS [pos][HK][8] fp16. Groups of 32 +// go through one normalized 32-point Hadamard (Sylvester order; exl3_cache.zig's policy) when ROT != 0: dots are +// unchanged by an orthogonal map, so queries are stored rotated by the prep, keys and values are cached rotated, and +// the merge rotates the output back (H is its own inverse). Code: clamp(floor(x / s * 128) + 128, 0, 255) - 128, +// value (c + 0.5) * s / 128, s = fp16(max |x| over the group) (the fp16 s is used on both sides). +// The attention maths after the tile load is upstream's attn_chunks unchanged (same mma tiles, same online softmax), +// so a bf16 cache that holds the dequantised values gives the same bits (tf-cuda-test exl3-kv kernel checks this). + +#include +#include +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace fn_kv8 { + +constexpr int TK = 32; // keys a tile holds +constexpr int HD = 256; // head width +constexpr int QR = 16; // query heads a KV head's tile holds +constexpr int GR = 32; // quant group +constexpr int NG = HD / GR; + +__device__ __forceinline__ float block_sum(float v, float* scratch) { + v = warp_sum(v); + const int warps = blockDim.x >> 5; + if ((threadIdx.x & 31) == 0) scratch[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < warps; ++w) total = __fadd_rn(total, scratch[w]); + __syncthreads(); + return total; +} + +__device__ __forceinline__ float inv_rms(float ss, int width, float eps) { + return __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ss, static_cast(width)), eps))); +} + +__device__ __forceinline__ float rope(float xn, float xpn, int d, int half, int pos, const float* __restrict__ inv) { + if (d >= 2 * half) return xn; + const int i = d < half ? d : d - half; + float s, c; + sincosf(__fmul_rn(static_cast(pos), inv[i]), &s, &c); + return d < half ? __fsub_rn(__fmul_rn(xn, c), __fmul_rn(xpn, s)) : __fadd_rn(__fmul_rn(xpn, s), __fmul_rn(xn, c)); +} + +__device__ __forceinline__ int partner(int d, int half) { + return d < half ? d + half : (d < 2 * half ? d - half : d); +} + +// The warp's 32 lanes hold one group (lane = index in the group): its normalized Sylvester Hadamard transform. +__device__ __forceinline__ float wht32(float x) { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { + const float y = __shfl_xor_sync(0xffffffffu, x, m); + x = (lane & m) ? __fsub_rn(y, x) : __fadd_rn(x, y); + } + return __fmul_rn(x, 0.17677669529663688f); // 1 / sqrt(32) +} + +__device__ __forceinline__ float warp_max(float v) { +#pragma unroll + for (int m = 16; m > 0; m >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, m)); + return v; +} + +// One lane's value of a group into its code; lane 0 writes the group's fp16 scale. +__device__ __forceinline__ void quant8(float x, int8_t* code, __half* scale) { + const float peak = warp_max(fabsf(x)); + const __half sh = __float2half_rn(fmaxf(peak, 6.103515625e-05f)); + const float s = __half2float(sh); + const float q = fminf(fmaxf(floorf(__fmul_rn(__fdiv_rn(x, s), 128.0f)) + 128.0f, 0.0f), 255.0f); + *code = static_cast(static_cast(q) - 128); + if ((threadIdx.x & 31) == 0) *scale = sh; +} + +__device__ __forceinline__ float dequant8(int8_t c, float s) { + return __fmul_rn(__fadd_rn(static_cast(c), 0.5f), __fmul_rn(s, 0.0078125f)); +} + +} // namespace fn_kv8 + +using namespace fn_kv8; + +// fn_attn_prep (text positions) with the K and V rows stored as int8 codes and group scales; queries rotated if ROT. +extern "C" __global__ void __launch_bounds__(256) fn_kv8_prep(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, + const float* __restrict__ QW, const float* __restrict__ KW, + const float* __restrict__ IW, const float* __restrict__ INV, + __nv_bfloat16* __restrict__ Q, int8_t* __restrict__ K8, + __half* __restrict__ KS, int8_t* __restrict__ V8, + __half* __restrict__ VS, __nv_bfloat16* __restrict__ IQ, + __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, + int NI, int IHD, int HALF, int ROT) { + __shared__ float scratch[32]; + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + const int pos = POS0[0] + r; + const long long row = static_cast(r) * PW; + if (head >= NQ + NKV + NI) { + if (d < IHD) IKC[static_cast(pos) * IHD + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; + return; + } + const bool index = head >= NQ + NKV; + const int width = index ? IHD : HD; + const bool live = d < width; + const long long src = row + (head < NQ ? head * 2 * HD : (head < NQ + NKV ? NQ * 2 * HD + (head - NQ) * HD : NQ * 2 * HD + 2 * NKV * HD + (head - NQ - NKV) * IHD)); + const float* w = head < NQ ? QW : (head < NQ + NKV ? KW : IW); + const float x = live ? bf(P[src + d]) : 0.0f; + const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), width, eps); + if (!live) return; + const int p = partner(d, HALF); + const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), w[d])); + const float xpn = rbf(__fmul_rn(__fmul_rn(bf(P[src + p]), rinv), w[p])); + const float roped = rope(xn, xpn, d, HALF, pos, INV); + if (index) { + IQ[(static_cast(r) * NI + head - NQ - NKV) * IHD + d] = __float2bfloat16_rn(roped); + return; + } + // the cache keeps upstream's bf16 rounding point, then the rotation and the code (a query: rotated, then rounded) + if (head < NQ) { + const float qv = ROT ? wht32(roped) : roped; + Q[(static_cast(r) * NQ + head) * HD + d] = __float2bfloat16_rn(qv); + return; + } + const long long slot = static_cast(pos) * NKV + head - NQ; + const float kv = rbf(roped); + const float vv = bf(P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]); + quant8(ROT ? wht32(kv) : kv, K8 + slot * HD + d, KS + slot * NG + d / GR); + quant8(ROT ? wht32(vv) : vv, V8 + slot * HD + d, VS + slot * NG + d / GR); +} + +// Upstream's attn_chunks with each tile's keys and values read as int8 codes + scales (cp.async into a raw buffer, NB +// raw buffers in flight), dequantised into one bf16 tile pair, then the same mma / online-softmax maths. +template +__device__ __forceinline__ void kv8_chunks(const __nv_bfloat16* __restrict__ Q, const int8_t* __restrict__ K8, + const __half* __restrict__ KS, const int8_t* __restrict__ V8, + const __half* __restrict__ VS, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, + int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + typedef __nv_bfloat16 Row[HD + 8]; + Row* qs = reinterpret_cast(smem_raw); // [QR] + Row* tiles_kv = qs + QR; // [2][TK] keys then values (bf16, dequantised) + __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + 2 * TK); // [QR] + int8_t (*raw)[TK][HD] = reinterpret_cast(ps + QR); // [NB * 2][TK][HD] + __half (*rsc)[TK][NG] = reinterpret_cast<__half (*)[TK][NG]>(raw + NB * 2); // [NB * 2][TK][NG] + __shared__ float red_m[4][QR], red_l[4][QR]; + __shared__ int key_ids[NB][TK]; + const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + int n = POS0[0] + r + 1; + bool sparse = false; + if (QSA) { + sparse = SPR[r] != 0; + if (sparse) n = NKR[r]; + } + const int start = c * CH; + if (start >= n) return; + for (int e = threadIdx.x; e < QR * HD / 8; e += blockDim.x) { + const int i = e / (HD / 8), v = e % (HD / 8); + qmm_frag::cp16z(&qs[i][8 * v], Q + (static_cast(r) * H + hk * G + min(i, G - 1)) * HD + 8 * v, i < G); + } + qmm_frag::commit(); + auto load = [&](int t0, int buf) { + if (threadIdx.x < TK) { + const int ki = start + t0 + threadIdx.x; + key_ids[buf][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + } + __syncthreads(); + for (int pass = 0; pass < 2; ++pass) { + const int8_t* src = pass == 0 ? K8 : V8; + const __half* ssrc = pass == 0 ? KS : VS; + for (int e = threadIdx.x; e < TK * HD / 16; e += blockDim.x) { + const int k = e / (HD / 16), v = e % (HD / 16), key = key_ids[buf][k]; + qmm_frag::cp16z(&raw[buf * 2 + pass][k][16 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 16 * v, key >= 0); + } + if (threadIdx.x < TK) { + const int key = key_ids[buf][threadIdx.x]; + qmm_frag::cp16z(&rsc[buf * 2 + pass][threadIdx.x][0], ssrc + (static_cast(max(key, 0)) * HK + hk) * NG, key >= 0); + } + } + qmm_frag::commit(); + }; + float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.0f, 0.0f}; + float o[8][4] = {}; + const int tiles = min(n - start, CH); + load(0, 0); + for (int t0 = 0, it = 0; t0 < tiles; t0 += TK, ++it) { + const int buf = NB == 2 ? it & 1 : 0; + const bool ahead = NB == 2 && t0 + TK < tiles; + if (ahead) load(t0 + TK, buf ^ 1); + if (ahead) qmm_frag::wait<1>(); else qmm_frag::wait<0>(); + __syncthreads(); + // dequantise the raw tile into the bf16 pair: a thread 16 codes of one key (half a group) at a time + for (int e = threadIdx.x; e < 2 * TK * HD / 16; e += blockDim.x) { + const int pass = e / (TK * HD / 16), rem = e % (TK * HD / 16), k = rem / (HD / 16), v = rem % (HD / 16); + const float s = __half2float(rsc[buf * 2 + pass][k][v / 2]); + const int4 codes = *reinterpret_cast(&raw[buf * 2 + pass][k][16 * v]); + const int8_t* cb = reinterpret_cast(&codes); + __align__(16) __nv_bfloat16 outv[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) outv[j] = __float2bfloat16_rn(dequant8(cb[j], s)); + int4* dst = reinterpret_cast(&tiles_kv[pass * TK + k][16 * v]); + dst[0] = reinterpret_cast(outv)[0]; + dst[1] = reinterpret_cast(outv)[1]; + } + __syncthreads(); + Row* kt = tiles_kv; + Row* vt = tiles_kv + TK; + const int* keys = key_ids[buf]; + constexpr int JT = TK / 32; + float s[JT][4] = {}; +#pragma unroll 4 + for (int ks = 0; ks < HD / 16; ks += 2) { + uint32_t a0[4], a1[4]; + qmm_frag::ldmatrix4(a0, &qs[lane % 16][ks * 16 + (lane / 16) * 8]); + qmm_frag::ldmatrix4(a1, &qs[lane % 16][(ks + 1) * 16 + (lane / 16) * 8]); +#pragma unroll + for (int jt = 0; jt < JT; ++jt) { + uint32_t b[4]; + qmm_frag::ldmatrix4(b, &kt[warp * 8 * JT + jt * 8 + lane % 8][ks * 16 + (lane / 8) * 8]); + qmm_frag::mma(s[jt], a0, b[0], b[1]); + qmm_frag::mma(s[jt], a1, b[2], b[3]); + } + } + float tile_m[2] = {-INFINITY, -INFINITY}; +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const bool valid = keys[warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] >= 0; + s[jt][e] = valid ? __fmul_rn(s[jt][e], SCALE) : -INFINITY; + tile_m[e >> 1] = fmaxf(tile_m[e >> 1], s[jt][e]); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 1)); + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 2)); + if (t == 0) red_m[warp][g + 8 * h] = tile_m[h]; + } + __syncthreads(); + float next_m[2], alpha[2], rows[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float tm = fmaxf(fmaxf(red_m[0][row], red_m[1][row]), fmaxf(red_m[2][row], red_m[3][row])); + const bool active = tm != -INFINITY; + next_m[h] = active ? fmaxf(m[h], tm) : m[h]; + alpha[h] = active ? (m[h] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m[h], next_m[h]))) : 1.0f; + rows[h] = 0.0f; + } +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int h = e >> 1; + const float p = s[jt][e] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(s[jt][e], next_m[h])); + rows[h] = __fadd_rn(rows[h], p); + ps[g + 8 * h][warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] = __float2bfloat16_rn(p); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 1)); + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 2)); + if (t == 0) red_l[warp][g + 8 * h] = rows[h]; + } + __syncthreads(); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float sum = __fadd_rn(__fadd_rn(__fadd_rn(red_l[0][row], red_l[1][row]), red_l[2][row]), red_l[3][row]); + l[h] = __fmaf_rn(l[h], alpha[h], sum); + m[h] = next_m[h]; + } +#pragma unroll + for (int q = 0; q < 8; ++q) +#pragma unroll + for (int e = 0; e < 4; ++e) o[q][e] = __fmul_rn(o[q][e], alpha[e >> 1]); + for (int ks = 0; ks < TK / 16; ++ks) { + uint32_t a[4]; + qmm_frag::ldmatrix4(a, &ps[lane % 16][ks * 16 + (lane / 16) * 8]); +#pragma unroll + for (int q = 0; q < 8; q += 2) { + uint32_t b[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b[0]), "=r"(b[1]), "=r"(b[2]), "=r"(b[3]) + : "r"(qmm_frag::smem(&vt[ks * 16 + lane % 16][warp * 64 + q * 8 + (lane / 16) * 8]))); + qmm_frag::mma(o[q], a, b[0], b[1]); + qmm_frag::mma(o[q + 1], a, b[2], b[3]); + } + } + __syncthreads(); + if (NB == 1 && t0 + TK < tiles) load(t0 + TK, 0); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = g + 8 * h; + if (i >= G) continue; + const long long base = (static_cast(r) * NCH + c) * H + hk * G + i; +#pragma unroll + for (int q = 0; q < 8; ++q) { + float2 v = make_float2(o[q][2 * h], o[q][2 * h + 1]); + *reinterpret_cast(PO + base * HD + warp * 64 + q * 8 + 2 * t) = v; + } + if (warp == 0 && t == 0) { + PM[base] = m[h]; + PL[base] = l[h]; + } + } +} + +extern "C" __global__ void __launch_bounds__(128) fn_kv8_chunks(const __nv_bfloat16* __restrict__ Q, const int8_t* __restrict__ K8, + const __half* __restrict__ KS, const int8_t* __restrict__ V8, + const __half* __restrict__ VS, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + kv8_chunks<1>(Q, K8, KS, V8, VS, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +extern "C" __global__ void __launch_bounds__(128) fn_kv8_chunks_ahead(const __nv_bfloat16* __restrict__ Q, const int8_t* __restrict__ K8, + const __half* __restrict__ KS, const int8_t* __restrict__ V8, + const __half* __restrict__ VS, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + kv8_chunks<2>(Q, K8, KS, V8, VS, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +// fn_attn_merge, then (ROT) the rotation back in fp32 before the one bf16 rounding. +extern "C" __global__ void __launch_bounds__(256) fn_kv8_merge(const float* __restrict__ PO, const float* __restrict__ PM, + const float* __restrict__ PL, const int* __restrict__ POS0, + __nv_bfloat16* __restrict__ OUT, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + int QSA, int ROT) { + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + int n = POS0[0] + r + 1; + if (QSA && SPR[r] != 0) n = NKR[r]; + const int chunks = (n + CH - 1) / CH; + float m = -INFINITY, l = 0.0f, o = 0.0f; + for (int c = 0; c < chunks; ++c) { + const long long base = (static_cast(r) * NCH + c) * H + head; + const float cm = PM[base], cl = PL[base], co = PO[base * HD + d]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; + o = __fmaf_rn(co, b, __fmul_rn(o, a)); + l = __fmaf_rn(cl, b, __fmul_rn(l, a)); + m = next; + } + float y = __fdiv_rn(o, l); + if (ROT) y = wht32(y); + OUT[(static_cast(r) * H + head) * HD + d] = __float2bfloat16_rn(y); +} + +// A bf16 cache's rows [n][HK][256] into int8 codes + scales (an existing cache, or a test's input); ROT as above. +extern "C" __global__ void __launch_bounds__(256) fn_kv8_pack(const __nv_bfloat16* __restrict__ X, int8_t* __restrict__ C, + __half* __restrict__ S, long long rows, int ROT) { + const long long row = blockIdx.x; // one (pos, head) row of 256 a block + if (row >= rows) return; + const float x = bf(X[row * HD + threadIdx.x]); + quant8(ROT ? wht32(x) : x, C + row * HD + threadIdx.x, S + row * NG + threadIdx.x / GR); +} + +// Gather (the explicit-staging alternative to reading a host cache in place): `n` keys' rows of `row_bytes` (a multiple +// of 16) from `src` (device or mapped host) at ids[k] into dst row k, a warp a key. ids < 0 copy nothing. +extern "C" __global__ void __launch_bounds__(256) fn_kv_gather(const uint4* __restrict__ src, uint4* __restrict__ dst, + const int* __restrict__ ids, int n, int row_vecs) { + const int k = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + if (k >= n) return; + const int id = ids[k]; + if (id < 0) return; + for (int v = threadIdx.x & 31; v < row_vecs; v += 32) dst[static_cast(k) * row_vecs + v] = src[static_cast(id) * row_vecs + v]; +} + +// Test support (tf-cuda-test exl3-kv full): a bf16 cache's rows [rows][256] replaced in place by what an int8 (BITS 8) +// or int4 (BITS 4) cache would hold, back in the model's domain and rounded to bf16. An emulation of the cache's +// error inside upstream's own attention kernels (any BITS 2..15: finer grids measure the run's sensitivity); the int8 kernels' own arithmetic is checked separately (kernel mode). +extern "C" __global__ void __launch_bounds__(256) fn_kv_roundtrip(__nv_bfloat16* __restrict__ X, long long rows, int ROT, int BITS) { + const long long row = blockIdx.x; + if (row >= rows) return; + float x = bf(X[row * HD + threadIdx.x]); + if (ROT) x = wht32(x); + const float peak = warp_max(fabsf(x)); + const float s = __half2float(__float2half_rn(fmaxf(peak, 6.103515625e-05f))); + // L = 2^(BITS-1) levels a side: BITS 8 is fn_kv8's grid ((c + 0.5) s / 128), BITS 4 exl3_cache.zig's int4 grid + const float lv = static_cast(1 << (BITS - 1)); + const float q = fminf(fmaxf(floorf(__fmul_rn(__fdiv_rn(x, s), lv)) + lv, 0.0f), 2.0f * lv - 1.0f); + float y = BITS == 8 ? dequant8(static_cast(static_cast(q) - 128), s) : __fmul_rn(__fmul_rn(__fadd_rn(__fsub_rn(q, lv), 0.5f), s), __fdiv_rn(1.0f, lv)); + if (ROT) y = wht32(y); + X[row * HD + threadIdx.x] = __float2bfloat16_rn(y); +} diff --git a/zig/kernels/cuda/flashnext_exl3_kvpage.cu b/zig/kernels/cuda/flashnext_exl3_kvpage.cu new file mode 100644 index 0000000000..c90a2d8d2b --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_kvpage.cu @@ -0,0 +1,250 @@ +// Dynamic KV paging for Flash Next attention (ours; cuda_exl3_kvpage.zig, docs/plans/exl3-kv-paging.md). +// +// A sequence's full K/V of each attention layer lives in (mapped, pinned) host memory, [pos][2][256] bf16 each. The +// card keeps one compact buffer for K and one for V, [12 windows of WR rows | S slots of 4 rows]: a layer's window +// holds its positions [win0, win0 + WR) densely; the slots, shared by the 12 layers, hold 4-key blocks the sparse +// indexer selected (slot2blk keys a slot by layer * nblocks + block). A layer's kernels see the buffer from its own +// window: window rows 0..WR-1, slot s at row SB + 4 s with SB = (12 - layer) * WR. After +// upstream's fn_qsa_select listed each row's keys (positions), these kernels +// kvp_plan (one block): the call's unique blocks outside the window; hits are stamped with the call's epoch; the +// misses take the least recently used slots (radix selection of the smallest stamps, as fn_qsa_select +// picks the largest scores; equal stamps claimed in any order), and the block <-> slot tables updated; +// kvp_fetch : each miss's 4 rows of K and V read from host memory into its slot; +// kvp_remap : each sparse row's key list rewritten from positions into rows of the compact buffer; +// so upstream's fn_attn_chunks / fn_attn_merge run unchanged on the compact buffer. Dense rows (<= 2051 keys) read +// rows 0..n-1, which is the window while win0 == 0 (the window is never moved before position WR - chunk). +// +// ctl (ints, one set a layer): [0] sequence clock (layer 0 only), [1] misses of the last call, [2] sticky failure, +// [3] misses in total, [4] needed blocks in total, [5] calls with a sparse row, [6] last executed epoch. +#include +#include + +// One block of 1024 threads. +template +__device__ __forceinline__ void kvp_plan_body(const int* __restrict__ ids, const int* __restrict__ nkr, + const int* __restrict__ spr, int m, int idw, int top, + int ratio, int win0, unsigned* clock, int* blk2slot, + int* blk2slot_all, int key_off, + int* __restrict__ slot2blk, unsigned* __restrict__ stamp, + unsigned* __restrict__ mark, int nslots, int* __restrict__ ctl, + int* __restrict__ need, int* __restrict__ miss, + int* __restrict__ fetch, unsigned long long* owner_generation, + const unsigned long long* lane_generation, + unsigned long long generation, int lane_stride, int all_blocks) { + __shared__ int s_need, s_miss, s_avail; + __shared__ int hist[256]; + __shared__ unsigned cut_s; + __shared__ int left_s; + __shared__ unsigned epoch; + (void)nkr; + if (threadIdx.x == 0) { + s_need = 0; + s_miss = 0; + s_avail = 0; + if (*clock == UINT32_MAX) { + ctl[2] = 1; + epoch = 0; + } else { + epoch = ++*clock; + } + ctl[6] = static_cast(epoch); + ctl[1] = 0; + } + __syncthreads(); + if (epoch == 0 || ctl[2] != 0) return; + if (Shared && lane_generation[key_off / lane_stride] != generation) { + if (threadIdx.x == 0) ctl[2] = 1; + return; + } + bool any = false; + for (int r = 0; r < m; ++r) any |= spr[r] != 0; + if (!any) { + if (threadIdx.x == 0) ctl[1] = 0; + return; + } + // the call's unique blocks before the window (a row lists its 512 blocks as 4 consecutive keys each) + for (int e = threadIdx.x; e < m * top; e += blockDim.x) { + const int r = e / top, j = e - r * top; + if (spr[r] == 0) continue; + const int b = ids[static_cast(r) * idw + j * ratio] / ratio; + if (b * ratio >= win0) continue; + if (atomicExch(&mark[b], epoch) != epoch) need[atomicAdd(&s_need, 1)] = b; + } + __syncthreads(); + const int n_need = s_need; + for (int i = threadIdx.x; i < n_need; i += blockDim.x) { + const int b = need[i]; + const int s = blk2slot[b]; + bool hit = s >= 0 && s < nslots; + if (Shared && hit) hit = owner_generation[s] == generation && slot2blk[s] == key_off + b; + if (hit) + stamp[s] = epoch; + else { + blk2slot[b] = -1; + miss[atomicAdd(&s_miss, 1)] = b; + } + } + __syncthreads(); + int n_miss = s_miss; + // slots this call does not use + for (int s = threadIdx.x; s < nslots; s += blockDim.x) + if (stamp[s] != epoch) atomicAdd(&s_avail, 1); + __syncthreads(); + if (threadIdx.x == 0) { + if (n_miss > s_avail) ctl[2] = 1; + ctl[1] = n_miss > s_avail ? 0 : n_miss; + unsigned* counters = reinterpret_cast(ctl); + counters[3] += ctl[1]; + counters[4] += n_need; + counters[5] += 1; + } + if (n_miss > s_avail) return; + n_miss = min(n_miss, s_avail); + if (n_miss == 0) return; + // the n_miss smallest stamps among unused slots: the largest keys of ~stamp (used slots key 0, never needed) + auto key = [&](int s) -> unsigned { const unsigned t = stamp[s]; return t == epoch ? 0u : ~t; }; + unsigned cut = 0; + int left = n_miss; + for (int p = 0; p < 4; ++p) { + const int shift = 24 - 8 * p; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int s = threadIdx.x; s < nslots; s += blockDim.x) { + const unsigned k = key(s); + if (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); + } + __syncthreads(); + if (threadIdx.x == 0) { + int above = 0; + int d = 255; + for (; d > 0; --d) { + if (above + hist[d] >= left) break; + above += hist[d]; + } + left_s = left - above; + cut_s = cut | (static_cast(d) << shift); + } + __syncthreads(); + cut = cut_s; + left = left_s; + __syncthreads(); + } + // every slot above the cut and `left` of those equal to it (any of them: equal stamps are equally old) take the + // misses; claims through two shared counters, so no block-wide scan + __shared__ int s_eq, s_place; + if (threadIdx.x == 0) { + s_eq = 0; + s_place = 0; + } + __syncthreads(); + for (int s = threadIdx.x; s < nslots; s += blockDim.x) { + const unsigned k = key(s); + bool chosen = k > cut; + if (!chosen && k == cut) chosen = atomicAdd(&s_eq, 1) < left; + if (!chosen) continue; + const int place = atomicAdd(&s_place, 1); + if (place >= n_miss) continue; + const int b = miss[place]; + const int old = slot2blk[s]; + if (Shared) { + if (old >= 0 && old < all_blocks && owner_generation[s] == lane_generation[old / lane_stride] && blk2slot_all[old] == s) blk2slot_all[old] = -1; + owner_generation[s] = generation; + } else if (old >= 0) blk2slot_all[old] = -1; + blk2slot[b] = s; + slot2blk[s] = key_off + b; + stamp[s] = epoch; + fetch[2 * place] = b; + fetch[2 * place + 1] = s; + } +} + +extern "C" __global__ void __launch_bounds__(1024) kvp_plan( + const int* ids, const int* counts, const int* sparse, int rows, int id_width, int top_blocks, + int ratio, int window_origin, unsigned* clock, int* inverse, int* all_inverse, int key_offset, + int* owners, unsigned* stamps, unsigned* marks, int slots, int* controls, int* needed, + int* missing, int* fetched) { + kvp_plan_body(ids, counts, sparse, rows, id_width, top_blocks, ratio, window_origin, + clock, inverse, all_inverse, key_offset, owners, stamps, marks, slots, controls, needed, + missing, fetched, nullptr, nullptr, 0, 1, 0); +} + +extern "C" __global__ void __launch_bounds__(1024) kvp_shared_plan( + const int* ids, const int* counts, const int* sparse, int rows, int id_width, int top_blocks, + int ratio, int window_origin, unsigned* clock, int* inverse, int* all_inverse, int key_offset, + int* owners, unsigned* stamps, unsigned* marks, int slots, int* controls, int* needed, + int* missing, int* fetched, unsigned long long* owner_generation, + const unsigned long long* lane_generation, unsigned long long generation, int lane_stride, + int all_blocks) { + kvp_plan_body(ids, counts, sparse, rows, id_width, top_blocks, ratio, window_origin, + clock, inverse, all_inverse, key_offset, owners, stamps, marks, slots, controls, needed, + missing, fetched, owner_generation, lane_generation, generation, lane_stride, all_blocks); +} + +extern "C" __global__ void __launch_bounds__(256) kvp_shared_release( + int* all_inverse, int* owners, unsigned* stamps, unsigned long long* owner_generation, + int slots, int key_begin, int key_end, unsigned long long generation) { + for (int slot = blockIdx.x * blockDim.x + threadIdx.x; slot < slots; slot += gridDim.x * blockDim.x) { + const int previous = owners[slot]; + if (previous < key_begin || previous >= key_end || owner_generation[slot] != generation) continue; + if (all_inverse[previous] == slot) all_inverse[previous] = -1; + owners[slot] = -1; + owner_generation[slot] = 0; + stamps[slot] = 0; + } +} + +// Each miss's block: 4 rows of K and of V (ROW16 16-byte words a row) from host memory into its slot (dk / dv: the +// slots' first row, WR 0; or a layer's window base with WR its slot base rows). +extern "C" __global__ void __launch_bounds__(256) kvp_fetch(const uint4* __restrict__ hk, const uint4* __restrict__ hv, + uint4* __restrict__ dk, uint4* __restrict__ dv, + const int* __restrict__ ctl, const int* __restrict__ fetch, + int wr, int row16, int ratio) { + const int n = ctl[1]; + const long long per = static_cast(ratio) * row16; + for (int i = blockIdx.x; i < n; i += gridDim.x) { + const long long b = fetch[2 * i], s = fetch[2 * i + 1]; + const uint4* sk = hk + b * per; + const uint4* sv = hv + b * per; + uint4* tk = dk + static_cast(wr) * row16 + s * per; + uint4* tv = dv + static_cast(wr) * row16 + s * per; + for (int e = threadIdx.x; e < per; e += blockDim.x) { + tk[e] = sk[e]; + tv[e] = sv[e]; + } + } +} + +// A row a block: its sparse key list from positions into rows from the layer's window base (WR: the slot base SB). +extern "C" __global__ void __launch_bounds__(256) kvp_remap(int* __restrict__ ids, int* __restrict__ nkr, + const int* __restrict__ spr, int idw, int win0, int wr, + const int* __restrict__ blk2slot, int ratio, int* ctl) { + const int r = blockIdx.x; + if (spr[r] == 0) return; + const int n = nkr[r]; + int* row = ids + static_cast(r) * idw; + __shared__ int missing; + if (threadIdx.x == 0) missing = atomicAdd(&ctl[2], 0) != 0; + __syncthreads(); + for (int k = threadIdx.x; k < n; k += blockDim.x) { + const int p = row[k]; + if (p < win0 && blk2slot[p / ratio] < 0) atomicExch(&missing, 1); + } + __syncthreads(); + if (missing) { + if (threadIdx.x == 0) { + atomicExch(&ctl[2], 1); + nkr[r] = 0; + } + return; + } + for (int k = threadIdx.x; k < n; k += blockDim.x) { + const int p = row[k]; + row[k] = p >= win0 ? p - win0 : wr + blk2slot[p / ratio] * ratio + p % ratio; + } +} + +// n rows of 16-byte words: dst[i] = src[i] for i < n (a window move within one buffer, issued in non-overlapping +// pieces by the host, or a write-out to host memory). +extern "C" __global__ void __launch_bounds__(256) kvp_copy(const uint4* __restrict__ src, uint4* __restrict__ dst, long long n) { + for (long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += static_cast(gridDim.x) * blockDim.x) dst[i] = src[i]; +} diff --git a/zig/kernels/cuda/flashnext_exl3_prefill.cu b/zig/kernels/cuda/flashnext_exl3_prefill.cu new file mode 100644 index 0000000000..438793fdc8 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_prefill.cu @@ -0,0 +1,327 @@ +// Flash Next EXL3 prefill experts: upstream's transposed 64-member prompt tiles (exl3_prompt.cuh, after #976 and +// MiaAI-Lab) for any per-expert width (K2 half-bits) of the mul1 codebook, on compacted pairs. +// +// A call takes items (expert, first, count <= 64) over a member list of pair ids (row * slots + slot) that the host +// built, expert by expert. The rotated inputs and the gate/up outputs are COMPACT: member c of the list owns row c of +// xg, xu and xd, so a call needs scratch for its own pairs only (an expert-major pass of a prefill chunk). The down +// writes each pair's fp32 output to y[pair], the layout Experts.runSlots leaves (y [rows * slots, D]). +// Per pair the maths is runSlots' (act_mode 1: fp32 SwiGLU): the same fp16 rotated input, the same epilogue formulas; +// only the fp32 order of the K sum differs (one pass over K here, split-K partials summed in order there). + +#include +#include +#include + +#include "exl3_decode.cuh" +#include "exl3_prompt.cuh" + +namespace tf_exl3pf { +using namespace tf_exl3; +using tf_exl3_prompt::cp_async16; +using tf_exl3_prompt::cp_async_commit; +using tf_exl3_prompt::cp_async_wait; +using tf_exl3_prompt::ldmatrix_x4; +using tf_exl3_prompt::tspill; +using tf_exl3_prompt::ttile; +using tf_exl3_prompt::TPM; +using tf_exl3_prompt::TNW; +using tf_exl3_prompt::TMT; +using tf_exl3_prompt::TNB; +using tf_exl3_prompt::TT; +using tf_exl3_prompt::TPH; +using tf_exl3_prompt::TBK; +using tf_exl3_prompt::TKT; +using tf_exl3_prompt::TSTAGES; +using tf_exl3_prompt::TXS; +using tf_exl3_prompt::GU_ES; +using tf_exl3_prompt::DN_ES; + +constexpr int KMAX2 = 10; // the widest tile a stage holds (5 bits): Flash Next's routed 3-bit and shared 5-bit experts +template +constexpr int stage_words() { return 16 * TKT * tile_words(); } +template +constexpr int stage_halves() { return XR * TXS + 2 * stage_words(); } +constexpr int smem_bytes() { + const int gu = TSTAGES * stage_halves<2 * TPM>() * 2; + const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +template +__device__ __forceinline__ void decode_tile_k(const uint32_t* tile, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t w[lane_words()]; + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = tile[(word + i) % tile_words()]; + decode_lane(w, lane, b0, b1); +} + +// acc[i][nb] = warp tile i x members nb * 8 .. + 7 over K; member i's input row is X[xfirst + i] (compact rows). +template +__device__ __forceinline__ void pmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, int xfirst, int cnt, int K, int NTILES, int cb, half* stage, float (&acc)[TMT][TNB][4]) { + constexpr int SH = stage_halves(); + constexpr int TW = tile_words(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < TNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; + const bool ok = i < cnt; + const half* src = (m ? X1 : X0) + (size_t)(xfirst + (ok ? i : 0)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, ok); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * K2; idx += TT) { // a tile is K2 pieces of 16 bytes + const int q = idx / (TKT * K2), kk = (idx / K2) % TKT, part = idx % K2; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)chunk * TKT + kk) * NTILES + nt) * TW + part * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * TW + part * 4, src, true); + } + }; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_k(ws + ((size_t)(warp * TMT + i) * TKT + kk) * TW, lane, b0, b1); + a[i][0] = b0[0]; + a[i][1] = b1[0]; + a[i][2] = b0[1]; + a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < TPM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +__device__ __forceinline__ float upTo(float x, float l) { return x > l ? l : x; } +__device__ __forceinline__ float within(float x, float l) { return x > l ? l : (x < -l ? -l : x); } +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +template +__device__ __forceinline__ void gateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* Tg, const uint32_t* Tu, + int e, int first, int cnt, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, + half* __restrict__ xd, int K, int N, float limit, int act_mode, unsigned char* smem_raw) { + const int cb = blockIdx.x; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + pmain(Xg, Xu, Tg, Tu, first, cnt, K, N >> 4, cb, reinterpret_cast(smem_raw), acc); + float* Eg = reinterpret_cast(smem_raw); + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); + tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float act; + if (act_mode == 0) { + float gg = bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])); + float uu = bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])); + if (limit > 0) { gg = upTo(gg, limit); uu = within(uu, limit); } + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { // runSlots' tf_exl3_gateup_epilogue, act_mode 1, operation for operation + float gg = gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j]); + float uu = uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j]); + if (limit > 0) { gg = upTo(gg, limit); uu = within(uu, limit); } + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)(first + hf * TPH + i) * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +template +__device__ __forceinline__ void down(const half* __restrict__ Xd, const uint32_t* Td, int e, int first, int cnt, const int* rows_sh, + const half* __restrict__ svh_d, float* __restrict__ y, int K, int D, unsigned char* smem_raw) { + const int cb = blockIdx.x; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + pmain(Xd, Xd, Td, Td, first, cnt, K, D >> 4, cb, reinterpret_cast(smem_raw), acc); + float* Ed = reinterpret_cast(smem_raw); + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); + tspill(Ed, DN_ES, warp * TMT * 16, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const int n = cb * 256 + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = Ed[(size_t)i * DN_ES + blk * 128 + 4 * lane + j]; + fwht128(v, lane); + float* o = y + (size_t)p * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + } + } + } +} + +} // namespace tf_exl3pf + +using namespace tf_exl3pf; + +// Program (compact member c, 128-block of D, matrix): xg/xu[c] = fp16((x[row] * suh[e]) @ H) for member c's pair. +extern "C" __global__ void tf_exl3pf_rotate(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, + const int* __restrict__ members, int slots, const half* __restrict__ suh_g, const half* __restrict__ suh_u, + half* __restrict__ xg, half* __restrict__ xu, int D) { + const int c = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int p = members[c]; + const int e = pick[p]; + const int row = p / slots; + const int lane = threadIdx.x; + const half* suh = (mat ? suh_u : suh_g) + (size_t)e * D + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? xu : xg) + (size_t)c * D + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +#define TF_EXL3PF_K2(X) \ + X(4) \ + X(6) \ + X(8) \ + X(10) + +// Block (128 width columns of gate and up, item): SwiGLU of the item's members into compact xd rows. +extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3pf_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, + const uint64_t* __restrict__ Tg, const uint64_t* __restrict__ Tu, const int* __restrict__ k2g, const int* __restrict__ items, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int K, + int N, float limit, int act_mode) { + __shared__ __align__(16) unsigned char smem_raw[smem_bytes()]; + const int item = blockIdx.y; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const uint32_t* G = reinterpret_cast(Tg[e]); + const uint32_t* U = reinterpret_cast(Tu[e]); + switch (k2g[e]) { +#define TF_CASE(K2_) \ + case K2_: gateup(Xg, Xu, G, U, e, first, cnt, svh_g, svh_u, suh_d, xd, K, N, limit, act_mode, smem_raw); break; + TF_EXL3PF_K2(TF_CASE) +#undef TF_CASE + default: __trap(); + } +} + +// Block (256 model-width columns, item): y[pair] = the down epilogue of the item's members. +extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3pf_down(const half* __restrict__ Xd, const uint64_t* __restrict__ Td, + const int* __restrict__ k2d, const int* __restrict__ items, const int* __restrict__ members, const half* __restrict__ svh_d, + float* __restrict__ y, int K, int D) { + __shared__ __align__(16) unsigned char smem_raw[smem_bytes()]; + __shared__ int rows_sh[TPM]; + const int item = blockIdx.y; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const uint32_t* T = reinterpret_cast(Td[e]); + switch (k2d[e]) { +#define TF_CASE(K2_) \ + case K2_: down(Xd, T, e, first, cnt, rows_sh, svh_d, y, K, D, smem_raw); break; + TF_EXL3PF_K2(TF_CASE) +#undef TF_CASE + default: __trap(); + } +} + +// Expert parallel (cuda_exl3_tp.zig): a rank's deep chunk computed only its own pairs (picks in [LO, HI)); row r's +// combined branch over those slots in slot order (fp32 fmaf from 0), into sum [rows, D] for the ranks' all-reduce. +// Another rank's slot is skipped by a select (its y is never read). Block (row), 256 threads over D. +extern "C" __global__ void __launch_bounds__(256) tf_exl3pf_tp_combine(const float* __restrict__ y, const float* __restrict__ wts, + const int* __restrict__ pick, float* __restrict__ sum, int slots, int D, int lo, int hi) { + __shared__ float w_sh[32]; + __shared__ int own_sh[32]; + const int r = blockIdx.x; + if ((int)threadIdx.x < slots) { + const int p = pick[r * slots + threadIdx.x]; + own_sh[threadIdx.x] = p >= lo && p < hi; + w_sh[threadIdx.x] = wts[r * slots + threadIdx.x]; + } + __syncthreads(); + const float* yr = y + (size_t)r * slots * D; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + float acc = 0.0f; + for (int k = 0; k < slots; ++k) + if (own_sh[k]) acc = fmaf(yr[(size_t)k * D + d], w_sh[k], acc); + sum[(size_t)r * D + d] = acc; + } +} + +// The ranks' summed branch back where the write-back reads it: y[r, 0, :] = sum[r], weights {1, 0, ..., 0}. The other +// slots' y keep finite values (zeroed when the deep buffers are made, then only expert outputs and sums), so the +// write-back's fma(y_k, 0, acc) adds exactly 0 for them. Block (row), 256 threads over D. +extern "C" __global__ void __launch_bounds__(256) tf_exl3pf_tp_place(const float* __restrict__ sum, float* __restrict__ y, + float* __restrict__ wts, int slots, int D) { + const int r = blockIdx.x; + float* out = y + (size_t)r * slots * D; + const float* in = sum + (size_t)r * D; + for (int d = threadIdx.x; d < D; d += blockDim.x) out[d] = in[d]; + if ((int)threadIdx.x < slots) wts[r * slots + threadIdx.x] = threadIdx.x == 0 ? 1.0f : 0.0f; +} diff --git a/zig/kernels/cuda/flashnext_exl3_tp.cu b/zig/kernels/cuda/flashnext_exl3_tp.cu new file mode 100644 index 0000000000..e0c8e8d02f --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_tp.cu @@ -0,0 +1,27 @@ +// The Flash Next family's EXL3 expert-parallel glue (cuda_exl3_tp.zig): a rank's masked picks, and the summed MoE +// branch placed where the pass's write-back reads the per-slot outputs. + +// A rank holds the expert ids [LO, HI) (global ids: rank 0 the routed 0..255, rank 1 the routed 256..511 and the +// shared expert 512). Another rank's slot becomes pick -1 (Experts.run skips it) with weight 0, so the rank's combined +// output is the fp32 fmaf chain over its own slots in slot order (a skipped slot adds weight 0 times a finite value). +extern "C" __global__ void fn_exl3_tp_mask(const int* __restrict__ pick, const float* __restrict__ wts, + int* __restrict__ mpick, float* __restrict__ mwts, int n, int lo, int hi) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const int p = pick[i]; + const bool own = p >= lo && p < hi; + mpick[i] = own ? p : -1; + mwts[i] = own ? wts[i] : 0.0f; +} + +// The ranks' summed branch [rows, D] into the write-back's slot layout: y[r, 0, :] = sum, the other slots 0, and the +// weights {1, 0, ..., 0}. hc_writeback (MODE 2) then reads fma(sum, 1, 0) = sum and adds 0 * 0 for every other slot, +// so it writes the sum's bf16 exactly. Program (row, slot), 256 threads over D. +extern "C" __global__ void __launch_bounds__(256) fn_exl3_tp_place(const float* __restrict__ sum, float* __restrict__ y, + float* __restrict__ wts, int d, int slots) { + const int r = blockIdx.x, k = blockIdx.y; + float* out = y + (static_cast(r) * slots + k) * d; + const float* in = sum + static_cast(r) * d; + for (int i = threadIdx.x; i < d; i += blockDim.x) out[i] = k == 0 ? in[i] : 0.0f; + if (k == 0 && threadIdx.x < slots) wts[r * slots + threadIdx.x] = threadIdx.x == 0 ? 1.0f : 0.0f; +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..27298e519a 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -336,6 +336,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .release = .{ .major = 13, .minor = 3 }, .usage = &.{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, } }, + .{ .release = .{ .major = 13, .minor = 4 }, .usage = &.{.{ .name = "exl3_candidate tf_exl3_prompt_2", .usage = .{ .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 } } }} }, }; pub fn newer(x: limits.Nvcc, y: limits.Nvcc) bool { return x.major > y.major or (x.major == y.major and x.minor > y.minor); diff --git a/zig/src/core/contract/registry_pending_test.zig b/zig/src/core/contract/registry_pending_test.zig index 8e0abe9a9f..e4d60ecad3 100644 --- a/zig/src/core/contract/registry_pending_test.zig +++ b/zig/src/core/contract/registry_pending_test.zig @@ -9,6 +9,19 @@ const cuda = registry.cuda; const checkCuda = registry.checkCuda; const targetDecision = registry.targetDecision; const TargetDecision = registry.TargetDecision; +test "EXL3 prompt resource records select measured CUDA 13.4 usage without changing older compiler records" { + const record = for (cuda) |candidate| { + if (std.mem.eql(u8, candidate.launch.name, "exl3_candidate tf_exl3_prompt_2")) break candidate; + } else return error.MissingExl3PromptRecord; + const older = registry.usageFor(record, .{ .major = 13, .minor = 3 }); + const current = registry.usageFor(record, .{ .major = 13, .minor = 4 }); + for (current, 0..) |usage, index| { + try std.testing.expectEqual(@as(u32, if (index < 3) 80 else 77), usage.regs); + try std.testing.expectEqual(@as(u32, 80), older[index].regs); + try std.testing.expectEqual(@as(u32, 32768), usage.static); + try std.testing.expectEqual(@as(u32, 256), usage.bound); + } +} test "pending 13.0 CUDA launches have no invented usage and cannot pass the measured gate" { for (cuda_pending) |r| { try std.testing.expect(r.launch.threads > 0); diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index 45ff280c07..4034af8db4 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -39,7 +39,8 @@ test "a compiler's records: its own set, else the newest below it, else 13.0's" .{ .{ .major = 13, .minor = 1 }, .{ .major = 13, .minor = 0 }, 91, 45 }, .{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 2 }, 91, 43 }, .{ .{ .major = 13, .minor = 3 }, .{ .major = 13, .minor = 3 }, 96, 43 }, - .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 3 }, 96, 43 }, + .{ .{ .major = 13, .minor = 4 }, .{ .major = 13, .minor = 4 }, 96, 43 }, + .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 4 }, 96, 43 }, }; for (cases) |c| { try std.testing.expectEqual(c[1], setFor(c[0])); diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index 10747bd88e..20846b7910 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -20,6 +20,7 @@ pub const error_not_ready: Result = 600; pub const stream_non_blocking: c_uint = 1; pub const event_disable_timing: c_uint = 2; +pub const event_wait_external: c_uint = 1; pub const host_alloc_portable: c_uint = 1; pub const host_alloc_devicemap: c_uint = 2; pub const host_register_devicemap: c_uint = 2; diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index c5bb710e5a..e2d0415bc1 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -52,6 +52,23 @@ pub fn symbol(comptime halves: u32, comptime codebook: Codebook, comptime warps: pub const rot_in_symbol = "tf_exl3_rot_in"; +const RotInArgs = extern struct { + x: u64, + suh: u64, + xh: u64, + x_dtype: i32, + m: i32, + k: i32, + counters: u64, + counter_count: i32, +}; + +comptime { + std.debug.assert(@sizeOf(RotInArgs) == 56); + std.debug.assert(@offsetOf(RotInArgs, "counters") == 40); + std.debug.assert(@offsetOf(RotInArgs, "counter_count") == 48); +} + pub fn variantIndex(halves: u32, codebook: Codebook, warps: u32) ?usize { for (variants, 0..) |v, i| { if (v.halves == halves and v.codebook == codebook and v.warps == warps) return i; @@ -130,28 +147,32 @@ pub const Linear = struct { return if (p.sk > 1) @as(usize, p.sk) * m * w.n * 4 else 0; } - /// Counter bytes one call needs: per pass and block, zeroed (fill8(0), never zeros()). + /// Counter bytes one call needs: per pass and block, reset by input rotation before each split-K launch. pub fn counterBytes(w: Weight) usize { return 8 * (w.n / 128) * 4; } - /// y [m, n] = x [m, k] through the weight; `xh` is fp16 [m, k] scratch, `z` the split partials, `counters` zeroed. + /// y [m, n] = x [m, k] through the weight; `xh` is fp16 [m, k] scratch, `z` the split partials, `counters` scratch. pub fn matmul(l: *const Linear, s: Stream, w: Weight, x: u64, x_dtype: DType, m: usize, y: u64, y_dtype: DType, xh: u64, z: u64, counters: u64) !void { // Shapes the loader admits (linear.py's k % 128 and n % 128, 1..128 rows) are returned errors here too. if (m < 1 or m > 128) return error.Exl3Shape; - if (w.k % 128 != 0 or w.n % 128 != 0) return error.Exl3Shape; + if (w.k == 0 or w.n == 0 or w.k % 128 != 0 or w.n % 128 != 0) return error.Exl3Shape; const p = plan(w.k, w.n, 192, 8); if ((w.k / 16) % (p.sk * p.wk) != 0) return error.Exl3Shape; const idx = variantIndex(w.halves, w.codebook, p.wk) orelse return error.UnsupportedExl3Width; - if (p.sk > 1 and z == 0) return error.Exl3Scratch; + if (p.sk > 1 and (z == 0 or counters == 0)) return error.Exl3Scratch; var r: launch_.Args = .{}; - r.add(x); - r.add(w.suh); - r.add(xh); - r.add(@as(i32, @backingInt(x_dtype))); - r.add(@as(i32, @intCast(m))); - r.add(@as(i32, @intCast(w.k))); + r.add(RotInArgs{ + .x = x, + .suh = w.suh, + .xh = xh, + .x_dtype = @backingInt(x_dtype), + .m = @intCast(m), + .k = @intCast(w.k), + .counters = counters, + .counter_count = if (p.sk > 1) @intCast(counterBytes(w) / 4) else 0, + }); try launch_.launch(l.rot_in, .{ .grid = .{ .x = @intCast((w.k / 128 + 3) / 4), .y = @intCast(m), .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, @@ -190,7 +211,7 @@ pub const Dense = struct { candidate_mod: ?Module = null, candidate_k: ?candidate.Kernels = null, implementation: enum { reference, tiles } = .reference, - counters: ?memory.DeviceBuffer = null, // per pass and block, zeroed before every call + counters: ?memory.DeviceBuffer = null, pub fn load(d: *const Driver) !Dense { return .{ .d = d, .lin = try Linear.load(d) }; @@ -264,7 +285,6 @@ pub const Dense = struct { while (r < rows) : (r += 128) { const m = @min(rows - r, 128); const xh = try de.scratch(&de.xh, m * w.k * 2); - try de.counters.?.fill8(0, s.handle); try de.lin.matmul(s, w, x + @as(u64, r) * w.k * xe, x_dtype, m, y + @as(u64, r) * w.n * ye, y_dtype, xh, z, counters); } } @@ -307,6 +327,74 @@ test "matmul refuses shapes the loader would refuse, before any launch" { nbad.n = 100; try std.testing.expectError(error.Exl3Shape, l.matmul(s, nbad, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); try std.testing.expectError(error.Exl3Scratch, l.matmul(s, w, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); // plan(4096, 1024) splits K + try std.testing.expectError(error.Exl3Scratch, l.matmul(s, w, 0, .bf16, 8, 0, .bf16, 0, 1, 0)); + kbad = w; + kbad.k = 0; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, kbad, 0, .bf16, 8, 0, .bf16, 0, 1, 1)); + nbad = w; + nbad.n = 0; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, nbad, 0, .bf16, 8, 0, .bf16, 0, 1, 1)); +} + +test "dense reference blocks fuse counter reset into rotation on every invocation" { + const abi = @import("abi.zig"); + const Fake = struct { + var rotations: usize = 0; + var products: usize = 0; + var clears: usize = 0; + var expected_counters: u64 = 0; + var expected_count: i32 = 0; + var last_rows: i32 = 0; + + fn launch(function: abi.Function, grid_x: c_uint, grid_y: c_uint, _: c_uint, threads: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result { + if (@intFromPtr(function.?) == 1) { + const args: *const RotInArgs = @ptrCast(@alignCast(parameters.?[0].?)); + std.debug.assert(args.counters == expected_counters and args.counter_count == expected_count); + std.debug.assert(grid_x == (@as(u32, @intCast(args.k)) / 128 + 3) / 4); + std.debug.assert(grid_y == args.m and threads == 128); + last_rows = args.m; + rotations += 1; + } else { + std.debug.assert(rotations == products + 1); + products += 1; + } + return abi.success; + } + + fn clear(_: u64, _: u8, _: usize, _: abi.Stream) callconv(.c) abi.Result { + clears += 1; + return abi.success; + } + }; + var test_driver = @import("exl3_mock.zig").init(); + test_driver.func_name = null; + test_driver.api.cuLaunchKernel = Fake.launch; + test_driver.api.cuMemsetD8Async = Fake.clear; + const rotation: Function = .{ .d = &test_driver, .handle = @ptrFromInt(1) }; + const product: Function = .{ .d = &test_driver, .handle = @ptrFromInt(2) }; + var dense: Dense = .{ + .d = &test_driver, + .lin = .{ .mod = undefined, .rot_in = rotation, .fns = @splat(product) }, + .xh = .{ .d = &test_driver, .ptr = 0x1000, .len = 128 * 4096 * 2 }, + .z = .{ .d = &test_driver, .ptr = 0x2000, .len = 8 * 128 * 8192 * 4 }, + .counters = .{ .d = &test_driver, .ptr = 0x3000, .len = 8 * 8192 / 128 * 4 }, + }; + Fake.expected_counters = dense.counters.?.ptr; + for ([_][2]u32{ .{ 2560, 2560 }, .{ 512, 512 }, .{ 4096, 8192 }, .{ 2560, 512 } }) |shape| { + const layout_strides = strides(shape[0], shape[1], 8, false); + const weight: Weight = .{ .trellis = 7, .suh = 8, .svh = 9, .k = shape[0], .n = shape[1], .halves = 8, .codebook = .mul1, .stride_k = layout_strides.k, .stride_nb = layout_strides.nb }; + Fake.expected_count = if (plan(weight.k, weight.n, 192, 8).sk > 1) @intCast(Linear.counterBytes(weight) / 4) else 0; + for ([_]usize{ 1, 17, 128, 129, 300, 3 }) |row_count| { + const before = Fake.products; + try dense.decode(.{ .d = &test_driver, .handle = null }, weight, 10, .bf16, row_count, 11, .bf16); + try std.testing.expectEqual((row_count + 127) / 128, Fake.products - before); + try std.testing.expectEqual(@as(i32, @intCast((row_count - 1) % 128 + 1)), Fake.last_rows); + try dense.prompt(.{ .d = &test_driver, .handle = null }, weight, 10, .bf16, row_count, 11, .f32); + try std.testing.expectEqual(2 * ((row_count + 127) / 128), Fake.products - before); + try std.testing.expectEqual(Fake.rotations, Fake.products); + try std.testing.expectEqual(@as(usize, 0), Fake.clears); + } + } } test "a failed scratch allocation keeps the old scratch owned" { diff --git a/zig/src/cuda/exl3_candidate.zig b/zig/src/cuda/exl3_candidate.zig index c894f33096..ce68fa0870 100644 --- a/zig/src/cuda/exl3_candidate.zig +++ b/zig/src/cuda/exl3_candidate.zig @@ -44,7 +44,7 @@ pub const MatrixArgs = extern struct { k: u32, n: u32, bits: u32, tiles_n: u32, /// Decode rows a launch; longer decode runs go window by window. pub const window = 16; /// Prompt rows a chunk: rotated to fp16 together, then 64 rows a block on the tensor units. -pub const prompt_chunk = 256; +pub const prompt_chunk = 2048; /// Blocks a launch aims for: a matrix with fewer 128-column blocks splits its input blocks into parts. const target_blocks = 96; diff --git a/zig/src/cuda/exl3_dma.zig b/zig/src/cuda/exl3_dma.zig new file mode 100644 index 0000000000..f2b78ffb6c --- /dev/null +++ b/zig/src/cuda/exl3_dma.zig @@ -0,0 +1,38 @@ +const std = @import("std"); +const Stream = @import("stream.zig").Stream; + +pub fn directLines(stream: Stream, dst: u64, capacity: usize, offsets: [3]usize, src: [3][]const u8, pinned_bank: ?[]const u8) !usize { + const Span = struct { dst: u64, bytes: []const u8, banked: bool }; + var spans: [3]Span = undefined; + var count: usize = 0; + var moved: usize = 0; + _ = std.math.add(u64, dst, capacity) catch return error.Invalid; + const bank_start = if (pinned_bank) |bank| @intFromPtr(bank.ptr) else 0; + const bank_end = if (pinned_bank) |bank| std.math.add(usize, bank_start, bank.len) catch return error.Invalid else 0; + for (src, offsets) |bytes, offset| { + if (offset > capacity or bytes.len > capacity - offset) return error.Invalid; + const start = @intFromPtr(bytes.ptr); + const end = std.math.add(usize, start, bytes.len) catch return error.Invalid; + moved = std.math.add(usize, moved, bytes.len) catch return error.Invalid; + const address = dst + offset; + const banked = pinned_bank != null and start >= bank_start and end <= bank_end; + if (count > 0 and bytes.len > 0 and banked) { + const previous = &spans[count - 1]; + if (previous.banked and previous.bytes.len > 0 and + @intFromPtr(previous.bytes.ptr) + previous.bytes.len == start and + previous.dst + previous.bytes.len == address) + { + previous.bytes = pinned_bank.?[@intFromPtr(previous.bytes.ptr) - bank_start ..][0 .. previous.bytes.len + bytes.len]; + continue; + } + } + spans[count] = .{ .dst = address, .bytes = bytes, .banked = banked }; + count += 1; + } + for (spans[0..count]) |span| try stream.d.check(stream.d.api.cuMemcpyHtoDAsync_v2(span.dst, span.bytes.ptr, span.bytes.len, stream.handle), "cuMemcpyHtoDAsync"); + return moved; +} + +test { + _ = @import("exl3_dma_test.zig"); +} diff --git a/zig/src/cuda/exl3_dma_test.zig b/zig/src/cuda/exl3_dma_test.zig new file mode 100644 index 0000000000..8cbafd61cb --- /dev/null +++ b/zig/src/cuda/exl3_dma_test.zig @@ -0,0 +1,216 @@ +const std = @import("std"); +const dma = @import("exl3_dma.zig"); +const pool = @import("exl3_pool.zig"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const Event = @import("stream.zig").Event; + +const Fake = struct { + const Call = enum { query, wait, copy, table, event }; + const Copy = struct { dst: u64, src: usize, len: usize, stream: abi.Stream }; + const Word = struct { dst: u64, value: u32 }; + var calls: [32]Call = undefined; + var count: usize = 0; + var copies: [3]Copy = undefined; + var copied: usize = 0; + var fail_copy: usize = 0; + var words: [18]Word = undefined; + var written: usize = 0; + + fn note(call: Call) void { + calls[count] = call; + count += 1; + } + + fn copy(dst: u64, src: ?*const anyopaque, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + note(.copy); + copies[copied] = .{ .dst = dst, .src = @intFromPtr(src), .len = len, .stream = stream }; + copied += 1; + return if (copied == fail_copy) abi.error_not_ready else abi.success; + } + + fn table(dst: u64, value: c_uint, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + std.debug.assert(len == 1 and stream == @as(abi.Stream, @ptrFromInt(7))); + note(.table); + words[written] = .{ .dst = dst, .value = value }; + written += 1; + return abi.success; + } + + fn event(_: abi.Event, _: abi.Stream) callconv(.c) abi.Result { + note(.event); + return abi.success; + } + + fn query(_: abi.Event) callconv(.c) abi.Result { + note(.query); + return abi.error_not_ready; + } + + fn wait(_: abi.Stream, _: abi.Event, _: c_uint) callconv(.c) abi.Result { + note(.wait); + return abi.success; + } + + fn driver() Driver { + count = 0; + copied = 0; + fail_copy = 0; + written = 0; + var result: Driver = undefined; + result.api.cuMemcpyHtoDAsync_v2 = copy; + result.api.cuMemsetD32Async = table; + result.api.cuEventRecord = event; + result.api.cuEventQuery = query; + result.api.cuStreamWaitEvent = wait; + return result; + } +}; + +fn run(src: [3][]const u8, offsets: [3]usize, bank: ?[]const u8, lengths: []const usize) !void { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = @ptrFromInt(7) }; + const moved = try dma.directLines(stream, 4096, 32, offsets, src, bank); + try std.testing.expectEqual(src[0].len + src[1].len + src[2].len, moved); + try std.testing.expectEqual(lengths.len, Fake.copied); + var projection: usize = 0; + for (Fake.copies[0..Fake.copied], lengths) |copy, len| { + try std.testing.expectEqual(len, copy.len); + try std.testing.expectEqual(4096 + offsets[projection], copy.dst); + try std.testing.expectEqual(@intFromPtr(src[projection].ptr), copy.src); + try std.testing.expectEqual(stream.handle, copy.stream); + var consumed: usize = 0; + while (projection < 3) { + const bytes = src[projection]; + const uploaded: [*]const u8 = @ptrFromInt(copy.src); + try std.testing.expectEqualSlices(u8, bytes, uploaded[consumed..][0..bytes.len]); + consumed += bytes.len; + projection += 1; + if (consumed == len) break; + } + try std.testing.expectEqual(len, consumed); + } + try std.testing.expectEqual(@as(usize, 3), projection); +} + +test "DMA coalescing: exact unequal widths merge three projections or either pair" { + const bank = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, &bank, &.{9}); + try run(src, .{ 0, 2, 8 }, &bank, &.{ 5, 4 }); + try run(src, .{ 0, 4, 7 }, &bank, &.{ 2, 7 }); + try run(.{ bank[0..2], bank[2..5], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 5, 4 }); + try run(.{ bank[0..2], bank[3..6], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 2, 7 }); +} + +test "DMA coalescing: absent provenance and source or destination gaps retain exact fallback copies" { + const bank: [32]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, null, &.{ 2, 3, 4 }); + try run(src, .{ 0, 4, 9 }, &bank, &.{ 2, 3, 4 }); + try run(.{ bank[0..2], bank[4..7], bank[9..13] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[10..20], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[0..4], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[2..9], &.{ 2, 7 }); + try run(src, .{ 0, 2, 5 }, bank[0..0], &.{ 2, 3, 4 }); + const foreign: [9]u8 = @splat(3); + try run(.{ foreign[0..2], foreign[2..5], foreign[5..9] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); +} + +test "DMA coalescing: zero lengths preserve fallback and never bridge a gap" { + const bank: [9]u8 = @splat(2); + try run(.{ bank[0..0], bank[0..3], bank[3..9] }, .{ 0, 0, 3 }, &bank, &.{ 0, 9 }); + try run(.{ bank[0..2], bank[2..2], bank[2..9] }, .{ 0, 2, 2 }, &bank, &.{ 2, 0, 7 }); + try run(.{ bank[0..2], bank[2..2], bank[3..9] }, .{ 0, 2, 3 }, &bank, &.{ 2, 0, 6 }); + try run(.{ bank[0..0], bank[0..0], bank[0..0] }, .{ 0, 0, 0 }, &bank, &.{ 0, 0, 0 }); +} + +test "DMA coalescing: bounds and overflowing source bank or destination fail before copying" { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const bank: [9]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 8, .{ 0, 2, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 10, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, std.math.maxInt(usize) }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, std.math.maxInt(u64) - 8, 9, .{ 0, 2, 5 }, src, &bank)); + const high: [*]const u8 = @ptrFromInt(std.math.maxInt(usize) - 1); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, .{ src[0], src[1], high[0..4] }, null)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, src, high[0..4])); + try std.testing.expectEqual(@as(usize, 0), Fake.count); +} + +test "DMA coalescing: production three-bit sizes issue one exact transfer" { + const width = 614400; + const bank = try std.testing.allocator.alloc(u8, 3 * width); + defer std.testing.allocator.free(bank); + @memset(bank, 1); + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const layout = pool.Layout.init(.{ width, width, width }); + const src: [3][]const u8 = .{ bank[0..width], bank[width .. 2 * width], bank[2 * width ..] }; + try std.testing.expectEqual(@as(usize, 3 * width), try dma.directLines(stream, 4096, layout.line_bytes, layout.off, src, bank)); + try std.testing.expectEqual(@as(usize, 1), Fake.copied); + try std.testing.expectEqual(@as(usize, 3 * width), Fake.copies[0].len); +} + +test "DMA coalescing: pool preserves reader waits table writes event and expert byte statistics" { + const bank: [1536]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..256], bank[256..768], bank[768..1536] }; + for ([_]?[]const u8{ null, &bank, bank[256..] }, [_]usize{ 3, 1, 2 }) |provenance, copies| { + for (0..copies + 1) |failure| { + var driver = Fake.driver(); + Fake.fail_copy = failure; + var epochs = [_]u64{1}; + var pending = [_]bool{false}; + var events = [_]Event{.{ .d = &driver, .handle = @ptrFromInt(2) }}; + var lines: pool.Lines = undefined; + lines.d = &driver; + lines.spec = .{ .experts = 2, .bytes = .{ 256, 512, 768 }, .lines = 1, .staging = 0 }; + lines.layout = pool.Layout.init(lines.spec.bytes); + lines.tlay = pool.Tables.init(2); + lines.mem = .{ .d = &driver, .ptr = 4096, .len = 1536 }; + lines.tables = .{ .d = &driver, .ptr = 8192, .len = lines.tlay.bytes }; + lines.staging = null; + lines.pinned_bank = provenance; + lines.copy = .{ .d = &driver, .handle = @ptrFromInt(7) }; + lines.line_ev = &events; + lines.line_pending = &pending; + lines.line_epoch = &epochs; + lines.rel_ev[1] = .{ .d = &driver, .handle = @ptrFromInt(3) }; + lines.done_epoch = 0; + lines.stats = .{}; + const result = lines.load(0, 1, src, .{ 3, 6, 9 }, 0, true, undefined); + try std.testing.expectEqualSlices(Fake.Call, &.{ .query, .wait }, Fake.calls[0..2]); + if (failure > 0) { + try std.testing.expectError(error.NotReady, result); + try std.testing.expectEqual(2 + failure, Fake.count); + try std.testing.expectEqual(failure, Fake.copied); + try std.testing.expect(!pending[0]); + try std.testing.expectEqual(@as(u64, 0), lines.stats.copies); + try std.testing.expectEqual(@as(u64, 0), lines.stats.bytes); + } else { + try result; + try std.testing.expectEqual(2 + copies + 18 + 1, Fake.count); + for (Fake.calls[2 + copies .. Fake.count - 1]) |call| try std.testing.expectEqual(Fake.Call.table, call); + try std.testing.expectEqual(Fake.Call.event, Fake.calls[Fake.count - 1]); + try std.testing.expect(pending[0]); + try std.testing.expectEqual(@as(u64, 1), lines.stats.copies); + try std.testing.expectEqual(@as(u64, bank.len), lines.stats.bytes); + try std.testing.expectEqual(@as(u64, 1), lines.stats.release_waits); + for (0..3) |which| { + const old = Fake.words[3 * which ..][0..3]; + const installed = Fake.words[9 + 3 * which ..][0..3]; + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which), .value = 0 }, old[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 4, .value = 0 }, old[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which), .value = 0 }, old[2]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 8, .value = @intCast(4096 + lines.layout.off[which]) }, installed[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 12, .value = 0 }, installed[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which) + 4, .value = @intCast(3 * (which + 1)) }, installed[2]); + } + } + for (Fake.calls[2 .. 2 + Fake.copied]) |call| try std.testing.expectEqual(Fake.Call.copy, call); + } + } +} diff --git a/zig/src/cuda/exl3_pool.zig b/zig/src/cuda/exl3_pool.zig new file mode 100644 index 0000000000..3cdfd03803 --- /dev/null +++ b/zig/src/cuda/exl3_pool.zig @@ -0,0 +1,349 @@ +//! Device lines for one layer's EXL3 routed-expert offload on upstream's `exl3_experts.Experts`. +//! +//! `Experts` reads a trellis POINTER per expert (int64 [E] for gate, up, down) and a half-bit width per expert (int32 +//! [E] each), all in device memory, in the stored layout. So a line is ONE device region holding one expert's gate, up +//! and down trellis back to back (no strips reorder), and making expert `id` resident in line `i` is: copy the line, +//! then write `id`'s six table words (three pointers, three widths) to point into line `i`. The expert that line held +//! has its entries zeroed (no stale pointer survives). The tables' addresses never change, so a CUDA graph that baked +//! them in replays after any number of swaps. +//! +//! Ordering, all on the device (the host never waits for a copy unless it reuses a staging slot): +//! copy stream: [wait release(epoch of the last pass that read the line)] -> H2D line -> table memsets -> line event +//! compute stream: [wait line event] for every line a pass reads -> experts -> release event (an epoch) +//! A miss is staged: the host copies the expert's three trellis tensors from the (mmapped) checkpoint into a pinned +//! staging line, then one H2D moves the whole line on the copy engine (measured 13.7 GB/s, 5-8x faster than kernels +//! reading mapped host memory). `direct` skips the staging memcpy when the source bytes are already pinned. +//! Which expert goes in which line is the policy's business (families/flashnext/exl3_offload.zig, exl3_round.zig). + +const std = @import("std"); +const abi = @import("abi.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +pub const dma = @import("exl3_dma.zig"); + +const Driver = @import("driver.zig").Driver; +const Error = @import("driver.zig").Error; +const DeviceBuffer = memory.DeviceBuffer; +const HostBuffer = memory.HostBuffer; +const Stream = stream_.Stream; +const Event = stream_.Event; + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +pub const Spec = struct { + /// Table length E (513: 512 routed experts and the shared expert as id 512). + experts: usize, + /// Largest trellis bytes of any expert, per projection (gate, up, down): a line fits every expert of the layer. + bytes: [3]usize, + lines: usize, + /// Pinned staging lines (a ring); 0 when every copy is `direct`. + staging: usize = 8, +}; + +/// Byte offsets of gate, up and down inside a line, each 256-aligned. +pub const Layout = struct { + off: [3]usize, + line_bytes: usize, + + pub fn init(bytes: [3]usize) Layout { + var l: Layout = .{ .off = undefined, .line_bytes = 0 }; + for (bytes, &l.off) |b, *o| { + o.* = l.line_bytes; + l.line_bytes += align256(b); + } + return l; + } +}; + +/// Where the six tables sit inside one allocation: three int64 [E] pointer tables, then three int32 [E] width tables. +pub const Tables = struct { + ptr: [3]usize, + k2: [3]usize, + bytes: usize, + + pub fn init(experts: usize) Tables { + var t: Tables = .{ .ptr = undefined, .k2 = undefined, .bytes = 0 }; + for (&t.ptr) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 8); + } + for (&t.k2) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 4); + } + return t; + } +}; + +pub const Stats = struct { + copies: u64 = 0, + bytes: u64 = 0, + staged_bytes: u64 = 0, // bytes the host memcpy'd into pinned staging + stage_ns: u64 = 0, // host time in those memcpys + staging_waits: u64 = 0, // a staging slot was still being copied from (the host blocked) + release_waits: u64 = 0, // a copy waited on the device for a pass that read its line + gate_waits: u64 = 0, // a pass waited on the device for a line copy +}; + +const release_ring = 64; + +pub const Lines = struct { + d: *const Driver, + spec: Spec, + layout: Layout, + tlay: Tables, + mem: DeviceBuffer, + tables: DeviceBuffer, + shared: ?DeviceBuffer = null, + pinned_bank: ?[]const u8 = null, + staging: ?HostBuffer, + stage_ev: []Event, + stage_busy: []bool, + stage_next: usize = 0, + copy: Stream, + line_ev: []Event, + line_pending: []bool, + /// Epoch of the last compute pass that read the line (0: never read). + line_epoch: []u64, + rel_ev: [release_ring]Event, + epoch: u64 = 0, // last epoch released + done_epoch: u64 = 0, // highest epoch known finished + stats: Stats = .{}, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, d: *const Driver, spec: Spec) !Lines { + if (spec.lines == 0 or spec.experts == 0) return error.Invalid; + const layout = Layout.init(spec.bytes); + const tlay = Tables.init(spec.experts); + var mem = try DeviceBuffer.alloc(d, spec.lines * layout.line_bytes); + errdefer mem.free(); + var tables = try DeviceBuffer.alloc(d, tlay.bytes); + errdefer tables.free(); + try tables.fill8(0, null); // every expert absent; synchronous, before any stream uses the tables + var staging: ?HostBuffer = if (spec.staging > 0) try HostBuffer.alloc(d, spec.staging * layout.line_bytes) else null; + errdefer if (staging) |*h| h.free(); + var copy = try Stream.init(d, true); + errdefer copy.deinit(); + const stage_ev = try makeEvents(gpa, d, spec.staging); + errdefer freeEvents(gpa, stage_ev); + const stage_busy = try gpa.alloc(bool, spec.staging); + errdefer gpa.free(stage_busy); + @memset(stage_busy, false); + const line_ev = try makeEvents(gpa, d, spec.lines); + errdefer freeEvents(gpa, line_ev); + const line_pending = try gpa.alloc(bool, spec.lines); + errdefer gpa.free(line_pending); + @memset(line_pending, false); + const line_epoch = try gpa.alloc(u64, spec.lines); + errdefer gpa.free(line_epoch); + @memset(line_epoch, 0); + var rel: [release_ring]Event = undefined; + var made: usize = 0; + errdefer for (rel[0..made]) |*e| e.deinit(); + for (&rel) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return .{ + .d = d, + .spec = spec, + .layout = layout, + .tlay = tlay, + .mem = mem, + .tables = tables, + .staging = staging, + .stage_ev = stage_ev, + .stage_busy = stage_busy, + .copy = copy, + .line_ev = line_ev, + .line_pending = line_pending, + .line_epoch = line_epoch, + .rel_ev = rel, + .gpa = gpa, + }; + } + + fn makeEvents(gpa: std.mem.Allocator, d: *const Driver, n: usize) ![]Event { + const evs = try gpa.alloc(Event, n); + var made: usize = 0; + errdefer { + for (evs[0..made]) |*e| e.deinit(); + gpa.free(evs); + } + for (evs) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return evs; + } + + fn freeEvents(gpa: std.mem.Allocator, evs: []Event) void { + for (evs) |*e| e.deinit(); + gpa.free(evs); + } + + pub fn deinit(l: *Lines) void { + l.copy.synchronize() catch {}; + for (&l.rel_ev) |*e| e.deinit(); + freeEvents(l.gpa, l.line_ev); + freeEvents(l.gpa, l.stage_ev); + l.gpa.free(l.stage_busy); + l.gpa.free(l.line_pending); + l.gpa.free(l.line_epoch); + l.copy.deinit(); + if (l.staging) |*h| h.free(); + if (l.shared) |*b| b.free(); + l.tables.free(); + l.mem.free(); + l.* = undefined; + } + + /// Device bytes: the lines, the tables and the resident shared expert. + pub fn deviceBytes(l: Lines) usize { + return l.mem.len + l.tables.len + if (l.shared) |b| b.len else 0; + } + + pub fn table(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.ptr[which]; + } + + pub fn widths(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.k2[which]; + } + + /// Device address of projection `which` (0 gate, 1 up, 2 down) of line `line`. + pub fn at(l: Lines, line: usize, which: usize) u64 { + return l.mem.ptr + line * l.layout.line_bytes + l.layout.off[which]; + } + + /// Uploads the never-offloaded expert `id` (the shared expert) once into its own region (it may be wider than a + /// routed line) and points its table entries at it. Synchronous: call before the first round. + pub fn setResident(l: *Lines, id: usize, src: [3][]const u8, k2: [3]i32) !void { + if (id >= l.spec.experts or l.shared != null) return error.Invalid; + const lay = Layout.init(.{ src[0].len, src[1].len, src[2].len }); + var b = try DeviceBuffer.alloc(l.d, lay.line_bytes); + errdefer b.free(); + for (src, lay.off) |s, o| try b.upload(o, s); + l.shared = b; + for (0..3) |w| try l.writeEntry(id, b.ptr + lay.off[w], k2[w], w); + try l.copy.synchronize(); + } + + fn memset32(l: *Lines, addr: u64, value: u32, st: Stream) !void { + try l.d.check(l.d.api.cuMemsetD32Async(addr, value, 1, st.handle), "cuMemsetD32Async"); + } + + /// One expert's projection `which`: its pointer (two 32-bit words) and width, on the copy stream. + fn writeEntry(l: *Lines, id: usize, ptr: u64, k2: i32, which: usize) !void { + const st = l.copy; + const p = l.table(which) + id * 8; + try l.memset32(p, @truncate(ptr), st); + try l.memset32(p + 4, @truncate(ptr >> 32), st); + try l.memset32(l.widths(which) + id * 4, @bitCast(k2), st); + } + + /// The copy stream waits until the last compute pass that read `line` has finished (no wait if known done). + fn afterReaders(l: *Lines, line: usize) !void { + const e = l.line_epoch[line]; + if (e == 0 or e <= l.done_epoch) return; + const slot = e % release_ring; + const ev = l.rel_ev[slot]; + if (try ev.done()) { + // the slot holds epoch e or a newer one; either way e has finished + l.done_epoch = @max(l.done_epoch, e); + return; + } + try l.copy.wait(ev); // a newer epoch in the slot is a conservative wait + l.stats.release_waits += 1; + } + + /// Makes `id` resident in `line`: the copy waits for the line's readers, then H2D (staged or direct), then `id`'s + /// table entries point into the line and `evicted`'s (if any) are zeroed, then the line's event. + pub fn load(l: *Lines, line: usize, id: usize, src: [3][]const u8, k2: [3]i32, evicted: ?usize, direct: bool, io: std.Io) !void { + if (line >= l.spec.lines or id >= l.spec.experts) return error.Invalid; + for (src, l.spec.bytes) |s, cap| if (s.len > cap) return error.Invalid; + try l.afterReaders(line); + const dst = l.mem.ptr + line * l.layout.line_bytes; + var moved: usize = 0; + if (direct or l.staging == null) { + moved = try dma.directLines(l.copy, dst, l.layout.line_bytes, l.layout.off, src, l.pinned_bank); + } else { + const slot = l.stage_next; + l.stage_next = (l.stage_next + 1) % l.spec.staging; + if (l.stage_busy[slot]) { + if (!try l.stage_ev[slot].done()) l.stats.staging_waits += 1; + try l.stage_ev[slot].synchronize(); + } + const buf = l.staging.?.bytes[slot * l.layout.line_bytes ..][0..l.layout.line_bytes]; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + for (src, l.layout.off) |s, o| @memcpy(buf[o..][0..s.len], s); + l.stats.stage_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t0); + // one copy for the line: up to the end of down's bytes (alignment padding rides along) + const n = l.layout.off[2] + src[2].len; + try l.d.check(l.d.api.cuMemcpyHtoDAsync_v2(dst, buf.ptr, n, l.copy.handle), "cuMemcpyHtoDAsync"); + try l.stage_ev[slot].record(l.copy); + l.stage_busy[slot] = true; + l.stats.staged_bytes += src[0].len + src[1].len + src[2].len; + moved = n; + } + if (evicted) |old| for (0..3) |w| try l.writeEntry(old, 0, 0, w); + for (0..3) |w| try l.writeEntry(id, dst + l.layout.off[w], k2[w], w); + try l.line_ev[line].record(l.copy); + l.line_pending[line] = true; + l.stats.copies += 1; + l.stats.bytes += moved; + } + + /// True once the line's last copy finished (or it never had one); clears the pending flag. + pub fn landed(l: *Lines, line: usize) !bool { + if (!l.line_pending[line]) return true; + if (!try l.line_ev[line].done()) return false; + l.line_pending[line] = false; + return true; + } + + /// `compute` waits on the device for the line's copy, if one may still be running. The host does not block. + pub fn gate(l: *Lines, compute: Stream, line: usize) !void { + if (try l.landed(line)) return; + try compute.wait(l.line_ev[line]); + l.stats.gate_waits += 1; + } + + /// After a compute pass that read `lines`: a release event on `compute`; a later copy into any of those lines + /// waits for it on the device. + pub fn release(l: *Lines, compute: Stream, lines: []const u32) !void { + l.epoch += 1; + const slot = l.epoch % release_ring; + try l.rel_ev[slot].record(compute); + for (lines) |i| l.line_epoch[i] = l.epoch; + } + + /// Blocks until every queued copy and table write finished. + pub fn drain(l: *Lines) !void { + try l.copy.synchronize(); + for (l.line_pending) |*p| p.* = false; + @memset(l.stage_busy, false); + } +}; + +test { + _ = dma; +} + +test "a line packs gate, up and down 256-aligned; the tables sit in one allocation" { + // Flash Next layer 0: 3-bit experts, K=2560 N=640 -> 160 x 40 tiles of 48 words + const tb = 160 * 40 * 48 * 2; + const l = Layout.init(.{ tb, tb, tb }); + try std.testing.expectEqual([3]usize{ 0, tb, 2 * tb }, l.off); + try std.testing.expectEqual(@as(usize, 3 * tb), l.line_bytes); + const odd = Layout.init(.{ 100, 300, 5 }); + try std.testing.expectEqual([3]usize{ 0, 256, 768 }, odd.off); + try std.testing.expectEqual(@as(usize, 1024), odd.line_bytes); + const t = Tables.init(513); + try std.testing.expectEqual([3]usize{ 0, 4352, 8704 }, t.ptr); + try std.testing.expectEqual(@as(usize, 13056), t.k2[0]); + try std.testing.expectEqual(@as(usize, 13056 + 3 * 2304), t.bytes); + for (t.ptr ++ t.k2) |o| try std.testing.expectEqual(@as(usize, 0), o % 256); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 80c5f50d6f..3bc05bf794 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -14,6 +14,12 @@ pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; +pub const flashnext_exl3: []const u8 = if (available) &Blob("fatbin_flashnext_exl3").bytes else &.{}; +pub const flashnext_exl3_tp: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_tp").bytes else &.{}; +pub const flashnext_exl3_dplan: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_dplan").bytes else &.{}; +pub const flashnext_exl3_kv: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_kv").bytes else &.{}; +pub const flashnext_exl3_kvpage: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_kvpage").bytes else &.{}; +pub const flashnext_exl3_prefill: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_prefill").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index b244b6ae2c..fcb9bcaa11 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -64,6 +64,7 @@ pub const exl3_experts = @import("exl3_experts.zig"); pub const exl3_upload = @import("exl3_upload.zig"); pub const exl3_candidate = @import("exl3_candidate.zig"); pub const exl3_grouped = @import("exl3_grouped.zig"); +pub const exl3_pool = @import("exl3_pool.zig"); pub const arena = @import("arena.zig"); pub const pack = @import("pack.zig"); @@ -104,6 +105,7 @@ test { _ = exl3; _ = exl3_experts; _ = exl3_upload; + _ = exl3_pool; _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); diff --git a/zig/src/cuda/stream.zig b/zig/src/cuda/stream.zig index 1733965193..bb7fbda772 100644 --- a/zig/src/cuda/stream.zig +++ b/zig/src/cuda/stream.zig @@ -54,6 +54,10 @@ pub const Stream = struct { pub fn wait(self: Stream, event: Event) Error!void { try self.d.check(self.d.api.cuStreamWaitEvent(self.handle, event.handle, 0), "cuStreamWaitEvent"); } + + pub fn waitExternal(self: Stream, event: Event) Error!void { + try self.d.check(self.d.api.cuStreamWaitEvent(self.handle, event.handle, abi.event_wait_external), "cuStreamWaitEvent"); + } }; pub const Event = struct { diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 5e79eee832..cb3bd3f1c6 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -7,6 +7,8 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); +const exl3_pass = @import("cuda_exl3_pass.zig"); +const exl3_kvpage = @import("cuda_exl3_kvpage.zig"); const layers_n = wt.layers_n; const linear_n = wt.linear_n; @@ -107,11 +109,12 @@ pub fn Pass(comptime Tri: type) type { prepared: bool = false, /// The pending write-back's branch sits in fp32 in moe_f32 (something added to it, or it is kept). pending_f32: bool = false, - + pending_f32_weights: ?u64 = null, // ---- launches ---------------------------------------------------------------------------------------- /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !void { + if (w.exl3) |e| return exl3_pass.lane(self, x, e, out, part, m, sk); if (w.native) |native| { if (sk == 1) return self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. @@ -207,6 +210,11 @@ pub fn Pass(comptime Tri: type) type { if (self.pending_f32) { // the fp32 branch rounded once, as MODE 2 rounds its sum (a slice of one) self.pending_f32 = false; + if (self.pending_f32_weights) |weights| { + self.pending_f32_weights = null; + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, s.moe_f32 + r0 * dims * 4, "*fp32", weights + r0 * 4, "*fp32", rs, n, dims, streams, 2, 0, 1, 1); + return; + } try self.tri.hcWriteback(h, h, pss, s.moe_f32 + r0 * dims * 4, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", rs, n, dims, streams, 4, 1, 1, 1); return; } @@ -344,6 +352,7 @@ pub fn Pass(comptime Tri: type) type { /// attnBlock, keeping its q|gate, k and v projections and o_proj's input. pub fn attnBlockKeep(self: *Self, lw: *const Layer, ai: usize, keep: LayerKeep) !void { + if (@import("cuda_exl3_heads.zig").isSplit(lw)) return @import("cuda_exl3_heads.zig").attnBlock(self, lw, ai); const s = self.s; const n = self.rows; const pa = s.proj; @@ -377,6 +386,8 @@ pub fn Pass(comptime Tri: type) type { /// `n` rows of self.seq from self.pos0: rope, cache writes and pooling, then scores, selection and attention. fn attnRows(self: *Self, lw: *const Layer, ai: usize, pa: u64, n: usize, at: RowsAt) !void { + if (self.seq.kv8) |c| return @import("cuda_exl3_kv.zig").attnRows(self, lw, ai, pa, n, at, c); // int8 KV (K10) + if (exl3_kvpage.takes(self.seq)) return exl3_kvpage.attnRows(self, lw, ai, pa, n, at); const s = self.s; const kc = self.seq.kc[ai].ptr; const vc = self.seq.vc[ai].ptr; @@ -426,6 +437,14 @@ pub fn Pass(comptime Tri: type) type { const s = self.s; const n = self.rows; const pairs = n * slots; + if (lw.exl3) |x| { + if (x.needsPool()) return error.Exl3PoolNotAttached; + if (keep.moe_x != 0) try self.copy(keep.moe_x, s.mixed, n * dims * 2); + try exl3_pass.moe(self, lw, x); + if (keep.moe_ids != 0) try self.copy(keep.moe_ids, s.picks, pairs * 4); + if (keep.moe_wts != 0) try self.copy(keep.moe_wts, s.wts, pairs * 4); + return; + } try self.copy(keep.moe_x, s.mixed, n * dims * 2); try self.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); try self.tri.topkMany(s.router_l, s.picks, s.wts, n); @@ -498,6 +517,7 @@ pub fn Pass(comptime Tri: type) type { self.pos0 = pos0; self.pending = null; self.pending_f32 = false; + self.pending_f32_weights = null; // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first if (!self.prepared) try s.ready(); const emb = s.hostEmb(n); @@ -551,7 +571,7 @@ pub fn Pass(comptime Tri: type) type { pub fn projectHead(self: *Self, out: u64, n: usize) !void { const s = self.s; if (self.m.head.n == vocab) return self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); - if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; + if (self.m.head.n != vocab / 2 or n > @min(s.cap, parts.head_local_rows)) return error.UnexpectedTensor; const c = self.m.collective orelse return error.TensorParallelUnavailable; try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index dc1749c657..56917247ef 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -47,6 +47,9 @@ pub const att_rows: usize = 256; pub const pocket_stride: usize = 128; /// Rows a pass holds at most unless asked for more. pub const default_chunk: usize = 1536; +/// Rows of `head_local` (a vocabulary-split head's local logits) a scratch holds at most: every row by default; an +/// EXL3 engine sets it to its widest logits pass (0 with a whole head, cuda_exl3_kv.zig setup, K10 lever 3). +pub var head_local_rows: usize = std.math.maxInt(usize); const ratio = wt.index_ratio; const top = wt.index_top; const idw = wt.index_idw; @@ -120,12 +123,18 @@ comptime { /// A caller's view into the pass. Without any, the pass is the server's prompt pass. pub const Hooks = struct { + pub const OutputTail = struct { + ptr: *anyopaque, + enqueue: *const fn (*anyopaque, cuda.Stream) anyerror!void, + }; + logits: Logits = .last, /// Device [rows, vocab] bf16 for `.all` (else the pass's own buffer for the last row). logits_out: u64 = 0, keep: ?*Keep = null, /// Device [rows, wide] bf16: every row's streams after the last MoE write-back (the MTP head's input); 0: none. streams: u64 = 0, + output_tail: ?OutputTail = null, /// A training pass's adds to each layer's MoE output, in place of the model's own (Model.pocket). moe_add: ?MoeAdd = null, }; @@ -204,6 +213,7 @@ pub const Uploads = struct { event: cuda.Event, pub fn copy(u: Uploads, dst: u64, bytes: []const u8) !void { + if (@import("cuda_exl3_graph.zig").capturing) return @import("cuda_exl3_graph.zig").pinnedCopy(u.stream, dst, bytes); // a copy node reads the pinned bytes at each replay try cuda.copy_source.enqueue(u.stream.d, dst, bytes, u.stream); } pub fn record(u: Uploads) !void { @@ -247,8 +257,11 @@ pub fn Staging(comptime Q: type) type { } /// Device scratch for chunks of up to `cap` rows. +pub const ScratchFormat = enum { classic, exl3 }; + pub const Scratch = struct { cap: usize, + format: ScratchFormat = .classic, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -364,103 +377,164 @@ pub const Scratch = struct { return @alignCast(std.mem.bytesAsSlice(f32, s.host.?.bytes[hostAt(s.cap)[3]..][0 .. n * (dims / 32) * 4])); } - /// Bytes `alloc(cap)` takes on the device. - pub fn bytes(cap: usize) usize { + pub fn allocationSizes(cap: usize, format: ScratchFormat) [63]?usize { const pairs = cap * slots; const nblk = cdiv(pairs, plan_small); const items = cdiv(pairs, tile) + experts_n + 1; - var n: usize = 0; - n += cap * 4 + (cap + 1) * 16; - n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; - n += cap * wide * 2 + 8 * cap * dims * 4; - n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; - n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; - n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; - n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; - n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; - n += att_rows * (idw + 2) * 4; - n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; - n += cap * moe_w * 6 + cap * dims * 4; - n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; - n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; - n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; - n += cap * dims * 4 + cap * pocket_stride * 4; - return n; + return .{ + cap * 4, + (cap + 1) * 16, + cap * wide * 2, + cap * pss_n * 4, + cap * wide * 2, + cap * (low + 4) * 2, + hc_sk * cap * (low + 4) * 4, + cap * wide * 2, + 8 * cap * dims * 4, + cap * low * 2, + cap * (low / 32) * 4, + cap * streams * 2, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * proj_row * 2, + out_sk * cap * dims * 4, + cap * gdn.value_dim * 2, + cap * (gdn.value_dim / 32) * 4, + cap * gdn.value_dim * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.value_dim * 4, + cap * gdn.nv * 2 * 4, + cap * q_heads * head_dim * 2, + cap * index_heads * index_dim * 2, + att_rows * nch * q_heads * head_dim * 4, + att_rows * nch * q_heads * 4, + att_rows * nch * q_heads * 4, + att_rows * idw * 4, + att_rows * 4, + att_rows * 4, + cap * out_k * 2, + cap * out_k * 2, + cap * (out_k / 32) * 4, + cap * (experts_n + 1) * 4, + pairs * 4, + pairs * 4, + if (format == .classic) pairs * 4 else null, + if (format == .classic) items * 3 * 4 else null, + if (format == .classic) 8 else null, + if (format == .classic) pairs * 4 else null, + if (format == .classic) nblk * (experts_n + 1) * 4 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * dims * 4 else null, + if (format == .classic) pairs * moe_w * 2 else null, + pairs * dims * 4, + cap * dims * 2, + cap * dims * 4, + cap * pocket_stride * 4, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * wide * 2, + cap * dims * 2, + 4 * cap * dims * 4, + cap * wide * 2, + cap * streams * 4, + cap * wide * 2, + vocab * 2, + @min(cap, head_local_rows) * vocab, + half * 4, + 4096, + }; + } + + pub fn bytes(cap: usize) usize { + return bytesForFormat(cap, .classic); + } + + pub fn bytesForFormat(cap: usize, format: ScratchFormat) usize { + var total: usize = 0; + for (allocationSizes(cap, format)) |optional| if (optional) |len| { + total += @max(len, 256); + }; + return total; } pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { - var s: Scratch = .{ .cap = cap }; + return allocForFormat(gpa, d, stream, cap, .classic); + } + + pub fn allocForFormat(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, format: ScratchFormat) !Scratch { + var s: Scratch = .{ .cap = cap, .format = format }; errdefer s.free(gpa); - const pairs = cap * slots; - const nblk = cdiv(pairs, plan_small); - const items = cdiv(pairs, tile) + experts_n + 1; - s.ids = try s.add(gpa, d, cap * 4); + const sizes = allocationSizes(cap, format); + s.ids = if (sizes[0]) |len| try s.add(gpa, d, len) else 0; // a position every 16 bytes: one a block of att_rows rows, or one a row of a decode round - s.pos = try s.add(gpa, d, (cap + 1) * 16); - s.h = try s.add(gpa, d, cap * wide * 2); - s.pss = try s.add(gpa, d, cap * pss_n * 4); - s.norm = try s.add(gpa, d, cap * wide * 2); - s.dn = try s.add(gpa, d, cap * (low + 4) * 2); - s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); - s.hc_up = try s.add(gpa, d, cap * wide * 2); - s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); - s.act = try s.add(gpa, d, cap * low * 2); - s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); - s.inj = try s.add(gpa, d, cap * streams * 2); - s.mixed = try s.add(gpa, d, cap * dims * 2); - s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); - s.proj = try s.add(gpa, d, cap * proj_row * 2); - s.part = try s.add(gpa, d, out_sk * cap * dims * 4); - s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); - s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); - s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.gq = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gk = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gv = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.ggb = try s.add(gpa, d, cap * gdn.nv * 2 * 4); - s.q = try s.add(gpa, d, cap * q_heads * head_dim * 2); - s.iq = try s.add(gpa, d, cap * index_heads * index_dim * 2); - s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); - s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.qids = try s.add(gpa, d, att_rows * idw * 4); - s.qnk = try s.add(gpa, d, att_rows * 4); - s.qspr = try s.add(gpa, d, att_rows * 4); - s.ao = try s.add(gpa, d, cap * out_k * 2); - s.gated = try s.add(gpa, d, cap * out_k * 2); - s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); - s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); - s.picks = try s.add(gpa, d, pairs * 4); - s.wts = try s.add(gpa, d, pairs * 4); - s.members = try s.add(gpa, d, pairs * 4); - s.items = try s.add(gpa, d, items * 3 * 4); - s.counts = try s.add(gpa, d, 8); - s.rank = try s.add(gpa, d, pairs * 4); - s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); - s.shared_gate = try s.add(gpa, d, cap * moe_w * 2); - s.shared_up = try s.add(gpa, d, cap * moe_w * 2); - s.shared_act = try s.add(gpa, d, cap * moe_w * 2); - s.shared_out = try s.add(gpa, d, cap * dims * 4); - s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); - s.y = try s.add(gpa, d, pairs * dims * 4); - s.branch = try s.add(gpa, d, cap * dims * 2); - s.moe_f32 = try s.add(gpa, d, cap * dims * 4); - s.pocket_act = try s.add(gpa, d, cap * pocket_stride * 4); - s.ple_emb = try s.add(gpa, d, cap * dims * 2); - s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); - s.ple_keys = try s.add(gpa, d, cap * wide * 2); - s.ple_vals = try s.add(gpa, d, cap * dims * 2); - s.ple_part = try s.add(gpa, d, 4 * cap * dims * 4); - s.ple_gated = try s.add(gpa, d, cap * wide * 2); - s.ple_pss = try s.add(gpa, d, cap * streams * 4); - s.ple_nrow = try s.add(gpa, d, cap * wide * 2); - s.logits = try s.add(gpa, d, vocab * 2); - s.head_local = try s.add(gpa, d, cap * vocab); + s.pos = if (sizes[1]) |len| try s.add(gpa, d, len) else 0; + s.h = if (sizes[2]) |len| try s.add(gpa, d, len) else 0; + s.pss = if (sizes[3]) |len| try s.add(gpa, d, len) else 0; + s.norm = if (sizes[4]) |len| try s.add(gpa, d, len) else 0; + s.dn = if (sizes[5]) |len| try s.add(gpa, d, len) else 0; + s.hc_part = if (sizes[6]) |len| try s.add(gpa, d, len) else 0; + s.hc_up = if (sizes[7]) |len| try s.add(gpa, d, len) else 0; + s.native_part = if (sizes[8]) |len| try s.add(gpa, d, len) else 0; + s.act = if (sizes[9]) |len| try s.add(gpa, d, len) else 0; + s.xs_act = if (sizes[10]) |len| try s.add(gpa, d, len) else 0; + s.inj = if (sizes[11]) |len| try s.add(gpa, d, len) else 0; + s.mixed = if (sizes[12]) |len| try s.add(gpa, d, len) else 0; + s.xsm = if (sizes[13]) |len| try s.add(gpa, d, len) else 0; + s.proj = if (sizes[14]) |len| try s.add(gpa, d, len) else 0; + s.part = if (sizes[15]) |len| try s.add(gpa, d, len) else 0; + s.gout = if (sizes[16]) |len| try s.add(gpa, d, len) else 0; + s.gxs = if (sizes[17]) |len| try s.add(gpa, d, len) else 0; + s.ys = if (sizes[18]) |len| try s.add(gpa, d, len) else 0; + s.gq = if (sizes[19]) |len| try s.add(gpa, d, len) else 0; + s.gk = if (sizes[20]) |len| try s.add(gpa, d, len) else 0; + s.gv = if (sizes[21]) |len| try s.add(gpa, d, len) else 0; + s.ggb = if (sizes[22]) |len| try s.add(gpa, d, len) else 0; + s.q = if (sizes[23]) |len| try s.add(gpa, d, len) else 0; + s.iq = if (sizes[24]) |len| try s.add(gpa, d, len) else 0; + s.po = if (sizes[25]) |len| try s.add(gpa, d, len) else 0; + s.pm = if (sizes[26]) |len| try s.add(gpa, d, len) else 0; + s.pl = if (sizes[27]) |len| try s.add(gpa, d, len) else 0; + s.qids = if (sizes[28]) |len| try s.add(gpa, d, len) else 0; + s.qnk = if (sizes[29]) |len| try s.add(gpa, d, len) else 0; + s.qspr = if (sizes[30]) |len| try s.add(gpa, d, len) else 0; + s.ao = if (sizes[31]) |len| try s.add(gpa, d, len) else 0; + s.gated = if (sizes[32]) |len| try s.add(gpa, d, len) else 0; + s.gxs_attn = if (sizes[33]) |len| try s.add(gpa, d, len) else 0; + s.router_l = if (sizes[34]) |len| try s.add(gpa, d, len) else 0; + s.picks = if (sizes[35]) |len| try s.add(gpa, d, len) else 0; + s.wts = if (sizes[36]) |len| try s.add(gpa, d, len) else 0; + s.members = if (sizes[37]) |len| try s.add(gpa, d, len) else 0; + s.items = if (sizes[38]) |len| try s.add(gpa, d, len) else 0; + s.counts = if (sizes[39]) |len| try s.add(gpa, d, len) else 0; + s.rank = if (sizes[40]) |len| try s.add(gpa, d, len) else 0; + s.hist = if (sizes[41]) |len| try s.add(gpa, d, len) else 0; + s.shared_gate = if (sizes[42]) |len| try s.add(gpa, d, len) else 0; + s.shared_up = if (sizes[43]) |len| try s.add(gpa, d, len) else 0; + s.shared_act = if (sizes[44]) |len| try s.add(gpa, d, len) else 0; + s.shared_out = if (sizes[45]) |len| try s.add(gpa, d, len) else 0; + s.moe_act = if (sizes[46]) |len| try s.add(gpa, d, len) else 0; + s.y = if (sizes[47]) |len| try s.add(gpa, d, len) else 0; + s.branch = if (sizes[48]) |len| try s.add(gpa, d, len) else 0; + s.moe_f32 = if (sizes[49]) |len| try s.add(gpa, d, len) else 0; + s.pocket_act = if (sizes[50]) |len| try s.add(gpa, d, len) else 0; + s.ple_emb = if (sizes[51]) |len| try s.add(gpa, d, len) else 0; + s.ple_xs = if (sizes[52]) |len| try s.add(gpa, d, len) else 0; + s.ple_keys = if (sizes[53]) |len| try s.add(gpa, d, len) else 0; + s.ple_vals = if (sizes[54]) |len| try s.add(gpa, d, len) else 0; + s.ple_part = if (sizes[55]) |len| try s.add(gpa, d, len) else 0; + s.ple_gated = if (sizes[56]) |len| try s.add(gpa, d, len) else 0; + s.ple_pss = if (sizes[57]) |len| try s.add(gpa, d, len) else 0; + s.ple_nrow = if (sizes[58]) |len| try s.add(gpa, d, len) else 0; + s.logits = if (sizes[59]) |len| try s.add(gpa, d, len) else 0; + s.head_local = if (sizes[60]) |len| try s.add(gpa, d, len) else 0; var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); - s.inv = try s.add(gpa, d, half * 4); + s.inv = if (sizes[61]) |len| try s.add(gpa, d, len) else 0; try s.bufs.items[s.bufs.items.len - 1].uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); - s.dummy = try s.add(gpa, d, 4096); + s.dummy = if (sizes[62]) |len| try s.add(gpa, d, len) else 0; try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); try stream.synchronize(); s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); @@ -484,6 +558,8 @@ pub const Scratch = struct { /// A sequence's state the pass reads and leaves as the row-by-row path would. pub const SeqState = struct { media: ?@import("cuda_media.zig").View = null, + /// An int8 K/V cache (kc/vc hold its codes; cuda_exl3_kv.zig); null: bf16 kc/vc. + kv8: ?@import("cuda_exl3_kv.zig").SeqKv8 = null, conv: []cuda.DeviceBuffer, state: []cuda.DeviceBuffer, kc: []cuda.DeviceBuffer, @@ -497,6 +573,8 @@ pub const SeqState = struct { cap: usize, /// Block scores for an attention block's rows [min(rows, att_rows), scoreStride(cap)] fp32 (the engine's). scores: u64, + /// The K/V paged to host memory with a device block cache (cuda_exl3_kvpage.zig), else kc/vc hold it all. + paged: ?*@import("cuda_exl3_kvpage.zig").Paged = null, }; /// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). @@ -560,4 +638,5 @@ pub const Gather = struct { test { _ = @import("cuda_batch_state_test.zig"); + _ = @import("cuda_scratch_layout_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index c66d20c2a4..709159df0a 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -10,6 +10,64 @@ const Advanced = struct { ncclCommGetAsyncError: *const fn (nccl.Comm, *nccl.Result) callconv(.c) nccl.Result, }; +/// How long a wait polls before it sleeps between polls (a decode step's work ends within a few milliseconds; each +/// 50 us sleep can add up to ~100 us to the step). +const spin_ns: u64 = 3 * std.time.ns_per_ms; + +pub const GraphState = enum(u32) { eager, capture, pending, ready, failed, timed_out }; +pub const GraphChoice = enum { eager, capture, replay }; +pub const GraphVote = struct { state: GraphState, cohort: u64 = 0 }; +pub const GraphDecision = struct { choice: GraphChoice, cohort: u64 }; + +pub fn selectGraphVictim(local: u64, peer: u64) u64 { + if (local == 0) return peer; + if (peer == 0) return local; + return @min(local, peer); +} + +test "exl3: selective retirement agrees stable union victims for asymmetric cohort sets" { + var left = [_]u64{ 3, 7, 11 }; + var right = [_]u64{ 5, 7, 11 }; + for ([_]u64{ 3, 5 }) |expected| { + var local: u64 = 0; + var peer: u64 = 0; + for (left) |cohort| local = selectGraphVictim(local, cohort); + for (right) |cohort| peer = selectGraphVictim(peer, cohort); + const victim = selectGraphVictim(local, peer); + try std.testing.expectEqual(expected, victim); + try std.testing.expectEqual(victim, selectGraphVictim(peer, local)); + for (&left) |*cohort| if (cohort.* == victim) { + cohort.* = 0; + }; + for (&right) |*cohort| if (cohort.* == victim) { + cohort.* = 0; + }; + } + try std.testing.expectEqual(@as(u64, 7), left[1]); + try std.testing.expectEqual(@as(u64, 11), right[2]); + try std.testing.expectEqual(@as(u64, 0), selectGraphVictim(0, 0)); +} + +fn chooseGraph(local: GraphVote, peer: GraphVote, step: u64) !GraphDecision { + if (local.state == .timed_out or peer.state == .timed_out) return error.GraphInstantiationTimeout; + if (local.state == .failed or peer.state == .failed) return error.GraphInstantiationFailed; + if (local.state == .eager or peer.state == .eager or local.state == .pending or peer.state == .pending) + return .{ .choice = .eager, .cohort = 0 }; + if (local.state == .ready and peer.state == .ready and local.cohort != 0 and local.cohort == peer.cohort) + return .{ .choice = .replay, .cohort = local.cohort }; + return .{ .choice = .capture, .cohort = step }; +} + +fn graphExchange(link: wire.Link, step: u64, signature: u64, vote: GraphVote, until: u64) !GraphDecision { + const magic: u32 = 0x47520000; + try link.sendHeader(.{ .kind = .ready, .sequence = step, .object = vote.cohort, .arg = signature, .flags = magic | @backingInt(vote.state) }, until); + const peer = try link.receiveHeader(until); + if (peer.kind != .ready or peer.sequence != step or peer.arg != signature or peer.bytes != 0 or peer.flags & 0xffff0000 != magic) + return error.GraphRendezvousMismatch; + const state = std.enums.fromInt(GraphState, peer.flags & 0xffff) orelse return error.GraphRendezvousMismatch; + return chooseGraph(vote, .{ .state = state, .cohort = peer.object }, step); +} + pub const Collective = struct { lib: nccl.Library, api: Advanced, @@ -17,6 +75,7 @@ pub const Collective = struct { link: wire.Link, timeout_ms: u64, failed: bool = false, + graph_step: u64 = 0, /// `link` has completed the source/settings handshake; it stays owned by the caller. pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { @@ -33,6 +92,11 @@ pub const Collective = struct { try link.send(std.mem.asBytes(&id), until); } else try link.receive(std.mem.asBytes(&id), until); var config: nccl.Config = .{}; + // TF_FLASHNEXT_NCCL_BLOCKING=1: a blocking communicator (NCCL enqueues on the calling thread; a nonblocking + // one hands each collective to its own thread, ~100 us of host time an eager collective) + if (std.c.getenv("TF_FLASHNEXT_NCCL_BLOCKING")) |v| if (std.mem.eql(u8, std.mem.span(v), "1")) { + config.blocking = 1; + }; errdefer c.abort(); try c.submitted(c.api.ncclCommInitRankConfig(&c.comm, 2, id, @intCast(rank), &config)); try c.wait(null, until); @@ -60,6 +124,8 @@ pub const Collective = struct { fn wait(c: *Collective, stream: ?cuda.Stream, until: u64) !void { errdefer c.abort(); + var polls: u32 = 0; + var spin_until: u64 = 0; while (true) { try c.link.checkPeer(); if (wire.now() >= until) return error.RankTimeout; @@ -71,7 +137,11 @@ pub const Collective = struct { if (try s.done()) return; } else return; } - wire.pause(); + // a nonblocking enqueue lands within tens of microseconds: poll that long before sleeping (each sleep + // of 50 us cost every eager or captured collective ~100 us of host time) + polls += 1; + if (polls == 1) spin_until = wire.now() + spin_ns; + if (wire.now() < spin_until) std.atomic.spinLoopHint() else wire.pause(); } } @@ -81,6 +151,34 @@ pub const Collective = struct { try c.wait(stream, wire.deadline(c.timeout_ms)); } + pub fn failGraph(c: *Collective) void { + c.abort(); + _ = std.c.shutdown(c.link.fd, std.posix.SHUT.RDWR); + } + + pub fn graphTimeoutMs(c: *const Collective) u64 { + return @min(c.timeout_ms, 30_000); + } + + pub fn graphAgree(c: *Collective, signature: u64, vote: GraphVote) !GraphDecision { + if (c.failed or c.comm == null) return error.RankDisconnected; + errdefer c.failGraph(); + c.graph_step = try std.math.add(u64, c.graph_step, 1); + return graphExchange(c.link, c.graph_step, signature, vote, wire.deadline(c.graphTimeoutMs())); + } + + pub fn graphVictim(c: *Collective, signature: u64, candidate: u64) !u64 { + if (c.failed or c.comm == null) return error.RankDisconnected; + errdefer c.failGraph(); + c.graph_step = try std.math.add(u64, c.graph_step, 1); + const until = wire.deadline(c.graphTimeoutMs()); + try c.link.sendHeader(.{ .kind = .ready, .sequence = c.graph_step, .object = candidate, .arg = signature, .flags = 0x47560000 }, until); + const peer = try c.link.receiveHeader(until); + if (peer.kind != .ready or peer.sequence != c.graph_step or peer.arg != signature or peer.bytes != 0 or peer.flags != 0x47560000) + return error.GraphRendezvousMismatch; + return selectGraphVictim(candidate, peer.object); + } + /// Disjoint BF16 vocabulary columns are concatenated without an arithmetic reduction. pub fn gather(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; @@ -109,3 +207,116 @@ pub const Collective = struct { try c.wait(null, wire.deadline(c.timeout_ms)); } }; + +test "exl3: graph readiness requires matching completed capture cohorts" { + const states = std.enums.values(GraphState); + for (states) |local| for (states) |peer| { + const left: GraphVote = .{ .state = local, .cohort = 7 }; + const right: GraphVote = .{ .state = peer, .cohort = 7 }; + if (local == .timed_out or peer == .timed_out) { + try std.testing.expectError(error.GraphInstantiationTimeout, chooseGraph(left, right, 10)); + } else if (local == .failed or peer == .failed) { + try std.testing.expectError(error.GraphInstantiationFailed, chooseGraph(left, right, 10)); + } else { + const choice: GraphChoice = if (local == .eager or peer == .eager or local == .pending or peer == .pending) + .eager + else if (local == .ready and peer == .ready) .replay else .capture; + const decision = try chooseGraph(left, right, 10); + try std.testing.expectEqual(choice, decision.choice); + try std.testing.expectEqualDeep(decision, try chooseGraph(right, left, 10)); + } + }; + try std.testing.expectEqual(GraphChoice.capture, (try chooseGraph(.{ .state = .ready, .cohort = 7 }, .{ .state = .ready, .cohort = 8 }, 10)).choice); +} + +fn graphTestPair() ![2]std.posix.socket_t { + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + errdefer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + return pair; +} + +test "exl3: graph rendezvous keeps a delayed rank eager until both workers finish" { + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const Peer = struct { + fn run(link: wire.Link) void { + const votes = [_]GraphVote{ .{ .state = .pending, .cohort = 3 }, .{ .state = .ready, .cohort = 3 }, .{ .state = .ready, .cohort = 3 } }; + for (votes, 1..) |vote, step| { + const decision = graphExchange(link, step, 42, vote, wire.deadline(1000)) catch @panic("peer rendezvous failed"); + std.debug.assert(decision.choice == (if (step < 3) GraphChoice.eager else GraphChoice.replay)); + } + } + }; + const thread = try std.Thread.spawn(.{}, Peer.run, .{wire.Link{ .fd = pair[1] }}); + defer thread.join(); + const votes = [_]GraphVote{ .{ .state = .ready, .cohort = 3 }, .{ .state = .pending, .cohort = 3 }, .{ .state = .ready, .cohort = 3 } }; + for (votes, 1..) |vote, step| { + const decision = try graphExchange(.{ .fd = pair[0] }, step, 42, vote, wire.deadline(1000)); + try std.testing.expectEqual(if (step < 3) GraphChoice.eager else GraphChoice.replay, decision.choice); + } +} + +test "exl3: graph rendezvous rejects stale frames failures timeouts and disconnected peers" { + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const local: wire.Link = .{ .fd = pair[0] }; + const peer: wire.Link = .{ .fd = pair[1] }; + const vote: GraphVote = .{ .state = .ready, .cohort = 2 }; + for ([_]wire.Header{ + .{ .kind = .ack, .sequence = 1, .arg = 42, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 2, .arg = 42, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 1, .arg = 43, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 1, .arg = 42, .flags = 0x475200ff }, + }) |header| { + try peer.sendHeader(header, wire.deadline(1000)); + try std.testing.expectError(error.GraphRendezvousMismatch, graphExchange(local, 1, 42, vote, wire.deadline(1000))); + _ = try peer.receiveHeader(wire.deadline(1000)); + } + try peer.sendHeader(.{ .kind = .ready, .sequence = 1, .arg = 42, .flags = 0x47520004 }, wire.deadline(1000)); + try std.testing.expectError(error.GraphInstantiationFailed, graphExchange(local, 1, 42, vote, wire.deadline(1000))); + _ = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectError(error.RankTimeout, graphExchange(local, 1, 42, vote, wire.deadline(5))); + _ = std.c.shutdown(peer.fd, std.posix.SHUT.RDWR); + try std.testing.expectError(error.RankDisconnected, graphExchange(local, 1, 42, vote, wire.deadline(1000))); +} + +test "exl3: graph rendezvous failure aborts the communicator and disconnects its control peer" { + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const Abort = struct { + var calls: usize = 0; + fn run(_: nccl.Comm) callconv(.c) nccl.Result { + calls += 1; + return 0; + } + }; + Abort.calls = 0; + var c: Collective = undefined; + c.comm = @ptrFromInt(1); + c.api.ncclCommAbort = Abort.run; + c.link = .{ .fd = pair[0] }; + c.failed = false; + c.timeout_ms = 5; + c.graph_step = 0; + try std.testing.expectError(error.RankTimeout, c.graphAgree(42, .{ .state = .pending })); + try std.testing.expect(c.failed and c.comm == null); + try std.testing.expectEqual(@as(usize, 1), Abort.calls); + try std.testing.expectError(error.RankDisconnected, c.graphAgree(42, .{ .state = .ready })); + const peer: wire.Link = .{ .fd = pair[1] }; + _ = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectError(error.RankDisconnected, peer.checkPeer()); +} diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 41b4099220..7dd3f047d9 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -3,6 +3,7 @@ const std = @import("std"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3 = @import("cuda_exl3_config.zig"); /// One config field and the value the kernels and buffers assume. const Want = struct { name: []const u8, value: i64 }; @@ -100,7 +101,7 @@ pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u return geometry(text); } -fn geometry(text: std.json.ObjectMap) ?[]const u8 { +pub fn geometry(text: std.json.ObjectMap) ?[]const u8 { // a folder Living Weights taught widens the shared expert by its pocket (a whole number of groups, 128 at most) const learned = int(text, learned_key) orelse 0; if (learned < 0 or learned > 128 or @mod(learned, 32) != 0) return learned_key; @@ -111,7 +112,7 @@ fn geometry(text: std.json.ObjectMap) ?[]const u8 { return semantics(text); } -pub const Mode = enum { affine, modelopt }; +pub const Mode = enum { affine, modelopt, exl3 }; fn projection(q: std.json.ObjectMap, name: []const u8, algo: []const u8, group: i64) bool { const p = object(q, name) orelse return false; @@ -201,8 +202,8 @@ pub fn mode(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Mode { const root = parsed.value.object; const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; - const selected: Mode = if (quant != null) .affine else .modelopt; - var field = if (selected == .affine) mismatch(text, quant) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; + const selected: Mode = if (quant != null) .affine else if (exl3.named(root)) .exl3 else .modelopt; + var field = if (selected == .affine) mismatch(text, quant) else if (selected == .exl3) exl3.mismatch(root, text, geometry) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { field = "tie_word_embeddings"; }; diff --git a/zig/src/families/flashnext/cuda_exl3.zig b/zig/src/families/flashnext/cuda_exl3.zig new file mode 100644 index 0000000000..1b9782685f --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3.zig @@ -0,0 +1,45 @@ +//! EXL3 checkpoints on the Flash Next CUDA family: what the upstream files' one-line hooks call. +pub const config = @import("cuda_exl3_config.zig"); +pub const store = @import("cuda_exl3_store.zig"); +pub const memory = @import("cuda_exl3_memory.zig"); +pub const load = @import("cuda_exl3_load.zig"); +pub const pass = @import("cuda_exl3_pass.zig"); +pub const tp = @import("cuda_exl3_tp.zig"); +pub const kv = @import("cuda_exl3_kv.zig"); +pub const kvpage = @import("cuda_exl3_kvpage.zig"); +pub const prefill_k = @import("cuda_exl3_prefill_k.zig"); +pub const prefill = @import("cuda_exl3_prefill.zig"); +pub const dplan = @import("cuda_exl3_dplan.zig"); +pub const graph = @import("cuda_exl3_graph.zig"); +pub const fork = @import("cuda_exl3_fork.zig"); +pub const chain = @import("cuda_exl3_chain.zig"); +pub const recycle = @import("cuda_exl3_recycle.zig"); +pub const telemetry = @import("cuda_exl3_telemetry.zig"); +pub const frequency = @import("cuda_exl3_frequency.zig"); +/// The upstream modules an EXL3 check drives (tf-cuda-test flashnext-exl3): the store, loaders and the batched pass. +pub const prompt_wt = @import("cuda_prompt_wt.zig"); +pub const prompt_load = @import("cuda_prompt_load.zig"); +pub const projection_load = @import("cuda_projection_load.zig"); +pub const batch = @import("cuda_batch.zig"); + +test { + _ = config; + _ = store; + _ = memory; + _ = load; + _ = pass; + _ = tp; + _ = kv; + _ = kvpage; + _ = prefill_k; + _ = prefill; + _ = @import("cuda_exl3_prefill_ring.zig"); + _ = dplan; + _ = graph; + _ = fork; + _ = chain; + _ = recycle; + _ = telemetry; + _ = frequency; + _ = @import("cuda_exl3_frequency_state.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_chain.zig b/zig/src/families/flashnext/cuda_exl3_chain.zig new file mode 100644 index 0000000000..a243749781 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_chain.zig @@ -0,0 +1,586 @@ +//! The MTP draft chain as one CUDA graph (TF_FLASHNEXT_EXL3_GRAPH=1, greedy drafts). +//! +//! Upstream's chain (cuda_mtp_ops.Ops.draft) picks a draft on the host after every head step: fn_row_best, a blocking +//! read of four floats, then the next step uploads that token. Here the whole chain is captured once per shape: each +//! pick's record stays on the device, fn_exl3_draft_token turns its column into the next step's input token, and the +//! head step (the same launches as forwardHead for one row, its experts on the device planner) runs on it; the host +//! reads the chain's records once and applies upstream's confidence rule to them. Every step past an early stop still +//! runs (its cache entries are past the chain and are trimmed by the next absorb, as drafted entries always are), so +//! the drafts and their probabilities are upstream's bit for bit (tf-cuda-test flashnext-exl3-graph and the serve +//! gate's drafted == undrafted). +//! +//! `enter`/`leave` (cuda_mtp_ops.forwardHead's hook) also put every head step of up to the planner's rows, eager or +//! captured, on the device planner, so the head's pool stays device-owned between decode steps. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); +const dplan = @import("cuda_exl3_dplan.zig"); +const graph = @import("cuda_exl3_graph.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const eps = wt.eps; +const out_sk = wt.out_sk; +const max_chain = mtp.max_chain; +const fc_sk = mtp.fc_sk; +const piece_rows = mtp.piece_rows; + +/// forwardHead's hook: the head's pool on the device planner for a step of `rows` rows; returns what to restore. +pub fn enter(e: *prompt.Engine, rows: usize) !?*dplan.DevPlan { + const was = dplan.active; + if (!graph.enabled() or rows == 0) return was; + const dp = (try dplan.forLayers(e.gpa, e.driver, e, e.layers[0..e.n_layers])) orelse return was; + const mp = dplan.mtp_pool orelse return was; + if (rows > dp.max_rows or dp.indexOf(mp) == null) return was; + try dp.toDevice(e.stream.*); + dplan.active = dp; + return was; +} + +pub fn leave(was: ?*dplan.DevPlan) void { + dplan.active = was; +} + +const Key = struct { + picks: usize, + heads: usize, + /// Attention seams of the first and the last step (chunk counts and score tiles move with the head's length). + seams: [4]u32, + head: usize, + kc: u64, + out: u64, + dlogits: u64, + cap: usize, + scratch: u64, + scratch_cap: usize, + scores: u64, + plan: u64, + confidence: u64, +}; + +const Entry = struct { key: Key, exec: cuda.graph.Exec = undefined, used: u64, pending: ?*graph.Pending = null, admission: ?graph.memory.Ticket = null, stream: ?cuda.Stream = null }; + +const State = struct { + head: usize, + ids: cuda.DeviceBuffer, + token: cuda.DeviceBuffer, + best: cuda.DeviceBuffer, + pos_host: cuda.HostBuffer, + best_host: cuda.HostBuffer, + module: cuda.Module, + token_fn: cuda.Function, + next_fn: cuda.Function, + /// [0] the chain's running decision, [1 + j] the decision after pick j (read back with the records). + alive: cuda.DeviceBuffer, + alive_host: cuda.HostBuffer, + /// Where a step's conditional body is captured (cuStreamBeginCaptureToGraph). + side: cuda.Stream, + cond: Cond, +}; + +/// The conditional-node driver entry points (CUDA 12.4+), looked up from libcuda. +const NodeParams = extern struct { + type: c_int, + reserved0: [3]c_int = .{ 0, 0, 0 }, + u: extern union { cond: CondParams, pad: [29]i64 }, + reserved2: i64 = 0, +}; +const CondParams = extern struct { handle: u64, type: c_int, size: c_uint, graphs: ?[*]?*anyopaque, ctx: ?*anyopaque }; +const Cond = struct { + lib: std.DynLib, + handle_create: *const fn (*u64, ?*anyopaque, ?*anyopaque, c_uint, c_uint) callconv(.c) c_int, + add_node: *const fn (*?*anyopaque, ?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, *NodeParams) callconv(.c) c_int, + capture_info: *const fn (?*anyopaque, *c_int, *u64, *?*anyopaque, *?[*]const ?*anyopaque, *?*const anyopaque, *usize) callconv(.c) c_int, + update_deps: *const fn (?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, c_uint) callconv(.c) c_int, + begin_to_graph: *const fn (?*anyopaque, ?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, c_int) callconv(.c) c_int, + ctx_current: *const fn (*?*anyopaque) callconv(.c) c_int, + + fn open() !Cond { + var lib = try std.DynLib.open("libcuda.so.1"); + errdefer lib.close(); + var c: Cond = undefined; + c.lib = lib; + c.handle_create = lib.lookup(@TypeOf(c.handle_create), "cuGraphConditionalHandleCreate") orelse return error.MissingSymbol; + c.add_node = lib.lookup(@TypeOf(c.add_node), "cuGraphAddNode_v2") orelse return error.MissingSymbol; + c.capture_info = lib.lookup(@TypeOf(c.capture_info), "cuStreamGetCaptureInfo_v3") orelse return error.MissingSymbol; + c.update_deps = lib.lookup(@TypeOf(c.update_deps), "cuStreamUpdateCaptureDependencies_v2") orelse return error.MissingSymbol; + c.begin_to_graph = lib.lookup(@TypeOf(c.begin_to_graph), "cuStreamBeginCaptureToGraph") orelse return error.MissingSymbol; + c.ctx_current = lib.lookup(@TypeOf(c.ctx_current), "cuCtxGetCurrent") orelse return error.MissingSymbol; + return c; + } + + fn ok(r: c_int, what: []const u8) !void { + if (r == 0) return; + std.log.err("flashnext exl3 chain: {s}: CUDA error {d}", .{ what, r }); + return error.CudaFailed; + } + + /// The capturing graph of `s` (main capture) and a new conditional handle on it. + fn handle(c: *const Cond, s: cuda.Stream) !u64 { + var status: c_int = 0; + var id: u64 = 0; + var g: ?*anyopaque = null; + var deps: ?[*]const ?*anyopaque = null; + var edges: ?*const anyopaque = null; + var n: usize = 0; + try ok(c.capture_info(@ptrCast(s.handle), &status, &id, &g, &deps, &edges, &n), "cuStreamGetCaptureInfo"); + var ctx: ?*anyopaque = null; + try ok(c.ctx_current(&ctx), "cuCtxGetCurrent"); + var h: u64 = 0; + try ok(c.handle_create(&h, g, ctx, 0, 1), "cuGraphConditionalHandleCreate"); + return h; + } + + /// An IF node on `h` after `s`'s captured work; `s` continues after it; returns its body graph. + fn ifNode(c: *const Cond, s: cuda.Stream, h: u64) !?*anyopaque { + var status: c_int = 0; + var id: u64 = 0; + var g: ?*anyopaque = null; + var deps: ?[*]const ?*anyopaque = null; + var edges: ?*const anyopaque = null; + var n: usize = 0; + try ok(c.capture_info(@ptrCast(s.handle), &status, &id, &g, &deps, &edges, &n), "cuStreamGetCaptureInfo"); + var ctx: ?*anyopaque = null; + try ok(c.ctx_current(&ctx), "cuCtxGetCurrent"); + var params: NodeParams = .{ .type = 13, .u = .{ .pad = @splat(0) } }; + params.u.cond = .{ .handle = h, .type = 0, .size = 1, .graphs = null, .ctx = ctx }; + var node: ?*anyopaque = null; + try ok(c.add_node(&node, g, deps, null, n, ¶ms), "cuGraphAddNode(conditional)"); + const one = [1]?*anyopaque{node}; + try ok(c.update_deps(@ptrCast(s.handle), &one, null, 1, 1), "cuStreamUpdateCaptureDependencies"); + return params.u.cond.graphs.?[0]; + } +}; + +var state: ?State = null; +var entries: [16]?Entry = @splat(null); +var clock: u64 = 0; +var broken = false; + +pub var stats: struct { chains: u64 = 0, replays: u64 = 0, captures: u64 = 0, updates: u64 = 0, eager: u64 = 0 } = .{}; + +fn stateFor(e: *prompt.Engine, h: *const mtp.Head) !*State { + if (state) |*s| if (s.head == @intFromPtr(h)) return s; + if (state) |*s| release(s); + var st: State = undefined; + st.head = @intFromPtr(h); + st.ids = try cuda.DeviceBuffer.fromHost(e.driver, std.mem.sliceAsBytes(h.ids)); + errdefer st.ids.free(); + st.token = try cuda.DeviceBuffer.alloc(e.driver, 256); + errdefer st.token.free(); + st.best = try cuda.DeviceBuffer.alloc(e.driver, max_chain * 16); + errdefer st.best.free(); + st.pos_host = try cuda.HostBuffer.alloc(e.driver, max_chain * 16); + errdefer st.pos_host.free(); + st.best_host = try cuda.HostBuffer.alloc(e.driver, max_chain * 16); + errdefer st.best_host.free(); + st.module = try cuda.Module.load(e.driver, cuda.kernels.flashnext_exl3_dplan); + errdefer st.module.unload(); + st.token_fn = try st.module.function("fn_exl3_draft_token"); + st.next_fn = try st.module.function("fn_exl3_draft_next"); + st.alive = try cuda.DeviceBuffer.alloc(e.driver, (1 + max_chain) * 4 + 256); + errdefer st.alive.free(); + st.alive_host = try cuda.HostBuffer.alloc(e.driver, (1 + max_chain) * 4); + errdefer st.alive_host.free(); + st.side = try cuda.Stream.init(e.driver, true); + errdefer st.side.deinit(); + st.cond = try Cond.open(); + state = st; + return &state.?; +} + +fn release(s: *State) void { + for (&entries) |*en| if (en.*) |*x| { + dropEntry(x); + en.* = null; + }; + s.cond.lib.close(); + s.side.deinit(); + s.alive_host.free(); + s.alive.free(); + s.module.unload(); + s.best_host.free(); + s.pos_host.free(); + s.best.free(); + s.token.free(); + s.ids.free(); + state = null; +} + +fn seams(end: usize, chunks: *u32, tiles: *u32) void { + chunks.* = @intCast((@min(end, wt.dense_keys) + 511) / 512); + tiles.* = if (end / wt.index_ratio > wt.index_top) @intCast(((end + 3) / 4 + 63) / 64) else 0; +} + +/// One head step for one row whose token is on the device: forwardHead's launches for `next.len == 1` (streams from +/// d.out, scores into d.dlogits), its position from the pinned row `j`. +fn step(comptime Tri: type, o: anytype, tri: Tri, p: *batch.Pass(Tri), d: *mtp.Draft, st: *State, j: usize, len0: usize) !void { + const e = o.e; + const h = o.h; + const s = p.s; + const stream = p.stream; + p.rows = 1; + p.pos0 = len0 + j; + p.pending = null; + try stream.d.check(stream.d.api.cuMemcpyDtoDAsync_v2(s.ids, st.token.ptr, 4, stream.handle), "cuMemcpyDtoDAsync"); + try graph.pinnedCopy(stream, s.pos, st.pos_host.bytes[j * 16 ..][0..16]); + if (e.table.dense) { + try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, 1); + } else { + try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, 1); + } + try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, 1, dims, dims); + try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, 1, fc_sk); + try p.sumSlices(s.part, s.branch, true, dims, fc_sk); + try tri.rmsnorm(d.out.ptr, h.norm_h.ptr, s.norm, s.gq, eps, wide, 1, wide, wide); + try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams, fc_sk); + try p.sumSlices(s.part, s.ple_gated, true, streams * dims, fc_sk); + try tri.addStreams(s.branch, s.ple_gated, s.h, 1, dims, streams); + try p.writeback(0, 1); + try p.readout(h.layer.attn_hc, 0, 1); + try p.attnBlock(&h.layer, 0); + try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + try p.readout(h.layer.mlp_hc, 0, 1); + try p.moeBlock(&h.layer); + p.pending = 0; + try p.writeback(0, 1); + try p.copy(d.out.ptr, s.h, wide * 2); + try p.readout(h.mixer, 0, 1); + try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); +} + +/// The chain's launches on `p`'s stream: the decision reset, every pick's record into st.best; after a pick that +/// continues, its token and decision (fn_exl3_draft_next) and a conditional node whose body (captured on st.side +/// through `ps`, a pass on that stream) is the next head step: a chain that stops skips its remaining steps. +fn chainBody(comptime Tri: type, o: anytype, p: *batch.Pass(Tri), ps: *batch.Pass(Tri), d: *mtp.Draft, st: *State, picks: usize, heads: usize, len0: usize, confidence: f64) !void { + const h = o.h; + const stream = p.stream; + try stream.d.check(stream.d.api.cuMemsetD32Async(st.alive.ptr, 1, 1 + max_chain, stream.handle), "cuMemsetD32Async"); + for (0..picks) |j| { + try h.best(stream, d.dlogits.ptr, h.ids.len, h.ids.len, 1, st.best.ptr + j * 16); + if (j >= heads) continue; + const handle = try st.cond.handle(stream); + var a: cuda.Args = .{}; + a.add(st.best.ptr + j * 16); + a.add(st.ids.ptr); + a.add(st.token.ptr); + a.add(st.alive.ptr); + a.add(st.alive.ptr + (1 + j) * 4); + a.add(handle); + a.add(confidence); + try cuda.launch.launch(st.next_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, stream, &a); + const body = try st.cond.ifNode(stream, handle); + try Cond.ok(st.cond.begin_to_graph(@ptrCast(st.side.handle), body, null, null, 0, 1), "cuStreamBeginCaptureToGraph"); + step(Tri, o, ps.tri, ps, d, st, j, len0) catch |err| { + var discarded: cuda.abi.Graph = null; + _ = stream.d.api.cuStreamEndCapture(st.side.handle, &discarded); + return err; + }; + var g: cuda.abi.Graph = null; + try stream.d.check(stream.d.api.cuStreamEndCapture(st.side.handle, &g), "cuStreamEndCapture(body)"); + } +} + +/// Ops.draft's hook (cuda_mtp_ops.zig): the chain for greedy drafts as one graph; false leaves it to upstream. +pub fn draft(comptime Tri: type, o: anytype, d: *mtp.Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?mtp.Sampling, confidence: f64, out: *mtp.Chain) !bool { + _ = position; + if (graph.enabled() and broken) return error.GraphCaptureFailed; + if (!graph.enabled() or count < 2 or d.media != null) return false; + if (sm) |s| if (s.temperature > 0) return false; + const e = o.e; + const dp = (try dplan.forLayers(e.gpa, e.driver, e, e.layers[0..e.n_layers])) orelse return false; + const mp = dplan.mtp_pool orelse return false; + if (dp.indexOf(mp) == null) return false; + const picks: usize = @min(count, max_chain); + // upstream runs a head step after pick j while j + 1 < count + const heads: usize = if (count > max_chain) max_chain else count - 1; + if (heads > picks) return false; + out.n = 0; + try o.absorb(d, next, rows, true); + if (d.len + heads > d.cap) { + // upstream's loop, which reports PromptTooLong where it would + stats.eager += 1; + return eagerChain(o, d, count, picks, confidence, out); + } + const st = try stateFor(e, o.h); + const len0 = d.len; + const pos: []i32 = @alignCast(std.mem.bytesAsSlice(i32, st.pos_host.bytes[0 .. max_chain * 16])); + @memset(pos, 0); + for (0..heads) |j| pos[4 * j] = @intCast(len0 + j); + const s = try e.reserve(@max(@as(usize, 32), 2)); + const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = o.tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&o.h.layer)[0..1], .mixer = o.h.mixer, .head = o.h.draft, .ngram = &e.ngram }, + .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .hooks = .{ .logits = .none }, + }; + var key: Key = .{ .picks = picks, .heads = heads, .seams = undefined, .head = @intFromPtr(o.h), .kc = d.kc.ptr, .out = d.out.ptr, .dlogits = d.dlogits.ptr, .cap = d.cap, .scratch = s.ids, .scratch_cap = s.cap, .scores = score_rows, .plan = dp.state.ptr, .confidence = @bitCast(confidence) }; + seams(len0 + 1, &key.seams[0], &key.seams[1]); + seams(len0 + @max(heads, 1), &key.seams[2], &key.seams[3]); + try dp.toDevice(p.stream); + const was = dplan.active; + dplan.active = dp; + defer dplan.active = was; + stats.chains += 1; + var tri_side = o.tri; + tri_side.s = st.side; + var ps = p; + ps.stream = st.side; + ps.tri = tri_side; + const launched = run(Tri, key, o, &p, &ps, d, st, picks, heads, len0, confidence) catch |err| { + std.log.err("flashnext exl3 chain: the draft chain failed ({t})", .{err}); + broken = true; + return err; + }; + if (!launched) { + stats.eager += 1; + return eagerChain(o, d, count, picks, confidence, out); + } + // the chain's records, once + const stream = p.stream; + try stream.d.check(stream.d.api.cuMemcpyDtoHAsync_v2(st.best_host.bytes.ptr, st.best.ptr, picks * 16, stream.handle), "cuMemcpyDtoHAsync"); + try stream.d.check(stream.d.api.cuMemcpyDtoHAsync_v2(st.alive_host.bytes.ptr, st.alive.ptr, (1 + max_chain) * 4, stream.handle), "cuMemcpyDtoHAsync"); + try stream.synchronize(); + const alive: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, st.alive_host.bytes[0 .. (1 + max_chain) * 4])); + d.len += heads; + d.drafted += heads; + const rec: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, st.best_host.bytes[0 .. picks * 16])); + for (0..picks) |j| { + const top = rec[4 * j]; + const lse = rec[4 * j + 1]; + const col: usize = @intCast(@as(i32, @bitCast(rec[4 * j + 2]))); + const got: mtp.Pick = .{ .token = o.h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) }; + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + // the device stopped the chain here (its decision is the same rule; a step past it did not run) + if (j < heads and alive[1 + j] == 0) break; + } + return true; +} + +/// upstream's loop (Ops.draft after its absorb), for chains that leave the graph path. +fn eagerChain(o: anytype, d: *mtp.Draft, count: usize, picks: usize, confidence: f64, out: *mtp.Chain) !bool { + var j: usize = 0; + while (j < picks) : (j += 1) { + const got = try o.pick(d, 0, null); + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + if (j + 1 < count) { + try o.forwardHead(d, &.{got.token}, d.out.ptr, true); + d.drafted += 1; + } + } + return true; +} + +fn sameShape(a: Key, b: Key) bool { + var x = a; + x.kc = b.kc; + x.out = b.out; + x.dlogits = b.dlogits; + return std.meta.eql(x, b); +} + +/// Replays the chain's graph (true), or captures it; a new graph instantiates on a worker thread and this chain runs +/// eagerly (false), as does every chain until it landed. +fn run(comptime Tri: type, key: Key, o: anytype, p: *batch.Pass(Tri), ps: *batch.Pass(Tri), d: *mtp.Draft, st: *State, picks: usize, heads: usize, len0: usize, confidence: f64) !bool { + clock += 1; + for (&entries) |*en| if (en.*) |*x| if (std.meta.eql(x.key, key)) { + x.used = clock; + if (x.pending) |pd| { + const got = graph.takePending(pd) catch |err| { + if (x.admission) |ticket| graph.memory.release(ticket) catch unreachable; + en.* = null; + return err; + }; + var exec = got orelse return false; + x.pending = null; + graph.uploadResident(&exec, p.stream) catch |err| { + exec.deinit(); + if (x.admission) |ticket| graph.memory.release(ticket) catch unreachable; + en.* = null; + return err; + }; + x.exec = exec; + if (x.admission) |ticket| try graph.memory.installed(ticket); + } + stats.replays += 1; + try x.exec.launchOn(p.stream); + return true; + }; + try p.stream.synchronize(); + const admission = try graph.reserveMemory(p.stream, .chain, 0); + var admitted = true; + defer if (admitted) graph.memory.release(admission) catch unreachable; + try cuda.graph.beginCapture(p.stream, .thread_local); + graph.capturing = true; + p.prepared = true; + chainBody(Tri, o, p, ps, d, st, picks, heads, len0, confidence) catch |err| { + graph.capturing = false; + if (cuda.graph.endCapture(p.stream)) |g| { + var x = g; + x.deinit(); + } else |_| {} + return err; + }; + graph.capturing = false; + var g = try cuda.graph.endCapture(p.stream); + var owned = true; + defer if (owned) g.deinit(); + stats.captures += 1; + // a stale graph of the same shape (another stream's buffers) is updated in place; else a new one + var oldest: usize = 0; + var oldest_used: u64 = std.math.maxInt(u64); + for (&entries, 0..) |*en, i| { + if (en.* == null) { + oldest = i; + oldest_used = 0; + break; + } + if (en.*.?.used < oldest_used) { + oldest_used = en.*.?.used; + oldest = i; + } + } + for (&entries) |*en| if (en.*) |*x| if (x.pending == null and sameShape(x.key, key) and clock - x.used > 64) { + if ((try x.exec.update(g)) == .success) { + try graph.uploadResident(&x.exec, p.stream); + x.key = key; + x.used = clock; + stats.updates += 1; + try x.exec.launchOn(p.stream); + return true; + } + }; + if (entries[oldest]) |*x| dropEntry(x); + entries[oldest] = null; + if (try graph.startInstantiate(g, false)) |pd| { + owned = false; + entries[oldest] = .{ .key = key, .used = clock, .pending = pd, .admission = admission, .stream = p.stream }; + admitted = false; + return false; + } + try installSync(&entries[oldest], g, .{ .key = key, .used = clock, .admission = admission, .stream = p.stream }); + admitted = false; + try entries[oldest].?.exec.launchOn(p.stream); + return true; +} + +fn installSync(target: *?Entry, template: cuda.graph.Graph, entry: Entry) !void { + var exec = try template.instantiate(); + errdefer exec.deinit(); + try graph.uploadResident(&exec, entry.stream.?); + try graph.memory.installed(entry.admission.?); + target.* = entry; + target.*.?.exec = exec; +} + +fn dropEntry(x: *Entry) void { + if (x.stream) |stream| stream.synchronize() catch @panic("chain retirement synchronization failed"); + defer if (x.admission) |ticket| { + graph.memory.release(ticket) catch unreachable; + x.admission = null; + }; + if (x.pending) |pd| { + while (true) { + const got = graph.takePending(pd) catch break; + if (got) |e| { + var ex = e; + ex.deinit(); + break; + } + std.Thread.yield() catch {}; + } + return; + } + x.exec.deinit(); +} + +pub fn evictOne() bool { + var oldest: ?usize = null; + for (entries, 0..) |maybe, index| if (maybe) |entry| { + if (oldest == null or entry.used < entries[oldest.?].?.used) oldest = index; + }; + const index = oldest orelse return false; + dropEntry(&entries[index].?); + entries[index] = null; + return true; +} + +test "exl3: synchronous chain launch failure retains ownership until retirement" { + const Mock = struct { + var drops: usize = 0; + + fn instantiate(out: *cuda.abi.GraphExec, _: cuda.abi.Graph, _: u64) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + fn upload(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn synchronize(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn launch(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return cuda.abi.error_not_ready; + } + + fn destroy(_: cuda.abi.GraphExec) callconv(.c) c_int { + drops += 1; + return 0; + } + + fn installAndLaunch(target: *?Entry, template: cuda.graph.Graph, entry: Entry) !void { + try installSync(target, template, entry); + try target.*.?.exec.launchOn(entry.stream.?); + } + }; + Mock.drops = 0; + var driver: cuda.Driver = undefined; + driver.api.cuGraphInstantiateWithFlags = Mock.instantiate; + driver.api.cuGraphUpload = Mock.upload; + driver.api.cuStreamSynchronize = Mock.synchronize; + driver.api.cuGraphLaunch = Mock.launch; + driver.api.cuGraphExecDestroy = Mock.destroy; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(1) }; + try graph.memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const ticket = try graph.memory.admit(.chain, 0, 100); + var ticket_owned = true; + defer if (ticket_owned) graph.memory.release(ticket) catch unreachable; + var entry: ?Entry = null; + defer if (entry) |*cached| dropEntry(cached); + const result = Mock.installAndLaunch(&entry, .{ .d = &driver, .handle = @ptrFromInt(1) }, .{ .key = undefined, .used = 1, .admission = ticket, .stream = stream }); + if (entry != null) ticket_owned = false; + try std.testing.expectError(error.NotReady, result); + try std.testing.expect(entry != null); + try std.testing.expectEqual(@as(usize, 0), Mock.drops); + try std.testing.expectEqual(@as(u64, 30), graph.memory.charged()); + dropEntry(&entry.?); + entry = null; + try std.testing.expectEqual(@as(usize, 1), Mock.drops); + try std.testing.expectEqual(@as(u64, 0), graph.memory.charged()); + try std.testing.expectError(error.StaleGraphTicket, graph.memory.release(ticket)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_config.zig b/zig/src/families/flashnext/cuda_exl3_config.zig new file mode 100644 index 0000000000..ad468aac09 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_config.zig @@ -0,0 +1,116 @@ +//! EXL3 admission for the Flash Next CUDA family: config.json's quantization_config names quant_method "exl3" and the +//! "mul1" codebook. The geometry and semantics are the ones cuda_config.zig fixes for every mode. The bit widths are +//! never read from `bits` (3.05 here is an average): each matrix's width comes from its own trellis (16 * bits words). + +const std = @import("std"); + +/// The one codebook the family reads (the kernels' `mul1`, cuda.exl3.Codebook.mul1). +pub const codebook = "mul1"; +/// The weight format the family lists for an EXL3 checkpoint. +pub const format = "exl3-mul1"; + +fn object(o: std.json.ObjectMap, name: []const u8) ?std.json.ObjectMap { + const v = o.get(name) orelse return null; + return if (v == .object) v.object else null; +} + +fn word(o: std.json.ObjectMap, name: []const u8) ?[]const u8 { + const v = o.get(name) orelse return null; + return if (v == .string) v.string else null; +} + +/// config.json's quantization_config: top level, else the text config's. +pub fn block(root: std.json.ObjectMap) ?std.json.ObjectMap { + if (object(root, "quantization_config")) |q| return q; + if (object(root, "text_config")) |t| if (object(t, "quantization_config")) |q| return q; + return null; +} + +/// Whether config.json declares an EXL3 checkpoint (quant_method "exl3"), whatever its codebook. +pub fn named(root: std.json.ObjectMap) bool { + const q = block(root) orelse return false; + const m = word(q, "quant_method") orelse return false; + return std.mem.eql(u8, m, "exl3"); +} + +/// The first field an EXL3 checkpoint differs in from what the family serves, else null. `geometry` is +/// cuda_config.geometry (shapes, then semantics), passed in so this file stays out of cuda_config's privates. +pub fn mismatch(root: std.json.ObjectMap, text: std.json.ObjectMap, comptime geometry: fn (std.json.ObjectMap) ?[]const u8) ?[]const u8 { + const q = block(root) orelse return "quantization_config"; + const m = word(q, "quant_method") orelse return "quantization_config.quant_method"; + if (!std.mem.eql(u8, m, "exl3")) return "quantization_config.quant_method"; + const book = word(q, "codebook") orelse return "quantization_config.codebook"; + if (!std.mem.eql(u8, book, codebook)) return "quantization_config.codebook"; + return geometry(text); +} + +const fixture = @embedFile("fixtures/exl3_config.json"); + +test "exl3: the Qwen3.8-Flash-Next EXL3 pack's config.json is admitted as Mode.exl3" { + const config = @import("cuda_config.zig"); + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = fixture }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expectEqual(config.Mode.exl3, try config.mode(a, std.testing.io, dir)); +} + +test "exl3: other codebooks and other methods are refused by name; bits are never read" { + const config = @import("cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: ?[]const u8 }; + const cases = [_]Case{ + .{ .field = "codebook", .value = .{ .string = "mcg" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .{ .string = "inst3" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .null, .refused = "quantization_config.codebook" }, + .{ .field = "quant_method", .value = .{ .string = "gptq" }, .refused = "quantization_config.quant_method" }, + // the average width and the head's are not a contract: the trellis says each matrix's own width + .{ .field = "bits", .value = .{ .float = 4.25 }, .refused = null }, + .{ .field = "bits", .value = .{ .integer = 2 }, .refused = null }, + .{ .field = "head_bits", .value = .{ .integer = 8 }, .refused = null }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = root.get("text_config").?.object; + try std.testing.expect(named(root)); + try std.testing.expectEqual(@as(?[]const u8, null), mismatch(root, text, config.geometry)); + try root.getPtr("quantization_config").?.object.put(a, c.field, c.value); + const got = mismatch(parsed.value.object, text, config.geometry); + if (c.refused) |want| try std.testing.expectEqualStrings(want, got orelse "accepted") else try std.testing.expectEqual(@as(?[]const u8, null), got); + } +} + +test "exl3: the shapes and semantics checks still apply to an EXL3 checkpoint" { + const config = @import("cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: []const u8 }; + const cases = [_]Case{ + .{ .field = "hidden_size", .value = .{ .integer = 4096 }, .refused = "hidden_size" }, + .{ .field = "num_experts_per_tok", .value = .{ .integer = 8 }, .refused = "num_experts_per_tok" }, + .{ .field = "rms_norm_eps", .value = .{ .float = 1e-5 }, .refused = "rms_norm_eps" }, + .{ .field = "output_gate_type", .value = .{ .string = "silu" }, .refused = "output_gate_type" }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = &root.getPtr("text_config").?.object; + try text.put(a, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, mismatch(root, text.*, config.geometry) orelse "accepted"); + } +} + +test "exl3: a non-EXL3 config is not named EXL3" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"quantization_config\":{\"quant_method\":\"modelopt\"}}", .{}); + try std.testing.expect(!named(parsed.value.object)); + const bare = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"text_config\":{}}", .{}); + try std.testing.expect(!named(bare.value.object)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_dispatch_test.zig b/zig/src/families/flashnext/cuda_exl3_dispatch_test.zig new file mode 100644 index 0000000000..b741687cd2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_dispatch_test.zig @@ -0,0 +1,74 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("cuda_batch.zig"); +const load = @import("cuda_prompt_load.zig"); +const exl3_load = @import("cuda_exl3_load.zig"); + +const FakeDriver = struct { + var copies: usize = 0; + + fn copy(_: u64, _: u64, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + copies += 1; + return 1; + } + + fn errorName(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn errorString(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn make() cuda.Driver { + var driver: cuda.Driver = undefined; + driver.api.cuMemcpyDtoDAsync_v2 = copy; + driver.api.cuGetErrorName = errorName; + driver.api.cuGetErrorString = errorString; + return driver; + } +}; + +const MockTri = struct { + own: ?*const @import("cuda_glue.zig").Glue = null, + + pub fn routerMany(_: MockTri, _: u64, _: u64, _: u64, _: usize, _: usize) !void {} + pub fn topkMany(_: MockTri, _: u64, _: u64, _: u64, _: usize) !void {} +}; + +test "exl3: Keep dispatch rejects an unattached expert bank before legacy copy" { + FakeDriver.copies = 0; + var driver = FakeDriver.make(); + var host_bytes: [0]u8 align(16) = .{}; + var experts: exl3_load.Experts = .{ + .gpa = std.testing.allocator, + .ops = undefined, + .layer = undefined, + .bank = .{ .d = &driver, .bytes = &host_bytes }, + }; + var layer: load.Layer = undefined; + layer.router.ptr = 1; + layer.exl3 = &experts; + var scratch: batch.Scratch = undefined; + scratch.mixed = 2; + scratch.router_l = 3; + scratch.picks = 4; + scratch.wts = 5; + var pass: batch.Pass(MockTri) = .{ + .gpa = std.testing.allocator, + .io = undefined, + .stream = .{ .d = &driver, .handle = null }, + .tri = .{}, + .k = undefined, + .s = &scratch, + .m = undefined, + .seq = undefined, + .hooks = .{}, + .rows = 1, + }; + + try std.testing.expectError(error.Exl3PoolNotAttached, pass.moeBlockKeep(&layer, .{ .moe_x = 9 })); + try std.testing.expectEqual(@as(usize, 0), FakeDriver.copies); +} diff --git a/zig/src/families/flashnext/cuda_exl3_dplan.zig b/zig/src/families/flashnext/cuda_exl3_dplan.zig new file mode 100644 index 0000000000..bc61fe931a --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_dplan.zig @@ -0,0 +1,724 @@ +//! The offload pool planned on the device (flashnext_exl3_dplan.cu): a decode round's MoE needs no host wait between +//! the router and the experts, so the whole step (48 layers, the head, at tp=2 the NCCL sums) is one stream of device +//! work that a CUDA graph can hold (cuda_exl3_graph.zig). +//! +//! Each decoder layer's pool keeps its residency twice: the host's (exl3_offload.Vram + lifetime counts, which +//! prefill, the MTP head and any eager round plan with) and this device copy. Optional recent-demand eviction state +//! stays with this planner across ownership transfers and is not request-snapshot state: +//! toDevice: the host's lines, stamps, counts and clock go up (one H2D on the compute stream, after the pools' +//! copy streams drained), then every decode round plans on the device; +//! toHost: the device copy comes back (one D2H after the stream finished) and rebuilds each pool's Vram, lifetime counts and +//! statistics, before any host plan (cuda_exl3_pass.pooled, the expert-major prefill). +//! By default, the device planner follows the host planner's rules for a one-pass round without prefetch or +//! double-buffer protection. Experimental recent eviction changes only miss ordering and victim scores; it still +//! computes every routed expert. The tables are the pool's own (fixed addresses): the device rewrites them in place. +//! +//! A miss is copied by fn_exl3_dcopy, a grid reading the pinned host bank through its mapped address (unified +//! addressing), instead of a DMA on the pool's copy stream: the copy is device work ordered on the compute stream. + +const std = @import("std"); +const cuda = @import("cuda"); +const pool_ = @import("cuda_exl3_pool.zig"); +const offload = pool_.offload; +const telemetry = @import("exl3_telemetry.zig"); +const recent_demand = @import("exl3_recent_demand.zig"); +const layout = @import("exl3_dplan_layout.zig"); +pub const demand_trace = @import("exl3_dplan_trace.zig"); +pub const replay = @import("exl3_dplan_replay.zig"); +pub const policy_eval = @import("exl3_dplan_policy_eval.zig"); + +const experts_n = layout.experts_n; +const recent_window = layout.recent_window; +const recent_words_per_row = experts_n / 32; +const misc_words = layout.misc_words; +const recentCountsOffset = layout.recentCountsOffset; +const recentOffset = layout.recentOffset; +pub const miscOffset = layout.miscOffset; +pub const PoolState = layout.PoolState; +const empty_id = offload.Cache.empty_id; + +comptime { + if (recent_demand.expert_count != experts_n) @compileError("recent expert history shape mismatch"); +} + +test { + _ = recent_demand; + _ = layout; + _ = demand_trace; + _ = replay; + _ = @import("exl3_dplan_trace_test.zig"); + _ = @import("exl3_dplan_trace_replay_test.zig"); + _ = policy_eval; +} + +test "exl3: demand trace diagnostic APIs reject disabled capture" { + var dp: DevPlan = undefined; + dp.trace_mem = null; + try std.testing.expectEqual(@as(usize, 0), dp.traceDeviceBytes()); + try std.testing.expectError(error.TraceDisabled, dp.readTraceAfterConsumers(0, &.{})); + try std.testing.expectError(error.TraceDisabled, dp.downloadTraceAfterConsumers(std.testing.allocator, 0)); + try dp.traceOwnership(undefined, true); + try dp.traceOwnership(undefined, false); +} + +test "exl3: ownership trace checkpoints use ragged pool state and serialized compute stream" { + const Probe = struct { + var calls: usize = 0; + var arguments: [4][4]u64 = undefined; + var streams: [4]cuda.abi.Stream = undefined; + + fn launch(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, stream: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + inline for (.{ u64, c_int, u64, c_int }, 0..) |Type, index| { + const value: *const Type = @ptrCast(@alignCast(parameters.?[index].?)); + arguments[calls][index] = @intCast(value.*); + } + streams[calls] = stream; + calls += 1; + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Probe.launch; + var entries = [_]PoolState{ try PoolState.init(10, 0), try PoolState.init(31, stateWords(10)) }; + var device: DevPlan = undefined; + device.pool_states = &entries; + device.state.ptr = 0x10000; + device.trace_mem = .{ .d = &driver, .ptr = 0x100000, .len = 2 * demand_trace.pool_bytes, .borrowed = true }; + device.trace_checkpoint_fn = .{ .d = &driver, .handle = @ptrFromInt(1) }; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(2) }; + Probe.calls = 0; + try device.traceOwnership(stream, true); + try device.traceOwnership(stream, false); + for (0..4) |call| { + const index = call % 2; + try std.testing.expectEqual([4]u64{ device.state.ptr + entries[index].byteOffset(), entries[index].lines, 0x100000 + index * demand_trace.pool_bytes, @intFromBool(call < 2) }, Probe.arguments[call]); + try std.testing.expectEqual(stream.handle, Probe.streams[call]); + } + device.trace_mem = null; + device.trace_checkpoint_fn = null; + try device.traceOwnership(stream, true); + try std.testing.expectEqual(@as(usize, 4), Probe.calls); +} + +/// u32 words of one layer's device state. +pub const stateWords = layout.stateWords; + +test "exl3: recent-demand planner state is separate from lifetime profile counts" { + try std.testing.expectEqual(@as(usize, 3 * experts_n), recentOffset(10) - 20); + try std.testing.expectEqual(@as(usize, recent_window * recent_words_per_row), miscOffset(10) - recentOffset(10)); + try std.testing.expectEqual(@as(usize, misc_words), stateWords(10) - miscOffset(10)); +} + +/// The planned copies a round can hold: a row's ten routed picks each. +fn workCap(max_rows: usize) usize { + return max_rows * 10; +} + +pub const Owner = enum { host, device }; + +/// Device statistics of a layer (the misc words after now). +pub const Misc = extern struct { now: u32, rounds: u32, hit_slots: u32, miss_slots: u32, copies: u32, evictions: u32, overflow: u32, inplace: u32, recent_cursor: u32, recent_filled: u32 }; + +pub const DevPlan = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + plan_fn: cuda.Function, + copy_fn: cuda.Function, + restore_fn: cuda.Function, + pools: []*pool_.Pool, + pool_states: []PoolState, + /// Rows a planned round may hold: rows * 10 distinct experts must fit the lines. + max_rows: usize, + state: cuda.DeviceBuffer, + mirror: cuda.HostBuffer, + /// Per pool: src [512 * 3] int64, lens [512 * 3] int64, k2 [512 * 3] int32 (static). + statics: cuda.DeviceBuffer, + work: cuda.DeviceBuffer, + owner: Owner = .host, + /// The stream the last device round ran on (toHost waits for it). + stream: ?cuda.Stream = null, + copy_blocks: u32 = 144, + /// Overlap (TF_FLASHNEXT_EXL3_OVERLAP): the held experts compute while the misses' copies run on `side`, then the + /// misses' experts (and the combine) follow; the per-slot outputs are the single pass's (each slot is written by + /// the one pass that holds its expert, as the host cut's split rounds do). + overlap: bool = false, + side: cuda.Stream, + ev_plan: cuda.Event, + ev_copy: cuda.Event, + /// The split's masked picks [max_rows * 11] each: held (and the shared expert), then the misses. + split: cuda.DeviceBuffer, + /// Device statistics already added to the pools' (a pull adds the difference). + seen: []Misc, + pushes: u64 = 0, + pulls: u64 = 0, + history_initialized: bool = false, + /// Experts read in place from the host bank (more misses than free lines in a round), summed at each pull. + inplace: u64 = 0, + trace_mem: ?cuda.DeviceBuffer = null, + trace_checkpoint_fn: ?cuda.Function = null, + + const static_bytes = layout.static_bytes; + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, pools: []const *pool_.Pool) !*DevPlan { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (pools.len == 0) return error.Invalid; + const pool_states = try gpa.alloc(PoolState, pools.len); + errdefer gpa.free(pool_states); + var total_words: usize = 0; + for (pools, pool_states) |p, *entry| { + if (p.lines.spec.experts < experts_n) return error.Invalid; + entry.* = try PoolState.init(p.opt.lines, total_words); + total_words = entry.endWords(); + } + var mod = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_dplan); + errdefer mod.unload(); + const dp = try gpa.create(DevPlan); + errdefer gpa.destroy(dp); + const own = try gpa.dupe(*pool_.Pool, pools); + errdefer gpa.free(own); + var state = try cuda.DeviceBuffer.alloc(d, total_words * @sizeOf(u32)); + errdefer state.free(); + var mirror = try cuda.HostBuffer.alloc(d, state.len); + errdefer mirror.free(); + var statics = try cuda.DeviceBuffer.alloc(d, pools.len * static_bytes); + errdefer statics.free(); + // rows a planned round may hold (MTP verify windows reach 16): more distinct misses than lines free of the + // round's experts are read in place from the mapped bank for that round (fn_exl3_dplan), so any row count + // up to this bound plans in one pass + const max_rows = maxRows(); + var trace_mem: ?cuda.DeviceBuffer = null; + errdefer if (trace_mem) |*buffer| buffer.free(); + if (traceEnabled()) { + for (pools, pool_states, 0..) |p, entry, index| _ = try demand_trace.header(index, p.opt.layer, entry.lines, max_rows); + trace_mem = try cuda.DeviceBuffer.alloc(d, try std.math.mul(usize, pools.len, demand_trace.pool_bytes)); + try trace_mem.?.fill32(0, null); + for (pools, pool_states, 0..) |p, entry, index| { + const header = try demand_trace.header(index, p.opt.layer, entry.lines, max_rows); + try trace_mem.?.upload(index * demand_trace.pool_bytes, std.mem.asBytes(&header)); + } + } + var work = try cuda.DeviceBuffer.alloc(d, workCap(max_rows) * 9 * 8 + 256 + (experts_n + 64) * 4); + errdefer work.free(); + const seen = try gpa.alloc(Misc, pools.len); + errdefer gpa.free(seen); + var side = try cuda.Stream.init(d, true); + errdefer side.deinit(); + var ev_plan = try cuda.Event.init(d, false); + errdefer ev_plan.deinit(); + var ev_copy = try cuda.Event.init(d, false); + errdefer ev_copy.deinit(); + var split = try cuda.DeviceBuffer.alloc(d, 2 * max_rows * 11 * 4); + errdefer split.free(); + // every expert's three host-bank addresses (mapped: unified addressing), byte counts and widths + const host = try gpa.alloc(u8, static_bytes); + defer gpa.free(host); + for (pools, 0..) |p, i| { + const src = std.mem.bytesAsSlice(u64, host[0 .. experts_n * 24]); + const lens = std.mem.bytesAsSlice(u64, host[experts_n * 24 .. experts_n * 48]); + const k2 = std.mem.bytesAsSlice(i32, host[experts_n * 48 ..]); + for (0..experts_n) |e| for (0..3) |w| { + const b = p.src.routed[e][w]; + src[e * 3 + w] = if (b.len == 0) 0 else @intFromPtr(b.ptr); + lens[e * 3 + w] = b.len; + k2[e * 3 + w] = p.k2[e][w]; + if (b.len != 0 and (b.len % 16 != 0 or @intFromPtr(b.ptr) % 16 != 0)) return error.Invalid; + }; + try statics.upload(i * static_bytes, host); + } + dp.* = .{ + .gpa = gpa, + .d = d, + .mod = mod, + .plan_fn = try mod.function("fn_exl3_dplan"), + .copy_fn = try mod.function("fn_exl3_dcopy"), + .restore_fn = try mod.function("fn_exl3_drestore"), + .pools = own, + .pool_states = pool_states, + .max_rows = max_rows, + .state = state, + .mirror = mirror, + .statics = statics, + .work = work, + .seen = seen, + .side = side, + .ev_plan = ev_plan, + .ev_copy = ev_copy, + .split = split, + .overlap = overlapFromEnv(pools[0]), + .trace_mem = trace_mem, + .trace_checkpoint_fn = if (trace_mem != null) try mod.function("fn_exl3_dcheckpoint") else null, + }; + if (dp.overlap) dp.copy_blocks = overlapCopyBlocks(d) catch dp.copy_blocks; + if (std.c.getenv("TF_FLASHNEXT_EXL3_DCOPY_BLOCKS")) |v| dp.copy_blocks = std.fmt.parseInt(u32, std.mem.span(v), 10) catch dp.copy_blocks; + return dp; + } + + /// With the trace on and TF_FLASHNEXT_EXL3_DPLAN_TRACE_DIR set: each pool's raw buffer as dplan--pool.bin. + fn dumpTrace(dp: *DevPlan) void { + const dir = std.c.getenv("TF_FLASHNEXT_EXL3_DPLAN_TRACE_DIR") orelse return; + const buffer = dp.trace_mem orelse return; + if (dp.stream) |s| s.synchronize() catch return; + const bytes = std.heap.page_allocator.alloc(u8, demand_trace.pool_bytes) catch return; + defer std.heap.page_allocator.free(bytes); + for (0..dp.pools.len) |index| { + buffer.download(index * demand_trace.pool_bytes, bytes) catch return; + var path: [256]u8 = undefined; + const name = std.fmt.bufPrint(&path, "{s}/dplan-{d}-pool{d}.bin\x00", .{ std.mem.span(dir), std.c.getpid(), index }) catch return; + const fd = std.c.open(@ptrCast(name.ptr), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o644)); + if (fd < 0) return; + defer _ = std.c.close(fd); + var sent: usize = 0; + while (sent < bytes.len) { + const n = std.c.write(fd, bytes[sent..].ptr, bytes.len - sent); + if (n <= 0) break; + sent += @intCast(n); + } + } + } + + pub fn deinit(dp: *DevPlan) void { + dp.side.synchronize() catch {}; + dumpTrace(dp); + if (dp.trace_mem) |*buffer| buffer.free(); + dp.split.free(); + dp.ev_copy.deinit(); + dp.ev_plan.deinit(); + dp.side.deinit(); + dp.work.free(); + dp.statics.free(); + dp.mirror.free(); + dp.state.free(); + dp.mod.unload(); + dp.gpa.free(dp.seen); + dp.gpa.free(dp.pool_states); + dp.gpa.free(dp.pools); + dp.gpa.destroy(dp); + } + + fn mirrorOf(dp: *DevPlan, i: usize) []u32 { + const all: []u32 = @alignCast(std.mem.bytesAsSlice(u32, dp.mirror.bytes[0..dp.state.len])); + return dp.pool_states[i].slice(all); + } + + /// The host's residency up to the device, on `s` (a no-op if the device owns it). Waits for every pool's copies. + pub fn toDevice(dp: *DevPlan, s: cuda.Stream) !void { + dp.stream = s; + if (dp.owner == .device) return; + for (dp.pools, 0..) |p, i| { + const lines = dp.pool_states[i].lines; + // copies in flight land, prefetched lines count as resident (the device knows no in-flight state) + try p.lines.drain(); + for (p.line_tab, p.slots_tab) |*state, id| if (state.* == .inflight) { + offload.reapAt(&p.v, p.opt.layer, id, p.now) catch {}; + }; + const m = dp.mirrorOf(i); + const line_id: []i32 = @ptrCast(m[0..lines]); + const stamp = m[lines .. 2 * lines]; + const line_of: []i32 = @ptrCast(m[2 * lines ..][0..experts_n]); + const counts = m[2 * lines + experts_n ..][0..experts_n]; + if (!dp.history_initialized) @memset(m[recentCountsOffset(lines)..miscOffset(lines)], 0); + @memset(line_of, -1); + for (0..lines) |ln| { + const id = p.slots_tab[ln]; + const held = p.line_tab[ln] == .resident and id != empty_id and id < experts_n; + line_id[ln] = if (held) id else -1; + stamp[ln] = p.stamp_tab[ln]; + if (held) line_of[id] = @intCast(ln); + } + @memcpy(counts, &p.counts[p.opt.layer]); + const misc: *Misc = @ptrCast(m[miscOffset(lines)..][0..misc_words]); + const recent_cursor = if (dp.history_initialized) misc.recent_cursor else 0; + const recent_filled = if (dp.history_initialized) misc.recent_filled else 0; + misc.* = .{ .now = p.now, .rounds = 0, .hit_slots = 0, .miss_slots = 0, .copies = 0, .evictions = 0, .overflow = 0, .inplace = 0, .recent_cursor = recent_cursor, .recent_filled = recent_filled }; + dp.seen[i] = misc.*; + } + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dp.state.ptr, dp.mirror.bytes.ptr, dp.state.len, s.handle), "cuMemcpyHtoDAsync"); + try dp.traceOwnership(s, true); + dp.history_initialized = true; + dp.owner = .device; + dp.pushes += 1; + } + + /// The device's residency back to the host pools (a no-op if the host owns it): waits for the last round's stream. + pub fn toHost(dp: *DevPlan) !void { + if (dp.owner == .host) return; + if (dp.stream) |s| try dp.traceOwnership(s, false); + if (dp.stream) |s| try s.synchronize(); + try dp.side.synchronize(); + try dp.d.check(dp.d.api.cuMemcpyDtoH_v2(dp.mirror.bytes.ptr, dp.state.ptr, dp.state.len), "cuMemcpyDtoH"); + var overflow: u64 = 0; + for (dp.pools, 0..) |p, i| { + const lines = dp.pool_states[i].lines; + const m = dp.mirrorOf(i); + const line_id: []const i32 = @ptrCast(m[0..lines]); + const stamp = m[lines .. 2 * lines]; + const counts = m[2 * lines + experts_n ..][0..experts_n]; + const misc: *const Misc = @ptrCast(m[miscOffset(lines)..][0..misc_words]); + for (0..lines) |ln| { + const id = line_id[ln]; + p.slots_tab[ln] = if (id < 0) empty_id else @intCast(id); + p.line_tab[ln] = if (id < 0) .empty else .resident; + p.stamp_tab[ln] = stamp[ln]; + p.layer_tab[ln] = p.opt.layer; + } + @memcpy(&p.counts[p.opt.layer], counts); + const was = dp.seen[i]; + telemetry.Device.record(&p.telemetry_device, &p.telemetry_valid, misc.*, was, dp.max_rows * 10); + p.now = misc.now; + p.stats.rounds += misc.rounds -% was.rounds; + p.stats.passes += misc.rounds -% was.rounds; + p.stats.hit_slots += misc.hit_slots -% was.hit_slots; + p.stats.miss_slots += misc.miss_slots -% was.miss_slots; + p.stats.misses += misc.copies -% was.copies; + p.stats.evictions += misc.evictions -% was.evictions; + overflow += misc.overflow; + dp.inplace += misc.inplace -% was.inplace; + dp.seen[i] = misc.*; + // the device's rounds finished (synchronized above): no line waits for a reader, nothing is in flight + @memset(p.lines.line_epoch, 0); + @memset(p.lines.line_pending, false); + p.prev_n = 0; + p.picks_pending = false; + } + dp.owner = .host; + dp.pulls += 1; + if (overflow != 0) { + std.log.err("flashnext exl3 dplan: {d} rounds had more distinct experts than the planner holds (256): results of those rounds are wrong", .{overflow}); + return error.PoolPasses; + } + } + + /// Device statistics so far, summed over the layers (synchronizes; does not change the owner). + pub fn peek(dp: *DevPlan) !Misc { + var t: Misc = std.mem.zeroes(Misc); + if (dp.owner == .host) return t; + if (dp.stream) |s| try s.synchronize(); + try dp.d.check(dp.d.api.cuMemcpyDtoH_v2(dp.mirror.bytes.ptr, dp.state.ptr, dp.state.len), "cuMemcpyDtoH"); + for (dp.pools, 0..) |_, i| { + const m = dp.mirrorOf(i); + const misc: *const Misc = @ptrCast(m[miscOffset(dp.pool_states[i].lines)..][0..misc_words]); + inline for (.{ "rounds", "hit_slots", "miss_slots", "copies", "evictions", "overflow", "inplace" }) |f| @field(t, f) += @field(misc, f) - @field(dp.seen[i], f); + } + return t; + } + + pub fn samplePoolAfterConsumers(dp: *const DevPlan, pool: *const pool_.Pool) !telemetry.Device { + if (!pool.telemetry_valid) return error.TelemetryCounterRegression; + var total = pool.telemetry_device; + if (dp.owner != .device) return total; + for (dp.pools, 0..) |candidate, index| { + if (candidate != pool) continue; + var observed: Misc = undefined; + const entry = dp.pool_states[index]; + const offset = entry.byteOffset() + miscOffset(entry.lines) * @sizeOf(u32); + try dp.state.download(offset, std.mem.asBytes(&observed)); + try telemetry.add(telemetry.Device, &total, try telemetry.Device.pendingBounded(observed, dp.seen[index], dp.max_rows * 10)); + return total; + } + return total; + } + + /// Exact additional device allocation; tracing is fixed at init and costs zero bytes when disabled. + pub fn traceDeviceBytes(dp: *const DevPlan) usize { + return if (dp.trace_mem) |buffer| buffer.len else 0; + } + + fn traceOwnership(dp: *const DevPlan, stream: cuda.Stream, transfer: bool) !void { + const buffer = dp.trace_mem orelse return; + for (dp.pool_states, 0..) |entry, index| { + var args: cuda.launch.Args = .{}; + args.add(dp.state.ptr + entry.byteOffset()); + args.add(@as(c_int, @intCast(entry.lines))); + args.add(buffer.ptr + index * demand_trace.pool_bytes); + args.add(@as(c_int, @intFromBool(transfer))); + try cuda.launch.launch(dp.trace_checkpoint_fn.?, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1 } }, stream, &args); + } + } + + /// Caller must join all planner/expert consumers and stop submissions before this blocking diagnostic read. + /// Outside capture only. Borrows `out` (pool_words u32s); reading does not consume records. + /// The initial seed, ownership checkpoints and ordered rounds include warmup/speculation. + pub fn readTraceAfterConsumers(dp: *const DevPlan, index: usize, out: []u32) !demand_trace.View { + const buffer = dp.trace_mem orelse return error.TraceDisabled; + if (index >= dp.pools.len or out.len != demand_trace.pool_words) return error.Invalid; + try buffer.download(index * demand_trace.pool_bytes, std.mem.sliceAsBytes(out)); + return demand_trace.View.init(out); + } + + /// Allocating diagnostic read with the same joined/quiescent contract. Caller owns the snapshot. + /// Snapshot.view().encode() serializes little-endian records and all full checkpoints. + pub fn downloadTraceAfterConsumers(dp: *const DevPlan, allocator: std.mem.Allocator, index: usize) !demand_trace.Snapshot { + if (dp.trace_mem == null) return error.TraceDisabled; + if (index >= dp.pools.len) return error.Invalid; + const words = try allocator.alloc(u32, demand_trace.pool_words); + errdefer allocator.free(words); + _ = try dp.readTraceAfterConsumers(index, words); + return .{ .allocator = allocator, .words = words }; + } + + /// The in-place experts of the last plan ([512] int32) and their count after it. + fn ovfList(dp: *const DevPlan) u64 { + return dp.work.ptr + workCap(dp.max_rows) * 9 * 8 + 256; + } + + /// After the round's experts (on `s`): the in-place experts' table entries cleared. + pub fn restore(dp: *DevPlan, s: cuda.Stream, i: usize) !void { + if (i >= dp.pools.len) return error.Invalid; + const p = dp.pools[i]; + var a: cuda.Args = .{}; + a.add(dp.ovfList()); + a.add(dp.ovfList() + experts_n * 4); + for (0..3) |w| a.add(p.lines.table(w)); + for (0..3) |w| a.add(p.lines.widths(w)); + try cuda.launch.launch(dp.restore_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + /// `s` waits for the last plan's copies (runLayer joins them itself; a caller of `plan` alone needs this). + pub fn join(dp: *DevPlan, s: cuda.Stream) !void { + if (dp.overlap) try s.wait(dp.ev_copy); + } + + /// The index of `p` among the planned pools, or null (the MTP layer's pool, another engine's). + pub fn indexOf(dp: *const DevPlan, p: *const pool_.Pool) ?usize { + const i = p.opt.layer; + if (i < dp.pools.len and dp.pools[i] == p) return i; + for (dp.pools, 0..) |q, j| if (q == p) return j; + return null; + } + + /// Plans pool `i`'s round of `rows` rows on the device and copies its misses into their lines, on `s`: after this + /// the pool's tables hold every expert `picks` [rows * 11] names. Graph-capturable (no host work depends on it). + /// All launches and graph replays using this DevPlan MUST be serialized on one compute timeline. + /// A stream handoff must join previous planner, copy and expert/restore consumers first. + /// The shared work/split scratch already requires this; trace headers have the same single writer. + /// The side stream runs copies only and never writes planner state or trace records. + pub fn plan(dp: *DevPlan, s: cuda.Stream, i: usize, picks: u64, rows: usize) !void { + if (rows == 0 or rows > dp.max_rows) return error.PoolPasses; + if (i >= dp.pools.len) return error.Invalid; + const p = dp.pools[i]; + const entry = dp.pool_states[i]; + const st = dp.state.ptr + entry.byteOffset(); + const sb = dp.statics.ptr + i * static_bytes; + const work = dp.work.ptr; + const work_n = dp.work.ptr + workCap(dp.max_rows) * 9 * 8; + var a: cuda.Args = .{}; + a.add(picks); + a.add(@as(c_int, @intCast(rows * 11))); + a.add(st); + a.add(@as(c_int, @intCast(entry.lines))); + a.add(@as(c_int, @intFromBool(recentEvictionEnabled()))); + for (0..3) |w| a.add(p.lines.table(w)); + for (0..3) |w| a.add(p.lines.widths(w)); + a.add(sb); + a.add(sb + experts_n * 24); + a.add(sb + experts_n * 48); + a.add(@as(u64, p.lines.mem.ptr)); + a.add(@as(u64, p.lines.layout.line_bytes)); + for (p.lines.layout.off) |o| a.add(@as(u64, o)); + a.add(work); + a.add(work_n); + a.add(if (dp.overlap) dp.split.ptr else @as(u64, 0)); + a.add(if (dp.overlap) dp.split.ptr + dp.max_rows * 11 * 4 else @as(u64, 0)); + a.add(dp.ovfList()); + a.add(dp.ovfList() + experts_n * 4); + a.add(if (dp.trace_mem) |buffer| buffer.ptr + i * demand_trace.pool_bytes else @as(u64, 0)); + try cuda.launch.launch(dp.plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, s, &a); + var c: cuda.Args = .{}; + c.add(work); + c.add(work_n); + if (!dp.overlap) return cuda.launch.launch(dp.copy_fn, .{ .grid = .{ .x = dp.copy_blocks }, .block = .{ .x = 256 } }, s, &c); + // the copies on the side stream (a fork in a captured graph); `s` joins them before the misses' experts + try dp.ev_plan.record(s); + try dp.side.wait(dp.ev_plan); + try cuda.launch.launch(dp.copy_fn, .{ .grid = .{ .x = dp.copy_blocks }, .block = .{ .x = 256 } }, dp.side, &c); + try dp.ev_copy.record(dp.side); + } +}; + +// ---- the engine's planner ------------------------------------------------------------------------------------------- + +/// The engine's device planner (one engine a process: a rank), made on the first device round. +var global: ?*DevPlan = null; +var global_owner: ?*const anyopaque = null; +/// Set while a decode round runs with device plans (cuda_exl3_graph.zig): cuda_exl3_pass.moe takes `layer` instead +/// of the host cut. +pub var active: ?*DevPlan = null; + +/// TF_FLASHNEXT_EXL3_DPLAN: 1 plans decode rounds on the device (implied by TF_FLASHNEXT_EXL3_GRAPH=1). +pub fn enabled() bool { + if (@import("cuda_exl3_graph.zig").force) |f| return f != .off; + if (std.c.getenv("TF_FLASHNEXT_EXL3_DPLAN")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + return @import("cuda_exl3_graph.zig").enabled(); +} + +/// Experimental recent-demand eviction policy. Lifetime counts still persist and seed startup placement. +pub fn recentEvictionEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_RECENT_EVICTION") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +/// TF_FLASHNEXT_EXL3_DPLAN_TRACE=1 enables bounded per-pool demand capture at planner initialization. +/// No per-round host reads, transfers, allocations, synchronization, or placement-policy changes. +fn traceEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_DPLAN_TRACE") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +test "exl3: demand trace planner arguments stay fixed and joined reads isolate pools" { + const Probe = struct { + var trace_address: u64 = 0; + var launches: usize = 0; + + fn launch(function: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + launches += 1; + if (@intFromPtr(function.?) == 1) { + const address: *const u64 = @ptrCast(@alignCast(parameters.?[25].?)); + trace_address = address.*; + } + return 0; + } + + fn download(destination: ?*anyopaque, source: u64, length: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrCast(destination.?); + const origin: [*]const u8 = @ptrFromInt(source); + @memcpy(target[0..length], origin[0..length]); + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Probe.launch; + driver.api.cuMemcpyDtoH_v2 = Probe.download; + const words = try std.testing.allocator.alloc(u32, 2 * demand_trace.pool_words); + defer std.testing.allocator.free(words); + for (0..2) |index| @memcpy(words[index * demand_trace.pool_words ..][0..demand_trace.header_words], &try demand_trace.header(index, index, 10, 16)); + var resident: pool_.Pool = undefined; + resident.lines.tables.ptr = 4096; + resident.lines.mem.ptr = 8192; + resident.lines.tlay = cuda.exl3_pool.Tables.init(513); + resident.lines.layout = cuda.exl3_pool.Layout.init(.{ 256, 256, 256 }); + var pools = [_]*pool_.Pool{ &resident, &resident }; + var states = [_]PoolState{ try PoolState.init(10, 0), try PoolState.init(10, stateWords(10)) }; + var device: DevPlan = undefined; + device.pools = &pools; + device.pool_states = &states; + device.max_rows = 16; + device.state.ptr = 16384; + device.statics.ptr = 32768; + device.work.ptr = 65536; + device.plan_fn = .{ .d = &driver, .handle = @ptrFromInt(1) }; + device.copy_fn = .{ .d = &driver, .handle = @ptrFromInt(2) }; + device.copy_blocks = 1; + device.overlap = false; + device.trace_mem = .{ .d = &driver, .ptr = @intFromPtr(words.ptr), .len = words.len * 4, .borrowed = true }; + const out = try std.testing.allocator.alloc(u32, demand_trace.pool_words); + defer std.testing.allocator.free(out); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + Probe.launches = 0; + for (0..2) |index| { + for (0..2) |_| { + try device.plan(stream, index, 1024, 1); + try std.testing.expectEqual(@intFromPtr(words.ptr) + index * demand_trace.pool_bytes, Probe.trace_address); + } + const view = try device.readTraceAfterConsumers(index, out); + try std.testing.expectEqual(@as(usize, 0), view.recorded()); + try std.testing.expectEqual(@as(u32, @intCast(index)), view.words[2]); + } + try std.testing.expectEqual(@as(usize, 8), Probe.launches); + try std.testing.expectEqual(words.len * 4, device.traceDeviceBytes()); + try std.testing.expectError(error.Invalid, device.readTraceAfterConsumers(2, out)); + try std.testing.expectError(error.Invalid, device.readTraceAfterConsumers(0, out[1..])); + device.trace_mem = null; + try device.plan(stream, 0, 1024, 1); + try std.testing.expectEqual(@as(u64, 0), Probe.trace_address); + try std.testing.expectEqual(@as(usize, 0), device.traceDeviceBytes()); + try std.testing.expectError(error.TraceDisabled, device.readTraceAfterConsumers(0, out)); +} + +/// The device planner over `layers`' pools (the decoder layers only), made once for `owner`; null when a layer has no +/// pool (resident experts). +pub fn forLayers(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const anyopaque, layers: []const @import("cuda_prompt_load.zig").Layer) !?*DevPlan { + if (global) |g| { + if (global_owner == owner) return g; + try g.toHost(); + g.deinit(); + global = null; + } + var pools: [64]*pool_.Pool = undefined; + var n: usize = 0; + for (layers) |*l| { + const x = l.exl3 orelse return null; + const p = x.pool orelse return null; + if (n == pools.len) return null; + pools[n] = p; + n += 1; + } + if (n == 0) return null; + // the MTP layer's pool last (its own lifetime counts; the draft chain plans it on the device) + if (mtp_pool) |mp| if (n < pools.len) { + pools[n] = mp; + n += 1; + }; + global = try DevPlan.init(gpa, d, pools[0..n]); + global_owner = owner; + return global; +} + +/// The draft head's (MTP layer's) pool, registered by cuda_exl3_pass.attach. +pub var mtp_pool: ?*pool_.Pool = null; + +/// The engine's planner if made (any owner). +pub fn current() ?*DevPlan { + return global; +} + +/// Before any host plan: the device's residency back to the pools (no-op without a device planner). +pub fn toHost() !void { + if (global) |g| try g.toHost(); +} + +pub fn releaseFor(owner: *const anyopaque) void { + if (global) |g| if (global_owner == owner) { + g.deinit(); + global = null; + global_owner = null; + active = null; + }; +} + +/// One offloaded layer's experts on the device plan: plan + copies, then `Experts.run` on the pool's tables, as the +/// host cut's single pass runs it (cuda_exl3_pass.pooled): same kernels, same inputs, so the same bits. +pub fn runLayer(dp: *DevPlan, st: cuda.Stream, i: usize, ex: *const cuda.exl3_experts.Experts, layer: cuda.exl3_experts.Layer, xin: u64, picks: u64, rows: usize, sc: cuda.exl3_experts.Scratch, wts: u64, out: u64) !void { + try dp.plan(st, i, picks, rows); + if (!dp.overlap) { + try ex.run(st, layer, xin, 2560, .bf16, picks, wts, rows, out, sc); + } else { + // held experts while the copies run, then the misses' experts and (expert parallel) the combine of every slot + try ex.run(st, layer, xin, 2560, .bf16, dp.split.ptr, 0, rows, 0, sc); + try st.wait(dp.ev_copy); + try ex.run(st, layer, xin, 2560, .bf16, dp.split.ptr + dp.max_rows * 11 * 4, wts, rows, out, sc); + } + try dp.restore(st, i); +} + +/// TF_FLASHNEXT_EXL3_GRAPH_ROWS: rows a device-planned round holds at most (default 16, the MTP window). +fn maxRows() usize { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_ROWS") orelse return 16; + return @max(1, std.fmt.parseInt(usize, std.mem.span(v), 10) catch 16); +} + +/// The split is on unless TF_FLASHNEXT_EXL3_OVERLAP=0: +3-4% decode on two RTX 5060 Ti (see the report). +fn overlapFromEnv(p: *const pool_.Pool) bool { + _ = p; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_OVERLAP") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// With the split the copies run beside the held experts: a quarter of the SMs, at least 8 blocks, keeps the link busy. +fn overlapCopyBlocks(d: *const cuda.Driver) !u32 { + var device: c_int = 0; + var sms: c_int = 0; + try d.check(d.api.cuDeviceGet(&device, 0), "cuDeviceGet"); + try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, device), "cuDeviceGetAttribute"); + return @max(8, @as(u32, @intCast(sms)) / 4); +} diff --git a/zig/src/families/flashnext/cuda_exl3_fork.zig b/zig/src/families/flashnext/cuda_exl3_fork.zig new file mode 100644 index 0000000000..0de3961231 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_fork.zig @@ -0,0 +1,127 @@ +//! A multi-part EXL3 projection's parts on parallel streams (a DeltaNet layer's q|k|v, z, b, a; attention's +//! q, k, v, indexer): each EXL3 part after the first runs on its own side stream with its own Dense scratch (rotated +//! input, split partials, counters, stitch buffer), forked from and joined back into the pass's stream with events, +//! so a captured graph holds them as parallel branches. Every part is the same kernels on the same inputs writing +//! its own columns: the same bits (tf-cuda-test flashnext-exl3-graph eq). A decode projection launches 80-192 +//! programs that leave an SM wave partly empty; the next part fills it. +//! TF_FLASHNEXT_EXL3_FORK=0 runs the parts one after another (upstream order). + +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("cuda_exl3_load.zig"); + +pub const branches = 3; +/// Rows a forked projection takes at most (decode rounds and verify windows; prompts keep one stream). +pub const max_rows: usize = 16; + +pub const Fork = struct { + side: [branches]cuda.Stream, + dense: [branches]cuda.exl3.Dense, + tmp: [branches]cuda.DeviceBuffer, + ev_fork: cuda.Event, + ev_join: [branches]cuda.Event, + tmp_bytes: usize, +}; + +var forks: std.AutoHashMapUnmanaged(*load.Ops, *Fork) = .empty; + +pub fn enabled() bool { + if (force) |f| return f; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_FORK") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// A test's override of TF_FLASHNEXT_EXL3_FORK. +pub var force: ?bool = null; + +/// Ops.reserve's hook: the branches' streams, kernels and scratch for `o`'s widest projection, before any pass. +pub fn reserve(o: *load.Ops) !void { + if (o.max_k == 0 or forks.contains(o)) return; + const f = try o.gpa.create(Fork); + errdefer o.gpa.destroy(f); + f.tmp_bytes = max_rows * @max(o.max_part_n, 128) * 2; + var made: usize = 0; + errdefer for (0..made) |i| { + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + }; + f.ev_fork = try cuda.Event.init(o.d, false); + errdefer f.ev_fork.deinit(); + for (0..branches) |i| { + f.side[i] = try cuda.Stream.init(o.d, true); + f.dense[i] = try cuda.exl3.Dense.load(o.d); + f.tmp[i] = try cuda.DeviceBuffer.alloc(o.d, f.tmp_bytes); + f.ev_join[i] = try cuda.Event.init(o.d, false); + made += 1; + const de = &f.dense[i]; + de.xh = try cuda.DeviceBuffer.alloc(o.d, 128 * o.max_k * 2); + if (o.max_z > 0) de.z = try cuda.DeviceBuffer.alloc(o.d, o.max_z); + de.counters = try cuda.DeviceBuffer.alloc(o.d, o.max_counters); + } + try forks.put(o.gpa, o, f); +} + +pub fn release(o: *load.Ops) void { + const kv = forks.fetchRemove(o) orelse return; + const f = kv.value; + for (0..branches) |i| { + f.side[i].synchronize() catch {}; + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + } + f.ev_fork.deinit(); + o.gpa.destroy(f); +} + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// Pass.lane()'s parts loop (bf16 out, sk 1) with the EXL3 parts after the first on side streams; false when the +/// projection does not fork (one EXL3 part, too many rows, no branches reserved), and the caller runs its loop. +/// `run(part, y)` runs a part on the pass's stream as the caller's loop would. +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, m: usize, comptime runMain: anytype) !bool { + if (!enabled() or m > max_rows) return false; + var n_exl3: usize = 0; + for (e.parts) |pt| n_exl3 += @intFromBool(pt.kind == .exl3); + if (n_exl3 < 2) return false; + const f = forks.get(e.ops) orelse return false; + for (e.parts) |pt| if (pt.kind == .exl3 and m * pt.n * 2 > f.tmp_bytes and pt.n != e.n and m != 1) return false; + const s = p.stream; + try f.ev_fork.record(s); + var b: usize = 0; + var first = true; + for (e.parts) |pt| { + const w = switch (pt.kind) { + .exl3 => |w| w, + .native => { + try runMain(p, e, pt, x, m, out); + continue; + }, + }; + if (first or b == branches) { + first = false; + try runMain(p, e, pt, x, m, out); + continue; + } + const st = f.side[b]; + try st.wait(f.ev_fork); + const de = &f.dense[b]; + if (pt.n == e.n or m == 1) { + try de.decode(st, w, x, .bf16, m, out + pt.col * 2, .bf16); + } else { + const tmp = f.tmp[b].ptr; + try de.decode(st, w, x, .bf16, m, tmp, .bf16); + try copy2d(st, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); + } + try f.ev_join[b].record(st); + b += 1; + } + for (f.ev_join[0..b]) |ev| try s.wait(ev); + return true; +} diff --git a/zig/src/families/flashnext/cuda_exl3_frequency.zig b/zig/src/families/flashnext/cuda_exl3_frequency.zig new file mode 100644 index 0000000000..2ff5ea05ee --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_frequency.zig @@ -0,0 +1,135 @@ +const std = @import("std"); +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const Identity = struct { + model: [32]u8, + world: u32, + rank: u32, +}; + +pub const decoder_layers = 48; +pub const experts = 512; +pub const format_version: u32 = 1; + +pub const Counts = struct { + decoder: [decoder_layers][experts]u32 = @splat(@splat(0)), + mtp: [experts]u32 = @splat(0), + mtp_present: bool = false, +}; + +const magic = "TFEXL3FQ"; +const header_bytes = 64; +const decoder_bytes = decoder_layers * experts * @sizeOf(u32); +const mtp_bytes = experts * @sizeOf(u32); +pub const max_file_bytes = header_bytes + decoder_bytes + mtp_bytes + Sha256.digest_length; + +pub fn read(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity) !Counts { + try validateIdentity(identity); + const file: Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDONLY, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + }, 0), + .flags = .{ .nonblocking = true }, + }; + defer file.close(io); + const stat = try file.stat(io); + if (stat.kind != .file) return error.ProfileNotRegular; + if (stat.size > max_file_bytes) return error.ProfileTooLarge; + const buffer = try gpa.alloc(u8, max_file_bytes + 1); + defer gpa.free(buffer); + var reader = file.reader(io, &.{}); + const length = reader.interface.readSliceShort(buffer) catch return reader.err.?; + if (length > max_file_bytes) return error.ProfileTooLarge; + return decode(buffer[0..length], identity); +} + +pub fn write(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity, counts: *const Counts) !void { + try validateIdentity(identity); + const basename = std.fs.path.basename(path); + if (path.len == 0 or std.mem.endsWith(u8, path, "/") or + std.mem.eql(u8, basename, ".") or std.mem.eql(u8, basename, "..")) return error.InvalidPath; + const length = header_bytes + decoder_bytes + (if (counts.mtp_present) @as(usize, mtp_bytes) else 0) + Sha256.digest_length; + const bytes = try gpa.alloc(u8, length); + defer gpa.free(bytes); + encode(bytes, identity, counts); + var atomic = try Io.Dir.cwd().createFileAtomic(io, path, .{ + .permissions = .fromMode(0o600), + .make_path = false, + .replace = true, + }); + defer atomic.deinit(io); + try atomic.file.writePositionalAll(io, bytes, 0); + try atomic.file.sync(io); + try atomic.replace(io); +} + +fn validateIdentity(identity: Identity) !void { + if (identity.world == 0 or identity.rank >= identity.world) return error.InvalidIdentity; +} + +fn encode(bytes: []u8, identity: Identity, counts: *const Counts) void { + @memcpy(bytes[0..8], magic); + putWord(bytes, 8, format_version); + putWord(bytes, 12, @intFromBool(counts.mtp_present)); + @memcpy(bytes[16..48], &identity.model); + putWord(bytes, 48, identity.world); + putWord(bytes, 52, identity.rank); + putWord(bytes, 56, decoder_layers); + putWord(bytes, 60, experts); + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |count| { + putWord(bytes, offset, count); + offset += 4; + }; + if (counts.mtp_present) for (&counts.mtp) |count| { + putWord(bytes, offset, count); + offset += 4; + }; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..offset], &checksum, .{}); + @memcpy(bytes[offset..], &checksum); +} + +fn decode(bytes: []const u8, identity: Identity) !Counts { + if (bytes.len < header_bytes + Sha256.digest_length) return error.TruncatedProfile; + if (!std.mem.eql(u8, bytes[0..8], magic)) return error.MalformedProfile; + if (getWord(bytes, 8) != format_version) return error.UnsupportedVersion; + const flags = getWord(bytes, 12); + if (flags > 1 or getWord(bytes, 56) != decoder_layers or getWord(bytes, 60) != experts) return error.MalformedProfile; + const payload_end = header_bytes + decoder_bytes + (if (flags == 1) @as(usize, mtp_bytes) else 0); + const expected_length = payload_end + Sha256.digest_length; + if (bytes.len < expected_length) return error.TruncatedProfile; + if (bytes.len != expected_length) return error.MalformedProfile; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..payload_end], &checksum, .{}); + if (!std.mem.eql(u8, bytes[payload_end..], &checksum)) return error.ChecksumMismatch; + if (!std.mem.eql(u8, bytes[16..48], &identity.model) or + getWord(bytes, 48) != identity.world or getWord(bytes, 52) != identity.rank) return error.IdentityMismatch; + var counts: Counts = .{ .mtp_present = flags == 1 }; + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |*count| { + count.* = getWord(bytes, offset); + offset += 4; + }; + if (counts.mtp_present) for (&counts.mtp) |*count| { + count.* = getWord(bytes, offset); + offset += 4; + }; + return counts; +} + +fn putWord(bytes: []u8, offset: usize, value: u32) void { + std.mem.writeInt(u32, bytes[offset..][0..4], value, .little); +} + +fn getWord(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +test { + _ = @import("cuda_exl3_frequency_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_frequency_state.zig b/zig/src/families/flashnext/cuda_exl3_frequency_state.zig new file mode 100644 index 0000000000..c92eaca83b --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_frequency_state.zig @@ -0,0 +1,541 @@ +//! Lifetime placement priors: routed/metadata fingerprint; startup warmth may be discarded by prefill. +const std = @import("std"); +const codec = @import("cuda_exl3_frequency.zig"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const namespace = "flashnext-exl3-placement-v1"; + +pub const State = struct { + gpa: std.mem.Allocator, + io: std.Io, + path: []u8, + identity: codec.Identity, + lock: std.Io.File, + + pub fn deinit(state: *State) void { + state.lock.close(state.io); + state.gpa.free(state.path); + state.* = undefined; + } + + pub fn save(state: *State, own: anytype) void { + const counts = snapshot(own, @import("cuda_exl3_dplan.zig").current()) catch |err| { + std.log.err("EXL3 placement profile rank {d}: unsafe snapshot skipped ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + codec.write(state.gpa, state.io, state.path, state.identity, &counts) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: save failed ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + std.log.info("EXL3 placement profile rank {d}: saved decoder/MTP lifetime placement snapshot to {s}", .{ state.identity.rank, state.path }); + } +}; + +pub const Prepared = struct { + state: State, + counts: ?codec.Counts, +}; + +pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, own: anytype, world: u32, rank: u32) !?State { + const prepared = try prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], own.head, world, rank) orelse return null; + var state = prepared.state; + errdefer state.deinit(); + if (prepared.counts) |*counts| try finish(own, counts); + return state; +} + +pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, head: anytype, world: u32, rank: u32) !?Prepared { + const root = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_DIR") orelse return null; + if (std.mem.span(root).len == 0) return null; + if (comptime @import("builtin").os.tag == .windows) { + std.log.warn("EXL3 placement profiles require POSIX storage; persistence disabled", .{}); + return null; + } + if (!available(layers)) { + std.log.warn("EXL3 placement profile rank {d}: unavailable decoder pools; persistence disabled", .{rank}); + return null; + } + const before = std.Io.Clock.awake.now(io).toNanoseconds(); + const model = fingerprint(gpa, io, dir, layers, head) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: fingerprint unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + std.log.info("EXL3 placement profile rank {d}: routed-bank/metadata fingerprint took {d:.3} s (excludes router/dense weights)", .{ rank, @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - before)) / 1e9 }); + var state = open(gpa, io, std.mem.span(root), .{ .model = model, .world = world, .rank = rank }) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: lock/storage unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + errdefer state.deinit(); + var counts: ?codec.Counts = codec.read(gpa, io, state.path, state.identity) catch |err| blk: { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ rank, @errorName(err) }); + break :blk null; + }; + if (counts) |*value| { + if (value.mtp_present != (mtpExperts(head) != null)) { + std.log.warn("EXL3 placement profile rank {d}: cold start (MTP schema mismatch)", .{rank}); + counts = null; + } else { + decay(value); + std.log.info("EXL3 placement profile rank {d}: loaded and decayed rank-local placement counts before pool allocation", .{rank}); + } + } + return .{ .state = state, .counts = counts }; +} + +pub fn finish(own: anytype, counts: *const codec.Counts) !void { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool orelse return error.ProfileUnavailable; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + if (mtpPool(own.head)) |pool| pool.counts[pool.opt.layer] = counts.mtp; + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + if (mtpPool(own.head)) |pool| warmed += try pool.warmProfile(); + std.log.info("EXL3 placement profile: imported and warmed {d} expert lines after pool allocation", .{warmed}); +} + +fn open(gpa: std.mem.Allocator, io: std.Io, root: []const u8, identity: codec.Identity) !State { + const private = try std.fs.path.join(gpa, &.{ root, namespace }); + defer gpa.free(private); + var directory = try std.Io.Dir.cwd().createDirPathOpen(io, private, .{ + .permissions = .fromMode(0o700), + .open_options = .{ .iterate = true }, + }); + defer directory.close(io); + try directory.setPermissions(io, .fromMode(0o700)); + const name = try std.fmt.allocPrint(gpa, "{s}-w{d}-r{d}.bin", .{ std.fmt.bytesToHex(&identity.model, .lower), identity.world, identity.rank }); + defer gpa.free(name); + const path = try std.fs.path.join(gpa, &.{ private, name }); + errdefer gpa.free(path); + const lock_path = try std.fmt.allocPrint(gpa, "{s}.lock", .{path}); + defer gpa.free(lock_path); + const lock = try openLock(io, lock_path); + errdefer lock.close(io); + try lock.setPermissions(io, .fromMode(0o600)); + return .{ .gpa = gpa, .io = io, .path = path, .identity = identity, .lock = lock }; +} + +fn openLock(io: std.Io, path: []const u8) !std.Io.File { + const file: std.Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(std.Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDWR, + .CREAT = true, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + .NOFOLLOW = true, + }, 0o600), + .flags = .{ .nonblocking = true }, + }; + errdefer file.close(io); + if ((try file.stat(io)).kind != .file) return error.ProfileLockNotRegular; + if (!try file.tryLock(io, .exclusive)) return error.WouldBlock; + return file; +} + +fn restore(state: *State, own: anytype) !bool { + var counts = codec.read(state.gpa, state.io, state.path, state.identity) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ state.identity.rank, @errorName(err) }); + return false; + }; + if (counts.mtp_present != (mtpPool(own.head) != null)) { + std.log.warn("EXL3 placement profile rank {d}: cold start (MTP schema mismatch)", .{state.identity.rank}); + return false; + } + decay(&counts); + const layers = own.engine.layers[0..own.engine.n_layers]; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + if (mtpPool(own.head)) |pool| pool.counts[pool.opt.layer] = counts.mtp; + std.log.info("EXL3 placement profile rank {d}: imported decoder/MTP lifetime counts, decayed once by 2 (nonzero minimum 1)", .{state.identity.rank}); + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + if (mtpPool(own.head)) |pool| warmed += try pool.warmProfile(); + std.log.info("EXL3 placement profile rank {d}: warmed {d} expert lines", .{ state.identity.rank, warmed }); + return true; +} + +fn available(layers: anytype) bool { + if (layers.len != 48) return false; + for (layers) |layer| { + const experts = layer.exl3 orelse return false; + if (experts.bank == null or experts.routed.len != codec.experts) return false; + if (experts.pool) |pool| if (pool.opt.layer >= 48) return false; + } + return true; +} + +fn mtpPool(head: anytype) @TypeOf(head.?.layer.exl3.?.pool) { + const present = head orelse return null; + const experts = present.layer.exl3 orelse return null; + return experts.pool; +} + +fn mtpExperts(head: anytype) @TypeOf(head.?.layer.exl3) { + const present = head orelse return null; + return present.layer.exl3; +} + +fn decay(counts: *codec.Counts) void { + for (&counts.decoder) |*layer| decayRow(layer); + decayRow(&counts.mtp); +} + +fn decayRow(row: *[512]u32) void { + for (row) |*count| { + if (count.* != 0) count.* = @max(1, count.* / 2); + } +} + +fn snapshot(own: anytype, planner: anytype) !codec.Counts { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + try own.stream.synchronize(); + if (planner) |current| { + if (current.indexOf(layers[0].exl3.?.pool.?) != null) { + for (current.pools) |pool| { + var ours = pool == mtpPool(own.head); + for (layers) |layer| ours = ours or pool == layer.exl3.?.pool.?; + if (!ours) return error.ProfilePlannerOwnerMismatch; + } + try current.toHost(); + } + } + var counts: codec.Counts = .{}; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + try pool.lines.drain(); + counts.decoder[index] = pool.counts[pool.opt.layer]; + } + if (mtpPool(own.head)) |pool| { + try pool.lines.drain(); + counts.mtp = pool.counts[pool.opt.layer]; + counts.mtp_present = true; + } + return counts; +} + +fn hashBytes(hash: *Sha256, bytes: []const u8) void { + var length: [8]u8 = undefined; + std.mem.writeInt(u64, &length, bytes.len, .little); + hash.update(&length); + hash.update(bytes); +} + +fn hashRouted(hash: *Sha256, routed: []const [3][]const u8) void { + for (routed) |expert| for (expert) |projection| hashBytes(hash, projection); +} + +fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, head: anytype) ![32]u8 { + var hash: Sha256 = .init(.{}); + hashBytes(&hash, namespace); + for ([_][]const u8{ "config.json", "model.safetensors.index.json" }) |name| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(bytes); + hashBytes(&hash, name); + hashBytes(&hash, bytes); + } + hashBytes(&hash, "decoder-routed-trellis"); + for (layers) |layer| hashRouted(&hash, expertRouted(layer.exl3.?)); + if (mtpExperts(head)) |experts| { + hashBytes(&hash, "mtp-routed-trellis"); + hashRouted(&hash, expertRouted(experts)); + } + return hash.finalResult(); +} + +fn expertRouted(experts: anytype) []const [3][]const u8 { + if (experts.pool) |pool| return pool.src.routed; + return experts.routed; +} + +const MockPool = struct { + counts: *[48][512]u32, + opt: struct { layer: u8, lines: usize = 10 }, + src: struct { routed: []const [3][]const u8 = &.{.{ "gate", "up", "down" }} } = .{}, + lines: struct { + drained: bool = false, + pub fn drain(copies: *@This()) !void { + copies.drained = true; + } + } = .{}, + warm_calls: usize = 0, + fail_warm: bool = false, + + pub fn warmProfile(pool: *MockPool) !usize { + pool.warm_calls += 1; + if (pool.fail_warm) return error.MockCudaWarmFailure; + return 1; + } +}; + +const mock_routed: [512][3][]const u8 = @splat(.{ "gate", "up", "down" }); +const MockExperts = struct { pool: ?*MockPool, bank: ?[]const u8 = "bank", routed: []const [3][]const u8 = &mock_routed }; +const MockLayer = struct { exl3: ?*MockExperts }; +const Mock = struct { + decoder: [48][512]u32 = @splat(@splat(0)), + mtp: [48][512]u32 = @splat(@splat(0)), + pools: [49]MockPool = undefined, + experts: [49]MockExperts = undefined, + head: struct { layer: MockLayer } = undefined, + engine: struct { layers: [48]MockLayer, n_layers: usize = 48 } = undefined, + + fn init(mock: *Mock) void { + mock.* = .{}; + mock.engine.n_layers = 48; + for (&mock.pools, &mock.experts, 0..) |*pool, *experts, index| { + pool.* = .{ .counts = if (index == 48) &mock.mtp else &mock.decoder, .opt = .{ .layer = if (index == 48) 0 else @intCast(index) }, .src = .{ .routed = &mock_routed } }; + experts.* = .{ .pool = pool }; + if (index < 48) mock.engine.layers[index] = .{ .exl3 = experts }; + } + mock.head = .{ .layer = .{ .exl3 = &mock.experts[48] } }; + } + + const Stream = struct { + syncs: usize = 0, + fail: bool = false, + pub fn synchronize(stream: *@This()) !void { + if (stream.fail) return error.MockCudaSyncFailure; + stream.syncs += 1; + } + }; + fn owner(mock: *Mock, drafting: bool) struct { engine: *@TypeOf(mock.engine), head: ?*@TypeOf(mock.head), stream: Stream } { + return .{ .engine = &mock.engine, .head = if (drafting) &mock.head else null, .stream = .{} }; + } +}; + +const MockPlan = struct { + pools: []const *MockPool, + pulls: usize = 0, + fail: bool = false, + pub fn indexOf(plan: *MockPlan, pool: *const MockPool) ?usize { + for (plan.pools, 0..) |held, index| if (held == pool) return index; + return null; + } + pub fn toHost(plan: *MockPlan) !void { + if (plan.fail) return error.MockCudaPullFailure; + plan.pulls += 1; + for (plan.pools, 0..) |pool, index| pool.counts[pool.opt.layer][7] = @intCast(100 + index); + } +}; + +const test_identity: codec.Identity = .{ .model = @splat(42), .world = 2, .rank = 0 }; + +test "exl3: placement state lifetime lock conflicts, rank separation and private storage" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + var held = true; + defer if (held) state.deinit(); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + const empty: codec.Counts = .{}; + try codec.write(allocator, io, state.path, state.identity, &empty); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + var rank_identity = test_identity; + rank_identity.rank = 1; + var other_rank = try open(allocator, io, root, rank_identity); + defer other_rank.deinit(); + try std.testing.expect(!std.mem.eql(u8, state.path, other_rank.path)); + const lock_stat = try state.lock.stat(io); + try std.testing.expectEqual(@as(u32, 0o600), lock_stat.permissions.toMode() & 0o777); + const directory = try std.Io.Dir.cwd().openDir(io, std.fs.path.dirname(state.path).?, .{}); + defer directory.close(io); + const directory_stat = try directory.stat(io); + try std.testing.expectEqual(@as(u32, 0o700), directory_stat.permissions.toMode() & 0o777); + state.deinit(); + held = false; + var reopened = try open(allocator, io, root, test_identity); + defer reopened.deinit(); +} + +test "exl3: placement state cold IO and MTP mismatch never import or warm; CUDA warm errors propagate" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(true); + try std.testing.expect(!try restore(&state, &own)); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = state.path, .data = "corrupt" }); + try std.testing.expect(!try restore(&state, &own)); + var counts: codec.Counts = .{}; + counts.decoder[0][1] = 123; + try codec.write(allocator, io, state.path, state.identity, &counts); + try std.testing.expect(!try restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 0), mock.decoder[0][1]); + try std.testing.expectEqual(@as(usize, 0), mock.pools[0].warm_calls); + counts.mtp_present = true; + try codec.write(allocator, io, state.path, state.identity, &counts); + mock.pools[0].fail_warm = true; + try std.testing.expectError(error.MockCudaWarmFailure, restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 61), mock.decoder[0][1]); + mock.engine.n_layers = 47; + try std.testing.expect(!available(mock.engine.layers[0..mock.engine.n_layers])); +} + +test "exl3: placement state decay and replacement snapshot do not add imported prior twice" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var counts: codec.Counts = .{ .mtp_present = true }; + counts.decoder[0][0..5].* = .{ 0, 1, 2, 3, std.math.maxInt(u32) }; + counts.mtp[1] = 9; + try codec.write(allocator, io, state.path, state.identity, &counts); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(true); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1, 1, 2147483647 }, mock.decoder[0][0..5]); + try std.testing.expectEqual(@as(u32, 4), mock.mtp[0][1]); + for (mock.pools) |pool| try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + mock.decoder[0][1] += 8; + mock.mtp[0][1] += 2; + const saved = try snapshot(&own, @as(?*MockPlan, null)); + try codec.write(allocator, io, state.path, state.identity, &saved); + const disk = try codec.read(allocator, io, state.path, state.identity); + try std.testing.expectEqual(@as(u32, 9), disk.decoder[0][1]); + mock.init(); + own = mock.owner(true); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 4), mock.decoder[0][1]); + try std.testing.expectEqual(@as(u32, 3), mock.mtp[0][1]); + const missing = try std.fs.path.join(allocator, &.{ root, "missing", "profile" }); + defer allocator.free(missing); + try std.testing.expectError(error.FileNotFound, codec.write(allocator, io, missing, state.identity, &saved)); +} + +test "exl3: placement state pulls only owning planner before copying decoder and MTP counts" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(true); + var owned: [49]*MockPool = undefined; + for (&mock.pools, &owned) |*pool, *target| target.* = pool; + var plan: MockPlan = .{ .pools = &owned }; + const counts = try snapshot(&own, @as(?*MockPlan, &plan)); + try std.testing.expectEqual(@as(usize, 1), own.stream.syncs); + try std.testing.expectEqual(@as(usize, 1), plan.pulls); + try std.testing.expectEqual(@as(u32, 100), counts.decoder[0][7]); + try std.testing.expectEqual(@as(u32, 147), counts.decoder[47][7]); + try std.testing.expectEqual(@as(u32, 148), counts.mtp[7]); + for (mock.pools) |pool| try std.testing.expect(pool.lines.drained); + plan.fail = true; + try std.testing.expectError(error.MockCudaPullFailure, snapshot(&own, @as(?*MockPlan, &plan))); + own.stream.fail = true; + try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own, @as(?*MockPlan, &plan))); + own.stream.fail = false; + var other: Mock = undefined; + other.init(); + const foreign = [_]*MockPool{&other.pools[0]}; + var unrelated: MockPlan = .{ .pools = &foreign, .fail = true }; + _ = try snapshot(&own, @as(?*MockPlan, &unrelated)); + try std.testing.expectEqual(@as(usize, 0), unrelated.pulls); + const mixed = [_]*MockPool{ &mock.pools[0], &other.pools[0] }; + unrelated.pools = &mixed; + try std.testing.expectError(error.ProfilePlannerOwnerMismatch, snapshot(&own, @as(?*MockPlan, &unrelated))); +} + +test "exl3: prepared placement counts import and warm only after pools attach" { + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(true); + var counts: codec.Counts = .{}; + for (&counts.decoder, 0..) |*row, layer| row[7] = @intCast(layer + 1); + counts.mtp[7] = 73; + counts.mtp_present = true; + + try finish(&own, &counts); + + for (mock.pools[0..48], 0..) |pool, layer| { + try std.testing.expectEqual(@as(u32, @intCast(layer + 1)), mock.decoder[layer][7]); + try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + } + try std.testing.expectEqual(@as(u32, 73), mock.mtp[0][7]); + try std.testing.expectEqual(@as(usize, 1), mock.pools[48].warm_calls); +} + +test "exl3: placement fingerprint reads routed and metadata bytes, ignores pool size and distinguishes MTP" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); + var mock: Mock = undefined; + mock.init(); + const own = mock.owner(false); + const original = try fingerprint(allocator, io, root, &mock.engine.layers, own.head); + for (&mock.experts) |*experts| experts.pool = null; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers, own.head)); + for (&mock.experts, &mock.pools) |*experts, *pool| experts.pool = pool; + mock.pools[0].opt.lines = 23; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers, own.head)); + mock.pools[47].src.routed = &.{.{ "changed gate", "up", "down" }}; + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); + mock.pools[47].src.routed = &.{.{ "gate", "up", "down" }}; + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "changed config" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "changed index" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); + const drafting = mock.owner(true); + const with_mtp = try fingerprint(allocator, io, root, &mock.engine.layers, drafting.head); + try std.testing.expect(!std.mem.eql(u8, &original, &with_mtp)); + mock.pools[48].src.routed = &.{.{ "changed mtp", "up", "down" }}; + try std.testing.expect(!std.mem.eql(u8, &with_mtp, &(try fingerprint(allocator, io, root, &mock.engine.layers, drafting.head)))); +} + +test "exl3: placement lock rejects nonregular storage without waiting for FIFO readers" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(path); + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, "/dev/null")); + const fifo = try std.fs.path.join(allocator, &.{ path, "profile.lock" }); + defer allocator.free(fifo); + const terminated = try allocator.dupeSentinel(u8, fifo, 0); + defer allocator.free(terminated); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "profile.lock", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, fifo)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_frequency_test.zig b/zig/src/families/flashnext/cuda_exl3_frequency_test.zig new file mode 100644 index 0000000000..7d2b2a0f7a --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_frequency_test.zig @@ -0,0 +1,316 @@ +const std = @import("std"); +const profile = @import("cuda_exl3_frequency.zig"); +const Io = std.Io; +const allocator = std.testing.allocator; +const io = std.testing.io; +const identity: profile.Identity = .{ .model = @splat(0xa7), .world = 2, .rank = 1 }; +const header_bytes = 64; +const checksum_bytes = 32; +const decoder_file_bytes = header_bytes + 48 * 512 * 4 + checksum_bytes; + +fn profilePath(directory: Io.Dir) ![]u8 { + const root = try directory.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + return std.fs.path.join(allocator, &.{ root, "frequency.bin" }); +} + +fn sample(counts: *profile.Counts, mtp_present: bool) void { + counts.* = .{ .mtp_present = mtp_present }; + for (&counts.decoder, 0..) |*layer, layer_index| for (layer, 0..) |*count, expert_index| { + count.* = @intCast((layer_index + 1) * 100_000 + expert_index); + }; + for (&counts.mtp, 0..) |*count, expert_index| count.* = @intCast(7_000_000 + expert_index); + counts.decoder[0][0] = 0; + counts.decoder[47][511] = std.math.maxInt(u32); + counts.mtp[511] = std.math.maxInt(u32); +} + +fn expectCounts(expected: *const profile.Counts, actual: *const profile.Counts) !void { + try std.testing.expectEqual(expected.mtp_present, actual.mtp_present); + for (&expected.decoder, &actual.decoder) |*expected_layer, *actual_layer| { + try std.testing.expectEqualSlices(u32, expected_layer, actual_layer); + } + if (expected.mtp_present) { + try std.testing.expectEqualSlices(u32, &expected.mtp, &actual.mtp); + } else { + for (actual.mtp) |count| try std.testing.expectEqual(@as(u32, 0), count); + } +} + +fn expectOnlyProfile(directory: Io.Dir) !void { + var iterator = directory.iterate(); + var entries: usize = 0; + while (try iterator.next(io)) |entry| { + entries += 1; + try std.testing.expectEqualStrings("frequency.bin", entry.name); + } + try std.testing.expectEqual(@as(usize, 1), entries); +} + +fn bytesAt(path: []const u8) ![]u8 { + return Io.Dir.cwd().readFileAlloc(io, path, allocator, .limited(profile.max_file_bytes + 2)); +} + +fn expectPreserved(path: []const u8, before: []const u8) !void { + const after = try bytesAt(path); + defer allocator.free(after); + try std.testing.expectEqualSlices(u8, before, after); +} + +test "exl3: frequency full decoder and optional MTP round trip with private permissions" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + for ([_]bool{ false, true, false }) |mtp_present| { + sample(&counts, mtp_present); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + const stat = try temporary.dir.statFile(io, "frequency.bin", .{}); + try std.testing.expectEqual(@as(u64, if (mtp_present) profile.max_file_bytes else decoder_file_bytes), stat.size); + try std.testing.expectEqual(@as(std.posix.mode_t, 0o600), stat.permissions.toMode() & 0o777); + const bytes = try bytesAt(path); + defer allocator.free(bytes); + try std.testing.expectEqualSlices(u8, "TFEXL3FQ", bytes[0..8]); + try std.testing.expectEqualSlices(u8, &identity.model, bytes[16..48]); + try std.testing.expectEqualSlices(u8, &.{ 1, 0, 0, 0 }, bytes[8..12]); + try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, bytes[64..68]); + try std.testing.expectEqualSlices(u8, &.{ 0xa1, 0x86, 0x01, 0x00 }, bytes[68..72]); + try std.testing.expectEqualSlices(u8, &.{ 0xff, 0xff, 0xff, 0xff }, bytes[64 + 48 * 512 * 4 - 4 ..][0..4]); + try expectOnlyProfile(temporary.dir); + } +} + +test "exl3: frequency missing profile and every exact identity mismatch are recoverable" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + for (0..32) |model_index| { + var different = identity; + different.model[model_index] ^= 1; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + } + var different = identity; + different.world = 3; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + different = identity; + different.rank = 0; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency rejects truncated oversized malformed and corrupt profiles" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts, true); + try profile.write(allocator, io, path, identity, &counts); + const original = try bytesAt(path); + defer allocator.free(original); + for ([_]usize{ 0, 7, 63, 64, 95, 96, decoder_file_bytes, profile.max_file_bytes - 1 }) |length| { + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = original[0..length] }); + try std.testing.expectError(error.TruncatedProfile, profile.read(allocator, io, path, identity)); + } + const changed = try allocator.dupe(u8, original); + defer allocator.free(changed); + const Mutation = struct { offset: usize, value: u8, expected: anyerror }; + for ([_]Mutation{ + .{ .offset = 0, .value = 0, .expected = error.MalformedProfile }, + .{ .offset = 8, .value = 2, .expected = error.UnsupportedVersion }, + .{ .offset = 12, .value = 2, .expected = error.MalformedProfile }, + .{ .offset = 56, .value = 47, .expected = error.MalformedProfile }, + .{ .offset = 60, .value = 1, .expected = error.MalformedProfile }, + .{ .offset = 16, .value = 0, .expected = error.ChecksumMismatch }, + .{ .offset = 48, .value = 3, .expected = error.ChecksumMismatch }, + .{ .offset = 64, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = 64 + 48 * 512 * 4, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = profile.max_file_bytes - 1, .value = original[profile.max_file_bytes - 1] ^ 1, .expected = error.ChecksumMismatch }, + }) |mutation| { + @memcpy(changed, original); + changed[mutation.offset] = mutation.value; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = changed }); + try std.testing.expectError(mutation.expected, profile.read(allocator, io, path, identity)); + } + const oversized = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(oversized); + @memcpy(oversized[0..original.len], original); + oversized[original.len] = 0; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = oversized }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(allocator, io, path, identity)); + counts.mtp_present = false; + try profile.write(allocator, io, path, identity, &counts); + const file = try temporary.dir.openFile(io, "frequency.bin", .{ .mode = .read_write }); + defer file.close(io); + try file.writePositionalAll(io, &.{0}, decoder_file_bytes); + try std.testing.expectError(error.MalformedProfile, profile.read(allocator, io, path, identity)); +} + +fn failWrite(userdata: ?*anyopaque, file: Io.File, header: []const u8, data: []const []const u8, splat: usize, offset: u64) Io.File.WritePositionalError!usize { + if (offset != 0) return error.NoSpaceLeft; + const prefix = data[0][0..13]; + return io.vtable.fileWritePositional(userdata, file, header, &.{prefix}, splat, offset); +} + +fn failSync(_: ?*anyopaque, file: Io.File) Io.File.SyncError!void { + const stat = file.stat(io) catch return error.InputOutput; + if (stat.permissions.toMode() & 0o777 != 0o600 or stat.size != profile.max_file_bytes) return error.AccessDenied; + return error.InputOutput; +} + +fn failRename(_: ?*anyopaque, _: Io.Dir, _: []const u8, _: Io.Dir, _: []const u8) Io.Dir.RenameError!void { + return error.AccessDenied; +} + +test "exl3: frequency partial write sync and rename failures preserve previous bytes and remove temporaries" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + sample(&counts, true); + for (0..3) |failure_index| { + var failing_vtable = io.vtable.*; + const expected: anyerror = switch (failure_index) { + 0 => failure: { + failing_vtable.fileWritePositional = failWrite; + break :failure error.NoSpaceLeft; + }, + 1 => failure: { + failing_vtable.fileSync = failSync; + break :failure error.InputOutput; + }, + else => failure: { + failing_vtable.dirRename = failRename; + break :failure error.AccessDenied; + }, + }; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(expected, profile.write(allocator, failing_io, path, identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); + const restored = try profile.read(allocator, io, path, identity); + try std.testing.expectEqual(@as(u32, 0), restored.decoder[47][511]); + } + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency allocation and invalid identity failures preserve previous profile" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, profile.write(failing.allocator(), io, path, identity, &counts)); + try std.testing.expectError(error.OutOfMemory, profile.read(failing.allocator(), io, path, identity)); + for ([_]profile.Identity{ + .{ .model = identity.model, .world = 0, .rank = 0 }, + .{ .model = identity.model, .world = 2, .rank = 2 }, + }) |invalid| { + try std.testing.expectError(error.InvalidIdentity, profile.write(allocator, io, path, invalid, &counts)); + try std.testing.expectError(error.InvalidIdentity, profile.read(allocator, io, path, invalid)); + } + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, ".", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "..", identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency read allocation stays fixed and oversized input needs no allocation" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const single: profile.Identity = .{ .model = @splat(0), .world = 1, .rank = 0 }; + var counts: profile.Counts = .{}; + sample(&counts, true); + try profile.write(allocator, io, path, single, &counts); + const backing = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(backing); + var bounded = std.heap.FixedBufferAllocator.init(backing); + const restored = try profile.read(bounded.allocator(), io, path, single); + try expectCounts(&counts, &restored); + const file = try temporary.dir.openFile(io, "frequency.bin", .{ .mode = .read_write }); + defer file.close(io); + try file.setLength(io, profile.max_file_bytes + 1); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(failing.allocator(), io, path, single)); + try std.testing.expect(!failing.has_induced_failure); +} + +test "exl3: frequency failed first save leaves no destination or temporary file" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts, true); + var failing_vtable = io.vtable.*; + failing_vtable.fileSync = failSync; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(error.InputOutput, profile.write(allocator, failing_io, path, identity, &counts)); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + var iterator = temporary.dir.iterate(); + try std.testing.expectEqual(@as(?Io.Dir.Entry, null), try iterator.next(io)); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency rejects a FIFO without a writer and other nonregular input" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const terminated_path = try allocator.dupeSentinel(u8, path, 0); + defer allocator.free(terminated_path); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "frequency.bin", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated_path, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileNotRegular, profile.read(failing.allocator(), io, path, identity)); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, "/dev/null", identity)); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, root, identity)); + try std.testing.expect(!failing.has_induced_failure); +} diff --git a/zig/src/families/flashnext/cuda_exl3_graph.zig b/zig/src/families/flashnext/cuda_exl3_graph.zig new file mode 100644 index 0000000000..97c6e11d30 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_graph.zig @@ -0,0 +1,2268 @@ +//! CUDA graphs for an EXL3 Flash Next decode step with the expert offload pool (TF_FLASHNEXT_EXL3_GRAPH=1). +//! +//! With the pool planned on the device (cuda_exl3_dplan.zig) a decode step has no host wait between its first launch +//! and its logits, so the step is captured once and replayed: a decode round (Engine.round, one row a stream) and an +//! MTP verify window (Engine.forward with every row's logits, up to the planner's row bound) each take ONE graph +//! launch; at tp=2 the graph also holds the step's NCCL sums and the head's gathers (both ranks capture the same +//! calls in the same order, the rank engine keeps them in lockstep). +//! +//! What a replay may change and how (the graph-safety rules of exl3_graph.zig): +//! - token ids, positions and the n-gram rows: pinned host bytes the graph's copies read, rewritten by `prepare` +//! before each launch (upstream's cuda_round_graph.zig does the same for its tp round); +//! - the experts each layer reads: the pool's pointer tables, rewritten in place by the device planner; +//! - everything else (weights, scratch, caches, the sequence's buffers, the logits buffer) sits at fixed addresses, +//! which the key holds: a different address, row count, attention chunk count or score-tile count captures anew. +//! A step the graph cannot take (more rows than the planner's bound, media) runs eagerly with the host cut. +//! +//! TF_FLASHNEXT_EXL3_DPLAN=1 alone plans on the device without capturing (the same kernels launched one by one). + +const std = @import("std"); +const cuda = @import("cuda"); +const dplan = @import("cuda_exl3_dplan.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const collective = @import("cuda_collective.zig"); +const wire = @import("cuda_rank_wire.zig"); +const profile_module = @import("exl3_graph_profile.zig"); +const state_binding = @import("exl3_state_binding.zig"); +const state_driver = @import("cuda_exl3_state_binding.zig"); +pub var profile: profile_module.Profile(64, 1024) = .{}; +fn profileEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_PROFILE") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} +pub const memory = @import("cuda_exl3_graph_memory.zig"); + +fn freeMemory(stream: cuda.Stream) !u64 { + var free: usize = 0; + var total: usize = 0; + try stream.d.check(stream.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +pub fn uploadResident(exec: *cuda.graph.Exec, stream: cuda.Stream) !void { + try exec.upload(stream); + try stream.synchronize(); +} + +pub fn reserveMemory(stream: cuda.Stream, class: memory.Class, cohort: u64) !memory.Ticket { + _ = try memory.startupReservation(); + while (true) { + return memory.admitWithQuery(class, cohort, stream, freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphMemoryHeadroom, error.GraphEvictionRequired => { + if (cohort != 0) return err; + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(graph retirement)"); + if (!evictOne()) return err; + continue; + }, + else => return err, + }; + } +} + +fn retirementBarrier(c: *collective.Collective, signature: u64, cohort: u64) !void { + const decision = try c.graphAgree(signature, .{ .state = .ready, .cohort = cohort }); + if (decision.choice != .replay) return error.GraphCohortMismatch; +} + +fn oldestCohort() u64 { + var cohort: u64 = 0; + var used: u64 = std.math.maxInt(u64); + for (entries) |maybe| if (maybe) |entry| { + if (entry.cohort != 0 and (entry.used < used or (entry.used == used and (cohort == 0 or entry.cohort < cohort)))) { + cohort = entry.cohort; + used = entry.used; + } + }; + return cohort; +} + +test "exl3: TP eviction proposes the least recently used cohort and ignores local graphs" { + entries = @splat(null); + defer entries = @splat(null); + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + entries[0] = .{ .key = key, .used = 30, .cohort = 3 }; + entries[1] = .{ .key = key, .used = 10, .cohort = 7 }; + entries[2] = .{ .key = key, .used = 0, .cohort = 0 }; + try std.testing.expectEqual(@as(u64, 7), oldestCohort()); + entries[0].?.used = 10; + try std.testing.expectEqual(@as(u64, 3), oldestCohort()); +} + +fn retireCohort(cohort: u64) void { + for (&entries) |*maybe| if (maybe.*) |*entry| { + if (entry.cohort == cohort) { + if (profileEnabled()) profile.evict(profile_module.Sample.fromKey(entry.key)); + drop(entry); + maybe.* = null; + } + }; +} + +fn reserveTogether(c: *collective.Collective, stream: cuda.Stream, signature: u64, cohort: u64) !memory.Ticket { + _ = try memory.startupReservation(); + var retired: usize = 0; + while (true) { + var ticket: ?memory.Ticket = memory.admitWithQuery(.decode, cohort, stream, freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphMemoryHeadroom, error.GraphEvictionRequired => null, + else => return err, + }; + errdefer if (ticket) |held| memory.release(held) catch unreachable; + var free_slot = false; + for (entries) |entry| free_slot = free_slot or entry == null; + const agreement = try c.graphAgree(signature ^ 0x61646d6974, .{ .state = if (ticket != null and free_slot) .ready else .eager, .cohort = cohort }); + if (agreement.choice == .replay) return ticket.?; + if (ticket) |held| { + try memory.release(held); + ticket = null; + } + const victim = try c.graphVictim(signature ^ 0x76696374696d, oldestCohort()); + if (victim == 0) { + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(TP local graph retirement)"); + try retirementBarrier(c, signature ^ 0x6c6f63616c7175, cohort); + const reclaimed = evictOne(); + const any_reclaimed = try c.graphVictim(signature ^ 0x6c6f63616c646f, @intFromBool(reclaimed)); + if (any_reclaimed == 0) return error.GraphBudgetExceeded; + continue; + } + if (retired >= entries.len * 2) return error.GraphBudgetExceeded; + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(TP graph retirement)"); + try retirementBarrier(c, signature ^ 0x71756965736365, victim); + retireCohort(victim); + try retirementBarrier(c, signature ^ 0x72657469726564, victim); + retired += 1; + } +} + +fn evictOne() bool { + var oldest: u64 = std.math.maxInt(u64); + var decode: ?usize = null; + var prefill: ?usize = null; + for (entries, 0..) |maybe, index| if (maybe) |entry| { + if (entry.cohort == 0 and entry.used < oldest) { + oldest = entry.used; + decode = index; + } + }; + for (pentries, 0..) |maybe, index| if (maybe) |entry| { + if (entry.used < oldest) { + oldest = entry.used; + decode = null; + prefill = index; + } + }; + if (decode) |index| { + if (profileEnabled()) profile.evict(profile_module.Sample.fromKey(entries[index].?.key)); + drop(&entries[index].?); + entries[index] = null; + return true; + } + if (prefill) |index| { + const entry = &pentries[index].?; + dropPrefill(entry); + pentries[index] = null; + return true; + } + return @import("cuda_exl3_chain.zig").evictOne(); +} + +const dims = wt.dims; +const vocab = wt.vocab; +const att_rows = @import("cuda_batch_state.zig").att_rows; + +/// TF_FLASHNEXT_EXL3_GRAPH=1: decode rounds and verify windows replay captured graphs. +pub fn enabled() bool { + if (force) |f| return f == .graph; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// A test's override of the environment: the host cut, the device plan launched eagerly, or graphs. +pub const Mode = enum { off, dplan, graph }; +pub var force: ?Mode = null; + +const max_lanes = 32; + +pub const Key = struct { + /// 1: a decode round; 2: a window. + kind: u8, + rows: usize, + /// Per lane (a window: lane 0): the sequence's buffers, its attention chunks and score tiles at the step's end. + seq: [max_lanes][5]u64 = @splat(@splat(0)), + paging: [max_lanes][6]u64 = @splat(@splat(0)), + paging_logical: ?[max_lanes][2]u64 = null, + chunks: [max_lanes]u32 = @splat(0), + score_tiles: [max_lanes]u32 = @splat(0), + /// The pass's scratch (its first buffer and its row capacity, which sets every other buffer's offset). + scratch: u64, + scratch_cap: usize = 0, + scores: u64, + output: u64, + streams: u64 = 0, + keep: u64 = 0, + keep_nonstate: u64 = 0, + /// The device planner's state (a new planner is new memory). + plan: u64 = 0, + + fn seams(end: usize, chunks: *u32, tiles: *u32) void { + chunks.* = @intCast((@min(end, wt.dense_keys) + 511) / 512); + tiles.* = if (end / wt.index_ratio > wt.index_top) @intCast(((end + 3) / 4 + 63) / 64) else 0; + } + + /// The DeltaNet state set is in the key: a drafting stream swaps it with its window's spare set (Draft.settle). + fn seqOf(q: anytype) [5]u64 { + return .{ if (q.kc.len > 0) q.kc[0].ptr else 0, if (q.conv.len > 0) q.conv[0].ptr else 0, q.tail.ptr, q.cap, if (q.state.len > 0) q.state[0].ptr else 0 }; + } +}; + +/// True while this module captures a step: the pass's pinned uploads become copy nodes (Uploads.copy). +pub var capturing: bool = false; + +/// A copy node from pinned host bytes (the pass's staging): the graph reads them at every replay, as `prepare` left +/// them. A pageable source fails the capture (the driver refuses it), so stale custody copies cannot hide here. +pub fn pinnedCopy(s: cuda.Stream, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +/// `template`: the graph an update took this exec's parameters from, kept while the exec uses them (a tp graph's +/// NCCL resources are tied to the graph they were captured in). +const Entry = struct { key: Key, exec: ?cuda.graph.Exec = null, used: u64, template: ?cuda.graph.Graph = null, pending: ?*Pending = null, cohort: u64 = 0, admission: ?memory.Ticket = null, stream: ?cuda.Stream = null, binding: ?state_binding.Binding = null, poisoned: bool = false }; + +fn stateRebindEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_STATE_REBIND") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +fn stateRequest(key: Key, pass: anytype) ?state_binding.Request { + if (!stateRebindEnabled()) return null; + return windowStateRequest(key, pass); +} + +fn windowStateRequest(key: Key, pass: anytype) ?state_binding.Request { + const Pass = @typeInfo(@TypeOf(pass)).pointer.child; + if (comptime !@hasField(Pass, "lanes")) return null; + if (key.kind != 2 or !state_binding.supportsWindow(key.rows) or pass.lanes.len != 0) return null; + const gdn = @import("cuda_gdn.zig"); + var request: state_binding.Request = .{ + .width = key.rows, + .chain = @intFromPtr(pass.k.chain.handle), + .chain_to = @intFromPtr(pass.k.chain_to.handle), + .geometry = .{ .grid = .{ gdn.nv, gdn.dv / 16, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = pass.s.gq, .gk = pass.s.gk, .gv = pass.s.gv, .ggb = pass.s.ggb, .rows = @intCast(key.rows), .ys = pass.s.ys }, + }; + for (pass.m.layers, 0..) |layer, index| { + if (!layer.linear) continue; + const linear = wt.linearIndex(index); + if (request.count == state_binding.max_layers or linear >= pass.seq.state.len) return null; + request.layers[request.count] = .{ + .layer = index, + .state_in = pass.seq.state[linear].ptr, + .state_out = if (pass.hooks.keep) |keep| keep.layers[index].state_out else 0, + }; + if (request.layers[request.count].state_out == 0) return null; + request.count += 1; + } + request.validate() catch return null; + return request; +} + +fn capturedBinding(key: Key, pass: anytype, graph: cuda.graph.Graph) ?state_binding.Binding { + const request = stateRequest(key, pass) orelse return null; + if (state_driver.nodeType(graph.d) == null) return null; + return state_binding.Binding.init(@intFromPtr(graph.handle), request) catch null; +} + +fn rebindCandidate(key: Key, request: state_binding.Request, tp: bool) ?*Entry { + var candidate: ?*Entry = null; + for (&entries) |*maybe| if (maybe.*) |*entry| { + if (entry.poisoned or entry.pending != null or entry.exec == null or entry.template == null or (tp and entry.cohort == 0)) continue; + const binding = entry.binding orelse continue; + if (binding.poisoned or !std.meta.eql(state_binding.normalized(entry.key), state_binding.normalized(key)) or !binding.request.compatible(request)) continue; + if (candidate == null or entry.cohort < candidate.?.cohort or (entry.cohort == candidate.?.cohort and entry.used > candidate.?.used)) candidate = entry; + }; + return candidate; +} + +fn prepareBinding(entry: *Entry) !?state_driver.Backend { + const backend = state_driver.Backend.init(entry.template.?, entry.exec.?) orelse return null; + entry.binding.?.load(backend) catch |err| switch (err) { + error.UnsupportedBinding => { + entry.binding.?.invalidate(); + entry.binding = null; + return null; + }, + else => return err, + }; + return backend; +} + +fn tryRebind(key: Key, pass: anytype) !bool { + const request = stateRequest(key, pass) orelse return false; + const entry = rebindCandidate(key, request, false) orelse { + stats.rebind_declined += 1; + return false; + }; + const prepared = prepareBinding(entry) catch |err| { + stats.rebind_failures += 1; + return err; + }; + const backend = prepared orelse { + stats.rebind_declined += 1; + return false; + }; + stats.rebind_attempts += 1; + entry.binding.?.apply(backend, request) catch |err| { + entry.poisoned = true; + stats.rebind_failures += 1; + return err; + }; + entry.key = key; + entry.used = clock; + stats.rebinds += 1; + stats.replays += 1; + note(key, .replay); + try entry.exec.?.launchOn(pass.stream); + return true; +} + +fn tryRebindTogether(c: *collective.Collective, signature: u64, key: Key, pass: anytype) !bool { + if (!stateRebindEnabled() or key.kind != 2 or !state_binding.supportsWindow(key.rows)) return false; + const request = stateRequest(key, pass); + const candidate = if (request) |target| rebindCandidate(key, target, true) else null; + var backend: ?state_driver.Backend = null; + var failed = false; + if (candidate) |entry| backend = prepareBinding(entry) catch blk: { + stats.rebind_failures += 1; + failed = true; + break :blk null; + }; + const selected = try c.graphAgree(signature ^ 0x737461746573656c, .{ + .state = if (failed) .failed else if (backend != null) .ready else .capture, + .cohort = if (candidate) |entry| entry.cohort else 0, + }); + if (selected.choice != .replay) { + stats.rebind_declined += 1; + return false; + } + const entry = candidate.?; + stats.rebind_attempts += 1; + entry.binding.?.applyTogether(backend.?, request.?, c, signature ^ 0x7374617465646f6e, selected.cohort) catch |err| { + entry.poisoned = true; + stats.rebind_failures += 1; + return err; + }; + entry.key = key; + entry.used = clock; + stats.rebinds += 1; + stats.replays += 1; + note(key, .replay); + try entry.exec.?.launchOn(pass.stream); + return true; +} + +test "exl3: state binding normalized identity preserves all nonstate key guards" { + const original: Key = .{ .kind = 2, .rows = 2, .scratch = 11, .scratch_cap = 16, .scores = 12, .output = 13, .keep = 14, .keep_nonstate = 15, .streams = 16, .plan = 17 }; + var flipped = original; + flipped.seq[0][4] = 100; + flipped.keep = 101; + try std.testing.expect(!std.meta.eql(original, flipped)); + try std.testing.expectEqualDeep(state_binding.normalized(original), state_binding.normalized(flipped)); + inline for (@typeInfo(Key).@"struct".field_names) |field| { + if (comptime !std.mem.eql(u8, field, "seq") and !std.mem.eql(u8, field, "paging") and + !std.mem.eql(u8, field, "paging_logical") and + !std.mem.eql(u8, field, "chunks") and !std.mem.eql(u8, field, "score_tiles") and !std.mem.eql(u8, field, "keep")) + { + var changed = original; + @field(changed, field) += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + } + for (0..max_lanes) |lane| { + for (0..5) |component| { + if (lane == 0 and component == 4) continue; + var changed = original; + changed.seq[lane][component] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + for (0..6) |component| { + var changed = original; + changed.paging[lane][component] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + var changed = original; + changed.chunks[lane] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + changed = original; + changed.score_tiles[lane] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + changed = original; + changed.paging_logical = @splat(@splat(0)); + changed.paging_logical.?[lane][0] = 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + changed.paging_logical.?[lane][1] = 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } +} + +test "exl3: state binding candidate excludes pending poisoned retired and incompatible generations" { + entries = @splat(null); + defer entries = @splat(null); + var request: state_binding.Request = .{ + .width = 2, + .chain = 1, + .chain_to = 2, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + for (2..state_binding.max_width + 1) |width| { + request.width = width; + request.nonstate.rows = @intCast(width); + const key: Key = .{ .kind = 2, .rows = width, .scratch = 1, .scores = 2, .output = 3, .keep = 4, .keep_nonstate = 5 }; + entries[0] = .{ + .key = key, + .used = 1, + .cohort = 7, + .exec = .{ .d = undefined, .handle = @ptrFromInt(1) }, + .template = .{ .d = undefined, .handle = @ptrFromInt(2) }, + .binding = try state_binding.Binding.init(2, request), + }; + var target = request; + std.mem.swap(u64, &target.layers[0].state_in, &target.layers[0].state_out); + var flipped = key; + flipped.seq[0][4] = 200; + flipped.keep = 9; + try std.testing.expect(find(flipped) == null); + try std.testing.expect(rebindCandidate(flipped, target, true) == &entries[0].?); + entries[0].?.pending = @ptrFromInt(@alignOf(Pending)); + try std.testing.expect(rebindCandidate(flipped, target, true) == null); + entries[0].?.pending = null; + entries[0].?.poisoned = true; + try std.testing.expect(rebindCandidate(flipped, target, true) == null); + entries[0].?.poisoned = false; + target.layers[0].state_in = 300; + try std.testing.expect(rebindCandidate(flipped, target, true) == null); + target = request; + flipped.output += 1; + try std.testing.expect(rebindCandidate(flipped, target, true) == null); + entries[0].?.binding.?.invalidate(); + try std.testing.expect(rebindCandidate(key, request, true) == null); + } +} + +test "exl3: state binding helper contracts" { + _ = @import("exl3_state_binding_test.zig"); + _ = state_driver; +} + +test "exl3: state binding transaction uses real graphAgree before publication" { + const Fake = struct { + calls: usize = 0, + fail: bool = false, + + pub fn set(self: *@This(), _: usize, _: usize, _: state_binding.Geometry, _: state_binding.Parameters) !void { + self.calls += 1; + if (self.fail and self.calls == 2) return error.SetterFailed; + } + }; + var request: state_binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 2, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + request.layers[1] = .{ .layer = 2, .state_in = 300, .state_out = 400 }; + var target = request; + for (target.layers[0..target.count]) |*mapping| std.mem.swap(u64, &mapping.state_in, &mapping.state_out); + for (2..state_binding.max_width + 1) |width| for (0..6) |scenario| { + request.width = width; + request.nonstate.rows = @intCast(width); + target.width = width; + target.nonstate.rows = @intCast(width); + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |descriptor| { + _ = std.c.close(descriptor); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |descriptor| { + const flags = std.c.fcntl(descriptor, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(descriptor, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, current: usize) void { + var control = GraphTestDriver.control(link); + const selected = control.graphAgree(51, .{ + .state = if (current == 2) .pending else if (current == 3) .capture else .ready, + .cohort = if (current == 1) 8 else 7, + }) catch @panic("binding selection failed"); + if (selected.choice != .replay) return; + const result = control.graphAgree(52, .{ .state = if (current == 5) .failed else .ready, .cohort = 7 }); + if (current == 4 or current == 5) { + if (result) |_| @panic("binding failure accepted") else |err| std.debug.assert(err == error.GraphInstantiationFailed); + } else { + std.debug.assert((result catch @panic("binding commit failed")).choice == .replay); + } + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, scenario }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + var state = try state_binding.Binding.init(42, request); + for (request.layers[0..request.count], 0..) |mapping, index| { + state.nodes[index] = index; + state.parameters[index] = .{ .nonstate = request.nonstate, .state_in = mapping.state_in, .state_out = mapping.state_out }; + } + state.loaded = true; + var fake: Fake = .{ .fail = scenario == 4 }; + const selected = try control.graphAgree(51, .{ .state = .ready, .cohort = 7 }); + var published = false; + if (selected.choice == .replay) { + if (scenario == 4 or scenario == 5) { + try std.testing.expectError(error.GraphInstantiationFailed, state.applyTogether(&fake, target, &control, 52, selected.cohort)); + try std.testing.expect(state.poisoned); + } else { + try state.applyTogether(&fake, target, &control, 52, selected.cohort); + published = true; + } + } + try std.testing.expectEqual(scenario == 0, published); + try std.testing.expectEqual(@as(usize, if (scenario == 0 or scenario >= 4) 2 else 0), fake.calls); + }; +} + +test "exl3: state binding current pass supports exact window layouts through the MTP row bound" { + try @import("cuda_exl3_state_request_test.zig").currentPass(Key, windowStateRequest); +} + +/// A captured graph being instantiated on a worker thread (20-35 ms that the step need not wait for: it runs +/// eagerly meanwhile, and replays from the next sighting after the instantiation landed). +pub const Pending = struct { + thread: std.Thread = undefined, + done: std.atomic.Value(bool) = .init(false), + graph: cuda.graph.Graph, + exec: ?cuda.graph.Exec = null, + ctx: ?*anyopaque, + keep_template: bool, + started_ns: u64 = 0, + + fn main(pd: *Pending) void { + const api = ctxApi() orelse { + pd.done.store(true, .release); + return; + }; + if (api.set(pd.ctx) != 0) { + pd.done.store(true, .release); + return; + } + pd.exec = pd.graph.instantiate() catch null; + pd.done.store(true, .release); + } +}; + +const CtxApi = struct { get: *const fn (*?*anyopaque) callconv(.c) c_int, set: *const fn (?*anyopaque) callconv(.c) c_int }; +var ctx_api: ?CtxApi = null; +var ctx_lib: ?std.DynLib = null; +fn ctxApi() ?CtxApi { + if (ctx_api) |a| return a; + var lib = std.DynLib.open("libcuda.so.1") catch return null; + const a: CtxApi = .{ + .get = lib.lookup(*const fn (*?*anyopaque) callconv(.c) c_int, "cuCtxGetCurrent") orelse return null, + .set = lib.lookup(*const fn (?*anyopaque) callconv(.c) c_int, "cuCtxSetCurrent") orelse return null, + }; + ctx_lib = lib; + ctx_api = a; + return a; +} + +/// TF_FLASHNEXT_EXL3_GRAPH_ASYNC=0 instantiates on the step's thread. +pub fn asyncInstantiate() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_ASYNC") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Starts instantiating `g` on a worker thread (it owns `g` from here); null when that is not possible. +pub fn startInstantiate(g: cuda.graph.Graph, keep_template: bool) !?*Pending { + if (!asyncInstantiate()) return null; + const api = ctxApi() orelse return null; + var ctx: ?*anyopaque = null; + if (api.get(&ctx) != 0) return null; + const pd = try std.heap.c_allocator.create(Pending); + pd.* = .{ .graph = g, .ctx = ctx, .keep_template = keep_template, .started_ns = wire.now() }; + pd.thread = std.Thread.spawn(.{}, Pending.main, .{pd}) catch |err| { + std.heap.c_allocator.destroy(pd); + return err; + }; + return pd; +} + +/// The finished instantiation's exec (the graph template freed), or null while it runs; error when it failed. +pub fn takePending(pd: *Pending) !?cuda.graph.Exec { + if (!pd.done.load(.acquire)) return null; + pd.thread.join(); + defer std.heap.c_allocator.destroy(pd); + pd.graph.deinit(); + return pd.exec orelse error.CudaFailed; +} + +/// A pending entry whose instantiation finished becomes a usable graph (or is dropped if it failed); true if usable. +fn settle(x: *Entry) bool { + if (x.poisoned) return false; + const pd = x.pending orelse return x.exec != null; + if (!pd.done.load(.acquire)) return false; + pd.thread.join(); + x.pending = null; + defer std.heap.c_allocator.destroy(pd); + if (pd.exec) |e| { + var exec = e; + if (x.stream) |stream| uploadResident(&exec, stream) catch { + exec.deinit(); + pd.graph.deinit(); + return false; + }; + x.exec = exec; + if (pd.keep_template) x.template = pd.graph else pd.graph.deinit(); + if (x.admission) |ticket| memory.installed(ticket) catch unreachable; + return true; + } + pd.graph.deinit(); + return false; +} + +fn drop(x: *Entry) void { + if (x.binding) |*binding| binding.invalidate(); + if (x.stream) |stream| stream.synchronize() catch @panic("graph retirement synchronization failed"); + defer if (x.admission) |ticket| { + memory.release(ticket) catch unreachable; + x.admission = null; + }; + if (x.pending) |pd| { + pd.thread.join(); + if (pd.exec) |*e| e.deinit(); + pd.graph.deinit(); + std.heap.c_allocator.destroy(pd); + return; + } + if (x.exec) |*e| e.deinit(); + if (x.template) |*g| g.deinit(); +} + +pub const Stats = struct { + captures: u64 = 0, + replays: u64 = 0, + rebind_attempts: u64 = 0, + rebinds: u64 = 0, + rebind_failures: u64 = 0, + rebind_declined: u64 = 0, + /// Steps on the device plan launched one by one (TF_FLASHNEXT_EXL3_DPLAN=1 without graphs). + eager_dplan: u64 = 0, + /// Steps that left for the host cut (too many rows, media, no pools). + declined: u64 = 0, + capture_ns: u64 = 0, + /// Captures that updated a cached graph of the same shape in place instead of instantiating one. + updates: u64 = 0, + /// Host time of `prepare` (ids, positions, n-gram rows) before a step's launch, and how many. + prepare_ns: u64 = 0, + prepares: u64 = 0, +}; + +pub var stats: Stats = .{}; + +/// Steps after which a graph nobody replayed counts as stale (another stream's, or a finished one's). +const stale_after: u64 = 512; + +/// Graphs kept (LRU): a round and up to `max_rows` window sizes per stream, at two seams. +const cache_n = 128; +var entries: [cache_n]?Entry = @splat(null); +var clock: u64 = 0; +/// A capture that failed once turns graphs off for the process (the steps keep the device plan). +var broken: bool = false; + +pub fn clear() void { + clearPrefill(); + for (&entries) |*e| if (e.*) |*x| { + drop(x); + e.* = null; + }; + seen = @splat(null); +} + +fn references(key: Key, tail: u64) bool { + if (tail == 0) return false; + for (key.seq) |seq| if (seq[2] == tail) return true; + return false; +} + +fn releaseSeqImpl(seq: anytype, keep_prefill: bool) void { + const tail = seq.tail.ptr; + if (tail == 0) return; + for (&entries) |*entry| if (entry.*) |*x| if (references(x.key, tail)) { + drop(x); + entry.* = null; + }; + for (&seen) |*entry| if (entry.*) |x| if (references(x.key, tail)) { + entry.* = null; + }; + for (&last_key) |*kinds| for (kinds) |*key| if (key.*) |k| if (references(k, tail)) { + key.* = null; + }; + if (!keep_prefill) { + for (&pentries) |*entry| if (entry.*) |*x| if (x.key.seq[2] == tail) { + dropPrefill(x); + entry.* = null; + }; + for (&pseen) |*entry| if (entry.*) |*x| if (x.key.seq[2] == tail) { + entry.* = null; + }; + } +} + +pub fn releaseSeq(seq: anytype) void { + releaseSeqImpl(seq, false); +} + +pub fn releaseDecodeSeq(seq: anytype) void { + releaseSeqImpl(seq, true); +} + +/// The engine's deinit: graphs and the device planner go. +pub fn releaseFor(owner: *const anyopaque) void { + clear(); + dplan.releaseFor(owner); +} + +fn find(key: Key) ?*Entry { + for (&entries) |*e| if (e.*) |*x| if (std.meta.eql(x.key, key)) return x; + return null; +} + +fn slot() *?Entry { + var best: usize = 0; + var best_used: u64 = std.math.maxInt(u64); + for (&entries, 0..) |*e, i| { + if (e.* == null) return e; + if (e.*.?.used < best_used) { + best_used = e.*.?.used; + best = i; + } + } + if (profileEnabled()) profile.evict(profile_module.Sample.fromKey(entries[best].?.key)); + drop(&entries[best].?); + entries[best] = null; + return &entries[best]; +} + +/// The key without the addresses a cudaGraphExecUpdate may change: equal shapes capture the same topology. +fn sameShape(a: Key, b: Key) bool { + return a.kind == b.kind and a.rows == b.rows and std.mem.eql(u32, &a.chunks, &b.chunks) and + std.mem.eql(u32, &a.score_tiles, &b.score_tiles) and a.scratch_cap == b.scratch_cap and + (a.keep != 0) == (b.keep != 0) and (a.streams != 0) == (b.streams != 0); +} + +/// Keys seen without a graph, and how often (a key is captured at its `after`-th sighting; one-off shapes, e.g. a +/// request's rare window sizes, stay eager). +const Seen = struct { key: Key, n: u32, used: u64 }; +var seen: [256]?Seen = @splat(null); + +fn sighting(key: Key) u32 { + var oldest: usize = 0; + var oldest_used: u64 = std.math.maxInt(u64); + for (&seen, 0..) |*e, i| { + if (e.*) |*x| { + if (std.meta.eql(x.key, key)) { + x.n += 1; + x.used = clock; + return x.n; + } + if (x.used < oldest_used) { + oldest_used = x.used; + oldest = i; + } + } else { + oldest = i; + oldest_used = 0; + } + } + seen[oldest] = .{ .key = key, .n = 1, .used = clock }; + return 1; +} + +/// TF_FLASHNEXT_EXL3_GRAPH_TP_UPDATE=0: tp graphs are always instantiated; default: updated in place like tp=1's, +/// their captured template kept alive with the exec (NCCL's captured resources belong to it). +fn tpUpdate() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_TP_UPDATE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// TF_FLASHNEXT_EXL3_GRAPH_AFTER: the sighting of a key that captures it (default 2; 1 captures at once). With a +/// blocking communicator a capture costs 3-7 ms and an in-place update 1-3 ms, and recycled stream buffers +/// (cuda_exl3_recycle.zig) keep keys stable across requests, so a key is captured once and replayed from then on. +fn after() u32 { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_AFTER") orelse return 2; + return @max(1, std.fmt.parseInt(u32, std.mem.span(v), 10) catch 2); +} + +/// Captures `pass.chunk(ids, pos0, null)` (prepared) under `key`, or replays the graph captured under it. A key not +/// yet seen `after()` times runs eagerly (prepared, uncaptured). A new key whose shape matches a cached graph's +/// updates that graph in place (cudaGraphExecUpdate) instead of instantiating one, except with a collective (NCCL's +/// captured resources belong to the graph they were captured in, so a tp graph is always instantiated). +/// Coverage by step kind (0 round, 1 window) and rows (1..16): replays, eager first sightings, captures that +/// instantiated, captures that updated in place. TF_FLASHNEXT_EXL3_GRAPH_LOG=N logs it every N steps (default 256). +pub const Outcome = enum { replay, eager, instantiated, updated }; +pub var coverage: [2][17][4]u32 = @splat(@splat(@splat(0))); +var steps_seen: u64 = 0; + +fn note(key: Key, o: Outcome) void { + coverage[key.kind - 1][@min(key.rows, 16)][@backingInt(o)] += 1; + steps_seen += 1; + const every: u64 = if (std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_LOG")) |v| std.fmt.parseInt(u64, std.mem.span(v), 10) catch 256 else 256; + if (every == 0 or steps_seen % every != 0) return; + if (profileEnabled()) logProfile(); + std.log.info("EXL3 state rebind attempts/succeeded/failed/declined {d}/{d}/{d}/{d}", .{ stats.rebind_attempts, stats.rebinds, stats.rebind_failures, stats.rebind_declined }); + var buf: [1024]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + for (0..2) |k| for (1..17) |r| { + const c = coverage[k][r]; + if (c[0] + c[1] + c[2] + c[3] == 0) continue; + w.print(" {s}{d}:{d}/{d}/{d}/{d}", .{ if (k == 0) "round" else "win", r, c[0], c[1], c[2], c[3] }) catch break; + }; + const rs = @import("cuda_exl3_recycle.zig").stats; + const cs = @import("cuda_exl3_chain.zig").stats; + std.log.info("flashnext exl3 graph coverage after {d} steps (replay/eager/instantiated/updated):{s}; recycled seqs {d}/{d} drafts {d}/{d} (reused/kept), shared lanes reassigned {d}; draft chains {d}: {d} replays, {d} captures ({d} updates), {d} eager; decode prepare {d} ms/{d}, capture {d} ms; prefill graphs {d}/{d}/{d}/{d}, capture {d} ms, paged replays {d}, origin rebinds {d}, sightings {d}/{d}, admission misses {d}", .{ steps_seen, w.buffered(), rs.seq_reused, rs.seq_kept, rs.draft_reused, rs.draft_kept, rs.shared_lane_reassigned, cs.chains, cs.replays, cs.captures, cs.updates, cs.eager, stats.prepare_ns / std.time.ns_per_ms, stats.prepares, stats.capture_ns / std.time.ns_per_ms, prefill_stats.replays, prefill_stats.eager, prefill_stats.captures, prefill_stats.updates, prefill_stats.capture_ns / std.time.ns_per_ms, prefill_stats.paged_replays, prefill_stats.rebinds, prefill_stats.sighting_hits, prefill_stats.sighting_misses, prefill_stats.admission_misses }); + std.log.info("EXL3 sequence retained-buffer evictions {d}", .{rs.seq_evicted}); +} + +fn logProfile() void { + std.log.info("EXL3 graph profile: misses mean absent keys, not missed replays; present-unready means pending or no exec at entry, before settlement; readiness is not replay proof; use runtime replay counters; fingerprint estimates, not allocation sizes; lane/key overflow {d}/{d}; changes overlap and compare within width except width transitions", .{ profile.lane_overflow, profile.key_overflow }); + for (profile.counters, 0..) |kind, kind_index| for (kind, 0..) |counts, width| { + if (counts.accesses == 0) continue; + std.log.info("EXL3 graph profile {s}{d}: accesses/key-presence-misses/present-unready/new/returning-key-misses/evicted/post-eviction-key-misses {d}/{d}/{d}/{d}/{d}/{d}/{d}; width/state/keep/seam/paging/allocation changes {d}/{d}/{d}/{d}/{d}/{d}; orientation overflow {d}", .{ if (kind_index == 0) "round" else "win", width, counts.accesses, counts.misses, counts.present_unready, counts.first_keys, counts.returning_misses, counts.pressure_evictions, counts.post_eviction_misses, counts.width_changes, counts.state_changes, counts.keep_changes, counts.seam_changes, counts.paging_changes, counts.allocation_changes, counts.orientation_overflow }); + }; + for (profile.lanes, 0..) |maybe, index| if (maybe) |lane| { + std.log.info("EXL3 graph profile lane {d}: width mask {x}; accesses/key-presence-misses/present-unready/new/returning-key-misses/evicted {d}/{d}/{d}/{d}/{d}/{d}; orientation A/B/other accesses {d}/{d}/{d}", .{ index, lane.width_mask, lane.counters.accesses, lane.counters.misses, lane.counters.present_unready, lane.counters.first_keys, lane.counters.returning_misses, lane.counters.pressure_evictions, lane.orientation_accesses[0], lane.orientation_accesses[1], lane.orientation_accesses[2] }); + std.log.info("EXL3 graph profile lane {d}: new fingerprints by width 0..16 {any}", .{ index, lane.new_keys_by_width }); + }; +} + +/// TF_FLASHNEXT_EXL3_GRAPH_DEBUG=1: which key fields changed since the last step of the same kind and rows. +var last_key: [2][17]?Key = @splat(@splat(null)); +fn debugKey(key: Key) void { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_DEBUG") orelse return; + if (!std.mem.eql(u8, std.mem.span(v), "1")) return; + const lk = &last_key[key.kind - 1][@min(key.rows, 16)]; + defer lk.* = key; + const prev = lk.* orelse return; + var buf: [256]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + inline for (@typeInfo(Key).@"struct".field_names) |f| { + if (!std.meta.eql(@field(prev, f), @field(key, f))) w.print(" {s}", .{f}) catch {}; + } + std.log.info("flashnext exl3 graph key ({s} of {d} rows) changed:{s} keep {x} seq {x} {x} {x} {d} {x} out {x} streams {x} scratch {x} scores {x}", .{ if (key.kind == 1) "round" else "window", key.rows, w.buffered(), key.keep, key.seq[0][0], key.seq[0][1], key.seq[0][2], key.seq[0][3], key.seq[0][4], key.output, key.streams, key.scratch, key.scores }); +} + +fn run(key: Key, pass: anytype, ids: []const u32, pos0: usize) !void { + if (profileEnabled()) { + const presence: profile_module.Presence = if (find(key)) |entry| + if (entry.pending != null or entry.exec == null) .present_unready else .present_ready + else + .absent; + profile.observeEntry(profile_module.Sample.fromKey(key), presence); + } + debugKey(key); + clock += 1; + if (pass.m.collective) |c| return runTogether(c, key, pass, ids, pos0); + if (find(key)) |e| { + e.used = clock; + if (!settle(e)) { + if (e.pending == null) { + // the instantiation failed: forget the key, run eagerly + for (&entries) |*en| if (en.*) |*x| if (x == e) { + drop(x); + en.* = null; + }; + return error.CudaFailed; + } + stats.eager_dplan += 1; + note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + stats.replays += 1; + note(key, .replay); + return e.exec.?.launchOn(pass.stream); + } + if (try tryRebind(key, pass)) return; + if (sighting(key) < after()) { + stats.eager_dplan += 1; + note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + const t0 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + try pass.stream.synchronize(); + const admission = try reserveMemory(pass.stream, .decode, 0); + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + try cuda.graph.beginCapture(pass.stream, .thread_local); + pass.prepared = true; + capturing = true; + defer capturing = false; + pass.chunk(ids, pos0, null) catch |err| { + if (cuda.graph.endCapture(pass.stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return err; + }; + var graph = try cuda.graph.endCapture(pass.stream); + var kept = false; + defer if (!kept) graph.deinit(); + const t1 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + var how: []const u8 = "instantiated"; + const tp = pass.m.collective != null; + const at: *Entry = blk: { + // a same-shape graph is updated in place only when it is stale (its key not used for a while, e.g. a finished + // stream's) or the cache is full: two live variants of one shape (a window's state ping-pong) keep two graphs + var free_slot = false; + for (&entries) |*e| free_slot = free_slot or e.* == null; + if ((!tp or tpUpdate()) and capturedBinding(key, pass, graph) == null) { + for (&entries) |*e| if (e.*) |*x| if (!x.poisoned and x.exec != null and x.pending == null and sameShape(x.key, key) and (!free_slot or clock - x.used > stale_after)) { + // every earlier launch finished (each step ends on the host), so the old parameters are free to go + const r = try x.exec.?.update(graph); + if (r == .success) { + if (x.binding) |*binding| binding.invalidate(); + x.binding = null; + try uploadResident(&x.exec.?, pass.stream); + x.key = key; + x.used = clock; + if (x.template) |*old| old.deinit(); + x.template = null; + if (tp) { + x.template = graph; + kept = true; + } + stats.updates += 1; + note(key, .updated); + how = "updated"; + break :blk x; + } + }; + } + if (asyncInstantiate()) if (ctxApi()) |api| { + var ctx: ?*anyopaque = null; + if (api.get(&ctx) == 0) { + const sl = slot(); + const pd = try std.heap.c_allocator.create(Pending); + const binding = capturedBinding(key, pass, graph); + pd.* = .{ .graph = graph, .ctx = ctx, .keep_template = tp or binding != null, .started_ns = wire.now() }; + kept = true; + pd.thread = std.Thread.spawn(.{}, Pending.main, .{pd}) catch |err| { + kept = false; + std.heap.c_allocator.destroy(pd); + return err; + }; + sl.* = .{ .key = key, .used = clock, .pending = pd, .admission = admission, .stream = pass.stream, .binding = binding }; + admitted = false; + note(key, .instantiated); + stats.captures += 1; + // this step runs eagerly while the graph instantiates + capturing = false; + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + }; + const sl = slot(); + var exec = try graph.instantiate(); + errdefer exec.deinit(); + try uploadResident(&exec, pass.stream); + try memory.installed(admission); + sl.* = .{ .key = key, .exec = exec, .used = clock, .admission = admission, .stream = pass.stream }; + sl.*.?.binding = capturedBinding(key, pass, graph); + if (sl.*.?.binding != null) { + sl.*.?.template = graph; + kept = true; + } + admitted = false; + note(key, .instantiated); + break :blk &sl.*.?; + }; + stats.captures += 1; + const t2 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + stats.capture_ns += @intCast(t2 - t0); + std.log.info("flashnext exl3 graph: captured a {s} of {d} rows in {d:.1} ms + {s} in {d:.1} ms ({d} captures, {d} updates, {d} replays, {d} eager, {d} declined steps, {d} residency pulls so far)", .{ if (key.kind == 1) "decode round" else "window", key.rows, @as(f64, @floatFromInt(t1 - t0)) / 1e6, how, @as(f64, @floatFromInt(t2 - t1)) / 1e6, stats.captures, stats.updates, stats.replays, stats.eager_dplan, stats.declined, if (dplan.current()) |dp| dp.pulls else 0 }); + try at.exec.?.launchOn(pass.stream); +} + +fn graphSignature(key: Key, pos0: usize) u64 { + var hash = std.hash.Wyhash.init(0); + std.hash.autoHash(&hash, .{ key.kind, key.rows, key.chunks, key.score_tiles, key.scratch_cap, key.keep != 0, key.streams != 0, pos0 }); + std.hash.autoHash(&hash, stateRebindEnabled()); + for (key.seq, key.paging, 0..) |seq, paging, index| { + const identity: [2]u64 = if (key.paging_logical) |logical| logical[index] else .{ paging[4], paging[5] }; + std.hash.autoHash(&hash, .{ seq[3], identity[0], identity[1] }); + } + return hash.final(); +} + +fn graphVote(entry: ?*Entry, timeout_ms: u64) collective.GraphVote { + const x = entry orelse return .{ .state = .capture }; + if (settle(x)) return .{ .state = .ready, .cohort = x.cohort }; + if (x.pending) |pd| { + const elapsed = wire.now() -| pd.started_ns; + return .{ .state = if (elapsed >= timeout_ms *| std.time.ns_per_ms) .timed_out else .pending, .cohort = x.cohort }; + } + return .{ .state = .failed, .cohort = x.cohort }; +} + +fn runTogether(c: *collective.Collective, key: Key, pass: anytype, ids: []const u32, pos0: usize) !void { + errdefer c.failGraph(); + const signature = graphSignature(key, pos0); + var entry = find(key); + if (entry) |x| x.used = clock; + const vote = if (entry != null) graphVote(entry, c.graphTimeoutMs()) else collective.GraphVote{ + .state = if (stateRequest(key, pass)) |request| + if (rebindCandidate(key, request, true) != null) .capture else if (sighting(key) < after()) .eager else .capture + else if (sighting(key) < after()) .eager else .capture, + }; + const decision = try c.graphAgree(signature, vote); + switch (decision.choice) { + .eager => { + stats.eager_dplan += 1; + note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + }, + .replay => { + stats.replays += 1; + note(key, .replay); + return entry.?.exec.?.launchOn(pass.stream); + }, + .capture => {}, + } + if (try tryRebindTogether(c, signature, key, pass)) return; + const started = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + try pass.stream.synchronize(); + const admission = try reserveTogether(c, pass.stream, signature, decision.cohort); + entry = find(key); + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + try cuda.graph.beginCapture(pass.stream, .thread_local); + pass.prepared = true; + capturing = true; + pass.chunk(ids, pos0, null) catch |err| { + capturing = false; + if (cuda.graph.endCapture(pass.stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return err; + }; + capturing = false; + var graph = try cuda.graph.endCapture(pass.stream); + var kept = false; + defer if (!kept) graph.deinit(); + var updated = false; + var candidate: ?*Entry = null; + if (tpUpdate() and capturedBinding(key, pass, graph) == null) for (&entries) |*maybe| if (maybe.*) |*current| { + if (!current.poisoned and current.cohort != 0 and current.exec != null and current.pending == null and sameShape(current.key, key) and + (std.meta.eql(current.key, key) or clock - current.used > stale_after)) + { + if (candidate == null or current.used < candidate.?.used) candidate = current; + } + }; + const update_agreement = try c.graphAgree(signature ^ 0x757064617465, if (candidate) |current| + .{ .state = .ready, .cohort = current.cohort } + else + .{ .state = .capture }); + if (update_agreement.choice == .replay) { + const current = candidate.?; + const result = try current.exec.?.update(graph); + if (result == .success) { + if (current.binding) |*binding| binding.invalidate(); + current.binding = null; + if (current.template) |*old| old.deinit(); + current.template = graph; + kept = true; + updated = true; + } + const result_agreement = try c.graphAgree(signature ^ 0x75706461746564, .{ .state = if (updated) .ready else .capture, .cohort = decision.cohort }); + if (result_agreement.choice == .replay) { + if (current.admission) |old| try memory.release(old); + current.admission = admission; + admitted = false; + try uploadResident(¤t.exec.?, pass.stream); + try memory.installed(admission); + current.key = key; + current.used = clock; + current.cohort = decision.cohort; + current.stream = pass.stream; + entry = current; + } else { + if (updated) { + current.template = null; + kept = false; + } + drop(current); + for (&entries) |*slot_entry| if (slot_entry.*) |*cached| if (cached == current) { + slot_entry.* = null; + }; + entry = find(key); + updated = false; + try retirementBarrier(c, signature ^ 0x72656275696c64, decision.cohort); + } + } + if (!updated) { + const target = if (entry) |old| blk: { + drop(old); + for (&entries) |*slot_entry| if (slot_entry.*) |*x| if (x == old) { + slot_entry.* = null; + break :blk slot_entry; + }; + unreachable; + } else blk: { + for (&entries) |*slot_entry| if (slot_entry.* == null) break :blk slot_entry; + return error.GraphBudgetBusy; + }; + if (try startInstantiate(graph, true)) |pending| { + target.* = .{ .key = key, .used = clock, .pending = pending, .cohort = decision.cohort, .admission = admission, .stream = pass.stream, .binding = capturedBinding(key, pass, graph) }; + } else { + var exec = try graph.instantiate(); + errdefer exec.deinit(); + try uploadResident(&exec, pass.stream); + try memory.installed(admission); + target.* = .{ .key = key, .exec = exec, .used = clock, .template = graph, .cohort = decision.cohort, .admission = admission, .stream = pass.stream, .binding = capturedBinding(key, pass, graph) }; + } + admitted = false; + kept = true; + entry = &target.*.?; + } + stats.captures += 1; + if (updated) stats.updates += 1; + note(key, if (updated) .updated else .instantiated); + stats.capture_ns += @intCast(std.Io.Clock.awake.now(pass.io).toNanoseconds() - started); + const installed = try c.graphAgree(signature ^ 0x696e7374616c6c, graphVote(entry, c.graphTimeoutMs())); + switch (installed.choice) { + .replay => try entry.?.exec.?.launchOn(pass.stream), + .eager => { + stats.eager_dplan += 1; + pass.prepared = true; + try pass.chunk(ids, pos0, null); + }, + .capture => return error.GraphRendezvousMismatch, + } +} + +/// A round's host inputs (what chunk's stage and gather would write), outside capture and replay. +fn prepareRound(pass: anytype, ids: []const u32) !void { + const s = pass.s; + try s.ready(); + for (s.hostIds(ids.len), ids) |*out, token| out.* = @intCast(token); + const positions = s.hostPos(ids.len * 4); + @memset(positions, 0); + for (pass.lanes, 0..) |lane, i| { + positions[i * 4] = @intCast(lane.pos); + try pass.m.ngram.rows(lane.seq.history[0..], ids[i..][0..1], s.hostEmb(ids.len)[i * dims ..][0..dims], s.hostXs(ids.len)[i * (dims / 32) ..][0 .. dims / 32]); + } +} + +/// A window's host inputs: ids, one position a block of attention rows, the n-gram rows after the sequence's history. +fn prepareWindow(pass: anytype, ids: []const u32, pos0: usize) !void { + const s = pass.s; + const n = ids.len; + try s.ready(); + for (s.hostIds(n), ids) |*out, token| out.* = @intCast(token); + const blocks = (n + att_rows - 1) / att_rows; + const positions = s.hostPos(blocks * 4); + @memset(positions, 0); + for (0..blocks) |b| positions[4 * b] = @intCast(pos0 + b * att_rows); + try pass.m.ngram.rows(pass.seq.history[0..], ids, s.hostEmb(n), s.hostXs(n)); +} + +/// The device planner for `e`, or null when its decode steps keep the host cut. +fn planner(e: anytype, rows: usize) !?*dplan.DevPlan { + if (!dplan.enabled()) return null; + const dp = (try dplan.forLayers(e.gpa, e.driver, e, e.layers[0..e.n_layers])) orelse return null; + if (rows > dp.max_rows) { + stats.declined += 1; + return null; + } + return dp; +} + +/// Engine.round's hook: true when this module ran the round (device plan, captured or not); false leaves it to the +/// engine. `p` is the round's pass (lanes set), `out` its logits [rows, vocab]. +pub fn round(e: anytype, p: anytype, seqs: anytype, ids: []const u32, scores: u64, out: u64) !bool { + for (seqs) |q| if (q.media != null) return false; + if (seqs.len > max_lanes) return false; + for (seqs) |q| if (q.paged) |pg| try pg.validateSharedLeaseOn(p.stream); + const dp = (try planner(e, ids.len)) orelse return false; + try dp.toDevice(p.stream); + dplan.active = dp; + defer dplan.active = null; + if (enabled() and broken) return error.GraphCaptureFailed; + if (!enabled()) { + stats.eager_dplan += 1; + try p.chunk(ids, 0, null); + return true; + } + for (seqs) |q| if (q.paged) |pg| for (0..wt.attn_n) |ai| try pg.validateRange(ai, q.fed, 1); + defer for (seqs) |q| if (q.paged) |pg| { + pg.graph_prepared = false; + }; + for (seqs) |q| if (q.paged) |pg| try pg.prepareGraph(p.stream, q.fed, 1); + var key: Key = .{ .kind = 1, .rows = ids.len, .scratch = p.s.ids, .scratch_cap = p.s.cap, .scores = scores, .output = out, .plan = dp.state.ptr }; + for (seqs, 0..) |q, i| { + key.seq[i] = Key.seqOf(q); + if (q.paged) |pg| { + key.paging[i] = pg.graphKey(); + if (key.paging_logical == null) key.paging_logical = @splat(@splat(0)); + key.paging_logical.?[i] = pg.logicalGraphKey(); + } + Key.seams(q.fed + 1, &key.chunks[i], &key.score_tiles[i]); + } + const tp0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try prepareRound(p, ids); + stats.prepare_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - tp0); + stats.prepares += 1; + run(key, p, ids, 0) catch |err| { + if (p.m.collective) |c| { + c.failGraph(); + return err; + } + broken = true; + return err; + }; + return true; +} + +/// Engine.forward's hook for a window: every row's logits into hooks.logits_out (an MTP verify window, or a one-row +/// step), at most the planner's rows. True when this module ran it (seq.fed and seq.logits advanced as forward does). +pub fn forward(e: anytype, p: anytype, seq: anytype, ids: []const u32) !bool { + const hooks = p.hooks; + if (hooks.logits != .all or hooks.logits_out == 0 or seq.media != null or ids.len == 0 or ids.len > att_rows) return false; + if (seq.fed + ids.len > seq.cap or ids.len > p.s.cap) return false; + if (seq.paged) |pg| try pg.validateSharedLeaseOn(p.stream); + const dp = (try planner(e, ids.len)) orelse return false; + try dp.toDevice(p.stream); + dplan.active = dp; + defer dplan.active = null; + const n = ids.len; + if (enabled() and broken) return error.GraphCaptureFailed; + if (!enabled()) { + stats.eager_dplan += 1; + try p.chunk(ids, seq.fed, null); + } else { + defer if (seq.paged) |pg| { + pg.graph_prepared = false; + }; + if (seq.paged) |pg| try pg.prepareGraph(p.stream, seq.fed, n); + var key: Key = .{ .kind = 2, .rows = n, .scratch = p.s.ids, .scratch_cap = p.s.cap, .scores = p.seq.scores, .output = hooks.logits_out, .streams = hooks.streams, .plan = dp.state.ptr }; + key.seq[0] = Key.seqOf(seq); + if (seq.paged) |pg| { + key.paging[0] = pg.graphKey(); + key.paging_logical = @splat(@splat(0)); + key.paging_logical.?[0] = pg.logicalGraphKey(); + } + Key.seams(seq.fed + n, &key.chunks[0], &key.score_tiles[0]); + if (hooks.keep) |kp| { + key.keep = std.hash.Wyhash.hash(1, std.mem.asBytes(kp)); + if (stateRebindEnabled()) key.keep_nonstate = state_binding.keepHash(kp.*); + } + const tp0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try prepareWindow(p, ids, seq.fed); + stats.prepare_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - tp0); + stats.prepares += 1; + run(key, p, ids, seq.fed) catch |err| { + if (p.m.collective) |c| { + c.failGraph(); + return err; + } + broken = true; + return err; + }; + } + if (hooks.output_tail) |tail| try tail.enqueue(tail.ptr, p.stream); + // the pass's head(.all) copies the last row's scores to the host after the step (outside the graph) + if (p.m.collective) |c| try c.finish(p.stream) else try p.stream.synchronize(); + const d = p.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(seq.logits.ptr, hooks.logits_out + (n - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + seq.fed += n; + return true; +} + +// ---- prefill: one layer's sub-blocks of a deep chunk ------------------------------------------------------------- +// +// The expert-major prefill (cuda_exl3_prefill.zig) runs each layer as its dense sub-blocks (HC read-outs, DeltaNet or +// attention, write-backs, router and top-k over every sub-block of the deep chunk) and then the layer's experts, whose +// passes the host plans from the layer's picks (the ring, the members, the copies): that half stays eager. The dense +// half has no host dependency and is captured per layer: a DeltaNet layer's graph depends only on the chunk's shape +// (rows, sub-block size) and replays for every later chunk of that shape; an attention layer's launch geometry follows +// the chunk's position, so each chunk re-captures it and updates the cached graph in place when the topology holds. + +pub const Prefill = enum { eager, replayed, capturing }; +const prefill_binding = @import("cuda_exl3_prefill_binding.zig"); + +const PBuffer = struct { + driver: *const cuda.Driver, + address: u64, + bytes: usize, + + fn of(buffer: cuda.DeviceBuffer) PBuffer { + return .{ .driver = buffer.d, .address = buffer.ptr, .bytes = buffer.len }; + } +}; + +const PGdn = struct { + conv: PBuffer, + state: PBuffer, + + fn of(sequence: anytype, layer: usize) PGdn { + const index = wt.linearIndex(layer); + return .{ .conv = PBuffer.of(sequence.conv[index]), .state = PBuffer.of(sequence.state[index]) }; + } +}; + +const PKey = struct { + layer: usize, + rows: usize, + sub: usize, + /// The chunk's first position (attention layers only; 0 for DeltaNet layers). + pos0: usize, + deep: u64, + base: u64, + base_cap: usize, + seq: [5]u64, + paging: [6]u64, + wait_count: usize, + wait_key: u64, + scores: u64, + tiles: bool, + positions: u64 = 0, + stream: u64 = 0, + gdn: ?PGdn = null, +}; + +const PEntry = struct { key: PKey, exec: cuda.graph.Exec, used: u64, admission: ?memory.Ticket = null, stream: ?cuda.Stream = null, binding: ?*prefill_binding.Binding = null, binding_admission: ?memory.Ticket = null }; + +fn setPrefillBinding(entry: *PEntry, binding: ?*prefill_binding.Binding, admission: ?memory.Ticket) void { + if (entry.binding) |old| old.deinit(std.heap.c_allocator); + if (entry.binding_admission) |ticket| memory.release(ticket) catch unreachable; + entry.binding = binding; + entry.binding_admission = admission; +} + +fn dropPrefill(entry: *PEntry) void { + if (entry.stream) |stream| stream.synchronize() catch @panic("prefill retirement synchronization failed"); + entry.exec.deinit(); + setPrefillBinding(entry, null, null); + if (entry.admission) |ticket| memory.release(ticket) catch unreachable; +} +const pcache_n = 112; +var pentries: [pcache_n]?PEntry = @splat(null); +var pseen: [pcache_n * 2]?struct { key: PKey, n: u32 } = @splat(null); +var pseen_at: usize = 0; +var pcur: ?PKey = null; +var pcur_admission: ?memory.Ticket = null; +var pcur_binding_admission: ?memory.Ticket = null; +var pcur_binding_request: ?prefill_binding.Request = null; + +pub var prefill_stats: struct { captures: u64 = 0, updates: u64 = 0, replays: u64 = 0, paged_replays: u64 = 0, rebinds: u64 = 0, sighting_hits: u64 = 0, sighting_misses: u64 = 0, eager: u64 = 0, admission_misses: u64 = 0, capture_ns: u64 = 0 } = .{}; + +fn prefillAdmissionMiss() Prefill { + prefill_stats.eager += 1; + prefill_stats.admission_misses += 1; + const misses = prefill_stats.admission_misses; + if (misses == 1 or misses % 256 == 0) std.log.info("flashnext exl3 prefill graph budget admission declines: {d}; retained cache preserved, layer runs eagerly", .{misses}); + return .eager; +} + +/// TF_FLASHNEXT_EXL3_GRAPH_PREFILL=1 captures prefill layers under TF_FLASHNEXT_EXL3_GRAPH=1 (default off pending +/// measurements). TP=2 graphs capture each incoming sub-block reduction as an external event wait. +fn prefillEnabled() bool { + if (!enabled() or broken) return false; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_PREFILL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +fn pshape(k: PKey) PKey { + var s = k; + s.pos0 = 0; + s.paging[5] = 0; + return s; +} + +fn prefillOriginCandidate(key: PKey, request: prefill_binding.Request) ?*PEntry { + for (&pentries) |*maybe| if (maybe.*) |*entry| { + const binding = entry.binding orelse continue; + if (!binding.poisoned and std.meta.eql(pshape(entry.key), pshape(key)) and binding.request.compatible(request)) return entry; + }; + return null; +} + +test "exl3: GDN prefill keys bind the actual layer buffers including driver and extent" { + var driver: cuda.Driver = undefined; + var other_driver: cuda.Driver = undefined; + var conv = [_]cuda.DeviceBuffer{ + .{ .d = &driver, .ptr = 100, .len = 16 }, + .{ .d = &driver, .ptr = 200, .len = 16 }, + .{ .d = &driver, .ptr = 300, .len = 16 }, + }; + var state = conv; + for (&state) |*buffer| buffer.ptr += 1000; + const sequence = .{ .conv = &conv, .state = &state }; + var key = std.mem.zeroes(PKey); + key.layer = 2; + key.seq = .{ conv[0].ptr, 0, 500, 4096, state[0].ptr }; + key.gdn = PGdn.of(sequence, key.layer); + const original = key.gdn.?; + try std.testing.expectEqual(@as(u64, 300), original.conv.address); + try std.testing.expectEqual(@as(u64, 1300), original.state.address); + for ([_]bool{ false, true }) |recurrent| { + const buffer = if (recurrent) &state[2] else &conv[2]; + const saved = buffer.*; + defer buffer.* = saved; + for (0..3) |field| { + buffer.* = saved; + switch (field) { + 0 => buffer.ptr += 1, + 1 => buffer.len += 1, + 2 => buffer.d = &other_driver, + else => unreachable, + } + var changed = key; + changed.gdn = PGdn.of(sequence, key.layer); + try std.testing.expectEqual(key.seq, changed.seq); + try std.testing.expect(!std.meta.eql(pshape(key), pshape(changed))); + } + } + try std.testing.expectEqualDeep(original, PGdn.of(sequence, key.layer)); +} + +test "exl3: prefill origin candidates preserve paging generation, incarnation and allocation keys" { + const key: PKey = .{ .layer = 7, .rows = 512, .sub = 256, .pos0 = 4096, .deep = 1, .base = 2, .base_cap = 4096, .seq = .{ 3, 4, 5, 4096, 6 }, .paging = .{ 7, 8, 9, 10, 11, 12 }, .wait_count = 0, .wait_key = 0, .scores = 13, .tiles = true, .positions = 14, .stream = 15 }; + const request: prefill_binding.Request = .{ .functions = .{ .prep = 1, .scores = 2, .chunks = 3, .ahead = 4, .plan = 5, .remap = 6 }, .rows = 512, .sub = 256, .pos0 = 4096, .positions = 14, .cap = 32768, .origin = 1024, .page_rows = 256, .shared = true, .generation = 1, .keys = 16, .values = 17 }; + var binding: prefill_binding.Binding = .{ .graph = undefined, .request = request, .records = &.{} }; + var driver: cuda.Driver = undefined; + pentries = @splat(null); + defer pentries = @splat(null); + pentries[0] = .{ .key = key, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .binding = &binding }; + var moved = key; + moved.pos0 = 8192; + moved.paging[5] += 1; + var target = request; + target.pos0 = 8192; + target.origin = 4096; + try std.testing.expect(prefillOriginCandidate(moved, target) != null); + for (0..5) |index| { + var changed = moved; + changed.paging[index] += 1; + try std.testing.expect(prefillOriginCandidate(changed, target) == null); + changed = moved; + changed.seq[index] += 1; + try std.testing.expect(prefillOriginCandidate(changed, target) == null); + } + inline for (.{ "deep", "base", "base_cap", "scores", "positions", "stream", "wait_key", "wait_count", "rows", "sub", "layer" }) |field| { + var changed = moved; + @field(changed, field) += 1; + try std.testing.expect(prefillOriginCandidate(changed, target) == null); + } + binding.poisoned = true; + try std.testing.expect(prefillOriginCandidate(moved, target) == null); +} + +test "exl3: prefill retained templates release both graph budget tickets on retirement" { + clearPrefill(); + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 0, .charge = .{ .exec = 10, .template = 10, .transient = 10 }, .prefill_charge = .{ .exec = 5, .template = 0, .transient = 10 } }); + const exec_ticket = try memory.admit(.prefill, 0, 100); + try memory.installed(exec_ticket); + const template_ticket = try memory.admit(.prefill, 0, 100); + try memory.installed(template_ticket); + try std.testing.expectEqual(@as(u64, 10), memory.charged()); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + const binding = try std.heap.c_allocator.create(prefill_binding.Binding); + binding.* = .{ .graph = .{ .d = &driver, .handle = @ptrFromInt(1) }, .request = undefined, .records = try std.heap.c_allocator.alloc(prefill_binding.Record, 0) }; + pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .admission = exec_ticket, .binding = binding, .binding_admission = template_ticket, .stream = .{ .d = &driver, .handle = @ptrFromInt(1) } }; + clearPrefill(); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.graph_drops); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); + try std.testing.expectError(error.StaleGraphTicket, memory.release(exec_ticket)); + try std.testing.expectError(error.StaleGraphTicket, memory.release(template_ticket)); +} + +test "prefill graph update candidates tolerate dynamic position and page-window origins only" { + const base: PKey = .{ .layer = 7, .rows = 4096, .sub = 256, .pos0 = 4096, .deep = 1, .base = 2, .base_cap = 4096, .seq = .{ 3, 4, 5, 4096, 6 }, .paging = .{ 7, 8, 9, 10, 11, 12 }, .wait_count = 0, .wait_key = 0, .scores = 13, .tiles = true }; + var moved = base; + moved.pos0 = 8192; + moved.paging[5] = 14; + try std.testing.expect(!std.meta.eql(base, moved)); + try std.testing.expect(std.meta.eql(pshape(base), pshape(moved))); + + moved.seq[0] += 1; + try std.testing.expect(!std.meta.eql(pshape(base), pshape(moved))); + moved = base; + moved.paging[0] += 1; + try std.testing.expect(!std.meta.eql(pshape(base), pshape(moved))); + moved = base; + moved.rows += 1; + try std.testing.expect(!std.meta.eql(pshape(base), pshape(moved))); + moved = base; + moved.wait_key += 1; + try std.testing.expect(!std.meta.eql(pshape(base), pshape(moved))); +} + +fn psighting(k: PKey) u32 { + const s = pshape(k); + for (&pseen) |*e| if (e.*) |*x| if (std.meta.eql(x.key, s)) { + x.n += 1; + prefill_stats.sighting_hits += 1; + return x.n; + }; + prefill_stats.sighting_misses += 1; + pseen[pseen_at % pseen.len] = .{ .key = s, .n = 1 }; + pseen_at += 1; + return 1; +} + +test "exl3: prefill warmup history retains two lanes and two chunk banks" { + pseen = @splat(null); + pseen_at = 0; + defer { + pseen = @splat(null); + pseen_at = 0; + } + var keys: [2 * 2 * wt.layers_n]PKey = undefined; + var count: usize = 0; + for (0..2) |sequence| for (0..2) |bank| for (0..wt.layers_n) |layer| { + keys[count] = .{ + .layer = layer, + .rows = 4096, + .sub = 256, + .pos0 = bank * 4096, + .deep = 100 + bank, + .base = 200, + .base_cap = 8192, + .seq = .{ 1, 2, 300 + sequence, 8192, 4 }, + .paging = .{ 5, 6, 7, 8, 9, sequence + bank + 1 }, + .wait_count = 0, + .wait_key = 0, + .scores = 10, + .tiles = true, + .positions = 11, + .stream = 12, + }; + count += 1; + }; + try std.testing.expectEqual(@as(usize, 192), count); + for (keys) |key| try std.testing.expectEqual(@as(u32, 1), psighting(key)); + for (keys) |key| try std.testing.expectEqual(@as(u32, 2), psighting(key)); +} + +/// TF_FLASHNEXT_EXL3_GRAPH_PREFILL_AFTER: the sighting of a layer shape that captures it (default 2: a chunk takes +/// seconds, a capture milliseconds). +fn prefillAfter() u32 { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_PREFILL_AFTER") orelse return 2; + return @max(1, std.fmt.parseInt(u32, std.mem.span(v), 10) catch 2); +} + +pub fn clearPrefill() void { + for (&pentries) |*e| if (e.*) |*x| { + dropPrefill(x); + e.* = null; + }; + pseen = @splat(null); +} + +/// Before a layer's sub-blocks: `.replayed` (the graph ran; skip the sub-blocks), `.capturing` (run them: they are +/// recorded; then `prefillEnd`), or `.eager` (run them; `prefillEnd` is a no-op). +/// External event waits for expert-parallel sub-blocks are captured in order with the dense work. +pub fn prefillBegin(p: anytype, linear: bool, layer: usize, rows: usize, sub: usize, pos0: usize, deep: u64, base: anytype, tiles: bool, waits: []const cuda.Event, positions: u64) !Prefill { + if (enabled() and broken) return error.GraphCaptureFailed; + if (p.seq.paged) |pg| try pg.validateSharedLeaseOn(p.stream); + if (!prefillEnabled()) return .eager; + const q = p.seq; + if (!linear) if (q.paged) |pg| { + try pg.prepareLayer(p.stream, wt.attnIndex(layer), pos0, rows); + pg.graph_prepared = true; + }; + errdefer if (q.paged) |pg| { + pg.graph_prepared = false; + }; + var wait_key: u64 = 0; + for (waits) |ev| { + const handle = if (ev.handle) |h| @intFromPtr(h) else 0; + wait_key = std.hash.Wyhash.hash(wait_key, std.mem.asBytes(&handle)); + } + const key: PKey = .{ + .layer = layer, + .rows = rows, + .sub = sub, + .pos0 = if (linear) 0 else pos0, + .deep = deep, + .base = base.ids, + .base_cap = base.cap, + .seq = .{ if (q.conv.len > 0) q.conv[0].ptr else 0, if (q.kc.len > 0) q.kc[0].ptr else 0, q.tail.ptr, q.cap, if (q.state.len > 0) q.state[0].ptr else 0 }, + .paging = if (!linear) (if (q.paged) |pg| pg.prefillGraphKey() else @splat(0)) else @splat(0), + .wait_count = waits.len, + .wait_key = wait_key, + .scores = q.scores, + .tiles = tiles, + .positions = positions, + .stream = @intFromPtr(p.stream.handle), + .gdn = if (linear) PGdn.of(q, layer) else null, + }; + clock += 1; + for (&pentries) |*e| if (e.*) |*x| if (std.meta.eql(x.key, key)) { + if (x.binding) |binding| if (binding.poisoned) continue; + x.used = clock; + prefill_stats.replays += 1; + if (x.binding != null) prefill_stats.paged_replays += 1; + try x.exec.launchOn(p.stream); + return .replayed; + }; + const request = prefill_binding.requestFor(p, linear, layer, rows, sub, pos0, positions); + if (request) |target| if (prefillOriginCandidate(key, target)) |entry| { + try p.stream.synchronize(); + try entry.binding.?.apply(prefill_binding.Backend{ .exec = entry.exec }, target); + entry.key = key; + entry.used = clock; + prefill_stats.rebinds += 1; + prefill_stats.replays += 1; + prefill_stats.paged_replays += 1; + try entry.exec.launchOn(p.stream); + return .replayed; + }; + if (psighting(key) < prefillAfter()) { + prefill_stats.eager += 1; + return .eager; + } + var free_slot = false; + for (pentries) |entry| free_slot = free_slot or entry == null; + if (!free_slot) return prefillAdmissionMiss(); + _ = try memory.startupReservation(); + const admission = memory.admitWithQuery(.prefill, 0, p.stream, freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphBudgetExceeded, error.GraphMemoryHeadroom, error.GraphEvictionRequired => return prefillAdmissionMiss(), + else => return err, + }; + errdefer memory.release(admission) catch unreachable; + const binding_admission: ?memory.Ticket = if (request != null) memory.admitWithQuery(.prefill, 0, p.stream, freeMemory) catch null else null; + errdefer if (binding_admission) |ticket| memory.release(ticket) catch unreachable; + try p.stream.synchronize(); + try cuda.graph.beginCapture(p.stream, .thread_local); + capturing = true; + pcur = key; + pcur_admission = admission; + pcur_binding_request = request; + pcur_binding_admission = binding_admission; + return .capturing; +} + +/// After a layer's sub-blocks ran under `.capturing`: the graph is instantiated (or a cached one of the same shape +/// updated) and launched. +pub fn prefillEnd(p: anytype, how: Prefill) !void { + defer if (p.seq.paged) |pg| { + pg.graph_prepared = false; + }; + if (how != .capturing) return; + const t0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + capturing = false; + const key = pcur.?; + pcur = null; + const admission = pcur_admission.?; + pcur_admission = null; + const request = pcur_binding_request; + pcur_binding_request = null; + const binding_admission = pcur_binding_admission; + pcur_binding_admission = null; + var binding_admitted = true; + defer if (binding_admitted) if (binding_admission) |ticket| memory.release(ticket) catch unreachable; + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + var graph = try cuda.graph.endCapture(p.stream); + const binding: ?*prefill_binding.Binding = if (binding_admission != null and request != null) prefill_binding.Binding.create(std.heap.c_allocator, graph, request.?) catch null else null; + var kept = false; + defer if (!kept) { + if (binding) |owned| owned.deinit(std.heap.c_allocator) else graph.deinit(); + }; + if (binding != null) try memory.installed(binding_admission.?); + const at: *PEntry = blk: { + for (&pentries) |*e| if (e.*) |*x| if (binding == null and std.meta.eql(pshape(x.key), pshape(key))) { + if ((try x.exec.update(graph)) == .success) { + try uploadResident(&x.exec, p.stream); + x.key = key; + x.used = clock; + setPrefillBinding(x, binding, if (binding != null) binding_admission else null); + kept = binding != null; + binding_admitted = binding == null; + prefill_stats.updates += 1; + break :blk x; + } + }; + var best: usize = 0; + var best_used: u64 = std.math.maxInt(u64); + var same_shape: ?usize = null; + for (pentries, 0..) |entry, index| if (entry) |cached| { + if (std.meta.eql(pshape(cached.key), pshape(key))) { + same_shape = index; + break; + } + }; + for (&pentries, 0..) |*e, i| { + if (e.* == null) { + best = i; + break; + } + if (e.*.?.used < best_used) { + best_used = e.*.?.used; + best = i; + } + } + if (same_shape) |index| best = index; + if (pentries[best]) |*old| { + dropPrefill(old); + pentries[best] = null; + } + var exec = try graph.instantiate(); + errdefer exec.deinit(); + try uploadResident(&exec, p.stream); + try memory.installed(admission); + pentries[best] = .{ .key = key, .exec = exec, .used = clock, .admission = admission, .stream = p.stream }; + setPrefillBinding(&pentries[best].?, binding, if (binding != null) binding_admission else null); + kept = binding != null; + binding_admitted = binding == null; + admitted = false; + break :blk &pentries[best].?; + }; + prefill_stats.captures += 1; + prefill_stats.capture_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - t0); + try at.exec.launchOn(p.stream); +} + +/// A failed capture: closed and dropped; prefill graphs stay off for the process (the chunk's error propagates). +pub fn prefillAbort(p: anytype, how: Prefill) void { + if (p.seq.paged) |pg| pg.graph_prepared = false; + if (how != .capturing) return; + capturing = false; + pcur = null; + pcur_binding_request = null; + defer if (pcur_binding_admission) |ticket| { + memory.release(ticket) catch unreachable; + pcur_binding_admission = null; + }; + defer if (pcur_admission) |ticket| { + memory.release(ticket) catch unreachable; + pcur_admission = null; + }; + if (cuda.graph.endCapture(p.stream)) |g| { + var d = g; + d.deinit(); + } else |_| {} + std.log.warn("flashnext exl3 graph: a prefill layer could not be captured; prefill and decode graphs are off", .{}); + broken = true; +} + +const GraphTestDriver = struct { + var exec_drops: usize = 0; + var graph_drops: usize = 0; + var launches: usize = 0; + var exec_updates: usize = 0; + + fn dropExec(_: cuda.abi.GraphExec) callconv(.c) c_int { + exec_drops += 1; + return 0; + } + + fn dropGraph(_: cuda.abi.Graph) callconv(.c) c_int { + graph_drops += 1; + return 0; + } + + fn instantiate(out: *cuda.abi.GraphExec, _: cuda.abi.Graph, _: u64) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + fn getContext(out: *?*anyopaque) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + fn noContext(_: *?*anyopaque) callconv(.c) c_int { + return 1; + } + + fn setContext(_: ?*anyopaque) callconv(.c) c_int { + return 0; + } + + fn driver() cuda.Driver { + var d: cuda.Driver = undefined; + d.api.cuGraphExecDestroy = dropExec; + d.api.cuGraphDestroy = dropGraph; + d.api.cuGraphLaunch = launch; + d.api.cuGraphInstantiateWithFlags = instantiate; + d.api.cuMemGetInfo_v2 = memoryInfo; + d.api.cuCtxSynchronize = synchronize; + d.api.cuStreamSynchronize = synchronizeStream; + d.api.cuStreamBeginCapture_v2 = beginCapture; + d.api.cuStreamEndCapture = endCapture; + d.api.cuGraphUpload = upload; + d.api.cuGraphExecUpdate_v2 = update; + return d; + } + + fn memoryInfo(free: *usize, total: *usize) callconv(.c) c_int { + free.* = 100; + total.* = 100; + return 0; + } + + fn synchronize() callconv(.c) c_int { + return 0; + } + + fn synchronizeStream(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn beginCapture(_: cuda.abi.Stream, _: cuda.abi.CaptureMode) callconv(.c) c_int { + return 0; + } + + fn endCapture(_: cuda.abi.Stream, out: *cuda.abi.Graph) callconv(.c) c_int { + out.* = @ptrFromInt(2); + return 0; + } + + fn upload(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn update(_: cuda.abi.GraphExec, _: cuda.abi.Graph, info: *cuda.abi.ExecUpdateResultInfo) callconv(.c) c_int { + exec_updates += 1; + info.result = .success; + return 0; + } + + fn launch(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + launches += 1; + return 0; + } + + fn abort(_: cuda.nccl.Comm) callconv(.c) cuda.nccl.Result { + return 0; + } + + fn control(link: wire.Link) collective.Collective { + var c: collective.Collective = undefined; + c.comm = @ptrFromInt(1); + c.api.ncclCommAbort = abort; + c.link = link; + c.failed = false; + c.timeout_ms = 1000; + c.graph_step = 0; + return c; + } + + fn finishPending(pd: *Pending, driver_: *cuda.Driver) void { + pd.exec = .{ .d = driver_, .handle = @ptrFromInt(1) }; + pd.done.store(true, .release); + } +}; + +test "exl3: sequence graph retirement removes mixed lanes and pending work without evicting a live owner" { + clear(); + defer clear(); + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + var driver = GraphTestDriver.driver(); + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0][2] = 100; + const owner_key = key; + entries[0] = .{ .key = key, .exec = exec, .used = 1 }; + key.seq[0][2] = 200; + const live_key = key; + entries[1] = .{ .key = key, .exec = exec, .used = 1 }; + key.rows = 2; + key.seq[1][2] = 100; + entries[2] = .{ .key = key, .exec = exec, .used = 1 }; + const pending = try std.heap.c_allocator.create(Pending); + pending.* = .{ .graph = .{ .d = &driver, .handle = @ptrFromInt(1) }, .ctx = null, .keep_template = true }; + pending.thread = try std.Thread.spawn(.{}, GraphTestDriver.finishPending, .{ pending, &driver }); + entries[3] = .{ .key = owner_key, .pending = pending, .used = 1 }; + const live_pending = try std.heap.c_allocator.create(Pending); + live_pending.* = .{ .graph = .{ .d = &driver, .handle = @ptrFromInt(1) }, .ctx = null, .keep_template = true }; + live_pending.thread = try std.Thread.spawn(.{}, GraphTestDriver.finishPending, .{ live_pending, &driver }); + var live_pending_key = live_key; + live_pending_key.rows = 3; + entries[4] = .{ .key = live_pending_key, .pending = live_pending, .used = 1 }; + const pkey: PKey = .{ .layer = 0, .rows = 32, .sub = 16, .pos0 = 0, .deep = 1, .base = 2, .base_cap = 32, .seq = .{ 1, 2, 100, 32, 3 }, .paging = @splat(0), .wait_count = 0, .wait_key = 0, .scores = 4, .tiles = false }; + pentries[0] = .{ .key = pkey, .exec = exec, .used = 1 }; + var live_pkey = pkey; + live_pkey.seq[2] = 200; + pentries[1] = .{ .key = live_pkey, .exec = exec, .used = 1 }; + seen[0] = .{ .key = owner_key, .n = 10, .used = 1 }; + seen[1] = .{ .key = live_key, .n = 10, .used = 1 }; + pseen[0] = .{ .key = pkey, .n = 10 }; + pseen[1] = .{ .key = live_pkey, .n = 10 }; + const Seq = struct { tail: struct { ptr: u64 } }; + releaseSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expect(find(owner_key) == null); + try std.testing.expect(find(live_key) != null); + try std.testing.expect(entries[2] == null and entries[3] == null); + try std.testing.expect(entries[4].?.pending == live_pending); + try std.testing.expect(pentries[0] == null and pentries[1] != null); + try std.testing.expect(seen[0] == null and seen[1] != null); + try std.testing.expect(pseen[0] == null and pseen[1] != null); + try std.testing.expectEqual(@as(usize, 4), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.graph_drops); + try std.testing.expectEqual(@as(u32, 1), sighting(owner_key)); + releaseSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expectEqual(@as(usize, 4), GraphTestDriver.exec_drops); +} + +test "exl3: snapshot graph retirement keeps prefill templates but drops decode graphs" { + clear(); + defer clear(); + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + var driver = GraphTestDriver.driver(); + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0][2] = 100; + entries[0] = .{ .key = key, .exec = exec, .used = 1 }; + const pkey: PKey = .{ .layer = 0, .rows = 32, .sub = 16, .pos0 = 0, .deep = 1, .base = 2, .base_cap = 32, .seq = .{ 1, 2, 100, 32, 3 }, .paging = @splat(0), .wait_count = 0, .wait_key = 0, .scores = 4, .tiles = false }; + pentries[0] = .{ .key = pkey, .exec = exec, .used = 1 }; + pseen[0] = .{ .key = pkey, .n = 2 }; + const Seq = struct { tail: struct { ptr: u64 } }; + releaseDecodeSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expect(entries[0] == null); + try std.testing.expect(pentries[0] != null); + try std.testing.expect(pseen[0] != null); +} + +test "exl3: pending graph deadlines and failed instantiation never expose an executable" { + var pending: Pending = .{ .graph = undefined, .ctx = null, .keep_template = true, .started_ns = wire.now() }; + var entry: Entry = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .pending = &pending, .used = 1, .cohort = 7 }; + try std.testing.expectEqual(collective.GraphState.pending, graphVote(&entry, 1000).state); + pending.started_ns = wire.now() -| 2 * std.time.ns_per_ms; + try std.testing.expectEqual(collective.GraphState.timed_out, graphVote(&entry, 1).state); + entry.pending = null; + try std.testing.expectEqual(collective.GraphState.failed, graphVote(&entry, 1000).state); + drop(&entry); +} + +test "exl3: asynchronous TP instantiation preserves its NCCL template until owner retirement" { + if (!asyncInstantiate()) return error.SkipZigTest; + const previous_api = ctx_api; + defer ctx_api = previous_api; + ctx_api = .{ .get = GraphTestDriver.getContext, .set = GraphTestDriver.setContext }; + var driver = GraphTestDriver.driver(); + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(1) }; + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + const pending = (try startInstantiate(graph, true)).?; + var entry: Entry = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .pending = pending, .used = 1, .cohort = 9 }; + defer drop(&entry); + const until = wire.deadline(1000); + while (!pending.done.load(.acquire)) { + if (wire.now() >= until) return error.RankTimeout; + wire.pause(); + } + const vote = graphVote(&entry, 1000); + try std.testing.expectEqual(collective.GraphState.ready, vote.state); + try std.testing.expectEqual(@as(u64, 9), vote.cohort); + try std.testing.expect(entry.pending == null and entry.exec != null and entry.template != null); + try std.testing.expectEqual(@as(usize, 0), GraphTestDriver.graph_drops); +} + +test "exl3: rendezvous signatures compare topology and paging origins without device addresses" { + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0] = .{ 1, 2, 100, 4096, 3 }; + key.paging[0] = .{ 10, 20, 30, 40, 50, 60 }; + const signature = graphSignature(key, 4096); + key.seq[0][0] = 1000; + key.seq[0][2] = 1001; + key.paging[0][0] = 1002; + key.output = 1003; + try std.testing.expectEqual(signature, graphSignature(key, 4096)); + try std.testing.expect(signature != graphSignature(key, 4097)); + key.paging[0][5] += 1; + try std.testing.expect(signature != graphSignature(key, 4096)); +} + +test "exl3: shared paging rendezvous ignores physical ownership but local graph keys retain it" { + var first: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + first.seq[0] = .{ 1, 2, 3, 4096, 4 }; + first.paging[0] = .{ 10, 20, 30, 40, 50, 60 }; + first.paging_logical = @splat(@splat(0)); + first.paging_logical.?[0] = .{ 100, 200 }; + var second = first; + second.paging[0] = .{ 1000, 2000, 3000, 4000, 5000, 60 }; + try std.testing.expect(!std.meta.eql(first, second)); + try std.testing.expect(!std.meta.eql(state_binding.normalized(first), state_binding.normalized(second))); + try std.testing.expectEqual(graphSignature(first, 4096), graphSignature(second, 4096)); + second.paging_logical.?[0][0] += 1; + try std.testing.expect(graphSignature(first, 4096) != graphSignature(second, 4096)); + second = first; + second.paging_logical.?[0][1] += 1; + try std.testing.expect(graphSignature(first, 4096) != graphSignature(second, 4096)); + second = first; + second.seq[0][3] += 1; + try std.testing.expect(graphSignature(first, 4096) != graphSignature(second, 4096)); +} + +test "exl3: TP decode dispatch stays eager while its peer instantiates and replays only after agreement" { + clear(); + defer clear(); + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + const Peer = struct { + fn main(link: wire.Link, signature: u64) void { + var c = GraphTestDriver.control(link); + const first = c.graphAgree(signature, .{ .state = .pending, .cohort = 7 }) catch @panic("first peer agreement failed"); + std.debug.assert(first.choice == .eager); + const second = c.graphAgree(signature, .{ .state = .ready, .cohort = 7 }) catch @panic("second peer agreement failed"); + std.debug.assert(second.choice == .replay); + } + }; + const thread = try std.Thread.spawn(.{}, Peer.main, .{ wire.Link{ .fd = pair[1] }, graphSignature(key, 0) }); + defer thread.join(); + var c = GraphTestDriver.control(.{ .fd = pair[0] }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.launches = 0; + entries[0] = .{ .key = key, .used = 1, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .cohort = 7 }; + const Pass = struct { + io: std.Io = std.testing.io, + stream: cuda.Stream, + m: struct { collective: ?*collective.Collective }, + prepared: bool = false, + chunks: usize = 0, + + fn chunk(p: *@This(), _: []const u32, _: usize, _: ?[]u8) !void { + p.chunks += 1; + } + }; + var pass: Pass = .{ .stream = cuda.Stream.borrowed(&driver, null), .m = .{ .collective = &c } }; + try run(key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); + try std.testing.expectEqual(@as(usize, 0), GraphTestDriver.launches); + try run(key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.launches); +} + +test "exl3: TP pressure on either rank retires cohorts before retrying admission" { + clear(); + defer clear(); + for ([_]bool{ true, false }) |local_pressure| { + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 }, .prefill_charge = .{ .exec = 1, .template = 0, .transient = 0 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + const local = try memory.admit(.prefill, 0, 100); + try memory.installed(local); + pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 0, .admission = local }; + if (local_pressure) for (0..2) |index| { + const ticket = try memory.admit(.decode, index + 7, 100); + try memory.installed(ticket); + entries[index] = .{ .key = .{ .kind = 1, .rows = index + 1, .scratch = 1, .scores = 2, .output = 3 }, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = index + 1, .cohort = index + 7, .admission = ticket }; + }; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, pressure: bool) void { + var control = GraphTestDriver.control(link); + const first = control.graphAgree(42 ^ 0x61646d6974, .{ .state = if (pressure) .ready else .eager, .cohort = 11 }) catch @panic("pressure agreement failed"); + std.debug.assert(first.choice == .eager); + const victim = control.graphVictim(42 ^ 0x76696374696d, if (pressure) 8 else 7) catch @panic("victim agreement failed"); + std.debug.assert(victim == 7); + retirementBarrier(&control, 42 ^ 0x71756965736365, victim) catch @panic("quiescence agreement failed"); + retirementBarrier(&control, 42 ^ 0x72657469726564, victim) catch @panic("retirement agreement failed"); + const admitted = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .ready, .cohort = 11 }) catch @panic("retry agreement failed"); + std.debug.assert(admitted.choice == .replay); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, local_pressure }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + const ticket = try reserveTogether(&control, cuda.Stream.borrowed(&driver, null), 42, 11); + try std.testing.expectEqual(@as(u64, if (local_pressure) 71 else 41), memory.charged()); + try std.testing.expectEqual(@as(usize, if (local_pressure) 1 else 0), GraphTestDriver.exec_drops); + try std.testing.expect(pentries[0] != null); + if (local_pressure) try std.testing.expectEqual(@as(u64, 8), entries[1].?.cohort); + try memory.release(ticket); + clear(); + } +} + +test "exl3: exhausted TP victims reclaim locals incrementally with asymmetric peer caches" { + clear(); + defer clear(); + for ([_]bool{ false, true }) |local_cached| for ([_]bool{ false, true }) |peer_cached| { + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + if (local_cached) { + const prefill = try memory.admit(.prefill, 0, 100); + try memory.installed(prefill); + pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .admission = prefill }; + const local = try memory.admit(.decode, 0, 100); + try memory.installed(local); + entries[0] = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 2, .admission = local }; + } + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, local_cached_peer: bool, cached: bool) void { + var control = GraphTestDriver.control(link); + const first = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .eager, .cohort = 11 }) catch @panic("pressure agreement failed"); + std.debug.assert(first.choice == .eager); + const victim = control.graphVictim(42 ^ 0x76696374696d, 0) catch @panic("victim agreement failed"); + std.debug.assert(victim == 0); + retirementBarrier(&control, 42 ^ 0x6c6f63616c7175, 11) catch @panic("local quiescence failed"); + const reclaimed = control.graphVictim(42 ^ 0x6c6f63616c646f, @intFromBool(cached)) catch @panic("local destruction acknowledgment failed"); + std.debug.assert(reclaimed == @as(u64, @intFromBool(local_cached_peer or cached))); + if (reclaimed == 0) return; + const admitted = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .ready, .cohort = 11 }) catch @panic("retry agreement failed"); + std.debug.assert(admitted.choice == .replay); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, local_cached, peer_cached }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + const result = reserveTogether(&control, cuda.Stream.borrowed(&driver, null), 42, 11); + if (local_cached or peer_cached) { + const ticket = try result; + try std.testing.expectEqual(@as(u64, if (local_cached) 70 else 40), memory.charged()); + try memory.release(ticket); + } else { + try std.testing.expectError(error.GraphBudgetExceeded, result); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); + } + try std.testing.expectEqual(@as(usize, @intFromBool(local_cached)), GraphTestDriver.exec_drops); + try std.testing.expect(pentries[0] == null); + try std.testing.expectEqual(local_cached, entries[0] != null); + clear(); + }; +} + +test "exl3: TP updates agree the old cohort and transfer template and admission together" { + try testTPUpdate(true); +} + +test "exl3: asymmetric TP update rejection retires both execs before re-instantiating" { + try testTPUpdate(false); +} + +fn testTPUpdate(peer_accept: bool) !void { + if (!tpUpdate()) return error.SkipZigTest; + clear(); + defer clear(); + const previous_api = ctx_api; + defer ctx_api = previous_api; + ctx_api = .{ .get = GraphTestDriver.noContext, .set = GraphTestDriver.setContext }; + try memory.configure(.{ .reservation_bytes = 120, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_updates = 0; + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + const old_ticket = try memory.admit(.decode, 7, 100); + try memory.installed(old_ticket); + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 2, .scores = 2, .output = 3 }; + var old_key = key; + old_key.scratch = 1; + entries[0] = .{ .key = old_key, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .template = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .cohort = 7, .admission = old_ticket }; + var old_binding_request: state_binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + old_binding_request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + entries[0].?.binding = try state_binding.Binding.init(1, old_binding_request); + seen[0] = .{ .key = key, .n = std.math.maxInt(u32) - 1, .used = 1 }; + const previous_clock = clock; + defer clock = previous_clock; + clock = 1024; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, signature: u64, accept: bool) void { + var control = GraphTestDriver.control(link); + control.graph_step = 10; + const capture = control.graphAgree(signature, .{ .state = .capture }) catch @panic("capture agreement failed"); + std.debug.assert(capture.choice == .capture and capture.cohort == 11); + retirementBarrier(&control, signature ^ 0x61646d6974, 11) catch @panic("admission failed"); + retirementBarrier(&control, signature ^ 0x757064617465, 7) catch @panic("old cohort agreement failed"); + const result = control.graphAgree(signature ^ 0x75706461746564, .{ .state = if (accept) .ready else .capture, .cohort = 11 }) catch @panic("update result agreement failed"); + std.debug.assert(result.choice == if (accept) collective.GraphChoice.replay else collective.GraphChoice.capture); + if (!accept) retirementBarrier(&control, signature ^ 0x72656275696c64, 11) catch @panic("rebuild retirement agreement failed"); + retirementBarrier(&control, signature ^ 0x696e7374616c6c, 11) catch @panic("installation agreement failed"); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, graphSignature(key, 0), peer_accept }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + control.graph_step = 10; + const Pass = struct { + io: std.Io = std.testing.io, + stream: cuda.Stream, + prepared: bool = false, + chunks: usize = 0, + fn chunk(pass: *@This(), _: []const u32, _: usize, _: ?[]u8) !void { + pass.chunks += 1; + } + }; + var pass: Pass = .{ .stream = cuda.Stream.borrowed(&driver, null) }; + try runTogether(&control, key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_updates); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.graph_drops); + try std.testing.expectEqual(@as(usize, if (peer_accept) 0 else 1), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(u64, 30), memory.charged()); + try std.testing.expectEqual(@as(u64, 11), find(key).?.cohort); + try std.testing.expect(find(key).?.binding == null); + try std.testing.expectError(error.StaleGraphTicket, memory.release(old_ticket)); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); +} diff --git a/zig/src/families/flashnext/cuda_exl3_graph_memory.zig b/zig/src/families/flashnext/cuda_exl3_graph_memory.zig new file mode 100644 index 0000000000..6da55cf54b --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_graph_memory.zig @@ -0,0 +1,406 @@ +const std = @import("std"); +const policy = @import("exl3_graph_budget.zig"); + +pub const Charge = policy.Charge; +pub const Class = policy.Class; +pub const Policy = struct { reservation_bytes: u64, headroom_bytes: u64, auxiliary_bytes: u64 = 0, charge: Charge, prefill_charge: ?Charge = null, chain_charge: ?Charge = null, two_chunk_reserved: bool = false }; +const mib: u64 = 1024 * 1024; + +pub fn defaultPolicy() Policy { + return workloadPolicy(8, 48) catch unreachable; +} + +pub fn workloadPolicy(streams: u64, layers: u64) !Policy { + if (streams == 0 or layers == 0) return error.InvalidGraphPolicy; + const decode = try std.math.mul(u64, try std.math.mul(u64, streams, 2), 32 * mib); + const prefill = try std.math.mul(u64, layers, mib); + const chain = try std.math.mul(u64, try std.math.mul(u64, streams, 2), mib); + const auxiliary = try std.math.add(u64, prefill, chain); + const retained = try std.math.add(u64, decode, auxiliary); + return .{ .reservation_bytes = try std.math.add(u64, retained, 192 * mib), .headroom_bytes = 128 * mib, .auxiliary_bytes = try std.math.add(u64, auxiliary, 32 * mib), .charge = .{ .exec = 16 * mib, .template = 16 * mib, .transient = 32 * mib }, .prefill_charge = .{ .exec = mib, .template = 0, .transient = 32 * mib }, .chain_charge = .{ .exec = mib, .template = 0, .transient = 32 * mib } }; +} + +pub fn policyFromEnv() !Policy { + const stream_text = std.mem.span(std.c.getenv("TF_FLASHNEXT_EXL3_STREAMS") orelse "1"); + return policyFromEnvForStreams(try std.fmt.parseInt(u64, stream_text, 10)); +} + +fn policyFromEnvForStreams(streams: u64) !Policy { + var settings = try workloadPolicy(streams, 48); + const window = @import("exl3_prefill_window.zig"); + const two_chunk = window.enabled(if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK")) |value| std.mem.span(value) else null); + const proposed = try withWindowReservation(settings, two_chunk, null); + settings = try withWindowReservation(settings, two_chunk, try envBytes("TF_FLASHNEXT_EXL3_GRAPH_RESERVATION_MIB", proposed.reservation_bytes)); + settings.headroom_bytes = try envBytes("TF_FLASHNEXT_EXL3_GRAPH_HEADROOM_MIB", settings.headroom_bytes); + settings.charge.exec = try envBytes("TF_FLASHNEXT_EXL3_GRAPH_EXEC_MIB", settings.charge.exec); + settings.charge.template = try envBytes("TF_FLASHNEXT_EXL3_GRAPH_TEMPLATE_MIB", settings.charge.template); + settings.charge.transient = try envBytes("TF_FLASHNEXT_EXL3_GRAPH_TRANSIENT_MIB", settings.charge.transient); + return settings; +} + +fn withWindowReservation(base: Policy, enabled: bool, limit: ?u64) !Policy { + var settings = base; + const charge = settings.prefill_charge orelse settings.charge; + const extra = if (enabled) try @import("exl3_prefill_window.zig").graphBytes(charge.exec, charge.template, charge.transient) else 0; + const required = try std.math.add(u64, base.reservation_bytes, extra); + settings.reservation_bytes = limit orelse required; + settings.two_chunk_reserved = enabled and settings.reservation_bytes >= required; + if (settings.two_chunk_reserved) settings.auxiliary_bytes = try std.math.add(u64, base.auxiliary_bytes, extra); + return settings; +} + +test "exl3: two-chunk graph reservation includes bindings peaks and declines explicit underfunding" { + const base = try workloadPolicy(2, 48); + try std.testing.expectEqualDeep(base, try withWindowReservation(base, false, null)); + const expanded = try withWindowReservation(base, true, null); + try std.testing.expect(expanded.two_chunk_reserved); + try std.testing.expectEqual(@as(u64, 288 * mib), expanded.reservation_bytes - base.reservation_bytes); + try std.testing.expectEqual(@as(u64, 288 * mib), expanded.auxiliary_bytes - base.auxiliary_bytes); + const short = try withWindowReservation(base, true, expanded.reservation_bytes - 1); + try std.testing.expect(!short.two_chunk_reserved); + try std.testing.expectEqual(base.auxiliary_bytes, short.auxiliary_bytes); + var overflow = base; + overflow.reservation_bytes = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, withWindowReservation(overflow, true, null)); +} + +fn envBytes(name: [:0]const u8, fallback: u64) !u64 { + const value = std.c.getenv(name) orelse return fallback; + return parseMib(std.mem.span(value)); +} + +fn parseMib(value: []const u8) !u64 { + if (value.len == 0) return error.InvalidGraphPolicy; + for (value) |byte| if (byte < '0' or byte > '9') return error.InvalidGraphPolicy; + return std.math.mul(u64, try std.fmt.parseInt(u64, value, 10), mib); +} +const Ledger = policy.Budget(256); +pub const Ticket = Ledger.Ticket; +var guard: std.atomic.Value(bool) = .init(false); +var ledger: ?Ledger = null; +var configured: ?Policy = null; +var free_memory_hook: ?*const fn () anyerror!u64 = null; + +fn lock() void { + while (guard.swap(true, .acquire)) std.atomic.spinLoopHint(); +} + +fn unlock() void { + guard.store(false, .release); +} + +pub fn configure(settings: Policy) !void { + lock(); + defer unlock(); + try configureLocked(settings); +} + +fn configureLocked(settings: Policy) !void { + if (ledger) |*current| if (current.charged() != 0) return error.GraphBudgetBusy; + var next = try Ledger.init(settings.reservation_bytes, settings.headroom_bytes); + if (settings.auxiliary_bytes > next.limit or try settings.charge.peak() > next.limit - settings.auxiliary_bytes) return error.InvalidGraphPolicy; + _ = try next.plan(settings.charge); + if (settings.prefill_charge) |charge| _ = try next.plan(charge); + if (settings.chain_charge) |charge| _ = try next.plan(charge); + if (ledger) |*current| next.generation = current.generation; + ledger = next; + configured = settings; +} + +pub fn startupReservation() !u64 { + return startupReservationWithStreams(null); +} + +pub fn twoChunkReservationBytes() u64 { + lock(); + defer unlock(); + return if (configured) |settings| (if (settings.two_chunk_reserved) settings.reservation_bytes else 0) else 0; +} + +pub fn declineTwoChunkReservation() !u64 { + lock(); + defer unlock(); + var settings = configured orelse return 0; + if (!settings.two_chunk_reserved) return settings.reservation_bytes; + const charge = settings.prefill_charge orelse settings.charge; + const extra = try @import("exl3_prefill_window.zig").graphBytes(charge.exec, charge.template, charge.transient); + settings.reservation_bytes = try std.math.sub(u64, settings.reservation_bytes, extra); + settings.auxiliary_bytes = try std.math.sub(u64, settings.auxiliary_bytes, extra); + settings.two_chunk_reserved = false; + try configureLocked(settings); + return settings.reservation_bytes; +} + +test "exl3: declined two-chunk allocation returns the original graph reservation before capture" { + const base = try workloadPolicy(2, 48); + try configure(try withWindowReservation(base, true, null)); + try std.testing.expect(twoChunkReservationBytes() > base.reservation_bytes); + try std.testing.expectEqual(base.reservation_bytes, try declineTwoChunkReservation()); + try std.testing.expectEqual(@as(u64, 0), twoChunkReservationBytes()); + try std.testing.expectEqual(base.reservation_bytes, try declineTwoChunkReservation()); +} + +pub fn startupReservationForStreams(streams: u32) !u64 { + if (streams == 0) return error.InvalidGraphPolicy; + return startupReservationWithStreams(streams); +} + +fn startupReservationWithStreams(streams: ?u32) !u64 { + lock(); + defer unlock(); + if (configured == null) { + const proposed = if (streams) |count| try policyFromEnvForStreams(count) else try policyFromEnv(); + try configureLocked(proposed); + const settings = configured.?; + std.log.info("EXL3 graph budget auxiliary logical allowance: {d} MiB (heuristic, only the unmaterialized remainder is protected from TP admission)", .{settings.auxiliary_bytes / mib}); + std.log.info("EXL3 graph budget heuristic allowances (not measured exec bytes): reservation {d} MiB, headroom {d} MiB, exec/template/transient {d}/{d}/{d} MiB", .{ settings.reservation_bytes / mib, settings.headroom_bytes / mib, settings.charge.exec / mib, settings.charge.template / mib, settings.charge.transient / mib }); + std.log.info("EXL3 workload heuristic follows requested server concurrency (or TF_FLASHNEXT_EXL3_STREAMS for standalone tests): two decode variants per stream, 48 prefill layer execs and two MTP chain variants per stream; local exec allowances are heuristic, not measured GPU sizes; local templates are released after instantiation", .{}); + } + return configured.?.reservation_bytes; +} + +pub fn setFreeMemoryHook(hook: ?*const fn () anyerror!u64) void { + lock(); + defer unlock(); + free_memory_hook = hook; +} + +pub fn admit(class: Class, cohort: u64, free_bytes: u64) !Ticket { + lock(); + defer unlock(); + const available = if (free_memory_hook) |hook| try hook() else free_bytes; + return admitLocked(class, cohort, available); +} + +pub fn admitWithQuery(class: Class, cohort: u64, context: anytype, comptime query: anytype) !Ticket { + lock(); + defer unlock(); + const available = if (free_memory_hook) |hook| try hook() else try query(context); + return admitLocked(class, cohort, available); +} + +fn admitLocked(class: Class, cohort: u64, available: u64) !Ticket { + const settings = configured orelse return error.GraphBudgetUnconfigured; + const current = &ledger.?; + const charge = switch (class) { + .decode => settings.charge, + .prefill => settings.prefill_charge orelse settings.charge, + .chain => settings.chain_charge orelse settings.charge, + }; + const peak = try charge.peak(); + if (class == .decode and cohort != 0) { + const unmaterialized = settings.auxiliary_bytes -| current.localCharged(); + const allowance = current.limit - unmaterialized; + if (peak > allowance or current.charged() > allowance - peak) return error.GraphBudgetBusy; + } + const needed = try std.math.add(u64, try std.math.add(u64, current.pending(), peak), settings.headroom_bytes); + if (available < needed) return error.GraphMemoryHeadroom; + return current.reserve(class, cohort, charge, 0); +} + +pub fn installed(ticket: Ticket) !void { + lock(); + defer unlock(); + const current = if (ledger) |*value| value else return error.GraphBudgetUnconfigured; + try current.installed(ticket); +} + +pub fn release(ticket: Ticket) !void { + lock(); + defer unlock(); + const current = if (ledger) |*value| value else return error.GraphBudgetUnconfigured; + const entry = try current.getForRelease(ticket); + if (entry) { + try current.touch(ticket, 0, true); + try current.beginRetire(ticket, try current.cohortOf(ticket)); + } + try current.destroyed(ticket); +} + +pub fn charged() u64 { + lock(); + defer unlock(); + return if (ledger) |*current| current.charged() else 0; +} + +test "exl3: runtime admissions charge pending overlap and reject reconfiguration" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const first = try admit(.decode, 1, 100); + defer release(first) catch unreachable; + try std.testing.expectError(error.GraphMemoryHeadroom, admit(.chain, 2, 99)); + try std.testing.expectError(error.GraphBudgetBusy, configure(configured.?)); + try installed(first); + try std.testing.expectEqual(@as(u64, 30), charged()); + try std.testing.expectError(error.GraphMemoryHeadroom, admit(.prefill, 2, 59)); + const second = try admit(.prefill, 2, 60); + try std.testing.expectEqual(@as(u64, 70), charged()); + try release(second); + try std.testing.expectError(error.StaleGraphTicket, release(second)); +} + +test "exl3: runtime policy defaults are bounded heuristics and MiB overrides parse strictly" { + const settings = defaultPolicy(); + const bounded = try Ledger.init(settings.reservation_bytes, settings.headroom_bytes); + _ = try bounded.plan(settings.charge); + try std.testing.expectEqual(@as(u64, 768 * mib), settings.reservation_bytes); + try std.testing.expectEqual(@as(u64, 128 * mib), settings.headroom_bytes); + try std.testing.expectEqual(@as(u64, 7 * mib), try parseMib("7")); + try std.testing.expectError(error.InvalidGraphPolicy, parseMib("")); + try std.testing.expectError(error.InvalidGraphPolicy, parseMib("-1")); + try std.testing.expectError(error.InvalidGraphPolicy, parseMib("1x")); + try std.testing.expectError(error.Overflow, parseMib("18446744073709551615")); +} + +test "exl3: graph reservation scales with the requested concurrent stream count" { + const one = try workloadPolicy(1, 48); + const two = try workloadPolicy(2, 48); + try std.testing.expectEqual(@as(u64, 66 * mib), two.reservation_bytes - one.reservation_bytes); +} + +test "exl3: startup reservation configures the environment policy without a separate configure call" { + try std.testing.expectEqual(@as(u64, 0), charged()); + const previous_settings = configured; + const previous_ledger = ledger; + defer { + configured = previous_settings; + ledger = previous_ledger; + } + configured = null; + ledger = null; + const expected = try policyFromEnv(); + try std.testing.expectEqual(expected.reservation_bytes, try startupReservation()); + try std.testing.expectEqual(expected.reservation_bytes, try startupReservation()); +} + +test "exl3: installed physical allocations are not subtracted twice by the free query" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const Query = struct { + fn free(available: u64) !u64 { + return available; + } + }; + const first = try admitWithQuery(.decode, 1, @as(u64, 100), Query.free); + defer release(first) catch unreachable; + try std.testing.expectError(error.GraphMemoryHeadroom, admitWithQuery(.chain, 0, @as(u64, 60), Query.free)); + try installed(first); + const second = try admitWithQuery(.chain, 0, @as(u64, 60), Query.free); + try release(second); +} + +test "exl3: runtime policy reconfiguration cannot reuse a released ticket generation" { + const settings: Policy = .{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }; + try configure(settings); + const old = try admit(.decode, 1, 100); + try release(old); + try configure(settings); + const current = try admit(.decode, 1, 100); + defer release(current) catch unreachable; + try std.testing.expectError(error.StaleGraphTicket, release(old)); +} + +test "exl3: simultaneous runtime admissions cannot spend the same free headroom" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const Worker = struct { + fn run(result: *?Ticket) void { + result.* = admit(.chain, 0, 100) catch null; + } + }; + var results: [8]?Ticket = @splat(null); + var threads: [8]std.Thread = undefined; + for (&threads, &results) |*thread, *result| thread.* = try std.Thread.spawn(.{}, Worker.run, .{result}); + for (threads) |thread| thread.join(); + var admitted: usize = 0; + for (results) |result| if (result) |ticket| { + admitted += 1; + try release(ticket); + }; + try std.testing.expectEqual(@as(usize, 2), admitted); + try std.testing.expectEqual(@as(u64, 0), charged()); +} + +test "exl3: default TP cohorts leave a complete local admission peak" { + const settings = defaultPolicy(); + try configure(settings); + var tickets: [256]?Ticket = @splat(null); + defer for (tickets) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + var count: usize = 0; + while (count < tickets.len) : (count += 1) { + const ticket = admit(.decode, count + 1, 1024 * mib) catch |err| { + try std.testing.expect(err == error.GraphBudgetBusy); + break; + }; + tickets[count] = ticket; + try installed(ticket); + } + try std.testing.expectEqual(@as(usize, 16), count); + for ([_]Class{ .prefill, .chain }) |class| { + const local = try admit(class, 0, 1024 * mib); + try release(local); + } +} + +test "exl3: bounded prefill admission retains reusable keys across 48 layer scans" { + try configure(defaultPolicy()); + var cached: [48]?Ticket = @splat(null); + defer for (cached) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + var replays: usize = 0; + var misses: usize = 0; + var captures: usize = 0; + for (0..3) |_| for (&cached) |*entry| { + if (entry.* != null) { + replays += 1; + continue; + } + const ticket = admit(.prefill, 0, 1024 * mib) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphEvictionRequired => { + misses += 1; + continue; + }, + else => return err, + }; + entry.* = ticket; + try installed(ticket); + captures += 1; + }; + try std.testing.expectEqual(@as(usize, 48), captures); + try std.testing.expectEqual(@as(usize, 96), replays); + try std.testing.expectEqual(@as(usize, 0), misses); +} + +test "exl3: workload heuristic scales with concurrency and shape with checked arithmetic" { + try std.testing.expectEqual(@as(u64, 768 * mib), (try workloadPolicy(8, 48)).reservation_bytes); + try std.testing.expectEqual(@as(u64, 504 * mib), (try workloadPolicy(4, 48)).reservation_bytes); + try std.testing.expectEqual(@as(u64, 792 * mib), (try workloadPolicy(8, 72)).reservation_bytes); + try std.testing.expectError(error.InvalidGraphPolicy, workloadPolicy(0, 48)); + try std.testing.expectError(error.InvalidGraphPolicy, workloadPolicy(8, 0)); + try std.testing.expectError(error.Overflow, workloadPolicy(std.math.maxInt(u64), 48)); + try std.testing.expectError(error.Overflow, workloadPolicy(8, std.math.maxInt(u64))); + var invalid = defaultPolicy(); + invalid.prefill_charge = .{ .exec = std.math.maxInt(u64), .template = 0, .transient = 1 }; + try std.testing.expectError(error.Overflow, configure(invalid)); +} + +test "exl3: saturated TP cache retains all local workload variants without double protection" { + try configure(defaultPolicy()); + var tickets: [80]?Ticket = @splat(null); + defer for (tickets) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + for (0..16) |index| { + const ticket = try admit(.decode, index + 1, 1024 * mib); + tickets[index] = ticket; + try installed(ticket); + } + for (16..80) |index| { + const ticket = try admit(if (index < 64) .prefill else .chain, 0, 1024 * mib); + tickets[index] = ticket; + try installed(ticket); + } + try std.testing.expectEqual(@as(u64, 576 * mib), charged()); + try std.testing.expectError(error.GraphBudgetBusy, admit(.decode, 99, 1024 * mib)); + try release(tickets[0].?); + tickets[0] = null; + tickets[0] = try admit(.decode, 99, 1024 * mib); + try installed(tickets[0].?); + try std.testing.expectEqual(@as(u64, 576 * mib), charged()); + const local = try admit(.chain, 0, 1024 * mib); + try release(local); +} diff --git a/zig/src/families/flashnext/cuda_exl3_heads.zig b/zig/src/families/flashnext/cuda_exl3_heads.zig new file mode 100644 index 0000000000..cb2cd2276d --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_heads.zig @@ -0,0 +1,282 @@ +//! The attention head split across the two expert-parallel ranks (docs/plans/exl3-kv-memory.md lever (b), §3): rank r +//! keeps query heads 12r..12r+11 and KV head r of every full-attention layer, so each rank's K/V cache is half of +//! upstream's replicated one. The indexer (its projection, pooled keys, scores and selection) stays replicated: both +//! ranks list the same blocks from the same keys. +//! +//! Loader (cuda_prompt_load.zig's attention branch, one hook): the stitched q|k|v|indexer row keeps q_proj's columns +//! [6144r, 6144r + 6144) (12 heads of [query 256 | gate 256]), k_proj's and v_proj's [256r, 256r + 256) and the whole +//! indexer projection (640): 7,296 columns, all whole 128-column EXL3 Hadamard blocks, so every kept column is the +//! whole matrix's column bit for bit (the output transform is block diagonal). o_proj keeps input rows +//! [3072r, 3072r + 3072): 24 whole 128-row blocks of the input transform, so each rank's product is a partial sum of +//! the whole one and one fp32 all-reduce of [rows, 2560] a layer adds the two (not bit-identical to tp=1: fp32 order). +//! +//! Pass (cuda_batch.zig attnBlock, one hook keyed on the layer: its projection is 7,296 wide): upstream's kernels with +//! the rank's head counts (fn_attn_prep / _chunks / _merge / _gate take PW, NQ, NKV, H, HK, G as arguments), then +//! o_proj into s.part slice 0 and the collective sum. bf16 K/V only (the int8 cache refuses a split layer). +//! TF_FLASHNEXT_ATTN_SPLIT=0 keeps upstream's replicated attention at tp=2. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_exl3_load.zig"); +const tp = @import("cuda_exl3_tp.zig"); +const projection = @import("cuda_projection.zig"); +const Glue = @import("cuda_glue.zig").Glue; + +const head_dim = wt.head_dim; +/// Query heads and KV heads a rank keeps at world 2. +pub const rank_q: usize = wt.q_heads / 2; +pub const rank_kv: usize = wt.kv_heads / 2; +const index_cols: usize = wt.index_heads * wt.index_dim + wt.index_dim; +/// A rank's projection row: its queries and gates, its K, its V, the whole indexer. +pub const rank_pw: usize = rank_q * 2 * head_dim + 2 * rank_kv * head_dim + index_cols; +/// A rank's o_proj inputs. +pub const rank_out_k: usize = rank_q * head_dim; + +comptime { + std.debug.assert(wt.q_heads * 2 * head_dim + 2 * wt.kv_heads * head_dim + index_cols == wt.proj_n); + std.debug.assert(rank_pw == 7296 and rank_out_k == 3072); +} + +/// Whether attention layers split at `world`: at 2 unless TF_FLASHNEXT_ATTN_SPLIT=0; never at 1. K/V paging +/// (TF_FLASHNEXT_KVPAGE=1) reads the replicated layout only, so it takes precedence over the split. +pub fn splitAt(world: usize) bool { + if (world != 2) return false; + if (@import("cuda_exl3_kvpage.zig").configFromEnv() != null) return false; + const v = std.c.getenv("TF_FLASHNEXT_ATTN_SPLIT") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Whether a loaded layer's attention is a rank's half (its projection row is the rank's). +pub fn isSplit(lw: anytype) bool { + return !lw.linear and lw.proj.n == rank_pw; +} + +/// KV heads an engine's sequences hold: 1 when its attention layers are split. +pub fn kvHeads(layers: anytype) usize { + for (layers) |*l| if (!l.linear) return if (isSplit(l)) rank_kv else wt.kv_heads; + return wt.kv_heads; +} + +// -- loader --------------------------------------------------------------------------------------------------------- + +/// Input rows [lo, hi) of an EXL3 matrix in the stored layout [K/16][N/16][tile words]: contiguous whole tile rows. +pub fn shardRows(m: load.Mat, lo: usize, hi: usize) ![]const u8 { + if (lo % 128 != 0 or hi % 128 != 0 or hi <= lo or hi > m.k) return error.BadHeadSharding; + const row = (m.n / 16) * @as(usize, m.halves) * 16; + if (m.trellis.len != (m.k / 16) * row) return error.UnexpectedTensor; + return m.trellis[(lo / 16) * row .. (hi / 16) * row]; +} + +fn upPart(p: *load.Proj, d: *cuda.Driver, trellis: []const u8, suh: []const u8, svh: []const u8, k: usize, n: usize, halves: u32) !cuda.exl3.Weight { + var ptrs: [3]u64 = undefined; + for (&ptrs, [3][]const u8{ trellis, suh, svh }) |*b, bytes| { + try p.bufs.ensureUnusedCapacity(p.gpa, 1); + const db = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(db); + b.* = db.ptr; + } + const st = cuda.exl3.strides(k, n, halves, false); + return .{ .trellis = ptrs[0], .suh = ptrs[1], .svh = ptrs[2], .k = @intCast(k), .n = @intCast(n), .halves = halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; +} + +fn wrapProj(d: *cuda.Driver, p: *load.Proj) projection.Projection { + var out = projection.Projection.empty(d); + out.n = p.n; + out.k = p.k; + out.exl3 = p; + return out; +} + +/// The rank's q|k|v|indexer row of attention `ap` (".self_attn" prefix): four EXL3 column shards side by side. +pub fn rankProjection(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ap: []const u8, rank: usize) !projection.Projection { + if (rank > 1) return error.BadRank; + const p = try gpa.create(load.Proj); + p.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .k = wt.dims, .n = rank_pw, .parts = &.{} }; + errdefer p.destroy(); + const Cut = struct { name: []const u8, lo: usize, hi: usize }; + const qw = rank_q * 2 * head_dim; + const kw = rank_kv * head_dim; + const cuts = [_]Cut{ + .{ .name = ".q_proj", .lo = rank * qw, .hi = (rank + 1) * qw }, + .{ .name = ".k_proj", .lo = rank * kw, .hi = (rank + 1) * kw }, + .{ .name = ".v_proj", .lo = rank * kw, .hi = (rank + 1) * kw }, + .{ .name = ".indexer.index_qk_proj", .lo = 0, .hi = index_cols }, + }; + var parts: [cuts.len]load.Part = undefined; + var col: usize = 0; + var buf: [256]u8 = undefined; + for (cuts, &parts) |c, *out| { + const m = try load.mat(store, try std.fmt.bufPrint(&buf, "{s}{s}", .{ ap, c.name })); + if (m.k != wt.dims) return error.UnexpectedTensor; + const trellis = try tp.shardTrellis(gpa, m, c.lo, c.hi); + defer gpa.free(trellis); + const n = c.hi - c.lo; + const w = try upPart(p, d, trellis, m.suh, m.svh[c.lo * 2 .. c.hi * 2], m.k, n, m.halves); + p.ops.noteWeight(w, rank_pw); + out.* = .{ .col = col, .n = n, .kind = .{ .exl3 = w } }; + col += n; + } + if (col != rank_pw) return error.UnexpectedTensor; + p.parts = try gpa.dupe(load.Part, &parts); + return wrapProj(d, p); +} + +/// The rank's o_proj rows: inputs [3072r, 3072r + 3072) of `prefix` (".self_attn.o_proj"), all 2560 outputs. +pub fn rankOut(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, rank: usize) !projection.Projection { + if (rank > 1) return error.BadRank; + const m = try load.mat(store, prefix); + if (m.k != wt.out_k or m.n != wt.dims) return error.UnexpectedTensor; + const lo = rank * rank_out_k; + const hi = lo + rank_out_k; + const p = try gpa.create(load.Proj); + p.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .k = rank_out_k, .n = wt.dims, .parts = &.{} }; + errdefer p.destroy(); + const w = try upPart(p, d, try shardRows(m, lo, hi), m.suh[lo * 2 .. hi * 2], m.svh, rank_out_k, m.n, m.halves); + p.ops.noteWeight(w, wt.dims); + p.parts = try gpa.alloc(load.Part, 1); + p.parts[0] = .{ .col = 0, .n = wt.dims, .kind = .{ .exl3 = w } }; + return wrapProj(d, p); +} + +/// cuda_prompt_load.zig's hook: the rank's attention projections when `world` splits them (an EXL3 pack), else null. +pub fn loadRank(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ap: []const u8, world: usize, rank: usize) !?[2]projection.Projection { + if (!splitAt(world) or !load.has(store, (std.fmt.bufPrint(&name_buf, "{s}.q_proj", .{ap}) catch return null))) return null; + var proj = try rankProjection(gpa, d, store, ap, rank); + errdefer proj.free(); + const out = try rankOut(gpa, d, store, try std.fmt.bufPrint(&name_buf, "{s}.o_proj", .{ap}), rank); + return .{ proj, out }; +} + +var name_buf: [256]u8 = undefined; + +// -- pass ----------------------------------------------------------------------------------------------------------- + +/// fn_attn_chunks' dynamic smem (cuda_glue.zig chunkSmem) and its read-ahead row bound. +fn chunkSmem(buffers: u32) u32 { + return (16 + buffers * 2 * 32) * (256 + 8) * 2 + 16 * (32 + 8) * 2; +} +const ahead_rows = 16; + +fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: u32, shared: u32, a: *cuda.Args) !void { + if (grid[0] == 0 or grid[1] == 0 or grid[2] == 0) return; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = block }, .shared = shared }, s, a); +} + +fn ci(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// cuda_batch.zig Pass.attnBlock for a split layer: the rank's projection row, its heads' attention on its KV head, +/// the gate, o_proj's partial sum into s.part slice 0 (the other slices zeroed by the EXL3 lane), then the ranks' sum. +pub fn attnBlock(p: anytype, lw: anytype, ai: usize) !void { + const s = p.s; + const n = p.rows; + const pa = s.proj; + const g: *const Glue = p.tri.own orelse return error.NativeGlueRequired; + if (p.tri.set != null) return error.NativeGlueRequired; + const c = p.m.collective orelse return error.TensorParallelUnavailable; + try p.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); + if (p.lanes.len > 0) { + const own_seq = p.seq; + const own_pos0 = p.pos0; + defer { + p.seq = own_seq; + p.pos0 = own_pos0; + } + for (p.lanes, 0..) |l, r| { + p.seq = l.seq; + p.pos0 = l.pos; + try rows(p, g, lw, ai, pa + r * rank_pw * 2, 1, s.pos + r * 16, s.q + r * rank_q * head_dim * 2, s.iq + r * wt.index_heads * wt.index_dim * 2, s.ao + r * rank_out_k * 2); + } + } else try rows(p, g, lw, ai, pa, n, s.pos, s.q, s.iq, s.ao); + try p.tri.attnGate(s.ao, pa, s.gated, s.gxs_attn, n, rank_pw, rank_q, head_dim); + try p.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, wt.out_sk); + try c.sum(p.stream, s.part, n * wt.dims); +} + +/// `n` rows of p.seq from p.pos0 on the rank's heads: upstream's attnRows with NQ 12, NKV 1, PW 7296. +fn rows(p: anytype, g: *const Glue, lw: anytype, ai: usize, pa: u64, n: usize, pos_at: u64, q_at: u64, iq_at: u64, ao_at: u64) !void { + const s = p.s; + if (p.seq.media != null or p.seq.kv8 != null) return error.SplitAttentionUnsupported; + const kc = p.seq.kc[ai].ptr; + const vc = p.seq.vc[ai].ptr; + const ikc = p.seq.ikc[ai].ptr; + const pooled = p.seq.pooled[ai].ptr; + if (p.seq.kc[ai].len < p.seq.cap * rank_kv * head_dim * 2) return error.UnexpectedTensor; + const st = p.stream; + { + var a: cuda.Args = .{}; + for ([_]u64{ pa, pos_at, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, q_at, kc, vc, iq_at, ikc }) |v| a.add(v); + a.add(@as(f32, wt.eps)); + for ([_]usize{ rank_pw, rank_q, rank_kv, wt.index_heads, wt.index_dim, wt.half }) |v| a.add(ci(v)); + try go(g.f.prep, st, .{ n, rank_q + rank_kv + wt.index_heads + 1, 1 }, 256, 0, &a); + } + try p.tri.qsaPool(ikc, pooled, pos_at, lw.ik_scale.ptr, s.inv, wt.eps, n); + const ratio = wt.index_ratio; + const top = wt.index_top; + const att_rows: usize = 256; + const nb = (cdiv(p.seq.cap, ratio) + 15) / 16 * 16; + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = p.pos0 + r0 + m; + if (ends > p.seq.cap) return error.UnexpectedTensor; + const pos = pos_at + b * 16; + const iq = iq_at + r0 * wt.index_heads * wt.index_dim * 2; + if (ends / ratio > top) try p.tri.qsaScores(iq, pooled, pos, p.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(p.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + for ([_]usize{ nb, wt.index_idw, top, ratio }) |v| a.add(ci(v)); + try cuda.launch.launch(p.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, st, &a); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + { + var ca: cuda.Args = .{}; + for ([_]u64{ q_at + r0 * rank_q * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr }) |v| ca.add(v); + for ([_]usize{ rank_q, rank_kv, rank_q / rank_kv, 512, wt.nch }) |v| ca.add(ci(v)); + ca.add(@as(f32, 0.0625)); + for ([_]usize{ wt.index_idw, 1 }) |v| ca.add(ci(v)); + const ahead = m <= ahead_rows; + try go(if (ahead) g.f.chunks_ahead else g.f.chunks, st, .{ m, rank_kv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &ca); + } + { + var ma: cuda.Args = .{}; + for ([_]u64{ s.po, s.pm, s.pl, pos, ao_at + r0 * rank_out_k * 2, s.qnk, s.qspr }) |v| ma.add(v); + for ([_]usize{ rank_q, rank_kv, rank_q / rank_kv, 512, wt.nch, 1 }) |v| ma.add(ci(v)); + try go(g.f.merge, st, .{ m, rank_q, 1 }, 256, 0, &ma); + } + } +} + +test "exl3: a rank's attention row is half the heads and the whole indexer, in whole 128-column blocks" { + try std.testing.expectEqual(@as(usize, 7296), rank_pw); + try std.testing.expectEqual(@as(usize, 0), rank_pw % 128); + try std.testing.expectEqual(@as(usize, 0), rank_out_k % 128); + try std.testing.expectEqual(@as(usize, 12), rank_q / rank_kv); +} + +test "exl3: o_proj's input-row shard is the contiguous tile rows" { + var tr: [4 * 2 * 4 * 16]u8 = undefined; // K 512 (32 tile rows), N 32 (2 tiles) at 2 halves... (shape-only check) + _ = &tr; + const words = try std.testing.allocator.alloc(u8, (256 / 16) * (128 / 16) * 2 * 16); + defer std.testing.allocator.free(words); + for (words, 0..) |*w, i| w.* = @truncate(i); + const m: load.Mat = .{ .trellis = words, .suh = &.{}, .svh = &.{}, .k = 256, .n = 128, .halves = 2 }; + const row = (128 / 16) * 2 * 16; + const hi = try shardRows(m, 128, 256); + try std.testing.expectEqual(@as(usize, 8 * row), hi.len); + try std.testing.expectEqual(words[8 * row], hi[0]); + try std.testing.expectError(error.BadHeadSharding, shardRows(m, 64, 128)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_kv.zig b/zig/src/families/flashnext/cuda_exl3_kv.zig new file mode 100644 index 0000000000..fb0666cc04 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_kv.zig @@ -0,0 +1,720 @@ +//! KV memory for 200k-token contexts on 16 GiB cards (K10, docs/plans/exl3-kv-memory.md): the per-sequence and +//! per-card byte arithmetic of each lever, and the int8-KV attention kernels (flashnext_exl3_kv.cu). +//! +//! The kernels are upstream's flashnext_attn.cu prep / chunks / merge with the cache stored as int8 codes and one fp16 +//! scale a 32-group (exl3_cache.zig's policy: an optional normalized 32-point Hadamard per group). Their pointers may be +//! device memory or mapped pinned host memory (`HostBuffer.allocMapped(..).device()`): the chunk kernel reads only the +//! keys the sparse indexer listed (<= 2,051 a row), so a host-resident cache costs PCIe reads, not device bytes. +//! Not wired into the pass yet: the hook is one branch in cuda_batch.zig attnRows (see the plan). + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const kvpage = @import("cuda_exl3_kvpage.zig"); + +pub const Kv = enum { bf16, int8, int4 }; + +const group: usize = 32; +const head_dim = wt.head_dim; +const kv_heads = wt.kv_heads; +const index_row: usize = wt.index_dim * 2; + +/// One token's K and V in one attention layer for `heads` KV heads (codes plus a 2-byte scale a 32-group). +pub fn kvRowBytes(kv: Kv, heads: usize) usize { + const per_head: usize = switch (kv) { + .bf16 => head_dim * 2, + .int8 => head_dim + head_dim / group * 2, + .int4 => head_dim / 2 + head_dim / group * 2, + }; + return 2 * heads * per_head; +} + +/// The levers that change a sequence's device bytes. +pub const Policy = struct { + kv: Kv = .bf16, + /// KV heads a rank keeps (2: replicated as upstream; 1: the attention head split across two ranks). + heads: usize = kv_heads, + /// K/V in mapped pinned host memory, read in place by the chunk kernel (only selected keys cross PCIe). + kv_host: bool = false, + /// The raw indexer keys [tokens][128] in mapped host memory (read only when a block of 4 completes and pools). + ikc_host: bool = false, +}; + +/// Bytes a sequence of `cap` tokens holds, split by where they live. +pub const SeqBytes = struct { + gdn: u64, + kv: u64, + ikc: u64, + pooled: u64, + tail: u64, + host: u64, + + pub fn device(s: SeqBytes) u64 { + return s.gdn + s.kv + s.ikc + s.pooled + s.tail; + } +}; + +pub fn pooledBytes(cap: usize) u64 { + return (cap + wt.index_ratio - 1) / wt.index_ratio * index_row; +} + +/// cuda_prompt.zig seqBytes split by lever (the default policy gives upstream's number: a host test checks it). +pub fn seqBytes(cap: usize, p: Policy) SeqBytes { + const attn: u64 = wt.attn_n; + const kv: u64 = attn * cap * kvRowBytes(p.kv, p.heads); + const ikc: u64 = attn * cap * index_row; + var s: SeqBytes = .{ + .gdn = wt.linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + wt.state_n * 4), + .kv = if (p.kv_host) 0 else kv, + .ikc = if (p.ikc_host) 0 else ikc, + .pooled = attn * pooledBytes(cap), + .tail = wt.ple_tail_n * wt.wide * 2, + .host = 0, + }; + if (p.kv_host) s.host += kv; + if (p.ikc_host) s.host += ikc; + return s; +} + +/// The MTP head's one attention layer of its draft cache (the part of cuda_mtp_draft.zig Draft.bytes that grows with +/// `cap`) under the same policy; the rest of Draft.bytes is fixed (streams, logits, saved states). +pub fn draftAttnBytes(cap: usize, p: Policy) u64 { + return (if (p.kv_host) 0 else @as(u64, cap) * kvRowBytes(p.kv, p.heads)) + (if (p.ikc_host) 0 else @as(u64, cap) * index_row) + pooledBytes(cap); +} + +pub const gib: f64 = 1024.0 * 1024.0 * 1024.0; + +// -- the int8 kernels ----------------------------------------------------------------------------------------------- + +const tile_heads: u32 = 16; +const tile_keys: u32 = 32; + +/// Dynamic shared bytes of fn_kv8_chunks with `nb` raw buffers (see the .cu: q rows, one bf16 tile pair, p, raw codes). +pub fn chunkSmem(nb: u32) u32 { + const row: u32 = @intCast((head_dim + 8) * 2); + return tile_heads * row + 2 * tile_keys * row + tile_heads * (tile_keys + 8) * 2 + nb * 2 * tile_keys * @as(u32, @intCast(head_dim)) + nb * 2 * tile_keys * @as(u32, @intCast(head_dim / group)) * 2; +} + +/// Rows up to which the chunk kernel keeps the next tile's codes in flight (upstream's bf16 kernel: 16). The int8 +/// tile's smem is one bf16 pair plus NB raw buffers, so NB 2 costs 17 KiB more, not a second bf16 pair. +pub var ahead_rows: usize = 16; + +pub const Kv8 = struct { + module: cuda.Module, + prep: cuda.Function, + chunks: cuda.Function, + chunks_ahead: cuda.Function, + merge: cuda.Function, + pack: cuda.Function, + gather: cuda.Function, + roundtrip: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kv8 { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kv8 = undefined; + k.module = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_kv); + errdefer k.module.unload(); + k.prep = try k.module.function("fn_kv8_prep"); + k.chunks = try k.module.function("fn_kv8_chunks"); + k.chunks_ahead = try k.module.function("fn_kv8_chunks_ahead"); + k.merge = try k.module.function("fn_kv8_merge"); + k.pack = try k.module.function("fn_kv8_pack"); + k.gather = try k.module.function("fn_kv_gather"); + k.roundtrip = try k.module.function("fn_kv_roundtrip"); + try k.chunks.allowDynamicShared(chunkSmem(1)); + try k.chunks_ahead.allowDynamicShared(chunkSmem(2)); + return k; + } + + pub fn deinit(k: *Kv8) void { + k.module.unload(); + } + + fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: u32, shared: u32, a: *cuda.Args) !void { + if (grid[0] == 0 or grid[1] == 0 or grid[2] == 0) return; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = block }, .shared = shared }, s, a); + } + + fn int(v: usize) c_int { + return @intCast(v); + } + + /// A layer's int8 cache: codes [cap][heads][256] and scales [cap][heads][8] for K and for V. + pub const Cache = struct { k8: u64, ks: u64, v8: u64, vs: u64 }; + + pub fn cacheBytes(cap: usize, heads: usize) struct { codes: usize, scales: usize } { + return .{ .codes = cap * heads * head_dim, .scales = cap * heads * (head_dim / group) * 2 }; + } + + /// fn_attn_prep with the K/V rows written as codes and scales; `nq`/`nkv` the rank's heads (proj row `pw` wide). + pub fn prepRows(k: *const Kv8, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, c: Cache, iq: u64, ikc: u64, eps: f32, rows: usize, pw: usize, nq: usize, nkv: usize, rot: bool) !void { + var a: cuda.Args = .{}; + for ([_]u64{ proj, pos, q_scale, k_scale, i_scale, inv, q, c.k8, c.ks, c.v8, c.vs, iq, ikc }) |v| a.add(v); + a.add(eps); + for ([_]usize{ pw, nq, nkv, wt.index_heads, wt.index_dim, wt.half, @intFromBool(rot) }) |v| a.add(int(v)); + try go(k.prep, s, .{ rows, nq + nkv + wt.index_heads + 1, 1 }, @intCast(head_dim), 0, &a); + } + + /// Upstream's attnChunks over an int8 cache (same grid, same 512-key chunks, `nq` query heads over `nkv` KV heads). + pub fn attnChunks(k: *const Kv8, s: cuda.Stream, q: u64, c: Cache, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, nq: usize, nkv: usize, chunks: usize, qsa: bool) !void { + const g = nq / nkv; + if (g > tile_heads) return error.RowShape; + var a: cuda.Args = .{}; + for ([_]u64{ q, c.k8, c.ks, c.v8, c.vs, pos, po, pm, pl, ids, nk, sparse }) |v| a.add(v); + for ([_]usize{ nq, nkv, g, 512, if (qsa) wt.nch else wt.dense_nch }) |v| a.add(int(v)); + a.add(@as(f32, 0.0625)); + for ([_]usize{ wt.index_idw, @intFromBool(qsa) }) |v| a.add(int(v)); + const ahead = rows <= ahead_rows; + try go(if (ahead) k.chunks_ahead else k.chunks, s, .{ rows, nkv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &a); + } + + pub fn attnMerge(k: *const Kv8, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, nq: usize, nkv: usize, qsa: bool, rot: bool) !void { + var a: cuda.Args = .{}; + for ([_]u64{ po, pm, pl, pos, out, nk, sparse }) |v| a.add(v); + for ([_]usize{ nq, nkv, nq / nkv, 512, if (qsa) wt.nch else wt.dense_nch, @intFromBool(qsa), @intFromBool(rot) }) |v| a.add(int(v)); + try go(k.merge, s, .{ rows, nq, 1 }, @intCast(head_dim), 0, &a); + } + + /// `rows` rows of 256 bf16 (a (pos, head) row each) into codes and scales. + pub fn packRows(k: *const Kv8, s: cuda.Stream, x: u64, codes: u64, scales: u64, rows: usize, rot: bool) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(codes); + a.add(scales); + a.add(@as(i64, @intCast(rows))); + a.add(int(@intFromBool(rot))); + try go(k.pack, s, .{ rows, 1, 1 }, @intCast(head_dim), 0, &a); + } + + /// Test support: `rows` bf16 rows of 256 replaced in place by an int8 (bits 8) or int4 (bits 4) cache's values. + pub fn roundTrip(k: *const Kv8, s: cuda.Stream, x: u64, rows: usize, rot: bool, bits: u32) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(rows))); + a.add(int(@intFromBool(rot))); + a.add(@as(c_int, @intCast(bits))); + try go(k.roundtrip, s, .{ rows, 1, 1 }, @intCast(head_dim), 0, &a); + } + + /// `n` rows of `row_bytes` from `src` at `ids` into `dst` (a warp a row, 8 rows a block). + pub fn gatherRows(k: *const Kv8, s: cuda.Stream, src: u64, dst: u64, ids: u64, n: usize, row_bytes: usize) !void { + if (row_bytes % 16 != 0) return error.RowShape; + var a: cuda.Args = .{}; + a.add(src); + a.add(dst); + a.add(ids); + a.add(int(n)); + a.add(int(row_bytes / 16)); + try go(k.gather, s, .{ (n + 7) / 8, 1, 1 }, 256, 0, &a); + } +}; + +// -- the pass's attention block over an int8 cache ------------------------------------------------------------------ + +/// cuda_batch.zig Pass.attnBlock with the layer's K/V in `c` (int8 codes + scales) instead of seq.kc/vc: the same +/// projection, pooling, block scores, selection, gate and output projection; prep, chunks and merge are ours. The hook +/// in the pass is one branch at the top of attnBlock (`if (self.seq.kv8) |c| return exl3_kv.attnBlock(...)`, plan). +/// One sequence's rows (a prompt chunk or one stream's decode rows); a decode round over several lanes is refused here. +pub fn attnBlock(p: anytype, k8: *const Kv8, lw: anytype, ai: usize, c: Kv8.Cache, rot: bool) !void { + const s = p.s; + const n = p.rows; + if (p.lanes.len > 0) return error.LanesUnsupported; + const pa = s.proj; + const att_rows: usize = 256; + const ratio = wt.index_ratio; + const top = wt.index_top; + const idw = wt.index_idw; + const proj_n = wt.proj_n; + try p.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); + const ikc = p.seq.ikc[ai].ptr; + const pooled = p.seq.pooled[ai].ptr; + try k8.prepRows(p.stream, pa, s.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, s.q, c, s.iq, ikc, 1e-6, n, proj_n, wt.q_heads, kv_heads, rot); + try p.tri.qsaPool(ikc, pooled, s.pos, lw.ik_scale.ptr, s.inv, 1e-6, n); + const nb = (cdiv(p.seq.cap, ratio) + 15) / 16 * 16; + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = p.pos0 + r0 + m; + if (ends > p.seq.cap) return error.UnexpectedTensor; + const pos = s.pos + b * 16; + const iq = s.iq + r0 * wt.index_heads * wt.index_dim * 2; + if (ends / ratio > top) try p.tri.qsaScores(iq, pooled, pos, p.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(p.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + for ([_]usize{ nb, idw, top, ratio }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(p.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, p.stream, &a); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + const q = s.q + r0 * wt.q_heads * head_dim * 2; + try k8.attnChunks(p.stream, q, c, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, wt.q_heads, kv_heads, chunks, true); + try k8.attnMerge(p.stream, s.po, s.pm, s.pl, pos, s.ao + r0 * wt.out_k * 2, s.qnk, s.qspr, m, wt.q_heads, kv_heads, true, rot); + } + try p.tri.attnGate(s.ao, pa, s.gated, s.gxs_attn, n, proj_n, wt.q_heads, head_dim); + try p.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, wt.out_sk); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +// -- the int8 cache in the family's engine (plan §8): `of(engine)`, Seq/SeqState.kv8, Pass.attnRows ------------- + +/// What an engine's new sequences hold (set once at load by `setup`, before any sequence; kept here by engine address +/// so prompt.Engine carries no field: upstream's tests build engines field by field). +pub const EngineKv = struct { + /// The int8 kernels, loaded: every new sequence keeps int8 K/V (kc/vc hold the codes, SeqKv8 the scales). + k8: ?*Kv8 = null, + /// Each new sequence's raw indexer keys [cap][128] in mapped pinned host memory (read only when a block pools). + ikc_host: bool = false, + /// The 32-point Hadamard on each group (K10: off by default; plain int8 sat at int12's noise floor). + rot: bool = false, + /// KV heads a sequence holds: 1 when the attention heads are split across two ranks (cuda_exl3_heads.zig). + heads: usize = kv_heads, + + pub fn policy(k: EngineKv) Policy { + return .{ .kv = if (k.k8 != null) .int8 else .bf16, .ikc_host = k.ikc_host, .heads = k.heads }; + } + + /// Whether new sequences differ from upstream's newSeq. + pub fn placed(k: EngineKv) bool { + return k.k8 != null or k.ikc_host or k.heads != kv_heads; + } + + /// A snapshot's rows of one layer a token: K (and V) bytes, and their int8 scales (0: bf16). + pub fn rows(k: EngineKv) Rows { + return snapRows(k.k8 != null, k.heads); + } +}; + +pub const Rows = struct { kv: usize, scales: usize }; + +/// A sequence's int8 cache beside its codes in kc/vc: each attention layer's K and V scales [cap][heads][8] fp16. +pub const SeqKv8 = struct { + k: *const Kv8, + ks: []cuda.DeviceBuffer, + vs: []cuda.DeviceBuffer, + rot: bool = false, + + pub fn cache(c: SeqKv8, kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ai: usize) Kv8.Cache { + return .{ .k8 = kc[ai].ptr, .ks = c.ks[ai].ptr, .v8 = vc[ai].ptr, .vs = c.vs[ai].ptr }; + } +}; + +/// TF_FLASHNEXT_KV=int8 (or int8r: rotated groups), TF_FLASHNEXT_IKC_HOST=1|0 (default: on with int8). +pub fn fromEnv(kv8_flag: bool) struct { int8: bool, rot: bool, ikc_host: bool } { + var int8 = kv8_flag; + var rot = false; + if (std.c.getenv("TF_FLASHNEXT_KV")) |v| { + const s = std.mem.span(v); + if (std.mem.eql(u8, s, "int8") or std.mem.eql(u8, s, "int8r")) int8 = true; + if (std.mem.eql(u8, s, "int8r")) rot = true; + if (std.mem.eql(u8, s, "bf16")) int8 = false; + } + var ikc_host = int8; + if (std.c.getenv("TF_FLASHNEXT_IKC_HOST")) |v| ikc_host = std.mem.eql(u8, std.mem.span(v), "1"); + return .{ .int8 = int8, .rot = rot, .ikc_host = ikc_host }; +} + +/// Loads the int8 kernels into `e` (an engine with no sequence yet) and sets its sequences' placement. +pub fn enable(e: anytype, int8: bool, rot: bool, ikc_host: bool) !void { + const kv = try entryOf(e); + if (int8 and kv.k8 == null) { + const k = try e.gpa.create(Kv8); + errdefer e.gpa.destroy(k); + k.* = try Kv8.load(e.driver); + kv.k8 = k; + } + kv.rot = rot; + kv.ikc_host = ikc_host; + std.log.info("flashnext: KV caches {s}{s}, indexer keys in {s} memory", .{ if (kv.k8 != null) "int8 (32-groups, fp16 scales)" else "bf16", if (rot and kv.k8 != null) " rotated" else "", if (ikc_host) "mapped pinned host" else "device" }); +} + +/// Widest logits pass the server runs (a decode round of 32 streams, an MTP window): `head_local`'s rows. +pub const head_local_serve_rows: usize = 64; +/// K10 lever 3: the pass's chunk for an EXL3 checkpoint unless TF_FLASHNEXT_CHUNK says (+3% prompt time a chunk at +/// 1024 vs 1536, 0.45 GiB less scratch). +pub const exl3_chunk: usize = 1024; + +/// The server's load hook (cuda_lanes_load.zig, right after the engine loads, before any scratch or sequence): the +/// KV placement from `--kv-cache fp8` (int8 here: our cache is int8 codes, not e4m3) and TF_FLASHNEXT_KV / +/// TF_FLASHNEXT_IKC_HOST; for an EXL3 checkpoint also the chunk (1024) and `head_local` sized for logits passes only. +pub fn setup(e: anytype, kv8_flag: bool, exl3: bool) !void { + const f = fromEnv(kv8_flag); + const heads = @import("cuda_exl3_heads.zig").kvHeads(e.layers[0..e.n_layers]); + if (kvpage.configFromEnv() != null and (f.int8 or heads != kv_heads)) return error.KvPagePlacement; + if (heads != kv_heads) (try entryOf(e)).heads = heads; + if (f.int8 or f.ikc_host) try enable(e, f.int8, f.rot, f.ikc_host); + if (heads != kv_heads) std.log.info("flashnext: attention heads split across the ranks: {d} query heads and {d} KV head a rank", .{ wt.q_heads / 2, heads }); + if (!exl3) return; + if (e.batch_s != null) return error.ScratchAlreadyReserved; + const decoder_exl3 = e.n_layers != 0 and for (e.layers[0..e.n_layers]) |layer| { + if (layer.exl3 == null) break false; + } else true; + const draft_exl3 = @import("cuda_exl3_load.zig").hasExperts(&e.store, "language_model.mtp.layers.0.mlp"); + e.scratch_format = if (decoder_exl3 and draft_exl3) .exl3 else .classic; + if (std.c.getenv("TF_FLASHNEXT_CHUNK") == null) e.chunk = exl3_chunk; + @import("cuda_batch_state.zig").head_local_rows = if (e.head_q.n == wt.vocab) 0 else head_local_serve_rows; +} + +/// One log line: a stream's device bytes at the window by piece, and the host bytes it maps. +pub fn logStream(e: anytype, window: usize, stream_bytes: usize) !void { + const kv = of(e); + const layout = try sequenceLayout(kv, window, e.chunk, kvpage.configFromEnv()); + const b = layout.bytes; + if (layout.paging) |cfg| std.log.info("flashnext: BF16 K/V RAM paging enabled: window {d}, slots {d}, rows {d}", .{ cfg.window, cfg.slots, cfg.rows }); + std.log.info("flashnext: a stream at a {d}-token window holds {d:.3} GiB on the device (K/V {d:.3} {s}, indexer keys {d:.3}, pooled {d:.3}, DeltaNet {d:.3}, draft and the rest {d:.3}) and {d:.3} GiB of mapped host memory; chunk {d} rows, scratch {d:.3} GiB", .{ + window, + @as(f64, @floatFromInt(stream_bytes)) / gib, + @as(f64, @floatFromInt(b.kv)) / gib, + if (kv.k8 != null) "int8" else "bf16", + @as(f64, @floatFromInt(b.ikc)) / gib, + @as(f64, @floatFromInt(b.pooled)) / gib, + @as(f64, @floatFromInt(b.gdn)) / gib, + @as(f64, @floatFromInt(stream_bytes -| b.device())) / gib, + @as(f64, @floatFromInt(b.host)) / gib, + e.chunk, + @as(f64, @floatFromInt(@import("cuda_batch_state.zig").Scratch.bytesForFormat(e.chunk, e.scratch_format))) / gib, + }); +} + +/// Whether kept prompt states go to pinned host memory (cuda_snap.on_host, set here): TF_FLASHNEXT_SNAP_HOST=1|0, +/// by default for an EXL3 checkpoint served past 32,768 tokens (its streams leave the card no room for 100+ MiB states +/// that grow by 27 KiB a token). +pub fn snapHost(context: usize, exl3: bool) bool { + var on = exl3 and context > 32768; + if (std.c.getenv("TF_FLASHNEXT_SNAP_HOST")) |v| on = std.mem.eql(u8, std.mem.span(v), "1"); + if (kvpage.configFromEnv() != null) on = true; + @import("cuda_snap.zig").on_host = on; + return on; +} + +/// cuda_lanes.zig rankRoom: the ranks agree how many STREAMS fit (each its own stream_bytes: rank 1 holds no MTP head +/// cache, so its streams are smaller), returned as this rank's bytes for that many streams. The native host then +/// admits the same count on both ranks. Fixed point (1/65536 of a stream) through Session.room's min. +pub fn rankRoom(r: anytype, local: u64, stream_bytes: usize) !u64 { + if (stream_bytes == 0) return r.room(local); + const sb: u64 = stream_bytes; + const q = (local / sb) << 16 | ((local % sb) << 16) / sb; + const m = try r.room(q); + return (m >> 16) * sb + ((m & 0xffff) * sb >> 16); +} + +/// The engine's deinit. +pub fn release(gpa: std.mem.Allocator, k: *EngineKv) void { + if (k.k8) |x| { + x.deinit(); + gpa.destroy(x); + } + k.* = .{}; +} + +pub const SeqLayout = struct { + cache_rows: usize, + row: Rows, + paging: ?kvpage.Config, + bytes: SeqBytes, +}; + +pub fn sequenceLayout(placement: EngineKv, cap: usize, chunk: usize, paging: ?kvpage.Config) !SeqLayout { + if (cap == 0) return error.UnexpectedTensor; + const page_bytes = if (paging) |requested| try kvpage.deviceBytesFor(cap, requested.forChunk(chunk)) else 0; + var layout: SeqLayout = .{ .cache_rows = cap, .row = placement.rows(), .paging = null, .bytes = seqBytes(cap, placement.policy()) }; + if (paging) |requested| { + if (placement.k8 != null or placement.heads != kv_heads) return error.KvPagePlacement; + const cfg = requested.forChunk(chunk); + try cfg.validate(); + layout.cache_rows = 1; + layout.paging = cfg; + layout.bytes.host = try std.math.add(u64, layout.bytes.host, layout.bytes.kv); + layout.bytes.kv = try std.math.add(u64, page_bytes, wt.attn_n * 2 * layout.row.kv); + } + return layout; +} + +/// Device bytes of the sequence caches under the same placement and paging shape used by newSeq. +pub fn engineSeqBytes(k: EngineKv, cap: usize, chunk: usize) !usize { + return @intCast((try sequenceLayout(k, cap, chunk, kvpage.configFromEnv())).bytes.device()); +} + +/// The snapshot rows of one layer of a sequence of `heads` KV heads: codes or bf16 K (and V) a token, and scales a +/// token (0: bf16). +pub fn snapRows(int8: bool, heads: usize) Rows { + return if (int8) .{ .kv = heads * head_dim, .scales = heads * (head_dim / group) * 2 } else .{ .kv = heads * head_dim * 2, .scales = 0 }; +} + +const Entry = struct { e: ?*const anyopaque = null, kv: EngineKv = .{} }; +var table: [8]Entry = @splat(.{}); + +/// The placement `setup` gave engine `e` (upstream's when none: bf16 on the device, two KV heads). +pub fn of(e: *const anyopaque) EngineKv { + for (&table) |*t| if (t.e == e) return t.kv; + return .{}; +} + +fn entryOf(e: *const anyopaque) !*EngineKv { + for (&table) |*t| if (t.e == e) return &t.kv; + for (&table) |*t| if (t.e == null) { + t.* = .{ .e = e }; + return &t.kv; + }; + return error.OutOfMemory; +} + +/// The engine's deinit: its placement and int8 kernels. +pub fn forget(gpa: std.mem.Allocator, e: *const anyopaque) void { + for (&table) |*t| if (t.e == e) { + release(gpa, &t.kv); + t.* = .{}; + }; +} + +/// prompt.Engine.newSeq under a non-default placement: the same buffers, K/V as codes (+ scales) and/or the indexer +/// keys in one mapped host buffer the 12 layers' DeviceBuffers borrow. +pub fn newSeq(e: anytype, comptime Seq: type, cap: usize) !Seq { + const gpa = e.gpa; + const d = e.driver; + const st = e.stream.*; + const attn = wt.attn_n; + const kv = of(e); + const layout = try sequenceLayout(kv, cap, e.chunk, e.page_config); + const paged = if (layout.paging) |cfg| if (cfg.shared_lanes != 0) try kvpage.Paged.initShared(gpa, d, cap, cfg, e.shared_arena orelse return error.SharedKvArenaRequired) else try kvpage.Paged.init(gpa, d, cap, cfg) else null; + errdefer if (paged) |pg| pg.destroy(); + const logits = try gpa.alloc(u8, wt.vocab * 2); + errdefer gpa.free(logits); + var held: Held = .{ .gpa = gpa }; + errdefer held.free(); + const conv = try held.bufs(d, st, wt.linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + const state = try held.bufs(d, st, wt.linear_n, wt.state_n * 4); + const int8 = kv.k8 != null; + if (int8 and kv.heads != kv_heads) return error.SplitAttentionUnsupported; + const row = layout.row; + const kc = try held.bufs(d, st, attn, layout.cache_rows * row.kv); + const vc = try held.bufs(d, st, attn, layout.cache_rows * row.kv); + var ks: []cuda.DeviceBuffer = &.{}; + var vs: []cuda.DeviceBuffer = &.{}; + if (int8) { + ks = try held.bufs(d, st, attn, cap * row.scales); + vs = try held.bufs(d, st, attn, cap * row.scales); + } + var host: ?cuda.HostBuffer = null; + errdefer if (host) |*h| h.free(); + const ikc_len = cap * index_row; + const ikc = if (kv.ikc_host) blk: { + const stride = std.mem.alignForward(usize, ikc_len, 256); + host = try cuda.HostBuffer.allocMapped(d, stride * attn); + @memset(host.?.bytes, 0); + const base = try host.?.device(); + const out = try gpa.alloc(cuda.DeviceBuffer, attn); + for (out, 0..) |*b, i| b.* = .{ .d = d, .ptr = base + i * stride, .len = ikc_len, .borrowed = true }; + break :blk out; + } else try held.bufs(d, st, attn, ikc_len); + errdefer if (host != null) gpa.free(ikc); + const pooled = try held.bufs(d, st, attn, @intCast(pooledBytes(cap))); + var tail = try cuda.DeviceBuffer.alloc(d, wt.ple_tail_n * wt.wide * 2); + errdefer tail.free(); + try tail.fill8(0, st.handle); + try st.synchronize(); + return .{ + .conv = conv, + .state = state, + .kc = kc, + .vc = vc, + .ikc = ikc, + .pooled = pooled, + .tail = tail, + .history = .{ e.eos, e.eos }, + .fed = 0, + .logits = logits, + .cap = cap, + .kv8 = if (kv.k8) |k| .{ .k = k, .ks = ks, .vs = vs, .rot = kv.rot } else null, + .ikc_host = host, + .paged = paged, + }; +} + +/// prompt.Seq.deinit's part for what newSeq adds (the scales and the indexer keys' host buffer). +pub fn deinitSeq(gpa: std.mem.Allocator, kv8: *?SeqKv8, host: *?cuda.HostBuffer) void { + if (kv8.*) |c| { + for (c.ks) |*b| b.free(); + for (c.vs) |*b| b.free(); + gpa.free(c.ks); + gpa.free(c.vs); + kv8.* = null; + } + if (host.*) |*h| { + h.free(); + host.* = null; + } +} + +/// Zeroed device buffers made so far, freed together on a failure. +const Held = struct { + gpa: std.mem.Allocator, + lists: [8][]cuda.DeviceBuffer = undefined, + n: usize = 0, + + fn bufs(h: *Held, d: *const cuda.Driver, st: cuda.Stream, count: usize, len: usize) ![]cuda.DeviceBuffer { + const out = try h.gpa.alloc(cuda.DeviceBuffer, count); + var made: usize = 0; + errdefer { + for (out[0..made]) |*b| b.free(); + h.gpa.free(out); + } + while (made < count) : (made += 1) { + out[made] = try cuda.DeviceBuffer.alloc(d, len); + out[made].fill8(0, st.handle) catch |err| { + out[made].free(); + return err; + }; + } + h.lists[h.n] = out; + h.n += 1; + return out; + } + + fn free(h: *Held) void { + for (h.lists[0..h.n]) |l| { + for (l) |*b| b.free(); + h.gpa.free(l); + } + h.n = 0; + } +}; + +/// cuda_batch.zig Pass.attnRows over an int8 cache (the hook: `if (self.seq.kv8) |c| return ...`): `n` rows of +/// p.seq from p.pos0, at `at` (positions, queries, indexer queries, outputs). The same projection rows, pooling, +/// block scores and selection as upstream's; prep, chunks and merge ours. A decode round calls it once a lane with that +/// lane's seq, as upstream's attnBlock does, so multi-lane rounds take int8 caches too. +pub fn attnRows(p: anytype, lw: anytype, ai: usize, pa: u64, n: usize, at: anytype, c: SeqKv8) !void { + const s = p.s; + if (p.seq.media != null) return error.Kv8MediaUnsupported; + const k8 = c.k; + const kv = c.cache(p.seq.kc, p.seq.vc, ai); + const ikc = p.seq.ikc[ai].ptr; + const pooled = p.seq.pooled[ai].ptr; + const ratio = wt.index_ratio; + const top = wt.index_top; + const att_rows: usize = 256; + try k8.prepRows(p.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kv, at.iq, ikc, 1e-6, n, wt.proj_n, wt.q_heads, kv_heads, c.rot); + try p.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, 1e-6, n); + const nb = (cdiv(p.seq.cap, ratio) + 15) / 16 * 16; + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = p.pos0 + r0 + m; + if (ends > p.seq.cap) return error.UnexpectedTensor; + const pos = at.pos + b * 16; + const iq = at.iq + r0 * wt.index_heads * wt.index_dim * 2; + if (ends / ratio > top) try p.tri.qsaScores(iq, pooled, pos, p.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(p.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + for ([_]usize{ nb, wt.index_idw, top, ratio }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(p.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, p.stream, &a); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + try k8.attnChunks(p.stream, at.q + r0 * wt.q_heads * head_dim * 2, kv, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, wt.q_heads, kv_heads, chunks, true); + try k8.attnMerge(p.stream, s.po, s.pm, s.pl, pos, at.ao + r0 * wt.out_k * 2, s.qnk, s.qspr, m, wt.q_heads, kv_heads, true, c.rot); + } +} + +// -- host tests ------------------------------------------------------------------------------------------------------- + +test "exl3: bf16 paging and host indexer placement share allocation and budget geometry" { + const cap: usize = 200_016; + const cfg: kvpage.Config = .{}; + for ([_]bool{ false, true }) |host_index| { + const placement: EngineKv = .{ .ikc_host = host_index }; + const resident = try sequenceLayout(placement, cap, 1024, null); + const paged = try sequenceLayout(placement, cap, 4096, cfg); + try std.testing.expectEqual(cap, resident.cache_rows); + try std.testing.expectEqual(@as(usize, 1), paged.cache_rows); + try std.testing.expectEqual(@as(usize, 4096), paged.paging.?.rows); + try std.testing.expectEqual(snapRows(false, kv_heads), paged.row); + const placeholders = wt.attn_n * 2 * paged.row.kv; + try std.testing.expectEqual(try kvpage.deviceBytesFor(cap, cfg.forChunk(4096)) + placeholders, paged.bytes.kv); + try std.testing.expectEqual(resident.bytes.host + resident.bytes.kv, paged.bytes.host); + try std.testing.expectEqual(resident.bytes.device() - resident.bytes.kv + paged.bytes.kv, paged.bytes.device()); + try std.testing.expectEqual(resident.bytes.ikc, paged.bytes.ikc); + try std.testing.expectEqual(@as(u64, if (host_index) 0 else wt.attn_n * cap * index_row), paged.bytes.ikc); + const small = try sequenceLayout(placement, cap, 1024, cfg); + try std.testing.expectEqual(@as(u64, 2 * wt.attn_n * (4096 - cfg.rows) * kvpage.row_bytes), paged.bytes.device() - small.bytes.device()); + } +} + +test "exl3: paging refuses incompatible cache layouts before allocating" { + const cfg: kvpage.Config = .{}; + var k8: Kv8 = undefined; + try std.testing.expectError(error.KvPagePlacement, sequenceLayout(.{ .k8 = &k8 }, 4096, 1024, cfg)); + try std.testing.expectError(error.KvPagePlacement, sequenceLayout(.{ .heads = 1 }, 4096, 1024, cfg)); + try std.testing.expectError(error.KvPageConfig, sequenceLayout(.{}, 4096, 1024, .{ .slots = 0 })); + try std.testing.expectError(error.KvPageConfig, sequenceLayout(.{}, 4096, 1024, .{ .window = wt.dense_keys - 1 })); +} + +test "exl3: newSeq with host indexer keys enters the paging allocator" { + if (cuda.kernels.available) return error.SkipZigTest; + const old_force = kvpage.force; + const old_off = kvpage.force_off; + defer { + kvpage.force = old_force; + kvpage.force_off = old_off; + } + kvpage.force = .{}; + kvpage.force_off = false; + var driver: cuda.Driver = undefined; + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: @import("cuda_prompt.zig").Engine = undefined; + engine.gpa = std.testing.allocator; + engine.driver = &driver; + engine.stream = &stream; + engine.chunk = 4096; + engine.round_graph = .{}; + engine.page_config = kvpage.force; + (try entryOf(&engine)).ikc_host = true; + defer forget(engine.gpa, &engine); + try std.testing.expectError(error.BuiltWithoutKernels, engine.newSeq(8192)); +} + +test "exl3: kv the default policy is upstream's seqBytes, split" { + const prompt = @import("cuda_prompt.zig"); + for ([_]usize{ 1, 5, 4096, 200_016, 262_144 }) |cap| { + const s = seqBytes(cap, .{}); + try std.testing.expectEqual(@as(u64, prompt.seqBytes(cap)), s.device()); + try std.testing.expectEqual(@as(u64, 0), s.host); + } +} + +test "exl3: kv int8 is 34 bytes a 32-group, int4 18; a head split halves K/V, not the indexer" { + try std.testing.expectEqual(@as(usize, 2048), kvRowBytes(.bf16, 2)); + try std.testing.expectEqual(@as(usize, 1088), kvRowBytes(.int8, 2)); + try std.testing.expectEqual(@as(usize, 576), kvRowBytes(.int4, 2)); + try std.testing.expectEqual(@as(usize, 544), kvRowBytes(.int8, 1)); + const cap = 200_000; + const a = seqBytes(cap, .{ .kv = .int8 }); + const b = seqBytes(cap, .{ .kv = .int8, .heads = 1 }); + try std.testing.expectEqual(a.kv, 2 * b.kv); + try std.testing.expectEqual(a.ikc, b.ikc); + const h = seqBytes(cap, .{ .kv = .int8, .kv_host = true, .ikc_host = true }); + try std.testing.expectEqual(a.kv + a.ikc, h.host); + try std.testing.expectEqual(a.device() - a.kv - a.ikc, h.device()); +} + +test "exl3: kv chunk smem fits the 99 KiB opt-in on sm_120" { + try std.testing.expect(chunkSmem(2) <= 99 * 1024); + try std.testing.expectEqual(@as(u32, 60928), chunkSmem(1)); + try std.testing.expectEqual(@as(u32, 78336), chunkSmem(2)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_kvpage.zig b/zig/src/families/flashnext/cuda_exl3_kvpage.zig new file mode 100644 index 0000000000..35477f7bbf --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_kvpage.zig @@ -0,0 +1,1220 @@ +//! Dynamic KV paging for long contexts (docs/plans/exl3-kv-paging.md): a sequence's full K/V of the 12 attention +//! layers in pinned host memory, the card holding a recent dense window and an LRU-like cache of the 4-key blocks +//! the sparse indexer selects. Upstream's attention kernels are unchanged: they read K/V through the rows' key lists, +//! so the paged path rewrites each list into positions of one compact device buffer [window | block slots]. +//! +//! Also the selection trace (TF_FLASHNEXT_KVPAGE_TRACE): every attention row's selected blocks, recorded after +//! fn_qsa_select, for the offline cache replay (tools/exl3_kvpage_replay.py). +//! +//! The hook in the pass is one line at the top of cuda_batch.zig attnRows (`takes` / `attnRows` here). + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); +const telemetry = @import("exl3_telemetry.zig"); +pub const shared_contract = @import("exl3_shared_kv.zig"); +pub const shared_device = @import("cuda_exl3_shared_kv.zig"); + +test "exl3: shared KV host ownership contracts" { + _ = shared_contract; + _ = shared_device; +} + +const head_dim = wt.head_dim; +const kv_heads = wt.kv_heads; +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const att_rows = parts.att_rows; +/// One token's K (or V) in one attention layer: [2][256] bf16. +pub const row_bytes: usize = kv_heads * head_dim * 2; + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +// -- the selection trace -------------------------------------------------------------------------------------------- + +/// Records: u32 layer (attention index), u32 first position, u32 rows, u32 kind (0 prompt, 1 decode round); then a +/// row each: u16 count (0: the row is dense, else 512) and that many u16 block ids (block = key / 4), in block order. +pub const Trace = struct { + gpa: std.mem.Allocator, + out: std.ArrayList(u8) = .empty, + /// Record prompt rows too (else only decode rounds). + prompt: bool = true, + rows: usize = 0, + ids: []i32 = &.{}, + nk: []i32 = &.{}, + sp: []i32 = &.{}, + + pub fn init(gpa: std.mem.Allocator, prompt: bool) !Trace { + return .{ .gpa = gpa, .prompt = prompt, .ids = try gpa.alloc(i32, att_rows * idw), .nk = try gpa.alloc(i32, att_rows), .sp = try gpa.alloc(i32, att_rows) }; + } + + pub fn deinit(t: *Trace) void { + t.out.deinit(t.gpa); + t.gpa.free(t.ids); + t.gpa.free(t.nk); + t.gpa.free(t.sp); + } + + fn word(t: *Trace, v: u32) !void { + try t.out.appendSlice(t.gpa, std.mem.asBytes(&v)); + } + + /// The `m` rows at `pos` the select kernel just listed in s.qids / qnk / qspr (one sync and three small copies). + fn record(t: *Trace, stream: cuda.Stream, ai: usize, pos: usize, m: usize, decode: bool, qids: u64, qnk: u64, qspr: u64) !void { + if (!decode and !t.prompt) return; + try stream.synchronize(); + const d = stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(t.ids.ptr, qids, m * idw * 4), "cuMemcpyDtoH"); + try d.check(d.api.cuMemcpyDtoH_v2(t.nk.ptr, qnk, m * 4), "cuMemcpyDtoH"); + try d.check(d.api.cuMemcpyDtoH_v2(t.sp.ptr, qspr, m * 4), "cuMemcpyDtoH"); + try t.word(@intCast(ai)); + try t.word(@intCast(pos)); + try t.word(@intCast(m)); + try t.word(@intFromBool(decode)); + for (0..m) |r| { + if (t.sp[r] == 0) { + try t.out.appendSlice(t.gpa, &[2]u8{ 0, 0 }); + continue; + } + const cnt: u16 = @intCast(top); + try t.out.appendSlice(t.gpa, std.mem.asBytes(&cnt)); + const row = t.ids[r * idw ..][0 .. top * ratio]; + for (0..top) |j| { + const b: u16 = @intCast(@divTrunc(row[j * ratio], @as(i32, ratio))); + try t.out.appendSlice(t.gpa, std.mem.asBytes(&b)); + } + } + t.rows += m; + } +}; + +/// The active trace, if any (a test sets it; the pass reads it). +pub var trace: ?*Trace = null; + +// -- the pass hook -------------------------------------------------------------------------------------------------- + +/// Whether attnRows runs here: a trace is on, or the sequence's K/V are paged. +pub fn takes(seq: anytype) bool { + return trace != null or seq.paged != null; +} + +/// cuda_batch.zig attnRows (rope, cache writes, pooling, then each 256-row block's scores, selection and attention) +/// with the trace after each selection, or with the sequence's paged K/V. `at` is the pass's RowsAt. +pub fn attnRows(p: anytype, lw: anytype, ai: usize, pa: u64, n: usize, at: anytype) !void { + if (p.seq.media != null) return error.KvPageMedia; + if (p.seq.paged) |pg| return pg.attnRows(p, lw, ai, pa, n, at); + const s = p.s; + const kc = p.seq.kc[ai].ptr; + const vc = p.seq.vc[ai].ptr; + const ikc = p.seq.ikc[ai].ptr; + const pooled = p.seq.pooled[ai].ptr; + try p.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, wt.eps, n, 0); + try p.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, wt.eps, n); + const nb = parts.scoreStride(p.seq.cap); + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = p.pos0 + r0 + m; + if (ends > p.seq.cap) return error.UnexpectedTensor; + const pos = at.pos + b * 16; + try select(p, at.iq + r0 * wt.index_heads * wt.index_dim * 2, pooled, pos, nb, m, ends); + if (trace) |t| try t.record(p.stream, ai, p.pos0 + r0, m, p.lanes.len > 0, s.qids, s.qnk, s.qspr); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + try p.tri.attnChunks(at.q + r0 * wt.q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true); + try p.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * wt.out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + } +} + +/// Block scores (rows past `top` complete blocks) and fn_qsa_select into s.qids / qnk / qspr, as upstream. +pub fn select(p: anytype, iq: u64, pooled: u64, pos: u64, nb: usize, m: usize, ends: usize) !void { + const s = p.s; + if (ends / ratio > top) try p.tri.qsaScores(iq, pooled, pos, p.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(p.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + for ([_]usize{ nb, idw, top, ratio }) |v| a.add(cint(v)); + try cuda.launch.launch(p.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, p.stream, &a); +} + +// -- the paged store ------------------------------------------------------------------------------------------------ + +/// Shape of a paged sequence's device side (TF_FLASHNEXT_KVPAGE=1 turns it on for every new sequence; the sizes +/// from TF_FLASHNEXT_KVPAGE_WINDOW / _SLOTS / _ROWS). +pub const Config = struct { + /// Tokens kept dense on the card before a call's first row (>= 2051, so a dense row's keys are all in it). + window: usize = 4096, + /// 4-key block slots, one pool for the 12 layers (the prompt's layer-major order gives one layer all of it). + slots: usize = 49152, + /// The most rows one attention call writes (the engine's chunk / a deep prefill's sub-block). + rows: usize = 1536, + shared_lanes: usize = 0, + + pub fn forChunk(cfg: Config, chunk: usize) Config { + var out = cfg; + out.rows = @max(out.rows, chunk); + return out; + } + + pub fn validate(cfg: Config) !void { + if (cfg.window < wt.dense_keys or cfg.slots == 0 or cfg.rows == 0) return error.KvPageConfig; + if (cfg.shared_lanes > shared_device.max_lanes) return error.KvPageConfig; + } +}; + +/// A test's override of the environment (null: read it). +pub var force: ?Config = null; +pub var force_off: bool = false; + +fn envUsize(name: [*:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} + +pub fn configFromEnv() ?Config { + if (force_off) return null; + if (force) |c| return c; + const v = std.c.getenv("TF_FLASHNEXT_KVPAGE") orelse return null; + if (!std.mem.eql(u8, std.mem.span(v), "1")) return null; + const d: Config = .{}; + var rows = envUsize("TF_FLASHNEXT_KVPAGE_ROWS", d.rows); + if (std.mem.eql(u8, std.mem.span(std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL") orelse "1"), "1")) { + rows = @max(rows, envUsize("TF_FLASHNEXT_EXL3_PREFILL_ROWS", 4096)); + } + const shared = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_FLASHNEXT_KVPAGE_SHARED") orelse "0"), "1"); + const shared_lanes = if (shared) envUsize("TF_FLASHNEXT_KVPAGE_SHARED_LANES", 8) else 0; + return .{ .window = envUsize("TF_FLASHNEXT_KVPAGE_WINDOW", d.window), .slots = envUsize("TF_FLASHNEXT_KVPAGE_SLOTS", d.slots), .rows = rows, .shared_lanes = if (shared and shared_lanes == 0) std.math.maxInt(usize) else shared_lanes }; +} + +pub fn configIdentity(cfg: ?Config) u64 { + const paging = cfg orelse return 0; + const shape = [_]u64{ paging.window, paging.rows, paging.slots, paging.shared_lanes, row_bytes, attn_n }; + return std.hash.Wyhash.hash(0, std.mem.asBytes(&shape)); +} + +pub fn admissionRoom(local: u64, stream_bytes: usize, cfg: ?Config) !u64 { + const paging = cfg orelse return local; + try paging.validate(); + if (paging.shared_lanes == 0) return local; + return @min(local, try std.math.mul(u64, stream_bytes, paging.shared_lanes)); +} + +const Kern = struct { + module: cuda.Module, + plan: cuda.Function, + fetch: cuda.Function, + remap: cuda.Function, + + fn load(d: *const cuda.Driver) !Kern { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kern = undefined; + k.module = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_kvpage); + errdefer k.module.unload(); + k.plan = try k.module.function("kvp_plan"); + k.fetch = try k.module.function("kvp_fetch"); + k.remap = try k.module.function("kvp_remap"); + return k; + } +}; + +const attn_n = wt.attn_n; +const ctl_n: usize = 8; +const work_n: usize = att_rows * top; + +pub const SnapshotLayout = struct { + cap: usize, + cfg: Config, + at: usize, + wr: usize, + nblocks: usize, + canonical: usize, + device: usize, + tables: usize, + total: usize, + + pub fn init(cap: usize, cfg: Config, at: usize) !SnapshotLayout { + try cfg.validate(); + if (cap == 0 or at > cap) return error.KvPageRows; + _ = try std.math.mul(usize, cap, attn_n * 2 * row_bytes); + if (cfg.shared_lanes != 0) _ = try shared_device.Layout.init(.{ .cap = cap, .lanes = cfg.shared_lanes, .window = cfg.window, .rows = cfg.rows, .slots = cfg.slots }); + const wr = try std.math.add(usize, try std.math.add(usize, cfg.window, cfg.rows), ratio); + const nblocks = try std.math.add(usize, (try std.math.add(usize, cap, ratio - 1)) / ratio, 1); + const canonical = try std.math.mul(usize, at, attn_n * 2 * row_bytes); + const rows = try std.math.add(usize, try std.math.mul(usize, attn_n, wr), if (cfg.shared_lanes == 0) try std.math.mul(usize, ratio, cfg.slots) else 0); + const device = try std.math.mul(usize, rows, 2 * row_bytes); + const words = if (cfg.shared_lanes == 0) try std.math.add(usize, try std.math.mul(usize, attn_n, try std.math.add(usize, try std.math.mul(usize, 2, nblocks), ctl_n)), try std.math.mul(usize, 2, cfg.slots)) else try std.math.mul(usize, attn_n, try std.math.add(usize, nblocks, ctl_n)); + const tables = try std.math.mul(usize, words, 4); + return .{ .cap = cap, .cfg = cfg, .at = at, .wr = wr, .nblocks = nblocks, .canonical = canonical, .device = device, .tables = tables, .total = try std.math.add(usize, canonical, try std.math.add(usize, device, tables)) }; + } +}; + +pub const Snapshot = struct { + layout: SnapshotLayout, + win0: [attn_n]usize, + rebases: [attn_n]u64, + written: [attn_n]u64, + + pub fn validate(image: Snapshot) !void { + const expected = try SnapshotLayout.init(image.layout.cap, image.layout.cfg, image.layout.at); + if (!std.meta.eql(expected, image.layout)) return error.KvPageSnapshotLayout; + for (image.win0) |origin| { + if (origin % ratio != 0 or origin > image.layout.at or image.layout.at - origin > image.layout.wr) return error.KvPageWindow; + } + } +}; + +/// Counters of one layer (kvp_plan's ctl) plus the host's window moves. +pub const Stats = struct { + misses: u64 = 0, + needed: u64 = 0, + calls: u64 = 0, + overflow: bool = false, + rebases: u64 = 0, + written_rows: u64 = 0, + epoch: u32 = 0, + + pub fn add(a: *Stats, b: Stats) void { + a.misses += b.misses; + a.needed += b.needed; + a.calls += b.calls; + a.overflow = a.overflow or b.overflow; + a.rebases += b.rebases; + a.written_rows += b.written_rows; + a.epoch = @max(a.epoch, b.epoch); + } +}; + +/// One sequence's K/V of the 12 attention layers: all of it in mapped pinned host memory, [layer][K|V][cap][2][256]; +/// on the card, for K and for V, [12 windows of wr rows | 4 * slots rows] [2][256] bf16 and the block tables. +pub const Paged = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + k: Kern, + cfg: Config, + cap: usize, + /// Window rows: the kept window, one call's rows, and a block's slack. + wr: usize, + nblocks: usize, + host: cuda.HostBuffer, + host_dev: u64, + status_host: cuda.HostBuffer, + dev: cuda.DeviceBuffer, + /// blk2slot [12][nblocks] i32, mark [12][nblocks] u32, ctl [12][8] i32, slot2blk [slots] i32, stamp [slots] u32. + tabs: cuda.DeviceBuffer, + /// need, miss [att_rows * 512] i32 and fetch [2 * att_rows * 512] i32 (one call's, shared by the layers). + work: cuda.DeviceBuffer, + win0: [attn_n]usize = @splat(0), + rebases: [attn_n]u64 = @splat(0), + written: [attn_n]u64 = @splat(0), + graph_prepared: bool = false, + telemetry_id: u64 = 0, + telemetry_revision: u64 = 0, + shared_arena: ?*shared_device.Arena = null, + shared_sequence: shared_contract.Sequence = undefined, + shared_attached: bool = true, + incarnation: u64 = 0, + retired_for_reuse: bool = false, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize, cfg: Config) !*Paged { + return initStore(gpa, d, cap, cfg, null); + } + + pub fn initShared(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize, cfg: Config, arena: *shared_device.Arena) !*Paged { + return initStore(gpa, d, cap, cfg, arena); + } + + fn initStore(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize, cfg: Config, shared: ?*shared_device.Arena) !*Paged { + const layout = try SnapshotLayout.init(cap, cfg, 0); + if ((cfg.shared_lanes != 0) != (shared != null)) return error.SharedKvArenaRequired; + if (shared) |arena| { + const options = arena.layout.options; + if (cap != options.cap) return error.SharedKvCommonCap; + if (d != arena.driver or cfg.shared_lanes != options.lanes or cfg.window != options.window or cfg.rows != options.rows or cfg.slots != options.slots) return error.SharedKvGeometry; + } + const pg = try gpa.create(Paged); + errdefer gpa.destroy(pg); + pg.* = .{ .gpa = gpa, .d = d, .k = undefined, .cfg = cfg, .cap = cap, .wr = layout.wr, .nblocks = layout.nblocks, .host = undefined, .host_dev = 0, .status_host = undefined, .dev = undefined, .tabs = undefined, .work = undefined }; + pg.telemetry_id = try telemetry.identity(); + pg.incarnation = try shared_contract.identity(); + pg.shared_arena = shared; + if (shared) |arena| pg.shared_sequence = try arena.register(cap) else pg.k = try Kern.load(d); + errdefer if (shared) |arena| arena.release(pg.shared_sequence, .{ .d = d, .handle = null }) catch unreachable else pg.k.module.unload(); + pg.host = try cuda.HostBuffer.allocMapped(d, attn_n * 2 * cap * row_bytes); + errdefer pg.host.free(); + pg.host_dev = try pg.host.device(); + pg.status_host = try cuda.HostBuffer.alloc(d, attn_n * ctl_n * 4); + errdefer pg.status_host.free(); + pg.dev = if (shared) |arena| arena.device else try cuda.DeviceBuffer.alloc(d, layout.device); + if (shared != null) pg.dev.borrowed = true; + errdefer pg.dev.free(); + pg.tabs = if (shared) |arena| arena.tables else try cuda.DeviceBuffer.alloc(d, layout.tables); + if (shared != null) pg.tabs.borrowed = true; + errdefer pg.tabs.free(); + pg.work = if (shared) |arena| arena.work else try cuda.DeviceBuffer.alloc(d, work_n * 4 * 4); + if (shared != null) pg.work.borrowed = true; + errdefer pg.work.free(); + if (shared == null) { + try pg.dev.fill8(0, null); + try pg.tabs.fill8(0, null); + try pg.resetTables(null); + } + return pg; + } + + pub fn destroy(pg: *Paged) void { + if (pg.cfg.shared_lanes != 0 and pg.shared_attached) pg.shared_arena.?.release(pg.shared_sequence, .{ .d = pg.d, .handle = null }) catch @panic("shared KV lane retirement failed"); + pg.work.free(); + pg.tabs.free(); + pg.dev.free(); + pg.status_host.free(); + pg.host.free(); + if (pg.cfg.shared_lanes == 0) pg.k.module.unload(); + pg.gpa.destroy(pg); + } + + pub fn retireForReuse(pg: *Paged) !void { + if (pg.retired_for_reuse) return error.KvPageRetired; + if (pg.shared_arena) |arena| if (arena.active != null) return error.SharedKvBusy; + try pg.d.synchronize(); + if (pg.shared_arena) |arena| { + try arena.release(pg.shared_sequence, .{ .d = pg.d, .handle = null }); + pg.shared_attached = false; + } + pg.graph_prepared = false; + pg.retired_for_reuse = true; + } + + pub fn resetForReuse(pg: *Paged, stream: cuda.Stream) !void { + if (!pg.retired_for_reuse) return error.KvPageNotRetired; + if (stream.d != pg.d) return error.SharedKvDriver; + const identity = try telemetry.identity(); + const incarnation = try shared_contract.identity(); + if (pg.shared_arena) |arena| { + if (pg.shared_attached) return error.SharedKvSequenceLive; + pg.shared_sequence = try arena.register(pg.cap); + pg.shared_attached = true; + } else { + try pg.dev.fill8(0, stream.handle); + try pg.work.fill8(0, stream.handle); + try pg.resetTables(stream); + } + @memset(pg.host.bytes, 0); + @memset(pg.status_host.bytes, 0); + pg.win0 = @splat(0); + pg.rebases = @splat(0); + pg.written = @splat(0); + pg.graph_prepared = false; + pg.telemetry_id = identity; + pg.telemetry_revision = 0; + pg.incarnation = incarnation; + pg.retired_for_reuse = false; + } + + /// Device bytes of the K/V side (windows and slots). + pub fn devBytes(pg: *const Paged) usize { + if (pg.cfg.shared_lanes != 0) return 0; + return 2 * pg.rowsDev() * row_bytes; + } + + /// Device bytes in all: K/V side, tables and the call's work lists (deviceBytesFor, a host test checks). + pub fn deviceBytes(pg: *const Paged) usize { + if (pg.cfg.shared_lanes != 0) return 0; + return pg.devBytes() + pg.tabWords() * 4 + work_n * 4 * 4; + } + + /// Rows of the K (or V) buffer: the 12 windows, then the slots. + fn rowsDev(pg: *const Paged) usize { + if (pg.cfg.shared_lanes != 0) return pg.shared_arena.?.layout.side_rows; + return attn_n * pg.wr + ratio * pg.cfg.slots; + } + + fn tabWords(pg: *const Paged) usize { + return attn_n * (2 * pg.nblocks + ctl_n) + 2 * pg.cfg.slots; + } + + /// Layer `ai`'s window base in the K buffer (its kernels index from here). + fn devK(pg: *const Paged, ai: usize) u64 { + if (pg.cfg.shared_lanes != 0) return pg.dev.ptr + (pg.shared_arena.?.layout.windowOffset(pg.shared_sequence.lane, ai) catch unreachable) * row_bytes; + return pg.dev.ptr + ai * pg.wr * row_bytes; + } + + fn devV(pg: *const Paged, ai: usize) u64 { + if (pg.cfg.shared_lanes != 0) return pg.devK(ai) + pg.rowsDev() * row_bytes; + return pg.dev.ptr + (pg.rowsDev() + ai * pg.wr) * row_bytes; + } + + /// Rows from layer `ai`'s window base to the first slot. + fn slotBase(pg: *const Paged, ai: usize) usize { + if (pg.cfg.shared_lanes != 0) return pg.shared_arena.?.layout.slotBase(pg.shared_sequence.lane, ai) catch unreachable; + return (attn_n - ai) * pg.wr; + } + + fn hostK(pg: *const Paged, ai: usize) u64 { + return pg.host_dev + (2 * ai) * pg.cap * row_bytes; + } + + fn hostV(pg: *const Paged, ai: usize) u64 { + return pg.host_dev + (2 * ai + 1) * pg.cap * row_bytes; + } + + const Tab = struct { blk2slot: u64, mark: u64, ctl: u64, slot2blk: u64, stamp: u64, blk2slot_all: u64 }; + + fn tab(pg: *const Paged, ai: usize) Tab { + if (pg.cfg.shared_lanes != 0) { + const arena = pg.shared_arena.?; + const layout = arena.layout; + const key = @as(usize, pg.shared_sequence.lane) * layout.lane_words + ai * layout.blocks; + return .{ .blk2slot = arena.tables.ptr + key * 4, .mark = arena.tables.ptr + layout.marks + key * 4, .ctl = arena.tables.ptr + layout.controls + (@as(usize, pg.shared_sequence.lane) * attn_n + ai) * ctl_n * 4, .slot2blk = arena.tables.ptr + layout.owners, .stamp = arena.tables.ptr + layout.stamps, .blk2slot_all = arena.tables.ptr }; + } + const base = pg.tabs.ptr; + const nb = pg.nblocks * 4; + const all = base; + const marks = base + attn_n * nb; + const ctls = marks + attn_n * nb; + const s2b = ctls + attn_n * ctl_n * 4; + return .{ .blk2slot = all + ai * nb, .mark = marks + ai * nb, .ctl = ctls + ai * ctl_n * 4, .slot2blk = s2b, .stamp = s2b + pg.cfg.slots * 4, .blk2slot_all = all }; + } + + fn memset32(pg: *const Paged, ptr: u64, value: u32, words: usize, stream: ?cuda.Stream) !void { + if (stream) |s| return pg.d.check(pg.d.api.cuMemsetD32Async(ptr, value, words, s.handle), "cuMemsetD32Async"); + try pg.d.check(pg.d.api.cuMemsetD32_v2(ptr, value, words), "cuMemsetD32"); + } + + /// An empty cache with a fresh device clock and every window at 0. + fn resetTables(pg: *Paged, stream: ?cuda.Stream) !void { + pg.advanceTelemetry(); + if (pg.cfg.shared_lanes != 0) { + const arena = pg.shared_arena.?; + try arena.clear(try arena.leaseFor(pg.shared_sequence), pg.shared_sequence, stream orelse .{ .d = pg.d, .handle = null }); + pg.win0 = @splat(0); + return; + } + const t = pg.tab(0); + try pg.memset32(t.blk2slot_all, 0xFFFF_FFFF, attn_n * pg.nblocks, stream); + try pg.memset32(t.slot2blk, 0xFFFF_FFFF, pg.cfg.slots, stream); + try pg.memset32(t.stamp, 0, pg.cfg.slots, stream); + try pg.memset32(pg.tabs.ptr + attn_n * pg.nblocks * 4, 0, attn_n * pg.nblocks, stream); + try pg.memset32(t.ctl, 0, attn_n * ctl_n, stream); + pg.win0 = @splat(0); + } + + fn copy(pg: *const Paged, stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { + if (n == 0) return; + try pg.d.check(pg.d.api.cuMemcpyDtoDAsync_v2(dst, src, n, stream.handle), "cuMemcpyDtoDAsync"); + } + + /// Moves the window so rows [pos0, pos0 + n) fit: the rows that leave it are written out to host memory, the + /// last `window` tokens before pos0 move to the window's start (pieces no longer than the shift: no overlap). + fn rebase(pg: *Paged, ai: usize, pos0: usize, stream: cuda.Stream) !void { + const w0 = pg.win0[ai]; + const new0 = (pos0 -| pg.cfg.window) / ratio * ratio; + if (new0 <= w0) return error.KvPageWindow; + const shift = new0 - w0; + if (pos0 < w0 or pos0 > w0 + pg.wr or shift > pg.wr) return error.KvPageWindow; + for ([_][2]u64{ .{ pg.devK(ai), pg.hostK(ai) }, .{ pg.devV(ai), pg.hostV(ai) } }) |dh| { + try pg.copy(stream, dh[1] + w0 * row_bytes, dh[0], shift * row_bytes); + const len = pos0 - new0; + var at: usize = 0; + while (at < len) { + const piece = @min(shift, len - at); + try pg.copy(stream, dh[0] + at * row_bytes, dh[0] + (shift + at) * row_bytes, piece * row_bytes); + at += piece; + } + } + pg.win0[ai] = new0; + pg.rebases[ai] += 1; + pg.written[ai] += shift; + } + + /// Every row the sequence holds written out to host memory (the window's rows too): host K/V then equal what an + /// all-on-device cache would hold, for a check or a snapshot. + pub fn flush(pg: *Paged, stream: cuda.Stream, fed: usize) !void { + _ = try pg.snapshot(fed); + for (0..attn_n) |ai| try pg.flushLayer(stream, ai, fed); + } + + pub fn snapshot(pg: *const Paged, at: usize) !Snapshot { + try pg.validateSharedGeometry(); + const image: Snapshot = .{ .layout = try SnapshotLayout.init(pg.cap, pg.cfg, at), .win0 = pg.win0, .rebases = pg.rebases, .written = pg.written }; + try image.validate(); + if (pg.wr != image.layout.wr or pg.nblocks != image.layout.nblocks or (pg.cfg.shared_lanes == 0 and (pg.dev.len != image.layout.device or pg.tabs.len != image.layout.tables)) or pg.host.bytes.len != attn_n * 2 * pg.cap * row_bytes) return error.KvPageSnapshotLayout; + return image; + } + + pub fn validateSnapshot(pg: *const Paged, image: Snapshot) !void { + try pg.validateSharedGeometry(); + try image.validate(); + const expected = try SnapshotLayout.init(pg.cap, pg.cfg, image.layout.at); + if (!std.meta.eql(expected, image.layout) or pg.wr != expected.wr or pg.nblocks != expected.nblocks or (pg.cfg.shared_lanes == 0 and (pg.dev.len != expected.device or pg.tabs.len != expected.tables)) or pg.host.bytes.len != attn_n * 2 * pg.cap * row_bytes) return error.KvPageSnapshotLayout; + } + + fn validateSharedGeometry(pg: *const Paged) !void { + if (pg.cfg.shared_lanes == 0) return; + if (!pg.shared_attached) return error.SharedKvSequence; + const arena = pg.shared_arena orelse return error.KvPageSnapshotLayout; + const expected = try shared_device.Layout.init(.{ .cap = pg.cap, .lanes = pg.cfg.shared_lanes, .window = pg.cfg.window, .rows = pg.cfg.rows, .slots = pg.cfg.slots }); + if (!std.meta.eql(expected, arena.layout) or pg.dev.ptr != arena.device.ptr or pg.dev.len != arena.device.len or pg.tabs.ptr != arena.tables.ptr or pg.tabs.len != arena.tables.len or pg.work.ptr != arena.work.ptr or pg.work.len != arena.work.len) return error.KvPageSnapshotLayout; + } + + pub fn moveSnapshot(pg: *Paged, stream: cuda.Stream, base: u64, image: Snapshot, saving: bool) !void { + var offset: usize = 0; + const prefix = image.layout.at * row_bytes; + for (0..attn_n) |ai| for ([_]u64{ pg.hostK(ai), pg.hostV(ai) }) |address| { + if (saving) try pg.copy(stream, base + offset, address, prefix) else try pg.copy(stream, address, base + offset, prefix); + offset += prefix; + }; + if (pg.cfg.shared_lanes != 0) { + for (0..attn_n) |ai| for ([_]u64{ pg.devK(ai), pg.devV(ai) }) |address| { + const size = pg.wr * row_bytes; + if (saving) try pg.copy(stream, base + offset, address, size) else try pg.copy(stream, address, base + offset, size); + offset += size; + }; + const tables = pg.tab(0); + for ([_][2]u64{ .{ tables.mark, attn_n * pg.nblocks * 4 }, .{ tables.ctl, attn_n * ctl_n * 4 } }) |section| { + if (saving) try pg.copy(stream, base + offset, section[0], @intCast(section[1])) else try pg.copy(stream, section[0], base + offset, @intCast(section[1])); + offset += @intCast(section[1]); + } + std.debug.assert(offset == image.layout.total); + return; + } + for ([_]cuda.DeviceBuffer{ pg.dev, pg.tabs }) |buffer| { + if (saving) try pg.copy(stream, base + offset, buffer.ptr, buffer.len) else try pg.copy(stream, buffer.ptr, base + offset, buffer.len); + offset += buffer.len; + } + std.debug.assert(offset == image.layout.total); + } + + pub fn publishSnapshot(pg: *Paged, image: Snapshot) void { + pg.advanceTelemetry(); + pg.win0 = image.win0; + pg.rebases = image.rebases; + pg.written = image.written; + pg.graph_prepared = false; + } + + fn advanceTelemetry(pg: *Paged) void { + if (pg.telemetry_revision == std.math.maxInt(u64)) pg.telemetry_id = 0 else pg.telemetry_revision += 1; + } + + pub fn flushLayer(pg: *Paged, stream: cuda.Stream, ai: usize, fed: usize) !void { + if (fed > pg.cap) return error.KvPageRows; + const w0 = pg.win0[ai]; + if (fed <= w0) return; + if (fed - w0 > pg.wr) return error.KvPageWindow; + try pg.copy(stream, pg.hostK(ai) + w0 * row_bytes, pg.devK(ai), (fed - w0) * row_bytes); + try pg.copy(stream, pg.hostV(ai) + w0 * row_bytes, pg.devV(ai), (fed - w0) * row_bytes); + } + + /// The host copy of layer `ai`'s K (v false) or V rows (after `flush`). + pub fn hostRows(pg: *const Paged, ai: usize, v: bool) []const u8 { + const off = (2 * ai + @intFromBool(v)) * pg.cap * row_bytes; + return pg.host.bytes[off..][0 .. pg.cap * row_bytes]; + } + + pub fn stats(pg: *const Paged, ai: usize) !Stats { + var c: [ctl_n]i32 = undefined; + try pg.d.check(pg.d.api.cuMemcpyDtoH_v2(&c, pg.tab(ai).ctl, ctl_n * 4), "cuMemcpyDtoH"); + return .{ .misses = @as(u32, @bitCast(c[3])), .needed = @as(u32, @bitCast(c[4])), .calls = @as(u32, @bitCast(c[5])), .overflow = c[2] != 0, .rebases = pg.rebases[ai], .written_rows = pg.written[ai], .epoch = @bitCast(c[6]) }; + } + + pub fn sampleClockAfterConsumers(pg: *const Paged) !u32 { + const address = if (pg.shared_arena) |arena| arena.tables.ptr + arena.layout.clock else pg.tab(0).ctl; + var clock: u32 = undefined; + try pg.d.check(pg.d.api.cuMemcpyDtoH_v2(&clock, address, @sizeOf(u32)), "cuMemcpyDtoH(KV telemetry clock)"); + return clock; + } + + pub fn total(pg: *const Paged) !Stats { + var t: Stats = .{}; + for (0..attn_n) |ai| t.add(try pg.stats(ai)); + return t; + } + + pub fn stageStatus(pg: *Paged, stream: cuda.Stream) !void { + const ctl = pg.tab(0).ctl; + try pg.d.check(pg.d.api.cuMemcpyDtoHAsync_v2(pg.status_host.bytes.ptr, ctl, attn_n * ctl_n * 4, stream.handle), "cuMemcpyDtoHAsync"); + } + + pub fn checkStatus(pg: *const Paged) !void { + const controls: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, pg.status_host.bytes[0 .. attn_n * ctl_n * 4])); + for (0..attn_n) |ai| if (controls[ai * ctl_n + 2] != 0) return error.KvPageOverflow; + } + + /// One call's paging after fn_qsa_select listed `m` rows' keys in qids / qnk / qspr: the plan (hits stamped, the + /// misses given the least recently used slots), the misses' rows fetched from host memory, the lists remapped. + pub fn page(pg: *Paged, stream: cuda.Stream, ai: usize, qids: u64, qnk: u64, qspr: u64, m: usize) !void { + if (ai >= attn_n) return error.KvPageRows; + if (try pagingRows(pg.cfg, pg.win0[ai], m) != m) return error.KvPageCapacity; + if (pg.cfg.shared_lanes != 0) { + const arena = pg.shared_arena.?; + return arena.pageFor(try arena.leaseFor(pg.shared_sequence), pg.shared_sequence, stream, ai, pg.win0[ai], m, qids, qnk, qspr, pg.hostK(ai), pg.hostV(ai)); + } + const t = pg.tab(ai); + const w0 = pg.win0[ai]; + const need = pg.work.ptr; + const miss = need + work_n * 4; + const fetch = miss + work_n * 4; + var a: cuda.Args = .{}; + for ([_]u64{ qids, qnk, qspr }) |v| a.add(v); + for ([_]usize{ m, idw, top, ratio, w0 }) |v| a.add(cint(v)); + a.add(pg.tab(0).ctl); + a.add(t.blk2slot); + a.add(t.blk2slot_all); + a.add(cint(ai * pg.nblocks)); + for ([_]u64{ t.slot2blk, t.stamp, t.mark }) |v| a.add(v); + a.add(cint(pg.cfg.slots)); + for ([_]u64{ t.ctl, need, miss, fetch }) |v| a.add(v); + try cuda.launch.launch(pg.k.plan, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream, &a); + var f: cuda.Args = .{}; + for ([_]u64{ pg.hostK(ai), pg.hostV(ai), pg.devK(ai), pg.devV(ai), t.ctl, fetch }) |v| f.add(v); + for ([_]usize{ pg.slotBase(ai), row_bytes / 16, ratio }) |v| f.add(cint(v)); + try cuda.launch.launch(pg.k.fetch, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, stream, &f); + var r: cuda.Args = .{}; + for ([_]u64{ qids, qnk, qspr }) |v| r.add(v); + for ([_]usize{ idw, w0, pg.slotBase(ai) }) |v| r.add(cint(v)); + r.add(t.blk2slot); + r.add(cint(ratio)); + r.add(t.ctl); + try cuda.launch.launch(pg.k.remap, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 256 } }, stream, &r); + } + + pub fn validateRange(pg: *const Paged, ai: usize, pos0: usize, n: usize) !void { + if (ai >= attn_n or n == 0 or n > pg.cfg.rows or pos0 > pg.cap or n > pg.cap - pos0) return error.KvPageRows; + _ = try windowPlan(pg.cfg, pg.wr, pg.win0[ai], pos0, n); + } + + pub fn prepareLayer(pg: *Paged, stream: cuda.Stream, ai: usize, pos0: usize, n: usize) !void { + if (pg.cfg.shared_lanes != 0) { + const arena = pg.shared_arena.?; + try arena.bindStream(try arena.leaseFor(pg.shared_sequence), stream); + } + try pg.validateRange(ai, pos0, n); + if (pos0 == 0 and ai == 0) try pg.resetTables(stream); + const w0 = try windowPlan(pg.cfg, pg.wr, pg.win0[ai], pos0, n); + if (w0 != pg.win0[ai]) try pg.rebase(ai, pos0, stream); + } + + pub fn prepareGraph(pg: *Paged, stream: cuda.Stream, pos0: usize, n: usize) !void { + for (0..attn_n) |ai| try pg.validateRange(ai, pos0, n); + for (0..attn_n) |ai| try pg.prepareLayer(stream, ai, pos0, n); + pg.graph_prepared = true; + } + + pub fn graphKey(pg: *const Paged) [6]u64 { + var shape: [9 + attn_n]usize = undefined; + shape[0..5].* = .{ pg.cap, pg.wr, pg.cfg.window, pg.cfg.rows, pg.cfg.slots }; + shape[5..8].* = if (pg.cfg.shared_lanes == 0) .{ 0, 0, @intCast(pg.incarnation) } else .{ pg.cfg.shared_lanes, pg.shared_sequence.lane, @intCast(pg.shared_sequence.generation) }; + shape[8] = if (pg.cfg.shared_lanes == 0) 0 else @intCast(pg.shared_arena.?.id); + for (pg.win0, 0..) |origin, ai| shape[9 + ai] = pagingRows(pg.cfg, origin, att_rows) catch 0; + return .{ pg.dev.ptr, pg.tabs.ptr, pg.work.ptr, pg.host_dev, std.hash.Wyhash.hash(0, std.mem.asBytes(&shape)), std.hash.Wyhash.hash(0, std.mem.asBytes(&pg.win0)) }; + } + + pub fn prefillGraphKey(pg: *const Paged) [6]u64 { + const strict = pg.graphKey(); + const logical = pg.logicalGraphKey(); + const lane: u64 = if (pg.cfg.shared_lanes == 0) 0 else pg.shared_sequence.lane; + const arena: u64 = if (pg.cfg.shared_lanes == 0) 0 else pg.shared_arena.?.id; + const shape = .{ logical[0], lane, arena }; + return .{ strict[0], strict[1], strict[2], strict[3], std.hash.Wyhash.hash(0, std.mem.asBytes(&shape)), std.hash.Wyhash.hash(strict[4], std.mem.asBytes(&strict[5])) }; + } + + pub fn logicalGraphKey(pg: *const Paged) [2]u64 { + var shape: [6 + attn_n]usize = undefined; + shape[0..6].* = .{ pg.cap, pg.wr, pg.cfg.window, pg.cfg.rows, pg.cfg.slots, pg.cfg.shared_lanes }; + for (pg.win0, 0..) |origin, layer| shape[6 + layer] = pagingRows(pg.cfg, origin, att_rows) catch 0; + return .{ std.hash.Wyhash.hash(0, std.mem.asBytes(&shape)), std.hash.Wyhash.hash(0, std.mem.asBytes(&pg.win0)) }; + } + + pub fn validateSharedLease(pg: *const Paged) !void { + if (pg.cfg.shared_lanes != 0) _ = try pg.shared_arena.?.leaseFor(pg.shared_sequence); + } + + pub fn validateSharedLeaseOn(pg: *const Paged, stream: cuda.Stream) !void { + if (pg.cfg.shared_lanes != 0) { + const arena = pg.shared_arena.?; + try arena.bindStream(try arena.leaseFor(pg.shared_sequence), stream); + } + } + + /// Test support: the window's start of layer `ai` set by hand (the rows before it then come from host memory). + pub fn setWindow(pg: *Paged, ai: usize, w0: usize) void { + pg.win0[ai] = w0; + } + + /// Test support: the device K (v false) or V buffer from layer `ai`'s window base: rows [0, slotRows(ai)). + pub fn devRows(pg: *const Paged, ai: usize, v: bool) u64 { + return if (v) pg.devV(ai) else pg.devK(ai); + } + + /// Test support: rows from layer `ai`'s window base to the end of the slots, and the first slot's row. + pub fn slotRows(pg: *const Paged, ai: usize) [2]usize { + return .{ pg.slotBase(ai) + ratio * pg.cfg.slots, pg.slotBase(ai) }; + } + + /// Test support: the host K/V of layer `ai` to write directly. + pub fn hostMut(pg: *Paged, ai: usize, v: bool) []u8 { + const off = (2 * ai + @intFromBool(v)) * pg.cap * row_bytes; + return pg.host.bytes[off..][0 .. pg.cap * row_bytes]; + } + + /// attnRows over the paged cache: prep writes the rows into the window (its base offset by -win0 rows), and + /// each 256-row block's key lists go through kvp_plan / kvp_fetch / kvp_remap before upstream's chunks and merge. + pub fn attnRows(pg: *Paged, p: anytype, lw: anytype, ai: usize, pa: u64, n: usize, at: anytype) !void { + try pg.validateSharedLeaseOn(p.stream); + const s = p.s; + const pos0 = p.pos0; + if (pg.graph_prepared) { + try pg.validateRange(ai, pos0, n); + if (pos0 + n > pg.win0[ai] + pg.wr) return error.KvPageWindow; + } else try pg.prepareLayer(p.stream, ai, pos0, n); + const w0 = pg.win0[ai]; + const kd = pg.devK(ai); + const vd = pg.devV(ai); + const ikc = p.seq.ikc[ai].ptr; + const pooled = p.seq.pooled[ai].ptr; + try p.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kd -% w0 * row_bytes, vd -% w0 * row_bytes, at.iq, ikc, s.dummy, s.dummy, wt.eps, n, 0); + try p.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, wt.eps, n); + const nb = parts.scoreStride(p.seq.cap); + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = pos0 + r0 + m; + const pos = at.pos + b * 16; + try select(p, at.iq + r0 * wt.index_heads * wt.index_dim * 2, pooled, pos, nb, m, ends); + if (trace) |tr| try tr.record(p.stream, ai, pos0 + r0, m, p.lanes.len > 0, s.qids, s.qnk, s.qspr); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + const page_rows = try pagingRows(pg.cfg, w0, m); + var first: usize = 0; + while (first < m) { + const count = @min(page_rows, m - first); + const ids = s.qids + first * idw * 4; + const nk = s.qnk + first * 4; + const spr = s.qspr + first * 4; + if (ends / ratio > top) { + try pg.page(p.stream, ai, ids, nk, spr, count); + } else if (w0 != 0) return error.KvPageWindow; + try p.tri.attnChunks(at.q + (r0 + first) * wt.q_heads * head_dim * 2, kd, vd, pos, s.po, s.pm, s.pl, ids, nk, spr, count, kv_heads, chunks, true); + try p.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + (r0 + first) * wt.out_k * 2, nk, spr, count, kv_heads, true); + first += count; + } + } + } +}; + +pub fn pagingRows(cfg: Config, origin: usize, rows: usize) !usize { + try cfg.validate(); + if (rows == 0 or rows > att_rows) return error.KvPageRows; + if (origin % ratio != 0) return error.KvPageWindow; + if (origin / ratio <= cfg.slots) return rows; + const bounded = @min(rows, cfg.slots / top); + if (bounded == 0) return error.KvPageCapacity; + return bounded; +} + +pub fn windowPlan(cfg: Config, wr: usize, w0: usize, pos0: usize, n: usize) !usize { + try cfg.validate(); + if (n == 0 or n > cfg.rows) return error.KvPageRows; + if (w0 % ratio != 0) return error.KvPageWindow; + if (pos0 == 0) { + _ = try pagingRows(cfg, 0, @min(n, att_rows)); + return 0; + } + if (pos0 < w0) return error.KvPageRewind; + if (pos0 - w0 > wr) return error.KvPageRows; + const end = try std.math.add(usize, pos0 - w0, n); + const next = if (end > wr) (pos0 -| cfg.window) / ratio * ratio else w0; + if (pos0 - next > wr or n > wr - (pos0 - next)) return error.KvPageRows; + _ = try pagingRows(cfg, next, @min(n, att_rows)); + return next; +} + +/// Device bytes of a paged sequence's K/V side for `cap` tokens under `cfg` (windows, slots, tables, work lists). +pub fn deviceBytesFor(cap: usize, cfg: Config) !usize { + const layout = try SnapshotLayout.init(cap, cfg, 0); + if (cfg.shared_lanes != 0) { + _ = try sharedBytesFor(cap, cfg); + return 0; + } + return std.math.add(usize, try std.math.add(usize, layout.device, layout.tables), work_n * 4 * 4); +} + +pub fn sharedBytesFor(cap: usize, cfg: Config) !usize { + try cfg.validate(); + if (cfg.shared_lanes == 0) return 0; + return (try shared_device.Layout.init(.{ .cap = cap, .lanes = cfg.shared_lanes, .window = cfg.window, .rows = cfg.rows, .slots = cfg.slots })).common_bytes; +} + +pub fn ownedSnapshotBytes(cfg: Config) !usize { + if (cfg.shared_lanes == 0) return 0; + try cfg.validate(); + return std.math.mul(usize, cfg.slots, 2 * ratio * row_bytes + @sizeOf(shared_contract.OwnedEntry)); +} + +/// A new sequence's paged store when the environment (or a test) turns paging on, else null. +pub fn forSeq(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize, chunk: usize) !?*Paged { + const cfg = configFromEnv() orelse return null; + return try Paged.init(gpa, d, cap, cfg.forChunk(chunk)); +} + +test "exl3: kvpage tables and device bytes" { + const cfg: Config = .{}; + try std.testing.expect(cfg.window >= wt.dense_keys); + // 200k tokens: the device side of the default shape against upstream's K/V for the same sequence + const wr = cfg.window + cfg.rows + ratio; + const dev = 2 * (attn_n * wr + ratio * cfg.slots) * row_bytes; + const full = attn_n * 2 * 200_016 * row_bytes; + try std.testing.expect(dev * 8 < full); + try std.testing.expectEqual(@as(usize, 1024), row_bytes); + // 200k: the default shape's device side is about half a GiB (upstream's K/V: 4.58 GiB) + try std.testing.expect(try deviceBytesFor(200_016, cfg) < 5 << 27); +} + +test "exl3: kvpage allocation and admission reject invalid checked sizes before CUDA access" { + const cfg: Config = .{}; + const layout = try SnapshotLayout.init(200_016, cfg, 0); + try std.testing.expectEqual(layout.device + layout.tables + work_n * 4 * 4, try deviceBytesFor(200_016, cfg)); + try std.testing.expectError(error.Overflow, deviceBytesFor(std.math.maxInt(usize), cfg)); + try std.testing.expectError(error.KvPageRows, deviceBytesFor(0, cfg)); + var driver: cuda.Driver = undefined; + try std.testing.expectError(error.Overflow, Paged.init(std.testing.allocator, &driver, std.math.maxInt(usize), cfg)); + var oversized = cfg; + oversized.rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, Paged.init(std.testing.allocator, &driver, 32, oversized)); +} + +test "exl3: kvpage window boundaries and conservative capacity" { + const cfg: Config = .{ .window = 2052, .rows = 256, .slots = 512 }; + const wr = cfg.window + cfg.rows + ratio; + try std.testing.expectEqual(@as(usize, 0), try windowPlan(cfg, wr, 0, wr - 1, 1)); + try std.testing.expectEqual(@as(usize, 260), try windowPlan(cfg, wr, 0, wr, 1)); + try std.testing.expectEqual(@as(usize, 0), try windowPlan(cfg, wr, 260, 0, 256)); + try std.testing.expectError(error.KvPageRewind, windowPlan(cfg, wr, 260, 259, 1)); + try std.testing.expectError(error.KvPageRows, windowPlan(cfg, wr, 0, wr + 1, 1)); + try std.testing.expectEqual(@as(usize, 8192), try windowPlan(cfg, wr, 8192, 8192, 256)); + try std.testing.expectEqual(@as(usize, 8192), try windowPlan(cfg, wr, 8192, 8192, 1)); +} + +test "exl3: kvpage graph key changes on rebase and host store replacement" { + var pg: Paged = undefined; + pg.dev.ptr = 1; + pg.tabs.ptr = 2; + pg.work.ptr = 3; + pg.host_dev = 4; + pg.cap = 200_016; + pg.cfg = .{}; + pg.wr = pg.cfg.window + pg.cfg.rows + ratio; + pg.win0 = @splat(0); + const first = pg.graphKey(); + try std.testing.expectEqual(first, pg.graphKey()); + pg.win0[0] = 4; + try std.testing.expect(!std.meta.eql(first, pg.graphKey())); + pg.win0[0] = 0; + pg.host_dev = 5; + try std.testing.expect(!std.meta.eql(first, pg.graphKey())); +} + +test "exl3: prefill graph key separates reusable geometry from generation and origin" { + var pg: Paged = undefined; + pg.dev.ptr = 1; + pg.tabs.ptr = 2; + pg.work.ptr = 3; + pg.host_dev = 4; + pg.cap = 200016; + pg.cfg = .{ .window = 4096, .rows = 2048, .slots = 20000 }; + pg.wr = pg.cfg.window + pg.cfg.rows + ratio; + pg.incarnation = 1; + pg.win0 = @splat(65536); + const first = pg.prefillGraphKey(); + pg.incarnation += 1; + const new_generation = pg.prefillGraphKey(); + try std.testing.expectEqual(first[4], new_generation[4]); + try std.testing.expect(first[5] != new_generation[5]); + pg.win0[0] += ratio; + const moved = pg.prefillGraphKey(); + try std.testing.expectEqual(first[4], moved[4]); + try std.testing.expect(moved[5] != new_generation[5]); + pg.cfg.slots += 1; + const changed_shape = pg.prefillGraphKey(); + try std.testing.expect(changed_shape[4] != moved[4]); +} + +test "exl3: kvpage capacity subdivision boundaries and checked admission" { + var cfg: Config = .{ .window = 4096, .rows = 2048, .slots = 16384 }; + try std.testing.expectEqual(@as(usize, 256), try pagingRows(cfg, 65536, 256)); + try std.testing.expectEqual(@as(usize, 32), try pagingRows(cfg, 65540, 256)); + try std.testing.expectEqual(@as(usize, 31), try pagingRows(cfg, 65540, 31)); + try std.testing.expectEqual(@as(usize, 32), try pagingRows(cfg, 65540, 33)); + cfg.slots = 16383; + try std.testing.expectEqual(@as(usize, 31), try pagingRows(cfg, 65536, 256)); + cfg.slots = top - 1; + try std.testing.expectEqual(@as(usize, 256), try pagingRows(cfg, (top - 1) * ratio, 256)); + try std.testing.expectError(error.KvPageCapacity, pagingRows(cfg, top * ratio, 1)); + const wr = cfg.window + cfg.rows + ratio; + try std.testing.expectError(error.KvPageCapacity, windowPlan(cfg, wr, top * ratio, top * ratio, 256)); + try std.testing.expectError(error.KvPageRows, pagingRows(cfg, 0, 0)); + try std.testing.expectError(error.KvPageRows, pagingRows(cfg, 0, att_rows + 1)); + try std.testing.expectError(error.KvPageWindow, pagingRows(cfg, 1, 1)); + try std.testing.expectError(error.Overflow, windowPlan(cfg, std.math.maxInt(usize), 0, std.math.maxInt(usize), 1)); + cfg.slots = 0; + try std.testing.expectError(error.KvPageConfig, windowPlan(cfg, wr, 0, 0, 1)); +} + +test "exl3: kvpage adversarial disjoint old blocks fit every subdivision" { + for ([_]usize{ 512, 513, 1023, 1024, 16383, 16384 }) |slots| { + const cfg: Config = .{ .slots = slots }; + const origin = (slots + 1) * ratio; + for ([_]usize{ 1, 2, 15, 16, 17, 30, 31, 32, 33, 127, 255, 256 }) |rows| { + const width = try pagingRows(cfg, origin, rows); + var first: usize = 0; + var calls: usize = 0; + while (first < rows) { + const count = @min(width, rows - first); + var selected = std.AutoHashMap(usize, void).init(std.testing.allocator); + defer selected.deinit(); + for (first..first + count) |row| for (0..top) |pick| { + const block = (row * top + pick) % (origin / ratio); + try selected.put(block, {}); + }; + try std.testing.expect(selected.count() <= slots); + try std.testing.expectEqual(count * top, selected.count()); + first += count; + calls += 1; + } + try std.testing.expectEqual(rows, first); + try std.testing.expectEqual(cdiv(rows, width), calls); + if (width < rows) try std.testing.expect(origin > wt.dense_keys); + } + } + const cfg: Config = .{ .slots = 511 }; + try std.testing.expectEqual(att_rows, try pagingRows(cfg, 511 * ratio, att_rows)); +} + +test "exl3: kvpage 18k and 200k schedules use identical admission and paging bounds" { + const cfg: Config = .{ .window = 4096, .rows = 2048, .slots = 16384 }; + const wr = cfg.window + cfg.rows + ratio; + for ([_]usize{ 18432, 18448, 200000, 200016 }) |cap| { + var origin: usize = 0; + var position: usize = 0; + var split_calls: usize = 0; + while (position < cap) { + const rows = @min(cfg.rows, cap - position); + origin = try windowPlan(cfg, wr, origin, position, rows); + var first: usize = 0; + while (first < rows) : (first += att_rows) { + const count = @min(att_rows, rows - first); + const width = try pagingRows(cfg, origin, count); + try std.testing.expect(@min(width * top, origin / ratio) <= cfg.slots); + if (origin / ratio <= cfg.slots) try std.testing.expectEqual(count, width); + if (width < count) split_calls += 1; + } + position += rows; + } + try std.testing.expectEqual(cap > 65536 + cfg.window, split_calls > 0); + try std.testing.expectEqual(@as(usize, 1), try pagingRows(cfg, origin, 1)); + try std.testing.expectEqual(@as(usize, 16), try pagingRows(cfg, origin, 16)); + } +} + +test "exl3: kvpage graph topology hash distinguishes subdivision across origins" { + var pg: Paged = undefined; + pg.dev.ptr = 1; + pg.tabs.ptr = 2; + pg.work.ptr = 3; + pg.host_dev = 4; + pg.cap = 200016; + pg.cfg = .{ .window = 4096, .rows = 2048, .slots = 16384 }; + pg.wr = pg.cfg.window + pg.cfg.rows + ratio; + pg.win0 = @splat(65536); + const full = pg.graphKey(); + pg.win0[0] += ratio; + const split = pg.graphKey(); + try std.testing.expect(full[4] != split[4]); + try std.testing.expect(full[5] != split[5]); + pg.win0[0] += ratio; + const moved = pg.graphKey(); + try std.testing.expectEqual(split[4], moved[4]); + try std.testing.expect(split[5] != moved[5]); +} + +test "exl3: kvpage real attention dispatch preserves selection and slices sparse rows" { + const Fake = struct { + const ids_base: u64 = 0x100000; + const nk_base: u64 = 0x200000; + const spr_base: u64 = 0x300000; + const query_base: u64 = 0x400000; + const output_base: u64 = 0x500000; + const pos_base: u64 = 0x600000; + var selected: usize = 0; + var consumed: usize = 0; + var planned: usize = 0; + var width: usize = 0; + var launches: usize = 0; + + fn value(parameters: ?[*]?*anyopaque, index: usize, comptime Value: type) Value { + const pointer: *const Value = @ptrCast(@alignCast(parameters.?[index].?)); + return pointer.*; + } + + fn launch(kernel: cuda.abi.Function, grid_x: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + launches += 1; + switch (@intFromPtr(kernel.?)) { + 1 => { + std.debug.assert(value(parameters, 1, u64) == pos_base + selected * 16); + std.debug.assert(value(parameters, 2, u64) == ids_base); + std.debug.assert(grid_x == @min(att_rows, 531 - consumed)); + selected += 1; + }, + 2 => { + const offset = consumed % att_rows; + std.debug.assert(value(parameters, 0, u64) == ids_base + offset * idw * 4); + std.debug.assert(value(parameters, 1, u64) == nk_base + offset * 4); + std.debug.assert(value(parameters, 2, u64) == spr_base + offset * 4); + planned = @intCast(value(parameters, 3, c_int)); + std.debug.assert(planned == @min(width, @min(att_rows - offset, 531 - consumed))); + }, + 3 => {}, + 4 => std.debug.assert(grid_x == planned), + else => unreachable, + } + return cuda.abi.success; + } + + fn attnPrep(_: @This(), _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: u64, _: f32, _: usize, _: usize) !void {} + fn qsaPool(_: @This(), _: u64, _: u64, _: u64, _: u64, _: u64, _: f32, _: usize) !void {} + fn qsaScores(_: @This(), _: u64, _: u64, _: u64, _: u64, _: usize, _: usize, _: usize) !void {} + fn attnChunks(_: @This(), query: u64, _: u64, _: u64, pos: u64, _: u64, _: u64, _: u64, ids: u64, nk: u64, spr: u64, rows: usize, _: usize, chunks: usize, qsa: bool) !void { + try std.testing.expectEqual(query_base + consumed * wt.q_heads * head_dim * 2, query); + try std.testing.expectEqual(pos_base + (selected - 1) * 16, pos); + try std.testing.expectEqual(ids_base + (consumed % att_rows) * idw * 4, ids); + try std.testing.expectEqual(nk_base + (consumed % att_rows) * 4, nk); + try std.testing.expectEqual(spr_base + (consumed % att_rows) * 4, spr); + try std.testing.expectEqual(planned, rows); + try std.testing.expectEqual(@as(usize, 5), chunks); + try std.testing.expect(qsa); + } + fn attnMerge(_: @This(), _: u64, _: u64, _: u64, pos: u64, output: u64, _: u64, _: u64, rows: usize, _: usize, qsa: bool) !void { + try std.testing.expectEqual(pos_base + (selected - 1) * 16, pos); + try std.testing.expectEqual(output_base + consumed * wt.out_k * 2, output); + try std.testing.expect(qsa); + consumed += rows; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Fake.launch; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const function: cuda.Function = .{ .d = &driver, .handle = @ptrFromInt(1) }; + var pg: Paged = undefined; + pg.d = &driver; + pg.cfg = .{ .window = 4096, .rows = 2048, .slots = 16384 }; + pg.cap = 200016; + pg.wr = pg.cfg.window + pg.cfg.rows + ratio; + pg.nblocks = cdiv(pg.cap, ratio) + 1; + pg.dev.ptr = 0x700000; + pg.tabs.ptr = 0x800000; + pg.work.ptr = 0x900000; + pg.host_dev = 0xa00000; + pg.graph_prepared = true; + pg.k.plan = .{ .d = &driver, .handle = @ptrFromInt(2) }; + pg.k.fetch = .{ .d = &driver, .handle = @ptrFromInt(3) }; + pg.k.remap = .{ .d = &driver, .handle = @ptrFromInt(4) }; + const Buffer = struct { ptr: u64 = 0 }; + const Seq = struct { cap: usize, ikc: [attn_n]Buffer = @splat(.{}), pooled: [attn_n]Buffer = @splat(.{}), scores: u64 = 0 }; + const State = struct { inv: u64 = 0, dummy: u64 = 0, qids: u64 = Fake.ids_base, qnk: u64 = Fake.nk_base, qspr: u64 = Fake.spr_base, po: u64 = 0, pm: u64 = 0, pl: u64 = 0 }; + const Selection = struct { select: cuda.Function }; + const Pass = struct { seq: *Seq, s: *State, stream: cuda.Stream, tri: Fake = .{}, k: *const Selection, pos0: usize, lanes: []const u8 = &.{} }; + var seq: Seq = .{ .cap = pg.cap }; + var state: State = .{}; + const kernels: Selection = .{ .select = function }; + const layer = .{ .q_scale = Buffer{}, .k_scale = Buffer{}, .i_scale = Buffer{}, .ik_scale = Buffer{} }; + const at = .{ .pos = Fake.pos_base, .q = Fake.query_base, .iq = @as(u64, 0), .ao = Fake.output_base }; + for ([_]usize{ 14336, 65536, 65540, 194560 }) |origin| { + pg.win0 = @splat(origin); + Fake.selected = 0; + Fake.consumed = 0; + Fake.planned = 0; + Fake.launches = 0; + Fake.width = try pagingRows(pg.cfg, origin, att_rows); + var pass: Pass = .{ .seq = &seq, .s = &state, .stream = stream, .k = &kernels, .pos0 = origin + pg.cfg.window }; + try pg.attnRows(&pass, layer, 0, 0, 531, at); + try std.testing.expectEqual(@as(usize, 3), Fake.selected); + try std.testing.expectEqual(@as(usize, 531), Fake.consumed); + const launches = Fake.launches; + pg.cfg.slots = top - 1; + try std.testing.expectError(error.KvPageCapacity, pg.attnRows(&pass, layer, 0, 0, 531, at)); + try std.testing.expectEqual(launches, Fake.launches); + pg.cfg.slots = 16384; + } +} + +test "exl3: kvpage host preparation writes back and compacts once before replay" { + const Fake = struct { + var copies: usize = 0; + fn copy(dst: u64, src: u64, bytes: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + std.debug.assert(dst + bytes <= src or src + bytes <= dst); + const target: [*]u8 = @ptrFromInt(dst); + const source: [*]const u8 = @ptrFromInt(src); + @memcpy(target[0..bytes], source[0..bytes]); + copies += 1; + return cuda.abi.success; + } + }; + Fake.copies = 0; + var driver: cuda.Driver = undefined; + driver.api.cuMemcpyDtoDAsync_v2 = Fake.copy; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var pg: Paged = undefined; + pg.d = &driver; + pg.cfg = .{ .window = 2052, .rows = 256, .slots = 512 }; + pg.cap = 8192; + pg.wr = pg.cfg.window + pg.cfg.rows + ratio; + pg.win0 = @splat(0); + pg.rebases = @splat(0); + pg.written = @splat(0); + const device = try std.testing.allocator.alloc(u8, pg.devBytes()); + defer std.testing.allocator.free(device); + pg.dev.ptr = @intFromPtr(device.ptr); + const host = try std.testing.allocator.alloc(u8, 2 * pg.cap * row_bytes); + defer std.testing.allocator.free(host); + @memset(host, 0); + pg.host_dev = @intFromPtr(host.ptr); + for (0..2) |side| { + const start = if (side == 0) pg.devK(0) else pg.devV(0); + const words: [*]u32 = @ptrFromInt(start); + for (0..pg.wr) |position| @memset(words[position * (row_bytes / 4) ..][0 .. row_bytes / 4], @as(u32, @intCast(position + side * pg.cap))); + } + const before = pg.win0[0]; + const pos0 = pg.wr; + try pg.prepareLayer(stream, 0, pos0, 1); + const moved = pg.win0[0]; + try std.testing.expect(moved > before); + try std.testing.expectEqual(@as(u64, 1), pg.rebases[0]); + try std.testing.expectEqual(@as(u64, moved), pg.written[0]); + for (0..2) |side| { + const written: [*]const u32 = @ptrFromInt(pg.host_dev + side * pg.cap * row_bytes); + const kept: [*]const u32 = @ptrFromInt(if (side == 0) pg.devK(0) else pg.devV(0)); + for (0..moved) |position| for (written[position * (row_bytes / 4) ..][0 .. row_bytes / 4]) |word| try std.testing.expectEqual(@as(u32, @intCast(position + side * pg.cap)), word); + for (moved..pos0) |position| for (kept[(position - moved) * (row_bytes / 4) ..][0 .. row_bytes / 4]) |word| try std.testing.expectEqual(@as(u32, @intCast(position + side * pg.cap)), word); + } + const copies = Fake.copies; + try pg.prepareLayer(stream, 0, pos0, 1); + try std.testing.expectEqual(copies, Fake.copies); + try std.testing.expectEqual(@as(u64, 1), pg.rebases[0]); +} diff --git a/zig/src/families/flashnext/cuda_exl3_load.zig b/zig/src/families/flashnext/cuda_exl3_load.zig new file mode 100644 index 0000000000..db9ce28eee --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_load.zig @@ -0,0 +1,566 @@ +//! EXL3 weights for the Flash Next CUDA family's own loaders (cuda_projection_load.zig, cuda_prompt_load.zig): +//! - `Proj`: a projection made of column parts, each an EXL3 matrix in its stored layout (cuda.exl3.Weight, no +//! strips reorder) or a dense BF16 lane weight (the F16 in_proj_a/b, widened exactly by the Store), written into +//! one row of `n` columns: what Projection.exl3 points at and Pass.lane() runs (cuda_exl3_pass.zig); +//! - `Experts`: one layer's 512 routed experts and the shared one as id 512 in the same 513-entry pointer and width +//! tables (cuda.exl3_experts.Layer), suh/svh stacked [513, ·], resident; +//! - `Ops`: the kernels and scratch every EXL3 projection and expert call of one engine shares (one exl3.Dense, one +//! exl3_experts.Experts), counted so the last owner unloads them. +//! Widths come from each matrix's own trellis (16 * bits int16 words a tile), never from config.json's `bits`. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); +const pool_ = @import("cuda_exl3_pool.zig"); + +const Tensor = core.safetensors.Tensor; +const Store = wt.Store; + +/// format.py's mul1 marker value (core.exl3_load.markers). +const mul1_magic: u32 = blk: { + for (core.exl3_load.markers) |m| if (m.codebook == .mul1) break :blk m.value; + unreachable; +}; + +/// The EXL3 kernels and scratch one engine's calls share, on the engine's one compute stream. +pub const Ops = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + refs: usize = 1, + dense: ?cuda.exl3.Dense = null, + experts: ?cuda.exl3_experts.Experts = null, + /// A part's rows before its columns land in the stitched row (bf16 [rows, part n]). + tmp: ?cuda.DeviceBuffer = null, + /// The routed experts' scratch for `moe_rows` rows of 11 slots. + moe: ?cuda.DeviceBuffer = null, + moe_rows: usize = 0, + /// fp16 rows split into bf16 hi and lo [2, rows, k], and their fp32 products [2, rows, n] (an fp16-input Proj). + halves: ?cuda.DeviceBuffer = null, + sums: ?cuda.DeviceBuffer = null, + glue: ?cuda.Module = null, + split_fn: cuda.Function = undefined, + /// Expert parallel's masks, sums and kernels (cuda_exl3_tp.zig), made on first use. + tp: ?*@import("cuda_exl3_tp.zig").Scratch = null, + /// Sizing maxima over what was loaded (`reserve` allocates every scratch for them once, before any round): the + /// widest EXL3 input, Dense's split partials and counters for 128 rows, the widest part narrower than its row, + /// an fp16-input projection's k and n, and one experts layer's shape. + max_k: usize = 0, + max_z: usize = 0, + max_counters: usize = 0, + max_part_n: usize = 0, + max_f16_k: usize = 0, + max_f16_n: usize = 0, + moe_layer: ?cuda.exl3_experts.Layer = null, + /// Pinned [moe_rows, slots] int32: one block's router picks on the host (the offload pool's A/B cut). + picks: ?cuda.HostBuffer = null, + + pub fn create(gpa: std.mem.Allocator, d: *const cuda.Driver) !*Ops { + const o = try gpa.create(Ops); + o.* = .{ .gpa = gpa, .d = d }; + return o; + } + + pub fn retain(o: *Ops) *Ops { + o.refs += 1; + return o; + } + + pub fn release(o: *Ops) void { + o.refs -= 1; + if (o.refs != 0) return; + @import("cuda_exl3_fork.zig").release(o); + if (o.dense) |*x| x.unload(); + if (o.experts) |*x| x.unload(); + if (o.tmp) |*b| b.free(); + if (o.moe) |*b| b.free(); + if (o.halves) |*b| b.free(); + if (o.sums) |*b| b.free(); + if (o.picks) |*b| b.free(); + if (o.glue) |*m| m.unload(); + if (o.tp) |t| t.destroy(); + o.gpa.destroy(o); + } + + /// The dense EXL3 linear, loaded on first use (no GPU work at checkpoint load). + pub fn linear(o: *Ops) !*cuda.exl3.Dense { + if (o.dense == null) o.dense = try cuda.exl3.Dense.load(o.d); + return &o.dense.?; + } + + /// fn_exl3_split_f16 (flashnext_exl3.cu), loaded on first use. + pub fn split(o: *Ops) !cuda.Function { + if (o.glue == null) { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var m = try cuda.Module.load(o.d, cuda.kernels.flashnext_exl3); + errdefer m.unload(); + o.split_fn = try m.function("fn_exl3_split_f16"); + o.glue = m; + } + return o.split_fn; + } + + pub fn grouped(o: *Ops) !*const cuda.exl3_experts.Experts { + if (o.experts == null) o.experts = try cuda.exl3_experts.Experts.load(o.d); + return &o.experts.?; + } + + pub fn noteWeight(o: *Ops, w: cuda.exl3.Weight, row_n: usize) void { + o.max_k = @max(o.max_k, w.k); + o.max_z = @max(o.max_z, cuda.exl3.Linear.zBytes(w, 128)); + o.max_counters = @max(o.max_counters, cuda.exl3.Linear.counterBytes(w)); + if (w.n != row_n) o.max_part_n = @max(o.max_part_n, w.n); + } + + /// Every scratch the EXL3 calls of a pass of up to `rows` rows read, allocated now (after `s` drains), so no call + /// inside a round grows a buffer (a growth synchronises, and would move an address a captured graph holds). + pub fn reserve(o: *Ops, s: cuda.Stream, rows: usize, moe_rows: usize, slots: usize) !void { + try s.synchronize(); + if (o.max_k > 0) { + const de = try o.linear(); + try grow(o.d, &de.xh, 128 * o.max_k * 2); + if (o.max_z > 0) try grow(o.d, &de.z, o.max_z); + try grow(o.d, &de.counters, o.max_counters); + try @import("cuda_exl3_fork.zig").reserve(o); // parallel parts' branches (decode projections) + } + if (o.max_part_n > 0) try grow(o.d, &o.tmp, rows * o.max_part_n * 2); + if (o.max_f16_k > 0) { + _ = try o.split(); + try grow(o.d, &o.halves, 2 * rows * o.max_f16_k * 2); + try grow(o.d, &o.sums, 2 * rows * o.max_f16_n * 4); + } + if (o.moe_layer) |l| { + _ = try o.grouped(); + try grow(o.d, &o.moe, try cuda.exl3_experts.scratchBytes(l, moe_rows, slots)); + if (o.picks == null) o.picks = try cuda.HostBuffer.alloc(o.d, moe_rows * slots * 4); + } + } + + fn grow(d: *const cuda.Driver, slot: *?cuda.DeviceBuffer, bytes: usize) !void { + if (bytes == 0) return; + if (slot.*) |b| if (b.len >= bytes) return; + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(d, bytes); + } + + /// Every scratch buffer's address and size: unchanged across a pass means nothing grew (nor synchronised) in it. + pub fn fingerprint(o: *const Ops) [12][2]u64 { + var f: [12][2]u64 = @splat(.{ 0, 0 }); + if (o.tp) |t| f[8..12].* = @import("cuda_exl3_tp.zig").fingerprint(t); + const de: ?cuda.exl3.Dense = o.dense; + const bufs = [8]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums, if (de) |x| x.xh else null, if (de) |x| x.z else null, if (de) |x| x.counters else null, null }; + for (bufs, f[0..8]) |b, *e| if (b) |x| { + e.* = .{ x.ptr, x.len }; + }; + if (o.picks) |h| f[7] = .{ @intFromPtr(h.bytes.ptr), h.bytes.len }; + return f; + } + + /// Device bytes `reserve` holds (for the pool's budget). + pub fn reservedBytes(o: *const Ops) usize { + var n: usize = 0; + for ([_]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums }) |b| if (b) |x| { + n += x.len; + }; + if (o.dense) |de| for ([_]?cuda.DeviceBuffer{ de.xh, de.z, de.counters }) |b| if (b) |x| { + n += x.len; + }; + if (o.tp) |t| n += @import("cuda_exl3_tp.zig").reservedBytes(t); + return n; + } + + /// At least `bytes` of part scratch; a growth waits for `s` first (an older call may still read the old buffer). + pub fn scratch(o: *Ops, s: cuda.Stream, slot: *?cuda.DeviceBuffer, bytes: usize) !u64 { + if (slot.*) |b| if (b.len >= bytes) return b.ptr; + try s.synchronize(); + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(o.d, bytes); + return slot.*.?.ptr; + } +}; + +/// The store's shared Ops, made on first use; every holder retains it. +pub fn opsOf(store: *Store, d: *const cuda.Driver) !*Ops { + if (store.exl3_ops == null) store.exl3_ops = try Ops.create(store.gpa, d); + return store.exl3_ops.?.retain(); +} + +/// One EXL3 matrix read from the store: words as stored, its width and shape. +pub const Mat = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u32, +}; + +/// Whether `prefix` names an EXL3 matrix (its trellis is indexed). +pub fn has(store: *const Store, prefix: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix}) catch return false); +} + +/// `prefix`'s trellis, scales and mul1 marker, checked: int16 [K/16, N/16, 16 * bits], fp16 suh [K] and svh [N]. +pub fn mat(store: *Store, prefix: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = try store.hold(try std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix})); + if (t.dtype != .i16 or t.rank != 3 or t.dim(2) % 8 != 0) return error.UnexpectedTensor; + const halves = t.dim(2) / 8; + // the mul1 half steps are odd halves; 2..16 halves (1 to 8 bits) are what the kernels instantiate + if (halves < 2 or halves > 16) return error.UnsupportedExl3Width; + const k = t.dim(0) * 16; + const n = t.dim(1) * 16; + if (k % 128 != 0 or n % 128 != 0) return error.UnexpectedTensor; + const marker = store.hold(try std.fmt.bufPrint(&buf, "{s}.mul1", .{prefix})) catch return error.NotExl3Mul1; + if (marker.dtype != .i32 or marker.bytes.len != 4 or std.mem.readInt(u32, marker.bytes[0..4], .little) != mul1_magic) return error.NotExl3Mul1; + const suh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.suh", .{prefix})); + const svh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.svh", .{prefix})); + if (!suh.is(.f16, &.{k}) or !svh.is(.f16, &.{n})) return error.UnexpectedTensor; + if (store.has(try std.fmt.bufPrint(&buf, "{s}.bias", .{prefix}))) return error.UnexpectedTensor; // none in the text model + return .{ .trellis = t.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = k, .n = n, .halves = @intCast(halves) }; +} + +/// One column range of a projection's output row. +pub const Part = struct { + col: usize, + n: usize, + kind: union(enum) { exl3: cuda.exl3.Weight, native: cuda.qmmf.Weight }, +}; + +/// A projection's parts and the device buffers they own. +pub const Proj = struct { + gpa: std.mem.Allocator, + ops: *Ops, + k: usize, + n: usize, + parts: []Part, + /// The rows the pass hands in: bf16, or fp16 bit patterns (an EXL3 pack's n-gram rows, cuda_exl3_ple.zig). + x: enum { bf16, f16 } = .bf16, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + natives: std.ArrayList(projection.Projection) = .empty, + + pub fn destroy(p: *Proj) void { + for (p.bufs.items) |*b| b.free(); + p.bufs.deinit(p.gpa); + for (p.natives.items) |*q| q.free(); + p.natives.deinit(p.gpa); + p.gpa.free(p.parts); + p.ops.release(); + p.gpa.destroy(p); + } + + fn up(p: *Proj, d: *const cuda.Driver, bytes: []const u8) !u64 { + try p.bufs.ensureUnusedCapacity(p.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// An EXL3 matrix as the stored-layout weight (trellis [K/16][N/16][W]: strides(k, n, halves, false)). + fn exl3Part(p: *Proj, d: *const cuda.Driver, m: Mat) !cuda.exl3.Weight { + const st = cuda.exl3.strides(m.k, m.n, m.halves, false); + return .{ .trellis = try p.up(d, m.trellis), .suh = try p.up(d, m.suh), .svh = try p.up(d, m.svh), .k = @intCast(m.k), .n = @intCast(m.n), .halves = m.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + } +}; + +/// A Projection whose rows come from `parts` (EXL3 matrices or dense BF16 rows), `n` columns of `k` inputs. +fn wrap(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = k, .n = n, .parts = &.{} }; + errdefer p.destroy(); + var parts: std.ArrayList(Part) = .empty; + errdefer parts.deinit(gpa); + var buf: [256]u8 = undefined; + var col: usize = 0; + var i: usize = 0; + while (i < names.len) { + const name = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, names[i] }); + if (has(store, name)) { + const m = try mat(store, name); + if (m.k != k) return error.UnexpectedTensor; + const w = try p.exl3Part(d, m); + p.ops.noteWeight(w, n); + try parts.append(gpa, .{ .col = col, .n = m.n, .kind = .{ .exl3 = w } }); + col += m.n; + i += 1; + continue; + } + // a run of dense BF16 parts (after the Store's exact widening) packs as one lane weight + var rows: usize = 0; + var j = i; + var bytes: std.ArrayList(u8) = .empty; + defer bytes.deinit(gpa); + while (j < names.len) : (j += 1) { + var nb: [256]u8 = undefined; + const dense_name = try std.fmt.bufPrint(&nb, "{s}{s}", .{ prefix, names[j] }); + if (has(store, dense_name)) break; + const t = try store.tensor(try std.fmt.bufPrint(&nb, "{s}{s}.weight", .{ prefix, names[j] })); + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k or t.bytes.len != t.dim(0) * k * 2) return error.UnexpectedTensor; + try bytes.appendSlice(gpa, t.bytes); + rows += t.dim(0); + } + try p.natives.ensureUnusedCapacity(gpa, 1); + const q = try projection.dense(gpa, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, k, 1, 1, 1 }, .bytes = bytes.items }); + p.natives.appendAssumeCapacity(q); + try parts.append(gpa, .{ .col = col, .n = rows, .kind = .{ .native = q.native.? } }); + col += rows; + i = j; + } + if (col != n) return error.UnexpectedTensor; + p.parts = try parts.toOwnedSlice(gpa); + var out = projection.Projection.empty(d); + out.n = n; + out.k = k; + out.exl3 = p; + return out; +} + +/// cuda_projection_load.load's EXL3 branch: `prefix` one EXL3 matrix of n outputs and k inputs. +pub fn load(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, &.{""}, n, k); +} + +/// Engine.initRank's head for an EXL3 pack: `lm_head` (5-bit, K = 2,560, N = 248,320) as one EXL3 projection that +/// Pass.lane() runs through exl3.Dense (bf16 logits, as the dense BF16 head writes them); null for any other pack. +/// At tp=2 the head is split by vocabulary by default (cuda_exl3_tp.headProjection: 124,160 columns a rank at whole +/// 128-column Hadamard blocks, joined by cuda_batch.zig projectHead's ncclAllGather); TF_FLASHNEXT_SHARD_HEAD=0 keeps +/// the whole head on both ranks. At tp=1 a split is refused. +pub fn head(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, world: usize, rank: usize) !?projection.Projection { + if (!has(store, "language_model.lm_head")) return null; + if (shardHead(world)) return try @import("cuda_exl3_tp.zig").headProjection(gpa, d, store, "language_model.lm_head", world, rank); + if (world == 1 and try @import("cuda_head.zig").enabled()) return error.BadHeadSharding; + return try load(gpa, d, store, "language_model.lm_head", wt.vocab, wt.dims); +} + +/// Whether an EXL3 head is split at `world`: at 2 unless TF_FLASHNEXT_SHARD_HEAD=0; never at 1. +pub fn shardHead(world: usize) bool { + if (world != 2) return false; + const v = std.c.getenv("TF_FLASHNEXT_SHARD_HEAD") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Whether any of `prefix`'s parts is an EXL3 matrix (cuda_projection_load.loadStitched's branch). +pub fn hasAny(store: *const Store, prefix: []const u8, names: []const []const u8) bool { + var buf: [256]u8 = undefined; + for (names) |name| if (has(store, std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name }) catch return false)) return true; + return false; +} + +/// cuda_projection_load.loadStitched's EXL3 branch: parts side by side in one output row, EXL3 or dense BF16. +pub fn stitched(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, names, n, k); +} + +/// cuda_prompt_load's PLE branch: an EXL3 pack's n-gram rows reach key_proj and value_proj as fp16 bit patterns +/// (cuda_exl3_ple.zig's default, as the reference multiplies fp16 rows by the F16 projections), so both dense +/// projections take fp16 rows: split exactly into two bf16 rows, two bf16 products summed in fp32 +/// (cuda_exl3_pass.zig). TF_FLASHNEXT_PLE_EMB=bf16 keeps the bf16 lanes (the rows then come rounded to bf16). +pub fn pleRows(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, key: *projection.Projection, value: *projection.Projection) !void { + if (store.exl3_ops == null) return; // not an EXL3 pack: the n-gram rows are bf16 + if (@import("cuda_exl3_ple.zig").optionsFromEnv().emb != .f16) return; + for ([_]*projection.Projection{ key, value }) |q| { + const native = q.native orelse return error.UnexpectedTensor; + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = q.k, .n = q.n, .parts = &.{}, .x = .f16 }; + errdefer p.destroy(); + p.ops.max_f16_k = @max(p.ops.max_f16_k, q.k); + p.ops.max_f16_n = @max(p.ops.max_f16_n, q.n); + p.parts = try gpa.alloc(Part, 1); + p.parts[0] = .{ .col = 0, .n = q.n, .kind = .{ .native = native } }; + try p.natives.append(gpa, q.*); + var out = projection.Projection.empty(d); + out.n = q.n; + out.k = q.k; + out.exl3 = p; + q.* = out; + } +} + +/// One layer's routed experts and the shared expert (id `wt.experts_n`) for exl3_experts.Experts. +/// Resident (TF_FLASHNEXT_EXL3_ROUTED=device): every trellis on the device, `layer` complete at load. +/// Offloaded (the default): the routed and shared trellis words in `bank` (pinned host memory, one region a layer), +/// the stacked suh/svh on the device, and `layer`'s six tables are the pool's (`attach`, once the family has sized it). +pub const Experts = struct { + gpa: std.mem.Allocator, + ops: *Ops, + layer: cuda.exl3_experts.Layer, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// Each expert's and projection's trellis words on the device, for a pool to rewrite (gate, up, down; 513 each). + trellis: [3][]u64 = .{ &.{}, &.{}, &.{} }, + halves: [3][]i32 = .{ &.{}, &.{}, &.{} }, + /// The ids this rank holds under expert parallel (cuda_exl3_tp.zig); null: one rank holds all 513. + tp: ?@import("cuda_exl3_tp.zig").Rank = null, + bank: ?cuda.HostBuffer = null, + /// The pool's Source: every routed expert's gate/up/down bytes inside `bank`, and the shared expert's. + routed: [][3][]const u8 = &.{}, + shared: [3][]const u8 = .{ &.{}, &.{}, &.{} }, + pool: ?*pool_.Pool = null, + /// Host time the bank's copies took at load (ns). + bank_ns: u64 = 0, + + pub fn destroy(x: *Experts) void { + if (x.pool) |p| { + p.deinit(); + x.gpa.destroy(p); + } + for (x.bufs.items) |*b| b.free(); + x.bufs.deinit(x.gpa); + for (x.trellis) |t| x.gpa.free(t); + for (x.halves) |h| x.gpa.free(h); + x.gpa.free(x.routed); + if (x.bank) |*b| b.free(); + x.ops.release(); + x.gpa.destroy(x); + } + + fn up(x: *Experts, d: *const cuda.Driver, bytes: []const u8) !u64 { + try x.bufs.ensureUnusedCapacity(x.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + x.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// Whether the routed experts wait for a pool (offloaded, none attached yet). + pub fn needsPool(x: *const Experts) bool { + return x.bank != null and x.pool == null; + } + + /// Bytes of one pool line of this layer (gate, up and down of its widest expert, 256-aligned). + pub fn lineBytes(x: *const Experts) usize { + var bytes: [3]usize = .{ 0, 0, 0 }; + for (x.routed) |e| for (0..3) |j| { + bytes[j] = @max(bytes[j], e[j].len); + }; + return cuda.exl3_pool.Layout.init(bytes).line_bytes; + } + + /// Makes the pool: `lines` device lines, `counts` the family's lifetime use table (null: the pool's own), `layer` + /// the policy's index; then `layer`'s tables are the pool's (fixed addresses, rewritten in place on a swap). + pub fn attach(x: *Experts, io: std.Io, d: *const cuda.Driver, which: u8, lines: usize, max_rows: usize, counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !void { + if (x.bank == null or x.pool != null) return error.Invalid; + const p = try x.gpa.create(pool_.Pool); + errdefer x.gpa.destroy(p); + const l = x.layer; + p.* = try pool_.Pool.init(x.gpa, io, d, .{ .routed = x.routed, .shared = if (x.shared[0].len == 0) null else x.shared, .pinned_bank = x.bank.?.bytes }, .{ .suh_g = l.suh_g, .suh_u = l.suh_u, .svh_g = l.svh_g, .svh_u = l.svh_u, .suh_d = l.suh_d, .svh_d = l.svh_d }, .{ + .layer = which, + .lines = lines, + .dims = wt.dims, + .width = wt.moe_w, + .max_rows = max_rows, + .slots = wt.slots, + .direct = true, // the bank is pinned: one H2D a projection on the copy engine, no staging memcpy + .codebook = .mul1, + .act_mode = l.act_mode, + .max_passes = pool_.passesFor(pool_.min_lines), + }, counts); + x.pool = p; + x.layer = p.tables(); + } +}; + +/// Whether `mlp` (".mlp" of a layer) holds EXL3 experts. +pub fn hasExperts(store: *const Store, mlp: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.trellis", .{mlp}) catch return false); +} + +/// TF_FLASHNEXT_EXL3_ROUTED=device keeps every routed expert on the device (cuda_exl3_memory.zig reads the same). +pub fn offloaded() bool { + return @import("cuda_exl3_memory.zig").options().routed_on_host; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// Every routed expert of `mlp`, the shared expert as id 512 (the same 2560 x 640 x 2560 shape): resident, or with +/// their trellis words in one pinned host region (see `Experts`). +pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, mlp: []const u8, world: usize, rank: usize) !*Experts { + if (world != 1 or rank != 0) return @import("cuda_exl3_tp.zig").experts(gpa, d, store, mlp, world, rank); // plan step 10 + const count = wt.experts_n + 1; + // TF_FLASHNEXT_EXL3_MTP_RESIDENT=1: the MTP layer's experts stay on the device (0.88 GiB) instead of a pool + const mtp_resident = std.mem.indexOf(u8, mlp, "mtp.") != null and if (std.c.getenv("TF_FLASHNEXT_EXL3_MTP_RESIDENT")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const offload = offloaded() and !mtp_resident; + const x = try gpa.create(Experts); + x.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .layer = undefined }; + errdefer x.destroy(); + for (&x.trellis, &x.halves) |*t, *h| { + t.* = try gpa.alloc(u64, count); + @memset(t.*, 0); + h.* = try gpa.alloc(i32, count); + @memset(h.*, 0); + } + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + const ks = [3]usize{ wt.dims, wt.dims, wt.moe_w }; + const ns = [3]usize{ wt.moe_w, wt.moe_w, wt.dims }; + var scales: [6]std.ArrayList(u8) = @splat(.empty); + defer for (&scales) |*l| l.deinit(gpa); + var buf: [256]u8 = undefined; + const nameOf = struct { + fn f(b: []u8, m: []const u8, e: usize, proj: []const u8) ![]const u8 { + if (e == wt.experts_n) return std.fmt.bufPrint(b, "{s}.shared_expert.{s}", .{ m, proj }); + return std.fmt.bufPrint(b, "{s}.experts.{d}.{s}", .{ m, e, proj }); + } + }.f; + var at: usize = 0; + if (offload) { + // the bank's size first (headers only), then one pinned region for the layer + var total: usize = 0; + for (0..count) |e| { + for (projs) |proj| total += align256((try mat(store, try nameOf(&buf, mlp, e, proj))).trellis.len); + store.release(); + } + x.bank = try cuda.HostBuffer.alloc(d, total); + x.routed = try gpa.alloc([3][]const u8, wt.experts_n); + } + const t0 = std.Io.Clock.awake.now(store.io).toNanoseconds(); + for (0..count) |e| { + for (projs, 0..) |proj, j| { + const m = try mat(store, try nameOf(&buf, mlp, e, proj)); + if (m.k != ks[j] or m.n != ns[j]) return error.UnexpectedTensor; + if (x.bank) |bank| { + const dst = bank.bytes[at..][0..m.trellis.len]; + @memcpy(dst, m.trellis); + at += align256(m.trellis.len); + if (e == wt.experts_n) x.shared[j] = dst else x.routed[e][j] = dst; + } else x.trellis[j][e] = try x.up(d, m.trellis); + x.halves[j][e] = @intCast(m.halves); + try scales[2 * j].appendSlice(gpa, m.suh); + try scales[2 * j + 1].appendSlice(gpa, m.svh); + } + // each expert's mapped pages are read once: release them so the 513 experts never pin a shard's worth + store.release(); + } + x.bank_ns = @intCast(std.Io.Clock.awake.now(store.io).toNanoseconds() - t0); + x.layer = .{ + // offloaded: the pool's tables replace these at attach (0 until then; a pass refuses an unattached layer) + .gate = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[0])), + .up = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[1])), + .down = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[2])), + .gate_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[0])), + .up_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[1])), + .down_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[2])), + .suh_g = try x.up(d, scales[0].items), + .svh_g = try x.up(d, scales[1].items), + .suh_u = try x.up(d, scales[2].items), + .svh_u = try x.up(d, scales[3].items), + .suh_d = try x.up(d, scales[4].items), + .svh_d = try x.up(d, scales[5].items), + .experts = count, + .dims = wt.dims, + .width = wt.moe_w, + .codebook = .mul1, + // fp32 SwiGLU between the grouped GEMMs (GLM's mode 0 rounds to bf16 there) + .act_mode = 1, + }; + x.ops.moe_layer = x.layer; + return x; +} + +test "exl3: the mul1 marker value is format.py's" { + try std.testing.expectEqual(@as(u32, 0x83DCD12D), mul1_magic); +} diff --git a/zig/src/families/flashnext/cuda_exl3_memory.zig b/zig/src/families/flashnext/cuda_exl3_memory.zig new file mode 100644 index 0000000000..8d8cdb5422 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_memory.zig @@ -0,0 +1,173 @@ +//! Device bytes of an EXL3 Flash Next checkpoint as cuda_exl3_load.zig places it: EXL3 matrices in their stored +//! layout (trellis, suh, svh; the codebook marker stays on the host), dense weights widened to BF16 and packed as +//! projection.dense packs them, norms and GDN scalars as fp32, the router's 513 rows as stored. Routed experts' trellis +//! words are host-resident under offload (`routed_on_host`; the pool decides which lines reach the device), while their +//! stacked suh/svh scales and the per-expert pointer and width tables stay resident. + +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const config = @import("cuda_config.zig"); +const Entry = core.safetensors.Entry; + +pub const Options = struct { + /// Routed experts' trellis words live in pinned host memory (the offload pool's), not counted here. + routed_on_host: bool = true, +}; + +/// TF_FLASHNEXT_EXL3_ROUTED=device counts every routed expert as resident; the default is the offload pool's layout. +pub fn options() Options { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_ROUTED") orelse return .{}; + return .{ .routed_on_host = !std.mem.eql(u8, std.mem.span(v), "device") }; +} + +/// Whether `dir`'s config.json admits as Mode.exl3 (any other mode, or a refusal: false). +pub fn isExl3(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + return (config.mode(gpa, io, dir) catch return false) == .exl3; +} + +/// One tensor's device bytes (0: kept on the host). +pub fn tensorBytes(name: []const u8, e: Entry, o: Options) !u64 { + const bytes: u64 = e.end - e.begin; + if (wt.hostOnly(name)) return 0; + // the codebook marker is checked at load and never uploaded + if (std.mem.endsWith(u8, name, ".mul1") or std.mem.endsWith(u8, name, ".mcg")) return 0; + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { + if (std.mem.endsWith(u8, name, ".trellis")) return if (o.routed_on_host) 0 else bytes; + return bytes; // stacked suh / svh, resident + } + if (std.mem.endsWith(u8, name, ".trellis") or std.mem.endsWith(u8, name, ".suh") or std.mem.endsWith(u8, name, ".svh")) return bytes; + if (!std.mem.endsWith(u8, name, ".weight") and e.rank == 1) return try std.math.mul(u64, e.dim(0), 4); // A_log, dt_bias as fp32 + if ((e.dtype == .f16 or e.dtype == .bf16) and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + // the router's rows and the embedding table go up as stored (BF16); every other dense matrix through projection.dense + if (std.mem.indexOf(u8, name, ".mlp.gate.") != null or std.mem.indexOf(u8, name, ".shared_expert_gate.") != null or std.mem.indexOf(u8, name, "embed_tokens") != null) return bytes; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(e.dim(0)), e.dim(1)), 2); + } + // norms the loaders scale go up as fp32 (cuda_rank_memory.zig's list), the GDN gated norm and convs as stored + if (e.rank == 1 and centered(name)) return try std.math.mul(u64, e.dim(0), 4); + return bytes; +} + +fn centered(name: []const u8) bool { + for ([_][]const u8{ ".hc_norm.weight", ".q_norm.weight", ".k_norm.weight", ".q_layernorm.weight", ".k_layernorm.weight", ".norm_key.weight", ".norm_query.weight", ".norm_conv.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return true; + } + return false; +} + +/// Per layer: three 513-entry pointer tables (int64) and three width tables (int32) for the experts' grouped GEMMs. +const tables_per_layer: u64 = 3 * (wt.experts_n + 1) * 8 + 3 * (wt.experts_n + 1) * 4; + +/// Every shard the index names (and the mixer patch, which Store.open adds), each read for its header alone. +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer * wt.layers_n; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| total = try std.math.add(u64, total, try tensorBytes(e.key_ptr.*, e.value_ptr.*, o)); + } + return total; +} + +/// The MTP head's device bytes: its `mtp.*` tensors placed as the main layers' are (its routed experts under the same +/// option), plus the draft vocabulary's bf16 rows dequantised from the EXL3 lm_head (cuda_mtp.zig's 79,591 ids). +pub fn headBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| if (std.mem.startsWith(u8, kv.key_ptr.*, "mtp.")) try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer + draft_rows * wt.dims * 2; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| { + const name = e.key_ptr.*; + if (!std.mem.startsWith(u8, name, "mtp.")) continue; + // the head's tensors as a decoder layer's: drop the `mtp.` prefix the host-only rule keys on + var buf: [256]u8 = undefined; + const as_layer = try std.fmt.bufPrint(&buf, "model.language_model.{s}", .{name["mtp.".len..]}); + total = try std.math.add(u64, total, try tensorBytes(as_layer, e.value_ptr.*, o)); + } + } + return total; +} + +/// The draft head's vocabulary rows (cuda_mtp.zig headBytes' count). +const draft_rows: u64 = 79_591; + +/// The family's deviceWeightBytes for an EXL3 checkpoint: the main model and the MTP head; 0 if unread. +pub fn weightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + const o = options(); + const main = deviceBytes(gpa, io, dir, o) catch return 0; + return main + (headBytes(gpa, io, dir, o) catch return 0); +} + +/// A rank's placement bytes; at world 2 the expert-parallel rank's (cuda_exl3_tp.zig, plan step 10). +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (world != 1) return @import("cuda_exl3_tp.zig").placement(gpa, io, dir, world, rank, drafts); + const o = options(); + return try deviceBytes(gpa, io, dir, o) + if (drafts) try headBytes(gpa, io, dir, o) else @as(u64, 0); +} + +fn entry(dtype: core.safetensors.DType, shape: []const usize, bytes: u64) Entry { + var e: Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = bytes }; + for (shape, 0..) |d, i| e.shape[i] = d; + return e; +} + +test "exl3: routed trellis stays on the host under offload, its scales and every other EXL3 matrix are resident" { + const tr = entry(.i16, &.{ 160, 40, 24 }, 160 * 40 * 24 * 2); + const routed = "model.language_model.layers.0.mlp.experts.7.gate_proj.trellis"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(routed, tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes(routed, tr, .{ .routed_on_host = false })); + const suh = entry(.f16, &.{2560}, 5120); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes("model.language_model.layers.0.mlp.experts.7.gate_proj.suh", suh, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("model.language_model.layers.0.mlp.shared_expert.gate_proj.trellis", tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("lm_head.mul1", entry(.i32, &.{}, 4), .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("mtp.layers.0.mlp.experts.3.up_proj.trellis", tr, .{ .routed_on_host = false })); +} + +test "exl3: dense F16 weights count as their BF16 device packing, norms as fp32" { + const router = entry(.f16, &.{ 512, 2560 }, 512 * 2560 * 2); + try std.testing.expectEqual(router.end, try tensorBytes("model.language_model.layers.3.mlp.gate.weight", router, .{})); + const down = entry(.f16, &.{ 320, 10240 }, 320 * 10240 * 2); + try std.testing.expectEqual(@as(u64, cuda.nvfp4.padded(320) * 10240 * 2), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.input_mix_weight_down.weight", down, .{})); + const norm = entry(.f16, &.{10240}, 20480); + try std.testing.expectEqual(@as(u64, 40960), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.hc_norm.weight", norm, .{})); + const gnorm = entry(.bf16, &.{128}, 256); + try std.testing.expectEqual(@as(u64, 256), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", gnorm, .{})); + try std.testing.expectEqual(@as(u64, 48 * 4), try tensorBytes("model.language_model.layers.0.linear_attn.A_log", entry(.bf16, &.{48}, 96), .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("model.visual.blocks.0.attn.qkv.weight", entry(.bf16, &.{ 3456, 1152 }, 3456 * 1152 * 2), .{})); +} + +test "exl3: placement names an EXL3 checkpoint and refuses worlds past 2 (world 2: cuda_exl3_tp.zig's test)" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = @embedFile("fixtures/exl3_config.json") }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expect(isExl3(a, std.testing.io, dir)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 3, 0, false)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_mtp.zig b/zig/src/families/flashnext/cuda_exl3_mtp.zig new file mode 100644 index 0000000000..e3bd4e1313 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_mtp.zig @@ -0,0 +1,317 @@ +//! An EXL3 pack's MTP head on v1.1.0's Flash Next family (cuda_mtp.zig): the draft vocabulary rows here, the rest a design. +//! +//! ## What the head computes (upstream cuda_mtp_ops.zig `forwardHead`, = tensorfold qwen4_exp mtp.py `FlashMTP`) +//! e = fc_embedding(norm_e(embed(token t+1))) norm over 2,560 (tri.rmsnorm, group = dims) +//! hs = fc_hidden(norm_h(streams_t) as 4 rows of 2,560) ONE RMS over the whole 10,240 row (group = wide), then +//! fc_hidden per stream: what the EXL3 reference does too +//! x = e + hs[s] per stream; one attention + MoE decoder layer (`mtp.layers.0`, its own KV, no PLE) +//! out = hyper_connection_mixer readout of x, then the draft head: bf16 rows of lm_head at the 79,591 draft ids. +//! Upstream's maths already matches the EXL3 reference; only the weights' formats differ. +//! +//! ## EXL3 tensors (pack on par1-cs13, index + mtp_hyper_connection_mixer_patch.safetensors, read 2026-10-19) +//! mtp.fc_embedding.{trellis,suh,svh,mul1}, mtp.fc_hidden.{...} EXL3 2560 -> 2560 (mtp_bits 3) +//! mtp.pre_fc_norm_embedding.weight [2560], mtp.pre_fc_norm_hidden.weight [10240] gamma - 1 (store.norm_offset) +//! mtp.layers.0.{self_attn.*, mlp.experts.N.*, mlp.shared_expert.*, mlp.gate, mlp.shared_expert_gate, +//! attn_hyper_connection.*, mlp_hyper_connection.*} as a main attention layer (F16 gates and HC) +//! mtp.hyper_connection_mixer.{hc_norm [10240], input_mix_weight_down [320,10240], _up [10240,320]} F16, in the +//! patch file, which model.safetensors.index.json does not name +//! lm_head.{trellis [160, 15520, 80], suh [2560], svh [248320], mul1} EXL3 5-bit (head_bits 5), K = 2560, N = vocab +//! +//! ## Hook points (one line each into upstream files; the work lives in cuda_exl3_*.zig) +//! 1. cuda_mtp_load.zig `draft` (after the mode check): `if (try exl3_mtp.draft(a, d, store, ids, affine)) |q| return q;` +//! DONE here: `draft` below dequantises the lm_head columns of the draft ids (exl3_format.dequantize on the +//! 128-column blocks that hold them) to bf16 rows [ids, 2560] and uploads them through projection.dense (or the +//! affine4 packer under TF_FLASHNEXT_DRAFT_HEAD=affine4). The forward's `p.lane(.., h.draft, ..)` is unchanged. +//! 2. cuda_mtp.zig Head.init `fc` (line ~198): `cuda_projection_load.load` refuses EXL3 (`UnsupportedQuantization`: +//! no `.weight`). E1's `Projection.exl3` load (plan step 4) covers it: in `load`, `if (store.has(prefix ++ +//! ".trellis")) return cuda_exl3_proj.load(...)`. `forwardHead` runs fc_e/fc_h through `Pass.lane(.., fc_sk = 4)`: +//! an EXL3 lane must honour `sk > 1` callers that read `part` (fp32 slices) and `sumSlices`, or return sk = 1 +//! with the sum already in `out` (then the two `sumSlices` calls at cuda_mtp_ops.zig:146 and :150 must be +//! skipped for EXL3: one branch each, `if (!h.fc_e.isExl3())`). +//! 3. cuda_mtp.zig Head.init `wt.loadHc(.., "language_model.mtp.hyper_connection_mixer", false)` (line ~179): the +//! Store (cuda_prompt_wt.zig `resolveKey`/`shard`) must fall back to the patch file for names the index lacks +//! (plan step 2), and `loadHc` needs the F16 dense HC path (plan step 3; `loadDenseHc` takes bf16 today, and +//! `asF32` refuses F16, so `normScale` of an F16 hc_norm fails too). +//! 4. cuda_mtp.zig Head.init `load.loadLayerAt(.., "language_model.mtp.layers.0", false, false)` (line ~185): the +//! alias `language_model.mtp.` -> `mtp.` already resolves; the layer needs E1's EXL3 attention projections and +//! `Experts` (plan steps 4-5) and the F16 router/shared gate (step 3). Its experts can stay resident (one layer, +//! 513 x 3-bit experts, about 0.95 GB) or ride on the offload pool as a 49th layer. +//! 5. cuda_mtp.zig `store.detectNorms()` (line ~164) reads the 48 decoder `attn_hyper_connection.hc_norm` (F16 in +//! this pack): same F16 `asF32` fix as 3. q_norm/k_norm of `mtp.layers.0.self_attn` average ~2.7 stored, still +//! gamma - 1 (the branch's exl3_mtp.zig used the single convention); upstream applies one offset to all: fine. +//! 6. cuda_mtp.zig `headBytes` (admission estimate) counts affine bytes; an EXL3 head is fc 2 x 2560 x 2560 x 3/8, +//! the layer's EXL3 bytes and the bf16 draft rows (79,591 x 2560 x 2 = 407.5 MB). A one-line branch on the format. +//! 7. Not needed: the draft chain, verify windows, rollback (cuda_mtp_draft.zig, cuda_mtp_ops.zig) are format-free. +//! +//! ## Draft rows: what is exact and what is not +//! EXL3 columns are not separable (a 128-column Hadamard mixes them), so a row is decoded from its whole column block +//! and rounded to bf16 once (round to nearest even from the float64 value). A draft logit then differs from the full +//! head's by at most 2^-9 * sum_k |x_k W[k,v]| (bf16 keeps 8 significant bits), the bound `tf-cuda-test exl3-draft` +//! checks against the full EXL3 head (upstream exl3.Dense, fp32 out) and an exact float64 product. Drafts are proposals: +//! the main head verifies every token, so acceptance (drafted == undrafted) is unaffected; only the acceptance rate is. + +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); +const fmt = core.exl3_format; + +const block_n: usize = fmt.hadamard_dim; + +/// lm_head as stored: trellis int16 [K/16, N/16, 8 * halves] (bytes, possibly unaligned), suh fp16 [K], svh fp16 [N]. +pub const Head = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u16, + + pub fn words(h: Head) usize { + return @as(usize, h.halves) * 8; + } + + /// From the stored tensors; `mul1` (when present) must hold the mul1 multiplier. + pub fn of(trellis: core.safetensors.Tensor, suh: core.safetensors.Tensor, svh: core.safetensors.Tensor, mul1: ?core.safetensors.Tensor) !Head { + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.dim(2) % 8 != 0) return error.UnexpectedTensor; + const h: Head = .{ .trellis = trellis.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = trellis.dim(0) * 16, .n = trellis.dim(1) * 16, .halves = @intCast(trellis.dim(2) / 8) }; + if (suh.dtype != .f16 or suh.numel() != h.k or svh.dtype != .f16 or svh.numel() != h.n) return error.UnexpectedTensor; + if (h.k % block_n != 0 or h.n % block_n != 0 or h.halves < 2 or h.halves > 16) return error.UnexpectedTensor; + if (mul1) |m| { + if (m.bytes.len != 4 or std.mem.readInt(u32, m.bytes[0..4], .little) != fmt.Mul1Mul) return error.UnexpectedTensor; + } + return h; + } +}; + +/// f64 -> bf16 bits, one rounding to nearest even (no f32 step, so no double rounding). +pub fn bf16Of(v: f64) u16 { + const near = wt.toBf16(@floatCast(v)); + if (std.math.isNan(v) or std.math.isInf(v)) return near; + var best = near; + var best_d = @abs(@as(f64, wt.promote(near)) - v); + for ([_]u16{ near -% 1, near +% 1 }) |cand| { + const cv: f64 = wt.promote(cand); + if (std.math.isNan(cv) or std.math.isInf(cv)) continue; + if (std.math.signbit(cv) != std.math.signbit(@as(f64, wt.promote(near))) and cv != 0) continue; + const d = @abs(cv - v); + if (d < best_d or (d == best_d and cand & 1 == 0)) { + best = cand; + best_d = d; + } + } + return best; +} + +const Work = struct { + head: Head, + ids: []const u32, + /// Indices into `ids`, sorted by id. + order: []const u32, + /// Start of each block's run in `order`; `blocks` entries plus one. + starts: []const u32, + out: []u16, + next: std.atomic.Value(usize) = .init(0), + failed: ?anyerror = null, +}; + +fn copyF16(bytes: []const u8, out: []f16) void { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); +} + +fn worker(w: *Work, gpa: std.mem.Allocator) void { + work(w, gpa) catch |err| { + w.failed = err; + }; +} + +fn work(w: *Work, gpa: std.mem.Allocator) !void { + const h = w.head; + const tiles_k = h.k / 16; + const tiles_n = h.n / 16; + const per = block_n / 16; + const words = h.words(); + const sub = try gpa.alloc(i16, tiles_k * per * words); + defer gpa.free(sub); + const vals = try gpa.alloc(f64, h.k * block_n); + defer gpa.free(vals); + const scratch = try gpa.alloc(f64, 2 * h.k * block_n); + defer gpa.free(scratch); + const suh = try gpa.alloc(f16, h.k); + defer gpa.free(suh); + copyF16(h.suh, suh); + var svh: [block_n]f16 = undefined; + while (true) { + const b = w.next.fetchAdd(1, .monotonic); + if (b + 1 >= w.starts.len) return; + const lo = w.starts[b]; + const hi = w.starts[b + 1]; + if (lo == hi) continue; + const nb = w.ids[w.order[lo]] / block_n; + // the block's tiles, k-major as a [tiles_k, per] trellis of their own + for (0..tiles_k) |tk| for (0..per) |j| { + const src = ((tk * tiles_n) + nb * per + j) * words * 2; + const dst = sub[(tk * per + j) * words ..][0..words]; + for (dst, 0..) |*d, i| d.* = std.mem.readInt(i16, h.trellis[src + 2 * i ..][0..2], .little); + }; + copyF16(h.svh[nb * block_n * 2 ..][0 .. block_n * 2], &svh); + fmt.dequantize(.{ .halves = h.halves }, .mul1, sub, tiles_k, per, suh, &svh, vals, scratch); + for (w.order[lo..hi]) |i| { + const col = w.ids[i] - nb * block_n; + const row = w.out[@as(usize, i) * h.k ..][0..h.k]; + for (row, 0..) |*o, k| o.* = bf16Of(vals[k * block_n + col]); + } + } +} + +/// Row i of `out` ([ids.len, K] bf16) is lm_head's column `ids[i]` (the weight row of token ids[i]), decoded in float64 +/// by exl3_format.dequantize on the 128-column block that holds it and rounded once to bf16. Up to `threads` threads. +pub fn draftRows(gpa: std.mem.Allocator, head: Head, ids: []const u32, out: []u16, threads: usize) !void { + if (out.len != ids.len * head.k) return error.UnexpectedTensor; + for (ids) |id| if (id >= head.n) return error.InvalidToken; + const order = try gpa.alloc(u32, ids.len); + defer gpa.free(order); + for (order, 0..) |*o, i| o.* = @intCast(i); + std.mem.sort(u32, order, ids, struct { + fn lt(x: []const u32, a: u32, b: u32) bool { + return x[a] < x[b]; + } + }.lt); + // runs of one block each + var starts: std.ArrayList(u32) = .empty; + defer starts.deinit(gpa); + var prev: ?usize = null; + for (order, 0..) |i, at| { + const nb = ids[i] / block_n; + if (prev == null or prev.? != nb) try starts.append(gpa, @intCast(at)); + prev = nb; + } + try starts.append(gpa, @intCast(order.len)); + var w: Work = .{ .head = head, .ids = ids, .order = order, .starts = starts.items, .out = out }; + const n_threads = @max(1, @min(threads, starts.items.len - 1)); + if (n_threads == 1) { + try work(&w, gpa); + return; + } + var pool: [64]std.Thread = undefined; + var spawned: usize = 0; + defer for (pool[0..spawned]) |t| t.join(); + for (0..@min(n_threads, pool.len)) |_| { + pool[spawned] = try std.Thread.spawn(.{}, worker, .{ &w, gpa }); + spawned += 1; + } + for (pool[0..spawned]) |t| t.join(); + spawned = 0; + if (w.failed) |err| return err; +} + +/// Threads for a load: the machine's cores, at most 32 (each holds about 8 MB of float64 scratch). +pub fn loadThreads() usize { + return @min(32, std.Thread.getCpuCount() catch 1); +} + +/// The draft head of an EXL3 pack (null for any other): bf16 rows uploaded as a dense projection, or affine4 rows. +pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []const u32, affine: bool) !?wt.Q4 { + if (!store.has("lm_head.trellis")) return null; + const head = try Head.of(try store.hold("lm_head.trellis"), try store.hold("lm_head.suh"), try store.hold("lm_head.svh"), if (store.has("lm_head.mul1")) try store.hold("lm_head.mul1") else null); + if (head.k != wt.dims or head.n != wt.vocab) return error.UnexpectedTensor; + const rows = try a.alloc(u16, ids.len * wt.dims); + defer a.free(rows); + try draftRows(a, head, ids, rows, loadThreads()); + const bytes = std.mem.sliceAsBytes(rows); + if (!affine) return try projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ ids.len, wt.dims, 1, 1, 1 }, .bytes = bytes }); + const words = try a.alloc(u8, ids.len * wt.dims / 2); + defer a.free(words); + const scales = try a.alloc(u8, ids.len * wt.dims / 16); + defer a.free(scales); + const biases = try a.alloc(u8, scales.len); + defer a.free(biases); + for (0..ids.len) |i| try @import("cuda_draft_affine.zig").row(bytes[i * wt.dims * 2 ..][0 .. wt.dims * 2], words[i * wt.dims / 2 ..][0 .. wt.dims / 2], scales[i * wt.dims / 16 ..][0 .. wt.dims / 16], biases[i * wt.dims / 16 ..][0 .. wt.dims / 16]); + return try wt.uploadQ4(a, d, words, scales, biases, ids.len, wt.dims / 8); +} + +// ---- tests -------------------------------------------------------------------------------------------------------- + +test "exl3 draft: bf16Of rounds float64 once, to nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), bf16Of(1.0)); + try std.testing.expectEqual(@as(u16, 0x3F80), bf16Of(1.0 + 1.0 / 256.0)); // tie: even (down) + try std.testing.expectEqual(@as(u16, 0x3F82), bf16Of(1.0 + 3.0 / 256.0)); // tie: even (up) + // just above a tie after an f32 step would round to the tie: float64 sees it above + const above = 1.0 + 1.0 / 256.0 + 1.0e-12; + try std.testing.expectEqual(@as(u16, 0x3F81), bf16Of(above)); + try std.testing.expectEqual(@as(u16, 0xBF80), bf16Of(-1.0)); + try std.testing.expectEqual(@as(u16, 0), bf16Of(0)); +} + +fn synthHead(gpa: std.mem.Allocator, r: std.Random, k: usize, n: usize, halves: u16) !struct { t: []i16, su: []f16, sv: []f16, head: Head } { + const t = try gpa.alloc(i16, (k / 16) * (n / 16) * halves * 8); + for (t) |*x| x.* = @bitCast(r.int(u16)); + const su = try gpa.alloc(f16, k); + for (su) |*x| x.* = @floatCast((r.float(f32) * 0.02 + 0.005) * (if (r.boolean()) @as(f32, 1) else -1)); + const sv = try gpa.alloc(f16, n); + for (sv) |*x| x.* = @floatCast((r.float(f32) * 0.5 + 0.5) * (if (r.boolean()) @as(f32, 1) else -1)); + return .{ .t = t, .su = su, .sv = sv, .head = .{ .trellis = std.mem.sliceAsBytes(t), .suh = std.mem.sliceAsBytes(su), .svh = std.mem.sliceAsBytes(sv), .k = k, .n = n, .halves = halves } }; +} + +test "exl3 draft: rows equal the whole matrix's float64 decode rounded to bf16, any id order, any thread count" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(0xd4af); + const r = prng.random(); + const k = 256; + const n = 384; + const s = try synthHead(gpa, r, k, n, 10); + defer gpa.free(s.t); + defer gpa.free(s.su); + defer gpa.free(s.sv); + const full = try gpa.alloc(f64, k * n); + defer gpa.free(full); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + fmt.dequantize(.{ .halves = 10 }, .mul1, s.t, k / 16, n / 16, s.su, s.sv, full, scratch); + const ids = [_]u32{ 300, 5, 127, 128, 383, 5, 0, 200, 255, 256 }; + for ([_]usize{ 1, 3 }) |threads| { + var rows: [ids.len * k]u16 = undefined; + try draftRows(gpa, s.head, &ids, &rows, threads); + for (ids, 0..) |id, i| for (0..k) |kk| { + try std.testing.expectEqual(bf16Of(full[kk * n + id]), rows[i * k + kk]); + }; + } + var bad: [k]u16 = undefined; + try std.testing.expectError(error.InvalidToken, draftRows(gpa, s.head, &.{384}, &bad, 1)); +} + +test "exl3 draft: a draft logit is the full head's within the bf16 weight bound" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(0x10917); + const r = prng.random(); + const k = 256; + const n = 256; + const s = try synthHead(gpa, r, k, n, 6); + defer gpa.free(s.t); + defer gpa.free(s.su); + defer gpa.free(s.sv); + const full = try gpa.alloc(f64, k * n); + defer gpa.free(full); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + fmt.dequantize(.{ .halves = 6 }, .mul1, s.t, k / 16, n / 16, s.su, s.sv, full, scratch); + var x: [k]f64 = undefined; + for (&x) |*v| v.* = (r.float(f64) - 0.5) * 4; + var ids: [n]u32 = undefined; + for (&ids, 0..) |*id, i| id.* = @intCast(i); + const rows = try gpa.alloc(u16, n * k); + defer gpa.free(rows); + try draftRows(gpa, s.head, &ids, rows, 2); + for (0..n) |v| { + var exact: f64 = 0; + var mag: f64 = 0; + var got: f64 = 0; + for (0..k) |kk| { + exact += x[kk] * full[kk * n + v]; + mag += @abs(x[kk] * full[kk * n + v]); + got += x[kk] * @as(f64, wt.promote(rows[v * k + kk])); + } + try std.testing.expect(@abs(got - exact) <= mag / 512.0 + 1e-12); + } +} diff --git a/zig/src/families/flashnext/cuda_exl3_pass.zig b/zig/src/families/flashnext/cuda_exl3_pass.zig new file mode 100644 index 0000000000..4ca9de47db --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_pass.zig @@ -0,0 +1,468 @@ +//! The batched pass's EXL3 branches (cuda_batch.zig): `lane` for a Projection whose rows are EXL3 parts, `moe` for a +//! layer whose experts are EXL3. `p` is the pass (`*cuda_batch.Pass(Tri)`): its stream, kernels and scratch. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_exl3_load.zig"); +const pool_ = @import("cuda_exl3_pool.zig"); +const dplan = @import("cuda_exl3_dplan.zig"); +const memory = @import("exl3_startup_memory.zig"); +const line_deltas = @import("exl3_line_deltas.zig"); +const profile_line_budget = @import("exl3_profile_line_budget.zig"); + +const dims = wt.dims; +const slots = wt.slots; + +/// Rows a routed-experts call takes: its scratch grows with rows * 11 (fp16 inputs, fp32 split sums). +pub const moe_rows: usize = 256; + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// One part's rows [m, part.n] into `y` (row pitch part.n), bf16 or fp32. +fn part(p: anytype, ops: *load.Ops, pt: load.Part, x: u64, k: usize, m: usize, y: u64, f32_out: bool) !void { + switch (pt.kind) { + .exl3 => |w| { + const de = try ops.linear(); + if (try @import("cuda_exl3_prefill.zig").tilesPart(de, p.stream, w, x, m, y, f32_out)) return; // prefill rows on #806's tiles + try de.decode(p.stream, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); + }, + .native => |w| { + if (f32_out) return p.k.native_lane.matmulF32(p.stream, x, k, m, w, y, p.s.native_part); + try p.k.native_lane.matmul(p.stream, x, k, m, w, y, p.s.native_part); + }, + } +} + +/// Pass.lane() for an EXL3 projection: out bf16 [m, n] (sk 1), or fp32 [m, n] in part's first slice with the other +/// sk - 1 slices zeroed (the slice-summing callers' contract, as the native lane keeps it). +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.x == .f16) return laneF16(p, x, e, out, slices, m, sk); + if (sk > 1) { + if (e.parts.len != 1) return error.UnexpectedTensor; + try part(p, e.ops, e.parts[0], x, e.k, m, slices, true); + const bytes = m * e.n * 4; + try s.d.check(s.d.api.cuMemsetD8Async(slices + bytes, 0, bytes * (sk - 1), s.handle), "cuMemsetD8Async"); + return; + } + if (try @import("cuda_exl3_fork.zig").lane(p, x, e, out, m, onePart)) return; // the parts on parallel streams + for (e.parts) |pt| try onePart(p, e, pt, x, m, out); +} + +/// One part of a bf16 sk-1 lane() on the pass's stream: one part over the whole row, or a single row, writes its +/// columns in place; otherwise through the stitch buffer. +fn onePart(p: anytype, e: *const load.Proj, pt: load.Part, x: u64, m: usize, out: u64) !void { + const s = p.stream; + if (pt.n == e.n or m == 1) return part(p, e.ops, pt, x, e.k, m, out + pt.col * 2, false); + const tmp = try e.ops.scratch(s, &e.ops.tmp, m * pt.n * 2); + try part(p, e.ops, pt, x, e.k, m, tmp, false); + try copy2d(s, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); +} + +/// fp16 rows through a dense BF16 lane weight: x = hi + lo exactly (fn_exl3_split_f16), hi . w and lo . w in fp32, +/// summed in fp32: bf16 [m, n] for sk 1 (fn_sum_slices), else the two products as K slices 0 and 1, the rest zeroed. +fn laneF16(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.parts.len != 1 or e.parts[0].kind != .native) return error.UnexpectedTensor; + const w = e.parts[0].kind.native; + const count = m * e.k; + const halves = try e.ops.scratch(s, &e.ops.halves, 2 * count * 2); + var a: cuda.Args = .{}; + a.add(x); + a.add(halves); + a.add(halves + count * 2); + a.add(@as(c_longlong, @intCast(count))); + try cuda.launch.launch(try e.ops.split(), .{ .grid = .{ .x = @intCast((count + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + const bytes = m * e.n * 4; + const sums = if (sk == 1) try e.ops.scratch(s, &e.ops.sums, 2 * bytes) else slices; + try p.k.native_lane.matmulF32(s, halves, e.k, m, w, sums, p.s.native_part); + try p.k.native_lane.matmulF32(s, halves + count * 2, e.k, m, w, sums + bytes, p.s.native_part); + if (sk == 1) return p.sumSlices(sums, out, true, m * e.n, 2); + if (sk > 2) try s.d.check(s.d.api.cuMemsetD8Async(slices + 2 * bytes, 0, bytes * (sk - 2), s.handle), "cuMemsetD8Async"); +} + +/// Pass.moeBlock() for EXL3 experts: the router's 513 logits (the BF16 rows the Store widened exactly), each row's ten +/// picks and the shared slot (id 512), then every slot's expert output, fp32 y [rows, 11, dims], in blocks of rows. +/// Offloaded layers cut each block in two: segment A (block, router, top-k) ends with the picks on the host, the pool +/// plans hits and misses, and segment B (copies on the copy engine, then the experts on its tables) follows. +pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + if (x.tp != null) return @import("cuda_exl3_tp.zig").moe(p, lw, x); // expert parallel (plan step 10) + const s = p.s; + const n = p.rows; + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try p.tri.topkMany(s.router_l, s.picks, s.wts, n); + if (x.needsPool()) return error.Exl3PoolNotAttached; + const ex = try x.ops.grouped(); + const block = @min(n, moe_rows); + const bytes = try cuda.exl3_experts.scratchBytes(x.layer, block, slots); + const base = try x.ops.scratch(p.stream, &x.ops.moe, bytes); + var r0: usize = 0; + while (r0 < n) : (r0 += block) { + const rows = @min(block, n - r0); + var sc = try cuda.exl3_experts.scratch(base, x.layer, block, slots); + // runSlots leaves each slot's fp32 output in sc.y [rows * slots, dims]: the pass's y rows are that layout + sc.y = s.y + r0 * slots * dims * 4; + const xin = s.mixed + r0 * dims * 2; + const picks = s.picks + r0 * slots * 4; + if (x.pool) |pool| { + if (dplan.active) |dp| if (dp.indexOf(pool)) |i| { // a decode round planned on the device (no host cut) + try dplan.runLayer(dp, p.stream, i, ex, x.layer, xin, picks, rows, sc, 0, 0); + continue; + }; + try pooled(p.stream, x.ops, pool, ex, x.layer, xin, picks, rows, sc, 0, 0); + continue; + } + try ex.runSlots(p.stream, x.layer, xin, dims, .bf16, picks, rows, sc); + } +} + +/// Host time of the A/B cut (the picks D2H and its wait) and of Pool.plan, summed over every pooled block (read by +/// tf-cuda-test flashnext-exl3-engine). +/// Offloaded expert layers of this process's engine (set when the pools are built); 0: none. +pub var offloaded_layers: usize = 0; + +/// Drafted rounds verify more rows, need more distinct experts and lose to plain decode on offloaded experts +/// (measured on two RTX 5060 Ti: about 21% slower on real text): requests decode plain unless TF_FLASHNEXT_EXL3_DRAFTS=1. +pub fn plainByDefault() bool { + if (offloaded_layers == 0) return false; + const value = std.c.getenv("TF_FLASHNEXT_EXL3_DRAFTS") orelse return true; + return !std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub var pool_time: struct { cut_ns: u64 = 0, plan_ns: u64 = 0, blocks: u64 = 0 } = .{}; + +/// One block of `rows` rows through `pool`: per-slot outputs in sc.y, and with `wts` != 0 the last pass also combines +/// every slot (sc.y keeps the earlier passes' slots; a slot no pass ran reads its stale, finite value at weight 0) into +/// fp32 `out` [rows, dims] (expert parallel's rank branch, cuda_exl3_tp.moe). +pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, ex: *const cuda.exl3_experts.Experts, layer: cuda.exl3_experts.Layer, xin: u64, picks: u64, rows: usize, sc: cuda.exl3_experts.Scratch, wts: u64, out: u64) !void { + const host = (ops.picks orelse return error.Exl3ScratchNotReserved).slice(i32)[0 .. rows * slots]; + // the device planned the last decode rounds of this pool: its residency back to the host pools first (the MTP + // layer's pool is never device-planned: a draft step leaves the decoder pools on the device) + if (dplan.current()) |dp| if (dp.indexOf(pool) != null) try dp.toHost(); + const io = pool.io; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + // end of segment A: the block's picks reach the host (the only host wait of the layer) + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host.ptr, picks, rows * slots * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); + const r = try pool.plan(host, rows); + pool_time.cut_ns += @intCast(t1 - t0); + pool_time.plan_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t1); + pool_time.blocks += 1; + // segment B, a pass at a time: its copies and table rewrites, then the experts on the pool's tables + for (0..r.n_passes) |j| { + try pool.stage(r, j, st); + const last = j + 1 == r.n_passes; + try ex.run(st, layer, xin, dims, .bf16, pool.picksFor(r, j, picks), if (last) wts else 0, rows, if (last) out else 0, sc); + try pool.release(r, j, st); + } +} + +/// Every offloaded experts layer of a loaded engine (the 48 decoder layers, then the MTP layer if any). +pub fn offloadedLayers(layers: []const @import("cuda_prompt_load.zig").Layer, mtp: ?*const @import("cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.needsPool()) { + out[n] = x; + n += 1; + }; + if (mtp) |l| if (l.exl3) |x| if (x.needsPool()) { + out[n] = x; + n += 1; + }; + return n; +} + +pub const Budget = struct { + graph_reservation: u64 = 0, + planner_rows: ?u64 = null, + /// Device bytes left free after the pools; null: the native server's own reserve (TENSORFOLD_MEMORY_RESERVE_GIB, + /// else the larger of 1 GiB and a 32nd of the card, native/cuda_memory.zig reserveBytes) plus `margin`. + reserve: ?u64 = null, + /// What the pools leave beyond the server's reserve: allocator rounding and the copy streams' own memory. + margin: u64 = 256 << 20, + /// Bytes the family still allocates after the pools: each admitted stream's caches (Loaded.stream_bytes) times + /// the streams the pools must leave room for. + other: u64 = 0, + /// Lines a layer, or null: from free memory (cuda_exl3_round.linesFor). + lines: ?usize = null, + /// Rows a planned round holds (the pass's block of routed rows). + max_rows: usize = moe_rows, +}; + +/// TF_FLASHNEXT_EXL3_LINES and TF_FLASHNEXT_EXL3_RESERVE_MB override pool sizing. `requested_streams` is the server's +/// actual concurrency; standalone test callers may omit it and use TF_FLASHNEXT_EXL3_STREAMS (default 1). +pub fn budgetFromEnv(stream_bytes: u64) !Budget { + return budgetFromEnvWithStreams(stream_bytes, null); +} + +pub fn budgetForStreams(stream_bytes: u64, requested_streams: u32) !Budget { + return budgetFromEnvWithStreams(stream_bytes, requested_streams); +} + +fn budgetFromEnvWithStreams(stream_bytes: u64, requested_streams: ?u32) !Budget { + const streams = requested_streams orelse 1; + var b: Budget = .{ .other = try streamReservation(stream_bytes, streams) }; + if (dplan.enabled()) { + b.planner_rows = if (std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_ROWS")) |v| @max(1, try std.fmt.parseInt(u64, std.mem.span(v), 10)) else 16; + } + if (std.c.getenv("TF_FLASHNEXT_EXL3_LINES")) |v| b.lines = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (std.c.getenv("TF_FLASHNEXT_EXL3_RESERVE_MB")) |v| b.reserve = try memory.mebibytes(std.mem.span(v)); + if (requested_streams == null) if (std.c.getenv("TF_FLASHNEXT_EXL3_STREAMS")) |v| { + const env_streams = try std.fmt.parseInt(u64, std.mem.span(v), 10); + if (env_streams == 0) return error.InvalidStreamCount; + b.other = try std.math.mul(u64, stream_bytes, env_streams); + }; + return b; +} + +fn streamReservation(stream_bytes: u64, streams: u32) !u64 { + if (streams == 0) return error.InvalidStreamCount; + return std.math.mul(u64, stream_bytes, streams); +} + +test "EXL3 pool reserves the requested number of concurrent stream caches" { + try std.testing.expectEqual(@as(u64, 200), try streamReservation(100, 2)); + try std.testing.expectError(error.InvalidStreamCount, streamReservation(100, 0)); +} + +pub const Attached = struct { layers: usize = 0, lines: usize = 0, line_bytes: usize = 0, device_bytes: u64 = 0, free_before: u64 = 0, free_after: u64 = 0, scratch_bytes: u64 = 0, graph_reservation: u64 = 0, after_pool_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0, prefill_window_bytes: u64 = 0 }; + +pub fn freeBytes(d: *const cuda.Driver) !u64 { + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +/// The sizing hook (cuda_lanes_load.zig, once everything else of the family is allocated): reserves every EXL3 +/// scratch of `rows`-row passes, then gives each offloaded layer the same number of lines from what is free after +/// `b.reserve` and `b.other`, with one lifetime use table shared by the decoder layers (the MTP layer keeps its own). +/// A no-op for a checkpoint without offloaded experts. +pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, layers: []const @import("cuda_prompt_load.zig").Layer, mtp: ?*const @import("cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { + var all: [wt.layers_n + 1]*load.Experts = undefined; + const n = offloadedLayers(layers, mtp, &all); + const delta_text = if (std.c.getenv("TF_FLASHNEXT_EXL3_LINE_DELTAS")) |value| std.mem.span(value) else null; + var out: Attached = .{}; + const scratch_before = try freeBytes(d); + out.graph_reservation = b.graph_reservation; + // every distinct Ops (the engine's store and the MTP head's) reserves its scratch now + var seen: [4]*load.Ops = undefined; + var n_seen: usize = 0; + for (layers) |*l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + if (mtp) |l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + // expert parallel (cuda_exl3_tp.zig): its masks, sums and block scratch, before the pools take what is free + for (layers) |*l| if (l.exl3) |x| if (x.tp != null) { + try @import("cuda_exl3_tp.zig").reserve(x.ops, st, x.layer, rows); + break; + }; + if (n == 0) { + if (delta_text != null) return error.InvalidLineDeltaIndex; + return out; + } + for (all[0..n]) |x| { + out.line_bytes = @max(out.line_bytes, x.lineBytes()); + out.bank_bytes += x.bank.?.bytes.len; + out.bank_ns += x.bank_ns; + } + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + out.free_before = free; + out.scratch_bytes = scratch_before -| free; + const reserve_floor = b.reserve orelse serverReserve(total); + const reserve = b.reserve orelse try std.math.add(u64, reserve_floor, b.margin); + // beside its lines a layer holds the shared expert's region, its six tables and its masked picks (sized for the + // fewest lines, the most passes); counted as fixed bytes before the lines share what is left + var fixed: u64 = 0; + for (all[0..n]) |x| { + for (x.shared) |t| fixed = try std.math.add(u64, fixed, try std.math.mul(u64, (try std.math.add(u64, t.len, 255)) / 256, 256)); + const picks_bytes = try std.math.mul(u64, try std.math.mul(u64, pool_.passesFor(pool_.min_lines), b.max_rows), slots * 4); + fixed = try std.math.add(u64, fixed, try std.math.add(u64, cuda.exl3_pool.Tables.init(wt.experts_n + 1).bytes + (64 << 10), picks_bytes)); + } + var future = try std.math.add(u64, b.other, b.graph_reservation); + const planner_fixed = if (b.planner_rows) |planner_rows| try memory.plannerBytes(n, 0, planner_rows) else 0; + const line_cost = try std.math.add(u64, out.line_bytes, if (b.planner_rows != null) 8 else 0); + const all_fixed = try std.math.add(u64, fixed, planner_fixed); + const prefill = @import("cuda_exl3_prefill.zig"); + const minimum_pools = try memory.poolBytes(b.lines orelse pool_.min_lines, line_cost, n, all_fixed); + const protected = try std.math.add(u64, try std.math.add(u64, reserve, future), minimum_pools); + out.prefill_window_bytes = prefill.reserveTwoChunk(d, free, protected, b.graph_reservation) catch |err| declined: { + std.log.warn("flashnext exl3 prefill: two-chunk reservation refused ({t}); keeping single-chunk prefill", .{err}); + break :declined 0; + }; + const graph_memory = @import("cuda_exl3_graph_memory.zig"); + const window_reservation = graph_memory.twoChunkReservationBytes(); + if (!prefill.twoChunkAllocated() and window_reservation != 0 and b.graph_reservation == window_reservation) { + out.graph_reservation = try graph_memory.declineTwoChunkReservation(); + future = try std.math.add(u64, b.other, out.graph_reservation); + } + if (out.prefill_window_bytes != 0) { + free = @intCast(try freeBytes(d)); + out.free_before = free; + } + const pool_room = try memory.available(free, reserve, future); + const lines = b.lines orelse try pool_.linesFor(free, reserve, try std.math.add(u64, future, all_fixed), line_cost, n, pool_.min_lines); + if (lines < pool_.min_lines) return error.PoolTooSmall; + const startup_bytes = try memory.poolBytes(lines, line_cost, n, all_fixed); + if (startup_bytes > pool_room) return error.InsufficientStartupMemory; + out.after_pool_bytes = if (b.planner_rows) |planner_rows| try memory.plannerBytes(n, lines, planner_rows) else 0; + out.lines = lines; + var pool_lines: [wt.layers_n + 1]usize = @splat(lines); + var line_bytes: [wt.layers_n + 1]usize = undefined; + var profile_sized = false; + var profile_min_lines: usize = std.math.maxInt(usize); + var profile_max_lines: usize = 0; + var decoder_pool_count: usize = 0; + for (all[0..n], 0..) |x, index| { + line_bytes[index] = x.lineBytes(); + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (!is_mtp) decoder_pool_count += 1; + } + if (delta_text == null and shouldProfileSize(profileSizingEnabled(), b.lines)) if (profile_counts) |prior| { + var count_rows: [wt.layers_n][]const u32 = undefined; + var allocation_costs: [wt.layers_n]usize = undefined; + var allocation_lines: [wt.layers_n]usize = undefined; + var decoder_index: usize = 0; + var has_profile_signal = false; + for (all[0..n], 0..) |x, pool_index| { + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (is_mtp) continue; + for (prior[decoder_index]) |count| { + has_profile_signal = has_profile_signal or count != 0; + } + count_rows[decoder_index] = &prior[decoder_index]; + allocation_costs[decoder_index] = try std.math.add(usize, line_bytes[pool_index], if (b.planner_rows != null) 8 else 0); + decoder_index += 1; + } + if (has_profile_signal and decoder_index == wt.layers_n) { + const rank = all[0].tp orelse try @import("cuda_exl3_tp.zig").Rank.init(1, 0); + const maybe_plan: ?profile_line_budget.Budget = profile_line_budget.allocate(gpa, .{ + .counts = &count_rows, + .rank = rank.rank, + .world = rank.world, + .line_bytes = &allocation_costs, + .uniform_lines = lines, + .min_lines = pool_.min_lines, + .max_rows = b.max_rows, + .slots = slots, + }, &allocation_lines) catch |err| blk: { + std.log.info("flashnext exl3: profile-sized pools unavailable ({s}); retaining uniform decoder pools", .{@errorName(err)}); + break :blk null; + }; + if (maybe_plan) |plan| { + decoder_index = 0; + for (all[0..n], 0..) |x, pool_index| { + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (is_mtp) continue; + pool_lines[pool_index] = allocation_lines[decoder_index]; + profile_min_lines = @min(profile_min_lines, allocation_lines[decoder_index]); + profile_max_lines = @max(profile_max_lines, allocation_lines[decoder_index]); + decoder_index += 1; + } + if (plan.uniform_fallback) { + @memset(pool_lines[0..n], lines); + std.log.info("flashnext exl3: profile candidate changed total variable CUDA bytes; retaining uniform decoder pools", .{}); + } else { + out.after_pool_bytes = if (b.planner_rows) |planner_rows| try memory.plannerBytesForPools(pool_lines[0..n], planner_rows) else 0; + profile_sized = true; + } + std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines, per-line planner state and device picks); saved lifetime picks are a proxy, not latency; MTP remains uniform", .{ plan.allocated_bytes, plan.baseline_bytes }); + std.log.info("flashnext exl3: profile decoder pool sizes range {d}..{d} lines; uniform-equivalent baseline was {d}", .{ profile_min_lines, profile_max_lines, lines }); + } + } + }; + if (delta_text) |text| { + const plan = try line_deltas.apply(text, lines, decoder_pool_count, line_bytes[0..n], pool_.min_lines, b.planner_rows, fixed, startup_bytes, pool_lines[0..n]); + out.after_pool_bytes = plan.planner_bytes; + std.log.info("flashnext exl3: experimental rank-local decoder pool line deltas {s}; base={d} planner_bytes={d} startup_bytes={d}", .{ text, lines, plan.planner_bytes, plan.pool_bytes }); + } + // the first decoder layer's pool owns the lifetime use table the other decoder layers share + var counts = profile_counts; + var decoder: u8 = 0; + for (all[0..n], pool_lines[0..n]) |x, count| { + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + try x.attach(io, d, if (is_mtp) 0 else decoder, count, b.max_rows, if (is_mtp) null else counts); + if (is_mtp) dplan.mtp_pool = x.pool.?; // the draft head's experts plan on the device too (cuda_exl3_chain.zig) + if (!is_mtp and counts == null) counts = x.pool.?.counts; + if (!is_mtp) decoder += 1; + out.device_bytes += x.pool.?.deviceBytes(); + out.layers += 1; + } + out.free_after = try freeBytes(d); + _ = memory.available(out.free_after, reserve_floor, try std.math.add(u64, future, out.after_pool_bytes)) catch |err| { + std.log.err("EXL3 pool reservation shortfall: free_before={d} free_after={d} pool_bytes={d} floor={d} allocator_margin={d} future={d} planner={d}", .{ free, out.free_after, out.device_bytes, reserve_floor, b.margin, future, out.after_pool_bytes }); + return err; + }; + const allocation = if (delta_text != null) "uniform-equivalent lines (experimental deltas applied)" else if (profile_sized) "uniform-equivalent lines (profile-sized decoder pools)" else "lines a layer"; + offloaded_layers = out.layers; + std.log.info("flashnext exl3: {d} offloaded expert layers, {d} {s} of {d:.2} MiB ({d:.2} GiB on the device, {d:.2} GiB free before); pinned host bank {d:.2} GiB copied in {d:.1} s", .{ out.layers, out.lines, allocation, @as(f64, @floatFromInt(out.line_bytes)) / (1 << 20), @as(f64, @floatFromInt(out.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.free_before)) / (1 << 30), @as(f64, @floatFromInt(out.bank_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.bank_ns)) / 1e9 }); + return out; +} + +test { + _ = line_deltas; + _ = profile_line_budget; +} + +fn profileSizingEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_SIZING") orelse return false; + return shouldProfileSize(std.mem.eql(u8, std.mem.span(value), "1"), null); +} + +fn shouldProfileSize(enabled: bool, explicit_lines: ?usize) bool { + return enabled and explicit_lines == null; +} + +test "exl3: profile sizing requires explicit opt-in and does not override fixed line budgets" { + try std.testing.expect(!shouldProfileSize(false, null)); + try std.testing.expect(!shouldProfileSize(true, 32)); + try std.testing.expect(shouldProfileSize(true, null)); +} + +/// native/cuda_memory.zig reserveBytes for a discrete card: TENSORFOLD_MEMORY_RESERVE_GIB, else max(1 GiB, total / 32). +fn serverReserve(total: u64) u64 { + if (std.c.getenv("TENSORFOLD_MEMORY_RESERVE_GIB")) |v| { + const g = std.fmt.parseFloat(f64, std.mem.trim(u8, std.mem.span(v), " ")) catch 0; + if (g >= 0.5) return @intFromFloat(g * (1 << 30)); + } + return @max(@as(u64, 1 << 30), total / 32); +} + +fn reserveOnce(o: *load.Ops, seen: []*load.Ops, n: *usize, st: cuda.Stream, rows: usize) !void { + for (seen[0..n.*]) |x| if (x == o) return; + if (n.* == seen.len) return error.Invalid; + seen[n.*] = o; + n.* += 1; + try o.reserve(st, rows, moe_rows, slots); +} + +/// Pool statistics summed over every attached layer. +pub fn poolStats(layers: []const @import("cuda_prompt_load.zig").Layer, mtp: ?*const @import("cuda_prompt_load.zig").Layer) pool_.Stats { + var t: pool_.Stats = .{}; + var all: [wt.layers_n + 1]*load.Experts = undefined; + const n = offloadedPools(layers, mtp, &all); + for (all[0..n]) |x| { + const s = x.pool.?.stats; + inline for (@typeInfo(pool_.Stats).@"struct".field_names) |f| { + @field(t, f) += @field(s, f); + } + } + return t; +} + +pub fn offloadedPools(layers: []const @import("cuda_prompt_load.zig").Layer, mtp: ?*const @import("cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.pool != null) { + out[n] = x; + n += 1; + }; + if (mtp) |l| if (l.exl3) |x| if (x.pool != null) { + out[n] = x; + n += 1; + }; + return n; +} diff --git a/zig/src/families/flashnext/cuda_exl3_ple.zig b/zig/src/families/flashnext/cuda_exl3_ple.zig new file mode 100644 index 0000000000..e512077007 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_ple.zig @@ -0,0 +1,556 @@ +//! An EXL3 pack's n-gram table (ExLlamaV3 `ngram_embedding.safetensors`) behind `cuda_ple.Ngram` (`format = .exl3`). +//! +//! The pack ships the table outside model.safetensors.index.json, as one file: 128 shards `shard_{i}.trellis` int16 +//! [rows, 51] (one fp16 scale word, then 160 tail-biting mul1 values of 5 bits), `head_bias` fp16 [16, 160], and the hash +//! constants `layer_multipliers`, `head_vocab_sizes`, `head_offsets` (checked here against `cuda_ple.Hash`, which derives +//! them from the config). Row ids and shard placement are upstream's (`Hash`, `lastIds`, `Hash.locate`): on the real +//! pack they are byte-equal to the ids of tensorfold's qwen4_exp `NGram.ids` (tf-cuda-test exl3-ple2). +//! +//! A row decodes as tensorfold's qwen4_exp/cuda/exl3_mm.py `_ple_rows` (ExLlamaV3 `ngram_dequant`): +//! out[i] = fp16(f32(fp16 mul1(state_i)) * f32(scale) + f32(head_bias[head][i])). +//! The reference keeps these values fp16 and multiplies them by the fp16 key/value projections (exl3.py `ple_layer`), +//! so `Ngram.rows` writes fp16 bit patterns by default; `.emb = .bf16` rounds them once more for a bf16-only lane. +//! +//! Two readers behind the same `row`: the file mapping (`.map`, madvise RANDOM, nothing populated), or explicit preads +//! (`.pread`: the 30 GiB file is never mapped) through a 4096-row direct-mapped hot cache (512 KiB). +//! TF_FLASHNEXT_PLE_READ=pread and TF_FLASHNEXT_PLE_EMB=bf16 select them at load. + +const std = @import("std"); +const core = @import("core"); +const ple = @import("cuda_ple.zig"); +const fmt = core.exl3_format; + +pub const file_name = "ngram_embedding.safetensors"; +/// Values a row holds: one head's slice of the 2,560-wide embedding. +pub const row_values: usize = 160; +pub const hot_slots: usize = 4096; +/// Bytes a hot slot holds (a row at 8 bits is 162). +pub const hot_cap: usize = 176; + +pub const Read = enum { map, pread }; +pub const Emb = enum { f16, bf16 }; +pub const Options = struct { read: Read = .map, emb: Emb = .f16 }; + +/// TF_FLASHNEXT_PLE_READ (map | pread) and TF_FLASHNEXT_PLE_EMB (f16 | bf16); anything else keeps the default. +pub fn optionsFromEnv() Options { + var o: Options = .{}; + if (std.c.getenv("TF_FLASHNEXT_PLE_READ")) |v| { + if (std.mem.eql(u8, std.mem.span(v), "pread")) o.read = .pread; + } + if (std.c.getenv("TF_FLASHNEXT_PLE_EMB")) |v| { + if (std.mem.eql(u8, std.mem.span(v), "bf16")) o.emb = .bf16; + } + return o; +} + +// ---- the row codec (ported from exl3-flashnext-window-wip exl3_ngram.zig) ------------------------------------------ + +/// Int16 words in one row at `bits`: the fp16 scale, then 160 values. `bits` outside 2..8 is `NgramWidth`. +pub fn wordsPerRow(bits: usize) error{NgramWidth}!usize { + if (bits < 2 or bits > 8) return error.NgramWidth; + return 1 + row_values * bits / 16; +} + +/// Mul1 state of value `index`: bit m comes from position `index - m / bits` (plus 160 when negative), bit `m % bits`. +pub fn state(words: []const i16, bits: usize, index: usize) u16 { + var acc: u16 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(index)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(row_values)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u16 = @bitCast(words[1 + (sb >> 4)]); + acc |= ((word >> @as(u4, @intCast(sb & 15))) & 1) << @as(u4, @intCast(m)); + } + return acc; +} + +/// One row into 160 fp16 values: one rounding of `f32(codebook) * f32(scale) + f32(bias[i])`. +pub fn decodeRow(words: []const i16, bits: usize, bias: []const f16, out: []f16) error{NgramWidth}!void { + const width = try wordsPerRow(bits); + if (words.len != width or bias.len != row_values or out.len != row_values) return error.NgramWidth; + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + // `state(words, bits, i)` for every i at once: value j is stream bits [j * bits, (j + 1) * bits); state i is + // values i, i - 1, ... (tail-biting) packed `bits` apart from bit 0, the low 16 bits (the same bits as `state`) + var vals: [row_values]u32 = undefined; + const mask: u32 = (@as(u32, 1) << @intCast(bits)) - 1; + for (&vals, 0..) |*v, j| { + const sb = j * bits; + const wi = 1 + (sb >> 4); + const lo: u32 = @as(u16, @bitCast(words[wi])); + const hi: u32 = if (wi + 1 < words.len) @as(u16, @bitCast(words[wi + 1])) else 0; + v.* = ((lo | (hi << 16)) >> @intCast(sb & 15)) & mask; + } + for (out, bias, 0..) |*o, b, i| { + var acc: u32 = 0; + var k: usize = 0; + while (k * bits < 16) : (k += 1) acc |= vals[(i + row_values - k) % row_values] << @intCast(k * bits); + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, @truncate(acc))); + o.* = @floatCast(cb * scale + @as(f32, @floatCast(b))); + } +} + +test "exl3 ple: the packed decode's states equal `state` for every value at every width" { + var prng = std.Random.DefaultPrng.init(11); + for (2..9) |bits| { + const n = try wordsPerRow(bits); + var words: [1 + row_values * 8 / 16]i16 = undefined; + for (words[0..n]) |*w| w.* = @bitCast(prng.random().int(u16)); + words[0] = @bitCast(@as(f16, 0.5)); + var bias: [row_values]f16 = undefined; + for (&bias) |*b| b.* = @floatCast(prng.random().float(f32) - 0.5); + var got: [row_values]f16 = undefined; + try decodeRow(words[0..n], bits, &bias, &got); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (got, bias, 0..) |g, b, i| { + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, state(words[0..n], bits, i))); + const want: f16 = @floatCast(cb * scale + @as(f32, @floatCast(b))); + try std.testing.expectEqual(@as(u16, @bitCast(want)), @as(u16, @bitCast(g))); + } + } +} + +fn bf16Of(v: f32) u16 { + return ple.toBf16(v); +} + +/// Group sums of 32 fp16 values (the tree `cuda_ple.sumGroups` adds bf16 values in), for an fp16 row. +pub fn sumGroupsF16(row: []const u16, out: []f32) void { + for (0..row.len / 32) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @floatCast(@as(f16, @bitCast(row[g * 32 + i]))); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +// ---- the table ---------------------------------------------------------------------------------------------------- + +const c = std.c; + +/// The open table. Heap-allocated (`open`/`destroy`): the pread reader's hot cache is 512 KiB. +pub const Table = struct { + gpa: std.mem.Allocator, + io: std.Io, + opts: Options, + heads: usize, + shards: usize, + bits: usize, + words: usize, + /// File byte offset of each shard's row 0, and its row count. + offset: [128]u64 = undefined, + rows: [128]u64 = undefined, + bias: []f16, + /// `.map`: the whole file's mapping (header parsed, nothing populated). + file: ?core.safetensors.File = null, + /// `.pread`: the descriptor; the file is never mapped past `open`. + fd: c.fd_t = -1, + hot: ?*Hot = null, + + pub const Hot = struct { + lock: std.atomic.Value(u32) = .init(0), + lookups: u64 = 0, + misses: u64 = 0, + tags: [hot_slots]u64 = @splat(std.math.maxInt(u64)), + stored: [hot_slots][hot_cap]u8 = undefined, + + fn acquire(h: *Hot) void { + while (h.lock.cmpxchgWeak(0, 1, .acquire, .monotonic) != null) std.atomic.spinLoopHint(); + } + fn release(h: *Hot) void { + h.lock.store(0, .release); + } + }; + + fn prefix(buf: []u8, part: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}", .{ ple.layer, part }); + } + + /// The pack's table when `dir` holds an EXL3 n-gram file whose constants equal `hash`'s; null when there is no such file. + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, hash: ple.Hash, opts: Options) !?*Table { + const path = try std.fs.path.join(gpa, &.{ dir, file_name }); + defer gpa.free(path); + var file = core.safetensors.File.open(gpa, io, path) catch |err| switch (err) { + error.FileNotFound => return null, + else => return err, + }; + var keep_file = false; + defer if (!keep_file) file.close(io); + var nb: [200]u8 = undefined; + const first = file.get(try prefix(&nb, "shard_0.trellis")) orelse return null; + if (first.dtype != .i16 or first.rank != 2) return error.UnexpectedTensor; + if (hash.shards > 128 or hash.dh != row_values) return error.UnexpectedTensor; + // the hash constants the pack ships against the ones the config derives + { + var got: [ple.max_heads]i64 = undefined; + try ple.readI64(file.get(try prefix(&nb, "layer_multipliers")) orelse return error.MissingTensor, got[0..hash.n]); + if (!std.mem.eql(i64, got[0..hash.n], hash.mult[0..hash.n])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_vocab_sizes")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.sizes[0..hash.heads])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_offsets")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.offsets[0..hash.heads])) return error.NgramConstants; + } + // the reference's EXL3 path applies no table scale: refuse a file that carries one + if (file.get(try prefix(&nb, "weight_scale")) != null) return error.UnexpectedTensor; + const bias_t = file.get(try prefix(&nb, "head_bias")) orelse return error.MissingTensor; + if (!bias_t.is(.f16, &.{ hash.heads, row_values })) return error.UnexpectedTensor; + const words = first.dim(1); + const bits = (words - 1) * 16 / row_values; + if ((wordsPerRow(bits) catch return error.UnexpectedTensor) != words) return error.UnexpectedTensor; + + const t = try gpa.create(Table); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .io = io, .opts = opts, .heads = hash.heads, .shards = hash.shards, .bits = bits, .words = words, .bias = try gpa.alloc(f16, hash.heads * row_values) }; + errdefer gpa.free(t.bias); + @memcpy(std.mem.sliceAsBytes(t.bias), bias_t.bytes); + for (0..hash.shards) |i| { + var part: [32]u8 = undefined; + const e = file.names.get(try prefix(&nb, try std.fmt.bufPrint(&part, "shard_{d}.trellis", .{i}))) orelse return error.MissingTensor; + const want: usize = @intCast(hash.starts[i + 1] - hash.starts[i]); + if (e.dtype != .i16 or e.rank != 2 or e.shape[0] != want or e.shape[1] != words) return error.UnexpectedTensor; + t.offset[i] = file.data + e.begin; + t.rows[i] = want; + } + switch (opts.read) { + .map => { + if (@import("builtin").os.tag == .linux) { + if (file.map.section != null) _ = std.os.linux.madvise(file.map.memory.ptr, file.map.memory.len, std.os.linux.MADV.RANDOM); + } + t.file = file; + keep_file = true; + }, + .pread => { + const z = try gpa.dupeSentinel(u8, path, 0); + defer gpa.free(z); + const fd = c.open(z.ptr, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + t.fd = fd; + errdefer _ = c.close(fd); + if (@import("builtin").os.tag == .linux) _ = std.os.linux.fadvise(fd, 0, 0, std.os.linux.POSIX_FADV.RANDOM); + const hot = try gpa.create(Hot); + hot.* = .{}; + t.hot = hot; + }, + } + return t; + } + + pub fn destroy(t: *Table) void { + if (t.file) |*f| f.close(t.io); + if (t.fd >= 0) _ = c.close(t.fd); + if (t.hot) |h| t.gpa.destroy(h); + t.gpa.free(t.bias); + t.gpa.destroy(t); + } + + pub fn rowBytes(t: *const Table) usize { + return t.words * 2; + } + + fn preadAll(t: *const Table, dest: []u8, at: u64) !void { + var done: usize = 0; + while (done < dest.len) { + const n = c.pread(t.fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + if (n <= 0) return error.ReadFailed; + done += @intCast(n); + } + } + + /// The stored bytes of row `local` of shard `shard` (`rowBytes` long), through the mapping or the hot cache. + pub fn raw(t: *const Table, shard: usize, local: usize, dest: []u8) !void { + if (shard >= t.shards or local >= t.rows[shard] or dest.len != t.rowBytes()) return error.UnexpectedTensor; + const at = t.offset[shard] + @as(u64, local) * t.rowBytes(); + if (t.file) |f| { + @memcpy(dest, f.map.memory[at..][0..dest.len]); + return; + } + const h = t.hot.?; + const slot: usize = @intCast(((at *% 0x9E3779B97F4A7C15) >> 20) % hot_slots); + h.acquire(); + h.lookups += 1; + if (h.tags[slot] == at) { + @memcpy(dest, h.stored[slot][0..dest.len]); + h.release(); + return; + } + h.release(); + try t.preadAll(dest, at); + h.acquire(); + h.misses += 1; + h.tags[slot] = at; + @memcpy(h.stored[slot][0..dest.len], dest); + h.release(); + } + + /// Row `local` of shard `shard`, decoded with head `head`'s bias into 160 fp16 (or bf16) bit patterns. + pub fn row(t: *const Table, shard: usize, local: usize, head: usize, out: []u16) !void { + if (head >= t.heads or out.len != row_values) return error.UnexpectedTensor; + var bytes: [hot_cap]u8 align(2) = undefined; + const rb = t.rowBytes(); + try t.raw(shard, local, bytes[0..rb]); + const words: []const i16 = @as([*]const i16, @ptrCast(&bytes))[0..t.words]; + var vals: [row_values]f16 = undefined; + try decodeRow(words, t.bits, t.bias[head * row_values ..][0..row_values], &vals); + switch (t.opts.emb) { + .f16 => for (out, vals) |*o, v| { + o.* = @bitCast(v); + }, + .bf16 => for (out, vals) |*o, v| { + o.* = bf16Of(@floatCast(v)); + }, + } + } + + /// Group sums of one decoded row in the table's element type. + pub fn sums(t: *const Table, row_bits: []const u16, out: []f32) void { + switch (t.opts.emb) { + .f16 => sumGroupsF16(row_bits, out), + .bf16 => ple.sumGroups(row_bits, out), + } + } + + pub fn hotStats(t: *const Table) struct { lookups: u64, misses: u64 } { + const h = t.hot orelse return .{ .lookups = 0, .misses = 0 }; + return .{ .lookups = h.lookups, .misses = h.misses }; + } +}; + +// ---- tests -------------------------------------------------------------------------------------------------------- + +/// The Triton `_ple_rows` formulation, a second decoder that shares only the three fp16 constants with `decodeRow`. +fn decodeRowRef(words: []const i16, bits: usize, bias: []const f16, out: []f16) void { + const k_inv: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1ScaleBits))); + const k_bias: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1BiasBits))); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (0..out.len) |i| { + var st: u64 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(i)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(out.len)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u64 = @as(u16, @bitCast(words[1 + (sb >> 4)])); + st |= ((word >> @intCast(sb & 15)) & 1) << @intCast(m); + } + const prod: u64 = (st * 0x83DCD12D) & 0xFFFFFFFF; + const hs: u64 = (prod & 255) + ((prod >> 8) & 255) + ((prod >> 16) & 255) + ((prod >> 24) & 255); + const cb: f16 = @floatCast(@as(f32, @floatFromInt(1024 + hs)) * k_inv + k_bias); + out[i] = @floatCast(@as(f32, @floatCast(cb)) * scale + @as(f32, @floatCast(bias[i]))); + } +} + +test "exl3 ngram: wordsPerRow is one scale word plus 160 K-bit values" { + try std.testing.expectEqual(@as(usize, 21), try wordsPerRow(2)); + try std.testing.expectEqual(@as(usize, 31), try wordsPerRow(3)); + try std.testing.expectEqual(@as(usize, 51), try wordsPerRow(5)); + try std.testing.expectEqual(@as(usize, 81), try wordsPerRow(8)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(1)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(9)); +} + +test "exl3 ngram: patterned rows read tail-biting states" { + var storage: [81]i16 = undefined; + for (&storage, 0..) |*w, i| w.* = @bitCast(@as(u16, @intCast(i))); + try std.testing.expectEqual(@as(u16, 0x5001), state(storage[0..21], 2, 0)); + try std.testing.expectEqual(@as(u16, 0x9001), state(storage[0..31], 3, 0)); + try std.testing.expectEqual(@as(u16, 0x8001), state(storage[0..51], 5, 0)); + try std.testing.expectEqual(@as(u16, 0x0020), state(storage[0..51], 5, 1)); + try std.testing.expectEqual(@as(u16, 0x0400), state(storage[0..51], 5, 2)); + try std.testing.expectEqual(@as(u16, 0x0001), state(storage[0..81], 8, 0)); + // value 0 at 5 bits: bit 5 wraps to position 159, word 50 bit 11 + var words: [51]i16 = @splat(0); + words[50] = @bitCast(@as(u16, 1) << 11); + try std.testing.expectEqual(@as(u16, 1) << 5, state(&words, 5, 0)); +} + +test "exl3 ngram: scale and bias round in f32, then once to f16" { + var words: [51]i16 = @splat(0); + words[0] = @bitCast(@as(u16, 0x3C01)); + var bias: [160]f16 = @splat(0); + bias[0] = @bitCast(@as(u16, 0x1400)); + var out: [160]f16 = undefined; + try decodeRow(&words, 5, &bias, &out); + // codebook(0) is 0xC2E8; times 0x3C01 plus 0x1400 is 0xC2E9 (rounding the product first would give 0xC2EA) + try std.testing.expectEqual(@as(u16, 0xC2E8), @as(u16, @bitCast(fmt.codebookValue(.mul1, 0)))); + try std.testing.expectEqual(@as(u16, 0xC2E9), @as(u16, @bitCast(out[0]))); + try std.testing.expectError(error.NgramWidth, decodeRow(words[0..50], 5, &bias, &out)); + try std.testing.expectError(error.NgramWidth, decodeRow(&words, 2, &bias, &out)); +} + +test "exl3 ngram: decodeRow equals the Triton formulation on random rows at every width" { + var rng = std.Random.DefaultPrng.init(0x91e); + const r = rng.random(); + var words: [81]i16 = undefined; + var bias: [160]f16 = undefined; + var a: [160]f16 = undefined; + var b: [160]f16 = undefined; + for ([_]usize{ 2, 3, 4, 5, 6, 8 }) |bits| { + const n = try wordsPerRow(bits); + for (0..300) |_| { + for (words[0..n]) |*x| x.* = @bitCast(r.int(u16)); + words[0] = @bitCast(@as(f16, @floatCast((r.float(f32) - 0.5) * 4))); + for (&bias) |*x| x.* = @floatCast((r.float(f32) - 0.5) * 0.25); + try decodeRow(words[0..n], bits, &bias, &a); + decodeRowRef(words[0..n], bits, &bias, &b); + try std.testing.expectEqualSlices(u16, @ptrCast(&b), @ptrCast(&a)); + } + } +} + +/// A small table file plus its config: 2 shards, 4 heads of 160 (n-gram 3, 2 heads an n-gram), 5-bit rows. +const Synth = struct { + const vocab = 100; + const eos = 99; + const base = 11; + + fn spec() ple.Hash { + return ple.Hash.init(.{ .vocab = vocab, .ngram = 3, .per = 2, .base = base, .divisor = 4, .shards = 2, .seed = 1234, .eos = eos, .ple_dim = 4 * row_values, .ple_index = 0 }) catch unreachable; + } + + fn put(list: *std.ArrayList(u8), gpa: std.mem.Allocator, bytes: []const u8) !void { + try list.appendSlice(gpa, bytes); + } + + /// Writes config.json, an empty index and ngram_embedding.safetensors under `dir`; row r's words are a hash of r. + fn write(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) !void { + const h = spec(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(gpa); + var header: std.ArrayList(u8) = .empty; + defer header.deinit(gpa); + try header.appendSlice(gpa, "{"); + const p = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding."; + var first = true; + const Entry = struct { + fn add(hd: *std.ArrayList(u8), a: std.mem.Allocator, f: *bool, name: []const u8, dtype: []const u8, shape: []const usize, lo: usize, hi: usize) !void { + if (!f.*) try hd.appendSlice(a, ","); + f.* = false; + var buf: [400]u8 = undefined; + var sb: [64]u8 = undefined; + var sl: usize = 0; + for (shape, 0..) |d, i| { + const s = try std.fmt.bufPrint(sb[sl..], "{s}{d}", .{ if (i == 0) "" else ",", d }); + sl += s.len; + } + try hd.appendSlice(a, try std.fmt.bufPrint(&buf, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ name, dtype, sb[0..sl], lo, hi })); + } + }; + var nb: [200]u8 = undefined; + for (0..h.shards) |s| { + const rows: usize = @intCast(h.starts[s + 1] - h.starts[s]); + const lo = data.items.len; + for (0..rows) |r| { + const global: u64 = @intCast(h.starts[s] + @as(i64, @intCast(r))); + var prng = std.Random.DefaultPrng.init(global + 1); + for (0..51) |w| { + var v: u16 = prng.random().int(u16); + if (w == 0) v = @bitCast(@as(f16, @floatCast(0.25 + 0.01 * @as(f32, @floatFromInt(global))))); + try data.appendSlice(gpa, &std.mem.toBytes(v)); + } + } + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}shard_{d}.trellis", .{ p, s }), "I16", &.{ rows, 51 }, lo, data.items.len); + } + var lo = data.items.len; + for (0..h.heads * row_values) |i| try data.appendSlice(gpa, &std.mem.toBytes(@as(f16, @floatCast(@as(f32, @floatFromInt(i % 7)) * 0.01)))); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_bias", .{p}), "F16", &.{ h.heads, row_values }, lo, data.items.len); + lo = data.items.len; + for (h.mult[0..h.n]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}layer_multipliers", .{p}), "I64", &.{h.n}, lo, data.items.len); + lo = data.items.len; + for (h.sizes[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_vocab_sizes", .{p}), "I64", &.{h.heads}, lo, data.items.len); + lo = data.items.len; + for (h.offsets[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_offsets", .{p}), "I64", &.{h.heads}, lo, data.items.len); + try header.appendSlice(gpa, "}"); + var out: std.ArrayList(u8) = .empty; + defer out.deinit(gpa); + try out.appendSlice(gpa, &std.mem.toBytes(@as(u64, header.items.len))); + try out.appendSlice(gpa, header.items); + try out.appendSlice(gpa, data.items); + try dir.writeFile(io, .{ .sub_path = file_name, .data = out.items }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{}}" }); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = + \\{"text_config":{"vocab_size":100,"eos_token_id":99,"num_hidden_layers":4,"hidden_size":640,"ple_layer_ids":[2], + \\"ngram_size":3,"heads_per_ngram":2,"ngram_vocab_size_base":11,"make_ngram_vocab_size_divisible_by":4, + \\"split_ngram_parts":2,"ple_embed_dim":640}} + }); + } +}; + +test "exl3 ngram: Ngram.rows on a synthetic EXL3 table: hook, map and pread readers, hot cache" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try Synth.write(gpa, io, tmp.dir); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + + var g = try ple.Ngram.init(gpa, io, dir); + defer g.deinit(); + try std.testing.expect(g.format == .exl3); + const h = g.hash; + try std.testing.expectEqual(@as(usize, 4), h.heads); + const tokens = [_]u32{ 5, 7, Synth.eos, 7, 7, 3 }; + const w = g.width(); + var out: [tokens.len * 4 * row_values]u16 = undefined; + var xs: [tokens.len * 4 * row_values / 32]f32 = undefined; + var history = [_]i64{ Synth.eos, Synth.eos }; + try g.rows(&history, &tokens, &out, &xs); + try std.testing.expectEqualSlices(i64, &.{ 7, 3 }, &history); + + // the same rows by hand: upstream's ids, the shard's bytes, decodeRow + var pread_t = (try Table.open(gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + var hist = [_]i64{ Synth.eos, Synth.eos }; + for (tokens, 0..) |tok, t| { + var ids: [4]i64 = undefined; + try ple.lastIds(h, &hist, tok, &ids); + hist = .{ hist[1], tok }; + for (ids, 0..) |id, head| { + const loc = try h.locate(id); + var bytes: [102]u8 align(2) = undefined; + try pread_t.raw(loc.shard, loc.local, &bytes); + var mapped: [102]u8 = undefined; + try g.x3.?.raw(loc.shard, loc.local, &mapped); + try std.testing.expectEqualSlices(u8, &mapped, &bytes); + var vals: [row_values]f16 = undefined; + try decodeRow(@as([*]const i16, @ptrCast(&bytes))[0..51], 5, pread_t.bias[head * row_values ..][0..row_values], &vals); + const got = out[t * w + head * row_values ..][0..row_values]; + try std.testing.expectEqualSlices(u16, @ptrCast(&vals), got); + var gs: [5]f32 = undefined; + sumGroupsF16(got, &gs); + try std.testing.expectEqualSlices(f32, &gs, xs[(t * w + head * row_values) / 32 ..][0..5]); + } + } + // tokens 2 and 3 after an eos: token 7 twice in a fresh segment reads the same rows; the cache serves the repeat + const before = pread_t.hotStats(); + var again: [102]u8 = undefined; + const loc = try h.locate(h.offsets[0]); + try pread_t.raw(loc.shard, loc.local, &again); + try pread_t.raw(loc.shard, loc.local, &again); + const after = pread_t.hotStats(); + try std.testing.expectEqual(before.lookups + 2, after.lookups); + try std.testing.expect(after.misses <= before.misses + 1); + // bf16 elements: the fp16 values rounded once more + var bf_t = (try Table.open(gpa, io, dir, h, .{ .emb = .bf16 })).?; + defer bf_t.destroy(); + var r16: [row_values]u16 = undefined; + var rbf: [row_values]u16 = undefined; + try g.x3.?.row(loc.shard, loc.local, 0, &r16); + try bf_t.row(loc.shard, loc.local, 0, &rbf); + for (r16, rbf) |a, b| try std.testing.expectEqual(ple.toBf16(@floatCast(@as(f16, @bitCast(a)))), b); +} + +test "exl3 ngram: a directory without the n-gram file is not an EXL3 table" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + try std.testing.expect((try Table.open(gpa, io, dir, Synth.spec(), .{})) == null); +} diff --git a/zig/src/families/flashnext/cuda_exl3_pool.zig b/zig/src/families/flashnext/cuda_exl3_pool.zig new file mode 100644 index 0000000000..1ae979846c --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_pool.zig @@ -0,0 +1,366 @@ +//! One layer's EXL3 routed-expert offload pool on upstream's `Experts` (plan step 6 of +//! docs/plans/exl3-on-upstream-1.1.0.md). Routed trellis stays in host memory (the mmapped checkpoint); `lines` device +//! lines hold the hot experts; the shared expert (id 512) and every expert's suh/svh stay resident. +//! +//! The segment cut is explicit. Per round of a layer (`rows` rows of `slots` picks, the shared expert in its slot): +//! graph A (block, router, top-k) -> D2H of the picks -> `plan(picks)` on the host (policy, no device work) +//! -> for each pass j of the returned round: `stage(round, j, s)` (copies + table rewrites on the copy stream, +//! `s` waits on the device), `Experts.runSlots(.., picksFor(round, j, own), ..)` (graph B), `release(round, j, s)`. +//! A hit needs no host work beyond the plan; an all-hit round (`round.allHits()`) has nothing to stage and its graph +//! B replays as captured: the tables (`tables()`) are device memory at fixed addresses, rewritten in place. +//! A round whose distinct experts do not fit the lines runs as several passes over masked picks (each slot's output is +//! written by exactly one pass; slots of other passes are left untouched, which is what runSlots does with pick -1). + +const std = @import("std"); +const cuda = @import("cuda"); +const round_ = @import("exl3_round.zig"); +const exl3e = cuda.exl3_experts; +const Lines = cuda.exl3_pool.Lines; +const telemetry = @import("exl3_telemetry.zig"); + +pub const offload = @import("exl3_offload.zig"); +pub const round = round_; +pub const graph = @import("exl3_graph.zig"); +pub const Round = round_.Round; +pub const linesFor = round_.linesFor; +pub const shape = offload.shape; +pub const shared_id = offload.shared_id; +/// The fewest lines a layer may have: one row's ten routed picks always fit one pass. +pub const min_lines: usize = 10; + +test { + _ = @import("exl3_frequency_seed.zig"); +} + +/// The layer's experts in host memory: routed [512] gate/up/down trellis bytes (stored layout), the shared expert's. +pub const Source = struct { + routed: []const [3][]const u8, + shared: ?[3][]const u8 = null, + pinned_bank: ?[]const u8 = null, +}; + +/// Device addresses of the resident, stacked fp16 scales ([E, D] or [E, I], E = 513). +pub const Scales = struct { suh_g: u64, suh_u: u64, svh_g: u64, svh_u: u64, suh_d: u64, svh_d: u64 }; + +pub const Options = struct { + layer: u8, + lines: usize, + dims: usize = 2560, + width: usize = 640, + /// Rows and slots a round may hold (sizes the masked-pick buffers). + max_rows: usize = 16, + slots: usize = 11, + staging: usize = 8, + /// Copy from the source bytes directly (they are pinned already); otherwise through pinned staging lines. + direct: bool = false, + codebook: cuda.exl3.Codebook = .mul1, + act_mode: i32 = 1, + swap_every: u16 = 4, + swap_cap: u16 = 96, + /// Bytes a swap round may prefetch (the copy budget); 0: no prefetch. + prefetch_bytes: u64 = 0, + /// Double buffering: a round that fits without them does not evict the previous round's lines, so its copies + /// need not wait for the previous round's compute. + double_buffer: bool = true, + /// Masked-pick buffers for this many passes of a split round (0: the round's own bound, uniq + 1). + max_passes: usize = 0, +}; + +/// Passes a round of a `lines`-line pool can split into when nothing prefetches (`prefetch_bytes` 0): pass 0 and one +/// pass a full pool of misses over the 512 routed experts, plus one for rounding. +pub fn passesFor(lines: usize) usize { + return 2 + (shape.experts + lines - 1) / lines; +} + +pub const Stats = struct { + rounds: u64 = 0, + passes: u64 = 0, + hit_slots: u64 = 0, + miss_slots: u64 = 0, + unique: u64 = 0, + misses: u64 = 0, // mandatory copies + prefetches: u64 = 0, + evictions: u64 = 0, + busy: u64 = 0, // plans that had to drain in-flight copies first + protected: u64 = 0, // rounds that spared the previous round's lines +}; + +pub const Pool = struct { + gpa: std.mem.Allocator, + io: std.Io, + opt: Options, + src: Source, + scales: Scales, + lines: Lines, + k2: [][3]i32, // per routed expert, from its trellis sizes + slots_tab: []u16, + line_tab: []offload.Line, + stamp_tab: []u32, + layer_tab: []u8, + v: offload.Vram, + counts: *[shape.layers][shape.experts]u32, + own_counts: bool, + round: Round, + now: u32 = 0, + round_no: u32 = 0, + picks_host: cuda.HostBuffer, + picks_dev: cuda.DeviceBuffer, + picks_ev: cuda.Event, + picks_pending: bool = false, + picks_n: usize = 0, // the planned round's picks (rows * slots) + prev: []u16, // the previous round's routed experts (double buffering) + prev_n: usize = 0, + stats: Stats = .{}, + telemetry_id: u64 = 0, + telemetry_valid: bool = true, + telemetry_device: telemetry.Device = .{}, + + /// `counts`: the family's lifetime use table shared by every layer's pool (null: the pool owns one). + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, src: Source, scales: Scales, opt: Options, counts: ?*[shape.layers][shape.experts]u32) !Pool { + if (opt.lines < min_lines) return error.PoolTooSmall; + if (src.routed.len != shape.experts) return error.Invalid; + const tiles = (opt.dims / 16) * (opt.width / 16); + var bytes: [3]usize = .{ 0, 0, 0 }; + const k2 = try gpa.alloc([3]i32, shape.experts); + errdefer gpa.free(k2); + var held: usize = 0; + for (src.routed, k2) |e, *w| { + // an expert another rank holds (expert parallel, cuda_exl3_tp.zig): no bytes, never picked here + if (e[0].len == 0 and e[1].len == 0 and e[2].len == 0) { + w.* = .{ 0, 0, 0 }; + continue; + } + held += 1; + for (0..3) |p| { + // trellis bytes = tiles * 16 * k2 (k2 half-bits a value, 8 * k2 int16 words a 16x16 tile) + if (e[p].len == 0 or e[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(e[p].len / (tiles * 16)); + bytes[p] = @max(bytes[p], e[p].len); + } + } + if (held == 0) return error.Invalid; + const table_e: usize = if (src.shared != null) shape.experts + 1 else shape.experts; + var lines = try Lines.init(gpa, d, .{ .experts = table_e, .bytes = bytes, .lines = opt.lines, .staging = if (opt.direct) 0 else opt.staging }); + errdefer lines.deinit(); + lines.pinned_bank = if (opt.direct) src.pinned_bank else null; + if (src.shared) |sh| { + var w: [3]i32 = undefined; + for (0..3) |p| { + if (sh[p].len == 0 or sh[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(sh[p].len / (tiles * 16)); + } + // the shared expert gets its own resident region (it may be wider than the routed experts) + try lines.setResident(shared_id, sh, w); + } + const n = opt.lines; + const slots_tab = try gpa.alloc(u16, n); + errdefer gpa.free(slots_tab); + const line_tab = try gpa.alloc(offload.Line, n); + errdefer gpa.free(line_tab); + const stamp_tab = try gpa.alloc(u32, n); + errdefer gpa.free(stamp_tab); + const layer_tab = try gpa.alloc(u8, n); + errdefer gpa.free(layer_tab); + const own = counts == null; + const c = counts orelse blk: { + const t = try gpa.create([shape.layers][shape.experts]u32); + t.* = @splat(@splat(0)); + break :blk t; + }; + errdefer if (own) gpa.destroy(c); + var rnd = try Round.init(gpa, n, opt.max_rows * opt.slots); + errdefer rnd.deinit(gpa); + const passes_n = if (opt.max_passes == 0) rnd.passes.len else @min(opt.max_passes, rnd.passes.len); + if (opt.max_passes != 0 and opt.prefetch_bytes != 0) return error.Invalid; // in-flight prefetches lengthen rounds + const pick_bytes = passes_n * opt.max_rows * opt.slots * 4; + var picks_host = try cuda.HostBuffer.alloc(d, pick_bytes); + errdefer picks_host.free(); + var picks_dev = try cuda.DeviceBuffer.alloc(d, pick_bytes); + errdefer picks_dev.free(); + var picks_ev = try cuda.Event.init(d, false); + errdefer picks_ev.deinit(); + const prev = try gpa.alloc(u16, shape.experts); + errdefer gpa.free(prev); + return .{ + .telemetry_id = try telemetry.identity(), + .prev = prev, + .gpa = gpa, + .io = io, + .opt = opt, + .src = src, + .scales = scales, + .lines = lines, + .k2 = k2, + .slots_tab = slots_tab, + .line_tab = line_tab, + .stamp_tab = stamp_tab, + .layer_tab = layer_tab, + .v = offload.Vram.init(slots_tab, line_tab, stamp_tab, layer_tab), + .counts = c, + .own_counts = own, + .round = rnd, + .picks_host = picks_host, + .picks_dev = picks_dev, + .picks_ev = picks_ev, + }; + } + + pub fn deinit(p: *Pool) void { + p.lines.deinit(); + p.picks_ev.deinit(); + p.picks_dev.free(); + p.picks_host.free(); + p.round.deinit(p.gpa); + if (p.own_counts) p.gpa.destroy(p.counts); + p.gpa.free(p.prev); + p.gpa.free(p.layer_tab); + p.gpa.free(p.stamp_tab); + p.gpa.free(p.line_tab); + p.gpa.free(p.slots_tab); + p.gpa.free(p.k2); + p.* = undefined; + } + + /// The `Experts` layer: the pool's device tables (fixed addresses) and the resident scales. + pub fn tables(p: *const Pool) exl3e.Layer { + return .{ + .gate = p.lines.table(0), + .up = p.lines.table(1), + .down = p.lines.table(2), + .gate_k2 = p.lines.widths(0), + .up_k2 = p.lines.widths(1), + .down_k2 = p.lines.widths(2), + .suh_g = p.scales.suh_g, + .suh_u = p.scales.suh_u, + .svh_g = p.scales.svh_g, + .svh_u = p.scales.svh_u, + .suh_d = p.scales.suh_d, + .svh_d = p.scales.svh_d, + .experts = p.lines.spec.experts, + .dims = p.opt.dims, + .width = p.opt.width, + .codebook = p.opt.codebook, + .act_mode = p.opt.act_mode, + }; + } + + /// Promotes prefetch lines whose copy landed (the policy's reap); never blocks. + fn reap(p: *Pool) !void { + for (p.line_tab, p.slots_tab, 0..) |state, id, i| { + if (state != .inflight) continue; + if (try p.lines.landed(i)) offload.reapAt(&p.v, p.opt.layer, id, p.now) catch unreachable; + } + } + + /// Host step of a round: `picks` [rows * slots] (int32, -1 skipped, 512 the shared expert), as the device holds + /// them. Decides hits, misses, victims and passes; enqueues nothing. The result stays valid until the next plan. + pub fn plan(p: *Pool, picks: []const i32, rows: usize) !*const Round { + return p.planInternal(picks, rows, true); + } + + pub fn planUnobserved(p: *Pool, picks: []const i32, rows: usize) !*const Round { + return p.planInternal(picks, rows, false); + } + + fn planInternal(p: *Pool, picks: []const i32, rows: usize, observe: bool) !*const Round { + if (rows == 0 or rows > p.opt.max_rows or picks.len != rows * p.opt.slots) return error.Invalid; + p.now += 1; + p.round_no += 1; + try p.reap(); + const ropt: round_.Options = .{ + .swap_every = p.opt.swap_every, + .swap_cap = p.opt.swap_cap, + .prefetch_bytes = p.opt.prefetch_bytes, + .line_bytes = p.lines.layout.line_bytes, + .protect = if (p.opt.double_buffer) p.prev[0..p.prev_n] else &.{}, + .observe = observe, + }; + round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round) catch |e| switch (e) { + error.Busy => { + // every line is held by a copy in flight: let them land, then plan again (nothing was written) + p.stats.busy += 1; + try p.lines.drain(); + try p.reap(); + try round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round); + }, + else => return e, + }; + const r = &p.round; + if (r.n_passes * p.opt.max_rows * p.opt.slots * 4 > p.picks_dev.len) return error.PoolPasses; + p.picks_n = picks.len; + p.prev_n = r.n_ids; + @memcpy(p.prev[0..r.n_ids], r.ids[0..r.n_ids]); + if (r.protected) p.stats.protected += 1; + if (r.masked()) { + if (p.picks_pending) try p.picks_ev.synchronize(); // the previous round's uploads read these bytes + p.picks_pending = false; + const stride = p.opt.max_rows * p.opt.slots; + const host = p.picks_host.slice(i32); + var member: [shape.experts + 1]bool = undefined; + for (0..r.n_passes) |j| { + @memset(&member, false); + for (r.passIds(j)) |id| member[id] = true; + member[shared_id] = j == 0; + for (picks, host[j * stride ..][0..picks.len]) |pk, *o| { + o.* = if (pk >= 0 and pk <= shared_id and member[@intCast(pk)]) pk else -1; + } + } + } + p.stats.rounds += 1; + p.stats.passes += r.n_passes; + p.stats.hit_slots += r.hit_slots; + p.stats.miss_slots += r.miss_slots; + p.stats.unique += r.unique; + p.stats.misses += r.misses_(); + p.stats.evictions += r.evictions; + return r; + } + + fn loadCopy(p: *Pool, c: round_.Copy) !void { + if (p.k2[c.id][0] == 0) return error.Invalid; // another rank's expert reached the planner unmasked + const ev: ?usize = if (c.evicted == offload.Cache.empty_id) null else c.evicted; + try p.lines.load(c.line, c.id, p.src.routed[c.id], p.k2[c.id], ev, p.opt.direct, p.io); + } + + /// Device step of pass `j`: its copies and table rewrites on the copy stream, its masked picks (if any) on `s`, and + /// `s` waits for every line the pass reads. The host only blocks to reuse a staging slot. + pub fn stage(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + for (r.passCopies(j)) |c| try p.loadCopy(c); + if (r.masked()) { + const stride = p.opt.max_rows * p.opt.slots * 4; + const n = p.picks_n * 4; + try p.picks_dev.d.check(p.picks_dev.d.api.cuMemcpyHtoDAsync_v2(p.picks_dev.ptr + j * stride, p.picks_host.bytes[j * stride ..].ptr, n, s.handle), "cuMemcpyHtoDAsync"); + try p.picks_ev.record(s); + p.picks_pending = true; + } + for (r.passLines(j)) |line| try p.lines.gate(s, line); + } + + /// The picks pass `j` reads: the round's own device picks, or the pool's masked copy for a split round. + pub fn picksFor(p: *const Pool, r: *const Round, j: usize, own: u64) u64 { + if (!r.masked()) return own; + return p.picks_dev.ptr + j * p.opt.max_rows * p.opt.slots * 4; + } + + /// After pass `j` was enqueued on `s`: its lines are protected until it finishes. After the last pass, the round's + /// prefetch copies (swap rounds) go on the copy stream. + pub fn release(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + try p.lines.release(s, r.passLines(j)); + if (j + 1 == r.n_passes) { + for (r.prefetch()) |c| try p.loadCopy(c); + p.stats.prefetches += r.n_prefetch; + } + } + + pub fn warmProfile(pool: *Pool) !usize { + return @import("exl3_frequency_seed.zig").warm(pool); + } + + pub fn deviceBytes(p: *const Pool) usize { + return p.lines.deviceBytes() + p.picks_dev.len; + } + + pub fn lineBytes(p: *const Pool) usize { + return p.lines.layout.line_bytes; + } +}; diff --git a/zig/src/families/flashnext/cuda_exl3_prefill.zig b/zig/src/families/flashnext/cuda_exl3_prefill.zig new file mode 100644 index 0000000000..2fc81e90c3 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill.zig @@ -0,0 +1,1892 @@ +//! Expert-major prefill for an EXL3 Flash Next checkpoint whose routed experts live in host memory (the offload pool). +//! +//! The batched pass (cuda_batch.zig `chunk`) takes up to `Engine.chunk` rows through every layer, and its MoE copies the +//! experts a 256-row block picked into the layer's pool lines: a long prompt copies every expert of every layer about +//! once per block. Here a prompt goes through in DEEP chunks of `rows` rows (thousands): every layer runs the pass's own +//! blocks (HC read-outs, DeltaNet / attention, write-backs, router and top-k) on sub-blocks of the engine's chunk size +//! over device rows that persist for the whole deep chunk, then the layer's routed experts run EXPERT-MAJOR over all +//! `rows` rows: each expert a pair picked is staged into a pool line once and computed on every pair that picked it +//! (flashnext_exl3_prefill.cu, 64-member tiles), in passes of at most half the layer's lines, so a pass's copies run on +//! the copy engine while the previous pass computes. A deep chunk thus copies each expert of a layer at most once. +//! +//! Bits: every block but the routed experts is the pass's own, row for row (its kernels are row-invariant), so the +//! streams stay those of the batched pass until the experts' fp32 sums, which run in another order (64-member tiles +//! over all of K instead of runSlots' split-K partials): per pair the output moves by fp32 rounding, measured at about +//! 4e-5 of a layer's peak (tf-cuda-test exl3-prefill prompt). The logits change accordingly (bf16 noise). +//! +//! The hook: Engine.forward calls `forward` first; it declines (false) for anything but a plain prompt on one rank of an +//! EXL3 checkpoint with pooled or resident experts, no media, no window keep, no logits for every row. +//! TF_FLASHNEXT_EXL3_PREFILL=0 turns it off; _ROWS sets the deep chunk (default 8192), _MIN the shortest prompt +//! that takes it (default 512 rows), _TILES=0 keeps the dense projections on the reference kernels, _RING=0 stages +//! each layer into its own lines, _KERNEL=slots runs the experts through runSlots (the batched pass's bits). + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("cuda_batch.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_exl3_load.zig"); +const pool_ = @import("cuda_exl3_pool.zig"); +const pk = @import("cuda_exl3_prefill_k.zig"); +const prompt = @import("cuda_prompt.zig"); +const ring_ = @import("cuda_exl3_prefill_ring.zig"); +const history_ = @import("cuda_exl3_prefill_history.zig"); +const tp_ = @import("cuda_exl3_tp.zig"); +const window_ = @import("exl3_prefill_window.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const eps = wt.eps; +const vocab = wt.vocab; +const out_sk = wt.out_sk; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const table_n = wt.experts_n + 1; +const routed_n_ = wt.experts_n; +const shared_id = wt.experts_n; +const att_rows = batch.att_rows; +const routed_guess: usize = 400; +const pss_n = batch.pss_n; + +pub const Options = struct { + enabled: bool = true, + /// Rows a deep chunk holds (whole attention blocks). + rows: usize = 8192, + /// The deep chunk the server reserves at load (reserveFor): its buffers come out of the pools' share of the card + /// (0.16 GiB a 1024 rows), so the server takes fewer than `rows` unless TF_FLASHNEXT_EXL3_PREFILL_ROWS says. + serve_rows: usize = 4096, + /// The shortest prompt that takes this path. + min_rows: usize = 512, + /// New experts a pass stages at most, as a fraction of the layer's lines (0.5: double buffering). + pass_lines_half: bool = true, + /// Stage the next layer's first pass while this layer's blocks compute. + prefetch: bool = true, + /// Run each pass through Experts.runSlots (256-row blocks over masked picks) instead of the 64-member prefill + /// kernel: the batched pass's bits exactly (a check of everything but the kernel; slow). + slots_kernel: bool = false, + /// Every decoder layer's pool lines as one ring the copy engine streams the coming layers' experts into + /// (cuda_exl3_prefill_ring.zig); false: each layer stages into its own lines, in passes of half of them. + ring: bool = true, + /// Experts a ring pass computes at most. + ring_pass: usize = 32, + /// Layers ahead the ring prefetches (the guess: the hottest experts, as many as the current layer used). + ring_ahead: usize = 1, + /// Dense EXL3 projections of a deep chunk's sub-blocks on #806's prompt tiles (Dense.prompt, .tiles) instead of + /// the reference 128-row blocks: 1.4-2x faster at 128+ rows, a different fp32 order (bf16-level changes; the + /// same float64 error, tf-cuda-test exl3-prefill dense). TF_FLASHNEXT_EXL3_PREFILL_TILES=0: the reference. + tiles: bool = true, + /// After a ring prefill, refill every pool with its layer's hottest experts of the prefills (measured: decode hit + /// +5%, ms/token unchanged, prefill +0.17 s; off by default, TF_FLASHNEXT_EXL3_PREFILL_RESEED=1). + reseed: bool = false, + reduced_branch: bool = false, + two_chunk: bool = false, +}; + +pub fn optionsFromEnv() Options { + var o: Options = .{}; + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL")) |v| o.enabled = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_ROWS")) |v| { + o.rows = std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.rows; + o.serve_rows = o.rows; + } + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_MIN")) |v| o.min_rows = std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.min_rows; + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_PREFETCH")) |v| o.prefetch = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_KERNEL")) |v| o.slots_kernel = std.mem.eql(u8, std.mem.span(v), "slots"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_TILES")) |v| o.tiles = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RESEED")) |v| o.reseed = std.mem.eql(u8, std.mem.span(v), "1"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_REDUCED_BRANCH")) |v| o.reduced_branch = std.mem.eql(u8, std.mem.span(v), "1"); + o.two_chunk = window_.enabled(if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK")) |value| std.mem.span(value) else null); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RING")) |v| o.ring = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RING_PASS")) |v| o.ring_pass = @max(1, std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.ring_pass); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD")) |v| o.ring_ahead = std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.ring_ahead; + return o; +} + +/// The options every call reads (a test may change them between calls). +pub var options: ?Options = null; + +fn opts() Options { + if (options == null) options = optionsFromEnv(); + return options.?; +} + +/// What the last deep forward did (read by tf-cuda-test exl3-prefill engine). +pub const Stats = struct { + chunks: u64 = 0, + rows: u64 = 0, + passes: u64 = 0, + copies: u64 = 0, + prefetched: u64 = 0, + members: u64 = 0, + host_wait_ns: u64 = 0, // the picks' D2H and its wait, every layer + plan_ns: u64 = 0, // host plan of the passes (members, items, pool rounds) + reduces: u64 = 0, // expert parallel: the ranks' all-reduces (one a layer a deep chunk) + reduce_ns: u64 = 0, // their combine + all-reduce + place, measured only with TF_FLASHNEXT_PROFILE=1 (a wait each) + head_rows: u64 = 0, // MTP head rows run expert-major (`head`) +}; +pub var stats: Stats = .{}; + +/// Device bytes of a deep chunk's own rows: streams, sums, gates, mixed rows, slot outputs, picks, n-gram rows, and +/// the experts' compact scratch and member lists. +pub fn deviceBytes(rows: usize) usize { + return Layout.of(rows, 1).total; +} + +/// The same at `world` ranks (expert parallel adds the ranks' sum, rows x 2560 fp32). +pub fn deviceBytesAt(rows: usize, world: usize) usize { + return Layout.of(rows, world).total; +} + +const Layout = struct { + h: usize, + pss: usize, + inj: usize, + mixed: usize, + xsm: usize, + y: usize, + wts: usize, + picks: usize, + ple_emb: usize, + ple_xs: usize, + ids: usize, + pos: usize, + members: usize, + items: usize, + kscratch: usize, + /// Expert parallel only: a rank's combined branch, then the ranks' sum [rows, D] fp32 (0 bytes at world 1). + sum: usize, + total: usize, + member_cap: usize, + + fn of(rows: usize, world: usize) Layout { + var l: Layout = undefined; + var at: usize = 0; + const take = struct { + fn f(a: *usize, n: usize) usize { + const r = a.*; + a.* = r + (n + 255) / 256 * 256; + return r; + } + }.f; + const pairs = rows * slots; + l.member_cap = 2 * rows; + l.h = take(&at, rows * wide * 2); + l.pss = take(&at, rows * pss_n * 4); + l.inj = take(&at, rows * streams * 2); + l.mixed = take(&at, rows * dims * 2); + l.xsm = take(&at, rows * (dims / 32) * 4); + l.y = take(&at, pairs * dims * 4); + l.wts = take(&at, pairs * 4); + l.picks = take(&at, pairs * 4); + l.ple_emb = take(&at, rows * dims * 2); + l.ple_xs = take(&at, rows * (dims / 32) * 4); + l.ids = take(&at, rows * 4); + l.pos = take(&at, (rows / att_rows + 1) * 16); + l.members = take(&at, pairs * 4); + l.items = take(&at, 3 * pk.maxItems(pairs, table_n) * 4); + l.kscratch = take(&at, pk.scratchBytes(l.member_cap, dims, wt.moe_w)); + l.sum = take(&at, if (world > 1) rows * dims * 4 else 0); + l.total = at; + return l; + } +}; + +/// Host side: pinned uploads (ids, positions, n-gram rows, members and items) and the picks' download. +const HostLayout = struct { + ids: usize, + pos: usize, + pos_stride: usize, + emb: usize, + xs: usize, + picks: usize, + members: usize, + items: usize, + total: usize, + + fn of(rows: usize) HostLayout { + var l: HostLayout = undefined; + var at: usize = 0; + const take = struct { + fn f(a: *usize, n: usize) usize { + const r = a.*; + a.* = r + (n + 63) / 64 * 64; + return r; + } + }.f; + const pairs = rows * slots; + l.ids = take(&at, rows * 4); + l.pos_stride = (rows / att_rows + 1) * 16; + l.pos = take(&at, l.pos_stride * (wt.layers_n + 1)); + l.emb = take(&at, rows * dims * 2); + l.xs = take(&at, rows * (dims / 32) * 4); + l.picks = take(&at, pairs * 4); + l.members = take(&at, pairs * 4); + l.items = take(&at, 3 * pk.maxItems(pairs, table_n) * 4); + l.total = at; + return l; + } +}; + +/// One pass of a layer: its experts (the pool round's routed ids), its members and items. +const Pass = struct { id_lo: u32, id_hi: u32, m_lo: u32, m_n: u32, i_lo: u32, i_n: u32 }; + +/// Sub-blocks a deep chunk holds at most (chunk's view array). +const max_sub = 64; + +const Persistent = struct { + h: usize, + pss: usize, + inj: usize, + y: usize, + wts: usize, + total: usize, + + fn of(rows: usize) Persistent { + const h = rows * wide * 2; + const pss = std.mem.alignForward(usize, rows * pss_n * 4, 256); + const inj = std.mem.alignForward(usize, rows * streams * 2, 256); + const y = rows * slots * dims * 4; + const wts = std.mem.alignForward(usize, rows * slots * 4, 256); + return .{ .h = 0, .pss = h, .inj = h + pss, .y = h + pss + inj, .wts = h + pss + inj + y, .total = h + pss + inj + y + wts }; + } +}; + +const Deep = struct { + gpa: std.mem.Allocator, + owner: *const anyopaque, + rows: usize, + world: usize = 1, + lay: Layout, + hl: HostLayout, + dev: cuda.DeviceBuffer, + sum_dev: ?cuda.DeviceBuffer = null, + host: cuda.HostBuffer, + kern: pk.Kernels, + aux: cuda.Stream, + /// Expert parallel: the ranks' all-reduces run on `red`, one a sub-block, so the next layer's sub-block b only waits + /// for its own rows' sum (`red_ev[b]`) and the others' reduce under its compute; `red_n` sub-blocks are in flight. + red: ?cuda.Stream = null, + red_ev: [max_sub]cuda.Event = undefined, + red_go: cuda.Event = undefined, + red_n: usize = 0, + reduced_branch: bool = false, + unit_weights: ?cuda.DeviceBuffer = null, + /// The layer's pass plan: expert ids in pass order, the passes, member counts and starts. + order: [table_n]u16 = undefined, + passes: [table_n + 1]Pass = undefined, + n_passes: usize = 0, + cnt: [table_n]u32 = undefined, + start: [table_n]u32 = undefined, + /// Experts the previous layer's prefetch staged for each layer (held at its MoE). + fake: [48 * slots]i32 = undefined, + /// The prefill line ring over the pools (made at the first prefill, once the pools exist). + ring: ?*ring_.Ring = null, + ring_failed: bool = false, + /// Experts each layer used in the current chunk (the next layers' prefetch size). + used: [wt.layers_n]usize = @splat(routed_guess), + history: [wt.layers_n]history_.Slot = @splat(.{}), + history_shape: ?history_.Shape = null, + second: ?*Deep = null, + shared: ?*const Deep = null, + window_graph_reservation: u64 = 0, + + fn dptr(dp: *const Deep, off: usize) u64 { + if (dp.shared) |source| { + const persistent = Persistent.of(dp.rows); + inline for (.{ "h", "pss", "inj", "y", "wts" }) |field| { + if (off == @field(dp.lay, field)) return dp.dev.ptr + @field(persistent, field); + } + if (off == dp.lay.sum) if (dp.sum_dev) |buffer| return buffer.ptr; + return source.dptr(off); + } + return dp.dev.ptr + off; + } + fn hbytes(dp: *const Deep, off: usize, n: usize) []u8 { + return dp.host.bytes[off..][0..n]; + } +}; + +var deep: ?*Deep = null; + +pub const RingTelemetry = struct { owner: u64, copies: u64, dma_requested_bytes: u64, unused_copies: u64, hit_picks: u64, miss_picks: u64, hit_expert_requests: u64, miss_expert_requests: u64, attribution: ?@import("cuda_exl3_prefill_attribution.zig").Totals }; + +pub fn ringTelemetry(owner: *const anyopaque) ?RingTelemetry { + const buffers = deep orelse return null; + if (buffers.owner != owner) return null; + const ring = buffers.ring orelse return null; + return .{ .owner = ring.telemetry_id, .copies = ring.stats.copies, .dma_requested_bytes = ring.stats.dma_requested_bytes, .unused_copies = ring.stats.wasted, .hit_picks = ring.stats.hit_picks, .miss_picks = ring.stats.miss_picks, .hit_expert_requests = ring.stats.hit_expert_requests, .miss_expert_requests = ring.stats.miss_expert_requests, .attribution = if (ring.attribution) |state| state.totals else null }; +} + +/// Frees the deep chunk's buffers (the engine's deinit, or a test). +pub fn release() void { + if (deep) |dp| { + @import("cuda_exl3_graph.zig").clearPrefill(); + if (dp.second) |second| { + dp.aux.synchronize() catch @panic("two-chunk aux retirement synchronization failed"); + dp.dev.d.synchronize() catch @panic("two-chunk compute retirement synchronization failed"); + second.host.free(); + if (dp.red != null) { + for (&second.red_ev) |*event| event.deinit(); + second.red_go.deinit(); + } + if (second.sum_dev) |*buffer| buffer.free(); + second.dev.free(); + dp.gpa.destroy(second); + } + if (dp.ring) |r| r.deinit(); + if (dp.red) |*rs| { + rs.synchronize() catch {}; + for (&dp.red_ev) |*ev| ev.deinit(); + dp.red_go.deinit(); + rs.deinit(); + } + if (dp.unit_weights) |*buffer| buffer.free(); + dp.kern.unload(); + dp.aux.deinit(); + dp.host.free(); + dp.dev.free(); + dp.gpa.destroy(dp); + deep = null; + } + if (agree_buf) |*b| b.free(); + agree_buf = null; +} + +/// The server's sizing hook (cuda_lanes_load.zig, before the pools take the free memory): the deep chunk's buffers for +/// an EXL3 checkpoint with offloaded experts on one rank, sized for min(TF_FLASHNEXT_EXL3_PREFILL_ROWS, context). +/// A failure leaves the batched pass in charge of prompts. +pub fn reserveFor(e: *prompt.Engine, context: usize, tp: usize) void { + const o = opts(); + if (!o.enabled or (tp != 1 and tp != 2) or e.n_layers != wt.layers_n) return; + for (e.layers[0..e.n_layers]) |*l| { + const x = l.exl3 orelse return; + if (x.bank == null) return; + // expert parallel: every layer's experts carry this world's rank + if ((if (x.tp) |r| r.world else 1) != tp) return; + } + const rows = @min(o.serve_rows, (context + att_rows - 1) / att_rows * att_rows); + if (rows < o.min_rows) return; + reserveAt(e.gpa, e.driver, e, rows, tp) catch |err| std.log.warn("flashnext exl3 prefill: deep chunks of {d} rows refused ({t}); prompts take the batched pass", .{ rows, err }); +} + +/// Rows a serving lane prefills a step (cuda_lanes.zig `step`): a deep chunk when this engine holds deep buffers, +/// else the engine's chunk; at most `draft_rows` (the lane's draft keeps that many rows of streams). +pub fn stepRows(e: *const prompt.Engine, draft_rows: ?usize) usize { + var n = e.chunk; + if (deep) |dp| if (dp.owner == @as(*const anyopaque, e) and opts().enabled) { + n = @max(n, stepCapacity(dp.rows, opts().two_chunk and dp.second != null)); + }; + if (draft_rows) |r| n = @min(n, @max(r, 1)); + return n; +} + +fn stepCapacity(rows: usize, paired: bool) usize { + return if (paired and rows <= window_.max_rows) rows * 2 else rows; +} + +/// Rows of streams a lane's MTP draft keeps (cuda_lanes.zig, cuda_lanes_load.zig): a whole step's. +pub fn draftRows(e: *const prompt.Engine) usize { + var n = stepRows(e, null); + if (deep) |dp| if (dp.owner == @as(*const anyopaque, e)) { + const o = opts(); + if (o.two_chunk and o.enabled and o.ring and !o.reduced_branch and dp.world == 2) { + n = @max(n, stepCapacity(dp.rows, true)); + } + }; + return n; +} + +/// The engine's deinit: frees the deep chunk's buffers if `owner` made them. +pub fn releaseFor(owner: *const anyopaque) void { + if (deep) |dp| if (dp.owner == owner) release(); +} + +/// Allocates the deep chunk's buffers for `rows` rows now (before the pool takes the free memory). A no-op if they +/// exist and are at least that large. +pub fn reserve(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const anyopaque, rows_in: usize) !void { + return reserveAt(gpa, d, owner, rows_in, 1); +} + +/// `reserve` for a rank of `world` (expert parallel at 2: the ranks' sum, and per-slot outputs zeroed once so a slot +/// another rank computed only ever holds finite values for the write-back's zero weights). +pub fn reserveAt(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const anyopaque, rows_in: usize, world: usize) !void { + const rows = rows_in / att_rows * att_rows; + if (rows == 0 or (world != 1 and world != 2)) return error.Invalid; + if (deep) |dp| { + if (dp.owner == owner and dp.rows >= rows and dp.world == world) return; + release(); + } + const dp = try gpa.create(Deep); + errdefer gpa.destroy(dp); + dp.* = .{ .gpa = gpa, .owner = owner, .rows = rows, .world = world, .lay = Layout.of(rows, world), .hl = HostLayout.of(rows), .dev = undefined, .host = undefined, .kern = undefined, .aux = undefined }; + dp.dev = try cuda.DeviceBuffer.alloc(d, dp.lay.total); + errdefer dp.dev.free(); + dp.host = try cuda.HostBuffer.alloc(d, dp.hl.total); + errdefer dp.host.free(); + dp.kern = try pk.Kernels.load(d); + errdefer dp.kern.unload(); + dp.aux = try cuda.Stream.init(d, true); + errdefer dp.aux.deinit(); + if (world > 1) { + try d.check(d.api.cuMemsetD8Async(dp.dptr(dp.lay.y), 0, rows * slots * dims * 4, dp.aux.handle), "cuMemsetD8Async"); + try dp.aux.synchronize(); + var rs = try cuda.Stream.init(d, true); + errdefer rs.deinit(); + var made: usize = 0; + errdefer for (dp.red_ev[0..made]) |*ev| ev.deinit(); + for (&dp.red_ev) |*ev| { + ev.* = try cuda.Event.init(d, false); + made += 1; + } + dp.red_go = try cuda.Event.init(d, false); + dp.red = rs; + } + deep = dp; + std.log.info("flashnext exl3 prefill: deep chunks of {d} rows, {d:.2} GiB on the device, {d:.1} MiB pinned (world {d})", .{ rows, @as(f64, @floatFromInt(dp.lay.total)) / (1 << 30), @as(f64, @floatFromInt(dp.hl.total)) / (1 << 20), world }); +} + +pub fn reserveTwoChunk(d: *const cuda.Driver, free: u64, protected: u64, graph_reservation: u64) !u64 { + const o = opts(); + if (!o.two_chunk or !o.enabled or !o.ring or o.reduced_branch) return 0; + const dp = deep orelse return 0; + if (dp.dev.d != d or dp.world != 2 or dp.rows > window_.max_rows) return 0; + if (dp.second != null) return 0; + const graph_required = @import("cuda_exl3_graph_memory.zig").twoChunkReservationBytes(); + if (@import("cuda_exl3_graph.zig").enabled() and (graph_required == 0 or graph_reservation < graph_required)) { + std.log.info("flashnext exl3 prefill: two-chunk window declined, incremental graph reservation unavailable (required={d} supplied={d})", .{ graph_required, graph_reservation }); + return 0; + } + const persistent = Persistent.of(dp.rows); + const sum_bytes = if (dp.world > 1) try std.math.mul(usize, dp.rows, dims * 4) else 0; + const device_bytes = try std.math.add(usize, persistent.total, sum_bytes); + if (!window_.admits(free, protected, device_bytes)) { + std.log.info("flashnext exl3 prefill: two-chunk window declined, extra_device={d} full_layout={d} extra_pinned=0 protected={d} free={d}", .{ device_bytes, dp.lay.total, protected, free }); + return 0; + } + const second = try dp.gpa.create(Deep); + second.* = .{ .gpa = dp.gpa, .owner = dp.owner, .rows = dp.rows, .world = dp.world, .lay = dp.lay, .hl = dp.hl, .dev = undefined, .host = undefined, .kern = dp.kern, .aux = dp.aux, .red = dp.red, .red_ev = undefined, .red_go = undefined, .shared = dp }; + var device_allocated = false; + var host_allocated = false; + var made_events: usize = 0; + var reduce_event_allocated = false; + errdefer { + if (reduce_event_allocated) second.red_go.deinit(); + for (second.red_ev[0..made_events]) |*event| event.deinit(); + if (host_allocated) second.host.free(); + if (second.sum_dev) |*buffer| buffer.free(); + if (device_allocated) second.dev.free(); + dp.gpa.destroy(second); + } + second.dev = try cuda.DeviceBuffer.alloc(d, persistent.total); + device_allocated = true; + if (sum_bytes != 0) { + second.sum_dev = try cuda.DeviceBuffer.alloc(d, sum_bytes); + } + second.host = try cuda.HostBuffer.alloc(d, second.hl.total); + host_allocated = true; + if (dp.red != null) { + for (&second.red_ev) |*event| { + event.* = try cuda.Event.init(d, false); + made_events += 1; + } + second.red_go = try cuda.Event.init(d, false); + reduce_event_allocated = true; + } else { + second.red_ev = dp.red_ev; + second.red_go = dp.red_go; + } + try d.check(d.api.cuMemsetD8Async(second.dptr(second.lay.y), 0, dp.rows * slots * dims * 4, dp.aux.handle), "cuMemsetD8Async"); + try dp.aux.synchronize(); + dp.second = second; + dp.window_graph_reservation = graph_reservation; + std.log.info("flashnext exl3 prefill: experimental two-chunk window, chunk_rows={d} hidden_bytes={d} full_layout={d} extra_device={d} extra_pinned={d}", .{ dp.rows, 2 * dp.rows * wide * 2, dp.lay.total, device_bytes, second.hl.total }); + return device_bytes; +} + +pub fn twoChunkAllocated() bool { + return if (deep) |buffers| buffers.second != null else false; +} + +fn applies(e: *prompt.Engine, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks, o: Options) bool { + if (!o.enabled or ids.len < o.min_rows) return false; + if (hooks.keep != null or hooks.logits == .all) return false; + if (seq.media != null or e.n_layers != wt.layers_n) return false; + // expert parallel (cuda_exl3_tp.zig) at two ranks over the engine's collective; the head whole or split by vocabulary + const world = worldOf(e); + if (e.head_q.n != vocab and !(world == 2 and e.head_q.n == vocab / 2)) return false; + for (e.layers[0..e.n_layers]) |*l| { + const x = l.exl3 orelse return false; + if ((if (x.tp) |r| r.world else 1) != world or x.needsPool()) return false; + if (x.layer.act_mode != 1 and x.layer.act_mode != 0) return false; + } + return e.chunk >= att_rows; +} + +fn worldOf(e: *const prompt.Engine) usize { + return if (e.collective != null) 2 else 1; +} + +/// Expert parallel: both ranks take the deep path or neither, with the same deep chunk (the smaller rank's): each +/// rank's room ([rank 0's, rank 1's] summed over the collective). 0: this prompt takes the batched pass on both. +fn agree(e: *prompt.Engine, mine: usize) !usize { + const c = e.collective orelse return mine; + const st = e.stream.*; + const rank = e.layers[0].exl3.?.tp.?.rank; + if (agree_buf == null) agree_buf = try cuda.DeviceBuffer.alloc(e.driver, 256); + const b = agree_buf.?; + var v: [2]f32 = .{ 0, 0 }; + v[rank] = @floatFromInt(mine); + try st.synchronize(); + try b.upload(0, std.mem.asBytes(&v)); + try c.sum(st, b.ptr, 2); + try c.finish(st); + try b.download(0, std.mem.asBytes(&v)); + if (v[0] < 1 or v[1] < 1) return 0; + return @intFromFloat(@min(v[0], v[1])); +} +var agree_buf: ?cuda.DeviceBuffer = null; +var tp_logged = false; + +fn now(io: std.Io) i96 { + return std.Io.Clock.awake.now(io).toNanoseconds(); +} + +/// The hook (Engine.forward): runs `ids` from `seq.fed` and returns true, or returns false having done nothing. +pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks) !bool { + const o = opts(); + if (!applies(e, seq, ids, hooks, o)) return false; + try @import("cuda_exl3_dplan.zig").toHost(); // the pools' residency on the host before planning a prompt + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + const s = try e.reserve(@min(ids.len, e.chunk)); + const sub = s.cap / att_rows * att_rows; + if (sub == 0) return false; + // the deep chunk: as large as asked (whole attention blocks), no larger than the prompt needs; its sub-blocks are + // the pass's chunks, the last one shorter + const rows = @max(att_rows, @min(o.rows, (ids.len + att_rows - 1) / att_rows * att_rows) / att_rows * att_rows); + // the buffers reserved at load (reserveFor) are kept whatever their size; a first prompt without them allocates + const world = worldOf(e); + const have = if (deep) |d| d.owner == @as(*const anyopaque, e) and d.world == world else false; + var room: usize = if (have) deep.?.rows else 0; + if (!have) { + if (reserveAt(e.gpa, e.driver, e, rows, world)) |_| { + room = deep.?.rows; + } else |err| std.log.warn("flashnext exl3 prefill: no room for deep chunks of {d} rows ({t}); the batched pass takes the prompt", .{ rows, err }); + } + // both ranks the same way (a collective a rank skips would hang the other) + room = try agree(e, room); + if (room == 0) return false; + if (world == 2 and !tp_logged) std.log.info("flashnext exl3 prefill: expert-parallel deep chunks of up to {d} rows (rank {d}: its own experts, one all-reduce a layer)", .{ room, e.layers[0].exl3.?.tp.?.rank }); + tp_logged = world == 2 or tp_logged; + const dp = deep.?; + try selectWriteback(dp, o.reduced_branch and world == 2); + const scores = try e.reserveScores(seq.cap, s.cap); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, + .seq = seq.view(scores), + .hooks = hooks, + }; + // equal deep chunks (a short last chunk would pay a whole chunk's expert copies for a few rows) + const cap_rows = @min(rows, room); + const n_chunks = (ids.len + cap_rows - 1) / cap_rows; + const deep_rows = @min(cap_rows, ((ids.len + n_chunks - 1) / n_chunks + att_rows - 1) / att_rows * att_rows); + const rg = if (o.ring) try ringOf(dp, e) else null; + var window_room: usize = 0; + if (o.two_chunk) { + var supported = dp.second != null and rg != null and world == 2 and !dp.reduced_branch and seq.kv8 == null and deep_rows <= window_.max_rows and ids.len > deep_rows; + if (@import("cuda_exl3_graph.zig").enabled()) { + const required = @import("cuda_exl3_graph_memory.zig").twoChunkReservationBytes(); + if (required == 0 or dp.window_graph_reservation < required) supported = false; + } + for (e.layers[0..e.n_layers]) |*layer| if (@import("cuda_exl3_heads.zig").isSplit(layer)) { + supported = false; + }; + window_room = try agree(e, if (supported) deep_rows else 0); + } + if (o.tiles and !o.slots_kernel) { + try warmTiles(e); + tiles_on = true; + } + defer tiles_on = false; + if (rg) |r| try r.begin(p.stream); + defer if (rg) |r| r.end(p.stream, o.reseed) catch |err| std.log.err("flashnext exl3 prefill: ring end: {t}", .{err}); + var at: usize = 0; + while (at < ids.len) { + if (window_room != 0) { + const plan_ = try window_.Plan.init(ids.len - at, window_room, att_rows); + try twoChunk(Tri, dp, &p, s, ids[at..][0..plan_.rows], seq.fed + at, sub, hooks, at, at + plan_.rows == ids.len, seq.logits, o, rg.?, plan_); + at += plan_.rows; + continue; + } + const n = @min(deep_rows, ids.len - at); + const last = at + n == ids.len; + const lg: batch.Logits = if (last) hooks.logits else .none; + const st: u64 = if (hooks.streams != 0) hooks.streams + at * wide * 2 else 0; + try chunk(Tri, dp, &p, s, ids[at..][0..n], seq.fed + at, sub, lg, st, if (last and hooks.logits != .none) seq.logits else null, o, rg); + at += n; + } + if (hooks.logits == .none) try p.stream.synchronize(); + seq.fed += ids.len; + return true; +} + +/// A sub-block's view of the pass's scratch: the per-row buffers a layer keeps across its sub-blocks point into the +/// deep chunk's rows [r0, r0 + n); everything else is the pass's own (reused by every sub-block). +fn view(base: *const batch.Scratch, dp: *const Deep, r0: usize) batch.Scratch { + var v = base.*; + const l = dp.lay; + v.h = dp.dptr(l.h) + r0 * wide * 2; + v.pss = dp.dptr(l.pss) + r0 * pss_n * 4; + v.inj = dp.dptr(l.inj) + r0 * streams * 2; + v.mixed = dp.dptr(l.mixed) + r0 * dims * 2; + v.xsm = dp.dptr(l.xsm) + r0 * (dims / 32) * 4; + v.y = dp.dptr(l.y) + r0 * slots * dims * 4; + v.wts = dp.dptr(l.wts) + r0 * slots * 4; + v.picks = dp.dptr(l.picks) + r0 * slots * 4; + if (dp.world > 1) v.moe_f32 = dp.dptr(l.sum) + r0 * dims * 4; + v.ple_emb = dp.dptr(l.ple_emb) + r0 * dims * 2; + v.ple_xs = dp.dptr(l.ple_xs) + r0 * (dims / 32) * 4; + v.ids = dp.dptr(l.ids) + r0 * 4; + v.pos = dp.dptr(l.pos) + (r0 / att_rows) * 16; + return v; +} + +fn selectWriteback(dp: *Deep, reduced_branch: bool) !void { + if (reduced_branch and dp.unit_weights == null) { + const weights: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(dp.hl.ids, dp.rows * 4))); + @memset(weights, 1); + dp.unit_weights = try cuda.DeviceBuffer.fromHost(dp.dev.d, std.mem.sliceAsBytes(weights)); + } + if (dp.reduced_branch != reduced_branch) @import("cuda_exl3_graph.zig").clearPrefill(); + dp.reduced_branch = reduced_branch; +} + +fn pendingWriteback(p: anytype, dp: *const Deep, layer: ?usize, r0: usize) void { + p.pending = layer; + p.pending_f32 = layer != null and dp.reduced_branch; + p.pending_f32_weights = if (p.pending_f32) dp.unit_weights.?.ptr + r0 * 4 else null; +} + +fn compactLast(p: anytype, reduced_branch: bool, row: usize) !void { + if (row == 0) return; + const s = p.s; + try p.copy(s.h, s.h + row * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + row * pss_n * 4, pss_n * 4); + try p.copy(s.inj, s.inj + row * streams * 2, streams * 2); + if (reduced_branch) { + try p.copy(s.moe_f32, s.moe_f32 + row * dims * 4, dims * 4); + } else { + try p.copy(s.y, s.y + row * slots * dims * 4, slots * dims * 4); + try p.copy(s.wts, s.wts + row * slots * 4, slots * 4); + } +} + +const Gather = struct { + ngram: *const @import("cuda_ple.zig").Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + + fn main(g: *Gather) void { + const t0 = now(g.io); + defer g.ns = now(g.io) - t0; + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + } +}; + +fn upload(st: cuda.Stream, dst: u64, bytes: []const u8) !void { + try st.d.check(st.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, st.handle), "cuMemcpyHtoDAsync"); +} + +fn windowStage(dp: *Deep, p: anytype, ids: []const u32, pos0: usize) !void { + const n = ids.len; + try p.stream.synchronize(); + const id32 = std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.ids, n * 4)); + for (id32, ids) |*destination, id| destination.* = @intCast(id); + const blocks = (n + att_rows - 1) / att_rows; + const positions = std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.pos, blocks * 16)); + @memset(positions, 0); + for (0..blocks) |block| positions[4 * block] = @intCast(pos0 + block * att_rows); + try upload(p.stream, dp.dptr(dp.lay.ids), dp.hbytes(dp.hl.ids, n * 4)); + try upload(p.stream, dp.dptr(dp.lay.pos), dp.hbytes(dp.hl.pos, blocks * 16)); + if (p.m.table_dense) { + try p.tri.embedDense(dp.dptr(dp.lay.ids), p.m.table_w, dp.dptr(dp.lay.h), dims, streams, n); + } else { + try p.tri.embed(dp.dptr(dp.lay.ids), p.m.table_w, p.m.table_s, p.m.table_b, dp.dptr(dp.lay.h), dims, streams, n); + } +} + +fn windowInputs(dp: *Deep, p: anytype, ids: []const u32, pos0: usize, layer: usize) !void { + const blocks = (ids.len + att_rows - 1) / att_rows; + const positions_offset = dp.hl.pos + (layer + 1) * dp.hl.pos_stride; + const positions: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(positions_offset, blocks * 16))); + @memset(positions, 0); + for (0..blocks) |block| positions[4 * block] = @intCast(pos0 + block * att_rows); + try upload(p.stream, dp.dptr(dp.lay.pos), dp.hbytes(positions_offset, blocks * 16)); + if (layer == ple_layer) { + const embedding: []u16 = @alignCast(std.mem.bytesAsSlice(u16, dp.hbytes(dp.hl.emb, ids.len * dims * 2))); + const sums: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(dp.hl.xs, ids.len * (dims / 32) * 4))); + try p.m.ngram.rows(p.seq.history[0..], ids, embedding, sums); + try upload(p.stream, dp.dptr(dp.lay.ple_emb), std.mem.sliceAsBytes(embedding)); + try upload(p.stream, dp.dptr(dp.lay.ple_xs), std.mem.sliceAsBytes(sums)); + } +} + +fn windowDense(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, layer: usize, n: usize, pos0: usize, sub: usize, waits: []const cuda.Event) !void { + const graph = @import("cuda_exl3_graph.zig"); + const lw = &p.m.layers[layer]; + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_sub) |block| views[block] = view(base, dp, block * sub); + defer p.s = base; + p.s = base; + p.rows = n; + p.pos0 = pos0; + const capture = try graph.prefillBegin(p, lw.linear, layer, n, sub, pos0, dp.dev.ptr, base, tiles_on, waits, dp.dptr(dp.lay.pos)); + errdefer graph.prefillAbort(p, capture); + if (capture != .replayed) for (0..n_sub) |block| { + const at = block * sub; + const count = @min(sub, n - at); + p.s = &views[block]; + p.rows = count; + p.pos0 = pos0 + at; + pendingWriteback(p, dp, if (layer == 0) null else layer - 1, at); + if (block < waits.len) switch (capture) { + .eager => try p.stream.wait(waits[block]), + .capturing => try p.stream.waitExternal(waits[block]), + .replayed => unreachable, + }; + try layerBlock(Tri, p, lw, layer, count); + }; + try graph.prefillEnd(p, capture); + dp.red_n = 0; +} + +fn windowFinish(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, n: usize, sub: usize, logits: batch.Logits, streams_out: u64, logits_host: ?[]u8) !void { + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_sub) |block| views[block] = view(base, dp, block * sub); + defer p.s = base; + stats.chunks += 1; + stats.rows += n; + const last_block = n_sub - 1; + const last_rows = n - last_block * sub; + if (streams_out != 0) { + for (0..n_sub) |block| { + const count = @min(sub, n - block * sub); + p.s = &views[block]; + p.rows = count; + pendingWriteback(p, dp, wt.layers_n - 1, block * sub); + try p.writeback(0, count); + try p.copy(streams_out + block * sub * wide * 2, views[block].h, count * wide * 2); + } + } + if (logits == .none) return; + p.s = &views[last_block]; + p.rows = last_rows; + const scratch = p.s; + const row = last_rows - 1; + if (streams_out != 0) { + if (row > 0) { + try p.copy(scratch.h, scratch.h + row * wide * 2, wide * 2); + try p.copy(scratch.pss, scratch.pss + row * pss_n * 4, pss_n * 4); + } + } else { + try compactLast(p, dp.reduced_branch, row); + pendingWriteback(p, dp, wt.layers_n - 1, last_block * sub); + try p.writeback(0, 1); + } + try p.readout(p.m.mixer, 0, 1); + if (p.m.head.n == vocab) { + try p.lane(scratch.mixed, scratch.xsm, p.m.head, scratch.logits, scratch.dummy, 1, 1); + } else { + const collective = p.m.collective orelse return error.TensorParallelUnavailable; + try p.lane(scratch.mixed, scratch.xsm, p.m.head, scratch.head_local, scratch.dummy, 1, 1); + try collective.gather(p.stream, scratch.head_local, scratch.logits, vocab / 2); + } + if (logits_host) |destination| { + if (p.m.collective) |collective| try collective.finish(p.stream); + try p.stream.synchronize(); + try p.stream.d.check(p.stream.d.api.cuMemcpyDtoH_v2(destination.ptr, scratch.logits, vocab * 2), "cuMemcpyDtoH"); + } +} + +fn twoChunk(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, ids: []const u32, pos0: usize, sub: usize, hooks: batch.Hooks, output_at: usize, last: bool, logits_host: []u8, o: Options, ring: *ring_.Ring, plan_: window_.Plan) !void { + const buffers = [_]*Deep{ dp, dp.second.? }; + for (plan_.segments[0..plan_.count]) |segment| { + if ((segment.rows + sub - 1) / sub > max_sub) return error.UnexpectedTensor; + } + try p.stream.synchronize(); + try ring.newChunk(); + history_.prepare(&dp.history, &dp.history_shape, .{ .sub = sub, .lo = ring.own_lo, .hi = ring.own_hi }); + for (0..@min(o.ring_ahead + 1, wt.layers_n)) |layer| _ = try ring.prefetch(layer, @min(routed_n_, dp.used[layer]), keepFree(ring, o), .initial_prefetch); + for (plan_.segments[0..plan_.count], 0..) |segment, index| { + try windowStage(buffers[index], p, ids[segment.offset..][0..segment.rows], pos0 + segment.offset); + } + var schedule = plan_.iterator(p.m.layers.len); + while (schedule.next()) |step| { + const active = buffers[step.chunk]; + const segment = plan_.segments[step.chunk]; + try windowInputs(active, p, ids[segment.offset..][0..segment.rows], pos0 + segment.offset, step.layer); + try windowDense(Tri, active, p, base, step.layer, segment.rows, pos0 + segment.offset, sub, active.red_ev[0..active.red_n]); + if (active != dp) { + active.used = dp.used; + active.history = dp.history; + } + const layer = &p.m.layers[step.layer]; + try ringMoe(Tri, active, ring, p, layer, step.layer, segment.rows, o); + if (active != dp) { + dp.used = active.used; + dp.history = active.history; + } + try rankSum(active, p, layer.exl3.?, segment.rows, sub); + if (step.layer + 1 == p.m.layers.len) { + for (active.red_ev[0..active.red_n]) |event| try p.stream.wait(event); + active.red_n = 0; + } + } + for (plan_.segments[0..plan_.count], 0..) |segment, index| { + const final = last and index + 1 == plan_.count; + const output = if (hooks.streams == 0) 0 else hooks.streams + (output_at + segment.offset) * wide * 2; + try windowFinish(Tri, buffers[index], p, base, segment.rows, sub, if (final) hooks.logits else .none, output, if (final and hooks.logits != .none) logits_host else null); + } +} + +fn chunk(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, ids: []const u32, pos0: usize, sub: usize, logits: batch.Logits, streams_out: u64, logits_host: ?[]u8, o: Options, rg: ?*ring_.Ring) !void { + const n = ids.len; + const st = p.stream; + const hl = dp.hl; + // the previous chunk's uploads read the pinned rows: they landed before its head (a synchronize), or below + try st.synchronize(); + if (rg) |r| { + try r.newChunk(); + history_.prepare(&dp.history, &dp.history_shape, .{ .sub = sub, .lo = r.own_lo, .hi = r.own_hi }); + // the first layers' experts start streaming now, under the embedding and layer 0's blocks + for (0..@min(o.ring_ahead + 1, wt.layers_n)) |l| _ = try r.prefetch(l, @min(routed_n_, dp.used[l]), keepFree(r, o), .initial_prefetch); + } + const id32 = std.mem.bytesAsSlice(i32, dp.hbytes(hl.ids, n * 4)); + for (id32, ids) |*x, v| x.* = @intCast(v); + const blocks = (n + att_rows - 1) / att_rows; + const pos = std.mem.bytesAsSlice(i32, dp.hbytes(hl.pos, blocks * 16)); + @memset(pos, 0); + for (0..blocks) |b| pos[4 * b] = @intCast(pos0 + b * att_rows); + try upload(st, dp.dptr(dp.lay.ids), dp.hbytes(hl.ids, n * 4)); + try upload(st, dp.dptr(dp.lay.pos), dp.hbytes(hl.pos, blocks * 16)); + const emb: []u16 = @alignCast(std.mem.bytesAsSlice(u16, dp.hbytes(hl.emb, n * dims * 2))); + const xs: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(hl.xs, n * (dims / 32) * 4))); + var gather: Gather = .{ .ngram = p.m.ngram, .history = p.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = p.io }; + var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); + defer if (thread) |th| th.join(); + const t0 = now(p.io); + defer if (profile()) { + st.synchronize() catch {}; + std.debug.print("exl3 prefill: {d} rows at {d} in {d:.1} ms (n-gram gather {d:.1} ms on the host)\n", .{ n, pos0, @as(f64, @floatFromInt(now(p.io) - t0)) / 1e6, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); + }; + p.rows = n; + p.pos0 = pos0; + pendingWriteback(p, dp, null, 0); + if (p.m.table_dense) { + try p.tri.embedDense(dp.dptr(dp.lay.ids), p.m.table_w, dp.dptr(dp.lay.h), dims, streams, n); + } else { + try p.tri.embed(dp.dptr(dp.lay.ids), p.m.table_w, p.m.table_s, p.m.table_b, dp.dptr(dp.lay.h), dims, streams, n); + } + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + if (n_sub > views.len) return error.UnexpectedTensor; + for (0..n_sub) |b| views[b] = view(base, dp, b * sub); + defer p.s = base; + for (p.m.layers, 0..) |*lw, layer| { + if (layer == ple_layer) { + if (thread) |th| th.join(); + thread = null; + if (gather.failed) |err| return err; + try upload(st, dp.dptr(dp.lay.ple_emb), std.mem.sliceAsBytes(emb)); + try upload(st, dp.dptr(dp.lay.ple_xs), std.mem.sliceAsBytes(xs)); + } + const pg = try @import("cuda_exl3_graph.zig").prefillBegin(p, lw.linear, layer, n, sub, pos0, dp.dev.ptr, base, tiles_on, dp.red_ev[0..dp.red_n], dp.dptr(dp.lay.pos)); + errdefer @import("cuda_exl3_graph.zig").prefillAbort(p, pg); + if (pg != .replayed) for (0..n_sub) |b| { + const m = @min(sub, n - b * sub); + p.s = &views[b]; + p.rows = m; + p.pos0 = pos0 + b * sub; + pendingWriteback(p, dp, if (layer > 0) layer - 1 else null, b * sub); + if (b < dp.red_n) switch (pg) { + .eager => try st.wait(dp.red_ev[b]), + .capturing => try st.waitExternal(dp.red_ev[b]), + .replayed => unreachable, + }; + try layerBlock(Tri, p, lw, layer, m); + }; + try @import("cuda_exl3_graph.zig").prefillEnd(p, pg); + dp.red_n = 0; + if (rg) |r| try ringMoe(Tri, dp, r, p, lw, layer, n, o) else try moe(Tri, dp, p, lw, layer, n, o); + try rankSum(dp, p, lw.exl3.?, n, sub); + } + for (dp.red_ev[0..dp.red_n]) |ev| try st.wait(ev); + dp.red_n = 0; + stats.chunks += 1; + stats.rows += n; + // the last layer's experts written back where the caller keeps the streams, then the head on the last row + const last_b = n_sub - 1; + const m_last = n - last_b * sub; + if (streams_out != 0) { + for (0..n_sub) |b| { + const m = @min(sub, n - b * sub); + p.s = &views[b]; + p.rows = m; + pendingWriteback(p, dp, wt.layers_n - 1, b * sub); + try p.writeback(0, m); + try p.copy(streams_out + b * sub * wide * 2, views[b].h, m * wide * 2); + } + } + if (logits == .none) return; + p.s = &views[last_b]; + p.rows = m_last; + const s = p.s; + const r = m_last - 1; + if (streams_out != 0) { + if (r > 0) { + try p.copy(s.h, s.h + r * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + } + } else { + try compactLast(p, dp.reduced_branch, r); + pendingWriteback(p, dp, wt.layers_n - 1, last_b * sub); + try p.writeback(0, 1); + } + try p.readout(p.m.mixer, 0, 1); + if (p.m.head.n == vocab) { + try p.lane(s.mixed, s.xsm, p.m.head, s.logits, s.dummy, 1, 1); + } else { + // the head split by vocabulary (expert parallel): each rank's half, gathered in vocabulary order (projectHead) + const c = p.m.collective orelse return error.TensorParallelUnavailable; + try p.lane(s.mixed, s.xsm, p.m.head, s.head_local, s.dummy, 1, 1); + try c.gather(st, s.head_local, s.logits, vocab / 2); + } + if (logits_host) |lh| { + if (p.m.collective) |c| try c.finish(st); + try st.synchronize(); + try st.d.check(st.d.api.cuMemcpyDtoH_v2(lh.ptr, s.logits, vocab * 2), "cuMemcpyDtoH"); + } +} + +fn profile() bool { + const v = std.c.getenv("TF_FLASHNEXT_PROFILE") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// One sub-block of one layer, as Pass.chunk runs it, up to the router's picks (the experts come after every +/// sub-block, expert-major). +fn layerBlock(comptime Tri: type, p: *batch.Pass(Tri), lw: *const @import("cuda_prompt_load.zig").Layer, layer: usize, m: usize) !void { + const s = p.s; + if (layer == ple_layer) { + if (p.pending != null) try p.writeback(0, m); + // Pass.pleBlock for a prompt's rows (no lanes) + try p.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, m, 1); + try p.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, m, 4); + try p.sumSlices(s.ple_part, s.ple_vals, true, m * dims, 4); + try p.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, m, dims, streams); + try p.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, p.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, m, dims, streams, 4, 3); + try p.shiftWindow(p.seq.tail.ptr, s.ple_nrow, wide, m, ple_tail_n, wide); + } + try p.writeback(0, m); + try p.readout(lw.attn_hc, 0, m); + if (lw.linear) { + try p.gdnBlock(lw, wt.linearIndex(layer), .{}); + } else { + try p.attnBlock(lw, wt.attnIndex(layer)); + } + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(lw.mlp_hc, 0, m); + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, m); + try p.tri.topkMany(s.router_l, s.picks, s.wts, m); +} + +/// Lines a pass may fill: half the layer's lines (the other half holds the previous pass, still computing). +fn passLines(pool: *const pool_.Pool, o: Options) usize { + const lines = pool.opt.lines; + return @max(1, if (o.pass_lines_half) lines / 2 else lines); +} + +/// The layer's experts over all `n` rows of the deep chunk, expert-major. +fn moe(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), lw: *const @import("cuda_prompt_load.zig").Layer, layer: usize, n: usize, o: Options) !void { + const x = lw.exl3.?; + const st = p.stream; + const pairs = n * slots; + const hl = dp.hl; + const t0 = now(p.io); + // the picks of every row on the host: the layer's only wait + const host_picks: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.picks, pairs * 4))); + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host_picks.ptr, dp.dptr(dp.lay.picks), pairs * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + maskOwned(host_picks, x.tp); + const t1 = now(p.io); + try plan(dp, x, host_picks, o); + // members and items of every pass in one upload (the pinned bytes stay put until the next layer's wait) + const members: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.members, pairs * 4))); + const items: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.items, 3 * pk.maxItems(pairs, table_n) * 4))); + const counts = build(dp, host_picks, members, items); + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(members[0..counts.members])); + try upload(st, dp.dptr(dp.lay.items), std.mem.sliceAsBytes(items[0 .. 3 * counts.items])); + const t2 = now(p.io); + stats.host_wait_ns += @intCast(t1 - t0); + stats.plan_ns += @intCast(t2 - t1); + stats.members += counts.members; + const pool = x.pool; + var fake_rows: usize = 0; + for (dp.passes[0..dp.n_passes]) |ps| { + var r: ?*const pool_.Round = null; + if (pool) |pl| { + const ids_ = dp.order[ps.id_lo..ps.id_hi]; + fake_rows = fakePicks(dp, ids_); + if (fake_rows > 0) { + r = try pl.planUnobserved(dp.fake[0 .. fake_rows * slots], fake_rows); + if (r.?.n_passes != 1) return error.Exl3PrefillPass; + try pl.stage(r.?, 0, st); + stats.copies += r.?.misses_(); + } + } + const l = if (pool) |pl| pl.tables() else x.layer; + if (o.slots_kernel) { + try slotsPass(dp, x, st, l, host_picks, ps, n); + if (r) |rr| try pool.?.release(rr, 0, st); + stats.passes += 1; + continue; + } + const dv: pk.Device = .{ .members = dp.dptr(dp.lay.members) + ps.m_lo * 4, .items = dp.dptr(dp.lay.items) + ps.i_lo * 12, .scratch = dp.dptr(dp.lay.kscratch) }; + try pk.run(&dp.kern, st, l, dp.dptr(dp.lay.mixed), dims, dp.dptr(dp.lay.picks), slots, ps.m_n, ps.i_n, dv, dp.dptr(dp.lay.y)); + if (r) |rr| try pool.?.release(rr, 0, st); + stats.passes += 1; + } + // the next layer's first pass, staged on the copy engine while this layer's experts and the next blocks run + if (o.prefetch and layer + 1 < p.m.layers.len) { + if (p.m.layers[layer + 1].exl3) |nx| if (nx.pool) |npl| try prefetch(dp, npl, x.tp, o); + } + p.pending = layer; +} + +/// Expert parallel: another rank's picks become -1 on the host (this rank plans, copies and computes its own pairs). +fn maskOwned(picks: []i32, r: ?tp_.Rank) void { + const rk = r orelse return; + for (picks) |*pk_| if (pk_.* >= 0 and !rk.owns(@intCast(pk_.*))) { + pk_.* = -1; + }; +} + +/// Expert parallel: after a layer's passes, each row's branch over this rank's slots (on the compute stream), then on +/// the reduce stream one all-reduce of [rows, 2560] fp32 a sub-block (both ranks issue the same calls with the same +/// rows in the same order) and its sum placed where the write-back reads it; the next layer's sub-block b waits for +/// red_ev[b] only. TF_FLASHNEXT_EXL3_PREFILL_OVERLAP=0: one all-reduce of all rows on the compute stream. +fn rankSum(dp: *Deep, p: anytype, x: *const load.Experts, n: usize, sub: usize) !void { + const r = x.tp orelse return; + if (r.world == 1) return; + const c = p.m.collective orelse return error.TensorParallelUnavailable; + const st = p.stream; + const sum = dp.dptr(dp.lay.sum); + const y = dp.dptr(dp.lay.y); + const wts = dp.dptr(dp.lay.wts); + const t0 = if (profile()) now(p.io) else 0; + try pk.tpCombine(&dp.kern, st, y, wts, dp.dptr(dp.lay.picks), sum, n, slots, dims, r.lo, r.hi); + if (!overlap() or dp.red == null) { + try c.sum(st, sum, n * dims); + if (!dp.reduced_branch) try pk.tpPlace(&dp.kern, st, sum, y, wts, n, slots, dims); + } else { + const rs = dp.red.?; + try dp.red_go.record(st); + try rs.wait(dp.red_go); + const n_sub = (n + sub - 1) / sub; + for (0..n_sub) |b| { + const r0 = b * sub; + const m = @min(sub, n - r0); + try c.sum(rs, sum + r0 * dims * 4, m * dims); + if (!dp.reduced_branch) try pk.tpPlace(&dp.kern, rs, sum + r0 * dims * 4, y + r0 * slots * dims * 4, wts + r0 * slots * 4, m, slots, dims); + try dp.red_ev[b].record(rs); + } + dp.red_n = n_sub; + } + if (profile()) { + if (dp.red) |rs| try c.finish(rs); + try c.finish(st); + stats.reduce_ns += @intCast(now(p.io) - t0); + } + stats.reduces += 1; +} + +fn overlap() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_OVERLAP") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Picks for a pool round that holds `ids` (routed, < 512): rows of 11 slots, -1 after the last. +fn fakePicks(dp: *Deep, ids: []const u16) usize { + var routed: usize = 0; + for (ids) |id| routed += @intFromBool(id != shared_id); + if (routed == 0) return 0; + const rows = (routed + slots - 1) / slots; + std.debug.assert(rows * slots <= dp.fake.len); + @memset(dp.fake[0 .. rows * slots], -1); + var k: usize = 0; + for (ids) |id| { + if (id == shared_id) continue; + dp.fake[k] = id; + k += 1; + } + return rows; +} + +/// Stages up to a pass's worth of the next layer's experts that its pool does not hold (lowest ids first): a deep +/// chunk uses nearly every expert, so any of them is a copy its MoE would make anyway. +fn prefetch(dp: *Deep, pool: *pool_.Pool, own: ?tp_.Rank, o: Options) !void { + const want = passLines(pool, o); + var ids: [table_n]u16 = undefined; + var k: usize = 0; + for (0..wt.experts_n) |id| { + if (k == want) break; + if (own) |rk| if (!rk.owns(id)) continue; + if (pool_.offload.lineOf(pool.v, pool.opt.layer, @intCast(id)) != null) continue; + ids[k] = @intCast(id); + k += 1; + } + const rows = fakePicks(dp, ids[0..k]); + if (rows == 0) return; + const r = try pool.plan(dp.fake[0 .. rows * slots], rows); + if (r.n_passes != 1) return error.Exl3PrefillPass; + // the copies wait for nothing on the compute stream: the aux stream takes the gates and the release + try pool.stage(r, 0, dp.aux); + try pool.release(r, 0, dp.aux); + stats.prefetched += r.misses_(); + stats.copies += r.misses_(); +} + +/// The layer's passes: pass 0 the shared expert and every expert the pool holds (no copy), then the others in id +/// order, at most `passLines` new experts and `member_cap` members a pass. Fills order, passes, cnt. +fn plan(dp: *Deep, x: *const load.Experts, picks: []const i32, o: Options) !void { + @memset(&dp.cnt, 0); + for (picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + const cap = dp.lay.member_cap; + var k: usize = 0; + dp.n_passes = 0; + var held: [table_n]bool = @splat(false); + held[shared_id] = true; + const per: usize = if (x.pool) |pl| passLines(pl, o) else table_n; + if (x.pool) |pl| { + for (0..wt.experts_n) |id| held[id] = pool_.offload.lineOf(pl.v, pl.opt.layer, @intCast(id)) != null; + } else { + @memset(&held, true); + } + // pass 0 (and more if the held experts' members exceed the cap): held experts, the shared one first + var lo: usize = 0; + var members: usize = 0; + const order_ids = struct { + fn close(d: *Deep, lo_: *usize, k_: usize, members_: *usize) void { + if (k_ == lo_.*) return; + d.passes[d.n_passes] = .{ .id_lo = @intCast(lo_.*), .id_hi = @intCast(k_), .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + d.n_passes += 1; + lo_.* = k_; + members_.* = 0; + } + }; + if (dp.cnt[shared_id] > 0) { + dp.order[k] = shared_id; + k += 1; + members += dp.cnt[shared_id]; + } + for (0..wt.experts_n) |id| { + if (!held[id] or dp.cnt[id] == 0) continue; + if (members + dp.cnt[id] > cap) order_ids.close(dp, &lo, k, &members); + dp.order[k] = @intCast(id); + k += 1; + members += dp.cnt[id]; + } + order_ids.close(dp, &lo, k, &members); + // the misses, `per` new experts a pass + var in_pass: usize = 0; + for (0..wt.experts_n) |id| { + if (held[id] or dp.cnt[id] == 0) continue; + if (in_pass == per or members + dp.cnt[id] > cap) { + order_ids.close(dp, &lo, k, &members); + in_pass = 0; + } + dp.order[k] = @intCast(id); + k += 1; + members += dp.cnt[id]; + in_pass += 1; + } + order_ids.close(dp, &lo, k, &members); +} + +/// Members (pairs grouped by expert in pass order, pairs ascending) and items (first relative to the pass's first +/// member) of every pass; fills each pass's ranges. +fn build(dp: *Deep, picks: []const i32, members: []i32, items: []i32) struct { members: usize, items: usize } { + var m_at: u32 = 0; + var i_at: u32 = 0; + @memset(&dp.start, 0); + for (dp.passes[0..dp.n_passes]) |*ps| { + ps.m_lo = m_at; + ps.i_lo = i_at; + for (dp.order[ps.id_lo..ps.id_hi]) |id| { + const c = dp.cnt[id]; + dp.start[id] = m_at; + var off: u32 = 0; + while (off < c) : (off += pk.tile) { + items[3 * i_at] = id; + items[3 * i_at + 1] = @intCast(m_at - ps.m_lo + off); + items[3 * i_at + 2] = @intCast(@min(@as(u32, pk.tile), c - off)); + i_at += 1; + } + m_at += c; + } + ps.m_n = m_at - ps.m_lo; + ps.i_n = i_at - ps.i_lo; + } + for (picks, 0..) |pk_, i| { + if (pk_ < 0 or pk_ >= table_n) continue; + const id: usize = @intCast(pk_); + members[dp.start[id]] = @intCast(i); + dp.start[id] += 1; + } + return .{ .members = m_at, .items = i_at }; +} + +/// One pass through Experts.runSlots: the deep chunk's picks with every expert outside the pass masked to -1, in the +/// batched pass's 256-row blocks (a check: these are the batched pass's kernels and bits). +fn slotsPass(dp: *Deep, x: *load.Experts, st: cuda.Stream, l: cuda.exl3_experts.Layer, picks: []const i32, ps: Pass, n: usize) !void { + var in_pass: [table_n]bool = @splat(false); + for (dp.order[ps.id_lo..ps.id_hi]) |id| in_pass[id] = true; + const pairs = n * slots; + // the members' pinned region holds the masked picks (members are not used in this mode); the previous pass's + // upload must have read it + try st.synchronize(); + const masked: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.members, pairs * 4))); + for (picks[0..pairs], masked) |pk_, *o| o.* = if (pk_ >= 0 and pk_ < table_n and in_pass[@intCast(pk_)]) pk_ else -1; + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(masked)); + const ex = try x.ops.grouped(); + const block: usize = @import("cuda_exl3_pass.zig").moe_rows; + const bytes = try cuda.exl3_experts.scratchBytes(l, block, slots); + const base = try x.ops.scratch(st, &x.ops.moe, bytes); + var r0: usize = 0; + while (r0 < n) : (r0 += block) { + const rows = @min(block, n - r0); + var sc = try cuda.exl3_experts.scratch(base, l, block, slots); + sc.y = dp.dptr(dp.lay.y) + r0 * slots * dims * 4; + try ex.runSlots(st, l, dp.dptr(dp.lay.mixed) + r0 * dims * 2, dims, .bf16, dp.dptr(dp.lay.members) + r0 * slots * 4, rows, sc); + } +} + +// ---------------------------------------------------------------------------------------------------- the MTP head + +/// The hook in cuda_mtp_ops.zig `forwardHead` (a lane's prompt chunk through the MTP head, after the main pass kept +/// its streams): the head's one layer EXPERT-MAJOR over the whole chunk. forwardHead runs the layer in 256-row pieces, +/// each piece's MoE copying nearly every expert of the head's pool (about 1.7 s per 8192 rows of copies); here every +/// piece runs up to the router's picks over the deep chunk's rows (the attention reads the head's cache in piece +/// order, as there), then the experts once over all rows (`moe`: each expert staged once), then each piece's +/// write-back. The head is one layer, so nothing a piece reads depends on another piece's experts: the same maths as +/// forwardHead but the experts' fp32 sum order (the 64-member kernel). Declines (false) unless the engine holds deep +/// buffers of at least `next.len` rows and the head's experts are pooled on this rank (whole: the MTP head is rank 0's +/// under expert parallel), for a prompt chunk (no scores) of at least TF_FLASHNEXT_EXL3_PREFILL_MIN rows without media. +/// TF_FLASHNEXT_EXL3_PREFILL_HEAD=0 keeps forwardHead. +pub fn head(comptime Tri: type, o: anytype, d: *@import("cuda_mtp.zig").Draft, next: []const u32, rows: u64, scores: bool) !bool { + const op = opts(); + const n = next.len; + // A paired window hands the head up to two deep chunks of rows: it runs them as two slices (headSplit). + if (!scores and deep != null and n > deep.?.rows) { + const first = headSplit(n, deep.?.rows, op.min_rows) orelse return false; + if (!try headSlice(Tri, o, d, next[0..first], rows, false)) return false; + if (!try headSlice(Tri, o, d, next[first..], rows + first * wide * 2, false)) return error.UnexpectedTensor; + return true; + } + return headSlice(Tri, o, d, next, rows, scores); +} + +/// Where a head call of `n` rows over a paired window splits into two deep-sized slices, or null when it cannot. +fn headSplit(n: usize, deep_rows: usize, min_rows: usize) ?usize { + if (n <= deep_rows or n > 2 * deep_rows) return null; + const first = ((n + 1) / 2 + att_rows - 1) / att_rows * att_rows; + if (first >= n or first > deep_rows or n - first < min_rows) return null; + return first; +} + +test "exl3: a paired window's head splits into two deep slices that cannot decline after the first runs" { + try std.testing.expectEqual(@as(?usize, 4096), headSplit(8192, 4096, 512)); + try std.testing.expectEqual(@as(?usize, null), headSplit(4096, 4096, 512)); + try std.testing.expectEqual(@as(?usize, null), headSplit(8193, 4096, 512)); + var n: usize = 4097; + while (n <= 8192) : (n += 61) { + const first = headSplit(n, 4096, 512) orelse continue; + try std.testing.expect(first <= 4096 and first % att_rows == 0 and n - first <= first and n - first >= 512); + } +} + +fn headSlice(comptime Tri: type, o: anytype, d: *@import("cuda_mtp.zig").Draft, next: []const u32, rows: u64, scores: bool) !bool { + const op = opts(); + const n = next.len; + if (!op.enabled or !headEnabled() or scores or n < op.min_rows or d.media != null) return false; + const e = o.e; + const dp = deep orelse return false; + if (dp.owner != @as(*const anyopaque, e) or dp.rows < n) return false; + const h = o.h; + const x = h.layer.exl3 orelse return false; + if (x.tp != null or x.pool == null or (x.layer.act_mode != 1 and x.layer.act_mode != 0)) return false; + if (d.len + n > d.cap) return error.PromptTooLong; + try @import("cuda_exl3_dplan.zig").toHost(); // the head's pool may be device-planned (cuda_exl3_chain.zig) + const mtp = @import("cuda_mtp.zig"); + const tri = o.tri; + const want = @max(@as(usize, 32), 2 * @min(n, mtp.piece_rows)); + const s = try e.reserve(want); + const piece = @max(@as(usize, 1), @min(mtp.piece_rows, s.cap / 2)) / att_rows * att_rows; + if (piece == 0) return false; + const n_pieces = (n + piece - 1) / piece; + if (n_pieces > max_sub) return false; + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, + .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .hooks = .{ .logits = .none }, + }; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_pieces) |b| views[b] = view(s, dp, b * piece); + defer p.s = s; + if (op.tiles and !op.slots_kernel) { + try warmTiles(e); + tiles_on = true; + } + defer tiles_on = false; + const lw = &h.layer; + for (0..n_pieces) |b| { + const at = b * piece; + const m = @min(piece, n - at); + p.s = &views[b]; + const v = p.s; + p.rows = m; + p.pos0 = d.len + at; + p.pending = null; + try p.stage(next[at..][0..m]); + // forwardHead's piece up to the router (the embedding branch, the hidden branch, attention over the head's cache) + if (e.table.dense) { + try tri.embedDense(v.ids, e.table.w.ptr, v.ple_keys, dims, streams, m); + } else { + try tri.embed(v.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, v.ple_keys, dims, streams, m); + } + try tri.rmsnorm(v.ple_keys, h.norm_e.ptr, v.mixed, v.xsm, eps, wide, m, dims, dims); + try p.lane(v.mixed, v.xsm, h.fc_e, v.dummy, v.part, m, mtp.fc_sk); + try p.sumSlices(v.part, v.branch, true, m * dims, mtp.fc_sk); + try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, v.norm, v.gq, eps, wide, m, wide, wide); + try p.lane(v.norm, v.gq, h.fc_h, v.dummy, v.part, streams * m, mtp.fc_sk); + try p.sumSlices(v.part, v.ple_gated, true, streams * m * dims, mtp.fc_sk); + try tri.addStreams(v.branch, v.ple_gated, v.h, m, dims, streams); + try p.writeback(0, m); + try p.readout(lw.attn_hc, 0, m); + try p.attnBlock(lw, 0); + try tri.hcWriteback(v.h, v.h, v.pss, v.part, "*fp32", v.inj, v.dummy, "*bf16", v.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(lw.mlp_hc, 0, m); + try tri.routerMany(v.mixed, lw.router.ptr, v.router_l, dims, m); + try tri.topkMany(v.router_l, v.picks, v.wts, m); + } + // the head's experts once over the chunk (its own pool, passes of half its lines) + try moe(Tri, dp, &p, lw, 0, n, op); + for (0..n_pieces) |b| { + const m = @min(piece, n - b * piece); + p.s = &views[b]; + p.rows = m; + p.pending = 0; + try p.writeback(0, m); + } + const last = &views[n_pieces - 1]; + try p.copy(d.out.ptr, last.h + (n - 1 - (n_pieces - 1) * piece) * wide * 2, wide * 2); + d.len += n; + stats.head_rows += n; + if (!head_logged) std.log.info("flashnext exl3 prefill: the MTP head runs expert-major over deep chunks ({d} rows first)", .{n}); + head_logged = true; + return true; +} +var head_logged = false; + +fn headEnabled() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_HEAD") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +// ---------------------------------------------------------------------------------------------------- the ring + +/// The ring over the engine's decoder pools, made once; null (and the per-layer passes) if they cannot form one. +fn ringOf(dp: *Deep, e: *prompt.Engine) !?*ring_.Ring { + if (dp.ring) |r| return r; + if (dp.ring_failed) return null; + var pools: [wt.layers_n]*pool_.Pool = undefined; + for (e.layers[0..e.n_layers], 0..) |*l, i| { + const x = l.exl3 orelse { + dp.ring_failed = true; + return null; + }; + pools[i] = x.pool orelse { + dp.ring_failed = true; + return null; + }; + } + dp.ring = ring_.Ring.init(dp.gpa, e.driver, &pools, e.telemetry.enabled) catch |err| { + std.log.warn("flashnext exl3 prefill: no line ring ({t}); each layer stages into its own lines", .{err}); + dp.ring_failed = true; + return null; + }; + if (e.layers[0].exl3.?.tp) |rk| { + dp.ring.?.own_lo = rk.lo; + dp.ring.?.own_hi = @min(rk.hi, routed_n_); + } + std.log.info("flashnext exl3 prefill: a ring of {d} expert slots over the {d} layers' pool lines", .{ dp.ring.?.slots(), wt.layers_n }); + return dp.ring; +} + +/// Slots a prefetch leaves free: a layer's worth of misses and two passes. +fn keepFree(r: *const ring_.Ring, o: Options) usize { + return @min(r.slots() / 2, routed_n_ + 2 * o.ring_pass); +} + +/// The layer's experts over all `n` rows of the deep chunk, through the ring: the experts the ring already holds +/// first (in load order), then the misses (hottest first), `ring_pass` a pass, each pass's copies queued a pass ahead. +fn ringMoe(comptime Tri: type, dp: *Deep, r: *ring_.Ring, p: *batch.Pass(Tri), lw: *const @import("cuda_prompt_load.zig").Layer, layer: usize, n: usize, o: Options) !void { + const x = lw.exl3.?; + const pool = x.pool.?; + const st = p.stream; + const pairs = n * slots; + const hl = dp.hl; + const t0 = now(p.io); + const host_picks: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.picks, pairs * 4))); + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host_picks.ptr, dp.dptr(dp.lay.picks), pairs * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + maskOwned(host_picks, x.tp); + const t1 = now(p.io); + @memset(&dp.cnt, 0); + for (host_picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + r.drop(layer, dp.cnt[0..routed_n_]); + // pass order: shared, then the held experts by load batch, then the misses hottest first + var held: [routed_n_]u16 = undefined; + var n_held: usize = 0; + var miss: [routed_n_]u16 = undefined; + var n_miss: usize = 0; + var used: usize = 0; + for (0..routed_n_) |id| { + if (dp.cnt[id] == 0) continue; + used += 1; + r.counts[layer][id] +|= dp.cnt[id]; + const present = r.isLoaded(layer, id); + r.stats.observeDemand(present, dp.cnt[id]); + if (present) { + held[n_held] = @intCast(id); + n_held += 1; + } else { + miss[n_miss] = @intCast(id); + n_miss += 1; + } + } + dp.used[layer] = used; + const ByBatch = struct { + r: *const ring_.Ring, + l: usize, + fn lt(c: @This(), a: u16, b: u16) bool { + return c.r.batch_of[c.l][a] < c.r.batch_of[c.l][b]; + } + }; + std.mem.sort(u16, held[0..n_held], ByBatch{ .r = r, .l = layer }, ByBatch.lt); + const ByCount = struct { + c: *const [table_n]u32, + fn gt(c: @This(), a: u16, b: u16) bool { + if (c.c[a] != c.c[b]) return c.c[a] > c.c[b]; + return a < b; + } + }; + std.mem.sort(u16, miss[0..n_miss], ByCount{ .c = &dp.cnt }, ByCount.gt); + // passes: at most ring_pass experts and member_cap members each; the shared expert opens pass 0 + var k: usize = 0; + var lo: usize = 0; + var members: usize = 0; + var in_pass: usize = 0; + dp.n_passes = 0; + const cap = dp.lay.member_cap; + const close = struct { + fn f(d: *Deep, lo_: *usize, k_: usize, members_: *usize, in_pass_: *usize) void { + if (k_ == lo_.*) return; + d.passes[d.n_passes] = .{ .id_lo = @intCast(lo_.*), .id_hi = @intCast(k_), .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + d.n_passes += 1; + lo_.* = k_; + members_.* = 0; + in_pass_.* = 0; + } + }.f; + if (dp.cnt[shared_id] > 0) { + dp.order[k] = shared_id; + k += 1; + members += dp.cnt[shared_id]; + } + for ([2][]const u16{ held[0..n_held], miss[0..n_miss] }) |list| for (list) |id| { + if (in_pass == o.ring_pass or members + dp.cnt[id] > cap) close(dp, &lo, k, &members, &in_pass); + dp.order[k] = id; + k += 1; + members += dp.cnt[id]; + in_pass += 1; + }; + close(dp, &lo, k, &members, &in_pass); + const staging0 = now(p.io); + const copies0 = r.stats.copies; + for (0..@min(2, dp.n_passes)) |j| try loadPass(dp, r, layer, j, false); + const staging_ns = now(p.io) - staging0; + const members_h: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.members, pairs * 4))); + const items_h: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.items, 3 * pk.maxItems(pairs, table_n) * 4))); + const counts = build(dp, host_picks, members_h, items_h); + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(members_h[0..counts.members])); + try upload(st, dp.dptr(dp.lay.items), std.mem.sliceAsBytes(items_h[0 .. 3 * counts.items])); + const t2 = now(p.io); + stats.host_wait_ns += @intCast(t1 - t0); + stats.plan_ns += @intCast(t2 - t1 - staging_ns); + stats.members += counts.members; + const l = pool.tables(); + // the first two passes' copies, then each pass: the next pass's copies, the gate, the kernel, the release + for (dp.passes[0..dp.n_passes], 0..) |ps, j| { + if (j + 2 < dp.n_passes) try loadPass(dp, r, layer, j + 2, true); + try loadPass(dp, r, layer, j, false); + const ids_ = dp.order[ps.id_lo..ps.id_hi]; + try r.gate(st, layer, ids_); + if (o.slots_kernel) { + try slotsPass(dp, x, st, l, host_picks, ps, n); + } else { + const dv: pk.Device = .{ .members = dp.dptr(dp.lay.members) + ps.m_lo * 4, .items = dp.dptr(dp.lay.items) + ps.i_lo * 12, .scratch = dp.dptr(dp.lay.kscratch) }; + try pk.run(&dp.kern, st, l, dp.dptr(dp.lay.mixed), dims, dp.dptr(dp.lay.picks), slots, ps.m_n, ps.i_n, dv, dp.dptr(dp.lay.y)); + } + try r.release(st, layer, ids_); + stats.passes += 1; + } + stats.copies += r.stats.copies - copies0; + dp.history[layer] = .{ .rows = n, .used = used }; + const c1 = r.stats.copies; + for (1..o.ring_ahead + 1) |ahead| { + const nl = layer + ahead; + if (nl >= wt.layers_n) break; + _ = try r.prefetch(nl, history_.want(dp.history[nl], n, used, routed_n_), keepFree(r, o), .cross_layer_prefetch); + } + stats.prefetched += r.stats.copies - c1; + stats.copies += r.stats.copies - c1; + p.pending = layer; +} + +/// Queues the copies of pass `j`'s routed experts the ring does not hold; `best_effort`: stop quietly when full. +fn loadPass(dp: *Deep, r: *ring_.Ring, layer: usize, j: usize, best_effort: bool) !void { + const ps = dp.passes[j]; + for (dp.order[ps.id_lo..ps.id_hi]) |id| { + if (id == shared_id) continue; + if (try r.load(layer, id, .demand)) continue; + if (best_effort) return; + return error.Exl3PrefillRingFull; + } + try r.flush(); +} + +// ---------------------------------------------------------------------------------------------------- dense tiles + +/// Set while a deep prefill runs with Options.tiles: cuda_exl3_pass.part sends prompt rows to Dense.prompt on tiles. +var tiles_on: bool = false; +/// Rows below which a projection keeps the reference kernels (tiles lose at a few rows). +const tiles_min_rows: usize = 64; + +/// The hook in cuda_exl3_pass.part: true when it ran the part on the tiles (bf16 or fp32 out). +pub fn tilesPart(de: *cuda.exl3.Dense, s: cuda.Stream, w: cuda.exl3.Weight, x: u64, m: usize, y: u64, f32_out: bool) !bool { + if (!tiles_on or m < tiles_min_rows) return false; + _ = cuda.exl3.Dense.candidateWeight(w) catch return false; + if (de.candidate_k == null or de.candidate_k.?.k_max < w.k) return false; + de.implementation = .tiles; + defer de.implementation = .reference; + try de.prompt(s, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); + return true; +} + +/// Loads the tiles once with the widest input any EXL3 projection has (no allocation inside a pass). +fn warmTiles(e: *prompt.Engine) !void { + const x = e.layers[0].exl3 orelse return; + const de = try x.ops.linear(); + const k_max = @max(x.ops.max_k, wt.dims * 3); + if (de.candidate_k != null and de.candidate_k.?.k_max >= k_max) return; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(de.d, cuda.kernels.exl3_candidate); + errdefer mod.unload(); + const k = try cuda.exl3_candidate.Kernels.resolve(de.d, mod, k_max); + if (de.candidate_k) |*old| old.deinit(); + if (de.candidate_mod) |*old| old.unload(); + de.candidate_mod = mod; + de.candidate_k = k; +} + +test "exl3: reduced prefill writeback preserves slot rounding sub-blocks and tail selection" { + try @import("cuda_exl3_prefill_writeback_test.zig").run(Deep, Layout.of, view, pendingWriteback, compactLast, selectWriteback); + try std.testing.expect(!(Options{}).reduced_branch); +} + +test "exl3: two-chunk scheduling and reservation contracts" { + _ = window_; + try std.testing.expect(!(Options{}).two_chunk); + try std.testing.expectEqual(@as(usize, 4096), stepCapacity(4096, false)); + try std.testing.expectEqual(@as(usize, 8192), stepCapacity(4096, true)); + try std.testing.expectEqual(@as(usize, 8192), stepCapacity(8192, true)); +} + +test "exl3: two-chunk arenas retain hidden and deferred HC operands without aliasing" { + const rows = 4096; + const layout = Layout.of(rows, 2); + const persistent = Persistent.of(rows); + try std.testing.expectEqual(@as(usize, 727399424), layout.total); + const host_layout = HostLayout.of(rows); + try std.testing.expectEqual(@as(usize, 160 << 20), 2 * rows * wide * 2); + var driver: cuda.Driver = undefined; + var first: Deep = undefined; + first.world = 2; + first.rows = rows; + first.shared = null; + first.lay = layout; + first.dev = .{ .d = &driver, .ptr = 0x100000, .len = layout.total }; + first.sum_dev = null; + var second = first; + second.dev.ptr += layout.total; + second.dev.len = persistent.total; + second.rows = rows; + second.shared = &first; + var base: batch.Scratch = undefined; + base.moe_f32 = 0; + const fields = .{ "h", "pss", "inj", "mixed", "xsm", "y", "wts", "picks", "ple_emb", "ple_xs", "ids", "pos", "members", "items", "kscratch", "sum" }; + inline for (fields, 0..) |field, index| { + const start = @field(layout, field); + const end = if (index + 1 < fields.len) @field(layout, fields[index + 1]) else layout.total; + try std.testing.expectEqual(@as(usize, 0), start % 256); + try std.testing.expect(start < end and end <= layout.total); + try std.testing.expect(first.dev.ptr + end <= second.dev.ptr); + } + for ([_]usize{ 0, 1024, 3072 }) |offset| { + const a = view(&base, &first, offset); + const b = view(&base, &second, offset); + inline for (.{ "h", "pss", "inj", "y", "wts" }) |field| { + const second_start = second.dev.ptr + @field(persistent, field); + try std.testing.expect(@field(b, field) >= second_start and @field(b, field) < second.dev.ptr + persistent.total); + try std.testing.expect(@field(a, field) < second.dev.ptr); + } + inline for (.{ "mixed", "xsm", "picks", "ple_emb", "ple_xs", "ids", "pos", "moe_f32" }) |field| { + try std.testing.expectEqual(@field(a, field), @field(b, field)); + } + } + inline for (.{ "members", "items", "kscratch", "sum" }) |field| try std.testing.expectEqual(first.dptr(@field(layout, field)), second.dptr(@field(layout, field))); + try std.testing.expectEqual(@as(usize, 546127872), persistent.total); + try std.testing.expect(persistent.total < layout.total); + try std.testing.expect(host_layout.total > rows * dims * 2); + try std.testing.expect(!window_.admits(layout.total - 1, 0, layout.total)); + try std.testing.expect(window_.admits(layout.total, 0, layout.total)); +} + +test "exl3: two-chunk memory refusal performs no CUDA allocation" { + const saved_deep = deep; + const saved_options = options; + const graph = @import("cuda_exl3_graph.zig"); + const saved_force = graph.force; + defer { + deep = saved_deep; + options = saved_options; + graph.force = saved_force; + } + graph.force = .off; + options = .{ .two_chunk = true }; + var driver: cuda.Driver = undefined; + var buffers: Deep = undefined; + buffers.dev = .{ .d = &driver, .ptr = 0x100000, .len = 1 }; + buffers.world = 2; + buffers.rows = 4096; + buffers.lay = Layout.of(4096, 2); + buffers.second = null; + deep = &buffers; + try std.testing.expectEqual(@as(u64, 0), try reserveTwoChunk(&driver, Persistent.of(4096).total - 1, 0, 0)); + try std.testing.expect(buffers.second == null); +} + +test "exl3: two-chunk PLE input staging advances mutable history only in chronological PLE segments" { + const Fixture = struct { + calls: usize = 0, + before: [2][2]i64 = undefined, + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + fn copy(_: cuda.abi.DevicePtr, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + pub fn rows(self: *@This(), history: []i64, ids: []const u32, embedding: []u16, sums: []f32) !void { + self.before[self.calls] = history[0..2].*; + self.calls += 1; + for (ids) |id| { + history[0] = history[1]; + history[1] = id; + } + @memset(embedding, 0); + @memset(sums, 0); + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuStreamSynchronize = Fixture.sync; + driver.api.cuMemcpyHtoDAsync_v2 = Fixture.copy; + const host_layout = HostLayout.of(256); + const host = try std.testing.allocator.alignedAlloc(u8, .@"16", host_layout.total); + defer std.testing.allocator.free(host); + var buffers: Deep = undefined; + buffers.rows = 256; + buffers.shared = null; + buffers.lay = Layout.of(256, 2); + buffers.hl = host_layout; + buffers.host = .{ .d = &driver, .bytes = host }; + buffers.dev = .{ .d = &driver, .ptr = 0x100000, .len = buffers.lay.total }; + var fixture: Fixture = .{}; + var history: [2]i64 = .{ -2, -1 }; + var pass = .{ .stream = cuda.Stream{ .d = &driver, .handle = null }, .m = .{ .ngram = &fixture }, .seq = .{ .history = &history } }; + var ids: [512]u32 = undefined; + for (&ids, 1..) |*id, value| id.* = @intCast(value); + for (0..2) |index| try windowInputs(&buffers, &pass, ids[index * 256 ..][0..256], index * 256, 0); + try std.testing.expectEqual(@as(usize, 0), fixture.calls); + try std.testing.expectEqualSlices(i64, &.{ -2, -1 }, &history); + for (0..2) |index| try windowInputs(&buffers, &pass, ids[index * 256 ..][0..256], index * 256, ple_layer); + try std.testing.expectEqual(@as(usize, 2), fixture.calls); + try std.testing.expectEqualSlices(i64, &.{ -2, -1 }, &fixture.before[0]); + try std.testing.expectEqualSlices(i64, &.{ 255, 256 }, &fixture.before[1]); + try std.testing.expectEqualSlices(i64, &.{ 511, 512 }, &history); +} + +test "exl3: prefill passes keep each expert's pairs together, items relative to their pass" { + var dp: Deep = undefined; + // pairs 0..8: experts 5, 512, 2 / 2, 512, 7 / 5, 512, 5 (-1 and out-of-table ids skipped) + const picks = [_]i32{ 5, 512, 2, 2, 512, 7, 5, 512, 5, -1, 600 }; + @memset(&dp.cnt, 0); + for (picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + // pass 0: the shared expert and 5; pass 1: 2 and 7 + dp.order[0] = shared_id; + dp.order[1] = 5; + dp.order[2] = 2; + dp.order[3] = 7; + dp.passes[0] = .{ .id_lo = 0, .id_hi = 2, .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + dp.passes[1] = .{ .id_lo = 2, .id_hi = 4, .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + dp.n_passes = 2; + var members: [16]i32 = undefined; + var items: [3 * 16]i32 = undefined; + const c = build(&dp, &picks, &members, &items); + try std.testing.expectEqual(@as(usize, 9), c.members); + try std.testing.expectEqual(@as(usize, 4), c.items); + try std.testing.expectEqualSlices(i32, &.{ 1, 4, 7, 0, 6, 8, 2, 3, 5 }, members[0..9]); + try std.testing.expectEqual(Pass{ .id_lo = 0, .id_hi = 2, .m_lo = 0, .m_n = 6, .i_lo = 0, .i_n = 2 }, dp.passes[0]); + try std.testing.expectEqual(Pass{ .id_lo = 2, .id_hi = 4, .m_lo = 6, .m_n = 3, .i_lo = 2, .i_n = 2 }, dp.passes[1]); + // (expert, first within its pass, count) + try std.testing.expectEqualSlices(i32, &.{ 512, 0, 3, 5, 3, 3, 2, 0, 2, 7, 2, 1 }, items[0..12]); +} + +test "exl3: expert-parallel prefill masks the other rank's picks and keeps the shared expert on rank 1" { + const r0 = try tp_.Rank.init(2, 0); + const r1 = try tp_.Rank.init(2, 1); + const picks = [_]i32{ 5, 300, 512, -1, 255, 256, 511, 0 }; + var a = picks; + maskOwned(&a, r0); + try std.testing.expectEqualSlices(i32, &.{ 5, -1, -1, -1, 255, -1, -1, 0 }, &a); + var b = picks; + maskOwned(&b, r1); + try std.testing.expectEqualSlices(i32, &.{ -1, 300, 512, -1, -1, 256, 511, -1 }, &b); + var c = picks; + maskOwned(&c, null); + try std.testing.expectEqualSlices(i32, &picks, &c); + // every pick lands on exactly one rank + for (picks, a, b) |p, x, y| if (p >= 0) try std.testing.expect((x == p) != (y == p)); + // the sum buffer exists only at world 2 + try std.testing.expectEqual(deviceBytesAt(4096, 1), deviceBytes(4096)); + try std.testing.expect(deviceBytesAt(4096, 2) >= deviceBytes(4096) + 4096 * dims * 4); +} + +test "exl3: prefetch history follows deep engine ownership reuse replacement and release" { + const Fixture = struct { + fn unload(_: cuda.abi.Module) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn stream(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn freeHost(_: ?*anyopaque) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn refuse(_: *cuda.abi.DevicePtr, _: usize) callconv(.c) cuda.abi.Result { + return cuda.abi.error_not_ready; + } + fn install(driver: *const cuda.Driver, owner: *const anyopaque) !*Deep { + const buffers = try std.testing.allocator.create(Deep); + buffers.* = .{ + .gpa = std.testing.allocator, + .owner = owner, + .rows = 4096, + .lay = undefined, + .hl = undefined, + .dev = .{ .d = driver, .ptr = 0, .len = 0 }, + .host = .{ .d = driver, .bytes = @constCast(&.{}) }, + .kern = undefined, + .aux = .{ .d = driver, .handle = null }, + }; + buffers.kern.mod = .{ .d = driver, .handle = null }; + deep = buffers; + return buffers; + } + }; + try std.testing.expect(deep == null); + defer release(); + var driver: cuda.Driver = undefined; + driver.api.cuModuleUnload = Fixture.unload; + driver.api.cuStreamSynchronize = Fixture.stream; + driver.api.cuStreamDestroy_v2 = Fixture.stream; + driver.api.cuMemFreeHost = Fixture.freeHost; + driver.api.cuMemAlloc_v2 = Fixture.refuse; + var owners: [2]u8 = .{ 0, 0 }; + const first = try Fixture.install(&driver, &owners[0]); + first.history[3] = .{ .rows = 4096, .used = 80 }; + first.history_shape = .{ .sub = 256, .lo = 0, .hi = 256 }; + releaseFor(&owners[1]); + try std.testing.expect(deep.? == first); + try reserveAt(std.testing.allocator, &driver, &owners[0], 2048, 1); + try std.testing.expectEqual(@as(usize, 96), history_.want(deep.?.history[3], 4096, 400, routed_n_)); + try std.testing.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &owners[1], 4096, 1)); + try std.testing.expect(deep == null); + const second = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(second.history[3], 4096, 400, routed_n_)); + try std.testing.expect(second.history_shape == null); + second.history[3] = .{ .rows = 4096, .used = 80 }; + try std.testing.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &owners[1], 4096, 2)); + try std.testing.expect(deep == null); + const third = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(third.history[3], 4096, 400, routed_n_)); + third.history[3] = .{ .rows = 4096, .used = 80 }; + releaseFor(&owners[1]); + try std.testing.expect(deep == null); + const fourth = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(fourth.history[3], 4096, 400, routed_n_)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_attribution.zig b/zig/src/families/flashnext/cuda_exl3_prefill_attribution.zig new file mode 100644 index 0000000000..c704649999 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_attribution.zig @@ -0,0 +1,49 @@ +const std = @import("std"); +const telemetry = @import("exl3_telemetry.zig"); + +pub const Source = enum { demand, initial_prefetch, cross_layer_prefetch }; +pub const Counts = struct { + copies: u64 = 0, + requested_bytes: u64 = 0, + unused_copies: u64 = 0, + unused_requested_bytes: u64 = 0, +}; +pub const Totals = [3]Counts; + +const Residency = struct { source: Source, bytes: u64 }; + +pub const State = struct { + slots: []Residency, + totals: Totals = @splat(.{}), + + pub fn init(allocator: std.mem.Allocator, slots: usize, enabled: bool) !?State { + if (!enabled) return null; + return .{ .slots = try allocator.alloc(Residency, slots) }; + } + + pub fn deinit(state: *State, allocator: std.mem.Allocator) void { + allocator.free(state.slots); + } + + pub fn copied(state: *State, slot: usize, source: Source, bytes: u64) void { + state.slots[slot] = .{ .source = source, .bytes = bytes }; + const counts = &state.totals[@backingInt(source)]; + counts.copies += 1; + counts.requested_bytes += bytes; + } + + pub fn unused(state: *State, slot: usize) void { + const residency = state.slots[slot]; + const counts = &state.totals[@backingInt(residency.source)]; + counts.unused_copies += 1; + counts.unused_requested_bytes += residency.bytes; + } +}; + +pub fn publish(totals: Totals, experts: *telemetry.Experts) void { + inline for (std.enums.values(Source), 0..) |source, index| { + inline for (@typeInfo(Counts).@"struct".field_names) |field| { + @field(experts.*, "ring_" ++ @tagName(source) ++ "_" ++ field) = @field(totals[index], field); + } + } +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_attribution_test.zig b/zig/src/families/flashnext/cuda_exl3_prefill_attribution_test.zig new file mode 100644 index 0000000000..2f9eb790e3 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_attribution_test.zig @@ -0,0 +1,310 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const attribution = @import("cuda_exl3_prefill_attribution.zig"); +const ring_module = @import("cuda_exl3_prefill_ring.zig"); +const pool_module = @import("cuda_exl3_pool.zig"); +const telemetry = @import("exl3_telemetry.zig"); +const State = std.meta.Child(@FieldType(ring_module.Ring, "state")); +const allocator = std.testing.allocator; + +const Calls = struct { + copies: usize = 0, + payload_bytes: usize = 0, + records: usize = 0, + queries: usize = 0, + waits: usize = 0, + syncs: usize = 0, +}; + +const Fixture = struct { + driver: cuda.Driver, + ring: *ring_module.Ring, + pool: *pool_module.Pool, + payload: [8]u8, + sources: [512][3][]const u8, + counts: [48][512]u32, + addresses: [2]u64, + offsets: [2][3]u32, + rooms: [2][3]u32, + states: [2]State, + owners: [2]u32, + epochs: [2]u64, + + var calls: Calls = .{}; + var copy_result: cuda.abi.Result = cuda.abi.success; + + fn copy(destination: u64, _: ?*const anyopaque, bytes: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + calls.copies += 1; + if (destination < 0x9000) calls.payload_bytes += bytes; + return copy_result; + } + + fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + calls.records += 1; + return cuda.abi.success; + } + + fn query(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + calls.queries += 1; + return cuda.abi.error_not_ready; + } + + fn wait(_: cuda.abi.Stream, _: cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result { + calls.waits += 1; + return cuda.abi.success; + } + + fn sync(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + calls.syncs += 1; + return cuda.abi.success; + } + + fn init(enabled: bool) !*Fixture { + const fixture = try allocator.create(Fixture); + errdefer allocator.destroy(fixture); + fixture.* = undefined; + fixture.ring = try allocator.create(ring_module.Ring); + errdefer allocator.destroy(fixture.ring); + fixture.pool = try allocator.create(pool_module.Pool); + errdefer allocator.destroy(fixture.pool); + fixture.pool.* = undefined; + fixture.pool.opt = .{ .layer = 0, .lines = 2 }; + fixture.pool.src = .{ .routed = &fixture.sources }; + fixture.pool.counts = &fixture.counts; + fixture.pool.lines.tables.ptr = 0x9000; + fixture.payload = @splat(7); + for (&fixture.sources, 0..) |*source, index| { + const length = index % 4 + 1; + source.* = .{ fixture.payload[0..length], fixture.payload[0 .. length + 1], fixture.payload[0 .. length + 2] }; + } + for (&fixture.counts) |*row| @memset(row, 0); + fixture.addresses = .{ 0x1000, 0x2000 }; + fixture.offsets = @splat(.{ 0, 8, 16 }); + fixture.rooms = @splat(.{ 8, 8, 8 }); + fixture.states = @splat(.free); + fixture.owners = @splat(0); + fixture.epochs = @splat(0); + fixture.driver = undefined; + fixture.driver.api.cuMemcpyHtoDAsync_v2 = copy; + fixture.driver.api.cuEventRecord = record; + fixture.driver.api.cuEventQuery = query; + fixture.driver.api.cuStreamWaitEvent = wait; + fixture.driver.api.cuEventSynchronize = sync; + const ring = fixture.ring; + ring.* = undefined; + ring.gpa = allocator; + ring.d = &fixture.driver; + ring.pools = @splat(fixture.pool); + ring.addr = &fixture.addresses; + ring.off = &fixture.offsets; + ring.room_of = &fixture.rooms; + ring.state = &fixture.states; + ring.owner = &fixture.owners; + ring.epoch_of = &fixture.epochs; + ring.next = 0; + ring.pending_n = 0; + ring.epoch = 0; + ring.copy_ordered_epoch = 0; + ring.batches = 0; + ring.open_layer = null; + ring.open_n = 0; + ring.own_lo = 0; + ring.own_hi = 512; + ring.copy = .{ .d = &fixture.driver, .handle = @ptrFromInt(1) }; + for (&ring.cev, 0..) |*event, index| event.* = .{ .d = &fixture.driver, .handle = @ptrFromInt(index + 1) }; + for (&ring.rev, 0..) |*event, index| event.* = .{ .d = &fixture.driver, .handle = @ptrFromInt(index + 1) }; + ring.tlay = cuda.exl3_pool.Tables.init(513); + const shadow = try allocator.alignedAlloc(u8, .@"16", 48 * ring.tlay.bytes); + errdefer allocator.free(shadow); + @memset(shadow, 0); + ring.shadow = .{ .d = &fixture.driver, .bytes = shadow }; + const snapshots = try allocator.alignedAlloc(u8, .@"16", 64 * ring.tlay.k2[0]); + errdefer allocator.free(snapshots); + @memset(snapshots, 0); + ring.snaps = .{ .d = &fixture.driver, .bytes = snapshots }; + for (&ring.slot) |*row| @memset(row, -1); + for (&ring.counts) |*row| @memset(row, 0); + ring.stats = .{}; + ring.attribution = try attribution.State.init(allocator, 2, enabled); + calls = .{}; + copy_result = cuda.abi.success; + return fixture; + } + + fn deinit(fixture: *Fixture) void { + if (fixture.ring.attribution) |*state| state.deinit(allocator); + allocator.free(fixture.ring.snaps.bytes); + allocator.free(fixture.ring.shadow.bytes); + allocator.destroy(fixture.pool); + allocator.destroy(fixture.ring); + allocator.destroy(fixture); + } + + fn checksum(fixture: *Fixture, unused_bytes: u64) !void { + var sum: attribution.Counts = .{}; + for (fixture.ring.attribution.?.totals) |counts| try telemetry.add(attribution.Counts, &sum, counts); + try std.testing.expectEqual(fixture.ring.stats.copies, sum.copies); + try std.testing.expectEqual(fixture.ring.stats.dma_requested_bytes, sum.requested_bytes); + try std.testing.expectEqual(fixture.ring.stats.wasted, sum.unused_copies); + try std.testing.expectEqual(unused_bytes, sum.unused_requested_bytes); + } + + fn snapshot(fixture: *Fixture) telemetry.Snapshot { + const stats = fixture.ring.stats; + var experts: telemetry.Experts = .{ .ring_dma_copies = stats.copies, .ring_dma_requested_bytes = stats.dma_requested_bytes, .ring_unused_copies = stats.wasted }; + attribution.publish(fixture.ring.attribution.?.totals, &experts); + return .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 1, .ring_owner = 1, .experts = experts, .kv = &.{} }; + } +}; + +test "exl3: prefill attribution counts actual copies once across held retries full rings drop and slot reuse" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(try ring.load(1, 1, .cross_layer_prefetch)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(usize, 15), Fixture.calls.payload_bytes); + const before_retry = ring.attribution.?.totals; + const calls_before_retry = Fixture.calls; + try std.testing.expect(!try ring.load(2, 2, .demand)); + try std.testing.expect(try ring.load(0, 0, .demand)); + try std.testing.expectEqualDeep(before_retry, ring.attribution.?.totals); + try std.testing.expectEqualDeep(calls_before_retry, Fixture.calls); + try std.testing.expectEqual(@as(i32, 0), ring.slot[0][0]); + try std.testing.expectEqual(@as(i32, 1), ring.slot[1][1]); + const used: [512]u32 = @splat(0); + ring.drop(0, &used); + ring.drop(0, &used); + try fixture.checksum(6); + try std.testing.expect(try ring.load(0, 0, .cross_layer_prefetch)); + try fixture.checksum(6); + try std.testing.expect(try ring.load(2, 2, .demand)); + try std.testing.expectEqual(@as(i32, -1), ring.slot[0][0]); + try std.testing.expectEqual(@as(i32, 1), ring.slot[1][1]); + try fixture.checksum(6); + try ring.newChunk(); + try fixture.checksum(27); + try ring.newChunk(); + try fixture.checksum(27); + const totals = ring.attribution.?.totals; + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 12, .unused_copies = 1, .unused_requested_bytes = 12 }, totals[0]); + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 6, .unused_copies = 1, .unused_requested_bytes = 6 }, totals[1]); + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 9, .unused_copies = 1, .unused_requested_bytes = 9 }, totals[2]); + try std.testing.expectEqualDeep(telemetry.Experts{ + .ring_dma_copies = 3, + .ring_dma_requested_bytes = 27, + .ring_unused_copies = 3, + .ring_demand_copies = 1, + .ring_demand_requested_bytes = 12, + .ring_demand_unused_copies = 1, + .ring_demand_unused_requested_bytes = 12, + .ring_initial_prefetch_copies = 1, + .ring_initial_prefetch_requested_bytes = 6, + .ring_initial_prefetch_unused_copies = 1, + .ring_initial_prefetch_unused_requested_bytes = 6, + .ring_cross_layer_prefetch_copies = 1, + .ring_cross_layer_prefetch_requested_bytes = 9, + .ring_cross_layer_prefetch_unused_copies = 1, + .ring_cross_layer_prefetch_unused_requested_bytes = 9, + }, fixture.snapshot().experts); + try std.testing.expectEqual(@as(usize, 0), Fixture.calls.syncs); +} + +test "exl3: prefill attribution retains outstanding provenance across snapshots releases and later chunk reset" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(try ring.load(0, 1, .cross_layer_prefetch)); + const before = fixture.snapshot(); + try ring.gate(ring.copy, 0, &.{ 0, 1 }); + try ring.release(ring.copy, 0, &.{0}); + const used: [512]u32 = @splat(0); + ring.drop(0, &used); + try fixture.checksum(9); + try std.testing.expect(try ring.load(0, 0, .demand)); + try ring.newChunk(); + try fixture.checksum(9); + const after = fixture.snapshot(); + const report = try telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 2048, .steps = 1 }); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_dma_copies); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_dma_requested_bytes); + try std.testing.expectEqual(@as(u64, 1), report.experts.ring_cross_layer_prefetch_unused_copies); + try std.testing.expectEqual(@as(u64, 9), report.experts.ring_cross_layer_prefetch_unused_requested_bytes); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_initial_prefetch_unused_copies); + const text = try report.formatAlloc(allocator); + defer allocator.free(text); + try std.testing.expect(std.mem.indexOf(u8, text, "ring_cross_layer_prefetch_unused_requested_bytes=9") != null); + try std.testing.expect(try ring.load(2, 2, .demand)); + try fixture.checksum(9); + try std.testing.expectEqual(@as(usize, 1), Fixture.calls.queries); + try std.testing.expectEqual(@as(usize, 2), Fixture.calls.waits); + try std.testing.expectEqual(@as(usize, 0), Fixture.calls.syncs); +} + +test "exl3: prefill attribution preserves prefetch ownership and first-copy source during cross-layer topups" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + ring.own_lo = 1; + ring.own_hi = 3; + ring.counts[4][0] = 100; + ring.counts[4][2] = 20; + ring.counts[4][1] = 10; + try std.testing.expectEqual(@as(usize, 1), try ring.prefetch(4, 1, 0, .initial_prefetch)); + try std.testing.expectEqual(@as(usize, 1), try ring.prefetch(4, 2, 0, .cross_layer_prefetch)); + try std.testing.expectEqual(@as(usize, 0), try ring.prefetch(4, 2, 0, .cross_layer_prefetch)); + try std.testing.expect(!ring.isLoaded(4, 0)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[1].copies); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[2].copies); + try ring.newChunk(); + try fixture.checksum(21); +} + +test "exl3: disabled prefill attribution allocates nothing and bookkeeping adds no driver operations" { + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expect((try attribution.State.init(failing.allocator(), 1000, false)) == null); + try std.testing.expectError(error.OutOfMemory, attribution.State.init(failing.allocator(), 1000, true)); + var previous_calls: Calls = undefined; + var previous_stats: ring_module.Stats = undefined; + var previous_image: u64 = 0; + for ([_]bool{ false, true }) |enabled| { + const fixture = try Fixture.init(enabled); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try ring.gate(ring.copy, 0, &.{0}); + try ring.release(ring.copy, 0, &.{0}); + try std.testing.expect(try ring.load(1, 1, .cross_layer_prefetch)); + try ring.newChunk(); + const image = std.hash.Wyhash.hash(0, ring.shadow.bytes); + if (enabled) { + try std.testing.expectEqualDeep(previous_calls, Fixture.calls); + try std.testing.expectEqualDeep(previous_stats, ring.stats); + try std.testing.expectEqual(previous_image, image); + try fixture.checksum(9); + } else { + try std.testing.expect(ring.attribution == null); + previous_calls = Fixture.calls; + previous_stats = ring.stats; + previous_image = image; + } + } +} + +test "exl3: prefill attribution ignores failed transfers without suppressing retries" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + Fixture.copy_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(!ring.isLoaded(0, 0)); + try fixture.checksum(0); + Fixture.copy_result = cuda.abi.success; + try std.testing.expect(try ring.load(0, 0, .demand)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[0].copies); + try std.testing.expectEqual(@as(u64, 0), ring.attribution.?.totals[1].copies); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_binding.zig b/zig/src/families/flashnext/cuda_exl3_prefill_binding.zig new file mode 100644 index 0000000000..471c66e82e --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_binding.zig @@ -0,0 +1,314 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); +const node_query = @import("cuda_exl3_state_binding.zig"); + +pub const Kind = enum { prep, scores, chunks, plan, remap }; +pub const Functions = struct { prep: usize, scores: usize, chunks: usize, ahead: usize, plan: usize, remap: usize }; +pub const Request = struct { + functions: Functions, + rows: usize, + sub: usize, + pos0: usize, + positions: u64, + cap: usize, + origin: usize, + page_rows: usize, + shared: bool, + generation: u64, + keys: u64, + values: u64, + + pub fn validate(self: Request) !void { + if (self.rows == 0 or self.sub == 0 or self.sub % parts.att_rows != 0 or + self.rows > 65536 or self.positions == 0 or self.page_rows == 0 or self.page_rows > parts.att_rows or + self.origin % wt.index_ratio != 0 or (self.shared and self.generation == 0) or self.origin > self.pos0 or self.pos0 > self.cap or self.rows > self.cap - self.pos0 or + self.cap > std.math.maxInt(c_int) or self.keys == 0 or self.values == 0) return error.UnsupportedPrefillBinding; + comptime var functions = std.meta.fieldNames(Functions); + inline for (functions, 0..) |field, index| { + if (@field(self.functions, field) == 0) return error.UnsupportedPrefillBinding; + inline for (functions[0..index]) |previous| { + if (@field(self.functions, field) == @field(self.functions, previous)) return error.UnsupportedPrefillBinding; + } + } + } + + pub fn end(self: Request, block: usize) usize { + return self.pos0 + @min(self.rows, (block + 1) * parts.att_rows); + } + + pub fn sparse(self: Request, block: usize) bool { + return self.end(block) / wt.index_ratio > wt.index_top; + } + + pub fn compatible(self: Request, target: Request) bool { + target.validate() catch return false; + var before = self; + var after = target; + before.pos0 = 0; + after.pos0 = 0; + before.origin = 0; + after.origin = 0; + before.generation = 0; + after.generation = 0; + if (!std.meta.eql(before, after)) return false; + for (0..divUp(self.rows, parts.att_rows)) |block| { + if (self.sparse(block) != target.sparse(block)) return false; + } + return true; + } + + pub fn chunks(self: Request, block: usize) u32 { + return @intCast(divUp(@min(self.end(block), wt.dense_keys), 512)); + } + + pub fn scores(self: Request, block: usize) u32 { + return @intCast(divUp(@min(parts.scoreStride(self.cap), divUp(self.end(block), wt.index_ratio)), 64)); + } +}; + +fn divUp(value: usize, divisor: usize) usize { + return (value + divisor - 1) / divisor; +} + +pub fn schema(kind: Kind, shared: bool) []const u8 { + return switch (kind) { + .prep => &.{ 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 4, 4, 4, 4, 4, 4, 4 }, + .scores => &.{ 8, 8, 8, 8, 4, 4, 4, 4 }, + .chunks => &.{ 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 4, 4, 4, 4, 4, 4, 4, 4 }, + .plan => if (shared) + &.{ 8, 8, 8, 4, 4, 4, 4, 4, 8, 8, 8, 4, 8, 8, 8, 4, 8, 8, 8, 8, 8, 8, 8, 4, 4 } + else + &.{ 8, 8, 8, 4, 4, 4, 4, 4, 8, 8, 8, 4, 8, 8, 8, 4, 8, 8, 8, 8 }, + .remap => &.{ 8, 8, 8, 4, 4, 4, 8, 4, 8 }, + }; +} + +pub const Record = struct { + node: cuda.graph.Node, + params: cuda.abi.KernelNodeParams, + kind: Kind, + block: usize = 0, + values: [25]u64 = @splat(0), + + pub fn read(node: cuda.graph.Node, params: cuda.abi.KernelNodeParams, kind: Kind, request: Request) !Record { + if (params.func == null or params.extra != null or params.params == null) return error.UnsupportedPrefillBinding; + var record: Record = .{ .node = node, .params = params, .kind = kind }; + record.params.params = null; + for (schema(kind, request.shared), 0..) |size, index| { + const source = params.params.?[index] orelse return error.UnsupportedPrefillBinding; + @memcpy(std.mem.asBytes(&record.values[index])[0..size], @as([*]const u8, @ptrCast(source))[0..size]); + } + const position_index: ?usize = switch (kind) { + .prep => 1, + .scores => 2, + .chunks => 3, + else => null, + }; + if (position_index) |index| { + const position = record.values[index]; + if (position < request.positions or (position - request.positions) % 16 != 0) return error.UnsupportedPrefillBinding; + record.block = @intCast((position - request.positions) / 16); + if (record.block >= divUp(request.rows, parts.att_rows)) return error.UnsupportedPrefillBinding; + } + const block_rows = @min(parts.att_rows, request.rows - record.block * parts.att_rows); + switch (kind) { + .prep => { + const offset = record.block * parts.att_rows; + if (offset % request.sub != 0 or params.grid_x != @min(request.sub, request.rows - offset) or + params.grid_y != wt.q_heads + wt.kv_heads + wt.index_heads + 1 or params.grid_z != 1 or params.block_x != 256 or + record.values[7] != request.keys -% (request.origin * row_bytes) or + record.values[8] != request.values -% (request.origin * row_bytes)) return error.UnsupportedPrefillBinding; + }, + .scores => { + if (!request.sparse(record.block) or params.grid_x != block_rows or params.grid_y != request.scores(record.block) or params.grid_z != 1 or params.block_x != 256 or + record.values[4] != parts.scoreStride(request.cap)) return error.UnsupportedPrefillBinding; + }, + .chunks => { + if (params.grid_z != request.chunks(record.block) or params.grid_y != wt.kv_heads or params.block_x != 128 or + params.grid_x == 0 or params.grid_x > @min(request.page_rows, block_rows)) return error.UnsupportedPrefillBinding; + }, + .plan => if (record.values[7] != request.origin or (request.shared and record.values[22] != request.generation) or record.values[3] == 0 or record.values[3] > request.page_rows or params.grid_x != 1 or params.block_x != 1024) return error.UnsupportedPrefillBinding, + .remap => if (record.values[4] != request.origin or params.grid_x == 0 or params.grid_x > request.page_rows or params.block_x != 256) return error.UnsupportedPrefillBinding, + } + if (params.block_y != 1 or params.block_z != 1) return error.UnsupportedPrefillBinding; + return record; + } + + pub fn changed(self: Record, before: Request, target: Request) bool { + return switch (self.kind) { + .prep, .remap => before.origin != target.origin, + .plan => before.origin != target.origin or before.generation != target.generation, + .scores => before.scores(self.block) != target.scores(self.block), + .chunks => before.chunks(self.block) != target.chunks(self.block), + }; + } + + pub fn rebound(self: Record, target: Request) Record { + var record = self; + switch (record.kind) { + .prep => { + record.values[7] = target.keys -% (target.origin * row_bytes); + record.values[8] = target.values -% (target.origin * row_bytes); + }, + .scores => record.params.grid_y = target.scores(record.block), + .chunks => record.params.grid_z = target.chunks(record.block), + .plan => { + record.values[7] = target.origin; + if (target.shared) record.values[22] = target.generation; + }, + .remap => record.values[4] = target.origin, + } + return record; + } +}; + +const row_bytes = wt.kv_heads * wt.head_dim * 2; + +pub fn requestFor(pass: anytype, linear: bool, layer: usize, rows: usize, sub: usize, pos0: usize, positions: u64) ?Request { + if (linear or pass.tri.set != null or pass.seq.media != null or pass.seq.kv8 != null or pass.lanes.len != 0) return null; + const paging = @import("cuda_exl3_kvpage.zig"); + if (paging.trace != null or @import("cuda_exl3_heads.zig").isSplit(&pass.m.layers[layer])) return null; + const pg = pass.seq.paged orelse return null; + const glue = pass.tri.own orelse return null; + const index = wt.attnIndex(layer); + const plan = if (pg.cfg.shared_lanes == 0) pg.k.plan else (pg.shared_arena.?.kernels orelse return null).plan; + const remap = if (pg.cfg.shared_lanes == 0) pg.k.remap else (pg.shared_arena.?.kernels orelse return null).remap; + const request: Request = .{ + .functions = .{ .prep = @intFromPtr(glue.f.prep.handle), .scores = @intFromPtr(glue.f.scores.handle), .chunks = @intFromPtr(glue.f.chunks.handle), .ahead = @intFromPtr(glue.f.chunks_ahead.handle), .plan = @intFromPtr(plan.handle), .remap = @intFromPtr(remap.handle) }, + .rows = rows, + .sub = sub, + .pos0 = pos0, + .positions = positions, + .cap = pass.seq.cap, + .origin = pg.win0[index], + .page_rows = paging.pagingRows(pg.cfg, pg.win0[index], parts.att_rows) catch return null, + .shared = pg.cfg.shared_lanes != 0, + .generation = if (pg.cfg.shared_lanes != 0) pg.shared_sequence.generation else 0, + .keys = pg.devRows(index, false), + .values = pg.devRows(index, true), + }; + request.validate() catch return null; + return request; +} + +pub const Binding = struct { + graph: cuda.graph.Graph, + request: Request, + records: []Record, + poisoned: bool = false, + + pub fn create(allocator: std.mem.Allocator, graph: cuda.graph.Graph, request: Request) !*Binding { + return createWithNodeType(allocator, graph, request, node_query.nodeType(graph.d)); + } + + pub fn createWithNodeType(allocator: std.mem.Allocator, graph: cuda.graph.Graph, request: Request, node_type: ?node_query.NodeType) !*Binding { + try request.validate(); + const query = node_type orelse return error.UnsupportedPrefillBinding; + const count = try graph.nodeCount(); + const nodes = try allocator.alloc(cuda.graph.Node, count); + defer allocator.free(nodes); + const actual = try graph.nodes(nodes); + if (actual.len != count) return error.UnsupportedPrefillBinding; + var records: std.array_list.Managed(Record) = .init(allocator); + defer records.deinit(); + for (actual) |node| { + var kind_id: c_uint = undefined; + try graph.d.check(query(node, &kind_id), "cuGraphNodeGetType(prefill binding)"); + if (kind_id == 4) return error.UnsupportedPrefillBinding; + if (kind_id != 0) continue; + var params: cuda.abi.KernelNodeParams = undefined; + try graph.d.check(graph.d.api.cuGraphKernelNodeGetParams_v2(node, ¶ms), "cuGraphKernelNodeGetParams(prefill binding)"); + const function = @intFromPtr(params.func); + const functions = request.functions; + const kind: Kind = if (function == functions.prep) .prep else if (function == functions.scores) .scores else if (function == functions.chunks or function == functions.ahead) .chunks else if (function == functions.plan) .plan else if (function == functions.remap) .remap else continue; + try records.append(try Record.read(node, params, kind, request)); + } + try validateRecords(request, records.items); + const binding = try allocator.create(Binding); + errdefer allocator.destroy(binding); + binding.* = .{ .graph = graph, .request = request, .records = try records.toOwnedSlice() }; + return binding; + } + + pub fn deinit(self: *Binding, allocator: std.mem.Allocator) void { + self.graph.deinit(); + allocator.free(self.records); + allocator.destroy(self); + } + + pub fn apply(self: *Binding, backend: anytype, target: Request) !void { + if (self.poisoned or !self.request.compatible(target)) return error.UnsupportedPrefillBinding; + errdefer self.poisoned = true; + for (self.records) |*record| { + if (!record.changed(self.request, target)) continue; + const next = record.rebound(target); + try backend.set(next, target.shared); + record.* = next; + } + self.request = target; + } +}; + +pub fn validateRecords(request: Request, records: []const Record) !void { + var preps: [256]usize = @splat(0); + var scores: [256]usize = @splat(0); + var chunks: [256]usize = @splat(0); + var chunk_rows: [256]usize = @splat(0); + var plans: usize = 0; + var remaps: usize = 0; + var plan_rows: usize = 0; + var remap_rows: usize = 0; + for (records) |record| switch (record.kind) { + .prep => preps[record.block] += 1, + .scores => scores[record.block] += 1, + .chunks => { + chunks[record.block] += 1; + chunk_rows[record.block] += record.params.grid_x; + }, + .plan => { + plans += 1; + plan_rows += @intCast(record.values[3]); + }, + .remap => { + remaps += 1; + remap_rows += record.params.grid_x; + }, + }; + var expected_pages: usize = 0; + var expected_rows: usize = 0; + for (0..divUp(request.rows, parts.att_rows)) |block| { + const rows = @min(parts.att_rows, request.rows - block * parts.att_rows); + const pages = divUp(rows, request.page_rows); + const sparse = request.sparse(block); + if (preps[block] != @intFromBool((block * parts.att_rows) % request.sub == 0) or + scores[block] != @intFromBool(sparse) or chunks[block] != pages or chunk_rows[block] != rows) return error.UnsupportedPrefillBinding; + if (sparse) { + expected_pages += pages; + expected_rows += rows; + } + } + if (plans != expected_pages or remaps != expected_pages or plan_rows != expected_rows or remap_rows != expected_rows) return error.UnsupportedPrefillBinding; +} + +pub const Backend = struct { + exec: cuda.graph.Exec, + + pub fn set(self: Backend, record: Record, shared: bool) !void { + var args: cuda.Args = .{}; + for (schema(record.kind, shared), 0..) |size, index| { + if (size == 8) args.add(record.values[index]) else args.add(@as(u32, @truncate(record.values[index]))); + } + const params = record.params; + try self.exec.setKernel(record.node, .{ .d = self.exec.d, .handle = params.func }, .{ + .grid = .{ .x = params.grid_x, .y = params.grid_y, .z = params.grid_z }, + .block = .{ .x = params.block_x, .y = params.block_y, .z = params.block_z }, + .shared = params.shared_bytes, + }, &args); + } +}; + +test { + _ = @import("cuda_exl3_prefill_binding_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_binding_test.zig b/zig/src/families/flashnext/cuda_exl3_prefill_binding_test.zig new file mode 100644 index 0000000000..407bede787 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_binding_test.zig @@ -0,0 +1,269 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const binding = @import("cuda_exl3_prefill_binding.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); + +fn request(shared: bool) binding.Request { + return .{ + .functions = .{ .prep = 10, .scores = 11, .chunks = 12, .ahead = 13, .plan = 14, .remap = 15 }, + .rows = 512, + .sub = 256, + .pos0 = 4096, + .positions = 0x1000000010, + .cap = 32768, + .origin = 1024, + .page_rows = 256, + .shared = shared, + .generation = if (shared) 0x1234567800000001 else 0, + .keys = 0x2000000020, + .values = 0x3000000030, + }; +} + +fn recordFor(req: binding.Request, kind: binding.Kind, block: usize) binding.Record { + const function = switch (kind) { + .prep => req.functions.prep, + .scores => req.functions.scores, + .chunks => req.functions.chunks, + .plan => req.functions.plan, + .remap => req.functions.remap, + }; + var record: binding.Record = .{ + .node = @ptrFromInt(1), + .params = .{ .func = @ptrFromInt(function), .grid_x = 256, .grid_y = 1, .grid_z = 1, .block_x = 256, .block_y = 1, .block_z = 1, .shared_bytes = 0, .params = null, .extra = null }, + .kind = kind, + .block = block, + }; + for (binding.schema(kind, req.shared), 0..) |size, index| { + record.values[index] = if (size == 8) 0x8000000000 + index else 100 + index; + } + const position = req.positions + block * 16; + switch (kind) { + .prep => { + record.params.grid_y = wt.q_heads + wt.kv_heads + wt.index_heads + 1; + record.values[1] = position; + record.values[7] = req.keys -% (req.origin * wt.kv_heads * wt.head_dim * 2); + record.values[8] = req.values -% (req.origin * wt.kv_heads * wt.head_dim * 2); + }, + .scores => { + record.values[2] = position; + record.values[4] = parts.scoreStride(req.cap); + record.params.grid_y = req.scores(block); + }, + .chunks => { + record.params.grid_y = wt.kv_heads; + record.params.block_x = 128; + record.values[3] = position; + record.params.grid_z = req.chunks(block); + }, + .plan => { + record.values[3] = 256; + record.values[7] = req.origin; + record.params.grid_x = 1; + record.params.block_x = 1024; + if (req.shared) record.values[22] = req.generation; + }, + .remap => record.values[4] = req.origin, + } + return record; +} + +fn recordsFor(req: binding.Request) [10]binding.Record { + var records: [10]binding.Record = undefined; + for (0..2) |block| for (std.enums.values(binding.Kind), 0..) |kind, index| { + records[block * 5 + index] = recordFor(req, kind, block); + records[block * 5 + index].node = @ptrFromInt(block * 5 + index + 1); + }; + return records; +} + +test "exl3: prefill origin binding keeps dense-sparse and paging-split topology explicit" { + const base = request(true); + try base.validate(); + var target = base; + target.pos0 = 8192; + target.origin = 4096; + target.generation += 1; + try std.testing.expect(base.compatible(target)); + target.page_rows = 32; + try std.testing.expect(!base.compatible(target)); + target = base; + target.pos0 = 1536; + try std.testing.expect(!base.compatible(target)); + target = base; + target.positions += 16; + try std.testing.expect(!base.compatible(target)); + target = base; + target.keys += 1024; + try std.testing.expect(!base.compatible(target)); + target = base; + target.functions.plan += 10; + try std.testing.expect(!base.compatible(target)); + target = base; + target.sub = 128; + try std.testing.expectError(error.UnsupportedPrefillBinding, target.validate()); + target = base; + target.pos0 = target.cap; + try std.testing.expectError(error.UnsupportedPrefillBinding, target.validate()); +} + +const Capture = struct { + var records: [10]binding.Record = undefined; + var args: cuda.Args = .{}; + var kind: c_uint = 0; + var count: usize = 10; + var destroyed: usize = 0; + var shared: bool = true; + + fn nodes(_: cuda.abi.Graph, output: ?[*]cuda.graph.Node, size: *usize) callconv(.c) cuda.abi.Result { + if (output) |storage| for (0..@min(size.*, count)) |index| { + storage[index] = @ptrFromInt(index + 1); + }; + size.* = count; + return 0; + } + + fn nodeType(_: cuda.graph.Node, output: *c_uint) callconv(.c) cuda.abi.Result { + output.* = kind; + return 0; + } + + fn params(node: cuda.graph.Node, output: *cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + const record = records[@intFromPtr(node) - 1]; + args = .{}; + for (binding.schema(record.kind, shared), 0..) |size, index| { + if (size == 8) args.add(record.values[index]) else args.add(@as(u32, @truncate(record.values[index]))); + } + output.* = record.params; + output.params = args.pointers(); + return 0; + } + + fn destroy(_: cuda.abi.Graph) callconv(.c) cuda.abi.Result { + destroyed += 1; + return 0; + } + + fn driver() cuda.Driver { + var result: cuda.Driver = undefined; + result.api.cuGraphGetNodes = nodes; + result.api.cuGraphKernelNodeGetParams_v2 = params; + result.api.cuGraphDestroy = destroy; + return result; + } +}; + +test "exl3: prefill binding validates captured schemas, missing nodes, duplicates and child graphs" { + var driver = Capture.driver(); + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(99) }; + for ([_]bool{ false, true }) |shared| { + const req = request(shared); + Capture.shared = shared; + Capture.records = recordsFor(req); + Capture.count = 10; + Capture.kind = 0; + const loaded = try binding.Binding.createWithNodeType(std.testing.allocator, graph, req, Capture.nodeType); + try std.testing.expectEqual(@as(usize, 10), loaded.records.len); + for (loaded.records, Capture.records) |actual, expected| try std.testing.expectEqualDeep(expected.values, actual.values); + loaded.deinit(std.testing.allocator); + Capture.count = 9; + try std.testing.expectError(error.UnsupportedPrefillBinding, binding.Binding.createWithNodeType(std.testing.allocator, graph, req, Capture.nodeType)); + Capture.count = 10; + Capture.records[9] = Capture.records[0]; + try std.testing.expectError(error.UnsupportedPrefillBinding, binding.Binding.createWithNodeType(std.testing.allocator, graph, req, Capture.nodeType)); + Capture.records = recordsFor(req); + Capture.records[0].values[7] += 1; + try std.testing.expectError(error.UnsupportedPrefillBinding, binding.Binding.createWithNodeType(std.testing.allocator, graph, req, Capture.nodeType)); + Capture.kind = 4; + try std.testing.expectError(error.UnsupportedPrefillBinding, binding.Binding.createWithNodeType(std.testing.allocator, graph, req, Capture.nodeType)); + try std.testing.expectError(error.UnsupportedPrefillBinding, binding.Binding.createWithNodeType(std.testing.allocator, graph, req, null)); + } +} + +const Setter = struct { + calls: usize = 0, + fail_at: ?usize = null, + records: [10]binding.Record = undefined, + + pub fn set(self: *@This(), record: binding.Record, _: bool) !void { + if (self.fail_at == self.calls) return error.InjectedSetterFailure; + self.calls += 1; + self.records[@intFromPtr(record.node) - 1] = record; + } +}; + +test "exl3: prefill rebind updates all origin arguments and grids without stale scalars" { + for ([_]bool{ false, true }) |shared| { + const req = request(shared); + var records = recordsFor(req); + var state: binding.Binding = .{ .graph = undefined, .request = req, .records = &records }; + var setter: Setter = .{}; + var target = req; + target.pos0 = 8192; + target.origin = 4096; + try state.apply(&setter, target); + try std.testing.expectEqual(@as(usize, 8), setter.calls); + for (records) |record| { + const before = recordFor(req, record.kind, record.block); + const expected = before.rebound(target); + try std.testing.expectEqualDeep(expected.values, record.values); + try std.testing.expectEqual(expected.params.grid_y, record.params.grid_y); + try std.testing.expectEqual(expected.params.grid_z, record.params.grid_z); + if (record.kind == .plan and shared) try std.testing.expectEqual(target.generation, record.values[22]); + } + setter.calls = 0; + try state.apply(&setter, target); + try std.testing.expectEqual(@as(usize, 0), setter.calls); + try state.apply(&setter, req); + try std.testing.expectEqual(@as(usize, 8), setter.calls); + for (records) |record| try std.testing.expectEqualDeep(recordFor(req, record.kind, record.block).values, record.values); + } +} + +test "exl3: prefill partial setter failures poison replay and topology declines do not mutate" { + const req = request(true); + var records = recordsFor(req); + var state: binding.Binding = .{ .graph = undefined, .request = req, .records = &records }; + var setter: Setter = .{ .fail_at = 1 }; + var target = req; + target.pos0 = 8192; + target.origin = 4096; + try std.testing.expectError(error.InjectedSetterFailure, state.apply(&setter, target)); + try std.testing.expect(state.poisoned); + try std.testing.expectEqual(req, state.request); + try std.testing.expectError(error.UnsupportedPrefillBinding, state.apply(&setter, req)); + state.poisoned = false; + target.pos0 = 1536; + setter.calls = 0; + try std.testing.expectError(error.UnsupportedPrefillBinding, state.apply(&setter, target)); + try std.testing.expectEqual(@as(usize, 0), setter.calls); +} + +test "exl3: prefill CUDA setter preserves scalar widths and 64-bit generations" { + const Fake = struct { + var expected: binding.Record = undefined; + var shared: bool = false; + var matches: bool = false; + + fn set(_: cuda.abi.GraphExec, node: cuda.graph.Node, params: *const cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + matches = node == expected.node and params.func == expected.params.func and params.grid_y == expected.params.grid_y and params.grid_z == expected.params.grid_z and params.kern == null and params.ctx == null and params.extra == null; + for (binding.schema(expected.kind, shared), 0..) |size, index| { + const value = @as([*]const u8, @ptrCast(params.params.?[index].?))[0..size]; + matches = matches and std.mem.eql(u8, value, std.mem.asBytes(&expected.values[index])[0..size]); + } + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphExecKernelNodeSetParams_v2 = Fake.set; + const backend: binding.Backend = .{ .exec = .{ .d = &driver, .handle = @ptrFromInt(99) } }; + for ([_]bool{ false, true }) |shared| for (std.enums.values(binding.Kind)) |kind| { + Fake.expected = recordFor(request(shared), kind, 0); + Fake.expected.params.kern = @ptrFromInt(123); + Fake.expected.params.ctx = @ptrFromInt(456); + Fake.shared = shared; + try backend.set(Fake.expected, shared); + try std.testing.expect(Fake.matches); + }; +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_history.zig b/zig/src/families/flashnext/cuda_exl3_prefill_history.zig new file mode 100644 index 0000000000..b762bdc296 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_history.zig @@ -0,0 +1,63 @@ +const std = @import("std"); + +pub const Slot = struct { + rows: usize = 0, + used: usize = 0, +}; + +pub const Shape = struct { sub: usize, lo: usize, hi: usize }; + +pub fn prepare(history: []Slot, previous: *?Shape, shape: Shape) void { + if (previous.* == null or !std.meta.eql(previous.*.?, shape)) @memset(history, .{}); + previous.* = shape; +} + +pub fn want(history: Slot, rows: usize, current_used: usize, cap: usize) usize { + const used = if (rows != 0 and history.rows == rows) history.used else current_used; + return @min(cap, used +| 16); +} + +test "exl3: prefetch history unseen and different rows retain the current-layer fallback" { + try std.testing.expectEqual(@as(usize, 416), want(.{}, 4096, 400, 512)); + try std.testing.expectEqual(@as(usize, 416), want(.{ .rows = 2048, .used = 50 }, 4096, 400, 512)); + try std.testing.expectEqual(@as(usize, 416), want(.{}, 0, 400, 512)); +} + +test "exl3: prefetch history uses the target layer's own unique demand including valid zero" { + var history: [4]Slot = @splat(.{}); + history[2] = .{ .rows = 4096, .used = 80 }; + history[0] = .{ .rows = 4096, .used = 300 }; + try std.testing.expectEqual(@as(usize, 316), want(history[1], 4096, 300, 512)); + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 300, 512)); + history[1] = .{ .rows = 4096, .used = 200 }; + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 200, 512)); + history[2] = .{ .rows = 4096, .used = 0 }; + try std.testing.expectEqual(@as(usize, 16), want(history[2], 4096, 200, 512)); + try std.testing.expectEqual(@as(usize, 216), want(history[2], 2048, 200, 512)); +} + +test "exl3: prefetch history caps both paths without overflow" { + for ([_]usize{ 0, 1, 16, 256, 512, std.math.maxInt(usize) }) |cap| { + try std.testing.expectEqual(cap, want(.{}, 4096, std.math.maxInt(usize), cap)); + try std.testing.expectEqual(cap, want(.{ .rows = 4096, .used = std.math.maxInt(usize) }, 4096, 0, cap)); + } + for (0..513) |used| { + try std.testing.expectEqual(@min(@as(usize, 512), used + 16), want(.{}, 4096, used, 512)); + } +} + +test "exl3: prefetch history resets for sub-block or rank shape changes" { + const shape: Shape = .{ .sub = 256, .lo = 0, .hi = 256 }; + var previous: ?Shape = null; + var history: [4]Slot = @splat(.{ .rows = 4096, .used = 80 }); + prepare(&history, &previous, shape); + for (history) |entry| try std.testing.expectEqual(@as(usize, 416), want(entry, 4096, 400, 512)); + history[2] = .{ .rows = 4096, .used = 80 }; + prepare(&history, &previous, shape); + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 400, 512)); + prepare(&history, &previous, .{ .sub = 512, .lo = 0, .hi = 256 }); + try std.testing.expectEqual(@as(usize, 416), want(history[2], 4096, 400, 512)); + history[2] = .{ .rows = 4096, .used = 80 }; + prepare(&history, &previous, .{ .sub = 512, .lo = 256, .hi = 512 }); + try std.testing.expectEqual(@as(usize, 416), want(history[2], 4096, 400, 512)); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_k.zig b/zig/src/families/flashnext/cuda_exl3_prefill_k.zig new file mode 100644 index 0000000000..c69379fabd --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_k.zig @@ -0,0 +1,227 @@ +//! The Flash Next EXL3 prefill experts (kernels/cuda/flashnext_exl3_prefill.cu): 64-member transposed prompt tiles on +//! upstream's pointer tables for any per-expert width, over a host-built member list (an expert-major pass). +//! +//! A pass is a set of experts and every (row, slot) pair of a chunk that picked one of them. The host groups the pairs +//! by expert (ascending ids, pairs in order) into `members` and cuts each expert's run into items of at most 64: +//! (expert, first, count). Member c owns the compact rows c of the rotated inputs and the SwiGLU outputs; the down +//! writes pair p's fp32 output to y[p] as Experts.runSlots does. +const std = @import("std"); +const cuda = @import("cuda"); + +const exl3e = cuda.exl3_experts; + +pub const tile: usize = 64; + +pub const Kernels = struct { + mod: cuda.Module, + rotate: cuda.Function, + gu: cuda.Function, + down: cuda.Function, + /// Expert parallel: a rank's combined branch over its own slots, and the ranks' sum placed for the write-back. + tp_combine: cuda.Function, + tp_place: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_prefill); + errdefer mod.unload(); + return .{ .mod = mod, .rotate = try mod.function("tf_exl3pf_rotate"), .gu = try mod.function("tf_exl3pf_gu"), .down = try mod.function("tf_exl3pf_down"), .tp_combine = try mod.function("tf_exl3pf_tp_combine"), .tp_place = try mod.function("tf_exl3pf_tp_place") }; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } +}; + +/// Members and items of one pass, host side (int32, the layouts the kernels read). +pub const Plan = struct { + members: []i32, + items: []i32, // 3 a item + n_members: usize = 0, + n_items: usize = 0, + + /// Room for `pairs` members and their items over `experts` table entries. + pub fn init(gpa: std.mem.Allocator, pairs: usize, experts: usize) !Plan { + const members = try gpa.alloc(i32, pairs); + errdefer gpa.free(members); + return .{ .members = members, .items = try gpa.alloc(i32, 3 * maxItems(pairs, experts)) }; + } + + pub fn deinit(p: *Plan, gpa: std.mem.Allocator) void { + gpa.free(p.members); + gpa.free(p.items); + p.* = undefined; + } + + pub fn itemBytes(p: *const Plan) []const u8 { + return std.mem.sliceAsBytes(p.items[0 .. 3 * p.n_items]); + } + + pub fn memberBytes(p: *const Plan) []const u8 { + return std.mem.sliceAsBytes(p.members[0..p.n_members]); + } +}; + +pub fn maxItems(pairs: usize, experts: usize) usize { + return @min(pairs, experts) + pairs / tile; +} + +/// Counting sort of `picks` (int32 [pairs], -1 or ids outside the table skipped) by expert, keeping only experts with +/// `take[id]` (null: all): members in (expert ascending, pair ascending) order, items of at most `tile` members. +/// `counts` is caller scratch of the table length. +pub fn build(picks: []const i32, take: ?[]const bool, counts: []u32, p: *Plan) void { + @memset(counts, 0); + const e_n = counts.len; + for (picks) |pk| { + if (pk < 0 or @as(usize, @intCast(pk)) >= e_n) continue; + const id: usize = @intCast(pk); + if (take) |t| if (!t[id]) continue; + counts[id] += 1; + } + // exclusive prefix into counts (the next free member slot of each expert), items as runs of `tile` + var at: u32 = 0; + var items: usize = 0; + for (counts, 0..) |*c, id| { + const n = c.*; + c.* = at; + var off: u32 = 0; + while (off < n) : (off += tile) { + p.items[3 * items] = @intCast(id); + p.items[3 * items + 1] = @intCast(at + off); + p.items[3 * items + 2] = @intCast(@min(@as(u32, tile), n - off)); + items += 1; + } + at += n; + } + for (picks, 0..) |pk, i| { + if (pk < 0 or @as(usize, @intCast(pk)) >= e_n) continue; + const id: usize = @intCast(pk); + if (take) |t| if (!t[id]) continue; + p.members[counts[id]] = @intCast(i); + counts[id] += 1; + } + p.n_members = at; + p.n_items = items; +} + +/// Device scratch a pass of `members` members needs: xg, xu [members, dims] and xd [members, width] fp16, 256-aligned. +pub fn scratchBytes(members: usize, dims: usize, width: usize) usize { + return 2 * align256(members * dims * 2) + align256(members * width * 2); +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +pub const Device = struct { + /// Device copies of the plan (members, items) and the compact scratch base (scratchBytes of the plan's members). + members: u64, + items: u64, + scratch: u64, +}; + +/// One pass: rotate every member's input, gate/up + SwiGLU, down into y [pairs, dims] fp32 (y[p] for member p). +/// `x` bf16 rows (stride in elements), `picks` the pairs' expert ids on the device. +pub fn run(k: *const Kernels, s: cuda.Stream, l: exl3e.Layer, x: u64, stride: usize, picks: u64, slots: usize, n_members: usize, n_items: usize, dv: Device, y: u64) !void { + if (n_members == 0) return; + const d = l.dims; + const w = l.width; + if (d % 256 != 0 or w % 128 != 0 or (l.act_mode != 1 and l.act_mode != 0)) return error.Exl3Shape; + const xg = dv.scratch; + const xu = xg + align256(n_members * d * 2); + const xd = xu + align256(n_members * d * 2); + var r: cuda.Args = .{}; + r.add(x); + r.add(@as(i32, @intCast(stride))); + r.add(picks); + r.add(dv.members); + r.add(@as(i32, @intCast(slots))); + r.add(l.suh_g); + r.add(l.suh_u); + r.add(xg); + r.add(xu); + r.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.rotate, .{ .grid = .{ .x = @intCast(n_members), .y = @intCast(d / 128), .z = 2 }, .block = .{ .x = 32 } }, s, &r); + var g: cuda.Args = .{}; + g.add(xg); + g.add(xu); + g.add(l.gate); + g.add(l.up); + g.add(l.gate_k2); + g.add(dv.items); + g.add(l.svh_g); + g.add(l.svh_u); + g.add(l.suh_d); + g.add(xd); + g.add(@as(i32, @intCast(d))); + g.add(@as(i32, @intCast(w))); + g.add(l.limit); + g.add(l.act_mode); + try cuda.launch.launch(k.gu, .{ .grid = .{ .x = @intCast(w / 128), .y = @intCast(n_items) }, .block = .{ .x = 256 } }, s, &g); + var dn: cuda.Args = .{}; + dn.add(xd); + dn.add(l.down); + dn.add(l.down_k2); + dn.add(dv.items); + dn.add(dv.members); + dn.add(l.svh_d); + dn.add(y); + dn.add(@as(i32, @intCast(w))); + dn.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.down, .{ .grid = .{ .x = @intCast(d / 256), .y = @intCast(n_items) }, .block = .{ .x = 256 } }, s, &dn); +} + +/// Expert parallel: sum[r] = the fp32 fmaf chain over row r's slots whose pick is in [lo, hi), in slot order. +pub fn tpCombine(k: *const Kernels, s: cuda.Stream, y: u64, wts: u64, picks: u64, sum: u64, rows: usize, slots: usize, d: usize, lo: usize, hi: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(y); + a.add(wts); + a.add(picks); + a.add(sum); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intCast(d))); + a.add(@as(i32, @intCast(lo))); + a.add(@as(i32, @intCast(hi))); + try cuda.launch.launch(k.tp_combine, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &a); +} + +/// Expert parallel: y[r, 0] = sum[r] and weights {1, 0, ...} (the write-back then writes the sum's bf16). +pub fn tpPlace(k: *const Kernels, s: cuda.Stream, sum: u64, y: u64, wts: u64, rows: usize, slots: usize, d: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(sum); + a.add(y); + a.add(wts); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.tp_place, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &a); +} + +test "exl3: prefill plan groups pairs by expert into items of at most 64" { + const gpa = std.testing.allocator; + // 3 rows of 3 slots: experts 2, 5, -1 / 5, 2, 7 / 2, 2, 513 (out of a 9-entry table: skipped) + const picks = [_]i32{ 2, 5, -1, 5, 2, 7, 2, 2, 513 }; + var counts: [9]u32 = undefined; + var p = try Plan.init(gpa, picks.len, 9); + defer p.deinit(gpa); + build(&picks, null, &counts, &p); + try std.testing.expectEqual(@as(usize, 7), p.n_members); + try std.testing.expectEqualSlices(i32, &.{ 0, 4, 6, 7, 1, 3, 5 }, p.members[0..7]); + try std.testing.expectEqual(@as(usize, 3), p.n_items); + try std.testing.expectEqualSlices(i32, &.{ 2, 0, 4, 5, 4, 2, 7, 6, 1 }, p.items[0..9]); + // a pass of expert 5 alone + var take: [9]bool = @splat(false); + take[5] = true; + build(&picks, &take, &counts, &p); + try std.testing.expectEqualSlices(i32, &.{ 1, 3 }, p.members[0..p.n_members]); + try std.testing.expectEqualSlices(i32, &.{ 5, 0, 2 }, p.items[0 .. 3 * p.n_items]); + // 150 pairs of one expert: items of 64, 64, 22 + var many: [150]i32 = @splat(1); + var p2 = try Plan.init(gpa, many.len, 9); + defer p2.deinit(gpa); + many[0] = -1; + build(&many, null, &counts, &p2); + try std.testing.expectEqual(@as(usize, 149), p2.n_members); + try std.testing.expectEqualSlices(i32, &.{ 1, 0, 64, 1, 64, 64, 1, 128, 21 }, p2.items[0 .. 3 * p2.n_items]); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_ring.zig b/zig/src/families/flashnext/cuda_exl3_prefill_ring.zig new file mode 100644 index 0000000000..a2f8db9de1 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_ring.zig @@ -0,0 +1,580 @@ +//! The prefill line ring: during an expert-major prefill every decoder layer's pool lines (cuda_exl3_pool.zig) form ONE +//! ring of expert slots, so the copy engine streams the experts of the coming layers while the current layer's blocks +//! compute. A layer's pool uses its own lines only while that layer runs; a prefill walks the layers in order, so the +//! other layers' lines are idle memory for it. Each layer's pointer tables (the pool's, fixed addresses) point into the +//! ring while its experts compute; at the end every pool is left empty (policy state cleared, routed table entries +//! zeroed) and refills itself on the next rounds. +//! +//! Ordering, on the device: +//! copy stream: [wait the release of the pass that last read the slot] -> H2D of the expert -> after a batch of +//! copies, one H2D of the layer's pointer table from its pinned shadow -> batch event +//! compute stream: [wait the batch event of the pass's last expert] -> prefill experts kernel -> release event +//! A slot that holds an expert no pass has read yet (`pending`) is never overwritten. + +const std = @import("std"); +const cuda = @import("cuda"); +const pool_ = @import("cuda_exl3_pool.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const attribution_ = @import("cuda_exl3_prefill_attribution.zig"); + +const layers_n = wt.layers_n; +const routed_n = wt.experts_n; +const table_n = routed_n + 1; +const batch_ring = 1024; +const release_ring = 4096; +/// Copies a batch holds before its table upload and event. +const batch_copies = 16; +/// Pinned snapshots of a layer's pointer tables, one a batch in flight (a batch uploads its snapshot, never the live +/// shadow, which the host keeps rewriting). +const snap_ring = 64; + +const SlotState = enum(u8) { free, pending }; + +pub const Stats = struct { + copies: u64 = 0, + dma_requested_bytes: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + hit_expert_requests: u64 = 0, + miss_expert_requests: u64 = 0, + wasted: u64 = 0, + ring_full: u64 = 0, + release_waits: u64 = 0, + compute_waits: u64 = 0, + reseeded: u64 = 0, + + pub fn observeDemand(stats: *Stats, present: bool, picks: u32) void { + if (picks == 0) return; + if (present) { + stats.hit_picks += picks; + stats.hit_expert_requests += 1; + } else { + stats.miss_picks += picks; + stats.miss_expert_requests += 1; + } + } +}; + +pub const Ring = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + pools: [layers_n]*pool_.Pool, + addr: []u64, + /// Each slot's projection offsets and room (its pool's line layout: under expert parallel a rank's layers hold + /// different widest experts, so their lines differ; an expert only goes into a slot it fits). + off: [][3]u32, + room_of: [][3]u32, + state: []SlotState, + owner: []u32, // (layer << 16 | id) + 1, 0: none + epoch_of: []u64, + next: usize = 0, + pending_n: usize = 0, + copy: cuda.Stream, + cev: [batch_ring]cuda.Event, + batches: u64 = 0, + rev: [release_ring]cuda.Event, + epoch: u64 = 0, + copy_ordered_epoch: u64 = 0, + tlay: cuda.exl3_pool.Tables, + /// Pinned images of each layer's tables (pointers and widths) as the device holds them during the prefill. + shadow: cuda.HostBuffer, + snaps: cuda.HostBuffer, + /// Slot of each layer's expert in the current chunk, or -1; the batch that copied it. + slot: [layers_n][routed_n]i32 = undefined, + batch_of: [layers_n][routed_n]u64 = undefined, + /// Prefill member counts a layer has seen (the guess order for the next chunks). + counts: [layers_n][routed_n]u32 = undefined, + /// The open batch: its layer and copies. + open_layer: ?usize = null, + open_n: usize = 0, + /// The routed ids this rank holds (expert parallel, cuda_exl3_tp.zig): a prefetch never guesses another rank's. + own_lo: usize = 0, + own_hi: usize = routed_n, + stats: Stats = .{}, + telemetry_id: u64 = 0, + attribution: ?attribution_.State = null, + + /// A ring over `pools` (one per decoder layer, every line the same layout); error.Invalid otherwise. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, pools: []const *pool_.Pool, telemetry: bool) !*Ring { + if (pools.len != layers_n) return error.Invalid; + // the tables' entries: 513 (routed and shared), or 512 on an expert-parallel rank without the shared expert + const te = pools[0].lines.spec.experts; + if (te != table_n and te != routed_n) return error.Invalid; + var n: usize = 0; + for (pools) |p| { + if (p.lines.spec.experts != te) return error.Invalid; + n += p.opt.lines; + } + const r = try gpa.create(Ring); + errdefer gpa.destroy(r); + r.* = .{ .gpa = gpa, .d = d, .pools = undefined, .off = undefined, .room_of = undefined, .addr = undefined, .state = undefined, .owner = undefined, .epoch_of = undefined, .copy = undefined, .cev = undefined, .rev = undefined, .tlay = cuda.exl3_pool.Tables.init(te), .shadow = undefined, .snaps = undefined }; + r.telemetry_id = try @import("exl3_telemetry.zig").identity(); + @memcpy(&r.pools, pools); + r.addr = try gpa.alloc(u64, n); + errdefer gpa.free(r.addr); + r.off = try gpa.alloc([3]u32, n); + errdefer gpa.free(r.off); + r.room_of = try gpa.alloc([3]u32, n); + errdefer gpa.free(r.room_of); + r.state = try gpa.alloc(SlotState, n); + errdefer gpa.free(r.state); + r.owner = try gpa.alloc(u32, n); + errdefer gpa.free(r.owner); + r.epoch_of = try gpa.alloc(u64, n); + errdefer gpa.free(r.epoch_of); + r.attribution = try attribution_.State.init(gpa, n, telemetry); + errdefer if (r.attribution) |*state| state.deinit(gpa); + var k: usize = 0; + for (pools) |p| for (0..p.opt.lines) |i| { + const lay = p.lines.layout; + r.addr[k] = p.lines.at(i, 0); + for (0..3) |w| { + const stop = if (w < 2) lay.off[w + 1] else lay.line_bytes; + r.off[k][w] = @intCast(lay.off[w]); + r.room_of[k][w] = @intCast(stop - lay.off[w]); + } + k += 1; + }; + r.copy = try cuda.Stream.init(d, true); + errdefer r.copy.deinit(); + var made_c: usize = 0; + errdefer for (r.cev[0..made_c]) |*e| e.deinit(); + for (&r.cev) |*e| { + e.* = try cuda.Event.init(d, false); + made_c += 1; + } + var made_r: usize = 0; + errdefer for (r.rev[0..made_r]) |*e| e.deinit(); + for (&r.rev) |*e| { + e.* = try cuda.Event.init(d, false); + made_r += 1; + } + r.shadow = try cuda.HostBuffer.alloc(d, layers_n * r.tlay.bytes); + errdefer r.shadow.free(); + r.snaps = try cuda.HostBuffer.alloc(d, snap_ring * r.tlay.k2[0]); + for (&r.counts) |*c| @memset(c, 0); + return r; + } + + pub fn deinit(r: *Ring) void { + r.copy.synchronize() catch {}; + r.snaps.free(); + r.shadow.free(); + for (&r.rev) |*e| e.deinit(); + for (&r.cev) |*e| e.deinit(); + r.copy.deinit(); + if (r.attribution) |*state| state.deinit(r.gpa); + r.gpa.free(r.epoch_of); + r.gpa.free(r.room_of); + r.gpa.free(r.off); + r.gpa.free(r.owner); + r.gpa.free(r.state); + r.gpa.free(r.addr); + r.gpa.destroy(r); + } + + pub fn slots(r: *const Ring) usize { + return r.addr.len; + } + + fn shadowOf(r: *Ring, layer: usize) []u8 { + return r.shadow.bytes[layer * r.tlay.bytes ..][0..r.tlay.bytes]; + } + + fn ptrEntry(r: *Ring, layer: usize, which: usize, id: usize) *align(1) u64 { + return @ptrCast(r.shadowOf(layer)[r.tlay.ptr[which] + id * 8 ..][0..8]); + } + + fn k2Entry(r: *Ring, layer: usize, which: usize, id: usize) *align(1) i32 { + return @ptrCast(r.shadowOf(layer)[r.tlay.k2[which] + id * 4 ..][0..4]); + } + + /// Before a prefill: nothing of the pools' may run (`s` drained), every slot free, each layer's table image read + /// (the shared expert's entries stay as the pool wrote them) with every routed entry absent, then the routed widths. + pub fn begin(r: *Ring, s: cuda.Stream) !void { + try s.synchronize(); + for (r.pools) |p| try p.lines.drain(); + @memset(r.state, .free); + @memset(r.owner, 0); + @memset(r.epoch_of, 0); + r.next = 0; + r.pending_n = 0; + r.copy_ordered_epoch = r.epoch; + r.open_layer = null; + r.open_n = 0; + for (&r.slot) |*row| @memset(row, -1); + for (r.pools, 0..) |p, l| { + const img = r.shadowOf(l); + try p.lines.tables.download(0, img); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + r.k2Entry(l, w, id).* = p.k2[id][w]; + }; + try upload(r.copy, p.lines.tables.ptr, img); + } + // the images above are read before the host rewrites them + try r.copy.synchronize(); + } + + /// A new deep chunk: every layer's slots of the last one are stale (nothing reads them any more). + pub fn newChunk(r: *Ring) !void { + try r.closeBatch(); + for (r.state, r.owner, 0..) |*st, *ow, index| { + if (st.* == .pending) { + r.stats.wasted += 1; + if (r.attribution) |*state| state.unused(index); + } + st.* = .free; + ow.* = 0; + } + r.pending_n = 0; + for (&r.slot, 0..) |*row, l| { + @memset(row, -1); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + }; + } + } + + /// After a prefill: every pool empty (policy state, routed table entries), the device idle; then, with `reseed`, + /// each pool refills its lines with its layer's hottest experts of the prefills so far (copies queued on the copy + /// engine, nothing waits for them but the first rounds that read those lines). + pub fn end(r: *Ring, s: cuda.Stream, reseed: bool) !void { + try r.closeBatch(); + try r.copy.synchronize(); + try s.synchronize(); + for (r.pools, 0..) |p, l| { + const img = r.shadowOf(l); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + r.k2Entry(l, w, id).* = 0; + }; + try p.lines.tables.upload(0, img); + p.v = pool_.offload.Vram.init(p.slots_tab, p.line_tab, p.stamp_tab, p.layer_tab); + p.prev_n = 0; + @memset(p.lines.line_pending, false); + @memset(p.lines.line_epoch, 0); + } + if (!reseed) return; + for (r.pools, 0..) |p, l| { + var order: [routed_n]u16 = undefined; + var n: usize = 0; + for (0..routed_n) |id| { + if (r.counts[l][id] == 0) continue; + order[n] = @intCast(id); + n += 1; + } + const Ctx = struct { + c: *const [routed_n]u32, + fn hotter(c: @This(), a: u16, b: u16) bool { + if (c.c[a] != c.c[b]) return c.c[a] > c.c[b]; + return a < b; + } + }; + std.mem.sort(u16, order[0..n], Ctx{ .c = &r.counts[l] }, Ctx.hotter); + const take = @min(n, p.opt.lines, p.opt.max_rows * p.opt.slots); + if (take == 0) continue; + const rows = (take + p.opt.slots - 1) / p.opt.slots; + var fake: [routed_n + 16]i32 = undefined; + const picks = fake[0 .. rows * p.opt.slots]; + @memset(picks, -1); + for (order[0..take], 0..) |id, i| picks[i] = id; + const rd = try p.plan(picks, rows); + // the pool's copy stream does the copies; the ring's stream takes the gates and the release + for (0..rd.n_passes) |j| { + try p.stage(rd, j, r.copy); + try p.release(rd, j, r.copy); + } + r.stats.reseeded += rd.misses_(); + } + } + + fn closeBatch(r: *Ring) !void { + const l = r.open_layer orelse return; + const p = r.pools[l]; + // the pointer tables only (widths were written at begin), from this batch's own snapshot + const n = r.tlay.k2[0]; + const b = r.batches; + if (b >= snap_ring) try r.cev[(b - snap_ring) % batch_ring].synchronize(); + const snap = r.snaps.bytes[(b % snap_ring) * n ..][0..n]; + @memcpy(snap, r.shadowOf(l)[0..n]); + try upload(r.copy, p.lines.tables.ptr, snap); + try r.cev[r.batches % batch_ring].record(r.copy); + r.batches += 1; + r.open_layer = null; + r.open_n = 0; + } + + /// The copy stream waits until the pass that last read slot `k` finished. + fn afterReaders(r: *Ring, k: usize) !void { + const e = r.epoch_of[k]; + if (e == 0 or e <= r.copy_ordered_epoch) return; + const ev = r.rev[e % release_ring]; + if (r.epoch - e >= release_ring) { + // its event was recorded over by a newer pass: wait for that one on the host (no device cycle) + try ev.synchronize(); + r.copy_ordered_epoch = e; + return; + } + if (try ev.done()) { + r.copy_ordered_epoch = e; + return; + } + try r.copy.wait(ev); + r.copy_ordered_epoch = e; + r.stats.release_waits += 1; + } + + pub fn isLoaded(r: *const Ring, layer: usize, id: usize) bool { + return r.slot[layer][id] >= 0; + } + + /// Free slots ahead of the cursor before a pending one (the ring's room). + pub fn room(r: *const Ring) usize { + return r.slots() - r.pending_n; + } + + /// Copies expert `id` of `layer` into the next free slot; false when the ring holds no free slot. + pub fn load(r: *Ring, layer: usize, id: usize, source: attribution_.Source) !bool { + if (r.slot[layer][id] >= 0) return true; + if (r.pending_n == r.slots()) { + r.stats.ring_full += 1; + return false; + } + const p = r.pools[layer]; + const src = p.src.routed[id]; + // the oldest slots are consumed first; skip any still pending, or too small for this expert + var tried: usize = 0; + while (r.state[r.next] == .pending or !fits(r.room_of[r.next], src)) { + r.next = (r.next + 1) % r.slots(); + tried += 1; + if (tried == r.slots()) { + r.stats.ring_full += 1; + return false; + } + } + const k = r.next; + r.next = (r.next + 1) % r.slots(); + // the open batch's snapshot first (it must not see the evicted expert gone), then the wait for the slot's + // last reader, then the evicted expert's entries (later snapshots are uploaded behind that wait) + if (r.open_layer != null and (r.open_layer.? != layer or r.open_n == batch_copies)) try r.closeBatch(); + try r.afterReaders(k); + if (r.owner[k] != 0) { + const o = r.owner[k] - 1; + const ol = o >> 16; + const oid = o & 0xffff; + r.slot[ol][oid] = -1; + for (0..3) |w| r.ptrEntry(ol, w, oid).* = 0; + } + const dst = r.addr[k]; + const offsets: [3]usize = .{ r.off[k][0], r.off[k][1], r.off[k][2] }; + const capacity = try std.math.add(usize, offsets[2], r.room_of[k][2]); + _ = try cuda.exl3_pool.dma.directLines(r.copy, dst, capacity, offsets, src, p.src.pinned_bank); + for (0..3) |w| r.ptrEntry(layer, w, id).* = dst + r.off[k][w]; + r.slot[layer][id] = @intCast(k); + r.batch_of[layer][id] = r.batches; + r.owner[k] = @intCast(((layer << 16) | id) + 1); + r.state[k] = .pending; + r.pending_n += 1; + r.open_layer = layer; + r.open_n += 1; + r.stats.copies += 1; + r.stats.dma_requested_bytes += src[0].len + src[1].len + src[2].len; + if (r.attribution) |*state| state.copied(k, source, src[0].len + src[1].len + src[2].len); + return true; + } + + /// The compute stream waits for every copy (and table write) of `ids` of `layer`. + pub fn gate(r: *Ring, s: cuda.Stream, layer: usize, ids: []const u16) !void { + var last: ?u64 = null; + for (ids) |id| { + if (id >= routed_n) continue; + const b = r.batch_of[layer][id]; + last = if (last) |x| @max(x, b) else b; + } + const b = last orelse return; + if (b == r.batches) try r.closeBatch(); // its batch is still open + const ev = r.cev[b % batch_ring]; + if (r.batches - b > batch_ring) { + try r.copy.synchronize(); + return; + } + try s.wait(ev); + r.stats.compute_waits += 1; + } + + /// After a pass that read `ids` of `layer` was enqueued on `s`: its slots become free once it finishes. + pub fn release(r: *Ring, s: cuda.Stream, layer: usize, ids: []const u16) !void { + r.epoch += 1; + try r.rev[r.epoch % release_ring].record(s); + for (ids) |id| { + if (id >= routed_n) continue; + const k: usize = @intCast(r.slot[layer][id]); + r.epoch_of[k] = r.epoch; + if (r.state[k] == .pending) { + r.state[k] = .free; + r.pending_n -= 1; + } + } + } + + /// Experts loaded for `layer` that its picks did not use: their slots are free again (never read). + pub fn drop(r: *Ring, layer: usize, used: []const u32) void { + for (0..routed_n) |id| { + if (used[id] != 0) continue; + const k = r.slot[layer][id]; + if (k < 0) continue; + const ku: usize = @intCast(k); + if (r.state[ku] == .pending) { + r.state[ku] = .free; + r.pending_n -= 1; + r.stats.wasted += 1; + if (r.attribution) |*state| state.unused(ku); + } + } + } + + /// Tops `layer` up to `want` loaded experts with the ones not loaded, hottest (prefill counts, then the pool's lifetime + /// counts) first, leaving `keep_free` slots; returns how many were queued. + pub fn prefetch(r: *Ring, layer: usize, want: usize, keep_free: usize, source: attribution_.Source) !usize { + var order: [routed_n]u16 = undefined; + var n: usize = 0; + var have: usize = 0; + for (0..routed_n) |id| { + if (r.slot[layer][id] >= 0) { + have += 1; + continue; + } + if (id < r.own_lo or id >= r.own_hi) continue; + order[n] = @intCast(id); + n += 1; + } + if (have >= want) return 0; + const Ctx = struct { + a: *const [routed_n]u32, + b: *const [routed_n]u32, + fn hotter(c: @This(), x: u16, y: u16) bool { + if (c.a[x] != c.a[y]) return c.a[x] > c.a[y]; + if (c.b[x] != c.b[y]) return c.b[x] > c.b[y]; + return x < y; + } + }; + const p = r.pools[layer]; + std.mem.sort(u16, order[0..n], Ctx{ .a = &r.counts[layer], .b = &p.counts[p.opt.layer] }, Ctx.hotter); + var queued: usize = 0; + for (order[0..n]) |id| { + if (have + queued == want or r.room() <= keep_free) break; + if (!try r.load(layer, id, source)) break; + queued += 1; + } + try r.closeBatch(); + return queued; + } + + pub fn flush(r: *Ring) !void { + try r.closeBatch(); + } +}; + +fn fits(room: [3]u32, src: [3][]const u8) bool { + for (room, src) |n, s| if (s.len > n) return false; + return true; +} + +fn upload(s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +test "exl3: a ring slot takes an expert only if each projection fits its line's room" { + const big: [300]u8 = @splat(0); + const room: [3]u32 = .{ 256, 512, 256 }; + try std.testing.expect(fits(room, .{ big[0..256], big[0..300], big[0..10] })); + try std.testing.expect(!fits(room, .{ big[0..300], big[0..10], big[0..10] })); + try std.testing.expect(fits(room, .{ &.{}, &.{}, &.{} })); +} + +const ReleaseFixture = @import("cuda_exl3_prefill_ring_fixture.zig").Fixture; + +test "exl3: prefill attribution contracts" { + _ = @import("cuda_exl3_prefill_attribution_test.zig"); +} + +test "exl3: prefill ring queues one release fence for a pass's pending slot copies" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs: [32]u64 = @splat(7); + ReleaseFixture.init(&driver, ring, &epochs, 7); + for (epochs, 0..) |_, index| { + try ring.afterReaders(index); + try upload(ring.copy, 42, &.{1}); + } + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait }, ReleaseFixture.calls[0..2]); + try std.testing.expectEqual(@as(usize, 34), ReleaseFixture.count); + for (ReleaseFixture.calls[2..ReleaseFixture.count]) |call| try std.testing.expectEqual(ReleaseFixture.Call.copy, call); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 1), ring.stats.release_waits); +} + +test "exl3: prefill ring queued fences cover older epochs but never a newer reader" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{ 4, 2, 4, 5, 0, 1, 6 }; + ReleaseFixture.init(&driver, ring, &epochs, 6); + for (0..6) |index| try ring.afterReaders(index); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait, .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(ring.rev[5].handle, ReleaseFixture.last_event); + ReleaseFixture.query_result = cuda.abi.success; + try ring.afterReaders(6); + try ring.afterReaders(3); + try std.testing.expectEqual(@as(usize, 5), ReleaseFixture.count); + try std.testing.expectEqual(@as(u64, 6), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 2), ring.stats.release_waits); +} + +test "exl3: prefill ring retains the host fallback at the release event wrap boundary" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{ 7, 7 }; + for ([_]u64{ release_ring - 1, release_ring, release_ring + 1 }) |distance| { + ReleaseFixture.init(&driver, ring, &epochs, 7 + distance); + try ring.afterReaders(0); + try ring.afterReaders(1); + if (distance < release_ring) { + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 1), ring.stats.release_waits); + } else { + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{.sync}, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 0), ring.stats.release_waits); + } + try std.testing.expectEqual(ring.rev[7].handle, ReleaseFixture.last_event); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + } +} + +test "exl3: prefill ring failed release dependencies cannot authorize subsequent copies" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{7}; + ReleaseFixture.init(&driver, ring, &epochs, 7); + ReleaseFixture.wait_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.afterReaders(0)); + try std.testing.expectEqual(@as(u64, 0), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 0), ring.stats.release_waits); + ReleaseFixture.wait_result = cuda.abi.success; + try ring.afterReaders(0); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait, .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + + ReleaseFixture.init(&driver, ring, &epochs, 7 + release_ring); + ReleaseFixture.sync_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.afterReaders(0)); + try std.testing.expectEqual(@as(u64, 0), ring.copy_ordered_epoch); + ReleaseFixture.sync_result = cuda.abi.success; + try ring.afterReaders(0); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .sync, .sync }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); +} diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_ring_fixture.zig b/zig/src/families/flashnext/cuda_exl3_prefill_ring_fixture.zig new file mode 100644 index 0000000000..62d11ad190 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_ring_fixture.zig @@ -0,0 +1,57 @@ +const cuda = @import("cuda"); + +pub const Fixture = struct { + pub const Call = enum { query, wait, sync, copy }; + pub var calls: [128]Call = undefined; + pub var count: usize = 0; + pub var query_result: cuda.abi.Result = cuda.abi.error_not_ready; + pub var wait_result: cuda.abi.Result = cuda.abi.success; + pub var sync_result: cuda.abi.Result = cuda.abi.success; + pub var last_event: cuda.abi.Event = null; + + fn note(call: Call, event: cuda.abi.Event) void { + calls[count] = call; + count += 1; + if (event != null) last_event = event; + } + + fn query(event: cuda.abi.Event) callconv(.c) cuda.abi.Result { + note(.query, event); + return query_result; + } + + fn wait(_: cuda.abi.Stream, event: cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result { + note(.wait, event); + return wait_result; + } + + fn sync(event: cuda.abi.Event) callconv(.c) cuda.abi.Result { + note(.sync, event); + return sync_result; + } + + fn copy(_: cuda.abi.DevicePtr, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + note(.copy, null); + return cuda.abi.success; + } + + pub fn init(driver: *cuda.Driver, ring: anytype, epochs: []u64, epoch: u64) void { + count = 0; + query_result = cuda.abi.error_not_ready; + wait_result = cuda.abi.success; + sync_result = cuda.abi.success; + last_event = null; + driver.* = undefined; + driver.api.cuEventQuery = query; + driver.api.cuStreamWaitEvent = wait; + driver.api.cuEventSynchronize = sync; + driver.api.cuMemcpyHtoDAsync_v2 = copy; + ring.* = undefined; + ring.copy = .{ .d = driver, .handle = @ptrFromInt(1) }; + ring.epoch_of = epochs; + ring.epoch = epoch; + ring.copy_ordered_epoch = 0; + ring.stats = .{}; + for (&ring.rev, 0..) |*event, index| event.* = .{ .d = driver, .handle = @ptrFromInt(index + 1) }; + } +}; diff --git a/zig/src/families/flashnext/cuda_exl3_prefill_writeback_test.zig b/zig/src/families/flashnext/cuda_exl3_prefill_writeback_test.zig new file mode 100644 index 0000000000..fd20a9320c --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_prefill_writeback_test.zig @@ -0,0 +1,131 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("cuda_batch.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const dims = wt.dims; +const slots = wt.slots; +const streams = wt.streams; + +const Call = struct { branch: u64, y: u64, weights: u64, mode: usize, topk: usize, slots: usize, rows: usize }; + +const Tri = struct { + call: *Call, + + pub fn hcWriteback(self: Tri, _: u64, _: u64, _: u64, branch: u64, _: []const u8, _: u64, y: u64, _: []const u8, weights: u64, _: []const u8, _: usize, rows: usize, _: usize, _: usize, mode: usize, topk: usize, count: usize, _: usize) !void { + self.call.* = .{ .branch = branch, .y = y, .weights = weights, .mode = mode, .topk = topk, .slots = count, .rows = rows }; + } +}; + +const Copies = struct { + s: *batch.Scratch, + calls: usize = 0, + last: struct { dst: u64, src: u64, bytes: usize } = undefined, + + pub fn copy(self: *Copies, dst: u64, src: u64, bytes: usize) !void { + self.calls += 1; + self.last = .{ .dst = dst, .src = src, .bytes = bytes }; + } +}; + +pub fn run(comptime Deep: type, comptime layout: anytype, comptime view: anytype, comptime pending: anytype, comptime compact: anytype, comptime select: anytype) !void { + var driver: cuda.Driver = undefined; + var dp: Deep = undefined; + dp.shared = null; + dp.world = 2; + dp.rows = 7; + dp.lay = layout(dp.rows, dp.world); + dp.dev = .{ .d = &driver, .ptr = 0x100000, .len = dp.lay.total }; + dp.reduced_branch = false; + dp.unit_weights = .{ .d = &driver, .ptr = 0x900000, .len = dp.rows * 4 }; + try select(&dp, true); + try std.testing.expect(dp.reduced_branch); + var base: batch.Scratch = undefined; + base.dummy = 0xdead; + var call: Call = undefined; + var pass: batch.Pass(Tri) = undefined; + pass.tri = .{ .call = &call }; + const sum = dp.dev.ptr + dp.lay.sum; + const sub: usize = 3; + for (0..3) |block| { + const at = block * sub; + const rows = @min(sub, dp.rows - at); + var scratch = view(&base, &dp, at); + pass.s = &scratch; + pass.rows = rows; + try std.testing.expectEqual(sum + at * dims * 4, scratch.moe_f32); + pending(&pass, &dp, null, at); + try std.testing.expect(!pass.pending_f32); + pending(&pass, &dp, 4, at); + try pass.writeback(0, rows); + try std.testing.expectEqual(Call{ .branch = base.dummy, .y = sum + at * dims * 4, .weights = dp.unit_weights.?.ptr + at * 4, .mode = 2, .topk = 0, .slots = 1, .rows = rows }, call); + try std.testing.expect(pass.pending == null and !pass.pending_f32 and pass.pending_f32_weights == null); + if (rows > 1) { + pending(&pass, &dp, 4, at); + try pass.writeback(1, rows - 1); + try std.testing.expectEqual(sum + (at + 1) * dims * 4, call.y); + try std.testing.expectEqual(dp.unit_weights.?.ptr + (at + 1) * 4, call.weights); + } + var copies: Copies = .{ .s = &scratch }; + try compact(&copies, true, rows - 1); + try std.testing.expectEqual(@as(usize, if (rows > 1) 4 else 0), copies.calls); + if (rows > 1) { + try std.testing.expectEqual(scratch.moe_f32, copies.last.dst); + try std.testing.expectEqual(scratch.moe_f32 + (rows - 1) * dims * 4, copies.last.src); + try std.testing.expectEqual(dims * 4, copies.last.bytes); + } + } + var scratch = view(&base, &dp, 0); + pass.s = &scratch; + pass.rows = sub; + try select(&dp, false); + pending(&pass, &dp, 4, 0); + try pass.writeback(0, sub); + try std.testing.expectEqual(@as(usize, slots), call.slots); + try std.testing.expectEqual(@as(usize, slots - 1), call.topk); + try std.testing.expectEqual(scratch.y, call.y); + pass.pending = 4; + pass.pending_f32 = true; + pass.pending_f32_weights = null; + try pass.writeback(0, sub); + try std.testing.expectEqual(@as(usize, 4), call.mode); + try std.testing.expectEqual(scratch.moe_f32, call.branch); + var copies: Copies = .{ .s = &scratch }; + try compact(&copies, false, 2); + try std.testing.expectEqual(@as(usize, 5), copies.calls); + try std.testing.expectEqual(scratch.wts, copies.last.dst); + try std.testing.expectEqual(scratch.wts + 2 * slots * 4, copies.last.src); + dp.world = 1; + base.moe_f32 = 0xbeef; + const single_rank = view(&base, &dp, sub); + try std.testing.expectEqual(base.moe_f32, single_rank.moe_f32); + try rounding(); +} + +fn bf16(value: f32) u16 { + const bits: u32 = @bitCast(value); + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} + +fn promote(value: u16) f32 { + return @bitCast(@as(u32, value) << 16); +} + +fn rounding() !void { + const values = [_]f32{ 0.0, -0.0, 0x1p-126, -0x1p-126, 0x1p-133, -0x1p-133, 0x1.01p0, -0x1.01p0, 0x1.03p0, -0x1.03p0, 0x1p40, -0x1p40 }; + const hidden = [_]u16{ 0, 0x8000, 0x3b34, 0xbe16, 0x3f80, 0xbf80 }; + const inject = [_]u16{ 0, 0x8000, 0x3c9a, 0x3f80, 0xbf80 }; + for (values) |sum| { + const direct = @mulAdd(f32, sum, 1, 0); + var old = direct; + for (1..slots) |slot| old = @mulAdd(f32, if (slot % 2 == 0) @as(f32, -7) else 13, 0, old); + for (hidden) |bits| for (inject) |gate| { + const want = bf16(@mulAdd(f32, promote(bf16(old)), promote(gate), promote(bits))); + const got = bf16(@mulAdd(f32, promote(bf16(direct)), promote(gate), promote(bits))); + try std.testing.expectEqual(want, got); + }; + } + const wrong = bf16(@mulAdd(f32, promote(bf16(-@as(f32, 0.0))), 1, -@as(f32, 0.0))); + const correct = bf16(@mulAdd(f32, @mulAdd(f32, -0.0, 1, 0), 1, -@as(f32, 0.0))); + try std.testing.expectEqual(@as(u16, 0x8000), wrong); + try std.testing.expectEqual(@as(u16, 0), correct); +} diff --git a/zig/src/families/flashnext/cuda_exl3_recycle.zig b/zig/src/families/flashnext/cuda_exl3_recycle.zig new file mode 100644 index 0000000000..3ef0f84540 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_recycle.zig @@ -0,0 +1,174 @@ +//! Stable buffers for captured decode graphs (cuda_exl3_graph.zig): a graph holds a sequence's and its draft window's +//! device addresses, and every request used to allocate new ones (its own capacity, its own buffers), so verify +//! windows rarely repeated a graph key and ran eagerly or paid a capture. With graphs on, a stream's capacity is +//! rounded up to a multiple of `bucket` tokens and a finished stream's sequence and draft buffers are kept (zeroed on +//! reuse exactly as a new one is) for the next stream of the same capacity, on both ranks. +//! The capacity only bounds the caches and strides the block scores; no kernel reads past a row's own positions +//! (fn_qsa_select uses the stride as a stride), so the bits do not change (tf-cuda-test flashnext-exl3-graph eq and +//! the serve gate's hash). TF_FLASHNEXT_EXL3_RECYCLE=0 turns it off. + +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const Draft = @import("cuda_mtp_draft.zig").Draft; + +pub const bucket: usize = 4096; +const pool_n = 4; + +pub fn enabled() bool { + if (!@import("cuda_exl3_graph.zig").enabled()) return false; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_RECYCLE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// A stream's capacity: `cap` rounded up to the bucket (at most `limit`) when recycling. +pub fn roundCap(cap: usize, limit: usize) usize { + if (!enabled()) return cap; + return @min(limit, (cap + bucket - 1) / bucket * bucket); +} + +var seqs: [pool_n]?prompt.Seq = @splat(null); +var drafts: [pool_n]?Draft = @splat(null); +var draining = false; +pub var stats: struct { seq_reused: u64 = 0, seq_kept: u64 = 0, seq_evicted: u64 = 0, shared_lane_reassigned: u64 = 0, draft_reused: u64 = 0, draft_kept: u64 = 0 } = .{}; + +fn zero(bufs: []cuda.DeviceBuffer, s: cuda.Stream) !void { + for (bufs) |b| try b.fill8(0, s.handle); +} + +/// Engine.newSeq's hook: a kept sequence of capacity `cap`, zeroed as newSeq zeroes a new one; null: allocate. +pub fn takeSeq(e: *prompt.Engine, cap: usize) !?prompt.Seq { + const reuse = enabled(); + for (&seqs) |*slot| if (slot.*) |*sequence| { + if (sequence.recycle_owner == e and (!reuse or !matches(sequence, e, cap))) { + discardSeq(slot, e.gpa); + stats.seq_evicted += 1; + } + }; + if (!reuse) return null; + for (&seqs) |*slot| if (slot.*) |*q| if (matches(q, e, cap)) { + var s = q.*; + slot.* = null; + errdefer { + s.tail.d.synchronize() catch @panic("recycled sequence reset did not drain"); + s.invalidated = true; + s.deinit(e.gpa); + } + const st = e.stream.*; + var previous_lane: ?u32 = null; + if (s.paged) |pg| { + if (pg.cfg.shared_lanes != 0) previous_lane = pg.shared_sequence.lane; + try pg.resetForReuse(st); + if (previous_lane) |lane| if (pg.shared_sequence.lane != lane) { + stats.shared_lane_reassigned += 1; + @import("cuda_exl3_graph.zig").releaseSeq(&s); + }; + } + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| try zero(bs, st); + try s.tail.fill8(0, st.handle); + @memset(s.logits, 0); + s.history = .{ e.eos, e.eos }; + s.fed = 0; + s.rank_id = 0; + s.media = null; + s.untaught = false; + s.invalidated = false; + stats.seq_reused += 1; + return s; + }; + return null; +} + +fn discardSeq(slot: *?prompt.Seq, allocator: std.mem.Allocator) void { + var sequence = slot.*.?; + slot.* = null; + sequence.recycle_owner = null; + sequence.deinit(allocator); +} + +fn canRetainSeq(s: *const prompt.Seq) bool { + if (s.invalidated or s.media != null or s.kv8 != null or s.recycle_owner == null) return false; + const placement = @import("cuda_exl3_kv.zig").of(s.recycle_owner.?); + return placement.k8 == null and placement.heads == @import("cuda_prompt_wt.zig").kv_heads and (s.ikc_host == null or s.paged != null); +} + +fn matches(sequence: *const prompt.Seq, engine: *prompt.Engine, cap: usize) bool { + if (sequence.recycle_owner != engine or sequence.tail.d != engine.driver or sequence.cap != cap) return false; + const placement = @import("cuda_exl3_kv.zig").of(engine); + if (placement.k8 != null or placement.heads != @import("cuda_prompt_wt.zig").kv_heads or placement.ikc_host != (sequence.ikc_host != null)) return false; + const cfg = if (engine.page_config) |config| config.forChunk(engine.chunk) else null; + if (sequence.paged) |pg| return cfg != null and std.meta.eql(pg.cfg, cfg.?) and pg.shared_arena == engine.shared_arena; + return cfg == null; +} + +/// Seq.deinit's hook: true when the sequence's BF16 allocations were kept for reuse. +pub fn giveSeq(s: *prompt.Seq) bool { + if (draining or !enabled() or !canRetainSeq(s)) return false; + for (&seqs) |*slot| if (slot.* == null) { + if (s.paged) |pg| { + @import("cuda_exl3_graph.zig").releaseDecodeSeq(s); + pg.retireForReuse() catch @panic("recycled shared KV owner did not retire"); + } else s.tail.d.synchronize() catch @panic("recycled sequence consumers did not drain"); + slot.* = s.*; + stats.seq_kept += 1; + return true; + }; + return false; +} + +test "exl3: invalidated restore destinations are never retained for reuse" { + var sequence: prompt.Seq = undefined; + sequence.invalidated = true; + try std.testing.expect(!canRetainSeq(&sequence)); +} + +/// Draft.init's hook: a kept draft window of the same shape, its caches zeroed and its state reset; null: allocate. +pub fn takeDraft(stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !?Draft { + if (!enabled()) return null; + for (&drafts) |*slot| if (slot.*) |*d| if (d.cap == cap and d.dlogits.len == draft_n * 2 and d.rows_cap == @max(rows, @import("cuda_mtp.zig").max_rows)) { + var x = d.*; + x.verify_cache.invalidate(); + slot.* = null; + for ([_]cuda.DeviceBuffer{ x.kc, x.vc, x.ikc, x.pooled }) |b| try b.fill8(0, stream.handle); + x.media = null; + x.len = 0; + x.drafted = 0; + x.rows = 0; + x.open = false; + x.prompt_last = 0; + x.fed0 = 0; + x.history = .{ 0, 0 }; + x.unused_history = .{ 0, 0 }; + stats.draft_reused += 1; + return x; + }; + return null; +} + +/// Draft.deinit's hook: true when the window was kept. +pub fn giveDraft(d: *Draft) bool { + if (draining or !enabled() or d.media != null) return false; + for (&drafts) |*slot| if (slot.* == null) { + slot.* = d.*; + stats.draft_kept += 1; + return true; + }; + return false; +} + +/// Frees everything kept (the engine's deinit). +pub fn drain(engine: *prompt.Engine) void { + draining = true; + defer draining = false; + for (&seqs) |*slot| if (slot.*) |*sequence| if (sequence.recycle_owner == engine) { + discardSeq(slot, engine.gpa); + }; + for (&drafts) |*slot| if (slot.*) |*d| { + d.deinit(); + slot.* = null; + }; +} + +test { + _ = @import("cuda_exl3_recycle_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_recycle_test.zig b/zig/src/families/flashnext/cuda_exl3_recycle_test.zig new file mode 100644 index 0000000000..277c78761e --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_recycle_test.zig @@ -0,0 +1,421 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); +const recycle = @import("cuda_exl3_recycle.zig"); +const graphs = @import("cuda_exl3_graph.zig"); +const kv = @import("cuda_exl3_kv.zig"); +const fixture = @import("cuda_snap_test_fixture.zig"); +const shared_fixture = @import("cuda_exl3_shared_kv_test.zig"); +const allocator = std.testing.allocator; + +const Recording = struct { + var api: cuda.abi.Api = undefined; + var allocations: usize = 0; + var fills: usize = 0; + var fail_fill_at: ?usize = null; + var host_allocations: [64]?[]align(16) u8 = @splat(null); + + fn hostAlloc(output: *?*anyopaque, len: usize, _: c_uint) callconv(.c) c_int { + for (&host_allocations) |*slot| if (slot.* == null) { + const bytes = allocator.alignedAlloc(u8, .@"16", len) catch return 2; + slot.* = bytes; + output.* = bytes.ptr; + return 0; + }; + return 2; + } + + fn hostFree(address: ?*anyopaque) callconv(.c) c_int { + fixture.Fake.freed_pending = fixture.Fake.freed_pending or fixture.Fake.pending != 0; + for (&host_allocations) |*slot| if (slot.*) |bytes| if (@intFromPtr(bytes.ptr) == @intFromPtr(address)) { + allocator.free(bytes); + slot.* = null; + return 0; + }; + return 1; + } + + fn alloc(output: *u64, len: usize) callconv(.c) c_int { + allocations += 1; + return api.cuMemAlloc_v2(output, len); + } + + fn fill(address: u64, value: u8, len: usize, stream: cuda.abi.Stream) callconv(.c) c_int { + fills += 1; + if (fail_fill_at == fills) return cuda.abi.error_not_ready; + return api.cuMemsetD8Async(address, value, len, stream); + } + + fn unload(_: cuda.abi.Module) callconv(.c) c_int { + return 0; + } + + fn driver() cuda.Driver { + fixture.Fake.reset(); + var driver_ = fixture.Fake.driver(); + api = driver_.api; + allocations = 0; + fills = 0; + fail_fill_at = null; + driver_.api.cuMemAlloc_v2 = alloc; + driver_.api.cuMemsetD8Async = fill; + driver_.api.cuModuleUnload = unload; + driver_.api.cuMemHostAlloc = hostAlloc; + driver_.api.cuMemFreeHost = hostFree; + return driver_; + } +}; + +fn engine(driver: *cuda.Driver, stream: *cuda.Stream, config: ?pager.Config, arena: ?*pager.shared_device.Arena) prompt.Engine { + var result: prompt.Engine = undefined; + result.gpa = allocator; + result.driver = driver; + result.stream = stream; + result.page_config = config; + result.shared_arena = arena; + result.chunk = 1; + result.eos = 42; + result.round_graph = .{}; + return result; +} + +fn buffer(driver: *const cuda.Driver, len: usize) !cuda.DeviceBuffer { + const result = try cuda.DeviceBuffer.alloc(driver, len); + try result.fill8(0x7f, null); + return result; +} + +fn buffers(driver: *const cuda.Driver) ![]cuda.DeviceBuffer { + const result = try allocator.alloc(cuda.DeviceBuffer, 1); + result[0] = try buffer(driver, 16); + return result; +} + +fn sequence(owner: *prompt.Engine, cap: usize) !prompt.Seq { + const driver = owner.driver; + var pg: ?*pager.Paged = null; + if (owner.page_config) |cfg| { + if (owner.shared_arena) |arena| { + pg = try pager.Paged.initShared(allocator, driver, cap, cfg, arena); + } else { + const layout = try pager.SnapshotLayout.init(cap, cfg, 0); + const store = try allocator.create(pager.Paged); + const host = try cuda.HostBuffer.allocMapped(driver, 24 * cap * pager.row_bytes); + store.* = .{ + .gpa = allocator, + .d = driver, + .k = undefined, + .cfg = cfg, + .cap = cap, + .wr = layout.wr, + .nblocks = layout.nblocks, + .host = host, + .host_dev = try host.device(), + .status_host = try cuda.HostBuffer.alloc(driver, 12 * 8 * 4), + .dev = try buffer(driver, layout.device), + .tabs = try buffer(driver, layout.tables), + .work = try buffer(driver, 16), + }; + store.k.module = .{ .d = driver, .handle = null }; + pg = store; + } + } + var result: prompt.Seq = .{ + .recycle_owner = owner, + .conv = try buffers(driver), + .state = try buffers(driver), + .kc = try buffers(driver), + .vc = try buffers(driver), + .ikc = try buffers(driver), + .pooled = try buffers(driver), + .tail = try buffer(driver, 16), + .history = .{ 1, 2 }, + .fed = cap, + .rank_id = 99, + .untaught = true, + .logits = try allocator.alloc(u8, 16), + .cap = cap, + .paged = pg, + }; + if (kv.of(owner).ikc_host) { + result.ikc[0].free(); + result.ikc_host = try cuda.HostBuffer.allocMapped(driver, 16); + result.ikc[0] = .{ .d = driver, .ptr = try result.ikc_host.?.device(), .len = 16, .borrowed = true }; + } + return result; +} + +fn zero(bytes: []const u8) !void { + for (bytes) |value| try std.testing.expectEqual(@as(u8, 0), value); +} + +test "exl3: paged reuse keeps allocations with fresh incarnations and a complete BF16 reset" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + for ([_]bool{ false, true }) |shared| for ([_]bool{ false, true }) |host_index| { + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + const cfg: pager.Config = .{ .window = 2052, .rows = 1, .slots = 3, .shared_lanes = if (shared) 2 else 0 }; + const arena = if (shared) try pager.shared_device.Arena.initStorage(allocator, &driver, .{ .cap = 16, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }) else null; + defer if (arena) |active| active.destroy() catch unreachable; + if (arena) |active| shared_fixture.Fake.install(active); + var owner = engine(&driver, &stream, cfg, arena); + defer kv.forget(allocator, &owner); + if (host_index) try kv.enable(&owner, false, false, true); + defer recycle.drain(&owner); + var current = try sequence(&owner, 16); + const pg = current.paged.?; + try std.testing.expectError(error.KvPageNotRetired, pg.resetForReuse(stream)); + const addresses = [_]u64{ current.tail.ptr, pg.host_dev, pg.dev.ptr, pg.tabs.ptr, pg.work.ptr }; + const count = Recording.allocations; + var write: [16]u8 = @splat(0x5a); + for (0..8) |_| { + current.fed = 16; + current.history = .{ 1, 2 }; + current.rank_id = 99; + current.untaught = true; + pg.win0 = @splat(4); + pg.rebases = @splat(9); + pg.written = @splat(7); + pg.graph_prepared = true; + @memset(pg.host.bytes, 0x6a); + @memset(pg.status_host.bytes, 0xff); + @memset(current.logits, 0xff); + const key = pg.graphKey(); + const identity = pg.telemetry_id; + const incarnation = pg.incarnation; + try driver.check(driver.api.cuMemcpyDtoDAsync_v2(pg.host_dev, @intFromPtr(&write), write.len, stream.handle), "queue prior request write"); + current.deinit(allocator); + try std.testing.expectEqual(@as(usize, 0), fixture.Fake.pending); + try std.testing.expectEqualSlices(u8, &write, pg.host.bytes[0..write.len]); + try std.testing.expect(pg.retired_for_reuse); + try std.testing.expectError(error.KvPageRetired, pg.retireForReuse()); + current = try owner.newSeq(16); + try std.testing.expectEqual(addresses, [_]u64{ current.tail.ptr, pg.host_dev, pg.dev.ptr, pg.tabs.ptr, pg.work.ptr }); + try std.testing.expectEqual(count, Recording.allocations); + try std.testing.expect(pg.telemetry_id != identity and pg.incarnation != incarnation); + try std.testing.expect(!std.meta.eql(key, pg.graphKey())); + try std.testing.expectEqual(@as(usize, 0), current.fed); + try std.testing.expectEqual([2]i64{ 42, 42 }, current.history); + try std.testing.expectEqual(@as(u64, 0), current.rank_id); + try std.testing.expect(!current.untaught and !current.invalidated and !pg.graph_prepared); + try std.testing.expect(!pg.retired_for_reuse); + try std.testing.expectEqual(@as(u64, 0), pg.telemetry_revision); + try std.testing.expectEqual(@as(usize, 0), pg.win0[0]); + try std.testing.expectEqual(@as(u64, 0), pg.rebases[0]); + try std.testing.expectEqual(@as(u64, 0), pg.written[0]); + try std.testing.expectEqual(@as(u64, 0), (try pg.total()).calls); + try zero(pg.host.bytes); + try zero(pg.status_host.bytes); + try zero(current.logits); + for ([_][]cuda.DeviceBuffer{ current.conv, current.state, current.kc, current.vc, current.ikc, current.pooled }) |group| for (group) |item| try zero(@as([*]const u8, @ptrFromInt(item.ptr))[0..item.len]); + try zero(@as([*]const u8, @ptrFromInt(current.tail.ptr))[0..current.tail.len]); + } + current.deinit(allocator); + try std.testing.expect(!fixture.Fake.freed_pending); + }; +} + +test "exl3: request allocation retention is bounded to four sequences" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream, null, null); + defer recycle.drain(&owner); + const before = cuda.memory.usage(false); + var requests: [5]prompt.Seq = undefined; + for (&requests) |*request| request.* = try sequence(&owner, 16); + for (&requests) |*request| request.deinit(allocator); + var retained: [4]prompt.Seq = undefined; + for (&retained) |*request| request.* = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + for (&retained) |*request| { + request.invalidated = true; + request.deinit(allocator); + } + const after = cuda.memory.usage(false); + try std.testing.expectEqual(before.device, after.device); + try std.testing.expectEqual(before.host, after.host); +} + +test "exl3: shared recycle releases slots and rejects stale owners while snapshots remain independent" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + const cfg: pager.Config = .{ .window = 2052, .rows = 1, .slots = 3, .shared_lanes = 2 }; + const arena = try pager.shared_device.Arena.initStorage(allocator, &driver, .{ .cap = 16, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + var owner = engine(&driver, &stream, cfg, arena); + defer recycle.drain(&owner); + var current = try sequence(&owner, 16); + const pg = current.paged.?; + const old = pg.shared_sequence; + const other = try arena.register(16); + defer arena.release(other, stream) catch unreachable; + shared_fixture.seed(arena, old, 0, 0, 0, 5, 0x71); + shared_fixture.seed(arena, other, 0, 1, 1, 6, 0x72); + const held = try arena.acquire(old); + try std.testing.expectError(error.SharedKvBusy, pg.retireForReuse()); + var saved = try arena.saveOwnedBlocks(allocator, held); + defer saved.deinit(allocator); + try arena.cancel(held); + fixture.Fake.fail_sync_at = fixture.Fake.sync_calls + 1; + try std.testing.expectError(error.NotReady, pg.retireForReuse()); + try std.testing.expect(!pg.retired_for_reuse and pg.shared_attached); + fixture.Fake.fail_sync_at = null; + const still_live = try arena.acquire(old); + try arena.cancel(still_live); + current.deinit(allocator); + try std.testing.expect(!pg.shared_attached); + try std.testing.expectError(error.SharedKvSequence, arena.acquire(old)); + try std.testing.expectError(error.SharedKvSequence, pg.snapshot(0)); + const replacement = try arena.register(16); + try std.testing.expectEqual(old.lane, replacement.lane); + try std.testing.expect(replacement.generation > old.generation); + try arena.release(replacement, stream); + current = try owner.newSeq(16); + try std.testing.expectEqual(pg, current.paged.?); + const fresh = current.paged.?.shared_sequence; + try std.testing.expect(fresh.generation > replacement.generation); + try std.testing.expectError(error.SharedKvSequence, arena.acquire(old)); + const lease = try arena.acquire(fresh); + try arena.restoreOwnedBlocks(lease, &saved); + try arena.complete(lease, stream); + const restored_lease = try arena.acquire(fresh); + var restored = try arena.saveOwnedBlocks(allocator, restored_lease); + defer restored.deinit(allocator); + try arena.cancel(restored_lease); + try std.testing.expectEqualSlices(u8, saved.payload, restored.payload); + try std.testing.expectEqual(old, saved.source); + try std.testing.expectEqual(fresh, restored.source); + const other_lease = try arena.acquire(other); + var other_image = try arena.saveOwnedBlocks(allocator, other_lease); + defer other_image.deinit(allocator); + try arena.cancel(other_lease); + try std.testing.expectEqual(@as(usize, 1), other_image.entries.len); + try std.testing.expectEqual(@as(u8, 0x72), other_image.payload[0]); + current.deinit(allocator); + const blocked = try arena.register(16); + try std.testing.expectError(error.SharedKvLanes, recycle.takeSeq(&owner, 16)); + try arena.release(blocked, stream); +} + +test "exl3: recycler isolates engine shape and driver and drains only its owner's sequences" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream, null, null); + var other = engine(&driver, &stream, null, null); + defer recycle.drain(&owner); + defer recycle.drain(&other); + var first = try sequence(&owner, 16); + first.deinit(allocator); + try std.testing.expect(try recycle.takeSeq(&other, 16) == null); + const evicted = recycle.stats.seq_evicted; + try std.testing.expect(try recycle.takeSeq(&owner, 17) == null); + try std.testing.expectEqual(evicted + 1, recycle.stats.seq_evicted); + first = try sequence(&owner, 16); + first.deinit(allocator); + owner.page_config = .{}; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + try std.testing.expectEqual(evicted + 2, recycle.stats.seq_evicted); + owner.page_config = null; + first = try sequence(&owner, 16); + first.deinit(allocator); + var wrong_driver: cuda.Driver = undefined; + owner.driver = &wrong_driver; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + try std.testing.expectEqual(evicted + 3, recycle.stats.seq_evicted); + owner.driver = &driver; + first = try sequence(&owner, 16); + const address = first.tail.ptr; + first.deinit(allocator); + var second = try sequence(&other, 16); + second.deinit(allocator); + recycle.drain(&other); + first = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expectEqual(address, first.tail.ptr); + first.invalidated = true; + try std.testing.expect(!recycle.giveSeq(&first)); + first.deinit(allocator); +} + +test "exl3: private c8 capacity churn evicts retained paging and host indexer allocations under the original limit" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + const cfg: pager.Config = .{ .window = 2052, .rows = 1, .slots = 3 }; + var owner = engine(&driver, &stream, cfg, null); + try kv.enable(&owner, false, false, true); + defer kv.forget(allocator, &owner); + defer recycle.drain(&owner); + var active: [8]?prompt.Seq = @splat(null); + defer for (&active) |*entry| if (entry.*) |*request| request.deinit(allocator); + for (&active) |*request| request.* = try sequence(&owner, 64); + const admitted = cuda.memory.usage(false); + const previous_limit = try cuda.memory.setAllocationLimit(admitted.device + admitted.host); + defer _ = cuda.memory.setAllocationLimit(previous_limit) catch unreachable; + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + const evicted = recycle.stats.seq_evicted; + for (active[0..4]) |*entry| { + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + entry.* = try sequence(&owner, 16); + } + try std.testing.expectEqual(evicted + 4, recycle.stats.seq_evicted); + const mixed = cuda.memory.usage(false); + try std.testing.expect(mixed.device + mixed.host <= admitted.device + admitted.host); + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + const count = Recording.allocations; + for (active[0..4]) |*entry| entry.* = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expectEqual(count, Recording.allocations); + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + graphs.force = .off; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + try std.testing.expectEqual(evicted + 8, recycle.stats.seq_evicted); +} + +test "exl3: a partial recycled reset drains and frees its allocations without retention" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + for ([_]bool{ false, true }) |shared| { + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + const cfg: pager.Config = .{ .window = 2052, .rows = 1, .slots = 3, .shared_lanes = if (shared) 2 else 0 }; + const arena = if (shared) try pager.shared_device.Arena.initStorage(allocator, &driver, .{ .cap = 16, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }) else null; + defer if (arena) |active| active.destroy() catch unreachable; + if (arena) |active| shared_fixture.Fake.install(active); + var owner = engine(&driver, &stream, cfg, arena); + defer recycle.drain(&owner); + const before = cuda.memory.usage(false); + var current = try sequence(&owner, 16); + current.deinit(allocator); + Recording.fail_fill_at = Recording.fills + 3; + try std.testing.expectError(error.NotReady, recycle.takeSeq(&owner, 16)); + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + const after = cuda.memory.usage(false); + try std.testing.expectEqual(before.device, after.device); + try std.testing.expectEqual(before.host, after.host); + try std.testing.expect(!fixture.Fake.freed_pending); + } +} diff --git a/zig/src/families/flashnext/cuda_exl3_shared_kv.zig b/zig/src/families/flashnext/cuda_exl3_shared_kv.zig new file mode 100644 index 0000000000..70c78b7a29 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_shared_kv.zig @@ -0,0 +1,516 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const contract = @import("exl3_shared_kv.zig"); +const parts = @import("cuda_batch_state.zig"); + +const row_bytes = wt.kv_heads * wt.head_dim * 2; +const ratio = wt.index_ratio; +const controls_words = 8; +const work_entries = parts.att_rows * wt.index_top; +pub const max_lanes = 32; + +pub const Options = struct { cap: usize, lanes: usize, window: usize, rows: usize, slots: usize }; + +pub const Layout = struct { + options: Options, + window_rows: usize, + blocks: usize, + lane_words: usize, + all_words: usize, + side_rows: usize, + device_bytes: usize, + marks: usize, + controls: usize, + owners: usize, + stamps: usize, + owner_generations: usize, + lane_generations: usize, + clock: usize, + table_bytes: usize, + work_bytes: usize, + common_bytes: usize, + + pub fn init(options: Options) !Layout { + if (options.cap == 0 or options.cap > wt.window or options.lanes == 0 or options.lanes > max_lanes or options.window < wt.dense_keys or options.rows == 0 or options.slots == 0) return error.SharedKvGeometry; + const window_rows = try std.math.add(usize, try std.math.add(usize, options.window, options.rows), ratio); + const blocks = (try std.math.add(usize, options.cap, ratio - 1)) / ratio + 1; + const lane_words = try std.math.mul(usize, wt.attn_n, blocks); + const all_words = try std.math.mul(usize, options.lanes, lane_words); + const window_count = try std.math.mul(usize, options.lanes, wt.attn_n); + const side_rows = try std.math.add(usize, try std.math.mul(usize, window_count, window_rows), try std.math.mul(usize, ratio, options.slots)); + if (all_words > std.math.maxInt(c_int) or side_rows > std.math.maxInt(c_int) or options.slots > std.math.maxInt(c_int)) return error.SharedKvGeometry; + const device_bytes = try std.math.mul(usize, side_rows, 2 * row_bytes); + const marks = try std.math.mul(usize, all_words, 4); + const controls = try std.math.mul(usize, marks, 2); + const owners = try std.math.add(usize, controls, try std.math.mul(usize, window_count, controls_words * 4)); + const stamps = try std.math.add(usize, owners, try std.math.mul(usize, options.slots, 4)); + const after_stamps = try std.math.add(usize, stamps, try std.math.mul(usize, options.slots, 4)); + const owner_generations = try std.math.add(usize, after_stamps, (8 - after_stamps % 8) % 8); + const lane_generations = try std.math.add(usize, owner_generations, try std.math.mul(usize, options.slots, 8)); + const clock = try std.math.add(usize, lane_generations, try std.math.mul(usize, options.lanes, 8)); + const table_bytes = try std.math.add(usize, clock, 4); + const work_bytes = work_entries * 4 * 4; + return .{ .options = options, .window_rows = window_rows, .blocks = blocks, .lane_words = lane_words, .all_words = all_words, .side_rows = side_rows, .device_bytes = device_bytes, .marks = marks, .controls = controls, .owners = owners, .stamps = stamps, .owner_generations = owner_generations, .lane_generations = lane_generations, .clock = clock, .table_bytes = table_bytes, .work_bytes = work_bytes, .common_bytes = try std.math.add(usize, device_bytes, try std.math.add(usize, table_bytes, work_bytes)) }; + } + + pub fn windowOffset(layout: Layout, lane: usize, layer: usize) !usize { + if (lane >= layout.options.lanes or layer >= wt.attn_n) return error.SharedKvKey; + return (lane * wt.attn_n + layer) * layout.window_rows; + } + + pub fn slotBase(layout: Layout, lane: usize, layer: usize) !usize { + return layout.options.lanes * wt.attn_n * layout.window_rows - try layout.windowOffset(lane, layer); + } +}; + +const Lane = struct { generation: u64 = 0, live: bool = false, failed: bool = false }; +pub const Lease = struct { arena: u64, serial: u64, sequence: contract.Sequence }; +const Active = struct { + lease: Lease, + owners: [max_lanes]contract.Sequence = undefined, + count: usize, + touched: [max_lanes]bool = @splat(false), + stream_bound: bool = false, + stream: cuda.abi.Stream = null, + submitted: bool = false, + failed: bool = false, +}; +pub const View = struct { keys: u64, values: u64, inverse: u64, marks: u64, controls: u64, slot_base: usize }; +pub const Kernels = struct { plan: cuda.Function, fetch: cuda.Function, remap: cuda.Function, release: cuda.Function }; + +pub const Arena = struct { + allocator: std.mem.Allocator, + driver: *const cuda.Driver, + layout: Layout, + id: u64, + device: cuda.DeviceBuffer, + tables: cuda.DeviceBuffer, + work: cuda.DeviceBuffer, + metadata: []align(8) u8, + lanes: [max_lanes]Lane = @splat(.{}), + serial: u64 = 0, + active: ?Active = null, + module: ?cuda.Module = null, + kernels: ?Kernels = null, + + pub fn init(allocator: std.mem.Allocator, driver: *const cuda.Driver, options: Options) !*Arena { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const arena = try initStorage(allocator, driver, options); + errdefer arena.destroy() catch unreachable; + var module = try cuda.Module.load(driver, cuda.kernels.flashnext_exl3_kvpage); + errdefer module.unload(); + arena.kernels = .{ .plan = try module.function("kvp_shared_plan"), .fetch = try module.function("kvp_fetch"), .remap = try module.function("kvp_remap"), .release = try module.function("kvp_shared_release") }; + arena.module = module; + return arena; + } + + pub fn initStorage(allocator: std.mem.Allocator, driver: *const cuda.Driver, options: Options) !*Arena { + const layout = try Layout.init(options); + const arena = try allocator.create(Arena); + errdefer allocator.destroy(arena); + var device = try cuda.DeviceBuffer.alloc(driver, layout.device_bytes); + errdefer device.free(); + var tables = try cuda.DeviceBuffer.alloc(driver, layout.table_bytes); + errdefer tables.free(); + var work = try cuda.DeviceBuffer.alloc(driver, layout.work_bytes); + errdefer work.free(); + const metadata = try allocator.alignedAlloc(u8, .@"8", layout.table_bytes); + errdefer allocator.free(metadata); + @memset(metadata, 0); + @memset(metadata[0..layout.marks], 0xff); + @memset(metadata[layout.owners..layout.stamps], 0xff); + errdefer cuda.completion.drain(driver, .{ .d = driver, .handle = null }, "shared KV initialization did not drain"); + try driver.check(driver.api.cuMemcpyHtoD_v2(tables.ptr, metadata.ptr, metadata.len), "cuMemcpyHtoD(shared KV tables)"); + try driver.check(driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV initialization)"); + arena.* = .{ .allocator = allocator, .driver = driver, .layout = layout, .id = try contract.identity(), .device = device, .tables = tables, .work = work, .metadata = metadata }; + return arena; + } + + pub fn destroy(arena: *Arena) !void { + if (arena.active != null) return error.SharedKvBusy; + for (arena.lanes) |lane| if (lane.live) return error.SharedKvSequenceLive; + if (arena.module) |*module| module.unload(); + arena.allocator.free(arena.metadata); + arena.work.free(); + arena.tables.free(); + arena.device.free(); + arena.allocator.destroy(arena); + } + + fn state(arena: *Arena, sequence: contract.Sequence) !*Lane { + if (sequence.arena != arena.id or sequence.lane >= arena.layout.options.lanes) return error.SharedKvSequence; + const lane = &arena.lanes[sequence.lane]; + if (!lane.live or lane.generation != sequence.generation or sequence.generation == 0) return error.SharedKvSequence; + return lane; + } + + pub fn register(arena: *Arena, cap: usize) !contract.Sequence { + if (arena.active != null) return error.SharedKvBusy; + if (cap != arena.layout.options.cap) return error.SharedKvCommonCap; + for (arena.lanes[0..arena.layout.options.lanes], 0..) |*lane, index| { + if (lane.live) continue; + const next_generation = try std.math.add(u64, lane.generation, 1); + errdefer cuda.completion.drain(arena.driver, .{ .d = arena.driver, .handle = null }, "shared KV registration did not drain"); + const lane_rows = wt.attn_n * arena.layout.window_rows; + for (0..2) |side| { + const start = arena.device.ptr + (side * arena.layout.side_rows + index * lane_rows) * row_bytes; + try arena.driver.check(arena.driver.api.cuMemsetD8_v2(start, 0, lane_rows * row_bytes), "cuMemsetD8(shared KV window)"); + } + const inverse = arena.tables.ptr + index * arena.layout.lane_words * 4; + try arena.driver.check(arena.driver.api.cuMemsetD32_v2(inverse, 0xffffffff, arena.layout.lane_words), "cuMemsetD32(shared KV inverse)"); + try arena.driver.check(arena.driver.api.cuMemsetD32_v2(arena.tables.ptr + arena.layout.marks + index * arena.layout.lane_words * 4, 0, arena.layout.lane_words), "cuMemsetD32(shared KV marks)"); + try arena.driver.check(arena.driver.api.cuMemsetD32_v2(arena.tables.ptr + arena.layout.controls + index * wt.attn_n * controls_words * 4, 0, wt.attn_n * controls_words), "cuMemsetD32(shared KV controls)"); + try arena.driver.check(arena.driver.api.cuMemcpyHtoD_v2(arena.tables.ptr + arena.layout.lane_generations + index * 8, &next_generation, 8), "cuMemcpyHtoD(shared KV generation)"); + try arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV registration)"); + lane.* = .{ .live = true, .generation = next_generation }; + return .{ .arena = arena.id, .lane = @intCast(index), .generation = next_generation }; + } + return error.SharedKvLanes; + } + + pub fn acquire(arena: *Arena, sequence: contract.Sequence) !Lease { + return arena.acquireMany(&.{sequence}); + } + + pub fn acquireMany(arena: *Arena, sequences: []const contract.Sequence) !Lease { + if (arena.active != null) return error.SharedKvBusy; + if (sequences.len == 0 or sequences.len > arena.layout.options.lanes) return error.SharedKvLeaseOwners; + for (sequences, 0..) |sequence, index| { + const lane = try arena.state(sequence); + if (lane.failed) return error.SharedKvSequenceFailed; + for (sequences[0..index]) |previous| if (previous.lane == sequence.lane) return error.SharedKvLeaseOwners; + } + const serial = try std.math.add(u64, arena.serial, 1); + const lease: Lease = .{ .arena = arena.id, .serial = serial, .sequence = sequences[0] }; + arena.serial = serial; + arena.active = .{ .lease = lease, .count = sequences.len }; + @memcpy(arena.active.?.owners[0..sequences.len], sequences); + return lease; + } + + fn held(arena: *Arena, lease: Lease) !*Active { + const active = if (arena.active) |*active| active else return error.SharedKvLease; + if (!std.meta.eql(active.lease, lease)) return error.SharedKvLease; + if (active.failed) return error.SharedKvPoisoned; + return active; + } + + pub fn submitted(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + active.submitted = true; + for (active.owners[0..active.count]) |sequence| active.touched[sequence.lane] = true; + } + + pub fn bindStream(arena: *Arena, lease: Lease, stream: cuda.Stream) !void { + const active = try arena.held(lease); + if (stream.d != arena.driver) return error.SharedKvDriver; + if (active.stream_bound and active.stream != stream.handle) return error.SharedKvStream; + active.stream_bound = true; + active.stream = stream.handle; + } + + pub fn submittedOn(arena: *Arena, lease: Lease, stream: cuda.Stream) !void { + try arena.bindStream(lease, stream); + try arena.submitted(lease); + } + + fn member(arena: *Arena, lease: Lease, sequence: contract.Sequence) !*Active { + const active = try arena.held(lease); + _ = try arena.state(sequence); + for (active.owners[0..active.count]) |owner| if (std.meta.eql(owner, sequence)) return active; + return error.SharedKvLeaseOwner; + } + + fn touch(arena: *Arena, lease: Lease, sequence: contract.Sequence) !void { + const active = try arena.member(lease, sequence); + active.submitted = true; + active.touched[sequence.lane] = true; + } + + pub fn leaseOwners(arena: *Arena, lease: Lease) ![]const contract.Sequence { + const active = try arena.held(lease); + return active.owners[0..active.count]; + } + + pub fn leaseFor(arena: *Arena, sequence: contract.Sequence) !Lease { + const active = arena.active orelse return error.SharedKvLease; + _ = try arena.member(active.lease, sequence); + return active.lease; + } + + pub fn clear(arena: *Arena, lease: Lease, sequence: contract.Sequence, stream: cuda.Stream) !void { + _ = try arena.member(lease, sequence); + try arena.bindStream(lease, stream); + const kernels = arena.kernels orelse return error.BuiltWithoutKernels; + var args: cuda.Args = .{}; + for ([_]u64{ arena.tables.ptr, arena.tables.ptr + arena.layout.owners, arena.tables.ptr + arena.layout.stamps, arena.tables.ptr + arena.layout.owner_generations }) |address| args.add(address); + for ([_]usize{ arena.layout.options.slots, @as(usize, sequence.lane) * arena.layout.lane_words, (@as(usize, sequence.lane) + 1) * arena.layout.lane_words }) |value| args.add(@as(c_int, @intCast(value))); + args.add(sequence.generation); + try arena.touch(lease, sequence); + errdefer arena.fail(lease) catch unreachable; + try cuda.launch.launch(kernels.release, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, stream, &args); + const inverse = arena.tables.ptr + @as(usize, sequence.lane) * arena.layout.lane_words * 4; + try arena.driver.check(arena.driver.api.cuMemsetD32Async(inverse, 0xffffffff, arena.layout.lane_words, stream.handle), "cuMemsetD32Async(shared KV inverse)"); + try arena.driver.check(arena.driver.api.cuMemsetD32Async(arena.tables.ptr + arena.layout.marks + @as(usize, sequence.lane) * arena.layout.lane_words * 4, 0, arena.layout.lane_words, stream.handle), "cuMemsetD32Async(shared KV marks)"); + try arena.driver.check(arena.driver.api.cuMemsetD32Async(arena.tables.ptr + arena.layout.controls + @as(usize, sequence.lane) * wt.attn_n * controls_words * 4, 0, wt.attn_n * controls_words, stream.handle), "cuMemsetD32Async(shared KV controls)"); + } + + pub fn cancel(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.submitted) return error.SharedKvPending; + arena.active = null; + } + + pub fn fail(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + active.failed = true; + for (active.touched, 0..) |touched, lane| if (touched) { + arena.lanes[lane].failed = true; + }; + } + + pub fn complete(arena: *Arena, lease: Lease, stream: cuda.Stream) !void { + return arena.finish(lease, stream, true); + } + + fn finish(arena: *Arena, lease: Lease, stream: cuda.Stream, check_status: bool) !void { + try arena.bindStream(lease, stream); + stream.synchronize() catch |err| { + try arena.fail(lease); + return err; + }; + arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV consumers)") catch |err| { + try arena.fail(lease); + return err; + }; + if (check_status) { + var controls: [wt.attn_n * controls_words]u32 = undefined; + const active = try arena.held(lease); + for (active.owners[0..active.count]) |owner| { + const control_base = arena.tables.ptr + arena.layout.controls + @as(usize, owner.lane) * controls.len * 4; + arena.driver.check(arena.driver.api.cuMemcpyDtoH_v2(&controls, control_base, controls.len * 4), "cuMemcpyDtoH(shared KV status)") catch |err| { + try arena.fail(lease); + return err; + }; + for (0..wt.attn_n) |layer| if (controls[layer * controls_words + 2] != 0) { + try arena.fail(lease); + return error.SharedKvOverflow; + }; + } + } + try arena.completeAfterConsumers(lease); + } + + pub fn completeAfterConsumers(arena: *Arena, lease: Lease) !void { + _ = try arena.held(lease); + arena.active = null; + } + + pub fn drainFailed(arena: *Arena, lease: Lease) !void { + const active = arena.active orelse return error.SharedKvLease; + if (!std.meta.eql(active.lease, lease) or !active.failed) return error.SharedKvLease; + try arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV failure)"); + arena.active = null; + } + + fn binding(arena: *Arena, sequence: contract.Sequence, layer: usize) !View { + _ = try arena.state(sequence); + const offset = try arena.layout.windowOffset(sequence.lane, layer); + const key_offset = @as(usize, sequence.lane) * arena.layout.lane_words + layer * arena.layout.blocks; + return .{ .keys = arena.device.ptr + offset * row_bytes, .values = arena.device.ptr + (arena.layout.side_rows + offset) * row_bytes, .inverse = arena.tables.ptr + key_offset * 4, .marks = arena.tables.ptr + arena.layout.marks + key_offset * 4, .controls = arena.tables.ptr + arena.layout.controls + (@as(usize, sequence.lane) * wt.attn_n + layer) * controls_words * 4, .slot_base = try arena.layout.slotBase(sequence.lane, layer) }; + } + + pub fn view(arena: *Arena, lease: Lease, sequence: contract.Sequence, layer: usize) !View { + _ = try arena.member(lease, sequence); + return arena.binding(sequence, layer); + } + + pub fn page(arena: *Arena, lease: Lease, stream: cuda.Stream, layer: usize, origin: usize, rows: usize, ids: u64, counts: u64, sparse: u64, host_keys: u64, host_values: u64) !void { + return arena.pageFor(lease, lease.sequence, stream, layer, origin, rows, ids, counts, sparse, host_keys, host_values); + } + + pub fn pageFor(arena: *Arena, lease: Lease, sequence: contract.Sequence, stream: cuda.Stream, layer: usize, origin: usize, rows: usize, ids: u64, counts: u64, sparse: u64, host_keys: u64, host_values: u64) !void { + _ = try arena.member(lease, sequence); + try arena.bindStream(lease, stream); + if (rows == 0 or rows > parts.att_rows or origin % ratio != 0 or origin > arena.layout.options.cap or @min(origin / ratio, try std.math.mul(usize, rows, wt.index_top)) > arena.layout.options.slots) return error.SharedKvCapacity; + const kernels = arena.kernels orelse return error.BuiltWithoutKernels; + const window = try arena.view(lease, sequence, layer); + const needed = arena.work.ptr; + const missing = needed + work_entries * 4; + const fetched = missing + work_entries * 4; + var args: cuda.Args = .{}; + for ([_]u64{ ids, counts, sparse }) |address| args.add(address); + for ([_]usize{ rows, wt.index_idw, wt.index_top, ratio, origin }) |value| args.add(@as(c_int, @intCast(value))); + for ([_]u64{ arena.tables.ptr + arena.layout.clock, window.inverse, arena.tables.ptr }) |address| args.add(address); + args.add(@as(c_int, @intCast(@as(usize, sequence.lane) * arena.layout.lane_words + layer * arena.layout.blocks))); + for ([_]u64{ arena.tables.ptr + arena.layout.owners, arena.tables.ptr + arena.layout.stamps, window.marks }) |address| args.add(address); + args.add(@as(c_int, @intCast(arena.layout.options.slots))); + for ([_]u64{ window.controls, needed, missing, fetched, arena.tables.ptr + arena.layout.owner_generations, arena.tables.ptr + arena.layout.lane_generations }) |address| args.add(address); + args.add(sequence.generation); + args.add(@as(c_int, @intCast(arena.layout.lane_words))); + args.add(@as(c_int, @intCast(arena.layout.all_words))); + try arena.touch(lease, sequence); + errdefer arena.fail(lease) catch unreachable; + try cuda.launch.launch(kernels.plan, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream, &args); + var fetch_args: cuda.Args = .{}; + for ([_]u64{ host_keys, host_values, window.keys, window.values, window.controls, fetched }) |address| fetch_args.add(address); + for ([_]usize{ window.slot_base, row_bytes / 16, ratio }) |value| fetch_args.add(@as(c_int, @intCast(value))); + try cuda.launch.launch(kernels.fetch, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, stream, &fetch_args); + var remap_args: cuda.Args = .{}; + for ([_]u64{ ids, counts, sparse }) |address| remap_args.add(address); + for ([_]usize{ wt.index_idw, origin, window.slot_base }) |value| remap_args.add(@as(c_int, @intCast(value))); + remap_args.add(window.inverse); + remap_args.add(@as(c_int, ratio)); + remap_args.add(window.controls); + try cuda.launch.launch(kernels.remap, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, stream, &remap_args); + } + + pub fn release(arena: *Arena, sequence: contract.Sequence, stream: cuda.Stream) !void { + const lane = try arena.state(sequence); + if (arena.active != null) return error.SharedKvBusy; + if (stream.d != arena.driver) return error.SharedKvDriver; + const kernels = arena.kernels orelse return error.BuiltWithoutKernels; + const failed = lane.failed; + lane.failed = false; + const lease = arena.acquire(sequence) catch |err| { + lane.failed = failed; + return err; + }; + errdefer if (arena.active != null and !arena.active.?.failed) arena.fail(lease) catch unreachable; + var args: cuda.Args = .{}; + for ([_]u64{ arena.tables.ptr, arena.tables.ptr + arena.layout.owners, arena.tables.ptr + arena.layout.stamps, arena.tables.ptr + arena.layout.owner_generations }) |address| args.add(address); + for ([_]usize{ arena.layout.options.slots, @as(usize, sequence.lane) * arena.layout.lane_words, (@as(usize, sequence.lane) + 1) * arena.layout.lane_words }) |value| args.add(@as(c_int, @intCast(value))); + args.add(sequence.generation); + try arena.submitted(lease); + try cuda.launch.launch(kernels.release, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, stream, &args); + try arena.finish(lease, stream, false); + arena.lanes[sequence.lane].live = false; + } + + fn word(arena: *Arena, offset: usize) *u32 { + return @ptrCast(@alignCast(arena.metadata[offset..][0..4])); + } + + fn generation(arena: *Arena, offset: usize) *u64 { + return @ptrCast(@alignCast(arena.metadata[offset..][0..8])); + } + + fn downloadMetadata(arena: *Arena) !void { + try arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV metadata)"); + try arena.driver.check(arena.driver.api.cuMemcpyDtoH_v2(arena.metadata.ptr, arena.tables.ptr, arena.metadata.len), "cuMemcpyDtoH(shared KV metadata)"); + } + + fn geometry(arena: *const Arena) contract.Geometry { + return .{ .lanes = @intCast(arena.layout.options.lanes), .layers = wt.attn_n, .blocks = @intCast(arena.layout.blocks), .slots = @intCast(arena.layout.options.slots), .row_bytes = row_bytes, .window_rows = arena.layout.window_rows, .work_words = arena.layout.work_bytes / 4 }; + } + + pub fn saveOwnedBlocks(arena: *Arena, allocator: std.mem.Allocator, lease: Lease) !contract.OwnedBlocks { + const active = try arena.held(lease); + if (active.count != 1) return error.SharedKvSnapshotLease; + try arena.downloadMetadata(); + const control_base = arena.layout.controls + @as(usize, lease.sequence.lane) * wt.attn_n * controls_words * 4; + for (0..wt.attn_n) |layer| if (arena.word(control_base + (layer * controls_words + 2) * 4).* != 0) return error.SharedKvOverflow; + const first_key = @as(usize, lease.sequence.lane) * arena.layout.lane_words; + var count: usize = 0; + for (0..arena.layout.options.slots) |slot| { + const owner = arena.word(arena.layout.owners + slot * 4).*; + if (owner >= first_key and owner - first_key < arena.layout.lane_words and arena.generation(arena.layout.owner_generations + slot * 8).* == lease.sequence.generation) count += 1; + } + const entries = try allocator.alloc(contract.OwnedEntry, count); + errdefer allocator.free(entries); + const block_bytes = ratio * row_bytes; + const payload = try allocator.alloc(u8, count * 2 * block_bytes); + errdefer allocator.free(payload); + const slot_rows = arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows; + var index: usize = 0; + for (0..arena.layout.options.slots) |slot| { + const owner = arena.word(arena.layout.owners + slot * 4).*; + if (owner < first_key or owner - first_key >= arena.layout.lane_words or arena.generation(arena.layout.owner_generations + slot * 8).* != lease.sequence.generation) continue; + if (arena.word(owner * 4).* != slot) return error.SharedKvInverse; + const key = owner - first_key; + entries[index] = .{ .key = .{ .layer = @intCast(key / arena.layout.blocks), .block = @intCast(key % arena.layout.blocks) }, .stamp = arena.word(arena.layout.stamps + slot * 4).*, .source_slot = @intCast(slot) }; + for (0..2) |side| { + const address = arena.device.ptr + (side * arena.layout.side_rows + slot_rows + ratio * slot) * row_bytes; + try arena.driver.check(arena.driver.api.cuMemcpyDtoH_v2(payload[(2 * index + side) * block_bytes ..].ptr, address, block_bytes), "cuMemcpyDtoH(shared KV owned block)"); + } + index += 1; + } + var image: contract.OwnedBlocks = .{ .source = lease.sequence, .geometry = arena.geometry(), .clock = arena.word(arena.layout.clock).*, .entries = entries, .payload = payload }; + try image.validate(arena.geometry()); + return image; + } + + fn evictMetadata(arena: *Arena, slot: usize) void { + const owner = arena.word(arena.layout.owners + slot * 4).*; + const owner_generation = arena.generation(arena.layout.owner_generations + slot * 8).*; + if (owner < arena.layout.all_words and owner_generation == arena.generation(arena.layout.lane_generations + owner / arena.layout.lane_words * 8).* and arena.word(owner * 4).* == slot) { + arena.word(owner * 4).* = 0xffffffff; + if (arena.active) |*active| active.touched[owner / arena.layout.lane_words] = true; + } + arena.word(arena.layout.owners + slot * 4).* = 0xffffffff; + arena.word(arena.layout.stamps + slot * 4).* = 0; + arena.generation(arena.layout.owner_generations + slot * 8).* = 0; + } + + pub fn restoreOwnedBlocks(arena: *Arena, lease: Lease, image: *const contract.OwnedBlocks) !void { + const active = try arena.held(lease); + if (active.count != 1) return error.SharedKvSnapshotLease; + try arena.validateOwnedBlocks(image); + try arena.downloadMetadata(); + const clock_base = @max(arena.word(arena.layout.clock).*, @as(u32, @intCast(image.clock))); + const next_clock = try std.math.add(u32, try std.math.add(u32, clock_base, 1), std.math.cast(u32, try image.maxAge()) orelse return error.Overflow); + const targets = try arena.allocator.alloc(u32, arena.layout.options.slots); + defer arena.allocator.free(targets); + const first_key = @as(usize, lease.sequence.lane) * arena.layout.lane_words; + for (0..arena.layout.options.slots) |slot| { + const owner = arena.word(arena.layout.owners + slot * 4).*; + if (owner >= first_key and owner - first_key < arena.layout.lane_words and arena.generation(arena.layout.owner_generations + slot * 8).* == lease.sequence.generation) arena.evictMetadata(slot); + } + for (targets, 0..) |*target, slot| target.* = @intCast(slot); + const Order = struct { + fn lessThan(owner: *Arena, left: u32, right: u32) bool { + const left_stamp = owner.word(owner.layout.stamps + @as(usize, left) * 4).*; + const right_stamp = owner.word(owner.layout.stamps + @as(usize, right) * 4).*; + return left_stamp < right_stamp or (left_stamp == right_stamp and left < right); + } + }; + std.mem.sort(u32, targets, arena, Order.lessThan); + for (image.entries, 0..) |entry, index| { + const slot = targets[index]; + arena.evictMetadata(slot); + arena.word(arena.layout.owners + slot * 4).* = @intCast(first_key + @as(usize, entry.key.layer) * arena.layout.blocks + entry.key.block); + arena.generation(arena.layout.owner_generations + slot * 8).* = lease.sequence.generation; + arena.word(arena.layout.stamps + slot * 4).* = next_clock - @as(u32, @intCast(image.clock - entry.stamp)); + } + arena.word(arena.layout.clock).* = next_clock; + try arena.submitted(lease); + errdefer arena.fail(lease) catch unreachable; + errdefer cuda.completion.drain(arena.driver, .{ .d = arena.driver, .handle = null }, "shared KV snapshot upload did not drain"); + try arena.driver.check(arena.driver.api.cuMemcpyHtoD_v2(arena.tables.ptr, arena.metadata.ptr, arena.metadata.len), "cuMemcpyHtoD(shared KV invalidated maps)"); + const block_bytes = ratio * row_bytes; + const slot_rows = arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows; + for (targets[0..image.entries.len], 0..) |slot, index| for (0..2) |side| { + const address = arena.device.ptr + (side * arena.layout.side_rows + slot_rows + ratio * slot) * row_bytes; + try arena.driver.check(arena.driver.api.cuMemcpyHtoD_v2(address, image.payload[(2 * index + side) * block_bytes ..].ptr, block_bytes), "cuMemcpyHtoD(shared KV owned block)"); + }; + try arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV block payload)"); + for (targets[0..image.entries.len], image.entries) |slot, entry| arena.word((first_key + @as(usize, entry.key.layer) * arena.layout.blocks + entry.key.block) * 4).* = slot; + try arena.driver.check(arena.driver.api.cuMemcpyHtoD_v2(arena.tables.ptr, arena.metadata.ptr, arena.layout.marks), "cuMemcpyHtoD(shared KV published maps)"); + try arena.driver.check(arena.driver.api.cuCtxSynchronize(), "cuCtxSynchronize(shared KV published maps)"); + } + + pub fn validateOwnedBlocks(arena: *const Arena, image: *const contract.OwnedBlocks) !void { + try image.validate(arena.geometry()); + if (image.clock > std.math.maxInt(u32)) return error.Overflow; + } + + pub fn graphKey(arena: *Arena, sequence: contract.Sequence, layer: usize, canonical: u64, origin_hash: u64) !contract.GraphKey { + const window = try arena.binding(sequence, layer); + return .{ .sequence = sequence, .geometry = .{ .lanes = @intCast(arena.layout.options.lanes), .layers = wt.attn_n, .blocks = @intCast(arena.layout.blocks), .slots = @intCast(arena.layout.options.slots), .row_bytes = row_bytes, .window_rows = arena.layout.window_rows, .work_words = arena.layout.work_bytes / 4 }, .bindings = .{ .keys = arena.device.ptr + arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows * row_bytes, .values = arena.device.ptr + (arena.layout.side_rows + arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows) * row_bytes, .owners = arena.tables.ptr + arena.layout.owners, .stamps = arena.tables.ptr + arena.layout.stamps, .clock = arena.tables.ptr + arena.layout.clock, .inverse = window.inverse, .marks = window.marks, .controls = window.controls, .work = arena.work.ptr, .window_keys = window.keys, .window_values = window.values, .canonical = canonical, .cap = arena.layout.options.cap, .rows = arena.layout.options.rows, .origin_hash = origin_hash } }; + } +}; + +test { + _ = @import("cuda_exl3_shared_kv_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_shared_kv_test.zig b/zig/src/families/flashnext/cuda_exl3_shared_kv_test.zig new file mode 100644 index 0000000000..5ee609e763 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_shared_kv_test.zig @@ -0,0 +1,519 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const shared = @import("cuda_exl3_shared_kv.zig"); +const contract = @import("exl3_shared_kv.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const options: shared.Options = .{ .cap = 8192, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }; + +pub const Fake = struct { + var allocations: [256]?[]u8 = @splat(null); + var launches: usize = 0; + var fail_launch: bool = false; + var fail_sync: bool = false; + var expected_generation: u64 = 0; + var expected_key: usize = 0; + var expected_base: usize = 0; + var copy_calls: usize = 0; + var fail_copy: ?usize = null; + pub var sync_calls: usize = 0; + pub var graph_calls: usize = 0; + pub var graph_pending: bool = false; + pub var fail_graph: bool = false; + pub var fail_alloc_at: ?usize = null; + var device_alloc_calls: usize = 0; + + fn alloc(address: *u64, length: usize) callconv(.c) c_int { + device_alloc_calls += 1; + if (fail_alloc_at == device_alloc_calls) return cuda.abi.error_not_ready; + for (&allocations) |*allocation| if (allocation.* == null) { + const memory = std.heap.page_allocator.alloc(u8, length) catch return 2; + allocation.* = memory; + address.* = @intFromPtr(memory.ptr); + return 0; + }; + return 2; + } + + fn free(address: u64) callconv(.c) c_int { + for (&allocations) |*allocation| if (allocation.*) |memory| if (@intFromPtr(memory.ptr) == address) { + std.heap.page_allocator.free(memory); + allocation.* = null; + return 0; + }; + return 1; + } + + fn upload(destination: u64, source: ?*const anyopaque, length: usize) callconv(.c) c_int { + copy_calls += 1; + if (fail_copy == copy_calls) return cuda.abi.error_not_ready; + const target: [*]u8 = @ptrFromInt(destination); + const origin: [*]const u8 = @ptrCast(source.?); + @memcpy(target[0..length], origin[0..length]); + return 0; + } + + fn download(destination: ?*anyopaque, source: u64, length: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrCast(destination.?); + const origin: [*]const u8 = @ptrFromInt(source); + @memcpy(target[0..length], origin[0..length]); + return 0; + } + + fn downloadAsync(destination: ?*anyopaque, source: u64, length: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return download(destination, source, length); + } + + fn memset(address: u64, fill: c_uint, count: usize) callconv(.c) c_int { + const target: [*]u32 = @ptrFromInt(address); + @memset(target[0..count], fill); + return 0; + } + + fn memset8(address: u64, fill: u8, count: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrFromInt(address); + @memset(target[0..count], fill); + return 0; + } + + fn memsetAsync(address: u64, fill: c_uint, count: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return memset(address, fill, count); + } + + fn memset8Async(address: u64, fill: u8, count: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return memset8(address, fill, count); + } + + fn sync() callconv(.c) c_int { + sync_calls += 1; + if (!fail_sync) graph_pending = false; + return if (fail_sync) cuda.abi.error_not_ready else 0; + } + + fn graphLaunch(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + graph_calls += 1; + graph_pending = true; + return if (fail_graph) cuda.abi.error_not_ready else 0; + } + + fn streamSync(_: cuda.abi.Stream) callconv(.c) c_int { + return sync(); + } + + fn value(parameters: ?[*]?*anyopaque, index: usize, comptime Value: type) Value { + const pointer: *const Value = @ptrCast(@alignCast(parameters.?[index].?)); + return pointer.*; + } + + fn launch(function: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + launches += 1; + if (fail_launch) return cuda.abi.error_not_ready; + switch (@intFromPtr(function.?)) { + 1 => { + std.debug.assert(value(parameters, 22, u64) == expected_generation); + std.debug.assert(value(parameters, 11, c_int) == expected_key); + }, + 2 => std.debug.assert(value(parameters, 6, c_int) == expected_base), + 3 => std.debug.assert(value(parameters, 5, c_int) == expected_base), + 4 => { + const inverse: [*]i32 = @ptrFromInt(value(parameters, 0, u64)); + const owners: [*]i32 = @ptrFromInt(value(parameters, 1, u64)); + const stamps: [*]u32 = @ptrFromInt(value(parameters, 2, u64)); + const generations: [*]u64 = @ptrFromInt(value(parameters, 3, u64)); + const slots: usize = @intCast(value(parameters, 4, c_int)); + const begin = value(parameters, 5, c_int); + const end = value(parameters, 6, c_int); + const generation = value(parameters, 7, u64); + for (0..slots) |slot| if (owners[slot] >= begin and owners[slot] < end and generations[slot] == generation) { + inverse[@intCast(owners[slot])] = -1; + owners[slot] = -1; + stamps[slot] = 0; + generations[slot] = 0; + }; + }, + else => unreachable, + } + return 0; + } + + pub fn driver() cuda.Driver { + launches = 0; + fail_launch = false; + fail_sync = false; + copy_calls = 0; + fail_copy = null; + sync_calls = 0; + graph_calls = 0; + graph_pending = false; + fail_graph = false; + fail_alloc_at = null; + device_alloc_calls = 0; + var result: cuda.Driver = undefined; + result.api.cuMemAlloc_v2 = alloc; + result.api.cuMemFree_v2 = free; + result.api.cuMemcpyHtoD_v2 = upload; + result.api.cuMemcpyDtoH_v2 = download; + result.api.cuMemcpyDtoHAsync_v2 = downloadAsync; + result.api.cuMemsetD32_v2 = memset; + result.api.cuMemsetD8_v2 = memset8; + result.api.cuMemsetD32Async = memsetAsync; + result.api.cuMemsetD8Async = memset8Async; + result.api.cuCtxSynchronize = sync; + result.api.cuStreamSynchronize = streamSync; + result.api.cuLaunchKernel = launch; + result.api.cuGraphLaunch = graphLaunch; + return result; + } + + pub fn install(arena: *shared.Arena) void { + arena.kernels = .{ .plan = .{ .d = arena.driver, .handle = @ptrFromInt(1) }, .fetch = .{ .d = arena.driver, .handle = @ptrFromInt(2) }, .remap = .{ .d = arena.driver, .handle = @ptrFromInt(3) }, .release = .{ .d = arena.driver, .handle = @ptrFromInt(4) } }; + } +}; + +pub fn seed(arena: *shared.Arena, sequence: contract.Sequence, layer: usize, block: usize, slot: usize, stamp: u32, value: u8) void { + const tables: [*]u8 = @ptrFromInt(arena.tables.ptr); + const key = @as(usize, sequence.lane) * arena.layout.lane_words + layer * arena.layout.blocks + block; + std.mem.writeInt(u32, tables[key * 4 ..][0..4], @intCast(slot), .little); + std.mem.writeInt(u32, tables[arena.layout.owners + slot * 4 ..][0..4], @intCast(key), .little); + std.mem.writeInt(u32, tables[arena.layout.stamps + slot * 4 ..][0..4], stamp, .little); + std.mem.writeInt(u64, tables[arena.layout.owner_generations + slot * 8 ..][0..8], sequence.generation, .little); + const old = std.mem.readInt(u32, tables[arena.layout.clock..][0..4], .little); + std.mem.writeInt(u32, tables[arena.layout.clock..][0..4], @max(old, stamp), .little); + const row_bytes = wt.kv_heads * wt.head_dim * 2; + const slot_rows = arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows; + for (0..2) |side| { + const payload: [*]u8 = @ptrFromInt(arena.device.ptr + (side * arena.layout.side_rows + slot_rows + wt.index_ratio * slot) * row_bytes); + @memset(payload[0 .. wt.index_ratio * row_bytes], value); + } +} + +test "exl3: shared CUDA arena fixed windows bound every upstream BF16 slot address" { + const layout = try shared.Layout.init(.{ .cap = 200016, .lanes = 8, .window = 4096, .rows = 2048, .slots = 16384 }); + for (0..layout.options.lanes) |lane| for (0..wt.attn_n) |layer| { + const offset = try layout.windowOffset(lane, layer); + const base = try layout.slotBase(lane, layer); + try std.testing.expectEqual(layout.options.lanes * wt.attn_n * layout.window_rows, offset + base); + try std.testing.expectEqual(layout.side_rows, offset + base + wt.index_ratio * layout.options.slots); + try std.testing.expect(base >= layout.window_rows); + }; + try std.testing.expectEqual(layout.device_bytes + layout.table_bytes + layout.work_bytes, layout.common_bytes); + var bad = options; + bad.rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, shared.Layout.init(bad)); + bad = options; + bad.cap = wt.window + 1; + try std.testing.expectError(error.SharedKvGeometry, shared.Layout.init(bad)); +} + +test "exl3: shared CUDA arena pager launch ABI uses rank maps tags and global relative slots" { + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const first = try arena.register(options.cap); + defer arena.release(first, stream) catch unreachable; + const second = try arena.register(options.cap); + defer arena.release(second, stream) catch unreachable; + try std.testing.expectError(error.SharedKvCommonCap, arena.register(options.cap - 1)); + const lease = try arena.acquire(second); + Fake.expected_generation = second.generation; + Fake.expected_key = arena.layout.lane_words + 3 * arena.layout.blocks; + Fake.expected_base = try arena.layout.slotBase(second.lane, 3); + try arena.page(lease, stream, 3, 0, 1, 1, 2, 3, 4, 5); + try std.testing.expectEqual(@as(usize, 3), Fake.launches); + try std.testing.expectError(error.SharedKvBusy, arena.acquire(first)); + try std.testing.expectError(error.SharedKvPending, arena.cancel(lease)); + try arena.complete(lease, stream); + const viewing = try arena.acquire(second); + const view = try arena.view(viewing, second, 3); + try arena.cancel(viewing); + const row_bytes = wt.kv_heads * wt.head_dim * 2; + try std.testing.expectEqual(arena.device.ptr + arena.layout.options.lanes * wt.attn_n * arena.layout.window_rows * row_bytes, view.keys + view.slot_base * row_bytes); +} + +test "exl3: shared CUDA owned block snapshots remap slots without copying foreign payload" { + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const source = try arena.register(options.cap); + const foreign = try arena.register(options.cap); + defer arena.release(foreign, stream) catch unreachable; + seed(arena, source, 0, 0, 0, 1, 0xa1); + seed(arena, source, 1, 2, 1, 2, 0xa2); + const saving = try arena.acquire(source); + var image = try arena.saveOwnedBlocks(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + try arena.release(source, stream); + seed(arena, foreign, 0, 0, 0, 3, 0xb1); + const destination = try arena.register(options.cap); + defer arena.release(destination, stream) catch unreachable; + try std.testing.expect(destination.generation > source.generation); + const restoring = try arena.acquire(destination); + try arena.restoreOwnedBlocks(restoring, &image); + try arena.complete(restoring, stream); + const saving_again = try arena.acquire(destination); + var again = try arena.saveOwnedBlocks(std.testing.allocator, saving_again); + defer again.deinit(std.testing.allocator); + try arena.cancel(saving_again); + try std.testing.expectEqual(@as(usize, 2), again.entries.len); + try std.testing.expectEqualSlices(u8, image.payload, again.payload); + try std.testing.expect(again.entries[0].source_slot != image.entries[0].source_slot); + const foreign_saving = try arena.acquire(foreign); + var foreign_image = try arena.saveOwnedBlocks(std.testing.allocator, foreign_saving); + defer foreign_image.deinit(std.testing.allocator); + try arena.cancel(foreign_saving); + try std.testing.expectEqual(@as(usize, 1), foreign_image.entries.len); + try std.testing.expect(std.mem.allEqual(u8, foreign_image.payload, 0xb1)); +} + +test "exl3: shared CUDA failed launches and completion retain allocations and lease until drain" { + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const sequence = try arena.register(options.cap); + const lease = try arena.acquire(sequence); + Fake.fail_launch = true; + try std.testing.expectError(error.NotReady, arena.page(lease, stream, 0, 0, 1, 1, 2, 3, 4, 5)); + try std.testing.expectError(error.SharedKvBusy, arena.destroy()); + try std.testing.expectError(error.SharedKvPoisoned, arena.cancel(lease)); + Fake.fail_sync = true; + try std.testing.expectError(error.NotReady, arena.drainFailed(lease)); + try std.testing.expect(arena.active != null); + Fake.fail_sync = false; + Fake.fail_launch = false; + try arena.drainFailed(lease); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.acquire(sequence)); + try arena.release(sequence, stream); + const reused = try arena.register(options.cap); + try std.testing.expect(reused.generation > sequence.generation); + try arena.release(reused, stream); +} + +test "exl3: shared CUDA batch lease interleaves owners by layer without a host wait" { + var driver = Fake.driver(); + var batch_options = options; + batch_options.lanes = 3; + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, batch_options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const first = try arena.register(options.cap); + defer arena.release(first, stream) catch unreachable; + const second = try arena.register(options.cap); + defer arena.release(second, stream) catch unreachable; + const excluded = try arena.register(options.cap); + defer arena.release(excluded, stream) catch unreachable; + try std.testing.expectError(error.SharedKvLeaseOwners, arena.acquireMany(&.{})); + try std.testing.expectError(error.SharedKvLeaseOwners, arena.acquireMany(&.{ first, first })); + const lease = try arena.acquireMany(&.{ first, second }); + try std.testing.expectEqualSlices(contract.Sequence, &.{ first, second }, try arena.leaseOwners(lease)); + try std.testing.expectError(error.SharedKvLeaseOwner, arena.view(lease, excluded, 0)); + try std.testing.expectError(error.SharedKvLeaseOwner, arena.pageFor(lease, excluded, stream, 0, 0, 1, 1, 2, 3, 4, 5)); + var stale = second; + stale.generation += 1; + try std.testing.expectError(error.SharedKvSequence, arena.view(lease, stale, 0)); + const before = Fake.sync_calls; + for (0..wt.attn_n) |layer| for ([_]contract.Sequence{ first, second }) |sequence| { + Fake.expected_generation = sequence.generation; + Fake.expected_key = @as(usize, sequence.lane) * arena.layout.lane_words + layer * arena.layout.blocks; + Fake.expected_base = try arena.layout.slotBase(sequence.lane, layer); + try arena.pageFor(lease, sequence, stream, layer, 0, 1, 1, 2, 3, 4, 5); + }; + try std.testing.expectEqual(before, Fake.sync_calls); + const alternate: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(2) }; + try std.testing.expectError(error.SharedKvStream, arena.pageFor(lease, second, alternate, 0, 0, 1, 1, 2, 3, 4, 5)); + try std.testing.expectError(error.SharedKvBusy, arena.release(second, stream)); + try std.testing.expectError(error.SharedKvSnapshotLease, arena.saveOwnedBlocks(std.testing.allocator, lease)); + try arena.completeAfterConsumers(lease); + try std.testing.expectEqual(before, Fake.sync_calls); + try std.testing.expectError(error.SharedKvLease, arena.completeAfterConsumers(lease)); + const next = try arena.acquire(excluded); + try std.testing.expectError(error.SharedKvLease, arena.pageFor(lease, first, stream, 0, 0, 1, 1, 2, 3, 4, 5)); + try arena.cancel(next); +} + +test "exl3: shared CUDA captured replay poisons every submitted batch owner until drain" { + var driver = Fake.driver(); + var batch_options = options; + batch_options.lanes = 3; + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, batch_options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const first = try arena.register(options.cap); + const second = try arena.register(options.cap); + const untouched = try arena.register(options.cap); + const lease = try arena.acquireMany(&.{ first, second }); + try arena.submitted(lease); + try arena.fail(lease); + try std.testing.expect(arena.lanes[first.lane].failed); + try std.testing.expect(arena.lanes[second.lane].failed); + try std.testing.expect(!arena.lanes[untouched.lane].failed); + try std.testing.expectError(error.SharedKvPoisoned, arena.completeAfterConsumers(lease)); + try std.testing.expectError(error.SharedKvBusy, arena.acquire(untouched)); + try arena.drainFailed(lease); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.acquireMany(&.{ first, second })); + const other = try arena.acquire(untouched); + try arena.cancel(other); + try arena.release(first, stream); + try arena.release(second, stream); + try arena.release(untouched, stream); +} + +test "exl3: shared CUDA constructor allocation and upload failures refund all device ownership" { + for ([_]usize{ 1, 2, 3 }) |failure| { + var driver = Fake.driver(); + const before = cuda.usage(false).device; + Fake.fail_alloc_at = failure; + try std.testing.expectError(error.NotReady, shared.Arena.initStorage(std.testing.allocator, &driver, options)); + try std.testing.expectEqual(before, cuda.usage(false).device); + } + var driver = Fake.driver(); + const before = cuda.usage(false).device; + Fake.fail_copy = 1; + try std.testing.expectError(error.NotReady, shared.Arena.initStorage(std.testing.allocator, &driver, options)); + try std.testing.expectEqual(before, cuda.usage(false).device); +} + +test "exl3: shared CUDA cross-arena native snapshots advance beyond copied mark epochs" { + const snapshots = @import("cuda_snap.zig"); + const fixtures = @import("cuda_snap_test_fixture.zig"); + fixtures.Fake.reset(); + var driver = fixtures.Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const geometry: shared.Options = .{ .cap = 64, .lanes = 2, .window = 2052, .rows = 256, .slots = 16 }; + const source_arena = try shared.Arena.initStorage(std.testing.allocator, &driver, geometry); + defer source_arena.destroy() catch unreachable; + Fake.install(source_arena); + const source = try fixtures.sharedFixture(&driver, source_arena); + defer source.deinit(); + seed(source_arena, source.pg.shared_sequence, 0, 0, 1, 100, 0xa1); + const source_tables: [*]u8 = @ptrFromInt(source_arena.tables.ptr); + const source_marks = source_tables[source_arena.layout.marks..][0 .. source_arena.layout.lane_words * 4]; + std.mem.writeInt(u32, source_marks[0..4], 100, .little); + std.mem.writeInt(u32, source_marks[4..8], 2, .little); + const saved = try snapshots.save(std.testing.allocator, &driver, stream, &source.seq, null); + defer snapshots.drop(std.testing.allocator, stream, saved); + try std.testing.expectEqual(@as(u64, 100), saved.owned_blocks.?.clock); + try std.testing.expectEqual(@as(u64, 0), try saved.owned_blocks.?.maxAge()); + for ([_]u32{ 0, 200 }) |current_clock| { + const destination_arena = try shared.Arena.initStorage(std.testing.allocator, &driver, geometry); + defer destination_arena.destroy() catch unreachable; + Fake.install(destination_arena); + const reserved = try destination_arena.register(geometry.cap); + defer destination_arena.release(reserved, stream) catch unreachable; + const destination = try fixtures.sharedFixture(&driver, destination_arena); + defer destination.deinit(); + try std.testing.expect(source.pg.shared_sequence.arena != destination.pg.shared_sequence.arena); + try std.testing.expect(source.pg.shared_sequence.lane != destination.pg.shared_sequence.lane); + const tables: [*]u8 = @ptrFromInt(destination_arena.tables.ptr); + std.mem.writeInt(u32, tables[destination_arena.layout.clock..][0..4], current_clock, .little); + try snapshots.restore(stream, &destination.seq, null, saved); + const first_key = @as(usize, destination.pg.shared_sequence.lane) * destination_arena.layout.lane_words; + const copied_marks = tables[destination_arena.layout.marks + first_key * 4 ..][0..source_marks.len]; + try std.testing.expectEqualSlices(u8, source_marks, copied_marks); + try std.testing.expectEqual(@as(u32, 0xffffffff), std.mem.readInt(u32, tables[(first_key + 1) * 4 ..][0..4], .little)); + const restored_clock = std.mem.readInt(u32, tables[destination_arena.layout.clock..][0..4], .little); + try std.testing.expectEqual(@max(current_clock, @as(u32, 100)) + 1, restored_clock); + for (0..destination_arena.layout.lane_words) |index| { + const mark = std.mem.readInt(u32, copied_marks[index * 4 ..][0..4], .little); + try std.testing.expect(restored_clock > mark); + try std.testing.expect(restored_clock + 1 != mark); + } + const lease = try destination_arena.acquire(destination.pg.shared_sequence); + var restored = try destination_arena.saveOwnedBlocks(std.testing.allocator, lease); + defer restored.deinit(std.testing.allocator); + try destination_arena.cancel(lease); + try std.testing.expectEqual(@as(usize, 1), restored.entries.len); + try std.testing.expect(restored.entries[0].source_slot != saved.owned_blocks.?.entries[0].source_slot); + try std.testing.expectEqual(restored.clock, restored.entries[0].stamp); + try std.testing.expectEqualSlices(u8, saved.owned_blocks.?.payload, restored.payload); + try std.testing.expect(destination_arena.active == null); + try std.testing.expect(!destination.seq.invalidated); + } + try std.testing.expect(!fixtures.Fake.freed_pending); +} + +test "exl3: shared CUDA restored clocks preserve ages and reject overflow before metadata mutation" { + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const source = try arena.register(options.cap); + seed(arena, source, 0, 0, 0, 1, 0xa1); + const saving = try arena.acquire(source); + var image = try arena.saveOwnedBlocks(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + try arena.release(source, stream); + const destination = try arena.register(options.cap); + defer arena.release(destination, stream) catch unreachable; + const Boundary = struct { current: u32, source: u32, age: u32, expected: ?u32 }; + const limit = std.math.maxInt(u32); + const boundaries = [_]Boundary{ + .{ .current = 0, .source = 100, .age = 10, .expected = 111 }, + .{ .current = 200, .source = 100, .age = 10, .expected = 211 }, + .{ .current = 0, .source = limit, .age = 0, .expected = null }, + .{ .current = limit, .source = 1, .age = 0, .expected = null }, + .{ .current = 0, .source = limit - 1, .age = 1, .expected = null }, + .{ .current = 0, .source = limit - 1, .age = 0, .expected = limit }, + }; + const tables: [*]u8 = @ptrFromInt(arena.tables.ptr); + for (boundaries) |boundary| { + image.clock = boundary.source; + image.entries[0].stamp = boundary.source - boundary.age; + const resetting = try arena.acquire(destination); + try arena.clear(resetting, destination, stream); + try arena.complete(resetting, stream); + std.mem.writeInt(u32, tables[arena.layout.clock..][0..4], boundary.current, .little); + const before = std.hash.Wyhash.hash(0, tables[0..arena.metadata.len]); + const copies_before = Fake.copy_calls; + const restoring = try arena.acquire(destination); + if (boundary.expected) |expected| { + try arena.restoreOwnedBlocks(restoring, &image); + try arena.complete(restoring, stream); + try std.testing.expectEqual(expected, std.mem.readInt(u32, tables[arena.layout.clock..][0..4], .little)); + const checking = try arena.acquire(destination); + var restored = try arena.saveOwnedBlocks(std.testing.allocator, checking); + defer restored.deinit(std.testing.allocator); + try arena.cancel(checking); + try std.testing.expectEqual(@as(usize, 1), restored.entries.len); + try std.testing.expectEqual(@as(u64, boundary.age), restored.clock - restored.entries[0].stamp); + } else { + try std.testing.expectError(error.Overflow, arena.restoreOwnedBlocks(restoring, &image)); + try std.testing.expect(!arena.active.?.submitted); + try std.testing.expectEqual(before, std.hash.Wyhash.hash(0, tables[0..arena.metadata.len])); + try std.testing.expectEqual(copies_before, Fake.copy_calls); + try arena.cancel(restoring); + } + } +} + +test "exl3: shared CUDA repeated empty snapshots advance one epoch without absolute clock amplification" { + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const sequence = try arena.register(options.cap); + defer arena.release(sequence, stream) catch unreachable; + const tables: [*]u8 = @ptrFromInt(arena.tables.ptr); + std.mem.writeInt(u32, tables[arena.layout.clock..][0..4], 1_000_000_000, .little); + for (0..8) |iteration| { + const saving = try arena.acquire(sequence); + var image = try arena.saveOwnedBlocks(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + const restoring = try arena.acquire(sequence); + try arena.restoreOwnedBlocks(restoring, &image); + try arena.complete(restoring, stream); + try std.testing.expectEqual(@as(u32, @intCast(1_000_000_001 + iteration)), std.mem.readInt(u32, tables[arena.layout.clock..][0..4], .little)); + } +} diff --git a/zig/src/families/flashnext/cuda_exl3_state_binding.zig b/zig/src/families/flashnext/cuda_exl3_state_binding.zig new file mode 100644 index 0000000000..0b49ba199b --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_state_binding.zig @@ -0,0 +1,239 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const binding = @import("exl3_state_binding.zig"); + +pub const NodeType = *const fn (cuda.graph.Node, *c_uint) callconv(.c) cuda.abi.Result; + +pub fn nodeType(driver: *const cuda.Driver) ?NodeType { + var library = driver.lib; + return library.lookup(NodeType, "cuGraphNodeGetType"); +} + +pub const Backend = struct { + graph: cuda.graph.Graph, + exec: cuda.graph.Exec, + node_type: NodeType, + + pub fn init(graph: cuda.graph.Graph, exec: cuda.graph.Exec) ?Backend { + return initWithNodeType(graph, exec, nodeType(graph.d)); + } + + pub fn initWithNodeType(graph: cuda.graph.Graph, exec: cuda.graph.Exec, node_type: ?NodeType) ?Backend { + return .{ .graph = graph, .exec = exec, .node_type = node_type orelse return null }; + } + + pub fn nodes(self: Backend, template: usize) ![]usize { + if (template != @intFromPtr(self.graph.handle)) return error.UnsupportedBinding; + const count = try self.graph.nodeCount(); + const storage = try std.heap.c_allocator.alloc(cuda.graph.Node, count); + defer std.heap.c_allocator.free(storage); + const actual = try self.graph.nodes(storage); + if (actual.len != count) return error.UnsupportedBinding; + const identifiers = try std.heap.c_allocator.alloc(usize, count); + for (actual, identifiers) |node, *identifier| identifier.* = @intFromPtr(node); + return identifiers; + } + + pub fn freeNodes(_: Backend, nodes_list: []usize) void { + std.heap.c_allocator.free(nodes_list); + } + + pub fn kernel(self: Backend, node: usize) !?binding.Header { + var kind: c_uint = undefined; + try self.graph.d.check(self.node_type(@ptrFromInt(node), &kind), "cuGraphNodeGetType(state binding)"); + if (kind == 4) return error.UnsupportedBinding; + if (kind != 0) return null; + var params: cuda.abi.KernelNodeParams = undefined; + try self.graph.d.check(self.graph.d.api.cuGraphKernelNodeGetParams_v2(@ptrFromInt(node), ¶ms), "cuGraphKernelNodeGetParams(state binding)"); + return .{ + .function = @intFromPtr(params.func), + .geometry = .{ .grid = .{ params.grid_x, params.grid_y, params.grid_z }, .block = .{ params.block_x, params.block_y, params.block_z }, .shared = params.shared_bytes }, + .params = params.params, + .extra = params.extra != null, + .alternate_kernel = params.func == null, + }; + } + + pub fn set(self: Backend, node: usize, function: usize, geometry: binding.Geometry, parameters: binding.Parameters) !void { + var args: cuda.Args = .{}; + args.add(parameters.nonstate.gq); + args.add(parameters.nonstate.gk); + args.add(parameters.nonstate.gv); + args.add(parameters.nonstate.ggb); + args.add(parameters.state_in); + if (parameters.state_out != 0) args.add(parameters.state_out); + args.add(parameters.nonstate.rows); + args.add(parameters.nonstate.ys); + try self.exec.setKernel(@ptrFromInt(node), .{ .d = self.exec.d, .handle = @ptrFromInt(function) }, .{ + .grid = .{ .x = geometry.grid[0], .y = geometry.grid[1], .z = geometry.grid[2] }, + .block = .{ .x = geometry.block[0], .y = geometry.block[1], .z = geometry.block[2] }, + .shared = geometry.shared, + }, &args); + } +}; + +test "exl3: state binding adapter declines absent type query and propagates query errors" { + std.testing.refAllDecls(Backend); + const Fake = struct { + var getter_calls: usize = 0; + var kind: c_uint = 1; + var result: cuda.abi.Result = 0; + + fn nodeType(_: cuda.graph.Node, output: *c_uint) callconv(.c) cuda.abi.Result { + output.* = kind; + return result; + } + + fn get(_: cuda.graph.Node, _: *cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + getter_calls += 1; + return cuda.abi.error_not_ready; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphKernelNodeGetParams_v2 = Fake.get; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + try std.testing.expect(Backend.initWithNodeType(graph, exec, null) == null); + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + Fake.getter_calls = 0; + Fake.kind = 1; + Fake.result = 0; + try std.testing.expect(try backend.kernel(1) == null); + try std.testing.expectEqual(@as(usize, 0), Fake.getter_calls); + Fake.result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, backend.kernel(1)); + try std.testing.expectEqual(@as(usize, 0), Fake.getter_calls); + Fake.result = 0; + Fake.kind = 0; + try std.testing.expectError(error.NotReady, backend.kernel(1)); + try std.testing.expectEqual(@as(usize, 1), Fake.getter_calls); + Fake.kind = 4; + try std.testing.expectError(error.UnsupportedBinding, backend.kernel(1)); +} + +test "exl3: state binding adapter setters retain exact schema7 and schema8 types" { + const Fake = struct { + var output: bool = false; + var captured: binding.Parameters = undefined; + var geometry: binding.Geometry = undefined; + var function: usize = 0; + var node: usize = 0; + + fn value(comptime Value: type, params: [*]?*anyopaque, index: usize) Value { + var copied: Value = undefined; + @memcpy(std.mem.asBytes(&copied), @as([*]const u8, @ptrCast(params[index].?))[0..@sizeOf(Value)]); + return copied; + } + + fn set(_: cuda.abi.GraphExec, target: cuda.graph.Node, params: *const cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + const args = params.params.?; + captured = .{ + .nonstate = .{ + .gq = value(u64, args, 0), + .gk = value(u64, args, 1), + .gv = value(u64, args, 2), + .ggb = value(u64, args, 3), + .rows = value(c_int, args, if (output) 6 else 5), + .ys = value(u64, args, if (output) 7 else 6), + }, + .state_in = value(u64, args, 4), + .state_out = if (output) value(u64, args, 5) else 0, + }; + geometry = .{ .grid = .{ params.grid_x, params.grid_y, params.grid_z }, .block = .{ params.block_x, params.block_y, params.block_z }, .shared = params.shared_bytes }; + function = @intFromPtr(params.func); + node = @intFromPtr(target); + return 0; + } + + fn nodeType(_: cuda.graph.Node, _: *c_uint) callconv(.c) cuda.abi.Result { + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphExecKernelNodeSetParams_v2 = Fake.set; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + const geometry: binding.Geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }; + for (1..binding.max_width + 1) |width| for ([_]bool{ false, true }) |output| for ([_]bool{ false, true }) |flipped| { + Fake.output = output; + const parameters: binding.Parameters = .{ + .nonstate = .{ .gq = 0x100000001, .gk = 0x200000002, .gv = 0x300000003, .ggb = 0x400000004, .rows = @intCast(width), .ys = 0x500000005 }, + .state_in = if (flipped) 2000 else 1000, + .state_out = if (!output) 0 else if (flipped) 1000 else 2000, + }; + try backend.set(9, if (output) 11 else 10, geometry, parameters); + try std.testing.expectEqualDeep(parameters, Fake.captured); + try std.testing.expectEqualDeep(geometry, Fake.geometry); + try std.testing.expectEqual(@as(usize, if (output) 11 else 10), Fake.function); + try std.testing.expectEqual(@as(usize, 9), Fake.node); + }; +} + +test "exl3: state binding whitelisted CUfunction ignores ctx and kern fields" { + const Fake = struct { + var args: cuda.Args = .{}; + + fn nodeType(_: cuda.graph.Node, output: *c_uint) callconv(.c) cuda.abi.Result { + output.* = 0; + return 0; + } + + fn get(_: cuda.graph.Node, params: *cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + params.* = .{ + .func = @ptrFromInt(11), + .grid_x = 16, + .grid_y = 8, + .grid_z = 1, + .block_x = 128, + .block_y = 1, + .block_z = 1, + .shared_bytes = 0, + .params = args.pointers(), + .extra = null, + .kern = @ptrFromInt(20), + .ctx = @ptrFromInt(30), + }; + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphKernelNodeGetParams_v2 = Fake.get; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + Fake.args = .{}; + for ([_]u64{ 1, 2, 3, 4, 100, 200 }) |value| Fake.args.add(value); + Fake.args.add(@as(c_int, 2)); + Fake.args.add(@as(u64, 5)); + const header = (try backend.kernel(1)).?; + try std.testing.expectEqual(@as(usize, 11), header.function); + try std.testing.expect(!header.alternate_kernel); + const Source = struct { + backend: Backend, + identifiers: [1]usize = .{1}, + + pub fn nodes(self: *@This(), _: usize) ![]usize { + return &self.identifiers; + } + + pub fn freeNodes(_: *@This(), _: []usize) void {} + + pub fn kernel(self: *@This(), node: usize) !?binding.Header { + return self.backend.kernel(node); + } + }; + var request: binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + var state = try binding.Binding.init(42, request); + var source: Source = .{ .backend = backend }; + try state.load(&source); + try std.testing.expect(state.loaded); +} diff --git a/zig/src/families/flashnext/cuda_exl3_state_request_test.zig b/zig/src/families/flashnext/cuda_exl3_state_request_test.zig new file mode 100644 index 0000000000..7f84aaa8aa --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_state_request_test.zig @@ -0,0 +1,72 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const mtp = @import("cuda_mtp.zig"); +const binding = @import("exl3_state_binding.zig"); + +pub fn currentPass(comptime Key: type, comptime windowRequest: anytype) !void { + try std.testing.expectEqual(mtp.max_rows, binding.max_width); + var pass: batch.Pass(struct {}) = undefined; + var kernels: batch.Kernels = undefined; + kernels.chain = .{ .d = undefined, .handle = @ptrFromInt(10) }; + kernels.chain_to = .{ .d = undefined, .handle = @ptrFromInt(11) }; + var scratch: batch.Scratch = undefined; + scratch.gq = 1; + scratch.gk = 2; + scratch.gv = 3; + scratch.ggb = 4; + scratch.ys = 5; + var states: [wt.linear_n]cuda.DeviceBuffer = undefined; + var layers: [wt.layers_n]@import("cuda_prompt_load.zig").Layer = undefined; + var keep: batch.Keep = .{}; + for (&layers, 0..) |*layer, index| layer.linear = wt.linearLayer(index); + pass.lanes = &.{}; + pass.k = &kernels; + pass.s = &scratch; + pass.seq.state = &states; + pass.m.layers = &layers; + pass.hooks.keep = &keep; + var key: Key = .{ .kind = 2, .rows = 2, .scratch = 1, .scores = 2, .output = 3 }; + for (2..mtp.max_rows + 1) |width| { + key.rows = width; + for ([_]bool{ false, true }) |flipped| { + for (&states, 0..) |*state, index| state.ptr = (if (flipped) @as(u64, 2000) else 1000) + index; + for (&layers, 0..) |layer, index| { + if (layer.linear) keep.layers[index].state_out = (if (flipped) @as(u64, 1000) else 2000) + wt.linearIndex(index); + } + const maybe = windowRequest(key, &pass); + try std.testing.expect(maybe != null); + const request = maybe.?; + try std.testing.expectEqual(width, request.width); + try std.testing.expectEqual(@as(c_int, @intCast(width)), request.nonstate.rows); + try std.testing.expectEqual(wt.linear_n, request.count); + for (request.layers[0..request.count], 0..) |mapping, index| { + try std.testing.expectEqual(index, wt.linearIndex(mapping.layer)); + try std.testing.expectEqual(states[index].ptr, mapping.state_in); + try std.testing.expectEqual(keep.layers[mapping.layer].state_out, mapping.state_out); + } + } + } + for ([_]usize{ 0, 1, mtp.max_rows + 1, batch.att_rows + 1 }) |width| { + key.rows = width; + try std.testing.expect(windowRequest(key, &pass) == null); + } + key.rows = 3; + key.kind = 1; + try std.testing.expect(windowRequest(key, &pass) == null); + key.kind = 2; + const lanes = [_]batch.Lane{undefined}; + pass.lanes = &lanes; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.lanes = &.{}; + pass.seq.state = states[0 .. wt.linear_n - 1]; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.seq.state = &states; + keep.layers[0].state_out = states[0].ptr; + try std.testing.expect(windowRequest(key, &pass) == null); + keep.layers[0].state_out = 0; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.hooks.keep = null; + try std.testing.expect(windowRequest(key, &pass) == null); +} diff --git a/zig/src/families/flashnext/cuda_exl3_store.zig b/zig/src/families/flashnext/cuda_exl3_store.zig new file mode 100644 index 0000000000..c4a5b1a3d2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_store.zig @@ -0,0 +1,129 @@ +//! The EXL3 pack's tensors as the Flash Next loaders read them (cuda_prompt_wt.zig `Store`): +//! - the MTP head's mixer lives in `mtp_hyper_connection_mixer_patch.safetensors`, outside the index: its names join +//! the store's weight map at open; +//! - its dense weights (router, shared gate, every hyper-connection tensor, in_proj_a/b, the PLE projections, conv +//! and norm, the MTP pre-norm) are stored F16 where the loaders read BF16. Each is widened to BF16 when every value +//! survives the round trip exactly (an F16 that came from a BF16 checkpoint always does: 8 significant bits, BF16's +//! exponent range). A tensor with any inexact value is refused by name (error.F16NotBf16Exact), never rounded. +//! EXL3's own parts (`.trellis`, `.suh`, `.svh`, the codebook marker) keep their stored dtype: only `.weight` widens. + +const std = @import("std"); +const core = @import("core"); + +/// The mixer file the index does not name (the MTP head's `mtp.hyper_connection_mixer.*`). +pub const patch_file = "mtp_hyper_connection_mixer_patch.safetensors"; + +/// Adds every tensor of `dir`'s mixer patch to the index's weight map (the file absent: nothing). Names already in the +/// index keep their shard. The strings live in the parsed index's arena. +pub fn addPatch(io: std.Io, dir: []const u8, parsed: *std.json.Parsed(std.json.Value)) !void { + const a = parsed.arena.allocator(); + const wm = (parsed.value.object.getPtr("weight_map") orelse return); + if (wm.* != .object) return; + const path = try std.fs.path.join(a, &.{ dir, patch_file }); + var file = core.safetensors.File.open(a, io, path) catch |err| switch (err) { + error.FileNotFound => return, + else => return err, + }; + defer file.close(io); + var it = file.names.iterator(); + while (it.next()) |e| { + if (wm.object.contains(e.key_ptr.*)) continue; + try wm.object.put(a, try a.dupe(u8, e.key_ptr.*), .{ .string = patch_file }); + } +} + +/// Whether the loaders read `name` widened: a dense F16 `.weight` (an EXL3 part never ends in `.weight`). +pub fn widens(name: []const u8, t: core.safetensors.Tensor) bool { + return t.dtype == .f16 and std.mem.endsWith(u8, name, ".weight"); +} + +/// RNE to BF16 of an F32 (cuda_prompt_wt.toBf16's rounding, without its NaN quieting: a NaN is inexact here). +fn bf16Of(v: f32) u16 { + const bits: u32 = @bitCast(v); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// The F16 words as BF16, exactly, or error.F16NotBf16Exact at the first value BF16 cannot hold (NaN included). +pub fn exactBf16(out: []u16, f16_bytes: []const u8) !void { + if (f16_bytes.len != out.len * 2) return error.UnexpectedTensor; + for (out, 0..) |*o, i| { + const h: f16 = @bitCast(std.mem.readInt(u16, f16_bytes[2 * i ..][0..2], .little)); + const v: f32 = h; + const b = bf16Of(v); + const back: f32 = @bitCast(@as(u32, b) << 16); + if (back != v or std.math.isNan(v)) return error.F16NotBf16Exact; + o.* = b; + } +} + +/// A widened copy of `t` as a BF16 tensor of the same shape, in `owned` (the caller frees it). +pub fn widen(gpa: std.mem.Allocator, name: []const u8, t: core.safetensors.Tensor, owned: *std.ArrayList([]u16)) !core.safetensors.Tensor { + const count = t.bytes.len / 2; + const words = try gpa.alloc(u16, count); + errdefer gpa.free(words); + exactBf16(words, t.bytes) catch |err| { + if (err == error.F16NotBf16Exact) std.log.err("flashnext: {s} is F16 with a value BF16 cannot hold exactly", .{name}); + return err; + }; + try owned.append(gpa, words); + var out = t; + out.dtype = .bf16; + out.bytes = std.mem.sliceAsBytes(words); + return out; +} + +test "exl3: an F16 that came from BF16 widens bit for bit and one extra mantissa bit is refused" { + const values = [_]f32{ 0, -0.0, 1, -2.5, 0.0078125, 65280, -3.0517578125e-05, 5.9604644775390625e-08, std.math.inf(f32) }; + var bytes: [values.len * 2]u8 = undefined; + for (values, 0..) |v, i| { + const b: f32 = @bitCast(@as(u32, bf16Of(v)) << 16); + try std.testing.expectEqual(v, b); // every case is a BF16 value + std.mem.writeInt(u16, bytes[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(v))), .little); + } + var out: [values.len]u16 = undefined; + try exactBf16(&out, &bytes); + for (values, out) |v, o| try std.testing.expectEqual(v, @as(f32, @bitCast(@as(u32, o) << 16))); + // 1 + 2^-10 needs 11 significant bits: F16 holds it, BF16 does not + std.mem.writeInt(u16, bytes[0..2], @bitCast(@as(f16, 1.0009765625)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); + std.mem.writeInt(u16, bytes[0..2], @bitCast(std.math.nan(f16)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); +} + +test "exl3: only dense F16 weights widen, EXL3 parts keep their dtype" { + const t: core.safetensors.Tensor = .{ .dtype = .f16, .rank = 1, .shape = .{ 4, 1, 1, 1, 1 }, .bytes = &.{} }; + try std.testing.expect(widens("model.language_model.layers.0.mlp.gate.weight", t)); + try std.testing.expect(widens("mtp.hyper_connection_mixer.hc_norm.weight", t)); + try std.testing.expect(!widens("model.language_model.layers.0.self_attn.q_proj.suh", t)); + try std.testing.expect(!widens("lm_head.svh", t)); + var b = t; + b.dtype = .bf16; + try std.testing.expect(!widens("model.language_model.layers.0.linear_attn.norm.weight", b)); +} + +test "exl3: the mixer patch's names join the weight map, an indexed name keeps its shard" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = "{\"mtp.hyper_connection_mixer.hc_norm.weight\":{\"dtype\":\"F16\",\"shape\":[1],\"data_offsets\":[0,2]},\"x\":{\"dtype\":\"U8\",\"shape\":[1],\"data_offsets\":[2,3]}}"; + var data: [8 + header.len + 3]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = patch_file, .data = &data }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var parsed = try std.json.parseFromSlice(std.json.Value, a, "{\"weight_map\":{\"x\":\"model-00001-of-00007.safetensors\"}}", .{}); + defer parsed.deinit(); + try addPatch(io, dir, &parsed); + const wm = parsed.value.object.get("weight_map").?.object; + try std.testing.expectEqualStrings(patch_file, wm.get("mtp.hyper_connection_mixer.hc_norm.weight").?.string); + try std.testing.expectEqualStrings("model-00001-of-00007.safetensors", wm.get("x").?.string); + // a folder without the patch file opens as before + var other = std.testing.tmpDir(.{}); + defer other.cleanup(); + const empty = try other.dir.realPathFileAlloc(io, ".", a); + defer a.free(empty); + try addPatch(io, empty, &parsed); +} diff --git a/zig/src/families/flashnext/cuda_exl3_telemetry.zig b/zig/src/families/flashnext/cuda_exl3_telemetry.zig new file mode 100644 index 0000000000..639b6af553 --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_telemetry.zig @@ -0,0 +1,119 @@ +const std = @import("std"); +const prompt = @import("cuda_prompt.zig"); +const load = @import("cuda_prompt_load.zig"); +const pool = @import("cuda_exl3_pool.zig"); +const pass = @import("cuda_exl3_pass.zig"); +const dplan = @import("cuda_exl3_dplan.zig"); +const prefill = @import("cuda_exl3_prefill.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); +const weights = @import("cuda_prompt_wt.zig"); +pub const counters = @import("exl3_telemetry.zig"); +pub const Completion = enum { all_rank_consumers_joined }; + +pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine, mtp: ?*const load.Layer, sequences: []const *prompt.Seq, rank: counters.Rank, completion: Completion) !counters.Snapshot { + _ = completion; + try rank.validate(); + if (engine.shared_arena) |arena| if (arena.active != null) return error.TelemetryConsumersActive; + var hash = std.hash.Wyhash.init(@intFromPtr(engine)); + var totals: counters.Experts = .{}; + var attached: [weights.layers_n + 1]*@import("cuda_exl3_load.zig").Experts = undefined; + const attached_count = pass.offloadedPools(engine.layers[0..engine.n_layers], mtp, &attached); + var layers: []counters.LayerSample = if (engine.telemetry.layers_enabled) try allocator.alloc(counters.LayerSample, attached_count) else @constCast(&.{}); + errdefer allocator.free(layers); + var seen: [weights.layers_n + 1]*pool.Pool = undefined; + var seen_count: usize = 0; + for (attached[0..attached_count]) |experts| { + const resident = experts.pool orelse continue; + var duplicate = false; + for (seen[0..seen_count]) |previous| if (previous == resident) { + duplicate = true; + break; + }; + if (duplicate) continue; + if (resident.telemetry_id == 0) return error.TelemetryOwnerChanged; + if (!resident.telemetry_valid) return error.TelemetryCounterRegression; + seen[seen_count] = resident; + seen_count += 1; + hash.update(std.mem.asBytes(&resident.telemetry_id)); + const device = if (dplan.current()) |planner| try planner.samplePoolAfterConsumers(resident) else resident.telemetry_device; + var minimum: u64 = std.math.maxInt(u64); + var maximum: u64 = 0; + for (resident.src.routed, resident.k2) |source, widths| { + if (widths[0] == 0) continue; + const bytes = try std.math.add(u64, try std.math.add(u64, source[0].len, source[1].len), source[2].len); + minimum = @min(minimum, bytes); + maximum = @max(maximum, bytes); + } + if (minimum == std.math.maxInt(u64)) { + if (device.copies != 0) return error.TelemetryInconsistentCounters; + minimum = 0; + } + const host = try counters.difference(counters.Device, resident.telemetry_device, .{ .plans = resident.stats.rounds, .hit_picks = resident.stats.hit_slots, .miss_picks = resident.stats.miss_slots, .copies = resident.stats.misses, .evictions = resident.stats.evictions, .inplace_experts = resident.telemetry_device.inplace_experts, .overflow_plans = resident.telemetry_device.overflow_plans }); + const contribution: counters.Experts = .{ + .host_plans = host.plans, + .host_hit_picks = host.hit_picks, + .host_miss_picks = host.miss_picks, + .host_mandatory_copies = host.copies, + .host_prefetch_copies = resident.stats.prefetches, + .pool_dma_copies = resident.lines.stats.copies, + .pool_dma_requested_bytes = resident.lines.stats.bytes, + .device_plans = device.plans, + .device_hit_picks = device.hit_picks, + .device_miss_picks = device.miss_picks, + .device_copies = device.copies, + .device_inplace_experts = device.inplace_experts, + .device_overflow_plans = device.overflow_plans, + .mapped_copy_payload_min_bytes = try std.math.mul(u64, device.copies, minimum), + .mapped_copy_payload_max_bytes = try std.math.mul(u64, device.copies, maximum), + }; + try counters.add(counters.Experts, &totals, contribution); + if (engine.telemetry.layers_enabled) { + var layer_sample: counters.LayerSample = .{ .owner = resident.telemetry_id, .kind = .mtp, .layer = 0, .lines = resident.opt.lines, .counts = counters.LayerCounts.fromExperts(contribution) }; + for (engine.layers[0..engine.n_layers], 0..) |layer, decoder_index| if (layer.exl3) |decoder| { + if (decoder.pool == resident) { + layer_sample.kind = .decoder; + layer_sample.layer = decoder_index; + break; + } + }; + layers[seen_count - 1] = layer_sample; + } + } + var count: usize = 0; + for (sequences, 0..) |sequence, index| { + if (sequence.invalidated) return error.TelemetryInvalidSequence; + for (sequences[0..index]) |previous| if (previous == sequence) return error.TelemetryDuplicateOwner; + if (sequence.paged != null) count = try std.math.add(usize, count, weights.attn_n); + } + const samples = try allocator.alloc(counters.KvSample, count); + errdefer allocator.free(samples); + var index: usize = 0; + for (sequences) |sequence| if (sequence.paged) |paged| { + if (paged.telemetry_id == 0) return error.TelemetryOwnerChanged; + const clock = try paged.sampleClockAfterConsumers(); + for (0..weights.attn_n) |layer| { + const status = try paged.stats(layer); + samples[index] = .{ .owner = paged.telemetry_id, .revision = paged.telemetry_revision, .layer = layer, .empty = sequence.fed == 0, .epoch = status.epoch, .clock = clock, .requests = @intCast(status.needed), .misses = @intCast(status.misses), .calls = @intCast(status.calls), .overflow = status.overflow, .max_blocks = @min(paged.cfg.slots, paged.nblocks), .block_bytes = weights.index_ratio * 2 * pager.row_bytes, .rebases = status.rebases, .written = status.written_rows }; + index += 1; + } + }; + const ring = prefill.ringTelemetry(engine); + if (ring) |observed| { + totals.ring_dma_copies = observed.copies; + totals.ring_dma_requested_bytes = observed.dma_requested_bytes; + totals.ring_unused_copies = observed.unused_copies; + totals.ring_hit_picks = observed.hit_picks; + totals.ring_miss_picks = observed.miss_picks; + totals.ring_hit_expert_requests = observed.hit_expert_requests; + totals.ring_miss_expert_requests = observed.miss_expert_requests; + if (observed.attribution) |by_source| @import("cuda_exl3_prefill_attribution.zig").publish(by_source, &totals); + } + const sampled_layers = if (engine.telemetry.layers_enabled) try allocator.realloc(layers, seen_count) else layers; + return .{ .rank = rank, .scope = hash.final(), .ring_owner = if (ring) |observed| observed.owner else 0, .experts = totals, .kv = samples, .expert_pools = seen_count, .layers = sampled_layers }; +} + +test "exl3: cache telemetry pure host contracts" { + _ = counters; + _ = @import("cuda_exl3_telemetry_test.zig"); + _ = @import("cuda_prompt_telemetry.zig"); +} diff --git a/zig/src/families/flashnext/cuda_exl3_telemetry_test.zig b/zig/src/families/flashnext/cuda_exl3_telemetry_test.zig new file mode 100644 index 0000000000..2d3e4fd14f --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_telemetry_test.zig @@ -0,0 +1,206 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const planner = @import("cuda_exl3_dplan.zig"); +const pool = @import("cuda_exl3_pool.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); +const telemetry = @import("exl3_telemetry.zig"); +const Fake = @import("cuda_exl3_shared_kv_test.zig").Fake; + +test "exl3: replay device sampling reads executed counters without pulling ownership or synchronizing" { + var driver = Fake.driver(); + var resident: pool.Pool = undefined; + resident.lines.pinned_bank = null; + resident.telemetry_device = .{ .plans = 5, .hit_picks = 20, .copies = 3 }; + resident.telemetry_valid = true; + var pools = [_]*pool.Pool{&resident}; + var state: [planner.stateWords(1)]u32 = @splat(0); + const executed: planner.Misc = .{ .now = 104, .rounds = 4, .hit_slots = 30, .miss_slots = 10, .copies = 8, .evictions = 2, .overflow = 0, .inplace = 1, .recent_cursor = 9, .recent_filled = 8 }; + const offset = planner.miscOffset(1) * 4; + @memcpy(std.mem.sliceAsBytes(&state)[offset..][0..@sizeOf(planner.Misc)], std.mem.asBytes(&executed)); + var seen = [_]planner.Misc{std.mem.zeroes(planner.Misc)}; + var device: planner.DevPlan = undefined; + device.owner = .device; + device.pools = &pools; + var pool_states = [_]planner.PoolState{try planner.PoolState.init(1, 0)}; + device.pool_states = &pool_states; + device.max_rows = 1; + device.state = .{ .d = &driver, .ptr = @intFromPtr(&state), .len = @sizeOf(@TypeOf(state)), .borrowed = true }; + device.seen = &seen; + const syncs = Fake.sync_calls; + const first = try device.samplePoolAfterConsumers(&resident); + const again = try device.samplePoolAfterConsumers(&resident); + try std.testing.expectEqualDeep(first, again); + try std.testing.expectEqual(@as(u64, 9), first.plans); + try std.testing.expectEqual(@as(u64, 50), first.hit_picks); + try std.testing.expectEqual(@as(u64, 11), first.copies); + try std.testing.expectEqual(@as(u64, 1), first.inplace_experts); + try std.testing.expectEqual(planner.Owner.device, device.owner); + try std.testing.expectEqual(@as(u32, 0), seen[0].copies); + try std.testing.expectEqual(syncs, Fake.sync_calls); + resident.telemetry_device = first; + seen[0] = executed; + device.owner = .host; + try std.testing.expectEqualDeep(first, try device.samplePoolAfterConsumers(&resident)); + device.owner = .device; + try std.testing.expectEqualDeep(first, try device.samplePoolAfterConsumers(&resident)); +} + +test "exl3: ragged planner telemetry and MTP lookup use each pool's state bounds" { + var driver = Fake.driver(); + const counts = [_]usize{ 10, 37, 12 }; + var residents: [counts.len]pool.Pool = undefined; + var pools: [counts.len]*pool.Pool = undefined; + var pool_states: [counts.len]planner.PoolState = undefined; + var seen: [counts.len]planner.Misc = @splat(std.mem.zeroes(planner.Misc)); + var words: usize = 0; + for (counts, 0..) |lines, index| { + residents[index].opt = .{ .layer = if (index == 2) 0 else @intCast(index), .lines = lines }; + residents[index].telemetry_valid = true; + residents[index].telemetry_device = .{}; + pools[index] = &residents[index]; + pool_states[index] = try planner.PoolState.init(lines, words); + words = pool_states[index].endWords(); + } + const state = try std.testing.allocator.alloc(u32, words); + defer std.testing.allocator.free(state); + const mirror = try std.testing.allocator.alignedAlloc(u8, .@"16", words * 4); + defer std.testing.allocator.free(mirror); + @memset(state, 0); + for (pool_states, 0..) |entry, index| { + const misc: *planner.Misc = @ptrCast(entry.slice(state)[planner.miscOffset(entry.lines)..][0..10]); + misc.* = std.mem.zeroes(planner.Misc); + misc.rounds = @intCast(index + 1); + misc.hit_slots = @intCast(10 * (index + 1)); + misc.copies = @intCast(index + 2); + } + var device: planner.DevPlan = undefined; + device.d = &driver; + device.owner = .device; + device.pools = &pools; + device.pool_states = &pool_states; + device.max_rows = 16; + device.state = .{ .d = &driver, .ptr = @intFromPtr(state.ptr), .len = state.len * 4, .borrowed = true }; + device.mirror = .{ .d = &driver, .bytes = mirror }; + device.stream = null; + device.seen = &seen; + const syncs = Fake.sync_calls; + for (&residents, 0..) |*resident, index| { + try std.testing.expectEqual(@as(?usize, index), device.indexOf(resident)); + const sample = try device.samplePoolAfterConsumers(resident); + try std.testing.expectEqual(@as(u64, index + 1), sample.plans); + try std.testing.expectEqual(@as(u64, 10 * (index + 1)), sample.hit_picks); + try std.testing.expectEqual(@as(u64, index + 2), sample.copies); + } + var outside: pool.Pool = undefined; + outside.opt = .{ .layer = 0, .lines = 18 }; + outside.telemetry_valid = true; + outside.telemetry_device = .{ .plans = 7 }; + try std.testing.expectEqual(@as(?usize, null), device.indexOf(&outside)); + try std.testing.expectEqualDeep(outside.telemetry_device, try device.samplePoolAfterConsumers(&outside)); + const totals = try device.peek(); + try std.testing.expectEqual(@as(u32, 6), totals.rounds); + try std.testing.expectEqual(@as(u32, 60), totals.hit_slots); + try std.testing.expectEqual(@as(u32, 9), totals.copies); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(state), mirror); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} + +test "exl3: raw KV counters preserve unsigned high bits and restore advances telemetry revision" { + var driver = Fake.driver(); + var tables: [12 * (2 + 8) + 2]u32 = @splat(0); + tables[24 + 3] = 0x80000001; + tables[24 + 4] = 0xfffffff0; + tables[24 + 5] = 0x90000002; + tables[24] = 0x80000010; + var paged: pager.Paged = undefined; + paged.d = &driver; + paged.shared_arena = null; + paged.cfg = .{ .slots = 1 }; + paged.wr = 4096; + paged.nblocks = 1; + paged.tabs = .{ .d = &driver, .ptr = @intFromPtr(&tables), .len = @sizeOf(@TypeOf(tables)), .borrowed = true }; + paged.rebases = @splat(0); + paged.written = @splat(0); + const stats = try paged.stats(0); + try std.testing.expectEqual(@as(u64, 0x80000001), stats.misses); + try std.testing.expectEqual(@as(u64, 0xfffffff0), stats.needed); + try std.testing.expectEqual(@as(u64, 0x90000002), stats.calls); + try std.testing.expectEqual(@as(u32, 0x80000010), try paged.sampleClockAfterConsumers()); + paged.telemetry_revision = 7; + paged.publishSnapshot(.{ .layout = undefined, .win0 = @splat(0), .rebases = @splat(0), .written = @splat(0) }); + try std.testing.expectEqual(@as(u64, 8), paged.telemetry_revision); +} + +test "exl3: shared telemetry samples the arena clock without completion waits" { + const shared = @import("cuda_exl3_shared_kv.zig"); + var driver = Fake.driver(); + const arena = try shared.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 8192, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }); + defer arena.destroy() catch unreachable; + const clock: *u32 = @ptrFromInt(arena.tables.ptr + arena.layout.clock); + clock.* = 524288; + var paged: pager.Paged = undefined; + paged.d = &driver; + paged.shared_arena = arena; + const syncs = Fake.sync_calls; + try std.testing.expectEqual(@as(u32, 524288), try paged.sampleClockAfterConsumers()); + clock.* += 1; + try std.testing.expectEqual(@as(u32, 524289), try paged.sampleClockAfterConsumers()); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} + +test "exl3: device counter widening handles one u32 rollover before summation" { + var previous = std.mem.zeroes(planner.Misc); + previous.hit_slots = std.math.maxInt(u32) - 2; + var current = previous; + current.hit_slots = 4; + const delta = telemetry.Device.pending(current, previous); + try std.testing.expectEqual(@as(u64, 7), delta.hit_picks); +} + +test "exl3: ring demand counts routed pick appearances separately from per-chunk expert requests and speculative copies" { + var stats: @import("cuda_exl3_prefill_ring.zig").Stats = .{}; + stats.observeDemand(true, 17); + stats.observeDemand(false, 31); + stats.observeDemand(false, 0); + try std.testing.expectEqual(@as(u64, 17), stats.hit_picks); + try std.testing.expectEqual(@as(u64, 31), stats.miss_picks); + try std.testing.expectEqual(@as(u64, 1), stats.hit_expert_requests); + try std.testing.expectEqual(@as(u64, 1), stats.miss_expert_requests); + try std.testing.expectEqual(@as(u64, 0), stats.copies); + try std.testing.expectEqual(@as(u64, 0), stats.dma_requested_bytes); +} + +test "exl3: aggregate telemetry finds attached pools deduplicates MTP aliases and separates committed device plans" { + var resident: pool.Pool = undefined; + resident.lines.pinned_bank = null; + resident.telemetry_id = 991; + resident.telemetry_valid = true; + resident.telemetry_device = .{ .plans = 5, .hit_picks = 20, .miss_picks = 10, .copies = 3, .inplace_experts = 1 }; + resident.stats = .{ .rounds = 7, .hit_slots = 25, .miss_slots = 13, .misses = 5 }; + resident.lines.stats = .{ .copies = 2, .bytes = 96 }; + const payload: [16]u8 = @splat(0); + var sources = [_][3][]const u8{.{ &payload, &payload, &payload }}; + var widths = [_][3]i32{.{ 1, 1, 1 }}; + resident.src = .{ .routed = &sources }; + resident.k2 = &widths; + var experts: @import("cuda_exl3_load.zig").Experts = undefined; + experts.pool = &resident; + var engine: @import("cuda_prompt.zig").Engine = undefined; + engine.shared_arena = null; + engine.telemetry = .{}; + engine.n_layers = 1; + engine.layers[0].exl3 = &experts; + const collector = @import("cuda_exl3_telemetry.zig"); + var snapshot = try collector.sampleAfterConsumers(std.testing.allocator, &engine, &engine.layers[0], &.{}, .{ .rank = 0, .world = 2 }, .all_rank_consumers_joined); + defer snapshot.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 1), snapshot.expert_pools); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_plans); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.host_hit_picks); + try std.testing.expectEqual(@as(u64, 3), snapshot.experts.host_miss_picks); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_mandatory_copies); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.device_plans); + try std.testing.expectEqual(@as(u64, 1), snapshot.experts.device_inplace_experts); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_min_bytes); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_max_bytes); + try std.testing.expectEqual(@as(u64, 96), snapshot.experts.pool_dma_requested_bytes); +} diff --git a/zig/src/families/flashnext/cuda_exl3_tp.zig b/zig/src/families/flashnext/cuda_exl3_tp.zig new file mode 100644 index 0000000000..96565971bd --- /dev/null +++ b/zig/src/families/flashnext/cuda_exl3_tp.zig @@ -0,0 +1,533 @@ +//! Two-rank expert parallel for an EXL3 Flash Next checkpoint (plan step 10 of docs/plans/exl3-on-upstream-1.1.0.md), +//! on upstream's rank engine (cuda_rank_engine.zig: one process a GPU, rank 1 replays rank 0's operations) and its +//! NCCL `Collective`. +//! +//! What a rank holds and runs: +//! - routed experts by id (exl3_tp.zig `routed_set`): rank 0 ids 0..255, rank 1 ids 256..511. An EXL3 expert cannot be +//! cut along its 640-wide intermediate (320 is not a multiple of the 128-wide Hadamard block), which is why +//! upstream's ModelOpt `Cut` (intermediate halves) does not apply; +//! - the shared expert (id 512) on rank 1, so each rank owns ONE contiguous id range: rank 0 [0, 256), rank 1 +//! [256, 513). Both ranks run the same kernel sequence (no rank branch in the pass: the other rank's slots are +//! masked, not skipped by a different code path), so rank 1 replays rank 0's operations and a captured round graph +//! has the same nodes on both. Upstream puts its shared expert on rank 0 behind an `if (rank == 0)` branch; here +//! rank 0 already carries the MTP head and the sampled logits, so the shared expert's compute and its 2.9 MB a layer +//! on rank 1 balance that; +//! - every expert table keeps its 513 global ids (the offload pool's tables, `shared_id` 512 and its policy are +//! unchanged); the stacked suh/svh scales hold only the rank's rows, their base pointers shifted by `lo` rows so the +//! kernels' `scale[e * width]` still indexes by global id; +//! - attention, GDN, hyper-connections, the router and KV are replicated, as upstream's ModelOpt tp does; +//! - offloaded (the default): each rank's routed trellis in its own pinned bank (its 256 experts, about 21 GiB a +//! process) and its own pool per layer (cuda_exl3_pass.attach sizes it from the rank's free memory), planning on the +//! rank's masked picks: the same lines cover half the experts, so a rank's hit rate is higher than tp=1's; +//! - the lm_head split by vocabulary (cuda_exl3_load.shardHead, default at tp=2): 124,160 columns a rank. +//! +//! The MoE branch (`moe`, the hook in cuda_exl3_pass.moe): the router and top-k run on both ranks (same rows, same +//! picks); `fn_exl3_tp_mask` turns the other rank's slots into pick -1, weight 0; `Experts.run` combines the rank's +//! slots in slot order (fp32 fmaf, as the write-back does) into [rows, 2560] fp32; ONE `Collective.sum` over that +//! (not upstream's per-slot [rows, 11, 2560]: 11 times less traffic); `fn_exl3_tp_place` writes the sum into slot 0 of +//! the pass's `y` with weights {1, 0, ...}, so the upstream write-back is unchanged and writes the sum's bf16. +//! tp=2 differs from tp=1 only in the fp32 grouping of the 11 terms: (rank 0's chain) + (rank 1's chain). +//! With two ranks the all-reduce is a single fp32 addition a + b, which is commutative, so both ranks get the same +//! bits whatever NCCL's algorithm (checked on the device by flashnext-exl3-tp). +//! +//! Masked slots and stale values: `Experts.run` leaves a skipped slot's per-slot output as it was in the scratch and +//! multiplies it by the slot's weight. The weight is 0 and the scratch is zeroed when allocated and only ever holds +//! finite expert outputs, so a skipped slot adds exactly 0. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_exl3_load.zig"); +const memory = @import("cuda_exl3_memory.zig"); +const projection = @import("cuda_projection.zig"); +const cuts = @import("exl3_tp.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const width = wt.moe_w; +/// Table entries: 512 routed experts and the shared one as id 512. +const table_n = wt.experts_n + 1; +/// The rank that holds the shared expert (id 512). +pub const shared_rank: usize = 1; +/// Rows an `Experts.run` call takes (cuda_exl3_pass.zig's block). +pub const moe_rows: usize = @import("cuda_exl3_pass.zig").moe_rows; + +/// The ids [lo, hi) rank `rank` of `world` holds. +pub const Rank = struct { + world: usize, + rank: usize, + lo: usize, + hi: usize, + + pub fn init(world: usize, rank: usize) !Rank { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (world == 1) return .{ .world = 1, .rank = 0, .lo = 0, .hi = table_n }; + const r = cuts.range(.routed_set, rank, world) catch return error.BadRank; + return .{ .world = world, .rank = rank, .lo = r.lo, .hi = if (rank == shared_rank) r.hi + 1 else r.hi }; + } + + pub fn owns(r: Rank, id: usize) bool { + return id >= r.lo and id < r.hi; + } +}; + +// ---- loading ----------------------------------------------------------------------------------------------------- + +fn up(x: *load.Experts, d: *const cuda.Driver, bytes: []const u8) !u64 { + try x.bufs.ensureUnusedCapacity(x.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + x.bufs.appendAssumeCapacity(b); + return b.ptr; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +fn expertName(b: []u8, mlp: []const u8, e: usize, proj: []const u8) ![]const u8 { + if (e == wt.experts_n) return std.fmt.bufPrint(b, "{s}.shared_expert.{s}", .{ mlp, proj }); + return std.fmt.bufPrint(b, "{s}.experts.{d}.{s}", .{ mlp, e, proj }); +} + +/// cuda_exl3_load.experts at world 2: the rank's experts, the 513-entry tables (other ids zero), the scales' rows +/// lo..hi with their bases shifted to global ids. Offloaded (the default, cuda_exl3_load.offloaded): only the rank's +/// own trellis words go to a pinned host bank (about half the tp=1 bank a process), the other rank's ids stay empty +/// in `routed` (the pool never stages them: their picks are masked to -1 before it plans), and the pool attached by +/// cuda_exl3_pass.attach owns the tables. Resident (TF_FLASHNEXT_EXL3_ROUTED=device): the rank's experts on the device. +pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, mlp: []const u8, world: usize, rank: usize) !*load.Experts { + const r = try Rank.init(world, rank); + const offload = load.offloaded(); + const x = try gpa.create(load.Experts); + x.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .layer = undefined, .tp = r }; + errdefer x.destroy(); + for (&x.trellis, &x.halves) |*t, *h| { + t.* = try gpa.alloc(u64, table_n); + @memset(t.*, 0); + h.* = try gpa.alloc(i32, table_n); + @memset(h.*, 0); + } + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + const ks = [3]usize{ dims, dims, width }; + const ns = [3]usize{ width, width, dims }; + var scales: [6]std.ArrayList(u8) = @splat(.empty); + defer for (&scales) |*l| l.deinit(gpa); + var buf: [256]u8 = undefined; + var at: usize = 0; + if (offload) { + // the bank's size from the headers of the rank's own experts, then one pinned region for the layer + var total: usize = 0; + for (r.lo..r.hi) |e| { + for (projs) |proj| total += align256((try load.mat(store, try expertName(&buf, mlp, e, proj))).trellis.len); + store.release(); + } + x.bank = try cuda.HostBuffer.alloc(d, total); + x.routed = try gpa.alloc([3][]const u8, wt.experts_n); + for (x.routed) |*e| e.* = .{ &.{}, &.{}, &.{} }; + } + const t0 = std.Io.Clock.awake.now(store.io).toNanoseconds(); + for (r.lo..r.hi) |e| { + for (projs, 0..) |proj, j| { + const m = try load.mat(store, try expertName(&buf, mlp, e, proj)); + if (m.k != ks[j] or m.n != ns[j]) return error.UnexpectedTensor; + if (x.bank) |bank| { + const dst = bank.bytes[at..][0..m.trellis.len]; + @memcpy(dst, m.trellis); + at += align256(m.trellis.len); + if (e == wt.experts_n) x.shared[j] = dst else x.routed[e][j] = dst; + } else x.trellis[j][e] = try up(x, d, m.trellis); + x.halves[j][e] = @intCast(m.halves); + try scales[2 * j].appendSlice(gpa, m.suh); + try scales[2 * j + 1].appendSlice(gpa, m.svh); + } + store.release(); + } + x.bank_ns = @intCast(std.Io.Clock.awake.now(store.io).toNanoseconds() - t0); + // a scale row's fp16 count: suh_g D, svh_g I, suh_u D, svh_u I, suh_d I, svh_d D + const rows = [6]usize{ dims, width, dims, width, width, dims }; + var base: [6]u64 = undefined; + for (&base, scales, rows) |*b, l, n| b.* = try up(x, d, l.items) - r.lo * n * 2; + x.layer = .{ + // offloaded: the pool's tables replace these at attach (0 until then; a pass refuses an unattached layer) + .gate = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[0])), + .up = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[1])), + .down = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[2])), + .gate_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[0])), + .up_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[1])), + .down_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[2])), + .suh_g = base[0], + .svh_g = base[1], + .suh_u = base[2], + .svh_u = base[3], + .suh_d = base[4], + .svh_d = base[5], + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + return x; +} + +// ---- the pass ---------------------------------------------------------------------------------------------------- + +/// The engine's expert-parallel scratch and kernels (held by the shared cuda_exl3_load.Ops, released with it). +pub const Scratch = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + mask_fn: cuda.Function, + place_fn: cuda.Function, + /// Masked picks (int32) and weights (fp32) [rows, 11]; the rank's combined branch, then the ranks' sum [rows, D]. + picks: ?cuda.DeviceBuffer = null, + wts: ?cuda.DeviceBuffer = null, + sum: ?cuda.DeviceBuffer = null, + rows: usize = 0, + /// Experts.run's scratch for `moe_cap` rows, zeroed when allocated (a skipped slot reads its stale, finite y). + moe: ?cuda.DeviceBuffer = null, + moe_cap: usize = 0, + + fn create(gpa: std.mem.Allocator, d: *const cuda.Driver) !*Scratch { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var m = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_tp); + errdefer m.unload(); + const t = try gpa.create(Scratch); + t.* = .{ .gpa = gpa, .d = d, .mod = m, .mask_fn = try m.function("fn_exl3_tp_mask"), .place_fn = try m.function("fn_exl3_tp_place") }; + return t; + } + + pub fn destroy(t: *Scratch) void { + inline for (.{ "picks", "wts", "sum", "moe" }) |f| if (@field(t, f)) |*b| b.free(); + t.mod.unload(); + t.gpa.destroy(t); + } + + /// Room for `rows` rows; a growth waits for `s` (an earlier call may still read the old buffers). + fn ensure(t: *Scratch, s: cuda.Stream, l: cuda.exl3_experts.Layer, rows: usize) !void { + if (rows > t.rows) { + try s.synchronize(); + inline for (.{ "picks", "wts", "sum" }) |f| if (@field(t, f)) |*b| { + b.free(); + @field(t, f) = null; + }; + t.rows = 0; + t.picks = try cuda.DeviceBuffer.alloc(t.d, rows * slots * 4); + t.wts = try cuda.DeviceBuffer.alloc(t.d, rows * slots * 4); + t.sum = try cuda.DeviceBuffer.alloc(t.d, rows * dims * 4); + t.rows = rows; + } + const cap = @min(rows, moe_rows); + if (cap > t.moe_cap) { + try s.synchronize(); + if (t.moe) |*b| b.free(); + t.moe = null; + t.moe_cap = 0; + const bytes = try cuda.exl3_experts.scratchBytes(l, cap, slots); + t.moe = try cuda.DeviceBuffer.alloc(t.d, bytes); + try s.d.check(s.d.api.cuMemsetD8Async(t.moe.?.ptr, 0, bytes, s.handle), "cuMemsetD8Async"); + t.moe_cap = cap; + } + } + + /// pick -> -1 and weight -> 0 outside [lo, hi), for `n` slots. + pub fn mask(t: *const Scratch, s: cuda.Stream, picks: u64, wts: u64, mpicks: u64, mwts: u64, n: usize, r: Rank) !void { + var a: cuda.Args = .{}; + a.add(picks); + a.add(wts); + a.add(mpicks); + a.add(mwts); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(r.lo))); + a.add(@as(c_int, @intCast(r.hi))); + try cuda.launch.launch(t.mask_fn, .{ .grid = .{ .x = @intCast((n + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + } + + /// The summed branch [rows, D] into y [rows, 11, D] slot 0 (others 0) and weights {1, 0, ...}. + pub fn place(t: *const Scratch, s: cuda.Stream, sum: u64, y: u64, wts: u64, rows: usize) !void { + var a: cuda.Args = .{}; + a.add(sum); + a.add(y); + a.add(wts); + a.add(@as(c_int, @intCast(dims))); + a.add(@as(c_int, @intCast(slots))); + try cuda.launch.launch(t.place_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(slots) }, .block = .{ .x = 256 } }, s, &a); + } +}; + +/// The Ops' expert-parallel scratch, made on first use. +pub fn scratchOf(ops: *load.Ops) !*Scratch { + if (ops.tp == null) ops.tp = try Scratch.create(ops.gpa, ops.d); + return ops.tp.?; +} + +/// Every expert-parallel scratch a pass of up to `rows` rows reads, allocated now (cuda_exl3_pass.attach, before the +/// pools size themselves from free memory), so no round grows a buffer (a growth synchronises the stream and moves an +/// address a captured graph would hold): the masked picks, weights and the ranks' sum for `rows` rows, Experts.run's +/// scratch for a block (zeroed), and the pinned picks of the pool's A/B cut. +pub fn reserve(ops: *load.Ops, s: cuda.Stream, l: cuda.exl3_experts.Layer, rows: usize) !void { + const t = try scratchOf(ops); + try t.ensure(s, l, rows); + if (ops.picks == null) ops.picks = try cuda.HostBuffer.alloc(ops.d, moe_rows * slots * 4); + try s.synchronize(); +} + +/// The expert-parallel buffers' addresses and sizes (cuda_exl3_load.Ops.fingerprint's tail). +pub fn fingerprint(t: *const Scratch) [4][2]u64 { + var f: [4][2]u64 = @splat(.{ 0, 0 }); + for ([4]?cuda.DeviceBuffer{ t.picks, t.wts, t.sum, t.moe }, &f) |b, *e| if (b) |x| { + e.* = .{ x.ptr, x.len }; + }; + return f; +} + +/// Device bytes the expert-parallel scratch holds. +pub fn reservedBytes(t: *const Scratch) usize { + var n: usize = 0; + for ([4]?cuda.DeviceBuffer{ t.picks, t.wts, t.sum, t.moe }) |b| if (b) |x| { + n += x.len; + }; + return n; +} + +/// The rank's half of the MoE branch, then the ranks' sum: cuda_exl3_pass.moe for a layer whose experts carry a rank. +/// `p` is the pass (`*cuda_batch.Pass(Tri)`); on return s.y / s.wts hold what the write-back reads. +pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + const r = x.tp orelse return error.UnexpectedTensor; + const s = p.s; + const n = p.rows; + const st = p.stream; + const c = p.m.collective orelse return error.TensorParallelUnavailable; + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try p.tri.topkMany(s.router_l, s.picks, s.wts, n); + const ex = try x.ops.grouped(); + const t = try scratchOf(x.ops); + try t.ensure(st, x.layer, n); + const mp = t.picks.?.ptr; + const mw = t.wts.?.ptr; + const sum = t.sum.?.ptr; + try t.mask(st, s.picks, s.wts, mp, mw, n * slots, r); + if (x.needsPool()) return error.Exl3PoolNotAttached; + const sc = try cuda.exl3_experts.scratch(t.moe.?.ptr, x.layer, t.moe_cap, slots); + var r0: usize = 0; + while (r0 < n) : (r0 += t.moe_cap) { + const rows = @min(t.moe_cap, n - r0); + const xin = s.mixed + r0 * dims * 2; + // offloaded: the rank's own pool plans on the rank's masked picks (the other rank's slots are -1) + if (x.pool) |pool| { + const dplan = @import("cuda_exl3_dplan.zig"); + if (dplan.active) |dp| if (dp.indexOf(pool)) |i| { // a decode round planned on the device (no host cut) + try dplan.runLayer(dp, st, i, ex, x.layer, xin, mp + r0 * slots * 4, rows, sc, mw + r0 * slots * 4, sum + r0 * dims * 4); + continue; + }; + try @import("cuda_exl3_pass.zig").pooled(st, x.ops, pool, ex, x.layer, xin, mp + r0 * slots * 4, rows, sc, mw + r0 * slots * 4, sum + r0 * dims * 4); + continue; + } + try ex.run(st, x.layer, xin, dims, .bf16, mp + r0 * slots * 4, mw + r0 * slots * 4, rows, sum + r0 * dims * 4, sc); + } + try c.sum(st, sum, n * dims); + try t.place(st, sum, s.y, s.wts, n); +} + +// ---- memory ------------------------------------------------------------------------------------------------------ + +/// The expert id a tensor belongs to (routed N, or 512 for the shared expert), else null (replicated). +pub fn expertOf(name: []const u8) ?usize { + if (std.mem.indexOf(u8, name, ".mlp.shared_expert.") != null) return wt.experts_n; + const key = ".mlp.experts."; + const at = std.mem.indexOf(u8, name, key) orelse return null; + const rest = name[at + key.len ..]; + const end = std.mem.indexOfScalar(u8, rest, '.') orelse return null; + return std.fmt.parseInt(usize, rest[0..end], 10) catch null; +} + +/// One tensor's device bytes on rank `r`: memory.tensorBytes for what the rank holds, 0 for the other rank's experts. +pub fn tensorBytes(name: []const u8, e: core.safetensors.Entry, o: memory.Options, r: Rank) !u64 { + if (expertOf(name)) |id| if (!r.owns(id)) return 0; + // the vocabulary-split head (load.shardHead): half the trellis and svh a rank, suh whole + if (r.world == 2 and load.shardHead(2) and (std.mem.eql(u8, name, "lm_head.trellis") or std.mem.eql(u8, name, "lm_head.svh"))) + return (try memory.tensorBytes(name, e, o)) / 2; + return memory.tensorBytes(name, e, o); +} + +/// Per layer: three 513-entry pointer tables (int64) and three width tables (int32), on every rank. +const tables_per_layer: u64 = 3 * table_n * 8 + 3 * table_n * 4; + +/// The main model's device bytes on rank `r` (the MTP head is memory.headBytes, rank 0's). +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: memory.Options, r: Rank) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer * wt.layers_n; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| total = try std.math.add(u64, total, try tensorBytes(e.key_ptr.*, e.value_ptr.*, o, r)); + } + return total; +} + +/// cuda_exl3_memory.placement at world 2: the rank's experts, everything else replicated, the MTP head on rank 0 +/// (cuda_lanes_load.zig loads it there only). +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + const r = try Rank.init(world, rank); + const o = memory.options(); + const head = if (drafts and rank == 0) try memory.headBytes(gpa, io, dir, o) else 0; + return try deviceBytes(gpa, io, dir, o, r) + head; +} + +// ---- the head's vocabulary split --------------------------------------------------------------------------------- +// +// cuda_head.zig splits a BF16 lm_head by output rows and cuda_batch.zig projectHead joins the halves with +// ncclAllGather. An EXL3 lm_head [K 2560, N 248,320] splits along N at 128-column blocks (Hadamard), 124,160 columns +// a rank; Pass.lane() already runs an EXL3 Projection of n = vocab / 2, so projectHead needs no change. What is missing +// is the hook in cuda_prompt.Engine.initRank that reads `lm_head` as EXL3 at all (the serve stage's); with it, +// `headProjection(..., world, rank)` is that hook's tp=2 branch under TF_FLASHNEXT_SHARD_HEAD=1. + +/// Rank `rank`'s output columns [lo, hi) of an `n`-column EXL3 matrix: whole 128-column blocks, equal shares. +pub fn headCut(n: usize, world: usize, rank: usize) ![2]usize { + if ((world != 1 and world != 2) or rank >= world or n % (128 * world) != 0) return error.BadHeadSharding; + const part = n / world; + return .{ rank * part, (rank + 1) * part }; +} + +/// The trellis tiles of output columns [lo, hi), in the stored layout [K/16][N/16][tile words] made dense. +pub fn shardTrellis(gpa: std.mem.Allocator, m: load.Mat, lo: usize, hi: usize) ![]u8 { + if (lo % 128 != 0 or hi % 128 != 0 or hi <= lo or hi > m.n) return error.BadHeadSharding; + const tile = @as(usize, m.halves) * 16; // bytes a 16x16 tile: 256 values at halves/2 bits + const row = (m.n / 16) * tile; + if (m.trellis.len != (m.k / 16) * row) return error.UnexpectedTensor; + const part = ((hi - lo) / 16) * tile; + const out = try gpa.alloc(u8, (m.k / 16) * part); + for (0..m.k / 16) |kt| @memcpy(out[kt * part ..][0..part], m.trellis[kt * row + (lo / 16) * tile ..][0..part]); + return out; +} + +/// The EXL3 head `prefix` (lm_head) as rank `rank`'s Projection of n / world columns (a dense shard on the device). +pub fn headProjection(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !projection.Projection { + const m = try load.mat(store, prefix); + const cut = try headCut(m.n, world, rank); + const n = cut[1] - cut[0]; + const trellis = try shardTrellis(gpa, m, cut[0], cut[1]); + defer gpa.free(trellis); + const p = try gpa.create(load.Proj); + p.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .k = m.k, .n = n, .parts = &.{} }; + errdefer p.destroy(); + var bufs: [3]u64 = undefined; + for (&bufs, [3][]const u8{ trellis, m.suh, m.svh[cut[0] * 2 .. cut[1] * 2] }) |*b, bytes| { + try p.bufs.ensureUnusedCapacity(gpa, 1); + const db = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(db); + b.* = db.ptr; + } + const st = cuda.exl3.strides(m.k, n, m.halves, false); + p.parts = try gpa.alloc(load.Part, 1); + p.parts[0] = .{ .col = 0, .n = n, .kind = .{ .exl3 = .{ .trellis = bufs[0], .suh = bufs[1], .svh = bufs[2], .k = @intCast(m.k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb } } }; + // the shard's Dense scratch (split partials, counters) is reserved with the rest at attach + p.ops.noteWeight(p.parts[0].kind.exl3, n); + store.release(); + var out = projection.Projection.empty(d); + out.n = n; + out.k = m.k; + out.exl3 = p; + return out; +} + +// ---- host tests -------------------------------------------------------------------------------------------------- + +test "exl3: expert-parallel ranks own one contiguous id range each, the shared expert on rank 1" { + const a = try Rank.init(2, 0); + const b = try Rank.init(2, 1); + try std.testing.expectEqual(@as(usize, 0), a.lo); + try std.testing.expectEqual(@as(usize, 256), a.hi); + try std.testing.expectEqual(a.hi, b.lo); + try std.testing.expectEqual(@as(usize, 513), b.hi); + for (0..table_n) |id| try std.testing.expect(a.owns(id) != b.owns(id)); // every id exactly once + try std.testing.expect(b.owns(wt.experts_n) and !a.owns(wt.experts_n)); + const one = try Rank.init(1, 0); + try std.testing.expect(one.owns(0) and one.owns(511) and one.owns(512)); + try std.testing.expectError(error.BadRank, Rank.init(2, 2)); + try std.testing.expectError(error.BadRank, Rank.init(4, 0)); +} + +test "exl3: expert tensors count on their own rank only; the router, the shared gate and HC on both" { + try std.testing.expectEqual(@as(?usize, 7), expertOf("model.language_model.layers.0.mlp.experts.7.gate_proj.trellis")); + try std.testing.expectEqual(@as(?usize, 511), expertOf("model.language_model.layers.47.mlp.experts.511.down_proj.svh")); + try std.testing.expectEqual(@as(?usize, 512), expertOf("model.language_model.layers.3.mlp.shared_expert.up_proj.suh")); + try std.testing.expectEqual(@as(?usize, null), expertOf("model.language_model.layers.3.mlp.shared_expert_gate.weight")); + try std.testing.expectEqual(@as(?usize, null), expertOf("model.language_model.layers.3.mlp.gate.weight")); + var e: core.safetensors.Entry = .{ .dtype = .f16, .rank = 1, .shape = .{ 2560, 1, 1, 1, 1 }, .begin = 0, .end = 5120 }; + const r0 = try Rank.init(2, 0); + const r1 = try Rank.init(2, 1); + const routed = "model.language_model.layers.0.mlp.experts.300.gate_proj.suh"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(routed, e, .{}, r0)); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes(routed, e, .{}, r1)); + const shared = "model.language_model.layers.0.mlp.shared_expert.gate_proj.suh"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared, e, .{}, r0)); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes(shared, e, .{}, r1)); + e.dtype = .f16; + e.rank = 2; + e.shape = .{ 512, 2560, 1, 1, 1 }; + e.end = 512 * 2560 * 2; + const router = "model.language_model.layers.0.mlp.gate.weight"; + try std.testing.expectEqual(try tensorBytes(router, e, .{}, r0), try tensorBytes(router, e, .{}, r1)); + try std.testing.expect(try tensorBytes(router, e, .{}, r0) > 0); +} + +test "exl3: the split head counts half its trellis and svh a rank" { + const r0 = try Rank.init(2, 0); + const tr: core.safetensors.Entry = .{ .dtype = .i16, .rank = 3, .shape = .{ 160, 15520, 80, 1, 1 }, .begin = 0, .end = 160 * 15520 * 80 * 2 }; + if (load.shardHead(2)) try std.testing.expectEqual(tr.end / 2, try tensorBytes("lm_head.trellis", tr, .{}, r0)); + try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{}, try Rank.init(1, 0))); +} + +test "exl3: tensor parallel placement opens at world 2 and still refuses other worlds" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = @embedFile("fixtures/exl3_config.json") }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + // past the gate: no index in the fixture dir, so it fails reading the checkpoint, not at the quantization + for ([_]usize{ 0, 1 }) |rank| { + if (memory.placement(a, std.testing.io, dir, 2, rank, rank == 0)) |_| return error.TestUnexpectedResult else |err| try std.testing.expect(err != error.UnsupportedQuantization and err != error.BadRank); + } + try std.testing.expectError(error.BadRank, memory.placement(a, std.testing.io, dir, 3, 0, false)); + try std.testing.expectError(error.BadRank, memory.placement(a, std.testing.io, dir, 2, 2, false)); +} + +test "exl3: the head splits at 128-column blocks and a shard's tiles are the full matrix's columns" { + try std.testing.expectEqual([2]usize{ 0, 124_160 }, try headCut(248_320, 2, 0)); + try std.testing.expectEqual([2]usize{ 124_160, 248_320 }, try headCut(248_320, 2, 1)); + try std.testing.expectError(error.BadHeadSharding, headCut(248_320 + 128, 2, 0)); + try std.testing.expectError(error.BadHeadSharding, headCut(256, 2, 2)); + // K 32 (2 tile rows), N 512 (32 tiles a row), 3 halves: 48 bytes a tile + const k = 32; + const n = 512; + const halves = 3; + const tile = halves * 16; + var bytes: [(k / 16) * (n / 16) * tile]u8 = undefined; + for (&bytes, 0..) |*v, i| v.* = @truncate(i * 7 + i / 255); + const m: load.Mat = .{ .trellis = &bytes, .suh = &.{}, .svh = &.{}, .k = k, .n = n, .halves = halves }; + const a = std.testing.allocator; + for (0..2) |rank| { + const cut = try headCut(n, 2, rank); + const got = try shardTrellis(a, m, cut[0], cut[1]); + defer a.free(got); + try std.testing.expectEqual(bytes.len / 2, got.len); + for (0..k / 16) |kt| for (0..(cut[1] - cut[0]) / 16) |j| { + const src = (kt * (n / 16) + cut[0] / 16 + j) * tile; + const dst = (kt * ((cut[1] - cut[0]) / 16) + j) * tile; + try std.testing.expectEqualSlices(u8, bytes[src..][0..tile], got[dst..][0..tile]); + }; + } + try std.testing.expectError(error.BadHeadSharding, shardTrellis(a, m, 0, 100)); +} diff --git a/zig/src/families/flashnext/cuda_images_test.zig b/zig/src/families/flashnext/cuda_images_test.zig index 38fe41a616..e053f6a2f7 100644 --- a/zig/src/families/flashnext/cuda_images_test.zig +++ b/zig/src/families/flashnext/cuda_images_test.zig @@ -42,6 +42,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, draf engine.chunk = 16; engine.eos = 7; engine.round_graph = .{}; + engine.page_config = null; var ids = [_]u32{1}; var head: mtp.Head = undefined; head.ids = &ids; @@ -138,6 +139,7 @@ fn release(comptime Cuda: type, comptime open: anytype) !void { engine.chunk = 16; engine.eos = 7; engine.round_graph = .{}; + engine.page_config = null; var self: Cuda = undefined; self.gpa = a; self.context = 64; diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 226c988b55..a8c11c600e 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -9,6 +9,7 @@ const mtp = @import("cuda_mtp.zig"); const snap = @import("cuda_snap.zig"); const ranks = @import("cuda_rank_engine.zig"); const config = @import("cuda_config.zig"); +const exl3_memory = @import("cuda_exl3_memory.zig"); const Glue = @import("cuda_glue.zig").Glue; const living = @import("living"); const slide_mod = @import("cuda_slide.zig"); @@ -27,8 +28,9 @@ pub const captured_head_entries = [_][]const u8{ "_rmsnorm", "_add_streams" }; pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; - pub const tp_draft_snapshots = false; - pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed" }; + // ours: rank 0 saves the head's rows with the sequence, rank 1 the sequence (cuda_rank_engine.Session.save) + pub const tp_draft_snapshots = true; + pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed", @import("cuda_exl3_config.zig").format }; // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments @@ -48,7 +50,8 @@ pub fn Namespace(comptime Tri: type) type { pub const min_prompt: u32 = 512; }; - pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 1, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551, kv8: bool = false }; + pub const preallocates_lanes = false; /// Living Weights' learner (--slide): lessons put neurons in each layer's pocket beside its shared expert. pub const slide = slide_mod.Slide(Tri); @@ -75,6 +78,8 @@ pub fn Namespace(comptime Tri: type) type { stream_bytes: usize, /// Every stream allocates its own sequence; none rides on the engine's. free_streams: u32 = 0, + /// Kept prompt states in pinned host memory (cuda_snap.on_host): the native host budgets them there. + host_snaps: bool = false, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, lone: ?LoneRun = null, @@ -89,6 +94,7 @@ pub fn Namespace(comptime Tri: type) type { /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + if (exl3_memory.isExl3(gpa, io, dir)) return exl3_memory.weightBytes(gpa, io, dir); const main = weights.deviceBytes(gpa, io, dir) catch return 0; return main + mtp.headBytes(); } @@ -96,6 +102,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const mode = try config.mode(gpa, io, dir); + if (mode == .exl3) return exl3_memory.placement(gpa, io, dir, world, rank, drafts); if (mode == .modelopt) { const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); // Dense draft projections and vocabulary rows exceed the affine head's estimate. @@ -119,8 +126,10 @@ pub fn Namespace(comptime Tri: type) type { } pub fn rankRoom(ptr: *anyopaque, local: u64) !u64 { - if (of(ptr).ranks) |r| return r.room(local); - return local; + const own = of(ptr); + const protected = try @import("cuda_exl3_kvpage.zig").admissionRoom(@import("exl3_startup_memory.zig").laneRoom(local, own.graph_reservation), own.stream_bytes, own.engine.page_config); + if (own.ranks) |r| return @import("cuda_exl3_kv.zig").rankRoom(r, protected, own.stream_bytes); + return protected; } /// A request or load this engine refuses, in words; null: none of its own. @@ -156,6 +165,8 @@ pub fn Namespace(comptime Tri: type) type { vision: ?*@import("cuda_images.zig").Vision = null, ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, + profile: ?@import("cuda_exl3_frequency_state.zig").State = null, + profile_counts: ?@import("cuda_exl3_frequency.zig").Counts = null, depth: usize = mtp.default_depth, confidence: f64 = mtp.default_confidence, /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. @@ -167,6 +178,9 @@ pub fn Namespace(comptime Tri: type) type { logits: cuda.DeviceBuffer, drawn: [ring]u32 = undefined, next: u64 = 0, + /// This rank's Loaded.stream_bytes (rankRoom agrees the ranks' rooms in streams, not bytes). + stream_bytes: usize = 0, + graph_reservation: u64 = 0, /// --slide: the pocket's layout, from config.json at load. lay: living.Layout = undefined, @@ -177,6 +191,7 @@ pub fn Namespace(comptime Tri: type) type { fn ops(self: *Cuda) Ops { return .{ .e = self.engine, .tri = self.tri(), .h = self.head.?, .main = if (self.ranks != null) .{ .ptr = self, + .rank = self.ranks.?.rank, .forward = mainForward, .verify = mainVerify, .rollback = mainRollback, @@ -188,9 +203,9 @@ pub fn Namespace(comptime Tri: type) type { return self.ranks.?.forward(Tri, self.tri(), q, ids, hooks); } - fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) anyerror!void { + fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32, tail: ?prompt.batch.Hooks.OutputTail) anyerror!void { const self = of(ptr); - return self.ranks.?.verify(Tri, self.tri(), q, d, ids); + return self.ranks.?.verifyWithTail(Tri, self.tri(), q, d, ids, tail); } fn mainRollback(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, keep: usize) anyerror!void { @@ -265,14 +280,14 @@ pub fn Namespace(comptime Tri: type) type { if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; const gop = try self.lanes.getOrPut(self.gpa, s); if (gop.found_existing) dropLane(self, gop.value_ptr); - const drafting = self.head != null and s.drafts; + const drafting = self.head != null and s.drafts and !@import("cuda_exl3_pass.zig").plainByDefault(); // without the head's drafts, one row a stream a round: only a drafted window can rewind its states if (!drafting) { s.drafts = false; s.proposer = null; } // a window's rows past the reply's last token (its drafts) still need cache slots - const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); + const cap = if (self.engine.page_config != null) @min(prompt.max_tokens, self.context + mtp.max_rows) else @import("cuda_exl3_recycle.zig").roundCap(@min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens), @min(prompt.max_tokens, self.context + mtp.max_rows)); gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; @@ -283,7 +298,7 @@ pub fn Namespace(comptime Tri: type) type { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, @import("cuda_exl3_prefill.zig").draftRows(self.engine)); // EXL3 deep prefill if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -295,14 +310,16 @@ pub fn Namespace(comptime Tri: type) type { const saved = s.reuse.saved orelse return l; const kept: *snap.Snap = @ptrCast(@alignCast(saved)); if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { + std.log.warn("flashnext: kept state at {d} not restored for reuse at {d} of {d} ids", .{ kept.at, s.reuse.at, ids.len }); s.reuse_failed = true; - } else if (self.ranks != null and l.draft != null) { - s.reuse_failed = true; - } else if (if (self.ranks) |r| r.restore(&l.seq, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { + } else if (if (self.ranks) |r| r.restore(&l.seq, if (l.draft) |*d| d else null, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { l.filled = kept.at; s.cached = @intCast(kept.at); } else |err| switch (err) { - error.SequenceTooShort, error.NoHeadState => s.reuse_failed = true, + error.SequenceTooShort, error.NoHeadState, error.KvPageSnapshotLayout => { + std.log.warn("flashnext: kept state at {d} not restored ({t})", .{ kept.at, err }); + s.reuse_failed = true; + }, else => return err, } return l; @@ -313,12 +330,15 @@ pub fn Namespace(comptime Tri: type) type { const ids = s.prompt(); const at = l.filled orelse return true; if (s.isCancelled()) return error.Cancelled; - var stop = @min(ids.len, at + self.engine.chunk); + var stop = @min(ids.len, at + @import("cuda_exl3_prefill.zig").stepRows(self.engine, if (l.draft) |*d| d.rows_cap else null)); // EXL3 deep prefill for (s.reuse.marks) |m| if (m > at and m < stop) { stop = m; break; }; const last = stop == ids.len; + const previous_telemetry_phase = self.engine.telemetry.forward_phase; + self.engine.telemetry.forward_phase = .prefill; + defer self.engine.telemetry.forward_phase = previous_telemetry_phase; if (l.draft) |*d| { try self.ops().prefillChunk(&l.seq, d, ids, at, stop); } else if (self.ranks) |r| { @@ -362,6 +382,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); + if (own.profile) |*profile| profile.save(own); var it = own.lanes.valueIterator(); while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); @@ -374,6 +395,7 @@ pub fn Namespace(comptime Tri: type) type { if (own.set) |*x| x.deinit(); own.glue.deinit(); own.stream.deinit(); + if (own.profile) |*profile| profile.deinit(); own.gpa.destroy(own); } @@ -440,7 +462,7 @@ pub fn Namespace(comptime Tri: type) type { var row_ids: [mtp.max_rows]u32 = undefined; row_ids[0] = w.pending; @memcpy(row_ids[1..rows], w.tokens); - try self.ops().verify(&l.seq, d, row_ids[0..rows]); + try self.ops().verifySampled(&l.seq, d, row_ids[0..rows], w.stream.sampling); try self.ops().draws(d, w.stream.sampling, o.sampled[0..rows]); } else { if (rows != 1) return error.CannotRewind; @@ -534,7 +556,9 @@ pub fn Namespace(comptime Tri: type) type { // -- the prompt cache's copies ------------------------------------------------------------------------ fn snapBytes(ptr: *anyopaque, at: u32) u64 { - return snap.bytes(at, of(ptr).head != null); + const self = of(ptr); + if (self.engine.page_config) |cfg| return snap.pagedBytes(@min(self.context + mtp.max_rows, prompt.max_tokens), cfg.forChunk(self.engine.chunk), at, self.head != null) catch std.math.maxInt(u64); + return snap.bytesKv(at, of(ptr).head != null, @import("cuda_exl3_kv.zig").of(of(ptr).engine).rows()); } /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). @@ -543,22 +567,24 @@ pub fn Namespace(comptime Tri: type) type { const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const l = try self.lane(s); if (l.seq.fed != at) return error.NotAtMark; - if (self.ranks) |r| { - if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; - return try r.save(&l.seq); - } + // ours: a drafting stream's state carries the head's rows on rank 0 (cuda_rank_engine.Session.save) + if (self.ranks) |r| return try r.save(&l.seq, if (l.draft) |*d| d else null); return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); } fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { const self = of(ptr); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (self.ranks) |r| { - if (l.draft != null) return error.NoHeadState; - return r.restore(&l.seq, @ptrCast(@alignCast(saved))); - } - try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); + const l = self.lane(s) catch |err| { + std.log.warn("flashnext: a kept state was not restored: no lane ({t})", .{err}); + return err; + }; + if (self.ranks) |r| return r.restore(&l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); + self.engine.round_graph.clear(); + snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))) catch |err| { + std.log.warn("flashnext: a kept state was not restored ({t})", .{err}); + return err; + }; } fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 5589587c80..a0805034d3 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -17,6 +17,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun engine.eos = 7; engine.chunk = 16; engine.round_graph = .{}; + engine.page_config = null; var ids = [_]u32{1}; var head: mtp.Head = undefined; head.ids = &ids; diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 69c300ef69..4271cf9283 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -7,6 +7,8 @@ const ranks = @import("cuda_rank_engine.zig"); const config = @import("cuda_config.zig"); const Glue = @import("cuda_glue.zig").Glue; const images = @import("cuda_images.zig"); +const memory = @import("exl3_startup_memory.zig"); +const exl3_pass = @import("cuda_exl3_pass.zig"); pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: F.Options) !F.Loaded { if (o.vision) |v| try @import("qwen_vision").media.checkLimits(v); const Cuda = F.Cuda; @@ -16,7 +18,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.segments != 1) return error.UnsupportedOption; if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; const mode = try config.mode(gpa, io, dir); - if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; + if (o.tp != 1 and ((mode != .modelopt and mode != .exl3) or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const driver: *cuda.Driver = @constCast(ctx.d); @@ -57,16 +59,23 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.engine.image_factory = own.vision; } errdefer if (own.vision) |v| v.deinit(); + var persistent_bytes = cuda.usage(false).device; + try @import("cuda_exl3_kv.zig").setup(own.engine, o.kv8, mode == .exl3); // K10: int8 KV, host index keys, chunk // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@max(@min(o.context, own.engine.chunk), if (o.slide) @as(usize, @import("cuda_train.zig").pack_rows) else 0)); // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) if (o.drafts and (o.tp == 1 or o.rank == 0)) { + const head_before = cuda.usage(false).device; own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); + persistent_bytes = try std.math.add(u64, persistent_bytes, cuda.usage(false).device -| head_before); own.depth = @max(1, @min(mtp.max_chain, F.envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; } errdefer if (own.head) |h| h.deinit(); const window = @min(o.context + mtp.max_rows, prompt.max_tokens); + const common_kv = try own.engine.reserveSharedKv(window); + persistent_bytes = try std.math.add(u64, persistent_bytes, common_kv); + if (common_kv != 0) std.log.info("flashnext rank {d}: shared BF16 KV common allocation {d} bytes, charged once before lane and expert-pool admission", .{ o.rank, common_kv }); // and an attention block's block scores at the window (64 MiB at 262,144 tokens) _ = try own.engine.reserveScores(window, @min(o.context, own.engine.chunk)); own.logits = try cuda.DeviceBuffer.alloc(driver, round_streams * prompt.vocabulary * 2); @@ -74,11 +83,40 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.lanes = .empty; own.next = 0; if (o.slide) own.lay = try @import("cuda_slide.zig").layoutOf(gpa, io, dir); - var stream_bytes = prompt.seqBytes(window); + const scratch_bytes = cuda.usage(false).device -| persistent_bytes; + const deep_before = cuda.usage(false).device; + @import("cuda_exl3_prefill.zig").reserveFor(own.engine, o.context, o.tp); // EXL3 offload: deep prefill rows before the pools + const deep_bytes = cuda.usage(false).device -| deep_before; + var stream_bytes = try @import("cuda_exl3_kv.zig").engineSeqBytes(@import("cuda_exl3_kv.zig").of(own.engine), window, own.engine.chunk); if (own.head) |h| { - stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); - } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); - if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; + stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, @import("cuda_exl3_prefill.zig").draftRows(own.engine))); + } else if (o.drafts and o.tp == 2) stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(1, 1, @import("cuda_exl3_prefill.zig").draftRows(own.engine))); // rank 1's window (no head cache) + if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); + try @import("cuda_exl3_kv.zig").logStream(own.engine, window, stream_bytes); // K10: what a stream holds, by piece + own.stream_bytes = stream_bytes; + var pool_budget = try exl3_pass.budgetForStreams(stream_bytes, o.lanes); + if (mode == .exl3 and @import("cuda_exl3_graph.zig").enabled()) pool_budget.graph_reservation = try @import("cuda_exl3_graph_memory.zig").startupReservationForStreams(o.lanes); + own.graph_reservation = pool_budget.graph_reservation; + const placement = @import("cuda_exl3_frequency_state.zig"); + if (try placement.prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], own.head, @intCast(o.tp), @intCast(o.rank))) |prepared| { + own.profile = prepared.state; + own.profile_counts = prepared.counts; + } + errdefer if (own.profile) |*profile| profile.deinit(); + const profile_counts = if (own.profile_counts) |*counts| &counts.decoder else null; + const attached = try exl3_pass.attach(gpa, io, driver, own.stream, own.engine.layers[0..own.engine.n_layers], if (own.head) |h| &h.layer else null, own.engine.chunk, pool_budget, profile_counts); + own.graph_reservation = attached.graph_reservation; + if (own.graph_reservation != 0) std.log.info("flashnext rank {d}: lane admission protects full graph reservation {d} bytes; already materialized graphs remain charged in free memory, so admission conservatively reserves those bytes twice", .{ o.rank, own.graph_reservation }); + if (own.profile_counts) |*counts| try placement.finish(own, counts); + const ledger: memory.Ledger = .{ + .persistent = persistent_bytes, + .scratch = scratch_bytes, + .deep_prefill = try std.math.add(u64, deep_bytes, attached.prefill_window_bytes), + .pool_scratch = attached.scratch_bytes, + .expert_pools = attached.device_bytes, + .reservation = .{ .stream_bytes = stream_bytes, .streams = if (stream_bytes == 0) 0 else pool_budget.other / stream_bytes, .graph_bytes = attached.graph_reservation, .after_pool_bytes = attached.after_pool_bytes }, + }; + std.log.info("flashnext rank {d} memory ledger: persistent={d} scratch={d} deep_prefill={d} pool_scratch={d} expert_pools={d} stream_bytes={d} streams={d} graph_reservation={d} planned_total={d} free_after_pools={d}", .{ o.rank, ledger.persistent, ledger.scratch, ledger.deep_prefill, ledger.pool_scratch, ledger.expert_pools, stream_bytes, ledger.reservation.streams, pool_budget.graph_reservation, try ledger.total(), attached.free_after }); return .{ .vision = own.vision, .backend = own.backend(), @@ -86,6 +124,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu .rows = if (own.head != null) mtp.max_rows else 2, .prompt_rows = @intCast(own.engine.chunk), .stream_bytes = stream_bytes, + .host_snaps = @import("cuda_exl3_kv.zig").snapHost(o.context, mode == .exl3), .ctx = own, .deinit = F.release, .learns = if (o.slide) own else null, diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig index e0e5967ecf..a4ab9ae6b5 100644 --- a/zig/src/families/flashnext/cuda_lanes_restore_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -61,6 +61,7 @@ pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytyp engine.eos = 7; engine.chunk = 16; engine.round_graph = .{}; + engine.page_config = null; var pair: [2]std.posix.socket_t = undefined; if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; defer _ = std.c.close(pair[0]); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 99e25c0c12..6580d2082b 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const exl3 = @import("cuda_exl3.zig"); pub const rank_engine = @import("cuda_rank_engine.zig"); pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); @@ -19,8 +20,12 @@ pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const ple_check = @import("cuda_ple_check.zig"); +/// An EXL3 pack's n-gram table and draft head rows (cuda_ple.Ngram's `.exl3` format, cuda_mtp_load.draft's EXL3 branch) +pub const exl3_ple = @import("cuda_exl3_ple.zig"); +pub const exl3_mtp = @import("cuda_exl3_mtp.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); +pub const snapshot = @import("cuda_snap.zig"); /// The row-by-row reference path, for flashnext-prompt pub const row_path = @import("cuda_rows.zig"); /// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback @@ -29,6 +34,7 @@ pub const mtp = @import("cuda_mtp.zig"); pub const pocket = @import("cuda_pocket.zig"); /// Living Weights' trainer over the batched pass pub const train = @import("cuda_train.zig"); +pub const exl3_pool = @import("cuda_exl3_pool.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); /// Our own glue kernels, which every launch takes when no captured set is loaded pub const Glue = glue_.Glue; @@ -385,6 +391,7 @@ test { _ = native; _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); + _ = @import("cuda_exl3.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_round_graph.zig"); @@ -394,6 +401,7 @@ test { _ = row_path; _ = @import("cuda_mtp_check.zig"); _ = @import("cuda_draft_affine.zig"); + _ = @import("cuda_exl3_mtp.zig"); _ = @import("cuda_head.zig"); _ = fn_back; _ = pocket; @@ -401,6 +409,13 @@ test { _ = @import("living.zig"); } +test "exl3: greedy verification host contracts" { + _ = @import("cuda_mtp_ops.zig"); + _ = @import("cuda_mtp_draft.zig"); + _ = @import("cuda_mtp_verify_cache.zig"); + _ = @import("cuda_exl3_dispatch_test.zig"); +} + test "embed grid is one program per row and per group of 32" { try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); } diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig index 435cf6c9da..b67ae0e8cc 100644 --- a/zig/src/families/flashnext/cuda_mtp_check.zig +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -47,7 +47,7 @@ pub fn decode(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, d: *Draft, out: @memcpy(ids[1..][0..ch.n], ch.slice()); const rows = 1 + ch.n; const tv = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); - try o.verify(seq, d, ids[0..rows]); + try o.verifySampled(seq, d, ids[0..rows], sm); try o.draws(d, sm, sampled[0..rows]); const td = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); st.verify_ms += @as(f64, @floatFromInt(td - tv)) / 1e6; diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index 1ae18ae2ab..ef0159b1cf 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -42,6 +42,18 @@ pub const Chain = struct { } }; +test "exl3: reset and settlement discard completed verification records" { + var window: Draft = undefined; + window.verify_cache = .{}; + window.verify_cache.publish(window.verify_cache.generation, 4, 10); + window.reset(); + try std.testing.expect(!window.verify_cache.take(4, 10, true)); + window.verify_cache.publish(window.verify_cache.generation, 4, 10); + window.open = false; + window.settle(undefined); + try std.testing.expect(!window.verify_cache.take(4, 10, true)); +} + /// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. pub const Draft = struct { media: ?@import("cuda_media.zig").View = null, @@ -81,6 +93,9 @@ pub const Draft = struct { /// The scores a sampled draw reads on the host (a window's rows, or the draft head's row). host: []u16, best_host: [max_rows * 4]f32 = undefined, + verify_host: ?cuda.HostBuffer = null, + verify_pending: ?cuda.Stream = null, + verify_cache: @import("cuda_mtp_verify_cache.zig").Cache = .{}, const hist_at: usize = 0; const proj_at: usize = hist_at + linear_n * conv_bytes; @@ -96,6 +111,7 @@ pub const Draft = struct { /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { + if (try @import("cuda_exl3_recycle.zig").takeDraft(stream, cap, draft_n, rows)) |kept| return kept; // graphs: stable addresses var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); errdefer d.kc.free(); @@ -147,6 +163,7 @@ pub const Draft = struct { /// An open window kept whole: its end states become the sequence's (buffers swapped, nothing copied). pub fn settle(d: *Draft, seq: *prompt.Seq) void { + d.verify_cache.invalidate(); if (!d.open) return; for (seq.state, &d.spare) |*a, *b| std.mem.swap(cuda.DeviceBuffer, a, b); d.aim(); @@ -154,6 +171,13 @@ pub const Draft = struct { } pub fn deinit(d: *Draft) void { + d.verify_cache.invalidate(); + if (d.verify_pending) |stream| { + stream.synchronize() catch {}; + d.verify_pending = null; + } + if (@import("cuda_exl3_recycle.zig").giveDraft(d)) return; // graphs: kept for the next stream + if (d.verify_host) |*host| host.free(); d.gpa.free(d.host); for (&d.spare) |*b| b.free(); d.saved.free(); @@ -170,6 +194,7 @@ pub const Draft = struct { /// Drop the head's cache (a new prompt). pub fn reset(d: *Draft) void { + d.verify_cache.invalidate(); d.len = 0; d.drafted = 0; d.rows = 0; diff --git a/zig/src/families/flashnext/cuda_mtp_load.zig b/zig/src/families/flashnext/cuda_mtp_load.zig index dc1a226d31..4ff61b9114 100644 --- a/zig/src/families/flashnext/cuda_mtp_load.zig +++ b/zig/src/families/flashnext/cuda_mtp_load.zig @@ -8,6 +8,7 @@ pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []con const mode = if (std.c.getenv("TF_FLASHNEXT_DRAFT_HEAD")) |v| std.mem.span(v) else "bf16"; const affine = std.mem.eql(u8, mode, "affine4"); if (!affine and !std.mem.eql(u8, mode, "bf16")) return error.BadDraftHead; + if (try @import("cuda_exl3_mtp.zig").draft(a, d, store, ids, affine)) |q| return q; const tensor = try store.tensor("language_model.lm_head.weight"); if (tensor.dtype == .bf16) { if (!tensor.is(.bf16, &.{ wt.vocab, wt.dims })) return error.UnexpectedTensor; diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index cfcd9fdd05..aefe8cb33f 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -42,6 +42,156 @@ fn cint(v: usize) c_int { return @intCast(v); } +test "exl3: greedy verification consumes completed records without touching the head or CUDA" { + var records: [max_rows * 16]u8 align(16) = undefined; + const values = std.mem.bytesAsSlice(f32, &records); + const tokens = [_]u32{ 0, 6124, 9564, @intCast(vocab - 1) }; + for (tokens, 0..) |token, row| values[4 * row + 2] = @bitCast(@as(i32, @intCast(token))); + var window: Draft = undefined; + window.rows = tokens.len; + window.fed0 = 37; + window.verify_cache = .{}; + window.verify_host = .{ .d = undefined, .bytes = &records }; + const ops: Ops(struct {}) = .{ .e = undefined, .tri = .{}, .h = undefined }; + var drawn: [tokens.len]u32 = undefined; + window.verify_cache.publish(window.verify_cache.generation, window.rows, window.fed0); + try ops.draws(&window, null, &drawn); + try std.testing.expectEqualSlices(u32, &tokens, &drawn); + try std.testing.expect(window.verify_cache.ready == null); + window.verify_cache.publish(window.verify_cache.generation, window.rows, window.fed0); + try ops.draws(&window, .{ .seed = 0, .temperature = 0 }, &drawn); + try std.testing.expectEqualSlices(u32, &tokens, &drawn); +} + +test "exl3: verification tail enqueues the existing top-1 kernel then only its records copy" { + const Fake = struct { + var launches: usize = 0; + var copies: usize = 0; + var rows: usize = 0; + var copied: usize = 0; + var fail_copy: bool = false; + + fn launch(_: cuda.abi.Function, grid: c_uint, _: c_uint, _: c_uint, block: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, params: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + std.debug.assert(block == 1024); + std.debug.assert(launches == copies); + const scores: *const u64 = @ptrCast(@alignCast(params.?[0].?)); + const width: *const c_int = @ptrCast(@alignCast(params.?[1].?)); + const stride: *const c_longlong = @ptrCast(@alignCast(params.?[2].?)); + const output: *const u64 = @ptrCast(@alignCast(params.?[3].?)); + std.debug.assert(scores.* == 0x2000 and width.* == vocab and stride.* == vocab and output.* == 0x4000); + rows = grid; + launches += 1; + return 0; + } + + fn copy(_: ?*anyopaque, src: u64, bytes: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + std.debug.assert(launches == copies + 1 and src == 0x4000); + copies += 1; + copied = bytes; + if (fail_copy) return cuda.abi.error_not_ready; + return 0; + } + }; + Fake.launches = 0; + Fake.copies = 0; + Fake.fail_copy = false; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Fake.launch; + driver.api.cuMemcpyDtoHAsync_v2 = Fake.copy; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var head: Head = undefined; + head.best_fn = .{ .d = &driver, .handle = null }; + var records: [max_rows * 16]u8 align(16) = undefined; + var window: Draft = undefined; + window.logits = .{ .d = &driver, .ptr = 0x2000, .len = max_rows * vocab * 2 }; + window.best = .{ .d = &driver, .ptr = 0x4000, .len = max_rows * 16 }; + window.verify_host = .{ .d = &driver, .bytes = &records }; + window.verify_cache = .{}; + for ([_]usize{ 1, 2, 3, 4, 8, 16 }) |rows| { + window.rows = rows; + window.verify_cache.invalidate(); + var tail: Ops(struct {}).VerifyTail = .{ .head = &head, .draft = &window }; + try Ops(struct {}).VerifyTail.enqueue(&tail, stream); + try std.testing.expectEqual(rows, Fake.rows); + try std.testing.expectEqual(rows * 16, Fake.copied); + try std.testing.expect(tail.enqueued); + try std.testing.expectEqual(window.verify_cache.generation, tail.generation); + try std.testing.expect(window.verify_cache.ready == null); + } + Fake.fail_copy = true; + var failed: Ops(struct {}).VerifyTail = .{ .head = &head, .draft = &window }; + try std.testing.expectError(error.NotReady, Ops(struct {}).VerifyTail.enqueue(&failed, stream)); + try std.testing.expect(failed.enqueued); + try std.testing.expect(window.verify_pending != null); + try std.testing.expect(window.verify_cache.ready == null); +} + +test "exl3: common verification scopes TP1 and TP2 validation and delegation without joins" { + const Tri = @import("cuda_launch.zig").Tri; + const Fake = @import("cuda_exl3_shared_kv_test.zig").Fake; + const Phase = @import("exl3_telemetry.zig").Phase; + const Delegate = struct { + engine: *prompt.Engine, + calls: usize = 0, + fail: bool = false, + + fn verify(ptr: *anyopaque, _: *prompt.Seq, _: *Draft, _: []const u32, _: ?batch.Hooks.OutputTail) !void { + const delegate: *@This() = @ptrCast(@alignCast(ptr)); + delegate.calls += 1; + try std.testing.expectEqual(Phase.verify, delegate.engine.telemetry.forward_phase); + if (delegate.fail) return error.VerificationFailed; + } + + fn forward(_: *anyopaque, _: *prompt.Seq, _: []const u32, _: batch.Hooks) !void { + return error.UnexpectedDelegation; + } + + fn rollback(_: *anyopaque, _: *prompt.Seq, _: *Draft, _: usize) !void { + return error.UnexpectedDelegation; + } + }; + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.stream = &stream; + engine.round_graph = .{}; + var sequence: prompt.Seq = undefined; + sequence.cap = max_rows; + var window: Draft = undefined; + window.verify_cache = .{}; + var delegate: Delegate = .{ .engine = &engine }; + var ops: Ops(Tri) = .{ .e = &engine, .tri = .{ .s = stream }, .h = undefined }; + const too_wide: [max_rows + 1]u32 = @splat(1); + const syncs = Fake.sync_calls; + for ([_]usize{ 1, 2 }) |world| { + for ([_]bool{ false, true }) |enabled| { + for ([_]Phase{ .unknown, .prefill, .decode, .verify }) |previous| { + engine.telemetry = .{ .enabled = enabled, .rank = .{ .rank = 0, .world = world }, .forward_phase = previous }; + sequence.fed = 0; + ops.main = null; + try std.testing.expectError(error.WindowTooWide, ops.verifyWithTail(&sequence, &window, &.{}, null)); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + try std.testing.expectError(error.WindowTooWide, ops.verifyWithTail(&sequence, &window, &too_wide, null)); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + sequence.fed = sequence.cap; + try std.testing.expectError(error.PromptTooLong, ops.verifyWithTail(&sequence, &window, &.{1}, null)); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + ops.main = .{ .ptr = &delegate, .rank = 0, .forward = Delegate.forward, .verify = Delegate.verify, .rollback = Delegate.rollback }; + const calls = delegate.calls; + delegate.fail = false; + try ops.verifyWithTail(&sequence, &window, &.{1}, null); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + delegate.fail = true; + try std.testing.expectError(error.VerificationFailed, ops.verifyWithTail(&sequence, &window, &.{1}, null)); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + try std.testing.expectEqual(calls + 2, delegate.calls); + try std.testing.expectEqual(enabled, engine.telemetry.enabled); + try std.testing.expectEqual(syncs, Fake.sync_calls); + } + } + } +} + fn copyAsync(stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { if (n == 0) return; const d = stream.d; @@ -57,8 +207,9 @@ fn download(stream: cuda.Stream, dst: []u8, src: u64) !void { /// An owner can mirror main-model mutations to a peer without exposing device pointers on the control wire. pub const Main = struct { ptr: *anyopaque, + rank: usize, forward: *const fn (*anyopaque, *prompt.Seq, []const u32, batch.Hooks) anyerror!void, - verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32) anyerror!void, + verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32, ?batch.Hooks.OutputTail) anyerror!void, rollback: *const fn (*anyopaque, *prompt.Seq, *Draft, usize) anyerror!void, }; @@ -83,6 +234,7 @@ pub fn Ops(comptime Tri: type) type { /// Prompt rows [at, stop): the pass with its streams kept, then the head over each row with a known next token. pub fn prefillChunk(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, at: usize, stop: usize) !void { + d.verify_cache.invalidate(); if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; d.media = if (seq.media) |*m| m.view() else null; @@ -90,7 +242,7 @@ pub fn Ops(comptime Tri: type) type { try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); const head_to = @min(stop, ids.len - 1); if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); - if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; + d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; } /// The whole prompt a chunk of the engine's rows at a time, from an empty head cache. @@ -106,11 +258,15 @@ pub fn Ops(comptime Tri: type) type { /// The head over `next.len` rows of streams `rows`: cache at d.len, last row in d.out, scores in d.dlogits. pub fn forwardHead(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + d.verify_cache.invalidate(); const e = o.e; const tri = o.tri; const h = o.h; if (next.len == 0) return; if (d.len + next.len > d.cap) return error.PromptTooLong; + if (try @import("cuda_exl3_prefill.zig").head(Tri, o, d, next, rows, scores)) return; // EXL3: the head expert-major over a deep chunk + const planned = try @import("cuda_exl3_chain.zig").enter(o.e, @min(next.len, piece_rows)); // graphs: the head's experts on the device planner + defer @import("cuda_exl3_chain.zig").leave(planned); const want = @max(@as(usize, 32), 2 * @min(next.len, piece_rows)); const s = try e.reserve(want); const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); @@ -175,6 +331,7 @@ pub fn Ops(comptime Tri: type) type { /// Python draft.absorb: the chain's entries trimmed, then the kept rows (main streams `rows`, next tokens). pub fn absorb(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + d.verify_cache.invalidate(); if (d.drafted > 0) { d.len -= d.drafted; d.drafted = 0; @@ -184,6 +341,7 @@ pub fn Ops(comptime Tri: type) type { /// The draft at `position` and its temperature-1 probability (Python Engine.sample_draft). pub fn pick(o: Self, d: *Draft, position: u64, sm: ?Sampling) !Pick { + d.verify_cache.invalidate(); const h = o.h; const n = h.ids.len; try h.best(o.stream(), d.dlogits.ptr, n, n, 1, d.best.ptr); @@ -208,6 +366,8 @@ pub fn Ops(comptime Tri: type) type { /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts, the first always kept. pub fn draft(o: Self, d: *Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?Sampling, confidence: f64, out: *Chain) !void { + d.verify_cache.invalidate(); + if (try @import("cuda_exl3_chain.zig").draft(Tri, o, d, next, rows, position, count, sm, confidence, out)) return; // graphs: the greedy chain as one graph out.n = 0; try o.absorb(d, next, rows, count > 0); var j: usize = 0; @@ -228,7 +388,47 @@ pub fn Ops(comptime Tri: type) type { /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams. pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { - if (o.main) |m| return m.verify(m.ptr, seq, d, ids); + return o.verifyWithTail(seq, d, ids, null); + } + + pub fn verifySampled(ops: Self, seq: *prompt.Seq, window: *Draft, ids: []const u32, sm: ?Sampling) !void { + window.verify_cache.invalidate(); + const greedy = if (sm) |sampling| sampling.temperature <= 0 else true; + const leader = if (ops.main) |main| main.rank == 0 else false; + const setting = std.c.getenv("TF_FLASHNEXT_EXL3_VERIFY_BEST") orelse return ops.verify(seq, window, ids); + if (!greedy or !leader or ops.e.collective == null or !@import("cuda_exl3_graph.zig").enabled() or !std.mem.eql(u8, std.mem.span(setting), "1")) return ops.verify(seq, window, ids); + if (window.verify_host == null) window.verify_host = try cuda.HostBuffer.alloc(ops.e.driver, max_rows * 16); + var tail: VerifyTail = .{ .head = ops.h, .draft = window }; + errdefer window.verify_cache.invalidate(); + try ops.verifyWithTail(seq, window, ids, .{ .ptr = &tail, .enqueue = VerifyTail.enqueue }); + if (tail.enqueued) { + window.verify_pending = null; + window.verify_cache.publish(tail.generation, window.rows, window.fed0); + } + } + + const VerifyTail = struct { + head: *Head, + draft: *Draft, + generation: u64 = 0, + enqueued: bool = false, + + fn enqueue(ptr: *anyopaque, st: cuda.Stream) !void { + const tail: *VerifyTail = @ptrCast(@alignCast(ptr)); + const window = tail.draft; + tail.generation = window.verify_cache.generation; + try tail.head.best(st, window.logits.ptr, vocab, vocab, window.rows, window.best.ptr); + tail.enqueued = true; + window.verify_pending = st; + try window.best.downloadAsync(0, window.verify_host.?.bytes[0 .. window.rows * 16], st.handle); + } + }; + + pub fn verifyWithTail(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, tail: ?batch.Hooks.OutputTail) !void { + const telemetry_scope = o.e.telemetry.enter(.verify); + defer telemetry_scope.restore(); + d.verify_cache.invalidate(); + if (o.main) |m| return m.verify(m.ptr, seq, d, ids, tail); o.e.round_graph.clear(); const n = ids.len; if (n == 0 or n > max_rows) return error.WindowTooWide; @@ -240,7 +440,7 @@ pub fn Ops(comptime Tri: type) type { @memcpy(d.ids[0..n], ids); const keeps = n > 1; if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); - try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); + try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .output_tail = tail }); if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); d.open = keeps; } @@ -249,9 +449,14 @@ pub fn Ops(comptime Tri: type) type { pub fn draws(o: Self, d: *Draft, sm: ?Sampling, out: []u32) !void { const n = d.rows; const greedy = if (sm) |s| s.temperature <= 0 else true; + const cached = d.verify_cache.take(n, d.fed0, greedy); if (greedy) { - try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); - try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + if (cached) { + @memcpy(d.best_host[0 .. 4 * n], d.verify_host.?.slice(f32)[0 .. 4 * n]); + } else { + try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + } for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); return; } @@ -262,6 +467,7 @@ pub fn Ops(comptime Tri: type) type { /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. pub fn rollback(o: Self, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + d.verify_cache.invalidate(); if (o.main) |m| return m.rollback(m.ptr, seq, d, keep); o.e.round_graph.clear(); if (keep == 0 or keep > d.rows) return error.CannotRewind; diff --git a/zig/src/families/flashnext/cuda_mtp_verify_cache.zig b/zig/src/families/flashnext/cuda_mtp_verify_cache.zig new file mode 100644 index 0000000000..c612943b7c --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_verify_cache.zig @@ -0,0 +1,65 @@ +const std = @import("std"); + +pub const Cache = struct { + generation: u64 = 0, + ready: ?u64 = null, + rows: usize = 0, + position: usize = 0, + + pub fn invalidate(cache: *Cache) void { + cache.generation +%= 1; + cache.ready = null; + } + + pub fn publish(cache: *Cache, generation: u64, rows: usize, position: usize) void { + if (generation != cache.generation) return; + cache.rows = rows; + cache.position = position; + cache.ready = generation; + } + + pub fn take(cache: *Cache, rows: usize, position: usize, greedy: bool) bool { + defer cache.invalidate(); + return greedy and cache.ready == cache.generation and cache.rows == rows and cache.position == position; + } +}; + +test "exl3: verification records require completion in the current generation and are consumed once" { + var cache: Cache = .{}; + try std.testing.expect(!cache.take(4, 10, true)); + const generation = cache.generation; + cache.publish(generation, 4, 10); + try std.testing.expect(cache.take(4, 10, true)); + try std.testing.expect(!cache.take(4, 10, true)); + cache.publish(generation, 4, 10); + try std.testing.expect(!cache.take(4, 10, true)); +} + +test "exl3: sampled draws, changed windows, rollback and reuse invalidate verification records" { + var cache: Cache = .{}; + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(4, 10, false)); + try std.testing.expect(!cache.take(4, 10, true)); + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(3, 10, true)); + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(4, 11, true)); + const prior = cache.generation; + cache.invalidate(); + cache.publish(prior, 4, 10); + try std.testing.expect(!cache.take(4, 10, true)); + cache.generation = std.math.maxInt(u64); + cache.publish(cache.generation, 4, 10); + cache.invalidate(); + try std.testing.expect(!cache.take(4, 10, true)); +} + +test "exl3: verification generations belong to each draft independently" { + var first: Cache = .{}; + var second: Cache = .{}; + first.publish(first.generation, 4, 10); + second.publish(second.generation, 8, 20); + first.invalidate(); + try std.testing.expect(!first.take(4, 10, true)); + try std.testing.expect(second.take(8, 20, true)); +} diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 459f56905d..ff01b86153 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const core = @import("core"); const embed = @import("cuda_embed.zig"); const hc = @import("cuda_hc.zig"); +const exl3_ple = @import("cuda_exl3_ple.zig"); const golden: u64 = 0x9E3779B97F4A7C15; const mix1: u64 = 0xBF58476D1CE4E5B9; @@ -341,8 +342,9 @@ pub const Ngram = struct { w: [max_shards][]const u8 = undefined, s: [max_shards][]const u8 = undefined, b: [max_shards][]const u8 = undefined, - format: enum { mlx, nvfp4 } = .mlx, + format: enum { mlx, nvfp4, exl3 } = .mlx, tables: [max_shards]core.nvfp4_rows.Table = undefined, + x3: ?*exl3_ple.Table = null, pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { const cfg = try readSpec(gpa, io, model_dir); @@ -359,6 +361,7 @@ pub const Ngram = struct { .ple_index = cfg.ple_index, }), .scale = 1 }; errdefer g.deinit(); + if (try exl3_ple.Table.open(gpa, io, model_dir, g.hash, exl3_ple.optionsFromEnv())) |t| return .{ .gpa = gpa, .hash = g.hash, .scale = 1, .format = .exl3, .x3 = t }; const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); defer gpa.free(index_path); const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); @@ -414,6 +417,7 @@ pub const Ngram = struct { } pub fn deinit(g: *Ngram) void { + if (g.x3) |t| t.destroy(); for (g.files.items) |*o| o.close(g.gpa); g.files.deinit(g.gpa); } @@ -432,6 +436,7 @@ pub const Ngram = struct { switch (g.format) { .mlx => try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, row), .nvfp4 => try g.tables[loc.shard].gather(&.{loc.local}, row), + .exl3 => try g.x3.?.row(loc.shard, loc.local, i % g.hash.heads, row), } } } @@ -458,7 +463,7 @@ pub const Ngram = struct { hist[history.len - 1] = token; } const max_threads = 128; - const n_threads = @min(max_threads, items / 4); + const n_threads = if (g.format == .exl3 and items <= 1024) 1 else @min(max_threads, items / 4); // EXL3 rows decode in about a microsecond: a decode step's few rows cost less than one thread spawn if (n_threads <= 1) { try g.fill(ids, out, 0, items); } else { @@ -485,7 +490,7 @@ pub const Ngram = struct { v.* = toBf16(f * g.scale); } } - sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); + if (g.x3) |x3| x3.sums(row, xs[t * (w / 32) ..][0 .. w / 32]) else sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); } @memcpy(history, hist[0..history.len]); } @@ -517,4 +522,5 @@ test "a fresh window hashes token 7 into sixteen rows" { test { _ = @import("cuda_ple_nvfp4_test.zig"); + _ = @import("cuda_exl3_ple.zig"); } diff --git a/zig/src/families/flashnext/cuda_projection.zig b/zig/src/families/flashnext/cuda_projection.zig index c0b96adb3b..3e718e6ce0 100644 --- a/zig/src/families/flashnext/cuda_projection.zig +++ b/zig/src/families/flashnext/cuda_projection.zig @@ -10,8 +10,11 @@ pub const Projection = struct { n: usize, k: usize, native: ?cuda.qmmf.Weight = null, + /// An EXL3 checkpoint's projection: stored-layout EXL3 parts (and dense BF16 ones) a row (cuda_exl3_load.zig). + exl3: ?*@import("cuda_exl3_load.zig").Proj = null, pub fn free(self: *Projection) void { + if (self.exl3) |e| e.destroy(); self.w.free(); self.s.free(); self.b.free(); diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig index c0c37f1b79..d602cfefb3 100644 --- a/zig/src/families/flashnext/cuda_projection_load.zig +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); const projection = @import("cuda_projection.zig"); +const exl3 = @import("cuda_exl3_load.zig"); pub const Micro = struct { n: usize, @@ -69,6 +70,7 @@ fn uploadMicro(a: std.mem.Allocator, d: *const cuda.Driver, value: Micro) !wt.Q4 } pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, n: usize, k: usize) !wt.Q4 { + if (exl3.has(store, prefix)) return exl3.load(a, d, store, prefix, n, k); var buf: [512]u8 = undefined; const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); switch (t.dtype) { @@ -96,6 +98,7 @@ pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []c pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !wt.Q4 { if (names.len == 0) return error.UnexpectedTensor; + if (exl3.hasAny(store, prefix, names)) return exl3.stitched(a, d, store, prefix, names, n, k); var buf: [512]u8 = undefined; const first = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, names[0] })); if (first.dtype == .f8_e4m3) { diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 0ff498c22d..dc4ccc2800 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -7,6 +7,8 @@ const gdn = @import("cuda_gdn.zig"); const ple = @import("cuda_ple.zig"); const weights = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); +const kvpage = @import("cuda_exl3_kvpage.zig"); +const shared_kv = @import("cuda_prompt_shared_kv.zig"); pub const batch = @import("cuda_batch.zig"); const fb = @import("cuda_fn_back.zig"); const pocket = @import("cuda_pocket.zig"); @@ -113,6 +115,7 @@ pub fn pooledBytes(cap: usize) usize { } const round_graph = @import("cuda_round_graph.zig"); +const cache_telemetry = @import("cuda_prompt_telemetry.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { @@ -133,6 +136,7 @@ pub const Engine = struct { ngram: ple.Ngram, /// The pass's scratch, grown to the widest chunk run. batch_s: ?batch.Scratch = null, + scratch_format: @import("cuda_batch_state.zig").ScratchFormat = .classic, /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). @@ -140,11 +144,16 @@ pub const Engine = struct { round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, + page_config: ?kvpage.Config = null, + shared_arena: ?*kvpage.shared_device.Arena = null, + telemetry: cache_telemetry.State = .{}, /// Each layer's Living Weights pocket, from a taught folder's shards or the learner: every pass adds it. pockets: pocket.Pockets = .{}, /// During a round of taught and untaught streams: its rows' pocket scales (0 on an untaught row); 0: none. taught_rows: u64 = 0, + const measured = cache_telemetry.Operations(Engine, Seq, batch.Hooks, .{ .forward = forwardUnmeasured, .round = roundUnmeasured }); + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { return initRank(gpa, io, driver, stream, model_dir, 1, 0); } @@ -159,8 +168,12 @@ pub const Engine = struct { e.driver = driver; e.stream = stream; e.n_layers = 0; + e.scratch_format = .classic; e.collective = null; e.image_factory = null; + e.page_config = kvpage.configFromEnv(); + e.shared_arena = null; + e.telemetry = cache_telemetry.State.fromEnv(rank, world); e.round_graph = .{}; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); @@ -180,8 +193,9 @@ pub const Engine = struct { e.eos = try ple.eosOf(gpa, io, e.model_dir); e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); errdefer e.mixer.free(); - const head_tensor = try e.store.tensor("language_model.lm_head.weight"); - if (head_tensor.dtype == .bf16) { + const exl3_head = try @import("cuda_exl3_load.zig").head(gpa, driver, &e.store, world, rank); + const head_tensor = if (exl3_head == null) try e.store.tensor("language_model.lm_head.weight") else undefined; + if (exl3_head) |q| e.head_q = q else if (head_tensor.dtype == .bf16) { if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, if (try @import("cuda_head.zig").enabled()) try @import("cuda_head.zig").shard(head_tensor, world, rank) else head_tensor); } else { @@ -229,11 +243,18 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { + @import("cuda_exl3_prefill.zig").releaseFor(e); + @import("cuda_exl3_graph.zig").releaseFor(e); + @import("cuda_exl3_recycle.zig").drain(e); + if (e.page_config) |cfg| if (cfg.shared_lanes != 0) { + if (e.shared_arena) |arena| arena.destroy() catch @panic("shared KV arena still owns live lanes"); + }; e.round_graph.clear(); e.pockets.deinit(); if (e.batch_s) |*bs| bs.free(e.gpa); if (e.scores) |*b| b.free(); e.batch_k.deinit(); + @import("cuda_exl3_kv.zig").forget(e.gpa, e); e.ngram.deinit(); for (e.layers[0..e.n_layers]) |*l| l.free(e.gpa); e.head_q.free(); @@ -289,15 +310,24 @@ pub const Engine = struct { pub fn newSeq(e: *Engine, cap: usize) !Seq { e.round_graph.clear(); if (cap == 0) return error.UnexpectedTensor; + if (e.page_config) |cfg| if (cfg.shared_lanes != 0) { + const arena = e.shared_arena orelse return error.SharedKvArenaRequired; + if (cap != arena.layout.options.cap) return error.SharedKvCommonCap; + }; + if (try @import("cuda_exl3_recycle.zig").takeSeq(e, cap)) |kept| return kept; + if (@import("cuda_exl3_kv.zig").of(e).placed()) return @import("cuda_prompt_lifecycle.zig").newPlacedSeq(e, cap); const logits = try e.gpa.alloc(u8, vocab * 2); errdefer e.gpa.free(logits); const conv = try zeros(e.gpa, e.driver, e.stream.*, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); errdefer freeOwned(e.gpa, conv); const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); errdefer freeOwned(e.gpa, state); - const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + // K/V paged to host memory (TF_FLASHNEXT_KVPAGE): kc/vc stay one row (cuda_exl3_kvpage.zig holds the rest) + const paged = if (e.page_config) |cfg| if (cfg.shared_lanes != 0) try kvpage.Paged.initShared(e.gpa, e.driver, cap, cfg.forChunk(e.chunk), e.shared_arena orelse return error.SharedKvArenaRequired) else try kvpage.Paged.init(e.gpa, e.driver, cap, cfg.forChunk(e.chunk)) else null; + errdefer if (paged) |pg| pg.destroy(); + const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, (if (paged != null) 1 else cap) * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, kc); - const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, (if (paged != null) 1 else cap) * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, vc); const ikc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); errdefer freeOwned(e.gpa, ikc); @@ -318,6 +348,8 @@ pub const Engine = struct { .fed = 0, .logits = logits, .cap = cap, + .paged = paged, + .recycle_owner = e, }; } @@ -326,9 +358,26 @@ pub const Engine = struct { return e.forward(Tri, tri, seq, ids, .{}); } - fn complete(e: *Engine) !void { - if (e.collective) |c| return c.finish(e.stream.*); - try e.stream.synchronize(); + const complete = measured.complete; + + pub fn reserveSharedKv(e: *Engine, cap: usize) !usize { + return shared_kv.reserve(e, cap); + } + + fn beginShared(e: *Engine, seqs: []const *Seq) !?kvpage.shared_device.Lease { + return shared_kv.begin(e, seqs); + } + + fn endShared(e: *Engine, lease: ?kvpage.shared_device.Lease) void { + shared_kv.end(e, lease); + } + + fn failShared(e: *Engine, lease: ?kvpage.shared_device.Lease, seqs: []const *Seq) void { + shared_kv.fail(e, lease, seqs); + } + + fn finishPaged(e: *Engine, seqs: []const *Seq) !void { + return shared_kv.finish(e, seqs, complete); } /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. @@ -348,19 +397,31 @@ pub const Engine = struct { /// Scratch for chunks of `rows` rows, reallocated only to grow. pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { if (e.batch_s) |*bs| { + if (bs.format != e.scratch_format) return error.ScratchFormatChanged; if (bs.cap >= rows) return bs; try e.complete(); e.round_graph.clear(); bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); + e.batch_s = try batch.Scratch.allocForFormat(e.gpa, e.driver, e.stream.*, rows, e.scratch_format); return &e.batch_s.?; } /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. - pub fn forward(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { - if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + pub const forward = measured.forward; + + fn forwardUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (seq.invalidated) return error.SnapshotDestinationInvalid; + if (ids.len == 0 or seq.fed > seq.cap or ids.len > seq.cap - seq.fed) return error.UnexpectedTensor; + const shared_sequences = [_]*Seq{seq}; + const shared_lease = try e.beginShared(&shared_sequences); + defer e.endShared(shared_lease); + errdefer e.failShared(shared_lease, &shared_sequences); + if (try @import("cuda_exl3_prefill.zig").forward(e, Tri, tri, seq, ids, hooks)) { + const one = [_]*Seq{seq}; + return e.finishPaged(&one); + } if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; const s = try e.reserve(@min(ids.len, e.chunk)); const scores = try e.reserveScores(seq.cap, s.cap); @@ -375,6 +436,10 @@ pub const Engine = struct { .seq = seq.view(scores), .hooks = hooks, }; + if (try @import("cuda_exl3_graph.zig").forward(e, &p, seq, ids)) { + const one = [_]*Seq{seq}; + return e.finishPaged(&one); + } var at: usize = 0; while (at < ids.len) { const n = @min(s.cap, ids.len - at); @@ -384,18 +449,27 @@ pub const Engine = struct { try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none) seq.logits else null); at += n; } - if (hooks.logits == .none) try e.complete(); + if (seq.paged != null) { + const one = [_]*Seq{seq}; + try e.finishPaged(&one); + } else if (hooks.logits == .none) try e.complete(); seq.fed += ids.len; } /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone. - pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { + pub const round = measured.round; + + fn roundUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor; var cap: usize = 0; for (seqs) |q| { + if (q.invalidated) return error.SnapshotDestinationInvalid; if (q.fed + 1 > q.cap) return error.UnexpectedTensor; cap = @max(cap, q.cap); } + const shared_lease = try e.beginShared(seqs); + defer e.endShared(shared_lease); + errdefer e.failShared(shared_lease, seqs); const rows = try e.gpa.alloc(batch.Lane, seqs.len); defer e.gpa.free(rows); const s = try e.reserve(seqs.len); @@ -426,11 +500,15 @@ pub const Engine = struct { const has_media = for (seqs) |q| { if (q.media != null) break true; } else false; - if (e.collective != null and round_graph.enabled() and !has_media) { + const has_paged = for (seqs) |q| { + if (q.paged != null) break true; + } else false; + if (try @import("cuda_exl3_graph.zig").round(e, &p, seqs, ids, scores, out)) {} else if (e.collective != null and round_graph.enabled() and !has_media and !has_paged) { const key = try round_graph.Key.init(seqs, s.ids, scores, out); try e.round_graph.run(key, &p, ids); } else try p.chunk(ids, 0, null); - try e.complete(); + try e.finishPaged(seqs); + if (!has_paged) try e.complete(); const d = e.stream.d; for (seqs, 0..) |q, r| { try d.check(d.api.cuMemcpyDtoH_v2(q.logits.ptr, out + r * vocab * 2, vocab * 2), "cuMemcpyDtoH"); @@ -439,8 +517,19 @@ pub const Engine = struct { } }; +test "exl3: Engine graph-hit lease covers every owner without paging callbacks or extra completion waits" { + try @import("cuda_prompt_test.zig").graphLease(Engine, Seq, .{ + .begin = Engine.beginShared, + .finish = Engine.finishPaged, + .end = Engine.endShared, + .fail = Engine.failShared, + }); +} + /// One stream's caches. `fed` tokens have already gone through the model. pub const Seq = struct { + recycle_owner: ?*Engine = null, + invalidated: bool = false, media: ?@import("cuda_media.zig").State = null, rank_id: u64 = 0, conv: []cuda.DeviceBuffer, @@ -454,30 +543,20 @@ pub const Seq = struct { fed: usize, logits: []u8, cap: usize, + /// int8 K/V scales beside the codes in kc/vc, and the indexer keys' mapped host buffer (cuda_exl3_kv.newSeq). + kv8: ?@import("cuda_exl3_kv.zig").SeqKv8 = null, + ikc_host: ?cuda.HostBuffer = null, + paged: ?*kvpage.Paged = null, /// The model as shipped: no learned neuron acts on its rows (a lesson writer's and its judge's reads). untaught: bool = false, /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). pub fn view(s: *Seq, scores: u64) batch.SeqState { - return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores }; + return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .kv8 = s.kv8, .paged = s.paged }; } pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { - if (s.media) |*m| m.deinit(); - freeBufs(s.conv); - freeBufs(s.state); - freeBufs(s.kc); - freeBufs(s.vc); - freeBufs(s.ikc); - freeBufs(s.pooled); - gpa.free(s.conv); - gpa.free(s.state); - gpa.free(s.kc); - gpa.free(s.vc); - gpa.free(s.ikc); - gpa.free(s.pooled); - if (s.tail.ptr != 0) s.tail.free(); - gpa.free(s.logits); + @import("cuda_prompt_lifecycle.zig").deinit(s, gpa); } }; @@ -491,52 +570,13 @@ fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { } test "linear layers skip every fourth index and attention starts at 3" { - try std.testing.expect(linearLayer(0)); - try std.testing.expect(!linearLayer(3)); - try std.testing.expect(!linearLayer(47)); - try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); - try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); - try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); - try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); - try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); + try @import("cuda_prompt_test.zig").layerIndices(); } test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { - const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); - try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); - try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); - try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); + try @import("cuda_prompt_test.zig").sequenceBytes(seqBytes); } test "a failed fill frees the buffer it was filling and every one before it" { - const Fake = struct { - live: *usize, - pub fn free(b: *@This()) void { - b.live.* -= 1; - } - }; - const Src = struct { - live: *usize, - fills: *usize, - fail_at: usize, - fn make(s: @This()) !Fake { - s.live.* += 1; - return .{ .live = s.live }; - } - fn fill(s: @This(), _: *Fake) !void { - s.fills.* += 1; - if (s.fills.* == s.fail_at) return error.FillFailed; - } - }; - for ([_]usize{ 1, 2, 4 }) |fail_at| { - var live: usize = 0; - var fills: usize = 0; - try std.testing.expectError(error.FillFailed, made(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); - try std.testing.expectEqual(@as(usize, 0), live); - } - var live: usize = 0; - var fills: usize = 0; - const all = try made(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); - defer std.testing.allocator.free(all); - try std.testing.expectEqual(@as(usize, 3), live); + try @import("cuda_prompt_test.zig").failedFill(made); } diff --git a/zig/src/families/flashnext/cuda_prompt_lifecycle.zig b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig new file mode 100644 index 0000000000..aec10b2014 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig @@ -0,0 +1,23 @@ +const std = @import("std"); +const prompt = @import("cuda_prompt.zig"); + +pub fn newPlacedSeq(engine: *prompt.Engine, cap: usize) !prompt.Seq { + var sequence = try @import("cuda_exl3_kv.zig").newSeq(engine, prompt.Seq, cap); + sequence.recycle_owner = engine; + return sequence; +} + +pub fn deinit(sequence: *prompt.Seq, allocator: std.mem.Allocator) void { + if (@import("cuda_exl3_recycle.zig").giveSeq(sequence)) return; + if (sequence.paged != null or sequence.invalidated) sequence.tail.d.synchronize() catch @panic("sequence consumers did not drain before freeing allocations"); + @import("cuda_exl3_graph.zig").releaseSeq(sequence); + if (sequence.media) |*media| media.deinit(); + @import("cuda_exl3_kv.zig").deinitSeq(allocator, &sequence.kv8, &sequence.ikc_host); + for ([_][]@import("cuda").DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }) |buffers| { + for (buffers) |*buffer| buffer.free(); + allocator.free(buffers); + } + if (sequence.tail.ptr != 0) sequence.tail.free(); + if (sequence.paged) |pager| pager.destroy(); + allocator.free(sequence.logits); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 45623e98cc..5d0540bc27 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -7,6 +7,7 @@ const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); const mixed = @import("cuda_mixed_experts.zig"); const projections = @import("cuda_projection_load.zig"); +const exl3_load = @import("cuda_exl3_load.zig"); const layers_n = wt.layers_n; const dims = wt.dims; @@ -80,9 +81,12 @@ pub const Layer = struct { experts_up: cuda.DeviceBuffer, experts_down: cuda.DeviceBuffer, native_experts: ?mixed.Owned = null, + /// EXL3 routed experts with the shared one as id 512 (cuda_exl3_load.zig). + exl3: ?*exl3_load.Experts = null, pub fn free(self: *Layer, gpa: std.mem.Allocator) void { _ = gpa; + if (self.exl3) |x| x.destroy(); self.attn_hc.free(); self.mlp_hc.free(); self.proj.free(); @@ -343,7 +347,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); - if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); defer gpa.free(a_log); @@ -361,8 +365,13 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S @memcpy(pre_buf[0..pre.len], pre); const ap = pre_buf[0..pre.len]; const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; - proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); - out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); + if (try @import("cuda_exl3_heads.zig").loadRank(gpa, driver, store, ap, world, rank)) |head_weights| { // EXL3 tp=2: this rank's heads + proj_q = head_weights[0]; + out_q = head_weights[1]; + } else { + proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); + out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); + } const qn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); defer gpa.free(qn); const kn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); @@ -384,6 +393,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const pp = pre_buf[0..pre.len]; key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); + try exl3_load.pleRows(gpa, driver, store, &key_q, &val_q); const nk = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); defer gpa.free(nk); const nq = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); @@ -395,7 +405,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); - if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); } @@ -405,7 +415,11 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S errdefer experts_up.free(); var experts_down = deadBuf(driver); errdefer experts_down.free(); - if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { + var exl3_experts: ?*exl3_load.Experts = null; + errdefer if (exl3_experts) |x| x.destroy(); + if (exl3_load.hasExperts(store, mlp_root)) { + exl3_experts = try exl3_load.experts(gpa, driver, store, mlp_root, world, rank); + } else if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); } else { if (world != 1 or rank != 0) return error.UnsupportedQuantization; @@ -468,5 +482,6 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S .experts_up = experts_up, .experts_down = experts_down, .native_experts = native_experts, + .exl3 = exl3_experts, }; } diff --git a/zig/src/families/flashnext/cuda_prompt_shared_kv.zig b/zig/src/families/flashnext/cuda_prompt_shared_kv.zig new file mode 100644 index 0000000000..6698163807 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_shared_kv.zig @@ -0,0 +1,61 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const kvpage = @import("cuda_exl3_kvpage.zig"); +const weights = @import("cuda_prompt_wt.zig"); + +pub fn reserve(engine: anytype, cap: usize) !usize { + const cfg = (engine.page_config orelse return 0).forChunk(engine.chunk); + const bytes = try kvpage.sharedBytesFor(cap, cfg); + if (cfg.shared_lanes == 0) return 0; + const placement = @import("cuda_exl3_kv.zig").of(engine); + if (placement.k8 != null or placement.heads != weights.kv_heads) return error.KvPagePlacement; + const options: kvpage.shared_device.Options = .{ .cap = cap, .lanes = cfg.shared_lanes, .window = cfg.window, .rows = cfg.rows, .slots = cfg.slots }; + if (engine.shared_arena) |arena| { + if (!std.meta.eql(arena.layout.options, options)) return error.SharedKvGeometry; + return bytes; + } + engine.shared_arena = try kvpage.shared_device.Arena.init(engine.gpa, engine.driver, options); + return bytes; +} + +pub fn begin(engine: anytype, sequences: anytype) !?kvpage.shared_device.Lease { + const cfg = engine.page_config orelse return null; + if (cfg.shared_lanes == 0) return null; + const arena = engine.shared_arena orelse return error.SharedKvArenaRequired; + if (sequences.len > kvpage.shared_device.max_lanes) return error.SharedKvLeaseOwners; + var owners: [kvpage.shared_device.max_lanes]kvpage.shared_contract.Sequence = undefined; + for (sequences, 0..) |sequence, index| { + const pg = sequence.paged orelse return error.SharedKvSequence; + if (pg.cfg.shared_lanes == 0 or pg.shared_arena != arena) return error.SharedKvSequence; + owners[index] = pg.shared_sequence; + } + const lease = try arena.acquireMany(owners[0..sequences.len]); + errdefer arena.cancel(lease) catch unreachable; + try arena.submittedOn(lease, engine.stream.*); + return lease; +} + +pub fn end(engine: anytype, lease: ?kvpage.shared_device.Lease) void { + if (lease) |held| if (engine.shared_arena.?.active != null) engine.shared_arena.?.completeAfterConsumers(held) catch unreachable; +} + +pub fn fail(engine: anytype, lease: ?kvpage.shared_device.Lease, sequences: anytype) void { + if (lease) |held| { + const arena = engine.shared_arena.?; + if (!arena.active.?.failed) arena.fail(held) catch unreachable; + for (sequences) |sequence| sequence.invalidated = true; + cuda.completion.drain(engine.driver, engine.stream.*, "shared KV consumers did not drain"); + arena.drainFailed(held) catch @panic("shared KV failure still owns consumers"); + } +} + +pub fn finish(engine: anytype, sequences: anytype, comptime complete: anytype) !void { + var any = false; + for (sequences) |q| if (q.paged) |pg| { + try pg.stageStatus(engine.stream.*); + any = true; + }; + if (!any) return; + try complete(engine); + for (sequences) |q| if (q.paged) |pg| try pg.checkStatus(); +} diff --git a/zig/src/families/flashnext/cuda_prompt_telemetry.zig b/zig/src/families/flashnext/cuda_prompt_telemetry.zig new file mode 100644 index 0000000000..d8208c56dd --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_telemetry.zig @@ -0,0 +1,288 @@ +const std = @import("std"); +const cache = @import("cuda_exl3_telemetry.zig"); +const counters = cache.counters; + +pub const State = struct { + enabled: bool = false, + layers_enabled: bool = false, + layer_cursor: counters.LayerCursor(@import("cuda_prompt_wt.zig").layers_n + 1) = .{}, + rank: counters.Rank = .{ .rank = 0, .world = 1 }, + forward_phase: counters.Phase = .unknown, + + pub fn fromEnv(rank: usize, world: usize) State { + const enabled = enabledFromEnv(); + const state: State = .{ .enabled = enabled, .layers_enabled = enabled and layersEnabledFromEnv(), .rank = .{ .rank = rank, .world = world } }; + if (state.enabled) std.log.warn("EXL3-CACHE-SAMPLING v=1 rank={d} world={d} optin_boundary_sync=1 bytes=logical_payload_not_pcie", .{ rank, world }); + return state; + } + + pub fn forwardPhase(state: State, fed: usize) counters.Phase { + return if (state.forward_phase != .unknown) state.forward_phase else if (fed == 0) .prefill else .unknown; + } + + pub fn enter(state: *State, phase: counters.Phase) Scope { + const previous = state.forward_phase; + state.forward_phase = phase; + return .{ .state = state, .previous = previous }; + } +}; + +pub const Scope = struct { + state: *State, + previous: counters.Phase, + + pub fn restore(scope: Scope) void { + scope.state.forward_phase = scope.previous; + } +}; + +pub fn forwardArg(state: State, fed: usize) u64 { + if (!state.enabled) return 0; + return switch (state.forwardPhase(fed)) { + .unknown => 0, + .prefill => 1, + .decode => 2, + .verify => 3, + }; +} + +pub fn phaseFromArg(value: u64) !counters.Phase { + return switch (value) { + 0 => .unknown, + 1 => .prefill, + 2 => .decode, + 3 => .verify, + else => error.BadRankFrame, + }; +} + +pub fn Operations(comptime Engine: type, comptime Seq: type, comptime Hooks: type, comptime underlying: anytype) type { + return struct { + pub fn complete(engine: *Engine) !void { + if (engine.collective) |collective| return collective.finish(engine.stream.*); + try engine.stream.synchronize(); + } + + pub fn forward(engine: *Engine, comptime Tri: type, tri: Tri, sequence: *Seq, ids: []const u32, hooks: Hooks) !void { + if (!engine.telemetry.enabled) return underlying.forward(engine, Tri, tri, sequence, ids, hooks); + const phase = engine.telemetry.forwardPhase(sequence.fed); + const sequences = [_]*Seq{sequence}; + var before = try begin(engine, &sequences, phase, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.forward(engine, Tri, tri, sequence, ids, hooks) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + if (before) |sample| try finish(engine, &sequences, sample, phase, ids.len, complete); + } + + pub fn round(engine: *Engine, comptime Tri: type, tri: Tri, sequences: []const *Seq, ids: []const u32, out: u64) !void { + if (!engine.telemetry.enabled) return underlying.round(engine, Tri, tri, sequences, ids, out); + var before = try begin(engine, sequences, .decode, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.round(engine, Tri, tri, sequences, ids, out) catch |failure| { + failed(engine, .decode, failure); + return failure; + }; + if (before) |sample| try finish(engine, sequences, sample, .decode, sequences.len, complete); + } + }; +} + +pub fn enabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_CACHE_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn layersEnabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_LAYER_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn failed(engine: anytype, phase: counters.Phase, failure: anyerror) void { + std.log.warn("EXL3-CACHE-INVALID v=1 rank={d} world={d} phase={s} reason={s}", .{ engine.telemetry.rank.rank, engine.telemetry.rank.world, @tagName(phase), @errorName(failure) }); +} + +pub fn begin(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype) !?counters.Snapshot { + return beginWithSampler(engine, sequences, phase, complete, cache.sampleAfterConsumers); +} + +pub fn beginWithSampler(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype, comptime sample: anytype) !?counters.Snapshot { + if (!engine.telemetry.enabled) return null; + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + return sample(engine.gpa, engine, null, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return null; + }; +} + +pub fn finish(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype) !void { + return finishWithSampler(engine, sequences, before, phase, rows, complete, cache.sampleAfterConsumers, publish); +} + +pub fn finishWithSampler(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype, comptime sample: anytype, comptime emit: anytype) !void { + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + var after = sample(engine.gpa, engine, null, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return; + }; + defer after.deinit(engine.gpa); + var report = counters.Report.between(before, after, .{ .phase = phase, .model_rows = rows, .steps = 1, .allow_initial_kv_reset = phase == .prefill }) catch |failure| { + failed(engine, phase, failure); + return; + }; + const layers = if (engine.telemetry.layers_enabled) engine.telemetry.layer_cursor.claim(engine.gpa, before, after) catch |failure| { + failed(engine, phase, failure); + return; + } else &.{}; + defer engine.gpa.free(layers); + report.layers = layers; + emit(engine, report) catch |failure| { + failed(engine, phase, failure); + return; + }; +} + +fn publish(engine: anytype, report: counters.Report) !void { + const text = try report.formatAlloc(engine.gpa); + defer engine.gpa.free(text); + std.debug.print("{s}", .{text}); +} + +test "exl3: serving telemetry disabled adds no completion or readback and opt in samples only completed boundaries" { + const Fake = struct { + gpa: std.mem.Allocator, + telemetry: State = .{ .rank = .{ .rank = 1, .world = 2 } }, + completions: u64 = 0, + reads: u64 = 0, + records: u64 = 0, + + fn complete(engine: *@This()) !void { + engine.completions += 1; + } + + fn sample(allocator: std.mem.Allocator, engine: *@This(), _: @TypeOf(null), _: @TypeOf(null), rank: counters.Rank, _: cache.Completion) !counters.Snapshot { + try std.testing.expect(engine.completions > engine.reads); + engine.reads += 1; + return .{ .rank = rank, .scope = 1, .experts = .{ .device_copies = engine.reads }, .kv = try allocator.alloc(counters.KvSample, 0) }; + } + + fn emit(engine: *@This(), report: counters.Report) !void { + try std.testing.expectEqual(counters.Phase.decode, report.interval.phase); + try std.testing.expectEqual(@as(u64, 2), report.interval.model_rows); + try std.testing.expectEqual(@as(u64, 1), report.experts.device_copies); + try std.testing.expectEqual(@as(usize, 1), report.rank.rank); + engine.records += 1; + } + }; + var engine: Fake = .{ .gpa = std.testing.allocator }; + try std.testing.expect((try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)) == null); + try std.testing.expectEqual(@as(u64, 0), engine.completions); + try std.testing.expectEqual(@as(u64, 0), engine.reads); + engine.telemetry.enabled = true; + var first = (try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)).?; + defer first.deinit(engine.gpa); + try finishWithSampler(&engine, null, first, .decode, 2, Fake.complete, Fake.sample, Fake.emit); + try std.testing.expectEqual(@as(u64, 2), engine.completions); + try std.testing.expectEqual(@as(u64, 2), engine.reads); + try std.testing.expectEqual(@as(u64, 1), engine.records); +} + +test "exl3: TP forward metadata has explicit bounded phase mapping independent of token count" { + const rank_engine = @import("cuda_rank_engine.zig"); + const wire = @import("cuda_rank_wire.zig"); + const logits = @import("cuda_prompt.zig").batch.Logits; + const phases = [_]counters.Phase{ .unknown, .prefill, .decode, .verify }; + var sender: State = .{ .enabled = true }; + var worker: State = .{ .enabled = true, .forward_phase = .verify }; + for (phases, 0..) |phase, index| { + sender.forward_phase = phase; + for ([_]u32{ 4, 8192 }) |bytes| { + const header: wire.Header = .{ .kind = .forward, .object = 17, .arg = forwardArg(sender, 16384), .bytes = bytes, .flags = @backingInt(logits.last) }; + try std.testing.expectEqual(@as(u64, index), header.arg); + const frame = try header.encode(); + const received = try wire.Header.decode(&frame); + try rank_engine.shape(received); + { + const scope = worker.enter(try phaseFromArg(received.arg)); + defer scope.restore(); + try std.testing.expectEqual(sender.forwardPhase(16384), worker.forwardPhase(16384)); + } + try std.testing.expectEqual(counters.Phase.verify, worker.forward_phase); + } + } + sender.enabled = false; + sender.forward_phase = .prefill; + try std.testing.expectEqual(@as(u64, 0), forwardArg(sender, 0)); + try rank_engine.shape(.{ .kind = .forward, .object = 17, .bytes = 4 }); + for ([_]u64{ 4, 255, std.math.maxInt(u64) }) |invalid_arg| { + try std.testing.expectError(error.BadRankFrame, phaseFromArg(invalid_arg)); + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .forward, .object = 17, .arg = invalid_arg, .bytes = 4 })); + } + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .verify, .object = 17, .arg = 3, .bytes = 4 })); + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .round, .arg = 1, .bytes = 12 })); +} + +test "exl3: forward and verify phase scopes restore nested caller state on error" { + const Call = struct { + fn fail(state: *State, phase: counters.Phase) !void { + const scope = state.enter(phase); + defer scope.restore(); + try std.testing.expectEqual(phase, state.forward_phase); + return error.ScopedOperationFailed; + } + }; + for ([_]counters.Phase{ .unknown, .prefill, .decode }) |previous| { + var state: State = .{ .forward_phase = previous }; + const parent = state.enter(.prefill); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .verify)); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, try phaseFromArg(1))); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + parent.restore(); + try std.testing.expectEqual(previous, state.forward_phase); + try std.testing.expect(!state.enabled); + } +} + +test "exl3: extracted disabled forward and round wrappers call underlying operations without joins" { + const prompt = @import("cuda_prompt.zig"); + const cuda = @import("cuda"); + const Fake = @import("cuda_exl3_shared_kv_test.zig").Fake; + const Calls = struct { + var forwards: usize = 0; + var rounds: usize = 0; + + fn forward(_: *prompt.Engine, comptime Tri: type, _: Tri, _: *prompt.Seq, _: []const u32, _: prompt.batch.Hooks) !void { + forwards += 1; + } + + fn round(_: *prompt.Engine, comptime Tri: type, _: Tri, _: []const *prompt.Seq, _: []const u32, _: u64) !void { + rounds += 1; + } + }; + Calls.forwards = 0; + Calls.rounds = 0; + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.telemetry = .{}; + engine.gpa = std.testing.allocator; + engine.collective = null; + engine.stream = &stream; + var sequence: prompt.Seq = undefined; + sequence.fed = 9000; + const ops = Operations(prompt.Engine, prompt.Seq, prompt.batch.Hooks, .{ .forward = Calls.forward, .round = Calls.round }); + const syncs = Fake.sync_calls; + try ops.forward(&engine, void, {}, &sequence, &.{123}, .{}); + try ops.round(&engine, void, {}, &.{&sequence}, &.{123}, 0); + try std.testing.expectEqual(@as(usize, 1), Calls.forwards); + try std.testing.expectEqual(@as(usize, 1), Calls.rounds); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} diff --git a/zig/src/families/flashnext/cuda_prompt_test.zig b/zig/src/families/flashnext/cuda_prompt_test.zig new file mode 100644 index 0000000000..badeba9078 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_test.zig @@ -0,0 +1,131 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const kvpage = @import("cuda_exl3_kvpage.zig"); +const weights = @import("cuda_prompt_wt.zig"); + +pub fn graphLease(comptime Engine: type, comptime Seq: type, comptime operations: anytype) !void { + const fixture = @import("cuda_exl3_shared_kv_test.zig"); + var driver = fixture.Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const arena = try kvpage.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 8192, .lanes = 2, .window = 2052, .rows = 1, .slots = 3 }); + defer arena.destroy() catch unreachable; + fixture.Fake.install(arena); + const first = try arena.register(8192); + defer arena.release(first, stream) catch unreachable; + const second = try arena.register(8192); + defer arena.release(second, stream) catch unreachable; + const cfg: kvpage.Config = .{ .window = 2052, .rows = 1, .slots = 3, .shared_lanes = 2 }; + var status_first: [weights.attn_n * 8 * 4]u8 align(16) = undefined; + var status_second: [weights.attn_n * 8 * 4]u8 align(16) = undefined; + var pager_first: kvpage.Paged = undefined; + var pager_second: kvpage.Paged = undefined; + for ([_]*kvpage.Paged{ &pager_first, &pager_second }, [_]kvpage.shared_contract.Sequence{ first, second }, [_][]align(16) u8{ &status_first, &status_second }) |pg, handle, status| { + pg.d = &driver; + pg.cfg = cfg; + pg.shared_arena = arena; + pg.shared_sequence = handle; + pg.status_host = .{ .d = &driver, .bytes = status }; + } + var sequence_first: Seq = undefined; + var sequence_second: Seq = undefined; + sequence_first.paged = &pager_first; + sequence_second.paged = &pager_second; + sequence_first.invalidated = false; + sequence_second.invalidated = false; + const seqs = [_]*Seq{ &sequence_first, &sequence_second }; + var engine: Engine = undefined; + engine.page_config = cfg; + engine.shared_arena = arena; + engine.driver = &driver; + engine.stream = &stream; + engine.collective = null; + var other_driver = fixture.Fake.driver(); + var wrong_stream: cuda.Stream = .{ .d = &other_driver, .handle = null }; + engine.stream = &wrong_stream; + { + defer if (arena.active) |active| arena.cancel(active.lease) catch unreachable; + const sync_before = fixture.Fake.sync_calls; + try std.testing.expectError(error.SharedKvDriver, operations.begin(&engine, &seqs)); + try std.testing.expect(arena.active == null); + try std.testing.expectEqual(sync_before, fixture.Fake.sync_calls); + try std.testing.expect(!sequence_first.invalidated and !sequence_second.invalidated); + } + const recovered = try arena.acquireMany(&.{ first, second }); + try arena.cancel(recovered); + engine.stream = &stream; + const lease = try operations.begin(&engine, &seqs); + try pager_first.validateSharedLease(); + try pager_second.validateSharedLease(); + const graph: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + const before = fixture.Fake.sync_calls; + try graph.launchOn(stream); + try std.testing.expect(fixture.Fake.graph_pending); + try std.testing.expectError(error.SharedKvBusy, arena.acquire(first)); + try operations.finish(&engine, &seqs); + try std.testing.expect(!fixture.Fake.graph_pending); + try std.testing.expectEqual(before + 1, fixture.Fake.sync_calls); + operations.end(&engine, lease); + try std.testing.expectEqual(before + 1, fixture.Fake.sync_calls); + try std.testing.expectEqual(@as(usize, 1), fixture.Fake.graph_calls); + try std.testing.expectError(error.SharedKvLease, pager_first.validateSharedLease()); + const failing = try operations.begin(&engine, &seqs); + fixture.Fake.fail_graph = true; + try std.testing.expectError(error.NotReady, graph.launchOn(stream)); + operations.fail(&engine, failing, &seqs); + try std.testing.expect(sequence_first.invalidated and sequence_second.invalidated); + try std.testing.expect(!fixture.Fake.graph_pending); + try std.testing.expect(arena.active == null); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.acquire(first)); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.acquire(second)); +} + +pub fn failedFill(comptime makeBuffers: anytype) !void { + const Fake = struct { + live: *usize, + pub fn free(buffer: *@This()) void { + buffer.live.* -= 1; + } + }; + const Src = struct { + live: *usize, + fills: *usize, + fail_at: usize, + pub fn make(source: @This()) !Fake { + source.live.* += 1; + return .{ .live = source.live }; + } + pub fn fill(source: @This(), _: *Fake) !void { + source.fills.* += 1; + if (source.fills.* == source.fail_at) return error.FillFailed; + } + }; + for ([_]usize{ 1, 2, 4 }) |fail_at| { + var live: usize = 0; + var fills: usize = 0; + try std.testing.expectError(error.FillFailed, makeBuffers(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); + try std.testing.expectEqual(@as(usize, 0), live); + } + var live: usize = 0; + var fills: usize = 0; + const all = try makeBuffers(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); + defer std.testing.allocator.free(all); + try std.testing.expectEqual(@as(usize, 3), live); +} + +pub fn layerIndices() !void { + try std.testing.expect(weights.linearLayer(0)); + try std.testing.expect(!weights.linearLayer(3)); + try std.testing.expect(!weights.linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), weights.linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), weights.linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), weights.linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), weights.attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), weights.attnIndex(47)); +} + +pub fn sequenceBytes(comptime sequenceSize: anytype) !void { + const per_token = weights.attn_n * (2 * weights.kv_heads * weights.head_dim * 2 + weights.index_dim * 2); + try std.testing.expectEqual(per_token, sequenceSize(6) - sequenceSize(5)); + try std.testing.expectEqual(per_token + weights.attn_n * weights.index_dim * 2, sequenceSize(5) - sequenceSize(4)); + try std.testing.expect(sequenceSize(weights.window) > weights.window * per_token); +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index ec27d953f1..8919af70f6 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -6,6 +6,7 @@ const core = @import("core"); const hc = @import("cuda_hc.zig"); const qmm = @import("cuda_qmm.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3_store = @import("cuda_exl3_store.zig"); pub const layers_n: usize = 48; pub const linear_n: usize = 36; @@ -89,13 +90,20 @@ pub const Store = struct { released: [2]Held = undefined, nreleased: usize = 0, norm_offset: f32 = 0, + /// Dense F16 weights an EXL3 pack stores, widened to BF16 exactly (cuda_exl3_store.zig), freed at release(). + widened: std.ArrayList([]u16) = .empty, + /// The EXL3 kernels and scratch this checkpoint's projections and experts share (cuda_exl3_load.zig), if any. + exl3_ops: ?*@import("cuda_exl3_load.zig").Ops = null, pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); defer gpa.free(index); const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); defer gpa.free(text); - return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + errdefer parsed.deinit(); + try exl3_store.addPatch(io, dir, &parsed); + return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = parsed }; } /// Infer one checkpoint-wide convention from every decoder HC norm, never from the final mixer's outlier. @@ -128,6 +136,7 @@ pub const Store = struct { } pub fn deinit(self: *Store) void { + if (self.exl3_ops) |o| o.release(); self.release(); self.closeReleased(); self.closeFile(); @@ -136,6 +145,8 @@ pub const Store = struct { /// Borrowers expire here; two indexed shards cover separate weight and calibration planes. pub fn release(self: *Store) void { + for (self.widened.items) |w| self.gpa.free(w); + self.widened.clearAndFree(self.gpa); for (self.held[0..self.nheld]) |h| { if (self.nreleased == self.released.len) { self.released[0].file.close(); @@ -162,7 +173,7 @@ pub const Store = struct { const key = try self.resolveKey(requested, &mapped); const file = try self.shard(key); for (self.held[0..self.nheld]) |*h| { - if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + if (std.mem.eql(u8, h.rel, file)) return self.dense(key, h.file.get(key) orelse return error.MissingTensor); } if (self.nheld == self.held.len) return error.UnexpectedTensor; for (self.released[0..self.nreleased], 0..) |cached, i| { @@ -171,7 +182,7 @@ pub const Store = struct { self.nreleased -= 1; if (i < self.nreleased) self.released[i] = self.released[self.nreleased]; self.nheld += 1; - return self.held[self.nheld - 1].file.get(key) orelse return error.MissingTensor; + return self.dense(key, self.held[self.nheld - 1].file.get(key) orelse return error.MissingTensor); } } const rel = try self.gpa.dupe(u8, file); @@ -180,10 +191,10 @@ pub const Store = struct { defer self.gpa.free(path); self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; self.nheld += 1; - return self.held[self.nheld - 1].file.get(key) orelse { + return self.dense(key, self.held[self.nheld - 1].file.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); } pub fn closeFile(self: *Store) void { @@ -234,10 +245,16 @@ pub const Store = struct { var mapped: [512]u8 = undefined; const key = try self.resolveKey(requested, &mapped); try self.use(key); - return self.file.?.get(key) orelse { + return self.dense(key, self.file.?.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); + } + + /// An EXL3 pack's dense F16 weight as the BF16 the loaders read (exact or refused); any other tensor as stored. + fn dense(self: *Store, key: []const u8, t: Tensor) !Tensor { + if (!exl3_store.widens(key, t)) return t; + return exl3_store.widen(self.gpa, key, t, &self.widened); } /// Check the exact or compatibility name without mapping its shard. diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index 35a1650541..24640fab45 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -7,6 +7,7 @@ const wire = @import("cuda_rank_wire.zig"); const Collective = @import("cuda_collective.zig").Collective; const wt = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); +const telemetry = @import("cuda_prompt_telemetry.zig"); pub const Identity = extern struct { config: [32]u8, @@ -20,6 +21,7 @@ pub const Identity = extern struct { head_sharded: u64 = 0, image_tokens: u64 = 0, max_images: u64 = 0, + pager: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -41,7 +43,7 @@ pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usiz defer a.free(index); const program = try std.process.executablePathAlloc(io, a); defer a.free(program); - return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()) }; + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()), .pager = @import("cuda_exl3_kvpage.zig").configIdentity(@import("cuda_exl3_kvpage.zig").configFromEnv()) }; } pub const Session = struct { @@ -58,6 +60,7 @@ pub const Session = struct { drafting: bool = false, stopped: bool = false, available: std.atomic.Value(bool) = .init(true), + restoring: ?struct { sequence: u64, snapshot: u64 } = null, pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity) !*Session { const s = try a.create(Session); @@ -187,7 +190,7 @@ pub const Session = struct { pub fn forward(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) !void { if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.logits == .all) return error.UnsupportedRankOperation; try tokenIds(ids); - const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); + const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .arg = telemetry.forwardArg(s.engine.telemetry, q.fed), .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); errdefer s.poison(); try s.engine.forward(Tri, tri, q, ids, hooks); try s.ack(h); @@ -212,13 +215,18 @@ pub const Session = struct { /// Both ranks keep their own recurrent window; only the leader consumes draft scores and hidden rows. pub fn verify(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) !void { - if (!s.drafting or ids.len == 0 or ids.len > mtp.max_rows or q.fed + ids.len > q.cap) return error.UnsupportedRankOperation; + return s.verifyWithTail(Tri, tri, q, d, ids, null); + } + + pub fn verifyWithTail(session: *Session, comptime Tri: type, tri: Tri, seq: *prompt.Seq, draft_state: *mtp.Draft, ids: []const u32, tail: ?prompt.batch.Hooks.OutputTail) !void { + if (tail != null and session.rank != 0) return error.UnsupportedRankOperation; + if (!session.drafting or ids.len == 0 or ids.len > mtp.max_rows or seq.fed + ids.len > seq.cap) return error.UnsupportedRankOperation; try tokenIds(ids); - const h = try s.command(.{ .kind = .verify, .object = q.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); - errdefer s.poison(); - const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; - try ops.verify(q, d, ids); - try s.ack(h); + const header = try session.command(.{ .kind = .verify, .object = seq.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); + errdefer session.poison(); + const ops: mtp.Ops(Tri) = .{ .e = session.engine, .tri = tri, .h = undefined }; + try ops.verifyWithTail(seq, draft_state, ids, tail); + try session.ack(header); } pub fn keep(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, count: usize) !void { @@ -235,26 +243,48 @@ pub const Session = struct { const item = try s.windows.getOrPut(s.gpa, q.rank_id); if (!item.found_existing) { errdefer _ = s.windows.remove(q.rank_id); - item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk); + // rank 1 never runs the head: its window keeps the rollback buffers and a one-entry head cache, not q.cap's + item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, 1, 1, s.engine.chunk); } return item.value_ptr; } - pub fn save(s: *Session, q: *prompt.Seq) !*snaps.Snap { - const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + /// `d`: rank 0's head cache (a drafting stream); rank 1 keeps the sequence alone (it never runs the head). + pub fn save(s: *Session, q: *prompt.Seq, d: ?*mtp.Draft) !*snaps.Snap { + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, d); errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); kept.rank_id = s.next_object; s.next_object = std.math.add(u64, s.next_object, 1) catch return error.BadRankFrame; const h = try s.command(.{ .kind = .save, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + errdefer s.poison(); try s.ack(h); return kept; } - pub fn restore(s: *Session, q: *prompt.Seq, kept: *const snaps.Snap) !void { + pub fn restore(s: *Session, q: *prompt.Seq, d: ?*mtp.Draft, kept: *const snaps.Snap) !void { + if (kept.paged != null) { + try snaps.validateRestore(q, d, kept); + const check = try s.command(.{ .kind = .restore_check, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + errdefer { + q.invalidated = true; + s.poison(); + } + try s.ack(check); + s.engine.round_graph.clear(); + const copying = try s.command(.{ .kind = .restore, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + try snaps.restoreData(s.engine.stream.*, q, d, kept); + try s.ack(copying); + const commit = try s.command(.{ .kind = .restore_commit, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + try s.ack(commit); + snaps.publishRestore(q, d, kept); + return; + } if (kept.at > q.cap) return error.SequenceTooShort; + // refused before rank 1 hears of it (snaps.restore's own check, made first) + if (d) |hd| if (!kept.head or kept.at > hd.cap) return error.NoHeadState; const h = try s.command(.{ .kind = .restore, .object = q.rank_id, .arg = kept.rank_id }, &.{}); errdefer s.poison(); - try snaps.restore(s.engine.stream.*, q, null, kept); + try snaps.restore(s.engine.stream.*, q, d, kept); try s.ack(h); } @@ -295,6 +325,7 @@ pub const Session = struct { fn apply(s: *Session, comptime Tri: type, tri: Tri, logits: u64, h: wire.Header, payload: []const u8) !bool { try shape(h); + if (s.restoring != null and h.kind != .restore and h.kind != .restore_commit) return error.BadRankFrame; switch (h.kind) { .create => { if (s.seqs.contains(h.object)) return error.BadRankFrame; @@ -315,6 +346,8 @@ pub const Session = struct { defer s.gpa.free(ids); for (ids, 0..) |*v, i| v.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); try tokenIds(ids); + const telemetry_scope = s.engine.telemetry.enter(try telemetry.phaseFromArg(h.arg)); + defer telemetry_scope.restore(); try s.engine.forward(Tri, tri, q, ids, .{ .logits = @fromBackingInt(@intCast(h.flags)) }); }, .round => { @@ -360,6 +393,7 @@ pub const Session = struct { .save => { if (s.saved.contains(h.arg)) return error.BadRankFrame; const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + if (s.windows.getPtr(h.object)) |d| if (d.open) return error.NotAtMark; const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); try s.saved.put(s.gpa, h.arg, kept); @@ -367,7 +401,43 @@ pub const Session = struct { .restore => { const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; - try snaps.restore(s.engine.stream.*, q, null, kept); + s.engine.round_graph.clear(); + if (kept.paged != null) { + const pending = s.restoring orelse return error.BadRankFrame; + if (pending.sequence != h.object or pending.snapshot != h.arg) return error.BadRankFrame; + if (s.windows.getPtr(h.object)) |d| d.verify_cache.invalidate(); + try snaps.restoreData(s.engine.stream.*, q, null, kept); + } else { + try snaps.restore(s.engine.stream.*, q, null, kept); + // a drafting sequence's verify window starts over at the restored mark, as rank 0's head cache does + if (s.windows.getPtr(h.object)) |d| { + d.reset(); + d.len = kept.at; + } + } + }, + .restore_check => { + if (s.restoring != null) return error.BadRankFrame; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; + if (kept.paged == null) return error.BadRankFrame; + try snaps.validateRestore(q, null, kept); + s.restoring = .{ .sequence = h.object, .snapshot = h.arg }; + }, + .restore_commit => { + const pending = s.restoring orelse return error.BadRankFrame; + if (pending.sequence != h.object or pending.snapshot != h.arg) return error.BadRankFrame; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + if (!q.invalidated) return error.BadRankFrame; + const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; + snaps.publishRestore(q, null, kept); + if (s.windows.getPtr(h.object)) |d| { + d.reset(); + d.open = false; + d.len = kept.at; + d.fed0 = kept.at; + } + s.restoring = null; }, .drop_snap => { const entry = s.saved.fetchRemove(h.object) orelse return error.UnknownRankSnapshot; @@ -406,11 +476,14 @@ pub fn shape(h: wire.Header) !void { switch (h.kind) { .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .media => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 64) return error.BadRankFrame, - .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, + .forward => { + _ = try telemetry.phaseFromArg(h.arg); + if (h.object == 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame; + }, .verify => if (h.object == 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 4 != 0 or h.bytes / 4 > mtp.max_rows) return error.BadRankFrame, .keep => if (h.object == 0 or h.arg == 0 or h.arg > mtp.max_rows or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, - .save, .restore => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .save, .restore, .restore_check, .restore_commit => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .drop_seq, .drop_snap => if (h.object == 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .stop => if (h.object != 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, else => return error.BadRankFrame, diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig index a49429688d..3d29eeef7b 100644 --- a/zig/src/families/flashnext/cuda_rank_wire.zig +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -5,7 +5,7 @@ const p = std.posix; const net = std.Io.net; pub const max_payload = 1 << 20; -pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep, media }; +pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep, media, restore_check, restore_commit }; pub const Header = struct { kind: Kind, sequence: u64 = 0, diff --git a/zig/src/families/flashnext/cuda_scratch_layout_test.zig b/zig/src/families/flashnext/cuda_scratch_layout_test.zig new file mode 100644 index 0000000000..a86fa3c09a --- /dev/null +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -0,0 +1,41 @@ +const std = @import("std"); +const Scratch = @import("cuda_batch_state.zig").Scratch; + +test "exl3: scratch omits only classic MoE allocations and keeps every common view" { + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const classic = Scratch.allocationSizes(cap, .classic); + const exl3 = Scratch.allocationSizes(cap, .exl3); + var removed: usize = 0; + var retained: usize = 0; + for (classic, exl3, 0..) |before, after, index| { + if (index >= 37 and index <= 46) { + try std.testing.expectEqual(@as(?usize, null), after); + removed += @max(before.?, 256); + } else { + try std.testing.expectEqual(before, after); + retained += @max(after.?, 256); + } + } + try std.testing.expectEqual(retained, Scratch.bytesForFormat(cap, .exl3)); + try std.testing.expectEqual(removed, Scratch.bytes(cap) - retained); + try std.testing.expectEqual(@as(?usize, cap * 11 * 2560 * 4), exl3[47]); + try std.testing.expectEqual(@as(?usize, cap * 2560 * 4), exl3[49]); + try std.testing.expectEqual(@as(?usize, cap * 128 * 4), exl3[50]); + } +} + +test "exl3: scratch savings include the allocator minimum and grow with requested rows" { + var previous: usize = 0; + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const pairs = cap * 11; + const items = (pairs + 15) / 16 + 513; + const blocks = (pairs + 1023) / 1024; + const expected = 2 * @max(pairs * 4, 256) + @max(items * 12, 256) + 256 + + @max(blocks * 513 * 4, 256) + 3 * @max(cap * 640 * 2, 256) + + @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256); + const saved = Scratch.bytes(cap) - Scratch.bytesForFormat(cap, .exl3); + try std.testing.expectEqual(expected, saved); + try std.testing.expect(saved > previous); + previous = saved; + } +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 0e551809b3..b8823f3a1e 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -1,4 +1,65 @@ //! Flash Next prompt states for the shared prompt cache: one stream's caches after `at` tokens, copied on the device. +//! +//! BF16 paged snapshot format and ordering contract (version 1, exact pager geometry). +//! +//! A paged sequence is not a dense kc/vc allocation. Its kc/vc views alias compact device windows. Dividing their +//! lengths by seq.cap, or copying at * row_bytes from them, cannot describe the logical history. The host store +//! alone is also insufficient: rows still in each device window have not necessarily been written back. +//! +//! Format: keep the dense/int8 format separate, with an explicit version and paging discriminator in Snap. The +//! paged header records at, history, head presence, BF16 row width, attention-layer count, index ratio, cap, +//! window/rows/slots, wr and nblocks, plus checked offsets/lengths for every section. No stored GPU pointers are +//! executable after restore. The initial exact-state format requires equal pager geometry and capacity on the +//! destination; reject a mismatch before touching it. A future capacity-changing format must explicitly re-encode +//! slot2blk's layer * nblocks keys and rebuild both inverse tables, rather than copying those keys unchanged. +//! +//! Payload, at one completed sequence boundary: +//! - Existing convolution, recurrent state, PLE tail, IKC [0,at), complete pooled blocks and optional MTP head +//! state. Include partially pooled raw IKC rows; they are needed when the next block completes. +//! - Canonical BF16 host K/V [layer][K|V][0,at), after flushing every pending window interval [win0[layer],at). +//! - The entire compact device K/V allocation, including all windows and shared slot payloads, and the entire +//! tabs allocation: blk2slot, mark, ctl (including global clock, sticky status and per-layer epoch), slot2blk, +//! stamp. Copying only host history and resetting tables would lose the requested mappings and LRU state. +//! - Host win0, rebases and written arrays. graph_prepared is transient execution state and restores as false. +//! work lists and status_host are scratch: rebuild/overwrite them before use; never restore stale completion +//! status as proof that an operation completed. Logical rows >= at are invalid irrespective of their bytes. +//! +//! Save ordering: +//! 1. Exclude scheduling/replay for this sequence. Establish an event dependency from every producer (including +//! prefill reduction and copy streams) onto the snapshot stream; synchronizing this stream alone is insufficient +//! unless the caller already joined those producers. Reject an outstanding draft or inconsistent head mark. +//! 2. Validate ALL layer ranges, geometry and checked sizes before any copy; require win0 aligned to index_ratio, +//! win0 <= at <= min(cap,win0+wr), and reject a sticky paging failure after a completed status transfer. +//! 3. Allocate the budgeted snapshot, enqueue window writebacks, then canonical host-prefix copies, device K/V, +//! tables and ordinary state on the ordered stream. Snapshot storage owns its copies; it never aliases a live +//! pager. Preserve counters as well as data; flush itself does not advance epochs or rebase windows. +//! 4. Wait for the copies before publishing the immutable snapshot or releasing source/storage ownership. On a +//! submission failure, drain already submitted copies before freeing their targets. If draining fails, retain +//! ownership until the device/session is quiesced. A partial snapshot must never reach the prompt cache. +//! +//! Restore ordering: +//! 1. Validate format, bounds, byte counts, pager geometry, sequence capacity and optional MTP head availability +//! before changing the destination. Drain prior uses and pending graph instantiation before overwriting state. +//! 2. Invalidate MTP verification generation before the first mutation, so previously published or delayed results +//! cannot be consumed. Restore into the destination's own allocations; never transplant snapshot addresses. +//! 3. Copy canonical host rows, compact device payload, tables, IKC/pooled/recurrent/tail/head state, then wait. +//! Restore host origins/counters only after device completion; clear graph_prepared and refresh paging status. +//! Matching tables, slot payloads and clock must become visible together before another attention call. +//! 4. Commit seq.fed/history and reset Draft's transient drafted/rows state and len only after successful completion. +//! On a failure after mutation, discard/poison the destination rather than publishing its previous position over +//! partially restored memory. The snapshot remains immutable and reusable. Destroy snapshots only after their +//! last asynchronous reader completes; dropping a source sequence must not invalidate saved storage. +//! +//! Integration: the engine freezes pager configuration at load; serving lanes use the configured, budgeted capacity. +//! snapBytes charges the canonical prefix, windows/slots, tables, ordinary state, next logits, optional head and header. +//! Snapshots own mapped host storage. The scheduler serializes operations; context completion joins all CUDA producers. +//! cuda_rank_engine uses prepare/copy/commit acknowledgements before returning a restored pair to scheduling. +//! releaseSeq joins pending graph instantiations and discards captures that reference the destination sequence. +//! Failed submissions drain through CUDA's completion policy before freeing storage; an undrainable context terminates +//! the rank through the existing completion handler. A partially restored sequence remains invalid for execution. +//! +//! Acceptance gates are specified in zig/tests/cuda/exl3_kvpage_tests.zig; host fakes alone do not establish GPU or +//! serving correctness. The snapshot command checks device lifecycle; flashnext-exl3-graph snapshot checks TP2 parity. const std = @import("std"); const cuda = @import("cuda"); @@ -6,12 +67,18 @@ const prompt = @import("cuda_prompt.zig"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); const Draft = @import("cuda_mtp_draft.zig").Draft; +const kv = @import("cuda_exl3_kv.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; const state_bytes = wt.state_n * 4; const index_row = wt.index_dim * 2; const tail_bytes = wt.ple_tail_n * wt.wide * 2; +test "exl3: paged snapshot host cases" { + _ = @import("cuda_snap_test.zig"); +} + /// One attention cache's first `at` key, value and indexer-key rows, and the pooled blocks they complete. fn attnBytes(at: usize) usize { return at * (2 * prompt.kv_row_bytes + index_row) + at / wt.index_ratio * index_row; @@ -22,24 +89,58 @@ pub fn bytes(at: usize, head: bool) usize { return wt.linear_n * (conv_bytes + state_bytes) + (wt.attn_n + @intFromBool(head)) * attnBytes(at) + tail_bytes; } +/// `bytes` for a sequence whose K/V rows are `r` (int8 codes and scales, or one KV head a rank: cuda_exl3_kv.zig); +/// the head's cache stays upstream's. +pub fn bytesKv(at: usize, head: bool, r: kv.Rows) usize { + const layer = at * (2 * (r.kv + r.scales) + index_row) + at / wt.index_ratio * index_row; + return wt.linear_n * (conv_bytes + state_bytes) + wt.attn_n * layer + @as(usize, @intFromBool(head)) * attnBytes(at) + tail_bytes; +} + +/// A sequence's K/V rows a token, read from its buffers (cap rows each). +fn seqRows(seq: *const prompt.Seq) kv.Rows { + return .{ .kv = seq.kc[0].len / seq.cap, .scales = if (seq.kv8) |c| c.ks[0].len / seq.cap else 0 }; +} + /// One kept state; `head`: it holds the MTP head's cache over the same rows, which a drafting stream needs. pub const Snap = struct { + version: u32 = 1, rank_id: u64 = 0, buf: cuda.DeviceBuffer, + /// The pinned host memory `buf` borrows (states parked on the host: `on_host`), or null (device memory). + host: ?cuda.HostBuffer = null, at: usize, history: [2]i64, head: bool, + paged: ?pager.Snapshot = null, + owned_blocks: ?pager.shared_contract.OwnedBlocks = null, }; +pub fn pagedBytes(cap: usize, cfg: pager.Config, at: usize, head: bool) !usize { + const layout = try pager.SnapshotLayout.init(cap, cfg, at); + const raw = try std.math.mul(usize, at, wt.attn_n * index_row); + const pooled = try std.math.mul(usize, at / wt.index_ratio, wt.attn_n * index_row); + const ordinary = try std.math.add(usize, wt.linear_n * (conv_bytes + state_bytes) + tail_bytes, try std.math.add(usize, raw, pooled)); + const head_bytes = if (head) try std.math.add(usize, try std.math.add(usize, try std.math.mul(usize, at, 2 * prompt.kv_row_bytes + index_row), try std.math.mul(usize, at / wt.index_ratio, index_row)), if (at == 0) 0 else wt.wide * 2) else 0; + return try std.math.add(usize, try pager.ownedSnapshotBytes(cfg), try std.math.add(usize, @sizeOf(Snap) + prompt.vocabulary * 2, try std.math.add(usize, layout.total, try std.math.add(usize, ordinary, head_bytes)))); +} + +fn payloadBytes(image: pager.Snapshot, head: bool) !usize { + return try pagedBytes(image.layout.cap, image.layout.cfg, image.layout.at, head) - @sizeOf(Snap) - try pager.ownedSnapshotBytes(image.layout.cfg); +} + const Way = enum { save, restore }; const Piece = struct { bufs: []cuda.DeviceBuffer, n: usize }; /// An attention cache's pieces for `at` rows: complete pooled blocks only (the rest pool again as later rows land). fn attnPieces(kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, at: usize) [4]Piece { + return attnPiecesRow(kc, vc, ikc, pooled, at, prompt.kv_row_bytes); +} + +fn attnPiecesRow(kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, at: usize, kv_row: usize) [4]Piece { return .{ - .{ .bufs = kc, .n = at * prompt.kv_row_bytes }, - .{ .bufs = vc, .n = at * prompt.kv_row_bytes }, + .{ .bufs = kc, .n = at * kv_row }, + .{ .bufs = vc, .n = at * kv_row }, .{ .bufs = ikc, .n = at * index_row }, .{ .bufs = pooled, .n = at / wt.index_ratio * index_row }, }; @@ -47,17 +148,25 @@ fn attnPieces(kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.Devi /// Each piece of the sequence (and the head's cache, with `d`) to or from its place in the buffer at `base`. fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, way: Way) !void { + if (seq.paged != null) return error.KvPageSnapshot; // paged K/V (cuda_exl3_kvpage.zig) has no snapshot yet const dr = stream.d; var off: u64 = 0; - var pieces: [11]Piece = undefined; + var pieces: [13]Piece = undefined; pieces[0] = .{ .bufs = seq.conv, .n = conv_bytes }; pieces[1] = .{ .bufs = seq.state, .n = state_bytes }; - pieces[2..6].* = attnPieces(seq.kc, seq.vc, seq.ikc, seq.pooled, at); + const rows = seqRows(seq); + pieces[2..6].* = attnPiecesRow(seq.kc, seq.vc, seq.ikc, seq.pooled, at, rows.kv); pieces[6] = .{ .bufs = (&seq.tail)[0..1], .n = tail_bytes }; var count: usize = 7; + if (seq.kv8) |c| { + // the int8 cache's scales (its codes went with kc/vc) + pieces[7] = .{ .bufs = c.ks, .n = at * rows.scales }; + pieces[8] = .{ .bufs = c.vs, .n = at * rows.scales }; + count = 9; + } if (d) |h| { - pieces[7..11].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); - count = 11; + pieces[count..][0..4].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); + count += 4; } for (pieces[0..count]) |p| for (p.bufs) |b| { if (p.n == 0) continue; @@ -65,41 +174,233 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, p.n, stream.handle), "cuMemcpyDtoDAsync"); off += p.n; }; - std.debug.assert(off == bytes(at, d != null)); + std.debug.assert(off == bytesKv(at, d != null, rows)); +} + +fn movePaged(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap, way: Way) !void { + const image = snap.paged.?; + const pg = seq.paged.?; + try pg.moveSnapshot(stream, snap.buf.ptr, image, way == .save); + var offset = image.layout.total; + const pieces = [_]Piece{ + .{ .bufs = seq.conv, .n = conv_bytes }, + .{ .bufs = seq.state, .n = state_bytes }, + .{ .bufs = seq.ikc, .n = snap.at * index_row }, + .{ .bufs = seq.pooled, .n = snap.at / wt.index_ratio * index_row }, + .{ .bufs = (&seq.tail)[0..1], .n = tail_bytes }, + }; + for (pieces) |piece| for (piece.bufs) |buffer| { + try copyPiece(stream, snap.buf.ptr + offset, buffer.ptr, piece.n, way); + offset += piece.n; + }; + if (snap.head) { + const head_n = attnBytes(snap.at) + (if (snap.at == 0) @as(usize, 0) else wt.wide * 2); + if (d) |head| { + const head_pieces = attnPieces((&head.kc)[0..1], (&head.vc)[0..1], (&head.ikc)[0..1], (&head.pooled)[0..1], snap.at); + for (head_pieces) |piece| for (piece.bufs) |buffer| { + try copyPiece(stream, snap.buf.ptr + offset, buffer.ptr, piece.n, way); + offset += piece.n; + }; + if (snap.at > 0) { + try copyPiece(stream, snap.buf.ptr + offset, if (way == .save) head.prompt_last else head.streams.ptr, wt.wide * 2, way); + offset += wt.wide * 2; + } + } else offset += head_n; + } + std.debug.assert(offset + seq.logits.len == snap.buf.len); +} + +fn copyPiece(stream: cuda.Stream, stored: u64, live: u64, len: usize, way: Way) !void { + if (len == 0) return; + const dst, const src = if (way == .save) .{ stored, live } else .{ live, stored }; + try stream.d.check(stream.d.api.cuMemcpyDtoDAsync_v2(dst, src, len, stream.handle), "cuMemcpyDtoDAsync"); +} + +fn quiesce(stream: cuda.Stream) !void { + try stream.d.synchronize(); +} + +fn validatePagedSeq(seq: *const prompt.Seq, d: ?*Draft, at: usize) !void { + if (seq.cap != seq.paged.?.cap or seq.conv.len != wt.linear_n or seq.state.len != wt.linear_n or seq.ikc.len != wt.attn_n or seq.pooled.len != wt.attn_n or seq.tail.len < tail_bytes) return error.KvPageSnapshotLayout; + const pieces = [_]Piece{ .{ .bufs = seq.conv, .n = conv_bytes }, .{ .bufs = seq.state, .n = state_bytes }, .{ .bufs = seq.ikc, .n = at * index_row }, .{ .bufs = seq.pooled, .n = at / wt.index_ratio * index_row } }; + for (pieces) |piece| for (piece.bufs) |buffer| if (buffer.len < piece.n) return error.KvPageSnapshotLayout; + if (d) |head| { + const head_pieces = attnPieces((&head.kc)[0..1], (&head.vc)[0..1], (&head.ikc)[0..1], (&head.pooled)[0..1], at); + for (head_pieces) |piece| for (piece.bufs) |buffer| if (buffer.len < piece.n) return error.NoHeadState; + } +} + +fn discard(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { + cuda.completion.drain(stream.d, stream, "snapshot copies did not drain"); + if (snap.owned_blocks) |*owned| owned.deinit(gpa); + if (snap.host) |*h| h.free() else snap.buf.free(); + gpa.destroy(snap); } /// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. +/// Kept states go to mapped pinned host memory (copied at ~14 GB/s by the same DtoD copies; the device keeps none of +/// them): set once at load (cuda_lanes_load.zig, TF_FLASHNEXT_SNAP_HOST), and the native host then budgets the prompt +/// cache in host memory (Loaded.host_snaps). +pub var on_host: bool = false; + pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft) !*Snap { const at = seq.fed; - if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; + if (seq.invalidated) return error.SnapshotDestinationInvalid; + if (at > seq.cap) return error.SequenceTooShort; + const shared_lease = try acquireShared(seq); + defer finishShared(seq, shared_lease); + errdefer failShared(stream, seq, shared_lease); + const image = if (seq.paged) |pg| try pg.snapshot(at) else null; + if (d) |h| if (h.len != at or h.drafted != 0 or h.open or at > h.cap) return error.NotAtMark; + if (image != null) { + try validatePagedSeq(seq, d, at); + if (seq.kv8 != null or seq.media != null or seq.logits.len != prompt.vocabulary * 2) return error.KvPageSnapshotLayout; + if (d) |h| if (at > 0 and (h.prompt_last < h.streams.ptr or h.streams.len < wt.wide * 2 or h.prompt_last - h.streams.ptr > h.streams.len - wt.wide * 2)) return error.NotAtMark; + errdefer { + cuda.completion.drain(stream.d, stream, "paged snapshot preparation did not drain"); + seq.invalidated = true; + } + try quiesce(stream); + try seq.paged.?.stageStatus(stream); + try stream.synchronize(); + try seq.paged.?.checkStatus(); + } + const n = if (image) |pg| try payloadBytes(pg, d != null) else bytesKv(at, d != null, seqRows(seq)); + const snap = try allocate(gpa, driver, n, at, seq.history, d != null, image); + errdefer discard(gpa, stream, snap); + errdefer if (image != null) { + seq.invalidated = true; + }; + if (image != null) { + if (shared_lease) |lease| try seq.paged.?.shared_arena.?.submitted(lease); + try seq.paged.?.flush(stream, at); + try movePaged(stream, seq, d, snap, .save); + if (shared_lease) |lease| snap.owned_blocks = try seq.paged.?.shared_arena.?.saveOwnedBlocks(gpa, lease); + } else try move(stream, seq, d, snap.buf.ptr, at, .save); + try stream.synchronize(); + if (image != null) { + const logits = snap.host.?.bytes[n - seq.logits.len .. n]; + if (at == 0) @memset(logits, 0) else @memcpy(logits, seq.logits); + } + return snap; +} + +fn allocate(gpa: std.mem.Allocator, driver: *const cuda.Driver, n: usize, at: usize, history: [2]i64, head: bool, image: ?pager.Snapshot) !*Snap { const snap = try gpa.create(Snap); errdefer gpa.destroy(snap); - snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; - errdefer snap.buf.free(); - try move(stream, seq, d, snap.buf.ptr, at, .save); - try stream.synchronize(); + snap.* = .{ .buf = undefined, .at = at, .history = history, .head = head, .paged = image }; + if (image) |pg| if (pg.layout.cfg.shared_lanes != 0) { + snap.version = 2; + }; + if (on_host or image != null) { + var host = try cuda.HostBuffer.allocMapped(driver, @max(n, 256)); + errdefer host.free(); + snap.buf = .{ .d = driver, .ptr = try host.device(), .len = n, .borrowed = true }; + snap.host = host; + } else snap.buf = try cuda.DeviceBuffer.alloc(driver, n); return snap; } /// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. -pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { +pub fn validateRestore(seq: *const prompt.Seq, d: ?*Draft, snap: *const Snap) !void { + const shared = if (snap.paged) |image| image.layout.cfg.shared_lanes != 0 else false; + if (snap.version != @as(u32, if (shared) 2 else 1)) return error.SnapshotVersion; + if (shared != (snap.owned_blocks != null)) return error.KvPageSnapshotLayout; + if (seq.invalidated) return error.SnapshotDestinationInvalid; if (snap.at > seq.cap) return error.SequenceTooShort; if (d) |h| if (!snap.head or snap.at > h.cap) return error.NoHeadState; - // a state with the head's rows serves a stream without drafts too: those rows stay in the buffer - try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); + if (snap.paged) |image| { + const pg = seq.paged orelse return error.KvPageSnapshotLayout; + try pg.validateSnapshot(image); + if (snap.owned_blocks) |*owned| { + try pg.shared_arena.?.validateOwnedBlocks(owned); + } + try validatePagedSeq(seq, d, snap.at); + if (snap.at != image.layout.at or snap.host == null or snap.buf.len != try payloadBytes(image, snap.head) or snap.host.?.bytes.len < snap.buf.len or seq.kv8 != null or seq.media != null or seq.logits.len != prompt.vocabulary * 2) return error.KvPageSnapshotLayout; + if (d) |h| if (h.streams.len < wt.wide * 2) return error.NoHeadState; + } else if (seq.paged != null) return error.KvPageSnapshotLayout; +} + +pub fn restoreData(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { + try validateRestore(seq, d, snap); + const shared_lease = try acquireShared(seq); + defer finishShared(seq, shared_lease); + errdefer failShared(stream, seq, shared_lease); + if (snap.paged != null) try quiesce(stream); + @import("cuda_exl3_graph.zig").releaseDecodeSeq(seq); + if (d) |h| h.verify_cache.invalidate(); + seq.invalidated = true; + if (shared_lease) |lease| try seq.paged.?.shared_arena.?.submitted(lease); + errdefer cuda.completion.drain(stream.d, stream, "partial snapshot restore did not drain"); + if (shared_lease) |lease| try seq.paged.?.shared_arena.?.restoreOwnedBlocks(lease, &snap.owned_blocks.?); + if (snap.paged != null) try movePaged(stream, seq, d, snap, .restore) else try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); + if (seq.paged) |pg| try pg.stageStatus(stream); try stream.synchronize(); + if (seq.paged) |pg| try pg.checkStatus(); +} + +fn acquireShared(seq: *const prompt.Seq) !?pager.shared_device.Lease { + const pg = seq.paged orelse return null; + if (pg.cfg.shared_lanes == 0) return null; + return try pg.shared_arena.?.acquire(pg.shared_sequence); +} + +fn finishShared(seq: *const prompt.Seq, lease: ?pager.shared_device.Lease) void { + if (lease) |held| { + const arena = seq.paged.?.shared_arena.?; + if (arena.active != null) arena.completeAfterConsumers(held) catch unreachable; + } +} + +fn failShared(stream: cuda.Stream, seq: *prompt.Seq, lease: ?pager.shared_device.Lease) void { + if (lease) |held| { + const arena = seq.paged.?.shared_arena.?; + if (!arena.active.?.submitted) { + arena.cancel(held) catch unreachable; + return; + } + if (!arena.active.?.failed) arena.fail(held) catch unreachable; + seq.invalidated = true; + cuda.completion.drain(stream.d, stream, "shared snapshot failed to drain"); + arena.drainFailed(held) catch @panic("shared snapshot still owns consumers"); + } +} + +pub fn publishRestore(seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) void { + std.debug.assert(seq.invalidated); + if (snap.paged) |image| { + seq.paged.?.publishSnapshot(image); + @memcpy(seq.logits, snap.host.?.bytes[snap.buf.len - seq.logits.len .. snap.buf.len]); + } seq.history = snap.history; seq.fed = snap.at; if (d) |h| { h.reset(); h.len = snap.at; + h.open = false; + h.fed0 = snap.at; + h.verify_pending = null; + if (snap.paged != null) h.prompt_last = if (snap.at == 0) 0 else h.streams.ptr; } + seq.invalidated = false; +} + +pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { + try restoreData(stream, seq, d, snap); + publishRestore(seq, d, snap); } pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { - stream.synchronize() catch {}; - snap.buf.free(); - gpa.destroy(snap); + discard(gpa, stream, snap); +} + +test "a placed state's bytes: int8 codes and scales or one KV head a token, the head's rows upstream's" { + const r = kv.snapRows(true, 2); + try std.testing.expectEqual(@as(usize, 12 * (2 * (512 + 32) + 128 * 2)), bytesKv(1, false, r) - bytesKv(0, false, r)); + try std.testing.expectEqual(bytesKv(1, false, r) - bytesKv(0, false, r) + (2 * 2 * 256 * 2 + 128 * 2), bytesKv(1, true, r) - bytesKv(0, true, r)); + try std.testing.expectEqual(bytes(7, true), bytesKv(7, true, kv.snapRows(false, 2))); + const half = kv.snapRows(false, 1); + try std.testing.expectEqual(@as(usize, 12 * (2 * 256 * 2 + 128 * 2)), bytesKv(1, false, half) - bytesKv(0, false, half)); } test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { diff --git a/zig/src/families/flashnext/cuda_snap_shared_test.zig b/zig/src/families/flashnext/cuda_snap_shared_test.zig new file mode 100644 index 0000000000..b50e6477d2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_snap_shared_test.zig @@ -0,0 +1,265 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const snap = @import("cuda_snap.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); +const prompt = @import("cuda_prompt.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const shared_fixture = @import("cuda_exl3_shared_kv_test.zig"); +const fixture_helpers = @import("cuda_snap_test_fixture.zig"); +const Fake = fixture_helpers.Fake; +const Fixture = fixture_helpers.Fixture; +const sharedFixture = fixture_helpers.sharedFixture; + +test "exl3: native shared snapshot version 2 isolates windows canonical data and remapped entries" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 64, .lanes = 3, .window = 2052, .rows = 256, .slots = 16 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + var source: ?*Fixture = try sharedFixture(&driver, arena); + defer if (source) |fixture| fixture.deinit(); + const foreign = try sharedFixture(&driver, arena); + defer foreign.deinit(); + source.?.seq.fed = 20; + for (0..wt.attn_n) |layer| for ([_]bool{ false, true }) |values| { + const target: [*]u8 = @ptrFromInt(source.?.pg.devRows(layer, values)); + @memset(target[0 .. source.?.pg.wr * pager.row_bytes], if (values) 0xa2 else 0xa1); + const other: [*]u8 = @ptrFromInt(foreign.pg.devRows(layer, values)); + @memset(other[0 .. foreign.pg.wr * pager.row_bytes], if (values) 0xb2 else 0xb1); + }; + shared_fixture.seed(arena, source.?.pg.shared_sequence, 0, 0, 0, 1, 0xc1); + shared_fixture.seed(arena, source.?.pg.shared_sequence, 1, 1, 1, 2, 0xc2); + const saved = try snap.save(std.testing.allocator, &driver, stream, &source.?.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + try std.testing.expectEqual(@as(u32, 2), saved.version); + try std.testing.expectEqual(@as(usize, 2), saved.owned_blocks.?.entries.len); + try std.testing.expect(try snap.pagedBytes(64, source.?.pg.cfg, 20, false) >= saved.buf.len + @sizeOf(snap.Snap) + saved.owned_blocks.?.payload.len + saved.owned_blocks.?.entries.len * @sizeOf(pager.shared_contract.OwnedEntry)); + const original_generation = source.?.pg.shared_sequence.generation; + source.?.deinit(); + source = null; + shared_fixture.seed(arena, foreign.pg.shared_sequence, 2, 2, 0, 3, 0xd1); + const destination = try sharedFixture(&driver, arena); + defer destination.deinit(); + try std.testing.expect(destination.pg.shared_sequence.generation > original_generation); + const pending_graph: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + try pending_graph.launchOn(stream); + try std.testing.expect(shared_fixture.Fake.graph_pending); + try snap.restore(stream, &destination.seq, null, saved); + try std.testing.expect(!shared_fixture.Fake.graph_pending); + try std.testing.expectEqual(@as(usize, 20), destination.seq.fed); + try std.testing.expect(!destination.seq.invalidated); + for (0..wt.attn_n) |layer| for ([_]bool{ false, true }) |values| { + const target: [*]const u8 = @ptrFromInt(destination.pg.devRows(layer, values)); + try std.testing.expect(std.mem.allEqual(u8, target[0 .. destination.pg.wr * pager.row_bytes], if (values) 0xa2 else 0xa1)); + const other: [*]const u8 = @ptrFromInt(foreign.pg.devRows(layer, values)); + try std.testing.expect(std.mem.allEqual(u8, other[0 .. foreign.pg.wr * pager.row_bytes], if (values) 0xb2 else 0xb1)); + try std.testing.expect(std.mem.allEqual(u8, foreign.pg.hostRows(layer, values), 0x80)); + try std.testing.expect(std.mem.allEqual(u8, destination.pg.hostRows(layer, values)[0 .. 20 * pager.row_bytes], if (values) 0xa2 else 0xa1)); + }; + const foreign_lease = try arena.acquire(foreign.pg.shared_sequence); + var foreign_blocks = try arena.saveOwnedBlocks(std.testing.allocator, foreign_lease); + defer foreign_blocks.deinit(std.testing.allocator); + try arena.cancel(foreign_lease); + try std.testing.expectEqual(@as(usize, 1), foreign_blocks.entries.len); + try std.testing.expectEqual(@as(u64, 3), foreign_blocks.entries[0].stamp); + try std.testing.expectEqual(@as(usize, 0), foreign.seq.fed); + try std.testing.expectEqual([2]i64{ 73, 91 }, foreign.seq.history); + try std.testing.expect(std.mem.allEqual(u8, foreign_blocks.payload, 0xd1)); + const again = try snap.save(std.testing.allocator, &driver, stream, &destination.seq, null); + defer snap.drop(std.testing.allocator, stream, again); + try std.testing.expectEqualSlices(u8, saved.host.?.bytes[0..saved.buf.len], again.host.?.bytes[0..again.buf.len]); + try std.testing.expectEqualSlices(u8, saved.owned_blocks.?.payload, again.owned_blocks.?.payload); + for (saved.owned_blocks.?.entries, again.owned_blocks.?.entries) |original_entry, restored_entry| { + try std.testing.expectEqual(original_entry.key, restored_entry.key); + try std.testing.expectEqual(saved.owned_blocks.?.clock - original_entry.stamp, again.owned_blocks.?.clock - restored_entry.stamp); + } + try std.testing.expect(saved.owned_blocks.?.entries[0].source_slot != again.owned_blocks.?.entries[0].source_slot); + try std.testing.expect(!Fake.freed_pending); +} + +test "exl3: native shared snapshot failures poison the destination after submitted work and drain its lease" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 64, .lanes = 3, .window = 2052, .rows = 256, .slots = 16 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + const source = try sharedFixture(&driver, arena); + defer source.deinit(); + const destination = try sharedFixture(&driver, arena); + defer destination.deinit(); + const saved = try snap.save(std.testing.allocator, &driver, stream, &source.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + const before = Fake.submitted; + saved.version = 1; + try std.testing.expectError(error.SnapshotVersion, snap.restore(stream, &destination.seq, null, saved)); + try std.testing.expectEqual(before, Fake.submitted); + try std.testing.expect(!destination.seq.invalidated); + saved.version = 2; + saved.owned_blocks.?.clock = @as(u64, std.math.maxInt(u32)) + 1; + try std.testing.expectError(error.Overflow, snap.restore(stream, &destination.seq, null, saved)); + try std.testing.expectEqual(before, Fake.submitted); + try std.testing.expect(!destination.seq.invalidated); + saved.owned_blocks.?.clock = 0; + Fake.fail_at = Fake.submitted + 2; + try std.testing.expectError(error.NotReady, snap.restore(stream, &destination.seq, null, saved)); + try std.testing.expect(destination.seq.invalidated); + try std.testing.expectEqual(@as(usize, 0), Fake.pending); + try std.testing.expect(arena.active == null); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.acquire(destination.pg.shared_sequence)); + const available = try arena.acquire(source.pg.shared_sequence); + try arena.cancel(available); + try std.testing.expect(!Fake.freed_pending); +} + +test "exl3: shared pager Engine sequence churn retains only the common allocation and reuses lane generations" { + Fake.reset(); + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 64, .lanes = 2, .window = 2052, .rows = 256, .slots = 16 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + const baseline = cuda.usage(false).device; + var engine: prompt.Engine = undefined; + engine.gpa = std.testing.allocator; + engine.driver = &driver; + engine.stream = &stream; + engine.chunk = 256; + engine.eos = 0; + engine.page_config = .{ .window = 2052, .rows = 256, .slots = 16, .shared_lanes = 2 }; + engine.shared_arena = arena; + engine.round_graph = .{}; + try std.testing.expectEqual(arena.layout.common_bytes, try engine.reserveSharedKv(64)); + try std.testing.expectEqual(baseline, cuda.usage(false).device); + var previous_generation: u64 = 0; + for (0..3) |_| { + var sequence = try engine.newSeq(64); + const pg = sequence.paged.?; + try std.testing.expect(pg.dev.borrowed and pg.tabs.borrowed and pg.work.borrowed); + try std.testing.expectEqual(arena.device.ptr, pg.dev.ptr); + try std.testing.expectEqual(@as(usize, 0), pg.deviceBytes()); + try std.testing.expect(pg.shared_sequence.generation > previous_generation); + previous_generation = pg.shared_sequence.generation; + sequence.deinit(std.testing.allocator); + try std.testing.expectEqual(baseline, cuda.usage(false).device); + } + try std.testing.expectError(error.SharedKvCommonCap, engine.newSeq(63)); + try std.testing.expectEqual(baseline, cuda.usage(false).device); +} + +test "exl3: shared pager constructor mapping failures retire registrations without touching foreign slots" { + Fake.reset(); + var driver = Fake.driver(); + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 64, .lanes = 2, .window = 2052, .rows = 256, .slots = 16 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const foreign = try arena.register(64); + defer arena.release(foreign, stream) catch unreachable; + shared_fixture.seed(arena, foreign, 0, 0, 0, 1, 0xb1); + const baseline = cuda.usage(false); + const cfg: pager.Config = .{ .window = 2052, .rows = 256, .slots = 16, .shared_lanes = 2 }; + Fake.fail_mapping = true; + try std.testing.expectError(error.NotReady, pager.Paged.initShared(std.testing.allocator, &driver, 64, cfg, arena)); + Fake.fail_mapping = false; + try std.testing.expectEqual(baseline.device, cuda.usage(false).device); + try std.testing.expectEqual(baseline.host, cuda.usage(false).host); + const available = try arena.acquire(foreign); + var image = try arena.saveOwnedBlocks(std.testing.allocator, available); + defer image.deinit(std.testing.allocator); + try arena.cancel(available); + try std.testing.expectEqual(@as(usize, 1), image.entries.len); + try std.testing.expect(std.mem.allEqual(u8, image.payload, 0xb1)); + const pg = try pager.Paged.initShared(std.testing.allocator, &driver, 64, cfg, arena); + pg.destroy(); + try std.testing.expectEqual(baseline.device, cuda.usage(false).device); + try std.testing.expectEqual(baseline.host, cuda.usage(false).host); +} + +test "exl3: shared snapshot mark preflight cancels its unsubmitted lease without invalidating the source" { + Fake.reset(); + var driver = Fake.driver(); + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, .{ .cap = 64, .lanes = 2, .window = 2052, .rows = 256, .slots = 16 }); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + const source = try sharedFixture(&driver, arena); + defer source.deinit(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var draft: @import("cuda_mtp_draft.zig").Draft = undefined; + draft.len = 1; + const before = Fake.submitted; + try std.testing.expectError(error.NotAtMark, snap.save(std.testing.allocator, &driver, stream, &source.seq, &draft)); + try std.testing.expectEqual(before, Fake.submitted); + try std.testing.expect(!source.seq.invalidated); + try std.testing.expect(arena.active == null); + const lease = try arena.acquire(source.pg.shared_sequence); + try arena.cancel(lease); +} + +test "exl3: shared pager common charge admission and snapshot formats are independently checked" { + const cfg: pager.Config = .{ .window = 2052, .rows = 256, .slots = 1024, .shared_lanes = 8 }; + const cap = 8192; + const common = try pager.sharedBytesFor(cap, cfg); + const layout = try pager.shared_device.Layout.init(.{ .cap = cap, .lanes = 8, .window = cfg.window, .rows = cfg.rows, .slots = cfg.slots }); + try std.testing.expectEqual(layout.common_bytes, common); + try std.testing.expectEqual(@as(usize, 0), try pager.deviceBytesFor(cap, cfg)); + try std.testing.expectEqual(@as(u64, 800), try pager.admissionRoom(10000, 100, cfg)); + try std.testing.expectEqual(@as(u64, 799), try pager.admissionRoom(799, 100, cfg)); + try std.testing.expectError(error.Overflow, pager.admissionRoom(10000, std.math.maxInt(usize), cfg)); + const shared_layout = try pager.SnapshotLayout.init(cap, cfg, 64); + var private_cfg = cfg; + private_cfg.shared_lanes = 0; + const private_layout = try pager.SnapshotLayout.init(cap, private_cfg, 64); + try std.testing.expect(shared_layout.device < private_layout.device); + try std.testing.expect(shared_layout.tables < private_layout.tables); + try std.testing.expect(pager.configIdentity(cfg) != pager.configIdentity(private_cfg)); + try std.testing.expect(try snap.pagedBytes(cap, cfg, 64, false) > shared_layout.total + try pager.ownedSnapshotBytes(cfg)); + var bad = cfg; + bad.rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, pager.sharedBytesFor(cap, bad)); +} + +test "exl3: shared pager local keys bind generations and arena incarnations while TP keys remain logical" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const options: pager.shared_device.Options = .{ .cap = 64, .lanes = 2, .window = 2052, .rows = 256, .slots = 16 }; + const arena = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, options); + defer arena.destroy() catch unreachable; + shared_fixture.Fake.install(arena); + const owner = try arena.register(64); + var pg: pager.Paged = undefined; + pg.cfg = .{ .window = 2052, .rows = 256, .slots = 16, .shared_lanes = 2 }; + pg.cap = 64; + pg.wr = arena.layout.window_rows; + pg.win0 = @splat(0); + pg.dev = arena.device; + pg.tabs = arena.tables; + pg.work = arena.work; + pg.host_dev = 1234; + pg.shared_arena = arena; + pg.shared_sequence = owner; + const local = pg.graphKey(); + const logical = pg.logicalGraphKey(); + try arena.release(owner, stream); + const replacement = try arena.register(64); + defer arena.release(replacement, stream) catch unreachable; + pg.shared_sequence = replacement; + try std.testing.expect(!std.meta.eql(local, pg.graphKey())); + try std.testing.expectEqual(logical, pg.logicalGraphKey()); + const other = try pager.shared_device.Arena.initStorage(std.testing.allocator, &driver, options); + defer other.destroy() catch unreachable; + shared_fixture.Fake.install(other); + const placeholder = try other.register(64); + defer other.release(placeholder, stream) catch unreachable; + const peer = try other.register(64); + defer other.release(peer, stream) catch unreachable; + pg.shared_arena = other; + pg.shared_sequence = peer; + try std.testing.expect(!std.meta.eql(local, pg.graphKey())); + try std.testing.expectEqual(logical, pg.logicalGraphKey()); + pg.shared_sequence = placeholder; + try std.testing.expect(!std.meta.eql(local, pg.graphKey())); + try std.testing.expectEqual(logical, pg.logicalGraphKey()); +} diff --git a/zig/src/families/flashnext/cuda_snap_test.zig b/zig/src/families/flashnext/cuda_snap_test.zig new file mode 100644 index 0000000000..87b844e672 --- /dev/null +++ b/zig/src/families/flashnext/cuda_snap_test.zig @@ -0,0 +1,246 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const snap = @import("cuda_snap.zig"); +const pager = @import("cuda_exl3_kvpage.zig"); +const prompt = @import("cuda_prompt.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const fixture_helpers = @import("cuda_snap_test_fixture.zig"); +const Fake = fixture_helpers.Fake; +const Fixture = fixture_helpers.Fixture; + +test { + _ = @import("cuda_snap_shared_test.zig"); +} + +test "exl3: paged snapshot byte admission is checked and includes fixed payload and header" { + const cfg: pager.Config = .{}; + const layout = try pager.SnapshotLayout.init(200_256, cfg, 200_000); + try std.testing.expect(try snap.pagedBytes(200_256, cfg, 200_000, true) > layout.total + @sizeOf(snap.Snap)); + try std.testing.expectError(error.Overflow, snap.pagedBytes(std.math.maxInt(usize), cfg, 1, false)); + try std.testing.expectError(error.KvPageRows, snap.pagedBytes(3, cfg, 4, false)); + var image: pager.Snapshot = .{ .layout = layout, .win0 = @splat(196_000), .rebases = @splat(3), .written = @splat(8192) }; + try image.validate(); + image.layout.total -= 1; + try std.testing.expectError(error.KvPageSnapshotLayout, image.validate()); +} + +test "exl3: paged snapshot restores distinct allocations and all state without source ownership" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var source: ?*Fixture = try Fixture.init(&driver); + defer if (source) |fixture| fixture.deinit(); + const destination = try Fixture.init(&driver); + defer destination.deinit(); + source.?.seq.fed = 4099; + source.?.pg.win0 = @splat(2048); + source.?.pg.win0[1] = 2052; + source.?.pg.rebases = @splat(8); + source.?.pg.written = @splat(2052); + const controls: [*]u32 = @ptrFromInt(source.?.pg.tabs.ptr + 12 * 2 * source.?.pg.nblocks * 4); + for (0..12) |layer| { + controls[layer * 8] = 73; + controls[layer * 8 + 6] = @intCast(17 + layer); + } + const saved = try snap.save(std.testing.allocator, &driver, stream, &source.?.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + try std.testing.expectEqual(saved.buf.len + @sizeOf(snap.Snap), try snap.pagedBytes(source.?.seq.cap, source.?.pg.cfg, saved.at, false)); + const expected = try std.testing.allocator.dupe(u8, saved.host.?.bytes[0..saved.buf.len]); + defer std.testing.allocator.free(expected); + source.?.deinit(); + source = null; + const device_address = destination.pg.dev.ptr; + for (0..2) |_| { + try snap.restore(stream, &destination.seq, null, saved); + try std.testing.expectEqual(device_address, destination.pg.dev.ptr); + try std.testing.expectEqual(saved.history, destination.seq.history); + try std.testing.expectEqual(saved.at, destination.seq.fed); + try std.testing.expect(!destination.pg.graph_prepared and !destination.seq.invalidated); + const again = try snap.save(std.testing.allocator, &driver, stream, &destination.seq, null); + defer snap.drop(std.testing.allocator, stream, again); + try std.testing.expectEqualDeep(saved.paged.?, again.paged.?); + try std.testing.expectEqualSlices(u8, expected, again.host.?.bytes[0..again.buf.len]); + } + try std.testing.expect(!Fake.freed_pending); +} + +test "exl3: paged snapshot submission failures drain and a partial restore stays invalid" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 4; + const saved = try snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + for (1..125) |failure| { + Fake.submitted = 0; + Fake.fail_at = failure; + fixture.seq.invalidated = false; + const previous = fixture.seq.fed; + try std.testing.expectError(error.NotReady, snap.restore(stream, &fixture.seq, null, saved)); + try std.testing.expect(fixture.seq.invalidated); + try std.testing.expectEqual(previous, fixture.seq.fed); + try std.testing.expectEqual(@as(usize, 0), Fake.pending); + } + Fake.fail_at = null; + fixture.seq.invalidated = false; + var incompatible = saved.*; + incompatible.version = 2; + const submissions = Fake.submitted; + try std.testing.expectError(error.SnapshotVersion, snap.restore(stream, &fixture.seq, null, &incompatible)); + try std.testing.expectEqual(submissions, Fake.submitted); + try std.testing.expect(!fixture.seq.invalidated); + fixture.pg.cfg.slots += 1; + try std.testing.expectError(error.KvPageSnapshotLayout, snap.restore(stream, &fixture.seq, null, saved)); + try std.testing.expectEqual(submissions, Fake.submitted); + try std.testing.expect(!Fake.freed_pending); +} + +test "exl3: paged snapshot preserves MTP head mark and rejects stale verification publication" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 3; + const allocator = fixture.arena.allocator(); + var head: @import("cuda_mtp_draft.zig").Draft = undefined; + head.cap = fixture.seq.cap; + head.kc = try Fixture.buffer(allocator, &driver, head.cap * prompt.kv_row_bytes, 0x80); + head.vc = try Fixture.buffer(allocator, &driver, head.cap * prompt.kv_row_bytes, 0x7F); + head.ikc = try Fixture.buffer(allocator, &driver, head.cap * wt.index_dim * 2, 0xFC); + head.pooled = try Fixture.buffer(allocator, &driver, prompt.pooledBytes(head.cap), 0xA5); + head.streams = try Fixture.buffer(allocator, &driver, 2 * wt.wide * 2, 0xFF); + head.prompt_last = head.streams.ptr + wt.wide * 2; + head.len = 3; + head.drafted = 0; + head.open = false; + head.verify_cache = .{}; + const saved = try snap.save(std.testing.allocator, &driver, stream, &fixture.seq, &head); + defer snap.drop(std.testing.allocator, stream, saved); + head.drafted = 1; + try std.testing.expectError(error.NotAtMark, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, &head)); + head.drafted = 0; + head.open = true; + try std.testing.expectError(error.NotAtMark, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, &head)); + head.open = false; + const generation = head.verify_cache.generation; + head.verify_cache.publish(generation, 3, 3); + try snap.restore(stream, &fixture.seq, &head, saved); + head.verify_cache.publish(generation, 3, 3); + try std.testing.expect(!head.verify_cache.take(3, 3, true)); + try std.testing.expectEqual(@as(usize, 3), head.len); + try std.testing.expect(head.drafted == 0 and head.rows == 0 and !head.open); + try std.testing.expectEqual(head.streams.ptr, head.prompt_last); +} + +test "exl3: paged restore stages data before publishing sequence metadata" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 3; + const saved = try snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + fixture.seq.fed = 7; + fixture.seq.history = .{ 111, 222 }; + try snap.restoreData(stream, &fixture.seq, null, saved); + try std.testing.expectEqual(@as(usize, 7), fixture.seq.fed); + try std.testing.expectEqual([2]i64{ 111, 222 }, fixture.seq.history); + try std.testing.expect(fixture.seq.invalidated); + snap.publishRestore(&fixture.seq, null, saved); + try std.testing.expectEqual(@as(usize, 3), fixture.seq.fed); + try std.testing.expectEqual(saved.history, fixture.seq.history); + try std.testing.expect(!fixture.seq.invalidated); +} + +test "exl3: paged save allocation, mapping and completion errors never publish or free pending copies" { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 3; + var allocator = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, snap.save(allocator.allocator(), &driver, stream, &fixture.seq, null)); + Fake.fail_alloc = true; + try std.testing.expectError(error.NotReady, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null)); + Fake.fail_alloc = false; + Fake.fail_mapping = true; + try std.testing.expectError(error.NotReady, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null)); + Fake.fail_mapping = false; + try std.testing.expect(!fixture.seq.invalidated); + Fake.sync_calls = 0; + Fake.fail_sync_at = 3; + try std.testing.expectError(error.NotReady, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null)); + try std.testing.expect(fixture.seq.invalidated); + try std.testing.expectEqual(@as(usize, 0), Fake.pending); + try std.testing.expect(!Fake.freed_pending); + for (Fake.allocations) |allocation| try std.testing.expect(allocation == null); +} + +test "exl3: paged snapshot preflight sync failures drain and invalidate the sequence" { + for (1..3) |failure_at| { + Fake.reset(); + var driver = Fake.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 3; + Fake.fail_sync_at = failure_at; + try std.testing.expectError(error.NotReady, snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null)); + try std.testing.expect(fixture.seq.invalidated); + try std.testing.expectEqual(@as(usize, 0), Fake.pending); + try std.testing.expect(!Fake.freed_pending); + } +} + +test "exl3: TP2 paged restore prepares both ranks and commits only after copy acknowledgements" { + const wire = @import("cuda_rank_wire.zig"); + const ranks = @import("cuda_rank_engine.zig"); + Fake.reset(); + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + const fixture = try Fixture.init(&driver); + defer fixture.deinit(); + fixture.seq.fed = 4; + fixture.seq.rank_id = 9; + const saved = try snap.save(std.testing.allocator, &driver, stream, &fixture.seq, null); + defer snap.drop(std.testing.allocator, stream, saved); + saved.rank_id = 13; + var engine: prompt.Engine = undefined; + engine.driver = &driver; + engine.stream = &stream; + engine.round_graph = .{}; + for ([_]bool{ false, true }) |reject| { + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = std.c.close(pair[0]); + defer _ = std.c.close(pair[1]); + const peer: wire.Link = .{ .fd = pair[1] }; + var session: ranks.Session = .{ .gpa = std.testing.allocator, .rank = 0, .engine = &engine, .link = .{ .fd = pair[0] }, .collective = undefined }; + session.collective.comm = null; + session.collective.failed = false; + fixture.seq.fed = 8; + fixture.seq.invalidated = false; + for (1..if (reject) @as(usize, 2) else 4) |sequence| try peer.sendHeader(.{ .kind = .ack, .sequence = sequence, .object = 9, .flags = if (reject) 2 else 0 }, wire.deadline(1000)); + const before = Fake.submitted; + if (reject) { + try std.testing.expectError(error.RemoteRank, session.restore(&fixture.seq, null, saved)); + try std.testing.expectEqual(before, Fake.submitted); + try std.testing.expectEqual(@as(usize, 8), fixture.seq.fed); + try std.testing.expect(fixture.seq.invalidated and !session.available.load(.acquire)); + } else { + try session.restore(&fixture.seq, null, saved); + try std.testing.expectEqual(saved.at, fixture.seq.fed); + for ([_]wire.Kind{ .restore_check, .restore, .restore_commit }, 1..) |kind, sequence| { + const command = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectEqual(kind, command.kind); + try std.testing.expectEqual(sequence, command.sequence); + try std.testing.expectEqual(@as(u64, 13), command.arg); + } + } + } +} diff --git a/zig/src/families/flashnext/cuda_snap_test_fixture.zig b/zig/src/families/flashnext/cuda_snap_test_fixture.zig new file mode 100644 index 0000000000..36c3e481ec --- /dev/null +++ b/zig/src/families/flashnext/cuda_snap_test_fixture.zig @@ -0,0 +1,174 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const pager = @import("cuda_exl3_kvpage.zig"); +const prompt = @import("cuda_prompt.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const shared_fixture = @import("cuda_exl3_shared_kv_test.zig"); + +pub const Fake = struct { + const Copy = struct { dst: u64, src: u64, len: usize }; + var queued: [512]Copy = undefined; + pub var pending: usize = 0; + pub var submitted: usize = 0; + pub var fail_at: ?usize = null; + pub var sync_calls: usize = 0; + pub var fail_sync_at: ?usize = null; + pub var fail_alloc: bool = false; + pub var fail_mapping: bool = false; + pub var allocations: [8]?[]align(16) u8 = @splat(null); + pub var freed_pending: bool = false; + + pub fn reset() void { + pending = 0; + submitted = 0; + fail_at = null; + sync_calls = 0; + fail_sync_at = null; + fail_alloc = false; + fail_mapping = false; + freed_pending = false; + } + + fn copy(dst: u64, src: u64, len: usize, _: cuda.abi.Stream) callconv(.c) c_int { + submitted += 1; + if (fail_at == submitted) return cuda.abi.error_not_ready; + queued[pending] = .{ .dst = dst, .src = src, .len = len }; + pending += 1; + return 0; + } + + fn read(dst: ?*anyopaque, src: u64, len: usize, stream: cuda.abi.Stream) callconv(.c) c_int { + return copy(@intFromPtr(dst), src, len, stream); + } + + fn sync(_: cuda.abi.Stream) callconv(.c) c_int { + sync_calls += 1; + if (fail_sync_at == sync_calls) return cuda.abi.error_not_ready; + for (queued[0..pending]) |item| { + const dst: [*]u8 = @ptrFromInt(item.dst); + const src: [*]const u8 = @ptrFromInt(item.src); + @memcpy(dst[0..item.len], src[0..item.len]); + } + pending = 0; + shared_fixture.Fake.graph_pending = false; + return 0; + } + + fn contextSync() callconv(.c) c_int { + return sync(null); + } + + fn alloc(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) c_int { + if (fail_alloc) return cuda.abi.error_not_ready; + for (&allocations) |*slot| if (slot.* == null) { + const bytes = std.testing.allocator.alignedAlloc(u8, .@"16", len) catch return 2; + slot.* = bytes; + out.* = bytes.ptr; + return 0; + }; + return 2; + } + + fn device(out: *u64, address: ?*anyopaque, _: c_uint) callconv(.c) c_int { + if (fail_mapping) return cuda.abi.error_not_ready; + out.* = @intFromPtr(address); + return 0; + } + + fn free(address: ?*anyopaque) callconv(.c) c_int { + freed_pending = freed_pending or pending != 0; + for (&allocations) |*slot| if (slot.*) |bytes| if (@intFromPtr(bytes.ptr) == @intFromPtr(address)) { + std.testing.allocator.free(bytes); + slot.* = null; + return 0; + }; + return 1; + } + + pub fn driver() cuda.Driver { + var driver_ = shared_fixture.Fake.driver(); + driver_.api.cuMemcpyDtoDAsync_v2 = copy; + driver_.api.cuMemcpyDtoHAsync_v2 = read; + driver_.api.cuStreamSynchronize = sync; + driver_.api.cuCtxSynchronize = contextSync; + driver_.api.cuMemHostAlloc = alloc; + driver_.api.cuMemHostGetDevicePointer_v2 = device; + driver_.api.cuMemFreeHost = free; + return driver_; + } +}; + +pub const Fixture = struct { + arena: std.heap.ArenaAllocator, + pg: pager.Paged, + seq: prompt.Seq, + + pub fn buffer(allocator: std.mem.Allocator, driver: *const cuda.Driver, len: usize, value: u8) !cuda.DeviceBuffer { + const bytes = try allocator.alignedAlloc(u8, .@"16", len); + @memset(bytes, value); + return .{ .d = driver, .ptr = @intFromPtr(bytes.ptr), .len = len, .borrowed = true }; + } + + fn buffers(allocator: std.mem.Allocator, driver: *const cuda.Driver, count: usize, len: usize) ![]cuda.DeviceBuffer { + const out = try allocator.alloc(cuda.DeviceBuffer, count); + for (out, 0..) |*item, index| item.* = try buffer(allocator, driver, len, @intCast(index)); + return out; + } + + pub fn init(driver: *const cuda.Driver) !*Fixture { + return initCap(driver, 8192); + } + + pub fn initCap(driver: *const cuda.Driver, cap: usize) !*Fixture { + const fixture = try std.testing.allocator.create(Fixture); + errdefer std.testing.allocator.destroy(fixture); + fixture.arena = .init(std.testing.allocator); + errdefer fixture.arena.deinit(); + const allocator = fixture.arena.allocator(); + const cfg: pager.Config = .{ .window = 2052, .rows = 256, .slots = 16 }; + const layout = try pager.SnapshotLayout.init(cap, cfg, 0); + const host = try allocator.alignedAlloc(u8, .@"16", 24 * cap * pager.row_bytes); + @memset(host, 0x80); + const status = try allocator.alignedAlloc(u8, .@"16", 12 * 8 * 4); + fixture.pg = .{ .gpa = allocator, .d = driver, .k = undefined, .cfg = cfg, .cap = cap, .wr = layout.wr, .nblocks = layout.nblocks, .host = .{ .d = driver, .bytes = host }, .host_dev = @intFromPtr(host.ptr), .status_host = .{ .d = driver, .bytes = status }, .dev = try buffer(allocator, driver, layout.device, 0xFF), .tabs = try buffer(allocator, driver, layout.tables, 0), .work = undefined }; + fixture.seq = .{ + .cap = cap, + .fed = 0, + .history = .{ 73, 91 }, + .paged = &fixture.pg, + .conv = try buffers(allocator, driver, wt.linear_n, (gdn.taps - 1) * gdn.conv_dim * 2), + .state = try buffers(allocator, driver, wt.linear_n, wt.state_n * 4), + .kc = &.{}, + .vc = &.{}, + .ikc = try buffers(allocator, driver, wt.attn_n, cap * wt.index_dim * 2), + .pooled = try buffers(allocator, driver, wt.attn_n, prompt.pooledBytes(cap)), + .tail = try buffer(allocator, driver, wt.ple_tail_n * wt.wide * 2, 0x7F), + .logits = try allocator.alloc(u8, prompt.vocabulary * 2), + }; + @memset(fixture.seq.logits, 0x81); + return fixture; + } + + pub fn deinit(fixture: *Fixture) void { + if (fixture.pg.cfg.shared_lanes != 0) fixture.pg.shared_arena.?.release(fixture.pg.shared_sequence, .{ .d = fixture.pg.d, .handle = null }) catch unreachable; + fixture.arena.deinit(); + std.testing.allocator.destroy(fixture); + } +}; + +pub fn sharedFixture(driver: *const cuda.Driver, arena: *pager.shared_device.Arena) !*Fixture { + const fixture = try Fixture.initCap(driver, arena.layout.options.cap); + errdefer fixture.deinit(); + const sequence = try arena.register(arena.layout.options.cap); + fixture.pg.cfg.shared_lanes = arena.layout.options.lanes; + fixture.pg.shared_arena = arena; + fixture.pg.shared_sequence = sequence; + fixture.pg.dev = arena.device; + fixture.pg.dev.borrowed = true; + fixture.pg.tabs = arena.tables; + fixture.pg.tabs.borrowed = true; + fixture.pg.work = arena.work; + fixture.pg.work.borrowed = true; + return fixture; +} diff --git a/zig/src/families/flashnext/exl3.zig b/zig/src/families/flashnext/exl3.zig new file mode 100644 index 0000000000..228ba3637f --- /dev/null +++ b/zig/src/families/flashnext/exl3.zig @@ -0,0 +1,39 @@ +//! Host-side EXL3 index for Qwen3.8-Flash-Next: the float64 oracles and header rules kept as tests under option C +//! (docs/plans/exl3-on-upstream-1.1.0.md). CUDA, Metal and the engine stay out of this file. +//! `exl3_format.zig` in this directory is a symlink to `zig/src/core/exl3_format.zig`: `zig test` of a file here uses +//! this directory as the module root, so a `../` import of the core decoder is rejected. +//! Run: `zig test zig/src/families/flashnext/exl3.zig`. + +pub const shape = @import("exl3_shape.zig"); +pub const scan = @import("exl3_scan.zig"); +pub const tp = @import("exl3_tp.zig"); +pub const cache = @import("exl3_cache.zig"); +pub const cpu = @import("exl3_cpu.zig"); +pub const place = @import("exl3_place.zig"); +pub const route = @import("exl3_route.zig"); +pub const router = @import("exl3_router.zig"); +pub const kv = @import("exl3_kv.zig"); +pub const norm = @import("exl3_norm.zig"); +pub const hc = @import("exl3_hc.zig"); +pub const moe = @import("exl3_moe.zig"); +pub const gdn_cpu = @import("exl3_gdn_cpu.zig"); +pub const attn_cpu = @import("exl3_attn_cpu.zig"); +pub const fwd_cpu = @import("exl3_fwd_cpu.zig"); + +test { + _ = shape; + _ = scan; + _ = tp; + _ = cache; + _ = cpu; + _ = place; + _ = route; + _ = router; + _ = kv; + _ = norm; + _ = hc; + _ = moe; + _ = gdn_cpu; + _ = attn_cpu; + _ = fwd_cpu; +} diff --git a/zig/src/families/flashnext/exl3_attn_cpu.zig b/zig/src/families/flashnext/exl3_attn_cpu.zig new file mode 100644 index 0000000000..4e429ab484 --- /dev/null +++ b/zig/src/families/flashnext/exl3_attn_cpu.zig @@ -0,0 +1,462 @@ +//! CPU oracle for one Qwen3.8-Flash-Next full-attention layer (layers 3, 7, ... 47), written against the model's own definition: +//! `qwen4_exp/model_layers.py` (`SparseAttention`, `Indexer`, `CenteredRMSNorm`) and the Python family's `glue.attn_prep` +//! and `attention._pool/_scores/_select`, read from the TensorFold Python family on par1-cs13. +//! +//! What was checked against that definition (and what each test below pins): +//! - q_proj has 2 * 24 * 256 outputs and is reshaped [heads, 2 * 256]: per head the first 256 are the query, the last 256 the +//! output gate (interleaved per head, not two halves). `prepare` reads q at `h * 512` and `attend` reads gate at `h * 512 + 256`. +//! - q_norm / k_norm / indexer q_layernorm / k_layernorm: fp32 RMSNorm, eps 1e-6, scale `1 + w` for centred storage (offset from +//! `exl3_scan.centredOffset`, the caller adds it), result rounded to bf16. +//! - RoPE is `mx.fast.rope(traditional=False)`: rotate-half over the first `rotary_dim = 0.25 * 256 = 64` dims, pairs (i, i + 32), +//! angle `pos * theta^(-2 i / 64)`, theta 1e7. mrope_interleaved sections [11, 11, 10] pick the t/h/w position per frequency index +//! (`i % 3 == 1 && i < 33` is h, `i % 3 == 2 && i < 30` is w, else t); for text every axis is the same 1D position, so the angle +//! is `pos * inv_freq[i]` (`mropeAxisCounts` below pins the 11/11/10 partition). +//! - Output gate: `bf16(bf16(o) / (1 + exp(-gate)))` before o_proj (fp16 here: the EXL3 linear reads fp16 and bf16 values fit). +//! - Indexer: `index_qk_proj` is 5 * 128 outputs, 4 query heads then one key. The key is stored raw (bf16). Blocks of 4 keys are +//! pooled (fp32 mean -> bf16 -> k_layernorm -> RoPE at the block's first position 4 * b -> bf16). Queries get q_layernorm + RoPE at +//! their own position. Score = sum over heads of relu(q . pooled) / sqrt(128). A row with at most 512 complete blocks (end <= 2051) +//! attends densely; otherwise it attends to the 512 best blocks (ties to the lower block id) and the 0..3 unfinished tail keys. +//! - Attention scale is `256 ** -0.5`; GQA maps q head h to kv head h / 12. +//! Arithmetic is f64 with the model's rounding points (bf16 after norm, after RoPE, after o; fp16 where a linear reads or writes). + +const std = @import("std"); +const shape = @import("exl3_shape.zig"); + +pub const q_heads: usize = shape.heads; +pub const kv_heads: usize = shape.kv_heads; +pub const head_dim: usize = shape.head_dim; +pub const rotary_dim: usize = 64; +pub const half: usize = rotary_dim / 2; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const ratio: usize = 4; +pub const top_blocks: usize = 512; +pub const theta: f64 = 1.0e7; +pub const eps: f64 = 1.0e-6; +/// Most keys one row ever reads: 512 selected blocks and a 3-key tail, or a dense row of up to 2051 keys. +pub const max_keys: usize = top_blocks * ratio + ratio; +pub const index_width: usize = (index_heads + 1) * index_dim; + +pub fn roundBf16(x: f64) f64 { + const f: f32 = @floatCast(x); + var u: u32 = @bitCast(f); + if (!std.math.isFinite(f)) return f; + // round to nearest even on the low 16 bits + const lsb = (u >> 16) & 1; + u +%= 0x7fff + lsb; + u &= 0xffff0000; + return @as(f32, @bitCast(u)); +} + +pub fn roundF16(x: f64) f64 { + const h: f16 = @floatCast(x); + return h; +} + +/// theta^(-2 i / rotary_dim) as the f32 table the kernels read. +pub fn invFreq() [half]f32 { + var t: [half]f32 = undefined; + for (&t, 0..) |*v, i| v.* = @floatCast(std.math.pow(f64, theta, -2.0 * @as(f64, @floatFromInt(i)) / @as(f64, rotary_dim))); + return t; +} + +pub const Norms = struct { + /// Scales with the centred offset already added: `w + offset`. + q: []const f32, + k: []const f32, + iq: []const f32, + ik: []const f32, + inv: [half]f32, +}; + +/// fp32 RMSNorm, bf16, rotate-half RoPE over the first 64 dims at `pos`, bf16. `x` and `out` are `width` long. +pub fn normRope(x: []const f64, scale: []const f32, pos: f64, inv: *const [half]f32, out: []f64) void { + const width = x.len; + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + var xn: [head_dim]f64 = undefined; + for (x, 0..) |v, i| xn[i] = roundBf16(v * rinv * @as(f64, scale[i])); + const posf: f32 = @floatCast(pos); + for (0..width) |d| { + if (d < half) { + const ang: f32 = posf * inv[d]; + out[d] = roundBf16(xn[d] * @cos(@as(f64, ang)) - xn[d + half] * @sin(@as(f64, ang))); + } else if (d < 2 * half) { + const ang: f32 = posf * inv[d - half]; + out[d] = roundBf16(xn[d - half] * @sin(@as(f64, ang)) + xn[d] * @cos(@as(f64, ang))); + } else out[d] = xn[d]; + } +} + +/// How many of the 32 frequency indices each mrope axis owns with `mrope_interleaved` sections [11, 11, 10]. +pub fn mropeAxisCounts() [3]usize { + var c: [3]usize = .{ 0, 0, 0 }; + for (0..half) |i| { + if (i % 3 == 1 and i < 3 * 11) { + c[1] += 1; + } else if (i % 3 == 2 and i < 3 * 10) { + c[2] += 1; + } else c[0] += 1; + } + return c; +} + +/// One full-attention layer's state: bf16-rounded keys and values, raw index keys, pooled blocks. +pub const Cache = struct { + gpa: std.mem.Allocator, + tokens: usize, + keys: []f64, // [pos][kv head][256] + values: []f64, + raw: []f64, // [pos][128] + pooled: []f64, // [pos / 4][128] + + pub fn init(gpa: std.mem.Allocator, tokens: usize) !Cache { + const keys = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(keys); + const values = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(values); + const raw = try gpa.alloc(f64, tokens * index_dim); + errdefer gpa.free(raw); + const pooled = try gpa.alloc(f64, (tokens / ratio + 1) * index_dim); + return .{ .gpa = gpa, .tokens = tokens, .keys = keys, .values = values, .raw = raw, .pooled = pooled }; + } + + pub fn deinit(c: *Cache) void { + c.gpa.free(c.keys); + c.gpa.free(c.values); + c.gpa.free(c.raw); + c.gpa.free(c.pooled); + } +}; + +/// Projections of one token (what the EXL3 linears emit, fp16): q_proj [24 * 512], k_proj [512], v_proj [512], index_qk_proj [640]. +pub const Proj = struct { + qp: []const f16, + kp: []const f16, + vp: []const f16, + ip: []const f16, +}; + +/// The normed, rotated queries (bf16 values, [24 * 256]) and indexer queries ([4 * 128]) of one token. Writes this token's +/// K, V and raw index key into the cache. +pub fn prepare(c: *Cache, pos: usize, p: Proj, n: *const Norms, q_out: []f64, iq_out: []f64) void { + var buf: [head_dim]f64 = undefined; + for (0..q_heads) |h| { + for (0..head_dim) |d| buf[d] = p.qp[h * 2 * head_dim + d]; + normRope(buf[0..head_dim], n.q, @floatFromInt(pos), &n.inv, q_out[h * head_dim ..][0..head_dim]); + } + for (0..kv_heads) |j| { + for (0..head_dim) |d| buf[d] = p.kp[j * head_dim + d]; + normRope(buf[0..head_dim], n.k, @floatFromInt(pos), &n.inv, c.keys[(pos * kv_heads + j) * head_dim ..][0..head_dim]); + for (0..head_dim) |d| c.values[(pos * kv_heads + j) * head_dim + d] = roundBf16(p.vp[j * head_dim + d]); + } + for (0..index_heads) |h| { + for (0..index_dim) |d| buf[d] = p.ip[h * index_dim + d]; + normRope(buf[0..index_dim], n.iq, @floatFromInt(pos), &n.inv, iq_out[h * index_dim ..][0..index_dim]); + } + for (0..index_dim) |d| c.raw[pos * index_dim + d] = roundBf16(p.ip[index_heads * index_dim + d]); +} + +/// Pool block `b` (keys 4b .. 4b + 3 must be in `raw`). +pub fn poolBlock(c: *Cache, b: usize, n: *const Norms) void { + var mean: [index_dim]f64 = undefined; + for (0..index_dim) |d| { + var acc: f32 = @floatCast(c.raw[(ratio * b) * index_dim + d]); + for (1..ratio) |k| acc = acc + @as(f32, @floatCast(c.raw[(ratio * b + k) * index_dim + d])); + mean[d] = roundBf16(@as(f64, acc / @as(f32, ratio))); + } + normRope(&mean, n.ik, @floatFromInt(ratio * b), &n.inv, c.pooled[b * index_dim ..][0..index_dim]); +} + +/// Score of every complete block for one indexer query: sum_h relu(q_h . pooled_b) / sqrt(128). +pub fn blockScores(c: *const Cache, iq: []const f64, complete: usize, out: []f64) void { + for (0..complete) |b| { + var total: f64 = 0; + for (0..index_heads) |h| { + var dot: f64 = 0; + for (0..index_dim) |d| dot += iq[h * index_dim + d] * c.pooled[b * index_dim + d]; + total += @max(dot, 0); + } + out[b] = total / @sqrt(@as(f64, index_dim)); + } +} + +/// The `top_blocks` best blocks, ties to the lower id, in ascending block order. `scores.len` is the complete-block count. +pub fn selectBlocks(gpa: std.mem.Allocator, scores: []const f64, out: *[top_blocks]usize) !void { + const order = try gpa.alloc(usize, scores.len); + defer gpa.free(order); + for (order, 0..) |*o, i| o.* = i; + const Ctx = struct { + s: []const f64, + fn less(ctx: @This(), a: usize, b: usize) bool { + if (ctx.s[a] != ctx.s[b]) return ctx.s[a] > ctx.s[b]; + return a < b; + } + }; + std.mem.sort(usize, order, Ctx{ .s = scores }, Ctx.less); + @memcpy(out, order[0..top_blocks]); + std.mem.sort(usize, out, {}, std.sort.asc(usize)); +} + +/// Keys a row at `end = pos + 1` reads. `blocks` is the selection for a sparse row (ignored when dense). +/// Dense rows read keys 0 .. end; sparse rows read 4 keys per selected block, then the tail. Returns the key count. +pub fn keyList(end: usize, blocks: ?*const [top_blocks]usize, out: *[max_keys]usize) usize { + const complete = end / ratio; + if (complete <= top_blocks) { + for (0..end) |i| out[i] = i; + return end; + } + const sel = blocks.?; + var n: usize = 0; + for (sel) |b| { + for (0..ratio) |k| { + out[n] = b * ratio + k; + n += 1; + } + } + var t = ratio * complete; + while (t < end) : (t += 1) { + out[n] = t; + n += 1; + } + return n; +} + +/// out[h * 256 + d] = fp16(bf16(bf16(o) / (1 + exp(-gate)))), o = softmax(q . K / 16) V over `keys`. +pub fn attend(c: *const Cache, q: []const f64, qp: []const f16, keys: []const usize, out: []f64) void { + var s: [max_keys]f64 = undefined; + const group = q_heads / kv_heads; + for (0..q_heads) |h| { + const kvh = h / group; + var m: f64 = -std.math.inf(f64); + for (keys, 0..) |key, j| { + var dot: f64 = 0; + for (0..head_dim) |d| dot += q[h * head_dim + d] * c.keys[(key * kv_heads + kvh) * head_dim + d]; + s[j] = dot * 0.0625; + m = @max(m, s[j]); + } + var l: f64 = 0; + for (keys, 0..) |_, j| { + s[j] = @exp(s[j] - m); + l += s[j]; + } + for (0..head_dim) |d| { + var acc: f64 = 0; + for (keys, 0..) |key, j| acc += s[j] * c.values[(key * kv_heads + kvh) * head_dim + d]; + const o = roundBf16(acc / l); + const g: f64 = qp[h * 2 * head_dim + head_dim + d]; + out[h * head_dim + d] = roundF16(roundBf16(o / (1.0 + @exp(-g)))); + } + } +} + +pub const RowInfo = struct { + sparse: bool = false, + nk: usize = 0, + /// The selected blocks of a sparse row (undefined when dense). + blocks: [top_blocks]usize = undefined, +}; + +/// The whole row: prepare, pool the blocks this token completes, select, attend. `scores` is scratch for tokens / 4 blocks. +pub fn row( + c: *Cache, + pos: usize, + p: Proj, + n: *const Norms, + scores: []f64, + q: []f64, + iq: []f64, + out: []f64, + info: *RowInfo, +) !void { + prepare(c, pos, p, n, q, iq); + const end = pos + 1; + if (end % ratio == 0) poolBlock(c, end / ratio - 1, n); + const complete = end / ratio; + info.sparse = complete > top_blocks; + if (info.sparse) { + blockScores(c, iq, complete, scores); + try selectBlocks(c.gpa, scores[0..complete], &info.blocks); + } + var keys: [max_keys]usize = undefined; + info.nk = keyList(end, if (info.sparse) &info.blocks else null, &keys); + attend(c, q, p.qp, keys[0..info.nk], out); +} + +// -- tests ----------------------------------------------------------------- + +test "mrope sections [11, 11, 10] interleave over 32 frequencies and need no per-axis table for text" { + const counts = mropeAxisCounts(); + try std.testing.expectEqual([3]usize{ 11, 11, 10 }, counts); + try std.testing.expectEqual(@as(usize, 32), counts[0] + counts[1] + counts[2]); + try std.testing.expectEqual(half, counts[0] + counts[1] + counts[2]); +} + +test "rotate-half RoPE matches an explicit rotation of the pairs (i, i + 32) and leaves dims 64.. alone" { + const inv = invFreq(); + try std.testing.expectApproxEqRel(@as(f32, 1.0), inv[0], 1e-7); + // theta^(-2*31/64) for the last frequency + try std.testing.expectApproxEqRel(@as(f32, @floatCast(std.math.pow(f64, 1e7, -62.0 / 64.0))), inv[31], 1e-6); + var x: [head_dim]f64 = undefined; + var rng = std.Random.DefaultPrng.init(7); + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var scale: [head_dim]f32 = @splat(1); + scale[3] = 2.5; + var out: [head_dim]f64 = undefined; + const pos: usize = 1234; + normRope(&x, &scale, @floatFromInt(pos), &inv, &out); + + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / head_dim + eps); + for (0..half) |i| { + const a = roundBf16(x[i] * rinv * scale[i]); + const b = roundBf16(x[i + half] * rinv * scale[i + half]); + const ang: f32 = @as(f32, @floatFromInt(pos)) * inv[i]; + try std.testing.expectEqual(roundBf16(a * @cos(@as(f64, ang)) - b * @sin(@as(f64, ang))), out[i]); + try std.testing.expectEqual(roundBf16(a * @sin(@as(f64, ang)) + b * @cos(@as(f64, ang))), out[i + half]); + } + for (rotary_dim..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); + // position 0 is the identity rotation + normRope(&x, &scale, 0, &inv, &out); + for (0..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); +} + +test "bf16 rounding is round-to-nearest-even" { + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 512.0)); // below half an ulp (ulp 2^-7) + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 256.0)); // tie goes to even (1.0) + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), roundBf16(1.0 + 3.0 / 256.0)); // tie goes up to even mantissa + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 128.0), roundBf16(1.0 + 1.0 / 128.0 + 1.0 / 1024.0)); +} + +test "selection takes the 512 best blocks, ties to the lower id, in block order" { + const gpa = std.testing.allocator; + const n: usize = 700; + const scores = try gpa.alloc(f64, n); + defer gpa.free(scores); + var rng = std.Random.DefaultPrng.init(11); + for (scores) |*v| v.* = @floor(rng.random().float(f64) * 40) / 4; // many exact ties + var sel: [top_blocks]usize = undefined; + try selectBlocks(gpa, scores, &sel); + try std.testing.expect(std.sort.isSorted(usize, &sel, {}, std.sort.asc(usize))); + // the cut: every unselected block scores no more than every selected one, and a tie at the cut keeps the lower id + var chosen: [700]bool = @splat(false); + var min_sel: f64 = std.math.inf(f64); + for (sel) |b| { + chosen[b] = true; + min_sel = @min(min_sel, scores[b]); + } + var count: usize = 0; + for (chosen) |c| count += @intFromBool(c); + try std.testing.expectEqual(top_blocks, count); + var max_out: f64 = -1; + var lowest_out_at_cut: usize = n; + for (0..n) |b| if (!chosen[b]) { + max_out = @max(max_out, scores[b]); + if (scores[b] == min_sel) lowest_out_at_cut = @min(lowest_out_at_cut, b); + }; + try std.testing.expect(max_out <= min_sel); + for (0..n) |b| if (chosen[b] and scores[b] == min_sel) try std.testing.expect(b < lowest_out_at_cut); +} + +test "key list: dense through 2051 keys, then 2048 selected plus the unfinished tail" { + var sel: [top_blocks]usize = undefined; + for (&sel, 0..) |*b, i| b.* = i + 1; // blocks 1..512 of 513 + var keys: [max_keys]usize = undefined; + try std.testing.expectEqual(@as(usize, 2051), keyList(2051, null, &keys)); + try std.testing.expectEqual(@as(usize, 2050), keys[2050]); + // end 2052: 513 complete blocks, no tail, block 0 dropped + try std.testing.expectEqual(@as(usize, 2048), keyList(2052, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 4), keys[0]); + try std.testing.expectEqual(@as(usize, 2051), keys[2047]); + // end 2055: three tail keys 2052, 2053, 2054 + try std.testing.expectEqual(@as(usize, 2051), keyList(2055, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 2052), keys[2048]); + try std.testing.expectEqual(@as(usize, 2054), keys[2050]); +} + +test "attend with one visible key returns that value gated" { + const gpa = std.testing.allocator; + var c = try Cache.init(gpa, 4); + defer c.deinit(); + for (0..kv_heads * head_dim) |i| { + c.keys[i] = 0.25; + c.values[i] = if (i < head_dim) 2.0 else -3.0; + } + var q: [q_heads * head_dim]f64 = @splat(0.5); + var qp: [q_heads * 2 * head_dim]f16 = @splat(0); + for (0..q_heads) |h| for (0..head_dim) |d| { + qp[h * 2 * head_dim + head_dim + d] = 0; // sigmoid(0) = 1/2 + }; + var out: [q_heads * head_dim]f64 = undefined; + attend(&c, &q, &qp, &.{0}, &out); + for (0..12) |h| try std.testing.expectEqual(@as(f64, 1.0), out[h * head_dim]); + for (12..q_heads) |h| try std.testing.expectEqual(@as(f64, -1.5), out[h * head_dim]); +} + +test "the device KV layout is exl3_kv.Table's for one layer" { + const kv = @import("exl3_kv.zig"); + const gpa = std.testing.allocator; + const tokens: usize = 8; + var t = try kv.Table.init(gpa, tokens, .bf16, 0, 1); + defer t.deinit(gpa); + try std.testing.expectEqual(@as(usize, kv_heads), t.heads); + const slot_bytes = 2 * head_dim * 2; + for ([_]usize{ 0, 1, 2 }) |layer| for ([_]usize{ 0, 5, 7 }) |token| for ([_]usize{ 0, 1 }) |head| { + var src: [slot_bytes]u8 = undefined; + @memset(&src, @intCast(1 + layer * 40 + token * 4 + head)); + try t.write(head, layer, token, &src); + // the kernel's byte offset: layer base + (token * kv_heads + head) * slot + const off = layer * tokens * kv_heads * slot_bytes + (token * kv_heads + head) * slot_bytes; + try std.testing.expectEqualSlices(u8, &src, t.bytes[off..][0..slot_bytes]); + }; +} + +test "a full row past 2048 keys reads the selected blocks, and the dense row at 2051 keys does not" { + const gpa = std.testing.allocator; + const tokens: usize = 2100; + var c = try Cache.init(gpa, tokens); + defer c.deinit(); + var rng = std.Random.DefaultPrng.init(3); + const inv = invFreq(); + const q_scale: [head_dim]f32 = @splat(1); + const i_scale: [index_dim]f32 = @splat(1); + const norms = Norms{ .q = &q_scale, .k = &q_scale, .iq = &i_scale, .ik = &i_scale, .inv = inv }; + const qp = try gpa.alloc(f16, q_heads * 2 * head_dim); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, index_width); + defer gpa.free(ip); + const scores = try gpa.alloc(f64, tokens / ratio + 1); + defer gpa.free(scores); + var q: [q_heads * head_dim]f64 = undefined; + var iq: [index_heads * index_dim]f64 = undefined; + var out: [q_heads * head_dim]f64 = undefined; + var info = RowInfo{}; + for (0..2056) |pos| { + for (qp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (kp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (vp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (ip) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + const proj = Proj{ .qp = qp, .kp = kp, .vp = vp, .ip = ip }; + switch (pos) { + 2050, 2051, 2054 => { + try row(&c, pos, proj, &norms, scores, &q, &iq, &out, &info); + switch (pos) { + 2050 => try std.testing.expect(!info.sparse and info.nk == 2051), + 2051 => try std.testing.expect(info.sparse and info.nk == 2048), + else => try std.testing.expect(info.sparse and info.nk == 2051), + } + }, + else => { + prepare(&c, pos, proj, &norms, &q, &iq); + if ((pos + 1) % ratio == 0) poolBlock(&c, (pos + 1) / ratio - 1, &norms); + }, + } + } + for (out) |v| try std.testing.expect(std.math.isFinite(v)); +} diff --git a/zig/src/families/flashnext/exl3_cache.zig b/zig/src/families/flashnext/exl3_cache.zig new file mode 100644 index 0000000000..a9d5b9f5df --- /dev/null +++ b/zig/src/families/flashnext/exl3_cache.zig @@ -0,0 +1,424 @@ +//! KV quant and the resident-set budget. +//! `layerBytes` is v0.6.6 `kvcache.py` `row_bytes` times tokens: K and V, scales included. +//! Groups of 32 go through one normalized 32-point Hadamard. Dequant stays in that domain. +//! Routed experts and the n-gram table are host-only; they are not part of the device sum. +//! `resumeAt` is the prefix rule. A later step calls it from the prompt cache. + +const std = @import("std"); +const shape = @import("exl3_shape.zig"); + +pub const Dtype = enum { bf16, int8, int4 }; + +const group: usize = 32; +const inv_sqrt32: f32 = 1.0 / @sqrt(32.0); + +pub fn layerBytes(tokens: usize, kv_heads: usize, head_dim: usize, dtype: Dtype) error{HeadDim}!u64 { + const t: u64 = @intCast(tokens); + const h: u64 = @intCast(kv_heads); + const d: u64 = @intCast(head_dim); + switch (dtype) { + .bf16 => return t * 2 * h * d * 2, + .int8, .int4 => { + if (head_dim % group != 0) return error.HeadDim; + const scales = 2 * h * (d / group) * 2; + const width: u64 = if (dtype == .int8) d else d / 2; + return t * (2 * h * width + scales); + }, + } +} + +/// A kept prefix stops one token before the prompt end. An empty prompt stays at 0. +pub fn resumeAt(prompt_len: u32) u32 { + if (prompt_len == 0) return 0; + return prompt_len - 1; +} + +/// In-place FWHT butterflies on 32 lanes, then one multiply by 1/sqrt(32). +fn hadamard(x: *[group]f32) void { + var lo: usize = 1; + while (lo <= 16) : (lo *= 2) { + const hi = group / (2 * lo); + var nxt: [group]f32 = undefined; + for (0..hi) |h| { + for (0..lo) |i| { + const a = x[(h * 2 + 0) * lo + i]; + const b = x[(h * 2 + 1) * lo + i]; + nxt[(h * 2 + 0) * lo + i] = a + b; + nxt[(h * 2 + 1) * lo + i] = a - b; + } + } + x.* = nxt; + } + for (x) |*v| v.* *= inv_sqrt32; +} + +pub fn rotate32(x: []const f32, out: []f32) error{Length}!void { + if (x.len != group or out.len != group) return error.Length; + var cur: [group]f32 = undefined; + @memcpy(cur[0..], x); + hadamard(&cur); + @memcpy(out, cur[0..]); +} + +fn groupScale(rot: []const f32) f32 { + var peak: f32 = 0; + for (rot) |v| peak = @max(peak, @abs(v)); + return peak + 1e-10; +} + +pub fn quantize8(x: []const f32, codes: []i8, scale: *f16) error{Length}!void { + if (x.len != group or codes.len != group) return error.Length; + var rot: [group]f32 = undefined; + try rotate32(x, &rot); + const s = groupScale(&rot); + for (rot, codes) |v, *c| { + const q = std.math.clamp(@floor(v / s * 128.0) + 128.0, 0.0, 255.0); + const as_i: i32 = @intFromFloat(q); + c.* = @intCast(as_i - 128); + } + scale.* = @floatCast(s); +} + +pub fn dequant8(codes: []const i8, scale: f16, out: []f32) error{Length}!void { + if (codes.len != group or out.len != group) return error.Length; + const s: f32 = @floatCast(scale); + for (codes, out) |c, *o| { + o.* = (@as(f32, @floatFromInt(c)) + 0.5) * s * (1.0 / 128.0); + } +} + +/// Low nibble is the even code, high nibble the odd code. `codes` are 32 values in 0..15. +pub fn packNibbles(codes: []const u8, out: []u8) error{Length}!void { + if (codes.len != group or out.len != group / 2) return error.Length; + for (out, 0..) |*byte, i| { + const lo = codes[i * 2] & 0xf; + const hi = codes[i * 2 + 1] & 0xf; + byte.* = lo | (hi << 4); + } +} + +pub fn unpackNibbles(bytes: []const u8, out: []u8) error{Length}!void { + if (bytes.len != group / 2 or out.len != group) return error.Length; + for (bytes, 0..) |byte, i| { + out[i * 2] = byte & 0xf; + out[i * 2 + 1] = byte >> 4; + } +} + +pub fn quantize4(x: []const f32, bytes: []u8, scale: *f16) error{Length}!void { + if (x.len != group or bytes.len != group / 2) return error.Length; + var rot: [group]f32 = undefined; + try rotate32(x, &rot); + const s = groupScale(&rot); + var codes: [group]u8 = undefined; + for (rot, &codes) |v, *c| { + const q = std.math.clamp(@floor(v / s * 8.0) + 8.0, 0.0, 15.0); + c.* = @intFromFloat(q); + } + try packNibbles(&codes, bytes); + scale.* = @floatCast(s); +} + +pub fn dequant4(bytes: []const u8, scale: f16, out: []f32) error{Length}!void { + if (bytes.len != group / 2 or out.len != group) return error.Length; + var codes: [group]u8 = undefined; + try unpackNibbles(bytes, &codes); + const s: f32 = @floatCast(scale); + for (codes, out) |q, *o| { + o.* = (@as(f32, @floatFromInt(q)) - 7.5) * s * 0.125; + } +} + +pub const Budget = struct { + device_bytes: u64, + /// Dense + shared + router + HC + MTP + graph workspace, already summed by the caller. + resident_bytes: u64, + kv_bytes: u64, + routed_expert_bytes: u64, + ngram_bytes: u64, +}; + +pub const Place = enum { host, refuse }; + +/// Resident weights plus KV. Routed experts and the n-gram table are not in this sum. +pub fn deviceSum(b: Budget) u64 { + return b.resident_bytes + b.kv_bytes; +} + +pub fn fits(b: Budget) error{ExceedsDevice}!void { + if (deviceSum(b) > b.device_bytes) return error.ExceedsDevice; +} + +/// One RTX 5060 Ti. The learned expert lines are whatever this does not already use. +pub const card_bytes: u64 = 16 << 30; +/// CUDA context and the batch-1 workspace. Not a weight, and not a second copy of the bank. +pub const device_reserve: u64 = 256 << 20; + +/// Bytes left for routed experts after pinned GEMMs, KV, and the reserve. +/// A zero result means the card is already full. The n-gram table is not an argument. +pub fn learnedBytes(device: u64, pinned: u64, kv: u64) u64 { + if (pinned >= device) return 0; + const after_pinned = device - pinned; + if (kv >= after_pinned) return 0; + const after_kv = after_pinned - kv; + if (after_kv <= device_reserve) return 0; + return after_kv - device_reserve; +} + +fn matrixBytes(k: usize, n: usize) u64 { + const kn: u64 = @as(u64, k) * @as(u64, n); + // 3.05 bpw trellis, plus one fp16 scale along K and one along N. + return kn * 3050 / 8000 + @as(u64, k + n) * 2; +} + +/// One routed expert as this rank stores it. Two ranks keep half of each matrix. +/// That half is a cut helper, not the bytes reserved for one resident expert. +pub fn expertBytes(ranks: usize) u64 { + const k = shape.hidden; + const n = shape.expert_intermediate; + if (ranks == 2) return matrixBytes(k / 2, n) * 2 + matrixBytes(n, k / 2); + return matrixBytes(k, n) * 2 + matrixBytes(n, k); +} + +pub fn expertLines(room: u64, ranks: usize) u64 { + return expertLinesOf(room, expertBytes(ranks)); +} + +/// Budget stride for one owned expert until a header span is larger. Not a line count. +pub const full_expert_bytes: u64 = 1_901_213; + +/// Lines that fit in `room` at `line` bytes. A zero stride is no lines, not a panic. +pub fn expertLinesOf(room: u64, line: u64) u64 { + if (line == 0) return 0; + return room / line; +} + +/// KV for every full-attention layer. GDN layers have no KV table. +/// Both heads: one rank keeps the whole head count. +pub fn contextKv(tokens: usize, dtype: Dtype) error{ HeadDim, Ranks }!u64 { + return contextKvRank(tokens, dtype, 1); +} + +/// This rank's KV. Heads split evenly. A zero rank, or a count that does not +/// divide the head count, is refused instead of charging zero bytes. +pub fn contextKvRank(tokens: usize, dtype: Dtype, ranks: usize) error{ HeadDim, Ranks }!u64 { + if (ranks == 0 or shape.kv_heads % ranks != 0) return error.Ranks; + const layers: u64 = shape.layers / shape.full_attention_interval; + const heads = shape.kv_heads / ranks; + return (try layerBytes(tokens, heads, shape.head_dim, dtype)) * layers; +} + +pub const KvFit = struct { + dtype: Dtype, + /// False when even int4 misses `expert_floor`. Not a dtype to start on. + floor_held: bool, +}; + +/// Widest KV that still leaves `expert_floor` bytes for learned experts. +/// Sized with one head per rank. bf16 is kept while it fits. `.int4` with +/// `floor_held` false missed the floor; the caller drops lines instead of +/// starting on int4. +pub fn kvFit(device: u64, pinned: u64, tokens: usize, expert_floor: u64, ranks: usize) error{ HeadDim, Ranks }!KvFit { + const order = [_]Dtype{ .bf16, .int8, .int4 }; + for (order) |dtype| { + const kv = try contextKvRank(tokens, dtype, ranks); + if (learnedBytes(device, pinned, kv) >= expert_floor) return .{ .dtype = dtype, .floor_held = true }; + } + return .{ .dtype = .int4, .floor_held = false }; +} + +/// Routed experts and the n-gram table always live on the host. +pub fn placeWeights(b: Budget) Place { + _ = b; + return .host; +} + +/// Neither the routed-expert bank nor the n-gram table may be counted as resident. +pub fn refusesResident(which: enum { routed, ngram }) bool { + _ = which; + return true; +} + +fn roundTrip8(x: []const f32) !void { + var codes: [group]i8 = undefined; + var scale: f16 = undefined; + try quantize8(x, &codes, &scale); + var back: [group]f32 = undefined; + try dequant8(&codes, scale, &back); + var rot: [group]f32 = undefined; + try rotate32(x, &rot); + const s: f32 = @floatCast(scale); + const limit = (0.5 * s / 128.0) + s * 0.01; + for (back, rot) |got, want| try std.testing.expect(@abs(got - want) <= limit); +} + +test "constant and ramp groups stay inside half an int8 bin of the rotation" { + var constant: [group]f32 = @splat(0.75); + try roundTrip8(&constant); + + var negative: [group]f32 = @splat(-1.5); + try roundTrip8(&negative); + + var ramp: [group]f32 = undefined; + for (&ramp, 0..) |*v, i| v.* = @as(f32, @floatFromInt(i)) * 0.17 - 2.4; + try roundTrip8(&ramp); + + var rot: [group]f32 = undefined; + try rotate32(&constant, &rot); + try std.testing.expectApproxEqAbs(rot[0], 0.75 * @sqrt(32.0), 1e-4); + for (rot[1..]) |v| try std.testing.expectApproxEqAbs(v, 0, 1e-5); +} + +test "nibbles pack low then high" { + var codes: [group]u8 = @splat(0); + codes[0] = 1; + codes[1] = 2; + codes[2] = 15; + codes[3] = 0; + var bytes: [group / 2]u8 = undefined; + try packNibbles(&codes, &bytes); + try std.testing.expectEqual(@as(u8, 0x21), bytes[0]); + try std.testing.expectEqual(@as(u8, 0x0f), bytes[1]); + + var back: [group]u8 = undefined; + try unpackNibbles(&bytes, &back); + try std.testing.expectEqualSlices(u8, &codes, &back); + + var x: [group]f32 = undefined; + for (&x, 0..) |*v, i| v.* = @as(f32, @floatFromInt(i)) - 8.0; + var qpacked: [group / 2]u8 = undefined; + var scale: f16 = undefined; + try quantize4(&x, &qpacked, &scale); + var qcodes: [group]u8 = undefined; + try unpackNibbles(&qpacked, &qcodes); + try std.testing.expectEqual(qcodes[0], qpacked[0] & 0xf); + try std.testing.expectEqual(qcodes[1], qpacked[0] >> 4); + + var restored: [group]f32 = undefined; + try dequant4(&qpacked, scale, &restored); + var rot: [group]f32 = undefined; + try rotate32(&x, &rot); + const s: f32 = @floatCast(scale); + const limit = (0.5 * s / 8.0) + s * 0.01; + for (restored, rot) |got, want| try std.testing.expect(@abs(got - want) <= limit); +} + +test "resumeAt stops one token before the end" { + try std.testing.expectEqual(@as(u32, 0), resumeAt(0)); + try std.testing.expectEqual(@as(u32, 7), resumeAt(8)); + try std.testing.expectEqual(@as(u32, 0), resumeAt(1)); +} + +test "layer bytes match v0.6.6 row_bytes" { + try std.testing.expectEqual(@as(u64, 2048), try layerBytes(1, 2, 256, .bf16)); + try std.testing.expectEqual(@as(u64, 1088), try layerBytes(1, 2, 256, .int8)); + try std.testing.expectEqual(@as(u64, 576), try layerBytes(1, 2, 256, .int4)); + try std.testing.expectEqual(@as(u64, 2048), try layerBytes(1, shape.kv_heads, shape.head_dim, .bf16)); + try std.testing.expectEqual(@as(u64, 1088), try layerBytes(1, shape.kv_heads, shape.head_dim, .int8)); + try std.testing.expectEqual(@as(u64, 576), try layerBytes(1, shape.kv_heads, shape.head_dim, .int4)); + try std.testing.expectEqual(@as(u64, 6144), try layerBytes(3, 2, 256, .bf16)); + try std.testing.expectEqual(@as(u64, 192), try layerBytes(1, 2, 24, .bf16)); + try std.testing.expectError(error.HeadDim, layerBytes(1, 2, 24, .int8)); + try std.testing.expectError(error.HeadDim, layerBytes(1, 2, 24, .int4)); +} + +test "device sum is resident plus kv, and a 16 GiB card refuses the overflow" { + const device: u64 = 16 << 30; + const resident: u64 = 6 << 30; + const kv: u64 = 8 << 30; + const held = Budget{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = kv, + .routed_expert_bytes = 70 << 30, + .ngram_bytes = 29 << 30, + }; + try std.testing.expectEqual(resident + kv, deviceSum(held)); + try fits(held); + try std.testing.expectEqual(Place.host, placeWeights(held)); + try std.testing.expect(refusesResident(.routed)); + try std.testing.expect(refusesResident(.ngram)); + + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = 11 << 30, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + const rest = device - resident; + try fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest, + .routed_expert_bytes = 1, + .ngram_bytes = 1, + }); + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest + 1, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + try std.testing.expectEqual(@as(u64, 1_893_120), expertBytes(1)); + try std.testing.expectEqual(@as(u64, 948_480), expertBytes(2)); + try std.testing.expectEqual(@as(u64, 1_901_213), full_expert_bytes); + try std.testing.expectEqual(@as(u64, 0), expertLinesOf(1, 0)); + try std.testing.expectEqual(@as(u64, 1), expertLinesOf(full_expert_bytes, full_expert_bytes)); + + // One head leaves 672 MiB at bf16, so 700 MiB steps to int8. 741 MiB is + // this pin's int4 room; one more MiB misses, and that int4 is not a start dtype. + const ranks: usize = 2; + const tight_pin: u64 = 15 << 30; + try std.testing.expectEqual(KvFit{ .dtype = .int8, .floor_held = true }, try kvFit(card_bytes, tight_pin, 8192, 700 << 20, ranks)); + try std.testing.expectEqual(KvFit{ .dtype = .int4, .floor_held = true }, try kvFit(card_bytes, tight_pin, 8192, 741 << 20, ranks)); + try std.testing.expectEqual(KvFit{ .dtype = .int4, .floor_held = false }, try kvFit(card_bytes, tight_pin, 8192, 742 << 20, ranks)); + try std.testing.expectEqual(@as(u64, 0), learnedBytes(card_bytes, card_bytes, 0)); + + // Counting the expert bank as resident is what blows the same card. + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident + (70 << 30), + .kv_bytes = kv, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); +} + +test "one kv head keeps bf16 through 8205 on the census floor" { + const pin: u64 = 2_072_396_202; + // Fixture floor only. Not a line count the loader hardcodes. + const floor = full_expert_bytes * 7752; + const ranks: usize = 2; + try std.testing.expectEqual(@as(u64, 100_663_296), try contextKvRank(8192, .bf16, ranks)); + + const at_8192 = try kvFit(card_bytes, pin, 8192, floor, ranks); + try std.testing.expectEqual(Dtype.bf16, at_8192.dtype); + try std.testing.expect(at_8192.floor_held); + const at_8205 = try kvFit(card_bytes, pin, 8205, floor, ranks); + try std.testing.expectEqual(Dtype.bf16, at_8205.dtype); + try std.testing.expect(at_8205.floor_held); + const at_8206 = try kvFit(card_bytes, pin, 8206, floor, ranks); + try std.testing.expectEqual(Dtype.int8, at_8206.dtype); + try std.testing.expect(at_8206.floor_held); + + // Both-heads divisor, 24,576 bytes per token. Not a call to kvFit. + const both: u64 = 24_576; + try std.testing.expectEqual(both, try contextKv(1, .bf16)); + const kv_room = card_bytes - pin - device_reserve - floor; + try std.testing.expectEqual(@as(u64, 100_834_350), kv_room); + try std.testing.expect(4_102 * both <= kv_room); + try std.testing.expect(4_103 * both > kv_room); + try std.testing.expectEqual(try contextKv(8192, .bf16), try contextKvRank(8192, .bf16, 1)); +} + +test "a zero rank does not panic and is not a held bf16 fit" { + // Floor 0 would be bf16 with the floor held if a 0-head charge were accepted. + try std.testing.expectError(error.Ranks, kvFit(card_bytes, 0, 8192, 0, 0)); + try std.testing.expectError(error.Ranks, contextKvRank(8192, .bf16, 0)); + try std.testing.expectError(error.Ranks, contextKvRank(8192, .bf16, 3)); + try std.testing.expectError(error.Ranks, kvFit(card_bytes, 0, 8192, 0, 4)); +} diff --git a/zig/src/families/flashnext/exl3_cpu.zig b/zig/src/families/flashnext/exl3_cpu.zig new file mode 100644 index 0000000000..fdfe227684 --- /dev/null +++ b/zig/src/families/flashnext/exl3_cpu.zig @@ -0,0 +1,393 @@ +//! CPU GEMV for one EXL3 matrix. No allocation. +//! The matrix is `exl3_format.dequantize`. y[j] = sum_i x[i] * W[i*N + j]. +//! `gemv` is the scalar oracle. `gemvAvx2` matches it with four column accumulators. + +const std = @import("std"); +const builtin = @import("builtin"); +const format = @import("exl3_format.zig"); + +/// AVX fp256, compiled only when this target actually has AVX2. No AVX-512. +const avx2_target = builtin.cpu.arch == .x86_64 and + builtin.zig_backend != .stage2_x86_64 and + builtin.cpu.has(.x86, .avx2); + +const Acc = @Vector(4, f64); + +/// fp32 SwiGLU from the v0.6.6 gate-up epilogue (`act_mode` 1). +/// `gate` is clamped to `limit` from above. `up` is clamped to `[-limit, limit]`. +/// `out[i] = silu(gate_i) * up_i`. A `limit` of +inf leaves both sides unclamped. +/// `error.Length` is returned before any store to `out`. +pub fn swiglu(gate: []const f64, up: []const f64, out: []f64, limit: f64) error{Length}!void { + if (gate.len != up.len or gate.len != out.len) return error.Length; + for (gate, up, out) |g, u, *y| { + const gg = @min(g, limit); + const uu = @min(@max(u, -limit), limit); + y.* = (gg / (1.0 + @exp(-gg))) * uu; + } +} + +/// Weighted sum of `weights.len` expert rows. `rows` is row-major `[topk][n]`. +/// `out[j] = sum_s weights[s] * rows[s * n + j]`. Length is checked before any store. +pub fn combine(rows: []const f64, weights: []const f64, n: usize, out: []f64) error{Length}!void { + if (n == 0 or out.len != n or weights.len == 0 or rows.len != weights.len * n) return error.Length; + @memset(out, 0); + for (weights, 0..) |w, s| { + const row = rows[s * n ..][0..n]; + for (out, row) |*y, v| y.* += w * v; + } +} + +/// K = tiles_k * 16 and N = tiles_n * 16, both multiples of 128 (one Hadamard block). +/// `scratch` holds 3*K*N f64 values: decoded W in `[0, K*N)`, then `dequantize`'s 2*K*N workspace. +/// `error.Length` is returned before any store to `y`. +pub fn gemv( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += x[i] * weight[i * n + j]; + y[j] = sum; + } +} + +/// Bit-exact with `gemv`. Lane `t` is column `j + t`, and that lane alone adds +/// `x[i] * W[i * N + j + t]` for increasing `i`. Folding the four products into +/// one sum is a different dot product. Mul and add stay separate: an FMA would +/// not match `gemv`. Scratch is the same `3 * K * N` f64 buffer `gemv` uses. +pub fn gemvAvx2( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + // N is a multiple of 128, so the four-column walk covers every column. + if (comptime avx2_target) { + accumulateAvx2(weight, x, y); + } else { + accumulatePortable(weight, x, y); + } +} + +/// Same four lanes as the AVX2 loop, one left-to-right sum per column. +/// Strict mode blocks a contracted FMA and a reassociated chain. +fn accumulatePortable(weight: []const f64, x: []const f64, y: []f64) void { + @setFloatMode(.strict); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: [4]f64 = .{ 0, 0, 0, 0 }; + var i: usize = 0; + while (i < k) : (i += 1) { + const xi = x[i]; + const row = i * n + j; + inline for (0..4) |lane| acc[lane] += xi * weight[row + lane]; + } + y[j..][0..4].* = acc; + } +} + +fn avx2Load(p: [*]const f64) Acc { + return asm ( + \\vmovupd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Broadcast(p: *const f64) Acc { + return asm ( + \\vbroadcastsd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Mul(a: Acc, b: Acc) Acc { + return asm ( + \\vmulpd %[b], %[a], %[out] + : [out] "=x" (-> Acc), + : [a] "x" (a), + [b] "x" (b), + ); +} + +/// `out = sum + prod`, the same operand order as `sum += prod`. +fn avx2Add(sum: Acc, prod: Acc) Acc { + return asm ( + \\vaddpd %[prod], %[sum], %[out] + : [out] "=x" (-> Acc), + : [sum] "x" (sum), + [prod] "x" (prod), + ); +} + +fn avx2Store(p: [*]f64, v: Acc) void { + asm volatile ( + \\vmovupd %[v], (%[p]) + : + : [v] "x" (v), + [p] "r" (p), + : .{ .memory = true }); +} + +/// ymm lanes are columns j..j+3. vmovupd: `scratch` is f64-aligned, not 32. +/// vmulpd then vaddpd. Not vfmadd, and not an AVX-512 instruction. +fn accumulateAvx2(weight: []const f64, x: []const f64, y: []f64) void { + comptime if (builtin.cpu.arch != .x86_64) @compileError("AVX2 GEMV is x86_64 only"); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: Acc = @splat(0); + var i: usize = 0; + while (i < k) : (i += 1) { + const prod = avx2Mul(avx2Broadcast(&x[i]), avx2Load(weight.ptr + i * n + j)); + acc = avx2Add(acc, prod); + } + avx2Store(y.ptr + j, acc); + } +} + +const sentinel: f64 = @bitCast(@as(u64, 0x7FF8_0000_0000_0001)); + +/// 128x128 mul1, width 2 halves. Trellis length is 8*8*Width.tileWords(). +const Square = struct { + allocator: std.mem.Allocator, + w: format.Width, + tiles_k: usize, + tiles_n: usize, + trellis: []i16, + suh: []f16, + svh: []f16, + x: []f64, + y: []f64, + scratch: []f64, + + fn init(allocator: std.mem.Allocator) !Square { + const w: format.Width = .{ .halves = 2 }; + const tiles_k: usize = 8; + const tiles_n: usize = 8; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + errdefer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + errdefer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + errdefer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + errdefer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + errdefer allocator.free(y); + for (y) |*v| v.* = sentinel; + + const scratch = try allocator.alloc(f64, 3 * k * n); + errdefer allocator.free(scratch); + return .{ + .allocator = allocator, + .w = w, + .tiles_k = tiles_k, + .tiles_n = tiles_n, + .trellis = trellis, + .suh = suh, + .svh = svh, + .x = x, + .y = y, + .scratch = scratch, + }; + } + + fn deinit(self: Square) void { + self.allocator.free(self.trellis); + self.allocator.free(self.suh); + self.allocator.free(self.svh); + self.allocator.free(self.x); + self.allocator.free(self.y); + self.allocator.free(self.scratch); + } +}; + +test "128x128 mul1 gemv matches dequantize then a naive matvec" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const k = s.tiles_k * 16; + const n = s.tiles_n * 16; + try std.testing.expectEqual(@as(usize, 128), k); + try std.testing.expectEqual(@as(usize, 128), n); + try std.testing.expectEqual(@as(usize, 8 * 8 * s.w.tileWords()), s.trellis.len); + + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + + const weight = try gpa.alloc(f64, k * n); + defer gpa.free(weight); + const dq_scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(dq_scratch); + format.dequantize(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, weight, dq_scratch); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += s.x[i] * weight[i * n + j]; + try std.testing.expectEqual(@as(u64, @bitCast(sum)), @as(u64, @bitCast(s.y[j]))); + } +} + +test "scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +test "swiglu is silu(gate) times up, and a finite limit clamps both sides" { + var out = [_]f64{ 9, 9 }; + const gate = [_]f64{ 0, 2 }; + const up = [_]f64{ 3, 4 }; + try swiglu(&gate, &up, &out, std.math.inf(f64)); + try std.testing.expectEqual(@as(f64, 0), out[0]); + const silu2 = 2.0 / (1.0 + @exp(-2.0)); + try std.testing.expectApproxEqAbs(silu2 * 4.0, out[1], 1e-12); + + const wide_gate = [_]f64{2}; + const wide_up = [_]f64{4}; + var clamped = [_]f64{9}; + try swiglu(&wide_gate, &wide_up, &clamped, 1); + const silu1 = 1.0 / (1.0 + @exp(-1.0)); + try std.testing.expectApproxEqAbs(silu1, clamped[0], 1e-12); + + var held = [_]f64{9}; + try std.testing.expectError(error.Length, swiglu(wide_gate[0..], wide_up[0..], held[0..0], 1)); + try std.testing.expectEqual(@as(f64, 9), held[0]); +} + +test "combine is the weighted sum of expert rows" { + const rows = [_]f64{ 1, 2, 3, 4 }; + const weights = [_]f64{ 0.25, 0.75 }; + var out = [_]f64{ 9, 9 }; + try combine(&rows, &weights, 2, &out); + try std.testing.expectApproxEqAbs(@as(f64, 2.5), out[0], 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 3.5), out[1], 1e-12); + try std.testing.expectError(error.Length, combine(rows[0..3], &weights, 2, &out)); +} + +test "x shorter than K returns Length" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short_x = s.x[0 .. s.x.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, short_x, s.y, s.scratch)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +fn expectGemvAvx2(allocator: std.mem.Allocator, tiles_k: usize, tiles_n: usize) !void { + const w: format.Width = .{ .halves = 2 }; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + defer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + defer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + defer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + defer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + defer allocator.free(y); + const y2 = try allocator.alloc(f64, n); + defer allocator.free(y2); + @memset(y, sentinel); + @memset(y2, sentinel); + + const scratch = try allocator.alloc(f64, 3 * k * n); + defer allocator.free(scratch); + try std.testing.expectEqual(@as(usize, 3 * k * n), scratch.len); + + try gemv(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y, scratch); + try gemvAvx2(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y2, scratch); + for (y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on the 128x128 mul1 tile" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const y2 = try gpa.alloc(f64, s.y.len); + defer gpa.free(y2); + @memset(s.y, sentinel); + @memset(y2, sentinel); + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + try gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, y2, s.scratch); + for (s.y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on synthetic 2560x640 and 640x2560" { + // 2560 * 640 * 3 f64. One buffer, not a second scratch shape. + try std.testing.expectEqual(@as(usize, 39_321_600), 3 * 2560 * 640 * @sizeOf(f64)); + const gpa = std.testing.allocator; + try expectGemvAvx2(gpa, 160, 40); + try expectGemvAvx2(gpa, 40, 160); +} + +test "gemvAvx2 scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} diff --git a/zig/src/families/flashnext/exl3_dplan_layout.zig b/zig/src/families/flashnext/exl3_dplan_layout.zig new file mode 100644 index 0000000000..e0d2da8d62 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_layout.zig @@ -0,0 +1,128 @@ +const std = @import("std"); + +pub const experts_n: usize = 512; +pub const recent_window: usize = 256; +pub const recent_words: usize = recent_window * (experts_n / 32); +pub const misc_words: usize = 10; +pub const static_bytes: usize = experts_n * 3 * (8 + 8 + 4); + +pub fn recentCountsOffset(lines: usize) usize { + return 2 * lines + 2 * experts_n; +} + +pub fn recentOffset(lines: usize) usize { + return 2 * lines + 3 * experts_n; +} + +pub fn miscOffset(lines: usize) usize { + return recentOffset(lines) + recent_words; +} + +pub fn stateWords(lines: usize) usize { + return miscOffset(lines) + misc_words; +} + +pub const PoolState = struct { + lines: usize, + offset_words: usize, + words: usize, + + pub fn init(lines: usize, offset_words: usize) !PoolState { + if (lines == 0 or lines > (std.math.maxInt(c_int) - stateWords(0)) / 2) return error.Invalid; + const words = stateWords(lines); + const end = try std.math.add(usize, offset_words, words); + _ = try std.math.mul(usize, end, @sizeOf(u32)); + return .{ .lines = lines, .offset_words = offset_words, .words = words }; + } + + pub fn endWords(self: PoolState) usize { + return self.offset_words + self.words; + } + + pub fn byteOffset(self: PoolState) usize { + return self.offset_words * @sizeOf(u32); + } + + pub fn slice(self: PoolState, all: []u32) []u32 { + return all[self.offset_words..self.endWords()]; + } +}; + +pub fn deviceBytes(layers: u64, total_lines: u64, rows: u64) !u64 { + const fixed = try std.math.mul(u64, layers, stateWords(0) * @sizeOf(u32) + static_bytes); + const state_lines = try std.math.mul(u64, total_lines, 2 * @sizeOf(u32)); + const work = try std.math.add(u64, try std.math.mul(u64, rows, 10 * 9 * 8 + 2 * 11 * 4), 256 + (experts_n + 64) * 4); + return std.math.add(u64, try std.math.add(u64, fixed, state_lines), work); +} + +test "exl3: ragged planner offsets are contiguous and exactly bounded" { + const counts = [_]usize{ 10, 73, 11, 32 }; + var end: usize = 0; + for (counts) |lines| { + const entry = try PoolState.init(lines, end); + try std.testing.expectEqual(end, entry.offset_words); + try std.testing.expectEqual(stateWords(lines), entry.words); + try std.testing.expectEqual(end * 4, entry.byteOffset()); + try std.testing.expectEqual(entry.words, miscOffset(lines) + misc_words); + end = entry.endWords(); + } + try std.testing.expectEqual(counts.len * stateWords(0) + 2 * (10 + 73 + 11 + 32), end); + try std.testing.expectError(error.Invalid, PoolState.init(0, 0)); + const max_lines = (std.math.maxInt(c_int) - stateWords(0)) / 2; + _ = try PoolState.init(max_lines, 0); + try std.testing.expectError(error.Invalid, PoolState.init(max_lines + 1, 0)); + try std.testing.expectError(error.Invalid, PoolState.init(std.math.maxInt(usize), 0)); + try std.testing.expectError(error.Overflow, PoolState.init(10, std.math.maxInt(usize))); + try std.testing.expectError(error.Overflow, PoolState.init(10, std.math.maxInt(usize) / 4)); + var uniform_end: usize = 0; + for (0..4) |index| { + const entry = try PoolState.init(32, uniform_end); + try std.testing.expectEqual(index * stateWords(32), entry.offset_words); + uniform_end = entry.endWords(); + } +} + +test "exl3: ragged planner state round trip isolates line maps history and counters" { + const counts = [_]usize{ 10, 57, 13 }; + var entries: [counts.len]PoolState = undefined; + var end: usize = 1; + for (counts, &entries) |lines, *entry| { + entry.* = try PoolState.init(lines, end); + end = entry.endWords(); + } + const host = try std.testing.allocator.alloc(u32, end + 1); + defer std.testing.allocator.free(host); + @memset(host, 0xfeedface); + for (entries, 0..) |entry, index| { + const state = entry.slice(host); + @memset(state, 0); + for (0..entry.lines) |line| { + const expert = (line + index * 100) % experts_n; + state[line] = @intCast(expert); + state[entry.lines + line] = @intCast(line + 1); + state[2 * entry.lines + expert] = @intCast(line); + } + state[recentCountsOffset(entry.lines)] = @intCast(index + 1); + state[recentOffset(entry.lines)] = @intCast(index + 10); + state[miscOffset(entry.lines)] = @intCast(index + 20); + state[state.len - 1] = @intCast(index + 30); + } + const device = try std.testing.allocator.dupe(u32, host); + defer std.testing.allocator.free(device); + @memset(host[1..end], 0); + for (entries) |entry| @memcpy(entry.slice(host), entry.slice(device)); + try std.testing.expectEqual(@as(u32, 0xfeedface), host[0]); + try std.testing.expectEqual(@as(u32, 0xfeedface), host[end]); + for (entries, 0..) |entry, index| { + const state = entry.slice(host); + for (0..entry.lines) |line| { + const expert = state[line]; + try std.testing.expectEqual(@as(u32, @intCast(line)), state[2 * entry.lines + expert]); + try std.testing.expectEqual(@as(u32, @intCast(line + 1)), state[entry.lines + line]); + } + try std.testing.expectEqual(@as(u32, @intCast(index + 1)), state[recentCountsOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 10)), state[recentOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 20)), state[miscOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 30)), state[state.len - 1]); + } +} diff --git a/zig/src/families/flashnext/exl3_dplan_policy_eval.zig b/zig/src/families/flashnext/exl3_dplan_policy_eval.zig new file mode 100644 index 0000000000..f92b1f87a8 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_policy_eval.zig @@ -0,0 +1,672 @@ +//! CPU counterfactuals for one ordered pool-demand timeline. Options.seed supplies +//! the trace's initial state when requested; policies share that initial +//! state and evolve independently. Raw picks determine demand and LRU order. +//! Misses are measured against residency at the start of each round. Every current +//! demand is pinned, including newly admitted experts; excess misses use the mapped +//! bank. One admission means one whole-expert copy, with no prefetch or byte model. +//! Counts describe this model only and do not predict runtime or transfer latency. +//! Lifetime uses saturating u32 counts, hottest-miss admission, then count/stamp/line +//! victim ties as in DPLAN. LRU admits most recently picked misses first and evicts +//! least recently picked residents, preserving raw last-pick order within a round. +//! Decayed frequency right-shifts every score at configured round boundaries before +//! adding appearances; it admits by score/appearances/ID and evicts by score/LRU/line. +//! Decayed admission additionally bypasses a newcomer colder than its eligible victim. +//! Seeds initialize scores for frequency policies; no pre-trace decay history +//! can be inferred. Seed stamps initialize LRU with numeric stamp order and line ties. +//! More than 256 unique routed experts is rejected because DPLAN overflows there. + +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); + +pub const expert_count: usize = 512; +pub const max_unique: usize = 256; + +pub const Policy = enum { lifetime, lru, decayed_frequency, decayed_admission }; + +pub const Seed = struct { + line_ids: []const i32, + stamps: []const u32, + counts: [expert_count]u32, + now: u32, +}; + +pub const Options = struct { + capacity: usize, + seed: ?Seed = null, + decay_interval_rounds: u64 = 64, + decay_shift: u5 = 1, +}; + +pub const Metrics = struct { + rounds: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + hit_experts: u64 = 0, + miss_experts: u64 = 0, + admissions: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + mapped_experts: u64 = 0, + mapped_picks: u64 = 0, + + fn plus(self: Metrics, other: Metrics) !Metrics { + var result: Metrics = undefined; + inline for (@typeInfo(Metrics).@"struct".field_names) |field| { + @field(result, field) = try std.math.add(u64, @field(self, field), @field(other, field)); + } + return result; + } +}; + +pub const Comparison = struct { + lifetime: Metrics, + lru: Metrics, + decayed_frequency: Metrics, + decayed_admission: Metrics, + belady_unique_misses_lower_bound: u64, +}; + +const Demand = struct { + counts: [expert_count]u32 = @splat(0), + ordered: [expert_count]usize = undefined, + unique: usize = 0, + valid_picks: u64 = 0, + picks: []const u32, +}; + +pub const State = struct { + allocator: std.mem.Allocator, + policy: Policy, + options: Options, + line_ids: []i32, + stamps: []u32, + scores: [expert_count]u32, + last_access: [expert_count]u64, + access_clock: u64, + now: u32, + total: Metrics = .{}, + + pub fn init(allocator: std.mem.Allocator, policy: Policy, options: Options) !State { + try validateOptions(options); + const line_ids = try allocator.alloc(i32, options.capacity); + errdefer allocator.free(line_ids); + const stamps = try allocator.alloc(u32, options.capacity); + @memset(line_ids, -1); + @memset(stamps, 0); + var self: State = .{ + .allocator = allocator, + .policy = policy, + .options = options, + .line_ids = line_ids, + .stamps = stamps, + .scores = @splat(0), + .last_access = @splat(0), + .access_clock = 0, + .now = 0, + }; + if (options.seed) |seed| { + @memcpy(line_ids, seed.line_ids); + @memcpy(stamps, seed.stamps); + self.scores = seed.counts; + self.now = seed.now; + self.access_clock = seed.now; + for (line_ids, stamps) |expert, stamp| { + self.access_clock = @max(self.access_clock, stamp); + if (expert >= 0) self.last_access[@intCast(expert)] = stamp; + } + } + self.options.seed = null; + return self; + } + + pub fn deinit(self: *State) void { + self.allocator.free(self.line_ids); + self.allocator.free(self.stamps); + self.* = undefined; + } + + pub fn observe(self: *State, record: trace.Round) !Metrics { + const demand = try readDemand(record); + const round_number = try std.math.add(u64, self.total.rounds, 1); + const next_clock = try std.math.add(u64, self.access_clock, demand.valid_picks); + var delta: Metrics = .{ .rounds = 1 }; + for (demand.ordered[0..demand.unique]) |expert| { + if (self.lineOf(expert) != null) { + delta.hit_picks += demand.counts[expert]; + delta.hit_experts += 1; + } else { + delta.miss_picks += demand.counts[expert]; + delta.miss_experts += 1; + } + } + _ = try self.total.plus(delta); + if ((self.policy == .decayed_frequency or self.policy == .decayed_admission) and round_number % self.options.decay_interval_rounds == 0) { + for (&self.scores) |*score| score.* >>= self.options.decay_shift; + } + self.now +%= 1; + for (demand.ordered[0..demand.unique]) |expert| { + self.scores[expert] +|= demand.counts[expert]; + if (self.lineOf(expert)) |line| self.stamps[line] = self.now; + } + for (demand.picks) |bits| { + const expert = routed(bits) orelse continue; + self.access_clock += 1; + self.last_access[expert] = self.access_clock; + } + std.debug.assert(self.access_clock == next_clock); + + var misses: [max_unique]usize = undefined; + var miss_count: usize = 0; + for (demand.ordered[0..demand.unique]) |expert| { + if (self.lineOf(expert) == null) { + misses[miss_count] = expert; + miss_count += 1; + } + } + var index: usize = 1; + while (index < miss_count) : (index += 1) { + const expert = misses[index]; + var position = index; + while (position > 0 and self.preferAdmission(expert, misses[position - 1], &demand)) : (position -= 1) { + misses[position] = misses[position - 1]; + } + misses[position] = expert; + } + for (misses[0..miss_count]) |expert| { + if (self.victim(&demand)) |line| { + if (self.policy == .decayed_admission and self.line_ids[line] >= 0 and self.scores[expert] < self.scores[@intCast(self.line_ids[line])]) { + delta.mapped_experts += 1; + delta.mapped_picks += demand.counts[expert]; + continue; + } + if (self.line_ids[line] >= 0) delta.evictions += 1; + self.line_ids[line] = @intCast(expert); + self.stamps[line] = self.now; + delta.admissions += 1; + delta.copies += 1; + } else { + delta.mapped_experts += 1; + delta.mapped_picks += demand.counts[expert]; + } + } + self.total = try self.total.plus(delta); + return delta; + } + + fn lineOf(self: *const State, expert: usize) ?usize { + for (self.line_ids, 0..) |resident, line| { + if (resident == @as(i32, @intCast(expert))) return line; + } + return null; + } + + fn preferAdmission(self: *const State, left: usize, right: usize, demand: *const Demand) bool { + if (self.policy == .lru) { + if (self.last_access[left] != self.last_access[right]) return self.last_access[left] > self.last_access[right]; + } else { + if (self.scores[left] != self.scores[right]) return self.scores[left] > self.scores[right]; + if (demand.counts[left] != demand.counts[right]) return demand.counts[left] > demand.counts[right]; + } + return left < right; + } + + fn victim(self: *const State, demand: *const Demand) ?usize { + var coldest: ?usize = null; + for (self.line_ids, 0..) |expert, line| { + if (expert < 0) return line; + if (demand.counts[@intCast(expert)] != 0) continue; + if (coldest == null or self.colder(line, coldest.?)) coldest = line; + } + return coldest; + } + + fn colder(self: *const State, left: usize, right: usize) bool { + const left_expert: usize = @intCast(self.line_ids[left]); + const right_expert: usize = @intCast(self.line_ids[right]); + if (self.policy != .lru and self.scores[left_expert] != self.scores[right_expert]) return self.scores[left_expert] < self.scores[right_expert]; + if (self.policy == .lifetime) { + if (self.stamps[left] != self.stamps[right]) return self.stamps[left] < self.stamps[right]; + } else { + if (self.last_access[left_expert] != self.last_access[right_expert]) return self.last_access[left_expert] < self.last_access[right_expert]; + } + return left > right; + } +}; + +pub fn evaluate(allocator: std.mem.Allocator, rounds: []const trace.Round, options: Options) !Comparison { + var lifetime = try State.init(allocator, .lifetime, options); + defer lifetime.deinit(); + var lru = try State.init(allocator, .lru, options); + defer lru.deinit(); + var decayed = try State.init(allocator, .decayed_frequency, options); + defer decayed.deinit(); + var decayed_admission = try State.init(allocator, .decayed_admission, options); + defer decayed_admission.deinit(); + for (rounds) |record| { + _ = try lifetime.observe(record); + _ = try lru.observe(record); + _ = try decayed.observe(record); + _ = try decayed_admission.observe(record); + } + return .{ + .lifetime = lifetime.total, + .lru = lru.total, + .decayed_frequency = decayed.total, + .decayed_admission = decayed_admission.total, + .belady_unique_misses_lower_bound = try beladyLowerBound(allocator, rounds, options), + }; +} + +pub fn evaluateTrace(allocator: std.mem.Allocator, view: trace.View) !Comparison { + const verification = try @import("exl3_dplan_replay.zig").verifyTrace(allocator, view); + const rounds = try allocator.alloc(trace.Round, view.recorded()); + defer allocator.free(rounds); + for (rounds, 0..) |*round, index| round.* = try view.round(index); + var options: Options = .{ .capacity = view.words[4], .seed = view.seedState(Seed) }; + var total: Comparison = .{ .lifetime = .{}, .lru = .{}, .decayed_frequency = .{}, .decayed_admission = .{}, .belady_unique_misses_lower_bound = 0 }; + var begin: usize = 0; + for (0..view.checkpoints() + 1) |index| { + const checkpoint = if (index < view.checkpoints()) try view.checkpoint(index) else null; + const end = if (checkpoint) |entry| entry.roundIndex() else view.recorded(); + if (end > begin) { + const segment = try evaluate(allocator, rounds[begin..end], options); + inline for (.{ "lifetime", "lru", "decayed_frequency", "decayed_admission" }) |field| { + @field(total, field) = try @field(total, field).plus(@field(segment, field)); + } + total.belady_unique_misses_lower_bound = try std.math.add(u64, total.belady_unique_misses_lower_bound, segment.belady_unique_misses_lower_bound); + } + if (checkpoint) |entry| options.seed = trace.seedFromState(Seed, view.words[4], entry.after()); + begin = end; + } + const baseline = total.lifetime; + const exact = verification.delta; + if (baseline.rounds != @as(u64, @intCast(verification.rounds)) or + baseline.hit_picks != exact.hit_picks or + baseline.miss_picks != exact.miss_picks or + baseline.copies != exact.copies or + baseline.evictions != exact.evictions or + baseline.mapped_experts != exact.inplace or + baseline.mapped_picks != exact.mapped_picks) + return error.PolicyBaselineMismatch; + return total; +} + +/// Belady with optional admission is optimal for uniform-cost sequential requests: +/// on a full-cache miss, discard the farthest-next-use of residents AND the incoming +/// expert (discarding the incoming expert is a bypass). Equal next uses may tie. +/// The stream contains each expert once per round, in raw first-appearance order. +/// It relaxes round pinning. Every feasible round policy has a feasible sequential +/// execution with the same unique misses, so this optimum is a lower bound on +/// miss_experts only. Repeated pick costs, copies and mapped reads are not bounded. +/// Future use ends at the supplied trace boundary; no unseen suffix is inferred. +pub fn beladyLowerBound(allocator: std.mem.Allocator, rounds: []const trace.Round, options: Options) !u64 { + try validateOptions(options); + var count: usize = 0; + for (rounds) |record| count = try std.math.add(usize, count, (try readDemand(record)).unique); + const events = try allocator.alloc(Event, count); + defer allocator.free(events); + var offset: usize = 0; + for (rounds) |record| { + const demand = try readDemand(record); + for (demand.ordered[0..demand.unique]) |expert| { + events[offset] = .{ .expert = expert }; + offset += 1; + } + } + return belady(events, options.capacity, if (options.seed) |seed| seed.line_ids else &.{}); +} + +const Event = struct { expert: usize, next: usize = std.math.maxInt(usize) }; + +fn belady(events: []Event, capacity: usize, initial: []const i32) u64 { + var next_use: [expert_count]usize = @splat(std.math.maxInt(usize)); + var position = events.len; + while (position > 0) { + position -= 1; + const event = &events[position]; + event.next = next_use[event.expert]; + next_use[event.expert] = position; + } + var resident: [expert_count]i32 = @splat(-1); + @memcpy(resident[0..initial.len], initial); + var misses: u64 = 0; + for (events) |event| { + var hit = false; + var empty: ?usize = null; + var farthest: ?usize = null; + for (resident[0..capacity], 0..) |expert, line| { + if (expert == @as(i32, @intCast(event.expert))) hit = true; + if (expert < 0) { + if (empty == null) empty = line; + } else if (farthest == null or next_use[@intCast(expert)] > next_use[@intCast(resident[farthest.?])]) { + farthest = line; + } + } + next_use[event.expert] = event.next; + if (hit) continue; + misses += 1; + if (empty) |line| { + resident[line] = @intCast(event.expert); + } else if (farthest) |line| { + if (event.next < next_use[@intCast(resident[line])]) resident[line] = @intCast(event.expert); + } + } + return misses; +} + +fn validateOptions(options: Options) !void { + if (options.capacity > expert_count or options.decay_interval_rounds == 0 or options.decay_shift == 0) return error.InvalidOptions; + if (options.seed) |seed| { + if (seed.line_ids.len != options.capacity or seed.stamps.len != options.capacity) return error.InvalidSeed; + var seen: [expert_count]bool = @splat(false); + for (seed.line_ids) |expert| { + if (expert < -1 or expert >= expert_count) return error.InvalidSeed; + if (expert >= 0) { + const index: usize = @intCast(expert); + if (seen[index]) return error.InvalidSeed; + seen[index] = true; + } + } + } +} + +fn readDemand(record: trace.Round) !Demand { + const shape = trace.header(0, 0, record.lines, record.max_rows) catch return error.InvalidRound; + if (record.words.len != shape[7]) return error.InvalidRound; + if (record.recentEnabled()) return error.UnsupportedRecentPolicy; + if (record.rows() == 0 or record.rows() > record.max_rows or record.words[3] > 1) return error.InvalidRound; + var demand: Demand = .{ .picks = record.picks() }; + for (demand.picks) |bits| { + const expert = routed(bits) orelse continue; + if (demand.counts[expert] == 0) { + demand.ordered[demand.unique] = expert; + demand.unique += 1; + } + demand.counts[expert] += 1; + demand.valid_picks += 1; + } + if (demand.valid_picks != record.validPicks() or demand.unique != record.uniqueCount()) return error.InvalidRound; + if (demand.unique > max_unique) return error.PlannerOverflow; + return demand; +} + +fn routed(bits: u32) ?usize { + const expert: i32 = @bitCast(bits); + return if (expert >= 0 and expert < expert_count) @intCast(expert) else null; +} + +const Synthetic = struct { + words: []u32, + max_rows: usize, + + fn init(picks: []const i32) !Synthetic { + const rows = @max(1, (picks.len + 10) / 11); + const shape = try trace.header(0, 0, 2, rows); + const words = try std.testing.allocator.alloc(u32, shape[7]); + @memset(words, 0); + words[0] = @intCast(rows); + const self: Synthetic = .{ .words = words, .max_rows = rows }; + const raw_offset = trace.record_header_words; + @memset(words[raw_offset..][0 .. rows * 11], std.math.maxInt(u32)); + var counts: [expert_count]u32 = @splat(0); + for (picks, 0..) |expert, index| { + words[raw_offset + index] = @bitCast(expert); + if (routed(@bitCast(expert))) |valid| { + if (counts[valid] == 0) words[2] += 1; + counts[valid] += 1; + words[1] += 1; + } + } + return self; + } + + fn round(self: Synthetic) trace.Round { + return .{ .words = self.words, .lines = 2, .max_rows = self.max_rows }; + } + + fn deinit(self: Synthetic) void { + std.testing.allocator.free(self.words); + } +}; + +test "exl3: policy evaluation pins round hits and admissions and separates mapped misses" { + var fixture = try Synthetic.init(&.{ 0, 2, 2, 2, 3, 512, -1, 999, -9 }); + defer fixture.deinit(); + var state = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = @splat(0), .now = 2 }, + }); + defer state.deinit(); + const result = try state.observe(fixture.round()); + try std.testing.expectEqual(Metrics{ .rounds = 1, .hit_picks = 1, .miss_picks = 4, .hit_experts = 1, .miss_experts = 2, .admissions = 1, .copies = 1, .evictions = 1, .mapped_experts = 1, .mapped_picks = 1 }, result); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, state.line_ids); +} + +test "exl3: lifetime admission saturation and victim ties match current planner" { + var fixture = try Synthetic.init(&.{ 2, 2, 2, 2, 2, 3 }); + defer fixture.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 5; + counts[1] = 3; + counts[2] = std.math.maxInt(u32) - 2; + var state = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = std.math.maxInt(u32) }, + }); + defer state.deinit(); + const result = try state.observe(fixture.round()); + try std.testing.expectEqual(@as(u64, 2), result.copies); + try std.testing.expectEqualSlices(i32, &.{ 3, 2 }, state.line_ids); + try std.testing.expectEqual(std.math.maxInt(u32), state.scores[2]); + try std.testing.expectEqual(@as(u32, 0), state.now); + + var tie = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 7, 7 }, .counts = @splat(0), .now = 7 }, + }); + defer tie.deinit(); + var single = try Synthetic.init(&.{9}); + defer single.deinit(); + _ = try tie.observe(single.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 9 }, tie.line_ids); +} + +test "exl3: LRU preserves raw last occurrence order when the round exceeds capacity" { + var first = try Synthetic.init(&.{ 1, 2, 1 }); + defer first.deinit(); + var second = try Synthetic.init(&.{ 1, 1, 2 }); + defer second.deinit(); + var left = try State.init(std.testing.allocator, .lru, .{ .capacity = 1 }); + defer left.deinit(); + var right = try State.init(std.testing.allocator, .lru, .{ .capacity = 1 }); + defer right.deinit(); + _ = try left.observe(first.round()); + _ = try right.observe(second.round()); + try std.testing.expectEqualSlices(i32, &.{1}, left.line_ids); + try std.testing.expectEqualSlices(i32, &.{2}, right.line_ids); + _ = try left.observe(second.round()); + try std.testing.expectEqualSlices(i32, &.{1}, left.line_ids); +} + +test "exl3: frequency decay halves scores at configured round boundaries before new demand" { + var fixture = try Synthetic.init(&.{ 0, 1, 1, 1 }); + defer fixture.deinit(); + var empty = try Synthetic.init(&.{}); + defer empty.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 64; + const options: Options = .{ + .capacity = 1, + .seed = .{ .line_ids = &.{-1}, .stamps = &.{0}, .counts = counts, .now = 0 }, + .decay_interval_rounds = 1, + .decay_shift = 6, + }; + var lifetime = try State.init(std.testing.allocator, .lifetime, options); + defer lifetime.deinit(); + var decayed = try State.init(std.testing.allocator, .decayed_frequency, options); + defer decayed.deinit(); + _ = try lifetime.observe(fixture.round()); + _ = try decayed.observe(fixture.round()); + try std.testing.expectEqualSlices(i32, &.{0}, lifetime.line_ids); + try std.testing.expectEqualSlices(i32, &.{1}, decayed.line_ids); + + var interval = try State.init(std.testing.allocator, .decayed_frequency, .{ .capacity = 0, .decay_interval_rounds = 2 }); + defer interval.deinit(); + _ = try interval.observe(fixture.round()); + try std.testing.expectEqual(@as(u32, 3), interval.scores[1]); + _ = try interval.observe(empty.round()); + try std.testing.expectEqual(@as(u32, 1), interval.scores[1]); +} + +test "exl3: LRU evicts by within-round recency and decayed eviction forgets old frequency" { + var touching = try Synthetic.init(&.{ 0, 1, 0 }); + defer touching.deinit(); + var incoming = try Synthetic.init(&.{2}); + defer incoming.deinit(); + var empty = try Synthetic.init(&.{}); + defer empty.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 128; + counts[1] = 4; + const options: Options = .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = 2 }, + .decay_interval_rounds = 1, + .decay_shift = 6, + }; + var lru = try State.init(std.testing.allocator, .lru, options); + defer lru.deinit(); + _ = try lru.observe(touching.round()); + _ = try lru.observe(incoming.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, lru.line_ids); + + var lifetime = try State.init(std.testing.allocator, .lifetime, options); + defer lifetime.deinit(); + var decayed = try State.init(std.testing.allocator, .decayed_frequency, options); + defer decayed.deinit(); + for (0..2) |_| { + _ = try lifetime.observe(empty.round()); + _ = try decayed.observe(empty.round()); + } + _ = try lifetime.observe(incoming.round()); + _ = try decayed.observe(incoming.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, lifetime.line_ids); + try std.testing.expectEqualSlices(i32, &.{ 2, 1 }, decayed.line_ids); +} + +test "exl3: frequency admission bypasses a colder newcomer and later admits it when hot" { + var warm = try Synthetic.init(&.{ 1, 1 }); + defer warm.deinit(); + var hot = try Synthetic.init(&.{ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 }); + defer hot.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 10; + var state = try State.init(std.testing.allocator, .decayed_admission, .{ + .capacity = 1, + .seed = .{ .line_ids = &.{0}, .stamps = &.{1}, .counts = counts, .now = 1 }, + }); + defer state.deinit(); + const cold = try state.observe(warm.round()); + try std.testing.expectEqual(@as(u64, 0), cold.copies); + try std.testing.expectEqual(@as(u64, 1), cold.mapped_experts); + try std.testing.expectEqualSlices(i32, &.{0}, state.line_ids); + const now_hot = try state.observe(hot.round()); + try std.testing.expectEqual(@as(u64, 1), now_hot.copies); + try std.testing.expectEqual(@as(u64, 1), now_hot.evictions); + try std.testing.expectEqualSlices(i32, &.{1}, state.line_ids); +} + +test "exl3: comparison evolves independently from the initial seed and supports zero capacity" { + var fixture = try Synthetic.init(&.{ 5, 5 }); + defer fixture.deinit(); + const rounds = [_]trace.Round{ fixture.round(), fixture.round() }; + const result = try evaluate(std.testing.allocator, &rounds, .{ .capacity = 1 }); + try std.testing.expectEqual(@as(u64, 2), result.lifetime.miss_picks); + try std.testing.expectEqual(@as(u64, 2), result.lifetime.hit_picks); + try std.testing.expectEqual(@as(u64, 1), result.lifetime.copies); + try std.testing.expectEqual(result.lifetime, result.lru); + try std.testing.expectEqual(result.lifetime, result.decayed_frequency); + try std.testing.expectEqual(result.lifetime, result.decayed_admission); + try std.testing.expectEqual(@as(u64, 1), result.belady_unique_misses_lower_bound); + const zero = try evaluate(std.testing.allocator, &rounds, .{ .capacity = 0 }); + try std.testing.expectEqual(@as(u64, 4), zero.lifetime.mapped_picks); + try std.testing.expectEqual(@as(u64, 0), zero.lifetime.copies); + try std.testing.expectEqual(@as(u64, 4), zero.decayed_admission.mapped_picks); + try std.testing.expectEqual(@as(u64, 0), zero.decayed_admission.copies); + try std.testing.expectEqual(@as(u64, 2), zero.belady_unique_misses_lower_bound); + const none = try evaluate(std.testing.allocator, &.{}, .{ .capacity = 2 }); + try std.testing.expectEqual(Metrics{}, none.lifetime); + try std.testing.expectEqual(@as(u64, 0), none.belady_unique_misses_lower_bound); +} + +test "exl3: Belady bypass retains a warm expert and bounds pinned round misses only" { + var first = try Synthetic.init(&.{ 1, 1, 2 }); + defer first.deinit(); + var second = try Synthetic.init(&.{ 0, 0 }); + defer second.deinit(); + const options: Options = .{ + .capacity = 1, + .seed = .{ .line_ids = &.{0}, .stamps = &.{0}, .counts = @splat(0), .now = 0 }, + }; + const result = try evaluate(std.testing.allocator, &.{ first.round(), second.round() }, options); + try std.testing.expectEqual(@as(u64, 2), result.belady_unique_misses_lower_bound); + try std.testing.expectEqual(@as(u64, 3), result.lifetime.miss_experts); + try std.testing.expect(result.belady_unique_misses_lower_bound <= result.lru.miss_experts); + try std.testing.expect(result.belady_unique_misses_lower_bound <= result.decayed_frequency.miss_experts); +} + +fn exhaustiveMisses(events: []const Event, position: usize, mask: u8, capacity: usize) u64 { + if (position == events.len) return 0; + const bit = @as(u8, 1) << @as(u3, @intCast(events[position].expert)); + if (mask & bit != 0) return exhaustiveMisses(events, position + 1, mask, capacity); + var best = exhaustiveMisses(events, position + 1, mask, capacity); + if (@popCount(mask) < capacity) { + best = @min(best, exhaustiveMisses(events, position + 1, mask | bit, capacity)); + } else { + for (0..3) |expert| { + const victim_bit = @as(u8, 1) << @as(u3, @intCast(expert)); + if (mask & victim_bit != 0) best = @min(best, exhaustiveMisses(events, position + 1, (mask & ~victim_bit) | bit, capacity)); + } + } + return 1 + best; +} + +test "exl3: Belady equals exhaustive optional-admission optimum for short cold and warm streams" { + for (0..243) |sequence| { + var encoded = sequence; + var events: [5]Event = undefined; + for (&events) |*event| { + event.* = .{ .expert = encoded % 3 }; + encoded /= 3; + } + for (0..3) |capacity| { + try std.testing.expectEqual(exhaustiveMisses(&events, 0, 0, capacity), belady(&events, capacity, &.{})); + if (capacity > 0) try std.testing.expectEqual(exhaustiveMisses(&events, 0, 1, capacity), belady(&events, capacity, &.{0})); + } + } +} + +test "exl3: invalid input is rejected before state changes including planner unique overflow" { + var state = try State.init(std.testing.allocator, .lifetime, .{ .capacity = 1 }); + defer state.deinit(); + var fixture = try Synthetic.init(&.{1}); + defer fixture.deinit(); + var truncated = fixture.round(); + truncated.words = truncated.words[0..3]; + try std.testing.expectError(error.InvalidRound, state.observe(truncated)); + fixture.words[1] = 2; + try std.testing.expectError(error.InvalidRound, state.observe(fixture.round())); + var picks: [257]i32 = undefined; + for (&picks, 0..) |*expert, index| expert.* = @intCast(index); + var overflow = try Synthetic.init(&picks); + defer overflow.deinit(); + try std.testing.expectError(error.PlannerOverflow, state.observe(overflow.round())); + try std.testing.expectEqual(Metrics{}, state.total); + try std.testing.expectEqualSlices(i32, &.{-1}, state.line_ids); + try std.testing.expectError(error.InvalidOptions, State.init(std.testing.allocator, .lru, .{ .capacity = 513 })); + try std.testing.expectError(error.InvalidOptions, State.init(std.testing.allocator, .lru, .{ .capacity = 1, .decay_interval_rounds = 0 })); + try std.testing.expectError(error.InvalidSeed, State.init(std.testing.allocator, .lru, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 1, 1 }, .stamps = &.{ 0, 0 }, .counts = @splat(0), .now = 0 }, + })); +} diff --git a/zig/src/families/flashnext/exl3_dplan_replay.zig b/zig/src/families/flashnext/exl3_dplan_replay.zig new file mode 100644 index 0000000000..e3855df0b0 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_replay.zig @@ -0,0 +1,392 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); + +pub const expert_count: usize = 512; +pub const max_unique: usize = 256; + +pub const Entry = struct { expert: u32, appearances: u32 }; +pub const Round = struct { + rows: u32, + valid_picks: u32, + entries: []const Entry, + planner_overflow: bool = false, +}; + +pub const Seed = struct { + line_ids: []const i32, + stamps: []const u32, + counts: [expert_count]u32, + now: u32, +}; + +pub const Delta = struct { + hit_picks: u64 = 0, + miss_picks: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + inplace: u64 = 0, + mapped_picks: u64 = 0, + overflow: u64 = 0, +}; + +pub const Verification = struct { rounds: usize, matched_transitions: usize, delta: Delta }; + +pub const State = struct { + allocator: std.mem.Allocator, + line_ids: []i32, + stamps: []u32, + counts: [expert_count]u32, + now: u32, + + pub fn init(allocator: std.mem.Allocator, seed: Seed) !State { + if (seed.line_ids.len == 0 or seed.line_ids.len != seed.stamps.len) return error.InvalidSeed; + var resident: [expert_count]bool = @splat(false); + for (seed.line_ids) |id| { + if (id < -1 or id >= @as(i32, expert_count)) return error.InvalidSeed; + if (id >= 0) { + const index: usize = @intCast(id); + if (resident[index]) return error.InvalidSeed; + resident[index] = true; + } + } + const line_ids = try allocator.dupe(i32, seed.line_ids); + errdefer allocator.free(line_ids); + const stamps = try allocator.dupe(u32, seed.stamps); + return .{ .allocator = allocator, .line_ids = line_ids, .stamps = stamps, .counts = seed.counts, .now = seed.now }; + } + + pub fn deinit(self: *State) void { + self.allocator.free(self.line_ids); + self.allocator.free(self.stamps); + self.* = undefined; + } + + pub fn replay(self: *State, round: Round) !Delta { + if (round.planner_overflow or round.rows == 0 or round.rows > std.math.maxInt(u32) / 11 or round.valid_picks > round.rows * 11 or round.entries.len > max_unique) return error.InvalidRound; + var seen: [expert_count]bool = @splat(false); + var appearances: u64 = 0; + for (round.entries) |entry| { + if (entry.expert >= @as(u32, expert_count) or entry.appearances == 0 or seen[@intCast(entry.expert)]) return error.InvalidRound; + seen[@intCast(entry.expert)] = true; + appearances += entry.appearances; + } + if (appearances != round.valid_picks) return error.InvalidRound; + + self.now +%= 1; + var delta: Delta = .{}; + for (round.entries) |entry| { + const expert: usize = @intCast(entry.expert); + self.counts[expert] = self.counts[expert] +| entry.appearances; + if (self.lineOf(entry.expert) != null) { + delta.hit_picks += entry.appearances; + } else { + delta.miss_picks += entry.appearances; + } + } + + for (round.entries) |entry| if (self.lineOf(entry.expert)) |line| { + self.stamps[line] = self.now; + }; + + var misses: [max_unique]Entry = undefined; + var miss_count: usize = 0; + for (round.entries) |entry| if (self.lineOf(entry.expert) == null) { + misses[miss_count] = entry; + miss_count += 1; + }; + var index: usize = 1; + while (index < miss_count) : (index += 1) { + const item = misses[index]; + var insert = index; + while (insert > 0 and hotter(self, item, misses[insert - 1])) : (insert -= 1) { + misses[insert] = misses[insert - 1]; + } + misses[insert] = item; + } + + for (misses[0..miss_count]) |entry| { + if (self.victim(&seen)) |line| { + if (self.line_ids[line] >= 0) delta.evictions += 1; + self.line_ids[line] = @intCast(entry.expert); + self.stamps[line] = self.now; + delta.copies += 1; + } else { + delta.inplace += 1; + delta.mapped_picks += entry.appearances; + } + } + return delta; + } + + fn lineOf(self: *const State, expert: u32) ?usize { + const target: i32 = @intCast(expert); + for (self.line_ids, 0..) |id, line| if (id == target) return line; + return null; + } + + fn victim(self: *const State, keep: *const [expert_count]bool) ?usize { + var empty: ?usize = null; + var coldest: ?usize = null; + for (self.line_ids, 0..) |id, line| { + if (id < 0) { + if (empty == null) empty = line; + continue; + } + if (keep[@intCast(id)]) continue; + if (coldest == null or colder(self, line, coldest.?)) coldest = line; + } + return empty orelse coldest; + } +}; + +pub const ExactState = struct { + allocator: std.mem.Allocator, + lines: usize, + words: []u32, + + pub fn init(allocator: std.mem.Allocator, lines: usize, seed: []const u32) !ExactState { + const entry = layout.PoolState.init(lines, 0) catch return error.InvalidSeed; + if (seed.len != entry.words) return error.InvalidSeed; + var reverse: [expert_count]u32 = @splat(std.math.maxInt(u32)); + for (seed[0..lines], 0..) |raw, line| { + const id: i32 = @bitCast(raw); + if (id < -1 or id >= expert_count) return error.InvalidSeed; + if (id >= 0) { + const expert: usize = @intCast(id); + if (reverse[expert] != std.math.maxInt(u32)) return error.InvalidSeed; + reverse[expert] = @intCast(line); + } + } + if (!std.mem.eql(u32, &reverse, seed[2 * lines ..][0..expert_count])) return error.InvalidSeed; + const misc = seed[layout.miscOffset(lines)..]; + if (misc[8] >= layout.recent_window or misc[9] > layout.recent_window) return error.InvalidSeed; + var recent_counts: [expert_count]u32 = @splat(0); + for (seed[layout.recentOffset(lines)..layout.miscOffset(lines)], 0..) |bits, index| { + for (0..32) |bit| { + if (bits & (@as(u32, 1) << @intCast(bit)) != 0) recent_counts[(index % 16) * 32 + bit] += 1; + } + } + if (!std.mem.eql(u32, &recent_counts, seed[layout.recentCountsOffset(lines)..][0..expert_count])) return error.InvalidSeed; + for (recent_counts) |count| if (count > misc[9]) return error.InvalidSeed; + return .{ .allocator = allocator, .lines = lines, .words = try allocator.dupe(u32, seed) }; + } + + pub fn deinit(self: *ExactState) void { + self.allocator.free(self.words); + self.* = undefined; + } + + pub fn replay(self: *ExactState, record: trace.Round) !Delta { + const lines = self.lines; + const ids = self.words[0..lines]; + const stamps = self.words[lines .. 2 * lines]; + const reverse = self.words[2 * lines ..][0..expert_count]; + const counts = self.words[2 * lines + expert_count ..][0..expert_count]; + const recent = self.words[layout.recentCountsOffset(lines)..][0..expert_count]; + const history = self.words[layout.recentOffset(lines)..layout.miscOffset(lines)]; + const misc = self.words[layout.miscOffset(lines)..]; + misc[0] +%= 1; + misc[1] +%= 1; + if (record.recentEnabled()) { + for (0..record.rows()) |row| { + const expired = history[misc[8] * 16 ..][0..16]; + for (expired, 0..) |bits, word| { + for (0..32) |bit| { + if (bits & (@as(u32, 1) << @intCast(bit)) != 0) recent[word * 32 + bit] -%= 1; + } + } + @memset(expired, 0); + for (record.picks()[row * 11 ..][0..11]) |raw| { + const expert: i32 = @bitCast(raw); + if (expert < 0 or expert >= expert_count) continue; + const index: usize = @intCast(expert); + const bit = @as(u32, 1) << @intCast(index & 31); + if (expired[index >> 5] & bit == 0) { + expired[index >> 5] |= bit; + recent[index] += 1; + } + } + misc[8] = (misc[8] + 1) % @as(u32, layout.recent_window); + misc[9] = @min(misc[9] + 1, @as(u32, layout.recent_window)); + } + } + var delta: Delta = .{}; + var ordered: [max_unique]usize = undefined; + var appearances: [expert_count]u16 = @splat(0); + var keep: [expert_count]bool = @splat(false); + var unique: usize = 0; + var overflow = false; + for (record.picks()) |raw| { + const expert: i32 = @bitCast(raw); + if (expert < 0 or expert >= expert_count) continue; + const index: usize = @intCast(expert); + counts[index] +|= 1; + if (reverse[index] != std.math.maxInt(u32)) delta.hit_picks += 1 else delta.miss_picks += 1; + if (!keep[index]) { + if (unique == max_unique) { + overflow = true; + continue; + } + keep[index] = true; + ordered[unique] = index; + unique += 1; + } + appearances[index] +%= 1; + } + var misses: [max_unique]usize = undefined; + var missing: usize = 0; + for (ordered[0..unique]) |expert| { + if (reverse[expert] != std.math.maxInt(u32)) { + stamps[reverse[expert]] = misc[0]; + } else { + misses[missing] = expert; + missing += 1; + } + } + const scores = if (record.recentEnabled()) recent else counts; + var cursor: usize = 1; + while (cursor < missing) : (cursor += 1) { + const item = misses[cursor]; + var insert = cursor; + while (insert > 0) : (insert -= 1) { + const previous = misses[insert - 1]; + const hotter_item = if (scores[item] != scores[previous]) scores[item] > scores[previous] else if (appearances[item] != appearances[previous]) appearances[item] > appearances[previous] else item < previous; + if (!hotter_item) break; + misses[insert] = previous; + } + misses[insert] = item; + } + if (overflow) { + delta.overflow = 1; + } else for (misses[0..missing]) |expert| { + var victim: ?usize = null; + for (ids, 0..) |raw, line| { + if (raw == std.math.maxInt(u32)) { + victim = line; + break; + } + if (keep[raw]) continue; + if (victim) |previous| { + const old = ids[previous]; + if (scores[raw] > scores[old] or (scores[raw] == scores[old] and (stamps[line] > stamps[previous] or (stamps[line] == stamps[previous] and line < previous)))) continue; + } + victim = line; + } + if (victim) |line| { + if (ids[line] != std.math.maxInt(u32)) { + reverse[ids[line]] = std.math.maxInt(u32); + delta.evictions += 1; + } + ids[line] = @intCast(expert); + reverse[expert] = @intCast(line); + stamps[line] = misc[0]; + delta.copies += 1; + } else { + delta.inplace += 1; + for (record.picks()) |raw| { + if (raw == expert) delta.mapped_picks += 1; + } + } + } + inline for (.{ .{ 2, "hit_picks" }, .{ 3, "miss_picks" }, .{ 4, "copies" }, .{ 5, "evictions" }, .{ 6, "overflow" }, .{ 7, "inplace" } }) |field| { + misc[field[0]] +%= @truncate(@field(delta, field[1])); + } + return delta; + } +}; + +pub fn verifyTrace(allocator: std.mem.Allocator, view: trace.View) !Verification { + return verifyTraceWithFinal(allocator, view, null); +} + +pub fn verifyTraceWithFinal(allocator: std.mem.Allocator, view: trace.View, final_words: ?[]const u32) !Verification { + if (view.dropped() != 0 or view.droppedCheckpoints() != 0) return error.TraceDroppedRecords; + if (view.words[12] != 0) return error.TraceDiscontinuous; + if (view.words[13] != 0) return error.TraceInvalidLaunch; + if (view.recorded() == 0) return error.EmptyTrace; + var state = try ExactState.init(allocator, view.words[4], view.initialSeed()); + defer state.deinit(); + var total: Delta = .{}; + var checkpoint_index: usize = 0; + for (0..view.recorded() + 1) |index| { + while (checkpoint_index < view.checkpoints()) { + const checkpoint = try view.checkpoint(checkpoint_index); + if (checkpoint.roundIndex() != index) break; + if (!std.mem.eql(u32, state.words, checkpoint.before())) return error.PlannerCheckpointMismatch; + var incoming = try ExactState.init(allocator, state.lines, checkpoint.after()); + std.mem.swap(ExactState, &state, &incoming); + incoming.deinit(); + checkpoint_index += 1; + } + if (index == view.recorded()) break; + const record = try view.round(index); + const delta = try state.replay(record); + if (!std.mem.eql(u32, state.words[layout.miscOffset(state.lines)..], record.misc()) or delta.mapped_picks != record.mappedPicks()) return error.PlannerCounterMismatch; + try add(&total, delta); + } + if (!std.mem.eql(u32, state.words, view.terminalSeed())) return error.PlannerFinalStateMismatch; + if (final_words) |expected| { + if (!std.mem.eql(u32, state.words, expected)) return error.PlannerFinalStateMismatch; + } + if (total.overflow != 0) return error.PlannerOverflow; + return .{ .rounds = view.recorded(), .matched_transitions = view.recorded(), .delta = total }; +} + +fn add(total: *Delta, delta: Delta) !void { + inline for (.{ "hit_picks", "miss_picks", "copies", "evictions", "inplace", "mapped_picks", "overflow" }) |field| { + @field(total, field) = try std.math.add(u64, @field(total, field), @field(delta, field)); + } +} + +fn hotter(state: *const State, a: Entry, b: Entry) bool { + const a_count = state.counts[@intCast(a.expert)]; + const b_count = state.counts[@intCast(b.expert)]; + if (a_count != b_count) return a_count > b_count; + if (a.appearances != b.appearances) return a.appearances > b.appearances; + if (a.expert != b.expert) return a.expert < b.expert; + return false; +} + +fn colder(state: *const State, a: usize, b: usize) bool { + const a_id: usize = @intCast(state.line_ids[a]); + const b_id: usize = @intCast(state.line_ids[b]); + if (state.counts[a_id] != state.counts[b_id]) return state.counts[a_id] < state.counts[b_id]; + if (state.stamps[a] != state.stamps[b]) return state.stamps[a] < state.stamps[b]; + return a > b; +} + +test "planner replay reproduces current lifetime-count admission and eviction" { + var counts: [expert_count]u32 = @splat(0); + counts[0] = 5; + counts[1] = 3; + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = 2 }); + defer state.deinit(); + + const delta = try state.replay(.{ .rows = 1, .valid_picks = 6, .entries = &.{ .{ .expert = 2, .appearances = 5 }, .{ .expert = 3, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .miss_picks = 6, .copies = 2, .evictions = 2 }, delta); + try std.testing.expectEqualSlices(i32, &.{ 3, 2 }, state.line_ids); + try std.testing.expectEqualSlices(u32, &.{ 3, 3 }, state.stamps); + try std.testing.expectEqual(@as(u32, 5), state.counts[2]); +} + +test "planner replay counts repeated slot misses and preserves held lines under overflow pressure" { + var counts: [expert_count]u32 = @splat(0); + counts[0] = 9; + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{0}, .stamps = &.{5}, .counts = counts, .now = 5 }); + defer state.deinit(); + + const delta = try state.replay(.{ .rows = 1, .valid_picks = 4, .entries = &.{ .{ .expert = 0, .appearances = 3 }, .{ .expert = 1, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .hit_picks = 3, .miss_picks = 1, .inplace = 1, .mapped_picks = 1 }, delta); + try std.testing.expectEqualSlices(i32, &.{0}, state.line_ids); + try std.testing.expectEqual(@as(u32, 6), state.stamps[0]); +} + +test "planner replay consumes empty lines first and rejects incomplete evidence" { + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{ -1, -1 }, .stamps = &.{ 90, 91 }, .counts = @splat(0), .now = 0 }); + defer state.deinit(); + const delta = try state.replay(.{ .rows = 1, .valid_picks = 2, .entries = &.{ .{ .expert = 10, .appearances = 1 }, .{ .expert = 11, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .miss_picks = 2, .copies = 2 }, delta); + try std.testing.expectEqualSlices(i32, &.{ 10, 11 }, state.line_ids); + try std.testing.expectError(error.InvalidRound, state.replay(.{ .rows = 1, .valid_picks = 1, .entries = &.{.{ .expert = 12, .appearances = 1 }}, .planner_overflow = true })); +} diff --git a/zig/src/families/flashnext/exl3_dplan_trace.zig b/zig/src/families/flashnext/exl3_dplan_trace.zig new file mode 100644 index 0000000000..791eb1a24a --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_trace.zig @@ -0,0 +1,198 @@ +//! DPLAN v4: default off, little-endian u32 words, fixed 1 MiB per device pool. +//! Header: magic, version, pool, layer, lines, max_rows, state_words, record_words, +//! capacity, recorded, dropped low/high, discontinuity, invalid launch, checkpoints, +//! dropped checkpoints. Round records grow forward; ownership checkpoints grow +//! backward from the buffer end. Each checkpoint is {round index, kind=1, +//! state_words, ordinal}, followed by the outgoing and incoming full states. +//! One full initial seed and one terminal checkpoint precede fixed-size records. +//! Records: rows, valid picks, unique experts, recent flag, ten POST-plan misc words, +//! mapped pick count, then max_rows * 11 raw i32 pick bits in original row order. +//! A pre-state unequal to the previous checkpoint invalidates the timeline. +//! First capacity rounds are retained; excess rounds saturate a u64 drop counter. +//! Counters describe planner decisions, not completed weight reads. Submissions +//! and joined diagnostic reads must obey DevPlan's serialized compute timeline. + +const std = @import("std"); +const layout = @import("exl3_dplan_layout.zig"); + +pub const magic: u32 = 0x44505431; +pub const version: u32 = 4; +pub const pool_bytes: usize = 1024 * 1024; +pub const pool_words: usize = pool_bytes / 4; +pub const header_words: usize = 16; +pub const record_header_words: usize = 4 + layout.misc_words + 1; +pub const checkpoint_header_words: usize = 4; + +pub fn checkpointWords(state_words: usize) usize { + return checkpoint_header_words + 2 * state_words; +} + +pub fn seedFromState(comptime Seed: type, lines: usize, seed: []const u32) Seed { + return .{ .line_ids = @ptrCast(seed[0..lines]), .stamps = seed[lines .. 2 * lines], .counts = seed[2 * lines + layout.experts_n ..][0..layout.experts_n].*, .now = seed[layout.miscOffset(lines)] }; +} + +pub const Checkpoint = struct { + words: []const u32, + + pub fn roundIndex(self: Checkpoint) usize { + return self.words[0]; + } + pub fn before(self: Checkpoint) []const u32 { + return self.words[checkpoint_header_words..][0..self.words[2]]; + } + pub fn after(self: Checkpoint) []const u32 { + return self.words[checkpoint_header_words + self.words[2] ..][0..self.words[2]]; + } +}; + +pub fn enabledValue(value: ?[]const u8) bool { + return if (value) |value_text| std.mem.eql(u8, value_text, "1") else false; +} + +pub fn header(pool_index: usize, layer: usize, lines: usize, max_rows: usize) ![header_words]u32 { + if (max_rows == 0 or max_rows > pool_words / 11) return error.TraceShapeTooLarge; + const entry = layout.PoolState.init(lines, 0) catch return error.TraceShapeTooLarge; + if (entry.words > (pool_words - header_words) / 2) return error.TraceShapeTooLarge; + const stride = record_header_words + max_rows * 11; + const available = pool_words - header_words - 2 * entry.words; + if (stride > available or pool_index > std.math.maxInt(u32) or layer > std.math.maxInt(u32)) return error.TraceShapeTooLarge; + return .{ magic, version, @intCast(pool_index), @intCast(layer), @intCast(lines), @intCast(max_rows), @intCast(entry.words), @intCast(stride), @intCast(available / stride), 0, 0, 0, 0, 0, 0, 0 }; +} + +pub const Round = struct { + words: []const u32, + lines: usize, + max_rows: usize, + + pub fn rows(self: Round) usize { + return self.words[0]; + } + pub fn validPicks(self: Round) u32 { + return self.words[1]; + } + pub fn uniqueCount(self: Round) usize { + return self.words[2]; + } + pub fn recentEnabled(self: Round) bool { + return self.words[3] == 1; + } + pub fn misc(self: Round) []const u32 { + return self.words[4..][0..layout.misc_words]; + } + pub fn mappedPicks(self: Round) u32 { + return self.words[4 + layout.misc_words]; + } + pub fn picks(self: Round) []const u32 { + return self.words[record_header_words..][0 .. self.rows() * 11]; + } + + fn validate(self: Round) !void { + if (self.rows() == 0 or self.rows() > self.max_rows or self.uniqueCount() > layout.experts_n or self.words[3] > 1) return error.InvalidTrace; + var seen: [layout.experts_n]bool = @splat(false); + var unique: usize = 0; + var valid: u32 = 0; + for (self.picks()) |bits| { + const expert: i32 = @bitCast(bits); + if (expert < 0 or expert >= layout.experts_n) continue; + const index: usize = @intCast(expert); + if (!seen[index]) unique += 1; + seen[index] = true; + valid += 1; + } + if (unique != self.uniqueCount() or valid != self.validPicks() or self.mappedPicks() > valid) return error.InvalidTrace; + if (self.misc()[8] >= layout.recent_window or self.misc()[9] > layout.recent_window) return error.InvalidTrace; + } +}; + +pub const View = struct { + words: []const u32, + + pub fn init(words: []const u32) !View { + if (words.len < header_words or words.len > pool_words or words[0] != magic or words[1] != version) return error.InvalidTrace; + const expected = header(words[2], words[3], words[4], words[5]) catch return error.InvalidTrace; + if (!std.mem.eql(u32, words[0..9], expected[0..9]) or words[9] > words[8]) return error.InvalidTrace; + if (words[12] > 1 or words[13] > 1) return error.InvalidTrace; + const self: View = .{ .words = words }; + const stride = checkpointWords(words[6]); + if (self.checkpoints() > (pool_words - self.recordsOffset()) / stride) return error.InvalidTrace; + const records_end = self.recordsOffset() + self.recorded() * words[7]; + const checkpoints_start = pool_words - self.checkpoints() * stride; + if (records_end > checkpoints_start) return error.InvalidTrace; + if (words.len < self.usedWords()) return error.InvalidTrace; + if (self.dropped() != 0 and self.droppedCheckpoints() == 0 and checkpoints_start - records_end >= words[7]) return error.InvalidTrace; + if (self.droppedCheckpoints() != 0 and self.dropped() == 0 and checkpoints_start - records_end >= stride) return error.InvalidTrace; + var previous_round: usize = 0; + for (0..self.checkpoints()) |index| { + const entry = try self.checkpoint(index); + if (entry.words[1] != 1 or entry.words[2] != words[6] or entry.words[3] != index or entry.roundIndex() > self.recorded() or entry.roundIndex() < previous_round) return error.InvalidTrace; + previous_round = entry.roundIndex(); + } + for (0..self.recorded()) |index| try (try self.round(index)).validate(); + return self; + } + + pub fn recorded(self: View) usize { + return self.words[9]; + } + pub fn dropped(self: View) u64 { + return @as(u64, self.words[10]) | (@as(u64, self.words[11]) << 32); + } + pub fn checkpoints(self: View) usize { + return self.words[14]; + } + pub fn droppedCheckpoints(self: View) u32 { + return self.words[15]; + } + pub fn checkpoint(self: View, index: usize) !Checkpoint { + if (index >= self.checkpoints()) return error.TraceCheckpointOutOfRange; + const stride = checkpointWords(self.words[6]); + return .{ .words = self.words[pool_words - (index + 1) * stride ..][0..stride] }; + } + pub fn initialSeed(self: View) []const u32 { + return self.words[header_words..][0..self.words[6]]; + } + pub fn terminalSeed(self: View) []const u32 { + return self.words[header_words + self.words[6] ..][0..self.words[6]]; + } + pub fn seedState(self: View, comptime Seed: type) Seed { + return seedFromState(Seed, self.words[4], self.initialSeed()); + } + pub fn recordsOffset(self: View) usize { + return header_words + 2 * @as(usize, self.words[6]); + } + pub fn round(self: View, index: usize) !Round { + if (index >= self.recorded()) return error.TraceRoundOutOfRange; + const stride = self.words[7]; + return .{ .words = self.words[self.recordsOffset() + index * stride ..][0..stride], .lines = self.words[4], .max_rows = self.words[5] }; + } + fn usedWords(self: View) usize { + if (self.checkpoints() != 0) return pool_words; + return self.recordsOffset() + self.recorded() * self.words[7]; + } + pub fn encode(self: View, allocator: std.mem.Allocator) ![]u8 { + const bytes = try allocator.alloc(u8, self.usedWords() * 4); + for (self.words[0..self.usedWords()], 0..) |word, index| std.mem.writeInt(u32, bytes[index * 4 ..][0..4], word, .little); + return bytes; + } +}; + +pub const Snapshot = struct { + allocator: std.mem.Allocator, + words: []u32, + + pub fn deinit(self: *Snapshot) void { + self.allocator.free(self.words); + self.* = undefined; + } + pub fn view(self: Snapshot) !View { + return View.init(self.words); + } + pub fn decode(allocator: std.mem.Allocator, bytes: []const u8) !Snapshot { + if (bytes.len % 4 != 0 or bytes.len > pool_bytes or bytes.len < header_words * 4) return error.InvalidTrace; + const words = try allocator.alloc(u32, bytes.len / 4); + errdefer allocator.free(words); + for (words, 0..) |*word, index| word.* = std.mem.readInt(u32, bytes[index * 4 ..][0..4], .little); + _ = try View.init(words); + return .{ .allocator = allocator, .words = words }; + } +}; diff --git a/zig/src/families/flashnext/exl3_dplan_trace_kernel_test.cpp b/zig/src/families/flashnext/exl3_dplan_trace_kernel_test.cpp new file mode 100644 index 0000000000..855468a8a4 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_trace_kernel_test.cpp @@ -0,0 +1,261 @@ +#include +#include +#include +#include +#include +#include + +#define __device__ +#define __global__ +#define __shared__ +#define __forceinline__ inline +#define __launch_bounds__(value) +struct TestDim { int x; }; +static TestDim threadIdx{0}, blockIdx{0}, blockDim{1}, gridDim{1}; +struct uint4 { unsigned x, y, z, w; }; +using cudaGraphConditionalHandle = unsigned long long; +static void __syncthreads() {} +static int __ffs(unsigned value) { return __builtin_ffs(value); } +static int __float_as_int(float value) { + int bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} +static void cudaGraphSetConditional(cudaGraphConditionalHandle, unsigned) {} + +#include "../../../kernels/cuda/flashnext_exl3_dplan.cu" + +static constexpr unsigned pool_words = 1024 * 1024 / 4; +static constexpr unsigned sentinel = 0xfeedface; + +static std::vector make_trace(unsigned lines, unsigned rows) { + std::vector trace(pool_words + 1, 0); + const unsigned seed_words = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10; + const unsigned stride = 15 + rows * 11; + auto* header = reinterpret_cast(trace.data()); + *header = {0x44505431, 4, 0, 0, lines, rows, seed_words, stride, + (pool_words - 16 - 2 * seed_words) / stride, 0, 0, 0, 0, 0, 0, 0}; + trace[pool_words] = sentinel; + return trace; +} + +static std::vector make_state(unsigned lines) { + std::vector state(2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10, 0); + for (unsigned line = 0; line < lines; ++line) state[line] = 0xffffffffu; + for (unsigned expert = 0; expert < DP_EXPERTS; ++expert) state[2 * lines + expert] = 0xffffffffu; + return state; +} + +struct PlanOutput { + long long tabs[3][513]{}; + int widths[3][513]{}; + long long work[512 * 9]{}; + int work_n = 0, overflow[512]{}, overflow_n = 0; +}; + +static void run_plan(std::vector& state, unsigned lines, const std::vector& picks, + unsigned* trace, int recent, PlanOutput& output) { + long long src[512 * 3], lens[512 * 3]; + int widths[512 * 3]; + for (int index = 0; index < 512 * 3; ++index) { + src[index] = 0x100000 + index * 256; + lens[index] = 256; + widths[index] = 3; + } + fn_exl3_dplan(picks.data(), static_cast(picks.size()), state.data(), lines, recent, + output.tabs[0], output.tabs[1], output.tabs[2], + output.widths[0], output.widths[1], output.widths[2], + src, lens, widths, 1024, 1024, 0, 256, 512, + output.work, &output.work_n, nullptr, nullptr, output.overflow, &output.overflow_n, trace); +} + +static void assert_same(const PlanOutput& first, const PlanOutput& second) { + assert(std::memcmp(first.tabs, second.tabs, sizeof(first.tabs)) == 0); + assert(std::memcmp(first.widths, second.widths, sizeof(first.widths)) == 0); + assert(std::memcmp(first.work, second.work, sizeof(first.work)) == 0); + assert(std::memcmp(first.overflow, second.overflow, sizeof(first.overflow)) == 0); + assert(first.work_n == second.work_n && first.overflow_n == second.overflow_n); +} + +static void test_bounds_and_flags() { + auto trace = make_trace(2, 47); + auto* header = reinterpret_cast(trace.data()); + auto state = make_state(2); + std::vector picks(47 * 11, 512); + for (int index = 0; index < 512; ++index) picks[index] = 511 - index; + picks[512] = -1; + picks[513] = 511; + PlanOutput output; + const auto seed = state; + run_plan(state, 2, picks, trace.data(), 1, output); + const unsigned offset = 16 + 2 * header->seed_words; + const unsigned* round = trace.data() + offset; + assert(round[0] == 47 && round[1] == 513 && round[2] == 512 && round[3] == 1); + assert(round[10] == 1 && round[8] == 0 && round[14] == 0); + assert(std::memcmp(trace.data() + 16, seed.data(), seed.size() * 4) == 0); + assert(std::memcmp(round + 15, picks.data(), picks.size() * 4) == 0); + assert(std::memcmp(trace.data() + 16 + seed.size(), state.data(), state.size() * 4) == 0); + const auto first = std::vector(round, round + header->record_words); + const unsigned capacity = header->capacity; + for (unsigned index = 1; index < capacity + 3; ++index) + run_plan(state, 2, picks, trace.data(), 1, output); + assert(header->recorded == capacity && header->dropped_lo == 3 && header->dropped_hi == 0); + assert(std::memcmp(trace.data() + offset, first.data(), first.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + const auto terminal = std::vector(trace.data() + 16 + seed.size(), trace.data() + offset); + header->dropped_lo = 0xffffffffu; + header->dropped_hi = 4; + run_plan(state, 2, picks, trace.data(), 0, output); + assert(header->dropped_lo == 0 && header->dropped_hi == 5); + header->dropped_lo = header->dropped_hi = 0xffffffffu; + run_plan(state, 2, picks, trace.data(), 0, output); + assert(header->dropped_lo == 0xffffffffu && header->dropped_hi == 0xffffffffu); + assert(std::memcmp(trace.data() + 16 + seed.size(), terminal.data(), terminal.size() * 4) == 0); + assert(dp_trace_begin(nullptr, 0, nullptr, nullptr, 0, 2) == nullptr); + auto invalid = make_trace(2, 1); + auto* invalid_header = reinterpret_cast(invalid.data()); + assert(dp_trace_begin(picks.data(), 12, invalid.data(), state.data(), 0, 2) == nullptr); + assert(invalid_header->invalid_launch == 1 && invalid_header->recorded == 0); + auto broken = make_trace(2, 47); + run_plan(state, 2, picks, broken.data(), 0, output); + state.back() ^= 1; + run_plan(state, 2, picks, broken.data(), 0, output); + assert(reinterpret_cast(broken.data())->discontinuity == 1); + + for (unsigned field : {0u, 1u, 4u, 5u, 6u, 7u, 8u, 9u, 12u, 13u, 14u}) { + auto corrupt = make_trace(2, 1); + corrupt[field] = 0xffffffffu; + assert(dp_trace_begin(nullptr, 11, corrupt.data(), nullptr, 0, 2) == nullptr); + assert(corrupt[13] == 1 && corrupt[pool_words] == sentinel); + corrupt[field] = 0xffffffffu; + fn_exl3_dcheckpoint(nullptr, 2, corrupt.data(), 1); + assert(corrupt[pool_words] == sentinel); + } + auto full = make_trace(2, 16); + auto* full_header = reinterpret_cast(full.data()); + const unsigned checkpoint_words = 4 + 2 * full_header->seed_words; + const unsigned checkpoint_capacity = (pool_words - 16 - 2 * full_header->seed_words) / checkpoint_words; + for (unsigned index = 0; index < checkpoint_capacity + 1; ++index) + fn_exl3_dcheckpoint(state.data(), 2, full.data(), 1); + assert(full_header->checkpoints == checkpoint_capacity && full_header->dropped_checkpoints == 1); + full_header->dropped_checkpoints = 0xffffffffu; + fn_exl3_dcheckpoint(state.data(), 2, full.data(), 1); + assert(full_header->dropped_checkpoints == 0xffffffffu && full[pool_words] == sentinel); + run_plan(state, 2, std::vector(11, -1), full.data(), 0, output); + assert(full_header->recorded == 0 && full_header->dropped_lo == 1); + + auto sealed = make_trace(2, 16); + fn_exl3_dcheckpoint(state.data(), 2, sealed.data(), 1); + state.back() ^= 1; + fn_exl3_dcheckpoint(state.data(), 2, sealed.data(), 0); + assert(reinterpret_cast(sealed.data())->discontinuity == 1); +} + +static std::vector test_checkpoint_capture() { + const unsigned lines = 10; + auto trace = make_trace(lines, 16); + auto* header = reinterpret_cast(trace.data()); + auto state = make_state(lines); + auto untraced = state; + PlanOutput captured, disabled; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + for (unsigned round = 0; round < 1024; ++round) { + if (round == 512 || round == 768) { + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 0); + for (unsigned host_round = 0; host_round < 3; ++host_round) { + std::vector host_picks(11, static_cast((round + host_round * 13) % 512)); + run_plan(state, lines, host_picks, nullptr, 0, captured); + run_plan(untraced, lines, host_picks, nullptr, 0, disabled); + } + const unsigned misc = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS; + for (unsigned counter = 1; counter < 8; ++counter) + state[misc + counter] = untraced[misc + counter] = 0; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + } + std::vector picks(11, static_cast((round * 7) % 512)); + picks.back() = 512; + run_plan(state, lines, picks, trace.data(), round % 3 != 0, captured); + run_plan(untraced, lines, picks, nullptr, round % 3 != 0, disabled); + assert(state == untraced); + assert_same(captured, disabled); + } + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 0); + auto* misc = reinterpret_cast(state.data() + 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS); + misc->now += 100; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + assert(header->recorded == 1024 && header->checkpoints == 4); + assert(header->dropped_lo == 0 && header->dropped_hi == 0 && header->dropped_checkpoints == 0); + assert(header->discontinuity == 0 && header->invalid_launch == 0); + assert(std::memcmp(trace.data() + 16 + header->seed_words, state.data(), state.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + return trace; +} + +static std::vector test_long_capture() { + const unsigned lines = 10; + auto trace = make_trace(lines, 16); + auto* header = reinterpret_cast(trace.data()); + assert(header->capacity >= 1200); + assert(16 + 2 * header->seed_words + header->capacity * header->record_words <= pool_words); + auto state = make_state(lines); + const unsigned misc_at = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS; + for (unsigned index = 0; index < 8; ++index) state[misc_at + index] = 0xfffffffeu; + state[2 * lines + DP_EXPERTS + 7] = 0xffffffffu; + auto untraced = state; + const auto initial = state; + PlanOutput captured, disabled; + for (int index = 0; index < 1200; ++index) { + std::vector picks(16 * 11); + for (unsigned slot = 0; slot < picks.size(); ++slot) { + switch (slot % 11) { + case 0: picks[slot] = -1; break; + case 1: picks[slot] = 512; break; + case 2: picks[slot] = -99; break; + case 3: picks[slot] = 800; break; + default: picks[slot] = (index * 7 + slot * 3) % 512; break; + } + } + const int recent = index % 3 != 0; + run_plan(state, lines, picks, trace.data(), recent, captured); + run_plan(untraced, lines, picks, nullptr, recent, disabled); + assert(state == untraced); + assert_same(captured, disabled); + assert(header->recorded == static_cast(index + 1)); + const unsigned* round = trace.data() + 16 + 2 * header->seed_words + index * header->record_words; + assert(std::memcmp(round + 4, state.data() + misc_at, sizeof(DpMisc)) == 0); + unsigned mapped = 0; + for (int pick : picks) + for (int overflow = 0; overflow < captured.overflow_n; ++overflow) + if (pick == captured.overflow[overflow]) ++mapped; + assert(round[14] == mapped); + assert(std::memcmp(round + 15, picks.data(), picks.size() * 4) == 0); + } + assert(header->discontinuity == 0 && header->invalid_launch == 0); + assert(std::memcmp(trace.data() + 16, initial.data(), initial.size() * 4) == 0); + assert(std::memcmp(trace.data() + 16 + header->seed_words, state.data(), state.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + return trace; +} + +static void write_trace(const std::vector& trace, const char* path) { + const auto* header = reinterpret_cast(trace.data()); + const unsigned used = header->checkpoints ? pool_words : 16 + 2 * header->seed_words + header->recorded * header->record_words; + FILE* output = std::fopen(path, "wb"); + assert(output != nullptr); + for (unsigned index = 0; index < used; ++index) { + const unsigned word = trace[index]; + const unsigned char bytes[4] = { + static_cast(word), static_cast(word >> 8), + static_cast(word >> 16), static_cast(word >> 24)}; + assert(std::fwrite(bytes, 1, 4, output) == 4); + } + assert(std::fclose(output) == 0); +} + +int main(int argc, char** argv) { + test_bounds_and_flags(); + const auto trace = test_long_capture(); + const auto checkpoint_trace = test_checkpoint_capture(); + if (argc >= 2) write_trace(trace, argv[1]); + if (argc >= 3) write_trace(checkpoint_trace, argv[2]); +} diff --git a/zig/src/families/flashnext/exl3_dplan_trace_replay_test.zig b/zig/src/families/flashnext/exl3_dplan_trace_replay_test.zig new file mode 100644 index 0000000000..4c5a41137d --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_trace_replay_test.zig @@ -0,0 +1,294 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); +const replay = @import("exl3_dplan_replay.zig"); +const policies = @import("exl3_dplan_policy_eval.zig"); +const fixture = @import("exl3_dplan_trace_test.zig"); + +fn capture(words: []u32, state: *replay.ExactState, picks: []const i32, recent: bool) !replay.Delta { + const before = try std.testing.allocator.dupe(u32, state.words); + defer std.testing.allocator.free(before); + const record_words = try std.testing.allocator.alloc(u32, words[7]); + defer std.testing.allocator.free(record_words); + @memset(record_words, 0); + record_words[0] = @intCast(picks.len / 11); + record_words[3] = @intFromBool(recent); + for (picks, 0..) |pick, index| record_words[trace.record_header_words + index] = @bitCast(pick); + const delta = try state.replay(.{ .words = record_words, .lines = state.lines, .max_rows = words[5] }); + fixture.append(words, before, state.words, picks, recent, @intCast(delta.mapped_picks)); + return delta; +} + +test "exl3: ownership checkpoints verify every segment including repeated and final transfers" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &.{ 2, 2, 3, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + @memcpy(state.words, seed); + state.words[0] = 77; + state.words[4 + 77] = 0; + state.words[4 + layout.experts_n + 77] = 90; + state.words[layout.miscOffset(2)] = 500; + fixture.checkpoint(words, state.words); + state.words[layout.miscOffset(2)] += 10; + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &.{ 77, 77, 88, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + state.words[layout.miscOffset(2)] += 10; + fixture.checkpoint(words, state.words); + const view = try trace.View.init(words); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(@as(usize, 2), result.matched_transitions); + try std.testing.expectEqual(@as(u64, 2), result.delta.hit_picks); + try std.testing.expectEqual(@as(u64, 4), result.delta.miss_picks); + const comparison = try policies.evaluateTrace(allocator, view); + try std.testing.expectEqual(result.delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(result.delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(result.delta.copies, comparison.lifetime.copies); + try std.testing.expectEqual(result.delta.evictions, comparison.lifetime.evictions); + try std.testing.expectEqual(result.delta.inplace, comparison.lifetime.mapped_experts); + try std.testing.expectEqual(result.delta.mapped_picks, comparison.lifetime.mapped_picks); + try std.testing.expectEqual(result.rounds, comparison.lifetime.rounds); + try std.testing.expectEqual(@as(usize, 4), view.checkpoints()); + for ([_]usize{ 0, 1, 1, 2 }, 0..) |round, index| try std.testing.expectEqual(round, (try view.checkpoint(index)).roundIndex()); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectEqual(trace.pool_bytes, bytes.len); + var snapshot = try trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + _ = try replay.verifyTraceWithFinal(allocator, try snapshot.view(), state.words); + state.words[0] ^= 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTraceWithFinal(allocator, view, state.words)); +} + +test "exl3: unexplained handoff and corrupted outgoing checkpoint cannot be reseeded away" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + const picks: [11]i32 = @splat(-1); + _ = try capture(words, &state, &picks, false); + state.words[layout.miscOffset(2)] += 100; + _ = try capture(words, &state, &picks, false); + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 0; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 1; + fixture.checkpoint(words, state.words); + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + + @memset(words, 0); + @memcpy(words[0..trace.header_words], &try trace.header(0, 0, 2, 1)); + @memcpy(state.words, seed); + _ = try capture(words, &state, &picks, false); + state.words[layout.miscOffset(2)] += 100; + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &picks, false); + const view = try trace.View.init(words); + _ = try replay.verifyTrace(allocator, view); + const offset = trace.pool_words - trace.checkpointWords(words[6]); + for (0..words[6]) |index| { + words[offset + 4 + index] ^= 1; + try std.testing.expectError(error.PlannerCheckpointMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[offset + 4 + index] ^= 1; + } + words[offset + 4 + words[6]] = 9999; + try std.testing.expectError(error.InvalidSeed, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: dropped ownership checkpoints invalidate exact replay even without decode rounds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 16); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + const capacity = (trace.pool_words - trace.header_words - 2 * words[6]) / trace.checkpointWords(words[6]); + for (0..capacity + 1) |_| fixture.checkpoint(words, seed); + try std.testing.expectError(error.TraceDroppedRecords, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: v4 exact warm replay matches known decisions counters and all terminal words" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + seed[0] = 0; + seed[1] = 1; + seed[2] = 10; + seed[3] = 20; + seed[4] = 0; + seed[5] = 1; + seed[4 + 512] = 5; + seed[5 + 512] = 3; + seed[layout.miscOffset(2)] = 20; + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + const first = try capture(words, &state, &.{ 2, 2, 2, 2, 2, 3, -1, -1, -1, -1, 512 }, false); + try std.testing.expectEqual(replay.Delta{ .miss_picks = 6, .copies = 2, .evictions = 2 }, first); + try std.testing.expectEqualSlices(u32, &.{ 3, 2, 21, 21 }, state.words[0..4]); + const second = try capture(words, &state, &.{ 2, 2, 4, -1, -1, -1, -1, -1, -1, -1, -1 }, false); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 1, .copies = 1, .evictions = 1 }, second); + const view = try trace.View.init(words); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(@as(usize, 2), result.matched_transitions); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 7, .copies = 3, .evictions = 3 }, result.delta); + const comparison = try policies.evaluate(allocator, &.{ try view.round(0), try view.round(1) }, .{ .capacity = 2, .seed = view.seedState(policies.Seed) }); + try std.testing.expectEqual(result.delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(result.delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(result.delta.copies, comparison.lifetime.copies); + try std.testing.expectEqual(result.delta.evictions, comparison.lifetime.evictions); + state.words[state.words.len - 1] = 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTraceWithFinal(allocator, view, state.words)); +} + +test "exl3: v4 1200 full rounds replay recent wrap saturated counts and modular counters" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 10, 16); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 10); + defer allocator.free(seed); + seed[20 + 512 + 7] = std.math.maxInt(u32); + const misc = layout.miscOffset(10); + @memset(seed[misc .. misc + 8], std.math.maxInt(u32) - 1); + var state = try replay.ExactState.init(allocator, 10, seed); + defer state.deinit(); + var total: replay.Delta = .{}; + for (0..1200) |round| { + var picks: [16 * 11]i32 = undefined; + for (&picks, 0..) |*pick, index| { + pick.* = switch (index % 11) { + 0 => -1, + 1 => 512, + 2 => -99, + 3 => 800, + else => @intCast((round * 7 + index * 3) % 512), + }; + } + const delta = try capture(words, &state, &picks, round % 3 != 0); + inline for (comptime std.meta.fieldNames(replay.Delta)) |field| @field(total, field) += @field(delta, field); + } + const view = try trace.View.init(words); + try std.testing.expectEqual(@as(usize, 1200), view.recorded()); + try std.testing.expectEqual(@as(u64, 0), view.dropped()); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(total, result.delta); + try std.testing.expectEqual(std.math.maxInt(u32), state.words[20 + 512 + 7]); + try std.testing.expectEqual(@as(u32, 1198), state.words[misc]); + try std.testing.expectEqual(@as(u32, 256), state.words[misc + 9]); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expect(bytes.len <= trace.pool_bytes); + var snapshot = try trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + _ = try replay.verifyTrace(allocator, try snapshot.view()); +} + +test "exl3: v4 starts from warm recent history and retains empty masked rounds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + seed[layout.recentCountsOffset(2) + 7] = 2; + seed[layout.recentOffset(2)] = 1 << 7; + seed[layout.recentOffset(2) + 16] = 1 << 7; + seed[layout.miscOffset(2) + 8] = 2; + seed[layout.miscOffset(2) + 9] = 2; + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + _ = try capture(words, &state, &.{ 7, 7, -1, 512, 999, -1, -1, -1, -1, -1, -1 }, true); + try std.testing.expectEqual(@as(u32, 3), state.words[layout.recentCountsOffset(2) + 7]); + try std.testing.expectEqual(@as(u32, 3), state.words[layout.miscOffset(2) + 8]); + const masked = try capture(words, &state, &@as([11]i32, @splat(-1)), true); + try std.testing.expectEqual(replay.Delta{}, masked); + const result = try replay.verifyTraceWithFinal(allocator, try trace.View.init(words), state.words); + try std.testing.expectEqual(@as(usize, 2), result.rounds); +} + +test "exl3: v4 mapped experts count repeated reads and mask shared and invalid slots" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 1, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 1); + defer allocator.free(seed); + seed[0] = 0; + seed[2] = 0; + var state = try replay.ExactState.init(allocator, 1, seed); + defer state.deinit(); + const delta = try capture(words, &state, &.{ 0, 0, 1, 1, 1, 512, -1, 700, -9, 512, -1 }, false); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 3, .inplace = 1, .mapped_picks = 3 }, delta); + const view = try trace.View.init(words); + _ = try replay.verifyTrace(allocator, view); + const comparison = try policies.evaluateTrace(allocator, view); + try std.testing.expectEqual(delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(delta.inplace, comparison.lifetime.mapped_experts); + try std.testing.expectEqual(delta.mapped_picks, comparison.lifetime.mapped_picks); + const offset = (try trace.View.init(words)).recordsOffset(); + words[offset + 14] = 2; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: v4 replay rejects altered production outcomes every terminal field and bad seeds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + _ = try capture(words, &state, &.{ 2, 2, 3, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + const view = try trace.View.init(words); + for (0..layout.misc_words) |index| { + const offset = view.recordsOffset() + 4 + index; + words[offset] ^= 1; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[offset] ^= 1; + } + const terminal_offset = trace.header_words + words[6]; + for (0..words[6]) |index| { + words[terminal_offset + index] ^= 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[terminal_offset + index] ^= 1; + } + for ([_]usize{ 0, 4, layout.recentCountsOffset(2), layout.miscOffset(2) + 8 }) |index| { + const saved = words[trace.header_words + index]; + words[trace.header_words + index] = 9999; + try std.testing.expectError(error.InvalidSeed, replay.verifyTrace(allocator, try trace.View.init(words))); + words[trace.header_words + index] = saved; + } +} + +test "exl3: v4 overflow dropped empty invalid and discontinuous timelines cannot pass" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 47); + defer allocator.free(words); + try std.testing.expectError(error.EmptyTrace, replay.verifyTrace(allocator, try trace.View.init(words))); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + var picks: [47 * 11]i32 = @splat(512); + for (0..512) |index| picks[index] = @intCast(index); + const delta = try capture(words, &state, &picks, false); + try std.testing.expectEqual(@as(u64, 512), delta.miss_picks); + try std.testing.expectEqual(@as(u64, 0), delta.copies); + try std.testing.expectEqual(@as(u64, 1), delta.overflow); + try std.testing.expectError(error.PlannerOverflow, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 1; + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 0; + words[13] = 1; + try std.testing.expectError(error.TraceInvalidLaunch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[13] = 0; + for (1..words[8] + 1) |_| _ = try capture(words, &state, &picks, false); + try std.testing.expectError(error.TraceDroppedRecords, replay.verifyTrace(allocator, try trace.View.init(words))); +} diff --git a/zig/src/families/flashnext/exl3_dplan_trace_test.zig b/zig/src/families/flashnext/exl3_dplan_trace_test.zig new file mode 100644 index 0000000000..bc9ee14235 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_trace_test.zig @@ -0,0 +1,262 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); + +pub fn emptySeed(allocator: std.mem.Allocator, lines: usize) ![]u32 { + const seed = try allocator.alloc(u32, layout.stateWords(lines)); + @memset(seed, 0); + @memset(seed[0..lines], std.math.maxInt(u32)); + @memset(seed[2 * lines ..][0..layout.experts_n], std.math.maxInt(u32)); + return seed; +} + +pub fn make(allocator: std.mem.Allocator, lines: usize, rows: usize) ![]u32 { + const words = try allocator.alloc(u32, trace.pool_words); + errdefer allocator.free(words); + @memset(words, 0); + @memcpy(words[0..trace.header_words], &try trace.header(3, 48, lines, rows)); + return words; +} + +pub fn append(words: []u32, before: []const u32, after: []const u32, picks: []const i32, recent: bool, mapped: u32) void { + const records_end = trace.header_words + 2 * words[6] + words[9] * words[7]; + const checkpoints_start = trace.pool_words - words[14] * trace.checkpointWords(words[6]); + if (checkpoints_start - records_end < words[7] or words[15] != 0) { + const dropped = (@as(u64, words[11]) << 32 | words[10]) +| 1; + words[10] = @truncate(dropped); + words[11] = @truncate(dropped >> 32); + return; + } + const seed = words[trace.header_words..][0..words[6]]; + const terminal = words[trace.header_words + words[6] ..][0..words[6]]; + if (words[9] == 0 and words[14] == 0) { + @memcpy(seed, before); + } else if (!std.mem.eql(u32, terminal, before)) { + words[12] = 1; + } + const offset = trace.header_words + 2 * words[6] + words[9] * words[7]; + const record = words[offset..][0..words[7]]; + @memset(record, 0); + record[0] = @intCast(picks.len / 11); + record[3] = @intFromBool(recent); + @memcpy(record[4..][0..layout.misc_words], after[layout.miscOffset(words[4])..]); + record[14] = mapped; + var seen: [layout.experts_n]bool = @splat(false); + for (picks, 0..) |expert, index| { + record[trace.record_header_words + index] = @bitCast(expert); + if (expert < 0 or expert >= layout.experts_n) continue; + const valid: usize = @intCast(expert); + record[1] += 1; + if (!seen[valid]) record[2] += 1; + seen[valid] = true; + } + @memcpy(terminal, after); + words[9] += 1; +} + +pub fn checkpoint(words: []u32, incoming: []const u32) void { + const stride = trace.checkpointWords(words[6]); + const end = trace.header_words + 2 * words[6] + words[9] * words[7]; + const start = trace.pool_words - words[14] * stride; + if (start - end < stride or words[10] != 0 or words[11] != 0 or words[15] != 0) { + words[15] +|= 1; + return; + } + const initial = words[trace.header_words..][0..words[6]]; + const terminal = words[trace.header_words + words[6] ..][0..words[6]]; + if (words[9] == 0 and words[14] == 0) { + @memcpy(initial, incoming); + @memcpy(terminal, incoming); + } + const entry = words[start - stride .. start]; + entry[0..4].* = .{ words[9], 1, words[6], words[14] }; + @memcpy(entry[4..][0..words[6]], terminal); + @memcpy(entry[4 + words[6] ..], incoming); + @memcpy(terminal, incoming); + words[14] += 1; +} + +test "exl3: checkpoints reject corrupt ordinals shapes order overlap and truncation" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 1); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const picks: [11]i32 = @splat(-1); + append(words, seed, seed, &picks, false, 0); + checkpoint(words, seed); + append(words, seed, seed, &picks, false, 0); + checkpoint(words, seed); + const view = try trace.View.init(words); + try std.testing.expectError(error.TraceCheckpointOutOfRange, view.checkpoint(2)); + const stride = trace.checkpointWords(words[6]); + const offset = trace.pool_words - stride; + for (0..4) |field| { + const saved = words[offset + field]; + words[offset + field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset + field] = saved; + } + words[offset - stride] = 0; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset - stride] = 2; + const saved_count = words[14]; + words[14] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[14] = @intCast((trace.pool_words - view.recordsOffset()) / stride); + words[9] = words[8]; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[14] = saved_count; + words[9] = 2; + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 4])); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words[0 .. words.len - 1])); + words[15] = 1; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); +} + +test "exl3: checkpoint exhaustion saturates without overwriting retained records or states" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 16); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const capacity = (trace.pool_words - trace.header_words - 2 * words[6]) / trace.checkpointWords(words[6]); + for (0..capacity) |_| checkpoint(words, seed); + const retained = try allocator.dupe(u32, words); + defer allocator.free(retained); + seed[layout.miscOffset(2)] += 20; + checkpoint(words, seed); + try std.testing.expectEqual(@as(u32, 1), (try trace.View.init(words)).droppedCheckpoints()); + try std.testing.expectEqualSlices(u32, retained[trace.header_words..], words[trace.header_words..]); + words[15] = std.math.maxInt(u32); + checkpoint(words, seed); + try std.testing.expectEqual(std.math.maxInt(u32), (try trace.View.init(words)).droppedCheckpoints()); + append(words, seed, seed, &@as([11]i32, @splat(-1)), false, 0); + try std.testing.expectEqual(@as(u64, 1), (try trace.View.init(words)).dropped()); +} + +test "exl3: v4 trace preserves raw row order seed checkpoint and post counters" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 2); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + seed[layout.miscOffset(2)] = 0xfffffffe; + const picks = [_]i32{ 7, 2, 7, -1, 512, 4, 2, 2, 0, -9, 800, 4, 7, 2, 511, -1, 512, 4, 0, 0, 7, 511 }; + append(words, seed, seed, &picks, true, 3); + const view = try trace.View.init(words); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + var decoded = try trace.Snapshot.decode(allocator, bytes); + defer decoded.deinit(); + const restored = try decoded.view(); + const record = try restored.round(0); + try std.testing.expectEqual(@as(usize, 2), record.rows()); + try std.testing.expectEqual(@as(u32, 16), record.validPicks()); + try std.testing.expectEqual(@as(usize, 5), record.uniqueCount()); + try std.testing.expect(record.recentEnabled()); + try std.testing.expectEqual(@as(u32, 3), record.mappedPicks()); + try std.testing.expectEqualSlices(u32, seed, restored.initialSeed()); + try std.testing.expectEqualSlices(u32, seed, restored.terminalSeed()); + try std.testing.expectEqual(@as(u32, 0xfffffffe), record.misc()[0]); + for (picks, record.picks()) |expert, bits| try std.testing.expectEqual(expert, @as(i32, @bitCast(bits))); + try std.testing.expectEqual(trace.magic, std.mem.readInt(u32, bytes[0..4], .little)); + try std.testing.expectEqual((trace.header_words + 2 * seed.len + words[7]) * 4, bytes.len); + try std.testing.expectError(error.TraceRoundOutOfRange, view.round(1)); +} + +test "exl3: v4 capacity exceeds 1000 with full 16-row records before ownership checkpoints" { + for ([_]usize{ 1, 10, 73, 256, 512 }) |lines| { + const shape = try trace.header(0, 0, lines, 16); + try std.testing.expect(shape[8] >= 1000); + const used = trace.header_words + 2 * shape[6] + shape[8] * shape[7]; + try std.testing.expect(used <= trace.pool_words); + try std.testing.expect(used + shape[7] > trace.pool_words); + try std.testing.expectEqual(@as(u32, 191), shape[7]); + } + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, trace.pool_words, 1)); + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, 2, std.math.maxInt(usize))); + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, 2, 0)); + try std.testing.expect(!trace.enabledValue(null)); + try std.testing.expect(!trace.enabledValue("true")); + try std.testing.expect(!trace.enabledValue("0")); + try std.testing.expect(trace.enabledValue("1")); +} + +test "exl3: v4 keeps initial seed detects discontinuity and saturates bounded drops" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 16); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const picks: [11]i32 = @splat(-1); + append(words, seed, seed, &picks, false, 0); + seed[layout.miscOffset(2)] = 700; + append(words, seed, seed, &picks, false, 0); + try std.testing.expectEqual(@as(u32, 1), words[12]); + try std.testing.expectEqual(@as(u32, 0), (try trace.View.init(words)).initialSeed()[layout.miscOffset(2)]); + for (2..words[8] + 3) |_| append(words, seed, seed, &picks, false, 0); + var view = try trace.View.init(words); + try std.testing.expectEqual(@as(usize, words[8]), view.recorded()); + try std.testing.expectEqual(@as(u64, 3), view.dropped()); + words[10] = std.math.maxInt(u32); + words[11] = 4; + append(words, seed, seed, &picks, false, 0); + try std.testing.expectEqual(@as(u64, 5) << 32, (try trace.View.init(words)).dropped()); + words[10] = std.math.maxInt(u32); + words[11] = std.math.maxInt(u32); + append(words, seed, seed, &picks, false, 0); + view = try trace.View.init(words); + try std.testing.expectEqual(std.math.maxInt(u64), view.dropped()); +} + +test "exl3: v4 rejects metadata corruption truncation old versions and false pick counts" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 1); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + append(words, seed, seed, &@as([11]i32, @splat(7)), false, 0); + const view = try trace.View.init(words); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 4])); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 1])); + const offset = view.recordsOffset(); + for ([_]usize{ 1, 2, 3, 14 }) |field| { + const saved = words[offset + field]; + words[offset + field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset + field] = saved; + } + for ([_]usize{ 1, 6, 7, 8, 9, 12, 13, 14, 15 }) |field| { + const saved = words[field]; + words[field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[field] = saved; + } + words[1] = 2; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[1] = 3; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[1] = trace.version; + words[10] = 1; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); +} + +test "exl3: v4 raw picks cover all 512 experts beyond the planner unique bound" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 47); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + var picks: [47 * 11]i32 = @splat(512); + for (0..512) |index| picks[index] = @intCast(511 - index); + append(words, seed, seed, &picks, false, 0); + const record = try (try trace.View.init(words)).round(0); + try std.testing.expectEqual(@as(usize, 512), record.uniqueCount()); + try std.testing.expectEqual(@as(u32, 512), record.validPicks()); + for (picks, record.picks()) |expert, raw| try std.testing.expectEqual(expert, @as(i32, @bitCast(raw))); +} diff --git a/zig/src/families/flashnext/exl3_dplan_trace_writer_test.zig b/zig/src/families/flashnext/exl3_dplan_trace_writer_test.zig new file mode 100644 index 0000000000..3ab0088f56 --- /dev/null +++ b/zig/src/families/flashnext/exl3_dplan_trace_writer_test.zig @@ -0,0 +1,31 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const replay = @import("exl3_dplan_replay.zig"); + +test "exl3: CUDA writer bytes replay all 1200 rounds and terminal words exactly on CPU" { + var snapshot = try trace.Snapshot.decode(std.testing.allocator, @embedFile("dplan_trace_fixture")); + defer snapshot.deinit(); + const view = try snapshot.view(); + const result = try replay.verifyTrace(std.testing.allocator, view); + try std.testing.expectEqual(@as(usize, 1200), result.rounds); + try std.testing.expectEqual(@as(usize, 1200), result.matched_transitions); + try std.testing.expectEqual(@as(u64, 0), result.delta.overflow); + try std.testing.expect(result.delta.mapped_picks > 0); + try std.testing.expect(result.delta.copies > 0); + try std.testing.expect(result.delta.evictions > 0); +} + +test "exl3: CUDA ownership checkpoints replay every segment and final transferred state" { + var snapshot = try trace.Snapshot.decode(std.testing.allocator, @embedFile("dplan_checkpoint_fixture")); + defer snapshot.deinit(); + const view = try snapshot.view(); + const result = try replay.verifyTraceWithFinal(std.testing.allocator, view, view.terminalSeed()); + try std.testing.expectEqual(@as(usize, 1024), result.rounds); + try std.testing.expectEqual(@as(usize, 1024), result.matched_transitions); + try std.testing.expectEqual(@as(usize, 4), view.checkpoints()); + try std.testing.expectEqual(@as(usize, 0), (try view.checkpoint(0)).roundIndex()); + try std.testing.expectEqual(@as(usize, 512), (try view.checkpoint(1)).roundIndex()); + try std.testing.expectEqual(@as(usize, 768), (try view.checkpoint(2)).roundIndex()); + try std.testing.expectEqual(@as(usize, 1024), (try view.checkpoint(3)).roundIndex()); + try std.testing.expect(result.delta.copies > 0); +} diff --git a/zig/src/families/flashnext/exl3_format.zig b/zig/src/families/flashnext/exl3_format.zig new file mode 120000 index 0000000000..c48a887b04 --- /dev/null +++ b/zig/src/families/flashnext/exl3_format.zig @@ -0,0 +1 @@ +../../core/exl3_format.zig \ No newline at end of file diff --git a/zig/src/families/flashnext/exl3_frequency_seed.zig b/zig/src/families/flashnext/exl3_frequency_seed.zig new file mode 100644 index 0000000000..274e8aa982 --- /dev/null +++ b/zig/src/families/flashnext/exl3_frequency_seed.zig @@ -0,0 +1,139 @@ +const std = @import("std"); + +pub fn select(counts: []const u32, widths: []const [3]i32, output: []u16) !usize { + if (counts.len != 512 or widths.len != counts.len or output.len > counts.len) return error.InvalidProfileShape; + var order: [512]u16 = undefined; + var used: usize = 0; + for (counts, widths, 0..) |count, width, expert| { + if (count == 0 or width[0] <= 0 or width[1] <= 0 or width[2] <= 0) continue; + order[used] = @intCast(expert); + used += 1; + } + const Ranking = struct { + fn hotter(histogram: []const u32, left: u16, right: u16) bool { + if (histogram[left] != histogram[right]) return histogram[left] > histogram[right]; + return left < right; + } + }; + std.mem.sort(u16, order[0..used], counts, Ranking.hotter); + const selected = @min(output.len, used); + @memcpy(output[0..selected], order[0..selected]); + return selected; +} + +pub fn warm(pool: anytype) !usize { + const policy = @import("exl3_offload.zig"); + for (pool.line_tab, pool.slots_tab) |state, expert| { + if (state != .empty or expert != policy.Cache.empty_id) return error.ProfilePoolNotEmpty; + } + if (pool.round_no != 0) return error.ProfilePoolNotEmpty; + var selected: [512]u16 = undefined; + const count = try select(&pool.counts[pool.opt.layer], pool.k2, selected[0..@min(pool.opt.lines, selected.len)]); + for (selected[0..count], 0..) |expert, held| { + const line = try policy.stageUsedAt(&pool.v, pool.opt.layer, expert, selected[0..held], pool.counts); + try pool.lines.load(line, expert, pool.src.routed[expert], pool.k2[expert], null, pool.opt.direct, pool.io); + } + try pool.lines.drain(); + for (selected[0..count]) |expert| try policy.reapAt(&pool.v, pool.opt.layer, expert, pool.now); + return count; +} + +test "exl3: frequency seed ranks local positive experts without heating observations" { + var histogram: [512]u32 = @splat(0); + var widths: [512][3]i32 = @splat(.{ 0, 0, 0 }); + histogram[2] = 19; + histogram[3] = 19; + histogram[4] = 20; + histogram[511] = std.math.maxInt(u32); + widths[2] = .{ 1, 1, 1 }; + widths[3] = .{ 1, 1, 1 }; + widths[4] = .{ 1, 0, 1 }; + var output: [2]u16 = undefined; + const before = histogram; + try std.testing.expectEqual(@as(usize, 2), try select(&histogram, &widths, &output)); + try std.testing.expectEqualSlices(u16, &.{ 2, 3 }, &output); + try std.testing.expectEqualSlices(u32, &before, &histogram); +} + +test "exl3: frequency seed leaves spare lines for unknown and unobserved experts" { + var histogram: [512]u32 = @splat(0); + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + var output: [512]u16 = undefined; + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, &output)); + histogram[511] = 1; + try std.testing.expectEqual(@as(usize, 1), try select(&histogram, &widths, &output)); + try std.testing.expectEqual(@as(u16, 511), output[0]); + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, output[0..0])); + try std.testing.expectError(error.InvalidProfileShape, select(histogram[0..511], &widths, &output)); +} + +test "exl3: frequency warm publishes residency only after copies drain and never heats counts" { + const policy = @import("exl3_offload.zig"); + const Copies = struct { + ids: [2]u16 = undefined, + used: usize = 0, + drained: bool = false, + states: []policy.Line, + fail: bool = false, + + pub fn load(copies: *@This(), line: usize, expert: usize, _: [3][]const u8, _: [3]i32, evicted: ?usize, _: bool, _: void) !void { + try std.testing.expect(evicted == null); + try std.testing.expectEqual(policy.Line.inflight, copies.states[line]); + if (copies.fail) return error.CopyFailed; + copies.ids[copies.used] = @intCast(expert); + copies.used += 1; + } + + pub fn drain(copies: *@This()) !void { + for (copies.states) |state| try std.testing.expectEqual(policy.Line.inflight, state); + copies.drained = true; + } + }; + const Fake = struct { + line_tab: []policy.Line, + slots_tab: []u16, + round_no: u32 = 0, + counts: *[48][512]u32, + opt: struct { layer: u8 = 7, lines: usize = 2, direct: bool = true } = .{}, + k2: []const [3]i32, + v: policy.Vram, + lines: Copies, + src: struct { routed: []const [3][]const u8 }, + io: void = {}, + now: u32 = 0, + }; + const allocator = std.testing.allocator; + const counts = try allocator.create([48][512]u32); + defer allocator.destroy(counts); + counts.* = std.mem.zeroes([48][512]u32); + counts[7][2] = 100; + counts[7][4] = 200; + const before = counts[7]; + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + const sources: [512][3][]const u8 = @splat(.{ &.{}, &.{}, &.{} }); + var ids: [2]u16 = undefined; + var states: [2]policy.Line = undefined; + var stamps: [2]u32 = undefined; + var layers: [2]u8 = undefined; + var fake: Fake = .{ + .line_tab = &states, + .slots_tab = &ids, + .counts = counts, + .k2 = &widths, + .v = policy.Vram.init(&ids, &states, &stamps, &layers), + .lines = .{ .states = &states }, + .src = .{ .routed = &sources }, + }; + try std.testing.expectEqual(@as(usize, 2), try warm(&fake)); + try std.testing.expect(fake.lines.drained); + try std.testing.expectEqualSlices(u16, &.{ 4, 2 }, &ids); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); + for (states) |state| try std.testing.expectEqual(policy.Line.resident, state); + try std.testing.expectError(error.ProfilePoolNotEmpty, warm(&fake)); + fake.v = policy.Vram.init(&ids, &states, &stamps, &layers); + fake.lines = .{ .states = &states, .fail = true }; + try std.testing.expectError(error.CopyFailed, warm(&fake)); + try std.testing.expect(!fake.lines.drained); + try std.testing.expectEqual(policy.Line.inflight, states[0]); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); +} diff --git a/zig/src/families/flashnext/exl3_fwd_cpu.zig b/zig/src/families/flashnext/exl3_fwd_cpu.zig new file mode 100644 index 0000000000..ba3a7bb642 --- /dev/null +++ b/zig/src/families/flashnext/exl3_fwd_cpu.zig @@ -0,0 +1,642 @@ +//! CPU f64 oracle for the whole Qwen3.8-Flash-Next EXL3 forward, assembled from the per-op oracles +//! (`exl3_hc` preF64/postF64, `exl3_gdn_cpu`, `exl3_attn_cpu`, `exl3_moe` routing, `exl3_cpu` swiglu). +//! It works layer-major: every prompt row goes through layer l before any row goes through layer l + 1, which is +//! exactly the causal order for a teacher-forced pass and lets each layer's weights be decoded once. +//! +//! `QMat` is the EXL3 matrix as the Zig decoder (`exl3_format`) defines it, W = diag(suh) H_K Wq H_N diag(svh), applied to a +//! vector without materialising W: y = ((x * suh) H_K) Wq H_N * svh, with the 128-wide Hadamard blocks done as a fast +//! Walsh-Hadamard transform. `exl3_format.dequantize` (dense 128x128 rotations) is the reference it is tested against. +//! Wq is the unpacked rotated-domain fp16 matrix, kept as f32 (exact), accumulated in f64. +//! No CUDA. + +const std = @import("std"); +const format = @import("exl3_format.zig"); +const hc = @import("exl3_hc.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); +const moe = @import("exl3_moe.zig"); +const cpu = @import("exl3_cpu.zig"); + +pub const hidden: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * hidden; +pub const topk: usize = 10; +pub const experts: usize = 512; +pub const inter: usize = 640; + +// -- parallel for ------------------------------------------------------------------------------------------ + +/// Runs `f(ctx, i)` for i in 0..n over up to `max_threads` threads (the caller included). `f` stores its own errors in `ctx`. +pub fn parallelFor(n: usize, max_threads: usize, ctx: anytype, comptime f: anytype) void { + const Ctx = @TypeOf(ctx); + const W = struct { + fn run(next: *std.atomic.Value(usize), count: usize, c: Ctx) void { + while (true) { + const i = next.fetchAdd(1, .monotonic); + if (i >= count) return; + f(c, i); + } + } + }; + var next = std.atomic.Value(usize).init(0); + const cpus = std.Thread.getCpuCount() catch 4; + const want = @min(@min(cpus, max_threads), n); + var handles: [128]std.Thread = undefined; + var spawned: usize = 0; + while (spawned + 1 < want and spawned < handles.len) : (spawned += 1) { + handles[spawned] = std.Thread.spawn(.{}, W.run, .{ &next, n, ctx }) catch break; + } + W.run(&next, n, ctx); + for (handles[0..spawned]) |h| h.join(); +} + +// -- QMat -------------------------------------------------------------------------------------------------- + +/// In-place orthonormal Walsh-Hadamard transform of each 128 block of `x` (Sylvester order, H[i][j] = (-1)^popcount(i&j) / sqrt(128)). +pub fn fwht128(x: []f64) void { + std.debug.assert(x.len % 128 == 0); + const inv = 1.0 / @sqrt(128.0); + var base: usize = 0; + while (base < x.len) : (base += 128) { + const b = x[base..][0..128]; + var h: usize = 1; + while (h < 128) : (h *= 2) { + var i: usize = 0; + while (i < 128) : (i += 2 * h) { + for (i..i + h) |j| { + const lo = b[j]; + const hi = b[j + h]; + b[j] = lo + hi; + b[j + h] = lo - hi; + } + } + } + for (b) |*v| v.* *= inv; + } +} + +pub const QMat = struct { + k: usize, + n: usize, + wq: []f32, + suh: []f32, + svh: []f32, + + /// `threads` parallelises the trellis unpack over tile rows. + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16, threads: usize) !QMat { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % 128 != 0 or n % 128 != 0 or suh.len != k or svh.len != n) return error.Length; + if (trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const wq = try gpa.alloc(f32, k * n); + errdefer gpa.free(wq); + const su = try gpa.alloc(f32, k); + errdefer gpa.free(su); + const sv = try gpa.alloc(f32, n); + errdefer gpa.free(sv); + for (su, suh) |*d, s| d.* = @floatCast(s); + for (sv, svh) |*d, s| d.* = @floatCast(s); + var job: UnpackJob = .{ .halves = halves, .trellis = trellis, .tiles_n = tiles_n, .wq = wq, .n = n }; + parallelFor(tiles_k, threads, &job, UnpackJob.row); + return .{ .k = k, .n = n, .wq = wq, .suh = su, .svh = sv }; + } + + pub fn deinit(m: *QMat, gpa: std.mem.Allocator) void { + gpa.free(m.wq); + gpa.free(m.suh); + gpa.free(m.svh); + m.* = undefined; + } + + /// y[0..n] = x[0..k] W. `tmp` is k long. + pub fn apply(m: QMat, x: []const f64, y: []f64, tmp: []f64) void { + std.debug.assert(x.len == m.k and y.len == m.n and tmp.len == m.k); + for (tmp, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(tmp); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp[i]; + if (u == 0) continue; + const row = m.wq[i * m.n ..][0..m.n]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh) |*o, s| o.* *= s; + } + + /// Same as `apply` with the output columns [lo, hi) only (the Hadamard blocks must be whole: lo and hi multiples of 128). + pub fn applyCols(m: QMat, tmp_h: []const f64, y: []f64, lo: usize, hi: usize) void { + std.debug.assert(lo % 128 == 0 and hi % 128 == 0 and y.len == hi - lo); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp_h[i]; + if (u == 0) continue; + const row = m.wq[i * m.n + lo ..][0 .. hi - lo]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh[lo..hi]) |*o, s| o.* *= s; + } + + /// The transformed input (x * suh) H_K that `applyCols` consumes. + pub fn prepare(m: QMat, x: []const f64, out: []f64) void { + for (out, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(out); + } +}; + +const UnpackJob = struct { + halves: u16, + trellis: []const i16, + tiles_n: usize, + wq: []f32, + n: usize, + fn row(j: *UnpackJob, tk: usize) void { + var placement: [256]format.Placement = undefined; + format.tilePlacement(&placement); + const w: format.Width = .{ .halves = j.halves }; + const words = w.tileWords(); + for (0..j.tiles_n) |tn| { + const tile = j.trellis[(tk * j.tiles_n + tn) * words ..][0..words]; + var states: [256]u32 = undefined; + format.tileStates(w, tile, &states); + for (0..256) |p| { + const v: f16 = format.codebookValue(.mul1, @intCast(states[p])); + j.wq[(tk * 16 + placement[p].row) * j.n + tn * 16 + placement[p].col] = @floatCast(v); + } + } + } +}; + +// -- small ops --------------------------------------------------------------------------------------------- + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f64 rows [t][hidden] -> mixed hidden of the hyper-connection (pre), per row. `inject` is [t][4] or null. +pub fn hcPreRows(gpa: std.mem.Allocator, w: hc.Weights, offset: f64, h: []const f64, rows: usize, mixed: []f64, inject: ?[]f64) !void { + const low = try gpa.alloc(f64, hc.real.lowrank); + defer gpa.free(low); + const mix = try gpa.alloc(f64, wide); + defer gpa.free(mix); + const normed = try gpa.alloc(f64, wide); + defer gpa.free(normed); + for (0..rows) |t| { + const inj: ?[]f64 = if (inject) |i| i[t * streams ..][0..streams] else null; + try hc.preF64(hc.real, w, offset, 1e-6, h[t * wide ..][0..wide], mixed[t * hidden ..][0..hidden], inj, low, mix, normed); + } +} + +pub fn hcPostRows(h: []f64, branch: []const f64, inject: []const f64, rows: usize) void { + for (0..rows) |t| { + for (0..streams) |s| { + const g = inject[t * streams + s]; + for (0..hidden) |d| h[t * wide + s * hidden + d] += branch[t * hidden + d] * g; + } + } +} + +// -- layer weights ----------------------------------------------------------------------------------------- + +pub const GdnW = struct { + qkv: QMat, + z: QMat, + out: QMat, + a_w: []const f64, // [48][2560] + b_w: []const f64, + params: gdn_cpu.Params, + + pub fn deinit(w: *GdnW, gpa: std.mem.Allocator) void { + w.qkv.deinit(gpa); + w.z.deinit(gpa); + w.out.deinit(gpa); + gpa.free(w.a_w); + gpa.free(w.b_w); + gpa.free(w.params.conv_w); + gpa.free(w.params.a_log); + gpa.free(w.params.dt_bias); + gpa.free(w.params.norm_w); + } +}; + +pub const AttnW = struct { + q: QMat, + k: QMat, + v: QMat, + o: QMat, + idx: QMat, + norms: attn_cpu.Norms, + + pub fn deinit(w: *AttnW, gpa: std.mem.Allocator) void { + w.q.deinit(gpa); + w.k.deinit(gpa); + w.v.deinit(gpa); + w.o.deinit(gpa); + w.idx.deinit(gpa); + gpa.free(w.norms.q); + gpa.free(w.norms.k); + gpa.free(w.norms.iq); + gpa.free(w.norms.ik); + } +}; + +/// Where the 512 routed experts of one layer come from: decode expert `e` (gate, up, down) into `out`. +pub const ExpertSource = struct { + ctx: *anyopaque, + load: *const fn (ctx: *anyopaque, e: usize, out: *[3]QMat) anyerror!void, +}; + +pub const MoeW = struct { + gate16: []const f16, // [512][2560] + sgate16: []const f16, // [2560] + shared: [3]QMat, + source: ExpertSource, + + pub fn deinit(w: *MoeW, gpa: std.mem.Allocator) void { + for (&w.shared) |*m| m.deinit(gpa); + } +}; + +pub const Mixer = union(enum) { gdn: *GdnW, attn: *AttnW }; + +pub const LayerW = struct { + attn_hc: hc.Weights, + mlp_hc: hc.Weights, + mixer: Mixer, + moe: *MoeW, +}; + +pub const LayerState = union(enum) { gdn: gdn_cpu.State, attn: attn_cpu.Cache }; + +/// What a layer pass records for comparison with the device. +pub const Trace = struct { + attn_branch: []f64, // [rows][2560] + mlp_branch: []f64, + picks: []u16, // [rows][10] + h_mid: []f64, // [rows][10240] streams after the attention block +}; + +// -- mixers ------------------------------------------------------------------------------------------------ + +fn roundF16Slice(dst: []f16, src: []const f64) void { + for (dst, src) |*d, s| d.* = @floatCast(s); +} + +const GdnProj = struct { + w: *const GdnW, + x: []const f64, + qkv: []f64, + z: []f64, + a: []f64, + b: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn row(j: *GdnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const xr = j.x[t * hidden ..][0..hidden]; + const c = gdn_cpu.qwen.channels(); + const zw = gdn_cpu.qwen.zWidth(); + j.w.qkv.apply(xr, j.qkv[t * c ..][0..c], tmp); + j.w.z.apply(xr, j.z[t * zw ..][0..zw], tmp); + gdn_cpu.rowsDot(j.w.a_w, xr, j.a[t * 48 ..][0..48]); + gdn_cpu.rowsDot(j.w.b_w, xr, j.b[t * 48 ..][0..48]); + } +}; + +const OutJob = struct { + m: *const QMat, + inp: []const f64, + in_w: usize, + out: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + fn row(j: *OutJob, t: usize) void { + const tmp = j.gpa.alloc(f64, j.m.k) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + j.m.apply(j.inp[t * j.in_w ..][0..j.in_w], j.out[t * j.m.n ..][0..j.m.n], tmp); + } +}; + +pub fn gdnBranch(gpa: std.mem.Allocator, w: *const GdnW, st: gdn_cpu.State, x: []const f64, rows: usize, out: []f64) !void { + const d = gdn_cpu.qwen; + const c = d.channels(); + const zw = d.zWidth(); + const qkv = try gpa.alloc(f64, rows * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, rows * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, rows * 48); + defer gpa.free(a); + const b = try gpa.alloc(f64, rows * 48); + defer gpa.free(b); + var pj: GdnProj = .{ .w = w, .x = x, .qkv = qkv, .z = z, .a = a, .b = b, .gpa = gpa }; + parallelFor(rows, 48, &pj, GdnProj.row); + if (pj.err) |e| return e; + const gated = try gpa.alloc(f64, rows * zw); + defer gpa.free(gated); + try gdn_cpu.tokens(gpa, d, w.params, st, rows, qkv, z, a, b, gated); + var oj: OutJob = .{ .m = &w.out, .inp = gated, .in_w = zw, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +const AttnProj = struct { + w: *const AttnW, + x: []const f64, + qp: []f16, + kp: []f16, + vp: []f16, + ip: []f16, + gpa: std.mem.Allocator, + err: ?anyerror = null, + const q_width = attn_cpu.q_heads * 2 * attn_cpu.head_dim; + const kv_width = attn_cpu.kv_heads * attn_cpu.head_dim; + + fn row(j: *AttnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const big = j.gpa.alloc(f64, q_width) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(big); + const xr = j.x[t * hidden ..][0..hidden]; + j.w.q.apply(xr, big, tmp); + roundF16Slice(j.qp[t * q_width ..][0..q_width], big); + j.w.k.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.kp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.v.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.vp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.idx.apply(xr, big[0..attn_cpu.index_width], tmp); + roundF16Slice(j.ip[t * attn_cpu.index_width ..][0..attn_cpu.index_width], big[0..attn_cpu.index_width]); + } +}; + +/// Attention rows pos0 .. pos0 + rows. The projections are rounded to fp16 where the EXL3 linears emit them. +pub fn attnBranch(gpa: std.mem.Allocator, w: *const AttnW, cache: *attn_cpu.Cache, pos0: usize, x: []const f64, rows: usize, out: []f64) !void { + const q_width = AttnProj.q_width; + const kv_width = AttnProj.kv_width; + const iw = attn_cpu.index_width; + const qp = try gpa.alloc(f16, rows * q_width); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, rows * iw); + defer gpa.free(ip); + var pj: AttnProj = .{ .w = w, .x = x, .qp = qp, .kp = kp, .vp = vp, .ip = ip, .gpa = gpa }; + parallelFor(rows, 48, &pj, AttnProj.row); + if (pj.err) |e| return e; + + const attn_width = attn_cpu.q_heads * attn_cpu.head_dim; + const o_in = try gpa.alloc(f64, rows * attn_width); + defer gpa.free(o_in); + const scores = try gpa.alloc(f64, cache.tokens / attn_cpu.ratio + 2); + defer gpa.free(scores); + var oq: [attn_width]f64 = undefined; + var oiq: [attn_cpu.index_heads * attn_cpu.index_dim]f64 = undefined; + var info = attn_cpu.RowInfo{}; + for (0..rows) |t| { + try attn_cpu.row(cache, pos0 + t, .{ + .qp = qp[t * q_width ..][0..q_width], + .kp = kp[t * kv_width ..][0..kv_width], + .vp = vp[t * kv_width ..][0..kv_width], + .ip = ip[t * iw ..][0..iw], + }, &w.norms, scores, &oq, &oiq, o_in[t * attn_width ..][0..attn_width], &info); + } + var oj: OutJob = .{ .m = &w.o, .inp = o_in, .in_w = attn_width, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +// -- MoE --------------------------------------------------------------------------------------------------- + +const Loaded = struct { + id: usize, + mats: [3]QMat, +}; + +const LoadJob = struct { + src: ExpertSource, + ids: []const usize, + slots: []Loaded, + err: ?anyerror = null, + fn one(j: *LoadJob, i: usize) void { + j.slots[i].id = j.ids[i]; + j.src.load(j.src.ctx, j.ids[i], &j.slots[i].mats) catch |e| { + j.err = e; + }; + } +}; + +const MoeRow = struct { + w: *const MoeW, + x: []const f64, + out: []f64, + picks: []const u16, // [rows][10] + weights: []const f64, // [rows][10] + shared_w: []const f64, + loaded: []const Loaded, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn find(j: *const MoeRow, id: usize) *const Loaded { + for (j.loaded) |*l| if (l.id == id) return l; + unreachable; + } + + fn row(j: *MoeRow, t: usize) void { + const buf = j.gpa.alloc(f64, hidden + 4 * inter + 3 * hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(buf); + const tmp = buf[0..hidden]; + const g = buf[hidden..][0..inter]; + const u = buf[hidden + inter ..][0..inter]; + const glu = buf[hidden + 2 * inter ..][0..inter]; + const y = buf[hidden + 3 * inter ..][0..hidden]; + const acc = j.out[t * hidden ..][0..hidden]; + @memset(acc, 0); + const xr = j.x[t * hidden ..][0..hidden]; + for (0..topk) |k| { + const l = j.find(j.picks[t * topk + k]); + expertApply(&l.mats, xr, tmp, g, u, glu, y); + const wk = j.weights[t * topk + k]; + for (acc, y) |*o, v| o.* += wk * v; + } + expertApply(&j.w.shared, xr, tmp, g, u, glu, y); + const sw = j.shared_w[t]; + for (acc, y) |*o, v| o.* += sw * v; + } +}; + +fn expertApply(m: *const [3]QMat, x: []const f64, tmp: []f64, g: []f64, u: []f64, glu: []f64, y: []f64) void { + m[0].apply(x, g, tmp); + m[1].apply(x, u, tmp); + for (glu, g, u) |*o, gv, uv| o.* = silu(gv) * uv; + m[2].apply(glu, y, tmp[0..inter]); +} + +/// Routed + shared experts for `rows` rows from the f64 hidden `x` (the f64 router, not the device's). Records the picks. +pub fn moeBranch(gpa: std.mem.Allocator, w: *const MoeW, x: []const f64, rows: usize, out: []f64, picks_out: []u16) !void { + const picks = picks_out; + const weights = try gpa.alloc(f64, rows * topk); + defer gpa.free(weights); + const shared_w = try gpa.alloc(f64, rows); + defer gpa.free(shared_w); + var logits: [experts + 1]f64 = undefined; + var ids: std.ArrayList(usize) = .empty; + defer ids.deinit(gpa); + for (0..rows) |t| { + try moe.logitsF16(w.gate16, experts, w.sgate16, x[t * hidden ..][0..hidden], &logits); + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, picks[t * topk ..][0..topk], weights[t * topk ..][0..topk], &sw); + shared_w[t] = sw; + for (picks[t * topk ..][0..topk]) |p| { + var seen = false; + for (ids.items) |i| seen = seen or i == p; + if (!seen) try ids.append(gpa, p); + } + } + const loaded = try gpa.alloc(Loaded, ids.items.len); + defer gpa.free(loaded); + var lj: LoadJob = .{ .src = w.source, .ids = ids.items, .slots = loaded }; + parallelFor(ids.items.len, 48, &lj, LoadJob.one); + defer if (lj.err == null) for (loaded) |*l| for (&l.mats) |*m| m.deinit(gpa); + if (lj.err) |e| return e; + var mj: MoeRow = .{ .w = w, .x = x, .out = out, .picks = picks, .weights = weights, .shared_w = shared_w, .loaded = loaded, .gpa = gpa }; + parallelFor(rows, 48, &mj, MoeRow.row); + if (mj.err) |e| return e; +} + +// -- one layer --------------------------------------------------------------------------------------------- + +/// `h` is [rows][10240] f64 streams, updated in place. `pos0` is the absolute position of row 0. +pub fn layerRows(gpa: std.mem.Allocator, w: LayerW, st: *LayerState, pos0: usize, offset: f64, h: []f64, rows: usize, trace: ?Trace) !void { + const mixed = try gpa.alloc(f64, rows * hidden); + defer gpa.free(mixed); + const inject = try gpa.alloc(f64, rows * streams); + defer gpa.free(inject); + const branch = try gpa.alloc(f64, rows * hidden); + defer gpa.free(branch); + + try hcPreRows(gpa, w.attn_hc, offset, h, rows, mixed, inject); + switch (w.mixer) { + .gdn => |g| try gdnBranch(gpa, g, st.gdn, mixed, rows, branch), + .attn => |a| try attnBranch(gpa, a, &st.attn, pos0, mixed, rows, branch), + } + if (trace) |tr| @memcpy(tr.attn_branch[0 .. rows * hidden], branch); + hcPostRows(h, branch, inject, rows); + if (trace) |tr| @memcpy(tr.h_mid[0 .. rows * wide], h[0 .. rows * wide]); + + try hcPreRows(gpa, w.mlp_hc, offset, h, rows, mixed, inject); + const picks = try gpa.alloc(u16, rows * topk); + defer gpa.free(picks); + try moeBranch(gpa, w.moe, mixed, rows, branch, picks); + if (trace) |tr| { + @memcpy(tr.mlp_branch[0 .. rows * hidden], branch); + @memcpy(tr.picks[0 .. rows * topk], picks); + } + hcPostRows(h, branch, inject, rows); +} + +// -- tests ------------------------------------------------------------------------------------------------- + +test "fwht128 is orthonormal and equals the dense Sylvester matrix" { + var rng = std.Random.DefaultPrng.init(3); + var x: [256]f64 = undefined; + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var y = x; + fwht128(&y); + for (0..2) |blk| { + for (0..128) |i| { + var s: f64 = 0; + for (0..128) |j| { + const sign: f64 = if (@popCount(i & j) & 1 == 1) -1 else 1; + s += sign * x[blk * 128 + j]; + } + try std.testing.expectApproxEqAbs(s / @sqrt(128.0), y[blk * 128 + i], 1e-12); + } + } + fwht128(&y); + for (x, y) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); +} + +test "QMat.apply equals the dense decoder's dequantize then matvec" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(17); + const r = rng.random(); + for ([_]u16{ 4, 6, 10 }) |halves| { + const tiles_k: usize = 8; + const tiles_n: usize = 16; + const k = tiles_k * 16; + const n = tiles_n * 16; + const words = @as(usize, halves) * 8; + const trellis = try gpa.alloc(i16, tiles_k * tiles_n * words); + defer gpa.free(trellis); + for (trellis) |*v| v.* = @bitCast(r.int(u16)); + const suh = try gpa.alloc(f16, k); + defer gpa.free(suh); + const svh = try gpa.alloc(f16, n); + defer gpa.free(svh); + for (suh) |*v| v.* = if (r.boolean()) 1.0 else -1.0; + for (svh) |*v| v.* = @floatCast(0.5 + r.float(f32)); + const dense = try gpa.alloc(f64, k * n); + defer gpa.free(dense); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, dense, scratch); + var q = try QMat.init(gpa, halves, trellis, tiles_k, tiles_n, suh, svh, 4); + defer q.deinit(gpa); + var x: [k]f64 = undefined; + for (&x) |*v| v.* = r.floatNorm(f64); + var want: [n]f64 = undefined; + for (0..n) |j| { + var s: f64 = 0; + for (0..k) |i| s += x[i] * dense[i * n + j]; + want[j] = s; + } + var got: [n]f64 = undefined; + var tmp: [k]f64 = undefined; + q.apply(&x, &got, &tmp); + for (want, got) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-9 * (1 + @abs(a))); + // column slices equal the whole + var prep: [k]f64 = undefined; + q.prepare(&x, &prep); + var part: [128]f64 = undefined; + q.applyCols(&prep, &part, 128, 256); + for (part, got[128..256]) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); + } +} + +test "parallelFor visits every index once" { + const Ctx = struct { + hits: []std.atomic.Value(u8), + fn f(c: *@This(), i: usize) void { + _ = c.hits[i].fetchAdd(1, .monotonic); + } + }; + var hits: [1000]std.atomic.Value(u8) = undefined; + for (&hits) |*h| h.* = std.atomic.Value(u8).init(0); + var c: Ctx = .{ .hits = &hits }; + parallelFor(hits.len, 8, &c, Ctx.f); + for (&hits) |*h| try std.testing.expectEqual(@as(u8, 1), h.load(.monotonic)); +} diff --git a/zig/src/families/flashnext/exl3_gdn_cpu.zig b/zig/src/families/flashnext/exl3_gdn_cpu.zig new file mode 100644 index 0000000000..35a23a5263 --- /dev/null +++ b/zig/src/families/flashnext/exl3_gdn_cpu.zig @@ -0,0 +1,417 @@ +//! CPU oracle for one Gated DeltaNet layer of Qwen3.8-Flash-Next (qwen4_exp_text), everything between the two EXL3 +//! projections: causal depthwise conv with state, SiLU, q/k L2 norm, the gated delta-rule recurrence, and the +//! gated RMS norm. All maths in f64; the device keeps the recurrent state in fp32. +//! +//! Verified against the model's own code, transformers `models/qwen4_exp/modeling_qwen4_exp.py` (5.8.0.dev0): +//! Qwen4ExpTextGatedDeltaNet.forward conv over [q|k|v] channels, silu, split [key_dim, key_dim, value_dim], +//! beta = sigmoid(b), g = -exp(A_log) * softplus(a + dt_bias), +//! q and k repeat_interleave by nv/nk (value head hv uses key head hv / (nv/nk)) +//! torch_recurrent_gated_delta_rule q,k = l2norm(eps 1e-6: x * rsqrt(sum x^2 + 1e-6)), q *= dk^-0.5, +//! S = S * exp(g); kv = S^T k; d = (v - kv) * beta; S += k d^T; o = S^T q +//! Qwen4ExpTextRMSNormGated (activation = config.output_gate_type = "sigmoid", NOT silu): +//! w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(z), w stored as is (not 1 + w) +//! and against the affine path's Metal kernel (kernels/metal/decode/fn_gdn.metal), which implements the same +//! arithmetic with the state laid out [value head][value row][key column]. +//! The prefill (chunked) rule in the model is mathematically this recurrence, so one token loop is the oracle for both. + +const std = @import("std"); + +pub const Dims = struct { + nk: usize, + nv: usize, + dk: usize, + dv: usize, + taps: usize, + + /// conv channels: q (nk*dk), k (nk*dk), v (nv*dv) + pub fn channels(d: Dims) usize { + return 2 * d.nk * d.dk + d.nv * d.dv; + } + pub fn zWidth(d: Dims) usize { + return d.nv * d.dv; + } + /// conv state rows kept per layer + pub fn keep(d: Dims) usize { + return d.taps - 1; + } + pub fn recElems(d: Dims) usize { + return d.nv * d.dv * d.dk; + } + pub fn convElems(d: Dims) usize { + return d.keep() * d.channels(); + } +}; + +/// Flash Next: 16 key heads, 48 value heads, 128 x 128, kernel 4. +pub const qwen: Dims = .{ .nk = 16, .nv = 48, .dk = 128, .dv = 128, .taps = 4 }; + +/// Per-sequence, per-layer state. `conv` is time-major, oldest row first: the last `taps - 1` raw projection rows. +/// `rec` is [value head][value row][key column] (the device layout). +pub const State = struct { + conv: []f64, + rec: []f64, + + pub fn init(gpa: std.mem.Allocator, d: Dims) !State { + const conv = try gpa.alloc(f64, d.convElems()); + errdefer gpa.free(conv); + const rec = try gpa.alloc(f64, d.recElems()); + @memset(conv, 0); + @memset(rec, 0); + return .{ .conv = conv, .rec = rec }; + } + pub fn deinit(s: *State, gpa: std.mem.Allocator) void { + gpa.free(s.conv); + gpa.free(s.rec); + } + pub fn reset(s: *State) void { + @memset(s.conv, 0); + @memset(s.rec, 0); + } +}; + +/// The layer's non-EXL3 parameters, as f64. +pub const Params = struct { + conv_w: []const f64, // [channels][taps] + a_log: []const f64, // [nv] + dt_bias: []const f64, // [nv] + norm_w: []const f64, // [dv] + eps: f64 = 1e-6, +}; + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +/// torch.nn.functional.softplus with threshold 20. +pub fn softplus(x: f64) f64 { + return if (x > 20.0) x else std.math.log1p(@exp(x)); +} + +/// y[j] = sum_i w[j * k + i] * x[i], for a dense [rows][k] weight (in_proj_a / in_proj_b). +pub fn rowsDot(w: []const f64, x: []const f64, y: []f64) void { + const k = x.len; + for (y, 0..) |*o, j| { + var s: f64 = 0; + for (x, w[j * k ..][0..k]) |xi, wi| s += xi * wi; + o.* = s; + } +} + +/// One token of the causal depthwise conv: out = silu(sum_tap w[c][tap] * x[t - (taps-1) + tap]), x before the +/// sequence is the stored state (zeros at a fresh sequence). Shifts the state in place. +pub fn convStep(d: Dims, p: Params, conv: []f64, x: []const f64, out: []f64) void { + const c = d.channels(); + const kp = d.keep(); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..d.taps) |tap| { + const v = if (tap < kp) conv[tap * c + ch] else x[ch]; + acc += p.conv_w[ch * d.taps + tap] * v; + } + out[ch] = silu(acc); + } + for (0..kp) |r| { + for (0..c) |ch| conv[r * c + ch] = if (r + 1 < kp) conv[(r + 1) * c + ch] else x[ch]; + } +} + +/// Decay exp(g) and beta for each value head from the raw a and b projections. +pub fn gates(d: Dims, p: Params, a: []const f64, b: []const f64, decay: []f64, beta: []f64) void { + for (0..d.nv) |h| { + const g = -@exp(p.a_log[h]) * softplus(a[h] + p.dt_bias[h]); + decay[h] = @exp(g); + beta[h] = sigmoid(b[h]); + } +} + +/// One token of the delta rule on all value heads. `act` is the conv output [q | k | v]; y is [nv][dv]. +pub fn deltaStep(d: Dims, rec: []f64, act: []const f64, decay: []const f64, beta: []const f64, y: []f64) void { + const group = d.nv / d.nk; + const kbase = d.nk * d.dk; + const vbase = 2 * kbase; + for (0..d.nv) |hv| { + const hk = hv / group; + const q = act[hk * d.dk ..][0..d.dk]; + const k = act[kbase + hk * d.dk ..][0..d.dk]; + var qq: f64 = 0; + var kk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(@as(f64, @floatFromInt(d.dk))); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..d.dv) |v| { + const row = rec[(hv * d.dv + v) * d.dk ..][0..d.dk]; + var kv: f64 = 0; + for (row, k) |*s, ki| { + s.* *= decay[hv]; + kv += s.* * ki * ik; + } + const delta = (act[vbase + hv * d.dv + v] - kv) * beta[hv]; + var o: f64 = 0; + for (row, k, q) |*s, ki, qi| { + s.* += ki * ik * delta; + o += s.* * qi * iq; + } + y[hv * d.dv + v] = o; + } + } +} + +/// Gated RMS norm per value head: w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(z). +pub fn gatedNorm(d: Dims, p: Params, y: []const f64, z: []const f64, out: []f64) void { + for (0..d.nv) |hv| { + const row = y[hv * d.dv ..][0..d.dv]; + var ss: f64 = 0; + for (row) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d.dv)) + p.eps); + for (row, 0..) |v, i| out[hv * d.dv + i] = p.norm_w[i] * (v * inv) * sigmoid(z[hv * d.dv + i]); + } +} + +/// `n` tokens, in order, from the projection outputs: qkv [n][channels], z [n][nv*dv], a and b [n][nv]. +/// out is [n][nv*dv], the out_proj input. The state is carried in `st`. +pub fn tokens(gpa: std.mem.Allocator, d: Dims, p: Params, st: State, n: usize, qkv: []const f64, z: []const f64, a: []const f64, b: []const f64, out: []f64) !void { + const c = d.channels(); + const zw = d.zWidth(); + const act = try gpa.alloc(f64, c); + defer gpa.free(act); + const y = try gpa.alloc(f64, zw); + defer gpa.free(y); + const decay = try gpa.alloc(f64, d.nv); + defer gpa.free(decay); + const beta = try gpa.alloc(f64, d.nv); + defer gpa.free(beta); + for (0..n) |t| { + convStep(d, p, st.conv, qkv[t * c ..][0..c], act); + gates(d, p, a[t * d.nv ..][0..d.nv], b[t * d.nv ..][0..d.nv], decay, beta); + deltaStep(d, st.rec, act, decay, beta, y); + gatedNorm(d, p, y, z[t * zw ..][0..zw], out[t * zw ..][0..zw]); + } +} + +/// Byte sizes of the device state table: conv rows are fp16 (they are raw EXL3 projection outputs), the recurrent +/// state is fp32. One entry per (slot, layer); a sequence's layers are contiguous. +pub const Table = struct { + dims: Dims, + layers: usize, + slots: usize, + + pub fn convBytesPer(t: Table) usize { + return t.dims.convElems() * 2; + } + pub fn recBytesPer(t: Table) usize { + return t.dims.recElems() * 4; + } + pub fn convBytes(t: Table) usize { + return t.convBytesPer() * t.layers * t.slots; + } + pub fn recBytes(t: Table) usize { + return t.recBytesPer() * t.layers * t.slots; + } + pub fn convOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.convBytesPer(); + } + pub fn recOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.recBytesPer(); + } +}; + +// ---- tests (small dims; the arithmetic does not depend on the Flash Next sizes) ---- + +const small: Dims = .{ .nk = 2, .nv = 4, .dk = 8, .dv = 8, .taps = 4 }; + +fn fill(rng: *std.Random.DefaultPrng, xs: []f64, scale: f64) void { + for (xs) |*x| x.* = (rng.random().float(f64) - 0.5) * 2.0 * scale; +} + +const TestParams = struct { + conv_w: []f64, + a_log: []f64, + dt_bias: []f64, + norm_w: []f64, + + fn init(gpa: std.mem.Allocator, d: Dims, rng: *std.Random.DefaultPrng) !TestParams { + const t: TestParams = .{ + .conv_w = try gpa.alloc(f64, d.channels() * d.taps), + .a_log = try gpa.alloc(f64, d.nv), + .dt_bias = try gpa.alloc(f64, d.nv), + .norm_w = try gpa.alloc(f64, d.dv), + }; + fill(rng, t.conv_w, 0.7); + fill(rng, t.a_log, 1.0); + fill(rng, t.dt_bias, 1.0); + for (t.norm_w) |*w| w.* = 0.5 + rng.random().float(f64); + return t; + } + fn deinit(t: TestParams, gpa: std.mem.Allocator) void { + gpa.free(t.conv_w); + gpa.free(t.a_log); + gpa.free(t.dt_bias); + gpa.free(t.norm_w); + } + fn view(t: TestParams) Params { + return .{ .conv_w = t.conv_w, .a_log = t.a_log, .dt_bias = t.dt_bias, .norm_w = t.norm_w }; + } +}; + +test "conv is a causal depthwise filter over the stored rows" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(1); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const c = small.channels(); + const n = 7; + const xs = try gpa.alloc(f64, n * c); + defer gpa.free(xs); + fill(&rng, xs, 1.0); + var st = try State.init(gpa, small); + defer st.deinit(gpa); + const out = try gpa.alloc(f64, c); + defer gpa.free(out); + for (0..n) |t| { + convStep(small, tp.view(), st.conv, xs[t * c ..][0..c], out); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..small.taps) |tap| { + // input index t - (taps - 1) + tap; before the sequence it is zero + const at: isize = @as(isize, @intCast(t)) - @as(isize, @intCast(small.taps - 1)) + @as(isize, @intCast(tap)); + if (at >= 0) acc += tp.conv_w[ch * small.taps + tap] * xs[@as(usize, @intCast(at)) * c + ch]; + } + try std.testing.expectApproxEqAbs(silu(acc), out[ch], 1e-12); + } + } + // the state ends as the last three raw rows + for (0..3) |r| for (0..c) |ch| try std.testing.expectEqual(xs[(n - 3 + r) * c + ch], st.conv[r * c + ch]); +} + +test "one scalar head follows S = g S + k (v - g S k) beta, o = S q" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 1, .taps = 4 }; + var rec = [_]f64{0.5}; + const act = [_]f64{ 2.0, 3.0, 0.75 }; // q, k, v + const decay = [_]f64{0.9}; + const beta = [_]f64{0.4}; + var y = [_]f64{0}; + deltaStep(d, &rec, &act, &decay, &beta, &y); + const ik = 1.0 / @sqrt(9.0 + 1e-6); + const iq = 1.0 / @sqrt(4.0 + 1e-6); + const k = 3.0 * ik; + const s0 = 0.5 * 0.9; + const delta = (0.75 - s0 * k) * 0.4; + const s1 = s0 + k * delta; + try std.testing.expectApproxEqAbs(s1, rec[0], 1e-14); + try std.testing.expectApproxEqAbs(s1 * 2.0 * iq, y[0], 1e-14); +} + +test "with beta 1 and no decay the state returns v for the key just written" { + const d: Dims = .{ .nk = 1, .nv = 2, .dk = 4, .dv = 3, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(2); + var rec: [2 * 3 * 4]f64 = undefined; + fill(&rng, &rec, 1.0); + var act: [2 * 4 + 2 * 3]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 1, 1 }; + const beta = [_]f64{ 1, 1 }; + var y: [6]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + var kk: f64 = 0; + for (act[4..8]) |k| kk += k * k; + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |hv| for (0..3) |v| { + var read: f64 = 0; + for (0..4) |i| read += rec[(hv * 3 + v) * 4 + i] * act[4 + i] * ik; + // S k = v * (|k|^2 / (|k|^2 + 1e-6)), the l2norm eps leaves a tiny shortfall + try std.testing.expectApproxEqAbs(act[8 + hv * 3 + v], read, 1e-5); + }; +} + +test "value heads share their key head in groups (repeat_interleave)" { + const d: Dims = .{ .nk = 2, .nv = 4, .dk = 4, .dv = 2, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(3); + var act: [2 * 2 * 4 + 4 * 2]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 0.8, 0.8, 0.8, 0.8 }; + const beta = [_]f64{ 0.5, 0.5, 0.5, 0.5 }; + var rec: [4 * 2 * 4]f64 = @splat(0); + var y: [8]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + // zero state: kv = 0, so S = k (v beta)^T and o = (k . q) v beta, with q, k of key head hv / 2 + for (0..4) |hv| { + const hk = hv / 2; + const q = act[hk * 4 ..][0..4]; + const k = act[8 + hk * 4 ..][0..4]; + var qq: f64 = 0; + var kk: f64 = 0; + var qk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + qk += a * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(4.0); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |v| try std.testing.expectApproxEqAbs(qk * iq * ik * act[16 + hv * 2 + v] * 0.5, y[hv * 2 + v], 1e-12); + } +} + +test "gated norm is w * rms-normalised y * sigmoid(z)" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 4, .taps = 4 }; + const p: Params = .{ .conv_w = &.{}, .a_log = &.{}, .dt_bias = &.{}, .norm_w = &.{ 1, 2, 3, 4 } }; + const y = [_]f64{ 1, -2, 3, -4 }; + const z = [_]f64{ 0, 1, -1, 2 }; + var out: [4]f64 = undefined; + gatedNorm(d, p, &y, &z, &out); + const inv = 1.0 / @sqrt((1.0 + 4 + 9 + 16) / 4.0 + 1e-6); + for (0..4) |i| try std.testing.expectApproxEqAbs(p.norm_w[i] * y[i] * inv * sigmoid(z[i]), out[i], 1e-14); +} + +test "processing tokens in two calls equals one call, state included" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(4); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const n = 9; + const c = small.channels(); + const zw = small.zWidth(); + const qkv = try gpa.alloc(f64, n * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, n * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, n * small.nv); + defer gpa.free(a); + const b = try gpa.alloc(f64, n * small.nv); + defer gpa.free(b); + fill(&rng, qkv, 1.0); + fill(&rng, z, 1.0); + fill(&rng, a, 1.0); + fill(&rng, b, 1.0); + var whole = try State.init(gpa, small); + defer whole.deinit(gpa); + var split = try State.init(gpa, small); + defer split.deinit(gpa); + const o1 = try gpa.alloc(f64, n * zw); + defer gpa.free(o1); + const o2 = try gpa.alloc(f64, n * zw); + defer gpa.free(o2); + try tokens(gpa, small, tp.view(), whole, n, qkv, z, a, b, o1); + const cut = 5; + try tokens(gpa, small, tp.view(), split, cut, qkv, z, a, b, o2); + try tokens(gpa, small, tp.view(), split, n - cut, qkv[cut * c ..], z[cut * zw ..], a[cut * small.nv ..], b[cut * small.nv ..], o2[cut * zw ..]); + try std.testing.expectEqualSlices(f64, o1, o2); + try std.testing.expectEqualSlices(f64, whole.rec, split.rec); + try std.testing.expectEqualSlices(f64, whole.conv, split.conv); +} + +test "state table addresses a (slot, layer) pair and Flash Next sizes" { + const t: Table = .{ .dims = qwen, .layers = 36, .slots = 2 }; + try std.testing.expectEqual(@as(usize, 10240), qwen.channels()); + try std.testing.expectEqual(@as(usize, 3 * 10240 * 2), t.convBytesPer()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), t.recBytesPer()); + try std.testing.expectEqual(t.recBytesPer() * 37, t.recOffset(1, 1)); + try std.testing.expectEqual(t.convBytesPer() * 36 * 2, t.convBytes()); +} diff --git a/zig/src/families/flashnext/exl3_graph.zig b/zig/src/families/flashnext/exl3_graph.zig new file mode 100644 index 0000000000..7bfd541eb5 --- /dev/null +++ b/zig/src/families/flashnext/exl3_graph.zig @@ -0,0 +1,436 @@ +//! Host-side description of what a CUDA graph may bake in versus what replay overwrites (ported from +//! exl3-flashnext-window-wip, retargeted to upstream's `Experts` pointer tables). +//! No CUDA calls. A replay binding is overwritten on each replay; a fixed binding stays +//! for the life of that executable graph. One capture per rank, layer, and window. +//! +//! On upstream's `Experts` the graph's kernel arguments are the ADDRESSES of the per-expert gate/up/down trellis +//! pointer tables and k2 tables, the picks and the weights. The offload pool rewrites the tables' CONTENTS in +//! device memory (memsets on its copy stream) when a line changes hands, so `expert_tables` and `picks` are replay +//! inputs whose addresses stay put, and the pool's line memory (`pool_lines`) is fixed. A round with a copy, or a line +//! whose copy has not landed, runs its staging before the replay (`roundReady`); the graph is never recaptured for a +//! swap. + +const std = @import("std"); +const offload = @import("exl3_offload.zig"); +const round_ = @import("exl3_round.zig"); + +/// The card and the device bytes kept free (were exl3_cache.zig's): a 16 GiB RTX 5060 Ti, 256 MiB reserve. +pub const card_bytes: u64 = 16 << 30; +pub const device_reserve: u64 = 256 << 20; + +pub const Binding = struct { + name: []const u8, + address: u64, + /// True: replay input. False: fixed for the life of the executable graph. + replay: bool, +}; + +pub const Capture = struct { + window: u32, + bindings: []const Binding, +}; + +/// One executable graph. Rank, layer, and window are the whole key. +pub const Key = struct { + rank: u8, + layer: u8, + window: u32, +}; + +const replay_names = [_][]const u8{ + "expert_ids", + "expert_tables", + "picks", + "ple_rows", + "kv_index", +}; + +const fixed_names = [_][]const u8{ + "dense_weights", + "shared_weights", + "pool_lines", + "expert_scales", +}; + +fn listed(name: []const u8, table: []const []const u8) bool { + for (table) |entry| if (std.mem.eql(u8, name, entry)) return true; + return false; +} + +/// Expert ids (pool line indices, the old engine), the expert pointer tables, picks, PLE rows, and the KV index must be +/// replay inputs; pool lines and the resident suh/svh stacks are fixed. +/// Slot pointers and the hit mask are not published. +pub fn isReplayName(name: []const u8) bool { + return listed(name, &replay_names); +} + +fn isFixedName(name: []const u8) bool { + return listed(name, &fixed_names); +} + +/// Known names must carry the flag this module requires. Unknown names keep the caller's flag. +pub fn valid(c: Capture) bool { + for (c.bindings) |b| { + if (isReplayName(b.name)) { + if (!b.replay) return false; + } else if (isFixedName(b.name)) { + if (b.replay) return false; + } + } + return true; +} + +/// Same window, same binding order, same replay flags, and the same address on every fixed binding. +/// Replay addresses may differ. +pub fn sameTopology(a: Capture, b: Capture) bool { + if (a.window != b.window or a.bindings.len != b.bindings.len) return false; + for (a.bindings, b.bindings) |ba, bb| { + if (!std.mem.eql(u8, ba.name, bb.name) or ba.replay != bb.replay) return false; + if (!ba.replay and ba.address != bb.address) return false; + } + return true; +} + +/// `sameTopology`, and every binding address matches. Eager and the graph stand-in share ids and weights. +pub fn argumentsEqual(a: Capture, b: Capture) bool { + if (!sameTopology(a, b)) return false; + for (a.bindings, b.bindings) |ba, bb| { + if (ba.address != bb.address) return false; + } + return true; +} + +fn sameKey(a: Key, b: Key) bool { + return a.rank == b.rank and a.layer == b.layer and a.window == b.window; +} + +/// Cold start stays eager through the first reap wave. That is the only eager window. +/// Ready only when `startup` is false, no decision is `.copy`, and no line is `.inflight`. +pub fn captureReady(startup: bool, decisions: []const offload.Decision, lines: []const offload.Line) bool { + if (startup) return false; + for (decisions) |d| if (d == .copy) return false; + for (lines) |line| if (line == .inflight) return false; + return true; +} + +/// Upstream `Experts` + offload pool: a captured experts graph replays as is when the round copies nothing. +/// Otherwise the pool stages the copies (and rewrites the tables) first; the replay still needs no recapture. +pub fn roundReady(startup: bool, r: *const round_.Round) bool { + if (startup) return false; + return r.allHits() and !r.masked(); +} + +pub const Promote = enum { + /// Rewrite the expert_ids line-index buffer and replay. + rewrite_ids, + /// This key has no capture yet. + recapture, +}; + +/// One capture per key. A stored key rewrites expert_ids and does not recapture. +pub fn promote(keys: []const Key, key: Key) Promote { + for (keys) |have| if (sameKey(have, key)) return .rewrite_ids; + return .recapture; +} + +/// One routed line. `count` is counts[layer][id]. +pub const LineUse = struct { + layer: u8, + id: u16, + count: u32, +}; + +fn colder(a: LineUse, b: LineUse) bool { + if (a.count != b.count) return a.count < b.count; + if (a.layer != b.layer) return a.layer < b.layer; + return a.id < b.id; +} + +/// `rank_pinned + kv_rank + n * stride + workspace` fits on the card. +/// Neither pin nor KV is reduced. Overflow does not fit. +fn cardHolds(rank_pinned: u64, kv_rank: u64, n: usize, stride: u64, workspace: u64) bool { + const rows: u64 = @intCast(n); + const line_bytes = std.math.mul(u64, rows, stride) catch return false; + const fixed = std.math.add(u64, rank_pinned, kv_rank) catch return false; + const with_lines = std.math.add(u64, fixed, line_bytes) catch return false; + const total = std.math.add(u64, with_lines, workspace) catch return false; + return total <= card_bytes; +} + +/// Lines kept once the batch workspace is placed. Kept rows are packed at the front. +/// `workspace <= device_reserve` keeps every line. +/// A larger workspace drops the coldest rows until +/// `rank_pinned + kv_rank + kept * stride + workspace <= card_bytes`. +/// Coldest is the smallest count, then the smaller layer, then the smaller id. +/// `rank_pinned` and `kv_rank` are not reduced. Zero rows that still miss the card +/// return `error.GraphWorkspace`. +pub fn fitWorkspace( + rank_pinned: u64, + kv_rank: u64, + lines: []LineUse, + stride: u64, + workspace: u64, +) error{GraphWorkspace}!usize { + if (workspace <= device_reserve) return lines.len; + var active: usize = lines.len; + while (active > 0 and !cardHolds(rank_pinned, kv_rank, active, stride, workspace)) { + var cold: usize = 0; + for (lines[1..active], 1..) |line, i| { + if (colder(line, lines[cold])) cold = i; + } + const last = active - 1; + const held = lines[cold]; + lines[cold] = lines[last]; + lines[last] = held; + active = last; + } + if (!cardHolds(rank_pinned, kv_rank, active, stride, workspace)) return error.GraphWorkspace; + return active; +} + +fn capture(window: u32, bindings: []const Binding) Capture { + return .{ .window = window, .bindings = bindings }; +} + +test "replay address change keeps the topology and drops argument equality" { + const fixed: u64 = 0x2000; + const a = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = fixed, .replay = false }, + }; + const b = [_]Binding{ + .{ .name = "expert_ids", .address = 0x3000, .replay = true }, + .{ .name = "dense_weights", .address = fixed, .replay = false }, + }; + const left = capture(4, &a); + const right = capture(4, &b); + try std.testing.expect(sameTopology(left, right)); + try std.testing.expect(!argumentsEqual(left, right)); +} + +test "a new fixed address is a different topology" { + const a = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = 0x2000, .replay = false }, + }; + const b = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = 0x4000, .replay = false }, + }; + try std.testing.expect(!sameTopology(capture(4, &a), capture(4, &b))); + try std.testing.expect(!argumentsEqual(capture(4, &a), capture(4, &b))); +} + +test "a known replay name marked fixed is not a valid capture" { + const bad = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = false }, + }; + try std.testing.expect(!valid(capture(1, &bad))); + try std.testing.expect(isReplayName("expert_ids")); +} + +test "the three replay names and the two fixed names" { + const good = [_]Binding{ + .{ .name = "expert_ids", .address = 1, .replay = true }, + .{ .name = "ple_rows", .address = 3, .replay = true }, + .{ .name = "kv_index", .address = 5, .replay = true }, + .{ .name = "dense_weights", .address = 6, .replay = false }, + .{ .name = "shared_weights", .address = 7, .replay = false }, + }; + try std.testing.expect(valid(capture(2, &good))); + for (replay_names) |name| try std.testing.expect(isReplayName(name)); + try std.testing.expect(!isReplayName("dense_weights")); + try std.testing.expect(!isReplayName("shared_weights")); + + const shared_replay = [_]Binding{ + .{ .name = "shared_weights", .address = 7, .replay = true }, + }; + try std.testing.expect(!valid(capture(2, &shared_replay))); + + const unknown = [_]Binding{ + .{ .name = "token_ids", .address = 9, .replay = true }, + }; + try std.testing.expect(valid(capture(2, &unknown))); + try std.testing.expect(!isReplayName("token_ids")); +} + +test "the pointer tables and picks are replay inputs; pool lines and scales are fixed" { + const t = [_]Binding{ + .{ .name = "expert_tables", .address = 0x100, .replay = true }, + .{ .name = "picks", .address = 0x200, .replay = true }, + .{ .name = "pool_lines", .address = 0x300, .replay = false }, + .{ .name = "expert_scales", .address = 0x400, .replay = false }, + }; + try std.testing.expect(valid(capture(1, &t))); + var bad = t; + bad[2].replay = true; + try std.testing.expect(!valid(capture(1, &bad))); + bad = t; + bad[0].replay = false; + try std.testing.expect(!valid(capture(1, &bad))); + // a pool that reallocated its lines is a new topology (recapture); new table contents are not + var moved = t; + moved[2].address += 256; + try std.testing.expect(!sameTopology(capture(1, &t), capture(1, &moved))); +} + +test "roundReady: an all-hit single-pass round replays straight away, a copy or a second pass stages first" { + const gpa = std.testing.allocator; + var slots: [4]u16 = undefined; + var line: [4]offload.Line = undefined; + var stamp: [4]u32 = undefined; + var lay: [4]u8 = undefined; + var v = offload.Vram.init(&slots, &line, &stamp, &lay); + var counts: [offload.shape.layers][offload.shape.experts]u32 = @splat(@splat(0)); + var r = try round_.Round.init(gpa, 4, 16); + defer r.deinit(gpa); + try round_.plan(&v, 0, &[_]i32{ 1, 2, 512 }, &counts, 1, 1, .{}, &r); + try std.testing.expect(!roundReady(false, &r)); + try round_.plan(&v, 0, &[_]i32{ 2, 1, 512 }, &counts, 2, 2, .{}, &r); + try std.testing.expect(roundReady(false, &r)); + try std.testing.expect(!roundReady(true, &r)); + try round_.plan(&v, 0, &[_]i32{ 1, 2, 3, 4, 5, 6 }, &counts, 3, 3, .{}, &r); + try std.testing.expect(r.masked() and !roundReady(false, &r)); +} + +test "slot_ptrs and hit_mask are not replay names" { + try std.testing.expect(!isReplayName("slot_ptrs")); + try std.testing.expect(!isReplayName("hit_mask")); + try std.testing.expect(isReplayName("ple_rows")); + try std.testing.expect(isReplayName("kv_index")); +} + +test "matching addresses are equal arguments" { + const rows = [_]Binding{ + .{ .name = "expert_ids", .address = 0x11, .replay = true }, + .{ .name = "slot_ptrs", .address = 0x22, .replay = true }, + .{ .name = "dense_weights", .address = 0x33, .replay = false }, + .{ .name = "shared_weights", .address = 0x44, .replay = false }, + }; + const again = [_]Binding{ + .{ .name = "expert_ids", .address = 0x11, .replay = true }, + .{ .name = "slot_ptrs", .address = 0x22, .replay = true }, + .{ .name = "dense_weights", .address = 0x33, .replay = false }, + .{ .name = "shared_weights", .address = 0x44, .replay = false }, + }; + try std.testing.expect(argumentsEqual(capture(1, &rows), capture(1, &again))); + + const other_window = capture(2, &again); + try std.testing.expect(!sameTopology(capture(1, &rows), other_window)); +} + +test "startup blocks capture when every decision is hit and every line is resident" { + const decisions = [_]offload.Decision{ .hit, .hit }; + const lines = [_]offload.Line{ .resident, .resident }; + try std.testing.expect(!captureReady(true, &decisions, &lines)); + try std.testing.expect(captureReady(false, &decisions, &lines)); +} + +test "a copy decision or an inflight line is not capture ready" { + const copy = [_]offload.Decision{ .hit, .copy, .cpu }; + const resident = [_]offload.Line{.resident}; + try std.testing.expect(!captureReady(false, ©, &resident)); + + const hits = [_]offload.Decision{.hit}; + const inflight = [_]offload.Line{ .resident, .inflight, .empty }; + try std.testing.expect(!captureReady(false, &hits, &inflight)); + + const cpu = [_]offload.Decision{.cpu}; + const quiet = [_]offload.Line{ .resident, .empty }; + try std.testing.expect(captureReady(false, &cpu, &quiet)); +} + +test "promote on a ready key is rewrite_ids, not recapture" { + const decisions = [_]offload.Decision{.hit}; + const resident = [_]offload.Line{.resident}; + const key: Key = .{ .rank = 2, .layer = 7, .window = 4 }; + try std.testing.expect(captureReady(false, &decisions, &resident)); + + const have = [_]Key{key}; + try std.testing.expectEqual(Promote.rewrite_ids, promote(&have, key)); + try std.testing.expect(promote(&have, key) != .recapture); + + const other_rank: Key = .{ .rank = 3, .layer = 7, .window = 4 }; + const other_layer: Key = .{ .rank = 2, .layer = 8, .window = 4 }; + const other_window: Key = .{ .rank = 2, .layer = 7, .window = 5 }; + try std.testing.expectEqual(Promote.recapture, promote(&have, other_rank)); + try std.testing.expectEqual(Promote.recapture, promote(&have, other_layer)); + try std.testing.expectEqual(Promote.recapture, promote(&have, other_window)); +} + +test "workspace at device_reserve keeps every line" { + var lines = [_]LineUse{ + .{ .layer = 1, .id = 2, .count = 1 }, + .{ .layer = 0, .id = 9, .count = 50 }, + }; + const stride: u64 = 8 << 30; + try std.testing.expect(2 * stride + device_reserve > card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, device_reserve); + try std.testing.expectEqual(@as(usize, 2), kept); + try std.testing.expectEqual(@as(u8, 1), lines[0].layer); + try std.testing.expectEqual(@as(u16, 2), lines[0].id); + try std.testing.expectEqual(@as(u8, 0), lines[1].layer); +} + +test "workspace above device_reserve drops the coldest line first" { + var lines = [_]LineUse{ + .{ .layer = 2, .id = 1, .count = 10 }, + .{ .layer = 4, .id = 3, .count = 1 }, + .{ .layer = 1, .id = 0, .count = 7 }, + }; + const stride: u64 = 6 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(3 * stride + workspace > card_bytes); + try std.testing.expect(2 * stride + workspace <= card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 2), kept); + try std.testing.expectEqual(@as(u32, 10), lines[0].count); + try std.testing.expectEqual(@as(u8, 2), lines[0].layer); + try std.testing.expectEqual(@as(u32, 7), lines[1].count); + try std.testing.expectEqual(@as(u8, 1), lines[1].layer); +} + +test "equal counts drop the smaller layer first" { + var lines = [_]LineUse{ + .{ .layer = 5, .id = 1, .count = 4 }, + .{ .layer = 2, .id = 9, .count = 4 }, + }; + const stride: u64 = 8 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(2 * stride + workspace > card_bytes); + try std.testing.expect(stride + workspace <= card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 1), kept); + try std.testing.expectEqual(@as(u8, 5), lines[0].layer); + try std.testing.expectEqual(@as(u16, 1), lines[0].id); + try std.testing.expectEqual(@as(u32, 4), lines[0].count); +} + +test "pins and kv stay in the sum" { + const pin: u64 = 8 << 30; + const kv: u64 = 6 << 30; + const stride: u64 = 1 << 30; + const workspace = device_reserve + 1; + // Original pin and kv. Two lines miss the card; one line fits. Dropping either input would keep both. + try std.testing.expect(pin + kv + 2 * stride + workspace > card_bytes); + try std.testing.expect(pin + kv + stride + workspace <= card_bytes); + try std.testing.expect(2 * stride + workspace <= card_bytes); + var lines = [_]LineUse{ + .{ .layer = 1, .id = 1, .count = 3 }, + .{ .layer = 0, .id = 2, .count = 1 }, + }; + const kept = try fitWorkspace(pin, kv, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 1), kept); + try std.testing.expectEqual(@as(u8, 1), lines[0].layer); + try std.testing.expectEqual(@as(u16, 1), lines[0].id); + try std.testing.expectEqual(@as(u32, 3), lines[0].count); +} + +test "zero lines past the card return GraphWorkspace" { + const pin: u64 = 10 << 30; + const kv: u64 = 6 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(pin + kv + workspace > card_bytes); + var lines = [_]LineUse{}; + try std.testing.expectError(error.GraphWorkspace, fitWorkspace(pin, kv, &lines, 1, workspace)); +} diff --git a/zig/src/families/flashnext/exl3_graph_budget.zig b/zig/src/families/flashnext/exl3_graph_budget.zig new file mode 100644 index 0000000000..2a24a7355b --- /dev/null +++ b/zig/src/families/flashnext/exl3_graph_budget.zig @@ -0,0 +1,206 @@ +const std = @import("std"); + +pub const Class = enum { decode, prefill, chain }; +pub const Charge = struct { + exec: u64, + template: u64, + transient: u64, + + pub fn peak(self: Charge) !u64 { + return std.math.add(u64, try std.math.add(u64, self.exec, self.template), self.transient); + } +}; + +pub fn Budget(comptime capacity: usize) type { + return struct { + const Self = @This(); + pub const Ticket = struct { slot: usize, generation: u64 }; + const Entry = struct { + generation: u64, + cohort: u64, + class: Class, + charge: Charge, + phase: enum { pending, retained, retiring } = .pending, + quiescent: bool = false, + used: u64, + }; + pub const Plan = struct { + victims: [capacity]Ticket = undefined, + len: usize = 0, + }; + + limit: u64, + headroom: u64, + entries: [capacity]?Entry = @splat(null), + generation: u64 = 0, + + pub fn init(reservation: u64, headroom: u64) !Self { + if (reservation <= headroom) return error.InvalidGraphReservation; + return .{ .limit = reservation - headroom, .headroom = headroom }; + } + + pub fn charged(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn pending(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + if (entry.phase == .pending) total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn localCharged(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + if (entry.cohort == 0) total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn plan(self: *const Self, charge: Charge) !Plan { + const needed = try charge.peak(); + if (charge.exec == 0 or needed > self.limit) return error.GraphBudgetExceeded; + var result: Plan = .{}; + var remaining = self.charged(); + var free_slot = false; + var selected: [capacity]bool = @splat(false); + for (self.entries) |entry| free_slot = free_slot or entry == null; + while (!free_slot or needed > self.limit - remaining) { + var oldest: ?usize = null; + for (self.entries, 0..) |maybe, index| if (maybe) |entry| { + if (selected[index] or entry.phase != .retained or !entry.quiescent) continue; + if (oldest == null or entry.used < self.entries[oldest.?].?.used) oldest = index; + }; + const index = oldest orelse return error.GraphBudgetBusy; + const entry = self.entries[index].?; + result.victims[result.len] = .{ .slot = index, .generation = entry.generation }; + result.len += 1; + selected[index] = true; + remaining -= try entry.charge.peak(); + free_slot = true; + } + return result; + } + + pub fn reserve(self: *Self, class: Class, cohort: u64, charge: Charge, used: u64) !Ticket { + const proposed = try self.plan(charge); + if (proposed.len != 0) return error.GraphEvictionRequired; + const generation = try std.math.add(u64, self.generation, 1); + for (&self.entries, 0..) |*entry, index| if (entry.* == null) { + entry.* = .{ .generation = generation, .cohort = cohort, .class = class, .charge = charge, .used = used }; + self.generation = generation; + return .{ .slot = index, .generation = generation }; + }; + unreachable; + } + + fn get(self: *Self, ticket: Ticket) !*Entry { + if (ticket.slot >= capacity) return error.StaleGraphTicket; + const entry = &(self.entries[ticket.slot] orelse return error.StaleGraphTicket); + if (entry.generation != ticket.generation) return error.StaleGraphTicket; + return &self.entries[ticket.slot].?; + } + + pub fn installed(self: *Self, ticket: Ticket) !void { + const entry = try self.get(ticket); + if (entry.phase != .pending) return error.InvalidGraphPhase; + entry.phase = .retained; + entry.charge.transient = 0; + } + + pub fn getForRelease(self: *Self, ticket: Ticket) !bool { + return (try self.get(ticket)).phase == .retained; + } + + pub fn cohortOf(self: *Self, ticket: Ticket) !u64 { + return (try self.get(ticket)).cohort; + } + + pub fn touch(self: *Self, ticket: Ticket, used: u64, quiescent: bool) !void { + const entry = try self.get(ticket); + if (entry.phase != .retained) return error.InvalidGraphPhase; + entry.used = used; + entry.quiescent = quiescent; + } + + pub fn beginRetire(self: *Self, ticket: Ticket, agreed_cohort: u64) !void { + const entry = try self.get(ticket); + if (entry.cohort != agreed_cohort) return error.GraphCohortMismatch; + if (entry.phase != .retained or !entry.quiescent) return error.GraphBudgetBusy; + entry.phase = .retiring; + } + + pub fn destroyed(self: *Self, ticket: Ticket) !void { + const entry = try self.get(ticket); + if (entry.phase != .retiring and entry.phase != .pending) return error.InvalidGraphPhase; + self.entries[ticket.slot] = null; + } + }; +} + +test "exl3: graph admission reserves pending peak and explicit headroom" { + var budget = try Budget(4).init(100, 20); + const ticket = try budget.reserve(.decode, 7, .{ .exec = 30, .template = 10, .transient = 20 }, 1); + try std.testing.expectEqual(@as(u64, 60), budget.charged()); + try std.testing.expectError(error.GraphBudgetBusy, budget.plan(.{ .exec = 30, .template = 0, .transient = 0 })); + try budget.installed(ticket); + try std.testing.expectEqual(@as(u64, 40), budget.charged()); + try std.testing.expectError(error.GraphBudgetExceeded, budget.plan(.{ .exec = 81, .template = 0, .transient = 0 })); +} + +test "exl3: pre eviction requires quiescence cohort agreement and completed destruction" { + var budget = try Budget(2).init(120, 20); + const first = try budget.reserve(.prefill, 11, .{ .exec = 40, .template = 10, .transient = 0 }, 1); + try budget.installed(first); + const second = try budget.reserve(.chain, 12, .{ .exec = 40, .template = 10, .transient = 0 }, 2); + try budget.installed(second); + try std.testing.expectError(error.GraphBudgetBusy, budget.plan(.{ .exec = 50, .template = 0, .transient = 0 })); + try budget.touch(first, 1, true); + const plan = try budget.plan(.{ .exec = 50, .template = 0, .transient = 0 }); + try std.testing.expectEqual(@as(usize, 1), plan.len); + try std.testing.expectEqual(first, plan.victims[0]); + try std.testing.expectError(error.GraphCohortMismatch, budget.beginRetire(first, 12)); + try budget.beginRetire(first, 11); + try std.testing.expectEqual(@as(u64, 100), budget.charged()); + try std.testing.expectError(error.GraphBudgetBusy, budget.reserve(.decode, 13, .{ .exec = 50, .template = 0, .transient = 0 }, 3)); + try budget.destroyed(first); + const replacement = try budget.reserve(.decode, 13, .{ .exec = 50, .template = 0, .transient = 0 }, 3); + try std.testing.expectError(error.StaleGraphTicket, budget.destroyed(first)); + try budget.destroyed(replacement); + try std.testing.expectEqual(@as(u64, 50), budget.charged()); +} + +test "exl3: graph estimates reject overflow and invalid reservations" { + try std.testing.expectError(error.InvalidGraphReservation, Budget(1).init(10, 10)); + var budget = try Budget(1).init(100, 10); + try std.testing.expectError(error.Overflow, budget.plan(.{ .exec = std.math.maxInt(u64), .template = 1, .transient = 0 })); + try std.testing.expectError(error.GraphBudgetExceeded, budget.plan(.{ .exec = 0, .template = 0, .transient = 0 })); +} + +test "exl3: cross class LRU planning is read only and launch revokes retirement eligibility" { + var budget = try Budget(3).init(130, 10); + const charge: Charge = .{ .exec = 30, .template = 10, .transient = 0 }; + const decode = try budget.reserve(.decode, 1, charge, 3); + const prefill = try budget.reserve(.prefill, 2, charge, 1); + const chain = try budget.reserve(.chain, 3, charge, 2); + for ([_]Budget(3).Ticket{ decode, prefill, chain }, [_]u64{ 3, 1, 2 }) |ticket, used| { + try budget.installed(ticket); + try budget.touch(ticket, used, true); + } + const request: Charge = .{ .exec = 70, .template = 0, .transient = 0 }; + const plan = try budget.plan(request); + try std.testing.expectEqual(@as(usize, 2), plan.len); + try std.testing.expectEqual(prefill, plan.victims[0]); + try std.testing.expectEqual(chain, plan.victims[1]); + try std.testing.expectEqual(@as(u64, 120), budget.charged()); + try std.testing.expectError(error.GraphEvictionRequired, budget.reserve(.decode, 4, request, 4)); + try budget.touch(prefill, 4, false); + try std.testing.expectError(error.GraphBudgetBusy, budget.beginRetire(prefill, 2)); + try std.testing.expectError(error.InvalidGraphPhase, budget.destroyed(chain)); +} diff --git a/zig/src/families/flashnext/exl3_graph_profile.zig b/zig/src/families/flashnext/exl3_graph_profile.zig new file mode 100644 index 0000000000..3d82e4c2fe --- /dev/null +++ b/zig/src/families/flashnext/exl3_graph_profile.zig @@ -0,0 +1,316 @@ +const std = @import("std"); + +pub const widths = 17; +pub const Presence = enum { absent, present_unready, present_ready }; +pub const Counters = struct { + accesses: u64 = 0, + misses: u64 = 0, + present_unready: u64 = 0, + first_keys: u64 = 0, + returning_misses: u64 = 0, + post_eviction_misses: u64 = 0, + pressure_evictions: u64 = 0, + width_changes: u64 = 0, + state_changes: u64 = 0, + keep_changes: u64 = 0, + seam_changes: u64 = 0, + paging_changes: u64 = 0, + allocation_changes: u64 = 0, + orientation_overflow: u64 = 0, +}; + +fn fingerprint(value: anytype) u64 { + var hash = std.hash.Wyhash.init(0); + std.hash.autoHash(&hash, value); + return hash.final(); +} + +pub const Sample = struct { + kind: usize, + width: usize, + lane: u64, + exact: u64, + state: u64, + keep: u64, + seam: u64, + paging: u64, + allocation: u64, + + pub fn fromKey(key: anytype) Sample { + var allocation = key; + allocation.rows = 0; + allocation.keep = 0; + allocation.chunks = @splat(0); + allocation.score_tiles = @splat(0); + allocation.paging = @splat(@splat(0)); + for (&allocation.seq) |*seq| seq[4] = 0; + var states = key.seq; + for (&states) |*seq| seq.* = .{ 0, 0, 0, 0, seq[4] }; + var origins = key.paging; + for (&origins) |*paging| paging.* = .{ 0, 0, 0, 0, 0, paging[5] }; + var paging_allocations = key.paging; + for (&paging_allocations) |*paging| paging[5] = 0; + var identities = key.seq; + for (&identities) |*seq| seq.* = .{ 0, 0, seq[2], 0, 0 }; + return .{ + .kind = key.kind - 1, + .width = @min(key.rows, widths - 1), + .lane = fingerprint(.{ key.kind, identities }), + .exact = fingerprint(key), + .state = fingerprint(states), + .keep = key.keep, + .seam = fingerprint(.{ key.chunks, key.score_tiles }), + .paging = fingerprint(origins), + .allocation = fingerprint(.{ allocation, paging_allocations }), + }; + } +}; + +pub fn Profile(comptime lane_capacity: usize, comptime key_capacity: usize) type { + return struct { + const Self = @This(); + const Lane = struct { + identity: u64, + last: Sample, + by_width: [2][widths]?Sample = @splat(@splat(null)), + orientations: [2]?u64 = @splat(null), + orientation_accesses: [3]u64 = @splat(0), + width_mask: u32 = 0, + new_keys_by_width: [widths]u64 = @splat(0), + counters: Counters = .{}, + }; + counters: [2][widths]Counters = @splat(@splat(.{})), + lanes: [lane_capacity]?Lane = @splat(null), + keys: [key_capacity]?struct { fingerprint: u64, evicted: bool = false } = @splat(null), + lane_overflow: u64 = 0, + key_overflow: u64 = 0, + + pub fn observeEntry(self: *Self, sample: Sample, presence: Presence) void { + self.observe(sample, presence != .absent); + if (presence != .present_unready) return; + self.counters[sample.kind][sample.width].present_unready +|= 1; + for (&self.lanes) |*lane| if (lane.*) |*value| { + if (value.identity == sample.lane) { + value.counters.present_unready +|= 1; + return; + } + }; + } + + pub fn observe(self: *Self, sample: Sample, hit: bool) void { + const counts = &self.counters[sample.kind][sample.width]; + counts.accesses +|= 1; + if (!hit) counts.misses +|= 1; + var known = false; + var empty: ?usize = null; + for (&self.keys, 0..) |*key, index| { + if (key.*) |*value| { + if (value.fingerprint == sample.exact) { + known = true; + if (value.evicted and !hit) counts.post_eviction_misses +|= 1; + value.evicted = false; + } + } else if (empty == null) empty = index; + } + if (known) { + if (!hit) counts.returning_misses +|= 1; + } else if (empty) |index| { + self.keys[index] = .{ .fingerprint = sample.exact }; + counts.first_keys +|= 1; + } else self.key_overflow +|= 1; + for (&self.lanes) |*lane| if (lane.*) |*value| { + if (value.identity == sample.lane) return compare(value, counts, sample, hit, !known and empty != null, known); + }; + for (&self.lanes) |*lane| if (lane.* == null) { + lane.* = .{ .identity = sample.lane, .last = sample }; + return compare(&lane.*.?, counts, sample, hit, !known and empty != null, known); + }; + self.lane_overflow +|= 1; + } + + fn compare(lane: *Lane, counts: *Counters, sample: Sample, hit: bool, new_key: bool, known: bool) void { + lane.counters.accesses +|= 1; + if (!hit) lane.counters.misses +|= 1; + if (new_key) lane.counters.first_keys +|= 1; + if (new_key) lane.new_keys_by_width[sample.width] +|= 1; + if (known and !hit) lane.counters.returning_misses +|= 1; + lane.width_mask |= @as(u32, 1) << @intCast(sample.width); + if (lane.last.width != sample.width) { + counts.width_changes +|= 1; + lane.counters.width_changes +|= 1; + } + if (lane.by_width[sample.kind][sample.width]) |previous| { + inline for (.{ .{ "state", "state_changes" }, .{ "keep", "keep_changes" }, .{ "seam", "seam_changes" }, .{ "paging", "paging_changes" }, .{ "allocation", "allocation_changes" } }) |fields| { + if (@field(previous, fields[0]) != @field(sample, fields[0])) { + @field(counts, fields[1]) +|= 1; + @field(lane.counters, fields[1]) +|= 1; + } + } + } + var orientation: usize = 2; + for (&lane.orientations, 0..) |*state, index| { + if (state.* == null) state.* = sample.state; + if (state.*.? == sample.state) { + orientation = index; + break; + } + } + lane.orientation_accesses[orientation] +|= 1; + if (orientation == 2) { + counts.orientation_overflow +|= 1; + lane.counters.orientation_overflow +|= 1; + } + lane.last = sample; + lane.by_width[sample.kind][sample.width] = sample; + } + + pub fn evict(self: *Self, sample: Sample) void { + self.counters[sample.kind][sample.width].pressure_evictions +|= 1; + for (&self.keys) |*key| if (key.*) |*value| { + if (value.fingerprint == sample.exact) value.evicted = true; + }; + for (&self.lanes) |*lane| if (lane.*) |*value| { + if (value.identity == sample.lane) { + value.counters.pressure_evictions +|= 1; + return; + } + }; + } + }; +} + +fn fixture() Sample { + return .{ .kind = 1, .width = 2, .lane = 1, .exact = 1, .state = 1, .keep = 1, .seam = 1, .paging = 1, .allocation = 1 }; +} + +test "exl3: profile entry presence does not imply replay readiness or replay coverage" { + var profile: Profile(2, 8) = .{}; + const sample = fixture(); + profile.observeEntry(sample, .absent); + profile.observeEntry(sample, .present_unready); + profile.observeEntry(sample, .present_ready); + try std.testing.expectEqual(@as(u64, 3), profile.counters[1][2].accesses); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].misses); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].present_unready); + try std.testing.expectEqual(@as(u64, 1), profile.lanes[0].?.counters.present_unready); + profile.counters[1][2].present_unready = std.math.maxInt(u64); + profile.observeEntry(sample, .present_unready); + try std.testing.expectEqual(std.math.maxInt(u64), profile.counters[1][2].present_unready); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].misses); +} + +test "exl3: profile distinguishes new keys returning misses and pressure evictions" { + var profile: Profile(2, 8) = .{}; + const sample = fixture(); + profile.observe(sample, false); + profile.observe(sample, true); + profile.evict(sample); + profile.observe(sample, false); + const counts = profile.counters[1][2]; + try std.testing.expectEqual(@as(u64, 3), counts.accesses); + try std.testing.expectEqual(@as(u64, 2), counts.misses); + try std.testing.expectEqual(@as(u64, 1), counts.first_keys); + try std.testing.expectEqual(@as(u64, 1), counts.returning_misses); + try std.testing.expectEqual(@as(u64, 1), counts.pressure_evictions); + try std.testing.expectEqual(@as(u64, 1), counts.post_eviction_misses); + try std.testing.expectEqual(@as(u64, 1), profile.lanes[0].?.new_keys_by_width[2]); +} + +test "exl3: profile couples state and keep and separates width baselines" { + var profile: Profile(2, 8) = .{}; + var sample = fixture(); + profile.observe(sample, false); + sample.width = 1; + sample.keep = 0; + sample.exact = 2; + profile.observe(sample, false); + sample.width = 2; + sample.keep = 2; + sample.state = 2; + sample.exact = 3; + profile.observe(sample, false); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].state_changes); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].keep_changes); + try std.testing.expectEqual(@as(u64, 0), profile.counters[1][1].keep_changes); + try std.testing.expectEqual([3]u64{ 2, 1, 0 }, profile.lanes[0].?.orientation_accesses); +} + +test "exl3: profile separates seam paging and allocation changes" { + var profile: Profile(2, 8) = .{}; + var sample = fixture(); + profile.observe(sample, false); + sample.seam = 2; + profile.observe(sample, false); + sample.paging = 2; + profile.observe(sample, false); + sample.allocation = 2; + profile.observe(sample, false); + const counts = profile.counters[1][2]; + try std.testing.expectEqual(@as(u64, 1), counts.seam_changes); + try std.testing.expectEqual(@as(u64, 1), counts.paging_changes); + try std.testing.expectEqual(@as(u64, 1), counts.allocation_changes); + try std.testing.expectEqual(@as(u64, 0), counts.state_changes); +} + +test "exl3: profile bounds lanes keys orientations and saturates counters" { + var profile: Profile(1, 1) = .{}; + var sample = fixture(); + profile.observe(sample, false); + sample.state = 2; + profile.observe(sample, true); + sample.state = 3; + profile.observe(sample, true); + sample.lane = 2; + sample.exact = 2; + profile.observe(sample, false); + try std.testing.expectEqual(@as(u64, 1), profile.lane_overflow); + try std.testing.expectEqual(@as(u64, 1), profile.key_overflow); + try std.testing.expectEqual(@as(u64, 1), profile.counters[1][2].orientation_overflow); + profile.counters[1][2].accesses = std.math.maxInt(u64); + profile.observe(sample, true); + try std.testing.expectEqual(std.math.maxInt(u64), profile.counters[1][2].accesses); +} + +test "exl3: profile key projection separates orientation seam paging and allocations without mutating key" { + const Key = struct { + kind: u8 = 2, + rows: usize = 2, + seq: [2][5]u64 = .{ .{ 1, 2, 3, 4, 5 }, @splat(0) }, + paging: [2][6]u64 = .{ .{ 6, 7, 8, 9, 10, 11 }, @splat(0) }, + chunks: [2]u32 = .{ 1, 0 }, + score_tiles: [2]u32 = @splat(0), + scratch: u64 = 12, + scratch_cap: usize = 32, + scores: u64 = 13, + output: u64 = 14, + streams: u64 = 15, + keep: u64 = 16, + plan: u64 = 17, + }; + const original: Key = .{}; + var key = original; + const base = Sample.fromKey(key); + key.seq[0][4] += 1; + key.keep += 1; + const orientation = Sample.fromKey(key); + try std.testing.expect(base.state != orientation.state); + try std.testing.expectEqual(base.allocation, orientation.allocation); + try std.testing.expectEqual(base.lane, orientation.lane); + key = original; + key.chunks[0] += 1; + const seam = Sample.fromKey(key); + try std.testing.expect(base.seam != seam.seam); + try std.testing.expectEqual(base.allocation, seam.allocation); + key = original; + key.paging[0][5] += 1; + const paging = Sample.fromKey(key); + try std.testing.expect(base.paging != paging.paging); + try std.testing.expectEqual(base.allocation, paging.allocation); + key = original; + key.paging[1][0] += 1; + try std.testing.expect(base.allocation != Sample.fromKey(key).allocation); + key = original; + key.scratch += 1; + try std.testing.expect(base.allocation != Sample.fromKey(key).allocation); + try std.testing.expectEqualDeep(original, Key{}); +} diff --git a/zig/src/families/flashnext/exl3_hc.zig b/zig/src/families/flashnext/exl3_hc.zig new file mode 100644 index 0000000000..a88aec4bb3 --- /dev/null +++ b/zig/src/families/flashnext/exl3_hc.zig @@ -0,0 +1,393 @@ +//! Host side of the unquantized hc_lowrank split. +//! Each rank already owns 160 of the 320-wide axis and produces a hidden partial. +//! Those partials meet here, rank 0 then rank 1. No CUDA. + +const std = @import("std"); +const shape = @import("exl3_shape.zig"); +const tp = @import("exl3_tp.zig"); + +/// One rank's dense product. `w` is row-major `[hc_lowrank/2][hidden]`: column `j` is strided by `hidden`, +/// matching `exl3_cpu.gemv`, so `out[j]` is the dot of `x` with that column. +/// A length other than those three sizes returns `error.Length` before any store. Does not allocate. +pub fn partial(w: []const f32, x: []const f32, out: []f32) error{Length}!void { + const part = shape.hc_lowrank / 2; + const n = shape.hidden; + if (w.len != part * n or x.len != part or out.len != n) return error.Length; + for (0..n) |j| { + var sum: f32 = 0; + for (0..part) |i| sum += x[i] * w[i * n + j]; + out[j] = sum; + } +} + +/// Hidden-sized sum of the two partials. Rank 0, then rank 1. Not inside `partial`. +/// A length other than `hidden` returns `error.Length` before any store. +pub fn add(partial0: []const f32, partial1: []const f32, out: []f32) error{Length}!void { + const n = shape.hidden; + if (partial0.len != n or partial1.len != n or out.len != n) return error.Length; + for (partial0, partial1, out) |a, b, *y| y.* = a + b; +} + +/// `exl3_bounce_bytes` for hc_lowrank only. +/// D2H of this partial (2560 values, 5120 bytes) plus H2D of the summed hidden, one layer. +/// Not multiplied by 48. MoE and GDN are not included. +pub fn bounceBytes(tokens: usize) u64 { + return @as(u64, @intCast(tokens)) * 5120 * 2; +} + +// --------------------------------------------------------------------------------------------- +// Hyper-connection pre / post, expand / collapse. Host f32, with an f64 oracle. +// +// Where it runs: on the host, like the lowrank split above. The three matrices are dense F16 in the pack +// (`input_mix_weight_down` [320][10240], `input_mix_weight_up` [10240][320], `block_inject_weight` [4][10240], +// all out-major, i.e. torch `nn.Linear` weights), they are not EXL3, and the existing design bounces the +// block output to the host (`bounceBytes`) and the mixed hidden back. 10240x320 MACs twice is ~7 M flops. +// +// Maths, from the tensorfold MLX reference (`families/qwen4_exp/model_layers.py` `HyperConnection`, +// `_write_back`, `model.py` `DecoderLayer`/`Body`/`Qwen4Exp.hidden`), with S streams of D, wide = S*D, stream-major +// (`h[s * D + d]`, the `reshape(S, D)` of `tile(embed, S)`): +// normed = per-stream centred RMSNorm: normed[s, d] = h[s, d] / sqrt(mean_d(h[s, :]^2) + eps) * (norm[s*D+d] + offset) +// low = silu((down @ normed) / S) // [lowrank] +// mix = sigmoid(up @ low) // [wide] +// mixed = mean_s(mix[s, d] * normed[s, d]) // block input [D] +// inject = 2 * sigmoid((block_inject @ normed) / S) // [S], absent for the final collapse +// post: h[s, d] += branch[d] * inject[s] +// expand (model start): h[s, d] = embed[d] for every s. collapse (model end): `pre` of +// `hyper_connection_mixer` (no inject), its `mixed` is the hidden the head reads (no final norm in the pack). +// `offset` is `scan.centredOffset` over the attention hc_norm means (1 when stored as gamma-1). +// --------------------------------------------------------------------------------------------- + +pub const Dims = struct { streams: usize, dim: usize, lowrank: usize }; +pub const real: Dims = .{ .streams = shape.hc_count, .dim = shape.hidden, .lowrank = shape.hc_lowrank }; + +/// Row-major, out-major weights as f32 (exact from F16). `inject` is null for the final mixer. +pub const Weights = struct { + norm: []const f32, + down: []const f32, + up: []const f32, + inject: ?[]const f32, +}; + +pub const Error = error{Length}; + +fn silu32(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} +fn sigmoid32(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f32 values `scratch` must hold: normed (wide), low (lowrank), mix (wide). +pub fn scratchLen(d: Dims) usize { + return 2 * d.streams * d.dim + d.lowrank; +} + +fn checkWeights(d: Dims, w: Weights) Error!void { + const wide = d.streams * d.dim; + if (w.norm.len != wide or w.down.len != d.lowrank * wide or w.up.len != wide * d.lowrank) return error.Length; + if (w.inject) |inj| if (inj.len != d.streams * wide) return error.Length; +} + +/// Block input `mixed` [dim] and, when the weights carry `block_inject_weight` and `inject` is non-null, the S gates. +/// `h` is [streams * dim]. Lengths are checked before any store. Does not allocate. +pub fn pre(d: Dims, w: Weights, offset: f32, eps: f32, h: []const f32, mixed: []f32, inject: ?[]f32, scratch: []f32) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or scratch.len < scratchLen(d)) return error.Length; + if (inject) |g| if (w.inject == null or g.len != d.streams) return error.Length; + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + const normed = scratch[0..wide]; + const low = scratch[wide .. wide + d.lowrank]; + const mix = scratch[wide + d.lowrank .. 2 * wide + d.lowrank]; + + for (0..d.streams) |s| { + const row = h[s * d.dim ..][0..d.dim]; + var sum_sq: f32 = 0; + for (row) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(d.dim)) + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = row[i] / rms * (w.norm[s * d.dim + i] + offset); + } + const inv_s: f32 = 1.0 / @as(f32, @floatFromInt(d.streams)); + for (0..d.lowrank) |r| { + var acc: f32 = 0; + const row = w.down[r * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + low[r] = silu32(acc * inv_s); + } + for (0..wide) |j| { + var acc: f32 = 0; + const row = w.up[j * d.lowrank ..][0..d.lowrank]; + for (row, low) |a, b| acc += a * b; + mix[j] = sigmoid32(acc); + } + for (0..d.dim) |i| { + var acc: f32 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc * inv_s; + } + if (inject) |g| { + const inj = w.inject.?; + for (0..d.streams) |s| { + var acc: f32 = 0; + const row = inj[s * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + g[s] = 2.0 * sigmoid32(acc * inv_s); + } + } +} + +/// `out[s, d] = h[s, d] + branch[d] * inject[s]`. `out` may be `h`. Lengths are checked before any store. +pub fn post(d: Dims, h: []const f32, branch: []const f32, inject: []const f32, out: []f32) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| { + for (0..d.dim) |i| out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + } +} + +/// Model start: the embedding row copied into every stream. +pub fn expand(d: Dims, embed: []const f32, out: []f32) Error!void { + if (embed.len != d.dim or out.len != d.streams * d.dim) return error.Length; + for (0..d.streams) |s| @memcpy(out[s * d.dim ..][0..d.dim], embed); +} + +/// The same maths in f64, written as plain loops over (stream, dim) with no scratch, for the oracle. +/// `mixed` is [dim]; `inject` [streams] or null. Allocation-free. +pub fn preF64(d: Dims, w: Weights, offset: f64, eps: f64, h: []const f64, mixed: []f64, inject: ?[]f64, low: []f64, mix: []f64, normed: []f64) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or low.len != d.lowrank or mix.len != wide or normed.len != wide) return error.Length; + for (0..d.streams) |s| { + var ms: f64 = 0; + for (0..d.dim) |i| ms += h[s * d.dim + i] * h[s * d.dim + i]; + ms /= @floatFromInt(d.dim); + const inv = 1.0 / @sqrt(ms + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = h[s * d.dim + i] * inv * (@as(f64, w.norm[s * d.dim + i]) + offset); + } + const sf: f64 = @floatFromInt(d.streams); + for (0..d.lowrank) |r| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, w.down[r * wide + c]) * normed[c]; + const z = acc / sf; + low[r] = z / (1.0 + @exp(-z)); + } + for (0..wide) |j| { + var acc: f64 = 0; + for (0..d.lowrank) |r| acc += @as(f64, w.up[j * d.lowrank + r]) * low[r]; + mix[j] = 1.0 / (1.0 + @exp(-acc)); + } + for (0..d.dim) |i| { + var acc: f64 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc / sf; + } + if (inject) |g| { + const inj = w.inject orelse return error.Length; + if (g.len != d.streams) return error.Length; + for (0..d.streams) |s| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, inj[s * wide + c]) * normed[c]; + g[s] = 2.0 / (1.0 + @exp(-acc / sf)); + } + } +} + +pub fn postF64(d: Dims, h: []const f64, branch: []const f64, inject: []const f64, out: []f64) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| for (0..d.dim) |i| { + out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + }; +} + +const tiny: Dims = .{ .streams = 2, .dim = 4, .lowrank = 3 }; + +fn tinyWeights(gpa: std.mem.Allocator, with_inject: bool) !struct { norm: []f32, down: []f32, up: []f32, inject: ?[]f32 } { + const wide = tiny.streams * tiny.dim; + var rng = std.Random.DefaultPrng.init(7); + const r = rng.random(); + const norm = try gpa.alloc(f32, wide); + const down = try gpa.alloc(f32, tiny.lowrank * wide); + const up = try gpa.alloc(f32, wide * tiny.lowrank); + for (norm) |*v| v.* = r.float(f32) - 0.5; + for (down) |*v| v.* = (r.float(f32) - 0.5) * 2; + for (up) |*v| v.* = (r.float(f32) - 0.5) * 2; + var inject: ?[]f32 = null; + if (with_inject) { + inject = try gpa.alloc(f32, tiny.streams * wide); + for (inject.?) |*v| v.* = (r.float(f32) - 0.5) * 2; + } + return .{ .norm = norm, .down = down, .up = up, .inject = inject }; +} + +test "pre and post match the f64 plain-loop oracle on a tiny shape" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + const h = [wide]f32{ 1, -2, 3, 0.5, 4, 1, -1, 2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 1e-6, &h, &mixed, &gates, &scratch); + + var h64: [wide]f64 = undefined; + for (&h64, h) |*a, b| a.* = b; + var mixed64: [tiny.dim]f64 = undefined; + var gates64: [tiny.streams]f64 = undefined; + var low: [tiny.lowrank]f64 = undefined; + var mix: [wide]f64 = undefined; + var normed: [wide]f64 = undefined; + try preF64(tiny, w, 1.0, 1e-6, &h64, &mixed64, &gates64, &low, &mix, &normed); + for (mixed, mixed64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + for (gates, gates64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + + const branch = [tiny.dim]f32{ 0.5, -1, 2, 0 }; + var out: [wide]f32 = undefined; + try post(tiny, &h, &branch, &gates, &out); + for (0..tiny.streams) |s| for (0..tiny.dim) |i| { + try std.testing.expectApproxEqAbs(h[s * tiny.dim + i] + branch[i] * gates[s], out[s * tiny.dim + i], 1e-6); + }; + // aliasing out and h is allowed + var again = h; + try post(tiny, &again, &branch, &gates, &again); + try std.testing.expectEqualSlices(f32, &out, &again); +} + +test "zero mixing weights give mix 1/2, inject 1, and the offset-only norm" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + @memset(t.down, 0); + @memset(t.up, 0); + @memset(t.inject.?, 0); + @memset(t.norm, 0); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + // stream 0 all 3s, stream 1 all -2s: each stream normalises to +-1 (offset 1, weight 0). + const h = [wide]f32{ 3, 3, 3, 3, -2, -2, -2, -2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 0.0, &h, &mixed, &gates, &scratch); + // mean_s(0.5 * (+1, -1)) = 0 + for (mixed) |v| try std.testing.expectApproxEqAbs(@as(f32, 0), v, 1e-6); + for (gates) |g| try std.testing.expectApproxEqAbs(@as(f32, 1), g, 1e-6); + // gamma stored directly: offset 0 and weight 0 zeroes the normed streams + try pre(tiny, w, 0.0, 0.0, &h, &mixed, &gates, &scratch); + for (mixed) |v| try std.testing.expectEqual(@as(f32, 0), v); +} + +test "pre, post and expand refuse bad lengths before a store" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, false); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = null }; + const wide = tiny.streams * tiny.dim; + const h: [wide]f32 = @splat(1); + var mixed: [tiny.dim]f32 = @splat(9); + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, h[1..], &mixed, null, &scratch)); + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, &gates, &scratch)); // no inject weights + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, null, scratch[0..3])); + try std.testing.expectEqual(@as(f32, 9), mixed[0]); + var out: [wide]f32 = undefined; + try std.testing.expectError(error.Length, post(tiny, &h, mixed[0..3], &gates, &out)); + try expand(tiny, &.{ 1, 2, 3, 4 }, &out); + try std.testing.expectEqualSlices(f32, &.{ 1, 2, 3, 4, 1, 2, 3, 4 }, &out); + try std.testing.expectError(error.Length, expand(tiny, &.{ 1, 2, 3 }, &out)); +} + +test "one-hot x picks the row of the hc_lowrank part" { + const cut0 = try tp.cut(.hc_lowrank, 0, 2); + try std.testing.expectEqual(@as(usize, 160), cut0.part); + const r0 = try tp.range(.hc_lowrank, 0, 2); + const r1 = try tp.range(.hc_lowrank, 1, 2); + try std.testing.expectEqual(@as(usize, 0), r0.lo); + try std.testing.expectEqual(@as(usize, 160), r0.hi); + try std.testing.expectEqual(@as(usize, 160), r1.lo); + try std.testing.expectEqual(@as(usize, 320), r1.hi); + try std.testing.expectEqual(cut0.part, r0.hi - r0.lo); + try std.testing.expectEqual(cut0.part, r1.hi - r1.lo); + + const part = cut0.part; + const n = shape.hidden; + try std.testing.expectEqual(@as(usize, 160 * 2560), part * n); + + const gpa = std.testing.allocator; + const w = try gpa.alloc(f32, part * n); + defer gpa.free(w); + @memset(w, 0); + + var x: [shape.hc_lowrank / 2]f32 = @splat(0); + try std.testing.expectEqual(part, x.len); + + var rank0: [shape.hidden]f32 = @splat(42); + const kept0 = rank0; + try std.testing.expectError(error.Length, partial(w[0..part], &x, &rank0)); + try std.testing.expectError(error.Length, partial(w, x[0 .. part - 1], &rank0)); + try std.testing.expectError(error.Length, partial(w, &x, rank0[0 .. n - 1])); + try std.testing.expectEqualSlices(f32, &kept0, &rank0); + + // Rank 0 owns global [0, 160). A one-hot at local 0 is that first row, not the poisoned neighbor. + const local0: usize = 0; + try std.testing.expectEqual(r0.lo + local0, local0); + for (0..n) |j| w[local0 * n + j] = @floatFromInt(j + 3); + for (0..n) |j| w[1 * n + j] = -1000; + x[local0] = 1; + try partial(w, &x, &rank0); + try std.testing.expectEqualSlices(f32, w[local0 * n .. local0 * n + n], &rank0); + + // Rank 1 owns global [160, 320). A one-hot at the last local row is global 319. + var rank1: [shape.hidden]f32 = @splat(-7); + @memset(w, 0); + const local1 = part - 1; + try std.testing.expectEqual(@as(usize, 319), r1.lo + local1); + try std.testing.expect(r1.lo + local1 < r1.hi); + for (0..n) |j| w[0 * n + j] = 80; + for (0..n) |j| w[local1 * n + j] = -@as(f32, @floatFromInt(j)) - 1; + @memset(&x, 0); + x[local1] = 1; + try partial(w, &x, &rank1); + try std.testing.expectEqualSlices(f32, w[local1 * n .. local1 * n + n], &rank1); + + // Two nonzeros so the product is a dot, not a copy of the hot row. + @memset(w, 0); + @memset(&x, 0); + const row_a: usize = 2; + const row_b: usize = 5; + for (0..n) |j| w[row_a * n + j] = 1; + for (0..n) |j| w[row_b * n + j] = 4; + x[row_a] = 2; + x[row_b] = 0.5; + var dotted: [shape.hidden]f32 = @splat(0); + try partial(w, &x, &dotted); + const fours: [shape.hidden]f32 = @splat(4); + try std.testing.expectEqualSlices(f32, &fours, &dotted); + + var out: [shape.hidden]f32 = @splat(7); + const kept = out; + try std.testing.expectError(error.Length, add(rank0[0 .. n - 1], &rank1, &out)); + try std.testing.expectError(error.Length, add(&rank0, rank1[0..0], &out)); + try std.testing.expectError(error.Length, add(&rank0, &rank1, out[0 .. n - 1])); + try std.testing.expectEqualSlices(f32, &kept, &out); + + // (j + 3) + (-j - 1) = 2. Rank 0, then rank 1. + try add(&rank0, &rank1, &out); + const twos: [shape.hidden]f32 = @splat(2); + try std.testing.expectEqualSlices(f32, &twos, &out); +} + +test "bounceBytes is the hc_lowrank exchange for one layer" { + try std.testing.expectEqual(@as(u64, 10240), bounceBytes(1)); + try std.testing.expectEqual(@as(u64, 0), bounceBytes(0)); + try std.testing.expectEqual(@as(u64, 10240) * 3, bounceBytes(3)); +} diff --git a/zig/src/families/flashnext/exl3_kv.zig b/zig/src/families/flashnext/exl3_kv.zig new file mode 100644 index 0000000000..e2bdde6b37 --- /dev/null +++ b/zig/src/families/flashnext/exl3_kv.zig @@ -0,0 +1,182 @@ +//! The buffer is `contextKvRank` for this rank's heads. + +const std = @import("std"); +const shape = @import("exl3_shape.zig"); +const cache = @import("exl3_cache.zig"); + +/// Full-attention layers only. GDN layers have no KV table. Same split as `contextKvRank`. +const full_layers: usize = shape.layers / shape.full_attention_interval; + +pub const Table = struct { + rank: usize, + ranks: usize, + head: usize, + /// Two ranks store one head each, not both on rank 0. + heads: usize, + tokens: usize, + dtype: cache.Dtype, + bytes: []u8, + + /// `contextKvRank` refuses a zero rank count, or a count that does not divide + /// the KV heads, before this divides. + pub fn init( + gpa: std.mem.Allocator, + tokens: usize, + dtype: cache.Dtype, + rank: usize, + ranks: usize, + ) error{ HeadDim, Ranks, RankOutOfRange, OutOfMemory }!Table { + const n = try cache.contextKvRank(tokens, dtype, ranks); + if (rank >= ranks) return error.RankOutOfRange; + const heads = shape.kv_heads / ranks; + const bytes = try gpa.alloc(u8, @intCast(n)); + @memset(bytes, 0); + return .{ + .rank = rank, + .ranks = ranks, + .head = rank * heads, + .heads = heads, + .tokens = tokens, + .dtype = dtype, + .bytes = bytes, + }; + } + + pub fn deinit(self: *Table, gpa: std.mem.Allocator) void { + gpa.free(self.bytes); + } + + /// One head, one token, one full-attention layer. `layer` is the table index + /// (not the model layer). A head this rank does not store is refused before any store. + pub fn write(self: *Table, head: usize, layer: usize, token: usize, src: []const u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, src.len); + @memcpy(self.bytes[off..][0..src.len], src); + } + + pub fn read(self: *const Table, head: usize, layer: usize, token: usize, dest: []u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, dest.len); + @memcpy(dest, self.bytes[off..][0..dest.len]); + } + + fn locate(self: *const Table, head: usize, layer: usize, token: usize, len: usize) error{ Head, Slot, Length }!usize { + if (self.heads == 0 or head < self.head or head >= self.head + self.heads) return error.Head; + if (self.tokens == 0 or full_layers == 0 or layer >= full_layers or token >= self.tokens) return error.Slot; + const denom = full_layers * self.tokens * self.heads; + if (self.bytes.len % denom != 0) return error.Slot; + const one = self.bytes.len / denom; + if (len != one) return error.Length; + const local = head - self.head; + return ((layer * self.tokens + token) * self.heads + local) * one; + } +}; + +test "one rank stores one head, sized by contextKvRank, not both heads" { + const gpa = std.testing.allocator; + const tokens: usize = 8192; + const ranks: usize = 2; + const both = try cache.contextKv(tokens, .bf16); + try std.testing.expectEqual(@as(u64, 201_326_592), both); + + { + var r0 = try Table.init(gpa, tokens, .bf16, 0, ranks); + defer r0.deinit(gpa); + var r1 = try Table.init(gpa, tokens, .bf16, 1, ranks); + defer r1.deinit(gpa); + try std.testing.expectEqual(@as(usize, 0), r0.head); + try std.testing.expectEqual(@as(usize, 1), r0.heads); + try std.testing.expectEqual(@as(usize, 1), r1.head); + try std.testing.expectEqual(@as(usize, 1), r1.heads); + try std.testing.expectEqual(@as(u64, 100_663_296), @as(u64, r0.bytes.len)); + try std.testing.expectEqual(r0.bytes.len, r1.bytes.len); + try std.testing.expectEqual(try cache.contextKvRank(tokens, .bf16, ranks), @as(u64, r0.bytes.len)); + try std.testing.expect(r0.bytes.len != both); + try std.testing.expect(r1.bytes.len != both); + } + + { + var t = try Table.init(gpa, tokens, .int8, 0, ranks); + defer t.deinit(gpa); + try std.testing.expectEqual(@as(usize, 1), t.heads); + try std.testing.expectEqual(@as(u64, 53_477_376), @as(u64, t.bytes.len)); + try std.testing.expectEqual(try cache.contextKvRank(tokens, .int8, ranks), @as(u64, t.bytes.len)); + try std.testing.expect(t.bytes.len != try cache.contextKv(tokens, .int8)); + } + + { + var t = try Table.init(gpa, tokens, .int4, 1, ranks); + defer t.deinit(gpa); + try std.testing.expectEqual(@as(usize, 1), t.head); + try std.testing.expectEqual(@as(usize, 1), t.heads); + try std.testing.expectEqual(@as(u64, 28_311_552), @as(u64, t.bytes.len)); + try std.testing.expectEqual(try cache.contextKvRank(tokens, .int4, ranks), @as(u64, t.bytes.len)); + try std.testing.expect(t.bytes.len != try cache.contextKv(tokens, .int4)); + } + + var none = try Table.init(gpa, 0, .bf16, 0, ranks); + defer none.deinit(gpa); + try std.testing.expectEqual(@as(usize, 0), none.bytes.len); + try std.testing.expectEqual(@as(usize, 1), none.heads); +} + +test "a bad rank is error.Ranks and not a zero-byte table" { + const gpa = std.testing.allocator; + try std.testing.expectError(error.Ranks, Table.init(gpa, 8192, .bf16, 0, 0)); + try std.testing.expectError(error.Ranks, Table.init(gpa, 8192, .int8, 0, 0)); + try std.testing.expectError(error.Ranks, Table.init(gpa, 8192, .int4, 1, 0)); + try std.testing.expectError(error.Ranks, Table.init(gpa, 8192, .bf16, 0, 3)); + try std.testing.expectError(error.Ranks, Table.init(gpa, 8192, .int8, 0, 3)); + try std.testing.expectError(error.RankOutOfRange, Table.init(gpa, 8, .bf16, 2, 2)); + try std.testing.expectError(error.RankOutOfRange, Table.init(gpa, 8, .int4, 1, 1)); +} + +test "a write on one rank does not change the other rank's head" { + const gpa = std.testing.allocator; + const ranks: usize = 2; + const tokens: usize = 2; + const cases = [_]struct { dtype: cache.Dtype, slot: usize, mark: u8 }{ + .{ .dtype = .bf16, .slot = 1024, .mark = 0xa5 }, + .{ .dtype = .int8, .slot = 544, .mark = 0x11 }, + .{ .dtype = .int4, .slot = 288, .mark = 0x22 }, + }; + var src: [1024]u8 = undefined; + var got: [1024]u8 = undefined; + + for (cases) |case| { + var r0 = try Table.init(gpa, tokens, case.dtype, 0, ranks); + defer r0.deinit(gpa); + var r1 = try Table.init(gpa, tokens, case.dtype, 1, ranks); + defer r1.deinit(gpa); + const slot = src[0..case.slot]; + @memset(slot, case.mark); + + try std.testing.expectError(error.Head, r0.write(1, 0, 0, slot)); + try std.testing.expectError(error.Head, r1.write(0, 0, 0, slot)); + try std.testing.expectError(error.Length, r0.write(0, 0, 0, slot[0 .. case.slot - 1])); + try std.testing.expectError(error.Slot, r0.write(0, 0, tokens, slot)); + try std.testing.expectError(error.Slot, r0.write(0, full_layers, 0, slot)); + try std.testing.expect(std.mem.allEqual(u8, r0.bytes, 0)); + try std.testing.expect(std.mem.allEqual(u8, r1.bytes, 0)); + + try r0.write(0, 0, 1, slot); + try std.testing.expect(std.mem.allEqual(u8, r1.bytes, 0)); + @memset(&got, 0); + try r0.read(0, 0, 1, got[0..case.slot]); + try std.testing.expectEqualSlices(u8, slot, got[0..case.slot]); + try std.testing.expectError(error.Head, r0.read(1, 0, 1, got[0..case.slot])); + + const snap = try gpa.dupe(u8, r0.bytes); + defer gpa.free(snap); + @memset(slot, case.mark +% 1); + try r1.write(1, 1, 0, slot); + try std.testing.expectEqualSlices(u8, snap, r0.bytes); + try r1.read(1, 1, 0, got[0..case.slot]); + try std.testing.expectEqualSlices(u8, slot, got[0..case.slot]); + try std.testing.expectError(error.Head, r1.read(0, 1, 0, got[0..case.slot])); + } + + var none = try Table.init(gpa, 0, .bf16, 0, ranks); + defer none.deinit(gpa); + var scratch: [0]u8 = .{}; + try std.testing.expectError(error.Slot, none.write(0, 0, 0, &scratch)); + try std.testing.expectError(error.Slot, none.read(0, 0, 0, &scratch)); +} diff --git a/zig/src/families/flashnext/exl3_line_deltas.zig b/zig/src/families/flashnext/exl3_line_deltas.zig new file mode 100644 index 0000000000..480414f23d --- /dev/null +++ b/zig/src/families/flashnext/exl3_line_deltas.zig @@ -0,0 +1,132 @@ +const std = @import("std"); +const memory = @import("exl3_startup_memory.zig"); +const layout = @import("exl3_dplan_layout.zig"); + +pub const Plan = struct { + pool_bytes: u64, + planner_bytes: u64, +}; + +pub fn apply(text: []const u8, uniform_lines: usize, decoder_pools: usize, line_bytes: []const usize, min_lines: usize, planner_rows: ?u64, fixed_bytes: u64, startup_reservation: u64, counts: []usize) !Plan { + if (counts.len != line_bytes.len or decoder_pools > counts.len or min_lines == 0) return error.InvalidLineDeltaIndex; + if (uniform_lines < min_lines) return error.PoolTooSmall; + @memset(counts, 0); + var added: u64 = 0; + var removed: u64 = 0; + var entries = std.mem.splitScalar(u8, text, ','); + while (entries.next()) |entry| { + const colon = std.mem.indexOfScalar(u8, entry, ':') orelse return error.InvalidLineDeltas; + const index = try decimal(entry[0..colon]); + if (index >= decoder_pools) return error.InvalidLineDeltaIndex; + if (counts[index] != 0) return error.DuplicateLineDelta; + const delta = entry[colon + 1 ..]; + if (delta.len < 2 or (delta[0] != '+' and delta[0] != '-')) return error.InvalidLineDeltas; + const magnitude = try decimal(delta[1..]); + const result = if (delta[0] == '+') result: { + added = try std.math.add(u64, added, magnitude); + break :result try std.math.add(usize, uniform_lines, magnitude); + } else result: { + removed = try std.math.add(u64, removed, magnitude); + if (magnitude > uniform_lines) return error.PoolTooSmall; + break :result uniform_lines - magnitude; + }; + if (result < min_lines) return error.PoolTooSmall; + counts[index] = result; + } + if (added != removed) return error.NonZeroLineDeltas; + var before: u64 = 0; + var after: u64 = 0; + var expert_bytes: u64 = 0; + for (counts, line_bytes) |*count, bytes| { + if (count.* == 0) count.* = uniform_lines; + _ = try layout.PoolState.init(count.*, 0); + const cost = try std.math.add(u64, bytes, 8); + before = try std.math.add(u64, before, try std.math.mul(u64, uniform_lines, cost)); + after = try std.math.add(u64, after, try std.math.mul(u64, count.*, cost)); + expert_bytes = try std.math.add(u64, expert_bytes, try std.math.mul(u64, count.*, bytes)); + } + if (before != after) return error.LineDeltaByteBudgetChanged; + const planner_bytes = if (planner_rows) |rows| try memory.plannerBytesForPools(counts, rows) else 0; + const pool_bytes = try std.math.add(u64, try std.math.add(u64, expert_bytes, fixed_bytes), planner_bytes); + if (pool_bytes > startup_reservation) return error.InsufficientStartupMemory; + return .{ .pool_bytes = pool_bytes, .planner_bytes = planner_bytes }; +} + +fn decimal(text: []const u8) !usize { + if (text.len == 0) return error.InvalidLineDeltas; + for (text) |digit| if (digit < '0' or digit > '9') return error.InvalidLineDeltas; + return std.fmt.parseInt(usize, text, 10); +} + +test "exl3: line deltas parse signed decoder ordinals and leave MTP unchanged" { + var counts: [22]usize = undefined; + const bytes: [22]usize = @splat(256); + const plan = try apply("20:+1,15:-1", 32, 21, &bytes, 10, 16, 100, std.math.maxInt(u64), &counts); + for (counts, 0..) |count, index| try std.testing.expectEqual(@as(usize, if (index == 20) 33 else if (index == 15) 31 else 32), count); + try std.testing.expectEqual(try memory.plannerBytes(22, 32, 16), plan.planner_bytes); + try std.testing.expectEqual(22 * 32 * 256 + 100 + plan.planner_bytes, plan.pool_bytes); + try std.testing.expectEqual(@as(usize, 32), counts[21]); +} + +test "exl3: line deltas reject malformed empty duplicate and out of range maps" { + var counts: [3]usize = undefined; + const bytes = [_]usize{ 256, 256, 512 }; + for ([_][]const u8{ "", "0", "0:1", "0:+", ":+1", "0:++1", "0:+1,", ",0:+1", "0:+1,,1:-1", " 0:+1,1:-1", "0:+1, 1:-1", "0:+1 ,1:-1", "-1:+1", "+0:+1", "0x0:+1", "0:+1_0", "0:+1:0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltas, apply(text, 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,0:-1", "0:+0,0:-0", "00:+0,0:+0" }) |text| { + try std.testing.expectError(error.DuplicateLineDelta, apply(text, 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "2:+0", "3:+0", "99:+0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltaIndex, apply(text, 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.InvalidLineDeltaIndex, apply("0:+0", 32, 0, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas conserve lines and actual bytes independently" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 256, 1024 }; + for ([_][]const u8{ "0:+1", "0:-1", "0:+2,2:-1" }) |text| { + try std.testing.expectError(error.NonZeroLineDeltas, apply(text, 32, 3, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,1:-1", "0:-1,1:+1" }) |text| { + try std.testing.expectError(error.LineDeltaByteBudgetChanged, apply(text, 32, 3, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + const expected_counts = [_]usize{ 33, 32, 31, 32 }; + const planner = try memory.plannerBytes(4, 32, 16); + const expected_bytes = 32 * (256 + 512 + 256 + 1024) + planner + 100; + const plan = try apply("0:+1,2:-1", 32, 3, &bytes, 10, 16, 100, expected_bytes, &counts); + try std.testing.expectEqualSlices(usize, &expected_counts, &counts); + try std.testing.expectEqual(expected_bytes, plan.pool_bytes); + try std.testing.expectEqual(planner, plan.planner_bytes); + try std.testing.expectError(error.InsufficientStartupMemory, apply("0:+1,2:-1", 32, 3, &bytes, 10, 16, 100, expected_bytes - 1, &counts)); + const without_planner = try apply("0:+1,2:-1", 32, 3, &bytes, 10, null, 100, expected_bytes - planner, &counts); + try std.testing.expectEqual(@as(u64, 0), without_planner.planner_bytes); + try std.testing.expectEqual(expected_bytes - planner, without_planner.pool_bytes); +} + +test "exl3: line deltas enforce floor before allocation and reject arithmetic overflow" { + var counts: [2]usize = undefined; + const bytes = [_]usize{ 256, 256 }; + _ = try apply("0:+22,1:-22", 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts); + try std.testing.expectEqualSlices(usize, &.{ 54, 10 }, &counts); + for ([_][]const u8{ "0:+23,1:-23", "0:+32,1:-32", "0:+33,1:-33" }) |text| { + try std.testing.expectError(error.PoolTooSmall, apply(text, 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.PoolTooSmall, apply("0:+0", 9, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+18446744073709551615,1:-18446744073709551615", 32, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+1,1:-1", 32, 2, &.{ std.math.maxInt(usize), std.math.maxInt(usize) }, 10, 16, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas allow byte neutral multiway transfers within conservative startup sizing" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 768, 1024 }; + const fixed: u64 = 100; + const reservation = try memory.poolBytes(32, 1024 + 8, 4, fixed + try memory.plannerBytes(4, 0, 16)); + const plan = try apply("0:+1,1:-2,2:+1", 32, 3, &bytes, 10, 16, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 33, 30, 33, 32 }, &counts); + try std.testing.expect(plan.pool_bytes <= reservation); + try std.testing.expectEqual(@as(u64, 32 * (256 + 512 + 768 + 1024) + fixed) + try memory.plannerBytes(4, 32, 16), plan.pool_bytes); + _ = try apply("0:+0", 32, 3, &bytes, 10, 16, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 32, 32, 32, 32 }, &counts); +} diff --git a/zig/src/families/flashnext/exl3_moe.zig b/zig/src/families/flashnext/exl3_moe.zig new file mode 100644 index 0000000000..e3bd05d859 --- /dev/null +++ b/zig/src/families/flashnext/exl3_moe.zig @@ -0,0 +1,154 @@ +//! CPU f64 oracle for one Flash Next MoE block: router, routed experts, shared expert, combine. +//! Built from `exl3_format.dequantize` (the Zig decoder) and `exl3_cpu.swiglu`, with the model's own +//! routing maths, not `exl3_route.select`: softmax over all 512 routed logits, take the top `topk` +//! (ties to the lower id), renormalise by their sum (tensorfold `families/qwen4_exp/model_layers.py` +//! `SparseMoE.route`). That equals a softmax over only the selected logits, which `select` computes. +//! Output = sum_k w_k * expert_k(x) + sigmoid(shared_gate . x) * shared(x), with +//! expert(x) = down(silu(gate x) * (up x)). No CUDA. + +const std = @import("std"); +const format = @import("exl3_format.zig"); +const cpu = @import("exl3_cpu.zig"); + +/// One decoded EXL3 matrix W[k][n] in f64. `y[j] = sum_i x[i] * w[i * n + j]`. +pub const Dense = struct { + k: usize, + n: usize, + w: []f64, + + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16) !Dense { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (suh.len != k or svh.len != n or trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const w = try gpa.alloc(f64, k * n); + errdefer gpa.free(w); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, w, scratch); + return .{ .k = k, .n = n, .w = w }; + } + + pub fn deinit(d: *Dense, gpa: std.mem.Allocator) void { + gpa.free(d.w); + d.* = undefined; + } + + pub fn apply(d: Dense, x: []const f64, y: []f64) error{Length}!void { + if (x.len != d.k or y.len != d.n) return error.Length; + for (0..d.n) |j| { + var sum: f64 = 0; + for (0..d.k) |i| sum += x[i] * d.w[i * d.n + j]; + y[j] = sum; + } + } +}; + +/// down(silu(gate x) * up x). `gate_buf`, `up_buf`, `glu_buf` are `gate.n` long; `out` is `down.n` long. +pub fn expert(gate: Dense, up: Dense, down: Dense, x: []const f64, gate_buf: []f64, up_buf: []f64, glu_buf: []f64, out: []f64) error{Length}!void { + if (gate.n != up.n or down.k != gate.n or down.n != gate.k) return error.Length; + if (gate_buf.len != gate.n or up_buf.len != gate.n or glu_buf.len != gate.n) return error.Length; + try gate.apply(x, gate_buf); + try up.apply(x, up_buf); + try cpu.swiglu(gate_buf, up_buf, glu_buf, std.math.inf(f64)); + try down.apply(glu_buf, out); +} + +/// `logits` is `experts + 1` values, the last being the shared-gate dot. Writes `picks[0..topk]` +/// (descending probability, ties to the lower id), their renormalised `weights`, and the shared sigmoid weight. +pub fn routeFromLogits(logits: []const f64, topk: usize, picks: []u16, weights: []f64, shared_weight: *f64) error{Length}!void { + if (logits.len < 2 or topk == 0 or topk > logits.len - 1 or picks.len != topk or weights.len != topk) return error.Length; + const experts = logits.len - 1; + var top = logits[0]; + for (logits[0..experts]) |v| top = @max(top, v); + var total: f64 = 0; + for (logits[0..experts]) |v| total += @exp(v - top); + var taken: [1024]bool = @splat(false); + if (experts > taken.len) return error.Length; + var kept: f64 = 0; + for (0..topk) |k| { + var best: usize = 0; + var best_v = -std.math.inf(f64); + for (logits[0..experts], 0..) |v, i| { + if (!taken[i] and v > best_v) { + best_v = v; + best = i; + } + } + taken[best] = true; + picks[k] = @intCast(best); + weights[k] = @exp(best_v - top) / total; + kept += weights[k]; + } + for (weights) |*w| w.* /= kept; + shared_weight.* = 1.0 / (1.0 + @exp(-logits[experts])); +} + +/// Row-major `[rows][x.len]` f16 gate dotted with `x`, then the one-row shared gate. +pub fn logitsF16(gate: []const f16, rows: usize, shared_gate: []const f16, x: []const f64, logits: []f64) error{Length}!void { + if (gate.len != rows * x.len or shared_gate.len != x.len or logits.len != rows + 1) return error.Length; + for (0..rows) |r| { + var sum: f64 = 0; + for (x, 0..) |v, i| sum += v * @as(f64, gate[r * x.len + i]); + logits[r] = sum; + } + var s: f64 = 0; + for (x, 0..) |v, i| s += v * @as(f64, shared_gate[i]); + logits[rows] = s; +} + +test "full softmax then renormalise is the softmax over the selected logits" { + const route = @import("exl3_route.zig"); + var rng = std.Random.DefaultPrng.init(11); + const r = rng.random(); + for (0..20) |_| { + var logits: [65]f64 = undefined; + var logits32: [65]f32 = undefined; + for (&logits, &logits32) |*a, *b| { + const v: f32 = (r.float(f32) - 0.5) * 8; + a.* = v; + b.* = v; + } + var picks: [6]u16 = undefined; + var w: [6]f64 = undefined; + var sw: f64 = 0; + try routeFromLogits(&logits, 6, &picks, &w, &sw); + var picks32: [7]u16 = undefined; + var w32: [7]f32 = undefined; + var scratch: [64]f32 = undefined; + try route.select(&logits32, 6, &picks32, &w32, &scratch); + try std.testing.expectEqualSlices(u16, &picks, picks32[0..6]); + var sum: f64 = 0; + for (w, w32[0..6]) |a, b| { + sum += a; + // select rounds through bf16: 8 bits of mantissa, relative 2^-8 + try std.testing.expectApproxEqRel(a, @as(f64, b), 1.0 / 128.0); + } + try std.testing.expectApproxEqAbs(@as(f64, 1), sum, 1e-12); + try std.testing.expectApproxEqRel(sw, @as(f64, w32[6]), 1.0 / 128.0); + try std.testing.expectEqual(@as(u16, 64), picks32[6]); + } +} + +test "ties take the lower id and bad lengths store nothing" { + const logits = [_]f64{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9 }; + var w = [_]f64{ 0, 0 }; + var sw: f64 = -1; + try routeFromLogits(&logits, 2, &picks, &w, &sw); + try std.testing.expectEqualSlices(u16, &.{ 1, 2 }, &picks); + try std.testing.expectEqual(@as(f64, 0.5), w[0]); + try std.testing.expectEqual(@as(f64, 0.5), sw); + var one = [_]u16{7}; + try std.testing.expectError(error.Length, routeFromLogits(&logits, 2, &one, &w, &sw)); + try std.testing.expectEqual(@as(u16, 7), one[0]); +} + +test "logitsF16 is a row dot and the shared gate is the last logit" { + const gate = [_]f16{ 1, 0, 0, 2, 0, 0 }; + const shared = [_]f16{ 0, 0, 3 }; + var logits: [3]f64 = undefined; + try logitsF16(&gate, 2, &shared, &.{ 1, 5, 7 }, &logits); + try std.testing.expectEqualSlices(f64, &.{ 1, 2, 21 }, &logits); + try std.testing.expectError(error.Length, logitsF16(&gate, 2, shared[0..2], &.{ 1, 5, 7 }, &logits)); +} diff --git a/zig/src/families/flashnext/exl3_norm.zig b/zig/src/families/flashnext/exl3_norm.zig new file mode 100644 index 0000000000..f80b9f2878 --- /dev/null +++ b/zig/src/families/flashnext/exl3_norm.zig @@ -0,0 +1,94 @@ +//! Centred RMSNorm apply. The addend is `scan.centredOffset`, not a second median test. +//! No module epsilon: the caller passes `eps`. + +const std = @import("std"); +const scan = @import("exl3_scan.zig"); + +/// `rms = sqrt(mean(x²) + eps)`, then `out[i] = x[i] / rms * (weight[i] + offset)`. +/// `offset` is 1 for gamma-1 storage and 0 for gamma. A non-finite or negative `eps` +/// is `error.Length`: the error set has no separate eps refusal. Every refusal +/// returns before a store to `out`. +pub fn apply(x: []const f32, weight: []const f32, means: []const f32, eps: f32, out: []f32) error{ Length, NormAmbiguous }!void { + const offset = try scan.centredOffset(means); + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + if (x.len == 0 or x.len != weight.len or x.len != out.len) return error.Length; + + var sum_sq: f32 = 0; + for (x) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(x.len)) + eps); + for (x, weight, out) |xi, wi, *yi| yi.* = xi / rms * (wi + offset); +} + +test "a unit vector changes by exactly the weight difference" { + // Same means `centredOffset` classifies as gamma-1 (1) and gamma (0). + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + // x = [1], eps = 0: mean(x²)+eps is 1, so the RMS is 1 and out = weight + offset. + const x = [_]f32{1}; + const light = [_]f32{2}; + const heavy = [_]f32{5}; + var zero_light = [_]f32{0}; + var zero_heavy = [_]f32{0}; + var one_light = [_]f32{0}; + var one_heavy = [_]f32{0}; + try apply(&x, &light, &near_zero, 0, &zero_light); + try apply(&x, &heavy, &near_zero, 0, &zero_heavy); + try apply(&x, &light, &near_one, 0, &one_light); + try apply(&x, &heavy, &near_one, 0, &one_heavy); + try std.testing.expectEqual(@as(f32, 3), zero_light[0]); + try std.testing.expectEqual(@as(f32, 6), zero_heavy[0]); + try std.testing.expectEqual(@as(f32, 2), one_light[0]); + try std.testing.expectEqual(@as(f32, 5), one_heavy[0]); + try std.testing.expectEqual(heavy[0] - light[0], zero_heavy[0] - zero_light[0]); + try std.testing.expectEqual(heavy[0] - light[0], one_heavy[0] - one_light[0]); + + // Two ones still have mean 1, so a forgotten division by length cannot pass. + const pair_x = [_]f32{ 1, 1 }; + const pair_w = [_]f32{ 2, 7 }; + var pair_out = [_]f32{ 0, 0 }; + try apply(&pair_x, &pair_w, &near_zero, 0, &pair_out); + try std.testing.expectEqual(@as(f32, 3), pair_out[0]); + try std.testing.expectEqual(@as(f32, 8), pair_out[1]); +} + +test "an ambiguous norm leaves out untouched" { + var out = [_]f32{ 9, 8 }; + const many: [65]f32 = @splat(0); + try std.testing.expectError(error.NormAmbiguous, apply(&.{ 1, 1 }, &.{ 1, 1 }, &many, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + // Length also disagrees. The locator runs first, so this is not Length. + try std.testing.expectError(error.NormAmbiguous, apply(&.{1}, &.{}, &.{ 0.4, 0.9 }, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a length mismatch leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{ 9, 8 }; + try std.testing.expectError(error.Length, apply(&.{}, &.{}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{ 1, 1 }, &.{1}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{1}, &.{1}, &near_one, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a negative eps leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{9}; + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_one, std.math.nan(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, std.math.inf(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); +} diff --git a/zig/src/families/flashnext/exl3_offload.zig b/zig/src/families/flashnext/exl3_offload.zig new file mode 100644 index 0000000000..e5be85111f --- /dev/null +++ b/zig/src/families/flashnext/exl3_offload.zig @@ -0,0 +1,1649 @@ +//! Host expert-cache policy for Flash Next EXL3 offload (ported from exl3-flashnext-window-wip, unchanged rules). +//! Routed experts stay in host RAM. A VRAM line is a copy that is in flight until +//! `reap`, then resident. Eviction is LRU among resident lines; an in-flight line +//! is not reusable until it is reaped. `planRound` does not write either table. +//! Shared expert 512 is always a hit and never occupies a routed slot. +//! Generate addresses a line as `(layer, id)` through `scheduleUsedAt` and `reapAt`. +//! `schedule` and `onDevice` stay id-only and are not used on a layer-tagged table. +//! On upstream's `Experts` a miss is always copied (the device pool, `cuda_exl3_pool.zig`, has no CPU path): the round +//! planner (`exl3_round.zig`) stages every miss with `stageUsedAt` and keeps the copy budget for optional prefetch. + +const std = @import("std"); +/// Flash Next's routed geometry (was exl3_shape.zig): 48 decoder layers, 512 routed experts, the shared expert is id 512. +pub const shape = struct { + pub const layers: usize = 48; + pub const experts: usize = 512; + pub const shared_expert: usize = 512; +}; + +pub const shared_id: u16 = @intCast(shape.shared_expert); + +pub const Decision = enum { hit, copy, cpu }; + +pub const Policy = struct { + slots: u16, + expert_bytes: u64, + /// Bytes the copy engine may move this round. copy_count = expert_bytes==0 ? 0 : copy_bytes / expert_bytes. + copy_bytes: u64, + swap_every: u16 = 4, + swap_cap: u16 = 96, +}; + +pub const Cache = struct { + slots: []u16, // expert id, or empty_id + freq: []u32, // frequency of the resident id, parallel to slots + pub const empty_id: u16 = 0xffff; +}; + +/// Clear every slot to `empty_id` and every frequency to 0. Lengths must match. +pub fn init(slots: []u16, freq: []u32) Cache { + std.debug.assert(slots.len == freq.len); + @memset(slots, Cache.empty_id); + @memset(freq, 0); + return .{ .slots = slots, .freq = freq }; +} + +/// Slot index of a resident expert, or null. `empty_id` is a vacant line, not an expert. +pub fn find(c: Cache, id: u16) ?usize { + if (id == Cache.empty_id) return null; + return indexOf(c.slots, id); +} + +/// One decision per id. Does not write the slot table. +/// Unique misses sort by count descending, then id ascending. The first +/// `copy_bytes / expert_bytes` of those (0 when `expert_bytes` is 0) are `copy`; +/// the rest are `cpu`. Duplicate ids in `ids` share one decision. +pub fn planRound( + c: Cache, + ids: []const u16, + policy: Policy, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (ids) |id| { + if (id == shared_id) continue; + if (find(c, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + for (ids, decisions) |id, *decision| { + if (id == shared_id or find(c, id) != null) { + decision.* = .hit; + continue; + } + const rank = indexOf(scratch[0..n], id).?; + decision.* = if (rank < copies) .copy else .cpu; + } +} + +/// Install on swap rounds only (`round > 0` and `round % swap_every == 0`). +/// Residents named in `requested` are heated before victims are chosen. +/// Candidates are unique non-resident, non-shared ids, hottest first, at most +/// `swap_cap`. Empty lines are filled before any eviction. A line placed earlier +/// in this call is not a victim of a later candidate. Returns how many experts +/// were installed. +pub fn swap( + c: *Cache, + requested: []const u16, + round: u32, + policy: Policy, + scratch: []u16, +) error{Scratch}!u16 { + if (policy.swap_every == 0 or round == 0 or round % @as(u32, policy.swap_every) != 0) return 0; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (requested) |id| { + if (id == shared_id) continue; + if (find(c.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], requested, hotterFirst); + heatResidents(c, requested); + + const cap = @min(n, @as(usize, policy.swap_cap)); + var installed: usize = 0; + for (scratch[0..cap]) |id| { + const slot = indexOf(c.slots, Cache.empty_id) orelse victim(c.*, scratch[0..installed]) orelse break; + c.slots[slot] = id; + c.freq[slot] = appearances(requested, id); + installed += 1; + } + return @intCast(installed); +} + +/// Fill slots from a routed histogram of length 512. Hottest count, then lower id. +/// `freq[slot] = histogram[id]`. Id 512 is not a histogram entry and is not stored. +pub fn installProfile(c: *Cache, histogram: []const u32) void { + std.debug.assert(histogram.len == shape.experts); + std.debug.assert(c.slots.len == c.freq.len); + @memset(c.slots, Cache.empty_id); + @memset(c.freq, 0); + + const n = @min(c.slots.len, histogram.len); + var filled: usize = 0; + while (filled < n) : (filled += 1) { + var best: u16 = 0; + var best_n: u32 = 0; + var any = false; + for (histogram, 0..) |count, i| { + const id: u16 = @intCast(i); + if (indexOf(c.slots[0..filled], id) != null) continue; + // Scan is id-ascending, so the first id that reaches the max count wins ties. + if (!any or count > best_n) { + best = id; + best_n = count; + any = true; + } + } + if (!any) break; + c.slots[filled] = best; + c.freq[filled] = best_n; + } +} + +fn copyCount(policy: Policy, unique_misses: usize) usize { + if (policy.expert_bytes == 0 or unique_misses == 0) return 0; + const raw = policy.copy_bytes / policy.expert_bytes; + if (raw >= unique_misses) return unique_misses; + return @intCast(raw); +} + +fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| { + if (x == id) n +|= 1; + } + return n; +} + +/// Count descending, id ascending. `ctx` is the multiset the counts are taken from. +fn hotterFirst(ctx: []const u16, a: u16, b: u16) bool { + const ca = appearances(ctx, a); + const cb = appearances(ctx, b); + if (ca != cb) return ca > cb; + return a < b; +} + +fn indexOf(ids: []const u16, id: u16) ?usize { + for (ids, 0..) |x, i| { + if (x == id) return i; + } + return null; +} + +fn pushUnique(dst: []u16, n: *usize, id: u16) error{Scratch}!void { + if (indexOf(dst[0..n.*], id) != null) return; + if (n.* >= dst.len) return error.Scratch; + dst[n.*] = id; + n.* += 1; +} + +fn heatResidents(c: *Cache, requested: []const u16) void { + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + const n = appearances(requested, id); + if (n == 0) continue; + c.freq[i] = c.freq[i] +| n; + } +} + +/// Lowest frequency, skipping lines already placed in this call. +/// Equal frequency evicts the higher slot index. +fn victim(c: Cache, skip: []const u16) ?usize { + var best: ?usize = null; + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + if (indexOf(skip, id) != null) continue; + if (best) |b| { + if (c.freq[i] < c.freq[b] or (c.freq[i] == c.freq[b] and i > b)) best = i; + } else best = i; + } + return best; +} + +pub const Line = enum { empty, inflight, resident }; + +/// VRAM working set. Host RAM still holds every routed expert. +/// `stamp` is the last reap or touch time. A higher stamp is newer. +pub const Vram = struct { + slots: []u16, + line: []Line, + stamp: []u32, + /// Layer of `slots[i]`. Parallel to the other three slices. + layer: []u8, + + pub fn init(slots: []u16, line: []Line, stamp: []u32, layer: []u8) Vram { + std.debug.assert(slots.len == line.len and slots.len == stamp.len and slots.len == layer.len); + @memset(slots, Cache.empty_id); + @memset(line, .empty); + @memset(stamp, 0); + @memset(layer, 0); + return .{ .slots = slots, .line = line, .stamp = stamp, .layer = layer }; + } +}; + +/// True once `id` has been reaped into VRAM. +/// Id-only: the first matching line wins. Generate calls `residentAt`. +pub fn onDevice(v: Vram, id: u16) bool { + const i = vramIndex(v, id) orelse return false; + return v.line[i] == .resident; +} + +/// True when this layer's copy of `id` is resident. Another layer's copy does not count. +pub fn residentAt(v: Vram, layer: u8, id: u16) bool { + const i = vramIndexAt(v, layer, id) orelse return false; + return v.line[i] == .resident; +} + +/// A resident hit. Moves that line to `now` so LRU keeps it. +pub fn touch(v: *Vram, id: u16, now: u32) bool { + const i = vramIndex(v.*, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +/// Start a host-to-VRAM copy of `id`. An empty line is used first. +/// Otherwise the resident line with the oldest stamp is replaced. +/// An in-flight line is left alone. `error.Busy` means every line is still in flight. +pub fn stage(v: *Vram, id: u16) error{ Busy, Shared }!usize { + return stageKeeping(v, id, &.{}); +} + +/// Same as `stage`, except a resident whose id is in `keep` is not a victim. +/// `error.Busy` means every line is in flight or held by `keep`. +/// Id-only. A newly assigned line is layer 0. Do not call this on a multi-layer table. +pub fn stageKeeping(v: *Vram, id: u16, keep: []const u16) error{ Busy, Shared }!usize { + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndex(v.*, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramLruKeep(v.*, keep) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = 0; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// One step of the VRAM cache. Does not reap. +/// A resident id is touched and decided `hit`. An in-flight id is `cpu` until reap. +/// Unique misses, hottest first, are staged until the copy budget runs out or every +/// free line is gone. A staged miss is `copy` and stays off-device until `reap`. +/// Residents named in `ids` are not eviction victims. Shared id 512 is `hit`. +/// A zero `expert_bytes` stages nothing: every miss stays `cpu`. +pub fn schedule( + v: *Vram, + ids: []const u16, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (onDevice(v.*, id)) { + _ = touch(v, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (vramIndex(v.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageKeeping(v, id, ids) catch break; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} + +/// The copy of `(layer, id)` has landed. Another layer's resident or inflight +/// line is not this copy. Generate calls this, not `reap`. +pub fn reapAt(v: *Vram, layer: u8, id: u16, now: u32) error{NotInflight}!void { + const i = vramIndexAt(v.*, layer, id) orelse return error.NotInflight; + if (v.line[i] != .inflight) return error.NotInflight; + v.line[i] = .resident; + v.stamp[i] = now; +} + +/// Layer-0 wrapper. Does not promote another layer's inflight copy. +/// Old tests are layer 0. Generate calls `reapAt`. +pub fn reap(v: *Vram, id: u16, now: u32) error{NotInflight}!void { + return reapAt(v, 0, id, now); +} + +/// First occupied slot with this id, on any layer. That alias is what +/// `schedule` and `onDevice` still call. Generate uses `vramIndexAt`. +fn vramIndex(v: Vram, id: u16) ?usize { + if (id == Cache.empty_id) return null; + for (v.slots, 0..) |slot, i| { + if (slot == id and v.line[i] != .empty) return i; + } + return null; +} + +/// Slot whose layer and id both match. The id-only search aliases two layers. +pub fn vramIndexAt(v: Vram, layer: u8, id: u16) ?usize { + if (id == Cache.empty_id) return null; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + for (v.slots, v.layer, 0..) |slot, lay, i| { + if (slot == id and lay == layer and v.line[i] != .empty) return i; + } + return null; +} + +pub fn touchAt(v: *Vram, layer: u8, id: u16, now: u32) bool { + const i = vramIndexAt(v.*, layer, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +fn vramEmpty(v: Vram) ?usize { + for (v.line, 0..) |state, i| if (state == .empty) return i; + return null; +} + +/// Oldest resident stamp. Equal stamps take the higher slot index. +fn vramLru(v: Vram) ?usize { + return vramLruKeep(v, &.{}); +} + +/// Oldest resident whose id is not in `keep`. Equal stamps take the higher slot index. +fn vramLruKeep(v: Vram, keep: []const u16) ?usize { + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (indexOf(keep, v.slots[i]) != null) continue; + if (best) |b| { + if (v.stamp[i] < v.stamp[b] or (v.stamp[i] == v.stamp[b] and i > b)) best = i; + } else best = i; + } + return best; +} + +fn countOf(counts: []const u32, id: u16) u32 { + if (id >= counts.len) return 0; + return counts[id]; +} + +pub fn countAt(counts: *const [shape.layers][shape.experts]u32, layer: u8, id: u16) u32 { + if (layer >= shape.layers or id >= shape.experts) return 0; + return counts[layer][id]; +} + +/// Lowest `counts[layer][id]`, then oldest stamp, then the higher slot index. +/// A keep entry is `(keep_layer, id)`. The same id on another layer is not protected. +/// Equal counts are the stamp LRU, so a zeroed table matches `vramLruKeep`. +fn vramCold( + v: Vram, + counts: *const [shape.layers][shape.experts]u32, + keep_layer: u8, + keep_ids: []const u16, +) ?usize { + std.debug.assert(v.layer.len == v.slots.len); + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (v.layer[i] == keep_layer and indexOf(keep_ids, v.slots[i]) != null) continue; + if (best) |b| { + const ci = countAt(counts, v.layer[i], v.slots[i]); + const cb = countAt(counts, v.layer[b], v.slots[b]); + const colder = ci < cb or + (ci == cb and v.stamp[i] < v.stamp[b]) or + (ci == cb and v.stamp[i] == v.stamp[b] and i > b); + if (colder) best = i; + } else best = i; + } + return best; +} + +const UsedOrder = struct { + counts: []const u32, + ids: []const u16, +}; + +/// Lifetime count descending, then this round's appearances, then lower id. +fn usedFirst(ctx: UsedOrder, a: u16, b: u16) bool { + const ca = countOf(ctx.counts, a); + const cb = countOf(ctx.counts, b); + if (ca != cb) return ca > cb; + const aa = appearances(ctx.ids, a); + const ab = appearances(ctx.ids, b); + if (aa != ab) return aa > ab; + return a < b; +} + +/// Add one to each routed id in `ids`. Shared id 512 is not a histogram entry. +/// `counts` is one cell per routed expert (`shape.experts`). +pub fn observe(counts: []u32, ids: []const u16) error{Length}!void { + if (counts.len != shape.experts) return error.Length; + for (ids) |id| { + if (id >= shape.experts) continue; + counts[id] +|= 1; + } +} + +/// Expert ids this rank stores. Rank 0 is `0..256`, rank 1 is `256..512`. +pub const Owned = struct { lo: u16, hi: u16 }; + +pub fn allOwned() Owned { + return .{ .lo = 0, .hi = @intCast(shape.experts) }; +} + +fn inOwned(owned: Owned, id: u16) bool { + return id >= owned.lo and id < owned.hi; +} + +/// Fill empty VRAM lines with the hottest counted ids. No eviction. +/// Ties take the lower id. Count 0 and ids already in the table are skipped. +/// Staged lines stay in flight until `reap`. +pub fn seedHottest(v: *Vram, counts: []const u32, limit: u16) error{Length}!u16 { + return seedSpan(v, counts, limit, allOwned()); +} + +fn seedSpan(v: *Vram, counts: []const u32, limit: u16, owned: Owned) error{Length}!u16 { + if (counts.len != shape.experts) return error.Length; + var seeded: u16 = 0; + while (seeded < limit) { + if (vramEmpty(v.*) == null) break; + var best: ?u16 = null; + var best_n: u32 = 0; + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + const count = counts[id]; + if (count == 0) continue; + if (vramIndex(v.*, id) != null) continue; + if (best == null or count > best_n) { + best = id; + best_n = count; + } + } + if (best == null) break; + _ = stage(v, best.?) catch break; + seeded += 1; + } + return seeded; +} + +/// Use every empty line. Hottest owned ids go in first. Lines that are still +/// empty take the remaining owned ids, low id first, so the card is not left +/// idle while the histogram warms up. No eviction. The other rank's ids are skipped. +pub fn fillOwned(v: *Vram, counts: []const u32, owned: Owned) error{Length}!u16 { + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + var n: u16 = try seedSpan(v, counts, std.math.maxInt(u16), owned); + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + if (vramEmpty(v.*) == null) break; + if (vramIndex(v.*, id) != null) continue; + _ = stage(v, id) catch break; + n += 1; + } + return n; +} + +fn quotaAt(lines: u16, layer: usize) u16 { + const n_layers: u16 = @intCast(shape.layers); + const base = lines / n_layers; + if (layer < @as(usize, lines % n_layers)) return base + 1; + return base; +} + +fn placeOwned(v: *Vram, layer: u8, id: u16) bool { + // Shared 512 is not a routed line. Callers only pass an owned id. + std.debug.assert(id != shared_id and id != Cache.empty_id); + const slot = vramEmpty(v.*) orelse return false; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return true; +} + +/// Hottest owned `(layer, id)` that still has stripe room, one line at a time. +fn placeHottest( + v: *Vram, + lines: u16, + counts: []const [shape.experts]u32, + owned: Owned, + used: *[shape.layers]u16, +) u16 { + var placed: u16 = 0; + while (placed < lines and vramEmpty(v.*) != null) { + var found = false; + var best_layer: u8 = 0; + var best_id: u16 = 0; + var best_c: u32 = 0; + for (counts, 0..) |row, layer_i| { + if (used[layer_i] >= quotaAt(lines, layer_i)) continue; + var id = owned.lo; + while (id < owned.hi) : (id += 1) { + const c = row[id]; + if (c == 0) continue; + if (vramIndexAt(v.*, @intCast(layer_i), id) != null) continue; + const layer: u8 = @intCast(layer_i); + const better = !found or c > best_c or + (c == best_c and id < best_id) or + (c == best_c and id == best_id and layer < best_layer); + if (!better) continue; + found = true; + best_layer = layer; + best_id = id; + best_c = c; + } + } + if (!found) break; + if (!placeOwned(v, best_layer, best_id)) break; + used[best_layer] += 1; + placed += 1; + } + return placed; +} + +/// Id, then layer. An uncapped layer-major walk would park every line on layer 0. +fn fillStripe(v: *Vram, lines: u16, owned: Owned, used: *[shape.layers]u16, placed: u16) u16 { + var n = placed; + var id = owned.lo; + while (id < owned.hi and n < lines) : (id += 1) { + var layer: usize = 0; + while (layer < shape.layers and n < lines) : (layer += 1) { + if (used[layer] >= quotaAt(lines, layer)) continue; + if (vramIndexAt(v.*, @intCast(layer), id) != null) continue; + if (!placeOwned(v, @intCast(layer), id)) return n; + used[layer] += 1; + n += 1; + } + } + return n; +} + +/// Stage up to `lines` owned experts. `lines` is the caller's live count. +/// Each layer gets `lines / shape.layers`, and the remainder goes one each to the low layers. +/// A zero histogram fills the lowest owned id on each of those layers. +/// A warm histogram takes the hottest owned pair that still fits that stripe, then +/// fills any line still empty. The other rank's ids are not staged. +pub fn fillOwnedLayers(v: *Vram, lines: u16, counts: []const [shape.experts]u32, owned: Owned) error{Length}!u16 { + if (counts.len != shape.layers) return error.Length; + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + var used: [shape.layers]u16 = @splat(0); + const hot = placeHottest(v, lines, counts, owned, &used); + return fillStripe(v, lines, owned, &used, hot); +} + +/// `stageKeeping`, but the victim is the coldest lifetime count. +/// Layer 0 only: a flat histogram does not score another layer's copy. +pub fn stageUsed(v: *Vram, id: u16, keep: []const u16, counts: []const u32) error{ Busy, Shared, Length }!usize { + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return stageUsedAt(v, 0, id, keep, &table); +} + +pub fn stageUsedAt( + v: *Vram, + layer: u8, + id: u16, + keep_ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, +) error{ Busy, Shared, Length }!usize { + if (layer >= shape.layers) return error.Length; + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndexAt(v.*, layer, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramCold(v.*, counts, layer, keep_ids) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// `schedule`, ordered by lifetime counts from `observe`. +/// A high count keeps its line even when its stamp is older. +/// Residents named in `ids` are not victims. `schedule` itself is unchanged. +/// This is the layer-0 wrapper. It is not the multi-layer path: the flat +/// histogram is copied into layer 0, so another layer's copy scores 0. +pub fn scheduleUsed( + v: *Vram, + ids: []const u16, + counts: []const u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return scheduleUsedAt(v, 0, ids, &table, policy, now, decisions, scratch, owned); +} + +/// Generate path. Misses sort by `counts[layer]`. A victim scores +/// `counts[victim.layer][victim.id]`. Keep is `(layer, id)` for this call's ids. +/// An id outside `owned` stays `.cpu` and is not staged. Shared 512 is a hit +/// and does not occupy a line. A landed copy is promoted with `reapAt`. +pub fn scheduleUsedAt( + v: *Vram, + layer: u8, + ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (layer >= shape.layers) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (residentAt(v.*, layer, id)) { + _ = touchAt(v, layer, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (!inOwned(owned, id)) continue; + if (vramIndexAt(v.*, layer, id) != null) continue; + try pushUnique(scratch, &n, id); + } + const row: []const u32 = &counts[layer]; + std.mem.sort(u16, scratch[0..n], UsedOrder{ .counts = row, .ids = ids }, usedFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageUsedAt(v, layer, id, ids, counts) catch |err| switch (err) { + error.Length => return error.Length, + error.Busy, error.Shared => break, + }; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} + +/// Line index of `(layer, id)` in any state, or null. The device pool addresses its buffers by this index. +pub fn lineOf(v: Vram, layer: u8, id: u16) ?usize { + return vramIndexAt(v, layer, id); +} + +/// A host-to-device copy the policy asked for: expert `id` into pool line `line`. +pub const Copy = struct { id: u16, line: u32 }; + +/// Per pick, the pool line the hit kernel reads, or -1 when the kernel must skip the slot +/// (a `.copy` or `.cpu` decision, the shared expert, or a line that is not resident). +/// The kernel ignores a negative id, so the host writes that slot's row into `y` itself. +pub fn lineIds(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []i32) error{Length}!void { + if (decisions.len != ids.len or out.len != ids.len) return error.Length; + for (ids, decisions, out) |id, d, *o| { + o.* = -1; + if (d != .hit or id == shared_id) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .resident) continue; + o.* = @intCast(i); + } +} + +/// Unique `(id, line)` pairs the round staged: a `.copy` decision whose line is in flight. +/// Returns how many were written. `out` shorter than the unique count is `error.Length`. +pub fn stagedCopies(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []Copy) error{Length}!usize { + if (decisions.len != ids.len) return error.Length; + var n: usize = 0; + for (ids, decisions) |id, d| { + if (d != .copy) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .inflight) continue; + var seen = false; + for (out[0..n]) |c| { + if (c.id == id) seen = true; + } + if (seen) continue; + if (n >= out.len) return error.Length; + out[n] = .{ .id = id, .line = @intCast(i) }; + n += 1; + } + return n; +} + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "a staged copy is not on device until reap, and LRU drops the older resident" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 4)); + try std.testing.expectEqual(@as(usize, 1), try stage(&vram, 7)); + try std.testing.expect(!onDevice(vram, 4)); + try std.testing.expectError(error.NotInflight, reap(&vram, 9, 1)); + try reap(&vram, 4, 1); + try reap(&vram, 7, 2); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expect(touch(&vram, 7, 5)); + try std.testing.expect(onDevice(vram, 4)); + + // 4 was reaped at 1 and not touched again. 7 was touched at 5. 4 is the victim. + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 8)); + try std.testing.expect(!onDevice(vram, 8)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[1]); + try reap(&vram, 8, 6); + try std.testing.expect(onDevice(vram, 8)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 4)); +} + +test "an in-flight line is not reused, and a full in-flight set is busy" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + _ = try stage(&vram, 1); + _ = try stage(&vram, 2); + try std.testing.expectError(error.Busy, stage(&vram, 3)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 2), slots[1]); + try std.testing.expectError(error.Shared, stage(&vram, shared_id)); +} + +test "schedule touches residents, stages the hottest miss, and will not evict a hit" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var decisions: [4]Decision = undefined; + var scratch: [4]u16 = undefined; + const policy = benchPolicy(2, 100, 100); + + const first = [_]u16{ 7, shared_id, 4, 7 }; + try schedule(&vram, &first, policy, 3, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[3]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(!onDevice(vram, 7)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(Line.empty, line[1]); + + try reap(&vram, 7, 4); + const held = [_]u16{7}; + try schedule(&vram, &held, policy, 9, decisions[0..1], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(@as(u32, 9), stamp[0]); + + // Both lines are residents this step asked for, so the miss cannot take either. + _ = try stage(&vram, 3); + try reap(&vram, 3, 10); + const protected = [_]u16{ 7, 3, 8 }; + try schedule(&vram, &protected, benchPolicy(2, 100, 10_000), 11, decisions[0..3], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(onDevice(vram, 3)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); +} + +test "a zero copy budget stages nothing, and an older resident is the one that moves" { + var slots = [_]u16{ 4, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 5 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try schedule(&vram, &miss, benchPolicy(2, 0, 10_000), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(@as(u16, 4), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + + try schedule(&vram, &miss, benchPolicy(2, 100, 100), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "resident ids are hits and planRound does not change slots" { + var slots = [_]u16{ 4, 1, 9 }; + var freq = [_]u32{ 3, 8, 1 }; + const cache = init(&slots, &freq); + slots = .{ 4, 1, 9 }; + freq = .{ 3, 8, 1 }; + const before_slots = slots; + const before_freq = freq; + + const ids = [_]u16{ 4, shared_id, 1, 9, 4 }; + var decisions: [5]Decision = undefined; + try planRound(cache, &ids, benchPolicy(3, 100, 10_000), &decisions, &[_]u16{}); + + for (decisions) |d| try std.testing.expectEqual(Decision.hit, d); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); + try std.testing.expectEqual(@as(?usize, 1), find(cache, 1)); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); +} + +test "misses past the copy budget are cpu and slots stay put" { + var slots = [_]u16{42}; + var freq = [_]u32{6}; + const cache = init(&slots, &freq); + slots[0] = 42; + freq[0] = 6; + const before_slots = slots; + const before_freq = freq; + + // Counts: 5 x4, 1 x3, 8 x1, 3 x1. Budget 250/100 = 2, so 5 and 1 are copied. + const ids = [_]u16{ 8, 5, 1, 5, 3, 5, 1, 42, 5, 1 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [8]u16 = undefined; + try planRound(cache, &ids, benchPolicy(1, 100, 250), &decisions, &scratch); + + const want = [_]Decision{ .cpu, .copy, .copy, .copy, .cpu, .copy, .copy, .hit, .copy, .copy }; + try std.testing.expectEqualSlices(Decision, &want, &decisions); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); +} + +test "duplicate miss ids share one decision" { + var slots: [1]u16 = undefined; + var freq: [1]u32 = undefined; + const cache = init(&slots, &freq); + const ids = [_]u16{ 6, 2, 6 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + // One copy. 6 appears twice, 2 once, so both 6s are copy and 2 is cpu. + try planRound(cache, &ids, benchPolicy(1, 100, 100), &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); +} + +test "equal miss counts break ties toward the lower id" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + // 10 x5, 6 x4, 2 x4. One copy would be 10; two copies are 10 and 2, not 6. + const ids = [_]u16{ 6, 10, 2, 6, 10, 2, 10, 6, 2, 10, 6, 10, 2 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 50, 100), &decisions, &scratch); + for (ids, decisions) |id, d| { + const want: Decision = if (id == 10 or id == 2) .copy else .cpu; + try std.testing.expectEqual(want, d); + } +} + +test "shared id is a hit and is never stored" { + var slots: [4]u16 = @splat(Cache.empty_id); + var freq: [4]u32 = @splat(0); + var cache = init(&slots, &freq); + slots[0] = 7; + freq[0] = 2; + + const ids = [_]u16{ shared_id, 4, shared_id }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(4, 100, 0), &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.hit, decisions[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + + const requested = [_]u16{ shared_id, 4, shared_id, 9 }; + const installed = try swap(&cache, &requested, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), installed); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); + + var hist: [shape.experts]u32 = @splat(0); + hist[4] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); +} + +test "round 2 installs nothing and round 4 installs min of cap and missing" { + var slots: [8]u16 = @splat(Cache.empty_id); + var freq: [8]u32 = @splat(0); + var cache = init(&slots, &freq); + const requested = [_]u16{ 3, 1, 3, 8 }; + var scratch: [8]u16 = undefined; + var policy = benchPolicy(8, 1, 0); + policy.swap_every = 4; + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 0, policy, &scratch)); + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 2, policy, &scratch)); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[0]); + try std.testing.expectEqual(@as(u32, 0), freq[0]); + + // Three unique misses, under the cap. + const n = try swap(&cache, &requested, 4, policy, &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual(@as(u16, 3), slots[0]); + try std.testing.expectEqual(@as(u32, 2), freq[0]); + try std.testing.expectEqual(@as(u16, 1), slots[1]); + try std.testing.expectEqual(@as(u16, 8), slots[2]); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[3]); +} + +test "a cold line is replaced by the missing id" { + var slots = [_]u16{41}; + var freq = [_]u32{0}; + var cache = init(&slots, &freq); + slots[0] = 41; + freq[0] = 0; + var scratch: [2]u16 = undefined; + const requested = [_]u16{9}; + const n = try swap(&cache, &requested, 4, benchPolicy(1, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(?usize, null), find(cache, 41)); +} + +test "eviction ties take the later slot" { + var slots = [_]u16{ 30, 10, 20 }; + var freq = [_]u32{ 4, 4, 4 }; + var cache = init(&slots, &freq); + slots = .{ 30, 10, 20 }; + freq = .{ 4, 4, 4 }; + var scratch: [2]u16 = undefined; + const n = try swap(&cache, &[_]u16{7}, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + // Higher id would have dropped 30. Higher slot index drops 20. + try std.testing.expectEqual([_]u16{ 30, 10, 7 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[2]); +} + +test "a resident named this round is heated before victims are chosen" { + var slots = [_]u16{ 2, 1 }; + var freq = [_]u32{ 0, 0 }; + var cache = init(&slots, &freq); + slots = .{ 2, 1 }; + freq = .{ 0, 0 }; + var scratch: [2]u16 = undefined; + // Without the heat, both freqs stay 0 and the later slot (id 1) would be the victim. + const n = try swap(&cache, &[_]u16{ 1, 99 }, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual([_]u16{ 99, 1 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "resident frequency adds this round's request count" { + var slots = [_]u16{ 5, Cache.empty_id }; + var freq = [_]u32{ 4, 0 }; + var cache = init(&slots, &freq); + slots = .{ 5, Cache.empty_id }; + freq = .{ 4, 0 }; + var scratch: [2]u16 = undefined; + const requested = [_]u16{ 5, 5, 5, 4 }; + const n = try swap(&cache, &requested, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 5), slots[0]); + try std.testing.expectEqual(@as(u32, 7), freq[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "empty lines are filled before a cold resident is evicted" { + var slots = [_]u16{ 30, Cache.empty_id, 40 }; + var freq = [_]u32{ 5, 0, 1 }; + var cache = init(&slots, &freq); + slots = .{ 30, Cache.empty_id, 40 }; + freq = .{ 5, 0, 1 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 7, 8 }, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual([_]u16{ 30, 7, 8 }, slots); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} + +test "a candidate placed this call is not evicted by a later one" { + var slots = [_]u16{ 1, 2, 3, 4 }; + var freq = [_]u32{ 1, 50, 2, 40 }; + var cache = init(&slots, &freq); + slots = .{ 1, 2, 3, 4 }; + freq = .{ 1, 50, 2, 40 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 10, 11, 12 }, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual([_]u16{ 10, 2, 11, 12 }, slots); + try std.testing.expectEqual(@as(u32, 50), freq[1]); +} + +test "installProfile keeps the hottest routed ids" { + var slots: [3]u16 = @splat(shared_id); + var freq = [_]u32{ 9, 9, 9 }; + var cache = init(&slots, &freq); + var hist: [shape.experts]u32 = @splat(0); + hist[10] = 1; + hist[4] = 7; + hist[8] = 7; + hist[2] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual([_]u16{ 4, 8, 2 }, slots); + try std.testing.expectEqual([_]u32{ 7, 7, 3 }, freq); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + + var wide_slots: [513]u16 = undefined; + var wide_freq: [513]u32 = undefined; + var wide = init(&wide_slots, &wide_freq); + var flat: [shape.experts]u32 = @splat(0); + flat[0] = 1; + installProfile(&wide, &flat); + for (wide_slots[0..shape.experts], 0..) |id, i| { + try std.testing.expectEqual(@as(u16, @intCast(i)), id); + } + try std.testing.expectEqual(Cache.empty_id, wide_slots[512]); + try std.testing.expectEqual(@as(u32, 1), wide_freq[0]); + try std.testing.expectEqual(@as(u32, 0), wide_freq[1]); + try std.testing.expectEqual(@as(?usize, null), find(wide, shared_id)); +} + +test "swap installs at most the cap even when slots and misses are larger" { + var slots: [200]u16 = @splat(Cache.empty_id); + var freq: [200]u32 = @splat(0); + var cache = init(&slots, &freq); + var requested: [200]u16 = undefined; + for (&requested, 0..) |*id, i| id.* = @intCast(i); + var scratch: [200]u16 = undefined; + const n = try swap(&cache, &requested, 4, benchPolicy(200, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 96), n); + for (slots[0..96], 0..) |id, i| try std.testing.expectEqual(@as(u16, @intCast(i)), id); + try std.testing.expectEqual(Cache.empty_id, slots[96]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 0), freq[96]); +} + +test "a high count keeps an old stamp and a one-off is the victim" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[1] = 50; + counts[4] = 1; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "zero lifetime counts fall back to stamp LRU" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var counts: [shape.experts]u32 = @splat(0); + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); +} + +test "lifetime count beats this round's appearances" { + var slots = [_]u16{Cache.empty_id}; + var line = [_]Line{.empty}; + var stamp = [_]u32{0}; + var layer = [_]u8{0}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[8] = 5; + const ids = [_]u16{ 3, 3, 3, 3, 8 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(1, 100, 100), 1, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[4]); + try std.testing.expectEqual(@as(u16, 8), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); +} + +test "a resident named this step is kept even when its count is low" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[4] = 1; + const ids = [_]u16{ 1, 4, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 9)); +} + +test "observed counts seed empty lines and a later round keeps the hot id" { + var counts: [shape.experts]u32 = @splat(0); + try observe(&counts, &.{ 1, shared_id, 1, 3 }); + try std.testing.expectEqual(@as(u32, 2), counts[1]); + try std.testing.expectEqual(@as(u32, 1), counts[3]); + try observe(&counts, &.{1}); + try std.testing.expectEqual(@as(u32, 3), counts[1]); + var short: [1]u32 = .{9}; + try std.testing.expectError(error.Length, observe(&short, &.{1})); + try std.testing.expectEqual(@as(u32, 9), short[0]); + + var seeded_counts: [shape.experts]u32 = @splat(0); + seeded_counts[3] = 10; + seeded_counts[1] = 10; + seeded_counts[8] = 2; + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + const n = try seedHottest(&vram, &seeded_counts, 2); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[1]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!onDevice(vram, 1)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); + try reap(&vram, 1, 1); + try std.testing.expect(onDevice(vram, 1)); + + var wide: [3]u16 = .{ Cache.empty_id, Cache.empty_id, Cache.empty_id }; + var wide_line: [3]Line = .{ .empty, .empty, .empty }; + var wide_stamp: [3]u32 = .{ 0, 0, 0 }; + var wide_layer: [3]u8 = .{ 0, 0, 0 }; + var wide_vram = Vram.init(&wide, &wide_line, &wide_stamp, &wide_layer); + try std.testing.expectError(error.Length, seedHottest(&wide_vram, counts[0..3], 1)); + try std.testing.expectEqual(Line.empty, wide_line[0]); +} + +test "a rank fills its own lines and does not copy the other rank's expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[400] = 100; + counts[7] = 3; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const n = try fillOwned(&vram, &counts, rank0); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(Line.inflight, line[0]); + try reap(&vram, 7, 1); + try reap(&vram, 0, 1); + + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + const asked = [_]u16{ 400, 9 }; + try scheduleUsed(&vram, &asked, &counts, benchPolicy(2, 100, 10_000), 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expect(vramIndex(vram, 9) != null); +} + +test "an empty slot table misses every routed id and installs nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + var cache = init(&slots, &freq); + const ids = [_]u16{ shared_id, 4, 4, 8 }; + var decisions: [4]Decision = undefined; + var scratch: [2]u16 = undefined; + const policy = benchPolicy(0, 10, 10); + try planRound(cache, &ids, policy, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(Decision.cpu, decisions[3]); + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &ids, 4, policy, &scratch)); + try std.testing.expectEqual(@as(usize, 0), slots.len); +} + +test "zero expert bytes copies nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + const ids = [_]u16{ 1, 2 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 0, 1000), &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); +} + +test "a short scratch or a short decision buffer is refused without writes" { + var slots = [_]u16{1}; + var freq = [_]u32{5}; + var cache = init(&slots, &freq); + slots[0] = 1; + freq[0] = 5; + + var decisions: [1]Decision = undefined; + try std.testing.expectError(error.Length, planRound(cache, &[_]u16{ 1, 2 }, benchPolicy(1, 1, 1), &decisions, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var one: [2]Decision = undefined; + try std.testing.expectError(error.Scratch, planRound(cache, &[_]u16{ 2, 3 }, benchPolicy(1, 1, 1), &one, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + + var tiny: [1]u16 = undefined; + try std.testing.expectError(error.Scratch, swap(&cache, &[_]u16{ 2, 3 }, 4, benchPolicy(1, 1, 0), &tiny)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var off = benchPolicy(1, 1, 0); + off.swap_every = 0; + var scratch: [2]u16 = undefined; + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &[_]u16{2}, 4, off, &scratch)); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} + +test "residentAt distinguishes the same expert id on two layers" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 2 }; + var layer = [_]u8{ 0, 40 }; + const vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 2 }; + layer = .{ 0, 40 }; + + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(!residentAt(vram, 1, 7)); + try std.testing.expect(!residentAt(vram, 0, 8)); + // Id-only lookup still aliases. Generate does not call it on this table. + try std.testing.expect(onDevice(vram, 7)); + + line[1] = .inflight; + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(!residentAt(vram, 40, 7)); +} + +test "reapAt promotes one inflight layer and leaves the other inflight" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .inflight, .inflight }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .inflight, .inflight }; + stamp = .{ 0, 0 }; + layer = .{ 0, 40 }; + + try reapAt(&vram, 0, 7, 3); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // The other layer's copy is still inflight, then promotable on its own reap. + try reapAt(&vram, 40, 7, 4); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(@as(u32, 4), stamp[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // Layer 0 is already resident. The wrapper must not treat that as the + // layer-40 inflight copy, and must not move layer 40. + line[1] = .inflight; + stamp[1] = 0; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 9)); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + try reapAt(&vram, 40, 7, 9); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 9), stamp[1]); + + // An inflight copy that exists only on layer 40 is not promoted by `reap`. + slots = .{ 7, Cache.empty_id }; + line = .{ .inflight, .empty }; + stamp = .{ 0, 0 }; + layer = .{ 40, 0 }; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 1)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try reapAt(&vram, 40, 7, 2); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u8, 40), layer[0]); +} + +test "stageKeeping tags a newly assigned line as layer 0" { + var slots = [_]u16{4}; + var line = [_]Line{.resident}; + var stamp = [_]u32{1}; + var layer = [_]u8{40}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{4}; + line = .{.resident}; + stamp = .{1}; + layer = .{40}; + + try std.testing.expectEqual(@as(usize, 0), try stageKeeping(&vram, 9, &.{})); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 40, 9)); + try std.testing.expect(slots[0] != shared_id); +} + +test "a keep on layer 0 does not protect that id on layer 40" { + var slots = [_]u16{ 7, 7, 3 }; + var line = [_]Line{ .resident, .resident, .resident }; + var stamp = [_]u32{ 4, 4, 4 }; + var layer = [_]u8{ 0, 40, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7, 3 }; + line = .{ .resident, .resident, .resident }; + stamp = .{ 4, 4, 4 }; + layer = .{ 0, 40, 0 }; + + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][7] = 10; + counts[40][7] = 1; + counts[0][3] = 8; + const ids = [_]u16{ 7, 9 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsedAt(&vram, 0, &ids, &counts, benchPolicy(3, 100, 100), 8, &decisions, &scratch, allOwned()); + + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 0, 3)); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(u32, 8), stamp[0]); + try std.testing.expectEqual(@as(u32, 4), stamp[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[2]); +} + +test "scheduleUsed is the layer-0 wrapper and is not the multi-layer path" { + var slots = [_]u16{ 1, 4 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 50 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + + // The flat vector says id 4 is hot. The wrapper stores that on layer 0 only, + // so the layer-40 copy scores 0 and is the victim. + var flat: [shape.experts]u32 = @splat(0); + flat[1] = 1; + flat[4] = 100; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &flat, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(@as(u8, 0), layer[1]); + try std.testing.expect(!residentAt(vram, 40, 4)); + + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][1] = 1; + counts[40][4] = 100; + try scheduleUsedAt(&vram, 0, &miss, &counts, benchPolicy(2, 100, 100), 7, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u8, 40), layer[1]); + try std.testing.expect(residentAt(vram, 40, 4)); + try std.testing.expect(!residentAt(vram, 0, 1)); +} + +test "scheduleUsedAt does not stage another rank or the shared expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[3][400] = 50; + counts[3][9] = 1; + const ids = [_]u16{ 400, shared_id, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const policy = benchPolicy(2, 100, 10_000); + try std.testing.expectError(error.Length, scheduleUsedAt(&vram, @intCast(shape.layers), &ids, &counts, policy, 4, &decisions, &scratch, rank0)); + try std.testing.expectEqual(Line.empty, line[0]); + + try scheduleUsedAt(&vram, 3, &ids, &counts, policy, 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, shared_id)); + try std.testing.expect(!residentAt(vram, 3, 9)); + try std.testing.expect(vramIndexAt(vram, 3, 9) != null); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 3), layer[0]); + try std.testing.expectEqual(Line.empty, line[1]); + try std.testing.expect(slots[0] != shared_id and slots[1] != shared_id); +} + +test "fillOwnedLayers stripes an explicit line count and skips the other rank" { + const wide: u16 = @intCast(shape.layers + 2); + var slots: [shape.layers + 2]u16 = @splat(Cache.empty_id); + var line: [shape.layers + 2]Line = @splat(.empty); + var stamp: [shape.layers + 2]u32 = @splat(0); + var layer: [shape.layers + 2]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const rank0 = Owned{ .lo = 0, .hi = 256 }; + var short: [2][shape.experts]u32 = @splat(@splat(0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &short, rank0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &counts, .{ .lo = 4, .hi = 2 })); + try std.testing.expectEqual(Line.empty, line[0]); + + const n = try fillOwnedLayers(&vram, wide, &counts, rank0); + try std.testing.expectEqual(wide, n); + // Id, then layer: slot 1 is layer 1's lowest id, not layer 0's second id. + try std.testing.expectEqual(@as(u16, 0), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(u8, 1), layer[1]); + try std.testing.expectEqual(@as(u16, 0), slots[shape.layers - 1]); + try std.testing.expectEqual(@as(u8, @intCast(shape.layers - 1)), layer[shape.layers - 1]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers]); + try std.testing.expectEqual(@as(u8, 0), layer[shape.layers]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers + 1]); + try std.testing.expectEqual(@as(u8, 1), layer[shape.layers + 1]); + var on_layer0: usize = 0; + for (slots, layer) |id, lay| { + try std.testing.expect(id < 256); + try std.testing.expect(id != shared_id); + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 2), on_layer0); + + // Fewer lines than layers: the remainder is layers 0..3, one lowest id each. + var few_slots: [4]u16 = @splat(Cache.empty_id); + var few_line: [4]Line = @splat(.empty); + var few_stamp: [4]u32 = @splat(0); + var few_layer: [4]u8 = @splat(0); + var few = Vram.init(&few_slots, &few_line, &few_stamp, &few_layer); + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&few, 4, &counts, rank0)); + try std.testing.expectEqual([_]u16{ 0, 0, 0, 0 }, few_slots); + try std.testing.expectEqual([_]u8{ 0, 1, 2, 3 }, few_layer); + + // Rank 1. A hotter id outside Owned is not staged. Layer 0 keeps its hottest owned id. + var rank_slots: [4]u16 = @splat(Cache.empty_id); + var rank_line: [4]Line = @splat(.empty); + var rank_stamp: [4]u32 = @splat(0); + var rank_layer: [4]u8 = @splat(0); + var rank = Vram.init(&rank_slots, &rank_line, &rank_stamp, &rank_layer); + counts[0][1] = 999; + counts[0][300] = 5; + const rank1 = Owned{ .lo = 256, .hi = 512 }; + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&rank, 4, &counts, rank1)); + try std.testing.expect(vramIndexAt(rank, 0, 300) != null); + try std.testing.expect(vramIndexAt(rank, 1, 256) != null); + try std.testing.expect(vramIndexAt(rank, 2, 256) != null); + try std.testing.expect(vramIndexAt(rank, 3, 256) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndex(rank, 1)); + for (rank_slots) |id| { + try std.testing.expect(id >= 256 and id < 512); + try std.testing.expect(id != shared_id); + } +} + +test "fillOwnedLayers keeps a hot id inside the stripe and does not give its layer a second line" { + const lines: u16 = @intCast(shape.layers); + var slots: [shape.layers]u16 = @splat(Cache.empty_id); + var line: [shape.layers]Line = @splat(.empty); + var stamp: [shape.layers]u32 = @splat(0); + var layer: [shape.layers]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][9] = 10; + counts[0][8] = 9; + const n = try fillOwnedLayers(&vram, lines, &counts, .{ .lo = 0, .hi = 256 }); + try std.testing.expectEqual(lines, n); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 0, 8)); + try std.testing.expect(vramIndexAt(vram, 1, 0) != null); + try std.testing.expect(vramIndexAt(vram, 40, 0) != null); + var on_layer0: usize = 0; + for (layer) |lay| { + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 1), on_layer0); +} + +test "lineIds maps hits to their line and everything else to -1, stagedCopies lists each staged line once" { + var slots: [3]u16 = undefined; + var line: [3]Line = undefined; + var stamp: [3]u32 = undefined; + var layer: [3]u8 = undefined; + var v = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const pol = Policy{ .slots = 3, .expert_bytes = 10, .copy_bytes = 20 }; + var ids = [_]u16{ 7, 9, 9, 11, 512 }; + var dec: [5]Decision = undefined; + var scratch: [8]u16 = undefined; + try observe(&counts[2], &ids); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 1, &dec, &scratch, allOwned()); + // 3 unique misses (9 twice is hottest), budget for two copies: 9 and 7 copy, 11 cpu, 512 hit. + try std.testing.expectEqual(Decision.copy, dec[0]); + try std.testing.expectEqual(Decision.copy, dec[1]); + try std.testing.expectEqual(Decision.copy, dec[2]); + try std.testing.expectEqual(Decision.cpu, dec[3]); + try std.testing.expectEqual(Decision.hit, dec[4]); + var copies: [4]Copy = undefined; + const n = try stagedCopies(v, 2, &ids, &dec, &copies); + try std.testing.expectEqual(@as(usize, 2), n); + var ln: [5]i32 = undefined; + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, -1, -1, -1 }, &ln); + try std.testing.expectError(error.Length, stagedCopies(v, 2, &ids, &dec, copies[0..1])); + for (copies[0..n]) |c| try reapAt(&v, 2, c.id, 2); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 3, &dec, &scratch, allOwned()); + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 7).?)), ln[0]); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 9).?)), ln[1]); + try std.testing.expectEqual(ln[1], ln[2]); + try std.testing.expect(ln[3] == -1 and ln[4] == -1); + try std.testing.expectEqual(@as(?usize, null), lineOf(v, 3, 7)); +} diff --git a/zig/src/families/flashnext/exl3_place.zig b/zig/src/families/flashnext/exl3_place.zig new file mode 100644 index 0000000000..e10f51c664 --- /dev/null +++ b/zig/src/families/flashnext/exl3_place.zig @@ -0,0 +1,332 @@ +//! Header-only home for one tensor name. No safetensors open, no download, no allocation. +//! First match wins: mixer patch, n-gram row, vision, shared expert, routed bank, centred norm, resident set. +//! Routed ids are `0 .. shape.experts`. The shared expert is the next id and is not this bank. + +const std = @import("std"); +const scan = @import("exl3_scan.zig"); +const shape = @import("exl3_shape.zig"); +const cache = @import("exl3_cache.zig"); +const tp = @import("exl3_tp.zig"); + +pub const Home = enum { + ngram_host, // PLE rows, page cache or SSD, never a VRAM alloc + routed_host, // routed expert bank, host RAM, cache policy decides GPU copies + shared_resident, // shared expert, replicated, stays on device + resident, // dense, router, HC, MTP head, unquantized fp16 linears + centred_norm, // add centredOffset in fp32 at load + mtp_patch, // the mixer file kept outside the index + vision_skip, // vision encoder, not this milestone + plain, +}; + +pub const Tensor = struct { + name: []const u8, + rank: usize, + /// Filename of the shard, used only for the MTP mixer patch. + file: []const u8 = "", +}; + +const centred_names = [_][]const u8{ + "hc_norm", + "q_norm", + "k_norm", + "q_layernorm", + "k_layernorm", + "norm_key", + "norm_query", + "norm_conv", + "pre_fc_norm_embedding", + "pre_fc_norm_hidden", +}; + +const resident_names = [_][]const u8{ + "lm_head", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "out_proj", + "router", + "hyper_connection", + "in_proj", + "conv1d", + "A_log", + "dt_bias", + "indexer", + "mtp", + // Router logits and the PLE pair. `k_proj` is not a substring of `key_proj`. + "mlp.gate", + "key_proj", + "value_proj", +}; + +comptime { + std.debug.assert(shape.experts == 512); + std.debug.assert(shape.shared_expert == shape.experts); +} + +pub fn home(t: Tensor) Home { + if (scan.isMtpMixerPatch(t.file)) return .mtp_patch; + + const kind = scan.classifyName(t.name, t.rank); + if (kind == .ngram_row) return .ngram_host; + if (isVision(t.name)) return .vision_skip; + if (contains(t.name, "shared_expert")) return .shared_resident; + if (isRouted(t.name)) return .routed_host; + + // Gamma is a rank-1 `.plain` weight. An EXL3 part of that module is not the + // centred tensor, and rule 7 must not take it: `embed` sits inside + // `pre_fc_norm_embedding`, and `mtp` sits beside `pre_fc_norm_*`. + if (isCentred(t.name)) { + if (kind == .exl3_part) return .plain; + return .centred_norm; + } + + // `mtp` pins the MTP head. `mtp.fc_embedding` is the host embed remainder + // (trellis, suh, svh, mul1), not that head and not `pre_fc_norm_embedding`. + if (contains(t.name, "mtp.fc_embedding")) return .plain; + + if ((kind == .exl3_part or kind == .plain) and isResident(t.name)) return .resident; + return .plain; +} + +/// True when `which` is the cache's host place and both host banks refuse a resident copy. +/// `placeWeights` is called with a zero budget; routed experts and the n-gram table are still host. +pub fn hostBank(which: cache.Place) bool { + const budget = cache.Budget{ + .device_bytes = 0, + .resident_bytes = 0, + .kv_bytes = 0, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + }; + const host = cache.refusesResident(.routed) and cache.refusesResident(.ngram) and cache.placeWeights(budget) == .host; + return host and which == .host; +} + +fn contains(hay: []const u8, needle: []const u8) bool { + return std.mem.indexOf(u8, hay, needle) != null; +} + +fn isVision(name: []const u8) bool { + return contains(name, "visual") or std.mem.startsWith(u8, name, "vision") or contains(name, ".vision."); +} + +/// `.experts.` / `experts.` gate, up, and down for ids 0..511. Any `shared` substring stays out. +fn isRouted(name: []const u8) bool { + if (contains(name, "shared")) return false; + return contains(name, ".experts.") or contains(name, "experts."); +} + +/// Token is a full dot-component, or the name's suffix. +pub fn isCentred(name: []const u8) bool { + for (centred_names) |token| { + if (std.mem.endsWith(u8, name, token)) { + if (name.len == token.len or name[name.len - token.len - 1] == '.') return true; + } + var it = std.mem.splitScalar(u8, name, '.'); + while (it.next()) |part| if (std.mem.eql(u8, part, token)) return true; + } + return false; +} + +fn isResident(name: []const u8) bool { + for (resident_names) |token| { + if (!contains(name, token)) continue; + // `mlp.switch_mlp.gate_proj` contains `mlp.gate` and is not the router logits. + if (std.mem.eql(u8, token, "mlp.gate") and contains(name, "gate_proj")) continue; + return true; + } + return false; +} + +test "a tensor name picks one home and the host banks stay off device" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Home.ngram_host, home(.{ .name = ngram, .rank = 2 })); + + const wide = home(.{ .name = ngram, .rank = 3 }); + try std.testing.expect(wide != .ngram_host); + try std.testing.expectEqual(Home.plain, wide); + + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.shared_resident, home(.{ + .name = "model.layers.0.mlp.shared_expert.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.attn_hyper_connection.hc_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "language_model.mtp.pre_fc_norm_embedding.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.layers.0.self_attn.q_norm.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.mtp_patch, home(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + })); + try std.testing.expectEqual(Home.vision_skip, home(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.resident, home(.{ + .name = "model.layers.0.linear_attn.out_proj.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "readme", .rank = 0 })); + try std.testing.expect(hostBank(.host)); +} + +/// Where the bytes for this rank live. A cut can still split a pinned tensor. +/// `vram_pinned` is a GEMM or a bias on the decode path of every token. +/// The embedding table is `host`: one row is 5 KiB, the table is about 1.2 GiB. +/// `mtp.fc_embedding` stays with that table. The `mtp` token does not pin it. +/// `vram_learned` is a routed expert this rank may stage. The other rank's ids are not. +pub const Stay = enum { vram_pinned, vram_learned, host, disk, other_rank, skip }; + +/// Integer after `.experts.`, or null when the name is not a routed expert. +pub fn expertId(name: []const u8) ?u16 { + const key = ".experts."; + const at = std.mem.indexOf(u8, name, key) orelse return null; + var i = at + key.len; + if (i >= name.len or name[i] < '0' or name[i] > '9') return null; + var n: u32 = 0; + while (i < name.len and name[i] >= '0' and name[i] <= '9') : (i += 1) { + n = n * 10 + (name[i] - '0'); + if (n > std.math.maxInt(u16)) return null; + } + if (i >= name.len or name[i] != '.') return null; + return @intCast(n); +} + +pub fn stay(which: Home, id: ?u16, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Stay { + if (ranks != 1 and ranks != 2) return error.IllegalCut; + if (rank >= ranks) return error.RankOutOfRange; + switch (which) { + .ngram_host => return .disk, + .vision_skip => return .skip, + .routed_host => { + const expert = id orelse return error.IllegalCut; + const span = try tp.range(.routed_set, rank, ranks); + if (@as(usize, expert) >= span.lo and @as(usize, expert) < span.hi) return .vram_learned; + return .other_rank; + }, + .plain => return .host, + .shared_resident, .resident, .centred_norm, .mtp_patch => return .vram_pinned, + } +} + +pub fn stayOf(t: Tensor, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Stay { + return stay(home(t), expertId(t.name), rank, ranks); +} + +test "tp 2 pins the decode path, leaves the n-gram on disk, and learns only this rank's experts" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 }, 1, 2)); + + const pinned = [_][]const u8{ + "model.layers.0.mlp.shared_expert.up_proj.trellis", + "lm_head.weight", + "model.layers.0.mlp.router.weight", + "model.layers.0.mlp_hyper_connection.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.linear_attn.in_proj_qkv.weight", + "model.layers.0.linear_attn.out_proj.weight", + "model.layers.0.linear_attn.conv1d.weight", + "model.layers.0.linear_attn.A_log", + "model.layers.0.linear_attn.dt_bias", + "model.language_model.layers.0.mlp.gate.weight", + "model.language_model.layers.1.ple.key_proj.weight", + "model.language_model.layers.1.ple.value_proj.weight", + }; + try std.testing.expectEqual(Stay.host, try stayOf(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + }, 0, 2)); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "readme", .rank = 0 }, 1, 2)); + for (pinned) |name| { + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 1, 2)); + } + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + }, 0, 2)); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + }, 1, 2)); + try std.testing.expectEqual(Stay.skip, try stayOf(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + }, 0, 2)); + + try std.testing.expectEqual(@as(?u16, 0), expertId("model.layers.0.mlp.experts.0.up_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 255), expertId("model.layers.0.mlp.experts.255.down_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 511), expertId("model.layers.0.mlp.experts.511.gate_proj.trellis")); + try std.testing.expectEqual(@as(?u16, null), expertId("model.layers.0.mlp.shared_expert.up_proj.trellis")); + + const low = "model.layers.0.mlp.experts.0.up_proj.trellis"; + const high = "model.layers.0.mlp.experts.400.up_proj.trellis"; + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = low, .rank = 3 }, 0, 2)); + try std.testing.expectEqual(Stay.other_rank, try stayOf(.{ .name = low, .rank = 3 }, 1, 2)); + try std.testing.expectEqual(Stay.other_rank, try stayOf(.{ .name = high, .rank = 3 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 }, 1, 2)); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 }, 0, 1)); + + try std.testing.expectError(error.IllegalCut, stayOf(.{ .name = low, .rank = 3 }, 0, 4)); + try std.testing.expectError(error.RankOutOfRange, stayOf(.{ .name = low, .rank = 3 }, 2, 2)); + try std.testing.expectError(error.IllegalCut, stay(.routed_host, null, 0, 2)); + try std.testing.expect(std.mem.indexOf(u8, "key_proj", "k_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "value_proj", "v_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "mlp.shared_expert_gate", "mlp.gate") == null); +} + +test "mlp.gate, key_proj, and value_proj stay pinned, and so does mlp.router.weight" { + const gate = "model.language_model.layers.0.mlp.gate.weight"; + const key = "model.language_model.layers.1.ple.key_proj.weight"; + const value = "model.language_model.layers.1.ple.value_proj.weight"; + const router = "model.layers.0.mlp.router.weight"; + const shared_gate = "model.layers.0.mlp.shared_expert_gate.weight"; + for ([_][]const u8{ gate, key, value, router }) |name| { + try std.testing.expectEqual(Home.resident, home(.{ .name = name, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 1, 2)); + } + // `shared_expert` matches before the resident tokens. This is not the 2560×640 expert. + try std.testing.expectEqual(Home.shared_resident, home(.{ .name = shared_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = shared_gate, .rank = 1 }, 0, 2)); + + const switch_gate = "model.layers.0.mlp.switch_mlp.gate_proj.weight"; + try std.testing.expectEqual(Home.plain, home(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = switch_gate, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.gate_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "mtp.fc_embedding.weight", .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Home.resident, home(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 }, 0, 2)); +} diff --git a/zig/src/families/flashnext/exl3_prefill_window.zig b/zig/src/families/flashnext/exl3_prefill_window.zig new file mode 100644 index 0000000000..23ea4e3a8c --- /dev/null +++ b/zig/src/families/flashnext/exl3_prefill_window.zig @@ -0,0 +1,104 @@ +const std = @import("std"); + +pub const max_rows = 4096; +pub const graph_entries = 112; + +pub fn graphBytes(exec: u64, template: u64, transient: u64) !u64 { + const retained = try std.math.mul(u64, graph_entries * 2, try std.math.add(u64, exec, template)); + return std.math.add(u64, retained, try std.math.mul(u64, 2, transient)); +} + +pub fn enabled(value: ?[]const u8) bool { + return if (value) |text| std.mem.eql(u8, text, "1") else false; +} + +pub fn admits(free: u64, protected: u64, extra: u64) bool { + const required = std.math.add(u64, protected, extra) catch return false; + return extra != 0 and free >= required; +} + +pub const Segment = struct { offset: usize, rows: usize }; +pub const Step = struct { layer: usize, chunk: usize }; + +pub const Plan = struct { + segments: [2]Segment, + count: usize, + rows: usize, + + pub fn init(remaining: usize, chunk_rows: usize, alignment: usize) !Plan { + if (remaining == 0 or alignment == 0 or chunk_rows == 0 or chunk_rows > max_rows or chunk_rows % alignment != 0) return error.InvalidPrefillWindow; + const first = @min(remaining, chunk_rows); + const second = @min(remaining - first, chunk_rows); + return .{ .segments = .{ .{ .offset = 0, .rows = first }, .{ .offset = first, .rows = second } }, .count = if (second == 0) 1 else 2, .rows = first + second }; + } + + pub fn iterator(plan: Plan, layers: usize) Iterator { + return .{ .count = plan.count, .layers = layers }; + } +}; + +pub const Iterator = struct { + count: usize, + layers: usize, + layer: usize = 0, + chunk: usize = 0, + + pub fn next(iterator: *Iterator) ?Step { + if (iterator.layer >= iterator.layers) return null; + const step: Step = .{ .layer = iterator.layer, .chunk = iterator.chunk }; + iterator.chunk += 1; + if (iterator.chunk == iterator.count) { + iterator.chunk = 0; + iterator.layer += 1; + } + return step; + } +}; + +test "exl3: two chunk windows require explicit opt-in and bounded aligned rows" { + for ([_]?[]const u8{ null, "", "0", "true", "2" }) |value| try std.testing.expect(!enabled(value)); + try std.testing.expect(enabled("1")); + for ([_]usize{ 0, 1, 255, 4097, 8192 }) |rows| try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(16384, rows, 256)); + try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(0, 4096, 256)); + try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(4096, 4096, 0)); +} + +test "exl3: two chunk scheduling keeps each layer chronological and retains chunk boundaries" { + for ([_]usize{ 1, 255, 256, 4095, 4096, 4097, 8192, 8193, 16384, 16641 }) |length| { + var position: [4]usize = @splat(0); + var cursor: usize = 0; + var windows: usize = 0; + while (cursor < length) { + const plan = try Plan.init(length - cursor, 4096, 256); + try std.testing.expect(plan.rows <= 8192); + var steps = plan.iterator(position.len); + var previous: ?Step = null; + while (steps.next()) |step| { + const segment = plan.segments[step.chunk]; + try std.testing.expectEqual(position[step.layer], cursor + segment.offset); + if (previous) |before| { + if (before.layer == step.layer) try std.testing.expectEqual(before.chunk + 1, step.chunk) else { + try std.testing.expectEqual(before.layer + 1, step.layer); + try std.testing.expectEqual(@as(usize, 0), step.chunk); + } + } + position[step.layer] += segment.rows; + previous = step; + } + cursor += plan.rows; + windows += 1; + } + for (position) |end| try std.testing.expectEqual(length, end); + try std.testing.expectEqual((length + 8191) / 8192, windows); + } +} + +test "exl3: two chunk admission protects all existing reservations and rejects overflow" { + try std.testing.expect(admits(100, 60, 40)); + try std.testing.expect(!admits(99, 60, 40)); + try std.testing.expect(!admits(100, 101, 1)); + try std.testing.expect(!admits(100, 0, 0)); + try std.testing.expect(!admits(std.math.maxInt(u64), std.math.maxInt(u64), 1)); + try std.testing.expectEqual(@as(u64, 288 << 20), try graphBytes(1 << 20, 0, 32 << 20)); + try std.testing.expectError(error.Overflow, graphBytes(std.math.maxInt(u64), 0, 0)); +} diff --git a/zig/src/families/flashnext/exl3_profile_line_budget.zig b/zig/src/families/flashnext/exl3_profile_line_budget.zig new file mode 100644 index 0000000000..955c48471e --- /dev/null +++ b/zig/src/families/flashnext/exl3_profile_line_budget.zig @@ -0,0 +1,468 @@ +const std = @import("std"); + +pub const heuristic = "Greedy marginal histogram coverage per line byte, not real miss-cost optimization. Equal ratios prefer fewer lines, then the lower decoder ordinal. Exact line completion uses bounded search, potentially exponential with many distinct costs. Candidates must also preserve exact device picks bytes; otherwise sizing stays uniform."; +pub const experts: usize = 512; + +pub const Request = struct { + counts: []const []const u32, + rank: usize, + world: usize, + line_bytes: []const usize, + uniform_lines: usize, + min_lines: usize, + max_rows: usize, + slots: usize, +}; + +pub const Budget = struct { + baseline_bytes: u64, + allocated_bytes: u64, + uniform_fallback: bool = false, +}; + +pub const Error = std.mem.Allocator.Error || error{ InvalidDimensions, InvalidRank, InvalidLineCost, InvalidBounds, Overflow, UnfillableBudget }; + +fn poolBytes(line_bytes: usize, lines: usize, max_rows: usize, slots: usize, min_lines: usize) Error!u64 { + if (lines == 0) return error.InvalidBounds; + if (min_lines == 0) return error.InvalidBounds; + if (max_rows == 0 or slots == 0) return error.InvalidDimensions; + const picks = try std.math.mul(usize, max_rows, slots); + const max_passes = 2 + experts / min_lines + @intFromBool(experts % min_lines != 0); + const passes = @min(max_passes, @min(picks, experts) + 1); + const picks_bytes = try std.math.mul(usize, try std.math.mul(usize, passes, picks), @sizeOf(i32)); + return std.math.add(u64, try std.math.mul(u64, lines, line_bytes), picks_bytes); +} + +const Group = struct { + bytes: u64, + available: usize, + + fn descending(_: void, left: Group, right: Group) bool { + return left.bytes > right.bytes; + } +}; + +pub fn allocate(gpa: std.mem.Allocator, request: Request, output: []usize) Error!Budget { + if ((request.world != 1 and request.world != 2) or request.rank >= request.world) return error.InvalidRank; + const local_experts = experts / request.world; + if (request.min_lines == 0 or request.min_lines > request.uniform_lines or request.uniform_lines > local_experts) return error.InvalidBounds; + if (request.counts.len == 0 or request.counts.len != request.line_bytes.len or output.len != request.counts.len) return error.InvalidDimensions; + var baseline_bytes: u64 = 0; + var line_budget_bytes: u64 = 0; + var floor_bytes: u64 = 0; + for (request.counts, request.line_bytes) |histogram, bytes| { + if (histogram.len != experts) return error.InvalidDimensions; + if (bytes == 0) return error.InvalidLineCost; + baseline_bytes = try std.math.add(u64, baseline_bytes, try poolBytes(bytes, request.uniform_lines, request.max_rows, request.slots, request.min_lines)); + line_budget_bytes = try std.math.add(u64, line_budget_bytes, try std.math.mul(u64, request.uniform_lines, bytes)); + floor_bytes = try std.math.add(u64, floor_bytes, try std.math.mul(u64, request.min_lines, bytes)); + } + const ranked = try gpa.alloc([experts]u32, request.counts.len); + defer gpa.free(ranked); + const lines = try gpa.alloc(usize, request.counts.len); + defer gpa.free(lines); + const storage = try gpa.alloc(Group, request.counts.len); + defer gpa.free(storage); + const layer_groups = try gpa.alloc(usize, request.counts.len); + defer gpa.free(layer_groups); + const attempted = try gpa.alloc(bool, request.counts.len); + defer gpa.free(attempted); + const suffix_length = try std.math.add(usize, request.counts.len, 1); + const suffix_bytes = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_bytes); + const suffix_gcd = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_gcd); + const first = request.rank * local_experts; + for (request.counts, ranked) |histogram, *ordered| { + @memcpy(ordered[0..local_experts], histogram[first..][0..local_experts]); + std.mem.sort(u32, ordered[0..local_experts], {}, std.sort.desc(u32)); + } + @memset(lines, request.min_lines); + var group_count: usize = 0; + for (request.line_bytes) |bytes| { + var found: ?usize = null; + for (storage[0..group_count], 0..) |group, index| { + if (group.bytes == bytes) { + found = index; + break; + } + } + const group_index = found orelse new: { + storage[group_count] = .{ .bytes = bytes, .available = 0 }; + group_count += 1; + break :new group_count - 1; + }; + storage[group_index].available = try std.math.add(usize, storage[group_index].available, local_experts - request.min_lines); + } + const groups = storage[0..group_count]; + std.mem.sort(Group, groups, {}, Group.descending); + for (request.line_bytes, layer_groups) |bytes, *group_index| { + for (groups, 0..) |group, index| { + if (group.bytes == bytes) { + group_index.* = index; + break; + } + } + } + var remaining = line_budget_bytes - floor_bytes; + while (remaining != 0) { + @memset(attempted, false); + var selected: ?usize = null; + for (0..lines.len) |_| { + var best: ?usize = null; + for (lines, request.line_bytes, 0..) |count, bytes, layer| { + if (attempted[layer] or count == local_experts or bytes > remaining) continue; + if (best == null or better(layer, best.?, ranked, lines, request.line_bytes)) best = layer; + } + const candidate = best orelse break; + attempted[candidate] = true; + const group_index = layer_groups[candidate]; + groups[group_index].available -= 1; + const possible = canFill(groups, remaining - request.line_bytes[candidate], suffix_bytes, suffix_gcd); + groups[group_index].available += 1; + if (possible) { + selected = candidate; + break; + } + } + const layer = selected orelse return error.UnfillableBudget; + lines[layer] += 1; + groups[layer_groups[layer]].available -= 1; + remaining -= request.line_bytes[layer]; + } + var allocated_bytes: u64 = 0; + for (lines, request.line_bytes) |count, bytes| allocated_bytes = try std.math.add(u64, allocated_bytes, try poolBytes(bytes, count, request.max_rows, request.slots, request.min_lines)); + const uniform_fallback = allocated_bytes != baseline_bytes; + if (uniform_fallback) { + @memset(lines, request.uniform_lines); + allocated_bytes = baseline_bytes; + } + @memcpy(output, lines); + return .{ .baseline_bytes = baseline_bytes, .allocated_bytes = allocated_bytes, .uniform_fallback = uniform_fallback }; +} + +fn better(left: usize, right: usize, ranked: []const [experts]u32, lines: []const usize, bytes: []const usize) bool { + const left_score = @as(u128, ranked[left][lines[left]]) * bytes[right]; + const right_score = @as(u128, ranked[right][lines[right]]) * bytes[left]; + if (left_score != right_score) return left_score > right_score; + if (lines[left] != lines[right]) return lines[left] < lines[right]; + return left < right; +} + +fn gcd(left: u64, right: u64) u64 { + var numerator = left; + var denominator = right; + while (denominator != 0) { + const remainder = numerator % denominator; + numerator = denominator; + denominator = remainder; + } + return numerator; +} + +fn canFill(groups: []const Group, target: u64, suffix_bytes: []u64, suffix_gcd: []u64) bool { + suffix_bytes[groups.len] = 0; + suffix_gcd[groups.len] = 0; + var index = groups.len; + while (index != 0) { + index -= 1; + const group = groups[index]; + const capacity = @as(u128, group.bytes) * group.available + suffix_bytes[index + 1]; + suffix_bytes[index] = @intCast(@min(capacity, target)); + suffix_gcd[index] = if (group.available == 0) suffix_gcd[index + 1] else gcd(group.bytes, suffix_gcd[index + 1]); + } + return fill(groups, 0, target, suffix_bytes, suffix_gcd); +} + +fn fill(groups: []const Group, index: usize, target: u64, suffix_bytes: []const u64, suffix_gcd: []const u64) bool { + if (target == 0) return true; + if (index == groups.len or target > suffix_bytes[index] or suffix_gcd[index] == 0 or target % suffix_gcd[index] != 0) return false; + const group = groups[index]; + if (group.available == 0) return fill(groups, index + 1, target, suffix_bytes, suffix_gcd); + const max_take = @min(group.available, target / group.bytes); + const required = target - @min(target, suffix_bytes[index + 1]); + const min_take = required / group.bytes + @intFromBool(required % group.bytes != 0); + if (min_take > max_take) return false; + var take = max_take; + while (true) { + if (fill(groups, index + 1, target - take * group.bytes, suffix_bytes, suffix_gcd)) return true; + if (take == min_take) break; + take -= 1; + } + return false; +} + +fn testRequest(counts: []const []const u32, bytes: []const usize, uniform: usize, minimum: usize) Request { + return .{ .counts = counts, .rank = 0, .world = 2, .line_bytes = bytes, .uniform_lines = uniform, .min_lines = minimum, .max_rows = 1, .slots = 1 }; +} + +test "profile budget: sorted marginal counts favor hot layers with exact bytes" { + var hot: [experts]u32 = @splat(0); + const cold: [experts]u32 = @splat(1); + @memset(hot[0..6], 100); + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 30 + 2 * 8), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + hot[0] = 0; + hot[255] = 100; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); +} + +test "profile budget: rank filtering handles asymmetric counts and TP1" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + @memset(first[0..4], 10); + @memset(second[257..260], 20); + second[511] = 20; + var output: [2]usize = undefined; + var request = testRequest(&.{ &first, &second }, &.{ 8, 8 }, 2, 1); + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 1 }, &output); + request.rank = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); + request.rank = 0; + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); +} + +test "profile budget: zero and flat histograms retain uniform sizing deterministically" { + for ([_]u32{ 0, 7, std.math.maxInt(u32) }) |value| { + const histogram: [experts]u32 = @splat(value); + var output: [3]usize = undefined; + for (0..3) |_| { + _ = try allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram, &histogram }, &.{ 16, 16, 16 }, 17, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 17, 17, 17 }, &output); + } + } +} + +test "profile budget: tied ratios prefer fewer lines then lower decoder ordinal" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + const third: [experts]u32 = @splat(0); + @memset(first[0..3], 10); + @memset(second[0..3], 10); + var output: [3]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &.{ 8, 8, 8 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 2, 1 }, &output); +} + +test "profile budget: ratios use wide products and avoid greedy dead ends" { + const hot: [experts]u32 = @splat(std.math.maxInt(u32) - 1); + const cold: [experts]u32 = @splat(std.math.maxInt(u32)); + const unit: usize = @as(usize, 1) << 32; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 2 * unit, 3 * unit }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 15) * unit + 2 * 8, plan.allocated_bytes); + const low: [experts]u32 = @splat(1); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &low }, &.{ 4, 6 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 2, 2 }, &output); +} + +test "profile budget: floors caps and single layer budgets" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(0); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 200, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 256, 144 }, &output); + var request = testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 400, 10); + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 512, 288 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 7, 11 }, 10, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 10, 10 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{&hot}, &.{13}, 47, 10), output[0..1]); + try std.testing.expectEqual(@as(usize, 47), output[0]); +} + +test "profile budget: malformed inputs and overflow leave output unchanged" { + const histogram: [experts]u32 = @splat(0); + const shared: [513]u32 = @splat(0); + var output = [_]usize{ 999, 999 }; + const valid = testRequest(&.{&histogram}, &.{8}, 20, 10); + for ([_][]const []const u32{ &.{}, &.{histogram[0..511]}, &.{&shared} }) |counts| { + var request = valid; + request.counts = counts; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + } + var request = valid; + request.line_bytes = &.{}; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, valid, &output)); + request = valid; + request.line_bytes = &.{0}; + try std.testing.expectError(error.InvalidLineCost, allocate(std.testing.allocator, request, output[0..1])); + for ([_][2]usize{ .{ 0, 0 }, .{ 3, 0 }, .{ 2, 2 }, .{ 1, 1 }, .{ 2, std.math.maxInt(usize) } }) |identity| { + request = valid; + request.world = identity[0]; + request.rank = identity[1]; + try std.testing.expectError(error.InvalidRank, allocate(std.testing.allocator, request, output[0..1])); + } + for ([_][2]usize{ .{ 20, 0 }, .{ 9, 10 }, .{ 257, 10 }, .{ 0, 0 } }) |bounds| { + request = valid; + request.uniform_lines = bounds[0]; + request.min_lines = bounds[1]; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, output[0..1])); + } + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{&histogram}, &.{std.math.maxInt(usize)}, 2, 1), output[0..1])); + const bytes = std.math.maxInt(usize) / 2 + 1; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram }, &.{ bytes, bytes }, 1, 1), &output)); + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); +} + +fn allocationFailureProbe(gpa: std.mem.Allocator) !void { + const histogram: [experts]u32 = @splat(7); + var output = [_]usize{ 999, 999 }; + const result = allocate(gpa, testRequest(&.{ &histogram, &histogram }, &.{ 4, 6 }, 3, 1), &output) catch |err| { + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); + return err; + }; + try std.testing.expectEqual(result.baseline_bytes, result.allocated_bytes); +} + +test "profile budget: allocation failures release scratch and preserve output" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, allocationFailureProbe, .{}); +} + +test "profile budget: feasibility matches exhaustive bounded enumeration" { + const groups = [_]Group{ .{ .bytes = 11, .available = 1 }, .{ .bytes = 7, .available = 2 }, .{ .bytes = 5, .available = 3 }, .{ .bytes = 3, .available = 4 } }; + var suffix_bytes: [5]u64 = undefined; + var suffix_gcd: [5]u64 = undefined; + for (0..65) |target| { + var expected = false; + for (0..2) |first| { + for (0..3) |second| { + for (0..4) |third| { + for (0..5) |fourth| { + if (11 * first + 7 * second + 5 * third + 3 * fourth == target) expected = true; + } + } + } + } + try std.testing.expectEqual(expected, canFill(&groups, target, &suffix_bytes, &suffix_gcd)); + } + try std.testing.expect(!canFill(&.{ .{ .bytes = 6, .available = 10 }, .{ .bytes = 4, .available = 10 } }, 2, &suffix_bytes, &suffix_gcd)); +} + +test "profile budget: heterogeneous allocations preserve the feasible uniform witness" { + const first: [experts]u32 = @splat(17); + const second: [experts]u32 = @splat(11); + const third: [experts]u32 = @splat(5); + var output: [3]usize = undefined; + for (1..6) |first_cost| { + for (1..6) |second_cost| { + for (1..6) |third_cost| { + const bytes = [_]usize{ first_cost, second_cost, third_cost }; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &bytes, 2, 1), &output); + try std.testing.expectEqual(@as(u64, 2 * (first_cost + second_cost + third_cost) + 3 * 8), plan.allocated_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + var observed: u64 = 0; + for (output, bytes) |count, cost| { + try std.testing.expect(count >= 1 and count <= 256); + observed += try poolBytes(cost, count, 1, 1, 1); + } + try std.testing.expectEqual(plan.allocated_bytes, observed); + } + } + } +} + +test "profile budget: device picks capacity stays uniform across pool sizes" { + const line_bytes = (1 << 20) + 8; + for ([_]usize{ 10, 11, 31, 32, 33, 64, 256, 512 }) |lines| { + const actual = try poolBytes(line_bytes, lines, 64, 11, 10); + try std.testing.expectEqual(@as(u64, lines * line_bytes + 54 * 64 * 11 * 4), actual); + try std.testing.expectEqual(@as(usize, 54), @divExact(@as(usize, @intCast(actual - lines * line_bytes)), 64 * 11 * 4)); + } + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 10, 1, 11, 10)); + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 2 * 4), try poolBytes(line_bytes, 10, 1, 1, 10)); + try std.testing.expectEqual(@as(u64, 256 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 256, 1, 11, 10)); +} + +test "profile budget: redistribution preserves expert planner and picks bytes exactly" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..55], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 54, 53); + request.max_rows = 64; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 55, 53 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 108 * line_bytes + 2 * 12 * 64 * 11 * 4), plan.baseline_bytes); + var observed: u64 = 0; + for (output) |lines| { + try std.testing.expect(lines >= request.min_lines); + observed += try poolBytes(line_bytes, lines, request.max_rows, request.slots, request.min_lines); + } + try std.testing.expectEqual(plan.baseline_bytes, observed); + try std.testing.expectEqual(observed, plan.allocated_bytes); +} + +test "profile budget: fixed picks capacity preserves exact bytes across pass thresholds" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..33], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 32, 31); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + request.max_rows = 64; + request.slots = 11; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 64 * line_bytes + 2 * 19 * 64 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: fewer picks bytes also retain the exact uniform budget" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(1); + var request = testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1); + request.max_rows = 16; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 30 + 2 * 177 * 16 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: TP2 line count above the local expert bank is rejected for caller fallback" { + const histogram: [experts]u32 = @splat(1); + var request = testRequest(&.{&histogram}, &.{8}, 257, 10); + var output: [1]usize = .{ 257 }; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 257), output[0]); + request.uniform_lines = 256; + _ = try allocate(std.testing.allocator, request, &output); +} + +test "profile budget: picks dimensions and byte overflow preserve output" { + const histogram: [experts]u32 = @splat(0); + var request = testRequest(&.{&histogram}, &.{8}, 20, 10); + var output = [_]usize{999}; + request.max_rows = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.max_rows = 16; + request.slots = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.slots = 11; + request.max_rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + request.slots = 1; + request.max_rows = std.math.maxInt(usize) / 4; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 999), output[0]); +} diff --git a/zig/src/families/flashnext/exl3_recent_demand.zig b/zig/src/families/flashnext/exl3_recent_demand.zig new file mode 100644 index 0000000000..a836408b20 --- /dev/null +++ b/zig/src/families/flashnext/exl3_recent_demand.zig @@ -0,0 +1,212 @@ +const std = @import("std"); + +pub const expert_count: usize = 512; +pub const Count = u32; +const Row = [expert_count / @bitSizeOf(u64)]u64; + +pub const RecentDemand = struct { + history: []Row, + counts: [expert_count]Count = @splat(0), + next_row: usize = 0, + retained_rows: usize = 0, + + pub fn init(allocator: std.mem.Allocator, window_rows: usize) !RecentDemand { + if (window_rows == 0 or window_rows > std.math.maxInt(Count)) return error.InvalidWindowSize; + const history = try allocator.alloc(Row, window_rows); + @memset(history, @splat(0)); + return .{ .history = history }; + } + + pub fn deinit(self: *RecentDemand, allocator: std.mem.Allocator) void { + allocator.free(self.history); + self.* = undefined; + } + + pub fn reset(self: *RecentDemand) void { + @memset(self.history, @splat(0)); + self.counts = @splat(0); + self.next_row = 0; + self.retained_rows = 0; + } + + pub fn windowRows(self: *const RecentDemand) usize { + return self.history.len; + } + + pub fn rowCount(self: *const RecentDemand) usize { + return self.retained_rows; + } + + pub fn scores(self: *const RecentDemand) *const [expert_count]Count { + return &self.counts; + } + + pub fn score(self: *const RecentDemand, expert_id: i32) Count { + if (expert_id < 0 or expert_id >= expert_count) return 0; + return self.counts[@intCast(expert_id)]; + } + + pub fn observeRow(self: *RecentDemand, picks: []const i32) void { + var incoming: Row = @splat(0); + for (picks) |expert_id| { + if (expert_id < 0 or expert_id >= expert_count) continue; + const expert: usize = @intCast(expert_id); + incoming[expert / @bitSizeOf(u64)] |= @as(u64, 1) << @as(u6, @intCast(expert % @bitSizeOf(u64))); + } + + const outgoing = self.history[self.next_row]; + for (outgoing, incoming, 0..) |expired_word, incoming_word, word_index| { + var removed = expired_word & ~incoming_word; + while (removed != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(removed); + self.counts[expert] -= 1; + removed &= removed - 1; + } + var added = incoming_word & ~expired_word; + while (added != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(added); + self.counts[expert] += 1; + added &= added - 1; + } + } + self.history[self.next_row] = incoming; + self.next_row += 1; + if (self.next_row == self.history.len) self.next_row = 0; + if (self.retained_rows < self.history.len) self.retained_rows += 1; + } +}; + +test "demand expires after exactly the configured number of rows" { + var demand = try RecentDemand.init(std.testing.allocator, 3); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{0}); + demand.observeRow(&.{ 1, 511 }); + demand.observeRow(&.{0}); + try std.testing.expectEqual(@as(Count, 2), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(511)); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); + + demand.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 1), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(1)); + demand.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(@as(Count, 0), demand.score(1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(511)); + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); +} + +test "duplicates contribute once and shared masked and invalid picks are ignored" { + var demand = try RecentDemand.init(std.testing.allocator, 2); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{ -1, 0, 0, 63, 64, 64, 511, 511, 512, 512, -2, 513, std.math.maxInt(i32), std.math.minInt(i32) }); + for (demand.scores(), 0..) |count, expert| { + const expected: Count = if (expert == 0 or expert == 63 or expert == 64 or expert == 511) 1 else 0; + try std.testing.expectEqual(expected, count); + } + try std.testing.expectEqual(@as(Count, 0), demand.score(-1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(512)); + try std.testing.expectEqual(@as(Count, 0), demand.score(std.math.maxInt(i32))); +} + +test "pick permutations produce identical state through wraparound" { + var ordered = try RecentDemand.init(std.testing.allocator, 3); + defer ordered.deinit(std.testing.allocator); + var permuted = try RecentDemand.init(std.testing.allocator, 3); + defer permuted.deinit(std.testing.allocator); + + for (0..10) |_| { + ordered.observeRow(&.{ 0, 63, 64, 511 }); + permuted.observeRow(&.{ 511, 64, 0, 63, 0, 512, -1 }); + ordered.observeRow(&.{ 7, 8 }); + permuted.observeRow(&.{ 8, 7, 8 }); + ordered.observeRow(&.{}); + permuted.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(ordered.counts, permuted.counts); + try std.testing.expectEqualSlices(Row, ordered.history, permuted.history); + try std.testing.expectEqual(ordered.next_row, permuted.next_row); + } +} + +test "expiration follows row order" { + var forward = try RecentDemand.init(std.testing.allocator, 2); + defer forward.deinit(std.testing.allocator); + var reverse = try RecentDemand.init(std.testing.allocator, 2); + defer reverse.deinit(std.testing.allocator); + + forward.observeRow(&.{1}); + forward.observeRow(&.{2}); + reverse.observeRow(&.{2}); + reverse.observeRow(&.{1}); + try std.testing.expectEqual(forward.counts, reverse.counts); + forward.observeRow(&.{}); + reverse.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 0), forward.score(1)); + try std.testing.expectEqual(@as(Count, 1), forward.score(2)); + try std.testing.expectEqual(@as(Count, 1), reverse.score(1)); + try std.testing.expectEqual(@as(Count, 0), reverse.score(2)); +} + +test "reset and replay reproduce every state transition" { + const rows = [_][]const i32{ &.{ 1, 2, 2 }, &.{511}, &.{}, &.{ -1, 512 }, &.{ 2, 511 }, &.{7}, &.{ 1, 7 } }; + var replayed = try RecentDemand.init(std.testing.allocator, 3); + defer replayed.deinit(std.testing.allocator); + var fresh = try RecentDemand.init(std.testing.allocator, 3); + defer fresh.deinit(std.testing.allocator); + + for (rows) |row| replayed.observeRow(row); + replayed.reset(); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(@as(usize, 0), replayed.rowCount()); + try std.testing.expectEqual(@as(usize, 0), replayed.next_row); + try std.testing.expectEqual(@as(usize, 3), replayed.windowRows()); + + for (rows) |row| { + replayed.observeRow(row); + fresh.observeRow(row); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(fresh.next_row, replayed.next_row); + try std.testing.expectEqual(fresh.rowCount(), replayed.rowCount()); + } + replayed.reset(); + replayed.reset(); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), replayed.scores().*); +} + +test "all expert scores stay exact at capacity over repeated replacements" { + var demand = try RecentDemand.init(std.testing.allocator, 1); + defer demand.deinit(std.testing.allocator); + var picks: [expert_count]i32 = undefined; + for (&picks, 0..) |*expert, index| expert.* = @intCast(index); + for (0..1024) |_| { + demand.observeRow(&picks); + try std.testing.expectEqual(@as([expert_count]Count, @splat(1)), demand.scores().*); + } + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + + var larger = try RecentDemand.init(std.testing.allocator, 4); + defer larger.deinit(std.testing.allocator); + for (0..16) |index| { + larger.observeRow(&picks); + const expected: Count = @intCast(@min(index + 1, larger.windowRows())); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } + for (0..4) |index| { + larger.observeRow(&.{}); + const expected: Count = @intCast(3 - index); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } +} + +test "invalid capacities are rejected before allocation" { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, 0)); + if (@bitSizeOf(usize) > @bitSizeOf(Count)) { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, @as(usize, std.math.maxInt(Count)) + 1)); + } +} diff --git a/zig/src/families/flashnext/exl3_round.zig b/zig/src/families/flashnext/exl3_round.zig new file mode 100644 index 0000000000..e160af6cac --- /dev/null +++ b/zig/src/families/flashnext/exl3_round.zig @@ -0,0 +1,621 @@ +//! One round's miss plan for a per-layer EXL3 expert pool (the host half of `cuda_exl3_pool.Pool.plan`). +//! Port of exl3_step.zig's round step onto upstream's `Experts`: there is no CPU expert any more, so every miss is +//! copied. A round is the picks of `rows` rows (`slots` a row, -1 skipped, the shared expert 512 never staged). +//! +//! The round splits into passes when its distinct routed experts do not fit the lines: pass 0 computes every expert +//! already held (resident, or a prefetch still in flight) plus the hottest misses that fit; each later pass takes the +//! next misses. A pass's copies may evict a line an EARLIER pass of the same round read: the device side orders that +//! (the copy waits for the earlier pass's release event), so after each pass its copies count as resident here. A +//! line of a copy that no compute waits for (a prefetch) stays in flight until `exl3_offload.reapAt`, and an +//! in-flight line is never a victim (the policy's in-flight protection). +//! +//! Prefetch: on swap rounds (`round % swap_every == 0`, round > 0) the hottest non-held experts by lifetime count +//! replace strictly colder residents (or fill empty lines) that this round does not use, at most `swap_cap` and at most +//! `prefetch_bytes / line_bytes` a round (the copy budget). No CUDA here. + +const std = @import("std"); +const offload = @import("exl3_offload.zig"); + +const shape = offload.shape; +const Vram = offload.Vram; +const empty_id = offload.Cache.empty_id; +pub const shared_id = offload.shared_id; + +/// A copy of routed expert `id` into pool line `line`; `evicted` is the expert the line held (empty_id: none). +pub const Copy = struct { id: u16, line: u32, evicted: u16 }; + +/// Ranges into `Round.ids` (the experts the pass computes) and `Round.copies` (what it stages first). +/// A round always has at least one pass; pass 0 also computes the shared expert. +pub const Pass = struct { id_lo: u32, id_hi: u32, copy_lo: u32, copy_hi: u32 }; + +pub const Options = struct { + swap_every: u16 = 4, + swap_cap: u16 = 96, + observe: bool = true, + /// Bytes a swap round may prefetch (the copy budget); 0 disables prefetch. + prefetch_bytes: u64 = 0, + /// Device bytes of one line (prefetch budget unit). + line_bytes: u64 = 1, + /// Double buffering: experts (of this layer) whose lines a single-pass round should not evict when it has room + /// without them, normally the previous round's. Their compute may still be running, so a copy into their lines + /// would wait for it; sparing them lets the copy overlap that compute. Ignored when the round would not fit. + protect: []const u16 = &.{}, +}; + +/// Caller-owned storage for one round's plan, sized once by `init` for the pool's line count and pick capacity. +pub const Round = struct { + ids: []u16, + /// The line each id of `ids` is read from in its pass (a later pass of the same round may reuse that line). + lines: []u32, + copies: []Copy, + passes: []Pass, + routed: []u16, // the round's routed picks as u16 (ordering scratch) + misses: []u16, + snap: []u16, // line -> id before a pass (empty_id for empty lines) + keep: []u16, // pass ids plus protected ids (stageUsedAt's keep list) + /// The round spared `protect`'s lines (double buffering held). + protected: bool = false, + n_ids: usize = 0, + n_copies: usize = 0, + n_passes: usize = 0, + n_prefetch: usize = 0, + /// Routed slots (pick in 0..511) whose expert was held / missing when the round started; shared slots counted apart. + hit_slots: u32 = 0, + miss_slots: u32 = 0, + shared_slots: u32 = 0, + /// Distinct routed experts, of which `held` were already in a line (`pending` of those still in flight). + unique: u32 = 0, + held: u32 = 0, + pending: u32 = 0, + /// Copies that replaced a held expert (mandatory and prefetch). + evictions: u32 = 0, + + pub fn init(gpa: std.mem.Allocator, lines: usize, max_picks: usize) !Round { + const uniq = @min(max_picks, shape.experts); + var r: Round = undefined; + r.ids = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.ids); + r.lines = try gpa.alloc(u32, uniq); + errdefer gpa.free(r.lines); + r.copies = try gpa.alloc(Copy, uniq + shape.experts); + errdefer gpa.free(r.copies); + r.passes = try gpa.alloc(Pass, uniq + 1); + errdefer gpa.free(r.passes); + r.routed = try gpa.alloc(u16, max_picks); + errdefer gpa.free(r.routed); + r.misses = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.misses); + r.snap = try gpa.alloc(u16, lines); + errdefer gpa.free(r.snap); + r.keep = try gpa.alloc(u16, uniq + lines); + r.reset(); + return r; + } + + pub fn deinit(r: *Round, gpa: std.mem.Allocator) void { + gpa.free(r.ids); + gpa.free(r.lines); + gpa.free(r.copies); + gpa.free(r.passes); + gpa.free(r.routed); + gpa.free(r.misses); + gpa.free(r.snap); + gpa.free(r.keep); + r.* = undefined; + } + + fn reset(r: *Round) void { + r.n_ids = 0; + r.n_copies = 0; + r.n_passes = 0; + r.n_prefetch = 0; + r.hit_slots = 0; + r.miss_slots = 0; + r.shared_slots = 0; + r.unique = 0; + r.held = 0; + r.pending = 0; + r.evictions = 0; + r.protected = false; + } + + pub fn passIds(r: *const Round, j: usize) []const u16 { + const p = r.passes[j]; + return r.ids[p.id_lo..p.id_hi]; + } + + pub fn passLines(r: *const Round, j: usize) []const u32 { + const p = r.passes[j]; + return r.lines[p.id_lo..p.id_hi]; + } + + pub fn passCopies(r: *const Round, j: usize) []const Copy { + const p = r.passes[j]; + return r.copies[p.copy_lo..p.copy_hi]; + } + + /// Copies no pass waits for (swap-round prefetch); their lines stay in flight until reaped. + pub fn prefetch(r: *const Round) []const Copy { + return r.copies[r.n_copies - r.n_prefetch .. r.n_copies]; + } + + /// Mandatory copies of all passes. + pub fn misses_(r: *const Round) usize { + return r.n_copies - r.n_prefetch; + } + + /// More than one pass: each pass reads its own masked picks (picks of other passes' experts set to -1). + pub fn masked(r: *const Round) bool { + return r.n_passes > 1; + } + + /// True when the round needs no copy at all (a graph of the experts may replay straight away). + pub fn allHits(r: *const Round) bool { + return r.n_copies == 0; + } +}; + +fn contains32(xs: []const u32, x: u32) bool { + return std.mem.indexOfScalar(u32, xs, x) != null; +} + +fn contains(xs: []const u16, id: u16) bool { + return std.mem.indexOfScalar(u16, xs, id) != null; +} + +const Order = struct { + counts: []const u32, + ids: []const u16, + fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| n += @intFromBool(x == id); + return n; + } + /// Lifetime count descending, then this round's appearances, then lower id (the policy's `usedFirst`). + fn hotter(o: Order, a: u16, b: u16) bool { + if (o.counts[a] != o.counts[b]) return o.counts[a] > o.counts[b]; + const aa = appearances(o.ids, a); + const ab = appearances(o.ids, b); + if (aa != ab) return aa > ab; + return a < b; + } +}; + +/// Lines whose copy is in flight and whose expert is not in `use`: they can be neither read nor evicted this pass. +fn blocked(v: Vram, use: []const u16) usize { + var n: usize = 0; + for (v.line, v.slots) |state, id| { + if (state == .inflight and !contains(use, id)) n += 1; + } + return n; +} + +fn snapshot(v: Vram, out: []u16) void { + for (v.slots, v.line, out) |id, state, *o| o.* = if (state == .empty) empty_id else id; +} + +/// Plans one round of `layer` (see the file header). Writes `v` (touches, staged lines) and `counts[layer]` (observe). +/// error.Busy: in-flight lines leave no line for a pass (reap or wait for the copies, then plan again; nothing was +/// written). error.Length: a pick above 512 or more picks than the round was sized for. +pub fn plan( + v: *Vram, + layer: u8, + picks: []const i32, + counts: *[shape.layers][shape.experts]u32, + now: u32, + round: u32, + opt: Options, + r: *Round, +) error{ Busy, Length }!void { + r.reset(); + if (picks.len > r.routed.len or layer >= shape.layers or v.slots.len != r.snap.len) return error.Length; + // the routed picks and the distinct experts, in first-appearance order + var n_routed: usize = 0; + var n_unique: usize = 0; + for (picks) |p| { + if (p < 0) continue; + if (p > shared_id) return error.Length; + const id: u16 = @intCast(p); + if (id == shared_id) { + r.shared_slots += 1; + continue; + } + r.routed[n_routed] = id; + n_routed += 1; + if (!contains(r.ids[0..n_unique], id)) { + r.ids[n_unique] = id; + n_unique += 1; + } + } + const routed = r.routed[0..n_routed]; + const uniq = r.ids[0..n_unique]; + // capacity check before anything is written: pass 0 holds the held experts, so at least one more line, or + // nothing to copy, must be free of in-flight copies of other experts + // (a pass after pass 0 can use every line not in flight, so one such line is enough) + var held_n: usize = 0; + for (uniq) |id| held_n += @intFromBool(offload.lineOf(v.*, layer, id) != null); + if (n_unique > held_n and blocked(v.*, &.{}) >= v.slots.len) return error.Busy; + + if (opt.observe) offload.observe(&counts[layer], routed) catch return error.Length; + const row: []const u32 = &counts[layer]; + + // held first (pass 0), misses hottest first; ids[] is rewritten in pass order + var n_held: usize = 0; + var n_miss: usize = 0; + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id)) |line| { + r.held += 1; + if (v.line[line] == .inflight) r.pending += 1 else _ = offload.touchAt(v, layer, id, now); + } else { + r.misses[n_miss] = id; + n_miss += 1; + } + } + for (routed) |id| { + if (offload.lineOf(v.*, layer, id) != null) r.hit_slots += 1 else r.miss_slots += 1; + } + r.unique = @intCast(n_unique); + // pack held ids at the front of ids (stable), then the sorted misses + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id) != null) { + r.ids[n_held] = id; + n_held += 1; + } + } + const misses = r.misses[0..n_miss]; + std.mem.sort(u16, misses, Order{ .counts = row, .ids = routed }, Order.hotter); + @memcpy(r.ids[n_held..][0..n_miss], misses); + + // passes + var at: usize = 0; // next id of ids[] to place + var mi: usize = 0; // next miss + while (at < n_unique) { + const first = r.n_passes == 0; + const lo = at; + const take_held: usize = if (first) n_held else 0; + at += take_held; + // lines usable this pass: all but in-flight lines of experts outside this pass + const usable = v.slots.len - blocked(v.*, r.ids[lo..at]); + const room = usable - take_held; + const n_new = @min(room, n_miss - mi); + std.debug.assert(n_new > 0 or take_held > 0); // the capacity check above + const hi = at + n_new; + // double buffering: also spare the protected experts' resident lines when the single pass still fits + var keep: []const u16 = r.ids[lo..hi]; + if (first and n_new == n_miss and opt.protect.len > 0) { + var spare: usize = 0; + var k: usize = hi - lo; + @memcpy(r.keep[0..k], r.ids[lo..hi]); + for (opt.protect) |id| { + if (contains(r.keep[0..k], id) or k == r.keep.len) continue; + const line = offload.lineOf(v.*, layer, id) orelse continue; + if (v.line[line] != .resident) continue; + r.keep[k] = id; + k += 1; + spare += 1; + } + if (room >= n_new + spare) { + keep = r.keep[0..k]; + r.protected = true; + } + } + snapshot(v.*, r.snap); + const copy_lo = r.n_copies; + for (r.ids[at..hi]) |id| { + const line = offload.stageUsedAt(v, layer, id, keep, counts) catch return error.Busy; + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.copies[r.n_copies] = .{ .id = id, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + } + // the device orders the next pass's copies after this pass's compute: these lines are resident from here on + for (r.copies[copy_lo..r.n_copies]) |c| offload.reapAt(v, layer, c.id, now) catch unreachable; + for (r.ids[lo..hi], r.lines[lo..hi]) |id, *ln| ln.* = @intCast(offload.lineOf(v.*, layer, id).?); + r.passes[r.n_passes] = .{ .id_lo = @intCast(lo), .id_hi = @intCast(hi), .copy_lo = @intCast(copy_lo), .copy_hi = @intCast(r.n_copies) }; + r.n_passes += 1; + at = hi; + mi += n_new; + } + r.n_ids = n_unique; + if (r.n_passes == 0) { // shared and -1 picks only: one pass with no routed expert (the shared expert still runs) + r.passes[0] = .{ .id_lo = 0, .id_hi = 0, .copy_lo = 0, .copy_hi = 0 }; + r.n_passes = 1; + } + + // swap-round prefetch: hottest non-held experts into empty or strictly colder lines this round does not use + if (opt.swap_every == 0 or round == 0 or round % opt.swap_every != 0 or opt.prefetch_bytes == 0 or opt.line_bytes == 0) return; + const budget = @min(@as(u64, opt.swap_cap), opt.prefetch_bytes / opt.line_bytes); + var staged: u64 = 0; + snapshot(v.*, r.snap); + while (staged < budget) { + // hottest id not held on this layer + var best: ?u16 = null; + var id: u16 = 0; + while (id < shape.experts) : (id += 1) { + if (row[id] == 0 or offload.lineOf(v.*, layer, id) != null) continue; + if (best == null or row[id] > row[best.?]) best = id; + } + const cand = best orelse break; + // the line it would take: an empty one, else the coldest resident outside this round + var victim: ?usize = null; + for (v.line, 0..) |state, i| { + if (state == .empty) { + victim = i; + break; + } + } + if (victim == null) { + for (v.line, v.slots, 0..) |state, held_id, i| { + if (state != .resident or contains(uniq_view(r), held_id)) continue; + if (victim == null or colder(v.*, layer, counts, i, victim.?)) victim = i; + } + const vi = victim orelse break; + if (offload.countAt(counts, layer, v.slots[vi]) >= row[cand]) break; // nothing colder than the candidate + } + const line = offload.stageUsedAt(v, layer, cand, uniq_view(r), counts) catch break; + std.debug.assert(line == victim.?); + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.snap[line] = cand; + r.copies[r.n_copies] = .{ .id = cand, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + r.n_prefetch += 1; + staged += 1; + } +} + +fn uniq_view(r: *const Round) []const u16 { + return r.ids[0..r.n_ids]; +} + +/// The policy's victim order (`vramCold`): lower lifetime count, then older stamp, then the higher line index. +fn colder(v: Vram, layer: u8, counts: *const [shape.layers][shape.experts]u32, i: usize, b: usize) bool { + const ci = offload.countAt(counts, layer, v.slots[i]); + const cb = offload.countAt(counts, layer, v.slots[b]); + return ci < cb or (ci == cb and v.stamp[i] < v.stamp[b]) or (ci == cb and v.stamp[i] == v.stamp[b] and i > b); +} + +/// Smallest pool that never needs a second pass for `rows` decode rows: every row's routed picks must fit at once. +pub fn minLines(top_k: usize) usize { + return top_k; +} + +/// Lines a layer gets from `free` device bytes after `reserve` and `other` (KV, scratch, graphs: the family's budget), +/// shared evenly by `layers` pools; refuses fewer than `min_lines` a layer. +pub fn linesFor(free: u64, reserve: u64, other: u64, line_bytes: u64, layers: usize, min_lines: usize) error{PoolTooSmall}!usize { + if (line_bytes == 0 or layers == 0) return error.PoolTooSmall; + const spend = std.math.sub(u64, free, reserve +| other) catch return error.PoolTooSmall; + const per = spend / (line_bytes * layers); + if (per < min_lines) return error.PoolTooSmall; + return @intCast(per); +} + +test { + _ = offload; + _ = @import("exl3_graph.zig"); +} + +// ------------------------------------------------------------------------------------------------ tests + +const testing = std.testing; +const Line = offload.Line; + +const Fix = struct { + slots: [16]u16 = undefined, + line: [16]Line = undefined, + stamp: [16]u32 = undefined, + lay: [16]u8 = undefined, + v: Vram = undefined, + counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)), + r: Round = undefined, + + fn init(f: *Fix, lines: usize, max_picks: usize) !void { + f.v = Vram.init(f.slots[0..lines], f.line[0..lines], f.stamp[0..lines], f.lay[0..lines]); + f.counts = @splat(@splat(0)); + f.r = try Round.init(testing.allocator, lines, max_picks); + } + fn deinit(f: *Fix) void { + f.r.deinit(testing.allocator); + } + fn go(f: *Fix, picks: []const i32, now: u32, round: u32, opt: Options) !void { + return plan(&f.v, 3, picks, &f.counts, now, round, opt, &f.r); + } + fn resident(f: *const Fix, id: u16) bool { + return offload.residentAt(f.v, 3, id); + } +}; + +test "misses fill empty lines in one pass, then the same picks are all hits" { + var f: Fix = .{}; + try f.init(10, 32); + defer f.deinit(); + const picks = [_]i32{ 5, 9, 5, 512, -1, 7 }; + try f.go(&picks, 1, 1, .{}); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 4), f.r.miss_slots); + try testing.expectEqual(@as(u32, 1), f.r.shared_slots); + try testing.expectEqual(@as(u32, 0), f.r.evictions); + for (f.r.passCopies(0)) |c| try testing.expectEqual(empty_id, c.evicted); + // 5 appears twice: hottest miss, copied first + try testing.expectEqual(@as(u16, 5), f.r.passCopies(0)[0].id); + try testing.expect(!f.r.masked()); + try testing.expect(f.resident(5) and f.resident(9) and f.resident(7)); + try testing.expectEqual(@as(u32, 2), f.counts[3][5]); + + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 4), f.r.hit_slots); + try testing.expectEqual(@as(u32, 3), f.r.held); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.passIds(0).len); + // a hit touched its line + try testing.expectEqual(@as(u32, 2), f.stamp[offload.lineOf(f.v, 3, 9).?]); +} + +test "synthetic prefill picks do not heat expert history" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 7, 8, 9 }, 1, 1, .{ .observe = false }); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 0), f.counts[3][7]); + try testing.expectEqual(@as(u32, 0), f.counts[3][8]); + try testing.expectEqual(@as(u32, 0), f.counts[3][9]); +} + +test "a miss evicts the coldest line outside the round, never a line the round reads" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // heat 1 and 3, leave 2 cold + try f.go(&[_]i32{ 1, 3 }, 2, 2, .{}); + try f.go(&[_]i32{ 1, 3 }, 3, 3, .{}); + try f.go(&[_]i32{ 1, 4 }, 4, 4, .{}); + try testing.expectEqual(@as(usize, 1), f.r.misses_()); + const c = f.r.passCopies(0)[0]; + try testing.expectEqual(@as(u16, 4), c.id); + try testing.expectEqual(@as(u16, 2), c.evicted); + try testing.expect(f.resident(1) and f.resident(3) and f.resident(4) and !f.resident(2)); + // 3 is hot but not in this round, 1 is in it; the cold one now is 4 (count 1) vs 3 (count 3) -> 4 goes + try f.go(&[_]i32{ 1, 3, 6 }, 5, 5, .{}); + try testing.expectEqual(@as(u16, 4), f.r.passCopies(0)[0].evicted); +} + +test "more distinct experts than lines split into passes that each fit, every expert computed once" { + var f: Fix = .{}; + try f.init(4, 64); + defer f.deinit(); + try f.go(&[_]i32{ 20, 21 }, 1, 1, .{}); + var picks: [30]i32 = undefined; + for (&picks, 0..) |*p, i| p.* = @intCast(10 + (i % 11)); // 10..20, 11 distinct; 20 held + picks[29] = 21; // held too + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.masked()); + try testing.expectEqual(@as(u32, 12), f.r.unique); + try testing.expectEqual(@as(u32, 2), f.r.held); + try testing.expectEqual(@as(usize, 10), f.r.misses_()); + var seen: [shape.experts]u8 = @splat(0); + for (0..f.r.n_passes) |j| { + const ids = f.r.passIds(j); + try testing.expect(ids.len >= 1 and ids.len <= 4); + // the lines a pass copies into are distinct + const cs = f.r.passCopies(j); + for (cs, 0..) |a, i| for (cs[i + 1 ..]) |b| try testing.expect(a.line != b.line); + for (ids) |id| seen[id] += 1; + } + for (10..22) |id| try testing.expectEqual(@as(u8, 1), seen[id]); + // each pass reads distinct lines; a copy of a later pass may land in a line an earlier pass read + var reused = false; + for (0..f.r.n_passes) |j| { + const ls = f.r.passLines(j); + for (ls, 0..) |a, i| for (ls[i + 1 ..]) |b| try testing.expect(a != b); + for (f.r.passCopies(j)) |c| { + try testing.expect(contains32(ls, c.line)); + if (j > 0) for (0..j) |k| { + if (contains32(f.r.passLines(k), c.line)) reused = true; + }; + } + } + try testing.expect(reused); + // pass 0 holds the two held experts first + try testing.expect(contains(f.r.passIds(0), 20) and contains(f.r.passIds(0), 21)); + try testing.expectEqual(@as(usize, 3), f.r.n_passes); // 2 held + 2 misses, then 4, then 4 + // the last pass's experts are resident after the round + for (f.r.passIds(f.r.n_passes - 1)) |id| try testing.expect(f.resident(id)); +} + +test "an in-flight prefetch line is not a victim and counts as held when the round asks for it" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // a prefetch of 9 into line of 2 (made by hand: stage it, leave it in flight) + const line = offload.lineOf(f.v, 3, 2).?; + f.slots[line] = 9; + f.line[line] = .inflight; + try f.go(&[_]i32{ 1, 4 }, 2, 2, .{}); + // 4 cannot take 9's in-flight line; it takes 3's (not in the round) + try testing.expectEqual(@as(u16, 3), f.r.passCopies(0)[0].evicted); + try testing.expectEqual(Line.inflight, f.line[line]); + try f.go(&[_]i32{ 9, 1 }, 3, 3, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 1), f.r.pending); + // every line in flight for another expert: busy, and nothing written + for (f.line[0..3]) |*l| l.* = .inflight; + const before = f.counts[3][7]; + try testing.expectError(error.Busy, f.go(&[_]i32{7}, 4, 4, .{})); + try testing.expectEqual(before, f.counts[3][7]); +} + +test "double buffering spares the previous round's lines when the round fits without them" { + var f: Fix = .{}; + try f.init(6, 16); + defer f.deinit(); + // previous round 1,2,3 (hot), older residents 4,5 (cold), one empty line + try f.go(&[_]i32{ 4, 5 }, 1, 1, .{}); + for (0..3) |i| try f.go(&[_]i32{ 1, 2, 3 }, @intCast(2 + i), @intCast(2 + i), .{}); + // two misses: with 1, 2, 3 spared they go to the empty line and the colder of 4, 5 + const prev = [_]u16{ 1, 2, 3 }; + try f.go(&[_]i32{ 7, 8 }, 5, 5, .{ .protect = &prev }); + try testing.expect(f.r.protected); + for (f.r.passCopies(0)) |c| try testing.expect(!contains(&prev, c.evicted)); + try testing.expect(f.resident(1) and f.resident(2) and f.resident(3)); + // a round that only fits by evicting a protected line ignores the protection + try f.go(&[_]i32{ 9, 10, 11, 12 }, 6, 6, .{ .protect = &[_]u16{ 7, 8, 1, 2, 3 } }); + try testing.expect(!f.r.protected); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); +} + +test "the shared expert and -1 are never staged; a pick above 512 is refused" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 512, -1, 512 }, 1, 1, .{}); + try testing.expectEqual(@as(usize, 0), f.r.n_copies); + try testing.expectEqual(@as(u32, 2), f.r.shared_slots); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 0), f.r.passIds(0).len); + for (f.slots[0..10]) |id| try testing.expect(id != shared_id); + try testing.expectError(error.Length, f.go(&[_]i32{513}, 2, 2, .{})); +} + +test "swap rounds prefetch the hottest non-held experts within the copy budget, only into colder lines" { + var f: Fix = .{}; + try f.init(4, 16); + defer f.deinit(); + const opt: Options = .{ .swap_every = 4, .prefetch_bytes = 200, .line_bytes = 100 }; + // lifetime counts: 30 hot (5), 31 warm (3), never in a line yet + f.counts[3][30] = 5; + f.counts[3][31] = 3; + f.counts[3][32] = 1; + try f.go(&[_]i32{ 1, 2 }, 1, 1, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); // round 1: not a swap round + try f.go(&[_]i32{ 1, 2 }, 4, 4, opt); + // two empty lines, budget 2: 30 and 31 in flight + try testing.expectEqual(@as(usize, 2), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 30), f.r.prefetch()[0].id); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[1].id); + try testing.expectEqual(Line.inflight, f.line[f.r.prefetch()[0].line]); + try testing.expect(!f.resident(30)); + try offload.reapAt(&f.v, 3, 30, 5); + try offload.reapAt(&f.v, 3, 31, 5); + // round 8: 32 (count 1) is not hotter than any resident outside the round (30: 5, 31: 3) -> nothing evicted + try f.go(&[_]i32{ 1, 2 }, 8, 8, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); + // heat 32 above 31: it replaces 31, never 1 or 2 (in this round) + f.counts[3][32] = 9; + try f.go(&[_]i32{ 1, 2 }, 12, 12, opt); + try testing.expectEqual(@as(usize, 1), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[0].evicted); +} + +test "line budget: free memory after reserve and the family's other bytes, at least min lines a layer" { + const line: u64 = 1_843_200; + // 12 GiB free, 1 GiB reserve, 5.5 GiB for KV and scratch: 5.5 GiB / (48 x 1,843,200 B) = 66 lines a layer + try testing.expectEqual(@as(usize, 66), try linesFor(12 << 30, 1 << 30, (11 << 30) / 2, line, 48, 10)); + try testing.expectEqual(@as(usize, 10), try linesFor(10 * 48 * line + 100, 100, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(10 * 48 * line - 1, 0, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(1 << 30, 2 << 30, 0, line, 48, 10)); +} diff --git a/zig/src/families/flashnext/exl3_route.zig b/zig/src/families/flashnext/exl3_route.zig new file mode 100644 index 0000000000..1bed1ce562 --- /dev/null +++ b/zig/src/families/flashnext/exl3_route.zig @@ -0,0 +1,103 @@ +//! One row of the Flash Next router. +//! v0.6.6 `src/tensorfold/cuda/moe.py` `_topk_rows`: fp32 logits, ties to the +//! lower id, softmax over the selected routed logits only, then the shared +//! expert in the last slot. Weights and the shared gate round through bf16 +//! the same way as `zig/src/core/contract/contract.zig` `toBf16` / `fromBf16`. +//! This file does not import that module: `zig test` here cannot use `../`. + +const std = @import("std"); + +/// `logits` is `experts + 1` fp32 values. The last value is the shared gate. +/// `picks` and `weights` hold `topk + 1` slots. The last slot is the shared +/// expert id (`logits.len - 1`) and its sigmoid weight. +/// `scratch` is `experts` long and is overwritten. +/// Length is checked before any store to `picks` or `weights`. +pub fn select( + logits: []const f32, + topk: usize, + picks: []u16, + weights: []f32, + scratch: []f32, +) error{Length}!void { + if (logits.len < 2) return error.Length; + const experts = logits.len - 1; + if (experts > std.math.maxInt(u16)) return error.Length; + if (topk == 0 or topk > experts) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != experts) return error.Length; + + @memcpy(scratch, logits[0..experts]); + var top = scratch[0]; + for (scratch[1..]) |v| top = @max(top, v); + + var total: f32 = 0; + for (0..topk) |k| { + var best_i: usize = 0; + var best = scratch[0]; + for (scratch[1..], 1..) |v, i| { + if (v > best) { + best = v; + best_i = i; + } + } + const ex = @exp(best - top); + picks[k] = @intCast(best_i); + weights[k] = ex; + total += ex; + scratch[best_i] = -std.math.inf(f32); + } + for (weights[0..topk]) |*w| w.* = bf16(w.* / total); + + const shared: u16 = @intCast(experts); + const gate = bf16(logits[experts]); + picks[topk] = shared; + weights[topk] = bf16(1.0 / (1.0 + @exp(-gate))); +} + +/// fp32 through bf16 and back. NaN stays a NaN. Same rounding as the contract helper. +fn bf16(v: f32) f32 { + const u: u32 = @bitCast(v); + const h: u16 = if (std.math.isNan(v)) + @truncate((u >> 16) | 0x40) + else + @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); + return @bitCast(@as(u32, h) << 16); +} + +test "ties take the lower id, and the shared expert is the last slot" { + const logits = [_]f32{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9, 9 }; + var weights = [_]f32{ 9, 9, 9 }; + var scratch: [4]f32 = undefined; + try select(&logits, 2, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(u16, 1), picks[0]); + try std.testing.expectEqual(@as(u16, 2), picks[1]); + try std.testing.expectEqual(@as(u16, 4), picks[2]); + try std.testing.expectEqual(@as(f32, 0.5), weights[0]); + try std.testing.expectEqual(@as(f32, 0.5), weights[1]); + try std.testing.expectEqual(@as(f32, 0.5), weights[2]); +} + +test "one winner keeps the whole routed weight" { + const logits = [_]f32{ 5, 5, 1, 2 }; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 9, 9 }; + var scratch: [3]f32 = undefined; + try select(&logits, 1, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(u16, 0), picks[0]); + try std.testing.expectEqual(@as(u16, 3), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); + const gate = bf16(2); + const want = bf16(1.0 / (1.0 + @exp(-gate))); + try std.testing.expectEqual(want, weights[1]); +} + +test "a short buffer is refused before picks change" { + const logits = [_]f32{ 1, 2, 0 }; + var picks = [_]u16{ 7, 7 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [2]f32 = undefined; + try std.testing.expectError(error.Length, select(&logits, 2, picks[0..1], &weights, &scratch)); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try std.testing.expectEqual(@as(f32, 3), weights[0]); +} diff --git a/zig/src/families/flashnext/exl3_router.zig b/zig/src/families/flashnext/exl3_router.zig new file mode 100644 index 0000000000..af1f4e3cbe --- /dev/null +++ b/zig/src/families/flashnext/exl3_router.zig @@ -0,0 +1,264 @@ +//! Rank 0 of the Flash Next router. +//! `mlp.gate` is a replicated bf16 matrix, 512 by 2560 (2,621,440 bytes). +//! A 513th row would be 2,626,560 bytes and is not read from that product. +//! `mlp.shared_expert_gate` is a second replicated vector. Its length is the +//! caller's header span in bf16 elements, not an expert-cut width. +//! Both dots stay fp32. `exl3_route.select` applies the sigmoid. + +const std = @import("std"); +const route = @import("exl3_route.zig"); + +pub const gate_rows: usize = 512; +pub const gate_cols: usize = 2560; +pub const gate_bytes: usize = gate_rows * gate_cols * @sizeOf(u16); + +pub const Error = error{ + NotRank0, + EmptyGate, + GateBytes, + MissingSharedGate, + Length, +}; + +comptime { + std.debug.assert(gate_bytes == 2_621_440); + std.debug.assert((gate_rows + 1) * gate_cols * @sizeOf(u16) == 2_626_560); +} + +/// Storage type of the two replicated gate tensors. The published 3.05bpw pack stores both as F16 +/// (safetensors header `F16`, verified: std 0.02 as F16, 1e-6 as bf16), although config.json says bfloat16. +/// `launch` keeps the bf16 reading; use `launchDtype(.f16, ...)` for the pack. +pub const Dtype = enum { bf16, f16 }; + +fn dot(dtype: Dtype, weight: []const u8, x: []const f32) f32 { + var sum: f32 = 0; + for (x, 0..) |v, i| { + const h = std.mem.readInt(u16, weight[i * 2 ..][0..2], .little); + const w: f32 = switch (dtype) { + .bf16 => @bitCast(@as(u32, h) << 16), + .f16 => @floatCast(@as(f16, @bitCast(h))), + }; + sum += w * v; + } + return sum; +} + +/// `gate` is row-major bf16, exactly `gate_bytes`, times `hidden` (`gate_cols`). +/// `shared_gate` is bf16; `shared_x.len` is its element count (the header span). +/// `logits` receives 512 routed dots, then the one shared-gate dot. +/// `select` writes `picks` and `weights`. Rank, empty, and size checks +/// return before any of those stores. An empty gate is the absent `mlp.gate` +/// token: nothing is invented. +pub fn launch( + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + return launchDtype(.bf16, rank, gate, hidden, shared_gate, shared_x, topk, logits, picks, weights, scratch); +} + +/// `launch` with the gate tensors' storage type named. +pub fn launchDtype( + dtype: Dtype, + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + if (rank != 0) return error.NotRank0; + if (gate.len == 0) return error.EmptyGate; + if (gate.len != gate_bytes) return error.GateBytes; + if (shared_gate.len == 0) return error.MissingSharedGate; + if (shared_gate.len % 2 != 0) return error.Length; + const shared_n = shared_gate.len / 2; + if (shared_x.len != shared_n) return error.Length; + if (hidden.len != gate_cols) return error.Length; + if (logits.len != gate_rows + 1) return error.Length; + if (topk == 0 or topk > gate_rows) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != gate_rows) return error.Length; + + const row_bytes = gate_cols * @sizeOf(u16); + for (0..gate_rows) |r| { + const row = gate[r * row_bytes ..][0..row_bytes]; + logits[r] = dot(dtype, row, hidden); + } + logits[gate_rows] = dot(dtype, shared_gate, shared_x); + try route.select(logits, topk, picks, weights, scratch); +} + +const one_bf16: u16 = 0x3f80; +const winner: usize = 42; +const shared_span: usize = 3; + +fn paint(gate: []u8) void { + @memset(gate, 0); + const at = winner * gate_cols * @sizeOf(u16); + std.mem.writeInt(u16, gate[at..][0..2], one_bf16, .little); +} + +fn sharedGate() [shared_span * 2]u8 { + var bytes: [shared_span * 2]u8 = @splat(0); + for (0..shared_span) |i| std.mem.writeInt(u16, bytes[i * 2 ..][0..2], one_bf16, .little); + return bytes; +} + +fn hiddenVec() [gate_cols]f32 { + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + return hidden; +} + +fn expectUntouched(logits: []const f32, picks: []const u16, weights: []const f32) !void { + try std.testing.expectEqual(@as(f32, 8), logits[0]); + try std.testing.expectEqual(@as(f32, 8), logits[gate_rows]); + try std.testing.expectEqualSlices(u16, &.{ 9, 9 }, picks); + try std.testing.expectEqualSlices(f32, &.{ 3, 3 }, weights); +} + +test "one nonzero gate row and a shared span of 3 select id 512" { + try std.testing.expect(shared_span != gate_cols); + try std.testing.expect(shared_span != 640); + try std.testing.expect(shared_span != 1280); + try std.testing.expect(shared_span != 61_440); + try std.testing.expectEqual(@as(usize, 2_621_440), gate_bytes); + + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = undefined; + try launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(f32, 0), logits[0]); + try std.testing.expectEqual(@as(f32, 1), logits[winner]); + try std.testing.expectEqual(@as(f32, 0), logits[gate_rows - 1]); + try std.testing.expectEqual(@as(f32, 9), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, winner), picks[0]); + try std.testing.expectEqual(@as(u16, 512), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); +} + +test "rank 1 leaves the pick buffer untouched" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.NotRank0, launch(1, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a missing shared gate does not call select" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.MissingSharedGate, launch(0, gate, &hidden, &.{}, &.{}, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a 513-row gate buffer is refused" { + const gpa = std.testing.allocator; + const wide_bytes = (gate_rows + 1) * gate_cols * @sizeOf(u16); + try std.testing.expectEqual(@as(usize, 2_626_560), wide_bytes); + const gate = try gpa.alloc(u8, wide_bytes); + defer gpa.free(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.GateBytes, launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "an empty gate does not launch" { + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.EmptyGate, launch(0, &.{}, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a shared-span mismatch stores nothing" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const short_x = [_]f32{1}; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.Length, launch(0, gate, &hidden, &shared_gate, &short_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "the same bits read as F16 are not the bf16 value" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + const one_f16: u16 = 0x3c00; + std.mem.writeInt(u16, gate[7 * gate_cols * 2 ..][0..2], one_f16, .little); + var shared: [gate_cols * 2]u8 = @splat(0); + std.mem.writeInt(u16, shared[0..2], one_f16, .little); + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + var logits: [gate_rows + 1]f32 = undefined; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 0, 0 }; + var scratch: [gate_rows]f32 = undefined; + try launchDtype(.f16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(f32, 1), logits[7]); + try std.testing.expectEqual(@as(f32, 1), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try launchDtype(.bf16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expect(logits[7] != 1); +} diff --git a/zig/src/families/flashnext/exl3_scan.zig b/zig/src/families/flashnext/exl3_scan.zig new file mode 100644 index 0000000000..30260a0a0f --- /dev/null +++ b/zig/src/families/flashnext/exl3_scan.zig @@ -0,0 +1,211 @@ +//! EXL3 safetensors header scan. Names and shapes only. No weight bytes. +//! Group rules follow the v0.6.6 `parse_group` checker. +//! `centredOffset` copies at most 64 means onto the stack. A longer slice is `NormAmbiguous`. + +const std = @import("std"); + +/// Same tags, in the same order, as `exl3_format.Codebook`. +/// This file is its own `zig test` root, so it cannot import `../../core/exl3_format.zig`. +pub const format = struct { + pub const Codebook = enum { inst3, mcg, mul1 }; +}; + +pub const Scales = enum { suh, su, svh, sv }; + +pub const Group = struct { + halves: u16, + codebook: format.Codebook, + k: usize, + n: usize, + in_scales: Scales, + out_scales: Scales, + bias: bool, +}; + +pub const Parts = struct { + trellis_words: usize, + k_tiles: usize, + n_tiles: usize, + suh_len: ?usize = null, + su_len: ?usize = null, + svh_len: ?usize = null, + sv_len: ?usize = null, + bias: bool = false, + mcg: bool = false, + mul1: bool = false, +}; + +pub const Kind = enum { exl3_part, ngram_row, plain }; + +const suffixes = [_][]const u8{ "trellis", "suh", "su", "svh", "sv", "bias", "mcg", "mul1" }; + +/// One EXL3 linear group from its header parts. Half-bit widths need the mul1 marker. +pub fn parseGroup(parts: Parts) error{NotExl3}!Group { + if (parts.trellis_words == 0 or parts.trellis_words % 8 != 0) return error.NotExl3; + const halves_usize = parts.trellis_words / 8; + if (halves_usize < 2 or halves_usize > 16) return error.NotExl3; + const halves: u16 = @intCast(halves_usize); + if (parts.mcg and parts.mul1) return error.NotExl3; + const codebook: format.Codebook = if (parts.mul1) .mul1 else if (parts.mcg) .mcg else .inst3; + if (halves % 2 == 1 and codebook != .mul1) return error.NotExl3; + + const k = parts.k_tiles * 16; + const n = parts.n_tiles * 16; + if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0) return error.NotExl3; + + const in_scales: Scales = if (parts.suh_len != null and parts.su_len == null) + .suh + else if (parts.su_len != null and parts.suh_len == null) + .su + else + return error.NotExl3; + if (in_scales == .suh and parts.suh_len.? != k) return error.NotExl3; + if (in_scales == .su and parts.su_len.? != k / 16) return error.NotExl3; + + const out_scales: Scales = if (parts.svh_len != null and parts.sv_len == null) + .svh + else if (parts.sv_len != null and parts.svh_len == null) + .sv + else + return error.NotExl3; + if (out_scales == .svh and parts.svh_len.? != n) return error.NotExl3; + if (out_scales == .sv and parts.sv_len.? != n / 16) return error.NotExl3; + + return .{ + .halves = halves, + .codebook = codebook, + .k = k, + .n = n, + .in_scales = in_scales, + .out_scales = out_scales, + .bias = parts.bias, + }; +} + +/// An n-gram shard is a rank-2 `ngram_embedding` `.trellis`, not an EXL3 tile tensor. +pub fn classifyName(name: []const u8, rank: usize) Kind { + if (rank == 2 and std.mem.indexOf(u8, name, "ngram_embedding") != null and std.mem.endsWith(u8, name, ".trellis")) { + return .ngram_row; + } + const dot = std.mem.lastIndexOfScalar(u8, name, '.') orelse return .plain; + const suffix = name[dot + 1 ..]; + for (suffixes) |part| if (std.mem.eql(u8, suffix, part)) return .exl3_part; + return .plain; +} + +/// The MTP mixer's extra file, the one turboderp keeps outside the index. +pub fn isMtpMixerPatch(filename: []const u8) bool { + return std.mem.eql(u8, filename, "mtp_hyper_connection_mixer_patch.safetensors"); +} + +fn codebookName(name: []const u8) error{QuantRefuse}!void { + if (std.mem.eql(u8, name, "mul1") or std.mem.eql(u8, name, "mcg") or std.mem.eql(u8, name, "3inst") or std.mem.eql(u8, name, "inst3")) return; + return error.QuantRefuse; +} + +fn integerBits(bits: f32) error{QuantRefuse}!void { + if (!std.math.isFinite(bits) or bits < 1 or bits > 8) return error.QuantRefuse; + if (bits != @round(bits)) return error.QuantRefuse; +} + +/// Admit a pack from config fields alone. Average bits may be fractional. Head and MTP bits may not. +pub fn admit(codebook: []const u8, average_bits: f32, head_bits: ?f32, mtp_bits: ?f32) error{QuantRefuse}!void { + try codebookName(codebook); + if (!std.math.isFinite(average_bits) or average_bits < 1 or average_bits > 8) return error.QuantRefuse; + if (head_bits) |bits| try integerBits(bits); + if (mtp_bits) |bits| try integerBits(bits); +} + +/// 1 when centred norms are stored as gamma-1, 0 when stored as gamma. +/// Empty input returns 1. More than 64 means is `NormAmbiguous` and is not allocated. +/// Median is the upper middle element, `sorted[len / 2]`. +pub fn centredOffset(means: []const f32) error{NormAmbiguous}!f32 { + if (means.len == 0) return 1; + if (means.len > 64) return error.NormAmbiguous; + var copy: [64]f32 = undefined; + @memcpy(copy[0..means.len], means); + var i: usize = 1; + while (i < means.len) : (i += 1) { + const v = copy[i]; + var j = i; + while (j > 0 and copy[j - 1] > v) : (j -= 1) copy[j] = copy[j - 1]; + copy[j] = v; + } + const median = copy[means.len / 2]; + var above: usize = 0; + for (means) |m| if (m > 0.5) { + above += 1; + }; + const frac = @as(f32, @floatFromInt(above)) / @as(f32, @floatFromInt(means.len)); + const around_zero = frac <= 0.1 and median >= -0.5 and median <= 0.25; + const around_one = frac >= 0.9 and median >= 0.75 and median <= 1.5; + if (around_zero == around_one) return error.NormAmbiguous; + return if (around_zero) 1 else 0; +} + +test "a 3-bit mul1 gate group parses and a half-bit group needs mul1" { + const gate = try parseGroup(.{ + .trellis_words = 48, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }); + try std.testing.expectEqual(@as(u16, 6), gate.halves); + try std.testing.expectEqual(format.Codebook.mul1, gate.codebook); + try std.testing.expectEqual(@as(usize, 2560), gate.k); + try std.testing.expectEqual(@as(usize, 640), gate.n); + try std.testing.expectEqual(Scales.suh, gate.in_scales); + try std.testing.expectEqual(Scales.svh, gate.out_scales); + + const half = Parts{ + .trellis_words = 40, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }; + try std.testing.expectEqual(@as(u16, 5), (try parseGroup(half)).halves); + var no_marker = half; + no_marker.mul1 = false; + try std.testing.expectError(error.NotExl3, parseGroup(no_marker)); + var both = half; + both.mcg = true; + try std.testing.expectError(error.NotExl3, parseGroup(both)); +} + +test "K must be a multiple of 128" { + try std.testing.expectError(error.NotExl3, parseGroup(.{ + .trellis_words = 48, + .k_tiles = 20, + .n_tiles = 40, + .suh_len = 320, + .svh_len = 640, + .mul1 = true, + })); +} + +test "names sort into n-gram rows, EXL3 parts and plain tensors" { + try std.testing.expectEqual(Kind.exl3_part, classifyName("model.language_model.layers.0.mlp.experts.0.gate_proj.trellis", 3)); + try std.testing.expectEqual(Kind.ngram_row, classifyName("ple.ngram_embedding.shard_3.trellis", 2)); + try std.testing.expectEqual(Kind.exl3_part, classifyName("ple.ngram_embedding.shard_3.trellis", 3)); + try std.testing.expectEqual(Kind.plain, classifyName("model.embed_tokens.weight", 2)); + try std.testing.expect(isMtpMixerPatch("mtp_hyper_connection_mixer_patch.safetensors")); + try std.testing.expect(!isMtpMixerPatch("model.safetensors")); +} + +test "the 3.05 config is admitted and a fractional head width is not" { + try admit("mul1", 3.05, 5, 3); + try admit("3inst", 4, null, null); + try std.testing.expectError(error.QuantRefuse, admit("mul1", 3.05, 3.05, null)); + try std.testing.expectError(error.QuantRefuse, admit("exl2", 3, null, null)); +} + +test "centred norms are gamma-1 when the means sit near zero" { + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{ -0.02, 0.01, 0.0, 0.04 })); + try std.testing.expectEqual(@as(f32, 0), try centredOffset(&.{ 1.0, 0.98, 1.02, 1.01 })); + try std.testing.expectError(error.NormAmbiguous, centredOffset(&.{ 0.4, 0.9 })); + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{})); +} diff --git a/zig/src/families/flashnext/exl3_shape.zig b/zig/src/families/flashnext/exl3_shape.zig new file mode 100644 index 0000000000..7a8c36fa3a --- /dev/null +++ b/zig/src/families/flashnext/exl3_shape.zig @@ -0,0 +1,72 @@ +//! Published text shape of turboderp/Qwen3.8-Flash-Next-exl3, revision 3.05bpw_h5_ng5. +//! Read from that revision's config.json on 2026-10-16 (text_config + quantization_config). +//! Per-tensor EXL3 widths still come from each trellis, never from `bits`. + +const std = @import("std"); + +pub const model = "turboderp/Qwen3.8-Flash-Next-exl3"; +pub const revision = "3.05bpw_h5_ng5"; + +pub const hidden: usize = 2560; +pub const layers: usize = 48; +pub const vocab: usize = 248320; +pub const heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const experts: usize = 512; +pub const topk: usize = 10; +pub const expert_intermediate: usize = 640; +pub const shared_intermediate: usize = 640; +/// Routed ids are 0..511. The shared expert is the next id in the same table. +pub const shared_expert: usize = experts; +pub const hc_count: usize = 4; +pub const hc_lowrank: usize = 320; + +pub const ngram: usize = 3; +pub const ngram_heads_per: usize = 8; +pub const ngram_dh: usize = 160; +/// Branch name `ng5`: the n-gram row codec is 5 bits. Confirmed by words-per-row once a shard header is read. +pub const ngram_bits: usize = 5; +pub const ngram_shards: usize = 128; +pub const ngram_vocab_base: usize = 20_000_000; + +pub const quant_method = "exl3"; +pub const quant_version = "1.4.4"; +pub const average_bits: f32 = 3.05; +pub const head_bits: usize = 5; +pub const mtp_bits: usize = 3; +pub const codebook = "mul1"; + +pub const tile: usize = 16; +pub const hadamard: usize = 128; +pub const tp_ranks: usize = 2; + +pub const linear_layers_per_cycle: usize = 3; +pub const full_attention_interval: usize = 4; + +/// Heads in one n-gram lookup: (ngram - 1) * heads_per_ngram. +pub fn ngramHeadCount() usize { + return (ngram - 1) * ngram_heads_per; +} + +/// q_proj outputs. 24 * 256 = 6144. +pub fn qOut() usize { + return heads * head_dim; +} + +/// One attention layer's kv width, keys or values. 2 * 256 = 512. +pub fn kvOut() usize { + return kv_heads * head_dim; +} + +test "the published 3.05 shape is the one the cuts and the budget are written against" { + try std.testing.expectEqual(@as(usize, 16), ngramHeadCount()); + try std.testing.expectEqual(@as(usize, 6144), qOut()); + try std.testing.expectEqual(@as(usize, 512), kvOut()); + try std.testing.expect(hidden % hadamard == 0); + try std.testing.expect(expert_intermediate % hadamard == 0); + try std.testing.expect(qOut() % hadamard == 0); + // Half of 640 is 320, which is not a Hadamard block. Column TP inside one expert is illegal. + try std.testing.expect(expert_intermediate / tp_ranks % hadamard != 0); + try std.testing.expect(hidden / tp_ranks % hadamard == 0); +} diff --git a/zig/src/families/flashnext/exl3_shared_kv.zig b/zig/src/families/flashnext/exl3_shared_kv.zig new file mode 100644 index 0000000000..214d505104 --- /dev/null +++ b/zig/src/families/flashnext/exl3_shared_kv.zig @@ -0,0 +1,486 @@ +const std = @import("std"); + +pub const Geometry = struct { + lanes: u32, + layers: u32, + blocks: u32, + slots: u32, + row_bytes: usize, + ratio: u32 = 4, + window_rows: usize, + work_words: usize, + + pub fn validate(geometry: Geometry) !void { + if (geometry.lanes == 0 or geometry.layers == 0 or geometry.blocks == 0 or geometry.slots == 0 or geometry.slots > std.math.maxInt(i32) or geometry.row_bytes == 0 or geometry.row_bytes % 2 != 0 or geometry.ratio == 0 or geometry.window_rows == 0) return error.SharedKvGeometry; + } + + pub fn blockBytes(geometry: Geometry) !usize { + try geometry.validate(); + return std.math.mul(usize, try std.math.mul(usize, geometry.row_bytes, geometry.ratio), 2); + } + + pub fn accounting(geometry: Geometry) !Accounting { + const block_bytes = try geometry.blockBytes(); + const payload = try std.math.mul(usize, geometry.slots, block_bytes); + const owner_bytes = try std.math.mul(usize, geometry.slots, @sizeOf(Owner)); + const stamps = try std.math.mul(usize, geometry.slots, @sizeOf(u64)); + const inverse_words = try std.math.mul(usize, geometry.layers, geometry.blocks); + const inverse = try std.math.mul(usize, inverse_words, @sizeOf(i32)); + const windows = try std.math.mul(usize, try std.math.mul(usize, geometry.layers, geometry.window_rows), try std.math.mul(usize, geometry.row_bytes, 2)); + const controls = try std.math.mul(usize, geometry.layers, 8 * @sizeOf(u32)); + const work = try std.math.mul(usize, geometry.work_words, @sizeOf(u32)); + const shared = try std.math.add(usize, payload, try std.math.add(usize, owner_bytes, stamps)); + const sequence = try std.math.add(usize, windows, try std.math.add(usize, try std.math.mul(usize, inverse, 2), try std.math.add(usize, controls, work))); + return .{ .payload = payload, .owners = owner_bytes, .stamps = stamps, .inverse = inverse, .windows = windows, .controls = controls, .work = work, .shared = shared, .sequence = sequence, .rank = try std.math.add(usize, shared, try std.math.mul(usize, sequence, geometry.lanes)) }; + } + + pub fn duplicatedLegacyBytes(geometry: Geometry) !usize { + const payload = try std.math.mul(usize, geometry.slots, try geometry.blockBytes()); + const metadata = try std.math.mul(usize, geometry.slots, 2 * @sizeOf(u32)); + return std.math.mul(usize, try std.math.add(usize, payload, metadata), geometry.lanes - 1); + } +}; + +pub const Accounting = struct { + payload: usize, + owners: usize, + stamps: usize, + inverse: usize, + windows: usize, + controls: usize, + work: usize, + shared: usize, + sequence: usize, + rank: usize, +}; + +pub const Sequence = struct { arena: u64, lane: u32, generation: u64 }; +pub const Key = struct { layer: u32, block: u32 }; +pub const Tag = struct { sequence: Sequence, key: Key }; +pub const Owner = extern struct { + generation: u64 = 0, + lane: u32 = 0, + layer: u32 = 0, + block: u32 = 0, + valid: u32 = 0, +}; + +pub const Lease = struct { arena: u64, serial: u64, sequence: Sequence }; +pub const Assignment = struct { key: Key, slot: u32, hit: bool, written: bool = false }; +pub const Phase = enum { acquired, planned, submitted, ready, poisoned }; +const Active = struct { lease: Lease, phase: Phase = .acquired, count: usize = 0 }; +const Lane = struct { generation: u64 = 0, live: bool = false, failed: bool = false }; + +pub const Bindings = struct { + keys: u64, + values: u64, + owners: u64, + stamps: u64, + clock: u64, + inverse: u64, + marks: u64, + controls: u64, + work: u64, + window_keys: u64, + window_values: u64, + canonical: u64, + cap: u64, + rows: u64, + origin_hash: u64, +}; + +pub const GraphKey = struct { + format: u32 = 2, + sequence: Sequence, + geometry: Geometry, + bindings: Bindings, +}; + +pub const OwnedEntry = struct { key: Key, stamp: u64, source_slot: u32 }; + +pub const OwnedBlocks = struct { + version: u32 = 2, + source: Sequence, + geometry: Geometry, + clock: u64, + entries: []OwnedEntry, + payload: []u8, + + pub fn maxAge(image: *const OwnedBlocks) !u64 { + var maximum: u64 = 0; + for (image.entries) |entry| { + if (entry.stamp == 0 or entry.stamp > image.clock) return error.SharedKvSnapshotEntry; + maximum = @max(maximum, image.clock - entry.stamp); + } + return maximum; + } + + pub fn deinit(image: *OwnedBlocks, allocator: std.mem.Allocator) void { + allocator.free(image.payload); + allocator.free(image.entries); + image.* = undefined; + } + + pub fn validate(image: *const OwnedBlocks, destination: Geometry) !void { + if (image.version != 2) return error.SharedKvSnapshotVersion; + try image.geometry.validate(); + try destination.validate(); + if (image.source.arena == 0 or image.source.generation == 0 or image.source.lane >= image.geometry.lanes or image.geometry.layers != destination.layers or image.geometry.blocks != destination.blocks or image.geometry.row_bytes != destination.row_bytes or image.geometry.ratio != destination.ratio) return error.SharedKvSnapshotGeometry; + if (image.entries.len > image.geometry.slots or image.entries.len > destination.slots) return error.SharedKvCapacity; + if (image.payload.len != try std.math.mul(usize, image.entries.len, try destination.blockBytes())) return error.SharedKvSnapshotLength; + for (image.entries, 0..) |entry, index| { + if (entry.key.layer >= destination.layers or entry.key.block >= destination.blocks or entry.source_slot >= image.geometry.slots or entry.stamp == 0 or entry.stamp > image.clock) return error.SharedKvSnapshotEntry; + for (image.entries[0..index]) |previous| if (std.meta.eql(previous.key, entry.key) or previous.source_slot == entry.source_slot) return error.SharedKvSnapshotDuplicate; + } + } +}; + +var identities: std.atomic.Value(u64) = .init(0); + +pub fn identity() !u64 { + var previous = identities.load(.monotonic); + while (true) { + const next = try std.math.add(u64, previous, 1); + previous = identities.cmpxchgWeak(previous, next, .monotonic, .monotonic) orelse return next; + } +} + +pub const Arena = struct { + allocator: std.mem.Allocator, + geometry: Geometry, + id: u64, + clock: u64 = 0, + serial: u64 = 0, + lanes: []Lane, + owners: []Owner, + stamps: []u64, + inverse: []i32, + payload: []u8, + assignments: []Assignment, + pinned: []bool, + active: ?Active = null, + + pub fn init(allocator: std.mem.Allocator, geometry: Geometry) !Arena { + const bytes = try geometry.accounting(); + const inverse_count = try std.math.mul(usize, bytes.inverse / @sizeOf(i32), geometry.lanes); + const lanes = try allocator.alloc(Lane, geometry.lanes); + errdefer allocator.free(lanes); + const owners = try allocator.alloc(Owner, geometry.slots); + errdefer allocator.free(owners); + const stamps = try allocator.alloc(u64, geometry.slots); + errdefer allocator.free(stamps); + const inverse = try allocator.alloc(i32, inverse_count); + errdefer allocator.free(inverse); + const payload = try allocator.alloc(u8, bytes.payload); + errdefer allocator.free(payload); + const assignments = try allocator.alloc(Assignment, geometry.slots); + errdefer allocator.free(assignments); + const pinned = try allocator.alloc(bool, geometry.slots); + errdefer allocator.free(pinned); + @memset(lanes, .{}); + @memset(owners, .{}); + @memset(stamps, 0); + @memset(inverse, -1); + @memset(payload, 0); + @memset(pinned, false); + return .{ .allocator = allocator, .geometry = geometry, .id = try identity(), .lanes = lanes, .owners = owners, .stamps = stamps, .inverse = inverse, .payload = payload, .assignments = assignments, .pinned = pinned }; + } + + pub fn deinit(arena: *Arena) !void { + if (arena.active != null) return error.SharedKvBusy; + arena.allocator.free(arena.pinned); + arena.allocator.free(arena.assignments); + arena.allocator.free(arena.payload); + arena.allocator.free(arena.inverse); + arena.allocator.free(arena.stamps); + arena.allocator.free(arena.owners); + arena.allocator.free(arena.lanes); + arena.* = undefined; + } + + fn state(arena: *const Arena, sequence: Sequence) !*Lane { + if (sequence.arena != arena.id or sequence.lane >= arena.geometry.lanes or sequence.generation == 0) return error.SharedKvSequence; + const lane = &arena.lanes[sequence.lane]; + if (!lane.live or lane.generation != sequence.generation) return error.SharedKvSequence; + return lane; + } + + pub fn register(arena: *Arena, lane: u32) !Sequence { + if (arena.active != null) return error.SharedKvBusy; + if (lane >= arena.geometry.lanes) return error.SharedKvSequence; + const entry = &arena.lanes[lane]; + if (entry.live) return error.SharedKvSequenceLive; + const generation = try std.math.add(u64, entry.generation, 1); + entry.* = .{ .generation = generation, .live = true }; + return .{ .arena = arena.id, .lane = lane, .generation = generation }; + } + + fn owns(owner: Owner, sequence: Sequence) bool { + return owner.generation != 0 and owner.generation == sequence.generation and owner.lane == sequence.lane; + } + + fn inverseIndex(arena: *const Arena, lane: u32, key: Key) usize { + return (@as(usize, lane) * arena.geometry.layers + key.layer) * arena.geometry.blocks + key.block; + } + + fn evict(arena: *Arena, slot: u32) void { + const owner = arena.owners[slot]; + if (owner.generation != 0) { + const lane = arena.lanes[owner.lane]; + if (lane.live and lane.generation == owner.generation) { + const inverse = &arena.inverse[arena.inverseIndex(owner.lane, .{ .layer = owner.layer, .block = owner.block })]; + if (inverse.* == @as(i32, @intCast(slot))) inverse.* = -1; + } + } + arena.owners[slot] = .{}; + arena.stamps[slot] = 0; + } + + fn clear(arena: *Arena, sequence: Sequence) void { + for (arena.owners, 0..) |owner, slot| if (owns(owner, sequence)) arena.evict(@intCast(slot)); + const start = arena.inverseIndex(sequence.lane, .{ .layer = 0, .block = 0 }); + const length = @as(usize, arena.geometry.layers) * arena.geometry.blocks; + @memset(arena.inverse[start..][0..length], -1); + } + + pub fn reset(arena: *Arena, sequence: Sequence) !Sequence { + if (arena.active != null) return error.SharedKvBusy; + const lane = try arena.state(sequence); + const generation = try std.math.add(u64, lane.generation, 1); + arena.clear(sequence); + lane.* = .{ .generation = generation, .live = true }; + return .{ .arena = arena.id, .lane = sequence.lane, .generation = generation }; + } + + pub fn release(arena: *Arena, sequence: Sequence) !void { + if (arena.active != null) return error.SharedKvBusy; + const lane = try arena.state(sequence); + arena.clear(sequence); + lane.live = false; + } + + pub fn begin(arena: *Arena, sequence: Sequence) !Lease { + const lane = try arena.state(sequence); + if (arena.active != null) return error.SharedKvBusy; + if (lane.failed) return error.SharedKvSequenceFailed; + const serial = try std.math.add(u64, arena.serial, 1); + const lease: Lease = .{ .arena = arena.id, .serial = serial, .sequence = sequence }; + arena.serial = serial; + arena.active = .{ .lease = lease }; + return lease; + } + + fn held(arena: *Arena, lease: Lease) !*Active { + const active = if (arena.active) |*active| active else return error.SharedKvLease; + if (!std.meta.eql(active.lease, lease)) return error.SharedKvLease; + return active; + } + + pub fn slotFor(arena: *const Arena, sequence: Sequence, key: Key) !?u32 { + _ = try arena.state(sequence); + if (key.layer >= arena.geometry.layers or key.block >= arena.geometry.blocks) return error.SharedKvKey; + const inverse = arena.inverse[arena.inverseIndex(sequence.lane, key)]; + if (inverse < 0) return null; + const slot: u32 = @intCast(inverse); + if (slot >= arena.geometry.slots) return error.SharedKvInverse; + const owner = arena.owners[slot]; + if (!owns(owner, sequence) or owner.layer != key.layer or owner.block != key.block or owner.valid == 0) return error.SharedKvInverse; + return slot; + } + + pub fn prepare(arena: *Arena, lease: Lease, keys: []const Key) ![]const Assignment { + return arena.prepareEntries(lease, keys); + } + + fn prepareEntries(arena: *Arena, lease: Lease, entries: anytype) ![]const Assignment { + const active = try arena.held(lease); + if (active.phase != .acquired) return error.SharedKvPhase; + @memset(arena.pinned, false); + var count: usize = 0; + for (entries) |entry| { + const key = if (@TypeOf(entry) == Key) entry else entry.key; + if (key.layer >= arena.geometry.layers or key.block >= arena.geometry.blocks) return error.SharedKvKey; + var duplicate = false; + for (arena.assignments[0..count]) |previous| if (std.meta.eql(previous.key, key)) { + duplicate = true; + break; + }; + if (duplicate) continue; + if (count == arena.assignments.len) return error.SharedKvCapacity; + const resident = try arena.slotFor(lease.sequence, key); + arena.assignments[count] = .{ .key = key, .slot = resident orelse 0, .hit = resident != null }; + if (resident) |slot| arena.pinned[slot] = true; + count += 1; + } + for (arena.assignments[0..count]) |*assignment| { + if (assignment.hit) continue; + var victim: ?u32 = null; + for (arena.pinned, 0..) |pinned, slot| { + if (pinned) continue; + if (victim == null or arena.stamps[slot] < arena.stamps[victim.?]) victim = @intCast(slot); + } + const slot = victim orelse return error.SharedKvCapacity; + assignment.slot = slot; + arena.pinned[slot] = true; + } + active.count = count; + active.phase = .planned; + return arena.assignments[0..count]; + } + + pub fn cancel(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .acquired and active.phase != .planned) return error.SharedKvPending; + arena.active = null; + } + + pub fn submit(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .planned) return error.SharedKvPhase; + const epoch = try std.math.add(u64, arena.clock, 1); + for (arena.assignments[0..active.count]) |assignment| { + if (!assignment.hit) { + arena.evict(assignment.slot); + arena.owners[assignment.slot] = .{ .generation = lease.sequence.generation, .lane = lease.sequence.lane, .layer = assignment.key.layer, .block = assignment.key.block }; + } + arena.stamps[assignment.slot] = epoch; + } + arena.clock = epoch; + active.phase = .submitted; + } + + fn block(arena: *Arena, slot: u32) []u8 { + const size = arena.geometry.blockBytes() catch unreachable; + return arena.payload[@as(usize, slot) * size ..][0..size]; + } + + pub fn writeFetched(arena: *Arena, lease: Lease, index: usize, payload: []const u8) !void { + const active = try arena.held(lease); + if (active.phase != .submitted) return error.SharedKvPhase; + if (index >= active.count or payload.len != try arena.geometry.blockBytes()) return error.SharedKvPayload; + const assignment = &arena.assignments[index]; + @memcpy(arena.block(assignment.slot), payload); + assignment.written = true; + } + + pub fn completeFetch(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .submitted) return error.SharedKvPhase; + for (arena.assignments[0..active.count]) |assignment| if (!assignment.hit and !assignment.written) return error.SharedKvFetchIncomplete; + for (arena.assignments[0..active.count]) |assignment| { + arena.owners[assignment.slot].valid = 1; + arena.inverse[arena.inverseIndex(lease.sequence.lane, assignment.key)] = @intCast(assignment.slot); + } + active.phase = .ready; + } + + pub fn read(arena: *Arena, lease: Lease, key: Key) ![]const u8 { + const active = try arena.held(lease); + if (active.phase != .ready) return error.SharedKvPhase; + const slot = try arena.slotFor(lease.sequence, key) orelse return error.SharedKvMissing; + return arena.block(slot); + } + + pub fn completeConsumers(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .ready) return error.SharedKvPhase; + arena.active = null; + } + + pub fn failSubmitted(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .submitted and active.phase != .ready) return error.SharedKvPhase; + active.phase = .poisoned; + arena.lanes[lease.sequence.lane].failed = true; + } + + pub fn recoverAfterDrain(arena: *Arena, lease: Lease) !void { + const active = try arena.held(lease); + if (active.phase != .poisoned) return error.SharedKvPhase; + arena.clear(lease.sequence); + arena.active = null; + } + + pub fn graphKey(arena: *Arena, sequence: Sequence, bindings: Bindings) !GraphKey { + _ = try arena.state(sequence); + return .{ .sequence = sequence, .geometry = arena.geometry, .bindings = bindings }; + } + + pub fn validate(arena: *const Arena) !void { + for (arena.owners, 0..) |owner, slot| { + if (owner.generation == 0) { + if (owner.valid != 0 or arena.stamps[slot] != 0) return error.SharedKvInverse; + continue; + } + if (owner.lane >= arena.geometry.lanes or owner.layer >= arena.geometry.layers or owner.block >= arena.geometry.blocks or arena.stamps[slot] == 0 or arena.stamps[slot] > arena.clock) return error.SharedKvInverse; + const lane = arena.lanes[owner.lane]; + if (!lane.live or lane.generation != owner.generation) return error.SharedKvInverse; + const inverse = arena.inverse[arena.inverseIndex(owner.lane, .{ .layer = owner.layer, .block = owner.block })]; + if ((owner.valid != 0 and inverse != @as(i32, @intCast(slot))) or (owner.valid == 0 and inverse != -1)) return error.SharedKvInverse; + } + for (arena.lanes, 0..) |lane, lane_index| { + for (0..arena.geometry.layers) |layer| for (0..arena.geometry.blocks) |logical_block| { + const key: Key = .{ .layer = @intCast(layer), .block = @intCast(logical_block) }; + const inverse = arena.inverse[arena.inverseIndex(@intCast(lane_index), key)]; + if (inverse == -1) continue; + if (!lane.live or inverse < 0 or inverse >= arena.geometry.slots) return error.SharedKvInverse; + const owner = arena.owners[@intCast(inverse)]; + if (owner.valid == 0 or owner.lane != lane_index or owner.generation != lane.generation or owner.layer != layer or owner.block != logical_block) return error.SharedKvInverse; + }; + } + } + + pub fn saveOwned(arena: *Arena, allocator: std.mem.Allocator, lease: Lease) !OwnedBlocks { + const active = try arena.held(lease); + if (active.phase != .acquired) return error.SharedKvPhase; + var count: usize = 0; + for (arena.owners) |owner| if (owns(owner, lease.sequence) and owner.valid != 0) { + count += 1; + }; + const size = try arena.geometry.blockBytes(); + const entries = try allocator.alloc(OwnedEntry, count); + errdefer allocator.free(entries); + const payload = try allocator.alloc(u8, try std.math.mul(usize, count, size)); + errdefer allocator.free(payload); + var position: usize = 0; + for (arena.owners, 0..) |owner, slot| { + if (!owns(owner, lease.sequence) or owner.valid == 0) continue; + const key: Key = .{ .layer = owner.layer, .block = owner.block }; + if ((try arena.slotFor(lease.sequence, key)) != @as(u32, @intCast(slot))) return error.SharedKvInverse; + entries[position] = .{ .key = key, .stamp = arena.stamps[slot], .source_slot = @intCast(slot) }; + @memcpy(payload[position * size ..][0..size], arena.block(@intCast(slot))); + position += 1; + } + return .{ .source = lease.sequence, .geometry = arena.geometry, .clock = arena.clock, .entries = entries, .payload = payload }; + } + + pub fn restoreOwned(arena: *Arena, lease: Lease, image: *const OwnedBlocks) !void { + const active = try arena.held(lease); + if (active.phase != .acquired) return error.SharedKvPhase; + try image.validate(arena.geometry); + const next_clock = try std.math.add(u64, try std.math.add(u64, arena.clock, 1), try image.maxAge()); + _ = try arena.prepareEntries(lease, image.entries); + try arena.submit(lease); + errdefer arena.failSubmitted(lease) catch unreachable; + for (arena.owners, 0..) |owner, slot| { + if (!owns(owner, lease.sequence)) continue; + var retained = false; + for (image.entries) |entry| if (entry.key.layer == owner.layer and entry.key.block == owner.block) { + retained = true; + break; + }; + if (!retained) arena.evict(@intCast(slot)); + } + const size = try arena.geometry.blockBytes(); + for (image.entries, 0..) |entry, index| { + try arena.writeFetched(lease, index, image.payload[index * size ..][0..size]); + arena.stamps[arena.assignments[index].slot] = next_clock - (image.clock - entry.stamp); + } + arena.clock = next_clock; + try arena.completeFetch(lease); + } +}; + +test { + _ = @import("exl3_shared_kv_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3_shared_kv_test.zig b/zig/src/families/flashnext/exl3_shared_kv_test.zig new file mode 100644 index 0000000000..26636cd314 --- /dev/null +++ b/zig/src/families/flashnext/exl3_shared_kv_test.zig @@ -0,0 +1,341 @@ +const std = @import("std"); +const shared = @import("exl3_shared_kv.zig"); + +const geometry: shared.Geometry = .{ .lanes = 3, .layers = 2, .blocks = 8, .slots = 3, .row_bytes = 4, .window_rows = 4, .work_words = 8 }; +const first: shared.Key = .{ .layer = 0, .block = 0 }; +const second: shared.Key = .{ .layer = 1, .block = 0 }; +const third: shared.Key = .{ .layer = 0, .block = 1 }; +const fourth: shared.Key = .{ .layer = 1, .block = 1 }; + +fn store(arena: *shared.Arena, sequence: shared.Sequence, keys: []const shared.Key, value: u8) !void { + const lease = try arena.begin(sequence); + const assignments = try arena.prepare(lease, keys); + try arena.submit(lease); + var payload: [32]u8 = @splat(value); + for (assignments, 0..) |assignment, index| if (!assignment.hit) try arena.writeFetched(lease, index, &payload); + try arena.completeFetch(lease); + try arena.validate(); + try arena.completeConsumers(lease); +} + +test "exl3: shared KV cross-lane eviction invalidates the evicted generation and layer" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence_a = try arena.register(0); + const sequence_b = try arena.register(1); + const sequence_c = try arena.register(2); + try store(&arena, sequence_a, &.{first}, 0x11); + try store(&arena, sequence_b, &.{first}, 0x22); + try store(&arena, sequence_a, &.{second}, 0x33); + const other_slot = (try arena.slotFor(sequence_b, first)).?; + try store(&arena, sequence_c, &.{first}, 0x44); + try std.testing.expectEqual(null, try arena.slotFor(sequence_a, first)); + try std.testing.expectEqual(other_slot, (try arena.slotFor(sequence_b, first)).?); + try std.testing.expect((try arena.slotFor(sequence_a, second)) != null); + try store(&arena, sequence_a, &.{ second, third }, 0x55); + try std.testing.expectEqual(null, try arena.slotFor(sequence_b, first)); + try std.testing.expect((try arena.slotFor(sequence_c, first)) != null); + try arena.validate(); +} + +test "exl3: shared KV reset release and lane reuse never alias old generations" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const previous = try arena.register(0); + const other = try arena.register(1); + try store(&arena, previous, &.{ first, second }, 0xa1); + try store(&arena, other, &.{first}, 0xb2); + const foreign_slot = (try arena.slotFor(other, first)).?; + const foreign_bytes = try std.testing.allocator.dupe(u8, arena.payload); + defer std.testing.allocator.free(foreign_bytes); + const current = try arena.reset(previous); + try std.testing.expect(current.generation > previous.generation); + try std.testing.expectError(error.SharedKvSequence, arena.begin(previous)); + try std.testing.expectError(error.SharedKvSequence, arena.release(previous)); + try std.testing.expectEqual(null, try arena.slotFor(current, first)); + try std.testing.expectEqual(foreign_slot, (try arena.slotFor(other, first)).?); + try std.testing.expectEqualSlices(u8, foreign_bytes, arena.payload); + try store(&arena, current, &.{first}, 0xc3); + try arena.release(current); + const reused = try arena.register(0); + try std.testing.expect(reused.generation > current.generation); + try std.testing.expectEqual(null, try arena.slotFor(reused, first)); + try std.testing.expectEqual(foreign_slot, (try arena.slotFor(other, first)).?); + try arena.validate(); +} + +test "exl3: shared KV planning pins later hits deduplicates and cancels without mutation" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + try store(&arena, sequence, &.{ first, second, third }, 0x15); + const hit_slot = (try arena.slotFor(sequence, first)).?; + const lease = try arena.begin(sequence); + const plan = try arena.prepare(lease, &.{ fourth, first, fourth }); + try std.testing.expectEqual(@as(usize, 2), plan.len); + try std.testing.expect(plan[0].slot != hit_slot); + try std.testing.expect(plan[1].hit); + try std.testing.expectEqual(hit_slot, plan[1].slot); + const clock = arena.clock; + try arena.cancel(lease); + try std.testing.expectEqual(clock, arena.clock); + try std.testing.expectEqual(null, try arena.slotFor(sequence, fourth)); + try std.testing.expect((try arena.slotFor(sequence, second)) != null); + try arena.validate(); + const another = try arena.begin(sequence); + try std.testing.expectError(error.SharedKvCapacity, arena.prepare(another, &.{ first, second, third, fourth })); + try std.testing.expectError(error.SharedKvKey, arena.prepare(another, &.{.{ .layer = 2, .block = 0 }})); + try arena.cancel(another); + try arena.validate(); +} + +test "exl3: shared KV lease spans fetch and consumers and cannot be released by stale tickets" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + const other = try arena.register(1); + const lease = try arena.begin(sequence); + _ = try arena.prepare(lease, &.{first}); + try arena.submit(lease); + try std.testing.expectError(error.SharedKvBusy, arena.begin(other)); + try std.testing.expectError(error.SharedKvBusy, arena.reset(other)); + try std.testing.expectError(error.SharedKvBusy, arena.release(sequence)); + try std.testing.expectError(error.SharedKvBusy, arena.deinit()); + try std.testing.expectError(error.SharedKvPending, arena.cancel(lease)); + try std.testing.expectError(error.SharedKvFetchIncomplete, arena.completeFetch(lease)); + try std.testing.expectEqual(null, try arena.slotFor(sequence, first)); + try arena.validate(); + const payload: [32]u8 = @splat(0x65); + try std.testing.expectError(error.SharedKvPayload, arena.writeFetched(lease, 0, payload[0..31])); + try arena.writeFetched(lease, 0, &payload); + try arena.completeFetch(lease); + try std.testing.expectEqualSlices(u8, &payload, try arena.read(lease, first)); + try std.testing.expectError(error.SharedKvBusy, arena.begin(other)); + try arena.completeConsumers(lease); + const next = try arena.begin(other); + try std.testing.expectError(error.SharedKvLease, arena.cancel(lease)); + try std.testing.expectError(error.SharedKvLease, arena.completeConsumers(lease)); + try arena.cancel(next); +} + +test "exl3: shared KV submission and consumer failures retain ownership until drained" { + for ([_]bool{ false, true }) |consumer_failure| { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + const other = try arena.register(1); + try store(&arena, other, &.{first}, 0x12); + const foreign = (try arena.slotFor(other, first)).?; + const lease = try arena.begin(sequence); + _ = try arena.prepare(lease, &.{second}); + try arena.submit(lease); + const payload: [32]u8 = @splat(0x78); + try arena.writeFetched(lease, 0, &payload); + if (consumer_failure) try arena.completeFetch(lease); + try arena.failSubmitted(lease); + try std.testing.expectError(error.SharedKvBusy, arena.begin(other)); + try std.testing.expectError(error.SharedKvBusy, arena.reset(sequence)); + try std.testing.expectError(error.SharedKvPending, arena.cancel(lease)); + try std.testing.expectError(error.SharedKvPhase, arena.completeConsumers(lease)); + try arena.recoverAfterDrain(lease); + try std.testing.expectError(error.SharedKvSequenceFailed, arena.begin(sequence)); + try std.testing.expectEqual(null, try arena.slotFor(sequence, second)); + try std.testing.expectEqual(foreign, (try arena.slotFor(other, first)).?); + const reset = try arena.reset(sequence); + try store(&arena, reset, &.{second}, 0x23); + try arena.validate(); + } +} + +test "exl3: shared KV owned-entry snapshots isolate foreign bytes and remap to a new generation" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const source = try arena.register(0); + const other = try arena.register(1); + try store(&arena, source, &.{first}, 0xa1); + try store(&arena, source, &.{second}, 0xa2); + const saving = try arena.begin(source); + var image = try arena.saveOwned(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + try std.testing.expectEqual(@as(usize, 2), image.entries.len); + try arena.release(source); + try store(&arena, other, &.{ first, third }, 0xb1); + const destination = try arena.register(2); + const foreign = (try arena.slotFor(other, third)).?; + const size = try geometry.blockBytes(); + const foreign_bytes = try std.testing.allocator.dupe(u8, arena.payload[@as(usize, foreign) * size ..][0..size]); + defer std.testing.allocator.free(foreign_bytes); + const restoring = try arena.begin(destination); + try arena.restoreOwned(restoring, &image); + for (image.entries, 0..) |entry, index| { + try std.testing.expectEqualSlices(u8, image.payload[index * size ..][0..size], try arena.read(restoring, entry.key)); + const slot = (try arena.slotFor(destination, entry.key)).?; + try std.testing.expectEqual(destination.generation, arena.owners[slot].generation); + try std.testing.expectEqual(destination.lane, arena.owners[slot].lane); + } + try std.testing.expectEqual(foreign, (try arena.slotFor(other, third)).?); + try std.testing.expectEqualSlices(u8, foreign_bytes, arena.payload[@as(usize, foreign) * size ..][0..size]); + try arena.completeConsumers(restoring); + try arena.validate(); + const first_slot = (try arena.slotFor(destination, first)).?; + const second_slot = (try arena.slotFor(destination, second)).?; + try std.testing.expect(arena.stamps[first_slot] < arena.stamps[second_slot]); + try std.testing.expect(first_slot != image.entries[0].source_slot or second_slot != image.entries[1].source_slot); + const image_copy = try std.testing.allocator.dupe(u8, image.payload); + defer std.testing.allocator.free(image_copy); + try store(&arena, destination, &.{third}, 0xd3); + const again = try arena.begin(destination); + try arena.restoreOwned(again, &image); + try std.testing.expectEqual(null, try arena.slotFor(destination, third)); + try arena.completeConsumers(again); + try std.testing.expectEqualSlices(u8, image_copy, image.payload); + try arena.validate(); +} + +test "exl3: shared KV empty owned snapshot clears only destination residency" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + const other = try arena.register(1); + const saving = try arena.begin(sequence); + var image = try arena.saveOwned(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + try store(&arena, sequence, &.{first}, 0x13); + try store(&arena, other, &.{second}, 0x24); + const restoring = try arena.begin(sequence); + try arena.restoreOwned(restoring, &image); + try arena.completeConsumers(restoring); + try std.testing.expectEqual(null, try arena.slotFor(sequence, first)); + try std.testing.expect((try arena.slotFor(other, second)) != null); + try arena.validate(); +} + +test "exl3: shared KV snapshot malformed versions geometry lengths duplicates and overflow are preflight errors" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + try store(&arena, sequence, &.{ first, second }, 0x13); + const saving = try arena.begin(sequence); + var image = try arena.saveOwned(std.testing.allocator, saving); + defer image.deinit(std.testing.allocator); + try arena.cancel(saving); + const restoring = try arena.begin(sequence); + var bad = image; + bad.version = 1; + try std.testing.expectError(error.SharedKvSnapshotVersion, arena.restoreOwned(restoring, &bad)); + bad = image; + bad.geometry.row_bytes *= 2; + try std.testing.expectError(error.SharedKvSnapshotGeometry, arena.restoreOwned(restoring, &bad)); + bad = image; + bad.payload = bad.payload[1..]; + try std.testing.expectError(error.SharedKvSnapshotLength, arena.restoreOwned(restoring, &bad)); + const original = image.entries[1]; + image.entries[1] = image.entries[0]; + try std.testing.expectError(error.SharedKvSnapshotDuplicate, arena.restoreOwned(restoring, &image)); + image.entries[1] = original; + image.entries[1].stamp = image.clock + 1; + try std.testing.expectError(error.SharedKvSnapshotEntry, arena.restoreOwned(restoring, &image)); + image.entries[1] = original; + arena.clock = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, arena.restoreOwned(restoring, &image)); + try std.testing.expectEqual(shared.Phase.acquired, arena.active.?.phase); + try arena.cancel(restoring); + try arena.validate(); +} + +test "exl3: shared KV byte accounting charges once and rejects overflow and non-BF16 rows" { + const rank: shared.Geometry = .{ .lanes = 8, .layers = 12, .blocks = 50005, .slots = 16384, .row_bytes = 1024, .window_rows = 6148, .work_words = 4 * 256 * 512 }; + const bytes = try rank.accounting(); + try std.testing.expectEqual(@as(usize, 128 * 1024 * 1024), bytes.payload); + try std.testing.expectEqual(bytes.shared + bytes.sequence * 8, bytes.rank); + try std.testing.expectEqual(@as(usize, 940441600), try rank.duplicatedLegacyBytes()); + try std.testing.expectEqual(@as(usize, 24), @sizeOf(shared.Owner)); + var bad = rank; + bad.row_bytes = std.math.maxInt(usize) - 1; + try std.testing.expectError(error.Overflow, bad.accounting()); + bad = rank; + bad.window_rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, bad.accounting()); + bad = rank; + bad.work_words = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, bad.accounting()); + bad = rank; + bad.row_bytes = 3; + try std.testing.expectError(error.SharedKvGeometry, bad.accounting()); + bad = rank; + bad.slots = std.math.maxInt(u32); + try std.testing.expectError(error.SharedKvGeometry, bad.accounting()); +} + +test "exl3: shared KV graph keys distinguish bindings geometry arena and reset generations" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + const bindings: shared.Bindings = .{ .keys = 1, .values = 2, .owners = 3, .stamps = 4, .clock = 5, .inverse = 6, .marks = 7, .controls = 8, .work = 9, .window_keys = 10, .window_values = 11, .canonical = 12, .cap = 32, .rows = 8, .origin_hash = 13 }; + const original = try arena.graphKey(sequence, bindings); + inline for (@typeInfo(shared.Bindings).@"struct".field_names) |field| { + var changed = bindings; + @field(changed, field) += 1; + try std.testing.expect(!std.meta.eql(original, try arena.graphKey(sequence, changed))); + } + inline for (@typeInfo(shared.Geometry).@"struct".field_names) |field| { + var changed = original; + @field(changed.geometry, field) += 1; + try std.testing.expect(!std.meta.eql(original, changed)); + } + try store(&arena, sequence, &.{first}, 0x89); + try std.testing.expectEqualDeep(original, try arena.graphKey(sequence, bindings)); + const reset = try arena.reset(sequence); + try std.testing.expect(!std.meta.eql(original, try arena.graphKey(reset, bindings))); + var another = try shared.Arena.init(std.testing.allocator, geometry); + defer another.deinit() catch unreachable; + const other = try another.register(0); + try std.testing.expect(!std.meta.eql(original, try another.graphKey(other, bindings))); + try std.testing.expectError(error.SharedKvSequence, another.begin(reset)); +} + +test "exl3: shared KV generation lease serial and clock exhaustion fail without wraparound" { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + var sequence = try arena.register(0); + arena.lanes[0].generation = std.math.maxInt(u64); + sequence.generation = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, arena.reset(sequence)); + try arena.release(sequence); + try std.testing.expectError(error.Overflow, arena.register(0)); + const other = try arena.register(1); + arena.serial = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, arena.begin(other)); + arena.serial = 0; + const lease = try arena.begin(other); + _ = try arena.prepare(lease, &.{first}); + arena.clock = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, arena.submit(lease)); + try std.testing.expectEqual(shared.Phase.planned, arena.active.?.phase); + try arena.cancel(lease); + try arena.validate(); +} + +fn allocationCase(allocator: std.mem.Allocator) !void { + var arena = try shared.Arena.init(allocator, geometry); + defer arena.deinit() catch unreachable; +} + +fn snapshotAllocationCase(allocator: std.mem.Allocator) !void { + var arena = try shared.Arena.init(std.testing.allocator, geometry); + defer arena.deinit() catch unreachable; + const sequence = try arena.register(0); + try store(&arena, sequence, &.{first}, 0x47); + const lease = try arena.begin(sequence); + defer arena.cancel(lease) catch unreachable; + var image = try arena.saveOwned(allocator, lease); + defer image.deinit(allocator); + try image.validate(geometry); +} + +test "exl3: shared KV allocation and snapshot allocation failures retain lease and clean ownership" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, allocationCase, .{}); + try std.testing.checkAllAllocationFailures(std.testing.allocator, snapshotAllocationCase, .{}); +} diff --git a/zig/src/families/flashnext/exl3_startup_memory.zig b/zig/src/families/flashnext/exl3_startup_memory.zig new file mode 100644 index 0000000000..d31f38ba03 --- /dev/null +++ b/zig/src/families/flashnext/exl3_startup_memory.zig @@ -0,0 +1,191 @@ +const std = @import("std"); +const planner_layout = @import("exl3_dplan_layout.zig"); +const demand_trace = @import("exl3_dplan_trace.zig"); + +pub const Reservation = struct { + stream_bytes: u64, + streams: u64, + graph_bytes: u64, + after_pool_bytes: u64 = 0, + + pub fn total(self: Reservation) !u64 { + return std.math.add(u64, try std.math.add(u64, try std.math.mul(u64, self.stream_bytes, self.streams), self.graph_bytes), self.after_pool_bytes); + } +}; + +pub const Ledger = struct { + persistent: u64, + scratch: u64, + deep_prefill: u64, + pool_scratch: u64, + expert_pools: u64, + reservation: Reservation, + + pub fn total(self: Ledger) !u64 { + var bytes = try self.reservation.total(); + inline for (.{ self.persistent, self.scratch, self.deep_prefill, self.pool_scratch, self.expert_pools }) |part| bytes = try std.math.add(u64, bytes, part); + return bytes; + } +}; + +pub fn mebibytes(text: []const u8) !u64 { + return std.math.mul(u64, try std.fmt.parseInt(u64, text, 10), 1 << 20); +} + +pub fn graphReservation(enabled: bool, text: ?[]const u8) !u64 { + if (!enabled) return 0; + const bytes = try mebibytes(text orelse return error.GraphReservationRequired); + if (bytes == 0) return error.GraphReservationRequired; + return bytes; +} + +pub fn available(free: u64, reserve: u64, future: u64) !u64 { + const held = try std.math.add(u64, reserve, future); + if (held > free) return error.InsufficientStartupMemory; + return free - held; +} + +pub fn laneRoom(local: u64, graph_reservation: u64) u64 { + return local -| graph_reservation; +} + +pub fn poolBytes(lines: u64, line_bytes: u64, layers: u64, fixed: u64) !u64 { + return std.math.add(u64, try std.math.mul(u64, try std.math.mul(u64, lines, line_bytes), layers), fixed); +} + +pub fn plannerBytes(layers: u64, lines: u64, rows: u64) !u64 { + const state = try planner_layout.deviceBytes(layers, try std.math.mul(u64, layers, lines), rows); + return std.math.add(u64, state, try plannerTraceBytes(layers, traceEnabled())); +} + +pub fn plannerBytesForPools(lines: []const usize, rows: u64) !u64 { + var total_lines: u64 = 0; + for (lines) |count| total_lines = try std.math.add(u64, total_lines, count); + const state = try planner_layout.deviceBytes(lines.len, total_lines, rows); + return std.math.add(u64, state, try plannerTraceBytes(lines.len, traceEnabled())); +} + +pub fn plannerTraceBytes(pools: u64, enabled: bool) !u64 { + return if (enabled) std.math.mul(u64, pools, demand_trace.pool_bytes) else 0; +} + +fn traceEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_DPLAN_TRACE") orelse return false; + return demand_trace.enabledValue(std.mem.span(value)); +} + +test "BF16 c1 c8 and graph on off preserve the requested stream bytes" { + const bf16_stream: u64 = 262144 * 256 * 2; + for ([_]u64{ 1, 8 }) |streams| { + for ([_]bool{ false, true }) |graphs| { + const graph = try graphReservation(graphs, "512"); + const reservation: Reservation = .{ .stream_bytes = bf16_stream, .streams = streams, .graph_bytes = graph, .after_pool_bytes = 4096 }; + try std.testing.expectEqual(bf16_stream * streams + graph + 4096, try reservation.total()); + const free = try available(16 << 30, 1 << 30, try reservation.total()); + try std.testing.expectEqual(@as(u64, 16 << 30), free + (1 << 30) + try reservation.total()); + } + } +} + +test "startup accounting fails closed on overflow and underfunded explicit pools" { + try std.testing.expectError(error.Overflow, (Reservation{ .stream_bytes = std.math.maxInt(u64), .streams = 8, .graph_bytes = 0 }).total()); + try std.testing.expectError(error.Overflow, mebibytes("18446744073709551615")); + try std.testing.expectError(error.GraphReservationRequired, graphReservation(true, null)); + try std.testing.expectError(error.GraphReservationRequired, graphReservation(true, "0")); + try std.testing.expectEqual(@as(u64, 0), try graphReservation(false, null)); + try std.testing.expectError(error.Overflow, poolBytes(std.math.maxInt(u64), 2, 49, 0)); + try std.testing.expectError(error.InsufficientStartupMemory, available(100, 20, try poolBytes(10, 1, 8, 1))); + try std.testing.expectError(error.Overflow, available(100, std.math.maxInt(u64), 1)); +} + +test "ledger includes all allocations after the pool checkpoint" { + const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .deep_prefill = 30, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8, .graph_bytes = 60, .after_pool_bytes = 70 } }; + try std.testing.expectEqual(@as(u64, 450), try ledger.total()); +} + +test "pool setup may consume its allocator margin but not lane graph or server reservations" { + const server: u64 = 100; + const margin: u64 = 25; + const future: u64 = 80; + const planner: u64 = 10; + const planned_free = server + margin + future + planner; + try std.testing.expectEqual(margin, try available(planned_free, server, future + planner)); + try std.testing.expectEqual(@as(u64, 0), try available(planned_free - margin, server, future + planner)); + try std.testing.expectError(error.InsufficientStartupMemory, available(planned_free - margin - 1, server, future + planner)); +} + +test "planner allocation ledger agrees with device state statics work and split buffers" { + const layers: u64 = 49; + const lines: u64 = 10; + const rows: u64 = 16; + const expected = layers * planner_layout.stateWords(lines) * 4 + layers * 512 * 3 * (8 + 8 + 4) + rows * 10 * 9 * 8 + 256 + (512 + 64) * 4 + 2 * rows * 11 * 4; + try std.testing.expectEqual(expected, try plannerBytes(layers, lines, rows)); + try std.testing.expectEqual(try plannerBytes(layers, 0, rows) + layers * lines * 8, try plannerBytes(layers, lines, rows)); + try std.testing.expectError(error.Overflow, plannerBytes(layers, std.math.maxInt(u64), rows)); +} + +test "opt-in demand traces are included in the startup planner reservation" { + try std.testing.expectEqual(@as(u64, 0), try plannerTraceBytes(48, false)); + try std.testing.expectEqual(@as(u64, 48 << 20), try plannerTraceBytes(48, true)); + const base = try planner_layout.deviceBytes(48, 0, 16); + try std.testing.expectEqual(base + (48 << 20), try std.math.add(u64, base, try plannerTraceBytes(48, true))); +} + +test "exl3: ragged planner redistribution conserves each rank's exact device budget" { + const uniform = [_]usize{ 32, 32, 32, 32 }; + const ragged = [_]usize{ 10, 73, 11, 34 }; + const other_rank = [_]usize{ 12, 28, 52, 36 }; + for ([_][]const usize{ &uniform, &ragged, &other_rank }) |counts| { + var state_words: usize = 0; + for (counts) |lines| { + const entry = try planner_layout.PoolState.init(lines, state_words); + state_words = entry.endWords(); + } + const rows = 16; + const work = rows * 10 * 9 * 8 + 256 + (512 + 64) * 4; + const split = 2 * rows * 11 * 4; + const allocated = state_words * 4 + counts.len * planner_layout.static_bytes + work + split; + try std.testing.expectEqual(allocated, try plannerBytesForPools(counts, rows)); + try std.testing.expectEqual(try plannerBytes(uniform.len, uniform[0], rows), try plannerBytesForPools(counts, rows)); + const line_bytes = 3 << 20; + const fixed = try plannerBytesForPools(counts, rows); + var total = fixed; + for (counts) |lines| total += lines * line_bytes; + try std.testing.expectEqual(try poolBytes(32, line_bytes, 4, fixed), total); + } + try std.testing.expectError(error.Overflow, plannerBytesForPools(&.{ std.math.maxInt(usize), std.math.maxInt(usize) }, 16)); + try std.testing.expectError(error.Overflow, plannerBytesForPools(&ragged, std.math.maxInt(u64))); +} + +test "lane admission protects the full graph reservation without changing BF16 stream sizes" { + const stream_bytes: u64 = 262144 * 256 * 2; + const graph_bytes: u64 = 512 << 20; + for ([_]u64{ 1, 8 }) |streams| { + const requested = try std.math.mul(u64, streams, stream_bytes); + try std.testing.expectEqual(requested, laneRoom(requested, 0)); + try std.testing.expectEqual(requested, laneRoom(try std.math.add(u64, requested, graph_bytes), graph_bytes)); + try std.testing.expect(laneRoom(requested, graph_bytes) < requested); + } + try std.testing.expectEqual(@as(u64, 0), laneRoom(graph_bytes - 1, graph_bytes)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(graph_bytes, graph_bytes)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(0, graph_bytes)); +} + +test "lane room subtraction cannot overflow at integer boundaries" { + const maximum = std.math.maxInt(u64); + try std.testing.expectEqual(maximum, laneRoom(maximum, 0)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(maximum, maximum)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(1, maximum)); + try std.testing.expectEqual(maximum - 1, laneRoom(maximum, 1)); + try std.testing.expectError(error.Overflow, (Reservation{ .stream_bytes = maximum, .streams = 1, .graph_bytes = 1 }).total()); +} + +test "per rank protection precedes agreement of lane counts" { + const graph_bytes: u64 = 512 << 20; + const rank_zero_stream: u64 = 128 << 20; + const rank_one_stream: u64 = 64 << 20; + const rank_zero_room = laneRoom(8 * rank_zero_stream + graph_bytes, graph_bytes); + const rank_one_room = laneRoom(8 * rank_one_stream + graph_bytes, graph_bytes); + try std.testing.expectEqual(@as(u64, 8), @min(rank_zero_room / rank_zero_stream, rank_one_room / rank_one_stream)); + try std.testing.expectEqual(@as(u64, 7), @min(rank_zero_room / rank_zero_stream, laneRoom(8 * rank_one_stream + graph_bytes - 1, graph_bytes) / rank_one_stream)); +} diff --git a/zig/src/families/flashnext/exl3_state_binding.zig b/zig/src/families/flashnext/exl3_state_binding.zig new file mode 100644 index 0000000000..c1ac69aa12 --- /dev/null +++ b/zig/src/families/flashnext/exl3_state_binding.zig @@ -0,0 +1,177 @@ +const std = @import("std"); + +pub const max_layers = 48; +pub const max_width = 16; + +pub fn supportsWindow(width: usize) bool { + return width >= 2 and width <= max_width; +} + +pub const Mapping = struct { layer: usize, state_in: u64, state_out: u64 }; +pub const Geometry = struct { + grid: [3]u32, + block: [3]u32, + shared: u32 = 0, +}; +pub const Nonstate = struct { gq: u64, gk: u64, gv: u64, ggb: u64, rows: c_int, ys: u64 }; +pub const Request = struct { + width: usize, + chain: usize, + chain_to: usize, + geometry: Geometry, + nonstate: Nonstate, + layers: [max_layers]Mapping = undefined, + count: usize = 0, + + pub fn validate(self: Request) !void { + if (self.width < 1 or self.width > max_width or self.nonstate.rows != self.width or + self.chain == 0 or self.chain_to == 0 or self.chain == self.chain_to or + self.count == 0 or self.count > max_layers) return error.UnsupportedBinding; + for (self.layers[0..self.count], 0..) |mapping, index| { + if (mapping.state_in == 0 or mapping.state_in == mapping.state_out) return error.UnsupportedBinding; + for (self.layers[0..index]) |previous| { + if (previous.layer == mapping.layer or previous.state_in == mapping.state_in or + previous.state_out == mapping.state_in or + (mapping.state_out != 0 and (previous.state_in == mapping.state_out or previous.state_out == mapping.state_out))) + return error.UnsupportedBinding; + } + } + } + + pub fn compatible(self: Request, target: Request) bool { + target.validate() catch return false; + if (self.width != target.width or self.chain != target.chain or self.chain_to != target.chain_to or + !std.meta.eql(self.geometry, target.geometry) or !std.meta.eql(self.nonstate, target.nonstate) or + self.count != target.count) return false; + for (self.layers[0..self.count], target.layers[0..target.count]) |before, after| { + if (before.layer != after.layer) return false; + const same = before.state_in == after.state_in and before.state_out == after.state_out; + const flipped = before.state_out != 0 and before.state_in == after.state_out and before.state_out == after.state_in; + if (!same and !flipped) return false; + } + return true; + } +}; + +pub const Header = struct { + function: usize, + geometry: Geometry, + params: ?[*]?*anyopaque, + extra: bool = false, + alternate_kernel: bool = false, +}; + +pub const Parameters = struct { + nonstate: Nonstate, + state_in: u64, + state_out: u64, +}; + +fn argument(comptime Value: type, params: [*]?*anyopaque, index: usize) !Value { + const pointer = params[index] orelse return error.UnsupportedBinding; + var value: Value = undefined; + @memcpy(std.mem.asBytes(&value), @as([*]const u8, @ptrCast(pointer))[0..@sizeOf(Value)]); + return value; +} + +pub const Binding = struct { + template: usize, + request: Request, + nodes: [max_layers]usize = undefined, + parameters: [max_layers]Parameters = undefined, + loaded: bool = false, + poisoned: bool = false, + + pub fn init(template: usize, request: Request) !Binding { + try request.validate(); + if (template == 0) return error.UnsupportedBinding; + return .{ .template = template, .request = request }; + } + + pub fn invalidate(self: *Binding) void { + self.loaded = false; + self.poisoned = true; + self.template = 0; + } + + pub fn load(self: *Binding, backend: anytype) !void { + if (self.poisoned or self.template == 0) return error.StateBindingPoisoned; + if (self.loaded) return; + var matched: [max_layers]bool = @splat(false); + const nodes = try backend.nodes(self.template); + defer backend.freeNodes(nodes); + for (nodes) |node| { + const header = (try backend.kernel(node)) orelse continue; + if (header.function != self.request.chain and header.function != self.request.chain_to) continue; + if (header.extra or header.alternate_kernel or !std.meta.eql(header.geometry, self.request.geometry)) return error.UnsupportedBinding; + const params = header.params orelse return error.UnsupportedBinding; + const output = header.function == self.request.chain_to; + const values: Parameters = .{ + .nonstate = .{ + .gq = try argument(u64, params, 0), + .gk = try argument(u64, params, 1), + .gv = try argument(u64, params, 2), + .ggb = try argument(u64, params, 3), + .rows = try argument(c_int, params, if (output) 6 else 5), + .ys = try argument(u64, params, if (output) 7 else 6), + }, + .state_in = try argument(u64, params, 4), + .state_out = if (output) try argument(u64, params, 5) else 0, + }; + if (!std.meta.eql(values.nonstate, self.request.nonstate)) return error.UnsupportedBinding; + var found: ?usize = null; + for (self.request.layers[0..self.request.count], 0..) |mapping, index| { + if (values.state_in == mapping.state_in and values.state_out == mapping.state_out) { + if (matched[index] or found != null) return error.UnsupportedBinding; + found = index; + } + } + const index = found orelse return error.UnsupportedBinding; + matched[index] = true; + self.nodes[index] = node; + self.parameters[index] = values; + } + for (matched[0..self.request.count]) |present| if (!present) return error.UnsupportedBinding; + self.loaded = true; + } + + pub fn apply(self: *Binding, backend: anytype, target: Request) !void { + if (self.poisoned) return error.StateBindingPoisoned; + if (!self.loaded or !self.request.compatible(target)) return error.UnsupportedBinding; + errdefer self.poisoned = true; + for (target.layers[0..target.count], 0..) |mapping, index| { + var parameters = self.parameters[index]; + parameters.state_in = mapping.state_in; + parameters.state_out = mapping.state_out; + try backend.set(self.nodes[index], if (mapping.state_out == 0) target.chain else target.chain_to, target.geometry, parameters); + } + } + + pub fn applyTogether(self: *Binding, backend: anytype, target: Request, peer: anytype, signature: u64, cohort: u64) !void { + var succeeded = true; + self.apply(backend, target) catch { + succeeded = false; + }; + errdefer self.poisoned = true; + const agreement = try peer.graphAgree(signature, .{ .state = if (succeeded) .ready else .failed, .cohort = cohort }); + if (!succeeded or agreement.choice != .replay or agreement.cohort != cohort) return error.StateRebindFailed; + } +}; + +pub fn normalized(key: anytype) @TypeOf(key) { + var candidate = key; + candidate.seq[0][4] = 0; + candidate.keep = key.keep_nonstate; + return candidate; +} + +pub fn keepHash(keep: anytype) u64 { + var hash = std.hash.Wyhash.init(0x73746174656b6565); + for (keep.layers) |layer| { + inline for (@typeInfo(@TypeOf(layer)).@"struct".field_names) |field| { + if (comptime !std.mem.eql(u8, field, "state_out")) std.hash.autoHash(&hash, @field(layer, field)); + } + } + std.hash.autoHash(&hash, keep.mixer); + return hash.final(); +} diff --git a/zig/src/families/flashnext/exl3_state_binding_test.zig b/zig/src/families/flashnext/exl3_state_binding_test.zig new file mode 100644 index 0000000000..bcaed32840 --- /dev/null +++ b/zig/src/families/flashnext/exl3_state_binding_test.zig @@ -0,0 +1,298 @@ +const std = @import("std"); +const binding = @import("exl3_state_binding.zig"); + +fn requestFor(width: usize, output: bool) binding.Request { + var request: binding.Request = .{ + .width = width, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 100, .gk = 200, .gv = 300, .ggb = 400, .rows = @intCast(width), .ys = 500 }, + .count = 2, + }; + request.layers[0] = .{ .layer = 0, .state_in = 1000, .state_out = if (output) 2000 else 0 }; + request.layers[1] = .{ .layer = 2, .state_in = 3000, .state_out = if (output) 4000 else 0 }; + return request; +} + +fn flip(request: binding.Request) binding.Request { + var target = request; + for (target.layers[0..target.count]) |*mapping| std.mem.swap(u64, &mapping.state_in, &mapping.state_out); + return target; +} + +const FakeNode = struct { + function: usize, + geometry: binding.Geometry, + values: binding.Parameters, + pointers: [8]?*anyopaque = @splat(null), + kernel: bool = true, + extra: bool = false, + + fn header(self: *FakeNode) binding.Header { + self.pointers[0] = &self.values.nonstate.gq; + self.pointers[1] = &self.values.nonstate.gk; + self.pointers[2] = &self.values.nonstate.gv; + self.pointers[3] = &self.values.nonstate.ggb; + self.pointers[4] = &self.values.state_in; + const output = self.values.state_out != 0; + if (output) self.pointers[5] = &self.values.state_out; + self.pointers[if (output) 6 else 5] = &self.values.nonstate.rows; + self.pointers[if (output) 7 else 6] = &self.values.nonstate.ys; + return .{ .function = self.function, .geometry = self.geometry, .params = &self.pointers, .extra = self.extra }; + } +}; + +const Fake = struct { + template: usize = 42, + records: [4]FakeNode = undefined, + ids: [4]usize = .{ 0, 1, 2, 3 }, + count: usize = 2, + calls: usize = 0, + fail_at: ?usize = null, + query_failure: bool = false, + writes: [2]binding.Parameters = undefined, + + fn init(request: binding.Request) Fake { + var fake: Fake = .{}; + for (request.layers[0..request.count], 0..) |mapping, index| { + fake.records[index] = .{ + .function = if (mapping.state_out == 0) request.chain else request.chain_to, + .geometry = request.geometry, + .values = .{ .nonstate = request.nonstate, .state_in = mapping.state_in, .state_out = mapping.state_out }, + }; + } + return fake; + } + + pub fn nodes(self: *Fake, template: usize) ![]usize { + if (template != self.template) return error.WrongTemplate; + return self.ids[0..self.count]; + } + + pub fn freeNodes(_: *Fake, _: []usize) void {} + + pub fn kernel(self: *Fake, node: usize) !?binding.Header { + if (self.query_failure) return error.PriorAsyncFailure; + if (!self.records[node].kernel) return null; + return self.records[node].header(); + } + + pub fn set(self: *Fake, node: usize, function: usize, geometry: binding.Geometry, parameters: binding.Parameters) !void { + const call = self.calls; + self.calls += 1; + if (self.fail_at == call) return error.SetterFailed; + try std.testing.expectEqual(self.records[node].function, function); + try std.testing.expectEqualDeep(self.records[node].geometry, geometry); + self.writes[node] = parameters; + } +}; + +test "exl3: state binding bounded widths and orientations preserve typed nonstate values" { + for (1..binding.max_width + 1) |width| { + const request = requestFor(width, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + const target = flip(request); + try state.apply(&fake, target); + for (target.layers[0..target.count], fake.writes) |mapping, written| { + try std.testing.expectEqualDeep(request.nonstate, written.nonstate); + try std.testing.expectEqual(mapping.state_in, written.state_in); + try std.testing.expectEqual(mapping.state_out, written.state_out); + } + fake.records[0].values.nonstate.rows = 99; + try state.apply(&fake, request); + for (request.layers[0..request.count], fake.writes) |mapping, written| { + try std.testing.expectEqualDeep(request.nonstate, written.nonstate); + try std.testing.expectEqual(mapping.state_in, written.state_in); + try std.testing.expectEqual(mapping.state_out, written.state_out); + } + try std.testing.expectEqual(@as(usize, 4), fake.calls); + } +} + +test "exl3: state binding schema7 reads c_int rows and u64 ys without a state output" { + for (1..binding.max_width + 1) |width| { + const request = requestFor(width, false); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + try state.apply(&fake, request); + try std.testing.expectEqualDeep(request.nonstate, fake.writes[0].nonstate); + var changed = request; + changed.layers[0].state_in += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + try std.testing.expectEqual(@as(usize, 2), fake.calls); + } +} + +test "exl3: state binding wrong functions are never inspected and missing duplicate mappings decline" { + const request = requestFor(2, true); + var fake = Fake.init(request); + fake.records[0].function = 999; + fake.records[0].extra = true; + var state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.count = 1; + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.records[1] = fake.records[0]; + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + var duplicate = request; + duplicate.layers[1].layer = duplicate.layers[0].layer; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, duplicate)); + duplicate = request; + duplicate.layers[1].state_in = duplicate.layers[0].state_out; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, duplicate)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); +} + +test "exl3: state binding full validation precedes every setter" { + const request = requestFor(1, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var target = flip(request); + target.layers[1].state_out += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.nonstate.gk += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.geometry.block[0] += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.chain_to += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); + fake = Fake.init(request); + fake.records[1].values.nonstate.ys += 1; + state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.query_failure = true; + try std.testing.expectError(error.PriorAsyncFailure, state.load(&fake)); +} + +test "exl3: state binding partial update poisons all future uses" { + const request = requestFor(2, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + fake.fail_at = 1; + try std.testing.expectError(error.SetterFailed, state.apply(&fake, flip(request))); + try std.testing.expect(state.poisoned); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, request)); + try std.testing.expectEqual(@as(usize, 2), fake.calls); +} + +const Peer = struct { + const Vote = struct { state: enum { ready, failed }, cohort: u64 }; + const Decision = struct { choice: enum { replay, capture, eager }, cohort: u64 }; + choice: Decision = .{ .choice = .replay, .cohort = 7 }, + observed: ?Vote = null, + disconnected: bool = false, + + pub fn graphAgree(self: *Peer, _: u64, vote: Vote) !Decision { + self.observed = vote; + if (self.disconnected) return error.RankDisconnected; + return self.choice; + } +}; + +test "exl3: state binding peer failure disagreement and disconnect poison successful local updates" { + const request = requestFor(2, true); + for (0..4) |scenario| { + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var peer: Peer = .{}; + switch (scenario) { + 0 => peer.choice.choice = .capture, + 1 => peer.choice.cohort = 8, + 2 => peer.disconnected = true, + 3 => fake.fail_at = 1, + else => unreachable, + } + if (scenario == 2) { + try std.testing.expectError(error.RankDisconnected, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + } else try std.testing.expectError(error.StateRebindFailed, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + try std.testing.expect(state.poisoned); + try std.testing.expectEqual(scenario == 3, peer.observed.?.state == .failed); + } + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var peer: Peer = .{}; + try state.applyTogether(&fake, flip(request), &peer, 3, 7); + try std.testing.expect(!state.poisoned); +} + +test "exl3: state binding extended windows reject incompatible shapes and poison partial rank updates" { + for (2..binding.max_width + 1) |width| { + const request = requestFor(width, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var changed = flip(request); + changed.width = if (width == binding.max_width) 2 else width + 1; + changed.nonstate.rows = @intCast(changed.width); + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.nonstate.rows += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.layers[1].state_in += 5000; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.layers[1].layer += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.nonstate.ys += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); + fake.fail_at = 1; + var peer: Peer = .{}; + try std.testing.expectError(error.StateRebindFailed, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + try std.testing.expect(state.poisoned); + try std.testing.expect(peer.observed.?.state == .failed); + try std.testing.expectEqual(@as(usize, 2), fake.calls); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, request)); + } + for ([_]usize{ 0, binding.max_width + 1, std.math.maxInt(usize) }) |width| { + var request = requestFor(2, true); + request.width = width; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, request)); + } +} + +test "exl3: state binding template replacement requires new metadata" { + const request = requestFor(1, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + state.invalidate(); + fake.template = 43; + try std.testing.expectError(error.StateBindingPoisoned, state.load(&fake)); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, flip(request))); + state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, flip(request))); + try state.load(&fake); + try state.apply(&fake, flip(request)); +} + +test "exl3: state binding keep hash excludes only layer state_out" { + const Layer = struct { proj: u64 = 3, state_out: u64 = 4, hist: u64 = 5 }; + const Keep = struct { layers: [2]Layer = @splat(.{}), mixer: u64 = 8 }; + const original: Keep = .{}; + var changed = original; + changed.layers[0].state_out += 1; + try std.testing.expectEqual(binding.keepHash(original), binding.keepHash(changed)); + changed.layers[1].hist += 1; + try std.testing.expect(binding.keepHash(original) != binding.keepHash(changed)); + changed = original; + changed.mixer += 1; + try std.testing.expect(binding.keepHash(original) != binding.keepHash(changed)); +} diff --git a/zig/src/families/flashnext/exl3_telemetry.zig b/zig/src/families/flashnext/exl3_telemetry.zig new file mode 100644 index 0000000000..f53a76b961 --- /dev/null +++ b/zig/src/families/flashnext/exl3_telemetry.zig @@ -0,0 +1,301 @@ +const std = @import("std"); + +pub const Phase = enum { prefill, decode, verify, unknown }; +pub const Rank = struct { + rank: usize, + world: usize, + + pub fn validate(owner: Rank) !void { + if (owner.world == 0 or owner.rank >= owner.world) return error.TelemetryRank; + } +}; + +var next_identity: std.atomic.Value(u64) = .init(0); + +pub fn identity() !u64 { + var previous = next_identity.load(.monotonic); + while (true) { + const next = try std.math.add(u64, previous, 1); + previous = next_identity.cmpxchgWeak(previous, next, .monotonic, .monotonic) orelse return next; + } +} + +pub const Device = struct { + plans: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + inplace_experts: u64 = 0, + overflow_plans: u64 = 0, + + pub fn pending(current: anytype, seen: @TypeOf(current)) Device { + return .{ .plans = current.rounds -% seen.rounds, .hit_picks = current.hit_slots -% seen.hit_slots, .miss_picks = current.miss_slots -% seen.miss_slots, .copies = current.copies -% seen.copies, .evictions = current.evictions -% seen.evictions, .inplace_experts = current.inplace -% seen.inplace, .overflow_plans = current.overflow -% seen.overflow }; + } + + pub fn pendingBounded(current: anytype, seen: @TypeOf(current), max_picks: u64) !Device { + const delta = pending(current, seen); + const bound = try std.math.mul(u64, delta.plans, max_picks); + if (max_picks == 0 or bound >= @as(u64, 1) << 32) return error.TelemetryExpertIntervalTooLong; + if (try std.math.add(u64, delta.hit_picks, delta.miss_picks) > bound or try std.math.add(u64, delta.copies, delta.inplace_experts) > bound or delta.overflow_plans > delta.plans) return error.TelemetryInconsistentCounters; + return delta; + } + + pub fn record(total: *Device, valid: *bool, current: anytype, seen: @TypeOf(current), max_picks: u64) void { + if (!valid.*) return; + const delta = pendingBounded(current, seen, max_picks) catch { + valid.* = false; + return; + }; + add(Device, total, delta) catch { + valid.* = false; + }; + } +}; + +pub fn add(comptime Counters: type, total: *Counters, contribution: Counters) !void { + var result = total.*; + inline for (@typeInfo(Counters).@"struct".field_names) |field| @field(result, field) = try std.math.add(u64, @field(result, field), @field(contribution, field)); + total.* = result; +} + +pub fn difference(comptime Counters: type, before: Counters, after: Counters) !Counters { + var result: Counters = undefined; + inline for (@typeInfo(Counters).@"struct".field_names) |field| { + if (@field(after, field) < @field(before, field)) return error.TelemetryCounterRegression; + @field(result, field) = @field(after, field) - @field(before, field); + } + return result; +} + +pub const Experts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, + ring_dma_copies: u64 = 0, + ring_dma_requested_bytes: u64 = 0, + ring_unused_copies: u64 = 0, + ring_hit_picks: u64 = 0, + ring_miss_picks: u64 = 0, + ring_hit_expert_requests: u64 = 0, + ring_miss_expert_requests: u64 = 0, + ring_demand_copies: u64 = 0, + ring_demand_requested_bytes: u64 = 0, + ring_demand_unused_copies: u64 = 0, + ring_demand_unused_requested_bytes: u64 = 0, + ring_initial_prefetch_copies: u64 = 0, + ring_initial_prefetch_requested_bytes: u64 = 0, + ring_initial_prefetch_unused_copies: u64 = 0, + ring_initial_prefetch_unused_requested_bytes: u64 = 0, + ring_cross_layer_prefetch_copies: u64 = 0, + ring_cross_layer_prefetch_requested_bytes: u64 = 0, + ring_cross_layer_prefetch_unused_copies: u64 = 0, + ring_cross_layer_prefetch_unused_requested_bytes: u64 = 0, +}; + +pub const Kv = struct { + block_requests: u64 = 0, + block_misses: u64 = 0, + sparse_calls: u64 = 0, + logical_fetch_bytes: u64 = 0, + rebases: u64 = 0, + written_rows_per_side: u64 = 0, +}; + +pub const KvSample = struct { + owner: u64, + revision: u64, + layer: usize, + empty: bool, + epoch: u32, + clock: u32, + requests: u32, + misses: u32, + calls: u32, + overflow: bool, + max_blocks: u64, + block_bytes: u64, + rebases: u64, + written: u64, + + pub fn delta(before: KvSample, after: KvSample, allow_initial_reset: bool) !Kv { + if (before.owner == 0 or before.owner != after.owner or before.layer != after.layer or before.max_blocks != after.max_blocks or before.block_bytes != after.block_bytes) return error.TelemetryOwnerChanged; + if (before.overflow or after.overflow) return error.TelemetryKvOverflow; + if (before.revision != after.revision) { + if (!allow_initial_reset or !before.empty or before.epoch != 0 or before.requests != 0 or before.misses != 0 or before.calls != 0 or after.revision != try std.math.add(u64, before.revision, 1)) return error.TelemetryKvReset; + } + if (after.clock < before.clock or after.rebases < before.rebases or after.written < before.written) return error.TelemetryCounterRegression; + const clock_delta = after.clock - before.clock; + const bound = try std.math.mul(u64, clock_delta, after.max_blocks); + if (bound >= @as(u64, 1) << 32) return error.TelemetryKvIntervalTooLong; + const requests: u64 = after.requests -% before.requests; + const misses: u64 = after.misses -% before.misses; + const calls: u64 = after.calls -% before.calls; + if (requests > bound or misses > requests or calls > clock_delta) return error.TelemetryInconsistentCounters; + return .{ .block_requests = requests, .block_misses = misses, .sparse_calls = calls, .logical_fetch_bytes = try std.math.mul(u64, misses, after.block_bytes), .rebases = after.rebases - before.rebases, .written_rows_per_side = after.written - before.written }; + } +}; + +pub const Snapshot = struct { + rank: Rank, + scope: u64, + ring_owner: u64 = 0, + experts: Experts = .{}, + kv: []KvSample, + expert_pools: usize = 0, + layers: []LayerSample = &.{}, + + pub fn deinit(snapshot: *Snapshot, allocator: std.mem.Allocator) void { + allocator.free(snapshot.kv); + allocator.free(snapshot.layers); + snapshot.* = undefined; + } +}; + +pub const LayerCounts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, + + pub fn fromExperts(experts: Experts) LayerCounts { + var result: LayerCounts = undefined; + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| @field(result, field) = @field(experts, field); + return result; + } +}; + +pub const LayerSample = struct { + owner: u64, + kind: enum { decoder, mtp }, + layer: usize, + lines: usize, + counts: LayerCounts = .{}, + + fn samePool(first: LayerSample, last: LayerSample) bool { + return first.owner != 0 and first.owner == last.owner and first.kind == last.kind and first.layer == last.layer and first.lines == last.lines; + } +}; + +pub fn LayerCursor(comptime capacity: usize) type { + return struct { + lock: std.atomic.Mutex = .unlocked, + scope: u64 = 0, + rank: Rank = .{ .rank = 0, .world = 1 }, + samples: [capacity]LayerSample = undefined, + len: usize = 0, + + pub fn claim(cursor: *@This(), allocator: std.mem.Allocator, before: Snapshot, after: Snapshot) ![]LayerSample { + while (!cursor.lock.tryLock()) std.atomic.spinLoopHint(); + defer cursor.lock.unlock(); + if (before.layers.len != after.layers.len or after.layers.len > capacity) return error.TelemetryOwnerChanged; + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope) return error.TelemetryOwnerChanged; + if (cursor.len != 0 and (cursor.scope != after.scope or !std.meta.eql(cursor.rank, after.rank) or cursor.len != after.layers.len)) return error.TelemetryOwnerChanged; + const result = try allocator.alloc(LayerSample, after.layers.len); + errdefer allocator.free(result); + for (before.layers, after.layers, result, 0..) |first, last, *delta, index| { + if (!first.samePool(last)) return error.TelemetryOwnerChanged; + for (after.layers[0..index]) |previous| if (previous.owner == last.owner) return error.TelemetryDuplicateOwner; + if (cursor.len != 0 and !cursor.samples[index].samePool(last)) return error.TelemetryOwnerChanged; + delta.* = last; + delta.counts = try difference(LayerCounts, first.counts, last.counts); + } + for (after.layers, 0..) |last, index| { + cursor.samples[index] = last; + } + cursor.scope = after.scope; + cursor.rank = after.rank; + cursor.len = after.layers.len; + return result; + } + }; +} + +pub const Interval = struct { + phase: Phase, + model_rows: u64, + steps: u64, + allow_initial_kv_reset: bool = false, +}; + +pub fn fraction(misses: u64, hits: u64) !?f64 { + const denominator = try std.math.add(u64, misses, hits); + return if (denominator == 0) null else @as(f64, @floatFromInt(misses)) / @as(f64, @floatFromInt(denominator)); +} + +pub const Report = struct { + rank: Rank, + interval: Interval, + experts: Experts, + kv: Kv, + expert_pools: usize = 0, + kv_layers: usize = 0, + layers: []const LayerSample = &.{}, + + pub fn write(report: Report, writer: *std.Io.Writer) !void { + try report.header(writer, "expert"); + inline for (@typeInfo(Experts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(report.experts, field) }); + try writer.writeByte('\n'); + for (report.layers) |layer| { + try report.header(writer, "expert-layer"); + try writer.print(" pool={d} kind={s} layer={d} lines={d}", .{ layer.owner, @tagName(layer.kind), layer.layer, layer.lines }); + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(layer.counts, field) }); + try writer.writeByte('\n'); + } + try report.header(writer, "bf16-kv"); + inline for (@typeInfo(Kv).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(report.kv, field) }); + try writer.writeByte('\n'); + } + + fn header(report: Report, writer: *std.Io.Writer, cache: []const u8) !void { + try writer.print("EXL3-CACHE v=1 rank={d} world={d} phase={s} rows={d} steps={d} expert_pools={d} kv_layers={d} cache={s}", .{ report.rank.rank, report.rank.world, @tagName(report.interval.phase), report.interval.model_rows, report.interval.steps, report.expert_pools, report.kv_layers, cache }); + } + + pub fn formatAlloc(report: Report, allocator: std.mem.Allocator) ![]u8 { + var writer: std.Io.Writer.Allocating = .init(allocator); + errdefer writer.deinit(); + try report.write(&writer.writer); + return writer.toOwnedSlice(); + } + + pub fn between(before: Snapshot, after: Snapshot, interval: Interval) !Report { + try before.rank.validate(); + try after.rank.validate(); + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope or before.kv.len != after.kv.len or before.expert_pools != after.expert_pools) return error.TelemetryOwnerChanged; + if (before.ring_owner != 0 and before.ring_owner != after.ring_owner) return error.TelemetryOwnerChanged; + if (interval.model_rows == 0 or interval.steps == 0) return error.TelemetryDenominator; + if (interval.allow_initial_kv_reset and interval.phase != .prefill) return error.TelemetryKvReset; + const experts = try difference(Experts, before.experts, after.experts); + if (experts.device_overflow_plans != 0) return error.TelemetryExpertOverflow; + var kv: Kv = .{}; + for (before.kv, after.kv) |first, last| try add(Kv, &kv, try first.delta(last, interval.allow_initial_kv_reset)); + return .{ .rank = after.rank, .interval = interval, .experts = experts, .kv = kv, .expert_pools = after.expert_pools, .kv_layers = after.kv.len }; + } +}; + +test "exl3: telemetry host contracts" { + _ = @import("exl3_telemetry_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3_telemetry_test.zig new file mode 100644 index 0000000000..b4706962db --- /dev/null +++ b/zig/src/families/flashnext/exl3_telemetry_test.zig @@ -0,0 +1,217 @@ +const std = @import("std"); +const telemetry = @import("exl3_telemetry.zig"); +const expectEqual = std.testing.expectEqual; +const expectError = std.testing.expectError; + +fn sample() telemetry.KvSample { + return .{ .owner = 12, .revision = 3, .layer = 0, .empty = false, .epoch = 100, .clock = 100, .requests = 40, .misses = 20, .calls = 10, .overflow = false, .max_blocks = 1024, .block_bytes = 8192, .rebases = 1, .written = 4096 }; +} + +test "exl3: KV telemetry unsigned wrap retains block denominator and BF16 payload units" { + var before = sample(); + before.requests = std.math.maxInt(u32) - 3; + before.misses = std.math.maxInt(u32) - 1; + var after = before; + after.epoch += 2; + after.clock += 2; + after.requests = 6; + after.misses = 2; + after.calls += 2; + const result = try before.delta(after, false); + try expectEqual(@as(u64, 10), result.block_requests); + try expectEqual(@as(u64, 4), result.block_misses); + try expectEqual(@as(u64, 32768), result.logical_fetch_bytes); + try expectEqual(@as(u64, 2), result.sparse_calls); +} + +test "exl3: KV telemetry rejects restore reset recycle and ambiguous wrap intervals" { + const before = sample(); + var after = before; + after.revision += 1; + after.requests += 1; + try expectError(error.TelemetryKvReset, before.delta(after, false)); + try expectError(error.TelemetryKvReset, before.delta(after, true)); + after = before; + after.owner += 1; + try expectError(error.TelemetryOwnerChanged, before.delta(after, false)); + after = before; + after.clock -= 1; + try expectError(error.TelemetryCounterRegression, before.delta(after, false)); + after = before; + after.clock += 1 << 22; + try expectError(error.TelemetryKvIntervalTooLong, before.delta(after, false)); + after = before; + after.overflow = true; + try expectError(error.TelemetryKvOverflow, before.delta(after, false)); + after = before; + after.calls += 1; + try expectError(error.TelemetryInconsistentCounters, before.delta(after, false)); +} + +test "exl3: KV initial empty reset needs explicit opt in and exactly one revision" { + var before = sample(); + before.empty = true; + before.epoch = 0; + before.clock = 0; + before.requests = 0; + before.misses = 0; + before.calls = 0; + var after = before; + after.revision += 1; + after.epoch = 10; + after.clock = 10; + after.requests = 90; + after.misses = 60; + after.calls = 4; + try expectError(error.TelemetryKvReset, before.delta(after, false)); + try expectEqual(@as(u64, 60), (try before.delta(after, true)).block_misses); + after.revision += 1; + try expectError(error.TelemetryKvReset, before.delta(after, true)); +} + +test "exl3: fresh shared lane at high arena clock uses completed boundary delta not last layer epoch" { + var before = sample(); + before.empty = true; + before.epoch = 0; + before.clock = 524288; + before.max_blocks = 8192; + before.requests = 0; + before.misses = 0; + before.calls = 0; + var after = before; + after.clock += 1; + after.epoch = after.clock; + after.requests = 512; + after.misses = 512; + after.calls = 1; + try std.testing.expect(@as(u64, after.epoch - before.epoch) * after.max_blocks >= @as(u64, 1) << 32); + const result = try before.delta(after, false); + try expectEqual(@as(u64, 512), result.block_requests); + try expectEqual(@as(u64, 512 * 8192), result.logical_fetch_bytes); + after.revision += 1; + try expectError(error.TelemetryKvReset, before.delta(after, false)); + try expectEqual(@as(u64, 512), (try before.delta(after, true)).block_misses); +} + +test "exl3: restored historical layer epoch cannot replace boundary clock or weaken wrap checks" { + var before = sample(); + before.epoch = 900000; + before.clock = 20; + before.requests = std.math.maxInt(u32) - 3; + before.misses = std.math.maxInt(u32) - 1; + var after = before; + after.clock = 21; + after.epoch = 21; + after.requests = 6; + after.misses = 2; + after.calls += 1; + const result = try before.delta(after, false); + try expectEqual(@as(u64, 10), result.block_requests); + try expectEqual(@as(u64, 4), result.block_misses); + after.clock = before.clock - 1; + try expectError(error.TelemetryCounterRegression, before.delta(after, false)); + after.clock = before.clock + (1 << 22); + try expectError(error.TelemetryKvIntervalTooLong, before.delta(after, false)); + before.clock = std.math.maxInt(u32); + after.clock = 0; + try expectError(error.TelemetryCounterRegression, before.delta(after, false)); +} + +test "exl3: checked telemetry aggregation exceeds u32 without overflow or partial mutation" { + var totals: telemetry.Experts = .{ .device_hit_picks = std.math.maxInt(u32) }; + try telemetry.add(telemetry.Experts, &totals, .{ .device_hit_picks = 10 }); + try expectEqual(@as(u64, std.math.maxInt(u32)) + 10, totals.device_hit_picks); + totals.device_hit_picks = std.math.maxInt(u64); + const before = totals; + try expectError(error.Overflow, telemetry.add(telemetry.Experts, &totals, .{ .host_plans = 5, .device_hit_picks = 1 })); + try std.testing.expectEqualDeep(before, totals); + try expectError(error.TelemetryCounterRegression, telemetry.difference(telemetry.Experts, totals, .{})); + try expectEqual(@as(?f64, null), try telemetry.fraction(0, 0)); + try expectEqual(@as(?f64, 0.25), try telemetry.fraction(1, 3)); + try expectError(error.Overflow, telemetry.fraction(std.math.maxInt(u64), 1)); +} + +test "exl3: invalid device telemetry cannot fail normal ownership transfer or publish partial totals" { + const Misc = struct { now: u32 = 0, rounds: u32 = 0, hit_slots: u32 = 0, miss_slots: u32 = 0, copies: u32 = 0, evictions: u32 = 0, overflow: u32 = 0, inplace: u32 = 0 }; + var total: telemetry.Device = .{ .copies = 10 }; + var valid = true; + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .hit_slots = 11 }, Misc{}, 10); + try std.testing.expect(!valid); + try expectEqual(@as(u64, 10), total.copies); + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .copies = 1 }, Misc{}, 10); + try expectEqual(@as(u64, 10), total.copies); + try expectError(error.TelemetryExpertIntervalTooLong, telemetry.Device.pendingBounded(Misc{ .rounds = 1 << 30 }, Misc{}, 10)); +} + +test "exl3: phase report stable schema rank and model-row denominators" { + var values = [_]telemetry.KvSample{sample()}; + var before: telemetry.Snapshot = .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 7, .kv = &values }; + var after = before; + after.experts.device_hit_picks = 40; + after.experts.device_miss_picks = 10; + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 16, .steps = 8 }); + const text = try report.formatAlloc(std.testing.allocator); + defer std.testing.allocator.free(text); + try std.testing.expect(std.mem.startsWith(u8, text, "EXL3-CACHE v=1 rank=1 world=2 phase=decode rows=16 steps=8 expert_pools=0 kv_layers=1 cache=expert ")); + try std.testing.expect(std.mem.indexOf(u8, text, "device_hit_picks=40 device_miss_picks=10") != null); + try std.testing.expect(std.mem.indexOf(u8, text, "cache=bf16-kv block_requests=0 block_misses=0 sparse_calls=0 logical_fetch_bytes=0") != null); + try expectError(error.TelemetryDenominator, telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 0, .steps = 1 })); + after.rank.rank = 0; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); + after = before; + after.scope += 1; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); + after = before; + after.ring_owner = 13; + _ = try telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 1, .steps = 1 }); + before.ring_owner = 12; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 1, .steps = 1 })); +} + +test "exl3: per-layer deltas sum to aggregate counters and overlapping reports claim counters once" { + const rank: telemetry.Rank = .{ .rank = 0, .world = 2 }; + var before_layers = [_]telemetry.LayerSample{.{ .owner = 11, .kind = .decoder, .layer = 4, .lines = 20 }}; + var after_layers = before_layers; + after_layers[0].counts.host_miss_picks = 5; + after_layers[0].counts.host_mandatory_copies = 2; + after_layers[0].counts.device_miss_picks = 3; + after_layers[0].counts.device_copies = 1; + var before: telemetry.Snapshot = .{ .rank = rank, .scope = 42, .kv = &.{}, .layers = &before_layers, .expert_pools = 1 }; + var after = before; + after.layers = &after_layers; + after.experts.host_miss_picks = 5; + after.experts.host_mandatory_copies = 2; + after.experts.device_miss_picks = 3; + after.experts.device_copies = 1; + var cursor: telemetry.LayerCursor(2) = .{}; + const first = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(first); + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(report.experts.host_miss_picks, first[0].counts.host_miss_picks); + try std.testing.expectEqual(report.experts.host_mandatory_copies, first[0].counts.host_mandatory_copies); + try std.testing.expectEqual(report.experts.device_miss_picks, first[0].counts.device_miss_picks); + try std.testing.expectEqual(report.experts.device_copies, first[0].counts.device_copies); + + var concurrent_before_layers = after_layers; + var concurrent_after_layers = after_layers; + concurrent_before_layers[0].counts.host_miss_picks = 3; + concurrent_before_layers[0].counts.host_mandatory_copies = 1; + concurrent_before_layers[0].counts.device_miss_picks = 2; + concurrent_before_layers[0].counts.device_copies = 0; + concurrent_after_layers[0].counts.host_miss_picks = 8; + concurrent_after_layers[0].counts.host_mandatory_copies = 4; + concurrent_after_layers[0].counts.device_miss_picks = 6; + concurrent_after_layers[0].counts.device_copies = 2; + before.layers = &concurrent_before_layers; + after.layers = &concurrent_after_layers; + before.experts = .{ .host_miss_picks = 5, .host_mandatory_copies = 2, .device_miss_picks = 3, .device_copies = 1 }; + after.experts = .{ .host_miss_picks = 10, .host_mandatory_copies = 5, .device_miss_picks = 7, .device_copies = 3 }; + const second = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(second); + const second_report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(second_report.experts.host_miss_picks, second[0].counts.host_miss_picks); + try std.testing.expectEqual(second_report.experts.host_mandatory_copies, second[0].counts.host_mandatory_copies); + try std.testing.expectEqual(second_report.experts.device_miss_picks, second[0].counts.device_miss_picks); + try std.testing.expectEqual(second_report.experts.device_copies, second[0].counts.device_copies); + try std.testing.expectEqual(@as(u64, 10), first[0].counts.host_miss_picks + second[0].counts.host_miss_picks); +} diff --git a/zig/src/families/flashnext/exl3_tp.zig b/zig/src/families/flashnext/exl3_tp.zig new file mode 100644 index 0000000000..3bd8d1604f --- /dev/null +++ b/zig/src/families/flashnext/exl3_tp.zig @@ -0,0 +1,151 @@ +//! Two-rank cuts for the Qwen3.8-Flash-Next EXL3 3.05 pack. +//! An EXL3 axis splits only when it is a multiple of 128 Hadamard blocks on every rank. +//! Gate and up are K-split (N=640 is not). Down is N-split (K=640 is not). +//! Routed experts are split by id. Reductions add rank 0, then rank 1. + +const std = @import("std"); +const shape = @import("exl3_shape.zig"); + +pub const Axis = enum { replicate, split_k, split_n, expert_parallel }; + +pub const Role = enum { gate, up, down, q, k, v, o, head, shared, routed_set, hc_lowrank }; + +pub const Cut = struct { + role: Role, + axis: Axis, + k: usize, + n: usize, + part: usize, +}; + +pub const Span = struct { lo: usize, hi: usize }; + +/// True when `axis_len` can be cut into `ranks` EXL3 pieces. +pub fn legal(axis_len: usize, ranks: usize) bool { + if (ranks == 0) return false; + return axis_len % (shape.hadamard * ranks) == 0; +} + +/// The 3.05 cut for `role` on `rank`. Ranks are 1 (whole tensor) or 2. +pub fn cut(role: Role, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Cut { + if (ranks != 1 and ranks != 2) return error.IllegalCut; + if (rank >= ranks) return error.RankOutOfRange; + if (ranks == 1) return unsplit(role); + return switch (role) { + .gate, .up => .{ .role = role, .axis = .split_k, .k = shape.hidden, .n = shape.expert_intermediate, .part = shape.hidden / 2 }, + .down => .{ .role = role, .axis = .split_n, .k = shape.expert_intermediate, .n = shape.hidden, .part = shape.hidden / 2 }, + .shared => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.expert_intermediate, .part = 0 }, + .routed_set => .{ .role = role, .axis = .expert_parallel, .k = 0, .n = 0, .part = 0 }, + .q => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.qOut(), .part = shape.qOut() / 2 }, + .k, .v => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.kvOut(), .part = shape.kvOut() / 2 }, + // o keeps the heads this rank produced in q, so the split is along K. + .o => .{ .role = role, .axis = .split_k, .k = shape.qOut(), .n = shape.hidden, .part = shape.qOut() / 2 }, + .head => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.vocab, .part = shape.vocab / 2 }, + // Unquantized. The 320-wide low-rank axis is split in half. 128 alignment does not apply. + .hc_lowrank => .{ .role = role, .axis = .split_n, .k = shape.hc_lowrank, .n = shape.hidden, .part = shape.hc_lowrank / 2 }, + }; +} + +fn unsplit(role: Role) Cut { + return switch (role) { + .gate, .up, .shared => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.expert_intermediate, .part = shape.hidden }, + .down => .{ .role = role, .axis = .replicate, .k = shape.expert_intermediate, .n = shape.hidden, .part = shape.hidden }, + .routed_set => .{ .role = role, .axis = .expert_parallel, .k = 0, .n = 0, .part = 0 }, + .q => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.qOut(), .part = shape.qOut() }, + .k, .v => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.kvOut(), .part = shape.kvOut() }, + .o => .{ .role = role, .axis = .replicate, .k = shape.qOut(), .n = shape.hidden, .part = shape.qOut() }, + .head => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.vocab, .part = shape.vocab }, + .hc_lowrank => .{ .role = role, .axis = .replicate, .k = shape.hc_lowrank, .n = shape.hidden, .part = shape.hc_lowrank }, + }; +} + +/// `[lo, hi)` on the split axis. Expert-parallel spans are expert ids. +pub fn range(role: Role, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Span { + const c = try cut(role, rank, ranks); + if (role == .routed_set) { + if (shape.experts % ranks != 0) return error.IllegalCut; + const part = shape.experts / ranks; + return .{ .lo = rank * part, .hi = (rank + 1) * part }; + } + const along: usize = switch (role) { + .gate, .up, .o => c.k, + .hc_lowrank => c.k, + else => c.n, + }; + if (c.axis == .replicate or ranks == 1) return .{ .lo = 0, .hi = along }; + const part = along / ranks; + return .{ .lo = rank * part, .hi = (rank + 1) * part }; +} + +/// True when an EXL3 N-split at `ranks` would break a Hadamard block. +pub fn refusesColumnSplit(k: usize, n: usize, ranks: usize) bool { + _ = k; + return !legal(n, ranks); +} + +/// `axis_len` is this tensor's header axis, not a class total. `hc_lowrank` is not a Hadamard cut. +pub fn share(role: Role, span: u64, axis_len: usize, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!u64 { + const c = try cut(role, rank, ranks); + if (span == 0 or c.axis == .replicate or c.axis == .expert_parallel) return span; + if (role == .hc_lowrank) { + if (axis_len == 0 or axis_len % ranks != 0) return span; + return span / @as(u64, ranks); + } + if (axis_len == 0 or !legal(axis_len, ranks)) return span; + return span / @as(u64, ranks); +} + +test "gate is K-split and down is N-split" { + const gate = try cut(.gate, 0, 2); + try std.testing.expectEqual(Axis.split_k, gate.axis); + try std.testing.expectEqual(@as(usize, 1280), gate.part); + try std.testing.expect(refusesColumnSplit(2560, 640, 2)); + try std.testing.expect(!legal(640, 2)); + try std.testing.expect(legal(2560, 2)); + try std.testing.expect(!legal(320, 2)); + + const down = try cut(.down, 1, 2); + try std.testing.expectEqual(Axis.split_n, down.axis); + try std.testing.expect(down.axis != .split_k); + const span = try range(.down, 1, 2); + try std.testing.expectEqual(@as(usize, 1280), span.lo); + try std.testing.expectEqual(@as(usize, 2560), span.hi); +} + +test "expert ids abut and q, k, hc take their parts" { + const a = try range(.routed_set, 0, 2); + const b = try range(.routed_set, 1, 2); + try std.testing.expectEqual(@as(usize, 0), a.lo); + try std.testing.expectEqual(@as(usize, 256), a.hi); + try std.testing.expectEqual(a.hi, b.lo); + try std.testing.expectEqual(@as(usize, 512), b.hi); + + const q = try cut(.q, 0, 2); + try std.testing.expectEqual(@as(usize, 3072), q.part); + try std.testing.expect(q.part % 128 == 0); + try std.testing.expectEqual(@as(usize, 256), (try cut(.k, 0, 2)).part); + try std.testing.expectEqual(@as(usize, 160), (try cut(.hc_lowrank, 0, 2)).part); + try std.testing.expect((try cut(.hc_lowrank, 0, 2)).part % 128 != 0); + try std.testing.expectError(error.RankOutOfRange, cut(.gate, 2, 2)); +} + +test "one rank keeps the whole tensor and both ranks of q stay 128-aligned" { + try std.testing.expectEqual(Axis.replicate, (try cut(.gate, 0, 1)).axis); + try std.testing.expectEqual(Axis.expert_parallel, (try cut(.routed_set, 0, 1)).axis); + const all = try range(.routed_set, 0, 1); + try std.testing.expectEqual(@as(usize, 0), all.lo); + try std.testing.expectEqual(@as(usize, 512), all.hi); + try std.testing.expectEqual(@as(usize, 124160), (try cut(.head, 0, 2)).part); + try std.testing.expect((try cut(.head, 0, 2)).part % 128 == 0); +} + +test "share follows the header axis and hc_lowrank still halves 320" { + try std.testing.expectEqual(@as(u64, 500), try share(.q, 1000, shape.qOut(), 0, 2)); + try std.testing.expectEqual(@as(u64, 1000), try share(.q, 1000, 100, 0, 2)); + try std.testing.expect(!legal(shape.hc_lowrank, 2)); + try std.testing.expectEqual(@as(u64, 160), try share(.hc_lowrank, 320, shape.hc_lowrank, 1, 2)); + try std.testing.expectEqual(@as(u64, 640), try share(.shared, 640, shape.expert_intermediate, 0, 2)); + try std.testing.expectEqual(@as(u64, 1000), try share(.q, 1000, shape.qOut(), 0, 1)); + try std.testing.expectError(error.IllegalCut, share(.q, 1000, shape.qOut(), 0, 0)); + try std.testing.expectError(error.IllegalCut, share(.gate, 1000, shape.hidden, 0, 4)); +} diff --git a/zig/src/families/flashnext/fixtures/exl3_config.json b/zig/src/families/flashnext/fixtures/exl3_config.json new file mode 100644 index 0000000000..ea4b025ebf --- /dev/null +++ b/zig/src/families/flashnext/fixtures/exl3_config.json @@ -0,0 +1,168 @@ +{ + "architectures": [ + "Qwen4ExpForConditionalGeneration" + ], + "image_token_id": 248056, + "language_model_only": false, + "model_type": "qwen4_exp", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 248044, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "hc_count": 4, + "hc_lowrank": 320, + "head_dim": 256, + "heads_per_ngram": 8, + "hidden_act": "silu", + "hidden_size": 2560, + "indexer_budget": 2048, + "indexer_compress_ratio": 4, + "indexer_head_dim": 128, + "indexer_kv_heads": 1, + "indexer_n_heads": 4, + "initializer_range": 0.02, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "make_ngram_vocab_size_divisible_by": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "model_type": "qwen4_exp_text", + "moe_intermediate_size": 640, + "mtp": { + "hybrid": true, + "layer_types": [ + "full_attention" + ], + "mtp_use_hidden_state_from_layer": null, + "num_hidden_layers": 1, + "rope_theta": 10000000 + }, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "ngram_size": 3, + "ngram_vocab_size_base": 20000000, + "num_attention_heads": 24, + "num_experts": 512, + "num_experts_per_tok": 10, + "num_hidden_layers": 48, + "num_key_value_heads": 2, + "output_gate_type": "sigmoid", + "output_router_logits": false, + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "ple_conv_kernel_size": 4, + "ple_embed_dim": 2560, + "ple_layer_ids": [ + 2 + ], + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "router_aux_loss_coef": 0.001, + "shared_expert_intermediate_size": 640, + "split_ngram_parts": 128, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 248320 + }, + "tie_word_embeddings": false, + "transformers_version": "5.8.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen4_exp", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2560, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053, + "quantization_config": { + "quant_method": "exl3", + "version": "1.4.4", + "bits": 3.05, + "head_bits": 5, + "calibration": { + "rows": 250, + "cols": 2048 + }, + "out_scales": "always", + "codebook": "mul1", + "vision_bits": 5, + "mtp_bits": 3 + } +} \ No newline at end of file diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a8463051a5..a205c027b8 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -29,7 +29,7 @@ const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, q /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; -pub const serves: []const []const u8 = &.{ "--mtp-drafts", "--mtp-confidence", "--policy" }; +pub const serves: []const []const u8 = &.{ "--tp", "--mtp-drafts", "--mtp-confidence", "--policy" }; pub const families: []const api.Family = blk: { var out: [registry.len + 2]api.Family = undefined; for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats }; @@ -345,7 +345,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); }; const free = loaded.free_streams; - const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); + const preallocates_lanes = if (@hasDecl(F, "preallocates_lanes")) F.preallocates_lanes else @hasField(F.Options, "lanes"); + const grow = budget.growRoom(room, !preallocates_lanes); const layout = kv_pages.layout(loaded); var paged: ?kv_pages.Plan = null; const streams = if (layout) |l| blk: { diff --git a/zig/src/native/cuda_flashnext_refusal_test.zig b/zig/src/native/cuda_flashnext_refusal_test.zig index 3ed47d9b34..3e7a1210c2 100644 --- a/zig/src/native/cuda_flashnext_refusal_test.zig +++ b/zig/src/native/cuda_flashnext_refusal_test.zig @@ -20,7 +20,7 @@ pub fn cache(comptime open: anytype) !void { .drafts = drafts, .prompt_cache_gib = gib, }; - if (tp == 2 and drafts and (gib orelse 1) > 0) { + if (!F.tp_draft_snapshots and tp == 2 and drafts and (gib orelse 1) > 0) { var problem: []const u8 = ""; try std.testing.expect((try open(F, a, a, std.testing.io, o, &problem)) == null); try std.testing.expectEqualStrings(message, problem); diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index d975185a82..bd23deeaeb 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -8,6 +8,7 @@ pub const deepseek_release_message = "DeepSeek-V4.1 is not served by this releas pub fn backend(o: api.Open) ?[]const u8 { if (std.mem.eql(u8, o.model_type, "deepseek_v4") or std.mem.eql(u8, o.model_type, "deepseek_v41")) return deepseek_release_message; + if (o.tp != 1 and o.tp != 2) return "--tp: the CUDA engine runs one or two ranks (one process a GPU): serve --tp 1 or --tp 2"; if (o.mtp_drafts) |n| if (n > nemotron.native.max_drafts) return std.fmt.comptimePrint("--mtp-drafts: the CUDA engine drafts at most {d}", .{nemotron.native.max_drafts}); if (o.mtp_confidence) |c| if (!(c >= 0 and c <= 1)) return "--mtp-confidence must be from 0 to 1"; if (o.mtp_drafts != null or o.mtp_confidence != null) { @@ -60,6 +61,15 @@ test "a family must consume rank options before tensor parallelism is admitted" try std.testing.expect(family(Ranked, .{ .dir = "d", .model_type = "m", .tp = 2 }) == null); } +test "--tp is served on CUDA for two ranks of a ranked family (Flash Next) only; 4 and 8 stay refused" { + const flashnext = @import("flashnext").native; + for ([_]u32{ 4, 8 }) |tp| try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = flashnext.model_type, .tp = tp }) orelse "", "--tp: the CUDA engine runs one or two ranks"); + for ([_]u32{ 1, 2 }) |tp| try std.testing.expect(backend(.{ .dir = "d", .model_type = flashnext.model_type, .tp = tp }) == null); + try std.testing.expect(family(flashnext, .{ .dir = "d", .model_type = flashnext.model_type, .tp = 2, .rank = 1, .drafts = false }) == null); + // a family without rank options (Nemotron) still refuses --tp 2 at its admission + try std.testing.expectStringStartsWith(family(nemotron.native, .{ .dir = "d", .model_type = nemotron.native.model_type, .tp = 2 }) orelse "", "this CUDA family has no tensor-parallel ranks"); +} + test "--learn is refused whatever the family, and --slide or a learned-disk size alone is not" { try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x" }) orelse "", "--learn learns shared prompt prefixes on Macs only"); try std.testing.expect(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x", .learn_conversations = true }) == null); diff --git a/zig/tests/cuda/exl3_attn_cpu.zig b/zig/tests/cuda/exl3_attn_cpu.zig new file mode 120000 index 0000000000..e1fb6f216b --- /dev/null +++ b/zig/tests/cuda/exl3_attn_cpu.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_attn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_cache.zig b/zig/tests/cuda/exl3_cache.zig new file mode 120000 index 0000000000..ab2377c26b --- /dev/null +++ b/zig/tests/cuda/exl3_cache.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_cache.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_cpu.zig b/zig/tests/cuda/exl3_cpu.zig new file mode 120000 index 0000000000..6107e5710c --- /dev/null +++ b/zig/tests/cuda/exl3_cpu.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_dense_tests.zig b/zig/tests/cuda/exl3_dense_tests.zig index 877d9ed791..7b40fa05f4 100644 --- a/zig/tests/cuda/exl3_dense_tests.zig +++ b/zig/tests/cuda/exl3_dense_tests.zig @@ -16,6 +16,14 @@ pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { var dense = try exl3.Dense.load(r.gpu.d); defer dense.unload(); const ql: qlinear.Linear = .{ .exl3 = &dense }; + { + var ok = true; + counterResetReuse(r, &dense.lin) catch |e| { + state.fail(e, "exl3.linear.counter-reset-reuse"); + ok = false; + }; + if (ok) check.pass("exl3.linear.counter-reset-reuse", .{}); + } { var ok = true; qlinearCheck(r, &ql, &dense.lin) catch |e| { @@ -50,6 +58,59 @@ pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { }) check.pass("exl3.load.redsnow", .{}); } +fn counterResetReuse(rig: *Rig, linear: *const exl3.Linear) !void { + for ([_][2]usize{ .{ 512, 512 }, .{ 2560, 2560 }, .{ 4096, 8192 }, .{ 2560, 512 }, .{ 1280, 128 }, .{ 1024, 8064 } }) |shape| { + const input_width = shape[0]; + const output_width = shape[1]; + const trellis, const suh, const svh = try tensors(rig, input_width, output_width, .{ .halves = 8 }); + const projection_weight = try weight(rig, trellis, suh, svh, input_width, output_width, 8, .mul1); + const input = try rig.bfs(128 * input_width, 1.0); + const device_input = try rig.dev(std.mem.sliceAsBytes(input)); + const counter_bytes = exl3.Linear.counterBytes(projection_weight); + const guarded = try rig.zeros(counter_bytes + 32); + const clean = try exl3_tests.zeroed(rig, counter_bytes); + const rotated = try rig.zeros(128 * input_width * 2); + const partial_bytes = exl3.Linear.zBytes(projection_weight, 128); + const partials = if (partial_bytes > 0) try rig.zeros(partial_bytes) else 0; + const expected_counters = try rig.a.alloc(u8, counter_bytes + 32); + @memset(expected_counters, 0xff); + if (partial_bytes > 0) @memset(expected_counters[16..][0..counter_bytes], 0); + for ([_]usize{ 1, 17, 128, 3, 127, 16 }) |row_count| { + for ([_]exl3.DType{ .bf16, .f32 }) |dtype| { + const bytes = row_count * output_width * @as(usize, if (dtype == .f32) 4 else 2); + const want = try rig.zeros(bytes); + const got = try rig.zeros(bytes + 32); + try linear.matmul(rig.s, projection_weight, device_input, .bf16, row_count, want, dtype, rotated, partials, clean); + try cuda.graph.beginCapture(rig.s, .thread_local); + linear.matmul(rig.s, projection_weight, device_input, .bf16, row_count, got + 16, dtype, rotated, partials, guarded + 16) catch |err| { + var discarded = cuda.graph.endCapture(rig.s) catch return err; + discarded.deinit(); + return err; + }; + var captured = try cuda.graph.endCapture(rig.s); + defer captured.deinit(); + try check.expect(try captured.nodeCount() == 2, "exl3.linear.counter-reset-reuse: expected rotation and GEMM only", .{}); + var exec = try captured.instantiate(); + defer exec.deinit(); + for (0..3) |repeat| { + try rig.s.d.check(rig.s.d.api.cuMemsetD8Async(guarded, 0xff, counter_bytes + 32, rig.s.handle), "cuMemsetD8Async"); + try rig.s.d.check(rig.s.d.api.cuMemsetD8Async(got, 0xff, bytes + 32, rig.s.handle), "cuMemsetD8Async"); + try rig.s.d.check(rig.s.d.api.cuMemsetD8Async(rotated, 0xff, 128 * input_width * 2, rig.s.handle), "cuMemsetD8Async"); + if (partial_bytes > 0) try rig.s.d.check(rig.s.d.api.cuMemsetD8Async(partials, 0xff, partial_bytes, rig.s.handle), "cuMemsetD8Async"); + if (repeat == 0) { + try linear.matmul(rig.s, projection_weight, device_input, .bf16, row_count, got + 16, dtype, rotated, partials, guarded + 16); + } else try exec.launchOn(rig.s); + try check.sameBytes("exl3.linear.counter-reset-reuse.output", try rig.back(got + 16, bytes), try rig.back(want, bytes)); + try check.sameBytes("exl3.linear.counter-reset-reuse.counters", try rig.back(guarded, counter_bytes + 32), expected_counters); + const canary: [16]u8 = @splat(0xff); + try check.sameBytes("exl3.linear.counter-reset-reuse.prefix", try rig.back(got, 16), &canary); + try check.sameBytes("exl3.linear.counter-reset-reuse.suffix", try rig.back(got + 16 + bytes, 16), &canary); + } + } + } + } +} + /// Through qlinear.Linear's decode and prompt dispatch, 1, 17 and 300 rows byte-equal the direct block loop. fn qlinearCheck(r: *Rig, ql: *const qlinear.Linear, lin: *const exl3.Linear) !void { const k = 1024; diff --git a/zig/tests/cuda/exl3_draft_tests.zig b/zig/tests/cuda/exl3_draft_tests.zig new file mode 100644 index 0000000000..12ae25f30d --- /dev/null +++ b/zig/tests/cuda/exl3_draft_tests.zig @@ -0,0 +1,166 @@ +//! exl3-draft: the MTP draft head rows of an EXL3 pack (cuda_exl3_mtp.draftRows) against the full EXL3 lm_head. +//! The 79,591 draft ids' rows are decoded from lm_head's trellis (exl3_format.dequantize per 128-column block, one bf16 +//! rounding). For hidden vectors x (a real embedding row and a random row, bf16), the draft logits sum_k x_k row_v[k] +//! (float64) are compared with +//! - the full head through upstream's exl3.Dense (fp32 out) at the draft columns, and +//! - the exact float64 product x · diag(suh) H Wq H diag(svh) at those columns (first x only: 636M products), +//! each within the bf16 bound 2^-9 sum_k |x_k W[k,v]| (plus the kernel's own error against float64 for the GPU). +//! Usage: tf-cuda-test exl3-draft [threads] + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const st = core.safetensors; +const exl3 = cuda.exl3; +const fmt = core.exl3_format; +const mtp3 = flash.exl3_mtp; +const Gpu = check.Gpu; + +fn bfVal(b: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, b) << 16))); +} + +fn bf16Bits(v: f32) u16 { + return flash.ple.toBf16(v); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const threads: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else mtp3.loadThreads(); + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = gpu.io; + const path = try std.fs.path.join(a, &.{ dir, "model-00007-of-00007.safetensors" }); + var file = try st.File.open(gpu.gpa, io, path); + defer file.close(io); + const tr = file.get("lm_head.trellis") orelse return error.MissingTensor; + const su = file.get("lm_head.suh") orelse return error.MissingTensor; + const sv = file.get("lm_head.svh") orelse return error.MissingTensor; + const head = try mtp3.Head.of(tr, su, sv, file.get("lm_head.mul1")); + const k = head.k; + const n = head.n; + std.debug.print("lm_head: K {d}, N {d}, halves {d} ({d}-bit), {d} MB trellis\n", .{ k, n, head.halves, head.halves / 2, tr.bytes.len >> 20 }); + + const ids = try flash.mtp.draftIds(a); + const rows = try a.alloc(u16, ids.len * k); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try mtp3.draftRows(gpu.gpa, head, ids, rows, threads); + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e6; + std.debug.print("draft rows: {d} ids in {d} column blocks, {d} threads, {d:.0} ms (host, ReleaseFast)\n", .{ ids.len, blk: { + var c: usize = 0; + var prev: usize = std.math.maxInt(usize); + for (ids) |id| { + if (id / 128 != prev) c += 1; + prev = id / 128; + } + break :blk c; + }, threads, ms }); + + // x: a real embedding row (bf16) and a random bf16 row + const embed_path = try std.fs.path.join(a, &.{ dir, "model-00001-of-00007.safetensors" }); + var ef = try st.File.open(gpu.gpa, io, embed_path); + defer ef.close(io); + const emb = ef.get("model.language_model.embed_tokens.weight") orelse return error.MissingTensor; + const xs_bits = try a.alloc(u16, 2 * k); + for (0..k) |i| xs_bits[i] = std.mem.readInt(u16, emb.bytes[(9338 * k + i) * 2 ..][0..2], .little); + var prng = std.Random.DefaultPrng.init(0xd1a); + for (xs_bits[k..]) |*b| b.* = bf16Bits((prng.random().float(f32) - 0.5) * 2); + + // full head on the GPU, fp32 out + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + const strides = exl3.strides(k, n, head.halves, false); + const trellis_al = try a.alloc(i16, tr.bytes.len / 2); + @memcpy(std.mem.sliceAsBytes(trellis_al), tr.bytes); + var d_tr = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(trellis_al)); + defer d_tr.free(); + var d_su = try cuda.DeviceBuffer.fromHost(gpu.d, su.bytes); + defer d_su.free(); + var d_sv = try cuda.DeviceBuffer.fromHost(gpu.d, sv.bytes); + defer d_sv.free(); + const wgt: exl3.Weight = .{ .trellis = d_tr.ptr, .suh = d_su.ptr, .svh = d_sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = head.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + var d_x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xs_bits)); + defer d_x.free(); + var d_y = try cuda.DeviceBuffer.alloc(gpu.d, 2 * n * 4); + defer d_y.free(); + try gpu.ctx.synchronize(); + try de.decode(s, wgt, d_x.ptr, .bf16, 2, d_y.ptr, .f32); + try s.synchronize(); + const full = try a.alloc(f32, 2 * n); + try d_y.download(0, std.mem.sliceAsBytes(full)); + + // exact float64 for the first x (the exl3-pack test's `exact` product, one row) + const exact = try a.alloc(f64, n); + { + const tiles_k = k / 16; + const tiles_n = n / 16; + const tw: []const i16 = trellis_al; + const wq = try a.alloc(f16, k * n); + fmt.unpack(.{ .halves = head.halves }, .mul1, tw, tiles_k, tiles_n, wq); + const xr = try a.alloc(f64, k); + for (xr, 0..) |*v, i| v.* = bfVal(xs_bits[i]) * @as(f64, @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, su.bytes[2 * i ..][0..2], .little))))); + fmt.rotate(xr, 1, k, 1); + @memset(exact, 0); + for (0..k) |ki| { + const xv = xr[ki]; + for (exact, wq[ki * n ..][0..n]) |*o, wv| o.* += xv * @as(f64, @floatCast(wv)); + } + fmt.rotate(exact, 1, n, 1); + for (exact, 0..) |*v, c| v.* *= @as(f64, @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, sv.bytes[2 * c ..][0..2], .little))))); + } + + for (0..2) |xi| { + const x = xs_bits[xi * k ..][0..k]; + var worst_gpu: f64 = 0; // in units of the bf16 bound + var worst_exact: f64 = 0; + var max_abs_gpu: f64 = 0; + var peak: f64 = 0; + var arg_draft: usize = 0; + var arg_full: usize = 0; + var best_draft = -std.math.inf(f64); + var best_full = -std.math.inf(f64); + var gpu_exact_rel: f64 = 0; + var peak_full: f64 = 0; + for (full[xi * n ..][0..n]) |v| peak_full = @max(peak_full, @abs(@as(f64, v))); + for (ids, 0..) |id, i| { + const row = rows[i * k ..][0..k]; + var got: f64 = 0; + var mag: f64 = 0; + for (row, x) |wb, xb| { + const wv = bfVal(wb); + got += wv * bfVal(xb); + mag += @abs(wv * bfVal(xb)); + } + const bound = mag / 512.0 * (1.0 + 1.0 / 128.0) + 1e-9; // |W - bf16(W)| <= 2^-9 |bf16(W)| (1 + 2^-7) + const g: f64 = full[xi * n + id]; + // plus the kernel's fp16-input error, measured at 1.75e-4 - 2.4e-4 of the row's peak (exl3-pack) + const tol_gpu = bound + 5e-4 * peak_full; + worst_gpu = @max(worst_gpu, @abs(got - g) / tol_gpu); + max_abs_gpu = @max(max_abs_gpu, @abs(got - g)); + peak = @max(peak, @abs(g)); + if (got > best_draft) { + best_draft = got; + arg_draft = id; + } + if (g > best_full) { + best_full = g; + arg_full = id; + } + if (xi == 0) { + worst_exact = @max(worst_exact, @abs(got - exact[id]) / bound); + gpu_exact_rel = @max(gpu_exact_rel, @abs(g - exact[id])); + } + } + std.debug.print("x {s}: draft vs full head (GPU Dense fp32): max |diff| {e:.3} at draft peak {d:.3} (row peak {d:.3}), worst diff / tolerance {d:.3}; argmax over draft ids: draft {d}, full {d}\n", .{ if (xi == 0) "embed row 9338" else "random", max_abs_gpu, peak, peak_full, worst_gpu, arg_draft, arg_full }); + if (xi == 0) std.debug.print(" draft vs exact float64: worst diff / bf16 bound {d:.3}; GPU fp32 vs exact max |diff| {e:.3}\n", .{ worst_exact, gpu_exact_rel }); + try check.expect(worst_gpu <= 1.0, "draft logits off the full EXL3 head beyond bf16 weight rounding", .{}); + if (xi == 0) try check.expect(worst_exact <= 1.0, "draft logits off the exact product beyond the bf16 bound", .{}); + } + check.pass("EXL3 draft head: {d} rows dequantised from lm_head; every draft logit equals the full head's (GPU Dense fp32 and exact f64) within bf16 weight rounding", .{ids.len}); +} diff --git a/zig/tests/cuda/exl3_format.zig b/zig/tests/cuda/exl3_format.zig new file mode 120000 index 0000000000..3d9f7ecc0c --- /dev/null +++ b/zig/tests/cuda/exl3_format.zig @@ -0,0 +1 @@ +../../src/core/exl3_format.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_frequency_tests.zig b/zig/tests/cuda/exl3_frequency_tests.zig new file mode 100644 index 0000000000..6e374519ca --- /dev/null +++ b/zig/tests/cuda/exl3_frequency_tests.zig @@ -0,0 +1,99 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const profile = flash.exl3.frequency; +const Pool = flash.exl3_pool.Pool; + +fn verify(gpu: check.Gpu, pool: *Pool, expected: usize) !void { + const tables = try gpu.gpa.alloc(u8, pool.lines.tables.len); + defer gpu.gpa.free(tables); + try pool.lines.tables.download(0, tables); + var resident: usize = 0; + var bytes: [16]u8 = undefined; + for (0..512) |expert| { + var found: ?usize = null; + for (pool.slots_tab, pool.line_tab, 0..) |held, state, line| { + if (held == expert and state == .resident) found = line; + } + if (found != null) resident += 1; + for (0..3) |projection| { + const pointer = std.mem.readInt(u64, tables[pool.lines.tlay.ptr[projection] + expert * 8 ..][0..8], .little); + const width = std.mem.readInt(i32, tables[pool.lines.tlay.k2[projection] + expert * 4 ..][0..4], .little); + if (found) |line| { + const offset = line * pool.lines.layout.line_bytes + pool.lines.layout.off[projection]; + try check.expect(pointer == pool.lines.mem.ptr + offset and width == 1, "frequency warm table expert {d} projection {d}", .{ expert, projection }); + try pool.lines.mem.download(offset, &bytes); + try check.sameBytes("frequency warm expert bytes", &bytes, pool.src.routed[expert][projection]); + } else try check.expect(pointer == 0 and width == 0, "unloaded frequency expert has no device entry", .{}); + } + } + try check.expect(resident == expected, "frequency resident count {d} expected {d}", .{ resident, expected }); +} + +pub fn run(gpu: check.Gpu, arguments: []const [:0]const u8) !void { + if (arguments.len != 2) return error.MissingArgument; + const rank = try std.fmt.parseInt(u32, arguments[0], 10); + if (rank >= 2) return error.BadRank; + const path = arguments[1]; + const identity: profile.Identity = .{ .model = @splat(0x57), .world = 2, .rank = rank }; + const start = @as(usize, rank) * 256; + var bank = try cuda.HostBuffer.alloc(gpu.d, 512 * 3 * 16); + defer bank.free(); + var source: [512][3][]const u8 = @splat(.{ &.{}, &.{}, &.{} }); + for (start..start + 256) |expert| for (0..3) |projection| { + const offset = (expert * 3 + projection) * 16; + const bytes = bank.bytes[offset..][0..16]; + @memset(bytes, @truncate(expert + projection)); + source[expert][projection] = bytes; + }; + var saved: profile.Counts = .{}; + for (0..20) |index| saved.decoder[7][start + index] = @intCast(100 + index); + saved.decoder[7][(start + 256) % 512] = std.math.maxInt(u32); + try profile.write(gpu.gpa, gpu.io, path, identity, &saved); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + const scales: flash.exl3_pool.Scales = .{ .suh_g = 0, .suh_u = 0, .svh_g = 0, .svh_u = 0, .suh_d = 0, .svh_d = 0 }; + for ([_]usize{ 10, 23 }, 0..) |capacity, restart| { + const restored = try profile.read(gpu.gpa, gpu.io, path, identity); + var pool = try Pool.init(gpu.gpa, gpu.io, gpu.d, .{ .routed = &source }, scales, .{ + .layer = 7, + .lines = capacity, + .dims = 16, + .width = 16, + .max_rows = 1, + .direct = true, + }, null); + defer pool.deinit(); + @memcpy(&pool.counts[7], &restored.decoder[7]); + const before = pool.counts[7]; + const heated = try pool.warmProfile(); + try check.expect(heated == @min(capacity, 20), "frequency restart fills known local experts only", .{}); + try check.expect(std.mem.eql(u32, &before, &pool.counts[7]), "warming preserves observation counts", .{}); + try verify(gpu, &pool, heated); + var picks: [11]i32 = @splat(-1); + for (0..10) |index| picks[index] = @intCast(start + 19 - index); + var device_picks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&picks)); + defer device_picks.free(); + var planner = try flash.exl3.dplan.DevPlan.init(gpu.gpa, gpu.d, &.{&pool}); + defer planner.deinit(); + try planner.toDevice(stream); + try planner.plan(stream, 0, device_picks.ptr, 1); + try planner.join(stream); + try planner.restore(stream, 0); + try planner.toHost(); + for (0..20) |index| { + const expected = before[start + index] + @as(u32, @intFromBool(index >= 10)); + try check.expect(pool.counts[7][start + index] == expected, "device learned count survives profile restart", .{}); + } + @memcpy(&saved.decoder[7], &pool.counts[7]); + try profile.write(gpu.gpa, gpu.io, path, identity, &saved); + std.debug.print("FREQUENCY rank={d} restart={d} capacity={d} warmed={d}\n", .{ rank, restart, capacity, heated }); + } + var other_rank = identity; + other_rank.rank = 1 - rank; + if (profile.read(gpu.gpa, gpu.io, path, other_rank)) |_| return error.AcceptedForeignProfile else |failure| { + try check.expect(failure == error.IdentityMismatch, "frequency rejects other TP rank", .{}); + } + check.pass("frequency disk restart, live device counts, rank isolation and resized warm pools", .{}); +} diff --git a/zig/tests/cuda/exl3_fwd_cpu.zig b/zig/tests/cuda/exl3_fwd_cpu.zig new file mode 120000 index 0000000000..3d93aa45bb --- /dev/null +++ b/zig/tests/cuda/exl3_fwd_cpu.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_fwd_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_gdn_cpu.zig b/zig/tests/cuda/exl3_gdn_cpu.zig new file mode 120000 index 0000000000..66f52e2263 --- /dev/null +++ b/zig/tests/cuda/exl3_gdn_cpu.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_gdn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_hc.zig b/zig/tests/cuda/exl3_hc.zig new file mode 120000 index 0000000000..475a7bc779 --- /dev/null +++ b/zig/tests/cuda/exl3_hc.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_hc.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_kv.zig b/zig/tests/cuda/exl3_kv.zig new file mode 120000 index 0000000000..4caa5eb91c --- /dev/null +++ b/zig/tests/cuda/exl3_kv.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_kv.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_kv_full.zig b/zig/tests/cuda/exl3_kv_full.zig new file mode 100644 index 0000000000..c45b9d5493 --- /dev/null +++ b/zig/tests/cuda/exl3_kv_full.zig @@ -0,0 +1,411 @@ +//! exl3-kv full: all 48 layers of the EXL3 pack through the family's own pass code, one layer resident at a time +//! (layer streaming: every routed expert of that layer resident, then freed), and the logits of the last rows through +//! the mixer and the EXL3 lm_head, once per KV variant. The variants differ only in the attention caches: after each +//! attention block the rows it wrote are replaced by what an int8 / int4 cache would hold (fn_kv_roundtrip), so every +//! later row and layer reads the quantised keys and values through upstream's kernels. +//! Usage: tf-cuda-test exl3-kv full DIR PROBES CHUNK VARIANTS FILE[,FILE...] +//! PROBES the prompt's last tokens fed one row at a time (decode rows); their logits are compared +//! CHUNK prompt chunk rows (<= 256 keeps one attention block a chunk) +//! VARIANTS comma list of bf16 | int8r | int8 | int4r | int4 (the first is the reference; bf16 recommended) +//! Layers 0 and 1 (GDN, the PLE) run once through Pass.chunk (they have no attention cache); 2..47 step block by block. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const fx = @import("flashnext_exl3_tests.zig"); +const rig = @import("flashnext_rig.zig"); + +const Kv = flash.exl3.kv; +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const bf = rig.bf; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const vocab = wt.vocab; + +const Side = struct { bits: u32 = 16, rot: bool = false }; +/// kv8: the attention caches really are int8 (cuda_exl3_kv.zig attnBlock and its kernels), not an emulation. +const Variant = struct { name: []const u8, k: Side = .{}, v: Side = .{}, chunk: usize = 0, kv8: bool = false, rot: bool = false }; + +/// bf16 | int8 | int8r | int4 | int4r (K and V alike) | kbf16v8 | k8v4r (K and V apart), each with an optional @CHUNK +/// (the prompt fed in chunks of CHUNK rows instead of the default: with bf16 it is the run's numerical noise floor). +fn parseVariant(a: std.mem.Allocator, t: []const u8) !Variant { + var v: Variant = .{ .name = try a.dupe(u8, t) }; + var base = t; + if (std.mem.indexOfScalar(u8, t, '@')) |at| { + v.chunk = try std.fmt.parseInt(usize, t[at + 1 ..], 10); + base = t[0..at]; + } + const both = struct { + fn of(bits: u32, rot: bool) [2]Side { + return .{ .{ .bits = bits, .rot = rot }, .{ .bits = bits, .rot = rot } }; + } + }; + if (std.mem.eql(u8, base, "kv8") or std.mem.eql(u8, base, "kv8r")) { + v.kv8 = true; + v.rot = base.len == 4; + return v; + } + const kv: [2]Side = if (std.mem.eql(u8, base, "bf16")) both.of(16, false) else if (std.mem.eql(u8, base, "int8")) both.of(8, false) else if (std.mem.eql(u8, base, "int8r")) both.of(8, true) else if (std.mem.eql(u8, base, "int4")) both.of(4, false) else if (std.mem.eql(u8, base, "int12")) both.of(12, false) else if (std.mem.eql(u8, base, "int10")) both.of(10, false) else if (std.mem.eql(u8, base, "int6")) both.of(6, false) else if (std.mem.eql(u8, base, "int4r")) both.of(4, true) else if (std.mem.eql(u8, base, "kbf16v8")) .{ .{}, .{ .bits = 8 } } else if (std.mem.eql(u8, base, "k8v4r")) .{ .{ .bits = 8 }, .{ .bits = 4, .rot = true } } else return error.UnknownVariant; + v.k = kv[0]; + v.v = kv[1]; + return v; +} + +/// One variant's sequence state on the device, laid out as Engine.newSeq makes it. +const VSeq = struct { + conv: [wt.linear_n]cuda.DeviceBuffer = undefined, + state: [wt.linear_n]cuda.DeviceBuffer = undefined, + kc: [wt.attn_n]cuda.DeviceBuffer = undefined, + vc: [wt.attn_n]cuda.DeviceBuffer = undefined, + ikc: [wt.attn_n]cuda.DeviceBuffer = undefined, + pooled: [wt.attn_n]cuda.DeviceBuffer = undefined, + tail: cuda.DeviceBuffer = undefined, + history: [2]i64 = .{ 0, 0 }, + cap: usize = 0, + /// The int8 caches of a kv8 variant: codes and scales for K and V, each attention layer. + c8: [wt.attn_n]Kv.Kv8.Cache = undefined, + c8_bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn zeroed(d: *const cuda.Driver, s: cuda.Stream, n: usize) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(d, n); + try b.fill8(0, s.handle); + return b; + } + + fn init(v: *VSeq, gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, cap: usize, eos: i64, kv8: bool) !void { + const gdn = flash.gdn; + if (kv8) { + const cb = Kv.Kv8.cacheBytes(cap, wt.kv_heads); + for (&v.c8) |*c| { + var parts: [4]u64 = undefined; + for (&parts, 0..) |*pp, i| { + const b = try zeroed(d, s, if (i % 2 == 0) cb.codes else cb.scales); + try v.c8_bufs.append(gpa, b); + pp.* = b.ptr; + } + c.* = .{ .k8 = parts[0], .ks = parts[1], .v8 = parts[2], .vs = parts[3] }; + } + } + v.cap = cap; + v.history = .{ eos, eos }; + for (&v.conv, &v.state) |*c, *x| { + c.* = try zeroed(d, s, (gdn.taps - 1) * gdn.conv_dim * 2); + x.* = try zeroed(d, s, wt.state_n * 4); + } + for (&v.kc, &v.vc, &v.ikc, &v.pooled) |*kc, *vc, *ikc, *pooled| { + kc.* = try zeroed(d, s, cap * wt.kv_heads * wt.head_dim * 2); + vc.* = try zeroed(d, s, cap * wt.kv_heads * wt.head_dim * 2); + ikc.* = try zeroed(d, s, cap * wt.index_dim * 2); + pooled.* = try zeroed(d, s, flash.prompt.pooledBytes(cap)); + } + v.tail = try zeroed(d, s, wt.ple_tail_n * wide * 2); + try s.synchronize(); + } + + fn deinit(v: *VSeq, gpa: std.mem.Allocator) void { + for (v.c8_bufs.items) |*b| b.free(); + v.c8_bufs.deinit(gpa); + for (&v.conv, &v.state) |*c, *x| { + c.free(); + x.free(); + } + for (&v.kc, &v.vc, &v.ikc, &v.pooled) |*a, *b, *c, *e| { + a.free(); + b.free(); + c.free(); + e.free(); + } + v.tail.free(); + } + + fn view(v: *VSeq, scores: u64) batch.SeqState { + return .{ .conv = &v.conv, .state = &v.state, .kc = &v.kc, .vc = &v.vc, .ikc = &v.ikc, .pooled = &v.pooled, .tail = v.tail, .history = &v.history, .cap = v.cap, .scores = scores }; + } +}; + +/// One decoder layer (>= 2: no PLE) for `n` rows at `pos0` from the streams in `h` (host, bf16), in place. +fn stepLayer(gpu: check.Gpu, p: *fx.Pass, k8: *const Kv.Kv8, v: Variant, vq: *const VSeq, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h: []u16) !void { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try fx.uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + if (lw.linear) { + try p.gdnBlock(lw, wt.linearIndex(l), .{}); + } else { + const ai = wt.attnIndex(l); + if (v.kv8) return stepRest(gpu, p, lw, l, n, h, Kv.attnBlock(p, k8, lw, ai, vq.c8[ai], v.rot)); + try p.attnBlock(lw, ai); + const row = wt.kv_heads * wt.head_dim * 2; + if (v.k.bits != 16) try k8.roundTrip(p.stream, p.seq.kc[ai].ptr + pos0 * row, n * wt.kv_heads, v.k.rot, v.k.bits); + if (v.v.bits != 16) try k8.roundTrip(p.stream, p.seq.vc[ai].ptr + pos0 * row, n * wt.kv_heads, v.v.rot, v.v.bits); + } + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + try p.readout(lw.mlp_hc, 0, n); + try p.moeBlock(lw); + p.pending = l; + try p.writeback(0, n); + const out = try fx.download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(out); + @memcpy(h, out); +} + +/// The rest of stepLayer after an attention block that returned `mixer` (the int8 block). +fn stepRest(gpu: check.Gpu, p: *fx.Pass, lw: *const pl.Layer, l: usize, n: usize, h: []u16, mixer: anyerror!void) !void { + try mixer; + const s = p.s; + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + try p.readout(lw.mlp_hc, 0, n); + try p.moeBlock(lw); + p.pending = l; + try p.writeback(0, n); + const out = try fx.download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(out); + @memcpy(h, out); +} + +/// The chunks a prompt of `total` tokens is fed in: prefill chunks of `chunk` rows, then `probes` single rows. +fn chunksOf(a: std.mem.Allocator, total: usize, probes: usize, chunk: usize) ![][2]usize { + var list: std.ArrayList([2]usize) = .empty; + var at: usize = 0; + while (at < total - probes) { + const n = @min(chunk, total - probes - at); + try list.append(a, .{ at, n }); + at += n; + } + while (at < total) : (at += 1) try list.append(a, .{ at, 1 }); + return list.toOwnedSlice(a); +} + +fn logSoftmax(row: []const u16, out: []f64) void { + var best = -std.math.inf(f64); + for (row) |x| best = @max(best, bf(x)); + var sum: f64 = 0; + for (row) |x| sum += @exp(bf(x) - best); + const lse = best + @log(sum); + for (row, out) |x, *o| o.* = bf(x) - lse; +} + +fn top5(row: []const u16) [5]u32 { + var best: [5]u32 = @splat(0); + var val: [5]f64 = @splat(-std.math.inf(f64)); + for (row, 0..) |x, i| { + const v = bf(x); + if (v <= val[4]) continue; + var j: usize = 4; + while (j > 0 and v > val[j - 1]) : (j -= 1) { + val[j] = val[j - 1]; + best[j] = best[j - 1]; + } + val[j] = v; + best[j] = @intCast(i); + } + return best; +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 5) return error.MissingArgument; + const dir = rest[0]; + const probes = try std.fmt.parseInt(usize, rest[1], 10); + const chunk = try std.fmt.parseInt(usize, rest[2], 10); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var variants: std.ArrayList(Variant) = .empty; + { + var it = std.mem.splitScalar(u8, rest[3], ','); + while (it.next()) |t| try variants.append(a, try parseVariant(a, t)); + } + var ids: std.ArrayList(u32) = .empty; + { + var it = std.mem.splitScalar(u8, rest[4], ','); + while (it.next()) |path| { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 26)); + var toks = std.mem.tokenizeAny(u8, text, ", \n\r\t"); + while (toks.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + } + } + const total = ids.items.len; + if (probes == 0 or probes >= total) return error.TooManyProbes; + const cap = total + 16; + const nv = variants.items.len; + // logits from row STATS_FROM (6th argument; default: every row of a short prompt, else the decode rows) to the end + const stats_from: usize = if (rest.len > 5) try std.fmt.parseInt(usize, rest[5], 10) else if (total <= 64) 0 else total - probes; + const block: usize = 64; + std.debug.print("== exl3-kv full: {d} tokens ({d} decode rows), chunk {d}, 48 layers streamed, variants:", .{ total, probes, chunk }); + for (variants.items) |v| std.debug.print(" {s};", .{v.name}); + std.debug.print("\n", .{}); + const t_all = check.now(gpu.io); + + var r: fx.Rig = undefined; + try r.init(gpu, dir, @max(chunk, block), 16); + defer r.deinit(); + const d: *cuda.Driver = @constCast(gpu.d); + var k8 = try Kv.Kv8.load(gpu.d); + defer k8.deinit(); + const eos = try flash.ple.eosOf(gpu.gpa, gpu.io, dir); + const vs = try a.alloc(VSeq, nv); + for (vs) |*v| v.* = .{}; + var made: usize = 0; + defer for (vs[0..made]) |*v| v.deinit(gpu.gpa); + for (vs, variants.items) |*v, vr| { + try v.init(gpu.gpa, gpu.d, r.stream, cap, eos, vr.kv8); + made += 1; + } + var max_chunk = chunk; + for (variants.items) |v| max_chunk = @max(max_chunk, v.chunk); + if (max_chunk > chunk) return error.VariantChunkAboveScratch; // a variant's chunk may only be smaller + var scores = try cuda.DeviceBuffer.alloc(gpu.d, batch.scoreBytes(cap, chunk)); + defer scores.free(); + const plan = try chunksOf(a, total, probes, chunk); + const plans = try a.alloc([][2]usize, nv); + for (variants.items, plans) |v, *pv| pv.* = if (v.chunk == 0) plan else try chunksOf(a, total, probes, v.chunk); + const hs = try a.alloc([]u16, nv); + for (hs) |*h| h.* = try a.alloc(u16, total * wide); + + // layers 0 and 1 (embedding, GDN, the PLE) through Pass.chunk, the streams after layer 1 kept: the same for every variant + { + var table = try flash.embed.Table.uploadDense(gpu.d, try r.store.tensor(flash.embed.weight_name)); + defer table.deinit(); + var ngram = try flash.ple.Ngram.init(gpu.gpa, gpu.io, dir); + defer ngram.deinit(); + var first: [2]pl.Layer = undefined; + first[0] = try pl.loadLayerRank(gpu.gpa, d, &r.stream, &r.store, r.pack_fn, 0, 1, 0, null); + defer first[0].free(gpu.gpa); + first[1] = try pl.loadLayerRank(gpu.gpa, d, &r.stream, &r.store, r.pack_fn, 1, 1, 0, null); + defer first[1].free(gpu.gpa); + r.store.release(); + var hbuf = try cuda.DeviceBuffer.alloc(gpu.d, chunk * wide * 2); + defer hbuf.free(); + var p = r.pass(&first); + p.seq = vs[0].view(scores.ptr); + p.m.table_w = table.w.ptr; + p.m.table_dense = table.dense; + p.m.ngram = &ngram; + p.hooks = .{ .logits = .none, .streams = hbuf.ptr }; + for (plan) |c| { + try p.chunk(ids.items[c[0]..][0..c[1]], c[0], null); + const got = try fx.download(u16, gpu, r.stream, hbuf.ptr, c[1] * wide); + defer gpu.gpa.free(got); + @memcpy(hs[0][c[0] * wide ..][0 .. c[1] * wide], got); + } + for (hs[1..]) |h| @memcpy(h, hs[0]); + std.debug.print(" layers 0-1 (embed, GDN, PLE) for {d} rows in {d:.1} s\n", .{ total, @as(f64, @floatFromInt(check.now(gpu.io) - t_all)) / 1e9 }); + } + + // layers 2..47, one resident at a time, every variant through it + var p = r.pass(&.{}); + p.hooks = .{ .logits = .none }; + var load_ns: i96 = 0; + var run_ns: i96 = 0; + for (2..wt.layers_n) |l| { + var t0 = check.now(gpu.io); + var lw = try pl.loadLayerRank(gpu.gpa, d, &r.stream, &r.store, r.pack_fn, l, 1, 0, null); + defer lw.free(gpu.gpa); + r.store.release(); + load_ns += check.now(gpu.io) - t0; + t0 = check.now(gpu.io); + for (variants.items, vs, hs, plans) |v, *vq, h, pv| { + p.seq = vq.view(scores.ptr); + for (pv) |c| try stepLayer(gpu, &p, &k8, v, vq, &lw, l, ids.items[c[0]..][0..c[1]], c[0], h[c[0] * wide ..][0 .. c[1] * wide]); + } + run_ns += check.now(gpu.io) - t0; + if (l % 8 == 7) std.debug.print(" layer {d}/47: loads {d:.1} s, passes {d:.1} s so far\n", .{ l, @as(f64, @floatFromInt(load_ns)) / 1e9, @as(f64, @floatFromInt(run_ns)) / 1e9 }); + } + + // the head: mixer read-out and the EXL3 lm_head over the kept rows, a variant at a time + var mixer = try wt.loadHc(gpu.gpa, d, &r.store, "language_model.model.hyper_connection_mixer", false); + defer mixer.free(); + var head = try X.projection_load.load(a, d, &r.store, "language_model.lm_head", vocab, dims); + defer head.free(); + r.store.release(); + var logits_d = try cuda.DeviceBuffer.alloc(gpu.d, block * vocab * 2); + defer logits_d.free(); + const lg = try a.alloc([]u16, nv); + for (lg) |*x| x.* = try a.alloc(u16, block * vocab); + const Acc = struct { rows: usize = 0, correct: usize = 0, agree: usize = 0, agree5: usize = 0, kl_sum: f64 = 0, kl_max: f64 = 0, flips: [4]usize = @splat(0), max_d: f64 = 0 }; + const acc = try a.alloc(Acc, nv); + for (acc) |*x| x.* = .{}; + const lp_ref = try a.alloc(f64, vocab); + const lp = try a.alloc(f64, vocab); + const margins = [_]f64{ 0.5, 1, 2, 4 }; + var b0 = stats_from; + while (b0 < total) : (b0 += block) { + const n = @min(block, total - b0); + for (hs, 0..) |h, vi| { + p.rows = n; + p.pending = null; + try fx.uploadTo(gpu, p.stream, p.s.h, std.mem.sliceAsBytes(h[b0 * wide ..][0 .. n * wide])); + try p.writeback(0, n); + try p.readout(mixer, 0, n); + try p.lane(p.s.mixed, p.s.xsm, head, logits_d.ptr, p.s.dummy, n, 1); + const got = try fx.download(u16, gpu, r.stream, logits_d.ptr, n * vocab); + defer gpu.gpa.free(got); + @memcpy(lg[vi][0 .. n * vocab], got); + } + for (0..n) |i| { + const pos = b0 + i; + const next: ?u32 = if (pos + 1 < total) ids.items[pos + 1] else null; + const ra = lg[0][i * vocab ..][0..vocab]; + const ta = top5(ra); + if (total <= 64) std.debug.print(" pos {d} -> {?d} | {any}\n", .{ pos, next, ta }); + logSoftmax(ra, lp_ref); + for (acc, lg, variants.items) |*ac, l, v| { + const rb = l[i * vocab ..][0..vocab]; + const tb = top5(rb); + ac.rows += 1; + if (next) |nx| ac.correct += @intFromBool(tb[0] == nx); + if (ta[0] == tb[0]) ac.agree += 1 else { + const m = bf(ra[ta[0]]) - bf(ra[tb[0]]); + for (margins, 0..) |mm, mi| ac.flips[mi] += @intFromBool(m > mm); + if (m > 1) std.debug.print(" FLIP {s} pos {d} (next {?d}): ref top-3 {any} {d:.3} {d:.3} {d:.3}; var top-3 {any} {d:.3} {d:.3} {d:.3}\n", .{ v.name, pos, next, ta[0..3], bf(ra[ta[0]]), bf(ra[ta[1]]), bf(ra[ta[2]]), tb[0..3], bf(rb[tb[0]]), bf(rb[tb[1]]), bf(rb[tb[2]]) }); + } + for (tb) |x| ac.agree5 += @intFromBool(x == ta[0]); + for (ra, rb) |x, y| ac.max_d = @max(ac.max_d, @abs(bf(x) - bf(y))); + logSoftmax(rb, lp); + var kl: f64 = 0; + for (lp_ref, lp) |x, y| kl += @exp(x) * (x - y); + ac.kl_sum += kl; + ac.kl_max = @max(ac.kl_max, kl); + } + } + } + // the prompt's last rows (a needle prompt's answer): every variant's top-1 against the next token + { + const tail = @min(6, total - stats_from); + const t0r = total - tail; + // recompute the logits of the tail block (the loop above kept only the last block) + for (hs, 0..) |h, vi| { + p.rows = tail; + p.pending = null; + try fx.uploadTo(gpu, p.stream, p.s.h, std.mem.sliceAsBytes(h[t0r * wide ..][0 .. tail * wide])); + try p.writeback(0, tail); + try p.readout(mixer, 0, tail); + try p.lane(p.s.mixed, p.s.xsm, head, logits_d.ptr, p.s.dummy, tail, 1); + const got = try fx.download(u16, gpu, r.stream, logits_d.ptr, tail * vocab); + defer gpu.gpa.free(got); + @memcpy(lg[vi][0 .. tail * vocab], got); + } + for (0..tail) |i| { + const pos = t0r + i; + std.debug.print(" tail pos {d} next {?d}:", .{ pos, if (pos + 1 < total) @as(?u32, ids.items[pos + 1]) else null }); + for (variants.items, lg) |v, l| std.debug.print(" {s}={d}", .{ v.name, top5(l[i * vocab ..][0..vocab])[0] }); + std.debug.print("\n", .{}); + } + } + const rows_n = total - stats_from; + for (acc, variants.items) |ac, v| { + std.debug.print("RESULT {s:<12} rows {d}..{d}: teacher-forced top-1 = next token {d}/{d}; vs {s}: top-1 equal {d}/{d}, its top-1 in our top-5 {d}/{d}, flips with reference margin > 0.5/1/2/4: {d}/{d}/{d}/{d}; KL mean {e:.3} max {e:.3}; max |logit diff| {d:.3}\n", .{ v.name, stats_from, total - 1, ac.correct, rows_n - 1, variants.items[0].name, ac.agree, rows_n, ac.agree5, rows_n, ac.flips[0], ac.flips[1], ac.flips[2], ac.flips[3], ac.kl_sum / @as(f64, @floatFromInt(rows_n)), ac.kl_max, ac.max_d }); + } + std.debug.print(" time: layer loads {d:.1} s, layer passes {d:.1} s, total {d:.1} s\n", .{ @as(f64, @floatFromInt(load_ns)) / 1e9, @as(f64, @floatFromInt(run_ns)) / 1e9, @as(f64, @floatFromInt(check.now(gpu.io) - t_all)) / 1e9 }); + check.pass("exl3-kv full: 48 layers streamed over {d} tokens, {d} variants; numbers above", .{ total, nv }); +} diff --git a/zig/tests/cuda/exl3_kv_real.zig b/zig/tests/cuda/exl3_kv_real.zig new file mode 100644 index 0000000000..6592c01900 --- /dev/null +++ b/zig/tests/cuda/exl3_kv_real.zig @@ -0,0 +1,436 @@ +//! exl3-kv real: the first LAYERS decoder layers of the EXL3 pack through the family's own pass (Pass.chunk: embed, +//! PLE, GDN, attention with the sparse indexer, EXL3 MoE) over a real long prompt, every routed expert resident. +//! Usage: tf-cuda-test exl3-kv real DIR LAYERS PROBES CHUNKS FILE[,FILE...] +//! LAYERS 4 or 8 (the last layer is attention: 3 or 7); PROBES the prompt's last tokens fed one row at a time; +//! CHUNKS prompt chunk sizes to time, e.g. 256,512,1536 (the first one drives the analysis); FILE comma id lists. +//! It reports, for the last attention layer: +//! - the indexer's selection overlap between consecutive rows (what a device cache of recent selections would hit) +//! and the distinct keys a 256-row block reads (what a host-resident cache would move over PCIe per block); +//! - at the probe rows, the attention output from real Q/K/V: upstream's kernels (re-run; bit-equal to the pass's), +//! int8 plain and rotated (flashnext_exl3_kv.cu), and int8/int4 caches in float64, each against float64 of the +//! bf16 cache; +//! - the prompt's wall time per chunk size over these layers (the scratch chunk lever). +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const fx = @import("flashnext_exl3_tests.zig"); +const kt = @import("exl3_kv_tests.zig"); +const rig = @import("flashnext_rig.zig"); + +const Kv = flash.exl3.kv; +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const bf = rig.bf; +const bf16 = rig.bf16; + +const q_heads = 24; +const kv_heads = 2; +const head_dim = 256; +const idw = wt.index_idw; +const top = wt.index_top; +const ratio = wt.index_ratio; +const nch = wt.nch; +const group = 32; +const ng = head_dim / group; + +fn readIds(gpa: std.mem.Allocator, io: std.Io, paths: []const u8, out: *std.ArrayList(u32)) !void { + var it = std.mem.splitScalar(u8, paths, ','); + while (it.next()) |p| { + const text = try std.Io.Dir.cwd().readFileAlloc(io, p, gpa, .limited(1 << 26)); + defer gpa.free(text); + var toks = std.mem.tokenizeAny(u8, text, ", \n\r\t"); + while (toks.next()) |t| try out.append(gpa, try std.fmt.parseInt(u32, t, 10)); + } +} + +/// A row's selected block set (sparse rows) as sorted block ids; dense rows: null. +fn blocksOf(ids: []const i32, nk: i32, spr: i32, out: []u32) ?[]u32 { + if (spr == 0) return null; + _ = nk; + for (0..top) |i| out[i] = @intCast(@divTrunc(ids[i * ratio], @as(i32, ratio))); + return out[0..top]; +} + +fn overlap(a: []const u32, b: []const u32) usize { + var i: usize = 0; + var j: usize = 0; + var n: usize = 0; + while (i < a.len and j < b.len) { + if (a[i] == b[j]) { + n += 1; + i += 1; + j += 1; + } else if (a[i] < b[j]) i += 1 else j += 1; + } + return n; +} + +const Stats = struct { + rows: usize = 0, + sum1: f64 = 0, + min1: f64 = 1, + sum4: f64 = 0, + rows4: usize = 0, + sum16: f64 = 0, + rows16: usize = 0, + hist: [11]usize = @splat(0), + blocks: usize = 0, + distinct: usize = 0, + max_distinct: usize = 0, +}; + +fn has(set: []const u32, x: u32) bool { + var lo: usize = 0; + var hi: usize = set.len; + while (lo < hi) { + const mid = (lo + hi) / 2; + if (set[mid] == x) return true; + if (set[mid] < x) lo = mid + 1 else hi = mid; + } + return false; +} + +/// Each sparse row against the previous row's set, and against the union of the last 4 and 16 rows' sets. +const Window = struct { + sets: std.ArrayList([]u32) = .empty, + + fn add(w: *Window, a: std.mem.Allocator, set: []const u32, st: *Stats) !void { + const n = w.sets.items.len; + if (n > 0) { + const o1 = @as(f64, @floatFromInt(overlap(set, w.sets.items[n - 1]))) / top; + st.sum1 += o1; + st.min1 = @min(st.min1, o1); + st.hist[@min(10, @as(usize, @intFromFloat(o1 * 10)))] += 1; + for ([_]usize{ 4, 16 }) |k| { + if (n < k) continue; + var hit: usize = 0; + for (set) |blk| { + for (w.sets.items[n - k ..]) |prev| { + if (has(prev, blk)) { + hit += 1; + break; + } + } + } + const f = @as(f64, @floatFromInt(hit)) / top; + if (k == 4) { + st.sum4 += f; + st.rows4 += 1; + } else { + st.sum16 += f; + st.rows16 += 1; + } + } + st.rows += 1; + } + try w.sets.append(a, try a.dupe(u32, set)); + } +}; + +const Probe = struct { pos: usize, q: []u16, ao: []u16, ids: []i32, nk: i32, spr: i32 }; + +fn errRow(got: []const f64, want: []const f64, e: *kt.Err) void { + for (got, want) |x, w| { + e.max_abs = @max(e.max_abs, @abs(x - w)); + e.peak = @max(e.peak, @abs(w)); + e.ss += (x - w) * (x - w); + e.ss_ref += w * w; + } +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 5) return error.MissingArgument; + const dir = rest[0]; + const n_layers = try std.fmt.parseInt(usize, rest[1], 10); + const probes = try std.fmt.parseInt(usize, rest[2], 10); + var chunk_list: std.ArrayList(usize) = .empty; + defer chunk_list.deinit(gpu.gpa); + { + var it = std.mem.splitScalar(u8, rest[3], ','); + while (it.next()) |t| try chunk_list.append(gpu.gpa, try std.fmt.parseInt(usize, t, 10)); + } + if (n_layers == 0 or n_layers > wt.layers_n or wt.linearLayer(n_layers - 1)) return error.LastLayerNotAttention; + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var ids: std.ArrayList(u32) = .empty; + try readIds(a, gpu.io, rest[4], &ids); + const total = ids.items.len; + if (probes >= total) return error.TooManyProbes; + const n_prefill = total - probes; + var max_chunk: usize = 0; + for (chunk_list.items) |c| max_chunk = @max(max_chunk, c); + const cap = total + 16; + std.debug.print("== exl3-kv real: {d} prompt tokens ({d} prefilled, {d} probe rows), layers 0..{d}, chunks {any}\n", .{ total, n_prefill, probes, n_layers - 1, chunk_list.items }); + + var r: fx.Rig = undefined; + try r.init(gpu, dir, max_chunk, cap); + defer r.deinit(); + const d: *cuda.Driver = @constCast(gpu.d); + var t0 = check.now(gpu.io); + var table = try flash.embed.Table.uploadDense(gpu.d, try r.store.tensor(flash.embed.weight_name)); + defer table.deinit(); + var ngram = try flash.ple.Ngram.init(gpu.gpa, gpu.io, dir); + defer ngram.deinit(); + const eos = try flash.ple.eosOf(gpu.gpa, gpu.io, dir); + const layers = try a.alloc(pl.Layer, n_layers); + var loaded: usize = 0; + defer for (layers[0..loaded]) |*lw| lw.free(gpu.gpa); + for (0..n_layers) |l| { + layers[l] = try pl.loadLayerRank(gpu.gpa, d, &r.stream, &r.store, r.pack_fn, l, 1, 0, null); + loaded += 1; + if (l % 16 == 15) r.store.release(); + } + r.store.release(); + var free_b: usize = 0; + var total_b: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free_b, &total_b), "cuMemGetInfo"); + std.debug.print(" loaded table, n-gram reader and {d} layers (every routed expert resident) in {d:.1} s; device free {d:.2} of {d:.2} GiB\n", .{ n_layers, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9, @as(f64, @floatFromInt(free_b)) / Kv.gib, @as(f64, @floatFromInt(total_b)) / Kv.gib }); + + var p = r.pass(layers); + p.m.table_w = table.w.ptr; + p.m.table_dense = table.dense; + p.m.ngram = &ngram; + p.hooks.logits = .none; + const ai = wt.attnIndex(n_layers - 1); + const s = p.s; + + // -- prompt passes: the first chunk size drives the analysis; the others only time it ------------------------------ + var st_prefill: Stats = .{}; + var win: Window = .{}; + var blk_buf: [top]u32 = undefined; + var times: std.ArrayList(f64) = .empty; + for (chunk_list.items, 0..) |chunk, ci| { + try r.reset(); + r.history = .{ eos, eos }; + try r.stream.synchronize(); + var at: usize = 0; + t0 = check.now(gpu.io); + var analysis_ns: i96 = 0; + while (at < n_prefill) { + const n = @min(chunk, n_prefill - at); + try p.chunk(ids.items[at..][0..n], at, null); + try r.stream.synchronize(); + if (ci == 0 and n <= X.batch.att_rows) { + // n <= att_rows: the selection buffers hold every row of this chunk (the last attention layer's) + const ta = check.now(gpu.io); + const qids = try fx.download(i32, gpu, r.stream, s.qids, n * idw); + defer gpu.gpa.free(qids); + const qnk = try fx.download(i32, gpu, r.stream, s.qnk, n); + defer gpu.gpa.free(qnk); + const qspr = try fx.download(i32, gpu, r.stream, s.qspr, n); + defer gpu.gpa.free(qspr); + var seen = try std.DynamicBitSetUnmanaged.initEmpty(a, cap); + defer seen.deinit(a); + var any_sparse = false; + for (0..n) |row| { + const nkr: usize = @intCast(qnk[row]); + if (qspr[row] != 0) { + any_sparse = true; + for (qids[row * idw ..][0..nkr]) |k| seen.set(@intCast(k)); + if (blocksOf(qids[row * idw ..][0..idw], qnk[row], qspr[row], &blk_buf)) |set| try win.add(a, set, &st_prefill); + } else for (0..nkr) |k| seen.set(k); + } + if (any_sparse) { + st_prefill.blocks += 1; + st_prefill.distinct += seen.count(); + st_prefill.max_distinct = @max(st_prefill.max_distinct, seen.count()); + } + analysis_ns += check.now(gpu.io) - ta; + } + at += n; + } + const ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0 - analysis_ns)) / 1e6; + try times.append(a, ms); + std.debug.print("RESULT prompt {d} tokens in chunks of {d} through layers 0..{d}: {d:.0} ms ({d:.3} ms a token)\n", .{ n_prefill, chunk, n_layers - 1, ms, ms / @as(f64, @floatFromInt(n_prefill)) }); + } + if (st_prefill.rows > 0) { + const rows_f: f64 = @floatFromInt(st_prefill.rows); + std.debug.print("RESULT layer {d} selection, {d} consecutive sparse prompt rows: overlap with the previous row mean {d:.3} min {d:.3}; covered by the last 4 rows' union {d:.3}, last 16 rows' {d:.3}\n", .{ n_layers - 1, st_prefill.rows, st_prefill.sum1 / rows_f, st_prefill.min1, st_prefill.sum4 / @as(f64, @floatFromInt(@max(1, st_prefill.rows4))), st_prefill.sum16 / @as(f64, @floatFromInt(@max(1, st_prefill.rows16))) }); + std.debug.print(" overlap histogram (tenths): {any}\n", .{st_prefill.hist}); + std.debug.print("RESULT layer {d}: a sparse 256-row prompt block reads {d:.0} distinct keys on average, {d} at most ({d} blocks)\n", .{ n_layers - 1, @as(f64, @floatFromInt(st_prefill.distinct)) / @as(f64, @floatFromInt(st_prefill.blocks)), st_prefill.max_distinct, st_prefill.blocks }); + } + + // -- probe rows: one token a pass (a decode row each), the pass's own q, selection and attention output kept ---------- + const probe = try a.alloc(Probe, probes); + var st_decode: Stats = .{}; + for (0..probes) |i| { + const pos = n_prefill + i; + try p.chunk(ids.items[pos..][0..1], pos, null); + try r.stream.synchronize(); + const q = try fx.download(u16, gpu, r.stream, s.q, q_heads * head_dim); + const ao = try fx.download(u16, gpu, r.stream, s.ao, q_heads * head_dim); + const qids = try fx.download(i32, gpu, r.stream, s.qids, idw); + const qnk = try fx.download(i32, gpu, r.stream, s.qnk, 1); + const qspr = try fx.download(i32, gpu, r.stream, s.qspr, 1); + probe[i] = .{ .pos = pos, .q = try a.dupe(u16, q), .ao = try a.dupe(u16, ao), .ids = try a.dupe(i32, qids), .nk = qnk[0], .spr = qspr[0] }; + if (blocksOf(qids, qnk[0], qspr[0], &blk_buf)) |set| try win.add(a, set, &st_decode); + gpu.gpa.free(q); + gpu.gpa.free(ao); + gpu.gpa.free(qids); + gpu.gpa.free(qnk); + gpu.gpa.free(qspr); + } + if (st_decode.rows > 0) { + const rows_f: f64 = @floatFromInt(st_decode.rows); + std.debug.print("RESULT layer {d} selection, {d} decode rows: overlap with the previous row mean {d:.3} min {d:.3}; last 4 rows' union {d:.3}, last 16 {d:.3}\n", .{ n_layers - 1, st_decode.rows, st_decode.sum1 / rows_f, st_decode.min1, st_decode.sum4 / @as(f64, @floatFromInt(@max(1, st_decode.rows4))), st_decode.sum16 / @as(f64, @floatFromInt(@max(1, st_decode.rows16))) }); + } + + // -- the cache variants at each probe row ------------------------------------------------------------------------ + var k8 = try Kv.Kv8.load(gpu.d); + defer k8.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + const nrow = cap * kv_heads; + const kc = try fx.download(u16, gpu, r.stream, r.kc[ai].ptr, nrow * head_dim); + defer gpu.gpa.free(kc); + const vc = try fx.download(u16, gpu, r.stream, r.vc[ai].ptr, nrow * head_dim); + defer gpu.gpa.free(vc); + const used = total * kv_heads; // rows past the prompt are zero + // the host's codes for every variant, and their float64 values in the model's domain + const Variant = struct { name: []const u8, rot: bool, int4: bool }; + const variants = [_]Variant{ .{ .name = "int8 plain", .rot = false, .int4 = false }, .{ .name = "int8 rotated", .rot = true, .int4 = false }, .{ .name = "int4 plain", .rot = false, .int4 = true }, .{ .name = "int4 rotated", .rot = true, .int4 = true } }; + const kf0 = try a.alloc(f64, nrow * head_dim); + const vf0 = try a.alloc(f64, nrow * head_dim); + for (kf0, kc) |*o, x| o.* = bf(x); + for (vf0, vc) |*o, x| o.* = bf(x); + var vals: [variants.len][2][]f64 = undefined; + for (variants, 0..) |v, vi| for (0..2) |kvi| { + const src = if (kvi == 0) kc else vc; + const out = try a.alloc(f64, nrow * head_dim); + @memset(out, 0); + for (0..used) |row| { + for (0..ng) |gi| { + var x: [group]f32 = undefined; + for (0..group) |j| x[j] = @floatCast(bf(src[row * head_dim + gi * group + j])); + if (v.rot) kt.wht32(&x); + var y: [group]f32 = undefined; + if (v.int4) kt.roundTrip4(&x, &y) else { + var codes: [group]i8 = undefined; + const sc = kt.quant8(&x, &codes); + for (0..group) |j| y[j] = kt.dequant8(codes[j], sc); + } + for (0..group) |j| { + var acc: f64 = 0; + if (v.rot) { + for (0..group) |m| acc += (if (@popCount(j & m) % 2 == 0) @as(f64, 1) else -1) * y[m]; + acc /= @sqrt(32.0); + } else acc = y[j]; + out[row * head_dim + gi * group + j] = acc; + } + } + } + vals[vi][kvi] = out; + }; + // the device's int8 caches (rotated and plain), packed from the pass's own bf16 cache + var codes: [2]Kv.Kv8.Cache = undefined; + var cbufs: [8]cuda.DeviceBuffer = undefined; + const cb = Kv.Kv8.cacheBytes(cap, kv_heads); + for (0..8) |i| cbufs[i] = try cuda.DeviceBuffer.alloc(gpu.d, if (i % 2 == 0) cb.codes else cb.scales); + defer for (&cbufs) |*x| x.free(); + for (0..2) |ri| { + codes[ri] = .{ .k8 = cbufs[ri * 4].ptr, .ks = cbufs[ri * 4 + 1].ptr, .v8 = cbufs[ri * 4 + 2].ptr, .vs = cbufs[ri * 4 + 3].ptr }; + try k8.packRows(r.stream, r.kc[ai].ptr, codes[ri].k8, codes[ri].ks, nrow, ri == 1); + try k8.packRows(r.stream, r.vc[ai].ptr, codes[ri].v8, codes[ri].vs, nrow, ri == 1); + } + var qb = try cuda.DeviceBuffer.alloc(gpu.d, q_heads * head_dim * 2); + defer qb.free(); + var ib = try cuda.DeviceBuffer.alloc(gpu.d, idw * 4 + 32); + defer ib.free(); + var posb = try cuda.DeviceBuffer.alloc(gpu.d, 16); + defer posb.free(); + var outb = try cuda.DeviceBuffer.alloc(gpu.d, q_heads * head_dim * 2); + defer outb.free(); + const want = try a.alloc(f64, q_heads * head_dim); + const got = try a.alloc(f64, q_heads * head_dim); + const sbuf = try a.alloc(f64, cap + idw); + const kbuf = try a.alloc(i32, cap); + var e_up: kt.Err = .{}; + var e_dev: [2]kt.Err = .{ .{}, .{} }; + var e_host: [variants.len]kt.Err = @splat(.{}); + var e_own: [2]kt.Err = .{ .{}, .{} }; + var replay_equal: usize = 0; + var flips_up: usize = 0; + for (probe) |pr| { + const n: usize = @intCast(pr.nk); + const keys: []const i32 = if (pr.spr != 0) pr.ids[0..n] else blk: { + for (0..n) |i| kbuf[i] = @intCast(i); + break :blk kbuf[0..n]; + }; + for (0..q_heads) |h| oracleRow(pr.q[h * head_dim ..][0..head_dim], kf0, vf0, keys, h, want[h * head_dim ..][0..head_dim], sbuf); + // the pass's own output (upstream's kernels inside Pass.attnRows) + for (got, pr.ao) |*o, x| o.* = bf(x); + errRow(got, want, &e_up); + // upstream's kernels re-run here on the same buffers: the same bits (the harness reproduces the pass) + try posb.upload(0, std.mem.sliceAsBytes(&[_]i32{ @intCast(pr.pos), 0, 0, 0 })); + try ib.upload(0, std.mem.sliceAsBytes(pr.ids)); + try ib.upload(idw * 4, std.mem.sliceAsBytes(&[_]i32{ pr.nk, pr.spr })); + try qb.upload(0, std.mem.sliceAsBytes(pr.q)); + const chunks = if (pr.spr != 0) nch else (pr.pos + 1 + 511) / 512; + try g.attnChunks(r.stream, qb.ptr, r.kc[ai].ptr, r.vc[ai].ptr, posb.ptr, s.po, s.pm, s.pl, ib.ptr, ib.ptr + idw * 4, ib.ptr + idw * 4 + 4, 1, kv_heads, chunks, true); + try g.attnMerge(r.stream, s.po, s.pm, s.pl, posb.ptr, outb.ptr, ib.ptr + idw * 4, ib.ptr + idw * 4 + 4, 1, kv_heads, true); + const again = try fx.download(u16, gpu, r.stream, outb.ptr, q_heads * head_dim); + replay_equal += @intFromBool(std.mem.eql(u16, again, pr.ao)); + gpu.gpa.free(again); + // int8 on the device: plain, and rotated (the query rotated from the pass's bf16 query, then rounded) + for (0..2) |ri| { + const rot = ri == 1; + const qr = try a.alloc(u16, q_heads * head_dim); + for (0..q_heads * ng) |gi| { + var x: [group]f32 = undefined; + for (0..group) |j| x[j] = @floatCast(bf(pr.q[gi * group + j])); + if (rot) kt.wht32(&x); + for (0..group) |j| qr[gi * group + j] = bf16(x[j]); + } + try qb.upload(0, std.mem.sliceAsBytes(qr)); + try k8.attnChunks(r.stream, qb.ptr, codes[ri], posb.ptr, s.po, s.pm, s.pl, ib.ptr, ib.ptr + idw * 4, ib.ptr + idw * 4 + 4, 1, q_heads, kv_heads, chunks, true); + try k8.attnMerge(r.stream, s.po, s.pm, s.pl, posb.ptr, outb.ptr, ib.ptr + idw * 4, ib.ptr + idw * 4 + 4, 1, q_heads, kv_heads, true, rot); + const o8 = try fx.download(u16, gpu, r.stream, outb.ptr, q_heads * head_dim); + defer gpu.gpa.free(o8); + for (got, o8) |*o, x| o.* = bf(x); + errRow(got, want, &e_dev[ri]); + // against the float64 of its own cache (the kernel's arithmetic alone) + const own = try a.alloc(f64, q_heads * head_dim); + const vi: usize = if (rot) 1 else 0; + for (0..q_heads) |h| oracleRow(pr.q[h * head_dim ..][0..head_dim], vals[vi][0], vals[vi][1], keys, h, own[h * head_dim ..][0..head_dim], sbuf); + errRow(got, own, &e_own[ri]); + } + // every variant's cache in float64 (no kernel): the cache's own error + for (variants, 0..) |_, vi| { + for (0..q_heads) |h| oracleRow(pr.q[h * head_dim ..][0..head_dim], vals[vi][0], vals[vi][1], keys, h, got[h * head_dim ..][0..head_dim], sbuf); + errRow(got, want, &e_host[vi]); + } + } + _ = &flips_up; + std.debug.print("RESULT layer {d} attention at {d} probe rows ({d}..{d}), against float64 of the pass's bf16 cache (max | rms, of the peak):\n", .{ n_layers - 1, probes, n_prefill, total - 1 }); + std.debug.print(" upstream bf16 (the pass's own output) {e:.3} | {e:.3} (re-run here bit-equal on {d}/{d} rows)\n", .{ kt.rel(e_up), kt.rms(e_up), replay_equal, probes }); + std.debug.print(" int8 plain device kernel {e:.3} | {e:.3} (kernel vs its own cache {e:.3})\n", .{ kt.rel(e_dev[0]), kt.rms(e_dev[0]), kt.rel(e_own[0]) }); + std.debug.print(" int8 rotated device kernel {e:.3} | {e:.3} (kernel vs its own cache {e:.3})\n", .{ kt.rel(e_dev[1]), kt.rms(e_dev[1]), kt.rel(e_own[1]) }); + for (variants, 0..) |v, vi| std.debug.print(" {s:<13} cache alone (float64) {e:.3} | {e:.3}\n", .{ v.name, kt.rel(e_host[vi]), kt.rms(e_host[vi]) }); + try check.expect(replay_equal == probes, "exl3-kv real: upstream's kernels re-run on the pass's buffers differ on {d} rows", .{probes - replay_equal}); + check.pass("exl3-kv real: layers 0..{d} over {d} tokens; numbers above", .{ n_layers - 1, total }); +} + +fn oracleRow(q: []const u16, kf: []const f64, vf: []const f64, keys: []const i32, h: usize, out: []f64, sbuf: []f64) void { + const kvh = h / (q_heads / kv_heads); + var best = -std.math.inf(f64); + for (keys, 0..) |key, j| { + const k: usize = @intCast(key); + var dot: f64 = 0; + for (0..head_dim) |dd| dot += bf(q[dd]) * kf[(k * kv_heads + kvh) * head_dim + dd]; + sbuf[j] = dot * 0.0625; + best = @max(best, sbuf[j]); + } + var l: f64 = 0; + for (sbuf[0..keys.len]) |*v| { + v.* = @exp(v.* - best); + l += v.*; + } + for (0..head_dim) |dd| { + var o: f64 = 0; + for (keys, 0..) |key, j| o += sbuf[j] * vf[(@as(usize, @intCast(key)) * kv_heads + kvh) * head_dim + dd]; + out[dd] = o / l; + } +} diff --git a/zig/tests/cuda/exl3_kv_tests.zig b/zig/tests/cuda/exl3_kv_tests.zig new file mode 100644 index 0000000000..bb3ce1ab6f --- /dev/null +++ b/zig/tests/cuda/exl3_kv_tests.zig @@ -0,0 +1,767 @@ +//! exl3-kv: the KV-memory levers for 200k contexts (docs/plans/exl3-kv-memory.md), on one card. +//! Usage: tf-cuda-test exl3-kv kernel int8 KV kernels (flashnext_exl3_kv.cu) vs the host codes, vs upstream's +//! chunk/merge kernels on the dequantised bf16 cache (bits), vs float64 +//! tf-cuda-test exl3-kv bench [CAP] decode and prompt attention over a CAP-token cache (200,016): bf16 +//! (upstream) and int8 (ours), device or mapped pinned host memory; the +//! copy-engine and gather alternatives; host round trip of the ids +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Kv = flash.exl3.kv; +const Rig = rig.Rig; +const bf = rig.bf; +const bf16 = rig.bf16; + +const q_heads = 24; +const kv_heads = 2; +const head_dim = 256; +const index_heads = 4; +const index_dim = 128; +const half = 32; +const proj_n = 13952; +const eps: f32 = 1e-6; +const idw = 2052; +const top = 512; +const ratio = 4; +const nch = 5; +const group = 32; +const ng = head_dim / group; + +// -- the host's codes, bit for bit what fn_kv8_pack / fn_kv8_prep write ----------------------------------------------- + +const inv_sqrt32: f32 = 0.17677669529663688; + +/// The normalized Sylvester Hadamard of 32 values in fp32, the kernel's butterfly order. +pub fn wht32(x: *[group]f32) void { + var m: usize = 1; + while (m < group) : (m <<= 1) { + var y: [group]f32 = undefined; + for (0..group) |i| y[i] = if (i & m != 0) x[i ^ m] - x[i] else x[i] + x[i ^ m]; + x.* = y; + } + for (x) |*v| v.* *= inv_sqrt32; +} + +/// One group of 32 values (already rotated if wanted) into its codes and fp16 scale. +pub fn quant8(x: *const [group]f32, codes: *[group]i8) f16 { + var peak: f32 = 0; + for (x) |v| peak = @max(peak, @abs(v)); + const sh: f16 = @floatCast(@max(peak, 6.103515625e-05)); + const s: f32 = @floatCast(sh); + for (x, codes) |v, *c| { + const q = std.math.clamp(@floor(v / s * 128.0) + 128.0, 0.0, 255.0); + c.* = @intCast(@as(i32, @intFromFloat(q)) - 128); + } + return sh; +} + +pub fn dequant8(c: i8, s: f16) f32 { + return (@as(f32, @floatFromInt(c)) + 0.5) * (@as(f32, @floatCast(s)) * 0.0078125); +} + +/// int4 (host only, for the quality comparison): exl3_cache.zig quantize4's grid, the scale rounded to fp16 first. +pub fn roundTrip4(x: *const [group]f32, out: *[group]f32) void { + var peak: f32 = 0; + for (x) |v| peak = @max(peak, @abs(v)); + const s: f32 = @floatCast(@as(f16, @floatCast(@max(peak, 6.103515625e-05)))); + for (x, out) |v, *o| { + const q = std.math.clamp(@floor(v / s * 8.0) + 8.0, 0.0, 15.0); + o.* = (q - 7.5) * s * 0.125; + } +} + +/// A row of 256 bf16 into codes and scales (rot: each group rotated first). +pub fn packRow(x: []const u16, rot: bool, codes: []i8, scales: []f16) void { + for (0..ng) |gi| { + var g: [group]f32 = undefined; + for (0..group) |j| g[j] = @floatCast(bf(x[gi * group + j])); + if (rot) wht32(&g); + scales[gi] = quant8(&g, codes[gi * group ..][0..group]); + } +} + +/// The dequantised row as the chunk kernel's bf16 tile holds it (still in the rotated domain when rot). +pub fn tileRow(codes: []const i8, scales: []const f16, out: []u16) void { + for (0..head_dim) |d| out[d] = bf16(dequant8(codes[d], scales[d / group])); +} + +/// The dequantised row back in the model's domain, in float64 (the oracle's view of what the cache holds). +pub fn valueRow(codes: []const i8, scales: []const f16, rot: bool, out: []f64) void { + for (0..ng) |gi| { + var g: [group]f64 = undefined; + for (0..group) |j| g[j] = dequant8(codes[gi * group + j], scales[gi]); + if (rot) { + // H is symmetric and orthonormal: its own inverse + var y: [group]f64 = undefined; + for (0..group) |i| { + var acc: f64 = 0; + for (0..group) |j| acc += (if (@popCount(i & j) % 2 == 0) @as(f64, 1) else -1) * g[j]; + y[i] = acc / @sqrt(32.0); + } + g = y; + } + @memcpy(out[gi * group ..][0..group], &g); + } +} + +// -- kernel mode ------------------------------------------------------------------------------------------------------- + +/// Keys with a few loud channels, as k_norm + rope leaves them (a lognormal scale a channel). +fn loudRows(r: *Rig, n: usize, width: usize) ![]u16 { + const scale = try r.a.alloc(f32, width); + for (scale) |*s| s.* = @exp(r.rng.floatNorm(f32) * 0.8); + for (0..4) |_| scale[r.rng.uintLessThan(usize, width)] *= 8; + const v = try r.a.alloc(u16, n * width); + for (v, 0..) |*x, i| x.* = bf16(r.rng.floatNorm(f32) * scale[i % width] * 0.5); + return v; +} + +/// Sorted distinct blocks: `count` of `complete`, then each block's 4 keys, then the tail (fn_qsa_select's layout). +fn sparseIds(r: *Rig, end: usize, out: []i32) usize { + const complete = end / ratio; + var chosen = std.DynamicBitSetUnmanaged.initEmpty(r.a, complete) catch unreachable; + var got: usize = 0; + while (got < top) { + const b = r.rng.uintLessThan(usize, complete); + if (chosen.isSet(b)) continue; + chosen.set(b); + got += 1; + } + var k: usize = 0; + var it = chosen.iterator(.{}); + while (it.next()) |b| { + for (0..ratio) |j| out[k * ratio + j] = @intCast(b * ratio + j); + k += 1; + } + var n: usize = top * ratio; + for (ratio * complete..end) |t| { + out[n] = @intCast(t); + n += 1; + } + return n; +} + +/// One row's attention in float64 over value rows `kf`/`vf` [key][head][256] (f64) and bf16 queries. +fn oracleRow(q: []const u16, kf: []const f64, vf: []const f64, keys: []const i32, h: usize, out: []f64, sbuf: []f64) void { + const kvh = h / (q_heads / kv_heads); + var best = -std.math.inf(f64); + for (keys, 0..) |key, j| { + const k: usize = @intCast(key); + var dot: f64 = 0; + for (0..head_dim) |d| dot += bf(q[d]) * kf[(k * kv_heads + kvh) * head_dim + d]; + sbuf[j] = dot * 0.0625; + best = @max(best, sbuf[j]); + } + var l: f64 = 0; + for (sbuf[0..keys.len]) |*v| { + v.* = @exp(v.* - best); + l += v.*; + } + for (0..head_dim) |d| { + var o: f64 = 0; + for (keys, 0..) |key, j| o += sbuf[j] * vf[(@as(usize, @intCast(key)) * kv_heads + kvh) * head_dim + d]; + out[d] = o / l; + } +} + +pub const Err = struct { max_abs: f64 = 0, peak: f64 = 0, ss: f64 = 0, ss_ref: f64 = 0 }; + +fn addErr(e: *Err, got: f64, want: f64) void { + e.max_abs = @max(e.max_abs, @abs(got - want)); + e.peak = @max(e.peak, @abs(want)); + e.ss += (got - want) * (got - want); + e.ss_ref += want * want; +} + +pub fn rel(e: Err) f64 { + return e.max_abs / e.peak; +} + +pub fn rms(e: Err) f64 { + return @sqrt(e.ss / e.ss_ref); +} + +fn kernelMode(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(0x6b76_38); + var r = try Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var k8 = try Kv.Kv8.load(gpu.d); + defer k8.deinit(); + const a = r.a; + const cap = 4096; + const kc = try loudRows(&r, cap * kv_heads, head_dim); + const vc = try loudRows(&r, cap * kv_heads, head_dim); + const nrow = cap * kv_heads; + var failed = false; + + // 1. fn_kv8_pack == the host's codes and scales, both domains + for ([_]bool{ false, true }) |rot| { + const codes = try a.alloc(i8, nrow * head_dim); + const scales = try a.alloc(f16, nrow * ng); + for (0..nrow) |i| packRow(kc[i * head_dim ..][0..head_dim], rot, codes[i * head_dim ..][0..head_dim], scales[i * ng ..][0..ng]); + const dx = try r.dev(std.mem.sliceAsBytes(kc)); + const dc = try r.out(nrow * head_dim); + const ds = try r.out(nrow * ng * 2); + try k8.packRows(r.s, dx, dc, ds, nrow, rot); + try check.sameBytes(if (rot) "pack codes (rotated)" else "pack codes", std.mem.sliceAsBytes(try r.back(i8, dc, codes.len)), std.mem.sliceAsBytes(codes)); + try check.sameBytes(if (rot) "pack scales (rotated)" else "pack scales", std.mem.sliceAsBytes(try r.back(f16, ds, scales.len)), std.mem.sliceAsBytes(scales)); + } + check.pass("kv8 pack: fn_kv8_pack's codes and fp16 scales equal the host's bit for bit ({d} rows of 256, plain and rotated)", .{nrow}); + + // 2. prep: upstream's Q, indexer rows, and its K/V cache rows packed == fn_kv8_prep's (plain); rotated K/V too + { + const rows = 17; + const pos0 = 700; + const p = r.bfs(rows * proj_n, 2); + const qw = r.f32s(head_dim, 0.5, 1.5); + const kw = r.f32s(head_dim, 0.5, 1.5); + const iw = r.f32s(index_dim, 0.5, 1.5); + const inv = try a.alloc(f32, half); + for (inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + const dp = try r.dev(std.mem.sliceAsBytes(p)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const dqw = try r.dev(std.mem.sliceAsBytes(qw)); + const dkw = try r.dev(std.mem.sliceAsBytes(kw)); + const diw = try r.dev(std.mem.sliceAsBytes(iw)); + const dinv = try r.dev(std.mem.sliceAsBytes(inv)); + const end = pos0 + rows; + const uq = try r.out(rows * q_heads * head_dim * 2); + const ukc = try r.out(end * kv_heads * head_dim * 2); + const uvc = try r.out(end * kv_heads * head_dim * 2); + const uiq = try r.out(rows * index_heads * index_dim * 2); + const uik = try r.out(end * index_dim * 2); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, uq, ukc, uvc, uiq, uik, eps, rows); + const ukc_h = try r.back(u16, ukc, end * kv_heads * head_dim); + const uvc_h = try r.back(u16, uvc, end * kv_heads * head_dim); + for ([_]bool{ false, true }) |rot| { + const oq = try r.out(rows * q_heads * head_dim * 2); + const c: Kv.Kv8.Cache = .{ .k8 = try r.out(end * kv_heads * head_dim), .ks = try r.out(end * kv_heads * ng * 2), .v8 = try r.out(end * kv_heads * head_dim), .vs = try r.out(end * kv_heads * ng * 2) }; + const oiq = try r.out(rows * index_heads * index_dim * 2); + const oik = try r.out(end * index_dim * 2); + try k8.prepRows(r.s, dp, dpos, dqw, dkw, diw, dinv, oq, c, oiq, oik, eps, rows, proj_n, q_heads, kv_heads, rot); + if (!rot) try check.sameBytes("prep q", std.mem.sliceAsBytes(try r.back(u16, oq, rows * q_heads * head_dim)), std.mem.sliceAsBytes(try r.back(u16, uq, rows * q_heads * head_dim))); + try check.sameBytes("prep iq", std.mem.sliceAsBytes(try r.back(u16, oiq, rows * index_heads * index_dim)), std.mem.sliceAsBytes(try r.back(u16, uiq, rows * index_heads * index_dim))); + try check.sameBytes("prep ik", std.mem.sliceAsBytes((try r.back(u16, oik, end * index_dim))[pos0 * index_dim ..]), std.mem.sliceAsBytes((try r.back(u16, uik, end * index_dim))[pos0 * index_dim ..])); + const n = rows * kv_heads; + const want_c = try a.alloc(i8, n * head_dim); + const want_s = try a.alloc(f16, n * ng); + for ([_]u64{ c.k8, c.v8 }, [_]u64{ c.ks, c.vs }, [_][]const u16{ ukc_h, uvc_h }) |dc, ds, src| { + for (0..n) |i| packRow(src[(pos0 * kv_heads + i) * head_dim ..][0..head_dim], rot, want_c[i * head_dim ..][0..head_dim], want_s[i * ng ..][0..ng]); + try check.sameBytes("prep k/v codes", std.mem.sliceAsBytes((try r.back(i8, dc, end * kv_heads * head_dim))[pos0 * kv_heads * head_dim ..]), std.mem.sliceAsBytes(want_c)); + try check.sameBytes("prep k/v scales", std.mem.sliceAsBytes((try r.back(f16, ds, end * kv_heads * ng))[pos0 * kv_heads * ng ..]), std.mem.sliceAsBytes(want_s)); + } + if (rot) { + // a rotated query, rotated back on the host, is upstream's query within one bf16 rounding of each side + const got = try r.back(u16, oq, rows * q_heads * head_dim); + const want = try r.back(u16, uq, rows * q_heads * head_dim); + var worst: f64 = 0; + for (0..rows * q_heads * ng) |gi| { + var x: [group]f64 = undefined; + for (0..group) |i| { + var acc: f64 = 0; + for (0..group) |j| acc += (if (@popCount(i & j) % 2 == 0) @as(f64, 1) else -1) * bf(got[gi * group + j]); + x[i] = acc / @sqrt(32.0); + } + var peak: f64 = 0; + for (0..group) |i| peak = @max(peak, @abs(bf(want[gi * group + i]))); + for (0..group) |i| worst = @max(worst, @abs(x[i] - bf(want[gi * group + i])) / (peak + 1e-30)); + } + std.debug.print(" rotated query rotated back vs upstream's query: worst {e:.3} of the group's peak\n", .{worst}); + try check.expect(worst < 1.6e-2, "rotated query off by {e:.3}", .{worst}); + } + } + check.pass("kv8 prep: upstream's q, index q/k bit-equal; K/V codes and scales == the host's packing of upstream's bf16 cache rows (plain and rotated)", .{}); + } + + // 3. attention: int8 kernels vs upstream's kernels on the dequantised bf16 cache (bits, plain), and vs float64 + for ([_]bool{ false, true }) |sparse| { + const rows = 20; // up to 16 rows take the read-ahead kernel; 20 one tile at a time (both run below) + const pos0: usize = if (sparse) 3000 else 1100; + const q = r.bfs(rows * q_heads * head_dim, 1); + const ids = try a.alloc(i32, rows * idw); + const nk = try a.alloc(i32, rows); + const spr = try a.alloc(i32, rows); + @memset(ids, 0); + for (0..rows) |row| { + const end = pos0 + row + 1; + spr[row] = @intFromBool(sparse); + nk[row] = @intCast(if (sparse) sparseIds(&r, end, ids[row * idw ..][0..idw]) else end); + if (!sparse) for (0..end) |i| { + ids[row * idw + i] = @intCast(@min(i, idw - 1)); + }; + } + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ @intCast(pos0), 0, 0, 0 })); + const dids = try r.dev(std.mem.sliceAsBytes(ids)); + const dnk = try r.dev(std.mem.sliceAsBytes(nk)); + const dspr = try r.dev(std.mem.sliceAsBytes(spr)); + const po = try r.out(rows * nch * q_heads * head_dim * 4); + const pm = try r.out(rows * nch * q_heads * 4); + const pl = try r.out(rows * nch * q_heads * 4); + const chunks = if (sparse) nch else (pos0 + rows + 511) / 512; + // the float64 reference over the ORIGINAL bf16 cache, shared by every variant + const kf0 = try a.alloc(f64, nrow * head_dim); + const vf0 = try a.alloc(f64, nrow * head_dim); + for (kf0, kc) |*o, x| o.* = bf(x); + for (vf0, vc) |*o, x| o.* = bf(x); + const want = try a.alloc(f64, rows * q_heads * head_dim); + const sbuf = try a.alloc(f64, pos0 + rows + idw); + const keys_of = struct { + fn get(ids_: []const i32, nk_: []const i32, row: usize, buf: []i32, sp: bool) []const i32 { + const n: usize = @intCast(nk_[row]); + if (sp) return ids_[row * idw ..][0..n]; + for (0..n) |i| buf[i] = @intCast(i); + return buf[0..n]; + } + }; + const kbuf = try a.alloc(i32, pos0 + rows); + for (0..rows) |row| for (0..q_heads) |h| oracleRow(q[(row * q_heads + h) * head_dim ..][0..head_dim], kf0, vf0, keys_of.get(ids, nk, row, kbuf, sparse), h, want[(row * q_heads + h) * head_dim ..][0..head_dim], sbuf); + // upstream on the original bf16 cache: its error is the baseline + const dkc = try r.dev(std.mem.sliceAsBytes(kc)); + const dvc = try r.dev(std.mem.sliceAsBytes(vc)); + const base_out = try r.out(rows * q_heads * head_dim * 2); + try g.attnChunks(r.s, try r.dev(std.mem.sliceAsBytes(q)), dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, rows, kv_heads, chunks, true); + try g.attnMerge(r.s, po, pm, pl, dpos, base_out, dnk, dspr, rows, kv_heads, true); + var e_base: Err = .{}; + for (try r.back(u16, base_out, want.len), want) |x, w| addErr(&e_base, bf(x), w); + for ([_]bool{ false, true }) |rot| { + const codes_k = try a.alloc(i8, nrow * head_dim); + const codes_v = try a.alloc(i8, nrow * head_dim); + const sc_k = try a.alloc(f16, nrow * ng); + const sc_v = try a.alloc(f16, nrow * ng); + for (0..nrow) |i| { + packRow(kc[i * head_dim ..][0..head_dim], rot, codes_k[i * head_dim ..][0..head_dim], sc_k[i * ng ..][0..ng]); + packRow(vc[i * head_dim ..][0..head_dim], rot, codes_v[i * head_dim ..][0..head_dim], sc_v[i * ng ..][0..ng]); + } + const c: Kv.Kv8.Cache = .{ .k8 = try r.dev(std.mem.sliceAsBytes(codes_k)), .ks = try r.dev(std.mem.sliceAsBytes(sc_k)), .v8 = try r.dev(std.mem.sliceAsBytes(codes_v)), .vs = try r.dev(std.mem.sliceAsBytes(sc_v)) }; + // queries as the prep stores them: rotated (fp32) then rounded once + const qr = try a.alloc(u16, q.len); + for (0..q.len / group) |gi| { + var x: [group]f32 = undefined; + for (0..group) |j| x[j] = @floatCast(bf(q[gi * group + j])); + if (rot) wht32(&x); + for (0..group) |j| qr[gi * group + j] = bf16(x[j]); + } + const dq = try r.dev(std.mem.sliceAsBytes(qr)); + var outs: [2]u64 = undefined; + for ([_]usize{ 16, rows }, 0..) |m, ci| { + outs[ci] = try r.out(m * q_heads * head_dim * 2); + try k8.attnChunks(r.s, dq, c, dpos, po, pm, pl, dids, dnk, dspr, m, q_heads, kv_heads, chunks, true); + try k8.attnMerge(r.s, po, pm, pl, dpos, outs[ci], dnk, dspr, m, q_heads, kv_heads, true, rot); + } + const got = try r.back(u16, outs[1], want.len); + try check.sameBytes("kv8 attention: 16 rows (read-ahead kernel) == the same rows of 20 (one tile at a time)", std.mem.sliceAsBytes(try r.back(u16, outs[0], 16 * q_heads * head_dim)), std.mem.sliceAsBytes(got[0 .. 16 * q_heads * head_dim])); + if (!rot) { + // upstream's kernels on a bf16 cache holding the tile values: the same bits + const tk = try a.alloc(u16, nrow * head_dim); + const tv = try a.alloc(u16, nrow * head_dim); + for (0..nrow) |i| { + tileRow(codes_k[i * head_dim ..][0..head_dim], sc_k[i * ng ..][0..ng], tk[i * head_dim ..][0..head_dim]); + tileRow(codes_v[i * head_dim ..][0..head_dim], sc_v[i * ng ..][0..ng], tv[i * head_dim ..][0..head_dim]); + } + const up = try r.out(rows * q_heads * head_dim * 2); + try g.attnChunks(r.s, dq, try r.dev(std.mem.sliceAsBytes(tk)), try r.dev(std.mem.sliceAsBytes(tv)), dpos, po, pm, pl, dids, dnk, dspr, rows, kv_heads, chunks, true); + try g.attnMerge(r.s, po, pm, pl, dpos, up, dnk, dspr, rows, kv_heads, true); + try check.sameBytes("kv8 attention == upstream's chunk/merge on the dequantised bf16 cache", std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(try r.back(u16, up, want.len))); + } + // vs the float64 of the original cache, and vs the float64 of what the int8 cache holds (its own kernel error) + var e_q: Err = .{}; + for (got, want) |x, w| addErr(&e_q, bf(x), w); + const kf = try a.alloc(f64, nrow * head_dim); + const vf = try a.alloc(f64, nrow * head_dim); + for (0..nrow) |i| { + valueRow(codes_k[i * head_dim ..][0..head_dim], sc_k[i * ng ..][0..ng], rot, kf[i * head_dim ..][0..head_dim]); + valueRow(codes_v[i * head_dim ..][0..head_dim], sc_v[i * ng ..][0..ng], rot, vf[i * head_dim ..][0..head_dim]); + } + const own = try a.alloc(f64, want.len); + for (0..rows) |row| for (0..q_heads) |h| oracleRow(q[(row * q_heads + h) * head_dim ..][0..head_dim], kf, vf, keys_of.get(ids, nk, row, kbuf, sparse), h, own[(row * q_heads + h) * head_dim ..][0..head_dim], sbuf); + var e_own: Err = .{}; + for (got, own) |x, w| addErr(&e_own, bf(x), w); + var e_cache: Err = .{}; + for (own, want) |x, w| addErr(&e_cache, x, w); + std.debug.print(" {s} {s}: vs float64 of the bf16 cache max {e:.3} rms {e:.3} of peak (upstream bf16: {e:.3} / {e:.3}); int8 cache alone (f64) {e:.3} / {e:.3}; kernel vs f64 of its own cache {e:.3} / {e:.3}\n", .{ if (sparse) "sparse" else "dense", if (rot) "int8 rotated" else "int8 plain", rel(e_q), rms(e_q), rel(e_base), rms(e_base), rel(e_cache), rms(e_cache), rel(e_own), rms(e_own) }); + // the kernel's own error stays at upstream's (bf16 tiles, bf16 p): within 2x of the baseline's + if (rel(e_own) > 2 * rel(e_base) + 1e-3) { + std.debug.print("FAIL kv8 attention: kernel error {e:.3} vs its own cache's float64 is more than twice upstream's {e:.3}\n", .{ rel(e_own), rel(e_base) }); + failed = true; + } + } + } + try check.expect(!failed, "exl3-kv kernel: a check failed", .{}); + check.pass("kv8 attention: plain int8 == upstream's kernels on the dequantised cache bit for bit (dense, sparse, read-ahead and one-tile); rotated within upstream's own error of its cache's float64", .{}); +} + +// -- bench mode --------------------------------------------------------------------------------------------------------- + +const Bench = struct { + s: cuda.Stream, + ev0: cuda.Event, + ev1: cuda.Event, + reps: usize = 20, + + fn time(b: *Bench, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try b.s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try b.ev0.record(b.s); + for (0..b.reps) |_| try f(ctx); + try b.ev1.record(b.s); + try b.ev1.synchronize(); + best = @min(best, @as(f64, try b.ev0.elapsedMs(b.ev1)) / @as(f64, @floatFromInt(b.reps))); + } + return best; + } +}; + +/// A cache of `cap` tokens in one place: bf16 K/V and its int8 codes + scales (same values, packed on the device). +const Place = struct { + kc: u64, + vc: u64, + c: Kv.Kv8.Cache, +}; + +const Attn = struct { + g: *const flash.Glue, + k8: *const Kv.Kv8, + s: cuda.Stream, + p: Place, + int8: bool, + q: u64, + pos: u64, + po: u64, + pm: u64, + pl: u64, + ids: u64, + nk: u64, + spr: u64, + out: u64, + rows: usize, + chunks: usize, + + fn go(a: *const Attn) anyerror!void { + if (a.int8) { + try a.k8.attnChunks(a.s, a.q, a.p.c, a.pos, a.po, a.pm, a.pl, a.ids, a.nk, a.spr, a.rows, q_heads, kv_heads, a.chunks, true); + try a.k8.attnMerge(a.s, a.po, a.pm, a.pl, a.pos, a.out, a.nk, a.spr, a.rows, q_heads, kv_heads, true, true); + } else { + try a.g.attnChunks(a.s, a.q, a.p.kc, a.p.vc, a.pos, a.po, a.pm, a.pl, a.ids, a.nk, a.spr, a.rows, kv_heads, a.chunks, true); + try a.g.attnMerge(a.s, a.po, a.pm, a.pl, a.pos, a.out, a.nk, a.spr, a.rows, kv_heads, true); + } + } +}; + +/// `rows` consecutive sparse rows ending at `end0 + rows`: row 0 picks `top` random blocks, each next row keeps a +/// fraction `keep` of the previous row's blocks and draws the rest anew (the selection overlap a real prompt shows). +fn overlapIds(rng: std.Random, a: std.mem.Allocator, end0: usize, rows: usize, keep: f64, ids: []i32, nk: []i32) !void { + const complete0 = (end0 + 1) / ratio; + var cur = try std.DynamicBitSetUnmanaged.initEmpty(a, complete0 + rows); + var got: usize = 0; + while (got < top) { + const b = rng.uintLessThan(usize, complete0); + if (cur.isSet(b)) continue; + cur.set(b); + got += 1; + } + for (0..rows) |row| { + const end = end0 + row + 1; + const complete = end / ratio; + if (row > 0) { + // drop (1 - keep) of the set, refill from all complete blocks + var drop: usize = @intFromFloat(@round((1 - keep) * @as(f64, top))); + while (drop > 0) { + const b = rng.uintLessThan(usize, complete); + if (!cur.isSet(b)) continue; + cur.unset(b); + drop -= 1; + } + while (cur.count() < top) { + const b = rng.uintLessThan(usize, complete); + if (!cur.isSet(b)) cur.set(b); + } + } + var k: usize = 0; + var it = cur.iterator(.{}); + while (it.next()) |b| { + for (0..ratio) |j| ids[row * idw + k * ratio + j] = @intCast(b * ratio + j); + k += 1; + } + var n: usize = top * ratio; + for (ratio * complete..end) |t| { + ids[row * idw + n] = @intCast(t); + n += 1; + } + nk[row] = @intCast(n); + } +} + +const Copy = struct { + s: cuda.Stream, + dst: u64, + src: [*]const u8, + bytes: usize, + pieces: usize, + fn go(c: *const Copy) anyerror!void { + const d = c.s.d; + const each = c.bytes / c.pieces; + for (0..c.pieces) |i| try d.check(d.api.cuMemcpyHtoDAsync_v2(c.dst + i * each, c.src + i * each, each, c.s.handle), "cuMemcpyHtoDAsync"); + } +}; + +const Gather = struct { + k8: *const Kv.Kv8, + s: cuda.Stream, + src_k: u64, + src_v: u64, + dst_k: u64, + dst_v: u64, + ids: u64, + n: usize, + row: usize, + fn go(g: *const Gather) anyerror!void { + try g.k8.gatherRows(g.s, g.src_k, g.dst_k, g.ids, g.n, g.row); + try g.k8.gatherRows(g.s, g.src_v, g.dst_v, g.ids, g.n, g.row); + } +}; + +fn benchMode(gpu: check.Gpu, cap: usize) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0xbe7c4); + const rng = prng.random(); + var r = try Rig.init(gpu, a, rng); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var k8 = try Kv.Kv8.load(gpu.d); + defer k8.deinit(); + var b: Bench = .{ .s = r.s, .ev0 = try cuda.Event.init(gpu.d, true), .ev1 = try cuda.Event.init(gpu.d, true) }; + { + var free_b: usize = 0; + var total_b: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free_b, &total_b), "cuMemGetInfo"); + std.debug.print(" device after context + the attention modules: free {d:.3} of {d:.3} GiB (context and modules {d:.3} GiB)\n", .{ @as(f64, @floatFromInt(free_b)) / Kv.gib, @as(f64, @floatFromInt(total_b)) / Kv.gib, @as(f64, @floatFromInt(total_b - free_b)) / Kv.gib }); + } + const nrow = cap * kv_heads; + const kv_bytes = nrow * head_dim * 2; + const cb = Kv.Kv8.cacheBytes(cap, kv_heads); + std.debug.print("== exl3-kv bench: cache of {d} tokens, one attention layer: bf16 K+V {d:.1} MiB, int8 codes+scales {d:.1} MiB\n", .{ cap, @as(f64, @floatFromInt(2 * kv_bytes)) / 1048576.0, @as(f64, @floatFromInt(2 * (cb.codes + cb.scales))) / 1048576.0 }); + + // the cache's values: random bf16 made on the device side once, then packed; the host copies get the same bytes + var dev_k = try cuda.DeviceBuffer.alloc(gpu.d, kv_bytes); + defer dev_k.free(); + var dev_v = try cuda.DeviceBuffer.alloc(gpu.d, kv_bytes); + defer dev_v.free(); + var host_k = try cuda.HostBuffer.allocMapped(gpu.d, kv_bytes); + defer host_k.free(); + var host_v = try cuda.HostBuffer.allocMapped(gpu.d, kv_bytes); + defer host_v.free(); + { + const hk = host_k.slice(u16); + const hv = host_v.slice(u16); + for (hk) |*x| x.* = bf16(rng.floatNorm(f32) * 0.5); + for (hv) |*x| x.* = bf16(rng.floatNorm(f32) * 0.5); + try dev_k.upload(0, host_k.bytes); + try dev_v.upload(0, host_v.bytes); + } + var dev_c: [4]cuda.DeviceBuffer = undefined; + var host_c: [4]cuda.HostBuffer = undefined; + for (0..4) |i| { + const n = if (i % 2 == 0) cb.codes else cb.scales; + dev_c[i] = try cuda.DeviceBuffer.alloc(gpu.d, n); + host_c[i] = try cuda.HostBuffer.allocMapped(gpu.d, n); + } + defer for (&dev_c, &host_c) |*x, *y| { + x.free(); + y.free(); + }; + try k8.packRows(r.s, dev_k.ptr, dev_c[0].ptr, dev_c[1].ptr, nrow, true); + try k8.packRows(r.s, dev_v.ptr, dev_c[2].ptr, dev_c[3].ptr, nrow, true); + try r.s.synchronize(); + for (0..4) |i| try dev_c[i].download(0, host_c[i].bytes); + const on_dev: Place = .{ .kc = dev_k.ptr, .vc = dev_v.ptr, .c = .{ .k8 = dev_c[0].ptr, .ks = dev_c[1].ptr, .v8 = dev_c[2].ptr, .vs = dev_c[3].ptr } }; + const on_host: Place = .{ .kc = try host_k.device(), .vc = try host_v.device(), .c = .{ .k8 = try host_c[0].device(), .ks = try host_c[1].device(), .v8 = try host_c[2].device(), .vs = try host_c[3].device() } }; + + const max_rows = 256; + const q = r.bfs(max_rows * q_heads * head_dim, 1); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const po = try r.out(max_rows * nch * q_heads * head_dim * 4); + const pm = try r.out(max_rows * nch * q_heads * 4); + const pl = try r.out(max_rows * nch * q_heads * 4); + const out = try r.out(max_rows * q_heads * head_dim * 2); + const ids = try a.alloc(i32, max_rows * idw); + const nk = try a.alloc(i32, max_rows); + const spr = try a.alloc(i32, max_rows); + + const Case = struct { name: []const u8, rows: usize, end0: usize, sparse: bool, keep: f64 }; + const cases = [_]Case{ + .{ .name = "decode 1 row, sparse at the cache end", .rows = 1, .end0 = cap - 2, .sparse = true, .keep = 0 }, + .{ .name = "decode window 4 rows (MTP verify), keep 0.95", .rows = 4, .end0 = cap - 5, .sparse = true, .keep = 0.95 }, + .{ .name = "prompt 256 rows, sparse, keep 0.95", .rows = 256, .end0 = cap - 257, .sparse = true, .keep = 0.95 }, + .{ .name = "prompt 256 rows, sparse, keep 0.80", .rows = 256, .end0 = cap - 257, .sparse = true, .keep = 0.80 }, + .{ .name = "prompt 256 rows, sparse, keep 0 (worst case)", .rows = 256, .end0 = cap - 257, .sparse = true, .keep = 0 }, + .{ .name = "prompt 256 rows, dense (keys 1,795..2,050)", .rows = 256, .end0 = 1794, .sparse = false, .keep = 0 }, + }; + for (cases) |c| { + @memset(ids, 0); + if (c.sparse) { + try overlapIds(rng, a, c.end0, c.rows, c.keep, ids, nk); + } else for (0..c.rows) |row| { + nk[row] = @intCast(c.end0 + row + 1); + } + for (spr[0..c.rows]) |*x| x.* = @intFromBool(c.sparse); + // distinct keys these rows read (what crosses PCIe at least once) + var seen = try std.DynamicBitSetUnmanaged.initEmpty(a, cap); + var reads: usize = 0; + for (0..c.rows) |row| { + const n: usize = @intCast(nk[row]); + reads += n; + for (0..n) |i| seen.set(if (c.sparse) @intCast(ids[row * idw + i]) else i); + } + const distinct = seen.count(); + const pos0: i32 = @intCast(c.end0 + 1); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const dids = try r.dev(std.mem.sliceAsBytes(ids[0 .. c.rows * idw])); + const dnk = try r.dev(std.mem.sliceAsBytes(nk[0..c.rows])); + const dspr = try r.dev(std.mem.sliceAsBytes(spr[0..c.rows])); + const chunks = if (c.sparse) nch else (c.end0 + c.rows + 1 + 511) / 512; + var ms: [4]f64 = undefined; + for ([_]bool{ false, true }, 0..) |int8, ii| for ([_]Place{ on_dev, on_host }, 0..) |p, pi| { + const at: Attn = .{ .g = &g, .k8 = &k8, .s = r.s, .p = p, .int8 = int8, .q = dq, .pos = dpos, .po = po, .pm = pm, .pl = pl, .ids = dids, .nk = dnk, .spr = dspr, .out = out, .rows = c.rows, .chunks = chunks }; + b.reps = if (pi == 1 and c.rows > 4) 3 else 20; + ms[ii * 2 + pi] = try b.time(&at, Attn.go); + }; + // the int8 kernel with the next tile in flight at every row count (Kv.ahead_rows raised) + var ms_ahead: f64 = ms[2]; + if (c.rows > Kv.ahead_rows) { + const keep_rows = Kv.ahead_rows; + Kv.ahead_rows = 1 << 20; + defer Kv.ahead_rows = keep_rows; + const at: Attn = .{ .g = &g, .k8 = &k8, .s = r.s, .p = on_dev, .int8 = true, .q = dq, .pos = dpos, .po = po, .pm = pm, .pl = pl, .ids = dids, .nk = dnk, .spr = dspr, .out = out, .rows = c.rows, .chunks = chunks }; + b.reps = 20; + ms_ahead = try b.time(&at, Attn.go); + } + const key_bf: f64 = 2 * kv_heads * head_dim * 2; // one key's K and V, bf16, both heads + const key_i8: f64 = @floatFromInt(Kv.kvRowBytes(.int8, kv_heads)); + const once_bf = @as(f64, @floatFromInt(distinct)) * key_bf / 1e6; + const once_i8 = @as(f64, @floatFromInt(distinct)) * key_i8 / 1e6; + std.debug.print("RESULT {s}: keys read {d}, distinct {d} ({d:.2} MB bf16 / {d:.2} MB int8 once)\n", .{ c.name, reads, distinct, once_bf, once_i8 }); + std.debug.print(" chunks+merge ms: bf16 device {d:.4} bf16 host-mapped {d:.4} ({d:.1} GB/s of distinct) int8 device {d:.4} (read-ahead kernel {d:.4}) int8 host-mapped {d:.4} ({d:.1} GB/s of distinct)\n", .{ ms[0], ms[1], once_bf / ms[1], ms[2], ms_ahead, ms[3], once_i8 / ms[3] }); + } + + // explicit staging alternatives for one decode row's 2,051 keys + { + const n = 2051; + const row_bytes = kv_heads * head_dim * 2; // one key, both heads, K (V the same) + const bytes = 2 * n * row_bytes; + var pinned = try cuda.HostBuffer.alloc(gpu.d, bytes); + defer pinned.free(); + var staging = try cuda.DeviceBuffer.alloc(gpu.d, bytes); + defer staging.free(); + const one: Copy = .{ .s = r.s, .dst = staging.ptr, .src = pinned.bytes.ptr, .bytes = bytes, .pieces = 1 }; + b.reps = 20; + const t_one = try b.time(&one, Copy.go); + const blocks: Copy = .{ .s = r.s, .dst = staging.ptr, .src = pinned.bytes.ptr, .bytes = 1024 * (bytes / 1024), .pieces = 1024 }; + const t_blocks = try b.time(&blocks, Copy.go); + try overlapIds(rng, a, cap - 2, 1, 0, ids, nk); + const dids = try r.dev(std.mem.sliceAsBytes(ids[0..n])); + const gk: Gather = .{ .k8 = &k8, .s = r.s, .src_k = on_host.kc, .src_v = on_host.vc, .dst_k = staging.ptr, .dst_v = staging.ptr + n * row_bytes, .ids = dids, .n = n, .row = row_bytes }; + const t_gather = try b.time(&gk, Gather.go); + // the selection's ids to the host and back to a waiting thread: one D2H and a stream synchronize + var hid = try cuda.HostBuffer.alloc(gpu.d, n * 4); + defer hid.free(); + var best: f64 = std.math.inf(f64); + for (0..50) |_| { + const t0 = check.now(gpu.io); + try gpu.d.check(gpu.d.api.cuMemcpyDtoHAsync_v2(hid.bytes.ptr, dids, n * 4, r.s.handle), "cuMemcpyDtoHAsync"); + try r.s.synchronize(); + best = @min(best, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6); + } + const mb = @as(f64, @floatFromInt(bytes)) / 1e6; + std.debug.print("RESULT staging one decode row's 2,051 keys (K+V bf16, {d:.2} MB): one pinned H2D {d:.4} ms ({d:.1} GB/s); 1,024 block copies {d:.4} ms; gather kernel from mapped host {d:.4} ms ({d:.1} GB/s); ids D2H + stream sync {d:.4} ms\n", .{ mb, t_one, mb / t_one, t_blocks, t_gather, mb / t_gather, best }); + } + // the indexer's raw keys in mapped host memory: upstream's pooling reads them only when a block of 4 completes + { + const ikc_bytes = cap * 128 * 2; + var ikc_dev = try cuda.DeviceBuffer.alloc(gpu.d, ikc_bytes); + defer ikc_dev.free(); + var ikc_host = try cuda.HostBuffer.allocMapped(gpu.d, ikc_bytes); + defer ikc_host.free(); + for (ikc_host.slice(u16)) |*x| x.* = bf16(rng.floatNorm(f32)); + try ikc_dev.upload(0, ikc_host.bytes); + const kw = try r.dev(std.mem.sliceAsBytes(r.f32s(128, 0.5, 1.5))); + const inv = try a.alloc(f32, half); + for (inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + const dinv = try r.dev(std.mem.sliceAsBytes(inv)); + const Pool = struct { + g: *const flash.Glue, + s: cuda.Stream, + ikc: u64, + pooled: u64, + pos: u64, + kw: u64, + inv: u64, + rows: usize, + fn go(p: *const @This()) anyerror!void { + try p.g.qsaPool(p.s, p.ikc, p.pooled, p.pos, p.kw, p.inv, 1e-6, p.rows); + } + }; + for ([_]usize{ 1, 256 }) |rows| { + const pos0: i32 = @intCast(cap - 256); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ if (rows == 1) pos0 + 3 else pos0, 0, 0, 0 })); + const pd = try r.out((cap / 4 + 1) * 128 * 2); + const ph = try r.out((cap / 4 + 1) * 128 * 2); + const on_d: Pool = .{ .g = &g, .s = r.s, .ikc = ikc_dev.ptr, .pooled = pd, .pos = dpos, .kw = kw, .inv = dinv, .rows = rows }; + const on_h: Pool = .{ .g = &g, .s = r.s, .ikc = try ikc_host.device(), .pooled = ph, .pos = dpos, .kw = kw, .inv = dinv, .rows = rows }; + b.reps = 20; + const td = try b.time(&on_d, Pool.go); + const th = try b.time(&on_h, Pool.go); + const n = (cap / 4 + 1) * 128; + const same = std.mem.eql(u16, try r.back(u16, pd, n), try r.back(u16, ph, n)); + std.debug.print("RESULT indexer pooling of {d} row(s) at the cache end: raw keys on the device {d:.4} ms, in mapped host memory {d:.4} ms; pooled rows {s}\n", .{ rows, td, th, if (same) "bit-equal" else "DIFFER" }); + try check.expect(same, "pooling from mapped host keys differs", .{}); + } + } + // a prompt-cache snapshot (cuda_snap.zig copies with cuMemcpyDtoDAsync) into mapped pinned host memory and back + { + const bytes: usize = 256 << 20; + var dev = try cuda.DeviceBuffer.alloc(gpu.d, bytes); + defer dev.free(); + var host = try cuda.HostBuffer.allocMapped(gpu.d, bytes); + defer host.free(); + const hp = try host.device(); + const Dtod = struct { + s: cuda.Stream, + dst: u64, + src: u64, + n: usize, + fn go(c: *const @This()) anyerror!void { + try c.s.d.check(c.s.d.api.cuMemcpyDtoDAsync_v2(c.dst, c.src, c.n, c.s.handle), "cuMemcpyDtoDAsync"); + } + }; + b.reps = 3; + const save: Dtod = .{ .s = r.s, .dst = hp, .src = dev.ptr, .n = bytes }; + const t_save = try b.time(&save, Dtod.go); + const back: Dtod = .{ .s = r.s, .dst = dev.ptr, .src = hp, .n = bytes }; + const t_back = try b.time(&back, Dtod.go); + const gb = @as(f64, @floatFromInt(bytes)) / 1e9; + std.debug.print("RESULT a 256 MiB snapshot through cuMemcpyDtoDAsync to/from mapped pinned host memory: save {d:.2} ms ({d:.1} GB/s), restore {d:.2} ms ({d:.1} GB/s)\n", .{ t_save, gb / t_save * 1e3, t_back, gb / t_back * 1e3 }); + } + check.pass("exl3-kv bench: timings above (an idle card under gpurun)", .{}); +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const mode = rest[0]; + if (std.mem.eql(u8, mode, "kernel")) return kernelMode(gpu); + if (std.mem.eql(u8, mode, "full")) return @import("exl3_kv_full.zig").run(gpu, rest[1..]); + if (std.mem.eql(u8, mode, "real")) return @import("exl3_kv_real.zig").run(gpu, rest[1..]); + if (std.mem.eql(u8, mode, "bench")) return benchMode(gpu, if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 200_016); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/exl3_kvpage_tests.zig b/zig/tests/cuda/exl3_kvpage_tests.zig new file mode 100644 index 0000000000..5e1e64c1b8 --- /dev/null +++ b/zig/tests/cuda/exl3_kvpage_tests.zig @@ -0,0 +1,803 @@ +//! exl3-kvpage: dynamic KV paging for long contexts (docs/plans/exl3-kv-paging.md) on the family's Engine with the +//! EXL3 offload pool, as serve attaches it (one card). +//! +//! exl3-kvpage DIR trace IDS N[,N..] DECODE OUT [prompt] +//! the first N ids of IDS (u32 little-endian, real text) prefilled, then DECODE greedy rounds; every attention +//! row's selected blocks recorded (cuda_exl3_kvpage.Trace) into OUT-N.bin (prompt rows only with `prompt`) +//! exl3-kvpage none kernel [SLOTS] +//! the paging kernels on synthetic key lists: every remapped key reads its own host row, misses against an +//! LRU model on the host, the plan's cost a call +//! exl3-kvpage DIR cmp IDS N[,N..] DECODE WINDOW SLOTS [NEEDLE_AT] +//! each prompt through the all-on-device cache, then through the paged one (window WINDOW tokens a layer, SLOTS +//! 4-key block slots shared by the 12 layers): the prefill's last logits and the DECODE greedy tokens must be +//! bit-equal, and every K/V row of every attention layer equal after the decode (paged rows flushed to host); +//! ms, copies and device bytes. +//! NEEDLE_AT (a fraction, e.g. 0.25): a sentence with a 5-digit code inserted that far into the text and a +//! question about it appended; the greedy answer must contain the code (both caches) +//! +//! BF16 paged snapshot acceptance plan (the format/order contract is in cuda_snap.zig): +//! 1. Host-driver fakes: allocate/save/restore failures at every section, queued copies drained before freeing, +//! truncated/unknown formats and checked-size overflow, exact byte admission, incompatible capacity/geometry, +//! no mutation on preflight rejection, and no successful publication after a partial copy failure. +//! 2. Kernel lifecycle: save at 0, 1, 3, 4, dense/sparse boundaries, and immediately before/after several rebases. +//! Seed distinct BF16 bit patterns (including NaNs/signed zero) per layer/K/V/row, dirty window tails and occupied +//! LRU slots; compare canonical rows, device windows/slots, both mapping directions, marks, epochs and stamps. +//! Test same-sequence rewind, restore into a distinct compatible sequence, source drop, repeated restore and +//! snapshot drop while two other sequences remain live. Host snapshots must remain immutable through each case. +//! 3. Model continuation: save, run a suffix, restore, then require bit-identical eager logits and flushed BF16 K/V +//! for that suffix; repeat across a window rebase and with partial pooled blocks and PLE n-gram history. +//! 4. MTP: require a settled head mark; reject outstanding drafts. Populate verification cache before restore and +//! inject a delayed pre-restore publication; neither may survive the generation change. Check restored head +//! rows, len/drafted/rows, recurrent rollback state and drafted/serial continuation equality across rebases. +//! 5. TP2/graphs: two distinct live histories, staggered rebases, capture/replay before AND after restore, and drop/ +//! recreation of one lane while the other remains live. Compare every rank's canonical K/V and outputs to eager; +//! require both attention and GDN graph reuse where compatible, and recapture when origin arguments differ. +//! 6. Inject one-rank overflow, copy failure, delayed instantiation and peer failure. Require bounded failure, +//! joined producers, no failed-batch results or stale verification results, no live-buffer reuse and clean drop. +//! 7. After those gates, two concurrent 200k TP2 requests with BF16 RAM paging, snapshot branch/reuse and MTP, graph +//! prefill/decode parity, host/device peak accounting and cold/warm timings. No gate above has been run here. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const flash = @import("flashnext"); +const Gpu = check.Gpu; + +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const xpre = flash.exl3.prefill; +const kvp = flash.exl3.kvpage; +const vocab = prompt.vocabulary; +const core = @import("core"); + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + attached: xpass.Attached, + + fn init(gpu: Gpu, dir: []const u8, cap: usize, deep_rows: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + g.seq = try g.e.newSeq(cap); + if (deep_rows > 0) try xpre.reserve(gpu.gpa, gpu.d, g.e, deep_rows); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, try xpass.budgetFromEnv(0), null); + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + std.debug.print("engine: loaded in {d:.1} s; cap {d}; pool {d} lines a layer; {d:.2} GiB free\n", .{ @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9, cap, g.attached.lines, @as(f64, @floatFromInt(free)) / (1 << 30) }); + return g; + } + + fn deinit(g: *Eng) void { + xpre.release(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn prefill(g: *Eng, ids: []const u32) !f64 { + try g.reset(); + const t0 = check.now(g.gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + return @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + } + + /// `n` greedy rounds from the prefill's logits; the tokens into `out`; ms a round. + fn decode(g: *Eng, out: []u32) !f64 { + var dev = try cuda.DeviceBuffer.alloc(g.gpu.d, vocab * 2); + defer dev.free(); + var tok = prompt.argmax(g.seq.logits); + const t0 = check.now(g.gpu.io); + for (out) |*o| { + o.* = tok; + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], dev.ptr); + tok = prompt.argmax(g.seq.logits); + } + return @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6 / @as(f64, @floatFromInt(@max(out.len, 1))); + } +}; + +fn readIds(gpu: Gpu, path: []const u8) ![]u32 { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 28)); + defer a.free(text); + const ids = try a.alloc(u32, text.len / 4); + for (ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + return ids; +} + +fn parseList(s: []const u8, out: []usize) !usize { + var n: usize = 0; + var it = std.mem.splitScalar(u8, s, ','); + while (it.next()) |t| { + out[n] = try std.fmt.parseInt(usize, t, 10); + n += 1; + } + return n; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + if (std.mem.eql(u8, mode, "trace")) return traceRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "cmp")) return cmpRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "kernel")) return kernelRun(gpu, if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 6144); + if (std.mem.eql(u8, mode, "kernel-graph")) return kernelReplay(gpu, 49152, true, false); + if (std.mem.eql(u8, mode, "window-graph")) return kernelReplay(gpu, 49152, true, true); + if (std.mem.eql(u8, mode, "overflow-graph")) return overflowReplay(gpu); + if (std.mem.eql(u8, mode, "snapshot")) return snapshotLifecycle(gpu); + if (std.mem.eql(u8, mode, "recycle")) return @import("exl3_recycle_tests.zig").run(gpu); + return error.UnknownMode; +} + +fn snapshotLifecycle(gpu: Gpu) !void { + const snaps = @import("flashnext").snapshot; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var engine: prompt.Engine = undefined; + engine.gpa = gpu.gpa; + engine.driver = @constCast(gpu.d); + engine.stream = &stream; + engine.chunk = 256; + engine.eos = 0; + engine.round_graph = .{}; + engine.page_config = .{ .window = 2052, .slots = 16, .rows = 256 }; + engine.shared_arena = null; + defer flash.exl3.recycle.drain(&engine); + var destination = try engine.newSeq(8192); + defer destination.deinit(gpu.gpa); + for ([_]usize{ 0, 1, 3, 4, 2051, 2052, 4096, 4099, 6147 }) |at| { + var source: ?prompt.Seq = try engine.newSeq(8192); + defer if (source) |*sequence| sequence.deinit(gpu.gpa); + const pg = source.?.paged.?; + source.?.fed = at; + source.?.history = .{ @intCast(at), 73 }; + for (0..12) |layer| { + const origin = (at -| 2052) / 4 * 4; + pg.win0[layer] = origin; + pg.rebases[layer] = layer + 3; + pg.written[layer] = origin; + for (0..2) |side| { + const host = pg.hostMut(layer, side == 1); + const words: []u16 = std.mem.bytesAsSlice(u16, @as([]align(2) u8, @alignCast(host))); + for (words, 0..) |*word, index| word.* = @truncate(index + layer * 997 + side * 32768); + } + } + try pg.dev.fill8(0xFF, stream.handle); + const words = try gpu.gpa.alloc(u32, pg.tabs.len / 4); + defer gpu.gpa.free(words); + try pg.tabs.download(0, std.mem.sliceAsBytes(words)); + const ctl_at = 24 * pg.nblocks; + const inverse_at = ctl_at + 12 * 8; + for (0..12) |layer| { + words[ctl_at + layer * 8] = 73; + words[ctl_at + layer * 8 + 6] = @intCast(31 + layer); + if (at >= 4) { + words[layer * pg.nblocks] = @intCast(layer); + words[12 * pg.nblocks + layer * pg.nblocks] = @intCast(31 + layer); + words[inverse_at + layer] = @intCast(layer * pg.nblocks); + words[inverse_at + pg.cfg.slots + layer] = @intCast(19 + layer); + } + } + try pg.tabs.upload(0, std.mem.sliceAsBytes(words)); + const saved = try snaps.save(gpu.gpa, gpu.d, stream, &source.?, null); + defer snaps.drop(gpu.gpa, stream, saved); + source.?.deinit(gpu.gpa); + source = null; + for (0..2) |_| { + try snaps.restore(stream, &destination, null, saved); + const again = try snaps.save(gpu.gpa, gpu.d, stream, &destination, null); + defer snaps.drop(gpu.gpa, stream, again); + try check.expect(std.meta.eql(saved.paged.?, again.paged.?), "snapshot pager counters at {d}", .{at}); + try check.expect(std.mem.eql(u8, saved.host.?.bytes[0..saved.buf.len], again.host.?.bytes[0..again.buf.len]), "snapshot exact BF16, windows, slots, tables and ordinary state at {d}", .{at}); + } + } + check.pass("BF16 paged snapshot source-drop, distinct destination, repeated restore at nine boundaries", .{}); +} + +fn traceRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 4) return error.MissingArgument; + const all = try readIds(gpu, rest[0]); + defer a.free(all); + var ns_buf: [16]usize = undefined; + const ns = ns_buf[0..try parseList(rest[1], &ns_buf)]; + const steps = try std.fmt.parseInt(usize, rest[2], 10); + const out_prefix = rest[3]; + const with_prompt = rest.len > 4 and std.mem.eql(u8, rest[4], "prompt"); + var max_n: usize = 0; + for (ns) |n| max_n = @max(max_n, n); + if (max_n > all.len) return error.NotEnoughIds; + const o = xpre.optionsFromEnv(); + const g = try Eng.init(gpu, dir, max_n + steps + 64, if (o.enabled) @min(o.rows, (max_n + 255) / 256 * 256) else 0); + defer g.deinit(); + const toks = try a.alloc(u32, steps); + defer a.free(toks); + for (ns) |n| { + var t = try kvp.Trace.init(a, with_prompt); + defer t.deinit(); + kvp.trace = &t; + defer kvp.trace = null; + const ms_p = try g.prefill(all[0..n]); + const prompt_rows = t.rows; + const ms_d = try g.decode(toks); + std.debug.print("TRACE {d} tokens: prefill {d:.0} ms (traced, {d} rows recorded), {d} decode rounds {d:.1} ms/round (traced: one sync a layer); {d} rows recorded; {d:.1} MiB\n", .{ n, ms_p, prompt_rows, steps, ms_d, t.rows, @as(f64, @floatFromInt(t.out.items.len)) / (1 << 20) }); + std.debug.print(" greedy: {any}\n", .{toks[0..@min(steps, 48)]}); + var buf: [512]u8 = undefined; + const path = try std.fmt.bufPrint(&buf, "{s}-{d}.bin", .{ out_prefix, n }); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = t.out.items }); + } + check.pass("exl3-kvpage trace written", .{}); +} + +fn freeMem(gpu: Gpu) !usize { + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +fn gib(x: usize) f64 { + return @as(f64, @floatFromInt(x)) / (1 << 30); +} + +const Run = struct { + prefill_ms: f64 = 0, + decode_ms: f64 = 0, + logits: []u8 = &.{}, + toks: []u32 = &.{}, + /// [layer][K|V] rows 0..fed (all-on-device: downloaded; paged: its host copy after a flush) + kv: [][]u8 = &.{}, + stats_prefill: kvp.Stats = .{}, + stats_decode: kvp.Stats = .{}, +}; + +fn oneRun(g: *Eng, ids: []const u32, steps: usize, keep_kv: bool) !Run { + const a = g.gpu.gpa; + var r: Run = .{}; + r.prefill_ms = try g.prefill(ids); + r.logits = try a.dupe(u8, g.seq.logits); + if (g.seq.paged) |pg| r.stats_prefill = try pg.total(); + r.toks = try a.alloc(u32, steps); + r.decode_ms = try g.decode(r.toks); + if (g.seq.paged) |pg| { + const t = try pg.total(); + r.stats_decode = .{ .misses = t.misses - r.stats_prefill.misses, .needed = t.needed - r.stats_prefill.needed, .calls = t.calls - r.stats_prefill.calls, .overflow = t.overflow, .rebases = t.rebases - r.stats_prefill.rebases, .written_rows = t.written_rows - r.stats_prefill.written_rows }; + } + if (!keep_kv) return r; + const fed = g.seq.fed; + const row = kvp.row_bytes; + r.kv = try a.alloc([]u8, 24); + for (0..12) |ai| for (0..2) |v| { + const dst = try a.alloc(u8, fed * row); + if (g.seq.paged) |pg| { + if (v == 0 and ai == 0) { + try pg.flush(g.stream, fed); + try g.stream.synchronize(); + } + @memcpy(dst, pg.hostRows(ai, v == 1)[0 .. fed * row]); + } else { + const src = if (v == 0) g.seq.kc[ai].ptr else g.seq.vc[ai].ptr; + try g.gpu.d.check(g.gpu.d.api.cuMemcpyDtoH_v2(dst.ptr, src, fed * row), "cuMemcpyDtoH"); + } + r.kv[2 * ai + v] = dst; + }; + return r; +} + +fn freeRun(a: std.mem.Allocator, r: *Run) void { + a.free(r.logits); + a.free(r.toks); + for (r.kv) |x| a.free(x); + if (r.kv.len > 0) a.free(r.kv); +} + +/// The first n ids, or (needle) the text with a code sentence at `at` and a question about it at the end. +fn promptOf(a: std.mem.Allocator, tok: *const core.tokenizer.Tokenizer, all: []const u32, n: usize, needle_at: ?f64, code: []const u8) ![]u32 { + if (needle_at == null) return a.dupe(u32, all[0..n]); + var buf: [256]u8 = undefined; + const fact = try tok.encodePlain(a, try std.fmt.bufPrint(&buf, "\n\nIMPORTANT NOTE: the secret access code for the blue vault is {s}. Remember this code.\n\n", .{code})); + defer a.free(fact); + const q = try tok.encodePlain(a, "\n\nQuestion: what is the secret access code for the blue vault, given in the important note above?\nAnswer: the secret access code for the blue vault is"); + defer a.free(q); + const body = n - fact.len - q.len; + const at: usize = @intFromFloat(@as(f64, @floatFromInt(body)) * needle_at.?); + var out: std.ArrayList(u32) = .empty; + try out.appendSlice(a, all[0..at]); + try out.appendSlice(a, fact); + try out.appendSlice(a, all[at..body]); + try out.appendSlice(a, q); + return out.toOwnedSlice(a); +} + +fn cmpRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 5) return error.MissingArgument; + const all = try readIds(gpu, rest[0]); + defer a.free(all); + var ns_buf: [16]usize = undefined; + const ns = ns_buf[0..try parseList(rest[1], &ns_buf)]; + const steps = try std.fmt.parseInt(usize, rest[2], 10); + const window = try std.fmt.parseInt(usize, rest[3], 10); + const slots = try std.fmt.parseInt(usize, rest[4], 10); + const needle_at: ?f64 = if (rest.len > 5) try std.fmt.parseFloat(f64, rest[5]) else null; + var max_n: usize = 0; + for (ns) |n| max_n = @max(max_n, n); + if (max_n > all.len) return error.NotEnoughIds; + var path_buf: [512]u8 = undefined; + var tok = try core.tokenizer.loadTokenizer(gpu.io, a, try std.fmt.bufPrint(&path_buf, "{s}/tokenizer.json", .{dir})); + defer tok.deinit(); + const o = xpre.optionsFromEnv(); + const cap = max_n + steps + 64; + kvp.force_off = true; + const g = try Eng.init(gpu, dir, cap, if (o.enabled) @min(o.rows, (max_n + 255) / 256 * 256) else 0); + defer g.deinit(); + kvp.force_off = false; + const cfg: kvp.Config = .{ .window = window, .slots = slots, .rows = g.e.chunk }; + var all_ok = true; + const codes = [_][]const u8{ "48613", "70259", "31846", "95072", "62418", "27530" }; + for (ns, 0..) |n, ni| { + const code = codes[ni % codes.len]; + const ids = try promptOf(a, &tok, all, n, needle_at, code); + defer a.free(ids); + // all on the device (the engine's own sequence) + var d = try oneRun(g, ids, steps, true); + defer freeRun(a, &d); + const free_dense = try freeMem(gpu); + g.seq.deinit(a); + const free_none = try freeMem(gpu); + kvp.force = cfg; + g.seq = try g.e.newSeq(cap); + kvp.force = null; + const free_paged = try freeMem(gpu); + const pg = g.seq.paged orelse return error.KvPageOff; + const pg_dev = pg.deviceBytes(); + const pg_host = pg.host.bytes.len; + var p = try oneRun(g, ids, steps, true); + defer freeRun(a, &p); + // back to an all-on-device sequence for the next prompt + g.seq.deinit(a); + kvp.force_off = true; + g.seq = try g.e.newSeq(cap); + kvp.force_off = false; + const logits_eq = std.mem.eql(u8, d.logits, p.logits); + const toks_eq = std.mem.eql(u32, d.toks, p.toks); + var kv_eq: usize = 0; + for (d.kv, p.kv) |x, y| kv_eq += @intFromBool(std.mem.eql(u8, x, y)); + const dense_seq = free_none - free_dense; + const paged_seq = free_none - free_paged; + std.debug.print("CMP {d} tokens{s}: prefill logits bit-equal {}; {d} greedy tokens equal {}; K/V rows 0..{d} equal in {d}/24 layer-sides\n", .{ ids.len, if (needle_at != null) " (needle)" else "", logits_eq, steps, toks_eq, ids.len + steps, kv_eq }); + std.debug.print(" device bytes of the sequence: all-on-device {d:.3} GiB, paged {d:.3} GiB (window {d} tokens a layer, {d} block slots shared by the layers: paged side {d:.3} GiB), freed {d:.3} GiB; host {d:.2} GiB pinned\n", .{ gib(dense_seq), gib(paged_seq), window, slots, gib(pg_dev), gib(dense_seq -| paged_seq), gib(pg_host) }); + std.debug.print(" prefill: all-on-device {d:.0} ms, paged {d:.0} ms; decode: all-on-device {d:.2} ms/token, paged {d:.2} ms/token\n", .{ d.prefill_ms, p.prefill_ms, d.decode_ms, p.decode_ms }); + const sp = p.stats_prefill; + const sd = p.stats_decode; + std.debug.print(" paged prefill: {d} calls, {d} blocks needed outside the window, {d} misses ({d:.1}% hit, {d:.1} MiB fetched, {d:.0} B a row a layer), {d} window moves, {d} rows written out; overflow {}\n", .{ sp.calls, sp.needed, sp.misses, 100 * (1 - @as(f64, @floatFromInt(sp.misses)) / @as(f64, @floatFromInt(@max(sp.needed, 1)))), @as(f64, @floatFromInt(sp.misses * 8192)) / (1 << 20), @as(f64, @floatFromInt(sp.misses * 8192)) / @as(f64, @floatFromInt(ids.len * 12)), sp.rebases, sp.written_rows, sp.overflow }); + std.debug.print(" paged decode: {d} calls, {d} needed outside the window, {d} misses ({d:.1}% hit), {d:.1} KiB fetched a token ({d:.2} KiB a layer); overflow {}\n", .{ sd.calls, sd.needed, sd.misses, 100 * (1 - @as(f64, @floatFromInt(sd.misses)) / @as(f64, @floatFromInt(@max(sd.needed, 1)))), @as(f64, @floatFromInt(sd.misses * 8192)) / 1024 / @as(f64, @floatFromInt(@max(steps, 1))), @as(f64, @floatFromInt(sd.misses * 8192)) / 1024 / @as(f64, @floatFromInt(@max(steps, 1) * 12)), sd.overflow }); + const td = try tok.decode(a, d.toks, false); + defer a.free(td); + const tp = try tok.decode(a, p.toks, false); + defer a.free(tp); + std.debug.print(" greedy (all-on-device): {any}\n text: \"{f}\"\n", .{ d.toks[0..@min(steps, 32)], std.zig.fmtString(td[0..@min(td.len, 240)]) }); + if (!toks_eq) std.debug.print(" greedy (paged): {any}\n text: \"{f}\"\n", .{ p.toks[0..@min(steps, 32)], std.zig.fmtString(tp[0..@min(tp.len, 240)]) }); + if (needle_at != null) { + const hit_d = std.mem.indexOf(u8, td, code) != null; + const hit_p = std.mem.indexOf(u8, tp, code) != null; + std.debug.print(" NEEDLE code {s} at {d:.2} of the text: all-on-device {s}, paged {s}\n", .{ code, needle_at.?, if (hit_d) "FOUND" else "missed", if (hit_p) "FOUND" else "missed" }); + if (!hit_p) all_ok = false; + } + if (!logits_eq or !toks_eq or kv_eq != 24 or sp.overflow or sd.overflow) all_ok = false; + } + try check.expect(all_ok, "paged K/V bit-equal to all-on-device (logits, greedy tokens, rows){s}", .{if (needle_at != null) " and the needle found" else ""}); + check.pass("exl3-kvpage cmp", .{}); +} + +/// The host's model of kvp_plan: exact LRU at call granularity, ties to the lower slot. +const Model = struct { + stamp: []u32, + slot_of: std.AutoHashMap(u32, u32), + blk: []i64, + epoch: u32 = 0, + + fn use(mo: *Model, a: std.mem.Allocator, need: []const u32) !usize { + mo.epoch += 1; + var misses: std.ArrayList(u32) = .empty; + defer misses.deinit(a); + for (need) |b| { + if (mo.slot_of.get(b)) |sl| mo.stamp[sl] = mo.epoch else try misses.append(a, b); + } + const order = try a.alloc(u32, mo.stamp.len); + defer a.free(order); + for (order, 0..) |*o, i| o.* = @intCast(i); + const Ctx = struct { + st: []const u32, + ep: u32, + fn key(c: @This(), i: u32) u64 { + const t = c.st[i]; + return if (t == c.ep) std.math.maxInt(u64) else (@as(u64, t) << 32) | i; + } + fn less(c: @This(), x: u32, y: u32) bool { + return c.key(x) < c.key(y); + } + }; + std.mem.sort(u32, order, Ctx{ .st = mo.stamp, .ep = mo.epoch }, Ctx.less); + for (misses.items, 0..) |b, i| { + const sl = order[i]; + if (mo.stamp[sl] == mo.epoch) return error.ModelOverflow; + if (mo.blk[sl] >= 0) _ = mo.slot_of.remove(@intCast(mo.blk[sl])); + mo.blk[sl] = b; + try mo.slot_of.put(b, sl); + mo.stamp[sl] = mo.epoch; + } + return misses.items.len; + } +}; + +fn kernelRun(gpu: Gpu, slots: usize) !void { + return kernelReplay(gpu, slots, false, false); +} + +fn kernelReplay(gpu: Gpu, slots: usize, graphed: bool, moving: bool) !void { + const a = gpu.gpa; + const idw = flash.exl3.prompt_wt.index_idw; + const row = kvp.row_bytes; + const cap: usize = 60_000; + const w0: usize = 40_000; + const nb_before = w0 / 4; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + const pg = try kvp.Paged.init(a, gpu.d, cap, .{ .window = 2052, .slots = slots, .rows = 256 }); + defer pg.destroy(); + pg.setWindow(0, w0); + // each host row's 256 words unique to (position, K or V) + for (0..2) |v| { + const h = std.mem.bytesAsSlice(u32, pg.hostMut(0, v == 1)); + for (0..cap) |p| for (0..row / 4) |i| { + h[p * (row / 4) + i] = @intCast((p *% 2654435761 +% i *% 40503 +% v *% 977) & 0xFFFF_FFFF); + }; + if (moving) try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(pg.devRows(0, v == 1), pg.hostRows(0, v == 1).ptr + w0 * row, pg.wr * row), "cuMemcpyHtoD"); + } + var ids_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * idw * 4); + defer ids_d.free(); + var nk_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * 4); + defer nk_d.free(); + var sp_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * 4); + defer sp_d.free(); + const ids = try a.alloc(i32, 256 * idw); + defer a.free(ids); + const back = try a.alloc(i32, 256 * idw); + defer a.free(back); + const nk = try a.alloc(i32, 256); + defer a.free(nk); + const sp = try a.alloc(i32, 256); + defer a.free(sp); + @memset(sp, 1); + const lay = pg.slotRows(0); + const dk = try a.alloc(u8, lay[0] * row); + defer a.free(dk); + const dv = try a.alloc(u8, lay[0] * row); + defer a.free(dv); + var model: Model = .{ .stamp = try a.alloc(u32, slots), .slot_of = .init(a), .blk = try a.alloc(i64, slots) }; + defer { + a.free(model.stamp); + a.free(model.blk); + model.slot_of.deinit(); + } + @memset(model.stamp, 0); + @memset(model.blk, -1); + var prng = std.Random.DefaultPrng.init(42); + const rng = prng.random(); + // a drifting selection: 512 distinct blocks of [0, nb_before), 15% replaced a call + var base: [512]u32 = undefined; + var used = std.AutoHashMap(u32, void).init(a); + defer used.deinit(); + var k: usize = 0; + while (k < 512) { + const b = rng.uintLessThan(u32, @intCast(nb_before)); + if ((try used.getOrPut(b)).found_existing) continue; + base[k] = b; + k += 1; + } + var gpu_miss: usize = 0; + var model_miss: usize = 0; + var bad: usize = 0; + var keys_checked: usize = 0; + var plan_ns: i96 = 0; + const calls: usize = 60; + var execs: [2]?cuda.graph.Exec = @splat(null); + defer for (&execs) |*exec| if (exec.*) |*e| e.deinit(); + var captured_window = w0; + var captures: usize = 0; + var replays: usize = 0; + for (0..calls) |c| { + const m: usize = if (c % 10 == 9) 64 else 1; + const pos0 = w0 + pg.wr - 70 + c; + if (moving) { + try pg.prepareLayer(stream, 0, pos0, m); + const win = pg.win0[0]; + for (0..2) |v| try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(pg.devRows(0, v == 1) + (pos0 - win) * row, pg.hostRows(0, v == 1).ptr + pos0 * row, m * row, stream.handle), "cuMemcpyHtoDAsync"); + if (win != captured_window) { + try stream.synchronize(); + for (&execs) |*exec| if (exec.*) |*e| { + e.deinit(); + exec.* = null; + }; + captured_window = win; + } + } + var need = std.AutoHashMap(u32, void).init(a); + defer need.deinit(); + for (0..m) |r| { + // drift: replace ~15% of the base blocks + for (0..77) |_| { + const j = rng.uintLessThan(usize, 512); + while (true) { + const b = rng.uintLessThan(u32, @intCast(nb_before)); + if ((try used.getOrPut(b)).found_existing) continue; + _ = used.remove(base[j]); + base[j] = b; + break; + } + } + var sorted = base; + std.mem.sort(u32, &sorted, {}, std.sort.asc(u32)); + const out = ids[r * idw ..][0..idw]; + for (sorted, 0..) |b, j| for (0..4) |t| { + out[j * 4 + t] = @intCast(b * 4 + t); + }; + // 3 tail keys inside the window + for (0..3) |t| out[2048 + t] = @intCast(if (moving) pos0 - 3 + t else w0 + 100 + c * 4 + t); + nk[r] = 2051; + for (sorted) |b| try need.put(b, {}); + } + var need_list: std.ArrayList(u32) = .empty; + defer need_list.deinit(a); + var it = need.keyIterator(); + while (it.next()) |b| try need_list.append(a, b.*); + model_miss += try model.use(a, need_list.items); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(ids_d.ptr, ids.ptr, m * idw * 4), "cuMemcpyHtoD"); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(nk_d.ptr, nk.ptr, m * 4), "cuMemcpyHtoD"); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(sp_d.ptr, sp.ptr, m * 4), "cuMemcpyHtoD"); + const before = try pg.stats(0); + const t0 = check.now(gpu.io); + if (graphed) { + const slot: usize = @intFromBool(m == 64); + if (execs[slot] == null) { + try cuda.graph.beginCapture(stream, .thread_local); + try pg.page(stream, 0, ids_d.ptr, nk_d.ptr, sp_d.ptr, m); + var template = try cuda.graph.endCapture(stream); + defer template.deinit(); + execs[slot] = try template.instantiate(); + captures += 1; + } else replays += 1; + try execs[slot].?.launchOn(stream); + } else try pg.page(stream, 0, ids_d.ptr, nk_d.ptr, sp_d.ptr, m); + try stream.synchronize(); + plan_ns += check.now(gpu.io) - t0; + const after = try pg.stats(0); + try check.expect(after.epoch == c + 1 and after.calls == c + 1, "pager device epoch/calls advance exactly once on launch {d}", .{c + 1}); + gpu_miss += after.misses - before.misses; + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(back.ptr, ids_d.ptr, m * idw * 4), "cuMemcpyDtoH"); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(dk.ptr, pg.devRows(0, false), dk.len), "cuMemcpyDtoH"); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(dv.ptr, pg.devRows(0, true), dv.len), "cuMemcpyDtoH"); + const hk = pg.hostRows(0, false); + const hv = pg.hostRows(0, true); + for (0..m) |r| for (0..2051) |ki| { + const p: usize = @intCast(ids[r * idw + ki]); + const q: usize = @intCast(back[r * idw + ki]); + keys_checked += 1; + if (p >= pg.win0[0]) { + bad += @intFromBool(q != p - pg.win0[0]); + continue; + } + if (q < lay[1] or q >= lay[0]) { + bad += 1; + continue; + } + bad += @intFromBool(!std.mem.eql(u8, dk[q * row ..][0..row], hk[p * row ..][0..row]) or !std.mem.eql(u8, dv[q * row ..][0..row], hv[p * row ..][0..row])); + }; + } + const st = try pg.stats(0); + if (graphed) { + std.debug.print("pager graphs: captures {d}, replays {d}, device epoch {d}, calls {d}, rebases {d}, written rows {d}\n", .{ captures, replays, st.epoch, st.calls, st.rebases, st.written_rows }); + try check.expect(replays > 0 and st.epoch == calls and st.calls == calls, "captured pager executes fresh epochs", .{}); + } + if (moving) { + const fed = w0 + pg.wr - 70 + calls + 63; + try pg.flushLayer(stream, 0, fed); + try stream.synchronize(); + for (0..2) |v| { + const words = std.mem.bytesAsSlice(u32, pg.hostRows(0, v == 1)); + for (0..fed) |position| for (0..row / 4) |column| { + const expected: u32 = @truncate(position *% 2654435761 +% column *% 40503 +% v *% 977); + if (words[position * (row / 4) + column] != expected) bad += 1; + }; + } + try check.expect(st.rebases > 0 and st.written_rows > 0 and captures > 2, "window moves flush rows and recapture pager graphs", .{}); + } + std.debug.print("kernel: {d} calls ({d} keys checked): bad {d}; misses GPU {d}, host LRU model {d}; needed {d}; overflow {}; plan+fetch+remap {d:.1} us a call (incl. launch + sync)\n", .{ calls, keys_checked, bad, gpu_miss, model_miss, st.needed, st.overflow, @as(f64, @floatFromInt(plan_ns)) / 1e3 / @as(f64, @floatFromInt(calls)) }); + try check.expect(bad == 0 and !st.overflow, "every remapped key reads its own host row", .{}); + const diff = @as(f64, @floatFromInt(@max(gpu_miss, model_miss) - @min(gpu_miss, model_miss))) / @as(f64, @floatFromInt(@max(model_miss, 1))); + try check.expect(diff < 0.05, "misses within 5% of the LRU model ({d:.3})", .{diff}); + check.pass("exl3-kvpage kernel", .{}); +} + +fn overflowReplay(gpu: Gpu) !void { + const idw = flash.exl3.prompt_wt.index_idw; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + const pg = try kvp.Paged.init(gpu.gpa, gpu.d, 4096, .{ .window = 2052, .rows = 1, .slots = 1 }); + defer pg.destroy(); + pg.setWindow(0, 2052); + var ids = try cuda.DeviceBuffer.alloc(gpu.d, idw * 4); + defer ids.free(); + var nk = try cuda.DeviceBuffer.alloc(gpu.d, 4); + defer nk.free(); + var sp = try cuda.DeviceBuffer.alloc(gpu.d, 4); + defer sp.free(); + var input: [idw]i32 = @splat(0); + for (0..2048) |index| input[index] = @intCast(index); + try ids.upload(0, std.mem.sliceAsBytes(&input)); + const count: i32 = 2048; + const sparse: i32 = 1; + try nk.upload(0, std.mem.asBytes(&count)); + try sp.upload(0, std.mem.asBytes(&sparse)); + try cuda.graph.beginCapture(stream, .thread_local); + try pg.page(stream, 0, ids.ptr, nk.ptr, sp.ptr, 1); + var template = try cuda.graph.endCapture(stream); + defer template.deinit(); + var exec = try template.instantiate(); + defer exec.deinit(); + for (0..2) |_| { + try nk.upload(0, std.mem.asBytes(&count)); + try exec.launchOn(stream); + try stream.synchronize(); + var actual: i32 = -1; + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(&actual, nk.ptr, 4), "cuMemcpyDtoH"); + const stats = try pg.stats(0); + try check.expect(actual == 0 and stats.overflow and stats.misses == 0, "overflow suppresses fetch and attention key reads on every replay", .{}); + } + check.pass("exl3-kvpage graph overflow fails closed before key reads", .{}); +} + +test "exl3: invalid paged snapshot geometry rejects before allocation, CUDA access or TP peer dispatch" { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var driver: cuda.Driver = undefined; + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.driver = &driver; + engine.stream = &stream; + var session: flash.rank_engine.Session = undefined; + session.gpa = failing.allocator(); + session.engine = &engine; + session.next_object = 77; + var pager: kvp.Paged = undefined; + pager.cfg = .{ .window = 0 }; + pager.cap = 200_016; + var seq: prompt.Seq = undefined; + seq.invalidated = false; + seq.paged = &pager; + var draft: flash.mtp.Draft = undefined; + draft.verify_cache = .{ .generation = 17, .ready = 17, .rows = 4, .position = 200_000 }; + const before = draft.verify_cache; + for ([_]usize{ 0, 1, 3, 4, 2051, 2052, 4096, 200_000 }) |at| { + seq.fed = at; + seq.cap = 200_016; + seq.history = .{ 81, 93 }; + try std.testing.expectError(error.KvPageConfig, session.save(&seq, null)); + try std.testing.expectError(error.KvPageConfig, session.save(&seq, &draft)); + try std.testing.expectEqual(at, seq.fed); + try std.testing.expectEqual([2]i64{ 81, 93 }, seq.history); + try std.testing.expectEqual(@as(u64, 77), session.next_object); + try std.testing.expectEqualDeep(before, draft.verify_cache); + } + try std.testing.expectEqual(@as(usize, 0), failing.alloc_index); + try std.testing.expect(!failing.has_induced_failure); +} + +test "exl3: paged flush orders all pending BF16 window rows before host snapshot reads" { + const Fake = struct { + const Copy = struct { dst: u64, src: u64, len: usize }; + var pending: [24]Copy = undefined; + var count: usize = 0; + + fn copy(dst: u64, src: u64, len: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + std.debug.assert(count < pending.len); + pending[count] = .{ .dst = dst, .src = src, .len = len }; + count += 1; + return cuda.abi.success; + } + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + for (pending[0..count]) |item| { + const dst: [*]u8 = @ptrFromInt(item.dst); + const src: [*]const u8 = @ptrFromInt(item.src); + @memcpy(dst[0..item.len], src[0..item.len]); + } + count = 0; + return cuda.abi.success; + } + + fn pattern(ai: usize, side: usize, position: usize) u8 { + return @intCast(ai * 7 + side * 3 + position); + } + }; + Fake.count = 0; + var driver: cuda.Driver = undefined; + driver.api.cuMemcpyDtoDAsync_v2 = Fake.copy; + driver.api.cuStreamSynchronize = Fake.sync; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var pg: kvp.Paged = undefined; + pg.d = &driver; + pg.cfg = .{ .window = 2052, .rows = 1, .slots = 1 }; + pg.cap = 32; + pg.wr = pg.cfg.window + pg.cfg.rows + 4; + const layout = try kvp.SnapshotLayout.init(pg.cap, pg.cfg, 20); + pg.nblocks = layout.nblocks; + pg.dev.len = layout.device; + pg.tabs.len = layout.tables; + pg.rebases = @splat(7); + pg.written = @splat(12); + const a = std.testing.allocator; + const device = try a.alloc(u8, pg.devBytes()); + defer a.free(device); + const host = try a.alignedAlloc(u8, .@"16", 24 * pg.cap * kvp.row_bytes); + defer a.free(host); + @memset(host, 0xFF); + pg.dev.ptr = @intFromPtr(device.ptr); + pg.host = .{ .d = &driver, .bytes = @alignCast(host) }; + pg.host_dev = @intFromPtr(host.ptr); + const fed: usize = 20; + for (0..12) |ai| { + pg.win0[ai] = ai % 3 * 4; + for (0..2) |side| { + const canonical = pg.hostMut(ai, side == 1); + for (0..pg.win0[ai]) |position| @memset(canonical[position * kvp.row_bytes ..][0..kvp.row_bytes], Fake.pattern(ai, side, position)); + const window: [*]u8 = @ptrFromInt(pg.devRows(ai, side == 1)); + for (pg.win0[ai]..fed) |position| @memset(window[(position - pg.win0[ai]) * kvp.row_bytes ..][0..kvp.row_bytes], Fake.pattern(ai, side, position)); + } + } + const origins = pg.win0; + try pg.flush(stream, fed); + try std.testing.expectEqual(@as(usize, 24), Fake.count); + for (0..12) |ai| for (0..2) |side| { + try std.testing.expectEqual(@as(u8, 0xFF), pg.hostRows(ai, side == 1)[pg.win0[ai] * kvp.row_bytes]); + }; + try stream.synchronize(); + try std.testing.expectEqual(@as(usize, 0), Fake.count); + for (0..12) |ai| for (0..2) |side| { + const canonical = pg.hostRows(ai, side == 1); + for (0..fed) |position| for (canonical[position * kvp.row_bytes ..][0..kvp.row_bytes]) |byte| { + try std.testing.expectEqual(Fake.pattern(ai, side, position), byte); + }; + for (canonical[fed * kvp.row_bytes ..]) |byte| try std.testing.expectEqual(@as(u8, 0xFF), byte); + }; + try std.testing.expectEqual(origins, pg.win0); + try std.testing.expectEqual(@as([12]u64, @splat(7)), pg.rebases); + try std.testing.expectEqual(@as([12]u64, @splat(12)), pg.written); +} diff --git a/zig/tests/cuda/exl3_moe.zig b/zig/tests/cuda/exl3_moe.zig new file mode 120000 index 0000000000..b491e13470 --- /dev/null +++ b/zig/tests/cuda/exl3_moe.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_moe.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_pack_tests.zig b/zig/tests/cuda/exl3_pack_tests.zig new file mode 100644 index 0000000000..f36ae098ca --- /dev/null +++ b/zig/tests/cuda/exl3_pack_tests.zig @@ -0,0 +1,584 @@ +//! exl3-pack: v1.1.0's EXL3 operators (exl3.Dense, exl3_experts.Experts) on a real Flash Next EXL3 pack's shapes. +//! Each projection: the error of rows 1 and 5 against an exact float64 product (x · diag(suh) H Wq H diag(svh)), +//! then the min over 5 batches of the mean call time (cudaEvent). The routed experts: one row's ten picks of +//! layer 0 (gate, up, SwiGLU in fp32, down, weighted sum) against float64 the same way. +//! Usage: tf-cuda-test exl3-pack [launches per batch, default 20] +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const exl3fmt = core.exl3_format; +const st = core.safetensors; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const shard_count = 7; + +const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + fn get(p: *const Pack, name: []const u8) ?st.Tensor { + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + return null; + } +}; + +/// One EXL3 matrix copied to aligned host memory: trellis int16 [K/16, N/16, 8 * halves], fp16 suh [K], svh [N]. +const Mat = struct { + trellis: []i16, + suh: []f16, + svh: []f16, + tiles_k: usize, + tiles_n: usize, + halves: u16, + + fn k(m: Mat) usize { + return m.tiles_k * 16; + } + fn n(m: Mat) usize { + return m.tiles_n * 16; + } +}; + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +fn mat(p: *const Pack, a: std.mem.Allocator, base: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = p.get(try std.fmt.bufPrint(&buf, "{s}.trellis", .{base})) orelse return error.MissingTensor; + const su = p.get(try std.fmt.bufPrint(&buf, "{s}.suh", .{base})) orelse return error.MissingTensor; + const sv = p.get(try std.fmt.bufPrint(&buf, "{s}.svh", .{base})) orelse return error.MissingTensor; + if (p.get(try std.fmt.bufPrint(&buf, "{s}.mul1", .{base})) == null) return error.NotMul1; + if (t.rank != 3 or t.dtype != .i16) return error.UnexpectedTensor; + return .{ + .trellis = try copyAs(i16, a, t.bytes), + .suh = try copyAs(f16, a, su.bytes), + .svh = try copyAs(f16, a, sv.bytes), + .tiles_k = t.dim(0), + .tiles_n = t.dim(1), + .halves = @intCast(t.dim(2) / 8), + }; +} + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +/// y [rows, N] = x [rows, K] · diag(suh) H Wq H diag(svh), all in float64 (x taken as given, no fp16 step). +fn exact(a: std.mem.Allocator, m: Mat, x: []const f64, rows: usize, y: []f64) !void { + const k = m.k(); + const n = m.n(); + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = m.halves }, .mul1, m.trellis, m.tiles_k, m.tiles_n, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(m.suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(m.svh[c])); + }; +} + +const Err = struct { max_abs: f64 = 0, peak: f64 = 0, worst_rel: f64 = 0 }; + +fn compareBf16(got: []const u16, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(bfValue(g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +fn compareF32(got: []const f32, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(@as(f64, g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +const Bench = struct { + gpu: Gpu, + s: cuda.Stream, + ev0: cuda.Event, + ev1: cuda.Event, + reps: usize, + total_ms: f64 = 0, + + fn time(b: *Bench, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try b.s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try b.ev0.record(b.s); + for (0..b.reps) |_| try f(ctx); + try b.ev1.record(b.s); + try b.ev1.synchronize(); + const ms = try b.ev0.elapsedMs(b.ev1); + best = @min(best, @as(f64, ms) / @as(f64, @floatFromInt(b.reps))); + } + return best; + } +}; + +const Call = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + fn go(c: Call) anyerror!void { + try c.de.decode(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +fn randX(a: std.mem.Allocator, rng: std.Random, n: usize) !struct { bits: []u16, vals: []f64 } { + const bits = try a.alloc(u16, n); + const vals = try a.alloc(f64, n); + for (bits, vals) |*b, *v| { + b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + v.* = bfValue(b.*); + } + return .{ .bits = bits, .vals = vals }; +} + +fn linear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8, impl: @TypeOf(de.implementation)) !void { + const m = mat(pack, a, name) catch |e| { + std.debug.print("{s}: skipped ({t})\n", .{ name, e }); + return; + }; + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + if (impl == .tiles) { + _ = exl3.Dense.candidateWeight(w) catch |e| { + std.debug.print("{s:<58} tiles: refused ({t})\n", .{ name, e }); + return; + }; + } + de.implementation = impl; + defer de.implementation = .reference; + for ([_]usize{ 1, 5 }) |rows| { + const x = try randX(a, rng, rows * k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, rows * n * 2); + defer yd.free(); + try yd.fill8(0xff, b.s.handle); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = rows }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(got)); + // the reference row count: lm_head's 636M products a row, so only its first row + const ref_rows: usize = if (n > 100_000) 1 else rows; + const want = try a.alloc(f64, ref_rows * n); + try exact(a, m, x.vals[0 .. ref_rows * k], ref_rows, want); + const e = compareBf16(got[0 .. ref_rows * n], want); + // a second call's bytes equal the first's (deterministic) + try call.go(); + try b.s.synchronize(); + const again = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(again)); + const same = std.mem.eql(u16, got, again); + const ms = try b.time(call, Call.go); + if (rows == 1) b.total_ms += ms; + std.debug.print("{s:<58} {s:<9} K={d:<5} N={d:<6} halves={d:<2} rows={d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2} repeat={s}\n", .{ name, @tagName(impl), k, n, m.halves, rows, ms, e.max_abs, e.peak, e.worst_rel, if (same) "same" else "DIFF" }); + if (rows == 1) { + // the same row with fp32 outputs: the accumulation's own error, without bf16's output rounding + var yf = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 4); + defer yf.free(); + try de.decode(b.s, w, xd.ptr, .bf16, 1, yf.ptr, .f32); + try b.s.synchronize(); + const gf = try a.alloc(f32, n); + try yf.download(0, std.mem.sliceAsBytes(gf)); + const ef = compareF32(gf, want[0..n]); + std.debug.print("{s:<58} {s:<9} fp32 out rows=1: max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ name, @tagName(impl), ef.max_abs, ef.peak, ef.worst_rel }); + } + } +} + +const ExpCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + pick: u64, + wts: u64, + out: u64, + sc: exl3e.Scratch, + fn go(c: ExpCall) anyerror!void { + try c.ex.run(c.s, c.l, c.x, c.l.dims, .bf16, c.pick, c.wts, 1, c.out, c.sc); + } +}; + +/// Ten experts of `layer` as E = 10, one row picking all ten (weights 0.1): fp32 out vs float64. The first `mapped` +/// experts' trellis words sit in mapped page-locked host memory (the kernel reads them over PCIe, no copy). +fn experts(b: *Bench, ex: *const exl3e.Experts, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, layer: usize, act_mode: i32, mapped: usize) !void { + const count = 10; + var buf: [160]u8 = undefined; + var g: [count]Mat = undefined; + var u: [count]Mat = undefined; + var dn: [count]Mat = undefined; + for (0..count) |e| { + g[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.gate_proj", .{ layer, e })); + u[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.up_proj", .{ layer, e })); + dn[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.down_proj", .{ layer, e })); + } + const dims = g[0].k(); + const width = g[0].n(); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*x| x.free(); + bufs.deinit(b.gpu.gpa); + } + const up = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(gpu.d, bytes); + try list.append(gpu.gpa, d); + return d.ptr; + } + }.f; + var hosts: std.ArrayList(cuda.HostBuffer) = .empty; + defer { + for (hosts.items) |*h| h.free(); + hosts.deinit(b.gpu.gpa); + } + const hostUp = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.HostBuffer), bytes: []const u8) !u64 { + var h = try cuda.HostBuffer.allocMapped(gpu.d, bytes.len); + errdefer h.free(); + @memcpy(h.bytes[0..bytes.len], bytes); + try list.append(gpu.gpa, h); + return try h.device(); + } + }.f; + var gp: [count]u64 = undefined; + var upp: [count]u64 = undefined; + var dp: [count]u64 = undefined; + var gk2: [count]i32 = undefined; + var uk2: [count]i32 = undefined; + var dk2: [count]i32 = undefined; + var suh_g: std.ArrayList(f16) = .empty; + var suh_u: std.ArrayList(f16) = .empty; + var svh_g: std.ArrayList(f16) = .empty; + var svh_u: std.ArrayList(f16) = .empty; + var suh_d: std.ArrayList(f16) = .empty; + var svh_d: std.ArrayList(f16) = .empty; + for (0..count) |e| { + if (e < mapped) { + gp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(dn[e].trellis)); + } else { + gp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(dn[e].trellis)); + } + gk2[e] = g[e].halves; + uk2[e] = u[e].halves; + dk2[e] = dn[e].halves; + try suh_g.appendSlice(a, g[e].suh); + try suh_u.appendSlice(a, u[e].suh); + try svh_g.appendSlice(a, g[e].svh); + try svh_u.appendSlice(a, u[e].svh); + try suh_d.appendSlice(a, dn[e].suh); + try svh_d.appendSlice(a, dn[e].svh); + } + const l: exl3e.Layer = .{ + .gate = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gp)), + .up = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&upp)), + .down = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dp)), + .gate_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gk2)), + .up_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&uk2)), + .down_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dk2)), + .suh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_g.items)), + .suh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_u.items)), + .svh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_g.items)), + .svh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_u.items)), + .suh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_d.items)), + .svh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_d.items)), + .experts = count, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = act_mode, + }; + const x = try randX(a, rng, dims); + var pick: [count]i32 = undefined; + for (&pick, 0..) |*p, i| p.* = @intCast((i * 7) % count); + var wts: [count]f32 = @splat(0.1); + const sbytes = try exl3e.scratchBytes(l, 1, count); + var sbuf = try cuda.DeviceBuffer.alloc(b.gpu.d, sbytes); + defer sbuf.free(); + var outd = try cuda.DeviceBuffer.alloc(b.gpu.d, dims * 4); + defer outd.free(); + try outd.fill8(0xff, b.s.handle); + const call: ExpCall = .{ + .ex = ex, + .s = b.s, + .l = l, + .x = try up(b.gpu, &bufs, std.mem.sliceAsBytes(x.bits)), + .pick = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&pick)), + .wts = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&wts)), + .out = outd.ptr, + .sc = try exl3e.scratch(sbuf.ptr, l, 1, count), + }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(f32, dims); + try outd.download(0, std.mem.sliceAsBytes(got)); + // float64: out = sum_e 0.1 * down_e(silu(gate_e(x)) * up_e(x)) + const want = try a.alloc(f64, dims); + @memset(want, 0); + const gy = try a.alloc(f64, width); + const uy = try a.alloc(f64, width); + const dy = try a.alloc(f64, dims); + for (pick) |pe| { + const e: usize = @intCast(pe); + try exact(a, g[e], x.vals, 1, gy); + try exact(a, u[e], x.vals, 1, uy); + for (gy, uy) |*gv, uv| gv.* = gv.* / (1.0 + @exp(-gv.*)) * uv; + try exact(a, dn[e], gy, 1, dy); + for (want, dy) |*o, v| o.* += 0.1 * v; + } + const err = compareF32(got, want); + const ms = try b.time(call, ExpCall.go); + // with a 64 MiB memset between calls (evicts the 32 MiB L2, so mapped host bytes come over PCIe each call), + // minus the memset alone + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const Flushed = struct { + c: ExpCall, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = true }, Flushed.go); + const alone = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = false }, Flushed.go); + std.debug.print(" cold L2 (64 MiB memset between calls, its time subtracted): {d:.4} ms\n", .{with - alone}); + std.debug.print("routed experts L{d} (1 row x 10 picks, act_mode {d}, {d} of 10 from mapped host memory) K={d} I={d} halves={d}/{d}/{d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ layer, act_mode, mapped, dims, width, g[0].halves, u[0].halves, dn[0].halves, ms, err.max_abs, err.peak, err.worst_rel }); + if (mapped > 0) { + // the copy-engine alternative: the same experts' trellis bytes, pinned host to device, one copy an expert matrix + var total: usize = 0; + for (0..mapped) |e| total += (g[e].trellis.len + u[e].trellis.len + dn[e].trellis.len) * 2; + var src = try cuda.HostBuffer.alloc(b.gpu.d, total); + defer src.free(); + var dst = try cuda.DeviceBuffer.alloc(b.gpu.d, total); + defer dst.free(); + const Copy = struct { + d: *const cuda.Driver, + s: cuda.Stream, + src: []u8, + dst: u64, + fn go(c: @This()) anyerror!void { + try c.d.check(c.d.api.cuMemcpyHtoDAsync_v2(c.dst, c.src.ptr, c.src.len, c.s.handle), "cuMemcpyHtoDAsync"); + } + }; + const c: Copy = .{ .d = b.gpu.d, .s = b.s, .src = src.bytes[0..total], .dst = dst.ptr }; + const cms = try b.time(c, Copy.go); + std.debug.print(" pinned H2D copy of those {d} experts' trellis ({d} bytes): {d:.4} ms ({d:.1} GB/s)\n", .{ mapped, total, cms, @as(f64, @floatFromInt(total)) / cms / 1e6 }); + } +} + +/// Rows=1 timing only, warm and with a 64 MiB memset between calls (the memset's own time subtracted). +fn coldLinear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8) !void { + const m = try mat(pack, a, name); + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const x = try randX(a, rng, k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 2); + defer yd.free(); + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = 1 }; + const warm = try b.time(call, Call.go); + const F = struct { + c: Call, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(F{ .c = call, .f = flush.ptr, .with_call = true }, F.go); + const alone = try b.time(F{ .c = call, .f = flush.ptr, .with_call = false }, F.go); + const bytes = m.trellis.len * 2; + std.debug.print("{s:<58} rows=1 bytes={d:>10} warm {d:.4} ms cold {d:.4} ms ({d:.0} GB/s cold)\n", .{ name, bytes, warm, with - alone, @as(f64, @floatFromInt(bytes)) / (with - alone) / 1e6 }); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const reps: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 20; + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var b: Bench = .{ .gpu = gpu, .s = s, .ev0 = try cuda.Event.init(gpu.d, true), .ev1 = try cuda.Event.init(gpu.d, true), .reps = reps }; + defer b.ev0.deinit(); + defer b.ev1.deinit(); + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.v_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + "model.language_model.layers.3.self_attn.indexer.index_qk_proj", + "model.language_model.layers.0.mlp.experts.0.gate_proj", + "model.language_model.layers.0.mlp.experts.0.down_proj", + "lm_head", + }; + if (rest.len > 2 and std.mem.eql(u8, rest[2], "cold")) { + for (names) |nm| try coldLinear(&b, &de, &pack, a, rng, nm); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + return; + } + if (rest.len > 2 and std.mem.eql(u8, rest[2], "experts")) { + for ([_]usize{ 0, 1, 4, 10 }) |mapped| try experts(&b, &ex, &pack, a, rng, 0, 1, mapped); + return; + } + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .reference); + std.debug.print("sum of rows=1 ms over the projections above (reference Dense): {d:.3}\n", .{b.total_ms}); + b.total_ms = 0; + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .tiles); + try experts(&b, &ex, &pack, a, rng, 0, 1, 0); + try experts(&b, &ex, &pack, a, rng, 0, 0, 0); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + check.pass("exl3-pack: v1.1.0 EXL3 operators ran on the pack's shapes (see the error and timing lines)", .{}); +} + +/// nccl-lat : two processes, one GPU each; fp32 all-reduce (sum) latency at a few sizes, cudaEvent +/// timed on rank 0 (min over 5 batches of 100). Rank 0 writes NCCL's unique id to the file, rank 1 reads it. +pub fn ncclLatency(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const rank = try std.fmt.parseInt(c_int, rest[0], 10); + const id_path = rest[1]; + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + const hex = std.fmt.bytesToHex(id.internal, .lower); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = id_path, .data = &hex }); + } else { + var hex: [256]u8 = undefined; + var tries: usize = 0; + while (true) : (tries += 1) { + const got = std.Io.Dir.cwd().readFile(gpu.io, id_path, &hex) catch { + if (tries > 200) return error.NoUniqueId; + try std.Io.sleep(gpu.io, .fromMilliseconds(50), .awake); + continue; + }; + if (got.len == 256) break; + if (tries > 200) return error.NoUniqueId; + try std.Io.sleep(gpu.io, .fromMilliseconds(50), .awake); + } + _ = try std.fmt.hexToBytes(&id.internal, &hex); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, id, rank), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var ev0 = try cuda.Event.init(gpu.d, true); + defer ev0.deinit(); + var ev1 = try cuda.Event.init(gpu.d, true); + defer ev1.deinit(); + const sizes = [_]usize{ 2560, 11 * 2560, 4 * 11 * 2560, 16 * 11 * 2560, 512 * 2560, 4096 * 2560 }; + var buf = try cuda.DeviceBuffer.alloc(gpu.d, sizes[sizes.len - 1] * 4); + defer buf.free(); + try buf.fill8(0, s.handle); + for (sizes) |count| { + for (0..10) |_| try lib.check(lib.api.ncclAllReduce(buf.ptr, buf.ptr, count, .f32, .sum, comm, s.handle), "ncclAllReduce"); + try s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try ev0.record(s); + for (0..100) |_| try lib.check(lib.api.ncclAllReduce(buf.ptr, buf.ptr, count, .f32, .sum, comm, s.handle), "ncclAllReduce"); + try ev1.record(s); + try ev1.synchronize(); + best = @min(best, @as(f64, try ev0.elapsedMs(ev1)) / 100.0); + } + if (rank == 0) std.debug.print("RESULT nccl all-reduce f32 {d:>9} values ({d:>9} bytes): {d:.4} ms\n", .{ count, count * 4, best }); + } +} diff --git a/zig/tests/cuda/exl3_place.zig b/zig/tests/cuda/exl3_place.zig new file mode 120000 index 0000000000..be915bde9a --- /dev/null +++ b/zig/tests/cuda/exl3_place.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_place.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_ple2_tests.zig b/zig/tests/cuda/exl3_ple2_tests.zig new file mode 100644 index 0000000000..a6c8b31d6b --- /dev/null +++ b/zig/tests/cuda/exl3_ple2_tests.zig @@ -0,0 +1,608 @@ +//! exl3-ple2: an EXL3 pack's n-gram path on v1.1.0's Flash Next family (cuda_ple.Ngram with `format = .exl3`). +//! 1. Row ids: upstream's `lastIds` against the numpy port of tensorfold qwen4_exp `NGram.ids` (crc32 recorded on the +//! exl3-flashnext-window-wip branch, qbench row 3, 24 tokens) and, element by element, against that branch's +//! `exl3_ple_ids.ids` (ported below as `oursIds`) over every token of all 22 qbench prompts. +//! 2. Rows: the stored bytes through the mapping and through pread + hot cache, byte-equal, crc32 equal to the numpy +//! reference; `Ngram.rows` decoded values (fp16) in both reader modes, crc32 equal to the numpy port of `_ple_rows`. +//! 3. RSS: the 30.4 GiB file is not resident (mapping and pread), and the hot cache's hit ratio over real text. +//! 4. The layer-1 injection through upstream's own kernels (fn_ple_gate, fn_ple_conv) on real EXL3 weights (F16 +//! key/value projections done on the host as the reference's F16 dense does, conv weights F16 -> bf16 as exl3.py +//! does), whole-prompt and token-by-token, against the float64 oracle and the f32/bf16 emulation of the reference. +//! Usage: tf-cuda-test exl3-ple2 [qbench row (3)] [tokens (24)] + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const st = core.safetensors; +const ple = flash.ple; +const x3 = flash.exl3_ple; +const Gpu = check.Gpu; + +const py_ids_crc: u32 = 0xd04d5584; +const py_raw_crc: u32 = 0x2d57eafe; +const py_decoded_crc: u32 = 0xdcf520dc; + +const streams = 4; +const hidden = 2560; +const wide = streams * hidden; +const taps = 4; +const dilation = 3; +const tail_rows = (taps - 1) * dilation; +const eps: f32 = 1e-6; +const prefix = "model.language_model.layers.1.ple."; + +// ---- RSS ---------------------------------------------------------------------------------------------------------- + +const Rss = struct { total: u64 = 0, file: u64 = 0, anon: u64 = 0 }; // kB + +fn rss() Rss { + var buf: [1 << 14]u8 = undefined; + const fd = std.c.open("/proc/self/status", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return .{}; + defer _ = std.c.close(fd); + var len: usize = 0; + while (len < buf.len) { + const n = std.c.read(fd, buf[len..].ptr, buf.len - len); + if (n <= 0) break; + len += @intCast(n); + } + var out: Rss = .{}; + var it = std.mem.splitScalar(u8, buf[0..len], '\n'); + while (it.next()) |line| { + const fields = [_]struct { key: []const u8, dst: *u64 }{ .{ .key = "VmRSS:", .dst = &out.total }, .{ .key = "RssFile:", .dst = &out.file }, .{ .key = "RssAnon:", .dst = &out.anon } }; + for (fields) |f| { + if (!std.mem.startsWith(u8, line, f.key)) continue; + var parts = std.mem.tokenizeAny(u8, line[f.key.len..], " \tkB"); + if (parts.next()) |n| f.dst.* = std.fmt.parseInt(u64, n, 10) catch 0; + } + } + return out; +} + +fn showRss(what: []const u8) Rss { + const r = rss(); + std.debug.print("RSS {s}: total {d} MiB, file-backed {d} MiB, anon {d} MiB\n", .{ what, r.total >> 10, r.file >> 10, r.anon >> 10 }); + return r; +} + +// ---- the branch's ids, ported from exl3-flashnext-window-wip zig/src/families/flashnext/exl3_ple_ids.zig ---------- + +/// Row ids [tokens.len][16] after `history` (incremental form; a token equal to eos starts a segment after itself). +fn oursIds(h: ple.Hash, history: [2]i64, tokens: []const u32, out: [][16]u64) void { + var seq: [3]i64 = undefined; + @memcpy(seq[0..2], &history); + var seg: usize = 0; + for (history) |t| seg = if (t == h.eos) 0 else @min(seg + 1, 2); + for (tokens, out[0..tokens.len]) |tok, *row| { + seq[2] = tok; + for (row, 0..) |*id, head| { + var mixed: u64 = 0; + for (0..2 + head / h.per) |q| { + const t = if (seg >= q) seq[2 - q] else h.eos; + const term = @as(u64, @bitCast(t)) *% @as(u64, @bitCast(h.mult[q])); + mixed = if (q == 0) term else mixed ^ term; + } + id.* = @as(u64, @intCast(h.offsets[head])) + mixed % @as(u64, @intCast(h.sizes[head])); + } + seg = if (tok == h.eos) 0 else @min(seg + 1, 2); + seq[0] = seq[1]; + seq[1] = seq[2]; + } +} + +// ---- weights -------------------------------------------------------------------------------------------------------- + +const Index = struct { + parsed: std.json.Parsed(std.json.Value), + files: std.StringHashMapUnmanaged(st.File) = .empty, + dir: []const u8, + + fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Index { + const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}), .dir = dir }; + } + + fn deinit(x: *Index, gpa: std.mem.Allocator, io: std.Io) void { + var it = x.files.valueIterator(); + while (it.next()) |f| f.close(io); + x.files.deinit(gpa); + x.parsed.deinit(); + } + + fn get(x: *Index, gpa: std.mem.Allocator, io: std.Io, name: []const u8) !st.Tensor { + const rel = (x.parsed.value.object.get("weight_map").?.object.get(name) orelse return error.MissingTensor).string; + if (x.files.get(rel)) |f| return f.get(name) orelse error.MissingTensor; + const path = try std.fs.path.join(gpa, &.{ x.dir, rel }); + defer gpa.free(path); + const f = try st.File.open(gpa, io, path); + try x.files.put(gpa, rel, f); + return f.get(name) orelse error.MissingTensor; + } +}; + +fn widen(gpa: std.mem.Allocator, t: st.Tensor, offset: f32) ![]f32 { + const out = try gpa.alloc(f32, t.numel()); + for (out, 0..) |*o, i| { + const w = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + const v: f32 = switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(w))), + .bf16 => @bitCast(@as(u32, w) << 16), + else => return error.UnexpectedTensor, + }; + o.* = v + offset; + } + return out; +} + +fn bf(x: f32) f32 { + return @bitCast(@as(u32, ple.toBf16(x)) << 16); +} + +fn bfBits(x: f32) u16 { + return ple.toBf16(x); +} + +fn bfVal(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +// ---- references of the injection (from exl3-flashnext-window-wip exl3_ple_cpu.zig) --------------------------------- + +const Weights = struct { + key: []const f32, // [wide][hidden] + value: []const f32, // [hidden][hidden] + norm_key: []const f32, // [wide], 1 + w + norm_query: []const f32, + norm_conv: []const f32, + conv: []const f32, // [wide][taps], as stored (F16) +}; + +fn dot32(w: []const f32, x: []const f32) f32 { + var acc: f32 = 0; + for (w, x) |a, b| acc += a * b; + return acc; +} + +/// The keys and values the reference's F16 dense projections store (fp32 sums, bf16 outputs). +fn project(w: Weights, emb: []const f32, keys: []f32, vals: []f32) void { + for (0..wide) |o| keys[o] = bf(dot32(w.key[o * hidden ..][0..hidden], emb)); + for (0..hidden) |o| vals[o] = bf(dot32(w.value[o * hidden ..][0..hidden], emb)); +} + +/// The reference's roundings (glue.py `_ple_gate` / `_ple_conv`) on the host, one token at a time. +const Branch = struct { + w: Weights, + conv_bf: []f32, + gated: []f32, + nrow: []f32, + tail: []f32, + + fn init(gpa: std.mem.Allocator, w: Weights) !Branch { + const conv_bf = try gpa.alloc(f32, w.conv.len); + for (conv_bf, w.conv) |*d, s| d.* = bf(s); + const tail = try gpa.alloc(f32, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .conv_bf = conv_bf, .gated = try gpa.alloc(f32, wide), .nrow = try gpa.alloc(f32, wide), .tail = tail }; + } + + fn step(b: *Branch, keys: []const f32, vals: []const f32, h: []const f32, out: []f32) void { + const w = b.w; + const d: f32 = hidden; + for (0..streams) |s| { + const k = keys[s * hidden ..][0..hidden]; + const q = h[s * hidden ..][0..hidden]; + var kss: f32 = 0; + var qss: f32 = 0; + for (k, q) |a, c| { + kss += a * a; + qss += c * c; + } + const krinv = 1.0 / @sqrt(kss / d + eps); + const qrinv = 1.0 / @sqrt(qss / d + eps); + var dot: f32 = 0; + for (0..hidden) |i| dot += bf(bf(k[i] * krinv * w.norm_key[s * hidden + i]) * bf(q[i] * qrinv * w.norm_query[s * hidden + i])); + var gate = bf(dot / @sqrt(d)); + const mag = @sqrt(@max(@abs(gate), 1e-6)); + gate = bf(if (gate > 0) mag else if (gate < 0) -mag else 0); + const sg = bf(1.0 / (1.0 + @exp(-gate))); + var gss: f32 = 0; + for (0..hidden) |i| { + const g = bf(sg * vals[i]); + b.gated[s * hidden + i] = g; + gss += g * g; + } + const rinv = 1.0 / @sqrt(gss / d + eps); + for (0..hidden) |i| b.nrow[s * hidden + i] = bf(b.gated[s * hidden + i] * rinv * w.norm_conv[s * hidden + i]); + } + for (0..wide) |e| { + var acc: f32 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += b.conv_bf[e * taps + t] * (if (at < tail_rows) b.tail[at * wide + e] else b.nrow[e]); + } + const cv = bf(acc); + out[e] = bf(b.gated[e] + bf(cv / (1.0 + @exp(-cv)))); + } + std.mem.copyForwards(f32, b.tail[0 .. (tail_rows - 1) * wide], b.tail[wide..]); + @memcpy(b.tail[(tail_rows - 1) * wide ..], b.nrow); + } +}; + +/// model.py `PLELayer` in float64 with no intermediate rounding (emb as decoded fp16, weights as stored). +const Oracle = struct { + w: Weights, + tail: []f64, + keys: []f64, + vals: []f64, + gated: []f64, + normed: []f64, + + fn init(gpa: std.mem.Allocator, w: Weights) !Oracle { + const tail = try gpa.alloc(f64, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .tail = tail, .keys = try gpa.alloc(f64, wide), .vals = try gpa.alloc(f64, hidden), .gated = try gpa.alloc(f64, wide), .normed = try gpa.alloc(f64, wide) }; + } + + fn centred(x: []const f64, scale: []const f32, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const r = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + @as(f64, eps)); + for (out, x, scale) |*o, v, s| o.* = v * r * s; + } + + fn step(o: *Oracle, emb: []const f32, h: []const f32, out: []f64) void { + const w = o.w; + for (0..wide) |r| { + var acc: f64 = 0; + for (w.key[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.keys[r] = acc; + } + for (0..hidden) |r| { + var acc: f64 = 0; + for (w.value[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.vals[r] = acc; + } + var kn: [hidden]f64 = undefined; + var qn: [hidden]f64 = undefined; + var hs: [hidden]f64 = undefined; + for (0..streams) |s| { + centred(o.keys[s * hidden ..][0..hidden], w.norm_key[s * hidden ..][0..hidden], &kn); + for (&hs, h[s * hidden ..][0..hidden]) |*d, v| d.* = v; + centred(&hs, w.norm_query[s * hidden ..][0..hidden], &qn); + var dot: f64 = 0; + for (kn, qn) |a, b| dot += a * b; + var gate = dot / @sqrt(@as(f64, hidden)); + const sgn: f64 = if (gate > 0) 1 else if (gate < 0) -1 else 0; + gate = sgn * @sqrt(@max(@abs(gate), 1e-6)); + const sg = 1.0 / (1.0 + @exp(-gate)); + for (0..hidden) |i| o.gated[s * hidden + i] = sg * o.vals[i]; + centred(o.gated[s * hidden ..][0..hidden], w.norm_conv[s * hidden ..][0..hidden], o.normed[s * hidden ..][0..hidden]); + } + for (0..wide) |e| { + var acc: f64 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += @as(f64, w.conv[e * taps + t]) * (if (at < tail_rows) o.tail[at * wide + e] else o.normed[e]); + } + out[e] = o.gated[e] + acc / (1.0 + @exp(-acc)); + } + std.mem.copyForwards(f64, o.tail[0 .. (tail_rows - 1) * wide], o.tail[wide..]); + @memcpy(o.tail[(tail_rows - 1) * wide ..], o.normed); + } +}; + +const Qbench = struct { rows: []const struct { input_ids: []const u32 } }; + +fn crc(bytes: []const u8) u32 { + return std.hash.Crc32.hash(bytes); +} + +fn bytesOf(comptime T: type, v: []const T) []const u8 { + return std.mem.sliceAsBytes(v); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const qrow: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 3; + const count: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24; + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = gpu.io; + const rss_start = showRss("start"); + + const qpath = try std.fs.path.join(a, &.{ dir, "qbench_prompts.json" }); + const qtext = try std.Io.Dir.cwd().readFileAlloc(io, qpath, a, .limited(1 << 26)); + const parsed = try std.json.parseFromSlice(Qbench, a, qtext, .{ .ignore_unknown_fields = true }); + const qrows = parsed.value.rows; + try check.expect(qrow < qrows.len and count <= qrows[qrow].input_ids.len, "qbench row {d} with {d} tokens", .{ qrow, count }); + const tokens = qrows[qrow].input_ids[0..count]; + + // ---- the table through upstream's Ngram (mapping) ----------------------------------------------------------- + var g = try ple.Ngram.init(gpu.gpa, io, dir); + defer g.deinit(); + try check.expect(g.format == .exl3, "Ngram.init took the EXL3 table", .{}); + const h = g.hash; + const heads = h.heads; + const dh = h.dh; + std.debug.print("hash: n {d}, heads {d} x {d}, eos {d}, {d} shards, {d} rows; table {d}-bit, {d} bytes a row, read {t}, emb {t}\n", .{ h.n, heads, dh, h.eos, h.shards, h.starts[h.shards], g.x3.?.bits, g.x3.?.rowBytes(), g.x3.?.opts.read, g.x3.?.opts.emb }); + check.pass("Ngram.init on the pack: config-derived multipliers, head sizes and offsets equal the pack's shipped tensors; 128 shard row counts equal Hash.starts", .{}); + const rss_open = showRss("after Ngram.init (30.4 GiB table mapped, header only)"); + + // ---- (1) ids --------------------------------------------------------------------------------------------------- + const ids = try a.alloc(i64, count * heads); + { + var hist = [_]i64{ h.eos, h.eos }; + for (tokens, 0..) |tok, t| { + try ple.lastIds(h, &hist, tok, ids[t * heads ..][0..heads]); + hist = .{ hist[1], tok }; + } + const got = crc(bytesOf(i64, ids)); + std.debug.print("ids[0][0..4] = {any}; ids crc32 {x:0>8} (numpy NGram.ids {x:0>8})\n", .{ ids[0..4], got, py_ids_crc }); + if (qrow == 3 and count == 24) try check.expect(got == py_ids_crc, "upstream lastIds crc32 {x:0>8} != numpy {x:0>8}", .{ got, py_ids_crc }); + // every token of every qbench prompt: upstream's lastIds against the branch's incremental ids, element by element + var compared: usize = 0; + var eos_seen: usize = 0; + for (qrows) |r| { + const ours = try gpu.gpa.alloc([16]u64, r.input_ids.len); + defer gpu.gpa.free(ours); + oursIds(h, .{ h.eos, h.eos }, r.input_ids, ours); + var hist2 = [_]i64{ h.eos, h.eos }; + for (r.input_ids, ours) |tok, want| { + if (tok == h.eos) eos_seen += 1; + var up: [16]i64 = undefined; + try ple.lastIds(h, &hist2, tok, up[0..heads]); + hist2 = .{ hist2[1], tok }; + for (up[0..heads], want[0..heads]) |u, o| { + if (u < 0 or @as(u64, @intCast(u)) != o) { + std.debug.print("FAIL ids differ at token {d}: upstream {any} ours {any}\n", .{ tok, up[0..heads], want[0..heads] }); + return error.TestFailed; + } + } + compared += 1; + } + } + check.pass("row ids: upstream Hash/lastIds byte-equal to the numpy NGram.ids crc{s} and to the branch's exl3_ple_ids on all {d} tokens of {d} qbench prompts ({d} eos inside)", .{ if (qrow == 3 and count == 24) "" else " (crc not checked: not row 3 x 24)", compared, qrows.len, eos_seen }); + } + + // ---- (2) raw rows, two readers; decoded rows through Ngram.rows, two readers ------------------------------------ + var pread_t = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + const rb = g.x3.?.rowBytes(); + const raw_map = try a.alloc(u8, ids.len * rb); + const raw_ssd = try a.alloc(u8, ids.len * rb); + for (ids, 0..) |id, i| { + const loc = try h.locate(id); + try g.x3.?.raw(loc.shard, loc.local, raw_map[i * rb ..][0..rb]); + try pread_t.raw(loc.shard, loc.local, raw_ssd[i * rb ..][0..rb]); + } + try check.sameBytes("rows: mapping vs pread", raw_ssd, raw_map); + const raw_crc = crc(raw_map); + std.debug.print("{d} rows x {d} bytes: raw crc32 {x:0>8} (numpy {x:0>8})\n", .{ ids.len, rb, raw_crc, py_raw_crc }); + if (qrow == 3 and count == 24) try check.expect(raw_crc == py_raw_crc, "raw rows crc32 {x:0>8} != numpy {x:0>8}", .{ raw_crc, py_raw_crc }); + const rss_map = showRss("after the mapped row reads"); + + const w = g.width(); + const emb_map = try a.alloc(u16, count * w); + const xs_map = try a.alloc(f32, count * w / 32); + { + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_map, xs_map); + } + const dec_crc = crc(bytesOf(u16, emb_map)); + var peak_emb: f32 = 0; + for (emb_map) |v| peak_emb = @max(peak_emb, @abs(@as(f32, @floatCast(@as(f16, @bitCast(v)))))); + std.debug.print("Ngram.rows: {d} fp16 values, max |value| {d:.4}, crc32 {x:0>8} (numpy _ple_rows {x:0>8})\n", .{ emb_map.len, peak_emb, dec_crc, py_decoded_crc }); + if (qrow == 3 and count == 24) try check.expect(dec_crc == py_decoded_crc, "decoded crc32 {x:0>8} != numpy {x:0>8}", .{ dec_crc, py_decoded_crc }); + // the same call with the pread reader swapped in behind Ngram + const emb_ssd = try a.alloc(u16, count * w); + const xs_ssd = try a.alloc(f32, count * w / 32); + { + const mapped_t = g.x3.?; + g.x3 = pread_t; + defer g.x3 = mapped_t; + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_ssd, xs_ssd); + } + try check.sameBytes("Ngram.rows: mapping vs pread", bytesOf(u16, emb_ssd), bytesOf(u16, emb_map)); + try check.sameBytes("Ngram.rows group sums: mapping vs pread", bytesOf(f32, xs_ssd), bytesOf(f32, xs_map)); + check.pass("rows: {d} ids, stored bytes and Ngram.rows fp16 values byte-identical through the mapping and through pread + hot cache{s}", .{ ids.len, if (qrow == 3 and count == 24) "; raw and decoded crc32 equal the numpy references" else "" }); + + // ---- (3) hot cache over real text; RSS ------------------------------------------------------------------------- + { + var fresh = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer fresh.destroy(); + const r0 = rss(); + var buf: [256]u8 = undefined; + var ntok: usize = 0; + for (qrows) |r| { + var hist = [_]i64{ h.eos, h.eos }; + for (r.input_ids) |tok| { + var row_ids: [16]i64 = undefined; + try ple.lastIds(h, &hist, tok, row_ids[0..heads]); + hist = .{ hist[1], tok }; + for (row_ids[0..heads]) |id| { + const loc = try h.locate(id); + try fresh.raw(loc.shard, loc.local, buf[0..rb]); + } + ntok += 1; + } + } + const s = fresh.hotStats(); + const r1 = rss(); + std.debug.print("pread + 4096-row hot cache over all {d} qbench tokens: {d} lookups, {d} hits ({d:.2}%); RSS file-backed {d} -> {d} MiB, anon {d} -> {d} MiB\n", .{ ntok, s.lookups, s.lookups - s.misses, 100.0 * @as(f64, @floatFromInt(s.lookups - s.misses)) / @as(f64, @floatFromInt(s.lookups)), r0.file >> 10, r1.file >> 10, r0.anon >> 10, r1.anon >> 10 }); + try check.expect(r1.file <= r0.file + 16 * 1024, "pread reads grew file-backed RSS by {d} MiB", .{(r1.file - r0.file) >> 10}); + } + const rss_rows = showRss("after rows (both readers) and the hot-cache pass"); + const table_mib = g.x3.?.file.?.map.memory.len >> 20; + std.debug.print("table {d} MiB; file-backed RSS: start {d}, after open {d}, after mapped reads {d}, now {d} MiB\n", .{ table_mib, rss_start.file >> 10, rss_open.file >> 10, rss_map.file >> 10, rss_rows.file >> 10 }); + try check.expect(rss_map.file <= rss_open.file + 128 * 1024 and rss_rows.file <= 512 * 1024, "the table became resident", .{}); + check.pass("30.4 GiB table stays out of memory: file-backed RSS {d} MiB after the mapped reads (+{d} MiB), anon {d} MiB", .{ rss_rows.file >> 10, (rss_map.file - @min(rss_map.file, rss_open.file)) >> 10, rss_rows.anon >> 10 }); + + // ---- (4) layer 1's injection through upstream's fn_ple_gate / fn_ple_conv --------------------------------------- + try injection(gpu, a, dir, tokens, emb_map, h.eos); +} + +fn injection(gpu: Gpu, a: std.mem.Allocator, dir: []const u8, tokens: []const u32, emb_bits: []const u16, eos: i64) !void { + _ = eos; + const io = gpu.io; + const n = tokens.len; + var index = try Index.open(gpu.gpa, io, dir); + defer index.deinit(gpu.gpa, io); + const key_t = try index.get(gpu.gpa, io, prefix ++ "key_proj.weight"); + const val_t = try index.get(gpu.gpa, io, prefix ++ "value_proj.weight"); + const conv_t = try index.get(gpu.gpa, io, prefix ++ "conv1d.weight"); + try check.expect(key_t.is(.f16, &.{ wide, hidden }) and val_t.is(.f16, &.{ hidden, hidden }) and conv_t.is(.f16, &.{ wide, 1, taps }), "PLE key/value/conv dtypes or shapes (want F16)", .{}); + var norms: [3][]f32 = undefined; + for ([_][]const u8{ "norm_key", "norm_query", "norm_conv" }, 0..) |nm, i| { + var nb: [96]u8 = undefined; + const t = try index.get(gpu.gpa, io, try std.fmt.bufPrint(&nb, prefix ++ "{s}.weight", .{nm})); + try check.expect(t.is(.bf16, &.{wide}), "{s} dtype/shape (want bf16 [10240])", .{nm}); + const raw = try widen(a, t, 0); + var mean: f64 = 0; + for (raw) |v| mean += v; + mean /= wide; + std.debug.print("{s} stored mean {d:.4} (gamma - 1 storage: scale 1 + w)\n", .{ nm, mean }); + try check.expect(@abs(mean) < 0.5, "{s} does not look gamma - 1", .{nm}); + norms[i] = try widen(a, t, 1.0); + } + const weights: Weights = .{ .key = try widen(a, key_t, 0), .value = try widen(a, val_t, 0), .norm_key = norms[0], .norm_query = norms[1], .norm_conv = norms[2], .conv = try widen(a, conv_t, 0) }; + // what upstream's loader (cuda_prompt_load.zig with_ple) would upload as CW: the raw bytes, read as bf16 + var raw_as_bf: usize = 0; + for (0..weights.conv.len) |i| { + const bits = std.mem.readInt(u16, conv_t.bytes[2 * i ..][0..2], .little); + if (bfVal(bits) != bf(weights.conv[i])) raw_as_bf += 1; + } + std.debug.print("conv1d.weight is F16: {d} of {d} values change if its bytes are read as bf16 (upstream's loader does), so it is converted (RNE) as exl3.py does\n", .{ raw_as_bf, weights.conv.len }); + + // inputs: the decoded rows (fp16), the layer input stand-in (token embedding tiled over 4 streams, bf16) + const embed_t = try index.get(gpu.gpa, io, "model.language_model.embed_tokens.weight"); + try check.expect(embed_t.dtype == .bf16 and embed_t.dim(1) == hidden, "embed_tokens", .{}); + const emb = try a.alloc(f32, n * hidden); + for (emb, emb_bits) |*e, b| e.* = @floatCast(@as(f16, @bitCast(b))); + const hq = try a.alloc(f32, n * wide); + const hq_bits = try a.alloc(u16, n * wide); + for (tokens, 0..) |tok, t| for (0..streams) |s| for (0..hidden) |d| { + const bits = std.mem.readInt(u16, embed_t.bytes[(@as(usize, tok) * hidden + d) * 2 ..][0..2], .little); + hq_bits[t * wide + s * hidden + d] = bits; + hq[t * wide + s * hidden + d] = bfVal(bits); + }; + // the projections the reference's F16 dense stores (bf16) + const keys = try a.alloc(f32, n * wide); + const vals = try a.alloc(f32, n * hidden); + for (0..n) |t| project(weights, emb[t * hidden ..][0..hidden], keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden]); + const keys_bits = try a.alloc(u16, keys.len); + for (keys_bits, keys) |*b, v| b.* = bfBits(v); + const vals_bits = try a.alloc(u16, vals.len); + for (vals_bits, vals) |*b, v| b.* = bfBits(v); + const cw_bits = try a.alloc(u16, weights.conv.len); + for (cw_bits, weights.conv) |*b, v| b.* = bfBits(v); + + // host references + const want = try a.alloc(f64, n * wide); + const emu = try a.alloc(f32, n * wide); + { + var oracle = try Oracle.init(a, weights); + var branch = try Branch.init(a, weights); + for (0..n) |t| { + oracle.step(emb[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], want[t * wide ..][0..wide]); + branch.step(keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], emu[t * wide ..][0..wide]); + } + } + + // the device: upstream's glue kernels + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const D = cuda.DeviceBuffer; + var d_keys = try D.fromHost(gpu.d, bytesOf(u16, keys_bits)); + defer d_keys.free(); + var d_vals = try D.fromHost(gpu.d, bytesOf(u16, vals_bits)); + defer d_vals.free(); + var d_h = try D.fromHost(gpu.d, bytesOf(u16, hq_bits)); + defer d_h.free(); + var d_nk = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_key)); + defer d_nk.free(); + var d_nq = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_query)); + defer d_nq.free(); + var d_nc = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_conv)); + defer d_nc.free(); + var d_cw = try D.fromHost(gpu.d, bytesOf(u16, cw_bits)); + defer d_cw.free(); + const zeros = try a.alloc(u8, n * wide * 2); + @memset(zeros, 0); + var d_zero = try D.fromHost(gpu.d, zeros); // conv's H: zero, so HOUT is the injection itself + defer d_zero.free(); + var d_tail = try D.fromHost(gpu.d, zeros[0 .. tail_rows * wide * 2]); + defer d_tail.free(); + var d_gated = try D.alloc(gpu.d, n * wide * 2); + defer d_gated.free(); + var d_pss = try D.alloc(gpu.d, n * streams * 4); + defer d_pss.free(); + var d_out = try D.alloc(gpu.d, n * wide * 2); + defer d_out.free(); + var d_nrow = try D.alloc(gpu.d, n * wide * 2); + defer d_nrow.free(); + try gpu.ctx.synchronize(); + + // whole prompt: one launch each, the conv reading its own earlier rows + try glue.pleGate(s, d_keys.ptr, d_vals.ptr, d_h.ptr, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, n, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, n, hidden, streams, taps, dilation); + try s.synchronize(); + const prompt_bits = try a.alloc(u16, n * wide); + try d_out.download(0, std.mem.sliceAsBytes(prompt_bits)); + + // token by token: the tail carried on the host from each launch's NROW, as shiftWindow keeps it + const tail = try a.alloc(u16, tail_rows * wide); + @memset(tail, 0); + const decode_bits = try a.alloc(u16, n * wide); + const nrow = try a.alloc(u16, wide); + for (0..n) |t| { + try d_tail.upload(0, bytesOf(u16, tail)); + try gpu.ctx.synchronize(); + try glue.pleGate(s, d_keys.ptr + t * wide * 2, d_vals.ptr + t * hidden * 2, d_h.ptr + t * wide * 2, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, 1, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, 1, hidden, streams, taps, dilation); + try s.synchronize(); + try d_out.download(0, std.mem.sliceAsBytes(decode_bits[t * wide ..][0..wide])); + try d_nrow.download(0, std.mem.sliceAsBytes(nrow)); + std.mem.copyForwards(u16, tail[0 .. (tail_rows - 1) * wide], tail[wide..]); + @memcpy(tail[(tail_rows - 1) * wide ..], nrow); + } + + var worst: f64 = 0; + var peak: f64 = 0; + var worst_rel_rms: f64 = 0; + var emu_worst: f64 = 0; + var emu_off: usize = 0; + var emu_vs_oracle: f64 = 0; + for (0..n) |t| { + var se: f64 = 0; + var sw: f64 = 0; + for (prompt_bits[t * wide ..][0..wide], want[t * wide ..][0..wide], emu[t * wide ..][0..wide]) |gb, o, e| { + const gv: f64 = bfVal(gb); + const err = gv - o; + worst = @max(worst, @abs(err)); + peak = @max(peak, @abs(o)); + se += err * err; + sw += o * o; + emu_worst = @max(emu_worst, @abs(gv - e)); + if (gv != e) emu_off += 1; + emu_vs_oracle = @max(emu_vs_oracle, @abs(@as(f64, e) - o)); + } + worst_rel_rms = @max(worst_rel_rms, @sqrt(se / sw)); + } + const same_modes = std.mem.eql(u16, prompt_bits, decode_bits); + std.debug.print("injection over {d} tokens: GPU max |gpu - f64| {e:.3} at peak {e:.3} ({d:.3}% of peak), worst per-token rms rel {e:.3}\n", .{ n, worst, peak, 100 * worst / peak, worst_rel_rms }); + std.debug.print(" GPU vs the reference's roundings emulated on the host: max |diff| {e:.3}, {d} of {d} values differ; emulation vs f64 max {e:.3}\n", .{ emu_worst, emu_off, n * wide, emu_vs_oracle }); + std.debug.print(" whole-prompt launch vs token-by-token launches: {s}\n", .{if (same_modes) "bit-identical" else "DIFFER"}); + try check.expect(peak > 0 and worst <= 0.03 * peak and worst_rel_rms <= 0.02, "injection beyond bf16 rounding of the f64 oracle", .{}); + try check.expect(same_modes, "prompt and decode launches differ", .{}); + check.pass("layer-1 PLE injection through upstream's fn_ple_gate/fn_ple_conv on real EXL3 weights: max |err| {e:.2} <= 3% of peak {e:.2}, rms rel {e:.2} <= 2e-2; prompt == decode bits", .{ worst, peak, worst_rel_rms }); +} diff --git a/zig/tests/cuda/exl3_pool_tests.zig b/zig/tests/cuda/exl3_pool_tests.zig new file mode 100644 index 0000000000..edf0ae2a5c --- /dev/null +++ b/zig/tests/cuda/exl3_pool_tests.zig @@ -0,0 +1,904 @@ +//! exl3-pool: the Flash Next EXL3 expert offload pool (families/flashnext/cuda_exl3_pool.zig, cuda/exl3_pool.zig) on +//! upstream's `exl3_experts.Experts`, with real tensors of the pack (port of exl3-flashnext-window-wip's exl3-offload). +//! Host bank: the memory-mapped shards. Reference: the same layer fully resident (513 experts on the device, tables +//! pointing at them). Every comparison is byte for byte: the pool only changes which device address a pointer table +//! entry holds, so the kernels must produce the same bits. +//! +//! exl3-pool equal [layers] pool at 10, 24, 64 lines vs fully resident over rounds forcing misses, evictions, +//! split passes (16-row chunks) and prefetch; line bytes, tables, staged-bytes ledger +//! exl3-pool graph [layers] runSlots captured once over the pool's tables, replayed after line swaps +//! exl3-pool bench [layer] miss cost a expert and a round (staged and pinned-direct), copy/compute overlap +//! exl3-pool hit ... hit rate vs line count: synthetic skewed 200 tokens, and real router picks +//! exl3-pool [layers] equal + graph +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const top_k = 10; +const slots = 11; +const max_rows = 16; +const shard_count = 7; +const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +// ---------------------------------------------------------------------------------------------------- one layer + +/// One layer's experts: host trellis slices (mmap), the stacked scales and the fully resident reference on the device. +const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + + fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + // stacked scales [E, ·]: gate suh [D] svh [I], up suh [D] svh [I], down suh [I] svh [D] + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + for (0..table_n) |e| { + for (projs, 0..) |pn, p| { + const tr = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.trellis", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.suh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.svh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (su.bytes.len != k * 2 or sv.bytes.len != n * 2) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + // the fully resident reference: every expert's trellis back to back, its own tables + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts mapped, fully resident reference {d:.0} MiB; routed k2 gate/up/down {d}/{d}/{d}, shared {d}/{d}/{d}\n", .{ + layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], k2s[0][512], k2s[1][512], k2s[2][512], + }); + return ld; + } + + fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } + + fn scales(ld: *const LayerData) fpool.Scales { + return .{ .suh_g = ld.scale_bufs[0].ptr, .suh_u = ld.scale_bufs[1].ptr, .svh_g = ld.scale_bufs[2].ptr, .svh_u = ld.scale_bufs[3].ptr, .suh_d = ld.scale_bufs[4].ptr, .svh_d = ld.scale_bufs[5].ptr }; + } + + fn source(ld: *const LayerData) fpool.Source { + return .{ .routed = &ld.routed, .shared = ld.shared }; + } +}; + +// ---------------------------------------------------------------------------------------------------- inputs + +/// Skewed expert choice: Zipf(s) over a fixed random permutation of the 512 ids. +const Zipf = struct { + perm: [routed_n]u16, + cdf: [routed_n]f64, + + fn init(seed: u64, s: f64) Zipf { + var z: Zipf = undefined; + var prng = std.Random.DefaultPrng.init(seed); + for (&z.perm, 0..) |*p, i| p.* = @intCast(i); + prng.random().shuffle(u16, &z.perm); + var sum: f64 = 0; + for (&z.cdf, 0..) |*c, i| { + sum += 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&z.cdf) |*c| c.* /= sum; + return z; + } + + /// One row: 10 distinct routed ids then the shared expert. + fn row(z: *const Zipf, r: std.Random, out: []i32) void { + var n: usize = 0; + while (n < top_k) { + const u = r.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and z.cdf[i] < u) i += 1; + const id: i32 = z.perm[i]; + if (std.mem.indexOfScalar(i32, out[0..n], id) != null) continue; + out[n] = id; + n += 1; + } + out[top_k] = 512; + } +}; + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// The device buffers of one engine side (pool or resident): x, picks, weights, scratch, out. +const Io = struct { + x: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + sbuf: cuda.DeviceBuffer, + sc: exl3e.Scratch, + + fn init(gpu: Gpu, l: exl3e.Layer) !Io { + var v: Io = undefined; + v.x = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 2); + v.picks = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.wts = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.out = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 4); + v.sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(l, max_rows, slots)); + v.sc = try exl3e.scratch(v.sbuf.ptr, l, max_rows, slots); + return v; + } + + fn free(v: *Io) void { + v.x.free(); + v.picks.free(); + v.wts.free(); + v.out.free(); + v.sbuf.free(); + } + + fn set(v: *Io, x: []const u16, picks: []const i32, wts: []const f32) !void { + try v.x.upload(0, std.mem.sliceAsBytes(x)); + try v.picks.upload(0, std.mem.sliceAsBytes(picks)); + try v.wts.upload(0, std.mem.sliceAsBytes(wts)); + } + + fn y(v: Io, gpa: std.mem.Allocator, rows: usize) ![]u8 { + const out = try gpa.alloc(u8, rows * slots * dims * 4); + try v.sbuf.d.check(v.sbuf.d.api.cuMemcpyDtoH_v2(out.ptr, v.sc.y, out.len), "cuMemcpyDtoH"); + return out; + } +}; + +fn call(ex: *const exl3e.Experts, s: cuda.Stream, l: exl3e.Layer, v: Io, picks: u64, rows: usize, combine: bool) !void { + if (combine) { + try ex.run(s, l, v.x.ptr, dims, .bf16, picks, v.wts.ptr, rows, v.out.ptr, v.sc); + } else { + try ex.runSlots(s, l, v.x.ptr, dims, .bf16, picks, rows, v.sc); + } +} + +/// One round through the pool: plan, then each pass staged, computed and released. +fn poolRound(pool: *fpool.Pool, ex: *const exl3e.Experts, s: cuda.Stream, v: Io, picks: []const i32, rows: usize, combine: bool) !*const fpool.Round { + const r = try pool.plan(picks, rows); + const l = pool.tables(); + for (0..r.n_passes) |j| { + try pool.stage(r, j, s); + try call(ex, s, l, v, pool.picksFor(r, j, v.picks.ptr), rows, combine); + try pool.release(r, j, s); + } + return r; +} + +fn genRound(z: *const Zipf, rng: std.Random, rows: usize, x: []u16, picks: []i32, wts: []f32) void { + for (x[0 .. rows * dims]) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + for (0..rows) |i| z.row(rng, picks[i * slots ..][0..slots]); + for (wts[0 .. rows * slots]) |*w| w.* = rng.float(f32) * 0.2; +} + +// ---------------------------------------------------------------------------------------------------- gate 2 + +fn verifyPool(gpu: Gpu, pool: *fpool.Pool, ld: *const LayerData) !usize { + const gpa = gpu.gpa; + try pool.lines.drain(); + const lb = pool.lines.layout.line_bytes; + const got = try gpa.alloc(u8, lb); + defer gpa.free(got); + // tables + const tl = pool.lines.tlay; + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + try pool.lines.tables.download(0, tab); + var checked: usize = 0; + var resident: [routed_n]?usize = @splat(null); + for (pool.slots_tab, pool.line_tab, 0..) |id, state, i| { + if (state == .empty) continue; + resident[id] = i; + try pool.lines.mem.download(i * lb, got); + for (0..3) |p| { + const want = ld.routed[id][p]; + try check.sameBytes("pool line vs mapped checkpoint", got[pool.lines.layout.off[p]..][0..want.len], want); + } + checked += 1; + } + for (0..table_n) |e| for (0..3) |p| { + const ptr = std.mem.readInt(u64, tab[tl.ptr[p] + e * 8 ..][0..8], .little); + const k2 = std.mem.readInt(i32, tab[tl.k2[p] + e * 4 ..][0..4], .little); + if (e == 512) { + try check.expect(ptr != 0 and k2 > 0, "shared expert entry missing", .{}); + } else if (resident[e]) |line| { + try check.expect(ptr == pool.lines.at(line, p) and k2 == pool.k2[e][p], "expert {d} proj {d}: table points at {x}, line {d} is at {x}", .{ e, p, ptr, line, pool.lines.at(line, p) }); + } else { + try check.expect(ptr == 0 and k2 == 0, "absent expert {d} has a stale table entry", .{e}); + } + }; + return checked; +} + +const Script = struct { rows: usize, combine: bool }; + +fn equalRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, lines: usize, prefetch_lines: usize, seed: u64) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, ld.source(), ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .prefetch_bytes = 0 }, null); + defer pool.deinit(); + pool.opt.prefetch_bytes = prefetch_lines * pool.lineBytes(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var vr = try Io.init(gpu, ld.res); + defer vr.free(); + const z = Zipf.init(0x5eed + ld.layer, 1.1); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // decode rows, two 16-row prompt chunks, a 4-row chunk; runSlots and run (with combine) alternate + var script: [40]Script = undefined; + for (&script, 0..) |*sc, i| sc.* = .{ .rows = 1, .combine = i % 2 == 1 }; + script[6] = .{ .rows = max_rows, .combine = false }; + script[17] = .{ .rows = max_rows, .combine = true }; + script[25] = .{ .rows = 4, .combine = false }; + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + var ledger_staged: u64 = 0; + var ledger_h2d: u64 = 0; + var ledger_copies: u64 = 0; + var split_rounds: usize = 0; + var max_passes: usize = 0; + var slot_bytes: usize = 0; + for (script, 0..) |sc, ri| { + genRound(&z, rng, sc.rows, &x, &picks, &wts); + const n = sc.rows * slots; + try vp.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vr.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vp.sbuf.fill8(0xff, s.handle); // a slot no pass wrote shows up + const r = try poolRound(&pool, ex, s, vp, picks[0..n], sc.rows, sc.combine); + for (r.copies[0..r.n_copies]) |c| { + const b = ld.routed[c.id]; + ledger_staged += b[0].len + b[1].len + b[2].len; + ledger_h2d += pool.lines.layout.off[2] + b[2].len; + ledger_copies += 1; + } + if (r.masked()) split_rounds += 1; + max_passes = @max(max_passes, r.n_passes); + try call(ex, s, ld.res, vr, vr.picks.ptr, sc.rows, sc.combine); + try s.synchronize(); + const yp = try vp.y(gpa, sc.rows); + defer gpa.free(yp); + const yr = try vr.y(gpa, sc.rows); + defer gpa.free(yr); + check.sameBytes("pool y vs resident y", yp, yr) catch |e| { + std.debug.print(" round {d} ({d} rows, {d} passes) differs\n", .{ ri, sc.rows, r.n_passes }); + return e; + }; + slot_bytes += yp.len; + if (sc.combine) { + const op = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(op); + const or_ = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(or_); + try vp.out.download(0, op); + try vr.out.download(0, or_); + try check.sameBytes("pool combined out vs resident", op, or_); + } + } + const lines_checked = try verifyPool(gpu, &pool, ld); + const ps = pool.stats; + const ls = pool.lines.stats; + const hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)); + std.debug.print(" [{d} lines, prefetch {d}/swap round] {d} rounds ({d} split into passes, up to {d}): routed slots hit {d} miss {d} ({d:.3} hit); {d} copies ({d} misses + {d} prefetch), {d} evictions, {d} protected rounds; H2D {d:.1} MiB, staged {d:.1} MiB (host memcpy {d:.2} ms/expert), staging waits {d}, release waits {d}, gate waits {d}; y {d} bytes compared\n", .{ + lines, prefetch_lines, ps.rounds, split_rounds, max_passes, ps.hit_slots, ps.miss_slots, hit, ls.copies, ps.misses, ps.prefetches, ps.evictions, ps.protected, + @as(f64, @floatFromInt(ls.bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.staged_bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(@max(ls.copies, 1))), ls.staging_waits, ls.release_waits, ls.gate_waits, slot_bytes, + }); + try check.expect(ls.copies == ledger_copies and ls.copies == ps.misses + ps.prefetches, "copy count: device {d}, plans {d}, misses+prefetch {d}", .{ ls.copies, ledger_copies, ps.misses + ps.prefetches }); + try check.expect(ls.staged_bytes == ledger_staged, "staged bytes {d} vs plan ledger {d}", .{ ls.staged_bytes, ledger_staged }); + try check.expect(ls.bytes == ledger_h2d, "H2D bytes {d} vs plan ledger {d}", .{ ls.bytes, ledger_h2d }); + try check.expect(ps.evictions > 0 and ps.miss_slots > 0 and ps.hit_slots > 0 and split_rounds > 0, "a path was not exercised (evictions {d}, misses {d}, hits {d}, split {d})", .{ ps.evictions, ps.miss_slots, ps.hit_slots, split_rounds }); + if (prefetch_lines > 0) try check.expect(ps.prefetches > 0, "no prefetch happened", .{}); + check.pass("layer {d}, {d} lines: every round's runSlots/run output is byte-identical to the fully resident layer; {d} lines and 513 table entries match the checkpoint; staged bytes match the plans", .{ ld.layer, lines, lines_checked }); +} + +// ---------------------------------------------------------------------------------------------------- gate 3 + +fn graphRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const gpa = gpu.gpa; + const lines = 10; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, ld.source(), ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var vr = try Io.init(gpu, ld.res); + defer vr.free(); + const z = Zipf.init(0xc0ffee + ld.layer, 0.6); + var prng = std.Random.DefaultPrng.init(0x96a9 + ld.layer); + const rng = prng.random(); + var x: [dims]u16 = undefined; + var picks: [slots]i32 = undefined; + var wts: [slots]f32 = undefined; + var exec: ?cuda.graph.Exec = null; + defer if (exec) |*e| e.deinit(); + var g: ?cuda.graph.Graph = null; + defer if (g) |*gg| gg.deinit(); + var swaps: usize = 0; + var last: ?[]u8 = null; + defer if (last) |b| gpa.free(b); + var distinct: usize = 0; + const states = 8; + for (0..states) |si| { + genRound(&z, rng, 1, &x, &picks, &wts); + if (si == 5) for (picks[0..5], 0..) |*p, i| { + p.* = @intCast(pool.slots_tab[i]); // half the row from what the lines hold: a partial hit + }; + try vp.set(&x, &picks, &wts); + try vr.set(&x, &picks, &wts); + const r = try pool.plan(&picks, 1); + try check.expect(!r.masked(), "one row must fit 10 lines", .{}); + try pool.stage(r, 0, s); + swaps += r.n_copies; + const l = pool.tables(); + if (exec == null) { + // capture once over the tables, with the outputs poisoned so a capture that ran shows up + try s.synchronize(); + try vp.sbuf.fill8(0xff, s.handle); + try s.synchronize(); + try cuda.graph.beginCapture(s, .thread_local); + try call(ex, s, l, vp, vp.picks.ptr, 1, false); + g = try cuda.graph.endCapture(s); + exec = try g.?.instantiate(); + try exec.?.upload(s); + const still = try vp.y(gpa, 1); + defer gpa.free(still); + try check.expect(std.mem.allEqual(u8, still, 0xff), "capture executed the kernels", .{}); + std.debug.print(" [graph] runSlots captured once: {d} nodes over the pool's fixed table addresses\n", .{try g.?.nodeCount()}); + } + try vp.sbuf.fill8(0xff, s.handle); + try exec.?.launchOn(s); + try pool.release(r, 0, s); + try s.synchronize(); + const yg = try vp.y(gpa, 1); + defer gpa.free(yg); + // eager over the same tables, then the fully resident layer + try vp.sbuf.fill8(0xff, s.handle); + try call(ex, s, l, vp, vp.picks.ptr, 1, false); + try call(ex, s, ld.res, vr, vr.picks.ptr, 1, false); + try s.synchronize(); + const ye = try vp.y(gpa, 1); + defer gpa.free(ye); + const yr = try vr.y(gpa, 1); + defer gpa.free(yr); + try check.sameBytes("graph replay vs eager", yg, ye); + try check.sameBytes("graph replay vs fully resident", yg, yr); + if (last) |b| { + if (!std.mem.eql(u8, b, yg)) distinct += 1; + gpa.free(b); + } + last = try gpa.dupe(u8, yg); + std.debug.print(" [graph] state {d}: {d} lines swapped by H2D before the replay ({d} held); replay == eager == resident, {d} slot rows byte for byte\n", .{ si, r.n_copies, r.held, slots }); + } + try check.expect(swaps > lines and distinct == states - 1, "the states did not exercise swaps ({d}) or distinct outputs ({d})", .{ swaps, distinct }); + check.pass("layer {d}: one captured runSlots graph replayed {d} states after {d} line swaps (no recapture), byte-identical to eager and to the resident layer", .{ ld.layer, states, swaps }); +} + +// ---------------------------------------------------------------------------------------------------- gate 4 + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// A row whose first `k` routed picks are not held by the pool and the rest are. +fn rowWithMisses(pool: *const fpool.Pool, rng: std.Random, k: usize, out: []i32) void { + var n: usize = 0; + // held experts first (up to 10 - k) + for (pool.slots_tab, pool.line_tab) |id, state| { + if (n == top_k - k) break; + if (state == .resident) { + out[k + n] = id; + n += 1; + } + } + var m: usize = 0; + while (m < k) { + const id: u16 = rng.intRangeLessThan(u16, 0, routed_n); + if (policy.lineOf(pool.v, pool.opt.layer, id) != null) continue; + if (std.mem.indexOfScalar(i32, out[0..m], id) != null) continue; + out[m] = id; + m += 1; + } + out[top_k] = 512; +} + +fn missCost(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, direct: bool, label: []const u8) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 64, .direct = direct, .double_buffer = false }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var prng = std.Random.DefaultPrng.init(0xbe7c4); + const rng = prng.random(); + var x: [dims]u16 = undefined; + for (&x) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + const wts: [slots]f32 = @splat(0.1); + var picks: [slots]i32 = undefined; + // warm: fill the pool with 64 experts + for (0..7) |_| { + rowWithMisses(&pool, rng, top_k, &picks); + try vp.set(&x, &picks, &wts); + _ = try poolRound(&pool, ex, s, vp, &picks, 1, false); + try s.synchronize(); + } + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + std.debug.print(" [{s}] one decode row (10 routed + shared), 64 lines, rounds of k misses (median of 9; idle card):\n", .{label}); + var base: f64 = 0; + for ([_]usize{ 0, 1, 2, 4, 10 }) |k| { + var wall: [9]f64 = undefined; + var copy: [9]f64 = undefined; + var comp: [9]f64 = undefined; + var stage_host: [9]f64 = undefined; + for (0..9) |it| { + rowWithMisses(&pool, rng, k, &picks); + try vp.set(&x, &picks, &wts); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, 1); + const t1 = check.now(gpu.io); + try c0.record(pool.lines.copy); + try pool.stage(r, 0, s); + try c1.record(pool.lines.copy); + const t2 = check.now(gpu.io); + try k0.record(s); + try call(ex, s, pool.tables(), vp, pool.picksFor(r, 0, vp.picks.ptr), 1, false); + try k1.record(s); + try pool.release(r, 0, s); + try s.synchronize(); + try pool.lines.copy.synchronize(); + const t3 = check.now(gpu.io); + wall[it] = ms(t0, t3); + copy[it] = try c0.elapsedMs(c1); + comp[it] = try k0.elapsedMs(k1); + stage_host[it] = ms(t1, t2); + } + const w = median(&wall); + if (k == 0) base = w; + std.debug.print(" k={d:>2}: round wall {d:.3} ms (+{d:.3} vs all-hit), copy stream {d:.3} ms ({d:.3} ms/expert), experts compute {d:.3} ms (measured after the stream's wait on the copies), host stage {d:.3} ms\n", .{ + k, w, w - base, median(©), if (k > 0) median(©) / @as(f64, @floatFromInt(k)) else 0, median(&comp), median(&stage_host), + }); + } + const ls = pool.lines.stats; + if (!direct) std.debug.print(" host staging memcpy (mmap -> pinned): {d:.3} ms/expert average over {d} copies ({d:.2} GB/s)\n", .{ @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(ls.copies)), ls.copies, @as(f64, @floatFromInt(ls.staged_bytes)) / @as(f64, @floatFromInt(ls.stage_ns)) }); +} + +/// A 16-row chunk through a 24-line pool (several passes): copies of pass j+1 overlap compute of pass j or not. +fn overlapChunk(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, serial: bool) !struct { wall: f64, copy: f64, comp: f64, passes: usize, copies: usize } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 24, .direct = true }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + const z = Zipf.init(0x0e1a9 + ld.layer, 0.8); + var prng = std.Random.DefaultPrng.init(0x0e1a9); + const rng = prng.random(); + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + genRound(&z, rng, max_rows, &x, &picks, &wts); + try vp.set(&x, &picks, &wts); + var cs: [64]cuda.Event = undefined; + for (&cs) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&cs) |*e| e.deinit(); + var ks: [64]cuda.Event = undefined; + for (&ks) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&ks) |*e| e.deinit(); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, max_rows); + if (r.n_passes > 32) return error.TooManyPasses; + const l = pool.tables(); + for (0..r.n_passes) |j| { + try cs[2 * j].record(pool.lines.copy); + try pool.stage(r, j, s); + try cs[2 * j + 1].record(pool.lines.copy); + if (serial) try pool.lines.copy.synchronize(); + try ks[2 * j].record(s); + try call(ex, s, l, vp, pool.picksFor(r, j, vp.picks.ptr), max_rows, false); + try ks[2 * j + 1].record(s); + try pool.release(r, j, s); + if (serial) try s.synchronize(); + } + try s.synchronize(); + try pool.lines.copy.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + var copy: f64 = 0; + var comp: f64 = 0; + for (0..r.n_passes) |j| { + copy += try cs[2 * j].elapsedMs(cs[2 * j + 1]); + comp += try ks[2 * j].elapsedMs(ks[2 * j + 1]); + } + return .{ .wall = wall, .copy = copy, .comp = comp, .passes = r.n_passes, .copies = r.n_copies }; +} + +/// 60 decode rounds (one row each, skewed picks known ahead), sequential (sync every round) vs pipelined (the host +/// stages round N+1 while round N computes), with and without double buffering. +fn decodePipe(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, lines: usize, pipelined: bool, double: bool) !struct { wall: f64, hit: f64, release_waits: u64, protected: u64 } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .direct = true, .double_buffer = double }, null); + defer pool.deinit(); + const rounds = 60; + var ios: [rounds]Io = undefined; + var made: usize = 0; + defer for (ios[0..made]) |*v| v.free(); + const z = Zipf.init(0xdec0 + ld.layer, 1.0); + var prng = std.Random.DefaultPrng.init(0xdec0); + const rng = prng.random(); + var picks: [rounds][slots]i32 = undefined; + var x: [dims]u16 = undefined; + var wts: [slots]f32 = undefined; + for (0..rounds) |i| { + ios[i] = try Io.init(gpu, pool.tables()); + made += 1; + genRound(&z, rng, 1, &x, &picks[i], &wts); + try ios[i].set(&x, &picks[i], &wts); + } + try s.synchronize(); + const t0 = check.now(gpu.io); + for (0..rounds) |i| { + _ = try poolRound(&pool, ex, s, ios[i], &picks[i], 1, false); + if (!pipelined) try s.synchronize(); + } + try s.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + const ps = pool.stats; + return .{ .wall = wall, .hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)), .release_waits = pool.lines.stats.release_waits, .protected = ps.protected }; +} + +fn benchRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const gpa = gpu.gpa; + // (a) miss cost, staged from the mapped checkpoint + try missCost(gpu, ex, s, ld, ld.source(), false, "staged: mmap -> pinned staging -> H2D"); + // the same experts in one pinned bank (what keeping the routed trellis in pinned host RAM gives): direct H2D + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try gpa.alloc([3][]const u8, routed_n); + defer gpa.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + std.debug.print(" pinned bank for layer {d}: {d:.0} MiB of routed trellis copied into page-locked host memory\n", .{ ld.layer, @as(f64, @floatFromInt(total)) / 1048576.0 }); + const psrc: fpool.Source = .{ .routed = pinned, .shared = ld.shared }; + try missCost(gpu, ex, s, ld, psrc, true, "pinned bank: direct H2D"); + + // (b) a 16-row chunk split into passes: serial vs overlapped copies + const ser = try overlapChunk(gpu, ex, s, ld, psrc, true); + const ovl = try overlapChunk(gpu, ex, s, ld, psrc, false); + const hidden_ms = (ser.wall - ovl.wall); + std.debug.print(" [16-row chunk, 24 lines, pinned bank] {d} passes, {d} copies: serial {d:.3} ms (copy {d:.3} + compute {d:.3}), overlapped {d:.3} ms; overlap hides {d:.3} ms = {d:.0}% of the smaller of copy/compute\n", .{ + ser.passes, ser.copies, ser.wall, ser.copy, ser.comp, ovl.wall, hidden_ms, 100.0 * hidden_ms / @min(ser.copy, ser.comp), + }); + // (c) decode rounds: pipelined staging, double buffering on/off + for ([_]usize{ 24, 48 }) |lines| { + const seq = try decodePipe(gpu, ex, s, ld, psrc, lines, false, true); + const pipe1 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, false); + const pipe2 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, true); + std.debug.print(" [60 decode rounds, {d} lines, pinned bank, hit {d:.3}] sequential {d:.2} ms; pipelined {d:.2} ms single-buffered ({d} copies waited for the previous round's compute), {d:.2} ms double-buffered ({d} waited, {d} rounds spared the previous round's lines)\n", .{ + lines, seq.hit, seq.wall, pipe1.wall, pipe1.release_waits, pipe2.wall, pipe2.release_waits, pipe2.protected, + }); + } + check.pass("layer {d}: miss cost, chunk overlap and decode pipelining measured (see lines above)", .{ld.layer}); +} + +// ---------------------------------------------------------------------------------------------------- hit rates + +const line_counts = [_]usize{ 10, 24, 48, 96, 160 }; + +const Sim = struct { + hit: u64 = 0, + miss: u64 = 0, + dec_hit: u64 = 0, + dec_miss: u64 = 0, + copies: u64 = 0, +}; + +/// Host-only replay of a pick sequence through one layer's policy (the same planner the pool runs). +/// picks [tokens][10] routed ids; the first `prompt` tokens go in 16-row rounds, the rest one row a round. +fn simulate(gpa: std.mem.Allocator, layer: u8, picks: []const u16, tokens: usize, prompt: usize, lines: usize, prefetch_lines: usize, counts: *[policy.shape.layers][policy.shape.experts]u32) !Sim { + const slots_tab = try gpa.alloc(u16, lines); + defer gpa.free(slots_tab); + const line_tab = try gpa.alloc(policy.Line, lines); + defer gpa.free(line_tab); + const stamp = try gpa.alloc(u32, lines); + defer gpa.free(stamp); + const lay = try gpa.alloc(u8, lines); + defer gpa.free(lay); + var v = policy.Vram.init(slots_tab, line_tab, stamp, lay); + var r = try fpool.Round.init(gpa, lines, max_rows * slots); + defer r.deinit(gpa); + var buf: [max_rows * slots]i32 = undefined; + var sim: Sim = .{}; + var prev: [routed_n]u16 = undefined; + var prev_n: usize = 0; + var t: usize = 0; + var round: u32 = 0; + while (t < tokens) { + const rows: usize = if (t < prompt) @min(max_rows, prompt - t) else 1; + for (0..rows) |i| { + for (0..top_k) |k| buf[i * slots + k] = picks[(t + i) * top_k + k]; + buf[i * slots + top_k] = 512; + } + round += 1; + try fpool.round.plan(&v, layer, buf[0 .. rows * slots], counts, round, round, .{ .prefetch_bytes = prefetch_lines, .line_bytes = 1, .protect = prev[0..prev_n] }, &r); + @memcpy(prev[0..r.n_ids], r.ids[0..r.n_ids]); + prev_n = r.n_ids; + // prefetch lines land before the next round (the device has a whole layer stack of time) + for (r.prefetch()) |c| policy.reapAt(&v, layer, c.id, round) catch {}; + sim.hit += r.hit_slots; + sim.miss += r.miss_slots; + sim.copies += r.n_copies; + if (t >= prompt) { + sim.dec_hit += r.hit_slots; + sim.dec_miss += r.miss_slots; + } + t += rows; + } + return sim; +} + +fn rate(h: u64, m: u64) f64 { + return if (h + m == 0) 0 else @as(f64, @floatFromInt(h)) / @as(f64, @floatFromInt(h + m)); +} + +fn hitRun(gpu: Gpu, files: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const counts = try gpa.create([policy.shape.layers][policy.shape.experts]u32); + defer gpa.destroy(counts); + // synthetic: 200 tokens of Zipf(1.1) picks, decode rows only + { + const z = Zipf.init(0x200, 1.1); + var prng = std.Random.DefaultPrng.init(0x200); + var row: [slots]i32 = undefined; + var picks: [200 * top_k]u16 = undefined; + for (0..200) |t| { + z.row(prng.random(), &row); + for (0..top_k) |k| picks[t * top_k + k] = @intCast(row[k]); + } + std.debug.print(" synthetic Zipf(1.1) over 512 experts, 200 decode tokens, one layer:\n", .{}); + for (line_counts) |lines| { + counts.* = @splat(@splat(0)); + const a = try simulate(gpa, 0, &picks, 200, 0, lines, 0, counts); + counts.* = @splat(@splat(0)); + const b = try simulate(gpa, 0, &picks, 200, 0, lines, 4, counts); + std.debug.print(" {d:>3} lines ({d:>5.1} MiB/layer at 1.76 MiB a line): hit {d:.3} (with 4-line prefetch every 4 rounds {d:.3})\n", .{ lines, @as(f64, @floatFromInt(lines)) * 1.7578, rate(a.hit, a.miss), rate(b.hit, b.miss) }); + } + } + for (files) |path| { + const blob = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpa, .limited(1 << 30)); + defer gpa.free(blob); + const hdr = std.mem.bytesAsSlice(u32, blob[0..12]); + const tokens: usize = hdr[0]; + const prompt: usize = hdr[1]; + const layers: usize = hdr[2]; + const raw = try gpa.alloc(u16, tokens * layers * top_k); + defer gpa.free(raw); + @memcpy(std.mem.sliceAsBytes(raw), blob[12..][0 .. raw.len * 2]); + std.debug.print(" real router picks {s}: {d} tokens ({d} prompt as 16-row chunks, {d} decode rows) x {d} layers\n", .{ path, tokens, prompt, tokens - prompt, layers }); + const per = try gpa.alloc(u16, tokens * top_k); + defer gpa.free(per); + // distinct experts a layer over the decode part (an upper bound on what any pool must hold) + var distinct_sum: usize = 0; + for (0..layers) |l| { + var seen: [routed_n]bool = @splat(false); + for (prompt..tokens) |t| for (0..top_k) |k| { + seen[raw[(t * layers + l) * top_k + k]] = true; + }; + for (seen) |b| distinct_sum += @intFromBool(b); + } + std.debug.print(" distinct experts a layer over the {d} decode tokens: {d:.1} on average\n", .{ tokens - prompt, @as(f64, @floatFromInt(distinct_sum)) / @as(f64, @floatFromInt(layers)) }); + for (line_counts) |lines| { + for ([_]usize{ 0, 4 }) |pf| { + counts.* = @splat(@splat(0)); + var tot: Sim = .{}; + var lo: f64 = 1; + var hi: f64 = 0; + var copies: u64 = 0; + for (0..layers) |l| { + for (0..tokens) |t| @memcpy(per[t * top_k ..][0..top_k], raw[(t * layers + l) * top_k ..][0..top_k]); + const s = try simulate(gpa, @intCast(l), per, tokens, prompt, lines, pf, counts); + tot.hit += s.hit; + tot.miss += s.miss; + tot.dec_hit += s.dec_hit; + tot.dec_miss += s.dec_miss; + copies += s.copies; + const dr = rate(s.dec_hit, s.dec_miss); + lo = @min(lo, dr); + hi = @max(hi, dr); + } + const dec_tokens = tokens - prompt; + const miss_per_tok = @as(f64, @floatFromInt(tot.dec_miss)) / @as(f64, @floatFromInt(dec_tokens)); + std.debug.print(" {d:>3} lines{s}: decode hit {d:.3} (layers {d:.3}..{d:.3}), {d:.1} routed misses a decode token over 48 layers; all rows hit {d:.3}; pool {d:.2} GiB for 48 layers\n", .{ + lines, if (pf > 0) " +prefetch" else " ", rate(tot.dec_hit, tot.dec_miss), lo, hi, miss_per_tok, rate(tot.hit, tot.miss), @as(f64, @floatFromInt(lines * 48)) * 1.7578 / 1024.0, + }); + } + } + } + check.pass("hit rates printed (host simulation of the pool's own planner)", .{}); +} + +// ---------------------------------------------------------------------------------------------------- driver + +fn parseLayers(spec: []const u8, out: []usize) !usize { + var n: usize = 0; + var it = std.mem.splitScalar(u8, spec, ','); + while (it.next()) |tok| { + if (n == out.len) return error.TooManyLayers; + out[n] = try std.fmt.parseInt(usize, tok, 10); + n += 1; + } + return n; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var mode: []const u8 = "all"; + var layer_spec: []const u8 = "0,17"; + if (rest.len > 1) { + if (std.ascii.isDigit(rest[1][0])) layer_spec = rest[1] else mode = rest[1]; + } + if (std.mem.eql(u8, mode, "hit")) return hitRun(gpu, rest[2..]); + if (rest.len > 2) layer_spec = rest[2]; + var layers: [8]usize = undefined; + const n_layers = try parseLayers(layer_spec, &layers); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + for (layers[0..n_layers]) |layer| { + std.debug.print("== layer {d}\n", .{layer}); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + const all = std.mem.eql(u8, mode, "all"); + if (all or std.mem.eql(u8, mode, "equal")) { + try equalRun(gpu, &ex, s, ld, 10, 0, 0x10 + layer); + try equalRun(gpu, &ex, s, ld, 24, 4, 0x24 + layer); + try equalRun(gpu, &ex, s, ld, 64, 0, 0x64 + layer); + } + if (all or std.mem.eql(u8, mode, "graph")) try graphRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "bench")) try benchRun(gpu, &ex, s, ld); + } +} diff --git a/zig/tests/cuda/exl3_prefill_origin_tests.zig b/zig/tests/cuda/exl3_prefill_origin_tests.zig new file mode 100644 index 0000000000..b24e9a054d --- /dev/null +++ b/zig/tests/cuda/exl3_prefill_origin_tests.zig @@ -0,0 +1,164 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const graph = flash.exl3.graph; +const paging = flash.exl3.kvpage; +const prefill = flash.exl3.prefill; +const step_rows = 256; +const steps = 24; + +fn require(name: [:0]const u8, expected: []const u8) !void { + const setting = std.c.getenv(name); + try check.expect(setting != null and std.mem.eql(u8, std.mem.span(setting.?), expected), "prefill-origin fixture requires {s}={s}", .{ name, expected }); +} + +fn stateBytes(sequence: *const flash.prompt.Seq) usize { + var bytes = sequence.logits.len + sequence.tail.len; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv, sequence.ikc, sequence.pooled }) |buffers| for (buffers) |buffer| { + bytes += buffer.len; + }; + const pager = sequence.paged.?; + const row_bytes = pager.hostMut(0, false).len / sequence.cap; + return bytes + 2 * sequence.ikc.len * sequence.fed * row_bytes; +} + +fn readState(sequence: *flash.prompt.Seq, stream: cuda.Stream, target: []u8) !void { + const pager = sequence.paged.?; + try pager.flush(stream, sequence.fed); + try stream.synchronize(); + var offset: usize = 0; + @memcpy(target[0..sequence.logits.len], sequence.logits); + offset += sequence.logits.len; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv, sequence.ikc, sequence.pooled }) |buffers| for (buffers) |buffer| { + try buffer.download(0, target[offset..][0..buffer.len]); + offset += buffer.len; + }; + try sequence.tail.download(0, target[offset..][0..sequence.tail.len]); + offset += sequence.tail.len; + for (0..sequence.ikc.len) |layer| for ([_]bool{ false, true }) |values| { + const canonical = pager.hostMut(layer, values); + const valid_bytes = sequence.fed * (canonical.len / sequence.cap); + @memcpy(target[offset..][0..valid_bytes], canonical[0..valid_bytes]); + offset += valid_bytes; + }; + try check.expect(offset == target.len, "prefill-origin state read covers every allocated comparison byte", .{}); +} + +fn eagerForward(engine: anytype, sequence: *flash.prompt.Seq, ids: []const u32, comptime operations: anytype) !void { + std.mem.swap(flash.prompt.Seq, &engine.seq, sequence); + defer std.mem.swap(flash.prompt.Seq, &engine.seq, sequence); + graph.force = .off; + try poisonScratch(engine); + try operations.forward(engine, ids, .{ .logits = .last }); +} + +fn poisonScratch(engine: anytype) !void { + try engine.e.scores.?.fill32(0x7fc00000, engine.stream.handle); + const scratch = engine.e.batch_s.?; + for (scratch.bufs.items) |buffer| { + if (buffer.ptr == scratch.po or buffer.ptr == scratch.pm or buffer.ptr == scratch.pl) try buffer.fill32(0x7fc00000, engine.stream.handle); + } +} + +fn finiteLogits(bytes: []const u8) bool { + var offset: usize = 0; + while (offset < bytes.len) : (offset += 2) { + if (std.mem.readInt(u16, bytes[offset..][0..2], .little) & 0x7f80 == 0x7f80) return false; + } + return true; +} + +pub fn run(gpu: check.Gpu, rank: usize, world: usize, port: u16, directory: []const u8, storage: []const u8, comptime operations: anytype) !void { + try require("TF_FLASHNEXT_EXL3_GRAPH_PREFILL", "1"); + try require("TF_FLASHNEXT_EXL3_GRAPH_PREFILL_AFTER", "1"); + try require("TF_FLASHNEXT_CHUNK", "256"); + try check.expect(world == 2 and !paging.force_off, "prefill-origin fixture requires TP2 and BF16 RAM paging", .{}); + const shared = std.mem.eql(u8, storage, "shared"); + try check.expect(shared or std.mem.eql(u8, storage, "private"), "prefill-origin storage must be private or shared", .{}); + const config: paging.Config = .{ .window = 2052, .rows = step_rows, .slots = 512, .shared_lanes = if (shared) 2 else 0 }; + const previous_paging = paging.force; + paging.force = config; + defer paging.force = previous_paging; + const previous_options = prefill.options; + prefill.options = .{ .rows = step_rows, .serve_rows = step_rows, .min_rows = step_rows, .ring = false, .prefetch = false, .tiles = false }; + defer prefill.options = previous_options; + const previous_mode = graph.force; + defer graph.force = previous_mode; + const engine = try operations.init(gpu, rank, world, port, directory, config); + defer operations.deinit(engine); + var eager = try engine.e.newSeq(engine.seq.cap); + defer eager.deinit(gpu.gpa); + try operations.reset(engine); + std.mem.swap(flash.prompt.Seq, &engine.seq, &eager); + operations.reset(engine) catch |failure| { + std.mem.swap(flash.prompt.Seq, &engine.seq, &eager); + return failure; + }; + std.mem.swap(flash.prompt.Seq, &engine.seq, &eager); + const pager = engine.seq.paged orelse return error.PrefillOriginPagingRequired; + try check.expect(engine.seq.kv8 == null and engine.e.chunk == step_rows and pager.cfg.rows == step_rows and pager.cfg.shared_lanes == config.shared_lanes, "prefill-origin fixture uses the exact native BF16 private/shared geometry", .{}); + const generation: ?paging.shared_contract.Sequence = if (shared) pager.shared_sequence else null; + graph.clearPrefill(); + var tokens: [steps * step_rows]u32 = undefined; + var generator = std.Random.DefaultPrng.init(0x70726566696c6c); + for (&tokens) |*token| token.* = 200 + generator.random().uintLessThan(u32, 60000); + var observed_chunk_grid = false; + var observed_score_grid = false; + var observed_origin = false; + var observed_sparse_decline = false; + var observed_split_decline = false; + const initial = graph.prefill_stats; + for (0..steps) |step| { + const first = step * step_rows; + const ids = tokens[first..][0..step_rows]; + try eagerForward(engine, &eager, ids, operations); + const reference = try gpu.gpa.alloc(u8, stateBytes(&eager)); + defer gpu.gpa.free(reference); + try readState(&eager, engine.stream, reference); + const before = graph.prefill_stats; + const before_deep = prefill.stats; + const old_origin = pager.win0[0]; + const old_page_rows = try paging.pagingRows(pager.cfg, old_origin, step_rows); + graph.force = .graph; + try poisonScratch(engine); + try operations.forward(engine, ids, .{ .logits = .last }); + const after = graph.prefill_stats; + const rebound = after.rebinds - before.rebinds; + const new_page_rows = try paging.pagingRows(pager.cfg, pager.win0[0], step_rows); + const actual = try gpu.gpa.alloc(u8, stateBytes(&engine.seq)); + defer gpu.gpa.free(actual); + try readState(&engine.seq, engine.stream, actual); + const exact = finiteLogits(engine.seq.logits) and std.mem.eql(u8, reference, actual) and eager.fed == engine.seq.fed and + std.meta.eql(eager.history, engine.seq.history) and std.meta.eql(eager.paged.?.win0, pager.win0); + const peer_exact = try operations.swap(engine, @intFromBool(exact)); + if (!exact) try check.sameBytes("prefill-origin logits, recurrent/conv/PLE, indexer and canonical BF16 K/V", actual, reference); + try check.expect(exact and peer_exact == 1, "rank {d}: {s} step {d} exact eager/rebound logits and state", .{ rank, storage, step }); + try check.expect(prefill.stats.chunks == before_deep.chunks + 1 and prefill.stats.rows == before_deep.rows + step_rows, "rank {d}: step {d} uses one actual deep-prefill chunk", .{ rank, step }); + const end = first + step_rows; + const sparse_transition = first / 4 <= 512 and end / 4 > 512; + const split_transition = old_page_rows != new_page_rows; + if (sparse_transition or split_transition) { + const declined = rebound == 0 and after.captures > before.captures; + const peer_declined = try operations.swap(engine, @intFromBool(declined)); + try check.expect(declined and peer_declined == 1, "rank {d}: step {d} topology transition recaptures with zero origin rebinds", .{ rank, step }); + observed_sparse_decline = observed_sparse_decline or sparse_transition; + observed_split_decline = observed_split_decline or split_transition; + } else if (step > 0 and rebound > 0) { + const old_chunks = (@min(first, 2051) + 511) / 512; + const new_chunks = (@min(end, 2051) + 511) / 512; + observed_chunk_grid = observed_chunk_grid or old_chunks != new_chunks; + observed_score_grid = observed_score_grid or (first > 2048 and (first + 255) / 256 != (end + 255) / 256); + observed_origin = observed_origin or old_origin != pager.win0[0]; + } + if (generation) |owner| try check.expect(std.meta.eql(owner, pager.shared_sequence), "rank {d}: origin replay preserves the exact paging generation/incarnation", .{rank}); + const status = try pager.total(); + try check.expect(!status.overflow, "rank {d}: step {d} paging stays within its slot capacity", .{ rank, step }); + std.debug.print("PREFILL-ORIGIN rank={d}/{d} store={s} step={d} end={d} origin={d} page_rows={d} rebinds={d} captures={d} exact=true\n", .{ rank, world, storage, step, end, pager.win0[0], new_page_rows, rebound, after.captures - before.captures }); + } + const covered = observed_chunk_grid and observed_score_grid and observed_origin and observed_sparse_decline and observed_split_decline; + const peer_covered = try operations.swap(engine, @intFromBool(covered)); + try check.expect(covered and peer_covered == 1, "rank {d}: real prefill covers chunk/score grids, sliding origins and both topology declines", .{rank}); + check.pass("rank {d}/{d}: native BF16 {s} prefill exact eager/rebound logits and state, {d} actual origin rebinds, both topology declines", .{ rank, world, storage, graph.prefill_stats.rebinds - initial.rebinds }); +} diff --git a/zig/tests/cuda/exl3_prefill_tests.zig b/zig/tests/cuda/exl3_prefill_tests.zig new file mode 100644 index 0000000000..1e0875a91c --- /dev/null +++ b/zig/tests/cuda/exl3_prefill_tests.zig @@ -0,0 +1,1080 @@ +//! exl3-prefill: the pieces of a Flash Next EXL3 prefill, timed on the pack's real tensors (one idle card). +//! +//! exl3-prefill DIR experts [LAYER] upstream's grouped experts (Experts.runSlots) on one fully resident layer: +//! ms a call at 16..1408 rows (Zipf and uniform picks), then one 1024-row block +//! split into expert-major passes over masked picks (bytes equal to one call) +//! exl3-prefill DIR copy [LAYER] expert lines from a pinned bank: Lines.load (copy + table writes) vs one bare +//! H2D of the same bytes, and copies running beside a 1024-row experts call +//! exl3-prefill DIR prompt [LAYER] the prefill experts kernel (flashnext_exl3_prefill.cu, 64-member tiles) vs +//! runSlots: y against y at 16..1408 rows, ms at 16..8192 rows, expert-major passes +//! exl3-prefill DIR engine IDS N,N.. [LINES] [cmp] the family's Engine with the offload pool (as serve attaches it): +//! prefill rows/s of the first N ids of IDS (u32 little-endian token ids, real text) +//! through the expert-major deep chunks; `cmp` also runs the batched pass on the +//! same prompt and compares the last row's logits +//! exl3-prefill DIR agree IDS FIRST STEP COUNT [LINES] last-row logits of COUNT prompt lengths FIRST, FIRST+STEP..: +//! the batched pass vs the deep path (env options as set) and vs the text's own +//! next token (top-1 agreement, next-token accuracy, |d| on the batched top-5) +//! exl3-prefill DIR dense exl3.Dense prompt rows: reference 128-row blocks vs #806 tiles, 16..2048 rows, +//! each against the other and rows 0-1 against an exact float64 product +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const flash = @import("flashnext"); +const pk_ = flash.exl3.prefill_k; + +const st = core.safetensors; +const exl3fmt = core.exl3_format; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const top_k = 10; +const slots = 11; +const shard_count = 7; +const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; +/// The group kernel stages a call's picks in 64 KiB of shared memory: rows * 11 * 4 bytes. +const max_call_rows: usize = 1408; + +const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +const Timer = struct { + a: cuda.Event, + b: cuda.Event, + + fn init(d: *const cuda.Driver) !Timer { + return .{ .a = try cuda.Event.init(d, true), .b = try cuda.Event.init(d, true) }; + } + fn deinit(t: *Timer) void { + t.a.deinit(); + t.b.deinit(); + } + /// Best of 5 batches of `reps` calls, ms a call (after 2 warm-up calls). + fn best(t: *Timer, s: cuda.Stream, reps: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..2) |_| try f(ctx); + try s.synchronize(); + var b: f64 = std.math.inf(f64); + for (0..5) |_| { + try t.a.record(s); + for (0..reps) |_| try f(ctx); + try t.b.record(s); + try t.b.synchronize(); + b = @min(b, @as(f64, try t.a.elapsedMs(t.b)) / @as(f64, @floatFromInt(reps))); + } + return b; + } +}; + +// ---------------------------------------------------------------------------------------------------- one layer + +/// One layer's 513 experts: the mmapped trellis, the stacked scales and every expert resident on the device. +const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + line_bytes: usize, + + fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + ld.line_bytes = 0; + for (0..table_n) |e| { + var line: usize = 0; + for (projs, 0..) |pn, p| { + const base = "model.language_model.layers.{d}.mlp."; + const tr = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.trellis", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.suh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.svh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + line += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + if (e < routed_n) ld.line_bytes = @max(ld.line_bytes, line); + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts resident ({d:.0} MiB), routed k2 gate/up/down {d}/{d}/{d}, a routed line {d:.3} MiB\n", .{ layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], @as(f64, @floatFromInt(ld.line_bytes)) / 1048576.0 }); + return ld; + } + + fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } +}; + +/// Picks for `rows` rows: ten distinct routed ids then the shared expert; Zipf(s) over a fixed permutation (s 0: uniform). +fn genPicks(rng: std.Random, s: f64, rows: usize, out: []i32) void { + var perm: [routed_n]u16 = undefined; + for (&perm, 0..) |*p, i| p.* = @intCast(i); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + prng.random().shuffle(u16, &perm); + var cdf: [routed_n]f64 = undefined; + var sum: f64 = 0; + for (&cdf, 0..) |*c, i| { + sum += if (s == 0) 1.0 else 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&cdf) |*c| c.* /= sum; + for (0..rows) |r| { + const row = out[r * slots ..][0..slots]; + var n: usize = 0; + while (n < top_k) { + const u = rng.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and cdf[i] < u) i += 1; + const id: i32 = perm[i]; + if (std.mem.indexOfScalar(i32, row[0..n], id) != null) continue; + row[n] = id; + n += 1; + } + row[top_k] = 512; + } +} + +const SlotsCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + picks: u64, + rows: usize, + sc: exl3e.Scratch, + fn go(c: SlotsCall) anyerror!void { + try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.picks, c.rows, c.sc); + } +}; + +/// Expert-major passes over one block: pass j runs the block's picks of experts [j*per, (j+1)*per) (others -1). +const PassesCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + masked: u64, // int32 [passes, rows * slots] + rows: usize, + passes: usize, + sc: exl3e.Scratch, + fn go(c: PassesCall) anyerror!void { + for (0..c.passes) |j| try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.masked + j * c.rows * slots * 4, c.rows, c.sc); + } +}; + +fn expertsRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + var t = try Timer.init(gpu.d); + defer t.deinit(); + const max_rows = max_call_rows; + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + const xh = try a.alloc(u16, max_rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, max_rows * slots); + defer a.free(picks); + var pk = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, max_rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, max_rows, slots); + std.debug.print(" runSlots, one resident layer (10 routed + shared a row), ms a call (best of 5 x 5):\n", .{}); + for ([_]f64{ 0.8, 0 }) |zs| { + genPicks(rng, zs, max_rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks)); + for ([_]usize{ 16, 64, 256, 512, 1024, max_rows }) |rows| { + var seen: [table_n]bool = @splat(false); + var uniq: usize = 0; + for (picks[0 .. rows * slots]) |p| if (!seen[@intCast(p)]) { + seen[@intCast(p)] = true; + uniq += 1; + }; + const ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" {s:<8} rows {d:5}: {d:8.3} ms {d:7.2} us/row {d:6.0} rows/s/layer ({d} experts incl. shared, {d:.1} member rows each)\n", .{ if (zs == 0) "uniform" else "zipf0.8", rows, ms, 1e3 * ms / @as(f64, @floatFromInt(rows)), @as(f64, @floatFromInt(rows)) / ms * 1e3, uniq, @as(f64, @floatFromInt(rows * slots)) / @as(f64, @floatFromInt(uniq)) }); + } + } + // expert-major passes over one 1024-row block (uniform picks): bytes equal to one call, and the time + const rows: usize = 1024; + genPicks(rng, 0, rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks[0 .. rows * slots])); + const one_bytes = try a.alloc(u8, rows * slots * dims * 4); + defer a.free(one_bytes); + try (SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }).go(); + try s.synchronize(); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(one_bytes.ptr, sc.y, one_bytes.len), "cuMemcpyDtoH"); + const base_ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" expert-major passes over one {d}-row block (uniform picks; pass j: experts [j*512/P, (j+1)*512/P) plus the shared one in pass 0):\n", .{rows}); + for ([_]usize{ 1, 2, 4, 8, 16, 32 }) |passes| { + const per = routed_n / passes; + const host = try a.alloc(i32, passes * rows * slots); + defer a.free(host); + for (0..passes) |j| for (picks[0 .. rows * slots], host[j * rows * slots ..][0 .. rows * slots]) |p, *o| { + const in = if (p == 512) j == 0 else @as(usize, @intCast(p)) / per == j; + o.* = if (in) p else -1; + }; + var masked = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer masked.free(); + try gpu.d.check(gpu.d.api.cuMemsetD8_v2(sc.y, 0xff, rows * slots * dims * 4), "cuMemsetD8"); + const call: PassesCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .masked = masked.ptr, .rows = rows, .passes = passes, .sc = sc }; + try call.go(); + try s.synchronize(); + const got = try a.alloc(u8, one_bytes.len); + defer a.free(got); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(got.ptr, sc.y, got.len), "cuMemcpyDtoH"); + const same = std.mem.eql(u8, got, one_bytes); + const ms = try t.best(s, 3, call, PassesCall.go); + std.debug.print(" {d:2} passes of {d:3} experts: {d:8.3} ms (+{d:.3} ms vs one call, {d:.3} ms a pass) y {s}\n", .{ passes, per, ms, ms - base_ms, ms / @as(f64, @floatFromInt(passes)), if (same) "bit-identical to one call" else "DIFFERS" }); + try check.expect(same, "masked passes reproduce one call's bytes", .{}); + } + check.pass("layer {d}: runSlots timed at 16..{d} rows; expert-major masked passes are bit-identical to one call", .{ ld.layer, max_rows }); +} + +// ---------------------------------------------------------------------------------------------------- prompt kernel + +const PromptCall = struct { + k: *const pk_.Kernels, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + picks: u64, + n_members: usize, + n_items: usize, + dv: pk_.Device, + y: u64, + fn go(c: PromptCall) anyerror!void { + try pk_.run(c.k, c.s, c.l, c.x, dims, c.picks, slots, c.n_members, c.n_items, c.dv, c.y); + } +}; + +/// P expert-major passes, each its own uploaded plan (members, items) over the same picks. +const PromptPasses = struct { + calls: []const PromptCall, + fn go(c: PromptPasses) anyerror!void { + for (c.calls) |call| try call.go(); + } +}; + +fn promptRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + var k = try pk_.Kernels.load(gpu.d); + defer k.unload(); + var t = try Timer.init(gpu.d); + defer t.deinit(); + const max_rows: usize = 8192; + const max_pairs = max_rows * slots; + var prng = std.Random.DefaultPrng.init(0x9a11); + const rng = prng.random(); + const xh = try a.alloc(u16, max_rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, max_pairs); + defer a.free(picks); + genPicks(rng, 0.8, max_rows, picks); + var pkd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer pkd.free(); + var plan = try pk_.Plan.init(a, max_pairs, table_n); + defer plan.deinit(a); + var counts: [table_n]u32 = undefined; + var mem_d = try cuda.DeviceBuffer.alloc(gpu.d, max_pairs * 4); + defer mem_d.free(); + var items_d = try cuda.DeviceBuffer.alloc(gpu.d, 3 * pk_.maxItems(max_pairs, table_n) * 4); + defer items_d.free(); + var scr = try cuda.DeviceBuffer.alloc(gpu.d, pk_.scratchBytes(max_pairs, dims, width)); + defer scr.free(); + var yk = try cuda.DeviceBuffer.alloc(gpu.d, max_pairs * dims * 4); + defer yk.free(); + // runSlots' reference (at most max_call_rows rows a call) + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, max_call_rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, max_call_rows, slots); + std.debug.print(" prefill experts kernel (64-member tiles, compact pairs) vs runSlots, one resident layer, zipf0.8 picks:\n", .{}); + for ([_]usize{ 16, 64, 256, 1024, 1408, 2048, 4096, 8192 }) |rows| { + const pairs = rows * slots; + pk_.build(picks[0..pairs], null, &counts, &plan); + try mem_d.upload(0, plan.memberBytes()); + try items_d.upload(0, plan.itemBytes()); + const dv: pk_.Device = .{ .members = mem_d.ptr, .items = items_d.ptr, .scratch = scr.ptr }; + const call: PromptCall = .{ .k = &k, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .n_members = plan.n_members, .n_items = plan.n_items, .dv = dv, .y = yk.ptr }; + try call.go(); + try s.synchronize(); + const ms_k = try t.best(s, 3, call, PromptCall.go); + var ms_r: f64 = std.math.nan(f64); + var rel: f64 = std.math.nan(f64); + var same_words: usize = 0; + if (rows <= max_call_rows) { + const rc: SlotsCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .rows = rows, .sc = sc }; + try rc.go(); + try s.synchronize(); + const want = try a.alloc(f32, pairs * dims); + defer a.free(want); + const got = try a.alloc(f32, pairs * dims); + defer a.free(got); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(want.ptr, sc.y, want.len * 4), "cuMemcpyDtoH"); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(got.ptr, yk.ptr, got.len * 4), "cuMemcpyDtoH"); + var dmax: f64 = 0; + var peak: f64 = 0; + for (got, want) |g, w| { + dmax = @max(dmax, @abs(@as(f64, g) - w)); + peak = @max(peak, @abs(@as(f64, w))); + same_words += @intFromBool(@as(u32, @bitCast(g)) == @as(u32, @bitCast(w))); + } + rel = dmax / peak; + ms_r = try t.best(s, 3, rc, SlotsCall.go); + try check.expect(rel < 1e-3, "prefill kernel within fp32 sum-order noise of runSlots ({e:.2})", .{rel}); + } + std.debug.print(" rows {d:5}: prefill kernel {d:8.3} ms ({d:6.2} us/row, {d} items) runSlots {d:8.3} ms ({d:6.2} us/row) speedup {d:.2}x max|dy|/peak {e:.2}, {d:.1}% of words bit-equal\n", .{ rows, ms_k, 1e3 * ms_k / @as(f64, @floatFromInt(rows)), plan.n_items, ms_r, 1e3 * ms_r / @as(f64, @floatFromInt(rows)), ms_r / ms_k, rel, 100 * @as(f64, @floatFromInt(same_words)) / @as(f64, @floatFromInt(pairs * dims)) }); + } + // expert-major passes over 8192 rows: the plan of each pass alone, back to back + const rows = max_rows; + const pairs = rows * slots; + std.debug.print(" expert-major passes over {d} rows (pass j: experts [j*512/P, (j+1)*512/P), the shared one in pass 0):\n", .{rows}); + for ([_]usize{ 1, 8, 16, 32 }) |passes| { + const per = routed_n / passes; + var calls: [32]PromptCall = undefined; + var mems: [32]cuda.DeviceBuffer = undefined; + var its: [32]cuda.DeviceBuffer = undefined; + var made: usize = 0; + defer for (mems[0..made], its[0..made]) |*m, *i| { + m.free(); + i.free(); + }; + for (0..passes) |j| { + var take: [table_n]bool = @splat(false); + for (j * per..(j + 1) * per) |id| take[id] = true; + take[512] = j == 0; + pk_.build(picks[0..pairs], &take, &counts, &plan); + mems[j] = try cuda.DeviceBuffer.fromHost(gpu.d, plan.memberBytes()); + its[j] = try cuda.DeviceBuffer.fromHost(gpu.d, plan.itemBytes()); + made += 1; + calls[j] = .{ .k = &k, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .n_members = plan.n_members, .n_items = plan.n_items, .dv = .{ .members = mems[j].ptr, .items = its[j].ptr, .scratch = scr.ptr }, .y = yk.ptr }; + } + const ms = try t.best(s, 2, PromptPasses{ .calls = calls[0..passes] }, PromptPasses.go); + std.debug.print(" {d:2} passes: {d:8.3} ms ({d:6.2} us/row)\n", .{ passes, ms, 1e3 * ms / @as(f64, @floatFromInt(rows)) }); + } + check.pass("layer {d}: prefill experts kernel within fp32 order noise of runSlots, timed to {d} rows", .{ ld.layer, max_rows }); +} + +// ---------------------------------------------------------------------------------------------------- copies + +fn copyRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + // the layer's routed trellis in one pinned bank (what S9 keeps), lines laid out as Lines does + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try a.alloc([3][]const u8, routed_n); + defer a.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + const n_lines: usize = 64; + var bytes: [3]usize = .{ 0, 0, 0 }; + var k2: [routed_n][3]i32 = undefined; + const tiles = (dims / 16) * (width / 16); + for (pinned, &k2) |e, *w| for (0..3) |p| { + bytes[p] = @max(bytes[p], e[p].len); + w[p] = @intCast(e[p].len / (tiles * 16)); + }; + var lines = try cuda.exl3_pool.Lines.init(a, gpu.d, .{ .experts = table_n, .bytes = bytes, .lines = n_lines, .staging = 0 }); + defer lines.deinit(); + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var raw = try cuda.DeviceBuffer.alloc(gpu.d, n_lines * lines.layout.line_bytes); + defer raw.free(); + const moved = blk: { + var m: usize = 0; + for (pinned[0..n_lines]) |e| m += e[0].len + e[1].len + e[2].len; + break :blk m; + }; + std.debug.print(" {d} expert lines ({d:.2} MiB, {d:.3} MiB a line) from the pinned bank, best of 5:\n", .{ n_lines, @as(f64, @floatFromInt(moved)) / 1048576.0, @as(f64, @floatFromInt(moved)) / 1048576.0 / n_lines }); + var best_load: f64 = std.math.inf(f64); + var best_raw: f64 = std.math.inf(f64); + var best_raw1: f64 = std.math.inf(f64); + for (0..5) |it| { + try lines.drain(); + try c0.record(lines.copy); + for (0..n_lines) |i| { + const id = (it * n_lines + i) % routed_n; + const prev = (id + routed_n - n_lines) % routed_n; + try lines.load(i, id, pinned[id], k2[id], if (it == 0) null else prev, true, gpu.io); + } + try c1.record(lines.copy); + try c1.synchronize(); + best_load = @min(best_load, try c0.elapsedMs(c1)); + // the same bytes as three H2D a line, no table writes + try c0.record(lines.copy); + for (0..n_lines) |i| for (0..3) |p| { + const e = pinned[i]; + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr + i * lines.layout.line_bytes + lines.layout.off[p], e[p].ptr, e[p].len, lines.copy.handle), "cuMemcpyHtoDAsync"); + }; + try c1.record(lines.copy); + try c1.synchronize(); + best_raw = @min(best_raw, try c0.elapsedMs(c1)); + // one H2D of the bank's first n_lines experts (contiguous in the bank) + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + best_raw1 = @min(best_raw1, try c0.elapsedMs(c1)); + } + const gbs = struct { + fn f(b: usize, ms: f64) f64 { + return @as(f64, @floatFromInt(b)) / ms / 1e6; + } + }.f; + std.debug.print(" Lines.load (3 H2D + 6-9 table memsets + event a line): {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_load, best_load / n_lines, gbs(moved, best_load) }); + std.debug.print(" 3 bare H2D a line: {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_raw, best_raw / n_lines, gbs(moved, best_raw) }); + std.debug.print(" one H2D of the same bytes: {d:.3} ms, {d:.2} GB/s\n", .{ best_raw1, gbs(moved, best_raw1) }); + + // copies beside compute: a 1024-row experts call on the resident layer while the copy stream moves 64 lines + const rows: usize = 1024; + var prng = std.Random.DefaultPrng.init(0xc0b1); + const rng = prng.random(); + const xh = try a.alloc(u16, rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + genPicks(rng, 0, rows, picks); + var pk = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, rows, slots); + const call: SlotsCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }; + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + var alone_c: f64 = std.math.inf(f64); + var alone_k: f64 = std.math.inf(f64); + var both_c: f64 = std.math.inf(f64); + var both_k: f64 = std.math.inf(f64); + var both_w: f64 = std.math.inf(f64); + for (0..5) |_| { + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + alone_k = @min(alone_k, try k0.elapsedMs(k1)); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + alone_c = @min(alone_c, try c0.elapsedMs(c1)); + // both at once + try s.synchronize(); + const t0 = check.now(gpu.io); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + try c1.synchronize(); + both_w = @min(both_w, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6); + both_c = @min(both_c, try c0.elapsedMs(c1)); + both_k = @min(both_k, try k0.elapsedMs(k1)); + } + std.debug.print(" overlap: {d}-row experts call alone {d:.3} ms, 64-line H2D alone {d:.3} ms; together: experts {d:.3} ms, H2D {d:.3} ms, wall {d:.3} ms (serial would be {d:.3})\n", .{ rows, alone_k, alone_c, both_k, both_c, both_w, alone_k + alone_c }); + check.pass("layer {d}: copy cost and copy/compute overlap measured", .{ld.layer}); +} + +// ---------------------------------------------------------------------------------------------------- dense + +const DenseCall = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + fn go(c: DenseCall) anyerror!void { + try c.de.prompt(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +/// y [rows, N] = x [rows, K] . diag(suh) H Wq H diag(svh) in float64. +fn exact(a: std.mem.Allocator, trellis: []const i16, suh: []const f16, svh: []const f16, tk: usize, tn: usize, halves: u16, x: []const f64, rows: usize, y: []f64) !void { + const k = tk * 16; + const n = tn * 16; + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = halves }, .mul1, trellis, tk, tn, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(svh[c])); + }; +} + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +fn denseRun(gpu: Gpu, pack: *const Pack, s: cuda.Stream) !void { + const a = gpu.gpa; + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var t = try Timer.init(gpu.d); + defer t.deinit(); + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + }; + const row_counts = [_]usize{ 16, 128, 256, 1024, 2048 }; + const max_rows: usize = 2048; + var prng = std.Random.DefaultPrng.init(0xde45e); + const rng = prng.random(); + var sum_ref: [row_counts.len]f64 = @splat(0); + var sum_tiles: [row_counts.len]f64 = @splat(0); + for (names) |name| { + const tr = try pack.get("{s}.trellis", .{name}); + const su = try pack.get("{s}.suh", .{name}); + const sv = try pack.get("{s}.svh", .{name}); + const tk = tr.dim(0); + const tn = tr.dim(1); + const halves: u16 = @intCast(tr.dim(2) / 8); + const k = tk * 16; + const n = tn * 16; + const trellis = try copyAs(i16, a, tr.bytes); + defer a.free(trellis); + const suh = try copyAs(f16, a, su.bytes); + defer a.free(suh); + const svh = try copyAs(f16, a, sv.bytes); + defer a.free(svh); + var trd = try cuda.DeviceBuffer.fromHost(gpu.d, tr.bytes); + defer trd.free(); + var sud = try cuda.DeviceBuffer.fromHost(gpu.d, su.bytes); + defer sud.free(); + var svd = try cuda.DeviceBuffer.fromHost(gpu.d, sv.bytes); + defer svd.free(); + const strides = exl3.strides(k, n, halves, false); + const w: exl3.Weight = .{ .trellis = trd.ptr, .suh = sud.ptr, .svh = svd.ptr, .k = @intCast(k), .n = @intCast(n), .halves = halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const tiles_ok = if (exl3.Dense.candidateWeight(w)) |_| true else |_| false; + const xh = try a.alloc(u16, max_rows * k); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + var yr = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yr.free(); + var yt = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yt.free(); + // f64 of rows 0 and 1 + const xf = try a.alloc(f64, 2 * k); + defer a.free(xf); + for (xf, xh[0 .. 2 * k]) |*o, b| o.* = bfValue(b); + const want = try a.alloc(f64, 2 * n); + defer a.free(want); + try exact(a, trellis, suh, svh, tk, tn, halves, xf, 2, want); + std.debug.print(" {s} K={d} N={d} halves={d}{s}\n", .{ name, k, n, halves, if (tiles_ok) "" else " (tiles refuse it)" }); + for (row_counts, 0..) |rows, ri| { + de.implementation = .reference; + const rc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yr.ptr, .rows = rows }; + const ms_ref = try t.best(s, 3, rc, DenseCall.go); + sum_ref[ri] += ms_ref; + var ms_tiles: f64 = std.math.nan(f64); + var diff: f64 = std.math.nan(f64); + var err_ref: f64 = 0; + var err_tiles: f64 = std.math.nan(f64); + var peak: f64 = 0; + const got_r = try a.alloc(u16, rows * n); + defer a.free(got_r); + try yr.download(0, std.mem.sliceAsBytes(got_r)); + for (got_r[0 .. 2 * n], want) |g, wv| { + err_ref = @max(err_ref, @abs(bfValue(g) - wv)); + peak = @max(peak, @abs(wv)); + } + if (tiles_ok) { + de.implementation = .tiles; + const tc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yt.ptr, .rows = rows }; + ms_tiles = try t.best(s, 3, tc, DenseCall.go); + sum_tiles[ri] += ms_tiles; + const got_t = try a.alloc(u16, rows * n); + defer a.free(got_t); + try yt.download(0, std.mem.sliceAsBytes(got_t)); + diff = 0; + var rpeak: f64 = 0; + for (got_r, got_t) |g, h| { + diff = @max(diff, @abs(bfValue(g) - bfValue(h))); + rpeak = @max(rpeak, @abs(bfValue(g))); + } + diff /= rpeak; + err_tiles = 0; + for (got_t[0 .. 2 * n], want) |g, wv| err_tiles = @max(err_tiles, @abs(bfValue(g) - wv)); + de.implementation = .reference; + } else sum_tiles[ri] += ms_ref; + std.debug.print(" rows {d:5}: reference {d:8.4} ms ({d:6.2} us/row) tiles {d:8.4} ms ({d:6.2} us/row) max|ref-tiles|/peak {e:.2} rows 0-1 vs f64: ref {e:.2} tiles {e:.2} (rel to peak {d:.3})\n", .{ rows, ms_ref, 1e3 * ms_ref / @as(f64, @floatFromInt(rows)), ms_tiles, 1e3 * ms_tiles / @as(f64, @floatFromInt(rows)), diff, err_ref / peak, err_tiles / peak, peak }); + } + } + std.debug.print(" sum over the {d} shapes (ms): ", .{names.len}); + for (row_counts, sum_ref, sum_tiles) |r, x, y| std.debug.print("rows {d}: reference {d:.3} best {d:.3}; ", .{ r, x, y }); + std.debug.print("\n", .{}); + check.pass("dense prompt kernels timed", .{}); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + if (std.mem.eql(u8, mode, "engine")) return engineRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "agree")) return agreeRun(gpu, dir, rest[2..]); + const layer: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 17; + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + if (std.mem.eql(u8, mode, "dense")) return denseRun(gpu, &pack, s); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + if (std.mem.eql(u8, mode, "experts")) return expertsRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "copy")) return copyRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "prompt")) return promptRun(gpu, &ex, s, ld); + return error.UnknownMode; +} + +// ---------------------------------------------------------------------------------------------------- engine + +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const xpre = flash.exl3.prefill; +const vocab = prompt.vocabulary; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + attached: xpass.Attached, + + fn init(gpu: Gpu, dir: []const u8, lines: ?usize, cap: usize, deep_rows: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + var free0: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free0, &total), "cuMemGetInfo"); + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + var free1: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free1, &total), "cuMemGetInfo"); + g.seq = try g.e.newSeq(cap); + var free2: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free2, &total), "cuMemGetInfo"); + if (deep_rows > 0) try xpre.reserve(gpu.gpa, gpu.d, g.e, deep_rows); + var free3: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free3, &total), "cuMemGetInfo"); + var b = try xpass.budgetFromEnv(0); + if (lines) |n| b.lines = n; + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, b, null); + var free4: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free4, &total), "cuMemGetInfo"); + const gib = struct { + fn f(x: usize) f64 { + return @as(f64, @floatFromInt(x)) / (1 << 30); + } + }.f; + std.debug.print("engine: loaded in {d:.1} s; card {d:.2} GiB; after weights {d:.2} free; pass scratch (chunk {d}) {d:.2} GiB; caches for {d} tokens {d:.2} GiB; deep chunks of {d} rows {d:.2} GiB; pool {d} lines a layer {d:.2} GiB; {d:.2} GiB free after\n", .{ @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9, gib(total), gib(free0), g.e.chunk, gib(free0 - free1), cap, gib(free1 - free2), deep_rows, gib(free2 - free3), g.attached.lines, gib(g.attached.device_bytes), gib(free4) }); + return g; + } + + fn deinit(g: *Eng) void { + xpre.release(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn poolStats(g: *Eng) flash.exl3_pool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } +}; + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +/// One prefill of `ids` from a fresh sequence: ms, and the last row's logits copied into `out`. +fn prefillOnce(g: *Eng, ids: []const u32, deep_on: bool, out: []u16) !f64 { + var o = xpre.optionsFromEnv(); + o.enabled = deep_on; + if (rows_override) |r| o.rows = r; + xpre.options = o; + try g.reset(); + const t0 = check.now(g.gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + const ms_ = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + @memcpy(std.mem.sliceAsBytes(out), g.seq.logits); + return ms_; +} + +/// Greedy decode rounds after each engine-mode prefill (TF_TEST_DECODE, default 32; 0: none). +var decode_steps: usize = 32; + +fn decodeAfter(g: *Eng, n: usize) !void { + var out = try cuda.DeviceBuffer.alloc(g.gpu.d, vocab * 2); + defer out.free(); + var tok = prompt.argmax(g.seq.logits); + const before = g.poolStats(); + const t0 = check.now(g.gpu.io); + var first: f64 = 0; + for (0..decode_steps) |i| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], out.ptr); + tok = prompt.argmax(g.seq.logits); + if (i == 7) first = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + } + const t = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + const after = g.poolStats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" decode after {d} tokens: {d:.2} ms/token over {d} steps (first 8: {d:.2} ms/token; hit {d:.1}%)\n", .{ n, t / @as(f64, @floatFromInt(decode_steps)), decode_steps, first / 8, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))) }); +} + +/// The deep chunk's rows for the next prefills (engine mode's ROWS list), else the environment's. +var rows_override: ?usize = null; + +fn engineRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 2) return error.MissingArgument; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, rest[0], a, .limited(1 << 28)); + defer a.free(text); + const all_ids = try a.alloc(u32, text.len / 4); + defer a.free(all_ids); + for (all_ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + var ns: [16]usize = undefined; + var n_ns: usize = 0; + var it = std.mem.splitScalar(u8, rest[1], ','); + while (it.next()) |tok| { + ns[n_ns] = try std.fmt.parseInt(usize, tok, 10); + n_ns += 1; + } + const lines: ?usize = if (rest.len > 2 and !std.mem.eql(u8, rest[2], "-")) try std.fmt.parseInt(usize, rest[2], 10) else null; + const cmp = rest.len > 3 and std.mem.eql(u8, rest[3], "cmp"); + // ROWS: deep chunk sizes to sweep (rest[3] when not "cmp", or rest[4]) + var rs: [16]usize = undefined; + var n_rs: usize = 0; + const rows_arg: ?[]const u8 = if (rest.len > 3 and !cmp) rest[3] else if (rest.len > 4) rest[4] else null; + if (rows_arg) |ra| { + var rit = std.mem.splitScalar(u8, ra, ','); + while (rit.next()) |tok| { + rs[n_rs] = try std.fmt.parseInt(usize, tok, 10); + n_rs += 1; + } + } + var max_n: usize = 0; + for (ns[0..n_ns]) |n| max_n = @max(max_n, n); + if (max_n > all_ids.len) return error.NotEnoughIds; + if (std.c.getenv("TF_TEST_DECODE")) |v| decode_steps = std.fmt.parseInt(usize, std.mem.span(v), 10) catch decode_steps; + const o = xpre.optionsFromEnv(); + var max_r = o.rows; + for (rs[0..n_rs]) |r| max_r = @max(max_r, r); + if (n_rs > 0) max_r = 0; + for (rs[0..n_rs]) |r| max_r = @max(max_r, r); + const deep_rows = @min(max_r, (max_n + 255) / 256 * 256); + const g = try Eng.init(gpu, dir, lines, max_n + 64, if (o.enabled) deep_rows else 0); + defer g.deinit(); + const old = try a.alloc(u16, vocab); + defer a.free(old); + const new = try a.alloc(u16, vocab); + defer a.free(new); + // warm-up: kernels, the pool's first lines + _ = try prefillOnce(g, all_ids[0..@min(max_n, 1024)], o.enabled, new); + var all_ok = true; + if (n_rs == 0) { + rs[0] = 0; + n_rs = 1; + } + for (rs[0..n_rs]) |rr| for (ns[0..n_ns]) |n| { + rows_override = if (rr == 0) null else rr; + if (rr != 0) std.debug.print("deep chunk rows {d}:\n", .{rr}); + const ids = all_ids[0..n]; + const p0 = g.poolStats(); + xpre.stats = .{}; + const t_new = try prefillOnce(g, ids, o.enabled, new); + const p1 = g.poolStats(); + const xs_ = xpre.stats; + std.debug.print("PREFILL {d:6} tokens ({s}): {d:9.1} ms = {d:7.1} rows/s; pool copies {d} ({d:.1} GiB, {d:.2} a layer per 1k rows), deep chunks {d}, passes {d}, prefetched {d}, members {d}; host: picks wait {d:.1} ms, plan {d:.1} ms\n", .{ n, if (o.enabled) "expert-major deep" else "batched pass", t_new, @as(f64, @floatFromInt(n)) / (t_new / 1e3), p1.misses - p0.misses, @as(f64, @floatFromInt(p1.misses - p0.misses)) * 1.758 / 1024, @as(f64, @floatFromInt(p1.misses - p0.misses)) / 48 / (@as(f64, @floatFromInt(n)) / 1000), xs_.chunks, xs_.passes, xs_.prefetched, xs_.members, @as(f64, @floatFromInt(xs_.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(xs_.plan_ns)) / 1e6 }); + // decode after the prefill: the pools were left empty by the ring (or hold the batched pass's lines) + if (decode_steps > 0) try decodeAfter(g, n); + if (!cmp) continue; + const q0 = g.poolStats(); + const t_old = try prefillOnce(g, ids, false, old); + const q1 = g.poolStats(); + std.debug.print(" same prompt, batched pass: {d:9.1} ms = {d:7.1} rows/s; pool copies {d}\n", .{ t_old, @as(f64, @floatFromInt(n)) / (t_old / 1e3), q1.misses - q0.misses }); + if (decode_steps > 0) try decodeAfter(g, n); + var dmax: f32 = 0; + var peak: f32 = 0; + var same: usize = 0; + for (old, new) |x, y| { + dmax = @max(dmax, @abs(promote(x) - promote(y))); + peak = @max(peak, @abs(promote(x))); + same += @intFromBool(x == y); + } + const to = top5(old); + const tn = top5(new); + var overlap: usize = 0; + for (to) |x| for (tn) |y| { + overlap += @intFromBool(x == y); + }; + const margin_old = promote(old[to[0]]) - promote(old[to[1]]); + std.debug.print("CMP {d:6} tokens: last-row logits max |d| {d:.4} (peak {d:.2}, {d:.1}% of logits bit-equal); top-1 {d} vs {d} ({s}, batched-pass margin {d:.4}); top-5 overlap {d}/5\n", .{ n, dmax, peak, 100 * @as(f64, @floatFromInt(same)) / vocab, to[0], tn[0], if (to[0] == tn[0]) "equal" else "DIFFERENT", margin_old, overlap }); + if (to[0] != tn[0] and margin_old > 0.25) all_ok = false; + }; + const ps = g.poolStats(); + std.debug.print("pool overall: {d} rounds, {d} copies, hit {d:.1}%\n", .{ ps.rounds, ps.misses, 100 * @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(@max(1, ps.hit_slots + ps.miss_slots))) }); + try check.expect(all_ok, "top-1 equal wherever the batched pass's margin exceeds bf16 noise", .{}); + check.pass("engine prefill measured", .{}); +} + +fn readIds(gpu: Gpu, path: []const u8) ![]u32 { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 28)); + defer a.free(text); + const ids = try a.alloc(u32, text.len / 4); + for (ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + return ids; +} + +fn agreeRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 4) return error.MissingArgument; + const all_ids = try readIds(gpu, rest[0]); + defer a.free(all_ids); + const first = try std.fmt.parseInt(usize, rest[1], 10); + const step = try std.fmt.parseInt(usize, rest[2], 10); + const count = try std.fmt.parseInt(usize, rest[3], 10); + const lines: ?usize = if (rest.len > 4) try std.fmt.parseInt(usize, rest[4], 10) else null; + const max_n = first + step * (count - 1); + if (max_n + 1 > all_ids.len) return error.NotEnoughIds; + const o = xpre.optionsFromEnv(); + const g = try Eng.init(gpu, dir, lines, max_n + 64, @min(o.rows, (max_n + 255) / 256 * 256)); + defer g.deinit(); + const old = try a.alloc(u16, vocab); + defer a.free(old); + const new = try a.alloc(u16, vocab); + defer a.free(new); + var agree: usize = 0; + var acc_old: usize = 0; + var acc_new: usize = 0; + var top5_same: usize = 0; + var sum_d: f64 = 0; + var max_d: f64 = 0; + var n_d: usize = 0; + var disagree_margin_max: f32 = 0; + var t_old: f64 = 0; + var t_new: f64 = 0; + for (0..count) |i| { + const n = first + step * i; + const ids = all_ids[0..n]; + t_old += try prefillOnce(g, ids, false, old); + t_new += try prefillOnce(g, ids, o.enabled, new); + const to = top5(old); + const tn = top5(new); + const want = all_ids[n]; + agree += @intFromBool(to[0] == tn[0]); + acc_old += @intFromBool(to[0] == want); + acc_new += @intFromBool(tn[0] == want); + top5_same += @intFromBool(std.mem.eql(u32, &to, &tn)); + for (to) |id| { + const d = @abs(@as(f64, promote(old[id])) - promote(new[id])); + sum_d += d; + max_d = @max(max_d, d); + n_d += 1; + } + if (to[0] != tn[0]) { + const m = promote(old[to[0]]) - promote(old[to[1]]); + disagree_margin_max = @max(disagree_margin_max, m); + std.debug.print(" length {d}: top-1 {d} (batched) vs {d} (deep), batched margin {d:.4}, deep margin {d:.4}\n", .{ n, to[0], tn[0], m, promote(new[tn[0]]) - promote(new[tn[1]]) }); + } + } + std.debug.print("AGREE {d} prompt lengths {d}..{d}: top-1 equal {d}/{d}, top-5 identical in order {d}/{d}; next-token accuracy batched {d}/{d}, deep {d}/{d}; |d| on the batched top-5: mean {d:.4}, max {d:.4}; largest batched margin where top-1 differs {d:.4}; time batched {d:.1} s, deep {d:.1} s\n", .{ count, first, max_n, agree, count, top5_same, count, acc_old, count, acc_new, count, sum_d / @as(f64, @floatFromInt(n_d)), max_d, disagree_margin_max, t_old / 1e3, t_new / 1e3 }); + check.pass("agreement measured", .{}); +} diff --git a/zig/tests/cuda/exl3_recycle_tests.zig b/zig/tests/cuda/exl3_recycle_tests.zig new file mode 100644 index 0000000000..e0bab2670c --- /dev/null +++ b/zig/tests/cuda/exl3_recycle_tests.zig @@ -0,0 +1,205 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; +const paging = flash.exl3.kvpage; +const weights = flash.exl3.prompt_wt; +const snapshots = flash.snapshot; +const cap = 8192; + +pub fn run(gpu: check.Gpu) !void { + const previous = flash.exl3.graph.force; + flash.exl3.graph.force = .graph; + defer flash.exl3.graph.force = previous; + if (!flash.exl3.recycle.enabled()) return error.RecyclingRequired; + for ([_]bool{ false, true }) |shared| for ([_]bool{ false, true }) |host_index| try lifecycle(gpu, shared, host_index); + check.pass("paged Engine.newSeq allocation reuse: private/shared, device/host indexer, real page fetch and snapshot round trips", .{}); +} + +fn zeroDevice(gpu: check.Gpu, buffer: cuda.DeviceBuffer) !void { + const scratch = try gpu.gpa.alloc(u8, @min(buffer.len, 1 << 20)); + defer gpu.gpa.free(scratch); + var offset: usize = 0; + while (offset < buffer.len) { + const size = @min(scratch.len, buffer.len - offset); + try buffer.download(offset, scratch[0..size]); + try check.expect(std.mem.allEqual(u8, scratch[0..size], 0), "reused allocation {x} has stale bytes at {d}", .{ buffer.ptr, offset }); + offset += size; + } +} + +fn zeroSequence(gpu: check.Gpu, stream: cuda.Stream, sequence: *prompt.Seq) !void { + try stream.synchronize(); + for ([_][]cuda.DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }) |group| for (group) |buffer| try zeroDevice(gpu, buffer); + try zeroDevice(gpu, sequence.tail); + const pg = sequence.paged.?; + try check.expect(sequence.fed == 0 and sequence.rank_id == 0 and !sequence.invalidated and !sequence.untaught and !pg.graph_prepared and !pg.retired_for_reuse, "new request resets ordinary and paging metadata", .{}); + try check.expect(std.meta.eql(sequence.history, [2]i64{ 73, 73 }) and std.mem.allEqual(u8, sequence.logits, 0), "new request resets history and logits", .{}); + try check.expect(std.mem.allEqual(u8, pg.host.bytes, 0) and std.mem.allEqual(u8, pg.status_host.bytes, 0), "new request clears pinned canonical rows and status", .{}); + const status = try pg.total(); + try check.expect(status.calls == 0 and status.misses == 0 and status.rebases == 0 and status.written_rows == 0 and !status.overflow, "new request clears pager counters", .{}); + for (0..weights.attn_n) |layer| { + try check.expect(pg.win0[layer] == 0, "layer {d} window resets", .{layer}); + for (0..2) |side| try zeroDevice(gpu, .{ .d = gpu.d, .ptr = pg.devRows(layer, side == 1), .len = pg.wr * paging.row_bytes, .borrowed = true }); + } +} + +fn seed(gpu: check.Gpu, stream: cuda.Stream, sequence: *prompt.Seq, at: usize, tag: u8) !void { + const pg = sequence.paged.?; + const origin = (at -| pg.cfg.window) / weights.index_ratio * weights.index_ratio; + sequence.fed = at; + sequence.history = .{ @intCast(at), tag }; + sequence.rank_id = 19; + sequence.untaught = true; + @memset(sequence.logits, tag); + for ([_][]cuda.DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }, 0..) |group, group_index| { + for (group, 0..) |buffer, index| try buffer.fill8(tag + @as(u8, @intCast(group_index + index)), stream.handle); + } + try sequence.tail.fill8(tag, stream.handle); + try stream.synchronize(); + for (0..weights.attn_n) |layer| { + pg.win0[layer] = origin; + pg.rebases[layer] = layer + 3; + pg.written[layer] = origin; + for (0..2) |side| { + const host = pg.hostMut(layer, side == 1); + const words = std.mem.bytesAsSlice(u16, @as([]align(2) u8, @alignCast(host))); + for (words, 0..) |*word, index| word.* = @truncate(index + layer * 997 + side * 32768 + tag); + try gpu.d.check(gpu.d.api.cuMemsetD8Async(pg.devRows(layer, side == 1), tag + @as(u8, @intCast(layer + side)), pg.wr * paging.row_bytes, stream.handle), "seed pager window"); + } + } + try stream.synchronize(); +} + +fn fetch(gpu: check.Gpu, stream: cuda.Stream, sequence: *prompt.Seq) !void { + const pg = sequence.paged.?; + var ids: [weights.index_idw]i32 = @splat(0); + for (0..weights.index_top) |block| for (0..weights.index_ratio) |offset| { + ids[block * weights.index_ratio + offset] = @intCast(block * weights.index_ratio + offset); + }; + var selected = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&ids)); + defer selected.free(); + const count = [_]i32{weights.index_top * weights.index_ratio}; + const sparse = [_]i32{1}; + var counts = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&count)); + defer counts.free(); + var kinds = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&sparse)); + defer kinds.free(); + const lease = if (pg.shared_arena) |arena| try arena.acquire(pg.shared_sequence) else null; + defer if (lease) |held| if (pg.shared_arena.?.active != null) pg.shared_arena.?.completeAfterConsumers(held) catch unreachable; + errdefer gpu.d.synchronize() catch @panic("fixture paging did not drain"); + try pg.page(stream, 0, selected.ptr, counts.ptr, kinds.ptr, 1); + if (lease) |held| try pg.shared_arena.?.complete(held, stream) else try stream.synchronize(); + const status = try pg.stats(0); + try check.expect(status.misses > 0 and status.calls > 0 and !status.overflow, "real pager fetches canonical BF16 rows", .{}); + try selected.download(0, std.mem.asBytes(&ids)); + const slot = std.math.cast(usize, ids[0]) orelse return error.InvalidRemappedKey; + const layout = pg.slotRows(0); + try check.expect(slot >= layout[1] and slot + 4 <= layout[0], "real sparse keys remap into owned slots", .{}); + var actual: [4 * paging.row_bytes]u8 = undefined; + for (0..2) |side| { + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(&actual, pg.devRows(0, side == 1) + slot * paging.row_bytes, actual.len), "read fetched block"); + try check.sameBytes("real fetched BF16 block", &actual, pg.hostRows(0, side == 1)[0..actual.len]); + } +} + +fn sameSnapshot(expected: *const snapshots.Snap, actual: *const snapshots.Snap) !void { + try check.expect(std.meta.eql(expected.paged, actual.paged) and expected.at == actual.at and std.meta.eql(expected.history, actual.history), "restored snapshot metadata remains exact", .{}); + try check.sameBytes("restored canonical, windows, tables, GDN, indexer, PLE and logits", actual.host.?.bytes[0..actual.buf.len], expected.host.?.bytes[0..expected.buf.len]); + if (expected.owned_blocks) |wanted| { + const found = actual.owned_blocks orelse return error.MissingOwnedBlocks; + try check.expect(wanted.entries.len == found.entries.len, "shared snapshot retains every owned block", .{}); + for (wanted.entries, found.entries) |before, after| try check.expect(std.meta.eql(before.key, after.key), "shared snapshot remaps logical block keys", .{}); + try check.sameBytes("restored shared BF16 blocks", found.payload, wanted.payload); + } +} + +fn identity(sequence: *const prompt.Seq) u64 { + const pg = sequence.paged.?; + var digest = std.hash.Wyhash.init(0); + const addresses = [_]u64{ @intFromPtr(sequence.logits.ptr), @intFromPtr(pg), @intFromPtr(pg.status_host.bytes.ptr), sequence.tail.ptr, pg.host_dev, pg.dev.ptr, pg.tabs.ptr, pg.work.ptr, pg.devRows(0, false) }; + digest.update(std.mem.asBytes(&addresses)); + for ([_][]cuda.DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }) |group| for (group) |buffer| { + const shape = [_]u64{ buffer.ptr, buffer.len }; + digest.update(std.mem.asBytes(&shape)); + }; + return digest.final(); +} + +fn digestSnapshot(saved: *const snapshots.Snap) u64 { + var digest = std.hash.Wyhash.hash(0, saved.host.?.bytes[0..saved.buf.len]); + if (saved.owned_blocks) |owned| digest ^= std.hash.Wyhash.hash(1, owned.payload); + return digest; +} + +fn lifecycle(gpu: check.Gpu, shared: bool, host_index: bool) !void { + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var engine: prompt.Engine = undefined; + engine.gpa = gpu.gpa; + engine.io = gpu.io; + engine.driver = @constCast(gpu.d); + engine.stream = &stream; + engine.chunk = 256; + engine.eos = 73; + engine.round_graph = .{}; + engine.page_config = .{ .window = 2052, .rows = 256, .slots = 1024, .shared_lanes = if (shared) 3 else 0 }; + engine.shared_arena = null; + try flash.exl3.kv.enable(&engine, false, false, host_index); + defer flash.exl3.kv.forget(gpu.gpa, &engine); + _ = try engine.reserveSharedKv(cap); + defer if (engine.shared_arena) |arena| arena.destroy() catch @panic("fixture still owns shared lanes"); + defer flash.exl3.recycle.drain(&engine); + var foreign: ?prompt.Seq = if (shared) try engine.newSeq(cap) else null; + defer if (foreign) |*sequence| sequence.deinit(gpu.gpa); + if (foreign) |*sequence| { + try seed(gpu, stream, sequence, 4099, 0x12); + try fetch(gpu, stream, sequence); + } + const foreign_saved = if (foreign) |*sequence| try snapshots.save(gpu.gpa, gpu.d, stream, sequence, null) else null; + defer if (foreign_saved) |saved| snapshots.drop(gpu.gpa, stream, saved); + var current = try engine.newSeq(cap); + var live = true; + defer if (live) current.deinit(gpu.gpa); + for ([_]usize{ 4099, 6147 }) |at| { + try seed(gpu, stream, ¤t, at, 0x31); + try fetch(gpu, stream, ¤t); + const saved = try snapshots.save(gpu.gpa, gpu.d, stream, ¤t, null); + defer snapshots.drop(gpu.gpa, stream, saved); + const digest = digestSnapshot(saved); + const addresses = identity(¤t); + const pg = current.paged.?; + const incarnation = pg.incarnation; + const key = pg.graphKey(); + const old_owner = if (shared) pg.shared_sequence else null; + const reused = flash.exl3.recycle.stats.seq_reused; + live = false; + current.deinit(gpu.gpa); + if (old_owner) |old| { + if (engine.shared_arena.?.acquire(old)) |held| { + try engine.shared_arena.?.cancel(held); + return error.StaleOwnerAccepted; + } else |err| if (err != error.SharedKvSequence) return err; + } + current = try engine.newSeq(cap); + live = true; + try check.expect(flash.exl3.recycle.stats.seq_reused == reused + 1 and std.meta.eql(addresses, identity(¤t)), "Engine.newSeq reuses every sequence/pager allocation address", .{}); + try check.expect(current.paged.?.incarnation != incarnation and !std.meta.eql(key, current.paged.?.graphKey()), "retained addresses receive a fresh strict paging identity", .{}); + if (old_owner) |old| try check.expect(current.paged.?.shared_sequence.generation > old.generation, "reattached shared owner advances generation", .{}); + try zeroSequence(gpu, stream, ¤t); + try seed(gpu, stream, ¤t, at, 0x47); + try fetch(gpu, stream, ¤t); + try snapshots.restore(stream, ¤t, null, saved); + const again = try snapshots.save(gpu.gpa, gpu.d, stream, ¤t, null); + defer snapshots.drop(gpu.gpa, stream, again); + try sameSnapshot(saved, again); + try check.expect(digestSnapshot(saved) == digest, "source snapshot remains immutable after allocation reuse", .{}); + if (foreign) |*sequence| { + const untouched = try snapshots.save(gpu.gpa, gpu.d, stream, sequence, null); + defer snapshots.drop(gpu.gpa, stream, untouched); + try sameSnapshot(foreign_saved.?, untouched); + } + std.debug.print("RECYCLE shared={} host_index={} at={d} tail={x} host={x} incarnation={d}->{d} snapshot={x}\n", .{ shared, host_index, at, current.tail.ptr, pg.host_dev, incarnation, pg.incarnation, digest }); + } +} diff --git a/zig/tests/cuda/exl3_route.zig b/zig/tests/cuda/exl3_route.zig new file mode 120000 index 0000000000..fb1ea4df16 --- /dev/null +++ b/zig/tests/cuda/exl3_route.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_route.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_router.zig b/zig/tests/cuda/exl3_router.zig new file mode 120000 index 0000000000..dc3d285b11 --- /dev/null +++ b/zig/tests/cuda/exl3_router.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_router.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_scan.zig b/zig/tests/cuda/exl3_scan.zig new file mode 120000 index 0000000000..168ffe1f6f --- /dev/null +++ b/zig/tests/cuda/exl3_scan.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_scan.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_shape.zig b/zig/tests/cuda/exl3_shape.zig new file mode 120000 index 0000000000..2b1c555b6c --- /dev/null +++ b/zig/tests/cuda/exl3_shape.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_shape.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_state_rebind_tests.zig b/zig/tests/cuda/exl3_state_rebind_tests.zig new file mode 100644 index 0000000000..930c20b645 --- /dev/null +++ b/zig/tests/cuda/exl3_state_rebind_tests.zig @@ -0,0 +1,113 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const graph = flash.exl3.graph; +const iterations = 4; +const prefix = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + +fn require(name: [:0]const u8, value: []const u8) !void { + const setting = std.c.getenv(name); + try check.expect(setting != null and std.mem.eql(u8, std.mem.span(setting.?), value), "rebind fixture requires {s}={s}", .{ name, value }); +} + +fn stateBytes(sequence: *const flash.prompt.Seq) usize { + var total = sequence.tail.len; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv }) |buffers| for (buffers) |buffer| { + total += buffer.len; + }; + return total; +} + +fn readState(sequence: *const flash.prompt.Seq, target: []u8) !void { + var offset: usize = 0; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv }) |buffers| for (buffers) |buffer| { + try buffer.download(0, target[offset..][0..buffer.len]); + offset += buffer.len; + }; + try sequence.tail.download(0, target[offset..][0..sequence.tail.len]); +} + +pub fn run(engine: anytype, comptime operations: anytype) !void { + try require("TF_FLASHNEXT_EXL3_GRAPH_STATE_REBIND", "1"); + try require("TF_FLASHNEXT_EXL3_GRAPH_AFTER", "1"); + try require("TF_FLASHNEXT_EXL3_GRAPH_ASYNC", "0"); + try check.expect(engine.seq.paged == null and engine.world == 2, "rebind fixture requires two ranks and the dense control KV layout", .{}); + var draft: flash.mtp.Draft = undefined; + draft.keep = .{}; + draft.verify_cache = .{}; + try check.expect(engine.seq.state.len == draft.spare.len, "rebind fixture requires the complete recurrent layer layout", .{}); + var made: usize = 0; + defer { + cuda.completion.drain(engine.gpu.d, engine.stream, "state rebind fixture consumers did not drain"); + graph.clear(); + for (draft.spare[0..made]) |*buffer| buffer.free(); + } + for (engine.seq.state) |state| { + draft.spare[made] = try cuda.DeviceBuffer.alloc(engine.gpu.d, state.len); + made += 1; + } + var linear: usize = 0; + for (engine.e.layers[0..engine.e.n_layers], 0..) |layer, index| { + if (!layer.linear) continue; + draft.keep.layers[index].state_out = draft.spare[linear].ptr; + linear += 1; + } + try check.expect(linear == made, "rebind fixture maps every physical recurrent layer", .{}); + const allocator = engine.gpu.gpa; + const state_bytes = stateBytes(&engine.seq); + const reference_state = try allocator.alloc(u8, iterations * state_bytes); + defer allocator.free(reference_state); + const actual_state = try allocator.alloc(u8, state_bytes); + defer allocator.free(actual_state); + const reference_logits = try allocator.alloc(u8, iterations * 16 * flash.prompt.vocabulary * 2); + defer allocator.free(reference_logits); + const actual_logits = try allocator.alloc(u8, 16 * flash.prompt.vocabulary * 2); + defer allocator.free(actual_logits); + const previous_mode = graph.force; + defer graph.force = previous_mode; + for ([_]usize{ 3, 4, 16 }) |width| { + const row_bytes = width * flash.prompt.vocabulary * 2; + var before: graph.Stats = undefined; + for ([_]bool{ false, true }) |captured| { + try flash.exl3.dplan.toHost(); + graph.clear(); + graph.force = .off; + try operations.reset(engine); + for (draft.spare) |buffer| try buffer.fill8(0, engine.stream.handle); + try operations.forward(engine, &prefix, .{ .logits = .last }); + graph.force = if (captured) .graph else .dplan; + before = graph.stats; + for (0..iterations) |step| { + try operations.forward(engine, prefix[0..width], .{ .logits = .all, .logits_out = engine.out.ptr, .keep = &draft.keep }); + try engine.out.download(0, actual_logits[0..row_bytes]); + const input_before = engine.seq.state[0].ptr; + const output_before = draft.spare[0].ptr; + draft.open = true; + draft.settle(&engine.seq); + try check.expect(engine.seq.state[0].ptr == output_before and draft.spare[0].ptr == input_before, "rank {d}: width {d} step {d} actually swaps the live state pair", .{ engine.rank, width, step }); + try readState(&engine.seq, actual_state); + const wanted_state = reference_state[step * state_bytes ..][0..state_bytes]; + const wanted_logits = reference_logits[step * row_bytes ..][0..row_bytes]; + if (!captured) { + @memcpy(wanted_state, actual_state); + @memcpy(wanted_logits, actual_logits[0..row_bytes]); + } + const equal = std.mem.eql(u8, wanted_state, actual_state) and std.mem.eql(u8, wanted_logits, actual_logits[0..row_bytes]); + const other_equal = try operations.swap(engine, @intFromBool(equal)); + try check.expect(equal and other_equal == 1, "rank {d}: width {d} step {d} exact eager/graph logits, recurrent state, convolution and PLE tail", .{ engine.rank, width, step }); + const digest = std.hash.Wyhash.hash(0, actual_logits[0..row_bytes]); + const other_digest = try operations.swap(engine, digest); + try check.expect(other_digest == digest, "rank {d}: width {d} step {d} cross-rank BF16 logits", .{ engine.rank, width, step }); + } + } + const after = graph.stats; + const succeeded = after.rebinds - before.rebinds; + const failures = after.rebind_failures - before.rebind_failures; + const proved = succeeded == iterations - 1 and failures == 0 and after.rebind_attempts - before.rebind_attempts == succeeded and after.captures - before.captures == 1 and after.replays - before.replays == succeeded; + const other_proved = try operations.swap(engine, @intFromBool(proved)); + try check.expect(proved and other_proved == 1, "rank {d}: width {d} needs one capture and three actual successful widened rebind replays, zero failures; got {d}/{d}", .{ engine.rank, width, succeeded, failures }); + check.pass("rank {d}/2 width {d}: real state flips; exact eager/graph logits and state; actual widened rebinds {d}, failures {d}", .{ engine.rank, width, succeeded, failures }); + } +} diff --git a/zig/tests/cuda/exl3_tp.zig b/zig/tests/cuda/exl3_tp.zig new file mode 120000 index 0000000000..483f3574f0 --- /dev/null +++ b/zig/tests/cuda/exl3_tp.zig @@ -0,0 +1 @@ +../../src/families/flashnext/exl3_tp.zig \ No newline at end of file diff --git a/zig/tests/cuda/exl3_tpprefill_tests.zig b/zig/tests/cuda/exl3_tpprefill_tests.zig new file mode 100644 index 0000000000..8d72d46a30 --- /dev/null +++ b/zig/tests/cuda/exl3_tpprefill_tests.zig @@ -0,0 +1,347 @@ +//! exl3-tpprefill: the expert-major deep-chunk prefill (cuda_exl3_prefill.zig) at tp=1 and at tp=2 expert parallel +//! (cuda_exl3_tp.zig), through the Flash Next family's Engine with each rank's own offload pool attached as the server +//! attaches it. One process a card; at tp=2 both ranks run the same calls in lockstep over the family's Collective. +//! +//! exl3-tpprefill RANK WORLD PORT DIR IDS N1,N2,.. ROWS CAP MODE DUMP +//! IDS int32 token ids (real text); each N takes its first N ids from a fresh sequence of CAP tokens +//! ROWS the deep chunk (rows, reserved before the pools); MODE deep | both (deep, then the batched pass) +//! DUMP path prefix: WORLD 1 writes DUMP-N.bin (deep last-row logits bf16 + GREEDY greedy ids), WORLD 2 compares +//! rank 0's against it (max |d|, top-1, greedy ids) +//! Env TF_TEST_GREEDY (default 16) greedy decode rounds after each prefill; TF_TEST_BOTH_MAX: the batched pass only +//! up to that many tokens. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const xpre = X.prefill; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !*Ranks { + const r = try gpu.gpa.create(Ranks); + errdefer gpu.gpa.destroy(r); + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 600_000); + errdefer link.close(); + r.* = .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + return r; + } + + fn close(r: *Ranks, gpa: std.mem.Allocator) void { + r.collective.deinit(); + r.link.close(); + gpa.destroy(r); + } + + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +const Eng = struct { + gpu: Gpu, + rank: usize, + world: usize, + ranks: ?*Ranks, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + attached: xpass.Attached, + + fn init(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, cap: usize, deep_rows: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.rank = rank; + g.world = world; + g.ranks = if (world == 2) try Ranks.open(gpu, rank, port) else null; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + if (g.ranks) |r| g.e.collective = &r.collective; + var free: [6]usize = undefined; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[0], &total), "cuMemGetInfo"); + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[1], &total), "cuMemGetInfo"); + g.seq = try g.e.newSeq(cap); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[2], &total), "cuMemGetInfo"); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, vocab * 2); + if (deep_rows > 0) try xpre.reserveAt(gpu.gpa, gpu.d, g.e, deep_rows, world); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[3], &total), "cuMemGetInfo"); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, try xpass.budgetFromEnv(0), null); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[4], &total), "cuMemGetInfo"); + std.debug.print("MEM rank {d}/{d}: loaded in {d:.1} s; card {d:.2} GiB; free after weights {d:.2}; pass scratch (chunk {d}) + scores {d:.2} GiB; caches for {d} tokens {d:.2} GiB; deep chunks of {d} rows {d:.2} GiB (layout {d:.2}); pool {d} lines a layer x {d} layers = {d:.2} GiB (line {d:.2} MiB, ring {d} slots); {d:.2} GiB free after; head columns {d}\n", .{ rank, world, ms(t0, check.now(gpu.io)) / 1e3, gib(total), gib(free[0]), g.e.chunk, gib(free[0] - free[1]), cap, gib(free[1] - free[2]), deep_rows, gib(free[2] - free[3]), gib(xpre.deviceBytesAt(deep_rows, world)), g.attached.lines, g.attached.layers, gib(g.attached.device_bytes), @as(f64, @floatFromInt(g.attached.line_bytes)) / (1 << 20), g.attached.lines * 48, gib(free[4]), g.e.head_q.n }); + return g; + } + + fn deinit(g: *Eng) void { + xpre.release(); + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + if (g.ranks) |r| r.close(g.gpu.gpa); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn swap(g: *Eng, v: u64) !u64 { + return if (g.ranks) |r| r.swap(v) else v; + } + + fn finish(g: *Eng) !void { + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + } + + fn forward(g: *Eng, ids: []const u32) !void { + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + try g.finish(); + } + + fn round(g: *Eng, tok: u32) !u32 { + const q = [_]*prompt.Seq{&g.seq}; + var t = tok; + try g.e.round(flash.Tri, g.tri(), &q, (&t)[0..1], g.out.ptr); + try g.finish(); + return prompt.argmax(g.seq.logits); + } +}; + +fn top1(row: []const u16) u32 { + var best: usize = 0; + for (row, 0..) |b, i| if (promote(b) > promote(row[best])) { + best = i; + }; + return @intCast(best); +} + +fn margin(row: []const u16) f32 { + const b = top1(row); + var second: f32 = -std.math.inf(f32); + for (row, 0..) |v, i| if (i != b) { + second = @max(second, promote(v)); + }; + return promote(row[b]) - second; +} + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +const Run = struct { ms: f64, logits: []u16, greedy: []u32, decode_ms: f64, stats: xpre.Stats }; + +/// A fresh sequence, `ids` prefilled (deep or the batched pass), then `steps` greedy rounds. +fn once(g: *Eng, ids: []const u32, deep: bool, rows: usize, steps: usize) !Run { + const a = g.gpu.gpa; + var o = xpre.optionsFromEnv(); + o.enabled = deep; + o.rows = rows; + xpre.options = o; + xpre.stats = .{}; + try g.reset(); + _ = try g.swap(ids.len); + const t0 = check.now(g.gpu.io); + try g.forward(ids); + const t = ms(t0, check.now(g.gpu.io)); + const st = xpre.stats; + const logits = try a.alloc(u16, vocab); + @memcpy(std.mem.sliceAsBytes(logits), g.seq.logits); + const greedy = try a.alloc(u32, steps + 1); + greedy[0] = top1(logits); + const t1 = check.now(g.gpu.io); + for (1..steps + 1) |i| greedy[i] = try g.round(greedy[i - 1]); + const dt = if (steps > 0) ms(t1, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)) else 0; + return .{ .ms = t, .logits = logits, .greedy = greedy, .decode_ms = dt, .stats = st }; +} + +fn freeRun(a: std.mem.Allocator, r: Run) void { + a.free(r.logits); + a.free(r.greedy); +} + +fn compare(what: []const u8, n: usize, ref: []const u16, got: []const u16, ref_g: []const u32, got_g: []const u32) bool { + var dmax: f32 = 0; + var peak: f32 = 0; + var same: usize = 0; + for (ref, got) |x, y| { + dmax = @max(dmax, @abs(promote(x) - promote(y))); + peak = @max(peak, @abs(promote(x))); + same += @intFromBool(x == y); + } + var first: ?usize = null; + for (ref_g, got_g, 0..) |x, y, i| if (x != y and first == null) { + first = i; + }; + const t_eq = top1(ref) == top1(got); + std.debug.print("CMP {s} {d} tokens: last-row logits max |d| {d:.4} (peak {d:.2}, {d:.1}% bit-equal); top-1 {d} vs {d} ({s}, reference margin {d:.4}); greedy {d}: {s}", .{ what, n, dmax, peak, 100 * @as(f64, @floatFromInt(same)) / vocab, top1(ref), top1(got), if (t_eq) "equal" else "DIFFERENT", margin(ref), ref_g.len, if (first == null) "identical" else "differ" }); + if (first) |f| std.debug.print(" from step {d} ({d} vs {d})", .{ f, ref_g[f], got_g[f] }); + std.debug.print("\n", .{}); + if (!t_eq) { + // the two candidates in both rows, and the top-5 overlap + const a_ = top1(ref); + const b_ = top1(got); + const r5 = top5(ref); + const g5 = top5(got); + var overlap: usize = 0; + for (r5) |x| for (g5) |y| { + overlap += @intFromBool(x == y); + }; + std.debug.print(" reference: logit[{d}] {d:.4}, logit[{d}] {d:.4}; ours: logit[{d}] {d:.4}, logit[{d}] {d:.4}; top-5 overlap {d}/5 (reference {any}, ours {any})\n", .{ a_, promote(ref[a_]), b_, promote(ref[b_]), a_, promote(got[a_]), b_, promote(got[b_]), overlap, r5, g5 }); + } + return t_eq or margin(ref) < 0.25; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 10) return error.MissingArgument; + const a = gpu.gpa; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const world = try std.fmt.parseInt(usize, rest[1], 10); + const port = try std.fmt.parseInt(u16, rest[2], 10); + const dir = rest[3]; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, rest[4], a, .limited(1 << 28)); + defer a.free(text); + const all_ids = try a.alloc(u32, text.len / 4); + defer a.free(all_ids); + for (all_ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + var ns: [16]usize = undefined; + var n_ns: usize = 0; + var it = std.mem.splitScalar(u8, rest[5], ','); + while (it.next()) |tok| { + ns[n_ns] = try std.fmt.parseInt(usize, tok, 10); + n_ns += 1; + } + const rows = try std.fmt.parseInt(usize, rest[6], 10); + const cap = try std.fmt.parseInt(usize, rest[7], 10); + const both = std.mem.eql(u8, rest[8], "both"); + const dump = rest[9]; + var steps: usize = 16; + if (std.c.getenv("TF_TEST_GREEDY")) |v| steps = std.fmt.parseInt(usize, std.mem.span(v), 10) catch steps; + var both_max: usize = std.math.maxInt(usize); + if (std.c.getenv("TF_TEST_BOTH_MAX")) |v| both_max = std.fmt.parseInt(usize, std.mem.span(v), 10) catch both_max; + for (ns[0..n_ns]) |n| if (n + steps + 1 > cap or n > all_ids.len) return error.NotEnoughRoom; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const g = try Eng.init(gpu, rank, world, port, dir, cap, rows); + defer g.deinit(); + // warm-up: kernels, tiles, the ring + { + const w = try once(g, all_ids[0..@min(1024, all_ids.len)], true, rows, 1); + freeRun(a, w); + } + var all_ok = true; + for (ns[0..n_ns]) |n| { + const ids = all_ids[0..n]; + const d = try once(g, ids, true, rows, steps); + defer freeRun(a, d); + const s = d.stats; + std.debug.print("PREFILL rank {d}/{d} {d:6} tokens deep: {d:9.1} ms = {d:7.1} rows/s; chunks {d}, passes {d}, copies {d} ({d} prefetched), all-reduces {d}; host picks wait {d:.1} ms, plan {d:.1} ms; decode after {d:.2} ms/token\n", .{ rank, world, n, d.ms, @as(f64, @floatFromInt(n)) / (d.ms / 1e3), s.chunks, s.passes, s.copies, s.prefetched, s.reduces, @as(f64, @floatFromInt(s.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(s.plan_ns)) / 1e6, d.decode_ms }); + if (s.reduce_ns > 0) std.debug.print(" rank {d}: combine + all-reduce + place {d:.1} ms over {d} layers (profiled, each waited)\n", .{ rank, @as(f64, @floatFromInt(s.reduce_ns)) / 1e6, s.reduces }); + if (both and n <= both_max) { + const b = try once(g, ids, false, rows, steps); + defer freeRun(a, b); + std.debug.print("PREFILL rank {d}/{d} {d:6} tokens batched: {d:9.1} ms = {d:7.1} rows/s; decode after {d:.2} ms/token\n", .{ rank, world, n, b.ms, @as(f64, @floatFromInt(n)) / (b.ms / 1e3), b.decode_ms }); + if (rank == 0) all_ok = compare("deep vs batched (same world)", n, b.logits, d.logits, b.greedy, d.greedy) and all_ok; + } + // both ranks hold the same logits and greedy ids + const h = std.hash.Wyhash.hash(std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(d.logits)), std.mem.sliceAsBytes(d.greedy)); + const other = try g.swap(h); + if (world == 2 and other != h) { + std.debug.print("rank {d}: the ranks' logits or greedy ids DIFFER at {d} tokens\n", .{ rank, n }); + all_ok = false; + } + if (rank != 0) continue; + var pb: [512]u8 = undefined; + const path = try std.fmt.bufPrint(&pb, "{s}-{d}.bin", .{ dump, n }); + if (world == 1) { + const bytes = try a.alloc(u8, vocab * 2 + d.greedy.len * 4); + defer a.free(bytes); + @memcpy(bytes[0 .. vocab * 2], std.mem.sliceAsBytes(d.logits)); + @memcpy(bytes[vocab * 2 ..], std.mem.sliceAsBytes(d.greedy)); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = bytes }); + continue; + } + const ref = std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 24)) catch |err| { + std.debug.print("no tp=1 dump at {s} ({t})\n", .{ path, err }); + continue; + }; + defer a.free(ref); + if (ref.len != vocab * 2 + d.greedy.len * 4) { + std.debug.print("tp=1 dump {s} has another greedy length; skipped\n", .{path}); + continue; + } + const rl = try a.alloc(u16, vocab); + defer a.free(rl); + @memcpy(std.mem.sliceAsBytes(rl), ref[0 .. vocab * 2]); + const rg = try a.alloc(u32, d.greedy.len); + defer a.free(rg); + @memcpy(std.mem.sliceAsBytes(rg), ref[vocab * 2 ..]); + all_ok = compare("tp=2 deep vs tp=1 deep", n, rl, d.logits, rg, d.greedy) and all_ok; + } + const both_ok = try g.swap(@intFromBool(all_ok)); + try check.expect(all_ok and both_ok == 1, "rank {d}: a top-1 differs past bf16 noise or the ranks disagree", .{rank}); + check.pass("rank {d}/{d}: deep prefill measured and compared", .{ rank, world }); +} diff --git a/zig/tests/cuda/flashnext_exl3_engine_tests.zig b/zig/tests/cuda/flashnext_exl3_engine_tests.zig new file mode 100644 index 0000000000..7935c20906 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3_engine_tests.zig @@ -0,0 +1,528 @@ +//! flashnext-exl3-engine: the Qwen3.8-Flash-Next EXL3 pack through the Flash Next family's own Engine +//! (families/flashnext/cuda_prompt.zig), with the offload pool attached exactly as the server attaches it +//! (cuda_exl3_pass.attach, called from cuda_lanes_load.open), checked against the independent reference golden +//! (zig/tests/fixtures/exl3_ref_golden.json: the Python engine's bf16 logits and greedy ids). +//! +//! flashnext-exl3-engine DIR tf GOLDEN [LINES] the golden's 16-id prompt teacher-forced in one pass: every +//! position's top-1 and top-5 against the golden's, the logit +//! difference on the golden's top-5 ids, and the EXL3 lm_head on the +//! pass's own head rows against an exact float64 product +//! flashnext-exl3-engine DIR tfc GOLDEN [LINES] each golden prompt and the golden's 24 greedy ids teacher-forced: +//! our top-1 at every continuation step, with both margins +//! flashnext-exl3-engine DIR gen GOLDEN [LINES] greedy 24 tokens for each golden prompt (prefill, then +//! Engine.round a token): ids equal to the golden's +//! flashnext-exl3-engine DIR bench [LINES] prefill rows/s at 16, 256 and 2048 tokens and decode ms/token, +//! with the pool's hit rate (a measurement on an idle card) +//! +//! LINES (or TF_FLASHNEXT_EXL3_LINES) fixes the pool's lines a layer; without it the pool takes free memory. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const Gpu = check.Gpu; +const st = core.safetensors; +const fmt = core.exl3_format; +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const vocab = prompt.vocabulary; +const dims: usize = 2560; +const cap_default: usize = 32768; +/// Rows of logits a teacher-forced pass returns at most. +const out_rows: usize = 64; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + cap: usize, + attached: xpass.Attached, + load_s: f64, + + fn init(gpu: Gpu, dir: []const u8, lines: ?usize, cap: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.cap = cap; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + const t1 = check.now(gpu.io); + // as cuda_lanes_load.open: the pass's scratch for a full chunk and the block scores, then the stream's caches + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + g.seq = try g.e.newSeq(cap); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + var b = try xpass.budgetFromEnv(0); + if (lines) |n| b.lines = n; + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, b, null); + const t2 = check.now(gpu.io); + g.load_s = ms(t0, t2) / 1e3; + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + std.debug.print("engine: loaded in {d:.1} s (Engine.initRank {d:.1} s incl. the pinned bank's copies {d:.1} s; pools {d:.1} s); {d} lines a layer x {d} layers ({d:.2} GiB), {d:.2} GiB free after, chunk {d}, context {d}\n", .{ g.load_s, ms(t0, t1) / 1e3, @as(f64, @floatFromInt(g.attached.bank_ns)) / 1e9, ms(t1, t2) / 1e3, g.attached.lines, g.attached.layers, @as(f64, @floatFromInt(g.attached.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(free)) / (1 << 30), g.e.chunk, cap }); + return g; + } + + fn deinit(g: *Eng) void { + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + /// The sequence as Engine.newSeq makes it, without a new allocation (the pools took the free memory). + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn stats(g: *Eng) flash.exl3_pool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } +}; + +const Golden = struct { + parsed: std.json.Parsed(std.json.Value), + + fn open(gpu: Gpu, path: []const u8) !Golden { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpu.gpa, text, .{}) }; + } + + fn prompts(g: *const Golden) []std.json.Value { + return g.parsed.value.object.get("prompts").?.array.items; + } +}; + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn num(v: std.json.Value) f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => std.math.nan(f64), + }; +} + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +// -- the EXL3 head in float64 ------------------------------------------------------------------------------------ + +const Raw = struct { + files: [7]st.File = undefined, + n: usize = 0, + + fn open(gpu: Gpu, dir: []const u8) !Raw { + var r: Raw = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + r.files[i] = try st.File.open(gpu.gpa, gpu.io, try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 })); + r.n += 1; + } + return r; + } + + fn close(r: *Raw, io: std.Io) void { + for (r.files[0..r.n]) |*f| f.close(io); + } + + fn get(r: *const Raw, name: []const u8) !st.Tensor { + for (r.files[0..r.n]) |*f| if (f.get(name)) |t| return t; + return error.MissingTensor; + } +}; + +const HeadWork = struct { + head: flash.exl3_mtp.Head, + x: []const f64, // [rows, K] + rows: usize, + out: []f64, // [rows, N] + next: std.atomic.Value(usize) = .init(0), + failed: ?anyerror = null, +}; + +fn headWorker(w: *HeadWork, gpa: std.mem.Allocator) void { + headWork(w, gpa) catch |err| { + w.failed = err; + }; +} + +/// logits[r][v] = sum_k x[r][k] * W[k][v], W decoded in float64 a 128-column block at a time (exl3_format.dequantize). +fn headWork(w: *HeadWork, gpa: std.mem.Allocator) !void { + const h = w.head; + const bn: usize = fmt.hadamard_dim; + const tiles_k = h.k / 16; + const tiles_n = h.n / 16; + const per = bn / 16; + const words = h.words(); + const sub = try gpa.alloc(i16, tiles_k * per * words); + defer gpa.free(sub); + const vals = try gpa.alloc(f64, h.k * bn); + defer gpa.free(vals); + const scratch = try gpa.alloc(f64, 2 * h.k * bn); + defer gpa.free(scratch); + const suh = try gpa.alloc(f16, h.k); + defer gpa.free(suh); + for (suh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.suh[2 * i ..][0..2], .little)); + var svh: [fmt.hadamard_dim]f16 = undefined; + while (true) { + const nb = w.next.fetchAdd(1, .monotonic); + if (nb >= h.n / bn) return; + for (0..tiles_k) |tk| for (0..per) |j| { + const src = ((tk * tiles_n) + nb * per + j) * words * 2; + const dst = sub[(tk * per + j) * words ..][0..words]; + for (dst, 0..) |*d, i| d.* = std.mem.readInt(i16, h.trellis[src + 2 * i ..][0..2], .little); + }; + for (&svh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.svh[2 * (nb * bn + i) ..][0..2], .little)); + fmt.dequantize(.{ .halves = h.halves }, .mul1, sub, tiles_k, per, suh, &svh, vals, scratch); + for (0..w.rows) |r| { + const xr = w.x[r * h.k ..][0..h.k]; + var acc: [fmt.hadamard_dim]f64 = @splat(0); + for (0..h.k) |k| { + const xv = xr[k]; + const row = vals[k * bn ..][0..bn]; + for (&acc, row) |*a, v| a.* += xv * v; + } + @memcpy(w.out[r * h.n + nb * bn ..][0..bn], &acc); + } + } +} + +fn headF64(gpa: std.mem.Allocator, raw: *const Raw, x: []const f64, rows: usize) ![]f64 { + const head = try flash.exl3_mtp.Head.of(try raw.get("lm_head.trellis"), try raw.get("lm_head.suh"), try raw.get("lm_head.svh"), raw.get("lm_head.mul1") catch null); + const out = try gpa.alloc(f64, rows * head.n); + errdefer gpa.free(out); + var w: HeadWork = .{ .head = head, .x = x, .rows = rows, .out = out }; + var pool: [64]std.Thread = undefined; + const n = @min(pool.len, @max(1, std.Thread.getCpuCount() catch 1)); + var spawned: usize = 0; + for (0..n) |_| { + pool[spawned] = try std.Thread.spawn(.{}, headWorker, .{ &w, gpa }); + spawned += 1; + } + for (pool[0..spawned]) |t| t.join(); + if (w.failed) |err| return err; + return out; +} + +// -- modes --------------------------------------------------------------------------------------------------------- + +fn tfMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const p = golden.prompts()[0]; + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const tf = p.object.get("teacher_forced").?.array.items; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + try g.reset(); + const n = ids.len; + const ops = g.e.layers[0].exl3.?.ops; + const before = ops.fingerprint(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const t1 = check.now(gpu.io); + const logits = try a.alloc(u16, n * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits)); + // the head rows the pass fed lm_head (the final mixer's read-out, bf16 [n, dims]) + const hrows = try a.alloc(u16, n * dims); + defer a.free(hrows); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(hrows.ptr, g.e.batch_s.?.mixed, n * dims * 2), "cuMemcpyDtoH"); + std.debug.print("teacher-forced pass: {d} rows in {d:.1} ms (first pass: includes every miss's copy)\n", .{ n, ms(t0, t1) }); + // one decode round too, then every EXL3 scratch buffer must be where Ops.reserve put it + { + var tok = prompt.argmax(g.seq.logits); + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const after = ops.fingerprint(); + try check.expect(std.mem.eql(u8, std.mem.sliceAsBytes(&before), std.mem.sliceAsBytes(&after)), "an EXL3 scratch buffer grew inside a pass", .{}); + check.pass("EXL3 scratch reserved up front: the {d} scratch buffers kept their addresses and sizes through a {d}-row prompt pass and a decode round", .{ blk: { + var k: usize = 0; + for (after) |e| k += @intFromBool(e[1] != 0); + break :blk k; + }, n }); + + std.debug.print(" pos input ours top1 (logit) golden top1 (logit) match top5 overlap max|dlogit| on golden top5\n", .{}); + var top1_ok: usize = 0; + var top5_same: usize = 0; + var worst: f64 = 0; + var sum_abs: f64 = 0; + var count: usize = 0; + for (tf, 0..) |row, pos| { + const r = logits[pos * vocab ..][0..vocab]; + const ours = top5(r); + const g5 = try ints(a, row.object.get("top5_ids").?); + defer a.free(g5); + const gl = row.object.get("top5_logits").?.array.items; + var overlap: usize = 0; + for (ours) |o| for (g5) |x| { + overlap += @intFromBool(o == x); + }; + var dmax: f64 = 0; + for (g5, gl) |id, lv| { + const d = @abs(@as(f64, promote(r[id])) - num(lv)); + dmax = @max(dmax, d); + sum_abs += d; + count += 1; + } + worst = @max(worst, dmax); + const ok = ours[0] == g5[0]; + top1_ok += @intFromBool(ok); + top5_same += @intFromBool(std.mem.eql(u32, &ours, g5)); + std.debug.print("{d:4} {d:7} {d:9} ({d:7.3}) {d:9} ({d:7.3}) {s:5} {d}/5 {s} {d:.4}\n", .{ pos, ids[pos], ours[0], promote(r[ours[0]]), g5[0], num(gl[0]), if (ok) "yes" else "NO", overlap, if (std.mem.eql(u32, &ours, g5)) "(same order)" else " ", dmax }); + } + std.debug.print("TF top1 {d}/{d}, top5 identical in order {d}/{d}, |our logit - golden bf16 logit| on the golden's top-5 ids: max {d:.4}, mean {d:.4}\n", .{ top1_ok, tf.len, top5_same, tf.len, worst, sum_abs / @as(f64, @floatFromInt(count)) }); + + // the EXL3 lm_head itself, on the same head rows, against an exact float64 product over all 248,320 columns + var raw = try Raw.open(gpu, dir); + defer raw.close(gpu.io); + const xf = try a.alloc(f64, n * dims); + defer a.free(xf); + for (xf, hrows) |*o, b| o.* = promote(b); + const t2 = check.now(gpu.io); + const exact = try headF64(a, &raw, xf, n); + defer a.free(exact); + const t3 = check.now(gpu.io); + var head_abs: f64 = 0; + var head_rel: f64 = 0; + var head_top1: usize = 0; + for (0..n) |r| { + var peak: f64 = 0; + var best: usize = 0; + var dmax: f64 = 0; + for (0..vocab) |v| { + const want = exact[r * vocab + v]; + peak = @max(peak, @abs(want)); + if (want > exact[r * vocab + best]) best = v; + dmax = @max(dmax, @abs(@as(f64, promote(logits[r * vocab + v])) - want)); + } + head_abs = @max(head_abs, dmax); + head_rel = @max(head_rel, dmax / peak); + const ours = top5(logits[r * vocab ..][0..vocab])[0]; + // a bf16 tie at the top is no disagreement + head_top1 += @intFromBool(ours == best or promote(logits[r * vocab + ours]) == promote(logits[r * vocab + best])); + } + std.debug.print("EXL3 lm_head (Dense, bf16 out) vs exact float64 on the pass's {d} head rows, {d} columns ({d:.1} s on the host): max |d| {d:.4} (rel to the row peak {e:.2}), top-1 equal on {d}/{d}\n", .{ n, vocab, ms(t2, t3) / 1e3, head_abs, head_rel, head_top1, n }); + // bf16 output rounding is 2^-9 of a value (0.06 at a logit of 30); the fp16 rotated input adds ~2.4e-4 of the peak + try check.expect(head_rel < 4e-3 and head_top1 == n, "lm_head within bf16 rounding of float64", .{}); + check.pass("lm_head: EXL3 hook through exl3.Dense, {d} rows, max |d| {d:.4} (rel {e:.2}) vs float64, top-1 equal on every row", .{ n, head_abs, head_rel }); + try check.expect(top1_ok == tf.len, "teacher-forced top-1 equals the golden on every position ({d}/{d})", .{ top1_ok, tf.len }); + check.pass("teacher-forced top-1 {d}/{d} vs the golden; max |dlogit| {d:.4} on the golden's top-5", .{ top1_ok, tf.len, worst }); + printPool(g, "after tf"); +} + +fn printPool(g: *Eng, what: []const u8) void { + const s = g.stats(); + const routed = s.hit_slots + s.miss_slots; + std.debug.print("pool {s}: {d} rounds, {d} passes, hit rate {d:.1}% ({d}/{d} routed slots), {d} copies, {d} evictions; host cut {d:.1} ms over {d} blocks, plan {d:.1} ms\n", .{ what, s.rounds, s.passes, 100 * @as(f64, @floatFromInt(s.hit_slots)) / @as(f64, @floatFromInt(@max(routed, 1))), s.hit_slots, routed, s.misses, s.evictions, @as(f64, @floatFromInt(xpass.pool_time.cut_ns)) / 1e6, xpass.pool_time.blocks, @as(f64, @floatFromInt(xpass.pool_time.plan_ns)) / 1e6 }); +} + +/// Greedy `steps` tokens after `ids`: the prefill's last row, then a decode round a token. +fn greedyRun(g: *Eng, ids: []const u32, out: []u32, step_ms: ?*f64) !void { + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + out[0] = prompt.argmax(g.seq.logits); + const t0 = check.now(g.gpu.io); + for (1..out.len) |i| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, out[i - 1 ..][0..1], g.out.ptr); + out[i] = prompt.argmax(g.seq.logits); + } + if (step_ms) |p| p.* = ms(t0, check.now(g.gpu.io)) / @as(f64, @floatFromInt(@max(out.len - 1, 1))); +} + +/// Each golden prompt followed by the golden's own 24 greedy ids, teacher-forced in one pass: our top-1 at every +/// continuation position against the golden's greedy id, with both margins (top-1 minus top-2). +fn tfcMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const cont = try ints(a, p.object.get("greedy_ids").?); + defer a.free(cont); + const steps = p.object.get("greedy_steps").?.array.items; + const all = try a.alloc(u32, ids.len + cont.len - 1); + defer a.free(all); + @memcpy(all[0..ids.len], ids); + @memcpy(all[ids.len..], cont[0 .. cont.len - 1]); + if (all.len > out_rows) return error.UnexpectedTensor; + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, all, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const logits = try a.alloc(u16, all.len * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits[0 .. all.len * vocab])); + std.debug.print("TFC {s}\n step golden id ours top1 our margin golden margin max|dlogit| golden top5\n", .{p.object.get("name").?.string}); + var ok_n: usize = 0; + for (cont, 0..) |want, k| { + const r = logits[(ids.len - 1 + k) * vocab ..][0..vocab]; + const t = top5(r); + const margin = promote(r[t[0]]) - promote(r[t[1]]); + const g5 = try ints(a, steps[k].object.get("top5_ids").?); + defer a.free(g5); + var dmax: f64 = 0; + for (g5, steps[k].object.get("top5_logits").?.array.items) |id, lv| dmax = @max(dmax, @abs(@as(f64, promote(r[id])) - num(lv))); + const ok = t[0] == want; + ok_n += @intFromBool(ok); + std.debug.print(" {d:4} {d:10} {d:9} {s} {d:10.4} {d:14.4} {d:10.4}\n", .{ k, want, t[0], if (ok) " " else "*", margin, num(steps[k].object.get("margin12").?), dmax }); + } + std.debug.print("TFC {s}: top-1 equals the golden's greedy id on {d}/{d} continuation positions\n", .{ p.object.get("name").?.string, ok_n, cont.len }); + if (ok_n != cont.len) all_ok = false; + } + printPool(g, "after tfc"); + try check.expect(all_ok, "teacher-forced continuation top-1 equals the golden everywhere", .{}); + check.pass("teacher-forced continuation: top-1 equals the golden's greedy id at every step of all {d} prompts", .{golden.prompts().len}); +} + +fn genMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const want = try ints(a, p.object.get("greedy_ids").?); + defer a.free(want); + const got = try a.alloc(u32, want.len); + defer a.free(got); + var step: f64 = 0; + try greedyRun(g, ids, got, &step); + var first_diff: ?usize = null; + for (got, want, 0..) |x, y, i| if (x != y and first_diff == null) { + first_diff = i; + }; + std.debug.print("GEN {s}: ", .{p.object.get("name").?.string}); + for (got) |x| std.debug.print("{d},", .{x}); + std.debug.print(" {s} ({d:.1} ms/step decode)\n", .{ if (first_diff != null) "DIFFERS" else "equal to the golden", step }); + if (first_diff) |d| { + std.debug.print(" first difference at step {d}: ours {d}, golden {d}; golden margin there {d:.4}\n", .{ d, got[d], want[d], num(p.object.get("greedy_steps").?.array.items[d].object.get("margin12").?) }); + all_ok = false; + } + } + printPool(g, "after gen"); + try check.expect(all_ok, "greedy 24 equals the golden on every prompt", .{}); + check.pass("greedy 24 tokens equal the golden on all {d} prompts", .{golden.prompts().len}); +} + +fn benchMode(gpu: Gpu, dir: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + const max_n: usize = 2048; + const ids = try a.alloc(u32, max_n); + defer a.free(ids); + // text-like ids: a fixed pseudo-random walk over the common part of the vocabulary + var prng = std.Random.DefaultPrng.init(0x5eed); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + // warm-up: one pass to load kernels, settle the pool + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..64], .{ .logits = .last }); + for ([_]usize{ 16, 256, 2048 }) |n| { + var best: f64 = std.math.inf(f64); + for (0..2) |_| { + try g.reset(); + const before = g.stats(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..n], .{ .logits = .last }); + const t = ms(t0, check.now(gpu.io)); + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" prefill {d:5} tokens: {d:8.1} ms = {d:7.1} rows/s (hit {d:.1}%, {d} copies)\n", .{ n, t, @as(f64, @floatFromInt(n)) / (t / 1e3), 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), after.misses - before.misses }); + best = @min(best, t); + } + std.debug.print("BENCH prefill {d} tokens: best {d:.1} ms, {d:.1} rows/s\n", .{ n, best, @as(f64, @floatFromInt(n)) / (best / 1e3) }); + // decode after this prompt: 32 tokens, greedy + var tok = prompt.argmax(g.seq.logits); + const before = g.stats(); + const t0 = check.now(gpu.io); + const steps: usize = 32; + for (0..steps) |_| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const t = ms(t0, check.now(gpu.io)) / steps; + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print("BENCH decode after {d} tokens: {d:.2} ms/token over {d} steps (hit {d:.1}%, {d:.1} copies a token)\n", .{ n, t, steps, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), @as(f64, @floatFromInt(after.misses - before.misses)) / steps }); + } + printPool(g, "after bench"); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + const lineArg = struct { + fn f(r: []const [:0]const u8, i: usize) !?usize { + return if (r.len > i) try std.fmt.parseInt(usize, r[i], 10) else null; + } + }.f; + if (std.mem.eql(u8, mode, "tf")) return tfMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "tfc")) return tfcMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "gen")) return genMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "bench")) return benchMode(gpu, dir, try lineArg(rest, 2)); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3_graph_tests.zig b/zig/tests/cuda/flashnext_exl3_graph_tests.zig new file mode 100644 index 0000000000..aebc87acbf --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3_graph_tests.zig @@ -0,0 +1,1053 @@ +//! flashnext-exl3-graph: the EXL3 offload pool planned on the device (cuda_exl3_dplan.zig) and the decode step as one +//! CUDA graph (cuda_exl3_graph.zig). +//! +//! flashnext-exl3-graph plan synthetic pools: the device planner's lines, stamps, +//! counts and tables equal the host planner's after random +//! rounds, every resident line holds its expert's bytes; +//! the mapped-bank copy kernel's bandwidth vs the copy engine +//! flashnext-exl3-graph RANK WORLD PORT DIR eq GOLDEN [STEPS] [MODES] +//! the real pack at tp=1/2: each golden prompt prefilled, +//! then STEPS greedy decode rounds and the same continuation +//! as verify windows (rows 3,1,4,2,8,16,...) under each mode +//! (off = host cut, dplan = device plan eager, graph): +//! every step's logits bit-equal to off's, ms/token, hit rate +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const dplan = X.dplan; +const graph = X.graph; +const fpool = flash.exl3_pool; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; +const cap_default: usize = 32768; +const out_rows: usize = 64; + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +// -- synthetic pools: device planner == host planner ----------------------------------------------------------------- + +const syn_dims = 256; +const syn_width = 64; + +fn synPool(gpu: Gpu, routed: [][3][]const u8, lines: usize, max_rows: usize) !fpool.Pool { + return fpool.Pool.init(gpu.gpa, gpu.io, gpu.d, .{ .routed = routed }, .{ .suh_g = 0, .suh_u = 0, .svh_g = 0, .svh_u = 0, .suh_d = 0, .svh_d = 0 }, .{ + .layer = 0, + .lines = lines, + .dims = syn_dims, + .width = syn_width, + .max_rows = max_rows, + .slots = 11, + .direct = true, + .double_buffer = false, + }, null); +} + +fn planCase(gpu: Gpu, lines: usize, max_rows: usize, rounds: usize, seed: u64) !void { + const a = gpu.gpa; + const tiles = (syn_dims / 16) * (syn_width / 16); + // a pinned bank of 512 experts, widths 2..5 half-bits, random bytes + var sizes: [512][3]usize = undefined; + var prng = std.Random.DefaultPrng.init(seed); + const rnd = prng.random(); + var total: usize = 0; + for (&sizes) |*e| for (e) |*b| { + b.* = tiles * 16 * (2 + rnd.uintLessThan(usize, 4)); + total += (b.* + 255) / 256 * 256; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + rnd.bytes(bank.bytes); + const routed = try a.alloc([3][]const u8, 512); + defer a.free(routed); + var at: usize = 0; + for (routed, sizes) |*r, sz| for (0..3) |w| { + r[w] = bank.bytes[at..][0..sz[w]]; + at += (sz[w] + 255) / 256 * 256; + }; + var host = try synPool(gpu, routed, lines, max_rows); + defer host.deinit(); + var dev = try synPool(gpu, routed, lines, max_rows); + defer dev.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var pools = [_]*fpool.Pool{&dev}; + const dp = try dplan.DevPlan.init(a, gpu.d, &pools); + defer dp.deinit(); + var picks_dev = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * 11 * 4); + defer picks_dev.free(); + const picks = try a.alloc(i32, max_rows * 11); + defer a.free(picks); + // skewed routing: a hot set of 24 experts takes most picks, so rounds hit, miss and evict + var host_ns: u64 = 0; + for (0..rounds) |round| { + const rows = 1 + rnd.uintLessThan(usize, max_rows); + for (0..rows) |r| { + var k: usize = 0; + while (k < 10) { + const id: i32 = if (rnd.uintLessThan(u32, 4) != 0) @intCast(rnd.uintLessThan(u32, 24) * 7) else @intCast(rnd.uintLessThan(u32, 512)); + var dup = false; + for (picks[r * 11 ..][0..k]) |x| dup = dup or x == id; + if (dup) continue; + picks[r * 11 + k] = id; + k += 1; + } + picks[r * 11 + 10] = if (round % 3 == 0) -1 else 512; + } + const n = rows * 11; + // host + const t0 = check.now(gpu.io); + const rp = try host.plan(picks[0..n], rows); + host_ns += @intCast(check.now(gpu.io) - t0); + try check.expect(rp.n_passes == 1, "a synthetic round split into {d} passes", .{rp.n_passes}); + try host.stage(rp, 0, s); + try host.release(rp, 0, s); + // device + try dp.toDevice(s); + try picks_dev.upload(0, std.mem.sliceAsBytes(picks[0..n])); + try dp.plan(s, 0, picks_dev.ptr, rows); + try dp.join(s); + try dp.restore(s, 0); + if (round % 37 == 36) try dp.toHost(); // owner switches mid-stream keep both views consistent + } + try s.synchronize(); + try host.lines.drain(); + try dp.toHost(); + // residency, stamps, counts + var same_lines: usize = 0; + for (host.slots_tab, dev.slots_tab, host.stamp_tab, dev.stamp_tab) |hs, ds, ht, dt| same_lines += @intFromBool(hs == ds and ht == dt); + const same_counts = std.mem.eql(u32, &host.counts[0], &dev.counts[0]); + std.debug.print("RESULT plan lines={d} max_rows={d} rounds={d}: lines+stamps equal {d}/{d}, counts equal {}, host hits {d}/{d} copies {d} evictions {d}; device hits {d}/{d} copies {d} evictions {d}; host plan {d:.1} us a round\n", .{ lines, max_rows, rounds, same_lines, lines, same_counts, host.stats.hit_slots, host.stats.hit_slots + host.stats.miss_slots, host.stats.misses, host.stats.evictions, dev.stats.hit_slots, dev.stats.hit_slots + dev.stats.miss_slots, dev.stats.misses, dev.stats.evictions, @as(f64, @floatFromInt(host_ns)) / 1e3 / @as(f64, @floatFromInt(rounds)) }); + try check.expect(same_lines == lines and same_counts, "device planner diverged from the host planner", .{}); + try check.expect(host.stats.hit_slots == dev.stats.hit_slots and host.stats.misses == dev.stats.misses and host.stats.evictions == dev.stats.evictions, "device planner statistics differ", .{}); + // tables: the same line (relative to each pool's lines) and width for every expert + const tb = host.lines.tables.len; + const th = try a.alloc(u8, tb); + defer a.free(th); + const td = try a.alloc(u8, tb); + defer a.free(td); + try host.lines.tables.download(0, th); + try dev.lines.tables.download(0, td); + var bad_tab: usize = 0; + for (0..512) |id| for (0..3) |w| { + const ph = std.mem.readInt(u64, th[host.lines.tlay.ptr[w] + id * 8 ..][0..8], .little); + const pd = std.mem.readInt(u64, td[dev.lines.tlay.ptr[w] + id * 8 ..][0..8], .little); + const kh = std.mem.readInt(i32, th[host.lines.tlay.k2[w] + id * 4 ..][0..4], .little); + const kd = std.mem.readInt(i32, td[dev.lines.tlay.k2[w] + id * 4 ..][0..4], .little); + const rh = if (ph == 0) 0 else ph - host.lines.mem.ptr + 1; + const rd = if (pd == 0) 0 else pd - dev.lines.mem.ptr + 1; + bad_tab += @intFromBool(rh != rd or kh != kd); + }; + // every resident line holds its expert's bytes (the copy kernel read the mapped bank) + const mem = try a.alloc(u8, dev.lines.mem.len); + defer a.free(mem); + try dev.lines.mem.download(0, mem); + var bad_lines: usize = 0; + var resident: usize = 0; + for (dev.slots_tab, 0..) |id, ln| { + if (id == fpool.offload.Cache.empty_id) continue; + resident += 1; + for (0..3) |w| { + const want = routed[id][w]; + const got = mem[ln * dev.lines.layout.line_bytes + dev.lines.layout.off[w] ..][0..want.len]; + if (!std.mem.eql(u8, want, got)) bad_lines += 1; + } + } + std.debug.print("RESULT plan tables: {d} of 1536 entries differ; {d} resident lines, {d} projections with wrong bytes\n", .{ bad_tab, resident, bad_lines }); + try check.expect(bad_tab == 0 and bad_lines == 0, "device tables or line bytes wrong", .{}); + check.pass("device planner == host planner over {d} rounds of up to {d} rows on {d} lines (tables, lines, stamps, counts, line bytes)", .{ rounds, max_rows, lines }); +} + +fn copyBench(gpu: Gpu) !void { + // Flash Next line: gate/up/down 3-bit, K 2560 x N 640: 160 x 40 tiles of 48 bytes = 614,400 bytes each + const seg: usize = 614_400; + const n_lines: usize = 30; + var bank = try cuda.HostBuffer.alloc(gpu.d, n_lines * 3 * seg); + defer bank.free(); + @memset(bank.bytes, 0x5a); + var dst = try cuda.DeviceBuffer.alloc(gpu.d, n_lines * 3 * seg); + defer dst.free(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_dplan); + defer mod.unload(); + const f = try mod.function("fn_exl3_dcopy"); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + for ([_]usize{ 1, 3, 10 }) |k| { + // work list: k lines of three segments + var work: [10 * 9]i64 = undefined; + for (0..k) |i| for (0..3) |w| { + const off = (i * 3 + w) * seg; + work[i * 9 + w * 3 + 0] = @intCast(@intFromPtr(bank.bytes.ptr) + off); + work[i * 9 + w * 3 + 1] = @intCast(dst.ptr + off); + work[i * 9 + w * 3 + 2] = @intCast(seg); + }; + var wbuf = try cuda.DeviceBuffer.alloc(gpu.d, 10 * 9 * 8 + 16); + defer wbuf.free(); + try wbuf.upload(0, std.mem.sliceAsBytes(work[0 .. k * 9])); + const cnt = [1]i32{@intCast(k)}; + try wbuf.upload(10 * 9 * 8, std.mem.sliceAsBytes(&cnt)); + for ([_]u32{ 36, 72, 144, 288, 576 }) |blocks| { + const reps = 20; + var args: cuda.Args = .{}; + args.add(wbuf.ptr); + args.add(wbuf.ptr + 10 * 9 * 8); + try cuda.launch.launch(f, .{ .grid = .{ .x = blocks }, .block = .{ .x = 256 } }, s, &args); + try e0.record(s); + for (0..reps) |_| try cuda.launch.launch(f, .{ .grid = .{ .x = blocks }, .block = .{ .x = 256 } }, s, &args); + try e1.record(s); + try s.synchronize(); + const t = (try e0.elapsedMs(e1)) / reps; + std.debug.print("BENCH dcopy {d} lines ({d:.2} MiB) with {d} blocks: {d:.3} ms = {d:.1} GB/s\n", .{ k, @as(f64, @floatFromInt(k * 3 * seg)) / (1 << 20), blocks, t, @as(f64, @floatFromInt(k * 3 * seg)) / (t * 1e6) }); + } + // the copy engine, one H2D a segment + try e0.record(s); + for (0..20) |_| for (0..k * 3) |j| try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst.ptr + j * seg, bank.bytes[j * seg ..].ptr, seg, s.handle), "cuMemcpyHtoDAsync"); + try e1.record(s); + try s.synchronize(); + const t = (try e0.elapsedMs(e1)) / 20; + std.debug.print("BENCH dma {d} lines ({d:.2} MiB): {d:.3} ms = {d:.1} GB/s\n", .{ k, @as(f64, @floatFromInt(k * 3 * seg)) / (1 << 20), t, @as(f64, @floatFromInt(k * 3 * seg)) / (t * 1e6) }); + } +} + +fn planMode(gpu: Gpu) !void { + try copyBench(gpu); + try planCase(gpu, 10, 1, 400, 1); + try planCase(gpu, 47, 4, 600, 2); + try planCase(gpu, 63, 6, 600, 3); +} + +// -- the engine at tp=1/2 ------------------------------------------------------------------------------------------- + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !*Ranks { + const r = try gpu.gpa.create(Ranks); + errdefer gpu.gpa.destroy(r); + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 600_000); + errdefer link.close(); + r.* = .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + return r; + } + + fn close(r: *Ranks, gpa: std.mem.Allocator) void { + r.collective.deinit(); + r.link.close(); + gpa.destroy(r); + } + + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +const Eng = struct { + gpu: Gpu, + rank: usize, + world: usize, + ranks: ?*Ranks, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + attached: xpass.Attached, + + fn init(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8) !*Eng { + return initPaged(gpu, rank, world, port, dir, null); + } + + fn initPaged(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, paging: ?X.kvpage.Config) !*Eng { + const graph_reservation = if ((paging != null or X.kvpage.force != null or X.kvpage.configFromEnv() != null) and X.graph.enabled()) try X.graph.memory.startupReservationForStreams(1) else 0; + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.rank = rank; + g.world = world; + g.ranks = if (world == 2) try Ranks.open(gpu, rank, port) else null; + errdefer if (g.ranks) |ranks| ranks.close(gpu.gpa); + g.stream = try cuda.Stream.init(gpu.d, true); + errdefer g.stream.deinit(); + g.glue = try flash.Glue.load(gpu.d); + errdefer g.glue.deinit(); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + errdefer g.e.deinit(); + if (g.ranks) |r| g.e.collective = &r.collective; + if (paging) |cfg| { + g.e.page_config = cfg; + try X.kv.setup(g.e, false, true); + } + _ = try g.e.reserve(@min(cap_default, g.e.chunk)); + _ = try g.e.reserveScores(cap_default, @min(cap_default, g.e.chunk)); + const common_before = cuda.usage(false).device; + const common_bytes = try g.e.reserveSharedKv(cap_default); + const common_resident = try std.math.sub(u64, cuda.usage(false).device, common_before); + if (paging) |cfg| if (cfg.shared_lanes != 0 and rank == 1) { + const arena = g.e.shared_arena orelse return error.SharedKvArenaRequired; + const retired = try arena.register(cap_default); + try arena.release(retired, g.stream); + }; + g.seq = try g.e.newSeq(cap_default); + errdefer g.seq.deinit(gpu.gpa); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + errdefer g.out.free(); + X.prefill.reserveFor(g.e, cap_default, world); + var budget = try xpass.budgetFromEnv(0); + budget.graph_reservation = graph_reservation; + if (paging) |_| { + const layout = try X.kv.sequenceLayout(X.kv.of(g.e), cap_default, g.e.chunk, g.e.page_config); + budget.other = try std.math.add(u64, budget.other, try std.math.mul(u64, layout.bytes.device(), 2)); + if (g.e.page_config.?.shared_lanes != 0) { + const arena = g.e.shared_arena orelse return error.SharedKvArenaRequired; + try check.expect(common_bytes == arena.layout.common_bytes and common_bytes == common_resident and common_bytes != 0, "rank {d}: shared common bytes charged once before expert pools", .{rank}); + std.debug.print("SHARED rank {d}/{d}: lanes {d}, cap {d}, window {d}+rows {d}, physical window rows {d}, slots {d}x4, K/V side rows {d}, device {d}, tables {d}, work {d}, common {d} bytes resident before pools; future two lane caches {d} bytes\n", .{ rank, world, arena.layout.options.lanes, cap_default, arena.layout.options.window, arena.layout.options.rows, arena.layout.window_rows, arena.layout.options.slots, arena.layout.side_rows, arena.layout.device_bytes, arena.layout.table_bytes, arena.layout.work_bytes, common_bytes, budget.other }); + } else try check.expect(common_bytes == 0 and g.e.shared_arena == null, "rank {d}: private control has no shared allocation", .{rank}); + } + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, budget, null); + std.debug.print("RESULT rank {d}/{d} load: {d:.1} s; {d} lines a layer x {d} layers ({d:.2} GiB of pool)\n", .{ rank, world, ms(t0, check.now(gpu.io)) / 1e3, g.attached.lines, g.attached.layers, gib(g.attached.device_bytes) }); + return g; + } + + fn deinit(g: *Eng) void { + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + if (g.ranks) |r| r.close(g.gpu.gpa); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn swap(g: *Eng, v: u64) !u64 { + return if (g.ranks) |r| r.swap(v) else v; + } + + fn forward(g: *Eng, ids: []const u32, hooks: prompt.batch.Hooks) !void { + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, hooks); + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + } + + fn round(g: *Eng, tok: u32) !u32 { + const q = [_]*prompt.Seq{&g.seq}; + var t = tok; + try g.e.round(flash.Tri, g.tri(), &q, (&t)[0..1], g.out.ptr); + if (g.ranks) |r| try r.collective.finish(g.stream); + return prompt.argmax(g.seq.logits); + } + + fn stats(g: *Eng) !fpool.Stats { + try dplan.toHost(); + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } +}; + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +const modes = [_]graph.Mode{ .off, .dplan, .graph }; + +const Run = struct { + ids: []u32, + /// Per decode step, the hash of the step's whole bf16 logits row; then per window row the same. + round_h: []u64, + window_h: []u64, + ms_token: f64 = 0, + ms_window_row: f64 = 0, + hit: f64 = 0, + copies: f64 = 0, +}; + +const win_sizes = [_]usize{ 3, 1, 4, 2, 8, 16 }; + +fn runMode(g: *Eng, ids: []const u32, steps: usize, r: *Run) !void { + const a = g.gpu.gpa; + // decode rounds + try dplan.toHost(); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var tok = prompt.argmax(g.seq.logits); + r.ids[0] = tok; + _ = try g.swap(1); + const before = try g.stats(); + const t0 = check.now(g.gpu.io); + for (0..steps) |i| { + tok = try g.round(tok); + r.ids[i + 1] = tok; + r.round_h[i] = std.hash.Wyhash.hash(0, g.seq.logits); + } + const t1 = check.now(g.gpu.io); + const after = try g.stats(); + r.ms_token = ms(t0, t1) / @as(f64, @floatFromInt(steps)); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + r.hit = 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))); + r.copies = @as(f64, @floatFromInt(after.misses - before.misses)) / @as(f64, @floatFromInt(steps)); + // the same continuation teacher-forced in windows of 3, 1, 4, 2 rows (every row's logits) + try dplan.toHost(); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var at: usize = 0; + var k: usize = 0; + const host = try a.alloc(u16, 16 * vocab); + defer a.free(host); + const tw = check.now(g.gpu.io); + while (at < steps) : (k += 1) { + const n = @min(win_sizes[k % win_sizes.len], steps - at); + try g.forward(r.ids[at..][0..n], .{ .logits = .all, .logits_out = g.out.ptr }); + try g.out.download(0, std.mem.sliceAsBytes(host[0 .. n * vocab])); + for (0..n) |j| r.window_h[at + j] = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(host[j * vocab ..][0..vocab])); + at += n; + } + r.ms_window_row = ms(tw, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)); +} + +/// Two sequences in one decode round (two rows, each its own caches): `steps` rounds after both prompts, every round's +/// two logits rows hashed into `out`; ms per round. +fn pairRun(g: *Eng, b: *prompt.Seq, ids_a: []const u32, ids_b: []const u32, steps: usize, out: []u64) !f64 { + try dplan.toHost(); + try g.reset(); + try g.forward(ids_a, .{ .logits = .last }); + for ([_][]cuda.DeviceBuffer{ b.conv, b.state, b.kc, b.vc, b.ikc, b.pooled }) |bs| for (bs) |x| try x.fill8(0, g.stream.handle); + try b.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + b.fed = 0; + b.history = .{ g.e.eos, g.e.eos }; + try g.e.forward(flash.Tri, g.tri(), b, ids_b, .{ .logits = .last }); + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + var t = [2]u32{ prompt.argmax(g.seq.logits), prompt.argmax(b.logits) }; + const q = [_]*prompt.Seq{ &g.seq, b }; + const t0 = check.now(g.gpu.io); + for (0..steps) |i| { + try g.e.round(flash.Tri, g.tri(), &q, &t, g.out.ptr); + if (g.ranks) |r| try r.collective.finish(g.stream); + out[i] = std.hash.Wyhash.hash(0, g.seq.logits) ^ std.hash.Wyhash.hash(7, b.logits); + t = .{ prompt.argmax(g.seq.logits), prompt.argmax(b.logits) }; + } + return ms(t0, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)); +} + +fn eqMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, steps: usize, which: []const u8) !void { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, golden_path, a, .limited(1 << 26)); + defer a.free(text); + var parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + var all_ok = true; + for (parsed.value.object.get("prompts").?.array.items) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + var runs: [modes.len]Run = undefined; + for (&runs) |*r| r.* = .{ .ids = try a.alloc(u32, steps + 1), .round_h = try a.alloc(u64, steps), .window_h = try a.alloc(u64, steps) }; + defer for (&runs) |*r| { + a.free(r.ids); + a.free(r.round_h); + a.free(r.window_h); + }; + // warm the pools once (first-touch copies are not any mode's cost) + graph.force = .off; + try runMode(g, ids, steps, &runs[0]); + for (modes, 0..) |m, mi| { + if (std.mem.indexOf(u8, which, @tagName(m)) == null and m != .off) continue; + graph.force = m; + const g0 = graph.stats; + try runMode(g, ids, steps, &runs[mi]); + const g1 = graph.stats; + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(runs[mi].round_h)) ^ std.hash.Wyhash.hash(1, std.mem.sliceAsBytes(runs[mi].window_h)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "rank {d}: the two ranks' logits differ ({s})", .{ rank, @tagName(m) }); + const rs = runs[mi]; + var same_r: usize = 0; + var same_w: usize = 0; + for (rs.round_h, runs[0].round_h) |x, y| same_r += @intFromBool(x == y); + for (rs.window_h, runs[0].window_h) |x, y| same_w += @intFromBool(x == y); + const ok = same_r == steps and same_w == steps and std.mem.eql(u32, rs.ids, runs[0].ids); + if (mi != 0) all_ok = all_ok and ok; + if (rank == 0) std.debug.print("RESULT tp={d} {s} {s}: decode {d:.2} ms/token (hit {d:.1}%, {d:.1} copies a token), windows {d:.2} ms/row; logits bit-equal to off: rounds {d}/{d}, window rows {d}/{d}; greedy ids {s}; graph captures {d} (updates {d}) replays {d}, eager-dplan steps {d}, declined {d}, capture {d:.1} ms\n", .{ world, p.object.get("name").?.string, @tagName(m), rs.ms_token, rs.hit, rs.copies, rs.ms_window_row, same_r, steps, same_w, steps, if (std.mem.eql(u32, rs.ids, runs[0].ids)) "equal" else "DIFFER", g1.captures - g0.captures, g1.updates - g0.updates, g1.replays - g0.replays, g1.eager_dplan - g0.eager_dplan, g1.declined - g0.declined, @as(f64, @floatFromInt(g1.capture_ns - g0.capture_ns)) / 1e6 }); + } + graph.force = null; + } + // two streams in one round + { + const ps = parsed.value.object.get("prompts").?.array.items; + const ia = try ints(a, ps[0].object.get("ids").?); + defer a.free(ia); + const ib = try ints(a, ps[ps.len - 1].object.get("ids").?); + defer a.free(ib); + var b = try g.e.newSeq(cap_default); + defer b.deinit(a); + const pair_steps = @min(steps, 32); + var ref: [64]u64 = undefined; + var got: [64]u64 = undefined; + graph.force = .off; + _ = try pairRun(g, &b, ia, ib, pair_steps, &ref); + for (modes) |m| { + if (std.mem.indexOf(u8, which, @tagName(m)) == null and m != .off) continue; + graph.force = m; + const t = try pairRun(g, &b, ia, ib, pair_steps, &got); + var same: usize = 0; + for (got[0..pair_steps], ref[0..pair_steps]) |x, y| same += @intFromBool(x == y); + if (m != .off) all_ok = all_ok and same == pair_steps; + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(got[0..pair_steps])); + if (world == 2) try check.expect(try g.swap(h) == h, "rank {d}: the two ranks' two-stream logits differ", .{rank}); + if (rank == 0) std.debug.print("RESULT tp={d} two streams {s}: {d:.2} ms a round of 2 rows; both rows' logits bit-equal to off in {d}/{d} rounds\n", .{ world, @tagName(m), t, same, pair_steps }); + } + graph.force = null; + } + if (dplan.current()) |dp| if (rank == 0) std.debug.print("RESULT dplan owner switches: {d} pushes, {d} pulls; {d} experts read in place (rounds with more misses than free lines)\n", .{ dp.pushes, dp.pulls, dp.inplace }); + try check.expect(all_ok, "rank {d}: a device-plan or graph step's logits differ from the host cut's", .{rank}); + check.pass("rank {d}/{d}: decode rounds and verify windows on the device plan and in graphs are bit-identical to the host cut ({d} steps a prompt)", .{ rank, world, steps }); +} + +/// A long prompt through the deep prefill (TF_FLASHNEXT_EXL3_PREFILL_ROWS sets its chunks), then `steps` greedy +/// rounds, twice per mode; every run's last-row logits and decode logits hashed and compared with the first eager run. +fn prefillMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, n: usize, steps: usize) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const ids = try a.alloc(u32, n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + const hs = try a.alloc(u64, steps + 1); + defer a.free(hs); + var ref: ?u64 = null; + var ok = true; + for ([_]graph.Mode{ .off, .off, .graph, .graph, .graph }) |m| { + graph.force = m; + try dplan.toHost(); + try g.reset(); + _ = try g.swap(3); + const p0 = graph.prefill_stats; + const f0 = X.prefill.stats; + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .last }); + const t1 = check.now(gpu.io); + const p1 = graph.prefill_stats; + const f1 = X.prefill.stats; + try check.expect(f1.chunks > f0.chunks and f1.rows >= f0.rows + n, "rank {d}: deep prefill ran instead of falling back", .{rank}); + if (m == .graph) try check.expect(p1.captures + p1.updates + p1.replays > p0.captures + p0.updates + p0.replays, "rank {d}: prefill graph captured or replayed a layer", .{rank}); + hs[0] = std.hash.Wyhash.hash(0, g.seq.logits); + var tok = prompt.argmax(g.seq.logits); + for (0..steps) |i| { + tok = try g.round(tok); + hs[i + 1] = std.hash.Wyhash.hash(0, g.seq.logits); + } + if (g.seq.paged) |pg| { + try g.stream.synchronize(); + const pager = try pg.total(); + try check.expect(!pager.overflow and pager.epoch > 0 and pager.calls > 0 and pager.rebases > 0, "rank {d}: paged prefill/decode advances epochs and crosses a window without overflow", .{rank}); + std.debug.print("PAGER rank {d}/{d} {s}: device epoch {d}, calls {d}, rebases {d}, written rows {d}, misses {d}\n", .{ rank, world, @tagName(m), pager.epoch, pager.calls, pager.rebases, pager.written_rows, pager.misses }); + } + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(hs)); + if (world == 2) try check.expect(try g.swap(h) == h, "rank {d}: the ranks' prefill logits differ", .{rank}); + if (ref == null) ref = h; + const same = ref.? == h; + ok = ok and same; + if (rank == 0) std.debug.print("RESULT tp={d} prefill {s}: {d} rows in {d:.1} ms = {d:.1} rows/s; layer graphs captured {d} (updates {d}) replayed {d} eager {d} (capture {d:.1} ms); last-row + {d} decode logits {s} the first eager run\n", .{ world, @tagName(m), n, ms(t0, t1), @as(f64, @floatFromInt(n)) / (ms(t0, t1) / 1e3), p1.captures - p0.captures, p1.updates - p0.updates, p1.replays - p0.replays, p1.eager - p0.eager, @as(f64, @floatFromInt(p1.capture_ns - p0.capture_ns)) / 1e6, steps, if (same) "bit-equal to" else "DIFFER from" }); + } + graph.force = null; + try check.expect(ok, "rank {d}: prefill under graphs differs from eager", .{rank}); + check.pass("rank {d}/{d}: deep prefill of {d} rows with layer graphs and the decode after it bit-identical to eager", .{ rank, world, n }); +} + +/// One warmed TP2 deep-prefill forward inside a CUDA-profiler API range for Nsight Systems. +/// The profiler window excludes graph warm-up and host setup; no per-layer synchronization is added. +fn prefillProfMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, n: usize, which: []const u8) !void { + var stage: []const u8 = "engine initialization"; + errdefer std.debug.print("prefill-prof rank={d} failed during {s}\n", .{ rank, stage }); + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + defer graph.force = null; + try check.expect(g.seq.paged != null, "prefill profiler requires BF16 KV paging", .{}); + const ids = try a.alloc(u32, n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + const use_graphs = std.mem.eql(u8, which, "graph"); + if (!use_graphs and !std.mem.eql(u8, which, "eager")) return error.BadMode; + graph.force = if (use_graphs) .graph else .off; + try dplan.toHost(); + stage = "graph warm-up"; + // Warm paging and, for graph mode, capture each layer shape before the measured forward. + for (0..if (use_graphs) @as(usize, 2) else 1) |_| { + try g.reset(); + _ = try g.swap(3); + try g.forward(ids, .{ .logits = .last }); + try g.stream.synchronize(); + } + const before = graph.prefill_stats; + const f_before = X.prefill.stats; + stage = "profile setup"; + try g.reset(); + _ = try g.swap(3); + try g.stream.synchronize(); + var prof = try Profiler.open(); + _ = prof.start(); + stage = "measured forward"; + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .last }); + const t1 = check.now(gpu.io); + _ = prof.stop(); + try g.stream.synchronize(); + const after = graph.prefill_stats; + const f_after = X.prefill.stats; + if (use_graphs and after.replays == before.replays) return error.NoPrefillGraphReplay; + const hash = std.hash.Wyhash.hash(0, g.seq.logits); + stage = "rank hash check"; + if (world == 2) try check.expect(try g.swap(hash) == hash, "rank {d}: profiled prefill logits differ across ranks", .{rank}); + std.debug.print("RESULT prefill-prof rank={d}/{d} mode={s} rows={d}: {d:.1} ms wall, graphs replayed {d}, eager {d}, captured {d}, updated {d}; expert copies {d} (prefetch {d}), passes {d}, pick wait {d:.1} ms, plan {d:.1} ms; logits {x}; Nsight range excludes warm-up\n", .{ rank, world, which, n, ms(t0, t1), after.replays - before.replays, after.eager - before.eager, after.captures - before.captures, after.updates - before.updates, f_after.copies - f_before.copies, f_after.prefetched - f_before.prefetched, f_after.passes - f_before.passes, @as(f64, @floatFromInt(f_after.host_wait_ns - f_before.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(f_after.plan_ns - f_before.plan_ns)) / 1e6, hash }); +} + +fn expectSnapshotState(rank: usize, previous: *const flash.snapshot.Snap, saved: *const flash.snapshot.Snap) !void { + const canonical = saved.paged.?.layout.canonical; + const state_offset = saved.paged.?.layout.total; + try check.expect(previous.at == saved.at and std.meta.eql(previous.history, saved.history), "rank {d}: exact snapshot position and history", .{rank}); + try check.expect(std.mem.eql(u8, previous.host.?.bytes[0..canonical], saved.host.?.bytes[0..canonical]), "rank {d}: exact canonical BF16 history", .{rank}); + try check.expect(std.mem.eql(u8, previous.host.?.bytes[state_offset..previous.buf.len], saved.host.?.bytes[state_offset..saved.buf.len]), "rank {d}: exact recurrent/index/PLE state and logits", .{rank}); +} + +fn expectOwnedRestore(engine: *Eng, sequence: *prompt.Seq, wanted: *const flash.snapshot.Snap) !void { + const saved = try flash.snapshot.save(engine.gpu.gpa, engine.gpu.d, engine.stream, sequence, null); + defer flash.snapshot.drop(engine.gpu.gpa, engine.stream, saved); + try expectSnapshotState(engine.rank, wanted, saved); + try check.expect(std.meta.eql(wanted.paged, saved.paged) and std.mem.eql(u8, wanted.host.?.bytes[0..wanted.buf.len], saved.host.?.bytes[0..saved.buf.len]), "rank {d}: exact destination windows, marks, controls and snapshot metadata", .{engine.rank}); + const previous = wanted.owned_blocks orelse return error.SharedKvSnapshotRequired; + const current = saved.owned_blocks orelse return error.SharedKvSnapshotRequired; + try check.expect(previous.entries.len > 0 and previous.entries.len == current.entries.len, "rank {d}: restore retains every owned resident entry, not a cold cache", .{engine.rank}); + try check.expect(std.meta.eql(current.source, sequence.paged.?.shared_sequence), "rank {d}: restored slots belong to destination generation", .{engine.rank}); + const block_bytes = try previous.geometry.blockBytes(); + for (previous.entries, 0..) |entry, old_index| { + const new_index = for (current.entries, 0..) |candidate, index| { + if (std.meta.eql(entry.key, candidate.key)) break index; + } else return error.SharedKvSnapshotMissingEntry; + try check.expect(previous.clock - entry.stamp == current.clock - current.entries[new_index].stamp, "rank {d}: owned entry retains relative LRU age", .{engine.rank}); + try check.expect(std.mem.eql(u8, previous.payload[old_index * block_bytes ..][0..block_bytes], current.payload[new_index * block_bytes ..][0..block_bytes]), "rank {d}: owned entry remaps exact K/V payload", .{engine.rank}); + } +} + +const SharedSnapshotProbe = struct { + sequence: prompt.Seq, + saved: *flash.snapshot.Snap, + continued: ?*flash.snapshot.Snap = null, + key: [6]u64, + buffers: [3]u64, + + fn init(engine: *Eng, ids: []const u32) !*SharedSnapshotProbe { + const probe = try engine.gpu.gpa.create(SharedSnapshotProbe); + errdefer engine.gpu.gpa.destroy(probe); + probe.* = .{ .sequence = try engine.e.newSeq(cap_default), .saved = undefined, .key = undefined, .buffers = undefined }; + errdefer probe.sequence.deinit(engine.gpu.gpa); + try engine.e.forward(flash.Tri, engine.tri(), &probe.sequence, ids, .{ .logits = .last }); + probe.saved = try flash.snapshot.save(engine.gpu.gpa, engine.gpu.d, engine.stream, &probe.sequence, null); + errdefer flash.snapshot.drop(engine.gpu.gpa, engine.stream, probe.saved); + const arena = engine.e.shared_arena orelse return error.SharedKvArenaRequired; + try check.expect(probe.saved.version == 2 and probe.saved.owned_blocks.?.entries.len > 0, "rank {d}: foreign probe holds shared resident blocks", .{engine.rank}); + probe.key = probe.sequence.paged.?.graphKey(); + probe.buffers = .{ arena.device.ptr, arena.tables.ptr, arena.work.ptr }; + return probe; + } + + fn deinit(probe: *SharedSnapshotProbe, engine: *Eng) void { + if (probe.continued) |saved| flash.snapshot.drop(engine.gpu.gpa, engine.stream, saved); + flash.snapshot.drop(engine.gpu.gpa, engine.stream, probe.saved); + probe.sequence.deinit(engine.gpu.gpa); + engine.gpu.gpa.destroy(probe); + } + + fn unchanged(probe: *SharedSnapshotProbe, engine: *Eng) !void { + const arena = engine.e.shared_arena.?; + try check.expect(std.meta.eql(probe.buffers, [3]u64{ arena.device.ptr, arena.tables.ptr, arena.work.ptr }), "rank {d}: arena addresses survive restore and sequence churn", .{engine.rank}); + try check.expect(std.meta.eql(probe.key, probe.sequence.paged.?.graphKey()), "rank {d}: foreign graph bindings remain legal after restore", .{engine.rank}); + const saved = try flash.snapshot.save(engine.gpu.gpa, engine.gpu.d, engine.stream, &probe.sequence, null); + defer flash.snapshot.drop(engine.gpu.gpa, engine.stream, saved); + try check.expect(std.meta.eql(probe.saved.paged, saved.paged), "rank {d}: foreign window origins and counters unchanged", .{engine.rank}); + try check.expect(probe.saved.at == saved.at and std.meta.eql(probe.saved.history, saved.history), "rank {d}: foreign logical position unchanged", .{engine.rank}); + try check.expect(std.mem.eql(u8, probe.saved.host.?.bytes[0..probe.saved.buf.len], saved.host.?.bytes[0..saved.buf.len]), "rank {d}: restore preserves foreign canonical history, full windows, marks, controls, ordinary state and logits", .{engine.rank}); + } + + fn continueAndRestore(probe: *SharedSnapshotProbe, engine: *Eng) !void { + try engine.e.forward(flash.Tri, engine.tri(), &probe.sequence, &.{ 417, 819, 1237 }, .{ .logits = .last }); + const saved = try flash.snapshot.save(engine.gpu.gpa, engine.gpu.d, engine.stream, &probe.sequence, null); + if (probe.continued) |previous| { + defer flash.snapshot.drop(engine.gpu.gpa, engine.stream, saved); + try expectSnapshotState(engine.rank, previous, saved); + } else probe.continued = saved; + try flash.snapshot.restore(engine.stream, &probe.sequence, null, probe.saved); + try expectOwnedRestore(engine, &probe.sequence, probe.saved); + try probe.unchanged(engine); + } +}; + +const SnapshotStorage = enum { legacy, private, shared }; + +fn snapshotStorageMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, steps: usize, storage: SnapshotStorage) !void { + if (world != 2) return error.SharedSnapshotRequiresTp2; + if (steps < 8 or steps > 64) return error.SharedSnapshotDecodeSteps; + if (X.kvpage.force_off) return error.SharedKvPagingDisabled; + var paging = X.kvpage.configFromEnv() orelse X.kvpage.Config{ .shared_lanes = 8 }; + if (storage == .shared) { + if (paging.shared_lanes < 3) return error.SharedSnapshotNeedsThreeLanes; + } else paging.shared_lanes = 0; + paging.slots = 1024; + paging.window = 4096; + paging.rows = 4096; + if ((X.prefill.options orelse X.prefill.optionsFromEnv()).serve_rows > paging.rows) return error.SharedSnapshotPrefillRows; + try paging.validate(); + const previous = X.kvpage.force; + X.kvpage.force = paging; + defer X.kvpage.force = previous; + return snapshotMode(gpu, rank, world, port, dir, steps, storage); +} + +fn snapshotMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, steps: usize, storage: SnapshotStorage) !void { + const snaps = flash.snapshot; + const allocator = gpu.gpa; + const shared = storage == .shared; + const paired = storage != .legacy; + const g = if (paired) try Eng.initPaged(gpu, rank, world, port, dir, X.kvpage.force.?) else try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + defer graph.force = null; + try check.expect(g.seq.paged != null, "snapshot test requires BF16 paging", .{}); + if (paired) try check.expect(g.seq.kv8 == null and g.seq.paged.?.cfg.slots == 1024 and g.seq.paged.?.cfg.rows == 4096, "rank {d}: paired fixture keeps fixed BF16 paging geometry", .{rank}); + if (shared) try check.expect(g.seq.paged.?.shared_arena == g.e.shared_arena and g.seq.paged.?.cfg.shared_lanes == X.kvpage.force.?.shared_lanes and g.seq.paged.?.cfg.slots == 1024 and g.seq.paged.?.cfg.rows == 4096, "rank {d}: shared snapshot fixture cannot silently use private or different paging geometry", .{rank}); + if (shared) { + const owner = g.seq.paged.?.shared_sequence; + const peer_generation = try g.swap(owner.generation); + try check.expect(owner.generation != peer_generation, "rank {d}: shared TP2 fixture requires legal cross-rank generation asymmetry", .{rank}); + const logical = g.seq.paged.?.logicalGraphKey(); + const logical_hash = std.hash.Wyhash.hash(0, std.mem.asBytes(&logical)); + try check.expect(try g.swap(logical_hash) == logical_hash, "rank {d}: asymmetric physical generations retain identical logical TP geometry", .{rank}); + std.debug.print("SHARED rank {d}/{d}: primary lane {d} generation {d}, peer generation {d}; logical TP key {x}\n", .{ rank, world, owner.lane, owner.generation, peer_generation, logical_hash }); + } + if (storage == .private) try check.expect(g.seq.paged.?.cfg.shared_lanes == 0 and g.seq.paged.?.shared_arena == null, "rank {d}: private oracle uses independent per-sequence slots", .{rank}); + var second = try g.e.newSeq(cap_default); + var second_live = true; + defer if (second_live) second.deinit(allocator); + const ids = try allocator.alloc(u32, if (paired) 24576 else 16384); + defer allocator.free(ids); + var random = std.Random.DefaultPrng.init(0x736e6170); + for (ids) |*token| token.* = 200 + random.random().uintLessThan(u32, 60000); + if (paired) std.debug.print("SNAPSHOT-ORACLE store={s} rank={d}/{d} seed=736e6170 ids={x:0>16} first=8707 second=4096+8709 continuation=8192 decode={d}; compare matching rank/run/lane hashes between snapshot-private and snapshot-shared, not a golden\n", .{ @tagName(storage), rank, world, std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(ids)), steps }); + graph.force = .graph; + const first_end: usize = if (paired) 8707 else 3075; + const second_end: usize = if (paired) 12805 else 6917; + try g.forward(ids[0..first_end], .{ .logits = .last }); + var first_saved: *snaps.Snap = undefined; + if (paired) first_saved = try snaps.save(allocator, gpu.d, g.stream, &g.seq, null); + defer if (paired) snaps.drop(allocator, g.stream, first_saved); + try g.e.forward(flash.Tri, g.tri(), &second, ids[4096..second_end], .{ .logits = .last }); + if (!paired) first_saved = try snaps.save(allocator, gpu.d, g.stream, &g.seq, null); + defer if (!paired) snaps.drop(allocator, g.stream, first_saved); + const second_saved = try snaps.save(allocator, gpu.d, g.stream, &second, null); + defer snaps.drop(allocator, g.stream, second_saved); + const source_key = second.paged.?.graphKey(); + const probe: ?*SharedSnapshotProbe = if (shared) try SharedSnapshotProbe.init(g, ids[14336..23045]) else null; + defer if (probe) |foreign| foreign.deinit(g); + const expected = try allocator.alloc(u8, (steps + 1) * 2 * vocab * 2); + defer allocator.free(expected); + var reference: [2]?*snaps.Snap = .{ null, null }; + defer for (reference) |item| if (item) |saved| snaps.drop(allocator, g.stream, saved); + for ([_]graph.Mode{ .off, .graph, .graph, .graph }, 0..) |mode, run_index| { + graph.force = mode; + try dplan.toHost(); + g.e.round_graph.clear(); + if (shared) { + const previous = second.paged.?.shared_sequence; + const window = second.paged.?.devRows(0, false); + const reuse_before = X.recycle.stats.seq_reused; + graph.force = .graph; + second_live = false; + second.deinit(allocator); + second = try g.e.newSeq(cap_default); + second_live = true; + graph.force = mode; + try check.expect(X.recycle.stats.seq_reused == reuse_before + 1, "rank {d}: request recycle reused the retained sequence allocations", .{rank}); + const current = second.paged.?.shared_sequence; + try check.expect(previous.arena == current.arena and previous.lane == current.lane and current.generation > previous.generation and current.generation > second_saved.owned_blocks.?.source.generation, "rank {d}: saved source restores into recycled lane with a new generation", .{rank}); + try check.expect(window == second.paged.?.devRows(0, false), "rank {d}: recycled lane retains its fixed window address", .{rank}); + try probe.?.unchanged(g); + std.debug.print("SHARED rank {d}/{d}: recycled lane {d}, sequence reuse count +1, source generation {d}, destination generation {d}, stable window {x}\n", .{ rank, world, current.lane, second_saved.owned_blocks.?.source.generation, current.generation, window }); + } + try snaps.restore(g.stream, &g.seq, null, first_saved); + if (shared) { + try expectOwnedRestore(g, &g.seq, first_saved); + try probe.?.unchanged(g); + } + try snaps.restore(g.stream, &second, null, second_saved); + if (shared) { + try expectOwnedRestore(g, &second, second_saved); + try check.expect(source_key[4] != second.paged.?.graphKey()[4], "rank {d}: identical restored geometry at same window address has a different strict graph generation key", .{rank}); + try probe.?.unchanged(g); + } + const before = graph.stats; + const prefill_before = graph.prefill_stats; + const deep_before = X.prefill.stats; + try g.forward(ids[first_end..][0..8192], .{ .logits = .last }); + try g.e.forward(flash.Tri, g.tri(), &second, ids[second_end..][0..8192], .{ .logits = .last }); + if (paired) try check.expect(X.prefill.stats.chunks > deep_before.chunks and X.prefill.stats.rows >= deep_before.rows + 16384, "rank {d}: both paired lanes actually execute deep prefill", .{rank}); + var tokens = [_]u32{ prompt.argmax(g.seq.logits), prompt.argmax(second.logits) }; + const sequences = [_]*prompt.Seq{ &g.seq, &second }; + for (0..steps + 1) |step| { + if (step > 0) { + try g.e.round(flash.Tri, g.tri(), &sequences, &tokens, g.out.ptr); + tokens = .{ prompt.argmax(g.seq.logits), prompt.argmax(second.logits) }; + } + for (sequences, 0..) |sequence, lane| { + const wanted = expected[(step * 2 + lane) * vocab * 2 ..][0 .. vocab * 2]; + if (run_index == 0) @memcpy(wanted, sequence.logits) else try check.expect(std.mem.eql(u8, wanted, sequence.logits), "rank {d} snapshot {s} lane {d} step {d} BF16 logits", .{ rank, @tagName(mode), lane, step }); + } + if (paired) { + for (sequences, 0..) |sequence, lane| std.debug.print("SNAPSHOT-ORACLE store={s} rank={d}/{d} mode={s} run={d} lane={d} step={d} logits={x:0>16}\n", .{ @tagName(storage), rank, world, @tagName(mode), run_index, lane, step, std.hash.Wyhash.hash(0, sequence.logits) }); + const digest = std.hash.Wyhash.hash(0, g.seq.logits) ^ std.hash.Wyhash.hash(1, second.logits); + try check.expect(try g.swap(digest) == digest, "rank {d}: paired TP2 logits agree at step {d}", .{ rank, step }); + } + } + for (sequences, 0..) |sequence, lane| { + const saved = try snaps.save(allocator, gpu.d, g.stream, sequence, null); + if (paired) { + const canonical_hash = std.hash.Wyhash.hash(0, saved.host.?.bytes[0..saved.paged.?.layout.canonical]); + const state_hash = std.hash.Wyhash.hash(1, saved.host.?.bytes[saved.paged.?.layout.total..saved.buf.len]); + const history_hash = std.hash.Wyhash.hash(2, std.mem.asBytes(&saved.history)) ^ std.hash.Wyhash.hash(3, std.mem.asBytes(&saved.at)); + std.debug.print("SNAPSHOT-ORACLE store={s} rank={d}/{d} mode={s} run={d} lane={d} at={d} canonical={x:0>16} state-logits={x:0>16} history={x:0>16} full={x:0>16}\n", .{ @tagName(storage), rank, world, @tagName(mode), run_index, lane, saved.at, canonical_hash, state_hash, history_hash, canonical_hash ^ state_hash ^ history_hash }); + } + if (run_index == 0) { + reference[lane] = saved; + } else { + defer snaps.drop(allocator, g.stream, saved); + const previous = reference[lane].?; + const canonical = saved.paged.?.layout.canonical; + const state_offset = saved.paged.?.layout.total; + try check.expect(std.mem.eql(u8, previous.host.?.bytes[0..canonical], saved.host.?.bytes[0..canonical]), "rank {d} snapshot lane {d} exact canonical BF16 history", .{ rank, lane }); + try check.expect(std.mem.eql(u8, previous.host.?.bytes[state_offset..previous.buf.len], saved.host.?.bytes[state_offset..saved.buf.len]), "rank {d} snapshot lane {d} exact recurrent/index/PLE state", .{ rank, lane }); + } + } + if (mode == .graph) { + std.debug.print("SNAPSHOT rank {d}/{d} graph cache run={d}: prefill captures +{d}, updates +{d}, replays +{d}, paged replays +{d}, rebinds +{d}, sighting hits +{d}, misses +{d}, eager +{d}; decode replays +{d}\n", .{ rank, world, run_index, graph.prefill_stats.captures - prefill_before.captures, graph.prefill_stats.updates - prefill_before.updates, graph.prefill_stats.replays - prefill_before.replays, graph.prefill_stats.paged_replays - prefill_before.paged_replays, graph.prefill_stats.rebinds - prefill_before.rebinds, graph.prefill_stats.sighting_hits - prefill_before.sighting_hits, graph.prefill_stats.sighting_misses - prefill_before.sighting_misses, graph.prefill_stats.eager - prefill_before.eager, graph.stats.replays - before.replays }); + if (run_index == 1) { + try check.expect(graph.prefill_stats.captures + graph.prefill_stats.updates + graph.prefill_stats.replays + graph.prefill_stats.eager > prefill_before.captures + prefill_before.updates + prefill_before.replays + prefill_before.eager, "rank {d} restored prefill graph warmup runs", .{rank}); + } else if (run_index == 2) { + try check.expect(graph.prefill_stats.captures + graph.prefill_stats.updates + graph.prefill_stats.replays > prefill_before.captures + prefill_before.updates + prefill_before.replays, "rank {d} restored prefill graph captures or rebinds after warmup", .{rank}); + } else { + try check.expect(graph.prefill_stats.paged_replays > prefill_before.paged_replays, "rank {d} paged-attention graph replay after snapshot restore", .{rank}); + } + try check.expect(graph.stats.replays > before.replays, "rank {d} concurrent decode graph replay after restore", .{rank}); + } + if (probe) |foreign| { + for (sequences) |sequence| { + const paged = sequence.paged.?; + const status = try paged.total(); + try check.expect(!status.overflow and status.misses > 0 and status.rebases > 0, "rank {d}: shared paging actually fetches blocks and rebases without overflow", .{rank}); + for (paged.win0) |origin| try check.expect(origin > 0 and try X.kvpage.pagingRows(paged.cfg, origin, 256) < 256, "rank {d}: slots1024 forces paging subdivision", .{rank}); + } + try foreign.unchanged(g); + try foreign.continueAndRestore(g); + } + std.debug.print("SNAPSHOT rank {d}/{d} {s}: two paged lanes, {d} decode steps, exact state and logits; prefill replays {d}, decode replays {d}\n", .{ rank, world, @tagName(mode), steps, graph.prefill_stats.replays - prefill_before.replays, graph.stats.replays - before.replays }); + } + if (shared) check.pass("rank {d}/{d}: shared slots1024 TP2 snapshot remapping, recycled generation, foreign windows/history/continuation, exact eager/graph state/logits and actual prefill/decode replays", .{ rank, world }) else check.pass("rank {d}/{d}: concurrent paged snapshots preserve BF16 continuation and graph prefill/decode", .{ rank, world }); +} + +/// cuProfilerStart/Stop from libcuda (nsys --capture-range=cudaProfilerApi), looked up at run time. +const Profiler = struct { + lib: std.DynLib, + start: *const fn () callconv(.c) c_int, + stop: *const fn () callconv(.c) c_int, + fn open() !Profiler { + var lib = try std.DynLib.open("libcuda.so.1"); + return .{ .start = lib.lookup(*const fn () callconv(.c) c_int, "cuProfilerStart") orelse return error.MissingSymbol, .stop = lib.lookup(*const fn () callconv(.c) c_int, "cuProfilerStop") orelse return error.MissingSymbol, .lib = lib }; + } +}; + +/// A profiling window for nsys: `steps` decode rounds (one row) then `steps` verify windows of `rows` rows, after a +/// warm-up that captures their graphs (TF_FLASHNEXT_EXL3_GRAPH_AFTER=1 recommended), inside cuProfilerStart/Stop. +fn profMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, steps: usize, rows: usize, which: []const u8) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const ids = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + graph.force = if (std.mem.eql(u8, which, "graph")) .graph else if (std.mem.eql(u8, which, "dplan")) .dplan else .off; + var prof = try Profiler.open(); + const toks = try a.alloc(u32, 4 * steps * rows + 64); + defer a.free(toks); + var phase: usize = 0; + while (phase < 2) : (phase += 1) { + try dplan.toHost(); + try g.reset(); + try g.forward(&ids, .{ .logits = .last }); + var tok = prompt.argmax(g.seq.logits); + _ = try g.swap(5); + if (phase == 1) _ = prof.start(); + const t0 = check.now(gpu.io); + for (0..steps) |i| { + tok = try g.round(tok); + toks[i] = tok; + } + const t1 = check.now(gpu.io); + for (0..steps) |i| { + for (0..rows) |j| toks[steps + i * rows + j] = toks[(i * rows + j) % steps]; + try g.forward(toks[steps + i * rows ..][0..rows], .{ .logits = .all, .logits_out = g.out.ptr }); + } + const t2 = check.now(gpu.io); + if (phase == 1) _ = prof.stop(); + if (rank == 0) std.debug.print("RESULT prof tp={d} {s} phase {d}: rounds {d:.2} ms each, windows of {d} rows {d:.2} ms each; graph captures {d} replays {d}; host prepare {d:.3} ms a step\n", .{ world, which, phase, ms(t0, t1) / @as(f64, @floatFromInt(steps)), rows, ms(t1, t2) / @as(f64, @floatFromInt(steps)), graph.stats.captures, graph.stats.replays, @as(f64, @floatFromInt(graph.stats.prepare_ns)) / 1e6 / @as(f64, @floatFromInt(@max(graph.stats.prepares, 1))) }); + } + graph.force = null; +} + +/// Host cost of the n-gram rows a step gathers (Ngram.rows), 1 and 3 tokens, fresh random histories. +fn ngramMode(gpu: Gpu, dir: []const u8) !void { + const a = gpu.gpa; + var ng = try flash.ple.Ngram.init(a, gpu.io, dir); + defer ng.deinit(); + const w = ng.width(); + const emb = try a.alloc(u16, 4 * w); + defer a.free(emb); + const xs = try a.alloc(f32, 4 * w / 32); + defer a.free(xs); + var prng = std.Random.DefaultPrng.init(7); + for ([_]usize{ 1, 3 }) |n| for (0..2) |rep| { + var tot: f64 = 0; + var worst: f64 = 0; + const reps = 64; + for (0..reps) |_| { + var hist = [2]i64{ prng.random().uintLessThan(u32, 240000), prng.random().uintLessThan(u32, 240000) }; + var ids: [3]u32 = undefined; + for (&ids) |*x| x.* = prng.random().uintLessThan(u32, 240000); + const h0 = hist; + const t0 = check.now(gpu.io); + try ng.rows(&hist, ids[0..n], emb[0 .. n * w], xs[0 .. n * w / 32]); + var t = ms(t0, check.now(gpu.io)); + if (rep == 1) { // the same rows again: their pages are resident now + hist = h0; + const t1 = check.now(gpu.io); + try ng.rows(&hist, ids[0..n], emb[0 .. n * w], xs[0 .. n * w / 32]); + t = ms(t1, check.now(gpu.io)); + } + tot += t; + worst = @max(worst, t); + } + std.debug.print("RESULT ngram rows of {d} tokens (pass {d}, 1 = warm repeat): {d:.3} ms mean, {d:.3} ms worst\n", .{ n, rep, tot / reps, worst }); + }; +} + +/// A scheduling change's bits: every golden prompt's decode rounds and windows (runMode) under `which` (off, dplan, +/// graph) with the parallel projection parts (X.fork) off, off again and on; logits compared bit for bit, ms/token. +fn forkMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, steps: usize, which: []const u8) !void { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, golden_path, a, .limited(1 << 26)); + defer a.free(text); + var parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + graph.force = if (std.mem.eql(u8, which, "graph")) .graph else if (std.mem.eql(u8, which, "dplan")) .dplan else .off; + defer graph.force = null; + defer X.fork.force = null; + var ok = true; + for (parsed.value.object.get("prompts").?.array.items) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + var runs: [3]Run = undefined; + for (&runs) |*r| r.* = .{ .ids = try a.alloc(u32, steps + 1), .round_h = try a.alloc(u64, steps), .window_h = try a.alloc(u64, steps) }; + defer for (&runs) |*r| { + a.free(r.ids); + a.free(r.round_h); + a.free(r.window_h); + }; + for ([_]bool{ false, false, true }, 0..) |on, i| { + X.fork.force = on; + try runMode(g, ids, steps, &runs[i]); + } + const same = std.mem.eql(u64, runs[1].round_h, runs[2].round_h) and std.mem.eql(u64, runs[1].window_h, runs[2].window_h) and std.mem.eql(u32, runs[1].ids, runs[2].ids); + ok = ok and same; + if (world == 2) { + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(runs[2].round_h)); + try check.expect(try g.swap(h) == h, "rank {d}: the ranks differ", .{rank}); + } + if (rank == 0) std.debug.print("RESULT tp={d} {s} {s}: decode {d:.2} -> {d:.2} ms/token, windows {d:.2} -> {d:.2} ms/row (serial -> parallel parts); logits {s}\n", .{ world, which, p.object.get("name").?.string, runs[1].ms_token, runs[2].ms_token, runs[1].ms_window_row, runs[2].ms_window_row, if (same) "bit-equal" else "DIFFER" }); + } + try check.expect(ok, "rank {d}: parallel parts changed the logits", .{rank}); + check.pass("rank {d}/{d}: parallel projection parts give the same bits ({s})", .{ rank, world, which }); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len >= 2 and std.mem.eql(u8, rest[0], "ngram")) return ngramMode(gpu, rest[1]); + if (rest.len < 1) return error.MissingArgument; + if (std.mem.eql(u8, rest[0], "plan")) return planMode(gpu); + if (std.mem.eql(u8, rest[0], "frequency")) return @import("exl3_frequency_tests.zig").run(gpu, rest[1..]); + if (rest.len < 6) return error.MissingArgument; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const world = try std.fmt.parseInt(usize, rest[1], 10); + const port = try std.fmt.parseInt(u16, rest[2], 10); + const dir = rest[3]; + const mode = rest[4]; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const steps = if (rest.len > 6 and !std.mem.eql(u8, mode, "prefill-prof")) try std.fmt.parseInt(usize, rest[6], 10) else 48; + const which: []const u8 = if (rest.len > 7) rest[7] else "dplan,graph"; + if (std.mem.eql(u8, mode, "prefill-origin")) return @import("exl3_prefill_origin_tests.zig").run(gpu, rank, world, port, dir, rest[5], .{ .init = Eng.initPaged, .deinit = Eng.deinit, .reset = Eng.reset, .forward = Eng.forward, .swap = Eng.swap }); + if (std.mem.eql(u8, mode, "rebind")) { + const engine = try Eng.init(gpu, rank, world, port, dir); + defer engine.deinit(); + return @import("exl3_state_rebind_tests.zig").run(engine, .{ .reset = Eng.reset, .forward = Eng.forward, .swap = Eng.swap }); + } + if (std.mem.eql(u8, mode, "snapshot")) return snapshotMode(gpu, rank, world, port, dir, steps, .legacy); + if (std.mem.eql(u8, mode, "snapshot-shared")) return snapshotStorageMode(gpu, rank, world, port, dir, steps, .shared); + if (std.mem.eql(u8, mode, "snapshot-private")) return snapshotStorageMode(gpu, rank, world, port, dir, steps, .private); + if (std.mem.eql(u8, mode, "fork")) return forkMode(gpu, rank, world, port, dir, rest[5], steps, which); + if (std.mem.eql(u8, mode, "prof")) return profMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), steps, which); + if (std.mem.eql(u8, mode, "prefill")) return prefillMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), steps); + if (std.mem.eql(u8, mode, "prefill-prof")) return prefillProfMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), if (rest.len > 6) rest[6] else "graph"); + if (std.mem.eql(u8, mode, "eq")) return eqMode(gpu, rank, world, port, dir, rest[5], steps, which); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3_tests.zig b/zig/tests/cuda/flashnext_exl3_tests.zig new file mode 100644 index 0000000000..c17846e002 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3_tests.zig @@ -0,0 +1,1184 @@ +//! flashnext-exl3: the Flash Next CUDA family on a real EXL3 pack (Qwen3.8-Flash-Next 3.05bpw), against our float64 +//! oracles (exl3_fwd_cpu.zig and the per-op ones it assembles), through the family's OWN loaders and pass code. +//! Usage: tf-cuda-test flashnext-exl3 names every tensor the family reads resolves; nothing +//! indexed is left unclaimed but the vision tower +//! tf-cuda-test flashnext-exl3 norms the one centred-norm convention and its cross-checks +//! tf-cuda-test flashnext-exl3 router [layers] [rows] router picks (device) vs float64 (default 0,3,17) +//! tf-cuda-test flashnext-exl3 ple PLE key/value from fp16 n-gram rows vs float64 +//! tf-cuda-test flashnext-exl3 layer [layers] projections, then one whole decoder layer (prefill +//! 8 rows + 4 decode rows) vs the float64 layer (0,3) +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("exl3_scan.zig"); +const place = @import("exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; + +/// "The capital of France is Paris. It is known for the Eiffel Tower." (exl3_fwd_tests.zig's default ids) +pub const default_ids = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + +// -- the raw pack (the oracle's view: tensors exactly as stored, no widening) ------------------------------------- + +pub const RawPack = struct { + files: std.ArrayList(st.File) = .empty, + + pub fn open(gpu: Gpu, dir: []const u8) !RawPack { + var p: RawPack = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 }); + try p.files.append(gpu.gpa, try st.File.open(gpu.gpa, gpu.io, path)); + } + return p; + } + + pub fn close(p: *RawPack, gpu: Gpu) void { + for (p.files.items) |*f| f.close(gpu.io); + p.files.deinit(gpu.gpa); + } + + pub fn get(p: *const RawPack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files.items) |*f| if (f.get(name)) |t| return t; + std.debug.print("raw pack: no {s}\n", .{name}); + return error.MissingTensor; + } +}; + +fn typed(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +fn valueOf(t: st.Tensor, i: usize) f64 { + const h = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + return switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(h))), + .bf16 => @as(f32, @bitCast(@as(u32, h) << 16)), + else => unreachable, + }; +} + +fn f64Of(a: std.mem.Allocator, t: st.Tensor) ![]f64 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f64, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = valueOf(t, i); + return out; +} + +fn f32Of(a: std.mem.Allocator, t: st.Tensor) ![]f32 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f32, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = @floatCast(valueOf(t, i)); + return out; +} + +fn f16Of(a: std.mem.Allocator, t: st.Tensor) ![]f16 { + if (t.dtype != .f16) return error.UnexpectedTensor; + return typed(f16, a, t.bytes); +} + +/// The EXL3 matrix `prefix` as the float64 oracle's QMat (W = diag(suh) H Wq H diag(svh), exl3_format's decoder). +fn qmat(a: std.mem.Allocator, raw: *const RawPack, prefix: []const u8, threads: usize) !fwd.QMat { + const t = try raw.get("{s}.trellis", .{prefix}); + const trellis = try typed(i16, a, t.bytes); + defer a.free(trellis); + const suh = try typed(f16, a, (try raw.get("{s}.suh", .{prefix})).bytes); + defer a.free(suh); + const svh = try typed(f16, a, (try raw.get("{s}.svh", .{prefix})).bytes); + defer a.free(svh); + return fwd.QMat.init(a, @intCast(t.dim(2) / 8), trellis, t.dim(0), t.dim(1), suh, svh, threads); +} + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bfValue(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +pub const Err = struct { max_abs: f64 = 0, peak: f64 = 0, rel: f64 = 0 }; + +pub fn errOf(got: anytype, want: []const f64) Err { + var e: Err = .{}; + for (want, 0..) |w, i| { + const g: f64 = switch (@TypeOf(got)) { + []const u16, []u16 => bfValue(got[i]), + else => @floatCast(got[i]), + }; + e.max_abs = @max(e.max_abs, @abs(g - w)); + e.peak = @max(e.peak, @abs(w)); + } + e.rel = if (e.peak > 0) e.max_abs / e.peak else 0; + return e; +} + +pub fn download(comptime T: type, gpu: Gpu, s: cuda.Stream, ptr: u64, n: usize) ![]T { + try s.synchronize(); + const out = try gpu.gpa.alloc(T, n); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// A pageable copy into the pass's scratch: the legacy-stream copy is ordered against nothing, so drain both sides. +pub fn uploadTo(gpu: Gpu, s: cuda.Stream, ptr: u64, bytes: []const u8) !void { + try s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); + try gpu.ctx.synchronize(); +} + +// -- names ---------------------------------------------------------------------------------------------------------- + +const Names = struct { + gpa: std.mem.Allocator, + store: *wt.Store, + claimed: std.StringHashMapUnmanaged(void) = .empty, + missing: usize = 0, + matrices: usize = 0, + dense: usize = 0, + widened: usize = 0, + by_halves: [17]usize = @splat(0), + + fn claim(n: *Names, requested: []const u8) !void { + var buf: [512]u8 = undefined; + const key = n.store.resolveKey(requested, &buf) catch { + n.missing += 1; + if (n.missing <= 20) std.debug.print(" missing: {s}\n", .{requested}); + return; + }; + if (n.claimed.contains(key)) return; + try n.claimed.put(n.gpa, try n.gpa.dupe(u8, key), {}); + } + + /// An EXL3 matrix: its four parts claimed and the matrix admitted by the family's own reader (marker, widths). + fn matrix(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const prefix = try std.fmt.bufPrint(&buf, fmt, args); + var pb: [280]u8 = undefined; + for ([_][]const u8{ "trellis", "suh", "svh", "mul1" }) |part| try n.claim(try std.fmt.bufPrint(&pb, "{s}.{s}", .{ prefix, part })); + if (!X.load.has(n.store, prefix)) { + n.missing += 1; + std.debug.print(" not EXL3: {s}\n", .{prefix}); + return; + } + const m = try X.load.mat(n.store, prefix); + n.by_halves[m.halves] += 1; + n.matrices += 1; + } + + /// A dense tensor as the loaders read it: through Store.tensor, so an F16 `.weight` is widened (exact or refused). + fn tensor(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + try n.claim(name); + if (!n.store.has(name)) return; + const before = n.store.widened.items.len; + _ = try n.store.tensor(name); + if (n.store.widened.items.len > before) n.widened += 1; + n.dense += 1; + } + + fn hcw(n: *Names, base: []const u8, inject: bool) !void { + try n.tensor("{s}.hc_norm.weight", .{base}); + try n.tensor("{s}.input_mix_weight_down.weight", .{base}); + try n.tensor("{s}.input_mix_weight_up.weight", .{base}); + if (inject) try n.tensor("{s}.block_inject_weight.weight", .{base}); + } + + /// What cuda_prompt_load.loadLayerAtRank (and the EXL3 hooks it reaches) read for one decoder layer at `base`. + fn layer(n: *Names, base: []const u8, linear: bool, with_ple: bool) !void { + var hb: [128]u8 = undefined; + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.attn_hyper_connection", .{base}), true); + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.mlp_hyper_connection", .{base}), true); + try n.tensor("{s}.mlp.gate.weight", .{base}); + try n.tensor("{s}.mlp.shared_expert_gate.weight", .{base}); + if (linear) { + for ([_][]const u8{ "in_proj_qkv", "in_proj_z", "out_proj" }) |p| try n.matrix("{s}.linear_attn.{s}", .{ base, p }); + for ([_][]const u8{ "in_proj_b.weight", "in_proj_a.weight", "conv1d.weight", "A_log", "dt_bias", "norm.weight" }) |p| try n.tensor("{s}.linear_attn.{s}", .{ base, p }); + } else { + for ([_][]const u8{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }) |p| try n.matrix("{s}.self_attn.{s}", .{ base, p }); + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }) |p| try n.tensor("{s}.self_attn.{s}.weight", .{ base, p }); + } + if (with_ple) { + for ([_][]const u8{ "key_proj", "value_proj", "norm_key", "norm_query", "norm_conv", "conv1d" }) |p| try n.tensor("{s}.ple.{s}.weight", .{ base, p }); + } + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.shared_expert.{s}", .{ base, p }); + for (0..experts_n) |e| { + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.experts.{d}.{s}", .{ base, e, p }); + if (e % 64 == 63) n.store.release(); + } + n.store.release(); + } +}; + +fn names(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var n: Names = .{ .gpa = gpa, .store = &store }; + defer { + var it = n.claimed.keyIterator(); + while (it.next()) |k| gpa.free(k.*); + n.claimed.deinit(gpa); + } + const t0 = check.now(gpu.io); + // the engine (cuda_prompt.zig Engine.initRank): embedding, the final mixer, the head (EXL3 here: plan step 9) + try n.tensor("language_model.model.embed_tokens.weight", .{}); + try n.hcw("language_model.model.hyper_connection_mixer", false); + try n.matrix("language_model.lm_head", .{}); + for (0..wt.layers_n) |l| { + var bb: [64]u8 = undefined; + try n.layer(try std.fmt.bufPrint(&bb, "language_model.model.layers.{d}", .{l}), wt.linearLayer(l), l == wt.ple_layer); + } + // the MTP head (cuda_mtp.zig Head.init): norms, two EXL3 fc projections, its mixer (the patch file), one layer + try n.tensor("language_model.mtp.pre_fc_norm_embedding.weight", .{}); + try n.tensor("language_model.mtp.pre_fc_norm_hidden.weight", .{}); + try n.matrix("language_model.mtp.fc_embedding", .{}); + try n.matrix("language_model.mtp.fc_hidden", .{}); + try n.hcw("language_model.mtp.hyper_connection_mixer", false); + try n.layer("language_model.mtp.layers.0", false, false); + // everything the index (and the patch) names: claimed, or the vision tower; and the header rules' homes agree + const wm = store.parsed.value.object.get("weight_map").?.object; + var vision: usize = 0; + var unclaimed: usize = 0; + var homes: [8]usize = @splat(0); + var disagree: usize = 0; + var it = wm.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + const home = place.home(.{ .name = name, .rank = if (std.mem.endsWith(u8, name, ".trellis")) 3 else 1, .file = kv.value_ptr.string }); + homes[@backingInt(home)] += 1; + if (n.claimed.contains(name)) { + // a routed expert's part must be the routed bank's, a shared expert's part the resident one's + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null and home != .routed_host) disagree += 1; + if (std.mem.indexOf(u8, name, ".shared_expert.") != null and home != .shared_resident) disagree += 1; + continue; + } + if (std.mem.startsWith(u8, name, "model.visual.")) { + vision += 1; + continue; + } + unclaimed += 1; + if (unclaimed <= 20) std.debug.print(" unclaimed: {s}\n", .{name}); + } + const ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + std.debug.print("names: {d} indexed (+ the mixer patch), {d} claimed, {d} vision, {d} unclaimed, {d} missing; {d} EXL3 matrices admitted (mul1 marker, widths from the trellis), {d} dense tensors read, {d} F16 weights widened to BF16 exactly; {d:.0} ms\n", .{ wm.count(), n.claimed.count(), vision, unclaimed, n.missing, n.matrices, n.dense, n.widened, ms }); + std.debug.print(" EXL3 widths (halves: count):", .{}); + for (n.by_halves, 0..) |c, h| if (c > 0) std.debug.print(" {d}:{d}", .{ h, c }); + std.debug.print("\n header homes (exl3_place):", .{}); + for (homes, 0..) |c, i| if (c > 0) std.debug.print(" {s}={d}", .{ @tagName(@as(place.Home, @fromBackingInt(@intCast(i)))), c }); + std.debug.print("\n", .{}); + try check.expect(n.missing == 0, "names: {d} tensors the family reads do not resolve", .{n.missing}); + try check.expect(unclaimed == 0, "names: {d} indexed tensors no family loader reads (vision excepted)", .{unclaimed}); + try check.expect(disagree == 0, "names: {d} expert tensors whose exl3_place home disagrees", .{disagree}); + // the family's device sizing (cuda_exl3_memory.zig) on the same pack + const gib = 1024.0 * 1024.0 * 1024.0; + const on_host = try X.memory.deviceBytes(gpa, gpu.io, dir, .{}); + const resident = try X.memory.deviceBytes(gpa, gpu.io, dir, .{ .routed_on_host = false }); + const head = try X.memory.headBytes(gpa, gpu.io, dir, .{}); + std.debug.print(" device bytes: {d:.2} GiB with routed experts on the host, {d:.2} GiB all resident; MTP head {d:.2} GiB (routed on the host)\n", .{ @as(f64, @floatFromInt(on_host)) / gib, @as(f64, @floatFromInt(resident)) / gib, @as(f64, @floatFromInt(head)) / gib }); + check.pass("names: every tensor the family reads resolves ({d}), nothing unclaimed but {d} vision tensors; {d} EXL3 matrices admitted; {d} F16 weights widened exactly", .{ n.claimed.count(), vision, n.matrices, n.widened }); +} + +// -- norms ---------------------------------------------------------------------------------------------------------- + +fn meanOf(t: st.Tensor) f64 { + var s: f64 = 0; + const count = t.bytes.len / 2; + for (0..count) |i| s += valueOf(t, i); + return s / @as(f64, @floatFromInt(count)); +} + +fn norms(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + try store.detectNorms(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var means: [wt.layers_n]f64 = undefined; + for (&means, 0..) |*m, l| m.* = meanOf(try raw.get("model.language_model.layers.{d}.attn_hyper_connection.hc_norm.weight", .{l})); + var sorted = means; + std.mem.sort(f64, &sorted, {}, std.sort.asc(f64)); + std.debug.print("norms: 48 attn hc_norm means min {d:.4} median {d:.4} max {d:.4} -> offset {d} (cuda_norms.offset via Store.detectNorms)\n", .{ sorted[0], (sorted[23] + sorted[24]) / 2, sorted[47], store.norm_offset }); + try check.expect(store.norm_offset == 1, "norms: the pack stores gamma - 1; the convention came out {d}", .{store.norm_offset}); + // the independent cross-check: the indexer norms' own 24 means locate the same storage + var idx: [24]f32 = undefined; + var c: usize = 0; + for (0..wt.layers_n) |l| { + if (wt.linearLayer(l)) continue; + for ([_][]const u8{ "q_layernorm", "k_layernorm" }) |nm| { + idx[c] = @floatCast(meanOf(try raw.get("model.language_model.layers.{d}.self_attn.indexer.{s}.weight", .{ l, nm }))); + c += 1; + } + } + const indexer = try scan.centredOffset(&idx); + std.debug.print("norms: indexer q/k_layernorm means over the 12 attention layers -> exl3_scan.centredOffset {d}\n", .{indexer}); + try check.expect(indexer == store.norm_offset, "norms: the indexer norms locate offset {d}, the hc_norm means {d}", .{ indexer, store.norm_offset }); + // every centred norm the family scales, MTP's included, loads through Store.normScale with that one offset + const mtp_names = [_][]const u8{ "language_model.mtp.pre_fc_norm_embedding.weight", "language_model.mtp.pre_fc_norm_hidden.weight", "language_model.mtp.hyper_connection_mixer.hc_norm.weight", "language_model.mtp.layers.0.attn_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.mlp_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.self_attn.q_norm.weight", "language_model.mtp.layers.0.self_attn.k_norm.weight", "language_model.mtp.layers.0.self_attn.indexer.q_layernorm.weight", "language_model.mtp.layers.0.self_attn.indexer.k_layernorm.weight", "language_model.model.layers.1.ple.norm_key.weight", "language_model.model.layers.1.ple.norm_query.weight", "language_model.model.layers.1.ple.norm_conv.weight", "language_model.model.layers.3.self_attn.q_norm.weight", "language_model.model.layers.3.self_attn.k_norm.weight", "language_model.model.hyper_connection_mixer.hc_norm.weight" }; + for (mtp_names) |name| { + const scales = try store.normScale(gpa, try store.tensor(name)); + defer gpa.free(scales); + var sum: f64 = 0; + var lo: f32 = std.math.inf(f32); + for (scales) |v| { + sum += v; + lo = @min(lo, v); + } + std.debug.print(" {s:<72} stored mean {d:>7.4} -> scale mean {d:.4} (min {d:.4})\n", .{ name, sum / @as(f64, @floatFromInt(scales.len)) - store.norm_offset, sum / @as(f64, @floatFromInt(scales.len)), lo }); + } + store.release(); + check.pass("norms: one convention (offset {d}) from the 48 hc_norm means, confirmed by the indexer norms; every centred norm incl. MTP's scales finite", .{store.norm_offset}); +} + +// -- router --------------------------------------------------------------------------------------------------------- + +fn router(gpu: Gpu, dir: []const u8, layers: []const usize, rows: usize) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x5eed_2026); + const rng = prng.random(); + var worst_rows_diff: usize = 0; + for (layers) |l| { + // the family's router buffer, as cuda_prompt_load builds it: mlp.gate's 512 rows then the shared gate (Store-widened) + var nb: [128]u8 = undefined; + const gate = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.gate.weight", .{l})); + try check.expect(gate.is(.bf16, &.{ experts_n, dims }), "router: layer {d} gate is not the loaders' bf16 [512, 2560]", .{l}); + const w = try gpa.alloc(u8, (experts_n + 1) * dims * 2); + defer gpa.free(w); + @memcpy(w[0 .. experts_n * dims * 2], gate.bytes); + const sg = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + @memcpy(w[experts_n * dims * 2 ..], sg.bytes); + store.release(); + var wd = try cuda.DeviceBuffer.fromHost(gpu.d, w); + defer wd.free(); + // the oracle's weights: the pack's F16 bits, never the widened copy + const g16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})); + defer gpa.free(g16); + const s16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + defer gpa.free(s16); + // rows: N(0, sigma) per row, sigma in {0.25, 1, 4}, rounded to bf16 as the pass's `mixed` is + const xb = try gpa.alloc(u16, rows * dims); + defer gpa.free(xb); + for (0..rows) |r| { + const sigma: f32 = switch (r % 3) { + 0 => 0.25, + 1 => 1, + else => 4, + }; + for (0..dims) |i| xb[r * dims + i] = bf16Bits(rng.floatNorm(f32) * sigma); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + var ld = try cuda.DeviceBuffer.alloc(gpu.d, rows * (experts_n + 1) * 4); + defer ld.free(); + var pd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer pd.free(); + var wdd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer wdd.free(); + try glue.routerMany(s, xd.ptr, wd.ptr, ld.ptr, dims, rows); + try glue.topkMany(s, ld.ptr, pd.ptr, wdd.ptr, rows); + const gl = try download(f32, gpu, s, ld.ptr, rows * (experts_n + 1)); + defer gpa.free(gl); + const gp = try download(i32, gpu, s, pd.ptr, rows * slots); + defer gpa.free(gp); + const gw = try download(f32, gpu, s, wdd.ptr, rows * slots); + defer gpa.free(gw); + var diff_rows: usize = 0; + var logit_err: f64 = 0; + var min_gap: f64 = std.math.inf(f64); + var worst_w: f64 = 0; + var worst_sw: f64 = 0; + var f32_agree: usize = 0; + const x = try gpa.alloc(f64, dims); + defer gpa.free(x); + const xf = try gpa.alloc(f32, dims); + defer gpa.free(xf); + var logits: [experts_n + 1]f64 = undefined; + for (0..rows) |r| { + for (x, xf, 0..) |*v, *vf, i| { + v.* = bfValue(xb[r * dims + i]); + vf.* = @floatCast(v.*); + } + try moe.logitsF16(g16, experts_n, s16, x, &logits); + var picks: [topk]u16 = undefined; + var wts: [topk]f64 = undefined; + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, &picks, &wts, &sw); + for (0..experts_n + 1) |i| logit_err = @max(logit_err, @abs(@as(f64, gl[r * (experts_n + 1) + i]) - logits[i])); + var sorted = logits[0..experts_n].*; + std.mem.sort(f64, &sorted, {}, std.sort.desc(f64)); + const gap = sorted[topk - 1] - sorted[topk]; + min_gap = @min(min_gap, gap); + var a: [topk]u16 = undefined; + for (&a, 0..) |*v, k| v.* = @intCast(gp[r * slots + k]); + var b = picks; + std.mem.sort(u16, &a, {}, std.sort.asc(u16)); + std.mem.sort(u16, &b, {}, std.sort.asc(u16)); + if (!std.mem.eql(u16, &a, &b)) { + diff_rows += 1; + if (diff_rows <= 10) std.debug.print(" layer {d} row {d}: picks differ; f64 logit gap between the 10th and 11th {e:.3}\n", .{ l, r, gap }); + } else { + for (0..topk) |k| { + // the same expert's weight: the device's bf16 against the float64 softmax + const id = gp[r * slots + k]; + for (picks, wts) |pk, wv| if (pk == id) { + worst_w = @max(worst_w, @abs(@as(f64, gw[r * slots + k]) - wv)); + }; + } + try check.expect(gp[r * slots + topk] == experts_n, "router: the shared slot holds {d}", .{gp[r * slots + topk]}); + worst_sw = @max(worst_sw, @abs(@as(f64, gw[r * slots + topk]) - sw)); + } + // the f32 host reference (exl3_router.launchDtype(.f16)) on the same row: its picks as a set + var l32: [experts_n + 1]f32 = undefined; + var p32: [topk + 1]u16 = undefined; + var w32: [topk + 1]f32 = undefined; + var scratch: [experts_n]f32 = undefined; + try router_cpu.launchDtype(.f16, 0, std.mem.sliceAsBytes(g16), xf, std.mem.sliceAsBytes(s16), xf, topk, &l32, &p32, &w32, &scratch); + var c32 = p32[0..topk].*; + std.mem.sort(u16, &c32, {}, std.sort.asc(u16)); + f32_agree += @intFromBool(std.mem.eql(u16, &c32, &b)); + } + std.debug.print("router layer {d}: {d} rows, picks equal to float64 on {d}/{d} (f32 host reference equal on {d}/{d}); max |device logit - f64| {e:.3}; smallest f64 10th/11th gap {e:.3}; weights max |bf16 - f64| {e:.3}, shared {e:.3}\n", .{ l, rows, rows - diff_rows, rows, f32_agree, rows, logit_err, min_gap, worst_w, worst_sw }); + worst_rows_diff = @max(worst_rows_diff, diff_rows); + try check.expect(diff_rows == 0 or min_gap <= 4 * logit_err, "router layer {d}: {d} rows pick differently with every gap above the logit error", .{ l, diff_rows }); + } + try check.expect(worst_rows_diff == 0, "router: picks differ from float64 on {d} rows (see the gaps above)", .{worst_rows_diff}); + check.pass("router: the exactly widened F16 router's top-10 equals the float64 oracle on every row of layers {any}", .{layers}); +} + +// -- the pass rig ----------------------------------------------------------------------------------------------------- + +pub const Pass = batch.Pass(flash.Tri); + +pub const Rig = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + k: batch.Kernels, + s: batch.Scratch, + store: wt.Store, + pack_mod: cuda.Module, + pack_fn: cuda.Function, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + conv: [wt.linear_n]cuda.DeviceBuffer = undefined, + state: [wt.linear_n]cuda.DeviceBuffer = undefined, + kc: [wt.attn_n]cuda.DeviceBuffer = undefined, + vc: [wt.attn_n]cuda.DeviceBuffer = undefined, + ikc: [wt.attn_n]cuda.DeviceBuffer = undefined, + pooled: [wt.attn_n]cuda.DeviceBuffer = undefined, + tail: cuda.DeviceBuffer = undefined, + history: [2]i64 = .{ 0, 0 }, + scores: cuda.DeviceBuffer = undefined, + cap: usize, + /// The ranks' NCCL owner when a check runs two ranks (flashnext_exl3_tp_tests.zig); null: one rank. + collective: ?*flash.collective.Collective = null, + + pub fn init(r: *Rig, gpu: Gpu, dir: []const u8, rows: usize, cap: usize) !void { + r.* = .{ .gpu = gpu, .stream = try cuda.Stream.init(gpu.d, true), .glue = undefined, .k = undefined, .s = undefined, .store = undefined, .pack_mod = undefined, .pack_fn = undefined, .cap = cap }; + r.glue = try flash.Glue.load(gpu.d); + r.k = try batch.Kernels.load(gpu.d); + r.s = try batch.Scratch.alloc(gpu.gpa, gpu.d, r.stream, rows); + r.store = try wt.Store.open(gpu.gpa, gpu.io, dir); + try r.store.detectNorms(); + r.pack_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + r.pack_fn = try r.pack_mod.function(wt.pack_symbol); + const gdn = flash.gdn; + const z: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = 0, .len = 0 }; + r.conv = @splat(z); + r.state = @splat(z); + r.kc = @splat(z); + r.vc = @splat(z); + r.ikc = @splat(z); + r.pooled = @splat(z); + // every layer's state, as Engine.newSeq makes them (a check may step any layer) + for (&r.conv, &r.state) |*c, *x| { + c.* = try r.zeroed((gdn.taps - 1) * gdn.conv_dim * 2); + x.* = try r.zeroed(wt.state_n * 4); + } + for (&r.kc, &r.vc, &r.ikc, &r.pooled) |*kc, *vc, *ikc, *pooled| { + kc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + vc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + ikc.* = try r.zeroed(cap * wt.index_dim * 2); + pooled.* = try r.zeroed(flash.prompt.pooledBytes(cap)); + } + r.tail = try r.zeroed(wt.ple_tail_n * wide * 2); + r.scores = try r.zeroed(batch.scoreBytes(cap, rows)); + } + + fn zeroed(r: *Rig, bytes: usize) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + return b; + } + + pub fn reset(r: *Rig) !void { + for (r.bufs.items) |b| try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + } + + pub fn deinit(r: *Rig) void { + r.stream.synchronize() catch {}; + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + r.pack_mod.unload(); + r.store.deinit(); + r.s.free(r.gpu.gpa); + r.k.deinit(); + r.glue.deinit(); + r.stream.deinit(); + } + + pub fn pass(r: *Rig, layers: []const pl.Layer) Pass { + return .{ + .gpa = r.gpu.gpa, + .io = r.gpu.io, + .stream = r.stream, + .tri = .{ .own = &r.glue, .s = r.stream }, + .k = &r.k, + .s = &r.s, + .m = .{ .table_w = 0, .table_s = 0, .table_b = 0, .layers = layers, .mixer = undefined, .head = undefined, .ngram = undefined, .collective = r.collective }, + .seq = .{ .conv = &r.conv, .state = &r.state, .kc = &r.kc, .vc = &r.vc, .ikc = &r.ikc, .pooled = &r.pooled, .tail = r.tail, .history = &r.history, .cap = r.cap, .scores = r.scores.ptr }, + .hooks = .{}, + }; + } +}; + +fn bf16Rows(a: std.mem.Allocator, rng: std.Random, n: usize, sigma: f32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*x| x.* = bf16Bits(rng.floatNorm(f32) * sigma); + return v; +} + +// -- projections (gate 4) -------------------------------------------------------------------------------------------- + +/// One Projection through Pass.lane(): `sk` 1 is bf16 [m, n]; `sk` > 1 fp32 [m, n] in the first slice, others zeroed. +fn laneOut(gpu: Gpu, p: *Pass, w: wt.Q4, x: u64, m: usize, sk: usize) !struct { bf: ?[]u16 = null, f: ?[]f32 = null, rest_zero: bool = true } { + const s = p.s; + if (sk == 1) { + try p.lane(x, s.dummy, w, s.proj, s.dummy, m, 1); + return .{ .bf = try download(u16, gpu, p.stream, s.proj, m * w.n) }; + } + // the slices in a buffer of their own: a 12,288-wide part's exceeds the pass's out_sk x dims `part` + var slices = try cuda.DeviceBuffer.alloc(gpu.d, sk * m * w.n * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + try p.lane(x, s.dummy, w, s.dummy, slices.ptr, m, sk); + const all = try download(f32, gpu, p.stream, slices.ptr, sk * m * w.n); + defer gpu.gpa.free(all); + var zero = true; + for (all[m * w.n ..]) |v| zero = zero and v == 0; + return .{ .f = try gpu.gpa.dupe(f32, all[0 .. m * w.n]), .rest_zero = zero }; +} + +/// The f64 product of `rows` bf16 rows through the EXL3 matrix `q`. +fn exactRows(a: std.mem.Allocator, q: *const fwd.QMat, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * q.n); + const x = try a.alloc(f64, q.k); + defer a.free(x); + const tmp = try a.alloc(f64, q.k); + defer a.free(tmp); + for (0..rows) |r| { + for (x, 0..) |*v, i| v.* = bfValue(xb[r * q.k + i]); + q.apply(x, y[r * q.n ..][0..q.n], tmp); + } + return y; +} + +/// The f64 product of `rows` bf16 rows through a dense row-major [n, k] weight. +fn denseRows(a: std.mem.Allocator, w: []const f64, n: usize, k: usize, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * n); + for (0..rows) |r| for (0..n) |j| { + var acc: f64 = 0; + for (0..k) |i| acc += bfValue(xb[r * k + i]) * w[j * k + i]; + y[r * n + j] = acc; + }; + return y; +} + +/// The fp16 rotated-input contract: x * suh after the Hadamard is rounded to fp16 (2^-11 of each term), so a column's +/// error is a few 1e-4 of the row's peak. exl3-pack measured 1.75e-4 .. 2.4e-4 (uniform rows); here N(0, 1) rows give +/// 1.3e-4 .. 3.04e-4 (layer 3 q_proj, one row, peak 4.5). +const proj_f32_tol = 4e-4; +const proj_bf16_tol = 4e-3; // bf16 output rounding: half an ulp is 2^-9 of a value, at most 1.95e-3 of the peak + +fn projections(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, a: std.mem.Allocator, rng: std.Random) !void { + var p = r.pass((lw)[0..1]); + const linear = wt.linearLayer(l); + const kind = if (linear) "linear_attn" else "self_attn"; + const e = lw.proj.exl3 orelse return error.NotExl3; + try check.expect(lw.out.exl3 != null and lw.proj.native == null and lw.out.native == null, "layer {d}: projections did not load as EXL3", .{l}); + var worst32: f64 = 0; + var worst16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xb = try bf16Rows(a, rng, m * dims, 1); + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + // the stitched input projection (bf16, as the pass reads it) and each of its parts alone (fp32) + const whole = try laneOut(gpu, &p, lw.proj, xd.ptr, m, 1); + for (e.parts) |pt| { + const want = switch (pt.kind) { + .exl3 => blk: { + var names_buf: [128]u8 = undefined; + const name = partName(&names_buf, l, kind, pt.col, linear); + var q = try qmat(a, raw, name, 48); + defer q.deinit(a); + const y = try exactRows(a, &q, xb, m); + // the part alone in fp32: Projection.exl3 over one part + var one = e.*; + var single = [1]X.load.Part{pt}; + single[0].col = 0; + one.parts = &single; + one.n = pt.n; + var view = wt.Q4.empty(gpu.d); + view.n = pt.n; + view.k = dims; + view.exl3 = &one; + const got = try laneOut(gpu, &p, view, xd.ptr, m, 2); + const ef = errOf(got.f.?, y); + try check.expect(got.rest_zero, "layer {d}: {s} left the other K slice unzeroed", .{ l, name }); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out: max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, name, m, ef.max_abs, ef.peak, ef.rel }); + worst32 = @max(worst32, ef.rel); + break :blk y; + }, + .native => blk: { + // in_proj_b then in_proj_a: the F16 rows the Store widened, the oracle reads the stored F16 + const b = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_b.weight", .{l})); + const av = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_a.weight", .{l})); + const w = try std.mem.concat(a, f64, &.{ b, av }); + break :blk try denseRows(a, w, pt.n, dims, xb, m); + }, + }; + // the part's columns of the stitched bf16 row + const got = try a.alloc(u16, m * pt.n); + for (0..m) |row| @memcpy(got[row * pt.n ..][0..pt.n], whole.bf.?[row * e.n + pt.col ..][0..pt.n]); + const eb = errOf(got, want); + std.debug.print(" L{d} stitched in-proj columns {d}..{d} ({s}) rows={d} bf16 out: rel {e:.3}\n", .{ l, pt.col, pt.col + pt.n, @tagName(pt.kind), m, eb.rel }); + worst16 = @max(worst16, eb.rel); + } + // the output projection as the pass runs it: fp32 into the first of out_sk K slices + const ok = lw.out.k; + const ob = try bf16Rows(a, rng, m * ok, 1); + var od = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(ob)); + defer od.free(); + var nb: [128]u8 = undefined; + const oname = try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, if (linear) "out_proj" else "o_proj" }); + var q = try qmat(a, raw, oname, 48); + defer q.deinit(a); + const want = try exactRows(a, &q, ob, m); + const got = try laneOut(gpu, &p, lw.out, od.ptr, m, wt.out_sk); + const eo = errOf(got.f.?, want); + try check.expect(got.rest_zero, "layer {d}: the output projection left its other K slices unzeroed", .{l}); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out (slice 0 of {d}): max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, oname, m, wt.out_sk, eo.max_abs, eo.peak, eo.rel }); + worst32 = @max(worst32, eo.rel); + } + try check.expect(worst32 <= proj_f32_tol, "layer {d}: an EXL3 projection's fp32 output is {e:.3} of peak from float64 (> {e})", .{ l, worst32, proj_f32_tol }); + try check.expect(worst16 <= proj_bf16_tol, "layer {d}: the stitched bf16 input projection is {e:.3} of peak from float64 (> {e})", .{ l, worst16, proj_bf16_tol }); + check.pass("layer {d} projections through Pass.lane(): fp32 out worst {e:.3} of peak, stitched bf16 in-proj worst {e:.3}, rows 1/5/8", .{ l, worst32, worst16 }); +} + +/// The EXL3 matrix behind a stitched part, from its first column. +fn partName(buf: []u8, l: usize, kind: []const u8, col: usize, linear: bool) []const u8 { + const name: []const u8 = if (linear) switch (col) { + 0 => "in_proj_qkv", + else => "in_proj_z", + } else switch (col) { + 0 => "q_proj", + 12288 => "k_proj", + 12800 => "v_proj", + else => "indexer.index_qk_proj", + }; + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, name }) catch unreachable; +} + +// -- the float64 layer -------------------------------------------------------------------------------------------- + +const HcHost = struct { + norm: []f32, + down: []f32, + up: []f32, + inject: ?[]f32, + + fn load(a: std.mem.Allocator, raw: *const RawPack, comptime base: []const u8, args: anytype, with_inject: bool) !HcHost { + return .{ + .norm = try f32Of(a, try raw.get(base ++ ".hc_norm.weight", args)), + .down = try f32Of(a, try raw.get(base ++ ".input_mix_weight_down.weight", args)), + .up = try f32Of(a, try raw.get(base ++ ".input_mix_weight_up.weight", args)), + .inject = if (with_inject) try f32Of(a, try raw.get(base ++ ".block_inject_weight.weight", args)) else null, + }; + } + + fn weights(h: HcHost) hc.Weights { + return .{ .norm = h.norm, .down = h.down, .up = h.up, .inject = h.inject }; + } +}; + +const ExpertCtx = struct { + gpa: std.mem.Allocator, + raw: *const RawPack, + layer: usize, + + fn load(ctx: *anyopaque, e: usize, out: *[3]fwd.QMat) anyerror!void { + const self: *ExpertCtx = @ptrCast(@alignCast(ctx)); + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var made: usize = 0; + errdefer for (0..made) |i| out[i].deinit(self.gpa); + while (made < 3) : (made += 1) { + var nb: [160]u8 = undefined; + out[made] = try qmat(self.gpa, self.raw, try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.mlp.experts.{d}.{s}", .{ self.layer, e, projs[made] }), 1); + } + } +}; + +/// The oracle's layer (exl3_fwd_cpu.LayerW) from the pack's stored tensors, decoded by the Zig f64 decoder. Arena-owned. +const Oracle = struct { + attn_hc: HcHost, + mlp_hc: HcHost, + gdn: fwd.GdnW = undefined, + attn: fwd.AttnW = undefined, + moe: fwd.MoeW, + ctx: ExpertCtx, + state: fwd.LayerState, + /// The mixer's state for the staged comparison (fed the device's own inputs), apart from the chained one. + staged: fwd.LayerState, + linear: bool, + + fn create(a: std.mem.Allocator, raw: *const RawPack, l: usize, offset: f32, cap: usize) !*Oracle { + const o = try a.create(Oracle); + o.linear = wt.linearLayer(l); + o.attn_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.attn_hyper_connection", .{l}, true); + o.mlp_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.mlp_hyper_connection", .{l}, true); + o.ctx = .{ .gpa = a, .raw = raw, .layer = l }; + var nb: [4][160]u8 = undefined; + o.moe = .{ + .gate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})), + .sgate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})), + .shared = .{ + try qmat(a, raw, try std.fmt.bufPrint(&nb[0], "model.language_model.layers.{d}.mlp.shared_expert.gate_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[1], "model.language_model.layers.{d}.mlp.shared_expert.up_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[2], "model.language_model.layers.{d}.mlp.shared_expert.down_proj", .{l}), 48), + }, + .source = .{ .ctx = &o.ctx, .load = ExpertCtx.load }, + }; + if (o.linear) { + const stem = "model.language_model.layers.{d}.linear_attn"; + o.gdn = .{ + .qkv = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".in_proj_qkv", .{l}), 48), + .z = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".in_proj_z", .{l}), 48), + .out = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".out_proj", .{l}), 48), + .a_w = try f64Of(a, try raw.get(stem ++ ".in_proj_a.weight", .{l})), + .b_w = try f64Of(a, try raw.get(stem ++ ".in_proj_b.weight", .{l})), + .params = .{ + .conv_w = try f64Of(a, try raw.get(stem ++ ".conv1d.weight", .{l})), + .a_log = try f64Of(a, try raw.get(stem ++ ".A_log", .{l})), + .dt_bias = try f64Of(a, try raw.get(stem ++ ".dt_bias", .{l})), + .norm_w = try f64Of(a, try raw.get(stem ++ ".norm.weight", .{l})), + }, + }; + o.state = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + o.staged = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + } else { + const stem = "model.language_model.layers.{d}.self_attn"; + var sc: [4][]f32 = undefined; + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }, 0..) |nm, i| { + var buf: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.self_attn.{s}.weight", .{ l, nm }); + sc[i] = try f32Of(a, try raw.get("{s}", .{name})); + for (sc[i]) |*v| v.* += offset; + } + o.attn = .{ + .q = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".q_proj", .{l}), 48), + .k = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".k_proj", .{l}), 48), + .v = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".v_proj", .{l}), 48), + .o = try qmat(a, raw, try std.fmt.bufPrint(&nb[3], stem ++ ".o_proj", .{l}), 48), + .idx = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".indexer.index_qk_proj", .{l}), 48), + .norms = .{ .q = sc[0], .k = sc[1], .iq = sc[2], .ik = sc[3], .inv = attn_cpu.invFreq() }, + }; + o.state = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + o.staged = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + } + return o; + } + + fn weights(o: *Oracle) fwd.LayerW { + return .{ .attn_hc = o.attn_hc.weights(), .mlp_hc = o.mlp_hc.weights(), .mixer = if (o.linear) .{ .gdn = &o.gdn } else .{ .attn = &o.attn }, .moe = &o.moe }; + } +}; + +// -- one decoder layer through the family's pass (gate 5) ---------------------------------------------------------- + +const Stage = struct { mixer: f64 = 0, streams: f64 = 0, moe: f64 = 0, out: f64 = 0, picks_same: usize = 0, rows: usize = 0 }; + +/// What the device fed each block: the read-outs' mixed rows and inject gates (bf16), attention's then the MoE's. +pub const Taps = struct { attn_mixed: []u16, attn_inj: []u16, mlp_mixed: []u16, mlp_inj: []u16 }; + +/// The pass's per-layer sequence (cuda_batch.zig Pass.chunk's loop body), stepped with read-backs between blocks. +pub fn deviceLayer(gpu: Gpu, p: *Pass, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h_in: []const u16, tr: fwd.Trace, h_out: []u16, taps: Taps) !void { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h_in)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + try tap(gpu, p, s.mixed, taps.attn_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.attn_inj[0 .. n * streams]); + if (lw.linear) try p.gdnBlock(lw, wt.linearIndex(l), .{}) else try p.attnBlock(lw, wt.attnIndex(l)); + const branch = try download(f32, gpu, p.stream, s.part, n * dims); + defer gpu.gpa.free(branch); + for (tr.attn_branch[0 .. n * dims], branch) |*d, v| d.* = v; + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + const mid = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(mid); + for (tr.h_mid[0 .. n * wide], mid) |*d, v| d.* = bfValue(v); + try p.readout(lw.mlp_hc, 0, n); + try tap(gpu, p, s.mixed, taps.mlp_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.mlp_inj[0 .. n * streams]); + try p.moeBlock(lw); + const picks = try download(i32, gpu, p.stream, s.picks, n * slots); + defer gpu.gpa.free(picks); + const wts = try download(f32, gpu, p.stream, s.wts, n * slots); + defer gpu.gpa.free(wts); + const y = try download(f32, gpu, p.stream, s.y, n * slots * dims); + defer gpu.gpa.free(y); + for (0..n) |r| { + for (0..topk) |k| tr.picks[r * topk + k] = @intCast(picks[r * slots + k]); + for (0..dims) |d| { + var acc: f64 = 0; + for (0..slots) |k| acc += @as(f64, wts[r * slots + k]) * y[(r * slots + k) * dims + d]; + tr.mlp_branch[r * dims + d] = acc; + } + } + p.pending = l; + try p.writeback(0, n); + const out = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(out); + @memcpy(h_out, out); +} + +fn tap(gpu: Gpu, p: *Pass, ptr: u64, out: []u16) !void { + const got = try download(u16, gpu, p.stream, ptr, out.len); + defer gpu.gpa.free(got); + @memcpy(out, got); +} + +fn toF64(a: std.mem.Allocator, bits: []const u16) ![]f64 { + const out = try a.alloc(f64, bits.len); + for (out, bits) |*o, b| o.* = bfValue(b); + return out; +} + +/// Each block against the float64 oracle fed the DEVICE's own input to that block: its error alone. +const Staged = struct { attn_mixed: f64 = 0, attn_inj: f64 = 0, mixer: f64 = 0, mlp_mixed: f64 = 0, mlp_inj: f64 = 0, moe: f64 = 0, writeback: f64 = 0 }; + +pub fn layerCheck(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, tokens: []const u32) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prefill = 8; + var t0 = check.now(gpu.io); + const o = try Oracle.create(a, raw, l, r.store.norm_offset, r.cap); + std.debug.print(" float64 layer {d} decoded in {d:.0} ms (offset {d})\n", .{ l, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6, r.store.norm_offset }); + try r.reset(); + var p = r.pass((lw)[0..1]); + const R = prefill; + const mk = struct { + fn trace(al: std.mem.Allocator) !fwd.Trace { + return .{ .attn_branch = try al.alloc(f64, R * dims), .mlp_branch = try al.alloc(f64, R * dims), .picks = try al.alloc(u16, R * topk), .h_mid = try al.alloc(f64, R * wide) }; + } + }; + const dt = try mk.trace(a); + const ot = try mk.trace(a); + const table = try raw.get("model.language_model.embed_tokens.weight", .{}); + const calls = [_][2]usize{ .{ 0, prefill }, .{ 8, 1 }, .{ 9, 1 }, .{ 10, 1 }, .{ 11, 1 } }; + var worst: Stage = .{}; + var alone: Staged = .{}; + var device_ms: f64 = 0; + const taps: Taps = .{ .attn_mixed = try a.alloc(u16, R * dims), .attn_inj = try a.alloc(u16, R * streams), .mlp_mixed = try a.alloc(u16, R * dims), .mlp_inj = try a.alloc(u16, R * streams) }; + for (calls) |c| { + const toks = tokens[c[0] .. c[0] + c[1]]; + const n = c[1]; + // the layer's input: each token's embedding row in all four streams (bf16, as the pass's expand writes it) + const hb = try a.alloc(u16, n * wide); + const ho = try a.alloc(f64, n * wide); + for (toks, 0..) |tk, t| for (0..dims) |i| { + const bits = std.mem.readInt(u16, table.bytes[(tk * dims + i) * 2 ..][0..2], .little); + for (0..streams) |sidx| { + hb[t * wide + sidx * dims + i] = bits; + ho[t * wide + sidx * dims + i] = bfValue(bits); + } + }; + const hg = try a.alloc(u16, n * wide); + t0 = check.now(gpu.io); + try deviceLayer(gpu, &p, lw, l, toks, c[0], hb, dt, hg, taps); + device_ms += @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + // staged: every block from the device's own input (before the chained oracle updates `ho` in place) + { + const w = o.weights(); + const off = r.store.norm_offset; + const mixed = try a.alloc(f64, n * dims); + const inj = try a.alloc(f64, n * streams); + try fwd.hcPreRows(a, w.attn_hc, off, ho, n, mixed, inj); + const e_am = errOf(@as([]const u16, taps.attn_mixed[0 .. n * dims]), mixed); + const e_ai = errOf(@as([]const u16, taps.attn_inj[0 .. n * streams]), inj); + const dev_mixed = try toF64(a, taps.attn_mixed[0 .. n * dims]); + const branch = try a.alloc(f64, n * dims); + switch (w.mixer) { + .gdn => |g| try fwd.gdnBranch(a, g, o.staged.gdn, dev_mixed, n, branch), + .attn => |at| try fwd.attnBranch(a, at, &o.staged.attn, c[0], dev_mixed, n, branch), + } + const e_mx = errOf(dt.attn_branch[0 .. n * dims], branch); + const dev_mid = dt.h_mid[0 .. n * wide]; + try fwd.hcPreRows(a, w.mlp_hc, off, dev_mid, n, mixed, inj); + const e_mm = errOf(@as([]const u16, taps.mlp_mixed[0 .. n * dims]), mixed); + const e_mi = errOf(@as([]const u16, taps.mlp_inj[0 .. n * streams]), inj); + const dev_mlp = try toF64(a, taps.mlp_mixed[0 .. n * dims]); + const moe_out = try a.alloc(f64, n * dims); + const picks = try a.alloc(u16, n * topk); + try fwd.moeBranch(a, w.moe, dev_mlp, n, moe_out, picks); + var same_picks = true; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same_picks = same_picks and std.mem.eql(u16, &x, &y); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], moe_out); + // the write-back alone: the device's streams after the mixer plus its MoE branch times its inject gates + const wb = try a.alloc(f64, n * wide); + for (0..n) |t| for (0..streams) |sidx| for (0..dims) |d| { + const at = t * wide + sidx * dims + d; + wb[at] = dev_mid[at] + dt.mlp_branch[t * dims + d] * bfValue(taps.mlp_inj[t * streams + sidx]); + }; + const e_wb = errOf(@as([]const u16, hg), wb); + std.debug.print(" staged (device input to each block): attn readout mixed rel {e:.3}, inject rel {e:.3}; mixer rel {e:.3}; mlp readout mixed rel {e:.3}, inject rel {e:.3}; MoE rel {e:.3} (picks {s}); write-back rel {e:.3}\n", .{ e_am.rel, e_ai.rel, e_mx.rel, e_mm.rel, e_mi.rel, e_moe.rel, if (same_picks) "equal" else "DIFFER", e_wb.rel }); + alone.attn_mixed = @max(alone.attn_mixed, e_am.rel); + alone.attn_inj = @max(alone.attn_inj, e_ai.rel); + alone.mixer = @max(alone.mixer, e_mx.rel); + alone.mlp_mixed = @max(alone.mlp_mixed, e_mm.rel); + alone.mlp_inj = @max(alone.mlp_inj, e_mi.rel); + if (same_picks) alone.moe = @max(alone.moe, e_moe.rel); + alone.writeback = @max(alone.writeback, e_wb.rel); + } + try fwd.layerRows(a, o.weights(), &o.state, c[0], r.store.norm_offset, ho, n, ot); + const e_mix = errOf(dt.attn_branch[0 .. n * dims], ot.attn_branch[0 .. n * dims]); + const e_mid = errOf(dt.h_mid[0 .. n * wide], ot.h_mid[0 .. n * wide]); + var same: usize = 0; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = ot.picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same += @intFromBool(std.mem.eql(u16, &x, &y)); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], ot.mlp_branch[0 .. n * dims]); + const e_out = errOf(@as([]const u16, hg), ho); + std.debug.print(" rows {d}..{d}: mixer branch rel {e:.3} (peak {d:.3}); streams after mixer rel {e:.3}; MoE picks same set {d}/{d}, MoE branch rel {e:.3} (peak {d:.3}); layer output streams rel {e:.3} (peak {d:.2}, max_abs {e:.2})\n", .{ c[0], c[0] + n - 1, e_mix.rel, e_mix.peak, e_mid.rel, same, n, e_moe.rel, e_moe.peak, e_out.rel, e_out.peak, e_out.max_abs }); + worst.mixer = @max(worst.mixer, e_mix.rel); + worst.streams = @max(worst.streams, e_mid.rel); + worst.picks_same += same; + worst.rows += n; + if (same == n) { + worst.moe = @max(worst.moe, e_moe.rel); + worst.out = @max(worst.out, e_out.rel); + } + } + std.debug.print(" device time for the 5 calls (read-backs between blocks included): {d:.0} ms\n", .{device_ms}); + std.debug.print(" worst staged: attn readout mixed {e:.3} inject {e:.3}; mixer {e:.3}; mlp readout mixed {e:.3} inject {e:.3}; MoE {e:.3}; write-back {e:.3}\n", .{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }); + std.debug.print(" worst chained: mixer {e:.3}, streams {e:.3}, MoE {e:.3}, output {e:.3}\n", .{ worst.mixer, worst.streams, worst.moe, worst.out }); + // Each block alone, fed the device's own input: its bf16 output and the roundings inside it (the read-out's act and + // inject, the projections' bf16 rows, the MoE weights' bf16) stay under 1% of the peak. + const staged = [_]f64{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }; + for (staged) |v| try check.expect(v <= 0.01, "layer {d}: a block alone differs from float64 by {d:.4} of its peak (> 0.01)", .{ l, v }); + // Chained from the embedding: the old engine's gates (mixer 1% GDN / 3% attention, streams 1%, MoE 2%) hold; its + // output gate was 1% with f32 residual streams, and upstream keeps the streams in bf16 between blocks, so the + // write-back's rounding lands on top of the chained MoE error: the output gate here is 2% (layer 0 measured 1.23e-2). + const mixer_tol: f64 = if (wt.linearLayer(l)) 0.01 else 0.03; + try check.expect(worst.mixer <= mixer_tol, "layer {d}: mixer branch differs from the oracle ({d:.4} > {d})", .{ l, worst.mixer, mixer_tol }); + try check.expect(worst.streams <= 0.01, "layer {d}: streams after the mixer differ ({d:.4})", .{ l, worst.streams }); + try check.expect(worst.moe <= 0.02, "layer {d}: MoE branch differs from the oracle ({d:.4})", .{ l, worst.moe }); + try check.expect(worst.out <= 0.02, "layer {d}: layer output streams differ ({d:.4})", .{ l, worst.out }); + if (worst.out > 0.01) std.debug.print(" NOTE layer {d}: chained output {e:.3} is above the old engine's 1% gate (f32 streams there)\n", .{ l, worst.out }); + check.pass("layer {d} ({s}) through the family's pass (writeback, readout, {s}, write-back, readout, EXL3 moeBlock, write-back) vs float64: mixer rel {e:.3}, streams rel {e:.3}, MoE rel {e:.3}, output rel {e:.3}; picks equal on {d}/{d} rows", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", if (wt.linearLayer(l)) "gdnBlock" else "attnBlock", worst.mixer, worst.streams, worst.moe, worst.out, worst.picks_same, worst.rows }); +} + +fn layerMode(gpu: Gpu, dir: []const u8, which: []const usize) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var prng = std.Random.DefaultPrng.init(0xe713_c03e); + var failed: usize = 0; + for (which) |l| { + const t0 = check.now(gpu.io); + var lw = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 1, 0, null); + defer lw.free(gpu.gpa); + std.debug.print("== layer {d} ({s}): the family's loader (cuda_prompt_load.loadLayerRank) in {d:.0} ms, every expert resident\n", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6 }); + try check.expect(lw.exl3 != null, "layer {d}: experts did not load as EXL3", .{l}); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + // every layer runs, a failed check reported and counted + projections(gpu, &rig, &raw, &lw, l, arena.allocator(), prng.random()) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + layerCheck(gpu, &rig, &raw, &lw, l, &default_ids) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + } + try check.expect(failed == 0, "flashnext-exl3 layer: {d} checks failed", .{failed}); +} + +/// The PLE key/value projections of an EXL3 pack take the n-gram rows as fp16 (cuda_exl3_ple.zig's default): through +/// Pass.lane() (exact split into two bf16 rows, two bf16 products, fp32 sums) against float64 of the fp16 rows times the +/// stored F16 weights; beside it, the error the rows would carry if rounded to bf16 first (TF_FLASHNEXT_PLE_EMB=bf16). +fn pleMode(gpu: Gpu, dir: []const u8) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const d: *cuda.Driver = @constCast(gpu.d); + // an EXL3 pack's store holds the shared Ops once its first EXL3 matrix loads; this check loads none + (try X.load.opsOf(&rig.store, gpu.d)).release(); + var key = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.key_proj", wide, dims); + defer key.free(); + var value = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.value_proj", dims, dims); + defer value.free(); + try X.load.pleRows(a, d, &rig.store, &key, &value); + try check.expect(key.exl3 != null and key.exl3.?.x == .f16 and value.exl3.?.x == .f16, "ple: key/value did not take fp16 rows", .{}); + var p = rig.pass(&.{}); + const kw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.key_proj.weight", .{})); + const vw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.value_proj.weight", .{})); + var prng = std.Random.DefaultPrng.init(0x91e_f16); + const rng = prng.random(); + var worst_key: f64 = 0; + var worst_value: f64 = 0; + var worst_bf16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xh = try a.alloc(f16, m * dims); + const xv = try a.alloc(f64, m * dims); + const xr = try a.alloc(u16, m * dims); // the same rows rounded to bf16 + for (xh, xv, xr) |*h, *v, *b| { + h.* = @floatCast(rng.floatNorm(f32) * 0.5); + v.* = @floatCast(h.*); + b.* = bf16Bits(@floatCast(h.*)); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer xd.free(); + const want_k = try a.alloc(f64, m * wide); + const want_v = try a.alloc(f64, m * dims); + for (0..m) |r| { + for (0..wide) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * kw[j * dims + i]; + want_k[r * wide + j] = acc; + } + for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * vw[j * dims + i]; + want_v[r * dims + j] = acc; + } + } + const got_k = try laneOut(gpu, &p, key, xd.ptr, m, 1); + const ek = errOf(got_k.bf.?, want_k); + // the value projection as pleBlock reads it: the sum of its 4 fp32 K slices + var slices = try cuda.DeviceBuffer.alloc(gpu.d, 4 * m * dims * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + try p.lane(xd.ptr, p.s.dummy, value, p.s.dummy, slices.ptr, m, 4); + const all = try download(f32, gpu, p.stream, slices.ptr, 4 * m * dims); + defer gpu.gpa.free(all); + const sum = try a.alloc(f64, m * dims); + for (sum, 0..) |*o, i| o.* = @as(f64, all[i]) + all[m * dims + i] + all[2 * m * dims + i] + all[3 * m * dims + i]; + const ev = errOf(sum, want_v); + // the bf16-row alternative, on the host in float64: rows rounded to bf16, weights exact + const alt = try a.alloc(f64, m * dims); + for (0..m) |r| for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += bfValue(xr[r * dims + i]) * vw[j * dims + i]; + alt[r * dims + j] = acc; + }; + const eb = errOf(alt, want_v); + std.debug.print(" ple rows={d}: key_proj bf16 out rel {e:.3}; value_proj fp32 slices rel {e:.3}; (bf16-rounded rows would give {e:.3})\n", .{ m, ek.rel, ev.rel, eb.rel }); + worst_key = @max(worst_key, ek.rel); + worst_value = @max(worst_value, ev.rel); + worst_bf16 = @max(worst_bf16, eb.rel); + } + // key: bf16 output rounding (2^-9 of a value); value: fp32 sums of exact products + try check.expect(worst_key <= 4e-3, "ple: key_proj from fp16 rows is {e:.3} of peak from float64", .{worst_key}); + try check.expect(worst_value <= 1e-5, "ple: value_proj from fp16 rows is {e:.3} of peak from float64", .{worst_value}); + check.pass("ple: fp16 n-gram rows through key/value (exact bf16 split): key {e:.3}, value {e:.3} of peak (bf16-rounded rows: {e:.3})", .{ worst_key, worst_value, worst_bf16 }); +} + +fn parseList(gpa: std.mem.Allocator, text: []const u8) ![]usize { + var list: std.ArrayList(usize) = .empty; + var it = std.mem.splitScalar(u8, text, ','); + while (it.next()) |tok| if (tok.len > 0) try list.append(gpa, try std.fmt.parseInt(usize, tok, 10)); + return list.toOwnedSlice(gpa); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + if (std.mem.eql(u8, mode, "names")) return names(gpu, dir); + if (std.mem.eql(u8, mode, "norms")) return norms(gpu, dir); + if (std.mem.eql(u8, mode, "ple")) return pleMode(gpu, dir); + if (std.mem.eql(u8, mode, "router")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3,17"); + defer gpu.gpa.free(which); + const rows = if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4096; + return router(gpu, dir, which, rows); + } + if (std.mem.eql(u8, mode, "layer")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3"); + defer gpu.gpa.free(which); + return layerMode(gpu, dir, which); + } + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3_tp_engine_tests.zig b/zig/tests/cuda/flashnext_exl3_tp_engine_tests.zig new file mode 100644 index 0000000000..d0a3a20f15 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3_tp_engine_tests.zig @@ -0,0 +1,788 @@ +//! flashnext-exl3-tpe: the Qwen3.8-Flash-Next EXL3 pack through the Flash Next family's Engine at tp=1 or tp=2 +//! (expert parallel, cuda_exl3_tp.zig) with each rank's own offload pool attached as the server attaches it +//! (cuda_exl3_pass.attach) and the lm_head split by vocabulary at tp=2 (cuda_exl3_load.shardHead). One process a card; +//! at tp=2 both ranks run the same calls in lockstep over the family's NCCL Collective (NCCL >= 2.30). +//! +//! flashnext-exl3-tpe RANK WORLD PORT DIR tf GOLDEN DUMP the golden's 16-id prompt teacher-forced; WORLD 1 writes +//! its bf16 logits to DUMP, WORLD 2 compares rank 0's to DUMP +//! (top-1 at every position, bit-equal count, max |d|) and +//! both to the golden +//! flashnext-exl3-tpe RANK WORLD PORT DIR tfc GOLDEN DUMP the same with the golden's 23 greedy ids after the prompt +//! (each continuation step teacher-forced, both margins) +//! flashnext-exl3-tpe RANK WORLD PORT DIR gen GOLDEN DUMP greedy 24 for each golden prompt; WORLD 1 writes the ids, +//! WORLD 2 compares (and the two ranks' ids) +//! flashnext-exl3-tpe RANK WORLD PORT DIR bench prefill 16/256/2048 rows, decode ms/token, each rank's +//! pool hit rate, lines and free memory +//! flashnext-exl3-tpe RANK 2 PORT DIR trace GOLDEN PREFIX [ROUNDS [SCHEDULE [REPEATS]]] +//! rounds per segment (default 1024); zero-based CSV prompt +//! schedule (default 0), repeated REPEATS times (default 1) +//! sequential shared pools, independent KV; not concurrency +//! binary trace per expert pool after exact lifetime-policy replay +//! flashnext-exl3-tpe 0 1 0 - policy FILES... compare lifetime, LRU, decayed frequency/admission and Belady +//! flashnext-exl3-tpe 0 1 0 - hit FILES... host replay of real router picks through the pool's planner +//! masked to each expert-parallel rank (RANKS: 1 or 2) +//! flashnext-exl3-tpe RANK 2 PORT - graph ncclAllReduce on the family's nonblocking communicator +//! captured in a CUDA graph and replayed on both ranks +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const tp = X.tp; +const fpool = flash.exl3_pool; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; +const cap_default: usize = 32768; +const out_rows: usize = 64; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +// -- the ranks' link ------------------------------------------------------------------------------------------------ + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !*Ranks { + const r = try gpu.gpa.create(Ranks); + errdefer gpu.gpa.destroy(r); + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 600_000); + errdefer link.close(); + r.* = .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + return r; + } + + fn close(r: *Ranks, gpa: std.mem.Allocator) void { + r.collective.deinit(); + r.link.close(); + gpa.destroy(r); + } + + /// Both ranks here; returns the other rank's `value`. + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +// -- the engine ----------------------------------------------------------------------------------------------------- + +const Eng = struct { + gpu: Gpu, + rank: usize, + world: usize, + ranks: ?*Ranks, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + seq_live: bool, + out: cuda.DeviceBuffer, + attached: xpass.Attached, + + fn init(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.rank = rank; + g.world = world; + g.ranks = if (world == 2) try Ranks.open(gpu, rank, port) else null; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + if (g.ranks) |r| g.e.collective = &r.collective; + const t1 = check.now(gpu.io); + _ = try g.e.reserve(@min(cap_default, g.e.chunk)); + _ = try g.e.reserveScores(cap_default, @min(cap_default, g.e.chunk)); + g.seq = try g.e.newSeq(cap_default); + g.seq_live = true; + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, try xpass.budgetFromEnv(0), null); + const t2 = check.now(gpu.io); + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + const head_n = g.e.head_q.n; + std.debug.print("RESULT rank {d}/{d} load: {d:.1} s (initRank {d:.1} s, pinned bank {d:.2} GiB copied in {d:.1} s); {d} lines a layer x {d} layers of {d:.2} MiB = {d:.2} GiB of pool; {d:.2} GiB free before the pools, {d:.2} GiB free after (card {d:.2} GiB); head columns {d}; chunk {d}\n", .{ rank, world, ms(t0, t2) / 1e3, ms(t0, t1) / 1e3, gib(g.attached.bank_bytes), @as(f64, @floatFromInt(g.attached.bank_ns)) / 1e9, g.attached.lines, g.attached.layers, @as(f64, @floatFromInt(g.attached.line_bytes)) / (1 << 20), gib(g.attached.device_bytes), gib(g.attached.free_before), gib(free), gib(total), head_n, g.e.chunk }); + return g; + } + + fn deinit(g: *Eng) void { + g.out.free(); + if (g.seq_live) g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + if (g.ranks) |r| r.close(g.gpu.gpa); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn stats(g: *Eng) fpool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } + + /// Both ranks reach this point; the other rank's value (or `v` at tp=1). + fn swap(g: *Eng, v: u64) !u64 { + return if (g.ranks) |r| r.swap(v) else v; + } + + fn forward(g: *Eng, ids: []const u32, hooks: prompt.batch.Hooks) !void { + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, hooks); + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + } + + fn round(g: *Eng, tok: u32) !u32 { + const q = [_]*prompt.Seq{&g.seq}; + var t = tok; + try g.e.round(flash.Tri, g.tri(), &q, (&t)[0..1], g.out.ptr); + if (g.ranks) |r| try r.collective.finish(g.stream); + return prompt.argmax(g.seq.logits); + } +}; + +fn printPool(g: *Eng, what: []const u8) void { + const s = g.stats(); + const routed = s.hit_slots + s.miss_slots; + std.debug.print("RESULT rank {d}/{d} pool {s}: {d} rounds, {d} passes, hit rate {d:.1}% ({d}/{d} routed slots of this rank), {d} copies, {d} evictions; host cut {d:.1} ms over {d} blocks, plan {d:.1} ms\n", .{ g.rank, g.world, what, s.rounds, s.passes, 100 * @as(f64, @floatFromInt(s.hit_slots)) / @as(f64, @floatFromInt(@max(routed, 1))), s.hit_slots, routed, s.misses, s.evictions, @as(f64, @floatFromInt(xpass.pool_time.cut_ns)) / 1e6, xpass.pool_time.blocks, @as(f64, @floatFromInt(xpass.pool_time.plan_ns)) / 1e6 }); +} + +// -- golden --------------------------------------------------------------------------------------------------------- + +const Golden = struct { + parsed: std.json.Parsed(std.json.Value), + + fn open(gpu: Gpu, path: []const u8) !Golden { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpu.gpa, text, .{}) }; + } + + fn prompts(g: *const Golden) []std.json.Value { + return g.parsed.value.object.get("prompts").?.array.items; + } +}; + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn num(v: std.json.Value) f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => std.math.nan(f64), + }; +} + +fn top1(row: []const u16) u32 { + var best: usize = 0; + for (row, 0..) |b, i| if (promote(b) > promote(row[best])) { + best = i; + }; + return @intCast(best); +} + +// -- modes ---------------------------------------------------------------------------------------------------------- + +fn margin(row: []const u16) f32 { + const b = top1(row); + var second: f32 = -std.math.inf(f32); + for (row, 0..) |v, i| if (i != b) { + second = @max(second, promote(v)); + }; + return promote(row[b]) - second; +} + +/// `cont`: the prompt followed by the golden's first 23 greedy ids (each continuation step teacher-forced). +fn tfMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, dump: []const u8, cont: bool) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const p = golden.prompts()[0]; + const prompt_ids = try ints(a, p.object.get("ids").?); + defer a.free(prompt_ids); + const greedy = try ints(a, p.object.get("greedy_ids").?); + defer a.free(greedy); + const ids = try a.alloc(u32, prompt_ids.len + if (cont) greedy.len - 1 else 0); + defer a.free(ids); + @memcpy(ids[0..prompt_ids.len], prompt_ids); + if (cont) @memcpy(ids[prompt_ids.len..], greedy[0 .. greedy.len - 1]); + const tf = p.object.get("teacher_forced").?.array.items; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + try g.reset(); + const n = ids.len; + const ops = g.e.layers[0].exl3.?.ops; + const before = ops.fingerprint(); + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .all, .logits_out = g.out.ptr }); + const t1 = check.now(gpu.io); + const logits = try a.alloc(u16, n * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits)); + std.debug.print("rank {d}/{d}: teacher-forced pass of {d} rows in {d:.1} ms (cold pool)\n", .{ rank, world, n, ms(t0, t1) }); + _ = try g.round(prompt.argmax(g.seq.logits)); + const after = ops.fingerprint(); + try check.expect(std.mem.eql(u8, std.mem.sliceAsBytes(&before), std.mem.sliceAsBytes(&after)), "rank {d}: an EXL3 scratch buffer grew inside a pass", .{rank}); + check.pass("rank {d}/{d}: EXL3 scratch (incl. expert parallel's) kept its addresses through a {d}-row prompt pass and a decode round", .{ rank, world, n }); + // both ranks hold the same gathered logits + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(logits)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "the two ranks' gathered logits differ", .{}); + if (rank != 0) return; + // vs the golden + var ok: usize = 0; + var worst: f64 = 0; + for (tf, 0..) |row, pos| { + const r = logits[pos * vocab ..][0..vocab]; + const g5 = try ints(a, row.object.get("top5_ids").?); + defer a.free(g5); + ok += @intFromBool(top1(r) == g5[0]); + for (g5, row.object.get("top5_logits").?.array.items) |id, lv| worst = @max(worst, @abs(@as(f64, promote(r[id])) - num(lv))); + } + std.debug.print("RESULT tp={d} TF vs golden: top-1 {d}/{d}, max |our logit - golden| on the golden top-5 {d:.4}\n", .{ world, ok, tf.len, worst }); + if (world == 1) { + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = dump, .data = std.mem.sliceAsBytes(logits) }); + check.pass("tp=1 logits written to {s}", .{dump}); + return; + } + const ref_bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, dump, a, .limited(1 << 30)); + defer a.free(ref_bytes); + if (ref_bytes.len != logits.len * 2) return error.UnexpectedTensor; + const ref = try a.alloc(u16, logits.len); + defer a.free(ref); + @memcpy(std.mem.sliceAsBytes(ref), ref_bytes); + var top_eq: usize = 0; + var bits: usize = 0; + var dmax: f64 = 0; + var dsum: f64 = 0; + std.debug.print(" pos tp=1 top1 (margin) tp=2 top1 (margin) max|d| over the row bit-equal of {d}\n", .{vocab}); + for (0..n) |pos| { + const x = ref[pos * vocab ..][0..vocab]; + const y = logits[pos * vocab ..][0..vocab]; + var row_d: f64 = 0; + var row_bits: usize = 0; + for (x, y) |u, v| { + row_bits += @intFromBool(u == v); + const d = @abs(@as(f64, promote(u)) - promote(v)); + row_d = @max(row_d, d); + dsum += d; + } + bits += row_bits; + dmax = @max(dmax, row_d); + const same = top1(x) == top1(y); + top_eq += @intFromBool(same); + std.debug.print(" {d:3} {d:9} ({d:6.3}) {d:9} ({d:6.3}) {d:.4} {d}{s}\n", .{ pos, top1(x), margin(x), top1(y), margin(y), row_d, row_bits, if (same) "" else " TOP-1 DIFFERS" }); + } + std.debug.print("RESULT tp=2 vs tp=1 logits: top-1 equal {d}/{d}; bf16 bit-equal {d}/{d} ({d:.3}%); max |d| {d:.4}, mean |d| {e:.3}\n", .{ top_eq, n, bits, n * vocab, 100 * @as(f64, @floatFromInt(bits)) / @as(f64, @floatFromInt(n * vocab)), dmax, dsum / @as(f64, @floatFromInt(n * vocab)) }); + try check.expect(top_eq == n, "tp=2 top-1 differs from tp=1", .{}); + check.pass("tp=2 teacher-forced top-1 equals tp=1 at all {d} positions", .{n}); + printPool(g, "after tf"); +} + +fn genMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, dump: []const u8) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + var text: std.ArrayList(u8) = .empty; + defer text.deinit(a); + var all_golden = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const want = try ints(a, p.object.get("greedy_ids").?); + defer a.free(want); + const got = try a.alloc(u32, want.len); + defer a.free(got); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + got[0] = prompt.argmax(g.seq.logits); + const t0 = check.now(gpu.io); + for (1..got.len) |i| got[i] = try g.round(got[i - 1]); + const step = ms(t0, check.now(gpu.io)) / @as(f64, @floatFromInt(got.len - 1)); + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(got)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "the two ranks' greedy ids differ", .{}); + var first: ?usize = null; + for (got, want, 0..) |x, y, i| if (x != y and first == null) { + first = i; + }; + if (first != null) all_golden = false; + for (got) |x| try text.print(a, "{d},", .{x}); + try text.append(a, '\n'); + if (rank == 0) { + std.debug.print("GEN tp={d} {s}: ", .{ world, p.object.get("name").?.string }); + for (got) |x| std.debug.print("{d},", .{x}); + std.debug.print(" {s} ({d:.1} ms/step)\n", .{ if (first) |_| "differs from the golden" else "equal to the golden", step }); + if (first) |d| std.debug.print(" first difference from the golden at step {d}: ours {d}, golden {d}\n", .{ d, got[d], want[d] }); + } + } + if (rank != 0) return; + printPool(g, "after gen"); + if (world == 1) { + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = dump, .data = text.items }); + check.pass("tp=1 greedy ids written to {s}", .{dump}); + return; + } + const ref = try std.Io.Dir.cwd().readFileAlloc(gpu.io, dump, a, .limited(1 << 20)); + defer a.free(ref); + const same = std.mem.eql(u8, ref, text.items); + std.debug.print("RESULT tp=2 greedy 24 on the {d} golden prompts: {s} tp=1's ids; all equal the golden: {}\n", .{ golden.prompts().len, if (same) "identical to" else "DIFFERENT from", all_golden }); + try check.expect(same, "tp=2 greedy ids differ from tp=1's", .{}); + check.pass("tp=2 greedy 24 tokens identical to tp=1 on all prompts", .{}); +} + +fn benchMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const max_n: usize = 2048; + const ids = try a.alloc(u32, max_n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x5eed); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + try g.reset(); + try g.forward(ids[0..64], .{ .logits = .last }); + for ([_]usize{ 16, 256, 2048 }) |n| { + var best: f64 = std.math.inf(f64); + for (0..2) |_| { + try g.reset(); + _ = try g.swap(0); + const before = g.stats(); + const t0 = check.now(gpu.io); + try g.forward(ids[0..n], .{ .logits = .last }); + const t = ms(t0, check.now(gpu.io)); + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" rank {d}/{d} prefill {d:5} tokens: {d:8.1} ms = {d:7.1} rows/s (hit {d:.1}%, {d} copies)\n", .{ rank, world, n, t, @as(f64, @floatFromInt(n)) / (t / 1e3), 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), after.misses - before.misses }); + best = @min(best, t); + } + std.debug.print("BENCH rank {d}/{d} prefill {d} tokens: best {d:.1} ms, {d:.1} rows/s\n", .{ rank, world, n, best, @as(f64, @floatFromInt(n)) / (best / 1e3) }); + var tok = prompt.argmax(g.seq.logits); + _ = try g.swap(1); + const before = g.stats(); + const t0 = check.now(gpu.io); + const steps: usize = 32; + for (0..steps) |_| tok = try g.round(tok); + const t = ms(t0, check.now(gpu.io)) / steps; + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print("BENCH rank {d}/{d} decode after {d} tokens: {d:.2} ms/token over {d} steps (this rank's hit {d:.1}%, {d:.1} copies a token)\n", .{ rank, world, n, t, steps, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), @as(f64, @floatFromInt(after.misses - before.misses)) / steps }); + } + printPool(g, "after bench"); +} + +fn traceSchedule(allocator: std.mem.Allocator, text: []const u8, prompts: usize) ![]usize { + var result: std.ArrayList(usize) = .empty; + errdefer result.deinit(allocator); + var fields = std.mem.splitScalar(u8, text, ','); + while (fields.next()) |field| { + const index = std.fmt.parseInt(usize, field, 10) catch return error.InvalidTraceSchedule; + if (index >= prompts) return error.InvalidTraceSchedule; + try result.append(allocator, index); + } + return result.toOwnedSlice(allocator); +} + +test "trace schedule preserves prompt order and repeated selections" { + const schedule = try traceSchedule(std.testing.allocator, "2,0,2,1", 3); + defer std.testing.allocator.free(schedule); + try std.testing.expectEqualSlices(usize, &.{ 2, 0, 2, 1 }, schedule); + for ([_][]const u8{ "", "0,", "3", "-1" }) |invalid| + try std.testing.expectError(error.InvalidTraceSchedule, traceSchedule(std.testing.allocator, invalid, 3)); +} + +fn traceMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, prefix: []const u8, rounds: usize, schedule_text: []const u8, repeats: usize) !void { + if (rounds == 0) return error.InvalidTraceRounds; + if (repeats == 0) return error.InvalidTraceSchedule; + if (world != 2) return error.TraceRequiresTp2; + if (!X.graph.enabled()) return error.TraceRequiresGraphs; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const a = gpu.gpa; + const schedule = try traceSchedule(a, schedule_text, golden.prompts().len); + defer a.free(schedule); + const segments = try std.math.mul(usize, schedule.len, repeats); + const total_rounds = try std.math.mul(usize, segments, rounds); + var record_bound: usize = 0; + for (schedule) |index| { + const count = golden.prompts()[index].object.get("ids").?.array.items.len; + if (count == 0 or count >= cap_default or rounds >= cap_default - count) return error.InvalidTraceRounds; + record_bound = try std.math.add(usize, record_bound, try std.math.add(usize, count, rounds)); + } + record_bound = try std.math.mul(usize, record_bound, repeats); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const initial_planner = (try X.dplan.forLayers(a, gpu.d, g.e, g.e.layers[0..g.e.n_layers])) orelse return error.TracePlannerMissing; + if (initial_planner.trace_mem == null) return error.TraceDisabled; + for (initial_planner.pools, initial_planner.pool_states, 0..) |pool, entry, index| { + const header = try X.dplan.demand_trace.header(index, pool.opt.layer, entry.lines, initial_planner.max_rows); + if (record_bound > header[8]) return error.TraceScheduleExceedsCapacity; + } + std.debug.print("TRACE sequential shared-pool workload (not concurrent): schedule={s}, repeats={d}, segments={d}, decode rounds={d}, conservative record bound={d}\n", .{ schedule_text, repeats, segments, total_rounds, record_bound }); + const generated = try a.alloc(u32, try std.math.add(usize, rounds, 1)); + defer a.free(generated); + var output_hash: u64 = 0; + for (0..segments) |segment| { + const selection = schedule[segment % schedule.len]; + if (segment != 0) { + try gpu.d.synchronize(); + g.seq.deinit(a); + g.seq_live = false; + g.seq = try g.e.newSeq(cap_default); + g.seq_live = true; + } + const ids = try ints(a, golden.prompts()[selection].object.get("ids").?); + defer a.free(ids); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var token = prompt.argmax(g.seq.logits); + const before = X.graph.coverage[0][1][@backingInt(X.graph.Outcome.replay)]; + generated[0] = token; + for (1..generated.len) |index| { + token = try g.round(token); + generated[index] = token; + } + if (X.graph.coverage[0][1][@backingInt(X.graph.Outcome.replay)] == before) return error.TraceRequiresGraphReplay; + const segment_hash = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(generated)); + try check.expect(try g.swap(segment_hash) == segment_hash, "trace segment {d}: tp=2 output hashes differ", .{segment}); + output_hash = std.hash.Wyhash.hash(output_hash, std.mem.sliceAsBytes(generated)); + std.debug.print("TRACE segment={d} prompt={d} rounds={d} output={x}\n", .{ segment, selection, rounds, segment_hash }); + } + try check.expect(try g.swap(output_hash) == output_hash, "trace run rank {d}: tp=2 output hashes differ", .{rank}); + if (X.dplan.current()) |planner| { + if (planner != initial_planner) return error.TracePlannerChanged; + try X.dplan.toHost(); + var decoder_pools: usize = 0; + for (planner.pools, 0..) |pool, index| { + var snapshot = try planner.downloadTraceAfterConsumers(a, index); + defer snapshot.deinit(); + const view = try snapshot.view(); + if (view.dropped() != 0) return error.TraceDroppedRecords; + if (pool.opt.layer < 48) { + decoder_pools += 1; + if (view.recorded() < total_rounds) return error.TraceMissingDecoderPool; + } else if (view.recorded() == 0) continue; + const entry = planner.pool_states[index]; + const terminal_words = try a.alloc(u32, entry.words); + defer a.free(terminal_words); + try planner.state.download(entry.byteOffset(), std.mem.sliceAsBytes(terminal_words)); + const replay = try X.dplan.replay.verifyTraceWithFinal(a, view, terminal_words); + const bytes = try view.encode(a); + defer a.free(bytes); + const path = try std.fmt.allocPrintSentinel(a, "{s}.rank-{d}.pool-{d}.exdpt", .{ prefix, rank, index }, 0); + defer a.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = bytes }); + std.debug.print("TRACE v{d} rank {d} pool {d} layer {d}: rounds={d}/{d}, checkpoints={d}, verified transitions={d}, picks hit/miss={d}/{d}, copies={d}, evictions={d}, in-place experts/picks={d}/{d}, bytes={d}, output={x}\n", .{ X.dplan.demand_trace.version, rank, index, pool.opt.layer, view.recorded(), view.words[8], view.checkpoints(), replay.matched_transitions, replay.delta.hit_picks, replay.delta.miss_picks, replay.delta.copies, replay.delta.evictions, replay.delta.inplace, replay.delta.mapped_picks, bytes.len, output_hash }); + } + if (decoder_pools != 48) return error.TraceDecoderPoolCount; + } else return error.TracePlannerMissing; +} + +// -- hit rates per rank (host only) --------------------------------------------------------------------------------- + +const top_k = 10; +const slots = 11; +const max_rows = 16; + +fn simulate(gpa: std.mem.Allocator, layer: u8, picks: []const u16, tokens: usize, prompt_n: usize, lines: usize, r: tp.Rank, counts: *[fpool.shape.layers][fpool.shape.experts]u32, dec: *[2]u64, all_: *[2]u64) !void { + const policy = fpool.offload; + const slots_tab = try gpa.alloc(u16, lines); + defer gpa.free(slots_tab); + const line_tab = try gpa.alloc(policy.Line, lines); + defer gpa.free(line_tab); + const stamp = try gpa.alloc(u32, lines); + defer gpa.free(stamp); + const lay = try gpa.alloc(u8, lines); + defer gpa.free(lay); + var v = policy.Vram.init(slots_tab, line_tab, stamp, lay); + var rnd = try fpool.Round.init(gpa, lines, max_rows * slots); + defer rnd.deinit(gpa); + var buf: [max_rows * slots]i32 = undefined; + var prev: [512]u16 = undefined; + var prev_n: usize = 0; + var t: usize = 0; + var round: u32 = 0; + while (t < tokens) { + const rows: usize = if (t < prompt_n) @min(max_rows, prompt_n - t) else 1; + for (0..rows) |i| { + for (0..top_k) |k| { + const id: usize = picks[(t + i) * top_k + k]; + buf[i * slots + k] = if (r.owns(id)) @intCast(id) else -1; + } + buf[i * slots + top_k] = if (r.owns(512)) 512 else -1; + } + round += 1; + try fpool.round.plan(&v, layer, buf[0 .. rows * slots], counts, round, round, .{ .protect = prev[0..prev_n] }, &rnd); + @memcpy(prev[0..rnd.n_ids], rnd.ids[0..rnd.n_ids]); + prev_n = rnd.n_ids; + all_[0] += rnd.hit_slots; + all_[1] += rnd.miss_slots; + if (t >= prompt_n) { + dec[0] += rnd.hit_slots; + dec[1] += rnd.miss_slots; + } + t += rows; + } +} + +fn rate(x: [2]u64) f64 { + return if (x[0] + x[1] == 0) 0 else @as(f64, @floatFromInt(x[0])) / @as(f64, @floatFromInt(x[0] + x[1])); +} + +fn hitMode(gpu: Gpu, files: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const counts = try gpa.create([fpool.shape.layers][fpool.shape.experts]u32); + defer gpa.destroy(counts); + for (files) |path| { + const blob = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpa, .limited(1 << 30)); + defer gpa.free(blob); + const hdr = std.mem.bytesAsSlice(u32, blob[0..12]); + const tokens: usize = hdr[0]; + const prompt_n: usize = hdr[1]; + const layers: usize = hdr[2]; + const raw = try gpa.alloc(u16, tokens * layers * top_k); + defer gpa.free(raw); + @memcpy(std.mem.sliceAsBytes(raw), blob[12..][0 .. raw.len * 2]); + const per = try gpa.alloc(u16, tokens * top_k); + defer gpa.free(per); + std.debug.print("real router picks {s}: {d} tokens ({d} prompt in 16-row rounds, {d} decode rows) x {d} layers\n", .{ path, tokens, prompt_n, tokens - prompt_n, layers }); + // routed slots a decode token each rank owns + var own: [2]u64 = .{ 0, 0 }; + for (prompt_n..tokens) |t| for (0..layers) |l| for (0..top_k) |k| { + own[@intFromBool(raw[(t * layers + l) * top_k + k] >= 256)] += 1; + }; + std.debug.print(" decode routed slots on rank 0 / rank 1: {d} / {d} ({d:.1}% on rank 0)\n", .{ own[0], own[1], 100 * @as(f64, @floatFromInt(own[0])) / @as(f64, @floatFromInt(own[0] + own[1])) }); + for ([_]usize{ 10, 24, 48, 56, 64, 96, 128, 160 }) |lines| { + var line_txt: [3][64]u8 = undefined; + var outs: [3][]const u8 = undefined; + const worlds = [3][2]usize{ .{ 1, 0 }, .{ 2, 0 }, .{ 2, 1 } }; + var miss_tok: [3]f64 = undefined; + for (worlds, 0..) |w, i| { + const r = try tp.Rank.init(w[0], w[1]); + counts.* = @splat(@splat(0)); + var dec: [2]u64 = .{ 0, 0 }; + var all_: [2]u64 = .{ 0, 0 }; + for (0..layers) |l| { + for (0..tokens) |t| @memcpy(per[t * top_k ..][0..top_k], raw[(t * layers + l) * top_k ..][0..top_k]); + try simulate(gpa, @intCast(l), per, tokens, prompt_n, lines, r, counts, &dec, &all_); + } + miss_tok[i] = @as(f64, @floatFromInt(dec[1])) / @as(f64, @floatFromInt(tokens - prompt_n)); + outs[i] = try std.fmt.bufPrint(&line_txt[i], "hit {d:.3} ({d:.1} misses/token)", .{ rate(dec), miss_tok[i] }); + } + std.debug.print(" {d:>3} lines a layer: tp=1 {s} | tp=2 rank 0 {s} | rank 1 {s} | slowest rank's misses/token {d:.1} vs tp=1 {d:.1}\n", .{ lines, outs[0], outs[1], outs[2], @max(miss_tok[1], miss_tok[2]), miss_tok[0] }); + } + } + check.pass("per-rank hit rates printed (host replay through the pool's own planner)", .{}); +} + +fn policyMode(gpu: Gpu, files: []const [:0]const u8) !void { + if (files.len == 0) return error.MissingTraceFiles; + const allocator = gpu.gpa; + for (files) |path| { + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, allocator, .limited(X.dplan.demand_trace.pool_bytes + 1)); + defer allocator.free(bytes); + var snapshot = try X.dplan.demand_trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + const view = try snapshot.view(); + if (view.dropped() != 0) return error.TraceDroppedRecords; + if (view.recorded() == 0) return error.EmptyTrace; + const first = try view.round(0); + const result = try X.dplan.policy_eval.evaluateTrace(allocator, view); + const lifetime = result.lifetime; + const lru = result.lru; + const decayed = result.decayed_frequency; + const admission = result.decayed_admission; + std.debug.print("POLICY {s} pool={d} layer={d} lines={d} rounds={d} drops={d}\n", .{ path, view.words[2], view.words[3], first.lines, view.recorded(), view.dropped() }); + if (view.checkpoints() != 0) std.debug.print(" {d} ownership checkpoints: each device segment starts from its recorded state; host intervals excluded\n", .{view.checkpoints()}); + std.debug.print(" lifetime picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ lifetime.miss_picks, lifetime.copies, lifetime.evictions, lifetime.mapped_experts }); + std.debug.print(" lru picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ lru.miss_picks, lru.copies, lru.evictions, lru.mapped_experts }); + std.debug.print(" decayed picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ decayed.miss_picks, decayed.copies, decayed.evictions, decayed.mapped_experts }); + std.debug.print(" guarded picks miss={d} copies={d} evictions={d} mapped={d} mapped_picks={d}\n", .{ admission.miss_picks, admission.copies, admission.evictions, admission.mapped_experts, admission.mapped_picks }); + std.debug.print(" Belady unique-miss lower bound={d} (not copy/time bound)\n", .{result.belady_unique_misses_lower_bound}); + } +} + +// -- NCCL in a captured graph --------------------------------------------------------------------------------------- + +fn ready(c: *Collective) !void { + while (true) { + var state: cuda.nccl.Result = 0; + try c.lib.check(c.api.ncclCommGetAsyncError(c.comm, &state), "ncclCommGetAsyncError"); + if (state == 0) return; + if (state != 7) return c.lib.check(state, "NCCL async"); + } +} + +fn allReduce(c: *Collective, s: cuda.Stream, src: u64, dst: u64, n: usize) !void { + const r = c.lib.api.ncclAllReduce(src, dst, n, .f32, .sum, c.comm, s.handle); + if (r != 0 and r != 7) return c.lib.check(r, "ncclAllReduce"); + try ready(c); +} + +fn graphMode(gpu: Gpu, rank: usize, port: u16) !void { + const a = gpu.gpa; + const rk = try Ranks.open(gpu, rank, port); + defer rk.close(a); + const c = &rk.collective; + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + var all_ok = true; + // one decode row's MoE sum (2,560 fp32), one row's per-slot y, and a 128-row block + for ([_]usize{ 2560, 28160, 327680 }) |n| for ([_]usize{ 1, 48 }) |per_graph| { + const mine = try a.alloc(f32, n * per_graph); + defer a.free(mine); + const theirs = try a.alloc(f32, n * per_graph); + defer a.free(theirs); + var p0 = std.Random.DefaultPrng.init(0x9a + rank); + var p1 = std.Random.DefaultPrng.init(0x9a + (1 - rank)); + for (mine) |*v| v.* = p0.random().floatNorm(f32); + for (theirs) |*v| v.* = p1.random().floatNorm(f32); + var src = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(mine)); + defer src.free(); + var dst = try cuda.DeviceBuffer.alloc(gpu.d, n * per_graph * 4); + defer dst.free(); + try dst.fill8(0, s.handle); + try s.synchronize(); + _ = try rk.swap(n); + // eager: per_graph all-reduces in a row + const reps: usize = if (per_graph == 1) 50 else 5; + for (0..3) |_| for (0..per_graph) |k| try allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n); + try s.synchronize(); + try e0.record(s); + for (0..reps) |_| for (0..per_graph) |k| try allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n); + try e1.record(s); + try e1.synchronize(); + const t_eager = (try e0.elapsedMs(e1)) / @as(f32, @floatFromInt(reps)); + try dst.fill8(0, s.handle); + try s.synchronize(); + _ = try rk.swap(n + 1); + // captured: the same calls in one graph, instantiated, replayed + const cap_ok = blk: { + cuda.graph.beginCapture(s, .thread_local) catch |err| { + std.debug.print("rank {d}: beginCapture failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + var failed = false; + for (0..per_graph) |k| allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n) catch |err| { + std.debug.print("rank {d}: ncclAllReduce under capture failed: {s}\n", .{ rank, @errorName(err) }); + failed = true; + break; + }; + var g = cuda.graph.endCapture(s) catch |err| { + std.debug.print("rank {d}: endCapture failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + defer g.deinit(); + if (failed) break :blk false; + const nodes = g.nodeCount() catch 0; + var exec = g.instantiate() catch |err| { + std.debug.print("rank {d}: instantiate failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + defer exec.deinit(); + exec.upload(s) catch {}; + for (0..3) |_| try exec.launchOn(s); + try s.synchronize(); + const got = try a.alloc(f32, n * per_graph); + defer a.free(got); + try dst.download(0, std.mem.sliceAsBytes(got)); + var eq: usize = 0; + for (got, mine, theirs) |x, m, t| { + const want = if (rank == 0) m + t else t + m; + eq += @intFromBool(@as(u32, @bitCast(x)) == @as(u32, @bitCast(want))); + } + try e0.record(s); + for (0..reps) |_| try exec.launchOn(s); + try e1.record(s); + try e1.synchronize(); + const t_graph = (try e0.elapsedMs(e1)) / @as(f32, @floatFromInt(reps)); + std.debug.print("RESULT rank {d} graph: {d} x ncclAllReduce of {d} fp32 ({d} graph nodes): eager {d:.4} ms, graph replay {d:.4} ms ({d:.4} vs {d:.4} ms an all-reduce); replayed sums == fp32 a+b on {d}/{d}\n", .{ rank, per_graph, n, nodes, t_eager, t_graph, t_eager / @as(f32, @floatFromInt(per_graph)), t_graph / @as(f32, @floatFromInt(per_graph)), eq, got.len }); + break :blk eq == got.len; + }; + all_ok = all_ok and cap_ok; + const both = try rk.swap(@intFromBool(cap_ok)); + all_ok = all_ok and both == 1; + }; + try check.expect(all_ok, "rank {d}: an all-reduce captured in a graph failed or gave other bits", .{rank}); + check.pass("rank {d}: ncclAllReduce on the nonblocking communicator captured in a CUDA graph replays on both ranks with the eager bits (SHM transport)", .{rank}); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 5) return error.MissingArgument; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const world = try std.fmt.parseInt(usize, rest[1], 10); + const port = try std.fmt.parseInt(u16, rest[2], 10); + const dir = rest[3]; + const mode = rest[4]; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (std.mem.eql(u8, mode, "hit")) return hitMode(gpu, rest[5..]); + if (std.mem.eql(u8, mode, "policy")) return policyMode(gpu, rest[5..]); + if (std.mem.eql(u8, mode, "graph")) return graphMode(gpu, rank, port); + if (std.mem.eql(u8, mode, "bench")) return benchMode(gpu, rank, world, port, dir); + if (std.mem.eql(u8, mode, "trace")) { + if (rest.len < 7) return error.MissingArgument; + const rounds = if (rest.len > 7) try std.fmt.parseInt(usize, rest[7], 10) else 1024; + if (rest.len > 10) return error.UnexpectedArgument; + const schedule = if (rest.len > 8) rest[8] else "0"; + const repeats = if (rest.len > 9) try std.fmt.parseInt(usize, rest[9], 10) else 1; + return traceMode(gpu, rank, world, port, dir, rest[5], rest[6], rounds, schedule, repeats); + } + if (rest.len < 7) return error.MissingArgument; + if (std.mem.eql(u8, mode, "tf")) return tfMode(gpu, rank, world, port, dir, rest[5], rest[6], false); + if (std.mem.eql(u8, mode, "tfc")) return tfMode(gpu, rank, world, port, dir, rest[5], rest[6], true); + if (std.mem.eql(u8, mode, "gen")) return genMode(gpu, rank, world, port, dir, rest[5], rest[6]); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3_tp_tests.zig b/zig/tests/cuda/flashnext_exl3_tp_tests.zig new file mode 100644 index 0000000000..e0cd566c0e --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3_tp_tests.zig @@ -0,0 +1,548 @@ +//! flashnext-exl3-tp: EXL3 expert parallel (cuda_exl3_tp.zig) on two ranks, one process a GPU, through the Flash Next +//! family's own loaders, pass and NCCL Collective (NCCL >= 2.30 on LD_LIBRARY_PATH). +//! Usage (run rank 1 and rank 0 together, CUDA_VISIBLE_DEVICES one card each): +//! tf-cuda-test flashnext-exl3-tp RANK PORT DIR layer [layers] whole decoder layers (default 0,3,17): tp=2 vs tp=1 on +//! rank 0 (picks, the exact fp32 regrouping, output bits) +//! and rank 0's tp=2 output vs float64 (the tp=1 gates) +//! tf-cuda-test flashnext-exl3-tp RANK PORT DIR sum the MoE sum at 1/128/512/2048 rows: all-reduce vs +//! rank-order send/recv + tf_rank_sum, bits and times +//! tf-cuda-test flashnext-exl3-tp 0 0 DIR head (one card) the lm_head vocabulary split vs the whole +//! tf-cuda-test flashnext-exl3-tp 0 0 DIR memory (one card) per-rank bytes and the 200k budgets +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const T = @import("flashnext_exl3_tests.zig"); + +const X = flash.exl3; +const tp = X.tp; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const Gpu = check.Gpu; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; + +// -- the two ranks' control link ----------------------------------------------------------------------------------- + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !Ranks { + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 120_000); + errdefer link.close(); + // rank 0 does float64 work between collectives: a generous deadline (the waits are bounded by the stream) + return .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + } + + fn close(r: *Ranks) void { + r.collective.deinit(); + r.link.close(); + } + + /// Both ranks here; rank 1's `value` reaches rank 0 (returned there), rank 0's reaches rank 1. + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +// -- one decoder layer, stepped ------------------------------------------------------------------------------------- + +/// What one call of a layer leaves: the output streams (bf16), and the MoE's picks, weights and per-slot y. +const Out = struct { h: []u16, picks: []i32, wts: []f32, y: []f32 }; + +fn grab(comptime E: type, a: std.mem.Allocator, gpu: Gpu, p: *T.Pass, ptr: u64, n: usize) ![]E { + const got = try T.download(E, gpu, p.stream, ptr, n); + defer gpu.gpa.free(got); + return a.dupe(E, got); +} + +/// cuda_batch.zig Pass.chunk's loop body for one layer (T.deviceLayer's sequence), the MoE's outputs read back. +fn step(gpu: Gpu, a: std.mem.Allocator, p: *T.Pass, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h_in: []const u16) !Out { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try T.uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h_in)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + if (lw.linear) try p.gdnBlock(lw, wt.linearIndex(l), .{}) else try p.attnBlock(lw, wt.attnIndex(l)); + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + try p.readout(lw.mlp_hc, 0, n); + try p.moeBlock(lw); + var o: Out = undefined; + o.picks = try grab(i32, a, gpu, p, s.picks, n * slots); + o.wts = try grab(f32, a, gpu, p, s.wts, n * slots); + o.y = try grab(f32, a, gpu, p, s.y, n * slots * dims); + p.pending = l; + try p.writeback(0, n); + o.h = try grab(u16, a, gpu, p, s.h, n * wide); + return o; +} + +/// prefill 8 rows then 4 single decode rows (T.layerCheck's calls) +const calls = [_][2]usize{ .{ 0, 8 }, .{ 8, 1 }, .{ 9, 1 }, .{ 10, 1 }, .{ 11, 1 } }; + +fn embedRows(a: std.mem.Allocator, raw: *const T.RawPack, toks: []const u32) ![]u16 { + const table = try raw.get("model.language_model.embed_tokens.weight", .{}); + const hb = try a.alloc(u16, toks.len * wide); + for (toks, 0..) |tk, t| for (0..dims) |i| { + const bits = std.mem.readInt(u16, table.bytes[(tk * dims + i) * 2 ..][0..2], .little); + for (0..streams) |sidx| hb[t * wide + sidx * dims + i] = bits; + }; + return hb; +} + +/// The 5 calls on `lw` from a reset rig. +fn sequence(gpu: Gpu, a: std.mem.Allocator, rig: *T.Rig, raw: *const T.RawPack, lw: *const pl.Layer, l: usize) ![calls.len]Out { + try rig.reset(); + var p = rig.pass(lw[0..1]); + var outs: [calls.len]Out = undefined; + for (calls, &outs) |c, *o| { + const toks = T.default_ids[c[0] .. c[0] + c[1]]; + o.* = try step(gpu, a, &p, lw, l, toks, c[0], try embedRows(a, raw, toks)); + } + return outs; +} + +fn hashOf(o: Out) u64 { + var h = std.hash.Wyhash.init(0x7e_2); + h.update(std.mem.sliceAsBytes(o.h)); + h.update(std.mem.sliceAsBytes(o.picks)); + // slot 0 of every row: the ranks' sum after placement + const n = o.picks.len / slots; + for (0..n) |r| h.update(std.mem.sliceAsBytes(o.y[r * slots * dims ..][0..dims])); + return h.final(); +} + +fn layerMode(gpu: Gpu, rk: *Ranks, dir: []const u8, which: []const usize) !void { + const rank = rk.rank; + var rig: T.Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try T.RawPack.open(gpu, dir); + defer raw.close(gpu); + const r0 = try tp.Rank.init(2, 0); + const r1 = try tp.Rank.init(2, 1); + var failed: usize = 0; + for (which) |l| { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var t0 = check.now(gpu.io); + var lw2 = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 2, rank, null); + defer lw2.free(gpu.gpa); + const x2 = lw2.exl3 orelse return error.UnexpectedTensor; + try check.expect(x2.tp != null and x2.tp.?.rank == rank, "layer {d}: rank {d} experts carry no rank", .{ l, rank }); + var held: usize = 0; + for (x2.trellis[0]) |ptr| held += @intFromBool(ptr != 0); + std.debug.print("== rank {d} layer {d} ({s}): tp=2 layer loaded in {d:.0} ms, {d} of 513 experts held (ids {d}..{d})\n", .{ rank, l, if (wt.linearLayer(l)) "GDN" else "attention", @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6, held, x2.tp.?.lo, x2.tp.?.hi - 1 }); + try check.expect(held == x2.tp.?.hi - x2.tp.?.lo, "layer {d}: rank {d} holds {d} experts", .{ l, rank, held }); + + // tp=1 on rank 0 alone (rank 1 waits at the swap) + var ref: [calls.len]Out = undefined; + if (rank == 0) { + t0 = check.now(gpu.io); + var lw1 = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 1, 0, null); + defer lw1.free(gpu.gpa); + rig.collective = null; + ref = try sequence(gpu, a, &rig, &raw, &lw1, l); + std.debug.print(" tp=1 reference (all 513 experts on this card) in {d:.0} ms\n", .{@as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6}); + } + _ = try rk.swap(0); + rig.collective = &rk.collective; + t0 = check.now(gpu.io); + const got = try sequence(gpu, a, &rig, &raw, &lw2, l); + const tp_ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + // both ranks must hold the same bits (replicated streams): rank 1's hashes to rank 0 + var same_ranks = true; + for (got) |o| { + const theirs = try rk.swap(hashOf(o)); + if (rank == 0) same_ranks = same_ranks and theirs == hashOf(o); + } + if (rank == 1) { + check.pass("rank 1 layer {d}: tp=2 mirrored the 5 calls in {d:.0} ms", .{ l, tp_ms }); + } else { + try check.expect(same_ranks, "layer {d}: the ranks' outputs or sums differ in bits", .{l}); + var picks_equal = true; + var sum_exact: usize = 0; + var sum_total: usize = 0; + var branch_abs: f64 = 0; + var branch_peak: f64 = 0; + var h_equal: usize = 0; + var h_total: usize = 0; + var h_abs: f64 = 0; + var h_peak: f64 = 0; + for (ref, got, calls) |want, have, c| { + const n = c[1]; + picks_equal = picks_equal and std.mem.eql(i32, want.picks, have.picks); + for (0..n) |r| for (0..dims) |d| { + // tp=1: the write-back's chain over the 11 slots; tp=2: each rank's chain over its own, then a + b + var one: f32 = 0; + var acc: [2]f32 = .{ 0, 0 }; + for (0..slots) |k| { + const at = r * slots + k; + const w = want.wts[at]; + const v = want.y[at * dims + d]; + one = @mulAdd(f32, v, w, one); + const id: usize = @intCast(want.picks[at]); + if (r0.owns(id)) acc[0] = @mulAdd(f32, w, v, acc[0]); + if (r1.owns(id)) acc[1] = @mulAdd(f32, w, v, acc[1]); + } + const predicted = acc[0] + acc[1]; + const placed = have.y[r * slots * dims + d]; + sum_exact += @intFromBool(placed == predicted); + sum_total += 1; + branch_abs = @max(branch_abs, @abs(@as(f64, placed) - one)); + branch_peak = @max(branch_peak, @abs(@as(f64, one))); + }; + for (want.h, have.h) |u, v| { + h_equal += @intFromBool(u == v); + h_total += 1; + h_abs = @max(h_abs, @abs(T.bfValue(u) - T.bfValue(v))); + h_peak = @max(h_peak, @abs(T.bfValue(u))); + } + } + std.debug.print(" tp=2 ({d:.0} ms for the 5 calls): picks {s}; summed MoE branch equals (rank 0's chain + rank 1's chain) of the tp=1 per-slot outputs on {d}/{d} values; vs tp=1's one chain max_abs {e:.3} (rel {e:.3}, peak {d:.3}); layer output bf16 equal {d}/{d}, max_abs {e:.3} (rel {e:.3})\n", .{ tp_ms, if (picks_equal) "equal" else "DIFFER", sum_exact, sum_total, branch_abs, branch_abs / branch_peak, branch_peak, h_equal, h_total, h_abs, h_abs / h_peak }); + var ok = true; + check.expect(picks_equal, "layer {d}: tp=2 router picks differ from tp=1", .{l}) catch { + ok = false; + }; + check.expect(sum_exact == sum_total, "layer {d}: tp=2's MoE sum is not the fp32 regrouping of tp=1's slots ({d}/{d})", .{ l, sum_exact, sum_total }) catch { + ok = false; + }; + // an fp32 regrouping of 11 terms: a few ulps of the branch's peak; bf16 outputs then differ by at most 1 ulp + check.expect(branch_abs <= 1e-5 * branch_peak, "layer {d}: tp=2 MoE branch off tp=1 by {e:.3} of peak", .{ l, branch_abs / branch_peak }) catch { + ok = false; + }; + check.expect(h_abs <= 8e-3 * h_peak, "layer {d}: tp=2 layer output off tp=1 by {e:.3} of peak", .{ l, h_abs / h_peak }) catch { + ok = false; + }; + if (ok) check.pass("layer {d} tp=2 vs tp=1: picks equal, MoE sum = fp32 regrouping of tp=1's slots on every value (bit-exact), output bf16 equal {d}/{d}, max diff {e:.3} of peak; ranks bit-identical", .{ l, h_equal, h_total, h_abs / h_peak }) else failed += 1; + } + // tp=2 against float64 (T.layerCheck on rank 0; rank 1 mirrors its device calls) + _ = try rk.swap(1); + if (rank == 0) { + T.layerCheck(gpu, &rig, &raw, &lw2, l, &T.default_ids) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + } else { + _ = try sequence(gpu, a, &rig, &raw, &lw2, l); + } + _ = try rk.swap(2); + } + const all = try rk.swap(failed); + try check.expect(failed == 0 and all == 0, "flashnext-exl3-tp layer: {d} checks failed on rank {d}", .{ failed, rank }); +} + +// -- the MoE sum's cost --------------------------------------------------------------------------------------------- + +fn fill(a: std.mem.Allocator, rows: usize, seed: u64) ![]f32 { + const out = try a.alloc(f32, rows * dims); + var prng = std.Random.DefaultPrng.init(seed); + for (out) |*v| v.* = prng.random().floatNorm(f32) * 0.05; + return out; +} + +fn waitComm(c: *Collective, s: cuda.Stream) !void { + try ready(c); + try s.synchronize(); +} + +/// A nonblocking communicator takes its next call only once the last one's enqueue finished (state 0, not 7). +fn ready(c: *Collective) !void { + while (true) { + var state: cuda.nccl.Result = 0; + try c.lib.check(c.api.ncclCommGetAsyncError(c.comm, &state), "ncclCommGetAsyncError"); + if (state == 0) break; + if (state != 7) { + std.debug.print("waitComm: NCCL async state {d}\n", .{state}); + return c.lib.check(state, "NCCL async"); + } + } +} + +/// Rank order: the peer's partial into `scratch` (send/recv), then out = rank 0's term + rank 1's (tf_rank_sum). +fn rankOrder(c: *Collective, k: *const cuda.rank_scatter.RankScatter, s: cuda.Stream, rank: usize, partial: u64, scratch: u64, out: u64, n: usize) !void { + const peer: c_int = @intCast(1 - rank); + const api = c.lib.api; + const ok = struct { + fn f(r: cuda.nccl.Result) !void { + if (r != 0 and r != 7) { + std.debug.print("rank-order: NCCL result {d}\n", .{r}); + return error.NcclFailed; + } + } + }.f; + try ok(api.ncclGroupStart()); + try ok(api.ncclSend(partial, n, .f32, peer, c.comm, s.handle)); + try ok(api.ncclRecv(scratch + (1 - rank) * n * 4, n, .f32, peer, c.comm, s.handle)); + try ok(api.ncclGroupEnd()); + try ready(c); + try k.rankSum(s, scratch, partial, 2, rank, n, out); +} + +fn timeIt(s: cuda.Stream, e0: *cuda.Event, e1: *cuda.Event, reps: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try e0.record(s); + for (0..reps) |_| try f(ctx); + try e1.record(s); + try e1.synchronize(); + best = @min(best, @as(f64, try e0.elapsedMs(e1.*)) / @as(f64, @floatFromInt(reps))); + } + return best; +} + +fn sumMode(gpu: Gpu, rk: *Ranks) !void { + const rank = rk.rank; + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + var rs = try cuda.rank_scatter.RankScatter.load(gpu.d); + defer rs.unload(); + var ops = try X.load.Ops.create(gpu.gpa, gpu.d); + defer ops.release(); + const t = try tp.scratchOf(ops); + const c = &rk.collective; + const sizes = [_]usize{ 1, 128, 512, 2048 }; + var all_equal = true; + for (sizes) |rows| { + const n = rows * dims; + const mine = try fill(a, rows, 0x5eed + rank); + const theirs = try fill(a, rows, 0x5eed + (1 - rank)); + var part = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(mine)); + defer part.free(); + var buf = try cuda.DeviceBuffer.alloc(gpu.d, n * 4); + defer buf.free(); + var scratch = try cuda.DeviceBuffer.alloc(gpu.d, 2 * n * 4); + defer scratch.free(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, n * 4); + defer out.free(); + // bits: the all-reduce, the rank-order sum and the host's rank0 + rank1 in fp32 + try s.d.check(s.d.api.cuMemcpyDtoDAsync_v2(buf.ptr, part.ptr, n * 4, s.handle), "cuMemcpyDtoDAsync"); + try c.sum(s, buf.ptr, n); + try c.finish(s); + try rankOrder(c, &rs, s, rank, part.ptr, scratch.ptr, out.ptr, n); + try waitComm(c, s); + const ar = try a.alloc(f32, n); + try buf.download(0, std.mem.sliceAsBytes(ar)); + const ro = try a.alloc(f32, n); + try out.download(0, std.mem.sliceAsBytes(ro)); + var eq_host: usize = 0; + var eq_ro: usize = 0; + for (0..n) |i| { + const want = if (rank == 0) mine[i] + theirs[i] else theirs[i] + mine[i]; + eq_host += @intFromBool(@as(u32, @bitCast(ar[i])) == @as(u32, @bitCast(want))); + eq_ro += @intFromBool(@as(u32, @bitCast(ar[i])) == @as(u32, @bitCast(ro[i]))); + } + all_equal = all_equal and eq_host == n and eq_ro == n; + // times: the pass's Collective.sum over [rows, 2560] fp32; rank order; upstream's per-slot [rows, 11, 2560] + const reps: usize = if (rows <= 128) 100 else if (rows <= 512) 20 else 5; + const Ctx = struct { c: *Collective, rs: *const cuda.rank_scatter.RankScatter, s: cuda.Stream, rank: usize, p: u64, sc: u64, o: u64, n: usize, t: *const tp.Scratch, picks: u64, wts: u64, y: u64, rows: usize }; + var per_slot = try cuda.DeviceBuffer.alloc(gpu.d, n * slots * 4); + defer per_slot.free(); + try per_slot.fill8(0, s.handle); + var picks = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer picks.free(); + try picks.fill8(0, s.handle); + var wts = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4 * 2); + defer wts.free(); + try wts.fill8(0, s.handle); + try s.synchronize(); + const ctx: Ctx = .{ .c = c, .rs = &rs, .s = s, .rank = rank, .p = part.ptr, .sc = scratch.ptr, .o = out.ptr, .n = n, .t = t, .picks = picks.ptr, .wts = wts.ptr, .y = per_slot.ptr, .rows = rows }; + const t_ar = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.c.sum(x.s, x.o, x.n); + } + }.f); + const t_raw = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + const r = x.c.lib.api.ncclAllReduce(x.o, x.o, x.n, .f32, .sum, x.c.comm, x.s.handle); + if (r != 0 and r != 7) return error.NcclFailed; + if (r == 7) try ready(x.c); + } + }.f); + const t_ro = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try rankOrder(x.c, x.rs, x.s, x.rank, x.p, x.sc, x.o, x.n); + } + }.f); + try waitComm(c, s); + const t_slot = try timeIt(s, &e0, &e1, @max(1, reps / 4), ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.c.sum(x.s, x.y, x.n * slots); + } + }.f); + const t_mask = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.t.mask(x.s, x.picks, x.wts, x.picks, x.wts + x.rows * slots * 4, x.rows * slots, try tp.Rank.init(2, x.rank)); + } + }.f); + const t_place = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.t.place(x.s, x.o, x.y, x.wts, x.rows); + } + }.f); + if (rank == 0) std.debug.print("RESULT tp2 MoE sum rows {d:>4} ({d:>9} B): all-reduce {d:.4} ms (raw ncclAllReduce, no peer/async checks: {d:.4}) | rank-order send/recv+sum {d:.4} ms | upstream per-slot all-reduce ({d:>9} B) {d:.4} ms | mask {d:.4} ms | place {d:.4} ms | bits: all-reduce == host fp32 a+b {d}/{d}, == rank-order {d}/{d}\n", .{ rows, n * 4, t_ar, t_raw, t_ro, n * slots * 4, t_slot, t_mask, t_place, eq_host, n, eq_ro, n }); + _ = try rk.swap(rows); + } + try check.expect(all_equal, "rank {d}: an all-reduce sum differs from the fp32 a + b or the rank-order sum", .{rank}); + check.pass("rank {d}: two-rank fp32 all-reduce of the combined MoE branch is bit-identical to rank-order a + b at 1..2048 rows", .{rank}); +} + +// -- the head's vocabulary split (one card) ------------------------------------------------------------------------- + +fn headMode(gpu: Gpu, dir: []const u8) !void { + var rig: T.Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + const d: *cuda.Driver = @constCast(gpu.d); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var full = try X.projection_load.load(gpu.gpa, d, &rig.store, "lm_head", wt.vocab, dims); + defer full.free(); + var halves: [2]flash.projection.Projection = undefined; + for (&halves, 0..) |*h, r| h.* = try tp.headProjection(gpu.gpa, d, &rig.store, "lm_head", 2, r); + defer for (&halves) |*h| h.free(); + var p = rig.pass(&.{}); + var prng = std.Random.DefaultPrng.init(0x4ead); + var all_equal = true; + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + for ([_]usize{ 1, 5 }) |m| { + const xb = try a.alloc(u16, m * dims); + for (xb) |*v| { + const f: f32 = prng.random().floatNorm(f32); + v.* = @truncate(@as(u32, @bitCast(f)) >> 16); + } + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer x.free(); + var yf = try cuda.DeviceBuffer.alloc(gpu.d, m * wt.vocab * 2); + defer yf.free(); + var yh = try cuda.DeviceBuffer.alloc(gpu.d, m * wt.vocab * 2); + defer yh.free(); + try p.lane(x.ptr, p.s.dummy, full, yf.ptr, p.s.dummy, m, 1); + try e0.record(p.stream); + try p.lane(x.ptr, p.s.dummy, full, yf.ptr, p.s.dummy, m, 1); + try e1.record(p.stream); + try e1.synchronize(); + const t_full = try e0.elapsedMs(e1); + var t_half: f32 = 0; + for (halves, 0..) |h, r| { + // the gathered layout: rank r's half at columns [r * N/2, (r + 1) * N/2) of every row + const tmp = try a.alloc(u16, m * h.n); + var yr = try cuda.DeviceBuffer.alloc(gpu.d, m * h.n * 2); + defer yr.free(); + try p.lane(x.ptr, p.s.dummy, h, yr.ptr, p.s.dummy, m, 1); + try e0.record(p.stream); + try p.lane(x.ptr, p.s.dummy, h, yr.ptr, p.s.dummy, m, 1); + try e1.record(p.stream); + try e1.synchronize(); + t_half = @max(t_half, try e0.elapsedMs(e1)); + try yr.download(0, std.mem.sliceAsBytes(tmp)); + for (0..m) |row| { + const dst = (row * wt.vocab + r * h.n) * 2; + try yh.upload(dst, std.mem.sliceAsBytes(tmp[row * h.n ..][0..h.n])); + } + } + const got = try a.alloc(u16, m * wt.vocab); + try yh.download(0, std.mem.sliceAsBytes(got)); + const want = try a.alloc(u16, m * wt.vocab); + try yf.download(0, std.mem.sliceAsBytes(want)); + var equal: usize = 0; + var worst: f64 = 0; + for (got, want) |g, w| { + equal += @intFromBool(g == w); + worst = @max(worst, @abs(T.bfValue(g) - T.bfValue(w))); + } + all_equal = all_equal and equal == got.len; + std.debug.print("RESULT head split rows {d}: halves (124,160 columns a rank) joined in vocabulary order equal the whole head on {d}/{d} bf16 logits (max diff {e:.3}); whole {d:.3} ms, one half {d:.3} ms\n", .{ m, equal, got.len, worst, t_full, t_half }); + } + try check.expect(all_equal, "head: the halves' logits differ from the whole head's", .{}); + check.pass("head: the EXL3 lm_head split at 128-column blocks (2 x 124,160) gives the whole head's bf16 logits bit for bit", .{}); +} + +// -- memory --------------------------------------------------------------------------------------------------------- + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +fn memoryMode(gpu: Gpu, dir: []const u8) !void { + const a = gpu.gpa; + const io = gpu.io; + const o = X.memory.options(); + const head = try X.memory.headBytes(a, io, dir, o); + const one = try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(1, 0)); + const ranks = [2]u64{ try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(2, 0)), try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(2, 1)) }; + const all_dev = try tp.deviceBytes(a, io, dir, .{ .routed_on_host = false }, try tp.Rank.init(2, 0)); + std.debug.print("RESULT memory (routed trellis {s}): tp=1 main {d:.3} GiB; tp=2 rank 0 {d:.3} GiB + MTP head {d:.3} GiB, rank 1 {d:.3} GiB; (rank 0 with its 256 routed experts resident: {d:.3} GiB)\n", .{ if (o.routed_on_host) "on the host (offload)" else "resident", gib(one), gib(ranks[0]), gib(head), gib(ranks[1]), gib(all_dev) }); + // the lm_head's split would save half its bytes on each rank + var store = try wt.Store.open(a, io, dir); + defer store.deinit(); + const lm = try X.load.mat(&store, "lm_head"); + const head_bytes: u64 = lm.trellis.len + lm.suh.len + lm.svh.len; + std.debug.print("RESULT memory lm_head {d:.3} GiB a rank replicated; split: {d:.3} GiB a rank\n", .{ gib(head_bytes), gib(head_bytes / 2) }); + const prompt = flash.prompt; + const window: usize = 200_000; + const seq = prompt.seqBytes(window); + const kv = wt.attn_n * window * 2 * prompt.kv_row_bytes; + const kv8 = kv / 2 + wt.attn_n * window * 2 * wt.kv_heads * 4; // int8 rows plus an fp32 scale a row and head + const draft = flash.mtp.Draft.bytes(window + flash.mtp.max_rows, 1, flash.prompt.batch.default_chunk); + const scratch = flash.prompt.batch.Scratch.bytes(flash.prompt.batch.default_chunk); + std.debug.print("RESULT memory a sequence at {d} tokens: {d:.3} GiB (bf16 K/V {d:.3} GiB of it; int8 K/V would be {d:.3}); MTP draft cache {d:.3} GiB; pass scratch at chunk {d}: {d:.3} GiB\n", .{ window, gib(seq), gib(kv), gib(kv8), gib(draft), flash.prompt.batch.default_chunk, gib(scratch) }); + check.pass("memory: per-rank placement read from the pack's headers", .{}); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 4) return error.MissingArgument; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const port = try std.fmt.parseInt(u16, rest[1], 10); + const dir = rest[2]; + const mode = rest[3]; + if (std.mem.eql(u8, mode, "head")) return headMode(gpu, dir); + if (std.mem.eql(u8, mode, "memory")) return memoryMode(gpu, dir); + var rk = try Ranks.open(gpu, rank, port); + defer rk.close(); + if (std.mem.eql(u8, mode, "sum")) return sumMode(gpu, &rk); + if (std.mem.eql(u8, mode, "layer")) { + var list: std.ArrayList(usize) = .empty; + defer list.deinit(gpu.gpa); + var it = std.mem.splitScalar(u8, if (rest.len > 4) rest[4] else "0,3,17", ','); + while (it.next()) |tok| if (tok.len > 0) try list.append(gpu.gpa, try std.fmt.parseInt(usize, tok, 10)); + return layerMode(gpu, &rk, dir, list.items); + } + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig index 974ce95d5e..5f590fe675 100644 --- a/zig/tests/cuda/flashnext_hc_tests.zig +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -369,6 +369,51 @@ fn writebackRounding(r: *Rig, g: *const flash.Glue) !void { } } +fn reducedWriteback(r: *Rig, g: *const flash.Glue) !void { + const rows: usize = 17; + const count = rows * dims; + const hidden = r.bfs(rows * wide, 1); + const inject = r.bfs(rows * streams, 1); + const branch = r.f32s(count, -2, 2); + const edges = [_]f32{ 0.0, -0.0, 0x1.01p0, -0x1.01p0, 0x1.03p0, -0x1.03p0, 0x1p-126, -0x1p-126 }; + for (branch, 0..) |*value, index| if (index % 7 == 0) { + value.* = edges[index % edges.len]; + }; + hidden[1] = 0x8000; + inject[0] = 0x3f80; + branch[1] = -0.0; + const y = r.f32s(count * 11, -13, 13); + const weights = try r.a.alloc(f32, rows * 11); + @memset(weights, 0); + const ones = try r.a.alloc(f32, rows); + @memset(ones, 1); + for (0..rows) |row| { + @memcpy(y[row * 11 * dims ..][0..dims], branch[row * dims ..][0..dims]); + weights[row * 11] = 1; + } + const dh = try r.dev(std.mem.sliceAsBytes(hidden)); + const di = try r.dev(std.mem.sliceAsBytes(inject)); + const dy = try r.dev(std.mem.sliceAsBytes(y)); + const dw = try r.dev(std.mem.sliceAsBytes(weights)); + const db = try r.dev(std.mem.sliceAsBytes(branch)); + const du = try r.dev(std.mem.sliceAsBytes(ones)); + const reference = try r.out(hidden.len * 2); + const result = try r.out(hidden.len * 2); + const reference_sums = try r.out(rows * (dims / 256) * streams * 4); + const result_sums = try r.out(rows * (dims / 256) * streams * 4); + try g.hcWriteback(r.s, dh, reference, reference_sums, 0, di, dy, dw, count, rows, dims, streams, 2, 10, 11, 1); + for ([_]usize{ 0, 3, 16 }, [_]usize{ 3, 13, 1 }) |at, nrows| { + try g.hcWriteback(r.s, dh + at * wide * 2, result + at * wide * 2, result_sums + at * (dims / 256) * streams * 4, 0, di + at * streams * 2, db + at * dims * 4, du + at * 4, nrows * dims, nrows, dims, streams, 2, 0, 1, 1); + } + const want = try r.back(u16, reference, hidden.len); + const got = try r.back(u16, result, hidden.len); + try check.expect(std.mem.eql(u16, want, got), "reduced writeback: hidden BF16 bits equal the placed 11-slot path", .{}); + const want_sums = try r.back(u32, reference_sums, rows * (dims / 256) * streams); + const got_sums = try r.back(u32, result_sums, want_sums.len); + try check.expect(std.mem.eql(u32, want_sums, got_sums), "reduced writeback: FP32 square-sum bits equal the placed 11-slot path", .{}); + check.pass("reduced writeback: exact hidden BF16 and FP32 square sums, finite stale slots, signed zero, BF16 ties and short tail", .{}); +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -377,6 +422,7 @@ pub fn run(gpu: check.Gpu) !void { defer r.deinit(); var g = try flash.Glue.load(gpu.d); defer g.deinit(); + try reducedWriteback(&r, &g); try writebackRounding(&r, &g); for ([_]usize{ 0, 2, 4 }) |mode| try writeback(&r, &g, mode); try readout(&r, &g); diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig index ebb937ae71..6f0d80d30e 100644 --- a/zig/tests/cuda/flashnext_model_rank_tests.zig +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -52,16 +52,16 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] const value: f32 = @bitCast(@as(u32, std.mem.readInt(u16, first.logits[i * 2 ..][0..2], .little)) << 16); try check.expect(std.math.isFinite(value), "finite initial logits at {d}", .{i}); } - const kept = try session.save(&first); + const kept = try session.save(&first, null); defer session.dropSaved(kept); const token = prompt.argmax(first.logits); try session.forward(flash.Tri, tri, &first, &.{token}, .{}); const expected = try gpu.gpa.dupe(u8, first.logits); defer gpu.gpa.free(expected); - try session.restore(&first, kept); + try session.restore(&first, null, kept); try session.forward(flash.Tri, tri, &first, &.{token}, .{}); try check.sameBytes("two-rank snapshot restore and replay", first.logits, expected); - try session.restore(&first, kept); + try session.restore(&first, null, kept); var second = try session.create(context); defer session.drop(&second); try session.forward(flash.Tri, tri, &second, ids.items, .{}); diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig index 999d773cd0..c915a19d2f 100644 --- a/zig/tests/cuda/flashnext_window_tests.zig +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -45,7 +45,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] var serial = try session.create(context); defer session.drop(&serial); try session.forward(flash.Tri, tri, &serial, ids.items, .{}); - const base = try session.save(&serial); + const base = try session.save(&serial, null); defer session.dropSaved(base); const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71, 99, 4, 80, 122, 7, 25, 61, 94 }; const row_bytes = prompt.vocabulary * 2; @@ -62,14 +62,14 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] const rows = try gpu.gpa.alloc(u8, expected.len); defer gpu.gpa.free(rows); for ([_]usize{ 1, 3, 7, 15, 16 }) |keep| { - try session.restore(&window, base); + try session.restore(&window, null, base); try session.verify(flash.Tri, tri, &window, &draft, &forced); try stream.synchronize(); try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(rows.ptr, draft.logits.ptr, rows.len), "cuMemcpyDtoH"); try check.sameBytes("each two-rank window row equals its serial teacher-forced row", rows, expected); try session.keep(flash.Tri, tri, &window, &draft, keep); try check.expect(window.fed == ids.items.len + keep, "kept prefix position is exact", .{}); - try session.restore(&serial, base); + try session.restore(&serial, null, base); for (forced[0..keep]) |token| try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); try check.expect(std.mem.eql(i64, &window.history, &serial.history), "kept host n-gram history matches serial", .{}); try session.forward(flash.Tri, tri, &serial, &.{13}, .{}); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..b626913306 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -37,6 +37,14 @@ const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); const exl3_tests = @import("exl3_tests.zig"); const exl3_tp2_tests = @import("exl3_tp2_tests.zig"); +const exl3_pack_tests = @import("exl3_pack_tests.zig"); +const exl3_ple2_tests = @import("exl3_ple2_tests.zig"); +const exl3_draft_tests = @import("exl3_draft_tests.zig"); +const exl3_pool_tests = @import("exl3_pool_tests.zig"); +const exl3_kv_tests = @import("exl3_kv_tests.zig"); +const exl3_kvpage_tests = @import("exl3_kvpage_tests.zig"); +const exl3_prefill_tests = @import("exl3_prefill_tests.zig"); +const exl3_tpprefill_tests = @import("exl3_tpprefill_tests.zig"); const exl3_candidate_tests = @import("exl3_candidate_tests.zig"); const experts3_tests = @import("experts3_tests.zig"); const moe_tests = @import("moe_tests.zig"); @@ -77,6 +85,11 @@ const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); +const flashnext_exl3_tests = @import("flashnext_exl3_tests.zig"); +const flashnext_exl3_tp_tests = @import("flashnext_exl3_tp_tests.zig"); +const flashnext_exl3_engine_tests = @import("flashnext_exl3_engine_tests.zig"); +const flashnext_exl3_tp_engine_tests = @import("flashnext_exl3_tp_engine_tests.zig"); +const flashnext_exl3_graph_tests = @import("flashnext_exl3_graph_tests.zig"); const ggml_tests = @import("ggml_tests.zig"); const paged_tests = @import("paged_tests.zig"); @@ -105,6 +118,16 @@ const usage = \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) + \\ flashnext-exl3 DIR names|norms|router [layers] [rows]|layer [layers] Flash Next on an EXL3 pack vs float64 + \\ flashnext-exl3-tp RANK PORT DIR layer [layers]|sum|head|memory EXL3 expert parallel on two ranks vs tp=1 and float64 + \\ flashnext-exl3-engine DIR tf|gen GOLDEN [LINES] | bench [LINES] the family's Engine with the offload pool vs the reference golden + \\ flashnext-exl3-tpe RANK 2 PORT DIR trace GOLDEN PREFIX [ROUNDS [SCHEDULE [REPEATS]]] + \\ graph-on TP2 v3; rounds/segment=1024, CSV zero-based prompt schedule=0, repeats=1 + \\ sequential shared-pool workload, independent sequence/KV; not concurrent requests + \\ e.g. trace GOLDEN PREFIX 128 0,1,2 2; rejects schedules exceeding 1 MiB/pool + \\ flashnext-exl3-tpe 0 1 0 - policy FILES... compare lifetime/LRU/frequency/bypass on expert demand + \\ flashnext-exl3-tpe RANK WORLD PORT DIR tf|gen GOLDEN DUMP | bench | hit FILES | graph the Engine at tp=1/2 with per-rank pools + \\ flashnext-exl3-graph plan | RANK WORLD PORT DIR eq GOLDEN [STEPS] [MODES] the pool planned on the device and decode graphs vs the host cut \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ info driver, device and library versions \\ smoke copies, fills, launches, argument packing, module globals @@ -185,6 +208,14 @@ const usage = \\ dense-rows bf16 rows times weight rows (routers, heads) against float64 references (synthetic) \\ exl3 [dir] EXL3 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) + \\ exl3-pack DIR [REPS] EXL3 Dense and grouped experts on a real EXL3 pack's shapes: float64 error and timing + \\ exl3-ple2 DIR [ROW] [N] EXL3 n-gram table through cuda_ple.Ngram (ids, rows, RSS) and the PLE kernels vs float64 + \\ exl3-draft DIR [THREADS] EXL3 draft head rows dequantised from lm_head vs the full head (Dense) and float64 + \\ exl3-prefill DIR experts|copy [LAYER]|dense EXL3 prefill pieces on the pack: grouped experts by rows, expert-major passes, line copies, prompt Dense + \\ exl3-tpprefill RANK WORLD PORT DIR IDS NS ROWS CAP deep|both DUMP EXL3 deep prefill at tp=1/2 vs the batched pass and tp=1 + \\ exl3-pool DIR [equal|graph|bench|hit FILES] [LAYERS] Flash Next EXL3 expert offload pool on Experts' pointer tables vs fully resident + \\ exl3-kv kernel|bench [CAP]|real DIR ... int8 KV attention kernels vs upstream and float64; KV placement timings; real-pass KV quality + \\ exl3-kvpage DIR trace|... dynamic KV paging: selection traces, the paged cache vs all-on-device \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) @@ -252,6 +283,11 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-exl3")) return flashnext_exl3_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-tp")) return flashnext_exl3_tp_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-engine")) return flashnext_exl3_engine_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-tpe")) return flashnext_exl3_tp_engine_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-graph")) return flashnext_exl3_graph_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); if (std.mem.eql(u8, cmd, "shared-context")) return runtime_tests.sharedContext(gpu); @@ -300,6 +336,15 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-pack")) return exl3_pack_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-ple2")) return exl3_ple2_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-draft")) return exl3_draft_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-pool")) return exl3_pool_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-kv")) return exl3_kv_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-kvpage")) return exl3_kvpage_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-prefill")) return exl3_prefill_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-tpprefill")) return exl3_tpprefill_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "nccl-lat")) return exl3_pack_tests.ncclLatency(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); diff --git a/zig/tests/exl3_prefill_body_compile.zig b/zig/tests/exl3_prefill_body_compile.zig new file mode 100644 index 0000000000..7302f7f285 --- /dev/null +++ b/zig/tests/exl3_prefill_body_compile.zig @@ -0,0 +1,5 @@ +const flash = @import("flashnext"); + +export fn exl3PrefillForwardBody(engine: *flash.prompt.Engine, sequence: *flash.prompt.Seq, launch: *const flash.Tri, ids: [*]const u32, count: usize) void { + engine.forward(flash.Tri, launch.*, sequence, ids[0..count], .{ .logits = .last }) catch {}; +} diff --git a/zig/tests/fixtures/exl3_ref_compare.json b/zig/tests/fixtures/exl3_ref_compare.json new file mode 100644 index 0000000000..4b99ba0651 --- /dev/null +++ b/zig/tests/fixtures/exl3_ref_compare.json @@ -0,0 +1,1907 @@ +{ + "gen": { + "ids16": { + "ref": [ + 561, + 6511, + 314, + 6124, + 369, + 25358, + 13, + 1049, + 369, + 3750, + 364, + 1141, + 6278, + 84713, + 13, + 561, + 6511, + 314, + 8071, + 369, + 66463, + 13, + 1049, + 557 + ], + "ours": [ + 561, + 6511, + 314, + 6124, + 369, + 25358, + 13, + 1049, + 369, + 3750, + 364, + 1141, + 5269, + 13, + 561, + 6511, + 314, + 14898, + 369, + 21047, + 13, + 1049, + 369, + 3750 + ], + "agree_prefix": 12, + "same_positions": 12, + "n": 24, + "first_divergence": 12, + "ref_top5_at_divergence": [ + [ + 6278, + 17.375 + ], + [ + 86088, + 17.125 + ], + [ + 5269, + 16.75 + ], + [ + 45092, + 16.625 + ], + [ + 81410, + 16.375 + ] + ], + "ref_logit_of_our_pick": 16.75, + "ref_margin_top1_minus_our_pick": 0.625 + }, + "The capital of France is": { + "ref": [ + 11751, + 13, + 561, + 6511, + 314, + 9564, + 369, + 19241, + 13, + 561, + 6511, + 314, + 14898, + 369, + 21047, + 13, + 561, + 6511, + 314, + 17163, + 369, + 23327, + 13, + 561 + ], + "ours": [ + 11751, + 13, + 561, + 6511, + 314, + 9564, + 369, + 19241, + 13, + 561, + 6511, + 314, + 14898, + 369, + 21047, + 13, + 561, + 6511, + 314, + 17163, + 369, + 23327, + 13, + 561 + ], + "agree_prefix": 24, + "same_positions": 24, + "n": 24, + "first_divergence": null + }, + "Once upon a time, in a small village by the sea": { + "ref": [ + 11, + 1017, + 11815, + 264, + 3777, + 7914, + 6725, + 8105, + 13, + 8105, + 557, + 264, + 21520, + 321, + 2993, + 66949, + 7914, + 852, + 9939, + 310, + 12804, + 279, + 1814, + 2094 + ], + "ours": [ + 11, + 1017, + 11815, + 264, + 3777, + 7914, + 6725, + 8105, + 13, + 8105, + 557, + 264, + 21520, + 321, + 2993, + 66949, + 7914, + 852, + 9939, + 310, + 12804, + 279, + 1814, + 2094 + ], + "agree_prefix": 24, + "same_positions": 24, + "n": 24, + "first_divergence": null + } + }, + "tf": { + "ids16@tf48.log": { + "summary": { + "positions": 16, + "ref_eq_gpu": 16, + "ref_eq_oracle": 16, + "gpu_eq_oracle": 16, + "top5_set_eq": 14, + "max_abs_ref_gpu_top1_logit": 0.35900000000000176, + "mean_abs_ref_gpu_top1_logit": 0.10731250000000025, + "max_abs_ref_oracle_top1_logit": 0.3359999999999985, + "mean_abs_ref_oracle_top1_logit": 0.10831250000000014 + }, + "rows": [ + { + "pos": 0, + "input": 760, + "gpu": 20438, + "gpu_logit": 12.414, + "oracle": 20438, + "oracle_logit": 12.411, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "ref_top1": 20438, + "ref_logit": 12.4375, + "ref_top5": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "ref_top5_logits": [ + 12.4375, + 12.0625, + 11.8125, + 11.3125, + 11.0625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0235000000000003, + "abs_ref_oracle": 0.026500000000000412, + "top5_set_eq": true + }, + { + "pos": 1, + "input": 6511, + "gpu": 314, + "gpu_logit": 16.234, + "oracle": 314, + "oracle_logit": 16.223, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "ref_top1": 314, + "ref_logit": 16.25, + "ref_top5": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "ref_top5_logits": [ + 16.25, + 15.1875, + 14.625, + 13.625, + 13.5625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.015999999999998238, + "abs_ref_oracle": 0.027000000000001023, + "top5_set_eq": true + }, + { + "pos": 2, + "input": 314, + "gpu": 279, + "gpu_logit": 15.016, + "oracle": 279, + "oracle_logit": 15.027, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 279, + 9338, + 8071, + 8100, + 13437 + ], + "ref_top1": 279, + "ref_logit": 15.0625, + "ref_top5": [ + 279, + 9338, + 8071, + 13437, + 6535 + ], + "ref_top5_logits": [ + 15.0625, + 12.75, + 12.25, + 12.0625, + 11.875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.046499999999999986, + "abs_ref_oracle": 0.03550000000000075, + "top5_set_eq": false + }, + { + "pos": 3, + "input": 9338, + "gpu": 369, + "gpu_logit": 17.766, + "oracle": 369, + "oracle_logit": 17.771, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 11, + 682, + 271, + 13 + ], + "ref_top1": 369, + "ref_logit": 17.875, + "ref_top5": [ + 369, + 11, + 682, + 271, + 13 + ], + "ref_top5_logits": [ + 17.875, + 17.5, + 14.9375, + 14.625, + 14.375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.1039999999999992, + "top5_set_eq": true + }, + { + "pos": 4, + "input": 369, + "gpu": 11751, + "gpu_logit": 16.391, + "oracle": 11751, + "oracle_logit": 16.398, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 11751, + 271, + 264, + 6924, + 1259 + ], + "ref_top1": 11751, + "ref_logit": 16.5, + "ref_top5": [ + 11751, + 271, + 6924, + 264, + 1259 + ], + "ref_top5_logits": [ + 16.5, + 14.3125, + 13.5, + 13.5, + 13.0 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.10200000000000031, + "top5_set_eq": true + }, + { + "pos": 5, + "input": 11751, + "gpu": 13, + "gpu_logit": 17.547, + "oracle": 13, + "oracle_logit": 17.572, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 13, + 11, + 271, + 321, + 198 + ], + "ref_top1": 13, + "ref_logit": 17.875, + "ref_top5": [ + 13, + 11, + 321, + 271, + 198 + ], + "ref_top5_logits": [ + 17.875, + 17.375, + 14.3125, + 14.1875, + 13.3125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.3279999999999994, + "abs_ref_oracle": 0.3030000000000008, + "top5_set_eq": true + }, + { + "pos": 6, + "input": 13, + "gpu": 561, + "gpu_logit": 15.695, + "oracle": 561, + "oracle_logit": 15.741, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 561, + 11751, + 1049, + 271, + 1061 + ], + "ref_top1": 561, + "ref_logit": 15.75, + "ref_top5": [ + 561, + 11751, + 1049, + 1061, + 198 + ], + "ref_top5_logits": [ + 15.75, + 15.625, + 13.8125, + 13.125, + 13.0 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.054999999999999716, + "abs_ref_oracle": 0.009000000000000341, + "top5_set_eq": false + }, + { + "pos": 7, + "input": 1049, + "gpu": 369, + "gpu_logit": 19.75, + "oracle": 369, + "oracle_logit": 19.739, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 579, + 682, + 557, + 725 + ], + "ref_top1": 369, + "ref_logit": 19.625, + "ref_top5": [ + 369, + 682, + 579, + 557, + 725 + ], + "ref_top5_logits": [ + 19.625, + 16.375, + 16.375, + 15.8125, + 15.4375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.125, + "abs_ref_oracle": 0.11400000000000077, + "top5_set_eq": true + }, + { + "pos": 8, + "input": 369, + "gpu": 799, + "gpu_logit": 19.484, + "oracle": 799, + "oracle_logit": 19.619, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 799, + 264, + 279, + 7172, + 29099 + ], + "ref_top1": 799, + "ref_logit": 19.375, + "ref_top5": [ + 799, + 279, + 264, + 7172, + 29099 + ], + "ref_top5_logits": [ + 19.375, + 17.0, + 17.0, + 16.75, + 16.625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.24399999999999977, + "top5_set_eq": true + }, + { + "pos": 9, + "input": 3750, + "gpu": 364, + "gpu_logit": 20.516, + "oracle": 364, + "oracle_logit": 20.627, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 364, + 430, + 14801, + 310, + 421 + ], + "ref_top1": 364, + "ref_logit": 20.625, + "ref_top5": [ + 364, + 430, + 14801, + 310, + 421 + ], + "ref_top5_logits": [ + 20.625, + 17.625, + 16.625, + 16.125, + 14.5 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.0019999999999988916, + "top5_set_eq": true + }, + { + "pos": 10, + "input": 364, + "gpu": 1141, + "gpu_logit": 19.828, + "oracle": 1141, + "oracle_logit": 19.862, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 1141, + 279, + 57902, + 25438, + 1599 + ], + "ref_top1": 1141, + "ref_logit": 19.75, + "ref_top5": [ + 1141, + 279, + 57902, + 25438, + 1599 + ], + "ref_top5_logits": [ + 19.75, + 18.25, + 16.5, + 16.25, + 15.6875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0779999999999994, + "abs_ref_oracle": 0.11199999999999832, + "top5_set_eq": true + }, + { + "pos": 11, + "input": 279, + "gpu": 242476, + "gpu_logit": 19.688, + "oracle": 242476, + "oracle_logit": 19.647, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 242476, + 9440, + 25438, + 2614, + 10919 + ], + "ref_top1": 242476, + "ref_logit": 19.75, + "ref_top5": [ + 242476, + 9440, + 25438, + 2614, + 10919 + ], + "ref_top5_logits": [ + 19.75, + 16.375, + 13.9375, + 13.625, + 13.3125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.062000000000001165, + "abs_ref_oracle": 0.10300000000000153, + "top5_set_eq": true + }, + { + "pos": 12, + "input": 242476, + "gpu": 300, + "gpu_logit": 20.766, + "oracle": 300, + "oracle_logit": 20.789, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 300, + 273, + 598, + 75, + 1076 + ], + "ref_top1": 300, + "ref_logit": 21.125, + "ref_top5": [ + 300, + 273, + 598, + 75, + 1076 + ], + "ref_top5_logits": [ + 21.125, + 14.4375, + 11.9375, + 11.875, + 11.4375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.35900000000000176, + "abs_ref_oracle": 0.3359999999999985, + "top5_set_eq": true + }, + { + "pos": 13, + "input": 300, + "gpu": 21262, + "gpu_logit": 20.094, + "oracle": 21262, + "oracle_logit": 20.022, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 21262, + 20618, + 198, + 13, + 271 + ], + "ref_top1": 21262, + "ref_logit": 20.125, + "ref_top5": [ + 21262, + 20618, + 198, + 13, + 271 + ], + "ref_top5_logits": [ + 20.125, + 14.9375, + 13.6875, + 12.875, + 12.625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.030999999999998806, + "abs_ref_oracle": 0.10300000000000153, + "top5_set_eq": true + }, + { + "pos": 14, + "input": 21262, + "gpu": 13, + "gpu_logit": 19.734, + "oracle": 13, + "oracle_logit": 19.794, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 13, + 321, + 11, + 248046, + 198 + ], + "ref_top1": 13, + "ref_logit": 19.75, + "ref_top5": [ + 13, + 321, + 11, + 248046, + 198 + ], + "ref_top5_logits": [ + 19.75, + 19.375, + 19.375, + 14.4375, + 14.1875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.015999999999998238, + "abs_ref_oracle": 0.04400000000000048, + "top5_set_eq": true + }, + { + "pos": 15, + "input": 13, + "gpu": 561, + "gpu_logit": 15.859, + "oracle": 561, + "oracle_logit": 15.932, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 561, + 271, + 3437, + 198, + 11751 + ], + "ref_top1": 561, + "ref_logit": 16.0, + "ref_top5": [ + 561, + 3437, + 271, + 198, + 11751 + ], + "ref_top5_logits": [ + 16.0, + 14.9375, + 14.8125, + 14.125, + 13.5625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.14100000000000001, + "abs_ref_oracle": 0.06799999999999962, + "top5_set_eq": true + } + ] + }, + "ids16@tf48_div.log": { + "summary": { + "positions": 28, + "ref_eq_gpu": 27, + "ref_eq_oracle": 27, + "gpu_eq_oracle": 26, + "top5_set_eq": 20, + "max_abs_ref_gpu_top1_logit": 0.35900000000000176, + "mean_abs_ref_gpu_top1_logit": 0.10070370370370355, + "max_abs_ref_oracle_top1_logit": 0.3359999999999985, + "mean_abs_ref_oracle_top1_logit": 0.11307407407407417 + }, + "rows": [ + { + "pos": 0, + "input": 760, + "gpu": 20438, + "gpu_logit": 12.414, + "oracle": 20438, + "oracle_logit": 12.411, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "ref_top1": 20438, + "ref_logit": 12.4375, + "ref_top5": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "ref_top5_logits": [ + 12.4375, + 12.0625, + 11.8125, + 11.3125, + 11.0625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0235000000000003, + "abs_ref_oracle": 0.026500000000000412, + "top5_set_eq": true + }, + { + "pos": 1, + "input": 6511, + "gpu": 314, + "gpu_logit": 16.234, + "oracle": 314, + "oracle_logit": 16.223, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "ref_top1": 314, + "ref_logit": 16.25, + "ref_top5": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "ref_top5_logits": [ + 16.25, + 15.1875, + 14.625, + 13.625, + 13.5625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.015999999999998238, + "abs_ref_oracle": 0.027000000000001023, + "top5_set_eq": true + }, + { + "pos": 2, + "input": 314, + "gpu": 279, + "gpu_logit": 15.016, + "oracle": 279, + "oracle_logit": 15.027, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 279, + 9338, + 8071, + 8100, + 13437 + ], + "ref_top1": 279, + "ref_logit": 15.0625, + "ref_top5": [ + 279, + 9338, + 8071, + 13437, + 6535 + ], + "ref_top5_logits": [ + 15.0625, + 12.75, + 12.25, + 12.0625, + 11.875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.046499999999999986, + "abs_ref_oracle": 0.03550000000000075, + "top5_set_eq": false + }, + { + "pos": 3, + "input": 9338, + "gpu": 369, + "gpu_logit": 17.766, + "oracle": 369, + "oracle_logit": 17.771, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 11, + 682, + 271, + 13 + ], + "ref_top1": 369, + "ref_logit": 17.875, + "ref_top5": [ + 369, + 11, + 682, + 271, + 13 + ], + "ref_top5_logits": [ + 17.875, + 17.5, + 14.9375, + 14.625, + 14.375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.1039999999999992, + "top5_set_eq": true + }, + { + "pos": 4, + "input": 369, + "gpu": 11751, + "gpu_logit": 16.391, + "oracle": 11751, + "oracle_logit": 16.398, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 11751, + 271, + 264, + 6924, + 1259 + ], + "ref_top1": 11751, + "ref_logit": 16.5, + "ref_top5": [ + 11751, + 271, + 6924, + 264, + 1259 + ], + "ref_top5_logits": [ + 16.5, + 14.3125, + 13.5, + 13.5, + 13.0 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.10200000000000031, + "top5_set_eq": true + }, + { + "pos": 5, + "input": 11751, + "gpu": 13, + "gpu_logit": 17.547, + "oracle": 13, + "oracle_logit": 17.572, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 13, + 11, + 271, + 321, + 198 + ], + "ref_top1": 13, + "ref_logit": 17.875, + "ref_top5": [ + 13, + 11, + 321, + 271, + 198 + ], + "ref_top5_logits": [ + 17.875, + 17.375, + 14.3125, + 14.1875, + 13.3125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.3279999999999994, + "abs_ref_oracle": 0.3030000000000008, + "top5_set_eq": true + }, + { + "pos": 6, + "input": 13, + "gpu": 561, + "gpu_logit": 15.695, + "oracle": 561, + "oracle_logit": 15.741, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 561, + 11751, + 1049, + 271, + 1061 + ], + "ref_top1": 561, + "ref_logit": 15.75, + "ref_top5": [ + 561, + 11751, + 1049, + 1061, + 198 + ], + "ref_top5_logits": [ + 15.75, + 15.625, + 13.8125, + 13.125, + 13.0 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.054999999999999716, + "abs_ref_oracle": 0.009000000000000341, + "top5_set_eq": false + }, + { + "pos": 7, + "input": 1049, + "gpu": 369, + "gpu_logit": 19.75, + "oracle": 369, + "oracle_logit": 19.739, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 579, + 682, + 557, + 725 + ], + "ref_top1": 369, + "ref_logit": 19.625, + "ref_top5": [ + 369, + 682, + 579, + 557, + 725 + ], + "ref_top5_logits": [ + 19.625, + 16.375, + 16.375, + 15.8125, + 15.4375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.125, + "abs_ref_oracle": 0.11400000000000077, + "top5_set_eq": true + }, + { + "pos": 8, + "input": 369, + "gpu": 799, + "gpu_logit": 19.484, + "oracle": 799, + "oracle_logit": 19.619, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 799, + 264, + 279, + 7172, + 29099 + ], + "ref_top1": 799, + "ref_logit": 19.375, + "ref_top5": [ + 799, + 279, + 264, + 7172, + 29099 + ], + "ref_top5_logits": [ + 19.375, + 17.0, + 17.0, + 16.75, + 16.625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.24399999999999977, + "top5_set_eq": true + }, + { + "pos": 9, + "input": 3750, + "gpu": 364, + "gpu_logit": 20.516, + "oracle": 364, + "oracle_logit": 20.627, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 364, + 430, + 14801, + 310, + 421 + ], + "ref_top1": 364, + "ref_logit": 20.625, + "ref_top5": [ + 364, + 430, + 14801, + 310, + 421 + ], + "ref_top5_logits": [ + 20.625, + 17.625, + 16.625, + 16.125, + 14.5 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.10900000000000176, + "abs_ref_oracle": 0.0019999999999988916, + "top5_set_eq": true + }, + { + "pos": 10, + "input": 364, + "gpu": 1141, + "gpu_logit": 19.828, + "oracle": 1141, + "oracle_logit": 19.862, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 1141, + 279, + 57902, + 25438, + 1599 + ], + "ref_top1": 1141, + "ref_logit": 19.75, + "ref_top5": [ + 1141, + 279, + 57902, + 25438, + 1599 + ], + "ref_top5_logits": [ + 19.75, + 18.25, + 16.5, + 16.25, + 15.6875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0779999999999994, + "abs_ref_oracle": 0.11199999999999832, + "top5_set_eq": true + }, + { + "pos": 11, + "input": 279, + "gpu": 242476, + "gpu_logit": 19.688, + "oracle": 242476, + "oracle_logit": 19.647, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 242476, + 9440, + 25438, + 2614, + 10919 + ], + "ref_top1": 242476, + "ref_logit": 19.75, + "ref_top5": [ + 242476, + 9440, + 25438, + 2614, + 10919 + ], + "ref_top5_logits": [ + 19.75, + 16.375, + 13.9375, + 13.625, + 13.3125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.062000000000001165, + "abs_ref_oracle": 0.10300000000000153, + "top5_set_eq": true + }, + { + "pos": 12, + "input": 242476, + "gpu": 300, + "gpu_logit": 20.766, + "oracle": 300, + "oracle_logit": 20.789, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 300, + 273, + 598, + 75, + 1076 + ], + "ref_top1": 300, + "ref_logit": 21.125, + "ref_top5": [ + 300, + 273, + 598, + 75, + 1076 + ], + "ref_top5_logits": [ + 21.125, + 14.4375, + 11.9375, + 11.875, + 11.4375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.35900000000000176, + "abs_ref_oracle": 0.3359999999999985, + "top5_set_eq": true + }, + { + "pos": 13, + "input": 300, + "gpu": 21262, + "gpu_logit": 20.094, + "oracle": 21262, + "oracle_logit": 20.022, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 21262, + 20618, + 198, + 13, + 271 + ], + "ref_top1": 21262, + "ref_logit": 20.125, + "ref_top5": [ + 21262, + 20618, + 198, + 13, + 271 + ], + "ref_top5_logits": [ + 20.125, + 14.9375, + 13.6875, + 12.875, + 12.625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.030999999999998806, + "abs_ref_oracle": 0.10300000000000153, + "top5_set_eq": true + }, + { + "pos": 14, + "input": 21262, + "gpu": 13, + "gpu_logit": 19.734, + "oracle": 13, + "oracle_logit": 19.794, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 13, + 321, + 11, + 248046, + 198 + ], + "ref_top1": 13, + "ref_logit": 19.75, + "ref_top5": [ + 13, + 321, + 11, + 248046, + 198 + ], + "ref_top5_logits": [ + 19.75, + 19.375, + 19.375, + 14.4375, + 14.1875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.015999999999998238, + "abs_ref_oracle": 0.04400000000000048, + "top5_set_eq": true + }, + { + "pos": 15, + "input": 13, + "gpu": 561, + "gpu_logit": 15.859, + "oracle": 561, + "oracle_logit": 15.932, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 561, + 271, + 3437, + 198, + 11751 + ], + "ref_top1": 561, + "ref_logit": 16.0, + "ref_top5": [ + 561, + 3437, + 271, + 198, + 11751 + ], + "ref_top5_logits": [ + 16.0, + 14.9375, + 14.8125, + 14.125, + 13.5625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.14100000000000001, + "abs_ref_oracle": 0.06799999999999962, + "top5_set_eq": true + }, + { + "pos": 16, + "input": 561, + "gpu": 6511, + "gpu_logit": 17.875, + "oracle": 6511, + "oracle_logit": 18.121, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 6511, + 9440, + 242476, + 3177, + 8323 + ], + "ref_top1": 6511, + "ref_logit": 18.0, + "ref_top5": [ + 6511, + 9440, + 242476, + 3177, + 8323 + ], + "ref_top5_logits": [ + 18.0, + 15.8125, + 15.3125, + 15.1875, + 13.75 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.125, + "abs_ref_oracle": 0.12099999999999866, + "top5_set_eq": true + }, + { + "pos": 17, + "input": 6511, + "gpu": 314, + "gpu_logit": 20.734, + "oracle": 314, + "oracle_logit": 20.808, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 314, + 369, + 198, + 695, + 271 + ], + "ref_top1": 314, + "ref_logit": 20.875, + "ref_top5": [ + 314, + 198, + 369, + 695, + 271 + ], + "ref_top5_logits": [ + 20.875, + 14.625, + 14.5, + 14.375, + 12.625 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.14099999999999824, + "abs_ref_oracle": 0.06700000000000017, + "top5_set_eq": true + }, + { + "pos": 18, + "input": 314, + "gpu": 6124, + "gpu_logit": 18.141, + "oracle": 6124, + "oracle_logit": 17.961, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 6124, + 9564, + 14898, + 17163, + 9169 + ], + "ref_top1": 6124, + "ref_logit": 18.125, + "ref_top5": [ + 6124, + 9564, + 14898, + 17163, + 9169 + ], + "ref_top5_logits": [ + 18.125, + 18.0, + 17.625, + 17.0, + 15.25 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.015999999999998238, + "abs_ref_oracle": 0.16400000000000148, + "top5_set_eq": true + }, + { + "pos": 19, + "input": 6124, + "gpu": 369, + "gpu_logit": 21.031, + "oracle": 369, + "oracle_logit": 21.08, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 271, + 25358, + 198, + 11 + ], + "ref_top1": 369, + "ref_logit": 20.875, + "ref_top5": [ + 369, + 198, + 25358, + 695, + 11 + ], + "ref_top5_logits": [ + 20.875, + 14.75, + 14.5, + 14.3125, + 14.125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.1559999999999988, + "abs_ref_oracle": 0.2049999999999983, + "top5_set_eq": false + }, + { + "pos": 20, + "input": 369, + "gpu": 25358, + "gpu_logit": 19.891, + "oracle": 25358, + "oracle_logit": 20.028, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 25358, + 271, + 198, + 695, + 239784 + ], + "ref_top1": 25358, + "ref_logit": 19.75, + "ref_top5": [ + 25358, + 198, + 695, + 271, + 239784 + ], + "ref_top5_logits": [ + 19.75, + 15.375, + 14.5, + 14.375, + 13.8125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.14099999999999824, + "abs_ref_oracle": 0.2779999999999987, + "top5_set_eq": true + }, + { + "pos": 21, + "input": 25358, + "gpu": 13, + "gpu_logit": 20.938, + "oracle": 13, + "oracle_logit": 20.945, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 13, + 11, + 271, + 321, + 198 + ], + "ref_top1": 13, + "ref_logit": 20.875, + "ref_top5": [ + 13, + 11, + 321, + 271, + 198 + ], + "ref_top5_logits": [ + 20.875, + 17.5, + 15.375, + 15.25, + 15.125 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.06299999999999883, + "abs_ref_oracle": 0.07000000000000028, + "top5_set_eq": true + }, + { + "pos": 22, + "input": 13, + "gpu": 1049, + "gpu_logit": 18.75, + "oracle": 1049, + "oracle_logit": 18.776, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 1049, + 271, + 3437, + 25358, + 561 + ], + "ref_top1": 1049, + "ref_logit": 18.75, + "ref_top5": [ + 1049, + 271, + 25358, + 3437, + 198 + ], + "ref_top5_logits": [ + 18.75, + 15.6875, + 15.375, + 15.25, + 15.1875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0, + "abs_ref_oracle": 0.0259999999999998, + "top5_set_eq": false + }, + { + "pos": 23, + "input": 1049, + "gpu": 369, + "gpu_logit": 20.688, + "oracle": 369, + "oracle_logit": 20.832, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 369, + 682, + 198, + 17882, + 271 + ], + "ref_top1": 369, + "ref_logit": 20.625, + "ref_top5": [ + 369, + 682, + 17882, + 198, + 20367 + ], + "ref_top5_logits": [ + 20.625, + 17.75, + 15.3125, + 15.3125, + 14.9375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.06299999999999883, + "abs_ref_oracle": 0.20700000000000074, + "top5_set_eq": false + }, + { + "pos": 24, + "input": 369, + "gpu": 3750, + "gpu_logit": 21.344, + "oracle": 3750, + "oracle_logit": 21.366, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 3750, + 10919, + 264, + 279, + 799 + ], + "ref_top1": 3750, + "ref_logit": 21.5, + "ref_top5": [ + 3750, + 10919, + 264, + 279, + 7172 + ], + "ref_top5_logits": [ + 21.5, + 18.875, + 17.875, + 16.875, + 16.375 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.1559999999999988, + "abs_ref_oracle": 0.13400000000000034, + "top5_set_eq": false + }, + { + "pos": 25, + "input": 3750, + "gpu": 364, + "gpu_logit": 21.156, + "oracle": 364, + "oracle_logit": 21.295, + "match": "yes", + "oracle_gap_of_gpu_pick": 0.0, + "gpu_top5": [ + 364, + 430, + 198, + 695, + 13 + ], + "ref_top1": 364, + "ref_logit": 21.25, + "ref_top5": [ + 364, + 430, + 198, + 695, + 13 + ], + "ref_top5_logits": [ + 21.25, + 15.5, + 14.0625, + 12.375, + 12.25 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": true, + "abs_ref_gpu": 0.0940000000000012, + "abs_ref_oracle": 0.045000000000001705, + "top5_set_eq": true + }, + { + "pos": 26, + "input": 364, + "gpu": 1141, + "gpu_logit": 17.328, + "oracle": 10113, + "oracle_logit": 17.281, + "match": "tie", + "oracle_gap_of_gpu_pick": 0.1151, + "gpu_top5": [ + 1141, + 10113, + 92045, + 279, + 39531 + ], + "ref_top1": 1141, + "ref_logit": 17.375, + "ref_top5": [ + 1141, + 10113, + 92045, + 279, + 1602 + ], + "ref_top5_logits": [ + 17.375, + 17.0, + 16.875, + 16.125, + 15.1875 + ], + "ref_eq_gpu": true, + "ref_eq_oracle": false, + "abs_ref_gpu": 0.0470000000000006, + "abs_ref_oracle": null, + "top5_set_eq": false + }, + { + "pos": 27, + "input": 1141, + "gpu": 5269, + "gpu_logit": 17.313, + "oracle": 6278, + "oracle_logit": 17.372, + "match": "tie", + "oracle_gap_of_gpu_pick": 0.9495, + "gpu_top5": [ + 5269, + 6278, + 86088, + 45092, + 31883 + ], + "ref_top1": 6278, + "ref_logit": 17.375, + "ref_top5": [ + 6278, + 86088, + 5269, + 45092, + 81410 + ], + "ref_top5_logits": [ + 17.375, + 17.125, + 16.75, + 16.625, + 16.375 + ], + "ref_eq_gpu": false, + "ref_eq_oracle": true, + "abs_ref_gpu": null, + "abs_ref_oracle": 0.0030000000000001137, + "top5_set_eq": false + } + ] + } + }, + "what": "tools/exl3_ref_compare.py: golden vs old engine ~/tensorfold-exl3-pr20 (tf-cuda-test exl3-forward gen 48 24 and tf 48 hc, GPU 1 under gpurun, 2026-10-19); tf rows include our f64 oracle" +} \ No newline at end of file diff --git a/zig/tests/fixtures/exl3_ref_golden.json b/zig/tests/fixtures/exl3_ref_golden.json new file mode 100644 index 0000000000..53424421cc --- /dev/null +++ b/zig/tests/fixtures/exl3_ref_golden.json @@ -0,0 +1,2252 @@ +{ + "what": "Independent reference: Python TensorFold CUDA engine (qwen4_exp/cuda) on the EXL3 pack; routed expert trellises in host-mapped pinned memory (only allocation site changed), serial decode path, no MTP, no graphs, bf16 KV, greedy argmax", + "pack": "/data/models/Qwen3.8-Flash-Next-exl3-3.05bpw", + "command": "tools/exl3_ref_golden.py --tf /home/chenco_adm/tf-final/src --pack /data/models/Qwen3.8-Flash-Next-exl3-3.05bpw --out golden.json", + "started_utc": "2026-10-18T22:40:21Z", + "finished_utc": "2026-10-18T22:42:31Z", + "versions": { + "tensorfold": "0.6.5", + "torch": "2.9.1+cu130", + "torch_cuda": "13.0", + "python": "3.12.3", + "gpu": "NVIDIA GeForce RTX 5060 Ti", + "capability": [ + 12, + 0 + ], + "driver": "615.71.09", + "nvcc": "Build cuda_13.4.r13.4/compiler.38855100_0" + }, + "load_seconds": 88.23916721343994, + "device_peak_gib": 5.463582515716553, + "centred_offset": 1.0, + "patched_expert_table_line": "big = _HOST_MAPPED(total)", + "notes": "teacher_forced[p] are the logits after position p (predicting ids[p+1]); greedy_steps[0] is the prefill chunk's last row, later steps serial decode", + "prompts": [ + { + "name": "ids16", + "ids": [ + 760, + 6511, + 314, + 9338, + 369, + 11751, + 13, + 1049, + 369, + 3750, + 364, + 279, + 242476, + 300, + 21262, + 13 + ], + "text": "The capital of France is Paris. It is known for the Eiffel Tower.", + "teacher_forced": [ + { + "pos": 0, + "top5_ids": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "top5_logits": [ + 12.4375, + 12.0625, + 11.8125, + 11.3125, + 11.0625 + ], + "logsumexp": 15.040582656860352, + "margin12": 0.375, + "next_actual": 6511 + }, + { + "pos": 1, + "top5_ids": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "top5_logits": [ + 16.25, + 15.1875, + 14.625, + 13.625, + 13.5625 + ], + "logsumexp": 17.2106876373291, + "margin12": 1.0625, + "next_actual": 314 + }, + { + "pos": 2, + "top5_ids": [ + 279, + 9338, + 8071, + 13437, + 6535 + ], + "top5_logits": [ + 15.0625, + 12.75, + 12.25, + 12.0625, + 11.875 + ], + "logsumexp": 16.259746551513672, + "margin12": 2.3125, + "next_actual": 9338 + }, + { + "pos": 3, + "top5_ids": [ + 369, + 11, + 682, + 271, + 13 + ], + "top5_logits": [ + 17.875, + 17.5, + 14.9375, + 14.625, + 14.375 + ], + "logsumexp": 18.581972122192383, + "margin12": 0.375, + "next_actual": 369 + }, + { + "pos": 4, + "top5_ids": [ + 11751, + 271, + 6924, + 264, + 1259 + ], + "top5_logits": [ + 16.5, + 14.3125, + 13.5, + 13.5, + 13.0 + ], + "logsumexp": 17.08510971069336, + "margin12": 2.1875, + "next_actual": 11751 + }, + { + "pos": 5, + "top5_ids": [ + 13, + 11, + 321, + 271, + 198 + ], + "top5_logits": [ + 17.875, + 17.375, + 14.3125, + 14.1875, + 13.3125 + ], + "logsumexp": 18.42790412902832, + "margin12": 0.5, + "next_actual": 13 + }, + { + "pos": 6, + "top5_ids": [ + 561, + 11751, + 1049, + 1061, + 198 + ], + "top5_logits": [ + 15.75, + 15.625, + 13.8125, + 13.125, + 13.0 + ], + "logsumexp": 16.809350967407227, + "margin12": 0.125, + "next_actual": 1049 + }, + { + "pos": 7, + "top5_ids": [ + 369, + 682, + 579, + 557, + 725 + ], + "top5_logits": [ + 19.625, + 16.375, + 16.375, + 15.8125, + 15.4375 + ], + "logsumexp": 19.795913696289062, + "margin12": 3.25, + "next_actual": 369 + }, + { + "pos": 8, + "top5_ids": [ + 799, + 279, + 264, + 7172, + 29099 + ], + "top5_logits": [ + 19.375, + 17.0, + 17.0, + 16.75, + 16.625 + ], + "logsumexp": 19.84423065185547, + "margin12": 2.375, + "next_actual": 3750 + }, + { + "pos": 9, + "top5_ids": [ + 364, + 430, + 14801, + 310, + 421 + ], + "top5_logits": [ + 20.625, + 17.625, + 16.625, + 16.125, + 14.5 + ], + "logsumexp": 20.72075843811035, + "margin12": 3.0, + "next_actual": 364 + }, + { + "pos": 10, + "top5_ids": [ + 1141, + 279, + 57902, + 25438, + 1599 + ], + "top5_logits": [ + 19.75, + 18.25, + 16.5, + 16.25, + 15.6875 + ], + "logsumexp": 20.07109832763672, + "margin12": 1.5, + "next_actual": 279 + }, + { + "pos": 11, + "top5_ids": [ + 242476, + 9440, + 25438, + 2614, + 10919 + ], + "top5_logits": [ + 19.75, + 16.375, + 13.9375, + 13.625, + 13.3125 + ], + "logsumexp": 19.80089569091797, + "margin12": 3.375, + "next_actual": 242476 + }, + { + "pos": 12, + "top5_ids": [ + 300, + 273, + 598, + 75, + 1076 + ], + "top5_logits": [ + 21.125, + 14.4375, + 11.9375, + 11.875, + 11.4375 + ], + "logsumexp": 21.127355575561523, + "margin12": 6.6875, + "next_actual": 300 + }, + { + "pos": 13, + "top5_ids": [ + 21262, + 20618, + 198, + 13, + 271 + ], + "top5_logits": [ + 20.125, + 14.9375, + 13.6875, + 12.875, + 12.625 + ], + "logsumexp": 20.1370906829834, + "margin12": 5.1875, + "next_actual": 21262 + }, + { + "pos": 14, + "top5_ids": [ + 13, + 321, + 11, + 248046, + 198 + ], + "top5_logits": [ + 19.75, + 19.375, + 19.375, + 14.4375, + 14.1875 + ], + "logsumexp": 20.629154205322266, + "margin12": 0.375, + "next_actual": 13 + }, + { + "pos": 15, + "top5_ids": [ + 561, + 3437, + 271, + 198, + 11751 + ], + "top5_logits": [ + 16.0, + 14.9375, + 14.8125, + 14.125, + 13.5625 + ], + "logsumexp": 16.994821548461914, + "margin12": 1.0625 + } + ], + "teacher_forced_top1": [ + 20438, + 314, + 279, + 369, + 11751, + 13, + 561, + 369, + 799, + 364, + 1141, + 242476, + 300, + 21262, + 13, + 561 + ], + "greedy_ids": [ + 561, + 6511, + 314, + 6124, + 369, + 25358, + 13, + 1049, + 369, + 3750, + 364, + 1141, + 6278, + 84713, + 13, + 561, + 6511, + 314, + 8071, + 369, + 66463, + 13, + 1049, + 557 + ], + "greedy_text": " The capital of Japan is Tokyo. It is known for its modern skyline. The capital of Australia is Canberra. It was", + "greedy_steps": [ + { + "pos": 0, + "top5_ids": [ + 561, + 3437, + 198, + 271, + 11751 + ], + "top5_logits": [ + 16.125, + 14.875, + 14.3125, + 14.1875, + 13.625 + ], + "logsumexp": 16.99021339416504, + "margin12": 1.25 + }, + { + "pos": 1, + "top5_ids": [ + 6511, + 9440, + 242476, + 3177, + 8323 + ], + "top5_logits": [ + 18.0, + 15.8125, + 15.3125, + 15.1875, + 13.75 + ], + "logsumexp": 18.37334442138672, + "margin12": 2.1875 + }, + { + "pos": 2, + "top5_ids": [ + 314, + 198, + 369, + 695, + 271 + ], + "top5_logits": [ + 20.875, + 14.625, + 14.5, + 14.375, + 12.625 + ], + "logsumexp": 20.883697509765625, + "margin12": 6.25 + }, + { + "pos": 3, + "top5_ids": [ + 6124, + 9564, + 14898, + 17163, + 9169 + ], + "top5_logits": [ + 18.125, + 18.0, + 17.625, + 17.0, + 15.25 + ], + "logsumexp": 19.219392776489258, + "margin12": 0.125 + }, + { + "pos": 4, + "top5_ids": [ + 369, + 198, + 25358, + 695, + 11 + ], + "top5_logits": [ + 20.875, + 14.75, + 14.5, + 14.3125, + 14.125 + ], + "logsumexp": 20.886533737182617, + "margin12": 6.125 + }, + { + "pos": 5, + "top5_ids": [ + 25358, + 198, + 695, + 271, + 239784 + ], + "top5_logits": [ + 19.75, + 15.375, + 14.5, + 14.375, + 13.8125 + ], + "logsumexp": 19.78594398498535, + "margin12": 4.375 + }, + { + "pos": 6, + "top5_ids": [ + 13, + 11, + 321, + 271, + 198 + ], + "top5_logits": [ + 20.875, + 17.5, + 15.375, + 15.25, + 15.125 + ], + "logsumexp": 20.929258346557617, + "margin12": 3.375 + }, + { + "pos": 7, + "top5_ids": [ + 1049, + 271, + 25358, + 3437, + 198 + ], + "top5_logits": [ + 18.75, + 15.6875, + 15.375, + 15.25, + 15.1875 + ], + "logsumexp": 18.974870681762695, + "margin12": 3.0625 + }, + { + "pos": 8, + "top5_ids": [ + 369, + 682, + 17882, + 198, + 20367 + ], + "top5_logits": [ + 20.625, + 17.75, + 15.3125, + 15.3125, + 14.9375 + ], + "logsumexp": 20.708698272705078, + "margin12": 2.875 + }, + { + "pos": 9, + "top5_ids": [ + 3750, + 10919, + 264, + 279, + 7172 + ], + "top5_logits": [ + 21.5, + 18.875, + 17.875, + 16.875, + 16.375 + ], + "logsumexp": 21.641921997070312, + "margin12": 2.625 + }, + { + "pos": 10, + "top5_ids": [ + 364, + 430, + 198, + 695, + 13 + ], + "top5_logits": [ + 21.25, + 15.5, + 14.0625, + 12.375, + 12.25 + ], + "logsumexp": 21.255651473999023, + "margin12": 5.75 + }, + { + "pos": 11, + "top5_ids": [ + 1141, + 10113, + 92045, + 279, + 1602 + ], + "top5_logits": [ + 17.375, + 17.0, + 16.875, + 16.125, + 15.1875 + ], + "logsumexp": 18.53413963317871, + "margin12": 0.375 + }, + { + "pos": 12, + "top5_ids": [ + 6278, + 86088, + 5269, + 45092, + 81410 + ], + "top5_logits": [ + 17.375, + 17.125, + 16.75, + 16.625, + 16.375 + ], + "logsumexp": 18.821073532104492, + "margin12": 0.25 + }, + { + "pos": 13, + "top5_ids": [ + 84713, + 81410, + 477, + 5269, + 17120 + ], + "top5_logits": [ + 20.0, + 19.125, + 18.625, + 18.375, + 17.875 + ], + "logsumexp": 20.74966812133789, + "margin12": 0.875 + }, + { + "pos": 14, + "top5_ids": [ + 13, + 321, + 271, + 1149, + 198 + ], + "top5_logits": [ + 20.75, + 18.125, + 15.4375, + 14.9375, + 14.75 + ], + "logsumexp": 20.837188720703125, + "margin12": 2.625 + }, + { + "pos": 15, + "top5_ids": [ + 561, + 271, + 3437, + 198, + 248046 + ], + "top5_logits": [ + 19.25, + 16.75, + 15.0, + 14.8125, + 14.4375 + ], + "logsumexp": 19.385690689086914, + "margin12": 2.5 + }, + { + "pos": 16, + "top5_ids": [ + 6511, + 3177, + 271, + 7526, + 198 + ], + "top5_logits": [ + 21.75, + 13.5, + 13.25, + 13.125, + 12.875 + ], + "logsumexp": 21.75299835205078, + "margin12": 8.25 + }, + { + "pos": 17, + "top5_ids": [ + 314, + 695, + 198, + 3177, + 14898 + ], + "top5_logits": [ + 21.875, + 13.1875, + 13.0, + 12.375, + 12.3125 + ], + "logsumexp": 21.87640380859375, + "margin12": 8.6875 + }, + { + "pos": 18, + "top5_ids": [ + 8071, + 15477, + 14898, + 14437, + 6650 + ], + "top5_logits": [ + 19.5, + 19.125, + 19.0, + 18.875, + 18.75 + ], + "logsumexp": 20.739377975463867, + "margin12": 0.375 + }, + { + "pos": 19, + "top5_ids": [ + 369, + 628, + 11, + 695, + 198 + ], + "top5_logits": [ + 21.25, + 14.3125, + 14.0625, + 13.9375, + 13.625 + ], + "logsumexp": 21.255847930908203, + "margin12": 6.9375 + }, + { + "pos": 20, + "top5_ids": [ + 66463, + 20620, + 32678, + 524, + 271 + ], + "top5_logits": [ + 20.125, + 15.0625, + 13.9375, + 13.875, + 12.9375 + ], + "logsumexp": 20.14453887939453, + "margin12": 5.0625 + }, + { + "pos": 21, + "top5_ids": [ + 13, + 271, + 11, + 198, + 321 + ], + "top5_logits": [ + 21.25, + 16.125, + 15.9375, + 13.9375, + 13.875 + ], + "logsumexp": 21.26676368713379, + "margin12": 5.125 + }, + { + "pos": 22, + "top5_ids": [ + 1049, + 271, + 561, + 66463, + 695 + ], + "top5_logits": [ + 19.75, + 16.25, + 15.4375, + 14.75, + 14.3125 + ], + "logsumexp": 19.827075958251953, + "margin12": 3.5 + }, + { + "pos": 23, + "top5_ids": [ + 557, + 369, + 17882, + 682, + 14533 + ], + "top5_logits": [ + 22.125, + 20.375, + 16.75, + 15.4375, + 15.375 + ], + "logsumexp": 22.295686721801758, + "margin12": 1.75 + } + ], + "prefill_vs_serial_last_top1_logit_absdiff": 0.125, + "seconds": 28.889628410339355 + }, + { + "name": "The capital of France is", + "ids": [ + 760, + 6511, + 314, + 9338, + 369 + ], + "text": "The capital of France is", + "teacher_forced": [ + { + "pos": 0, + "top5_ids": [ + 20438, + 220, + 2614, + 1414, + 3049 + ], + "top5_logits": [ + 12.4375, + 12.0625, + 11.8125, + 11.3125, + 11.0625 + ], + "logsumexp": 15.040582656860352, + "margin12": 0.375, + "next_actual": 6511 + }, + { + "pos": 1, + "top5_ids": [ + 314, + 3177, + 2695, + 5759, + 1954 + ], + "top5_logits": [ + 16.25, + 15.1875, + 14.625, + 13.625, + 13.5625 + ], + "logsumexp": 17.2106876373291, + "margin12": 1.0625, + "next_actual": 314 + }, + { + "pos": 2, + "top5_ids": [ + 279, + 9338, + 8071, + 13437, + 6535 + ], + "top5_logits": [ + 15.0625, + 12.75, + 12.25, + 12.0625, + 11.875 + ], + "logsumexp": 16.259746551513672, + "margin12": 2.3125, + "next_actual": 9338 + }, + { + "pos": 3, + "top5_ids": [ + 369, + 11, + 682, + 271, + 13 + ], + "top5_logits": [ + 17.875, + 17.5, + 14.9375, + 14.625, + 14.375 + ], + "logsumexp": 18.581972122192383, + "margin12": 0.375, + "next_actual": 369 + }, + { + "pos": 4, + "top5_ids": [ + 11751, + 271, + 6924, + 264, + 1259 + ], + "top5_logits": [ + 16.5, + 14.3125, + 13.5, + 13.5, + 13.0 + ], + "logsumexp": 17.08510971069336, + "margin12": 2.1875 + } + ], + "teacher_forced_top1": [ + 20438, + 314, + 279, + 369, + 11751 + ], + "greedy_ids": [ + 11751, + 13, + 561, + 6511, + 314, + 9564, + 369, + 19241, + 13, + 561, + 6511, + 314, + 14898, + 369, + 21047, + 13, + 561, + 6511, + 314, + 17163, + 369, + 23327, + 13, + 561 + ], + "greedy_text": " Paris. The capital of Germany is Berlin. The capital of Italy is Rome. The capital of Spain is Madrid. The", + "greedy_steps": [ + { + "pos": 0, + "top5_ids": [ + 11751, + 271, + 264, + 6924, + 1259 + ], + "top5_logits": [ + 16.5, + 14.1875, + 13.375, + 13.25, + 12.9375 + ], + "logsumexp": 17.03568458557129, + "margin12": 2.3125 + }, + { + "pos": 1, + "top5_ids": [ + 13, + 11, + 321, + 271, + 198 + ], + "top5_logits": [ + 17.625, + 17.375, + 14.375, + 14.1875, + 13.375 + ], + "logsumexp": 18.29311752319336, + "margin12": 0.25 + }, + { + "pos": 2, + "top5_ids": [ + 561, + 11751, + 1049, + 271, + 9338 + ], + "top5_logits": [ + 15.875, + 15.625, + 13.9375, + 13.4375, + 13.1875 + ], + "logsumexp": 16.879920959472656, + "margin12": 0.25 + }, + { + "pos": 3, + "top5_ids": [ + 6511, + 3177, + 7526, + 1379, + 6823 + ], + "top5_logits": [ + 19.875, + 15.0625, + 13.6875, + 13.375, + 13.25 + ], + "logsumexp": 19.910140991210938, + "margin12": 4.8125 + }, + { + "pos": 4, + "top5_ids": [ + 314, + 369, + 682, + 321, + 11 + ], + "top5_logits": [ + 20.875, + 15.6875, + 13.75, + 13.6875, + 13.5 + ], + "logsumexp": 20.88703727722168, + "margin12": 5.1875 + }, + { + "pos": 5, + "top5_ids": [ + 9564, + 17163, + 14898, + 9169, + 6124 + ], + "top5_logits": [ + 16.875, + 15.5625, + 15.5, + 15.25, + 15.0 + ], + "logsumexp": 17.757217407226562, + "margin12": 1.3125 + }, + { + "pos": 6, + "top5_ids": [ + 369, + 19241, + 11, + 557, + 271 + ], + "top5_logits": [ + 20.75, + 14.75, + 14.625, + 14.4375, + 14.1875 + ], + "logsumexp": 20.76364517211914, + "margin12": 6.0 + }, + { + "pos": 7, + "top5_ids": [ + 19241, + 271, + 30, + 1048, + 11751 + ], + "top5_logits": [ + 20.125, + 14.375, + 13.375, + 13.25, + 12.8125 + ], + "logsumexp": 20.141525268554688, + "margin12": 5.75 + }, + { + "pos": 8, + "top5_ids": [ + 13, + 11, + 321, + 271, + 198 + ], + "top5_logits": [ + 21.125, + 15.625, + 14.8125, + 14.5, + 13.0625 + ], + "logsumexp": 21.13551139831543, + "margin12": 5.5 + }, + { + "pos": 9, + "top5_ids": [ + 561, + 3437, + 19592, + 11751, + 15767 + ], + "top5_logits": [ + 19.625, + 15.5, + 14.8125, + 14.5625, + 14.4375 + ], + "logsumexp": 19.698711395263672, + "margin12": 4.125 + }, + { + "pos": 10, + "top5_ids": [ + 6511, + 695, + 1994, + 17828, + 198 + ], + "top5_logits": [ + 21.75, + 13.625, + 13.5, + 13.1875, + 12.8125 + ], + "logsumexp": 21.753225326538086, + "margin12": 8.125 + }, + { + "pos": 11, + "top5_ids": [ + 314, + 695, + 198, + 14898, + 413 + ], + "top5_logits": [ + 21.375, + 14.0, + 13.75, + 12.5625, + 12.375 + ], + "logsumexp": 21.3778133392334, + "margin12": 7.375 + }, + { + "pos": 12, + "top5_ids": [ + 14898, + 6124, + 17163, + 31868, + 279 + ], + "top5_logits": [ + 18.125, + 16.5, + 16.375, + 14.625, + 14.5625 + ], + "logsumexp": 18.635133743286133, + "margin12": 1.625 + }, + { + "pos": 13, + "top5_ids": [ + 369, + 30, + 21047, + 271, + 11 + ], + "top5_logits": [ + 20.5, + 15.1875, + 14.1875, + 14.1875, + 13.8125 + ], + "logsumexp": 20.515628814697266, + "margin12": 5.3125 + }, + { + "pos": 14, + "top5_ids": [ + 21047, + 271, + 30, + 198, + 1092 + ], + "top5_logits": [ + 18.75, + 15.8125, + 15.625, + 14.75, + 14.25 + ], + "logsumexp": 18.914113998413086, + "margin12": 2.9375 + }, + { + "pos": 15, + "top5_ids": [ + 13, + 11, + 271, + 198, + 321 + ], + "top5_logits": [ + 21.0, + 14.875, + 14.8125, + 14.1875, + 14.125 + ], + "logsumexp": 21.010028839111328, + "margin12": 6.125 + }, + { + "pos": 16, + "top5_ids": [ + 561, + 3437, + 15767, + 19592, + 271 + ], + "top5_logits": [ + 19.0, + 16.25, + 15.8125, + 15.5625, + 14.75 + ], + "logsumexp": 19.208843231201172, + "margin12": 2.75 + }, + { + "pos": 17, + "top5_ids": [ + 6511, + 198, + 695, + 89798, + 1994 + ], + "top5_logits": [ + 22.25, + 13.875, + 13.6875, + 13.625, + 13.5625 + ], + "logsumexp": 22.252941131591797, + "margin12": 8.375 + }, + { + "pos": 18, + "top5_ids": [ + 314, + 695, + 369, + 198, + 17163 + ], + "top5_logits": [ + 21.25, + 12.9375, + 12.875, + 12.8125, + 12.5625 + ], + "logsumexp": 21.25246810913086, + "margin12": 8.3125 + }, + { + "pos": 19, + "top5_ids": [ + 17163, + 6124, + 32208, + 279, + 8100 + ], + "top5_logits": [ + 18.625, + 15.875, + 14.5625, + 14.1875, + 14.125 + ], + "logsumexp": 18.790557861328125, + "margin12": 2.75 + }, + { + "pos": 20, + "top5_ids": [ + 369, + 30, + 271, + 23327, + 11 + ], + "top5_logits": [ + 20.875, + 15.875, + 14.8125, + 14.375, + 14.0 + ], + "logsumexp": 20.890579223632812, + "margin12": 5.0 + }, + { + "pos": 21, + "top5_ids": [ + 23327, + 271, + 30, + 198, + 13 + ], + "top5_logits": [ + 19.375, + 15.625, + 15.125, + 13.6875, + 13.4375 + ], + "logsumexp": 19.435102462768555, + "margin12": 3.75 + }, + { + "pos": 22, + "top5_ids": [ + 13, + 271, + 30, + 11, + 198 + ], + "top5_logits": [ + 20.875, + 15.4375, + 14.1875, + 14.1875, + 14.0 + ], + "logsumexp": 20.885868072509766, + "margin12": 5.4375 + }, + { + "pos": 23, + "top5_ids": [ + 561, + 3437, + 15767, + 271, + 19592 + ], + "top5_logits": [ + 18.5, + 16.375, + 15.8125, + 15.6875, + 15.3125 + ], + "logsumexp": 18.88610076904297, + "margin12": 2.125 + } + ], + "prefill_vs_serial_last_top1_logit_absdiff": 0.0, + "seconds": 4.493253707885742 + }, + { + "name": "Once upon a time, in a small village by the sea", + "ids": [ + 12162, + 5028, + 264, + 854, + 11, + 303, + 264, + 2526, + 13721, + 539, + 279, + 9117 + ], + "text": "Once upon a time, in a small village by the sea", + "teacher_forced": [ + { + "pos": 0, + "top5_ids": [ + 5028, + 279, + 488, + 264, + 1495 + ], + "top5_logits": [ + 18.75, + 15.4375, + 14.75, + 14.625, + 13.9375 + ], + "logsumexp": 18.875125885009766, + "margin12": 3.3125, + "next_actual": 5028 + }, + { + "pos": 1, + "top5_ids": [ + 264, + 449, + 854, + 279, + 328 + ], + "top5_logits": [ + 26.125, + 16.25, + 16.0, + 14.75, + 14.4375 + ], + "logsumexp": 26.125137329101562, + "margin12": 9.875, + "next_actual": 264 + }, + { + "pos": 2, + "top5_ids": [ + 854, + 944, + 2404, + 96341, + 7047 + ], + "top5_logits": [ + 22.625, + 12.75, + 12.0625, + 12.0, + 11.5625 + ], + "logsumexp": 22.62544059753418, + "margin12": 9.875, + "next_actual": 854 + }, + { + "pos": 3, + "top5_ids": [ + 11, + 303, + 1017, + 264, + 471 + ], + "top5_logits": [ + 21.75, + 18.25, + 15.4375, + 13.625, + 13.1875 + ], + "logsumexp": 21.783252716064453, + "margin12": 3.5, + "next_actual": 11 + }, + { + "pos": 4, + "top5_ids": [ + 1017, + 303, + 264, + 383, + 23492 + ], + "top5_logits": [ + 20.5, + 20.25, + 15.25, + 14.75, + 14.6875 + ], + "logsumexp": 21.09103012084961, + "margin12": 0.25, + "next_actual": 303 + }, + { + "pos": 5, + "top5_ids": [ + 264, + 279, + 1004, + 220, + 2086 + ], + "top5_logits": [ + 22.25, + 20.75, + 17.25, + 15.75, + 14.625 + ], + "logsumexp": 22.461238861083984, + "margin12": 1.5, + "next_actual": 264 + }, + { + "pos": 6, + "top5_ids": [ + 22960, + 1814, + 4128, + 2944, + 6037 + ], + "top5_logits": [ + 20.375, + 19.125, + 19.0, + 19.0, + 18.625 + ], + "logsumexp": 21.255910873413086, + "margin12": 1.25, + "next_actual": 2526 + }, + { + "pos": 7, + "top5_ids": [ + 6092, + 13721, + 321, + 33101, + 18508 + ], + "top5_logits": [ + 19.375, + 17.125, + 16.375, + 15.4375, + 14.6875 + ], + "logsumexp": 19.610164642333984, + "margin12": 2.25, + "next_actual": 13721 + }, + { + "pos": 8, + "top5_ids": [ + 11, + 85624, + 2512, + 303, + 3043 + ], + "top5_logits": [ + 19.5, + 18.875, + 18.25, + 18.0, + 17.375 + ], + "logsumexp": 20.446653366088867, + "margin12": 0.625, + "next_actual": 539 + }, + { + "pos": 9, + "top5_ids": [ + 279, + 264, + 11342, + 449, + 9117 + ], + "top5_logits": [ + 20.375, + 16.875, + 12.25, + 11.4375, + 10.6875 + ], + "logsumexp": 20.406328201293945, + "margin12": 3.5, + "next_actual": 279 + }, + { + "pos": 10, + "top5_ids": [ + 9117, + 14367, + 21128, + 13255, + 89579 + ], + "top5_logits": [ + 19.625, + 17.125, + 15.0625, + 14.875, + 14.8125 + ], + "logsumexp": 19.758014678955078, + "margin12": 2.5, + "next_actual": 9117 + }, + { + "pos": 11, + "top5_ids": [ + 11, + 303, + 2512, + 11815, + 6725 + ], + "top5_logits": [ + 22.875, + 16.25, + 15.9375, + 15.5625, + 14.8125 + ], + "logsumexp": 22.879135131835938, + "margin12": 6.625 + } + ], + "teacher_forced_top1": [ + 5028, + 264, + 854, + 11, + 1017, + 264, + 22960, + 6092, + 11, + 279, + 9117, + 11 + ], + "greedy_ids": [ + 11, + 1017, + 11815, + 264, + 3777, + 7914, + 6725, + 8105, + 13, + 8105, + 557, + 264, + 21520, + 321, + 2993, + 66949, + 7914, + 852, + 9939, + 310, + 12804, + 279, + 1814, + 2094 + ], + "greedy_text": ", there lived a young boy named Tom. Tom was a curious and kind-hearted boy who loved to explore the world around", + "greedy_steps": [ + { + "pos": 0, + "top5_ids": [ + 11, + 2512, + 303, + 11815, + 6725 + ], + "top5_logits": [ + 22.75, + 16.125, + 16.125, + 15.75, + 14.75 + ], + "logsumexp": 22.7548885345459, + "margin12": 6.625 + }, + { + "pos": 1, + "top5_ids": [ + 1017, + 11815, + 264, + 449, + 6107 + ], + "top5_logits": [ + 21.625, + 19.625, + 16.125, + 14.8125, + 13.4375 + ], + "logsumexp": 21.75886344909668, + "margin12": 2.0 + }, + { + "pos": 2, + "top5_ids": [ + 11815, + 557, + 998, + 995, + 3571 + ], + "top5_logits": [ + 22.875, + 21.375, + 19.125, + 13.25, + 12.75 + ], + "logsumexp": 23.09602165222168, + "margin12": 1.5 + }, + { + "pos": 3, + "top5_ids": [ + 264, + 449, + 1330, + 2250, + 1599 + ], + "top5_logits": [ + 21.5, + 20.5, + 18.5, + 18.375, + 17.625 + ], + "logsumexp": 21.89972496032715, + "margin12": 1.0 + }, + { + "pos": 4, + "top5_ids": [ + 3777, + 1809, + 2545, + 7914, + 21520 + ], + "top5_logits": [ + 19.0, + 18.125, + 17.375, + 17.0, + 16.5 + ], + "logsumexp": 19.82315444946289, + "margin12": 0.875 + }, + { + "pos": 5, + "top5_ids": [ + 7914, + 855, + 7397, + 95704, + 89343 + ], + "top5_logits": [ + 18.25, + 17.75, + 17.0, + 16.125, + 16.0 + ], + "logsumexp": 19.209653854370117, + "margin12": 0.5 + }, + { + "pos": 6, + "top5_ids": [ + 6725, + 2512, + 13, + 321, + 852 + ], + "top5_logits": [ + 21.5, + 16.625, + 14.8125, + 14.0, + 13.8125 + ], + "logsumexp": 21.510845184326172, + "margin12": 4.875 + }, + { + "pos": 7, + "top5_ids": [ + 8105, + 7366, + 9076, + 14163, + 1587 + ], + "top5_logits": [ + 17.5, + 16.0, + 15.6875, + 12.8125, + 12.1875 + ], + "logsumexp": 17.89229393005371, + "margin12": 1.5 + }, + { + "pos": 8, + "top5_ids": [ + 13, + 321, + 852, + 440, + 11 + ], + "top5_logits": [ + 22.75, + 17.875, + 17.75, + 14.0625, + 13.5625 + ], + "logsumexp": 22.764793395996094, + "margin12": 4.875 + }, + { + "pos": 9, + "top5_ids": [ + 8105, + 3648, + 5133, + 1216, + 271 + ], + "top5_logits": [ + 21.5, + 18.5, + 16.875, + 14.5, + 13.6875 + ], + "logsumexp": 21.561349868774414, + "margin12": 3.0 + }, + { + "pos": 10, + "top5_ids": [ + 557, + 995, + 9939, + 579, + 11815 + ], + "top5_logits": [ + 23.75, + 21.25, + 20.625, + 19.625, + 17.875 + ], + "logsumexp": 23.88846206665039, + "margin12": 2.5 + }, + { + "pos": 11, + "top5_ids": [ + 264, + 1066, + 2590, + 1040, + 1546 + ], + "top5_logits": [ + 22.625, + 21.125, + 19.875, + 19.875, + 18.625 + ], + "logsumexp": 22.975915908813477, + "margin12": 1.5 + }, + { + "pos": 12, + "top5_ids": [ + 21520, + 2993, + 1546, + 2503, + 46569 + ], + "top5_logits": [ + 20.625, + 18.875, + 18.5, + 17.75, + 17.5 + ], + "logsumexp": 21.013933181762695, + "margin12": 1.75 + }, + { + "pos": 13, + "top5_ids": [ + 321, + 13136, + 10061, + 7914, + 1623 + ], + "top5_logits": [ + 22.375, + 18.125, + 18.125, + 18.0, + 17.0 + ], + "logsumexp": 22.42906379699707, + "margin12": 4.25 + }, + { + "pos": 14, + "top5_ids": [ + 2993, + 46569, + 64821, + 7539, + 43361 + ], + "top5_logits": [ + 20.5, + 20.0, + 18.375, + 17.625, + 17.0 + ], + "logsumexp": 21.187562942504883, + "margin12": 0.5 + }, + { + "pos": 15, + "top5_ids": [ + 66949, + 17207, + 4593, + 7914, + 40842 + ], + "top5_logits": [ + 23.75, + 16.875, + 15.9375, + 15.625, + 14.9375 + ], + "logsumexp": 23.75210189819336, + "margin12": 6.875 + }, + { + "pos": 16, + "top5_ids": [ + 7914, + 10061, + 13136, + 30501, + 1623 + ], + "top5_logits": [ + 20.75, + 18.875, + 18.375, + 18.0, + 17.875 + ], + "logsumexp": 21.082473754882812, + "margin12": 1.875 + }, + { + "pos": 17, + "top5_ids": [ + 852, + 11, + 13, + 440, + 2590 + ], + "top5_logits": [ + 26.0, + 24.5, + 19.0, + 17.5, + 16.875 + ], + "logsumexp": 26.202503204345703, + "margin12": 1.5 + }, + { + "pos": 18, + "top5_ids": [ + 9939, + 2590, + 557, + 995, + 40300 + ], + "top5_logits": [ + 24.375, + 21.625, + 20.75, + 20.125, + 17.625 + ], + "logsumexp": 24.477510452270508, + "margin12": 2.75 + }, + { + "pos": 19, + "top5_ids": [ + 310, + 780, + 23218, + 9862, + 6618 + ], + "top5_logits": [ + 23.0, + 20.0, + 19.375, + 18.375, + 17.875 + ], + "logsumexp": 23.097925186157227, + "margin12": 3.0 + }, + { + "pos": 20, + "top5_ids": [ + 12804, + 3827, + 8070, + 8577, + 1432 + ], + "top5_logits": [ + 23.875, + 23.0, + 19.0, + 18.75, + 18.375 + ], + "logsumexp": 24.24313735961914, + "margin12": 0.875 + }, + { + "pos": 21, + "top5_ids": [ + 279, + 491, + 321, + 780, + 4156 + ], + "top5_logits": [ + 25.375, + 22.875, + 21.75, + 20.375, + 19.0 + ], + "logsumexp": 25.487751007080078, + "margin12": 2.5 + }, + { + "pos": 22, + "top5_ids": [ + 1814, + 54705, + 37771, + 6037, + 13721 + ], + "top5_logits": [ + 24.625, + 19.125, + 18.5, + 17.875, + 17.5 + ], + "logsumexp": 24.63690757751465, + "margin12": 5.5 + }, + { + "pos": 23, + "top5_ids": [ + 2094, + 321, + 13, + 4732, + 314 + ], + "top5_logits": [ + 27.375, + 23.125, + 18.375, + 16.0, + 16.0 + ], + "logsumexp": 27.38933563232422, + "margin12": 4.25 + } + ], + "prefill_vs_serial_last_top1_logit_absdiff": 0.125, + "seconds": 5.046553134918213 + } + ], + "harness_changes": [ + "exl3.expert_table: 'big' (routed-expert trellis buffer) in cudaHostRegister'd host memory read through UVA (42.32 GiB); no other line changed", + "exl3.PREFILL_ROWS=64 (scratch rows; prompts <= 40 tokens), decode.Engine(capacity=128, max_rows=1, prefill_rows=64, graphs=False), load(mtp=False, draft_vocab=None); capacity admission not run", + "run on par1-cs13 GPU 1 under ~/bin/gpurun 1, extensions JIT-built with TORCH_EXTENSIONS_DIR=~/tensorfold-c-ref/ext; reference source ~/tf-final/src read-only (not a git checkout)", + "logits are the engine's bf16 head output: resolution 0.125 in [16,32), 0.0625 in [8,16)" + ], + "remote_paths": { + "pack": "/data/models/Qwen3.8-Flash-Next-exl3-3.05bpw", + "python": "/data/src/kt-venv/bin/python", + "reference_src": "/home/chenco_adm/tf-final/src" + } +} \ No newline at end of file