mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

merge into: mtfyun/ds4:glm5.2
mtfyun/ds4:main mtfyun/ds4:glm-5.3-flash mtfyun/ds4:ds4f-mxfp4 mtfyun/ds4:laguna-s2.1 mtfyun/ds4:glm5.2 mtfyun/ds4:responses-api
..
pull from: mtfyun/ds4:main
mtfyun/ds4:main mtfyun/ds4:glm-5.3-flash mtfyun/ds4:ds4f-mxfp4 mtfyun/ds4:laguna-s2.1 mtfyun/ds4:glm5.2 mtfyun/ds4:responses-api
741 Commits
glm5.2 .. main
Author SHA1 Message Date
antirez b3f50180ec Read half-aligned Qwen Q2 weight rows correctly on Metal 
The ordinary Q2 row dot treated packed codes as word-aligned. Metal read the wrong codes when the expert row started two bytes into a word. Read the same two words via packed bytes and preserve the accumulation order.

Compile the two production down kernels independently in the comparison test so the intervening fused reduction does not pull in unrelated dependencies. The remote M5 Max passes 19,955,180 comparisons across 1,650 fixtures, including half-aligned and address-table weights.
 2026-11-03 16:55:02 +01:00
antirez 5e3a3449e2 Check both staging widths in the Qwen compact expert dispatch test 
The half and float wrappers now forward a fourth template argument for the K tile. Exercise both widths, including negative controls for a dropped or incorrectly fixed width. The M5 Max passes all 3,200 schedule cases under ASan and UBSan.
 2026-11-03 16:52:44 +01:00
antirez cd1c5507b6 Preserve resident Qwen prefill speed with speculative decoding 
Prepare predictor history by default for SSD streaming. When weights fit in RAM, initialize historical draft rows to zero and keep target verification enabled. Full prompt history remains available with DS4_QWEN4_MTP_PREFILL=on.

With 7,809 prompt tokens and 64 generated tokens on an M5 Max, the resident default took 6.47 s versus 7.53 s with eager history. Generated text matched previous main in both modes. The same policy recovered prefill throughput on the DGX Spark.
 2026-11-03 16:50:59 +01:00
antirez 2cd0ae7064 Keep M6 Qwen tuning separate from DeepSeek and GLM kernel selection 2026-11-03 16:41:21 +01:00
antirez 5c9ed63c16 Keep generated Qwen research dumps out of the release tree 
The original experiments remain available in the contributor history. Keep the production code, regression tests and concise benchmarks in the release checkout.
 2026-11-03 16:38:47 +01:00
antirez 7bba0f7d2c Load Qwen version 2 checkpoints and rebuild missing predictor history on demand 2026-11-03 16:38:47 +01:00
antirez a2798605b7 Use resident Qwen arithmetic when experts are streamed from SSD 
Keep the resident projection, attention and shared-expert reductions at each batch size. The storage-only change now preserves full logits across 29-token prompts, appends and the sparse boundary on the M5 Max. Compare complete resident and streamed vocabularies in the model-backed generation test.
 2026-11-03 16:37:20 +01:00
antirez 0dbfd05e16 Keep Metal-only Qwen predictor kernel tests out of the CUDA binary 2026-11-03 16:32:40 +01:00
antirez afd33efab2 Keep current M5 expert tuning alongside the streamed Metal path 2026-11-03 16:29:53 +01:00
antirez c942a90fdd Restore Qwen predictor scratch views after the fused resident expert path 2026-11-03 16:28:41 +01:00
antirez 4ef9e4d0d0 Merge #1056: Giorgio Oppo’s Qwen3.8 Metal SSD expert streaming 
Integrate the streamed expert cache with the current resident IQ2 kernels, projection dispatch and asynchronous n-gram reads. Preserve the original contributor commits.
 2026-11-03 16:28:10 +01:00
Giorgio Oppo 9eb4a12f25 Merge upstream/main into qwen-kernel-opt 
Integrate antirez/main at 890898cf and resolve the shared Metal Q8 conflicts.
Keep Qwen's plain NR2/NR4 matvec and reduction arithmetic alongside upstream's
BF16, RoPE, HC and row variants. Preserve Qwen dispatch guards and prefill
unpack policy, with BF16 rounding disabled for Qwen and NR4 disabled for the
BF16 entry points. Keep upstream's concurrent-unpack and cluster2 shape guards.

Update the Q8 oracle source extraction for the dedicated plain implementation;
the independent frozen references remain unchanged.

Validated CPU/Metal builds; complete Qwen kernels and SSD expert tests;
resident/SSD arithmetic in four math/buffer modes; memory, MTP/device policies
and invalid-token staging; frontend, Engram and session/TP host tests.
Frozen Q8 oracles pass 3072 reduction cases and 1024 full-GEMV fixtures with
zero non-NaN bit differences.

V4.1 synthetic coverage passes up to its decode-topk dispatch-count assertion
at tests/test_deepseek41_metal.c:1618. The same assertion also fails on an
isolated build of unmodified upstream/main on this M1 Max. CUDA/ROCm, M5/M6
hardware and full-model inference were not tested.
 2026-11-02 21:35:23 +01:00
antirez e4709b44be Share DeepSeek V4 Metal index queries across four keys 
Adapt the query reuse idea from Tarjei Mandt and contributors in #1073 to the 64-head Flash scorer. Each SIMD group retains its original dot-product reduction. Store the rounded contribution from each head before summing them in their original order.

On M5 Max, V4 Flash 0731 Q2 steady decode improved from 39.34 to 40.27 tokens/s at 16K and 35.31 to 37.80 at 66K. All full-vocabulary logits matched on the same 96-token continuation at both lengths. Keep the one-key schedules for short rows, tensor parallelism and strict math.
 2026-11-02 11:36:48 +01:00
antirez 890898cfb4 Apply Metal weight tracking to the model being opened 
Select Mark Shank’s V4.1 read-only weight policy after validating the GGUF family. A fresh process otherwise used the V4 Flash default, and reopening the engine inherited the previous model’s setting.
 2026-11-02 11:26:14 +01:00
antirez 9fe88cd312 Integrate the per-model Metal weight policy fix 2026-11-02 11:26:14 +01:00
antirez a4bbb38ce4 Integrate the V4.1 memory-aware Metal TP fix 2026-11-02 11:17:03 +01:00
antirez c54841002d Keep V4.1 TP sharded when either Mac needs the smaller working set 2026-11-02 11:17:02 +01:00
antirez 8c54d7fb01 Merge #1073 with local Metal and Spark optimizations 2026-11-02 11:12:11 +01:00
antirez 0accf8b395 Merge #1073: speed up DeepSeek V4.1 Flash on Metal 
Integrate Tarjei Mandt (kernelpool) and contributors’ Metal decode, prefill and tensor-parallel work. Keep the existing four-row index scorer for short and non-FP4 inputs, and retain the Spark small-batch projection.
 2026-11-02 11:09:09 +01:00
antirez 6075227f1e Preserve the Spark small-batch shared-expert projection 2026-11-02 11:04:22 +01:00
antirez 224e77e791 Keep Metal-only V4.1 decode helpers out of the Spark build 2026-11-02 11:04:22 +01:00
antirez a88711246c Merge reviewed server PRs #1165 and #1245 from upstream 2026-11-01 18:07:54 +01:00
Salvatore Sanfilippo b4e3faf526 Merge #1245: allow up to 256 images per server request 
Allow up to 256 images per server request
 2026-11-01 18:07:47 +01:00
Salvatore Sanfilippo da2a102dd0 Merge #1165: advertise the loaded model modalities 
Server: advertise loaded model modalities
 2026-11-01 18:06:35 +01:00
antirez b7615c141f Keep the default CUDA unroll policy for multirow Qwen Q8 
Only the single-token Q8 loop needs a four-way unroll. Express the multirow loop without a pragma so NVCC does not warn about and ignore a zero unroll count. On GB10, all 552 executable sections and their constants and shared-memory metadata match the previous build byte for byte.
 2026-11-01 09:46:06 +01:00
antirez 36e2680676 Link ROCm fallback symbols into standalone GPU test tools 
The release QA link on Strix Halo exposed two unresolved V4.1 projection symbols from ds4_rocm.o. Add the fallback object already used by the engine to the four standalone consumers. All four now link on gfx1151 and the MXFP4, GLM attention and KDA tests pass.
 2026-11-01 09:43:04 +01:00
antirez ae40c5c4f0 Carry the complete V4.1 raw window through medium prefill pruning 
Keep all 128 terminal rows in the dependency calculation. A one-row seed was enough for the fresh benchmark but left early keys of the final decoder window different on a CUDA 4096-to-10240 append. The compressed caches matched; differences in the raw window also changed the next-token logits.

The same Spark append now matches the entire 74 MiB session and all five vocabulary vectors at 202.7 vs 154.0 tokens/s. Cover an exactly aligned 4096-token starting frontier, continued prefill, a ragged final chunk and ordinary decode in the regression.
 2026-10-31 20:43:40 +01:00
antirez c1bce729fc Cross the legacy GLM sparse boundary in one SSD layer sweep 
For full single-host compact IQ2 graphs, keep aligned dense and sparse ranges in the same layer sweep. Let the existing attention slices switch regimes and score only the suffix that consumes explicit IDs. Keep a split for partial tiles, two-row verification, TP, imatrix and quality graphs.

Full GLM 5.2/5.3 M5 Max IQ2 checkpoints with a 40 GiB cache: 4096-token 5.2 prefill reached 113 tokens/s with identical saved state and logits (the pre-change runs were 83–87). A full 5.3 1792-to-2304 append improved from 13.0 to 25.6 tokens/s, also byte-exact.
 2026-10-31 20:42:03 +01:00
antirez ef8e536501 Reuse Q2_K bit planes in Qwen Metal expert prefills 
Extend the Q4_K load-reuse approach from Matthijs van den Bos (#1198) to Q2_K. Four K steps select planes from the same 16 bytes, and eight share the same scale pair. Keep four-byte alignment and the original dequantization.

The M5 Max IQ2/Q2 prefill expert benchmark averaged 7114 to 7032 microseconds across four runs per binary. The Metal specialization fixture passed and an 8192-token Q2 prompt reproduced the complete saved session and continuation logits.
 2026-10-31 20:35:44 +01:00
antirez 524c874506 Skip wired GLM weights during full-layer SSD read-ahead 
Record successful locking of the complete static weight set. In whole-layer reads for GLM, prefetch only the three routed expert tensors when all static tensors are wired. Dense layers then need no worker. Preserve normal prefetch after partial or failed mlock.

M5 Max full GLM 5.2, 1024-token SSD prefill: the saved session and five full logit vectors were identical in four runs with a 40 GiB cache. This removes redundant pread traffic, including calls for weights that were already resident.
 2026-10-31 20:32:49 +01:00
antirez 46ade1bb3e Reuse rotated compact keys across legacy GLM prefill queries 
Prepare the live FP16 key cache once per sparse attention slice with the original RoPE helper, then share the FP32 results across queries and heads. Bound scratch at 16 MiB, retain the original path for small or unsupported launches, and fall back on allocation failure. Keep the serialized cache format unchanged.

On M5 Max the full GLM 5.3 IQ2 4096-to-6144 SSD append averaged 80.12 to 83.21 tokens/s. GLM 5.2 also preserved all frontier and continuation logits. Add a direct rollback comparison that catches differences from folding the rotation constants.
 2026-10-31 20:31:53 +01:00
antirez 1c6d058d9f Gather Qwen BF16 prompt embeddings directly on the GPU 
Reuse the idle token-routing and block-output tensors for batched BF16 gather and HC expansion. Send the token IDs in place of a host-expanded FP32 residual, preserving position handling and validation. Keep short inputs, other embedding formats, image spans and ROCm on their original paths.

Four-run 4096-token Q2/Q4 comparisons on M5 Max and GB10 matched the saved state and the prefill and four continuation logit vectors. Q2 on GB10 measured 725–734 tokens/s vs 629–680 before.
 2026-10-31 20:30:32 +01:00
antirez d9fc9cf30c Merge #1147: validate sparse selection prefixes in Metal GLM prefill 
Credit to Emilian Bold. Full eight-head groups can use the valid-row specialization up to the first padding ID when every later slot is invalid. Interior holes keep the guarded path.

On M5 Max the first four exact 8192-to-10240 GLM 5.3 Flash Q2 append runs measured 410–419 tokens/s with the prefix path and 382–385 with it disabled. The attention fixtures pass under Metal API validation.
 2026-10-31 20:30:04 +01:00
antirezandMarcello Costagliola b799c90e3a Resize independent image rows in parallel before vision prefill 
Keep every bicubic tap, its accumulation order, canvas stride and rounding. Partition large resizes across at most eight workers and finish unassigned rows synchronously if thread creation fails. Small images stay serial.

The flat padded-canvas regression comes from Marcello Costagliola (#1008). The tap-table implementation changed pixels on GB10. The row scheduler matches all 78 differential cases on M5 Max and DGX Spark, including adversarial rounding ties. A 4032x3024 to 1024x768 resize takes 29.7 vs 190.5 ms on M5 Max and 54.2 vs 227.7 ms on GB10.

Co-authored-by: Marcello Costagliola <marcello.costagliola1@gmail.com>
 2026-10-31 20:28:10 +01:00
antirez 773247a881 Skip whole decoder chunks in medium V4.1 prefills 
Adapt the aligned dependency suffix and 128-row ring warmup identified by HDYS-TTBYS in #1272 to 4096–8191-row single-host text sweeps. Preserve matrix partitions and keep TP, DSpark, multimodal, ROCm and the existing long-sweep policy on their prior routes.

At 6144 tokens the 1017 IQ2 SSD run matched the saved state and nine full logit vectors on M5 Max and GB10. Spark improved from 154.6 to 198.7 tokens/s. The default-kernel regression covers a ragged 4353-row prefill and a nonzero-start 6144-row append.
 2026-10-31 20:21:25 +01:00
antirez 8c9b1986a4 Merge #1127: overlap short callback-driven Metal Flash prefills 
Preserve emilianbold’s two-layer submission bound and report progress only after the queued layers finish. SSD streaming, tensor parallelism, visual input, profiling and large chunks keep their existing scheduling.

On M5 Max with Flash-0731 Q2, 12 warmed 256-token prefills per schedule averaged 296.09 to 300.08 tokens/s; all 3,102,720 checked logits were identical. Four 511-token fresh runs also matched the serialized cache, all logits after prefill and eight continuation steps, and all 44 progress notifications.
 2026-10-31 20:15:45 +01:00
antirez c2779bb4e8 Keep the four-head Flash indexer schedule on Metal TP 
The one-barrier schedule improves single-host decode but slowed an 8192-token Flash-0731 RDMA decode on two M5 Max machines. Reverse-order repeats gave 42.57/42.49 tok/s on the new schedule versus 43.56/42.95 on the previous schedule. Select a separate pipeline at runtime and preserve the four-head schedule for the pair.

After the guard, the final and baseline builds averaged 42.98 and 43.07 tok/s with byte-identical FP32 logits throughout all four runs. Single-host raw scores and guards still match at all five widths.
 2026-10-31 17:49:26 +01:00
antirez 32604bc6de Issue Qwen Q8 weights ahead during single-token CUDA decode 
Unroll the stock four-quant lane loop by four only for the one-row kernel. Inspired by matrixfede’s load-scheduling work in #1207. Keep the existing dot-product order and GGUF weights.

DGX Spark, a real 2048-token prefix and 96 teacher-forced steps, eight interleaved snapshot replays: warm median 21.95 to 22.20 tok/s on Q4 and 23.36 to 23.63 on Q2. All 248320 logits were identical after every token. The Qwen CUDA kernel suite and a rebuilt one-row-specialized binary passed.
 2026-10-31 17:46:11 +01:00
antirez dc9580daf4 Merge #1213: stop discarded GLM MTP drafts after their KV write 
Integrate Flor1an-B’s catch-up optimization and use a null draft destination to share normal completion and cleanup. The accepted pair still publishes its first next-token KV entry before the real draft consumes it.

On a 128 GiB M5 Max with a 40 GiB SSD expert cache, four runs of the full GLM 5.3 IQ2 checkpoint averaged 4.802 to 4.848 tok/s (+1.0%). All 80 output tokens and all 23 multi-token cycles matched. The same workload on GLM 5.2 matched its output and acceptance counts, with timing within noise. The final build also passed the streamed GLM 5.2 boundary case and GPU session-state suite.
 2026-10-31 17:38:36 +01:00
antirez babe24db9f Fill single-token GLM indexer tensor tiles with heads on Spark 
In the one-token path, put 16 of the 32 query heads in the WMMA tile that normally spans 16 tokens. Keep the FP16 conversion, eight ordered MMA steps and serial FP32 head-weight accumulation used by the batched kernel. Leave causal, quality and other GPU paths on their existing dispatch.

Sixteen FP16/FP32 cache tests compare every score to the original batched kernel, including ties, ragged key counts and an output guard. CUDA indexer regressions and Compute Sanitizer passed. Full GLM 5.3 Flash Q2 on one Spark, real 16384-token prefix and 96 decoded tokens: steady 14.94 to 15.06 tok/s (+0.8%) with bit-identical full logits in four ABBA runs.
 2026-10-31 17:33:15 +01:00
antirez db67455bed Reduce Metal Flash and Pro index scores with one barrier 
Each SIMD group retains its key and writes independent head contributions. Sum them in the original head order. On M5 Max the raw score arrays were byte-identical at five widths, including tied keys and a guarded tail. At 32768 prompt tokens on Flash-0731 Q2, ABBA runs averaged 38.81 to 39.38 steady tok/s (+1.5%); all logits matched across the 96 decode steps.
 2026-10-31 17:33:15 +01:00
antirezandtrueimage efd018cc43 Tile single-host GLM 5.3 Flash indexed attention on Metal 
Adapt trueimage’s four-phase attention kernels from f4023085 in #1090. Reuse the existing split-attention buffers for scores and normalization, and reuse the consumed low-rank query for value accumulation. Retain the established paths for TP, RoPE and unsupported cache and weight shapes.

M5 Max, full Q2 model, 2048-token prefix and 80 decode steps: four fresh runs averaged 29.91 to 33.43 steady tok/s (+11.8%). Raw logits matched at each step. The ragged, repeated-index GPU oracle also matches the generic kernel with the production scratch aliases under Metal validation.

Co-authored-by: trueimage <11846060+trueimage@users.noreply.github.com>
 2026-10-31 17:32:53 +01:00
antirez b00bbf42da Remove per-head barriers in the CUDA Flash and Pro decode indexer 
Each warp writes all its head contributions before the one block barrier. The final sum retains the original groups of four. On GB10, 32769 keys took 0.301 ms versus 0.409 ms with identical raw score bytes, including tied keys and an output guard. On Flash-0731 Q2, four fresh 32768-prompt runs averaged 18.47 to 18.65 tok/s; every logit matched at all 96 decode steps.
 2026-10-31 17:29:39 +01:00
antirez 79e59d7ac2 Merge #1270: keep legacy GLM MTP live across the sparse boundary 
Credit Luo Hanjun (Elvis-Luo). On an M5 Max with the real GLM 5.2 IQ2 checkpoint and a 40 GiB SSD expert cache, a 2046-token prompt reproduces the pre-merge verifier failure at position 2047. After the fix the next 32 greedy tokens match ordinary decoding exactly, including nine multi-token cycles. The GPU session-state tests pass.
 2026-10-31 17:26:35 +01:00
antirez 95976d737f Merge #1238: speed up Qwen3.8 Metal decode 
Integrate paradise-runner’s expert, attention, hyper-connection, and PLE scheduling work. Preserve main’s validation and chunk prefetch history. Limit the new PLE overlap to single-row Metal decode. Omit the process-lifetime router copy and speculative model-range prefetch.

On a 128 GiB M5 Max, four fresh runs with a 2048-token prompt and 80 decode steps measured 52.54 to 59.68 tok/s on Q2 (+13.6%) and 54.18 to 56.77 on Q4 (+4.8%). Every FP32 logit matched the pre-merge binary, including greedy Q4. The PLE prefetch and error-recovery fixture passed.
 2026-10-31 17:16:16 +01:00
antirez d61aa0c139 Reuse V4.1 index queries across eight keys per Spark warp 
Keep eight independent keys in each warp's registers and reuse the
query loads across them. The per-key FP32 operations, reduction tree,
mask sentinel and original score positions stay the same. Below 8192
keys the single-key warp supplies more parallelism. Multi-row and
partially visible inputs also use the original shape.

The Metal four-row scorer from HDYS-TTBYS (#1267) inspired the grouping.
On GB10, the isolated dense scorer goes from 0.37 to 0.24 ms at 65K
and 5.7 to 3.3 ms at 1M keys. On two physical Sparks over RDMA with a
real 65536-token prompt and full Q2 weights, steady ABBA runs measured
19.75/20.22/20.30/20.14 t/s (+1.6%). All 129280 logits match for all
64 steps in all four runs. One Spark at a 1048512-token synthetic cache
frontier with full streamed Flash-1017 weights went from 4.10 to 4.26
t/s (+3.9%, warm medians), with all logits identical for 32 steps.

Cover a ragged 8197-key row against the original shared reduction,
alternately retaining and rejecting the tail and checking an output
guard. The CUDA and ROCm V4.1 kernel suites pass, including CUDA memcheck.
 2026-10-31 16:27:39 +01:00
antirezandbadelon e3137ce6bc Batch disk Engram reads for Metal verifier rows 
Adapt the Engram read batching in #1240 to consecutive rows of one session. Read directly into the existing staging allocation with strided row IDs. Independent sessions and mixed batches keep their per-session reads.

M5 Max, streamed V4.1 Q2 with DSpark, 2048-token prompt and 192 generated tokens: four fresh ABBA runs averaged 14.07 to 14.28 steady tokens/s. Generated text matches, and all accepted-prefix, rewind, exact sampling and snapshot checks pass with Metal validation enabled.

Co-authored-by: badelon <211747211+badelon@users.noreply.github.com>
 2026-10-31 15:54:28 +01:00
antirez 23c333bd1c Skip rejected V4.1 index keys in Metal SSD decode 
A streaming single-token sweep owns the candidate blocks for exactly the same lifetime as a resident sweep. Enable the GPU-only scoring from #1262 here too.

On an M5 Max with Flash-1017 Q2, a 40 GiB expert budget and 131072 prompt tokens, the warm median across three paired 96-step replays goes from 16.65 to 16.99 tokens/s. All 129280 logits match bit for bit at every step.
 2026-10-31 15:52:26 +01:00
antirez c7bb2d7656 Skip masked V4.1 index scores on DGX Spark 
The layer-20 shortlist contains at most 2048 eight-key blocks. On later
reindex layers, one CUDA CTA owns exactly one of those blocks. Write the
existing rejection sentinel before loading its keys when the mask excludes
it. Admitted keys use the same warp reduction and original score slots.

On a GB10, Flash-1017 Q2 with a 40 GiB SSD expert budget, a real 65536-token
prefix and 64 teacher-forced decode steps, the warm control runs measured
12.65/12.64/12.64 t/s and the corresponding optimized runs 12.80/12.83/
12.85 t/s. All vocabulary logits matched at every step. The CUDA primitive
compares scores with a separate full score/filter and guards the final
partial block; Compute Sanitizer reported zero errors.

The corresponding Metal candidate-selection idea came from HDYS-TTBYS
in #1262. This CUDA variant consumes the existing GPU mask directly.
 2026-10-31 15:49:39 +01:00
antirez 4c29e82b3e Merge #1034: queue single-host V4.1 Metal SSD decode by default 
Bring in Dango233's bounded streaming schedule and forced-eviction
oracle. Expert-ID reads and slot replacement already order the stream.
The separate Engram inputs from #1261 remove the mid-token overwrite.
Retain the old waits for quality/imatrix and as an A/B control.

M5 Max 128 GiB, Flash-1017 Q2, 40 GiB SSD expert budget, 2048 prompt
tokens and 128 greedy output tokens, fresh processes in ABBA order:
steady control 14.65/14.51, queued 15.77/15.47 t/s (+7.1%).
The output matched byte for byte. In same-session teacher-forced replay,
all 129280 logits matched after each of 128 tokens. Metal API validation
also passed the forced-eviction oracle, including logits, history and
all live cache spans, and the public snapshot/cancellation fixture.
 2026-10-31 15:47:40 +01:00
antirez 47400ed008 Merge #1262: skip rejected V4.1 keys without changing Metal top-k ties 
Keep HDYS-TTBYS's prefix merge and alternating partial buffers, resolved
against the four-row scorer from #1267. A repeated-key oracle exposed a
difference in the PR's compact sort: equal scores produced different
original row IDs (3224 vs 3612 at rank 16). Ascending candidate indices
do not make the original bitonic sort stable.

Score layer 20's selected blocks directly on the GPU, writing scores at
their original cache indices. Reuse the usual mask and full-width top-k
so equal scores encounter the same comparison network as before. This
also removes the CPU readback, qsort, upload and remapping buffers.

M5 Max: raw scores and all ordered top-512 IDs match on a tied, ragged
32773-key fixture. The bounded resident graph matched every logit for
40 consecutive positions at 16K, 64K and 131K. Warm 131K throughput
increased from 27.39 to 28.74 t/s. The full resident checkpoint exceeds
this host's memory; the timing fixture shares one routed bank.
 2026-10-31 15:42:39 +01:00
antirez 32f4243a09 Merge #1261: pipeline resident V4.1 Metal decode submissions 
Preserve HDYS-TTBYS's four scheduling and Engram buffer commits. On an
M5 Max, a bounded resident 40-layer fixture (one shared routed-expert
bank) gave 23.04–23.08 t/s drained and 30.61–30.64 t/s queued after
warmup. All 129280 logits matched at each of 40 positions. Two- and
six-row fixtures also matched the pre-PR binary byte for byte. This is
a scheduling measurement, not throughput for the full resident model.
The real Q2 SSD session-state test and Metal kernel suite passed.
 2026-10-31 15:37:19 +01:00
Giorgio Oppo f2e60304a3 Merge upstream/main into qwen-kernel-opt 
Resolve conflicts while preserving Qwen SSD/MTP optimizations and the new
upstream PLE prefetch path. Validate tokens and repair the predictor boundary
before either input-staging path, and preserve projection phases for prefill,
decode and verification. Keep both sets of build targets and kernel tests,
adapting the raw16 fixture to the branch's M5/M6 device helper.

Validated Metal and CPU builds, Qwen kernels, SSD cache/read/fallback tests,
resident/SSD arithmetic in four math/buffer modes, memory/MTP/device policies,
PLE prefetch/cancel/fallback/recovery with the local synthetic fixture,
Metal buffer lifetimes, V4.1 indexer parity, and model-independent host tests.
The complete make test suite requires the unavailable ds4flash.gguf model;
CUDA/ROCm and M5/M6 hardware were not tested.
 2026-10-31 09:39:36 +01:00
antirez e14d7f8b1e Release cached engines between regression tests 
Close each test's cached engines before the next entry opens an engine with its own options. This prevents make test and combinations such as session-snapshot plus logprob-vectors from colliding with the runtime's single-instance lock.

Validated the formerly failing combined tests and the full make test suite on M5 Max with the release Flash 0731 GGUF and DSpark support. Legacy MTP passes separately because SSD cache-pressure tests intentionally use an incompatible engine configuration.
 2026-10-30 19:32:07 +01:00
antirez c97964b8aa Validate Qwen tokens before prefetched input staging 
Keep the existing vocabulary check ahead of embedding reads on both staging paths. Reject malformed chunks before changing live PLE history, and test invalid tokens in the first chunk, a prefetched chunk, and the partial tail with exact recovery.
 2026-10-30 18:10:33 +01:00
antirez 21178462db Test indexer row-group tails and output bounds 2026-10-30 17:57:19 +01:00
HDYS-TTBYS 345d7f44cb Score four indexer rows per SIMD group 
kernel_glm_indexer_score_one_direct spends one 128-thread threadgroup per row:
it stages the 512 B key row in threadgroup memory and closes every head group
with two threadgroup barriers. It also re-reads the whole query vector
(32 heads x 512 B = 16 KB) for every row, even though that vector is the same
for all rows of a token.

Give one SIMD group four rows instead. Each lane holds four float4 key values
in registers, and the head loop reads q4[lane] once and reuses it for all four
dots and all four simd_sum calls. No threadgroup memory and no barriers, and a
threadgroup covers 16 rows instead of 4.

On this tree (without candidate-row scoring) 1M index rows go from 7.93 ms to
2.89 ms, 68 -> 186 GB/s, 2.75x on the score pass alone. In a tree that also
carries #1262 the same kernel runs the stream at 304 GB/s instead of 143 and
the 130K-context score pass drops from 1413 to 657 us per token.

The arithmetic per row is untouched: dot(float4), simd_sum, then
max(s*scale,0)*w[head] accumulated on lane 0 in head order. The accumulate
carries #pragma clang fp contract(off) because the library builds with fast
math and the scalar kernel rounds the product through threadgroup memory.
Without it 334,121 of 1M rows differ by one fma rounding, and --dump-logits
still passes: these scores only feed a top-k, so a one-ulp change is invisible
until it flips a selection at the boundary. That is why this adds a kernel-level
memcmp test rather than relying on the logits gate.

tests/test_v41_indexer_score compares both kernels over 1M random rows,
fp4-representable keys (the cache holds fp4-dequantized values) and a
full-range f16 cache, and prints each variant's rate so the number reproduces
without a model. DS4_METAL_DISABLE_V41_SCORE_ROWS4=1 selects the scalar kernel
for A/B.

Prefill parity is exact at every frontier, the session fixture passes, and
--dump-logits plus generated text are byte-identical on and off at 32K.

Scope: the shared host entry point also serves GLM 5.2/5.3 indexer scoring
(same 32 heads x 128 dim shape), so those paths take the new kernel as well.
Bit-identity is established at kernel level including f16 caches; no GLM model
was available to run end to end.
 2026-10-30 17:56:53 +01:00
Matthijs van den Bos e73e516e75 metal: reuse Q4_K raw loads in Qwen NAX prefill 2026-10-30 17:56:53 +01:00
Matthijs van den Bos 64baaa0678 Test Qwen PLE prefetch parity and recovery 
Exercise 256-token prefill chunks and a partial tail with the real model,
checking exact logits at each frontier and after decode against rollback.
Verify cancellation/resume, background read failure, and unavailable worker
fallback with test-local injection.

Validated on Apple Metal with Qwen3.8-Flash-Next-Q4.gguf:
make -j4 tests/test_qwen4_ple_prefetch test-qwen4-ngrams
./tests/test_qwen4_ple_prefetch gguf/Qwen3.8-Flash-Next-Q4.gguf

CUDA not run.
 2026-10-30 17:56:10 +01:00
Matthijs van den Bos 347c2e1fbe Overlap Qwen PLE reads with prefill GPU work 2026-10-30 17:56:10 +01:00
Emilian Bold 82aa7e8e60 metal: use shuffle network for non-causal top-k block sort 
The non-causal kernel_argsort_f32_i32_desc used the threadgroup-memory
bitonic network; the simd_shuffle_xor variant added for DSv4.1 was only
wired to the causal path. Instantiate it for the plain kernel too.
Selected indices are bit-identical (154 shape/tie/-inf/odd-size cases).

The 1024-wide block sort is ~85% of top-k time, so whole
ds4_gpu_indexer_topk_tensor calls (M5 Max, 256 rows, wall clock incl.
~250us fixed commit/wait), main vs this change:

  8192/512     973 -> 754 us  (-22%)
  16384/1024  2024 -> 1565 us (-23%)
  32768/2048  4275 -> 3424 us (-20%)
  32768/512   3646 -> 2793 us (-23%)
 2026-10-30 17:56:04 +01:00
nomish 5382eac5e0 metal: process medium V4.1 SSD appends in one layer sweep 
e9e1baa made CUDA process medium SSD appends (2048 < N < 8192 with a
final tile of at least 256 rows) in one layer sweep that keeps the
2048-row arithmetic partitions. Enable the same rule on Metal, where
these appends still streamed the experts twice. Below 8192 rows a Metal
sweep already runs in 2048-row encoder chunks and skips the decoder
suffix, so the partitions are unchanged.

DS4_METAL_DISABLE_V41_MEDIUM_SWEEP turns it off. The dispatch test
expects the one-sweep count on Apple, and the Metal model check runs its
control session with the switch set so state_equal covers both paths.
 2026-10-30 17:55:43 +01:00
Tim Lawrenz a8ce74afd5 server: render blank reasoning as absent 
A must-echo client (DeepSeek-family) replays an assistant turn that produced
no reasoning as a single space, because the hosted API rejects an empty
string. The renderers echoed that body verbatim, so the prompt carried a token
between the think markers that the model never sampled.

KV reuse is a prefix match, so one extra token invalidates every turn after
it. Measured on a live 36k conversation: the token tier reports
first_mismatch at the space, and the text tier — which normally rescues a
token mismatch by comparing rendered text instead of ids — fails as well,
because the rendered text now differs by a byte. The request then falls back
to a 4,096-token disk checkpoint and re-prefills 32,818 tokens (3.4 min).
The same failure at 200k tokens costs 201,906 tokens of prefill (~19 min),
and 14 of 14 such turns in one day took that path.

Normalise a whitespace-only reasoning body to absent at parse time so every
renderer treats it exactly like a turn that never had reasoning. Measured
after the change with --ds41-render: " ", "" and an absent key now produce
byte-identical prompts.

Adds test_render_blank_reasoning_is_absent(), which asserts a padded turn and
a reasoning-free turn render identically.

Suite status: `./ds4_test --server` has one pre-existing failure
(test_json_int_handles_non_finite_values, json_int("NaN")) that reproduces on
5379c1c without this change and is unaffected by it.
 2026-10-30 17:52:11 +01:00
nomish 97e32768e6 metal: use exact divisions in the V4.1 FP4 activation quantizer 
kernel_dsv41_quantize computes the E4M3 block scale as amax / 6 and the E2M1
value as |x| / scale with the default fast-math division. On an M5 Max the
quotient can land one ulp past an exact E2M1 rounding midpoint, so the value
rounds up instead of to even and make test-deepseek41-metal stops at:

    quantization mode=3 index=5428: 11.25 != 7.5

(two adjacent E2M1 steps at that block scale; deterministic across runs).

Use precise::divide for those two divisions so the kernel matches the
round-to-nearest-even reference. The E2M1 division is shared with the
power-of-two scale mode, where dividing by a power of two was already exact;
the FP8 path is unchanged.
 2026-10-30 17:50:17 +01:00
antirez 381808d839 Test CPU-touched Metal buffer reclamation 
Exercise tensor contents, read and write from a C caller without an outer autorelease pool. Check physical footprint after warmup, since ordinary heap statistics miss retained shared Metal buffers.

The bounded 32-allocation regression grows by 512 MiB without the lifetime fix and passes with it. Include the check in the existing Metal command-memory target.
 2026-10-30 17:49:58 +01:00
Mark Shank a4c922c560 Fix a Metal session memory leak on /ctx 
Under ARC, -contents retains and autoreleases its buffer. The C threads
calling ds4_gpu_tensor_contents/write/read have no pool, so a freed
session kept its CPU-touched buffers alive. Every /ctx in the CLI leaked
about 0.14 GiB with GLM 5.3 Flash at 64K context, and about 1 GiB per
session with V4.1 after a 2K prefill. Numerics are unchanged.
 2026-10-30 17:48:55 +01:00
Giulio Davide Carparelli 2768d0ed3c fix(jpeg): validate SOS spectral parameters in progressive scans 
jpeg_load_mem copied Ss/Se/Ah/Al from the SOS segment unchecked. The
progressive AC decoders loop with k <= Se and index jpeg_zigzag[k], so a
scan with Se > 63 read past the 64-entry table and used the result as a
write index into the coefficient block. This was reachable from request
images on the server when --vision is enabled.

Reject progressive scans that violate T.81 G.1.1.1.1: Se > 63, Ss > Se,
DC scans with Se != 0, interleaved AC scans, and Ah/Al > 13.

Add the issue's reproducer bytes as a fixture plus an invalid DC scan
case to tests/test_image_decode.

Closes #1292
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
 2026-10-30 17:47:37 +01:00
Dango233 3459d5546b docs: revalidate streaming decode queue on Flash 1017 2026-10-30 18:57:29 +08:00
Dango233 fc69b626ae Merge remote-tracking branch 'origin/main' into dango-codex/queue-1017-validation 2026-10-30 15:33:45 +08:00
HDYS-TTBYS 1e6a0e9602 Decode: give batched Engram tables separate row buffers 2026-10-26 23:16:44 +09:00
paradise-runner 694620956b Qwen decode: DS4_QWEN4_NO_PLE_OVERLAP restores up-front PLE staging; name the moe down_reduce n_hc slot 2026-10-26 06:36:53 -06:00
HDYS-TTBYS 1be7b4db2a Decode: give each Engram table its own input buffer 
The queued single-token decode drained after layer 13 only because the
layer-14 Engram upload reused the buffer layer 1 was still reading. Give
table 1 a second scratch buffer and point the layer at the buffer that was
just written, so the queue runs from layer 0 to the final drain without a
mid-token wait. TP keeps its layer-13 drain: that path cannot be exercised
on this machine and retains the scheduling it was validated with.

Bit-identical: logits at 3776 tokens, 64-token greedy output, and
--decode-consistency 16 all match the previous binary exactly
(max_abs 1.63013864, rms 0.310800667 at token 16066).

ds4-bench, 64 forced tokens, alternating runs:
  4K:  gen tps 37.69 -> 38.07 (+1.0%), first token 28.95 -> 28.45 ms
  32K: gen tps 35.94 -> 36.67 (+2.0%), first token 33.4 -> 28.8 ms
Prefill unchanged.
 2026-10-26 21:14:32 +09:00
Luo Hanjun d9f2d15b40 glm: bridge MTP verification across the indexed attention boundary 2026-10-25 15:22:31 +08:00
HDYS-TTBYS 5fd33a4e53 Set the candidate-decode flag after the Engram reads 
Early returns before the layer loop left candidate_decode set, which would
re-enable the per-row candidate build in a following prefill sweep. Move the
assignment past those paths and disable candidate scoring for SSD streaming,
where the compact path's CPU/GPU round trip has not been exercised against
the expert cache.
 2026-10-24 00:00:24 +09:00
HDYS-TTBYS 9d7a00b4fb Restrict candidate-row scoring to single-token decode 
The candidate list is only valid for the token that built it, and one shared
buffer cannot describe several rows. Copied-row paths (prefill sweep
fallback, batch index, step_batch) could therefore see a stale candidate
count and score layers 24-36 against another token's rows; the prefill row
loop also paid a drain and CPU sort for every row past 16K.

Build candidates only while ds41_graph_step runs, record candidate_pos, and
require candidate_decode && candidate_pos == pos before taking the compact
path. Every other path keeps the full-index scoring it had before.
 2026-10-23 23:04:12 +09:00
HDYS-TTBYS 01772c087a Optimize V4.1 long-context decode 
Three independent changes, all verified byte-identical against the previous
path:

- Decode top-k uses the prefix-merge kernel for the 512-row selection, so
  each merge reads only a run's top 512 entries instead of the whole run.
- Layers after the block-selecting layer 20 score only the candidate rows:
  the layer-20 block ids are sorted, expanded into a 16K-row list, scored by
  a row-indexed variant of the direct score kernel, reduced with top-k over
  the compact scores, and mapped back to original row ids. Layers 2/8/14/20
  and short contexts keep the full-index path.
  DS4_V41_DISABLE_CANDIDATE_SCORING=1 restores it.
- The direct score kernel alternates two partial-sum buffers, removing one
  barrier per head group without changing the accumulation order.

ds4-bench with 64 forced tokens, M5 Ultra + 1017 IQ2: ctx 4K 37.90,
64K 35.95, 256K 33.06, 512K 30.73, 1M 27.19 t/s. The 1M drop versus 4K is
28.3%, down from 49.6%. Logits are byte-identical with candidate scoring
on/off at 32K/64K/256K, prefill-parity and session-fixture pass, and
decode-consistency reports the same max_abs/rms.
 2026-10-23 22:31:35 +09:00
HDYS-TTBYS 57df7a5bf4 Pipeline V4.1 single-host batched decode command buffers 
Commit each layer without waiting in the batched step too. SSD streaming and
TP are excluded; the layer-13 drain is unchanged. The batched step also
carries short prefill and DSpark verification, and prefill-parity plus
session-fixture pass with the new scheduling.

Session-batch oracle: same max_abs (3.8147e-06), 45.5 vs 43.5 rows/s.
ds4-server --batched-session 8: 43.5/43.7 vs 40.7/41.1 tok/s. DSpark output
unchanged. DS4_METAL_DISABLE_V41_BATCH_DECODE_PIPELINE=1 reverts.
 2026-10-23 18:24:30 +09:00
HDYS-TTBYS be2773bc6f Pipeline V4.1 single-host decode command buffers 
Commit each layer's command buffer without waiting so the GPU runs layer N
while the CPU encodes layer N+1. The required drains stay at the Engram
handoff (layer 13) and at the final layer, and streaming, imatrix and TP
paths keep their existing scheduling.

Measured on M5 Ultra with the 1017 IQ2 GGUF: ctx 4096, 256 tokens,
27.35 -> 38.10 t/s; 8K prompt with 512 tokens, 27.08 -> 37.53 t/s; ctx 64K,
25.43 -> 34.53 t/s. Logits and greedy output are byte-identical, and
session-fixture, prefill-parity, decode-consistency and MTL_DEBUG_LAYER
pass. DS4_METAL_DISABLE_V41_DECODE_PIPELINE=1 restores the per-layer wait.
 2026-10-23 18:24:30 +09:00
Tarjei Mandt 72ccf36317 Revert the V4.1 decode pre-encode: held command buffer, per-token GPU spin 2026-10-22 15:15:36 +11:00
Tarjei Mandt 51fef8784a Merge pull request #9 from IngeniousIdiocy/v41-decode-on-kp7 2026-10-22 15:04:09 +11:00
Tarjei Mandt 3650da359e Turn off the V4.1 TP concurrent shared expert with the fixed-route switch 2026-10-22 11:54:16 +11:00
Tarjei Mandt ce5b718be8 Start V4.1 short-sweep suffix pruning at 512 rows 2026-10-22 11:54:06 +11:00
Tarjei Mandt 5b83cdfe8c Merge pull request #7 from ivanfioravanti/v41-tp-decode-engram-warm 2026-10-22 11:53:51 +11:00
Mark Shank e1d8114167 Encode the next V4.1 decode token while this one finishes 
At the end of each decode step the CPU waits for the GPU to finish the last command buffer. It now uses that wait to encode the next token's first command buffer (embedding, RoPE table and layer 0) ahead of time. That buffer reads the token id from GPU memory, so it can be encoded before the token is known. The next step submits it if its inputs are unchanged (same graph and position, and no other pass over the graph in between) and throws it away otherwise. A small kernel queued behind it keeps the GPU from going idle between the two steps. Sessions with a DSpark drafter don't use it, because they draft between steps.

The kernels, their order and the command buffer boundaries are the same as without it, so the results are bit-identical. A test that needs the model file checks this across two sessions, a prompt append and a session restore.

DS4_METAL_DISABLE_V41_PREENCODE=1 turns it off.
 2026-10-21 15:50:18 -04:00
Mark Shank 4ff3af604f Speed up the V4.1 decode top-512 
V4.1 decode picks the 512 best index scores for each token. For that case (V4 Flash is unchanged), each block is sorted with SIMD shuffles, each merge of two sorted runs keeps only the best 512, the 512 winners of the radix select are sorted together in one pass, and when scoring is limited to candidate blocks, the candidate list is built once per token. The selected indices and their order are the same as before on any input without NaNs. This builds on Tarjei Mandt's radix select and candidate-block scoring (af7c02b5, 2d5ef0d3).

DS4_METAL_DISABLE_V41_DECODE_TOPK=1 uses the previous kernels.
 2026-10-21 15:50:18 -04:00
Mark Shank 64ceeea4f5 Score long V4.1 decode index rows from bf16 keys 
The indexer's keys are FP4 values times a power-of-two scale, and every kernel that writes them rounds them to bf16, so a bf16 copy holds them exactly in half the space of the f32 key cache. Those kernels now also write a bf16 copy, and once there are 8192 or more index keys, decode scores them from the copy. The scoring math and its order are unchanged, so the results are bit-identical. Restoring a saved session rebuilds the copy, and if any restored key isn't exact in bf16, scoring uses the f32 keys. This is Metal only.

DS4_METAL_DISABLE_V41_INDEX_BF16=1 skips the copy entirely.
 2026-10-21 15:50:17 -04:00
Mark Shank e99dc21751 Spread two single-core V4.1 decode reductions over the GPU 
Two reductions in V4.1 decode ran on a tiny part of the GPU: the Engram gate used one 32-thread SIMD group per stream, and the hyper-connection collapse used a single threadgroup. The gate now uses eight SIMD groups per stream, and the collapse is split across 42 threadgroups in two kernels. Each thread still adds its values in the original order, so the results are bit-identical.

DS4_METAL_DISABLE_V41_SPLIT_REDUCE=1 goes back to the single-group kernels.
 2026-10-21 15:50:13 -04:00
Mark Shank f0b4f55fcc Compute one RoPE table per V4.1 decode token 
Every layer computed the same RoPE cos/sin values for the current position, in both the attention and the indexer. Decode now computes them once per token into a small table at the start of the step, and the layers read it.

With the table, the kernel that finishes attention can also undo the RoPE rotation and round to bf16 as it writes its output, so the next projection becomes a plain Q8_0 matrix-vector product. In the indexer, applying RoPE and quantizing the query to FP4 are now one kernel instead of two. The results are bit-identical.

DS4_METAL_DISABLE_V41_ROPE_TABLE=1 computes RoPE in each layer as before.
 2026-10-21 15:49:56 -04:00
Mark Shank 55b1e75e64 Narrow more V4.1 decode weights and read them with wider loads 
The indexer's q_b weights are FP8 (E4M3) in the original checkpoint, but the GGUF widens them to F16, as it does for the Engram projection. Decode now reads an exact FP8 copy of them, built when the model loads, which halves the bytes read. If any layer can't be converted without loss, every layer keeps reading F16.

The BF16 router and FP8 matvec kernels now load 16 bytes at a time, and the FP8 kernel reads each scale once for both of its rows. They multiply and add in the same order as before, so the results are bit-identical.

The copies are allocated like V4.1's model weights, without Metal hazard tracking. This builds on Ivan Fioravanti's FP8 and BF16 weight copies.

DS4_METAL_DISABLE_V41_INDEXER_E4M3=1 keeps the F16 reads. DS4_METAL_MODEL_TRACKED=1 turns hazard tracking on for the model weights and the copies.
 2026-10-21 15:49:41 -04:00
Mark Shank b66c926046 Compile V4.1's fixed decode shapes into its decode matvecs 
Some V4.1 decode matrix-vector products always run with the same sizes: the attention q_a and kv projections, the MoE gate and up projections, the attention output projection and the FFN's hyper-connection merge. They now use kernel variants with those sizes written in as constants, so the compiler can unroll the loops and fold the index math. The results are bit-identical. This extends Ivan Fioravanti's fixed-size kernels to more calls.

DS4_METAL_DISABLE_V41_LITERAL_SHAPES=1 uses the general kernels.
 2026-10-21 15:49:18 -04:00
Mark Shank 95606b5045 Fix a data race in the V4.1 FP8 weight-copy builder 
The FP8 weight copies are built by several threads at once, and any of them can set job->failed when a block can't be converted exactly. The flag was a plain bool written and read across threads. Make its accesses atomic, and have the other threads stop early once one has failed.
 2026-10-21 15:49:03 -04:00
Mark Shank 26335c7044 Overlap V4.1 decode hyper-connection weights with the projections 
Each V4.1 block takes its input as a weighted sum of the four residual streams, with weights from the previous block, and also computes its own hyper-connection weights (the mixer projection and its Sinkhorn split), which are first needed after the block's attention or FFN output. Single-Mac decode now computes the input first and the weights alongside the next projections: in attention, inside the query projection's dispatch, which also stores the layer's KV row; in the FFN, beside the routed experts' gate and up projections. Each part does the same arithmetic as before, so the results are bit-identical.

This comes from our earlier V4.1 decode work, with the short-K Q8 row helper that its query kernel uses.

This commit also adds two things the later commits use. Each V4.1 decode option is read from the environment once at startup, and with the DS4_GPU_TEST_COUNT_KERNELS test flag a test can count how often a kernel ran, to check that the new code was actually used.

DS4_METAL_DISABLE_V41_DEFERRED_HC=1 turns it off.
 2026-10-21 15:49:00 -04:00
ivanfioravanti f5127dfc82 Allow up to 256 images per server request 
ds4-server rejected requests with more than 16 images ("too many images;
at most 16 are allowed"). Clients that keep screenshots in an agent
conversation pass that within a few turns. The 16 was only the server's
own cap and the size of two stack arrays; the engine, session
checkpoints, TP mirroring and the KV disk cache all size image state
dynamically.

SERVER_MAX_IMAGES (256) now sets the cap, the error message and both
arrays (about 22 KB of stack at most, in request threads).
 2026-10-21 19:43:58 +02:00
ivanfioravanti 5a588dbf12 Share weight loads across two rows in the exact TP K-slice matvec 
Two-Mac TP prefill on pre-M5 GPUs computes each rank's half of wo_b with
kernel_mul_mv_q8_0_f32_rows8, one threadgroup per activation row. The
new kernel_mul_mv_q8_0_f32_rows8_tok2 takes two activation rows per
threadgroup: each lane dequantizes a weight block once into registers
and runs the one-row kernel's walk for both rows, and each row then takes
the one-row reduction. The partials, and so the TP logits, are unchanged.

TP prefill, 2x M3 Ultra, Q4 MXFP4, 2K/4K/8K appends: 445/435/538 ->
475/462/574 t/s, decode logits bit-identical. Other shapes measured
(4x4, 4x2, 8x4, 8x8, 16x2, 2x8, rows grouped through the cache) were
slower or equal. DS4_METAL_DISABLE_V41_TP_KSLICE_TOK2=1 restores the
one-row kernel.
 2026-10-21 17:54:49 +02:00
ivanfioravanti 565db43559 Prune the V4.1 decoder suffix of short sweeps without changing the output 
Layer il of a sweep only needs its last 1 + (39 - il) * 127 rows. Sweeps
of 2541+ rows already skip whole 2048-row tiles of the rest; shorter
sweeps and appends ran every row through all 40 layers.

Under 8192 rows, a sweep now also prunes inside the plain schedule's
chunks while keeping each kept row's arithmetic:

- chunks before the cut are skipped, and later ones run unchanged;
- the cut chunk starts at a 128-row boundary of the sequence (the
  compressor blocks) and keeps at least 256 rows, below which other
  router and expert kernels run;
- its F32 router picks the kernel the whole chunk would use
  (ds4_gpu_matmul_f32_plan_tensor): the tiled kernel gives each row the
  same sums for any row count, the row-count switch was the only change.

Decode logits after 2048 -> 8192, 777 -> 8192, 1000 -> 8000, 4100,
6500 and 8192 -> 16384 sweeps hash identically to the previous build.
Q4 MXFP4 on M3 Ultra: 2K/4K/8K appends 406/390/484 -> 479/467/576 t/s.
DS4_METAL_DISABLE_V41_SUFFIX_ALL=1 restores whole-tile pruning only.
 2026-10-21 17:20:24 +02:00
ivanfioravanti 187f46b4fa ds4-bench: per-step decode logits hash and dump 
DS4_BENCH_DECODE_LOGITS_HASH=1 prints an FNV-1a hash of every decode
step's full logits, outside the timed region, to compare arms bit for bit;
DS4_BENCH_DECODE_LOGITS_DUMP=path also appends the raw F32 logits for KL
checks. The encoder timeline keeps longer kernel names.
 2026-10-21 16:03:25 +02:00
ivanfioravanti bdae7640ef Warm the V4.1 Engram page cache when the tables fit in memory 
Prefill reads 2 x 24 Engram rows per token, spread across 189 GiB. That
random pattern is the worst case for external SSDs: a USB drive served
about 6K rows/s over the whole span, uncached or buffered alike, so a 2K
prompt stalled about 12 s before layers 1 and 14. Q2 single-Mac prefill
ran at 205-290 t/s while its GPU work alone supported ~420.

When the resident model, graph buffers, a 64 GiB reserve and the tables
fit in physical memory (Q2 on a 512 GB Mac; not Q4), a low-priority
background reader streams the tables once into the kernel page cache at
the drive's sequential rate (~1 GB/s over USB, about 3.5 minutes). The
uncached row readers then hit memory. No table is allocated or wired, the
pages stay reclaimable, and the warmer stops at memory pressure. The same
bytes are read either way, so logits do not change.

Q2 from a USB SSD, cold prompts: 2K 198 -> 420 t/s, 8K 214 -> 518 t/s.
DS4_DISABLE_V41_ENGRAM_WARM=1 skips the warmer; DS4_V41_ENGRAM_WARM_WAIT=1
blocks until it finishes, for benchmarks.
 2026-10-21 16:03:25 +02:00
ivanfioravanti e91c453ec2 Overlap more of V4.1 two-Mac TP decode 
- A balanced TP rank, which binds every expert and runs three of the six
  routed, takes the fixed-route static MoE kernels; their per-row walks
  equal the generic ones for any slot count.
- The rank's half of the shared expert runs in the routed experts'
  concurrent encoder: its lanes' BF16 mid, then a K-slice down matvec that
  rounds its input to BF16 as it loads (kernel_mul_mv_q8_0_f32_bf16in).
- out_b's K slice likewise rounds low to BF16 on load, dropping a pass.

TP logits are bit-identical. Switches:
DS4_METAL_DISABLE_V41_TP_FIXED_ROUTE, DS4_METAL_DISABLE_V41_TP_CONCURRENT_FFN,
DS4_METAL_DISABLE_V41_TP_KSLICE_ROUND_IN.
 2026-10-21 16:03:25 +02:00
antirez ee7e840523 cuda: improve SSD-streamed prefill with ordered expert reads 
Read selected experts in file order and coalesce adjacent transfers during
batched prefill. Make V4.1 whole-next-layer prefetch opt-in, while preserving
the original order for single-token decode and cache seeding.

Keep router remapping and cache budgets unchanged. Publish cache entries only
after uploads complete, and drain partial uploads before releasing failed
reservations. Retain and document the switches for restoring the old policy.

Validate the SSD-only integration with Metal and CUDA builds, 44 CUDA cache
checks in each of three modes, outstanding-reader teardown, and exact
129280-value frontier logits for both 1017 GGUFs in paired 2K prefill runs.
The earlier two-Spark default qualification covers larger contexts, tight
caches, greedy output parity and sanitizer checks for the same loader.
 2026-10-21 09:30:33 +02:00
ivanfioravanti 0f9b03c9df Speed up exact V4.1 prefill matmuls and indexer scores on M3 
- V4.1 BF16 Q8 matmuls on M3 Ultra take the existing unpack-then-F16-GEMM
  path when the batch is wide: +0.6-0.8% prefill, bit-identical.
- kernel_dsv41_indexer_scores_mma keeps each simdgroup's K^T fragments in
  registers and reads Q and head weights from device memory, so the head
  loop runs without threadgroup memory or barriers. Products, block order
  and head sum order match the tiled F32 kernel: bit-identical scores,
  1.31 -> 0.77 s at 16K.

Together, cold prefill 32K 860 -> 892 t/s and 64K 853 -> 898 t/s.
DS4_QWEN4_Q8_PREFILL_UNPACK=0 and DS4_METAL_DISABLE_V41_INDEX_MMA=1 roll
them back.
 2026-10-21 08:10:02 +02:00
ivanfioravanti 7d9903b2ed Read narrow decode copies of the V4.1 router and Engram projection 
The GGUF widens two decode weights: the BF16-source router is stored as
F32 and the Engram projection, an E4M3 checkpoint, as F16. Decode now
reads a BF16 copy of the router and the exact E4M3 codes with their
32x32 block scales, built once at graph allocation and freed with the
engine. The E4M3 matvec maps each code to an F16 bit pattern (exactly
value / 256, subnormals included) instead of decoding it.

Logits are bit-identical. Decode at 2K on M3 Ultra: 38.5 -> 39.2 t/s;
the Engram projection drops from 849 to 658 us per token.

DS4_METAL_DISABLE_V41_ROUTER_BF16=1 and DS4_METAL_DISABLE_V41_ENGRAM_E4M3=1
restore the wide reads.
 2026-10-21 08:09:09 +02:00
paradise-runner b5aedb9b96 Qwen hyper-connections: load the norm squares and inject partials ahead of their in-order sums 
Result: {"status":"keep","speedup":1.1865,"cand_tps":60.36,"ctrl_tps":50.98,"ratio_sd":0.0024}
 2026-10-20 09:00:20 -06:00
paradise-runner f768d146e1 Qwen attention decode: two query heads per simdgroup over six simdgroups 
Result: {"status":"keep","speedup":1.1736,"cand_tps":60.07,"ctrl_tps":51.06,"ratio_sd":0.0056}
 2026-10-20 09:00:20 -06:00
paradise-runner 6605c3cefb Qwen attention merge: load the split maxima sixteen at a time 
Result: {"status":"keep","speedup":1.1637,"cand_tps":59.4,"ctrl_tps":51.03,"ratio_sd":0.0017}
 2026-10-20 09:00:20 -06:00
paradise-runner 0f12ea4577 Qwen prefetch: factor the pending-range binding 2026-10-20 09:00:20 -06:00
paradise-runner 9000fb74f9 Qwen decode: hc norm prefetches its down projection into the system cache 2026-10-20 09:00:20 -06:00
paradise-runner 6b7cb44613 Qwen decode router: read a BF16 copy of the BF16-valued F32 rows 2026-10-20 09:00:20 -06:00
paradise-runner a3a7e2de6b Qwen Q2_K row dot: two aligned word loads per lane instead of eight bytes 2026-10-20 09:00:20 -06:00
paradise-runner 4ed25e4f07 Qwen single-token MoE: fold the weighted reduce into the down kernel 2026-10-20 09:00:19 -06:00
paradise-runner b440c3c20e Qwen F32 row dot: unroll so decode GDN alpha/beta loads issue ahead 2026-10-20 09:00:19 -06:00
paradise-runner d2ae13eac0 Qwen decode: read PLE n-grams while the GPU runs layer 0; flush layer 1 separately 2026-10-20 09:00:19 -06:00
paradise-runner fd16a72f52 Specialize Qwen IQ2_XXS/Q2_K decode expert rows on M5 2026-10-20 09:00:19 -06:00
paradise-runner 47c34fb444 Qwen IQ2_XXS decode mid: stage grid/sign tables in threadgroup memory 2026-10-20 09:00:19 -06:00
Tarjei Mandt e2d09d1b05 Schedule DeepSeek V4.1 DSpark drafts with the V4 scheduler 2026-10-20 11:59:28 +11:00
Tarjei Mandt f41806dcd2 Keep the DeepSeek V4.1 prefill partition when a DSpark drafter is loaded 2026-10-20 11:59:28 +11:00
Tarjei Mandt 732af32f9c Merge remote-tracking branch 'origin/main' into ds41-perf 2026-10-20 09:57:24 +11:00
Tarjei Mandt 75408d41a6 Remove the DeepSeek V4.1 DSpark runtime before merging upstream's 2026-10-20 09:09:49 +11:00
antirez a101246338 Add verified V4.1 Flash MXFP4 downloads and resumable assembly 
Download all nine 1017 MXFP4 shards, verify their published sizes and SHA-256 hashes, and restore the original page-aligned GGUF before selecting it. Preserve the input shards and check peak disk space before downloading.

Correct the 1017 Q2 checksum and document both targets. Validation: 16 offline downloader tests, shell syntax, and exact original-header reconstruction from the actual 1017 shard metadata.
 2026-10-19 14:14:09 +02:00
antirez 5379c1c9ac Download script now defaults to new DS4.1 Flash checkpoint. 2026-10-19 12:55:24 +02:00
antirez 902baebb70 Merge PR #1135: fix Metal compressed-attention causality 2026-10-19 10:03:07 +02:00
Tarjei Mandt 31786cb594 Begin DeepSeek V4.1 TP verify blocks as two gates per layer 2026-10-18 23:39:36 +11:00
Tarjei Mandt 5958ac718c Keep a restored DeepSeek V4.1 session's drafter off the ring rows it did not write 2026-10-18 22:22:59 +11:00
Tarjei Mandt 8c2b536717 Align DeepSeek V4.1 DSpark sampling and --dspark-strict with the DSpark paper 2026-10-18 22:22:58 +11:00
Tarjei Mandt f547bf91ca Rewind DeepSeek V4.1 sessions inside the last DSpark block 2026-10-18 22:22:58 +11:00
Tarjei Mandt c25ff9a4d4 Cut DeepSeek V4.1 DSpark per-cycle overhead 2026-10-18 22:22:58 +11:00
Tarjei Mandt c91c968b52 Sample DeepSeek V4.1 tokens with position-keyed Gumbel noise 2026-10-18 22:22:58 +11:00
Tarjei Mandt a3a231ce40 Verify DeepSeek V4.1 DSpark rows on the decode attention kernels 2026-10-18 22:22:58 +11:00
Tarjei Mandt 62c6684eb2 Keep the DeepSeek V4.1 matrix rows kernels to DSpark proposals and verify 2026-10-18 22:22:58 +11:00
Tarjei Mandt 2d5ef0d3c1 Score only the candidate blocks of later DeepSeek V4.1 index layers 2026-10-18 22:22:58 +11:00
Tarjei Mandt b71e7bcb8c Decode DeepSeek V4.1 without DSpark proposals under --dspark-strict and --quality 2026-10-18 22:22:58 +11:00
Tarjei Mandt 78d5d4d7fa Keep the DeepSeek V4.1 prefill partition when a DSpark drafter is attached 2026-10-18 22:22:57 +11:00
Tarjei Mandt 0a02a3cd56 Project DeepSeek V4.1 DSpark prefill rows without half-precision inputs 2026-10-18 22:22:57 +11:00
Tarjei Mandt 15a3d710d3 End the tensor parallelism bind line 2026-10-18 22:22:57 +11:00
Tarjei Mandt cbb687462b Detach the TP model residency set from the queue only on request 2026-10-18 22:22:57 +11:00
Tarjei Mandt b90fce71f4 Keep DeepSeek V4.1 short-prefill tails under 256 rows on their own tile 2026-10-18 22:22:57 +11:00
Tarjei Mandt aa5692e985 Merge pull request #6 from ivanfioravanti/ds41-perf 2026-10-18 21:48:33 +11:00
ivanfioravanti 8f194c8158 Use the wider V4.1 short-prefill tiles under TP too 
Each wider encoder tile also crosses fewer TP partial-sum gates. On two
M3 Ultras over RDMA, 7168-token prefill improves 603 to 621 tok/s; 14336
tokens is unchanged (748 to 750). Frontier logits match the narrow-tile
control byte for byte at both sizes.
 2026-10-18 00:19:36 +02:00
ivanfioravanti 559616bb19 Merge branch 'pr4-wider-tiles' into ds41-perf 2026-10-18 00:17:19 +02:00
ivanfioravanti 724336c8e9 Give TP big-gate arrival its own grace 
Layers 1 and 14 of each V4.1 prefill block first read the Engram rows for
up to 8192 tokens from the disk table on the host. When one rank's read
outlasted the other's gate timeout plus 2 s, the peer failed the big-gate
header exchange (layer 1 gate 1) and the prefill aborted: with the default
750 ms timeout, 64K and 128K prompts failed on two M3 Ultras.

The arrival handshake now waits up to 30 s (DS4_TP_BIG_GATE_ARRIVAL_MS).
The decode timeout and the payload exchange stay strict, and a dead peer
still fails at once on the closed socket. 64K and 128K prompts complete at
767 and 724 tok/s with the default timeout.
 2026-10-18 00:09:29 +02:00
ivanfioravanti 52780d4165 Detach the model residency set from the queue under TP 
With the residency set attached to the command queue, TP prefill on two
512 GB M3 Ultras hit sporadic stalls in the V4.1 decoder suffix: a command
buffer started 2-2.6 s after commit (GPU span ~10 ms), the cost of a full
residency request for the 285 GB of model views, while the peer rank waited
at the next gate. 8K prefill swung between 130 and 460 tok/s.

The views are made resident at load, so keep the committed set but detach
it from the queue in ds4_gpu_tp_init. 8K prefill is now 697-701 tok/s with
no stalls. DS4_METAL_TP_QUEUE_RESIDENCY_SET=1 keeps the set attached.
Tested with iogpu.wired_limit_mb raised and the wired collector disabled.

Also add DS4_TP_STALL_TRACE=MS to report gates whose GPU wait or exchange
exceeds the threshold.
 2026-10-17 23:19:19 +02:00
ivanfioravanti e255ed02d8 Use the fused MXFP4 routed MoE kernels under the TP expert split 
The pair-SwiGLU mm_id kernels lacked expert ownership, so TP ranks fell
back to the unfused path along with the compact tile, map scatter, tail
culls, down half LUT and pair half-scale variants. Teach the pair kernels
to zero the mid rows of unowned experts and to index the rank-local expert
slice, as kernel_mul_mm_id does, and drop the g_tp_split_world gates.

TP logits match the previous build exactly; single-Mac logits are
unchanged. MoE time per 4096-row chunk drops from 72 to 65 ms.
 2026-10-17 23:19:19 +02:00
ivanfioravanti 7e7e595455 Merge remote-tracking branch 'upstream/main' into ds41-perf 
# Conflicts:
#	docs/SPECULATIVE_DECODING.md
#	ds4.c
#	ds4_tp.c
#	tests/test_deepseek41_metal.c
 2026-10-17 13:30:19 +02:00
LEFBE 9f365a1748 glm: skip MTP head/attention for discarded-output catch-up forwards 
The accept path calls glm_graph_mtp_step() twice: once to advance the
MTP sub-model's KV cache to the just-committed position (output
discarded into a variable never read), then again to produce the real
next draft. The first call ran the full pipeline -- Q/RoPE/indexed
attention/attn_out/ffn_norm/router/routed-experts/shared-expert/residual/
head_norm/lm_head -- for a result that was thrown away immediately.

Everything from qk_low onward only feeds *draft_out; nothing past
kv_store (which writes the only state that outlives the call, the MTP
compact KV cache at this position) is needed when the caller doesn't
read the output. Add a need_output parameter and skip straight to
end_commands/cleanup when false, matching llama.cpp's headless-graph
pruning for the equivalent GLM5-Next MTP catch-up forward.

Verified byte-exact: ds4_test --session-snapshot DS4_TEST_GLM_MTP=1
reports identical accept/reject sequence (cycles=16 single=11 double=5
tokens=21) on real GLM-5.3-Flash-Q2 before/after. Two live-server
scenarios (short deterministic list, long 400-token
prose+code+explanation) are byte-identical against the unpatched
binary with --mtp on.

Does not fix the known MTP decode regression (see memory): per-cycle
timing (--mtp-timing) shows the catch-up call this touches was already
cheap (~4ms) on the accept path; the dominant cost is the unavoidable
extra full forward pass on every cycle (reject: ~32-36ms; the verify2
batch-of-2 step itself: ~35-41ms, constant) -- both untouched by this
change since their output is genuinely needed. Kept anyway: real,
verified-correct removal of wasted GPU work at zero risk.
 2026-10-14 08:36:24 +02:00
antirez d329b3c523 Reuse the RDMA verification window for V4.1 DSpark 2026-10-10 13:40:50 +02:00
antirez f38445cf46 Use batched CUDA F16 projections in the V4.1 drafter 2026-10-10 13:15:21 +02:00
antirez 0df64b67a9 Batch V4.1 verification attention outputs on Metal and CUDA 2026-10-10 13:15:21 +02:00
antirez 97079d7dcf Enable DeepSeek V4.1 DSpark over Metal and CUDA TP 2026-10-10 13:10:32 +02:00
Giorgio Oppo 8d543ecc95 Merge antirez/main into qwen-kernel-opt 
Resolve the two ds4.c conflicts while retaining upstream CUDA TP batching
and Spark confidence defaults, the branch's M5/M6 device policy, and
Qwen SSD memory admission without reapplying generic cache limits.

Validation on Apple M1 Max, 32 GiB, in an isolated clean checkout:
- Full Metal application build and CPU application build.
- Qwen memory, MTP policy and Metal device-policy tests.
- Qwen resident arithmetic in all four math/command-buffer modes.
- Qwen kernels, exact Q2, prefill reuse, Q8 prefill variants, HC math,
  and SSD selected-expert/cache/failure-retry tests.
- DSpark confidence, window, EOS and rewind host contract tests.
- git diff --check.

CUDA/ROCm hardware, distributed transport and full-model performance
were not tested by this merge validation.
 2026-10-10 12:26:41 +02:00
antirez 792e7e6aad Include the pinned revision in DSpark artifact QA 2026-10-09 18:28:56 +02:00
antirez b290d83bc5 Skip V4.1 candidate sorting when every block is retained 
Build the identical all-allowed mask directly below the 2048-block cap. This avoids the costly arbitrary-size sort on ROCm without changing sparse attention or routing. Verify exact masks across the boundary on Metal, CUDA and ROCm.
 2026-10-09 15:35:18 +02:00
antirez 784617c8f7 Document V4.1 DSpark conversion, usage and quality checks 2026-10-09 15:22:37 +02:00
antirez 4c64ae44c1 Enable V4.1 DSpark on Metal, CUDA and ROCm 
Add single-host batched verification, exact stochastic acceptance, bounded prefix rollback and session restoration. Keep the native drafter resident alongside target SSD streaming and account for its memory. Enable the V4.1 ROCm target graph and reuse batched causal index selection. Add native-weight, runtime and official-score verification coverage. Network TP speculation remains explicitly unsupported.
 2026-10-09 15:07:43 +02:00
antirez f3ee00a1d2 Keep ROCm exact-row Q8 projections on decode arithmetic 
Do not switch explicit decode-row calls to prefill activation precision. Reuse batched Q8 activation kernels, or scalar rows in quality mode. Verified attention output boundaries and the 100-case Flash 0731 official continuation suite.
 2026-10-09 14:04:04 +02:00
antirez 5c3b0a49ea Preserve ROCm router probabilities and signed-zero sort ties 2026-10-09 13:49:12 +02:00
antirez 2dcead8ce2 Allow DSpark conversion from only the native draft shards 2026-10-09 13:18:45 +02:00
antirez 87a95e0c11 Add the V4.1 DSpark draft graph and native-weight reference tests 2026-10-09 12:51:07 +02:00
antirez da5a1272cb Convert native DeepSeek V4.1 Flash DSpark weights to GGUF 2026-10-09 12:35:04 +02:00
antirez a615eb6f07 Batch V4.1 CUDA TP attention output projections 
Reuse the exact grouped Q8 projection for short prefills and concurrent sessions while keeping causal attention row-local. Preserve BF16 boundaries and use the original per-row path for unsupported layouts. Physical five-, seven- and eight-session and mixed-batch tests retain exact logits.
 2026-10-08 13:34:02 +02:00
antirez cef4a4d906 Balance V4.1 shared experts across CUDA TP peers 
Assign the shared expert to the rank with fewer selected routed experts, with deterministic ties. Preserve the same addition grouping in short batches and reject mismatched ownership settings at TP startup. Cover graph replay, router scratch reuse, batched masking and physical mode negotiation.
 2026-10-08 13:12:38 +02:00
antirez 3920942d7a Overlap V4.1 disk n-gram reads with CUDA decode 
Use four persistent readers per graph and await each table only at its consuming layer. Keep tables disk-backed, disable Linux read-ahead, and join pending reads on every exit. Cover reader reuse, failures, cancellation by destruction and concurrent readers.
 2026-10-08 13:12:15 +02:00
antirez 68ecb1b0c4 Fuse V4.1 CUDA normalization and HC BF16 stores 
Preserve the original reductions and every BF16 rounding boundary while eliminating separate rounding launches. Keep the unfused path available for diagnostics.

Add bitwise kernel tests covering ragged widths, multiple rows, in-place normalization and output guards. Dual-Spark runs measured a repeatable roughly one percent decode gain; CUDA memory checks passed.
 2026-10-08 12:43:22 +02:00
antirez 8a522702bd Use exact Q8 matrix kernels for CUDA TP prefill slices 
Respect the parent weight row stride when computing a K slice. Reuse the existing integer matrix path for prefill-sized batches without changing scalar block reductions. Extend grouped and sliced projection tests through 129 rows and the V4.1 output dimensions.

Dual-Spark V4.1 Q2 warm measurements: 438 to 548 tokens/s at 32K, and 337 to 412 tokens/s for an 8K append. Frontier logits are byte-identical; focused bitwise tests and CUDA memcheck pass.
 2026-10-08 11:38:32 +02:00
Salvatore Sanfilippo fc80bd695d Avoid full context rebuilds at DSpark response boundaries 
Retain the latest verifier frontier and restore a bounded token prefix when a response ends inside an accepted block. Preserve valid next-token logits for sampling-mode changes, and keep the existing rebuild fallback for unavailable or stale snapshots.

Coordinate TP restores with an acknowledged state-only command before replay. Add rewind, EOS, transport, and two-turn streaming regression tests.
 2026-10-07 21:43:05 +02:00
Salvatore Sanfilippo d6b516abb4 Score API top choices independently of alternative order 
Select maximum reported logprobs, count tied maxima once, and skip unmappable maxima without substituting lower-ranked tokens. Cover ordering, ties and damaged token bytes in the quality fixture.
 2026-10-06 18:35:04 +02:00
Salvatore Sanfilippo b14bf6220f Repair DSpark and ROCm regression tests 
Compile the HIP MMVQ test separately from its object-only link and use the existing ROCm fold stub. Update the EOS contract fixture for TP proposals and split verifier heads while retaining the historical negative control.
 2026-10-06 17:48:04 +02:00
Tarjei Mandt cad79332b3 Keep the DeepSeek V4.1 matrix rows, idle prefill and DSpark checks off other models 2026-10-06 14:58:14 +11:00
Tarjei Mandt ff95f53e87 Reset the radix top-k scratch for new row counts and keep it to DeepSeek V4.1 2026-10-06 00:08:04 +11:00
antirez f0962e3757 Stop tensor-parallel inference after a transport failure 
Refuse commands and scalar, batched and bulk gates once the link is
failed. Keep failed command sends and verify-window setup/completion
failures fatal, rather than attempting replay on unsynchronized queues.

A one-sided head-row failure previously allowed later commands and
decode exchanges to consume peer drain traffic as model partial sums.
Recovery now requires a fresh connection. Healthy inference arithmetic
and protocol messages are unchanged.

Add failed-context no-I/O checks, mocked drain/post/completion failures,
and a physical mismatched-head-row regression in both link directions.
The old code fails the new unit regressions; ASAN/UBSAN, TCP/RDMA fallbacks,
peer exits, physical Q2/MXFP4 speculation and server coding checks pass.
 2026-10-05 14:17:41 +02:00
antirez b3afc8e991 Document dual-Spark inference and its release checks 
Describe Q2/MXFP4 network tensor parallelism, DSpark defaults, matching
peer settings, and measured performance with workload and timing scope.
Document native drafter conversion without changing the shipped default.

Extend future QA with physical RDMA, memory ownership, active-stream and
graph checks, official continuation scoring, context boundaries, server
coding sessions, cancellation and fallback coverage.
 2026-10-05 14:17:41 +02:00
antirez 63059b654c Optimize Spark attention and support native MXFP4 DSpark experts 
Coalesce bounded F16 projection copies, specialize Q8 vocabulary work,
tile indexer scoring and compact top-k sorts, and accelerate indexed
attention. The single-row three-pass path preserves scalar arithmetic,
including runtime reciprocal-square-root rounding and reduction order.

Fix divergent top-k compaction decisions. Admit native MXFP4 drafter
experts without changing the default support format. Keep unsupported
shapes and allocation failures on their existing fallbacks.

Strengthen expert-identity fixtures and wire the retained kernel and
history regressions into QA targets. Final single-row tests cover graph
replay, masks, tails and independent references; memory, race and sync
sanitizers pass. Official and continued-context full logits match the
preceding control. Physical Q2/MXFP4 speculation, server tool calls,
prefix reuse, cancellation and V4.1 regression checks are covered.
 2026-10-05 14:17:41 +02:00
antirez 43a91aaa3b Distribute and accelerate Flash speculative decoding on Spark 
Split the drafter and verifier vocabulary head across peers, negotiate
their state and memory requirements, and retain local drafting when the
split is unavailable. Use an ordered RDMA command mailbox for verification
and commits. Preserve accepted-prefix history through restores and skips.

Balance MXFP4 experts, stage short expert blocks, overlap independent
projections and shared experts, batch compressor and frontier copies, and
specialize measured Q8 projection shapes. Fix active-stream ordering,
paired projection row coverage and the model's indexer top-k threshold.

Add independent numerical references, graph and stream tests, transport
failure tests, and speculative session lifecycle coverage. The fused HC
path preserves the mathematical graph but changes rounding; official
continuation scores are mixed, not universally improved. A pre-existing
scalar-versus-batched book oracle discrepancy remains under investigation.
 2026-10-05 14:17:41 +02:00
antirez 77570b465c Add two-host CUDA tensor parallelism for DeepSeek V4 Flash 
Support resident Q2 and MXFP4 Flash across two CUDA peers. Split attention,
shared experts and short speculative work; reuse bounded CUDA scratch and
keep short embedding and Markov operations on the GPU.

Correct shared/routed contribution ordering and quantized vector bounds.
Keep eight-bit activations for MXFP4 matrix kernels. Align DSpark target
positions, confidence inputs and rolling history across GPU backends,
including continued prefill and skipped proposals. This extends the
reference alignment introduced by Donato Capitella in 394865c5.

Include kernel references, bounds checks, cache-history tests and physical
TP coverage. The cross-backend history changes still require physical
Metal and ROCm release QA.
 2026-10-05 14:17:41 +02:00
antirez 962ac1913d Improve V4.1 CUDA attention and Spark RDMA exchanges 
Batch causal attention selection, preserve exact FP8 Engram conversion,
and fuse shared-expert activation stages. Fix token-tile shared addressing.

Negotiate larger RDMA messages and let CUDA use registered exchange
buffers directly. Retain transport and allocation fallbacks, with focused
attention, buffer-ownership and physical-link regression tests.
 2026-10-05 14:17:41 +02:00
Tarjei Mandt 4f806387f7 Keep the DeepSeek V4.1 wide index scorer exact and off the GLM path 2026-10-05 11:21:26 +11:00
Giorgio Oppo 37f66ba3c9 Metal: restore main arithmetic for resident Qwen prefill 
Scope the legacy half/unsplit projection policy and parent-chunk attention
dispatch to SSD streaming graphs. Resident prompts again use main's FP32
small-batch F16/Q8 projections, F32 split-K reductions and per-range attention
at the dense/sparse boundary. This corrects a policy regression that changed
GDN state and subsequent greedy tokens. Decode and SSD arithmetic are retained.

Add a bounded regression test with a frozen 0aaea5a dispatch oracle, independent
attention references, negative controls, guard regions and immutable inputs
and weights. Exercise retained/unretained command buffers and fast/safe Metal
math. Require SSD mode for the existing legacy-unsplit generation reference.

Archive source provenance, replay helpers, a frozen HC shader oracle, the
resident fixture generator and full-vocabulary/greedy A/B captures. Exclude
models, executables, object files and generated source checkouts. Pending
experimental production kernel changes remain outside this commit.

Validation on Apple M1 Max, 32 GiB:
- Build ds4, ds4-server and regression tests from exactly the staged sources.
- Four Metal validation/math/command-buffer configurations pass; the old
  resident policy fails the regression test as expected.
- Resident repeated-weight Q4 fixture: all 248320 logits are bit-identical to
  main for 29, 575 and 5942-token prompts; both 64-step greedy dumps match.
- Full Q4 SSD: 29 and 575-token prompt dumps are byte-identical before/after.
- HC shader reference: 71793 comparisons, no differences.
- One-shot/session generation and SSD legacy-reference tests match all logits
  across the prompt frontier and eight continuation tokens.
- Memory planning and whitespace checks pass.

The resident fixture uses repeated real layer weights and zero PLE. Results
are arithmetic regressions, not quality or performance claims; the tester's
full Q2/M4 Max configuration still needs external confirmation.
 2026-10-03 12:02:43 +02:00
Giorgio Oppo 575369b19d Research: isolate Qwen prefill arithmetic and token drift 
Archive the investigation of PR #1056 comment 5963229415 comparing main
0aaea5a with 02313128 on M4 Max/Q2 resident. Trace the device-independent
prefill policy introduced by 2d6a207b: small F16/Q8 prompt batches change
operand precision, while GDN alpha/beta F32 projections stop using split-K.
Distinguish this policy from the numerically neutral M2-M4 tuning gate,
89986c3c attention partitions and the secondary HC arithmetic candidate.

Record ten full-vocabulary captures and four 64-step greedy comparisons
from one frozen HEAD snapshot on M1 Max/Q4/SSD, with inherited DS4 flags
removed. Restoring only F32 split-K changes the 575-token frontier by
0.41702342; restoring the complete dense policy produces identical logits
on that fixture. The 29/575 repeated HEAD controls are byte-identical.
The 5942-token tail-only split ablation changes logits by 0.33132219.

With ordinary T1 decode policy held constant, the dense-policy ablation
changes the first generated token at step 59 for Rome29 and step 20 for
book575. Preserve raw logits, top-20 logprobs, logs, token-ID fixtures,
commands, source hashes and independent provenance/numerical audits.
These fixtures do not reproduce the tester's exact M4/Q2 prompts or
attribute every component of the reported deltas; they are not quality
or performance evaluations.

Provide a preparation script that recreates 103 tracked source files from
0231312820, applies the isolated diagnostic
patch and verifies all source hashes. Add fresh snapshot/output arguments
to both runners and document serial reproduction using the local Q4 GGUF.
Keep generated snapshot copies, executables and object files outside Git.

Validation:
- Independent recalculation of all ten 248320-logit dumps and greedy
  common prefixes, selected token IDs and bytes passed.
- Preparation in a fresh temporary directory reproduced all 103 source
  hashes and applied only the reviewed diagnostic patch.
- Python AST/JSON parsing and both runner help commands passed.
- git diff --cached --check passed.

This commit archives research only; pending production optimization
experiments remain outside its scope. No production dispatch is changed.
 2026-10-03 09:21:17 +02:00
Donato Capitella 67f2ca912c Server: advertise loaded model modalities 2026-10-01 10:44:24 +01:00
Giorgio Oppo 0231312820 Metal: extend Qwen M1 Max tuning to M2, M3 and M4 
Share the established M1 Max dispatch policy with the M2-M4 families
without changing kernel arithmetic, supported tensor shapes, math-mode
guards or diagnostic overrides.

- Centralize device eligibility with generation-bounded name matching.
  Preserve M1 base/Pro/Ultra and M5/M6 defaults, and retain M3 Ultra's
  dedicated MoE specialization, geometry and prefill tile policy.
- Enable wide Q8 decode row reuse, HC gate/mix reuse, short MTP Q8
  projections and the existing IQ2/Q2/MXFP4 expert paths on M2-M4.
- Extend small HC prefill tiles, large Q8 SSD prefill, sparse IQ2/Q2
  SSD tiles, Q4/MXFP4 K32 tiles, large-batch half staging and Q4 decode
  scratch release under the same existing shape and batch guards.
- Update SSD dispatch probes, cover 47 device-name cases and document
  the supported families and the limits of current hardware evidence.

Validation: build all application targets and run the device-policy,
Qwen kernel, safe-math reuse/projection, MoE specialization and SSD
expert tests (retained and unretained command buffers) on an isolated
HEAD export containing only this commit. Local hardware is M1 Max;
M2/M3/M4 throughput still requires validation on those devices.
 2026-10-01 08:04:59 +02:00
Giorgio Oppo f32a16a561 Research: archive Qwen Q4 SSD scheduling and decode experiments 
Preserve the September 30 radical optimization prototypes and their measured
limits without enabling a new production path.

- Archive frozen MoE shaders and a standalone Metal oracle for compact tile
  scheduling with heavily used experts first. The paired screens measure
  +8.51% and +10.07% MoE throughput on captured layer 9/47 routing, with
  unchanged FP32 outputs. The projected whole-prefill benefit is only about
  2% and has not been demonstrated end to end.
- Include raw timing samples, exactness results, source hashes, and a lossless
  495 KB routing fixture containing indices only. Document direct rebuild
  commands and the limitations of the historical preparation script.
- Record twelve plain/MTP model runs with matching output bytes. MTP improves
  Roma decode by 12.59% versus the ordinary CLI; same-session comparisons
  show 19.44% on Roma and 24.24% on code. Preserve the extra memory/I/O costs
  and distinguish these results from untested architecture proposals.
- Document elastic prefill-to-decode workspace, inter-layer SSD prefetch,
  prompt-lookup drafting, and the rejected lossless weight-compression screen.
- Bundle both prompts and retain raw stdout bytes and frozen source hashes
  without line-ending or whitespace normalization.

Validation for this archive: standalone Objective-C oracle compiles; Python
syntax and JSON parse; frozen source SHA-256 values match; routing gzip,
record hashes and expert/pair membership validate; copied Roma prompt matches.
The GPU and model results are the recorded September 30 runs, not rerun here.

Earlier local production bundles, requantization work, executables and model
weights are outside this commit. Recorded model source/binary hashes identify
the working tree used; this archive alone does not reconstruct that binary.
 2026-10-01 05:18:07 +02:00
Emilian Bold 7172b249a8 Recover Metal GLM prefill speed with validated selection prefixes 
Scan selected IDs on the GPU before using the valid-row specialization. Trim only invalid tail slots, retain guarded handling for interior holes, and preserve the physical selection stride and attention order.

Enable the fast path for full head groups outside TP, with a diagnostic rollback. Add CPU-reference and bitwise rollback tests for padded, interior-invalid, empty and offset selections.

On M5 Max with resident GLM-5.3-Flash-Q2, main-based interleaved continued-prefill measurements improve 7.2% at 16K and 7.1% at 24K, with bit-exact full logits. Addresses the dispatch regression identified in #1029 without reverting invalid-row masking.
 2026-09-29 03:40:50 +03:00
Giorgio Oppo b96a12be7b Metal: optimize Qwen Q4 SSD prefill on M1 Max 
Add K32 NT4 expert matrix kernels for Q4_K/MXFP4 SSD prefill.
Reuse Q4 headers across eight staging steps and reduce threadgroup
storage while retaining the original sequence of FP32 MMA operations.
Keep the existing K64 entry points for other devices, resident inference
and other token tile sizes.

Use half activation shadows for the Q4/MXFP4 pair at 8192+ tokens,
preserving the public float intermediate. Enable the existing Q8 unpack
plus F16 GEMM path for large Qwen SSD batches on M1 Max. Retain scratch
allocations through unretained command buffers and fall back to direct
Q8 if optional allocation fails. Release prefill-only Q4 scratch at the
existing synchronized transition to row kernels, without adding a sync.

Extend validation with a frozen ee79f22 Metal oracle, exact Q4/MXFP4
float/half comparisons, SSD dispatch probes, failed-read retries,
scratch growth and reuse, prefill/decode/prefill transitions, and Q8
scratch lifetime checks in both command-buffer ownership modes.

Validation:
- Forced full build; no compiler warnings.
- Q4 oracle in default and safe math: 331,727,600 exact float values and
  121,446,912 exact half values per mode; IQ2/Q2 oracle also passes.
- SSD selected-expert and Q8 prefill variant suites pass.
- Resident K64 oracle remains exact and shows no observed slowdown.

Full-model results and diagnostic methodology are recorded in
speed-bench/qwen38-m1-q4-k32.{md,json}. All compared outputs are identical;
expert-cache capacity and requested SSD bytes are unchanged. The report
retains the earlier negative decode results and the subsequent controlled
comparison, rather than treating the exploratory gain as guaranteed.

Eight controlled ABBA+BAAB runs on M1 Max/32 GiB, same executable and
shader, toggling only the new prefill defaults in a diagnostic copy:
- Prefill: 154.31 -> 161.49 token/s (+4.65%).
- Decode: 4.0656 -> 4.0569 token/s (-0.21%), within observed variation.
- Later 50-forward decode window: 226.931 -> 221.873 ms/token (+2.28%).
The earlier -8.06% decode result did not reproduce in this control. Scratch
release is a verified memory reduction, not an established explanation
for that earlier result. No diagnostic rollback flag enters production.
 2026-09-27 17:23:11 +02:00
Giorgio Oppo ee79f22b90 Metal: accelerate Qwen MXFP4 prefill and batched down rows 
Decode each MXFP4 prefill half-block with one exponent read and four
byte-aligned packed loads. Keep the original LUT multiplication and final
half/float conversion, without changing accumulation or tile geometry.

Reuse activations across two MXFP4 down rows for M1 Max multi-token NR2
dispatches, preserving each row's block, FMA and reduction order. Use the
existing paired Q8 helper for shared experts and guard odd output tails.
Specialize this through function constant 907 and include it in the
pipeline cache key, compiling the extra register footprint out of T1.
Retain the existing NR1/NR4, other-device and safe-math fallback policies.

Add a standalone staging oracle against two unchanged stage8 decoders,
covering all exponent bytes, nibble values and byte alignments with both
storage types and math modes. Extend exact down tests with odd row counts,
block tails, T1/T2/T3, shared Q8/MXFP4 slots, NR1/2/4 and output guards.
Register the oracle build/test/clean targets and document measured gains.

Validation on Apple M1 Max:
- Build all binaries and the affected Metal test targets successfully.
- Staging oracle: 1,048,576 comparisons, zero NaN payload differences.
- Exact MV tests pass with DS4_METAL_MATH_SAFE=0 and 1.
- MoE MM specialization: 132 passing cases per math mode.
- SSD experts: 94 passing cases per math mode, including overlap,
  selected-only reads, cache reuse, failure/retry and fallback paths.
- Original/current intermediate output fingerprints match in both modes.
- Eight full Q4 SSD runs, including MTP, produce identical output within
  each workload; git diff --check passes.

Against 46425263, alternating resident prefill measurements reduce median
kernel time by 5.2-6.1%. Same-process two-token down measurements improve
throughput by 14.7-17.2%, with bit-identical results and no T1 GPU regression.
Full SSD prefill improves from 22.94 to 23.41 t/s in the measured workload;
median MTP generation stays at 7.65 t/s. These local results do not imply
the same gains for complete inference or other GPUs.
 2026-09-27 10:24:37 +02:00
Giorgio Oppo 4642526317 Metal: optimize Qwen Q4 decode and verify SSD parity 
Reduce Q4_K gate/up memory traffic by sharing Q8 activation reads across
the shared expert's two projections. Extract the existing IQ2 helper for
reuse, and widen routed activation and scale loads with packed float4 and
half2 operations. Preserve the scalar quantized-code loads and the original
per-lane FP32 accumulation and SIMD reduction order.

Enable MXFP4 down prefetch automatically on M1 Max for the measured
640-to-2560 projection. Preserve existing M5/M6 fast-math policies and the
diagnostic disable switch. Track the compiled library's math mode so
strict math keeps the plain noncontracted kernel even when prefetch is
explicitly requested; the prefetched kernel's explicit FMA otherwise
introduces a one-bit difference in that mode.

Strengthen tests with a plain-kernel oracle independent of automatic
prefetch, explicit/default dispatch checks, and Q4 SSD coverage at the
32/33 and 128-token boundaries. Compare addressed and masked Q4 decode
against the generic resident kernel across cache reuse, fallback and retry.
Add MXFP4 down microbenchmarks and document the full-model measurements.
Keep the existing prefill and row-dispatch policies: the experimental
smaller Q2 sparse-prefill tiles did not show a stable Q4 benefit.

Validation on Apple M1 Max 32 GiB:
- make -j2 all and the affected Metal test targets; git diff --check.
- Q4 ordered geometry and MoE decode exactness with MATH_SAFE=0 and 1.
- IQ2 shared-projection regression: 180 fixtures, 1,119,600 comparisons.
- MoE prefill specialization: 132 passing cases per math mode.
- SSD expert regression: 94 passing cases per math mode, including
  missing-only reads, overlap, output guards, failure and exact retry.
- Q4 SSD, ctx=32768, chunk=1024, temp=0, nothink, 100 generated tokens:
  three alternating pairs improve median decode from 6.07 to 6.42 t/s
  (+5.8%), with identical output in every run. Prefill medians are 6.94
  and 6.91 t/s. SSD/system variation remains substantial; these are local
  measurements and do not establish gains on other devices.
- Q4 MTP with two draft tokens: before/after 64-token output identical.
 2026-09-26 15:35:37 +02:00
Giorgio Oppo 81ee4badb8 Qwen: honor Metal's SSD streaming memory budget 
Avoid discounting Metal's recommended working set twice when sizing the
Qwen SSD expert cache. Use the smaller of 87.5% of physical RAM and the
reported recommendation, then retain the existing graph, staging and
2 GiB runtime reservations.

Let Qwen's complete admission calculation handle both automatic and
explicit cache budgets. Keep admitted static weights eligible for mlock
instead of applying the generic cache cap again after admission. Other
model families and backends retain their existing checks.

Add coverage for host- and Metal-limited budgets, unavailable capacities
and overflow-safe arithmetic, and document the resulting memory policy.

Validation:
- make test-qwen4-memory
- Qwen Q4 SSD runs on M1 Max 32 GiB: 6.32 GiB static weights locked,
  11.58 GiB expert cache, 2.59 GiB staging and 22.96 GiB planned total
  at ctx=32768 and prefill_chunk=1024.
 2026-09-26 15:35:21 +02:00
Sebastian Christiansen e5b19e280e Fix compressed-key causality in Metal visual attention 
Intersect compressed-key selection with each query's causal frontier while preserving raw image bidirectionality. Add an analytic regression covering selection patterns, compression boundaries, FP16/FP32 keys, and tensor view offsets.

Validated on M5 Max 128 GiB with the Vision-Exp IQ2XXS model: baseline fails the new regression; patched Metal kernels, 60 image generations with exact replay, resident/SSD session checks, long-context recall, tool-call quality, frontend tests, and CPU/Metal builds pass. Strict image grading remains 10/12 per arm due to screenshot wording; existing SSD output differences are unchanged. Balanced text throughput changed +0.0033% prefill and -0.2598% generation. Official Flash-0731 vectors, CUDA runtime, and distributed execution were not validated.
 2026-09-26 15:11:33 +02:00
Giorgio Oppo 089d966069 Qwen: make MTP prompt preparation configurable 
Resident short-context workloads can pay for preparing every historical
predictor row without recovering that cost during decode. Add an explicit
DS4_QWEN4_MTP_PREFILL policy so that this tradeoff can be controlled while
MTP draft generation and target verification remain enabled:

- on (default, also 1): prepare predictor history for all deployments.
- off (also 0): skip historical projections and initialize those rows to zero.
- auto: prepare history only when using SSD streaming.

Keep the existing default until resident acceptance and total wall time
are measured. A cold predictor prefix changes draft context and can reduce
acceptance; it is not equivalent to fully prepared history. Reject invalid
settings before mutating the Qwen session, and ignore the setting when MTP
is disabled.

Initialize skipped K/V and raw indexer rows deterministically, synchronize
before Metal host writes, and recompute completed pooled indexer blocks to
preserve real keys from an existing partial block. Defer each cold chunk's
last row to normal boundary preparation so uninterrupted prefill, append,
and checkpoint restore share the same predictor state. Apply the policy
only to new prompt rows; retain reused and restored cache contents.

Document the modes, acceptance tradeoff and checkpoint behavior. Add 64
model-free policy cases and extend the Metal prefill suite with poisoned
caches, reuse, invalid settings, append, mixed real/zero pooled blocks,
snapshots, cancellation/resume, target-logit comparisons and active MTP
including forced depth-3 verification. Keep the full-history oracle on.

Validation:
- Built ds4, ds4-server and the Qwen MTP prefill test.
- All 64 model-free policy cases passed; git diff --check passed.
- Short Metal model suite passed on M1 Max with SSD streaming and the local
  Qwen Q2-AProjQ4 GGUF, including greedy and exact-sampling continuations.
  The model test used a temporary source copy with existing Q4_K attention
  loader support and generic GEMV from the separate Q4 branch; those Q4
  compatibility changes are not included here.
- Resident 3.5k performance and CUDA execution were not measured; no
  throughput or acceptance improvement is claimed.
 2026-09-26 13:59:29 +02:00
Emilian Bold c71ca56d81 Overlap bounded Metal prefill layer submissions 2026-09-25 19:23:59 +03:00
Giorgio Oppo b1af94b438 Merge antirez/main into qwen-kernel-opt 2026-09-23 09:40:35 +02:00
Giorgio Oppo 028b43f3a1 Metal: stabilize Qwen HC arithmetic and extend M5 tuning to M6 2026-09-23 09:36:00 +02:00
Giorgio Oppo 47ca2eb5d9 Metal: enable Qwen SSD prefill overlap on all devices 2026-09-23 09:17:35 +02:00
Tarjei Mandt 3d3c83bbbd Merge remote-tracking branch 'origin/main' into ds41-perf 
# Conflicts:
#	ds4.c
#	ds4_engram.c
 2026-09-21 08:48:08 +10:00
Tarjei Mandt 798c64f444 Round a short DeepSeek V4.1 sweep down only for a token-major tail 2026-09-21 08:14:17 +10:00
ivanfioravanti 6bf43ee8da Amortize V4.1 short prefill with wider encoder tiles 
Port the short encoder tile selection from 4f9a2e0 onto KernelPool ds41-perf a15028e. Limit it to non-draft, single-node pre-M5 Metal inference and preserve decoder partitions.

Four alternating-order Q4 trials on M3 Ultra: 7956-token initial prefill improves 572.31 to 592.74 t/s (+3.57%); a 15894-token continuation improves 694.63 to 700.44 t/s (+0.84%). Exact live state, history and logits match in resident and SSD modes, including following decode. Add reproducible graph checks.

Builds and CPU syntax check pass. make test reproduces the same 14 V4 fixture assertions seen on previous baseline main; seven remaining recipe checks pass separately. CUDA hardware, distributed inference and DSpark were not rerun.
 2026-09-20 13:23:48 +02:00
Tarjei Mandt 92197579f8 Keep DeepSeek V4.1 CUDA row batches on the per-row decode kernels 2026-09-20 21:00:59 +10:00
Tarjei Mandt 2a55193649 Guard the DeepSeek V4.1 Metal-only helpers for CUDA builds 2026-09-20 21:00:59 +10:00
Tarjei Mandt 64240fb9c2 Prefill DeepSeek V4.1 prompts under 8192 tokens on the plain schedule 2026-09-20 21:00:59 +10:00
Salvatore Sanfilippo 0aaea5a238 Document physical Metal TP QA preflight and native batch coverage 
Check per-boot RDMA addressing and wired residency limits, match prefill chunk overrides on both ranks, and account for the extra sessions used by numerical oracles. Require at least three Metal TP rows when checking the native V4.1 path; two rows intentionally use ordered decoding.
 2026-09-20 03:21:24 +02:00
Salvatore Sanfilippo 6687c6f90b Guard ROCm memory baseline out of CPU-only builds 
The CPU-only placement tests also define DS4_ROCM_BUILD, but never use the GPU memory baseline. Exclude its declaration there, matching the existing guards around its uses.

Verified warning-free ROCm applications and test builds on Strix Halo, with placement, memory, session, server and agent tests passing. The production ds4.o is byte-identical after stripping debug metadata.
 2026-09-20 03:12:38 +02:00
GiorgioOppo a6ad6361ca test(qwen): add sparse and long-context MTP replay diagnostics 
Add an optional 17408-token synthetic fixture to the real-model Metal
MTP prefill regression. Run one trunk graph with SSD streaming and probe
actual residual rows at the first sparse query and near the final prefix.
Save only predictor state to keep the test practical on a 32 GiB M1 Max.

Replay cache-only, optimized last-row and recursive predictor paths
against the frozen full-layer reference using identical saved history,
trunk rows and lookahead IDs. Cover T1/T2/T3 and exact T3 verification;
poison destination cache rows before replay and compare complete logits,
EH projections, residuals and cache publication. Restore predictor state,
RoPE positions, verification policy and the device target head afterward.

Provide bounded chunk/prefix options and optional token, tensor and JSONL
dumps for comparisons across builds. Keep the original short regression.
Document usage, the actual 8192-row threshold of the M1 half-operand path,
measured results and limits of the acceptance investigation.

Validation on M1 Max, 32 GiB, Qwen3.8-Flash-Next-Q2, SSD cache 1024:
- Short MTP regression, memory tests and production build pass.
- The long fixture passes at chunks 128/2048 on 11811b9 and at chunk
  2048 on 04c0867 and 89986c3, with Metal API validation.
- At chunk 2048 all 36 captures from 89986c3 and 11811b9 are byte-identical.
  Differences from 04c0867 already occur in the corrected trunk inputs.
- Between current chunks 128/2048, captured trunk rows and final target
  logits match exactly. Maximum predictor delta is 0.0001266003; the
  16 full/recursive proposals retain the same argmax.
- A separate same-row grouping diagnostic finds pre-existing HC projection
  rounding differences, with identical stage reports on 04c0867/11811b9.
- The independent half-kernel oracle passes 31 fixtures including 8192
  rows under API validation. A 24-fixture NT1/2/4 subset passes shader
  validation; the NT8 mid stress case exceeds instrumented M1 shared
  memory capacity (36864 bytes versus 32768).
- git diff --check passes. No production arithmetic or dispatch changes.

These synthetic checks do not reproduce the reviewer's acceptance counts
without the original prompt. Full-model chunk-8192 acceptance is not
covered by these runs. Temporary diagnostics and captures are excluded.

Refs: https://github.com/antirez/ds4/pull/1056#issuecomment-5744306216
 2026-09-19 21:19:31 +02:00
GiorgioOppo fd0c24af54 test(qwen): repair compact MoE checks after half-kernel extraction 
The compact-dispatch regression still expected expert mapping in the
float gate/up entry point after that logic moved to the shared float/half
implementation. Follow the shared helper and validate both wrappers,
including template arguments, grid coordinates and shadow-buffer wiring.
Check that expert bounds precede counts reads and cover every half-down
NT instantiation. Ignore comments when inspecting shader source.

Add 29 negative controls for broken mappings, missing or late guards,
incorrect wrapper arguments and missing half-down variants. Extend the
CPU streaming stub with the half-pass association used by the extracted
production guards; verify invalidation on empty passes and preservation
on nonempty passes with and without active streaming state.

Validation: make test-qwen4-moe-mm-compact passes with ASan/UBSan across
1033 mapping cases, 3200 schedules, 961600 matching events, 3 pointer
lifetime cases, 12 NAX empty-state cases and 4 streamed half-state cases.
This is host/dispatch coverage; production kernels are unchanged.

Refs: https://github.com/antirez/ds4/pull/1056#issuecomment-5744306216
 2026-09-19 21:19:15 +02:00
Salvatore Sanfilippo 06f7135d0f Keep Metal-only Qwen batch tests out of the CUDA suite 
Guard native cross-session attention, grouped expert, and Q8 batch GEMM tests with __APPLE__. CUDA retains its supported scalar, prefill, and recurrent kernel oracles instead of failing to link against Metal-only entry points.
 2026-09-19 20:23:37 +02:00
Tarjei Mandt a15028ee40 Round DeepSeek V4.1 vectors with the shared bf16 helper 2026-09-20 01:51:24 +10:00
Tarjei Mandt 8be1005089 Fold the one-row DeepSeek V4.1 entry points into their rows form 2026-09-20 01:51:24 +10:00
Emilian Bold 0719a0baf8 Improve Metal router softplus accuracy for small logits 
Use a four-term log1p series for exp(x) below 1/32 in all four DeepSeek router transform sites. Clamp the polynomial input so its unselected branch stays finite, and compare Metal probabilities against the shared accurate reference without widening the tolerance.

Refs antirez/ds4#1039. Validated on M5 Max with a fresh default build, test_deepseek41_metal, test_deepseek41_graph --vision-routing, and full Metal test runs with DS4_METAL_DISABLE_PRE_M5_ROUTER_TRANSFORM_FINALIZE_FUSION=1 and DS4_METAL_MATH_SAFE=1. No model required. Pre-M5 validation and model-level speed/SSD checks remain outstanding.

(cherry picked from commit 7f5f1b77ae)
 2026-09-19 17:00:45 +02:00
Salvatore Sanfilippo d9bc45dbda Add fresh official V4.1 Flash continuations and router quality comparisons 2026-09-19 17:00:38 +02:00
Tarjei Mandt 0ac42d6fe0 Tidy the DeepSeek V4.1 comments and switches 2026-09-20 00:50:13 +10:00
Tarjei Mandt 8387625991 Drop the DeepSeek V4.1 tensor-parallel gate trace 2026-09-20 00:50:13 +10:00
Tarjei Mandt 5432223d60 Keep the resident Q4_K expert tiles to the DeepSeek V4.1 shape 2026-09-20 00:50:13 +10:00
Tarjei Mandt 7d89c59c4d Pick the DeepSeek V4.1 index kernel under test through the GPU test flags 2026-09-20 00:50:13 +10:00
Tarjei Mandt c6f4086f6d Fold the DeepSeek V4.1 RoPE quantize test into the Metal test 2026-09-20 00:50:13 +10:00
Tarjei Mandt ed43e1a56b Shorten the DeepSeek V4.1 CUDA stubs 2026-09-20 00:50:13 +10:00
Tarjei Mandt c4b66b9032 Remove the unused DeepSeek V4.1 split HC weighted sum 2026-09-20 00:50:13 +10:00
Tarjei Mandt 66dff86949 Remove the DeepSeek V4.1 exact rows attention path 2026-09-20 00:50:12 +10:00
Tarjei Mandt 258ad331dc Drop the embedded DSpark stages from the DeepSeek V4.1 engine 2026-09-20 00:50:12 +10:00
Tarjei Mandt 4fe6a78029 Guard the DeepSeek V4.1 tensor-parallel flag folds for non-Apple builds 2026-09-20 00:50:12 +10:00
GiorgioOppo 11811b96b2 perf(metal): reuse half operands for large Qwen SSD MoE prefill 
Add half-storage input and intermediate variants of the existing
simdgroup MoE gate/up/activation and down kernels. Preserve FP32 matrix
accumulation, activation arithmetic and the public float intermediate;
move the existing float-to-half operand conversions to a shared input
copy and a half shadow written by the mid epilogue.

Enable the path automatically only for M1 Max SSD streaming with
IQ2_XXS gate/up, Q2_K down and chunks of at least 8192 tokens. Keep the
original path for smaller chunks: a 2048-token experiment did not show
a reliable full-model benefit despite a faster isolated kernel pair.
No production opt-in flag is added.

Scope the half shadow to a single streaming invocation and validate the
producer tensors and dimensions before consuming it. Invalidate state
on entry/failure, separate resident and missing passes, retain scratch
buffers across submitted commands and reconvert reused input contents.
Allocate both optional buffers before encoding and fall back to float
storage on allocation failure. Account explicitly for scratch storage:
140 MiB at the tested 8192-token model shape, already covered by the
existing conservative workspace estimate.

Add a standalone frozen-source oracle/benchmark with generic weight
function constants, NT1/2/4/8, padded Q2_K rows, sparse routing, guarded
buffers and poisoned outputs. Check both FP32 outputs and the actual
half payloads against CPU binary16 conversion. Add Make targets and
extend SSD tests across the 8191/8192 policy boundary, scratch growth to
16384 and reuse, split resident/missing passes, failed reads and retries.

Validation on Apple M1 Max with Qwen3.8-Flash-Next-Q2:
- Build, full Metal kernel suite, CPU memory tests and SSD tests passed.
- Standalone oracle: 31 fixtures, 1015326720 float comparisons and
  100325376 half comparisons, all exact against the independent reference.
- Generic MoE pair at 8192 tokens: 322.620 -> 294.798 ms (+9.44%
  throughput), including conversion and dual writes; nine ABBA/BAAB rounds.
- Complete model, SSD cache 1024 experts, context 16384, MTP off,
  8192-token prompt plus 128 decode evaluations, four runs per arm:
  prefill 181.993 -> 198.252 t/s (+8.93% aggregate; +5.84% comparing
  per-run medians). All four prefill pairs improve; a slow baseline run
  raises the aggregate. A stable 10% overall gain is not demonstrated.
- Decode at chunk 8192 is inconclusive (+0.34% aggregate).
- All timed/validation replays preserve complete vocabulary logits bit
  for bit. Final Metal-debug replays separately validate the retained
  2048 fallback and 8192 optimized path, each comparing 32033280 logits.

The recorded 2048 optimized-path experiment is not enabled by this
commit. Only the automatic threshold/comment changed after measurement;
the retained 8192 execution path is unchanged. Other Apple chips keep
the previous dispatch and are unbenchmarked; CUDA/ROCm are not modified.
 2026-09-19 11:04:08 +02:00
GiorgioOppo 9bff1ca4e1 perf(metal): reuse Qwen decode activations across four Q8 output rows 
Add a four-row specialization of the existing Q8_0 matvec, preserving its
K traversal, eight-term dot product, SIMD-group count and reduction tree.
Select it for single-row Qwen projections on M1 Max with NSG=4, outside
TP2, for the measured 2560->6144/10240/12288 and 6144->2560 shapes.
Complete four-row tiles avoid out-of-bounds weight reads; vocabulary,
other shapes and other devices retain their existing dispatch.

Extend the frozen e37f185 oracle to compare two- and four-row variants
under default/safe math, NSG variants, offset views, guarded outputs and
explicitly padded tail weights. Add a public-dispatch regression covering
the four model shapes plus non-four-row and NSG fallback cases.

Validation on Apple M1 Max:
- 1024 frozen-reference fixtures, 571136 values, zero differing bits.
- Full Qwen Metal kernel suite and public-dispatch checks passed with
  Metal API validation enabled.
- Nine alternating GPU benchmark rounds with 32 dispatches per sample
  show median paired kernel throughput gains of 11.6-28.7%.
- Combined candidate/full-model replay retains bit-exact logits through
  8192 prompt tokens and 128 decode evaluations.

The isolated kernel speedup is not a claim of equivalent SSD decode
throughput: the combined 8192-chunk experiment measures only +0.34%
aggregate decode throughput, within substantial run-to-run variability.
 2026-09-19 11:03:41 +02:00
Tarjei Mandt b3d4a267c0 Open only the DSpark shards for a drafter-only DeepSeek V4.1 conversion 2026-09-19 09:46:34 +10:00
Tarjei Mandt c9bfdcc951 Release DeepSeek V4.1 tensor-parallel decode gates inside the command buffer 2026-09-19 09:46:34 +10:00
Salvatore Sanfilippo aafc65b444 Bound loopback packet size in the tiny-buffer TCP regression 
Negotiate 512-byte TCP segments on Linux before connecting. Otherwise shrinking loopback sockets to 1 KiB receive buffers leaves 32 KiB segments, making zero-window probing dominate large test transfers and sometimes trip the idle timeout. Retain the payload sizes, bidirectional backpressure, descriptor-mode checks and timeout/EOF cases. Linux and macOS transport tests pass; production transport is unchanged.
 2026-09-18 23:27:51 +02:00
Salvatore Sanfilippo 077a25760e Keep decode-sized n-gram concurrency on Apple dispatch workers 
Retain the measured Metal win without imposing per-token pthread creation on CUDA and ROCm. Tables remain disk-only. Warm M5 Max SSD decode improved from 12.62 to 13.39 tokens/s with byte-identical output; row and error tests also pass under ASAN/UBSAN.
 2026-09-18 23:27:51 +02:00
Jake Maness 6c00e2d535 Let V4.1 single-token decode use the concurrent Engram reader 
ds4_engram_read_batch() already fans row reads across ENGRAM_READERS
workers, with the comment "Fixed concurrency hides random-read latency
without caching the table" -- but only above a 256-row threshold.

ds41_graph_step() reads DS4_ENGRAM_COLS rows per table (48 rows for the
two tables) through the serial ds4_engram_read(), so ordinary decode
never reached that path. Those 48 preads run in series, and they run
before ds4_gpu_begin_commands(), so the GPU is idle for all of them.
Only bulk prefill ever cleared 256.

This routes the decode read through the batched entry point and lowers
the threshold to 8 so decode-sized batches engage the existing workers.
Same rows, same order, same values -- only the fetch is concurrent.

Measured on a Mac Studio M3 Ultra (80 GPU cores, 256 GB, macOS 26.6.2),
DeepSeek-V4.1-Flash-Q2, resident (no --ssd-streaming), --ctx 131072,
temperature 0, 200 generated tokens, this branch vs its merge-base:

  main            19.18, 19.14 tok/s                    median 19.16
  main + patch    20.42, 21.00, 21.11, 21.18 tok/s      median 21.06

  +9.9%, about 4.7 ms/token. Generated text was byte-identical between
  the two builds. (The 20.42 was a cold first run after rebuild.)

The underlying cost is storage latency, not bandwidth -- a standalone
pread benchmark against the same GGUF, 48 random 264-byte rows with
F_NOCACHE and readahead disabled, 150 batches:

  serial    5.177 ms/batch  (107.9 us/row)
  parallel  0.614 ms/batch  ( 12.8 us/row)

The saving is roughly constant per token, so it is worth more on a
faster tree: on ivanfioravanti's ds41f-optimizations branch, where this
was originally profiled, the same change measured 26.8 -> 30.5 tok/s
(Engram wall 5.55 -> 1.12 ms/token) over three alternating pairs, also
byte-identical.

Credit where it is due:

  - ivanfioravanti, whose V4.1 optimization work is where this was
    profiled. His ds41f-dspark branch already carries a better version
    of this fix, ds41_engram_parallel(), which reads both tables in a
    single dispatch -- but gated to speculative sessions, so ordinary
    decode still falls through to the serial loop there too. Ungating
    that is probably the better long-term shape of this change.

  - The oMLX / drowzeys V4.1-on-Mac work, which is what pointed at
    Engram as a decode bottleneck at all. Their approach was a RAM
    cache of hot rows; measured row reuse here was only 22.5% over a
    600-token coding generation, worth roughly +1 tok/s, so the win
    turned out to be concurrency rather than caching. Wrong idea,
    right place to look.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit 92a61612a2)
 2026-09-18 23:27:51 +02:00
Salvatore Sanfilippo 5994f062de Validate image positions before token-prefix slot reuse 
Keep fingerprint-only matching for authenticated visible histories that rebase spans, but require exact historical image positions for token and rewind reuse. Reproduce the old false hit with a synthetic checkpoint. Fix the live eviction oracle to include all competing slots. Four-slot Qwen and GLM vision/tool/eviction and concurrent-request tests pass. Follow-up to #765.
 2026-09-18 23:27:51 +02:00
antirez 4067320f06 Merge remote-tracking branch 'pr-review/765' 2026-09-18 23:27:47 +02:00
Salvatore Sanfilippo 25334fef16 Test numeric reasoning strings at the API boundary 
Cover 0 through 100, malformed values, overflow-sized input and unchanged named controls. Follow-up to #1074.
 2026-09-18 23:27:31 +02:00
GiorgioOppo 89986c3c89 fix(qwen): retain prefill attention arithmetic across Metal partitions 
A Qwen prefill crossing the 512-block sparse-attention budget leaves three
internal dense rows at absolute positions 2048 through 2050. With chunk 128
or chunk 2048, these rows previously selected the short-tail decode kernel
because their subrange contained fewer than nine rows. A single-pass
prefill selected matrix attention for the same rows. The different
rounding introduced drift despite identical input tokens and queries.

Carry the parent prefill policy through both dense and sparse attention
subranges. Add a Metal prefill entry point backed by the existing shared
implementation so internal short partitions retain matrix attention when
the parent prefill has more than eight rows. Preserve the existing shape
guards and scalar diagnostic override. Genuine short prefill tails,
decode, speculative verification and non-Apple dispatch retain their
previous policies; shader arithmetic and cache storage are unchanged.

Add regression coverage comparing a whole 2051-row attention reference
with the three-row dense partition, and sparse partitions of one through
eight rows. Require bit-identical outputs, bounded tensor views and intact
output guards. Include the previous scalar dispatch as a negative control
and verify that the diagnostic override still selects its arithmetic.

Validation on Apple M1 Max, 32 GiB, Qwen3.8-Flash-Next-Q2, SSD streaming:
- Production and affected test executables build successfully.
- Complete Qwen Metal kernel suite passes with MTL_DEBUG_LAYER=1.
- Memory/workspace, generation and MTP prefill suites pass. Generation
  includes a genuine three-token tail and the legacy prefill reference;
  MTP checks use the suite's existing numerical tolerances.
- For a 5760-token prompt and 128 teacher-forced decode steps, requested
  chunks 128, 2048 and 8192 each match the original single-pass reference
  bit for bit at all 129 full-vocabulary frontiers: 96,099,840 logit-pair
  comparisons, with MTP disabled and a 1024-expert cache.
- The original chunked prefill had maximum absolute logit error 0.597853661
  against that single pass; the candidate has zero on this fixture.
- Per-layer traces locate the first divergence at layer 3 attention output
  for token 2048, after identical inputs and queries. Traced final captures
  reproduce the uninstrumented captures exactly.
- git diff --check passes.

This targets internal attention partitions rather than promising equality
for arbitrary chunk sizes: genuine small tails and very short sparse
indexer scoring retain their existing policies. The single-pass comparison
is a numerical consistency check, not a model-quality or throughput claim.

Refs: https://github.com/antirez/ds4/pull/1056#issuecomment-5731630396
 2026-09-18 17:38:43 +02:00
Tarjei Mandt af7c02b594 Select DeepSeek V4.1 index candidates with a radix select on Metal 2026-09-18 23:32:47 +10:00
Tarjei Mandt 4b1b1519ed Score DeepSeek V4.1 index keys with one thread per key on Metal 2026-09-18 23:32:47 +10:00
GiorgioOppo 04c08672b5 fix(metal): preserve Q8 dot arithmetic to eliminate late decode drift 
The paired char2 quant loop changed Q8 GEMV results under Metal's
standard fast-math compilation despite preserving the apparent source
accumulation order. Relative to e37f185, the Qwen long-prompt replay
amplified these differences to a maximum logit error of 0.513016701 at
decode step 65, with another increase at step 106.

Restore the original unrolled eight-term qs[i] dot loop while keeping
packed_float4 activation loads and the optimized two-stage reduction.
Buffer layouts, dispatch geometry, quantization formats, and cache policy
remain unchanged. HC and MoE ablations reproduced the unfixed outputs
exactly; isolating the paired quant loop removed the drift. CUDA and ROCm
are unchanged.

Add a complete Metal Q8 GEMV regression oracle frozen independently from
e37f185. Exercise the dot product, scale application, K traversal,
reduction, output tails, scalar-aligned views, varied scales, outliers,
cancellation, and signed zeros. Verify input immutability, output guards,
and independent CPU dot samples. Cover NSG 1/2/4/8 with default and safe
math. Add make targets for execution and compilation-only checks, plus
binary cleanup and ignore rules.

Validation on Apple M1 Max:
- Full Q8 oracle: 480 fixtures and 35,712 outputs match reference bits.
- Negative control using the parent shader: 1,549 output bit differences,
  confirming that the new regression test detects the original defect.
- Existing Q8 reduction oracle: 3,072 fixtures / 4,096 output rows pass.
- Complete Qwen kernel suite, production build, and diff checks pass.
- Three full-vocabulary Qwen replays with 29, 575, and 5,760 prompt tokens
  and 128 teacher-forced decode steps each match e37f185 byte for byte:
  387 frontiers / 96,099,840 logits, including the former late spikes.
  Configuration: Qwen3.8-Flash-Next-Q2, Metal SSD streaming, 1,024 cached
  experts, context 8,192, chunks 128/2,048, MTP off, API validation on.

These results establish parity for the tested M1 Max fixtures and model
trajectories. Kernel timing was variable; no end-to-end throughput claim
is made. Temporary reports and captures are not included in this commit.
 2026-09-18 15:29:17 +02:00
GiorgioOppo 2d6a207bbb fix(qwen): preserve reference prefill arithmetic with explicit projection phases 
Short prompt chunks and tails were selecting kernels intended for independent
decode batches. Float-operand F16/Q8 tiles and split-K FP32 reductions changed
prefill arithmetic and introduced substantial logit drift against e37f185,
including different results for the same prompt at different chunk sizes.

Track prefill/decode projection policy explicitly on each Qwen graph:
- Route contiguous prompt chunks, including single-token tails, through the
  reference projection paths and an unsplit Metal dense GEMM entry point.
- Preserve existing optimized projection dispatch for ordinary decode, native
  session batches, MTP verification, predictor cache updates and draft chains.
- Scope and restore the phase on success and failure; mark copied batch row
  graphs as decode rather than inferring phase from row count or all_rows.
- Keep CUDA dispatch and the separate vision encoder unchanged. The new CUDA
  prefill entry point delegates to the existing dense function.

Add --prefill-reference to the real-model generation regression test. Compare
normal one-shot/session dispatch against legacy unsplit direct prefill, then
restore automatic decode on both sides. Cover 29-token short prompts,
575-token prompts with a 63-token tail, and two complete 128-token chunks.
Reject inherited reference overrides that could hide a dispatch regression.
Update the private MTP oracle calls for the graph-aware projection helpers.

Validation on Apple M1 Max, 32 GiB, Metal SSD streaming:
- ds4/server and relevant test builds pass; memory planner, Metal kernel,
  generation/reference and MTP cache/continuation suites pass.
- All 27 generation/reference vocabulary frontiers are bit-exact. The same
  test compiled against original 0024ca0 fails at the first frontier, as
  expected, confirming that it detects the prefill regression.
- Five configurations compared with e37f185 cover 85 frontiers and 21,107,200
  FP32 logits, using identical prompts and 16 teacher-forced continuation IDs.
  Worst prefill error falls from 0.48730135 to 1.90734863e-6; worst decode
  error falls from 1.37674665 to 0.000828742981. No argmax disagreements.
  Residual numerical differences remain; this is not bit-exact parity with
  e37f185 or a model-quality evaluation.
- Patched 5,760-token captures are bit-exact across chunk sizes 128 and 2048
  at all 17 frontiers, including the subsequent decode steps.
- Sixteen alternating SSD timing runs against 0024ca0 (four per variant per
  workload) show median prefill changes of -0.71% at 575/chunk128 and +2.54%
  at 5760/chunk2048. Decode medians change +1.59% and +8.12%, respectively;
  short 16-step samples and overlapping ranges do not establish a general
  speedup or isolate decode kernel performance.

CUDA/ROCm GPU execution and resident-mode validation were not performed.
Generated reports, binary captures and temporary logs are excluded.
 2026-09-18 14:15:14 +02:00
GiorgioOppo 0024ca0803 fix(qwen): unify prefill chunk resolution and test generation parity 
An explicit prefill chunk at or above the context size previously fell back
to DS4_QWEN4_PREFILL_CHUNK or the default in one-shot and session allocation,
while the memory estimator clamped it to the context. This could change the
actual prefill boundaries and make execution disagree with memory admission.

Use one resolver for one-shot generation, session allocation, memory
estimation and SSD admission. Explicit requests always take precedence over
the environment and clamp to the context size. Preserve existing defaults
when no explicit chunk is supplied.

Add CPU coverage for chunk precedence and context-boundary cases. Add a
real-model Metal regression test that checks public one-shot token emission
and progress against sessions, then compares every vocabulary logit between
an independently replayed direct graph and the session through eight decode
steps. Cover conflicting environment settings, oversized requests and long
prompts with configurable token counts; bound the SSD expert cache.

Document the test and the upstream 8db1d1d one-shot behavior that ignores the
CLI prefill chunk. Cross-revision comparisons must also set the same
DS4_QWEN4_PREFILL_CHUNK value to align effective chunk boundaries.

Validation on Apple M1 Max, 32 GiB:
- Built ds4, ds4-server and both affected tests without warnings.
- CPU memory/chunk-resolution tests passed; git diff --check passed.
- Five Qwen Q2 SSD cases passed with Metal API validation, including 5,760
  prompt tokens at chunks 128 and 2,048 and an oversized explicit request.
- All public greedy tokens matched; 45 complete vocabulary frontiers were
  bit-exact between direct replay and sessions (11,174,400 float comparisons).
- The oversized-request regression test fails against f47db4f with the
  observed first frontier 32 rather than the expected 247, and passes here.

These checks establish same-build parity at identical chunk boundaries.
They do not establish main-versus-branch resident parity on the tester's
64 GiB machine, or token invariance across different chunk sizes.
 2026-09-18 13:22:40 +02:00
Alessandro Checco ecf7a93148 Enhance reasoning effort parsing logic 
With this fix the server with respond to e.g.
"reasoning_effort": "50"

so it becomes compatible with pi.dev that doens't accept numerical value (and cannot send think_level by default)

(cherry picked from commit de418c0861)
 2026-09-18 09:54:45 +02:00
Spencer Gilbert 63bbba4ba7 Advertise the loaded GLM model id in /v1/models 
(cherry picked from commit 28c06c88e0)
 2026-09-18 09:54:41 +02:00
Mark Shank 064bc229c6 Fix Metal object header dependencies 
(cherry picked from commit a32a030e1b)
 2026-09-18 09:54:38 +02:00
Salvatore Sanfilippo 62867e913a Make JPEG regression failure diagnostics safe 
Always compute the decoded fingerprint before checking dimensions, so the error path cannot print an uninitialized buffer. Follow-up to #981.
 2026-09-18 09:54:34 +02:00
korale77 ac43e900d8 Fix progressive JPEG AC refinement ZRL handling in Iris decoder 
jpeg_prog_decode_ac_refine treated ZRL (run=15, size=0) as run=16 and
reused the skip loop, which kept consuming correction bits for non-zero
coefficients after the 16th zero-history coefficient. Those bits belong
after the next Huffman symbol (ITU T.81 G.1.2.3, libjpeg
decode_mcu_AC_refine), so the bitstream desynchronised: many progressive
JPEGs, including plain `cjpeg -progressive` output, failed to load or
decoded with wrong pixels.

Skip `run` zeros while refining non-zeros on the way, write the new value
(0 for ZRL) at the next zero, and stop, matching libjpeg's structure.

Add tests/test_image_decode with two cjpeg-generated fixtures that fail
on the old decoder (wrong pixels for grayscale, NULL for 4:2:0) and are
bit-exact with djpeg after the fix. Wire it into `make test`.

(cherry picked from commit 50ed0a4c73)
 2026-09-18 09:53:23 +02:00
Salvatore Sanfilippo 39bc628cc3 Cover split Qwen reasoning markers in both streaming APIs 
Verify tool-enabled answers are emitted before generation finishes for OpenAI and Anthropic clients. Keep the existing DeepSeek repeated-reasoning regressions. Follow-up to #1077.
 2026-09-18 09:52:20 +02:00
Jeff d31089d5ad Fix Qwen tool content streaming 
(cherry picked from commit 77a39ebf0ea81ff2f7767af1bf9631890215482e)
 2026-09-18 09:50:23 +02:00
ayuan1129 fff391e6c7 GLM: render the tool call block inline, as the chat template does 
append_glm_tool_calls_text() wrapped the reconstructed block in a newline on
each side.  GLM's tokenizer.chat_template does not: it emits

  <|assistant|><think></think><tool_call>name<arg_key>k</arg_key>...</tool_call>

with no surrounding whitespace in any variant (empty content, non-empty
content, reasoning present, followed by an observation or by a user turn).
The block is only wrapped when it has to be rebuilt from parsed arguments;
when the raw sampled bytes are available they are emitted verbatim, so the two
paths disagreed.

That disagreement is invisible while a conversation stays in one process: the
live tool turn feeds the raw sampled bytes into the prompt, so the running KV
prefix still lines up.  It surfaces when a client replays the history, because
the replay goes through the reconstruction path and picks up the extra
newlines.  Any key or checkpoint derived from the visible transcript then fails
to match the replay: in the restart smoke test a conversation that should
re-prefill ~14 tokens re-prefilled 3800 of 11992, because the visible-key
checkpoint it wrote was three bytes off.

Cold prefill is affected too, and independently of any cache: a replayed GLM
tool conversation was being rendered in a shape the model never trained on.

Drop the two newlines.  test_glm_tool_calls_render_matches_template asserts the
template-true block and that the raw and reconstructed renders are byte
identical; test_render_glm_preserves_reasoning_with_tools had pinned the old
bytes and is corrected.

Reproduced with the restart-restore smoke test on GLM-5.3-Flash-Q2:
  before: FAIL, restored 8192 of 11992 tokens (3800 re-prefilled)
  after:  OK,   restored 11998 of 12012 tokens (14 re-prefilled)

(cherry picked from commit d186421d96)
 2026-09-18 09:49:08 +02:00
Salvatore Sanfilippo 9775e26ccd Test GLM pooled decode padding against an independent reference 
Cover all pool tail lengths after 2048 selected rows. Repeat the split Metal decode with padded IDs and compare its value projection to a double-precision attention reference. Follow-up to #1051.
 2026-09-18 09:48:57 +02:00
Mark Shank 582d809fd3 Fix GLM-5.3 Metal decode masking for padded selections 
Track the pooled selection’s guaranteed prefix instead of promising that all selected rows are valid. Padded selections use the existing bounds-checked kernel.

Adapted from 101d20c23f3a27456bdc61e0da35e2bea97eb04e.

Validation: make -j1 and git diff --check passed on macOS M3 Ultra (Metal build). No model was loaded; GPU runtime and performance tests were not rerun.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
(cherry picked from commit 4e166b1fb4)
 2026-09-18 09:45:20 +02:00
Tarjei Mandt fd7091916a Score DeepSeek V4.1 verify rows with the tiled indexer kernel 2026-09-18 12:53:34 +10:00
Tarjei Mandt 1c6920666c Stop proposing DeepSeek V4.1 drafts while they lose to plain decoding 2026-09-18 12:53:34 +10:00
Tarjei Mandt 0b2d6e28cb Write only the DeepSeek V4.1 text tensors from the converter 2026-09-18 12:53:34 +10:00
Mikhail Korobov 0edbd9cd64 Merge upstream/main: vision live-state accounting in the reuse probe 
Upstream 81a7658 ("Preserve live vision context across server tool calls")
rewrote the live-continuation helpers this branch replaced with
slot_probe_reuse(): image markers carry request-local nonces and are
normalized for visible replay, image tokens are spliced into continuations
instead of being tokenized as text, and image appends are validated by
fingerprint and row count rather than by exact span equality.  The merge
keeps the probe architecture and folds the new conditions into it, so slot
routing and the execution ladder still cannot disagree:

- The vision gate is fingerprint-prefix (new
  ds4_session_vision_fingerprint_prefix_matches): historical images must
  match by row count and fingerprint, positions may shift because
  build_live_prompt_suffix() rebases them, and appended images are allowed.
  An exact-state gate rejected valid appends, and a position-sensitive one
  rejected replays that omit hidden reasoning.
- Every live tier materializes through build_live_prompt_suffix(), so new
  images are spliced into the continuation and a failed rebase is a safe
  miss rather than a malformed prompt.
- responses-visible and thinking-visible compare in visible_prompt_key
  space, so a replay with fresh marker nonces still binds while a changed or
  shifted image does not.
- Tool-output-only continuations keep upstream's contract: if the selected
  live tier cannot be materialized, the request is answered with 409
  (live_continuation_unavailable) instead of being generated from a
  context-free prompt built from tool outputs alone.
- Single-session mode refreshes the slot's rendered-text cache like the
  batched worker does, so the memory-text tier is probe-authoritative on
  both worker paths; the redundant ladder fallback and the four superseded
  upstream helpers are removed.

Tests: test_slot_probe_vision_tiers covers image appends, pixel mismatches
and marker-nonce visible replay at the probe/routing level;
test_live_continuation_contract pins the 409 predicate;
tests/test_session_state.c gains fingerprint-prefix cases; and the new live
suite tests/test_server_vision_live_state.py checks slot routing, eviction
victims and the tool-result-only 409 against a real vision model.

Verified on DeepSeek V4 Flash Vision Experimental (M5 Max): ./ds4_test
--server, ds4_agent_test, test-session-state, live test_server_vision_cache.py
(1- and 4-session), test_server_story.py, test_server_batching.py and the new
vision live-state suite all pass.
 2026-09-18 04:07:30 +05:00
Tarjei Mandt aea1a0bae9 Document DeepSeek V4.1 DSpark support files, two-rank drafting and idle-slot prefill 2026-09-17 14:27:09 +10:00
Tarjei Mandt 5c2dac7e37 End the tensor parallelism bind line 2026-09-17 14:27:09 +10:00
Tarjei Mandt 29ce271717 Give DeepSeek V4.1 prefill rows their own second Engram table 2026-09-17 14:27:09 +10:00
Tarjei Mandt b7bcccdea2 Stop DeepSeek V4.1 drafts before the end-of-sequence token 2026-09-17 12:13:32 +10:00
Tarjei Mandt 7da9535f2d Sync idle server slots in the engine's prefill chunk 2026-09-17 12:13:32 +10:00
Tarjei Mandt bd85ba6aca Keep the row kernels on the exact DeepSeek V4.1 rows path 2026-09-17 11:20:10 +10:00
Tarjei Mandt 1e3bcf8d87 Release DeepSeek V4.1 verify-block gates through the spin kernel 2026-09-17 10:38:31 +10:00
Tarjei Mandt c674b198d4 Verify DeepSeek V4.1 DSpark blocks across two tensor-parallel nodes 2026-09-17 10:38:30 +10:00
Tarjei Mandt a6b50ff6c4 Multiply the DeepSeek V4.1 shared expert and attention output rows through simdgroup matrices 2026-09-17 08:54:56 +10:00
Tarjei Mandt 262b4a66e3 Multiply DeepSeek V4.1 verify rows through simdgroup matrices 2026-09-17 08:54:56 +10:00
Tarjei Mandt e49643fff9 Run the DeepSeek V4.1 shared expert beside the routed experts on Metal 2026-09-17 08:54:56 +10:00
Tarjei Mandt 867465cc66 Fold the DeepSeek V4.1 tensor-parallel gate sums into the hyper-connection expansion 2026-09-17 08:54:56 +10:00
Tarjei Mandt a13b513e63 Split the DeepSeek V4.1 shared expert across the tensor-parallel ranks 2026-09-17 08:54:56 +10:00
Tarjei Mandt 8d73a0effa Balance the DeepSeek V4.1 tensor-parallel decode experts by list position 2026-09-17 08:54:56 +10:00
Tarjei Mandt a3f6f313a6 Port the decode-control regression test and fix its router batch call 2026-09-17 08:54:55 +10:00
Tarjei Mandt 1923131d4c Match the rope contraction of the DeepSeek V4.1 fused rope-quantize store 2026-09-17 08:54:55 +10:00
Tarjei Mandt dc81436da2 Bind the owned expert range once when a DeepSeek V4.1 batch falls back to one-token MoE 2026-09-17 08:54:55 +10:00
Tarjei Mandt 9f3892d2fe Write DeepSeek V4.1 tensor-parallel partials into the slab slot and batch the kslice rows 2026-09-17 08:54:55 +10:00
Tarjei Mandt 4fbbc4a45b Encode the DeepSeek V4.1 vocabulary head before the last decode drain and start the wide sweep at 3072 tokens 2026-09-17 08:54:55 +10:00
Tarjei Mandt d8d1523b33 Speed up DeepSeek V4.1 Metal decode and run DSpark verify rows on its decode kernels 2026-09-17 08:54:55 +10:00
Tarjei Mandt a60edc61f5 Write the DeepSeek V4.1 DSpark stages to a separate support GGUF 2026-09-17 08:54:55 +10:00
Tarjei Mandt c016ea7730 Keep the released MXFP4 experts in the DeepSeek V4.1 converter 2026-09-17 08:54:55 +10:00
Tarjei Mandt 3b7f8f2242 Select every key for short DeepSeek V4.1 index rows 2026-09-17 08:54:55 +10:00
Tarjei Mandt d744ebb9e7 Run Q4_K DeepSeek V4.1 experts on the compact resident tiles 2026-09-17 08:54:55 +10:00
Tarjei Mandt ce5a812fd3 Select DeepSeek V4.1 candidate blocks in batches 2026-09-17 08:54:55 +10:00
Tarjei Mandt ccbf2c083e Run the DeepSeek V4.1 decoder suffix from 2541 tokens 2026-09-17 08:54:55 +10:00
Tarjei Mandt e768396179 Skip DeepSeek V4.1 candidate selection while every block is kept 2026-09-17 08:54:55 +10:00
Tarjei Mandt 10118742ae Publish DeepSeek V4.1 ratio-1 keys in batches on Metal 2026-09-17 08:54:54 +10:00
Tarjei Mandt edceb7ab7f Select DeepSeek V4.1 experts in one dispatch 2026-09-17 08:54:54 +10:00
Tarjei Mandt d95f8b6101 Round DeepSeek V4.1 activations inside the producing kernels 2026-09-17 08:54:54 +10:00
Tarjei Mandt 2a281b0802 Queue DeepSeek V4.1 decode layers on a single node 2026-09-17 08:54:54 +10:00
GiorgioOppo f47db4fb83 Merge antirez/main at 8db1d1d into qwen-kernel-opt 
Integrate upstream Qwen CUDA support, Metal multi-session decode and
speculative batching, shared prefill workspaces and recurrent-state pools,
grouped resident MoE kernels, dense/Q8 batch projections, and server state
handling. Retain the Qwen SSD-streaming and Metal kernel optimizations from
this branch.

Resolve the host/shader MoE argument layout by retaining the compact SSD
slot masks and appending the upstream grouped-list fields. Update layout
assertions and direct gate/up and down regression fixtures to the resulting
72-byte ABI. Keep SSD address resources and in-flight buffer lifetimes in
the common dispatch path alongside the upstream resident row cache.

Keep SSD sessions on the ordered ordinary/speculative fallback with private
streaming workspaces. Resident sessions use the new upstream shared arena.
Preserve the existing SSD prefill tile threshold and guard Metal-only SSD,
cache-only attention preparation, and MTP projection calls from CUDA builds.

Reconcile retained MTP tail rows and snapshot ownership with native batching:
prepare pending predictor boundaries, copy per-session tails and rejection
snapshots, maintain tail metadata after commit, and complete the historical
cache row after an accepted two-token verification when next-draft creation
is skipped. Preserve checkpoint payload compatibility within this branch.

Retain both upstream and branch-specific build/test targets. Extend the SSD
failure/recovery test to reject native resident dispatch explicitly and
verify exact logits across ordered ordinary and speculative batches. Clarify
resident versus SSD batching in the Qwen and server documentation.

Validation on Apple M1 Max, 32 GiB:
- Build ds4, ds4-server, and all selected test binaries without warnings.
- Pass Qwen memory planning and compact-MoE CPU tests with ASan/UBSan.
- Pass the complete Qwen kernel suite with Metal API validation, including
  upstream grouped MoE, batch projections, and attention paths.
- Pass direct MoE mid (180 fixtures, 1,119,600 comparisons), MoE down
  (1,650 fixtures, 19,955,180 comparisons), HC (92 fixtures), and Q8
  reduction (3,072 cases) regression suites.
- Pass SSD expert cache/read/retry tests and MoE specialization tests.
- Pass real Qwen Q2 SSD n-gram failure/recovery, ordered batch parity, MTP
  prefill, snapshot/payload replay, and continuation checks.
- Pass a 50-token greedy SSD generation for 'narrami la storia di roma'.

CUDA execution and full resident-model multi-session integration were not
run on this 32 GiB Mac. This merge does not make new performance claims.
 2026-09-16 19:51:07 +02:00
Salvatore Sanfilippo 8db1d1d155 Document Qwen Metal batching and its quality checks 2026-09-16 15:17:26 +02:00
Salvatore Sanfilippo ab38bc6b7a Serve Qwen batched MTP with safe session state and output limits 
Keep exact sampled requests on ordinary batches. Clear stale verification snapshots, honor logical context limits, and reset speculative timing estimates when batch width changes. Exercise forced rejection, mixed decode, snapshot restore, disk failures and arena replacement.
 2026-09-16 15:17:25 +02:00
Salvatore Sanfilippo d1620ba061 Reject incomplete concurrency benchmarks and oversized requests 2026-09-16 15:02:09 +02:00
Salvatore Sanfilippo 8e608a66e0 Trim Qwen batching controls and benchmark documentation 2026-09-16 14:46:28 +02:00
davidandClaude Fable 5.1 9ba464d529 Price the batched cycle kinds after eight drafted cycles, not sixteen 
Sixteen speculative cycles before the first plain one made a short
generation on prose pay the losing kind for most of its length (C=16 at
32 steps: 174 tok/s against ~195 once the policy has settled).  Eight
cycles give p over 128 draft rows at sixteen streams, enough to price.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:18:48 +02:00
davidandClaude Fable 5.1 307778657f Decide per cycle whether a speculative batch drafts 
A draft row's cost is not per row: the seventeenth row of a batched cycle
opens a second tile of every dense projection, after which rows are
nearly free (16 rows 80 ms, 23 rows 126, 32 rows 149), so gating streams
by their draft's confidence bought nothing (prose 184.2 vs 183.9 at a
margin of 2, code 210.7 vs 213.6).  A batch speculates for everyone or
for nobody, and it pays when (1 + p) c_plain > c_spec.

The engine keeps p and both cycle times as running averages (the same
shape as the session's depth policy), starts by pricing both kinds and
then probes the other kind for two cycles, every thirty-two while plain
and every hundred and twenty-eight while speculating.  DS4_QWEN4_SPEC_TRACE
prints the state per cycle.

C=16 prose 184 -> 195 tok/s (plain ~200, mostly plain cycles now), code
213.6 -> 214.3; C=8 prose 131 plain -> 158, code 131 -> 179 (the drafts
fill the one tile a C=8 batch has).  C=2 verifies lossless.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:18:47 +02:00
antirez 548cdaf897 Bound speculative batches and invalidate failed predictor state 2026-09-16 14:18:17 +02:00
davidandClaude Fable 5.1 e2cba0b7be Batch the predictor layer across a speculative batch's sessions 
The batched cycle drafted with one predictor pass per session, sixteen
passes of the nextn layer over one or two rows each, every one reading
the layer's weights again.  One pass now covers every row of the batch:
the trunk streams sit in the batch's R at the positions the trunk rows
had, so the layer's attention takes the same row table as the trunk's,
the staged inputs and residual go where the prompt priming puts them,
and the sessions' last rows are gathered for one pass of the head mixer
before the batched head and the per-session argmax.  A rejected draft's
row runs along with its parent as the token that follows; the next pass
rewrites its cache entry before anything reads it.

C=16 code 204 -> 213.5 tok/s cumulative (cycle 156 -> 148 ms), prose
176 -> 183; C=2 verifies lossless against plain greedy.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:16:26 +02:00
davidandClaude Fable 5.1 d569d05aec Run the speculative batch's recurrent and attention steps as rows 
The speculative batch ran the GDN convolution, scan and the attention of
each session through the single-session kernels, one to two rows at a
time: sixteen sessions, forty-eight layers, ~2300 short dispatches a
cycle.  The GDN rows kernels now come in a variant addressed through a
table (state, history and the verify snapshots of each session by GPU
address, its first row and its row count), which walks a session's one or
two tokens in order and writes the snapshot after the first, with the
per-token arithmetic of the verify scan; the attention rows kernels take
one entry per token, so a draft row is simply a second entry at pos + 1
that sees the first token's keys.  The trunk of a cycle with no drafts now
costs what a plain step does (89 ms against 92).

With drafts the cycle is 158 ms for 1.50 tokens per stream (152 tok/s at
C=16 on the prose corpus): every draft row costs its own MoE pass, so the
marginal row (~4.4 ms) is close to a plain token (5.6 ms) and greedy
speculation cannot pay at this width on this model unless acceptance is
high.  Lossless against plain greedy decode (harness --spec --verify).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:15:22 +02:00
davidandClaude Fable 5.1 a504136766 Speculate over a batch of Qwen3.8 sessions 
ds4_sessions_eval_batch_speculative_argmax runs one MTP cycle for a batch:
every session contributes its next token and, when the predictor left it a
draft for that token, the draft as a second row.  The dense work of the
trunk runs once over all the rows; the recurrent, n-gram and attention
steps run per session over its one or two rows with the single-session
verify kernels, snapshots included, so a session's own arithmetic is that
of the C=1 speculative cycle.  A draft stands when it is the target's
argmax for its row; a rejected one restores the row-0 snapshot.  The
predictor then runs its layer per session over the committed rows and
every session's head as one product.  The Q8 batch GEMM serves 9..32 rows
in sixteen-row halves.

The harness gains --spec (with --verify checking the committed tokens
against plain greedy decode, which speculation must reproduce).  On an M5
Max at C=16 the first cut commits 1.49 tokens per stream and cycle on the
prose corpus but the cycle costs 171 ms against 92 for a plain step: the
per-session kernels and the sequential predictor passes are the next
thing to batch.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:15:13 +02:00
antirez 467b6f05a0 Harden Qwen Metal batch integration and quality checks 2026-09-16 14:14:36 +02:00
davidandClaude Fable 5.1 9e91baccda Keep prefill chunks on the generic F16 tile 
The fp32 tile for F16 batches was meant for decode batches, where the
generic tile's half-staged activations carried rounding into the logits;
its rule also caught every prefill chunk's narrow hyper-connection
projections.  There it is slower (1305 vs 1410 tok/s at 4k) and it moved
a single stream's tokens off the reference's: with the generic tile back
for chunks past 64 rows, a C=1 stream generates the same 32 tokens as
upstream main, and prefill returns to 1408 tok/s at 4k, 1455 at 16k.
Decode batches keep the fp32 tile.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:07:24 +02:00
davidandClaude Fable 5.1 b4da720207 Count a prefill chunk's token tiles in the k-split rule 
The k-split of the dense GEMM counted only the row tiles, so a prefill
chunk's narrow projections split k thirteen ways although their 256
token tiles already filled the machine, and their partial planes grew
with the chunk: past the first context size of a run they no longer fit
the scratch allocated once, the GEMM failed mid-batch and every call fell
back to the per-token kernel, costing 8% of prefill at 8k and 16k.  The
rule now counts both tile dimensions: a decode batch splits as before, a
prefill chunk does not.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:07:24 +02:00
davidandClaude Fable 5.1 65d2eb3e15 Replace the shared session arena when a wider session finds it idle 
The shared transients arena took the shape of the first session, and a
later session needing a wider one (a larger context, hence more indexer
blocks) silently fell back to private transients, ~4-5 GiB each: sixteen
of them at the next context size of a sweep exceeded memory and the
process died without a word.  The engine now counts the sessions
borrowing the arena and, when a wider session arrives while none does,
frees the arena so that session builds the wider one; while sessions
borrow it, a wider one keeps private transients as before.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:07:24 +02:00
davidandClaude Fable 5.1 643d4cb5bc Give the batched output head the Q8 tile 
The plain batched decode sized its logits scratch to the batch width, so
the Q8 tile's 32-row padding did not fit and the output head, 248320 rows
read once per four tokens by the per-token matvec, cost 3.9 ms of a
76 ms step.  The scratch now holds the tile's rows, which the speculative
path already did, and the head takes the tile like every other Q8
projection.

C=16 plain prose 196.9 -> 203 tok/s cumulative; batched vs sequential
unchanged (same worst logit deviation, tokens agree).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:59 +02:00
davidandClaude Fable 5.1 891c0053c7 Keep verify rows on the few-row matvec 
The batch tile took every F16 projection past one row, which also caught
the two- and three-row speculative verify of a single session: its
hyper-connection down products cost three times what the few-row matvec
charged them, and that was their original path.  The tile now starts at
four rows.  C=1 with MTP on the prose corpus: 50.5 -> 55.6 tok/s (plain
51.3); C=16 is unchanged.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:38 +02:00
davidandClaude Fable 5.1 a154af774b Walk both rows of a grouped down simdgroup together 
The grouped down pass ran its two rows one after the other, each a
dependent chain of block products.  Both rows now walk their blocks in
the same loop, two independent chains per pair in flight; each chain's
own order is the per-token kernel's and the test pins it byte for bit.
+2.3% on the C=16 step (1.16x over the per-token kernels).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:38 +02:00
davidandClaude Fable 5.1 3c0cfc3108 Specialize the grouped MoE down pass on its pair count 
Same shape as the mid pass: a template on the pair count so the block
loop carries no per-pair branch, one named accumulator per pair, the
per-token kernel's pinned chain untouched and pinned byte for bit by the
test.  +1.4% on the C=16 step (1.14x over the per-token kernels).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:38 +02:00
davidandClaude Fable 5.1 54a39d7fa2 Specialize the grouped MoE mid pass on its pair count 
The grouped mid kernel tested the pair count inside its innermost loop,
three uniform branches per weight, and read each block's scales and
nibbles a byte at a time.  A pass is now a template on the pair count
(one to four), so the inner loops carry no branch, and the block's words
come in as three loads (the byte values are the same); the per-token
arithmetic and its order are untouched, the test still pins it byte for
bit against the per-token kernel.  +7% on the C=16 step of an M5 Max
(1.12x over the per-token kernels, from 1.047x).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqNkMw6d
 2026-09-16 14:06:38 +02:00
davidandClaude Fable 5.1 d36dd5d592 Allocate the k-split scratch once 
The scratch for the k-split planes grew on demand, and the command buffer
holds unretained references: a plane buffer replaced while earlier
dispatches of the same batch still pointed at it was read after its
release.  That is what made splits of sixteen and more look wrong (a
logit difference of 1.3 that never moved).  The scratch is now allocated
once, 64 MiB or the first request if larger, and a request it cannot serve
fails instead of reallocating.  Splits of 32 and 64 now measure exact; they
are no faster, so the default width stays.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:24 +02:00
davidandClaude Fable 5.1 8378039eb4 Multiply Q8 decode batches on fp32 simdgroup matrices 
At sixteen rows the per-token Q8 matvec read every matrix once per four
rows and streamed the batch's activations once per four weight rows: a
third of the step for 3.4 GB of weights.  The batch GEMM gives each
simdgroup 32 weight rows and every token: each lane stages one q8_0 block
of its row as floats (the matvec's d * q values) into the simdgroup's own
threadgroup memory, so simdgroups never wait on each other, and the rows
multiply on fp32 simdgroup matrices against token tiles read straight from
the activations.  Narrow projections split k across the grid and sum their
planes with the tile reduce.  The sums stay within 2e-6 of the matvec's;
tests pin the kernel against a double reference and the matvec.

+10% on the C=16 step of an M5 Max (paired A/B).  Every simdgroup matrix
index is a compile-time constant: with the token-tile count a runtime
value the same kernel ran twice slower.  DS4_QWEN4_NO_BATCH_MM keeps the
matvec for A/B.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:24 +02:00
davidandClaude Fable 5.1 fe457518a6 Group a decode batch's routed experts and run its shared expert densely 
Sixteen rows pick ten experts each but only about a third as many
distinct ones, and the per-row kernels read an expert once per row that
chose it.  The grouped kernels give each expert to the threadgroup of the
first (row, slot) pair in its list, which walks the expert's rows once per
four pairs with a named accumulator per pair, so every row's block walk and
chain order are the per-row kernel's and its outputs match bit for bit
(tests/test_qwen4_kernels.c pins mid and down under heavy reuse).  The
shared expert, read once per row as an extra slot, becomes dense
projections over the batch as in prefill.

Together +7% on the C=16 step of an M5 Max (paired A/B); single rows and
verify rows keep the per-row kernels.  DS4_QWEN4_MOE_NO_GROUP and
DS4_QWEN4_MOE_SHARED_DENSE_MIN keep the old paths for A/B.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:24 +02:00
davidandClaude Fable 5.1 2c5b1aab14 Keep F16 decode-batch projections on the fp32 tile 
Past eight rows the F16 projections reached the generic tile, which stages
the activations as halves; the hyper-connection up projections carried
that rounding into the logits of a batched step, flipping near-tied tokens
against sequential decode.  The k-split fp32 tile keeps float operands at
the same cost (1.01x on the C=16 step), so it takes every F16 batch.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:24 +02:00
davidandClaude Fable 5.1 4e01ef9dec Take F32 decode-batch projections through the tiled GEMM 
The router and the GDN gate projections are F32 and narrow (512 and 48
rows), so at sixteen rows the per-token kernel gives the machine a few
dozen threadgroups and reads the matrix once per four tokens.  The
k-split tile spreads them over the machine: +2% on the C=16 step (paired
A/B on an M5 Max).  Q8 weights measured the other way (the tile loses 5%)
and stay on the per-token kernel.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:24 +02:00
davidandClaude Fable 5.1 3e068bb0e7 Batch Qwen3.8 attention rows into one dispatch per stage 
The decode batch ran attention as sixteen chains of per-row dispatches:
prep, block key, score, select, expand, decode and merge for each session
in turn.  The rows kernels take the row as a grid axis and reach each
session's caches through a staged table of GPU addresses, so a layer costs
one dispatch per stage.  Every kernel keeps the single-row arithmetic, the
row's own split count included, and the single-row kernels call the same
bodies; tests/test_qwen4_kernels.c pins the rows path byte-exact against
the per-row dispatches for dense, block-completing and sparse rows on both
selector widths.

On an M5 Max at C=16, ctx 4096 the step is unchanged (1.00x): those grids
were not latency-bound after all, the attention stages take ~6.5 ms of a
110 ms step.  DS4_QWEN4_NO_BATCH_ATTN keeps the per-row path for A/B.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:23 +02:00
davidandClaude Fable 5.1 6efb642d53 Report the logit distance when a batched verify fails 
The distance gathered up to the failing step says whether the flip was a
near tie under a changed reduction order or a wrong result.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:23 +02:00
davidandClaude Fable 5.1 ba3b0c23e3 Batch Qwen3.8 Flash Next decode natively across sessions 
Qwen3.8 sessions were excluded from native session batching: each decode
step ran the whole graph once per session, so sixteen sessions cost
sixteen times the weight traffic.  The batched encoder now runs the dense
work (hyper-connection mixers, projections, MoE, output head) once over
the stacked rows and keeps only the recurrent state, the caches and the
PLE history per session, with batched GDN scan and conv kernels reading
the pooled state through per-row slots.

Session transients (prefill scratch) move to one shared arena sized for
the largest session, so each further session costs its caches only
(memory per session divided by ~3.6 at ctx 4096).  The dense GEMM gains a
split-K path for the narrow-batch regime, and the dense/MoE/mixer GEMM
thresholds, tuned for T=1 or T>=8192, are moved to serve decode batches
of 4..64 rows.  Every step keeps the tokens of the sequential path;
float reduction order changes where the batched kernels do.

C=16 at ctx 4096 on an M5 Max goes from ~50 to ~143 tok/s aggregate.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:06:23 +02:00
davidandClaude Fable 5.1 147b2632e0 Add a session-concurrency benchmark for the engine and the server 
Serving several requests at once is a different regime from the single
session walk ds4-bench measures.  session_concurrency_bench drives the
engine API directly over a concurrency x context grid and reports the
aggregate decode throughput, the per-stream loss as the width grows, the
cost a resumed prefill imposes on the decoders, and a teacher-forced check
that a batched step emits the tokens the sequential path emits.  Cells are
priced from the engine's own memory accounting against a budget sampled
at startup, because macOS swaps long before an allocation fails.
concurrency_sweep.sh runs one process per cell; serve_concurrency_bench.py
measures the HTTP server with vLLM-style metrics.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:05:35 +02:00
davidandClaude Fable 5.1 d0b74340bf Read a decode token's n-gram rows concurrently 
The PLE n-gram table is read from disk uncached, one 320-byte pread per
row, and a token needs sixteen rows.  Requests under 256 rows were read
in sequence to spare the worker dispatch, which put 1.3 ms of SSD latency
per token on the critical path before the GPU could start: 7% of a
19 ms decode step, 10% of a speculative cycle.  Any request of two rows
or more now goes through the parallel readers, capped at the row count.

C=1 plain 51.3 -> 54.5 tok/s (stage 1.31 -> 0.25 ms); C=1 MTP code
69.9 -> 78.9, prose 58.3 -> 66.0.  Logits bit-identical to the previous
build; prefill unchanged.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SxdJvkAEqcDT2o7tqKMw6d
 2026-09-16 14:04:25 +02:00
GiorgioOppo 9a8462a19f Metal: streamline Qwen Q2 down loads and shared expert row reuse 
Widen Q2_K activation, scale and quant loads with packed types while keeping
the original FP32 accumulation chain. Process the first two complete blocks
without per-lane tail checks, then process the final 128 live values only in
the lower sixteen lanes. Never form padded activation addresses for logical
K640 weights stored at physical K768.

Reuse shared Q8 activation loads across two output rows, preserving each
row's block order, four-term expression and simd_sum. Keep the existing host
dispatch and fallback for other shared formats.

Add a source-derived GPU regression against the generic row-dot kernel:
resident and separate SSD payloads, full/sparse/empty masks, T1/2/3, odd row
tails, minimum alignments, nonzero padding, default/safe math and edge values.
Validation: 1650 fixtures and 19955180 comparisons pass with Metal API
validation and zero NaN payload differences. Host Q2 dispatch tests pass
for T1/2/3, output2560/2561 and shared present/absent.

The tail/shared changes, measured against the preceding packed-load candidate,
improve isolated full-slot kernel throughput by 6.36%, 7.77% and 7.09% at
T1/T2/T3. Eight ordinary SSD requests show decode medians of 9.71 versus
9.67 t/s; no end-to-end gain is demonstrated. Output and cache work match.
 2026-09-16 08:21:04 +02:00
GiorgioOppo 13c53d922b Metal: specialize Qwen IQ2 decode and reuse shared Q8 gate/up inputs 
Export NR1 and NR2 variants from one IQ2XXS gate/up kernel body. On M1 Max,
select NR1 for the shipped single-token K2560/F640 shape with at least eight
active slots and the default four-SIMD-group policy. Preserve NR2 for small
split passes, multiple tokens and explicit generic-specialization settings.
Register both exports with the SSD address-table specialization.

For the shared Q8 expert, load each activation group once for gate and up
while keeping independent ordered accumulators and SIMD reductions. Preserve
the fallback for other shared weight types and the original slot addressing.

Add a standalone GPU test for both exports, resident/SSD addressing, compact
masks, tails, floating-point edge cases and default/safe math. Compare shared
outputs against separate row-dot projections, independently of both exports.

Validation: the current M1 Max run passes 180 fixtures and 1119600 comparisons,
including the independent shared oracle, with zero NaN payload differences.
These correctness results do not imply an end-to-end throughput gain.
 2026-09-16 08:21:04 +02:00
GiorgioOppo a37fd7f53d Metal: reuse IQ2 headers when staging MoE prefill tiles 
Handle IQ2XXS directly in qwen4_mm_stage16 instead of decoding two independent
eight-value halves. Read the shared 32-value header once through a packed
ushort4, then load aligned four-byte codebook slices for both halves.

Respect the two-byte alignment of 66-byte GGUF superblocks and avoid reading
past codebook entries. Preserve the stage8 scale expression, sign mapping,
element order and final destination conversion for each staged value.
Other quantization formats retain their existing paths.

Validation: strict/fast CPU decoding oracles and sanitizer checks from the
optimization review pass; the current Metal MoE MM specialization suite
also passes with exact mid/down output and intact padding on M1 Max.
 2026-09-16 08:21:03 +02:00
GiorgioOppo bbbc01207e Metal: tune Qwen hyper-connection dispatches for M1 Max 
HC-down's narrow 320-channel output leaves too few threadgroups in ordinary
prefill chunks. Add a 32-channel by 16-token F16 tile using 64 threads and
3072 bytes of shared memory, preserving half RHS staging and K32 MMA order.
Select it only on M1 Max for K=10240, output=320 and 9..128 tokens.

Enable the existing F16 RMS-reuse path for measured M1 Max chunks of 48..256
tokens at the shipped HC shape. Use sixteen SIMD groups for the single-token
reuse mixer without changing per-row lane mapping or reduction order.
Retain the established device and shape fallbacks.

Add a direct HC-down GPU oracle with partial tiles, bounds checks, default
and safe math; add host norm-dispatch boundary cases at 47/48 and 256/257
as well as ordinary chunk/tail sizes. Wire the test into Makefile and clean.

Validation: the direct HC suite passes all 16 pipelines and 92 fixtures on
M1 Max with Metal API validation. Existing host norm/mixer parity checks
cover the changed policies; no cross-device speedup is claimed.
 2026-09-16 08:21:03 +02:00
GiorgioOppo a50fecce6f Metal: streamline Q8 matrix-vector loads and reduction 
Use packed float4 activation loads and aligned char2 quant loads while
preserving each row's eight-operand FP32 accumulation sequence.

Give Q8 a single-use reduction helper: SIMD groups publish only their own
partials, zero padding stays in registers, and group zero performs the final
reduction. Preserve both simd_sum stages even for NSG=1, including positive
zero padding, while removing the shared-memory clear and its extra barrier.
Document the scratch lifetime requirement at the helper.

Add a standalone Metal regression using a frozen legacy reduction oracle,
with compile-only support and Makefile/clean integration. Exercise grouped
reductions, row tails and floating-point edge cases.

Validation: 3072 cases and 4096 output rows pass on M1 Max with Metal API
validation; non-NaN results match bitwise and NaN classifications match,
with zero NaN payload differences in this run.
 2026-09-16 08:21:03 +02:00
GiorgioOppo 2ba92ab234 Metal: compact streamed MoE prefill dispatches by active expert 
SSD split passes previously dispatched a full expert grid even when many
experts had no work in that pass. Copy the active original expert IDs into
each encoder's arguments and launch only those experts for mid/down, regular
and tail tiles, including the NAX variants.

Keep unfiltered routing frequencies for tile-policy decisions and preserve
original expert IDs for weights, counts, token lists and outputs. Resident
and missing passes carry immutable CPU pass-count views; GPU arguments
contain copied IDs rather than CPU pointers. Reject inconsistent pass counts
and skip empty passes while invalidating stale NAX half-mid associations.

Update host and shader layouts together, with size/offset assertions for the
2112-byte structure and the 4096-byte setBytes limit. Add a source-derived CPU
regression for mappings, schedules, split-pass lifetimes and empty-state
handling, including ASan/UBSan execution.

Validation: sanitized compact-dispatch regression and the Metal MoE MM
specialization suite pass. Existing full-model SSD comparisons retain
identical generated output and expert-cache work.
 2026-09-16 08:21:03 +02:00
GiorgioOppo 29607116c3 Qwen: skip unused indexer queries in dense attention prefixes 
Dense-prefix attention consumes every complete visible block, so its indexer
query projection and normalization have no consumer. Derive the first sparse
row from one shared, 64-bit causal-block boundary and omit only those unused
queries when an entire ordinary chunk remains dense.

Continue preparing and pooling every indexer key for later sparse attention.
Chunks crossing the boundary keep the original full query projection and
batch shape, preserving rounding. MTP query preparation and cache-only
predictor population retain their existing behavior.

Add row-by-row CPU boundary tests, including the shipped position-2051 cutoff
and unsigned limits. Extend Metal preparation tests to T=1/2/3/128, both RoPE
modes, poisoned IQ destinations, exact retained Q/gate/K/V/IK outputs, and
surrounding cache guards.

Validation: test_qwen4_memory and the targeted Metal attention-preparation
suite pass on M1 Max with Metal API validation enabled.
 2026-09-16 08:21:03 +02:00
Salvatore Sanfilippo f396751952 Optimize Qwen CUDA projections and expert kernels 
Accumulate dense component products directly in the output and fuse short F16 projections while retaining FP32 correction sums and cuBLAS fallbacks. Tune expert tile sizes by weight format and unpack IQ2 signs in packed bytes without changing represented weights. Keep BF16 n-grams disk-only.

Extend independent kernel references for large outputs, partial tiles, alignment, strides, repeated writes and output guards. Refresh Spark performance references and repeatable sanitizer checks.

Paired Spark runs improve Q2/Q4 large prefill from 685/670 to 745/755 tokens/s, Q2 decoding from 21.0 to 22.6 tokens/s, and Q2 English MTP from 26.8 to 30.3 tokens/s. Official continuation scores remain unchanged; kernel references, sanitizers, session/MTP, mixed-prefill and vision checks pass.
 2026-09-16 00:26:12 +02:00
antirez bb6a7485f3 Keep release QA focused on repeatable checks 2026-09-15 17:04:07 +02:00
antirez 442982dd87 Update Qwen Spark performance references 2026-09-15 17:04:06 +02:00
antirez f74d410962 Check Qwen CUDA split attention against CPU references 2026-09-15 16:44:01 +02:00
antirez 02a2558225 Share Qwen CUDA prefill attention across grouped query heads 
Use tensor-core products with scaled residual correction for FP32 queries and probabilities. Keep scalar and split attention fallbacks, and extend coverage with independent double-precision references, sparse masks, empty rows and head-group tails.
 2026-09-15 16:23:23 +02:00
antirez 7fe838b8cb Reuse Qwen CUDA prefill normalization across injection chunks 
Keep the decode reduction order and verify byte-exact normalized rows and injection partials against one-token controls, including long batches and guarded tensor views.
 2026-09-15 15:40:42 +02:00
Salvatore Sanfilippo 63b5dca93e Add official Qwen quality references and validated API scoring 
Collect references from Alibaba through OpenRouter with an explicit empty system message and no provider fallback. Validate token byte alignment and reject lossy Unicode metadata for API agreement while retaining continuation likelihood scoring.

Track 100 short and 12 long reference cases, Q2/Q4 default and quality controls, continued-prefill results, and independent long-context answers. Include collector and scorer regressions and document the hosted-checkpoint and precision limitations.
 2026-09-15 15:00:11 +02:00
Salvatore Sanfilippo df6eedd972 Add optimized CUDA inference for Qwen3.8 Flash Next 
Share the Metal session graph with CUDA for text and vision, MTP and exact sampling, steering, checkpoints, and disk-only BF16 n-grams. Add bounded concurrent Linux n-gram reads, vectorized dense and expert kernels, register-held recurrent prefill state, and efficient small speculative batches.

Keep Qwen out of DeepSeek batch dispatch and exclude disk n-grams from resident memory accounting. Use normal file-backed weight loading in diagnostics and link n-gram support into the CUDA quality scorer.

Include independent kernel and memory-safety tests, mixed-prefill and recovery checks, long-context server and agent coverage, and user documentation with measured DGX Spark performance.
 2026-09-15 15:00:11 +02:00
GiorgioOppo 5cbdb6ef81 Metal: pipeline Qwen SSD MoE prefill and reuse Q2_K down rows 
Overlap Qwen batched MoE work with reads of the selected experts on M1 Max.
Previously the MM path waited for the missing payloads before running its
routed projections. Run resident experts first, then submit missing gate/up
work as soon as those weights arrive while the read pool loads down weights.
Only the remaining down projection is submitted after the second read phase.

Keep per-expert token lists and routing frequencies unchanged. Separate
resident/missing count buffers and immutable address snapshots preserve tile
selection and caller-owned routing state. Skip zero-count experts before the
MM tile logic, including the NAX helpers.

Retain pending payloads and indirect-address resources across command-buffer
flushes, including unretained command buffers. Install cache entries only
after all projections have been read, protect them with the current inflight
sequence, and drain submitted work before reclaiming failed reservations.
Recycle uninstalled slab slots or explicit buffers after read/install errors.
Report the two read phases accurately in the streaming diagnostics.

Add a Q2_K down kernel that reuses each activation load for two output rows
at the Qwen width of 640. Preserve the lane mapping, block/element order and
SIMD reduction, and ignore the unused tail of the physical 768-column row.
Select it automatically on M1 Max with the existing two-row dispatch policy;
keep the generic fallback for other shapes, devices and diagnostic overrides.

Extend regression coverage for Q2_K T1/T2/T3, odd output widths, shared rows,
nonzero quantized padding and guarded tensor views. Add SSD MM cold/mixed
fixtures and 18 down-read failure cases across three quantization pairs and
T29/T33/T128. Each new case exercises four consecutive failures followed by
an exact retry and a warm-cache pass. Verify unfinished outputs, input lists,
counts, cache ownership and batch restoration on failure.

Validation on Apple M1 Max, 32 GiB:
- Metal builds and git diff --check passed.
- Full Qwen kernel and MoE specialization suites passed, including
  T8191/T8192/T8193 boundaries.
- SSD tests passed with ordinary and unretained command buffers and with
  separate expert buffers; full-model SSD MTP cache/logit checks passed.
- 96 exploratory/final performance runs retained identical output bytes.
  The final 32-run ABBA/BAAB comparison used Qwen3.8-Flash-Next-Q2,
  SSD streaming, ctx 4096, chunk 128, temp 0 and --nothink, with short/long
  prompts and MTP on/off (four observations per arm and configuration).
- Relative to the already optimized resident-overlap/Q2-reuse baseline,
  the additional SSD pipeline improved median prefill by 6.4-10.9%.
  Median decode changes ranged from -4.4% to +3.1%; this is a prefill
  optimization with a measured decode tradeoff, not a uniform speedup.
  Cache budgets, hits/misses and read volumes matched within each case.

The experimental Q8 row variants and FP16 activation staging were removed
after full-model measurements. No model-format or precision change, new
permanent opt-in, benchmark report or temporary artifact is included.
 2026-09-15 11:15:31 +02:00
GiorgioOppo 81b9ebbf7a Metal: overlap Qwen MoE down and MTP verification with SSD reads 
The existing SSD prefix handled resident/shared gate-up work for one token.
Two- and three-token MTP verification waited for all missing experts, and
the down projection remained after I/O even when its inputs were resident.

Give MoE launches independent slot masks for T1/T2/T3, with explicitly empty
rows and compact grids mapped back to the original routed/shared slots.
Apply the masks to generic, IQ2 and Q4 mid kernels and both down variants.
Submit resident/shared mid and down work before reading missing experts,
then finish only missing slots before the unchanged reduction and HC update.

Snapshot all three GPU-address tables so cache installation cannot race the
prefix. Retain their buffers and mark referenced expert resources inflight;
keep failure drains, batch ownership and off-size/full-layer fallback paths.
Numerical dot/reduction order and slot layout are preserved. No new opt-in
is required. Timing totals count shared submission/read intervals once.

Extend SSD fixtures with independent T2/T3 masks, empty prefix/suffix rows,
duplicates, shared/no-shared experts, Q2 padding and Q4/MXFP4/Q8 formats.
Verify exact mid/down/reduce/HC output, guards, missing-only reads, warm
operation with an unreadable fd, failure cleanup and retry. Also exercise
24 disjoint cache replacements through the public T128 streaming path.

Validation: make all; memory and cache-state tests; full Qwen Metal kernel
suite; SSD default, unretained, Q4-NR1, MXFP4-prefetch and specialized MV
variants; real-model MTP prefix/checkpoint/continuation and n-gram recovery
oracles all passed. SSD stage comparisons and recovery logits are exact.
The independent differently batched continuation oracle remained within its
existing tolerance (maximum absolute logit error 0.000974178).

Combined-series benchmark against b86c8ae: M1 Max 32 GiB, Qwen Q2, Metal SSD,
ctx 4096, chunk 128, temp 0, no thinking. Four runs per version/scenario,
ABBA/BAAB order, 32 runs total; every generated output was byte-identical.
Short prompt: 29 input/100 output tokens; long: 575 input/50 output tokens.
Median prefill/decode token rates, baseline -> complete series:
  short MTP:    9.695/8.250  -> 9.575/10.170  (-1.24%/+23.27%)
  long MTP:    39.865/8.405  -> 42.840/11.015 (+7.46%/+31.05%)
  short normal: 9.455/7.415  -> 9.375/7.955   (-0.85%/+7.28%)
  long normal:40.940/8.620  -> 44.890/9.445  (+9.65%/+9.57%)

These measure all three changes together on this host, not this commit in
isolation. Baseline decode was variable (short MTP 7.71-10.82 t/s; long MTP
8.25-10.78 t/s); filesystem caches were not purged. The short prefill shows
no stable gain. Intermediate reports, logs and model files are excluded.
 2026-09-15 09:19:17 +02:00
GiorgioOppo e2b4a47e07 Metal: plan Qwen SSD expert cache evictions once per batch 
A full Qwen expert cache scanned all 40,960 metadata entries for each missing
expert in a selected batch. Large prefill unions repeated that scan hundreds
of times before parallel reads could start.

Collect at most 512 eligible candidates with a bounded heap and sort them by
hotness, last use, layer and expert ID. Consume each victim at the original
buffer-allocation point, preserving the scalar selector's tie order, lazy
eviction, buffer offsets and accounting. Planning itself never evicts an
unused candidate if a subsequent allocation or read fails.

Limit the path to owned-buffer selected batches and invalid target entries.
Keep the existing selector for inflight/service-thread states, mapped
batches, live target replacement and recovery after a stale or exhausted
plan. GLM/DeepSeek mapped-batch callers retain their current path.

Add a CPU-only Metal-cache oracle test for 1/16/17/511/512 victims, a full
40,960-entry cache, protected IDs, numeric ties, sentinel values, interrupted
consumption, stale/inflight plans, live replacement and exact counters.
Register the make target and ignore its generated executable.

Validation: make test-metal-ssd-reuse passed. The full series additionally
passed the public Qwen SSD fixture with 24 disjoint victims. Combined-series
M1 Max measurements reduced long-prompt cache scan time by 81% without MTP
and 86% with MTP; these are nested preparation times, not end-to-end gains.
 2026-09-15 09:19:02 +02:00
GiorgioOppo 9e1429b8b6 Qwen: bound MTP staging to recover SSD expert cache capacity 
The memory admission path reserved two complete copies of the largest routed
layer, including the differently quantized MTP predictor. Ordinary draft and
chain calls select only ten predictor experts, so the reserve unnecessarily
reduced the trunk expert cache when MTP was enabled.

Compute the staging bound from the maximum of two full trunk windows, two
compact predictor payloads with their three address tables, and one full
predictor layer for the older-Metal/no-gpuAddress fallback. Preserve the full
fallback bound: streamed MoE drains previous transient allocations before
loading another payload, and discards compact buffers before full fallback.

For Qwen3.8-Flash-Next-Q2 this changes staging from 2,778,726,400 to
1,525,678,080 bytes. On the tested M1 Max/4096 context, the automatic expert
cache grows from 7,371 to 8,212 entries without increasing the planned budget.

Add metadata-only tests for real padded Q2 tensor dimensions, disabled or
absent MTP, compact/full fallback dominance, address tables, invalid shapes,
and saturated or overflowing sizes.

Validation: make all and test_qwen4_memory passed. The complete three-change
series also passed Metal kernel/SSD fixtures and the real-model MTP and
n-gram recovery oracles; performance of the series is recorded in the final
commit, rather than attributed to this memory change alone.
 2026-09-15 09:18:35 +02:00
GiorgioOppo b86c8ae1db Metal: remove redundant Qwen MTP projections and expert work 
Preparing historical predictor cache rows previously computed query/gate
and indexer-query projections that were never consumed. Catch-up batches
also ran the predictor MoE for every row even though only the final
residual feeds the vocabulary head or the next recursive draft.

Add a cache-only attention preparation kernel that writes K, V and raw
indexer keys using the original normalization, RoPE and cache arithmetic.
Skip query projections at graph level while retaining fused K/V/IK GEMV
for short batches and the existing pooled-key updates. Ordinary attention
continues through the complete preparation path.

Add a Q8 EH projection for the staged [embedding|0] and [0|hidden] rows.
Separate embedding and hidden dispatches within one encoder, traversing
only the active weight half while preserving full row strides and the
original float4 dot products and SIMD reduction order. Check the inactive
half's scale words to preserve NaN/Inf classification for zero products.
Validate tensor spans, alignment and output aliasing before dispatch.
Select this kernel on M1 Max for E2560, four HC streams and T1..3; keep
the existing path on other devices/layouts, non-Q8 weights and diagnostic
GEMM or Q8 prefill-profile selections. No quantization change, new buffer
allocation or production opt-in is introduced.

Keep batched attention, HC mixing and router arithmetic during MTP
catch-up, then bind last-row views for the independent routed/shared MoE
and residual reduction. Restore graph bindings and release views on all
exit paths. Earlier rows still build the full predictor KV prefix; only
the final residual is consumed. Ordinary trunk MoE remains unchanged.

Extend the focused Metal kernel regression with T1/T2/T3, actual Qwen
projection dimensions, cache boundaries and guards, nonzero multimodal
RoPE positions, signed zero, nonfinite Q8 scales, diagnostic fallbacks,
input immutability and invalid-span/alias rejection. Add a frozen original
MTP input/layer oracle to the model-backed prefix test, including recursive
drafts and exact three-row verification. Preserve diagnostic environment
settings and cover unfused execution without assuming speculation runs.

Validation on M1 Max 32 GiB with Qwen3.8-Flash-Next-Q2.gguf:
- Application and relevant test targets built; git diff --check passed.
- Focused synthetic cache/projection kernel tests and memory checks passed.
- Real SSD predictor comparisons matched EH, KV/pooled keys, final residual,
  all 248320 logits and recursive drafts bit for bit for T1/T2/T3.
- Prefix, chunk/append, payload and unfused tests passed. Session checkpoint,
  restore/rewind, exact-sampling continuation, n-gram failure recovery and
  small-buffer regressions passed. Independent target continuation stayed
  within the established batch tolerance (max absolute error 0.000974178).

Performance versus 92f57fc, frozen binaries/shaders, SSD streaming,
ctx4096/chunk128, temp0, no thinking, MTP enabled via --mtp-timing:
- Four runs per version and prompt in ABBA BAAB order, harmonic-mean rates.
- Rome, 29 input/100 output tokens: prefill 9.85 -> 9.88 t/s (+0.3%),
  decode 9.01 -> 9.57 t/s (+6.2%); mean process time increased 0.8%.
- Long prompt, 575 input/50 output tokens: prefill 41.28 -> 42.12 t/s
  (+2.0%), decode 9.40 -> 9.11 t/s (-3.0%); process time decreased 0.5%.
- Generated outputs and first-draft acceptance counts were identical.
- Isolated EH microbenchmark: 16 samples per arm, 64 dispatches per sample,
  including host encoding and GPU completion. T1: 161.28 -> 127.36 us;
  T2: 309.67 -> 188.75 us; T3: 438.84 -> 277.03 us (+27..64% throughput).
- A separate no-MTP control measured 7.58 -> 7.33 t/s decode (-3.3%).
  End-to-end variation remains material; no stable global 10% improvement
  is established. Final runs had no concurrent model processes, but OS
  file-cache and desktop background activity were not controlled.

CUDA/ROCm were not exercised; graph changes are in the Metal Qwen path.
Temporary benchmarks, logs, notes and model files are excluded.
 2026-09-14 20:53:23 +02:00
GiorgioOppo 92f57fc4d7 Qwen: prepare MTP prefix caches and preserve predictor state 
Qwen prompt sync previously ran only the trunk. The first predictor step
then wrote at the decode position while attention read the entire nextn
prefix, including unwritten rows or data left by a previous session.
Target verification still gated the generated tokens, but the drafter's
context was incomplete and could reduce proposal acceptance.

Build the predictor K/V, raw indexer keys and pooled block keys while each
prefill chunk's trunk hidden states are available. Share the existing MTP
input projection and attention preparation, using at most three rows of
draft scratch. Skip the predictor attention output, MoE and vocabulary
head, avoiding routed-expert SSD reads for this cache preparation.

Retain the last committed trunk row until its actual successor is known:
- Repair the last predictor row before prompt append, ordinary decode or
  verification, including a sampled token that differs from the argmax.
- Produce the first draft from this retained row rather than requiring an
  initial ordinary decode cycle to seed the predictor.
- Preserve the row in all three verifier snapshots and restore it during
  rejection, exact-sampling replacement and caller-requested rewind.
- Track a contiguous predictor prefix and invalidate provisional chain
  rows and staged inputs when restoring or preparing only the cache.
- Initialize future RoPE positions for recursive drafts without advancing
  the trunk's multimodal position state.

Qwen payload v3 stores only the canonical predictor prefix plus its pending
trunk row. Reject old payload tags that cannot prove a complete prefix;
old Qwen disk checkpoints must be rebuilt. Account for the retained row
and verifier copies in memory estimates: 40 KiB per row on this model,
up to 160 KiB including all snapshots. Clarify that timing statistics count
first-draft acceptance, and document MTP usage with Metal SSD streaming.

Add a model-backed Metal regression comparing cache-only preparation with
the complete MTP layer on identical trunk batches. Cover poisoned caches,
one-token prompts, chunk lookahead, non-argmax append, reused sessions,
MTP on/off payload interoperability, legacy-tag rejection and depth-2,
depth-3 and exact-sampling continuations. Restore FORCE_ACCEPT after the
rewind test and remove assumptions that the first MTP call commits only
one token. Allow the small-buffer CLI regression to use SSD streaming.

Validation on M1 Max 32 GiB, Qwen3.8-Flash-Next-Q2.gguf:
- Standard make and Qwen memory-estimate test passed.
- MTP prefix oracle passed with bitwise cache, residual and prefill logits
  agreement; continuation logits stayed within the established tolerance
  (maximum absolute error below 0.001 across 248320 vocabulary entries).
- Prefill checkpoints, reused restore, rewind, exact-sampling resample,
  n-gram read-failure recovery and small/unfused buffer tests passed.
- Two runs per arm in ABBA order, ctx4096/chunk128, SSD streaming, temp0:
  short-prompt decode 8.13 -> 8.09 t/s; 575-token-prompt decode
  8.87 -> 9.88 t/s (+11.4%), using harmonic means. Output was identical.
  Long-prompt first-draft acceptance rose from 45.5% to 56.2%; prefill
  throughput fell 2.5% on the long prompt and 6.4% on the short prompt.
  Mean total process time improved only 1.1% on the long prompt and
  increased 3.4% on the short prompt, including startup and prefill.
  These are local observations with uncontrolled OS file-cache/background
  variation, not a general 10% speed guarantee. CUDA/ROCm and vision were
  not exercised; the affected Qwen execution path is Metal.
 2026-09-14 19:45:36 +02:00
GiorgioOppo 496b153eed Metal: adapt Qwen SSD MoE tiles and specialize low-bit prefill 
Short Qwen SSD prefill batches route relatively few tokens to each of 512
experts. On M1 Max, the previous 32-token matrix tiles waste matrix work
on padding and the generic dequantizer retains run-time format selection
inside the K loop.

Use the CPU routing frequencies already validated against selected IDs and
GPU lists to choose eight-token tiles with one grid slice when their padded
token work is less than half that of the 32-token tiles. Preserve wider
tiles for concentrated routing so weights are reused across more tokens.
This adds no allocation, GPU readback, or synchronization.

Enable the existing type-specialized IQ2_XXS/Q2_K matrix kernels on this
measured SSD path. Function constant 900 removes dynamic format selection;
dequantization, FP16 operands, accumulation order, and output layouts stay
unchanged. Preserve the existing diagnostic overrides and PSO cache keys.

Scope automatic selection to M1 Max, 512 experts, and batches of 33..128
tokens. Batches of at most 32 tokens showed no stable end-to-end benefit
and retain their old pipelines. Resident inference, other devices/types,
large prefills, and the single-token decode path retain their prior policy.

Extend SSD regression coverage with a forced generic NT4/grid4 resident
oracle and repeated generic NT1/grid4 versus automatic SSD dispatch on the
same buffers. Verify frozen router-list membership, per-expert counts,
offsets and guards, skewed routing, the 32/33-token boundary, cache hits with
an unreadable backing fd, cache overflow, and full-selection fallback.
Compare mid, down, reduction, and HC outputs bit for bit; restore all five
diagnostic environment settings after each case.

Validation:
- make all; full Qwen Metal kernel suite and MoE specialization matrix.
- Qwen SSD fixture: default, unretained command buffers, Q4 NR1 fallback;
  final 32/33-token boundary fixture passes default and unretained modes.
- Metal SSD regression and real Qwen ngram state/failure recovery tests.
- Real SSD model: all 248,320 logits match the old generic NT4 path after
  a 342-token prefill and each of 25 subsequent decode steps.

Performance versus cdfc0d5, Qwen3.8-Flash-Next-Q2, M1 Max 32 GiB, SSD,
ctx 4096, prefill chunk 128, temp 0, no thinking, max 50 output tokens:
- 352-token prompt: 32.58 -> 36.92 t/s prefill (+13.3%).
- 575-token natural prompt: 38.96 -> 46.53 t/s prefill (+19.4%).
Each long prompt uses two alternated A/B pairs with frozen binaries and
identical generated output; rates aggregate equal token counts over elapsed
time. Application expert caches start empty; the OS cache is uncontrolled.
All measured long-prompt chunks are above 32 tokens and exercise the final
policy. No 10% claim is made for short prompts, decode, or other devices.
 2026-09-14 19:12:04 +02:00
GiorgioOppo cdfc0d51e1 Metal: reuse Qwen IQ2 and HC inputs on M1 Max 
Reduce repeated activation work in Qwen decode without changing the model's
quantization or per-row arithmetic. Select the new paths automatically on
M1 Max and retain the existing diagnostic overrides for direct comparison.

Add an IQ2_XXS gate/up kernel with two rows per SIMD group. Load each lane's
eight activation values once per quantization block and reuse them across
both rows and projections. Preserve the original partial sums, block order,
SIMD reduction, shared-expert fallback and output layout. Read 66-byte block
headers through aligned ushort loads. Carry the existing SSD address-table
function constant and compact slot mapping through the new dispatch, with
guards for odd rows and invalid expert IDs.

Add an F16 hyper-connection mixer that stores each low-rank input and its
sigmoid in threadgroup memory for reuse across streams and output rows.
Keep the two values separate to preserve (x * weight) * sigmoid(x) rounding;
pin reassociation and contraction at the original arithmetic boundaries.
Allocate 2560 bytes of scratch for rank 320 and place the barrier before
partial-group exits. Keep the paired two-token mixer and other-device
defaults unchanged. No public API or production environment flag is added.

Extend numerical tests with an explicit generic oracle versus automatic
dispatch, retaining all previous NR/NSG geometry checks. Cover T1/T2/T3,
F640/F641, shared and unshared experts, E33/E2560 HC outputs and prefix/suffix
guards. Add a focused 96 MiB test/benchmark arena using packed IQ2 weights
without dense shadows. Time seven alternating pairs on identical buffers,
after 20 warmups per arm, with 100 IQ2 or 1000 HC calls per measured batch.
Timings include CPU encoding and the final GPU completion wait.

Validation on Apple M1 Max, 32 GiB:
- Built all application targets and relevant Metal tests.
- Passed the complete Qwen kernel suite and MoE exact comparisons.
- Passed Qwen SSD fixtures in default, unretained-command-buffer and Q4_K
  NR1 configurations, plus existing Metal SSD expert regressions.
- Passed real Qwen3.8 Q2 n-gram failure/recovery tests for prefill, decode
  and MTP. Two paired model runs preserved all logits bit for bit across
  410 decode steps, with independent session state and alternating order.

Resident microbenchmarks measured IQ2 gate/up at 173.78 -> 163.91 us for
T1 (+6.0% throughput), with similar gains at T2/T3. The HC mixer measured
36.57 -> 17.64 us for T1 and 101.22 -> 41.28 us for T3.

Full-model gains are smaller and variable. Two paired runs with a shared
7708-expert SSD cache measured +2.51% and -0.08%, or +1.20% pooled over
400 measured steps after warmup. Independent CLI runs did not establish a
stable end-to-end speedup; these results do not claim a 10% gain, a prefill
improvement, or a benefit on untested devices. Temporary logs, model files
and intermediate benchmark harnesses are excluded from this commit.
 2026-09-14 18:18:33 +02:00
GiorgioOppo d6cbc77ca6 Metal: overlap Qwen decode expert reads with cached gate/up 
Qwen SSD decode previously waited for every selected expert payload before
starting the MoE. For cache-eligible single-token calls with misses, submit
gate/up and SwiGLU for cached routed slots and the shared expert before the
existing parallel pread batch. Compute missing slots after installation,
then execute the original full down projection and ordered reduction once.

Implementation:
- Add a private before-read callback after missing buffers are reserved and
  selected entries are protected. Existing DeepSeek/GLM callers pass NULL.
- Snapshot the resident gate/up GPU address tables so cache installation can
  update its global tables while the submitted prefix is still executing.
  Retain snapshots through GPU completion, including unretained command
  buffers, and mark cached payloads inflight before dispatch.
- Add a private slot mask to the generic and ordered Q4_K mid kernels.
  Compact the dispatch grid to popcount(mask), then recover original slot
  indices before reading IDs or weights. Preserve shared placement, output
  strides, dot-product order, and the zero-mask full-dispatch behavior.
- Drain submitted work before failure cleanup clears cache entries, and
  restore the caller's command-batch ownership for exact retries.
- Keep all-hit, batched prefill, off-size MTP and overflow staging schedules
  unchanged. No public API, CLI option, or quantization change is introduced.

Regression coverage:
- Extend the sparse-file SSD fixture with partial/cold/warm T1 selections,
  duplicate IDs, optional shared experts, and 10/16 routed slots.
- Compare mid/down/reduce/HC outputs and guards bit for bit with resident
  kernels. Restrict actual pread requests to unique missing expert ranges;
  verify failed-read prefix completion, untouched missing slots, recovery,
  and zero-read warm hits with an unreadable file descriptor.
- Passed the fixture in default, unretained-buffer and Q4_K NR1 modes;
  resident MoE specialization and ordered Q4_K tests; the existing Metal SSD
  expert regression; and real Qwen n-gram failure/recovery tests for prefill,
  decode and MTP with exact logits. Built all application targets.

Measured against 85d37ae on an M1 Max with 32 GiB, Qwen3.8 Q2, automatic
expert cache, ctx 4096, prefill chunk 128, 100 generated tokens, temp 0 and
nothink, using "narrami la storia di roma". Three alternating pairs yielded
decode medians 8.25 -> 8.44 t/s (+2.3%) and prefill 9.98 -> 10.05 t/s.
Generated output was identical. Profiled cache hits, misses, evictions and
12.80 GiB of expert reads were unchanged. Document the observed variability:
these warm-file measurements do not establish a 10% or cross-device gain.
 2026-09-14 17:50:40 +02:00
GiorgioOppo 85d37aec2a Metal: stage only selected Qwen experts for MTP and cache overflow 
Avoid reading all 512 experts when a Qwen layer has a different size from
the persistent cache or a prefill selection exceeds its capacity. The Q2
pack's Q4/MXFP4 MTP layer previously read about 1.29 GiB on each invocation,
even when the router selected only ten experts.

Allocate compact gate/up/down payloads for the distinct selected experts
and build private GPU-address tables indexed by their original IDs. Keep
the existing row/MM/NAX arithmetic, Q2 physical down-row padding and final
reductions unchanged. Preserve the trunk cache and its address tables.
Retain all six staging buffers through GPU completion, including batches
using unretained command buffers, and declare the indirect payloads as
encoder resources without registering nonexistent cache entries.

Use the original three contiguous reads when every expert is selected or
GPU buffer addresses are unavailable. Choose that fallback before I/O;
actual read failures propagate without rereading the complete layer.
Keep the existing conservative memory admission and staging reservation.

Extend the SSD fixture with a test-only pread interceptor, linked through a
separate Metal backend object. It performs real reads, rejects accesses to
unselected expert ranges, checks per-projection byte totals and verifies
zero reads for resident/cache-hit paths. Cover compact cache overflow,
off-size MTP with two/three tokens, full selection of all 512 experts,
failed reads, retries and retention of the original trunk cache.
Update the Qwen documentation to describe selective temporary staging.

Validation on Apple M1 Max, 32 GiB:
- Full build, host Qwen memory test and existing resident MoE exact tests.
- SSD bitwise mid/down/reduce/HC parity and read-range assertions, including
  DS4_METAL_UNRETAINED_COMMAND_BUFFERS=1 and existing Metal SSD regressions.
- Real Q2 n-gram recovery: exact logits after prefill/decode/MTP read errors
  across two SSD sessions.
- Three alternating before/after runs of the Rome prompt, ctx4096, chunk128,
  50 generated tokens, temperature0 and no thinking: median MTP generation
  4.55 -> 7.86 t/s (+72.7%). All generated output matched, with 27 verifier
  cycles and 21 accepted drafts in every MTP run.
- A separate profiling comparison reduced expert pread bytes from
  48.40 -> 12.02 GiB (-75.2%); these include filesystem-cache reads.

The speed result is specific to this short MTP workload. Expert I/O still
finishes before routed computation; no overlap or next-layer prefetch is
introduced, and CUDA/ROCm or M5 NAX execution was not tested here.
 2026-09-14 17:16:37 +02:00
GiorgioOppo e37f18576c Metal: add bounded SSD expert streaming for Qwen3.8 Flash Next 
Allow the existing --ssd-streaming mode to run the Qwen Metal graph without
making every routed expert resident. Keep dense, shared, GDN, attention and
hyper-connection weights mapped; retain the separate disk-only n-gram path.

Reuse the explicit-pread expert cache with 512 expert IDs and top-10 routing.
Dispatch the existing row, grouped-MM and NAX kernels through specialized
expert-address tables, preserving their accumulation and reduction order.
Keep Q2_K physical down-row padding separate from logical activations, and
pass the shared gate/up and down quantization types independently.

Stage complete layers in owned buffers when a batch exceeds cache capacity
or the layer has a different expert size, including the Q2 pack's Q4/MXFP4
MTP layer. Never fall back to mapping routed weights for Qwen. Retain these
buffers through GPU completion, including unretained command buffers.
Validate router IDs, lists, tensor ranges and aliases before dispatch, and
preserve command-batch ownership across read failures and retries.

Size the cache after accounting for static weights, context, recurrent and
speculative state, prefill workspace, staging and runtime headroom. Share
the Qwen estimator between GPU and CPU-only builds, fix its large-context
integer overflow, and check additional sessions and direct-generation
graphs against the same budget. Reject unsupported preload/full-layer
options instead of silently applying another model's policy.

Add sparse-file GPU parity tests for cold/warm caches, eviction, overflow,
mixed expert sizes, invalid input and failed-read recovery. Add host memory
regressions and enable the real-model n-gram recovery test in SSD mode.
Document cache sizing, commands and current I/O/MTP limitations.

Validation: full macOS build; Qwen SSD parity with ordinary and unretained
command buffers; existing Qwen resident decode/Q4 geometry tests; existing
Metal SSD IQ2/Q4/MXFP4 tests; host cache and Qwen memory tests. A real Q2
Metal SSD run completed the Rome prompt with a 50-token generation limit.
The real-model n-gram test also passed prefill, decode and MTP read-failure
recovery with exact logits across two SSD sessions.
 2026-09-14 17:00:10 +02:00
antirez 9139e2ae58 Document and download self-contained Qwen BF16 n-gram releases 
Use one verified GGUF per quantization, remove sidecar setup, and document disk versus resident size. Record native-reader, model, agent/server, steering, fallback, and continued-prefill QA with explicit limits.
 2026-09-14 13:46:19 +02:00
antirez 67bd55c0ec Honor steering in standalone greedy Qwen generation 
Pass the requested vector and scales into the standalone graph, as the session runner already does. The native-BF16 QA exposed the missing load because ordinary generation silently ignored attention steering. Require the enabled diagnostic and retain exact ordinary/MTP text checks for both Q2 and Q4.
 2026-09-14 13:41:49 +02:00
antirez 846428239d Keep short Qwen prefill tails on FP32-query attention 
The native-BF16 Q4 fixture exposed half-query rounding amplified by nearly tied MoE routing scores at a four-row tail. CPU double-precision attention checks confirmed the matrix formula, but using decode arithmetic for at most eight rows keeps this boundary stable without changing large-prefill kernels. Extend kernel and full-logit coverage through the eight-row tail.
 2026-09-14 13:33:34 +02:00
antirez 97ba327538 Read Qwen native BF16 n-grams directly from the GGUF 
Keep the table outside model mappings and Metal residency. Use bounded explicit reads, reject failed lookups before GPU consumption, and remove the runtime sidecar option and paging controls. Preserve Qwen hash and speculative state handling; cover malformed files and exact recovery after prefill, decode, and MTP read failures.
 2026-09-14 13:32:48 +02:00
antirez e68bcc8fab Package Qwen with original BF16 n-grams 2026-09-14 13:00:36 +02:00
Salvatore Sanfilippo 15e88cfe91 Link CUDA regression tests with their shared helpers 
Include the image helper object used by the CUDA visual-attention code in standalone test binaries. Let the engine runtime test reuse the normal core dependencies, excluding its separately instrumented engine object.

Verified all affected targets link on DGX Spark sm_121a; CUDA long-context, Q8 row and DeepSeek 4.1 kernel regressions pass.
 2026-09-14 11:44:47 +02:00
Salvatore Sanfilippo 948307e189 Keep Qwen steering capture compatible with Python 3.9 
Defer annotation evaluation so the new optional PLE argument can be imported by the system Python. Verified steering, MTP, all 48 layer captures, and prompt-only dumps on M5 Max.
 2026-09-14 11:13:31 +02:00
Salvatore Sanfilippo cbaf6cd336 Use explicit Qwen scoring inputs and remove an unused CUDA stub 
Add --ple to the official continuation scorer and use its existing --rendered-prompt option instead of rewriting hardcoded token IDs through an environment variable. Remove the unreferenced multi-shard mapping stub introduced by the PR.
 2026-09-14 11:06:16 +02:00
Salvatore Sanfilippo 923b819abf Harden Qwen checkpoints, speculative decode and tool calls 
Restore matching logits on verifier rewinds, replay stale recurrent state, discard old snapshots on restore, and invalidate partial payload reads. Keep lazy snapshots all-or-nothing and respect small prefill buffers and the unfused diagnostic path.

Preserve schema-declared string arguments. Repair all-row reference dumps and teacher-forced chunking. Add model-backed regressions and fix test-fixture leaks exposed by ASAN.
 2026-09-14 11:06:16 +02:00
Salvatore Sanfilippo efc27d6d86 Keep literal reasoning and tool tags inside Qwen arguments 
Use the structural reasoning delimiter on the server and wait for the parameter delimiter while streaming in ds4-agent. Add server and bytewise agent regressions; both fail before the fixes and pass afterward on M5 Max.
 2026-09-14 10:34:10 +02:00
Salvatore Sanfilippo bab202990e Keep the Qwen Metal graph out of CUDA and ROCm builds 
Reject unsupported Qwen execution modes before warming or sharding the model. Preserve the CPU first-token reference. The ROCm build now links without Metal-only symbols and is warning-free.
 2026-09-14 10:32:26 +02:00
Salvatore Sanfilippo f06d3ebd52 Merge Qwen3.8 Flash Next support from PR #991 
Closes #991. Preserve the original implementation commit and all author credits.
 2026-09-14 10:13:08 +02:00
Dango233 0a21d1abb8 bench: revalidate Metal decode queue against current upstream 2026-09-14 15:37:08 +08:00
Dango233 7a18d09149 Align with upstream DeepSeek V4.1 CUDA support 2026-09-14 15:11:44 +08:00
Salvatore Sanfilippo 6e4c285a92 Record CUDA SSD prefill gains and regression checks 
Document the measured README, fresh and continued prefill gains, exact state and continuation checks, memory bounds and agent regression. Include the small automatic-cache decode cost instead of claiming a decoding speedup.
 2026-09-14 00:13:55 +02:00
Salvatore Sanfilippo e9e1baaff1 Process medium CUDA SSD appends in one layer sweep 
Reuse the carry workspace for medium appends with a substantial final tile. Preserve the existing 2048-row arithmetic partitions instead of visiting every layer twice. Retain the old dispatch for tiny tails and explicitly smaller prefill chunks.

Extend dispatch and mixed-state controls. Full logits, caches, snapshots and next decode remain exact through 60K context and fresh 49K deferred prefill. Together with read-ahead, the 3241-token README prefill improves from about 55 to 98 tokens/s with a 64 GiB cache hint.
 2026-09-14 00:01:08 +02:00
Salvatore Sanfilippo c2c3ce324c Overlap CUDA SSD expert reads with V4.1 prefill 
Read one next layer into reserved slots of the admitted expert cache using two 8 MiB staging buffers. Protect current inputs, publish only completed copies, keep unused read-ahead cold, and join before cancellation or cache/model teardown.

Enable read-ahead from 2K tokens; smaller prefills were slower with full-layer reads. Add eviction, truncated-source, buffered-I/O, cancellation and early-exit tests. Validate exact logits and saved state on the Sparks, including Compute Sanitizer.
 2026-09-14 00:00:54 +02:00
Salvatore Sanfilippo a04f46fa42 DeepSeek v4.1 Flash support for CUDA 
Add bounded SSD streaming, exact short prefills, accelerated large prefills, and native session batching. Support resident tensor parallel inference across two DGX Sparks over RoCE, with aligned expert shards, split vocabulary projection, decode graphs, and shared-expert overlap.

Include primitive, state, transport, and model regression tests, plus usage documentation and measured QA results.
 2026-09-13 17:32:48 +02:00
ivanfioravantiAndreas ParuseldavidMichael KuebbelerOgilthorp3anemll
 		ccea768827 Add Qwen3.8 Flash Next model support 
Implement qwen4exp loading, tokenizer and chat formatting, PLE sidecars,
Metal text and vision graphs, built-in MTP, and recurrent KV persistence.
Integrate the CLI, agent, server, download targets, and conversion tools.

Preserve the native Q4_1 encoder, Q4_K accumulation guard, and shared
quantizer build target. Keep exact-MTP rewind and vision-cache row-width
fixes. Persist valid logits and recurrent state at completed prefill chunks,
and reuse Qwen text-answer history when tools clients omit reasoning.
Retain exact reasoning replay and test continued checkpoints, cancellation,
cache eviction, and restart reuse with and without reasoning_content. Include
reproducible charts and measurements comparing throughput before and after
the checkpoint fixes on M3 Ultra.

Add Qwen trunk steering and prompt activation capture, including GDN layers.
Keep the MTP predictor outside the 48-row bank and preserve prompt dumps
through one-token prefill tails and multi-row speculative verification.
Recover Chrome target creation and retain its last page target. Preserve
full agent transcripts in stripped-session format when KV no longer fits,
while retaining the newer automatic compaction and continuation logic.
Add steering/MTP/capture, Chrome CDP recovery, and full-session save/reload
regressions. Adapt the intended changes from fork PRs #13 and #7.

Reconcile with antirez main bd66c40, preserving DeepSeek v4.1 model IDs,
tool syntax, session state, and Metal kernel dispatch. Exclude superseded
engine code and prototypes.

Co-authored-by: Andreas Parusel <aparuael@7170c824-6152-4cf1-985e-9503832ce673.fritz.box>
Co-authored-by: david <contact@davidberard.fr>
Co-authored-by: Michael Kuebbeler <cropduster@users.noreply.github.com>
Co-authored-by: Ogilthorp3 <bertrand@nepveu.name>
Co-authored-by: anemll <realanemll@gmail.com>
 2026-09-13 14:57:29 +02:00
Dango233 9b50495847 DeepSeek V4.1 Flash: queue bounded single-host Metal SSD decode 2026-09-13 11:19:19 +08:00
Salvatore Sanfilippo bd66c40207 DeepSeek v4.1 Flash support for Metal 2026-09-12 11:47:54 +02:00
Salvatore Sanfilippo 6289c51627 Give agent hints a distinct terminal label 
Render the existing hint marker with a white-on-teal badge and a slim left rule. Keep body text, redirected output and the model transcript unchanged.

Test every stream split, unfinished markers, quoted continuations, literal code and both tool syntaxes. Verify terminal colors and UTF-8 at multiple widths, including redraws between fragments.
 2026-09-08 17:26:35 +02:00
Salvatore Sanfilippo 74fbe546ff Add vision client and compaction recovery release checks 
Record the focused M5 Max results, long-prefix image latency, actual Pi tasks across three APIs, and the remaining pre-existing DSpark stop-boundary rebuild. Distinguish this pass from untested release platforms.
 2026-09-08 14:12:56 +02:00
Salvatore Sanfilippo e33a5b2538 Keep agent tasks running across context compaction 
Reserve summary space, recover already-full conversations by summarizing a bounded prefix, and retain every unsummarized token. Reject oversized input without corrupting the current conversation.

Resume unfinished responses in the same assistant message while preserving thinking state and the output budget. Discard incomplete tool calls before compacting and retry without executing partial arguments.

Cover real-model compaction, generated-code correctness, budget accounting, and recovery in a PTY regression; add deterministic boundary tests.
 2026-09-08 14:12:35 +02:00
Salvatore Sanfilippo 81a76585ff Preserve live vision context across server tool calls 
Reuse bounded image embeddings, preserve exact image identities while appending new images, and splice image tokens into authenticated live continuations. Handle image tool results in Responses and Anthropic requests as well as Chat Completions. Normalize request-local image markers for visible replay without trusting text placeholders as image identities.

Add prefix and encoder-cache unit coverage plus real-model image replay and Pi coding-client regressions. Validated GLM 5.3 Flash and DeepSeek Vision-Exp on M5 Max, including 53K/56K prefixes, batched sessions, all three API formats, fresh logit replay, and frontend ASan/UBSan.
 2026-09-08 13:57:09 +02:00
Rowan TrollopeSalvatore SanfilippoClaude Fable 5
 		6546ae5473 agent: add opt-in programming hints 
Adapt teaching-mode PR #391 to /hints on|off. Append the full
instruction once per context and short state changes at safe
assistant boundaries. New and resumed sessions start off;
compaction reinstates the current policy.

Use ordinary Markdown asides and the existing renderer, without
a learner profile, new tool, or persistent preference.

Tested on M5 Max IT with DeepSeek Flash and GLM Flash coding sessions,
prefix reuse, queued changes, compaction, save/switch/new, frontend
and PTY tests, and agent ASAN/UBSAN checks.

Closes #391

Co-authored-by: Salvatore Sanfilippo <antirez@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
 2026-09-07 18:16:55 +02:00
antirez f62ca29a30 README: intro section updated. 2026-09-07 13:23:08 +02:00
Salvatore Sanfilippo b6af0adf8c Bound GLM SSD prefill cache priorities so decode can adapt 
The recent-token bonus selects the initial cache contents; it is not a count of observed decode routes. Passing that inflated score into eviction hotness kept the seed effectively pinned after long prefills, starving demand-loaded experts.

Use the existing bounded rank policy without changing the seeded expert list or inference arithmetic. On M5 Max IT, full GLM 5.3 SSD agent decode at 100K context improved from 0.48 to 4.38 t/s. Matched 128-token greedy replays stayed identical, and the Metal SSD cache tests passed under API validation.
 2026-09-07 01:39:49 +02:00
Salvatore Sanfilippo c0a6119f36 Preserve the final byte of in-memory session snapshots 
Reserve a separate byte for the terminator written by fmemopen, without changing the serialized payload length. Add model-free grow/reuse coverage and compare real-model snapshots against file serialization.
 2026-09-06 23:59:05 +02:00
Salvatore Sanfilippo 49ec955c8c Keep ROCm QA builds clean for partial tiles and CPU-only fixtures 
Do not force full unrolling of runtime-bounded IQ2 tail loads. Keep the ROCm graph memory-baseline initialization out of DS4_NO_GPU builds, matching the declaration and allowing the model-free session-state suite to compile.
 2026-09-06 23:41:42 +02:00
Salvatore Sanfilippo d12898f4e2 Initialize the terminal escape buffer for GCC dataflow checks 
Make the empty escape buffer explicit without changing incremental key handling. This removes the maybe-uninitialized warning found by the DGX Spark QA build.
 2026-09-06 23:41:41 +02:00
Salvatore Sanfilippo 78c522e96b Keep DSpark EOS fixture compatible with pre-fix source 
Use a macro for the scheduler policy mock so historical verifier extraction does not fail -Werror on an unused helper. Check the current verifier with sanitizers and confirm that the pre-fix source reproduces the excess committed EOS row.
 2026-09-06 23:06:12 +02:00
Salvatore Sanfilippo 5ad935d7bb Merge PR #996: validate saved quality scores and frontier logits 
Retain the original commit by Donato Capitella. Add backend-independent offline validation for complete scorer coverage and full-vector numerical comparisons without changing inference.

Closes #996
 2026-09-06 23:03:45 +02:00
Salvatore Sanfilippo bff84aec32 Adapt DSpark EOS regression fixture to current scheduler helper 
Mock the existing short-fallback policy referenced by the extracted verifier. Keep scheduler behavior outside this state-ownership fixture, and retain compatibility with pre-fix negative controls.
 2026-09-06 23:03:45 +02:00
Salvatore Sanfilippo 909d6560cb Merge PR #992: ROCm gfx1151 prefill and DSpark improvements 
Preserve all five original commits by Donato Capitella. Resolve the DSpark fallback conflict by retaining the existing Metal/DGX Spark short-proposal behavior and adding the ROCm reference-alignment condition.

Closes #992
 2026-09-06 22:58:22 +02:00
Salvatore Sanfilippo b36e26ed13 Make agent file tools and background shell jobs reliable 
Replace files atomically with metadata and concurrent-change checks. Stream bounded reads with exact continuation, report incomplete searches, preserve literal [upto] edits, and avoid silently truncating user text.

Drain shell output and enforce deadlines independently of model progress. Synchronize observations and cleanup, bound log tails, honor status refresh requests, and filter terminal commands from displayed shell output. Align tool prompts and schemas with actual behavior and vision availability.

Add file, job, schema and PTY regressions. Combined fixes passed frontend tests, ASan/UBSan, ThreadSanitizer and DeepSeek/GLM editing smoke tests on M5 Max IT. Linux-specific metadata preservation still needs native QA.
 2026-09-06 22:50:54 +02:00
Salvatore Sanfilippo db1942148c Fix agent terminal redraws, Unicode input and Markdown rendering 
Buffer complete terminal updates, repaint status without erasing live input, and retain pending redraws across throttling. Preserve fragmented UTF-8, navigation and paste input, and bound tab completion safely.

Use shared Unicode width handling for input, output and queued prompts. Preserve unmatched Markdown delimiters and publish code recoloring as one update. Add focused terminal regression tests.
 2026-09-06 22:50:20 +02:00
Donato Capitella 85a4f0e611 tests: validate complete score tables and frontier logits 
Add two offline standard-library utilities with 30 host tests. Require complete API-covered official-score manifests, finite consistent values and matched denominators; optionally enforce per-case identity. Validate every expected ds4-bench frontier and complete float32 vector, preserving signed zero and reporting drift separately from invalid artifacts.

The existing comparator silently accepts missing rows, duplicate rows and NaN in reproduced controls; the validator rejects them. Saved Flash/Vision 100-case pairs pass unchanged. The frontier tool detects full-vector drift despite matching argmax IDs at all four measured frontiers. Document use and scope without changing inference or introducing model/GPU dependencies.
 2026-09-06 21:37:42 +01:00
Donato Capitella 394865c5d1 rocm: align DSpark target features and retain their rolling window 
Pair the previous target feature with the next seed at their respective positions and construct target KV directly from the projected normalized feature. Merge trusted target captures into the model support window instead of exposing the full prefill or replacing history with each verifier batch. Keep temporary draft rows outside persistent history and reset safely across capture gaps.

Scope the correction to physical gfx1151 independently of diagnostic fast-path overrides. Add 162 CPU cache-range and physical-ring assertions, and document the tested five-draft/confidence-0.7 configuration. Exact-sampling history-gap performance remains unqualified.

Sampled coding continuation improves 14.17/14.38 to 25.10/24.46 tokens/s at 8K/16K. Long coding passes 24/24 checks, sampled coding 7/7, and an actual tool round trip passes. Both 100-case official reports remain identical; the 64K native continuation completes at 22.61 tokens/s.
 2026-09-06 21:26:09 +01:00
Donato Capitella e83b106309 rocm: accelerate gfx1151 DSpark verifier attention projection 
Use independent version-matched hipBLASLt 2379 plans for batched attention-A at verifier widths 2 through 6. Preserve handle-stream ordering, shape/model guards, failure isolation and the original fallback.

Production-shape full-output, input and canary checks pass with documented rounding differences. The 5/6-row operator is about 2.55x faster; native and objective coding checks retain the prefill gain.
 2026-09-06 21:26:09 +01:00
Donato Capitella 1ef9bba608 rocm: select version-matched gfx1151 Lt prefill plans 
Use zero-workspace F32-output hipBLASLt plans 2537/2539 for the qualified DeepSeek projection shapes. Require matching gfx1151, model, library versions and supported descriptors; retain the original fallback.

Warm 16K native prefill improves 262.64 to 301.30 tokens/s. Operator safety and official scoring pass; long-prompt NLL changes +0.65% at default 4K chunks while the four coding tasks pass 24/24 checks.
 2026-09-06 21:26:09 +01:00
Salvatore Sanfilippo 90aa4b560a Document full GLM SSD append gains and validated coding QA 2026-09-06 17:29:02 +02:00
Salvatore Sanfilippo d7ade9fd65 Reuse cached IQ2 experts for small Metal GLM prefills 
Keep the existing MV, SIMDgroup and TensorOps arithmetic while reading selected experts from the bounded cache. Enable this by default for single-host eight-expert IQ2 batches below 256 tokens when a full layer fits; retain the existing larger-prefill and fallback paths.

Test exact intermediate and final outputs at both synthetic and full GLM dimensions, cache eviction, dispatch boundaries and undersized budgets. Full-model logits, coding sessions and MTP snapshot checks pass on M5 Max.
 2026-09-06 17:28:31 +02:00
Salvatore Sanfilippo fb2abb29e4 Align Metal GLM attention scratch for arbitrary context lengths 
Metal API validation aborted on full-attention and indexed-decode scratch sizes that were not multiples of 16 bytes. Pad only the allocation; kernel indexing, precision and arithmetic remain unchanged. Cover all context-length remainders, both cache formats and invalid selected rows in the attention regression.
 2026-09-06 16:29:03 +02:00
Salvatore Sanfilippo 4814f44fed Remove managed-storage notifications from shared Metal buffers 
Shared buffers do not use didModifyRange. Metal API validation rejects these calls; remove them and their now-empty timing metric. Require the SSD primitive test under API validation.
 2026-09-06 15:49:05 +02:00
Salvatore Sanfilippo 31d1992633 Record Metal SSD cache gains and the outstanding full GLM server stall 2026-09-06 15:42:09 +02:00
Salvatore Sanfilippo 437dbcba8e Read full IQ2 GLM layers ahead during large continued prefills 
Keep selected-expert prefill for typed Flash layouts. Generic IQ2 GLM reuses the sequential whole-layer path with one-layer lookahead for continued chunks of at least 256 tokens, without changing prompt chunks or expert arithmetic.

On M5 Max, full GLM 5.3 Q2 at a fixed 35.36 GiB expert cache improves 8K/4K-append prefill from 63.23/61.56 to 90.70/84.31 t/s together with bounded model views. All 154880 logits at both frontiers are byte-identical. The 512-token append improves from 14.64 to 26.01 t/s in the larger-cache test.
 2026-09-06 13:49:51 +02:00
Salvatore Sanfilippo 660e1d4b64 Bound Metal streaming model views without losing auxiliary mappings 
Replace single-range target views just like disjoint layer spans. Preserve separately mapped vision and drafter weights. Explicitly map all indexer normalization weights before GLM compact-cache warmup instead of relying on old layers remaining mapped.

Add repeated mapping and auxiliary-weight checks to the SSD expert regression, and label wrapper allocation counters as cumulative. Full GLM 8K/12K streaming completes with a 61.35 GiB expert cache without the swap growth seen in the old mapping control; Flash full logits remain byte-identical.
 2026-09-06 13:44:00 +02:00
Salvatore Sanfilippo 715ed7a865 Age Metal expert caches from the first routed layer 
Use the same cache clock for typed and generic GLM routes, including models with leading dense layers. Report every selected expert in routing diagnostics. Add an eviction oracle that checks all eight slots and byte-identical output across different cache capacities.
 2026-09-06 13:16:43 +02:00
Salvatore Sanfilippo 91dda0bc88 Let Metal expert caching adapt beyond the built-in preload ranking 
Keep the preload list but bound its initial eviction preference. Fixed-code SSD reads fall from 94.65 to 78.18 GiB; matching Vision Exp continuation scores are unchanged. No routing, precision or arithmetic changes.
 2026-09-06 13:09:45 +02:00
Salvatore Sanfilippo 85837d02e1 Use more of the safe Metal working set for DeepSeek expert caching 2026-09-06 12:19:01 +02:00
Salvatore Sanfilippo edcafe4d57 Reuse GLM router readback during IQ2 expert streaming 2026-09-06 12:09:20 +02:00
Donato Capitella 7a8c1b27c8 dspark: stop verification at the first EOS 
A greedy draft can contain EOS followed by more tokens. The verifier previously retained that suffix even when the returned token vector stopped at EOS, leaving session state ahead of the caller. Bound the validated draft inclusively at its first EOS before verification, capture and commit. Preserve upstream seed accounting, repeated prefix restoration, invalid-token and budget checks, ignored-EOS handling, and the separate exact stochastic helper.

Add a CPU regression executing the actual greedy helper against checked target-state mocks. Unmodified main returns two tokens while advancing three on the interior-EOS probe; the repair returns and advances two. All 7,217 cases and 240,503 assertions pass on the repaired helper, covering positions through the 16-token capacity, partial commits, rollback/replay, fused and ordinary seeds, output/context limits and transmitted TP tokens. No-EOS and ignored-EOS subsets pass on both versions.

This is shared verification control flow without backend arithmetic edits. The host fixture does not establish GPU numerical, public-dispatch or multi-GPU qualification; batches containing a suffix after EOS may become shorter.
 2026-09-06 11:01:26 +01:00
Donato Capitella 64e490bc1a rocm: specialize gfx1151 IQ2 prefill input staging 
Use direct input loads for full IQ2_XXS tiles with width at least 32 on HIP/gfx1151. For partial tiles, initialize the invalid suffix once and reload only the valid prefix at each K stage. Preserve scratch lifetime, barriers and invalid-tail handling; other backend and quantization templates retain their existing path.

Keep small IQ2 and Q2_K templates on their established masked-load path, where the wider load-scheduling change altered compiler accumulation order during development. The specialization changes input staging without changing the selected matrix operations or their source-level accumulation.
 2026-09-06 11:01:25 +01:00
Salvatore Sanfilippo e965352906 Preserve RoPE in full GLM compact prefill 2026-09-06 12:00:03 +02:00
Salvatore Sanfilippo 54fc4f691c Document measured Metal SSD gains and scoped coding-session QA 2026-09-06 03:44:30 +02:00
Salvatore Sanfilippo 6ccc1e7e76 Use the available Metal cache and overlap the first GLM streaming prefill 2026-09-06 03:39:00 +02:00
Salvatore Sanfilippo 22e417286c Keep static Metal streaming weights resident within the runtime budget 2026-09-06 03:39:00 +02:00
Salvatore Sanfilippo d6580a5ffc Let parallel Metal prefill reads run without serial read-ahead 
Batch expert loads already submit all requested tensors to the pread worker pool. Avoid reading those pages serially first; keep decode read-ahead unchanged. M5 Max GLM Q4 initial/continued prefill improved from 73/77 to 93/97 tokens/s in the isolated comparison, with identical full frontier logits.
 2026-09-06 02:54:46 +02:00
Salvatore Sanfilippo aa8fef4a16 Require real server coding sessions and prefix reuse checks in QA 2026-09-06 01:59:27 +02:00
Salvatore Sanfilippo 46e46e7d7c Extend tool-call QA with speculative boundaries and real coding clients 2026-09-06 01:50:26 +02:00
Salvatore Sanfilippo 42da390e2f Run frontend regressions without loading a model 2026-09-06 01:37:21 +02:00
Salvatore Sanfilippo 6eac69e0aa Bound tool recovery by the request budget and explicit stops 2026-09-06 01:34:57 +02:00
Salvatore Sanfilippo 930ab73279 Keep valid opportunistic drafts across tool sampling boundaries 2026-09-06 01:27:29 +02:00
Salvatore Sanfilippo e1711c21d9 Report tool observation errors without compacting the conversation 
Distinguish image/rendering failures from results that exceed the context budget. Preserve the transcript and return the original error. Add an invalid-image regression and run the agent suite under ASan and UBSan.
 2026-09-06 01:18:40 +02:00
Salvatore Sanfilippo d108ae4e8d Preserve literal tool data across agent and server rendering 
Do not HTML-unescape argument bodies. Share a narrow, reversible escape for each body closing delimiter, including its literal escaped spelling. Keep streaming fragments intact across those escapes. Preserve protocol-looking text inside GLM agent arguments until the actual argument terminator. Add round-trip, streaming and agent regressions.
 2026-09-06 01:14:31 +02:00
Salvatore Sanfilippo fc6414cdfc Keep literal tool and thinking markers inside argument data 
Use the incremental grammar state to detect completed calls, and skip argument/key bodies when searching final control markers. Validate complete calls before recovering unclosed reasoning. Cover DeepSeek and GLM arguments containing both tool endings and thinking tags without truncating or losing the call.
 2026-09-06 01:07:50 +02:00
Salvatore Sanfilippo 759dd7cd69 Do not repair truncated tool actions or exceed output limits 
Only repair an omitted DeepSeek wrapper around already closed invocations. Never supply a missing argument or invocation terminator. Preserve length finishes and do not retry beyond the client output budget. Log the actual stop cause, token ID and available budget when a tool call is incomplete.
 2026-09-06 01:03:59 +02:00
Salvatore Sanfilippo 5b3cc8bfd1 Respect GLM tool syntax and speculative server boundaries 
Track GLM function names and argument bodies with the existing incremental recognizer. Keep reasoning quotes out of syntax sampling. Discard unconsumed speculative tokens at tool and control stops, and resample after a structural/payload boundary with restored logits. Rebuild invalid recurrent prefixes with their image conditioning. Add GLM fragment and byte-boundary coverage.
 2026-09-06 00:58:00 +02:00
Salvatore Sanfilippo d077fa67e5 Keep split tool delimiters in the server sampling tracker 
Retain a lone opening angle bracket until the next token, without forcing code or HTML greedy. Previously the tracker consumed it and missed a parameter close split across tokens. Add byte-by-byte regressions for all three DeepSeek spellings and string/JSON bodies; all twelve failing checks now pass.
 2026-09-06 00:50:14 +02:00
Salvatore Sanfilippo aa1e71cf03 Document DGX Spark DSpark speed and focused QA 2026-09-06 00:44:53 +02:00
Salvatore Sanfilippo 2803820dc5 Use aligned Q8 vector kernels for small Spark batches 
Extend the existing resident DeepSeek Q8 dispatch through eight rows on GB10. Keep GLM and other devices on their current paths. Check each batch against separate vector calls, including reused scratch and graph replay. The combined DSpark coding workload improves beyond the aligned expert dispatch; ordinary 100-case continuation scores remain unchanged.
 2026-09-06 00:42:59 +02:00
Salvatore Sanfilippo 45f54e9d11 Use aligned expert vector kernels for small Spark batches 
Reuse the existing fused IQ2 gate/up and Q2 down paths for up to eight resident DeepSeek rows on GB10. Keep GLM, streaming and multi-GPU dispatch unchanged. Test overlapping experts, distinct gate/up weights, invalid IDs and exact agreement with separate rows; CUDA memory, race and synchronization checks pass.
 2026-09-06 00:40:24 +02:00
Salvatore Sanfilippo 558cd34311 Restore text-only agent tool results for DeepSeek 
Use each model family’s normal message formatter when a multimodal observation contains no images. Non-vision Flash and PRO no longer reject every tool result as an apparent context overflow, and text results retain their native wrapper and escaping. Add model-free Flash/PRO/GLM regression coverage; CUDA coding/edit tasks pass with ordinary opportunistic and exact DSpark sampling.
 2026-09-06 00:36:26 +02:00
Salvatore Sanfilippo f6c169f825 Retain all five CUDA DSpark partial prefixes 
Save the fifth compressor frontier so accepting five of six verified rows does not replay the last token. Size verifier test sessions to their prompts and continuations instead of allocating two unnecessary 32K contexts. The prefix and serial-target oracles pass across continued contexts through 8K.
 2026-09-05 23:48:13 +02:00
Salvatore Sanfilippo df80a87da1 Batch DSpark seed tokens on resident DGX Spark 
Use the existing seed verifier for resident IQ2/Q2 Spark sessions. Avoid duplicate declined proposals, check short drafts after normal seed decode, and exclude the seed from scheduler acceptance. Pause longer for unproductive prose, but retry sooner after useful long drafts. Keep exact sampling, streaming and other hardware defaults unchanged.
 2026-09-05 23:48:13 +02:00
Salvatore Sanfilippo 9ab705347c Record long GLM quality controls and Spark decode gains 
Replace the expired-credential note with completed Z.AI FP8 long-context comparisons. Record matched Flash-0731 timing medians, exact before/after quality, scratch graph tests and warning-free backend builds.
 2026-09-05 20:55:01 +02:00
Salvatore Sanfilippo 4a14eb83ba Remove unused CUDA and ROCm Q8 matvec kernels 
Delete the two obsolete kernels that quantized activations separately for every output row. Neither has callers or a testing role; the active prequantized and reference paths remain unchanged. CUDA sm121a and ROCm gfx1151 builds pass without warnings.
 2026-09-05 20:53:22 +02:00
Salvatore Sanfilippo b42682f12f Reuse aligned Q8 scratch during CUDA decoding 
Avoid per-call pool allocations in dense and paired aligned Q8 kernels. Reuse the existing device-checked buffer, retaining the pool fallback for missing or undersized scratch.

Three interleaved Spark Flash Q2 runs improve 2K/4K/8K steady decode from 17.90/15.23/15.01 to 19.25/16.22/15.97 t/s. Prefill is unchanged and full prefill and post-decode logits match exactly. Add scratch fallback and graph replay tests.
 2026-09-05 20:36:08 +02:00
Salvatore Sanfilippo 189ff0b302 Add long GLM FP8 references with reasoning-aware scoring 2026-09-05 20:26:22 +02:00
Salvatore Sanfilippo eb847ac568 Fix official-score logprob parsing for null API responses 2026-09-05 20:03:38 +02:00
Salvatore Sanfilippo f4d03f6cf9 Record backend QA results and preserve failure diagnostics 
Keep failed DSpark command diagnostics and report snapshot prompt/context settings. Extend release checks for attention boundaries, usable ROCm memory, repeated vision embeddings and bounded speculative replay.

Record focused GLM and Flash quality, SSD, batching and performance results, including the cold-prefill cost and remaining long-context API and release-matrix gaps.
 2026-09-05 19:42:02 +02:00
Salvatore Sanfilippo 96f547a245 Remove unused ROCm kernels and helpers 
Drop unreferenced Q8, attention, MoE and GLM projection implementations while retaining active fallbacks, diagnostics and performance tests.
 2026-09-05 19:42:01 +02:00
Salvatore Sanfilippo 0e9cc2d43e Fix GLM attention masking and reductions across GPU backends 
Use the model-defined 2051-token dense boundary, handle padded selections, and synchronize reused reduction storage in Metal, CUDA and ROCm, including vision normalization.

Enable bounded ROCm selected-attention GEMM for GLM 5.3 without RoPE. Add independent attention references, boundary tests and finite, repeatable vision checks. Focused Strix continued prefill improves from about 21 to 68 tokens per second; stored official scores are unchanged.
 2026-09-05 19:42:01 +02:00
Salvatore Sanfilippo 9d9e1296d2 Budget ROCm allocations against usable pinned memory 
Exclude CMA pages from host availability, guard pinned and integrated-GPU allocations, and reserve space for the OS and required runtime buffers. Defer optional Q8 expansion and avoid wasting large model arenas.

Add memory parsing and allocation-admission tests. Focused Strix checks cover matched Flash/DSpark, resident GLM and automatic, oversized and one-expert SSD cache budgets.
 2026-09-05 19:41:28 +02:00
Salvatore Sanfilippo ade4d623f9 Fix Linux builds after Metal TP changes 
Guard Apple-only TP hooks and verbs-only timing variables. Link standalone CUDA and ROCm tests with their image and MMQ dependencies.
 2026-09-05 19:40:33 +02:00
Salvatore Sanfilippo 6c25bba129 Put common getting-started instructions in the README 
Keep platform-specific details in their guides, update the cross-links, and remove the release QA matrix link from the README.
 2026-09-05 19:38:18 +02:00
Salvatore Sanfilippo 2f0bc1e3bf Split README into platform setup and focused reference guides 
Keep the introduction and everyday usage in the README, with setup guides near the top and detailed references at the end. Document Ada/L40S support explicitly and replace stale model, memory, caching and performance claims.

Retain the root Strix Halo page as a link to its new guide.
 2026-09-05 14:29:52 +02:00
Salvatore Sanfilippo 233eeb85a7 Fix session rollback, TP restore and GLM memory accounting 
Restore GLM recurrent and pooled-indexer state after rejected speculation. Invalidate unsafe DeepSeek rewinds and rebuild retained agent prefixes with their images before decoding again.

Rebuild saved TP token prefixes on both ranks. Give GLM speculative cycles an explicit command and shared token limit, with protocol version 10 requiring peers to update together.

Account for live and planned Metal GLM sessions when reserving graph and streaming memory. Add model-free state and protocol tests plus a small GPU rollback regression.
 2026-09-05 14:29:18 +02:00
Salvatore Sanfilippo 932896fd0a Fix standalone Metal test linkage 
Link the KDA and MXFP4 Metal tests with the image helpers required by the Metal backend.
 2026-09-05 14:28:35 +02:00
Salvatore Sanfilippo bbec2006f8 Finish deferred Metal TP residuals before mixed-batch output heads 
The mixed prefill/decode path omitted the final deferred residual expansion used by ordinary TP decode. Finish each session head before another session can reuse transient exchange buffers.

The existing full-logit oracle reproduced a 23.7409 maximum difference on the baseline despite identical argmax tokens. Two, four and six MXFP4 TP sessions now match serial logits exactly, including continued prefill; mixed Q2/Q4 also passes with six sessions in TP and on one host. Document the required regression checks.
 2026-09-05 03:25:06 -07:00
Salvatore Sanfilippo 5806d0994d Batch Metal TP DSpark experts and reuse MXFP4 prefill activations 
Use the M5 seed schedule under TP and evaluate verifier rows together instead of running each row separately. Reuse the existing static MXFP4 kernels with row-aware addressing, and enable packed MXFP4 activations for two-rank M5 prefill.

Keep streaming, exact-sampling scheduling and unmeasured device defaults unchanged. Add synthetic ownership/static-shape checks and a physical TP speculative oracle. Record focused quality, speed comparisons and transport caveats in the QA guide.
 2026-09-05 03:24:49 -07:00
Salvatore Sanfilippo be76012927 Batch confident DSpark seeds on M5 and back off on weak drafts 2026-09-04 18:52:56 -07:00
Salvatore Sanfilippo 8d309ab062 Capture partial cache frontiers for six-token DSpark verification 2026-09-04 18:07:16 -07:00
Salvatore Sanfilippo 8fcd61d6bf Fix Metal scratch size for double-buffered dense kernels 2026-09-04 17:51:24 -07:00
Salvatore Sanfilippo 2ea3839ef3 Reuse routed activations in Metal tensor prefill 2026-09-04 17:33:17 -07:00
Salvatore Sanfilippo 29aa7267af Extend ds4-eval with hard capability suite 2026-09-03 15:28:24 -07:00
Mikhail Korobov 469ed2f4b5 Merge upstream/main: fold multimodal live-KV hardening into the reuse probe 
Upstream 98b9135 ('server: reuse exact multimodal live KV safely') and
b0982a1 ('Harden multimodal cache fallback') restructured the execution
ladder this PR replaced with slot_probe_reuse().  Resolution keeps the
probe architecture and moves the new safety conditions into it, so
execution and routing still cannot disagree:

- The live_vision_match gate (image fingerprints are invisible to text/
  token prefixes) is now checked once at the top of
  slot_probe_reuse_locked(); job_slot_score inherits it, subsuming
  upstream's -1 vision guard.
- The GLM memory-rewind materialization re-validates the frontier after
  ds4_session_rewind() (common prefix lands on the rewind target and the
  vision state still matches) and falls back to disk cache / cold prefill
  when validation fails, as upstream's ladder did.
- Upstream's ladder reorder (memory-text before disk cache) is subsumed
  by the probe's tier order; the deleted ladder helpers stay deleted.
- Tests from both sides kept: the slot probe/routing tests and upstream's
  test_multimodal_prefill_resume_frontier.
 2026-09-03 12:27:00 +05:00
Salvatore Sanfilippo b0a147a7fb Add conversation prefix conditioning 2026-09-03 01:25:25 +02:00
antirez 5d41a80fd2 Remove dead Metal TP experiments 2026-09-02 19:39:22 +02:00
antirez c185c7aabf Fix Metal TP ARC cleanup warnings 2026-09-02 15:12:29 +02:00
antirez 4886c56815 Merge remote-tracking branch 'macm5us/mac-tp-1-sep-2026' 2026-09-02 15:10:31 +02:00
Salvatore SanfilippoandClaude Fable 5.1 e0b8c8fdf9 Remove dead experiments, keep diagnostics 
Dropped the measured-no-benefit experiments: DS4_TP_FAKE_BIG_GATE (timing
diagnostic with garbage results), DS4_TP_KICK_SETTLE_US, the FFN sub-gate
row halves (DS4_TP_FFN_SUBGATE, byte-exact but net-negative), the fused
DSpark seed batch under TP (DS4_DSPARK_SEED_BATCH_TP, slower), and
DS4_METAL_MM_ID_MIN_ROWS (grouped GEMM below 32 rows, slower).  Kept:
ds4_gpu_tp_big_gate_kick_flush (used by the attention sub-gate option),
the worker-side prefix commit hardening, and all env-gated diagnostics
(DS4_METAL_CB_TIMES, DS4_TP_BIG_GATE_DEBUG, DS4_DSPARK_CYCLE_TRACE,
DS4_BENCH_SPEC_TRACE, DS4_DSPARK_EXPERT_OVERLAP, batch-gate profile,
tp_ffn_recv/tp_ffn_sum dumps).

Validated: frontier hashes e45a0e5953cb0e64 (2048) and 8f0470ff68881053
(128) unchanged; DSpark verify window ok.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 06:04:16 -07:00
Salvatore SanfilippoandClaude Fable 5.1 b941e98aad README: measured DSpark numbers on a single M5 Max with the resident Q2 model 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 05:23:31 -07:00
Salvatore SanfilippoandClaude Fable 5.1 cffcd64ae9 DSpark diagnostics: DS4_DSPARK_EXPERT_OVERLAP (distinct experts per verify block), DS4_METAL_MM_ID_MIN_ROWS (grouped-GEMM row threshold experiment), cycle trace requires stats timing 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 05:22:28 -07:00
Salvatore SanfilippoandClaude Fable 5.1 7fc640c1bb README: DSpark under two-Mac TP is slower than plain decode 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 05:17:19 -07:00
Salvatore SanfilippoandClaude Fable 5.1 b65ad7dfbc Metal TP DSpark: pre-posted RDMA verify window, seed-batch under TP, tracing 
- ds4_tp.c: ds4_tp_batch_block_begin/end arm a receive window for a whole
  speculative verify block (one batch gate per layer): the decode window is
  drained once, the first layer's receives are posted behind a single
  control-byte barrier, and every gate posts the next layer's receives
  before sending its own rows, so no send can reach an unposted queue and
  no per-gate TCP handshake is needed (the old path cost two round trips
  per gate).  Both ranks call it from the verify path (ds4.c, via the TP
  context stored at ds4_engine_tp_bind).  DS4_TP_DISABLE_VERIFY_WINDOW=1
  keeps the old per-gate protocol.  -4 to -5 ms per verify block.
- ds4.c: the worker's prefix commit no longer tries to preserve a DSpark
  capture it never has (it fails with DS4_DSPARK_SEED_BATCH set); the fused
  seed batch can be enabled under TP with DS4_DSPARK_SEED_BATCH_TP=1
  (experiment: on this pair it verifies more rows per cycle and is slower).
- Diagnostics: DS4_DSPARK_CYCLE_TRACE=1 (per-cycle propose/verify/commit
  timing), DS4_BENCH_SPEC_TRACE=1 (bench per-call cost by accepted count),
  DS4_TP_GATE_PROFILE now also reports verify-block batch gates (arrival
  wait, exchange, release per gate).

Measured on the M5 Max pair (README prompt, 128 context): a verify block
costs ~27 ms + ~5.7 ms per draft row of GPU time (each row re-streams its
own routed experts at bandwidth) plus ~3 ms of gates; DSpark decode stays
at ~35 t/s vs ~48 t/s plain on prose (24-32% draft acceptance).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 05:15:45 -07:00
Salvatore SanfilippoandClaude Fable 5.1 daa0324569 README: two-Mac prefill numbers after the TensorOps routed GEMM and queue keepalive 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 04:41:57 -07:00
Salvatore SanfilippoandClaude Fable 5.1 a06a1bed27 Metal TP prefill: pipelined big-gate kicks (opt-in FFN sub-gate), diagnostics 
- ds4_gpu_tp_big_gate_kick_flush: a kick that ends the command buffer so
  the shared-event signal is published while the caller keeps encoding.
  The service thread then waits for that buffer's completion before it
  reads the payload: the event fires when the GPU is done, but the stores
  only land for the CPU at buffer completion (measured: the first ~10 MB
  of the second half's rows were stale without the wait).  The attention
  sub-gate pipeline (DS4_TP_SUBGATE_PIPELINE=1) now uses it too.
- FFN sub-gate (DS4_TP_FFN_SUBGATE=1, off by default): the routed partial
  is produced in two row halves with this rank's shared-expert rows folded
  per half before each kick, so the exchanged rows are what the peer adds.
  Byte-identical to the single gate but net-negative on the M5 Max pair
  (781 vs ~850 t/s at 2048 tokens: two 1024-row routed calls cost ~27%
  more kernel time than one 2048-row call).
- Diagnostics: DS4_TP_BIG_GATE_DEBUG=2 logs every big gate (period since
  the previous release, handshake, transfer); DS4_TP_FAKE_BIG_GATE=1 skips
  the exchange for a compute-only timing bound (results are garbage);
  DS4_TP_KICK_SETTLE_US adds a delay after the kick-buffer wait;
  tp_ffn_recv / tp_ffn_sum tensor dumps after the FFN all-reduce.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 04:40:58 -07:00
Salvatore SanfilippoandClaude Fable 5.1 8cbdd7624f Metal prefill: expand the compressor position table with one tiled-row kernel 
ds4_gpu_compressor_store_batch_tensor and the score-with-APE path used to
encode one f16->f32 copy per ratio-row segment (n_tokens/4 single-threadgroup
encoders per layer: 21.5k of the 25k encoders of a 2048-token prefill).  A
tiled-row kernel (kernel_cpy_tile_rows_f16_f32 / _f32_f32) writes the whole
batch in one dispatch with identical values.  README: updated two-Mac numbers
and the new opt-out env vars.

Two-Mac TP prefill, byte-exact vs the TensorOps-MoE hashes: 2048 tokens
785 -> 830 t/s, 128 tokens 242 -> 253 t/s.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 04:13:29 -07:00
Salvatore SanfilippoandClaude Fable 5.1 22402728c3 Metal prefill: MXFP4 routed experts on the TensorOps (nax) mm_id kernel 
The routed-expert prefill GEMMs for MXFP4 experts now use the existing
kernel_mul_mm_id_mpp TensorOps kernel (new instantiations
kernel_mul_mm_id_mxfp4_f32_mpp / _f16_mpp), as IQ2_XXS and Q2_K experts
already did.  Same half weights x half activations with f32 accumulation
as the simdgroup kernel; the accumulation order differs, so frontier
logits are no longer byte-identical to the previous kernel (2048-token
prefill on the two-Mac TP pair: mean |dlogit| 0.17, max 1.48, same
argmax, 9/10 top-10 overlap, KL 1e-4).  DS4_METAL_DISABLE_MXFP4_MM_ID_MPP=1
restores the simdgroup kernels.

Two-Mac TP prefill: 2048 tokens 673 -> 785 t/s, 128 tokens 164 -> 242 t/s.
Decode unchanged.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 04:07:56 -07:00
Salvatore SanfilippoandClaude Fable 5.1 e2c39074f4 Metal TP prefill: staged windowed RDMA big gates, main-queue idle keepalive 
- ds4_tp.c: the prefill big-gate exchange now moves the partial through the
  slab batch regions in windows of up to 256 16 KB messages, receives posted
  in chains of 64 behind a one-byte control barrier, sends in chains of 64
  with at most four signaled chains in flight (32 MB in 4-6 ms instead of
  ~26 ms).  Queue pairs ask for 1024/1024 work requests (fallback 256/64)
  and the device limits are printed at open.  Payload-buffer registration
  stays opt-in (DS4_TP_RDMA_REG_PAYLOAD): Apple's verbs cap a memory region
  at 16 MB and registering Metal buffers gave LOC_PROT_ERR completions.
  DS4_TP_BIG_GATE_DEBUG=1 prints handshake/transfer times every 32 gates.
- ds4_metal.m: a command buffer submitted after the engine's queue idled
  for ~3 s starts 600-800 ms late on M5 Max (measured; other queues do not
  prevent it).  A main-queue keepalive submits one 1 KB kernel per second
  (DS4_METAL_DISABLE_QUEUE_KEEPALIVE=1 opts out) and the sharded load pays
  the first submission up front (ds4_gpu_warm_command_queue).  The big-gate
  arrival spin sleep-polls (DS4_TP_BIG_GATE_SPIN_US).  Diagnostics: blit
  passes in the encoder timeline, DS4_METAL_CB_TIMES prints encode/commit/
  GPU-span per owned command buffer with a backtrace on stalls.

Two-Mac TP prefill, byte-exact frontier logits: 2048 tokens 565 -> 673 t/s,
128 tokens 93 -> 164 t/s; decode unchanged.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 04:02:10 -07:00
antirezandClaude Fable 5.1 1cb18a2c2b README: DeepSeek V4 Flash tensor parallelism across two MacBooks (setup, commands, measured rates, knobs) 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 03:13:44 -07:00
antirezandClaude Fable 5.1 399e48a73b Metal TP: attention split count back to 32; opt-in rank-speed lane bias 
The 8-group attention split adopted after the sweep saved 1 us in the
attention kernels but made the layer's later concurrent FFN block ~50 us
slower (114 -> 162 us; cause not understood, measured by toggling
DS4_METAL_FLASH_NWG on one build), costing ~4 t/s.  Back to 32 groups;
the env override stays for experiments.  With it the default output is
byte-identical to the pre-session reference again, which also validates
the kv-task fold and the two-rank router-select fusion against it.

The two-rank router-select branch no longer sets router_only_done (that
enabled the single-machine parallel FFN path and computed the whole
shared expert on the coordinator); it only skips the plain router matvec
and the standalone select.

Rank-speed lane bias, opt-in (DS4_TP_LANE_BIAS=1): the gate service
probes the peer's RDMA arrival while it spins on its own flag, giving a
signed per-gate skew (peer early > 0) without any protocol change; the
coordinator turns the FFN-gate skew EWMA into a lane bias (default 19
lanes per us of skew, smoothed, +-shared_dim/2, 64-lane steps), carries
it in the EVAL command's reserved word, and both ranks add it to the
shared-expert split.  The logits-half frame carries one extra float (the
sender's FFN-gate wait) for future use.  Measured on this pair: skew 2-5
us, 64-128 lanes moved, no token-rate change (48.0 vs 48.2, 47.0 vs
46.4), so it stays off by default because it makes the output depend on
timing.  DS4_TP_LANE_BIAS_GAIN and DS4_TP_LANE_BIAS_DEBUG for experiments.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 02:25:52 -07:00
antirezandClaude Fable 5.1 5024711a7d Metal TP: kv norm task concurrent with q_b; fused router select for two ranks 
The fused q/kv norm + KV RoPE + FP8 store dispatch ran its two tasks as
two threadgroups and cost 13 us per layer, most of it the serial kv task.
The kernel now takes a task mask; on the two-rank decode path the q task
runs alone (3.8 us) and the kv task is deferred and encoded in a
concurrent encoder together with the q_b projection, whose 17.8 MB stream
hides it completely (the section measures 35.8 us, q_b alone 36.6).  The
KV staging and the raw-only attention path flush a still-pending task
before reading the cache, so every path sees the row.  Same arithmetic,
byte-identical output.  DS4_METAL_DISABLE_KV_NORM_DEFER=1 opts out.

The existing fused router projection + exact select kernel (128
threadgroups, last-arriving threadgroup selects) was gated to
single-machine runs; enable it for two ranks, where it replaces the f16
router matvec (7.6 us) + select (6.0) with one 11.1 us dispatch and the
shared expert keeps running in the concurrent FFN encoder.  Byte-identical.
DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE=1 opts out.

Validated on the M5 Max pair: identical text with every toggle at ctx 128
(256 tokens) and ctx 2050 (partial pad block), prefill logits unchanged;
attention-side command buffer 226.1 -> 214.7 us of GPU compute per layer.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 02:06:00 -07:00
antirezandClaude Fable 5.1 aef6daebbd Metal TP: adaptive flash split count, slice-parallel HC step as experiment 
The vec attention kernel's split-K work group count is a pipeline
constant; make it adaptive (8 groups per head for <= 256 keys, else 32;
DS4_METAL_FLASH_NWG overrides).  Measured at 128 keys: 8 groups -1 us
net, 4 groups +9 us.

Add kernel_dsv4_hc_expand4_mix_norm_slices_f16: the HC expand + gate
matvec + collapse + input norm as 64 slice threadgroups with fixed-order
partial sums and a last-arriver tail, no spin barrier.  Measured 16.7 us
against the six-threadgroup cluster kernel's 14.5 (the cluster kernel is
launch + serial tail bound, not barrier bound), so it stays behind
DS4_METAL_HC_EXPAND_SLICES=1 and the cluster kernel remains the default.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-02 01:46:49 -07:00
antirezandClaude Fable 5.1 313bebc639 TP RDMA: warm-up round with posted-receive barrier, queue pair recreation, diagnostics 
Apple's Thunderbolt RDMA (UC queue pairs, two-sided SEND/RECV) has a
failure mode UC never reports: a send that reaches a queue pair with no
receive posted does not merely get lost, that direction of the pair
stops delivering for good (the pair then reads RTR on query and later
sends never complete).  Runs used to discover this only at the first
bulk prefill round, as "timeout waiting for bulk RDMA round (0/64
recvs, send=0)", typically on the first run after a role swap.

After the ready barrier both sides now run a warm-up round: post one
full-size receive, exchange a "receives posted" control frame
(DS4_TP_FRAME_RDMA_POSTED), then resend a tagged 16 KB message until the
peer confirms receipt over the control channel (DS4_TP_FRAME_RDMA_WARM,
100 ms per attempt).  A failed warm-up recreates the queue pair and
completion queue (new number, new PSN), redoes the info exchange and
warm-up, up to four setups; both sides take the same path because the
status exchange gives them the same outcome.  Bulk-round receive
completions check the byte length, and a bulk timeout reports the queue
pair state and chunk layout.  DS4_TP_DISABLE_RDMA_WARMUP=1 skips the
round.

Measured on the M5 Max pair: with the posted barrier the warm-up passes
on the first attempt in every run (10/10); without the barrier the
warm-up itself raced the peer's send and poisoned pairs.  The bulk
round keeps its pre-existing race (a peer that finished the previous
round earlier can send before our receives are posted); a control-frame
barrier there collides with the eval protocol, so the fix is to pre-post
the next round's receives into a second staging area before sending the
current round.  Noted in misc/MAC_TP_SPEED.md.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 20:29:58 -07:00
antirezandClaude Fable 5.1 23bbbf36e3 Metal TP: touch-table helper, four-row Q8 matvec experiment knob, fold diagnostics 
Factor the one-dispatch touch-table encode out of the gate prefetch
(ds4_gpu_encode_touch_table).  Add kernel_mul_mv_q8_0_f32_r4 and its
flag-folded twin (four rows per threadgroup, bit-exact with the two-row
kernels) behind DS4_METAL_Q8_MV_NR0=4 for the pass-through dispatch
sites; measured slower on the M5 Max pair (48.1 vs 49.3, 45.1 vs 48.4
t/s), so the default stays at two rows.  DS4_TP_FOLD_DEBUG=1 prints the
first flag-fold request/take decisions.

An attention-window prefetch (touching this rank's out_a half in a
concurrent encoder with the flash-attention vec kernel) was measured and
not kept: the projection got 10 us faster but the latency-bound vec
kernel slowed by the touch's duration under saturated DRAM, token rate
unchanged.  Notes in misc/MAC_TP_SPEED.md.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 19:10:31 -07:00
antirezandClaude Fable 5.1 75787d9bb9 Metal TP: fold the gate flag into the payload producers; one-dispatch prefetch 
The checked poll-gate flag was its own kernel after each partial (4-7 us
per gate).  The payload producers now publish it themselves from the
last-arriving threadgroup: kernel_dsv4_add2_f32_tp_flag_checked for the
FFN gate (local shared+routed sum into the slab slot) and
kernel_dsv4_mul_mv_q8_0_f32_tp_flag_checked for the attention gate (the
attention-output K-slice matvec).  Each threadgroup adds the integer sum
of the words it stored to a per-slot accumulator, fences, then arrives;
the last arriver exchanges the accumulator, resets both counters, and
stores the checksum and the flag.  Integer sums are order independent,
so the checksum equals the standalone kernel's over the same payload and
the stored values are the plain kernels' values: text byte-identical.
ds4_gpu_tp_gate_encode skips its flag kernel when the producer
prepublished the sequence number; ds4_gpu_add_tensor_tp_flag and
ds4_gpu_tp_flag_fold_request (taken by the K-slice matvec when its
output is the slot) are the host entry points.  DS4_TP_DISABLE_FLAG_FOLD=1
opts out.  Timeline: fused add+flag 3.2 us vs add 2.1 + flag 4.0.

Gate-time prefetch was one dispatch per plan range (up to 7) at ~3 us
fixed cost each, 16-22 us per gate and exposed whenever the peer's
partial was already there.  It is now one dispatch over a table of up to
8 bindless views (gpuAddress + useResources), threads assigned to ranges
by cumulative line counts (kernel_touch_u8_stride_table).  The
attention-gate plan lists the other shared half last for the dynamic
shared split.  DS4_TP_PREFETCH_MARGIN_US / DS4_TP_PREFETCH_BYTES_PER_US
override the sizing constants for experiments.

Measured on two M5 Max over TB5 RDMA (README prompt, 256 tokens,
interleaved pairs): 47.35 vs 46.78 and 46.70 vs 45.86 t/s with the folds
off; eight repeated default runs identical in text and frontier logits.
Encoder timeline: attention-side command buffer 233.5 -> 226.1 us,
FFN-side 173.1 -> 168.6 us of GPU compute.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 18:48:54 -07:00
antirezandClaude Fable 5.1 c69e57fe91 Metal TP: GPU-decided shared-expert lane split for rank balancing 
Routed experts are owned by expert id, so each rank streams 0..6 of the
6 selected experts per token while the shared expert was split in fixed
halves; the rank owning more selected experts was the critical path of
the layer and the other rank idled at the gate.  The new split kernels
(kernel_dsv4_shared_gate_up_swiglu_q8_0_split and
kernel_dsv4_shared_down_q8_0_split) read the selected ids and take a
complementary lane range per rank sized to balance bytes: lanes0 =
shared_dim/2 + (n1 - n0) * shift * shared_dim with shift = routed expert
bytes / (2 * shared expert bytes), rounded to 64 lanes.  On the V4 Flash
MXFP4/Q8 file one extra routed expert moves half of the shared expert to
the other rank.  A rank with zero lanes writes a zero partial, so the
local routed+shared add is unchanged.  The kernels keep the args in
constant memory with explicit row/K counts (a by-value args copy cost
+9.8 us per layer in the matvec bodies); with shift 0 they are bit-exact
with the static halves (+1.9 us per layer).  The attention-gate prefetch
plan lists the other shared half last so a rank with a long measured
wait streams it in.

Default for two-rank Metal TP decode.  DS4_TP_STATIC_SHARED_SPLIT=1
restores the fixed halves; DS4_TP_SHARED_SPLIT_SHIFT_Q16 overrides the
shift (0 = static halves through the split kernels, for exactness
checks).  Measured on two M5 Max over TB5 RDMA (README prompt, 256
tokens, interleaved pairs): ctx 128 46.6-47.4 vs 45.2-46.7 t/s, ctx 2048
45.6-46.2 vs 45.1 t/s; both ranks wait less at FFN gates.  Output is
deterministic and differs from the static split only by floating-point
grouping of the partial sums.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 18:19:11 -07:00
antirezandClaude Fable 5.1 6cfc278706 TP: RDMA gate post/peer-wait profile, MXFP4 MoE microbench 
- ds4_tp.c: under DS4_TP_GATE_PROFILE print per gate kind the RDMA post cost
  and the wait for the peer's partial (the attention gate sits at the
  ~11 us wire floor; FFN-gate averages include prefill batch gates).
- misc/mxfp4_moe_bench.m: standalone microbench of the exact-shape TP MXFP4
  pair-SwiGLU and sum6 kernels on rotating expert copies with a warm-up
  pass; COMPACT=1 measures the unowned-slot launch cost.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 17:32:46 -07:00
antirezandClaude Fable 5.1 1a976b7dbf Metal TP: poll gates, checksum-verified flags, HC expand fusion, gate prefetch 
Two-Mac DeepSeek V4 Flash tensor-parallel decode on M5 Max pairs.

- Poll gates (default; DS4_TP_DISABLE_POLL_GATES opts out): each flag gate
  ends the command buffer and the next queued buffer opens with a bounded
  fresh-line poll released by the service thread, instead of parking on a
  shared event whose CPU->GPU wake-up costs 27-43 us. Ring of 16 regions,
  ~600 ms budget, timeout reported at the next gate.
- Checksum-verified flags: kernel_dsv4_tp_flag_set_checked publishes an
  exact integer checksum of the 16 KiB partial; the service thread verifies
  the payload in memory before the RDMA post (the flag line can land ~35 us
  before the buffer completes; unverified sends were nondeterministic).
- HC post/expand folded into the compound HC producer with a device-scope
  barrier across its six threadgroups; the engine defers the post-gate
  expand and flushes it before any other reader
  (DS4_METAL_DISABLE_HC_EXPAND_PRODUCER_FUSE opts out).
- Gate-time prefetch: the next phase's weights (rank slices, priority
  order, cut at this rank's measured peer-wait window) are touched into the
  cache before the poll (DS4_TP_DISABLE_GATE_PREFETCH opts out).
- Concurrent routed/shared FFN branches, safe decode command-buffer split,
  exact-shape TP MXFP4 static kernels, DSpark direct partial commit, and
  split gate profiling from the morning session.
- Diagnostics: DS4_METAL_ENCODER_TIMELINE per-pass GPU timeline,
  DS4_METAL_CB_TIMES per-command-buffer driver/GPU timing.

README ctx 128, 256 greedy tokens: 45.5 vs 43.3 t/s (event gates, same
build); frontier logits and generated text byte-identical throughout.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WjmBSeiH6tPDtJuiQNFSpq
 2026-09-01 17:26:50 -07:00
antirez 110afdd888 Use random completion response IDs 2026-09-01 17:32:21 +02:00
antirez b0982a1b4e Harden multimodal cache fallback 2026-09-01 13:33:00 +02:00
Jordi Posthumus 98b91351ad server: reuse exact multimodal live KV safely 2026-09-01 07:46:20 -03:00
antirez 3da2da96b2 Document DeepSeek Vision-Exp usage 2026-09-01 12:23:44 +02:00
antirez f9c459ce2f Fix ROCm MMQ scratch lifetime and tail reads 2026-09-01 03:26:23 +02:00
antirez 200b9f45bf Fix CUDA vision build 2026-09-01 00:32:06 +02:00
antirez 409080a623 Fix ROCm vision build 2026-09-01 00:23:37 +02:00
antirez e976a54331 Add DeepSeek Vision-Exp support to ROCm 2026-09-01 00:22:23 +02:00
Salvatore Sanfilippo fc8bf3c39c Add DeepSeek V4 Flash vision support 2026-09-01 00:11:58 +02:00
antirez ad4ac91a39 Speed up GLM 5.3 ROCm prefill 2026-08-31 23:58:50 +02:00
antirez ab06d196da Speed up GLM 5.3 decode on ROCm 2026-08-31 20:14:12 +02:00
antirez ec7642cdd9 README: GLM 5.3 is supported as well 2026-08-31 15:36:35 +02:00
antirez 4b321cf1ec Make ROCm release builds warning-free 2026-08-31 15:12:02 +02:00
antirez f5ee85b341 Make GLM SSD streaming memory budgets adaptive 
Measure the runtime decode map before selecting the global Metal path, account its mapped bytes once, and fall back to per-layer mapping when it cannot fit. Treat explicit cache sizes as performance targets that may be reduced after model, graph, context, and platform reserves are known. Share the final guard-budget calculation between fitting and refusal on Metal and ROCm.
 2026-08-31 14:54:16 +02:00
Salvatore Sanfilippo e4593c374f Fix GLM streaming resident layer maps 2026-08-31 13:45:02 +02:00
antirez 5da32061ba Merge GLM 5.3 SSD streaming and full checkpoint support 
# Conflicts:
#	ds4.c
#	tests/test_glm53_kda.c
 2026-08-31 13:25:05 +02:00
Salvatore Sanfilippo dcc70b41ef Generalize GLM quantization for full checkpoints 2026-08-31 03:33:22 -07:00
Salvatore Sanfilippo 81c8c39a5c Support FP32 GLM compact attention caches 2026-08-31 03:33:22 -07:00
Salvatore Sanfilippo c8b6e4f4d8 Improve GLM routed imatrix coverage 2026-08-31 03:33:21 -07:00
Salvatore Sanfilippo 01f3316c9c Fix full GLM SSD streaming on Metal 2026-08-31 03:33:00 -07:00
antirez a547fdff4a Merge pull request #887 from kyuz0/perf/rocm-gfx1151-mmq-kernel-lab 
# Conflicts:
#	Makefile
#	ds4.c
#	ds4_bench.c
#	ds4_help.c
#	tests/test_gpu_args_cli.sh
 2026-08-31 11:48:21 +02:00
antirez b1b4ea0364 cuda: accelerate GLM prefill on DGX Spark 2026-08-30 23:47:28 +02:00
Salvatore Sanfilippo 6cf658a4da Speed up GLM 5.3 BF16 decode on M3 Ultra 2026-08-30 15:19:14 +02:00
Salvatore Sanfilippo 10badf696c Quantize GLM 5.3 Q4 dense weights to Q8 2026-08-30 15:19:00 +02:00
Donato Capitella 9254c3daf4 tests: honor no-GPU ROCm test builds 2026-08-30 12:43:58 +01:00
Donato Capitella e1147cdcbd server: support fixed-length greedy decode 2026-08-30 12:40:52 +01:00
antirez 87495f6cc9 Add live steering controls 2026-08-30 13:33:09 +02:00
antirez 639d4eb8d1 Fix GLM 5.3 continued prefill 2026-08-30 13:08:26 +02:00
Donato Capitella 742adcaf14 docs: update Strix Halo ROCm setup 2026-08-29 21:51:45 +01:00
kyuz0 c012d5cedf fix(rocm): stabilize gfx1151 MMQ routing maps 2026-08-29 21:51:45 +01:00
Donato Capitella 17c2855987 Revert "bench: add untimed context baseline" 
This reverts commit e344bbef57b4713136ab703af392017ac80c7927.
 2026-08-29 21:51:45 +01:00
Donato Capitella 985cacd07b bench: add untimed context baseline 2026-08-29 21:51:45 +01:00
Donato Capitella d208b2b762 rocm: generalize validated gfx1151 DSpark defaults 2026-08-29 21:51:45 +01:00
Donato Capitella f9e4c0fe2c rocm: accelerate gfx1151 DSpark verification 2026-08-29 21:51:45 +01:00
Donato Capitella 9dfb291b0b rocm: tile large IQ2 MMQ prefill batches 2026-08-29 21:51:45 +01:00
Donato Capitella e532136631 fix(rocm): adapt DSpark scheduling on gfx1151 2026-08-29 21:51:45 +01:00
Donato Capitella 24665f8192 fix(rocm): preserve target cache with DSpark support 2026-08-29 21:51:45 +01:00
Donato Capitella e73b367db9 feat(bench): add basic DSpark throughput support 2026-08-29 21:51:45 +01:00
Donato Capitella d94a54e612 perf(rocm): retune DSpark F16 solutions for rocBLAS 5.6 2026-08-29 21:51:45 +01:00
Donato Capitella 0efd822b9b fix(rocm): stop retrying rejected DSpark GEMM solutions 2026-08-29 21:51:45 +01:00
Donato Capitella e8686b9d04 perf(dspark): fuse gfx1151 seed and draft verification 2026-08-29 21:51:45 +01:00
Donato Capitella 9ca8642d01 perf(dspark): skip zero-threshold confidence probe 2026-08-29 21:51:45 +01:00
Donato Capitella 227fa190aa perf(rocm): accelerate gfx1151 DSpark verification 2026-08-29 21:51:45 +01:00
Donato Capitella 307dc093d9 perf(rocm): specialize tiny Q8 verifier batches 2026-08-29 21:51:45 +01:00
Donato Capitella a4e015e004 perf(rocm): accelerate tiny-batch DSpark verification 2026-08-29 21:51:45 +01:00
Donato Capitella ab3346a695 bench: finalize gfx1151 prefill result 2026-08-29 21:51:45 +01:00
Donato Capitella 1ef0680b79 perf(rocm): default to tuned gfx1151 prefill paths 2026-08-29 21:51:45 +01:00
Donato Capitella 49a30025d6 perf(rocm): reduce gfx1151 prefill staging overhead 2026-08-29 21:51:45 +01:00
Donato Capitella b7f2f2d37f bench: close gfx1151 attention key tile sweep 2026-08-29 21:51:45 +01:00
Donato Capitella 86128aae00 bench: record 290.79 tok/s gfx1151 prefill lead 2026-08-29 21:51:45 +01:00
Donato Capitella 4e23cf1bcc rocm: use 80-key WMMA attention tiles on gfx1151 2026-08-29 21:51:45 +01:00
Donato Capitella becf9f96d8 bench: record large F16 gfx1151 win 2026-08-29 21:51:45 +01:00
Donato Capitella 8c1629c95c perf(rocm): accelerate large F16 prefill projection 2026-08-29 21:51:45 +01:00
Donato Capitella 491e7e3d3f perf(rocm): remove rejected paired IQ2 kernel 2026-08-29 21:51:45 +01:00
Donato Capitella 9fd5e399e9 bench: profile paired gfx1151 IQ2 kernel 2026-08-29 21:51:45 +01:00
Donato Capitella c40e8a8e09 perf(rocm): pair IQ2 gate and up tiles 2026-08-29 21:51:45 +01:00
Donato Capitella 11e1eb31b0 bench: record compact gfx1151 IQ2 launch 2026-08-29 21:51:45 +01:00
Donato Capitella 4bf7c47b71 perf(rocm): compact IQ2 MoE tile launches 2026-08-29 21:51:45 +01:00
Donato Capitella 540ccbafa7 bench: record 277 tok/s gfx1151 lead 2026-08-29 21:51:45 +01:00
Donato Capitella 9ab922f50e perf(rocm): tighten IQ2 MoE launch bound 2026-08-29 21:51:45 +01:00
Donato Capitella 45dbb31505 docs: record gfx1151 small-M F16 results 2026-08-29 21:51:45 +01:00
Donato Capitella 6bdc68bccf perf(rocm): specialize small-M F16 prefill GEMMs 2026-08-29 21:51:45 +01:00
kyuz0 b0d80b13c8 bench: record 263 tok/s gfx1151 lead 2026-08-29 21:51:45 +01:00
kyuz0 e4cd58705b rocm: specialize tiny-m f16 prefill gemm 2026-08-29 21:51:45 +01:00
kyuz0 d8d5fa679a bench: record 256 tok/s gfx1151 lead 2026-08-29 21:51:45 +01:00
kyuz0 4d89fe9e84 rocm: vectorize indexed attention kv staging 2026-08-29 21:51:45 +01:00
kyuz0 df1269414d bench: close unsafe f16 kv mirror experiment 2026-08-29 21:51:45 +01:00
kyuz0 64b22d74f1 bench: map remaining gfx1151 prefill pools 2026-08-29 21:51:45 +01:00
kyuz0 300c5814b7 bench: record indexed attention tile rejection 2026-08-29 21:51:45 +01:00
kyuz0 68e0ca1e44 bench: record 248 tok/s prefill lead 2026-08-29 21:51:45 +01:00
kyuz0 e43738a0ab rocm: specialize attention output B for gfx1151 2026-08-29 21:51:45 +01:00
kyuz0 d104a31b14 bench: map attention output GEMM 2026-08-29 21:51:44 +01:00
kyuz0 c2b67652e9 bench: record Q2 n4 rejection 2026-08-29 21:51:44 +01:00
kyuz0 f6022ee42d bench: record IQ2 sign-mask optimization 2026-08-29 21:51:44 +01:00
kyuz0 0610a6577e rocm: use cached IQ2 sign masks in MMQ 2026-08-29 21:51:44 +01:00
kyuz0 dd356c3fad bench: record gfx1151 prefill optimization results 2026-08-29 21:51:44 +01:00
kyuz0 7470dd60a7 build: link MMQ objects in ROCm regression target 2026-08-29 21:51:44 +01:00
kyuz0 c53b2ef6fa rocm: accelerate sparse attention with WMMA 2026-08-29 21:51:44 +01:00
kyuz0 b93fb08937 rocm: tile Q8 prefill WMMA by output shape 2026-08-29 21:51:44 +01:00
kyuz0 9933e01f50 rocm: wire tuned MMQ IQ2 prefill experiment 2026-08-29 21:51:44 +01:00
kyuz0 61e9b991c3 rocm: add gfx1151 MMQ kernel lab 2026-08-29 21:51:44 +01:00
kyuz0 e6e7e6a5d7 rocm: reuse LDS output tile in Q2_K MoE down kernel 2026-08-29 21:51:44 +01:00
Salvatore Sanfilippo 8adb22178c Add GLM 5.3 directional steering 2026-08-29 22:31:32 +02:00
antirez 767e517639 Map vision weights directly on integrated CUDA 2026-08-29 17:55:38 +02:00
antirez c6433cb51b Add GLM 5.3 vision on CUDA and ROCm 2026-08-29 17:31:33 +02:00
antirez ab993cb2a7 Accept ROCm backend in ds4-agent 2026-08-29 13:06:04 +02:00
antirez b8cfc1431f Harden ROCm GLM snapshot replay test 2026-08-29 13:05:59 +02:00
antirez 04478e0ff6 Fix streamed terminal output at wrap boundaries 2026-08-29 11:36:22 +02:00
antirez 1df8ff0d9f Strengthen GLM 5.3 vision replay QA 2026-08-29 11:08:52 +02:00
Salvatore Sanfilippo fea665ce46 Fix ROCm GLM 5.3 causal attention race 2026-08-29 09:44:39 +02:00
Salvatore Sanfilippo 0415d3b0d6 Add GLM 5.3 ROCm text inference 2026-08-29 05:47:30 +02:00
antirez 747eafc2c2 Add GLM 5.3 vision quality gates 2026-08-29 02:54:01 +02:00
antirez 9f95d9f5ce Fix GLM 5.3 vision in compact prefill 2026-08-29 02:24:55 +02:00
antirez b0c31af367 Improve GLM 5.3 attention memory and batching 2026-08-29 02:15:31 +02:00
Salvatore Sanfilippo 47713296d5 Add GLM 5.3 Flash vision support 2026-08-28 23:26:29 +02:00
antirez 147109a3cc Add GLM 5.3 Flash inference 
Implement GLM 5.3 Q2 and Q4 inference across Metal and CUDA, including tensor parallelism, RDMA, batching, MTP, snapshots, SSD streaming, and GLM 5.2 shared-path improvements.

Add quantization and validation tools, official quality fixtures, model downloads, regression coverage, release QA, and user documentation.
 2026-08-28 23:26:29 +02:00
esvhd 8db89fe083 download: add GLM 5.3 Flash models (#888) 
Adds Hugging Face download targets for the Q2, Q4_K, and native FP8 artifacts.

Co-authored-by: esvhd <7330385+esvhd@users.noreply.github.com>
 2026-08-28 23:25:07 +02:00
antirez c1d4597a80 qa: update DGX Spark host addresses 2026-08-23 15:26:19 +02:00
antirez c35cf38ba9 models: download the DeepSeek V4 Pro 0813 quant 
Point pro-q2-imatrix at the published dated GGUF and update the README examples and native PRO conversion description.
 2026-08-23 15:05:13 +02:00
antirez 3c63c06c5a qa: add DeepSeek V4 Pro 0813 oracle 
Keep the preview oracle intact, add the 100-case official 0813 fixture, and record the source, conversion, safety, quality, hash, and publication results.
 2026-08-23 15:05:13 +02:00
antirez ed5c1a92c6 gguf: harden DeepSeek PRO quantization 
Load compression ratios from the source checkpoint instead of retaining stale template metadata. Stop Q2_K/Q4_K refinement when an accepted fit collapses the search range.
 2026-08-23 15:05:13 +02:00
antirez abd4ed3885 cli: expose exact MTP sampling mode 2026-08-23 15:05:13 +02:00
antirez 8d683d6616 dspark: add opportunistic temperature sampling 2026-08-23 15:05:13 +02:00
antirez 769a8ba42d dspark: support exact stochastic speculative decoding 2026-08-23 15:05:13 +02:00
Salvatore Sanfilippo ecdea42ee7 Merge pull request #656 from jyatesdotdev/rocm-mxfp4-pr 
Add ROCm MXFP4 routed expert kernels
 2026-08-23 15:03:33 +02:00
Jonathan Yates babca3d0fe Makefile: restore Darwin link guard for ds4_test/ds4_agent_test 2026-08-20 21:23:53 -07:00
Jonathan Yates 7891f278d0 Makefile: ROCm core-test target and portable scorer link 
Add 'make test-rocm', mirroring 'make test' for ROCm-only hosts (the
CUDA-specific sampling/session oracles stay in 'make test'), and link
ds4_test/ds4_agent_test/score_official through DS4_LINK like the main
binaries so they build without nvcc.  score_official now compiles its
translation unit in a separate step with ROCM_HOST_CFLAGS: hipcc's link
driver injects -x c ahead of the object list and would otherwise try to
parse relocatables as C source, and the Strix host objects are built
position-independent so the scorer TU must match.
 2026-08-20 18:45:38 -07:00
Jonathan Yates 47ecd56841 Launch MXFP4 tile kernels on the default stream in this branch 
The r9/r11 campaign trees pass ds4_rocm_stream() to these launches; that
helper belongs to the HIP graph-capture work, which is intentionally not
part of this PR.  It returns the default stream whenever capture is not
armed, so passing 0 here is behaviorally identical for this branch.
 2026-08-20 18:45:13 -07:00
Jonathan Yates affa9c3054 MXFP4 prefill MoE occupancy variants (env-gated, bit-exact, hw-verified) 
The tile8 gate/up kernel is latency-bound, not format-bound: 57 KiB of
staged-activation LDS pins residency to 1-2 blocks/WGP and each block
has only 8 warps to hide the serial block-loop latency.  The down
kernel has the opposite shape - ~917k tiny blocks per layer-chunk at
4096 tokens, each re-staging the same 16.6 KiB of activations.

Three variants, all preserving the exact per-(token,row) accumulation
order (same 8-lane q8_K block split, same dp4a helpers, same
quarter-warp reduction, same epilogue - bit-exact by construction):

- DS4_ROCM_ENABLE_MXFP4_TILE4: 4-token expert tiles for gate/up
  (templated tileN<TILE_TOKENS, QWARPS> kernel) - 29 KiB LDS at 7168
  dims, 4 blocks/WGP.
- DS4_ROCM_ENABLE_MXFP4_ROW64: 512-thread gate/up blocks (64 rows x
  8 tokens against one staged copy) - 2 blocks/WGP = 32 resident warps.
- DS4_ROCM_MXFP4_DOWN_RGROUP=1..8: down kernel loops over consecutive
  32-row groups against one staged activation copy - staging traffic
  and block count divide by R.

tests/test_mxfp4_rocm now snapshots the default path's out tensor at
n=128/512 and memcmp-asserts each variant against it (cross-path
bitwise equality, stronger than the tolerance checks).

Compile-verified on gfx1151; hardware A/B and the acceptance gate run
in the next test window.  The WMMA rewrite is deferred until these
say whether occupancy or format is the binding constraint.
 2026-08-20 18:45:13 -07:00
Jonathan Yates df1a49a50d MXFP4 prefill MoE: L2-friendly grid order + wider A-staging (r9) 
Production kernel trace of the 10.5k-token prefill (rocprofv3 on the
coordinator) showed the MXFP4 routed-MoE tile kernels dominate prefill:
gate/up 247.7 ms and down 174.9 ms per layer-chunk, bound not by
compute (~0.6% of fp32 peak) but by expert-weight DRAM re-reads - every
8-token tile re-reads its experts' full gate+up rows.

Changes (all bit-exact; per-element accumulation order untouched):
- gate/up tile8 kernel grid swapped to (tile, row-block): consecutive
  blocks re-walk the same 32-row weight slice across an expert's token
  tiles, so the slice stays L2-resident.  Production: 247.7 -> 222.3 ms
  per layer-chunk (-10.3%), 10.5k prefill 61.5 s -> 59.9 s.
- A-activation LDS staging threshold raised 16 -> 28 q8_K blocks (the
  7168-dim case; 56.9 KiB dynamic smem with the >48 KiB opt-in).
  Neutral in production (L2 was already absorbing the re-reads), large
  in the isolated bench; staged and unstaged dot helpers are bit-equal
  (integer bsum order is exact; verified in-tree and at 7168 dims).
- down-kernel grid swap measured NEGATIVE (174.9 -> 184.8 ms) and was
  reverted.
- Two experimental kernels parked default-off behind env gates with
  honest numbers: tile32 B-register reuse (DS4_ROCM_ENABLE_MXFP4_TILE32;
  register spills, 2155 ms bench) and row8-ldsB B-staging
  (DS4_ROCM_ENABLE_MXFP4_LDSB; 154 s production - scattered A reads).
- tests/bench_mxfp4_rocm: production-shape MoE bench harness.  NOTE:
  it is a poor proxy for production kernel time (isolated process
  memory/L2 behavior differs ~5x); production traces are the arbiter.
- DS4_ROCM_MOE_PATH_DEBUG: launch-path diagnostics.

Validation: tests/test_mxfp4_rocm PASS (default dims + 7168/top-8
variant), acceptance gate df07199e5a292872 holds, decode 15.38 t/s,
10.5k probe 59.9 s (176 t/s avg).
 2026-08-20 18:45:13 -07:00
Jonathan Yates bfd0682494 Vectorize MXFP4 tile kernel activation traffic and fuse gate/up 
The MXFP4 tile kernels ran at roughly three percent of dp4a peak and
were insensitive to occupancy, scratch, DRAM weight traffic, and chain
depth: the limiter was instruction issue, with four dword activation
loads and loop control paired against every eight dp4a.

Stage activations as 16-byte-aligned quant slices with the Q8_K scales
in a separate array, so the dot helper loads one b128 per half
sub-block instead of four dwords; the layout also drops the unused
32-byte bsums from every staged chunk.  Weight bytes load as one
misaligned b128 per block, the pair loop is a fixed predicated eight,
and the gate/up kernel computes both matrices from a single activation
load through a fused helper, halving its dominant LDS traffic.  Low and
high nibble sums use separate integer accumulators; all integer
reassociation is exact and the outputs stay bit-identical, with the
struct-layout helper retained for the unstaged wide-activation path.

gfx1151, DS4 Flash shape (4096/2048, 256 experts, 6 used), rocprofv3
min-of-3 on the synthetic sweep: gate/up tile kernel 106.0ms -> 50.1ms
(2.11x), down tile kernel 69.6ms -> 34.0ms (2.05x).  Distributed pair
(coordinator 0:21 + worker 22:output, dist chunk 4096): prefill 70.4 ->
120.6 tok/s serial single-chunk (1.71x) and 103.9 -> 176.5 tok/s on a
three-chunk pipelined frontier (1.70x); decode unchanged at 7.7 tok/s.
 2026-08-20 18:45:13 -07:00
Jonathan Yates 5d90ebace6 Test ROCm MXFP4 routed experts at the DS4 Flash shape 
The synthetic MoE test ran at 512/256 dims, which never reaches the
16-chunk activation shape the resident kernels see in production, and
its down-projection reference was built from the analytic mid.  The GPU
quantizes its own gate/up mid before the down stage, and a legitimate
one-LSB rounding flip against the analytic mid moves one down dot by up
to d * 2^(e-127) * 6: at real-shape scales that exceeds any fixed
tolerance and failed tokens=3/4 despite both results being correct.

Default the test to the DS4 Flash routed shape (4096/2048, 256 experts,
6 used) and judge the down stage against a CPU requantization of the
GPU's own mid, so the comparison shares quantized inputs and the
tolerances stay tight enough to catch real kernel bugs.  Expected rows
are cached per pattern and reused after a bitwise mid comparison, which
keeps the CPU reference cost flat across the 512-token sweep.
 2026-08-20 18:45:13 -07:00
Jonathan Yates f52dc1ca91 Keep ROCm MXFP4 tile kernel state in registers and right-size LDS 
The MXFP4 expert-tile kernels built their per-tile pair state with
runtime-count loops, so pair/tok/slot/xqb were dynamically indexed and
spilled: the compiler reported 48 bytes/lane of scratch on the gate/up
kernel, pairing private-memory traffic with the inner dp4a loop.  Setup
and finalize now run a fixed, fully unrolled 8 iterations with a uniform
p < np guard, which lets every array scalarize.  tok/slot are gone
entirely since pair == tok * n_expert + slot already indexes the route
weights.

The gate/up kernel also staged activations into a static [8][16] Q8_K
tile even when xq_blocks is smaller; the buffer is now dynamic shared
memory sized by the launch.

gfx1151 compiler-reported resources for the gate/up tile kernel go from
99 VGPRs / 48 B scratch / 12 waves to 83 VGPRs / no scratch / 16 waves.
On the DS4 Flash shape (xq_blocks=16) the LDS footprint is unchanged and
the kernel stays weight-stream bound, so end-to-end prefill is neutral
there; the change is codegen hygiene plus wins for smaller shapes.
Numerics are identical since the summation order is preserved.
 2026-08-20 18:45:13 -07:00
Jonathan Yates 4ac740ad23 Optimize ROCm MXFP4 decode and tiny batches 
Coalesce MXFP4 weight reads across wave32 lanes and unpack nibbles with register permutations. Keep the gate/up and down reductions register-resident, route two-to-four-token batches through the direct kernels, and extend the ROCm CPU-oracle coverage across that dispatch boundary.
 2026-08-20 18:45:13 -07:00
Jonathan Yates 27180ec19f Document ROCm MXFP4 pipeline validation 2026-08-20 18:45:13 -07:00
Jonathan Yates 39a8f18b03 Add ROCm MXFP4 routed expert kernels 
Implement resident MXFP4 gate/up and down paths for decode and batched routed MoE on ROCm. Add a synthetic CPU-oracle regression covering full, partial, and large expert tiles, report MXFP4 as a 4-bit routed quant, and make Strix host objects position independent for TheRock linking.
 2026-08-20 18:45:13 -07:00
Mikhail Korobov ada805ad59 server: staleness-aware eviction tiers in slot routing 
The reuse-aware router sent a request that reuses nothing to the slot
with the shortest resident checkpoint (-live).  On a long-running
server this leaks: a finished long conversation is never the cheapest
victim, so such checkpoints pile up, the pool of usable slots shrinks,
and subagent/parallel requests end up evicting shorter *active*
conversations - the prefill thrashing this routing set out to fix
comes back.

Split the no-reuse victim selection into staleness tiers:

- empty slots still win (nothing to evict);
- a checkpoint idle past SLOT_STALE_AFTER_SEC (1h) is most likely a
  finished conversation: victims are picked most-idle-first,
  regardless of size, so forgotten sessions get reclaimed;
- fresher checkpoints belong to active conversations: keep the -live
  rule (evict the shortest, lose the least prefilled work).

Staleness only lowers a checkpoint's priority as a *victim*; it never
blocks reuse - the owner of a stale checkpoint who continues the
conversation still pins the slot via the reuse tier (which now scores
the actual reusable prefix instead of the full checkpoint length).

Also log evictions at dispatch time (slot, evicted tokens, idle time,
stale/protected), so routing decisions are visible in the server log
instead of being inferable only from the next request's cache miss.
 2026-08-14 12:26:49 +05:00
Mikhail Korobov 4d7028e066 server: reuse-aware slot routing via shared live-state probe 
With --batched-session N, job_slot_score routed requests to the slot
with the longest raw common token prefix, with no check that the slot's
checkpoint is actually reusable. Parallel subagents sharing a system
prompt were attracted to the slot holding the main agent's long context,
evicting it with zero reuse benefit and forcing a full re-prefill when
main resumed (masked only by --kv-disk-dir round-trips). Even an alien
request sharing no prefix at all evicted slot[0] on the tie-break.

Split every RAM reuse tier into a cheap probe predicate and its
materialization, and define tier order and conditions exactly once in
slot_probe_reuse_locked(): responses-visible, responses/anthropic
tool-output, GLM memory-rewind, memory-token, thinking-visible,
memory-text. The execution ladder (generate_job) probes once and
materializes only the selected tier; job_slot_score consumes the same
probe, so routing and execution cannot disagree. A slot scores +live
when the request would reuse its live state, -live otherwise: an empty
slot always beats an eviction, and the cheapest (shortest) unusable
checkpoint becomes the victim. Text-tier bindings are part of the probe,
so single-session chats whose resent history retokenizes differently or
omits hidden thinking keep pinning to their slot.

The memory-text probe needs rendered checkpoint text; cache it per slot
and refresh in the slot worker after each job, so probing stays a memcmp
and never detokenizes under tool_mu.

Verified: unit tests for every probe tier plus dispatch routing
(./ds4_test --server), full make test green, and live checks showing the
subagent deflected to the empty slot, main resuming via memory-token,
and thinking-visible / anthropic-tool-output / responses-tool-output
continuations unchanged.
 2026-08-10 16:05:17 +05:00
antirez 84cc882352 rocm: enable DSpark speculative decoding 2026-08-09 19:53:31 +02:00
antirez 5186e2cdc3 build: keep Metal decode changes portable 2026-08-09 19:53:22 +02:00
antirez 42033ee3f4 metal: pipeline DFlash verification 
Submit tiny target-verifier batches incrementally on Metal and retune the Metal confidence default for direct verifier-state commits.
 2026-08-09 18:35:35 +02:00
Salvatore Sanfilippo 0ad494e514 metal: make exact fast paths automatic and trim experiments 2026-08-09 17:48:01 +02:00
antirez fd2d5e05a1 Merge Ivan Fioravanti's exact Metal MXFP4 speedups 2026-08-09 16:57:08 +02:00
ivanfioravanti 5b738048fc chore: drop dead abandoned Metal env allowlist names 
Remove allowlist-only OUTPUT_Q8_NR8/NR4X2/EXACT_VIEW/MAX_BYTES knobs from
the decode bench after auditing newly introduced DS4_METAL_* gates.
 2026-08-09 11:38:24 +02:00
ivanfioravanti acb66c66c4 chore: tighten speed-bench ignores and LUT drift check 
Drop duplicate ignore entries, nest the Metal bench binaries, clarify the
historical Round-15 test note, and gate test-mxfp4-metal on the generator
--check so the checked-in MXFP4 half LUT cannot drift silently.
 2026-08-09 11:38:24 +02:00
ivanfioravanti c6743a4fbf docs: index campaign results and fold M5 fusion contract 
Keep final campaign reports as the review surface and record mechanism,
eligibility, and rollback notes in the M5 results document.
 2026-08-09 11:38:24 +02:00
ivanfioravanti 023614eab7 metal: remove broken stream512 top-k path 
Drop the non-reproducible Metal streaming top-512 selector and always use
the deterministic argsort/merge path.
 2026-08-09 11:38:24 +02:00
ivanfioravanti 319614935f server: canonicalize OpenAI tool schema JSON spelling 
Render tool schemas with stable Python-style separators and escapes so
semantically identical client JSON tokenizes the same way in the prompt.
 2026-08-09 11:38:02 +02:00
ivanfioravanti 0658dac845 metal: fix long-context prefill and decode correctness 2026-08-09 11:38:02 +02:00
ivanfioravanti 0ead8a87b0 server: recover truncated DSML tool calls 2026-08-09 11:38:02 +02:00
ivanfioravanti e6a6ab518a metal: port exact M5 decode paths to pre-M5 
Promote the positive ordinary-Metal M5 decode compounds on the existing M1-M4 path: both QKV/compressor forms, compressor finalization and exact pool, packed FlashAttention, the HC producer continuation, and the exact-shape concurrent Q2 FFN. Keep all existing fallbacks and add dominant per-feature plus aggregate rollback controls.

Revalidate the resident IQ2XXS/Q2_K 2/32 command-buffer schedule and admit it only for the exact single-device tape and measured position window. Explicit split values, including second=0, remain authoritative.

On M3 Ultra, the official interleaved median moves from 41.08 to 44.18 tok/s (+7.55%; steady 41.44 to 44.52). The 1024-token full-vocabulary A/B is bit exact over 1041 rows, 134580480 F32 logits, and 1040 selected IDs. MXFP4 also improves 39.97 to 42.67 tok/s. Full make test and focused Metal/Q2 tests pass.
 2026-08-09 11:38:02 +02:00
Ivan Fioravanti 0e8fb54487 metal: push exact M5 q2 decode past 45 tok/s 2026-08-09 11:38:02 +02:00
Ivan Fioravanti 4a0b305789 docs: record M5 decode passing 42 tok/s 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 1be41c85f4 metal: pack exact split-K FlashAttention reduction on M5 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 92d83e6ecd metal: cluster exact HC norm mix producers on M5 2026-08-09 11:38:01 +02:00
Ivan Fioravanti ad4d05c824 metal: fuse exact ratio4 compressor pool on M5 2026-08-09 11:38:01 +02:00
Ivan Fioravanti ecec0cc31d docs: M5 round-3 final numbers - decode 40.96 med / 41.21 steady 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 3826f52175 metal: admit M5 to compressor exact softmax+product fusion (+0.1-0.2%) 
Round-1 screen measured this flat (-0.16%) before the emit-path
finalize fusion; on the current tree two forward A/B runs are positive
(+0.15%/+0.10%) and the inverse (disable) is negative (-0.24%), all
blocks bit-exact.  Keeps DS4_METAL_DISABLE_PRE_M5_*/_M3_* rollbacks.
 2026-08-09 11:38:01 +02:00
Ivan Fioravanti c8163cb1a8 docs: M5 round-3 results - pair+quad and emit-finalize fusions, steady decode 41.05 2026-08-09 11:38:01 +02:00
Ivan Fioravanti b64d3e0689 metal: fuse emit-path compressor finalize on M5 decode (+0.4-0.5%) 
Every ratio-th decode token, each layer finalizes one freshly pooled
compressor row per compressor: RMS norm, RoPE tail, then the FP8
round-trip + F16 commit copy (attention row, 512 floats) or the
Hadamard+FP4 QAT (indexer row, 128 floats), plus the ratio-4 state
shifts for both.  Those were eleven tiny single-row dispatches per
layer per emit token; kernel_dsv4_comp_row_finalize_f32 runs them as
ONE dispatch: threadgroup 0 finalizes the attention row, threadgroup 1
the indexer row, threadgroups 2-21 do both state shifts.

Bit-exact by construction: the norm phases replicate
kernel_rms_norm_mul_f32_4's float4/simd_sum/zero-padded cross-simdgroup
tree at its exact thread counts (128 for the 512-row, 32 for the
128-row); the RoPE phases call ds4_rope_tail_pair_affine_row verbatim
(lanes 0-63, nthreads=64); the FP8 phase replicates
kernel_dsv4_fp8_kv_quantize_f32's 64-lane shmem max tree; the QAT phase
replicates kernel_dsv4_indexer_hadamard_fp4_f32; the commit copy and
state shifts are elementwise.  Threads outside a phase's virtual width
still execute every barrier, so barriers stay uniform.

The compressor update wrapper gains a defer_finalize flag that stops
after the pool (the shift is deferred too and rides the fused
dispatch).  Default on M5 (DS4_METAL_ENABLE_COMP_FINALIZE_FUSE opts in
elsewhere; DS4_METAL_DISABLE_M5_COMP_FINALIZE_FUSE rolls back).

Evidence on M5 Max, ds4f-q2: six controlled inverse A/B runs, all
positive (+0.36%/+0.41%/+0.42%/+0.42%/+0.48%/+0.59%), every retained
block exact (exact_rows=1041 exact_floats=134580480
exact_selected_ids=1040); official frontier hashes canonical.
 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 71d8c2a072 metal: fuse q_a/kv pair + quad compressor store on M5 decode (+0.5%) 
One dispatch now covers the q_a/kv Q8 pair projection and the four F16
compressor projections (attention + indexer) with their state-store
epilogue in the resident FULL-phase decode layer.  Both stages read the
same normalized attention input and write disjoint outputs.

kernel_dsv4_qkv_pair_quad_compressor_store_q8_0 hosts two virtual NSG=4
cohorts per threadgroup for the Q8 pair range (exact replica of
kernel_mul_mv_q8_0_f32_pair per-lane K walk and reduction tree, same
cohort trick as kernel_dsv4_router_shared_gate_up_q8_0); the compressor
ranges run kernel_mul_mv_f16_f32_pair_4_impl<2> and the paired store
epilogue verbatim.  Every output bit matches the separate dispatches:
controlled A/B exact_rows=1041 exact_floats=134580480
exact_selected_ids=1040 in every retained block, and the official
frontier hash stays eb794f497861d7d9e373665f6e115d8ebe4e17c13c431aa7e318282f16a0f21d.

Default on M5 Apple silicon only (DS4_METAL_ENABLE_QKV_PAIR_QUAD_FUSE
opts in elsewhere; DS4_METAL_DISABLE_M5_QKV_PAIR_QUAD_FUSE rolls back).
Measurements on M5 Max, ds4f-q2: inverse A/B (disable) -0.44%/-0.57%;
interleaved official replicas +0.40%/+1.48% decode.
 2026-08-09 11:38:01 +02:00
Ivan Fioravanti bfeb61d178 docs: M5 round-2 beyond-41 attempt - negative results and exactness caveat 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 1bc0f44b86 docs: M5 decode fusion port results (ds4f-q2, M5 Max) 2026-08-09 11:38:01 +02:00
Ivan Fioravanti 0e125e6e54 metal: admit M5 to seven exact decode fusions, decode 39.51 -> 40.54 tok/s 
Ports the pre-M5 bit-exact decode fusions to M5 (Metal 4 / TensorOps)
per speed-bench/ds4_m5_fusion_port_handoff.md, measured on M5 Max 128 GiB
with the resident ds4f-q2 (IQ2XXS experts) model, official benchmark
command, frontier SHA-256 eb794f497861d7d9e373665f6e115d8ebe4e17c13c431aa7e318282f16a0f21d
stable across every retained replica (baseline and all variants).

Baseline (branch tip, fusions off, 3 replicas, medians):
  prefill 784.13 tok/s, decode 39.39 tok/s, steady 39.60 tok/s,
  first token 28.33 ms. GPU busy during decode ~98%, so the lever is
  per-dispatch overhead inside a busy timeline, as on pre-M5.

Promoted to M5 default (new ds4_gpu_device_is_m5_apple_silicon() gate,
existing DS4_METAL_DISABLE_* rollbacks unchanged):

  HC norm+mix fusion (attn + ffn HC-pre):
      forward A/B +0.13%/+0.10%, inverse -0.61% when disabled.
  Compressor quad projection:
      forward +0.20%/+0.17%, inverse -0.42%.
  Router + shared-expert gate/up fusion:
      forward +0.30%/+0.28%, inverse -1.04%.
  qkv norm + KV RoPE + FP8 store triple fusion:
      forward +0.64%/+0.51%, inverse -0.58%.
  Attention inverse-RoPE -> flash-attention reduce fusion:
      forward +0.28%/+0.14%, inverse -0.29%.
  Compressor ratio-4 decode pack fusion:
      forward +0.31%/+0.39%, inverse -0.42%.
  Router transform+finalize+weights one-simd fusion:
      forward +0.12%/+0.07%/+0.15%, inverse -0.21%.

Every controlled A/B retained block exact (exact_rows=1041
exact_floats=134580480 exact_selected_ids=1040).

Interleaved same-window official replicas (drift-controlled):
  ON  decode 40.59/40.41/40.54 (median 40.54), steady 40.7 med,
      prefill 787.6 med, first token 27.9 ms med.
  OFF decode 39.59/39.55/39.48/39.48 (median 39.51), steady 39.7 med,
      prefill 786.9 med, first token 28.1 ms med.
  Decode +2.6% same-window, +2.9% vs the original baseline median,
  no prefill regression.

Screened and NOT ported (flat or negative on M5):
  5/16 decode split schedule (-0.12%; M5 encode/GPU overlap already
  covered by the 4/0 schedule at ~98% busy), compressor exact reduction
  fusion (-0.16%), head RMS+RoPE static pipeline (-0.06%), raw zero
  attention mask (-0.06%). MXFP4-only pre-M5 paths (nsg=1 MoE decode,
  fixed-route, pipeline fast lookup, compact-tile prefill) do not engage
  with IQ2XXS experts and were not applicable.
 2026-08-09 11:38:01 +02:00
ivanfioravanti 4e401a10e1 metal: four exact pre-M5 decode fusions, decode 38.65 -> 40.20 tok/s 2026-08-09 11:38:01 +02:00
ivanfioravanti f75a065678 metal: four exact MXFP4 decode and prefill speedups 
Decode MoE static trip counts, prefill pair-tile half-domain scaling, KV
RoPE tail fused into the FP8 store, and inverse RoPE folded into the flash
attention reduce. All four are bit-exact with PRE_M5 disable rollbacks.
 2026-08-09 11:38:01 +02:00
ivanfioravanti bcd1e57c5e metal: exact MXFP4 speedups for rounds 7-9 
Integrate the MXFP4 Metal decode/prefill optimizations (rounds 7-9) as one
commit. Keep both the MXFP4 half_lut/tail_cull down-pipeline selection and
main's mpp_mask override for IQ2_XXS gate/up and Q2_K/IQ2_XXS down; they
target disjoint quant types and coexist.
 2026-08-09 11:38:00 +02:00
Salvatore Sanfilippo 0e89a0eeff dspark: commit accepted verifier state directly 2026-08-07 13:01:51 +02:00
Salvatore Sanfilippo 8a703b6707 server: test real client disconnects 2026-08-07 10:40:02 +02:00
Matteo Collina e9ded97afe Cancel server work when clients disconnect 2026-08-05 20:45:38 +00:00
antirez b030961104 agent: disable upto by default for new DS4F 
The new DS4F checkpoint struggles with anchored [upto] edits. Keep exact old/new replacement as the default and make [upto] prompting and automatic insertion an explicit --edit-upto opt-in.
 2026-08-05 14:03:22 +02:00
Salvatore Sanfilippo e0c63d9345 docs: refresh M5 Max and DGX Spark benchmarks 2026-08-05 13:44:49 +02:00
Salvatore Sanfilippo 532ec8b784 metal: accelerate DeepSeek routed MoE prefill 2026-08-05 12:44:52 +02:00
Salvatore Sanfilippo 222b2cb9f9 metal: accelerate DeepSeek indexed prefill 2026-08-05 12:44:52 +02:00
Salvatore Sanfilippo 6747e7718d Reject unsupported GLM TP expert layouts early 2026-08-04 21:32:16 +02:00
Salvatore Sanfilippo 0ab7aaf0f8 Explain the project as an AI-era implementation template 2026-08-04 16:26:20 +02:00
Salvatore Sanfilippo 7fb28303df Replay partial DSpark accepts through ordinary decode 
The partial-accept shortcut retained the verifier batch compressor frontier and could change later greedy tokens. Restore the snapshot and replay partial accepts just like full accepts.

Closes #658

Closes #659
 2026-08-04 14:28:32 +02:00
Salvatore Sanfilippo d14ce35661 Keep streamed prefill maps consistent with expert kernels 
Select decode-only model maps per layer and skip mapped-hotlist seeding when the selected-address kernel owns expert loading. This prevents short SSD-streamed prefills from reading expert tensors outside the active Metal views.
 2026-08-04 14:23:51 +02:00
Salvatore Sanfilippo 464c1f6b62 Keep Q4K regression test warning-free 2026-08-04 13:57:21 +02:00
Salvatore Sanfilippo 7adb2a497d Add critical-fix release gates 2026-08-04 13:48:58 +02:00
Salvatore Sanfilippo 355da75342 Add focused regressions for critical parser fixes 2026-08-04 13:47:47 +02:00
Salvatore Sanfilippo a968c08d76 Make model conversion size checks overflow-safe 
Validate FP8 and FP4 shape products before comparing buffer sizes or allocating converted tensors, including the packed FP4 dimension expansion.
 2026-08-04 13:43:58 +02:00
Salvatore Sanfilippo 51a1c14f35 Recover complete tool calls from unclosed reasoning 
Stop relying on the model to restart a call after an injected reasoning close. Wait for a complete call, keep the preceding prose as reasoning, return the call structurally, and keep its protocol text out of streaming output.

Closes #675
 2026-08-04 13:40:20 +02:00
Salvatore Sanfilippo fe2d3b0c12 Keep repeated reasoning out of streamed answers 
Stop asking tool-enabled models to open a reasoning block that the prompt already opened. Hold the tentative answer until a tool call, response end, or second close marker proves whether it is answer text or another reasoning pass, for both OpenAI and Anthropic streams.

Closes #678
 2026-08-04 13:35:02 +02:00
Salvatore Sanfilippo 7694112475 Reuse repeated GLM prompts by rewinding live state 
When an incoming GLM prompt exactly matches the beginning of the live session, rewind one token before its boundary and reevaluate that final prompt token. This restores the correct logits without rebuilding a long prompt.

Closes #668
 2026-08-04 13:32:49 +02:00
Salvatore Sanfilippo 24fa85eafa Keep disk KV checkpoints after successful loads 
The cold-prompt save limit no longer deletes deeper cold or continued checkpoints after their first restore. Disk-budget eviction remains responsible for reclaiming files.

Closes #673
 2026-08-04 13:31:31 +02:00
Salvatore Sanfilippo b1f88933ce Build Blackwell sm_120 with architecture-specific features 
Compile RTX 50 and RTX PRO Blackwell targets as compute_120a/sm_120a so the vendored block-scaled MXFP4 instructions are accepted. Accept both sm_120 and sm_120a spellings.

Closes #672

Closes #679
 2026-08-04 13:30:47 +02:00
nexus-cw af80694e95 Restore greedy identity after DSpark acceptance 
Replay accepted tokens through the ordinary single-token path so speculative verification cannot leave numerically different compressor state behind.

Closes #658

Closes #659
 2026-08-04 13:26:38 +02:00
Luigi Colluto a169cffa4e Validate tool histories in linear time 
Build the nearest prior call lookup while scanning forward so large Responses and Anthropic histories do not trigger repeated full-history searches.

Closes #547
 2026-08-04 13:26:28 +02:00
AAB 5c318ed804 Replace owned JSON fields only after successful parsing 
Parse duplicate request fields into temporary storage before releasing the previous value, and cover malformed duplicates across OpenAI, Anthropic, Responses, and tool schemas.

Closes #433
 2026-08-04 13:25:50 +02:00
Luigi Colluto 8340f351a6 Reject NaN before JSON integer conversion 
Clamp non-numeric and negative values before converting request fields to int, avoiding undefined floating-point conversion behavior.

Closes #541
 2026-08-04 13:25:17 +02:00
Luigi Colluto 0fa15c6c67 Bound agent cache strings by the remaining file 
Reject truncated cache records before allocating memory for their declared text or title lengths.

Closes #544
 2026-08-04 13:25:03 +02:00
Gilbert Barajas 422ceb1a0b Skip compressor-less layers during strict Flash capture 
The first two Flash layers have no compressor state to capture, so treat them as valid no-op layers in strict speculative verification.

Closes #648
 2026-08-04 13:24:46 +02:00
Joel Teply 727836a79d Include float.h for CUDA FLT_MAX 
Make the CUDA source independent of indirect standard-library includes in newer NVCC releases.

Closes #666

Closes #667
 2026-08-04 13:24:34 +02:00
Kingsley Wong 9735ad8014 Route unterminated reasoning to reasoning_content 
Keep truncated DeepSeek and GLM reasoning out of the final-answer field in non-streaming responses.

Closes #509

Closes #665
 2026-08-04 13:24:22 +02:00
Luigi Colluto c7689dbcf0 Reject impossible model-file sizes 
Bound safetensors and GGUF length fields by the bytes remaining in the file, cross-check FP8/FP4 tensor storage against declared shapes, and reject impossible metadata and tensor counts before allocation.

Closes #542

Closes #543

Closes #545
 2026-08-04 13:24:12 +02:00
Luigi Colluto f02d79cb6b Reject oversized snapshots before allocation 
Discard invalid distributed snapshot bodies before reserving token memory or entering the token read loop.

Closes #546
 2026-08-04 13:23:59 +02:00
Luigi Colluto 8e49a5cc2f Fix double-free in JSON request parser 
Clear string outputs before parsing so malformed duplicate fields cannot leave a freed pointer behind.

Closes #539
 2026-08-04 13:23:48 +02:00
antirez ae504c963e Update Flash downloads and DSpark support for 0731 2026-08-04 13:00:44 +02:00
antirez 8e0f997333 Harden Metal streamed-expert mlock fallback 2026-08-04 11:39:05 +02:00
antirez 55b17c6a80 Merge pull request #578 from gilbert-barajas/fix/streaming-expert-mlock-pinning 2026-08-04 11:28:49 +02:00
antirez 0dd0d36617 Add release speed regression gate 2026-08-04 10:58:38 +02:00
antirez 6841f07239 Merge pull request #623 from kyuz0/main 2026-08-04 10:43:04 +02:00
antirez 7b55707c00 Default L40S server to Flash 0731 MXFP4 2026-08-03 19:40:35 +02:00
antirez b7e9f00911 Version DeepSeek Flash fixtures by checkpoint 2026-08-03 19:25:05 +02:00
antirez 6537ba089b Account for batched GLM session memory during placement 2026-08-03 14:48:05 +02:00
antirez 67d97ab48d Add native MXFP4 inference to CUDA 2026-08-03 13:13:50 +02:00
antirez df341214a9 Merge branch 'ds4f-mxfp4' 2026-08-03 10:12:26 +02:00
antirez 8579ec82c0 Make batched server prefill quantum configurable 2026-08-03 10:07:30 +02:00
antirez d0a04acead build: link CUDA test targets with MMQ 2026-08-03 00:20:35 +02:00
antirez 1464c0f202 cuda: accelerate small Q4 tensor-parallel batches 2026-08-03 00:15:41 +02:00
antirez 73dd0f9004 cuda: fix generic MMQ linkage and warnings 2026-08-03 00:01:55 +02:00
antirez 1c5e24b448 Add managed CUDA TP server restart 2026-08-02 23:58:29 +02:00
antirez 108eb3abe3 Fix CUDA TP server launch and toolkit detection 2026-08-02 23:58:29 +02:00
Entrpi 3b3b29eb3c cuda: build and route aligned self-load artifacts 2026-08-02 23:48:38 +02:00
Entrpi fa4a475eca cuda: route wide zero-prefix attention through HMMA 
Treat contiguous initial mixed-attention chunks as degenerate raw-cache rings and send eligible unmasked shapes through the shared dense token-tile dispatcher.
 2026-08-02 13:12:49 +02:00
Entrpi 9f3256179f cuda: drain and invalidate before scratch growth 
Synchronize rare global scratch resizes and discard decode graphs that may hold the retiring pointer before freeing it.
 2026-08-02 13:06:48 +02:00
antirez e21f84db12 qa: link official scorer with CUDA MMQ 2026-08-02 13:05:39 +02:00
Entrpi 1728313187 cuda: extend token-tiled attention to dense ranges 
Use the shared HMMA path for wide non-indexed attention batches. Build their causal compressed-row records arithmetically and retain the established path for masked, narrow, multi-GPU, or captured work.
 2026-08-02 13:00:02 +02:00
Entrpi bff6f5eb01 cuda: narrow token-tiled attention geometry 
Use four-token tiles with eight head rows so deep indexed attention stages a smaller compressed-row union while retaining the same 32-row MMA shape.
 2026-08-02 12:52:05 +02:00
Entrpi cfaee47995 cuda: add token-tiled HMMA indexed attention 
Port the fork token-tile kernel to the single-sequence indexed-prefill path. Keep its scratch private and retain the existing attention implementation as the fallback for other shapes and graph capture.
 2026-08-02 12:49:02 +02:00
Entrpi 574ccbba32 cuda: fold HC RMS norm into FP16 projection input 2026-08-02 12:22:53 +02:00
Entrpi 0976bdd68b cuda: pin heads8 attention occupancy on SM121 2026-08-02 12:16:23 +02:00
Entrpi 56ec892607 cuda: use SM121 MXFP4 for wide indexer scoring 2026-08-02 12:04:58 +02:00
Entrpi 45f4ef9fec cuda: stream exact top-512 selection for wide prefill 
Port the exact streaming selector from entrpi/batched-serving commit f01c5a1d. It retains the existing score/index total order and leaves decode capture on the established specialized tiers.
 2026-08-02 11:49:47 +02:00
Entrpi 08fecd939e cuda: register-block the prefill indexer scorer 
Port the 32-token WMMA scorer from entrpi/batched-serving commit ff435fea. Keep per-element accumulation order while sharing staged index fragments and retaining accumulators in registers.
 2026-08-02 11:43:47 +02:00
Entrpi c784026f67 cuda: remove indexer scorer shared-memory bank conflicts 
Port the 136-half WMMA staging stride from entrpi/batched-serving commit 3352032d. The padding changes only shared-memory addressing and preserves the scorer arithmetic.
 2026-08-02 11:34:09 +02:00
EntrpiandClaude Fable 5 6f087f4258 build: cuda-spark uses native arch 
The empty CUDA_ARCH left nvcc on its default target, which compiles the
mmq tensor-core paths out (measured 39 vs 423 prefill tok/s on GB10).
The fork's packaging notes flagged the same trap.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011CsqQzZEeuz9biUSEkGm8b
 2026-08-02 00:54:09 +02:00
EntrpiandClaude Fable 5 0f62c620e8 cuda: vendored mmq prefill tier (raw entries) from batched-serving fork 
Vendors the Entrpi/ds4 fork's cuda/mmq module at its v0.5.1 state
(llama.cpp mmq kernels pinned in cuda/mmq/VENDOR.md, MIT, plus the
fork's ds4_mmq C ABI, D2R kernels and repack library; fused prefill
pipeline portions carry Marco Palaferri's Portions-Copyright headers)
and wires the RAW-layout tier only: dense Q8_0 GEMMs and the IQ2_XXS
gate/up pair + Q2_K down routed-MoE pipeline at n_tok >= 2. The
aligned-SoA / D2R / producer-quantized tiers and the weight server are
deliberately not wired (they need the repack artifact pipeline).
Decode (n_tok == 1) is untouched. mmq changes FP32 reduction order vs
the cublas+dequant pipeline, so prefill logits drift at ULP scale:
gated on quality mode off, DS4_CUDA_MMQ=0 restores legacy dispatch.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011CsqQzZEeuz9biUSEkGm8b
 2026-08-02 00:47:10 +02:00
EntrpiandClaude Fable 5 e50104f60f cuda: decode-island CUDA graph capture for serial decode 
Port of the Entrpi/ds4 batched-serving fork's per-layer decode graph
design (fork commits abe26575, 2200c675, 9caaa508; author Entrpi
<entrpi@proton.me>) onto upstream's decode body. Two
position-independent islands of each decode layer are captured once
per activation-buffer variant and replayed on later tokens: the hc-pre
stage (exactly the TO_QKV phase prefix) and the
attention-output-to-layer-end tail (exactly the FROM_ATTN_TO_FFN phase
suffix). Capture and replay reuse the existing phase machinery, so
there is no duplicated encode path.

Island kernels ride cuda_decode_stream(): the legacy NULL stream in
eager mode (bit-identical behavior to before) or a dedicated blocking
stream while capturing/replaying, which preserves the legacy stream's
implicit ordering; cuBLAS is pointed at the capture stream for the
duration of a capture. First sight of a key runs eager to warm lazy
allocators, second sight captures + instantiates + launches, later
sights replay. Any failure retires the entry and re-encodes the island
eagerly, so a broken capture cannot lose a token's work. TP,
multi-tier, SSD streaming, debug dumps, profiling, reference kernels,
and the hash-router layers (token-dependent router argument) all
disqualify capture and stay eager. DS4_CUDA_DECODE_GRAPHS=0 disables;
DS4_CUDA_DECODE_GRAPH_LOG=1 traces captures.

Measured on DGX Spark GB10 (sm_121), DeepSeek V4 Flash q2 imatrix
gguf: 300-token greedy output byte-identical to main with graphs on
(86 island-0 + 80 island-1 graphs live); decode 17.2-17.9 t/s vs
17.0-17.5 baseline (+1-2%). The modest win shows launch overhead is
not the dominant serial-decode cost in upstream's body at 2k context;
the capture substrate is groundwork for extending coverage to the
position-dependent middle section (fork-style device-scalar
substrates), where the remaining launches live.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011CsqQzZEeuz9biUSEkGm8b
 2026-08-01 21:43:23 +02:00
antirez 80df56af40 Refresh DeepSeek V4 Flash 0731 quality fixtures 2026-08-01 11:46:40 +02:00
Salvatore Sanfilippo 4893e0c40f Seed streamed expert cache from mapped prefill layers 2026-08-01 09:50:30 +02:00
Salvatore Sanfilippo 7bec128aaa Add portable Metal MXFP4 expert inference 2026-08-01 09:50:30 +02:00
Salvatore Sanfilippo 1f862bd282 Add MXFP4 GGUF format and scalar reference 2026-08-01 09:50:29 +02:00
Salvatore Sanfilippo 860e673123 Scope strict imatrix checks to weighted quants 2026-08-01 09:50:05 +02:00
Salvatore Sanfilippo 725b084db3 Add lossless DeepSeek MXFP4 GGUF conversion 2026-08-01 09:50:05 +02:00
Salvatore Sanfilippo 0a62b39686 Fix imatrix prompt delimiter parsing 2026-08-01 09:50:05 +02:00
Salvatore Sanfilippo 48392e3d36 Use aligned Q8 tiles for arbitrary DeepSeek prefills 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo ac5cb2c85f Compact fused DeepSeek routed prefill 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo 24903f33e4 Fuse DeepSeek Q-head normalization and RoPE 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo 81f8534e3b Compact routed MoE prefill dispatches 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo d69a017181 Accelerate partial Q4 prefill tiles 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo 96c3ba4777 Accelerate Q4 attention output prefill with Metal 4 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo f284165dc8 Keep DSpark final hidden work in the stage command buffer 2026-08-01 09:49:56 +02:00
Salvatore Sanfilippo c0fc1f704e Fix DeepSeek DenseQ4 decode dispatch 2026-08-01 09:49:55 +02:00
Salvatore Sanfilippo a51e6ec6e9 Tune the default DSpark confidence threshold 2026-08-01 09:49:55 +02:00
Salvatore Sanfilippo 4591cb175e Make DSpark scheduling deterministic by default 2026-08-01 09:49:55 +02:00
Salvatore Sanfilippo 74e8f11918 Commit partial DSpark verifier blocks without replay 2026-08-01 09:49:55 +02:00
Salvatore Sanfilippo df048879fa Parallelize DeepSeek indexed decode attention 2026-08-01 09:49:55 +02:00
Donato Capitella d250a7c07c Optimize ROCm two-token MTP verification 2026-08-01 08:29:22 +01:00
Donato Capitella 2f49b27edd Fix ROCm IQ2 SSD prefill deadlock 2026-07-29 09:27:35 +01:00
Donato Capitella 6882a0f2f4 Tighten ROCm DeepSeek large model arenas 2026-07-28 21:44:11 +01:00
Donato Capitella 9e1988bcd1 Fix ROCm DeepSeek Q4 SSD expert staging 2026-07-28 19:22:22 +01:00
Donato Capitella f23f13d5c3 Merge ROCm DeepSeek prequant decode restoration 2026-07-28 16:34:28 +01:00
Donato Capitella 1e8f16ca05 Restore ROCm DeepSeek prequant decode kernels 2026-07-28 13:33:07 +01:00
Salvatore Sanfilippo 54b36ed9ba Merge pull request #617 from kyuz0/fix/rocm-distributed-glm 
Add GLM 5.2 inference (SSD streaming and distributed) on ROCm for GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf
 2026-07-28 14:04:52 +02:00
Donato Capitella 11da8b334f Merge ROCm batched HC norm fix 2026-07-28 10:17:01 +01:00
Donato Capitella 45a9dcb0d3 Fix ROCm batched HC norm fusion 2026-07-28 10:16:47 +01:00
Donato Capitella 2510636cb0 Remove internal GLM validation artifacts 2026-07-27 09:14:54 +01:00
Donato Capitella 4f11251b5f Remove ineffective GLM group8 decode experiment 2026-07-26 22:53:46 +01:00
Donato Capitella 4ab75720eb Add GLM indexed decode tuning path 2026-07-26 22:25:05 +01:00
Donato Capitella 4e80dd7502 Default GLM selected attention to tile 16 2026-07-26 21:26:15 +01:00
Donato Capitella 367a64c1bc Extend GLM selected attention head tiling 2026-07-26 20:58:15 +01:00
Donato Capitella 4d7fbbf5f3 Add GLM selected attention head tiling 2026-07-26 20:25:30 +01:00
Donato Capitella cfb1ff9a07 Record GLM causal attention deployment 2026-07-26 19:57:09 +01:00
Donato Capitella dd837f033d Default GLM causal attention GEMM 2026-07-26 19:35:26 +01:00
Donato Capitella 764987bc2c Record GLM selected attention deployment 2026-07-26 19:27:15 +01:00
Donato Capitella ea2e76681c Default GLM selected attention GEMM 2026-07-26 19:07:56 +01:00
Donato Capitella 34b3862d80 Default GLM wave value projection 2026-07-26 18:33:29 +01:00
Donato Capitella 7a84671c07 Add opt-in GLM wave value projection 2026-07-26 14:56:43 +01:00
Donato Capitella 41a3c79d47 Remove ineffective ROCm GLM QKV experiment 2026-07-26 11:29:24 +01:00
Donato Capitella 5920d6808b Add opt-in ROCm GLM performance experiments 2026-07-26 10:32:37 +01:00
Donato Capitella 5e61ad5ca9 Remove ineffective ROCm GLM F16 compact cache 2026-07-26 10:10:48 +01:00
Donato Capitella 312934c14a Add opt-in ROCm GLM F16 compact cache 2026-07-25 21:20:49 +01:00
Donato Capitella b7f5cddc61 remove FP32 hipblas path 2026-07-25 21:05:31 +01:00
Donato Capitella 0e609ab85b Add F32 ROCm GLM causal attention GEMMs 2026-07-25 19:48:43 +01:00
Donato Capitella c9e1013e71 Add optimized ROCm GLM slice decode path 2026-07-25 19:14:45 +01:00
Donato Capitella 26e8e88665 Add ROCm GLM causal prefill GEMM path 2026-07-25 18:45:22 +01:00
Donato Capitella 9463fd73d9 Rename DS4_ROCM_GLM_GROUPED_QK_LOW to DS4_ROCM_GLM_GROUPED_QK_LOW 2026-07-25 18:29:37 +01:00
Donato Capitella f6f5b92631 Optimize GLM ROCm prefill projections 2026-07-25 17:28:46 +01:00
Donato Capitella fa8b0b9301 Optimize ROCm indexer top-k 2048 2026-07-25 15:26:19 +01:00
Donato Capitella 7a424a456a Add grouped GLM value projection path 2026-07-25 13:44:57 +01:00
antirez 0a7ad776b9 Fix batched server session recovery race 2026-07-23 09:07:10 +00:00
Donato Capitella 2dc022d90d Tune GLM resident ROCm slice memory reserve 2026-07-22 15:10:34 +01:00
Donato Capitella 80afbe7836 Fix ROCm IQ2 expert tile LUT initialization 2026-07-22 14:23:46 +01:00
Donato Capitella 507fd7c6a9 Fix resident ROCm IQ2 routed MoE correctness 2026-07-22 14:18:05 +01:00
Donato Capitella 784bc77b21 Add resident IQ2 MoE sorted-path isolation 2026-07-22 13:51:29 +01:00
Donato Capitella 34d8046bbf Add GLM batch FFN stage dumps 2026-07-22 13:21:33 +01:00
Donato Capitella fd6d056f5c Add non-invasive ROCm graph dumps 2026-07-22 12:33:40 +01:00
Donato Capitella 84d2de4b3f Allow resident ROCm GLM distributed slices 2026-07-22 04:40:25 +01:00
Donato Capitella d7a89093fa Fix ROCm distributed GLM slice residency 2026-07-22 03:40:37 +01:00
Donato Capitella 24f2c83d50 Add ROCm GLM routed IQ2 down support 2026-07-22 02:33:15 +01:00
Gilbert Barajas 3ce677710e metal: restore streamed-expert mlock pinning dropped in the GLM 5.2 refactor 
The "Add GLM 5.2 runtime support" refactor removed all mlock pinning of
streamed experts -- the mlock calls, slab_lock_slot/unlock_slot, and the
budget cap all went to zero. Without pinning, in the SSD-streaming regime
(model >> RAM) the OS evicts the expert cache under memory pressure and
every token re-streams from disk. That is the DeepSeek-Flash decode
regression reported in #532 (~2-3x slower, ~6s first-token).

Port the pinning back, adapted to the refactored slab allocator
(slab_backed/slab_slot, 3 gate/up/down buffers per entry): whole-buffer
mlock at alloc, per-slot lock on fill / unlock under relief, the
routed-memory budget cap for graceful degradation, cap-on-failure, and a
relief routine that unpins the coldest ~10% of victims under pressure.

Bench (M5 Max 128GB, DeepSeek-V4-Flash IQ2XXS, --simulate-used-memory 72
to force the streaming regime):
  before (broken): 0.3-2.5 gen t/s, ~6.0s first-token
  after  (this):   ~7.0 gen t/s,    ~1.7s first-token
  baseline 80ebbc3: ~6.1 gen t/s
Recovers past baseline (the 427e281 kernel work pays off once experts
stay resident).

Fixes #532
 2026-07-20 11:18:21 -05:00
antirez efdadd41e2 README: update the AI assisted models list. 2026-07-20 16:50:40 +02:00
antirez 646db66656 README: remove title now that there is the logo. 2026-07-20 16:41:44 +02:00
antirez cb5103449b Add DwarfStar logo and credits 2026-07-20 16:40:44 +02:00
Ivan Fioravanti 427e281d24 Optimize Metal prefill and decode kernels 2026-07-20 11:09:49 +02:00
antirez e0824ddcaa Make release builds warning-free 2026-07-20 11:09:28 +02:00
antirez a185c36dc4 Stabilize CUDA and Metal session batching 2026-07-18 18:50:59 +02:00
antirez ef8d923ad7 Add ROCm GLM 5.2 support 2026-07-18 18:50:30 +02:00
antirez 66cbce5289 Unify distributed inference CLI and documentation 2026-07-18 18:50:24 +02:00
antirez 36ee8c1407 Add CUDA tensor parallelism and session batching 2026-07-18 18:50:03 +02:00
antirez 36cd0ca915 Add native Metal session batching 2026-07-18 18:49:00 +02:00
antirez 005afedcd9 Add GLM 5.2 inference and quality fixtures 2026-07-18 18:48:45 +02:00
antirez 63d9874c54 Add two-machine Metal tensor parallelism 2026-07-18 18:47:58 +02:00
antirez fc9efd1cc7 Add DSpark speculative decoding 2026-07-18 18:47:26 +02:00
antirez 519c4d8fe7 Fix PRO streaming and sampling correctness 2026-07-18 18:46:19 +02:00
2729 changed files with 4035861 additions and 139815 deletions
Expand all files Collapse all files

No files matched your search

.gitignore
+117 -1
View File
@@ -4,15 +4,131 @@
/ds4-eval
/ds4-agent
/ds4_native
/ds4_agent_test
/ds4_server_test
/ds4_test
/ds4flash.gguf
/TODO.md
/gguf/
/gguf
/core
/core.*
/tests/test_engine_mgpu_placement
/tests/test_engine_mgpu_runtime
/tests/test_engine_mgpu_refusal
/tests/test_engine_correctness
/tests/test_gpu_args
/tests/test_gpu_lookup_cache_strict
/tests/test_gpu_model_cache
/tests/test_gpu_xdev
/tests/test_cuda_mixed_batch
/tests/test_cuda_session_batch
/tests/cuda_long_context_smoke
/tests/test_layer_pack
/tests/test_metal_session_batch
/tests/test_metal_tp_spec
/tests/test_metal_tp_cancel
/tests/test_metal_moe_prefill
/tests/test_metal_ssd_experts
/tests/test_metal_ssd_reuse
/tests/test_metal_command_memory
/tests/test_metal_dense_mpp
/tests/test_mxfp4_cuda
/tests/test_mmvq_cuda
/tests/test_mmvq_rocm
/tests/test_mxfp4_dot
/tests/test_mxfp4_metal
/tests/test_q4k_dot
/tests/test_sampling
/tests/test_web_recovery
/tests/test_quality_api
/tests/test_cuda_q8_scratch
/tests/test_cuda_dspark_moe
/tests/test_session_state
/tests/test_session_state_gpu
/tests/test_tp_commands
/tests/test_tp_rdma
/tests/test_tp_tcp
/tests/test_tp_link
/tests/test_cuda_tp
/tests/test_cuda_tp_link
/tests/test_cuda_verify_head
/tests/test_cuda_shared_side
/tests/test_cuda_tokentile
/tests/test_cuda_attn_side
/tests/test_cuda_compressor
/tests/test_cuda_copy_spans
/tests/test_cuda_dspark
/tests/test_cuda_dspark_tp_markov
/tests/test_cuda_dspark_tp_split
/tests/test_cuda_embed
/tests/test_cuda_f16_chunks
/tests/test_cuda_grouped_q8
/tests/test_cuda_hc_norm
/tests/test_cuda_hc_pre
/tests/test_cuda_indexed_one
/tests/test_cuda_indexer
/tests/test_cuda_mxfp4_split
/tests/test_cuda_pool
/tests/test_cuda_q8_0_gemv
/tests/test_cuda_q8_pair
/tests/test_cuda_router_stream
/tests/test_cuda_topk_stream
/tests/test_deepseek_indexer
/tests/test_mxfp4_staged_cuda
/tests/test_cuda_tp_repack
/tests/test_metal_tp_bulk
/tests/test_glm53_kda
/tests/test_glm53_vision_engine
/tests/test_glm_attention
/tests/test_glm_attention_rocm
/tests/test_linux_memory
/tests/test_rocm_memory
/tests/test_qwen4_kernels
/tests/test_metal_device_policy
/tests/test_metal_q8_reduction
/tests/test_qwen4_moe_mm_specialize
/tests/test_qwen4_ssd_experts
/tests/test_qwen4_memory
/tests/test_qwen4_generation
/tests/test_qwen4_vision
/tests/test_deepseek4_vision_image
/tests/test_image_decode
/tests/test_deepseek41_metal
/tests/test_deepseek41_cuda
/tests/test_cuda_q8_rows
/tests/test_cuda_reductions
/tests/test_cuda_shared
/tests/test_cuda_ssd_cache
/tests/test_cuda_ssd_batch
/tests/test_cuda_ssd_repack
/tests/test_deepseek41_gguf
/tests/test_qwen4_ngrams
/tests/test_qwen4_ngram_state
/tests/test_qwen4_mtp_prefill
/tests/test_qwen4_mtp_policy
/tests/test_qwen4_ple_prefetch
/tests/test_deepseek41_graph
/tests/test_v41_indexer_score
/tests/test_deepseek41_prefill
/tests/test_deepseek41_cli
/tests/test_engram
*.o
*.dSYM/
__pycache__/
*.pyc
/gguf-tools/libds4quants.dylib
/gguf-tools/libds4quants.so
/misc/
.*.swp
.DS_Store
/tests/test_qwen4_conv_parallel
/tests/test_q8_prefill_variants
/tests/test_metal_qwen4_hc
/tests/test_metal_qwen4_moe_mid
/tests/test_metal_qwen4_moe_down
/tests/test_metal_qwen4_mxfp4_stage
/tests/test_metal_q8_gemv_reference
/tests/test_metal_qwen4_moe_half
/tests/test_qwen4_resident_arithmetic
EVAL_DATA.md
+54
View File
@@ -0,0 +1,54 @@
# Evaluation data
`ds4-eval` contains small, curated subsets of external benchmarks. They are
integration tests for DwarfStar, not official benchmark distributions or
leaderboard scores. Case IDs are retained so each prompt can be checked against
its source.
## Hard suite
The 50-case `hard` suite was assembled on 2026-09-03:
| Source | Cases | Source revision | License |
|---|---:|---|---|
| [MMLU-Pro](https://github.com/TIGER-AI-Lab/MMLU-Pro) | 30 | `f418b116db00b065c2aea046518d8fcf74d39872` | Apache-2.0 |
| [OlympiadBench](https://huggingface.co/datasets/Hothan/OlympiadBench) | 10 | `91184b52131e7fc9455fef848035173aea8cc01a` | Apache-2.0 |
| [LiveBench](https://huggingface.co/datasets/livebench/reasoning) | 5 | `6fc6498a5dfba553f69f4413feabade1f1a2d384` | Apache-2.0 |
| [NIST Juliet 1.3](https://samate.nist.gov/SARD/test-suites/112) | 5 | archive SHA-256 `ada9d7e1c323d283446df3f55bdee0d00bda1fed786785fe98764d58688f38eb` | CC0-1.0 / public domain |
Selected source IDs:
- MMLU-Pro: `86`, `671`, `899`, `1467`, `1995`, `2398`, `2826`, `3024`,
`3546`, `3696`, `4674`, `4699`, `5078`, `5848`, `6034`, `6500`, `6844`,
`6955`, `7738`, `7754`, `8771`, `8863`, `9055`, `9487`, `10360`, `10567`,
`10775`, `11120`, `11292`, `11346`.
- OlympiadBench: `1606`, `1613`, `1614`, `1618`, `1697`, `1709`, `1716`,
`1736`, `1750`, `1766`.
- LiveBench: `c833574b0ff6e8609bee19b36fabc48501cdcd9a64b99cd0422f54ada1acc53d`,
`b7b7e331df4614775210ef5e524ed7f20c5fec6655fb0dd179c94e4d71170862`,
`0daa7ca38beec4441b9d5c04d0b98912322926f0a3ac28a5097889d4ed83506f`,
`9ee37b9a04ab050936c86b2c5bb7abbaa0bc0e737d59a7bff9ba11e9b4069c1d`,
`d7071c9ff5d9779e7ab955366d0ae8db40f785aadfe7ff0b5a7ede98c05c44ea`.
- NIST Juliet reductions use the local IDs beginning with `juliet-` in
`ds4_eval_cases.c`.
The MMLU-Pro choices and keys are unchanged. OlympiadBench cases are English,
text-only rows; only their published final answers are embedded. LiveBench's
benchmark-specific output wrapper was removed because `ds4-eval` supplies one
uniform final-answer instruction. The NIST cases are short, locally written
reductions of Juliet weakness patterns rather than copies of complete Juliet
test cases.
Multiple-choice order is retained for MMLU-Pro. Open answers declare their
grading type and, where needed, a small list of equivalent surface forms. The
12-case `hard-smoke` suite is a fixed subset of `hard`.
Apache-2.0 terms for redistributed benchmark material are included in
[`licenses/Apache-2.0.txt`](licenses/Apache-2.0.txt). DwarfStar's own source
remains under the repository's top-level MIT license.
## Core suite
The original 92-case `core` suite contains 25 GPQA Diamond, 25 SuperGPQA, 25
AIME 2025, and 17 reduced defensive code-review cases. Its existing source IDs,
ordering, prompts, and answer keys are unchanged by the hard-suite addition.
LICENSE
+1
View File
@@ -2,6 +2,7 @@ MIT License
Copyright (c) 2026 The ds4.c authors
Copyright (c) 2023-2026 The ggml authors
Copyright (c) 2023 DeepSeek
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
Makefile
+1430 -62
View File
File diff suppressed because it is too large. Load diff
QA_BEFORE_RELEASES.md
+2053 -31
View File
File diff suppressed because it is too large. Load diff
README.md
+198 -1205
View File
File diff suppressed because it is too large. Load diff
STRIXHALO.md
+3 -133
View File
@@ -1,136 +1,6 @@
# DS4 on Strix Halo
This is the minimal setup for DS4 ROCm inference on a
Strix Halo machine with 128 GB RAM and Radeon 8060S (`gfx1151`).
The setup guide has moved to [docs/STRIX_HALO.md](docs/STRIX_HALO.md).
## 1. Install ROCm
On Ubuntu 26.04 LTS, install the ROCm compiler/runtime and libraries used by the Strix Halo backend:
```sh
sudo apt-get update
sudo apt-get install -y \
hipcc rocminfo rocm-smi \
libamdhip64-dev \
libhipblas-dev libhipblaslt-dev \
librocblas-dev \
librocwmma-dev \
libhipcub-dev
```
The backend uses rocWMMA. On this Ubuntu 26.04 setup, `librocwmma-dev`
installs the top-level rocWMMA headers but misses `rocwmma/internal/`.
No Ubuntu package currently provides those internal headers. Install a complete
matching rocWMMA header tree:
```sh
git clone --depth 1 --branch rocm-7.1.0 https://github.com/ROCm/rocWMMA.git /tmp/rocWMMA-rocm-7.1.0
sudo mkdir -p /usr/local/include
sudo cp -a /tmp/rocWMMA-rocm-7.1.0/library/include/rocwmma /usr/local/include/
```
If ROCm is installed under `/usr` but tooling expects `/opt/rocm`, add these
compatibility links:
```sh
sudo mkdir -p /opt/rocm/bin
sudo ln -sf /usr/bin/hipcc /opt/rocm/bin/hipcc
sudo ln -sfn /usr/lib/x86_64-linux-gnu /opt/rocm/lib
sudo ln -sfn /usr/include /opt/rocm/include
```
## 2. Enable ROCm access
The user running DS4 must be able to open `/dev/kfd` and the DRM render node:
```sh
sudo usermod -aG render,video "$USER"
```
Log out and back in, or reboot. Verify:
```sh
rocminfo | grep -A80 'Name: gfx1151'
```
If DS4 says `no ROCm-capable device is detected`, check that `rocminfo` can open
`/dev/kfd` and that `groups` includes `render`.
## 3. Increase GPU-visible memory
A 128 GB Strix Halo system may initially expose only about 62 GB of GPU-visible
memory. DS4 needs the larger GTT aperture for the 80.76 GiB model plus runtime
buffers.
Use these kernel parameters:
```text
amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856
```
On Ubuntu with GRUB:
```sh
sudo cp /etc/default/grub /etc/default/grub.bak
sudoedit /etc/default/grub
```
Set:
```text
GRUB_CMDLINE_LINUX_DEFAULT="quiet splash amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856"
```
Then:
```sh
sudo update-grub
sudo reboot
```
After reboot, verify:
```sh
cat /proc/cmdline
sudo dmesg | grep -Ei 'GTT|gttsize|TTM|VRAM'
rocminfo | grep -A80 'Name: gfx1151'
```
Expected signs:
```text
amdgpu: 126976M of GTT memory ready
rocminfo gfx1151 pool: 130023424 KB
```
## 4. Build DS4
Use the normal Strix Halo target. It builds the standard binary names:
```sh
make strix-halo -j"$(nproc)"
```
`make rocm` is an alias for `make strix-halo`.
## 5. Use the right GGUF
Use the standard IQ2XXS/Q2K/Q8 imatrix GGUF:
```text
DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
```
Avoid the mixed IQ2/IQ4 or IQ2/Q4 GGUFs on this machine for now. They put much
more memory pressure on the ROCm path and can trigger system OOM instead of a
clean DS4 failure.
## 6. Run DS4
Run it normally:
```sh
./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
```
The ROCm build uses the Strix Halo backend automatically.
It covers ROCm prerequisites, GPU-visible memory, current model downloads,
and first runs for DeepSeek Flash and GLM 5.3 Flash.
cuda/mmq/VENDOR.md
+157
View File
@@ -0,0 +1,157 @@
# Vendored llama.cpp mmq kernels
This directory contains source files copied verbatim from
[llama.cpp's `ggml-cuda` backend](https://github.com/ggml-org/llama.cpp/tree/master/ggml/src/ggml-cuda),
plus a thin ds4-side adapter (`ds4_ggml_stubs.{h,cu}` and `ds4_mmq.{h,cu}`)
that lets the templated CUDA kernels compile and link without the full ggml
runtime.
## Why these files are vendored, not submoduled
ds4 is a flat, self-contained C/CUDA codebase with no third-party build
dependencies. A submodule would force consumers to clone, configure, and
build all of llama.cpp just to use a few quantized-matmul kernels. Vendoring
keeps ds4 self-contained at the cost of a periodic re-sync.
## Upstream pin
| Field | Value |
|---------------|--------------------------------------------------------------------------------------------------------------------|
| Source | https://github.com/ggml-org/llama.cpp |
| Commit | `5c0e9468378eba6bf3cc1989ff5d62fbbe4d9e3a` |
| Commit date | 2026-05-14 |
| Commit title | `ggml-hexagon: cpy: add contiguous fast-path in reshape copy (#23076)` |
| License | MIT (`LICENSE` at the repository root, copyright "2023-2026 The ggml authors"). Compatible with ds4's MIT license. |
## File inventory
| File | Origin in llama.cpp | Status | Lines |
|-----------------------|----------------------------------------------|--------------------------------------------------------------------------|-------|
| `mmq.cuh` | `ggml/src/ggml-cuda/mmq.cuh` | verbatim | 4176 |
| `mma.cuh` | `ggml/src/ggml-cuda/mma.cuh` | verbatim | 1456 |
| `vecdotq.cuh` | `ggml/src/ggml-cuda/vecdotq.cuh` | verbatim | 1317 |
| `quantize.cuh` | `ggml/src/ggml-cuda/quantize.cuh` | verbatim | 41 |
| `quantize.cu` | `ggml/src/ggml-cuda/quantize.cu` | verbatim | 443 |
| `mmid.cuh` | `ggml/src/ggml-cuda/mmid.cuh` | verbatim | 5 |
| `mmid.cu` | `ggml/src/ggml-cuda/mmid.cu` | verbatim | 164 |
| `mmvq.cuh` | `ggml/src/ggml-cuda/mmvq.cuh` | patched (Step 6): `mul_mat_vec_q_switch_type` proto exposed; ggml-tensor entries gated on `DS4_MMVQ_INCLUDE_GGML_ENTRIES` | ~36 |
| `mmvq.cu` | `ggml/src/ggml-cuda/mmvq.cu` | patched: `mul_mat_vec_q_switch_type` promoted from `static`; `ggml_cuda_mul_mat_vec_q` + `ggml_cuda_op_mul_mat_vec_q` gated on `DS4_MMVQ_INCLUDE_GGML_ENTRIES` | 1163 |
| `unary.cuh` | `ggml/src/ggml-cuda/unary.cuh` | verbatim (needed by `mmvq.cu` for inline GLU epilogues) | 114 |
| `common.cuh` | `ggml/src/ggml-cuda/common.cuh` | verbatim | 1489 |
| `ggml-common.h` | `ggml/src/ggml-common.h` | verbatim | 1900 |
| `vendors/cuda.h` | `ggml/src/ggml-cuda/vendors/cuda.h` | verbatim | 28 |
| `ggml.h` | (new) | redirect to `ds4_ggml_stubs.h` | 5 |
| `ggml-impl.h` | (new) | redirect to `ds4_ggml_stubs.h` | 5 |
| `ggml-cuda.h` | (new) | redirect to `ds4_ggml_stubs.h` | 5 |
| `ds4_ggml_stubs.h` | (new) | shim: ggml_type enum, macros, info struct, type_size lookups | ~280 |
| `ds4_ggml_stubs.cu` | (new) | shim impls: `ggml_cuda_info`, naive pool, `ggml_backend_cuda_context::*` | ~110 |
| `ds4_mmq.h` | (new) | public C ABI for ds4 to call | ~70 |
| `ds4_mmq.cu` | (new) | host wrappers; Phase 0 instantiates `mul_mat_q_case<Q8_0>` only | ~120 |
**Total vendored:** ~11,000 lines of CUDA. **Total shim/adapter:** ~600 lines.
## What llama.cpp's `mmq.cu` does that we don't vendor
The upstream `ggml/src/ggml-cuda/mmq.cu` (372 lines) is the ggml-backend
dispatch entry. It talks to `ggml_tensor` / `ggml_backend_buffer_get_usage`,
does its own activation Q8_1 quantization, and wires into the ggml op graph.
We **don't vendor it.** Instead `ds4_mmq.cu` provides ds4-style entries
that bypass the ggml graph and call the per-type `mul_mat_q_case<T>` directly.
## Symbol-resolution table
Symbols the vendored files reference, and how they resolve in this directory:
| Symbol category | Resolution |
|------------------------------------------------|--------------------------------------------------------------------------------------------------|
| `GGML_ASSERT`, `GGML_ABORT`, `GGML_UNUSED` | Macros in `ds4_ggml_stubs.h` |
| `GGML_PAD`, `GGML_UNUSED_VARS` | Macros in `ds4_ggml_stubs.h` |
| `GGML_CUDA_CC_*` constants | Defined in vendored `common.cuh` |
| `GGML_TYPE_*` enum values | Defined in `ds4_ggml_stubs.h::enum ggml_type` |
| `GGML_CUDA_MAX_DEVICES`, `GGML_CUDA_NAME` | Macros in `ds4_ggml_stubs.h` |
| `GGML_MAX_DIMS`, `GGML_MAX_SRC` | Macros in `ds4_ggml_stubs.h` |
| `block_q8_0`, `block_q2_K`, `block_iq2_xxs` | Defined in vendored `ggml-common.h` (gated by `GGML_COMMON_IMPL_CUDA`) |
| `ggml_half`, `ggml_half2` | `uint16_t` typedefs in `ds4_ggml_stubs.h` |
| `ggml_type_size`, `ggml_blck_size` | Inline implementations in `ds4_ggml_stubs.h` |
| `ggml_cuda_info()` | Singleton in `ds4_ggml_stubs.cu`, populated via `cudaGetDeviceProperties` |
| `ggml_cuda_get_device`, `ggml_cuda_set_device` | Thin wrappers in `ds4_ggml_stubs.cu` |
| `ggml_time_us` | `std::chrono::steady_clock` in `ds4_ggml_stubs.cu` |
| `ggml_cuda_pool`, `ggml_cuda_pool_alloc` | Defined in vendored `common.cuh`. Concrete `ds4_naive_pool` subclass in `ds4_ggml_stubs.cu` |
| `ggml_backend_cuda_context` | Defined in vendored `common.cuh`. `new_pool_for_device` and `~dtor` provided in `ds4_ggml_stubs.cu` |
| `ggml_tensor` | Forward declaration in `ds4_ggml_stubs.h`. Never dereferenced - only held as pointer. |
| `ggml_glu_op` | Enum stub in `ds4_ggml_stubs.h` (fusion args present in common.cuh but unused by ds4) |
| `CUDA_CHECK`, `CUBLAS_CHECK` | Macros in `ds4_ggml_stubs.h` |
| `ggml_cuda_launch_mm_ids_helper` | Defined in vendored `mmid.cu` |
| `ggml_cuda_should_use_mmq` | Defined in vendored `mmq.cu` - **we re-implement in `ds4_mmq.cu`** since we don't vendor mmq.cu |
| `ggml_cuda_mul_mat_q*` | Upstream's host entry. **Replaced** by `ds4_mmq_q8_0_dense` and family. |
## Things we deliberately do NOT support
- **CUDA graphs.** ds4 manages its own streams; `USE_CUDA_GRAPH` is undefined.
- **HIP / MUSA / AMD backends.** vendored code's HIP/MUSA `#ifdef` branches are dead in our build but kept for upstream diff-cleanliness.
- **The full `ggml_tensor` type.** No tensor introspection - shapes and strides come in via raw arguments to `ds4_mmq_*`.
- **`ggml_op` graph evaluation.** We call kernels directly.
## Re-syncing with upstream
Local MMVQ bounds fix: pass the logical row count to the generic vector
kernel instead of using its output stride as a row bound. Both generic and
multi-token MoE kernels also guard tail-row weight loads.
MXFP4 MMQ uses Q8_1 activations on Blackwell as on older GPUs. Its tile
loader, K iteration and activation-block stride use the integer-MMA layout;
the stored MXFP4 weights are unchanged.
When upstream lands a bugfix or perf improvement we want, the procedure is:
```sh
cd /tmp/llama-research && git fetch && git checkout <NEW_COMMIT>
cd /Users/ent/code/ds4-mmq-lift/cuda/mmq
cp /tmp/llama-research/ggml/src/ggml-cuda/{mmq.cuh,mma.cuh,vecdotq.cuh,quantize.cuh,quantize.cu,mmid.cuh,mmid.cu,common.cuh} .
cp /tmp/llama-research/ggml/src/ggml-common.h .
cp /tmp/llama-research/ggml/src/ggml-cuda/vendors/cuda.h vendors/
# Verify the shim still covers all referenced symbols:
grep -hoE 'GGML_[A-Z_]+|ggml_[a-z_]+' *.cuh *.cu *.h | sort -u > /tmp/symbols.new
diff /tmp/symbols.last /tmp/symbols.new # check for newly-introduced names
# Update this VENDOR.md's commit pin and run `make cuda CUDA_ARCH=sm_120`.
```
If the new symbols are minor (e.g., new `GGML_CUDA_CC_*` constants), they
likely come from `common.cuh` which we vendor and don't need any shim
changes. If they're new `ggml_*` host functions (rare, but possible if
upstream adds a new helper), extend `ds4_ggml_stubs.h`.
## Testing matrix
| Test | Status | Phase |
|-----------------------------------------------|------------------------------|-----------|
| `nvcc -c cuda/mmq/ds4_mmq.cu` builds cleanly | **passes** (sm_120, nvcc 13) | Phase 0 |
| Q8_0 dense parity vs CPU reference | **passes** (4 shapes) | Phase 2 |
| Q2_K dense parity | **passes** (4 shapes) | Phase 3 |
| IQ2_XXS dense parity | **passes** (4 shapes) | Phase 3 |
| MoE `_id` Q8_0 parity | **passes** (3 shapes) | Phase 4 |
| MoE `_id` Q2_K parity | **passes** (3 shapes) | Phase 4 |
| MoE `_id` IQ2_XXS parity | **passes** (4 shapes) | Phase 4 |
| `make ds4-bench` with mmq integration | **builds and runs** | Phase 5/6 |
| Frontier sweep, ctx 2k-16k, V4 Flash IQ2XXS | **see results below** | Phase 7 |
## Validated performance
PRO 6000 Blackwell (sm_120), CUDA 13.0, V4 Flash IQ2_XXS GGUF (86.7 GB),
default `DS4_CUDA_MMQ_MOE_MIN_TOKENS=2` (legacy decode preserved):
| ctx | baseline pf t/s | mmq pf t/s | speedup | baseline gen t/s | mmq gen t/s | gen ratio |
|--------|-----------------|------------|------------|------------------|-------------|-----------|
| 2048 | 373.21 | 1033.42 | **2.77x** | 40.48 | 39.33 | 0.972x |
| 4096 | 366.25 | 1041.39 | **2.84x** | 39.64 | 38.64 | 0.975x |
| 6144 | 364.81 | 1025.24 | **2.81x** | 39.50 | 38.55 | 0.976x |
| 8192 | 364.01 | 1026.71 | **2.82x** | 38.81 | 37.88 | 0.976x |
| 10240 | 361.75 | 1019.53 | **2.82x** | 38.45 | 37.57 | 0.977x |
| 12288 | 360.52 | 1013.17 | **2.81x** | 38.31 | 37.22 | 0.972x |
| 14336 | 359.15 | 1004.45 | **2.80x** | 38.09 | 37.18 | 0.976x |
| 16384 | 357.99 | 1001.29 | **2.80x** | 38.71 | 37.86 | 0.978x |
Sustained ~2.80x prefill speedup across the swept context range; gen
within 2.5% of baseline (run-to-run variance). See `local/docs/`
(auto-round companion repo) for the full Phase 0-7 execution log,
parity-test output, and detailed plan.
cuda/mmq/common.cuh
+1489
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/ds4_ggml_stubs.cu
+238
View File
@@ -0,0 +1,238 @@
// SPDX-License-Identifier: MIT
// Implementations of the ggml-API stubs declared in ds4_ggml_stubs.h plus
// the bodies of ggml_backend_cuda_context / ggml_cuda_info /
// new_pool_for_device that the vendored common.cuh declares without
// defining.
//
// CUDA uses stream-ordered allocation. HIP retains stream-local allocations
// because returning live MMQ scratch to its asynchronous pool is unreliable.
#include "common.cuh" // pulls in ds4_ggml_stubs.h via redirect headers
#if defined(GGML_USE_HIP)
#include "vendors/hip.h"
#include "../../ds4_rocm_memory.h"
#else
#include <cuda_runtime.h>
#include <cublas_v2.h>
#endif
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <memory>
#include <mutex>
#include <vector>
// ----------------------------------------------------------------------------
// Device info singleton.
//
// Common.cuh declares `const ggml_cuda_device_info & ggml_cuda_info();` -
// we provide the body. The struct layout (common.cuh:1091) is:
// { int device_count; cuda_device_info devices[GGML_CUDA_MAX_DEVICES];
// std::array<float, GGML_CUDA_MAX_DEVICES> default_tensor_split; }
// where cuda_device_info has { cc, nsm, smpb, smpbo, integrated, vmm,
// vmm_granularity, total_vram, warp_size, supports_cooperative_launch }.
// ----------------------------------------------------------------------------
const ggml_cuda_device_info & ggml_cuda_info() {
static ggml_cuda_device_info info;
static std::once_flag once;
std::call_once(once, []{
int count = 0;
cudaError_t err = cudaGetDeviceCount(&count);
if (err != cudaSuccess) {
fprintf(stderr, "ggml_cuda_info: cudaGetDeviceCount failed: %s\n", cudaGetErrorString(err));
count = 0;
}
if (count > GGML_CUDA_MAX_DEVICES) count = GGML_CUDA_MAX_DEVICES;
info.device_count = count;
for (int i = 0; i < count; i++) {
cudaDeviceProp p;
CUDA_CHECK(cudaGetDeviceProperties(&p, i));
#if defined(GGML_USE_HIP)
unsigned gfx_id = 0;
if (sscanf(p.gcnArchName, "gfx%x", &gfx_id) == 1) {
info.devices[i].cc = GGML_CUDA_CC_OFFSET_AMD + (int)gfx_id;
} else {
info.devices[i].cc = GGML_CUDA_CC_OFFSET_AMD + p.major * 0x100 + p.minor * 0x10;
}
#else
info.devices[i].cc = p.major * 100 + p.minor * 10;
#endif
info.devices[i].nsm = p.multiProcessorCount;
info.devices[i].smpb = p.sharedMemPerBlock;
#if defined(GGML_USE_HIP)
info.devices[i].smpbo = p.sharedMemPerBlock;
#else
info.devices[i].smpbo = p.sharedMemPerBlockOptin;
#endif
info.devices[i].integrated = p.integrated != 0;
info.devices[i].vmm = false;
info.devices[i].vmm_granularity = 0;
info.devices[i].total_vram = p.totalGlobalMem;
info.devices[i].warp_size = p.warpSize;
info.devices[i].supports_cooperative_launch = p.cooperativeLaunch != 0;
}
});
return info;
}
// ggml_cuda_get_device / ggml_cuda_set_device are declared (not defined) in
// common.cuh. We provide thin wrappers.
int ggml_cuda_get_device() {
int dev = 0;
(void)cudaGetDevice(&dev);
return dev;
}
void ggml_cuda_set_device(int device) {
int cur = -1;
(void)cudaGetDevice(&cur);
if (cur != device) {
CUDA_CHECK(cudaSetDevice(device));
}
}
int64_t ggml_time_us() {
using clock = std::chrono::steady_clock;
static const auto t0 = clock::now();
return std::chrono::duration_cast<std::chrono::microseconds>(clock::now() - t0).count();
}
// ----------------------------------------------------------------------------
// ggml_cuda_error: invoked by the CUDA_CHECK / CUBLAS_CHECK macros defined
// in common.cuh on the error path. Marked [[noreturn]] in the declaration
// (common.cuh:155) - abort() satisfies that contract.
// ----------------------------------------------------------------------------
[[noreturn]] void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) {
fprintf(stderr, "CUDA error: %s\n call: %s\n in: %s at %s:%d\n", msg, stmt, func, file, line);
fflush(stderr);
abort();
}
// ----------------------------------------------------------------------------
// Concrete device scratch pool.
// ----------------------------------------------------------------------------
namespace {
/* Thread-local stream that ds4_naive_pool uses for cudaMallocAsync /
* cudaFreeAsync. Defaults to cudaStreamPerThread (preserves prior
* behaviour). Step 8 / CUDA Graphs sets this to the capture stream
* before allocating, so the alloc node lives on the captured stream
* and capture is not invalidated. Set via ds4_pool_set_stream() from
* ds4_mmq.cu wrappers; an explicit stream=0 means the legacy default
* stream so pool ops stay ordered with legacy-stream kernels. */
static thread_local cudaStream_t t_ds4_pool_stream = cudaStreamPerThread;
} // anonymous namespace
extern "C" void ds4_pool_set_stream(cudaStream_t stream) {
t_ds4_pool_stream = stream;
}
extern "C" cudaStream_t ds4_pool_get_stream(void) {
return t_ds4_pool_stream;
}
namespace {
struct ds4_naive_pool : public ggml_cuda_pool {
int device;
explicit ds4_naive_pool(int device) : device(device) {}
#if defined(GGML_USE_HIP)
struct cached_block {
void *ptr;
size_t size;
cudaStream_t stream;
};
std::mutex mutex;
std::vector<cached_block> available;
std::vector<void *> owned;
#endif
void * alloc(size_t size, size_t * actual_size) override {
ggml_cuda_set_device(device);
#if defined(GGML_USE_HIP)
{
std::lock_guard<std::mutex> guard(mutex);
size_t best = available.size();
for (size_t i = 0; i < available.size(); i++) {
if (available[i].stream != t_ds4_pool_stream ||
available[i].size < size) {
continue;
}
if (best == available.size() ||
available[i].size < available[best].size) {
best = i;
}
}
if (best != available.size()) {
const cached_block block = available[best];
available[best] = available.back();
available.pop_back();
if (actual_size) *actual_size = block.size;
return block.ptr;
}
}
void * ptr = nullptr;
CUDA_CHECK(ds4_rocm_allocation_fits(size) ? hipSuccess : hipErrorOutOfMemory);
CUDA_CHECK(cudaMallocAsync(&ptr, size, t_ds4_pool_stream));
{
std::lock_guard<std::mutex> guard(mutex);
owned.push_back(ptr);
}
if (actual_size) *actual_size = size;
return ptr;
#else
void * ptr = nullptr;
CUDA_CHECK(cudaMallocAsync(&ptr, size, t_ds4_pool_stream));
if (actual_size) *actual_size = size;
return ptr;
#endif
}
void free(void * ptr, size_t size) override {
if (!ptr) return;
ggml_cuda_set_device(device);
#if defined(GGML_USE_HIP)
std::lock_guard<std::mutex> guard(mutex);
available.push_back({ptr, size, t_ds4_pool_stream});
#else
CUDA_CHECK(cudaFreeAsync(ptr, t_ds4_pool_stream));
#endif
}
#if defined(GGML_USE_HIP)
~ds4_naive_pool() override {
ggml_cuda_set_device(device);
(void)cudaDeviceSynchronize();
for (void *ptr : owned) (void)cudaFree(ptr);
}
#endif
};
} // anonymous namespace
std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, int /*stream_no*/) {
return std::unique_ptr<ggml_cuda_pool>(new ds4_naive_pool(device));
}
ggml_backend_cuda_context::~ggml_backend_cuda_context() {
if (copy_event) {
(void)cudaEventDestroy(copy_event);
copy_event = nullptr;
}
// streams[][], cublas_handles[], and pools[][] are owned-by-value
// (cudaStream_t and cublasHandle_t are opaque handles - destroying the
// context "should" tear them down, but in our shim ds4 manages streams
// externally and we leave them alone. The pools auto-destruct via
// unique_ptr.).
}
cuda/mmq/ds4_ggml_stubs.h
+303
View File
@@ -0,0 +1,303 @@
// SPDX-License-Identifier: MIT
// ds4_ggml_stubs.h - minimal ggml-API stubs for ds4's vendored mmq kernels.
//
// The mmq.cuh / mma.cuh / vecdotq.cuh / quantize.cuh / mmid.cuh / common.cuh
// files in this directory are vendored verbatim from llama.cpp's ggml-cuda
// backend (MIT, copyright 2023-2026 The ggml authors). They transitively
// #include "ggml.h", "ggml-impl.h", "ggml-cuda.h" - in ds4 those names
// resolve to thin redirect headers in this directory which all #include this
// stubs file.
//
// This file declares the minimum surface of the ggml API that the vendored
// CUDA code references, EXCLUDING what's already provided by common.cuh
// itself (compute-capability constants, MMA flags, ggml_cuda_device_info,
// ggml_cuda_pool, ggml_cuda_pool_alloc, ggml_backend_cuda_context, the
// CUDA_CHECK / CUBLAS_CHECK macros, ggml_cuda_get_device, ggml_cuda_set_device,
// ggml_cuda_info). Those names live in common.cuh and we let it own them.
//
// Things this header DOES provide:
// * GGML_ASSERT / GGML_ABORT / GGML_UNUSED / GGML_UNUSED_VARS / GGML_PAD
// * GGML_MAX_DIMS / GGML_MAX_SRC / GGML_CUDA_NAME / GGML_CUDA_MAX_DEVICES /
// GGML_CUDA_MAX_STREAMS / GGML_LOG_DEBUG
// * enum ggml_type (all 21 mmq type codes - we only USE a subset for V4
// Flash but the switch in mmq.cu's downstream replacement must compile)
// * enum ggml_glu_op (just for the unused mm_fusion_args fields)
// * struct ggml_tensor (complete enough for common.cuh's
// ggml_cuda_concurrent_event::is_valid() to compile - we never call it)
// * int64_t ggml_nbytes(const ggml_tensor *) (stub - never called)
// * int64_t ggml_time_us() (used by USE_CUDA_GRAPH paths we disable)
// * inline ggml_type_size() / ggml_blck_size() lookups
//
// Things ggml-common.h (vendored) owns:
// * ggml_half / ggml_half2 typedefs
// * GGML_EXTENSION macro
// * block_q*, block_iq* struct definitions
#pragma once
#include <cassert>
#include <cstdarg>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
// ----------------------------------------------------------------------------
// Macros
// ----------------------------------------------------------------------------
#ifndef GGML_ASSERT
#define GGML_ASSERT(cond) \
do { if (!(cond)) { \
fprintf(stderr, "GGML_ASSERT(%s) failed at %s:%d\n", #cond, __FILE__, __LINE__); \
abort(); \
} } while (0)
#endif
#ifndef GGML_ABORT
#define GGML_ABORT(fmt, ...) \
do { \
fprintf(stderr, "GGML_ABORT: " fmt " at %s:%d\n", ##__VA_ARGS__, __FILE__, __LINE__); \
abort(); \
} while (0)
#endif
#ifndef GGML_UNUSED
#define GGML_UNUSED(x) ((void)(x))
#endif
// Variadic GGML_UNUSED_VARS: drop up to 12 unused names without warnings.
#ifndef GGML_UNUSED_VARS
#define GGML_UNUSED_VARS_1(_1) GGML_UNUSED(_1)
#define GGML_UNUSED_VARS_2(_1,_2) GGML_UNUSED(_1); GGML_UNUSED(_2)
#define GGML_UNUSED_VARS_3(_1,_2,_3) GGML_UNUSED_VARS_2(_1,_2); GGML_UNUSED(_3)
#define GGML_UNUSED_VARS_4(_1,_2,_3,_4) GGML_UNUSED_VARS_3(_1,_2,_3); GGML_UNUSED(_4)
#define GGML_UNUSED_VARS_5(_1,_2,_3,_4,_5) GGML_UNUSED_VARS_4(_1,_2,_3,_4); GGML_UNUSED(_5)
#define GGML_UNUSED_VARS_6(_1,_2,_3,_4,_5,_6) GGML_UNUSED_VARS_5(_1,_2,_3,_4,_5); GGML_UNUSED(_6)
#define GGML_UNUSED_VARS_7(_1,_2,_3,_4,_5,_6,_7) GGML_UNUSED_VARS_6(_1,_2,_3,_4,_5,_6); GGML_UNUSED(_7)
#define GGML_UNUSED_VARS_8(_1,_2,_3,_4,_5,_6,_7,_8) GGML_UNUSED_VARS_7(_1,_2,_3,_4,_5,_6,_7); GGML_UNUSED(_8)
#define GGML_UNUSED_VARS_9(_1,_2,_3,_4,_5,_6,_7,_8,_9) GGML_UNUSED_VARS_8(_1,_2,_3,_4,_5,_6,_7,_8); GGML_UNUSED(_9)
#define GGML_UNUSED_VARS_10(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10) GGML_UNUSED_VARS_9(_1,_2,_3,_4,_5,_6,_7,_8,_9); GGML_UNUSED(_10)
#define GGML_UNUSED_VARS_11(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11) GGML_UNUSED_VARS_10(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10); GGML_UNUSED(_11)
#define GGML_UNUSED_VARS_12(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11,_12) GGML_UNUSED_VARS_11(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11); GGML_UNUSED(_12)
#define GGML_UNUSED_VARS_PICK(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11,_12,NAME,...) NAME
#define GGML_UNUSED_VARS(...) \
GGML_UNUSED_VARS_PICK(__VA_ARGS__, \
GGML_UNUSED_VARS_12, GGML_UNUSED_VARS_11, GGML_UNUSED_VARS_10, \
GGML_UNUSED_VARS_9, GGML_UNUSED_VARS_8, GGML_UNUSED_VARS_7, \
GGML_UNUSED_VARS_6, GGML_UNUSED_VARS_5, GGML_UNUSED_VARS_4, \
GGML_UNUSED_VARS_3, GGML_UNUSED_VARS_2, GGML_UNUSED_VARS_1)(__VA_ARGS__)
#endif
#ifndef GGML_PAD
#define GGML_PAD(x, n) (((x) + (n) - 1) / (n) * (n))
#endif
#ifndef GGML_MAX_DIMS
#define GGML_MAX_DIMS 4
#endif
#ifndef GGML_MAX_SRC
#define GGML_MAX_SRC 10
#endif
#ifndef GGML_CUDA_NAME
#define GGML_CUDA_NAME "DS4_CUDA"
#endif
#ifndef GGML_CUDA_MAX_DEVICES
#define GGML_CUDA_MAX_DEVICES 16
#endif
#ifndef GGML_CUDA_MAX_STREAMS
#define GGML_CUDA_MAX_STREAMS 8
#endif
#ifndef GGML_LOG_DEBUG
#define GGML_LOG_DEBUG(...) ((void)0)
#endif
// Cuda-graphs are explicitly disabled - ds4 manages its own streams.
#undef GGML_CUDA_USE_GRAPHS
#undef GGML_HIP_GRAPHS
#undef GGML_MUSA_GRAPHS
// GGML_EXTENSION: ggml-common.h provides the canonical definition. We leave
// it undefined here so the vendored header's `#define GGML_EXTENSION
// __extension__` wins.
// ----------------------------------------------------------------------------
// Quantization type enum.
//
// Order matches llama.cpp's enum ggml_type. Values are pinned because the
// mmq switch uses them as case labels.
// ----------------------------------------------------------------------------
enum ggml_type {
GGML_TYPE_F32 = 0,
GGML_TYPE_F16 = 1,
GGML_TYPE_Q4_0 = 2,
GGML_TYPE_Q4_1 = 3,
// GGML_TYPE_Q4_2 / Q4_3 deprecated
GGML_TYPE_Q5_0 = 6,
GGML_TYPE_Q5_1 = 7,
GGML_TYPE_Q8_0 = 8,
GGML_TYPE_Q8_1 = 9,
GGML_TYPE_Q2_K = 10,
GGML_TYPE_Q3_K = 11,
GGML_TYPE_Q4_K = 12,
GGML_TYPE_Q5_K = 13,
GGML_TYPE_Q6_K = 14,
GGML_TYPE_Q8_K = 15,
GGML_TYPE_IQ2_XXS = 16,
GGML_TYPE_IQ2_XS = 17,
GGML_TYPE_IQ3_XXS = 18,
GGML_TYPE_IQ1_S = 19,
GGML_TYPE_IQ4_NL = 20,
GGML_TYPE_IQ3_S = 21,
GGML_TYPE_IQ2_S = 22,
GGML_TYPE_IQ4_XS = 23,
GGML_TYPE_I8 = 24,
GGML_TYPE_I16 = 25,
GGML_TYPE_I32 = 26,
GGML_TYPE_I64 = 27,
GGML_TYPE_F64 = 28,
GGML_TYPE_IQ1_M = 29,
GGML_TYPE_BF16 = 30,
GGML_TYPE_MXFP4 = 39,
GGML_TYPE_NVFP4 = 40,
GGML_TYPE_Q1_0 = 41,
GGML_TYPE_COUNT = 42,
};
enum ggml_glu_op {
GGML_GLU_OP_REGLU,
GGML_GLU_OP_GEGLU,
GGML_GLU_OP_SWIGLU,
GGML_GLU_OP_SWIGLU_OAI, // referenced by mmvq.cu's fused-GLU epilogue
GGML_GLU_OP_COUNT,
};
// ----------------------------------------------------------------------------
// ggml_tensor: complete enough for common.cuh's
// ggml_cuda_concurrent_event::is_valid() to compile cleanly. We NEVER
// instantiate or dereference one of these - the concurrent path is
// disabled.
//
// Field set matches the upstream order/types so cudaGraph node_properties
// (which holds a `ggml_tensor node` by value inside `#ifdef USE_CUDA_GRAPH`)
// also compiles. Sizes are conservative for storage; ds4 never copies into
// these.
// ----------------------------------------------------------------------------
struct ggml_tensor {
enum ggml_type type;
int32_t op; // enum ggml_op (opaque to us)
int32_t flags;
int64_t ne[GGML_MAX_DIMS]; // shape
size_t nb[GGML_MAX_DIMS]; // stride bytes
int32_t op_params[16]; // GGML_MAX_OP_PARAMS / sizeof(int32_t)
struct ggml_tensor * src[GGML_MAX_SRC];
struct ggml_tensor * view_src;
size_t view_offs;
void * data;
char name[64]; // GGML_MAX_NAME
void * extra;
char padding[8];
};
// ggml_nbytes: byte size of tensor data. We never call this; provide a
// stub so common.cuh's is_valid() compiles. If anything does call it the
// returned 0 will surface as an immediate logic error.
static inline int64_t ggml_nbytes(const struct ggml_tensor * /*t*/) { return 0; }
// Microsecond timer (used only inside USE_CUDA_GRAPH paths we disable).
int64_t ggml_time_us();
// ----------------------------------------------------------------------------
// Inline size traits.
//
// Lookup tables hand-aligned with llama.cpp's ggml_type_traits in
// ggml/src/ggml.c. Only types referenced by the vendored switch are needed.
// ----------------------------------------------------------------------------
inline size_t ggml_type_size(enum ggml_type t) {
switch (t) {
case GGML_TYPE_F32: return 4;
case GGML_TYPE_F16: return 2;
case GGML_TYPE_BF16: return 2;
case GGML_TYPE_I8: return 1;
case GGML_TYPE_I16: return 2;
case GGML_TYPE_I32: return 4;
case GGML_TYPE_I64: return 8;
case GGML_TYPE_F64: return 8;
case GGML_TYPE_Q4_0: return 18;
case GGML_TYPE_Q4_1: return 20;
case GGML_TYPE_Q5_0: return 22;
case GGML_TYPE_Q5_1: return 24;
case GGML_TYPE_Q8_0: return 34;
case GGML_TYPE_Q8_1: return 36;
case GGML_TYPE_Q2_K: return 84;
case GGML_TYPE_Q3_K: return 110;
case GGML_TYPE_Q4_K: return 144;
case GGML_TYPE_Q5_K: return 176;
case GGML_TYPE_Q6_K: return 210;
case GGML_TYPE_Q8_K: return 292;
case GGML_TYPE_IQ2_XXS: return 66;
case GGML_TYPE_IQ2_XS: return 74;
case GGML_TYPE_IQ2_S: return 82;
case GGML_TYPE_IQ3_XXS: return 98;
case GGML_TYPE_IQ3_S: return 110;
case GGML_TYPE_IQ1_S: return 50;
case GGML_TYPE_IQ1_M: return 56;
case GGML_TYPE_IQ4_NL: return 18;
case GGML_TYPE_IQ4_XS: return 136;
case GGML_TYPE_MXFP4: return 17;
case GGML_TYPE_NVFP4: return 18;
case GGML_TYPE_Q1_0: return 36;
default: return 0;
}
}
inline int64_t ggml_blck_size(enum ggml_type t) {
switch (t) {
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_I8:
case GGML_TYPE_I16:
case GGML_TYPE_I32:
case GGML_TYPE_I64:
case GGML_TYPE_F64:
return 1;
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
case GGML_TYPE_Q8_1:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
return 32;
case GGML_TYPE_NVFP4:
return 64;
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
case GGML_TYPE_Q8_K:
case GGML_TYPE_IQ2_XXS:
case GGML_TYPE_IQ2_XS:
case GGML_TYPE_IQ2_S:
case GGML_TYPE_IQ3_XXS:
case GGML_TYPE_IQ3_S:
case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ1_M:
case GGML_TYPE_IQ4_XS:
case GGML_TYPE_Q1_0:
return 256;
default:
return 1;
}
}
cuda/mmq/ds4_mmq.cu
+5706
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/ds4_mmq.h
+897
View File
@@ -0,0 +1,897 @@
// SPDX-License-Identifier: MIT
// ds4_mmq.h - public C ABI for ds4's quantized matmul kernels.
//
// All functions are extern "C" so ds4.c / ds4_cuda.cu can call them
// without C++ compilation. Functions return 0 on success and non-zero on
// failure (with stderr error message). Device pointers are caller-owned.
//
// Phase 0: skeleton only. Q8_0 dense entry compiles and instantiates
// mul_mat_q_case<Q8_0> but is not yet wired into ds4_cuda.cu.
// Phase 1: Q8_1 activation quantizer wrapper added.
// Phase 2: Q8_0 dense entry verified against cublas+dequant baseline.
// Phase 3: Q2_K + IQ2_XXS dense entries.
// Phase 4: MoE _id variants of all three.
#pragma once
#if defined(GGML_USE_HIP)
#include "vendors/hip.h"
#else
#include <cuda_runtime.h>
#endif
#include <stddef.h>
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
// One-time init. Sets the current CUDA device and triggers lazy population
// of the device-info singleton. Safe to call repeatedly.
//
// device: CUDA device ordinal (0 for the primary GPU).
// Returns 0 on success.
int ds4_mmq_init(int device);
void ds4_mmq_set_aligned_q81_scratch(void *ptr, size_t bytes);
// Query whether ds4_mmq is willing to handle a given matmul. Returns
// 1 if mmq is faster than dequant+cublas for this shape on this device,
// 0 otherwise (caller should fall back to its existing dequant+cublas path).
//
// Wraps ggml_cuda_should_use_mmq. type_x uses ds4 quant codes which match
// ggml's enum:
// 8 = Q8_0
// 10 = Q2_K
// 16 = IQ2_XXS
//
// ne11: batch dimension (number of activation columns / tokens).
// n_experts: 0 for dense matmul, >0 for MoE (e.g. 256 for V4 Flash).
int ds4_mmq_should_use(int type_x, int64_t ne11, int64_t n_experts);
// Dense matmul entry points. Per-type wrappers that all share the same
// underlying mul_mat_q template, parameterised by the weight quant type.
//
// All three variants compute:
//
// out[col, row] = sum_k W[row, k] * X[k, col] 0 <= row < M, 0 <= col < N
//
// Layouts (matching ggml + llama.cpp mmq conventions, all on device):
// W: [M rows, K cols], row-major, packed in the type-specific block
// format. K must be a multiple of 256.
// X_f32: [N rows, K cols] F32 row-major (logical [K, N] with K
// innermost - i.e. for each "column" col of the logical [K, N]
// matrix, K contiguous floats live at X[col*K .. col*K + K]).
// out_f32: caller-allocated, M*N floats. mmq writes in column-major:
// out[col*M + row]. Callers expecting row-major must transpose.
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_q8_0_dense(
const void * W_q8_0,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
// v0.5 flat-pool p5a: same contract as ds4_mmq_q8_0_dense but the
// activation arrives PRE-QUANTIZED in the block_q8_1_mmq D4 layout
// (producer-emitted, op-for-op equal to quantize_mmq_q8_1<D4>; the fused
// own out_a kernel is the producer). Y_q8_mmq must hold N*(K/128)
// 144-byte blocks (ib = kseg*N + row) plus mmq_x_max tail slack, which
// this entry zeroes. Requires K such that GGML_PAD(K, MATRIX_ROW_PADDING)
// == K. Returns 0 on success; callers fall back to the quantizing entry
// on any non-zero.
int ds4_mmq_q8_0_dense_preq(
const void * W_q8_0,
const void * Y_q8_mmq,
size_t y_bytes,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
// p5a verify instrument: reference quantize (dense_impl parameters) into a
// caller buffer, for byte-diffing producer emits.
int ds4_mmq_q8_0_quantize_ref(
const float * X,
void * y,
size_t y_bytes,
int N,
int K,
cudaStream_t stream);
// Dense Q8_0 D2R on the kind-5 aligned artifact (weight server
// --repack-q8-aligned). Same in/out contract as ds4_mmq_q8_0_dense but W is
// the ALIGNED artifact base ([half dq[nblk]][pad64][int8 qs]), and the shape
// must satisfy M % 128 == 0 && K % 1024 == 0. Callers gate on n_tok scale
// and K <= 4096 (o_proj's K=8192 measured faster on mmq).
int ds4_mmq_q8_0_dense_d2r(
const void * W_aligned,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
// flat-pool p5c: D2R dense over a producer-quantized token-major Y
// (block_q8_1_mmq D4, ib = kseg*N + row, requires GGML_PAD(K) == K).
// y_bytes must cover the payload plus max(mmq_x_max, 128) slack blocks;
// the slack is zeroed here each call (S1.1a). Returns 0 on success;
// callers fall back to the quantizing entry on any non-zero.
int ds4_mmq_q8_0_dense_d2r_preq(
const void * W_aligned,
const void * Y_q8_mmq,
size_t y_bytes,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_q2_K_dense(
const void * W_q2_K,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_dense(
const void * W_iq2_xxs,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_q4_K_dense(
const void * W_q4_K,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_mxfp4_dense(
const void * W_mxfp4,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
// MoE matmul entry points. For each (token, slot-within-token's-top-k) pair
// the kernel computes:
//
// out[col, row] = sum_k W[ids[token, slot], row, k] * X[token, k]
//
// where col = token * n_expert_used + slot, row in [0, M). The caller is
// responsible for any downstream sum-weighted-by-router-weights reduction
// across the n_expert_used dimension (Phase 5 wires this into ds4's
// existing moe_sum_kernel).
//
// Layouts:
// W: device pointer, [n_experts, M rows, K cols] in the
// type-specific block format. Per-expert slab is M*K/blck
// blocks stored contiguously; experts are stacked.
// X_f32: device pointer, [n_tokens, K] F32 row-major (K innermost).
// ids: device pointer, [n_tokens, n_expert_used] int32_t row-major.
// ids[t*n_expert_used + s] is the expert id for token t's
// s-th routing slot. Values must be in [0, n_experts).
// out_f32: caller-allocated, M * n_tokens * n_expert_used floats.
// Column-major: out[col*M + row].
//
// K must be a multiple of 256. n_expert_used must be one of the values
// the vendored mm_ids_helper template specialises on: 2, 4, 6, 8, 16, 32
// (or any other value, which falls back to the generic path). For V4
// Flash, n_expert_used = 6.
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_q8_0_moe(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q2_K_moe(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_moe(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// ds4 (P4 Inc3): same contract as ds4_mmq_q2_K_moe but W_soa is the aligned
// row-pair-SoA artifact (weight server --repack-q2k-aligned, layout in
// ds4_mmq_q2_k_aligned_bytes' comment) instead of the raw block stream. The
// tile loader reads the SoA sections directly -- bit-identical output to the
// raw path, no derepack scratch.
// CONTRACT DIFFERENCE (P3): unlike the raw entries, the output is NOT
// nonfinite-sanitized; the routed-MoE consumers (moe_mmq_swiglu / moe_sum
// with guard_nonfinite=1) sanitize at read, so the standalone whole-buffer
// pass is skipped. New callers must sanitize at consumption.
int ds4_mmq_q2_K_moe_soa(
const void * W_soa,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q4_K_moe(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_mxfp4_moe(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Paired MoE entries. Compute gate AND up over the same activation in a
// single call so the Q8_1 quantize of X (and the mm_ids_helper bookkeeping)
// happens once instead of twice. Both weights must be the same quant type
// and the same shape (M, K, n_experts); out_a / out_b have the same layout
// as a single ds4_mmq_<type>_moe call. Saves one launch of
// quantize_mmq_q8_1_cuda and one ggml_cuda_launch_mm_ids_helper per MoE
// block. See ds4_mmq.cu / routed_moe_launch for the wiring.
//
// Returns 0 on success; on error neither output is guaranteed valid.
int ds4_mmq_iq2_xxs_moe_pair(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// ds4 (P4 Inc3): same contract as ds4_mmq_iq2_xxs_moe_pair but over the
// aligned-SoA artifacts (weight server --repack-iq2-aligned); see
// ds4_mmq_q2_K_moe_soa.
int ds4_mmq_iq2_xxs_moe_pair_soa(
const void * Wa_soa,
const void * Wb_soa,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
/* v0.5 inc-9 (F7): fused target-prefill pipeline over the aligned-SoA
* IQ2_XXS gate/up and Q2_K down artifacts. Builds the expert-major
* assignment map once, runs the paired gate/up MMQs, computes clamp +
* SwiGLU + router weighting in mid_f32, then gathers and quantizes those
* rows for the Q2_K down MMQ through the same ids_dst/expert_bounds. No
* second mm_ids_helper; gate/up/mid/down keep the pair-major layout.
* preserve_reduction retains the raw MMQ accumulation order instead of D2R. */
int ds4_mmq_iq2_xxs_q2_K_moe_fused_soa(
const void * W_gate_soa,
const void * W_up_soa,
const void * W_down_soa,
const float * X_f32,
const int32_t * ids,
const float * router_weights,
float * gate_f32,
float * up_f32,
float * mid_f32,
float * down_f32,
int expert_mid_dim,
int expert_in_dim,
int out_dim,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
int preserve_reduction,
cudaStream_t stream);
/* Aligned-artifact production fast path: gate/up stay in registers, weighted
* SwiGLU is quantized directly into down_q8_scratch, and only the pair-major
* down output is materialized. Caller-owned scratch keeps this hot path free
* of stream-ordered allocations; all three ranges must be distinct, remain
* live until return, and be sized to their LOGICAL segments (never an owning
* arena's capacity). */
/* flat-pool p5c: input_q8_ext, when non-NULL, is a producer-emitted
* TOKEN-COMPACT block_q8_1_mmq buffer of X (ib = kseg*n_tokens + row,
* n_tokens rows) and the internal activation quantize is skipped; the
* gate/up D2R kernel reads it through the p5b ids_src indirection.
* Ignored (classic gathered quantize) when DS4_MMQ_NO_YIND is set. */
int ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_soa(
const void * W_gate_soa,
const void * W_up_soa,
const void * W_down_soa,
const float * X_f32,
const int32_t * ids,
const float * router_weights,
void * input_q8_scratch,
size_t input_q8_scratch_bytes,
void * down_q8_scratch,
size_t down_q8_scratch_bytes,
void * work_scratch,
size_t work_scratch_bytes,
const void * input_q8_ext,
size_t input_q8_ext_bytes,
float * down_f32,
int expert_mid_dim,
int expert_in_dim,
int out_dim,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_scratch_sizes(
int expert_mid_dim,
int expert_in_dim,
int n_tokens,
int n_experts,
int n_expert_used,
size_t *input_q8_bytes,
size_t *down_q8_bytes,
size_t *work_bytes);
int ds4_mmq_q4_K_moe_pair(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_mxfp4_moe_pair(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// MoE vector matmul entries (Step 6). Same signature and semantics as the
// ds4_mmq_<type>_moe entries above, but route through llama.cpp's mmvq
// kernels instead of mmq. mmvq is structurally optimised for small batch
// counts (single-token decode, short prefill), where mmq's tile-based
// approach wastes work on empty columns.
//
// Constraints:
// - n_tokens * something must fit under mmvq's per-arch batch cap
// (MMVQ_MAX_BATCH_SIZE = 8 on Blackwell). Specifically, ncols_dst as
// computed by the wrapper must be <= 8. The wrapper rejects with -1
// if the request is too large.
// - K must be a multiple of 256 (same as the mmq path).
//
// Unlike the mmq path, mmvq consumes a CANONICAL block_q8_1 buffer (not
// the interleaved block_q8_1_mmq the mmq path uses). The wrapper builds
// the canonical buffer internally; callers cannot reuse a Q8_1 buffer
// previously built for the mmq path.
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_q8_0_moe_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q2_K_moe_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_moe_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q4_K_moe_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_mxfp4_moe_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Aligned-SoA IQ2_XXS decode matvec (megakernel program M1-Inc1).
//
// block_iq2_xxs is 66 bytes, so the raw expert stream is only 2-byte aligned
// and every 32-bit code word costs two 16-bit loads. W_aligned is a repacked
// copy of the SAME bytes with the block scales split out and the code stream
// 64B-aligned:
//
// [ __half dq[nblk] ][ pad to 64B ][ uint2 qs[nblk * 8] ]
//
// where nblk = n_experts * M * (K / 256) and block linear order matches the
// raw tensor byte order (expert-major, then row, then block). The weight
// server builds this layout (--repack-iq2-aligned); use
// ds4_mmq_iq2_xxs_aligned_bytes to size or validate an artifact.
//
// Semantics and output layout are identical to ds4_mmq_iq2_xxs_moe_vec at
// n_tokens == 1 (the only supported width; other widths return non-zero so
// the caller can fall back).
uint64_t ds4_mmq_iq2_xxs_aligned_bytes(int M, int K, int n_experts);
// M1-Inc2b: exact inverse of the weight-server repack. Fills raw_out
// (nblk * 66 bytes, raw block_iq2_xxs byte stream) from an aligned artifact,
// device->device on `stream`. Lets the batched/mmq raw-layout consumers run
// from a device scratch instead of the client mmap when the raw spans were
// excluded from the upload.
int ds4_mmq_iq2_xxs_aligned_derepack(
const void * W_aligned,
void * raw_out,
int M,
int K,
int n_experts,
cudaStream_t stream);
// M2 moe-down: aligned row-pair-SoA Q2_K routed-expert decode matvec. Twin
// of mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> at the down-leg call shape
// (n_expert_used == 1: each (token, slot) assignment is its own "token") with
// bit-identical outputs; only the weight layout changes. W_aligned is the
// weight-server --repack-q2k-aligned artifact (DERIVED_Q2_K_ALIGNED_MOE,
// REPLACES the raw range; byte-neutral):
//
// npair = n_experts * (M/2) * (K/256)
// [ uint2 dm2[npair] ][ pad 64B ][ int4 sc4[npair*2] ][ pad 64B ]
// [ uint2 qs2[npair*16] ]
//
// keyed to the kernel's rows_per_block == 2 (see ds4_mmq.cu for the exact
// field packing). Use ds4_mmq_q2_k_aligned_bytes to size or validate an
// artifact. Unsupported shapes return non-zero so the caller can fall back.
uint64_t ds4_mmq_q2_k_aligned_bytes(int M, int K, int n_experts);
int ds4_mmq_q2_K_aligned_moe_vec(
const void * W_aligned,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Exact inverse of the weight-server repack: fills raw_out (nblk * 84 bytes,
// raw block_q2_K stream) from an aligned artifact, device->device on
// `stream`, for the batched/mmq raw-layout consumers (same role as
// ds4_mmq_iq2_xxs_aligned_derepack).
int ds4_mmq_q2_K_aligned_derepack(
const void * W_aligned,
void * raw_out,
int M,
int K,
int n_experts,
cudaStream_t stream);
// M1-Inc3: aligned-SoA Q8_0 dense decode matvec. Artifact layout
// [__half dq[nblk]][pad to 64B][int8 qs[nblk*32]], nblk = M * (K/32), block
// order equal to the raw tensor byte order (weight server
// --repack-q8-aligned; raw spans stay served — dense artifacts duplicate,
// they do not replace). n_tokens == 1 and K % 1024 == 0 only; other shapes
// return non-zero so the caller can fall back to ds4_mmq_q8_0_dense_vec.
uint64_t ds4_mmq_q8_0_aligned_bytes(int M, int K);
int ds4_mmq_q8_0_aligned_dense_vec(
const void * W_aligned,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
// Decode-only pair: quantize the shared activation once and compute both
// aligned projections in one row-concatenated launch.
int ds4_mmq_q8_0_aligned_dense_vec_pair(
const void * W0_aligned,
const void * W1_aligned,
const float * X_f32,
float * out0_f32,
float * out1_f32,
int M0,
int M1,
int K,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_aligned_moe_vec(
const void * W_aligned,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// M1-Inc2 variants over the same aligned artifacts (n_tokens == 1 only).
//
// _pair_vec: one activation quantize + one launch computes both raw gate and
// up outputs (nonfinite zeroed in-kernel; no sanitize pass needed). The
// caller still runs its clamp-aware SwiGLU.
//
// _gate_up_mid_vec: additionally folds the clamp/SwiGLU/router-weight
// epilogue (identical semantics to ds4_mmq_moe_gate_up_mid_q8_1_qwarp32) and
// writes mid[slot * M + row] directly.
int ds4_mmq_iq2_xxs_aligned_moe_pair_vec(
const void * W_gate_aligned,
const void * W_up_aligned,
const float * X_f32,
const int32_t * ids,
float * gate_out,
float * up_out,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(
const void * W_gate_aligned,
const void * W_up_aligned,
const float * X_f32,
const int32_t * ids,
const float * weights,
float * mid_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
cudaStream_t stream);
// Fused down+sum vector entries for routed MoE with top_k=6. These preserve
// the canonical Q8_1 activation quantization used by the regular _moe_vec
// path, but avoid materializing [token, slot, out_dim] down results and avoid
// the separate slot-sum kernel. The caller must already have baked router
// weights into X_f32, so the helper computes:
//
// out[token, row] = sum_slot W[ids[token, slot], row, :] @ X[token, slot, :]
//
// Layouts:
// W: [n_experts, M rows, K cols] in Q2_K or Q4_K blocks
// X_f32: [n_tokens * 6, K] F32 row-major
// ids: [n_tokens, 6] int32 row-major
// out_f32: [n_tokens, M] F32 row-major
//
// Constraints:
// - n_expert_used must be 6
// - K must be a multiple of 256
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_q2_K_moe_down_sum6_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q4_K_moe_down_sum6_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_mxfp4_moe_down_sum6_vec(
const void * W,
const float * X_f32,
const int32_t * ids,
float * out_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Fused gate+up+SwiGLU vector entries for routed MoE with top_k=6. These use
// canonical Q8_1 activation quantization and compute weighted mid directly:
//
// mid[token, slot, row] =
// silu(clamp_gate(W_gate[expert,row] @ X[token]))
// * clamp_up(W_up[expert,row] @ X[token])
// * router_weight[token, slot]
//
// Clamp semantics match ds4_cuda.cu: gate is capped above, up is clamped to
// [-clamp, clamp], and no clamp is applied when clamp <= 1e-6.
//
// Layouts:
// W_gate/W_up: [n_experts, M rows, K cols] in IQ2_XXS or Q4_K blocks
// X_f32: [n_tokens, K] F32 row-major
// ids: [n_tokens, 6] int32 row-major
// weights: [n_tokens, 6] F32 row-major
// mid_f32: [n_tokens * 6, M] F32 row-major
//
// Constraints:
// - n_expert_used must be 6
// - K must be a multiple of 256
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_iq2_xxs_moe_gate_up_mid_vec(
const void * W_gate,
const void * W_up,
const float * X_f32,
const int32_t * ids,
const float * weights,
float * mid_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
cudaStream_t stream);
int ds4_mmq_q4_K_moe_gate_up_mid_vec(
const void * W_gate,
const void * W_up,
const float * X_f32,
const int32_t * ids,
const float * weights,
float * mid_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
cudaStream_t stream);
int ds4_mmq_mxfp4_moe_gate_up_mid_vec(
const void * W_gate,
const void * W_up,
const float * X_f32,
const int32_t * ids,
const float * weights,
float * mid_f32,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
float clamp,
cudaStream_t stream);
// Pair-fused MoE vector matmul entries (Step 6). Computes
//
// out[col, row] = (W_a[ids, row, :] @ X[token, :])
// * silu(W_b[ids, row, :] @ X[token, :])
//
// in a SINGLE mmvq launch via mmvq's built-in fusion (fusion.gate = W_b,
// fusion.glu_op = GGML_GLU_OP_SWIGLU). The kernel applies silu to the
// fusion.gate matmul and multiplies into the main matmul: pass the
// SwiGLU "up" weights as W_a and the SwiGLU "gate" weights as W_b to
// match ds4's expected silu(gate)*up semantics. The DeepSeek V4 clamp
// and router-weight multiplication are NOT applied by the kernel - the
// caller is expected to apply them as a small post-process (or to skip
// clamp if clamp==0).
//
// Constraints:
// - n_tokens = 1 ONLY. mmvq supports fusion only at ncols_dst = 1.
// - K must be a multiple of 256.
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_iq2_xxs_moe_pair_vec(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_silu,
int M,
int K,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q4_K_moe_pair_vec(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_silu,
int M,
int K,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Raw pair MoE vector entries. They quantize X to canonical Q8_1 once, then
// run the same mmvq matvec kernel twice to produce the unfused gate and up
// outputs. This preserves the caller's clamp-aware SwiGLU epilogue while
// avoiding the duplicate Q8_1 quantize/scratch setup of two separate
// ds4_mmq_*_moe_vec calls.
//
// Layout and constraints match ds4_mmq_*_moe_vec:
// out_[token * n_expert_used + slot, row]
// K must be a multiple of 256.
int ds4_mmq_iq2_xxs_moe_pair_raw_vec(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
int ds4_mmq_q4_K_moe_pair_raw_vec(
const void * W_a,
const void * W_b,
const float * X_f32,
const int32_t * ids,
float * out_a,
float * out_b,
int M,
int K,
int n_tokens,
int n_experts,
int n_expert_used,
cudaStream_t stream);
// Dense vector matmul entry (Step 6). Same shape semantics as
// ds4_mmq_q8_0_dense but routed through mmvq for batch counts that
// favour the vec path (n_tokens <= 8 on Blackwell).
//
// Returns 0 on success, non-zero on validation or launch failure.
int ds4_mmq_q8_0_dense_vec(
const void * W_q8_0,
const float * X_f32,
float * out_f32,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_q4_K_dense_pair_vec(
const void * W0_q4_K,
const void * W1_q4_K,
const float * X_f32,
float * out0_f32,
float * out1_f32,
int M,
int K,
cudaStream_t stream);
// Set the thread-local stream used for internal stream-ordered scratch.
// Defaults to cudaStreamPerThread.
// Step 8 (CUDA Graphs) calls this with the capture stream so pool
// allocations land on the captured stream and don't invalidate capture.
// Pass cudaStreamPerThread to restore the default; NULL selects the legacy
// default stream.
void ds4_pool_set_stream(cudaStream_t stream);
#ifdef __cplusplus
} // extern "C"
#endif
cuda/mmq/ds4_mmq_d2r.cu
+3374
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/ds4_mmq_d2r.cuh
+95
View File
@@ -0,0 +1,95 @@
// SPDX-License-Identifier: MIT
// Internal launcher for the gated D2R Q2_K MoE down-GEMM path.
#pragma once
#if defined(GGML_USE_HIP)
#include "vendors/hip.h"
#else
#include <cuda_runtime.h>
#endif
#include <stddef.h>
#include <stdint.h>
bool ds4_mmq_q2_K_moe_d2r_available(int cc);
bool ds4_mmq_iq2_xxs_moe_d2r_available(int cc);
size_t ds4_mmq_q2_K_moe_d2r_scratch_bytes(int64_t ncols_max, int n_experts);
size_t ds4_mmq_iq2_xxs_moe_d2r_pair_scratch_bytes(int64_t ncols_max, int n_experts);
size_t ds4_mmq_iq2_xxs_moe_d2r_fused_scratch_bytes(
int64_t ncols_max, int n_experts);
int ds4_mmq_q2_K_moe_d2r_launch(
const void * W_soa,
int64_t soa_blocks,
const void * q8,
const int32_t * ids_dst,
const int32_t * expert_bounds,
float * out,
int M,
int K,
int64_t ne_get_rows,
int n_experts,
void * worklist_scratch,
size_t worklist_scratch_bytes,
cudaStream_t stream);
bool ds4_mmq_q8_0_dense_d2r_available(int cc);
// Dense Q8_0 D2R on the kind-5 aligned artifact (--repack-q8-aligned).
// q8 = block_q8_1_mmq D4 activation buffer ([k128][col], stride N cols,
// over-allocated >= 128 blocks past N*K/128 for the guarded last col tile).
int ds4_mmq_q8_0_dense_d2r_launch(
const void * W_aligned,
const void * q8,
float * out,
int M,
int N,
int K,
cudaStream_t stream);
int ds4_mmq_iq2_xxs_moe_d2r_pair_launch(
const void * gate_soa,
const void * up_soa,
int64_t soa_blocks,
const void * q8,
const int32_t * ids_dst,
const int32_t * expert_bounds,
float * out_gate,
float * out_up,
int M,
int K,
int64_t ne_get_rows,
int n_experts,
void * worklist_scratch,
size_t worklist_scratch_bytes,
cudaStream_t stream);
// Complete target-prefill gate/up path: both IQ2_XXS projections share one
// activation tile, then sanitize + clamp + SwiGLU + routing weight are folded
// directly into the expert-major Q8_1 D2S6 input consumed by Q2_K down.
//
// flat-pool p5b: with ids_src != NULL, input_q8 is TOKEN-COMPACT (n_tokens
// rows quantized once, no gather) and each assignment column stages its
// token's row through ids_src[col]; with ids_src == NULL, input_q8 is the
// classic slot-gathered buffer (ne_get_rows rows) and n_tokens is ignored.
int ds4_mmq_iq2_xxs_moe_d2r_fused_launch(
const void * gate_soa,
const void * up_soa,
int64_t soa_blocks,
const void * input_q8,
const int32_t * ids_src,
int n_tokens,
const int32_t * ids_dst,
const int32_t * expert_bounds,
const float * router_weights,
void * down_q8,
int M,
int K,
int64_t ne_get_rows,
int n_experts,
float clamp,
void * worklist_scratch,
size_t worklist_scratch_bytes,
cudaStream_t stream);
cuda/mmq/ds4_repack.cu
+1160
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/ds4_repack.h
+133
View File
@@ -0,0 +1,133 @@
/* ds4_repack: the aligned-SoA derived-weight layout library shared by the
* weight server (tools/ds4_weight_server.cu) and the engine's in-process
* self-load build (ds4_cuda.cu). The GGUF tensor catalog walk, the repack
* candidate predicates, the S5 parallel repack driver, the FNV-1a artifact
* hash, and the three aligned repack builds (IQ2_XXS gate/up, Q2_K down,
* Q8_0 dense) live here so both producers emit bit-identical artifacts
* (gate: --repack-hash / DS4_WS_REPACK_HASH FNV lines). Layout contracts
* are shared with the consuming kernels in cuda/mmq/ds4_mmq.h. */
#ifndef DS4_REPACK_H
#define DS4_REPACK_H
#include <cstdint>
#include <string>
#include <vector>
struct ds4_repack_file {
int fd = -1;
int direct_fd = -1;
const uint8_t *data = nullptr;
uint64_t size = 0;
uint64_t direct_align = 1;
};
struct ds4_repack_span {
uint64_t off = 0;
uint64_t end = 0;
};
struct ds4_repack_tensor {
std::string name;
uint32_t type = 0;
uint32_t ndim = 0;
uint64_t dims[8] = {};
uint64_t elements = 0;
uint64_t off = 0;
uint64_t bytes = 0;
};
/* Aligned-artifact kinds. Numbering is shared with the weight-server
* manifest format and the engine's cuda_derived_kind (ds4_cuda.cu); the
* layout comments live on the builders in ds4_repack.cu. */
enum {
DS4_REPACK_Q8_0_F16_COLMAJOR = 2,
DS4_REPACK_IQ2_XXS_ALIGNED_MOE = 4,
DS4_REPACK_Q8_0_ALIGNED_DENSE = 5,
DS4_REPACK_Q2_K_ALIGNED_MOE = 6,
};
struct ds4_repack_artifact {
const ds4_repack_tensor *t = nullptr; /* borrowed from the caller's catalog */
uint32_t kind = 0; /* DS4_REPACK_* */
uint64_t bytes = 0; /* artifact byte size */
uint64_t in_dim = 0; /* dims[0] */
uint64_t out_dim = 0; /* dims[1] */
uint32_t group_count = 0; /* dims[2] for MoE kinds, 1 for dense,
0 for f16 colmajor (the engine's
lazy-cache lookup passes 0) */
void *dev = nullptr; /* device artifact buffer */
};
/* Optional device-allocation hooks (weight server VMM backend). alloc must
* fill art->dev with at least art->bytes of device memory and return true;
* free releases an allocation made by alloc. Leave both NULL for plain
* cudaMalloc/cudaFree. */
typedef bool (*ds4_repack_alloc_fn)(void *ctx, ds4_repack_artifact *art);
typedef void (*ds4_repack_free_fn)(void *ctx, ds4_repack_artifact *art);
struct ds4_repack_build_args {
const char *log_prefix = "ds4"; /* stderr line prefix ("ds4_weight_server" | "ds4") */
const char *model_id = "base"; /* log label */
const char *path = nullptr; /* GGUF path (each builder re-maps it) */
const std::vector<ds4_repack_tensor> *records = nullptr;
int device = 0;
uint64_t copy_chunk_bytes = 0; /* staged read chunk (weight server --copy-chunk-mb) */
ds4_repack_alloc_fn alloc_fn = nullptr;
ds4_repack_free_fn free_fn = nullptr;
void *alloc_ctx = nullptr;
};
bool ds4_repack_map_file(const char *log_prefix, const char *path, ds4_repack_file &m);
void ds4_repack_unmap_file(ds4_repack_file &m);
/* Parse the GGUF header + tensor table into whole-tensor spans and records
* with absolute file offsets. Either out pointer may be NULL. */
bool ds4_repack_collect_catalog(const char *log_prefix,
const ds4_repack_file &m,
std::vector<ds4_repack_span> *spans,
std::vector<ds4_repack_tensor> *records);
bool ds4_repack_pread_full(int fd, void *buf, uint64_t bytes, uint64_t offset);
/* Read [file_off, file_off+bytes) into stage, preferring the O_DIRECT fd;
* *payload points at the requested bytes inside stage. */
bool ds4_repack_read_stage(const ds4_repack_file &m, void *stage, uint64_t stage_bytes,
uint64_t file_off, uint64_t bytes, const char **payload);
/* Candidate predicates for the three aligned repacks. */
bool ds4_repack_iq2_candidate(const ds4_repack_tensor &t);
bool ds4_repack_q2k_candidate(const ds4_repack_tensor &t);
bool ds4_repack_q8_candidate(const ds4_repack_tensor &t);
bool ds4_repack_q8_f16_candidate(const ds4_repack_tensor &t);
/* S5 driver configuration. Unset (0 / false) defers to the env knobs
* DS4_WS_REPACK_THREADS and DS4_WS_REPACK_HASH. */
void ds4_repack_set_threads(int n);
void ds4_repack_set_hash(bool on);
uint64_t ds4_repack_fnv1a(const unsigned char *p, uint64_t n, uint64_t h);
/* Gather raw cached experts into one aligned artifact on the default stream.
* slots is a device array of count valid source expert indices. The caller
* sizes dst using the matching ds4_mmq_*_aligned_bytes layout contract. */
bool ds4_repack_selected_experts(void *dst, const void *raw, const int32_t *slots,
uint32_t kind, uint32_t rows, uint32_t cols,
uint32_t count);
/* Build every candidate's aligned artifact. On success appends the built
* artifacts to out, adds their byte total to *repacked_bytes_out (may be
* NULL), and returns true; on any failure releases every artifact this call
* allocated and returns false without touching out. */
bool ds4_repack_build_q8_aligned(const ds4_repack_build_args &a,
std::vector<ds4_repack_artifact> &out,
uint64_t *repacked_bytes_out);
bool ds4_repack_build_iq2_aligned(const ds4_repack_build_args &a,
std::vector<ds4_repack_artifact> &out,
uint64_t *repacked_bytes_out);
bool ds4_repack_build_q2k_aligned(const ds4_repack_build_args &a,
std::vector<ds4_repack_artifact> &out,
uint64_t *repacked_bytes_out);
bool ds4_repack_build_q8_f16(const ds4_repack_build_args &a,
std::vector<ds4_repack_artifact> &out,
uint64_t *repacked_bytes_out);
#endif /* DS4_REPACK_H */
cuda/mmq/ggml-common.h
+1900
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/ggml-cuda.h
+4
View File
@@ -0,0 +1,4 @@
// SPDX-License-Identifier: MIT
// Thin redirect - see ggml.h.
#pragma once
#include "ds4_ggml_stubs.h"
cuda/mmq/ggml-impl.h
+4
View File
@@ -0,0 +1,4 @@
// SPDX-License-Identifier: MIT
// Thin redirect - see ggml.h.
#pragma once
#include "ds4_ggml_stubs.h"
cuda/mmq/ggml.h
+5
View File
@@ -0,0 +1,5 @@
// SPDX-License-Identifier: MIT
// Thin redirect so the vendored common.cuh / vecdotq.cuh / mmid.cu / etc.
// can write #include "ggml.h" unchanged.
#pragma once
#include "ds4_ggml_stubs.h"
cuda/mmq/mma.cuh
+1456
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/mmid.cu
+339
View File
@@ -0,0 +1,339 @@
#include "common.cuh"
#include "mmid.cuh"
// To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each.
struct mm_ids_helper_store {
uint32_t data;
__device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
data = (it & 0x003FFFFF) | (iex_used << 22);
}
__device__ uint32_t it() const {
return data & 0x003FFFFF;
}
__device__ uint32_t iex_used() const {
return data >> 22;
}
};
static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store");
// Helper function for mul_mat_id, converts ids to a more convenient format.
// ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert.
// ids_dst describes the same mapping but for the dst tensor.
// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
template <int n_expert_used_template>
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
extern __shared__ char data_mm_ids_helper[];
mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
int nex_prev = 0; // Number of columns for experts with a lower index.
int it_compact = 0; // Running index for the compact slice of this expert.
if constexpr (n_expert_used_template == 0) {
// Generic implementation:
for (int it = 0; it < n_tokens; ++it) {
int iex_used = -1; // The index at which the expert is used, if any.
for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
const int expert_used = ids[it*si1 + iex];
nex_prev += expert_used < expert;
if (expert_used == expert) {
iex_used = iex;
}
}
if (iex_used != -1) {
store[it_compact] = mm_ids_helper_store(it, iex_used);
}
if (warp_reduce_any<warp_size>(iex_used != -1)) {
it_compact++;
}
}
} else {
// Implementation optimized for specific numbers of experts used:
static_assert(n_expert_used == 6 || warp_size % n_expert_used == 0, "bad n_expert_used");
const int neu_padded = n_expert_used == 6 ? 8 : n_expert_used; // Padded to next higher power of 2.
for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
const int it = it0 + threadIdx.x / neu_padded;
const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
ids[it*si1 + iex] : INT_MAX;
const int iex_used = expert_used == expert ? iex : -1;
nex_prev += expert_used < expert;
// Whether the threads at this token position have used the expert:
const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
// Do a scan over threads at lower token positions in warp to get the correct index for writing data:
int it_compact_add_lower = 0;
#pragma unroll
for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
if (threadIdx.x >= static_cast<unsigned int>(offset)) {
it_compact_add_lower += tmp;
}
}
if (iex_used != -1) {
store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
}
// The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
}
}
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
// Post-Volta independent thread scheduling: the store[] writes above are done by
// some warp lanes and read below by other lanes. Without an explicit warp barrier
// those shared-memory writes are not guaranteed visible to the cross-lane reads,
// so a lane could read a stale/uninitialized store[] slot -> wrong compacted
// expert ids -> nondeterministic MoE routing. On GB10 (sm_121) this realized as
// the cont-multiseq non-determinism + BOS-spam (compute-sanitizer racecheck:
// "RAW hazard at __shared__" between the store writes and these reads).
__syncwarp();
for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
const mm_ids_helper_store store_it = store[itc];
const int it = store_it.it();
const int iex_used = store_it.iex_used();
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
}
if (threadIdx.x != 0) {
return;
}
expert_bounds[expert] = nex_prev;
if (expert < static_cast<int>(gridDim.x) - 1) {
return;
}
expert_bounds[gridDim.x] = nex_prev + it_compact;
}
// ds4 local (P5): large-n variant of mm_ids_helper with NO shared-memory
// staging. The smem kernel above stages the compact (it, iex_used) list in
// n_tokens*4 B of dynamic shared memory, capping n_tokens at smpbo/4 (~25k on
// GB10); the routed-MoE down matmul passes n_tokens = assignment rows
// (6x the forward width), so 8192-row prefill chunks hit 48384 "tokens" and
// the whole MoE block used to fall back to the pre-mmq expert-tile kernels.
// This variant runs the SAME per-expert scan twice: pass 0 only counts
// (nex_prev, bucket size), pass 1 re-scans and writes ids_src1/ids_dst
// directly at their final offsets. Bit-identical to the smem kernel by
// construction: identical (it, iex_used) tuples in identical token order and
// identical output expressions — the (lossless) 22/10-bit store round-trip is
// simply removed. The rescan is cheap: the ids array at the target shape is
// ~190 KB and L2-resident across all expert blocks. Like the smem kernel,
// behavior is undefined if one token lists the same expert in multiple slots
// (the router's top-k is without replacement, so this cannot occur).
template <int n_expert_used_template>
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper_global(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
int nex_prev = 0; // Number of columns for experts with a lower index.
int it_compact_count = 0; // Bucket size for this expert (pass-0 result).
#pragma unroll 1
for (int pass = 0; pass < 2; ++pass) {
int it_compact = 0; // Running index for the compact slice of this expert.
if constexpr (n_expert_used_template == 0) {
// Generic implementation:
for (int it = 0; it < n_tokens; ++it) {
int iex_used = -1; // The index at which the expert is used, if any.
for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
const int expert_used = ids[it*si1 + iex];
if (pass == 0) {
nex_prev += expert_used < expert;
}
if (expert_used == expert) {
iex_used = iex;
}
}
if (pass == 1 && iex_used != -1) {
ids_src1[nex_prev + it_compact] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + it_compact] = it*n_expert_used + iex_used;
}
if (warp_reduce_any<warp_size>(iex_used != -1)) {
it_compact++;
}
}
} else {
// Implementation optimized for specific numbers of experts used:
static_assert(n_expert_used == 6 || warp_size % n_expert_used == 0, "bad n_expert_used");
const int neu_padded = n_expert_used == 6 ? 8 : n_expert_used; // Padded to next higher power of 2.
for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
const int it = it0 + threadIdx.x / neu_padded;
const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
ids[it*si1 + iex] : INT_MAX;
const int iex_used = expert_used == expert ? iex : -1;
if (pass == 0) {
nex_prev += expert_used < expert;
}
// Whether the threads at this token position have used the expert:
const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
// Do a scan over threads at lower token positions in warp to get the correct index for writing data:
int it_compact_add_lower = 0;
#pragma unroll
for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
if (threadIdx.x >= static_cast<unsigned int>(offset)) {
it_compact_add_lower += tmp;
}
}
if (pass == 1 && iex_used != -1) {
const int itc = it_compact + it_compact_add_lower;
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
}
// The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
}
}
if (pass == 0) {
it_compact_count = it_compact;
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
}
}
if (threadIdx.x != 0) {
return;
}
expert_bounds[expert] = nex_prev;
if (expert < static_cast<int>(gridDim.x) - 1) {
return;
}
expert_bounds[gridDim.x] = nex_prev + it_compact_count;
}
// ds4 local: kill switch for the large-n global path (DS4_MMID_LARGE=0).
// With the switch off the ds4_mmq.cu callers refuse past-cap shapes exactly
// as before (whole-MoE fallback to the expert-tile kernels).
bool ds4_mmid_large_enabled(void) {
static int cached = -1;
if (cached < 0) {
const char * env = getenv("DS4_MMID_LARGE");
cached = !(env && env[0] == '0');
}
return cached != 0;
}
template <int n_expert_used_template>
static void launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
GGML_ASSERT(n_tokens < (1 << 22) && "too few bits in mm_ids_helper_store");
GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
const int id = ggml_cuda_get_device();
const int warp_size = ggml_cuda_info().devices[id].warp_size;
const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
const dim3 num_blocks(n_experts, 1, 1);
const dim3 block_size(warp_size, 1, 1);
const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
// ds4 local (P5): past the smem cap, take the two-pass global variant
// (bit-identical outputs, see mm_ids_helper_global). One-shot stderr
// line = path proof for the gate harnesses.
if (nbytes_shared > smpbo) {
static bool logged = false;
if (!logged) {
logged = true;
fprintf(stderr, "ds4: mm_ids_helper large-n global path engaged (P5, n_tokens=%d > cap %zu)\n",
n_tokens, smpbo / sizeof(mm_ids_helper_store));
}
mm_ids_helper_global<n_expert_used_template><<<num_blocks, block_size, 0, stream>>>
(ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
return;
}
CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
(ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
}
// ds4 local: kill switch for the case-1 fast path below (DS4_MMID_CASE1=0).
static bool ds4_mmid_case1_enabled() {
static int cached = -1;
if (cached < 0) {
const char * env = getenv("DS4_MMID_CASE1");
cached = !(env && env[0] == '0');
}
return cached != 0;
}
void ggml_cuda_launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
switch (n_expert_used) {
case 1:
// ds4 local: the routed-MoE down matmul reinterprets (token, slot)
// assignment rows as single-expert "tokens" (ds4_mmq.cu
// ds4_mmq_moe_impl, n_expert_used=1). Without this case it fell to
// the generic <0> template: one warp per expert scanning all
// assignment rows with a single active lane -> 22.5 ms/launch at
// W4096 prefill (2.90 s of a 12k admission). The optimized template
// at neu_padded=1 covers 32 rows/iteration and emits bit-identical
// id maps (proto_mm_ids.cu: parity on uniform/skewed/2%-invalid/
// decode shapes, 20.4x at the W4096 shape). DS4_MMID_CASE1=0
// reverts to the generic path.
if (ds4_mmid_case1_enabled()) {
launch_mm_ids_helper< 1>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
} else {
launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
}
break;
case 2:
launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
case 4:
launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
case 6:
launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
case 8:
launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
case 16:
launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
case 32:
launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
default:
launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
break;
}
}
cuda/mmq/mmid.cuh
+9
View File
@@ -0,0 +1,9 @@
#pragma once
void ggml_cuda_launch_mm_ids_helper(
const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, cudaStream_t stream);
// ds4 local (P5): whether the large-n global-memory mm_ids path is enabled
// (default on; DS4_MMID_LARGE=0 reverts callers to the past-cap refusal).
bool ds4_mmid_large_enabled(void);
cuda/mmq/mmq.cuh
+4500
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/mmvq.cu
+1186
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/mmvq.cuh
+34
View File
@@ -0,0 +1,34 @@
#pragma once
#include "common.cuh"
#define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels.
// Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID,
// based on the quantization type and GPU architecture (compute capability).
int get_mmvq_mmid_max_batch(ggml_type type, int cc);
// ds4: direct switch over the per-type templated mmvq dispatch. Was
// `static` upstream; we promote it so ds4_mmq.cu can call it without
// constructing ggml_tensor structs. Args match the original signature
// exactly.
void mul_mat_vec_q_switch_type(
const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
const int ncols_x, const int nrows_x, const int ncols_dst,
const int stride_row_x, const int stride_col_y, const int stride_col_dst,
const int nchannels_x, const int nchannels_y, const int nchannels_dst,
const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
const int ids_stride, cudaStream_t stream);
#ifdef DS4_MMVQ_INCLUDE_GGML_ENTRIES
// Original upstream ggml-graph entries. Disabled by default in the ds4
// vendor; re-enable for parity comparisons against llama.cpp.
void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx,
const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr);
void ggml_cuda_op_mul_mat_vec_q(
ggml_backend_cuda_context & ctx,
const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
const int64_t src1_padded_row_size, cudaStream_t stream);
#endif
cuda/mmq/quantize.cu
+443
View File
@@ -0,0 +1,443 @@
#include "quantize.cuh"
#include <cstdint>
__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
static __global__ void quantize_q8_1(
const float * __restrict__ x, void * __restrict__ vy,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const uint32_t ne1, const uint3 ne2) {
const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
if (i0 >= ne0) {
return;
}
const int64_t i3 = fastdiv(blockIdx.z, ne2);
const int64_t i2 = blockIdx.z - i3*ne2.z;
const int64_t i1 = blockIdx.y;
const int64_t & i00 = i0;
const int64_t & i01 = i1;
const int64_t & i02 = i2;
const int64_t & i03 = i3;
const int64_t i_cont = ((i3*ne2.z + i2) * ne1 + i1) * ne0 + i0;
block_q8_1 * y = (block_q8_1 *) vy;
const int64_t ib = i_cont / QK8_1; // block index
const int64_t iqs = i_cont % QK8_1; // quant index
const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
float amax = fabsf(xi);
float sum = xi;
amax = warp_reduce_max<QK8_1>(amax);
sum = warp_reduce_sum<QK8_1>(sum);
const float d = amax / 127.0f;
const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
y[ib].qs[iqs] = q;
if (iqs > 0) {
return;
}
y[ib].ds = make_half2(d, sum);
}
__device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
if (!(amax > 0.0f)) {
return 0;
}
// FP4 E2M1: max exponent (unbiased) is 2.
constexpr int FP4_E2M1_EMAX = 2;
const float e = log2f(amax);
// "even" -> round-to-nearest integer, ties-to-even
const int e_int = __float2int_rn(e);
const int shared_exp = e_int - FP4_E2M1_EMAX;
int biased = shared_exp + 127;
biased = max(biased, 0);
biased = min(biased, 254);
return static_cast<uint8_t>(biased);
}
static __global__ void quantize_mmq_nvfp4(
const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2) {
#if defined(BLACKWELL_MMA_AVAILABLE)
const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
if (i0_base >= ne0) {
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
const int64_t i01 = ids ? ids[i1] : i1;
const int64_t k_block = i0_base / QK_K;
const int64_t blocks_per_col = (ne0 + QK_K - 1) / QK_K;
if (k_block >= blocks_per_col) {
return;
}
const int64_t ib = blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x;
block_fp4_mmq * y = (block_fp4_mmq *) vy;
block_fp4_mmq * yb = y + ib;
const int sub = (i0_base % QK_K) / QK_NVFP4_SUB;
float vals_raw[QK_NVFP4_SUB];
float amax_raw = 0.0f;
const int64_t base_idx = i3 * s03 + i2 * s02 + i01 * s01;
#pragma unroll
for (int k = 0; k < QK_NVFP4_SUB; k++) {
const int64_t i00 = i0_base + k;
if (i00 < ne00) {
const float v = x[base_idx + i00];
vals_raw[k] = v;
amax_raw = fmaxf(amax_raw, fabsf(v));
} else {
vals_raw[k] = 0.0f;
}
}
static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2};
const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
float best_err = FLT_MAX;
uint8_t fp8_code = 0;
float subblock_scale = 0.0f;
#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
for (int i = 0; i < 5; i++) {
const int test_code = first_fp8_code + test_offsets[i];
if (test_code < 0 || test_code > 0x7e) {
continue;
}
const uint8_t code = (uint8_t) test_code;
const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
float cur_err = 0.0f;
#pragma unroll
for (int k = 0; k < QK_NVFP4_SUB; ++k) {
const float v = vals_raw[k];
const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
cur_err = fmaf(err_diff, err_diff, cur_err);
}
if (cur_err < best_err) {
best_err = cur_err;
fp8_code = test_code;
subblock_scale = test_scale;
}
}
const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
uint32_t q0 = 0;
uint32_t q1 = 0;
#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 0], inv_scale) << (8 * k);
q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 8], inv_scale) << (8 * k + 4);
q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 4], inv_scale) << (8 * k);
q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
}
uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
yqs[2 * sub + 0] = q0;
yqs[2 * sub + 1] = q1;
reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
#else
NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
#endif // defined(BLACKWELL_MMA_AVAILABLE)
}
// quantize values in the format mxfp4 is stored which is interleaved nibbles
// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const int32_t * __restrict__ ids,
void * __restrict__ vy,
const int64_t ne00,
const int64_t s01,
const int64_t s02,
const int64_t s03,
const int64_t ne0,
const int ne1,
const int ne2) {
constexpr int vals_per_scale = 32;
constexpr int vals_per_warp = 2 * vals_per_scale; // Each warp processes 2 blocks of 32 = 64 values
const int warp_id = threadIdx.y;
const int lane_id_32 = threadIdx.x;
const int nwarps = blockDim.y;
const int64_t warp_start_offset = (blockIdx.y * nwarps + warp_id) * vals_per_warp;
if (warp_start_offset >= ne0) {
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
const int64_t i01 = ids ? ids[i1] : i1;
const int64_t i02 = i2;
const int64_t i03 = i3;
block_fp4_mmq * y = (block_fp4_mmq *) vy;
const int64_t block_fp4_mmq_size = 8 * QK_MXFP4; // 256 values
const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
const int64_t quad_idx_in_block = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
const int group_id = lane_id_32 / 4;
const int lane_in_group = lane_id_32 % 4;
const int base = group_id * 2;
char2 * yqs2 = (char2 *) y[ib].qs;
const int64_t base_pos = i03 * s03 + i02 * s02 + i01 * s01;
uint8_t scales[2];
#pragma unroll
for (int b = 0; b < 2; ++b) {
const int64_t i0 = warp_start_offset + b * vals_per_scale + lane_id_32;
const float xi = (i0 < ne00) ? x[base_pos + i0] : 0.0f;
float amax = fabsf(xi);
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, WARP_SIZE));
}
const uint8_t e = compute_e8m0_scale(amax);
scales[b] = e;
const float inv_s = (amax == 0.0f) ? 0.0f : __frcp_rn(ggml_cuda_e8m0_to_fp32(e));
#if CUDART_VERSION >= 12080
const float scaled_val = xi * inv_s;
const float val0 = __shfl_sync(0xFFFFFFFF, scaled_val, base, WARP_SIZE);
const float val1 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 16, WARP_SIZE);
const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
if (lane_in_group == 0) {
__nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = *(char2 *) &fp4_packed;
}
#else
// Fallback: manual FP4 conversion using LUT
const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
const uint8_t q_lo_0 = __shfl_sync(0xFFFFFFFF, q_val, base, WARP_SIZE);
const uint8_t q_lo_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 1, WARP_SIZE);
const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
if (lane_in_group == 0) {
char2 q;
q.x = (q_hi_0 << 4) | q_lo_0;
q.y = (q_hi_1 << 4) | q_lo_1;
yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = q;
}
#endif // CUDART_VERSION >= 12080
}
if (lane_id_32 == 0) {
// Store 2 scales packed into 1 uint32
y[ib].d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
}
}
template <mmq_q8_1_ds_layout ds_layout>
static __global__ void quantize_mmq_q8_1(
const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int ne1, const int ne2) {
constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4;
if (i0 >= ne0) {
return;
}
const int64_t i1 = blockIdx.x;
const int64_t i2 = blockIdx.z % ne2;
const int64_t i3 = blockIdx.z / ne2;
const int64_t i00 = i0;
const int64_t i01 = ids ? ids[i1] : i1;
const int64_t i02 = i2;
const int64_t i03 = i3;
const float4 * x4 = (const float4 *) x;
block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel
const int64_t iqs = i0 % (4*QK8_1); // quant index in block
// Load 4 floats per thread and calculate max. abs. value between them:
const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float amax = fabsf(xi.x);
amax = fmaxf(amax, fabsf(xi.y));
amax = fmaxf(amax, fabsf(xi.z));
amax = fmaxf(amax, fabsf(xi.w));
// Exchange max. abs. value between vals_per_scale/4 threads.
#pragma unroll
for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) {
amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
}
float sum;
if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
sum = xi.x + xi.y + xi.z + xi.w;
// Calculate sums across vals_per_sum/4 threads.
#pragma unroll
for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
}
}
const float d_inv = 127.0f / amax;
char4 q;
q.x = roundf(xi.x*d_inv);
q.y = roundf(xi.y*d_inv);
q.z = roundf(xi.z*d_inv);
q.w = roundf(xi.w*d_inv);
// Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
char4 * yqs4 = (char4 *) y[ib].qs;
yqs4[iqs/4] = q;
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
if (iqs % 16 != 0 || iqs >= 96) {
return;
}
y[ib].d2s6[2 + iqs/16] = sum;
if (iqs % 64 != 0) {
return;
}
const float d = 1.0f / d_inv;
y[ib].d2s6[iqs/64] = d;
return;
}
if (iqs % 32 != 0) {
return;
}
const float d = 1.0f / d_inv;
if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
y[ib].ds4[iqs/32] = make_half2(d, sum);
} else {
y[ib].d4[iqs/32] = d;
}
}
void quantize_row_q8_1_cuda(
const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(!ids);
GGML_ASSERT(ne0 % QK8_1 == 0);
const uint3 ne2_fastdiv = init_fastdiv_values(ne2);
const int64_t block_num_x = (ne0 + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
const dim3 num_blocks(block_num_x, ne1, ne2*ne3);
const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);
quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
GGML_UNUSED(type_src0);
}
void quantize_mmq_q8_1_cuda(
const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(ne00 % 4 == 0);
GGML_ASSERT(ne0 % (4*QK8_1) == 0);
// ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid:
const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
const dim3 num_blocks(ne1, block_num_y, ne2*ne3);
const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
switch (mmq_get_q8_1_ds_layout(type_src0)) {
case MMQ_Q8_1_DS_LAYOUT_D4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
break;
case MMQ_Q8_1_DS_LAYOUT_DS4:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
break;
case MMQ_Q8_1_DS_LAYOUT_D2S6:
quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>
<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
break;
default:
GGML_ABORT("fatal error");
break;
}
}
void quantize_mmq_fp4_cuda(
const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
GGML_ASSERT(ne0 > 0);
if (type_src0 == GGML_TYPE_NVFP4) {
GGML_ASSERT(ne00 % QK_NVFP4 == 0);
constexpr int nvfp4_block_size = 128;
const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
const dim3 block_size(nvfp4_block_size, 1, 1);
const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
quantize_mmq_nvfp4<<<num_blocks, block_size, 0, stream>>>(
x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
} else {
GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
constexpr int nwarps = 8;
constexpr int vals_per_warp = 2 * QK_MXFP4;
constexpr int vals_per_block = nwarps * vals_per_warp;
const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
const dim3 block_size(WARP_SIZE, nwarps, 1);
quantize_mmq_mxfp4<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
}
}
cuda/mmq/quantize.cuh
+41
View File
@@ -0,0 +1,41 @@
#pragma once
#include "common.cuh"
#include "mmq.cuh"
#include <cstdint>
#define CUDA_QUANTIZE_BLOCK_SIZE 256
#define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128
static_assert(MATRIX_ROW_PADDING % CUDA_QUANTIZE_BLOCK_SIZE == 0, "Risk of out-of-bounds access.");
static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access.");
typedef void (*quantize_cuda_t)(
const float * x, const int32_t * ids, void * vy,
ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
void quantize_row_q8_1_cuda(
const float * x, const int32_t * ids, void * vy,
ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
void quantize_mmq_q8_1_cuda(
const float * x, const int32_t * ids, void * vy,
ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
void quantize_mmq_fp4_cuda(const float * x,
const int32_t * ids,
void * vy,
ggml_type type_src0,
int64_t ne00,
int64_t s01,
int64_t s02,
int64_t s03,
int64_t ne0,
int64_t ne1,
int64_t ne2,
int64_t ne3,
cudaStream_t stream);
cuda/mmq/test/d2r_stubs.cu
+36
View File
@@ -0,0 +1,36 @@
#include "ds4_mmq_d2r.cuh"
extern "C" int ds4_cuda_q8_fold_take_q81(
const void *src, uint64_t in_dim, const void **q81) {
(void)src;
(void)in_dim;
(void)q81;
return 0;
}
// The D2R kernels use NVIDIA cp.async and MMA assembly. Disable them in ROCm
// builds and use the portable raw/SoA paths instead.
bool ds4_mmq_q2_K_moe_d2r_available(int) { return false; }
bool ds4_mmq_iq2_xxs_moe_d2r_available(int) { return false; }
bool ds4_mmq_q8_0_dense_d2r_available(int) { return false; }
size_t ds4_mmq_q2_K_moe_d2r_scratch_bytes(int64_t, int) { return 0; }
size_t ds4_mmq_iq2_xxs_moe_d2r_pair_scratch_bytes(int64_t, int) { return 0; }
size_t ds4_mmq_iq2_xxs_moe_d2r_fused_scratch_bytes(int64_t, int) { return 0; }
int ds4_mmq_q2_K_moe_d2r_launch(
const void *, int64_t, const void *, const int32_t *, const int32_t *,
float *, int, int, int64_t, int, void *, size_t, cudaStream_t) { return 1; }
int ds4_mmq_q8_0_dense_d2r_launch(
const void *, const void *, float *, int, int, int, cudaStream_t) { return 1; }
int ds4_mmq_iq2_xxs_moe_d2r_pair_launch(
const void *, const void *, int64_t, const void *, const int32_t *,
const int32_t *, float *, float *, int, int, int64_t, int, void *, size_t,
cudaStream_t) { return 1; }
int ds4_mmq_iq2_xxs_moe_d2r_fused_launch(
const void *, const void *, int64_t, const void *, const int32_t *, int,
const int32_t *, const int32_t *, const float *, void *, int, int, int64_t,
int, float, void *, size_t, cudaStream_t) { return 1; }
cuda/mmq/test/iq2_host_tables.h
+105
View File
@@ -0,0 +1,105 @@
// SPDX-License-Identifier: MIT
// iq2_host_tables.h - host-accessible copies of the IQ2 lookup tables.
// Regenerate from cuda/mmq/ggml-common.h whenever upstream changes them:
//
// awk '/GGML_TABLE_BEGIN.*iq2xxs_grid, 256/,/GGML_TABLE_END/' \
// cuda/mmq/ggml-common.h | grep -v GGML_TABLE > iq2xxs_grid.lines
//
// In the vendored ggml-common.h the IQ2 tables are gated on
// GGML_COMMON_IMPL_CUDA and emitted as static const __device__, so they
// are not host-accessible. Reproducing them here as plain const arrays
// gives the CPU parity reference something to look up against.
#pragma once
#include <cstdint>
namespace iq2_host {
// Bit positions for sign extraction (1, 2, 4, ..., 128).
static const uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128};
// 128-entry sign lookup. Index is a 7-bit pattern of which positions
// are negative; value is the 8-bit pattern with the 8th bit set to
// maintain even parity (matches ggml-common.h:503).
static const uint8_t ksigns_iq2xs[128] = {
0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15,
144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159,
160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175,
48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63,
192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207,
80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95,
96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
};
// 256-entry IQ2_XXS grid: each uint64 holds 8 uint8 quantized values.
// Indexed by the 8-bit aux8[l] from the encoded super-block.
// (matches ggml-common.h:550).
static const uint64_t iq2xxs_grid[256] = {
0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
};
} // namespace iq2_host
cuda/mmq/test/proto_gemm_dense_q8_d2r.cu
+773
View File
@@ -0,0 +1,773 @@
// SPDX-License-Identifier: MIT
// proto_gemm_dense_q8_d2r.cu - dense Q8_0 decode-to-registers prototype.
//
// Targets the dense-q8 census shapes (cmdq8shape15 2026-07-09, N = chunk
// tokens 4096) where production mul_mat_q<8,128,0> runs far under the
// cuBLASLt-int8 ceiling at the same shapes:
// sexp gate/up M= 2048 K=4096 mmq 65 TF vs Lt-int8 156
// sexp down M= 4096 K=2048 mmq 24 TF vs Lt-int8 111
// attn q_b M=32768 K=1024 mmq 56 TF vs Lt-int8 79-85
// o_proj M= 4096 K=8192 mmq 96 TF vs Lt-int8 145
//
// Weight layout: byte-neutral SoA repack of raw Q8_0 (34 B/block of 32):
// [half d[M*nb]] [int8 qs[M*K]] (d row-major [M][nb], qs row-major [M][K])
// Total bytes == raw bytes. NOTE a "nicer" fully-tiled contiguous-stage
// layout ([row_tile][k64][h][row][16B]) measured 27% SLOWER at identical
// instruction counts: every col-tile CTA of a row group reads the same 8 KiB
// block concurrently and the contiguous span lives on a few L2 slices (slice
// camping). Row-major keeps each stage's 128 rows K bytes apart -> spread
// across slices; its 2-way ldgsts store conflicts are the cheaper poison.
// This is the REPLACING-artifact candidate (IQ2/Q2K precedent); raw consumers
// fall back to the host mmap.
//
// Kernel: m128n64 CTA (8 warps), q8 activations = production block_q8_1_mmq
// D4 layout (f32 scale per 32) double-buffered at k128; weights double-
// buffered at k64 (128 rows x 64 B + 16 B row pad for ldmatrix bank rotation);
// mma.m16n8k32.s8 with per-k32 float fold acc += (s32)C * d_w(row) * d_a(col).
// Q8_0 has no mins: no bias terms, no fix staging (d4 read straight from the
// q8 block header). Grid is 1-D with triton-style grouped supertiling
// (group_m row-tiles hot in L2 while col tiles stream) - group_m is a runtime
// arg, swept by the harness.
//
// Build on the GB10/Blackwell box after `make -j6 cuda-spark`:
// ar rcs libds4mmq.a cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o \
// cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o
// /usr/local/cuda/bin/nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
// -Icuda/mmq cuda/mmq/test/proto_gemm_dense_q8_d2r.cu libds4mmq.a \
// -lcudart -lcublas -lcuda -o proto_gemm_dense_q8_d2r
// argv: [reps=20] [warmup=5] [N=4096]
#include <type_traits>
#include "ds4_mmq.h"
#include "quantize.cuh"
extern "C" int ds4_cuda_q8_fold_take_q81(const void *src, uint64_t in_dim, void *out) {
(void)src;
(void)in_dim;
(void)out;
return 0;
}
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <random>
#include <vector>
#define CHECK(call) do { \
cudaError_t err__ = (call); \
if (err__ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d: %s\n", #call, __FILE__, __LINE__, \
cudaGetErrorString(err__)); \
exit(1); \
} \
} while (0)
namespace {
constexpr int kMTile = 128;
#ifndef PROTO_Q8_NTILE
#define PROTO_Q8_NTILE 128
#endif
constexpr int kNTile = PROTO_Q8_NTILE;
// 16-warp CTA, one per SM: 8 row-warps x 2 col-warps. Same 16 warps/SM
// residency as 2x8, but per-thread acc halves to 8 fragments (32 regs) --
// the 2x8-warp NT128 shape spilled ~76 B and the spill reloads were the top
// stall (long_sb 3.02, local_ld 2.65M).
constexpr int kRowWarps = 8;
constexpr int kColWarps = 2;
constexpr int kWarps = kRowWarps * kColWarps;
constexpr int kThreads = 32 * kWarps;
// 3-stage rings: issues happen at the TOP of each iter for stage i+2 (a
// buffer no in-flight reader touches), so the post-compute __syncthreads of
// the 2-stage schedule disappears -- one barrier per k64. 69 KiB total needs
// the dynamic-smem opt-in (1 CTA/SM is already forced by the 512-thread
// register budget, so no occupancy cost).
constexpr int kStages = 3; // q8 qs ring (k64 halves) + header ring (k128)
constexpr int kWStages = 3; // weight ring, k64 granularity
constexpr int kNFrag = kNTile / 8;
constexpr int kWRowPad = 80; // 64 B payload + 16 B pad: 20-int ldmatrix stride rotates banks
// q8 staging is SPLIT so NT=128 fits under 48 KiB static smem: the 16 B
// header (d4 scales) double-buffered at k128 cadence, the 128 B qs payload
// staged as k64 halves (64 B rows, same padded geometry as the weight ring).
constexpr int kQ8QsChunks = kNTile * 4; // cols x 4 16B-chunks per k64 half
constexpr int kQ8QsTrips = (kQ8QsChunks + kThreads - 1) / kThreads;
constexpr int kWQChunks = kMTile * 4; // 128 rows x 4 16B-chunks (64 B payload)
constexpr int kWQTrips = kWQChunks / kThreads;
static_assert(kQ8QsTrips == 1, "unexpected q8 qs issue trip count");
static_assert(kWQTrips == 1, "unexpected weight issue trip count");
static_assert(kNTile == 128, "16-warp decomposition assumes NT=128");
static_assert(kNTile <= kThreads, "q8 header prefetch assumes one trip");
constexpr int kNFragPerWarp = kNFrag / kColWarps;
// Loop invariants live in plain (warp-uniform) registers: everything here
// derives arithmetically from blockIdx + kernel args, so ptxas can promote
// the math to the uniform register file. (The Q2K proto's shared invariant
// table exists for its memory-sourced expert-worklist lookups - dense has
// none, and routing these through volatile smem forced them into regular
// registers, feeding the 128-reg wall.)
struct DenseQ8Params {
const char *wd_row0; // d plane at cta_row0 (row stride nb*2 = K/16)
const char *wq_row0; // qs plane at cta_row0 (row stride K)
const char *q8_tile; // q8 blocks at col_lo
uint32_t q8_k128_stride_bytes;
uint32_t wq_row_stride; // K bytes
int k64_iters;
int k128_iters;
};
uint16_t sane_half(std::mt19937 &rng) {
const uint16_t mant = (uint16_t)(rng() & 0x03FFu);
const uint16_t sign = (uint16_t)((rng() & 1u) << 15);
return (uint16_t)(sign | (14u << 10) | mant);
}
float *device_random_floats(std::mt19937 &rng, size_t n) {
std::vector<float> h(n);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < n; ++i) {
h[i] = dist(rng);
}
float *d = nullptr;
CHECK(cudaMalloc(&d, n * sizeof(float)));
CHECK(cudaMemcpy(d, h.data(), n * sizeof(float), cudaMemcpyHostToDevice));
return d;
}
template <typename F>
float best_time_ms(F &&fn, int warmup, int reps) {
cudaEvent_t a, b;
CHECK(cudaEventCreate(&a));
CHECK(cudaEventCreate(&b));
for (int i = 0; i < warmup; ++i) {
fn();
}
CHECK(cudaDeviceSynchronize());
float best = std::numeric_limits<float>::infinity();
for (int r = 0; r < reps; ++r) {
CHECK(cudaEventRecord(a));
fn();
CHECK(cudaEventRecord(b));
CHECK(cudaEventSynchronize(b));
float ms = 0.0f;
CHECK(cudaEventElapsedTime(&ms, a, b));
best = std::min(best, ms);
}
CHECK(cudaEventDestroy(a));
CHECK(cudaEventDestroy(b));
return best;
}
__device__ __forceinline__ void zero_16B(void *dst) {
*reinterpret_cast<int4 *>(dst) = make_int4(0, 0, 0, 0);
}
__device__ __forceinline__ void cp_async_16B(void *dst, const void *src, bool pred) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
if (pred) {
const unsigned smem = static_cast<unsigned>(__cvta_generic_to_shared(dst));
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
:: "r"(smem), "l"(src));
} else {
zero_16B(dst);
}
#else
if (pred) {
*reinterpret_cast<int4 *>(dst) = *reinterpret_cast<const int4 *>(src);
} else {
zero_16B(dst);
}
#endif
}
__device__ __forceinline__ void cp_async_4B(void *dst, const void *src, bool pred) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
if (pred) {
const unsigned smem = static_cast<unsigned>(__cvta_generic_to_shared(dst));
asm volatile("cp.async.ca.shared.global [%0], [%1], 4;"
:: "r"(smem), "l"(src));
} else {
*reinterpret_cast<uint32_t *>(dst) = 0u;
}
#else
if (pred) {
*reinterpret_cast<uint32_t *>(dst) = *reinterpret_cast<const uint32_t *>(src);
} else {
*reinterpret_cast<uint32_t *>(dst) = 0u;
}
#endif
}
__device__ __forceinline__ void cp_async_commit() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("cp.async.commit_group;");
#endif
}
template <int KeepGroups>
__device__ __forceinline__ void cp_async_wait_group() {
static_assert(KeepGroups >= 0 && KeepGroups <= 7, "bad cp.async wait_group depth");
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
#endif
}
// CTA is launched as dim3(32, kWarps): the ggml mma primitives
// (load_ldmatrix, tile get_i/get_j) use threadIdx.x as the LANE.
__device__ __forceinline__ int d2r_lane() {
#if defined(__CUDA_ARCH__)
return (int)threadIdx.x;
#else
return 0;
#endif
}
__device__ __forceinline__ int d2r_warp() {
#if defined(__CUDA_ARCH__)
return (int)threadIdx.y;
#else
return 0;
#endif
}
__device__ __forceinline__ int d2r_tid() {
#if defined(__CUDA_ARCH__)
return (int)(threadIdx.y * 32 + threadIdx.x);
#else
return 0;
#endif
}
constexpr size_t kSmemQ8HBytes = (size_t)kStages * kNTile * 16; // headers
constexpr size_t kSmemQ8QBytes = (size_t)kStages * kNTile * kWRowPad; // qs k64 halves
constexpr size_t kSmemWQBytes = (size_t)kWStages * kMTile * kWRowPad;
constexpr size_t kSmemWDBytes = (size_t)kWStages * kMTile * sizeof(uint32_t);
constexpr size_t kSmemWQOff = 0;
constexpr size_t kSmemQ8QOff = kSmemWQOff + kSmemWQBytes;
constexpr size_t kSmemQ8HOff = kSmemQ8QOff + kSmemQ8QBytes;
constexpr size_t kSmemWDOff = kSmemQ8HOff + kSmemQ8HBytes;
constexpr size_t kSmemTotalBytes = kSmemWDOff + kSmemWDBytes;
static_assert(kSmemTotalBytes <= 99ull * 1024ull, "dynamic shared memory exceeds sm_121 limit");
// ---- q8 activation prefetch, split staging ----
// qs k64 halves: 128 cols x 64 B (bytes 16 + (k64&1)*64 .. of each col's
// 144 B block), padded 80 B rows — same geometry as the weight ring.
// headers: 128 cols x 16 B (the d4 scale header), k128 cadence.
__device__ __forceinline__ void issue_q8_qs_prefetch(
int8_t (&s_q8q)[kStages][kNTile][kWRowPad],
const DenseQ8Params &p,
int k64_iter) {
const bool pred = k64_iter < p.k64_iters;
const char *base = p.q8_tile +
(uint64_t)(k64_iter >> 1) * p.q8_k128_stride_bytes +
16u + (uint64_t)(k64_iter & 1) * 64u;
const int buf = k64_iter % kStages;
#pragma unroll
for (int trip = 0; trip < kQ8QsTrips; ++trip) {
const int chunk = d2r_tid() + trip * kThreads;
if (kQ8QsChunks % kThreads != 0 && chunk >= kQ8QsChunks) {
break;
}
const int col = chunk >> 2;
const int h = chunk & 3;
void *dst = &s_q8q[buf][col][h * 16];
const void *src = base + (uint64_t)col * sizeof(block_q8_1_mmq) + (uint64_t)h * 16u;
cp_async_16B(dst, src, pred);
}
}
__device__ __forceinline__ void issue_q8_hdr_prefetch(
float (&s_q8h)[kStages][kNTile][4],
const DenseQ8Params &p,
int k128_iter) {
if (d2r_tid() >= kNTile) {
return;
}
const bool pred = k128_iter < p.k128_iters;
const char *base = p.q8_tile +
(uint64_t)k128_iter * p.q8_k128_stride_bytes;
const int col = d2r_tid();
void *dst = &s_q8h[k128_iter % kStages][col][0];
const void *src = base + (uint64_t)col * sizeof(block_q8_1_mmq);
cp_async_16B(dst, src, pred);
}
// ---- weight prefetch (one k64 stage = 128 rows x 64 B qs + 128 x 4 B d) ----
__device__ __forceinline__ void issue_w_prefetch(
int8_t (&s_wq)[kWStages][kMTile][kWRowPad],
uint32_t (&s_wd)[kWStages][kMTile],
const DenseQ8Params &p,
int k64_iter) {
const bool pred = k64_iter < p.k64_iters;
const int wst = k64_iter % kWStages;
// Row-major planes (see layout note at top: contiguous-stage tiling was
// 27% slower from L2 slice camping).
const char *wq = p.wq_row0 + (uint64_t)k64_iter * 64u;
const char *wd = p.wd_row0 + (uint64_t)k64_iter * 4u;
const uint32_t wq_row_stride = p.wq_row_stride;
const uint32_t wd_row_stride = p.wq_row_stride >> 4; // nb*2 = K/16
#pragma unroll
for (int trip = 0; trip < kWQTrips; ++trip) {
const int chunk = d2r_tid() + trip * kThreads;
const int row = chunk >> 2;
const int h = chunk & 3;
void *dst = &s_wq[wst][row][h * 16];
const void *src = wq + (uint64_t)row * wq_row_stride + (uint64_t)h * 16u;
cp_async_16B(dst, src, pred);
}
if (d2r_tid() < kMTile) {
const int row = d2r_tid();
void *dst = &s_wd[wst][row];
const void *src = wd + (uint64_t)row * wd_row_stride;
cp_async_4B(dst, src, pred);
}
}
// ---- mainloop ----
template <typename TileA, typename TileB, typename TileC>
__device__ __forceinline__ void dense_q8_d2r_mainloop(
float (&acc)[kNFragPerWarp][4],
float (&s_q8h)[kStages][kNTile][4],
int8_t (&s_q8q)[kStages][kNTile][kWRowPad],
int8_t (&s_wq)[kWStages][kMTile][kWRowPad],
uint32_t (&s_wd)[kWStages][kMTile],
const DenseQ8Params &p) {
static_assert(TileC::ne == 4, "expected m16n8 s32 accumulator fragment");
const int k64_iters = p.k64_iters;
const int lane = d2r_lane();
const int group = lane >> 2;
const int wrow = (d2r_warp() >> 1) * 16; // row-warp
const int nf0 = (d2r_warp() & 1) * kNFragPerWarp; // col-warp half
const int c0 = TileC::get_j(0);
const int c1 = TileC::get_j(1);
// Prologue: stages 0 and 1 in flight as two groups. Loop iters run
// wait -> barrier -> issue(i+2) -> compute: with 3 buffers the issue
// never aliases the two live stages, so the 2-stage schedule's
// post-compute barrier disappears (one barrier per k64).
issue_w_prefetch(s_wq, s_wd, p, 0);
issue_q8_qs_prefetch(s_q8q, p, 0);
issue_q8_hdr_prefetch(s_q8h, p, 0);
cp_async_commit();
issue_w_prefetch(s_wq, s_wd, p, 1);
issue_q8_qs_prefetch(s_q8q, p, 1);
issue_q8_hdr_prefetch(s_q8h, p, 1);
cp_async_commit();
for (int i = 0; i < k64_iters; ++i) {
// wait -> barrier -> issue: the barrier both publishes stage i's
// copies CTA-wide and proves every warp finished reading stage i-1,
// whose buffer ((i-1)%3 == (i+2)%3) the issues below overwrite.
cp_async_wait_group<1>();
__syncthreads();
issue_w_prefetch(s_wq, s_wd, p, i + 2);
issue_q8_qs_prefetch(s_q8q, p, i + 2);
if ((i & 1) == 0) {
issue_q8_hdr_prefetch(s_q8h, p, (i >> 1) + 2);
}
cp_async_commit();
const int wst = i % kWStages;
const int qbuf = i % kStages; // qs halves ride the k64 cadence
const int hbuf = (i >> 1) % kStages; // headers ride the k128 cadence
const uint32_t dw0_bits = s_wd[wst][wrow + group];
const uint32_t dw1_bits = s_wd[wst][wrow + group + 8];
const float2 dw0 = __half22float2(*reinterpret_cast<const half2 *>(&dw0_bits));
const float2 dw1 = __half22float2(*reinterpret_cast<const half2 *>(&dw1_bits));
// t-phased: each A load overlaps the previous phase's B/MMA/fold tail
// (fusing both A loads up front measured -27%: hard A dependency at the
// k64 head + burstier MIO).
#pragma unroll
for (int t = 0; t < 2; ++t) {
TileA A;
ggml_cuda_mma::load_ldmatrix(
A, reinterpret_cast<const int *>(&s_wq[wst][wrow][t * 32]),
kWRowPad / (int)sizeof(int));
const int tq8 = (i & 1) * 2 + t; // k32 index within the q8 k128 block
const float dwr0 = t ? dw0.y : dw0.x;
const float dwr1 = t ? dw1.y : dw1.x;
#pragma unroll
for (int nf = 0; nf < kNFragPerWarp; ++nf) {
TileB B;
TileC C;
ggml_cuda_mma::load_ldmatrix(
B, reinterpret_cast<const int *>(&s_q8q[qbuf][(nf0 + nf) * 8][t * 32]),
kWRowPad / (int)sizeof(int));
ggml_cuda_mma::mma(C, A, B);
const float da0 = s_q8h[hbuf][(nf0 + nf) * 8 + c0][tq8];
const float da1 = s_q8h[hbuf][(nf0 + nf) * 8 + c1][tq8];
acc[nf][0] = fmaf((float)C.x[0], dwr0 * da0, acc[nf][0]);
acc[nf][1] = fmaf((float)C.x[1], dwr0 * da1, acc[nf][1]);
acc[nf][2] = fmaf((float)C.x[2], dwr1 * da0, acc[nf][2]);
acc[nf][3] = fmaf((float)C.x[3], dwr1 * da1, acc[nf][3]);
}
}
}
}
__global__ __launch_bounds__(kThreads, 1)
void dense_q8_d2r_kernel(const char * __restrict__ wd_plane,
const char * __restrict__ wq_plane,
const block_q8_1_mmq * __restrict__ q8,
float * __restrict__ outf,
__half * __restrict__ outh,
int M, int N, int K, int group_m) {
#if defined(TURING_MMA_AVAILABLE)
using tile_A = ggml_cuda_mma::tile<16, 8, int>;
using tile_B = ggml_cuda_mma::tile<8, 8, int>;
using tile_C = ggml_cuda_mma::tile<16, 8, int>;
extern __shared__ __align__(16) char smem_dyn[];
auto &s_wq = *reinterpret_cast<int8_t (*)[kWStages][kMTile][kWRowPad]>(smem_dyn + kSmemWQOff);
auto &s_q8q = *reinterpret_cast<int8_t (*)[kStages][kNTile][kWRowPad]>(smem_dyn + kSmemQ8QOff);
auto &s_q8h = *reinterpret_cast<float (*)[kStages][kNTile][4]>(smem_dyn + kSmemQ8HOff);
auto &s_wd = *reinterpret_cast<uint32_t (*)[kWStages][kMTile]>(smem_dyn + kSmemWDOff);
// Triton-style grouped supertile order: group_m row-tiles stay L2-hot
// while the col tiles stream past them.
const int num_m = M / kMTile;
const int num_n = (N + kNTile - 1) / kNTile;
const int width = group_m * num_n;
const int g = (int)blockIdx.x / width;
const int rem = (int)blockIdx.x - g * width;
int gsize = num_m - g * group_m;
if (gsize > group_m) {
gsize = group_m;
}
const int pid_m = g * group_m + rem % gsize;
const int pid_n = rem / gsize;
const int cta_row0 = pid_m * kMTile;
const int col_lo = pid_n * kNTile;
DenseQ8Params p;
p.wd_row0 = wd_plane + (uint64_t)cta_row0 * (uint64_t)(K >> 4);
p.wq_row0 = wq_plane + (uint64_t)cta_row0 * (uint64_t)K;
p.q8_tile = reinterpret_cast<const char *>(q8) +
(uint64_t)col_lo * sizeof(block_q8_1_mmq);
p.q8_k128_stride_bytes = (uint32_t)((uint64_t)N * sizeof(block_q8_1_mmq));
p.wq_row_stride = (uint32_t)K;
p.k64_iters = K >> 6;
p.k128_iters = K >> 7;
float acc[kNFragPerWarp][tile_C::ne] = {};
dense_q8_d2r_mainloop<tile_A, tile_B, tile_C>(acc, s_q8h, s_q8q, s_wq, s_wd, p);
const int out_col_lo = col_lo + (d2r_warp() & 1) * (kNFragPerWarp * 8);
const int out_N = N;
const int out_M = M;
const int out_row0 = cta_row0 + ((d2r_warp() >> 1) << 4);
float *outf_base = outf;
__half *outh_base = outh;
#pragma unroll
for (int nf = 0; nf < kNFragPerWarp; ++nf) {
const int col_frag0 = out_col_lo + nf * 8;
#pragma unroll
for (int l = 0; l < tile_C::ne; ++l) {
const int row = out_row0 + tile_C::get_i(l);
const int col = col_frag0 + tile_C::get_j(l);
if (col < out_N) {
const float v = isfinite(acc[nf][l]) ? acc[nf][l] : 0.0f;
if (outh_base) {
outh_base[(uint64_t)col * (uint64_t)out_M + (uint64_t)row] = __float2half(v);
} else {
outf_base[(uint64_t)col * (uint64_t)out_M + (uint64_t)row] = v;
}
}
}
}
#else
(void)wd_plane;
(void)wq_plane;
(void)q8;
(void)outf;
(void)outh;
(void)M;
(void)N;
(void)K;
(void)group_m;
#endif
}
int launch_dense_q8_d2r(const void *wd_plane, const void *wq_plane,
const block_q8_1_mmq *q8, float *outf, __half *outh,
int M, int N, int K, int group_m, cudaStream_t stream) {
if (!wd_plane || !wq_plane || !q8 || (!outf && !outh) ||
M <= 0 || (M & (kMTile - 1)) != 0 || K % 256 != 0 || N <= 0 || group_m <= 0) {
fprintf(stderr, "launch_dense_q8_d2r: bad args M=%d N=%d K=%d group_m=%d\n",
M, N, K, group_m);
return -1;
}
const int num_m = M / kMTile;
const int num_n = (N + kNTile - 1) / kNTile;
static int smem_opted = 0;
if (!smem_opted) {
cudaFuncSetAttribute(dense_q8_d2r_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)kSmemTotalBytes);
smem_opted = 1;
}
const dim3 grid((unsigned)(num_m * num_n), 1, 1);
const dim3 block(32, kWarps, 1);
dense_q8_d2r_kernel<<<grid, block, kSmemTotalBytes, stream>>>(
(const char *)wd_plane, (const char *)wq_plane, q8, outf, outh, M, N, K, group_m);
const cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
fprintf(stderr, "dense_q8_d2r_kernel launch failed: %s\n", cudaGetErrorString(err));
return -2;
}
return 0;
}
void quantize_q8_mmq_dense(const float *x, block_q8_1_mmq *q8, int K, int N,
cudaStream_t stream) {
quantize_mmq_q8_1_cuda(
x, /*ids=*/nullptr, (void *)q8,
GGML_TYPE_Q8_0, /*ne00=*/K, /*s01=*/(int64_t)K, /*s02=*/0, /*s03=*/0,
/*ne0=*/K, /*ne1=*/N, /*ne2=*/1, /*ne3=*/1, stream);
CHECK(cudaGetLastError());
}
struct ParityStats {
double max_abs = 0.0;
double max_rel_rms = 0.0;
int max_row = -1;
int max_col = -1;
};
ParityStats check_parity(const float *ref, const float *got, int M, int N) {
std::vector<double> ss((size_t)M, 0.0);
for (int c = 0; c < N; ++c) {
const float *rp = ref + (size_t)c * M;
for (int r = 0; r < M; ++r) {
ss[(size_t)r] += (double)rp[r] * (double)rp[r];
}
}
ParityStats st;
for (int c = 0; c < N; ++c) {
const float *rp = ref + (size_t)c * M;
const float *gp = got + (size_t)c * M;
for (int r = 0; r < M; ++r) {
const double abs_err = std::abs((double)gp[r] - (double)rp[r]);
const double rms = std::sqrt(ss[(size_t)r] / std::max(1, N));
const double rel = abs_err / std::max(1.0e-6, rms);
if (abs_err > st.max_abs) {
st.max_abs = abs_err;
st.max_row = r;
st.max_col = c;
}
st.max_rel_rms = std::max(st.max_rel_rms, rel);
}
}
return st;
}
void run_shape(std::mt19937 &rng, const char *label, int M, int K, int N,
int warmup, int reps, bool do_parity) {
const int nb = K / 32;
const uint64_t raw_bytes = (uint64_t)M * nb * 34u;
const uint64_t d_bytes = (uint64_t)M * nb * 2u;
// Raw Q8_0 blocks (34 B: half d + 32 int8) for the production reference,
// and the byte-neutral SoA planes for the D2R kernel, from the same data.
std::vector<uint8_t> h_raw(raw_bytes);
std::vector<uint8_t> h_soa(raw_bytes);
{
// Row-major SoA planes: d [M][nb] halves, qs [M][K] int8.
uint8_t *raw = h_raw.data();
uint16_t *dp = reinterpret_cast<uint16_t *>(h_soa.data());
int8_t *qp = reinterpret_cast<int8_t *>(h_soa.data() + d_bytes);
for (int r = 0; r < M; ++r) {
for (int b = 0; b < nb; ++b) {
uint8_t *blk = raw + ((uint64_t)r * nb + b) * 34u;
const uint16_t d = sane_half(rng);
memcpy(blk, &d, 2);
for (int j = 0; j < 32; ++j) {
blk[2 + j] = (uint8_t)(rng() & 0xFFu);
}
dp[(uint64_t)r * nb + b] = d;
memcpy(qp + (uint64_t)r * K + (uint64_t)b * 32u, blk + 2, 32);
}
}
}
void *d_raw = nullptr;
void *d_soa = nullptr;
CHECK(cudaMalloc(&d_raw, raw_bytes));
CHECK(cudaMalloc(&d_soa, raw_bytes));
CHECK(cudaMemcpy(d_raw, h_raw.data(), raw_bytes, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(d_soa, h_soa.data(), raw_bytes, cudaMemcpyHostToDevice));
const char *d_wd = (const char *)d_soa;
const char *d_wq = (const char *)d_soa + d_bytes;
float *d_x = device_random_floats(rng, (size_t)N * (size_t)K);
// +kNTile block slack so a guarded last col tile can read past col N-1
// without faulting (values discarded by the epilogue guard).
const size_t q8_blocks = (size_t)N * (size_t)(K / 128) + kNTile;
block_q8_1_mmq *d_q8 = nullptr;
CHECK(cudaMalloc(&d_q8, q8_blocks * sizeof(block_q8_1_mmq)));
CHECK(cudaMemset(d_q8, 0, q8_blocks * sizeof(block_q8_1_mmq)));
quantize_q8_mmq_dense(d_x, d_q8, K, N, 0);
float *d_ref = nullptr;
float *d_new = nullptr;
__half *d_newh = nullptr;
CHECK(cudaMalloc(&d_ref, (size_t)M * N * sizeof(float)));
CHECK(cudaMalloc(&d_new, (size_t)M * N * sizeof(float)));
CHECK(cudaMalloc(&d_newh, (size_t)M * N * sizeof(__half)));
if (ds4_mmq_q8_0_dense(d_raw, d_x, d_ref, M, N, K, 0) != 0) {
fprintf(stderr, "reference ds4_mmq_q8_0_dense failed M=%d K=%d\n", M, K);
exit(1);
}
if (launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, 8, 0) != 0) {
exit(1);
}
CHECK(cudaDeviceSynchronize());
if (do_parity) {
std::vector<float> ref((size_t)M * N), got((size_t)M * N);
CHECK(cudaMemcpy(ref.data(), d_ref, ref.size() * sizeof(float), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(got.data(), d_new, got.size() * sizeof(float), cudaMemcpyDeviceToHost));
const ParityStats st = check_parity(ref.data(), got.data(), M, N);
const bool pass = st.max_rel_rms <= 2.0e-3;
printf(" parity %-8s M=%5d K=%4d N=%d: max_abs=%.6g at col=%d row=%d max_abs/rms(row)=%.6g %s\n",
label, M, K, N, st.max_abs, st.max_col, st.max_row, st.max_rel_rms,
pass ? "PASS" : "FAIL");
if (!pass && getenv("PROTO_Q8_DEBUG")) {
// Error structure: per (row-tile, col-tile) fraction of bad elements,
// plus per row%16 / col%8 / warp-row histograms of bad positions.
const int ntm = M / 128, ntn = (N + 63) / 64;
size_t bad_total = 0;
std::vector<size_t> bad_rowmod(16, 0), bad_colmod(8, 0), bad_warp(8, 0);
std::vector<size_t> tile_bad((size_t)ntm * ntn, 0);
for (int c = 0; c < N; ++c) {
const float *rp = ref.data() + (size_t)c * M;
const float *gp = got.data() + (size_t)c * M;
for (int r = 0; r < M; ++r) {
const double e = std::abs((double)gp[r] - (double)rp[r]);
const double mag = std::max(1.0, std::abs((double)rp[r]));
if (e / mag > 1.0e-2) {
++bad_total;
++bad_rowmod[r & 15];
++bad_colmod[c & 7];
++bad_warp[(r >> 4) & 7];
++tile_bad[(size_t)(r / 128) * ntn + (c / 64)];
}
}
}
printf(" bad(rel>1e-2) %zu / %zu\n", bad_total, (size_t)M * N);
printf(" by row%%16:"); for (int i = 0; i < 16; ++i) printf(" %zu", bad_rowmod[i]); printf("\n");
printf(" by col%%8: "); for (int i = 0; i < 8; ++i) printf(" %zu", bad_colmod[i]); printf("\n");
printf(" by warp(row/16%%8):"); for (int i = 0; i < 8; ++i) printf(" %zu", bad_warp[i]); printf("\n");
printf(" tile map (rows=row-tile, cols=col-tile, %% bad, first 8x16):\n");
for (int tm = 0; tm < std::min(ntm, 8); ++tm) {
printf(" ");
for (int tn = 0; tn < std::min(ntn, 16); ++tn) {
printf("%3zu", tile_bad[(size_t)tm * ntn + tn] * 100 / (128 * 64));
}
printf("\n");
}
// A few sample mismatches with ratios (scale bugs show as clean ratios).
int shown = 0;
for (int c = 0; c < N && shown < 6; ++c) {
for (int r = 0; r < M && shown < 6; ++r) {
const float rv = ref[(size_t)c * M + r], gv = got[(size_t)c * M + r];
if (std::abs(gv - rv) / std::max(1.0f, std::abs(rv)) > 1.0e-2) {
printf(" sample r=%d c=%d ref=%.6g got=%.6g got/ref=%.4f\n",
r, c, rv, gv, rv != 0.0f ? gv / rv : 0.0f);
++shown;
}
}
}
}
if (!pass) {
exit(2);
}
}
const double flop = 2.0 * (double)M * (double)K * (double)N;
const float ref_ms = best_time_ms([&] {
ds4_mmq_q8_0_dense(d_raw, d_x, d_ref, M, N, K, 0);
}, warmup, reps);
printf(" %-8s M=%5d K=%4d N=%d mmq-ref %7.3f ms %5.1f TF\n",
label, M, K, N, ref_ms, flop / (ref_ms * 1.0e-3) / 1.0e12);
float best_kernel_ms = std::numeric_limits<float>::infinity();
int best_g = 0;
const int groups[] = {1, 4, 8, 16};
for (int gi = 0; gi < 4; ++gi) {
const int g = groups[gi];
const float ms = best_time_ms([&] {
launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, g, 0);
}, warmup, reps);
printf(" d2r f32-out group_m=%2d %7.3f ms %5.1f TF\n",
g, ms, flop / (ms * 1.0e-3) / 1.0e12);
if (ms < best_kernel_ms) {
best_kernel_ms = ms;
best_g = g;
}
}
const float f16_ms = best_time_ms([&] {
launch_dense_q8_d2r(d_wd, d_wq, d_q8, nullptr, d_newh, M, N, K, best_g, 0);
}, warmup, reps);
const float full_ms = best_time_ms([&] {
quantize_q8_mmq_dense(d_x, d_q8, K, N, 0);
launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, best_g, 0);
}, warmup, reps);
printf(" d2r BEST group_m=%d: f32-out %7.3f ms %5.1f TF | f16-out %7.3f ms %5.1f TF | +q8 %7.3f ms %5.1f TF | vs mmq %.2fx\n",
best_g,
best_kernel_ms, flop / (best_kernel_ms * 1.0e-3) / 1.0e12,
f16_ms, flop / (f16_ms * 1.0e-3) / 1.0e12,
full_ms, flop / (full_ms * 1.0e-3) / 1.0e12,
ref_ms / best_kernel_ms);
CHECK(cudaFree(d_raw));
CHECK(cudaFree(d_soa));
CHECK(cudaFree(d_x));
CHECK(cudaFree(d_q8));
CHECK(cudaFree(d_ref));
CHECK(cudaFree(d_new));
CHECK(cudaFree(d_newh));
}
} // namespace
int main(int argc, char **argv) {
const int reps = (argc > 1) ? atoi(argv[1]) : 20;
const int warmup = (argc > 2) ? atoi(argv[2]) : 5;
const int N = (argc > 3) ? atoi(argv[3]) : 4096;
printf("### proto_gemm_dense_q8_d2r reps=%d warmup=%d N=%d ###\n", reps, warmup, N);
printf("### CTA m%dn%d, warps=%d, q8 x%d stages, weights x%d stages, dyn smem %zu B ###\n",
kMTile, kNTile, kWarps, kStages, kWStages, kSmemTotalBytes);
printf("### Lt-int8 ceilings at these shapes: gate/up 156, down 111, q_b 79-85, o_proj 145 TF ###\n");
std::mt19937 rng(0xD84D8021u + 13u * (unsigned)N);
run_shape(rng, "gate/up", 2048, 4096, N, warmup, reps, true);
run_shape(rng, "down", 4096, 2048, N, warmup, reps, true);
run_shape(rng, "q_b", 32768, 1024, N, warmup, reps, false);
run_shape(rng, "o_proj", 4096, 8192, N, warmup, reps, false);
return 0;
}
cuda/mmq/test/proto_gemm_lib.cu
+127
View File
@@ -0,0 +1,127 @@
// SPDX-License-Identifier: MIT
// proto_gemm_lib.cu - NVIDIA-library dense GEMM ceiling at the exact MoE-leg
// shapes (cuBLASLt heuristics, up to 16 algos, best-of; TN col-major):
// down-shape: m=4096 n=24576 k=2048 (== one W4096 down launch's total FLOP)
// gateup-shape: m=2048 n=24576 k=4096 (== one gate or up GEMM of the pair)
// Kinds: int8->int32 (COMPUTE_32I) | f16->f16 @ f32-acc (COMPUTE_32F) |
// f16->f16 @ f16-acc (COMPUTE_16F).
//
// This is "what NVIDIA's best dense kernel does on this silicon at our shape" —
// the practical ceiling any MoE grouped-GEMM rewrite chases (grouping + quant
// decode overhead comes on top of it).
//
// Build: /usr/local/cuda/bin/nvcc -O3 -std=c++17 -arch=sm_121 \
// cuda/mmq/test/proto_gemm_lib.cu -o proto_gemm_lib -lcublasLt
// argv: [reps=10]
#include <cublasLt.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>
#define CHECK(call) do { cudaError_t e_=(call); if(e_!=cudaSuccess){ \
fprintf(stderr,"CUDA %s @%d: %s\n",#call,__LINE__,cudaGetErrorString(e_)); exit(1);} } while(0)
#define LT(call) do { cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
fprintf(stderr,"cublasLt %s @%d: status %d\n",#call,__LINE__,(int)s_); goto done; } } while(0)
namespace {
__global__ void fillk(unsigned *p, size_t n, unsigned seed) {
for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < n;
i += (size_t)gridDim.x * blockDim.x)
p[i] = ((unsigned)i * 2654435761u + seed) & 0x3BFF3BFFu; // finite small f16 pairs; valid int8 bytes
}
void bench_case(cublasLtHandle_t lt, void *dA, void *dB, void *dC, void *ws,
size_t wsSize, int m, int n, int k, int kind, int reps,
const char *label) {
const cudaDataType ab = (kind == 0) ? CUDA_R_8I : CUDA_R_16F;
const cudaDataType cd = (kind == 0) ? CUDA_R_32I : CUDA_R_16F;
const cublasComputeType_t ct = (kind == 0) ? CUBLAS_COMPUTE_32I
: (kind == 1) ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_16F;
const cudaDataType st = (kind == 0) ? CUDA_R_32I
: (kind == 1) ? CUDA_R_32F : CUDA_R_16F;
cublasLtMatmulDesc_t op = nullptr;
cublasLtMatrixLayout_t La = nullptr, Lb = nullptr, Lc = nullptr;
cublasLtMatmulPreference_t pref = nullptr;
cublasLtMatmulHeuristicResult_t res[16];
int nres = 0;
const cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
const int32_t ai = 1, bi = 0; const float af = 1.f, bf = 0.f;
const __half ah = __float2half(1.f), bh = __float2half(0.f);
const void *alpha = (kind == 0) ? (const void *)&ai : (kind == 1) ? (const void *)&af : (const void *)&ah;
const void *beta = (kind == 0) ? (const void *)&bi : (kind == 1) ? (const void *)&bf : (const void *)&bh;
float best = 1e30f; int besti = -1;
LT(cublasLtMatmulDescCreate(&op, ct, st));
LT(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
LT(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
LT(cublasLtMatrixLayoutCreate(&La, ab, k, m, k)); // stored kxm; opA=T -> mxk
LT(cublasLtMatrixLayoutCreate(&Lb, ab, k, n, k));
LT(cublasLtMatrixLayoutCreate(&Lc, cd, m, n, m));
LT(cublasLtMatmulPreferenceCreate(&pref));
LT(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&wsSize, sizeof(wsSize)));
LT(cublasLtMatmulAlgoGetHeuristic(lt, op, La, Lb, Lc, Lc, pref, 16, res, &nres));
if (!nres) { printf(" %-14s m=%5d n=%5d k=%5d: NOALGO\n", label, m, n, k); goto done; }
for (int i = 0; i < nres; ++i) {
if (res[i].state != CUBLAS_STATUS_SUCCESS) continue;
if (cublasLtMatmul(lt, op, alpha, dA, La, dB, Lb, beta, dC, Lc, dC, Lc,
&res[i].algo, ws, wsSize, 0) != CUBLAS_STATUS_SUCCESS) continue;
CHECK(cudaDeviceSynchronize());
cudaEvent_t e0, e1; CHECK(cudaEventCreate(&e0)); CHECK(cudaEventCreate(&e1));
CHECK(cudaEventRecord(e0));
for (int r = 0; r < reps; ++r)
cublasLtMatmul(lt, op, alpha, dA, La, dB, Lb, beta, dC, Lc, dC, Lc,
&res[i].algo, ws, wsSize, 0);
CHECK(cudaEventRecord(e1)); CHECK(cudaEventSynchronize(e1));
float ms = 0; CHECK(cudaEventElapsedTime(&ms, e0, e1)); ms /= (float)reps;
if (ms < best) { best = ms; besti = i; }
CHECK(cudaEventDestroy(e0)); CHECK(cudaEventDestroy(e1));
}
if (besti < 0) { printf(" %-14s m=%5d n=%5d k=%5d: ALL-ALGOS-FAILED (n=%d)\n", label, m, n, k, nres); goto done; }
printf(" %-14s m=%5d n=%5d k=%5d: best %7.3f ms %6.1f TFLOPS (algo %d of %d)\n",
label, m, n, k, best, 2.0 * m * (double)n * k / (best * 1e-3) / 1e12, besti, nres);
done:
if (pref) cublasLtMatmulPreferenceDestroy(pref);
if (Lc) cublasLtMatrixLayoutDestroy(Lc);
if (Lb) cublasLtMatrixLayoutDestroy(Lb);
if (La) cublasLtMatrixLayoutDestroy(La);
if (op) cublasLtMatmulDescDestroy(op);
}
} // namespace
int main(int argc, char **argv) {
const int reps = (argc > 1) ? atoi(argv[1]) : 10;
printf("### proto_gemm_lib cublasLt %zu reps=%d ###\n", cublasLtGetVersion(), reps);
cublasLtHandle_t lt; if (cublasLtCreate(&lt) != CUBLAS_STATUS_SUCCESS) { fprintf(stderr, "lt create fail\n"); return 1; }
const size_t maxA = (size_t)32768 * 8192 * 2; // covers kxm all shapes, f16
const size_t maxB = (size_t)8192 * 24576 * 2;
const size_t maxC = (size_t)32768 * 24576 * 4;
const size_t wsSize = 256u << 20;
void *dA, *dB, *dC, *ws;
CHECK(cudaMalloc(&dA, maxA)); CHECK(cudaMalloc(&dB, maxB));
CHECK(cudaMalloc(&dC, maxC)); CHECK(cudaMalloc(&ws, wsSize));
fillk<<<512, 256>>>((unsigned *)dA, maxA / 4, 0xA5A5u);
fillk<<<512, 256>>>((unsigned *)dB, maxB / 4, 0x5A5Au);
CHECK(cudaDeviceSynchronize());
/* MoE-equiv shapes + the dense-q8 census shapes (cmdq8shape15 2026-07-09):
* q_b [32768x1024], o_proj [4096x8192], sexp gate/up [2048x4096],
* sexp down [4096x2048], all at N=4096 chunk tokens. */
const int shapes[][3] = { {4096, 24576, 2048}, {2048, 24576, 4096},
{32768, 4096, 1024}, {4096, 4096, 8192},
{2048, 4096, 4096}, {4096, 4096, 2048} };
const char *kn[] = { "int8.s32acc", "f16.f32acc", "f16.f16acc" };
for (const auto &s : shapes) {
printf("--- shape m=%d n=%d k=%d ---\n", s[0], s[1], s[2]);
for (int kind = 0; kind < 3; ++kind)
bench_case(lt, dA, dB, dC, ws, wsSize, s[0], s[1], s[2], kind, reps, kn[kind]);
}
cublasLtDestroy(lt);
CHECK(cudaFree(dA)); CHECK(cudaFree(dB)); CHECK(cudaFree(dC)); CHECK(cudaFree(ws));
return 0;
}
cuda/mmq/test/proto_indexer_score.cu
+2430
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/test/proto_iq2_aligned.cu
+262
View File
@@ -0,0 +1,262 @@
// proto_iq2_aligned.cu — M1-MoE Inc1 prototype (megakernel program).
//
// HYPOTHESIS: the IQ2_XXS decode matvec runs at ~142 GB/s (vs the ~200
// engine ceiling) because block_iq2_xxs is 66 bytes, so the code stream is
// only 2-byte aligned and every 32-bit weight word costs two 16-bit loads
// (get_int_b2). Repacking into an aligned SoA layout (d[] halves separate,
// qs[] 64B-aligned per block) allows full-width aligned loads at identical
// byte count.
//
// A/B at the exact production decode shape (per layer, per token):
// gate/up experts: M=2048 rows, K=4096, top-6 of n_experts, n_tokens=1
// baseline = ds4_mmq_iq2_xxs_moe_vec (the production vec-tier entry)
// aligned = warp-per-row kernel, same dp4a/grid/sign math, SoA weights
//
// Correctness: aligned vs baseline output, tolerance (float accum order
// differs across blocks; per-pair integer math is identical).
//
// Build (box):
// nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
// proto_iq2_aligned.cu ../ds4_mmq.o -lcudart -lcuda -o proto_iq2_aligned
#include "ds4_mmq.h"
#define GGML_COMMON_DECL_CUDA
#define GGML_COMMON_IMPL_CUDA
#include "../ggml-common.h"
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
namespace {
constexpr int QKK = 256; // values per IQ2_XXS block
constexpr int BLOCK_BYTES = 66; // sizeof(block_iq2_xxs) = 2 + 32*2
constexpr int PAIRS_PER_BLOCK = 8; // 8 (q2,aux32) uint32 pairs, 32 values each
// ---------------------------------------------------------------- host q8_1
struct q8_1_block { __half2 ds; int8_t qs[32]; }; // mirrors block_q8_1
static void quantize_q8_1_host(const float *x, q8_1_block *out, int k) {
for (int b = 0; b < k / 32; b++) {
float amax = 0.0f;
for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
const float d = amax / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
int sum = 0;
for (int i = 0; i < 32; i++) {
int q = (int)roundf(x[b * 32 + i] * id);
if (q > 127) q = 127;
if (q < -127) q = -127;
out[b].qs[i] = (int8_t)q;
sum += q;
}
out[b].ds = __floats2half2_rn(d, d * (float)sum);
}
}
// ------------------------------------------------------------ device kernel
__device__ __forceinline__ uint32_t proto_unpack_ksigns(const uint8_t v) {
const uint32_t p = __popc(v) & 1;
const uint32_t s = v ^ p << 7;
return s * 0x01010101;
}
__device__ __forceinline__ int proto_dp4a(const int a, const int b, int c) {
return __dp4a(a, b, c);
}
// Aligned SoA weights: for (expert e, row r) with nb blocks per row:
// dq[(e*M + r)*nb + b] : __half block scale
// qs[((e*M + r)*nb + b) * 8 + p] : uint2 (q2, aux32) pair p
// qs base is 16B-aligned; each block's 8 uint2 = 64B, 64B-aligned stride.
__global__ void iq2_aligned_moe_vec_kernel(
float *out, // [n_slots, M]
const uint2 *qs, // aligned code pairs
const __half *dq, // block scales
const q8_1_block *x8, // [K/32] quantized activation (shared by slots)
const int32_t *ids, // [n_slots] expert ids
const float *slot_w, // [n_slots] router weights (1.0 in proto)
int M,
int nb) // blocks per row = K/256
{
const int row = blockIdx.x;
const int slot = blockIdx.y;
const int lane = threadIdx.x; // 32 lanes: lane covers (block b, pair p)
const long long rbase = ((long long)ids[slot] * M + row) * nb;
float acc = 0.0f;
// 32 lanes cover 4 blocks x 8 pairs per pass.
for (int b0 = 0; b0 < nb; b0 += 4) {
const int b = b0 + (lane >> 3);
const int p = lane & 7;
const uint2 cw = qs[(rbase + b) * PAIRS_PER_BLOCK + p]; // aligned 8B
const uint32_t q2 = cw.x, aux32 = cw.y;
const uint8_t *aux8 = (const uint8_t *)&q2;
int sumi = 0;
const int q8i = (b * QKK + p * 32) / 32; // q8_1 block covering these 32 values
const int *u = (const int *)x8[q8i].qs;
#pragma unroll
for (int k0 = 0; k0 < 8; k0 += 2) {
const uint2 grid_pos = ((const uint2 *)iq2xxs_grid)[aux8[k0 / 2]];
const uint32_t signs = proto_unpack_ksigns(aux32 >> (7 * k0 / 2));
const int signs0 = __vcmpne4(signs & 0x08040201, 0);
const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0);
sumi = proto_dp4a(grid0, u[k0 + 0], sumi);
const int signs1 = __vcmpne4(signs & 0x80402010, 0);
const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1);
sumi = proto_dp4a(grid1, u[k0 + 1], sumi);
}
const int ls = aux32 >> 27 | 1;
sumi = sumi * ls / 8;
const float d = __half2float(dq[rbase + b]) * __low2float(x8[q8i].ds);
acc += d * (float)sumi;
}
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
acc += __shfl_down_sync(0xffffffffu, acc, off);
if (lane == 0) out[(long long)slot * M + row] = acc * slot_w[slot];
}
} // namespace
int main(int argc, char **argv) {
const int M = 2048, K = 4096;
const int n_experts = 256;
const int n_slots = 6;
const int n_idsets = 32;
const int iters = 2000;
const int nb = K / QKK; // 16 blocks/row
const long long row_bytes = (long long)nb * BLOCK_BYTES; // 1056
const double slot_gb = (double)n_slots * M * row_bytes / 1e9;
std::mt19937 rng(argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u);
// Random IQ2_XXS blocks (same trick as test_mmq_parity: random bytes are
// valid IQ2_XXS content; both kernels see identical bits).
std::vector<uint8_t> W((size_t)n_experts * M * row_bytes);
for (auto &b : W) b = (uint8_t)(rng() & 0xff);
// keep block d halves sane (avoid inf/nan): overwrite each block's d
for (long long blk = 0; blk < (long long)n_experts * M * nb; blk++) {
uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff)); // ~0.0625..0.1 range fp16
memcpy(&W[blk * BLOCK_BYTES], &h, 2);
}
std::vector<float> X(K);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<q8_1_block> X8(K / 32);
quantize_q8_1_host(X.data(), X8.data(), K);
std::vector<int32_t> ids(n_idsets * n_slots);
for (int s = 0; s < n_idsets; s++)
for (int j = 0; j < n_slots; j++)
ids[s * n_slots + j] = (s * n_slots + j) % n_experts;
// Aligned SoA repack
const long long nblocks = (long long)n_experts * M * nb;
std::vector<uint2> QS(nblocks * PAIRS_PER_BLOCK);
std::vector<__half> DQ(nblocks);
for (long long blk = 0; blk < nblocks; blk++) {
const uint8_t *src = &W[blk * BLOCK_BYTES];
uint16_t h; memcpy(&h, src, 2);
DQ[blk] = __ushort_as_half(h);
memcpy(&QS[blk * PAIRS_PER_BLOCK], src + 2, 64);
}
// Device buffers
uint8_t *dW; float *dX, *dOutBase, *dOutAl; int32_t *dIds; uint2 *dQS; __half *dDQ;
q8_1_block *dX8; float *dSlotW;
CK(cudaMalloc(&dW, W.size()));
CK(cudaMalloc(&dX, sizeof(float) * K * n_slots));
CK(cudaMalloc(&dOutBase, sizeof(float) * n_slots * M));
CK(cudaMalloc(&dOutAl, sizeof(float) * n_slots * M));
CK(cudaMalloc(&dIds, sizeof(int32_t) * ids.size()));
CK(cudaMalloc(&dQS, QS.size() * sizeof(uint2)));
CK(cudaMalloc(&dDQ, DQ.size() * sizeof(__half)));
CK(cudaMalloc(&dX8, X8.size() * sizeof(q8_1_block)));
CK(cudaMalloc(&dSlotW, sizeof(float) * n_slots));
CK(cudaMemcpy(dW, W.data(), W.size(), cudaMemcpyHostToDevice));
// moe_vec's X layout: [n_tokens*n_expert_used rows? or one row]; the vec
// entry treats X as [n_tokens, K] with router weights applied later; we
// pass the same activation replicated per slot to keep both paths equal.
for (int s = 0; s < n_slots; s++)
CK(cudaMemcpy(dX + (size_t)s * K, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
CK(cudaMemcpy(dIds, ids.data(), sizeof(int32_t) * ids.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dQS, QS.data(), QS.size() * sizeof(uint2), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dDQ, DQ.data(), DQ.size() * sizeof(__half), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dX8, X8.data(), X8.size() * sizeof(q8_1_block), cudaMemcpyHostToDevice));
std::vector<float> slotw(n_slots, 1.0f);
CK(cudaMemcpy(dSlotW, slotw.data(), sizeof(float) * n_slots, cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
// ---- baseline: production vec entry -----------------------------------
int rc = ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
if (rc != 0) { printf("baseline moe_vec rc=%d\n", rc); return 1; }
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds + (i % n_idsets) * n_slots, dOutBase,
M, K, 1, n_experts, n_slots, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
ms_base /= iters;
// ---- aligned prototype -------------------------------------------------
dim3 grid(M, n_slots, 1);
iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8, dIds, dSlotW, M, nb);
CK(cudaGetLastError());
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8,
dIds + (i % n_idsets) * n_slots,
dSlotW, M, nb);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
ms_al /= iters;
// ---- correctness (vs baseline, last idset used by both loops) ----------
// rerun both once on idset 0 for a clean compare
(void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8, dIds, dSlotW, M, nb);
CK(cudaStreamSynchronize(stream));
std::vector<float> ob(n_slots * M), oa(n_slots * M);
CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * ob.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * oa.size(), cudaMemcpyDeviceToHost));
double max_rel = 0.0, max_abs = 0.0; int bad = 0;
for (size_t i = 0; i < ob.size(); i++) {
const double a = ob[i], b = oa[i];
const double ad = fabs(a - b);
const double rd = ad / (fabs(a) > 1e-3 ? fabs(a) : 1e-3);
if (ad > max_abs) max_abs = ad;
if (rd > max_rel) max_rel = rd;
if (rd > 2e-2 && ad > 1e-2) bad++;
}
printf("PROTO_IQ2_ALIGNED M=%d K=%d slots=%d experts=%d iters=%d (weights/iter %.1f MB)\n",
M, K, n_slots, n_experts, iters, slot_gb * 1000.0);
printf(" baseline moe_vec : %.4f ms -> %6.1f GB/s\n", ms_base, slot_gb / (ms_base / 1e3));
printf(" aligned proto : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_al, slot_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
printf(" parity: max_rel=%.3e max_abs=%.3e bad=%d -> %s\n",
max_rel, max_abs, bad, bad == 0 ? "PASS" : "FAIL");
return bad == 0 ? 0 : 2;
}
cuda/mmq/test/proto_kv_reencode_idem.cu
+858
View File
@@ -0,0 +1,858 @@
// SPDX-License-Identifier: MIT
// proto_kv_reencode_idem.cu - packed-KV restore re-encode idempotence probe.
//
// v0.2.4 ships REFUSED packed disk-KV restores: the restore path reads the
// F32 committed rows from the file and regenerates the fp8/fp4 mirrors by
// re-encoding, and the kv_crossmode gate observed the resulting packed-server
// continuations DRIFTING (fp8-only and fp4-only independently) despite the
// pow2-scale idempotence argument saying re-encode of committed values is
// value-exact. This proto isolates the kernel math from the pipeline:
// fp8 : quantize(x) -> C1,codes1,scales1 ; dequant==C1 ;
// quantize(C1) -> C2,codes2,scales2 ; C2 ?= C1 (bitwise) ; dequant==C1
// fp4 : hadamard_qat(x) -> C1 ; fp4_dequant==C1 ;
// noHad_encode(C1) -> C2 ; C2 ?= C1 ; dequant==C1
// plus a third-iteration fixed-point check and a fast-math boundary scan of
// scale = exp2f(ceilf(log2f(amax/448 or /6))) at exact power-of-two inputs
// (the one corner where fast-math approximation could flip ceilf and shift
// the whole block one exponent).
// Code/scale mismatches with VALUE identity are tallied separately: a block
// whose committed max sits exactly on a half-boundary legitimately re-picks
// scale/2 with all codes shifted one exponent, values bit-identical.
//
// The first battery runs verbatim copies of the v0.2.4 (PRE-fix) ds4_cuda.cu
// bodies and is EXPECTED to fail -- kept as the regression repro. The _fx
// battery is the fix as shipped in v0.3 (dsv4_pow2_ceil_scale + signbit) and
// is the gate (exit status). Build MUST mirror the engine's fast-math flags:
// /usr/local/cuda/bin/nvcc -O3 -std=c++17 --use_fast_math -arch=sm_121 \
// cuda/mmq/test/proto_kv_reencode_idem.cu -o proto_kv_reencode_idem
// argv: [n_rows_fp8=131072] [n_rows_fp4=131072]
#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <cstdint>
#define CHECK(call) do { cudaError_t e_=(call); if(e_!=cudaSuccess){ \
fprintf(stderr,"CUDA %s @%d: %s\n",#call,__LINE__,cudaGetErrorString(e_)); exit(1);} } while(0)
// ---------------------------------------------------------------------------
// Verbatim from ds4_cuda.cu (v0.2.4 tag content).
// ---------------------------------------------------------------------------
#define DS4_OPP_C_FP8_NOPE_DEV 448u /* 512 - 64 */
#define DS4_OPP_C_FP8_BLOCKS_DEV 7u /* 448 / 64 */
#define DS4_OPP_C_FP8_ROW_BYTES_DEV 704u /* 448 codes + 64 FP32 tail */
__device__ static float dsv4_e4m3fn_value_dev(int i) {
int exp = (i >> 3) & 15;
int mant = i & 7;
if (exp == 0) return (float)mant * 0.001953125f;
return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
}
__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
float sign = x < 0.0f ? -1.0f : 1.0f;
float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return sign * dsv4_e4m3fn_value_dev(best);
}
__device__ static unsigned char dsv4_e4m3fn_encode_dev(float x) {
const unsigned int sign_bit = (x < 0.0f) ? 0x80u : 0x00u;
const float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
const int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
const float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
const float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return (unsigned char)(sign_bit | (unsigned int)best);
}
__device__ float dsv4_e4m3fn_decode_table[128];
static int fp8_decode_table_init(void) {
float table[128];
for (int i = 0; i < 128; i++) {
const int exp = (i >> 3) & 15;
const int mant = i & 7;
table[i] = (exp == 0)
? (float)mant * 0.001953125f
: (1.0f + (float)mant * 0.125f) * ldexpf(1.0f, exp - 7);
}
return cudaMemcpyToSymbol(dsv4_e4m3fn_decode_table, table,
sizeof(table), 0, cudaMemcpyHostToDevice) == cudaSuccess;
}
__device__ static float fp8_kv_read(
const unsigned char * __restrict__ codes_base,
const float * __restrict__ scale_base,
uint32_t c,
uint32_t d) {
const uint64_t row_off = (uint64_t)c * DS4_OPP_C_FP8_ROW_BYTES_DEV;
if (d < DS4_OPP_C_FP8_NOPE_DEV) {
const unsigned char code = codes_base[row_off + d];
const float scale = scale_base[(uint64_t)c * DS4_OPP_C_FP8_BLOCKS_DEV + (d >> 6)];
const float mag = dsv4_e4m3fn_decode_table[code & 0x7fu];
return ((code & 0x80u) ? -mag : mag) * scale;
}
const float *tail = (const float *)(codes_base + row_off + DS4_OPP_C_FP8_NOPE_DEV);
return tail[d - DS4_OPP_C_FP8_NOPE_DEV];
}
__global__ static void fp8_kv_dequant_rows_kernel(
float * __restrict__ dst,
const unsigned char * __restrict__ codes_base,
const float * __restrict__ scale_base,
uint32_t n_rows,
uint32_t head_dim) {
const uint32_t r = blockIdx.x;
if (r >= n_rows) return;
float *out = dst + (uint64_t)r * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
out[d] = fp8_kv_read(codes_base, scale_base, r, d);
}
}
__global__ static void fp8_kv_quantize_kernel(
float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
unsigned char * __restrict__ codes_base,
float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
uint32_t n_nope = head_dim - n_rot;
const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
float *xr = x + (uint64_t)row * head_dim;
unsigned char *codes_row = codes_base
? codes_base + (uint64_t)row * row_stride
: (unsigned char *)0;
float *scale_row = scale_base
? scale_base + (uint64_t)row * scale_stride
: (float *)0;
__shared__ float scratch[64];
for (uint32_t off = 0; off < n_nope; off += 64) {
float v = 0.0f;
if (off + tid < n_nope) v = xr[off + tid];
scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
__syncthreads();
}
float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
if (off + tid < n_nope) {
float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
xr[off + tid] = dsv4_e4m3fn_dequant_dev(clamp) * scale;
if (codes_row) codes_row[off + tid] = dsv4_e4m3fn_encode_dev(clamp);
}
if (scale_row && tid == 0) scale_row[off / 64u] = scale;
__syncthreads();
}
if (codes_row && tid < n_rot) {
float *tail = (float *)(codes_row + (uint64_t)n_nope);
tail[tid] = xr[(uint64_t)n_nope + tid];
}
}
__device__ static float dsv4_e2m1fn_value_dev(int i) {
switch (i & 7) {
case 0: return 0.0f;
case 1: return 0.5f;
case 2: return 1.0f;
case 3: return 1.5f;
case 4: return 2.0f;
case 5: return 3.0f;
case 6: return 4.0f;
default: return 6.0f;
}
}
__device__ static int dsv4_e2m1fn_level_dev(float ax) {
ax = fminf(ax, 6.0f);
int best = 0;
float best_diff = fabsf(ax - dsv4_e2m1fn_value_dev(0));
for (int i = 1; i < 8; i++) {
float diff = fabsf(ax - dsv4_e2m1fn_value_dev(i));
if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
best = i;
best_diff = diff;
}
}
return best;
}
__device__ static float indexer_fp4_read(
const unsigned char * __restrict__ codes_base,
const float * __restrict__ scale_base,
uint32_t row, uint32_t d) {
const unsigned char byte = codes_base[(uint64_t)row * 64u + (d >> 1u)];
const uint32_t nib = ((uint32_t)byte >> (4u * (d & 1u))) & 0xFu;
const float sign = (nib & 8u) ? -1.0f : 1.0f;
const float scale = scale_base[(uint64_t)row * 4u + (d >> 5u)];
return sign * dsv4_e2m1fn_value_dev((int)(nib & 7u)) * scale;
}
__global__ static void indexer_fp4_dequant_rows_kernel(
float * __restrict__ dst,
const unsigned char * __restrict__ codes_base,
const float * __restrict__ scale_base,
uint32_t n_rows, uint32_t head_dim) {
const uint32_t r = blockIdx.x;
if (r >= n_rows) return;
float *out = dst + (uint64_t)r * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
out[d] = indexer_fp4_read(codes_base, scale_base, r, d);
}
}
__global__ static void indexer_fp4_encode_rows_kernel(
float * __restrict__ x, uint32_t n_rows, uint32_t head_dim,
unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
__shared__ float absbuf[128];
float *xr = x + (uint64_t)row * head_dim;
float v = xr[tid]; /* committed value -- NO Hadamard */
uint32_t fp4_block = tid >> 5u;
uint32_t lane = tid & 31u;
uint32_t block_base = fp4_block * 32u;
absbuf[tid] = fabsf(v);
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride)
absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
__syncthreads();
}
float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
float scale = exp2f(ceilf(log2f(amax / 6.0f)));
float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
xr[tid] = ((clamp < 0.0f ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
uint32_t nib = (clamp < 0.0f ? 8u : 0u) | (uint32_t)level;
uint32_t hi = __shfl_down_sync(0xffffffffu, nib, 1u);
if ((tid & 1u) == 0u) codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
}
__global__ static void indexer_hadamard_fp4_kernel(float *x, uint32_t n_rows, uint32_t head_dim,
unsigned char * __restrict__ codes_base,
float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
__shared__ float vals[128];
__shared__ float absbuf[128];
float *xr = x + (uint64_t)row * head_dim;
vals[tid] = xr[tid];
__syncthreads();
for (uint32_t stride = 1u; stride < 128u; stride <<= 1u) {
if ((tid & stride) == 0u) {
uint32_t base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
float a = vals[base];
float b = vals[base + stride];
vals[base] = a + b;
vals[base + stride] = a - b;
}
__syncthreads();
}
float v = vals[tid] * 0.08838834764831845f;
uint32_t fp4_block = tid >> 5u;
uint32_t lane = tid & 31u;
uint32_t block_base = fp4_block * 32u;
absbuf[tid] = fabsf(v);
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride) {
absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
}
__syncthreads();
}
float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
float scale = exp2f(ceilf(log2f(amax / 6.0f)));
float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
xr[tid] = ((clamp < 0.0f ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
if (codes_base) {
uint32_t nib = (clamp < 0.0f ? 8u : 0u) | (uint32_t)level;
uint32_t hi = __shfl_down_sync(0xffffffffu, nib, 1u);
if ((tid & 1u) == 0u) codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
}
}
// ---------------------------------------------------------------------------
// FIXED variants (candidate engine change): exact pow2-ceil scale via frexpf
// (bit ops, immune to fast-math lg2.approx one-sided error at 2^-n) and
// signbit-based sign handling (preserves -0.0 across the round trip).
// ---------------------------------------------------------------------------
__device__ static float dsv4_pow2_ceil_scale(float r) {
int e; float m = frexpf(r, &e); /* r = m * 2^e, m in [0.5, 1) */
return ldexpf(1.0f, m == 0.5f ? e - 1 : e);
}
__device__ static float dsv4_e4m3fn_dequant_fx(float x) {
float sign = signbit(x) ? -1.0f : 1.0f;
float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return sign * dsv4_e4m3fn_value_dev(best);
}
__device__ static unsigned char dsv4_e4m3fn_encode_fx(float x) {
const unsigned int sign_bit = signbit(x) ? 0x80u : 0x00u;
const float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
const int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
const float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
const float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return (unsigned char)(sign_bit | (unsigned int)best);
}
__global__ static void fp8_kv_quantize_kernel_fx(
float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
unsigned char * __restrict__ codes_base,
float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
uint32_t n_nope = head_dim - n_rot;
const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
float *xr = x + (uint64_t)row * head_dim;
unsigned char *codes_row = codes_base
? codes_base + (uint64_t)row * row_stride
: (unsigned char *)0;
float *scale_row = scale_base
? scale_base + (uint64_t)row * scale_stride
: (float *)0;
__shared__ float scratch[64];
for (uint32_t off = 0; off < n_nope; off += 64) {
float v = 0.0f;
if (off + tid < n_nope) v = xr[off + tid];
scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
__syncthreads();
}
float scale = dsv4_pow2_ceil_scale(fmaxf(scratch[0], 1.0e-4f) / 448.0f);
if (off + tid < n_nope) {
float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
xr[off + tid] = dsv4_e4m3fn_dequant_fx(clamp) * scale;
if (codes_row) codes_row[off + tid] = dsv4_e4m3fn_encode_fx(clamp);
}
if (scale_row && tid == 0) scale_row[off / 64u] = scale;
__syncthreads();
}
if (codes_row && tid < n_rot) {
float *tail = (float *)(codes_row + (uint64_t)n_nope);
tail[tid] = xr[(uint64_t)n_nope + tid];
}
}
__global__ static void indexer_fp4_encode_rows_kernel_fx(
float * __restrict__ x, uint32_t n_rows, uint32_t head_dim,
unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
__shared__ float absbuf[128];
float *xr = x + (uint64_t)row * head_dim;
float v = xr[tid];
uint32_t fp4_block = tid >> 5u;
uint32_t lane = tid & 31u;
uint32_t block_base = fp4_block * 32u;
absbuf[tid] = fabsf(v);
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride)
absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
__syncthreads();
}
float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
float scale = dsv4_pow2_ceil_scale(amax / 6.0f);
float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
xr[tid] = ((signbit(clamp) ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
uint32_t nib = (signbit(clamp) ? 8u : 0u) | (uint32_t)level;
uint32_t hi = __shfl_down_sync(0xffffffffu, nib, 1u);
if ((tid & 1u) == 0u) codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
}
__global__ static void indexer_hadamard_fp4_kernel_fx(float *x, uint32_t n_rows, uint32_t head_dim,
unsigned char * __restrict__ codes_base,
float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
__shared__ float vals[128];
__shared__ float absbuf[128];
float *xr = x + (uint64_t)row * head_dim;
vals[tid] = xr[tid];
__syncthreads();
for (uint32_t stride = 1u; stride < 128u; stride <<= 1u) {
if ((tid & stride) == 0u) {
uint32_t base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
float a = vals[base];
float b = vals[base + stride];
vals[base] = a + b;
vals[base + stride] = a - b;
}
__syncthreads();
}
float v = vals[tid] * 0.08838834764831845f;
uint32_t fp4_block = tid >> 5u;
uint32_t lane = tid & 31u;
uint32_t block_base = fp4_block * 32u;
absbuf[tid] = fabsf(v);
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride) {
absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
}
__syncthreads();
}
float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
float scale = dsv4_pow2_ceil_scale(amax / 6.0f);
float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
xr[tid] = ((signbit(clamp) ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
if (codes_base) {
uint32_t nib = (signbit(clamp) ? 8u : 0u) | (uint32_t)level;
uint32_t hi = __shfl_down_sync(0xffffffffu, nib, 1u);
if ((tid & 1u) == 0u) codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
}
}
// ---------------------------------------------------------------------------
// Proto scaffolding (not engine code).
// ---------------------------------------------------------------------------
// Deterministic per-element fill: xorshift hash of (row,lane) -> [-1,1),
// scaled per row-band regime. Bands (of 8) for the fp8 input:
// 0 uniform | 1 x1e-5 | 2 x1e4 | 3 pow2 magnitude ladder 2^((row%40)-20)
// 4 block amax forced EXACTLY 448*2^k (k=(row%40)-20; lane0 of block)
// 5 block amax forced EXACTLY 224*2^k (post-encode halving boundary)
// 6 block amax 433*2^k (e4m3 rounds it UP to 448*2^k -> boundary amax')
// 7 tiny/zero: x1e-6, every 3rd row all-zero (sub-1e-4-floor blocks)
// Exact 2^k as a float, immune to any exp2f approximation question.
__device__ static float exact_pow2(int k) {
return __uint_as_float((uint32_t)(127 + k) << 23);
}
__device__ static float hash_unit(uint32_t row, uint32_t lane, uint32_t salt) {
uint32_t h = row * 2654435761u ^ (lane + 1u) * 2246822519u ^ salt * 3266489917u;
h ^= h >> 15; h *= 2246822519u; h ^= h >> 13; h *= 3266489917u; h ^= h >> 16;
return ((float)(h & 0xFFFFFFu) / 8388608.0f) - 1.0f; // [-1,1)
}
__global__ static void fill_fp8_input_kernel(float *x, uint32_t n_rows, uint32_t salt) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x; // 512 threads
if (row >= n_rows || tid >= 512u) return;
uint32_t band = row / (n_rows / 8u + 1u);
float u = hash_unit(row, tid, salt);
float v = u;
int k = (int)(row % 40u) - 20;
switch (band) {
case 0: break;
case 1: v = u * 1e-5f; break;
case 2: v = u * 1e4f; break;
case 3: v = u * exact_pow2(k); break;
case 4: v = (tid < 448u && (tid & 63u) == 0u) ? 448.0f * exact_pow2(k)
: u * 100.0f * exact_pow2(k); break;
case 5: v = (tid < 448u && (tid & 63u) == 0u) ? 224.0f * exact_pow2(k)
: u * 100.0f * exact_pow2(k); break;
case 6: v = (tid < 448u && (tid & 63u) == 0u) ? 433.0f * exact_pow2(k)
: u * 100.0f * exact_pow2(k); break;
default:
v = (row % 3u == 0u) ? 0.0f : u * 1e-6f;
break;
}
x[(uint64_t)row * 512u + tid] = v;
}
__global__ static void fill_fp4_input_kernel(float *x, uint32_t n_rows, uint32_t salt) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x; // 128 threads
if (row >= n_rows || tid >= 128u) return;
uint32_t band = row / (n_rows / 6u + 1u);
float u = hash_unit(row, tid, salt ^ 0x9e3779b9u);
float v = u;
int k = (int)(row % 60u) - 30;
switch (band) {
case 0: break;
case 1: v = u * 1e-6f; break;
case 2: v = u * 1e5f; break;
case 3: v = u * exact_pow2(k); break;
case 4: v = (u > 0.0f ? 1.0f : -1.0f) * exact_pow2(k); break; // pure pow2 lanes
default:
v = (row % 3u == 0u) ? 0.0f : u * 1e-30f;
break;
}
x[(uint64_t)row * 128u + tid] = v;
}
struct mismatch_rec { uint32_t idx_hi, idx_lo, a_bits, b_bits; };
#define MAX_RECS 8
__global__ static void compare_f32_kernel(const float *a, const float *b, uint64_t n,
unsigned long long *count,
mismatch_rec *recs, int *nrec) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
uint32_t ab = __float_as_uint(a[i]);
uint32_t bb = __float_as_uint(b[i]);
if (ab != bb) {
atomicAdd(count, 1ull);
int slot = atomicAdd(nrec, 1);
if (slot < MAX_RECS) {
recs[slot].idx_hi = (uint32_t)(i >> 32);
recs[slot].idx_lo = (uint32_t)(i & 0xFFFFFFFFu);
recs[slot].a_bits = ab;
recs[slot].b_bits = bb;
}
}
}
__global__ static void compare_u8_kernel(const unsigned char *a, const unsigned char *b,
uint64_t n, unsigned long long *count) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
if (a[i] != b[i]) atomicAdd(count, 1ull);
}
// Fast-math boundary scan: the exact-pow2 inputs the re-encode presents when
// a block's committed max lands on 448*2^k (fp8) or {6,3,1.5}*2^k (fp4).
// out[j] = scale computed by the shipped formula; host prints deviations.
__global__ static void boundary_scan_kernel(float *out) {
int idx = threadIdx.x; // 61 k values: -30..30
if (idx >= 61) return;
int k = idx - 30;
float p2k = exact_pow2(k);
// fp8 formula at amax = 448*2^k, 224*2^k, and one-ulp neighbours of 448*2^k
float a448 = 448.0f * p2k;
float a224 = 224.0f * p2k;
float a448u = __uint_as_float(__float_as_uint(a448) + 1u);
float a448d = __uint_as_float(__float_as_uint(a448) - 1u);
out[idx*8 + 0] = exp2f(ceilf(log2f(fmaxf(a448, 1.0e-4f) / 448.0f)));
out[idx*8 + 1] = exp2f(ceilf(log2f(fmaxf(a224, 1.0e-4f) / 448.0f)));
out[idx*8 + 2] = exp2f(ceilf(log2f(fmaxf(a448u, 1.0e-4f) / 448.0f)));
out[idx*8 + 3] = exp2f(ceilf(log2f(fmaxf(a448d, 1.0e-4f) / 448.0f)));
// fp4 formula at amax = 6*2^k, 3*2^k, 1.5*2^k, one-ulp above 6*2^k
float a6 = 6.0f * p2k;
float a3 = 3.0f * p2k;
float a15 = 1.5f * p2k;
float a6u = __uint_as_float(__float_as_uint(a6) + 1u);
out[idx*8 + 4] = exp2f(ceilf(log2f(fmaxf(a6, 7.052966104933725e-38f) / 6.0f)));
out[idx*8 + 5] = exp2f(ceilf(log2f(fmaxf(a3, 7.052966104933725e-38f) / 6.0f)));
out[idx*8 + 6] = exp2f(ceilf(log2f(fmaxf(a15, 7.052966104933725e-38f) / 6.0f)));
out[idx*8 + 7] = exp2f(ceilf(log2f(fmaxf(a6u, 7.052966104933725e-38f) / 6.0f)));
}
static unsigned long long report_f32(const char *tag, const float *a, const float *b,
uint64_t n, uint32_t row_dim) {
unsigned long long *d_count; mismatch_rec *d_recs; int *d_nrec;
CHECK(cudaMalloc(&d_count, sizeof(*d_count)));
CHECK(cudaMalloc(&d_recs, MAX_RECS * sizeof(mismatch_rec)));
CHECK(cudaMalloc(&d_nrec, sizeof(int)));
CHECK(cudaMemset(d_count, 0, sizeof(*d_count)));
CHECK(cudaMemset(d_nrec, 0, sizeof(int)));
uint64_t blocks = (n + 255) / 256;
compare_f32_kernel<<<(uint32_t)blocks, 256>>>(a, b, n, d_count, d_recs, d_nrec);
CHECK(cudaDeviceSynchronize());
unsigned long long count; mismatch_rec recs[MAX_RECS]; int nrec;
CHECK(cudaMemcpy(&count, d_count, sizeof(count), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(recs, d_recs, sizeof(recs), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(&nrec, d_nrec, sizeof(nrec), cudaMemcpyDeviceToHost));
printf(" %-28s mismatches=%llu / %llu\n", tag, count, (unsigned long long)n);
for (int i = 0; i < nrec && i < MAX_RECS; i++) {
uint64_t idx = ((uint64_t)recs[i].idx_hi << 32) | recs[i].idx_lo;
float fa, fb;
memcpy(&fa, &recs[i].a_bits, 4); memcpy(&fb, &recs[i].b_bits, 4);
printf(" [%d] row=%llu lane=%u a=%08x(%.9g) b=%08x(%.9g)\n",
i, (unsigned long long)(idx / row_dim), (uint32_t)(idx % row_dim),
recs[i].a_bits, fa, recs[i].b_bits, fb);
}
cudaFree(d_count); cudaFree(d_recs); cudaFree(d_nrec);
return count;
}
static unsigned long long report_u8(const char *tag, const unsigned char *a,
const unsigned char *b, uint64_t n) {
unsigned long long *d_count;
CHECK(cudaMalloc(&d_count, sizeof(*d_count)));
CHECK(cudaMemset(d_count, 0, sizeof(*d_count)));
uint64_t blocks = (n + 255) / 256;
compare_u8_kernel<<<(uint32_t)blocks, 256>>>(a, b, n, d_count);
CHECK(cudaDeviceSynchronize());
unsigned long long count;
CHECK(cudaMemcpy(&count, d_count, sizeof(count), cudaMemcpyDeviceToHost));
printf(" %-28s byte-mismatches=%llu / %llu (level-shift benign iff values match)\n",
tag, count, (unsigned long long)n);
cudaFree(d_count);
return count;
}
int main(int argc, char **argv) {
uint32_t n_fp8 = argc > 1 ? (uint32_t)atoi(argv[1]) : 131072u;
uint32_t n_fp4 = argc > 2 ? (uint32_t)atoi(argv[2]) : 131072u;
if (!fp8_decode_table_init()) { fprintf(stderr, "decode table init failed\n"); return 1; }
int fails = 0; /* FIXED-variant failures: the actual gate */
int orig_fail = 0; /* shipped-kernel failures: expected repro of the bug */
// ---------------- fp8 comp path ----------------
{
const uint32_t HD = 512u, NROT = 64u;
uint64_t xn = (uint64_t)n_fp8 * HD;
uint64_t cb = (uint64_t)n_fp8 * DS4_OPP_C_FP8_ROW_BYTES_DEV;
uint64_t sn = (uint64_t)n_fp8 * DS4_OPP_C_FP8_BLOCKS_DEV;
float *x, *d1, *x2, *d2, *s1, *s2;
unsigned char *c1, *c2;
CHECK(cudaMalloc(&x, xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
fill_fp8_input_kernel<<<n_fp8, 512>>>(x, n_fp8, 7u);
fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x, n_fp8, HD, NROT, c1, s1); // -> C1
fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d1, c1, s1, n_fp8, HD); // D1
CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2); // -> C2
fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d2, c2, s2, n_fp8, HD); // D2
CHECK(cudaDeviceSynchronize());
printf("FP8 comp path (%u rows, head_dim %u):\n", n_fp8, HD);
unsigned long long m_read = report_f32("read-identity D1 vs C1", d1, x, xn, HD);
unsigned long long m_idem = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x, xn, HD);
report_u8("codes2 vs codes1", c1, c2, cb);
report_f32("scales2 vs scales1", s1, s2, sn, DS4_OPP_C_FP8_BLOCKS_DEV);
// third iteration: fixed point
fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
CHECK(cudaDeviceSynchronize());
unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
printf("REENC-IDEM FP8: %s\n\n", ok ? "PASS" : "FAIL (repro)");
if (!ok) orig_fail++;
cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
}
// ---------------- fp4 indexer path ----------------
{
const uint32_t HD = 128u;
uint64_t xn = (uint64_t)n_fp4 * HD;
uint64_t cb = (uint64_t)n_fp4 * 64u;
uint64_t sn = (uint64_t)n_fp4 * 4u;
float *x, *d1, *x2, *d2, *s1, *s2;
unsigned char *c1, *c2;
CHECK(cudaMalloc(&x, xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
fill_fp4_input_kernel<<<n_fp4, 128>>>(x, n_fp4, 7u);
indexer_hadamard_fp4_kernel<<<n_fp4, 128>>>(x, n_fp4, HD, c1, s1); // QAT -> C1
indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d1, c1, s1, n_fp4, HD); // D1
CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
indexer_fp4_encode_rows_kernel<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2); // load-path re-encode -> C2
indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d2, c2, s2, n_fp4, HD); // D2
CHECK(cudaDeviceSynchronize());
printf("FP4 indexer path (%u rows, head_dim %u):\n", n_fp4, HD);
unsigned long long m_read = report_f32("read-identity D1 vs C1", d1, x, xn, HD);
unsigned long long m_idem = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x, xn, HD);
report_u8("codes2 vs codes1", c1, c2, cb);
report_f32("scales2 vs scales1", s1, s2, sn, 4u);
indexer_fp4_encode_rows_kernel<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
CHECK(cudaDeviceSynchronize());
unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
printf("REENC-IDEM FP4: %s\n\n", ok ? "PASS" : "FAIL (repro)");
if (!ok) orig_fail++;
cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
}
// ---------------- FIXED fp8 comp path ----------------
{
const uint32_t HD = 512u, NROT = 64u;
uint64_t xn = (uint64_t)n_fp8 * HD;
uint64_t cb = (uint64_t)n_fp8 * DS4_OPP_C_FP8_ROW_BYTES_DEV;
uint64_t sn = (uint64_t)n_fp8 * DS4_OPP_C_FP8_BLOCKS_DEV;
float *x, *d1, *x2, *d2, *s1, *s2;
unsigned char *c1, *c2;
CHECK(cudaMalloc(&x, xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
fill_fp8_input_kernel<<<n_fp8, 512>>>(x, n_fp8, 7u);
fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x, n_fp8, HD, NROT, c1, s1);
fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d1, c1, s1, n_fp8, HD);
CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d2, c2, s2, n_fp8, HD);
CHECK(cudaDeviceSynchronize());
printf("FIXED FP8 comp path (frexpf scale + signbit; %u rows):\n", n_fp8);
unsigned long long m_read = report_f32("read-identity D1 vs C1", d1, x, xn, HD);
unsigned long long m_idem = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x, xn, HD);
unsigned long long m_code = report_u8("codes2 vs codes1", c1, c2, cb);
unsigned long long m_scale = report_f32("scales2 vs scales1", s1, s2, sn, DS4_OPP_C_FP8_BLOCKS_DEV);
fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
CHECK(cudaDeviceSynchronize());
unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
printf("REENC-IDEM FP8-FIXED: %s%s\n\n", ok ? "PASS" : "FAIL",
(ok && (m_code || m_scale)) ? " (codes level-shift, values exact)" : "");
if (!ok) fails++;
cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
}
// ---------------- FIXED fp4 indexer path ----------------
{
const uint32_t HD = 128u;
uint64_t xn = (uint64_t)n_fp4 * HD;
uint64_t cb = (uint64_t)n_fp4 * 64u;
uint64_t sn = (uint64_t)n_fp4 * 4u;
float *x, *d1, *x2, *d2, *s1, *s2;
unsigned char *c1, *c2;
CHECK(cudaMalloc(&x, xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
fill_fp4_input_kernel<<<n_fp4, 128>>>(x, n_fp4, 7u);
indexer_hadamard_fp4_kernel_fx<<<n_fp4, 128>>>(x, n_fp4, HD, c1, s1);
indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d1, c1, s1, n_fp4, HD);
CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
indexer_fp4_encode_rows_kernel_fx<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d2, c2, s2, n_fp4, HD);
CHECK(cudaDeviceSynchronize());
printf("FIXED FP4 indexer path (frexpf scale + signbit; %u rows):\n", n_fp4);
unsigned long long m_read = report_f32("read-identity D1 vs C1", d1, x, xn, HD);
unsigned long long m_idem = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x, xn, HD);
unsigned long long m_code = report_u8("codes2 vs codes1", c1, c2, cb);
unsigned long long m_scale = report_f32("scales2 vs scales1", s1, s2, sn, 4u);
indexer_fp4_encode_rows_kernel_fx<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
CHECK(cudaDeviceSynchronize());
unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
printf("REENC-IDEM FP4-FIXED: %s%s\n\n", ok ? "PASS" : "FAIL",
(ok && (m_code || m_scale)) ? " (codes level-shift, values exact)" : "");
if (!ok) fails++;
cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
}
// ---------------- fast-math boundary scan ----------------
{
float *d_out;
CHECK(cudaMalloc(&d_out, 61 * 8 * sizeof(float)));
boundary_scan_kernel<<<1, 64>>>(d_out);
CHECK(cudaDeviceSynchronize());
float out[61 * 8];
CHECK(cudaMemcpy(out, d_out, sizeof(out), cudaMemcpyDeviceToHost));
cudaFree(d_out);
int devs = 0, ulps = 0;
printf("Fast-math boundary scan (exact-pow2 amax inputs, k=-30..30):\n");
for (int idx = 0; idx < 61; idx++) {
int k = idx - 30;
float p2k = ldexpf(1.0f, k);
float p2k1 = ldexpf(1.0f, k - 1);
// expectations from exact real arithmetic; skip cases where the
// 1e-4 amax floor binds instead. `ulp` cases (one-ulp-off amax)
// probe approximation sensitivity, not a re-encode reachable
// state -- tallied separately as INFO.
struct { const char *name; float got, want; int skip, ulp; } cases[8] = {
{"fp8 448*2^k ", out[idx*8+0], p2k, 448.0f*p2k < 1.0e-4f, 0},
{"fp8 224*2^k ", out[idx*8+1], p2k1, 224.0f*p2k < 1.0e-4f, 0},
{"fp8 448u ", out[idx*8+2], 2.0f*p2k, 448.0f*p2k < 1.0e-4f, 1},
{"fp8 448d ", out[idx*8+3], p2k, 448.0f*p2k < 1.0e-4f, 1},
{"fp4 6*2^k ", out[idx*8+4], p2k, 0, 0},
{"fp4 3*2^k ", out[idx*8+5], p2k1, 0, 0},
{"fp4 1.5*2^k ", out[idx*8+6], ldexpf(1.0f, k-2), 0, 0},
{"fp4 6u ", out[idx*8+7], 2.0f*p2k, 0, 1},
};
for (int c = 0; c < 8; c++) {
if (cases[c].skip) continue;
if (cases[c].got != cases[c].want) {
printf(" %s k=%3d %s got=%.9g want=%.9g\n",
cases[c].ulp ? "ulp-info " : "DEVIATION",
k, cases[c].name, cases[c].got, cases[c].want);
if (cases[c].ulp) ulps++; else devs++;
}
}
}
printf("boundary scan: exact-case deviations=%d%s, ulp-info=%d\n\n", devs,
devs ? " (MECHANISM CANDIDATE -- fast-math log2f/ceilf at pow2)" : "", ulps);
}
printf("PROTO-KV-REENC-IDEM: shipped kernels %s, fixed variants %s\n",
orig_fail ? "REPRODUCE the bug" : "PASS (bug not reproduced?!)",
fails == 0 ? "PASS" : "FAIL");
return fails != 0;
}
cuda/mmq/test/proto_m2_comp.cu
+537
View File
@@ -0,0 +1,537 @@
/* proto_m2_comp.cu — M2-Inc5 prototype: fused compressor pair+store.
* Replaces the per-compressor-event 5-launch chain
* matmul_f16_splitk(kv) + combine + matmul_f16_splitk(sc) + combine
* + compressor_store
* with ONE cooperative kernel, at the three decode widths
* out 1024 (primary ratio-4), out 512 (primary ratio-2),
* out 256 (indexer ratio-4, head_dim 128); in_dim 4096, ksplit 2/4/8.
*
* Gate: kv_cur/sc_cur/state_kv/state_score all BIT-IDENTICAL to the ref
* chain across widths x ape f16/f32 x pos0 sweep; capture==eager; timing.
*
* Bit-exactness: phase 1 replicates the ref split-K tile exactly like
* M2-Inc3 (proto_m2_router.cu), generalized to seg in {2048,1024,512}:
* a lane's A = seg/256 accumulators hold chunks lane+32a (identical values
* at identical lanes to ref warps a), each reduced by the verbatim
* warp_sum_f32; the ref's cross-warp shfl_down tree over [s0..s7,0-padded]
* imposes the parenthesization
* A=2: s0+s1
* A=4: (s0+s2)+(s1+s3)
* A=8: ((s0+s4)+(s2+s6))+((s1+s5)+(s3+s7))
* (identical up to x+0.0f==x on exact zeros). The tail combines partials
* in kseg order (verbatim combine kernel) and applies the verbatim store
* body (state copy + ape add) to the just-combined values.
*
* Build (GB10):
* nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_comp.cu -o proto_m2_comp
*/
#include <cooperative_groups.h>
#include <cuda_fp16.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cmath>
#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
#define N_IN 4096u
/* ------------------------------------------------------------------ */
/* Vendored production device functions + kernels (ds4_cuda.cu, verbatim) */
struct ds4_decode_scalars {
uint32_t pos0;
uint32_t raw_row;
uint32_t raw_start;
uint32_t n_raw;
uint32_t n_comp;
uint32_t emit_phase;
uint32_t comp_row;
uint32_t index_row;
uint32_t flags;
uint32_t token;
};
__device__ static float warp_sum_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v += __shfl_down_sync(0xffffffffu, v, offset);
}
return v;
}
__device__ static float block_reduce_sum_256(float v) {
v = warp_sum_f32(v);
__shared__ float s[8];
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
if (lane == 0u) s[warp] = v;
__syncthreads();
float r = 0.0f;
if (warp == 0u) {
r = (lane < 8u) ? s[lane] : 0.0f;
r += __shfl_down_sync(0xffffffffu, r, 4);
r += __shfl_down_sync(0xffffffffu, r, 2);
r += __shfl_down_sync(0xffffffffu, r, 1);
}
return r;
}
__global__ static void matmul_f16_splitk_kernel(
float *partial,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint32_t ksplit,
int vec_ok) {
const uint64_t row = (uint64_t)blockIdx.x;
const uint32_t kseg = blockIdx.y;
if (row >= out_dim) return;
uint64_t seg = (in_dim + (uint64_t)ksplit - 1u) / (uint64_t)ksplit;
seg = (seg + 7u) & ~(uint64_t)7u;
const uint64_t k0 = (uint64_t)kseg * seg;
if (k0 >= in_dim) {
if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = 0.0f;
return;
}
uint64_t k1 = k0 + seg;
if (k1 > in_dim) k1 = in_dim;
const __half *wr = w + row * in_dim;
const uint32_t tid = threadIdx.x;
const uint32_t nt = blockDim.x;
float sum = 0.0f;
if (vec_ok) {
const uint64_t vend = k0 + ((k1 - k0) & ~(uint64_t)7u);
for (uint64_t i = k0 + (uint64_t)tid * 8u; i < vend; i += (uint64_t)nt * 8u) {
const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
const float4 xa = *reinterpret_cast<const float4 *>(x + i);
const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
for (uint64_t i = vend + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
} else {
for (uint64_t i = k0 + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
}
sum = block_reduce_sum_256(sum);
if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = sum;
}
__global__ static void matmul_f16_splitk_combine_kernel(
float *out,
const float *partial,
uint64_t out_dim,
uint32_t ksplit) {
const uint64_t row = (uint64_t)blockIdx.x * (uint64_t)blockDim.x + (uint64_t)threadIdx.x;
if (row >= out_dim) return;
const float *p = partial + row * (uint64_t)ksplit;
float s = 0.0f;
for (uint32_t k = 0u; k < ksplit; k++) s += p[k];
out[row] = s;
}
__device__ static float model_scalar_dev(const void *base, uint64_t offset, uint32_t type, uint64_t idx) {
const char *p = (const char *)base + offset;
if (type == 1u) return __half2float(((const __half *)p)[idx]);
return ((const float *)p)[idx];
}
__global__ static void compressor_store_kernel(
const float *kv,
const float *sc,
float *state_kv,
float *state_score,
const void *model_map,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens,
const struct ds4_decode_scalars * __restrict__ s_override) {
if (s_override) {
pos0 = s_override->pos0;
}
uint32_t coff = ratio == 4u ? 2u : 1u;
uint32_t width = coff * head_dim;
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * width;
if (gid >= n) return;
uint32_t t = gid / width;
uint32_t j = gid - (uint64_t)t * width;
uint32_t pos_mod = (pos0 + t) % ratio;
uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
state_kv[(uint64_t)dst_row * width + j] = kv[(uint64_t)t * width + j];
state_score[(uint64_t)dst_row * width + j] =
sc[(uint64_t)t * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)pos_mod * width + j);
}
/* ------------------------------------------------------------------ */
/* FUSED: pair matmul + combine + store as one cooperative kernel.
* Template KS = ksplit in {2,4,8}; SEG = 4096/KS; A = SEG/256 accumulators
* per lane. */
template <uint32_t KS>
__global__ static void comp_pair_store_fused_kernel(
float *kv_cur, /* [width] */
float *sc_cur, /* [width] */
float *partials, /* [2*width*KS] dedicated scratch */
float *state_kv,
float *state_score,
const __half *w_kv, /* [width][4096] */
const __half *w_sc, /* [width][4096] */
const float *x, /* [4096] */
const void *ape_map,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t width, /* == coff*head_dim == out_dim */
uint32_t pos0,
const struct ds4_decode_scalars * __restrict__ s_override) {
constexpr uint32_t SEG = N_IN / KS;
constexpr uint32_t A = SEG / 256u;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
/* Phase 1: one warp per (matmul m, row, kseg) tile. */
const uint32_t tiles_per_m = width * KS;
const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
for (uint32_t t = blockIdx.x * (blockDim.x >> 5u) + warp;
t < 2u * tiles_per_m; t += nwarp) {
const uint32_t m = t < tiles_per_m ? 0u : 1u;
const uint32_t tt = m ? t - tiles_per_m : t;
const uint32_t row = tt / KS;
const uint32_t kseg = tt % KS;
const uint64_t k0 = (uint64_t)kseg * SEG;
const __half *wr = (m ? w_sc : w_kv) + (uint64_t)row * N_IN;
float s[A];
#pragma unroll
for (uint32_t a = 0; a < A; a++) {
const uint64_t i = k0 + ((uint64_t)lane + 32u * a) * 8u;
const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
const float4 xa = *reinterpret_cast<const float4 *>(x + i);
const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
float sum = 0.0f;
sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
s[a] = warp_sum_f32(sum);
}
/* ref block_reduce_sum_256 cross-warp tree over [s0..s(A-1), 0...]:
* shfl_down 4,2,1 -> fixed parenthesization (zero-adds are identity). */
float partial;
if (KS == 8u) { /* A == 2 */
partial = s[0] + s[1 % A];
} else if (KS == 4u) { /* A == 4 */
partial = (s[0] + s[2 % A]) + (s[1 % A] + s[3 % A]);
} else { /* KS == 2, A == 8 */
partial = ((s[0] + s[4 % A]) + (s[2 % A] + s[6 % A]))
+ ((s[1 % A] + s[5 % A]) + (s[3 % A] + s[7 % A]));
}
if (lane == 0u) partials[(uint64_t)m * tiles_per_m + (uint64_t)row * KS + kseg] = partial;
}
cooperative_groups::this_grid().sync();
/* Phase 2: combine (kseg ascending, verbatim) + verbatim store body. */
const uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t n2 = 2u * width;
if (gid >= n2) return;
const uint32_t m = gid < width ? 0u : 1u;
const uint32_t row = m ? gid - width : gid;
const float *p = partials + (uint64_t)m * tiles_per_m + (uint64_t)row * KS;
float s = 0.0f;
for (uint32_t k = 0u; k < KS; k++) s += p[k];
uint32_t p0 = s_override ? s_override->pos0 : pos0;
const uint32_t pos_mod = p0 % ratio;
const uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
if (m == 0u) {
kv_cur[row] = s;
state_kv[(uint64_t)dst_row * width + row] = s;
} else {
sc_cur[row] = s;
state_score[(uint64_t)dst_row * width + row] =
s + model_scalar_dev(ape_map, ape_offset, ape_type, (uint64_t)pos_mod * width + row);
}
}
/* ------------------------------------------------------------------ */
/* Host harness */
struct cfg_t {
const char *name;
uint32_t head_dim, ratio, width, ksplit;
};
static const cfg_t CFGS[3] = {
{ "primary-r4 (1024,k2)", 512, 4, 1024, 2 },
{ "primary-r2 (512,k4)", 512, 2, 512, 4 },
{ "index-r4 (256,k8)", 128, 4, 256, 8 },
};
static void fill_rand(float *p, size_t n, unsigned seed, float scale) {
unsigned s = seed;
for (size_t i = 0; i < n; i++) {
s = s * 1664525u + 1013904223u;
p[i] = ((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 2.0f * scale;
}
}
struct bufs {
__half *w_kv, *w_sc;
float *x, *kv_cur, *sc_cur, *state_kv, *state_score, *part_ref, *part_fu;
char *ape; /* f32 or f16 depending on ape_type */
ds4_decode_scalars *scalars;
};
static void *fused_ptr(uint32_t ks) {
if (ks == 2u) return (void *)comp_pair_store_fused_kernel<2u>;
if (ks == 4u) return (void *)comp_pair_store_fused_kernel<4u>;
return (void *)comp_pair_store_fused_kernel<8u>;
}
static void run_ref(const bufs &b, const cfg_t &c, uint32_t ape_type, uint32_t pos0, int via_scalars) {
dim3 sg1(c.width, c.ksplit, 1);
matmul_f16_splitk_kernel<<<sg1, 256>>>(b.part_ref, b.w_kv, b.x, N_IN, c.width, c.ksplit, 1);
matmul_f16_splitk_combine_kernel<<<(c.width + 255u) / 256u, 256>>>(b.kv_cur, b.part_ref, c.width, c.ksplit);
matmul_f16_splitk_kernel<<<sg1, 256>>>(b.part_ref, b.w_sc, b.x, N_IN, c.width, c.ksplit, 1);
matmul_f16_splitk_combine_kernel<<<(c.width + 255u) / 256u, 256>>>(b.sc_cur, b.part_ref, c.width, c.ksplit);
compressor_store_kernel<<<(c.width + 255u) / 256u, 256>>>(
b.kv_cur, b.sc_cur, b.state_kv, b.state_score,
b.ape, 0, ape_type, c.head_dim, c.ratio,
via_scalars ? 0xdeadbeefu : pos0, 1,
via_scalars ? b.scalars : NULL);
CHECK(cudaGetLastError());
}
static int run_fused(const bufs &b, const cfg_t &c, uint32_t ape_type, uint32_t pos0, int via_scalars,
unsigned grid_blocks, cudaStream_t stream) {
float *kv_cur = b.kv_cur, *sc_cur = b.sc_cur, *partials = b.part_fu;
float *state_kv = b.state_kv, *state_score = b.state_score, *x = b.x;
const __half *w_kv = b.w_kv, *w_sc = b.w_sc;
const void *ape = b.ape;
uint64_t ape_off = 0;
uint32_t at = ape_type, hd = c.head_dim, ra = c.ratio, wi = c.width;
uint32_t p0 = via_scalars ? 0xdeadbeefu : pos0;
const ds4_decode_scalars *sc_ptr = via_scalars ? b.scalars : NULL;
void *args[] = {
(void *)&kv_cur, (void *)&sc_cur, (void *)&partials,
(void *)&state_kv, (void *)&state_score,
(void *)&w_kv, (void *)&w_sc, (void *)&x,
(void *)&ape, (void *)&ape_off, (void *)&at,
(void *)&hd, (void *)&ra, (void *)&wi, (void *)&p0, (void *)&sc_ptr };
cudaError_t err = cudaLaunchCooperativeKernel(
fused_ptr(c.ksplit), dim3(grid_blocks, 1, 1), dim3(256, 1, 1), args, 0, stream);
if (err != cudaSuccess) {
fprintf(stderr, "coop launch failed (%s, %u blk): %s\n", c.name, grid_blocks, cudaGetErrorString(err));
return 0;
}
return 1;
}
int main() {
int dev = 0, coop = 0, n_sm = 0;
CHECK(cudaGetDevice(&dev));
CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev));
CHECK(cudaDeviceGetAttribute(&n_sm, cudaDevAttrMultiProcessorCount, dev));
printf("coop=%d n_sm=%d\n", coop, n_sm);
if (!coop) return 1;
unsigned cap[3];
for (int i = 0; i < 3; i++) {
int bps = 0;
CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps, fused_ptr(CFGS[i].ksplit), 256, 0));
cap[i] = (unsigned)(bps * n_sm);
printf(" %s: blocks/sm=%d cap=%u\n", CFGS[i].name, bps, cap[i]);
}
const uint32_t WMAX = 1024;
bufs b = {};
CHECK(cudaMalloc(&b.w_kv, (size_t)WMAX * N_IN * sizeof(__half)));
CHECK(cudaMalloc(&b.w_sc, (size_t)WMAX * N_IN * sizeof(__half)));
CHECK(cudaMalloc(&b.x, N_IN * sizeof(float)));
CHECK(cudaMalloc(&b.kv_cur, WMAX * sizeof(float)));
CHECK(cudaMalloc(&b.sc_cur, WMAX * sizeof(float)));
CHECK(cudaMalloc(&b.state_kv, 8u * WMAX * sizeof(float)));
CHECK(cudaMalloc(&b.state_score, 8u * WMAX * sizeof(float)));
CHECK(cudaMalloc(&b.part_ref, (size_t)WMAX * 8u * sizeof(float)));
CHECK(cudaMalloc(&b.part_fu, 2u * (size_t)WMAX * 8u * sizeof(float)));
CHECK(cudaMalloc(&b.ape, 4u * WMAX * sizeof(float)));
CHECK(cudaMalloc(&b.scalars, sizeof(ds4_decode_scalars)));
float *hbuf = (float *)malloc((size_t)WMAX * N_IN * sizeof(float));
__half *hh = (__half *)malloc((size_t)WMAX * N_IN * sizeof(__half));
float *hape = (float *)malloc(4u * WMAX * sizeof(float));
__half *hape16 = (__half *)malloc(4u * WMAX * sizeof(__half));
const size_t OUTB = WMAX * sizeof(float);
const size_t STB = 8u * WMAX * sizeof(float);
float *ref_out = (float *)malloc(2 * OUTB + 2 * STB);
float *fu_out = (float *)malloc(2 * OUTB + 2 * STB);
int total_bad = 0, n_cases = 0;
const unsigned grids[] = { 48u, 96u, 128u };
const uint32_t poss[] = { 0u, 1u, 2u, 3u, 7u, 1023u };
for (int ci = 0; ci < 3; ci++) {
const cfg_t &c = CFGS[ci];
for (unsigned seed = 1; seed <= 3; seed++) {
fill_rand(hbuf, (size_t)c.width * N_IN, seed * 7919u + ci, 0.05f);
for (size_t i = 0; i < (size_t)c.width * N_IN; i++) hh[i] = __float2half(hbuf[i]);
CHECK(cudaMemcpy(b.w_kv, hh, (size_t)c.width * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
fill_rand(hbuf, (size_t)c.width * N_IN, seed * 104729u + ci, 0.05f);
for (size_t i = 0; i < (size_t)c.width * N_IN; i++) hh[i] = __float2half(hbuf[i]);
CHECK(cudaMemcpy(b.w_sc, hh, (size_t)c.width * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
fill_rand(hbuf, N_IN, seed * 31u, 1.0f);
CHECK(cudaMemcpy(b.x, hbuf, N_IN * sizeof(float), cudaMemcpyHostToDevice));
fill_rand(hape, (size_t)c.ratio * c.width, seed * 71u, 0.5f);
for (size_t i = 0; i < (size_t)c.ratio * c.width; i++) hape16[i] = __float2half(hape[i]);
for (uint32_t ape_type = 0; ape_type <= 1; ape_type++) {
if (ape_type == 0) CHECK(cudaMemcpy(b.ape, hape, (size_t)c.ratio * c.width * sizeof(float), cudaMemcpyHostToDevice));
else CHECK(cudaMemcpy(b.ape, hape16, (size_t)c.ratio * c.width * sizeof(__half), cudaMemcpyHostToDevice));
for (uint32_t pi = 0; pi < 6; pi++) {
const uint32_t pos0 = poss[pi];
const int via_scalars = (int)(pi & 1u);
ds4_decode_scalars sc = {};
sc.pos0 = pos0;
CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
/* ref */
CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
CHECK(cudaMemset(b.sc_cur, 0xCB, OUTB));
CHECK(cudaMemset(b.state_kv, 0xCB, STB));
CHECK(cudaMemset(b.state_score, 0xCB, STB));
run_ref(b, c, ape_type, pos0, via_scalars);
CHECK(cudaDeviceSynchronize());
CHECK(cudaMemcpy(ref_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(ref_out + WMAX, b.sc_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(ref_out + 2 * WMAX, b.state_kv, STB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(ref_out + 10 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
for (unsigned g : grids) {
if (g > cap[ci]) continue;
CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
CHECK(cudaMemset(b.sc_cur, 0xCB, OUTB));
CHECK(cudaMemset(b.state_kv, 0xCB, STB));
CHECK(cudaMemset(b.state_score, 0xCB, STB));
if (!run_fused(b, c, ape_type, pos0, via_scalars, g, 0)) return 1;
CHECK(cudaDeviceSynchronize());
CHECK(cudaMemcpy(fu_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(fu_out + WMAX, b.sc_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(fu_out + 2 * WMAX, b.state_kv, STB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(fu_out + 10 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
n_cases++;
if (memcmp(ref_out, fu_out, 2 * OUTB + 2 * STB)) {
total_bad++;
printf("MISMATCH %s seed%u ape%u pos%u g%u\n", c.name, seed, ape_type, pos0, g);
for (int i = 0; i < 6; i++)
printf(" kv[%d] ref=%.9g fu=%.9g sc ref=%.9g fu=%.9g\n", i,
ref_out[i], fu_out[i], ref_out[WMAX + i], fu_out[WMAX + i]);
}
}
}
}
}
}
printf("parity: %d/%d cases bit-identical\n", n_cases - total_bad, n_cases);
if (total_bad) return 1;
/* capture/replay bit-identity (one config) */
{
const cfg_t &c = CFGS[2];
const unsigned g = cap[2] < 128u ? cap[2] : 128u;
ds4_decode_scalars sc = {};
sc.pos0 = 5;
CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
if (!run_fused(b, c, 1, 0, 1, g, 0)) return 1;
CHECK(cudaDeviceSynchronize());
CHECK(cudaMemcpy(ref_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(ref_out + 2 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
cudaStream_t cs;
CHECK(cudaStreamCreate(&cs));
cudaGraph_t graph;
cudaGraphExec_t gexec;
CHECK(cudaStreamBeginCapture(cs, cudaStreamCaptureModeGlobal));
if (!run_fused(b, c, 1, 0, 1, g, cs)) return 1;
CHECK(cudaStreamEndCapture(cs, &graph));
CHECK(cudaGraphInstantiate(&gexec, graph, NULL, NULL, 0));
for (int r = 0; r < 3; r++) {
CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
CHECK(cudaGraphLaunch(gexec, cs));
CHECK(cudaStreamSynchronize(cs));
CHECK(cudaMemcpy(fu_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(fu_out + 2 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
if (memcmp(ref_out, fu_out, OUTB) || memcmp(ref_out + 2 * WMAX, fu_out + 2 * WMAX, STB)) {
printf("capture-replay MISMATCH\n");
return 1;
}
}
printf("capture/replay: bit-identical (3 replays)\n");
CHECK(cudaGraphExecDestroy(gexec));
CHECK(cudaGraphDestroy(graph));
CHECK(cudaStreamDestroy(cs));
}
/* timing: rotating 16-layer weight sets (L2-defeated) per config */
{
const int NL = 16, ITERS = 2000, WARM = 200;
cudaEvent_t e0, e1;
CHECK(cudaEventCreate(&e0));
CHECK(cudaEventCreate(&e1));
float ms;
for (int ci = 0; ci < 3; ci++) {
const cfg_t &c = CFGS[ci];
__half *wl[NL][2];
for (int l = 0; l < NL; l++)
for (int m = 0; m < 2; m++) {
CHECK(cudaMalloc(&wl[l][m], (size_t)c.width * N_IN * sizeof(__half)));
CHECK(cudaMemcpy(wl[l][m], m ? b.w_sc : b.w_kv,
(size_t)c.width * N_IN * sizeof(__half), cudaMemcpyDeviceToDevice));
}
bufs bl = b;
for (int it = 0; it < WARM; it++) {
bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
run_ref(bl, c, 1, 3, 1);
}
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(e0));
for (int it = 0; it < ITERS; it++) {
bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
run_ref(bl, c, 1, 3, 1);
}
CHECK(cudaEventRecord(e1));
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventElapsedTime(&ms, e0, e1));
printf("%s ref 5-launch : %8.3f us/event\n", c.name, ms * 1000.0f / ITERS);
for (unsigned g : grids) {
if (g > cap[ci]) continue;
for (int it = 0; it < WARM; it++) {
bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
run_fused(bl, c, 1, 3, 1, g, 0);
}
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(e0));
for (int it = 0; it < ITERS; it++) {
bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
run_fused(bl, c, 1, 3, 1, g, 0);
}
CHECK(cudaEventRecord(e1));
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventElapsedTime(&ms, e0, e1));
printf("%s fused g=%3u : %8.3f us/event\n", c.name, g, ms * 1000.0f / ITERS);
}
for (int l = 0; l < NL; l++)
for (int m = 0; m < 2; m++) CHECK(cudaFree(wl[l][m]));
}
}
printf("PROTO_M2_COMP OK\n");
return 0;
}
cuda/mmq/test/proto_m2_hc.cu
+1111
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/test/proto_m2_q2k.cu
+741
View File
@@ -0,0 +1,741 @@
// proto_m2_q2k.cu — M2 megakernel program, moe-down Q2K aligned-SoA proto.
//
// Target: mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> (cuda/mmq/mmvq.cu:608), the
// production decode down-proj kernel: 43 launches/step at 87.3 us = 3.755
// ms/step reading raw 84-byte block_q2_K stacks at ~190 GB/s vs the 216-233
// GB/s device substrate (nsys33d). This proto vendors the production kernel
// as the bit-exact reference and races two aligned-SoA layout twins that keep
// the SAME lane->(block,iqs) mapping and float accumulation order (layout-only
// change => bit-identical outputs, M1-Inc4 gate recipe):
//
// V1: plain SoA sections [half2 dm[nblk]] [pad64] [u8 scales[nblk*16]]
// [pad64] [u32 qs[nblk*16]], block order = raw tensor byte order.
// V2: row-pair SoA keyed to rows_per_block=2 (row pair = 2*blockIdx.x):
// [uint2 dm2[npair*nb]] [pad64] [int4 sc4[npair*nb*2]] [pad64]
// [uint2 qs2[npair*nb*16]]; per lane-iter one 8B qs load (both rows),
// one 16B scales load (both rows' 8B window), one 8B dm load.
//
// Parity: full-buffer bitwise compare vs the raw reference across seeds x id
// patterns (random experts, duplicates, -1 router ids, all -1) x ncols_dst
// {1, 6, 8}, plus graph capture/replay == eager for both variants.
//
// Timing: rotating L-layer weight pool (defeats L2 the way the 43-layer model
// does), production shape M=4096 K=2048 top-6, 6 distinct experts per launch.
//
// build (GB10): nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
// -o proto_m2_q2k proto_m2_q2k.cu
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <vector>
#include <random>
#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
// ---------------------------------------------------------------------------
// Vendored ggml pieces (verbatim math).
// ---------------------------------------------------------------------------
#define QK_K 256
#define QI2_K 16 // QK_K / (4*QR2_K)
#define QR2_K 4
#define QK8_1 32
#define QI8_1 8 // QK8_1 / (4*QR8_1)
typedef struct {
__half2 ds; // d (delta), s (d * sum of quants)
int8_t qs[QK8_1];
} block_q8_1;
static_assert(sizeof(block_q8_1) == 36, "block_q8_1 size");
typedef struct {
uint8_t scales[QK_K/16];
uint8_t qs[QK_K/4];
__half2 dm; // d, dmin
} block_q2_K;
static_assert(sizeof(block_q2_K) == 84, "block_q2_K size");
static __device__ __forceinline__ int get_int_b4(const void *x, const int &i32) {
return ((const int *)x)[i32]; // 4-byte aligned
}
static __device__ __forceinline__ int dp4a_(const int a, const int b, int c) {
return __dp4a(a, b, c);
}
template<int width = 32>
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x += __shfl_xor_sync(0xffffffff, x, offset, width);
}
return x;
}
// vec_dot_q2_K_q8_1_impl_mmvq, verbatim (vecdotq.cuh:364).
static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales,
const __half2 & dm2, const float * __restrict__ d8) {
float sumf_d = 0.0f;
float sumf_m = 0.0f;
#pragma unroll
for (int i = 0; i < QR2_K; ++i) {
const int sc = scales[2*i];
const int vi = (v >> (2*i)) & 0x03030303;
sumf_d += d8[i] * (dp4a_(vi, u[i], 0) * (sc & 0xF));
int m = sc >> 4;
m |= m << 8;
m |= m << 16;
sumf_m += d8[i] * dp4a_(m, u[i], 0);
}
const float2 dm2f = __half22float2(dm2);
return dm2f.x*sumf_d - dm2f.y*sumf_m;
}
// Same impl but the four scale bytes come from two pre-loaded 32-bit window
// words (V2's 16B scales load). sc value chain identical: byte (lo + 2*i) of
// the 8-byte window == scales[scale_offset + 2*i] of the raw block.
static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq_w(
const int & v, const int * __restrict__ u, const uint32_t w0, const uint32_t w1,
const int lo, const __half2 & dm2, const float * __restrict__ d8) {
float sumf_d = 0.0f;
float sumf_m = 0.0f;
#pragma unroll
for (int i = 0; i < QR2_K; ++i) {
const int bidx = lo + 2*i; // 0..7 within the window
const uint32_t w = (bidx < 4) ? w0 : w1;
const int sc = (int)((w >> ((bidx & 3) * 8)) & 0xFFu);
const int vi = (v >> (2*i)) & 0x03030303;
sumf_d += d8[i] * (dp4a_(vi, u[i], 0) * (sc & 0xF));
int m = sc >> 4;
m |= m << 8;
m |= m << 16;
sumf_m += d8[i] * dp4a_(m, u[i], 0);
}
const float2 dm2f = __half22float2(dm2);
return dm2f.x*sumf_d - dm2f.y*sumf_m;
}
// vec_dot_q2_K_q8_1, verbatim (vecdotq.cuh:814).
static __device__ __forceinline__ float vec_dot_q2_K_q8_1(
const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx;
const int bq8_offset = QR2_K * (iqs / QI8_1);
const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
const uint8_t * scales = bq2_K->scales + scale_offset;
const int v = get_int_b4(bq2_K->qs, iqs);
int u[QR2_K];
float d8[QR2_K];
#pragma unroll
for (int i = 0; i < QR2_K; ++ i) {
u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
}
return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8);
}
// ---------------------------------------------------------------------------
// Reference: mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> (mmvq.cu:608), verbatim up
// to Q2_K constant folding (qk 256, qi 16, vdr 1) and fastmodulo -> % (index
// math only). GB10 (sm_121 >= Ada): launch bounds MMVQ_MAX_BATCH_SIZE(8)*32.
// ---------------------------------------------------------------------------
template <int c_rows_per_block>
__launch_bounds__(8*32, 1)
static __global__ void ref_mul_mat_vec_q_moe(
const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids,
float * __restrict__ dst,
const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
const uint32_t ncols_dst, const uint32_t ids_stride) {
constexpr int qk = QK_K;
constexpr int qi = QI2_K;
constexpr int vdr = 1; // VDR_Q2_K_Q8_1_MMVQ
constexpr int warp_size = 32;
const uint32_t token_idx = threadIdx.y;
const int row0 = c_rows_per_block*blockIdx.x;
const int blocks_per_row_x = ncols_x / qk;
constexpr int blocks_per_iter = vdr * warp_size / qi; // 2
const uint32_t channel_dst = blockIdx.y;
if (token_idx >= ncols_dst) {
return;
}
const int32_t id_raw = ids[channel_dst + token_idx * ids_stride];
const bool invalid_id = id_raw < 0;
const uint32_t channel_x = invalid_id ? 0u : (uint32_t)id_raw;
const uint32_t channel_y = channel_dst % nchannels_y;
const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
const int kbx_offset = channel_x*stride_channel_x + row0*stride_row_x;
float tmp[c_rows_per_block] = {0.0f};
for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
const int kqs = vdr * (threadIdx.x % (qi/vdr));
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] += vec_dot_q2_K_q8_1(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
}
}
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
}
if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
}
}
// ---------------------------------------------------------------------------
// V1: plain SoA twin. Same skeleton, same lane mapping, same accumulation
// order; only the three weight streams are re-sourced from aligned sections.
// ---------------------------------------------------------------------------
template <int c_rows_per_block>
__launch_bounds__(8*32, 1)
static __global__ void v1_soa_moe_kernel(
const __half2 * __restrict__ dm_soa, // [nblk]
const uint8_t * __restrict__ scales_soa, // [nblk*16]
const int * __restrict__ qs_soa, // [nblk*16]
const void * __restrict__ vy, const int32_t * __restrict__ ids,
float * __restrict__ dst,
const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
const uint32_t ncols_dst, const uint32_t ids_stride) {
constexpr int qk = QK_K;
constexpr int qi = QI2_K;
constexpr int vdr = 1;
constexpr int warp_size = 32;
const uint32_t token_idx = threadIdx.y;
const int row0 = c_rows_per_block*blockIdx.x;
const int blocks_per_row_x = ncols_x / qk;
constexpr int blocks_per_iter = vdr * warp_size / qi;
const uint32_t channel_dst = blockIdx.y;
if (token_idx >= ncols_dst) {
return;
}
const int32_t id_raw = ids[channel_dst + token_idx * ids_stride];
const bool invalid_id = id_raw < 0;
const uint32_t channel_x = invalid_id ? 0u : (uint32_t)id_raw;
const uint32_t channel_y = channel_dst % nchannels_y;
const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
const int kbx_offset = channel_x*stride_channel_x + row0*stride_row_x;
float tmp[c_rows_per_block] = {0.0f};
for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
const int kqs = vdr * (threadIdx.x % (qi/vdr));
const int iqs = kqs;
const int bq8_offset = QR2_K * (iqs / QI8_1);
const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
const block_q8_1 * bq8_1 = &y[kby];
int u[QR2_K];
float d8[QR2_K];
#pragma unroll
for (int i = 0; i < QR2_K; ++i) {
u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
}
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
const size_t blk = (size_t)(kbx_offset + i*stride_row_x + kbx);
const int v = qs_soa[blk*16u + (unsigned)iqs];
tmp[i] += vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales_soa + blk*16u + (unsigned)scale_offset,
dm_soa[blk], d8);
}
}
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
}
if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
}
}
// ---------------------------------------------------------------------------
// V2: row-pair SoA twin (rows_per_block == 2 baked into the layout).
// Per lane-iter: one uint2 qs load (rows 0/1), one int4 scales-window load
// (both rows' 8B window for this lane's iqs half), one uint2 dm load. Scale
// BYTES, qs ints, dm halves and the float chain are value-identical to the
// reference; only load shapes change.
// ---------------------------------------------------------------------------
__launch_bounds__(8*32, 1)
static __global__ void v2_pair_soa_moe_kernel(
const uint2 * __restrict__ dm2_soa, // [npair*nb] {row0 half2, row1 half2}
const int4 * __restrict__ sc4_soa, // [npair*nb*2] {row0 w0,w1, row1 w0,w1} per half
const uint2 * __restrict__ qs2_soa, // [npair*nb*16] {row0 int, row1 int}
const void * __restrict__ vy, const int32_t * __restrict__ ids,
float * __restrict__ dst,
const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
const uint32_t ncols_dst, const uint32_t ids_stride) {
constexpr int qk = QK_K;
constexpr int qi = QI2_K;
constexpr int vdr = 1;
constexpr int warp_size = 32;
const uint32_t token_idx = threadIdx.y;
const int row0 = 2*blockIdx.x;
const int blocks_per_row_x = ncols_x / qk;
constexpr int blocks_per_iter = vdr * warp_size / qi;
const uint32_t channel_dst = blockIdx.y;
if (token_idx >= ncols_dst) {
return;
}
const int32_t id_raw = ids[channel_dst + token_idx * ids_stride];
const bool invalid_id = id_raw < 0;
const uint32_t channel_x = invalid_id ? 0u : (uint32_t)id_raw;
const uint32_t channel_y = channel_dst % nchannels_y;
const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
// pair-block base for (expert, row pair): pairs are laid out expert-major,
// pair-major, block-minor -- same linear order as raw blocks per pair.
const size_t pair_base = ((size_t)channel_x * (nrows_x/2u) + (size_t)blockIdx.x) * (size_t)blocks_per_row_x;
float tmp[2] = {0.0f, 0.0f};
for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
const int kqs = vdr * (threadIdx.x % (qi/vdr));
const int iqs = kqs;
const int bq8_offset = QR2_K * (iqs / QI8_1);
const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
const int whalf = iqs / QI8_1; // 0: window at scales[0..7], 1: at [8..15]
const int lo = scale_offset - 8*whalf; // 0 or 1 within the window
const block_q8_1 * bq8_1 = &y[kby];
int u[QR2_K];
float d8[QR2_K];
#pragma unroll
for (int i = 0; i < QR2_K; ++i) {
u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
}
const size_t pblk = pair_base + (size_t)kbx;
const uint2 v2 = qs2_soa[pblk*16u + (unsigned)iqs];
const uint2 dmw = dm2_soa[pblk];
const int4 scw = sc4_soa[pblk*2u + (unsigned)whalf];
const __half2 dm0 = *(const __half2 *)&dmw.x;
const __half2 dm1 = *(const __half2 *)&dmw.y;
tmp[0] += vec_dot_q2_K_q8_1_impl_mmvq_w((int)v2.x, u, (uint32_t)scw.x, (uint32_t)scw.y, lo, dm0, d8);
tmp[1] += vec_dot_q2_K_q8_1_impl_mmvq_w((int)v2.y, u, (uint32_t)scw.z, (uint32_t)scw.w, lo, dm1, d8);
}
#pragma unroll
for (int i = 0; i < 2; ++i) {
tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
}
if (threadIdx.x < 2 && uint32_t(row0 + threadIdx.x) < nrows_x) {
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
}
}
// ---------------------------------------------------------------------------
// Host helpers.
// ---------------------------------------------------------------------------
static const int M = 4096; // out rows (down out_dim)
static const int K = 2048; // in cols (expert_mid_dim)
static const int NB = K / QK_K; // 8 superblocks per row
static const int NYB = K / QK8_1; // 64 q8_1 blocks per activation row
static const int NEXP = 256;
static const int NUSED = 6;
static size_t nblk_layer() { return (size_t)NEXP * M * NB; }
static void fill_raw_layer(std::vector<block_q2_K> &w, std::mt19937 &rng) {
std::uniform_int_distribution<int> byte(0, 255);
std::uniform_real_distribution<float> mag(0.001f, 0.05f);
for (block_q2_K &b : w) {
for (int i = 0; i < 16; i++) b.scales[i] = (uint8_t)byte(rng);
for (int i = 0; i < 64; i++) b.qs[i] = (uint8_t)byte(rng);
b.dm = __floats2half2_rn(mag(rng), mag(rng));
}
}
static void quantize_q8_1(const float *x, block_q8_1 *y, int n) {
for (int ib = 0; ib < n / QK8_1; ib++) {
const float *xb = x + ib*QK8_1;
float amax = 0.0f;
for (int j = 0; j < QK8_1; j++) { float a = fabsf(xb[j]); if (a > amax) amax = a; }
const float d = amax / 127.0f;
const float idf = d != 0.0f ? 1.0f/d : 0.0f;
int sum = 0;
for (int j = 0; j < QK8_1; j++) {
int q = (int)roundf(xb[j] * idf);
if (q > 127) q = 127;
if (q < -127) q = -127;
y[ib].qs[j] = (int8_t)q;
sum += q;
}
y[ib].ds = __floats2half2_rn(d, d * (float)sum);
}
}
// V1 repack: [dm nblk*4][pad64][scales nblk*16][pad64][qs nblk*64]
static void repack_v1(const block_q2_K *raw, size_t nblk,
__half2 *dm, uint8_t *scales, int *qs) {
for (size_t b = 0; b < nblk; b++) {
dm[b] = raw[b].dm;
memcpy(scales + b*16u, raw[b].scales, 16);
memcpy(qs + b*16u, raw[b].qs, 64);
}
}
// V2 repack: pair p covers raw blocks of rows (2p, 2p+1) within an expert.
// Raw linear block index of (expert e, row r, block b) = (e*M + r)*NB + b.
static void repack_v2(const block_q2_K *raw,
uint2 *dm2, int4 *sc4, uint2 *qs2) {
const size_t npair_blk = (size_t)NEXP * (M/2) * NB;
for (size_t pb = 0; pb < npair_blk; pb++) {
const size_t b = pb % NB;
const size_t pr = (pb / NB) % (M/2);
const size_t e = pb / ((size_t)NB * (M/2));
const block_q2_K *r0 = raw + ((e*M + 2*pr + 0)*NB + b);
const block_q2_K *r1 = raw + ((e*M + 2*pr + 1)*NB + b);
uint2 dmw;
memcpy(&dmw.x, &r0->dm, 4);
memcpy(&dmw.y, &r1->dm, 4);
dm2[pb] = dmw;
// sc4[pb*2 + h] = {row0 w(h)[0:4], row0 w(h)[4:8], row1 w(h)[0:4], row1 w(h)[4:8]}
for (int h = 0; h < 2; h++) {
int4 s;
memcpy(&s.x, r0->scales + 8*h + 0, 4);
memcpy(&s.y, r0->scales + 8*h + 4, 4);
memcpy(&s.z, r1->scales + 8*h + 0, 4);
memcpy(&s.w, r1->scales + 8*h + 4, 4);
sc4[pb*2 + h] = s;
}
for (int i = 0; i < 16; i++) {
uint2 q;
memcpy(&q.x, r0->qs + 4*i, 4);
memcpy(&q.y, r1->qs + 4*i, 4);
qs2[pb*16 + i] = q;
}
}
}
struct dev_layer {
block_q2_K *raw = nullptr;
__half2 *v1_dm = nullptr; uint8_t *v1_sc = nullptr; int *v1_qs = nullptr;
uint2 *v2_dm = nullptr; int4 *v2_sc = nullptr; uint2 *v2_qs = nullptr;
};
// Production launch config for the down leg (routed_moe_launch vec tier):
// ncols_dst = n_assignments, nchannels_dst = 1, ids_stride = 1,
// stride_col_dst = M, stride_channel_dst = M, stride_col_y = NYB.
static void launch_ref(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
float *dst, int ncols, cudaStream_t s) {
dim3 grid(M/2, 1);
dim3 block(32, ncols);
ref_mul_mat_vec_q_moe<2><<<grid, block, 0, s>>>(
L.raw, y, ids, dst,
(uint32_t)K, 1u, (uint32_t)M,
(uint32_t)NB, (uint32_t)NYB, (uint32_t)M,
(uint32_t)(M*NB), (uint32_t)NYB, (uint32_t)M,
(uint32_t)ncols, 1u);
}
static void launch_v1(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
float *dst, int ncols, cudaStream_t s) {
dim3 grid(M/2, 1);
dim3 block(32, ncols);
v1_soa_moe_kernel<2><<<grid, block, 0, s>>>(
L.v1_dm, L.v1_sc, L.v1_qs, y, ids, dst,
(uint32_t)K, 1u, (uint32_t)M,
(uint32_t)NB, (uint32_t)NYB, (uint32_t)M,
(uint32_t)(M*NB), (uint32_t)NYB, (uint32_t)M,
(uint32_t)ncols, 1u);
}
static void launch_v2(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
float *dst, int ncols, cudaStream_t s) {
dim3 grid(M/2, 1);
dim3 block(32, ncols);
v2_pair_soa_moe_kernel<<<grid, block, 0, s>>>(
L.v2_dm, L.v2_sc, L.v2_qs, y, ids, dst,
(uint32_t)K, 1u, (uint32_t)M,
(uint32_t)NYB, (uint32_t)M,
(uint32_t)NYB, (uint32_t)M,
(uint32_t)ncols, 1u);
}
int main(int argc, char **argv) {
int n_layers = 6;
int timing_iters = 300;
if (argc > 1) n_layers = atoi(argv[1]);
if (argc > 2) timing_iters = atoi(argv[2]);
const size_t nblk = nblk_layer();
const size_t raw_bytes = nblk * sizeof(block_q2_K);
printf("proto_m2_q2k: M=%d K=%d E=%d nb=%d layers=%d raw=%.1f MiB/layer\n",
M, K, NEXP, NB, n_layers, raw_bytes / 1048576.0);
std::mt19937 rng(0xC0FFEE);
// Host-side raw layer template + repacks (one host copy reused per layer
// with a per-layer byte perturbation so layers differ without n_layers
// full random generations).
std::vector<block_q2_K> hraw(nblk);
fill_raw_layer(hraw, rng);
std::vector<__half2> h_v1dm(nblk);
std::vector<uint8_t> h_v1sc(nblk * 16u);
std::vector<int> h_v1qs(nblk * 16u);
std::vector<uint2> h_v2dm(nblk / 2u);
std::vector<int4> h_v2sc(nblk); // nblk/2 pairs * 2
std::vector<uint2> h_v2qs((nblk / 2u) * 16u);
std::vector<dev_layer> layers(n_layers);
for (int l = 0; l < n_layers; l++) {
if (l > 0) {
// cheap per-layer variation: rotate scales/qs bytes by l
for (size_t b = 0; b < nblk; b += 977u) {
hraw[b].qs[(l*7) & 63] ^= (uint8_t)(0x5A + l);
hraw[b].scales[(l*3) & 15] ^= (uint8_t)(0xA5 - l);
}
}
repack_v1(hraw.data(), nblk, h_v1dm.data(), h_v1sc.data(), h_v1qs.data());
repack_v2(hraw.data(), h_v2dm.data(), h_v2sc.data(), h_v2qs.data());
dev_layer &L = layers[l];
CHECK(cudaMalloc(&L.raw, raw_bytes));
CHECK(cudaMalloc(&L.v1_dm, nblk * 4u));
CHECK(cudaMalloc(&L.v1_sc, nblk * 16u));
CHECK(cudaMalloc(&L.v1_qs, nblk * 64u));
CHECK(cudaMalloc(&L.v2_dm, (nblk/2u) * 8u));
CHECK(cudaMalloc(&L.v2_sc, nblk * 16u));
CHECK(cudaMalloc(&L.v2_qs, (nblk/2u) * 16u * 8u));
CHECK(cudaMemcpy(L.raw, hraw.data(), raw_bytes, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v1_dm, h_v1dm.data(), nblk * 4u, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v1_sc, h_v1sc.data(), nblk * 16u, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v1_qs, h_v1qs.data(), nblk * 64u, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v2_dm, h_v2dm.data(), (nblk/2u) * 8u, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v2_sc, h_v2sc.data(), nblk * 16u, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(L.v2_qs, h_v2qs.data(), (nblk/2u) * 16u * 8u, cudaMemcpyHostToDevice));
}
// Activations: MAX_COLS q8_1 rows.
const int MAX_COLS = 8;
std::vector<float> hx(MAX_COLS * K);
std::vector<block_q8_1> hy(MAX_COLS * NYB);
std::uniform_real_distribution<float> xf(-1.0f, 1.0f);
for (float &v : hx) v = xf(rng);
for (int c = 0; c < MAX_COLS; c++) quantize_q8_1(hx.data() + c*K, hy.data() + c*NYB, K);
block_q8_1 *dy = nullptr;
CHECK(cudaMalloc(&dy, hy.size() * sizeof(block_q8_1)));
CHECK(cudaMemcpy(dy, hy.data(), hy.size() * sizeof(block_q8_1), cudaMemcpyHostToDevice));
int32_t *dids = nullptr;
CHECK(cudaMalloc(&dids, MAX_COLS * sizeof(int32_t)));
float *dref = nullptr, *dv1 = nullptr, *dv2 = nullptr;
const size_t out_bytes = (size_t)MAX_COLS * M * sizeof(float);
CHECK(cudaMalloc(&dref, out_bytes));
CHECK(cudaMalloc(&dv1, out_bytes));
CHECK(cudaMalloc(&dv2, out_bytes));
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// ------------------------------------------------------------------
// Parity: seeds x id patterns x ncols on rotating layers.
// ------------------------------------------------------------------
std::vector<float> href(MAX_COLS * M), hv1(MAX_COLS * M), hv2(MAX_COLS * M);
const int ncols_cases[3] = {1, 6, 8};
int checks = 0, fails = 0;
std::uniform_int_distribution<int> expd(0, NEXP - 1);
for (int seed = 0; seed < 10; seed++) {
for (int pat = 0; pat < 4; pat++) {
for (int nc = 0; nc < 3; nc++) {
const int ncols = ncols_cases[nc];
const dev_layer &L = layers[(seed + pat + nc) % n_layers];
int32_t ids[MAX_COLS];
for (int i = 0; i < ncols; i++) {
switch (pat) {
case 0: ids[i] = expd(rng); break; // random
case 1: ids[i] = (i < 2) ? expd(rng) : ids[i % 2]; // duplicates
break;
case 2: ids[i] = (i % 3 == 1) ? -1 : expd(rng); break; // sparse -1
case 3: ids[i] = -1; break; // all -1
}
}
CHECK(cudaMemcpy(dids, ids, ncols * sizeof(int32_t), cudaMemcpyHostToDevice));
CHECK(cudaMemset(dref, 0xCB, out_bytes));
CHECK(cudaMemset(dv1, 0xCB, out_bytes));
CHECK(cudaMemset(dv2, 0xCB, out_bytes));
launch_ref(L, dy, dids, dref, ncols, stream);
launch_v1(L, dy, dids, dv1, ncols, stream);
launch_v2(L, dy, dids, dv2, ncols, stream);
CHECK(cudaStreamSynchronize(stream));
CHECK(cudaGetLastError());
const size_t cmp = (size_t)ncols * M * sizeof(float);
CHECK(cudaMemcpy(href.data(), dref, cmp, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(hv1.data(), dv1, cmp, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(hv2.data(), dv2, cmp, cudaMemcpyDeviceToHost));
checks += 2;
if (memcmp(href.data(), hv1.data(), cmp) != 0) {
fails++;
printf("PARITY FAIL V1 seed=%d pat=%d ncols=%d\n", seed, pat, ncols);
}
if (memcmp(href.data(), hv2.data(), cmp) != 0) {
fails++;
printf("PARITY FAIL V2 seed=%d pat=%d ncols=%d\n", seed, pat, ncols);
}
}
}
}
printf("parity: %d/%d bit-identical%s\n", checks - fails, checks, fails ? " <-- FAIL" : "");
// ------------------------------------------------------------------
// Capture/replay parity: capture each variant once, replay, compare to
// the eager output just validated.
// ------------------------------------------------------------------
{
int32_t ids[MAX_COLS];
for (int i = 0; i < NUSED; i++) ids[i] = expd(rng);
ids[3] = -1;
CHECK(cudaMemcpy(dids, ids, NUSED * sizeof(int32_t), cudaMemcpyHostToDevice));
const dev_layer &L = layers[0];
const size_t cmp = (size_t)NUSED * M * sizeof(float);
launch_ref(L, dy, dids, dref, NUSED, stream);
CHECK(cudaStreamSynchronize(stream));
CHECK(cudaMemcpy(href.data(), dref, cmp, cudaMemcpyDeviceToHost));
int replay_fails = 0;
for (int variant = 0; variant < 2; variant++) {
float *dout = variant == 0 ? dv1 : dv2;
CHECK(cudaMemset(dout, 0xCB, out_bytes));
cudaGraph_t graph;
cudaGraphExec_t exec;
CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
if (variant == 0) launch_v1(L, dy, dids, dout, NUSED, stream);
else launch_v2(L, dy, dids, dout, NUSED, stream);
CHECK(cudaStreamEndCapture(stream, &graph));
CHECK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0));
CHECK(cudaGraphLaunch(exec, stream));
CHECK(cudaStreamSynchronize(stream));
CHECK(cudaMemcpy(hv1.data(), dout, cmp, cudaMemcpyDeviceToHost));
if (memcmp(href.data(), hv1.data(), cmp) != 0) {
replay_fails++;
printf("REPLAY FAIL V%d\n", variant + 1);
}
CHECK(cudaGraphExecDestroy(exec));
CHECK(cudaGraphDestroy(graph));
}
printf("capture/replay: %d/2 bit-identical%s\n", 2 - replay_fails, replay_fails ? " <-- FAIL" : "");
fails += replay_fails;
}
// ------------------------------------------------------------------
// Timing: rotating layers, 6 distinct experts per launch (production
// decode shape), events around the whole loop.
// ------------------------------------------------------------------
{
// Pre-generate rotating id sets on device (one buffer per iteration
// slot would perturb the graph; production re-writes ids in place, so
// do the same: ids uploaded once per timing pass, kernels rotate
// layers only -- weight reads dominate).
int32_t ids[MAX_COLS];
bool used[NEXP] = {false};
int n = 0;
while (n < NUSED) {
int e = expd(rng);
if (used[e]) continue;
used[e] = true;
ids[n++] = e;
}
CHECK(cudaMemcpy(dids, ids, NUSED * sizeof(int32_t), cudaMemcpyHostToDevice));
const double gib = (double)NUSED * M * NB * 84.0; // distinct-expert weight bytes/launch
cudaEvent_t e0, e1;
CHECK(cudaEventCreate(&e0));
CHECK(cudaEventCreate(&e1));
for (int variant = 0; variant < 3; variant++) {
const char *name = variant == 0 ? "ref-raw" : (variant == 1 ? "v1-soa " : "v2-pair");
float *dout = variant == 0 ? dref : (variant == 1 ? dv1 : dv2);
// warm
for (int it = 0; it < 30; it++) {
const dev_layer &L = layers[it % n_layers];
if (variant == 0) launch_ref(L, dy, dids, dout, NUSED, stream);
else if (variant == 1) launch_v1(L, dy, dids, dout, NUSED, stream);
else launch_v2(L, dy, dids, dout, NUSED, stream);
}
CHECK(cudaStreamSynchronize(stream));
CHECK(cudaEventRecord(e0, stream));
for (int it = 0; it < timing_iters; it++) {
const dev_layer &L = layers[it % n_layers];
if (variant == 0) launch_ref(L, dy, dids, dout, NUSED, stream);
else if (variant == 1) launch_v1(L, dy, dids, dout, NUSED, stream);
else launch_v2(L, dy, dids, dout, NUSED, stream);
}
CHECK(cudaEventRecord(e1, stream));
CHECK(cudaStreamSynchronize(stream));
float ms = 0.0f;
CHECK(cudaEventElapsedTime(&ms, e0, e1));
const double us = ms * 1000.0 / timing_iters;
printf("%s: %8.2f us/launch %7.1f GB/s (x43 = %.3f ms/step)\n",
name, us, gib / (us * 1e-6) / 1e9, us * 43.0 / 1000.0);
}
CHECK(cudaEventDestroy(e0));
CHECK(cudaEventDestroy(e1));
}
printf(fails ? "RESULT: FAIL\n" : "RESULT: PASS\n");
return fails ? 1 : 0;
}
cuda/mmq/test/proto_m2_qkv.cu
+604
View File
@@ -0,0 +1,604 @@
/* proto_m2_qkv.cu — M2-Inc2 prototypes: QKV-post chain fusions.
*
* Three kill-switch-ready fusions over the decode QKV-post pool, each proven
* BIT-EXACT here against verbatim copies of the production kernels:
*
* 2a dsv4_qkv_rms_norm_rows + q8_1 emission of the q row (kills the q_b
* mmvq quantize prelude, quantize_q8_1 grid=4).
* 2b head_rms_norm + rope_tail(q) fused, device-scalars pos source (the
* existing fused kernel is dead code because it bakes host pos0 into
* captures; this is its capture-safe twin).
* 2c rope_tail(kv) + fp8_kv_quantize + store_raw_kv fused into ONE
* 8-block kernel: rope touches only the 64-float rotary tail, fp8
* covers only the 448 nope elems (disjoint), and the production fp8
* kernel's 5.9 us is 7 SEQUENTIAL 64-elem groups in a single 64-thread
* block — the groups are independent, so they become 7 parallel blocks
* (+1 rope/tail block). Raw-store f16 round-trip preserved.
*
* Build: nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_qkv.cu -o proto_qkv
*/
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cmath>
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
printf("CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
#define N_HEAD 64u
#define HEAD_DIM 512u
#define N_ROT 64u
#define Q_RANK 1024u
#define KV_DIM 512u
#define RAW_CAP 1024u
#define RMS_EPS 1e-6f
/* mirror of the production decode-scalars substrate (only pos0/raw_row read) */
struct ds4_decode_scalars {
uint32_t pos0, raw_row, raw_start, n_raw, n_comp, emit_phase,
comp_row, index_row, flags, token;
};
typedef struct { __half2 ds; int8_t qs[32]; } block_q8_1_t;
/* ============ verbatim production device helpers ============ */
__device__ static float rope_yarn_ramp_dev(float low, float high, int i0) {
float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
}
__device__ static float dsv4_e4m3fn_value_dev(int i) {
int exp = (i >> 3) & 15;
int mant = i & 7;
if (exp == 0) return (float)mant * 0.001953125f;
return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
}
__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
float sign = x < 0.0f ? -1.0f : 1.0f;
float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return sign * dsv4_e4m3fn_value_dev(best);
}
/* ============ verbatim production kernels (references) ============ */
__global__ static void dsv4_qkv_rms_norm_rows_kernel(
float *q_out, const float *q, const float *q_w, uint32_t q_n,
float *kv_out, const float *kv, const float *kv_w, uint32_t kv_n,
uint32_t rows, float eps) {
const uint32_t row = blockIdx.x;
const uint32_t which = blockIdx.y;
if (row >= rows || which > 1u) return;
const uint32_t n = which == 0u ? q_n : kv_n;
const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
const float *w = which == 0u ? q_w : kv_w;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
const float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)n + eps);
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale * w[i];
}
}
__global__ static void head_rms_norm_kernel(float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
uint32_t row = blockIdx.x;
if (row >= n_tok * n_head) return;
float *xr = x + (uint64_t)row * head_dim;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
float scale = rsqrtf(partial[0] / (float)head_dim + eps);
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) xr[i] *= scale;
}
__global__ static void rope_tail_scalars_kernel(
float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot,
const struct ds4_decode_scalars * __restrict__ scalars,
int32_t pos_offset, uint32_t pos_stride, uint32_t n_ctx_orig, int inverse,
float freq_base, float freq_scale, float ext_factor, float attn_factor,
float beta_fast, float beta_slow) {
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t pairs = n_tok * n_head * (n_rot / 2);
if (gid >= pairs) return;
uint32_t pair = gid % (n_rot / 2);
uint32_t tmp = gid / (n_rot / 2);
uint32_t h = tmp % n_head;
uint32_t t = tmp / n_head;
uint32_t n_nope = head_dim - n_rot;
uint32_t i = pair * 2;
const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
float theta_extrap = (float)(pos0 + t * pos_stride) * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float *tail = x + ((uint64_t)t * n_head + h) * head_dim + n_nope;
float x0 = tail[i];
float x1 = tail[i + 1];
tail[i] = x0 * c - x1 * s;
tail[i + 1] = x0 * s + x1 * c;
}
__global__ static void fp8_kv_quantize_kernel(
float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
uint32_t n_nope = head_dim - n_rot;
const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
float *xr = x + (uint64_t)row * head_dim;
unsigned char *codes_row = codes_base
? codes_base + (uint64_t)row * row_stride : (unsigned char *)0;
float *scale_row = scale_base
? scale_base + (uint64_t)row * scale_stride : (float *)0;
__shared__ float scratch[64];
for (uint32_t off = 0; off < n_nope; off += 64) {
float v = 0.0f;
if (off + tid < n_nope) v = xr[off + tid];
scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
__syncthreads();
}
float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
if (off + tid < n_nope) {
float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
xr[off + tid] = dsv4_e4m3fn_dequant_dev(clamp) * scale;
if (codes_row) codes_row[off + tid] = 0; /* unused in decode path */
}
if (scale_row && tid == 0) scale_row[off / 64u] = scale;
__syncthreads();
}
(void)row_stride;
}
__global__ static void store_raw_kv_batch_kernel(
float *raw, const float *kv,
uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim,
const int32_t * __restrict__ positions, const int32_t * __restrict__ seq_id,
const struct ds4_decode_scalars * __restrict__ s_override) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * head_dim;
if (gid >= n) return;
uint32_t d = gid % head_dim;
uint32_t t = gid / head_dim;
uint32_t row;
if (s_override != NULL && n_tokens == 1u) {
row = s_override->raw_row;
} else {
uint32_t pos = positions ? (uint32_t)positions[t] : pos0 + t;
uint32_t seq_base = seq_id ? (uint32_t)seq_id[t] * raw_cap : 0u;
row = seq_base + (pos % raw_cap);
}
raw[(uint64_t)row * head_dim + d] = __half2float(__float2half(kv[(uint64_t)t * head_dim + d]));
}
/* q8_1 reference (vendored quantize_q8_1 semantics, contiguous case) */
__global__ static void ref_quantize_q8_1_kernel(block_q8_1_t *y, const float *x, uint32_t n) {
const uint32_t i = blockIdx.x * 32u + threadIdx.x;
if (i >= n) return;
const float xi = x[i];
float amax = fabsf(xi);
float sum = xi;
for (uint32_t off = 16u; off > 0u; off >>= 1u) {
amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off, 32));
sum += __shfl_xor_sync(0xffffffffu, sum, off, 32);
}
const float d = amax / 127.0f;
const int8_t q = amax == 0.0f ? (int8_t)0 : (int8_t)roundf(xi / d);
y[blockIdx.x].qs[threadIdx.x] = q;
if (threadIdx.x == 0) y[blockIdx.x].ds = __floats2half2_rn(d, sum);
}
/* ============ 2a: qkv_rms_norm_rows + q8_1(q row) ============ */
__global__ static void dsv4_qkv_rms_norm_rows_q81_kernel(
float *q_out, const float *q, const float *q_w, uint32_t q_n,
float *kv_out, const float *kv, const float *kv_w, uint32_t kv_n,
uint32_t rows, float eps, block_q8_1_t *q81) {
const uint32_t row = blockIdx.x;
const uint32_t which = blockIdx.y;
if (row >= rows || which > 1u) return;
const uint32_t n = which == 0u ? q_n : kv_n;
const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
const float *w = which == 0u ? q_w : kv_w;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
const float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)n + eps);
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale * w[i];
}
/* q8_1 emission of the q row (row 0 decode): re-read the just-written
* orow values (identical bits) in warp-contiguous granules. */
if (which == 0u && q81 != NULL && row == 0u && (n & 31u) == 0u) {
__syncthreads();
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t nwarp = blockDim.x >> 5u;
for (uint32_t qb = warp; qb < n / 32u; qb += nwarp) {
const float v = orow[qb * 32u + lane];
float amax = fabsf(v);
float s = v;
for (uint32_t off = 16u; off > 0u; off >>= 1u) {
amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off, 32));
s += __shfl_xor_sync(0xffffffffu, s, off, 32);
}
const float d1 = amax / 127.0f;
const int8_t qv = amax == 0.0f ? (int8_t)0 : (int8_t)roundf(v / d1);
q81[qb].qs[lane] = qv;
if (lane == 0u) q81[qb].ds = __floats2half2_rn(d1, s);
}
}
}
/* ============ 2b: head_rms + rope_tail(q), scalars pos ============ */
__global__ static void head_rms_norm_rope_tail_scalars_kernel(
float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot,
const struct ds4_decode_scalars * __restrict__ scalars,
int32_t pos_offset, uint32_t pos_stride, uint32_t n_ctx_orig, int inverse,
float freq_base, float freq_scale, float ext_factor, float attn_factor,
float beta_fast, float beta_slow, float eps) {
uint32_t row = blockIdx.x;
if (row >= n_tok * n_head) return;
uint32_t t = row / n_head;
float *xr = x + (uint64_t)row * head_dim;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)head_dim + eps);
const uint32_t n_nope = head_dim - n_rot;
for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) {
xr[i] *= scale;
}
const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) {
uint32_t i = pair * 2u;
float theta_extrap = (float)(pos0 + t * pos_stride) * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float *tail = xr + n_nope;
float x0 = tail[i] * scale;
float x1 = tail[i + 1] * scale;
tail[i] = x0 * c - x1 * s;
tail[i + 1] = x0 * s + x1 * c;
}
}
/* ============ 2c: rope_tail(kv) + fp8 quantize + raw store ============ */
/* Grid: (n_nope/64 + 1) blocks x 64 threads. Blocks [0, n_nope/64) each own
* one independent fp8 group (identical 64-wide shared-tree reduce as the
* production kernel, so scales and rounded values are bit-exact) and store
* their post-quant region to the raw row. The last block ropes the rotary
* tail (32 pairs, verbatim rope_tail_scalars math at n_head=1/t=0) and
* stores the rotated tail. Disjoint regions -> no cross-block ordering. */
__global__ static void kv_rope_fp8_store_scalars_kernel(
float *kv, float *raw, uint32_t raw_cap, uint32_t head_dim, uint32_t n_rot,
const struct ds4_decode_scalars * __restrict__ scalars,
int32_t pos_offset, uint32_t n_ctx_orig, int inverse,
float freq_base, float freq_scale, float ext_factor, float attn_factor,
float beta_fast, float beta_slow) {
const uint32_t n_nope = head_dim - n_rot;
const uint32_t b = blockIdx.x;
const uint32_t tid = threadIdx.x;
const uint32_t row = scalars->raw_row;
float *raw_row = raw + (uint64_t)(row % raw_cap) * head_dim;
if (b < n_nope / 64u) {
/* fp8 group b: verbatim group math from fp8_kv_quantize_kernel */
const uint32_t off = b * 64u;
float v = 0.0f;
if (off + tid < n_nope) v = kv[off + tid];
__shared__ float scratch[64];
scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
__syncthreads();
}
float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
if (off + tid < n_nope) {
float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
const float qv = dsv4_e4m3fn_dequant_dev(clamp) * scale;
kv[off + tid] = qv;
raw_row[off + tid] = __half2float(__float2half(qv));
}
} else {
/* rope the rotary tail then store it */
const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
float *tail = kv + n_nope;
if (tid < n_rot / 2u) {
const uint32_t i = tid * 2u;
float theta_extrap = (float)pos0 * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float x0 = tail[i];
float x1 = tail[i + 1];
const float r0 = x0 * c - x1 * s;
const float r1 = x0 * s + x1 * c;
tail[i] = r0;
tail[i + 1] = r1;
raw_row[n_nope + i] = __half2float(__float2half(r0));
raw_row[n_nope + i + 1] = __half2float(__float2half(r1));
}
}
}
/* ============ harness ============ */
static float *dmalloc_f(uint64_t n) { float *p; CK(cudaMalloc(&p, n * sizeof(float))); return p; }
static void fill_rand(float *d, uint64_t n, unsigned seed) {
float *h = (float *)malloc(n * sizeof(float));
srand(seed);
for (uint64_t i = 0; i < n; i++) h[i] = ((float)rand() / RAND_MAX - 0.5f) * 4.0f;
CK(cudaMemcpy(d, h, n * sizeof(float), cudaMemcpyHostToDevice));
free(h);
}
static int cmp_dev(const char *what, const float *a, const float *b, uint64_t n) {
float *ha = (float *)malloc(n * sizeof(float)), *hb = (float *)malloc(n * sizeof(float));
CK(cudaMemcpy(ha, a, n * sizeof(float), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(hb, b, n * sizeof(float), cudaMemcpyDeviceToHost));
int bad = memcmp(ha, hb, n * sizeof(float)) != 0;
if (bad) {
uint64_t k = 0;
for (; k < n; k++) if (ha[k] != hb[k]) break;
printf(" %s: MISMATCH at %llu (%a vs %a)\n", what, (unsigned long long)k, ha[k], hb[k]);
} else printf(" %s: bit-identical (%llu floats)\n", what, (unsigned long long)n);
free(ha); free(hb);
return bad;
}
int main() {
cudaStream_t st;
CK(cudaStreamCreate(&st));
int total_bad = 0;
/* rope parameter sets: plain and yarn-extended */
struct RP { float fb, fs, ef, af, bf, bs; uint32_t ctxo; const char *name; };
RP rps[2] = {
{ 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 0u, "plain" },
{ 10000.0f, 0.25f, 1.0f, 1.0f, 32.0f, 1.0f, 4096u, "yarn" },
};
uint32_t poss[3] = { 0u, 999u, 32767u };
ds4_decode_scalars hs;
memset(&hs, 0, sizeof(hs));
ds4_decode_scalars *ds;
CK(cudaMalloc(&ds, sizeof(hs)));
/* ---- 2a ---- */
{
float *q_in = dmalloc_f(Q_RANK), *kv_in = dmalloc_f(KV_DIM);
float *qw = dmalloc_f(Q_RANK), *kvw = dmalloc_f(KV_DIM);
float *qo_r = dmalloc_f(Q_RANK), *kvo_r = dmalloc_f(KV_DIM);
float *qo_f = dmalloc_f(Q_RANK), *kvo_f = dmalloc_f(KV_DIM);
block_q8_1_t *q81_r, *q81_f;
CK(cudaMalloc(&q81_r, (Q_RANK / 32) * sizeof(block_q8_1_t)));
CK(cudaMalloc(&q81_f, (Q_RANK / 32) * sizeof(block_q8_1_t)));
fill_rand(q_in, Q_RANK, 11); fill_rand(kv_in, KV_DIM, 12);
fill_rand(qw, Q_RANK, 13); fill_rand(kvw, KV_DIM, 14);
dim3 g(1, 2, 1);
dsv4_qkv_rms_norm_rows_kernel<<<g, 256, 0, st>>>(qo_r, q_in, qw, Q_RANK, kvo_r, kv_in, kvw, KV_DIM, 1, RMS_EPS);
ref_quantize_q8_1_kernel<<<Q_RANK / 32, 32, 0, st>>>(q81_r, qo_r, Q_RANK);
dsv4_qkv_rms_norm_rows_q81_kernel<<<g, 256, 0, st>>>(qo_f, q_in, qw, Q_RANK, kvo_f, kv_in, kvw, KV_DIM, 1, RMS_EPS, q81_f);
CK(cudaStreamSynchronize(st));
printf("2a qkv_rms + q8_1 emission:\n");
total_bad += cmp_dev("q_out", qo_r, qo_f, Q_RANK);
total_bad += cmp_dev("kv_out", kvo_r, kvo_f, KV_DIM);
char *h1 = (char *)malloc((Q_RANK / 32) * sizeof(block_q8_1_t));
char *h2 = (char *)malloc((Q_RANK / 32) * sizeof(block_q8_1_t));
CK(cudaMemcpy(h1, q81_r, (Q_RANK / 32) * sizeof(block_q8_1_t), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(h2, q81_f, (Q_RANK / 32) * sizeof(block_q8_1_t), cudaMemcpyDeviceToHost));
int bad = memcmp(h1, h2, (Q_RANK / 32) * sizeof(block_q8_1_t)) != 0;
printf(" q8_1 codes: %s\n", bad ? "MISMATCH" : "bit-identical");
total_bad += bad;
free(h1); free(h2);
}
/* ---- 2b ---- */
for (int r = 0; r < 2; r++) for (int p = 0; p < 3; p++) {
RP rp = rps[r];
hs.pos0 = poss[p];
CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
const uint64_t qn = (uint64_t)N_HEAD * HEAD_DIM;
float *qa = dmalloc_f(qn), *qb = dmalloc_f(qn);
fill_rand(qa, qn, 100 + r * 10 + p);
CK(cudaMemcpy(qb, qa, qn * sizeof(float), cudaMemcpyDeviceToDevice));
head_rms_norm_kernel<<<N_HEAD, 256, 0, st>>>(qa, 1, N_HEAD, HEAD_DIM, RMS_EPS);
uint32_t pairs = N_HEAD * (N_ROT / 2);
rope_tail_scalars_kernel<<<(pairs + 255) / 256, 256, 0, st>>>(qa, 1, N_HEAD, HEAD_DIM, N_ROT,
ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
head_rms_norm_rope_tail_scalars_kernel<<<N_HEAD, 256, 0, st>>>(qb, 1, N_HEAD, HEAD_DIM, N_ROT,
ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs, RMS_EPS);
CK(cudaStreamSynchronize(st));
printf("2b head_rms+rope (%s, pos=%u):\n", rp.name, poss[p]);
total_bad += cmp_dev("q", qa, qb, qn);
CK(cudaFree(qa)); CK(cudaFree(qb));
}
/* ---- 2c ---- */
for (int r = 0; r < 2; r++) for (int p = 0; p < 3; p++) {
RP rp = rps[r];
hs.pos0 = poss[p];
hs.raw_row = poss[p] % RAW_CAP;
CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
float *kva = dmalloc_f(KV_DIM), *kvb = dmalloc_f(KV_DIM);
float *rawa = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM), *rawb = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM);
CK(cudaMemset(rawa, 0, (uint64_t)RAW_CAP * HEAD_DIM * sizeof(float)));
CK(cudaMemset(rawb, 0, (uint64_t)RAW_CAP * HEAD_DIM * sizeof(float)));
fill_rand(kva, KV_DIM, 200 + r * 10 + p);
CK(cudaMemcpy(kvb, kva, KV_DIM * sizeof(float), cudaMemcpyDeviceToDevice));
/* reference chain: rope(kv) -> fp8 -> store */
uint32_t pairs = 1 * (N_ROT / 2);
rope_tail_scalars_kernel<<<(pairs + 255) / 256, 256, 0, st>>>(kva, 1, 1, HEAD_DIM, N_ROT,
ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
fp8_kv_quantize_kernel<<<1, 64, 0, st>>>(kva, 1, HEAD_DIM, N_ROT, NULL, NULL);
store_raw_kv_batch_kernel<<<(HEAD_DIM + 255) / 256, 256, 0, st>>>(rawa, kva, RAW_CAP, 0, 1, HEAD_DIM, NULL, NULL, ds);
/* fused */
kv_rope_fp8_store_scalars_kernel<<<(HEAD_DIM - N_ROT) / 64 + 1, 64, 0, st>>>(kvb, rawb, RAW_CAP, HEAD_DIM, N_ROT,
ds, 0, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
CK(cudaStreamSynchronize(st));
printf("2c kv rope+fp8+store (%s, pos=%u, row=%u):\n", rp.name, poss[p], hs.raw_row);
total_bad += cmp_dev("kv", kva, kvb, KV_DIM);
total_bad += cmp_dev("raw row", rawa + (uint64_t)hs.raw_row * HEAD_DIM,
rawb + (uint64_t)hs.raw_row * HEAD_DIM, HEAD_DIM);
CK(cudaFree(kva)); CK(cudaFree(kvb)); CK(cudaFree(rawa)); CK(cudaFree(rawb));
}
/* ---- timing: reference chains vs fused, 43-layer sweeps ---- */
{
const int REPS = 500, LAYERS = 43;
hs.pos0 = 1000u; hs.raw_row = 1000u % RAW_CAP;
CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
RP rp = rps[1];
const uint64_t qn = (uint64_t)N_HEAD * HEAD_DIM;
float *q = dmalloc_f(qn), *kv = dmalloc_f(KV_DIM);
float *raw = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM);
float *qi = dmalloc_f(Q_RANK), *kvi = dmalloc_f(KV_DIM);
float *qw = dmalloc_f(Q_RANK), *kvw = dmalloc_f(KV_DIM);
float *qo = dmalloc_f(Q_RANK), *kvo = dmalloc_f(KV_DIM);
block_q8_1_t *q81; CK(cudaMalloc(&q81, (Q_RANK / 32) * sizeof(block_q8_1_t)));
fill_rand(q, qn, 31); fill_rand(kv, KV_DIM, 32);
fill_rand(qi, Q_RANK, 33); fill_rand(kvi, KV_DIM, 34);
fill_rand(qw, Q_RANK, 35); fill_rand(kvw, KV_DIM, 36);
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
uint32_t pairs_q = N_HEAD * (N_ROT / 2), pairs_kv = N_ROT / 2;
dim3 g2(1, 2, 1);
float ms;
/* reference chain */
for (int w = 0; w < 2; w++) {
if (w) CK(cudaEventRecord(e0, st));
for (int rr = 0; rr < (w ? REPS : 10); rr++) for (int l = 0; l < LAYERS; l++) {
dsv4_qkv_rms_norm_rows_kernel<<<g2, 256, 0, st>>>(qo, qi, qw, Q_RANK, kvo, kvi, kvw, KV_DIM, 1, RMS_EPS);
ref_quantize_q8_1_kernel<<<Q_RANK / 32, 32, 0, st>>>(q81, qo, Q_RANK);
head_rms_norm_kernel<<<N_HEAD, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, RMS_EPS);
rope_tail_scalars_kernel<<<(pairs_q + 255) / 256, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
rope_tail_scalars_kernel<<<1, 256, 0, st>>>(kv, 1, 1, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
fp8_kv_quantize_kernel<<<1, 64, 0, st>>>(kv, 1, HEAD_DIM, N_ROT, NULL, NULL);
store_raw_kv_batch_kernel<<<(HEAD_DIM + 255) / 256, 256, 0, st>>>(raw, kv, RAW_CAP, 0, 1, HEAD_DIM, NULL, NULL, ds);
(void)pairs_kv;
}
if (w) { CK(cudaEventRecord(e1, st)); CK(cudaStreamSynchronize(st)); }
else CK(cudaStreamSynchronize(st));
}
CK(cudaEventElapsedTime(&ms, e0, e1));
printf("timing reference chain (7 launches): %.2f us/layer (x43 = %.3f ms/step)\n",
1000.0f * ms / (REPS * LAYERS), ms / REPS);
/* fused chain */
for (int w = 0; w < 2; w++) {
if (w) CK(cudaEventRecord(e0, st));
for (int rr = 0; rr < (w ? REPS : 10); rr++) for (int l = 0; l < LAYERS; l++) {
dsv4_qkv_rms_norm_rows_q81_kernel<<<g2, 256, 0, st>>>(qo, qi, qw, Q_RANK, kvo, kvi, kvw, KV_DIM, 1, RMS_EPS, q81);
head_rms_norm_rope_tail_scalars_kernel<<<N_HEAD, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs, RMS_EPS);
kv_rope_fp8_store_scalars_kernel<<<(HEAD_DIM - N_ROT) / 64 + 1, 64, 0, st>>>(kv, raw, RAW_CAP, HEAD_DIM, N_ROT, ds, 0, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
}
if (w) { CK(cudaEventRecord(e1, st)); CK(cudaStreamSynchronize(st)); }
else CK(cudaStreamSynchronize(st));
}
CK(cudaEventElapsedTime(&ms, e0, e1));
printf("timing fused chain (3 launches): %.2f us/layer (x43 = %.3f ms/step)\n",
1000.0f * ms / (REPS * LAYERS), ms / REPS);
}
printf(total_bad ? "PROTO_M2_QKV FAILED (%d)\n" : "PROTO_M2_QKV_DONE all bit-exact\n", total_bad);
return total_bad != 0;
}
cuda/mmq/test/proto_m2_router.cu
+939
View File
@@ -0,0 +1,939 @@
/* proto_m2_router.cu — M2-Inc3 prototype: fused router
* (f16 split-K logits matmul + combine + top-6 select) as ONE cooperative
* kernel, replacing the production 3-kernel chain
* matmul_f16_splitk_kernel (grid 256x8, 256 thr)
* + matmul_f16_splitk_combine_kernel (1 blk, 256 thr)
* + router_select_warp_topk_kernel (1 blk, 32x4)
* at the router shape in_dim=4096, out_dim=256 (=n_expert), n_tok=1, ksplit=8.
*
* Gate: all four outputs (logits, probs, selected, weights) BIT-IDENTICAL to
* the reference chain, in topk-bias / no-bias / hash modes, incl. exact-tie
* logits; capture==eager bit-identity; then per-layer timing vs the chain.
*
* Bit-exactness argument for the matmul phase: the ref (row,kseg) block gives
* each thread tid<64 exactly one 8-half chunk (seg 512 = 64 chunks); its
* block_reduce_sum_256 reduces warp 0 (chunks 0..31) and warp 1 (chunks
* 32..63) with warp_sum_f32, then the cross-warp tree only adds those two
* warp sums and exact zeros. The fused warp-per-tile form computes the SAME
* chunk values at the SAME lanes, applies the SAME warp_sum_f32 trees, and
* adds s0+s1 -- identical up to x+0.0f==x (breaks only for a partial that is
* exactly -0.0f, i.e. an all-zero activation segment; unreachable for
* rms-normed activations, and select-invariant even then).
*
* Build (GB10):
* nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_router.cu -o proto_m2_router
*/
#include <cooperative_groups.h>
#include <cuda_fp16.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cmath>
#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
#define N_IN 4096u
#define N_EXPERT 256u
#define KSPLIT 8u
#define SEG 512u
#define N_TOP 6u
/* ------------------------------------------------------------------ */
/* Vendored production device functions + kernels (ds4_cuda.cu, verbatim) */
struct ds4_decode_scalars {
uint32_t pos0;
uint32_t raw_row;
uint32_t raw_start;
uint32_t n_raw;
uint32_t n_comp;
uint32_t emit_phase;
uint32_t comp_row;
uint32_t index_row;
uint32_t flags;
uint32_t token;
};
__device__ static float warp_sum_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v += __shfl_down_sync(0xffffffffu, v, offset);
}
return v;
}
__device__ static float block_reduce_sum_256(float v) {
v = warp_sum_f32(v);
__shared__ float s[8];
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
if (lane == 0u) s[warp] = v;
__syncthreads();
float r = 0.0f;
if (warp == 0u) {
r = (lane < 8u) ? s[lane] : 0.0f;
r += __shfl_down_sync(0xffffffffu, r, 4);
r += __shfl_down_sync(0xffffffffu, r, 2);
r += __shfl_down_sync(0xffffffffu, r, 1);
}
return r;
}
__global__ static void matmul_f16_splitk_kernel(
float *partial,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint32_t ksplit,
int vec_ok) {
const uint64_t row = (uint64_t)blockIdx.x;
const uint32_t kseg = blockIdx.y;
if (row >= out_dim) return;
uint64_t seg = (in_dim + (uint64_t)ksplit - 1u) / (uint64_t)ksplit;
seg = (seg + 7u) & ~(uint64_t)7u;
const uint64_t k0 = (uint64_t)kseg * seg;
if (k0 >= in_dim) {
if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = 0.0f;
return;
}
uint64_t k1 = k0 + seg;
if (k1 > in_dim) k1 = in_dim;
const __half *wr = w + row * in_dim;
const uint32_t tid = threadIdx.x;
const uint32_t nt = blockDim.x;
float sum = 0.0f;
if (vec_ok) {
const uint64_t vend = k0 + ((k1 - k0) & ~(uint64_t)7u);
for (uint64_t i = k0 + (uint64_t)tid * 8u; i < vend; i += (uint64_t)nt * 8u) {
const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
const float4 xa = *reinterpret_cast<const float4 *>(x + i);
const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
for (uint64_t i = vend + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
} else {
for (uint64_t i = k0 + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
}
sum = block_reduce_sum_256(sum);
if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = sum;
}
__global__ static void matmul_f16_splitk_combine_kernel(
float *out,
const float *partial,
uint64_t out_dim,
uint32_t ksplit) {
const uint64_t row = (uint64_t)blockIdx.x * (uint64_t)blockDim.x + (uint64_t)threadIdx.x;
if (row >= out_dim) return;
const float *p = partial + row * (uint64_t)ksplit;
float s = 0.0f;
for (uint32_t k = 0u; k < ksplit; k++) s += p[k];
out[row] = s;
}
__device__ static float softplus_dev(float x) {
if (x > 20.0f) return x;
if (x < -20.0f) return expf(x);
return log1pf(expf(x));
}
__device__ __forceinline__ static bool router_score_better(float av, uint32_t ai, float bv, uint32_t bi) {
return av > bv || (av == bv && ai < bi);
}
__global__ static void router_select_warp_topk_kernel(
int32_t *selected,
float *weights,
float *probs,
const float *bias,
const int32_t *hash,
const float *logits,
const int32_t *tokens,
int32_t token_scalar,
uint32_t hash_rows,
uint32_t n_tokens,
int has_bias,
int hash_mode,
const struct ds4_decode_scalars *scalars) {
const uint32_t lane = threadIdx.x;
const uint32_t row_in_block = threadIdx.y;
const uint32_t t = blockIdx.x * blockDim.y + row_in_block;
if (t >= n_tokens || lane >= 32u) return;
const float *log = logits + (uint64_t)t * 256u;
float *prob = probs + (uint64_t)t * 256u;
int32_t *sel = selected + (uint64_t)t * 6u;
float *w = weights + (uint64_t)t * 6u;
__shared__ float sprob[4][256];
float local_prob[8];
float local_score[8];
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float p = sqrtf(softplus_dev(log[e]));
local_prob[j] = p;
local_score[j] = p + (has_bias ? bias[e] : 0.0f);
sprob[row_in_block][e] = p;
prob[e] = p;
}
__syncwarp();
if (hash_mode) {
if (lane == 0) {
int32_t tok = tokens ? tokens[t]
: (scalars ? (int32_t)scalars->token : token_scalar);
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6u;
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
const int32_t e = row[j];
sel[j] = e;
const float v = (e >= 0 && e < 256) ? sprob[row_in_block][(uint32_t)e] : 0.0f;
w[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
}
return;
}
float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
#pragma unroll
for (uint32_t k = 0; k < 6u; k++) {
float best_score = -INFINITY;
float best_prob = 0.0f;
uint32_t best_idx = UINT32_MAX;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float s = local_score[j];
if (router_score_better(s, e, best_score, best_idx)) {
best_score = s;
best_prob = local_prob[j];
best_idx = e;
}
}
#pragma unroll
for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
if (router_score_better(other_score, other_idx, best_score, best_idx)) {
best_score = other_score;
best_prob = other_prob;
best_idx = other_idx;
}
}
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
if (e == best_idx) local_score[j] = -INFINITY;
}
if (lane == 0) {
out_idx[k] = best_idx;
out_prob[k] = best_prob;
}
}
if (lane == 0) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
sel[j] = (int32_t)out_idx[j];
w[j] = out_prob[j];
sum += out_prob[j];
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
}
}
/* ------------------------------------------------------------------ */
/* FUSED: one cooperative kernel for the whole router stage. */
__global__ static void router_fused_coop_kernel(
float *logits, /* [256] */
float *partials, /* [256*8] dedicated scratch */
int32_t *selected, /* [6] */
float *weights, /* [6] */
float *probs, /* [256] */
const __half *w, /* [256][4096] f16 */
const float *x, /* [4096] */
const float *bias, /* [256] or NULL */
const int32_t *hash, /* [hash_rows][6] or NULL */
const int32_t *tokens, /* decode: NULL */
int32_t token_scalar,
uint32_t hash_rows,
int has_bias,
int hash_mode,
const struct ds4_decode_scalars *scalars) {
__shared__ float sbias[256];
__shared__ float sprob[256];
__shared__ int32_t shash[6];
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
/* Phase 0 (block 0): prefetch the tail's small model-map reads so their
* latency hides behind the matmul phase. Verbatim copies -- the select
* below consumes identical values. */
if (blockIdx.x == 0u) {
if (has_bias) sbias[threadIdx.x] = bias[threadIdx.x];
if (hash_mode && threadIdx.x == 0u) {
int32_t tok = tokens ? tokens[0]
: (scalars ? (int32_t)scalars->token : token_scalar);
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6u;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) shash[j] = row[j];
}
}
/* Phase 1: split-K logits matmul, one warp per (row, kseg) tile.
* Chunk values, lane placement, and both warp_sum_f32 trees replicate the
* ref block's warps 0/1 exactly; s0+s1 matches its cross-warp tree, which
* otherwise only adds exact zeros. */
const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
for (uint32_t t = blockIdx.x * (blockDim.x >> 5u) + warp;
t < N_EXPERT * KSPLIT; t += nwarp) {
const uint32_t row = t >> 3u;
const uint32_t kseg = t & 7u;
const uint64_t k0 = (uint64_t)kseg * SEG;
const __half *wr = w + (uint64_t)row * N_IN;
float sum0 = 0.0f;
float sum1 = 0.0f;
{
const uint64_t i = k0 + (uint64_t)lane * 8u;
const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
const float4 xa = *reinterpret_cast<const float4 *>(x + i);
const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
{
const uint64_t i = k0 + ((uint64_t)lane + 32u) * 8u;
const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
const float4 xa = *reinterpret_cast<const float4 *>(x + i);
const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
const float s0 = warp_sum_f32(sum0);
const float s1 = warp_sum_f32(sum1);
if (lane == 0u) partials[(uint64_t)row * KSPLIT + kseg] = s0 + s1;
}
cooperative_groups::this_grid().sync();
if (blockIdx.x != 0u) return;
/* Phase 2: combine partials in fixed kseg order (verbatim combine kernel;
* out_dim 256 == one 256-thread block). */
{
const uint64_t row = (uint64_t)threadIdx.x;
const float *p = partials + row * KSPLIT;
float s = 0.0f;
for (uint32_t k = 0u; k < KSPLIT; k++) s += p[k];
logits[row] = s;
}
__syncthreads();
/* Phase 3: verbatim router_select_warp_topk_kernel warp body (t = 0,
* row_in_block = 0); bias/hash reads swapped for the prefetched copies. */
if (warp != 0u) return;
const float *log = logits;
float local_prob[8];
float local_score[8];
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float p = sqrtf(softplus_dev(log[e]));
local_prob[j] = p;
local_score[j] = p + (has_bias ? sbias[e] : 0.0f);
sprob[e] = p;
probs[e] = p;
}
__syncwarp();
if (hash_mode) {
if (lane == 0u) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
const int32_t e = shash[j];
selected[j] = e;
const float v = (e >= 0 && e < 256) ? sprob[(uint32_t)e] : 0.0f;
weights[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
}
return;
}
float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
#pragma unroll
for (uint32_t k = 0; k < 6u; k++) {
float best_score = -INFINITY;
float best_prob = 0.0f;
uint32_t best_idx = UINT32_MAX;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float s = local_score[j];
if (router_score_better(s, e, best_score, best_idx)) {
best_score = s;
best_prob = local_prob[j];
best_idx = e;
}
}
#pragma unroll
for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
if (router_score_better(other_score, other_idx, best_score, best_idx)) {
best_score = other_score;
best_prob = other_prob;
best_idx = other_idx;
}
}
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
if (e == best_idx) local_score[j] = -INFINITY;
}
if (lane == 0) {
out_idx[k] = best_idx;
out_prob[k] = best_prob;
}
}
if (lane == 0) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
selected[j] = (int32_t)out_idx[j];
weights[j] = out_prob[j];
sum += out_prob[j];
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
}
}
/* V2: identical numerics, but each warp owns TWO tiles per pass with all four
* DRAM-latency w-vector loads hoisted before any math, doubling per-warp
* memory-level parallelism (x reads stay L1-resident after warmup). */
__global__ static void router_fused_coop_kernel_v2(
float *logits,
float *partials,
int32_t *selected,
float *weights,
float *probs,
const __half *w,
const float *x,
const float *bias,
const int32_t *hash,
const int32_t *tokens,
int32_t token_scalar,
uint32_t hash_rows,
int has_bias,
int hash_mode,
const struct ds4_decode_scalars *scalars) {
__shared__ float sbias[256];
__shared__ float sprob[256];
__shared__ int32_t shash[6];
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
if (blockIdx.x == 0u) {
if (has_bias) sbias[threadIdx.x] = bias[threadIdx.x];
if (hash_mode && threadIdx.x == 0u) {
int32_t tok = tokens ? tokens[0]
: (scalars ? (int32_t)scalars->token : token_scalar);
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6u;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) shash[j] = row[j];
}
}
const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
const uint32_t gwarp = blockIdx.x * (blockDim.x >> 5u) + warp;
for (uint32_t t0 = gwarp; t0 < N_EXPERT * KSPLIT; t0 += 2u * nwarp) {
const uint32_t t1 = t0 + nwarp;
const int have1 = t1 < N_EXPERT * KSPLIT;
const uint64_t k0a = (uint64_t)(t0 & 7u) * SEG;
const __half *wra = w + (uint64_t)(t0 >> 3u) * N_IN;
const uint64_t ia0 = k0a + (uint64_t)lane * 8u;
const uint64_t ia1 = k0a + ((uint64_t)lane + 32u) * 8u;
const uint64_t k0b = have1 ? (uint64_t)(t1 & 7u) * SEG : k0a;
const __half *wrb = have1 ? w + (uint64_t)(t1 >> 3u) * N_IN : wra;
const uint64_t ib0 = k0b + (uint64_t)lane * 8u;
const uint64_t ib1 = k0b + ((uint64_t)lane + 32u) * 8u;
/* all four w vectors in flight before any dependent math */
const uint4 wva0 = *reinterpret_cast<const uint4 *>(wra + ia0);
const uint4 wva1 = *reinterpret_cast<const uint4 *>(wra + ia1);
const uint4 wvb0 = *reinterpret_cast<const uint4 *>(wrb + ib0);
const uint4 wvb1 = *reinterpret_cast<const uint4 *>(wrb + ib1);
float s0, s1;
{
float sum0 = 0.0f;
float sum1 = 0.0f;
{
const float4 xa = *reinterpret_cast<const float4 *>(x + ia0);
const float4 xb = *reinterpret_cast<const float4 *>(x + ia0 + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wva0);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
{
const float4 xa = *reinterpret_cast<const float4 *>(x + ia1);
const float4 xb = *reinterpret_cast<const float4 *>(x + ia1 + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wva1);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
s0 = warp_sum_f32(sum0);
s1 = warp_sum_f32(sum1);
}
float s2, s3;
if (have1) {
float sum0 = 0.0f;
float sum1 = 0.0f;
{
const float4 xa = *reinterpret_cast<const float4 *>(x + ib0);
const float4 xb = *reinterpret_cast<const float4 *>(x + ib0 + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wvb0);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
{
const float4 xa = *reinterpret_cast<const float4 *>(x + ib1);
const float4 xb = *reinterpret_cast<const float4 *>(x + ib1 + 4);
const __half2 *h = reinterpret_cast<const __half2 *>(&wvb1);
const float2 w0 = __half22float2(h[0]);
const float2 w1 = __half22float2(h[1]);
const float2 w2 = __half22float2(h[2]);
const float2 w3 = __half22float2(h[3]);
sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+ w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
}
s2 = warp_sum_f32(sum0);
s3 = warp_sum_f32(sum1);
}
if (lane == 0u) {
partials[(uint64_t)(t0 >> 3u) * KSPLIT + (t0 & 7u)] = s0 + s1;
if (have1) partials[(uint64_t)(t1 >> 3u) * KSPLIT + (t1 & 7u)] = s2 + s3;
}
}
cooperative_groups::this_grid().sync();
if (blockIdx.x != 0u) return;
{
const uint64_t row = (uint64_t)threadIdx.x;
const float *p = partials + row * KSPLIT;
float s = 0.0f;
for (uint32_t k = 0u; k < KSPLIT; k++) s += p[k];
logits[row] = s;
}
__syncthreads();
if (warp != 0u) return;
const float *log = logits;
float local_prob[8];
float local_score[8];
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float p = sqrtf(softplus_dev(log[e]));
local_prob[j] = p;
local_score[j] = p + (has_bias ? sbias[e] : 0.0f);
sprob[e] = p;
probs[e] = p;
}
__syncwarp();
if (hash_mode) {
if (lane == 0u) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
const int32_t e = shash[j];
selected[j] = e;
const float v = (e >= 0 && e < 256) ? sprob[(uint32_t)e] : 0.0f;
weights[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
}
return;
}
float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
#pragma unroll
for (uint32_t k = 0; k < 6u; k++) {
float best_score = -INFINITY;
float best_prob = 0.0f;
uint32_t best_idx = UINT32_MAX;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
const float s = local_score[j];
if (router_score_better(s, e, best_score, best_idx)) {
best_score = s;
best_prob = local_prob[j];
best_idx = e;
}
}
#pragma unroll
for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
if (router_score_better(other_score, other_idx, best_score, best_idx)) {
best_score = other_score;
best_prob = other_prob;
best_idx = other_idx;
}
}
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const uint32_t e = lane + j * 32u;
if (e == best_idx) local_score[j] = -INFINITY;
}
if (lane == 0) {
out_idx[k] = best_idx;
out_prob[k] = best_prob;
}
}
if (lane == 0) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
selected[j] = (int32_t)out_idx[j];
weights[j] = out_prob[j];
sum += out_prob[j];
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
}
}
/* ------------------------------------------------------------------ */
/* Host harness */
struct outs {
float logits[N_EXPERT];
float probs[N_EXPERT];
int32_t sel[N_TOP];
float wgt[N_TOP];
};
static void fill_rand(float *p, size_t n, unsigned seed, float scale) {
unsigned s = seed;
for (size_t i = 0; i < n; i++) {
s = s * 1664525u + 1013904223u;
p[i] = ((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 2.0f * scale;
}
}
struct dev_bufs {
__half *w;
float *x, *bias, *logits, *probs, *partials_ref, *partials_fu, *wgt;
int32_t *hash, *sel;
ds4_decode_scalars *scalars;
uint32_t hash_rows;
};
static void run_ref(const dev_bufs &b, int has_bias, int hash_mode, int32_t token_via_scalars) {
dim3 sgrid(N_EXPERT, KSPLIT, 1);
matmul_f16_splitk_kernel<<<sgrid, 256>>>(b.partials_ref, b.w, b.x, N_IN, N_EXPERT, KSPLIT, 1);
matmul_f16_splitk_combine_kernel<<<1, 256>>>(b.logits, b.partials_ref, N_EXPERT, KSPLIT);
dim3 block(32, 4, 1);
router_select_warp_topk_kernel<<<1, block>>>(b.sel, b.wgt, b.probs,
has_bias && !hash_mode ? b.bias : NULL,
hash_mode ? b.hash : NULL,
b.logits, NULL, /*token_scalar=*/-1, b.hash_rows, 1,
has_bias && !hash_mode, hash_mode,
token_via_scalars ? b.scalars : NULL);
CHECK(cudaGetLastError());
}
static int launch_fused_k(void *kernel, const dev_bufs &b, int has_bias, int hash_mode,
int32_t token_via_scalars, unsigned grid_blocks, cudaStream_t stream) {
float *logits = b.logits, *partials = b.partials_fu, *probs = b.probs, *wgt = b.wgt, *x = b.x;
const __half *w = b.w;
const float *bias = has_bias && !hash_mode ? b.bias : NULL;
const int32_t *hash = hash_mode ? b.hash : NULL;
int32_t *sel = b.sel;
const int32_t *tokens = NULL;
int32_t token_scalar = -1;
uint32_t hash_rows = b.hash_rows;
int hb = has_bias && !hash_mode;
int hm = hash_mode;
const ds4_decode_scalars *scalars = token_via_scalars ? b.scalars : NULL;
void *args[] = {
(void *)&logits, (void *)&partials, (void *)&sel, (void *)&wgt, (void *)&probs,
(void *)&w, (void *)&x, (void *)&bias, (void *)&hash, (void *)&tokens,
(void *)&token_scalar, (void *)&hash_rows, (void *)&hb, (void *)&hm, (void *)&scalars };
cudaError_t err = cudaLaunchCooperativeKernel(
kernel,
dim3(grid_blocks, 1, 1), dim3(256, 1, 1), args, 0, stream);
if (err != cudaSuccess) {
fprintf(stderr, "coop launch failed (%u blk): %s\n", grid_blocks, cudaGetErrorString(err));
return 0;
}
return 1;
}
static int launch_fused(const dev_bufs &b, int has_bias, int hash_mode, int32_t token_via_scalars,
unsigned grid_blocks, cudaStream_t stream) {
return launch_fused_k((void *)router_fused_coop_kernel, b, has_bias, hash_mode,
token_via_scalars, grid_blocks, stream);
}
static void read_outs(const dev_bufs &b, outs *o) {
CHECK(cudaMemcpy(o->logits, b.logits, sizeof(o->logits), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(o->probs, b.probs, sizeof(o->probs), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(o->sel, b.sel, sizeof(o->sel), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(o->wgt, b.wgt, sizeof(o->wgt), cudaMemcpyDeviceToHost));
}
static int cmp_outs(const outs &a, const outs &bb, const char *tag) {
int bad = 0;
if (memcmp(a.logits, bb.logits, sizeof(a.logits))) { printf(" [%s] logits MISMATCH\n", tag); bad = 1; }
if (memcmp(a.probs, bb.probs, sizeof(a.probs))) { printf(" [%s] probs MISMATCH\n", tag); bad = 1; }
if (memcmp(a.sel, bb.sel, sizeof(a.sel))) { printf(" [%s] selected MISMATCH\n", tag); bad = 1; }
if (memcmp(a.wgt, bb.wgt, sizeof(a.wgt))) { printf(" [%s] weights MISMATCH\n", tag); bad = 1; }
if (bad) {
for (int i = 0; i < 8; i++) printf(" logit[%d] ref=%.9g fused=%.9g\n", i, a.logits[i], bb.logits[i]);
for (int i = 0; i < 6; i++) printf(" sel[%d] ref=%d fused=%d w ref=%.9g fused=%.9g\n",
i, a.sel[i], bb.sel[i], a.wgt[i], bb.wgt[i]);
}
return bad;
}
int main() {
int dev = 0, coop = 0, n_sm = 0, bps = 0;
CHECK(cudaGetDevice(&dev));
CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev));
CHECK(cudaDeviceGetAttribute(&n_sm, cudaDevAttrMultiProcessorCount, dev));
CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps, router_fused_coop_kernel, 256, 0));
const unsigned max_blocks = (unsigned)(bps * n_sm);
printf("coop=%d n_sm=%d blocks/sm=%d max co-resident=%u\n", coop, n_sm, bps, max_blocks);
if (!coop) { printf("NO cooperative launch support\n"); return 1; }
dev_bufs b = {};
b.hash_rows = 1000;
CHECK(cudaMalloc(&b.w, (size_t)N_EXPERT * N_IN * sizeof(__half)));
CHECK(cudaMalloc(&b.x, N_IN * sizeof(float)));
CHECK(cudaMalloc(&b.bias, N_EXPERT * sizeof(float)));
CHECK(cudaMalloc(&b.logits, N_EXPERT * sizeof(float)));
CHECK(cudaMalloc(&b.probs, N_EXPERT * sizeof(float)));
CHECK(cudaMalloc(&b.partials_ref, N_EXPERT * KSPLIT * sizeof(float)));
CHECK(cudaMalloc(&b.partials_fu, N_EXPERT * KSPLIT * sizeof(float)));
CHECK(cudaMalloc(&b.sel, N_TOP * sizeof(int32_t)));
CHECK(cudaMalloc(&b.wgt, N_TOP * sizeof(float)));
CHECK(cudaMalloc(&b.hash, (size_t)b.hash_rows * N_TOP * sizeof(int32_t)));
CHECK(cudaMalloc(&b.scalars, sizeof(ds4_decode_scalars)));
float *hw = (float *)malloc((size_t)N_EXPERT * N_IN * sizeof(float));
__half *hwh = (__half *)malloc((size_t)N_EXPERT * N_IN * sizeof(__half));
float *hx = (float *)malloc(N_IN * sizeof(float));
float *hb = (float *)malloc(N_EXPERT * sizeof(float));
int32_t *hh = (int32_t *)malloc((size_t)b.hash_rows * N_TOP * sizeof(int32_t));
/* -------- parity sweep -------- */
int total_bad = 0, n_cases = 0;
const unsigned grids[] = { 48u, 96u, 128u, 192u, 256u };
for (unsigned seed = 1; seed <= 10; seed++) {
fill_rand(hw, (size_t)N_EXPERT * N_IN, seed * 7919u, 0.05f);
if (seed == 3) {
/* exact-tie case: duplicate rows -> equal logits (and equal bias) ->
* tie-break by lower index must survive fusion */
for (int r = 0; r < 6; r++)
memcpy(hw + (size_t)(100 + r) * N_IN, hw + (size_t)(10 + r) * N_IN, N_IN * sizeof(float));
}
for (size_t i = 0; i < (size_t)N_EXPERT * N_IN; i++) hwh[i] = __float2half(hw[i]);
fill_rand(hx, N_IN, seed * 104729u, 1.0f);
if (seed == 5) for (size_t i = 0; i < N_IN; i++) hx[i] = 0.0f; /* all-equal logits */
fill_rand(hb, N_EXPERT, seed * 31u, 0.3f);
if (seed == 3) for (int r = 0; r < 6; r++) hb[100 + r] = hb[10 + r];
unsigned s = seed * 2654435761u;
for (size_t i = 0; i < (size_t)b.hash_rows * N_TOP; i++) {
s = s * 1664525u + 1013904223u;
hh[i] = (int32_t)(s % 300u) - 20; /* includes out-of-range ids */
}
CHECK(cudaMemcpy(b.w, hwh, (size_t)N_EXPERT * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(b.x, hx, N_IN * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(b.bias, hb, N_EXPERT * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(b.hash, hh, (size_t)b.hash_rows * N_TOP * sizeof(int32_t), cudaMemcpyHostToDevice));
ds4_decode_scalars sc = {};
sc.token = (seed * 37u) % 1200u; /* sometimes >= hash_rows -> clamp path */
CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
struct { int has_bias, hash_mode, via_scalars; const char *name; } cases[] = {
{ 1, 0, 0, "bias" }, { 0, 0, 0, "nobias" }, { 0, 1, 1, "hash" },
};
for (auto &c : cases) {
outs oref, ofu;
run_ref(b, c.has_bias, c.hash_mode, c.via_scalars);
CHECK(cudaDeviceSynchronize());
read_outs(b, &oref);
for (unsigned g : grids) {
if (g > max_blocks) continue;
/* scrub outputs so a silent non-write is caught */
CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
CHECK(cudaMemset(b.probs, 0xCB, N_EXPERT * sizeof(float)));
CHECK(cudaMemset(b.sel, 0xCB, N_TOP * sizeof(int32_t)));
CHECK(cudaMemset(b.wgt, 0xCB, N_TOP * sizeof(float)));
if (!launch_fused(b, c.has_bias, c.hash_mode, c.via_scalars, g, 0)) return 1;
CHECK(cudaDeviceSynchronize());
read_outs(b, &ofu);
char tag[64];
snprintf(tag, sizeof(tag), "seed%u %s g%u", seed, c.name, g);
total_bad += cmp_outs(oref, ofu, tag);
n_cases++;
CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
CHECK(cudaMemset(b.probs, 0xCB, N_EXPERT * sizeof(float)));
CHECK(cudaMemset(b.sel, 0xCB, N_TOP * sizeof(int32_t)));
CHECK(cudaMemset(b.wgt, 0xCB, N_TOP * sizeof(float)));
if (!launch_fused_k((void *)router_fused_coop_kernel_v2,
b, c.has_bias, c.hash_mode, c.via_scalars, g, 0)) return 1;
CHECK(cudaDeviceSynchronize());
read_outs(b, &ofu);
snprintf(tag, sizeof(tag), "V2 seed%u %s g%u", seed, c.name, g);
total_bad += cmp_outs(oref, ofu, tag);
n_cases++;
}
}
}
printf("parity: %d/%d cases bit-identical\n", n_cases - total_bad, n_cases);
if (total_bad) return 1;
/* -------- capture/replay bit-identity -------- */
{
const unsigned g = max_blocks < 256u ? max_blocks : 256u;
outs oeager, oreplay;
if (!launch_fused(b, 1, 0, 0, g, 0)) return 1;
CHECK(cudaDeviceSynchronize());
read_outs(b, &oeager);
cudaStream_t cs;
CHECK(cudaStreamCreate(&cs));
cudaGraph_t graph;
cudaGraphExec_t gexec;
CHECK(cudaStreamBeginCapture(cs, cudaStreamCaptureModeGlobal));
if (!launch_fused(b, 1, 0, 0, g, cs)) return 1;
CHECK(cudaStreamEndCapture(cs, &graph));
CHECK(cudaGraphInstantiate(&gexec, graph, NULL, NULL, 0));
for (int r = 0; r < 3; r++) {
CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
CHECK(cudaGraphLaunch(gexec, cs));
CHECK(cudaStreamSynchronize(cs));
read_outs(b, &oreplay);
if (cmp_outs(oeager, oreplay, "capture-replay")) return 1;
}
printf("capture/replay: bit-identical (3 replays, %u blk)\n", g);
CHECK(cudaGraphExecDestroy(gexec));
CHECK(cudaGraphDestroy(graph));
CHECK(cudaStreamDestroy(cs));
}
/* -------- timing: rotating 68-layer weight set (L2-defeated) -------- */
{
const int NL = 68;
__half *wl[NL];
for (int l = 0; l < NL; l++) {
CHECK(cudaMalloc(&wl[l], (size_t)N_EXPERT * N_IN * sizeof(__half)));
CHECK(cudaMemcpy(wl[l], b.w, (size_t)N_EXPERT * N_IN * sizeof(__half), cudaMemcpyDeviceToDevice));
}
const int ITERS = 3000, WARM = 300;
cudaEvent_t e0, e1;
CHECK(cudaEventCreate(&e0));
CHECK(cudaEventCreate(&e1));
float ms;
/* ref chain */
dev_bufs bl = b;
for (int it = 0; it < WARM; it++) { bl.w = wl[it % NL]; run_ref(bl, 1, 0, 0); }
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(e0));
for (int it = 0; it < ITERS; it++) { bl.w = wl[it % NL]; run_ref(bl, 1, 0, 0); }
CHECK(cudaEventRecord(e1));
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventElapsedTime(&ms, e0, e1));
printf("ref chain : %8.3f us/layer\n", ms * 1000.0f / ITERS);
int bps2 = 0;
CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps2, router_fused_coop_kernel_v2, 256, 0));
const unsigned max_blocks2 = (unsigned)(bps2 * n_sm);
printf("V2 blocks/sm=%d max co-resident=%u\n", bps2, max_blocks2);
struct { void *k; const char *name; unsigned cap; } variants[] = {
{ (void *)router_fused_coop_kernel, "V1", max_blocks },
{ (void *)router_fused_coop_kernel_v2, "V2", max_blocks2 },
};
for (auto &v : variants) {
for (unsigned g : grids) {
if (g > v.cap) continue;
for (int it = 0; it < WARM; it++) { bl.w = wl[it % NL]; launch_fused_k(v.k, bl, 1, 0, 0, g, 0); }
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(e0));
for (int it = 0; it < ITERS; it++) { bl.w = wl[it % NL]; launch_fused_k(v.k, bl, 1, 0, 0, g, 0); }
CHECK(cudaEventRecord(e1));
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventElapsedTime(&ms, e0, e1));
printf("fused %s g=%3u : %8.3f us/layer\n", v.name, g, ms * 1000.0f / ITERS);
}
}
for (int l = 0; l < NL; l++) CHECK(cudaFree(wl[l]));
}
printf("PROTO_M2_ROUTER OK\n");
return 0;
}
cuda/mmq/test/proto_mm_ids.cu
+475
View File
@@ -0,0 +1,475 @@
/* proto_mm_ids.cu — mm_ids_helper parity harness.
*
* P4 Inc1 (landed 86c5d4d): mm_ids_helper<1> fast path for the routed-MoE
* down matmul (ds4_mmq.cu ds4_mmq_moe_impl: n_expert_used=1, si1=1, sis1=1,
* nchannels_y=1, n_experts=256). Without a case-1 dispatch the down matmul
* fell to the generic <0> template: one warp per expert scanning all
* assignment rows with a single active lane — 22.5 ms/launch at W4096.
*
* P5 (this extension): mm_ids_helper_global — the large-n variant with NO
* shared-memory staging. The smem kernel stages the compact (it, iex_used)
* list in n_tokens*4 B of dynamic shared memory, capping n_tokens at smpbo/4
* (~25k on GB10); the down matmul passes n_tokens = assignment rows = 6x the
* forward width, so 8192-row prefill chunks hit 48384 "tokens", the mmq
* entries refused, and the WHOLE MoE block fell back to the pre-mmq
* expert-tile kernels (the W8192 prefill cliff, trace cmtp5w8.sqlite). The
* global variant runs the same per-expert scan twice: pass 0 counts
* (nex_prev, bucket size), pass 1 re-scans and writes ids_src1/ids_dst
* directly at final offsets.
*
* Gate: ids_src1 / ids_dst / expert_bounds BIT-IDENTICAL across
* (a) smem <0> vs smem <1>/<6> (the landed Inc1 claim),
* (b) smem vs global at every sub-cap shape,
* (c) global vs a host reference at EVERY shape, including past-cap legs
* (48384/49152 assignments = the W8192 down shape; a >cap
* n_expert_used=6 leg; the cap boundary +/-1)
* with the production pre-zero of both id maps (dropped -1 router ids leave
* slots unwritten; expert_bounds[0] counts them, so the zeros sit below
* expert 0's bucket and are never consumed).
*
* Bit-exactness argument for the global variant: identical (it, iex_used)
* tuples in identical token order and identical output expressions; the
* (lossless) 22/10-bit store round-trip is simply removed. Both variants
* are undefined if one token lists the same expert twice (router top-k is
* without replacement); the generators here sample distinct experts.
*
* Build (GB10):
* nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_mm_ids.cu -o proto_mm_ids
*/
#include <cuda_runtime.h>
#include <climits>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cassert>
#include <vector>
#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
static constexpr int WARP_SIZE = 32;
/* ------------------------------------------------------------------ */
/* Vendored from cuda/mmq/mmid.cu + common.cuh (verbatim semantics) */
template <int width = WARP_SIZE>
static __device__ __forceinline__ int warp_reduce_any(int x) {
if (width == WARP_SIZE) {
return __any_sync(0xffffffff, x);
} else {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x = __shfl_xor_sync(0xffffffff, x, offset, width) || x;
}
return x;
}
}
template <int width = WARP_SIZE>
static __device__ __forceinline__ int warp_reduce_sum(int x) {
#pragma unroll
for (int offset = width/2; offset > 0; offset >>= 1) {
x += __shfl_xor_sync(0xffffffff, x, offset, width);
}
return x;
}
struct mm_ids_helper_store {
uint32_t data;
__device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
data = (it & 0x003FFFFF) | (iex_used << 22);
}
__device__ uint32_t it() const { return data & 0x003FFFFF; }
__device__ uint32_t iex_used() const { return data >> 22; }
};
static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store");
template <int n_expert_used_template>
__launch_bounds__(WARP_SIZE, 1)
static __global__ void mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
constexpr int warp_size = WARP_SIZE;
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
extern __shared__ char data_mm_ids_helper[];
mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
int nex_prev = 0;
int it_compact = 0;
if constexpr (n_expert_used_template == 0) {
for (int it = 0; it < n_tokens; ++it) {
int iex_used = -1;
for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
const int expert_used = ids[it*si1 + iex];
nex_prev += expert_used < expert;
if (expert_used == expert) {
iex_used = iex;
}
}
if (iex_used != -1) {
store[it_compact] = mm_ids_helper_store(it, iex_used);
}
if (warp_reduce_any<warp_size>(iex_used != -1)) {
it_compact++;
}
}
} else {
static_assert(n_expert_used_template == 6 || warp_size % n_expert_used_template == 0, "bad n_expert_used");
constexpr int neu_padded = n_expert_used_template == 6 ? 8 : n_expert_used_template;
for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
const int it = it0 + threadIdx.x / neu_padded;
const int iex = threadIdx.x % neu_padded;
const int expert_used = (neu_padded == n_expert_used_template || iex < n_expert_used_template) && it < n_tokens ?
ids[it*si1 + iex] : INT_MAX;
const int iex_used = expert_used == expert ? iex : -1;
nex_prev += expert_used < expert;
const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
int it_compact_add_lower = 0;
#pragma unroll
for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
if (threadIdx.x >= static_cast<unsigned int>(offset)) {
it_compact_add_lower += tmp;
}
}
if (iex_used != -1) {
store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
}
it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
}
}
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
__syncwarp();
for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
const mm_ids_helper_store store_it = store[itc];
const int it = store_it.it();
const int iex_used = store_it.iex_used();
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
}
if (threadIdx.x != 0) {
return;
}
expert_bounds[expert] = nex_prev;
if (expert < static_cast<int>(gridDim.x) - 1) {
return;
}
expert_bounds[gridDim.x] = nex_prev + it_compact;
}
/* Vendored mm_ids_helper_global (P5 large-n variant, mmid.cu) */
template <int n_expert_used_template>
__launch_bounds__(WARP_SIZE, 1)
static __global__ void mm_ids_helper_global(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
constexpr int warp_size = WARP_SIZE;
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
int nex_prev = 0;
int it_compact_count = 0;
#pragma unroll 1
for (int pass = 0; pass < 2; ++pass) {
int it_compact = 0;
if constexpr (n_expert_used_template == 0) {
for (int it = 0; it < n_tokens; ++it) {
int iex_used = -1;
for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
const int expert_used = ids[it*si1 + iex];
if (pass == 0) {
nex_prev += expert_used < expert;
}
if (expert_used == expert) {
iex_used = iex;
}
}
if (pass == 1 && iex_used != -1) {
ids_src1[nex_prev + it_compact] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + it_compact] = it*n_expert_used + iex_used;
}
if (warp_reduce_any<warp_size>(iex_used != -1)) {
it_compact++;
}
}
} else {
static_assert(n_expert_used_template == 6 || warp_size % n_expert_used_template == 0, "bad n_expert_used");
constexpr int neu_padded = n_expert_used_template == 6 ? 8 : n_expert_used_template;
for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
const int it = it0 + threadIdx.x / neu_padded;
const int iex = threadIdx.x % neu_padded;
const int expert_used = (neu_padded == n_expert_used_template || iex < n_expert_used_template) && it < n_tokens ?
ids[it*si1 + iex] : INT_MAX;
const int iex_used = expert_used == expert ? iex : -1;
if (pass == 0) {
nex_prev += expert_used < expert;
}
const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
int it_compact_add_lower = 0;
#pragma unroll
for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
if (threadIdx.x >= static_cast<unsigned int>(offset)) {
it_compact_add_lower += tmp;
}
}
if (pass == 1 && iex_used != -1) {
const int itc = it_compact + it_compact_add_lower;
ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
}
it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
}
}
if (pass == 0) {
it_compact_count = it_compact;
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
}
}
if (threadIdx.x != 0) {
return;
}
expert_bounds[expert] = nex_prev;
if (expert < static_cast<int>(gridDim.x) - 1) {
return;
}
expert_bounds[gridDim.x] = nex_prev + it_compact_count;
}
/* NOTE: the proto instantiations replace the vendored kernels' runtime-const
* neu_padded with a constexpr derived from the template arg (identical value
* per instantiation; the vendored code relies on the same constant folding). */
/* ------------------------------------------------------------------ */
static constexpr int N_EXPERTS = 256;
/* Host reference implementing the mm_ids_helper output contract:
* expert_bounds[e] = #{assignment values < e} (signed: -1 counts for all
* experts, so invalid rows occupy never-consumed slots below expert 0);
* bucket e = token-ascending (it, iex) pairs with ids[it*si1+iex] == e;
* ids_src1[pos] = it*sis1 + iex % nchannels_y; ids_dst[pos] = it*neu + iex.
* src1/dst must be pre-zeroed by the caller (production memset semantics). */
static void cpu_ref(const int32_t *ids, int nt, int neu, int nchannels_y, int si1, int sis1,
int32_t *src1, int32_t *dst, int32_t *eb) {
std::vector<int32_t> h(N_EXPERTS + 2, 0); // h[b+1] counts value b, b in [-1, N_EXPERTS]
for (int it = 0; it < nt; ++it) {
for (int iex = 0; iex < neu; ++iex) {
int v = ids[it*si1 + iex];
if (v < -1) v = -1;
if (v > N_EXPERTS) v = N_EXPERTS;
h[v + 1]++;
}
}
std::vector<int32_t> below(N_EXPERTS + 1, 0); // below[e] = #{v < e}
int32_t acc = h[0];
for (int e = 0; e <= N_EXPERTS; ++e) {
below[e] = acc;
if (e < N_EXPERTS) acc += h[e + 1];
}
for (int e = 0; e <= N_EXPERTS; ++e) {
eb[e] = below[e];
}
std::vector<int32_t> fill(N_EXPERTS, 0);
for (int it = 0; it < nt; ++it) {
for (int iex = 0; iex < neu; ++iex) {
const int e = ids[it*si1 + iex];
if (e < 0 || e >= N_EXPERTS) continue;
const int64_t pos = (int64_t)below[e] + fill[e]++;
src1[pos] = it*sis1 + iex % nchannels_y;
dst [pos] = it*neu + iex;
}
}
}
struct leg {
const char *name;
int n_tokens; // assignment rows (neu=1) or tokens (neu=6)
int neu; // n_expert_used: 1 = down shape, 6 = gate/up shape
int invalid_pct; // % of slots with -1 (router NaN path)
bool skew; // hot-expert distribution
};
int main() {
int smpbo = 0;
CHECK(cudaDeviceGetAttribute(&smpbo, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0));
CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<0>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<1>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<6>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
const int cap_tokens = smpbo / 4;
printf("smpbo=%d (cap %d tokens)\n", smpbo, cap_tokens);
std::vector<leg> legs = {
// down shape (neu=1): n_tokens = rows*6 assignments
{"W4096x6dn uniform", 4096*6, 1, 0, false},
{"W4096x6dn 2%invalid", 4096*6, 1, 2, false},
{"W4096x6dn skewed", 4096*6, 1, 1, true},
{"W2048x6dn uniform", 2048*6, 1, 2, false},
{"decode w1x6", 6, 1, 0, false},
{"decode w8x6", 48, 1, 0, false},
// past-cap down shapes (the P5 cliff: W8192 -> 8064*6 = 48384 traced, 8192*6 = 49152 max)
{"W8192x6dn uniform", 48384, 1, 0, false},
{"W8192x6dn 2%invalid", 48384, 1, 2, false},
{"W8192x6dn skewed", 48384, 1, 1, true},
{"W8192x6dn max", 49152, 1, 2, false},
// gate/up shape (neu=6): n_tokens = rows, si1 = 6, distinct experts/token
{"W8192gu uniform", 8064, 6, 0, false},
{"W8192gu 2%invalid", 8064, 6, 2, false},
{"gu past-cap", 30000, 6, 2, false},
};
// cap-boundary seam legs (runtime-sized)
leg seam_at = {"cap-boundary", cap_tokens, 1, 2, false};
leg seam_over = {"cap-boundary+1", cap_tokens + 1, 1, 2, false};
legs.push_back(seam_at);
legs.push_back(seam_over);
int failures = 0;
for (const leg &L : legs) {
const int nt = L.n_tokens;
const int neu = L.neu;
const int si1 = neu, sis1 = 1, nchannels_y = 1;
const int64_t n_ids = (int64_t)nt * neu;
const size_t smem = (size_t)nt * 4;
const bool smem_fits = smem <= (size_t)smpbo;
int32_t *h_ids = (int32_t*)malloc(n_ids * sizeof(int32_t));
srand(20260708);
if (neu == 1) {
for (int i = 0; i < nt; i++) {
if (L.invalid_pct && rand() % 100 < L.invalid_pct) { h_ids[i] = -1; continue; }
h_ids[i] = L.skew ? (rand() % 100 < 60 ? rand() % 8 : rand() % N_EXPERTS)
: rand() % N_EXPERTS;
}
} else {
// router top-k semantics: distinct experts per token
for (int it = 0; it < nt; it++) {
int chosen[8]; int c = 0;
while (c < neu) {
const int e = L.skew ? (rand() % 100 < 60 ? rand() % 8 : rand() % N_EXPERTS)
: rand() % N_EXPERTS;
bool dup = false;
for (int j = 0; j < c; j++) dup = dup || chosen[j] == e;
if (!dup) chosen[c++] = e;
}
for (int iex = 0; iex < neu; iex++) {
h_ids[(int64_t)it*neu + iex] =
(L.invalid_pct && rand() % 100 < L.invalid_pct) ? -1 : chosen[iex];
}
}
}
// host reference
int32_t *r_src1 = (int32_t*)calloc(n_ids, sizeof(int32_t));
int32_t *r_dst = (int32_t*)calloc(n_ids, sizeof(int32_t));
int32_t r_eb[N_EXPERTS+1];
cpu_ref(h_ids, nt, neu, nchannels_y, si1, sis1, r_src1, r_dst, r_eb);
int32_t *d_ids, *d_src1, *d_dst, *d_eb;
CHECK(cudaMalloc(&d_ids, n_ids * sizeof(int32_t)));
CHECK(cudaMalloc(&d_src1, n_ids * sizeof(int32_t)));
CHECK(cudaMalloc(&d_dst, n_ids * sizeof(int32_t)));
CHECK(cudaMalloc(&d_eb, (N_EXPERTS+1) * sizeof(int32_t)));
CHECK(cudaMemcpy(d_ids, h_ids, n_ids * sizeof(int32_t), cudaMemcpyHostToDevice));
const dim3 grid(N_EXPERTS, 1, 1), block(WARP_SIZE, 1, 1);
int32_t *h_src1 = (int32_t*)malloc(n_ids * sizeof(int32_t));
int32_t *h_dst = (int32_t*)malloc(n_ids * sizeof(int32_t));
int32_t h_eb[N_EXPERTS+1];
// candidate kernels for this leg: smem generic, smem fast, global fast
struct cand { const char *name; int variant; }; // 0=<0> smem, 1=<neu> smem, 2=<neu> global
std::vector<cand> cands;
if (smem_fits) { cands.push_back({"smem<0>", 0}); cands.push_back({"smem<t>", 1}); }
cands.push_back({"global<t>", 2});
bool ok = true;
float ms[3] = {0, 0, 0};
for (const cand &C : cands) {
// production pre-zero of both id maps
CHECK(cudaMemset(d_src1, 0, n_ids * sizeof(int32_t)));
CHECK(cudaMemset(d_dst, 0, n_ids * sizeof(int32_t)));
CHECK(cudaMemset(d_eb, 0xFF, (N_EXPERTS+1) * sizeof(int32_t)));
auto launch = [&](void) {
switch (C.variant) {
case 0: mm_ids_helper<0><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1); break;
case 1:
if (neu == 1) mm_ids_helper<1><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
else mm_ids_helper<6><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
break;
case 2:
if (neu == 1) mm_ids_helper_global<1><<<grid, block, 0>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
else mm_ids_helper_global<6><<<grid, block, 0>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
break;
}
};
launch();
CHECK(cudaGetLastError());
CHECK(cudaDeviceSynchronize());
CHECK(cudaMemcpy(h_src1, d_src1, n_ids*4, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(h_dst, d_dst, n_ids*4, cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(h_eb, d_eb, (N_EXPERTS+1)*4, cudaMemcpyDeviceToHost));
const bool m1 = memcmp(h_src1, r_src1, n_ids*4) == 0;
const bool m2 = memcmp(h_dst, r_dst, n_ids*4) == 0;
const bool m3 = memcmp(h_eb, r_eb, (N_EXPERTS+1)*4) == 0;
if (!(m1 && m2 && m3)) {
ok = false;
printf(" %-18s %-9s MISMATCH src1=%d dst=%d eb=%d\n", L.name, C.name, m1, m2, m3);
}
// timing (50 reps after 5 warmups)
cudaEvent_t t0, t1;
CHECK(cudaEventCreate(&t0)); CHECK(cudaEventCreate(&t1));
for (int w = 0; w < 5; w++) launch();
CHECK(cudaEventRecord(t0));
for (int r = 0; r < 50; r++) launch();
CHECK(cudaEventRecord(t1)); CHECK(cudaEventSynchronize(t1));
float m = 0; CHECK(cudaEventElapsedTime(&m, t0, t1));
ms[C.variant] = m / 50;
CHECK(cudaEventDestroy(t0)); CHECK(cudaEventDestroy(t1));
}
if (smem_fits) {
printf("%-20s nt=%-6d neu=%d %s smem<0> %8.3f ms smem<t> %8.3f ms global<t> %8.3f ms\n",
L.name, nt, neu, ok ? "PARITY-OK " : "MISMATCH!!", ms[0], ms[1], ms[2]);
} else {
printf("%-20s nt=%-6d neu=%d %s (past cap) global<t> %8.3f ms\n",
L.name, nt, neu, ok ? "PARITY-OK " : "MISMATCH!!", ms[2]);
}
if (!ok) failures++;
free(h_ids); free(h_src1); free(h_dst); free(r_src1); free(r_dst);
cudaFree(d_ids); cudaFree(d_src1); cudaFree(d_dst); cudaFree(d_eb);
}
printf(failures ? "PROTO FAIL (%d legs)\n" : "PROTO PASS\n", failures);
return failures ? 1 : 0;
}
cuda/mmq/test/proto_q8_aligned.cu
+296
View File
@@ -0,0 +1,296 @@
// proto_q8_aligned.cu — M1-Inc3 prototype (megakernel program).
//
// HYPOTHESIS: the dense Q8_0 decode GEMVs run below the ~200-202 GB/s kernel
// ceiling (attrib15b: attn_q_b 1024x32768 @164, 2048x4096 @157, head
// 4096x129280 @146) because block_q8_0 is 34 bytes, so the int8 code stream
// is only 2-byte aligned (same misalignment class as IQ2_XXS, proven +12%
// by proto_iq2_aligned). Repacking into an aligned SoA layout (d[] halves
// separate, qs[] 64B-aligned 32B per block) allows int4 (16B) aligned loads
// at identical byte count.
//
// A/B at the three production decode shapes (K = in_dim, M = out rows):
// attn_q_b : K=1024, M=32768 (34.0 MiB/call, 164 GB/s baseline)
// mid : K=2048, M=4096 ( 8.5 MiB/call, 157 GB/s baseline)
// head : K=4096, M=129280 (536.6 MiB/call, 146 GB/s baseline)
// baseline = ds4_mmq_q8_0_dense_vec (production decode entry, quantizes X
// internally; bare-kernel aligned numbers exclude that ~small
// cost — Inc1 lesson: judge the KERNEL here, entry design later)
//
// Build (box):
// nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
// proto_q8_aligned.cu ../*.o -lcudart -lcuda -o proto_q8_aligned
#include "ds4_mmq.h"
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
namespace {
constexpr int QK = 32; // values per Q8_0 block
constexpr int BLOCK_BYTES = 34; // sizeof(block_q8_0) = 2 + 32
// ---------------------------------------------------------------- host q8_1
struct q8_1_block { __half2 ds; int8_t qs[32]; }; // mirrors block_q8_1
static void quantize_q8_1_host(const float *x, q8_1_block *out, int k) {
for (int b = 0; b < k / 32; b++) {
float amax = 0.0f;
for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
const float d = amax / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
int sum = 0;
for (int i = 0; i < 32; i++) {
int q = (int)roundf(x[b * 32 + i] * id);
if (q > 127) q = 127;
if (q < -127) q = -127;
out[b].qs[i] = (int8_t)q;
sum += q;
}
out[b].ds = __floats2half2_rn(d, d * (float)sum);
}
}
// ------------------------------------------------------------ device kernel
// Aligned SoA weights: dq[blk] __half scale; qs stream 64B-aligned, 32B per
// block (int4 x2). Warp-per-row GEMV: lane covers one block per pass.
__global__ void q8_aligned_dense_vec_kernel(
float *out, // [M]
const int4 *qs, // aligned codes, 2 int4 per block
const __half *dq, // block scales
const q8_1_block *x8, // [K/32] quantized activation
int M,
int nb) // blocks per row = K/32
{
const int row = blockIdx.x;
const int lane = threadIdx.x;
const long long rbase = (long long)row * nb;
float acc = 0.0f;
for (int b0 = 0; b0 < nb; b0 += 32) {
const int b = b0 + lane;
const int4 w0 = qs[(rbase + b) * 2 + 0]; // aligned 16B
const int4 w1 = qs[(rbase + b) * 2 + 1];
const int *u = (const int *)x8[b].qs;
int sumi = 0;
sumi = __dp4a(w0.x, u[0], sumi);
sumi = __dp4a(w0.y, u[1], sumi);
sumi = __dp4a(w0.z, u[2], sumi);
sumi = __dp4a(w0.w, u[3], sumi);
sumi = __dp4a(w1.x, u[4], sumi);
sumi = __dp4a(w1.y, u[5], sumi);
sumi = __dp4a(w1.z, u[6], sumi);
sumi = __dp4a(w1.w, u[7], sumi);
acc += __half2float(dq[rbase + b]) * __low2float(x8[b].ds) * (float)sumi;
}
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
acc += __shfl_down_sync(0xffffffffu, acc, off);
if (lane == 0) out[row] = acc;
}
// Variant R2: two rows per 64-thread block (one warp each) — probes whether
// higher occupancy per SM changes the DRAM rate at skinny K.
__global__ void q8_aligned_dense_vec_r2_kernel(
float *out,
const int4 *qs,
const __half *dq,
const q8_1_block *x8,
int M,
int nb)
{
const int row = blockIdx.x * 2 + (threadIdx.x >> 5);
const int lane = threadIdx.x & 31;
if (row >= M) return;
const long long rbase = (long long)row * nb;
float acc = 0.0f;
for (int b0 = 0; b0 < nb; b0 += 32) {
const int b = b0 + lane;
const int4 w0 = qs[(rbase + b) * 2 + 0];
const int4 w1 = qs[(rbase + b) * 2 + 1];
const int *u = (const int *)x8[b].qs;
int sumi = 0;
sumi = __dp4a(w0.x, u[0], sumi);
sumi = __dp4a(w0.y, u[1], sumi);
sumi = __dp4a(w0.z, u[2], sumi);
sumi = __dp4a(w0.w, u[3], sumi);
sumi = __dp4a(w1.x, u[4], sumi);
sumi = __dp4a(w1.y, u[5], sumi);
sumi = __dp4a(w1.z, u[6], sumi);
sumi = __dp4a(w1.w, u[7], sumi);
acc += __half2float(dq[rbase + b]) * __low2float(x8[b].ds) * (float)sumi;
}
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
acc += __shfl_down_sync(0xffffffffu, acc, off);
if (lane == 0) out[row] = acc;
}
static void run_shape(const char *label, int M, int K, int iters, uint32_t seed) {
const int nb = K / QK;
const long long row_bytes = (long long)nb * BLOCK_BYTES;
const double call_gb = (double)M * row_bytes / 1e9;
std::mt19937 rng(seed);
// Random bytes are valid Q8_0 content; keep block scales sane.
std::vector<uint8_t> W((size_t)M * row_bytes);
for (auto &b : W) b = (uint8_t)(rng() & 0xff);
const long long nblk = (long long)M * nb;
for (long long blk = 0; blk < nblk; blk++) {
uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
memcpy(&W[blk * BLOCK_BYTES], &h, 2);
}
std::vector<float> X(K);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<q8_1_block> X8(K / 32);
quantize_q8_1_host(X.data(), X8.data(), K);
// Aligned SoA repack: [half dq[nblk]][pad to 64B][int8 qs[nblk*32]]
const uint64_t dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
std::vector<uint8_t> ART(dq_bytes + (uint64_t)nblk * 32u);
for (long long blk = 0; blk < nblk; blk++) {
memcpy(&ART[blk * 2], &W[blk * BLOCK_BYTES], 2);
memcpy(&ART[dq_bytes + (uint64_t)blk * 32u], &W[blk * BLOCK_BYTES + 2], 32);
}
// L2 trap (proven on the iq2 pair test AND the first run of this proto:
// "mid" printed 931 GB/s, 4x the DRAM ceiling): a weight buffer smaller
// than ~L2 stays resident across timing iters. Rotate through enough
// copies that consecutive iters never re-hit cache (~256 MB footprint).
const int n_copies = (int)((256ull * 1024 * 1024) / (W.size() + 1)) + 1;
uint8_t *dW, *dArt; float *dX, *dOutBase, *dOutAl; q8_1_block *dX8;
CK(cudaMalloc(&dW, W.size() * n_copies));
CK(cudaMalloc(&dArt, ART.size() * n_copies));
CK(cudaMalloc(&dX, sizeof(float) * K));
CK(cudaMalloc(&dOutBase, sizeof(float) * M));
CK(cudaMalloc(&dOutAl, sizeof(float) * M));
CK(cudaMalloc(&dX8, X8.size() * sizeof(q8_1_block)));
for (int c = 0; c < n_copies; c++) {
CK(cudaMemcpy(dW + (size_t)c * W.size(), W.data(), W.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dArt + (size_t)c * ART.size(), ART.data(), ART.size(), cudaMemcpyHostToDevice));
}
CK(cudaMemcpy(dX, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
CK(cudaMemcpy(dX8, X8.data(), X8.size() * sizeof(q8_1_block), cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
// ---- baseline: production dense vec entry ------------------------------
int rc = ds4_mmq_q8_0_dense_vec(dW, dX, dOutBase, M, 1, K, stream);
if (rc != 0) { printf("%s: baseline dense_vec rc=%d\n", label, rc); exit(1); }
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_q8_0_dense_vec(dW + (size_t)(i % n_copies) * W.size(),
dX, dOutBase, M, 1, K, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
ms_base /= iters;
// ---- aligned prototype (warp per row) ----------------------------------
const int4 *qs = (const int4 *)(dArt + dq_bytes);
const __half *dq = (const __half *)dArt;
q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
CK(cudaGetLastError());
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++) {
const uint8_t *art_i = dArt + (size_t)(i % n_copies) * ART.size();
q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(
dOutAl, (const int4 *)(art_i + dq_bytes), (const __half *)art_i, dX8, M, nb);
}
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
ms_al /= iters;
// ---- variant R2 ---------------------------------------------------------
q8_aligned_dense_vec_r2_kernel<<<(M + 1) / 2, 64, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
CK(cudaGetLastError());
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++) {
const uint8_t *art_i = dArt + (size_t)(i % n_copies) * ART.size();
q8_aligned_dense_vec_r2_kernel<<<(M + 1) / 2, 64, 0, stream>>>(
dOutAl, (const int4 *)(art_i + dq_bytes), (const __half *)art_i, dX8, M, nb);
}
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_r2 = 0.0f; CK(cudaEventElapsedTime(&ms_r2, e0, e1));
ms_r2 /= iters;
// ---- correctness: both kernels vs a DOUBLE host reference --------------
// The per-block integer dot is exact on both paths; only the float
// accumulation order differs (256 block terms at K=8192), so
// kernel-vs-kernel tolerance false-fails on cancellation-heavy rows.
// Judge each kernel by its distance from the double reference instead.
(void)ds4_mmq_q8_0_dense_vec(dW, dX, dOutBase, M, 1, K, stream);
q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
CK(cudaStreamSynchronize(stream));
std::vector<float> ob(M), oa(M);
CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * M, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * M, cudaMemcpyDeviceToHost));
double err_b = 0.0, err_a = 0.0; int bad = 0;
for (int i = 0; i < M; i++) {
double ref = 0.0;
const uint8_t *rw = &W[(size_t)i * row_bytes];
for (int b = 0; b < nb; b++) {
const uint8_t *blk = rw + (size_t)b * BLOCK_BYTES;
uint16_t h; memcpy(&h, blk, 2);
const float dw = __half2float(__ushort_as_half(h));
const int8_t *q = (const int8_t *)(blk + 2);
const int8_t *u = X8[b].qs;
int s = 0;
for (int j = 0; j < 32; j++) s += (int)q[j] * (int)u[j];
ref += (double)dw * (double)__half2float(__low2half(X8[b].ds)) * (double)s;
}
const double eb = fabs((double)ob[i] - ref);
const double ea = fabs((double)oa[i] - ref);
if (eb > err_b) err_b = eb;
if (ea > err_a) err_a = ea;
// aligned must not be meaningfully further from ref than baseline
if (ea > eb * 4.0 + 1e-2) bad++;
}
printf("%-9s K=%-5d M=%-7d (%.1f MB/call, iters=%d)\n", label, K, M, call_gb * 1000.0, iters);
printf(" baseline dense_vec : %.4f ms -> %6.1f GB/s\n", ms_base, call_gb / (ms_base / 1e3));
printf(" aligned warp/row : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_al, call_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
printf(" aligned 2row/blk : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_r2, call_gb / (ms_r2 / 1e3), 100.0 * (ms_base / ms_r2 - 1.0));
printf(" parity vs double ref: max_abs_err base=%.3e aligned=%.3e bad=%d -> %s\n",
err_b, err_a, bad, bad == 0 ? "PASS" : "FAIL");
if (bad != 0) exit(2);
CK(cudaFree(dW)); CK(cudaFree(dArt)); CK(cudaFree(dX));
CK(cudaFree(dOutBase)); CK(cudaFree(dOutAl)); CK(cudaFree(dX8));
CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1));
CK(cudaStreamDestroy(stream));
}
} // namespace
int main(int argc, char **argv) {
const uint32_t seed = argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u;
if (ds4_mmq_init(0) != 0) { printf("ds4_mmq_init failed\n"); return 1; }
printf("PROTO_Q8_ALIGNED (block_q8_0 34B -> aligned SoA [dq][pad64][qs 32B])\n");
run_shape("attn_q_b", 32768, 1024, 2000, seed);
run_shape("mid", 4096, 2048, 2000, seed + 1);
run_shape("out_a", 4096, 8192, 2000, seed + 3);
run_shape("head", 129280, 4096, 200, seed + 2);
return 0;
}
cuda/mmq/test/proto_q8_warp8.cu
+665
View File
@@ -0,0 +1,665 @@
// proto_q8_warp8.cu — Inc4 prototype (megakernel program).
//
// The nsys truth table (nsys33) put the three CUSTOM warp8 Q8_0 decode
// kernels at 21.2 ms/step combined (hc_expand 9.09 / grouped_a 6.98 /
// pair 5.13) running at 150-195 GB/s on the raw 34-byte block layout.
// proto_q8_aligned proved the aligned SoA layout ([half dq[nblk]][pad64]
// [int8 qs[nblk*32]], derived kind 5) reaches 235-245 GB/s at these K.
//
// This proto A/Bs each warp8 kernel VERBATIM against an aligned variant
// with the SAME warp8 structure, prequant inputs (xq int8 + xscale float,
// quantize_q8_0_f32 semantics), accumulation order, and epilogue — only
// the weight load path changes (unaligned 34B-stride assembles -> two
// aligned int4 per block). Production Flash shapes:
// hc_expand : out=4096 in=8192 (attn_output_b) + out=4096 in=2048 (shared_down)
// pair : in=4096 out0=1024 out1=512 (q_a+kv) + out0=out1=2048 (shared gate+up)
// grouped : group_dim=4096 rank=1024 n_groups=8 (attn_output_a)
//
// L2-defeating rotation (>=256 MB of weight copies) + double host reference
// parity, per the proto_q8_aligned method rules.
//
// Build (box):
// nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
// proto_q8_warp8.cu -lcudart -lcuda -o proto_q8_warp8
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
namespace {
// ---- device helpers copied verbatim from ds4_cuda.cu ----------------------
__device__ __forceinline__ int32_t load_i8x4_i32_unaligned(const int8_t *p) {
const uint8_t *u = (const uint8_t *)p;
return (int32_t)((uint32_t)u[0] |
((uint32_t)u[1] << 8) |
((uint32_t)u[2] << 16) |
((uint32_t)u[3] << 24));
}
__device__ __forceinline__ int32_t load_i8x4_i32_aligned(const int8_t *p) {
return *(const int32_t *)p;
}
__device__ __forceinline__ int32_t dot_i8x32_dp4a(const int8_t *a, const int8_t *b) {
int32_t dot = 0;
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
dot = __dp4a(load_i8x4_i32_unaligned(a + i), load_i8x4_i32_aligned(b + i), dot);
}
return dot;
}
__device__ __forceinline__ int32_t dot_i8_block(const int8_t *a, const int8_t *b, uint64_t n, int use_dp4a) {
if (use_dp4a && n == 32u) return dot_i8x32_dp4a(a, b);
int32_t dot = 0;
for (uint64_t i = 0; i < n; i++) dot += (int32_t)a[i] * (int32_t)b[i];
return dot;
}
__device__ float warp_sum_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v += __shfl_down_sync(0xffffffffu, v, offset);
}
return v;
}
// Aligned per-block dot: two 16B loads + 8 dp4a, same integer result.
__device__ __forceinline__ int32_t dot_i8x32_aligned(const int8_t *qs32, const int8_t *xqb) {
const int4 w0 = *(const int4 *)qs32;
const int4 w1 = *(const int4 *)(qs32 + 16);
const int *u = (const int *)xqb;
int sumi = 0;
sumi = __dp4a(w0.x, u[0], sumi);
sumi = __dp4a(w0.y, u[1], sumi);
sumi = __dp4a(w0.z, u[2], sumi);
sumi = __dp4a(w0.w, u[3], sumi);
sumi = __dp4a(w1.x, u[4], sumi);
sumi = __dp4a(w1.y, u[5], sumi);
sumi = __dp4a(w1.z, u[6], sumi);
sumi = __dp4a(w1.w, u[7], sumi);
return sumi;
}
// ==== 1. hc_expand ==========================================================
// Baseline: verbatim matmul_q8_0_hc_expand_preq_warp8_kernel (ds4_cuda.cu).
__global__ void hc_expand_base_kernel(
float *out_hc, float *block_out, const float *block_add,
const float *residual_hc, const float *split, const unsigned char *w,
const int8_t *xq, const float *xscale,
uint64_t in_dim, uint64_t out_dim, uint32_t n_embd, uint32_t n_hc,
uint64_t blocks, int has_add, int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const unsigned char *wr = w + row * blocks * 34;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xq + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xscale[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) {
const uint32_t d = (uint32_t)row;
block_out[d] = acc;
float block_v = acc;
if (has_add) block_v += block_add[d];
const float *post = split + n_hc;
const float *comb = split + 2u * n_hc;
for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
float hc_acc = block_v * post[dst_hc];
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
hc_acc += comb_v * res_v;
}
out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
}
}
}
// Aligned variant: identical loop/epilogue; weight loads via [dq|qs] artifact.
__global__ void hc_expand_aligned_kernel(
float *out_hc, float *block_out, const float *block_add,
const float *residual_hc, const float *split,
const int8_t *qs_art, const __half *dq_art,
const int8_t *xq, const float *xscale,
uint64_t in_dim, uint64_t out_dim, uint32_t n_embd, uint32_t n_hc,
uint64_t blocks, int has_add) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const uint64_t rbase = row * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const int sumi = dot_i8x32_aligned(qs_art + (rbase + b) * 32u, xq + b * 32u);
acc += __half2float(dq_art[rbase + b]) * xscale[b] * (float)sumi;
}
acc = warp_sum_f32(acc);
if (lane == 0) {
const uint32_t d = (uint32_t)row;
block_out[d] = acc;
float block_v = acc;
if (has_add) block_v += block_add[d];
const float *post = split + n_hc;
const float *comb = split + 2u * n_hc;
for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
float hc_acc = block_v * post[dst_hc];
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
hc_acc += comb_v * res_v;
}
out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
}
}
}
// ==== 2. pair ===============================================================
__global__ void pair_base_kernel(
float *out0, float *out1, const unsigned char *w0, const unsigned char *w1,
const int8_t *xq, const float *xscale,
uint64_t in_dim, uint64_t out0_dim, uint64_t out1_dim, uint64_t blocks,
int use_dp4a) {
uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
uint32_t lane = threadIdx.x & 31u;
if (row >= out0_dim && row >= out1_dim) return;
float acc0 = 0.0f;
float acc1 = 0.0f;
const unsigned char *wr0 = row < out0_dim ? w0 + row * blocks * 34 : NULL;
const unsigned char *wr1 = row < out1_dim ? w1 + row * blocks * 34 : NULL;
for (uint64_t b = lane; b < blocks; b += 32u) {
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const int8_t *xqb = xq + b * 32;
const float xs = xscale[b];
if (wr0) {
const __half *scale_h = (const __half *)(wr0 + b * 34);
const int8_t *qs = (const int8_t *)(wr0 + b * 34 + 2);
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc0 += __half2float(*scale_h) * xs * (float)dot;
}
if (wr1) {
const __half *scale_h = (const __half *)(wr1 + b * 34);
const int8_t *qs = (const int8_t *)(wr1 + b * 34 + 2);
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc1 += __half2float(*scale_h) * xs * (float)dot;
}
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
if (lane == 0) {
if (row < out0_dim) out0[row] = acc0;
if (row < out1_dim) out1[row] = acc1;
}
}
__global__ void pair_aligned_kernel(
float *out0, float *out1,
const int8_t *qs0, const __half *dq0,
const int8_t *qs1, const __half *dq1,
const int8_t *xq, const float *xscale,
uint64_t in_dim, uint64_t out0_dim, uint64_t out1_dim, uint64_t blocks) {
uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
uint32_t lane = threadIdx.x & 31u;
if (row >= out0_dim && row >= out1_dim) return;
const uint64_t rbase = row * blocks;
const int in0 = row < out0_dim;
const int in1 = row < out1_dim;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const int8_t *xqb = xq + b * 32;
const float xs = xscale[b];
if (in0) {
const int sumi = dot_i8x32_aligned(qs0 + (rbase + b) * 32u, xqb);
acc0 += __half2float(dq0[rbase + b]) * xs * (float)sumi;
}
if (in1) {
const int sumi = dot_i8x32_aligned(qs1 + (rbase + b) * 32u, xqb);
acc1 += __half2float(dq1[rbase + b]) * xs * (float)sumi;
}
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
if (lane == 0) {
if (row < out0_dim) out0[row] = acc0;
if (row < out1_dim) out1[row] = acc1;
}
}
// ==== 3. grouped ============================================================
__global__ void grouped_base_kernel(
float *low, const unsigned char *w, const int8_t *xq, const float *xscale,
uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint32_t n_tokens,
uint64_t blocks, int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
const uint32_t lane = threadIdx.x & 31u;
const uint64_t low_dim = (uint64_t)n_groups * rank;
if (row >= low_dim || tok >= n_tokens) return;
const uint64_t group = row / rank;
const uint64_t row_in_group = row - group * rank;
const unsigned char *wr = w + (group * rank + row_in_group) * blocks * 34;
const uint64_t xrow = tok * (uint64_t)n_groups + group;
const int8_t *xqr = xq + xrow * blocks * 32;
const float *xsr = xscale + xrow * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = group_dim - i0 < 32 ? group_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xqr + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) low[tok * low_dim + row] = acc;
}
__global__ void grouped_aligned_kernel(
float *low, const int8_t *qs_art, const __half *dq_art,
const int8_t *xq, const float *xscale,
uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint32_t n_tokens,
uint64_t blocks) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
const uint32_t lane = threadIdx.x & 31u;
const uint64_t low_dim = (uint64_t)n_groups * rank;
if (row >= low_dim || tok >= n_tokens) return;
const uint64_t group = row / rank;
const uint64_t rbase = row * blocks;
const uint64_t xrow = tok * (uint64_t)n_groups + group;
const int8_t *xqr = xq + xrow * blocks * 32;
const float *xsr = xscale + xrow * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const int sumi = dot_i8x32_aligned(qs_art + (rbase + b) * 32u, xqr + b * 32u);
acc += __half2float(dq_art[rbase + b]) * xsr[b] * (float)sumi;
}
acc = warp_sum_f32(acc);
if (lane == 0) low[tok * low_dim + row] = acc;
}
// ==== host helpers ==========================================================
// Mirror quantize_q8_0_f32_kernel semantics (lrintf, clamp 127/-128).
static void quantize_q8_0_host(const float *x, int8_t *xq, float *xscale, int K) {
for (int b = 0; b < K / 32; b++) {
float amax = 0.0f;
for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
const float d = amax / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
xscale[b] = d;
for (int i = 0; i < 32; i++) {
int v = (int)lrintf(x[b * 32 + i] * id);
v = v > 127 ? 127 : (v < -128 ? -128 : v);
xq[b * 32 + i] = (int8_t)v;
}
}
}
struct RawArt {
std::vector<uint8_t> W; // raw 34B-block rows
std::vector<uint8_t> ART; // [dq halves][pad64][qs 32B]
uint64_t dq_bytes;
};
static RawArt make_weights(std::mt19937 &rng, long long rows, int blocks) {
RawArt r;
const long long nblk = rows * blocks;
r.W.resize((size_t)nblk * 34);
for (auto &b : r.W) b = (uint8_t)(rng() & 0xff);
for (long long blk = 0; blk < nblk; blk++) {
uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
memcpy(&r.W[blk * 34], &h, 2);
}
r.dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
r.ART.resize(r.dq_bytes + (uint64_t)nblk * 32u);
for (long long blk = 0; blk < nblk; blk++) {
memcpy(&r.ART[blk * 2], &r.W[blk * 34], 2);
memcpy(&r.ART[r.dq_bytes + (uint64_t)blk * 32u], &r.W[blk * 34 + 2], 32);
}
return r;
}
// Double reference GEMV row: sum_b dq[b] * xscale[b] * (int dot)
static double ref_row(const uint8_t *wrow, const int8_t *xq, const float *xs, int blocks) {
double ref = 0.0;
for (int b = 0; b < blocks; b++) {
const uint8_t *blk = wrow + (size_t)b * 34;
uint16_t h; memcpy(&h, blk, 2);
const float dw = __half2float(__ushort_as_half(h));
const int8_t *q = (const int8_t *)(blk + 2);
const int8_t *u = xq + (size_t)b * 32;
int s = 0;
for (int j = 0; j < 32; j++) s += (int)q[j] * (int)u[j];
ref += (double)dw * (double)xs[b] * (double)s;
}
return ref;
}
struct DevCopies {
uint8_t *dW, *dArt;
int n_copies;
size_t w_sz, art_sz;
const uint8_t *w(int i) const { return dW + (size_t)(i % n_copies) * w_sz; }
const uint8_t *art(int i) const { return dArt + (size_t)(i % n_copies) * art_sz; }
const int8_t *art_qs(int i, uint64_t dqb) const { return (const int8_t *)(art(i) + dqb); }
const __half *art_dq(int i) const { return (const __half *)art(i); }
};
static DevCopies upload_rotating(const RawArt &r, size_t min_footprint = 256ull << 20) {
DevCopies d;
d.w_sz = r.W.size(); d.art_sz = r.ART.size();
d.n_copies = (int)(min_footprint / (d.w_sz + 1)) + 1;
CK(cudaMalloc(&d.dW, d.w_sz * d.n_copies));
CK(cudaMalloc(&d.dArt, d.art_sz * d.n_copies));
for (int c = 0; c < d.n_copies; c++) {
CK(cudaMemcpy(d.dW + (size_t)c * d.w_sz, r.W.data(), d.w_sz, cudaMemcpyHostToDevice));
CK(cudaMemcpy(d.dArt + (size_t)c * d.art_sz, r.ART.data(), d.art_sz, cudaMemcpyHostToDevice));
}
return d;
}
static float time_loop(cudaStream_t stream, cudaEvent_t e0, cudaEvent_t e1, int iters,
void (*launch)(int, void *), void *ud) {
launch(0, ud);
CK(cudaGetLastError());
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++) launch(i, ud);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms = 0.0f; CK(cudaEventElapsedTime(&ms, e0, e1));
return ms / iters;
}
static int check_pair_tol(const std::vector<float> &base, const std::vector<float> &al,
const std::vector<double> &ref, const char *tag) {
double eb_max = 0.0, ea_max = 0.0; int bad = 0;
for (size_t i = 0; i < ref.size(); i++) {
const double eb = fabs((double)base[i] - ref[i]);
const double ea = fabs((double)al[i] - ref[i]);
if (eb > eb_max) eb_max = eb;
if (ea > ea_max) ea_max = ea;
if (ea > eb * 4.0 + 1e-2) bad++;
}
printf(" parity(%s): max_abs_err base=%.3e aligned=%.3e bad=%d -> %s\n",
tag, eb_max, ea_max, bad, bad == 0 ? "PASS" : "FAIL");
return bad;
}
// ==== section runners =======================================================
struct HcCtx {
cudaStream_t stream; DevCopies dev; uint64_t dq_bytes;
float *out_hc, *block_out; const float *block_add, *residual, *split;
const int8_t *xq; const float *xs;
uint64_t in_dim, out_dim, blocks; uint32_t n_embd, n_hc; int base;
};
static void hc_launch(int i, void *p) {
HcCtx *c = (HcCtx *)p;
const unsigned grid = ((unsigned)c->out_dim + 7u) / 8u;
if (c->base) {
hc_expand_base_kernel<<<grid, 256, 0, c->stream>>>(
c->out_hc, c->block_out, c->block_add, c->residual, c->split,
c->dev.w(i), c->xq, c->xs, c->in_dim, c->out_dim, c->n_embd, c->n_hc,
c->blocks, 1, 1);
} else {
hc_expand_aligned_kernel<<<grid, 256, 0, c->stream>>>(
c->out_hc, c->block_out, c->block_add, c->residual, c->split,
c->dev.art_qs(i, c->dq_bytes), c->dev.art_dq(i), c->xq, c->xs,
c->in_dim, c->out_dim, c->n_embd, c->n_hc, c->blocks, 1);
}
}
static int run_hc_expand(const char *label, int out_dim, int in_dim, int iters, uint32_t seed) {
const int blocks = in_dim / 32, n_hc = 4;
std::mt19937 rng(seed);
RawArt r = make_weights(rng, out_dim, blocks);
std::vector<float> X(in_dim);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<int8_t> XQ(in_dim); std::vector<float> XS(blocks);
quantize_q8_0_host(X.data(), XQ.data(), XS.data(), in_dim);
std::vector<float> ADD(out_dim), RES((size_t)n_hc * out_dim), SPL(2 * n_hc + n_hc * n_hc);
for (auto &v : ADD) v = ((float)(rng() % 2000) - 1000.0f) / 1000.0f;
for (auto &v : RES) v = ((float)(rng() % 2000) - 1000.0f) / 1000.0f;
for (auto &v : SPL) v = ((float)(rng() % 2000) - 1000.0f) / 2000.0f;
DevCopies dev = upload_rotating(r);
float *dOutHc, *dBlk, *dAdd, *dRes, *dSpl, *dXs; int8_t *dXq;
CK(cudaMalloc(&dOutHc, sizeof(float) * n_hc * out_dim));
CK(cudaMalloc(&dBlk, sizeof(float) * out_dim));
CK(cudaMalloc(&dAdd, sizeof(float) * out_dim));
CK(cudaMalloc(&dRes, sizeof(float) * RES.size()));
CK(cudaMalloc(&dSpl, sizeof(float) * SPL.size()));
CK(cudaMalloc(&dXq, XQ.size()));
CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
CK(cudaMemcpy(dAdd, ADD.data(), sizeof(float) * ADD.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dRes, RES.data(), sizeof(float) * RES.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dSpl, SPL.data(), sizeof(float) * SPL.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
HcCtx c = {stream, dev, r.dq_bytes, dOutHc, dBlk, dAdd, dRes, dSpl,
dXq, dXs, (uint64_t)in_dim, (uint64_t)out_dim, (uint64_t)blocks,
(uint32_t)out_dim, (uint32_t)n_hc, 1};
const float ms_base = time_loop(stream, e0, e1, iters, hc_launch, &c);
std::vector<float> bBlk(out_dim), bHc((size_t)n_hc * out_dim);
CK(cudaMemcpy(bBlk.data(), dBlk, sizeof(float) * out_dim, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(bHc.data(), dOutHc, sizeof(float) * bHc.size(), cudaMemcpyDeviceToHost));
c.base = 0;
const float ms_al = time_loop(stream, e0, e1, iters, hc_launch, &c);
std::vector<float> aBlk(out_dim), aHc((size_t)n_hc * out_dim);
CK(cudaMemcpy(aBlk.data(), dBlk, sizeof(float) * out_dim, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(aHc.data(), dOutHc, sizeof(float) * aHc.size(), cudaMemcpyDeviceToHost));
// double refs
std::vector<double> rBlk(out_dim), rHc((size_t)n_hc * out_dim);
for (int d = 0; d < out_dim; d++) {
rBlk[d] = ref_row(&r.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
const double bv = rBlk[d] + (double)ADD[d];
for (int dst = 0; dst < n_hc; dst++) {
double hc = bv * (double)SPL[n_hc + dst];
for (int src = 0; src < n_hc; src++)
hc += (double)SPL[2 * n_hc + dst + src * n_hc] * (double)RES[(size_t)src * out_dim + d];
rHc[(size_t)dst * out_dim + d] = hc;
}
}
const double gb = (double)out_dim * blocks * 34 / 1e9;
printf("%-22s out=%d in=%d (%.1f MB/call, iters=%d)\n", label, out_dim, in_dim, gb * 1000, iters);
printf(" baseline warp8 : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
printf(" aligned warp8 : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
int bad = check_pair_tol(bBlk, aBlk, rBlk, "block_out");
bad += check_pair_tol(bHc, aHc, rHc, "out_hc");
CK(cudaFree(dev.dW)); CK(cudaFree(dev.dArt));
CK(cudaFree(dOutHc)); CK(cudaFree(dBlk)); CK(cudaFree(dAdd)); CK(cudaFree(dRes));
CK(cudaFree(dSpl)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
return bad;
}
struct PairCtx {
cudaStream_t stream; DevCopies d0, d1; uint64_t dqb0, dqb1;
float *out0, *out1; const int8_t *xq; const float *xs;
uint64_t in_dim, out0_dim, out1_dim, blocks; int base;
};
static void pair_launch(int i, void *p) {
PairCtx *c = (PairCtx *)p;
const uint64_t max_out = c->out0_dim > c->out1_dim ? c->out0_dim : c->out1_dim;
const unsigned grid = ((unsigned)max_out + 7u) / 8u;
if (c->base) {
pair_base_kernel<<<grid, 256, 0, c->stream>>>(
c->out0, c->out1, c->d0.w(i), c->d1.w(i), c->xq, c->xs,
c->in_dim, c->out0_dim, c->out1_dim, c->blocks, 1);
} else {
pair_aligned_kernel<<<grid, 256, 0, c->stream>>>(
c->out0, c->out1,
c->d0.art_qs(i, c->dqb0), c->d0.art_dq(i),
c->d1.art_qs(i, c->dqb1), c->d1.art_dq(i),
c->xq, c->xs, c->in_dim, c->out0_dim, c->out1_dim, c->blocks);
}
}
static int run_pair(const char *label, int out0, int out1, int in_dim, int iters, uint32_t seed) {
const int blocks = in_dim / 32;
std::mt19937 rng(seed);
RawArt r0 = make_weights(rng, out0, blocks);
RawArt r1 = make_weights(rng, out1, blocks);
std::vector<float> X(in_dim);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<int8_t> XQ(in_dim); std::vector<float> XS(blocks);
quantize_q8_0_host(X.data(), XQ.data(), XS.data(), in_dim);
DevCopies d0 = upload_rotating(r0, 128ull << 20);
DevCopies d1 = upload_rotating(r1, 128ull << 20);
float *dOut0, *dOut1, *dXs; int8_t *dXq;
CK(cudaMalloc(&dOut0, sizeof(float) * out0));
CK(cudaMalloc(&dOut1, sizeof(float) * out1));
CK(cudaMalloc(&dXq, XQ.size()));
CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
PairCtx c = {stream, d0, d1, r0.dq_bytes, r1.dq_bytes, dOut0, dOut1,
dXq, dXs, (uint64_t)in_dim, (uint64_t)out0, (uint64_t)out1, (uint64_t)blocks, 1};
const float ms_base = time_loop(stream, e0, e1, iters, pair_launch, &c);
std::vector<float> b0(out0), b1(out1);
CK(cudaMemcpy(b0.data(), dOut0, sizeof(float) * out0, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(b1.data(), dOut1, sizeof(float) * out1, cudaMemcpyDeviceToHost));
c.base = 0;
const float ms_al = time_loop(stream, e0, e1, iters, pair_launch, &c);
std::vector<float> a0(out0), a1(out1);
CK(cudaMemcpy(a0.data(), dOut0, sizeof(float) * out0, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(a1.data(), dOut1, sizeof(float) * out1, cudaMemcpyDeviceToHost));
std::vector<double> ref0(out0), ref1(out1);
for (int d = 0; d < out0; d++) ref0[d] = ref_row(&r0.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
for (int d = 0; d < out1; d++) ref1[d] = ref_row(&r1.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
const double gb = ((double)out0 + out1) * blocks * 34 / 1e9;
printf("%-22s out0=%d out1=%d in=%d (%.1f MB/call, iters=%d)\n", label, out0, out1, in_dim, gb * 1000, iters);
printf(" baseline warp8 : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
printf(" aligned warp8 : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
int bad = check_pair_tol(b0, a0, ref0, "out0");
bad += check_pair_tol(b1, a1, ref1, "out1");
CK(cudaFree(d0.dW)); CK(cudaFree(d0.dArt)); CK(cudaFree(d1.dW)); CK(cudaFree(d1.dArt));
CK(cudaFree(dOut0)); CK(cudaFree(dOut1)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
return bad;
}
struct GroupCtx {
cudaStream_t stream; DevCopies dev; uint64_t dqb;
float *low; const int8_t *xq; const float *xs;
uint64_t group_dim, rank; uint32_t n_groups; uint64_t blocks; int base;
};
static void grouped_launch(int i, void *p) {
GroupCtx *c = (GroupCtx *)p;
const uint64_t low_dim = (uint64_t)c->n_groups * c->rank;
const unsigned grid = ((unsigned)low_dim + 7u) / 8u;
if (c->base) {
grouped_base_kernel<<<grid, 256, 0, c->stream>>>(
c->low, c->dev.w(i), c->xq, c->xs, c->group_dim, c->rank,
c->n_groups, 1u, c->blocks, 1);
} else {
grouped_aligned_kernel<<<grid, 256, 0, c->stream>>>(
c->low, c->dev.art_qs(i, c->dqb), c->dev.art_dq(i), c->xq, c->xs,
c->group_dim, c->rank, c->n_groups, 1u, c->blocks);
}
}
static int run_grouped(const char *label, int group_dim, int rank, int n_groups, int iters, uint32_t seed) {
const int blocks = group_dim / 32;
const long long low_dim = (long long)n_groups * rank;
std::mt19937 rng(seed);
RawArt r = make_weights(rng, low_dim, blocks);
// per-group activation rows
std::vector<float> X((size_t)n_groups * group_dim);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<int8_t> XQ(X.size()); std::vector<float> XS((size_t)n_groups * blocks);
for (int g = 0; g < n_groups; g++)
quantize_q8_0_host(&X[(size_t)g * group_dim], &XQ[(size_t)g * group_dim],
&XS[(size_t)g * blocks], group_dim);
DevCopies dev = upload_rotating(r);
float *dLow, *dXs; int8_t *dXq;
CK(cudaMalloc(&dLow, sizeof(float) * low_dim));
CK(cudaMalloc(&dXq, XQ.size()));
CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
GroupCtx c = {stream, dev, r.dq_bytes, dLow, dXq, dXs,
(uint64_t)group_dim, (uint64_t)rank, (uint32_t)n_groups, (uint64_t)blocks, 1};
const float ms_base = time_loop(stream, e0, e1, iters, grouped_launch, &c);
std::vector<float> bl(low_dim);
CK(cudaMemcpy(bl.data(), dLow, sizeof(float) * low_dim, cudaMemcpyDeviceToHost));
c.base = 0;
const float ms_al = time_loop(stream, e0, e1, iters, grouped_launch, &c);
std::vector<float> al(low_dim);
CK(cudaMemcpy(al.data(), dLow, sizeof(float) * low_dim, cudaMemcpyDeviceToHost));
std::vector<double> ref(low_dim);
for (long long d = 0; d < low_dim; d++) {
const int g = (int)(d / rank);
ref[d] = ref_row(&r.W[(size_t)d * blocks * 34], &XQ[(size_t)g * group_dim],
&XS[(size_t)g * blocks], blocks);
}
const double gb = (double)low_dim * blocks * 34 / 1e9;
printf("%-22s gdim=%d rank=%d groups=%d (%.1f MB/call, iters=%d)\n",
label, group_dim, rank, n_groups, gb * 1000, iters);
printf(" baseline warp8 : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
printf(" aligned warp8 : %.4f ms -> %6.1f GB/s (%+.1f%%)\n",
ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
int bad = check_pair_tol(bl, al, ref, "low");
CK(cudaFree(dev.dW)); CK(cudaFree(dev.dArt));
CK(cudaFree(dLow)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
return bad;
}
} // namespace
int main(int argc, char **argv) {
const uint32_t seed = argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u;
printf("PROTO_Q8_WARP8 (Inc4: aligned-SoA variants of the custom warp8 q8_0 kernels)\n");
int bad = 0;
bad += run_hc_expand("hc_expand(output_b)", 4096, 8192, 1500, seed);
bad += run_hc_expand("hc_expand(shr_down)", 4096, 2048, 2000, seed + 1);
bad += run_pair("pair(q_a+kv)", 1024, 512, 4096, 2000, seed + 2);
bad += run_pair("pair(shr_gate+up)", 2048, 2048, 4096, 2000, seed + 3);
bad += run_grouped("grouped(output_a)", 4096, 1024, 8, 1500, seed + 4);
printf(bad == 0 ? "ALL PASS\n" : "FAILURES: %d\n", bad);
return bad == 0 ? 0 : 2;
}
cuda/mmq/test/test_f16_gemm_lt.cu
+295
View File
@@ -0,0 +1,295 @@
/* P3.3 proto (2026-07-07, CLOSED "cuBLAS knows best"): are the cutlass_80_*
* (Ampere-tile) f16 GEMMs on sm_121 leaving throughput on the table? NO:
* - The 6.25 ms cutlass_80_64x256 launch (129x/12k admission) is the
* ATTENTION OUTPUT projection (M=8192=8 groups x 1024 rank, K=4096,
* N=4096 -> 47 TFLOPS, near roof). The "11 TFLOPS" that chartered this
* recon assumed it was the indexer q_b (K=1024) -- wrong dims. The real
* indexer q_b GEMM is producer-gated (n_comp > top-k; last chunk only at
* 12k) and runs cutlass_80_128x256 at 1.3 ms / 53 TFLOPS.
* - cublasLt heuristics beat GemmEx by only ~6-10% on these shapes (not
* worth the plumbing); TF32 math mode, operand alignment (A/B/C +16/+32B),
* L2 warmth (8-set rotation) and host-vs-device A were all bracketed:
* only host-mapped A costs real time (2x), and production weights already
* resolve to device-resident copies (HBM cache / WS import) -- verified
* in-server via DS4_F16_PTR_PROBE (recon-only, not in tree).
* - hc mix (m=24) is latency-floor at any tile; comp_proj m=256 is fine.
*
* Shapes measured (production call form: A=weights f16 OP_T lda=k, B=acts
* f16 OP_N ldb=k, C f32 ldc=m, computeType F32, 32 MiB handle workspace):
*
* Build (on a GB10 box):
* /usr/local/cuda/bin/nvcc -O3 -arch=sm_121 -std=c++17 \
* -o /home/ent/dspark_work/test_f16_gemm_lt cuda/mmq/test/test_f16_gemm_lt.cu \
* -lcublasLt -lcublas
* Run: ./test_f16_gemm_lt (all shapes, 100 iters each)
* Kernel names: nsys profile -o /tmp/lt ./test_f16_gemm_lt && inspect sqlite.
*/
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cuda_fp16.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#define CK(x) do { auto _e = (x); if (_e != cudaSuccess) { \
fprintf(stderr, "CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(_e)); exit(1); } } while (0)
#define CB(x) do { auto _e = (x); if (_e != CUBLAS_STATUS_SUCCESS) { \
fprintf(stderr, "CUBLAS ERR %s @%d: %d\n", #x, __LINE__, (int)_e); exit(1); } } while (0)
static void fill_half(__half *p, size_t nelem, unsigned seed) {
__half *h = (__half *)malloc(nelem * sizeof(__half));
unsigned s = seed;
for (size_t i = 0; i < nelem; i++) {
s = s * 1664525u + 1013904223u;
h[i] = __float2half(((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 0.25f);
}
CK(cudaMemcpy(p, h, nelem * sizeof(__half), cudaMemcpyHostToDevice));
free(h);
}
static float time_iters(cudaStream_t st, int iters, void (*fn)(void *), void *ud) {
cudaEvent_t e0, e1;
CK(cudaEventCreate(&e0));
CK(cudaEventCreate(&e1));
for (int i = 0; i < 10; i++) fn(ud); /* warmup */
CK(cudaStreamSynchronize(st));
CK(cudaEventRecord(e0, st));
for (int i = 0; i < iters; i++) fn(ud);
CK(cudaEventRecord(e1, st));
CK(cudaEventSynchronize(e1));
float ms = 0.f;
CK(cudaEventElapsedTime(&ms, e0, e1));
cudaEventDestroy(e0);
cudaEventDestroy(e1);
return ms / iters;
}
struct GemmCase {
int m, n, k;
const __half *A, *B;
float *C;
cublasHandle_t bl;
cublasLtHandle_t lt;
void *ws;
size_t ws_bytes;
cublasLtMatmulDesc_t op;
cublasLtMatrixLayout_t la, lb, lc;
cublasLtMatmulAlgo_t algo;
int has_algo;
};
static void run_gemmex(void *ud) {
GemmCase *g = (GemmCase *)ud;
const float alpha = 1.f, beta = 0.f;
CB(cublasGemmEx(g->bl, CUBLAS_OP_T, CUBLAS_OP_N, g->m, g->n, g->k,
&alpha, g->A, CUDA_R_16F, g->k, g->B, CUDA_R_16F, g->k,
&beta, g->C, CUDA_R_32F, g->m, CUDA_R_32F, CUBLAS_GEMM_DEFAULT));
}
static void run_lt(void *ud) {
GemmCase *g = (GemmCase *)ud;
const float alpha = 1.f, beta = 0.f;
CB(cublasLtMatmul(g->lt, g->op, &alpha, g->A, g->la, g->B, g->lb,
&beta, g->C, g->lc, g->C, g->lc,
g->has_algo ? &g->algo : NULL, g->ws, g->ws_bytes, 0));
}
int main(void) {
struct { const char *name; int m, n, k; } shapes[] = {
{ "indexer_q_b m=8192 n=4096 k=1024", 8192, 4096, 1024 },
{ "hc_mix m=24 n=4096 k=16384", 24, 4096, 16384 },
{ "comp_proj m=256 n=4096 k=4096", 256, 4096, 4096 },
};
cublasHandle_t bl;
CB(cublasCreate(&bl));
void *hws = NULL;
const size_t hws_bytes = 32u * 1024u * 1024u;
CK(cudaMalloc(&hws, hws_bytes));
CB(cublasSetWorkspace(bl, hws, hws_bytes));
cublasLtHandle_t lt;
CB(cublasLtCreate(&lt));
void *ltws = NULL;
const size_t ltws_bytes = 32u * 1024u * 1024u;
CK(cudaMalloc(&ltws, ltws_bytes));
for (auto &s : shapes) {
__half *A, *B;
float *C;
CK(cudaMalloc(&A, (size_t)s.m * s.k * sizeof(__half)));
CK(cudaMalloc(&B, (size_t)s.n * s.k * sizeof(__half)));
CK(cudaMalloc(&C, (size_t)s.m * s.n * sizeof(float)));
fill_half(A, (size_t)s.m * s.k, 1234);
fill_half(B, (size_t)s.n * s.k, 5678);
GemmCase g = {};
g.m = s.m; g.n = s.n; g.k = s.k;
g.A = A; g.B = B; g.C = C;
g.bl = bl; g.lt = lt; g.ws = ltws; g.ws_bytes = ltws_bytes;
const double flops = 2.0 * s.m * s.n * s.k;
const float ms_ex = time_iters(0, 100, run_gemmex, &g);
printf("%-36s GemmEx DEFAULT : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_ex, flops / (ms_ex * 1e9));
/* Production handle state (ds4_gpu_init): deprecated TF32 math mode.
* On CUDA 13 this may shunt the heuristic onto a legacy kernel table
* -- the suspected source of the 64x256 cutlass_80 pick. */
CB(cublasSetMathMode(bl, CUBLAS_TF32_TENSOR_OP_MATH));
const float ms_tf = time_iters(0, 100, run_gemmex, &g);
printf("%-36s GemmEx TF32-mode : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_tf, flops / (ms_tf * 1e9));
CB(cublasSetMathMode(bl, CUBLAS_DEFAULT_MATH));
/* Production pointer classes: A = gguf tensor offset inside a mapped
* model (32B alignment class, not cudaMalloc's 256B); B = cuda_tmp
* suballocation. Heuristics gate kernels on operand alignment. */
{
__half *A32, *B32;
CK(cudaMalloc(&A32, (size_t)s.m * s.k * sizeof(__half) + 256));
CK(cudaMalloc(&B32, (size_t)s.n * s.k * sizeof(__half) + 256));
CK(cudaMemcpy((char *)A32 + 32, A, (size_t)s.m * s.k * sizeof(__half), cudaMemcpyDeviceToDevice));
CK(cudaMemcpy((char *)B32 + 32, B, (size_t)s.n * s.k * sizeof(__half), cudaMemcpyDeviceToDevice));
GemmCase g2 = g;
g2.A = (const __half *)((const char *)A32 + 32);
const float ms_a = time_iters(0, 100, run_gemmex, &g2);
printf("%-36s GemmEx A+32B : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_a, flops / (ms_a * 1e9));
GemmCase g3 = g;
g3.B = (const __half *)((const char *)B32 + 32);
const float ms_b = time_iters(0, 100, run_gemmex, &g3);
printf("%-36s GemmEx B+32B : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_b, flops / (ms_b * 1e9));
GemmCase g4 = g2;
g4.B = g3.B;
const float ms_ab = time_iters(0, 100, run_gemmex, &g4);
printf("%-36s GemmEx A+32B B+32B : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_ab, flops / (ms_ab * 1e9));
float *C32;
CK(cudaMalloc(&C32, (size_t)s.m * s.n * sizeof(float) + 256));
GemmCase g5 = g;
g5.C = (float *)((char *)C32 + 32);
const float ms_c = time_iters(0, 100, run_gemmex, &g5);
printf("%-36s GemmEx C+32B : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_c, flops / (ms_c * 1e9));
GemmCase g6 = g;
g6.C = (float *)((char *)C32 + 16);
const float ms_c16 = time_iters(0, 100, run_gemmex, &g6);
printf("%-36s GemmEx C+16B : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_c16, flops / (ms_c16 * 1e9));
CK(cudaFree(C32));
CK(cudaFree(A32));
CK(cudaFree(B32));
}
/* Lt setup: same TN f16/f16/f32-compute layout */
CB(cublasLtMatmulDescCreate(&g.op, CUBLAS_COMPUTE_32F, CUDA_R_32F));
cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
CB(cublasLtMatmulDescSetAttribute(g.op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
CB(cublasLtMatmulDescSetAttribute(g.op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
CB(cublasLtMatrixLayoutCreate(&g.la, CUDA_R_16F, s.k, s.m, s.k)); /* A: k x m, lda=k (OP_T) */
CB(cublasLtMatrixLayoutCreate(&g.lb, CUDA_R_16F, s.k, s.n, s.k));
CB(cublasLtMatrixLayoutCreate(&g.lc, CUDA_R_32F, s.m, s.n, s.m));
/* Substrate x L2-warmth matrix. Production weight pointers are
* host-mapped (client mmap-registered / WS IPC import -- cmtp3d
* measured both at the same speed); the proto default is cudaMalloc
* A reused every iteration (L2-warm). 8-set rotation defeats L2
* reuse; host-alloc A reproduces the mapped-weight class. */
{
enum { NSETS = 8 };
static __half *As[NSETS], *Bs[NSETS];
static float *Cs[NSETS];
static GemmCase gsets;
static int set_idx;
struct Rot {
static void run(void *ud) {
GemmCase *g = (GemmCase *)ud;
(void)g;
GemmCase tmp = gsets;
tmp.A = As[set_idx % NSETS];
tmp.B = Bs[set_idx % NSETS];
tmp.C = Cs[set_idx % NSETS];
set_idx++;
run_gemmex(&tmp);
}
};
gsets = g;
/* dev-cold */
for (int i = 0; i < NSETS; i++) {
CK(cudaMalloc(&As[i], (size_t)s.m * s.k * sizeof(__half)));
CK(cudaMalloc(&Bs[i], (size_t)s.n * s.k * sizeof(__half)));
CK(cudaMalloc(&Cs[i], (size_t)s.m * s.n * sizeof(float)));
fill_half(As[i], (size_t)s.m * s.k, 100 + i);
fill_half(Bs[i], (size_t)s.n * s.k, 200 + i);
}
set_idx = 0;
const float ms_dc = time_iters(0, 100, Rot::run, &g);
printf("%-36s GemmEx devA cold : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_dc, flops / (ms_dc * 1e9));
for (int i = 0; i < NSETS; i++) { CK(cudaFree(As[i])); CK(cudaFree(Bs[i])); CK(cudaFree(Cs[i])); }
/* hostA warm */
__half *Ah;
CK(cudaHostAlloc(&Ah, (size_t)s.m * s.k * sizeof(__half), cudaHostAllocMapped));
{
unsigned sd = 42;
for (size_t i = 0; i < (size_t)s.m * s.k; i++) {
sd = sd * 1664525u + 1013904223u;
Ah[i] = __float2half(((float)(sd >> 8) / (float)(1u << 24) - 0.5f) * 0.25f);
}
}
GemmCase gh = g;
gh.A = Ah;
const float ms_hw = time_iters(0, 100, run_gemmex, &gh);
printf("%-36s GemmEx hostA warm : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_hw, flops / (ms_hw * 1e9));
/* hostA cold: rotate device B/C, host A stays (weights are the
* suspect; hostA likely bypasses L2 anyway) */
for (int i = 0; i < NSETS; i++) {
CK(cudaMalloc(&Bs[i], (size_t)s.n * s.k * sizeof(__half)));
CK(cudaMalloc(&Cs[i], (size_t)s.m * s.n * sizeof(float)));
fill_half(Bs[i], (size_t)s.n * s.k, 300 + i);
As[i] = Ah;
}
gsets = gh;
set_idx = 0;
const float ms_hc = time_iters(0, 100, Rot::run, &g);
printf("%-36s GemmEx hostA cold : %8.3f ms %6.1f TFLOPS\n",
s.name, ms_hc, flops / (ms_hc * 1e9));
for (int i = 0; i < NSETS; i++) { CK(cudaFree(Bs[i])); CK(cudaFree(Cs[i])); }
CK(cudaFreeHost(Ah));
}
cublasLtMatmulPreference_t pref;
CB(cublasLtMatmulPreferenceCreate(&pref));
CB(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ltws_bytes, sizeof(ltws_bytes)));
cublasLtMatmulHeuristicResult_t res[8];
int nres = 0;
CB(cublasLtMatmulAlgoGetHeuristic(lt, g.op, g.la, g.lb, g.lc, g.lc, pref, 8, res, &nres));
printf("%-36s Lt heuristics: %d candidates\n", s.name, nres);
float best_ms = 1e30f;
int best_i = -1;
for (int i = 0; i < nres; i++) {
g.algo = res[i].algo;
g.has_algo = 1;
const float ms = time_iters(0, 100, run_lt, &g);
int algo_id = -1, tile = -1, splitk = -1;
size_t sz;
cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_ID, &algo_id, sizeof(algo_id), &sz);
cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_TILE_ID, &tile, sizeof(tile), &sz);
cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &splitk, sizeof(splitk), &sz);
printf("%-36s Lt[%d] algo=%d tile=%d splitk=%d ws=%zu: %8.3f ms %6.1f TFLOPS\n",
s.name, i, algo_id, tile, splitk, res[i].workspaceSize,
ms, flops / (ms * 1e9));
if (ms < best_ms) { best_ms = ms; best_i = i; }
}
if (best_i >= 0)
printf("%-36s Lt BEST[%d] : %8.3f ms %6.1f TFLOPS (GemmEx/Lt = %.2fx)\n",
s.name, best_i, best_ms, flops / (best_ms * 1e9), ms_ex / best_ms);
cublasLtMatmulPreferenceDestroy(pref);
cublasLtMatrixLayoutDestroy(g.la);
cublasLtMatrixLayoutDestroy(g.lb);
cublasLtMatrixLayoutDestroy(g.lc);
cublasLtMatmulDescDestroy(g.op);
CK(cudaFree(A)); CK(cudaFree(B)); CK(cudaFree(C));
}
printf("LT-GEMM DONE\n");
return 0;
}
cuda/mmq/test/test_iq2_aligned_entry.cu
+424
View File
@@ -0,0 +1,424 @@
// Production IQ2/Q2 MoE parity, including small verifier batches.
// Compare raw and aligned entries, fused gate/up and separate token rows,
// and fused down/sum against the sum of independent expert projections.
// Both entries use the same Q8_1 activation quantizer.
// Build/run: make test-cuda-dspark-moe CUDA_ARCH=sm_121
// Use --check-only to skip timing loops under Compute Sanitizer.
#include "ds4_mmq.h"
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
static int test_q2_down(cudaStream_t stream, std::mt19937 &rng) {
const int rows = 256, columns = 2048, experts = 64;
const size_t pairs = (size_t)experts * (rows / 2) * (columns / 256);
std::vector<uint8_t> artifact(ds4_mmq_q2_k_aligned_bytes(rows, columns, experts));
for (auto &byte : artifact) byte = (uint8_t)rng();
for (size_t i = 0; i < pairs * 8; i += 2) {
const uint16_t scale = 0x2400;
memcpy(artifact.data() + i, &scale, 2);
}
void *aligned, *raw;
CK(cudaMalloc(&aligned, artifact.size()));
CK(cudaMalloc(&raw, (size_t)experts * rows * (columns / 256) * 84));
CK(cudaMemcpy(aligned, artifact.data(), artifact.size(), cudaMemcpyHostToDevice));
if (ds4_mmq_q2_K_aligned_derepack(aligned, raw, rows, columns, experts, stream)) return 1;
int bad = 0;
for (int tokens = 1; tokens <= 8; tokens++) {
const int assignments = tokens * 6;
std::vector<float> x((size_t)assignments * columns);
for (float &v : x) v = ((int)(rng() % 2000) - 1000) / 500.0f;
std::vector<int32_t> ids(assignments);
for (int i = 0; i < assignments; i++) ids[i] = (i * 7) % experts;
ids.back() = -1;
float *dx, *a, *b;
int32_t *di;
const size_t out_bytes = (size_t)assignments * rows * sizeof(float);
CK(cudaMalloc(&dx, x.size() * sizeof(float)));
CK(cudaMalloc(&di, ids.size() * sizeof(int32_t)));
CK(cudaMalloc(&a, out_bytes));
CK(cudaMalloc(&b, out_bytes));
CK(cudaMemcpyAsync(dx, x.data(), x.size() * sizeof(float), cudaMemcpyHostToDevice, stream));
CK(cudaMemcpyAsync(di, ids.data(), ids.size() * sizeof(int32_t), cudaMemcpyHostToDevice, stream));
if (ds4_mmq_q2_K_aligned_moe_vec(aligned, dx, di, a, rows, columns,
assignments, experts, 1, stream) ||
ds4_mmq_q2_K_moe_vec(raw, dx, di, b, rows, columns,
assignments, experts, 1, stream)) return 1;
CK(cudaStreamSynchronize(stream));
std::vector<float> got(out_bytes / sizeof(float)), ref(got.size());
CK(cudaMemcpy(got.data(), a, out_bytes, cudaMemcpyDeviceToHost));
CK(cudaMemcpy(ref.data(), b, out_bytes, cudaMemcpyDeviceToHost));
if (memcmp(got.data(), ref.data(), out_bytes)) bad++;
if (ds4_mmq_q2_K_moe_down_sum6_vec(raw, dx, di, a, rows, columns,
tokens, experts, 6, stream)) return 1;
CK(cudaStreamSynchronize(stream));
CK(cudaMemcpy(got.data(), a, (size_t)tokens * rows * sizeof(float), cudaMemcpyDeviceToHost));
for (int t = 0; t < tokens; t++) {
for (int row = 0; row < rows; row++) {
float expected = 0.0f;
for (int s = 0; s < 6; s++) expected += ref[((size_t)t * 6 + s) * rows + row];
const float actual = got[(size_t)t * rows + row];
if (!std::isfinite(actual) || fabsf(actual - expected) > 1e-3f + fabsf(expected) * 2e-5f) bad++;
}
}
CK(cudaFree(dx)); CK(cudaFree(di)); CK(cudaFree(a)); CK(cudaFree(b));
}
CK(cudaFree(aligned)); CK(cudaFree(raw));
printf("Q2 down assignments 6..48: layout/sum failures=%d\n", bad);
return bad;
}
int main(int argc, char **argv) {
const bool parity_only = argc > 1 && !strcmp(argv[1], "--check-only");
const int M = 2048, K = 4096;
const int n_experts = 256;
const int n_slots = 6;
const int n_idsets = 32;
const int iters = parity_only ? 1 : 2000;
const int nb = K / 256;
const long long row_bytes = (long long)nb * 66;
const double slot_gb = (double)n_slots * M * row_bytes / 1e9;
if (ds4_mmq_init(0) != 0) { printf("ds4_mmq_init failed\n"); return 1; }
std::mt19937 rng(argc > 1 && !parity_only ? (uint32_t)atoi(argv[1]) : 1234u);
// Random bytes are valid IQ2_XXS content; keep block scales sane.
std::vector<uint8_t> W((size_t)n_experts * M * row_bytes);
for (auto &b : W) b = (uint8_t)(rng() & 0xff);
const long long nblk = (long long)n_experts * M * nb;
for (long long blk = 0; blk < nblk; blk++) {
uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
memcpy(&W[blk * 66], &h, 2);
}
std::vector<float> X(K);
for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<int32_t> ids(n_idsets * n_slots);
for (int s = 0; s < n_idsets; s++)
for (int j = 0; j < n_slots; j++)
ids[s * n_slots + j] = (s * n_slots + j) % n_experts;
// Artifact layout, exactly as the weight server repack kernel builds it:
// [__half dq[nblk]][pad to 64B][uint2 qs[nblk*8]]
const uint64_t expect_bytes = ds4_mmq_iq2_xxs_aligned_bytes(M, K, n_experts);
const uint64_t dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
const uint64_t art_bytes = dq_bytes + (uint64_t)nblk * 64u;
if (expect_bytes != art_bytes) {
printf("aligned_bytes mismatch: helper=%llu local=%llu\n",
(unsigned long long)expect_bytes, (unsigned long long)art_bytes);
return 1;
}
std::vector<uint8_t> ART(art_bytes);
for (long long blk = 0; blk < nblk; blk++) {
memcpy(&ART[blk * 2], &W[blk * 66], 2);
memcpy(&ART[dq_bytes + (uint64_t)blk * 64u], &W[blk * 66 + 2], 64);
}
// dArt2: identical bytes at a DIFFERENT address, used as the "up" stream
// in the pair/fused entries. Passing dArt twice lets the second stream's
// reads ride the L2 lines the first stream just fetched, inflating the
// effective rate past the DRAM ceiling; production gate/up are distinct
// 553 MB artifacts.
uint8_t *dW, *dArt, *dArt2; float *dX, *dOutBase, *dOutAl; int32_t *dIds;
CK(cudaMalloc(&dW, W.size()));
CK(cudaMalloc(&dArt, ART.size()));
CK(cudaMalloc(&dArt2, ART.size()));
CK(cudaMalloc(&dX, sizeof(float) * K));
CK(cudaMalloc(&dOutBase, sizeof(float) * n_slots * M));
CK(cudaMalloc(&dOutAl, sizeof(float) * n_slots * M));
CK(cudaMalloc(&dIds, sizeof(int32_t) * ids.size()));
CK(cudaMemcpy(dW, W.data(), W.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dArt, ART.data(), ART.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dArt2, dArt, ART.size(), cudaMemcpyDeviceToDevice));
CK(cudaMemcpy(dX, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
CK(cudaMemcpy(dIds, ids.data(), sizeof(int32_t) * ids.size(), cudaMemcpyHostToDevice));
cudaStream_t stream; CK(cudaStreamCreate(&stream));
cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
// ---- de-repack round-trip (M1-Inc2b): artifact -> raw must be byte-exact
{
uint8_t *dRaw2; CK(cudaMalloc(&dRaw2, W.size()));
CK(cudaMemset(dRaw2, 0xAB, W.size()));
const int drc = ds4_mmq_iq2_xxs_aligned_derepack(dArt, dRaw2, M, K, n_experts, stream);
if (drc != 0) { printf("derepack rc=%d\n", drc); return 1; }
CK(cudaStreamSynchronize(stream));
std::vector<uint8_t> RT(W.size());
CK(cudaMemcpy(RT.data(), dRaw2, W.size(), cudaMemcpyDeviceToHost));
if (memcmp(RT.data(), W.data(), W.size()) != 0) {
size_t off = 0; while (off < W.size() && RT[off] == W[off]) off++;
printf("derepack round-trip MISMATCH at byte %zu\n", off);
return 1;
}
CK(cudaEventRecord(e0, stream));
const int dr_iters = parity_only ? 1 : 50;
for (int i = 0; i < dr_iters; i++)
(void)ds4_mmq_iq2_xxs_aligned_derepack(dArt, dRaw2, M, K, n_experts, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_dr = 0.0f; CK(cudaEventElapsedTime(&ms_dr, e0, e1));
ms_dr /= dr_iters;
printf(" derepack round-trip: byte-exact PASS, %.3f ms/fill -> %.1f GB/s (%.1f MB tensor)\n",
ms_dr, (double)(ART.size() + W.size()) / ms_dr / 1e6,
(double)W.size() / 1e6);
CK(cudaFree(dRaw2));
}
// ---- baseline: raw-layout production vec entry -------------------------
int rc = ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
if (rc != 0) { printf("baseline moe_vec rc=%d\n", rc); return 1; }
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds + (i % n_idsets) * n_slots, dOutBase,
M, K, 1, n_experts, n_slots, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
ms_base /= iters;
// ---- aligned production entry ------------------------------------------
rc = ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 1, n_experts, n_slots, stream);
if (rc != 0) { printf("aligned entry rc=%d\n", rc); return 1; }
// width > 16 (beyond the vec-tier envelope) must be rejected so the
// caller can fall back; 1..16 are accepted since the multi-token fix.
if (ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 17, n_experts, n_slots, stream) == 0) {
printf("aligned entry accepted n_tokens=17 (must reject)\n");
return 1;
}
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds + (i % n_idsets) * n_slots, dOutAl,
M, K, 1, n_experts, n_slots, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
ms_al /= iters;
// ---- Inc2 variant P: pair (one quantize, one z=2 launch) ----------------
float *dOutUpP; CK(cudaMalloc(&dOutUpP, sizeof(float) * n_slots * M));
rc = ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX, dIds, dOutAl, dOutUpP,
M, K, 1, n_experts, n_slots, stream);
if (rc != 0) { printf("pair entry rc=%d\n", rc); return 1; }
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX,
dIds + (i % n_idsets) * n_slots,
dOutAl, dOutUpP, M, K, 1, n_experts, n_slots, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_pair = 0.0f; CK(cudaEventElapsedTime(&ms_pair, e0, e1));
ms_pair /= iters;
// ---- Inc2 variant F: fused gate+up+swiglu -> mid ------------------------
float *dMid, *dSlotW; CK(cudaMalloc(&dMid, sizeof(float) * n_slots * M));
CK(cudaMalloc(&dSlotW, sizeof(float) * n_slots));
std::vector<float> slotw(n_slots);
for (int j = 0; j < n_slots; j++) slotw[j] = 0.5f + 0.1f * j;
CK(cudaMemcpy(dSlotW, slotw.data(), sizeof(float) * n_slots, cudaMemcpyHostToDevice));
const float clampv = 10.0f;
rc = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX, dIds, dSlotW, dMid,
M, K, 1, n_experts, n_slots, clampv, stream);
if (rc != 0) { printf("fused entry rc=%d\n", rc); return 1; }
CK(cudaStreamSynchronize(stream));
CK(cudaEventRecord(e0, stream));
for (int i = 0; i < iters; i++)
(void)ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX,
dIds + (i % n_idsets) * n_slots,
dSlotW, dMid, M, K, 1, n_experts, n_slots,
clampv, stream);
CK(cudaEventRecord(e1, stream));
CK(cudaStreamSynchronize(stream));
float ms_fused = 0.0f; CK(cudaEventElapsedTime(&ms_fused, e0, e1));
ms_fused /= iters;
// NOTE: pair/fused per-iter weight bytes are 2x slot_gb (gate + up); the
// baseline-equivalent cost of one iter is 2x the single-entry time.
// ---- parity on idset 0 --------------------------------------------------
(void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
(void)ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 1, n_experts, n_slots, stream);
float *dPairG; CK(cudaMalloc(&dPairG, sizeof(float) * n_slots * M));
(void)ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX, dIds, dPairG, dOutUpP,
M, K, 1, n_experts, n_slots, stream);
(void)ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX, dIds, dSlotW, dMid,
M, K, 1, n_experts, n_slots, clampv, stream);
CK(cudaStreamSynchronize(stream));
std::vector<float> ob(n_slots * M), oa(n_slots * M), opg(n_slots * M), opu(n_slots * M), om(n_slots * M);
CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * ob.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * oa.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(opg.data(), dPairG, sizeof(float) * opg.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(opu.data(), dOutUpP, sizeof(float) * opu.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(om.data(), dMid, sizeof(float) * om.size(), cudaMemcpyDeviceToHost));
auto count_bad = [](const std::vector<float> &ref, const std::vector<float> &got,
double *mr, double *ma) {
double max_rel = 0.0, max_abs = 0.0; int bad = 0;
for (size_t i = 0; i < ref.size(); i++) {
const double a = ref[i], b = got[i];
if (!std::isfinite(a) || !std::isfinite(b)) { bad++; continue; }
const double ad = fabs(a - b);
const double rd = ad / (fabs(a) > 1e-3 ? fabs(a) : 1e-3);
if (ad > max_abs) max_abs = ad;
if (rd > max_rel) max_rel = rd;
if (rd > 2e-2 && ad > 1e-2) bad++;
}
if (mr) *mr = max_rel;
if (ma) *ma = max_abs;
return bad;
};
double max_rel = 0.0, max_abs = 0.0;
int bad = count_bad(ob, oa, &max_rel, &max_abs);
// pair: gate and up both computed from dArt, so both must match the
// aligned single-entry output exactly (same kernel math, same inputs).
int bad_pair = count_bad(ob, opg, nullptr, nullptr) + count_bad(ob, opu, nullptr, nullptr);
// fused: expected mid from the baseline outputs + the exact epilogue.
std::vector<float> em(n_slots * M);
for (int s = 0; s < n_slots; s++) {
for (int r = 0; r < M; r++) {
float g = ob[(size_t)s * M + r], u = g;
if (!std::isfinite(g)) g = 0.0f;
if (!std::isfinite(u)) u = 0.0f;
if (clampv > 1.0e-6f) {
if (g > clampv) g = clampv;
if (u > clampv) u = clampv;
if (u < -clampv) u = -clampv;
}
const float silu = g / (1.0f + expf(-g));
em[(size_t)s * M + r] = silu * u * slotw[s];
}
}
int bad_mid = count_bad(em, om, nullptr, nullptr);
// Tiny verifier batches: distinct, shared and partly overlapping experts.
// Deduplicating gate/up must be exact against separate one-token calls.
// Use different gate/up matrices so swapped operands cannot pass.
std::vector<uint8_t> up_art = ART;
for (size_t i = dq_bytes; i < up_art.size(); i++) up_art[i] ^= 0x35;
uint8_t *dRawUp;
CK(cudaMalloc(&dRawUp, W.size()));
CK(cudaMemcpyAsync(dArt2, up_art.data(), up_art.size(), cudaMemcpyHostToDevice, stream));
if (ds4_mmq_iq2_xxs_aligned_derepack(dArt2, dRawUp, M, K, n_experts, stream)) return 1;
int bad4 = 0, bad4_mid = 0, bad_dedup = test_q2_down(stream, rng);
for (int n_tok = 1; n_tok <= 8; n_tok++) {
for (int overlap = 0; overlap < 3; overlap++) {
const int n_asg = n_tok * n_slots;
std::vector<float> X4((size_t)n_tok * K);
for (auto &v : X4) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
std::vector<int32_t> ids4(n_asg);
for (int a = 0; a < n_asg; a++) {
ids4[a] = overlap == 0 ? (a * 7 + 3) % n_experts :
(a % n_slots + (overlap == 2 ? 2 * (a / n_slots) : 0)) % n_experts;
}
if (n_asg > 13) ids4[13] = -1;
std::vector<float> w4(n_asg);
for (int a = 0; a < n_asg; a++) w4[a] = 0.25f + 0.05f * a;
float *dX4, *dOut4B, *dOut4A, *dUp4B, *dMid4, *dW4; int32_t *dIds4;
CK(cudaMalloc(&dX4, sizeof(float) * X4.size()));
CK(cudaMalloc(&dOut4B, sizeof(float) * n_asg * M));
CK(cudaMalloc(&dOut4A, sizeof(float) * n_asg * M));
CK(cudaMalloc(&dUp4B, sizeof(float) * n_asg * M));
CK(cudaMalloc(&dMid4, sizeof(float) * n_asg * M));
CK(cudaMalloc(&dW4, sizeof(float) * n_asg));
CK(cudaMalloc(&dIds4, sizeof(int32_t) * n_asg));
CK(cudaMemcpy(dX4, X4.data(), sizeof(float) * X4.size(), cudaMemcpyHostToDevice));
CK(cudaMemcpy(dW4, w4.data(), sizeof(float) * n_asg, cudaMemcpyHostToDevice));
CK(cudaMemcpy(dIds4, ids4.data(), sizeof(int32_t) * n_asg, cudaMemcpyHostToDevice));
int r1 = ds4_mmq_iq2_xxs_moe_vec(dW, dX4, dIds4, dOut4B, M, K, n_tok, n_experts, n_slots, stream);
int r2 = ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX4, dIds4, dOut4A, M, K, n_tok, n_experts, n_slots, stream);
int r3 = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX4, dIds4, dW4, dMid4,
M, K, n_tok, n_experts, n_slots, clampv, stream);
int r4 = ds4_mmq_iq2_xxs_moe_vec(dRawUp, dX4, dIds4, dUp4B,
M, K, n_tok, n_experts, n_slots, stream);
if (r1 != 0 || r2 != 0 || r3 != 0 || r4 != 0) {
printf("multi-token entries rc=%d/%d/%d/%d\n", r1, r2, r3, r4);
return 1;
}
CK(cudaStreamSynchronize(stream));
std::vector<float> ob4((size_t)n_asg * M), oa4((size_t)n_asg * M), om4((size_t)n_asg * M);
CK(cudaMemcpy(ob4.data(), dOut4B, sizeof(float) * ob4.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(oa4.data(), dOut4A, sizeof(float) * oa4.size(), cudaMemcpyDeviceToHost));
CK(cudaMemcpy(om4.data(), dMid4, sizeof(float) * om4.size(), cudaMemcpyDeviceToHost));
std::vector<float> up4(om4.size());
CK(cudaMemcpy(up4.data(), dUp4B, sizeof(float) * up4.size(), cudaMemcpyDeviceToHost));
bad4 += count_bad(ob4, oa4, nullptr, nullptr);
for (int t = 0; t < n_tok; t++) {
if (ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(
dArt, dArt2, dX4 + t * K, dIds4 + t * n_slots,
dW4 + t * n_slots, dMid4 + t * n_slots * M,
M, K, 1, n_experts, n_slots, clampv, stream) != 0) return 1;
}
CK(cudaStreamSynchronize(stream));
std::vector<float> serial_mid(om4.size());
CK(cudaMemcpy(serial_mid.data(), dMid4, sizeof(float) * serial_mid.size(),
cudaMemcpyDeviceToHost));
if (memcmp(om4.data(), serial_mid.data(), sizeof(float) * om4.size())) bad_dedup++;
std::vector<float> em4((size_t)n_asg * M);
for (int a = 0; a < n_asg; a++) {
for (int r = 0; r < M; r++) {
float g = ob4[(size_t)a * M + r], u = up4[(size_t)a * M + r];
if (!std::isfinite(g)) g = 0.0f;
if (!std::isfinite(u)) u = 0.0f;
if (clampv > 1.0e-6f) {
if (g > clampv) g = clampv;
if (u > clampv) u = clampv;
if (u < -clampv) u = -clampv;
}
const float silu = g / (1.0f + expf(-g));
em4[(size_t)a * M + r] = silu * u * w4[a];
}
}
bad4_mid += count_bad(em4, om4, nullptr, nullptr);
if (ds4_mmq_iq2_xxs_moe_gate_up_mid_vec(
dW, dRawUp, dX4, dIds4, dW4, dMid4,
M, K, n_tok, n_experts, n_slots, clampv, stream)) return 1;
CK(cudaStreamSynchronize(stream));
CK(cudaMemcpy(om4.data(), dMid4, sizeof(float) * om4.size(), cudaMemcpyDeviceToHost));
bad4_mid += count_bad(em4, om4, nullptr, nullptr);
for (int t = 0; t < n_tok; t++) {
if (ds4_mmq_iq2_xxs_moe_gate_up_mid_vec(
dW, dRawUp, dX4 + t * K, dIds4 + t * n_slots,
dW4 + t * n_slots, dMid4 + t * n_slots * M,
M, K, 1, n_experts, n_slots, clampv, stream)) return 1;
}
CK(cudaStreamSynchronize(stream));
CK(cudaMemcpy(serial_mid.data(), dMid4, sizeof(float) * serial_mid.size(), cudaMemcpyDeviceToHost));
if (memcmp(om4.data(), serial_mid.data(), sizeof(float) * om4.size())) bad_dedup++;
CK(cudaFree(dX4)); CK(cudaFree(dOut4B)); CK(cudaFree(dOut4A));
CK(cudaFree(dMid4)); CK(cudaFree(dW4)); CK(cudaFree(dIds4)); CK(cudaFree(dUp4B));
}
}
CK(cudaFree(dRawUp));
printf("TEST_IQ2_ALIGNED_ENTRY M=%d K=%d slots=%d experts=%d iters=%d (weights/iter %.1f MB)\n",
M, K, n_slots, n_experts, iters, slot_gb * 1000.0);
printf(" baseline moe_vec : %.4f ms -> %6.1f GB/s (x2 for gate+up: %.4f ms)\n",
ms_base, slot_gb / (ms_base / 1e3), 2.0f * ms_base);
printf(" aligned entry : %.4f ms -> %6.1f GB/s (%+.1f%%; x2: %.4f ms)\n",
ms_al, slot_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0), 2.0f * ms_al);
printf(" pair entry : %.4f ms -> %6.1f GB/s (vs 2x aligned %+.1f%%)\n",
ms_pair, 2.0 * slot_gb / (ms_pair / 1e3), 100.0 * (2.0f * ms_al / ms_pair - 1.0));
printf(" fused entry : %.4f ms -> %6.1f GB/s (vs 2x aligned %+.1f%%)\n",
ms_fused, 2.0 * slot_gb / (ms_fused / 1e3), 100.0 * (2.0f * ms_al / ms_fused - 1.0));
printf(" parity: aligned max_rel=%.3e max_abs=%.3e bad=%d | pair bad=%d | fused-mid bad=%d"
" | tiny-batch bad=%d mid bad=%d dedup-exact bad=%d -> %s\n",
max_rel, max_abs, bad, bad_pair, bad_mid, bad4, bad4_mid, bad_dedup,
(bad == 0 && bad_pair == 0 && bad_mid == 0 && bad4 == 0 && bad4_mid == 0 && bad_dedup == 0) ? "PASS" : "FAIL");
return (bad == 0 && bad_pair == 0 && bad_mid == 0 && bad4 == 0 && bad4_mid == 0 && bad_dedup == 0) ? 0 : 2;
}
cuda/mmq/test/test_mmq_parity.cu
+1290
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/test/test_mmq_soa_tiles.cu
+290
View File
@@ -0,0 +1,290 @@
// SPDX-License-Identifier: MIT
// test_mmq_soa_tiles.cu - P4 Inc3 decision gate: SoA-direct mmq tile loaders
// vs the raw-block path, at the production V4-Flash MoE shapes.
//
// Method: generate a RANDOM aligned-SoA artifact (the weight-server format),
// derive the raw block stream from it with the production derepack kernel
// (ds4_mmq_*_aligned_derepack -- byte-exactly what the scratch path serves),
// then run the same MoE matmul through both entries and require bit-identical
// outputs:
//
// Q2_K down leg: ds4_mmq_q2_K_moe(raw) vs ds4_mmq_q2_K_moe_soa(soa)
// IQ2_XXS gate/up leg: ds4_mmq_iq2_xxs_moe_pair(raw) vs ds4_mmq_iq2_xxs_moe_pair_soa(soa)
//
// Shapes are the V4-Flash production calls at a W4096 prefill chunk:
// down: M=4096 (n_embd), K=2048 (n_ff_exp), 24576 assignments x 1 expert
// gate/up: M=2048, K=4096, 4096 tokens x 6 experts
// with 256 experts and uniform routing (~96 columns/expert, the trace shape).
//
// Timing: cudaEvent over REPS calls of each entry (identical quantize +
// mm_ids overhead on both sides, so the delta is the tile-load effect), plus
// the standalone derepack kernel cost (what the SoA path deletes per layer
// switch).
//
// Build (on the box, after `make -j6 cuda-spark` built the mmq objects):
// ar rcs libds4mmq.a cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o \
// cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o
// /usr/local/cuda/bin/nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
// -Icuda/mmq cuda/mmq/test/test_mmq_soa_tiles.cu libds4mmq.a \
// -lcudart -lcublas -lcuda -o test_mmq_soa_tiles
#include "ds4_mmq.h"
#if defined(GGML_USE_HIP)
#include "vendors/hip.h"
#else
#include <cuda_runtime.h>
#endif
// libds4mmq.a references this ds4_cuda.cu symbol from the q8-fold vec paths
// (C3 Inc4); the entries under test never reach it, so a "no fold available"
// stub satisfies the link.
extern "C" int ds4_cuda_q8_fold_take_q81(const void *src, uint64_t in_dim, void *out) {
(void)src; (void)in_dim; (void)out;
return 0;
}
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
#define CHECK(call) do { \
cudaError_t err_ = (call); \
if (err_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d: %s\n", #call, __FILE__, __LINE__, \
cudaGetErrorString(err_)); \
exit(1); \
} \
} while (0)
namespace {
int g_failures = 0;
// Random bytes with half-precision words nudged into a sane exponent range
// (bit-parity would hold for any bits -- both paths read the same words --
// but normal-range scales keep the timing legs representative).
uint16_t sane_half(std::mt19937 &rng) {
const uint16_t mant = (uint16_t)(rng() & 0x03FFu);
const uint16_t sign = (uint16_t)((rng() & 1u) << 15);
return (uint16_t)(sign | (14u << 10) | mant); // +/- [0.25, 0.5)
}
void fill_random(std::mt19937 &rng, void *dst, size_t bytes) {
uint32_t *p = (uint32_t *)dst;
for (size_t i = 0; i < bytes / 4; ++i) p[i] = rng();
}
float *device_random_floats(std::mt19937 &rng, size_t n) {
std::vector<float> h(n);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < n; ++i) h[i] = dist(rng);
float *d = nullptr;
CHECK(cudaMalloc(&d, n * sizeof(float)));
CHECK(cudaMemcpy(d, h.data(), n * sizeof(float), cudaMemcpyHostToDevice));
return d;
}
bool compare_outputs(const char *what, const float *d_a, const float *d_b, size_t n) {
std::vector<float> a(n), b(n);
CHECK(cudaMemcpy(a.data(), d_a, n * sizeof(float), cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(b.data(), d_b, n * sizeof(float), cudaMemcpyDeviceToHost));
if (memcmp(a.data(), b.data(), n * sizeof(float)) == 0) {
printf(" %-28s PARITY-OK (bit-identical, %zu floats)\n", what, n);
return true;
}
size_t first = 0, count = 0;
for (size_t i = 0; i < n; ++i) {
if (memcmp(&a[i], &b[i], 4) != 0) { if (!count) first = i; ++count; }
}
printf(" %-28s PARITY-FAIL: %zu/%zu differ, first at %zu (raw=%.9g soa=%.9g)\n",
what, count, n, first, a[first], b[first]);
++g_failures;
return false;
}
template <typename F>
float time_ms(F &&fn, int reps) {
cudaEvent_t e0, e1;
CHECK(cudaEventCreate(&e0));
CHECK(cudaEventCreate(&e1));
fn(); // warmup
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(e0));
for (int r = 0; r < reps; ++r) fn();
CHECK(cudaEventRecord(e1));
CHECK(cudaEventSynchronize(e1));
float ms = 0.0f;
CHECK(cudaEventElapsedTime(&ms, e0, e1));
CHECK(cudaEventDestroy(e0));
CHECK(cudaEventDestroy(e1));
return ms / reps;
}
// ---------------------------------------------------------------------------
// Q2_K down leg.
// ---------------------------------------------------------------------------
void run_q2k_leg(std::mt19937 &rng, int reps, int prompt_tokens) {
const int E = 256, M = 4096, K = 2048;
const int n_tokens = prompt_tokens * 6; // assignments, n_expert_used == 1
const uint64_t nb_row = (uint64_t)K / 256;
const uint64_t nblk = (uint64_t)E * M * nb_row;
const uint64_t npair = nblk / 2;
const uint64_t soa_bytes = ds4_mmq_q2_k_aligned_bytes(M, K, E);
const uint64_t raw_bytes = nblk * 84ull;
printf("Q2_K down leg: E=%d M=%d K=%d cols=%d (soa %.1f MiB, raw %.1f MiB)\n",
E, M, K, n_tokens, soa_bytes / 1048576.0, raw_bytes / 1048576.0);
// Random SoA artifact; dm halves get sane exponents.
std::vector<uint8_t> h_soa(soa_bytes);
fill_random(rng, h_soa.data(), soa_bytes);
uint16_t *dm = (uint16_t *)h_soa.data();
for (uint64_t i = 0; i < npair * 4; ++i) dm[i] = sane_half(rng);
void *d_soa = nullptr, *d_raw = nullptr;
CHECK(cudaMalloc(&d_soa, soa_bytes));
CHECK(cudaMalloc(&d_raw, raw_bytes));
CHECK(cudaMemcpy(d_soa, h_soa.data(), soa_bytes, cudaMemcpyHostToDevice));
// Raw twin = production derepack of the artifact.
if (ds4_mmq_q2_K_aligned_derepack(d_soa, d_raw, M, K, E, 0) != 0) {
fprintf(stderr, "q2k derepack failed\n"); exit(1);
}
CHECK(cudaDeviceSynchronize());
float *d_x = device_random_floats(rng, (size_t)n_tokens * K);
std::vector<int32_t> h_ids(n_tokens);
for (int i = 0; i < n_tokens; ++i) h_ids[i] = (int32_t)(rng() % E);
int32_t *d_ids = nullptr;
CHECK(cudaMalloc(&d_ids, n_tokens * sizeof(int32_t)));
CHECK(cudaMemcpy(d_ids, h_ids.data(), n_tokens * sizeof(int32_t), cudaMemcpyHostToDevice));
const size_t out_n = (size_t)M * n_tokens;
float *d_out_raw = nullptr, *d_out_soa = nullptr;
CHECK(cudaMalloc(&d_out_raw, out_n * sizeof(float)));
CHECK(cudaMalloc(&d_out_soa, out_n * sizeof(float)));
if (ds4_mmq_q2_K_moe(d_raw, d_x, d_ids, d_out_raw, M, K, n_tokens, E, 1, 0) != 0) {
fprintf(stderr, "q2k raw moe failed\n"); exit(1);
}
if (ds4_mmq_q2_K_moe_soa(d_soa, d_x, d_ids, d_out_soa, M, K, n_tokens, E, 1, 0) != 0) {
fprintf(stderr, "q2k soa moe failed\n"); exit(1);
}
CHECK(cudaDeviceSynchronize());
compare_outputs("q2k down W4096", d_out_raw, d_out_soa, out_n);
const float ms_raw = time_ms([&] {
ds4_mmq_q2_K_moe(d_raw, d_x, d_ids, d_out_raw, M, K, n_tokens, E, 1, 0);
}, reps);
const float ms_soa = time_ms([&] {
ds4_mmq_q2_K_moe_soa(d_soa, d_x, d_ids, d_out_soa, M, K, n_tokens, E, 1, 0);
}, reps);
const float ms_dr = time_ms([&] {
ds4_mmq_q2_K_aligned_derepack(d_soa, d_raw, M, K, E, 0);
}, reps);
printf(" q2k down: raw %.3f ms soa %.3f ms (%+.1f%%) [derepack alone %.3f ms/launch]\n",
ms_raw, ms_soa, 100.0 * (ms_soa - ms_raw) / ms_raw, ms_dr);
CHECK(cudaFree(d_soa)); CHECK(cudaFree(d_raw));
CHECK(cudaFree(d_x)); CHECK(cudaFree(d_ids));
CHECK(cudaFree(d_out_raw)); CHECK(cudaFree(d_out_soa));
}
// ---------------------------------------------------------------------------
// IQ2_XXS gate/up pair leg.
// ---------------------------------------------------------------------------
void run_iq2_leg(std::mt19937 &rng, int reps, int prompt_tokens) {
const int E = 256, M = 2048, K = 4096;
const int n_tokens = prompt_tokens, n_used = 6;
const uint64_t nb_row = (uint64_t)K / 256;
const uint64_t nblk = (uint64_t)E * M * nb_row;
const uint64_t soa_bytes = ds4_mmq_iq2_xxs_aligned_bytes(M, K, E);
const uint64_t raw_bytes = nblk * 66ull;
printf("IQ2_XXS gate/up leg: E=%d M=%d K=%d tokens=%dx%d (soa %.1f MiB x2, raw %.1f MiB x2)\n",
E, M, K, n_tokens, n_used, soa_bytes / 1048576.0, raw_bytes / 1048576.0);
void *d_soa[2] = {nullptr, nullptr}, *d_raw[2] = {nullptr, nullptr};
for (int t = 0; t < 2; ++t) {
std::vector<uint8_t> h_soa(soa_bytes);
fill_random(rng, h_soa.data(), soa_bytes);
uint16_t *dq = (uint16_t *)h_soa.data();
for (uint64_t i = 0; i < nblk; ++i) dq[i] = sane_half(rng);
CHECK(cudaMalloc(&d_soa[t], soa_bytes));
CHECK(cudaMalloc(&d_raw[t], raw_bytes));
CHECK(cudaMemcpy(d_soa[t], h_soa.data(), soa_bytes, cudaMemcpyHostToDevice));
if (ds4_mmq_iq2_xxs_aligned_derepack(d_soa[t], d_raw[t], M, K, E, 0) != 0) {
fprintf(stderr, "iq2 derepack failed\n"); exit(1);
}
}
CHECK(cudaDeviceSynchronize());
float *d_x = device_random_floats(rng, (size_t)n_tokens * K);
// Production routing: n_used DISTINCT experts per token.
std::vector<int32_t> h_ids((size_t)n_tokens * n_used);
for (int i = 0; i < n_tokens; ++i) {
int picked[6];
for (int j = 0; j < n_used; ++j) {
int e;
bool dup;
do {
e = (int)(rng() % E);
dup = false;
for (int q = 0; q < j; ++q) dup |= (picked[q] == e);
} while (dup);
picked[j] = e;
h_ids[(size_t)i * n_used + j] = e;
}
}
int32_t *d_ids = nullptr;
CHECK(cudaMalloc(&d_ids, h_ids.size() * sizeof(int32_t)));
CHECK(cudaMemcpy(d_ids, h_ids.data(), h_ids.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
const size_t out_n = (size_t)M * n_tokens * n_used;
float *d_out[4] = {nullptr, nullptr, nullptr, nullptr}; // raw a/b, soa a/b
for (int i = 0; i < 4; ++i) CHECK(cudaMalloc(&d_out[i], out_n * sizeof(float)));
if (ds4_mmq_iq2_xxs_moe_pair(d_raw[0], d_raw[1], d_x, d_ids, d_out[0], d_out[1],
M, K, n_tokens, E, n_used, 0) != 0) {
fprintf(stderr, "iq2 raw pair failed\n"); exit(1);
}
if (ds4_mmq_iq2_xxs_moe_pair_soa(d_soa[0], d_soa[1], d_x, d_ids, d_out[2], d_out[3],
M, K, n_tokens, E, n_used, 0) != 0) {
fprintf(stderr, "iq2 soa pair failed\n"); exit(1);
}
CHECK(cudaDeviceSynchronize());
compare_outputs("iq2 gate W4096", d_out[0], d_out[2], out_n);
compare_outputs("iq2 up W4096", d_out[1], d_out[3], out_n);
const float ms_raw = time_ms([&] {
ds4_mmq_iq2_xxs_moe_pair(d_raw[0], d_raw[1], d_x, d_ids, d_out[0], d_out[1],
M, K, n_tokens, E, n_used, 0);
}, reps);
const float ms_soa = time_ms([&] {
ds4_mmq_iq2_xxs_moe_pair_soa(d_soa[0], d_soa[1], d_x, d_ids, d_out[2], d_out[3],
M, K, n_tokens, E, n_used, 0);
}, reps);
const float ms_dr = time_ms([&] {
ds4_mmq_iq2_xxs_aligned_derepack(d_soa[0], d_raw[0], M, K, E, 0);
}, reps);
printf(" iq2 pair: raw %.3f ms soa %.3f ms (%+.1f%%) [derepack alone %.3f ms/launch, x2 tensors]\n",
ms_raw, ms_soa, 100.0 * (ms_soa - ms_raw) / ms_raw, ms_dr);
for (int t = 0; t < 2; ++t) { CHECK(cudaFree(d_soa[t])); CHECK(cudaFree(d_raw[t])); }
CHECK(cudaFree(d_x)); CHECK(cudaFree(d_ids));
for (int i = 0; i < 4; ++i) CHECK(cudaFree(d_out[i]));
}
} // namespace
int main(int argc, char **argv) {
const int reps = (argc > 1) ? atoi(argv[1]) : 20;
const int prompt_tokens = (argc > 2) ? atoi(argv[2]) : 4096;
std::mt19937 rng(0x5A0Au);
run_q2k_leg(rng, reps, prompt_tokens);
run_iq2_leg(rng, reps, prompt_tokens);
printf(g_failures ? "SOA-TILES FAIL (%d)\n" : "SOA-TILES PASS\n", g_failures);
return g_failures ? 1 : 0;
}
cuda/mmq/unary.cuh
+114
View File
@@ -0,0 +1,114 @@
#pragma once
#include "common.cuh"
#define CUDA_NEG_BLOCK_SIZE 256
#define CUDA_STEP_BLOCK_SIZE 256
#define CUDA_GELU_BLOCK_SIZE 256
#define CUDA_SILU_BLOCK_SIZE 256
#define CUDA_SILU_BACK_BLOCK_SIZE 256
#define CUDA_TANH_BLOCK_SIZE 256
#define CUDA_RELU_BLOCK_SIZE 256
#define CUDA_SIGMOID_BLOCK_SIZE 256
#define CUDA_HARDSIGMOID_BLOCK_SIZE 256
#define CUDA_EXP_BLOCK_SIZE 256
#define CUDA_HARDSWISH_BLOCK_SIZE 256
#define CUDA_SQR_BLOCK_SIZE 256
#define CUDA_SQRT_BLOCK_SIZE 256
#define CUDA_SIN_BLOCK_SIZE 256
#define CUDA_COS_BLOCK_SIZE 256
#define CUDA_GLU_BLOCK_SIZE 256
#define CUDA_XIELU_BLOCK_SIZE 256
void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_neg(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_silu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_silu_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_gelu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_gelu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_tanh(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_exp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_hardswish(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_leaky_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_sqrt(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_sin(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_cos(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_log(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_elu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_floor(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_ceil(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node);
void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node);
__device__ __forceinline__ float ggml_cuda_op_silu_single(float x) {
return x / (1.0f + expf(-x));
}
__device__ __forceinline__ float ggml_cuda_op_gelu_single(float x) {
const float GELU_COEF_A = 0.044715f;
const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
return 0.5f * x * (1.0f + tanhf(SQRT_2_OVER_PI * x * (1.0f + GELU_COEF_A * x * x)));
}
__device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) {
x = fminf(x, limit);
g = fmaxf(fminf(g, limit), -limit);
float out_glu = x / (1.0f + expf(-x * alpha));
out_glu = out_glu * (1.0f + g);
return out_glu;
}
cuda/mmq/vecdotq.cuh
+1317
View File
File diff suppressed because it is too large. Load diff
cuda/mmq/vendors/cuda.h
Vendored
+28
View File
@@ -0,0 +1,28 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda.h>
#include <cublas_v2.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#ifdef GGML_USE_NCCL
#include <nccl.h>
#endif // GGML_USE_NCCL
#if CUDART_VERSION >= 11080
#include <cuda_fp8.h>
#define FP8_AVAILABLE
#endif // CUDART_VERSION >= 11080
#if CUDART_VERSION >= 12080
#include <cuda_fp4.h>
#endif // CUDART_VERSION >= 12080
#if CUDART_VERSION < 11020
#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH
#define CUBLAS_COMPUTE_16F CUDA_R_16F
#define CUBLAS_COMPUTE_32F CUDA_R_32F
#define cublasComputeType_t cudaDataType_t
#endif // CUDART_VERSION < 11020
cuda/mmq/vendors/hip.h
Vendored
+308
View File
@@ -0,0 +1,308 @@
#pragma once
#define HIP_DISABLE_WARP_SYNC_BUILTINS 1
#include <hip/hip_runtime.h>
#include <hipblas/hipblas.h>
#include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
#ifdef GGML_USE_NCCL
#include <rccl/rccl.h>
#endif // GGML_USE_NCCL
#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT
#define CUBLAS_OP_N HIPBLAS_OP_N
#define CUBLAS_OP_T HIPBLAS_OP_T
#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
#define CUBLAS_TF32_TENSOR_OP_MATH 0
#define CUDA_R_16F HIPBLAS_R_16F
#define CUDA_R_16BF HIPBLAS_R_16B
#define CUDA_R_32F HIPBLAS_R_32F
#define CUBLAS_SIDE_RIGHT HIPBLAS_SIDE_RIGHT
#define CUBLAS_FILL_MODE_UPPER HIPBLAS_FILL_MODE_UPPER
#define CUBLAS_DIAG_NON_UNIT HIPBLAS_DIAG_NON_UNIT
#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported
#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended
#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned
#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice
#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite
#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }}
#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width)
#define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width)
#define __shfl_down_sync(mask, var, laneMask, ...) __shfl_down(var, laneMask, ##__VA_ARGS__)
#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)
#define __all_sync(mask, var) __all(var)
#define __any_sync(mask, var) __any(var)
#define cublasStrsmBatched hipblasStrsmBatched
#define cublasCreate hipblasCreate
#define cublasDestroy hipblasDestroy
#define cublasGemmEx hipblasGemmEx
#define cublasGemmBatchedEx hipblasGemmBatchedEx
#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx
#define cublasHandle_t hipblasHandle_t
#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS
#define cublasSetStream hipblasSetStream
#define cublasSgemm hipblasSgemm
#define cublasSgemmBatched hipblasSgemmBatched
#define cublasSgemmStridedBatched hipblasSgemmStridedBatched
#define cublasStatus_t hipblasStatus_t
#define cublasOperation_t hipblasOperation_t
#define cudaDevAttrCooperativeLaunch hipDeviceAttributeCooperativeLaunch
#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer
#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess
#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess
#define cudaDeviceGetAttribute hipDeviceGetAttribute
#define cudaDeviceGetPCIBusId hipDeviceGetPCIBusId
#define cudaDeviceProp hipDeviceProp_t
#define cudaDeviceSynchronize hipDeviceSynchronize
#define cudaError_t hipError_t
#define cudaErrorMemoryAllocation hipErrorOutOfMemory
#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
#define cudaEventCreateWithFlags hipEventCreateWithFlags
#define cudaEventDisableTiming hipEventDisableTiming
#define cudaEventRecord hipEventRecord
#define cudaEventSynchronize hipEventSynchronize
#define cudaEvent_t hipEvent_t
#define cudaEventDestroy hipEventDestroy
#define cudaFree hipFree
#define cudaFreeAsync hipFreeAsync
#define cudaFreeHost hipHostFree
#define cudaGetDevice hipGetDevice
#define cudaGetDeviceCount hipGetDeviceCount
#define cudaGetDeviceProperties hipGetDeviceProperties
#define cudaGetErrorString hipGetErrorString
#define cudaGetLastError hipGetLastError
#define cudaHostRegister hipHostRegister
#define cudaHostRegisterPortable hipHostRegisterPortable
#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
#define cudaHostUnregister hipHostUnregister
#define cudaLaunchCooperativeKernel hipLaunchCooperativeKernel
#define cudaLaunchHostFunc hipLaunchHostFunc
#define cudaMalloc hipMalloc
#define cudaMallocAsync hipMallocAsync
#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault)
#define cudaMallocManaged hipMallocManaged
#define cudaMemAdvise hipMemAdvise
#define cudaMemcpy hipMemcpy
#define cudaMemcpyAsync hipMemcpyAsync
#define cudaMemcpyPeerAsync hipMemcpyPeerAsync
#define cudaMemcpy2DAsync hipMemcpy2DAsync
#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
#define cudaMemcpyKind hipMemcpyKind
#define cudaMemset hipMemset
#define cudaMemsetAsync hipMemsetAsync
#define cudaMemGetInfo hipMemGetInfo
#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize
#define cudaSetDevice hipSetDevice
#define cuDeviceGet hipDeviceGet
#define CUdevice hipDevice_t
#define CUdeviceptr hipDeviceptr_t
#define cuMemUnmap hipMemUnmap
#define CUmemAccessDesc hipMemAccessDesc
#define cuMemAddressFree hipMemAddressFree
#define cuMemRelease hipMemRelease
#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t
#define cuMemCreate hipMemCreate
#define cuMemAddressReserve hipMemAddressReserve
#define cuMemMap hipMemMap
#define cuMemSetAccess hipMemSetAccess
#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity
#define CUmemAllocationProp hipMemAllocationProp
#define cuDeviceGetAttribute hipDeviceGetAttribute
#define cudaStreamCreateWithFlags hipStreamCreateWithFlags
#define cudaStreamCreate hipStreamCreate
#define cudaStreamDestroy hipStreamDestroy
#define cudaStreamFireAndForget hipStreamFireAndForget
#define cudaStreamNonBlocking hipStreamNonBlocking
#define cudaStreamPerThread hipStreamPerThread
#define cudaStreamSynchronize hipStreamSynchronize
#define cudaStreamCaptureStatus hipStreamCaptureStatus
#define cudaStreamCaptureStatusNone hipStreamCaptureStatusNone
#define cudaStreamIsCapturing hipStreamIsCapturing
#define cudaStreamWaitEvent hipStreamWaitEvent
#define cudaGraphExec_t hipGraphExec_t
#define cudaGraphNode_t hipGraphNode_t
#define cudaKernelNodeParams hipKernelNodeParams
#define cudaKernelNodeParams hipKernelNodeParams
#define cudaGraphExecDestroy hipGraphExecDestroy
#define cudaGraphLaunch hipGraphLaunch
#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure
#define cudaGraphExecUpdateResult hipGraphExecUpdateResult
#define cudaGraphNodeType hipGraphNodeType
#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel
#define cudaGraphInstantiate hipGraphInstantiate
#define cudaStreamEndCapture hipStreamEndCapture
#define cudaGraphDestroy hipGraphDestroy
#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams
#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction
#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams
#define cudaGraphNodeGetType hipGraphNodeGetType
#define cudaGraphGetNodes hipGraphGetNodes
#define cudaGraphExecUpdate hipGraphExecUpdate
#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed
#define cudaStreamBeginCapture hipStreamBeginCapture
#define cudaGraph_t hipGraph_t
#define cudaStream_t hipStream_t
#define cudaSuccess hipSuccess
#define cudaOccupancyMaxActiveBlocksPerMultiprocessor hipOccupancyMaxActiveBlocksPerMultiprocessor
#define cudaFuncSetAttribute hipFuncSetAttribute
#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
#define __trap() do { abort(); __builtin_unreachable(); } while(0)
#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED
#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED
#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE
#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH
#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR
#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED
#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR
#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED
#if HIP_VERSION >= 60500000
#define CUBLAS_COMPUTE_16F HIPBLAS_COMPUTE_16F
#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F
#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_COMPUTE_32F_FAST_16F
#define cublasComputeType_t hipblasComputeType_t
#define cudaDataType_t hipDataType
#else
#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F
#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F
#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F
#define cublasComputeType_t hipblasDatatype_t
#define cudaDataType_t hipblasDatatype_t
#endif // HIP_VERSION >= 6050000
#if !defined(__HIP_PLATFORM_AMD__)
#error "The HIP backend supports only AMD targets"
#endif // !defined(__HIP_PLATFORM_AMD__)
#define __CUDA_ARCH__ 1300
#if defined(__gfx900__) || defined(__gfx906__)
#define GCN5
#endif // defined(__gfx900__) || defined(__gfx906__)
#if defined(__gfx803__)
#define GCN4
#endif // defined(__gfx803__)
#if defined(GCN5) || defined(GCN4)
#define GCN
#endif // defined(GCN5) || defined(GCN4)
#if defined(__gfx950__)
#define CDNA4
#endif // defined(__gfx950__)
#if defined(__gfx942__)
#define CDNA3
#endif // defined(__gfx942__)
#if defined(__gfx90a__)
#define CDNA2
#endif // defined(__gfx90a__)
#if defined(__gfx908__)
#define CDNA1
#endif // defined(__gfx908__)
#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1)
#define CDNA // For the entire family
#endif // defined(CDNA4) || defined(CDNA3) || defined(CDNA2) || defined(CDNA1)
#if defined(__GFX12__)
#define RDNA4
#endif // defined(__GFX12__)
#if defined(__GFX11__)
#define RDNA3
#endif // defined(__GFX11__)
#if defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__)
#define RDNA3_5
#endif // defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__)
#if defined(RDNA3) && !defined(RDNA3_5)
#define RDNA3_0
#endif // defined(RDNA3) && !defined(RDNA3_5)
#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \
defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__)
#define RDNA2
#endif
#if defined(__gfx1010__) || defined(__gfx1012__)
#define RDNA1
#endif // defined(__gfx1010__) || defined(__gfx1012__)
#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
#define RDNA // For the entire family
#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
#ifndef __has_builtin
#define __has_builtin(x) 0
#endif
typedef __hip_bfloat16 nv_bfloat16;
typedef __hip_bfloat162 nv_bfloat162;
#if HIP_VERSION >= 60200000
#include <hip/hip_fp8.h>
typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
#define FP8_AVAILABLE
#endif // HIP_VERSION >= 60200000
typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4)));
static __device__ __forceinline__ int __vsubss4(const int a, const int b) {
const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
#if __has_builtin(__builtin_elementwise_sub_sat)
const int8x4_t c = __builtin_elementwise_sub_sat(va, vb);
return reinterpret_cast<const int &>(c);
#else
int8x4_t c;
int16_t tmp;
#pragma unroll
for (int i = 0; i < 4; i++) {
tmp = va[i] - vb[i];
if(tmp > std::numeric_limits<int8_t>::max()) tmp = std::numeric_limits<int8_t>::max();
if(tmp < std::numeric_limits<int8_t>::min()) tmp = std::numeric_limits<int8_t>::min();
c[i] = tmp;
}
return reinterpret_cast<int &>(c);
#endif // __has_builtin(__builtin_elementwise_sub_sat)
}
static __device__ __forceinline__ int __vsub4(const int a, const int b) {
return __vsubss4(a, b);
}
static __device__ __forceinline__ unsigned int __vcmpeq4(unsigned int a, unsigned int b) {
const uint8x4_t& va = reinterpret_cast<const uint8x4_t&>(a);
const uint8x4_t& vb = reinterpret_cast<const uint8x4_t&>(b);
unsigned int c;
uint8x4_t& vc = reinterpret_cast<uint8x4_t&>(c);
#pragma unroll
for (int i = 0; i < 4; ++i) {
vc[i] = va[i] == vb[i] ? 0xff : 0x00;
}
return c;
}
static __device__ __forceinline__ unsigned int __vcmpne4(unsigned int a, unsigned int b) {
const uint8x4_t& va = reinterpret_cast<const uint8x4_t&>(a);
const uint8x4_t& vb = reinterpret_cast<const uint8x4_t&>(b);
unsigned int c;
uint8x4_t& vc = reinterpret_cast<uint8x4_t&>(c);
#pragma unroll
for (int i = 0; i < 4; ++i) {
vc[i] = va[i] == vb[i] ? 0x00 : 0xff;
}
return c;
}
dir-steering/README.md
+61 -3
View File
@@ -1,8 +1,9 @@
# Directional Steering
Directional steering is a runtime activation edit for DS4. A steering file is a
flat `f32` matrix with one normalized 4096-wide direction per layer. During
inference, ds4 can apply the edit after attention outputs, FFN outputs, or both:
flat `f32` matrix with one normalized hidden-width direction per normal
transformer layer. During inference, ds4 can apply the edit after attention
outputs, FFN outputs, or both:
```text
y = y - scale * direction[layer] * dot(direction[layer], y)
@@ -11,10 +12,20 @@ y = y - scale * direction[layer] * dot(direction[layer], y)
Positive scale removes the represented direction. Negative scale amplifies it.
With no steering file or zero scales, ds4 follows the normal inference path.
The file shape depends on the model:
- DeepSeek V4 Flash: `43 x 4096`.
- GLM 5.3 Flash: `45 x 4096`. The separate MTP predictor layer is omitted.
- Qwen3.8 Flash Next: `48 x 2560`. FFN steering is applied to each
hyper-connection branch of the residual; dumps average those branches
at the last prompt token.
GLM 5.2 steering is not implemented.
## Runtime Options
```text
--dir-steering-file FILE load a 43 x 4096 f32 direction file
--dir-steering-file FILE load one f32 direction per normal model layer
--dir-steering-ffn F apply steering after FFN outputs; default is 1 when a file is provided
--dir-steering-attn F apply steering after attention outputs; default is 0
```
@@ -23,6 +34,27 @@ The FFN output is usually the best first target because it is late enough in
each layer to represent behavior, style, and topic signals. Attention steering
is available for experiments, but it can be more fragile.
## GLM 5.3 Example
Build a GLM 5.3 direction from paired target and control prompt lists:
```sh
python3 dir-steering/tools/build_direction.py \
--profile glm-5.3-flash \
--ds4 ./ds4 \
--model gguf/GLM-5.3-Flash-Q2.gguf \
--good-file /path/to/target-prompts.txt \
--bad-file /path/to/control-prompts.txt \
--out dir-steering/out/glm53-direction.json \
--component ffn_out \
--ctx 512
```
Generated `.f32` vectors are local artifacts and are not stored in the
repository. GLM 5.3 steering works with `--mtp`, `ds4-server`, native session
batching, and two-Mac tensor parallelism. For tensor parallelism, pass the same
steering file and scales to both the worker and coordinator.
## Verbosity Example
The bundled example builds a style direction from 100 paired prompts. Each pair
@@ -39,6 +71,7 @@ Build the vector:
```sh
python3 dir-steering/tools/build_direction.py \
--profile deepseek-v4-flash \
--ds4 ./ds4 \
--model ds4flash.gguf \
--good-file dir-steering/examples/succinct.txt \
@@ -149,3 +182,28 @@ Style control:
The method is not a fine-tune. It is a low-rank runtime edit, so it works best
for coarse behavior, topic, or style directions that are consistently present in
the activation captures.
## Qwen3.8 Flash Next
Capture uses `--think` / `--nothink`
(not `--think-high`). Dumps track the prompt phase explicitly, including
one-token tails, and retain the last prompt token during ordinary and MTP
decode. `attn_out` captures the output projection of both GDN and full-attention
layers, giving one row for each of the 48 trunk layers:
```sh
python3 dir-steering/tools/build_direction.py \
--profile qwen3.8-flash-next \
--ds4 ./ds4 \
--model gguf/Qwen3.8-Flash-Next-Q4.gguf \
--good-file /path/to/target-prompts.txt \
--bad-file /path/to/control-prompts.txt \
--out dir-steering/out/qwen38-direction.json \
--component ffn_out \
--ctx 512
```
Qwen steering is Metal-only. `--mtp-model`, SSD streaming, and `--power`
remain unsupported for this graph. The bank contains only the 48 trunk layers;
the embedded MTP predictor remains unsteered. Its drafts are verified by the
steered target trunk, so `--mtp` remains supported.
dir-steering/tools/build_direction.py
+51 -49
View File
@@ -1,8 +1,8 @@
#!/usr/bin/env python3
"""Build a DS4 directional-steering vector from paired prompt sets.
The extractor asks ds4 to dump one 4096-wide activation row per layer, averages
the target and control rows, and writes a flat f32 file with 43 layer vectors.
The extractor asks ds4 to dump one hidden activation row per normal transformer
layer, averages the target and control rows, and writes a flat f32 file.
At runtime ds4 applies:
y = y - scale * direction[layer] * dot(direction[layer], y)
@@ -10,6 +10,8 @@ At runtime ds4 applies:
Positive scale suppresses the target direction. Negative scale amplifies it.
"""
from __future__ import annotations
import argparse
import array
import json
@@ -20,15 +22,10 @@ import tempfile
from pathlib import Path
N_LAYER = 43
N_EMBD = 4096
SPECIALS = {
"bos": "<｜begin▁of▁sentence｜>",
"user": "<｜User｜>",
"assistant": "<｜Assistant｜>",
"think": "<think>",
"nothink": "</think>",
MODEL_PROFILES = {
"deepseek-v4-flash": (43, 4096),
"glm-5.3-flash": (45, 4096),
"qwen3.8-flash-next": (48, 2560),
}
@@ -45,20 +42,6 @@ def read_prompt_file(path: Path) -> list[str]:
return prompts
def render_ds4_prompt(system: str, user: str, think: bool) -> str:
"""Render the minimal DS4 chat prefix used for activation capture."""
pieces = [SPECIALS["bos"]]
if system:
pieces.append(system)
pieces += [
SPECIALS["user"],
user,
SPECIALS["assistant"],
SPECIALS["think"] if think else SPECIALS["nothink"],
]
return "".join(pieces)
def normalize(v: list[float]) -> list[float]:
n2 = sum(x * x for x in v)
if n2 <= 0.0:
@@ -79,17 +62,20 @@ def run_capture(
think: bool,
ctx: int,
component: str,
n_layer: int,
n_embd: int,
work: Path,
) -> list[list[float]]:
"""Run ds4 once and return the last prompt-row dump for every layer."""
prompt_path = work / "prompt.txt"
prompt_path.write_text(render_ds4_prompt(system, prompt, think), encoding="utf-8")
prompt_path.write_text(prompt, encoding="utf-8")
dump_prefix = work / "dump"
env = os.environ.copy()
env["DS4_METAL_GRAPH_DUMP_PREFIX"] = str(dump_prefix)
env["DS4_METAL_GRAPH_DUMP_NAME"] = component
env["DS4_METAL_GRAPH_DUMP_POS"] = "0"
env["DS4_QWEN4_PREFILL_CHUNK"] = str(max(ctx, 1024))
cmd = [
str(ds4),
@@ -97,24 +83,31 @@ def run_capture(
"--ctx", str(ctx),
"--prompt-file", str(prompt_path),
"-n", "1",
"--prefill-chunk", str(max(ctx, 1024)),
]
subprocess.run(cmd, cwd=ds4.parent, env=env, check=True,
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
if system:
cmd += ["--system", system]
cmd.append("--think" if think else "--nothink")
result = subprocess.run(cmd, cwd=ds4.parent, env=env, check=False,
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
if result.returncode != 0:
error = result.stderr.decode("utf-8", errors="replace")
raise RuntimeError(f"ds4 activation capture failed:\n{error}")
rows: list[list[float]] = []
for layer in range(N_LAYER):
for layer in range(n_layer):
path = work / f"dump_{component}-{layer}_pos0.bin"
data = array.array("f")
with path.open("rb") as f:
data.fromfile(f, path.stat().st_size // 4)
if len(data) < N_EMBD or len(data) % N_EMBD != 0:
if len(data) < n_embd or len(data) % n_embd != 0:
raise RuntimeError(f"bad dump shape for {path}: {len(data)} floats")
rows.append(list(data[-N_EMBD:]))
rows.append(list(data[-n_embd:]))
return rows
def add_rows(total: list[list[float]], rows: list[list[float]]) -> None:
for layer in range(N_LAYER):
def add_rows(total: list[list[float]], rows: list[list[float]], n_layer: int) -> None:
for layer in range(n_layer):
dst = total[layer]
src = rows[layer]
for i, value in enumerate(src):
@@ -125,17 +118,22 @@ def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--ds4", default="./ds4", help="path to the ds4 CLI")
ap.add_argument("--model", default="ds4flash.gguf", help="GGUF model path")
ap.add_argument("--profile", choices=MODEL_PROFILES,
default="deepseek-v4-flash",
help="model shape and chat-template profile")
ap.add_argument("--good-file", required=True,
help="desired/target prompts, one per line")
ap.add_argument("--bad-file", required=True,
help="contrast/control prompts, one per line")
ap.add_argument("--out", default="dir-steering/out/direction.json",
help="metadata JSON path; .f32 is written next to it")
ap.add_argument("--source", default="",
help="optional corpus source or provenance note")
ap.add_argument("--ctx", type=int, default=512)
ap.add_argument("--system", default="You are a helpful assistant.")
ap.add_argument("--component", default="ffn_out",
choices=("ffn_out", "attn_out"),
help="runtime-editable 4096-wide activation stream")
help="runtime-editable activation stream at the profile's embedding width")
ap.add_argument("--think", action="store_true",
help="capture after <think>; default captures direct answers")
ap.add_argument("--pair-normalize", action="store_true",
@@ -145,16 +143,18 @@ def main() -> None:
args = ap.parse_args()
ds4 = Path(args.ds4).resolve()
model = Path(args.model).resolve()
model_arg = Path(args.model)
model = model_arg.resolve()
n_layer, n_embd = MODEL_PROFILES[args.profile]
good_prompts = read_prompt_file(Path(args.good_file))
bad_prompts = read_prompt_file(Path(args.bad_file))
n = min(len(good_prompts), len(bad_prompts))
good_prompts = good_prompts[:n]
bad_prompts = bad_prompts[:n]
good_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
bad_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
pair_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
good_sum = [[0.0] * n_embd for _ in range(n_layer)]
bad_sum = [[0.0] * n_embd for _ in range(n_layer)]
pair_sum = [[0.0] * n_embd for _ in range(n_layer)]
with tempfile.TemporaryDirectory(prefix="ds4-dir-steer-") as td:
root = Path(td)
@@ -165,22 +165,22 @@ def main() -> None:
gw.mkdir()
bw.mkdir()
good_rows = run_capture(ds4, model, good, args.system, args.think,
args.ctx, args.component, gw)
args.ctx, args.component, n_layer, n_embd, gw)
bad_rows = run_capture(ds4, model, bad, args.system, args.think,
args.ctx, args.component, bw)
add_rows(good_sum, good_rows)
add_rows(bad_sum, bad_rows)
args.ctx, args.component, n_layer, n_embd, bw)
add_rows(good_sum, good_rows, n_layer)
add_rows(bad_sum, bad_rows, n_layer)
if args.pair_normalize:
for layer in range(N_LAYER):
for layer in range(n_layer):
diff = normalize([
good_rows[layer][j] - bad_rows[layer][j]
for j in range(N_EMBD)
for j in range(n_embd)
])
for j, value in enumerate(diff):
pair_sum[layer][j] += value
layers = []
for layer in range(N_LAYER):
for layer in range(n_layer):
good_mean = [x / n for x in good_sum[layer]]
bad_mean = [x / n for x in bad_sum[layer]]
if args.pair_normalize:
@@ -188,14 +188,14 @@ def main() -> None:
else:
direction = normalize([
good_mean[i] - bad_mean[i]
for i in range(N_EMBD)
for i in range(n_embd)
])
if not args.no_orthogonalize:
base = normalize(bad_mean)
projection = dot(direction, base)
direction = normalize([
direction[i] - projection * base[i]
for i in range(N_EMBD)
for i in range(n_embd)
])
layers.append(direction)
@@ -203,14 +203,16 @@ def main() -> None:
out.parent.mkdir(parents=True, exist_ok=True)
payload = {
"format": "ds4-directional-steering-v1",
"shape": [N_LAYER, N_EMBD],
"shape": [n_layer, n_embd],
"profile": args.profile,
"component": args.component,
"thinking": bool(args.think),
"pair_normalize": bool(args.pair_normalize),
"orthogonalize_control_mean": not args.no_orthogonalize,
"good_file": str(Path(args.good_file)),
"bad_file": str(Path(args.bad_file)),
"model": str(model),
"model": str(model_arg),
"source": args.source,
"note": "runtime positive scale suppresses this direction; negative scale amplifies it",
}
out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
docs/CLIENTS.md
+125
View File
@@ -0,0 +1,125 @@
# Coding Agent Clients
[README](../README.md) | [Server](SERVER.md)
Start the server first:
```sh
./ds4-server --ctx 100000 --kv-disk-dir /tmp/ds4-kv --kv-disk-space-mb 8192
```
Set the client's context limit no higher than the server's. Output tokens also
consume that context; a client limit does not enlarge the server allocation.
The examples use Flash and localhost. Change the address for a trusted remote
server. The `dsv4-local` strings are placeholders, not server authentication.
## Pi
Add this provider to `~/.pi/agent/models.json`:
```json
{
"providers": {
"ds4": {
"baseUrl": "http://127.0.0.1:8000/v1",
"api": "openai-completions",
"apiKey": "dsv4-local",
"compat": {
"supportsStore": false,
"supportsDeveloperRole": false,
"supportsReasoningEffort": true,
"supportsUsageInStreaming": true,
"maxTokensField": "max_tokens",
"supportsStrictMode": false,
"thinkingFormat": "deepseek",
"requiresReasoningContentOnAssistantMessages": true
},
"models": [{
"id": "deepseek-v4-flash",
"name": "DwarfStar Flash",
"reasoning": true,
"thinkingLevelMap": {
"off": null, "minimal": "low", "low": "low",
"medium": "medium", "high": "high", "xhigh": "xhigh"
},
"input": ["text"],
"contextWindow": 100000,
"maxTokens": 16384,
"cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0}
}]
}
}
}
```
Select `ds4/deepseek-v4-flash` in Pi. This is a text-only client configuration;
image use also requires the appropriate client input declaration and server
vision encoder.
## OpenCode
Merge the provider into `~/.config/opencode/opencode.json`:
```json
{
"$schema": "https://opencode.ai/config.json",
"provider": {
"ds4": {
"name": "DwarfStar",
"npm": "@ai-sdk/openai-compatible",
"options": {
"baseURL": "http://127.0.0.1:8000/v1",
"apiKey": "dsv4-local"
},
"models": {
"deepseek-v4-flash": {
"name": "DwarfStar Flash",
"limit": {"context": 100000, "output": 16384}
}
}
}
}
}
```
Select `ds4/deepseek-v4-flash` as the model.
## Codex CLI
Use the Responses API. Add a provider to the Codex configuration:
```toml
[model_providers.ds4]
name = "DwarfStar"
base_url = "http://127.0.0.1:8000/v1"
wire_api = "responses"
stream_idle_timeout_ms = 1000000
```
```sh
codex --model deepseek-v4-flash -c model_provider=ds4
```
## Claude Code
Use the Anthropic-compatible endpoint. A shell wrapper can select the local
model for the main agent and its secondary model roles:
```sh
#!/bin/sh
unset ANTHROPIC_API_KEY
export ANTHROPIC_BASE_URL="http://127.0.0.1:8000"
export ANTHROPIC_AUTH_TOKEN="dsv4-local"
export ANTHROPIC_MODEL="deepseek-v4-flash"
export ANTHROPIC_DEFAULT_SONNET_MODEL="deepseek-v4-flash"
export ANTHROPIC_DEFAULT_HAIKU_MODEL="deepseek-v4-flash"
export ANTHROPIC_DEFAULT_OPUS_MODEL="deepseek-v4-flash"
export CLAUDE_CODE_SUBAGENT_MODEL="deepseek-v4-flash"
export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1
export CLAUDE_STREAM_IDLE_TIMEOUT_MS=600000
exec claude "$@"
```
Give the wrapper a different name from `claude`. Agent clients can send large
initial prompts, so the first prefill may take time. Disk caching helps reuse
compatible prefixes on later sessions.
docs/CUDA_MULTI_GPU.md
+90
View File
@@ -0,0 +1,90 @@
# NVIDIA CUDA GPUs
[README](../README.md) | [Getting started](../README.md#start-here)
**Ada Lovelace is supported, including 48 GB L40S cards.** This is an important
use case: a model whose other inference implementations require newer GPUs
can still run here when DwarfStar supports its GGUF layout. The Flash kernels
do not require Blackwell's native FP4 instructions; Ada uses an appropriate
CUDA path. Support remains model-specific, not a promise to run arbitrary GGUFs.
## Build
Install the NVIDIA driver and CUDA toolkit, including `nvcc` and cuBLAS.
Use the local GPU architecture, or select Ada explicitly:
```sh
make cuda-generic
# For an Ada build:
make cuda CUDA_ARCH=sm_89
```
For a GB10 machine use the [DGX Spark](DGX_SPARK.md) target instead.
## Normal layer placement
Without tensor parallelism, DwarfStar places complete layers on the selected
GPUs. `--gpu-vram auto` uses reported free VRAM, reserving space for the graph
and context. Explicit budgets are comma-separated GiB values, one per device.
```sh
./download_model.sh ds4f-q2
./ds4 --cuda --gpu-devices 0,1,2,3 --gpu-vram auto --ctx 32768
```
This is also the multi-GPU placement mode for GLM. Startup refuses a layout
that would require unsupported CPU execution; reduce context or choose a
smaller model if necessary.
## Flash tensor parallelism
`--cuda-tensor-parallel` pairs GPUs and divides routed-expert work inside
each pair. Pairs also own successive layer ranges. This is an in-process
configuration, not the network `--role coordinator` / `--role worker` mode.
The device order matters. List all layer-home GPUs first, then their partners.
For physical pairs `(0,1)`, `(2,3)`, `(4,5)`, `(6,7)`, use:
```text
0,2,4,6,1,3,5,7
```
Check your machine's topology with `nvidia-smi topo -m`; do not assume this
ordering gives the best pairs on another server. Each pair stores one half
of the routed experts per GPU. Dense attention, routers, and shared experts
are replicated within the pair; the output head is vocabulary-sharded.
For the tested eight-L40S setup:
```sh
./download_model.sh ds4f-q4
./ds4-agent --cuda --cuda-tensor-parallel \
--gpu-devices 0,2,4,6,1,3,5,7 --gpu-vram auto --ctx 100000
```
Q4 has native grouped routed kernels for multi-user throughput. Q2 needs less
memory, but unsupported grouped shapes use a slower, correct fallback.
Four 48 GB cards with Q2 and eight with Q4 are tested configurations. An even
device count alone does not guarantee memory fit.
## Serve multiple users
```sh
./ds4-server --cuda --cuda-tensor-parallel \
--gpu-devices 0,2,4,6,1,3,5,7 --gpu-vram auto \
--ctx 100000 --batched-session 16 --host 0.0.0.0
```
The supported eight-L40S Flash configuration has reached roughly 126 aggregate
generation tokens/s with 16 decode rows. This is total throughput, not the
speed of one user. The recorded conditions and regression requirements are in
[the QA guide](../QA_BEFORE_RELEASES.md).
CUDA TP defaults to 2048-token prefill chunks. Keep that default for the
16-session, 100k-context setup; larger chunks need more scratch memory.
Reduce session count or context if all KV states do not fit.
The repository's [server launcher](../run-nvidia-tp-server.sh) is an example
from the L40S deployment, not a portable default: it uses host-specific model
and cache paths and defaults to MXFP4. The commands above need no launcher or
environment tuning. See [serving](SERVER.md) for disk caches and API access.
docs/DGX_SPARK.md
+191
View File
@@ -0,0 +1,191 @@
# DGX Spark
[README](../README.md) | [Getting started](../README.md#start-here)
Use the NVIDIA driver and CUDA development
toolkit for the machine; the build needs `nvcc` and cuBLAS.
Check that `nvidia-smi` sees the GPU before building.
## Build and run
```sh
make cuda-spark
./download_model.sh ds4f-q2
./ds4 --cuda
```
The build targets GB10 (`sm_121a`) and enables the Blackwell-specific kernels.
Do not use `--cuda-tensor-parallel` on a single Spark.
Stop other inference services before loading a model so they do not compete
for memory. Restore any services you stopped when finished.
## DeepSeek V4.1 Flash
Q2 needs SSD streaming on one 128 GB Spark:
```sh
./download_model.sh ds41f-q2
./ds4 --cuda -m gguf/DeepSeek-V4.1-Flash-Q2.gguf \
--ssd-streaming --ctx 32768
```
The file is 341 GiB, including 189 GiB of disk-only Engram tables. Keep it on
a fast local SSD. The expert cache is sized automatically; leave memory for
other programs and the context. Use the same options with `ds4-agent` or
`ds4-server`.
Two Sparks can instead keep half the experts each, using
[network tensor parallelism](DISTRIBUTED.md#tensor-parallelism-between-two-sparks).
Both need the complete GGUF on their local SSD. Single-host V4.1
[DSpark](SPECULATIVE_DECODING.md#deepseek-v41-flash-dspark) also works with SSD
streaming; TP speculation and CUDA vision are not supported for V4.1.
The [model guide](MODELS.md#deepseek-v41-flash) covers
thinking levels and the separate Metal configurations.
September 13-14, 2026, Q2 SSD with automatic cache sizing: a 3,241-token
`/read README.md` prefill reached 93-96 t/s, up from 49-54 t/s, with 32K
allocated context. No extra flag is needed. A separate 64 GiB cache-budget
run with 64K allocated context reached 384 t/s for a 32K initial prefill and
88 t/s for a 3.2K append. These are SSD-streamed V4.1 results, not resident
V4 Flash numbers; speed depends on prompt length and the expert cache.
The automatic-cache 256-token reply test decoded about 5% slower, but
prefill plus the reply fell from about 87 to 59 seconds.
See the [QA record](../QA_BEFORE_RELEASES.md#cuda-ssd-streaming) for conditions
and longer-context measurements.
On two Sparks over the direct RoCE link, Q2 measured **21.8 t/s** for one
session on October 2, 2026: a 1K prompt, 64K allocated context and 2,048
teacher-forced decode tokens, without speculation.
| Work | Speed |
| --- | ---: |
| Fresh 32K prefill | 411 t/s |
| Add 8K after 32K | 314 t/s |
Both runs used 64K allocated context and 512 decode tokens per frontier;
decode at these longer contexts was about 18 t/s. Engram tables stayed on disk.
Decode batching starts at five ready sessions; smaller groups run in order.
The earlier eight-session reference is about **28 aggregate t/s**, not per
client. See the [network TP QA record](../QA_BEFORE_RELEASES.md#cuda-network-tensor-parallelism)
for the workloads and memory limits.
## Flash 0731 on two Sparks
[Network tensor parallelism](DISTRIBUTED.md#tensor-parallelism-between-two-sparks)
also supports Flash 0731 Q2 and MXFP4, including DSpark with the matching
support file. Each host holds about 45 GiB of Q2 weights or 77 GiB of MXFP4,
plus context and runtime buffers. Download `ds4f-mxfp4` on both for MXFP4.
October 2, 2026, Q2, direct 200 Gb/s RoCE, 64K allocated context: a 105-token
English prose prompt decoded 512 tokens at **26.0 t/s** without speculation.
A fresh 32K prefill reached **618 t/s**; appending 8K reached **610 t/s**.
Ordinary decode at those longer frontiers was about 19 t/s. These are single
measurements, not release medians.
October 4, 2026, matching 0731 DSpark support, default speculation settings,
temperature zero, 64K allocated context and 512 generated tokens. Both formats
use the drafter on both ranks over RDMA. These are medians of two runs per case
(four for MXFP4 prose), including the first token or block. Benchmark processes
ran on the performance cores (`taskset -c 5-9,15-19`):
| Weights | Prompt | DSpark |
| --- | --- | ---: |
| Q2 | English prose, 105 tokens | 44.0 t/s |
| Q2 | C hash table, 31 tokens | 83.3 t/s |
| MXFP4 | English prose, 105 tokens | 41.7 t/s |
| MXFP4 | C hash table, 31 tokens | 75.6 t/s |
On one Spark, the same Q2 prose and coding prompts reached 23.4 and 44.0 t/s
with DSpark. These are workload-specific measurements, not a promised speedup;
poor draft acceptance can still make it slower. The measured settings are
automatic once DSpark is enabled; speculation itself remains opt-in.
## Qwen3.8 Flash Next
Both Q2 and Q4 fit on one Spark. After building:
```sh
./download_model.sh qwen38-q2
./ds4 --mtp --ctx 32768
```
Q2 keeps 41.73 GiB of weights resident; Q4 keeps 69.74 GiB. Both GGUFs also
contain 95.37 GiB of BF16 n-grams, read directly from the SSD. Leave room
for context and runtime buffers. Vision, the native agent, server APIs and
text-session checkpoints work on CUDA; see [Qwen setup](QWEN38_FLASH_NEXT.md).
September 15, 2026, Q2 on one Spark: a 16K prefill reached 250 t/s, extending
it to 30K reached 243 t/s, and ordinary decode stayed around 17.3 t/s.
These are single runs with 8,192-token prefill chunks and 32 decode tokens.
A separate 256-token prose run measured 17.1 t/s ordinarily and 19.9 t/s with
MTP, at temperature zero and 8K allocated context. Speculation is opt-in;
its benefit depends on the prompt.
## GLM 5.3 Flash
Q2 is the resident target for one Spark:
```sh
./download_model.sh glm53-q2
./ds4 --cuda -m gguf/GLM-5.3-Flash-Q2.gguf --ctx 16384
```
Q4 does not fit resident. GLM Spark-to-Spark tensor parallelism is not
implemented; the two-Mac RDMA instructions do not apply to GLM on CUDA.
## Vision and speculative decoding
DeepSeek Vision Experimental needs its matching text model and encoder:
```sh
./download_model.sh ds4f-vision-q2
./ds4 --cuda --vision gguf/DeepSeek-V4-Flash-Vision-Encoder.gguf
```
Use `/read image.png` in the CLI. GLM vision also works on this backend; see
[models and vision](MODELS.md#vision).
For Flash 0731, [DSpark](SPECULATIVE_DECODING.md) uses the separate 0731 support
file. Vision Experimental has a different drafter. GLM uses its built-in MTP
block with `--mtp`. None is enabled by default.
## Flash Q2 with DSpark
September 6, 2026, fully resident Flash 0731: median generation speed from
three runs of 256 tokens, 4K allocated context and a 512-token prefill chunk.
| Prompt | Temperature | Ordinary decode | DSpark |
| --- | ---: | ---: | ---: |
| C hash table | 0 | 19.72 t/s | 31.41 t/s |
| C hash table | 1 | 19.53 t/s | 29.98 t/s |
| Unpredictable prose | 1 | 19.53 t/s | 18.81 t/s |
Longer coding runs reached about 33 t/s. Poor draft acceptance can still make
DSpark slower. Temperature 1 uses the default opportunistic policy; these are
not exact-sampling results. DSpark does not accelerate prefill.
To reproduce the greedy coding case after downloading `ds4f-q2` and
`ds4f-dspark`:
```sh
./ds4 --cuda -m ds4flash.gguf \
--dspark --mtp-model gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf \
--ctx 4096 --prefill-chunk 512 --nothink --temp 0 --seed 12345 -n 256 \
-p 'Write a complete C hash table implementation with string keys, insert, find, delete, and a test main. Output only C code.'
```
For the temperature-1 rows, use `--temp 1 --top-p 0.95 --min-p 0.05`.
The [QA record](../QA_BEFORE_RELEASES.md#16-speed-regression) includes the previous
implementation, continued-context checks and quality comparisons.
## Larger models and serving
CUDA also has [SSD streaming](SSD_STREAMING.md) paths for larger weights.
Memory fit and speed depend on the model layout; start with Q2 for normal use.
V4.1 Q2 SSD serving batches up to eight decode rows together. An eight-session
test reached 11.0 aggregate t/s versus 8.4 with ordered execution; this is
total throughput, not speed per client. Other single-Spark model paths retain
their existing scheduling.
See [serving](SERVER.md) and the recorded [benchmarks](PERFORMANCE.md).
docs/DISTRIBUTED.md
+208
View File
@@ -0,0 +1,208 @@
# Inference Across Machines
[README](../README.md)
There are two modes:
| Mode | Split | Main use |
| --- | --- | --- |
| Tensor parallelism | Routed experts and per-layer work across two Macs or two Sparks | Resident inference with lower per-token work on each GPU |
| Pipeline parallelism | Complete layer ranges across several machines | Fit larger models and overlap long prefills |
These are separate from [tensor parallelism across CUDA cards](CUDA_MULTI_GPU.md).
Network protocols have no authentication or encryption. Use trusted machines
and a trusted network; run the same commit on every peer. Model paths and
artifacts must agree. Update all TP peers together when changing versions.
## Tensor parallelism between two Macs
This is a 50/50 split with exactly one worker. Do not pass `--layers`.
Routed experts are sharded; attention partitioning depends on the model. Both
GPUs work on the same token and exchange partial results. This can reduce generation
latency, but the gain depends on the model, link, and comparison setup.
Two 128 GB Macs are useful for V4 Flash Q4/MXFP4 or GLM 5.3 Flash Q4.
GLM 5.2 IQ2_XXS is another tested capacity setup. V4.1 Flash Q2 also runs
on two 128 GB Macs, with disk-only Engram tables.
A larger quant may need larger machines even though its tensor layout is supported.
### Link setup
Use a Thunderbolt cable. RDMA requires an active verbs device with an
IPv4-mapped GID; a working ping alone does not establish that.
```sh
rdma_ctl status
ibv_devinfo -v
```
Addresses must be on the cabled member interfaces, not only the Thunderbolt
bridge. For example, after checking which interfaces are active:
```sh
# Machine A, example member interface en1.
sudo ifconfig en1 inet 10.99.0.2/30 alias
# Machine B, example member interface en6.
sudo ifconfig en6 inet 10.99.0.1/30 alias
```
For the large tested shards on otherwise idle 128 GB Macs, the setup raised
the per-boot GPU wired-memory limit on both machines:
```sh
sudo sysctl iogpu.wired_limit_mb=120000
```
This is specific to that memory configuration. It grants a larger GPU budget;
it does not create more RAM. Check model, context, and system headroom before
raising a limit on your machine.
### Start the pair
Download the same model on both machines. For GLM 5.3 Flash Q4:
```sh
./download_model.sh glm53-q4
```
Start the worker first; it retries while the coordinator loads:
```sh
# Machine B.
./ds4 --tensor-parallel --role worker \
--coordinator 10.99.0.2 9911 --transport rdma --ctx 8192
# Machine A.
./ds4 --tensor-parallel --role coordinator \
--listen 10.99.0.2 9911 --transport rdma --ctx 8192
```
The verbs device and GID are selected automatically. If ambiguous, specify
`--rdma-device` and `--rdma-gid-index` from `ibv_devinfo`. Use `--transport tcp`
on both peers when RDMA is unavailable. Do not probe the waiting coordinator
with `curl` or `nc`: it may treat the connection as a worker handshake.
Keep workers running in a terminal or managed session and retain both logs.
Do not treat repeated handshake or RDMA timeouts as successful QA merely
because a retry works.
The coordinator can be `ds4`, `ds4-agent`, `ds4-server`, or `ds4-bench`;
workers run `ds4`. For models with vision support, pass the same `--vision FILE`
to both for image input.
For GLM MTP, enable `--mtp` on both. For DeepSeek DSpark, both need the
matching support model and DSpark options. V4.1 supports vision but not
speculative decoding.
TP disk-cache restore currently rebuilds the exact saved token prefix on both
ranks rather than restoring the coordinator alone. Expect prefill on restore.
See [speculation](SPECULATIVE_DECODING.md) and [serving](SERVER.md).
## Tensor parallelism between two Sparks
V4 Flash 0731 Q2/MXFP4 and V4.1 Flash Q2 text inference support one GPU per rank,
with a 50/50 expert split. Each Spark holds about 45 GiB of Flash 0731 Q2
weights, 77 GiB of Flash MXFP4, or 81 GiB for V4.1, plus context and runtime
buffers. V4.1 Engram tables stay on disk. Do not add `--ssd-streaming` or
`--cuda-tensor-parallel`: those select different memory/execution modes.
Build the same commit with `make cuda-spark` on both machines and download
`ds41f-q2` on both. RDMA needs the libibverbs development headers at build time,
its runtime library, and an active RoCEv2 link. Check `ibv_devinfo -v` and use
the addresses of the directly connected ports, not the management network.
For example, with the coordinator at `172.31.250.1` on the direct link:
```sh
# Worker.
./ds4 --cuda -m gguf/DeepSeek-V4.1-Flash-Q2.gguf --ctx 32768 \
--tensor-parallel --role worker --coordinator 172.31.250.1 9911 \
--transport rdma
# Coordinator.
./ds4 --cuda -m gguf/DeepSeek-V4.1-Flash-Q2.gguf --ctx 32768 \
--tensor-parallel --role coordinator --listen 172.31.250.1 9911 \
--transport rdma
```
The link address selects the matching verbs device and GID. If selection is
ambiguous, use `--rdma-device` and `--rdma-gid-index`. TCP is also available
with `--transport tcp` on both peers.
For Flash 0731, download `ds4f-q2` or `ds4f-mxfp4` on both hosts and use its
GGUF path in the same commands. For DSpark, download `ds4f-dspark` on both and
add `--dspark --mtp-model PATH` to both commands. With MXFP4 over RDMA, both
GPUs also share the drafter work. This needs about 2.5 GiB extra per rank,
beside the support model and context. If the worker has no drafter, or the
configuration cannot use the split, drafting stays on the coordinator.
Speed depends on the prompt.
The coordinator can be `ds4-agent`, `ds4-server` or `ds4-bench`. Match context
sizes and any `--prefill-chunk` overrides on both sides. For V4.1, five or more
ready sessions use native batched decode; smaller groups run in order.
Use `--batched-session 8` on the server and allow memory for all eight contexts.
Flash 0731 network sessions currently run in order, without native batching.
CUDA network TP does not support vision, V4.1 DSpark or other model/quant layouts.
## Pipeline parallelism
Each process maps only its assigned layers, retaining that slice of the KV
state. Layer ranges are inclusive. `N:output` includes the final layer and
output head. Activations travel from one stage to the next over TCP.
For Flash Q4 on two machines, download `ds4f-q4` on both, then start each side.
Replace the example address with your coordinator's reachable address:
```sh
# Machine A.
./ds4 --role coordinator --layers 0:19 --listen 10.99.0.2 9911
# Machine B.
./ds4 --role worker --layers 20:output --coordinator 10.99.0.2 9911
```
Normally give the output head to the final worker. With several workers,
choose non-overlapping ranges covering the entire model. Workers register
their ranges with the coordinator; intermediate workers forward activations
directly to the next stage.
Long prefill chunks can occupy different stages simultaneously. A single
generation stream cannot use that overlap: each token must finish the route
before the next one is sampled. Use pipeline mode primarily for capacity and
long-prefill throughput, not as a guaranteed decode speedup.
### Full PRO Q4
For two 512 GB Mac Studios, use the split artifacts:
```sh
# Machine A.
./download_model.sh pro-q4-layers00-30
./ds4 -m gguf/DeepSeek-V4-Pro-Q4K-Layers00-30.gguf \
--role coordinator --layers 0:30 --listen 10.99.0.2 9911
# Machine B.
./download_model.sh pro-q4-layers31-output
./ds4 -m gguf/DeepSeek-V4-Pro-Q4K-Layers-31-output.gguf \
--role worker --layers 31:output --coordinator 10.99.0.2 9911
```
These downloads do not change `ds4flash.gguf`. Startup is expensive because
each side must make its model slice resident.
### Tuning and recovery
Keep default chunk sizes first. `--dist-prefill-window N` controls the number
of chunks in flight; `--dist-prefill-chunk N` overrides the session-derived
chunk size. `--debug` shows route and per-hop timings.
Activations use 32-bit transport by default. `--dist-activation-bits 16` halves
the payload; `8` is more aggressive. These change numerical precision on the
wire, not weights or KV storage. Validate output when changing them.
A disconnected worker invalidates the route. In-flight work can fail; later
requests need a complete route before proceeding, and the coordinator can
replay the saved token prefix to rebuild worker state. Pipeline snapshots
serialize all layer slices into one payload and redistribute them when loaded.
For protocol details, see [ds4_distributed.c](../ds4_distributed.c)
and [ds4_tp.c](../ds4_tp.c).
docs/METAL.md
+71
View File
@@ -0,0 +1,71 @@
# Metal on Apple Silicon
[README](../README.md) | [Getting started](../README.md#start-here)
## Build and run
Install Apple's command-line developer tools if needed:
```sh
xcode-select --install
```
From the repository root:
```sh
make
./download_model.sh ds4f-q2
./ds4
```
The same build supports M3 and M5 Macs. Hardware-specific fast paths are
selected automatically; no environment variable is needed to enable them.
Leave other GPU and memory-heavy applications idle when comparing performance.
## Choose a model
| Memory | Starting point |
| --- | --- |
| 64 GB | Flash Q2 with `--ssd-streaming` |
| 96 GB | Flash Q2; leave room for the context and other applications |
| 128 GB | Flash Q2, or GLM 5.3 Flash Q2 with modest initial context |
| 256 GB | Flash Q4/MXFP4 or GLM 5.3 Flash Q4 |
| 512 GB | Larger models, including PRO Q2 |
Here, Flash means DeepSeek V4 Flash. V4.1 Flash has different memory
requirements; see its [model guide](MODELS.md#deepseek-v41-flash).
These are starting points, not guarantees that every context or session count
will fit. GLM 5.3 Flash Q2 is about 90 GiB before runtime allocations.
Stop other memory-heavy workloads before loading it resident.
```sh
./download_model.sh glm53-q2
./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf --ctx 32768
```
For a model larger than RAM, start with automatic cache sizing:
```sh
./ds4 --ssd-streaming
```
See [SSD streaming](SSD_STREAMING.md) before increasing the expert cache.
Do not bypass the memory guard just to make an oversized resident model start.
## Two Macs
Two 128 GB Macs can run a larger model fully resident with a 50/50 routed-expert
split. Thunderbolt RDMA is the low-latency option; TCP is also supported.
Follow [tensor parallel setup](DISTRIBUTED.md#tensor-parallelism-between-two-macs).
For more than two machines, use [pipeline parallelism](DISTRIBUTED.md#pipeline-parallelism).
## Next steps
- [Vision and GLM models](MODELS.md)
- [DSpark and GLM MTP](SPECULATIVE_DECODING.md)
- [Batched serving](SERVER.md#multiple-sessions)
- [Benchmarking](PERFORMANCE.md)
For DeepSeek V4 Flash and PRO, `--power 70` trades throughput for lower
sustained GPU load. V4.1 and GLM currently require `--power 100`.
docs/MODELS.md
+273
View File
@@ -0,0 +1,273 @@
# Models and Vision
[README](../README.md) | [Getting started](../README.md#start-here)
DwarfStar is not a general GGUF runner. Use the download targets below: other
GGUFs may have unsupported tensor layouts, metadata, or quantization mixes.
Run `./download_model.sh --help` for filenames and all available targets.
Main-model downloads update `ds4flash.gguf`. Encoders, draft models, packaged
FP8 weights, and split PRO pieces do not. Pass `-m FILE` to avoid depending on
which model was downloaded last.
Some downloads require the Hugging Face CLI; the script prints installation
instructions when needed. Authentication is optional for public weights;
your cached Hugging Face token or `HF_TOKEN` is used when present.
## DeepSeek V4
| Target | Use |
| --- | --- |
| `ds4f-q2` | Flash 0731, about 81 GiB; starting point for 96/128 GB systems |
| `ds4f-q2-q4` | Mostly Q2, with the last six routed-expert layers at Q4; needs more memory |
| `ds4f-q4` | Flash 0731 Q4; larger-memory or distributed systems |
| `ds4f-mxfp4` | Native MXFP4 routed experts; larger-memory or distributed systems |
| `pro-q2-imatrix` | PRO 0813; 512 GB resident target, or SSD streaming |
| `pro-q4-split` | Both PRO Q4 pieces for pipeline execution |
```sh
./download_model.sh ds4f-q2
./ds4
```
The Flash Q2 recipe spends most of its compression on routed experts:
IQ2_XXS gate/up and Q2_K down. Other components use higher precision, including
Q8 projections, shared experts and output, plus F16/F32 tensors. They are not
all untouched source weights. The imatrix guides the routed quantization.
The MXFP4 recipe preserves DeepSeek's released MXFP4 routed experts without
requantizing them. Metal and CUDA support it. Blackwell uses native FP4 matrix
instructions and FP4 activations for batched expert work; CUDA decode and
other CUDA architectures use Q8 activations. ROCm also has a resident MXFP4
path, including pipeline execution for models too large for one host.
To build weights rather than download them, see [GGUF tools](../gguf-tools/README.md).
## DeepSeek V4.1 Flash
V4.1 Flash text and vision inference work on Metal. CUDA supports text with
Q2 SSD streaming on one Spark or resident shards across two Sparks. ROCm
supports single-host text inference with SSD streaming. It needs its own
GGUF, tokenizer and inference graph; V4 Flash weights and DSpark support files
are not interchangeable with it.
| Target | File size | Main weights |
| --- | ---: | ---: |
| `ds41f-q2` | 341 GiB | 152 GiB |
| `ds41f-q4` | 483 GiB | 294 GiB |
| `ds41f-mxfp4` | 467 GiB | 278 GiB |
Q2 and Q4 use imatrix-calibrated routed experts; MXFP4 uses OCP MXFP4 routed
experts with higher-precision attention, shared experts and output weights.
All include 189 GiB of Engram tables. Engram rows are read directly from the
file as needed in every mode, never loaded as a resident table. Keep the GGUF
on a fast local SSD.
On one 128 GB Mac, use SSD streaming. Leave the expert cache budget automatic:
```sh
./download_model.sh ds41f-q2
./ds4 -m gguf/DeepSeek-V4.1-Flash-1017-IQ2_XXS-Q2_K-imatrix.gguf \
--ssd-streaming --ctx 32768
```
Use `ds4-agent` or `ds4-server` with the same model and memory options.
On a DGX Spark, add `--cuda`; see the [Spark guide](DGX_SPARK.md#deepseek-v41-flash).
`--think-level 25` sets reasoning effort explicitly; the range is 1 to 100,
with 0 disabling thinking. `/think 25` changes it in the CLI or native agent.
`--think` selects 75 and `--think-max` selects 100.
For two 128 GB Macs or Sparks, follow the [TP/RDMA setup](DISTRIBUTED.md), passing this
GGUF with `-m` on both ranks and omitting `--ssd-streaming`. Each rank holds
about 81 GiB of main weights, plus context and runtime buffers. Both machines
need the complete GGUF on disk. A 256 GB or larger Mac can instead hold all
main weights; full residency has been tested on an M3 Ultra with 512 GB.
Q4 needs SSD streaming on smaller Macs, or a 512 GB Mac for full residency.
It does not fit resident TP across two 128 GB Macs. Download it with
`./download_model.sh ds41f-q4` and select `gguf/DeepSeek-V4.1-Flash-Q4.gguf`.
The download comes in two parts; the script joins and verifies them automatically.
Allow another 37 GiB of free disk space while joining. Rerun the command to resume an
interrupted download or join.
For the 1017 MXFP4 build, run:
```sh
./download_model.sh ds41f-mxfp4
```
The target downloads all nine standard GGUF shards, verifies each size and
SHA-256, and merges them into `gguf/DeepSeek-V4.1-Flash-1017-MXFP4-experts.gguf`.
DwarfStar's V4.1 loader needs this single file with page-aligned Engram tables;
the script restores the original 16 KiB alignment and verifies the original
whole-file checksum before updating `ds4flash.gguf`. The shard files are kept,
so allow about **935 GiB** for shards plus the merged output. Free space is
checked before downloading. Rerun the same command to resume downloads or an
interrupted merge; existing files and the previous default model are preserved
if verification fails. Q2 and MXFP4 use the 1017 checkpoint; Q4 remains the
previous release.
Large SSD prefills process layers in wide batches. Metal overlaps computation
with the next layer's reads; CUDA stages experts into its bounded device cache.
Short appends keep using the expert cache.
Resident and TP inference also batch continued prefills automatically.
For concurrent serving, see [session batching](SERVER.md#multiple-sessions).
Each slot needs its own context memory; start with `--ctx 4096` before
increasing both context and slot count. CUDA Q2 SSD mode batches up to eight
decode rows; CUDA network TP currently serves sessions in order.
[DSpark](SPECULATIVE_DECODING.md#deepseek-v41-flash-dspark) uses a
separate V4.1 support GGUF and currently requires one host. Pipeline execution
is not implemented for V4.1; vision requires Metal.
Scalar, batched and tensor-parallel execution are not numerically identical.
Q4 batched prefill shows a small probability-score loss on the short official
continuation test, with unchanged overall top-token agreement. See the
[QA results](../QA_BEFORE_RELEASES.md#17-deepseek-v41-flash) for details
and remaining differences.
For images, download the matching encoder and add it to the same command:
```sh
./download_model.sh ds41f-vision
./ds4-agent -m gguf/DeepSeek-V4.1-Flash-1017-IQ2_XXS-Q2_K-imatrix.gguf \
--ssd-streaming --vision gguf/DeepSeek-V4.1-Flash-Vision.gguf
```
Vision works with SSD streaming, full residency and two-Mac TP. Pass the encoder
on both TP ranks. Use `/read image.png` in `ds4`, `view_image` in `ds4-agent`,
or the [server image API](SERVER.md#images). V4 Flash vision encoders do not
work with V4.1. See [conversion](../gguf-tools/README.md#convert-deepseek-v41-flash)
to build the GGUFs from safetensors.
## Qwen3.8 Flash Next
`./download_model.sh qwen38-q2` downloads one **137.10 GiB** GGUF: 41.73 GiB
of main/MTP weights and 95.37 GiB of original BF16 n-grams kept on disk.
Its gate/up experts use IQ2_XXS; down experts use Q2_K with 640 logical inputs
padded to 768 in the weight file. It replaces the larger MXFP4-down IQ2 release.
For 64 GB Macs, start at 8K context with a 1,024-token prefill chunk; runtime
allocations add to the main weights, but the n-gram table is not mapped or
preloaded. Keep the GGUF on a fast local SSD.
The larger `qwen38-q4k` target uses 165.11 GiB on disk and 69.74 GiB for
resident weights, before runtime buffers.
This model runs on Metal and single-GPU CUDA, including DGX Spark.
The script links `ds4flash.gguf` to the combined GGUF:
```sh
./ds4 --ctx 8192 --prefill-chunk 1024
```
Add `--mtp` for speculation; no second file is needed.
See [Qwen setup](QWEN38_FLASH_NEXT.md)
for memory, conversion, vision, and sampling details.
Vision uses a separate encoder. `./download_model.sh qwen38-vision` downloads
llama.cpp's Q8_0 mmproj from
[ggml-org/Qwen3.8-Flash-Next-GGUF](https://huggingface.co/ggml-org/Qwen3.8-Flash-Next-GGUF);
pass it at runtime with `--vision`.
## GLM 5.3 Flash
| Target | Approximate file size | Use |
| --- | ---: | --- |
| `glm53-q2` | 90 GiB | One 128 GB Mac or DGX Spark; ROCm also supported |
| `glm53-q4` | 178 GiB | Larger Mac, two 128 GB Macs, or SSD streaming |
| `glm53-fp8` | 305 GiB | Packaged native weights only; inference not implemented |
```sh
./download_model.sh glm53-q2
./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf --ctx 32768
```
GLM 5.3 Flash has recurrent KDA layers, sparse DSA attention, hyper-connections,
and a built-in MTP block. The Q2 file uses imatrix-guided IQ2_XXS gate/up and
Q2_K down experts. Q4 is the higher-precision alternative.
Q2 is close enough to a 128 GB machine's memory budget that other workloads
and context size matter. Follow the [Metal](METAL.md), [Spark](DGX_SPARK.md),
or [Strix Halo](STRIX_HALO.md) starting configuration for your host.
Ordinary decode is the default. Enable the embedded draft block with `--mtp`:
```sh
./ds4-agent -m gguf/GLM-5.3-Flash-Q2.gguf --mtp --ctx 50000
```
No second model file is needed. See [sampling behavior](SPECULATIVE_DECODING.md)
before choosing between the default opportunistic mode and exact sampling.
## Full GLM 5.3 and GLM 5.2
Full GLM 5.3 Q2 is about 197 GiB. Use a sufficiently large machine or streaming:
```sh
./download_model.sh glm53-full-q2
./ds4 --ssd-streaming
```
GLM 5.2 downloads are `glm-antirez-iq2xxs`, `glm-antirez-q2`,
`glm-antirez-q4`, and the 11-shard `glm-unsloth-q4`. They are much larger
than GLM 5.3 Flash; choose memory capacity before choosing the quantization.
GLM runs on Metal, CUDA and ROCm. The routed paths include IQ2_XXS, Q2_K, and
Q4_K, with additional mixed layouts supported by the tested GGUFs. Two-Mac
ownership-aware TP accepts IQ2_XXS, Q2_K, and Q4_K gate/up layouts; this does
not mean every Q4 model fits two 128 GB machines. Use a tested artifact, not
an arbitrary combination of supported tensor types.
GLM uses graph-selected prefill chunks and does not accept `--prefill-chunk`
or an external `--mtp-model`. It currently requires `--power 100`.
Directional steering is supported for GLM 5.3, not GLM 5.2.
## Vision
PNG and JPEG input works in the CLI, native agent, and HTTP server on Metal,
single-GPU CUDA, and ROCm. The encoder must match the model.
V4.1 Flash vision is currently Metal-only; its setup is [above](#deepseek-v41-flash).
### DeepSeek Flash Vision Experimental
Vision Experimental is a different language checkpoint from Flash 0731.
The main download includes its encoder:
```sh
./download_model.sh ds4f-vision-q2
./ds4 --vision gguf/DeepSeek-V4-Flash-Vision-Encoder.gguf
```
Larger targets are `ds4f-vision-q2-q4` and `ds4f-vision-mxfp4`.
`ds4f-vision-encoder` downloads just the encoder when the language GGUF is
already present. Vision Experimental has its own [DSpark drafter](SPECULATIVE_DECODING.md).
### GLM 5.3 Flash
The text GGUF stays the same. Download and add the encoder explicitly:
```sh
./download_model.sh glm53-vision
./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf \
--vision gguf/GLM-5.3-Flash-Vision-Encoder.gguf
```
Use `/read image.png` in `ds4`, or start `ds4-agent` with the same `--vision`
argument to enable `view_image`. Agent sessions containing images cannot yet
be saved with `/save`.
For two-Mac TP, pass the same encoder on both ranks. The coordinator encodes
the image and sends the projected visual tokens to the worker.
For HTTP image formats and limits, see [serving](SERVER.md#images).
### Qwen3.8 Flash Next
The text GGUF stays the same. Download and add the encoder explicitly:
```sh
./download_model.sh qwen38-vision
./ds4 --mtp \
--vision gguf/mmproj-Qwen3.8-Flash-Next-Q8_0.gguf
```
The encoder is llama.cpp's Q8_0 mmproj conversion of the model's Qwen3-VL
tower. Use `/read image.png` in `ds4` or `image_url` parts over HTTP; see
[Qwen setup](QWEN38_FLASH_NEXT.md) for image limits and resize behavior.
docs/PERFORMANCE.md
+72
View File
@@ -0,0 +1,72 @@
# Performance and Benchmarking
[README](../README.md)
Compare the same checkpoint, quantization, context, and sampling settings.
Record the commit and whether weights were resident, streamed, or distributed.
Keep other GPU workloads idle and repeat in alternating order: one favorable
run is not a speed result.
## Context sweeps
`ds4-bench` measures prefill and generation at successive context frontiers:
```sh
./ds4-bench -m ds4flash.gguf \
--prompt-file speed-bench/promessi_sposi.txt \
--ctx-start 2048 --ctx-max 65536 --step-incr 2048 --gen-tokens 128
```
Each prefill number measures the newly added interval. Generation uses a fixed
greedy, non-EOS probe. The benchmark normally restores a memory snapshot after
each probe; network TP/pipeline runs and snapshots beyond its memory limit use
prefix replay instead. Do not interpret replay time as continued-prefill speed.
Use `--step-mul F` for exponential context spacing. Output is CSV, including
prefill throughput, generation throughput, and snapshot size when available.
The prompt is the cleaned public-domain *I Promessi Sposi* text in
[speed-bench](../speed-bench/README.md).
Prefill defaults are configuration-specific: ordinary DeepSeek uses 4096-token
chunks, CUDA TP uses 2048, and long PRO prompts may use 8192. GLM chooses its
own chunks and rejects `--prefill-chunk`. The strict DeepSeek API-vector test
pins 2048; do not generalize that setting to every benchmark.
## Recorded Flash Q2 baseline
These existing sweeps use 2048-token intervals and 128 generation tokens per
frontier. They are recorded baselines, not measurements of every subsequent
commit. Full data: [M5 Max](../speed-bench/m5_max.csv) and
[DGX Spark](../speed-bench/gb10.csv).
| Machine | Context | Prefill | Generation |
| --- | ---: | ---: | ---: |
| M5 Max, 128 GB | 2048 | 790.18 t/s | 39.35 t/s |
| M5 Max, 128 GB | 16384 | 572.53 t/s | 36.14 t/s |
| M5 Max, 128 GB | 32768 | 557.04 t/s | 34.36 t/s |
| M5 Max, 128 GB | 65536 | 398.50 t/s | 27.64 t/s |
| DGX Spark, 128 GB | 2048 | 825.76 t/s | 18.05 t/s |
| DGX Spark, 128 GB | 16384 | 872.44 t/s | 15.10 t/s |
| DGX Spark, 128 GB | 32768 | 855.94 t/s | 14.43 t/s |
| DGX Spark, 128 GB | 65536 | 822.98 t/s | 13.84 t/s |
![M5 Max Flash Q2 throughput](../speed-bench/m5_max_ts.svg)
Historical PRO Q2 measurements on the M3 Ultra are retained in this chart:
![PRO Q2 on M3 Ultra](../speed-bench/pro_model_m3_ultra_ts.svg)
## What to compare next
- For SSD streaming, record the effective cache and distinguish cold startup
from a warm cache.
- For multiple sessions, report both individual latency and aggregate
throughput. Ordered fallback is not native batching.
- For DSpark/MTP, compare plain decode too. A faster drafter path can still be
slower than ordinary decoding on an unpredictable prompt.
- For TP, keep quantization and prompt equal; comparing resident Q4 on two
machines with streamed Q4 on one measures capacity benefits as well as parallelism.
Recent focused TP and DSpark comparisons, including their limitations, live in
[QA_BEFORE_RELEASES.md](../QA_BEFORE_RELEASES.md). Use its speed-regression
procedure rather than accumulating one-off timings in the main README.
docs/QWEN38_FLASH_NEXT.md
+551
View File
@@ -0,0 +1,551 @@
# Qwen3.8 Flash Next
[Back to README](../README.md)
Qwen3.8-Flash-Next uses the `qwen4exp` GGUF architecture and a dedicated
Metal/CUDA graph with gated delta-net, gated GQA, block-sparse attention,
hyper-connections, n-gram embeddings, MoE, and MTP.
## Download and run
```sh
make
./download_model.sh qwen38-q2
./ds4 --ctx 8192 --prefill-chunk 1024
```
The download installs one GGUF containing the main model, MTP and original
BF16 n-grams, and updates `ds4flash.gguf`:
| Target | File size | Resident weights |
| --- | ---: | ---: |
| `qwen38-q2` | 137.10 GiB | 41.73 GiB |
| `qwen38-q4k` | 165.11 GiB | 69.74 GiB |
The 95.37 GiB n-gram table stays on disk. The runtime reads selected rows
directly from the GGUF, without mapping or preloading the table. Keep the
file on a fast local SSD. No sidecar or n-gram option is needed.
The Q2 model uses IQ2_XXS gate/up experts and Q2_K down projections padded
from 640 to 768 columns, calibrated with an imatrix. Q4 uses calibrated
Q4_K gate/up and MXFP4 down experts. Context and runtime buffers still need
RAM beyond the resident weights; start Q2 with the context and chunk above
on a 64 GB machine. Q4 needs a larger machine.
On Metal, SSD streaming keeps a bounded cache of routed experts instead of
requiring all of them in RAM:
```sh
./ds4 --ssd-streaming --ctx 4096 --prefill-chunk 128
```
The Q2 and Q4 packs support streaming during prefill, decode and MTP. Dense
weights, shared experts, context and runtime buffers still need RAM; the
n-gram table continues to use its existing disk reads. Missing routed experts
are read into owned buffers. If a prefill batch needs more experts than the
cache can hold, or an MTP layer uses a different expert size, it uses explicit
temporary staging while retaining the normal kernel arithmetic. On devices
with GPU buffer addresses, this staging contains only the distinct selected
experts. Selecting every expert, or running without buffer-address support,
retains the full-layer fallback.
During single-token decode, cache misses overlap with gate/up computation
for cached experts and the shared expert. Once the reads finish, the missing
gate/up slots are computed, followed by one down pass and the original ordered
reduction. Private address snapshots keep the running GPU work independent of
cache updates. Fully cached selections use one gate/up pass. Cached batched
prefill uses overlap on every Metal device: resident experts run while missing
weights load, then missing gate/up computation overlaps the down-weight reads.
Temporary staging retains its existing schedule; an MTP layer whose
expert size differs from the main model is staged again on each invocation,
without replacing the main model's cache entries.
A local M1 Max 32 GiB check compared this compact split schedule with
`85d37ae`, using the Q2 pack, automatic expert cache, `--ctx 4096`,
`--prefill-chunk 128`, `-n 100`, `--temp 0 --nothink`, and the prompt
`narrami la storia di roma`. Three alternating pairs produced:
| Decode | Runs (t/s) | Median (t/s) |
| --- | --- | --- |
| Before | 8.40, 8.25, 7.26 | 8.25 |
| Compact split | 8.44, 8.51, 8.14 | 8.44 |
All generated output was identical. The observed median gain was 2.3%; the
variation between runs does not establish a 10% improvement. Prefill medians
were 9.98 and 10.05 t/s; its execution path is unchanged. These are local
warm-file measurements, not a guarantee for other Macs or cold SSD reads.
Q4_K gate/up shares activation reads between the shared expert's Q8
projections and uses packed activation/scale loads while preserving the
original FP32 accumulation order. On M1 Max, the 640-to-2560 MXFP4 down
projection also enables prefetch by default. `DS4_QWEN4_MOE_DOWN_PREFETCH=0`
restores the plain down kernel for diagnostics. Safe math always retains that
plain kernel, including when prefetch is explicitly requested, because its
noncontracted arithmetic differs from the prefetched kernel's explicit FMA.
A local Q4 SSD check on M1 Max 32 GiB used `--ctx 32768`,
`--prefill-chunk 1024`, `--temp 0 --nothink`, `-n 100`, and the prompt
`Narrami in italiano la storia di Roma`. Three alternating pairs comparing
the previous kernels with these changes produced:
| Decode | Runs (t/s) | Median (t/s) |
| --- | --- | --- |
| Before | 6.02, 6.07, 6.26 | 6.07 |
| Q4 kernel updates | 6.42, 6.03, 7.27 | 6.42 |
All generated output was identical. The observed median gain was 5.8%; SSD
and system variation remained substantial. Prefill medians were 6.94 and
6.91 t/s. These changes retain the existing prefill and row-dispatch policies;
Q2's smaller sparse-prefill tiles did not show a stable benefit for Q4 in
local comparisons. The measurements do not establish gains on other devices.
MXFP4 prefill also decodes each 16-value half-block with one exponent read
and packed byte loads, retaining the original multiplication and final
half/float conversion. Compared with `46425263`, alternating resident-kernel
benchmarks on M1 Max at 2048 tokens and 10 routed slots reduced median
gate/up-plus-down time from 70.65 to 66.31 ms with 32 experts and from 81.57
to 77.34 ms with 256 experts (6.1% and 5.2% less time).
The staging oracle checks all exponent bytes, nibble codes and byte
alignments against two original 8-value decoders in both math modes.
On M1 Max, batched MXFP4 down rows also reuse activations across two output
rows. A function constant enables this only for multi-token NR2 dispatches;
single-token decode and NR1/NR4 retain the original row loop. A warmed
same-process comparison with 16 alternating samples per variant, 128 calls
per sample, width 640, 2560 output rows and 10 routed slots measured:
| Two-token down kernel | Before GPU time | After GPU time |
| --- | --- | --- |
| Q8 shared expert | 176.62 us | 150.70 us |
| MXFP4 shared expert | 216.20 us | 188.48 us |
Outputs were bit-identical, and single-token GPU timings did not regress.
Full Q4 SSD runs remain sensitive to I/O: two alternating pairs at context
32768 and chunk 1024 gave prefill medians of 22.94 and 23.41 t/s on a
201-token prompt, with all 100 generated tokens unchanged. With MTP and two
draft tokens on the Rome prompt, median generation remained 7.65 t/s despite
the faster down kernel. These checks compare against `46425263`; they do
not imply the same percentage gain for complete inference or other devices.
Leave the expert-cache budget automatic initially. A plain count passed to
`--ssd-streaming-cache-experts` requests dynamic cache slots; an `NGB` budget
also includes the reserved layer-staging space. Startup accounts for the
context, prefill chunk, static weights, vision weights when present, and
staging before choosing the cache size. Reduce the context or prefill chunk
if those fixed requirements leave insufficient memory.
The total budget is the smaller of 87.5% of physical RAM and Metal's
recommended working set, with a further 2 GiB reserved for runtime overhead.
The recommendation is not discounted again, and admitted static weights stay
eligible for locking in RAM as the expert cache grows.
Qwen fills its expert cache on demand, so `--ssd-streaming-cold` does not
change its preload behavior. Popularity preloading and the GLM-specific
`--ssd-streaming-full-layers` option are not supported for Qwen.
On M1 Max, the shared half-operand MoE prefill path requires IQ2_XXS gate/up,
Q2_K down and at least 8192 actual tokens in one streamed batch. The configured
`--prefill-chunk` is only an upper limit: a 5760-token prompt with chunk 8192
does not activate this path. Comparing that prompt at chunks 2048 and 8192
measures the change in batching, not an isolated half-operand kernel gain.
On a DGX Spark, build with `make cuda-spark` instead of `make`. Both Q2 and
Q4 fit resident on an otherwise idle 128 GB Spark; the n-grams stay on SSD.
The same commands support text, vision, MTP, steering and text-session
checkpoints in `ds4`, `ds4-agent` and `ds4-server`. No CUDA-specific model
conversion or feature flags are needed.
Leave `--prefill-chunk` unset on the Spark to use the faster 8192-token
default. The 1024-token example above saves memory on smaller Macs.
Use the same model options with `ds4-agent` or `ds4-server`.
Add `--mtp` for speculative decoding using the built-in MTP weights.
`DS4_QWEN4_MTP_PREFILL` controls how prefill prepares the predictor’s prompt history.
| Value | Predictor history for new prompt rows |
| --- | --- |
| `auto` (default) | Prepare it only with SSD streaming |
| `on` (also `1`) | Prepare it for both resident and SSD streaming runs |
| `off` (also `0`) | Skip historical projections and initialize those cache rows to zero |
To prepare the full predictor history when the weights fit in RAM, run
```sh
DS4_QWEN4_MTP_PREFILL=on ./ds4 -m gguf/Qwen3.8-Flash-Next-Q2.gguf --mtp --ctx 8192
```
Use `on` for resident workloads where the decode
benefit outweighs the prefill cost, and compare total wall time on your prompts.
`off` and resident `auto` keep MTP drafts and target verification enabled, but
the cold predictor history can change drafts and reduce their acceptance rate.
Zero keys still participate in predictor attention. Target prefill caches and logits are unchanged.
Decode continues to prepare real predictor rows. Each prefill chunk also
retains its final row for real preparation at the next boundary, keeping
checkpoint restore consistent with uninterrupted prefill.
The setting applies to newly prefilled rows. Reused or restored checkpoints
retain their existing predictor cache contents.
Without `--mtp`, this setting has no effect.
`ds4-bench` benchmarks ordinary decoding; it does not accept `--mtp`.
`--nothink` disables thinking. The server exposes
`qwen3.8-flash-next`, `qwen3.8-flash-next-chat`, and
`qwen3.8-flash-next-reasoner` aliases. Tool calls use the native
`<tool_call><function=...><parameter=...>` format.
With resident weights on Metal, `ds4-server --batched-session N` decodes
the slots together. Work that
only reads weights runs once for the whole batch, while the delta-net
recurrence, the attention caches and the n-gram convolution stay per session.
The slots also share one prefill workspace, so an extra slot costs its caches
rather than another few GiB of transients. As with the other natively batched
models, grouping changes the order of floating point reductions, so a batched
reply can differ from the same prompt decoded alone.
Add `--mtp` to use batched speculative decoding as well. The server accepts
matching greedy drafts by default, including at nonzero temperature. With
`--mtp-exact-sampling`, sampled requests use ordinary batched decoding;
temperature-zero requests can still speculate. Default greedy acceptance also
means a seed need not reproduce a reply under a different batching schedule.
Speculative batches above 16 sessions, images and steering use the ordered
fallback. Metal SSD expert streaming also uses the ordered fallback and
keeps each session's streaming workspace; it does not use the resident batch
arena. CUDA currently decodes sessions in order.
Disk KV checkpoints include recurrent state. Rewinding to an earlier position
replays the retained prefix on the next evaluation. The native context is
262144 tokens; `DS4_QWEN4_YARN_FACTOR=2` or `=4` enables static YaRN for
longer contexts, with a possible quality cost on shorter prompts.
## Metal tuning on M1 Max and M2-M4
The established M1 Max simdgroup optimizations also apply to Apple M2, M3
and M4, including their Pro, Max and Ultra variants, within the same tensor
shapes, batch sizes, math modes and diagnostic overrides:
- Wide Q8 decode projections, HC gate/mix reuse and short MTP Q8 projections.
- IQ2 gate/up, Q2_K down-row reuse and the Q4 model's MXFP4 down prefetch.
- Small HC-down and HC-normalization prefill batches, and large Q8 SSD prefill.
- Sparse IQ2/Q2 SSD tiles, Q4/MXFP4 K32 tiles, optional half-operand staging
for large SSD batches and release of that scratch before Q4 decode.
M3 Ultra retains its dedicated MoE specialization, row/group geometry and
prefill tile policy; the other shared Q8, HC and MTP choices still apply.
M5/M6 retain their existing policy. The selection does not expand to M1
base/Pro/Ultra or to future generation names.
This widens device selection, not the kernel arithmetic or GGUF formats.
Local validation uses M1 Max. An independent M4 Max Q2 resident test reported
about 7% more decode throughput after widening the decode gates, with matching
outputs; it did not validate M2/M3, other M4 variants, or SSD performance.
See the [M4 Max test report](https://github.com/antirez/ds4/pull/1056#issuecomment-5922543210).
Device eligibility therefore does not imply a measured speedup on every chip.
## DGX Spark performance
Measured on a single Spark with resident weights and disk-only n-grams:
| Model | First 1024 tokens | Next 7168 tokens | Decode at 8192 tokens |
| --- | ---: | ---: | ---: |
| Q2 | 516 t/s | 745 t/s | 22.6 t/s |
| Q4 | 513 t/s | 755 t/s | 21.2 t/s |
These are averages of two runs using `speed-bench/promessi_sposi.txt`,
8192-token prefill chunks and 128 teacher-forced decode tokens, without MTP.
Loading the model is excluded. Q2 also reached about 771 t/s on a fresh
32K-token prefix with `--prefill-chunk 32768`. Small continuations have lower
throughput: adding 32 tokens after an 8K Q4 prefix took about 273 ms.
MTP speed depends on how often drafts are accepted; the table measures
ordinary decoding.
## Conversion
See [GGUF conversion](../gguf-tools/README.md) for conversion from the HF
checkpoint and calibrated low-bit expert quantization.
`gguf-tools/qwen4_pack_to_qwen4exp.py` migrates old DS4 fast packs to the
active `qwen4exp` schema. Finish older main-only conversions with
`gguf-tools/qwen4_native_ngrams.py`, using the original HF n-gram shards.
Expanding an old quantized table to BF16 does not restore its lost precision.
## Vision
Images go through the model's Qwen3-VL vision tower. Download the encoder
with `./download_model.sh qwen38-vision`, then use
`--vision gguf/mmproj-Qwen3.8-Flash-Next-Q8_0.gguf`. The CLI accepts
`/read image.png`; the server accepts `image_url` parts. Each image is resized to
multiples of 32 pixels within 64 to 1024 tokens (`DS4_QWEN4_IMAGE_MAX_TOKENS`
raises the cap), encoded on the GPU, and takes the model's 3D rope positions;
live KV reuse keys on the image fingerprints. `make test-qwen4-vision`
compares the tower with the Hugging Face implementation using the same GGUF
weights, dequantized to float32. It separately reports GGUF-versus-original
checkpoint quality and GPU-versus-original agreement. Both comparisons use
the unchanged minimum per-token cosine threshold of 0.99; implementation
parity gates the default exit status. Add `--require-quality` to also fail on
GGUF-versus-original quality loss. A passing implementation check alone does
not mean the quantized encoder matches the original checkpoint.
For a small suite including OCR, diagrams, a photograph, and resizing:
```sh
make tests/test_qwen4_vision
uv run --with numpy --with torch --with torchvision --with pillow \
--with safetensors --with transformers --with gguf tests/qwen4_vision_ref.py \
--snapshot /path/to/HF-checkpoint \
--mmproj gguf/mmproj-Qwen3.8-Flash-Next-Q8_0.gguf \
--image tests/vision-fixtures/qwen38/orbit.png \
--image tests/vision-fixtures/qwen38/maple.png \
--image tests/vision-fixtures/glm53/diagram.png \
--image tests/vision-fixtures/glm53/text.png \
--image tests/vision-fixtures/glm53/earth.jpg \
--image tests/vision-fixtures/glm53/screenshot.png \
--json-report /tmp/qwen38-vision-parity.json
```
Transformers must include `qwen4_exp`; the reference runs on CPU. Repeat
`--image` to add cases. The original single-image `--out` embedding dump is
still supported. Metric checks without model weights run with
`uv run --with numpy python -m unittest discover -s tests -p test_qwen4_vision_ref.py`.
JPEG decoder agreement with Pillow, including progressive scans, subsampling,
and image edges, is checked separately with
`uv run --with numpy --with pillow python -m unittest discover -s tests -p test_jpeg_decode.py`.
To check CLI `/read` image turns and text follow-ups with ordinary and MTP
decode, use two different PNG or JPEG images with the model-backed regression:
```sh
make ds4
uv run tests/test_qwen4_cli_vision.py --model /path/to/main-with-mtp.gguf \
--vision gguf/mmproj-Qwen3.8-Flash-Next-Q8_0.gguf \
--image /path/to/first.png --image /path/to/second.png
```
The test checks that all four turns complete in each mode, including errors
that the interactive CLI can report without a nonzero process exit. It saves
the responses and diagnostics for inspection; it does not grade image content.
The Metal and CUDA graphs accept Q8_0, Q4_0, F16, BF16 and F32
dense weights, Q8_0/MXFP4/Q4_0/Q4_K/Q2_K/IQ2_XXS experts, F16/F32/Q8_0
hyper-connection mixers and the original BF16 n-gram table.
Tensor parallelism and pipeline execution are not implemented for this model
yet. SSD expert streaming is supported on Metal only. ROCm is not supported.
CPU code is a correctness reference, not a general inference backend.
M6 uses the same Metal dispatch defaults as M5, including the Qwen indexer,
attention merge, HC mixers and MoE scheduling. Existing M5 diagnostic rollback
switches apply to both generations; tensor kernels still require runtime
capability support. The F16 HC mixers explicitly preserve their accumulation
order in both normal and safe shader math modes. This does not guarantee
bit-identical logits across GPU generations: tensor matmuls retain their own
accumulation order.
## Validation
```sh
make test-qwen4-kernels test-qwen4-q2 test-qwen4-prefill-reuse test-q8-prefill-variants
make test-qwen4-ssd-experts test-qwen4-memory
make test-qwen4-hc-math test-metal-device-policy
make test-frontends
make test-qwen4-ngrams
make tests/test_qwen4_ngram_state
make -B -C gguf-tools quants-shared
python3 -m unittest discover -s gguf-tools/tests -p test_qwen4_pack.py
python3 -m unittest discover -s gguf-tools/tests -p test_qwen4_native_ngrams.py
```
The kernel tests exercise the active Metal API. Vision and end-to-end model
checks additionally require the checkpoints described above.
The SSD test needs no GGUF: a sparse temporary file uses a 512-expert table,
including indices 383, 384 and 511, with cold and warm selections, eviction,
prefill staging, different-size MTP layers and failed reads followed by exact recovery. It
compares routed intermediates, reduced outputs and hyper-connection residuals
bit for bit against the resident kernels, including the padded Q2 down rows.
Its test-only backend checks actual `pread` requests against selected expert
ranges and counts returned bytes: ten experts for off-size MTP, 24 for cache
overflow, zero reads for cache hits, and the full 512-expert fallback.
Single-token tests also cover partially cached selections, duplicate IDs,
optional shared experts and all 16 routed slots. A failed SSD read must leave
the completed cached/shared gate/up slots exact, preserve the untouched slots,
and allow an exact retry after draining GPU work.
The memory-estimate test runs without a GPU or model.
To compare ordinary one-shot generation with the session path at identical
prefill boundaries, including every vocabulary logit during continuation:
```sh
make tests/test_qwen4_generation
./tests/test_qwen4_generation gguf/Qwen3.8-Flash-Next-Q2.gguf --ssd-streaming
# Exercise sparse attention and many chunk boundaries with a longer fixture:
./tests/test_qwen4_generation gguf/Qwen3.8-Flash-Next-Q2.gguf \
--ssd-streaming --chunk 128 --ctx 8192 --tokens 5760
```
Resident and SSD inference use the same arithmetic for prompt projections,
attention and shared experts. `make test-qwen4-resident-arithmetic` checks
these operations against the resident reference without loading a GGUF.
On a Mac that can hold the model, add `--prefill-reference` to the streamed
generation test to compare every vocabulary score with a resident run.
Omit `--ssd-streaming` on a machine that can keep the model resident. The
test checks explicit chunk precedence, progress frontiers, greedy tokens and
full logits. Different chunk sizes can change reduction order and routed
expert choices; compare the actual chunk boundaries before diagnosing token
drift. In particular, upstream `8db1d1d` ignores `--prefill-chunk` in the
ordinary one-shot Qwen path, although its sessions honor the option. For an
A/B comparison with that revision, also set `DS4_QWEN4_PREFILL_CHUNK` to the
same value in both binaries. This branch uses an explicit chunk in preference
to that environment setting and clamps it to the context size.
On CUDA, use `make test-qwen4-cuda` for the kernel tests. They compare the
active kernels with independent CPU references, without model weights.
Vision and end-to-end checks additionally require the checkpoints above.
Run `tests/test_qwen4_ngram_state MODEL.gguf` on either backend to check
failed disk reads during prefill, decode and MTP, then exact recovery.
For a smaller memory footprint, run
`tests/test_qwen4_ngram_state MODEL.gguf --ssd-streaming`. This uses a
256-token context, an eight-token prefill chunk and a 1024-expert cache request
while keeping both comparison sessions and MTP enabled.
To compare the predictor's cache-only prefill with a complete MTP pass:
```sh
make tests/test_qwen4_mtp_prefill
./tests/test_qwen4_mtp_prefill gguf/Qwen3.8-Flash-Next-Q2.gguf --ssd-streaming
```
This checks one-token prompts, chunk boundaries, non-argmax continuation,
stale cache contents and payload restore, followed by greedy, depth-three
and exact-sampling continuation checks. It uses a 128-token context and a
1024-expert cache request. The small-buffer MTP CLI regression also accepts
`--ssd-streaming`.
For the sparse-attention boundary and a 17408-token prefix, use the longer
predictor regression:
```sh
./tests/test_qwen4_mtp_prefill gguf/Qwen3.8-Flash-Next-Q2.gguf \
--ssd-streaming --long-context --prefill-chunk 2048 \
--dump-prefix /tmp/qwen-mtp-long-2048
```
Repeat with `--prefill-chunk 128` and a different dump prefix. The fixture
uses repeated, fixed chat token IDs, one graph, a 17416-token context and a
1024-expert cache request. It runs the trunk once and probes actual residual
rows across the first sparse query and at the final prefix. Each probe
compares cache-only preparation, the optimized last-row predictor and a
recursive draft with the frozen full-layer reference, using identical input
rows and lookahead tokens. It checks complete logits, cache contents, EH
projections, residuals and unchanged trunk state for T1/T2/T3 and exact T3
verification. Destination cache rows are poisoned before each replay.
Bit-identical comparisons are between the candidate and frozen reference
replayed from the same saved predictor state within each run. Equality between
chunk sizes is checked separately: chunk-dependent grouping can change the
arithmetic used to build the historical predictor cache.
`--dump-prefix` writes the fixture IDs, target/trunk captures, reference and
candidate predictor logits, and a JSONL manifest for comparisons between
builds. `--prefix-tokens N` changes the fixture length; it must be a multiple
of four covering the sparse boundary. These are implementation checks on a
synthetic sequence, not a quality evaluation or an acceptance-rate benchmark.
Comparing MTP acceptance between commits additionally requires the same
prompt, sampling settings and draft-depth policy. A trunk arithmetic change
can alter predictor inputs even when the predictor itself remains correct.
On an M1 Max with 32 GiB and Q2 SSD streaming, the 17408-token fixture passed
at chunks 128 and 2048 on `11811b9`, and at chunk 2048 on `04c0867` and
`89986c3`. At chunk 2048, all 36 captured tensors from `89986c3` and `11811b9`
were byte-identical. Differences from `04c0867` already appeared in the trunk
residuals after the attention fix in `89986c3`; both predictor implementations
followed those changed inputs.
Between chunks 128 and 2048 on `11811b9`, the two captured trunk slices and
final target logits were byte-identical. Predictor logits differed by at most
0.0001266003, with no argmax change in the 16 full/recursive comparisons. A
separate replay of identical rows as T3, T1+1+1, T2+1 and T1+2 found arithmetic
differences beginning at the HC down projection; its per-stage differences
were identical on `04c0867` and `11811b9`. This identifies an existing source
of grouping-dependent predictor cache values, not a demonstrated quality or
acceptance regression. The full-model runs above do not exercise the
8192-row half-operand path.
The independent `test_metal_qwen4_moe_half` comparison against the shader
source from `89986c3` also passed 31 fixtures, including 8192 rows, with Metal
API validation. Full shader validation exceeded the M1 threadgroup-memory
limit for the NT8 gate/up stress case (36864 bytes reported, 32768 available).
A separate NT1/2/4 subset passed with shader validation, including the
8192-row fixture at mid NT4/down NT4. These kernel checks do not replace an
8192-chunk full-model acceptance comparison.
Official Alibaba continuations are tracked for 100 short prompts and 12
archive/code prompts from 2K to 24K tokens. Build the quality scorer, then run
from the repository root:
```sh
gguf-tools/quality-testing/score_official MODEL.gguf \
gguf-tools/quality-testing/data/qwen38-flash-alibaba-100/manifest.tsv /tmp/qwen-short.tsv 4096
gguf-tools/quality-testing/score_official MODEL.gguf \
gguf-tools/quality-testing/data/qwen38-flash-alibaba-long/manifest.tsv /tmp/qwen-long.tsv 32768
```
Repeat with `--quality` and, for the long set, `--continued-prefill 1` and
`--continued-prefill 256`. These fixtures match the no-thinking template;
no rendered-prompt flag is needed. See [quality testing](../gguf-tools/quality-testing/README.md)
for collection settings, measurements and the hosted-checkpoint limitations.
[Checkpoint-fix benchmark charts and measurements](../speed-bench/qwen38-checkpoints/README.md)
compare prefill, ordinary decode, and MTP decode against the preceding PR head.
Qwen prefill checkpoints include matching logits at each completed chunk,
including cancellation frontiers. To test save/restore and continuation:
```sh
DS4_TEST_MODEL=/path/to/main-with-mtp.gguf \
DS4_TEST_GLM_MTP=1 DS4_TEST_MTP_EXACT=1 \
./ds4_test --qwen4-prefill-checkpoints --qwen4-restore-reuse \
--session-snapshot --session-rewind --session-rewind-resample
python3 tests/test_qwen4_checkpoint_replay.py \
--model /path/to/main-with-mtp.gguf
python3 tests/test_qwen4_mtp_limits.py \
--model /path/to/main-with-mtp.gguf
python3 tests/test_qwen4_logit_dump.py \
--model /path/to/main-with-mtp.gguf
```
The server regression checks continued disk saves, cache-budget eviction,
and text-answer history with and without tool schemas. It covers clients
that omit reasoning and clients that echo `reasoning_content`, including
server restart after a continuation. For tools-enabled histories, the first
answer keeps the exact disk key; after a client continues without reasoning,
the next checkpoint uses the visible history key for restart reuse.
The MTP tests cover small prefill buffers, rewinds, reused sessions, and
truncated checkpoints. The logit-dump test requires NumPy and compares all-row
prefill with teacher-forced decode across a chunk boundary. For official
continuation scoring, use `--rendered-prompt`
when a fixture already contains the complete model chat template.
`tests/test_qwen4_prefill MODEL PROMPT 8192` checks mixed prefill sizes,
progress callbacks and exact replay through the sparse-attention boundary.
It also reports differences against a fresh prefill followed by individual
decodes. Those schedules can round differently and exchange nearly tied
experts; their full logits need not match. Use model-quality checks as well,
not only state replay:
```sh
python3 tests/test_server_story.py --url http://127.0.0.1:8000 \
--model qwen3.8-flash-next --output /tmp/qwen-story
```
Start the server with at least `--ctx 49152`. This checks all sixteen facts
in a 31K-token story, then a correction turn that must reuse the prefix.
Qwen directional steering and activation capture support all 48 trunk layers;
see [directional steering](../dir-steering/README.md). Model-backed regressions:
```sh
python3 tests/test_qwen4_steering.py --model /path/to/main-with-mtp.gguf
./ds4_agent_test --full-context-save /path/to/main-with-mtp.gguf
make test-web-recovery
```
The agent can save an already-full text transcript without a KV payload.
Restart with a larger `--ctx`, then `/switch` to the saved session and
`/compact` to recover room. Sessions containing images still use the existing
save restriction. Automatic compaction reserves summary space and continues
interrupted generation through the existing agent compaction path.
docs/QWEN38_Q4_REGISTER_REUSE.md
+149
View File
@@ -0,0 +1,149 @@
# Qwen Q4_K raw-load register reuse
Measured on 9 October 2026: warmed fresh 16K prefill improved **2.38%** on the
standard text fixture and **2.27%** on C source, with exact old/new output checks.
## Change and provenance
The NAX gate/up and down lookahead loops previously reloaded a Q4_K header and
quant chunk for every 32-weight group. The new next-step loader retains the
header within each 256-weight superblock and the quant words across each
low/high-nibble pair. Initial loads and non-Q4 loaders are unchanged. The
existing dequantizer, arithmetic order, barriers and dispatch geometry are
unchanged. No additional buffer or runtime flag is introduced.
Inspired by the Q4 portion of [StarForge 08ea4d7](https://github.com/Chida82/sf-q3-8flash/commit/08ea4d7eb2e780e2d31690766600880e8dea8718).
Only this register reuse is adopted; the donor's other optimizations are excluded.
The measured host baseline was patched-source commit
`b51ae1ce23a59e5e5f7031c0144e76ac35f13e6c`. It already included
[PLE prefetch #1154](https://github.com/antirez/ds4/pull/1154), common to both arms.
**These gains are incremental to #1154, not a measurement of that PR.** The PR
is based on upstream `fc80bd695da76ee14bcb86820bf900a5cdfda806`; its unmodified
shader, kernel test and Metal host source match the campaign baseline. The
complete upstream host integration was built separately, but was not rerun
through the full model campaign.
Shader SHA256:
- Control: `13224aeae23802defef678877146ae8047a45c2d1be824361efa2f4f5aa25b28`
- Candidate: `dff0b9e48e5d55b68a818d4ccc28ea3468ecd464b8eff83cf426dcb9a102b4cf`
## Measurement
M5 Max, 128 GiB, macOS 26.6.2, resident Qwen3.8 Flash Next Q4 GGUF, Q4_K routed
gate/up and MXFP4 routed down, 8192-token chunks, native context allocation
262144, no batching. Each process warmed weights and a disposable complete
prefill/decode session, freed that session, then measured fresh KV on the same
engine/source. This avoids first-use specialized Metal pipeline compilation.
Only one model process/session ran at a time. Timed runs did not add stage
synchronization or state dumps. Decode measurements use 64 ordinary greedy
tokens; they do not measure MTP throughput.
| Corpus | Runs / balanced blocks | Prefill throughput gain | Conditional 95% block-bootstrap interval | Ordinary decode change |
| --- | --- | --- | --- | --- |
| Standard text, 16384 tokens | 16 / 4 | +2.376% | +2.172 to +2.660% | -0.061% |
| C source, 16384 tokens | 8 / 2 | +2.274% | +2.147 to +2.401% | -0.622% |
Each block uses ABBA or BAAB (A=control, B=candidate). Its ratio is mean B time
/ mean A time; the summary is the geometric mean of block ratios, converted
with `gain = 100 * (1 / ratio - 1)`. The interval enumerates ordered block
resamples with replacement. It is conditional on these few blocks, not a
complete estimate of host uncertainty. A warmed A/A block showed a false
-1.83% difference, and each arm's text prefill times drifted about 6–7%
first-to-last. All four text blocks and both code blocks favored the candidate.
Cold-process results are excluded from this estimate. Do not extrapolate this
gain to every context length.
Historical model SHA256:
`680944460a8cbe93ba8b6d7b6107213ffb7e22320bd913000e563ca0a0f25a8a`.
The 177280286720-byte model was not rehashed for this campaign; its recorded
size, inode and modification time matched the previously verified file.
### Per-run durations
Seconds reconstructed from CSV throughput rounded to two decimal places;
printed precision below is not measurement precision. Each row measures
16384 prefill tokens and 64 ordinary decode tokens.
| Corpus | Run | Arm | Block | Prefill seconds | Decode seconds |
| --- | --- | --- | --- | --- | --- |
| Text | 1 | A | 1 | 10.618895 | 1.210974 |
| Text | 2 | B | 1 | 10.405314 | 1.201652 |
| Text | 3 | B | 1 | 10.576054 | 1.200075 |
| Text | 4 | A | 1 | 10.947627 | 1.197829 |
| Text | 5 | B | 2 | 11.428253 | 1.231954 |
| Text | 6 | A | 2 | 11.613763 | 1.234806 |
| Text | 7 | A | 2 | 11.543149 | 1.238151 |
| Text | 8 | B | 2 | 11.247340 | 1.234806 |
| Text | 9 | A | 3 | 11.488353 | 1.211892 |
| Text | 10 | B | 3 | 11.199748 | 1.226054 |
| Text | 11 | B | 3 | 11.201050 | 1.218351 |
| Text | 12 | A | 3 | 11.421640 | 1.221141 |
| Text | 13 | B | 4 | 11.157419 | 1.220675 |
| Text | 14 | A | 4 | 11.397010 | 1.217193 |
| Text | 15 | A | 4 | 11.355461 | 1.213730 |
| Text | 16 | B | 4 | 11.078879 | 1.218120 |
| Code | 1 | A | 1 | 10.207209 | 1.199850 |
| Code | 2 | B | 1 | 10.157848 | 1.220210 |
| Code | 3 | B | 1 | 10.268109 | 1.225115 |
| Code | 4 | A | 1 | 10.657300 | 1.221840 |
| Code | 5 | B | 2 | 10.232773 | 1.218351 |
| Code | 6 | A | 2 | 10.552621 | 1.219048 |
| Code | 7 | A | 2 | 10.505931 | 1.217193 |
| Code | 8 | B | 2 | 10.331954 | 1.224646 |
## Correctness and regression checks
Campaign checks passed on the patched host baseline and final local integration:
- The focused fixture directly byte-compared 87,492,020 old/new output bytes,
covering both changed call sites, NAX levels 0–6, Q4/Q4, Q4/MXFP4, all-MXFP4
and IQ2XXS/Q2_K, K512/2560, signed/zero finite Q4 headers, shuffled routes and
counts around 8/16/32/64 boundaries. Live values must be finite and nonzero;
reserved output slots and trailing write sentinels must remain untouched.
- All 24 warmed runs matched each corpus's full 248320-float frontier row.
- Complete persisted payloads at 16384 and 16896 tokens, save/load/resave,
128 ordinary decode vocabulary rows, token IDs and positions matched exactly.
- Near-native retained prefixes at 131072, 245760 and 246272 tokens matched
full frontier logits, including the final 512-token append. These are
correctness observations, not statistically repeated speed measurements.
- `make`, `make test`, the full Qwen kernel suite, PLE rollback/prefetch/failure
tests, checkpoint/restore/snapshot/rewind with opportunistic MTP, ordinary/MTP
vision completion tests and synthetic Metal resident/SSD Q4 comparisons passed.
- An isolated 11-request sampled tool/history/restart fixture passed at
temperature 0.6 with non-exact MTP: 102 verification cycles, 94 drafts accepted.
This bounded smoke is not a multi-hour agent soak or an MTP speed comparison.
CUDA sources are untouched; no external CUDA or distributed runtime test is
claimed. The finite fixture is not exhaustive over all possible bit patterns.
Serialized state comparisons do not cover every ephemeral MTP scratch buffer.
The existing masked-row loaders can form inactive weight-row addresses past a
65-row logical tensor; this is present in both arms. Output sentinels check
writes, not logical weight-read bounds. Register reuse preserves/subsets those
read addresses; deployed row dimensions are multiples of 64. A bounds repair
needs separate validation and is not included in this patch.
### Reproduce the focused exact comparison
On M5, from the PR checkout, run the same test executable in separate processes
against the baseline and candidate shaders. A fixture pass alone is not the
exactness oracle; compare the complete dumps. Run GPU/model tests only when the
managed server is free and exclusively reserved.
```sh
make tests/test_qwen4_kernels
git show fc80bd695da76ee14bcb86820bf900a5cdfda806:metal/qwen4.metal > /tmp/qwen4-control.metal
DS4_METAL_QWEN4_SOURCE=/tmp/qwen4-control.metal DS4_TEST_QWEN4_RAW16_ONLY=1 DS4_TEST_QWEN4_RAW16_DUMP=/tmp/qwen4-control.raw ./tests/test_qwen4_kernels
DS4_METAL_QWEN4_SOURCE="$PWD/metal/qwen4.metal" DS4_TEST_QWEN4_RAW16_ONLY=1 DS4_TEST_QWEN4_RAW16_DUMP=/tmp/qwen4-candidate.raw ./tests/test_qwen4_kernels
cmp /tmp/qwen4-control.raw /tmp/qwen4-candidate.raw
```
Expected dump size: 87492020 bytes; expected SHA256 for each:
`6cd1882c9ede246f78cad98f0e6d13b63a045a5d0ba332e186e25e57e5dba554`.
Timing used a diagnostic warmed-session harness rather than the stock first
prefill invocation. Raw local manifests, shaders, dumps and generator scripts
are retained outside the PR; the durations above allow the reported timing
ratios to be independently recalculated.
docs/SERVER.md
+154
View File
@@ -0,0 +1,154 @@
# Serving Models
[README](../README.md) | [Client setup](CLIENTS.md)
## Start locally
```sh
./ds4-server --ctx 32768
```
The default address is `http://127.0.0.1:8000`. Use `--host 0.0.0.0` to listen
on other interfaces. Restrict access to trusted clients; for an Internet-facing
deployment, put authentication and TLS in front of the server.
`--cors` enables browser cross-origin headers. It does not change the listening
address or provide access control.
The selected build chooses the backend. Pass `-m FILE` for an explicit model.
Use `--chdir /path/to/ds4` when starting outside the project directory so
relative runtime files such as Metal kernels can be found.
## APIs
| Endpoint | Use |
| --- | --- |
| `GET /v1/models` | Loaded model information |
| `POST /v1/chat/completions` | OpenAI-style chat |
| `POST /v1/responses` | Responses-style requests and continuations |
| `POST /v1/completions` | Text completions |
| `POST /v1/messages` | Anthropic-style messages |
The Flash and PRO names accepted by the model endpoints are compatibility
aliases, not separate loaded models. The GGUF passed at startup selects the model.
`GET /v1/models` and `GET /v1/models/{id}` include `architecture.input_modalities` and `architecture.output_modalities`. Input modalities are `["text", "image"]` when the loaded engine has vision enabled, otherwise `["text"]`; output modalities are always `["text"]`. Compatibility aliases report the same loaded-engine capabilities.
```sh
curl http://127.0.0.1:8000/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{"model":"deepseek-v4-flash","messages":[{"role":"user","content":"Explain Redis streams."}],"stream":true}'
```
Chat, Responses, and Anthropic support tools and SSE streaming. Reasoning is
returned separately from visible text in each API's native form. Standard
sampling and output-budget fields are supported; explicit request parameters
take precedence over defaults.
The default sampling settings are temperature 1, top-p 1, and min-p 0.05.
For DeepSeek, thinking is on by default. `reasoning_effort=max` selects Think
Max only with sufficient context; otherwise it falls back to normal thinking.
`xhigh` maps to normal thinking, not Think Max. Use `think:false`, a disabled
thinking object, or a non-thinking model alias for direct answers.
## Multiple sessions
```sh
./ds4-server --ctx 4096 --batched-session 4
```
Without `--batched-session`, there is one resident session. With it, the server
preallocates independent KV states and queues requests when all slots are busy.
Choose context and slot count together: a context that fits once may not fit
four times. Idle slots can be cached before reuse; active requests are not evicted.
Where the model supports it, the slots share one prefill workspace instead of
each keeping its own, so an extra slot costs only its caches. That matters most
for Qwen3.8 Flash Next, whose transients are sized by the prefill chunk rather
than by the context: at the default chunk they run to several GiB per session.
The startup line reports both figures.
| Backend/model | Decode execution |
| --- | --- |
| Metal, resident Flash | Native shared-expert/QKV batching where supported |
| Metal, resident V4.1 Flash | Native decoding for 2-8 sessions |
| Metal RDMA TP, V4.1 Flash | Native decoding for 3-8 sessions; ordered fallback for two |
| Metal SSD streaming, V4.1 Flash | Ordered fallback |
| Metal, GLM 5.2 | Ordered fallback |
| Metal, GLM 5.3 | Native batching through 2051 visible tokens; ordered fallback afterward |
| Metal, Qwen3.8 Flash Next with resident weights | Native batching of the shared work; recurrent state, caches and PLE history stay per session |
| CUDA, supported multi-GPU Flash TP layout | Native grouped decode and mixed prefill/decode |
| Metal, Qwen3.8 Flash Next with SSD expert streaming | Ordered fallback with per-session streaming workspaces |
| Single-GPU CUDA, including Spark | Ordered fallback |
Fallback executes the rows separately. It provides concurrency and scheduling
fairness, not the aggregate speedup of native batching. Native grouping may
change floating-point reduction order slightly. V4.1 sessions containing images
use the ordered fallback.
Long prefills yield to active decoders in bounded intervals, normally 128
tokens. `--mixed-prefill-quantum N` changes that interval for testing. While
nothing is decoding, a V4.1 slot prefills in the engine's own chunk.
Session-batched serving uses ordinary target decoding, except Qwen3.8 on
Metal, where `--mtp` also batches speculative decoding with resident weights
and uses ordered speculative cycles with SSD expert streaming. Its
`--mtp-exact-sampling` mode uses ordinary batches for nonzero-temperature
requests. Other models do not use MTP/DSpark while session batching is active.
For the eight-L40S example, see [CUDA GPUs](CUDA_MULTI_GPU.md#serve-multiple-users).
## Images
Start with the matching language GGUF and `--vision FILE`; see
[model-specific instructions](MODELS.md#vision).
OpenAI chat and Responses accept inline PNG/JPEG data URIs. Anthropic accepts
base64 image sources. Remote URLs and server-side file paths are rejected.
Image blocks preserve their order in the request. The limit is 16 images and
a 64 MiB HTTP body.
## Disk KV cache
Disk caching saves useful prefixes across slot reuse and server restarts:
```sh
./ds4-server --ctx 100000 \
--kv-disk-dir /tmp/ds4-kv --kv-disk-space-mb 8192
```
Clients can resend complete conversation histories. The server first tries
the live token prefix, then compatible rendered-text prefixes from disk, and
prefills the new suffix. With multiple slots, each slot has its own live state;
disk is the additional persistence layer, not the only way to retain sessions.
TP cache loading rebuilds the saved token prefix on both ranks. It is not an
instantaneous restoration of both GPUs. Pipeline loading redistributes the
saved layer state over its route.
Defaults are intended to avoid saving fragile token boundaries. For unusual
workloads, the controls are `--kv-cache-min-tokens`,
`--kv-cache-cold-max-tokens`, `--kv-cache-continued-interval-tokens`,
`--kv-cache-boundary-trim-tokens`, and `--kv-cache-boundary-align-tokens`.
Check `./ds4-server --help` for their defaults.
Quantization variants may share compatible prefixes. Add
`--kv-cache-reject-different-quant` for same-quant reuse only.
Cache files contain prompt text and model state: treat the directory as
private. It is disposable; stop the server before clearing it.
## Tool history and debugging
For DeepSeek, the server preserves sampled DSML tool blocks and assigns
unguessable tool IDs. Replaying those IDs avoids retokenizing a differently
formatted JSON history. The bounded replay map can be stored in cache files.
When exact replay is unavailable, canonical rendering may require rebuilding
part of the prefix.
`--tool-memory-max-ids` bounds this map.
`--disable-exact-dsml-tool-replay` disables it for diagnostic comparisons.
Use `--trace /tmp/ds4-trace.txt` to record prompt rendering, cache decisions,
generated text, and tool-parser events. Traces can contain sensitive content.
Cache formats are implementation details. The current header and extension
definitions are in [ds4_kvstore.h](../ds4_kvstore.h) and
[ds4_kvstore.c](../ds4_kvstore.c); model-specific payload handling is in
[ds4.c](../ds4.c).
docs/SPECULATIVE_DECODING.md
+179
View File
@@ -0,0 +1,179 @@
# Speculative Decoding
[README](../README.md)
Speculation proposes future tokens with a smaller draft block, then checks
them with the main model. An accepted prefix advances generation by several
tokens in one verification pass. It does not accelerate prefill.
It is opt-in. Gains depend on the prompt, model, backend, and context length;
poor acceptance can make it slower. Measure your workload rather than assuming
that a draft model always helps.
## DeepSeek Flash: DSpark
DSpark is a separate support GGUF, not a standalone language model. It proposes
up to five future tokens. Match its checkpoint to the main model:
| Main checkpoint | Download | Support file |
| --- | --- | --- |
| Flash 0731 | `ds4f-dspark` | `gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf` |
| Flash Vision Experimental | `ds4f-vision-dspark` | `gguf/DeepSeek-V4-Flash-Vision-Exp-DSpark-support.gguf` |
For the 0731 Q2 model:
```sh
./download_model.sh ds4f-q2
./download_model.sh ds4f-dspark
./ds4 --dspark --mtp-model gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf
```
For Vision Experimental, substitute its matching main model and support file.
Do not mix the two checkpoints. DSpark is not supported for PRO.
The same flags work in `ds4-agent` and non-batched `ds4-server` requests.
The support file adds about 5.6 GiB of weights plus runtime state. The main
model must be resident; external DSpark is not supported with SSD streaming.
DSpark replaces the legacy one-stage MTP drafter for that run; the two are
not stacked.
Resident M5 paths batch supported verifier expert rows, including two-Mac TP.
On DGX Spark, resident Flash 0731 Q2 and two-Spark MXFP4 also batch the seed
with longer drafts and use small-batch Q8 and expert kernels. No extra flags
are needed.
Single-Spark Q2 uses a confidence-pruning threshold of 0.4. Two-Spark
Flash 0731 Q2 and MXFP4 leave confidence pruning off and split the drafter
over RDMA when both peers load its support file. `--dspark-confidence`
overrides the threshold. Exact sampling keeps its default threshold of 0.8.
The scheduler can back off when drafting is unproductive. Defaults select the
fast paths; diagnostic environment variables are not needed for normal use.
Recorded comparisons are in [the QA guide](../QA_BEFORE_RELEASES.md).
For the tested Strix Halo coding configuration, use `--dspark --dspark-confidence 0.7` with the default five-token draft cap and scheduler. Client sampling is temperature `1.0`, `top_p=0.95`, `min_p=0`, and `top_k=0`; high reasoning was also checked on coding and tool-use requests. This uses opportunistic sampling as described below; exact-mode throughput is not qualified by these measurements. `--mtp-draft` controls legacy autoregressive MTP, not the DSpark draft width.
## DeepSeek V4.1 Flash: DSpark
V4.1 has its own 8.04 GiB support GGUF. It retains the native MXFP4 experts
and uses a different draft graph from V4 Flash. Do not interchange the files.
See [conversion](../gguf-tools/README.md#convert-v41-dspark) to build it.
```sh
./ds4 -m gguf/DeepSeek-V4.1-Flash-Q2.gguf \
--ssd-streaming --dspark \
--mtp-model gguf/DeepSeek-V4.1-Flash-DSpark.gguf
```
The same options work with `ds4-agent` and `ds4-server`. Metal, CUDA and ROCm
support single-host execution, with either resident target weights or SSD
streaming. The drafter stays resident, so leave another 8 GiB for it.
Two-host tensor parallelism works on Metal and CUDA: add the same `--dspark
--mtp-model FILE` options to both the coordinator and worker. Each peer keeps
the drafter resident; the target and its verification run across both GPUs.
Up to five drafts are checked together. The sampling choices below apply;
`--mtp-exact-sampling` uses sampled proposals and the full acceptance/correction
rule for V4.1. Restored sessions warm the draft window with 128 ordinary tokens
before speculation resumes. Measure speed on your prompts: acceptance and SSD
cache size both matter. V4.1 DSpark can still be slower than ordinary decoding,
including in TP mode.
## GLM: built-in MTP
GLM's draft block is already in its main GGUF:
```sh
./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf --mtp
```
`--mtp-timing` also enables it and prints acceptance and timing counters.
The current GLM cycle commits up to two tokens. No external support file is
needed, and ordinary decode remains the default.
## Qwen3.8: built-in MTP
Both Qwen downloads include MTP and native BF16 n-grams in the main GGUF:
```sh
./download_model.sh qwen38-q4k
./ds4 --mtp
```
Ordinary decode uses the same file with `--mtp` omitted. For non-zero
temperature, add `--mtp-exact-sampling` to preserve the target sampling
distribution. See [Qwen setup](QWEN38_FLASH_NEXT.md) for Metal and CUDA.
For the Q2 model with Metal SSD streaming:
```sh
./ds4 --metal -m gguf/Qwen3.8-Flash-Next-Q2.gguf \
--ssd-streaming --ctx 4096 --prefill-chunk 128 \
--mtp-timing --temp 0 --nothink -n 100 \
-p "narrami la storia di roma"
```
`--mtp-timing` enables MTP and reports verification cycles and first-draft
acceptance. Measure generation throughput against the same command without
that option: speculative decoding can lose time when too few proposals are
accepted, especially when loading the predictor's experts from SSD.
Qwen's [official model configuration](https://huggingface.co/Qwen/Qwen3.8-Flash-Next/blob/main/config.json)
contains one MTP layer with no dedicated token embeddings. The same layer
can be called repeatedly to propose multiple tokens; the number of MTP
layers is not the draft length. The
[official blog](https://www.alibabacloud.com/blog/qwen3-8-flash-next-a-new-architecture-towardsultimate-cost-efficiency_603501)
describes training this module over multiple steps and using QSA in it.
These weights are already included in the GGUF; no separate draft model is
needed.
With SSD streaming, prompt prefill prepares the predictor's K/V and indexer
cache from the trunk hidden states and the next prompt token by default. It skips the predictor's
attention output, MoE and vocabulary head, so this preparation does not
stream predictor experts from SSD. The final hidden row is retained until
the actual next token is known. Snapshot restore and speculative rollback
retain that row too. This costs 40 KiB per row for this model, at most
160 KiB including the three verifier snapshots. Qwen payload version 3
stores the initialized prefix and retained row. Version 2 checkpoints and
plain checkpoints restore their target state directly. When MTP is enabled
without a saved predictor tail, the engine replays the saved tokens once
before continuing. With resident weights, prefill skips historical predictor
projections and initializes those cache rows to zero. MTP and target
verification remain enabled. Preparing the full history can improve draft
acceptance for long generations. Use `DS4_QWEN4_MTP_PREFILL=on` to enable
it for resident weights too. [Qwen MTP settings](QWEN38_FLASH_NEXT.md)
The cycle drafts one token ahead by default and engages a **second, chained
draft** (one extra nextn-layer step conditioned on the predictor's own
stream, verified in a 3-row pass) while recent first-draft acceptance is
perfect, disengaging after repeated second-draft rejections.
`DS4_QWEN4_MTP_DEPTH=2` or `=3` fixes the depth;
`0` (default) is the adaptive policy.
## Sampling and reproducibility
At temperature zero, accepted drafts must match the target's greedy
continuation. At non-zero temperature, the default mode is opportunistic:
ordinary tokens use the requested sampling settings, but matching greedy
drafts are accepted directly. Sampling resumes when the proposed suffix does
not match. This is deliberately more deterministic than ordinary sampling.
Use `--mtp-exact-sampling` to preserve the ordinary target sampling
distribution. Greedy proposals are accepted with their target probability;
sampled proposals use the target-to-draft probability ratio. Rejections sample
from the remaining target probability mass.
When a verified block crosses a tool sampling-mode boundary (for example
entering tool-call syntax during server decoding), the server rewinds to the
block start and re-evaluates the boundary token so the next sample uses the
new mode. Under exact sampling that rewind restores a pre-verify snapshot of
the recurrent state instead of resetting the graph, so long retained
contexts are not replayed at every boundary.
Accepted tokens keep the state produced by the batched verifier. Floating-point
reduction order can differ from one-token decode, so long greedy continuations
need not be byte-identical. For DeepSeek comparisons against the ordinary
target-only path, use `--quality` or `--dspark-strict`; these disable the
speculative acceptance path. They do not promise identical output across
different hardware or execution configurations.
Session-batched serving uses ordinary target decoding instead of combining
DSpark/MTP with the session batch. See [serving](SERVER.md#multiple-sessions).
docs/SSD_STREAMING.md
+118
View File
@@ -0,0 +1,118 @@
# Models Larger Than RAM
[README](../README.md)
Use resident inference when the model and runtime fit: it is faster.
SSD streaming keeps a bounded cache of routed experts and reads missing
experts from the GGUF. It trades speed for capacity; it does not remove the
memory needed for other weights, activations, scratch, and the context.
Metal supports streaming for DeepSeek, GLM and [Qwen3.8](QWEN38_FLASH_NEXT.md).
CUDA has streaming paths too, and ROCm supports GLM 5.2/5.3 streaming.
Do not infer support for every model
and tensor layout from the existence of the flag.
## Start with the automatic budget
```sh
./ds4 -m ds4flash.gguf --ssd-streaming
```
The startup report shows the effective cache, resident layers where applicable,
and memory requirements. Prefer the automatic budget for a first run. Use a
local SSD and leave memory for the OS and other applications.
Examples:
```sh
# Flash Q2 on a smaller Mac.
./download_model.sh ds4f-q2
./ds4 --ssd-streaming --ctx 32768 --nothink
# GLM 5.3 Flash Q4 on a 128 GB Mac.
./download_model.sh glm53-q4
./ds4 --ssd-streaming --ctx 4096
# PRO Q2 on a 128 GB Mac: usable for inspection, but slow.
./download_model.sh pro-q2-imatrix
./ds4 --ssd-streaming --ctx 32768 --nothink
```
Generation is usually more sensitive to cache misses than prefill. A large
model that starts successfully can still be too slow for interactive work.
Use a short generation before committing to a long task.
## Adjusting the cache
To leave more room for context or other sessions:
```sh
./ds4 --ssd-streaming --ssd-streaming-cache-experts 32GB
```
A byte budget is a target, not a guaranteed allocation. DwarfStar reserves
routed-prefill headroom and fits the cache to the remaining model, graph,
context, and backend budget. The effective value may be smaller than requested.
Non-routed weights and KV state are additional to that expert-cache budget.
A plain number, such as `--ssd-streaming-cache-experts 4000`, requests dynamic
expert slots rather than a byte budget. It is also subject to memory limits.
By default, GLM spends the cache budget on selected experts across all layers.
`--ssd-streaming-full-layers N` reserves full routed prefix layers instead.
Metal also keeps the non-routed weights resident when it can; they are needed
by every token, and paging them back in delays generation after a prompt.
An oversized expert cache can displace those weights and slow decoding.
More cache helps only while the rest of the working set still fits.
More context and more server sessions consume more memory. Metal GLM budgeting
includes planned and already resident sessions; a request can still be refused
when its mandatory allocations cannot fit. Disabling the memory guard is not
a remedy for insufficient RAM.
Leave expert preloading enabled for normal use. `--ssd-streaming-cold` and
`--ssd-streaming-preload-experts N` are mainly useful for controlled measurements.
See [benchmarking](PERFORMANCE.md) and the [release QA guide](../QA_BEFORE_RELEASES.md).
## CUDA expert reads
DeepSeek V4.1 CUDA prefill reads selected experts in file order and combines
adjacent reads by default, instead of prefetching every expert in the next
layer. Single-token decode retains its original read order, which avoids a
slowdown with small expert caches. The weight formats, routing, and cache
budget are unchanged; no extra flags are needed.
For comparisons on other storage or workloads,
`DS4_CUDA_ENABLE_SSD_PREFETCH=1` restores whole-next-layer prefetch.
`DS4_CUDA_DISABLE_SSD_SORT_SELECTED_READS=1` restores the original selected-read
order. Set both to reproduce the previous policy.
`DS4_CUDA_DISABLE_SSD_PREFETCH=1` still overrides prefetch enablement.
For experiments, `DS4_CUDA_SSD_SORT_SELECTED_READS=1` also orders single-token
loads; the explicit sort-disable switch takes precedence.
Unset that old experimental sort flag to use the new prefill-only default.
These switches use presence semantics: unset them to return to the default;
setting a switch to `0` still enables it.
## M5 Max measurements
On a 128 GB M5 Max, with automatic cache sizing and no speculative decoding:
| Model | Initial prefill | Continued prefill | Generation after each | Runs |
| --- | ---: | ---: | ---: | --- |
| GLM 5.3 Flash Q4_K, 177.77 GiB | 121 t/s | 104 t/s | 11.9 / 14.9 t/s | Three-run median |
| DeepSeek Flash Vision Exp MXFP4, 145.26 GiB | 300 t/s | 263 t/s | 11.9 / 19.3 t/s | Single run |
September 6, 2026. GLM used a 2K prompt and a 1K append; DeepSeek used 8K
and a 4K append. Both generated 128 tokens at each frontier. DeepSeek's latest
run uses an 86.2 GiB expert cache and updated eviction priorities; it is not
directly comparable to the earlier 64-token measurement. Generation includes
the first-token wait. Cache reuse depends on the prompt, so these are
workload references, not a speed
guarantee for every model larger than RAM.
Full GLM 5.3 IQ2_XXS (196.58 GiB) also benefits from caching short tool-result
prompts. With an 8K context and a 61.35 GiB effective expert cache, a 16-token
append fell from 30.8 to 2.9 seconds; generation afterward rose from 4.09 to
5.11 t/s (three-run medians, 64 generated tokens). All compared logits and
generated text were identical. This is a short-append improvement, not an
equivalent gain in initial prefill or every generation workload.
docs/STRIX_HALO.md
+101
View File
@@ -0,0 +1,101 @@
# AMD Strix Halo
[README](../README.md) | [Getting started](../README.md#start-here)
The reference system is a 128 GB Strix Halo with Radeon 8060S (`gfx1151`),
such as the Framework Desktop. The ROCm build uses the standard binary names
and selects the ROCm backend by default.
## Prerequisites
For a container setup, see the maintained
[ROCm toolbox](https://github.com/kyuz0/strix-halo-ds4-toolbox/blob/main/toolboxes/Dockerfile.rocm-10.0).
It can also be managed with
[AI Toolbox Cockpit](https://github.com/kyuz0/ai-toolbox-cockpit).
For a native Ubuntu build you need HIP, hipBLAS, hipBLASLt, rocBLAS, rocWMMA,
and hipCUB development files. The Ubuntu 26.04 setup used these packages:
```sh
sudo apt-get update
sudo apt-get install -y hipcc rocminfo rocm-smi \
libamdhip64-dev libhipblas-dev libhipblaslt-dev librocblas-dev \
librocwmma-dev libhipcub-dev
sudo usermod -aG render,video "$USER"
```
Log out and back in after changing groups. `rocminfo` must report `gfx1151`
and be able to open `/dev/kfd` before DwarfStar can run.
Some packaged rocWMMA headers omit `rocwmma/internal/`. If compilation fails
there, install the complete headers matching your ROCm installation, or use
the container. Do not mix header versions as a general workaround.
## GPU-visible memory
Check the memory pool reported by `rocminfo`. Some 128 GB configurations expose
only about 62 GB to the GPU, which is insufficient for resident Flash Q2 plus
runtime buffers. Firmware and kernel GTT/TTM settings control this limit.
The native reference setup used these memory parameters:
```text
amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856
```
They are a system-specific starting point, not an allocation budget for
DwarfStar. Preserve existing boot options and consult your kernel's settings
before changing them. Keep RAM available for the OS even when the GPU can
address most of it. Do not disable the IOMMU merely to copy another host's
configuration; doing so changes device isolation.
## Build and run Flash
```sh
make strix-halo
./download_model.sh ds4f-q2
./ds4 --rocm
```
`make rocm` is an alias. Use the current 0731 Q2 download for a first run;
larger mixed and Q4 models have substantially higher memory requirements.
Flash's ROCm resident and pipeline paths should not be confused with the GLM
SSD-streaming path.
## DeepSeek V4.1 Flash
Use Q2 with SSD streaming on a 128 GB machine:
```sh
./download_model.sh ds41f-q2
./ds4 --rocm -m gguf/DeepSeek-V4.1-Flash-Q2.gguf \
--ssd-streaming --ctx 4096
```
The 341 GiB GGUF includes disk-only Engram tables. Keep it on a fast local
SSD and start with automatic cache sizing. The matching
[V4.1 DSpark](SPECULATIVE_DECODING.md#deepseek-v41-flash-dspark) support file
is optional. V4.1 vision and network TP are not supported on ROCm.
Large prefills are not yet tuned on this backend and can be slow.
## GLM 5.3 Flash
The reference Q2 setup uses SSD streaming to leave room for its graph and KV
state. Begin with automatic cache sizing and a small context:
```sh
./download_model.sh glm53-q2
./ds4 --rocm -m gguf/GLM-5.3-Flash-Q2.gguf \
--ssd-streaming --ctx 4096
```
GLM 5.2 also supports ROCm streaming. Full-model GLM 5.2 inference requires it;
distributed layer slices can be resident. See [SSD streaming](SSD_STREAMING.md)
before adjusting the cache budget.
Both GLM 5.3 Flash and DeepSeek Flash Vision Experimental support images on
ROCm. Add the matching encoder with `--vision FILE`, as described in
[models and vision](MODELS.md#vision).
For a model-free routed-kernel check, use `make test-mxfp4-rocm`.
Full-model validation is described in [testing](TESTING.md).
docs/TESTING.md
+66
View File
@@ -0,0 +1,66 @@
# Testing and Development
[README](../README.md)
Read [CONTRIBUTING.md](../CONTRIBUTING.md) before proposing an inference change.
The [release QA guide](../QA_BEFORE_RELEASES.md) defines the hardware matrix,
checkpoint-matched quality tests, long-context agent tasks, and speed checks.
It also records which checks were not completed. A smoke test is not full QA.
## Focused checks
Model-free checks include:
```sh
make ds4_test ds4_agent_test test-session-state
./ds4_test --server
./ds4_agent_test
```
On Metal, small GPU tensor tests are available without loading a full GGUF:
```sh
make tests/test_session_state_gpu tests/test_glm53_kda tests/test_mxfp4_metal
./tests/test_session_state_gpu
./tests/test_glm53_kda
./tests/test_mxfp4_metal
```
`make test` also includes model-backed tests. Select the right GGUF and ensure
that it fits before running it; do not accidentally load a large model on a
single device during multi-GPU QA. ROCm has `make test-rocm`.
Official-vector tests must use continuations from the same checkpoint as the
GGUF. Flash 0731 and Vision Experimental are not interchangeable fixtures.
See [test vectors](../tests/test-vectors/README.md) and
[quality scoring](../gguf-tools/quality-testing/README.md).
## Investigating output
```sh
./ds4 --dump-tokens -p "..."
./ds4 --dump-logprobs /tmp/out.json --logprobs-top-k 20 --temp 0 -p "..."
./ds4 --dump-logits /tmp/logits.json --nothink --prompt-file prompt.txt
./ds4-server --trace /tmp/ds4-trace.txt
```
Token dumps catch template differences without inference. Logits and
continuations help distinguish sampling changes from graph errors. Server
traces include cache and tool-parser decisions. Keep traces private when they
contain real conversations.
## Tools and source references
- [GGUF conversion and quantization](../gguf-tools/README.md)
- [Imatrix collection](../gguf-tools/imatrix/README.md) and
[calibration corpus](../gguf-tools/imatrix/dataset/README.md)
- [Steering vectors](../dir-steering/README.md)
- [Benchmark scripts and charts](../speed-bench/README.md)
- [Evaluation data and licenses](../EVAL_DATA.md)
- [Session payload implementation](../ds4.c) and
[cache header definitions](../ds4_kvstore.h)
- [Pipeline protocol](../ds4_distributed.c) and [TP protocol](../ds4_tp.c)
For changes to state handling, cover rewind/replay, save/load, images, and
multiple sessions as well as a fresh prompt. For distributed changes, exercise
both ranks and failures; a local command-parser test is not physical TP QA.
download_model.sh
+503 -34
View File
@@ -1,23 +1,45 @@
#!/bin/sh
set -e
DEEPSEEK_REPO="antirez/deepseek-v4-gguf"
GLM_UNSLOTH_REPO="unsloth/GLM-5.2-GGUF"
GLM_ANTIREZ_REPO="antirez/GLM-5.2-GGUF"
REPO=$DEEPSEEK_REPO
Q2_IMATRIX_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf"
Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf"
Q2_Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf"
PRO_Q2_IMATRIX_FILE="DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf"
GLM53_REPO="antirez/glm-5.3-flash-gguf"
GLM53_FULL_REPO="antirez/glm-5.3-gguf"
QWEN38_REPO="antirez/qwen3.8-flash-next-gguf"
QWEN38_MMPROJ_REPO="ggml-org/Qwen3.8-Flash-Next-GGUF"
REPO="antirez/deepseek-v4-gguf"
DS41_REPO="antirez/deepseek-v4.1-flash-gguf"
DS41_Q2_FILE="DeepSeek-V4.1-Flash-1017-IQ2_XXS-Q2_K-imatrix.gguf"
DS41_Q4_FILE="DeepSeek-V4.1-Flash-Q4.gguf"
DS41_MXFP4_BASE="DeepSeek-V4.1-Flash-1017-MXFP4-experts"
DS41_VISION_FILE="DeepSeek-V4.1-Flash-Vision.gguf"
DS4F_Q2_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf"
DS4F_Q4_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix-0731.gguf"
DS4F_MXFP4_FILE="DeepSeek-V4-Flash-MXFP4Experts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-mxfp4-0731.gguf"
DS4F_Q2_Q4_FILE="DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed-0731.gguf"
PRO_Q2_IMATRIX_FILE="DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf"
PRO_Q4_LAYERS00_30_FILE="DeepSeek-V4-Pro-Q4K-Layers00-30.gguf"
PRO_Q4_LAYERS31_OUTPUT_FILE="DeepSeek-V4-Pro-Q4K-Layers-31-output.gguf"
MTP_FILE="DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf"
DS4F_DSPARK_FILE="DeepSeek-V4-Flash-DSpark-support-0731.gguf"
DS4F_VISION_Q2_FILE="DeepSeek-V4-Flash-Vision-Exp-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8.gguf"
DS4F_VISION_Q2_Q4_FILE="DeepSeek-V4-Flash-Vision-Exp-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8.gguf"
DS4F_VISION_MXFP4_FILE="DeepSeek-V4-Flash-Vision-Exp-MXFP4Experts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out.gguf"
DS4F_VISION_ENCODER_FILE="DeepSeek-V4-Flash-Vision-Encoder.gguf"
DS4F_VISION_DSPARK_FILE="DeepSeek-V4-Flash-Vision-Exp-DSpark-support.gguf"
GLM_UNSLOTH_Q4_REMOTE_BASE="UD-Q4_K_XL/GLM-5.2-UD-Q4_K_XL"
GLM_UNSLOTH_Q4_LOCAL_BASE="GLM-5.2-UD-Q4_K_XL"
GLM_UNSLOTH_Q4_FIRST_FILE="$GLM_UNSLOTH_Q4_LOCAL_BASE-00001-of-00011.gguf"
GLM_ANTIREZ_IQ2XXS_FILE="GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf"
GLM_ANTIREZ_Q2_FILE="GLM-5.2-UD-Q2_K_RoutedQ2K.gguf"
GLM_ANTIREZ_Q4_FILE="GLM-5.2-UD-Q4_K_RoutedQ4K.gguf"
GLM53_FULL_Q2_FILE="GLM-5.3-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf"
GLM53_Q2_FILE="GLM-5.3-Flash-Q2.gguf"
GLM53_Q4_FILE="GLM-5.3-Flash-Q4_K.gguf"
GLM53_FP8_FILE="GLM-5.3-Flash-FP8.gguf"
GLM53_VISION_FILE="GLM-5.3-Flash-Vision-Encoder.gguf"
QWEN38_Q4_FILE="Qwen3.8-Flash-Next-Q4.gguf"
QWEN38_Q2_FILE="Qwen3.8-Flash-Next-Q2.gguf"
QWEN38_VISION_FILE="mmproj-Qwen3.8-Flash-Next-Q8_0.gguf"
ROOT=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
OUT_DIR=${DS4_GGUF_DIR:-"$ROOT/gguf"}
@@ -32,38 +54,106 @@ usage() {
DwarfStar GGUF downloader
Usage:
./download_model.sh q2-imatrix [--token TOKEN]
./download_model.sh q2-q4-imatrix [--token TOKEN]
./download_model.sh q4-imatrix [--token TOKEN]
./download_model.sh ds4f-q2 [--token TOKEN]
./download_model.sh ds4f-q2-q4 [--token TOKEN]
./download_model.sh ds4f-q4 [--token TOKEN]
./download_model.sh ds4f-mxfp4 [--token TOKEN]
./download_model.sh ds4f-dspark [--token TOKEN]
./download_model.sh ds4f-vision-q2 [--token TOKEN]
./download_model.sh ds4f-vision-q2-q4 [--token TOKEN]
./download_model.sh ds4f-vision-mxfp4 [--token TOKEN]
./download_model.sh ds4f-vision-encoder [--token TOKEN]
./download_model.sh ds4f-vision-dspark [--token TOKEN]
./download_model.sh ds41f-q2 [--token TOKEN]
./download_model.sh ds41f-q4 [--token TOKEN]
./download_model.sh ds41f-mxfp4 [--token TOKEN]
./download_model.sh ds41f-vision [--token TOKEN]
./download_model.sh pro-q2-imatrix [--token TOKEN]
./download_model.sh pro-q4-layers00-30 [--token TOKEN]
./download_model.sh pro-q4-layers31-output [--token TOKEN]
./download_model.sh pro-q4-split [--token TOKEN]
./download_model.sh mtp [--token TOKEN]
./download_model.sh glm-unsloth-q4 [--token TOKEN]
./download_model.sh glm-antirez-iq2xxs [--token TOKEN]
./download_model.sh glm-antirez-q2 [--token TOKEN]
./download_model.sh glm-antirez-q4 [--token TOKEN]
./download_model.sh glm53-full-q2 [--token TOKEN]
./download_model.sh glm53-q2 [--token TOKEN]
./download_model.sh glm53-q4 [--token TOKEN]
./download_model.sh glm53-fp8 [--token TOKEN]
./download_model.sh glm53-vision [--token TOKEN]
./download_model.sh qwen38-q2 [--token TOKEN]
./download_model.sh qwen38-q4k [--token TOKEN]
./download_model.sh qwen38-vision [--token TOKEN]
Targets:
q2-imatrix
ds4f-q2
2-bit routed experts, about 81 GB on disk.
Recommended model for 96 and 128 GB RAM machines.
q2-q4-imatrix
ds4f-q2-q4
Mixed Flash quant: mostly q2 routed experts, with the last 6 layers
using q4 routed experts. About 98 GB on disk. Good for higher
quality inference for 128 GB MacBooks. Works on DGX Spark but loading
may struggle compared to q2-imatrix.
may struggle compared to ds4f-q2.
q4-imatrix
ds4f-q4
4-bit routed experts, about 153 GB on disk.
Recommended model for machines with 256 GB RAM or more.
ds4f-mxfp4
Native DeepSeek V4 Flash MXFP4 routed experts, about 156 GB on disk.
Supported by Metal and CUDA; Blackwell uses FP4 tensor cores for batched
expert work, while CUDA decode keeps Q8 activations.
ds4f-dspark
Optional DSpark speculative decoding support GGUF for Flash 0731, about
6 GB. Enable it with --dspark and --mtp-model when running ds4 or ds4-server.
ds4f-vision-q2
DeepSeek V4 Flash Vision Experimental with 2-bit routed experts, about
81 GiB, plus its vision encoder. Recommended for 96 and 128 GB Macs.
ds4f-vision-q2-q4
Mixed Vision Experimental quant with routed experts in layers 37-42 at
Q4_K and the others at 2 bits, about 91 GiB, plus its vision encoder.
ds4f-vision-mxfp4
Native MXFP4 Vision Experimental model, about 145 GiB, plus its vision
encoder. Intended for CUDA or two 128 GB Macs using tensor parallelism.
ds4f-vision-encoder
Standalone Vision Experimental encoder, about 0.9 GiB. Use this when
the matching language GGUF is already present.
ds4f-vision-dspark
Matching DSpark speculative decoding support for Vision Experimental,
about 5.6 GiB. It is not compatible with the 0731 language checkpoint.
ds41f-q2
DeepSeek V4.1 Flash calibrated Q2, about 341 GiB on disk. Main weights
occupy 152 GiB; Engram tables stay on disk. Metal only: use SSD streaming
on one 128 GB Mac, tensor parallelism on two, or a larger resident Mac.
ds41f-vision
Matching V4.1 Flash vision encoder, about 0.9 GiB. Add --vision FILE
to the language-model command. Does not update ./ds4flash.gguf.
ds41f-q4
DeepSeek V4.1 Flash calibrated Q4, about 483 GiB on disk. Main weights
occupy 294 GiB; Engram tables stay on disk. Metal only: use SSD streaming
on smaller Macs or full residency on a 512 GB Mac. Downloads two parts
and joins them automatically; allow another 37 GiB of free disk space.
ds41f-mxfp4
DeepSeek V4.1 Flash 1017 MXFP4 routed experts, about 467 GiB on disk.
Downloads and verifies all nine GGUF shards, then restores the original
page-aligned GGUF for DwarfStar. Retains the shards: allow about 935 GiB
of disk space for shards plus output. Downloads and merging can resume.
pro-q2-imatrix
DeepSeek V4 PRO q2 imatrix quant, as a single GGUF file. About 430 GB
on disk; intended for 512 GB RAM machines.
DeepSeek V4 PRO 0813 q2 imatrix quant, as a single GGUF file. About
430 GB on disk; intended for 512 GB RAM machines.
pro-q4-layers00-30
First half of the DeepSeek V4 PRO Q4 routed-expert quant, layers 0..30.
@@ -78,10 +168,6 @@ Targets:
Downloads both PRO Q4 split files into the download directory. About
838 GB total. This target does not update ./ds4flash.gguf.
mtp Optional speculative decoding component, about 3.5 GB on disk.
It is useful with q2-imatrix, q2-q4-imatrix, and q4-imatrix, but must be
enabled explicitly with --mtp when running ds4 or ds4-server.
glm-unsloth-q4
GLM 5.2 Unsloth UD-Q4_K_XL quant from unsloth/GLM-5.2-GGUF.
Downloads all 11 shards and links ./ds4flash.gguf to the first shard.
@@ -98,6 +184,46 @@ Targets:
GLM 5.2 antirez routed Q4_K GGUF from antirez/GLM-5.2-GGUF.
About 434 GB on disk.
glm53-full-q2
Full GLM 5.3 routed IQ2_XXS/Q2_K GGUF, about 197 GiB on disk.
Intended for 256 GB machines or SSD streaming on smaller systems.
glm53-q2
GLM 5.3 Flash imatrix Q2 GGUF, about 90 GiB on disk. Intended for
resident inference on 128 GB Macs.
glm53-q4
GLM 5.3 Flash Q4_K GGUF, about 178 GiB on disk. Intended for two-Mac
tensor parallelism or single-Mac SSD streaming.
glm53-fp8
Text-only GLM 5.3 Flash native FP8 GGUF, about 305 GiB on disk. It
preserves the released weights without requantization. DwarfStar
inference support for this paired FP8-code/scale format is pending.
glm53-vision
GLM 5.3 Flash vision encoder, about 1.1 GB on disk. Load it separately
with --vision; this target does not update ./ds4flash.gguf.
qwen38-q2 (alias: qwen38-iq2)
Qwen3.8-Flash-Next Q2: one 137.10 GiB GGUF. Main/MTP weights
occupy 41.73 GiB; original BF16 n-grams stay on disk.
Imatrix IQ2_XXS gate/up and Q2_K down experts (640 logical inputs,
padded to 768 on disk), with higher-precision dense/control tensors.
Smaller option for 64 GB Macs: start with --ctx 8192 and
--prefill-chunk 1024. Context buffers need additional RAM.
Keep the file on a local SSD. Add --mtp for speculation.
qwen38-q4k
Qwen3.8-Flash-Next Q4: one 165.11 GiB GGUF. Main/MTP weights
occupy 69.74 GiB; original BF16 n-grams stay on disk.
Imatrix Q4_K gate/up and MXFP4 down routed experts. Fits 128 GB
Macs with room for context buffers. Add --mtp for speculation.
qwen38-vision
Qwen3.8-Flash-Next vision encoder, about 0.6 GB on disk. Load it
with --vision; this target does not update ./ds4flash.gguf.
Options:
--token TOKEN Hugging Face token. Otherwise HF_TOKEN or the local HF token
cache is used if present.
@@ -113,10 +239,13 @@ Then the default commands work:
./ds4 -p "Hello"
./ds4-server --ctx 100000
After downloading mtp, enable it explicitly, for example:
./ds4 --mtp <download directory>/$MTP_FILE --mtp-draft 2
Qwen3.8 includes its n-grams and MTP; add --mtp to enable speculation:
./ds4 --mtp
PRO and GLM files are downloaded with the official Hugging Face downloader
After downloading DSpark support, enable it explicitly:
./ds4 --dspark --mtp-model <download directory>/$DS4F_DSPARK_FILE
PRO, V4.1, GLM and Qwen files use the official Hugging Face downloader
because they are too large, sharded, or nested for the curl path used by the
smaller DeepSeek Flash GGUF files.
EOF
@@ -135,9 +264,60 @@ FORCE_HF_DOWNLOAD=0
FLATTEN_DOWNLOADS=0
case "$MODEL" in
q2-imatrix) MODEL_FILE=$Q2_IMATRIX_FILE ;;
q2-q4-imatrix) MODEL_FILE=$Q2_Q4_IMATRIX_FILE ;;
q4-imatrix) MODEL_FILE=$Q4_IMATRIX_FILE ;;
ds4f-q2) MODEL_FILE=$DS4F_Q2_FILE ;;
ds4f-q2-q4) MODEL_FILE=$DS4F_Q2_Q4_FILE ;;
ds4f-q4) MODEL_FILE=$DS4F_Q4_FILE ;;
ds4f-mxfp4) MODEL_FILE=$DS4F_MXFP4_FILE; FORCE_HF_DOWNLOAD=1 ;;
ds4f-dspark) MODEL_FILE=$DS4F_DSPARK_FILE; LINK_MODEL=0 ;;
ds4f-vision-q2)
MODEL_FILE=$DS4F_VISION_Q2_FILE
MODEL_FILES="$MODEL_FILE $DS4F_VISION_ENCODER_FILE"
FORCE_HF_DOWNLOAD=1
;;
ds4f-vision-q2-q4)
MODEL_FILE=$DS4F_VISION_Q2_Q4_FILE
MODEL_FILES="$MODEL_FILE $DS4F_VISION_ENCODER_FILE"
FORCE_HF_DOWNLOAD=1
;;
ds4f-vision-mxfp4)
MODEL_FILE=$DS4F_VISION_MXFP4_FILE
MODEL_FILES="$MODEL_FILE $DS4F_VISION_ENCODER_FILE"
FORCE_HF_DOWNLOAD=1
;;
ds4f-vision-encoder)
MODEL_FILE=$DS4F_VISION_ENCODER_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
ds4f-vision-dspark)
MODEL_FILE=$DS4F_VISION_DSPARK_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
ds41f-q2)
REPO=$DS41_REPO
MODEL_FILE=$DS41_Q2_FILE
FORCE_HF_DOWNLOAD=1
;;
ds41f-q4)
REPO=$DS41_REPO
MODEL_FILE=$DS41_Q4_FILE
FORCE_HF_DOWNLOAD=1
;;
ds41f-mxfp4)
REPO=$DS41_REPO
MODEL_FILE="$DS41_MXFP4_BASE.gguf"
for part in 00001 00002 00003 00004 00005 00006 00007 00008 00009; do
MODEL_FILES="$MODEL_FILES $DS41_MXFP4_BASE-${part}-of-00009.gguf"
done
FORCE_HF_DOWNLOAD=1
;;
ds41f-vision)
REPO=$DS41_REPO
MODEL_FILE=$DS41_VISION_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
pro-q2-imatrix) MODEL_FILE=$PRO_Q2_IMATRIX_FILE ;;
pro-q4-layers00-30) MODEL_FILE=$PRO_Q4_LAYERS00_30_FILE; LINK_MODEL=0 ;;
pro-q4-layers31-output) MODEL_FILE=$PRO_Q4_LAYERS31_OUTPUT_FILE; LINK_MODEL=0 ;;
@@ -145,7 +325,6 @@ case "$MODEL" in
MODEL_FILES="$PRO_Q4_LAYERS00_30_FILE $PRO_Q4_LAYERS31_OUTPUT_FILE"
LINK_MODEL=0
;;
mtp) MODEL_FILE=$MTP_FILE; LINK_MODEL=0 ;;
glm-unsloth-q4)
REPO=$GLM_UNSLOTH_REPO
MODEL_FILE=$GLM_UNSLOTH_Q4_FIRST_FILE
@@ -171,6 +350,49 @@ case "$MODEL" in
MODEL_FILE=$GLM_ANTIREZ_Q4_FILE
FORCE_HF_DOWNLOAD=1
;;
glm53-full-q2)
REPO=$GLM53_FULL_REPO
MODEL_FILE=$GLM53_FULL_Q2_FILE
FORCE_HF_DOWNLOAD=1
;;
glm53-q2)
REPO=$GLM53_REPO
MODEL_FILE=$GLM53_Q2_FILE
FORCE_HF_DOWNLOAD=1
;;
glm53-q4)
REPO=$GLM53_REPO
MODEL_FILE=$GLM53_Q4_FILE
FORCE_HF_DOWNLOAD=1
;;
glm53-fp8)
REPO=$GLM53_REPO
MODEL_FILE=$GLM53_FP8_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
glm53-vision)
REPO=$GLM53_REPO
MODEL_FILE=$GLM53_VISION_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
qwen38-q2|qwen38-iq2)
REPO=$QWEN38_REPO
MODEL_FILE=$QWEN38_Q2_FILE
FORCE_HF_DOWNLOAD=1
;;
qwen38-q4k)
REPO=$QWEN38_REPO
MODEL_FILE=$QWEN38_Q4_FILE
FORCE_HF_DOWNLOAD=1
;;
qwen38-vision)
REPO=$QWEN38_MMPROJ_REPO
MODEL_FILE=$QWEN38_VISION_FILE
FORCE_HF_DOWNLOAD=1
LINK_MODEL=0
;;
-h|--help|help)
usage
exit 0
@@ -224,6 +446,12 @@ find_hf_command() {
printf '%s\n' hf
return 0
fi
for dir in "$HOME"/Library/Python/*/bin "$HOME"/.local/bin; do
if [ -x "$dir/hf" ]; then
printf '%s\n' "$dir/hf"
return 0
fi
done
return 1
}
@@ -235,6 +463,99 @@ local_download_name() {
fi
}
artifact_identity() {
case "$1" in
"$QWEN38_Q2_FILE")
expected_bytes=147207127040
expected_sha=b1b93fa69aca5f187b0fb813aca8f3ec1beb5cf8cf0bd38cf041b93e0b6ccac9
;;
"$QWEN38_Q4_FILE")
expected_bytes=177280286720
expected_sha=680944460a8cbe93ba8b6d7b6107213ffb7e22320bd913000e563ca0a0f25a8a
;;
"$DS41_Q2_FILE")
expected_bytes=365713686528
expected_sha=5296e85ff64b42ecb62c8d0c0c54e1caa07e3df8e93112f2365ebbaed62ce93a
;;
"$DS41_Q4_FILE")
expected_bytes=518596067328
expected_sha=a5e2e2c3ada4b2e98d9f9e4b50f6d9c2a12c2c96f5da165c07e13aff9264984e
;;
"$DS41_Q4_FILE.part1")
expected_bytes=480000000000
expected_sha=6442b1f9224079662c02003c0ef9ef6be6e2aff509510f681dab9e6cc41df246
;;
"$DS41_Q4_FILE.part2")
expected_bytes=38596067328
expected_sha=7c3e10646c918eeaffbc39305a75ec96117450262c61454ff194cef00d7617f0
;;
"$DS41_MXFP4_BASE.gguf")
expected_bytes=501609136128
expected_sha=79608a69cd5bb919d8140c37105edc3a33419692d4e6caf4e68c901e4b359141
;;
"$DS41_MXFP4_BASE-00001-of-00009.gguf")
expected_bytes=49353784544
expected_sha=636918b09de6253bf41af29f06aac3ceeda220e46379d8a2058ef92c4e433ed8
;;
"$DS41_MXFP4_BASE-00002-of-00009.gguf")
expected_bytes=49412635712
expected_sha=b26d3fc1e6ae1be6c52665542f3c1fee6a33791478f26e6cd90b5a75940a9f96
;;
"$DS41_MXFP4_BASE-00003-of-00009.gguf")
expected_bytes=49582960352
expected_sha=20c7fe7adb0578336d8a5e5f9eed521aeb88e8c119ba46e0916cdaab28931c80
;;
"$DS41_MXFP4_BASE-00004-of-00009.gguf")
expected_bytes=49418206432
expected_sha=67592ba15870c4292a615666adcc2736b53cc345b2cccb1c36710a0a6caf15b4
;;
"$DS41_MXFP4_BASE-00005-of-00009.gguf")
expected_bytes=49412829632
expected_sha=ecd04e0aa71bfab8cfe79ea7fef7f35b6be0b6772455139e2081d6ce1dcc18ed
;;
"$DS41_MXFP4_BASE-00006-of-00009.gguf")
expected_bytes=49257603776
expected_sha=88214e0cc80fbae578b80c114527e6ce725199684e0b6849475c7d00b3d0f54c
;;
"$DS41_MXFP4_BASE-00007-of-00009.gguf")
expected_bytes=2406482112
expected_sha=38939faf3e4258f3ae333f0fa1954ce89515b75ceee90d37bf4ebe5fafe285e6
;;
"$DS41_MXFP4_BASE-00008-of-00009.gguf")
expected_bytes=101377628544
expected_sha=a9d1634da2ac57459a64b10b69ddedfadda7dcde0e71ca6e55bea7c293c4a9f0
;;
"$DS41_MXFP4_BASE-00009-of-00009.gguf")
expected_bytes=101380404256
expected_sha=b9832b355288c005da6ef444d02424464dd27f6ebdefeadbec2244264eb65285
;;
"$DS41_VISION_FILE")
expected_bytes=970555552
expected_sha=cc283f032b3e8b8d78aeb5fccaa14e97b859b0c53aae3cd6bffa690ddf0e9e15
;;
*) return 1 ;;
esac
}
verify_download() {
artifact_identity "$1" || return 0
if [ "$(wc -c < "$2")" -ne "$expected_bytes" ]; then
echo "Incorrect file size: $2. Move the incomplete file aside and retry." >&2
exit 1
fi
echo "Verifying SHA-256: $2"
if command -v sha256sum >/dev/null 2>&1; then
actual_sha=$(sha256sum < "$2")
else
actual_sha=$(shasum -a 256 < "$2")
fi
actual_sha=${actual_sha%% *}
if [ "$actual_sha" != "$expected_sha" ]; then
echo "Checksum mismatch: $2. The file was not accepted." >&2
exit 1
fi
}
download_one_hf() {
file=$1
local_file=$(local_download_name "$file")
@@ -245,6 +566,7 @@ download_one_hf() {
mkdir -p "$(dirname "$out")"
if [ -s "$out" ]; then
verify_download "$file" "$out"
echo "Already downloaded: $out"
return
fi
@@ -252,7 +574,7 @@ download_one_hf() {
if [ -e "$part" ]; then
echo "Found curl partial download: $part" >&2
echo "The Hugging Face downloader cannot resume curl .part files." >&2
echo "Move or remove that partial download before retrying this PRO target." >&2
echo "Move or remove that partial download before retrying this target." >&2
exit 1
fi
@@ -284,6 +606,7 @@ download_one_hf() {
echo "Hugging Face download finished but expected file is missing: $out" >&2
exit 1
fi
verify_download "$file" "$out"
}
download_one() {
@@ -325,7 +648,138 @@ download_one() {
mv "$part" "$out"
}
if [ -n "$MODEL_FILES" ]; then
download_ds41_q4() {
q4_out="$OUT_DIR/$DS41_Q4_FILE"
if [ -e "$q4_out" ]; then
verify_download "$DS41_Q4_FILE" "$q4_out"
echo "Already downloaded: $q4_out"
return
fi
if ! command -v python3 >/dev/null 2>&1; then
echo "Joining the Q4 download requires Python 3." >&2
exit 1
fi
if [ ! -e "$q4_out.assembling" ]; then
download_one_hf "$DS41_Q4_FILE.part1"
fi
download_one_hf "$DS41_Q4_FILE.part2"
artifact_identity "$DS41_Q4_FILE"
q4_bytes=$expected_bytes
q4_sha=$expected_sha
artifact_identity "$DS41_Q4_FILE.part1"
python3 - "$q4_out" "$expected_bytes" "$q4_bytes" "$q4_sha" <<'PY'
import fcntl
import hashlib
import os
from pathlib import Path
import shutil
import sys
out = Path(sys.argv[1])
boundary, total = map(int, sys.argv[2:4])
expected = sys.argv[4]
first, second, pending = (Path(str(out) + suffix)
for suffix in (".part1", ".part2", ".assembling"))
def verify(path):
if path.stat().st_size != total:
sys.exit("Incorrect file size: " + str(path))
print("Verifying SHA-256: " + str(path), flush=True)
digest = hashlib.sha256()
with path.open("rb") as fp:
for block in iter(lambda: fp.read(16 << 20), b""):
digest.update(block)
if digest.hexdigest() != expected:
sys.exit("Checksum mismatch: " + str(path) +
". Move this file aside before retrying; it was not accepted.")
# Keep this lock file: unlinking it could allow two different locks for the
# same download. The first part becomes the output without a second full copy.
with Path(str(pending) + ".lock").open("a") as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
if out.exists():
verify(out)
sys.exit(0)
if not pending.exists():
if first.stat().st_size != boundary:
sys.exit("Incorrect file size: " + str(first))
first.rename(pending)
size = pending.stat().st_size
if not boundary <= size <= total:
sys.exit("Invalid partial assembly: " + str(pending) +
". Move it aside before retrying.")
if second.stat().st_size != total - boundary:
sys.exit("Incorrect file size: " + str(second))
# An interrupted append restarts only the smaller tail, never the prefix.
with pending.open("r+b") as dst:
dst.truncate(boundary)
dst.seek(boundary)
if shutil.disk_usage(out.parent).free < total - boundary + (1 << 30):
sys.exit("Not enough disk space to join Q4; keep the parts and retry.")
print("Joining Q4 download parts", flush=True)
with second.open("rb") as src:
shutil.copyfileobj(src, dst, 16 << 20)
dst.flush()
os.fsync(dst.fileno())
verify(pending)
pending.rename(out)
second.unlink()
PY
}
download_ds41_mxfp4() {
mx_out="$OUT_DIR/$MODEL_FILE"
if [ -e "$mx_out" ]; then
verify_download "$MODEL_FILE" "$mx_out"
echo "Already downloaded: $mx_out"
return
fi
if ! command -v python3 >/dev/null 2>&1; then
echo "Merging the MXFP4 shards requires Python 3." >&2
exit 1
fi
if [ ! -f "$ROOT/gguf-tools/merge_deepseek41_shards.py" ]; then
echo "Missing gguf-tools/merge_deepseek41_shards.py; use a complete checkout." >&2
exit 1
fi
mkdir -p "$OUT_DIR"
artifact_identity "$MODEL_FILE"
mx_bytes=$expected_bytes
mx_sha=$expected_sha
set -- "$mx_out.assembling" "$mx_bytes"
for shard in $MODEL_FILES; do
artifact_identity "$shard"
set -- "$@" "$OUT_DIR/$shard" "$expected_bytes"
done
# Check peak storage before downloading hundreds of GB. Retain the input
# shards and account for a resumable output; HF partial caches are extra.
python3 - "$OUT_DIR" "$@" <<'PY'
from pathlib import Path
import shutil
import sys
needed = 1 << 30
for name, size in zip(sys.argv[2::2], sys.argv[3::2]):
path = Path(name)
needed += max(0, int(size) - (path.stat().st_size if path.exists() else 0))
if shutil.disk_usage(sys.argv[1]).free < needed:
sys.exit("Not enough disk space for MXFP4 shards and merged GGUF: "
f"need {needed / (1 << 30):.1f} GiB free. Existing files were kept.")
PY
set --
for shard in $MODEL_FILES; do
download_one_hf "$shard"
set -- "$@" "$OUT_DIR/$shard"
done
python3 "$ROOT/gguf-tools/merge_deepseek41_shards.py" \
--output "$mx_out" --expected-bytes "$mx_bytes" --sha256 "$mx_sha" "$@"
}
if [ "$MODEL" = "ds41f-mxfp4" ]; then
download_ds41_mxfp4
elif [ "$MODEL" = "ds41f-q4" ]; then
download_ds41_q4
elif [ -n "$MODEL_FILES" ]; then
for file in $MODEL_FILES; do
download_one "$file"
done
@@ -333,11 +787,14 @@ else
download_one "$MODEL_FILE"
fi
if [ "$MODEL" = "mtp" ]; then
if [ "$MODEL" = "ds4f-dspark" ]; then
echo
echo "MTP is an optional component for q2-imatrix, q2-q4-imatrix, and q4-imatrix."
echo "Enable it explicitly, for example:"
echo " ./ds4 --mtp $OUT_DIR/$MTP_FILE --mtp-draft 2"
echo "DSpark support downloaded. Enable it explicitly:"
echo " ./ds4 --dspark -m ./ds4flash.gguf --mtp-model $OUT_DIR/$DS4F_DSPARK_FILE"
elif [ "$MODEL" = "ds4f-vision-dspark" ]; then
echo
echo "Vision Experimental DSpark support downloaded. Use it only with the matching checkpoint:"
echo " ./ds4 --dspark -m ./ds4flash.gguf --mtp-model $OUT_DIR/$DS4F_VISION_DSPARK_FILE"
elif [ "$MODEL" = "pro-q4-layers00-30" ] || [ "$MODEL" = "pro-q4-layers31-output" ] || [ "$MODEL" = "pro-q4-split" ]; then
echo
echo "Downloaded PRO Q4 distributed split file(s). Use them with --layers,"
@@ -350,3 +807,15 @@ fi
echo
echo "Done."
if [ "$MODEL" = qwen38-q2 ] || [ "$MODEL" = qwen38-iq2 ]; then
echo "Run with an 8K starting context:"
echo " ./ds4 --ctx 8192 --prefill-chunk 1024"
echo "Add --mtp to enable speculation."
elif [ "$MODEL" = qwen38-q4k ]; then
echo "Run ./ds4. Add --mtp to enable speculation."
fi
if [ "$MODEL" = qwen38-vision ]; then
echo
echo "Qwen3.8 vision encoder downloaded. Pass it with --vision, for example:"
printf ' ./ds4 --vision "%s/%s"\n' "$OUT_DIR" "$QWEN38_VISION_FILE"
fi
ds4.c
+51497 -3515
View File
File diff suppressed because it is too large. Load diff
ds4.h
+319 -4
View File
@@ -26,7 +26,11 @@ typedef enum {
DS4_THINK_NONE,
DS4_THINK_HIGH,
DS4_THINK_MAX,
DS4_THINK_LOW, /* Qwen3.8 reasoning_effort low/medium; other models render them as HIGH */
DS4_THINK_MEDIUM,
} ds4_think_mode;
/* Explicit numeric effort lives outside the stable named-mode values. */
#define DS4_THINK_LEVEL_BASE 1000
typedef enum {
DS4_LOG_DEFAULT,
@@ -91,15 +95,50 @@ typedef struct {
bool debug;
} ds4_distributed_options;
/* Tensor parallelism: two identical machines run the model in lockstep and
* split the heavy per-layer matvecs, exchanging partial sums at gates inside
* the graph (see misc/METAL_TENSOR_PARALLELISM.md). Each rank keeps one
* contiguous half of the routed experts resident; dense and shared weights
* remain replicated. The leader owns prompt/sampling and listens; the worker
* dials in and mirrors every session sync/eval. */
typedef enum {
DS4_TP_NONE = 0,
DS4_TP_LEADER,
DS4_TP_WORKER,
} ds4_tp_role;
typedef enum {
DS4_TP_TRANSPORT_AUTO = 0,
DS4_TP_TRANSPORT_RDMA,
DS4_TP_TRANSPORT_TCP,
} ds4_tp_transport;
typedef struct {
ds4_tp_role role;
bool requested; /* --tensor-parallel with shared role options */
const char *listen_host; /* leader listens here for the worker */
int listen_port;
const char *leader_host; /* worker dials the leader */
int leader_port;
ds4_tp_transport transport;
const char *rdma_device;
int rdma_gid_index;
bool rdma_gid_index_set;
bool glm_token_prefill;
int debug_hash; /* cross-check hidden state every N tokens */
} ds4_tp_options;
typedef struct {
const char *model_path;
const char *mtp_path;
const char *vision_path;
ds4_backend backend;
int n_threads;
int context_size;
uint32_t prefill_chunk;
int mtp_draft_tokens;
float mtp_margin;
float dspark_confidence_threshold;
const char *directional_steering_file;
const char *expert_profile_path;
float directional_steering_attn;
@@ -112,10 +151,23 @@ typedef struct {
uint64_t simulate_used_memory_bytes;
bool warm_weights;
bool quality;
bool glm_mtp;
bool glm_mtp_timing;
bool dspark;
bool dspark_strict;
bool dspark_exact_sampling;
bool dspark_confidence_threshold_set;
bool cuda_tensor_parallel;
bool ssd_streaming;
bool ssd_streaming_cold;
bool ssd_streaming_full_layers_set;
bool inspect_only;
/* Multi-GPU placement uses this to price per-layer KV storage. */
int placement_ctx_hint;
/* Number of independently allocated session graphs/caches to reserve. */
int placement_session_count_hint;
/* Server batch mode serializes execution and can share prefill scratch. */
bool share_session_prefill_workspace;
bool first_token_test;
bool metal_graph_test;
bool load_slice;
@@ -123,8 +175,27 @@ typedef struct {
uint32_t load_layer_end;
bool load_output;
ds4_distributed_options distributed;
ds4_tp_options tp;
} ds4_engine_options;
typedef struct {
float *data;
uint32_t token_count;
uint32_t layout;
uint32_t grid_width;
uint32_t grid_height;
uint32_t width;
uint32_t height;
uint32_t content_width;
uint32_t content_height;
uint8_t fingerprint[32];
} ds4_vision_embedding;
typedef struct {
uint32_t token_start;
ds4_vision_embedding embedding;
} ds4_vision_span;
typedef void (*ds4_token_emit_fn)(void *ud, int token);
typedef void (*ds4_generation_done_fn)(void *ud);
@@ -150,15 +221,78 @@ typedef struct {
} ds4_session_payload_file;
int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt);
/* Multi-GPU pipeline-parallel entry point (wave 2).
*
* Accepts an optional ds4_gpu_config (defined in ds4_gpu_mgpu.h) that
* lets callers describe a multi-GPU placement target. Passing NULL is
* back-compatible with ds4_engine_open and produces identical engine
* state — bit-equivalent execution at runtime.
*
* When a non-NULL config is supplied AND the computed placement spans
* more than one tier (either multiple GPUs or any CPU-spill), this
* wave-2 implementation prints the layout and refuses to open: full
* multi-tier execution wiring lands in a follow-up task
* (mgpu-graph-session-execution). Callers receive a non-zero return
* and a documented stderr notice. */
/* ds4_gpu_config is declared in ds4_gpu_mgpu.h, which callers should
* include separately. We forward-declare it here so this header can be
* used as-is (callers passing NULL don't need the struct definition). */
struct ds4_gpu_config;
int ds4_engine_create_with_gpu_config(ds4_engine **out,
const ds4_engine_options *opt,
const struct ds4_gpu_config *gpu_cfg);
void ds4_engine_close(ds4_engine *e);
void ds4_engine_summary(ds4_engine *e);
int ds4_engine_vocab_size(ds4_engine *e);
uint32_t ds4_engine_prefill_chunk(ds4_engine *e);
int ds4_engine_power(ds4_engine *e);
int ds4_engine_set_power(ds4_engine *e, int power_percent);
const char *ds4_engine_model_name(ds4_engine *e);
int ds4_engine_layer_count(ds4_engine *e);
/* Decode gate schedule for the TP transport; see ds4_tp_identity. */
enum { DS4_TP_GATE_MASK_WORDS = 3 };
void ds4_engine_tp_gate_schedule(ds4_engine *e,
uint32_t *start,
uint32_t *step,
uint32_t *per_token,
uint64_t mask[DS4_TP_GATE_MASK_WORDS]);
uint32_t ds4_engine_layer_compress_ratio(ds4_engine *e, uint32_t layer);
uint64_t ds4_engine_hidden_f32_values(ds4_engine *e);
int ds4_engine_embd_dim(ds4_engine *e);
uint64_t ds4_engine_model_bytes(ds4_engine *e);
bool ds4_engine_has_vision(ds4_engine *e);
int ds4_engine_vision_encode_file(ds4_engine *e,
const char *path,
ds4_vision_embedding *out,
char *error,
size_t error_cap);
int ds4_engine_vision_encode_memory(ds4_engine *e,
const uint8_t *encoded,
size_t encoded_len,
ds4_vision_embedding *out,
char *error,
size_t error_cap);
void ds4_vision_embedding_free(ds4_vision_embedding *embedding);
int ds4_prompt_append_vision(ds4_engine *e,
ds4_tokens *tokens,
ds4_vision_span *span,
ds4_vision_embedding *embedding,
char *error,
size_t error_cap);
/* Append one user or tool message whose text parts alternate with images.
* text_parts must contain image_count + 1 entries. On success ownership of
* each embedding is transferred to the corresponding output span. */
int ds4_chat_append_multimodal_message(ds4_engine *e,
ds4_tokens *tokens,
const char *role,
const char *const *text_parts,
ds4_vision_embedding *embeddings,
size_t image_count,
ds4_vision_span *spans,
char *error,
size_t error_cap);
int ds4_engine_tp_vocab_split(ds4_engine *e);
bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e,
uint32_t layer,
uint32_t full_live,
@@ -172,8 +306,14 @@ bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e,
* Pro and later shapes must use nonzero ids. */
int ds4_engine_model_id(ds4_engine *e);
bool ds4_engine_is_glm_dsa(ds4_engine *e);
bool ds4_engine_is_glm53(ds4_engine *e);
bool ds4_engine_is_qwen4(ds4_engine *e);
/* Qwen3.8 reasoning-effort system instruction for a think mode (NULL when none) */
const char *ds4_qwen4_reasoning_effort_text(ds4_think_mode mode);
const char *ds4_backend_name(ds4_backend backend);
bool ds4_think_mode_enabled(ds4_think_mode mode);
int ds4_think_mode_level(ds4_think_mode mode);
bool ds4_think_mode_parse_level(const char *text, ds4_think_mode *out);
const char *ds4_think_mode_name(ds4_think_mode mode);
const char *ds4_think_max_prefix(void);
const char *ds4_glm_reasoning_effort_text(ds4_think_mode mode);
@@ -205,10 +345,20 @@ int ds4_engine_collect_imatrix(ds4_engine *e,
const char *output_path,
int ctx_size,
int max_prompts,
int max_tokens);
int max_tokens,
int min_expert_samples);
void ds4_engine_dump_tokens(ds4_engine *e, const ds4_tokens *tokens);
int ds4_dump_text_tokenization(const char *model_path, const char *text, FILE *fp);
int ds4_dump_chat_tokenization(const char *model_path,
const char *system,
const char *prompt,
ds4_think_mode think_mode,
int ctx_size,
FILE *fp);
int ds4_engine_head_test(ds4_engine *e, const ds4_tokens *prompt);
bool ds4_engine_is_glm_dsa(ds4_engine *e);
bool ds4_engine_is_deepseek41(ds4_engine *e);
const char *ds4_deepseek41_reasoning_effort_text(ds4_think_mode mode);
int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt);
int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt);
int ds4_engine_metal_graph_full_test(ds4_engine *e, const ds4_tokens *prompt);
@@ -229,6 +379,7 @@ void ds4_encode_chat_prompt(
ds4_think_mode think_mode,
ds4_tokens *out);
void ds4_chat_append_max_effort_prefix(ds4_engine *e, ds4_tokens *tokens);
void ds4_chat_append_think_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode mode);
void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role, const char *content);
void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode);
@@ -242,18 +393,33 @@ bool ds4_token_is_stop_for_think_mode(ds4_engine *e,
int ds4_token_user(ds4_engine *e);
int ds4_token_assistant(ds4_engine *e);
/* Tensor-parallel binding: allocates the GPU gate slab, registers it with
* the transport and arms the per-layer gate machinery. Call once, after
* ds4_tp_create() and before any session work. Transport lifecycle stays
* with the caller. */
struct ds4_tp;
int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen);
/* Release gate resources before freeing a bound transport. Sessions must
* already be closed. Also called by ds4_engine_close(). */
void ds4_engine_tp_unbind(ds4_engine *e);
int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size);
void ds4_session_free(ds4_session *s);
int ds4_session_power(ds4_session *s);
int ds4_session_set_power(ds4_session *s, int power_percent);
float ds4_session_directional_steering_ffn(ds4_session *s);
/* Change steering for future evaluation without rebuilding the existing KV
* state. Live changes are currently limited to non-distributed sessions. */
int ds4_session_set_directional_steering_ffn(ds4_session *s, float scale);
bool ds4_session_is_distributed(ds4_session *s);
void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
/* UI-only progress. It may report fine-grained progress inside a prefill chunk;
* callers must not treat it as a durable KV checkpoint boundary. */
void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
/* Optional cooperative cancellation. ds4_session_sync() checks it only at
* safe boundaries where the live checkpoint is either unchanged or represents a
* valid token prefix, and returns DS4_SESSION_SYNC_INTERRUPTED when it stops. */
/* Cooperative cancellation for ds4_session_sync(), which drains pending work
* before returning DS4_SESSION_SYNC_INTERRUPTED. A complete prefix may remain
* usable, but a partial layer-major pass is invalidated and the next sync
* rebuilds it. Do not assume an interrupted session can be saved or decoded. */
void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud);
void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total);
/* Distributed coordinator sessions return 1 when the full layer route is
@@ -271,27 +437,175 @@ typedef enum {
/* Synchronize the live session to a full prompt token prefix. If the current
* checkpoint is a prefix, only the suffix is evaluated; otherwise the backend
* state is refilled from scratch. */
#define DS4_SESSION_SYNC_INTERRUPTED 2
int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen);
int ds4_session_sync_multimodal(ds4_session *s,
const ds4_tokens *prompt,
const ds4_vision_span *images,
size_t image_count,
char *err,
size_t errlen);
/* A reusable image prefix has unchanged spans/fingerprints for all historical
* images, with any new images starting at or after the live token frontier.
* The caller must also check the token prefix. Invalid checkpoints never match. */
bool ds4_session_vision_prefix_matches(const ds4_session *s,
const ds4_vision_span *images,
size_t image_count);
/* Like the prefix check, but also require exactly the same image count. */
bool ds4_session_vision_state_matches(const ds4_session *s,
const ds4_vision_span *images,
size_t image_count);
/* Fingerprint-only image prefix: every historical image matches by row count
* and fingerprint, ignoring token positions (which rebase repairs) and
* allowing appended request images. Lets the server routing probe admit
* image-append continuations while still rejecting pixel mismatches. */
bool ds4_session_vision_fingerprint_prefix_matches(const ds4_session *s,
const ds4_vision_span *images,
size_t image_count);
/* Restore image positions from an independently authenticated live continuation
* (for example, matching tool-call IDs). Checks every fingerprint and row count;
* on failure, leaves spans unchanged. This does not verify the text history. */
bool ds4_session_rebase_vision_state(const ds4_session *s,
ds4_vision_span *images, size_t image_count);
/* True while a session contains, or is actively syncing, image-conditioned
* state. Such state must not be written to the text-keyed disk KV cache. */
bool ds4_session_has_vision_state(const ds4_session *s);
bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common);
ds4_session_rewrite_result ds4_session_rewrite_from_common(
ds4_session *s, const ds4_tokens *prompt, int common,
char *err, size_t errlen);
int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt);
bool ds4_session_checkpoint_valid(const ds4_session *s);
/* Test helpers (ds4-test): allocate a session shell holding only the given
* checkpoint tokens, for server-side routing/probe unit tests. Not usable
* for inference; free with ds4_session_free_test_checkpoint(). */
ds4_session *ds4_session_new_test_checkpoint(const int *tokens, int n);
void ds4_session_free_test_checkpoint(ds4_session *s);
/* Attach synthetic image identities to a test checkpoint (copies
* token_start/row-count/fingerprint per span). Not usable for inference. */
void ds4_session_set_test_images(ds4_session *s,
const ds4_vision_span *images, size_t n);
int ds4_session_argmax(ds4_session *s);
int ds4_session_argmax_excluding(ds4_session *s, int excluded_id);
int ds4_session_argmax_ignoring_eos(ds4_session *s,
ds4_think_mode think_mode);
int ds4_sample_logits(const float *logits, int n_vocab, float temperature,
int top_k, float top_p, float min_p, uint64_t *rng);
int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng);
#ifdef DS4_TEST_HOOKS
/* Teacher-force a V4.1 verification block, returning logits after each token. */
int ds4_test_ds41_verify_tokens(ds4_session *s, const int *tokens, int count,
float *logits, char *err, size_t errlen);
int ds4_test_sample_logits(const float *logits, uint32_t n_vocab,
float temperature, int top_k,
float top_p, float min_p, uint64_t *rng,
float *prob_scratch);
int ds4_test_sampling_probabilities(const float *logits, uint32_t n_vocab,
float temperature, int top_k,
float top_p, float min_p, float *probs);
int ds4_test_speculative_sample(const float *target_logits,
const float *draft_logits,
uint32_t n_vocab,
float temperature,
int top_k,
float top_p,
float min_p,
uint64_t *rng,
float *target_probs,
float *draft_probs);
int ds4_test_speculative_delta_sample(const float *target_logits,
uint32_t n_vocab,
int draft_token,
float temperature,
int top_k,
float top_p,
float min_p,
uint64_t *rng,
float *target_probs);
int ds4_test_argmax_excluding_logits(const float *logits, uint32_t n_vocab,
int excluded_id);
uint64_t ds4_test_mixed_native_count(void);
uint64_t ds4_test_ds41_batch_count(void);
#endif
int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k);
int ds4_session_token_logprob(ds4_session *s, int token, ds4_token_score *out);
int ds4_session_copy_logits(ds4_session *s, float *out, int cap);
int ds4_session_set_logits(ds4_session *s, const float *logits, int n);
/* Pay the one-time first-submission GPU cost outside any measured window;
* used by the TP worker right after session create (no-op on CPU/GLM). */
void ds4_session_gpu_warmup(ds4_session *s);
int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen);
typedef struct {
ds4_session *session;
int token;
} ds4_decode_item;
/* Advance independent sessions by one token each. Batch size one is exactly
* ds4_session_eval(). Backends without native batching use a correctness-first
* sequential fallback. */
int ds4_sessions_eval_batch(ds4_decode_item *items, int count,
char *err, size_t errlen);
/* One speculative cycle for a batch of sessions (greedy acceptance, Qwen3.8
* with --mtp): each item feeds its token; a pending draft rides along as a
* second row and is committed when it is the target's argmax. accepted[i]
* lists the tokens committed for item i (the fed token, then the draft) and
* n_accepted[i] how many; the session's logits then follow its last
* committed token. Engines without native batching run one cycle per
* session in turn. */
int ds4_sessions_eval_batch_speculative_argmax(ds4_decode_item *items, int count,
int (*accepted)[2], int *n_accepted,
char *err, size_t errlen);
/* Advance one resumed prefill suffix and an independent decode batch as one
* scheduling step. Unsupported combinations use the ordinary serialized
* session operations. */
int ds4_sessions_eval_batch_with_prefill(
ds4_decode_item *items, int count,
ds4_session *prefill_session, const ds4_tokens *prefill_prompt,
char *err, size_t errlen);
int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
int max_tokens, int eos_token,
int *accepted, int accepted_cap,
char *err, size_t errlen);
int ds4_session_eval_speculative_argmax_ignoring_eos(
ds4_session *s, int first_token, int max_tokens, int eos_token,
ds4_think_mode think_mode,
int *accepted, int accepted_cap, char *err, size_t errlen);
/* Evaluate one already-sampled target token and speculatively extend it.
* Positive-temperature DSpark normally commits greedily verified draft
* tokens; dspark_exact_sampling selects exact stochastic p/q acceptance for
* DSpark or an internal GLM MTP block. */
int ds4_session_eval_speculative(ds4_session *s, int first_token,
int max_tokens, int eos_token,
float temperature, int top_k,
float top_p, float min_p, uint64_t *rng,
int *accepted, int accepted_cap,
char *err, size_t errlen);
/* TP worker side of a mirrored speculative-verify block: run its half of the
* batch verify for KV side effects, then obey the leader's commit frame
* (keep, or roll back and replay). flags are the VERIFY frame's
* DS4_TP_VERIFY_* bits. Only called from ds4_tp_worker_run. */
int ds4_session_tp_spec_cycle(ds4_session *s, const int *drafts, int draft_n,
uint32_t flags, char *err, size_t errlen);
int ds4_session_glm_tp_spec_cycle(ds4_session *s, int token, int limit,
char *err, size_t errlen);
/* TP worker side of a lockstep DSpark DRAFT/MAINTAIN frame (mode is a
* DS4_TP_DSPARK_* value). Only called from ds4_tp_worker_run. */
int ds4_session_tp_dspark_draft(ds4_session *s, int token, int pos,
uint32_t mode, char *err, size_t errlen);
/* Internal TP restore: no collectives and no logits exchange. Return 0 after
* restoring a recent verifier prefix, 1 if unavailable, -1 on copy failure.
* Logits are valid only at the block start; otherwise replay before sampling. */
int ds4_session_restore_speculative_prefix(ds4_session *s, int pos);
void ds4_session_invalidate(ds4_session *s);
/* Try a recent speculative rewind, including valid next-token logits. False
* means the caller must use ordinary rewind/sync before continuing; a failed
* restore may have invalidated the checkpoint. */
bool ds4_session_rewind_speculative(ds4_session *s, int pos);
/* Keep the token prefix, restoring recurrent state where possible. This may
* copy GPU state and replay a bounded speculative suffix on both TP ranks.
* Otherwise the checkpoint becomes invalid: sync the retained prefix before eval.
* Callers retaining images must use sync_multimodal for that rebuild. */
void ds4_session_rewind(ds4_session *s, int pos);
int ds4_session_pos(ds4_session *s);
int ds4_session_ctx(ds4_session *s);
@@ -300,6 +614,7 @@ int ds4_engine_routed_quant_bits(ds4_engine *e);
bool ds4_engine_has_output_head(ds4_engine *e);
bool ds4_engine_has_mtp(ds4_engine *e);
int ds4_engine_mtp_draft_tokens(ds4_engine *e);
bool ds4_engine_mtp_exact_sampling(ds4_engine *e);
const ds4_tokens *ds4_session_tokens(ds4_session *s);
/* Low-level graph slice entry points used by distributed inference. The
ds4_agent.c
+3233 -674
View File
File diff suppressed because it is too large. Load diff
ds4_bench.c
+331 -24
View File
@@ -1,6 +1,8 @@
#include "ds4.h"
#include "ds4_distributed.h"
#include "ds4_gpu_args.h"
#include "ds4_help.h"
#include "ds4_tp.h"
/* Purpose-built throughput benchmark.
*
@@ -22,15 +24,26 @@
#include <string.h>
#include <time.h>
/* CUDA builds resolve these weak symbols through libcudart; other builds
* remain independent of CUDA headers and libraries. */
#if defined(__GNUC__) && !defined(__APPLE__)
#define DS4_BENCH_HAVE_CUDA_PROFILER 1
extern int cudaProfilerStart(void) __attribute__((weak));
extern int cudaProfilerStop(void) __attribute__((weak));
#endif
#define DS4_BENCH_DEFAULT_SNAPSHOT_MAX_BYTES (UINT64_C(1) << 30)
typedef struct {
const char *model_path;
const char *mtp_path;
const char *prompt_path;
const char *chat_prompt_path;
const char *system;
const char *csv_path;
const char *expert_profile_path;
const char *gpu_vram_arg;
const char *gpu_devices_arg;
ds4_backend backend;
int threads;
int ctx_start;
@@ -48,11 +61,18 @@ typedef struct {
double step_mul;
const char *dump_frontier_logits_dir;
ds4_dist_options dist;
ds4_tp_options tp;
bool warm_weights;
bool quality;
bool ssd_streaming;
bool ssd_streaming_cold;
bool ssd_streaming_full_layers_set;
bool cuda_tensor_parallel;
bool show_output;
bool teacher_forced_decode;
bool dspark;
bool dspark_confidence_threshold_set;
float dspark_confidence_threshold;
} bench_config;
static double bench_now_sec(void) {
@@ -229,8 +249,38 @@ static bench_config parse_options(int argc, char **argv) {
}
if (dist_parse == DS4_DIST_CLI_MATCHED) continue;
char tp_parse_err[256] = {0};
ds4_tp_cli_parse_result tp_parse =
ds4_tp_parse_cli_arg(arg,
&i,
argc,
argv,
&c.tp,
tp_parse_err,
sizeof(tp_parse_err));
if (tp_parse == DS4_TP_CLI_ERROR) {
fprintf(stderr,
"ds4-bench: %s\n",
tp_parse_err[0] ? tp_parse_err : "invalid tensor-parallel option");
exit(2);
}
if (tp_parse == DS4_TP_CLI_MATCHED) continue;
if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
c.model_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--mtp-model")) {
c.mtp_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--dspark")) {
c.dspark = true;
} else if (!strcmp(arg, "--dspark-confidence")) {
const double v = parse_double_arg(need_arg(&i, argc, argv, arg), arg);
if (v < 0.0 || v > 1.0) {
fprintf(stderr, "ds4-bench: --dspark-confidence must be between 0 and 1\n");
exit(2);
}
c.dspark = true;
c.dspark_confidence_threshold = (float)v;
c.dspark_confidence_threshold_set = true;
} else if (!strcmp(arg, "--prompt-file")) {
c.prompt_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--chat-prompt-file")) {
@@ -268,6 +318,12 @@ static bench_config parse_options(int argc, char **argv) {
} else if (!strcmp(arg, "--cuda")) {
c.backend = DS4_BACKEND_CUDA;
#endif
} else if (!strcmp(arg, "--gpu-vram")) {
c.gpu_vram_arg = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--gpu-devices")) {
c.gpu_devices_arg = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--cuda-tensor-parallel")) {
c.cuda_tensor_parallel = true;
} else if (!strcmp(arg, "--cpu")) {
c.backend = DS4_BACKEND_CPU;
} else if (!strcmp(arg, "--quality")) {
@@ -315,6 +371,10 @@ static bench_config parse_options(int argc, char **argv) {
}
} else if (!strcmp(arg, "--warm-weights")) {
c.warm_weights = true;
} else if (!strcmp(arg, "--show-output")) {
c.show_output = true;
} else if (!strcmp(arg, "--teacher-forced-decode")) {
c.teacher_forced_decode = true;
} else {
fprintf(stderr, "ds4-bench: unknown option: %s\n", arg);
usage(stderr, NULL);
@@ -326,6 +386,16 @@ static bench_config parse_options(int argc, char **argv) {
fprintf(stderr, "ds4-bench: specify exactly one of --prompt-file or --chat-prompt-file\n");
exit(2);
}
if (c.dspark && !c.mtp_path) {
fprintf(stderr, "ds4-bench: --dspark requires --mtp-model FILE\n");
exit(2);
}
if (c.dspark && c.teacher_forced_decode) {
fprintf(stderr,
"ds4-bench: --dspark cannot be combined with "
"--teacher-forced-decode\n");
exit(2);
}
if (c.ctx_start > c.ctx_max) {
fprintf(stderr, "ds4-bench: --ctx-start must be <= --ctx-max\n");
exit(2);
@@ -347,12 +417,20 @@ static bench_config parse_options(int argc, char **argv) {
fprintf(stderr, "ds4-bench: --ctx-alloc must be greater than ctx-max + gen-tokens\n");
exit(2);
}
char tp_err[256];
if (!ds4_tp_adopt_distributed_options(&c.tp,
&c.dist,
tp_err,
sizeof(tp_err))) {
fprintf(stderr, "ds4-bench: %s\n", tp_err);
exit(2);
}
char dist_err[256];
if (ds4_dist_prepare_engine_options(&c.dist, NULL, dist_err, sizeof(dist_err)) != 0) {
fprintf(stderr, "ds4-bench: %s\n", dist_err);
exit(2);
}
if (c.dist.role == DS4_DISTRIBUTED_WORKER) {
if (c.dist.role == DS4_DISTRIBUTED_WORKER || c.tp.role == DS4_TP_WORKER) {
fprintf(stderr, "ds4-bench: --role worker is a serving mode; start workers with ./ds4\n");
exit(2);
}
@@ -538,11 +616,38 @@ static void maybe_warn_distributed_step_shape(const bench_config *cfg, ds4_sessi
}
}
static void close_engine(ds4_engine *engine, ds4_tp *tp) {
if (tp) ds4_tp_send_stop(tp);
ds4_engine_close(engine);
ds4_tp_free(tp);
}
int main(int argc, char **argv) {
bench_config cfg = parse_options(argc, argv);
/* Hint the packer at the largest ctx this bench run will exercise
* so per-layer KV bytes are priced for the real session size, not
* a stale 4096 default. Single-tier and CPU paths ignore this. */
int placement_ctx_hint = cfg.ctx_max;
if (cfg.ctx_alloc > placement_ctx_hint) placement_ctx_hint = cfg.ctx_alloc;
ds4_gpu_config gpu_cfg = {0};
bool skip_cuda = false;
const bool have_gpu_config = cfg.gpu_vram_arg || cfg.gpu_devices_arg;
if (have_gpu_config) {
char gpu_err[256];
if (parse_gpu_vram_arg(cfg.gpu_vram_arg, cfg.gpu_devices_arg,
&gpu_cfg, &skip_cuda,
gpu_err, sizeof(gpu_err)) != 0) {
fprintf(stderr, "ds4-bench: %s\n", gpu_err);
return 2;
}
cfg.backend = skip_cuda ? DS4_BACKEND_CPU : DS4_BACKEND_CUDA;
}
ds4_engine_options opt = {
.model_path = cfg.model_path,
.mtp_path = cfg.mtp_path,
.backend = cfg.backend,
.n_threads = cfg.threads,
.context_size = cfg.ctx_alloc,
@@ -555,22 +660,73 @@ int main(int argc, char **argv) {
.power_percent = cfg.power_percent,
.warm_weights = cfg.warm_weights,
.quality = cfg.quality,
.dspark = cfg.dspark,
.dspark_confidence_threshold = cfg.dspark_confidence_threshold,
.dspark_confidence_threshold_set = cfg.dspark_confidence_threshold_set,
.cuda_tensor_parallel = cfg.cuda_tensor_parallel,
.ssd_streaming = cfg.ssd_streaming,
.ssd_streaming_cold = cfg.ssd_streaming_cold,
.ssd_streaming_full_layers_set = cfg.ssd_streaming_full_layers_set,
.expert_profile_path = cfg.expert_profile_path,
.distributed = cfg.dist,
.tp = cfg.tp,
};
char dist_err[256];
if (ds4_dist_prepare_engine_options(&cfg.dist, &opt, dist_err, sizeof(dist_err)) != 0) {
fprintf(stderr, "ds4-bench: %s\n", dist_err);
return 2;
}
char tp_err[256];
if (!ds4_tp_validate_engine_options(&opt, tp_err, sizeof(tp_err))) {
fprintf(stderr, "ds4-bench: %s\n", tp_err);
return 2;
}
ds4_engine *engine = NULL;
if (ds4_engine_open(&engine, &opt) != 0) return 1;
log_context_memory(cfg.backend,
if (have_gpu_config && !skip_cuda) {
const bool was_auto =
(cfg.gpu_vram_arg && !strcmp(cfg.gpu_vram_arg, "auto")) ||
(!cfg.gpu_vram_arg && cfg.gpu_devices_arg);
char layout[256];
if (format_gpu_layout_line(&gpu_cfg, was_auto,
layout, sizeof(layout)) > 0) {
fprintf(stdout, "%s\n", layout);
fflush(stdout);
}
if (ds4_engine_create_with_gpu_config(
&engine, &opt, &gpu_cfg) != 0) return 1;
} else if (ds4_engine_open(&engine, &opt) != 0) {
return 1;
}
ds4_tp *tp_leader = NULL;
if (cfg.tp.role == DS4_TP_LEADER) {
ds4_tp_identity tp_id = {
.gguf_bytes = ds4_engine_model_bytes(engine),
.model_id = (uint32_t)ds4_engine_model_id(engine),
.n_layer = (uint32_t)ds4_engine_layer_count(engine),
.n_embd = (uint32_t)ds4_engine_embd_dim(engine),
.n_vocab = (uint32_t)ds4_engine_vocab_size(engine),
.quant_bits = (uint32_t)ds4_engine_routed_quant_bits(engine),
.ctx_size = (uint32_t)cfg.ctx_alloc,
};
ds4_engine_tp_gate_schedule(engine,
&tp_id.gate_slot_start,
&tp_id.gate_slot_step,
&tp_id.gates_per_token,
tp_id.gate_slot_mask);
if (!ds4_tp_create(&tp_leader,
&cfg.tp,
&tp_id,
tp_err,
sizeof(tp_err)) ||
!ds4_engine_tp_bind(engine, tp_leader, tp_err, sizeof(tp_err))) {
fprintf(stderr, "ds4-bench: %s\n", tp_err);
close_engine(engine, tp_leader);
return 1;
}
}
log_context_memory(opt.backend,
cfg.ctx_alloc,
cfg.prefill_chunk,
ds4_engine_prefill_chunk(engine),
cfg.ssd_streaming);
char *text = read_file(cfg.prompt_path ? cfg.prompt_path : cfg.chat_prompt_path);
@@ -582,13 +738,16 @@ int main(int argc, char **argv) {
}
free(text);
if (prompt.len < cfg.ctx_max) {
const int prompt_tokens_needed = cfg.ctx_max +
(cfg.teacher_forced_decode ? cfg.gen_tokens : 0);
if (prompt.len < prompt_tokens_needed) {
fprintf(stderr,
"ds4-bench: prompt has %d tokens, need at least --ctx-max=%d\n",
"ds4-bench: prompt has %d tokens, need at least %d%s\n",
prompt.len,
cfg.ctx_max);
prompt_tokens_needed,
cfg.teacher_forced_decode ? " for teacher-forced decode" : "");
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
close_engine(engine, tp_leader);
return 1;
}
@@ -596,7 +755,7 @@ int main(int argc, char **argv) {
if (ds4_session_create(&session, engine, cfg.ctx_alloc) != 0) {
fprintf(stderr, "ds4-bench: failed to create session\n");
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
close_engine(engine, tp_leader);
return 1;
}
if (cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR &&
@@ -604,7 +763,7 @@ int main(int argc, char **argv) {
{
ds4_session_free(session);
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
close_engine(engine, tp_leader);
return 1;
}
maybe_warn_distributed_step_shape(&cfg, session);
@@ -616,7 +775,7 @@ int main(int argc, char **argv) {
fprintf(stderr, "ds4-bench: failed to open %s: %s\n", cfg.csv_path, strerror(errno));
ds4_session_free(session);
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
close_engine(engine, tp_leader);
return 1;
}
}
@@ -624,7 +783,23 @@ int main(int argc, char **argv) {
fflush(out);
const int eos = ds4_token_eos(engine);
const bool distributed = cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR;
const bool distributed =
cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR ||
cfg.tp.role == DS4_TP_LEADER;
const bool speculative = cfg.dspark && ds4_engine_mtp_draft_tokens(engine) > 1;
if (cfg.dspark && !speculative) {
fprintf(stderr, "ds4-bench: DSpark support model did not enable speculative decoding\n");
if (out != stdout) fclose(out);
ds4_session_free(session);
ds4_tokens_free(&prompt);
close_engine(engine, tp_leader);
return 1;
}
if (speculative) {
fprintf(stderr,
"ds4-bench: DSpark enabled with draft width %d; frontier restoration uses session snapshots\n",
ds4_engine_mtp_draft_tokens(engine));
}
ds4_session_snapshot snap = {0};
const uint64_t snapshot_max_bytes = bench_snapshot_max_bytes();
bool warned_large_snapshot = false;
@@ -640,13 +815,27 @@ int main(int argc, char **argv) {
};
const double prefill_t0 = bench_now_sec();
if (ds4_session_sync(session, &prefix, err, sizeof(err)) != 0) {
#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
const bool cuda_profile_prefill =
getenv("DS4_BENCH_CUDA_PROFILE_PREFILL") != NULL &&
cudaProfilerStart && cudaProfilerStop;
if (cuda_profile_prefill) (void)cudaProfilerStart();
#endif
const int prefill_rc =
ds4_session_sync(session, &prefix, err, sizeof(err));
#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
if (cuda_profile_prefill) (void)cudaProfilerStop();
#endif
if (prefill_rc != 0) {
fprintf(stderr, "ds4-bench: prefill to %d failed: %s\n", frontier, err);
rc = 1;
break;
}
const double prefill_t1 = bench_now_sec();
const double prefill_sec = prefill_t1 - prefill_t0;
if (getenv("DS4_METAL_CB_TIMES"))
fprintf(stderr, "ds4-bench: prefill window mono %.1f .. %.1f ms\n",
prefill_t0 * 1e3, prefill_t1 * 1e3);
const int prefill_tokens = frontier - previous;
if (write_frontier_logits_json(&cfg, engine, session, frontier, previous) != 0) {
@@ -684,31 +873,148 @@ int main(int argc, char **argv) {
const double gen_t0 = bench_now_sec();
double gen_first_sec = 0.0;
double gen_steady_sec = 0.0;
int gen_first_tokens = 0;
int gen_done = 0;
for (int i = 0; i < cfg.gen_tokens; i++) {
int *gen_token_buf = cfg.show_output && cfg.gen_tokens > 0
? malloc((size_t)cfg.gen_tokens * sizeof(gen_token_buf[0]))
: NULL;
int gen_token_count = 0;
int cuda_profile_start = -1;
int cuda_profile_tokens = 0;
const char *cuda_profile_range =
getenv("DS4_BENCH_CUDA_PROFILE_RANGE");
if (cuda_profile_range &&
(sscanf(cuda_profile_range, "%d:%d",
&cuda_profile_start, &cuda_profile_tokens) != 2 ||
cuda_profile_start < 0 || cuda_profile_tokens <= 0)) {
fprintf(stderr,
"ds4-bench: invalid DS4_BENCH_CUDA_PROFILE_RANGE=%s "
"(expected START_TOKEN:TOKEN_COUNT)\n",
cuda_profile_range);
cuda_profile_start = -1;
cuda_profile_tokens = 0;
}
bool generation_stop = false;
while (gen_done < cfg.gen_tokens && !generation_stop) {
if (ds4_session_pos(session) + 1 >= ds4_session_ctx(session)) {
fprintf(stderr, "ds4-bench: generation would exceed allocated context at frontier %d\n", frontier);
rc = 1;
break;
}
const int token = ds4_session_argmax_excluding(session, eos);
const int token = cfg.teacher_forced_decode
? prompt.v[frontier + gen_done]
: ds4_session_argmax_excluding(session, eos);
if (token < 0) {
fprintf(stderr, "ds4-bench: failed to choose non-EOS token at frontier %d\n", frontier);
rc = 1;
break;
}
const double token_t0 = bench_now_sec();
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
fprintf(stderr, "ds4-bench: decode at frontier %d failed: %s\n", frontier, err);
rc = 1;
break;
#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
if (gen_done == cuda_profile_start && cudaProfilerStart) {
(void)cudaProfilerStart();
}
#endif
int toks[17];
int ntok = 1;
if (speculative) {
const double spec_t0 = bench_now_sec();
ntok = ds4_session_eval_speculative_argmax(
session, token, cfg.gen_tokens - gen_done, eos,
toks, (int)(sizeof(toks) / sizeof(toks[0])),
err, sizeof(err));
if (getenv("DS4_BENCH_SPEC_TRACE") && ntok > 0 && ntok < 17) {
static double call_ms[17];
static unsigned call_n[17];
static unsigned calls_total;
call_ms[ntok] += (bench_now_sec() - spec_t0) * 1e3;
call_n[ntok]++;
if (++calls_total % 32u == 0u) {
fprintf(stderr, "ds4-bench: spec calls %u:", calls_total);
for (int k = 1; k < 17; k++)
if (call_n[k])
fprintf(stderr, " %d-tok x%u avg %.1f ms", k, call_n[k], call_ms[k] / call_n[k]);
fprintf(stderr, "\n");
}
}
if (ntok < 0) {
fprintf(stderr, "ds4-bench: DSpark decode at frontier %d failed: %s\n", frontier, err);
rc = 1;
break;
}
if (ntok == 0) {
fprintf(stderr, "ds4-bench: DSpark decode at frontier %d accepted no tokens\n", frontier);
rc = 1;
break;
}
} else {
toks[0] = token;
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
fprintf(stderr, "ds4-bench: decode at frontier %d failed: %s\n", frontier, err);
rc = 1;
break;
}
}
#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
if (cuda_profile_start >= 0 &&
gen_done + ntok >=
cuda_profile_start + cuda_profile_tokens &&
cudaProfilerStop) {
(void)cudaProfilerStop();
}
#endif
const double token_t1 = bench_now_sec();
if (i == 0) gen_first_sec = token_t1 - token_t0;
else gen_steady_sec += token_t1 - token_t0;
gen_done++;
/* Diagnostic: DS4_BENCH_DECODE_LOGITS_HASH prints an FNV-1a hash of
* every decode step's full logits, to compare arms bit for bit. */
if (!speculative && getenv("DS4_BENCH_DECODE_LOGITS_HASH")) {
const int vocab = ds4_engine_vocab_size(engine);
float *lg = vocab > 0 ? malloc((size_t)vocab * sizeof(lg[0])) : NULL;
if (lg && ds4_session_copy_logits(session, lg, vocab) == vocab) {
uint64_t h = 1469598103934665603ull;
const unsigned char *b = (const unsigned char *)lg;
for (size_t k = 0; k < (size_t)vocab * sizeof(lg[0]); k++) h = (h ^ b[k]) * 1099511628211ull;
fprintf(stderr, "ds4-bench: decode-hash frontier=%d step=%d %016llx\n",
frontier, gen_done, (unsigned long long)h);
const char *dump = getenv("DS4_BENCH_DECODE_LOGITS_DUMP");
if (dump) {
FILE *fp = fopen(dump, "ab");
if (fp) { fwrite(lg, sizeof(lg[0]), (size_t)vocab, fp); fclose(fp); }
}
}
free(lg);
}
int cycle_tokens = 0;
for (int j = 0; j < ntok && gen_done < cfg.gen_tokens; j++) {
if (toks[j] == eos) {
generation_stop = true;
break;
}
if (gen_token_buf) gen_token_buf[gen_token_count++] = toks[j];
gen_done++;
cycle_tokens++;
}
if (gen_first_tokens == 0) {
gen_first_sec = token_t1 - token_t0;
gen_first_tokens = cycle_tokens;
} else {
gen_steady_sec += token_t1 - token_t0;
}
}
const double gen_t1 = bench_now_sec();
if (cfg.show_output && gen_token_buf && gen_token_count > 0) {
fprintf(stderr, "ds4-bench: gen[ctx=%d] decoded text: \"", frontier);
for (int i = 0; i < gen_token_count; i++) {
size_t tlen = 0;
char *txt = ds4_token_text(engine, gen_token_buf[i], &tlen);
if (txt) {
fwrite(txt, 1, tlen, stderr);
free(txt);
}
}
fprintf(stderr, "\"\n");
fflush(stderr);
}
free(gen_token_buf);
if (rc != 0) break;
if (!need_restore_after_generation) {
@@ -728,7 +1034,8 @@ int main(int argc, char **argv) {
}
const double gen_sec = gen_t1 - gen_t0;
const int gen_steady_tokens = gen_done > 1 ? gen_done - 1 : 0;
const int gen_steady_tokens = gen_done > gen_first_tokens ?
gen_done - gen_first_tokens : 0;
fprintf(out,
"%d,%d,%.2f,%d,%.2f,%.3f,%d,%.2f,%llu\n",
frontier,
@@ -750,6 +1057,6 @@ int main(int argc, char **argv) {
ds4_session_snapshot_free(&snap);
ds4_session_free(session);
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
close_engine(engine, tp_leader);
return rc;
}
ds4_cli.c
+494 -66
View File
@@ -1,6 +1,9 @@
#include "ds4.h"
#include "ds4_distributed.h"
#include "ds4_gpu_args.h"
#include "ds4_tp.h"
#include "ds4_help.h"
#include "ds4_prompt_prefix.h"
#include "linenoise.h"
/* ds4 CLI.
@@ -25,9 +28,41 @@
#include <time.h>
#include <unistd.h>
static bool cli_env_flag_enabled(const char *name, bool defval) {
const char *v = getenv(name);
if (!v || !v[0]) return defval;
return strcmp(v, "0") != 0;
}
static bool cli_splitkv_spec_requested(void) {
if (cli_env_flag_enabled("DS4_CUDA_NO_SPLITKV_SPEC", false)) return false;
return cli_env_flag_enabled("DS4_CUDA_SPLITKV_SPEC", false);
}
static bool cli_greedy_fast_attention_requested(void) {
if (!cli_env_flag_enabled("DS4_CUDA_NO_GREEDY_SPLITKV", false) &&
cli_env_flag_enabled("DS4_CUDA_GREEDY_SPLITKV", false))
{
return true;
}
if (!cli_env_flag_enabled("DS4_CUDA_NO_GREEDY_VEC4", false) &&
cli_env_flag_enabled("DS4_CUDA_GREEDY_VEC4", false))
{
return true;
}
return false;
}
static bool cli_greedy_argmax_requested(bool speculative_requested) {
if (cli_greedy_fast_attention_requested()) return true;
if (speculative_requested) return false;
return cli_env_flag_enabled("DS4_CUDA_GREEDY_TOP1", true);
}
typedef struct {
const char *prompt;
const char *system;
ds4_prompt_prefix prefix;
bool raw_prompt;
int n_predict;
int ctx_size;
@@ -48,6 +83,7 @@ typedef struct {
const char *imatrix_output_path;
int imatrix_max_prompts;
int imatrix_max_tokens;
int imatrix_min_expert_samples;
ds4_think_mode think_mode;
bool head_test;
bool first_token_test;
@@ -62,6 +98,10 @@ typedef struct {
cli_generation_options gen;
char *prompt_owned;
bool inspect;
/* CLI flag wiring: raw argv values for --gpu-vram and --gpu-devices.
* Resolved post-parse via parse_gpu_vram_arg(). */
const char *gpu_vram_arg;
const char *gpu_devices_arg;
} cli_config;
static volatile sig_atomic_t cli_interrupted;
@@ -460,14 +500,29 @@ static void print_generated_token(void *ud, int token) {
free(text);
}
static void build_chat_prompt(ds4_engine *engine,
const cli_generation_options *gen,
ds4_tokens *out) {
const ds4_think_mode think_mode = cli_effective_think_mode(gen);
ds4_chat_begin(engine, out);
ds4_chat_append_think_prefix(engine, out, think_mode);
if (gen->system && gen->system[0])
ds4_chat_append_message(engine, out, "system", gen->system);
ds4_prompt_prefix_append(engine, out, &gen->prefix);
ds4_chat_append_message(engine, out, "user", gen->prompt ? gen->prompt : "");
ds4_chat_append_assistant_prefix(engine, out, think_mode);
}
static void build_prompt(ds4_engine *engine, const cli_generation_options *gen, ds4_tokens *out) {
if (gen->raw_prompt) {
ds4_tokenize_text(engine, gen->prompt ? gen->prompt : "", out);
} else if (is_rendered_chat_prompt(gen->prompt)) {
ds4_tokenize_rendered_chat(engine, gen->prompt, out);
} else {
} else if (gen->prefix.count == 0) {
ds4_encode_chat_prompt(engine, gen->system, gen->prompt,
cli_effective_think_mode(gen), out);
} else {
build_chat_prompt(engine, gen, out);
}
}
@@ -491,6 +546,9 @@ static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, con
ds4_session_free(session);
return 1;
}
/* Pay the one-time first-submission GPU cost before the prefill timer
* starts (matches the TP worker's startup warmup). */
ds4_session_gpu_warmup(session);
char err[160];
ds4_think_mode think_mode = cli_effective_think_mode(&cfg->gen);
@@ -535,25 +593,37 @@ static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, con
uint64_t rng = cfg->gen.seed ? cfg->gen.seed :
((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());
int generated = 0;
const bool speculative_argmax = cfg->gen.temperature <= 0.0f &&
((ds4_engine_mtp_draft_tokens(engine) > 1 &&
getenv("DS4_MTP_SPEC_DISABLE") == NULL) ||
cli_splitkv_spec_requested());
const bool greedy_argmax = cfg->gen.temperature <= 0.0f &&
cli_greedy_argmax_requested(speculative_argmax);
bool have_greedy_next = false;
int greedy_next = -1;
const double t_decode0 = cli_now_sec();
while (generated < max_tokens && !cli_interrupt_requested()) {
int token = ds4_session_sample(session, cfg->gen.temperature, 0,
int token;
if (greedy_argmax && have_greedy_next) {
token = greedy_next;
have_greedy_next = false;
} else {
token = ds4_session_sample(session, cfg->gen.temperature, 0,
cfg->gen.top_p, cfg->gen.min_p, &rng);
}
if (ds4_token_is_stop_for_think_mode(engine, token, think_mode)) break;
int toks[17];
int ntok = 0;
if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&
if (ds4_engine_mtp_draft_tokens(engine) > 1 &&
getenv("DS4_MTP_SPEC_DISABLE") == NULL) {
cli_dist_busy_set(cfg, true);
ntok = ds4_session_eval_speculative_argmax(session,
token,
max_tokens - generated,
ds4_token_eos(engine),
toks,
(int)(sizeof(toks) / sizeof(toks[0])),
err,
sizeof(err));
ntok = ds4_session_eval_speculative(
session, token, max_tokens - generated,
ds4_token_eos(engine), cfg->gen.temperature, 0,
cfg->gen.top_p, cfg->gen.min_p, &rng,
toks, (int)(sizeof(toks) / sizeof(toks[0])),
err, sizeof(err));
cli_dist_busy_set(cfg, false);
if (ntok < 0) {
fprintf(stderr, "ds4: decode failed: %s\n", err);
@@ -1168,8 +1238,14 @@ static int run_generation(ds4_engine *engine, const cli_config *cfg) {
}
} else {
if (cfg->engine.distributed.role == DS4_DISTRIBUTED_COORDINATOR ||
cfg->engine.tp.role == DS4_TP_LEADER ||
getenv("DS4_CLI_FORCE_SESSION") != NULL ||
cfg->gen.temperature > 0.0f ||
ds4_engine_mtp_draft_tokens(engine) > 1) {
/* TP leaders always drive the session path: the sync/eval
* mirroring that keeps the worker in lockstep lives there.
* The env override exists so TP-vs-single-node validation
* compares like with like. */
rc = run_sampled_generation(engine, cfg, &prompt);
} else {
token_printer printer = {
@@ -1210,12 +1286,13 @@ static char *trim_inplace(char *s) {
static void print_repl_help(void) {
puts("Commands:");
puts(" /help Show this help.");
puts(" /think Use normal thinking mode.");
puts(" /think-max Use Think Max only when context is at least 393216 tokens.");
puts(" /think [N] Use normal thinking, or V4.1 effort 0..100 (0 disables thinking).");
puts(" /think-max Use maximum thinking (V4.1: 100; V4: requires ctx >= 393216).");
puts(" /nothink Disable thinking mode.");
puts(" /ctx N Set context size for following prompts.");
puts(" /power N Set GPU duty cycle percentage, 1..100.");
puts(" /read FILE Read a prompt from FILE and run it.");
puts(" /steer F Set FFN steering for subsequent tokens; no value shows it.");
puts(" /read FILE Submit a text file, PNG, or JPEG.");
puts(" /quit, /exit Leave the prompt.");
puts(" Ctrl+C Stop generation and return to the prompt.");
}
@@ -1228,6 +1305,18 @@ static bool parse_power_percent(const char *arg, int *out) {
return true;
}
static bool parse_steering_level(const char *arg, float *out) {
char *end = NULL;
errno = 0;
float v = strtof(arg, &end);
if (!arg[0] || *end != '\0' || errno == ERANGE || !isfinite(v) ||
v < -100.0f || v > 100.0f) {
return false;
}
*out = v;
return true;
}
static void history_file_path(char *buf, size_t len) {
const char *home = getenv("HOME");
if (!home || !home[0]) home = ".";
@@ -1237,11 +1326,38 @@ static void history_file_path(char *buf, size_t len) {
typedef struct {
ds4_session *session;
ds4_tokens transcript;
ds4_vision_span *images;
size_t image_count;
size_t image_cap;
int ctx_size;
int think_prefix_pos;
int think_prefix_tokens;
bool initial_system_text;
} repl_chat;
static void repl_chat_free(repl_chat *chat);
static void repl_chat_trim_images(repl_chat *chat, size_t count) {
while (chat->image_count > count) {
chat->image_count--;
ds4_vision_embedding_free(&chat->images[chat->image_count].embedding);
}
}
static ds4_vision_span *repl_chat_add_image(repl_chat *chat) {
if (chat->image_count == chat->image_cap) {
size_t cap = chat->image_cap ? chat->image_cap * 2u : 4u;
if (cap > SIZE_MAX / sizeof(chat->images[0])) return NULL;
void *next = realloc(chat->images, cap * sizeof(chat->images[0]));
if (!next) return NULL;
chat->images = next;
chat->image_cap = cap;
}
ds4_vision_span *span = &chat->images[chat->image_count++];
memset(span, 0, sizeof(*span));
return span;
}
static void tokens_insert(ds4_tokens *dst, int pos, const ds4_tokens *src) {
if (!src || src->len <= 0) return;
if (pos < 0) pos = 0;
@@ -1269,34 +1385,22 @@ static void tokens_remove(ds4_tokens *dst, int pos, int n) {
dst->len -= n;
}
static const char *repl_glm_reasoning_effort_text(ds4_think_mode mode) {
switch (mode) {
case DS4_THINK_HIGH: return "Reasoning Effort: High";
case DS4_THINK_MAX: return "Reasoning Effort: Max";
case DS4_THINK_NONE: return NULL;
}
return NULL;
}
static void repl_chat_build_think_prefix(ds4_engine *engine,
ds4_think_mode mode,
ds4_tokens *prefix) {
if (ds4_engine_is_glm_dsa(engine)) {
const char *effort = repl_glm_reasoning_effort_text(mode);
if (effort) ds4_chat_append_message(engine, prefix, "system", effort);
} else if (mode == DS4_THINK_MAX) {
ds4_chat_append_max_effort_prefix(engine, prefix);
}
ds4_chat_append_think_prefix(engine, prefix, mode);
}
/* Insert/replace the model-family thinking prefix inside the existing
* transcript. It lives immediately after the BOS sequence and before any
* user/system text, matching the GGUF chat templates. */
static void repl_chat_apply_think_prefix(ds4_engine *engine,
static bool repl_chat_apply_think_prefix(ds4_engine *engine,
repl_chat *chat,
ds4_think_mode mode) {
ds4_tokens prefix = {0};
repl_chat_build_think_prefix(engine, mode, &prefix);
if (ds4_engine_is_deepseek41(engine) && chat->initial_system_text && !prefix.len)
ds4_chat_append_message(engine, &prefix, "system", "");
bool same = chat->think_prefix_tokens == prefix.len;
if (same && prefix.len > 0) {
@@ -1305,13 +1409,25 @@ static void repl_chat_apply_think_prefix(ds4_engine *engine,
(size_t)prefix.len * sizeof(prefix.v[0]));
}
if (!same) {
const int old_prefix_tokens = chat->think_prefix_tokens;
const int shift = prefix.len - old_prefix_tokens;
if (chat->ctx_size > 0 && (int64_t)chat->transcript.len + shift >= chat->ctx_size) {
fprintf(stderr, "ds4: no context room to change the thinking prefix\n");
ds4_tokens_free(&prefix);
return false;
}
tokens_remove(&chat->transcript, chat->think_prefix_pos,
chat->think_prefix_tokens);
tokens_insert(&chat->transcript, chat->think_prefix_pos, &prefix);
chat->think_prefix_tokens = prefix.len;
for (size_t i = 0; i < chat->image_count; i++) {
chat->images[i].token_start =
(uint32_t)((int64_t)chat->images[i].token_start + shift);
}
if (chat->session) ds4_session_invalidate(chat->session);
}
ds4_tokens_free(&prefix);
return true;
}
static int repl_chat_create_session(ds4_engine *engine, repl_chat *chat, int ctx_size) {
@@ -1328,19 +1444,55 @@ static int repl_chat_create_session(ds4_engine *engine, repl_chat *chat, int ctx
static int repl_chat_init(ds4_engine *engine, repl_chat *chat, const cli_config *cfg) {
memset(chat, 0, sizeof(*chat));
chat->initial_system_text = cfg->gen.system && cfg->gen.system[0];
ds4_chat_begin(engine, &chat->transcript);
chat->think_prefix_pos = chat->transcript.len;
repl_chat_apply_think_prefix(engine, chat, cli_effective_think_mode(&cfg->gen));
if (cfg->gen.system && cfg->gen.system[0]) {
ds4_chat_append_message(engine, &chat->transcript, "system", cfg->gen.system);
}
return repl_chat_create_session(engine, chat, cfg->gen.ctx_size);
ds4_prompt_prefix_append(engine, &chat->transcript, &cfg->gen.prefix);
if (repl_chat_create_session(engine, chat, cfg->gen.ctx_size) != 0) {
ds4_tokens_free(&chat->transcript);
return 1;
}
if (cfg->gen.prefix.count == 0) return 0;
if (cli_wait_distributed_route(cfg, chat->session) != 0) {
repl_chat_free(chat);
return 1;
}
char err[160] = {0};
cli_prefill_progress progress = {
.base_tokens = 0,
.input_tokens = chat->transcript.len,
.use_color = ds4_log_is_tty(stderr),
};
ds4_session_set_progress(chat->session, cli_prefill_progress_cb, &progress);
ds4_session_set_display_progress(chat->session,
progress.use_color ? cli_prefill_progress_cb : NULL,
progress.use_color ? &progress : NULL);
cli_dist_busy_set(cfg, true);
int rc = ds4_session_sync(chat->session, &chat->transcript,
err, sizeof(err));
cli_dist_busy_set(cfg, false);
ds4_session_set_progress(chat->session, NULL, NULL);
ds4_session_set_display_progress(chat->session, NULL, NULL);
if (rc != 0) {
fprintf(stderr, "ds4: prefix prefill failed: %s\n",
err[0] ? err : "unknown error");
repl_chat_free(chat);
return 1;
}
return 0;
}
static void repl_chat_free(repl_chat *chat) {
if (!chat) return;
ds4_session_free(chat->session);
ds4_tokens_free(&chat->transcript);
repl_chat_trim_images(chat, 0);
free(chat->images);
memset(chat, 0, sizeof(*chat));
}
@@ -1359,7 +1511,9 @@ static bool repl_chat_assistant_turn_uses_eos(ds4_engine *engine) {
* and assistant markers, then ds4_session_sync() decides whether this is a KV
* continuation. If prompt processing fails, the transcript rolls back before
* returning to the prompt. */
static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, const char *user_text) {
static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat,
const char *user_text,
ds4_vision_embedding *image) {
if (!chat->session) {
fprintf(stderr, "ds4: no active interactive KV cache\n");
return 1;
@@ -1367,9 +1521,25 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
ds4_think_mode think_mode = ds4_think_mode_for_context(cfg->gen.think_mode,
chat->ctx_size);
repl_chat_apply_think_prefix(engine, chat, think_mode);
if (!repl_chat_apply_think_prefix(engine, chat, think_mode)) return 1;
const int rollback_len = chat->transcript.len;
ds4_chat_append_message(engine, &chat->transcript, "user", user_text);
const size_t rollback_images = chat->image_count;
if (image) {
ds4_vision_span *span = repl_chat_add_image(chat);
const char *text_parts[] = {"", user_text ? user_text : ""};
char image_error[160] = {0};
if (!span || !ds4_chat_append_multimodal_message(
engine, &chat->transcript, "user", text_parts,
image, 1, span, image_error, sizeof(image_error))) {
repl_chat_trim_images(chat, rollback_images);
chat->transcript.len = rollback_len;
fprintf(stderr, "ds4: failed to add image: %s\n",
image_error[0] ? image_error : "out of memory");
return 1;
}
} else {
ds4_chat_append_message(engine, &chat->transcript, "user", user_text);
}
ds4_chat_append_assistant_prefix(engine, &chat->transcript, think_mode);
const int old_pos = ds4_session_pos(chat->session);
@@ -1389,12 +1559,18 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
progress.use_color ? cli_prefill_progress_cb : NULL,
progress.use_color ? &progress : NULL);
cli_dist_busy_set(cfg, true);
int sync_rc = ds4_session_sync(chat->session, &chat->transcript, err, sizeof(err));
int sync_rc = ds4_session_sync_multimodal(chat->session,
&chat->transcript,
chat->images,
chat->image_count,
err,
sizeof(err));
cli_dist_busy_set(cfg, false);
if (sync_rc != 0) {
ds4_session_set_progress(chat->session, NULL, NULL);
ds4_session_set_display_progress(chat->session, NULL, NULL);
chat->transcript.len = rollback_len;
repl_chat_trim_images(chat, rollback_images);
fprintf(stderr, "ds4: prompt processing failed: %s\n", err);
return 1;
}
@@ -1419,29 +1595,41 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
uint64_t rng = cfg->gen.seed ? cfg->gen.seed :
((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());
int generated = 0;
const bool speculative_argmax = cfg->gen.temperature <= 0.0f &&
((ds4_engine_mtp_draft_tokens(engine) > 1 &&
getenv("DS4_MTP_SPEC_DISABLE") == NULL) ||
cli_splitkv_spec_requested());
const bool greedy_argmax = cfg->gen.temperature <= 0.0f &&
cli_greedy_argmax_requested(speculative_argmax);
bool have_greedy_next = false;
int greedy_next = -1;
const double t_decode0 = cli_now_sec();
while (generated < max_tokens && !cli_interrupt_requested()) {
int token = ds4_session_sample(chat->session,
int token;
if (greedy_argmax && have_greedy_next) {
token = greedy_next;
have_greedy_next = false;
} else {
token = ds4_session_sample(chat->session,
cfg->gen.temperature,
0,
cfg->gen.top_p,
cfg->gen.min_p,
&rng);
}
if (ds4_token_is_stop_for_think_mode(engine, token, think_mode)) break;
int toks[17];
int ntok = 0;
if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&
if (ds4_engine_mtp_draft_tokens(engine) > 1 &&
getenv("DS4_MTP_SPEC_DISABLE") == NULL) {
cli_dist_busy_set(cfg, true);
ntok = ds4_session_eval_speculative_argmax(chat->session,
token,
max_tokens - generated,
ds4_token_eos(engine),
toks,
(int)(sizeof(toks) / sizeof(toks[0])),
err,
sizeof(err));
ntok = ds4_session_eval_speculative(
chat->session, token, max_tokens - generated,
ds4_token_eos(engine), cfg->gen.temperature, 0,
cfg->gen.top_p, cfg->gen.min_p, &rng,
toks, (int)(sizeof(toks) / sizeof(toks[0])),
err, sizeof(err));
cli_dist_busy_set(cfg, false);
if (ntok < 0) {
fprintf(stderr, "ds4: decode failed: %s\n", err);
@@ -1491,6 +1679,7 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
const bool interrupted = cli_interrupt_requested();
if (interrupted && generated == 0) {
chat->transcript.len = rollback_len;
repl_chat_trim_images(chat, rollback_images);
ds4_session_invalidate(chat->session);
} else if (repl_chat_assistant_turn_uses_eos(engine)) {
ds4_tokens_push(&chat->transcript, ds4_token_eos(engine));
@@ -1507,6 +1696,16 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
return 0;
}
static bool cli_file_has_image_magic(const char *path) {
unsigned char magic[8] = {0};
FILE *fp = fopen(path, "rb");
if (!fp) return false;
size_t n = fread(magic, 1, sizeof(magic), fp);
fclose(fp);
return (n >= 8 && !memcmp(magic, "\x89PNG\r\n\x1a\n", 8)) ||
(n >= 2 && magic[0] == 0xff && magic[1] == 0xd8);
}
static int run_repl(ds4_engine *engine, cli_config *cfg) {
repl_chat chat;
if (repl_chat_init(engine, &chat, cfg) != 0) return 1;
@@ -1547,22 +1746,31 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
if (!strcmp(cmd, "/help")) {
print_repl_help();
} else if (!strcmp(cmd, "/think")) {
cfg->gen.think_mode = DS4_THINK_HIGH;
repl_chat_apply_think_prefix(engine, &chat, DS4_THINK_HIGH);
puts("Thinking mode: high.");
} else if (!strncmp(cmd, "/think", 6) &&
(cmd[6] == '\0' || isspace((unsigned char)cmd[6]))) {
const char *arg = trim_inplace(cmd + 6);
ds4_think_mode mode = DS4_THINK_HIGH;
if (arg[0] && (!ds4_engine_is_deepseek41(engine) ||
!ds4_think_mode_parse_level(arg, &mode))) {
fprintf(stderr, "ds4: /think N requires V4.1 and a level from 0 to 100\n");
} else if (repl_chat_apply_think_prefix(engine, &chat, mode)) {
cfg->gen.think_mode = mode;
printf("Thinking mode: %s.\n", ds4_think_mode_name(mode));
}
} else if (!strcmp(cmd, "/think-max")) {
cfg->gen.think_mode = DS4_THINK_MAX;
bool active = ds4_think_mode_for_context(cfg->gen.think_mode,
bool active = ds4_think_mode_for_context(DS4_THINK_MAX,
chat.ctx_size) == DS4_THINK_MAX;
repl_chat_apply_think_prefix(engine, &chat,
active ? DS4_THINK_MAX : DS4_THINK_HIGH);
cli_warn_think_max_downgraded(&cfg->gen, "/think-max");
printf("Thinking mode: %s.\n", active ? "max" : "high (ctx below 393216)");
if (repl_chat_apply_think_prefix(engine, &chat,
active ? DS4_THINK_MAX : DS4_THINK_HIGH)) {
cfg->gen.think_mode = DS4_THINK_MAX;
cli_warn_think_max_downgraded(&cfg->gen, "/think-max");
printf("Thinking mode: %s.\n", active ? "max" : "high (ctx below 393216)");
}
} else if (!strcmp(cmd, "/nothink")) {
cfg->gen.think_mode = DS4_THINK_NONE;
repl_chat_apply_think_prefix(engine, &chat, DS4_THINK_NONE);
puts("Thinking mode: none.");
if (repl_chat_apply_think_prefix(engine, &chat, DS4_THINK_NONE)) {
cfg->gen.think_mode = DS4_THINK_NONE;
puts("Thinking mode: none.");
}
} else if (!strncmp(cmd, "/power", 6) && (cmd[6] == '\0' || isspace((unsigned char)cmd[6]))) {
char *arg = trim_inplace(cmd + 6);
if (!arg[0]) {
@@ -1578,6 +1786,22 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
printf("Power: %d%%.\n", power);
}
}
} else if (!strncmp(cmd, "/steer", 6) &&
(cmd[6] == '\0' || isspace((unsigned char)cmd[6]))) {
char *arg = trim_inplace(cmd + 6);
if (!arg[0]) {
printf("Steering FFN: %g.\n",
(double)ds4_session_directional_steering_ffn(chat.session));
} else {
float scale = 0.0f;
if (!parse_steering_level(arg, &scale)) {
fprintf(stderr, "ds4: /steer must be between -100 and 100\n");
} else if (ds4_session_set_directional_steering_ffn(
chat.session, scale) == 0) {
cfg->engine.directional_steering_ffn = scale;
printf("Steering FFN: %g.\n", (double)scale);
}
}
} else if (!strncmp(cmd, "/ctx", 4) && (cmd[4] == '\0' || isspace((unsigned char)cmd[4]))) {
char *arg = trim_inplace(cmd + 4);
if (!arg[0]) {
@@ -1586,7 +1810,7 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
cfg->gen.ctx_size = parse_int(arg, "/ctx");
log_context_memory(cfg->engine.backend,
cfg->gen.ctx_size,
cfg->engine.prefill_chunk,
ds4_engine_prefill_chunk(engine),
cfg->engine.ssd_streaming);
rc = repl_chat_set_ctx(engine, &chat, cfg->gen.ctx_size);
if (rc != 0) {
@@ -1605,10 +1829,24 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
char *path = trim_inplace(cmd + 5);
if (!path[0]) {
fprintf(stderr, "ds4: /read needs a file path\n");
} else if (cli_file_has_image_magic(path)) {
char image_error[256] = {0};
ds4_vision_embedding image = {0};
if (!ds4_engine_vision_encode_file(engine, path, &image,
image_error,
sizeof(image_error))) {
fprintf(stderr, "ds4: /read image failed: %s\n", image_error);
} else {
fprintf(stderr,
"ds4: image %ux%u, %u image tokens\n",
image.width, image.height, image.token_count);
rc = run_chat_turn(engine, cfg, &chat, "", &image);
ds4_vision_embedding_free(&image);
}
} else {
char *prompt = read_prompt_file(path, false);
if (prompt) {
rc = run_chat_turn(engine, cfg, &chat, prompt);
rc = run_chat_turn(engine, cfg, &chat, prompt, NULL);
free(prompt);
}
}
@@ -1616,7 +1854,7 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
fprintf(stderr, "ds4: unknown command: %s\n", cmd);
fprintf(stderr, "ds4: type /help for commands\n");
} else {
rc = run_chat_turn(engine, cfg, &chat, cmd);
rc = run_chat_turn(engine, cfg, &chat, cmd, NULL);
}
linenoiseFree(line);
}
@@ -1730,6 +1968,20 @@ static cli_config parse_options(int argc, char **argv) {
}
if (dist_parse == DS4_DIST_CLI_MATCHED) continue;
char tp_parse_err[256] = {0};
ds4_tp_cli_parse_result tp_parse = ds4_tp_parse_cli_arg(arg,
&i,
argc,
argv,
&c.engine.tp,
tp_parse_err,
sizeof(tp_parse_err));
if (tp_parse == DS4_TP_CLI_ERROR) {
fprintf(stderr, "ds4: %s\n", tp_parse_err[0] ? tp_parse_err : "invalid tensor-parallel option");
exit(2);
}
if (tp_parse == DS4_TP_CLI_MATCHED) continue;
if (!strcmp(arg, "-p") || !strcmp(arg, "--prompt")) {
if (c.gen.prompt) {
fprintf(stderr, "ds4: specify only one prompt source\n");
@@ -1743,18 +1995,50 @@ static cli_config parse_options(int argc, char **argv) {
}
c.prompt_owned = read_prompt_file(need_arg(&i, argc, argv, arg), true);
c.gen.prompt = c.prompt_owned;
} else if (!strcmp(arg, "--prefix-file")) {
if (c.gen.prefix.count != 0) {
fprintf(stderr, "ds4: specify --prefix-file only once\n");
exit(2);
}
const char *path = need_arg(&i, argc, argv, arg);
char error[256] = {0};
if (ds4_prompt_prefix_load(&c.gen.prefix, path,
error, sizeof(error)) != 0) {
fprintf(stderr, "ds4: %s\n",
error[0] ? error : "invalid prefix file");
exit(2);
}
} else if (!strcmp(arg, "-sys") || !strcmp(arg, "--system")) {
c.gen.system = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--raw") || !strcmp(arg, "--raw-prompt")) {
c.gen.raw_prompt = true;
} else if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
c.engine.model_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--vision")) {
c.engine.vision_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--mtp")) {
c.engine.glm_mtp = true;
} else if (!strcmp(arg, "--mtp-model")) {
c.engine.mtp_path = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--mtp-draft")) {
c.engine.mtp_draft_tokens = parse_int(need_arg(&i, argc, argv, arg), arg);
} else if (!strcmp(arg, "--mtp-margin")) {
c.engine.mtp_margin = parse_float_range(need_arg(&i, argc, argv, arg), arg, 0.0f, 1000.0f);
} else if (!strcmp(arg, "--mtp-timing")) {
c.engine.glm_mtp = true;
c.engine.glm_mtp_timing = true;
} else if (!strcmp(arg, "--dspark")) {
c.engine.dspark = true;
} else if (!strcmp(arg, "--dspark-confidence")) {
c.engine.dspark = true;
c.engine.dspark_confidence_threshold =
parse_float_range(need_arg(&i, argc, argv, arg), arg, 0.0f, 1.0f);
c.engine.dspark_confidence_threshold_set = true;
} else if (!strcmp(arg, "--dspark-strict")) {
c.engine.dspark = true;
c.engine.dspark_strict = true;
} else if (!strcmp(arg, "--mtp-exact-sampling")) {
c.engine.dspark_exact_sampling = true;
} else if (!strcmp(arg, "-n") || !strcmp(arg, "--tokens")) {
c.gen.n_predict = parse_int(need_arg(&i, argc, argv, arg), arg);
} else if (!strcmp(arg, "-c") || !strcmp(arg, "--ctx")) {
@@ -1843,6 +2127,12 @@ static cli_config parse_options(int argc, char **argv) {
} else if (!strcmp(arg, "--cuda")) {
c.engine.backend = DS4_BACKEND_CUDA;
#endif
} else if (!strcmp(arg, "--gpu-vram")) {
c.gpu_vram_arg = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--gpu-devices")) {
c.gpu_devices_arg = need_arg(&i, argc, argv, arg);
} else if (!strcmp(arg, "--cuda-tensor-parallel")) {
c.engine.cuda_tensor_parallel = true;
} else if (!strcmp(arg, "--dump-tokens")) {
c.gen.dump_tokens = true;
} else if (!strcmp(arg, "--dump-logits")) {
@@ -1864,6 +2154,14 @@ static cli_config parse_options(int argc, char **argv) {
c.gen.imatrix_max_prompts = parse_int(need_arg(&i, argc, argv, arg), arg);
} else if (!strcmp(arg, "--imatrix-max-tokens")) {
c.gen.imatrix_max_tokens = parse_int(need_arg(&i, argc, argv, arg), arg);
} else if (!strcmp(arg, "--imatrix-min-expert-samples")) {
c.gen.imatrix_min_expert_samples =
parse_int(need_arg(&i, argc, argv, arg), arg);
} else if (!strcmp(arg, "--think-level")) {
if (!ds4_think_mode_parse_level(need_arg(&i, argc, argv, arg), &c.gen.think_mode)) {
fprintf(stderr, "ds4: --think-level requires an integer from 0 to 100\n");
exit(2);
}
} else if (!strcmp(arg, "--think")) {
c.gen.think_mode = DS4_THINK_HIGH;
} else if (!strcmp(arg, "--think-max")) {
@@ -1923,15 +2221,38 @@ static cli_config parse_options(int argc, char **argv) {
fprintf(stderr, "ds4: --imatrix-dataset requires --imatrix-out\n");
exit(2);
}
if (c.gen.imatrix_min_expert_samples < 0) {
fprintf(stderr, "ds4: --imatrix-min-expert-samples must not be negative\n");
exit(2);
}
if (c.gen.perplexity_file_path && c.gen.prompt) {
fprintf(stderr, "ds4: --perplexity-file does not use -p/--prompt-file\n");
exit(2);
}
if (c.gen.prefix.count != 0 && c.gen.raw_prompt) {
fprintf(stderr, "ds4: --prefix-file cannot be combined with --raw-prompt\n");
exit(2);
}
if (c.gen.prefix.count != 0 && is_rendered_chat_prompt(c.gen.prompt)) {
fprintf(stderr,
"ds4: --prefix-file cannot be combined with an already-rendered prompt\n");
exit(2);
}
char tp_err[256];
if (!ds4_tp_adopt_distributed_options(&c.engine.tp, c.dist,
tp_err, sizeof(tp_err))) {
fprintf(stderr, "ds4: %s\n", tp_err);
exit(2);
}
char dist_err[256];
if (ds4_dist_prepare_engine_options(c.dist, &c.engine, dist_err, sizeof(dist_err)) != 0) {
fprintf(stderr, "ds4: %s\n", dist_err);
exit(2);
}
if (!ds4_tp_validate_engine_options(&c.engine, tp_err, sizeof(tp_err))) {
fprintf(stderr, "ds4: %s\n", tp_err);
exit(2);
}
return c;
}
@@ -1939,15 +2260,35 @@ static cli_config parse_options(int argc, char **argv) {
int main(int argc, char **argv) {
cli_config cfg = parse_options(argc, argv);
if (cfg.gen.dump_tokens) {
if (cfg.gen.prompt == NULL) {
fprintf(stderr, "ds4: --dump-tokens requires -p or --prompt-file\n");
if (cfg.gen.prefix.count != 0) {
fprintf(stderr, "ds4: --dump-tokens does not support --prefix-file\n");
ds4_prompt_prefix_free(&cfg.gen.prefix);
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 2;
}
int rc = ds4_dump_text_tokenization(cfg.engine.model_path,
if (cfg.gen.prompt == NULL) {
fprintf(stderr, "ds4: --dump-tokens requires -p or --prompt-file\n");
ds4_prompt_prefix_free(&cfg.gen.prefix);
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 2;
}
int rc;
if (cfg.gen.raw_prompt || is_rendered_chat_prompt(cfg.gen.prompt)) {
rc = ds4_dump_text_tokenization(cfg.engine.model_path,
cfg.gen.prompt,
stdout);
} else {
rc = ds4_dump_chat_tokenization(cfg.engine.model_path,
cfg.gen.system,
cfg.gen.prompt,
cfg.gen.think_mode,
cfg.gen.ctx_size,
stdout);
}
ds4_dist_options_free(cfg.dist);
ds4_prompt_prefix_free(&cfg.gen.prefix);
free(cfg.prompt_owned);
return rc;
}
@@ -1955,13 +2296,92 @@ int main(int argc, char **argv) {
cfg.engine.first_token_test = cfg.gen.first_token_test;
cfg.engine.metal_graph_test = cfg.gen.metal_graph_test;
cfg.engine.context_size = cfg.gen.ctx_size;
cfg.engine.placement_ctx_hint = cfg.gen.ctx_size;
ds4_engine *engine = NULL;
if (ds4_engine_open(&engine, &cfg.engine) != 0) {
if (cfg.gpu_vram_arg || cfg.gpu_devices_arg) {
ds4_gpu_config gpu_cfg = {0};
bool skip_cuda = false;
char errbuf[256];
if (parse_gpu_vram_arg(cfg.gpu_vram_arg, cfg.gpu_devices_arg,
&gpu_cfg, &skip_cuda,
errbuf, sizeof(errbuf)) != 0) {
fprintf(stderr, "ds4: %s\n", errbuf);
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 2;
}
cfg.engine.backend = skip_cuda ? DS4_BACKEND_CPU : DS4_BACKEND_CUDA;
if (skip_cuda) {
if (ds4_engine_open(&engine, &cfg.engine) != 0) {
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 1;
}
} else {
const bool was_auto =
(cfg.gpu_vram_arg && !strcmp(cfg.gpu_vram_arg, "auto")) ||
(!cfg.gpu_vram_arg && cfg.gpu_devices_arg);
char layout[256];
if (format_gpu_layout_line(&gpu_cfg, was_auto,
layout, sizeof(layout)) > 0) {
fprintf(stdout, "%s\n", layout);
fflush(stdout);
}
if (ds4_engine_create_with_gpu_config(&engine, &cfg.engine,
&gpu_cfg) != 0) {
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 1;
}
}
} else if (ds4_engine_open(&engine, &cfg.engine) != 0) {
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 1;
}
if (ds4_think_mode_level(cfg.gen.think_mode) >= 0 && !ds4_engine_is_deepseek41(engine)) {
fprintf(stderr, "ds4: --think-level requires a DeepSeek V4.1 model\n");
ds4_engine_close(engine);
ds4_dist_options_free(cfg.dist);
ds4_prompt_prefix_free(&cfg.gen.prefix);
free(cfg.prompt_owned);
return 2;
}
cli_apply_model_sampling_defaults(engine, &cfg.gen);
if (cfg.engine.tp.role == DS4_TP_WORKER) {
int rc = ds4_tp_worker_run(engine, &cfg.engine.tp);
ds4_engine_close(engine);
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return rc;
}
ds4_tp *tp_leader = NULL;
if (cfg.engine.tp.role == DS4_TP_LEADER) {
char tp_err[256] = "";
ds4_tp_identity tp_id = {
.gguf_bytes = ds4_engine_model_bytes(engine),
.model_id = (uint32_t)ds4_engine_model_id(engine),
.n_layer = (uint32_t)ds4_engine_layer_count(engine),
.n_embd = (uint32_t)ds4_engine_embd_dim(engine),
.n_vocab = (uint32_t)ds4_engine_vocab_size(engine),
.quant_bits = (uint32_t)ds4_engine_routed_quant_bits(engine),
.ctx_size = (uint32_t)cfg.gen.ctx_size,
};
ds4_engine_tp_gate_schedule(engine,
&tp_id.gate_slot_start,
&tp_id.gate_slot_step,
&tp_id.gates_per_token,
tp_id.gate_slot_mask);
if (!ds4_tp_create(&tp_leader, &cfg.engine.tp, &tp_id, tp_err, sizeof(tp_err)) ||
!ds4_engine_tp_bind(engine, tp_leader, tp_err, sizeof(tp_err))) {
fprintf(stderr, "ds4: %s\n", tp_err);
ds4_tp_free(tp_leader);
ds4_engine_close(engine);
ds4_dist_options_free(cfg.dist);
free(cfg.prompt_owned);
return 1;
}
}
if (cfg.dist && cfg.dist->role == DS4_DISTRIBUTED_WORKER) {
ds4_dist_generation_options dist_gen = {
.prompt = cfg.gen.prompt,
@@ -1984,6 +2404,10 @@ int main(int argc, char **argv) {
return rc;
}
if (!cfg.inspect) {
log_context_memory(cfg.engine.backend,
cfg.gen.ctx_size,
ds4_engine_prefill_chunk(engine),
cfg.engine.ssd_streaming);
cli_warn_think_max_downgraded(&cfg.gen, "--think-max");
}
int rc = 0;
@@ -1995,7 +2419,8 @@ int main(int argc, char **argv) {
cfg.gen.imatrix_output_path,
cfg.gen.ctx_size,
cfg.gen.imatrix_max_prompts,
cfg.gen.imatrix_max_tokens);
cfg.gen.imatrix_max_tokens,
cfg.gen.imatrix_min_expert_samples);
} else if (cfg.gen.perplexity_file_path) {
rc = run_perplexity_file(engine, &cfg);
} else if (cfg.gen.prompt == NULL) {
@@ -2003,8 +2428,11 @@ int main(int argc, char **argv) {
} else {
rc = run_generation(engine, &cfg);
}
if (tp_leader) ds4_tp_send_stop(tp_leader);
ds4_engine_close(engine);
ds4_tp_free(tp_leader);
ds4_dist_options_free(cfg.dist);
ds4_prompt_prefix_free(&cfg.gen.prefix);
free(cfg.prompt_owned);
return rc;
}
ds4_cuda.cu
+26366 -3208
View File
File diff suppressed because it is too large. Load diff
ds4_deepseek41_cuda.cuh
+1017
View File
File diff suppressed because it is too large. Load diff
ds4_deepseek41_dspark.h
+471
View File
@@ -0,0 +1,471 @@
/* Internal V4.1 drafter graph. Included after the target graph in ds4.c. */
#ifndef DS4_DEEPSEEK41_DSPARK_H
#define DS4_DEEPSEEK41_DSPARK_H
#define DS41_DRAFT_WINDOW 128u
#define DS41_DRAFT_ROWS 5u
#define DS41_VERIFY_ROWS (DS41_DRAFT_ROWS + 1u)
struct ds41_dspark {
const ds4_dspark_weights *weights;
ds41_prefill_row batch;
ds4_gpu_tensor *target[3], *window[3];
uint32_t positions[3][DS41_DRAFT_WINDOW];
uint32_t kv_end;
bool kv_valid;
ds4_gpu_tensor *mean_weights, *mean, *input, *projected, *main_x, *main_kv;
ds4_gpu_tensor *logits, *tokens, *markov_embed, *markov_logits, *confidence_input, *confidence;
float *host_logits, *draft_probs;
ds4_gpu_tensor *verify_residual, *verify_pre, *verify_logits, *verify_pool;
ds4_gpu_tensor *old_raw, *old_features;
uint32_t old_positions[3][DS41_VERIFY_ROWS];
ds41_gpu_graph verified[DS41_VERIFY_ROWS + 1u];
float *host_verified;
uint32_t verify_start, verify_count, rewind_end;
};
static void ds41_dspark_reset(ds41_dspark *d) {
if (!d) return;
memset(d->positions, 0xff, sizeof(d->positions));
d->kv_valid = false;
d->kv_end = 0;
d->verify_count = 0;
d->rewind_end = 0;
}
static void ds41_dspark_free(ds41_dspark *d) {
if (!d) return;
#define DS41_DRAFT_FREE(name, count) ds4_gpu_tensor_free(d->batch.name);
DS41_PREFILL_ROWS(DS41_DRAFT_FREE)
#undef DS41_DRAFT_FREE
for (unsigned i = 0; i < 3; i++) {
ds4_gpu_tensor_free(d->target[i]);
ds4_gpu_tensor_free(d->window[i]);
}
ds4_gpu_tensor_free(d->mean_weights); ds4_gpu_tensor_free(d->mean);
ds4_gpu_tensor_free(d->input); ds4_gpu_tensor_free(d->projected);
ds4_gpu_tensor_free(d->main_x); ds4_gpu_tensor_free(d->main_kv);
ds4_gpu_tensor_free(d->logits); ds4_gpu_tensor_free(d->markov_embed);
ds4_gpu_tensor_free(d->tokens);
ds4_gpu_tensor_free(d->markov_logits); ds4_gpu_tensor_free(d->confidence_input);
ds4_gpu_tensor_free(d->confidence);
for (unsigned i = 0; i <= DS41_VERIFY_ROWS; i++) {
ds41_gpu_graph *g = &d->verified[i];
ds4_gpu_tensor_free(g->residual); ds4_gpu_tensor_free(g->pre);
ds4_gpu_tensor_free(g->logits);
ds4_gpu_tensor_free(g->tp_logits_half);
for (unsigned j = 0; j < 4; j++) {
ds4_gpu_tensor_free(g->verify_previous_kv[j]);
ds4_gpu_tensor_free(g->verify_previous_score[j]);
}
}
ds4_gpu_tensor_free(d->verify_residual); ds4_gpu_tensor_free(d->verify_pre);
ds4_gpu_tensor_free(d->verify_logits); ds4_gpu_tensor_free(d->verify_pool);
ds4_gpu_tensor_free(d->old_raw); ds4_gpu_tensor_free(d->old_features);
free(d->host_verified);
free(d->host_logits);
free(d->draft_probs);
free(d);
}
static DS4_MAYBE_UNUSED ds41_dspark *ds41_dspark_alloc(const ds4_dspark_weights *dw) {
if (!dw || !dw->v41 || dw->n_stages != 3 || dw->block_size != DS41_DRAFT_ROWS ||
dw->missing_tensors || dw->invalid_tensors || dw->metadata_errors) return NULL;
ds41_dspark *d = calloc(1, sizeof(*d));
if (!d) return NULL;
d->weights = dw;
const ds41_gpu_graph shape = {.ctx = DS41_DRAFT_WINDOW}, *g = &shape;
#define DS41_DRAFT_ALLOC(name, width) \
if (!(d->batch.name = ds4_gpu_tensor_alloc((uint64_t)(width) * DS41_DRAFT_ROWS * sizeof(float)))) goto fail;
DS41_PREFILL_ROWS(DS41_DRAFT_ALLOC)
#undef DS41_DRAFT_ALLOC
for (unsigned i = 0; i < 3; i++) {
d->target[i] = ds4_gpu_tensor_alloc((uint64_t)DS41_DRAFT_WINDOW * DS4_N_EMBD * 4u);
d->window[i] = ds4_gpu_tensor_alloc((uint64_t)(DS41_DRAFT_WINDOW + DS41_DRAFT_ROWS) * DS4_N_HEAD_DIM * 4u);
if (!d->target[i] || !d->window[i]) goto fail;
}
#define DS41_DRAFT_A(name, count) \
if (!(d->name = ds4_gpu_tensor_alloc((uint64_t)(count) * sizeof(float)))) goto fail;
DS41_DRAFT_A(mean_weights, DS41_DRAFT_WINDOW * DS4_N_HC)
DS41_DRAFT_A(mean, DS41_DRAFT_WINDOW * DS4_N_EMBD)
DS41_DRAFT_A(input, DS41_DRAFT_WINDOW * 3u * DS4_N_EMBD)
DS41_DRAFT_A(projected, DS41_DRAFT_WINDOW * DS4_N_EMBD)
DS41_DRAFT_A(main_x, DS41_DRAFT_WINDOW * DS4_N_EMBD)
DS41_DRAFT_A(main_kv, DS41_DRAFT_WINDOW * DS4_N_HEAD_DIM)
DS41_DRAFT_A(logits, DS41_DRAFT_ROWS * DS4_N_VOCAB)
DS41_DRAFT_A(tokens, DS41_DRAFT_ROWS)
DS41_DRAFT_A(markov_embed, dw->markov_rank)
DS41_DRAFT_A(markov_logits, DS4_N_VOCAB)
DS41_DRAFT_A(confidence_input, DS4_N_EMBD + dw->markov_rank)
DS41_DRAFT_A(confidence, 1)
DS41_DRAFT_A(verify_residual, (DS41_VERIFY_ROWS + 1u) * DS4_N_HC * DS4_N_EMBD)
DS41_DRAFT_A(verify_pre, (DS41_VERIFY_ROWS + 1u) * DS4_N_HC)
DS41_DRAFT_A(verify_logits, (DS41_VERIFY_ROWS + 1u) * DS4_N_VOCAB)
DS41_DRAFT_A(verify_pool, (DS41_VERIFY_ROWS + 1u) * 8u * DS4_N_HEAD_DIM)
DS41_DRAFT_A(old_raw, 40u * DS41_VERIFY_ROWS * DS4_N_HEAD_DIM)
DS41_DRAFT_A(old_features, 3u * DS41_VERIFY_ROWS * DS4_N_EMBD)
#undef DS41_DRAFT_A
d->host_logits = malloc((uint64_t)DS41_DRAFT_ROWS * DS4_N_VOCAB * sizeof(float));
d->draft_probs = malloc((uint64_t)DS41_DRAFT_ROWS * DS4_N_VOCAB * sizeof(float));
d->host_verified = malloc((uint64_t)(DS41_VERIFY_ROWS + 1u) * DS4_N_VOCAB * sizeof(float));
if (!d->host_logits || !d->draft_probs || !d->host_verified || !ds4_gpu_tensor_fill_f32(d->mean_weights,
1.0f / DS4_N_HC, DS41_DRAFT_WINDOW * DS4_N_HC) ||
!ds4_gpu_tensor_fill_f32(d->tokens, 0, DS41_DRAFT_ROWS)) goto fail;
ds41_dspark_reset(d);
return d;
fail:
ds41_dspark_free(d);
return NULL;
}
static bool ds41_dspark_capture(ds41_dspark *d, uint32_t layer, uint32_t pos,
const ds4_gpu_tensor *residual, uint32_t rows) {
if (!d) return true;
d->rewind_end = 0;
unsigned stage = 0;
while (stage < 3 && d->weights->target_layers[stage] != layer) stage++;
if (stage == 3) return true;
if (!rows || rows > UINT32_MAX - pos) return false;
const uint32_t count = rows < DS41_DRAFT_WINDOW ? rows : DS41_DRAFT_WINDOW;
const uint32_t first = rows - count, start = pos + first;
const uint64_t bytes = (uint64_t)DS4_N_EMBD * sizeof(float);
ds4_gpu_tensor *src = ds4_gpu_tensor_view((ds4_gpu_tensor *)residual,
first * DS4_N_HC * bytes, count * DS4_N_HC * bytes);
ds4_gpu_tensor *out = ds4_gpu_tensor_view(d->mean, 0, count * bytes);
ds4_gpu_tensor *mix = ds4_gpu_tensor_view(d->mean_weights, 0, count * DS4_N_HC * sizeof(float));
bool ok = src && out && mix && ds4_gpu_hc_weighted_sum_tensor(out, src, mix, DS4_N_EMBD, DS4_N_HC) &&
ds4_gpu_dsv41_quantize(out, DS4_N_EMBD, count, DS4_V41_BF16);
const uint32_t slot = start % DS41_DRAFT_WINDOW;
const uint32_t part = count < DS41_DRAFT_WINDOW - slot ? count : DS41_DRAFT_WINDOW - slot;
if (ok) ok = ds4_gpu_tensor_copy(d->target[stage], slot * bytes, out, 0, part * bytes) &&
(part == count || ds4_gpu_tensor_copy(d->target[stage], 0, out, part * bytes, (count - part) * bytes));
ds4_gpu_tensor_free(mix); ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(src);
if (ok) for (uint32_t i = 0; i < count; i++)
d->positions[stage][(start + i) % DS41_DRAFT_WINDOW] = start + i;
return ok;
}
static bool ds41_dspark_matmul_batch(ds4_gpu_tensor *out, const ds4_model *m,
const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t count, bool round) {
#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD)
/* Draft rows share a GEMM. Target verification retains its scalar-reduction
* kernels; proposal arithmetic need not reproduce their reduction order. */
if (count > 1 && weight->type == DS4_TENSOR_F16 &&
!getenv("DS4_DISABLE_V41_DRAFT_GEMM"))
return metal_graph_matmul_plain_tensor(out, m, weight, weight->dim[0],
weight->dim[1], in, count) && (!round ||
ds4_gpu_dsv41_quantize(out, (uint32_t)weight->dim[1], count, DS4_V41_BF16));
#endif
return ds41_matmul_batch(out, m, weight, in, count, round);
}
/* Only accepted target features enter these rings. A fresh prefill seeds the
* last window once; subsequent calls project only the newly accepted rows. */
static bool ds41_dspark_seed(ds41_dspark *d, const ds4_model *m, uint32_t pos) {
if (!pos) return false;
const uint32_t oldest = pos > DS41_DRAFT_WINDOW ? pos - DS41_DRAFT_WINDOW : 0;
const uint32_t start = d->kv_valid && d->kv_end >= oldest && d->kv_end <= pos ? d->kv_end : oldest;
const uint32_t count = pos - start;
if (!count) return true;
const uint64_t row_bytes = (uint64_t)DS4_N_EMBD * sizeof(float);
for (uint32_t i = 0; i < count; i++) {
const uint32_t slot = (start + i) % DS41_DRAFT_WINDOW;
for (unsigned stage = 0; stage < 3; stage++) {
if (d->positions[stage][slot] != start + i ||
!ds4_gpu_tensor_copy(d->input, (i * 3u + stage) * row_bytes,
d->target[stage], slot * row_bytes, row_bytes)) return false;
}
}
const ds4_dspark_stage_weights *first = &d->weights->stage[0];
if (!ds41_dspark_matmul_batch(d->projected, m, first->main_proj, d->input, count, true) ||
!ds41_norm_batch(d->main_x, d->projected, m, first->main_norm, count)) return false;
const uint64_t kv_bytes = DS4_N_HEAD_DIM * sizeof(float);
for (unsigned stage = 0; stage < 3; stage++) {
const ds4_layer_weights *l = &d->weights->stage[stage].block;
if (!ds41_dspark_matmul_batch(d->main_kv, m, l->attn_kv, d->main_x, count, true) ||
!ds41_norm_batch(d->main_kv, d->main_kv, m, l->attn_kv_a_norm, count) ||
!ds4_gpu_dsv41_rope(d->main_kv, DS4_N_HEAD_DIM, 1, count, start, false, false) ||
!ds4_gpu_dsv41_quantize(d->main_kv, DS4_N_HEAD_DIM, count, DS4_V41_FP8_E8M0)) return false;
const uint32_t slot = start % DS41_DRAFT_WINDOW;
const uint32_t part = count < DS41_DRAFT_WINDOW - slot ? count : DS41_DRAFT_WINDOW - slot;
if (!ds4_gpu_tensor_copy(d->window[stage], slot * kv_bytes, d->main_kv, 0, part * kv_bytes) ||
(part < count && !ds4_gpu_tensor_copy(d->window[stage], 0, d->main_kv,
part * kv_bytes, (count - part) * kv_bytes))) return false;
}
d->kv_end = pos;
d->kv_valid = true;
return true;
}
static bool ds41_dspark_attention_output(ds41_prefill_row *b, const ds4_model *m,
const ds4_layer_weights *l) {
const uint32_t groups = DS4_N_OUT_GROUP, width = DS4_N_HEAD * DS4_N_HEAD_DIM / groups;
uint64_t weight_row;
if (!tensor_nbytes(l->attn_output_a->type, width, &weight_row)) return false;
for (uint32_t row = 0; row < DS41_DRAFT_ROWS; row++) {
for (uint32_t group = 0; group < groups; group++) {
ds4_tensor slice = *l->attn_output_a;
slice.abs_offset += group * DS4_N_LORA_O * weight_row;
slice.dim[1] = DS4_N_LORA_O;
ds4_gpu_tensor *x = ds4_gpu_tensor_view(b->heads,
((uint64_t)row * groups + group) * width * 4u, width * 4u);
ds4_gpu_tensor *y = ds4_gpu_tensor_view(b->low,
((uint64_t)row * groups + group) * DS4_N_LORA_O * 4u, DS4_N_LORA_O * 4u);
const bool ok = x && y && ds41_matmul(y, m, &slice, x, true);
ds4_gpu_tensor_free(y); ds4_gpu_tensor_free(x);
if (!ok) return false;
}
}
return ds41_dspark_matmul_batch(b->block, m, l->attn_output_b, b->low, DS41_DRAFT_ROWS, true);
}
static bool ds41_dspark_stage(ds41_dspark *d, const ds4_model *m, uint32_t stage, uint32_t pos) {
ds41_prefill_row *b = &d->batch;
const ds4_layer_weights *l = &d->weights->stage[stage].block;
const uint32_t n = DS41_DRAFT_ROWS, past = pos < DS41_DRAFT_WINDOW ? pos : DS41_DRAFT_WINDOW;
if (!ds41_hc_mix_batch(b, m, l, false, n) ||
!(stage ? ds4_gpu_hc_weighted_sum_split_tensor(b->x, b->residual, b->ffn_split, DS4_N_EMBD, DS4_N_HC) :
ds4_gpu_hc_weighted_sum_tensor(b->x, b->residual, b->pre, DS4_N_EMBD, DS4_N_HC)) ||
!ds4_gpu_dsv41_quantize(b->x, DS4_N_EMBD, n, DS4_V41_BF16) ||
!ds41_norm_batch(b->norm, b->x, m, l->attn_norm, n) ||
!ds41_dspark_matmul_batch(b->qr, m, l->attn_q_a, b->norm, n, true) ||
!ds41_norm_batch(b->qr, b->qr, m, l->attn_q_a_norm, n) ||
!ds41_dspark_matmul_batch(b->q, m, l->attn_q_b, b->qr, n, true) ||
!ds4_gpu_dsv41_rope(b->q, DS4_N_HEAD_DIM, DS4_N_HEAD, n, pos, false, false) ||
!ds41_dspark_matmul_batch(b->kv, m, l->attn_kv, b->norm, n, true) ||
!ds41_norm_batch(b->kv, b->kv, m, l->attn_kv_a_norm, n) ||
!ds4_gpu_dsv41_rope(b->kv, DS4_N_HEAD_DIM, 1, n, pos, false, false) ||
!ds4_gpu_dsv41_quantize(b->kv, DS4_N_HEAD_DIM, n, DS4_V41_FP8_E8M0) ||
!ds4_gpu_tensor_copy(d->window[stage], (uint64_t)past * DS4_N_HEAD_DIM * 4u,
b->kv, 0, (uint64_t)n * DS4_N_HEAD_DIM * 4u) ||
!ds4_gpu_attention_noncausal_raw_batch_heads_tensor(b->heads, m->map, m->size,
l->attn_sinks->abs_offset, b->q, d->window[stage], n, past + n,
DS41_DRAFT_WINDOW + n, 0, DS4_N_HEAD, DS4_N_HEAD_DIM) ||
!ds4_gpu_dsv41_quantize(b->heads, DS4_N_HEAD * DS4_N_HEAD_DIM, n, DS4_V41_BF16) ||
!ds4_gpu_dsv41_rope(b->heads, DS4_N_HEAD_DIM, DS4_N_HEAD, n, pos, false, true) ||
!ds41_dspark_attention_output(b, m, l) ||
!ds41_after_attention_batch(b, m, l, n, false)) {
fprintf(stderr, "ds4: V4.1 draft attention failed\n");
return false;
}
const uint32_t experts = d->weights->n_experts, active = d->weights->n_experts_used;
const uint64_t gate_row = routed_expert_row_bytes(l->ffn_gate_exps);
const uint64_t down_row = routed_expert_row_bytes(l->ffn_down_exps);
bool mid_f16 = false;
if (!(ds41_dspark_matmul_batch(b->route_logits, m, l->ffn_gate_inp, b->norm, n, false) &&
ds4_gpu_router_select_batch_tensor(b->selected, b->route_weights, b->route_probs,
m->map, m->size, l->ffn_exp_probs_b->abs_offset, 0, 0, 0, 0, true, false,
b->route_logits, d->tokens, experts, active, DS4_EXPERT_WEIGHT_SCALE, n))) {
fprintf(stderr, "ds4: V4.1 draft routing failed\n"); return false;
}
if (!(ds41_dspark_matmul_batch(b->shared_gate, m, l->ffn_gate_shexp, b->norm, n, true) &&
ds41_dspark_matmul_batch(b->shared_up, m, l->ffn_up_shexp, b->norm, n, true) &&
ds4_gpu_swiglu_tensor(b->shared_mid, b->shared_gate, b->shared_up,
n * DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP, 1.0f) &&
ds4_gpu_dsv41_quantize(b->shared_mid, DS4_N_FF_EXP, n, DS4_V41_BF16) &&
ds41_dspark_matmul_batch(b->shared, m, l->ffn_down_shexp, b->shared_mid, n, true))) {
fprintf(stderr, "ds4: V4.1 draft shared expert failed\n"); return false;
}
if (!ds4_gpu_routed_moe_batch_tensor(b->routed, b->gate, b->up, b->mid, b->experts,
m->map, m->size, l->ffn_gate_exps->abs_offset, l->ffn_up_exps->abs_offset,
l->ffn_down_exps->abs_offset, l->ffn_gate_exps->type, l->ffn_down_exps->type,
gate_row * DS4_N_FF_EXP, gate_row, down_row * DS4_N_EMBD, down_row,
DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, b->selected, b->route_weights,
experts, active, DS4_SWIGLU_CLAMP_EXP, b->norm, stage, n, &mid_f16, true)) {
fprintf(stderr, "ds4: V4.1 draft routed experts failed\n"); return false;
}
return ds4_gpu_add_tensor(b->block, b->routed, b->shared, n * DS4_N_EMBD) &&
ds4_gpu_dsv41_quantize(b->block, DS4_N_EMBD, n, DS4_V41_BF16) &&
ds4_gpu_hc_expand_split_tensor(b->residual, b->block, b->after_attn,
b->ffn_split, DS4_N_EMBD, DS4_N_HC) &&
ds4_gpu_dsv41_quantize(b->residual, DS4_N_HC * DS4_N_EMBD, n, DS4_V41_BF16);
}
/* Leaves the five base-logit rows and unnormalized collapsed features on the
* GPU. Markov sampling is sequential and performed after this stage chain. */
static DS4_MAYBE_UNUSED bool ds41_dspark_forward(ds41_dspark *d, const ds4_model *target,
const ds4_weights *tw, const ds4_model *draft,
uint32_t pos, int token, ds41_gpu_graph *target_graph) {
if (!d || !pos || token < 0 || (uint32_t)token >= DS4_N_VOCAB) return false;
ds41_prefill_row *b = &d->batch;
const bool local_experts = target_graph && target_graph->tp_world == 2;
if (local_experts) ds4_gpu_tp_suspend_expert_sharding(1);
bool ok = ds4_gpu_begin_commands() && ds41_dspark_seed(d, draft, pos);
if (!ok) fprintf(stderr, "ds4: V4.1 DSpark seed failed\n");
const float pre[DS41_DRAFT_ROWS][4] = {{1, 0, 0, 0}, {1, 0, 0, 0}, {1, 0, 0, 0}, {1, 0, 0, 0}, {1, 0, 0, 0}};
if (ok) ok = ds4_gpu_tensor_write(b->pre, 0, pre, sizeof(pre));
for (uint32_t row = 0; ok && row < DS41_DRAFT_ROWS; row++) {
const uint64_t bytes = (uint64_t)DS4_N_EMBD * DS4_N_HC * sizeof(float);
ds4_gpu_tensor *out = ds4_gpu_tensor_view(b->residual, row * bytes, bytes);
ok = out && ds4_gpu_embed_token_hc_tensor(out, target->map, target->size,
tw->token_embd->abs_offset, DS4_N_VOCAB, row ? d->weights->noise_token_id : (uint32_t)token,
DS4_N_EMBD, DS4_N_HC);
ds4_gpu_tensor_free(out);
}
for (uint32_t stage = 0; ok && stage < 3; stage++) {
ok = ds41_dspark_stage(d, draft, stage, pos);
if (!ok) fprintf(stderr, "ds4: V4.1 DSpark stage %u failed\n", stage);
}
if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(b->x, b->residual,
b->ffn_split, DS4_N_EMBD, DS4_N_HC) &&
ds4_gpu_dsv41_quantize(b->x, DS4_N_EMBD, DS41_DRAFT_ROWS, DS4_V41_BF16) &&
ds41_norm_batch(b->norm, b->x, draft, d->weights->stage[2].norm, DS41_DRAFT_ROWS);
if (ok) ok = target_graph ?
ds41_output_projection(target_graph, d->logits, target, tw, b->norm, DS41_DRAFT_ROWS) :
ds41_matmul_batch(d->logits, target, tw->output, b->norm, DS41_DRAFT_ROWS, false);
if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
if (local_experts) ds4_gpu_tp_suspend_expert_sharding(0);
if (!ok) d->kv_valid = false;
return ok;
}
static DS4_MAYBE_UNUSED bool ds41_dspark_markov(ds41_dspark *d, const ds4_model *m,
uint32_t row, int previous, float *confidence) {
if (!d || row >= DS41_DRAFT_ROWS || previous < 0 ||
(uint32_t)previous >= DS4_N_VOCAB || !confidence) return false;
const ds4_dspark_stage_weights *last = &d->weights->stage[2];
float embedding[256];
const uint32_t rank = d->weights->markov_rank;
if (rank != 256 || !dspark_dense_row_to_f32(embedding, m, last->markov_w1, (uint32_t)previous))
return false;
const uint64_t bytes = (uint64_t)DS4_N_VOCAB * sizeof(float);
ds4_gpu_tensor *logits = ds4_gpu_tensor_view(d->logits, row * bytes, bytes);
bool ok = logits && ds4_gpu_begin_commands() &&
ds4_gpu_tensor_write(d->markov_embed, 0, embedding, sizeof(embedding)) &&
ds41_matmul(d->markov_logits, m, last->markov_w2, d->markov_embed, false) &&
ds4_gpu_add_tensor(logits, logits, d->markov_logits, DS4_N_VOCAB) &&
ds4_gpu_tensor_copy(d->confidence_input, 0, d->batch.x,
(uint64_t)row * DS4_N_EMBD * 4u, (uint64_t)DS4_N_EMBD * 4u) &&
ds4_gpu_tensor_copy(d->confidence_input, (uint64_t)DS4_N_EMBD * 4u,
d->markov_embed, 0, rank * 4u) &&
ds41_matmul(d->confidence, m, last->confidence_proj, d->confidence_input, false);
if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
if (ok) ok = ds4_gpu_tensor_read(logits, 0, d->host_logits + (uint64_t)row * DS4_N_VOCAB, bytes) &&
ds4_gpu_tensor_read(d->confidence, 0, confidence, sizeof(*confidence));
ds4_gpu_tensor_free(logits);
return ok;
}
/* Cache appends beyond the accepted frontier are invisible. Only the six
* overwritten raw-ring slots and pooling carries need rollback; no compressed
* cache proportional to the context is copied. */
static DS4_MAYBE_UNUSED bool ds41_dspark_verify(ds41_dspark *d, ds41_gpu_graph *g,
const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t count) {
if (!d || !g->valid || count < 2 || count > DS41_VERIFY_ROWS ||
g->pos > g->ctx || count > g->ctx - g->pos || !ds41_tp_batch_enabled(g)) return false;
/* Prepost the complete small-row exchange window, as in V4.0 verification.
* Bulk-prefill handshakes at every attention/FFN gate are too expensive. */
const bool tp_block = g->tp_world == 2;
uint32_t gates = 2;
if (tp_block && (!g->tp_batch_out || !g->tp_batch_in || !g_tp_block_ctx ||
!ds4_tp_batch_block_begin_gates(g_tp_block_ctx, count, DS4_N_LAYER, &gates) ||
gates != 2)) return false;
const uint64_t kv = DS4_N_HEAD_DIM * 4u, feature = DS4_N_EMBD * 4u;
ds41_gpu_graph *graphs[DS41_VERIFY_ROWS];
bool ok = ds4_gpu_begin_commands();
d->verify_start = g->pos;
d->verify_count = 0;
d->rewind_end = 0;
for (uint32_t i = 0; ok && i <= DS41_VERIFY_ROWS; i++) {
if (i >= count && i != DS41_VERIFY_ROWS) continue;
const uint32_t slot = (g->pos + i) % DS41_DRAFT_WINDOW;
for (uint32_t il = 0; ok && i < count && il < 40; il++)
ok = ds4_gpu_tensor_copy(d->old_raw, (il * DS41_VERIFY_ROWS + i) * kv,
g->window[il], slot * kv, kv);
for (uint32_t stage = 0; ok && i < count && stage < 3; stage++) {
d->old_positions[stage][i] = d->positions[stage][slot];
ok = ds4_gpu_tensor_copy(d->old_features, (stage * DS41_VERIFY_ROWS + i) * feature,
d->target[stage], slot * feature, feature);
}
ds41_gpu_graph *row = &d->verified[i];
ds4_gpu_tensor *res = row->residual, *pre = row->pre, *logits = row->logits;
ds4_gpu_tensor *half = row->tp_logits_half;
ds4_gpu_tensor *pool_k[4], *pool_s[4];
memcpy(pool_k, row->verify_previous_kv, sizeof(pool_k));
memcpy(pool_s, row->verify_previous_score, sizeof(pool_s));
*row = *g;
row->residual = res ? res : ds4_gpu_tensor_view(d->verify_residual, i * DS4_N_HC * feature, DS4_N_HC * feature);
row->pre = pre ? pre : ds4_gpu_tensor_view(d->verify_pre, i * DS4_N_HC * 4u, DS4_N_HC * 4u);
row->logits = logits ? logits : ds4_gpu_tensor_view(d->verify_logits,
(uint64_t)i * DS4_N_VOCAB * 4u, (uint64_t)DS4_N_VOCAB * 4u);
row->tp_logits_half = NULL;
if (g->tp_logits_half) {
const uint64_t bytes = (uint64_t)DS4_N_VOCAB * 2u;
row->tp_logits_half = half ? half : ds4_gpu_tensor_view(row->logits,
g->tp_rank * bytes, bytes);
ok = ok && row->tp_logits_half;
}
row->verify_prefix = true;
for (uint32_t owner = 0; owner < 4; owner++) {
const uint64_t off = (i * 8u + owner * 2u) * kv;
row->verify_previous_kv[owner] = pool_k[owner] ? pool_k[owner] :
ds4_gpu_tensor_view(d->verify_pool, off, kv);
row->verify_previous_score[owner] = pool_s[owner] ? pool_s[owner] :
ds4_gpu_tensor_view(d->verify_pool, off + kv, kv);
ok = ok && row->verify_previous_kv[owner] && row->verify_previous_score[owner];
}
ok = ok && row->residual && row->pre && row->logits;
if (i < count) graphs[i] = row;
else {
/* Keep the block start as well: exact sampling can change mode
* at the first token, before any verified token is consumed. */
ok = ok && ds4_gpu_tensor_copy(row->residual, 0, g->residual, 0, DS4_N_HC * feature) &&
ds4_gpu_tensor_copy(row->pre, 0, g->pre, 0, DS4_N_HC * 4u) &&
ds4_gpu_tensor_copy(row->logits, 0, g->logits, 0, (uint64_t)DS4_N_VOCAB * 4u);
for (uint32_t owner = 0; ok && owner < 4; owner++)
ok = ds4_gpu_tensor_copy(row->verify_previous_kv[owner], 0, g->previous_kv[owner], 0, kv) &&
ds4_gpu_tensor_copy(row->verify_previous_score[owner], 0, g->previous_score[owner], 0, kv);
}
}
if (ok) ok = ds41_graph_step_batch(graphs, tokens, (int)count, count, m, w);
if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
if (tp_block && (!ok || !ds4_tp_batch_block_end(g_tp_block_ctx))) {
ds4_tp_mark_failed(g_tp_block_ctx);
ok = false;
}
if (ok) ok = ds4_gpu_tensor_read(d->verify_logits, 0, d->host_verified,
(uint64_t)count * DS4_N_VOCAB * 4u) &&
ds4_gpu_tensor_read(d->verify_logits, (uint64_t)DS41_VERIFY_ROWS * DS4_N_VOCAB * 4u,
d->host_verified + (uint64_t)DS41_VERIFY_ROWS * DS4_N_VOCAB, DS4_N_VOCAB * 4u);
if (ok) d->verify_count = count;
else g->valid = false;
return ok;
}
static DS4_MAYBE_UNUSED bool ds41_dspark_commit(ds41_dspark *d, ds41_gpu_graph *g,
uint32_t keep, float *logits) {
if (!d || !d->verify_count || keep > d->verify_count ||
(g->pos != d->verify_start && g->pos != d->rewind_end) ||
(d->rewind_end && d->verify_start + keep > d->rewind_end)) return false;
const uint64_t kv = DS4_N_HEAD_DIM * 4u, feature = DS4_N_EMBD * 4u;
const uint32_t row = keep ? keep - 1u : DS41_VERIFY_ROWS;
ds41_gpu_graph *last = &d->verified[row];
bool ok = ds4_gpu_begin_commands();
for (uint32_t i = keep; ok && i < d->verify_count; i++) {
const uint32_t slot = (d->verify_start + i) % DS41_DRAFT_WINDOW;
for (uint32_t il = 0; ok && il < 40; il++)
ok = ds4_gpu_tensor_copy(g->window[il], slot * kv,
d->old_raw, (il * DS41_VERIFY_ROWS + i) * kv, kv);
for (uint32_t stage = 0; ok && stage < 3; stage++) {
ok = ds4_gpu_tensor_copy(d->target[stage], slot * feature,
d->old_features, (stage * DS41_VERIFY_ROWS + i) * feature, feature);
d->positions[stage][slot] = d->old_positions[stage][i];
}
}
for (uint32_t owner = 0; ok && owner < 4; owner++)
ok = ds4_gpu_tensor_copy(g->previous_kv[owner], 0, last->verify_previous_kv[owner], 0, kv) &&
ds4_gpu_tensor_copy(g->previous_score[owner], 0, last->verify_previous_score[owner], 0, kv);
if (ok) ok = ds4_gpu_tensor_copy(g->residual, 0, last->residual, 0, DS4_N_HC * feature) &&
ds4_gpu_tensor_copy(g->pre, 0, last->pre, 0, DS4_N_HC * 4u) &&
ds4_gpu_tensor_copy(g->logits, 0, last->logits, 0, (uint64_t)DS4_N_VOCAB * 4u);
if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false;
if (ok) {
g->pos = last->pos;
g->history = last->history;
memcpy(logits, d->host_verified + (uint64_t)row * DS4_N_VOCAB, DS4_N_VOCAB * 4u);
} else g->valid = false;
d->rewind_end = ok ? g->pos : 0;
return ok;
}
#endif
ds4_deepseek41_gpu.h
+315
View File
@@ -0,0 +1,315 @@
#ifndef DS4_DEEPSEEK41_GPU_H
#define DS4_DEEPSEEK41_GPU_H
#include <stdbool.h>
#include <stdint.h>
#ifndef DS4_GPU_TENSOR_DEFINED
#define DS4_GPU_TENSOR_DEFINED
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
#endif
#ifdef __cplusplus
extern "C" {
#endif
/* V4.1 activation/cache formats. Buffers are float-addressable but the
* rounded values follow the released BF16/FP8/FP4 inference graph. */
typedef enum {
DS4_V41_BF16 = 0,
DS4_V41_FP8_E8M0 = 1,
DS4_V41_FP4_E8M0 = 2,
DS4_V41_FP4_E4M3 = 3,
} ds4_v41_activation_format;
/* One-token router: probabilities, biased top-k and normalised weights in one dispatch. */
int ds4_gpu_dsv41_router_one(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs, const ds4_gpu_tensor *logits,
const void *model_map, uint64_t model_size, uint64_t bias_offset,
uint32_t n_expert, uint32_t top_k, float scale);
/* The same for a few rows (Metal only), and the rows' F32 logits with the
* single-row matvec's reduction. */
int ds4_gpu_dsv41_router_rows(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs, const ds4_gpu_tensor *logits,
const void *model_map, uint64_t model_size, uint64_t bias_offset,
uint32_t n_expert, uint32_t top_k, float scale, uint32_t rows);
int ds4_gpu_dsv41_matmul_f32_rows(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, uint32_t rows);
int ds4_gpu_dsv41_quantize(ds4_gpu_tensor *x, uint32_t width, uint32_t rows,
ds4_v41_activation_format format);
#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU)
/* Same reductions as the ordinary kernels, with the required BF16 store
* folded into the operation. hc_sum reads either 4- or 24-float row weights. */
int ds4_gpu_dsv41_norm(ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *map, uint64_t size, uint64_t offset,
uint32_t width, uint32_t rows, float eps);
int ds4_gpu_dsv41_hc_sum(ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const ds4_gpu_tensor *weights, uint32_t width, uint32_t rows,
uint32_t weight_stride);
int ds4_gpu_dsv41_hc_expand(ds4_gpu_tensor *out, const ds4_gpu_tensor *block,
const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split,
uint32_t width, uint32_t rows);
/* CUDA scalar Q8 shared expert. 1: queued; 0: unsupported, no work queued;
* -1: failure. After 1, keep the input/output tensors alive and unchanged
* until join, which orders the result before subsequent main-stream work.
* Other work may run between start and join using separate tensors.
* Optional ownership picks the rank with fewer selected routed experts.
* IDs are read on the main stream before later owned-expert filtering. */
typedef struct {
const ds4_gpu_tensor *selected;
uint32_t experts, active, rank, tie_owner;
} ds4_v41_shared_owner;
int ds4_gpu_dsv41_shared_start(
ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up,
ds4_gpu_tensor *mid, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size,
uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset,
uint32_t width, uint32_t hidden, float clamp,
const ds4_v41_shared_owner *owner);
int ds4_gpu_dsv41_shared_join(void);
/* Mask a computed short batch before the owned router filter mutates IDs. */
int ds4_gpu_dsv41_shared_mask(ds4_gpu_tensor *shared, const ds4_gpu_tensor *selected,
uint32_t width, uint32_t rows, uint32_t experts, uint32_t active,
uint32_t rank, uint32_t tie_owner);
#endif
/* Full-head prefill, with BF16 rounding between the two Q8 projections. */
int ds4_gpu_dsv41_attention_output_batch(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
const ds4_gpu_tensor *heads, uint32_t n_tokens);
/* Packed local 32-head input and BF16 low projection; output is an unrounded
* rank partial. The graph sums ranks before rounding the attention block. */
int ds4_gpu_dsv41_attention_output_tp_batch(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
const ds4_gpu_tensor *heads, uint32_t n_tokens, uint32_t tp_rank);
/* Adjacent-pair, unit-magnitude RoPE with the released V4.1 frequencies. */
int ds4_gpu_dsv41_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, bool compressed, bool inverse);
/* Compressed pairs advance two absolute token positions per stored row. */
int ds4_gpu_dsv41_rope_bf16(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, bool compressed, bool inverse);
int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, uint32_t stride,
bool compressed, bool inverse);
int ds4_gpu_dsv41_engram_add(ds4_gpu_tensor *residual,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *q_weight,
const ds4_gpu_tensor *k_weight,
const ds4_gpu_tensor *mask,
uint32_t width, uint32_t rows, float eps);
/* Pool complete pairs and retain the last unpaired projection in previous_*.
* start is the absolute token position, including earlier chunks. */
int ds4_gpu_dsv41_pool2(ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv, const ds4_gpu_tensor *scores,
ds4_gpu_tensor *previous_kv, ds4_gpu_tensor *previous_scores,
uint32_t width, uint32_t rows, uint32_t start);
/* Candidate blocks contain eight compressed positions. Produce causal block
* maxima, pinning the newest block; filter consumes a 0/-inf block mask. */
int ds4_gpu_dsv41_candidate_blocks(ds4_gpu_tensor *blocks,
const ds4_gpu_tensor *scores,
uint32_t width, uint32_t rows,
uint32_t start, uint32_t ratio);
int ds4_gpu_dsv41_candidate_filter(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *block_mask,
uint32_t width, uint32_t rows,
uint32_t start, uint32_t ratio,
uint32_t mask_stride);
/* Per-row block top-2048 with the causal block count of ratio-1 rows, and
* the 0/-inf block masks it selects at mask_stride floats per row. */
int ds4_gpu_dsv41_candidate_topk_batch(ds4_gpu_tensor *selected,
const ds4_gpu_tensor *blocks,
uint32_t width, uint32_t rows, uint32_t start);
int ds4_gpu_dsv41_candidate_mask_batch(ds4_gpu_tensor *mask,
const ds4_gpu_tensor *selected,
uint32_t width, uint32_t rows,
uint32_t mask_stride);
/* Causal index scores over ratio-1/2 compressed keys, without an extra cast
* of the already quantized FP4 queries/keys. Scores have source_rows stride. */
int ds4_gpu_dsv41_indexer_scores_batch(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *keys,
uint32_t source_rows, uint32_t rows,
uint32_t start, uint32_t ratio);
#ifndef __APPLE__
int ds4_gpu_dsv41_indexer_scores_masked_one(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys,
const ds4_gpu_tensor *block_mask, uint32_t width);
#endif
/* One-token index scores of FP4 queries/keys: the GLM scorer, plus its wide
* kernel for long rows (Metal only). */
int ds4_gpu_dsv41_indexer_score_one_tensor(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows, uint32_t n_head,
uint32_t head_dim, float scale,
bool cache_f16);
/* Rows [first, first + rows) of the index-key cache into its bf16 view; with flag, a key that is not bf16-exact
* sets it nonzero. score_row: one decode index row (32 heads x 128, scale 1/64) over the cache, reading view, when
* given, where the wide scorer runs (the same bits). */
int ds4_gpu_dsv41_index_keys_bf16(ds4_gpu_tensor *view, const ds4_gpu_tensor *cache,
uint32_t first, uint32_t rows, ds4_gpu_tensor *flag);
int ds4_gpu_dsv41_indexer_score_row(ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *cache, const ds4_gpu_tensor *view, uint32_t rows);
/* The top-k of a row filtered by a block mask of more than 2048 blocks, from
* the scores of its 2048 candidate blocks only: the same ids in the same order
* (Metal only). Returns 0, leaving the full row to the caller, where unavailable. */
int ds4_gpu_dsv41_indexer_candidates_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *scores,
ds4_gpu_tensor *list, const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys,
const ds4_gpu_tensor *block_mask, uint32_t n_comp,
float scale, uint32_t top_k, int list_ready);
/* list_ready != 0: list already holds block_mask's candidates (built by an earlier call since the mask was written),
* so it is not rebuilt. */
int ds4_gpu_dsv41_tensor_ops_available(void);
/* Reuse exact BF16 views of FP4 queries/keys across score tiles. Invalid
* casts retain the F32 arithmetic for the affected tile. */
uint64_t ds4_gpu_dsv41_indexer_packed_bytes(uint32_t source_rows, uint32_t rows);
int ds4_gpu_dsv41_indexer_pack(ds4_gpu_tensor *packed,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *keys,
uint32_t source_rows, uint32_t rows);
int ds4_gpu_dsv41_indexer_scores_packed(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *keys,
const ds4_gpu_tensor *packed,
uint32_t source_rows, uint32_t rows,
uint32_t start, uint32_t ratio,
uint32_t packed_rows, uint32_t offset);
/* Exact row-sort ordering with independent causal widths; more than 512
* visible keys per row. Output stride is 512 indices. */
int ds4_gpu_dsv41_indexer_topk_batch(ds4_gpu_tensor *selected,
const ds4_gpu_tensor *scores,
uint32_t width, uint32_t rows,
uint32_t start, uint32_t ratio);
/* The indexer top-k, plus the radix select for wide rows (Metal); other models
* keep the merge ladder and its tie order. */
int ds4_gpu_dsv41_indexer_topk_tensor(ds4_gpu_tensor *selected,
const ds4_gpu_tensor *scores,
uint32_t n_comp, uint32_t n_tokens,
uint32_t top_k);
/* Rows with at most 512 visible keys select all of them, in key order. */
int ds4_gpu_dsv41_indexer_all_batch(ds4_gpu_tensor *selected, uint32_t rows,
uint32_t start, uint32_t ratio);
enum { DS4_V41_CARRY_BF16, DS4_V41_CARRY_MASK, DS4_V41_CARRY_F32 };
/* Lossless storage for already-BF16 activations or 0/-inf candidate masks.
* Packed rows are padded to whole uint32_t words. Plain rows remain F32. */
int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset,
ds4_gpu_tensor *plain, uint32_t width, uint32_t rows,
uint32_t format, bool pack);
/* Batched F16 projections with the same arithmetic as individual matvecs. */
int ds4_gpu_dsv41_projection_rows(ds4_gpu_tensor *out,
const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t width,
uint32_t outputs, uint32_t rows,
const ds4_gpu_tensor *in);
/* The HC mixer projection of one row with its norm fused (0: not fused). */
int ds4_gpu_dsv41_hc_project(ds4_gpu_tensor *mix, const ds4_gpu_tensor *residual,
const void *model_map, uint64_t model_size, uint64_t fn_offset,
uint32_t n, uint32_t mix_dim, float eps);
/* One block's whole HC input in one dispatch: the mixer projection with its
* norm, then split, weighted sum with `pre`'s first n_hc weights and norm
* (0: not fused). */
/* The same for a few rows: `pre` advances pre_stride floats per row. */
int ds4_gpu_dsv41_hc_block_input_rows(ds4_gpu_tensor *mix, ds4_gpu_tensor *x, ds4_gpu_tensor *norm,
ds4_gpu_tensor *split, const ds4_gpu_tensor *stream,
const ds4_gpu_tensor *pre, const void *model_map, uint64_t model_size,
uint64_t fn_offset, uint64_t scale_offset, uint64_t base_offset,
uint64_t norm_offset, uint32_t n, uint32_t mix_dim, uint32_t n_embd,
uint32_t n_hc, uint32_t sinkhorn_iters, float hc_eps, float norm_eps,
uint32_t rows, uint32_t pre_stride);
/* One block's input: split, weighted sum with `pre`'s first n_hc weights, norm. */
int ds4_gpu_dsv41_hc_input(ds4_gpu_tensor *x, ds4_gpu_tensor *norm, ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix, const ds4_gpu_tensor *residual,
const ds4_gpu_tensor *pre, const void *model_map, uint64_t model_size,
uint64_t scale_offset, uint64_t base_offset, uint64_t norm_offset,
uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters,
float hc_eps, float norm_eps);
/* Two Q8_0 projections of one row, outputs rounded to bf16 (0: run them apart). */
int ds4_gpu_dsv41_project_pair_q8(ds4_gpu_tensor *out_a, ds4_gpu_tensor *out_b,
const void *model_map, uint64_t model_size,
uint64_t offset_a, uint64_t offset_b, uint32_t in_dim,
uint32_t out_a_dim, uint32_t out_b_dim, const ds4_gpu_tensor *x);
/* The shared expert's SwiGLU input from the Q8_0 gate and up weights, as the
* bf16 matvecs and the SwiGLU pass would produce it; gate/up are scratch. */
int ds4_gpu_dsv41_shared_swiglu(ds4_gpu_tensor *mid, ds4_gpu_tensor *gate, ds4_gpu_tensor *up,
const void *model_map, uint64_t model_size, uint64_t gate_offset,
uint64_t up_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, float clamp, uint32_t rows);
/* Rope, quantize and store one row into dst at dst_offset bytes; x may be
* left roped and quantized in place. */
int ds4_gpu_dsv41_rope_quantize(const ds4_gpu_tensor *x, ds4_gpu_tensor *dst, uint64_t dst_offset,
uint32_t width, uint32_t start, bool compressed,
ds4_v41_activation_format format);
/* One row's F16 projection rounded to bf16 in the kernel (0: not fused). */
int ds4_gpu_dsv41_project_f16_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x);
/* Decode attention over the raw ring and the compressed rows selected by
* ids, staged in one pass (0: not available; `selected` is gather scratch). */
int ds4_gpu_dsv41_attention_decode(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size,
uint64_t sinks_offset, const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv, uint32_t n_raw, uint32_t raw_cap,
uint32_t raw_start, const ds4_gpu_tensor *comp_cache,
const ds4_gpu_tensor *ids, ds4_gpu_tensor *selected,
uint32_t n_comp, uint32_t attended, uint32_t n_head, uint32_t head_dim);
/* Single-Mac decode: the block input's collapse and norm now, then its hyper-connection weights inside the query
* projection's dispatch, which also ropes and stores the layer's KV row. 0 when DS4_GPU_V41_DEFERRED_HC is off. */
int ds4_gpu_dsv41_deferred_collapse(ds4_gpu_tensor *x, ds4_gpu_tensor *norm,
const ds4_gpu_tensor *residual, const ds4_gpu_tensor *pre,
const void *model_map, uint64_t model_size, uint64_t norm_offset, float eps);
int ds4_gpu_dsv41_query_deferred_hc(ds4_gpu_tensor *q, ds4_gpu_tensor *mix, ds4_gpu_tensor *split,
const ds4_gpu_tensor *qr, const ds4_gpu_tensor *residual,
const void *model_map, uint64_t model_size, uint64_t q_offset, uint64_t fn_offset,
uint64_t scale_offset, uint64_t base_offset, uint32_t pos, bool compressed,
uint32_t iters, float hc_eps, float norm_eps, ds4_gpu_tensor *x, ds4_gpu_tensor *norm,
const ds4_gpu_tensor *pre, uint64_t norm_offset,
const ds4_gpu_tensor *kv, ds4_gpu_tensor *cache, uint64_t cache_offset);
/* The Q8_0 q projection of one row with the layer's rope on its outputs. */
int ds4_gpu_dsv41_project_q(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, uint32_t pos, bool compressed);
/* A bf16 Q8_0 matvec of one row expanded straight into the n_hc residual
* streams; `add` (optional) joins the row before the expand. */
int ds4_gpu_dsv41_matmul_expand_rows(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, const ds4_gpu_tensor *add,
const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split,
uint32_t n_hc, uint32_t rows);
/* The attention output low projection of one row with the heads' inverse
* rope (and bf16 rounding) folded into the load; heads may be roped in place. */
int ds4_gpu_dsv41_attention_low(ds4_gpu_tensor *low, const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint32_t group_dim, uint32_t rank,
uint32_t n_groups, ds4_gpu_tensor *heads, uint32_t pos, bool compressed);
/* Metal: one decode step's RoPE coefficient table (64 float2: (cos, sin) of float(pos) * frequency, entries 0..31
* raw, 32..63 compressed). rope_quantize_heads_tab: rope + quantization of heads rows of width in place, in one
* dispatch reading it. heads_once_arm: the next single-row decode reduce of 64 heads x 512 writes its heads already
* bf16-rounded and inverse-roped with the table; take returns 1 when it did (and disarms), and attention_low_plain
* then projects them. reduce_check (test entry): that reduce on tmp's partials, plain or transformed. */
int ds4_gpu_dsv41_rope_coeff_table(ds4_gpu_tensor *tab, uint32_t pos);
int ds4_gpu_dsv41_rope_quantize_heads_tab(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, uint32_t start,
bool compressed, ds4_v41_activation_format format, const ds4_gpu_tensor *tab);
void ds4_gpu_dsv41_heads_once_arm(const ds4_gpu_tensor *tab, bool compressed);
int ds4_gpu_dsv41_heads_once_take(void);
int ds4_gpu_dsv41_attention_low_plain(ds4_gpu_tensor *low, const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint32_t group_dim, uint32_t rank,
uint32_t n_groups, const ds4_gpu_tensor *heads);
int ds4_gpu_dsv41_flash_attn_reduce_check(ds4_gpu_tensor *heads, const ds4_gpu_tensor *tmp, uint32_t nrows,
int transform, const ds4_gpu_tensor *tab, bool compressed);
/* Two weighted RMS norms rounded to bf16 in one dispatch. */
int ds4_gpu_dsv41_norm_pair_rows(ds4_gpu_tensor *out0, const ds4_gpu_tensor *x0, uint64_t weight0_offset, uint32_t n0,
ds4_gpu_tensor *out1, const ds4_gpu_tensor *x1, uint64_t weight1_offset, uint32_t n1,
const void *model_map, uint64_t model_size, float eps, uint32_t rows);
/* Gather 512-wide F32 KV rows; IDs must come from top-k over source_rows. */
int ds4_gpu_dsv41_gather_kv(ds4_gpu_tensor *out, const ds4_gpu_tensor *source,
const ds4_gpu_tensor *ids, uint32_t source_rows,
uint32_t selected_rows);
#ifdef __cplusplus
}
#endif
#endif
ds4_deepseek4_vision_gpu.cuh
+444
View File
@@ -0,0 +1,444 @@
/* DeepSeek V4 Flash Vision-Exp operations shared by CUDA and ROCm. */
#ifndef DS4_DEEPSEEK4_VISION_TYPES_DEFINED
#define DS4_DEEPSEEK4_VISION_TYPES_DEFINED
#define DS4_DEEPSEEK4_VISION_LAYERS 32u
#define DS4_DEEPSEEK4_LANGUAGE_LAYERS 43u
#define DS4_DEEPSEEK4_MTP_LAYERS 3u
typedef struct {
uint64_t norm1;
uint64_t qkv_weight;
uint64_t qkv_bias;
uint64_t attn_proj_weight;
uint64_t attn_proj_bias;
uint64_t norm2;
uint64_t mlp_w1;
uint64_t mlp_w2;
} ds4_deepseek4_vision_layer_weights;
typedef struct {
uint64_t patch_weight;
uint64_t patch_bias;
uint64_t post_norm;
uint64_t aligner_w1;
uint64_t aligner_w1_bias;
uint64_t aligner_w2;
uint64_t aligner_w2_bias;
uint64_t image_start;
uint64_t image_pad;
uint64_t image_newline;
uint64_t image_end;
uint64_t visual_router_bias[DS4_DEEPSEEK4_LANGUAGE_LAYERS];
uint64_t mtp_visual_router_bias[DS4_DEEPSEEK4_MTP_LAYERS];
uint64_t hash_router_bias[3];
ds4_deepseek4_vision_layer_weights layer[DS4_DEEPSEEK4_VISION_LAYERS];
} ds4_deepseek4_vision_weights;
#endif
#ifndef DS4_DEEPSEEK4_VISION_STREAM
#define DS4_DEEPSEEK4_VISION_STREAM 0
#endif
__device__ __forceinline__ static float deepseek4_vision_round_bf16_dev(
float value) {
uint32_t bits = __float_as_uint(value);
if ((bits & 0x7f800000u) == 0x7f800000u) return value;
bits += 0x00007fffu + ((bits >> 16u) & 1u);
return __uint_as_float(bits & 0xffff0000u);
}
__global__ static void deepseek4_vision_round_bf16_kernel(
float *x,
uint64_t count) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < count) x[i] = deepseek4_vision_round_bf16_dev(x[i]);
}
__global__ static void deepseek4_vision_qkv_rope_kernel(
float *q,
float *k,
float *v,
const float *qkv,
const uint16_t *bias,
uint32_t rows,
uint32_t grid_w) {
const uint32_t row = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t lane = threadIdx.x;
if (row >= rows || head >= 16u || lane >= 32u) return;
const uint64_t qkv_base = (uint64_t)row * 3072u +
(uint64_t)head * 64u;
const uint64_t out_base = (uint64_t)row * 1024u +
(uint64_t)head * 64u;
const float q0 = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + lane] +
glm53_vision_bf16(bias + (uint64_t)head * 64u + lane));
const float q1 = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + lane + 32u] +
glm53_vision_bf16(bias + (uint64_t)head * 64u + lane + 32u));
const float k0 = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + 1024u + lane] +
glm53_vision_bf16(bias + 1024u +
(uint64_t)head * 64u + lane));
const float k1 = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + 1024u + lane + 32u] +
glm53_vision_bf16(bias + 1024u +
(uint64_t)head * 64u + lane + 32u));
const uint32_t y = row / grid_w;
const uint32_t x = row - y * grid_w;
const uint32_t pos = lane < 16u ? y : x;
const uint32_t freq = lane & 15u;
const float inv_freq = powf(10000.0f, -(float)freq / 16.0f);
const float angle = (float)pos * inv_freq;
const float cs = cosf(angle);
const float sn = sinf(angle);
q[out_base + lane] = deepseek4_vision_round_bf16_dev(q0 * cs - q1 * sn);
q[out_base + lane + 32u] =
deepseek4_vision_round_bf16_dev(q1 * cs + q0 * sn);
k[out_base + lane] = deepseek4_vision_round_bf16_dev(k0 * cs - k1 * sn);
k[out_base + lane + 32u] =
deepseek4_vision_round_bf16_dev(k1 * cs + k0 * sn);
v[out_base + lane] = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + 2048u + lane] +
glm53_vision_bf16(bias + 2048u +
(uint64_t)head * 64u + lane));
v[out_base + lane + 32u] = deepseek4_vision_round_bf16_dev(
qkv[qkv_base + 2048u + lane + 32u] +
glm53_vision_bf16(bias + 2048u +
(uint64_t)head * 64u + lane + 32u));
}
__global__ static void deepseek4_vision_swiglu_split_kernel(
float *out,
const float *gate_up,
uint64_t count,
uint32_t width) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count) return;
const uint32_t d = i % width;
const uint64_t row = i / width;
const uint64_t source = row * width * 2u + d;
const float gate = gate_up[source];
const float up = gate_up[source + width];
out[i] = deepseek4_vision_round_bf16_dev(
(gate / (1.0f + expf(-gate))) * up);
}
__global__ static void deepseek4_vision_add_residual_kernel(
float *x,
const float *residual,
uint64_t count) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < count)
x[i] = deepseek4_vision_round_bf16_dev(x[i] + residual[i]);
}
__global__ static void deepseek4_vision_aligner_reorder_kernel(
float *out,
const float *x,
uint32_t grid_h,
uint32_t grid_w,
uint32_t output_rows) {
const uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t row = blockIdx.y;
if (d >= 9216u || row >= output_rows) return;
const uint32_t output_w = (grid_w + 2u) / 3u;
const uint32_t block_y = row / output_w;
const uint32_t block_x = row - block_y * output_w;
const uint32_t channel = d / 9u;
const uint32_t within = d - channel * 9u;
const uint32_t source_y = block_y * 3u + within / 3u;
const uint32_t source_x = block_x * 3u + within % 3u;
float value = 0.0f;
if (source_y < grid_h && source_x < grid_w) {
const uint64_t source_row = (uint64_t)source_y * grid_w + source_x;
value = x[source_row * 1024u + channel];
}
out[(uint64_t)row * 9216u + d] = value;
}
__global__ static void deepseek4_vision_gelu_bias_kernel(
float *out,
const float *x,
const uint16_t *bias,
uint64_t count,
uint32_t width) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count) return;
const uint32_t d = i % width;
const float value = deepseek4_vision_round_bf16_dev(
x[i] + glm53_vision_bf16(bias + d));
out[i] = deepseek4_vision_round_bf16_dev(
0.5f * value *
(1.0f + glm53_vision_erf_approx(
value * 0.7071067811865475f)));
}
static int deepseek4_vision_round_tensor(
ds4_gpu_tensor *x,
uint64_t count,
const char *label) {
deepseek4_vision_round_bf16_kernel<<<
(unsigned)((count + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>((float *)x->ptr, count);
return glm53_vision_launch_ok(label);
}
extern "C" int ds4_gpu_deepseek4_vision_encode(
float *out,
const float *patches,
uint32_t grid_h,
uint32_t grid_w,
const void *model_map,
uint64_t model_size,
const ds4_deepseek4_vision_weights *weights) {
if (!out || !patches || !model_map || !weights || grid_h == 0u ||
grid_w == 0u || grid_h > UINT32_MAX / grid_w) return 0;
const uint32_t rows = grid_h * grid_w;
const uint32_t aligned_rows =
((grid_h + 2u) / 3u) * ((grid_w + 2u) / 3u);
const uint64_t row1024 = (uint64_t)rows * 1024u;
const uint64_t row2816 = (uint64_t)rows * 2816u;
const uint64_t row3072 = (uint64_t)rows * 3072u;
const uint64_t row5632 = (uint64_t)rows * 5632u;
const uint64_t aligned4096 = (uint64_t)aligned_rows * 4096u;
const uint64_t aligned9216 = (uint64_t)aligned_rows * 9216u;
if (row5632 > SIZE_MAX / sizeof(float) ||
aligned9216 > SIZE_MAX / sizeof(float)) return 0;
ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL;
ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL;
ds4_gpu_tensor *mlp_w1 = NULL, *mlp_mid = NULL;
ds4_gpu_tensor *align_in = NULL, *align_a = NULL, *align_b = NULL;
ds4_gpu_tensor *cur = NULL, *tmp = NULL;
const uint16_t *bias = NULL;
int ok = 0;
#define DSV4_VISION_ALLOC(name_, count_) do { \
name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \
if (!(name_)) goto cleanup; \
} while (0)
DSV4_VISION_ALLOC(patch, (uint64_t)rows * 588u);
DSV4_VISION_ALLOC(a, row1024);
DSV4_VISION_ALLOC(b, row1024);
DSV4_VISION_ALLOC(qkv, row3072);
DSV4_VISION_ALLOC(q, row1024);
DSV4_VISION_ALLOC(k, row1024);
DSV4_VISION_ALLOC(v, row1024);
DSV4_VISION_ALLOC(attn, row1024);
DSV4_VISION_ALLOC(mlp_w1, row5632);
DSV4_VISION_ALLOC(mlp_mid, row2816);
DSV4_VISION_ALLOC(align_in, aligned9216);
DSV4_VISION_ALLOC(align_a, aligned4096);
DSV4_VISION_ALLOC(align_b, aligned4096);
#undef DSV4_VISION_ALLOC
if (!ds4_gpu_tensor_write(
patch, 0, patches,
(uint64_t)rows * 588u * sizeof(float)) ||
!ds4_gpu_begin_commands()) goto cleanup;
ok = ds4_gpu_glm53_matmul_bf16(
a, model_map, model_size, weights->patch_weight,
588u, 1024u, patch, rows);
if (ok) {
bias = glm53_vision_weight(model_map, model_size,
weights->patch_bias, 1024u, "DeepSeek vision patch bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)a->ptr, bias, NULL, row1024, 1024u);
ok = glm53_vision_launch_ok("DeepSeek vision patch bias");
}
if (ok) ok = deepseek4_vision_round_tensor(
a, row1024, "DeepSeek vision patch round");
cur = a;
tmp = b;
for (uint32_t il = 0; ok && il < DS4_DEEPSEEK4_VISION_LAYERS; il++) {
const ds4_deepseek4_vision_layer_weights *w = &weights->layer[il];
const uint16_t *norm = glm53_vision_weight(
model_map, model_size, w->norm1, 1024u,
"DeepSeek vision norm1");
if (!norm) { ok = 0; break; }
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-6f);
ok = glm53_vision_launch_ok("DeepSeek vision norm1");
if (ok) ok = deepseek4_vision_round_tensor(
tmp, row1024, "DeepSeek vision norm1 round");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
qkv, model_map, model_size, w->qkv_weight,
1024u, 3072u, tmp, rows);
const uint16_t *qkv_bias = NULL;
if (ok) {
qkv_bias = glm53_vision_weight(
model_map, model_size, w->qkv_bias, 3072u,
"DeepSeek vision QKV bias");
if (!qkv_bias) ok = 0;
}
if (ok) {
deepseek4_vision_qkv_rope_kernel<<<
dim3(rows, 16u, 1u), 32u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)q->ptr, (float *)k->ptr, (float *)v->ptr,
(const float *)qkv->ptr, qkv_bias, rows, grid_w);
ok = glm53_vision_launch_ok("DeepSeek vision QKV");
}
if (ok) {
glm53_vision_attention_kernel<<<
dim3(rows, 16u, 1u), 32u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)attn->ptr, (const float *)q->ptr,
(const float *)k->ptr, (const float *)v->ptr, rows);
ok = glm53_vision_launch_ok("DeepSeek vision attention");
}
if (ok) ok = deepseek4_vision_round_tensor(
attn, row1024, "DeepSeek vision attention round");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->attn_proj_weight,
1024u, 1024u, attn, rows);
if (ok) {
bias = glm53_vision_weight(
model_map, model_size, w->attn_proj_bias, 1024u,
"DeepSeek vision attention bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)tmp->ptr, bias, (const float *)cur->ptr,
row1024, 1024u);
ok = glm53_vision_launch_ok(
"DeepSeek vision attention residual");
}
if (ok) ok = deepseek4_vision_round_tensor(
tmp, row1024, "DeepSeek vision attention residual round");
ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap;
if (ok) {
norm = glm53_vision_weight(
model_map, model_size, w->norm2, 1024u,
"DeepSeek vision norm2");
if (!norm) ok = 0;
}
if (ok) {
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-6f);
ok = glm53_vision_launch_ok("DeepSeek vision norm2");
}
if (ok) ok = deepseek4_vision_round_tensor(
tmp, row1024, "DeepSeek vision norm2 round");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
mlp_w1, model_map, model_size, w->mlp_w1,
1024u, 5632u, tmp, rows);
if (ok) ok = deepseek4_vision_round_tensor(
mlp_w1, row5632, "DeepSeek vision MLP input round");
if (ok) {
deepseek4_vision_swiglu_split_kernel<<<
(unsigned)((row2816 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)mlp_mid->ptr, (const float *)mlp_w1->ptr,
row2816, 2816u);
ok = glm53_vision_launch_ok("DeepSeek vision SwiGLU");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->mlp_w2,
2816u, 1024u, mlp_mid, rows);
if (ok) ok = deepseek4_vision_round_tensor(
tmp, row1024, "DeepSeek vision MLP output round");
if (ok) {
deepseek4_vision_add_residual_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr, row1024);
ok = glm53_vision_launch_ok("DeepSeek vision MLP residual");
}
swap = cur; cur = tmp; tmp = swap;
}
if (ok) {
const uint16_t *norm = glm53_vision_weight(
model_map, model_size, weights->post_norm, 1024u,
"DeepSeek vision final norm");
if (!norm) ok = 0;
else {
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-6f);
ok = glm53_vision_launch_ok("DeepSeek vision final norm");
}
}
if (ok) ok = deepseek4_vision_round_tensor(
tmp, row1024, "DeepSeek vision final norm round");
if (ok) {
dim3 grid((9216u + 255u) / 256u, aligned_rows, 1u);
deepseek4_vision_aligner_reorder_kernel<<<
grid, 256u, 0, DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)align_in->ptr, (const float *)tmp->ptr,
grid_h, grid_w, aligned_rows);
ok = glm53_vision_launch_ok("DeepSeek vision aligner reorder");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
align_a, model_map, model_size, weights->aligner_w1,
9216u, 4096u, align_in, aligned_rows);
if (ok) {
bias = glm53_vision_weight(
model_map, model_size, weights->aligner_w1_bias, 4096u,
"DeepSeek vision aligner hidden bias");
if (!bias) ok = 0;
}
if (ok) {
deepseek4_vision_gelu_bias_kernel<<<
(unsigned)((aligned4096 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)align_b->ptr, (const float *)align_a->ptr,
bias, aligned4096, 4096u);
ok = glm53_vision_launch_ok("DeepSeek vision aligner GELU");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
align_a, model_map, model_size, weights->aligner_w2,
4096u, 4096u, align_b, aligned_rows);
if (ok) {
bias = glm53_vision_weight(
model_map, model_size, weights->aligner_w2_bias, 4096u,
"DeepSeek vision aligner output bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((aligned4096 + 255u) / 256u), 256u, 0,
DS4_DEEPSEEK4_VISION_STREAM>>>(
(float *)align_a->ptr, bias, NULL, aligned4096, 4096u);
ok = glm53_vision_launch_ok(
"DeepSeek vision aligner output bias");
}
if (ok) ok = deepseek4_vision_round_tensor(
align_a, aligned4096, "DeepSeek vision aligner output round");
if (ds4_gpu_end_commands() == 0) ok = 0;
if (ok) ok = ds4_gpu_tensor_read(
align_a, 0, out, aligned4096 * sizeof(float));
cleanup:
ds4_gpu_tensor_free(align_b);
ds4_gpu_tensor_free(align_a);
ds4_gpu_tensor_free(align_in);
ds4_gpu_tensor_free(mlp_mid);
ds4_gpu_tensor_free(mlp_w1);
ds4_gpu_tensor_free(attn);
ds4_gpu_tensor_free(v);
ds4_gpu_tensor_free(k);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(qkv);
ds4_gpu_tensor_free(b);
ds4_gpu_tensor_free(a);
ds4_gpu_tensor_free(patch);
return ok;
}
#undef DS4_DEEPSEEK4_VISION_STREAM
ds4_distributed.c
+11
View File
@@ -7012,6 +7012,17 @@ static int dist_worker_handle_snapshot_load(
snprintf(err, sizeof(err), "invalid distributed snapshot load header");
}
/* Reject a token_count larger than the worker context up front, before the
* allocation and the read loop below. token_count is only bounded to
* ~UINT32_MAX/4 by the header check above, so without this a peer could
* drive a multi-GB malloc and a full-length socket read before the
* matching-state check further down (which also tests token_count > ctx_size)
* finally rejects the request. */
if (!err[0] && begin.token_count > (uint32_t)state->ctx_size) {
dist_discard_bytes(upstream->fd, body_bytes);
snprintf(err, sizeof(err), "snapshot token count exceeds worker context size");
}
int *tokens = NULL;
if (!err[0]) {
tokens = malloc((size_t)begin.token_count * sizeof(tokens[0]));
ds4_engram.c
+680
View File
@@ -0,0 +1,680 @@
#define _DARWIN_C_SOURCE
#define _POSIX_C_SOURCE 200809L
#include "ds4_engram.h"
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <math.h>
#include <pthread.h>
#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
#ifdef __APPLE__
#include <dispatch/dispatch.h>
#include <pthread/qos.h>
#include <sys/sysctl.h>
#endif
bool ds4_engram_layout_valid(const ds4_engram_layout *l) {
if (!l || !l->token_map || !l->vocab_size ||
!l->compressed_vocab_size || l->compressed_vocab_size > INT32_MAX ||
l->pad_id >= l->compressed_vocab_size) return false;
for (uint32_t i = 0; i < l->vocab_size; i++)
if (l->token_map[i] >= l->compressed_vocab_size) return false;
for (int layer = 0; layer < DS4_ENGRAM_LAYERS; layer++) {
for (int i = 0; i < DS4_ENGRAM_NGRAM; i++) {
uint64_t m = l->multipliers[layer][i];
if (!(m & 1) || m > (uint64_t)INT64_MAX / l->compressed_vocab_size)
return false;
}
uint64_t total = 0;
for (int i = 0; i < DS4_ENGRAM_COLS; i++) {
if (l->primes[layer][i] < 2) return false;
total += l->primes[layer][i];
}
if (total != l->rows[layer]) return false;
}
return true;
}
void ds4_engram_history_reset(ds4_engram_history *h) {
for (int i = 0; i < DS4_ENGRAM_NGRAM - 1; i++) h->tail[i] = DS4_ENGRAM_DEAD;
}
bool ds4_engram_hash(const ds4_engram_layout *l, ds4_engram_history *h,
const int *tokens, const uint8_t *mask, size_t count,
uint32_t *rows) {
if (!l || !h || !l->token_map || (count && (!tokens || !rows)) ||
count > SIZE_MAX / (DS4_ENGRAM_LAYERS * DS4_ENGRAM_COLS * sizeof(*rows)))
return false;
for (int i = 0; i < DS4_ENGRAM_NGRAM - 1; i++) {
if (h->tail[i] < DS4_ENGRAM_DEAD ||
(h->tail[i] >= 0 && (uint32_t)h->tail[i] >= l->compressed_vocab_size))
return false;
}
for (size_t i = 0; i < count; i++) {
if (tokens[i] < 0 || (uint32_t)tokens[i] >= l->vocab_size) return false;
}
for (size_t i = 0; i < count; i++) {
int32_t current = mask && !mask[i] ? DS4_ENGRAM_DEAD :
(int32_t)l->token_map[tokens[i]];
uint32_t ids[DS4_ENGRAM_NGRAM];
bool blocked = false;
for (int j = 0; j < DS4_ENGRAM_NGRAM; j++) {
int32_t id = j ? h->tail[j - 1] : current;
blocked |= id == DS4_ENGRAM_DEAD;
ids[j] = blocked ? l->pad_id : (uint32_t)id;
}
for (int layer = 0; layer < DS4_ENGRAM_LAYERS; layer++) {
uint64_t hash = (uint64_t)ids[0] * l->multipliers[layer][0];
uint32_t offset = 0;
for (int j = 1; j < DS4_ENGRAM_NGRAM; j++) {
hash ^= (uint64_t)ids[j] * l->multipliers[layer][j];
for (int head = 0; head < DS4_ENGRAM_HEADS; head++) {
int col = (j - 1) * DS4_ENGRAM_HEADS + head;
uint32_t prime = l->primes[layer][col];
*rows++ = (uint32_t)(hash % prime) + offset;
offset += prime;
}
}
}
for (int j = DS4_ENGRAM_NGRAM - 2; j > 0; j--) h->tail[j] = h->tail[j - 1];
h->tail[0] = current;
}
return true;
}
bool ds4_engram_table_open(ds4_engram_table *t, const char *path,
uint64_t offset, uint32_t rows) {
if (!t) return false;
*t = (ds4_engram_table){.fd = -1};
uint64_t bytes = (uint64_t)rows * DS4_ENGRAM_ROW_BYTES;
if (!path || !rows || offset > INT64_MAX || bytes > INT64_MAX - offset) {
errno = EINVAL;
return false;
}
int fd = open(path, O_RDONLY | O_CLOEXEC);
if (fd < 0) return false;
struct stat st;
if (fstat(fd, &st) != 0) goto fail;
if (!S_ISREG(st.st_mode) || st.st_size < 0 || offset + bytes > (uint64_t)st.st_size) {
errno = EINVAL;
goto fail;
}
#ifdef __APPLE__
if (fcntl(fd, F_NOCACHE, 1) != 0 || fcntl(fd, F_RDAHEAD, 0) != 0) goto fail;
#elif defined(POSIX_FADV_RANDOM)
/* Hashed rows are scattered across the table. Read-ahead fetches unused
* pages and competes with the model's limited memory headroom. */
(void)posix_fadvise(fd, (off_t)offset, (off_t)bytes, POSIX_FADV_RANDOM);
#endif
*t = (ds4_engram_table){.fd = fd, .offset = offset, .rows = rows};
return true;
fail: {
int saved = errno;
close(fd);
errno = saved;
return false;
}
}
void ds4_engram_table_close(ds4_engram_table *t) {
if (!t) return;
if (t->fd >= 0) close(t->fd);
*t = (ds4_engram_table){.fd = -1};
}
static bool read_row(int fd, uint64_t offset, uint8_t row[DS4_ENGRAM_ROW_BYTES]) {
size_t done = 0;
while (done < DS4_ENGRAM_ROW_BYTES) {
ssize_t n = pread(fd, row + done, DS4_ENGRAM_ROW_BYTES - done,
(off_t)(offset + done));
if (n < 0 && errno == EINTR) continue;
if (n <= 0) {
if (n == 0) errno = EIO;
return false;
}
done += (size_t)n;
}
return true;
}
static float e4m3(uint8_t byte) {
int exponent = (byte >> 3) & 15, mantissa = byte & 7;
float value = exponent ? ldexpf((float)(8 + mantissa), exponent - 10) :
ldexpf((float)mantissa, -9);
return byte & 128 ? -value : value;
}
/* All code/scale pairs fit in 128 KiB; this is a format lookup, not a cache
* of model rows. Preserve the reference's F32 scaling and BF16 rounding. */
static uint16_t engram_bf16[256][256];
static pthread_once_t engram_bf16_once = PTHREAD_ONCE_INIT;
static void init_engram_bf16(void) {
for (unsigned scale = 0; scale < 256; scale++) {
for (unsigned code = 0; code < 256; code++) {
uint16_t result = UINT16_MAX;
if ((code & 127u) != 127u && scale != 255u) {
float value = ldexpf(e4m3((uint8_t)code), (int)scale - 127);
uint32_t bits;
memcpy(&bits, &value, sizeof(bits));
bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u;
if ((bits & 0x7f800000u) != 0x7f800000u) result = (uint16_t)(bits >> 16);
}
engram_bf16[scale][code] = result;
}
}
}
bool ds4_engram_read(const ds4_engram_table *t, const uint32_t *rows,
size_t count, float *out) {
if (!t || t->fd < 0 || (count && (!rows || !out)) ||
count > SIZE_MAX / (DS4_ENGRAM_DIM * sizeof(*out))) {
errno = EINVAL;
return false;
}
for (size_t i = 0; i < count; i++) {
if (rows[i] >= t->rows) {
errno = EINVAL;
return false;
}
}
pthread_once(&engram_bf16_once, init_engram_bf16);
uint8_t raw[DS4_ENGRAM_ROW_BYTES];
for (size_t i = 0; i < count; i++) {
if (!read_row(t->fd, t->offset + (uint64_t)rows[i] * sizeof(raw), raw)) return false;
for (int j = 0; j < DS4_ENGRAM_DIM; j++) {
const uint16_t value = engram_bf16[raw[DS4_ENGRAM_DIM + j / 32]][raw[j]];
if (value == UINT16_MAX) {
errno = EDOM;
return false;
}
const uint32_t bits = (uint32_t)value << 16;
memcpy(out + i * DS4_ENGRAM_DIM + j, &bits, sizeof(bits));
}
}
return true;
}
typedef struct {
uint32_t row, output;
} engram_request;
static int request_order(const void *a, const void *b) {
const engram_request *x = a, *y = b;
return (x->row > y->row) - (x->row < y->row);
}
enum { ENGRAM_READERS = 16 };
#ifdef __APPLE__
enum { ENGRAM_PARALLEL_MIN_ROWS = 8 };
#else
/* Unlike dispatch's shared pool, this path creates threads for each batch. */
enum { ENGRAM_PARALLEL_MIN_ROWS = 256 };
#endif
typedef struct {
const ds4_engram_table *table;
const engram_request *request;
float *out;
size_t count, readers;
int error[ENGRAM_READERS];
} engram_batch;
static void read_batch_part(void *context, size_t part) {
engram_batch *batch = context;
const engram_request *request = batch->request;
const size_t begin = batch->count * part / batch->readers;
const size_t end = batch->count * (part + 1) / batch->readers;
const float *previous = NULL;
for (size_t i = begin; i < end; i++) {
float *dst = batch->out + (size_t)request[i].output * DS4_ENGRAM_DIM;
if (i > begin && request[i].row == request[i - 1].row) {
memcpy(dst, previous, DS4_ENGRAM_DIM * sizeof(*dst));
} else {
if (!ds4_engram_read(batch->table, &request[i].row, 1, dst)) {
batch->error[part] = errno ? errno : EIO;
return;
}
previous = dst;
}
}
}
#ifndef __APPLE__
typedef struct {
engram_batch *batch;
size_t part;
} engram_reader;
static void *read_batch_thread(void *context) {
engram_reader *reader = context;
read_batch_part(reader->batch, reader->part);
return NULL;
}
#endif
bool ds4_engram_read_batch(const ds4_engram_table *t, const uint32_t *rows,
size_t tokens, size_t stride, float *out) {
if (!t || t->fd < 0 || (tokens && (!rows || !out || stride < DS4_ENGRAM_COLS)) ||
tokens > SIZE_MAX / (DS4_ENGRAM_COLS * DS4_ENGRAM_DIM * sizeof(*out)) ||
(tokens && tokens - 1 > (SIZE_MAX / sizeof(*rows) - DS4_ENGRAM_COLS) / stride)) {
errno = EINVAL;
return false;
}
for (size_t i = 0; i < tokens; i++) {
for (size_t j = 0; j < DS4_ENGRAM_COLS; j++) {
if (rows[i * stride + j] >= t->rows) {
errno = EINVAL;
return false;
}
}
}
if (!tokens) return true;
enum { BATCH_TOKENS = 2048 };
const size_t cap = tokens < BATCH_TOKENS ? tokens : BATCH_TOKENS;
engram_request *request = malloc(cap * DS4_ENGRAM_COLS * sizeof(*request));
if (!request) return false;
bool ok = true;
for (size_t start = 0; ok && start < tokens; start += cap) {
const size_t n = tokens - start < cap ? tokens - start : cap;
const size_t count = n * DS4_ENGRAM_COLS;
for (size_t i = 0; i < count; i++) {
request[i] = (engram_request){
rows[(start + i / DS4_ENGRAM_COLS) * stride + i % DS4_ENGRAM_COLS],
(uint32_t)i
};
}
qsort(request, count, sizeof(*request), request_order);
engram_batch batch = {.table = t, .request = request, .count = count,
.out = out + start * DS4_ENGRAM_COLS * DS4_ENGRAM_DIM, .readers = 1};
/* Fixed concurrency hides random-read latency without caching the table.
* Each worker owns disjoint output rows; all finish before GPU use. */
if (count >= ENGRAM_PARALLEL_MIN_ROWS) {
batch.readers = count < ENGRAM_READERS ? count : ENGRAM_READERS;
#ifdef __APPLE__
dispatch_apply_f(batch.readers,
dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), &batch, read_batch_part);
#else
pthread_t threads[ENGRAM_READERS - 1];
engram_reader readers[ENGRAM_READERS - 1];
size_t started = 0;
for (size_t part = 1; part < batch.readers; part++) {
readers[started] = (engram_reader){&batch, part};
if (pthread_create(&threads[started], NULL, read_batch_thread,
&readers[started])) break;
started++;
}
read_batch_part(&batch, 0);
/* Thread exhaustion only reduces concurrency, not correctness. */
for (size_t part = started + 1; part < batch.readers; part++)
read_batch_part(&batch, part);
for (size_t part = 0; part < started; part++)
if (pthread_join(threads[part], NULL)) abort();
#endif
} else
read_batch_part(&batch, 0);
for (size_t i = 0; i < batch.readers; i++) {
if (batch.error[i]) {
errno = batch.error[i];
ok = false;
break;
}
}
}
int saved = errno;
free(request);
errno = saved;
return ok;
}
struct ds4_engram_prefetch {
const ds4_engram_table *table;
const uint32_t *rows;
size_t tokens, stride;
float *out;
bool ok;
int error;
#ifdef __APPLE__
dispatch_group_t group;
#else
pthread_t thread;
bool threaded;
#endif
};
static void prefetch_run(void *context) {
ds4_engram_prefetch *p = context;
p->ok = ds4_engram_read_batch(p->table, p->rows, p->tokens, p->stride, p->out);
if (!p->ok) p->error = errno ? errno : EIO;
}
#ifndef __APPLE__
static void *prefetch_thread(void *context) {
prefetch_run(context);
return NULL;
}
#endif
ds4_engram_prefetch *ds4_engram_read_batch_start(const ds4_engram_table *table, const uint32_t *rows,
size_t tokens, size_t stride, float *out) {
ds4_engram_prefetch *p = calloc(1, sizeof(*p));
if (!p) return NULL;
*p = (ds4_engram_prefetch){.table = table, .rows = rows, .tokens = tokens, .stride = stride, .out = out};
#ifdef __APPLE__
p->group = dispatch_group_create();
dispatch_group_async_f(p->group, dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), p, prefetch_run);
#else
p->threaded = pthread_create(&p->thread, NULL, prefetch_thread, p) == 0;
if (!p->threaded) prefetch_run(p);
#endif
return p;
}
/* Page-cache warmer. Sequential reads fill the kernel's cache with the table
* once; the uncached row readers then hit it instead of the device. Spread
* prefill rows are the worst access pattern for external SSDs: a USB drive
* served ~6K random rows/s over the 189 GiB span, ~12 s per 2K prefill, but
* streams the whole span at ~1 GB/s. Clean cached pages stay reclaimable, so
* the warmer is a hint, never a reservation: it stops at memory pressure. */
enum { WARM_THREADS = 4 };
#define WARM_CHUNK ((uint64_t)8 << 20)
struct ds4_engram_warm {
int fd;
uint64_t begin, end, next;
_Atomic bool stop;
bool done, ok;
double started, seconds;
pthread_mutex_t lock;
pthread_cond_t cond;
unsigned running;
pthread_t thread[WARM_THREADS];
unsigned threads;
};
static double warm_now(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + ts.tv_nsec * 1e-9;
}
static bool warm_pressure(void) {
#ifdef __APPLE__
int level = 0;
size_t len = sizeof(level);
/* 1 normal, 2 warn, 4 critical. */
return sysctlbyname("kern.memorystatus_vm_pressure_level", &level, &len, NULL, 0) == 0 &&
level > 1;
#else
return false;
#endif
}
static void *warm_thread(void *arg) {
ds4_engram_warm *w = arg;
#ifdef __APPLE__
pthread_set_qos_class_self_np(QOS_CLASS_UTILITY, 0);
#endif
uint8_t *buf = malloc(WARM_CHUNK);
bool ok = buf != NULL;
for (unsigned n = 0; ok && !w->stop; n++) {
const uint64_t off = __atomic_fetch_add(&w->next, WARM_CHUNK, __ATOMIC_RELAXED);
if (off >= w->end) break;
if (!(n & 63) && warm_pressure()) {
w->stop = true;
ok = false;
break;
}
size_t left = (size_t)(w->end - off < WARM_CHUNK ? w->end - off : WARM_CHUNK);
const uint64_t end = off + left;
while (left) {
/* A short read still warms those pages; only the file offset advances. */
const ssize_t got = pread(w->fd, buf, left, (off_t)(end - left));
if (got < 0 && errno == EINTR) continue;
if (got <= 0 || (size_t)got > left) {
ok = false;
break;
}
left -= (size_t)got;
}
}
free(buf);
pthread_mutex_lock(&w->lock);
if (!ok) w->ok = false;
if (--w->running == 0) {
w->done = true;
w->seconds = warm_now() - w->started;
pthread_cond_broadcast(&w->cond);
}
pthread_mutex_unlock(&w->lock);
return NULL;
}
uint64_t ds4_engram_physical_memory(void) {
#ifdef __APPLE__
uint64_t bytes = 0;
size_t len = sizeof(bytes);
return sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) == 0 ? bytes : 0;
#else
const long pages = sysconf(_SC_PHYS_PAGES), page = sysconf(_SC_PAGESIZE);
return pages > 0 && page > 0 ? (uint64_t)pages * (uint64_t)page : 0;
#endif
}
ds4_engram_warm *ds4_engram_warm_start(const char *path, uint64_t offset, uint64_t bytes) {
if (!path || !bytes || offset > INT64_MAX || bytes > INT64_MAX - offset) return NULL;
ds4_engram_warm *w = calloc(1, sizeof(*w));
if (!w) return NULL;
w->fd = open(path, O_RDONLY | O_CLOEXEC);
if (w->fd < 0) {
free(w);
return NULL;
}
w->begin = w->next = offset;
w->end = offset + bytes;
w->ok = true;
w->started = warm_now();
pthread_mutex_init(&w->lock, NULL);
pthread_cond_init(&w->cond, NULL);
for (unsigned i = 0; i < WARM_THREADS; i++) {
pthread_mutex_lock(&w->lock);
w->running++;
pthread_mutex_unlock(&w->lock);
if (pthread_create(&w->thread[i], NULL, warm_thread, w)) {
pthread_mutex_lock(&w->lock);
w->running--;
pthread_mutex_unlock(&w->lock);
break;
}
w->threads++;
}
if (!w->threads) {
ds4_engram_warm_stop(w);
return NULL;
}
return w;
}
bool ds4_engram_warm_wait(ds4_engram_warm *w, double *seconds) {
if (!w) return false;
pthread_mutex_lock(&w->lock);
while (!w->done) pthread_cond_wait(&w->cond, &w->lock);
const bool ok = w->ok;
if (seconds) *seconds = w->seconds;
pthread_mutex_unlock(&w->lock);
return ok;
}
void ds4_engram_warm_stop(ds4_engram_warm *w) {
if (!w) return;
w->stop = true;
for (unsigned i = 0; i < w->threads; i++)
if (pthread_join(w->thread[i], NULL)) abort();
pthread_mutex_destroy(&w->lock);
pthread_cond_destroy(&w->cond);
close(w->fd);
free(w);
}
bool ds4_engram_read_batch_finish(ds4_engram_prefetch *p) {
if (!p) return false;
#ifdef __APPLE__
dispatch_group_wait(p->group, DISPATCH_TIME_FOREVER);
dispatch_release(p->group);
#else
if (p->threaded && pthread_join(p->thread, NULL)) abort();
#endif
const bool ok = p->ok;
if (!ok) errno = p->error;
free(p);
return ok;
}
enum { TOKEN_READERS = 4 };
typedef struct {
pthread_t thread;
ds4_engram_reader *reader;
unsigned index;
} engram_read_worker;
struct ds4_engram_reader {
engram_read_worker workers[TOKEN_READERS];
unsigned threads;
pthread_mutex_t mutex;
pthread_cond_t work, done;
bool active, shutdown;
uint64_t generation;
unsigned completed[2];
int errors[2];
ds4_engram_table tables[2];
uint32_t ids[2][DS4_ENGRAM_COLS];
float *out;
};
static void *read_token_thread(void *context) {
const engram_read_worker *worker = context;
ds4_engram_reader *r = worker->reader;
const unsigned index = worker->index;
uint64_t seen = 0;
pthread_mutex_lock(&r->mutex);
for (;;) {
while (r->generation == seen && !r->shutdown) pthread_cond_wait(&r->work, &r->mutex);
if (r->shutdown) break;
seen = r->generation;
pthread_mutex_unlock(&r->mutex);
for (unsigned table = 0; table < 2; table++) {
const unsigned first = index * DS4_ENGRAM_COLS / TOKEN_READERS;
const unsigned end = (index + 1u) * DS4_ENGRAM_COLS / TOKEN_READERS;
const bool ok = ds4_engram_read(&r->tables[table], r->ids[table] + first,
end - first, r->out + (table * DS4_ENGRAM_COLS + first) * DS4_ENGRAM_DIM);
const int error = ok ? 0 : errno ? errno : EIO;
pthread_mutex_lock(&r->mutex);
if (!r->errors[table]) r->errors[table] = error;
if (++r->completed[table] == TOKEN_READERS) {
if (table == 1) r->active = false;
pthread_cond_broadcast(&r->done);
}
pthread_mutex_unlock(&r->mutex);
}
pthread_mutex_lock(&r->mutex);
}
pthread_mutex_unlock(&r->mutex);
return NULL;
}
ds4_engram_reader *ds4_engram_reader_create(void) {
ds4_engram_reader *r = calloc(1, sizeof(*r));
if (!r) return NULL;
int error = pthread_mutex_init(&r->mutex, NULL);
if (error) goto fail;
error = pthread_cond_init(&r->work, NULL);
if (error) goto mutex;
error = pthread_cond_init(&r->done, NULL);
if (error) goto work;
for (unsigned i = 0; i < TOKEN_READERS; i++) {
r->workers[i].reader = r;
r->workers[i].index = i;
error = pthread_create(&r->workers[i].thread, NULL, read_token_thread, &r->workers[i]);
if (error) break;
r->threads++;
}
if (!error) return r;
pthread_mutex_lock(&r->mutex);
r->shutdown = true;
pthread_cond_broadcast(&r->work);
pthread_mutex_unlock(&r->mutex);
for (unsigned i = 0; i < r->threads; i++)
if (pthread_join(r->workers[i].thread, NULL)) abort();
pthread_cond_destroy(&r->done);
work:
pthread_cond_destroy(&r->work);
mutex:
pthread_mutex_destroy(&r->mutex);
fail:
free(r);
errno = error;
return NULL;
}
bool ds4_engram_reader_start(ds4_engram_reader *r,
const ds4_engram_table tables[2], const uint32_t ids[2][DS4_ENGRAM_COLS], float *out) {
if (!r || !tables || !ids || !out) { errno = EINVAL; return false; }
for (unsigned t = 0; t < 2; t++) {
if (tables[t].fd < 0) { errno = EINVAL; return false; }
for (unsigned col = 0; col < DS4_ENGRAM_COLS; col++)
if (ids[t][col] >= tables[t].rows) { errno = EINVAL; return false; }
}
pthread_mutex_lock(&r->mutex);
if (r->active || r->shutdown) {
pthread_mutex_unlock(&r->mutex);
errno = EBUSY;
return false;
}
memcpy(r->tables, tables, sizeof(r->tables));
memcpy(r->ids, ids, sizeof(r->ids));
memset(r->errors, 0, sizeof(r->errors));
r->out = out;
memset(r->completed, 0, sizeof(r->completed));
r->generation++;
r->active = true;
pthread_cond_broadcast(&r->work);
pthread_mutex_unlock(&r->mutex);
return true;
}
bool ds4_engram_reader_wait(ds4_engram_reader *r, unsigned table) {
if (!r || table >= 2) { errno = EINVAL; return false; }
pthread_mutex_lock(&r->mutex);
if (!r->active && !r->completed[1]) {
pthread_mutex_unlock(&r->mutex);
errno = EINVAL;
return false;
}
while (r->completed[table] < TOKEN_READERS) pthread_cond_wait(&r->done, &r->mutex);
const int error = r->errors[table];
pthread_mutex_unlock(&r->mutex);
if (error) errno = error;
return !error;
}
void ds4_engram_reader_free(ds4_engram_reader *r) {
if (!r) return;
pthread_mutex_lock(&r->mutex);
r->shutdown = true;
pthread_cond_broadcast(&r->work);
pthread_mutex_unlock(&r->mutex);
for (unsigned i = 0; i < r->threads; i++)
if (pthread_join(r->workers[i].thread, NULL)) abort();
pthread_cond_destroy(&r->done);
pthread_cond_destroy(&r->work);
pthread_mutex_destroy(&r->mutex);
free(r);
}
ds4_engram.h
+90
View File
@@ -0,0 +1,90 @@
#ifndef DS4_ENGRAM_H
#define DS4_ENGRAM_H
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
enum {
DS4_ENGRAM_LAYERS = 2,
DS4_ENGRAM_NGRAM = 4,
DS4_ENGRAM_HEADS = 8,
DS4_ENGRAM_COLS = 24,
DS4_ENGRAM_DIM = 256,
DS4_ENGRAM_ROW_BYTES = 264,
DS4_ENGRAM_DEAD = -1
};
typedef struct {
const uint32_t *token_map;
uint32_t vocab_size, compressed_vocab_size, pad_id;
uint32_t rows[DS4_ENGRAM_LAYERS];
uint64_t multipliers[DS4_ENGRAM_LAYERS][DS4_ENGRAM_NGRAM];
uint32_t primes[DS4_ENGRAM_LAYERS][DS4_ENGRAM_COLS];
} ds4_engram_layout;
/* Newest first; copying this value also snapshots the complete hash state. */
typedef struct {
int32_t tail[DS4_ENGRAM_NGRAM - 1];
} ds4_engram_history;
bool ds4_engram_layout_valid(const ds4_engram_layout *layout);
void ds4_engram_history_reset(ds4_engram_history *history);
/* Validate the layout once at model load. Mask zero breaks all n-grams that
* cross that position. Output is [token][layer][column], including masked rows;
* the graph must suppress Engram at masked positions, as in the reference. */
bool ds4_engram_hash(const ds4_engram_layout *layout,
ds4_engram_history *history,
const int *tokens, const uint8_t *mask, size_t count,
uint32_t *rows);
typedef struct {
int fd;
uint64_t offset;
uint32_t rows;
} ds4_engram_table;
/* A separate uncached file descriptor, never an mmap or Metal model view.
* Each GGUF I8 row is 256 E4M3 bytes followed by 8 original E8M0 scales. */
bool ds4_engram_table_open(ds4_engram_table *table, const char *path,
uint64_t offset, uint32_t rows);
void ds4_engram_table_close(ds4_engram_table *table);
/* Output uses F32 storage for the reference's BF16-rounded values. No whole
* table allocation; caller owns count * DIM floats. Failure invalidates output. */
bool ds4_engram_read(const ds4_engram_table *table, const uint32_t *rows,
size_t count, float *out);
/* Read COLS rows per token, restoring token order after deduplicated disk reads.
* Input stride is in row IDs; output is packed [token][COLS][DIM]. Temporary
* storage is bounded to 384 KiB, independent of the table and prefix size.
* On macOS, large batches use bounded concurrent pread readers. */
bool ds4_engram_read_batch(const ds4_engram_table *table, const uint32_t *rows,
size_t tokens, size_t stride, float *out);
/* The same read on a worker: `rows` and `out` stay the caller's until finish,
* which joins the worker and returns its result. */
typedef struct ds4_engram_prefetch ds4_engram_prefetch;
ds4_engram_prefetch *ds4_engram_read_batch_start(const ds4_engram_table *table, const uint32_t *rows,
size_t tokens, size_t stride, float *out);
bool ds4_engram_read_batch_finish(ds4_engram_prefetch *prefetch);
/* One reusable reader for a token's two tables. Tables/IDs are copied at
* start; descriptors and output must remain valid until wait(reader, 1) or
* free, even if wait(reader, 0) reports an error. No table data is cached.
* One caller owns the reader. Creation failure permits synchronous fallback. */
typedef struct ds4_engram_reader ds4_engram_reader;
ds4_engram_reader *ds4_engram_reader_create(void);
bool ds4_engram_reader_start(ds4_engram_reader *reader,
const ds4_engram_table tables[2], const uint32_t ids[2][DS4_ENGRAM_COLS], float *out);
bool ds4_engram_reader_wait(ds4_engram_reader *reader, unsigned table);
void ds4_engram_reader_free(ds4_engram_reader *reader);
/* Background sequential read of [offset, offset + bytes) into the kernel page
* cache, so the uncached row readers above hit memory. No table allocation;
* stops early at memory pressure. Wait reports whether it covered the span. */
typedef struct ds4_engram_warm ds4_engram_warm;
uint64_t ds4_engram_physical_memory(void);
ds4_engram_warm *ds4_engram_warm_start(const char *path, uint64_t offset, uint64_t bytes);
bool ds4_engram_warm_wait(ds4_engram_warm *warm, double *seconds);
void ds4_engram_warm_stop(ds4_engram_warm *warm);
#endif
ds4_eval.c
+788 -102
View File
File diff suppressed because it is too large. Load diff
ds4_eval_cases.c
+1118
View File
File diff suppressed because it is too large. Load diff
ds4_eval_cases.h
+51
View File
@@ -0,0 +1,51 @@
#ifndef DS4_EVAL_CASES_H
#define DS4_EVAL_CASES_H
#include <stddef.h>
#include <stdint.h>
#define EVAL_MAX_CHOICES 10
#define EVAL_MAX_ALIASES 4
typedef enum {
EVAL_ANSWER_AUTO = 0,
EVAL_ANSWER_CHOICE,
EVAL_ANSWER_INTEGER,
EVAL_ANSWER_RATIONAL,
EVAL_ANSWER_EXACT_TEXT,
EVAL_ANSWER_ORDERED_SEQUENCE,
EVAL_ANSWER_LINE_SET,
} eval_answer_kind;
enum {
EVAL_SUITE_CORE = 1u << 0,
EVAL_SUITE_HARD = 1u << 1,
EVAL_SUITE_HARD_SMOKE = 1u << 2,
};
typedef struct {
const char *source;
const char *id;
const char *domain;
const char *title;
const char *question;
const char *choice[EVAL_MAX_CHOICES];
const char *answer;
const char *alias[EVAL_MAX_ALIASES];
eval_answer_kind answer_kind;
uint32_t suites;
uint32_t max_tokens;
} eval_case;
typedef struct {
const char *name;
const char *url;
const char *license;
} eval_source;
extern const eval_case eval_hard_cases[];
extern const size_t eval_hard_case_count;
extern const eval_source eval_hard_sources[];
extern const size_t eval_hard_source_count;
#endif
ds4_glm53_vision_gpu.cuh
+607
View File
@@ -0,0 +1,607 @@
/* GLM-5.3 vision operations shared by CUDA and ROCm. */
#ifndef DS4_GLM53_VISION_TYPES_DEFINED
#define DS4_GLM53_VISION_TYPES_DEFINED
#define DS4_GLM53_VISION_LAYERS 24u
typedef struct {
uint64_t norm1;
uint64_t qkv_weight;
uint64_t qkv_bias;
uint64_t q_norm;
uint64_t k_norm;
uint64_t attn_proj_weight;
uint64_t attn_proj_bias;
uint64_t norm2;
uint64_t gate_weight;
uint64_t gate_bias;
uint64_t up_weight;
uint64_t up_bias;
uint64_t down_weight;
uint64_t down_bias;
} ds4_glm53_vision_layer_weights;
typedef struct {
uint64_t patch_weight;
uint64_t patch_bias;
uint64_t post_norm;
uint64_t downsample_weight;
uint64_t downsample_bias;
uint64_t merger_proj;
uint64_t merger_norm;
uint64_t merger_norm_bias;
uint64_t merger_gate;
uint64_t merger_up;
uint64_t merger_down;
ds4_glm53_vision_layer_weights layer[DS4_GLM53_VISION_LAYERS];
} ds4_glm53_vision_weights;
#endif
#ifndef DS4_GLM53_VISION_STREAM
#define DS4_GLM53_VISION_STREAM 0
#endif
__device__ __forceinline__ static float glm53_vision_bf16(
const uint16_t *p) {
return __uint_as_float((uint32_t)(*p) << 16);
}
__device__ __forceinline__ static float glm53_vision_erf_approx(float x) {
const float sign = x < 0.0f ? -1.0f : 1.0f;
const float a = fabsf(x);
const float t = 1.0f / (1.0f + 0.3275911f * a);
const float p = (((((1.061405429f * t - 1.453152027f) * t) +
1.421413741f) * t - 0.284496736f) * t +
0.254829592f) * t;
return sign * (1.0f - p * expf(-a * a));
}
__global__ static void glm53_vision_bias_kernel(
float *x,
const uint16_t *bias,
const float *residual,
uint64_t count,
uint32_t width) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count) return;
float v = x[i] + glm53_vision_bf16(bias + i % width);
if (residual) v += residual[i];
x[i] = v;
}
__global__ static void glm53_vision_rms_kernel(
float *out,
const float *x,
const uint16_t *weight,
uint32_t width,
float eps) {
__shared__ float partial[256];
const uint32_t row = blockIdx.x;
const uint32_t tid = threadIdx.x;
const float *xr = x + (uint64_t)row * width;
float *yr = out + (uint64_t)row * width;
float sum = 0.0f;
for (uint32_t d = tid; d < width; d += blockDim.x) {
sum = fmaf(xr[d], xr[d], sum);
}
partial[tid] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
if (tid < stride) partial[tid] += partial[tid + stride];
__syncthreads();
}
const float inv = rsqrtf(partial[0] / (float)width + eps);
for (uint32_t d = tid; d < width; d += blockDim.x) {
yr[d] = xr[d] * inv * glm53_vision_bf16(weight + d);
}
}
__global__ static void glm53_vision_qkv_rope_kernel(
float *q,
float *k,
float *v,
const float *qkv,
const uint16_t *bias,
const uint16_t *q_weight,
const uint16_t *k_weight,
uint32_t rows,
uint32_t grid_w,
float eps) {
__shared__ float qsum[32];
__shared__ float ksum[32];
const uint32_t row = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t lane = threadIdx.x;
if (row >= rows || head >= 16u || lane >= 32u) return;
const uint64_t qkv_base = (uint64_t)row * 3072u +
(uint64_t)head * 64u;
const uint64_t out_base = (uint64_t)row * 1024u +
(uint64_t)head * 64u;
float q0 = qkv[qkv_base + lane] +
glm53_vision_bf16(bias + (uint64_t)head * 64u + lane);
float q1 = qkv[qkv_base + lane + 32u] +
glm53_vision_bf16(bias + (uint64_t)head * 64u + lane + 32u);
float k0 = qkv[qkv_base + 1024u + lane] +
glm53_vision_bf16(bias + 1024u +
(uint64_t)head * 64u + lane);
float k1 = qkv[qkv_base + 1024u + lane + 32u] +
glm53_vision_bf16(bias + 1024u +
(uint64_t)head * 64u + lane + 32u);
qsum[lane] = fmaf(q0, q0, q1 * q1);
ksum[lane] = fmaf(k0, k0, k1 * k1);
__syncthreads();
for (uint32_t stride = 16u; stride != 0u; stride >>= 1u) {
if (lane < stride) {
qsum[lane] += qsum[lane + stride];
ksum[lane] += ksum[lane + stride];
}
__syncthreads();
}
const float qinv = rsqrtf(qsum[0] / 64.0f + eps);
const float kinv = rsqrtf(ksum[0] / 64.0f + eps);
q0 *= qinv * glm53_vision_bf16(q_weight + lane);
q1 *= qinv * glm53_vision_bf16(q_weight + lane + 32u);
k0 *= kinv * glm53_vision_bf16(k_weight + lane);
k1 *= kinv * glm53_vision_bf16(k_weight + lane + 32u);
const uint32_t merge_w = grid_w / 2u;
const uint32_t group_index = row / 4u;
const uint32_t within = row & 3u;
const uint32_t py = (group_index / merge_w) * 2u + within / 2u;
const uint32_t px = (group_index % merge_w) * 2u + within % 2u;
const uint32_t freq_index = lane & 15u;
const uint32_t pos = lane < 16u ? py : px;
const float inv_freq = powf(10000.0f, -(float)freq_index / 16.0f);
const float angle = (float)pos * inv_freq;
const float cs = cosf(angle);
const float sn = sinf(angle);
q[out_base + lane] = q0 * cs - q1 * sn;
q[out_base + lane + 32u] = q1 * cs + q0 * sn;
k[out_base + lane] = k0 * cs - k1 * sn;
k[out_base + lane + 32u] = k1 * cs + k0 * sn;
v[out_base + lane] = qkv[qkv_base + 2048u + lane] +
glm53_vision_bf16(bias + 2048u +
(uint64_t)head * 64u + lane);
v[out_base + lane + 32u] = qkv[qkv_base + 2048u + lane + 32u] +
glm53_vision_bf16(
bias + 2048u +
(uint64_t)head * 64u + lane + 32u);
}
__global__ static void glm53_vision_attention_kernel(
float *out,
const float *q,
const float *k,
const float *v,
uint32_t rows) {
__shared__ float dot[32];
const uint32_t row = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t lane = threadIdx.x;
if (row >= rows || head >= 16u || lane >= 32u) return;
const uint64_t base = (uint64_t)row * 1024u +
(uint64_t)head * 64u;
const float q0 = q[base + lane];
const float q1 = q[base + lane + 32u];
float acc0 = 0.0f;
float acc1 = 0.0f;
float max_score = -INFINITY;
float denom = 0.0f;
for (uint32_t key_row = 0; key_row < rows; key_row++) {
const uint64_t kb = (uint64_t)key_row * 1024u +
(uint64_t)head * 64u;
dot[lane] = q0 * k[kb + lane] + q1 * k[kb + lane + 32u];
__syncthreads();
for (uint32_t stride = 16u; stride != 0u; stride >>= 1u) {
if (lane < stride) dot[lane] += dot[lane + stride];
__syncthreads();
}
const float score = dot[0] * 0.125f;
const float next_max = fmaxf(max_score, score);
const float old_scale = key_row == 0u ? 0.0f :
expf(max_score - next_max);
const float new_scale = expf(score - next_max);
denom = denom * old_scale + new_scale;
acc0 = acc0 * old_scale + new_scale * v[kb + lane];
acc1 = acc1 * old_scale + new_scale * v[kb + lane + 32u];
max_score = next_max;
__syncthreads();
}
out[base + lane] = acc0 / denom;
out[base + lane + 32u] = acc1 / denom;
}
__global__ static void glm53_vision_swiglu_bias_kernel(
float *out,
const float *gate,
const uint16_t *gate_bias,
const float *up,
const uint16_t *up_bias,
uint64_t count,
uint32_t width) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count) return;
const uint32_t d = (uint32_t)(i % width);
const float g = fminf(gate[i] + glm53_vision_bf16(gate_bias + d), 10.0f);
const float u = fminf(fmaxf(up[i] + glm53_vision_bf16(up_bias + d),
-10.0f), 10.0f);
out[i] = (g / (1.0f + expf(-g))) * u;
}
__global__ static void glm53_vision_reorder_kernel(
float *out,
const float *x,
uint64_t count) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count) return;
const uint32_t d = (uint32_t)(i % 4096u);
const uint64_t row = i / 4096u;
const uint32_t channel = d / 4u;
const uint32_t within = d & 3u;
out[i] = x[(row * 4u + within) * 1024u + channel];
}
__global__ static void glm53_vision_layernorm_gelu_kernel(
float *out,
const float *x,
const uint16_t *weight,
const uint16_t *bias,
uint32_t width,
float eps) {
__shared__ float partial[256];
const uint32_t row = blockIdx.x;
const uint32_t tid = threadIdx.x;
const float *xr = x + (uint64_t)row * width;
float *yr = out + (uint64_t)row * width;
float sum = 0.0f;
for (uint32_t d = tid; d < width; d += blockDim.x) sum += xr[d];
partial[tid] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
if (tid < stride) partial[tid] += partial[tid + stride];
__syncthreads();
}
const float mean = partial[0] / (float)width;
__syncthreads();
float var = 0.0f;
for (uint32_t d = tid; d < width; d += blockDim.x) {
const float centered = xr[d] - mean;
var = fmaf(centered, centered, var);
}
partial[tid] = var;
__syncthreads();
for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
if (tid < stride) partial[tid] += partial[tid + stride];
__syncthreads();
}
const float inv = rsqrtf(partial[0] / (float)width + eps);
const float inv_sqrt2 = 0.7071067811865475f;
for (uint32_t d = tid; d < width; d += blockDim.x) {
float value = (xr[d] - mean) * inv * glm53_vision_bf16(weight + d) +
glm53_vision_bf16(bias + d);
yr[d] = 0.5f * value *
(1.0f + glm53_vision_erf_approx(value * inv_sqrt2));
}
}
__global__ static void glm53_vision_scatter_hc_kernel(
float *hc,
const float *image,
uint32_t dst_row,
uint32_t image_row,
uint32_t rows,
uint32_t total_rows,
uint32_t width,
uint32_t n_hc) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t count = (uint64_t)rows * n_hc * width;
if (i >= count) return;
const uint32_t d = (uint32_t)(i % width);
const uint64_t linear_row = i / width;
const uint32_t image_delta = (uint32_t)(linear_row / n_hc);
const uint32_t hc_index = (uint32_t)(linear_row % n_hc);
if (dst_row + image_delta >= total_rows) return;
const uint64_t dst = ((uint64_t)(dst_row + image_delta) * n_hc +
hc_index) * width + d;
const uint64_t src = (uint64_t)(image_row + image_delta) * width + d;
hc[dst] = image[src];
}
static const uint16_t *glm53_vision_weight(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t elements,
const char *label) {
if (!model_map || elements > UINT64_MAX / sizeof(uint16_t) ||
offset > model_size) return NULL;
const uint64_t bytes = elements * sizeof(uint16_t);
if (bytes > model_size - offset) return NULL;
#if defined(__HIP_PLATFORM_AMD__)
return (const uint16_t *)cuda_model_range_ptr(
model_map, offset, bytes, label);
#else
return (const uint16_t *)cuda_resolve_weight_ptr(
model_map, offset, bytes, 0, label);
#endif
}
static int glm53_vision_launch_ok(const char *label) {
return cuda_ok(cudaGetLastError(), label);
}
extern "C" int ds4_gpu_glm53_vision_encode(
float *out,
const float *patches,
uint32_t grid_h,
uint32_t grid_w,
const void *model_map,
uint64_t model_size,
const ds4_glm53_vision_weights *weights) {
if (!out || !patches || !model_map || !weights || grid_h == 0u ||
grid_w == 0u || (grid_h & 1u) != 0u || (grid_w & 1u) != 0u ||
grid_h > UINT32_MAX / grid_w) return 0;
const uint32_t rows = grid_h * grid_w;
const uint32_t merged_rows = rows / 4u;
const uint64_t row1024 = (uint64_t)rows * 1024u;
const uint64_t row3072 = (uint64_t)rows * 3072u;
const uint64_t row4096 = (uint64_t)rows * 4096u;
const uint64_t merged10240 = (uint64_t)merged_rows * 10240u;
if (row4096 > SIZE_MAX / sizeof(float) ||
merged10240 > SIZE_MAX / sizeof(float) ||
merged10240 > UINT32_MAX) return 0;
ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL;
ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL;
ds4_gpu_tensor *gate = NULL, *up = NULL, *mid = NULL;
ds4_gpu_tensor *cur = NULL, *tmp = NULL;
const uint16_t *bias = NULL;
const uint16_t *merger_norm = NULL, *merger_bias = NULL;
const uint64_t merged4096 = (uint64_t)merged_rows * 4096u;
int ok = 0;
#define VISION_ALLOC(name_, count_) do { \
name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \
if (!(name_)) goto cleanup; \
} while (0)
VISION_ALLOC(patch, (uint64_t)rows * 1176u);
VISION_ALLOC(a, row1024);
VISION_ALLOC(b, row1024);
VISION_ALLOC(qkv, row3072);
VISION_ALLOC(q, row1024);
VISION_ALLOC(k, row1024);
VISION_ALLOC(v, row1024);
VISION_ALLOC(attn, row1024);
VISION_ALLOC(gate, row4096);
VISION_ALLOC(up, row4096);
VISION_ALLOC(mid, row4096);
#undef VISION_ALLOC
if (!ds4_gpu_tensor_write(patch, 0, patches,
(uint64_t)rows * 1176u * sizeof(float)) ||
!ds4_gpu_begin_commands()) goto cleanup;
ok = ds4_gpu_glm53_matmul_bf16(
a, model_map, model_size, weights->patch_weight,
1176u, 1024u, patch, rows);
if (ok) {
bias = glm53_vision_weight(model_map, model_size,
weights->patch_bias, 1024u, "vision patch bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)a->ptr, bias, NULL, row1024, 1024u);
ok = glm53_vision_launch_ok("GLM-5.3 vision patch bias");
}
cur = a;
tmp = b;
for (uint32_t il = 0; ok && il < DS4_GLM53_VISION_LAYERS; il++) {
const ds4_glm53_vision_layer_weights *w = &weights->layer[il];
const uint16_t *norm = glm53_vision_weight(
model_map, model_size, w->norm1, 1024u, "vision norm1");
if (!norm) { ok = 0; break; }
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-5f);
ok = glm53_vision_launch_ok("GLM-5.3 vision norm1");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
qkv, model_map, model_size, w->qkv_weight,
1024u, 3072u, tmp, rows);
const uint16_t *qkv_bias = NULL, *q_norm = NULL, *k_norm = NULL;
if (ok) {
qkv_bias = glm53_vision_weight(model_map, model_size,
w->qkv_bias, 3072u, "vision QKV bias");
q_norm = glm53_vision_weight(model_map, model_size,
w->q_norm, 64u, "vision Q norm");
k_norm = glm53_vision_weight(model_map, model_size,
w->k_norm, 64u, "vision K norm");
if (!qkv_bias || !q_norm || !k_norm) ok = 0;
}
if (ok) {
glm53_vision_qkv_rope_kernel<<<dim3(rows, 16u, 1u), 32u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)q->ptr, (float *)k->ptr, (float *)v->ptr,
(const float *)qkv->ptr, qkv_bias, q_norm, k_norm,
rows, grid_w, 1.0e-5f);
ok = glm53_vision_launch_ok("GLM-5.3 vision QKV");
}
if (ok) {
glm53_vision_attention_kernel<<<dim3(rows, 16u, 1u), 32u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)attn->ptr, (const float *)q->ptr,
(const float *)k->ptr, (const float *)v->ptr, rows);
ok = glm53_vision_launch_ok("GLM-5.3 vision attention");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->attn_proj_weight,
1024u, 1024u, attn, rows);
if (ok) {
bias = glm53_vision_weight(model_map, model_size,
w->attn_proj_bias, 1024u, "vision attention bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, bias, (const float *)cur->ptr,
row1024, 1024u);
ok = glm53_vision_launch_ok("GLM-5.3 vision attention residual");
}
ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap;
if (ok) {
norm = glm53_vision_weight(model_map, model_size,
w->norm2, 1024u, "vision norm2");
if (!norm) ok = 0;
}
if (ok) {
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-5f);
ok = glm53_vision_launch_ok("GLM-5.3 vision norm2");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
gate, model_map, model_size, w->gate_weight,
1024u, 4096u, tmp, rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
up, model_map, model_size, w->up_weight,
1024u, 4096u, tmp, rows);
const uint16_t *gate_bias = NULL, *up_bias = NULL;
if (ok) {
gate_bias = glm53_vision_weight(model_map, model_size,
w->gate_bias, 4096u, "vision gate bias");
up_bias = glm53_vision_weight(model_map, model_size,
w->up_bias, 4096u, "vision up bias");
if (!gate_bias || !up_bias) ok = 0;
}
if (ok) {
glm53_vision_swiglu_bias_kernel<<<
(unsigned)((row4096 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)mid->ptr, (const float *)gate->ptr, gate_bias,
(const float *)up->ptr, up_bias, row4096, 4096u);
ok = glm53_vision_launch_ok("GLM-5.3 vision SwiGLU");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->down_weight,
4096u, 1024u, mid, rows);
if (ok) {
bias = glm53_vision_weight(model_map, model_size,
w->down_bias, 1024u, "vision down bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((row1024 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, bias, (const float *)cur->ptr,
row1024, 1024u);
ok = glm53_vision_launch_ok("GLM-5.3 vision down residual");
}
swap = cur; cur = tmp; tmp = swap;
}
if (ok) {
const uint16_t *norm = glm53_vision_weight(model_map, model_size,
weights->post_norm, 1024u, "vision post norm");
if (!norm) ok = 0;
else {
glm53_vision_rms_kernel<<<rows, 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
norm, 1024u, 1.0e-5f);
ok = glm53_vision_launch_ok("GLM-5.3 vision post norm");
}
}
if (ok) {
glm53_vision_reorder_kernel<<<
(unsigned)((merged4096 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)cur->ptr, (const float *)tmp->ptr, merged4096);
ok = glm53_vision_launch_ok("GLM-5.3 vision downsample reorder");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, weights->downsample_weight,
4096u, 4096u, cur, merged_rows);
if (ok) {
bias = glm53_vision_weight(model_map, model_size,
weights->downsample_bias, 4096u, "vision downsample bias");
if (!bias) ok = 0;
}
if (ok) {
glm53_vision_bias_kernel<<<
(unsigned)((merged4096 + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, bias, NULL, merged4096, 4096u);
ok = glm53_vision_launch_ok("GLM-5.3 vision downsample bias");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
cur, model_map, model_size, weights->merger_proj,
4096u, 4096u, tmp, merged_rows);
if (ok) {
merger_norm = glm53_vision_weight(model_map, model_size,
weights->merger_norm, 4096u, "vision merger norm");
merger_bias = glm53_vision_weight(model_map, model_size,
weights->merger_norm_bias, 4096u, "vision merger norm bias");
if (!merger_norm || !merger_bias) ok = 0;
}
if (ok) {
glm53_vision_layernorm_gelu_kernel<<<merged_rows, 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)tmp->ptr, (const float *)cur->ptr,
merger_norm, merger_bias, 4096u, 1.0e-5f);
ok = glm53_vision_launch_ok("GLM-5.3 vision merger norm");
}
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
gate, model_map, model_size, weights->merger_gate,
4096u, 10240u, tmp, merged_rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
up, model_map, model_size, weights->merger_up,
4096u, 10240u, tmp, merged_rows);
if (ok) ok = ds4_gpu_swiglu_tensor(
mid, gate, up, (uint32_t)merged10240, 10.0f, 1.0f);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
cur, model_map, model_size, weights->merger_down,
10240u, 4096u, mid, merged_rows);
if (ds4_gpu_end_commands() == 0) ok = 0;
if (ok) ok = ds4_gpu_tensor_read(
cur, 0, out, merged4096 * sizeof(float));
cleanup:
ds4_gpu_tensor_free(mid);
ds4_gpu_tensor_free(up);
ds4_gpu_tensor_free(gate);
ds4_gpu_tensor_free(attn);
ds4_gpu_tensor_free(v);
ds4_gpu_tensor_free(k);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(qkv);
ds4_gpu_tensor_free(b);
ds4_gpu_tensor_free(a);
ds4_gpu_tensor_free(patch);
return ok;
}
extern "C" int ds4_gpu_glm53_scatter_image_hc(
ds4_gpu_tensor *hc,
const ds4_gpu_tensor *image,
uint32_t dst_row,
uint32_t image_row,
uint32_t rows,
uint32_t total_rows,
uint32_t n_embd,
uint32_t n_hc) {
if (!hc || !image || rows == 0u || n_embd == 0u || n_hc == 0u ||
dst_row > total_rows || rows > total_rows - dst_row) return 0;
const uint64_t count = (uint64_t)rows * n_hc * n_embd;
glm53_vision_scatter_hc_kernel<<<
(unsigned)((count + 255u) / 256u), 256u, 0,
DS4_GLM53_VISION_STREAM>>>(
(float *)hc->ptr, (const float *)image->ptr,
dst_row, image_row, rows, total_rows, n_embd, n_hc);
return glm53_vision_launch_ok("GLM-5.3 vision scatter");
}
#undef DS4_GLM53_VISION_STREAM
ds4_gpu.h
+2204 -3
View File
File diff suppressed because it is too large. Load diff
ds4_gpu_args.c
+230
View File
@@ -0,0 +1,230 @@
/* ds4_gpu_args.c — shared parser for --gpu-vram and --gpu-devices.
*
* See ds4_gpu_args.h. The actual CUDA probing for the "auto" branch
* lives in ds4_cuda.cu so this .c file never includes <cuda_runtime.h>
* directly (mirroring the rest of the tree's separation between .c and
* .cu compilation units).
*/
#include "ds4_gpu_args.h"
#include <ctype.h>
#include <errno.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define DS4_GPU_ARGS_DEFAULT_SAFETY_MARGIN ((size_t)512 * 1024 * 1024)
/* Helper: copy at most outlen-1 bytes of msg into out and NUL-terminate. */
static void errbuf_set(char *out, size_t outlen, const char *msg) {
if (!out || outlen == 0) return;
size_t n = strlen(msg);
if (n >= outlen) n = outlen - 1;
memcpy(out, msg, n);
out[n] = '\0';
}
static void errbuf_setf(char *out, size_t outlen, const char *fmt, ...) {
if (!out || outlen == 0) return;
va_list ap;
va_start(ap, fmt);
vsnprintf(out, outlen, fmt, ap);
va_end(ap);
}
/* Maximum per-entry values. Generous, but small enough to keep all
* downstream arithmetic safely inside size_t/int. */
#define DS4_GPU_ARGS_MAX_VRAM_GB 16384 /* 16 TiB per device */
#define DS4_GPU_ARGS_MAX_DEVICE_INDEX 1023 /* CUDA visible-device ID cap */
/* Parse a comma-separated list of nonnegative integers into out_count
* entries of out[]. Returns 0 on success, nonzero on parse error.
* Empty entries (",,") or trailing/leading commas error. Per-entry
* upper bound enforced via max_value (use LONG_MAX for "no cap"). */
static int parse_csv_int_list(const char *s,
long *out,
int *out_count,
int max_count,
long max_value,
char *errbuf,
size_t errbuflen,
const char *what) {
*out_count = 0;
if (!s || !*s) {
errbuf_setf(errbuf, errbuflen, "%s: empty value", what);
return 1;
}
const char *p = s;
while (*p) {
/* Disallow leading whitespace inside the list. */
while (*p == ' ' || *p == '\t') p++;
char *end = NULL;
errno = 0;
long v = strtol(p, &end, 10);
if (end == p) {
errbuf_setf(errbuf, errbuflen, "%s: not a number: '%s'", what, p);
return 1;
}
if (errno != 0 || v < 0 || v > max_value) {
errbuf_setf(errbuf, errbuflen,
"%s: out-of-range value: '%s' (allowed 0..%ld)",
what, p, max_value);
return 1;
}
if (*out_count >= max_count) {
errbuf_setf(errbuf, errbuflen, "%s: too many entries (max %d)", what, max_count);
return 1;
}
out[*out_count] = v;
(*out_count)++;
while (*end == ' ' || *end == '\t') end++;
if (*end == ',') {
end++;
if (*end == '\0' || *end == ',') {
errbuf_setf(errbuf, errbuflen, "%s: empty entry near '%s'", what, end);
return 1;
}
} else if (*end != '\0') {
errbuf_setf(errbuf, errbuflen, "%s: junk after value: '%s'", what, end);
return 1;
}
p = end;
}
if (*out_count == 0) {
errbuf_setf(errbuf, errbuflen, "%s: no values parsed", what);
return 1;
}
return 0;
}
int parse_gpu_vram_arg(const char *vram_arg,
const char *devices_arg,
ds4_gpu_config *out,
bool *out_skip_cuda,
char *errbuf,
size_t errbuflen) {
if (!out || !out_skip_cuda) {
errbuf_set(errbuf, errbuflen, "internal: NULL out parameter");
return 1;
}
*out_skip_cuda = false;
/* Caller zeroed *out — just confirm. */
out->n_gpus = 0;
out->safety_margin_bytes = DS4_GPU_ARGS_DEFAULT_SAFETY_MARGIN;
/* Parse devices_arg first so we have the filter when handling "auto". */
long dev_list[DS4_MAX_GPUS] = {0};
int dev_count = 0;
if (devices_arg) {
if (parse_csv_int_list(devices_arg, dev_list, &dev_count,
DS4_MAX_GPUS,
DS4_GPU_ARGS_MAX_DEVICE_INDEX,
errbuf, errbuflen,
"--gpu-devices") != 0) {
return 1;
}
}
/* Case 1: vram_arg unset. */
if (vram_arg == NULL) {
if (devices_arg == NULL) {
errbuf_set(errbuf, errbuflen, "internal: parse_gpu_vram_arg called with both args NULL");
return 1;
}
/* --gpu-devices given without --gpu-vram: implicit auto with filter. */
vram_arg = "auto";
}
/* Case 2: --gpu-vram 0 — short-circuit. */
if (!strcmp(vram_arg, "0")) {
if (devices_arg) {
errbuf_set(errbuf, errbuflen,
"--gpu-vram 0 conflicts with --gpu-devices");
return 1;
}
*out_skip_cuda = true;
return 0;
}
/* Case 3: --gpu-vram auto. */
if (!strcmp(vram_arg, "auto")) {
#if !defined(DS4_NO_GPU) && !defined(__APPLE__)
int filter_buf[DS4_MAX_GPUS];
const int *filter_ptr = NULL;
if (dev_count > 0) {
for (int i = 0; i < dev_count; i++) filter_buf[i] = (int)dev_list[i];
filter_ptr = filter_buf;
}
if (ds4_gpu_args_probe_auto_cuda(filter_ptr, dev_count, out,
out->safety_margin_bytes,
errbuf, errbuflen) != 0) {
return 1;
}
return 0;
#else
(void)dev_list; (void)dev_count;
errbuf_set(errbuf, errbuflen,
"--gpu-vram auto requires a CUDA build");
return 1;
#endif
}
/* Case 4: explicit GB budgets — "N[,N...]" */
long vram_list[DS4_MAX_GPUS] = {0};
int vram_count = 0;
if (parse_csv_int_list(vram_arg, vram_list, &vram_count,
DS4_MAX_GPUS,
DS4_GPU_ARGS_MAX_VRAM_GB,
errbuf, errbuflen,
"--gpu-vram") != 0) {
return 1;
}
if (dev_count > 0 && dev_count != vram_count) {
errbuf_setf(errbuf, errbuflen,
"--gpu-devices count (%d) does not match --gpu-vram count (%d)",
dev_count, vram_count);
return 1;
}
out->n_gpus = vram_count;
for (int i = 0; i < vram_count; i++) {
int dev_id = (dev_count > 0) ? (int)dev_list[i] : i;
out->device_indices[i] = dev_id;
out->vram_bytes[i] = (size_t)vram_list[i] * (size_t)1024 * 1024 * 1024;
}
return 0;
}
int format_gpu_layout_line(const ds4_gpu_config *cfg,
bool was_auto,
char *out,
size_t outlen) {
if (!out || outlen == 0 || !cfg) return -1;
/* "ds4: GPU config: N devices [d0,d1,...] requested, budgets G0,G1,... GB; auto=<bool>" */
int written = snprintf(out, outlen,
"ds4: GPU config: %d device%s [",
cfg->n_gpus, cfg->n_gpus == 1 ? "" : "s");
if (written < 0 || (size_t)written >= outlen) return written;
for (int i = 0; i < cfg->n_gpus; i++) {
int n = snprintf(out + written, outlen - (size_t)written,
"%s%d", i ? "," : "", cfg->device_indices[i]);
if (n < 0 || (size_t)(written + n) >= outlen) return written + n;
written += n;
}
int n = snprintf(out + written, outlen - (size_t)written,
"] requested, budgets ");
if (n < 0 || (size_t)(written + n) >= outlen) return written + n;
written += n;
for (int i = 0; i < cfg->n_gpus; i++) {
double gb = (double)cfg->vram_bytes[i] / (1024.0 * 1024.0 * 1024.0);
int m = snprintf(out + written, outlen - (size_t)written,
"%s%.0f", i ? "," : "", gb);
if (m < 0 || (size_t)(written + m) >= outlen) return written + m;
written += m;
}
int t = snprintf(out + written, outlen - (size_t)written,
" GB; auto=%s", was_auto ? "true" : "false");
if (t < 0) return t;
written += t;
return written;
}
ds4_gpu_args.h
+84
View File
@@ -0,0 +1,84 @@
/* ds4_gpu_args.h — shared CLI argument parser for --gpu-vram and
* --gpu-devices flags. Used by ds4_cli, ds4_server, ds4_agent, and
* ds4_bench. See mgpu-cli-wiring task (wave 2) for design.
*
* The parser does not link CUDA headers directly. On CUDA builds the
* "auto" path is gated and dispatches to ds4_gpu_args_probe_auto_cuda,
* which is defined in ds4_cuda.cu. On Metal/CPU builds, "auto" returns
* an error.
*/
#ifndef DS4_GPU_ARGS_H
#define DS4_GPU_ARGS_H
#include <stdbool.h>
#include <stddef.h>
#include "ds4_gpu_mgpu.h"
#ifdef __cplusplus
extern "C" {
#endif
/* Single canonical parser. Populates *out and *out_skip_cuda on
* success; returns 0 on success, nonzero on parse error (errbuf
* populated). Caller zeroes out and out_skip_cuda before calling.
*
* vram_arg may be:
* NULL -- flag unset
* "0" -- skip CUDA entirely; sets *out_skip_cuda = true
* "auto" -- query free VRAM via CUDA (CUDA build only)
* "N[,N...]" -- explicit GB budgets per device
*
* devices_arg may be:
* NULL -- 0..n-1 of visible devices (or 0..vram_count-1)
* "N[,N...]" -- explicit CUDA device indices
*
* Constraints:
* - If both args are given and counts differ, error.
* - Negative or non-integer values error.
* - Counts > DS4_MAX_GPUS error.
*/
int parse_gpu_vram_arg(
const char *vram_arg,
const char *devices_arg,
ds4_gpu_config *out,
bool *out_skip_cuda,
char *errbuf,
size_t errbuflen);
/* Format the layout echo line written to stdout right before
* ds4_engine_create_with_gpu_config is called.
*
* Example output:
* ds4: GPU config: 2 devices [0,1] requested, budgets 40,12 GB; auto=true
*
* Returns the number of bytes written (excluding the trailing NUL), or
* -1 if out is NULL/outlen is 0.
*/
int format_gpu_layout_line(const ds4_gpu_config *cfg,
bool was_auto,
char *out,
size_t outlen);
/* CUDA-only auto probe. Defined in ds4_cuda.cu under the same gate.
* Returns 0 on success, nonzero on failure (errbuf populated).
* device_filter may be NULL (use all visible devices) or a list of
* length filter_len (use only those device indices).
*
* Declared here for the parser's use; not called by application code
* directly. The Mac/CPU builds skip this branch via the
* #if !defined(DS4_NO_GPU) && !defined(__APPLE__) gate in the parser. */
#if !defined(DS4_NO_GPU) && !defined(__APPLE__)
int ds4_gpu_args_probe_auto_cuda(const int *device_filter,
int filter_len,
ds4_gpu_config *out,
size_t safety_margin_bytes,
char *errbuf,
size_t errbuflen);
#endif
#ifdef __cplusplus
}
#endif
#endif /* DS4_GPU_ARGS_H */
ds4_gpu_copy.h
+24
View File
@@ -0,0 +1,24 @@
#ifndef DS4_GPU_COPY_H
#define DS4_GPU_COPY_H
#include <stdint.h>
/* Device copy spans, shared by ds4_gpu.h and the CUDA backend, which does
* not include ds4_gpu.h. The tensor stays opaque here. */
#ifndef DS4_GPU_TENSOR_DEFINED
#define DS4_GPU_TENSOR_DEFINED
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
#endif
/* One device copy, with the arguments of ds4_gpu_tensor_copy. */
typedef struct {
ds4_gpu_tensor *dst;
uint64_t dst_offset;
const ds4_gpu_tensor *src;
uint64_t src_offset;
uint64_t bytes;
} ds4_gpu_copy_span;
#define DS4_GPU_COPY_SPANS_MAX 316u
#endif
ds4_gpu_mgpu.h
+237
View File
@@ -0,0 +1,237 @@
/* ds4_gpu_mgpu.h — multi-GPU plumbing types and APIs (v0).
*
* This header carries the new multi-GPU additions for the multi-GPU plumbing PP work
* (device-aware CUDA). It is included from ds4_cuda.cu and from
* downstream tasks that need access to g_gpu[], g_n_gpus, g_gpu_peer_ok[],
* the ds4_gpu_config struct, and the new tensor APIs.
*
* Why not in ds4_gpu.h? The legacy ds4_gpu.h is included from C-only
* callers (ds4.c, ds4_cli.c, etc.) and from the Metal build, but is NOT
* included from ds4_cuda.cu historically. That asymmetry hid pre-existing
* signature mismatches between the legacy header and ds4_cuda.cu. We keep
* the legacy header opaque and put the new shared types here, so this
* file is the single source of truth for both ds4_cuda.cu and downstream
* multi-GPU tasks without disturbing the legacy contract.
*
* The struct definitions reference CUDA-specific handle types via void *
* placeholders so the header is safe to include from C builds, Metal
* builds, and the CUDA build (where ds4_cuda.cu casts the void * back
* to cudaStream_t / cublasHandle_t / cudaEvent_t internally).
*/
#ifndef DS4_GPU_MGPU_H
#define DS4_GPU_MGPU_H
#include <stddef.h>
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
#define DS4_MAX_GPUS 16
/* Complete definition of the previously-opaque ds4_gpu_tensor, plus a
* typedef so the new API prototypes below can use the bare name
* `ds4_gpu_tensor *` in both C and C++ without forcing callers to
* include ds4_gpu.h first. Callers that include this header can
* stack-allocate or struct-embed tensors and pass to
* ds4_gpu_tensor_alloc_on. */
struct ds4_gpu_tensor {
void *ptr;
uint64_t bytes;
int owner;
int device_id; /* -1 means legacy/untagged → treat as device 0 */
};
#ifndef DS4_GPU_TENSOR_DEFINED
#define DS4_GPU_TENSOR_DEFINED
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
#endif
#ifndef DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
#define DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
#define DS4_GPU_ATTENTION_DECODE_BATCH_MAX 32u
typedef struct {
uint64_t raw_kv;
uint64_t comp_kv;
uint64_t topk;
uint32_t pos;
uint32_t n_raw;
uint32_t raw_cap;
uint32_t raw_start;
uint32_t n_comp;
uint32_t top_k;
uint32_t window;
uint32_t ratio;
uint32_t indexed;
} ds4_gpu_attention_decode_row;
#endif
/* Tagged so headers (notably ds4.h) can forward-declare `struct
* ds4_gpu_config` without dragging in this entire header. */
typedef struct ds4_gpu_config {
int device_indices[DS4_MAX_GPUS]; /* CUDA device IDs to use */
/* Explicit per-device budget in bytes. The engine does NOT auto-fill
* missing budgets - a value of 0 means "zero bytes of budget for
* this slot" and (combined with reserves) will push placement to
* CPU spill. Auto-detection (e.g. mapping --gpu-vram auto to
* cudaMemGetInfo) is the caller's job; see CLI flag wiring for the
* canonical CLI path. The engine emits a clear stderr and refuses
* if n_gpus > 0 and every vram_bytes[] is 0 (almost certainly a
* caller bug from zero-initializing the struct). */
size_t vram_bytes[DS4_MAX_GPUS];
int n_gpus;
size_t safety_margin_bytes; /* per-device reserve */
} ds4_gpu_config;
typedef struct {
int device_id;
void *stream; /* cudaStream_t under CUDA */
void *cublas; /* cublasHandle_t under CUDA */
int cublas_ready;
void *scratch;
size_t scratch_bytes;
size_t budget_bytes;
size_t used_bytes;
void *boundary_event; /* cudaEvent_t under CUDA */
} ds4_gpu_ctx;
extern ds4_gpu_ctx g_gpu[DS4_MAX_GPUS];
extern int g_n_gpus;
extern int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS];
/* Primary multi-device init. The existing ds4_gpu_init (declared in
* ds4_gpu.h) is a thin shim that builds a single-device config for
* device 0 and calls this. */
int ds4_gpu_init_multi(const ds4_gpu_config *cfg);
/* Caller-supplied struct alloc on a specific device. Returns 0 on
* success, nonzero on error. Pair with ds4_gpu_tensor_free_in_place. */
int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, uint64_t bytes);
void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t);
/* Heap-allocated tensor on a specific logical tier; mirrors the legacy
* ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *) but with a tier
* parameter. Returns NULL on failure. Used by the multi-tier graph
* allocations in ds4.c. Single-tier callers can continue using the
* legacy ds4_gpu_tensor_alloc(bytes) which is equivalent to
* ds4_gpu_tensor_alloc_ptr_on(0, bytes). */
ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes);
int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst, const ds4_gpu_tensor *src, uint64_t bytes);
void ds4_gpu_enable_q8_dequant_gemm(void);
/* Heap-allocated MANAGED-memory tensor on a specific logical tier. Used
* for the KV-cache pool. Stamps tier on the tensor so subsequent
* ds4_gpu_tensor_free runs under the correct device. Note managed memory
* pages on first-touch and is not strictly device-bound; the stamped
* tier records the home tier for accounting + free. */
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes);
/* Cross-device tensor copy. Same-device → cudaMemcpyAsync; peer-capable
* cross-device → cudaMemcpyPeerAsync with event sync; non-peer → pinned
* host bounce (per src→dst pair). Honors DS4_FORCE_HOST_BOUNCE=1. */
int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes);
int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes);
/* Grouped default-stream handoff whose copies execute on the destination
* device. The source default stream records readiness; the destination waits,
* performs all three copies, then naturally orders its following kernels.
* This allows source-device work submitted afterward to overlap the handoff. */
int ds4_gpu_tensor_copy_xdev3_default_dst(
ds4_gpu_tensor *dst0,
const ds4_gpu_tensor *src0,
uint64_t bytes0,
ds4_gpu_tensor *dst1,
const ds4_gpu_tensor *src1,
uint64_t bytes1,
ds4_gpu_tensor *dst2,
const ds4_gpu_tensor *src2,
uint64_t bytes2);
/* Cross-device copy of three buffers between the same source/destination
* tiers, recording a single readiness event for the destination stream. This
* is meant for tiny grouped activation handoffs where three independent
* ds4_gpu_tensor_copy_xdev calls would spend more time in event plumbing than
* in the copies themselves. Falls back internally when a grouped peer copy is
* not applicable. */
int ds4_gpu_tensor_copy_xdev3(ds4_gpu_tensor *dst0,
const ds4_gpu_tensor *src0,
uint64_t bytes0,
ds4_gpu_tensor *dst1,
const ds4_gpu_tensor *src1,
uint64_t bytes1,
ds4_gpu_tensor *dst2,
const ds4_gpu_tensor *src2,
uint64_t bytes2);
/* Cross-device copy that also orders against prior work on the destination
* stream before overwriting dst. This is needed for pipelined prefill slots
* where the next producer reuses a destination buffer that the destination
* stage just consumed. */
int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes);
/* Order a peer read without copying: records readiness on src's stream and
* makes dst_tier's stream wait for it. */
int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier);
int ds4_gpu_tensor_wait_xdev_default(const ds4_gpu_tensor *src, int dst_tier);
/* Cross-device float add used by CUDA tensor-parallel reductions.
* See ds4_gpu.h for the full contract. */
int ds4_gpu_add_xdev_tensor(ds4_gpu_tensor *out,
const ds4_gpu_tensor *local,
const ds4_gpu_tensor *remote,
ds4_gpu_tensor *remote_tmp,
uint32_t n);
/* Returns the device_id recorded on the tensor; -1 if untagged. */
int ds4_gpu_tensor_device(const ds4_gpu_tensor *t);
/* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1).
*
* This is the canonical shim for per-layer device routing in the
* multi-tier execution path. The caller passes a logical tier index;
* the shim internally indexes g_gpu[tier].device_id and calls
* cudaSetDevice. Returns 0 on success, nonzero on error or if the
* tier index is out of range.
*
* Wave-2 multi-GPU placement scaffolding adds this shim but does not
* exercise the multi-tier execution path; multi-GPU execution
* (follow-up) is its first caller. */
int ds4_gpu_set_current_device(int logical_tier);
int ds4_gpu_set_current_device_fenced(int logical_tier);
/* Register the mmap'd host model pointer for selective-cache lookups
* WITHOUT triggering any device-side copy. This bypasses the
* DS4_CUDA_COPY_MODEL environment-variable branch that
* ds4_gpu_set_model_map normally honors, which is essential for
* multi-tier startup: we want only per-device selective tensor caches,
* never the whole-model copy. Returns 1 on success, 0 on error. */
int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size);
/* Strict per-device selective-cache lookup (no fallback).
*
* Returns 1 only if a covering entry exists whose device_id matches the
* caller-supplied expected_device (a PHYSICAL CUDA device id). Otherwise
* returns 0 — no host-pointer fallback, no different-device match. The
* caller is expected to have cudaSetDevice'd to expected_device before
* invoking; the returned pointer is valid to consume from that device's
* kernel.
*
* Used by multi-tier kernel-dispatch resolvers in
* multi-GPU execution (multi-GPU execution). Single-tier callers should
* keep using ds4_gpu_lookup_cache for back-compat behavior. */
int ds4_gpu_lookup_cache_strict(uint64_t source_offset,
uint64_t bytes,
int expected_device,
void **out_device_ptr);
#ifdef __cplusplus
} /* extern "C" */
#endif
#endif /* DS4_GPU_MGPU_H */
ds4_gpu_tp.h
+63
View File
@@ -0,0 +1,63 @@
#ifndef DS4_GPU_TP_H
#define DS4_GPU_TP_H
#include <stdint.h>
#ifndef DS4_GPU_TENSOR_DEFINED
#define DS4_GPU_TENSOR_DEFINED
typedef struct ds4_gpu_tensor ds4_gpu_tensor;
#endif
#ifdef __cplusplus
extern "C" {
#endif
/* Metal queues each gate on its service thread and orders GPU arrival/release
* with events or flags. CUDA waits for its local stream and calls the exchange
* synchronously. Callbacks return nonzero on success; both ranks must issue
* the same gate sequence. Shutdown precedes transport/slab destruction. */
typedef int (*ds4_gpu_tp_exchange_fn)(void *ud, uint32_t layer, uint32_t gate, uint64_t seq);
typedef int (*ds4_gpu_tp_batch_exchange_fn)(void *ud, uint32_t layer,
uint32_t rows, uint64_t seq);
typedef int (*ds4_gpu_tp_big_exchange_fn)(void *ud, uint32_t layer,
uint64_t seq, const void *out,
void *in, uint64_t bytes);
/* CUDA: optional pinned slab shared by an integrated GPU and the transport.
* NULL requests the ordinary device tensor plus host staging fallback. */
ds4_gpu_tensor *ds4_gpu_tp_slab_alloc(uint64_t bytes);
int ds4_gpu_tp_init(uint32_t rank, ds4_gpu_tensor *slab,
uint64_t gpu_flags_off, uint64_t out_off, uint64_t vec_bytes,
ds4_gpu_tp_exchange_fn fn, void *ud);
void ds4_gpu_tp_shutdown(void);
int ds4_gpu_tp_failed(void);
int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate);
void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn);
int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows);
void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn);
int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t, uint64_t bytes);
/* Metal multi-session tapes reuse slab slots and therefore require event
* arrival. Single-session flag gates may flush in layer order. */
void ds4_gpu_tp_set_session_batch_mode(int enabled);
/* Metal decode gates publish and wait inside the command buffer (no flush per gate). */
int ds4_gpu_tp_decode_inline_gates(void);
int ds4_gpu_tp_decode_split_flush_safe(void);
/* Weight ranges consumed by the next Metal poll gate of the given kind. */
int ds4_gpu_tp_gate_prefetch_plan(uint32_t gate,
const void *model_map, uint64_t model_size,
const uint64_t *offsets, const uint64_t *bytes,
uint32_t count);
/* Coordinator-only drafting is replicated; verification remains sharded. */
void ds4_gpu_tp_suspend_expert_sharding(int suspend);
/* No-op when TP is not bound. */
void ds4_gpu_tp_keepalive_pause(int paused);
/* GLM attention head ownership; the caller zeros unowned heads and exchanges
* the output-projection partials at the big gate. */
void ds4_gpu_tp_set_attn_head_split(int enabled);
#ifdef __cplusplus
}
#endif
#endif
ds4_help.c
+67 -19
View File
@@ -147,12 +147,20 @@ static void print_model_runtime(FILE *fp, const help_colors *c,
ds4_help_tool tool, bool full) {
title(fp, c, "Model And Runtime");
opt(fp, c, "-m, --model FILE", "GGUF model path. Default: ds4flash.gguf");
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) {
opt(fp, c, "--vision FILE", "Vision encoder GGUF for the selected model.");
}
#ifdef DS4_ROCM_BUILD
opt(fp, c, "--metal | --rocm | --cpu", "Select the backend explicitly.");
opt(fp, c, "--backend NAME", "Backend name: metal, rocm, or cpu.");
#else
opt(fp, c, "--metal | --cuda | --cpu", "Select the backend explicitly.");
opt(fp, c, "--backend NAME", "Backend name: metal, cuda, or cpu.");
opt(fp, c, "--gpu-vram N[,N,...]|auto", "CUDA VRAM budgets per device, in GiB, or auto-detect free VRAM.");
opt(fp, c, "--gpu-devices N[,N,...]", "CUDA device indices used by multi-GPU placement.");
if (tool != DS4_HELP_EVAL) {
opt(fp, c, "--cuda-tensor-parallel", "Enable the paired DeepSeek tensor/expert path on an even multi-GPU CUDA placement.");
}
#endif
if (tool != DS4_HELP_BENCH) {
opt(fp, c, "-c, --ctx N", "Allocated context tokens.");
@@ -164,21 +172,31 @@ static void print_model_runtime(FILE *fp, const help_colors *c,
opt(fp, c, "--power N", "GPU duty-cycle target, 1..100. Default: 100");
opt(fp, c, "--ssd-streaming", "Metal/CUDA/ROCm: opt in to SSD-backed model streaming instead of full residency.");
opt(fp, c, "--ssd-streaming-cold", "SSD streaming: skip default popularity-based expert-cache preload.");
opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming: N is an exact dynamic expert count; NGB is a routed memory budget that also reserves two full prefill layers. Auto: 80% working set minus non-routed weights; GLM Metal caps lower.");
opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming cache target. N requests dynamic expert slots; NGB also reserves two full prefill layers. Either may be reduced to fit the model, graph, context, and backend working set.");
opt(fp, c, "--ssd-streaming-full-layers N", "GLM Metal streaming: keep the first N routed layers fully resident. Default: auto from NGB expert budget; use 0 to disable.");
opt(fp, c, "--ssd-streaming-preload-experts N", "SSD streaming: upfront popularity preload count. DeepSeek auto-seeds by default; GLM demand-fills unless N is explicit.");
opt(fp, c, "--simulate-used-memory NGB", "Diagnostic: lock N GiB before model load to simulate a smaller-memory machine.");
opt(fp, c, "--prefill-chunk N", "Metal graph prefill chunk size. Default: auto (PRO long prompts use 8192; others use 4096).");
opt(fp, c, "--prefill-chunk N", "Graph prefill chunk size. Default: CUDA TP 2048; PRO long prompts 8192; others 4096.");
if (full) {
if (tool != DS4_HELP_BENCH) {
opt(fp, c, "--mtp FILE", "Optional MTP support GGUF used for draft-token probes.");
if (tool == DS4_HELP_EVAL || tool == DS4_HELP_BENCH) {
opt(fp, c, "--mtp-model FILE", "External MTP or DSpark support GGUF.");
}
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) {
opt(fp, c, "--mtp", "Enable model-embedded MTP speculation.");
opt(fp, c, "--mtp-model FILE", "External MTP or DSpark support GGUF.");
opt(fp, c, "--mtp-draft N", "Maximum autoregressive MTP draft tokens. Default: 1");
opt(fp, c, "--mtp-margin F", "Verifier confidence margin for fast MTP acceptance. Default: 3");
opt(fp, c, "--mtp-timing", "Enable embedded MTP and print acceptance/timing counters.");
opt(fp, c, "--dspark", "Enable DSpark using the support GGUF passed with --mtp-model.");
opt(fp, c, "--dspark-confidence F", "Enable DSpark with confidence pruning threshold 0..1. Greedy/opportunistic default: resident Spark Q2 0.4, two-Spark Flash Q2/MXFP4 0, Metal 0.6, other CUDA/ROCm 0.7; exact sampling: 0.8");
opt(fp, c, "--mtp-exact-sampling", "Preserve the ordinary temperature distribution instead of accepting target-matching greedy drafts directly.");
opt(fp, c, "--dspark-strict", "Load DSpark support but keep target-only decode.");
} else if (tool == DS4_HELP_BENCH) {
opt(fp, c, "--dspark", "Benchmark greedy DSpark using the support GGUF passed with --mtp-model.");
opt(fp, c, "--dspark-confidence F", "DSpark confidence pruning threshold 0..1.");
}
opt(fp, c, "--quality", "Prefer exact kernels where faster approximate paths exist.");
opt(fp, c, "--warm-weights", "Touch mapped tensor pages at startup to reduce first-use stalls.");
opt(fp, c, "--warm-weights", "Touch resident weights at startup to reduce first-use stalls.");
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_BENCH) {
opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON.");
}
@@ -186,7 +204,7 @@ static void print_model_runtime(FILE *fp, const help_colors *c,
fputc('\n', fp);
}
static void print_sampling(FILE *fp, const help_colors *c, bool full) {
static void print_sampling(FILE *fp, const help_colors *c, bool full, ds4_help_tool tool) {
title(fp, c, "Prompt And Sampling");
opt(fp, c, "-n, --tokens N", "Maximum generated tokens.");
opt(fp, c, "--temp F", "Sampling temperature. 0 is greedy/deterministic.");
@@ -194,8 +212,10 @@ static void print_sampling(FILE *fp, const help_colors *c, bool full) {
opt(fp, c, "--min-p F", "Keep tokens scoring at least F times the top token.");
opt(fp, c, "--seed N", "Sampling seed for reproducible non-greedy runs.");
para(fp, c, "GLM CLI and agent runs default to temperature 1.0, top-p 0.95, and min-p 0 unless those options are set explicitly.");
opt(fp, c, "--think", "Use normal thinking mode.");
opt(fp, c, "--think-max", "Use Think Max when context is large enough.");
opt(fp, c, "--think", "Use normal thinking mode (V4.1: effort 75).");
opt(fp, c, "--think-max", "Use maximum thinking (V4.1: 100; V4: requires ctx >= 393216).");
if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT)
opt(fp, c, "--think-level N", "V4.1 thinking effort, 1..100; 0 disables thinking.");
opt(fp, c, "--nothink", "Disable thinking and ask for direct replies.");
if (full) {
opt(fp, c, "-sys, --system TEXT", "System prompt. Empty string disables the default where supported.");
@@ -229,6 +249,17 @@ static void print_distributed(FILE *fp, const help_colors *c) {
opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits.");
opt(fp, c, "--debug", "Print coordinator route/debug logs.");
fputc('\n', fp);
title(fp, c, "Tensor Parallelism");
fputc('\n', fp);
para(fp, c, "Tensor parallelism uses the same coordinator/worker addresses as distributed mode, but always runs one 50/50 worker. Add --tensor-parallel, omit --layers, start the worker, then start the coordinator.");
fputc('\n', fp);
opt(fp, c, "--tensor-parallel", "Switch --role/--listen/--coordinator to two-machine tensor parallelism.");
opt(fp, c, "--transport auto|rdma|tcp", "Tensor gate transport. Default: auto");
opt(fp, c, "--rdma-device NAME", "Select a verbs device when auto-detection is ambiguous.");
opt(fp, c, "--rdma-gid-index N", "Select the local verbs GID index.");
opt(fp, c, "--tensor-parallel-token-prefill", "GLM diagnostic: prefill one token at a time for exact arithmetic.");
opt(fp, c, "--debug-hash N", "Cross-check hidden state every N tokens.");
fputc('\n', fp);
}
static void print_cli_diagnostics(FILE *fp, const help_colors *c);
@@ -238,6 +269,7 @@ static void print_cli_specific(FILE *fp, const help_colors *c, bool full) {
opt(fp, c, "ds4", "Start the interactive prompt.");
opt(fp, c, "ds4 -p TEXT", "Run one prompt and exit.");
opt(fp, c, "ds4 --prompt-file FILE", "Run a long prompt from a file and exit.");
opt(fp, c, "--prefix-file FILE", "Preload complete alternating USER:/ASSISTANT: turns before the live conversation.");
fputc('\n', fp);
if (full) {
print_cli_diagnostics(fp, c);
@@ -247,7 +279,7 @@ static void print_cli_specific(FILE *fp, const help_colors *c, bool full) {
static void print_cli_diagnostics(FILE *fp, const help_colors *c) {
title(fp, c, "Diagnostics And Data Collection");
opt(fp, c, "--inspect", "Load the model and print a summary only.");
opt(fp, c, "--dump-tokens", "Tokenize the prompt exactly as written, then exit.");
opt(fp, c, "--dump-tokens", "Print the exact CLI prompt token stream, then exit. Use --raw for literal text.");
opt(fp, c, "--dump-logits FILE", "Write full next-token logits as JSON.");
opt(fp, c, "--dump-logprobs FILE", "Write greedy continuation top-logprobs as JSON.");
opt(fp, c, "--logprobs-top-k N", "Alternatives stored by --dump-logprobs. Default: 20");
@@ -258,6 +290,7 @@ static void print_cli_diagnostics(FILE *fp, const help_colors *c) {
opt(fp, c, "--imatrix-out FILE", "Write llama-compatible routed-MoE imatrix .dat.");
opt(fp, c, "--imatrix-max-prompts N", "Stop imatrix collection after N prompts.");
opt(fp, c, "--imatrix-max-tokens N", "Stop imatrix collection after N prompt tokens.");
opt(fp, c, "--imatrix-min-expert-samples N", "Continue until every routed expert has N samples.");
opt(fp, c, "--head-test", "Run the output HC/logits head after the native slice.");
opt(fp, c, "--first-token-test", "Run exact CPU whole-model pass for the first prompt token.");
opt(fp, c, "--metal-graph-test", "Compare first GPU-resident graph stages with CPU.");
@@ -269,10 +302,10 @@ static void print_cli_diagnostics(FILE *fp, const help_colors *c) {
static void print_cli_commands(FILE *fp, const help_colors *c) {
title_red(fp, c, "Interactive Commands");
opt(fp, c, "/help", "Show interactive commands.");
opt(fp, c, "/think, /think-max, /nothink", "Switch thinking mode.");
opt(fp, c, "/think [N], /think-max, /nothink", "Switch thinking mode; V4.1 accepts effort 0..100.");
opt(fp, c, "/ctx N", "Restart the interactive session with a new context size.");
opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
opt(fp, c, "/read FILE", "Read FILE and submit it as the next user message.");
opt(fp, c, "/read FILE", "Submit a text file, PNG, or JPEG as the next user message.");
opt(fp, c, "/quit, /exit", "Leave the prompt.");
opt(fp, c, "Ctrl+C", "Stop current generation and return to ds4>.");
fputc('\n', fp);
@@ -281,8 +314,11 @@ static void print_cli_commands(FILE *fp, const help_colors *c) {
static void print_agent_specific(FILE *fp, const help_colors *c) {
title(fp, c, "Agent Options");
opt(fp, c, "-p, --prompt TEXT", "Submit an initial prompt after startup.");
opt(fp, c, "--non-interactive", "Run without TUI. With -p: one turn; without -p: repeated stdin prompts.");
opt(fp, c, "--raw-prompt", "Non-interactive -p only: tokenize prompt without agent chat/tool text.");
opt(fp, c, "--prompt-file FILE", "Read the initial prompt from FILE.");
opt(fp, c, "--prefix-file FILE", "Preload complete alternating USER:/ASSISTANT: turns before the live task.");
opt(fp, c, "--non-interactive", "Run without TUI. With an initial prompt: one turn; otherwise: repeated stdin prompts.");
opt(fp, c, "--raw-prompt", "Non-interactive initial prompt only: omit agent chat/tool text.");
opt(fp, c, "--edit-upto", "Enable anchored [upto] edits and automatic marker insertion.");
opt(fp, c, "-sys, --system TEXT", "Extra system prompt. Empty disables extra text.");
opt(fp, c, "--trace FILE", "Write prompt, token, and DSML debug trace.");
opt(fp, c, "--chdir DIR", "Change working directory before loading runtime assets.");
@@ -298,6 +334,8 @@ static void print_agent_sessions(FILE *fp, const help_colors *c) {
opt(fp, c, "/del ID", "Delete a saved session.");
opt(fp, c, "/strip ID", "Remove KV payload; the text history can be rebuilt later.");
opt(fp, c, "/history [N]", "Show N recent user turns from the current session.");
opt(fp, c, "/hints on|off", "Enable or disable brief programming hints. New and resumed sessions start off.");
opt(fp, c, "/think [N]", "V4.1: set effort 0..100 (default 75) after this turn; rebuilds the prefix.");
opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
opt(fp, c, "/new", "Start a fresh session from the system prompt.");
opt(fp, c, "/quit, /exit", "Exit.");
@@ -310,6 +348,8 @@ static void print_server_api(FILE *fp, const help_colors *c) {
opt(fp, c, "--port N", "Bind port. Default: 8000");
opt(fp, c, "--cors", "Add Access-Control-Allow-* headers for browser JS clients.");
opt(fp, c, "--trace FILE", "Write prompts, cache decisions, output, and tool calls.");
opt(fp, c, "--batched-session N", "Keep N resident sessions and batch decode-ready requests.");
opt(fp, c, "--mixed-prefill-quantum N", "Prefill chunk while generations are active. Default: 128; GLM-5.3 minimum: 1024");
para(fp, c, "Endpoints: /v1/chat/completions, /v1/responses, /v1/completions, and /v1/messages.");
para(fp, c, "Model endpoint aliases include deepseek-v4-flash and deepseek-v4-pro; both serve the loaded GGUF.");
fputc('\n', fp);
@@ -353,6 +393,7 @@ static void print_bench_specific(FILE *fp, const help_colors *c) {
opt(fp, c, "--step-mul F", "Multiplicative step. Default: 1");
opt(fp, c, "--step-incr N", "Linear step when --step-mul is 1. Default: 2048");
opt(fp, c, "--gen-tokens N", "Greedy decode tokens per frontier. 0 for pure prefill. Default: 128");
opt(fp, c, "--teacher-forced-decode", "Decode the following prompt tokens instead of each predicted argmax.");
opt(fp, c, "--csv FILE", "Write CSV there instead of stdout.");
opt(fp, c, "--dump-frontier-logits-dir DIR", "Write one full-logit JSON file per frontier.");
fputc('\n', fp);
@@ -360,9 +401,16 @@ static void print_bench_specific(FILE *fp, const help_colors *c) {
static void print_eval_specific(FILE *fp, const help_colors *c) {
title(fp, c, "Evaluation");
opt(fp, c, "-n, --tokens N", "Max generated tokens per question. Default: 16000");
opt(fp, c, "--questions N", "Run only the first N embedded questions.");
opt(fp, c, "--suite NAME", "core, hard, all, or hard-smoke. Default: core");
opt(fp, c, "--source NAME", "Run only cases from this source.");
opt(fp, c, "--domain NAME", "Run only cases in this domain.");
opt(fp, c, "--case-id ID", "Run the case with this source ID.");
opt(fp, c, "--list-cases", "List selected cases without loading a model.");
opt(fp, c, "--validate-cases", "Validate all embedded cases and exit.");
opt(fp, c, "-n, --tokens N", "Override the generation budget for every question.");
opt(fp, c, "--questions N", "Run only the first N selected questions.");
opt(fp, c, "--case-sequence LIST", "Run 1-based case numbers in this comma-separated order.");
opt(fp, c, "--retry-incomplete", "Retry a missing final answer once with twice the budget.");
opt(fp, c, "--trace FILE", "Write questions, outputs, and grading decisions.");
opt(fp, c, "--regrade-trace FILE", "Regrade a prior trace without loading the model.");
opt(fp, c, "--soft-limit-reply-budget N", "Soft close thinking near the end of reply budget. Default: 1024");
@@ -483,7 +531,7 @@ static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, c
static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) {
if (streq(topic, "all")) {
print_model_runtime(fp, c, tool, true);
if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true);
if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true, tool);
if (tool_has_topic(tool, "steering")) print_steering(fp, c);
print_distributed(fp, c);
if (tool == DS4_HELP_DS4) {
@@ -505,7 +553,7 @@ static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, cons
}
if (streq(topic, "runtime")) print_model_runtime(fp, c, tool, true);
else if (streq(topic, "sampling")) print_sampling(fp, c, true);
else if (streq(topic, "sampling")) print_sampling(fp, c, true, tool);
else if (streq(topic, "steering")) print_steering(fp, c);
else if (streq(topic, "distributed")) print_distributed(fp, c);
else if (tool == DS4_HELP_DS4 && streq(topic, "diagnostics")) print_cli_diagnostics(fp, c);
@@ -519,7 +567,7 @@ static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, cons
title(fp, c, "Agent Tool System");
para(fp, c, "The agent can read, search, write, edit, run bash, and browse through Chrome-backed web tools.");
para(fp, c, "DeepSeek-family models emit DSML tool calls; GLM models use native <tool_call> syntax. Both are rendered live in the terminal.");
para(fp, c, "Edit uses exact old/new replacement; [upto] can bridge a unique head and tail for large anchored edits.");
para(fp, c, "Edit uses exact old/new replacement. --edit-upto enables anchored replacements between a unique head and tail.");
fputc('\n', fp);
} else if (tool == DS4_HELP_BENCH && streq(topic, "benchmark")) print_bench_specific(fp, c);
else if (tool == DS4_HELP_EVAL && streq(topic, "evaluation")) print_eval_specific(fp, c);
@@ -530,7 +578,7 @@ static void print_default(FILE *fp, const help_colors *c, ds4_help_tool tool) {
if (tool == DS4_HELP_DS4) {
print_cli_specific(fp, c, true);
print_sampling(fp, c, false);
print_sampling(fp, c, false, tool);
} else if (tool == DS4_HELP_SERVER) {
print_server_api(fp, c);
print_kv_cache(fp, c);
ds4_image.c
+1204
View File
File diff suppressed because it is too large. Load diff
ds4_image.h
+157
View File
@@ -0,0 +1,157 @@
#ifndef DS4_IMAGE_H
#define DS4_IMAGE_H
#include <stddef.h>
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
#define DS4_IMAGE_MAX_ENCODED_BYTES (64u * 1024u * 1024u)
#define DS4_IMAGE_MAX_DIMENSION 16384u
#define DS4_IMAGE_MAX_PIXELS (64u * 1024u * 1024u)
typedef struct {
uint32_t width;
uint32_t height;
uint8_t *rgb;
uint8_t fingerprint[32];
} ds4_image;
typedef struct {
uint32_t content_width;
uint32_t content_height;
uint32_t padded_width;
uint32_t padded_height;
uint32_t grid_width;
uint32_t grid_height;
uint32_t patch_count;
uint32_t image_token_count;
float *patches;
} ds4_image_patches;
typedef struct {
uint32_t content_width;
uint32_t content_height;
uint32_t padded_width;
uint32_t padded_height;
uint32_t grid_width;
uint32_t grid_height;
uint32_t llm_grid_width;
uint32_t llm_grid_height;
uint32_t patch_count;
float *patches;
} ds4_deepseek4_image_patches;
typedef enum {
DS4_DEEPSEEK4_IMAGE_START = 0,
DS4_DEEPSEEK4_IMAGE_PAD = 1,
DS4_DEEPSEEK4_IMAGE = 2,
DS4_DEEPSEEK4_IMAGE_NEWLINE = 3,
DS4_DEEPSEEK4_IMAGE_END = 4,
} ds4_deepseek4_image_token_type;
typedef struct {
uint32_t token_count;
uint32_t image_count;
uint8_t *types;
uint32_t *perm;
} ds4_deepseek4_image_layout;
int ds4_image_decode_memory(
ds4_image *out,
const uint8_t *encoded,
size_t encoded_len,
char *error,
size_t error_cap);
int ds4_image_decode_file(
ds4_image *out,
const char *path,
char *error,
size_t error_cap);
void ds4_image_free(ds4_image *image);
int ds4_image_preprocess_glm53(
ds4_image_patches *out,
const ds4_image *image,
uint32_t min_image_tokens,
uint32_t max_image_tokens,
char *error,
size_t error_cap);
/* Qwen3-VL style: resize to multiples of 32 within the token budget, normalize
* with mean/std 0.5, emit 3*16*16 patches in 2x2 merge-window order. */
int ds4_image_preprocess_qwen4(
ds4_image_patches *out,
const ds4_image *image,
uint32_t min_image_tokens,
uint32_t max_image_tokens,
char *error,
size_t error_cap);
void ds4_image_patches_free(ds4_image_patches *patches);
int ds4_image_preprocess_deepseek4(
ds4_deepseek4_image_patches *out,
const ds4_image *image,
char *error,
size_t error_cap);
int ds4_image_preprocess_deepseek41(
ds4_deepseek4_image_patches *out, const ds4_image *image,
char *error, size_t error_cap);
int ds4_deepseek41_image_layout_build(
ds4_deepseek4_image_layout *out, uint32_t height, uint32_t width,
char *error, size_t error_cap);
int ds4_deepseek4_image_layout_build(
ds4_deepseek4_image_layout *out,
uint32_t grid_height,
uint32_t grid_width,
uint32_t start_pos,
char *error,
size_t error_cap);
void ds4_deepseek4_image_patches_free(ds4_deepseek4_image_patches *patches);
void ds4_deepseek4_image_layout_free(ds4_deepseek4_image_layout *layout);
/* Find the next DeepSeek image block at or after *cursor. A block includes
* the alignment PAD tokens immediately before IMAGE_START. Returns 1 when a
* block was found, 0 at end of input, and -1 for malformed synthetic tokens. */
int ds4_deepseek4_next_image_span(
const int *tokens,
uint32_t token_count,
uint32_t vocab_size,
uint32_t *cursor,
uint32_t *block_start,
uint32_t *image_start,
uint32_t *image_end);
/* Choose a prefill chunk without splitting a DeepSeek image block. */
int ds4_deepseek4_prefill_chunk(
const int *tokens,
uint32_t token_count,
uint32_t vocab_size,
uint32_t pos0,
uint32_t cap,
uint32_t *chunk);
/* Fill [lo, hi] absolute raw-key bounds for every query row. */
int ds4_deepseek4_attention_bounds(
const int *tokens,
uint32_t token_count,
uint32_t vocab_size,
uint32_t pos0,
uint32_t n_raw,
uint32_t window,
uint32_t *bounds);
#ifdef __cplusplus
}
#endif
#endif
ds4_kvstore.c
+16 -23
View File
@@ -414,6 +414,12 @@ void ds4_kvstore_fill_header(uint8_t h[DS4_KVSTORE_FIXED_HEADER],
kv_le_put64(h + 40, payload_bytes);
}
/* routed-expert quantizations whose checkpoints the store keys on */
bool ds4_kvstore_quant_bits_supported(int quant_bits) {
return quant_bits == 2 || quant_bits == 4 || quant_bits == 5 ||
quant_bits == 6 || quant_bits == 8;
}
bool ds4_kvstore_read_header(FILE *fp, ds4_kvstore_entry *e,
uint32_t *text_bytes) {
uint8_t h[DS4_KVSTORE_FIXED_HEADER];
@@ -436,7 +442,7 @@ bool ds4_kvstore_read_header(FILE *fp, ds4_kvstore_entry *e,
if (fread(tb, 1, sizeof(tb), fp) != sizeof(tb)) return false;
*text_bytes = ds4_kvstore_le_get32(tb);
e->text_bytes = *text_bytes;
return e->tokens != 0 && (e->quant_bits == 2 || e->quant_bits == 4);
return e->tokens != 0 && ds4_kvstore_quant_bits_supported(e->quant_bits);
}
bool ds4_kvstore_read_entry_file(const char *path, const char sha[41],
@@ -940,7 +946,7 @@ bool ds4_kvstore_store_live_prefix_text(ds4_kvstore *kc,
ds4_kvstore_tokens_copy_prefix(&store_tokens, tokens, store_len);
const int quant_bits = ds4_engine_routed_quant_bits(engine);
if (quant_bits != 2 && quant_bits != 4) {
if (!ds4_kvstore_quant_bits_supported(quant_bits)) {
ds4_tokens_free(&store_tokens);
return false;
}
@@ -1224,7 +1230,7 @@ int ds4_kvstore_try_load_text(ds4_kvstore *kc,
if (effective_prompt) effective_prompt->len = 0;
if (!kc->enabled || !prompt_text) return 0;
const int quant_bits = ds4_engine_routed_quant_bits(engine);
if (quant_bits != 2 && quant_bits != 4) return 0;
if (!ds4_kvstore_quant_bits_supported(quant_bits)) return 0;
const int model_id = ds4_engine_model_id(engine);
const size_t prompt_bytes = strlen(prompt_text);
int idx = ds4_kvstore_find_text_prefix(kc, prompt_text, model_id, quant_bits,
@@ -1318,32 +1324,19 @@ int ds4_kvstore_try_load_text(ds4_kvstore *kc,
const double load_ms = (kv_now_sec() - load_t0) * 1000.0;
kc->continued_last_store_tokens = loaded;
const char *key_kind = ds4_kvstore_key_kind(hdr.ext_flags);
bool consumed = false;
if (kc->opt.cold_max_tokens > 0 && loaded > kc->opt.cold_max_tokens) {
unlink(path);
consumed = true;
kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
"%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms consumed file=%s",
kv_log_name(kc),
responses_protocol ? " " : "",
responses_protocol ? "RESPPROTO" : "",
loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
} else {
ds4_kvstore_touch_file(path, hdr.hits + 1);
kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
"%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms file=%s",
kv_log_name(kc),
responses_protocol ? " " : "",
responses_protocol ? "RESPPROTO" : "",
loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
}
ds4_kvstore_touch_file(path, hdr.hits + 1);
kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
"%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms file=%s",
kv_log_name(kc),
responses_protocol ? " " : "",
responses_protocol ? "RESPPROTO" : "",
loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
if (result) {
result->tokens = loaded;
result->text_bytes = text_bytes;
result->quant_bits = hdr.quant_bits;
result->ext_flags = hdr.ext_flags;
result->load_ms = load_ms;
result->consumed = consumed;
result->path = kv_xstrdup(path);
}
}
ds4_kvstore.h
+2 -1
View File
@@ -103,7 +103,6 @@ typedef struct {
uint8_t quant_bits;
uint8_t ext_flags;
double load_ms;
bool consumed;
char *path;
} ds4_kvstore_load_result;
@@ -215,4 +214,6 @@ char *ds4_kvstore_path_for_sha(ds4_kvstore *kc, const char sha[41]);
void ds4_kvstore_le_put32(uint8_t *p, uint32_t v);
uint32_t ds4_kvstore_le_get32(const uint8_t *p);
bool ds4_kvstore_quant_bits_supported(int quant_bits);
#endif
ds4_layer_pack.c
+149
View File
@@ -0,0 +1,149 @@
/* Monotonic-contiguous multi-tier layer placement packer.
*
* Pure C99: no CUDA, no platform-specific code. Used by both the CUDA and
* Metal/CPU builds. See ds4_layer_pack.h for the API contract and the
* design doc docs/superpowers/specs/2026-05-26-multi-gpu-pp-v0-design.md
* for the rationale. */
#include "ds4_layer_pack.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
int ds4_compute_layer_placement(const size_t *entry_bytes,
int n_entries,
const ds4_layer_pack_config *cfg,
int *device_for_entry) {
if (!entry_bytes || !cfg || !device_for_entry) return 1;
if (n_entries < 0) return 2;
if (cfg->n_gpus < 0 || cfg->n_gpus > DS4_LAYER_PACK_MAX_GPUS) return 3;
/* Work over a local copy of budgets so the caller's struct is untouched. */
size_t budget[DS4_LAYER_PACK_MAX_GPUS];
for (int d = 0; d < cfg->n_gpus; d++) {
budget[d] = cfg->gpu_budget_bytes[d];
}
int d = 0;
for (int e = 0; e < n_entries; e++) {
/* Advance until the current device can hold this entry, or we run
* out of devices. Strict greater-than: exact fits stay on d. */
while (d < cfg->n_gpus && entry_bytes[e] > budget[d]) {
d++;
}
if (d < cfg->n_gpus) {
device_for_entry[e] = d;
budget[d] -= entry_bytes[e];
} else {
device_for_entry[e] = DS4_LAYER_PACK_CPU;
/* d stays at cfg->n_gpus so every subsequent entry also lands
* on the CPU tier. */
}
}
return 0;
}
/* Internal helper: append " + embedding" or " + output head" to a buffer if
* the matching pseudo-layer lives on this tier. */
static void append_pseudo_layer_tags(char *buf, size_t buflen,
const int *device_for_entry,
int tier,
int n_entries,
int n_layers) {
/* entry 0 -> embedding, entry n_layers+1 -> output head */
if (n_entries > 0 && device_for_entry[0] == tier) {
strncat(buf, " + embedding", buflen - strlen(buf) - 1);
}
if (n_entries > 1 && n_layers + 1 < n_entries &&
device_for_entry[n_layers + 1] == tier) {
strncat(buf, " + output head", buflen - strlen(buf) - 1);
}
}
/* Emit "L-R" for a contiguous transformer-layer range assigned to a tier.
* Returns 1 if anything was printed (i.e. tier owned at least one layer). */
static int print_layer_range_for_tier(FILE *out, const int *device_for_entry,
int tier, int n_entries, int n_layers) {
/* Transformer layers live at indices 1..n_layers in device_for_entry.
* Their human-facing numbers are 0..n_layers-1 (i.e. entry index - 1). */
int first = -1;
int last = -1;
for (int i = 1; i <= n_layers && i < n_entries; i++) {
if (device_for_entry[i] == tier) {
if (first < 0) first = i - 1; /* human-facing index */
last = i - 1;
}
}
if (first < 0) {
fprintf(out, "(no transformer layers)");
return 0;
}
if (first == last) {
fprintf(out, "layer %d", first);
} else {
fprintf(out, "layers %d-%d", first, last);
}
return 1;
}
void ds4_layer_pack_print(FILE *out,
const int *device_for_entry,
int n_entries,
int n_layers,
const size_t *entry_bytes,
const size_t *gpu_used_bytes,
const size_t *gpu_budget_bytes,
int n_gpus) {
if (!out || !device_for_entry) return;
(void)entry_bytes; /* not needed for the textual layout */
fprintf(out, "multi-GPU layout:\n");
for (int d = 0; d < n_gpus; d++) {
fprintf(out, " GPU%d: ", d);
int have_layers = print_layer_range_for_tier(out, device_for_entry,
d, n_entries, n_layers);
/* Pseudo-layer tags. */
char tags[64];
tags[0] = '\0';
append_pseudo_layer_tags(tags, sizeof(tags), device_for_entry, d,
n_entries, n_layers);
if (tags[0]) {
if (!have_layers) {
/* Erase the "(no transformer layers)" stub and just write
* the tag with leading "+ " stripped. */
/* The stub was already written; appending " + embedding"
* is still informative. */
}
fprintf(out, "%s", tags);
}
/* Usage line. */
if (gpu_used_bytes && gpu_budget_bytes) {
const double used_gb = (double)gpu_used_bytes[d] / 1073741824.0;
const double budget_gb = (double)gpu_budget_bytes[d] / 1073741824.0;
/* Two-space pad before the parens matches the design doc. */
fprintf(out, " (%.1f / %.1f GB)", used_gb, budget_gb);
}
fputc('\n', out);
}
/* CPU tier: only print if at least one entry is on CPU. */
int cpu_present = 0;
for (int i = 0; i < n_entries; i++) {
if (device_for_entry[i] == DS4_LAYER_PACK_CPU) { cpu_present = 1; break; }
}
if (cpu_present) {
fprintf(out, " CPU : ");
int have_layers = print_layer_range_for_tier(out, device_for_entry,
DS4_LAYER_PACK_CPU,
n_entries, n_layers);
char tags[64];
tags[0] = '\0';
append_pseudo_layer_tags(tags, sizeof(tags), device_for_entry,
DS4_LAYER_PACK_CPU, n_entries, n_layers);
if (tags[0]) {
(void)have_layers;
fprintf(out, "%s", tags);
}
fputc('\n', out);
}
}
ds4_layer_pack.h
+77
View File
@@ -0,0 +1,77 @@
#ifndef DS4_LAYER_PACK_H
#define DS4_LAYER_PACK_H
#include <stddef.h>
#include <stdio.h>
#define DS4_LAYER_PACK_MAX_GPUS 16
#define DS4_LAYER_PACK_CPU (-1)
typedef struct {
/* Per-device budget in BYTES (already net of per-device overhead such as
* cuBLAS workspace, global scratch, logits buffer, MTP, steering, and the
* configured safety margin). The packer treats these as raw capacity. */
size_t gpu_budget_bytes[DS4_LAYER_PACK_MAX_GPUS];
int n_gpus;
} ds4_layer_pack_config;
#ifdef __cplusplus
extern "C" {
#endif
/* Compute a monotonic-contiguous layer placement.
*
* Inputs:
* entry_bytes[] — per-entry byte footprint, in FORWARD ORDER:
* entry 0 = embedding pseudo-layer
* entry 1..n_layers = transformer layers
* entry n_layers+1 = output-head pseudo-layer
* n_entries — total number of entries (typically n_layers + 2)
* cfg — per-device budgets
*
* Output:
* device_for_entry[] — filled with target device for each entry; CPU is
* DS4_LAYER_PACK_CPU. Caller-owned, must be
* large enough for n_entries int slots.
*
* Returns:
* 0 on success.
* Nonzero on configuration errors (null pointer, n_entries < 0,
* n_gpus out of range).
*
* Per the design doc, an entry that exceeds every GPU budget spills to CPU;
* by the monotonicity rule every entry after it also goes to CPU. There is
* no error path for "entry too large" — the CPU tier is always available. */
int ds4_compute_layer_placement(const size_t *entry_bytes,
int n_entries,
const ds4_layer_pack_config *cfg,
int *device_for_entry);
/* Print a human-readable layout summary matching the design doc's format:
*
* multi-GPU layout:
* GPU0: layers 0-21 + embedding (38.4 / 40.0 GB)
* GPU1: layers 22-31 (11.7 / 12.0 GB)
* CPU : layers 32-42 + output head
*
* Entry-naming convention:
* entry 0 -> "embedding"
* entry i in 1..n_layers -> transformer layer (numbered i-1)
* entry n_layers+1 -> "output head"
*
* gpu_used_bytes[] and gpu_budget_bytes[] may be NULL — in that case the
* "(used / budget)" line is omitted for the GPU lines. */
void ds4_layer_pack_print(FILE *out,
const int *device_for_entry,
int n_entries,
int n_layers,
const size_t *entry_bytes,
const size_t *gpu_used_bytes,
const size_t *gpu_budget_bytes,
int n_gpus);
#ifdef __cplusplus
}
#endif
#endif /* DS4_LAYER_PACK_H */
ds4_linux_memory.h
+37
View File
@@ -0,0 +1,37 @@
#ifndef DS4_LINUX_MEMORY_H
#define DS4_LINUX_MEMORY_H
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
/* CMA pages can back movable allocations, but not ordinary pinned GPU pages. */
static inline bool ds4_linux_nonmovable_memory_parse(FILE *fp, uint64_t *bytes) {
char line[256];
uint64_t available = 0, cma = 0;
bool found = false;
while (fgets(line, sizeof(line), fp)) {
unsigned long long kib;
if (sscanf(line, "MemAvailable: %llu kB", &kib) == 1) {
if (kib > UINT64_MAX / 1024u) return false;
available = (uint64_t)kib * 1024u;
found = true;
} else if (sscanf(line, "CmaFree: %llu kB", &kib) == 1) {
if (kib > UINT64_MAX / 1024u) return false;
cma = (uint64_t)kib * 1024u;
}
}
if (!found || ferror(fp)) return false;
*bytes = available > cma ? available - cma : 0;
return true;
}
static inline bool ds4_linux_nonmovable_memory(uint64_t *bytes) {
FILE *fp = fopen("/proc/meminfo", "r");
if (!fp) return false;
const bool ok = ds4_linux_nonmovable_memory_parse(fp, bytes);
fclose(fp);
return ok;
}
#endif
Loaded 100 of 2729 files, more files were not shown because too many files have changed in this diff. Show more