diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ea7aeb5..d5f45991 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ # Unreleased +- Use one egglog rewrite engine and calibrated whole-program search instead of live attention/submission retuning. - Broadcast scalar gradients directly and eliminate single-row RoPE at static position zero. - Fix split-attention partial indexing and clear empty splits when reusing a KV cache. - Measure mapped versus staged readback per allocation and size; reuse bounded staging for faster CPU reads. @@ -22,7 +23,7 @@ is a pure function under test, so the config tests don't flip env vars. The six stragglers that still read the environment inline are typed options now: `MEGANEURA_OPTIMIZER` neighbors `GREEDY_PACK_SWIGLU` - (`OptimizeConfig::greedy_pack_swiglu`), `MATMUL_K_STAGE` and + (`OptimizeConfig::pack_swiglu`), `MATMUL_K_STAGE` and `MEGANEURA_INTERLEAVE_COLUMNS` ride `TuningKnobs` into the matmul codegen, and `MEGANEURA_DEVICE_PARAMETERS` / `MEGANEURA_REUSE_UPLOAD` are `SessionOptions` fields. All are registered and documented in the diff --git a/Cargo.toml b/Cargo.toml index 193eabf0..0f204ee3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,9 +79,7 @@ profiler = ["dep:tracing-subscriber"] blade-graphics = { version = "0.9", git = "https://github.com/kvark/blade", rev = "eaff5092096aab136f11fa728b81c1bed3c0dcd4" } blade-macros = "0.3" # `default-features = false` drops egglog's `bin` stack (clap, mimalloc, -# env_logger, chrono, graphviz). Greedy mode never constructs an EGraph -# but the crate is still linked for the optional equality-saturation -# modes. +# env_logger, chrono, graphviz); library builds use the shared rewrite engine. egglog = { version = "2.0", default-features = false } bytemuck = { version = "1", features = ["derive"] } # Match Blade's Naga major: generated modules cross the public API boundary. diff --git a/README.md b/README.md index 80e15c65..a7e3accf 100644 --- a/README.md +++ b/README.md @@ -15,8 +15,8 @@ Define a graph, call `build_session`, train. Meganeura handles autodiff, graph rewrites, WGSL specialization, Naga parsing and validation, and GPU -dispatch automatically. The rewrite engine supports a fast deterministic -greedy mode and experimental equality-saturation modes. +dispatch automatically. Graph rewriting uses bounded equality saturation +through egglog. ```rust use meganeura::{Graph, Trainer, TrainConfig, build_session}; @@ -149,6 +149,11 @@ MEGANEURA_DEVICE_ID=0x744c cargo run --release --example mnist All of the environment variables are resolved in `SessionConfig::from_env()` and never visible to the core modules directly. +For joint graph and implementation search with representative inputs, use +`train::build_measured` with initialization and numerical-qualification callbacks. +It replaces live attention/submission retuning; ordinary `build` does not execute +an uninitialized model. See [compiler search](docs/compiler-search.md). + | Variable | Effect | |---|---| | `MEGANEURA_DISABLE_COOP` | Force the portable scalar matmul path (regression diagnosis). | @@ -162,10 +167,10 @@ All of the environment variables are resolved in `SessionConfig::from_env()` and | `MEGANEURA_PIN_BUFS=3,25-40` | Force-pin logical buffers to bisect aliasing corruption. | | `MEGANEURA_DUMP_PLAN` | Dump dispatch order, provenance, and the alias map at build. | | `MEGANEURA_DUMP_WGSL=` | Write every generated shader into ``. | -| `MEGANEURA_OPTIMIZER` | Rewrite mode: `off` \| `greedy` \| `egglog-windowed` \| `egglog-outlined` \| `egglog-whole`. | +| `MEGANEURA_OPTIMIZER` | Rewrite mode: `off` \| `egglog-windowed` \| `egglog-outlined` (default) \| `egglog-whole`. Legacy `greedy` maps to outlined egglog. | | `MEGANEURA_EGRAPH_COST` | Extraction objective: `ast-size` \| `tensor-traffic`. | | `MEGANEURA_EGRAPH_CUTOFF=` | Saturation segment-size ceiling (default 300). | -| `MEGANEURA_GREEDY_PACK_SWIGLU=0` | Skip packing consecutive SwiGLU ops into one parameter buffer during the greedy sweep. | +| `MEGANEURA_GREEDY_PACK_SWIGLU=0` | Skip packing consecutive SwiGLU ops into one parameter buffer during graph optimization (legacy variable name). | | `MEGANEURA_DEVICE_PARAMETERS` | Experimental placement of unaliased parameter buffers on the device: `1` → device-transient, `device-buddy` → device. Default is host-visible. | | `MEGANEURA_REUSE_UPLOAD` | Reuse one staging buffer across `set_parameter` uploads instead of restaging per parameter. | | `MEGANEURA_TUNE` | Opt-in bounded matmul, convolution and GEMV search at build (`SessionConfig { tune: true }`), using private scratch. Scalar tiles and GEMV shapes also support reduced-storage weights. | diff --git a/bench/gguf_latency.md b/bench/gguf_latency.md index af2076c7..37f14785 100644 --- a/bench/gguf_latency.md +++ b/bench/gguf_latency.md @@ -1,29 +1,34 @@ # Matched Vulkan GGUF diagnostic -`gguf_latency.rs` and `gguf_latency.cpp` compare the same GGUF tensors and -token IDs through Meganeura and llama.cpp. This is a kernel/runtime diagnostic, -not a replacement for the frozen Inferena paper cohort or a language-quality -evaluation. - -Both use one sequence: a 128-token prefill, 32 subsequent cached decode steps, context capacity -256, f32 K/V caches, three warmups and seven measured sequences. Prefill produces -only the last row of logits. Every measured call includes submission, waiting, -and copying the full vocabulary's logits to CPU memory. The first measured -sequence saves 33 logit vectors for comparison. llama.cpp uses Vulkan with all -layers offloaded and flash attention enabled; it refuses a missing Vulkan device. -The helpers intentionally do not tokenize text. - -Decode's token batch is one. Prefill's 128 tokens belong to that same sequence, -not 128 concurrent requests. SmolLM2-135M at batch one exposes host recording -and readback overhead; it is an interactive-latency diagnostic, not batched -throughput evidence. Larger token blocks amortize recording, but increasing -prefill length alone does not establish batched-decode performance. Report -GPU work and whole-call latency separately before generalizing to larger models -or request batches. Command-buffer reuse is not a production optimization goal. - -Build against a recorded llama.cpp checkout: +`gguf_latency.rs` and `gguf_latency.cpp` compare identical GGUF tensors and +fixed token IDs. Both use one sequence, 128-token prefill, 32 cached decode steps, +context capacity 256, F32 K/V caches, three warmup sequences and seven measured +sequences. Prefill returns only the last row of logits. Whole-call times include +recording, submission, waiting and copying the full vocabulary to CPU memory. +The first measured sequence saves 33 logit vectors. + +This is interactive batch-one latency, not batched throughput or language +quality. A 128-token prompt is not 128 concurrent requests. llama.cpp uses +Vulkan with all layers offloaded and flash attention enabled. Meganeura records +fresh commands on every step. This diagnostic does not change the paper cohort. + +## Reproduction + +The measured model is `HuggingFaceTB/SmolLM2-135M` at +`93efa2f097d58c2a74874c7e644dbc9b0cee75a2`, converted to F16 with the official +converter from llama.cpp `05f2dcfdba3879c55f735efa0f124b1a56f7ed11`. +Use that checkout for both conversion and the comparison binary. The resulting +GGUF SHA-256 is +`56df427e4aa9a57d67b207d45e1a2b51e80b956b11cd9fbf0c162429f95168c7`. +The converter's Q/K permutation matters; an earlier hand-converted file was +incorrect and its timings were discarded. ```sh +hf download HuggingFaceTB/SmolLM2-135M \ + --revision 93efa2f097d58c2a74874c7e644dbc9b0cee75a2 --local-dir hf-smollm2 +# Install the pinned llama.cpp converter's requirements in a separate venv. +python ../llama.cpp/convert_hf_to_gguf.py hf-smollm2 \ + --outtype f16 --outfile model.gguf cargo build --release --features gguf --example gguf_latency cmake -S ../llama.cpp -B ../llama.cpp/build -DGGML_VULKAN=ON cmake --build ../llama.cpp/build -j2 @@ -32,264 +37,103 @@ c++ -O2 -std=c++17 bench/gguf_latency.cpp \ -L ../llama.cpp/build/bin -Wl,-rpath,"$PWD/../llama.cpp/build/bin" \ -lllama -lggml -lggml-base -o /tmp/llama-latency -# Select exactly one GPU's ICD. For Intel use intel_icd.json instead. +# Use intel_icd.json for Intel; select exactly one GPU. export VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json export GGML_VK_VISIBLE_DEVICES=0 MEGANEURA_COOP_F16=1 target/release/examples/gguf_latency model.gguf /tmp/meg 30 /tmp/llama-latency model.gguf /tmp/llama ``` -The optional final argument bounds tuning to 30 seconds **per session**; -the diagnostic allows 256 MiB of scratch so the vocabulary projection is not -silently excluded by the default 64 MiB limit. Reports record candidates, -qualification, selections, preparation time and CPU record/wait/readback stages. -No precision tolerances are changed. Run engines and GPUs sequentially, without -profilers or builds competing with the timed runs. Repeat in fresh processes -and reverse their order. First-read mapped/staged qualification is absorbed -by warmup, not included in steady-state latency. - -For separate attribution captures, add `MEGANEURA_GPU_TIMING=1` to the Meganeura -command or `GGML_VK_PERF_LOGGER=1` to llama.cpp. Profiled numbers are not substitute -latencies: instrumentation changes execution, and Meganeura's pass intervals -include the transition to the next pass. In particular, wall time minus their -sum is **not** a measurement of barrier cost. - -## September 19, 2026 checkpoint - -Source: Meganeura `4195be1b869cbb8b9eaf4b0403c3f48875d1e675`, based on -`da0e28424395904ae9a8cb53e7da32e15012ec1a`. llama.cpp: -`05f2dcfdba3879c55f735efa0f124b1a56f7ed11`. The original unfinished experiments -are preserved separately at `experiment/llama-prototypes-2026-09-17` -(`78e2ff448030fdbd7932cd06c76ae282bcef6542`), not in this branch's ancestry. - -Model: `HuggingFaceTB/SmolLM2-135M` at -`93efa2f097d58c2a74874c7e644dbc9b0cee75a2`, converted with the recorded -llama.cpp revision's official converter: - -```sh -hf download HuggingFaceTB/SmolLM2-135M \ - --revision 93efa2f097d58c2a74874c7e644dbc9b0cee75a2 --local-dir hf-smollm2 -# Install llama.cpp's converter requirements in a separate venv first. -python ../llama.cpp/convert_hf_to_gguf.py hf-smollm2 \ - --outtype f16 --outfile smollm2-official-f16.gguf -``` - -GGUF SHA-256: -`56df427e4aa9a57d67b207d45e1a2b51e80b956b11cd9fbf0c162429f95168c7`. -All 272 tensors were checked against the HF source after the converter's Q/K -permutation and storage conversion. An inherited local GGUF lacked that -permutation; its preliminary results were discarded, not mixed into this table. - -The split-attention combine on the upstream base multiplied a partial-row -stride twice. Correcting that and clearing empty partials restored all 33 -next-token predictions against llama.cpp on both GPUs. The regression exercises -multiple heads, query rows, cache reuse and a split context; the submitted paper's -older Inferena path is not this newly added GGUF path. - -An independent CPU f32 eager run of the pinned HF model produced the same 33 -next-token choices. Maximum per-row relative L2 logit errors against it were -0.0000093 (Meganeura) and 0.01121 (llama.cpp) on the 5070, and 0.000175 and -0.01314 on the B570. Matching weight storage and f32 caches does not guarantee -matching intermediate arithmetic. Top-1 agreement on these fixed tokens is a -sanity check, not a language-quality evaluation. - -### Unprofiled latency - -Milliseconds; median of three fresh-process medians, each with seven prefill -and 224 decode observations. Process order was rotated. The control is -`d0bb9669bab48942a3ab873e38911090e27e9002`, including the necessary attention -fix; timing incorrect upstream outputs would not be a useful baseline. - -| GPU | Phase | Corrected control | This branch | llama.cpp | -| --- | --- | ---: | ---: | ---: | -| RTX 5070 | Prefill, 128 tokens | 13.47 | 11.57 | 7.12 | -| RTX 5070 | Decode, per token | 5.83 | 4.35 | 1.38 | -| Arc B570 | Prefill, 128 tokens | 86.39 | 27.62 | 15.65 | -| Arc B570 | Decode, per token | 22.81 | 6.18 | 3.45 | - -This is not parity: the remaining decode ratios are 3.16x and 1.79x. -The B570 control was variable (61.26--86.64 ms prefill, 15.17--22.85 ms decode); -the new branch was steadier (27.62--27.64 and 6.16--6.18 ms). CPU/GPU clocks were -not fixed. Both engines used the same six physical i5-12400F cores, pinned with -`taskset -c 0,2,4,6,8,10`; the CPU retained its powersave governor. NVIDIA driver: -595.91.07. The secondary B570's reported link was PCIe 2.5 GT/s x1, so its -host-readback numbers must not be generalized to a full-bandwidth installation. - -The branch's median preparation times, including model loading, both sessions -and tuning, were 3.45 s (5070) and 9.79 s (B570). llama.cpp took 0.26 and 0.50 s. -Neither Meganeura session exhausted its 30 s tuning budget. The control used -the previous 64 MiB tuning scratch cap, which omitted the vocabulary GEMV; -the branch allowed 256 MiB. These are process starts with existing driver caches, -not cache-cold compilation measurements. - -### What changed, and what remains - -Mapped VRAM was the main avoidable host-read cost. The runtime now measures direct -and staged reads, checks bit-exact agreement, caches the choice per allocation -and byte count, and reuses at most 16 MiB of download staging. Dense F16 and packed -weights can also enter the existing scalar tile search with unchanged decoding -and numerical qualification. - -The 5070's median CPU-side readback stage fell from 3.20 to 0.37 ms. The B570's -fell from 13.78 to 0.51 ms. The whole-call improvements are smaller than these -differences suggest: CPU record/submit rose from 0.37 to 1.71 ms on NVIDIA and -0.62 to 1.36 ms on Intel, while Intel's wait stage also changed. These overlapping -CPU/GPU stages under variable clocks are not an additive causal gap budget. - -Separate instrumented Meganeura captures locate the remaining GPU work: - -| GPU / phase | Matrix pass intervals | Cached-attention intervals | Other intervals | +Run engines and GPUs sequentially, without competing builds or profilers. +Repeat in fresh processes and reverse engine order. The optional budget is a soft +construction-time limit in seconds per session; `0` disables measured search. +It includes compilation, initialization, qualification and measurements. In-flight +driver calls cannot be preempted. Private kernel probes retain their full +qualification and the diagnostic's original 256 MiB scratch limit. + +`build_measured` compares graph forms, kernel choices and submission chunks in one +complete-plan search. Each candidate checks every logit and cache element against +an untuned reference (`atol=1e-5`, `rtol=1e-4`), with full staged readback. +Decode calibration primes that cache through the middle of the decode range +(position 144). Trials reset private mutable state; +selected prefill/decode sessions share caches only after search. Reference +preparation and priming are outside the search budget but inside `prepare_ms`. +The old scope argument and separate live submission tuner are removed. Historical +measurements below retain their original protocol and revision. + +Known staged output downloads are queued before the CPU wait. Initial probes +and mapped reads still wait first. `decode_record_finish_ms` reports CPU +record/submit and combined wait/copy intervals, not separate GPU costs. +For separate captures, use `MEGANEURA_GPU_TIMING=1` or llama.cpp's +`GGML_VK_PERF_LOGGER=1`. Instrumentation changes execution; wall time minus +summed pass intervals is not a measurement of barrier cost. + +## Calibrated-construction checkpoint, 2026-09-20 + +Exact `789bcb60e82d8a2993a440acb500bf53b797cbeb` versus +`b06629cb19467ac00e75f3b383bd3ef70442d2a7`, two fresh processes per arm and +GPU, reversing the second pair. Same model, inputs, precision and host/device +setup as below; GPUs and builds ran sequentially. These are diagnostics, not a +new publication cohort or a new llama.cpp comparison. + +Median of process medians; execution in milliseconds, preparation in seconds: + +| GPU | Prefill before / after | Decode before / after | Preparation before / after | | --- | ---: | ---: | ---: | -| 5070 / decode | 51% | 30% | 19% | -| 5070 / prefill | 57% | 35% | 8% | -| B570 / decode | 39% | 31% | 30% | - -Decode still issues 394 dispatches in 274 dependency groups. Cached attention -uses 30 split/combine pairs. llama.cpp's own Vulkan timestamp logger reports -about 0.18 ms for its 30 decode attention operations on NVIDIA, versus about -0.98 ms in Meganeura's attention pass intervals. The instrumentation differs, -so those numbers identify a kernel target, not a directly subtractable barrier -cost. An Nsight Systems 2025.5.2 Vulkan/OS-runtime capture of Meganeura completed. -The llama.cpp injection exited abnormally; that capture was excluded from claims. -Uninstrumented runs and llama.cpp's own timestamp logger completed on both GPUs. - -Cooperative matrices remain a concrete implementation gap. This B570 advertises -8x16x16 f16-input/f32-accumulator tiles; Blade's capability filter and pinned -Naga WGSL frontend accept only square 8x8 or 16x16 tiles. llama.cpp reports -`KHR_coopmat` there, and `NV_coopmat2` on the 5070. Meganeura also excludes -reduced-storage weights from its cooperative matmul path. Scalar tile tuning -does not resolve either limitation. Cooperative coverage is a prefill target; -decode additionally needs better cached attention, GEMV and host submission. -Sharing Vulkan alone does not make these kernel implementations equivalent. - -A tilewise softmax-rescaling trial is retained at -`experiment/cached-attention-tile-softmax-2026-09-19` -(`bd25688a585ee6e093ea0beb8ae98edfe56fca88`), not included in the production -branch. It passed the numerical checks but gave only a small NVIDIA prefill -gain (11.64 to 11.17 ms) and no Intel gain in three reversed-order pairs. - -### Fixed-head attention follow-up - -`dd9f56f9f9d334f791ce64ae7b0163d6ed631058` compiles cached attention at the graph's -known head dimension, using the existing per-head pipeline mechanism. It removes -unused per-thread values and dynamic head-width branches without changing the -algorithm or precision. -Three fresh-process pairs, reversing order, compared it with `5aeb856`: - -| GPU | Phase | Before | Fixed head width | -| --- | --- | ---: | ---: | -| RTX 5070 | Prefill | 11.58 | 8.82 | -| RTX 5070 | Decode | 4.29 | 3.20 | -| Arc B570 | Prefill | 27.63 | 22.79 | -| Arc B570 | Decode | 6.09 | 5.66 | - -Units and sampling are unchanged. All 33 token predictions match the independent -CPU reference on both devices; maximum per-row relative L2 error remains below -0.000010 on NVIDIA and 0.000178 on Intel. The existing cache regression now also -covers head widths 64, 80 and 512, including partially occupied thread lanes. -NVIDIA attention pass intervals fell from 0.98 to 0.49 ms for decode and from -3.80 to 1.66 ms for prefill in separate captures. This is still not parity with -llama.cpp; matrix kernels and host submission remain substantial costs. - -### Serialized command replay: attribution mini-study, not production - -The experiment at Meganeura `6672ef385d79e73accae5cc7c7d2eefca8ff50fd` -pins Blade `4befad4f5fbd427c1aca4b1b001fb8b7acd8e109` and reuses an unchanged -Vulkan inference recording after waiting for its previous execution. -Input contents remain dynamic. Rebinding, tuning, profiling, changing submission -chunks, and other uses of the encoder invalidate the recording. Training, -timestamped sessions, Metal and GLES keep ordinary recording. The normal -`Session::step` path in that experiment selects this automatically; no numerical -change is involved. Production keeps ordinary command recording. The replay -API and its invalidation machinery were rejected as a poor fit for Blade and -for the project's intended workloads, despite helping this small decode case. - -Three fresh-process trials per engine, with rotated order and the same diagnostic: - -| GPU | Phase | Fixed head, recording | Fixed head, replay | llama.cpp | -| --- | --- | ---: | ---: | ---: | -| RTX 5070 | Prefill | 8.65 | 7.23 | 7.11 | -| RTX 5070 | Decode | 3.13 | 1.82 | 1.38 | -| Arc B570 | Prefill | 22.89 | 21.31 | 15.61 | -| Arc B570 | Decode | 5.64 | 4.48 | 3.42 | - -The NVIDIA recording control varied from 2.28 to 3.15 ms per decode; replay -ranged from 1.81 to 1.85 ms. Intel replay ranged from 4.45 to 4.54 ms. CPU clocks -remain uncontrolled. The CPU record/submit stage fell from 1.26 to 0.040 ms on -NVIDIA and from 1.42 to 0.044 ms on Intel. Replay does not remove GPU barriers -or change kernel arithmetic. All 33 next-token choices still match the CPU -reference; maximum relative L2 errors are 0.0000102 and 0.000177 respectively. -Median preparation remains 3.46 s and 9.79 s, including tuning. - -The replay arm gets NVIDIA prefill within 2% of llama.cpp on this case, but -those are not production results. The ordinary-recording column remains the -relevant baseline: 8.65/3.13 ms prefill/decode on NVIDIA and 22.89/5.64 ms on -Intel. Replay improves CPU recording, not GPU kernels, and does not establish -an improvement in the paper's GPU-resident training workloads. Keep the frozen -paper cohort unchanged until kernel work and measurements on its own graphs -justify recollection. - -The complete source is preserved on Meganeura -`experiment/serialized-replay-2026-09-19` (the revision above) and Blade -`perf/compute-command-replay` (the pinned revision above). The timing-only -precursor remains on both repositories' `experiment/llama-replay-2026-09-19`. -CPU-stage instrumentation is on Meganeura -`experiment/llama-cpu-record-2026-09-19` (`29e3cd4`). No binary artifacts are needed. - -### Further kernel probes - -`b6d6971cce673c559fbcbcc162b5fd5a8313379d` adds fused transposed matmuls to -the existing qualified tile search. It changes neither precision policy nor -the search budget. Existing CPU reference/layout checks and the opt-in GPU -tuning regression cover both transpose directions and nonzero addends. - -Three source-only probes are separate from production: - -- `experiment/rmsnorm-subgroup-2026-09-19` (`04421dd`): using subgroup sums - for the fused norm as well as the matrix reduction gave less than 1% decode - improvement, smaller than process variation. Not retained. -- `experiment/attention-split-counts-2026-09-19` (`a3fbd89`): a pilot over - 1/2/4/8/16 splits found different preferred geometry for prefill and decode. - It does not justify a replacement fixed threshold. The intended production - route is a bounded, numerically qualified search over complete split/combine - sequences, including scratch and combine cost, keyed by shape and device. - This search is not implemented yet; the existing matrix tuner does not select - attention split counts. Cache position and sliding-window lengths also vary - at runtime, so a winner must not be chosen from one unrepresentative length. -- `experiment/gguf-row-weights-2026-09-19` (`621de44`): retaining dense GGUF - rows and using transposed GEMV reduced decode from 1.814 to 1.609 ms on - NVIDIA and 4.549 to 4.284 ms on Intel in three reversed-order pairs. - Prefill stayed near 7.27 and 21.2 ms after extending tile search; without - that fix, Intel prefill took 27.3 ms in a pilot. All saved logits passed - the same reference checks. The layout still loses existing norm/packing - fusions and adds temporary GEMV-plus-add dispatches, so it remains an - experiment, not part of the production latency table above. - -All three probes above used serialized replay. Their improvements must be -rechecked with ordinary recording before making production latency claims. - -Each experiment branch contains its source and a short result summary, without -raw timings or binaries. The next layout work is to preserve the applicable -fusions, not to add model-specific kernels or change the paper cohort. - -### Verification - -Formatting and all-target/all-feature Clippy passed. CPU unit tests: 449 passed, -three ignored. Readback bit preservation, cached-attention reference checks and -the opt-in reduced-storage tile-tuning regression passed on both GPUs. The -existing cache regression now includes nonuniform scores, sliding windows and -long-to-short cache reuse rather than adding another test binary. - -NVIDIA's GGUF-feature regression/smoke run passed 183/81 tests, with 13 ignored. -On this driver, repeated context destruction/recreation eventually fails with -`ERROR_INCOMPATIBLE_DRIVER`, including on unchanged main. Keeping the ICD loaded -with a process-local `LD_PRELOAD=libGLX_nvidia.so.0` made the suite pass. This is -a test workaround, not a production global-context cache or a profiling setting. -Intel's broader model-feature run passed 184 regression and 94 smoke tests; -one pairwise-distance gradient comparison failed by 2.1e-7 and reproduced on -unchanged main with the same dependency lock. Its tolerance was not weakened. - -Raw timings, logits, profiles and binaries stay outside Git. Do not update the -paper or launch another full cohort on the strength of this one model and shape. -In particular, the frozen Inferena cohort times GPU-resident work; these -host-readback savings do not automatically improve that timed window. +| RTX 5070 | 7.251 / 7.404 | 1.425 / 1.463 | 15.40 / 37.77 | +| Arc B570 | 14.060 / 14.062 | 3.871 / 3.786 | 31.75 / 63.76 | + +NVIDIA is about 2–3% slower in these aggregate medians. Decode process medians +range from 1.392–1.458 before to 1.460–1.466 after. Intel's ranges are +3.865–3.877 before and 3.541–4.032 after; the apparent aggregate improvement is +not a stable gain. No performance-neutrality claim is made. + +Preparation is not measured under an identical policy: the old diagnostic gave +each session 30 seconds of isolated kernel tuning plus a separate submission +search. The new 30-second construction budget includes allocation, initialization, +full output/state qualification and whole-program comparisons. Reference priming +is additional. Full staged readback avoids repeated mapped-read probes without +dropping checked values. Rebuilding candidates remains costly, particularly on +Intel; avoiding that cost is unfinished work. + +All eight full-logit sets are finite and match all 33 independent CPU-reference +token choices. Maximum per-row relative L2 is 1.33e-5 on NVIDIA and 1.77e-4 on +Intel. Some NVIDIA prefill alternatives fail the stricter whole-output gate and +are discarded. The selected programs retain that gate. + +This GGUF graph's outlined region contains cache writes, so logical alternative +extraction is explicitly skipped. This check exercises the shared rewrite engine, +kernel choices, pre-allocation attention choices and submission search; it does +not demonstrate a graph-alternative win. NVIDIA exhausts the generated choices; +Intel reaches the time bound. See [the design and remaining limits](../docs/compiler-search.md). + +## Earlier measured checkpoint and experiment history + +The source-only branch `experiment/llama-catchup-pre-cleanup-2026-09-20` at +`137f93364afe06cfff89c888efddb6bb05c54d33` preserves the implementations, +intermediate reports and links to isolated experiments. It contains no raw +data or binaries. In particular, [queued readback](https://github.com/kvark/meganeura/blob/137f93364afe06cfff89c888efddb6bb05c54d33/bench/queued-readback.md) +records the final pre-cleanup checkpoint and [scalar-layout tuning](https://github.com/kvark/meganeura/blob/137f93364afe06cfff89c888efddb6bb05c54d33/bench/scalar-matmul-autotune.md) +isolates the 29% Intel prefill improvement. + +Milliseconds, median of three fresh-process medians, rotated engine order: + +| GPU / phase | Meganeura | llama.cpp | +| --- | ---: | ---: | +| RTX 5070 prefill | 7.210 | 7.126 | +| RTX 5070 decode/token | 1.389 | 1.361 | +| Arc B570 prefill | 14.048 | 15.650 | +| Arc B570 decode/token | 3.666 | 3.431 | + +These are measurements of `7c95707b9a69478c1bd498f940b7dde6e851a7e1`, not a +new cohort on the cleanup branch. All saved logits are finite and retain the +33 independent CPU-reference token choices. Maximum per-row relative L2: +0.00000954/0.000178 for Meganeura, 0.01121/0.01314 for llama.cpp. Preparation +takes 9.61/31.45 seconds versus 0.25/0.51 seconds, with existing driver caches. + +The i5-12400F used physical cores `0,2,4,6,8,10`, unfixed clocks, NVIDIA +595.91.07 and Intel Mesa 26.0.3. B570 was on a secondary PCIe x1 link, limiting +readback generalization. Decode process medians vary with measured split/chunk +choices; this is close performance, not stable parity. Reusable command +recording remains an attribution experiment, not a production path. Known +Naga Workgroup ArrayStride diagnostics remain unresolved. diff --git a/bench/gguf_latency.rs b/bench/gguf_latency.rs index b9e7b80a..cf31ccd9 100644 --- a/bench/gguf_latency.rs +++ b/bench/gguf_latency.rs @@ -13,13 +13,13 @@ const DECODE: usize = 32; const CONTEXT: usize = 256; const SAMPLES: usize = 7; -fn run( +fn set_inputs( session: &mut Session, model: &gguf::GgufModel, config: &gguf::arch::ModelConfig, position: usize, count: usize, -) -> (Vec, [f64; 3]) { +) { let tokens: Vec = (position..position + count) .map(|i| 42 + (i % 31) as u32) .collect(); @@ -38,25 +38,36 @@ fn run( .expect("per-layer embedding gather"); session.set_input("ple", &ple); } +} + +fn run( + session: &mut Session, + model: &gguf::GgufModel, + config: &gguf::arch::ModelConfig, + position: usize, + count: usize, +) -> (Vec, [f64; 2]) { + set_inputs(session, model, config, position, count); let start = Instant::now(); session.step(); let submitted = Instant::now(); - session.wait(); - let finished = Instant::now(); let mut logits = vec![0.0; config.vocab_size]; - session.read_output_by_index(0, &mut logits); + session.wait_read_output(0, &mut logits); assert!(logits.iter().all(|x| x.is_finite())); let read = Instant::now(); ( logits, [ submitted.duration_since(start).as_secs_f64() * 1000.0, - finished.duration_since(submitted).as_secs_f64() * 1000.0, - read.duration_since(finished).as_secs_f64() * 1000.0, + read.duration_since(submitted).as_secs_f64() * 1000.0, ], ) } +fn read_outputs(session: &Session) -> Vec> { + session.read_buffers(&session.plan().output_buffers) +} + fn main() -> Result<(), Box> { env_logger::init(); let args: Vec<_> = std::env::args().collect(); @@ -71,7 +82,7 @@ fn main() -> Result<(), Box> { let f32_activations = std::env::var_os("MEGANEURA_F32_ACTIVATIONS").is_some(); let mut sessions = Vec::new(); let mut tuning = Vec::new(); - for block in [1, PROMPT] { + for block in [PROMPT, 1] { let mut graph = Graph::new(); let built = gguf::graph::build(&mut graph, &model, &config, block, CONTEXT)?; graph.set_outputs(built.outputs()); @@ -84,27 +95,103 @@ fn main() -> Result<(), Box> { cfg.options.quantized_activations = false; } let mut session = meganeura::build(&graph, cfg).0; - session.set_submission_chunks(1); + gguf::weights::load(&mut session, &model, &config)?; + gguf::weights::reset_caches(&mut session, &built, &config); + let position = if block == 1 { PROMPT + DECODE / 2 } else { 0 }; + let mut initial_caches = Vec::new(); + if let Some(prefill) = sessions.first_mut() { + run(prefill, &model, &config, 0, PROMPT); + for (name, _) in &prefill.plan().param_buffers { + if name.starts_with("cache.") { + let values = prefill.read_params(&[name])[0].clone(); + session.set_parameter(name, &values); + initial_caches.push((name.clone(), values)); + } + } + } if tune_seconds != 0 { - tuning.push(session.tune_with(meganeura::tune::TuneOptions { - max_time: Duration::from_secs(tune_seconds), - max_classes: 64, - max_scratch_bytes: 256 * 1024 * 1024, - ..Default::default() - })?); + if block == 1 { + for pos in PROMPT..position { + run(&mut session, &model, &config, pos, 1); + } + for (name, values) in &mut initial_caches { + *values = session.read_params(&[name])[0].clone(); + } + } + run(&mut session, &model, &config, position, block); + let expected = read_outputs(&session); + let mut cfg = SessionConfig::inference_from_env_on(session.context()); + cfg.tune = false; + cfg.cache = None; + if f32_activations { + cfg.options.quantized_activations = false; + } + drop(session); + let (selected, report) = meganeura::train::build_measured( + &graph, + cfg, + meganeura::train::BuildSearchOptions { + max_time: Duration::from_secs(tune_seconds), + max_plan_bytes: 4 << 30, + tuning: meganeura::TuneOptions { + max_classes: 64, + max_scratch_bytes: 256 << 20, + min_improvement: 0.01, + ..Default::default() + }, + ..Default::default() + }, + |s, donor| { + if let Some(source) = donor.or_else(|| sessions.first_mut()).filter(|source| { + s.plan().param_buffers.iter().all(|(name, _)| { + name.starts_with("cache.") || source.has_parameter(name) + }) + }) { + for (name, _) in s.plan().param_buffers.clone() { + if !name.starts_with("cache.") { + s.share_parameter_from(source, &name) + .map_err(|e| e.to_string())?; + } + } + } else { + gguf::weights::load(s, &model, &config).map_err(|e| e.to_string())?; + } + gguf::weights::reset_caches(s, &built, &config); + for (name, values) in &initial_caches { + s.set_parameter(name, values); + } + set_inputs(s, &model, &config, position, block); + Ok(()) + }, + |s| { + let actual = read_outputs(s); + for (index, (a, b)) in actual.iter().zip(&expected).enumerate() { + if a.len() != b.len() + || a.iter().zip(b).any(|(&a, &b)| { + !a.is_finite() || (a - b).abs() > 1e-5 + 1e-4 * b.abs() + }) + { + return Err(format!( + "full output/cache {index} differs from untuned reference" + )); + } + } + Ok(()) + }, + )?; + session = selected; + tuning.push(report); } - if let Some(decode) = sessions.first_mut() { + if let Some(prefill) = sessions.first_mut() { for (name, _) in session.plan().param_buffers.clone() { - if decode.has_parameter(&name) { - session.share_parameter_from(decode, &name)?; + if prefill.has_parameter(&name) { + session.share_parameter_from(prefill, &name)?; } } - } else { - gguf::weights::load(&mut session, &model, &config)?; - gguf::weights::reset_caches(&mut session, &built, &config); } sessions.push(session); } + sessions.swap(0, 1); let prepare_ms = started.elapsed().as_secs_f64() * 1000.0; let mut prefill_ms = Vec::new(); let mut decode_ms = Vec::new(); @@ -143,7 +230,7 @@ fn main() -> Result<(), Box> { "vocab": config.vocab_size, "prepare_ms": prepare_ms, "activations": if f32_activations { "f32" } else { "q8_1" }, "prefill_ms": prefill_ms, "decode_ms": decode_ms, - "decode_record_wait_read_ms": decode_parts_ms, + "decode_record_finish_ms": decode_parts_ms, "dispatches": [sessions[1].plan().dispatches.len(), sessions[0].plan().dispatches.len()], "tuning": tuning, }); diff --git a/bench/optimizer_ablation.rs b/bench/optimizer_ablation.rs index 1e5b2fe6..dd33a0c8 100644 --- a/bench/optimizer_ablation.rs +++ b/bench/optimizer_ablation.rs @@ -23,7 +23,6 @@ enum Phase { fn parse_mode(value: &str) -> OptimizeMode { match value { "off" => OptimizeMode::Off, - "greedy" => OptimizeMode::Greedy, "egglog-windowed" | "windowed" => OptimizeMode::EgglogWindowed, "egglog-outlined" | "outlined" => OptimizeMode::EgglogOutlined, "egglog-whole" | "whole" => OptimizeMode::EgglogWhole, diff --git a/docs/compiler-search.md b/docs/compiler-search.md new file mode 100644 index 00000000..cebc5887e --- /dev/null +++ b/docs/compiler-search.md @@ -0,0 +1,311 @@ +# Compiler search: direction and alternatives + +Architecture discussion and implementation notes, 2026-09-20, for PR #200. +The submitted P3HPC results use their recorded revisions and are unaffected by +this design work. + +## Decision + +Use egglog to retain equivalent structures and a bounded measurement loop to +choose their implementations. Keep the Rust/WGSL/Blade backend. Replace the +competing selection paths rather than adding another optimizer above them. + +The immediate goal is to search the implementations we can already generate. +Generating fundamentally new kernels from lower-level primitives is a separate +project. Neither a larger search space nor equality saturation guarantees a +faster program under a finite compilation budget. + +## Construction-time search + +For a matrix product followed by an epilogue, keep fused, unfused and split-K +forms available together. Each form has a legal domain of scalar/cooperative +implementations, tile sizes and reduction choices. Do not first select a graph +by estimated memory traffic and only then tune that graph's kernels. + +1. Retain equivalent forms using one set of rewrite rules. Shapes, precision + constraints and effects delimit legal transformations. +2. Extract bounded, diverse implementation families. Leave numerical schedule + parameters symbolic until needed; do not expand their whole Cartesian product. +3. Lower candidates through the normal compiler, scheduler and allocator. +4. Initialize private representative inputs/state, qualify, then measure. +5. Return the qualified incumbent when the budget ends. Application state must + not have advanced during search. + +These are distinct responsibilities, not competing optimizers. Cheap and +thorough construction should differ in budget, not duplicate transformation +implementations. Identical lowered programs and repeated implementation classes +can reuse work. Semantic equivalence alone is not performance equivalence. + +Structural changes belong before allocation. A live session should not need +buffer growth, alias-map edits, dispatch insertion or profiling-index remapping +to install a tuned structure. Rebuilding explicitly is an acceptable tradeoff. +Existing layout-preserving kernel probes remain useful shared infrastructure. + +Measured construction needs initialized inputs. An ordinary build must not +silently execute an uninitialized model. Calibrated construction makes its input +and state contract explicit; ordinary construction still uses the same rewrite +and lowering machinery. Synthetic isolated-kernel probes are not a substitute +for whole-program qualification on representative inputs. + +## Bounds, correctness and observability + +- Budget compilation, allocation, initialization, qualification and measurements, + not just GPU kernel time. Deadlines are soft around in-flight driver calls. +- Preserve a legal incumbent. An incomplete or invalid comparison cannot win. + Measurement noise and unrepresentative inputs still limit performance claims. +- Retain different physical interfaces, such as layouts, until their consumers + and conversion costs have been considered. One cheapest logical expression is + not necessarily the cheapest complete execution. +- Treat outputs, gradients and persistent updates as observable. Reset private + state between trials; effects cannot be commuted as if they were pure values. +- Numerical tests complement transformation legality. A close result on a sample + is not a universal equivalence proof, and real-arithmetic identities need not + preserve a floating-point policy. +- Use compact choice/measurement reports, including skipped and unexplored work. + Do not require serialization of runtime dispatch objects for search itself. +- Confirm selections end to end. Isolated timings omit interactions with memory, + synchronization and CPU submission. Do not infer barrier cost by subtraction. + +Start with dense projection families and cached attention, including shared +subgraphs and stateful validation on NVIDIA and Intel. Assess construction time, +memory, execution time and deleted machinery. Broad invariant/regression tests +and inspectable plans are preferable to multiplying timing-sensitive tests. + +## API and current boundary + +`train::build_measured(graph, config, options, initialize, qualify)` is the +calibrated entry point. `BuildSearchOptions` bounds graph forms, complete plans, +time and declared plan/state bytes. Its `tuning` field configures the existing +private kernel probes and paired noise guard. `BuildSearchReport` records the +extracted expressions, trial descriptions, phase times, rejections, selected +trial and unfinished work. The ordinary `build` API remains input-free. + +The initializer supplies representative inputs and weights to a new private +session. The qualifier reads all observable outputs, gradients and state updates +after exactly one step and checks the application's numerical contract. The +runner resets persistent writes between steps and before returning the winner. +Configure runtime optimizers and external/shared writable bindings afterward. +An invalid challenger is discarded; an invalid incumbent aborts the search. +An initializer may share compatible immutable weights from the idle incumbent. + +Implemented choices include fused/unfused graph forms, dispatch fusion, +cached-attention splits and fresh submission chunk counts. Existing +layout-preserving kernel choices are tuned before +whole-plan comparison; completed comparisons can be reused within that build. +Attention splits use the existing compiler, with no session-buffer patching. +There is no reusable command recording or separate live structural tuner. + +This is a bounded first implementation, not exhaustive graph scheduling. Small +graphs are searched together; large graphs currently explore one verified +repeated region. Stateful or mixed-precision regions remain opaque. Physical +settings are currently applied uniformly to eligible operations within a plan, +not independently to every site. The report must be read with those limits. +There is no persistent measured-plan cache yet. Dense split-K and alternate row +reductions remain prototypes on the experiment branch, not additions to this +refactoring. The broader alternatives below +remain research directions, not additional dependencies or hidden search paths. + +## Refactor checkpoint + +Relative to PR #200 at `789bcb6`, this removes roughly 370 lines of non-test +compiler/runtime/shader source. The separate greedy rewrite implementation and +the live attention/submission tuner modules are gone. Kernel generation and +pipeline preparation remain shared with ordinary execution. + +This is not a large total-line-count reduction: the documentation, benchmark +adapter and broad invariant tests more than offset that saving. The main benefit +is removing live allocation/alias/dispatch repair, not reducing every subsystem +to fewer lines. Rebuilding candidate sessions also costs more than private +isolated probes; construction time must be reported alongside execution time. + +Precision domains are opaque cut edges for rewriting. Preserving them also keeps +forward/backward shared-output associations intact, including cross-entropy's +logits gradient. The existing end-to-end gradient tests caught this integration +issue when egglog became the default; no numerical tolerance was relaxed. + +## Why not retain greedy rewriting indefinitely? + +The historical small, mostly locally profitable rule set did not demonstrate an +execution advantage for equality saturation. The recorded SmolLM ablation used +0.089 ms for greedy rewriting, 2.94 ms for outlined egglog and 56.2 ms for +whole-graph inference saturation; whole differentiated-graph saturation took +7.43 s. These are historical observations, not measurements of the proposed +joint search. See the ablation in `paper/main.tex`. + +Equality saturation preserves represented alternatives, but resource limits can +prevent discovering a derivation. Extraction optimizes a supplied cost model, +not actual execution time automatically. Tree extraction can also overcharge +shared work in a DAG. Keeping two separately implemented rule sets adds drift +without resolving any of these problems. + +## CPU overhead and review follow-up + +The September 20 review found avoidable work in our integration, not just in +egglog. Ordinary extraction rebuilt the same cost table for every escaping +root. It now computes one extractor and shares one term DAG per saturated +segment. Node bindings use direct e-graph lookup rather than evaluating new +expressions. The segment report is passed as one object instead of eleven +independent arguments. + +Alternative extraction uses an explicit postorder stack, term-ID memoization +and directly interned integer literals. It no longer creates an AST to recover +literal values. Exclusion lookup borrows function names and argument slices, +without allocating an edge on every cost query. The queue, visited set and +extractor share immutable exclusion sets. Candidates are deduplicated by term +ID in a shared DAG; expression strings are generated only for retained reports. +Egglog's public terms and function lookup still use constructor names; replacing +those with a second local operator registry would add another mapping to maintain. + +CPU-only release measurements on zork's i5-12400F, pinned to physical CPU 0, +compare `dd743f4` with this follow-up. Each cell is the median of five warm +invocations after one discarded invocation, in a fresh process per model/arm. +No GPU context, shader compilation or GPU tuning is involved. Training totals +include automatic differentiation and the ablation harness's graph copies; +they are not complete deployment preparation times. + +| Workload | Inference before / after (ms) | Training before / after (ms) | +|---|---:|---:| +| SmolLM2-135M | 4.92 / 4.31 | 734.19 / 91.85 | +| SmolVLA | 6.61 / 5.98 | 387.24 / 51.54 | +| StableDiffusion | 28.05 / 17.65 | 521.42 / 88.39 | +| ResNet-50 | 11.61 / 10.27 | 541.29 / 450.10 | +| Whisper-tiny | 6.01 / 5.49 | 175.81 / 125.51 | + +The bounded four-form repeated-region search separately takes 4.19 / 3.57 ms +for SmolLM2, 4.24 / 3.59 ms for SmolVLA, and 2.54 / 2.27 ms for Whisper. +Whisper has one represented candidate here; the other two hit the four-form +bound. This measures enumeration, not the subsequent physical-program search. +The same small timing harness was applied to the baseline. Ordinary graph +node/fusion counts, e-graph sizes and extraction-failure counts match in all +ten cases. Peak process RSS does not increase (8.5--325.3 MiB after the change). + +Reproduce with `cargo run --release --features models --example +optimizer_ablation -- --model SmolVLA --phase training --repeats 6`, discarding +the first returned sample. The opt-in CPU search probe is `cargo test --release +--features models --lib cpu_search_overhead -- --ignored --nocapture`. +Neither adds a default CI timing test or another test executable. + +The remaining ResNet/Whisper preparation cost is mostly outside the reported +egglog/stamping intervals, which total about 40/22 ms in the final training +samples. Differentiation and graph-copy costs need separate attribution before +changing egglog again. Linux sampling was unavailable (`perf_event_paranoid=4`); +these are elapsed-time experiments, not a claim about sampled CPU hotspots. + +## Next cohort gate + +Do not start a distributed cohort merely by merging this PR. Inferena at +`fa5a04e1` still calls ordinary `build` and then `tune_with`, before model inputs +and weights are initialized. A pin update alone would miss joint structural +search. Its `hub` feature and direct `Dispatch.use_coop` field access also need +the current API spellings. + +First adapt the runner to initialized, qualified `build_measured` sessions; +retain graph/plan coverage and skipped-region receipts, explicit total budgets, +and peak-memory bounds for two candidate sessions. Then run all five models in +both contracts on the local NVIDIA and Intel devices, check outputs/gradients, +and compare held-out latency and preparation costs before freezing one pin. +The previous GGUF check still has a small NVIDIA regression, variable Intel +decode timing and higher preparation cost; these CPU improvements do not +establish a GPU speedup or resolve calibration stability. + +Platforms without a usable PyTorch GPU path belong in the separate +[qualification workflow](../paper/p3hpc/QUALIFICATION.md), never a CPU/GPU speed +comparison. RPL-U has retained qualification evidence; Mendocino is pending. + +## Alternatives for a later session + +The judgments below concern fit for Meganeura, not a ranking of published +speedups across different workloads, devices, precision policies and budgets. + +### Halide + +Separates computation from scheduling: tiling, producer placement, storage, +recomputation and parallelism. Its GPU autoscheduler uses hierarchical sampling +and memoization. This is a strong route toward composing kernels instead of +writing another fused template. It does not alone supply all alternative +algebraic formulations. Adoption would require translating our operations and +integrating another compiler with our context/buffer ownership. + +Borrow scheduling concepts if we broaden code generation. Do not adopt a new +compiler solely to simplify selection among existing WGSL kernels. Adams2019 +is CPU-only; Anderson2021 is the relevant full-GPU autoscheduler. + +- [GPU autoscheduler paper](https://arxiv.org/abs/2012.07145) +- [Official autoscheduler integration](https://halide-lang.org/docs/HalideCMakePackage.html) + +### Ansor + +Generates coarse implementation sketches, samples complete configurations and +improves them with evolutionary search and a learned cost model. It allocates +tuning effort across subgraphs. This helps separate structural and numerical +choices without a complete hand-written template per operator. + +The original pipeline partitions graphs before scheduling; broader joint graph +optimization was left as future work. Learned models and evolutionary search +also have cold-start costs, and hardware intrinsics still need derivation rules. +Borrow sketch generation and budget allocation, not necessarily the TVM stack. + +- [Ansor, OSDI 2020](https://www.usenix.org/system/files/osdi20-zheng.pdf) + +### MetaSchedule + +TVM separates schedule rules, candidate generation, search strategies, +builders/runners and a tuning database. This is a useful engineering reference +for reusable legality and measurement, and for compact reproducible decisions. +It cannot recover graph alternatives discarded before workload extraction. + +Direct adoption needs TVM integration; recreating all its extension points in +Rust could itself become bloat. Borrow the separation of responsibilities and +decision records. Initially use one concrete policy, not a plugin framework. + +- [MetaSchedule architecture](https://tvm.apache.org/docs/deep_dive/tensor_ir/tutorials/meta_schedule.html) + +### Cascades + +Maintains memoized groups of equivalent expressions and interleaves logical and +physical optimization on demand under required physical properties. This maps +naturally to implementation families and layout-sensitive consumers: retain the +best alternative per required property rather than one universal winner. + +It is the strongest alternative to egglog for our immediate scope. However, +implementing matching, memoization and search ourselves risks rebuilding similar +infrastructure. GPU costs are not cleanly compositional, and adding context for +sharing, allocation and neighboring work can enlarge the search state. Revisit +if bounded egglog matching/extraction remains the dominant problem. + +- [Cascades framework, 1995](https://15721.courses.cs.cmu.edu/spring2019/papers/22-optimizer1/graefe-ieee1995.pdf) + +### Mirage + +Represents programs across kernels, thread blocks and threads, searching both +algebraic and scheduling transformations, including new kernel structures. Its +restricted-domain probabilistic equivalence checking is supplemented with +floating-point tests. This can expose implementations absent from our templates. + +Adopting its scope means substantially more compiler and verification work. +The published NVIDIA evaluation does not establish transfer to our Vulkan/Metal +backend. Its mathematical guarantees do not cover arbitrary stateful training or +all floating-point behavior. Use it as an offline research reference first. + +- [Mirage, OSDI 2025](https://www.usenix.org/system/files/osdi25-wu-mengdi.pdf) + +### EquiForge + +This September 11, 2026 preprint puts tensor expressions, tiled computation, +reductions and storage alternatives into one e-graph. It extracts implementation +families, prunes redundant partial candidates and progressively tunes schedules. +It is the closest research blueprint to retaining structure through measurement. + +Its evaluation allows four hours per configuration, uses Triton and tests NVIDIA +A100/RTX 5090. Algebraic rules assume real arithmetic; numerical validation is +separate. Neither short startup nor our precision policy is solved by adopting +the design. Borrow representation and deduplication ideas selectively. + +- [EquiForge preprint](https://arxiv.org/html/2609.12330v1) + +### Supporting references + +- [Egglog ruleset scheduling](https://egraphs-good.github.io/egglog-tutorial/04-scheduling.html) +- [Egglog extraction and cost models](https://egraphs-good.github.io/egglog-tutorial/05-cost-model-and-extraction.html) +- [TENSAT: shared-DAG extraction](https://proceedings.mlsys.org/paper_files/paper/2021/file/cc427d934a7f6c0663e5923f49eba531-Paper.pdf) +- [Guided equality saturation](https://arxiv.org/abs/2111.13040) diff --git a/examples/diagnose_fusion.rs b/examples/diagnose_fusion.rs index f16b7ed5..a3eafb28 100644 --- a/examples/diagnose_fusion.rs +++ b/examples/diagnose_fusion.rs @@ -34,7 +34,7 @@ fn shader_name(s: &ShaderEntry) -> String { /// epilogue absorption on the scalar plan inspected here. fn is_scalar_matmul(d: &Dispatch) -> bool { use ShaderEntry::*; - !d.use_coop + !d.use_coop() && matches!( d.shader, MatMul | MatMulAT | MatMulBT | FusedMatMulAdd | FusedMatMulATAdd | FusedMatMulBTAdd @@ -71,9 +71,6 @@ fn count_consumers(plan: &ExecutionPlan) -> HashMap> { for b in &d.input_buffers { m.entry(*b).or_default().push(i); } - for b in &d.epilogue_buffers { - m.entry(*b).or_default().push(i); - } } m } @@ -142,7 +139,7 @@ fn diagnose(plan: &ExecutionPlan) -> Vec { if matches!( d.shader, ShaderEntry::MatMul | ShaderEntry::MatMulBT | ShaderEntry::MatMulAT - ) && !d.use_coop + ) && !d.use_coop() && d.input_buffers.len() >= 2 { for (slot_idx, in_buf) in d.input_buffers[..2].iter().enumerate() { @@ -164,7 +161,7 @@ fn diagnose(plan: &ExecutionPlan) -> Vec { // Pattern 3: MatMul → (single consumer) Add/BiasAdd with a matmul-fused variant // already existing (FusedMatMulAdd). Count cases where this is being done in // two dispatches. - if matches!(d.shader, ShaderEntry::Add | ShaderEntry::BiasAdd) && d.pointwise.is_none() { + if matches!(d.shader, ShaderEntry::Add | ShaderEntry::BiasAdd) && d.pointwise().is_none() { for (slot_idx, in_buf) in d.input_buffers.iter().enumerate() { if !external.contains(in_buf) && let Some(&prod_i) = producer.get(in_buf) diff --git a/examples/gpu_compare.rs b/examples/gpu_compare.rs index 8df4037c..c2125311 100644 --- a/examples/gpu_compare.rs +++ b/examples/gpu_compare.rs @@ -26,9 +26,9 @@ fn bench_matmul(n: usize, warmup: usize, iters: usize) -> (f64, &'static str) { .iter() .find(|d| matches!(d.shader, meganeura::compile::ShaderEntry::MatMul)) .map(|d| { - if d.use_coop { + if d.use_coop() { "coop" - } else if d.use_small_tiles { + } else if d.use_small_tiles() { "small" } else { "tile" diff --git a/examples/matmul_throughput.rs b/examples/matmul_throughput.rs index a9170b50..e909c9ba 100644 --- a/examples/matmul_throughput.rs +++ b/examples/matmul_throughput.rs @@ -98,13 +98,13 @@ fn bench_shape( .map(|d| { let kernel = match d.shader { meganeura::compile::ShaderEntry::MatMulGemv => "gemv", - _ if d.use_coop => "coop", - _ if d.use_small_tiles => "small", + _ if d.use_coop() => "coop", + _ if d.use_small_tiles() => "small", _ => "tile", }; ( - if d.use_coop { 1 } else { 0 }, - if d.use_small_tiles { 1 } else { 0 }, + if d.use_coop() { 1 } else { 0 }, + if d.use_small_tiles() { 1 } else { 0 }, d.workgroups[0] * d.workgroups[1] * d.workgroups[2], kernel, ) diff --git a/paper/p3hpc/QUALIFICATION.md b/paper/p3hpc/QUALIFICATION.md new file mode 100644 index 00000000..3ea5d971 --- /dev/null +++ b/paper/p3hpc/QUALIFICATION.md @@ -0,0 +1,77 @@ +# Graphics-only platform qualification + +These runs establish that Meganeura executes and meets the numerical gates on +a GPU without a usable PyTorch GPU path. CPU PyTorch is a correctness oracle, +not a performance competitor. Do not publish CPU/GPU timing ratios or include +these runs in GPU performance, preparation-time or portability-score aggregates. + +The existing collector already supports this; no new harness is needed. +Commands below use Inferena's `experiment/p3hpc-cuda-graphs` branch. Its current +engine pin is the submitted cohort's `428fc2d2`, **not PR #200**. A run now is +bring-up evidence at that pin. Repeat qualification at the next frozen pin +before describing it as coverage of the new engine. + +## AMD Mendocino on rubik + +From the Inferena checkout, with Rust, uv and a working RADV Vulkan driver: + +```sh +git switch experiment/p3hpc-cuda-graphs +git pull --ff-only +bash scripts/setup.sh cpu .venv-p3hpc-cpu +cargo run --release --locked -p inferena-meganeura -- --list-devices +.venv-p3hpc-cpu/bin/python scripts/p3hpc.py \ + --backend cpu --gpu MENDOCINO --qualify-only --eager +cp ../inferena-results/latest.tgz ../inferena-results/rubik-mendocino-qualification.tgz +``` + +Reuse `.venv-p3hpc-cpu` if already installed; setup refuses to overwrite it. +The device list must show the intended hardware as available and not software +emulated. `--gpu MENDOCINO` matches a RADV name containing that substring +(including `RAPHAEL_MENDOCINO`). If the list uses a different name, use that +exact name instead. If more than one matches, stop and disambiguate; do not +drop the filter. The collector selects its device ID and checks the executed +device in every result. It refuses software-renderer or ambiguous matches. + +For multiple Mesa devices with identical names, first inspect +`MESA_VK_DEVICE_SELECT=list vulkaninfo`. Mesa's +[device-selection layer](https://docs.mesa3d.org/envvars.html#vulkan-mesa-device-select-layer-environment-variables) +can restrict enumeration with `MESA_VK_DEVICE_SELECT=vendor:device!`, using +the IDs actually listed on rubik. Do not guess a PCI ID from the product name. + +`--eager` avoids spending time compiling the CPU oracle. It does not disable +Meganeura tuning, change precision, or relax forward/backward checks. Both +arithmetic contracts and all five models run, with one fresh qualification +process per condition rather than three 20-sample measurement processes. +No ROCm installation, architecture override or `--no-max-autotune` is needed +for this CPU-oracle/Vulkan path. Integrated GPUs remain subject to ordinary +device-memory and numerical checks; incomplete runs are not passes. + +## Intel RPL-U + +Use the same CPU environment and replace the collection command with: + +```sh +.venv-p3hpc-cpu/bin/python scripts/p3hpc.py \ + --backend cpu --gpu RPL-U --qualify-only --eager +cp ../inferena-results/latest.tgz ../inferena-results/intel-rpl-u-qualification.tgz +``` + +The submitted archives already establish ten passing workload/precision +conditions, each in three processes, on RPL-U. The paper retains this evidence +in the qualification table and no longer reports the CPU-reference timings. + +## Evidence and archive lifetime + +Require `campaign.json` to finish with `status: complete`, `args.collect: false`, +the intended `native_device`, and ten valid qualification runs. Preserve the +archive, driver/hardware identity and exact source pin. A CPU-only wheel alone +does not prove that a GPU PyTorch path is unavailable: retain the GPU-backend +probe failure or vendor support evidence separately. Mendocino is still pending, +not a passing or unavailable platform inferred merely from this command. + +Each attempt creates a new timestamped directory. `latest.tgz` is atomically +replaced on **completion or failure**, so copy it before another attempt or +platform run. A failed archive is diagnostic evidence, not a qualification pass. +The timestamped directories remain available; never run two collectors sharing +the same `latest.tgz` destination at once. diff --git a/paper/p3hpc/RESULTS.md b/paper/p3hpc/RESULTS.md index 02f08155..b55075a1 100644 --- a/paper/p3hpc/RESULTS.md +++ b/paper/p3hpc/RESULTS.md @@ -1,6 +1,6 @@ # Final P3HPC cohort: evidence guide -Updated September 14, 2026. The manuscript now uses the completed v9 cohort, +Updated September 20, 2026. The manuscript uses the completed v9 cohort, not the superseded v7 tables. No new timing was collected, condition changed, or outlier removed during analysis. @@ -26,14 +26,22 @@ preserves the measured engine revision, separately from the manuscript branch. | Radeon 780M | 30 / 30 | Same, with recorded ROCm overrides | | Arc B570 | 30 / 30 | Default compilation + XPU graph; math SDPA and embedding workaround | | Apple M3, macOS | 30 / 30 | Compiled MPS; no equivalent public whole-phase replay API | -| Intel RPL-U | 30 / 30 | Vulkan versus explicitly selected compiled CPU | | H100 360M/1.7B extension | 12 / 12 | Both contracts, three processes each, CUDA Graph | -Eight main campaigns contribute 240 pairs; the extension adds 12. There are -no interrupted campaigns or failed pairs. Seven configurations have GPU -references; RPL-U stays separate. RTX 5070 and B570 share a host and were +Seven main GPU-reference campaigns contribute 210 pairs; the extension adds +12. There are no interrupted campaigns or failed pairs. RTX 5070 and B570 share a host and were measured sequentially; B570 uses the secondary PCIe 3.0 x1 link. +| Qualified platform without a usable PyTorch GPU path | Native path | Qualified workload/precision conditions | +|---|---|---:| +| Intel RPL-U | Vulkan, ANV Mesa 26.0.3 | 10 / 10, each in three processes | + +The 30 retained RPL-U processes supply numerical qualification evidence, using +CPU PyTorch only as an oracle. Their timings are excluded from all current +performance, preparation and search tables and the generated condition CSV. +Original records remain unchanged. AMD Mendocino on `rubik` is a prospective +addition, not yet a qualified result; see [qualification instructions](QUALIFICATION.md). + Every pair passes the frozen Inferena checker and an independent audit of raw/joined agreement, timing medians, numerical errors, executed policies, full replay statistics and replication. All nine replication reports agree @@ -92,7 +100,6 @@ Consequently this is a cross-cohort comparison, not a controlled ablation. | Radeon 780M | 1.01 / 1.00 / 1.02 | | Arc B570 | 1.27 / 1.03 / 1.17 | | Apple M3 | 1.09 / 1.00 / 0.92 | -| Intel RPL-U | 1.11 / 1.03 / 1.19 | Each entry is the median of five old/new native time ratios, not a ratio of aggregate times. Particularly useful improvements are strict ResNet training: @@ -125,11 +132,11 @@ them in all five strict workloads (2,150 dispatch instances over 15 processes). No measured Vulkan device exposes native-f32 tiles through this stack; NVIDIA gains therefore cannot be attributed to strict cooperative f32. -Of 708 sessions, 684 visit every eligible class (including zero-class -sessions); 13,322 / 14,085 class instances are visited. The 24 truncated -sessions are ResNet training on 780M, B570, M3 and RPL-U. The longest search -is 60.441 seconds. Candidate comparisons record 13,417 FasterCandidate, -31,456 KeepBaseline, 330 InvalidOutput and 24 TimeBudget decisions. +Of 624 sessions in paired GPU campaigns, 606 visit every eligible class +(including zero-class sessions); 11,870 / 12,243 class instances are visited. +The 18 truncated sessions are ResNet training on 780M, B570 and M3. The longest +search is 60.112 seconds. Candidate comparisons record 12,290 FasterCandidate, +28,053 KeepBaseline, 330 InvalidOutput and 18 TimeBudget decisions. Numerically rejected candidates are not installed. Counts include repeated processes and candidate comparisons, not distinct kernels or graph speedups. GEMV, reduced-input cooperative variants and arbitrary graph representations @@ -178,7 +185,7 @@ claimed. [Diagnostic source and analysis](https://github.com/kvark/inferena/blob The MI300X report remains separately identified (SHA-256 e9be97e695140e8a36e09da0f0dd850113b0b22359182921dbc34b689d9776e8). -Its driver experiments are not timing cells. RPL-U CPU, XPU workarounds, +Its driver experiments are not timing cells. RPL-U qualification, XPU workarounds, and the 780M ROCm overrides remain explicit portability evidence; old Windows/H100 failures are not attributed to the complete final cohort. diff --git a/paper/p3hpc/artifact/README.md b/paper/p3hpc/artifact/README.md index c2e1d329..80af6c53 100644 --- a/paper/p3hpc/artifact/README.md +++ b/paper/p3hpc/artifact/README.md @@ -2,8 +2,10 @@ ## Camera-ready cohort -The September 14 final analysis uses v9: eight complete 30-pair device -campaigns and a complete 12-pair H100 extension, 252 valid pairs total. +The v9 archives contain eight complete 30-pair device campaigns and a complete +12-pair H100 extension. The September 20 presentation separates 222 GPU-reference +pairs from 30 RPL-U qualification processes. CPU-reference timings are not +reported; the original records remain intact for numerical/provenance checks. Inferena fa5a04e1 and Meganeura 428fc2d2 are the measured revisions. [RESULTS.md](../RESULTS.md) records scope, findings and interpretation; [cohort.sha256](cohort.sha256) identifies the nine original archives. @@ -16,8 +18,10 @@ From the repository root, with Python 3.11+: The audit checks source/checkpoint identity, raw/joined agreement, 20-sample medians, cross-engine numerical gates, fixed per-tensor/whole-gradient replay bounds, native search/precision policies, compiled reference/replay -receipts and all three-process gradient reports. Eight LaTeX table/figure -fragments and an 84-row condition CSV are regenerated without a GPU/network. +receipts and all three-process gradient reports. Nine LaTeX table/figure +fragments and a 74-row GPU-reference condition CSV are regenerated without a +GPU/network. A separate qualification table lists RPL-U's ten passing +workload/precision conditions without CPU-versus-GPU timing ratios. The supplementary ZIP is self-contained for this audit. Its records.jsonl.xz is a lossless compression of every original JSON value, including raw/joined diff --git a/paper/p3hpc/artifact/cohort.py b/paper/p3hpc/artifact/cohort.py index 8099b0de..c199b39e 100644 --- a/paper/p3hpc/artifact/cohort.py +++ b/paper/p3hpc/artifact/cohort.py @@ -28,7 +28,7 @@ "amd-igpu": "Radeon 780M", "intel-b570": "Arc B570", "apple-m3": "Apple M3", - "intel-igpu": "Intel RPL-U (CPU ref.)", + "intel-igpu": "Intel RPL-U", "nvidia-h100-large": "H100 extension", } ENGINES = ("meganeura", "pytorch") @@ -313,8 +313,10 @@ def ratio(value): def tables(campaigns, rows, groups): output = {} + paired_devices = {device: label for device, label in DEVICES.items() + if campaigns[device]["args"]["backend"] != "cpu"} lines = [] - for device, label in DEVICES.items(): + for device, label in paired_devices.items(): if device == "nvidia-h100-large": continue c = campaigns[device] @@ -327,12 +329,18 @@ def tables(campaigns, rows, groups): output["devices.tex"] = tex_table("llllr", "Device & PyTorch path & Graphics driver & Explicit replay & Valid/selected pairs", lines) lines = [] for device, label in DEVICES.items(): + if device in paired_devices: + continue + conditions = len(groups[device]) + lines.append([label, "Vulkan", f"{conditions}/{conditions}", "Unavailable"]) + output["qualification.tex"] = tex_table("llrl", + "Device & Meganeura path & Qualified conditions & PyTorch GPU", lines) + lines = [] + for device, label in paired_devices.items(): c = campaigns[device] if device == "nvidia-h100-large": continue condition = primary_condition(c) - if device == "intel-igpu": - lines.append([r"\multicolumn{8}{l}{\emph{GPU-versus-CPU support comparison; excluded from GPU aggregates}}"]) for model in MODELS: values = [ratio(rows[device][precision, model, condition]["ratio_" + phase]) for precision in ("strict", "accelerated") for phase in PHASES] @@ -342,7 +350,7 @@ def tables(campaigns, rows, groups): r" & & Inf. & Min. & F+L+B & Inf. & Min. & F+L+B") output["ratios.tex"] = tex_table("llrrrrrr", ratio_heading, lines) lines = [] - for device, label in DEVICES.items(): + for device, label in paired_devices.items(): runs = [run for batch in groups[device].values() for run in batch] compile_s = [sum(run["pair"][engine]["timings"]["compile_s"] for run in runs) for engine in ENGINES] graph_s = [sum(phase.get(part + "_s", 0) for run in runs @@ -352,7 +360,7 @@ def tables(campaigns, rows, groups): output["preparation.tex"] = tex_table("lrrrrr", r"Device & Pairs & M compile+tune & P compile & P graph prep. & P qualification", lines) lines = [] - for device, label in DEVICES.items(): + for device, label in paired_devices.items(): searches = [session["search"] for batch in groups[device].values() for run in batch for session in run["pair"]["meganeura"]["optimizer"]["sessions"]] decisions = sum((Counter(s["decisions"]) for s in searches), Counter()) @@ -458,7 +466,8 @@ def main(): for name, digest in c["sha256"].items(): name = name.replace("\\", "/") require(input_hashes.setdefault(name, digest) == digest, "input identity differs: " + name) - campaigns[device], rows[device], all_groups[device] = c, aggregate(groups), groups + campaigns[device], all_groups[device] = c, groups + rows[device] = aggregate(groups) if c["args"]["backend"] != "cpu" else {} if bundle: bundle.write(json.dumps(records, separators=(",", ":")) + "\n") print(device, c["status"], dict(Counter(run["status"] for run in c["runs"])), "failed:", failed) @@ -488,6 +497,16 @@ def main(): errors = [(run["errors"], device, key) for device, groups in all_groups.items() for key, runs in groups.items() for run in runs] for key in errors[0][0]: print("maximum", key, max((e[key], device, group) for e, device, group in errors)) + searches = [session["search"] for device, groups in all_groups.items() + if campaigns[device]["args"]["backend"] != "cpu" + for runs in groups.values() for run in runs + for session in run["pair"]["meganeura"]["optimizer"]["sessions"]] + print("GPU-paired native search:", len(searches), "sessions;", + sum(s["visited_classes"] == s["eligible_classes"] for s in searches), "complete;", + sum(s["visited_classes"] for s in searches), "/", + sum(s["eligible_classes"] for s in searches), "classes; longest", + max(s["elapsed_seconds"] for s in searches), "seconds") + print("GPU-paired decisions:", sum((Counter(s["decisions"]) for s in searches), Counter())) if __name__ == "__main__": diff --git a/paper/p3hpc/main.tex b/paper/p3hpc/main.tex index 6e15ee4e..85801c98 100644 --- a/paper/p3hpc/main.tex +++ b/paper/p3hpc/main.tex @@ -43,16 +43,16 @@ units (GPUs). Meganeura is a compact native compiler and runtime that lowers tensor graphs and automatic differentiation to Vulkan and Metal kernels. We evaluate five workloads -on eight device configurations, including an NVIDIA H100 and Windows, with one PyTorch +on seven GPU-reference configurations, including an NVIDIA H100 and Windows, with one PyTorch source revision, compiled references, and independently qualified CUDA, HIP, and XPU graph replay wherever those APIs are available. Meganeura always performs bounded kernel search during compilation, and strict arithmetic permits native-f32 cooperative matrices. -All 252 paired process runs pass the numerical gates, including replicated +All 222 paired process runs pass the numerical gates, including replicated 360-million- and 1.7-billion-parameter H100 extensions. Across seven complete GPU-reference configurations, median strict 32-bit floating-point inference and training times are $1.83\times$ and -$2.47\times$ PyTorch's. Search covers every eligible class in 684 of 708 +$2.47\times$ PyTorch's. Search covers every eligible class in 606 of 624 native sessions; preparation cost and performance gains remain workload-dependent. Meganeura executes on an Intel integrated GPU without a usable PyTorch GPU path; a separate AMD MI300X attempt @@ -92,7 +92,7 @@ \section{Introduction} The evaluation compares \system{} with PyTorch \cite{paszke2019pytorch} at one source commit across CUDA, ROCm, XPU, MPS, -and CPU wheels. Every reference uses default compilation. CUDA, HIP, and +and CPU-oracle wheels. Every timed reference uses default compilation. CUDA, HIP, and XPU execute explicitly captured, numerically qualified whole-phase graphs; MPS has no equivalent public replay interface. Native empirical selection is always enabled, independently of the reference's compiler mode. @@ -126,8 +126,8 @@ \section{Introduction} An arithmetic contract (Section~\ref{sec:contracts}) specifies permitted numerical fast paths. Each device--workload--contract--policy comparison -is validated before its timings are admitted. GPU-versus-CPU support -comparisons remain visible without entering the GPU performance aggregate. +is validated before its timings are admitted. Graphics-only qualification +remains visible without entering the GPU performance aggregate. This study includes datacenter hardware and up to 1.7-billion-parameter models, but not production-scale large language model (LLM) training, optimizer updates, distributed execution, or scientific simulation workloads. @@ -252,7 +252,7 @@ \subsection{Revisions and reference conditions} All wheels report PyTorch 2.13.0 at source commit \code{cf30153}, with Python 3.13.13. NVIDIA uses \code{+cu130}, AMD \code{+rocm7.2}, Arc \code{+xpu}, Apple the MPS -wheel, and the RPL-U comparison \code{+cpu}. +wheel. RPL-U qualification uses \code{+cpu} as a correctness oracle only. Checkpoint revisions and hashes are recorded in the manifests. Driver, host, and backend-library differences remain part of the platform; a common PyTorch revision does not make those implementations identical. @@ -285,7 +285,7 @@ \subsection{Revisions and reference conditions} scaled-dot-product-attention (SDPA) backend to accommodate graph capture; the setting and a qualified embedding-backward workaround are recorded (Section~\ref{sec:availability}). Other backends retain automatic SDPA selection. -MPS and the explicitly selected CPU reference also compile all requested +MPS also compiles all requested phases. MPS's Inductor path is exercised \cite{pytorch2026mpssource}, but the pinned MPS API offers no equivalent public whole-phase replay interface. MPSGraph execution is not CUDA Graph capture. No failed compiled or @@ -387,11 +387,11 @@ \subsection{Machines and coverage} \input{tables/devices} \end{table*} -Table~\ref{tab:devices} lists eight device configurations; the seven complete -GPU-reference configurations are RTX~5070, H100, RTX~3050, RX~7900~XT, +Table~\ref{tab:devices} lists the seven complete +GPU-reference configurations: RTX~5070, H100, RTX~3050, RX~7900~XT, Radeon~780M, Arc~B570, and Apple~M3. -The complete RPL-U campaign compares Vulkan with compiled CPU and is -reported separately. RTX~3050 runs natively on Windows with driver 591.86; +RPL-U has no usable PyTorch GPU path and appears only in the qualification +table (Table~\ref{tab:qualification}). RTX~3050 runs natively on Windows with driver 591.86; the other Vulkan campaigns run on Linux, and M3 uses macOS. RTX~5070 and B570 share the same host and are measured sequentially, not concurrently. B570 uses a secondary PCIe~3.0 $\times$1 link; @@ -399,9 +399,11 @@ \subsection{Machines and coverage} Valid/selected counts whole paired processes across both arithmetic contracts under the common default-compilation policy (Section~\ref{sec:reference-policy}). -Eight completed campaigns contribute 240 pairs, including all 30 +Seven completed GPU-reference campaigns contribute 210 pairs, including all 30 selected Windows pairs. The H100 extension contributes 12 valid -pairs: 252 valid pairs in total, with no interrupted campaign. +pairs: 222 valid GPU pairs in total, with no interrupted campaign. +The 30 retained RPL-U processes establish numerical qualification only; +their CPU-reference timings are not reported or aggregated. All forward and backward pairwise gates replay from the retained evidence. \section{Results} @@ -420,8 +422,8 @@ \section{Results} contracts. Each group of three columns is full inference, minimal-shape latency, and F+L+B. Ratios below one are bold; a rounded 1.00 is a near tie, not evidence of a significant win. CUDA, HIP, and XPU use qualified replay; -MPS and CPU are compiled without explicit replay. -The CPU support block is excluded from GPU counts and scores. +MPS is compiled without explicit replay. Graphics-only qualification +has its own table and contributes no timing ratios. Absolute times and process ranges are retained in the supplementary CSV. Figure~\ref{fig:smollm2} visualizes the SmolLM2 subset in milliseconds. @@ -489,11 +491,11 @@ \subsection{What compile-time search actually explores} \end{table*} Table~\ref{tab:search} reports search coverage. -Of 708 sessions, 684 visit every eligible class, including sessions -with no eligible class. In total, 13,322 of 14,085 class instances are -visited. The 24 deadline-limited sessions are ResNet training on 780M, -B570, M3, and RPL-U; all other sessions finish their declared search. -The longest search lasts 60.44\,s against the soft 60-second ceiling. +Of 624 sessions in the paired GPU campaigns, 606 visit every eligible class, +including sessions with no eligible class. In total, 11,870 of 12,243 class +instances are visited. The 18 deadline-limited sessions are ResNet training +on 780M, B570, and M3; all other sessions finish their declared search. +The longest search lasts 60.11\,s against the soft 60-second ceiling. The table counts class instances across processes, not distinct algorithms. \emph{Faster} counts candidate comparisons that replace an incumbent, not an end-to-end graph speedup; all 330 numerically rejected @@ -579,7 +581,7 @@ \subsection{Conditional performance portability} Meganeura/PyTorch means are 0.51/0.93 for inference, 0.60/0.85 for minimal shapes, and 0.39/0.98 for training. -This is explicitly conditional on shared GPU support. It excludes the CPU +This is explicitly conditional on shared GPU support. It excludes the RPL-U reference, additional H100 model sizes, and MI300X availability attempt; it cannot stand in for universal support. If support is required across every attempted GPU, both stacks have a @@ -789,7 +791,7 @@ \subsection{Preparation as a deployment budget} Whisper reuses one inference session for its two timing shapes. PyTorch \code{compile\_s} includes compilation and the first specializations/executions of forward, backward, and minimal shapes, -including on MPS and CPU. +including on MPS. Graph preparation and full-tensor qualification are recorded separately; model loading, cross-engine validation, and other process startup are outside these fields. These are declared preparation components, @@ -865,14 +867,21 @@ \subsection{Availability is part of portability} workloads through Vulkan and passes the numerical gates. This is a concrete graphics-API deployment advantage: PyTorch's available reference on this machine is CPU, not a GPU path awaiting benchmark -configuration. The campaign explicitly selects compiled CPU PyTorch. -Meganeura is faster in all five -strict full-inference comparisons (ratios 0.51--0.76), and in four -training comparisons. CPU wins minimal SmolVLA (1.37) and Whisper training (1.41). -These are GPU-versus-CPU results and are labeled as such; they do not -claim an equal-device speedup. Other Intel GPU generations can support +configuration. CPU PyTorch supplies only the numerical oracle for the +qualification in Table~\ref{tab:qualification}; no CPU-versus-GPU timing +comparison is included. Other Intel GPU generations can support XPU, as the separate Arc comparison demonstrates. +\begin{table}[t] +\caption{Qualified graphics-API platforms without a usable PyTorch GPU path +in the evaluated stack. Conditions are five workloads in two arithmetic contracts.} +\label{tab:qualification} +\centering +\scriptsize +\setlength{\tabcolsep}{3pt} +\input{tables/qualification} +\end{table} + Arc~B570 provides the separate, working XPU comparison. Its pinned native dense embedding backward passes only 19,968 of 73,728 expected elements in the shape-derived exact probe. @@ -1028,7 +1037,7 @@ \section{Threats to Validity} All primary references compile; CUDA/HIP/XPU use qualified replay, while MPS has no equivalent public whole-phase API. XPU's recorded embedding and math-attention accommodations are not an unmodified stock execution path. -The RPL-U CPU comparison, additional H100 models, and MI300X report +The RPL-U qualification, additional H100 models, and MI300X report cannot be pooled into the seven-system, five-workload GPU aggregate. Conditional scores would overstate universal portability if support holes were omitted. @@ -1067,7 +1076,8 @@ \section{Threats to Validity} Cross-engine sampled outputs and gradient norms can miss localized or sign errors. Full-element replay qualification checks PyTorch against itself; it does not turn the cross-engine gate into an elementwise gradient -comparison. All 252 pairs individually pass the 5\% gradient bound as well +comparison. All 222 GPU pairs and 30 RPL-U qualification processes individually +pass the 5\% gradient bound as well as their required three-process rule. Replay's fixed tensor maximum/RMS bounds and their coefficients are explicit; pointwise mismatch counts are diagnostics, not a second acceptance rule. @@ -1112,14 +1122,14 @@ \section{Conclusion} A single Vulkan/Metal compiler and runtime executes matched ML inference and backward workloads across heterogeneous consumer systems and H100, including Windows and an Intel GPU lacking a usable PyTorch GPU path. -All 252 paired processes validate at common source revisions. Every native +All 222 paired GPU processes validate at common source revisions. Every native session performs bounded measured search, strict arithmetic permits native-f32 cooperative matrices, and the compiled reference receives qualified CUDA/HIP/XPU replay wherever available. The performance result remains mixed: median strict inference and training times are $1.83\times$ and $2.47\times$ PyTorch's across seven shared GPU -configurations. Search visits every eligible class in 684 of 708 sessions, +configurations. Search visits every eligible class in 606 of 624 paired-GPU sessions, but the remaining gap includes limited kernel parallelism, excluded search dimensions, and backend-specific implementation choices. Replicated 1.7B results narrow the strict training gap without establishing @@ -1183,8 +1193,9 @@ \section*{Artifact Description} Python and pinned vendor requirements. \code{python scripts/p3hpc.py} collects the five-model cohort with native tuning, default reference compilation, and applicable replay. -No \code{--no-max-autotune} flag is needed; \code{--backend cpu} explicitly -selects the CPU support comparison, and \code{--gpu} disambiguates multiple +No \code{--no-max-autotune} flag is needed. Use \code{--backend cpu --qualify-only} +for a CPU correctness oracle on GPUs without a usable PyTorch path; +\code{--gpu} disambiguates multiple installed GPUs. The optional \code{--models SmolLM2-360M SmolLM2-1.7B} command collects the extension on H100. Each campaign writes a manifest and a sibling \code{latest.tgz}; diff --git a/paper/p3hpc/tables/devices.tex b/paper/p3hpc/tables/devices.tex index fd79ab75..54188b26 100644 --- a/paper/p3hpc/tables/devices.tex +++ b/paper/p3hpc/tables/devices.tex @@ -9,6 +9,5 @@ Radeon 780M & ROCM/compiled & Mesa 25.2.8 & HIP graph & 30/30 \\ Arc B570 & XPU/compiled & Mesa 26.0.3 & XPU graph & 30/30 \\ Apple M3 & MPS/compiled & Metal & no public replay & 30/30 \\ -Intel RPL-U (CPU ref.) & CPU/compiled & Mesa 26.0.3 & none & 30/30 \\ \bottomrule \end{tabular} diff --git a/paper/p3hpc/tables/preparation.tex b/paper/p3hpc/tables/preparation.tex index 62f306ba..bb6805af 100644 --- a/paper/p3hpc/tables/preparation.tex +++ b/paper/p3hpc/tables/preparation.tex @@ -9,7 +9,6 @@ Radeon 780M & 30 & 11.78 & 21.57 & 0.58 & 1.20 \\ Arc B570 & 30 & 12.86 & 26.53 & 0.09 & 1.73 \\ Apple M3 & 30 & 14.69 & 6.34 & 0.00 & 0.00 \\ -Intel RPL-U (CPU ref.) & 30 & 23.95 & 17.52 & 0.00 & 0.00 \\ H100 extension & 12 & 1.91 & 9.88 & 0.08 & 27.32 \\ \bottomrule \end{tabular} diff --git a/paper/p3hpc/tables/qualification.tex b/paper/p3hpc/tables/qualification.tex new file mode 100644 index 00000000..4fa9aee0 --- /dev/null +++ b/paper/p3hpc/tables/qualification.tex @@ -0,0 +1,7 @@ +\begin{tabular}{llrl} +\toprule +Device & Meganeura path & Qualified conditions & PyTorch GPU \\ +\midrule +Intel RPL-U & Vulkan & 10/10 & Unavailable \\ +\bottomrule +\end{tabular} diff --git a/paper/p3hpc/tables/ratios.tex b/paper/p3hpc/tables/ratios.tex index 08133991..79d65d1d 100644 --- a/paper/p3hpc/tables/ratios.tex +++ b/paper/p3hpc/tables/ratios.tex @@ -37,11 +37,5 @@ & StableDiffusion & \textbf{0.46} & \textbf{0.47} & 1.17 & \textbf{0.47} & \textbf{0.47} & 1.17 \\ & ResNet-50 & 3.06 & 1.92 & 2.04 & 2.62 & 1.82 & 2.15 \\ & Whisper-tiny & 1.79 & 1.09 & 3.14 & 1.82 & 1.15 & 2.51 \\ -\multicolumn{8}{l}{\emph{GPU-versus-CPU support comparison; excluded from GPU aggregates}} \\ -Intel RPL-U (CPU ref.) & SmolLM2-135M & \textbf{0.72} & \textbf{0.88} & \textbf{0.93} & \textbf{0.72} & 1.07 & \textbf{0.93} \\ - & SmolVLA & \textbf{0.76} & 1.37 & \textbf{0.80} & \textbf{0.76} & 1.20 & \textbf{0.79} \\ - & StableDiffusion & \textbf{0.67} & \textbf{0.67} & \textbf{0.92} & \textbf{0.71} & \textbf{0.67} & \textbf{0.85} \\ - & ResNet-50 & \textbf{0.51} & \textbf{0.52} & \textbf{0.98} & \textbf{0.51} & \textbf{0.52} & \textbf{0.98} \\ - & Whisper-tiny & \textbf{0.72} & \textbf{0.68} & 1.41 & \textbf{0.70} & \textbf{0.68} & 1.41 \\ \bottomrule \end{tabular} diff --git a/paper/p3hpc/tables/search.tex b/paper/p3hpc/tables/search.tex index 7b735b69..6b3903ae 100644 --- a/paper/p3hpc/tables/search.tex +++ b/paper/p3hpc/tables/search.tex @@ -9,7 +9,6 @@ Radeon 780M & 84 & 1533/1707 & 6 & 876 & 30 \\ Arc B570 & 84 & 1747/1842 & 6 & 1659 & 30 \\ Apple M3 & 84 & 1594/1698 & 6 & 1143 & 30 \\ -Intel RPL-U (CPU ref.) & 84 & 1452/1842 & 6 & 1127 & 0 \\ H100 extension & 36 & 168/168 & 0 & 119 & 0 \\ \bottomrule \end{tabular} diff --git a/src/cache.rs b/src/cache.rs index 8ef328de..e1fee5b9 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -9,9 +9,8 @@ use std::{io, path::Path}; /// Increment whenever the serialized execution plan or build pipeline changes /// in a way that can make an older plan unsafe to reuse. -// Version 7 invalidates plans that may have replaced a generated pointwise DAG -// with its legacy shader sentinel while fusing a matmul epilogue. -const CACHE_FORMAT_VERSION: u32 = 7; +// Version 10 changes default extraction and adds pre-allocation attention choices. +const CACHE_FORMAT_VERSION: u32 = 10; /// Cached execution plan with a graph fingerprint for invalidation. #[derive(Serialize, Deserialize)] @@ -423,7 +422,7 @@ mod tests { ..defaults }, OptimizeConfig { - mode: OptimizeMode::EgglogOutlined, + mode: OptimizeMode::EgglogWhole, ..defaults }, OptimizeConfig { @@ -438,6 +437,10 @@ mod tests { no_winograd: true, ..defaults }, + OptimizeConfig { + pack_swiglu: false, + ..defaults + }, ] { assert!( load_build_plan(&graph, hash(&config), &path) diff --git a/src/codegen.rs b/src/codegen.rs index 1dd1d9b8..ccf4ba4f 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -20,28 +20,48 @@ pub enum GemvReduction { /// and their barriers, so the wider the wave the more it removes — six /// levels on AMD's 64-wide wave against five on a 32-wide one. /// - /// Subgroup leaders claim workgroup slots dynamically, so this does not - /// assume a subgroup width or a mapping from local lanes to subgroups. + /// Subgroup IDs/counts index the partials without assuming a width or a + /// mapping from local lanes. A one-subgroup workgroup needs no barrier. Subgroup, } -/// Workgroup width and reduction style for the K-split GEMV family. +/// Workgroup geometry and reduction style for the K-split GEMV family. /// -/// Every GEMV kernel — plain, fused-add, transposed-B, f16, block-packed and -/// RmsNorm-folded — is derived from `matmul_gemv.wgsl` and shaped here, so -/// this is one axis across all of them rather than a choice per kernel. -/// Which shape wins is a property of the device, not of the graph, so it is -/// measured rather than predicted: see `Session::tune_with`. +/// Plain, fused-add, reduced-storage and RMSNorm-folded kernels share the +/// width/reduction generator. Dense transposed-B kernels also share row +/// grouping. The winner depends on both device and shape; `Session::tune_with` +/// measures the alternatives without changing the arithmetic or bindings. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] pub struct GemvShape { /// Threads per workgroup: 32, 64, 128 or 256. pub threads: u32, pub reduction: GemvReduction, + /// Contiguous rows of transposed B per workgroup: 1, 2 or 4. + /// Other layouts keep their existing four output columns per workgroup. + #[serde(default = "GemvShape::one_row")] + pub bt_rows: u32, } impl GemvShape { /// Widths a K-split GEMV can be generated at. pub const WIDTHS: [u32; 4] = [32, 64, 128, 256]; + /// Contiguous row groups supported by the transposed-B GEMV source. + pub const BT_ROWS: [u32; 3] = [1, 2, 4]; + + fn one_row() -> u32 { + 1 + } + + pub(crate) fn for_group(mut self, group: ShaderGroup) -> Self { + self.validate(); + if !matches!( + group, + ShaderGroup::MatMulGemvBT | ShaderGroup::MatMulGemvBTAdd + ) { + self.bt_rows = 1; + } + self + } /// The shape a group is generated with when nothing has chosen one. /// @@ -55,16 +75,19 @@ impl GemvShape { pub(crate) fn initial(group: ShaderGroup) -> Self { let threads = match group { ShaderGroup::MatMulGemv => 256, - ShaderGroup::MatMulGemvAdd | ShaderGroup::MatMulGemvBT => 32, + ShaderGroup::MatMulGemvAdd + | ShaderGroup::MatMulGemvBT + | ShaderGroup::MatMulGemvBTAdd => 32, _ => panic!("{group:?} is not a GEMV group"), }; Self { threads, reduction: GemvReduction::Tree, + bt_rows: 1, } } - /// Reject a width no GEMV source can be generated at. + /// Reject geometry no GEMV source can be generated at. pub(crate) fn validate(self) { assert!( Self::WIDTHS.contains(&self.threads), @@ -72,6 +95,10 @@ impl GemvShape { Self::WIDTHS, self.threads ); + assert!( + Self::BT_ROWS.contains(&self.bt_rows), + "unsupported GEMV row count" + ); } } @@ -152,53 +179,8 @@ fn parse_source(source: &str) -> Result { naga::front::wgsl::parse_str(source) } -/// Generate WGSL declarations and body for a fused epilogue chain. -/// -/// Returns (declarations, body) where declarations are `var` -/// lines for extra buffers, and body is a sequence of WGSL statements -/// that transform `val` (the matmul result for one output element). -pub fn epilogue_to_wgsl(epilogue: &[crate::compile::EpilogueOp]) -> (String, String) { - use crate::compile::EpilogueOp; - let mut decls = Vec::new(); - let mut body = Vec::new(); - let mut declared = std::collections::HashSet::new(); - - for op in epilogue { - #[allow(clippy::pattern_type_mismatch)] - match op { - EpilogueOp::Add(buf_idx) => { - let name = format!("epi_buf_{}", buf_idx); - if declared.insert(*buf_idx) { - decls.push(format!("var {}: array;", name)); - } - body.push(format!("val = val + {}[idx];", name)); - } - EpilogueOp::BiasAdd(buf_idx) => { - let name = format!("epi_buf_{}", buf_idx); - if declared.insert(*buf_idx) { - decls.push(format!("var {}: array;", name)); - } - body.push(format!("val = val + {}[col];", name)); - } - EpilogueOp::Relu => { - body.push("val = max(val, 0.0);".to_string()); - } - EpilogueOp::Silu => { - body.push("val = val / (1.0 + exp(-val));".to_string()); - } - EpilogueOp::Sigmoid => { - body.push("val = 1.0 / (1.0 + exp(-val));".to_string()); - } - EpilogueOp::Neg => { - body.push("val = -val;".to_string()); - } - } - } - (decls.join("\n"), body.join("\n ")) -} - -/// Generate WGSL for a [`crate::compile::MatMulEpilogue`] — the PointwiseDAG-based -/// replacement for `epilogue_to_wgsl`. Returns (declarations, body). +/// Generate WGSL for a [`crate::compile::MatMulEpilogue`]. +/// Returns (declarations, body). /// /// The DAG's `LoadInput(0)` maps to `val` (the matmul accumulator). /// `LoadInput(1+)` maps to `epi_buf_{n}` indexed by either `idx` @@ -291,14 +273,6 @@ pub fn matmul_prologue_to_wgsl( (decls.join("\n"), cache_decls.join("\n"), cache_init, expr) } -/// Where the fused epilogue statements come from. -pub enum EpilogueSource<'a> { - /// PointwiseDAG epilogue — what the compiler emits today. - Dag(&'a crate::compile::MatMulEpilogue), - /// Flat op chain, kept for plans cached before the DAG migration. - Ops(&'a [crate::compile::EpilogueOp]), -} - /// Tuning knobs for the register-tiled scalar matmul codegen. /// /// Defaults are what the plain kernel ships with: 32-row K staging and @@ -329,6 +303,14 @@ impl Default for MatmulKnobs { } } +/// Measured scalar layout for plain F32/F16 weights, with F32 accumulation. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub struct ScalarMatmulShape { + pub tile_size: u32, + pub k_stage: u32, + pub interleave_columns: bool, +} + /// How to specialize the matmul the epilogue is fused into. /// /// [`Default`] is the plain f32 64×64 kernel, so a caller that only wants @@ -341,7 +323,7 @@ pub struct MatMulOptions { /// computed for, or the grid and the kernel disagree about coverage. pub tile: MatMulTile, /// K staging depth and column layout for the tiled skeleton; only - /// consulted for f32 B storage, quantized formats have their own. + /// consulted for F32/F16 B storage; block-quantized formats have their own. pub knobs: MatmulKnobs, } @@ -353,7 +335,7 @@ pub struct MatMulOptions { /// `$STORE_BODY` hook serves every weight format. pub fn generate_matmul_with_epilogue( group: ShaderGroup, - epilogue: EpilogueSource<'_>, + epilogue: Option<&crate::compile::MatMulEpilogue>, options: MatMulOptions, ) -> ShaderModule { // Store-side fusion compiles through this generator rather than @@ -366,10 +348,7 @@ pub fn generate_matmul_with_epilogue( run along the parameter's first dimension, which is N here, not K" ); } - let (epi_decl, epi_body) = match epilogue { - EpilogueSource::Dag(dag) => matmul_epilogue_to_wgsl(dag), - EpilogueSource::Ops(ops) => epilogue_to_wgsl(ops), - }; + let (epi_decl, epi_body) = epilogue.map(matmul_epilogue_to_wgsl).unwrap_or_default(); let MatMulOptions { tile, .. } = options; let (a_idx, b_idx, fused_decl, fused_expr) = match group { ShaderGroup::MatMul => (MATMUL_A_FWD, MATMUL_B_FWD, "", ""), @@ -475,6 +454,7 @@ pub enum ShaderGroup { /// M=1 MatMulBT (`B` stored `[N,K]`): `C[1,N] = A × Bᵀ`. K-split with /// coalesced contiguous-K vec4 loads. MatMulGemvBT, + MatMulGemvBTAdd, Reduce, Softmax, CrossEntropy, @@ -582,7 +562,10 @@ pub fn generate_module(group: ShaderGroup, knobs: MatmulKnobs) -> ShaderModule { } ShaderGroup::MatMulATAdd => gen_matmul_at_add(knobs), ShaderGroup::MatMulBTAdd => gen_matmul_bt_add(knobs), - ShaderGroup::MatMulGemv | ShaderGroup::MatMulGemvAdd | ShaderGroup::MatMulGemvBT => { + ShaderGroup::MatMulGemv + | ShaderGroup::MatMulGemvAdd + | ShaderGroup::MatMulGemvBT + | ShaderGroup::MatMulGemvBTAdd => { generate_module_gemv(group, WeightFormat::F32, GemvShape::initial(group)) } ShaderGroup::Reduce => ShaderModule::new(include_str!("shaders/reduce.wgsl")), @@ -826,13 +809,20 @@ pub fn generate_horizontal_matmul( fn_src = fn_src.replacen("@workgroup_size(16, 16)", "", 1); fn_src = fn_src.replace("@builtin(workgroup_id) ", ""); fn_src = fn_src.replace("@builtin(local_invocation_id) ", ""); + fn_src = fn_src.replace("@builtin(subgroup_id) ", ""); fn_src = fn_src.replace("matrix_b[", &format!("matrix_b{i}[")); fn_src = fn_src.replace("matrix_c[", &format!("matrix_c{i}[")); bodies.push_str(&fn_src); bodies.push('\n'); } + let subgroup_arg = if coop.is_some() { + ", @builtin(subgroup_id) sg: u32" + } else { + "" + }; + let subgroup_call = if coop.is_some() { ", sg" } else { "" }; let mut dispatch = format!( - "{compute_attr}\nfn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) {{\n" + "{compute_attr}\nfn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3{subgroup_arg}) {{\n" ); for i in 0..count { let cond = if i + 1 == count { @@ -842,7 +832,9 @@ pub fn generate_horizontal_matmul( } else { format!("else if wgid.z == {i}u") }; - dispatch.push_str(&format!(" {cond} {{ horiz_{i}(wgid, lid); }}\n")); + dispatch.push_str(&format!( + " {cond} {{ horiz_{i}(wgid, lid{subgroup_call}); }}\n" + )); } dispatch.push_str("}\n"); ShaderModule::new(&format!("{header}{bodies}{dispatch}")) @@ -855,6 +847,7 @@ pub fn generate_wgsl(group: ShaderGroup) -> String { ShaderGroup::Conv2dGemmCoop | ShaderGroup::Conv2dGradInputGemmCoop => { naga::valid::Capabilities::COOPERATIVE_MATRIX | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP } ShaderGroup::ToF16 => naga::valid::Capabilities::SHADER_FLOAT16, _ => naga::valid::Capabilities::empty(), @@ -1358,10 +1351,9 @@ fn matmul_vars_tiled( }; let bm = tile.bm(); let tm = tile.tm(); - // The knobs only define the f32 skeleton; quantized and f16 B storage - // have their own layouts. + // Block decoders have fixed layouts; plain F32/F16 share the same skeleton. let k_tile = match b_mode { - WeightFormat::F32 => knobs.k_stage, + WeightFormat::F32 | WeightFormat::F16 => knobs.k_stage, _ => 32, }; assert!( @@ -1369,7 +1361,7 @@ fn matmul_vars_tiled( "unsupported scalar matmul K stage: {}", knobs.k_stage ); - let interleave_columns = b_mode == WeightFormat::F32 && knobs.interleave_columns; + let interleave_columns = !b_mode.is_quantized() && knobs.interleave_columns; let (acc_decl, compute_body, acc_array) = tiled_matmul_body(tile, k_tile, interleave_columns); let output_column = if interleave_columns { "tx + j * 16u".to_string() @@ -1883,7 +1875,10 @@ pub fn generate_module_weighted( // Every block-packed format takes the same K-split GEMV with its // own decoder substituted, so the format picks the helper rather // than the arm. - ShaderGroup::MatMulGemv | ShaderGroup::MatMulGemvAdd | ShaderGroup::MatMulGemvBT => { + ShaderGroup::MatMulGemv + | ShaderGroup::MatMulGemvAdd + | ShaderGroup::MatMulGemvBT + | ShaderGroup::MatMulGemvBTAdd => { generate_module_gemv(group, mode, GemvShape::initial(group)) } // Unsupported packed routes must fail closed: falling through would @@ -1917,11 +1912,16 @@ fn substitute(source: &str, old: &str, new: &str) -> String { /// Derived from `matmul_gemv.wgsl` by substitution so shape and reduction /// changes reach it automatically. Each consuming workgroup recomputes the /// small sum-of-squares prologue, avoiding a separate dispatch and boundary. -fn gemv_rmsnorm_source(format: WeightFormat) -> String { +fn gemv_rmsnorm_source(group: ShaderGroup, format: WeightFormat) -> String { + assert!(matches!( + group, + ShaderGroup::MatMulGemv | ShaderGroup::MatMulGemvBT + )); + let transposed = group == ShaderGroup::MatMulGemvBT; // `_pad` carries eps; the fused kernel needs no other new parameter. // Start from the already-format-specialized GEMV so packed decoders // compose with the prologue rather than being overwritten by it. - let src = gemv_source(ShaderGroup::MatMulGemv, format); + let src = gemv_source(group, format); let src = substitute(&src, " _pad: u32,", " eps_bits: u32,"); let src = if src.contains("var matrix_b: array>;") { substitute( @@ -1944,19 +1944,27 @@ fn gemv_rmsnorm_source(format: WeightFormat) -> String { }; let src = substitute( &src, - "var reduce_buf: array, LANES>;", - "var reduce_buf: array, LANES>;\n\ - var scale_buf: array;\n\ - var inv_rms: f32;", + "var reduce_buf:", + "var scale_buf: array;\n\ + var inv_rms: f32;\nvar reduce_buf:", ); // The early-out must not precede the prologue's barriers, which every // lane has to reach. The dispatch is exactly N/4 workgroups, so it // never fires in practice, but keep it uniform regardless. - let src = substitute( - &src, - " if col4 >= n_v4 { return; }\n let k = params.k;\n", - " let k = params.k;\n", - ); + let src = if transposed { + let src = substitute(&src, " if col >= params.n { return; }\n", ""); + substitute( + &src, + " var acc = 0.0;", + " let k = params.k;\n // Each thread accumulates a partial sum over its K-stride slice.\n var acc = 0.0;", + ) + } else { + substitute( + &src, + " if col4 >= n_v4 { return; }\n let k = params.k;\n", + " let k = params.k;\n", + ) + }; let src = substitute( &src, " // Each thread accumulates a partial sum over its K-stride slice.", @@ -1987,19 +1995,106 @@ fn gemv_rmsnorm_source(format: WeightFormat) -> String { \n\ // Each thread accumulates a partial sum over its K-stride slice.", ); - substitute( - &src, - " let a = matrix_a[kk];", - " let a = matrix_a[kk] * rs * norm_w[kk];", - ) + if transposed { + let src = substitute( + &src, + "if col4 >= n_v4 { return; }", + "if col >= params.n { return; }", + ); + let src = substitute( + &src, + "let v = matrix_a[si];", + "let v = matrix_a[si / 4u][si % 4u];", + ); + substitute( + &src, + " let a = matrix_a[kk_v4];", + " let at = kk_v4 * 4u;\n let a = matrix_a[kk_v4] * rs * vec4(norm_w[at], norm_w[at+1u], norm_w[at+2u], norm_w[at+3u]);", + ) + } else { + substitute( + &src, + " let a = matrix_a[kk];", + " let a = matrix_a[kk] * rs * norm_w[kk];", + ) + } } -pub fn generate_module_gemv_rmsnorm(shape: GemvShape, format: WeightFormat) -> ShaderModule { - ShaderModule::new(&gemv_shape_source(&gemv_rmsnorm_source(format), shape)) +pub fn generate_module_gemv_rmsnorm( + group: ShaderGroup, + shape: GemvShape, + format: WeightFormat, +) -> ShaderModule { + ShaderModule::new(&gemv_shape_source( + &gemv_row_source( + &gemv_rmsnorm_source(group, format), + shape.for_group(group).bt_rows, + ), + shape, + )) } const LANES_PREFIX: &str = "const LANES: u32 = "; +/// Share A loads and the norm prologue across contiguous B rows. Vector +/// accumulators reuse the same tree/subgroup reduction as single-row GEMV. +fn gemv_row_source(source: &str, rows: u32) -> String { + if rows == 1 { + return source.to_owned(); + } + assert!(matches!(rows, 2 | 4)); + let source = substitute( + source, + "let col = wgid.x + grid.x * wgid.y;", + &format!("let col = (wgid.x + grid.x * wgid.y) * {rows}u;"), + ); + let source = substitute( + &source, + "reduce_buf: array", + &format!("reduce_buf: array, LANES>"), + ); + let mut source = substitute( + &source, + "var acc = 0.0;", + &format!("var acc = vec{rows}(0.0);"), + ); + let start = source.find(" let b = ").expect("GEMV weight load"); + let end = start + + source[start..] + .find(" kk_v4 +=") + .expect("GEMV K stride"); + let load_end = start + source[start..].find(';').unwrap(); + let load = source[start + " let b = ".len()..load_end].to_owned(); + let mut body = String::new(); + for (row, component) in ["x", "y", "z", "w"].iter().take(rows as usize).enumerate() { + let value = substitute( + &load, + "row_off + kk_v4", + &format!("row_off + {row}u * k_v4 + kk_v4"), + ); + body.push_str(&format!( + " if col + {row}u < params.n {{ acc.{component} += dot(a, {value}); }}\n" + )); + } + source.replace_range(start..end, &body); + let start = source + .find(" matrix_c[col] = ") + .expect("GEMV output store"); + let end = start + source[start..].find(';').unwrap() + 1; + let addend = source[start..end].contains("src[col]"); + let mut body = " let total = reduce_buf[0] + reduce_buf[1];\n".to_owned(); + for (row, component) in ["x", "y", "z", "w"].iter().take(rows as usize).enumerate() { + let residual = if addend { + format!(" + src[col + {row}u]") + } else { + String::new() + }; + body.push_str(&format!(" if col + {row}u < params.n {{ matrix_c[col + {row}u] = total.{component}{residual}; }}\n")); + } + source.replace_range(start..end, &body); + source +} + /// The width a GEMV source is written against, from its `LANES` constant. /// /// Every kernel in the family declares that constant once and derives its @@ -2060,41 +2155,23 @@ fn gemv_shape_source(source: &str, shape: GemvShape) -> String { } body } - // One partial per wave reaches workgroup memory, and lane 0 sums the - // few that do. - // - // Neither the slot nor the leader may be derived from the local - // invocation id. WGSL and Vulkan both decline to relate - // `local_invocation_id` to subgroup membership, so `lane / sg_size` - // is not a subgroup index and `sg_id == 0` need not name exactly one - // invocation per subgroup — a wave holding the odd local ids, or a - // partially populated one, breaks both. Instead each wave's leader, - // elected by `subgroupBroadcastFirst`, claims a slot with an atomic, - // and the count comes back from the same counter. That costs one - // extra barrier to zero the counter, so two rather than the tree's - // one per halving level. - GemvReduction::Subgroup => { - let _ = threads; - " if lane == 0u { atomicStore(&wave_slots, 0u); }\n\ - \x20 workgroupBarrier();\n\ - \x20 let wave_total = subgroupAdd(acc);\n\ - \x20 if sg_id == subgroupBroadcastFirst(sg_id) {\n\ - \x20 reduce_buf[atomicAdd(&wave_slots, 1u)] = wave_total;\n\ - \x20 }\n\ - \x20 workgroupBarrier();\n\ - \x20 if lane == 0u {\n\ - \x20 let waves = atomicLoad(&wave_slots);\n\ - \x20 var total = reduce_buf[0];\n\ - \x20 var g = 1u;\n\ - \x20 loop {\n\ - \x20 if g >= waves { break; }\n\ - \x20 total = total + reduce_buf[g];\n\ - \x20 g = g + 1u;\n\ + // Use actual subgroup IDs, never lane/width. Elect a participating + // leader even for a partially populated subgroup. The subgroup count + // is workgroup-uniform, so the conditional barrier is convergent. + GemvReduction::Subgroup => " var group_total = subgroupAdd(acc);\n\ + \x20 if wave_count > 1u {\n\ + \x20 if sg_id == subgroupBroadcastFirst(sg_id) {\n\ + \x20 reduce_buf[wave_id] = group_total;\n\ + \x20 }\n\ + \x20 workgroupBarrier();\n\ + \x20 if lane == 0u {\n\ + \x20 group_total = reduce_buf[0];\n\ + \x20 for (var g = 1u; g < wave_count; g += 1u) {\n\ + \x20 group_total += reduce_buf[g];\n\ + \x20 }\n\ \x20 }\n\ - \x20 reduce_buf[0] = total;\n\ \x20 }\n" - .to_owned() - } + .to_owned(), }; source.replace_range(start..end, &reduction); @@ -2103,28 +2180,19 @@ fn gemv_shape_source(source: &str, shape: GemvShape) -> String { // not fold in a second slot. The store expression is otherwise left // alone, which is what keeps the fused-add and transposed-B forms // working without their own reduction code. - let folded = source.replace("reduce_buf[0] + reduce_buf[1]", "reduce_buf[0]"); + let folded = source.replace("reduce_buf[0] + reduce_buf[1]", "group_total"); assert_ne!(folded, source, "GEMV store did not fold two reduce slots"); source = folded; let signature = "@builtin(local_invocation_id) lid: vec3)"; let with_builtins = source.replace( signature, "@builtin(local_invocation_id) lid: vec3, \ - @builtin(subgroup_invocation_id) sg_id: u32)", + @builtin(subgroup_invocation_id) sg_id: u32, \ + @builtin(subgroup_id) wave_id: u32, \ + @builtin(num_subgroups) wave_count: u32)", ); assert_ne!(with_builtins, source, "GEMV entry point signature changed"); source = with_builtins; - - // The slot counter lives next to the buffer it indexes, and only the - // subgroup form declares it: the tree has no use for it and should - // not spend workgroup memory on it. - let anchor = "var reduce_buf:"; - let with_counter = source.replace( - anchor, - &format!("var wave_slots: atomic;\n{anchor}"), - ); - assert_ne!(with_counter, source, "GEMV workgroup buffer declaration"); - source = with_counter; } source } @@ -2173,17 +2241,21 @@ fn gemv_source(group: ShaderGroup, mode: WeightFormat) -> String { let base = match group { ShaderGroup::MatMulGemv => include_str!("shaders/matmul_gemv.wgsl"), ShaderGroup::MatMulGemvAdd => include_str!("shaders/matmul_gemv_add.wgsl"), - ShaderGroup::MatMulGemvBT => include_str!("shaders/matmul_gemv_bt.wgsl"), + ShaderGroup::MatMulGemvBT | ShaderGroup::MatMulGemvBTAdd => { + include_str!("shaders/matmul_gemv_bt.wgsl") + } _ => panic!("{group:?} is not a GEMV group"), }; - match (group, mode) { + let source = match (group, mode) { (_, WeightFormat::F32) => base.to_owned(), - (ShaderGroup::MatMulGemvBT, WeightFormat::F16) => gemv_bt_f16_source(base), + (ShaderGroup::MatMulGemvBT | ShaderGroup::MatMulGemvBTAdd, WeightFormat::F16) => { + gemv_bt_f16_source(base) + } (_, WeightFormat::F16) => gemv_f16_source(base), // Blocks run along the parameter's first dimension, which is N for a // transposed B, while every packed decoder indexes along K. The // kernel would return plausible but wrong numbers, so refuse. - (ShaderGroup::MatMulGemvBT, _) => panic!( + (ShaderGroup::MatMulGemvBT | ShaderGroup::MatMulGemvBTAdd, _) => panic!( "no {mode:?} variant for {group:?}; block-quantized weights run \ their blocks along K and cannot serve a transposed B" ), @@ -2193,6 +2265,20 @@ fn gemv_source(group: ShaderGroup, mode: WeightFormat) -> String { }); gemv_packed_source(base, helpers.as_str(), call) } + }; + if group == ShaderGroup::MatMulGemvBTAdd { + let source = substitute( + &source, + "var params:", + "var src: array;\nvar params:", + ); + substitute( + &source, + "matrix_c[col] = reduce_buf[0] + reduce_buf[1];", + "matrix_c[col] = reduce_buf[0] + reduce_buf[1] + src[col];", + ) + } else { + source } } @@ -2289,7 +2375,10 @@ pub(crate) fn generate_module_gemv( mode: WeightFormat, shape: GemvShape, ) -> ShaderModule { - ShaderModule::new(&gemv_shape_source(&gemv_source(group, mode), shape)) + ShaderModule::new(&gemv_shape_source( + &gemv_row_source(&gemv_source(group, mode), shape.for_group(group).bt_rows), + shape, + )) } fn gemv_packed_source(src: &str, helpers: &str, call: &str) -> String { @@ -2862,6 +2951,8 @@ fn gen_matmul_coop_wgsl_full( let shared_size_s = format!("{}", shared_size); let result_shared_size = output_tile * output_tile; let (result_shared_decl, result_store) = if epilogue.is_some() { + // Cooperative matrices are subgroup-scoped. Only one subgroup may + // store these shared tiles; every invocation still reaches the barrier. let store_iters = result_shared_size.div_ceil(wg_size); ( format!( @@ -2869,10 +2960,12 @@ fn gen_matmul_coop_wgsl_full( result_shared_size ), format!( - "coopStoreT(acc00, &shared_c[0], {output_tile}u);\n\ + "if sg == 0u {{\n\ + \x20 coopStoreT(acc00, &shared_c[0], {output_tile}u);\n\ \x20 coopStoreT(acc01, &shared_c[{tile}u], {output_tile}u);\n\ \x20 coopStoreT(acc10, &shared_c[{}u], {output_tile}u);\n\ \x20 coopStoreT(acc11, &shared_c[{}u], {output_tile}u);\n\ + \x20 }}\n\ \x20 workgroupBarrier();\n\ \n\ \x20 for (var e = 0u; e < {store_iters}u; e++) {{\n\ @@ -2897,7 +2990,8 @@ fn gen_matmul_coop_wgsl_full( } else { ( String::new(), - "coopStoreT(acc00, &matrix_c[c00], n);\n\ + "if sg == 0u {\n\ + \x20 coopStoreT(acc00, &matrix_c[c00], n);\n\ \x20 if n1_valid {\n\ \x20 coopStoreT(acc01, &matrix_c[c01], n);\n\ \x20 }\n\ @@ -2906,6 +3000,7 @@ fn gen_matmul_coop_wgsl_full( \x20 }\n\ \x20 if n1_valid && m1_valid {\n\ \x20 coopStoreT(acc11, &matrix_c[c11], n);\n\ + \x20 }\n\ \x20 }" .to_string(), ) @@ -3693,7 +3788,7 @@ pub fn generate_flash_attention_coop_module(head_dim: u32) -> ShaderModule { src.push('\n'); let _ = writeln!(src, "@compute @workgroup_size({wg_size})"); - src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) {\n"); + src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3, @builtin(subgroup_id) sg: u32) {\n"); let _ = writeln!(src, " let pos_base = wgid.x * {bq}u;"); src.push_str(" let head = wgid.y;\n"); src.push_str(" let q_seq = params.q_seq;\n"); @@ -3724,8 +3819,9 @@ pub fn generate_flash_attention_coop_module(head_dim: u32) -> ShaderModule { // outer KV loop). src.push_str(" let last_pos = min(pos_base + 15u, q_seq - 1u);\n"); src.push_str(" let max_kv_len = select(kv_seq, last_pos + 1u, kv_seq == 0u);\n"); + src.push_str(" let first_kv_len = select(kv_seq, pos_base + 1u, kv_seq == 0u);\n"); src.push_str( - " let min_kv_start = select(0u, max_kv_len - min(max_kv_len, window_size), window_size > 0u);\n\n", + " let min_kv_start = select(0u, first_kv_len - min(first_kv_len, window_size), window_size > 0u);\n\n", ); // Per-thread O accumulator and softmax state — REGISTERS. @@ -3795,7 +3891,7 @@ pub fn generate_flash_attention_coop_module(head_dim: u32) -> ShaderModule { src.push_str(" }\n"); let _ = writeln!( src, - " coopStoreT(score_acc, &shared_score[0], {bkv}u);" + " if sg == 0u {{ coopStoreT(score_acc, &shared_score[0], {bkv}u); }}" ); src.push_str(" workgroupBarrier();\n\n"); @@ -3974,7 +4070,7 @@ pub fn generate_flash_grad_q_coop_module(head_dim: u32) -> ShaderModule { src.push('\n'); let _ = writeln!(src, "@compute @workgroup_size({wg_size})"); - src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) {\n"); + src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3, @builtin(subgroup_id) sg: u32) {\n"); let _ = writeln!(src, " let pos_base = wgid.x * {bq}u;"); src.push_str(" let head = wgid.y;\n"); src.push_str(" let q_seq = params.q_seq;\n"); @@ -3991,8 +4087,9 @@ pub fn generate_flash_grad_q_coop_module(head_dim: u32) -> ShaderModule { // Workgroup-wide KV iteration bounds. src.push_str(" let last_pos = min(pos_base + 15u, q_seq - 1u);\n"); src.push_str(" let max_kv_len = select(kv_seq, last_pos + 1u, kv_seq == 0u);\n"); + src.push_str(" let first_kv_len = select(kv_seq, pos_base + 1u, kv_seq == 0u);\n"); src.push_str( - " let min_kv_start = select(0u, max_kv_len - min(max_kv_len, window_size), window_size > 0u);\n\n", + " let min_kv_start = select(0u, first_kv_len - min(first_kv_len, window_size), window_size > 0u);\n\n", ); // ---- Stage Q + dO into shared (once per workgroup) ---- @@ -4100,7 +4197,7 @@ pub fn generate_flash_grad_q_coop_module(head_dim: u32) -> ShaderModule { src.push_str(" }\n"); let _ = writeln!( src, - " coopStoreT(score_acc, &shared_score[0], {bkv}u);" + " if sg == 0u {{ coopStoreT(score_acc, &shared_score[0], {bkv}u); }}" ); // dp = dO @ V^T. @@ -4119,7 +4216,10 @@ pub fn generate_flash_grad_q_coop_module(head_dim: u32) -> ShaderModule { ); src.push_str(" dp_acc = coopMultiplyAdd(a, b, dp_acc);\n"); src.push_str(" }\n"); - let _ = writeln!(src, " coopStoreT(dp_acc, &shared_dp[0], {bkv}u);"); + let _ = writeln!( + src, + " if sg == 0u {{ coopStoreT(dp_acc, &shared_dp[0], {bkv}u); }}" + ); src.push_str(" workgroupBarrier();\n\n"); // ds = p * (dp - row_sum). 64 threads × 4 elements each = 256 entries @@ -4311,7 +4411,7 @@ pub fn generate_flash_grad_kv_coop_module(head_dim: u32) -> ShaderModule { src.push('\n'); let _ = writeln!(src, "@compute @workgroup_size({wg_size})"); - src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) {\n"); + src.push_str("fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3, @builtin(subgroup_id) sg: u32) {\n"); let _ = writeln!(src, " let kv_base = wgid.x * {bkv}u;"); src.push_str(" let kv_head = wgid.y;\n"); src.push_str(" let q_seq = params.q_seq;\n"); @@ -4472,7 +4572,7 @@ pub fn generate_flash_grad_kv_coop_module(head_dim: u32) -> ShaderModule { src.push_str(" }\n"); let _ = writeln!( src, - " coopStoreT(score_acc, &shared_score[0], {bq}u);" + " if sg == 0u {{ coopStoreT(score_acc, &shared_score[0], {bq}u); }}" ); // dp = V @ dO^T (BKV x BQ). @@ -4491,7 +4591,10 @@ pub fn generate_flash_grad_kv_coop_module(head_dim: u32) -> ShaderModule { ); src.push_str(" dp_acc = coopMultiplyAdd(a_v, b_dot, dp_acc);\n"); src.push_str(" }\n"); - let _ = writeln!(src, " coopStoreT(dp_acc, &shared_dp[0], {bq}u);"); + let _ = writeln!( + src, + " if sg == 0u {{ coopStoreT(dp_acc, &shared_dp[0], {bq}u); }}" + ); src.push_str(" workgroupBarrier();\n\n"); // p[kv, q] = exp(score * scale - lse[q]); ds[kv, q] = p * (dp - row_sum[q]). @@ -5306,7 +5409,7 @@ pub fn generate_conv2d_coop_module( // Main function let _ = writeln!(src, "@compute @workgroup_size(64)"); src.push_str( - "fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) {\n", + "fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3, @builtin(subgroup_id) sg: u32) {\n", ); if backward { @@ -5561,6 +5664,7 @@ pub fn generate_conv2d_coop_module( if config.use_f16_input || !backward { // Forward and f16 cooperative kernels retain the direct-store // alignment requirement enforced by runtime selection. + src.push_str(" if sg == 0u {\n"); src.push_str(" coopStoreT(acc00, &dst[c00], n_total);\n"); src.push_str(" if n1_valid {\n"); src.push_str(" coopStoreT(acc01, &dst[c01], n_total);\n"); @@ -5571,23 +5675,28 @@ pub fn generate_conv2d_coop_module( src.push_str(" if n1_valid && m1_valid {\n"); src.push_str(" coopStoreT(acc11, &dst[c11], n_total);\n"); src.push_str(" }\n"); + src.push_str(" }\n"); } else { // A direct cooperative store at a partial right edge crosses the // logical NCHW row boundary. Full column tiles keep the fast path; // only the final partial workgroup stages through the now-dead f32 // input tiles and performs bounds-checked scalar stores. let _ = writeln!(src, " if (tile_col + {output_tile}u) <= n_total {{"); + src.push_str(" if sg == 0u {\n"); src.push_str(" coopStoreT(acc00, &dst[c00], n_total);\n"); src.push_str(" coopStoreT(acc01, &dst[c01], n_total);\n"); src.push_str(" if m1_valid {\n"); src.push_str(" coopStoreT(acc10, &dst[c10], n_total);\n"); src.push_str(" coopStoreT(acc11, &dst[c11], n_total);\n"); src.push_str(" }\n"); + src.push_str(" }\n"); src.push_str(" } else {\n"); + src.push_str(" if sg == 0u {\n"); let _ = writeln!(src, " coopStoreT(acc00, &shared_b0[0], {tile}u);"); let _ = writeln!(src, " coopStoreT(acc01, &shared_b1[0], {tile}u);"); let _ = writeln!(src, " coopStoreT(acc10, &shared_a0[0], {tile}u);"); let _ = writeln!(src, " coopStoreT(acc11, &shared_a1[0], {tile}u);"); + src.push_str(" }\n"); src.push_str(" workgroupBarrier();\n"); let _ = writeln!( src, @@ -6048,6 +6157,10 @@ mod tests { ShaderGroup::MatMulGemvBT, gemv_caps(ShaderGroup::MatMulGemvBT), ), + ( + ShaderGroup::MatMulGemvBTAdd, + gemv_caps(ShaderGroup::MatMulGemvBTAdd), + ), (ShaderGroup::Reduce, naga::valid::Capabilities::empty()), (ShaderGroup::Softmax, naga::valid::Capabilities::empty()), ( @@ -6075,17 +6188,20 @@ mod tests { ( ShaderGroup::FlashAttentionCoop, naga::valid::Capabilities::COOPERATIVE_MATRIX - | naga::valid::Capabilities::SHADER_FLOAT16, + | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP, ), ( ShaderGroup::FlashGradQCoop, naga::valid::Capabilities::COOPERATIVE_MATRIX - | naga::valid::Capabilities::SHADER_FLOAT16, + | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP, ), ( ShaderGroup::FlashGradKVCoop, naga::valid::Capabilities::COOPERATIVE_MATRIX - | naga::valid::Capabilities::SHADER_FLOAT16, + | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP, ), ( ShaderGroup::MultiHeadAttnGradQ, @@ -6170,7 +6286,8 @@ mod tests { // Cooperative execution is a modifier rather than a group, so its // modules are reached through the scalar group they derive from. let coop_caps = naga::valid::Capabilities::COOPERATIVE_MATRIX - | naga::valid::Capabilities::SHADER_FLOAT16; + | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP; let config = CoopConfig { tile_size: 16, use_f16_input: true, @@ -6331,12 +6448,16 @@ mod tests { ( "coop", coop, - Capabilities::COOPERATIVE_MATRIX | Capabilities::SHADER_FLOAT16, + Capabilities::COOPERATIVE_MATRIX + | Capabilities::SHADER_FLOAT16 + | Capabilities::SUBGROUP, ), ( "compensated", compensated, - Capabilities::COOPERATIVE_MATRIX | Capabilities::SHADER_FLOAT16, + Capabilities::COOPERATIVE_MATRIX + | Capabilities::SHADER_FLOAT16 + | Capabilities::SUBGROUP, ), ] { let module = generate_horizontal_matmul(ShaderGroup::MatMul, count, Some(&cfg)); @@ -6368,7 +6489,9 @@ mod tests { ); Validator::new( ValidationFlags::all() ^ ValidationFlags::BINDINGS, - Capabilities::COOPERATIVE_MATRIX | Capabilities::SHADER_FLOAT16, + Capabilities::COOPERATIVE_MATRIX + | Capabilities::SHADER_FLOAT16 + | Capabilities::SUBGROUP, ) .validate(&module.module) .unwrap_or_else(|error| panic!("{group:?} compensated coop failed: {error:#?}")); @@ -6394,7 +6517,9 @@ mod tests { use_f16_input: true, compensated: false, }; - let capabilities = Capabilities::COOPERATIVE_MATRIX | Capabilities::SHADER_FLOAT16; + let capabilities = Capabilities::COOPERATIVE_MATRIX + | Capabilities::SHADER_FLOAT16 + | Capabilities::SUBGROUP; let flags = ValidationFlags::all() ^ ValidationFlags::BINDINGS; for group in [ @@ -6423,7 +6548,8 @@ mod tests { let empty = naga::valid::Capabilities::empty(); let f16 = naga::valid::Capabilities::SHADER_FLOAT16; let coop = naga::valid::Capabilities::COOPERATIVE_MATRIX - | naga::valid::Capabilities::SHADER_FLOAT16; + | naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SUBGROUP; let groups: &[(ShaderGroup, naga::valid::Capabilities)] = &[ (ShaderGroup::Unary, empty), (ShaderGroup::Binary, empty), @@ -6558,7 +6684,7 @@ mod tests { | ShaderEntry::MatMulGemvBT => { vec!["matrix_a", "matrix_b", "matrix_c", "params"] } - ShaderEntry::MatMulGemvAdd => { + ShaderEntry::MatMulGemvAdd | ShaderEntry::MatMulGemvBTAdd => { vec!["matrix_a", "matrix_b", "matrix_c", "src", "params"] } ShaderEntry::FusedMatMulAdd @@ -6744,6 +6870,7 @@ mod tests { ShaderEntry::MatMulGemv, ShaderEntry::MatMulGemvAdd, ShaderEntry::MatMulGemvBT, + ShaderEntry::MatMulGemvBTAdd, ShaderEntry::FusedMatMulAdd, ShaderEntry::FusedMatMulATAdd, ShaderEntry::FusedMatMulBTAdd, @@ -6890,7 +7017,9 @@ mod tests { fn generated_conv2d_coop_modules_are_valid() { use naga::valid::{Capabilities, ValidationFlags, Validator}; - let coop_caps = Capabilities::COOPERATIVE_MATRIX | Capabilities::SHADER_FLOAT16; + let coop_caps = Capabilities::COOPERATIVE_MATRIX + | Capabilities::SHADER_FLOAT16 + | Capabilities::SUBGROUP; let flags = ValidationFlags::all() ^ ValidationFlags::BINDINGS; let configs = [ CoopConfig { @@ -6996,6 +7125,9 @@ mod tests { /// fails here rather than at pipeline creation on someone's GPU. #[test] fn every_gemv_shape_composes_with_every_weight_format() { + let legacy: GemvShape = + serde_json::from_str(r#"{"threads":32,"reduction":"Tree"}"#).unwrap(); + assert_eq!(legacy.bt_rows, 1); let formats = [ (WeightFormat::F32, "matrix_b: array>"), (WeightFormat::F16, "array>"), @@ -7008,60 +7140,92 @@ mod tests { (WeightFormat::Q3K, "dequant_q3k("), ]; for (format, marker) in formats { - for group in [ShaderGroup::MatMulGemv, ShaderGroup::MatMulGemvAdd] { + for group in [ + ShaderGroup::MatMulGemv, + ShaderGroup::MatMulGemvAdd, + ShaderGroup::MatMulGemvBT, + ShaderGroup::MatMulGemvBTAdd, + ] { + if format.is_quantized() + && matches!( + group, + ShaderGroup::MatMulGemvBT | ShaderGroup::MatMulGemvBTAdd + ) + { + continue; + } for threads in [32, 64, 128, 256] { for reduction in [GemvReduction::Tree, GemvReduction::Subgroup] { - let shape = GemvShape { threads, reduction }; - let module = generate_module_gemv(group, format, shape); - let caps = match reduction { - GemvReduction::Tree => naga::valid::Capabilities::empty(), - GemvReduction::Subgroup => naga::valid::Capabilities::SUBGROUP, - } | match format { - // f16 storage reads real `f16` values. - WeightFormat::F16 => naga::valid::Capabilities::SHADER_FLOAT16, - // Block scales are f16 bit patterns decoded with - // `unpack2x16float` into f32, which is the weaker - // capability — it needs no f16 arithmetic type. - f if f.is_quantized() => { - naga::valid::Capabilities::SHADER_FLOAT16_IN_FLOAT32 + for bt_rows in GemvShape::BT_ROWS { + let shape = GemvShape { + threads, + reduction, + bt_rows, } - _ => naga::valid::Capabilities::empty(), - }; - let flags = naga::valid::ValidationFlags::all() - ^ naga::valid::ValidationFlags::BINDINGS; - naga::valid::Validator::new(flags, caps) - .validate(&module.module) - .unwrap_or_else(|e| { - panic!("{format:?} {group:?} {shape:?} failed validation: {e:#?}") - }); - let source = module.source; - assert!( - source.contains(marker), - "{format:?} {group:?} {shape:?} lost its B representation" - ); - assert!( - source.contains(&format!("{LANES_PREFIX}{threads}u;")), - "{format:?} {group:?} {shape:?} kept the declared width" - ); - let subgroup = reduction == GemvReduction::Subgroup; - assert_eq!( - source.contains("subgroupAdd"), - subgroup, - "{format:?} {group:?} {shape:?} reduction mismatch" - ); - // The tree walks the workgroup in halves; the subgroup - // form must leave the total in slot 0 alone. - assert_eq!( - source.contains("reduce_buf[0] + reduce_buf[1]"), - !subgroup, - "{format:?} {group:?} {shape:?} store expression mismatch" - ); - // The fused add's own term survives the rewrite. - if group == ShaderGroup::MatMulGemvAdd { + .for_group(group); + if shape.bt_rows != bt_rows { + continue; + } + let module = generate_module_gemv(group, format, shape); + let caps = match reduction { + GemvReduction::Tree => naga::valid::Capabilities::empty(), + GemvReduction::Subgroup => naga::valid::Capabilities::SUBGROUP, + } | match format { + // f16 storage reads real `f16` values. + WeightFormat::F16 => naga::valid::Capabilities::SHADER_FLOAT16, + // Block scales are f16 bit patterns decoded with + // `unpack2x16float` into f32, which is the weaker + // capability — it needs no f16 arithmetic type. + f if f.is_quantized() => { + naga::valid::Capabilities::SHADER_FLOAT16_IN_FLOAT32 + } + _ => naga::valid::Capabilities::empty(), + }; + let flags = naga::valid::ValidationFlags::all() + ^ naga::valid::ValidationFlags::BINDINGS; + naga::valid::Validator::new(flags, caps) + .validate(&module.module) + .unwrap_or_else(|e| { + panic!( + "{format:?} {group:?} {shape:?} failed validation: {e:#?}" + ) + }); + let source = module.source; assert!( - source.contains("src[col4]"), - "{format:?} {shape:?} dropped the fused addend" + source.contains(marker), + "{format:?} {group:?} {shape:?} lost its B representation" ); + assert!( + source.contains(&format!("{LANES_PREFIX}{threads}u;")), + "{format:?} {group:?} {shape:?} kept the declared width" + ); + let subgroup = reduction == GemvReduction::Subgroup; + assert_eq!( + source.contains("subgroupAdd"), + subgroup, + "{format:?} {group:?} {shape:?} reduction mismatch" + ); + // The tree walks the workgroup in halves; the subgroup + // form must leave the total in slot 0 alone. + assert_eq!( + source.contains("reduce_buf[0] + reduce_buf[1]"), + !subgroup, + "{format:?} {group:?} {shape:?} store expression mismatch" + ); + // The fused add's own term survives the rewrite. + if group == ShaderGroup::MatMulGemvAdd { + assert!( + source.contains("src[col4]"), + "{format:?} {shape:?} dropped the fused addend" + ); + } + if group == ShaderGroup::MatMulGemvBTAdd { + assert!(source.contains(if bt_rows == 1 { + "src[col]" + } else { + "src[col + 0u]" + })); + } } } } @@ -7096,7 +7260,11 @@ mod tests { }; for threads in [32, 64, 128, 256] { for reduction in [GemvReduction::Tree, GemvReduction::Subgroup] { - let shape = GemvShape { threads, reduction }; + let shape = GemvShape { + threads, + reduction, + bt_rows: 1, + }; for packed_dot in [true, false] { for norm in [false, true] { let module = generate_module_gemv_int_dot( @@ -7158,31 +7326,44 @@ mod tests { fn packed_gemv_rmsnorm_keeps_the_decoder() { for format in [ WeightFormat::F32, + WeightFormat::F16, WeightFormat::Q40, WeightFormat::Q4K, WeightFormat::Q8, ] { - let sm = generate_module_gemv_rmsnorm( - GemvShape { - threads: 64, - reduction: GemvReduction::Subgroup, - }, - format, - ); - assert!( - sm.source.contains("inv_rms"), - "{format:?} fused GEMV lost the RmsNorm prologue" - ); - assert!( - sm.source.contains("norm_w"), - "{format:?} fused GEMV lost the norm-weight binding" - ); - match format { - WeightFormat::F32 => assert!(sm.source.contains("matrix_b: array>")), - WeightFormat::Q40 => assert!(sm.source.contains("dequant_q40(")), - WeightFormat::Q4K => assert!(sm.source.contains("dequant_q4k(")), - WeightFormat::Q8 => assert!(sm.source.contains("dequant_q8(")), - _ => {} + for group in [ShaderGroup::MatMulGemv, ShaderGroup::MatMulGemvBT] { + if group == ShaderGroup::MatMulGemvBT && format.is_quantized() { + continue; + } + let sm = generate_module_gemv_rmsnorm( + group, + GemvShape { + threads: 64, + reduction: GemvReduction::Subgroup, + bt_rows: 4, + }, + format, + ); + let flags = + naga::valid::ValidationFlags::all() ^ naga::valid::ValidationFlags::BINDINGS; + naga::valid::Validator::new(flags, naga::valid::Capabilities::all()) + .validate(&sm.module) + .unwrap(); + assert!( + sm.source.contains("inv_rms"), + "{format:?} fused GEMV lost the RmsNorm prologue" + ); + assert!( + sm.source.contains("norm_w"), + "{format:?} fused GEMV lost the norm-weight binding" + ); + match format { + WeightFormat::F32 => assert!(sm.source.contains("matrix_b: array>")), + WeightFormat::Q40 => assert!(sm.source.contains("dequant_q40(")), + WeightFormat::Q4K => assert!(sm.source.contains("dequant_q4k(")), + WeightFormat::Q8 => assert!(sm.source.contains("dequant_q8(")), + _ => {} + } } } } @@ -7199,17 +7380,18 @@ mod tests { let tree = barriers(GemvShape { threads, reduction: GemvReduction::Tree, + bt_rows: 1, }); let subgroup = barriers(GemvShape { threads, reduction: GemvReduction::Subgroup, + bt_rows: 1, }); - // One barrier to zero the slot counter and one after the - // leaders have claimed their slots. Constant in the width, - // which is the point. + // The only barrier gathers partials when there is more than + // one subgroup; a single subgroup keeps its sum in registers. assert_eq!( - subgroup, 2, - "the subgroup reduction needs exactly two barriers at {threads} threads" + subgroup, 1, + "the subgroup reduction needs one conditional barrier at {threads} threads" ); assert_eq!( tree, @@ -7256,7 +7438,7 @@ mod tests { }; let sm = generate_matmul_with_epilogue( ShaderGroup::MatMul, - EpilogueSource::Dag(&epi), + Some(&epi), MatMulOptions { format: WeightFormat::Q4, ..Default::default() @@ -7283,10 +7465,9 @@ mod tests { #[test] #[should_panic(expected = "does not support block-quantized")] fn quantized_bt_epilogue_is_refused() { - let relu = [crate::compile::EpilogueOp::Relu]; let _ = generate_matmul_with_epilogue( ShaderGroup::MatMulBTAdd, - EpilogueSource::Ops(&relu), + None, MatMulOptions { format: WeightFormat::Q4K, ..Default::default() @@ -7503,20 +7684,26 @@ mod tests { ShaderGroup::MatMulAT, ShaderGroup::MatMulBT, ] { - let relu = [crate::compile::EpilogueOp::Relu]; + let relu = crate::compile::MatMulEpilogue { + dag: crate::schedule::PointwiseDAG { + n_inputs: 1, + ops: vec![ + crate::schedule::Pw::LoadInput(0), + crate::schedule::Pw::Relu(0), + ], + output: 1, + }, + inputs: Vec::new(), + }; let small = generate_matmul_with_epilogue( group, - EpilogueSource::Ops(&relu), + Some(&relu), MatMulOptions { tile: MatMulTile::Small, ..Default::default() }, ); - let large = generate_matmul_with_epilogue( - group, - EpilogueSource::Ops(&relu), - MatMulOptions::default(), - ); + let large = generate_matmul_with_epilogue(group, Some(&relu), MatMulOptions::default()); assert_ne!( small.source, large.source, "{group:?}: small and large epilogue shaders must differ" diff --git a/src/compile.rs b/src/compile.rs index 45a48aed..733ff92c 100644 --- a/src/compile.rs +++ b/src/compile.rs @@ -1,3 +1,4 @@ +use crate::codegen::ShaderGroup; use crate::graph::{DType, Graph, Node, NodeId, Op, PairwiseGradKind}; use crate::schedule::{PointwiseDAG, Pw, ReductionEpilogue, ReductionKernel}; use serde::{Deserialize, Serialize}; @@ -6,6 +7,33 @@ use std::collections::HashMap; mod softplus; mod split_k; +/// Host layout of the cached block-attention WGSL uniform. Dispatch encoding, +/// runtime binding and tuning all use this layout rather than indexing words. +#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod)] +#[repr(C)] +pub(crate) struct CachedBlockAttentionParams { + pub window_size: u32, + pub num_heads: u32, + pub num_kv_heads: u32, + pub head_dim: u32, + pub block_len: u32, + pub max_seq: u32, + pub splits: u32, + pub chunk: u32, +} + +impl CachedBlockAttentionParams { + pub fn from_words(words: &[u32]) -> Option { + bytemuck::try_from_bytes(bytemuck::cast_slice(words)) + .ok() + .copied() + } + + pub fn to_words(self) -> Vec { + bytemuck::cast_slice(std::slice::from_ref(&self)).to_vec() + } +} + /// Weight storage format for matmul B operands. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum WeightFormat { @@ -72,7 +100,7 @@ impl WeightFormat { /// participates in the plan-cache fingerprint automatically. Defaults come /// from capability-signature heuristics plus `MEGANEURA_FLASH_*` env /// overrides; a session-build tuner can substitute measured values instead. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct TuningKnobs { /// Elements-per-thread cap for flash-attention forward codegen. pub flash_ept_cap: u32, @@ -138,7 +166,7 @@ pub struct CompileOptions { /// Enable the experimental reduced-precision cooperative flash /// backward kernels. pub flash_backward_coop: bool, - /// Starting workgroup width and cross-lane reduction for the K-split + /// Starting workgroup geometry and cross-lane reduction for the K-split /// GEMV family, overriding each group's own default. /// /// This is a starting point, not a decision: `Session::tune_with` @@ -146,6 +174,10 @@ pub struct CompileOptions { /// faster. Set it to pin a shape for a benchmark or a reproduction. #[serde(default, skip_serializing_if = "Option::is_none")] pub gemv_shape: Option, + /// Cached-attention implementation: 1 is unsplit, 2..=16 use partials. + /// None retains the ordinary lowering; measured construction searches this. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cached_attention_splits: Option, /// Quantize the activation row to Q8_1 inside the K-split GEMV and do /// the inner product with integer dot products, for the weight formats /// that have an int-dot kernel (GGML Q4_0 and Meganeura Q8). @@ -170,11 +202,34 @@ impl Default for CompileOptions { flash_forward_coop: true, flash_backward_coop: false, gemv_shape: None, + cached_attention_splits: None, quantized_activations: true, } } } +impl CompileOptions { + fn gemv_kernel(&self, group: ShaderGroup, format: WeightFormat) -> Kernel { + Kernel::Gemv { + shape: self.gemv_shape.map_or_else( + || crate::codegen::GemvShape::initial(group), + |shape| shape.for_group(group), + ), + integer_dot: self.quantized_activations + && matches!(group, ShaderGroup::MatMulGemv | ShaderGroup::MatMulGemvAdd) + && matches!( + format, + WeightFormat::Q40 + | WeightFormat::Q8 + | WeightFormat::Q4K + | WeightFormat::Q5K + | WeightFormat::Q6K + | WeightFormat::Q3K + ), + } + } +} + /// Identifies which shader and entry point to use. #[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum ShaderEntry { @@ -194,6 +249,8 @@ pub enum ShaderEntry { /// M=1 MatMulBT specialization (`B` stored `[N,K]`). K-split with /// naturally coalesced vec4 reads along the contiguous K axis. MatMulGemvBT, + /// M=1 transposed-B GEMV with a fused addend. + MatMulGemvBTAdd, FusedMatMulAdd, FusedMatMulATAdd, FusedMatMulBTAdd, @@ -362,6 +419,7 @@ impl ShaderEntry { | ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvAdd | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd | ShaderEntry::FusedMatMulAdd | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd => "matrix", @@ -490,6 +548,7 @@ impl ShaderEntry { ShaderEntry::MatMulGemv => ShaderGroup::MatMulGemv, ShaderEntry::MatMulGemvAdd => ShaderGroup::MatMulGemvAdd, ShaderEntry::MatMulGemvBT => ShaderGroup::MatMulGemvBT, + ShaderEntry::MatMulGemvBTAdd => ShaderGroup::MatMulGemvBTAdd, ShaderEntry::FusedMatMulAdd => ShaderGroup::MatMulAdd, ShaderEntry::FusedMatMulATAdd => ShaderGroup::MatMulATAdd, ShaderEntry::FusedMatMulBTAdd => ShaderGroup::MatMulBTAdd, @@ -606,6 +665,7 @@ impl ShaderEntry { | ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvAdd | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd | ShaderEntry::FusedMatMulAdd | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd @@ -724,9 +784,9 @@ impl Dispatch { pub fn profile_family(&self) -> &'static str { if self.is_row_data_movement() { "data_movement" - } else if self.reduction.is_some() { + } else if self.reduction().is_some() { "normalization_reduction" - } else if self.pointwise.is_some() { + } else if self.pointwise().is_some() { "pointwise" } else { self.shader.profile_family() @@ -743,7 +803,7 @@ impl Dispatch { } fn is_zero_fill(&self) -> bool { - self.pointwise.as_ref().is_some_and(|dag| { + self.pointwise().is_some_and(|dag| { dag.n_inputs == 1 && dag.ops == [Pw::const_f32(0.0)] && dag.output == 0 }) } @@ -755,18 +815,6 @@ impl Dispatch { } } -/// Legacy enum — kept for serde backward compat of cached plans. -/// New code should use `MatMulEpilogue` (a `PointwiseDAG`). -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub enum EpilogueOp { - Add(u8), - BiasAdd(u8), - Relu, - Silu, - Sigmoid, - Neg, -} - /// A RmsNorm folded into a GEMV's A operand. /// /// A modifier on `ShaderEntry::MatMulGemv` rather than a shader entry of @@ -791,8 +839,8 @@ pub enum EpilogueLoadKind { } /// A fused epilogue applied in the matmul store loop, expressed as a -/// [`PointwiseDAG`]. Replaces the closed `EpilogueOp` enum so arbitrary -/// per-element transforms can be fused without new enum variants. +/// [`PointwiseDAG`], so arbitrary per-element transforms can be fused +/// without new enum variants. /// /// `LoadInput(0)` in the DAG = `val` (the matmul accumulator result). /// `LoadInput(1+)` indexes into `inputs`, each with its own buffer + @@ -895,21 +943,20 @@ fn can_horizontal_fuse(a: &Dispatch, b: &Dispatch) -> bool { && a.workgroups == b.workgroups && a.workgroups[2] == 1 && a.params == b.params - && a.use_coop == b.use_coop - && a.use_coop_compensated == b.use_coop_compensated - && a.use_small_tiles == b.use_small_tiles + && a.kernel == b.kernel + && matches!( + a.kernel, + Kernel::Default + | Kernel::SmallTile + | Kernel::Cooperative + | Kernel::CooperativeCompensated + ) && !a.weight_format.uses_reduced_storage() && a.weight_format == b.weight_format && a.matmul_prologue.is_none() && b.matmul_prologue.is_none() && a.matmul_epilogue.is_none() && b.matmul_epilogue.is_none() - && a.epilogue.is_empty() - && b.epilogue.is_empty() - && a.pointwise.is_none() - && b.pointwise.is_none() - && a.reduction.is_none() - && b.reduction.is_none() && a.gemv_rmsnorm.is_none() && b.gemv_rmsnorm.is_none() && a.input_buffers.len() == 2 @@ -922,10 +969,6 @@ fn merge_horizontal(dispatches: &[Dispatch], batch: &[usize]) -> Dispatch { let n = batch.len() as u32; merged.horizontal_batch = n; merged.workgroups[2] = n; - // The original fallback describes one output with z=1. The packed - // pipeline has different bindings, and only its selected variant is - // compiled. A future cooperative search must construct a packed fallback. - merged.scalar_fallback = None; merged.input_buffers = vec![merged.input_buffers[0]]; merged.extra_outputs.clear(); for (k, &idx) in batch.iter().enumerate() { @@ -978,46 +1021,15 @@ pub struct Dispatch { pub extra_outputs: Vec, /// Extra params to upload as a uniform buffer. pub params: Vec, - /// When true, this dispatch uses the cooperative matrix pipeline - /// (set at runtime based on per-dispatch eligibility). - #[serde(default)] - pub use_coop: bool, - /// Cooperative f16 path with hi/lo residual staging (C1). Retained as an - /// explicit experimental variant; automatic selection does not use it - /// for `requires_full_precision` work because it cannot preserve f32's - /// exponent range. + /// Exactly one implementation of the shader's binding contract. #[serde(default)] - pub use_coop_compensated: bool, + pub kernel: Kernel, /// Number of same-A sibling matmuls packed into this dispatch (D1). /// 0/1 = not packed. Extra B operands follow A in `input_buffers`; /// extra C outputs are `extra_outputs`. `workgroups[2]` is the pack /// count when this is ≥ 2 (only applied when the original Z was 1). #[serde(default)] pub horizontal_batch: u32, - /// When true, use the 32×32 small-tile matmul pipeline instead of 64×64. - #[serde(default)] - pub use_small_tiles: bool, - /// Shape-specialized scalar convolution, with this many K elements staged. - /// Set by measured selection; None keeps the shared uniform-parameter kernel. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub conv_k_tile: Option, - /// The K-split GEMV quantizes its activation row to Q8_1 and uses - /// integer dot products. GGML Q4_0, Meganeura Q8 and every K-quant - /// layout have a kernel for it. - /// - /// Set when the plan is compiled, from the weight format and - /// [`CompileOptions::quantized_activations`]; changes the numbers, so - /// it is never turned on by measurement. - #[serde(default)] - pub gemv_int_dot: bool, - /// Workgroup width and cross-lane reduction for a K-split GEMV. - /// - /// Set by measured selection; None keeps the group's initial shape. Which - /// shape wins is a property of the device — how wide its waves are, how - /// much a `workgroupBarrier` costs — so it is measured rather than - /// predicted. See [`crate::codegen::GemvShape`]. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gemv_shape: Option, /// The dispatch belongs to numerically sensitive derivative work and may /// not be promoted to a reduced-input-precision implementation. Native /// f32 cooperative kernels remain eligible. @@ -1040,12 +1052,6 @@ pub struct Dispatch { /// `$PROLOGUE_DECL` template variables from the prologue's factors. #[serde(default, skip_serializing_if = "Option::is_none")] pub matmul_prologue: Option, - /// Legacy fields — kept for serde backward compat of cached plans. - /// New code uses `matmul_epilogue` instead. - #[serde(default)] - pub epilogue: Vec, - #[serde(default)] - pub epilogue_buffers: Vec, /// Human-readable label for profiling (e.g. `"MatMul[50,720,960]"`). #[serde(default)] pub label: String, @@ -1055,37 +1061,113 @@ pub struct Dispatch { /// `Session::read_node` do. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub origin: Vec, - /// The scalar `(shader, workgroups)` this dispatch had before - /// cooperative-matrix promotion. Retained for diagnostics and future - /// complete cooperative candidates; the current tuner searches scalar - /// tiles only. Not serialized: cached plans are stored pre-selection and - /// re-selected per session. - #[serde(skip)] - pub scalar_fallback: Option<(ShaderEntry, [u32; 3])>, - /// When `Some`, this dispatch uses a schedule-template-generated - /// pointwise kernel. The runtime compiles a dedicated pipeline from the - /// DAG (keyed by `PointwiseDAG::hash_key`) and binds it using the same - /// `UnaryData` / `BinaryData` layout that `shader` already selects. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub pointwise: Option, - /// When `Some`, this dispatch uses a schedule-template-generated - /// reduction kernel. Mutually exclusive with `pointwise`. The runtime - /// compiles a dedicated pipeline from the kernel spec (keyed by - /// `ReductionKernel::hash_key`) and picks the binding layout based on - /// the kernel's buffer-input arity. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reduction: Option, /// Storage format of the B (weight) input buffer. #[serde(default)] pub weight_format: WeightFormat, } +/// Mutually exclusive implementations. Bindings and launch geometry live on +/// the dispatch; shader-specific configuration lives only in its variant. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum Kernel { + #[default] + Default, + SmallTile, + ScalarMatmul(crate::codegen::ScalarMatmulShape), + SpecializedConv { + k_tile: u32, + }, + Cooperative, + /// Experimental f16 hi/lo staging, not a full-range f32 implementation. + CooperativeCompensated, + Gemv { + shape: crate::codegen::GemvShape, + /// Precision policy, never enabled by measurement. + integer_dot: bool, + }, + Pointwise(PointwiseDAG), + Reduction(ReductionKernel), +} + +impl Dispatch { + pub fn use_coop(&self) -> bool { + matches!( + self.kernel, + Kernel::Cooperative | Kernel::CooperativeCompensated + ) + } + + pub fn use_coop_compensated(&self) -> bool { + matches!(self.kernel, Kernel::CooperativeCompensated) + } + + pub fn use_small_tiles(&self) -> bool { + matches!( + self.kernel, + Kernel::SmallTile + | Kernel::ScalarMatmul(crate::codegen::ScalarMatmulShape { tile_size: 32, .. }) + ) + } + + pub fn scalar_matmul(&self) -> Option { + match self.kernel { + Kernel::ScalarMatmul(shape) => Some(shape), + _ => None, + } + } + + pub fn conv_k_tile(&self) -> Option { + match self.kernel { + Kernel::SpecializedConv { k_tile } => Some(k_tile), + _ => None, + } + } + + pub fn gemv_shape(&self) -> Option { + match self.kernel { + Kernel::Gemv { shape, .. } => Some(shape), + _ => None, + } + } + + pub fn gemv_int_dot(&self) -> bool { + matches!( + self.kernel, + Kernel::Gemv { + integer_dot: true, + .. + } + ) + } + + pub fn pointwise(&self) -> Option<&PointwiseDAG> { + match self.kernel { + Kernel::Pointwise(ref dag) => Some(dag), + _ => None, + } + } + + pub fn reduction(&self) -> Option<&ReductionKernel> { + match self.kernel { + Kernel::Reduction(ref kernel) => Some(kernel), + _ => None, + } + } + + pub fn reduction_mut(&mut self) -> Option<&mut ReductionKernel> { + match self.kernel { + Kernel::Reduction(ref mut kernel) => Some(kernel), + _ => None, + } + } +} + /// Reference to a GPU buffer in the execution plan. #[derive(Clone, Debug, Default, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct BufferRef(pub u32); /// The complete execution plan: a static sequence of dispatches. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ExecutionPlan { /// Buffer sizes in bytes, indexed by BufferRef. pub buffers: Vec, @@ -1343,9 +1425,6 @@ fn fuse_row_scaled_scatters(plan: &mut ExecutionPlan) { for buffer in &dispatch.input_buffers { *reads.entry(*buffer).or_insert(0usize) += 1; } - for buffer in &dispatch.epilogue_buffers { - *reads.entry(*buffer).or_insert(0usize) += 1; - } } let mut candidate = None; @@ -1364,8 +1443,8 @@ fn fuse_row_scaled_scatters(plan: &mut ExecutionPlan) { let mul = &plan.dispatches[mul_index]; let plain_mul = mul.shader == ShaderEntry::Mul && mul.input_buffers.len() == 2 - && mul.reduction.is_none() - && match mul.pointwise.as_ref() { + && mul.reduction().is_none() + && match mul.pointwise() { None => true, Some(dag) => { dag.n_inputs == 2 @@ -1453,7 +1532,7 @@ fn fuse_row_scaled_scatters(plan: &mut ExecutionPlan) { if small_row { scatter.workgroups = [scatter.params[1].div_ceil(256), 1, 1]; } - scatter.pointwise = None; + scatter.kernel = Kernel::Default; scatter.label = format!("ScatterAddAtomicRowMul[{total}]"); // The zero entry point does not read `src`, but its shared binding // layout still requires a valid buffer. Stop it from retaining the @@ -1475,7 +1554,7 @@ fn fuse_row_scaled_scatters(plan: &mut ExecutionPlan) { /// carry the DAG the pass needs). /// /// Conservative criteria — a producer P is fused into consumer C only when: -/// 1. Both `P.pointwise` and `C.pointwise` are `Some`. +/// 1. Both `P.pointwise()` and `C.pointwise()` are `Some`. /// 2. P's output buffer is read by exactly one dispatch (C) and appears /// in no plan-level role (output/loss/param/input/constant/extra). /// 3. C's workgroups match P's (same output length). @@ -1525,15 +1604,12 @@ fn fuse_pointwise_chains(plan: &mut ExecutionPlan) { for b in &d.extra_outputs { protected.insert(*b); } - for b in &d.epilogue_buffers { - *reads.entry(*b).or_default() += 1; - } } let mut fused_any = false; for ci in 0..n { let c = &plan.dispatches[ci]; - if c.pointwise.is_none() || c.fusion_barrier { + if c.pointwise().is_none() || c.fusion_barrier { continue; } @@ -1551,7 +1627,7 @@ fn fuse_pointwise_chains(plan: &mut ExecutionPlan) { continue; } let p = &plan.dispatches[pi]; - if p.pointwise.is_none() || p.fusion_barrier { + if p.pointwise().is_none() || p.fusion_barrier { continue; } if reads.get(buf).copied().unwrap_or(0) != 1 { @@ -1589,13 +1665,9 @@ fn fuse_pointwise_chains(plan: &mut ExecutionPlan) { let producer_d = plan.dispatches[pi].clone(); let consumer_d = &mut plan.dispatches[ci]; - let p_dag = producer_d.pointwise.expect("checked above"); - let c_dag = consumer_d - .pointwise - .as_ref() - .expect("checked above") - .clone(); - let fused_dag = c_dag.fuse_input(input_idx, &p_dag); + let p_dag = producer_d.pointwise().expect("checked above"); + let c_dag = consumer_d.pointwise().expect("checked above").clone(); + let fused_dag = c_dag.fuse_input(input_idx, p_dag); // Rebuild consumer input_buffers: producer inputs, then // consumer inputs with the fused slot removed, in order. @@ -1606,7 +1678,7 @@ fn fuse_pointwise_chains(plan: &mut ExecutionPlan) { } } consumer_d.input_buffers = new_inputs; - consumer_d.pointwise = Some(fused_dag); + consumer_d.kernel = Kernel::Pointwise(fused_dag); consumer_d.origin.extend(producer_d.origin.iter().copied()); // The consumer now reads from more buffers; its ShaderEntry // (used only to pick the data layout) must reflect the new @@ -1644,9 +1716,6 @@ fn shared_pointwise_consumers_are_foldable( ) -> bool { let mut foldable_reads = 0usize; for dispatch in &plan.dispatches { - if dispatch.epilogue_buffers.contains(&buffer) { - return false; - } let occurrences = dispatch .input_buffers .iter() @@ -1658,7 +1727,7 @@ fn shared_pointwise_consumers_are_foldable( if occurrences > 1 { return false; } - let Some(ref kernel) = dispatch.reduction else { + let Some(kernel) = dispatch.reduction() else { return false; }; if kernel.input_row_repeats.iter().any(|&factor| factor != 1) { @@ -1693,9 +1762,6 @@ fn shared_embedding_consumers_are_foldable( ) -> bool { let mut foldable_reads = 0usize; for dispatch in &plan.dispatches { - if dispatch.epilogue_buffers.contains(&buffer) { - return false; - } let occurrences = dispatch .input_buffers .iter() @@ -1707,7 +1773,7 @@ fn shared_embedding_consumers_are_foldable( if occurrences > 1 { return false; } - let Some(ref kernel) = dispatch.reduction else { + let Some(kernel) = dispatch.reduction() else { return false; }; if kernel.input_row_repeats.iter().any(|&factor| factor != 1) { @@ -1797,9 +1863,6 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { for b in &d.input_buffers { *reads.entry(*b).or_default() += 1; } - for b in &d.epilogue_buffers { - *reads.entry(*b).or_default() += 1; - } } (producer, reads) }; @@ -1812,7 +1875,7 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { 'outer: for ci in 0..plan.dispatches.len() { let c = &plan.dispatches[ci]; - let Some(kernel) = c.reduction.as_ref() else { + let Some(kernel) = c.reduction() else { continue; }; // Phase 1 invariant: no gather streams yet, so input_buffers @@ -1838,7 +1901,7 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { continue; } let p = &plan.dispatches[pi]; - if p.pointwise.is_none() || p.reduction.is_some() || p.fusion_barrier { + if p.pointwise().is_none() || p.fusion_barrier { continue; } // Producer must cover the per-element domain (outer*inner). @@ -1859,12 +1922,12 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { } let producer_d = plan.dispatches[pi].clone(); - let p_dag = producer_d.pointwise.clone().expect("checked"); + let p_dag = producer_d.pointwise().expect("checked"); let c = &mut plan.dispatches[ci]; - let kernel = c.reduction.as_mut().expect("checked"); - kernel.prologue = kernel.prologue.fuse_input(s as u8, &p_dag); + let kernel = c.reduction_mut().expect("checked"); + kernel.prologue = kernel.prologue.fuse_input(s as u8, p_dag); for prologue in &mut kernel.extra_prologues { - *prologue = prologue.fuse_input(s as u8, &p_dag); + *prologue = prologue.fuse_input(s as u8, p_dag); } // The reduction epilogue sees the same per-element streams // before its per-column and reduced-value inputs. If it @@ -1872,7 +1935,7 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { // well; otherwise its declared arity no longer matches // n_per_elem and lowering aborts. if let Some(epilogue) = kernel.epilogue.as_mut() { - epilogue.dag = epilogue.dag.fuse_input(s as u8, &p_dag); + epilogue.dag = epilogue.dag.fuse_input(s as u8, p_dag); } kernel.n_per_elem = new_n_per_elem as u8; kernel.gather_elem = Vec::new(); @@ -1906,7 +1969,7 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { 'outer2: for ci in 0..plan.dispatches.len() { let c = &plan.dispatches[ci]; - let Some(kernel) = c.reduction.as_ref() else { + let Some(kernel) = c.reduction() else { continue; }; if kernel.input_row_repeats.iter().any(|&factor| factor != 1) { @@ -1947,8 +2010,8 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { // Plain Embedding dispatch: indexed load, gathered axis == // reduced axis (params [seq, hidden] = [outer, inner]). let is_embedding = p.shader == ShaderEntry::Embedding - && p.reduction.is_none() - && p.pointwise.is_none() + && p.reduction().is_none() + && p.pointwise().is_none() && p.params.first().copied() == Some(outer) && p.params.get(1).copied() == Some(inner) && p.input_buffers.len() == 2; @@ -1966,7 +2029,7 @@ fn fuse_reduction_chains(plan: &mut ExecutionPlan) { let producer_origin = p.origin.clone(); let c = &mut plan.dispatches[ci]; - let kernel = c.reduction.as_mut().expect("checked"); + let kernel = c.reduction_mut().expect("checked"); if kernel.gather_elem.is_empty() { kernel.gather_elem = vec![false; per_elem]; } @@ -2051,14 +2114,17 @@ pub fn fuse_rmsnorm_into_gemv(plan: &mut ExecutionPlan) { // Folding is a variant of MatMulGemv: the fused pipeline is // `Variant::GemvRmsNorm` — or its int-dot form, `GemvRmsNormIntDot` // — keyed by weight format and shape, so a packed GEMV keeps its - // decoder. Only the plain GEMV path is rewritten; fused-add and BT - // stay separate kernels. Int-dot GEMVs fold too: their activation + // decoder. Plain GEMV and dense transposed-B GEMV are eligible; + // fused-add remains separate. Int-dot GEMVs fold too: their activation // quantizer runs inside the same workgroup as the prologue, so the // integer arithmetic sees exactly the row the unfused path would. if consumers.is_empty() || consumers.iter().any(|&c| { let d = &plan.dispatches[c]; - d.shader != ShaderEntry::MatMulGemv || d.input_buffers.first() != Some(&normed) + !matches!( + d.shader, + ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvBT + ) || d.input_buffers.first() != Some(&normed) }) { continue; @@ -2151,8 +2217,7 @@ pub fn fuse_rmsnorm_into_add(plan: &mut ExecutionPlan) { // RmsNormAdd has no pointwise epilogue, so replacing a non-canonical // Add dispatch would silently drop the absorbed consumer. let has_fused_epilogue = add - .pointwise - .as_ref() + .pointwise() .is_some_and(|pointwise| pointwise != &plain_add); if add.shader != ShaderEntry::Add || add.output_buffer == normed || has_fused_epilogue { continue; @@ -2181,7 +2246,7 @@ pub fn fuse_rmsnorm_into_add(plan: &mut ExecutionPlan) { // rewrite, pipeline selection prefers that kernel over RmsNormAdd // and silently drops the residual. Route the fused operation through // its dedicated shader and restore its one-workgroup-per-row shape. - d.reduction = None; + d.kernel = Kernel::Default; d.workgroups = [d.params[0], 1, 1]; d.input_buffers.push(residual); d.output_buffer = out; @@ -2229,7 +2294,7 @@ pub fn fuse_rmsnorm_prologues(plan: &mut ExecutionPlan) { // single-consumer RmsNorm. The scalar pipeline cannot execute a // matmul prologue, so only transform dispatches that runtime policy // has already selected for cooperative matrices. - if !d.use_coop + if !d.use_coop() || !matches!( d.shader, ShaderEntry::MatMul @@ -2243,7 +2308,10 @@ pub fn fuse_rmsnorm_prologues(plan: &mut ExecutionPlan) { // Skip GEMV variants (M=1) — those use a different kernel path. if matches!( d.shader, - ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvAdd | ShaderEntry::MatMulGemvBT + ShaderEntry::MatMulGemv + | ShaderEntry::MatMulGemvAdd + | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd ) { continue; } @@ -2292,8 +2360,7 @@ pub fn fuse_rmsnorm_prologues(plan: &mut ExecutionPlan) { output_buffer: rsqrt_buf, extra_outputs: vec![], params: vec![rows, cols, eps_bits, 0], - use_coop: false, - use_small_tiles: false, + origin: norm_origin.clone(), label: norm_label, ..Default::default() @@ -2385,11 +2452,11 @@ fn fuse_epilogues(plan: &mut ExecutionPlan) { // DAG when one is present instead of silently changing its meaning. use crate::schedule::Pw; let d_shader = d.shader.clone(); - let (pw_op, legacy_op) = match d_shader { - ShaderEntry::Relu => (Pw::Relu(0), EpilogueOp::Relu), - ShaderEntry::Sigmoid => (Pw::Sigmoid(0), EpilogueOp::Sigmoid), - ShaderEntry::Neg => (Pw::Neg(0), EpilogueOp::Neg), - ShaderEntry::Silu => (Pw::Silu(0), EpilogueOp::Silu), + let pw_op = match d_shader { + ShaderEntry::Relu => Pw::Relu(0), + ShaderEntry::Sigmoid => Pw::Sigmoid(0), + ShaderEntry::Neg => Pw::Neg(0), + ShaderEntry::Silu => Pw::Silu(0), _ => continue, }; let canonical_dag = PointwiseDAG { @@ -2397,10 +2464,10 @@ fn fuse_epilogues(plan: &mut ExecutionPlan) { ops: vec![Pw::LoadInput(0), pw_op], output: 1, }; - let (epilogue_dag, has_legacy_equivalent) = match d.pointwise.as_ref() { - Some(dag) if dag.n_inputs == 1 => (dag.clone(), *dag == canonical_dag), + let epilogue_dag = match d.pointwise() { + Some(dag) if dag.n_inputs == 1 => dag.clone(), Some(_) => continue, - None => (canonical_dag, true), + None => canonical_dag, }; let primary_buf = d.input_buffers[0]; let elem_output = d.output_buffer; @@ -2426,7 +2493,7 @@ fn fuse_epilogues(plan: &mut ExecutionPlan) { | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd ); - if !is_matmul || prod.use_coop { + if !is_matmul || prod.use_coop() { continue; } @@ -2451,13 +2518,6 @@ fn fuse_epilogues(plan: &mut ExecutionPlan) { inputs: vec![], }); } - // Maintain the old flat field only when it describes the same op. - // New readers use `matmul_epilogue`; inventing a legacy Relu for a - // Clamp DAG is actively misleading to old diagnostics and caches. - if has_legacy_equivalent { - dispatches[prod_idx].epilogue.push(legacy_op); - } - dispatches[prod_idx].requires_full_precision |= consumer_requires_full_precision; dispatches[prod_idx].output_buffer = elem_output; let absorbed_origin = dispatches[i].origin.clone(); @@ -2556,8 +2616,14 @@ fn binary_shader_to_pointwise(shader: &ShaderEntry) -> Option { const MAX_COMPUTE_WORKGROUPS_PER_DIMENSION: u32 = 65_535; -/// Tile a scalar matmul across Y and Z without exceeding the portable -/// per-dimension compute-dispatch limit. +/// Tile row-GEMV workgroups across X and Y within the portable limit. +pub(crate) fn row_gemv_workgroups(n: u32) -> [u32; 3] { + // Large vocabularies can exceed the portable X workgroup limit. Spread + // rows over Y as well; the kernel flattens the actual dispatch grid. + let y = n.div_ceil(65_535).max(1); + [n.div_ceil(y), y, 1] +} + fn matmul_workgroups(m: u32, n: u32, tile: u32) -> [u32; 3] { let columns = n.div_ceil(tile); assert!( @@ -2603,6 +2669,9 @@ impl<'a> Compiler<'a> { coop_caps: crate::codegen::CoopCaps, allow_reduced_precision_attention_backward: bool, ) -> Self { + if let Some(shape) = options.gemv_shape { + shape.validate(); + } Self { graph, plan: ExecutionPlan { @@ -2792,9 +2861,8 @@ impl<'a> Compiler<'a> { output_buffer: output, extra_outputs: vec![], params: vec![rows, inner, 1.0_f32.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -2810,8 +2878,7 @@ impl<'a> Compiler<'a> { output_buffer: output, extra_outputs: vec![], params: vec![total, inner, 1, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -2936,6 +3003,7 @@ impl<'a> Compiler<'a> { ShaderEntry::MatMulAT | ShaderEntry::MatMulBT | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd => { format!( @@ -3052,7 +3120,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - pointwise: Some(PointwiseDAG { + kernel: Kernel::Pointwise(PointwiseDAG { n_inputs: 1, ops: vec![Pw::LoadInput(0)], output: 0, @@ -3089,25 +3157,9 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, k, n, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, - // The int-dot kernels read GGML's split-nibble Q4_0, - // Meganeura's whole-word Q8, and every K-quant - // superblock layout directly; no other weight - // format has a layout they can feed without a - // shuffle. - gemv_int_dot: self.options.quantized_activations - && matches!( - wf, - WeightFormat::Q40 - | WeightFormat::Q8 - | WeightFormat::Q4K - | WeightFormat::Q5K - | WeightFormat::Q6K - | WeightFormat::Q3K - ), - gemv_shape: self.options.gemv_shape, + kernel: self.options.gemv_kernel(ShaderGroup::MatMulGemv, wf), ..Default::default() }); } else { @@ -3118,8 +3170,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, k, n, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -3143,8 +3194,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, k, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -3184,15 +3234,16 @@ impl<'a> Compiler<'a> { // sees packed data. self.plan.dispatches.push(Dispatch { shader: ShaderEntry::MatMulGemvBT, - workgroups: [n, 1, 1], + workgroups: row_gemv_workgroups( + n.div_ceil(self.options.gemv_shape.map_or(1, |s| s.bt_rows)), + ), input_buffers: vec![a, b], output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, k, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, - gemv_shape: self.options.gemv_shape, + kernel: self.options.gemv_kernel(ShaderGroup::MatMulGemvBT, wf), ..Default::default() }); } else { @@ -3203,8 +3254,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, k, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -3270,7 +3320,11 @@ impl<'a> Compiler<'a> { input_buffers: vec![a, b], output_buffer: out_buf, params: vec![m, n, k, groups], - use_small_tiles: small, + kernel: if small { + Kernel::SmallTile + } else { + Kernel::Default + }, ..Default::default() }); } @@ -3294,22 +3348,9 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, k, n, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, - // Keep a fused residual from silently disabling the - // requested Q8_1-activation path. - gemv_int_dot: self.options.quantized_activations - && matches!( - wf, - WeightFormat::Q40 - | WeightFormat::Q8 - | WeightFormat::Q4K - | WeightFormat::Q5K - | WeightFormat::Q6K - | WeightFormat::Q3K - ), - gemv_shape: self.options.gemv_shape, + kernel: self.options.gemv_kernel(ShaderGroup::MatMulGemvAdd, wf), ..Default::default() }); } else { @@ -3320,8 +3361,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, k, n, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -3346,8 +3386,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, k, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -3358,8 +3397,8 @@ impl<'a> Compiler<'a> { let a = self.get_buffer(node.inputs[0]); let b = self.get_buffer(node.inputs[1]); let d = self.get_buffer(node.inputs[2]); - // Same block-axis mismatch as `Op::MatMulBT`. Greedy mode - // rewrites `Add(MatMulBT, ?)` into this op before compile, + // Same block-axis mismatch as `Op::MatMulBT`. Graph rewrites + // can fuse `Add(MatMulBT, ?)` into this op before compile, // so the unfused assert would never see a quantized B. assert!( !WeightFormat::from_dtype(self.graph.node(node.inputs[1]).ty.dtype) @@ -3373,16 +3412,31 @@ impl<'a> Compiler<'a> { let m = a_shape[0] as u32; let k = a_shape[1] as u32; let n = b_shape[0] as u32; + let gemv = m == 1 && k.is_multiple_of(4); self.plan.dispatches.push(Dispatch { - shader: ShaderEntry::FusedMatMulBTAdd, - workgroups: matmul_workgroups(m, n, 64), + shader: if gemv { + ShaderEntry::MatMulGemvBTAdd + } else { + ShaderEntry::FusedMatMulBTAdd + }, + workgroups: if gemv { + row_gemv_workgroups( + n.div_ceil(self.options.gemv_shape.map_or(1, |s| s.bt_rows)), + ) + } else { + matmul_workgroups(m, n, 64) + }, input_buffers: vec![a, b, d], output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, k, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, + kernel: if gemv { + self.options.gemv_kernel(ShaderGroup::MatMulGemvBTAdd, wf) + } else { + Kernel::Default + }, ..Default::default() }); } @@ -3409,8 +3463,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, bias_len, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3427,8 +3480,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, scale_len, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3467,7 +3519,7 @@ impl<'a> Compiler<'a> { input_buffers: vec![input], output_buffer: out_buf, params: vec![len, 0, 0, 0], - pointwise: Some(pointwise), + kernel: Kernel::Pointwise(pointwise), ..Default::default() }); } @@ -3494,7 +3546,7 @@ impl<'a> Compiler<'a> { input_buffers: vec![input], output_buffer: out_buf, params: vec![len, 0, 0, 0], - pointwise: Some(pointwise), + kernel: Kernel::Pointwise(pointwise), ..Default::default() }); } @@ -3512,7 +3564,7 @@ impl<'a> Compiler<'a> { input_buffers: vec![input], output_buffer: out_buf, params: vec![len, 0, 0, 0], - pointwise: Some(pointwise), + kernel: Kernel::Pointwise(pointwise), ..Default::default() }); } @@ -3527,7 +3579,7 @@ impl<'a> Compiler<'a> { input_buffers: vec![grad_output, input], output_buffer: out_buf, params: vec![len, 0, 0, 0], - pointwise: Some(pointwise), + kernel: Kernel::Pointwise(pointwise), ..Default::default() }); } @@ -3542,8 +3594,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3558,8 +3609,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3578,8 +3628,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3598,8 +3647,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, 3, u32::from(reverse)], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3619,8 +3667,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, n, 2, offset as u32], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3689,9 +3736,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, inner, 1.0_f32.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -3761,9 +3807,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, inner, 1.0_f32.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -3803,9 +3848,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![total, inner, 1.0_f32.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -3828,8 +3872,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![total, inner, pairs, mode], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3877,9 +3920,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![vector_rows, inner, 1.0_f32.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -3899,8 +3941,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, features, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3936,8 +3977,7 @@ impl<'a> Compiler<'a> { output_buffer: softmax_buf, extra_outputs: vec![], params: vec![batch, features, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); let len = batch * features; @@ -3948,8 +3988,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3976,8 +4015,7 @@ impl<'a> Compiler<'a> { output_buffer: grad_buf, extra_outputs: vec![out_buf], params: vec![batch, features, write_grad, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -3994,8 +4032,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4012,8 +4049,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![m, n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4042,8 +4078,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![out_len, half_n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4059,8 +4094,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![out_len, half_n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4078,8 +4112,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![grad_out_len, half_n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4096,8 +4129,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![grad_out_len, half_n, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4118,8 +4150,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4147,8 +4178,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, hidden, 0, 0], - use_coop: false, - use_small_tiles: false, + weight_format: wf, ..Default::default() }); @@ -4166,8 +4196,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4188,13 +4217,12 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![total, 0, 0, 0], - pointwise: Some(PointwiseDAG { + kernel: Kernel::Pointwise(PointwiseDAG { n_inputs: 1, ops: vec![Pw::const_f32(0.0)], output: 0, }), - use_coop: false, - use_small_tiles: false, + ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -4206,8 +4234,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params, - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -4218,8 +4245,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params, - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4246,8 +4272,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, dim, theta.to_bits(), 0, head_dim, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else if node.inputs.len() == 2 { @@ -4260,8 +4285,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, dim, theta.to_bits(), 0, head_dim, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -4272,8 +4296,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, dim, theta.to_bits(), pos_offset, head_dim, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4292,8 +4315,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, dim, theta.to_bits(), 0, head_dim, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4324,8 +4346,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![lse_buf], params: vec![seq, 0, (num_heads << 16) | num_kv_heads, head_dim], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4357,8 +4378,7 @@ impl<'a> Compiler<'a> { head_dim, window_size, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4379,8 +4399,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![seq, dim, theta.to_bits(), pos_offset, head_dim, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4411,8 +4430,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, spatial, num_groups, eps.to_bits(), 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -4438,9 +4456,8 @@ impl<'a> Compiler<'a> { output_buffer: partials, extra_outputs: vec![], params: vec![slices, group_size / chunks, 0, 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(group_norm_stats_kernel()), + + kernel: Kernel::Reduction(group_norm_stats_kernel()), ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -4450,8 +4467,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params, - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4479,8 +4495,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, spatial, num_groups, eps.to_bits(), 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -4504,9 +4519,8 @@ impl<'a> Compiler<'a> { output_buffer: partials, extra_outputs: vec![], params: vec![slices, group_size / chunks, 0, 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(group_norm_stats_kernel()), + + kernel: Kernel::Reduction(group_norm_stats_kernel()), ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -4516,8 +4530,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params, - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4541,8 +4554,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, spatial, num_groups, eps.to_bits(), 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4564,8 +4576,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, spatial, num_groups, eps.to_bits(), 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4586,8 +4597,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels_a, channels_b, spatial], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4607,8 +4617,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels_a, channels_b, spatial], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4628,8 +4637,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels_a, channels_b, spatial], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4649,8 +4657,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, in_h, in_w], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4670,8 +4677,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, channels, in_h, in_w], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4736,8 +4742,7 @@ impl<'a> Compiler<'a> { out_w, padding_w, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } // else (non-1x1 conv) @@ -4754,8 +4759,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, spatial, channels, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4771,8 +4775,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, spatial, channels, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4803,8 +4806,7 @@ impl<'a> Compiler<'a> { batch, channels, in_h, in_w, kernel_h, kernel_w, stride, padding_h, out_h, out_w, padding_w, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4842,8 +4844,7 @@ impl<'a> Compiler<'a> { output_buffer: weight_xform, extra_outputs: vec![], params: vec![out_channels, in_channels, 0, 0, 0, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); @@ -4864,8 +4865,7 @@ impl<'a> Compiler<'a> { tiles_w, total_tiles, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); @@ -4878,8 +4878,7 @@ impl<'a> Compiler<'a> { output_buffer: mm_out_buf, extra_outputs: vec![], params: vec![out_channels, total_tiles, in_channels, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); @@ -4900,8 +4899,7 @@ impl<'a> Compiler<'a> { total_tiles, 0, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -4959,8 +4957,7 @@ impl<'a> Compiler<'a> { out_w, padding_w, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5019,8 +5016,7 @@ impl<'a> Compiler<'a> { out_w, padding_w, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5039,8 +5035,7 @@ impl<'a> Compiler<'a> { output_buffer: cache, extra_outputs: vec![], params: vec![dim, 0, 0, 0], // kv_pos read from input buffer at runtime - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5063,8 +5058,7 @@ impl<'a> Compiler<'a> { output_buffer: cache, extra_outputs: vec![], params: vec![dim, block_len, max_seq, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5091,8 +5085,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![q_seq, num_heads, num_kv_heads, head_dim], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5110,15 +5103,30 @@ impl<'a> Compiler<'a> { let valid_len_input = self.get_buffer(node.inputs[4]); let block_len = self.graph.node(node.inputs[0]).ty.shape[0] as u32; let max_seq = self.graph.node(node.inputs[1]).ty.shape[0] as u32; - // Flash-decoding split-K: for a single decode query past a - // threshold, split the KV range across workgroups per head - // and combine their online-softmax partials. Batched prefill - // stays on CachedBlockAttention: the split shader's slices - // are a decode-only layout, and applying them independently - // to every query produces incorrect causal attention. - if block_len == 1 && max_seq > 64 { - let splits = (max_seq.div_ceil(32)).clamp(2, 16); - let chunk = max_seq.div_ceil(splits); + let mut params = CachedBlockAttentionParams { + window_size, + num_heads, + num_kv_heads, + head_dim, + block_len, + max_seq, + splits: 0, + chunk: 0, + }; + let splits = self.options.cached_attention_splits.unwrap_or_else(|| { + if block_len == 1 && max_seq > 64 { + (max_seq.div_ceil(32)).clamp(2, 16) + } else { + 1 + } + }); + assert!( + (1..=16).contains(&splits), + "cached attention needs 1..=16 splits" + ); + if splits > 1 { + params.splits = splits; + params.chunk = max_seq.div_ceil(splits); let scratch_idx = self.plan.buffers.len() as u32; self.plan.buffers.push( block_len as usize @@ -5134,18 +5142,8 @@ impl<'a> Compiler<'a> { input_buffers: vec![q, k_cache, v_cache, kv_pos_input, valid_len_input], output_buffer: partials, extra_outputs: vec![], - params: vec![ - window_size, - num_heads, - num_kv_heads, - head_dim, - block_len, - max_seq, - splits, - chunk, - ], - use_coop: false, - use_small_tiles: false, + params: params.to_words(), + ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -5154,18 +5152,8 @@ impl<'a> Compiler<'a> { input_buffers: vec![partials], output_buffer: out_buf, extra_outputs: vec![], - params: vec![ - window_size, - num_heads, - num_kv_heads, - head_dim, - block_len, - max_seq, - splits, - chunk, - ], - use_coop: false, - use_small_tiles: false, + params: params.to_words(), + ..Default::default() }); } else { @@ -5175,18 +5163,8 @@ impl<'a> Compiler<'a> { input_buffers: vec![q, k_cache, v_cache, kv_pos_input, valid_len_input], output_buffer: out_buf, extra_outputs: vec![], - params: vec![ - window_size, - num_heads, - num_kv_heads, - head_dim, - block_len, - max_seq, - 0, - 0, - ], - use_coop: false, - use_small_tiles: false, + params: params.to_words(), + ..Default::default() }); } @@ -5219,8 +5197,7 @@ impl<'a> Compiler<'a> { 0, 0, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5238,8 +5215,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![cols, rows, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5269,8 +5245,7 @@ impl<'a> Compiler<'a> { batch, channels, in_h, in_w, kernel_h, kernel_w, stride, padding, out_h, out_w, 0, 0, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5285,8 +5260,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![channels, spatial, total_out, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5304,8 +5278,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![total, spatial, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5334,8 +5307,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5361,8 +5333,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![lse_buf], params: vec![seq, seq, (num_heads << 16) | num_kv_heads, head_dim], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5387,8 +5358,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![lse_buf], params: vec![q_seq, kv_seq, (num_heads << 16) | num_kv_heads, head_dim], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5413,8 +5383,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![lse_buf], params: vec![q_seq, kv_seq, (num_heads << 16) | num_kv_heads, head_dim], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5490,8 +5459,7 @@ impl<'a> Compiler<'a> { head_dim, window_size, ], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5573,8 +5541,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![dv_buf], params: attention_params, - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5606,8 +5573,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5624,8 +5590,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5642,8 +5607,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5667,8 +5631,7 @@ impl<'a> Compiler<'a> { output_buffer: temp_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -5678,8 +5641,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -5691,8 +5653,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5721,8 +5682,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), lanes_per_row], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5744,8 +5704,7 @@ impl<'a> Compiler<'a> { output_buffer: temp_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); self.plan.dispatches.push(Dispatch { @@ -5755,8 +5714,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, 0, 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } else { @@ -5767,8 +5725,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5788,8 +5745,7 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, + ..Default::default() }); } @@ -5868,9 +5824,8 @@ impl<'a> Compiler<'a> { output_buffer: row_max, extra_outputs: vec![], params: vec![batch, features, 0, 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(max_kernel), + + kernel: Kernel::Reduction(max_kernel), ..Default::default() }); @@ -5927,9 +5882,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![batch, features, 0, 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(sum_kernel), + + kernel: Kernel::Reduction(sum_kernel), ..Default::default() }); } @@ -6021,9 +5975,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -6112,9 +6065,8 @@ impl<'a> Compiler<'a> { // bits in dispatch metadata as well: runtime may rewrite this // dispatch to RmsNormRsqrt for a cooperative matmul prologue. params: vec![rows, cols, eps.to_bits(), 0], - use_coop: false, - use_small_tiles: false, - reduction: Some(kernel), + + kernel: Kernel::Reduction(kernel), ..Default::default() }); } @@ -6129,9 +6081,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, - pointwise: Some(PointwiseDAG { + + kernel: Kernel::Pointwise(PointwiseDAG { n_inputs: 1, ops: vec![Pw::LoadInput(0), op], output: 1, @@ -6155,9 +6106,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, - pointwise, + + kernel: pointwise.map_or(Kernel::Default, Kernel::Pointwise), ..Default::default() }); } @@ -6178,9 +6128,8 @@ impl<'a> Compiler<'a> { output_buffer: out_buf, extra_outputs: vec![], params: vec![len, 0, 0, 0], - use_coop: false, - use_small_tiles: false, - pointwise, + + kernel: pointwise.map_or(Kernel::Default, Kernel::Pointwise), ..Default::default() }); } @@ -6290,7 +6239,13 @@ mod tests { // MatMul with Relu fused into epilogue (epilogue fusion pass) assert_eq!(plan.dispatches.len(), 1); assert_eq!(plan.dispatches[0].shader, ShaderEntry::MatMul); - assert_eq!(plan.dispatches[0].epilogue, vec![EpilogueOp::Relu]); + assert_eq!( + plan.dispatches[0].matmul_epilogue.as_ref().unwrap().dag.ops, + [ + crate::schedule::Pw::LoadInput(0), + crate::schedule::Pw::Relu(0) + ] + ); } #[test] @@ -6309,7 +6264,7 @@ mod tests { assert_eq!(plan.dispatches[1].shader, ShaderEntry::Sigmoid); assert_eq!(plan.dispatches[2].shader, ShaderEntry::Neg); assert_eq!(plan.dispatches[3].shader, ShaderEntry::Relu); - assert!(plan.dispatches[3].pointwise.is_some()); + assert!(plan.dispatches[3].pointwise().is_some()); // All unary ops: params = [len, 0, 0, 0] for d in &plan.dispatches { assert_eq!(d.params[0], 32); // 4*8 @@ -6335,7 +6290,7 @@ mod tests { assert_eq!(copy.workgroups, [1, 1, 1]); assert_eq!(copy.input_buffers.len(), 1); assert_ne!(copy.input_buffers[0], copy.output_buffer); - assert!(copy.pointwise.is_some()); + assert!(copy.pointwise().is_some()); assert!(copy.fusion_barrier); let split = &plan.dispatches[1]; @@ -6406,14 +6361,7 @@ mod tests { let narrow_plan = compile(&narrow); let narrow_dispatch = &narrow_plan.dispatches[0]; assert_eq!(narrow_dispatch.workgroups, [1, 1, 1]); - assert_eq!( - narrow_dispatch - .reduction - .as_ref() - .unwrap() - .rows_per_workgroup, - 256 - ); + assert_eq!(narrow_dispatch.reduction().unwrap().rows_per_workgroup, 256); let mut wide = Graph::new(); let input = wide.input("input", &[100, 33]); @@ -6422,10 +6370,7 @@ mod tests { let wide_plan = compile(&wide); let wide_dispatch = &wide_plan.dispatches[0]; assert_eq!(wide_dispatch.workgroups, [100, 1, 1]); - assert_eq!( - wide_dispatch.reduction.as_ref().unwrap().rows_per_workgroup, - 1 - ); + assert_eq!(wide_dispatch.reduction().unwrap().rows_per_workgroup, 1); let mut product = Graph::new(); let a = product.input("a", &[100, 9]); @@ -6439,7 +6384,7 @@ mod tests { 1, "narrow reductions should fold their pointwise producer" ); - let product_kernel = product_plan.dispatches[0].reduction.as_ref().unwrap(); + let product_kernel = product_plan.dispatches[0].reduction().unwrap(); assert_eq!(product_kernel.n_per_elem, 2); } @@ -6453,7 +6398,7 @@ mod tests { let forward_plan = compile(&forward); assert_eq!(forward_plan.dispatches.len(), 1); let reduction = &forward_plan.dispatches[0]; - assert!(reduction.reduction.is_some()); + assert!(reduction.reduction().is_some()); assert_eq!(reduction.params[..2], [100, 3]); assert_eq!(reduction.input_buffers.len(), 1); @@ -6477,7 +6422,7 @@ mod tests { let non_unit_plan = compile(&non_unit); assert_eq!(non_unit_plan.dispatches.len(), 1); assert_eq!(non_unit_plan.dispatches[0].shader, ShaderEntry::MatMul); - assert!(non_unit_plan.dispatches[0].reduction.is_none()); + assert!(non_unit_plan.dispatches[0].reduction().is_none()); } #[test] @@ -6510,7 +6455,7 @@ mod tests { let narrow_plan = compile(&narrow); let forward = &narrow_plan.dispatches[0]; assert_eq!(forward.workgroups, [2, 1, 1]); - assert_eq!(forward.reduction.as_ref().unwrap().rows_per_workgroup, 64); + assert_eq!(forward.reduction().unwrap().rows_per_workgroup, 64); let mut narrow_grad = Graph::new(); let dy = narrow_grad.input("dy", &[100, 3]); @@ -6683,7 +6628,7 @@ mod tests { assert_eq!(plan.dispatches[0].params[1], 10); // features/inner for dispatch in &plan.dispatches { assert_eq!(dispatch.workgroups, [7, 1, 1]); - assert_eq!(dispatch.reduction.as_ref().unwrap().rows_per_workgroup, 16); + assert_eq!(dispatch.reduction().unwrap().rows_per_workgroup, 16); } } @@ -6698,7 +6643,7 @@ mod tests { let plan = compile(&g); assert!(!plan.dispatches.is_empty()); for dispatch in &plan.dispatches { - if let Some(reduction) = dispatch.reduction.as_ref() + if let Some(reduction) = dispatch.reduction() && let Some(epilogue) = reduction.epilogue.as_ref() { assert_eq!( @@ -6744,7 +6689,7 @@ mod tests { let softmax_dispatches = plan .dispatches .iter() - .filter(|d| d.shader == ShaderEntry::Softmax || d.reduction.is_some()) + .filter(|d| d.shader == ShaderEntry::Softmax || d.reduction().is_some()) .count(); assert_eq!( softmax_dispatches, 0, @@ -6818,11 +6763,10 @@ mod tests { } #[test] - fn horizontal_fusion_drops_unpacked_fallback_and_preserves_precision() { + fn horizontal_fusion_preserves_precision() { let mut dispatches = vec![mm_dispatch(0, 1, 2, 1, 32), mm_dispatch(0, 3, 4, 1, 32)]; for d in &mut dispatches { - d.use_coop = true; - d.scalar_fallback = Some((d.shader.clone(), [1, 1, 1])); + d.kernel = crate::compile::Kernel::Cooperative; } dispatches[1].requires_full_precision = true; let mut groups = Vec::new(); @@ -6834,9 +6778,8 @@ mod tests { let packed = &dispatches[0]; assert_eq!(packed.horizontal_batch, 2); assert_eq!(packed.workgroups, [2, 2, 2]); - assert!(packed.use_coop); + assert!(packed.use_coop()); assert!(packed.requires_full_precision); - assert!(packed.scalar_fallback.is_none()); assert_eq!(packed.extra_outputs, [BufferRef(4)]); } @@ -6958,7 +6901,7 @@ mod tests { .iter() .position(|dispatch| dispatch.shader == ShaderEntry::MatMul) .expect("matmul dispatch"); - coop_plan.dispatches[matmul_index].use_coop = true; + coop_plan.dispatches[matmul_index].kernel = crate::compile::Kernel::Cooperative; fuse_rmsnorm_prologues(&mut coop_plan); let rsqrt = coop_plan @@ -7128,7 +7071,6 @@ mod tests { output: 4, } ); - assert!(plan.dispatches[0].epilogue.is_empty()); } #[test] @@ -7333,12 +7275,12 @@ mod tests { assert!( plan.dispatches .iter() - .any(|dispatch| dispatch.reduction.is_some()) + .any(|dispatch| dispatch.reduction().is_some()) ); assert!( plan.dispatches .iter() - .filter(|dispatch| dispatch.reduction.is_some()) + .filter(|dispatch| dispatch.reduction().is_some()) .all(|dispatch| dispatch.profile_family() == "normalization_reduction") ); } @@ -7364,7 +7306,7 @@ mod tests { let large = make_plan(8202); assert_eq!(large.dispatches.len(), 2); - assert!(large.dispatches[0].reduction.is_some()); + assert!(large.dispatches[0].reduction().is_some()); assert_eq!(large.dispatches[1].shader, ShaderEntry::GroupNormApply); let chunks = large.dispatches[1].params[5]; assert_eq!(chunks, 3); diff --git a/src/compile/split_k.rs b/src/compile/split_k.rs index 1e824f6c..d631164b 100644 --- a/src/compile/split_k.rs +++ b/src/compile/split_k.rs @@ -34,7 +34,7 @@ impl ExecutionPlan { let mut class = TuneClass::from_dispatch(dispatch, None) .filter(|class| { class.shader == ShaderEntry::Conv2dGradWeightGemm - && dispatch.conv_k_tile.is_none() + && dispatch.conv_k_tile().is_none() }) .ok_or(TuneError( "split-K requires an unmodified legal scalar weight gradient", @@ -230,7 +230,7 @@ mod tests { } rejected(plan.clone(), &[(index, 2)], 0); for change in [ - |d: &mut Dispatch| d.use_coop = true, + |d: &mut Dispatch| d.kernel = crate::compile::Kernel::Cooperative, |d: &mut Dispatch| d.workgroups[2] = 2, |d: &mut Dispatch| d.params[6] = 0, |d: &mut Dispatch| d.input_buffers[0] = d.output_buffer, diff --git a/src/config.rs b/src/config.rs index 7b91e108..83732582 100644 --- a/src/config.rs +++ b/src/config.rs @@ -163,14 +163,14 @@ registry! { DUMP_WGSL: "MEGANEURA_DUMP_WGSL", Text, Diagnostic, "Directory to write every generated/parsed WGSL shader into."; OPTIMIZER: "MEGANEURA_OPTIMIZER", Text, Diagnostic, - "Rewrite mode: off | greedy | egglog-windowed | egglog-outlined | egglog-whole."; + "Rewrite mode: off | egglog-windowed | egglog-outlined (default) | egglog-whole."; EGRAPH_COST: "MEGANEURA_EGRAPH_COST", Text, Diagnostic, "Extraction objective: ast-size | tensor-traffic."; EGRAPH_CUTOFF: "MEGANEURA_EGRAPH_CUTOFF", U32, Diagnostic, "Saturation segment-size ceiling (default 300)."; GREEDY_PACK_SWIGLU: "MEGANEURA_GREEDY_PACK_SWIGLU", Bool, Diagnostic, "Set to 0 to skip packing consecutive SwiGLU ops into one parameter buffer \ - during the greedy sweep."; + during graph optimization (legacy environment variable name)."; DEVICE_PARAMETERS: "MEGANEURA_DEVICE_PARAMETERS", Text, Diagnostic, "Experimental placement of unaliased parameter buffers on the device: \ 1 → device-transient, device-buddy → device (default: host-visible)."; @@ -215,7 +215,7 @@ registry! { impl OptimizeConfig { /// Read benchmark-oriented overrides while retaining production defaults. /// - /// - `MEGANEURA_OPTIMIZER=off|greedy|egglog-windowed|egglog-outlined|egglog-whole` + /// - `MEGANEURA_OPTIMIZER=off|egglog-windowed|egglog-outlined|egglog-whole` /// - `MEGANEURA_EGRAPH_COST=ast-size|tensor-traffic` /// - `MEGANEURA_EGRAPH_CUTOFF=` /// - `MEGANEURA_NO_WINOGRAD` @@ -228,13 +228,18 @@ impl OptimizeConfig { if let Some(value) = OPTIMIZER.text() { config.mode = match value.as_str() { "off" => OptimizeMode::Off, - "greedy" => OptimizeMode::Greedy, + "greedy" => { + log::warn!( + "MEGANEURA_OPTIMIZER=greedy now uses the shared outlined egglog optimizer" + ); + OptimizeMode::EgglogOutlined + } "egglog-windowed" | "windowed" => OptimizeMode::EgglogWindowed, "egglog-outlined" | "outlined" => OptimizeMode::EgglogOutlined, "egglog-whole" | "whole" => OptimizeMode::EgglogWhole, _ => { - log::warn!("unknown MEGANEURA_OPTIMIZER={value:?}; using greedy"); - OptimizeMode::Greedy + log::warn!("unknown MEGANEURA_OPTIMIZER={value:?}; using outlined egglog"); + OptimizeMode::EgglogOutlined } }; } @@ -255,7 +260,7 @@ impl OptimizeConfig { log::warn!("MEGANEURA_EGRAPH_CUTOFF must be > 0; using the default"); } } - config.greedy_pack_swiglu = GREEDY_PACK_SWIGLU.bool_or(true); + config.pack_swiglu = GREEDY_PACK_SWIGLU.bool_or(true); config } } diff --git a/src/graph.rs b/src/graph.rs index 116f2826..79e6b31e 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -831,10 +831,12 @@ pub struct Node { } /// How a derived parameter is computed from its source(s). -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum ParamTransform { /// Horizontal concatenation: interleave source columns per row. HorizontalConcat, + /// Vertical concatenation: append dense source rows without transposing. + VerticalConcat, /// Winograd F(2,3) weight transform: [Co, Ci, 3, 3] → [16, Co, Ci]. Winograd3x3 { out_channels: usize, @@ -848,9 +850,9 @@ pub enum ParamTransform { pub struct DerivedParam { /// Name of the new parameter (e.g. "gate_proj.weight+up_proj.weight") pub name: String, - /// Source parameters to concatenate horizontally: (name, cols) + /// Source parameters and their extent along the concatenated axis. pub sources: Vec<(String, usize)>, - /// Total rows (shared across all sources) + /// Rows in the derived parameter. pub rows: usize, /// How to compute this parameter from sources. pub transform: ParamTransform, diff --git a/src/lib.rs b/src/lib.rs index 09090d8e..3852a2d4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,7 @@ //! Meganeura: graph-optimized neural network framework on blade-graphics. //! //! Models are defined as declarative computation graphs, optimized with -//! greedy rewrites by default (or optional equality saturation via egglog), +//! bounded equality saturation through egglog, //! and compiled to static GPU dispatch sequences — no manual CUDA-graphing //! needed. diff --git a/src/load/gguf/graph.rs b/src/load/gguf/graph.rs index 9cdfe3e3..31002ed6 100644 --- a/src/load/gguf/graph.rs +++ b/src/load/gguf/graph.rs @@ -191,7 +191,7 @@ pub fn build( ple_size, )?; let proj_norm = g.parameter(proj_norm_name, &[ple_size]); - let layered = g.matmul(x, w_proj); + let layered = project(g, x, w_proj); let layered = g.scale(layered, 1.0 / (hidden as f32).sqrt()); let layered = g.reshape(layered, &[ple_rows, ple_size]); let layered = g.rms_norm(layered, proj_norm, eps); @@ -231,7 +231,7 @@ pub fn build( &format!("{p}.attn_q.weight"), &[hidden, q_dim], )?; - let mut q = g.matmul(attn_in, q); + let mut q = project(g, attn_in, q); // Optional rather than gated on the architecture: Qwen2 biases Q, // K and V but not the attention output, and requiring all four @@ -298,8 +298,8 @@ pub fn build( &format!("{p}.attn_v.weight"), &[hidden, kv_dim], )?; - let mut k = g.matmul(attn_in, w_k); - let mut v = g.matmul(attn_in, w_v); + let mut k = project(g, attn_in, w_k); + let mut v = project(g, attn_in, w_v); if arch.qk_norm() { k = per_head_norm( g, @@ -386,7 +386,7 @@ pub fn build( &format!("{p}.attn_output.weight"), &[q_dim, hidden], )?; - let mut attn_out = g.matmul(attn, wo); + let mut attn_out = project(g, attn, wo); attn_out = optional_bias(g, model, &format!("{p}.attn_output.bias"), attn_out, hidden)?; // Gemma2 norms each block's output before it rejoins the residual. if arch.post_block_norms() { @@ -454,10 +454,10 @@ pub fn build( let selected = g.matmul(sel, ple); let ple_layer = g.reshape(selected, &[block_size, ple_size]); - let gated = g.matmul(x, w_gate); + let gated = project(g, x, w_gate); let gated = g.gelu(gated); let mixed = g.mul(gated, ple_layer); - let mixed = g.matmul(mixed, w_proj); + let mixed = project(g, mixed, w_proj); let mixed = g.rms_norm(mixed, post, eps); x = g.add(x, mixed); } @@ -483,7 +483,7 @@ pub fn build( g.matmul_bt(last, embed) } else { let head = projection(g, model, config, OUTPUT, &[hidden, config.vocab_size])?; - g.matmul(last, head) + project(g, last, head) }; if let Some(cap) = config.final_logit_softcap { @@ -697,8 +697,8 @@ fn feed_forward( &format!("{prefix}.ffn_up.weight"), &[hidden, ffn], )?; - let gate = g.matmul(input, w_gate); - let up = g.matmul(input, w_up); + let gate = project(g, input, w_gate); + let up = project(g, input, w_up); if arch.gates_with_gelu() { // Gemma gates with GELU where llama gates with SiLU. The // multiply is the same; only the activation differs, so this @@ -716,7 +716,7 @@ fn feed_forward( &format!("{prefix}.ffn_up.weight"), &[hidden, ffn], )?; - let mut up = g.matmul(input, w_up); + let mut up = project(g, input, w_up); up = optional_bias(g, model, &format!("{prefix}.ffn_up.bias"), up, ffn)?; g.gelu(up) }; @@ -728,7 +728,7 @@ fn feed_forward( &format!("{prefix}.ffn_down.weight"), &[ffn, hidden], )?; - let mut out = g.matmul(hidden_act, w_down); + let mut out = project(g, hidden_act, w_down); out = optional_bias(g, model, &format!("{prefix}.ffn_down.bias"), out, hidden)?; Ok(out) } @@ -910,10 +910,8 @@ fn split_block_leaf(param: &str) -> Option<(&str, &str)> { /// Declare a projection weight in the dtype the file stores it in. /// -/// GGUF names dimensions fastest-first, so a `[K, N]` tensor there is a -/// `[K, N]` Meganeura parameter — the orientations agree for weights, and -/// [`super::GgufTensor::to_packed`] performs the transpose that the packed -/// layouts imply. +/// Dense weights retain GGUF's contiguous `[N, K]` rows for `MatMulBT`. +/// Block formats keep the existing `[K, N]` logical shape and decoder layout. fn projection( g: &mut Graph, model: &GgufModel, @@ -943,7 +941,21 @@ fn projection( source.tensor, tensor.dims, ))); } - Ok(parameter_of(g, name, shape, weight_dtype(tensor)?)) + let dtype = weight_dtype(tensor)?; + let shape = if matches!(dtype, DType::F32 | DType::F16) { + [shape[1], shape[0]] + } else { + *shape + }; + Ok(parameter_of(g, name, &shape, dtype)) +} + +fn project(g: &mut Graph, input: NodeId, weight: NodeId) -> NodeId { + if matches!(g.node(weight).ty.dtype, DType::F32 | DType::F16) { + g.matmul_bt(input, weight) + } else { + g.matmul(input, weight) + } } /// The dtype a projection weight is declared — and so must be *filled* — @@ -1239,11 +1251,11 @@ mod tests { let declared = declared_parameters(&g); assert_eq!( declared["blk.0.attn_k.weight"], - vec![config.hidden_size, config.kv_dim_at(0)] + vec![config.kv_dim_at(0), config.hidden_size] ); assert_eq!( declared["blk.1.attn_k.weight"], - vec![config.hidden_size, config.kv_dim_at(1)] + vec![config.kv_dim_at(1), config.hidden_size] ); let kv_heads: Vec<_> = g .nodes() @@ -1355,20 +1367,20 @@ mod tests { let declared = declared_parameters(&g); assert_eq!( declared["blk.0.attn_q.weight"], - vec![config.hidden_size, config.q_dim()] + vec![config.q_dim(), config.hidden_size] ); assert_eq!( declared["blk.0.attn_k.weight"], - vec![config.hidden_size, config.kv_dim_at(0)] + vec![config.kv_dim_at(0), config.hidden_size] ); assert_eq!( declared["blk.0.ffn_down.weight"], - vec![config.intermediate_size, config.hidden_size] + vec![config.hidden_size, config.intermediate_size] ); assert_eq!( declared[TOKEN_EMBD], vec![config.vocab_size, config.hidden_size], - "the table is a list of rows, not a [K, N] weight" + "the table retains its token-major rows" ); } @@ -1492,7 +1504,7 @@ mod tests { build(&mut g, &model, &config, 4, 16).unwrap(); assert_eq!( declared_parameters(&g)[OUTPUT], - vec![config.hidden_size, config.vocab_size] + vec![config.vocab_size, config.hidden_size] ); } diff --git a/src/load/gguf/weights.rs b/src/load/gguf/weights.rs index d2299362..7ef09b27 100644 --- a/src/load/gguf/weights.rs +++ b/src/load/gguf/weights.rs @@ -11,8 +11,8 @@ //! model rather than of the container: //! //! - The **embedding table** is dequantized to f16, because the gather has -//! no block-quantized variant. It is also the one tensor read in GGUF's -//! own row order rather than transposed — see +//! no block-quantized variant. Like dense projection weights it retains +//! GGUF's own row order — see //! [`GgufTensor::to_f32_rows`](super::GgufTensor::to_f32_rows). //! - **Gemma's norm weights** are stored centred on zero and applied as //! `1 + w`. Folding the `+1` in here keeps every shader unaware of it. @@ -255,7 +255,7 @@ fn load_one( crate::graph::DType::F32 | crate::graph::DType::F16 => { // Declared f32 or f16; either way `set_parameter` takes f32 and // the runtime narrows it if the buffer is half-width. - let values = tensor.to_f32()?; + let values = tensor.to_f32_rows()?; session.set_parameter(name, &values); report.dequantized += 1; } @@ -586,6 +586,8 @@ mod tests { for arch in ["llama", "qwen2", "qwen3", "gemma", "gemma2", "gemma4"] { let model = fixture::model(arch); let config = ModelConfig::from_gguf(&model).unwrap(); + let mut g = crate::Graph::new(); + graph::build(&mut g, &model, &config, 4, 16).unwrap(); for name in graph::parameter_names(&config) { let tensor = &model.tensors[&name]; if name == graph::TOKEN_EMBD || tensor.dims.len() == 1 { @@ -597,6 +599,13 @@ mod tests { crate::graph::DType::F32, "{arch}/{name}" ); + let param = g.nodes().iter().find(|node| matches!(&node.op, crate::graph::Op::Parameter { name: param } if param == &name)).unwrap(); + assert_eq!( + param.ty.shape, + [tensor.dims[1], tensor.dims[0]], + "{arch}/{name}" + ); + assert_eq!(param.ty.dtype, crate::graph::DType::F32); } } } diff --git a/src/memplan.rs b/src/memplan.rs index ceba6695..2c2f6b6e 100644 --- a/src/memplan.rs +++ b/src/memplan.rs @@ -391,9 +391,6 @@ fn compute_pinned( for b in &d.input_buffers { uses[b.0 as usize].read(g); } - for b in &d.epilogue_buffers { - uses[b.0 as usize].read(g); - } if let Some(epi) = d.matmul_epilogue.as_ref() { for &(b, _) in &epi.inputs { uses[b.0 as usize].read(g); @@ -418,6 +415,30 @@ fn compute_pinned( _ => {} } } + // Keep an attention operation's external bindings live across its whole + // split/combine sequence. Tuning can then replace it with a single dispatch + // without moving an output onto storage still occupied by its query. + let partials: HashMap<_, _> = plan + .dispatches + .iter() + .enumerate() + .filter(|&(_, d)| d.shader == ShaderEntry::CachedBlockAttentionSplit) + .map(|(i, d)| (d.output_buffer, i)) + .collect(); + for (i, d) in plan.dispatches.iter().enumerate() { + if d.shader != ShaderEntry::CachedBlockAttentionCombine { + continue; + } + let Some(&producer) = d.input_buffers.first().and_then(|b| partials.get(b)) else { + continue; + }; + let first = group_of[producer]; + let last = group_of[i]; + uses[d.output_buffer.0 as usize].write(first); + for b in &plan.dispatches[producer].input_buffers { + uses[b.0 as usize].read(last); + } + } // Debug aid: MEGANEURA_PIN_BUFS="3,17,25-40" force-pins logical // buffers, excluding them from aliasing. Used to bisect aliasing // corruption down to a single buffer. @@ -567,6 +588,14 @@ mod tests { "output stays host-visible" ); assert_eq!(alias.device_local.len(), alias.sizes.len()); + + // A split attention result may not reuse its query's allocation: + // the tuner must be free to evaluate the same operation unsplit. + p.dispatches[1].shader = ShaderEntry::CachedBlockAttentionSplit; + p.dispatches[2].shader = ShaderEntry::CachedBlockAttentionCombine; + let alias = plan_buffer_aliasing(&p, &groups, None); + assert_ne!(alias.map[1], alias.map[3]); + check_disjoint(&p, &groups, &alias); } #[test] diff --git a/src/optimize.rs b/src/optimize.rs index 31794861..03750a1a 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -1,8 +1,6 @@ -//! Graph optimization with deterministic greedy rewrites by default and -//! optional equality saturation through egglog. The current local rule set -//! reaches the same useful forms with greedy rewriting at much lower build -//! cost. Equality saturation retains alternatives for extraction using an -//! expression-size or estimated tensor-traffic objective. +//! Graph optimization through bounded equality saturation. Ordinary construction +//! extracts a deterministic candidate; calibrated construction retains alternatives +//! until their complete implementations can be measured. //! //! In outlined mode, graphs over `SATURATION_CUTOFF` are split into segments: //! repeated regions @@ -18,7 +16,8 @@ //! `Graph::toposort` restores it after passes that append nodes. use crate::graph::{Graph, Node, NodeId, Op, TensorType}; -use egglog::{Term, TermDag, TermId, ast::Literal}; +pub(crate) mod search; +use egglog::{Term, TermDag, TermId, ast::Literal, extract::Extractor}; use std::collections::{HashMap, HashSet}; use std::{fmt, time::Instant}; @@ -26,20 +25,16 @@ use std::{fmt, time::Instant}; /// graph is segmented (see module docs). Shared-parameter graphs create /// large e-classes that make pattern matching superlinear: the SmolVLA /// training graph (~750 nodes) takes minutes unsegmented. -const SATURATION_CUTOFF: usize = 300; +pub(crate) const SATURATION_CUTOFF: usize = 300; /// Rewrite strategy used by the graph optimizer. /// -/// Equality-saturation variants are primarily useful for controlled compiler -/// ablations and future global rewrites. `Greedy` is the production strategy: -/// for the current local rewrite set it extracts the same useful forms with -/// far less compile-time overhead. +/// Outlined saturation is the production strategy. Windowed and whole-graph +/// modes expose the region-boundary tradeoff for compiler diagnostics. #[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum OptimizeMode { /// Preserve the graph as written (apart from dead-code elimination). Off, - /// Apply the same rewrite patterns deterministically to a fixed point. - Greedy, /// Run equality saturation in fixed-size windows, without outlining. EgglogWindowed, /// Outline repeated regions, then saturate regions and residual windows. @@ -52,7 +47,6 @@ impl OptimizeMode { pub fn as_str(self) -> &'static str { match self { Self::Off => "off", - Self::Greedy => "greedy", Self::EgglogWindowed => "egglog-windowed", Self::EgglogOutlined => "egglog-outlined", Self::EgglogWhole => "egglog-whole", @@ -91,19 +85,19 @@ pub struct OptimizeConfig { /// with few channels over a large image sits near its boundary; this /// makes which side it should be on measurable without a rebuild. pub no_winograd: bool, - /// Pack consecutive SwiGLU ops into one packed parameter buffer during - /// the greedy sweep. - pub greedy_pack_swiglu: bool, + /// Pack SwiGLU projections into one derived parameter buffer. + #[serde(alias = "greedy_pack_swiglu")] + pub pack_swiglu: bool, } impl Default for OptimizeConfig { fn default() -> Self { Self { - mode: OptimizeMode::Greedy, + mode: OptimizeMode::EgglogOutlined, extraction_cost: ExtractionCost::TensorTraffic, saturation_cutoff: SATURATION_CUTOFF, no_winograd: false, - greedy_pack_swiglu: true, + pack_swiglu: true, } } } @@ -179,7 +173,7 @@ impl egglog::extract::CostModel for FusionCostModel { // back to constants that keep fused ops preferred. match name { "FusedMatMulAdd" | "FusedMatMulATAdd" | "FusedMatMulBTAdd" | "SwiGLUPacked" - | "GeGLUPacked" => 9, + | "GeGLUPacked" | "SwiGLUPackedBT" | "GeGLUPackedBT" => 9, _ => 10, } } @@ -215,9 +209,9 @@ pub struct OptimizeReport { pub nodes_after: usize, /// Fusions applied: list of (fusion_name, node_index) pairs. pub fusions_applied: Vec<(String, u32)>, - /// Wall-clock time for egglog saturation. + /// Encoding, egglog parsing/saturation, extraction and e-graph statistics. pub egglog_time: std::time::Duration, - /// Wall-clock time for extraction + term stamping. + /// Term stamping and dead-code elimination. pub extract_time: std::time::Duration, /// Repeated regions outlined for per-block saturation (0 when the /// whole graph fit under the saturation cutoff). @@ -342,7 +336,6 @@ pub(crate) fn optimize_owned_with_config( ) -> (Graph, OptimizeReport) { match config.mode { OptimizeMode::Off => optimize_off(graph, config), - OptimizeMode::Greedy => optimize_greedy(graph, config), OptimizeMode::EgglogWindowed | OptimizeMode::EgglogOutlined | OptimizeMode::EgglogWhole => { optimize_egglog(graph, config) } @@ -357,36 +350,26 @@ fn optimize_egglog(mut g: Graph, config: OptimizeConfig) -> (Graph, OptimizeRepo let segment_count = segments.len(); let max_segment_nodes = segments.iter().map(|s| s.ids.len()).max().unwrap_or(0); - let mut fusions: Vec<(String, u32)> = Vec::new(); let mut index = build_structural_index(&g); - let mut first_program = String::new(); - let mut num_eclasses = 0; - let mut num_enodes = 0; - let mut egglog_time = std::time::Duration::ZERO; - let mut extract_time = std::time::Duration::ZERO; - let mut extraction_failures = 0; + let mut report = OptimizeReport { + mode: config.mode, + extraction_cost: config.extraction_cost, + nodes_before, + outlined_regions, + segments: segment_count, + max_segment_nodes, + ..OptimizeReport::empty() + }; for seg in &segments { - process_segment( - &mut g, - seg, - &mut index, - &mut fusions, - &mut first_program, - &mut num_eclasses, - &mut num_enodes, - &mut egglog_time, - &mut extract_time, - &mut extraction_failures, - config.extraction_cost, - ); + process_segment(&mut g, seg, &mut index, &mut report, config); } let dce_start = Instant::now(); sweep_dead_nodes(&mut g); - extract_time += dce_start.elapsed(); + report.extract_time += dce_start.elapsed(); - let nodes_after = g + report.nodes_after = g .nodes() .iter() .filter(|n| !matches!(n.op, Op::Nop)) @@ -394,38 +377,20 @@ fn optimize_egglog(mut g: Graph, config: OptimizeConfig) -> (Graph, OptimizeRepo log::info!( "optimizer: {} fusions on {} nodes", - fusions.len(), - nodes_after + report.fusions_applied.len(), + report.nodes_after ); - let mut rules_fired: Vec<(String, usize)> = Vec::new(); - for fusion in &fusions { - if let Some(entry) = rules_fired.iter_mut().find(|e| e.0 == fusion.0) { + for fusion in &report.fusions_applied { + if let Some(entry) = report.rules_fired.iter_mut().find(|e| e.0 == fusion.0) { entry.1 += 1; } else { - rules_fired.push((fusion.0.clone(), 1)); + report.rules_fired.push((fusion.0.clone(), 1)); } } - for &(ref name, count) in &rules_fired { + for &(ref name, count) in &report.rules_fired { log::info!(" {}x {}", count, name); } - let report = OptimizeReport { - mode: config.mode, - extraction_cost: config.extraction_cost, - egglog_program: first_program, - num_eclasses, - num_enodes, - rules_fired, - nodes_before, - nodes_after, - fusions_applied: fusions, - egglog_time, - extract_time, - outlined_regions, - segments: segment_count, - max_segment_nodes, - extraction_failures, - }; (g.into_toposort(), report) } @@ -461,361 +426,77 @@ fn optimize_off(mut g: Graph, config: OptimizeConfig) -> (Graph, OptimizeReport) ) } -fn optimize_greedy(mut g: Graph, config: OptimizeConfig) -> (Graph, OptimizeReport) { - let nodes_before = g.nodes().len(); - let start = Instant::now(); - let mut fusions = Vec::new(); - - loop { - let before = fusions.len(); - apply_greedy_unary_simplifications(&mut g, &mut fusions); - apply_greedy_matmul_add(&mut g, &mut fusions); - apply_greedy_silu(&mut g, &mut fusions); - apply_greedy_swiglu(&mut g, &mut fusions); - if config.greedy_pack_swiglu { - apply_greedy_swiglu_packed(&mut g, &mut fusions); - } - apply_greedy_geglu(&mut g, &mut fusions); - apply_greedy_geglu_packed(&mut g, &mut fusions); - if fusions.len() == before { - break; - } - } - sweep_dead_nodes(&mut g); - let extract_time = start.elapsed(); - let nodes_after = g - .nodes() - .iter() - .filter(|node| !matches!(node.op, Op::Nop)) - .count(); - let rules_fired = summarize_fusions(&fusions); - - ( - g.into_toposort(), - OptimizeReport { - mode: config.mode, - extraction_cost: config.extraction_cost, - egglog_program: String::new(), - num_eclasses: 0, - num_enodes: 0, - rules_fired, - nodes_before, - nodes_after, - fusions_applied: fusions, - egglog_time: std::time::Duration::ZERO, - extract_time, - outlined_regions: 0, - segments: 0, - max_segment_nodes: 0, - extraction_failures: 0, - }, - ) -} - -fn summarize_fusions(fusions: &[(String, u32)]) -> Vec<(String, usize)> { - let mut summary = Vec::new(); - for fusion in fusions { - let name = &fusion.0; - if let Some(entry) = summary - .iter_mut() - .find(|entry: &&mut (String, usize)| entry.0 == *name) - { - entry.1 += 1; - } else { - summary.push((name.clone(), 1)); - } - } - summary -} - -fn apply_greedy_unary_simplifications(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (outer, inner_id) = { - let node = &graph.nodes()[id]; - if node.inputs.len() != 1 { - continue; - } - (node.op.clone(), node.inputs[0]) - }; - let inner = graph.node(inner_id); - let (replacement, label) = match (outer, inner.op.clone()) { - (Op::Neg, Op::Neg) => (inner.inputs[0], "Neg(Neg(x))→x"), - (Op::Transpose, Op::Transpose) => (inner.inputs[0], "Transpose(Transpose(x))→x"), - (Op::Relu, Op::Relu) => (inner_id, "Relu(Relu(x))→Relu(x)"), - ( - Op::RoPE { - theta, - pos_offset: 0, - .. - } - | Op::RoPEGrad { - theta, - pos_offset: 0, - .. - }, - _, - ) if theta.is_finite() - && theta > 0.0 - && graph.nodes()[id].ty.shape.first() == Some(&1) => - { - (inner_id, "RoPE(position=0)→x") - } - _ => continue, - }; - graph.nodes_mut()[id].op = Op::Identity; - graph.nodes_mut()[id].inputs = vec![replacement]; - fusions.push((label.to_string(), id as u32)); - } -} - -fn apply_greedy_matmul_add(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (lhs, rhs) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::Add) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let (mm_id, addend) = - if matches!(graph.node(lhs).op, Op::MatMul | Op::MatMulAT | Op::MatMulBT) { - (lhs, rhs) - } else if matches!(graph.node(rhs).op, Op::MatMul | Op::MatMulAT | Op::MatMulBT) { - (rhs, lhs) - } else { - continue; - }; - let mm = graph.node(mm_id); - let (op, label) = match mm.op { - Op::MatMul => (Op::FusedMatMulAdd, "MatMul+Add→FusedMatMulAdd"), - Op::MatMulAT => (Op::FusedMatMulATAdd, "MatMulAT+Add→FusedMatMulATAdd"), - Op::MatMulBT => (Op::FusedMatMulBTAdd, "MatMulBT+Add→FusedMatMulBTAdd"), - _ => unreachable!(), - }; - let inputs = vec![mm.inputs[0], mm.inputs[1], addend]; - graph.nodes_mut()[id].op = op; - graph.nodes_mut()[id].inputs = inputs; - fusions.push((label.to_string(), id as u32)); - } -} - -fn apply_greedy_silu(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (a, b) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::Mul) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let x = if matches!(graph.node(b).op, Op::Sigmoid) && graph.node(b).inputs[0] == a { - a - } else if matches!(graph.node(a).op, Op::Sigmoid) && graph.node(a).inputs[0] == b { - b - } else { - continue; - }; - graph.nodes_mut()[id].op = Op::Silu; - graph.nodes_mut()[id].inputs = vec![x]; - fusions.push(("Mul+Sigmoid→Silu".to_string(), id as u32)); - } -} - -fn apply_greedy_swiglu(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (a, b) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::Mul) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let (gate, up) = if matches!(graph.node(a).op, Op::Silu) { - (graph.node(a).inputs[0], b) - } else if matches!(graph.node(b).op, Op::Silu) { - (graph.node(b).inputs[0], a) - } else { - continue; - }; - graph.nodes_mut()[id].op = Op::SwiGLU; - graph.nodes_mut()[id].inputs = vec![gate, up]; - fusions.push(("Silu+Mul→SwiGLU".to_string(), id as u32)); - } -} - -fn apply_greedy_swiglu_packed(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (gate_id, up_id) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::SwiGLU) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let (h, wg, wu) = { - let gate = graph.node(gate_id); - let up = graph.node(up_id); - if !matches!(gate.op, Op::MatMul) - || !matches!(up.op, Op::MatMul) - || gate.inputs[0] != up.inputs[0] - { - continue; - } - (gate.inputs[0], gate.inputs[1], up.inputs[1]) - }; - let (gate_name, up_name, in_features, out_features, dtype) = { - let gate_weight = graph.node(wg); - let up_weight = graph.node(wu); - let (gate_name, up_name) = match (gate_weight.op.clone(), up_weight.op.clone()) { - (Op::Parameter { name: gate }, Op::Parameter { name: up }) => (gate, up), - _ => continue, - }; - if gate_weight.ty.shape.len() != 2 - || gate_weight.ty.shape != up_weight.ty.shape - || gate_weight.ty.dtype != up_weight.ty.dtype - || graph.node(h).ty.shape.len() != 2 - { - continue; - } - ( - gate_name, - up_name, - gate_weight.ty.shape[0], - gate_weight.ty.shape[1], - gate_weight.ty.dtype, - ) - }; - - let concat_name = format!("{gate_name}+{up_name}"); - graph.derived_params.push(crate::graph::DerivedParam { - name: concat_name.clone(), - sources: vec![(gate_name, out_features), (up_name, out_features)], - rows: in_features, - transform: crate::graph::ParamTransform::HorizontalConcat, - }); - let requires_full_precision = graph.node(id as NodeId).requires_full_precision; - let concat_w = graph.add_raw_node_with_precision( - Op::Parameter { name: concat_name }, - vec![], - TensorType::new(vec![in_features, 2 * out_features], dtype), - requires_full_precision, - ); - let m = graph.node(h).ty.shape[0]; - let wide_mm = graph.add_raw_node_with_precision( - Op::MatMul, - vec![h, concat_w], - TensorType::f32(vec![m, 2 * out_features]), - requires_full_precision, - ); - graph.nodes_mut()[id].op = Op::SwiGLUConcat; - graph.nodes_mut()[id].inputs = vec![wide_mm]; - fusions.push(( - "SwiGLU(MatMul,MatMul)→SwiGLUConcat(MatMul)".to_string(), - id as u32, - )); - } -} - -fn apply_greedy_geglu(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (a, b) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::Mul) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let (gate, up) = if matches!(graph.node(a).op, Op::Gelu) { - (graph.node(a).inputs[0], b) - } else if matches!(graph.node(b).op, Op::Gelu) { - (graph.node(b).inputs[0], a) +fn pack_glu_matmul( + graph: &mut Graph, + h: NodeId, + wg: NodeId, + wu: NodeId, + transposed: bool, + requires_full_precision: bool, +) -> Option { + let (gate, up) = (graph.node(wg), graph.node(wu)); + let Op::Parameter { + name: ref gate_name, + } = gate.op + else { + return None; + }; + let Op::Parameter { name: ref up_name } = up.op else { + return None; + }; + if gate.ty.shape.len() != 2 + || gate.ty.shape != up.ty.shape + || gate.ty.dtype != up.ty.dtype + || graph.node(h).ty.shape.len() != 2 + || (transposed + && !matches!( + gate.ty.dtype, + crate::graph::DType::F32 | crate::graph::DType::F16 + )) + { + return None; + } + let (in_features, out_features) = if transposed { + (gate.ty.shape[1], gate.ty.shape[0]) + } else { + (gate.ty.shape[0], gate.ty.shape[1]) + }; + let dtype = gate.ty.dtype; + let concat_name = format!( + "{gate_name}+{up_name}{}", + if transposed { ":rows" } else { "" } + ); + let shape = if transposed { + vec![2 * out_features, in_features] + } else { + vec![in_features, 2 * out_features] + }; + graph.derived_params.push(crate::graph::DerivedParam { + name: concat_name.clone(), + sources: vec![ + (gate_name.clone(), out_features), + (up_name.clone(), out_features), + ], + rows: shape[0], + transform: if transposed { + crate::graph::ParamTransform::VerticalConcat } else { - continue; - }; - graph.nodes_mut()[id].op = Op::GeGLU; - graph.nodes_mut()[id].inputs = vec![gate, up]; - fusions.push(("Gelu+Mul→GeGLU".to_string(), id as u32)); - } -} - -fn apply_greedy_geglu_packed(graph: &mut Graph, fusions: &mut Vec<(String, u32)>) { - let node_ids: Vec = (0..graph.nodes().len()).collect(); - for id in node_ids { - let (gate_id, up_id) = { - let node = &graph.nodes()[id]; - if !matches!(node.op, Op::GeGLU) { - continue; - } - (node.inputs[0], node.inputs[1]) - }; - let (h, wg, wu) = { - let gate = graph.node(gate_id); - let up = graph.node(up_id); - if !matches!(gate.op, Op::MatMul) - || !matches!(up.op, Op::MatMul) - || gate.inputs[0] != up.inputs[0] - { - continue; - } - (gate.inputs[0], gate.inputs[1], up.inputs[1]) - }; - let (gate_name, up_name, in_features, out_features, dtype) = { - let gate_weight = graph.node(wg); - let up_weight = graph.node(wu); - let (gate_name, up_name) = match (gate_weight.op.clone(), up_weight.op.clone()) { - (Op::Parameter { name: gate }, Op::Parameter { name: up }) => (gate, up), - _ => continue, - }; - if gate_weight.ty.shape.len() != 2 - || gate_weight.ty.shape != up_weight.ty.shape - || gate_weight.ty.dtype != up_weight.ty.dtype - || graph.node(h).ty.shape.len() != 2 - { - continue; - } - ( - gate_name, - up_name, - gate_weight.ty.shape[0], - gate_weight.ty.shape[1], - gate_weight.ty.dtype, - ) - }; - - let concat_name = format!("{gate_name}+{up_name}"); - graph.derived_params.push(crate::graph::DerivedParam { - name: concat_name.clone(), - sources: vec![(gate_name, out_features), (up_name, out_features)], - rows: in_features, - transform: crate::graph::ParamTransform::HorizontalConcat, - }); - let requires_full_precision = graph.node(id as NodeId).requires_full_precision; - let concat_w = graph.add_raw_node_with_precision( - Op::Parameter { name: concat_name }, - vec![], - TensorType::new(vec![in_features, 2 * out_features], dtype), - requires_full_precision, - ); - let m = graph.node(h).ty.shape[0]; - let wide_mm = graph.add_raw_node_with_precision( - Op::MatMul, - vec![h, concat_w], - TensorType::f32(vec![m, 2 * out_features]), - requires_full_precision, - ); - graph.nodes_mut()[id].op = Op::GeGLUConcat; - graph.nodes_mut()[id].inputs = vec![wide_mm]; - fusions.push(( - "GeGLU(MatMul,MatMul)→GeGLUConcat(MatMul)".to_string(), - id as u32, - )); - } + crate::graph::ParamTransform::HorizontalConcat + }, + }); + let concat_w = graph.add_raw_node_with_precision( + Op::Parameter { name: concat_name }, + vec![], + TensorType::new(shape, dtype), + requires_full_precision, + ); + let m = graph.node(h).ty.shape[0]; + Some(graph.add_raw_node_with_precision( + if transposed { Op::MatMulBT } else { Op::MatMul }, + vec![h, concat_w], + TensorType::f32(vec![m, 2 * out_features]), + requires_full_precision, + )) } /// Dump the whole-graph egglog program (for standalone debugging). @@ -831,7 +512,7 @@ pub fn dump_egglog_program(graph: &Graph) -> String { ids, shifts: vec![0], }; - segment_program(graph, &seg).0 + segment_program(graph, &seg, true).0 } // --------------------------------------------------------------------------- @@ -854,19 +535,11 @@ fn plan_segments(g: &Graph, mode: OptimizeMode, saturation_cutoff: usize) -> Vec .iter() .filter(|n| !matches!(n.op, Op::Nop)) .count(); - if mode == OptimizeMode::EgglogWhole { - return vec![Segment { - ids: g - .nodes() - .iter() - .filter(|node| !matches!(node.op, Op::Nop)) - .map(|node| node.id as usize) - .collect(), - shifts: vec![0], - }]; - } - - let saturation_cutoff = saturation_cutoff.max(1); + let saturation_cutoff = if mode == OptimizeMode::EgglogWhole { + n.max(1) + } else { + saturation_cutoff.max(1) + }; let mut segments = Vec::new(); let mut covered = vec![false; n]; if mode == OptimizeMode::EgglogOutlined && active > saturation_cutoff { @@ -906,7 +579,30 @@ fn plan_segments(g: &Graph, mode: OptimizeMode, saturation_cutoff: usize) -> Vec shifts: vec![0], }); } + // A derivative may read a forward result, including an output shared with + // another kernel (for example cross-entropy logits gradients). Keep such + // cut edges opaque: reconstructing a forward expression under the backward + // root's precision policy can duplicate its producer and lose that identity. segments + .into_iter() + .flat_map(|seg| { + let (full, relaxed) = seg + .ids + .into_iter() + .partition(|&id| g.nodes()[id].requires_full_precision); + [ + Segment { + ids: relaxed, + shifts: seg.shifts.clone(), + }, + Segment { + ids: full, + shifts: seg.shifts, + }, + ] + }) + .filter(|seg| !seg.ids.is_empty()) + .collect() } // --------------------------------------------------------------------------- @@ -918,7 +614,7 @@ fn plan_segments(g: &Graph, mode: OptimizeMode, saturation_cutoff: usize) -> Vec /// ones added later — encodes through the arity-generic `Op1..Op6` /// constructors, tagged with the node id so ops with different /// attributes (eps, strides, head counts) never unify. -fn egglog_prelude(prog: &mut String) { +fn egglog_prelude(prog: &mut String, pack_swiglu: bool) { prog.push_str( "\ (datatype Op @@ -938,9 +634,11 @@ fn egglog_prelude(prog: &mut String) { (Silu Op) (SwiGLU Op Op) (SwiGLUPacked Op Op Op) + (SwiGLUPackedBT Op Op Op) (Gelu Op) (GeGLU Op Op) (GeGLUPacked Op Op Op) + (GeGLUPackedBT Op Op Op) (Op1 i64 Op) (Op2 i64 Op Op) (Op3 i64 Op Op Op) @@ -963,6 +661,9 @@ fn egglog_prelude(prog: &mut String) { (rewrite (Add ?d (MatMulAT ?a ?b)) (FusedMatMulATAdd ?a ?b ?d)) (rewrite (Add (MatMulBT ?a ?b) ?d) (FusedMatMulBTAdd ?a ?b ?d)) (rewrite (Add ?d (MatMulBT ?a ?b)) (FusedMatMulBTAdd ?a ?b ?d)) +(rewrite (FusedMatMulAdd ?a ?b ?d) (Add (MatMul ?a ?b) ?d)) +(rewrite (FusedMatMulATAdd ?a ?b ?d) (Add (MatMulAT ?a ?b) ?d)) +(rewrite (FusedMatMulBTAdd ?a ?b ?d) (Add (MatMulBT ?a ?b) ?d)) ; --- ONNX decomposed op recognition --- ; PyTorch decomposes compound ops when exporting to ONNX. These rules @@ -975,19 +676,21 @@ fn egglog_prelude(prog: &mut String) { ; SwiGLU: silu(gate) * up (rewrite (Mul (Silu ?gate) ?up) (SwiGLU ?gate ?up)) - -; Packed SwiGLU: gate and up projections sharing the input become one -; wide matmul over a concatenated weight (the derived parameter is -; created at stamp time; stamping falls back to the unpacked form when -; the weights are not plain 2D parameters). -(rewrite (SwiGLU (MatMul ?h ?wg) (MatMul ?h ?wu)) (SwiGLUPacked ?h ?wg ?wu)) +(rewrite (Mul ?up (Silu ?gate)) (SwiGLU ?gate ?up)) ; GeGLU: gelu(gate) * up, then the same HorizontalConcat packing as SwiGLU. (rewrite (Mul (Gelu ?gate) ?up) (GeGLU ?gate ?up)) +(rewrite (Mul ?up (Gelu ?gate)) (GeGLU ?gate ?up)) (rewrite (GeGLU (MatMul ?h ?wg) (MatMul ?h ?wu)) (GeGLUPacked ?h ?wg ?wu)) +(rewrite (GeGLU (MatMulBT ?h ?wg) (MatMulBT ?h ?wu)) (GeGLUPackedBT ?h ?wg ?wu)) ", ); + if pack_swiglu { + // Stamping creates the derived weight, or retains separate projections + // when the operands are not compatible 2D parameters. + prog.push_str("(rewrite (SwiGLU (MatMul ?h ?wg) (MatMul ?h ?wu)) (SwiGLUPacked ?h ?wg ?wu))\n(rewrite (SwiGLU (MatMulBT ?h ?wg) (MatMulBT ?h ?wu)) (SwiGLUPackedBT ?h ?wg ?wu))\n"); + } // Saturation is bounded: the deepest rewrite chain is three rules // (Mul(x, Sigmoid(x)) -> Silu, Mul(Silu, up) -> SwiGLU, then // SwiGLU(MatMul, MatMul) -> SwiGLUPacked), so three iterations reach @@ -1023,6 +726,22 @@ fn node_to_egglog_expr(node: &Node) -> String { Op::Input { .. } | Op::Parameter { .. } | Op::Constant { .. } => { format!("(Leaf {})", node.id) } + Op::RoPE { + theta, + pos_offset: 0, + .. + } + | Op::RoPEGrad { + theta, + pos_offset: 0, + .. + } if theta.is_finite() + && theta > 0.0 + && node.ty.shape.first() == Some(&1) + && node.inputs.len() == 1 => + { + format!("$n{}", node.inputs[0]) + } Op::Nop => unreachable!("Nop nodes are filtered before encoding"), ref op => { let args: Vec = node.inputs.iter().map(|i| format!("$n{}", i)).collect(); @@ -1045,7 +764,7 @@ fn node_to_egglog_expr(node: &Node) -> String { /// opaque `Leaf` terms, segment nodes are encoded in id order. Returns /// the program and the external node ids (needed to size their e-classes /// for traffic-aware extraction). -fn segment_program(g: &Graph, seg: &Segment) -> (String, Vec) { +fn segment_program(g: &Graph, seg: &Segment, pack_swiglu: bool) -> (String, Vec) { let idset: HashSet = seg.ids.iter().copied().collect(); let mut externals: Vec = Vec::new(); let mut seen = HashSet::new(); @@ -1064,7 +783,7 @@ fn segment_program(g: &Graph, seg: &Segment) -> (String, Vec) { externals.sort_unstable(); let mut prog = String::new(); - egglog_prelude(&mut prog); + egglog_prelude(&mut prog, pack_swiglu); for &e in &externals { prog.push_str(&format!("(let $n{} (Leaf {}))\n", e, e)); } @@ -1088,7 +807,7 @@ fn segment_program(g: &Graph, seg: &Segment) -> (String, Vec) { /// of their own. fn eclass_sizes( graph: &Graph, - egraph: &mut egglog::EGraph, + egraph: &egglog::EGraph, ids: impl Iterator, ) -> HashMap { let mut sizes = HashMap::new(); @@ -1098,9 +817,7 @@ fn eclass_sizes( continue; } let var = format!("$n{}", node.id); - if let Ok((_sort, value)) = - egraph.eval_expr(&egglog::ast::Expr::Var(egglog::ast::Span::Panic, var)) - { + if let Some(value) = egraph.lookup_function(&var, &[]) { sizes.insert(value, node.ty.size_bytes() as u64); } } @@ -1190,24 +907,17 @@ fn instance_ext_map( Some(map) } -#[allow(clippy::too_many_arguments)] fn process_segment( g: &mut Graph, seg: &Segment, index: &mut HashMap<(&'static str, Vec, bool), NodeId>, - fusions: &mut Vec<(String, u32)>, - first_program: &mut String, - num_eclasses: &mut usize, - num_enodes: &mut usize, - egglog_time: &mut std::time::Duration, - extract_time: &mut std::time::Duration, - extraction_failures: &mut usize, - extraction_cost: ExtractionCost, + report: &mut OptimizeReport, + config: OptimizeConfig, ) { let egglog_start = Instant::now(); - let (program, externals) = segment_program(g, seg); - if first_program.is_empty() { - first_program.clone_from(&program); + let (program, externals) = segment_program(g, seg, config.pack_swiglu); + if report.egglog_program.is_empty() { + report.egglog_program.clone_from(&program); } let mut egraph = egglog::EGraph::default(); if let Err(e) = egraph.parse_and_run_program(None, &program) { @@ -1216,43 +926,46 @@ fn process_segment( seg.ids.len(), e ); - *extraction_failures += 1; - *egglog_time += egglog_start.elapsed(); + report.extraction_failures += 1; + report.egglog_time += egglog_start.elapsed(); return; } - let cm = match extraction_cost { + let cm = match config.extraction_cost { ExtractionCost::AstSize => FusionCostModel::ast_size(), ExtractionCost::TensorTraffic => { let size_ids = externals.iter().copied().chain(seg.ids.iter().copied()); - FusionCostModel::with_sizes(eclass_sizes(g, &mut egraph, size_ids)) + FusionCostModel::with_sizes(eclass_sizes(g, &egraph, size_ids)) } }; let roots = segment_roots(g, seg); - let mut terms: Vec<(usize, TermDag, TermId)> = Vec::new(); + let sort = egraph.get_sort_by_name("Op").unwrap().clone(); + let extractor = Extractor::compute_costs_from_rootsorts(Some(vec![sort]), &egraph, cm); + let mut dag = TermDag::default(); + let mut terms = Vec::new(); for &root in &roots { let var = format!("$n{}", root); - match egraph.eval_expr(&egglog::ast::Expr::Var(egglog::ast::Span::Panic, var)) { - Ok((sort, value)) => { + match egraph.lookup_function(&var, &[]) { + Some(value) => { let extraction = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - egraph.extract_value_with_cost_model(&sort, value, cm.clone()) + extractor.extract_best(&egraph, &mut dag, value) })); match extraction { - Ok(Ok((dag, term_id, cost))) => { + Ok(Some((cost, term_id))) => { log::debug!( "extracted $n{} (cost {}): {}", root, cost, dag.to_string(term_id) ); - terms.push((root, dag, term_id)); + terms.push((root, term_id)); } - Ok(Err(e)) => { - *extraction_failures += 1; - log::warn!("extraction failed for $n{}: {}", root, e); + Ok(None) => { + report.extraction_failures += 1; + log::warn!("extraction failed for $n{}", root); } Err(_) => { - *extraction_failures += 1; + report.extraction_failures += 1; log::warn!( "egglog panicked while reconstructing $n{} — root left unchanged", root @@ -1260,16 +973,16 @@ fn process_segment( } } } - Err(e) => { - *extraction_failures += 1; - log::warn!("failed to eval $n{}: {}", root, e); + None => { + report.extraction_failures += 1; + log::warn!("missing e-class for $n{}", root); } } } let serialized = egraph.serialize(egglog::SerializeConfig::default()); - *num_eclasses += serialized.egraph.class_data.len(); - *num_enodes += serialized.egraph.nodes.len(); - *egglog_time += egglog_start.elapsed(); + report.num_eclasses += serialized.egraph.class_data.len(); + report.num_enodes += serialized.egraph.nodes.len(); + report.egglog_time += egglog_start.elapsed(); // Stamping. All instance translations are computed before any // mutation: stamping overwrites root inputs, which may be the very @@ -1290,7 +1003,7 @@ fn process_segment( ); continue; }; - for &(root, ref dag, term_id) in &terms { + for &(root, term_id) in &terms { let requires_full_precision = g.node((root + shift) as NodeId).requires_full_precision; let mut stamper = Stamper { g, @@ -1298,16 +1011,16 @@ fn process_segment( seg_ids: &idset, shift, ext_map, - fusions, + fusions: &mut report.fusions_applied, memo: HashMap::new(), requires_full_precision, }; - if let Err(e) = stamper.stamp_root(root + shift, dag, term_id) { + if let Err(e) = stamper.stamp_root(root + shift, &dag, term_id) { log::warn!("stamping $n{} (+{}) failed: {}", root, shift, e); } } } - *extract_time += stamp_start.elapsed(); + report.extract_time += stamp_start.elapsed(); } /// Rebuilds extracted terms in the graph IR. Interior nodes whose @@ -1339,7 +1052,7 @@ struct Stamper<'a> { impl Stamper<'_> { /// Overwrite the root node in place with the extracted term. fn stamp_root(&mut self, root: usize, dag: &TermDag, term_id: TermId) -> Result<(), String> { - match dag.get(term_id).clone() { + match *dag.get(term_id) { Term::App(ref name, ref children) if named_constructor_exists(name) => { let inputs = self.resolve_children(dag, children)?; // Unchanged term → nothing to do. @@ -1380,7 +1093,7 @@ impl Stamper<'_> { if let Some(&id) = self.memo.get(&term_id) { return Ok(id); } - let id = match dag.get(term_id).clone() { + let id = match *dag.get(term_id) { Term::App(ref name, ref children) if name == "Leaf" => { self.translate(lit_node_id(dag, children[0])?)? } @@ -1406,7 +1119,7 @@ impl Stamper<'_> { None => self.build_named(name, inputs, None)?, } } - other => return Err(format!("unexpected term {:?}", other)), + ref other => return Err(format!("unexpected term {:?}", other)), }; self.memo.insert(term_id, id); Ok(id) @@ -1431,11 +1144,24 @@ impl Stamper<'_> { inputs: Vec, target: Option, ) -> Result { - if name == "SwiGLUPacked" { - return self.build_glu_packed(&inputs, target, Op::SwiGLUConcat, "SwiGLUPacked"); - } - if name == "GeGLUPacked" { - return self.build_glu_packed(&inputs, target, Op::GeGLUConcat, "GeGLUPacked"); + match name { + "SwiGLUPacked" | "SwiGLUPackedBT" => { + return self.build_glu_packed( + &inputs, + target, + Op::SwiGLUConcat, + static_constructor(name)?, + ); + } + "GeGLUPacked" | "GeGLUPackedBT" => { + return self.build_glu_packed( + &inputs, + target, + Op::GeGLUConcat, + static_constructor(name)?, + ); + } + _ => {} } let shape = |id: NodeId| self.g.node(id).ty.shape.clone(); let ty_of = |id: NodeId| self.g.node(id).ty.clone(); @@ -1519,55 +1245,22 @@ impl Stamper<'_> { packed_key: &'static str, ) -> Result { let (h, wg, wu) = (inputs[0], inputs[1], inputs[2]); - let packable = { - let (g_node, u_node) = (self.g.node(wg), self.g.node(wu)); - matches!(g_node.op, Op::Parameter { .. }) - && matches!(u_node.op, Op::Parameter { .. }) - && g_node.ty.shape.len() == 2 - && g_node.ty.shape == u_node.ty.shape - && g_node.ty.dtype == u_node.ty.dtype - && self.g.node(h).ty.shape.len() == 2 - }; + let transposed = packed_key.ends_with("BT"); + let matmul = if transposed { "MatMulBT" } else { "MatMul" }; let unpacked = match concat_op { Op::SwiGLUConcat => "SwiGLU", Op::GeGLUConcat => "GeGLU", _ => unreachable!("glu pack only for SwiGLU/GeGLU concat"), }; - if !packable { - let gate = self.lookup_or_build("MatMul", vec![h, wg])?; - let up = self.lookup_or_build("MatMul", vec![h, wu])?; + let Some(wide_mm) = + pack_glu_matmul(self.g, h, wg, wu, transposed, self.requires_full_precision) + else { + let gate = self.lookup_or_build(matmul, vec![h, wg])?; + let up = self.lookup_or_build(matmul, vec![h, wu])?; return self.build_named(unpacked, vec![gate, up], target); - } - let param_name = |id: NodeId| match self.g.node(id).op { - Op::Parameter { ref name } => name.clone(), - _ => unreachable!(), }; - let (gate_name, up_name) = (param_name(wg), param_name(wu)); - let in_features = self.g.node(wg).ty.shape[0]; - let out_features = self.g.node(wg).ty.shape[1]; - let m = self.g.node(h).ty.shape[0]; - let concat_name = format!("{}+{}", gate_name, up_name); - // Record the derivation so the runtime fills the packed buffer - // from the original parameters. - self.g.derived_params.push(crate::graph::DerivedParam { - name: concat_name.clone(), - sources: vec![(gate_name, out_features), (up_name, out_features)], - rows: in_features, - transform: crate::graph::ParamTransform::HorizontalConcat, - }); - let concat_dtype = self.g.node(wg).ty.dtype; - let concat_w = self.g.add_raw_node_with_precision( - Op::Parameter { name: concat_name }, - vec![], - TensorType::new(vec![in_features, 2 * out_features], concat_dtype), - self.requires_full_precision, - ); - let wide_mm = self.g.add_raw_node_with_precision( - Op::MatMul, - vec![h, concat_w], - TensorType::f32(vec![m, 2 * out_features]), - self.requires_full_precision, - ); + let shape = &self.g.node(wide_mm).ty.shape; + let (m, out_features) = (shape[0], shape[1] / 2); let id = self.place( concat_op, vec![wide_mm], @@ -1579,7 +1272,7 @@ impl Stamper<'_> { id, ); self.fusions.push(( - format!("{unpacked}(MatMul,MatMul)→{unpacked}Concat(MatMul)"), + format!("{unpacked}({matmul},{matmul})→{unpacked}Concat({matmul})"), id, )); Ok(id) @@ -1644,9 +1337,11 @@ fn static_constructor(name: &str) -> Result<&'static str, String> { "Silu" => "Silu", "SwiGLU" => "SwiGLU", "SwiGLUPacked" => "SwiGLUPacked", + "SwiGLUPackedBT" => "SwiGLUPackedBT", "Gelu" => "Gelu", "GeGLU" => "GeGLU", "GeGLUPacked" => "GeGLUPacked", + "GeGLUPackedBT" => "GeGLUPackedBT", other => return Err(format!("unknown constructor {}", other)), }) } @@ -1883,8 +1578,8 @@ mod tests { use super::*; #[test] - fn greedy_is_the_production_default() { - assert_eq!(OptimizeConfig::default().mode, OptimizeMode::Greedy); + fn outlined_egglog_is_the_production_default() { + assert_eq!(OptimizeConfig::default().mode, OptimizeMode::EgglogOutlined); } #[test] @@ -2055,11 +1750,13 @@ mod tests { let opt = optimize(&g); assert_eq!( matches!(opt.node(opt.outputs()[0]).op, Op::Identity), - rows == 1 && offset == 0 && !dynamic + rows == 1 && offset == 0 && !dynamic, + "rows={rows}, offset={offset}, dynamic={dynamic}: {opt}" ); assert_eq!( matches!(opt.node(opt.outputs()[1]).op, Op::Identity), - rows == 1 && offset == 0 + rows == 1 && offset == 0, + "rows={rows}, offset={offset}, dynamic={dynamic}: {opt}" ); } } @@ -2232,7 +1929,6 @@ mod tests { assert_eq!(off_report.mode, OptimizeMode::Off); for mode in [ - OptimizeMode::Greedy, OptimizeMode::EgglogWindowed, OptimizeMode::EgglogOutlined, OptimizeMode::EgglogWhole, @@ -2260,73 +1956,88 @@ mod tests { } } - /// SwiGLU(MatMul, MatMul) → SwiGLUConcat(MatMul) fusion. #[test] - fn test_swiglu_concat_fusion() { - let mut g = Graph::new(); - let h = g.input("h", &[50, 720]); - let w_gate = g.parameter("w_gate", &[720, 2048]); - let w_up = g.parameter("w_up", &[720, 2048]); - let gate = g.matmul(h, w_gate); - let up = g.matmul(h, w_up); - let out = g.swiglu(gate, up); - g.set_outputs(vec![out]); - - let (opt, report) = optimize_with_report(&g); - let output_node = opt.node(opt.outputs()[0]); - assert!( - matches!(output_node.op, Op::SwiGLUConcat), - "expected SwiGLUConcat, got {:?}", - output_node.op - ); - assert!( - report - .fusions_applied - .iter() - .any(|entry| entry.0.contains("SwiGLU")), - "no SwiGLU fusion in report: {:?}", - report.fusions_applied - ); - // The fused matmul should have shape [50, 4096] (2*2048) - let mm_id = output_node.inputs[0]; - let mm_node = opt.node(mm_id); - assert!(matches!(mm_node.op, Op::MatMul)); - assert_eq!(mm_node.ty.shape, vec![50, 4096]); - assert_eq!(opt.derived_params.len(), 1); - } - - /// GeGLU(MatMul, MatMul) → GeGLUConcat(MatMul) fusion. - #[test] - fn test_geglu_concat_fusion() { - let mut g = Graph::new(); - let h = g.input("h", &[50, 720]); - let w_gate = g.parameter("w_gate", &[720, 2048]); - let w_up = g.parameter("w_up", &[720, 2048]); - let gate = g.matmul(h, w_gate); - let up = g.matmul(h, w_up); - let out = g.geglu(gate, up); - g.set_outputs(vec![out]); - - let (opt, report) = optimize_with_report(&g); - let output_node = opt.node(opt.outputs()[0]); - assert!( - matches!(output_node.op, Op::GeGLUConcat), - "expected GeGLUConcat, got {:?}", - output_node.op - ); - assert!( - report - .fusions_applied - .iter() - .any(|entry| entry.0.contains("GeGLU")), - "no GeGLU fusion in report: {:?}", - report.fusions_applied - ); - let mm_id = output_node.inputs[0]; - let mm_node = opt.node(mm_id); - assert!(matches!(mm_node.op, Op::MatMul)); - assert_eq!(mm_node.ty.shape, vec![50, 4096]); - assert_eq!(opt.derived_params.len(), 1); + fn glu_concat_fusion_preserves_orientation_and_format() { + for transposed in [false, true] { + for gelu in [false, true] { + for dtype in [crate::graph::DType::F32, crate::graph::DType::F16] { + let mut g = Graph::new(); + let h = g.input("h", &[3, 12]); + let shape = if transposed { vec![7, 12] } else { vec![12, 7] }; + let mut project = |name: &str| { + let w = g.add_raw_node( + Op::Parameter { name: name.into() }, + vec![], + TensorType::new(shape.clone(), dtype), + ); + if transposed { + g.matmul_bt(h, w) + } else { + g.matmul(h, w) + } + }; + let gate = project("gate"); + let up = project("up"); + let out = if gelu { + g.geglu(gate, up) + } else { + g.swiglu(gate, up) + }; + g.set_outputs(vec![out]); + if !gelu { + let (unpacked, _) = optimize_with_config( + &g, + OptimizeConfig { + pack_swiglu: false, + ..Default::default() + }, + ); + assert!(unpacked.derived_params.is_empty()); + assert!(matches!( + unpacked.node(unpacked.outputs()[0]).op, + Op::SwiGLU + )); + } + for mode in [OptimizeMode::EgglogWhole, OptimizeMode::EgglogOutlined] { + for extraction_cost in + [ExtractionCost::AstSize, ExtractionCost::TensorTraffic] + { + let (opt, _) = optimize_with_config( + &g, + OptimizeConfig { + mode, + extraction_cost, + ..Default::default() + }, + ); + let output = opt.node(opt.outputs()[0]); + assert!(matches!( + (&output.op, gelu), + (Op::SwiGLUConcat, false) | (Op::GeGLUConcat, true) + )); + let mm = opt.node(output.inputs[0]); + assert!(matches!( + (&mm.op, transposed), + (Op::MatMul, false) | (Op::MatMulBT, true) + )); + assert_eq!(mm.ty.shape, [3, 14]); + let weight = opt.node(mm.inputs[1]); + assert_eq!(weight.ty.dtype, dtype); + assert_eq!( + weight.ty.shape, + if transposed { [14, 12] } else { [12, 14] } + ); + assert_eq!(opt.derived_params.len(), 1); + assert!(matches!( + (&opt.derived_params[0].transform, transposed), + (crate::graph::ParamTransform::HorizontalConcat, false) + | (crate::graph::ParamTransform::VerticalConcat, true) + )); + } + } + } + } + } } /// Backward ops are encoded into egglog (not skipped). diff --git a/src/optimize/search.rs b/src/optimize/search.rs new file mode 100644 index 00000000..5ea8d64a --- /dev/null +++ b/src/optimize/search.rs @@ -0,0 +1,584 @@ +//! Bounded equivalent-region extraction for measured implementation selection. +//! +//! This uses the existing egglog rules and graph reconstruction. It deliberately +//! accepts a bounded region and extracts its observable roots together. Timing +//! roots independently would double-count shared work. Each candidate is tuned before +//! comparing its complete execution, not ranked by its untuned kernel timings. + +use super::{FusionCostModel, Segment, Stamper}; + +use crate::{Graph, graph::Op}; +use egglog::{ + Term, TermDag, TermId, Value, + extract::{CostModel, Extractor}, +}; +use std::{ + collections::{BTreeMap, HashMap, HashSet, VecDeque}, + ops::Range, + sync::Arc, +}; + +pub(crate) struct Candidate { + pub graph: Graph, + pub expression: String, +} + +pub(crate) struct SearchSpace { + pub candidates: Vec, + pub truncated: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +struct Edge { + head: String, + inputs: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct Cost { + forbidden: usize, + estimate: u64, +} + +impl egglog::extract::Cost for Cost { + fn identity() -> Self { + Self { + forbidden: 0, + estimate: 0, + } + } + fn unit() -> Self { + Self { + forbidden: 0, + estimate: 1, + } + } + fn combine(self, other: &Self) -> Self { + Self { + forbidden: self.forbidden.saturating_add(other.forbidden), + estimate: self.estimate.saturating_add(other.estimate), + } + } +} + +#[derive(Clone)] +struct Excluding { + costs: FusionCostModel, + forbidden: Arc<[Edge]>, +} + +impl CostModel for Excluding { + fn fold(&self, _: &str, children: &[Cost], head: Cost) -> Cost { + use egglog::extract::Cost as _; + children + .iter() + .fold(head, |cost, child| cost.combine(child)) + } + + fn enode_cost( + &self, + egraph: &egglog::EGraph, + func: &egglog::Function, + row: &egglog::FunctionRow, + ) -> Cost { + Cost { + forbidden: usize::from( + self.forbidden + .binary_search_by(|edge| { + edge.head.as_str().cmp(func.name()).then_with(|| { + edge.inputs.as_slice().cmp(&row.vals[..row.vals.len() - 1]) + }) + }) + .is_ok(), + ), + estimate: self.costs.enode_cost(egraph, func, row).max(1), + } + } +} + +fn edges( + egraph: &egglog::EGraph, + terms: &TermDag, + root: TermId, +) -> Result, String> { + let mut result = Vec::new(); + let mut values = vec![None; terms.size()]; + let mut stack = vec![(root, false)]; + while let Some((id, expanded)) = stack.pop() { + if values[id].is_some() { + continue; + } + let value = match *terms.get(id) { + Term::App(ref head, ref args) => { + if !expanded { + stack.push((id, true)); + stack.extend(args.iter().rev().map(|&arg| (arg, false))); + continue; + } + let inputs: Vec<_> = args.iter().map(|&arg| values[arg].unwrap()).collect(); + let value = egraph + .lookup_function(head, &inputs) + .ok_or("extracted term is missing from the e-graph")?; + if head != "Leaf" { + result.push(( + value, + Edge { + head: head.clone(), + inputs, + }, + )); + } + value + } + Term::Lit(egglog::ast::Literal::Int(value)) => egraph.base_to_value(value), + _ => return Err("expected an operator or node-id literal".into()), + }; + values[id] = Some(value); + } + Ok(result) +} + +/// Retain equivalent implementations before lowering or kernel tuning. +/// +/// Egglog's extractor reconstructs each candidate. Excluding a selected e-node +/// exposes alternative choices, including inside children. This is bounded +/// enumeration, not a globally optimal or k-best schedule search. The estimate +/// orders exploration only; callers must measure complete lowered candidates. +/// `truncated` reports an unfinished search. No GPU measurement happens here. +pub(crate) fn candidates( + graph: &Graph, + config: super::OptimizeConfig, + limit: usize, +) -> Result { + region_candidates(graph, 0..graph.nodes().len(), config, limit) +} + +/// Explore a contiguous region without changing the surrounding graph. Its +/// escaping values are extracted together; dependencies outside the region stay +/// opaque. Returned graphs preserve the original inputs, parameters and outputs. +pub(crate) fn region_candidates( + graph: &Graph, + region: Range, + config: super::OptimizeConfig, + limit: usize, +) -> Result { + if limit == 0 + || region.is_empty() + || region.len() > super::SATURATION_CUTOFF + || region.end > graph.nodes().len() + { + return Err("expected a bounded region and a positive candidate limit".into()); + } + let segment = Segment { + ids: region + .filter(|&id| !matches!(graph.nodes()[id].op, Op::Nop)) + .collect(), + shifts: vec![0], + }; + segment_candidates(graph, segment, config, limit) +} + +/// Apply each extracted alternative to all verified instances, then return the +/// complete model. Parameters and cut-edge placement remain model properties, +/// not newly introduced host-visible region inputs. +pub(crate) fn repeated_candidates( + graph: &Graph, + region: crate::outline::Region, + config: super::OptimizeConfig, + limit: usize, +) -> Result { + if limit == 0 + || !crate::outline::detect_repeated_regions(graph).contains(®ion) + || region.period > super::SATURATION_CUTOFF + { + return Err("expected a verified repeated region and a positive candidate limit".into()); + } + segment_candidates( + graph, + Segment { + ids: (region.start..region.start + region.period) + .filter(|&id| !matches!(graph.nodes()[id].op, Op::Nop)) + .collect(), + shifts: (0..region.count).map(|i| i * region.period).collect(), + }, + config, + limit, + ) +} + +fn segment_candidates( + graph: &Graph, + segment: Segment, + config: super::OptimizeConfig, + limit: usize, +) -> Result { + if graph + .nodes() + .iter() + .any(|node| node.inputs.iter().any(|&id| id >= node.id)) + { + return Err("region search requires topologically ordered nodes".into()); + } + let roots = super::segment_roots(graph, &segment); + if roots.is_empty() { + return Err("region has no observable output".into()); + } + let full_precision = graph.node(roots[0] as u32).requires_full_precision; + for id in segment + .shifts + .iter() + .flat_map(|shift| segment.ids.iter().map(move |id| id + shift)) + { + let node = &graph.nodes()[id]; + match node.op { + Op::Input { .. } | Op::Parameter { .. } | Op::Constant { .. } => continue, + Op::CacheWrite | Op::CacheWritePrefix | Op::ScatterAdd { .. } => { + return Err("stateful regions need an explicit mutation contract".into()); + } + _ => {} + } + if node.requires_full_precision != full_precision { + return Err("region crosses a precision boundary".into()); + } + if !(1..=6).contains(&node.inputs.len()) { + return Err("region contains an unsupported operator arity".into()); + } + } + let (mut program, externals) = super::segment_program(graph, &segment, config.pack_swiglu); + let root_name = if roots.len() == 1 { + format!("$n{}", roots[0]) + } else { + program.push_str(&format!( + "(constructor SearchOutputs ({}) Op)\n(let $outputs (SearchOutputs {}))\n", + vec!["Op"; roots.len()].join(" "), + roots + .iter() + .map(|id| format!("$n{id}")) + .collect::>() + .join(" "), + )); + "$outputs".into() + }; + let mut egraph = egglog::EGraph::default(); + egraph + .parse_and_run_program(None, &program) + .map_err(|e| e.to_string())?; + let sort = egraph.get_sort_by_name("Op").unwrap().clone(); + let value = egraph + .lookup_function(&root_name, &[]) + .ok_or("missing extraction root")?; + let costs = match config.extraction_cost { + super::ExtractionCost::AstSize => FusionCostModel::ast_size(), + super::ExtractionCost::TensorTraffic => FusionCostModel::with_sizes(super::eclass_sizes( + graph, + &egraph, + segment.ids.iter().chain(&externals).copied(), + )), + }; + let ids: HashSet<_> = segment.ids.iter().copied().collect(); + let uses = super::external_uses(graph, &segment); + let ext_maps = segment + .shifts + .iter() + .map(|&shift| super::instance_ext_map(graph, &segment, &uses, shift)) + .collect::>>() + .ok_or("region instances have ambiguous external edges")?; + let empty = Arc::<[Edge]>::from([]); + let mut pending = VecDeque::from([empty.clone()]); + let mut visited = HashSet::from([empty]); + let mut expressions = HashSet::new(); + let mut terms = TermDag::default(); + let mut choices = HashMap::new(); + let mut result = Vec::new(); + let attempts = limit.saturating_mul(segment.ids.len()); + let mut bounded = false; + for _ in 0..attempts { + let Some(forbidden) = pending.pop_front() else { + break; + }; + let extractor = Extractor::compute_costs_from_rootsorts( + Some(vec![sort.clone()]), + &egraph, + Excluding { + costs: costs.clone(), + forbidden: forbidden.clone(), + }, + ); + let Some((cost, term)) = extractor.extract_best(&egraph, &mut terms, value) else { + continue; + }; + if cost.forbidden != 0 { + continue; + } + let mut branches = Vec::new(); + let mut families = BTreeMap::>::new(); + for (value, edge) in edges(&egraph, &terms, term)? { + let branching = *choices.entry(value).or_insert_with(|| { + extractor + .extract_variants(&egraph, &mut TermDag::default(), value, 2) + .len() + > 1 + }); + if !branching { + continue; + } + families + .entry(edge.head.clone()) + .or_default() + .push(edge.clone()); + branches.push(edge); + } + // Visit whole-constructor alternatives before their individual sites. + // Otherwise a small bound explores many nearly identical partial + // unfusions and can miss the fully unfused family entirely. + let exclusions = families + .into_values() + .filter(|edges| edges.len() > 1) + .chain(branches.into_iter().map(|edge| vec![edge])); + for excluded in exclusions { + let mut next = forbidden.to_vec(); + next.extend(excluded); + next.sort_unstable(); + next.dedup(); + if visited.contains(next.as_slice()) { + continue; + } + if visited.len() == attempts { + bounded = true; + } else { + let next = Arc::<[Edge]>::from(next); + visited.insert(next.clone()); + pending.push_back(next); + } + } + if !expressions.insert(term) { + continue; + } + let mut candidate = graph.deep_clone(); + let mut index = super::build_structural_index(&candidate); + let terms_to_stamp = if roots.len() == 1 { + vec![term] + } else { + match *terms.get(term) { + Term::App(ref head, ref args) if head == "SearchOutputs" => args.clone(), + _ => return Err("missing joint extraction roots".into()), + } + }; + for (&shift, ext_map) in segment.shifts.iter().zip(&ext_maps) { + for (&root, &term) in roots.iter().zip(&terms_to_stamp) { + Stamper { + g: &mut candidate, + index: &mut index, + seg_ids: &ids, + shift, + ext_map, + fusions: &mut Vec::new(), + memo: HashMap::new(), + requires_full_precision: full_precision, + } + .stamp_root(root + shift, &terms, term)?; + } + } + super::sweep_dead_nodes(&mut candidate); + result.push(Candidate { + graph: candidate.into_toposort(), + expression: terms.to_string(term), + }); + if result.len() == limit { + break; + } + } + Ok(SearchSpace { + candidates: result, + truncated: bounded || !pending.is_empty(), + }) +} + +#[cfg(test)] +mod tests { + use super::candidates; + use crate::{Graph, graph::Op}; + + #[test] + #[ignore = "CPU search timing, run separately from correctness tests"] + #[cfg(feature = "models")] + fn cpu_search_overhead() { + use crate::models::{smollm2, smolvla, whisper}; + for model in ["SmolLM2-135M", "SmolVLA", "Whisper-tiny"] { + let mut graph = Graph::new(); + let output = match model { + "SmolLM2-135M" => { + smollm2::build_graph(&mut graph, &smollm2::Config::smollm2_135m(), 128) + } + "SmolVLA" => smolvla::build_action_expert( + &mut graph, + &smolvla::Config::smolvla_base(), + 50, + 16, + ), + _ => whisper::build_encoder(&mut graph, &whisper::Config::whisper_tiny(), 1, 3000), + }; + graph.set_outputs(vec![output]); + let region = crate::outline::detect_repeated_regions(&graph)[0]; + for sample in 0..6 { + let start = std::time::Instant::now(); + let space = + super::repeated_candidates(&graph, region, Default::default(), 4).unwrap(); + println!( + "{model} sample={sample} region_nodes={} candidates={} truncated={} ms={:.3}", + region.period, + space.candidates.len(), + space.truncated, + start.elapsed().as_secs_f64() * 1000.0 + ); + assert!(!space.candidates.is_empty()); + } + } + } + + #[test] + fn keeps_fused_and_unfused_implementations_before_kernel_tuning() { + let mut graph = Graph::new(); + let a = graph.input("a", &[3, 7]); + let b = graph.parameter("b", &[7, 5]); + let c = graph.input("c", &[3, 5]); + let product = graph.matmul(a, b); + let out = graph.add(product, c); + graph.set_outputs(vec![out]); + let space = candidates(&graph, Default::default(), 8).unwrap(); + assert!(!space.truncated); + let choices = space.candidates; + assert!( + choices + .iter() + .any(|c| matches!(c.graph.node(c.graph.outputs()[0]).op, Op::Add)) + ); + assert!( + choices + .iter() + .any(|c| matches!(c.graph.node(c.graph.outputs()[0]).op, Op::FusedMatMulAdd)) + ); + assert!( + choices + .iter() + .all(|c| c.graph.node(c.graph.outputs()[0]).ty.shape == [3, 5]) + ); + assert!(candidates(&graph, Default::default(), 0).is_err()); + let optimized = crate::optimize::optimize(&graph); + let recovered = candidates(&optimized, Default::default(), 8).unwrap(); + assert!(recovered.candidates.iter().any(|candidate| { + matches!( + candidate.graph.node(candidate.graph.outputs()[0]).op, + Op::Add + ) + })); + graph.set_outputs(vec![product, out]); + let space = candidates(&graph, Default::default(), 8).unwrap(); + assert!(!space.truncated); + assert!( + space + .candidates + .iter() + .all(|c| c.graph.outputs().len() == 2) + ); + graph.nodes_mut()[out as usize].requires_full_precision = true; + assert!(candidates(&graph, Default::default(), 8).is_err()); + } + + #[test] + fn explores_inner_choices_without_committing_to_the_cheapest_child() { + let mut graph = Graph::new(); + let a = graph.input("a", &[3, 7]); + let b = graph.parameter("b", &[7, 5]); + let c = graph.input("c", &[3, 5]); + let mm = graph.matmul(a, b); + let add = graph.add(mm, c); + let out = graph.neg(add); + graph.set_outputs(vec![out]); + let space = candidates(&graph, Default::default(), 8).unwrap(); + assert!(!space.truncated); + assert!( + space + .candidates + .iter() + .any(|c| c.expression.starts_with("(Neg (FusedMatMulAdd")) + ); + assert!( + space + .candidates + .iter() + .any(|c| c.expression.starts_with("(Neg (Add (MatMul")) + ); + assert!(candidates(&graph, Default::default(), 1).unwrap().truncated); + + let d = graph.input("d", &[3, 7]); + let e = graph.parameter("e", &[7, 5]); + let f = graph.input("f", &[3, 5]); + let mm2 = graph.matmul(d, e); + let add2 = graph.add(mm2, f); + let out = graph.mul(add, add2); + graph.set_outputs(vec![out]); + let space = candidates(&graph, Default::default(), 16).unwrap(); + assert!(!space.truncated); + let forms: std::collections::HashSet<_> = space + .candidates + .iter() + .map(|candidate| candidate.expression.matches("FusedMatMulAdd").count()) + .collect(); + assert_eq!(forms, [0, 1, 2].into_iter().collect()); + assert_eq!(space.candidates.len(), 4); + let bounded = candidates(&graph, Default::default(), 2).unwrap(); + assert!(bounded.truncated); + assert_eq!( + bounded + .candidates + .iter() + .map(|c| c.expression.matches("FusedMatMulAdd").count()) + .collect::>(), + [2, 0], + ); + + // Search only the second pair. External inputs keep their identities, + // and the first independent pair is not rewritten as a side effect. + let region = super::region_candidates( + &graph, + mm2 as usize..add2 as usize + 1, + Default::default(), + 8, + ) + .unwrap(); + assert_eq!(region.candidates.len(), 2); + assert!( + region + .candidates + .iter() + .all(|c| c.expression.matches("FusedMatMulAdd").count() <= 1) + ); + assert!(super::region_candidates(&graph, 0..0, Default::default(), 8).is_err()); + + let mut graph = Graph::new(); + let mut h = graph.input("x", &[3, 8]); + for layer in 0..10 { + let w = graph.parameter(&format!("w{layer}"), &[8, 8]); + let c = graph.parameter(&format!("c{layer}"), &[3, 8]); + let mm = graph.matmul(h, w); + h = graph.add(mm, c); + } + graph.set_outputs(vec![h]); + let region = crate::outline::detect_repeated_regions(&graph)[0]; + let space = super::repeated_candidates(&graph, region, Default::default(), 8).unwrap(); + assert!(!space.truncated); + assert_eq!(space.candidates.len(), 2); + assert!(space.candidates.iter().any(|candidate| { + candidate + .graph + .nodes() + .iter() + .filter(|node| matches!(node.op, Op::FusedMatMulAdd)) + .count() + == region.count + })); + graph.nodes_mut()[region.start + region.period + 2].requires_full_precision = true; + assert!(super::repeated_candidates(&graph, region, Default::default(), 8).is_err()); + } +} diff --git a/src/outline.rs b/src/outline.rs index 7ab25f60..1c74d476 100644 --- a/src/outline.rs +++ b/src/outline.rs @@ -13,11 +13,14 @@ //! candidate (start, period, count) lattices by sequence periodicity, //! then exact verification checks op/type equality (parameter and input //! names wildcarded) and edge isomorphism — every edge must either shift -//! with the instance (in-block and chain edges) or point at the same -//! shared global node for all instances. +//! with the instance (in-block and chain edges), point at the same shared +//! global node, or consistently rebind a dependency outside the region. use crate::graph::{Graph, Op}; -use std::hash::{Hash, Hasher}; +use std::{ + collections::HashMap, + hash::{Hash, Hasher}, +}; /// Longest block period considered. Also bounds the outlined egglog /// program size, keeping saturation fast (the full-graph cutoff is 300). @@ -74,6 +77,7 @@ fn node_signature(graph: &Graph, id: usize) -> u64 { } node.ty.shape.hash(&mut h); format!("{:?}", node.ty.dtype).hash(&mut h); + node.requires_full_precision.hash(&mut h); node.inputs.len().hash(&mut h); h.finish() } @@ -92,9 +96,9 @@ fn ops_equivalent(a: &Op, b: &Op) -> bool { /// Check that instance `m` and instance `m+1` of the lattice are exact /// structural copies: equivalent ops and types at each offset, and every -/// input edge either shifts by one period (in-block and chain edges — -/// instance m+1 reading from instance m is `ia + period`) or points at -/// the same shared global node (embeddings, masks, …). +/// internal or chain edge shifts by one period. Dependencies before the entire +/// region can be rebound consistently by type: topological sorting may have +/// hoisted independent parameter projections ahead of the compute blocks. /// /// Comparing *consecutive* pairs matters: the first instance's incoming /// chain edge points at whatever pre-region node produced the initial @@ -106,22 +110,32 @@ fn pair_isomorphic(graph: &Graph, start: usize, period: usize, m: usize) -> bool if start + (m + 2) * period > nodes.len() { return false; } + let left = start + m * period..start + (m + 1) * period; + let right = left.end..left.end + period; + let mut external = HashMap::new(); for off in 0..period { let a = &nodes[start + m * period + off]; let b = &nodes[start + (m + 1) * period + off]; - if !ops_equivalent(&a.op, &b.op) || a.ty != b.ty || a.inputs.len() != b.inputs.len() { + if !ops_equivalent(&a.op, &b.op) + || a.ty != b.ty + || a.inputs.len() != b.inputs.len() + || a.requires_full_precision != b.requires_full_precision + { return false; } for (&ia, &ib) in a.inputs.iter().zip(b.inputs.iter()) { let ia = ia as usize; let ib = ib as usize; - if ib == ia + period { - continue; // lattice edge (in-block or chain) + if left.contains(&ia) || right.contains(&ib) { + if !left.contains(&ia) || !right.contains(&ib) || ib != ia + period { + return false; + } + } else if nodes[ia].ty != nodes[ib].ty + || external.insert(ia, ib).is_some_and(|old| old != ib) + || (ib != ia + period && ib != ia && !(ia < start && ib < start)) + { + return false; } - if ib == ia { - continue; // shared global - } - return false; } } true @@ -203,6 +217,17 @@ pub fn detect_repeated_regions(graph: &Graph) -> Vec { if cand.period > MAX_PERIOD || cand.len() < MIN_COVERAGE { continue; } + if graph.nodes()[cand.start..cand.start + cand.period] + .iter() + .all(|node| { + matches!( + node.op, + Op::Input { .. } | Op::Parameter { .. } | Op::Constant { .. } | Op::Nop + ) + }) + { + continue; + } if accepted.iter().any(|r| r.overlaps(&cand)) { continue; } @@ -244,6 +269,42 @@ mod tests { assert_eq!(r.period, 5); assert_eq!(r.count, 12); assert_eq!(r.start, 1); // node 0 is the input + + let sorted = g.toposort(); + let regions = detect_repeated_regions(&sorted); + assert!(!regions.is_empty()); + assert!(regions.iter().all(|region| { + sorted.nodes()[region.start..region.start + region.period] + .iter() + .any(|node| !matches!(node.op, Op::Parameter { .. } | Op::Input { .. })) + })); + let mut split_precision = sorted; + let region = regions[0]; + let at = region.start + region.period; + split_precision.nodes_mut()[at].requires_full_precision = true; + assert!(!pair_isomorphic( + &split_precision, + region.start, + region.period, + 0 + )); + + let mut projected = Graph::new(); + let mut h = projected.input("x", &[4, 16]); + let weights: Vec<_> = (0..12) + .map(|i| { + let weight = projected.parameter(&format!("w{i}"), &[16, 16]); + projected.neg(weight) + }) + .collect(); + for weight in weights { + h = projected.matmul(h, weight); + h = projected.relu(h); + h = projected.neg(h); + } + projected.set_outputs(vec![h]); + let regions = detect_repeated_regions(&projected.toposort()); + assert!(regions.iter().any(|r| r.period == 3 && r.count >= 10)); } #[test] diff --git a/src/profiler.rs b/src/profiler.rs index e5e8aefc..8248125d 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -689,12 +689,12 @@ fn capture_windows( .product(), input_buffer_bytes, output_buffer_bytes, - cooperative: dispatch.use_coop, - small_tile: dispatch.use_small_tiles, + cooperative: dispatch.use_coop(), + small_tile: dispatch.use_small_tiles(), requires_full_precision: dispatch.requires_full_precision, weight_format: format!("{:?}", dispatch.weight_format), has_prologue: dispatch.matmul_prologue.is_some(), - has_epilogue: dispatch.matmul_epilogue.is_some() || !dispatch.epilogue.is_empty(), + has_epilogue: dispatch.matmul_epilogue.is_some(), timing_samples_ms: timing_samples[index].clone(), median_ms, p25_ms: quantile(&timing_samples[index], 0.25), diff --git a/src/runtime.rs b/src/runtime.rs index 28855f20..8e680e80 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,11 +1,13 @@ -use crate::compile::{BufferRef, Dispatch, ExecutionPlan, ShaderEntry}; +use crate::compile::{BufferRef, CachedBlockAttentionParams, Dispatch, ExecutionPlan, ShaderEntry}; use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use std::sync::Arc; mod checkpoint; +pub(crate) mod search_state; mod tuning; pub use crate::tune::TuneOutcome; +pub(crate) use tuning::KernelMemo; type Gpu = blade_graphics::Context; @@ -876,22 +878,6 @@ struct CachedAttentionData { params: MatMulParams, // queries, num_heads, num_kv_heads, head_dim } -#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod)] -#[repr(C)] -struct CachedBlockAttentionParams { - window_size: u32, - num_heads: u32, - num_kv_heads: u32, - head_dim: u32, - block_len: u32, - max_seq: u32, - // Split-K: the split count and per-split token chunk. Zero splits on - // the fused single-kernel form; the spare slots avoid a new params - // struct and layout. - splits: u32, - chunk: u32, -} - #[derive(blade_macros::ShaderData)] struct CachedBlockAttentionData { src_a: blade_graphics::BufferPiece, @@ -1055,33 +1041,19 @@ struct MultiHeadAttnGradKVData { /// f32 shader over packed blocks, or a 64×64 shader under a workgroup /// count computed for 32×32 tiles. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -enum EpiloguePipelineKey { - Dag( - crate::compile::MatMulEpilogue, - crate::compile::WeightFormat, - crate::codegen::MatMulTile, - ), - Legacy( - Vec, - crate::compile::WeightFormat, - crate::codegen::MatMulTile, - ), -} +struct EpiloguePipelineKey( + crate::compile::MatMulEpilogue, + crate::compile::WeightFormat, + crate::codegen::MatMulTile, +); fn epilogue_pipeline_key(dispatch: &Dispatch) -> Option { let format = dispatch.weight_format; let tile = epilogue_tile(dispatch); - if let Some(ref epilogue) = dispatch.matmul_epilogue { - Some(EpiloguePipelineKey::Dag(epilogue.clone(), format, tile)) - } else if !dispatch.epilogue.is_empty() { - Some(EpiloguePipelineKey::Legacy( - dispatch.epilogue.clone(), - format, - tile, - )) - } else { - None - } + dispatch + .matmul_epilogue + .as_ref() + .map(|epilogue| EpiloguePipelineKey(epilogue.clone(), format, tile)) } /// Tile geometry the epilogue shader must be generated for. @@ -1092,7 +1064,7 @@ fn epilogue_pipeline_key(dispatch: &Dispatch) -> Option { /// check keeps the result correct, but three quarters of the workgroups /// do nothing. fn epilogue_tile(dispatch: &Dispatch) -> crate::codegen::MatMulTile { - if dispatch.use_small_tiles { + if dispatch.use_small_tiles() { crate::codegen::MatMulTile::Small } else { crate::codegen::MatMulTile::Large @@ -1103,13 +1075,16 @@ fn epilogue_tile(dispatch: &Dispatch) -> crate::codegen::MatMulTile { /// /// A dispatch names a `ShaderEntry`, but the pipeline it actually runs /// also depends on the modifiers the plan attached to it - weight format, -/// tiling, a fused epilogue, and so on. Each modifier axis is one variant -/// here, and [`Pipelines::candidates`] lists them most-specific-first, so -/// adding an axis costs a variant, a `candidates` line, and a `label` -/// arm rather than a map, a struct field, and four parallel match chains. +/// tiling, a fused epilogue, and so on. [`Pipelines::key`] resolves that +/// implementation once; preparation builds exactly the selected pipeline. #[derive(Clone, Debug, PartialEq, Eq, Hash)] enum Variant { SpecializedConv(ShaderEntry, Vec, u32), + ScalarMatmul( + ShaderEntry, + crate::compile::WeightFormat, + crate::codegen::ScalarMatmulShape, + ), /// Schedule-template kernels, keyed by kernel content hash. These are /// generated from a DAG rather than a shader group, so no `ShaderEntry` /// identifies them. @@ -1120,8 +1095,7 @@ enum Variant { /// containing different attention widths run every dispatch through /// whichever width happened to be encountered last. Attention(ShaderEntry, u32), - /// Epilogue-fused matmuls, keyed by their actual DAG (or the legacy - /// closed op list for deserialized old plans). The cooperative form + /// Epilogue-fused matmuls, keyed by their actual DAG. The cooperative form /// uses workgroup memory to expose accumulator lanes to the epilogue. Epilogue(ShaderEntry, EpiloguePipelineKey), CoopEpilogue(ShaderEntry, EpiloguePipelineKey), @@ -1163,9 +1137,7 @@ enum Variant { /// Non-f32 weight storage (f16, Q4, Q8). Weight(ShaderEntry, crate::compile::WeightFormat), WeightSmall(ShaderEntry, crate::compile::WeightFormat), - /// Cooperative-matrix and small-tile (32×32) forms. Unlike every other - /// axis these are pure performance: the scalar pipeline computes the - /// same thing, so falling back to it is safe. + /// Cooperative-matrix implementation qualified for the session's precision policy. Coop(ShaderEntry), /// Cooperative f16 with hi/lo residual staging (C1). CoopCompensated(ShaderEntry), @@ -1174,7 +1146,7 @@ enum Variant { /// a backward pack of the same arity (compensated). Horizontal(ShaderEntry, u32, HorizMatMulKind), SmallTile(ShaderEntry), - /// The unmodified pipeline. Always compiled. + /// The unmodified pipeline. Scalar(ShaderEntry), } @@ -1186,9 +1158,9 @@ enum HorizMatMulKind { } fn horiz_kind(dispatch: &Dispatch) -> HorizMatMulKind { - if dispatch.use_coop_compensated { + if dispatch.use_coop_compensated() { HorizMatMulKind::CoopCompensated - } else if dispatch.use_coop { + } else if dispatch.use_coop() { HorizMatMulKind::Coop } else { HorizMatMulKind::Scalar @@ -1203,6 +1175,7 @@ impl Variant { Variant::Reduction(_) | Variant::Pointwise(_) => None, Variant::Attention(ref e, _) | Variant::SpecializedConv(ref e, _, _) + | Variant::ScalarMatmul(ref e, _, _) | Variant::Epilogue(ref e, _) | Variant::CoopEpilogue(ref e, _) | Variant::CoopPrologue(ref e, _) @@ -1223,6 +1196,9 @@ impl Variant { /// Name used by the profiler and by pipeline-statistics dumps. fn label(&self) -> String { match *self { + Variant::ScalarMatmul(ref e, format, shape) => { + format!("{e:?}:scalar-{format:?}-{shape:?}") + } Variant::SpecializedConv(ref e, ref params, k_tile) => { format!("{e:?}:fixed-native-div-k{k_tile}-{params:?}") } @@ -1234,22 +1210,16 @@ impl Variant { Variant::CoopPrologue(ref e, ref kinds) => { format!("{e:?}:cooperative-prologue:{kinds:?}") } - Variant::GemvRmsNorm(ref e, format, shape) => format!( - "{e:?}:rmsnorm-{format:?}-{}t-{:?}", - shape.threads, shape.reduction - ), - Variant::GemvIntDot(ref e, format, shape) => format!( - "{e:?}:gemv-intdot-{format:?}-{}t-{:?}", - shape.threads, shape.reduction - ), - Variant::GemvRmsNormIntDot(ref e, format, shape) => format!( - "{e:?}:gemv-rmsnorm-intdot-{format:?}-{}t-{:?}", - shape.threads, shape.reduction - ), - Variant::Gemv(ref e, format, shape) => format!( - "{e:?}:gemv-{format:?}-{}t-{:?}", - shape.threads, shape.reduction - ), + Variant::GemvRmsNorm(ref e, format, shape) => { + format!("{e:?}:rmsnorm-{format:?}-{shape:?}") + } + Variant::GemvIntDot(ref e, format, shape) => { + format!("{e:?}:gemv-intdot-{format:?}-{shape:?}") + } + Variant::GemvRmsNormIntDot(ref e, format, shape) => { + format!("{e:?}:gemv-rmsnorm-intdot-{format:?}-{shape:?}") + } + Variant::Gemv(ref e, format, shape) => format!("{e:?}:gemv-{format:?}-{shape:?}"), Variant::Weight(ref e, format) => format!("{e:?}:weight-{format:?}"), Variant::WeightSmall(ref e, format) => format!("{e:?}:weight-{format:?}-small-tile"), Variant::Coop(ref e) => format!("{e:?}:cooperative"), @@ -1263,8 +1233,10 @@ impl Variant { struct Pipelines { map: HashMap, - /// Matmul codegen knobs the plan was compiled with. - matmul_knobs: crate::codegen::MatmulKnobs, + /// Resolved after compilation or tuning, never while recording a step. + selected: Vec, + /// Codegen knobs the plan was compiled with. + knobs: crate::compile::TuningKnobs, /// Where to write every WGSL the pipeline layer compiles — [`SessionOptions::wgsl_dump_dir`]. dump_dir: Option, } @@ -1277,14 +1249,15 @@ fn create_gen_shader( gpu: &Gpu, module: crate::codegen::ShaderModule, dump_dir: Option<&str>, -) -> blade_graphics::Shader { +) -> Result { if let Some(dir) = dump_dir { module.dump(dir); } - gpu.create_shader(blade_graphics::ShaderDesc { + gpu.try_create_shader(blade_graphics::ShaderDesc { source: &module.source, naga_module: Some(module.module), }) + .map_err(|error| error.to_string()) } fn create_profiled_pipeline( @@ -1313,236 +1286,152 @@ impl Pipelines { coop_config: Option<&crate::codegen::CoopConfig>, wgsl_dump_dir: Option<&str>, ) -> Self { - // Generated attention kernels must agree with the plan's dispatch - // geometry, so their knobs come from the plan itself. - let knobs = plan.knobs; - - // Collect which shader groups are needed. - // For matmul entries, compile BOTH scalar and coop if any dispatch uses coop. - use crate::codegen::ShaderGroup; - let mut needed: HashSet = HashSet::new(); - let mut needed_coop: HashSet = HashSet::new(); - let mut needed_coop_compensated: HashSet = HashSet::new(); - let mut needed_weighted: HashMap> = - HashMap::new(); - let mut needed_small: HashSet = HashSet::new(); - let mut entries_for_group: HashMap> = HashMap::new(); - let mut attention_entries: HashSet<(ShaderEntry, u32)> = HashSet::new(); - + let mut pipelines = Self { + map: HashMap::new(), + selected: Vec::new(), + knobs: plan.knobs, + dump_dir: wgsl_dump_dir.map(str::to_string), + }; for dispatch in &plan.dispatches { - if dispatch.conv_k_tile.is_some() { - continue; - } - let group = dispatch.shader.shader_group(); - // `pipeline_variants` resolves a fused epilogue before it looks - // at the small-tile or weight-format modifiers, so such a - // dispatch only ever runs its epilogue pipeline. Keep it from - // requesting the variants it cannot select — a group that also - // holds a plain dispatch still gets them from that one. - let resolves_to_epilogue = - dispatch.horizontal_batch < 2 && epilogue_pipeline_key(dispatch).is_some(); - // Generated conv2d coop entries are coop-only; skip for non-coop map. - let is_gen_coop = matches!( - dispatch.shader, - ShaderEntry::Conv2dGradInputGemmCoopGen(..) | ShaderEntry::Conv2dGemmCoopGen(..) - ); - if !is_gen_coop && !dispatch.gemv_int_dot { - needed.insert(group); - } - entries_for_group - .entry(group) - .or_default() - .insert(dispatch.shader.clone()); - // Retain the recorded scalar pipeline for cooperative fallback - // provenance. The current scratch tuner only searches scalar tiles. - if let Some((ref fb_shader, _)) = dispatch.scalar_fallback { - let fb_group = fb_shader.shader_group(); - needed.insert(fb_group); - entries_for_group - .entry(fb_group) - .or_default() - .insert(fb_shader.clone()); - } - // Attention dispatches store head_dim at params[3]. - if matches!( - group, - ShaderGroup::MultiHeadAttn - | ShaderGroup::FlashAttention - | ShaderGroup::FlashAttentionCoop - | ShaderGroup::FlashGradQ - | ShaderGroup::FlashGradQCoop - | ShaderGroup::FlashGradKV - | ShaderGroup::FlashGradKVCoop - | ShaderGroup::CachedBlockAttention - | ShaderGroup::CachedBlockAttentionSplit - | ShaderGroup::CachedBlockAttentionCombine - ) && dispatch.params.len() >= 4 - { - attention_entries.insert((dispatch.shader.clone(), dispatch.params[3])); - } - if dispatch.use_small_tiles - && !dispatch.weight_format.uses_reduced_storage() - && !resolves_to_epilogue - { - needed_small.insert(group); - entries_for_group - .entry(group) - .or_default() - .insert(dispatch.shader.clone()); - } - // Cooperative execution is a modifier on the dispatch, so the - // group stays the scalar one and only the generated module - // differs. Conv2d is the exception: `select_variants` already - // rewrote those dispatches to their per-kernel generated - // entries, which carry their own group. - if dispatch.use_coop - && (crate::codegen::coop_shape(group).is_some() - || matches!( - group, - ShaderGroup::Conv2dGemmCoop | ShaderGroup::Conv2dGradInputGemmCoop - )) - { - needed_coop.insert(group); - if dispatch.use_coop_compensated { - needed_coop_compensated.insert(group); - } - } - if dispatch.weight_format.uses_reduced_storage() - && !resolves_to_epilogue - && !dispatch.gemv_int_dot - { - needed_weighted - .entry(dispatch.weight_format) - .or_default() - .insert(group); - entries_for_group - .entry(group) - .or_default() - .insert(dispatch.shader.clone()); + pipelines + .prepare(gpu, dispatch, coop_config) + .expect("selected shader was rejected"); + } + if !plan.param_grad_pairs.is_empty() { + for shader in [ + ShaderEntry::SgdUpdate, + ShaderEntry::AdamUpdate, + ShaderEntry::GradClipZero, + ShaderEntry::GradClipNormSq, + ShaderEntry::GradClipScale, + ShaderEntry::AdaptiveGradClip, + ShaderEntry::GradAccum, + ] { + pipelines + .prepare( + gpu, + &Dispatch { + shader, + ..Default::default() + }, + None, + ) + .expect("optimizer shader was rejected"); } } + pipelines.select(&plan.dispatches); + pipelines + } - // Always compile SGD and Adam if the plan has trainable parameters. - if !plan.param_grad_pairs.is_empty() { - needed.insert(ShaderGroup::Sgd); - entries_for_group - .entry(ShaderGroup::Sgd) - .or_default() - .insert(ShaderEntry::SgdUpdate); - needed.insert(ShaderGroup::Adam); - entries_for_group - .entry(ShaderGroup::Adam) - .or_default() - .insert(ShaderEntry::AdamUpdate); - // Grad-clip shaders (zero, norm-sq accum, scale). Compiled - // unconditionally so `set_grad_clip_norm` works at runtime - // without a session rebuild. - needed.insert(ShaderGroup::GradClipZero); - entries_for_group - .entry(ShaderGroup::GradClipZero) - .or_default() - .insert(ShaderEntry::GradClipZero); - needed.insert(ShaderGroup::GradClipNormSq); - entries_for_group - .entry(ShaderGroup::GradClipNormSq) - .or_default() - .insert(ShaderEntry::GradClipNormSq); - needed.insert(ShaderGroup::GradClipScale); - entries_for_group - .entry(ShaderGroup::GradClipScale) - .or_default() - .insert(ShaderEntry::GradClipScale); - needed.insert(ShaderGroup::AdaptiveGradClip); - entries_for_group - .entry(ShaderGroup::AdaptiveGradClip) - .or_default() - .insert(ShaderEntry::AdaptiveGradClip); - // Temporal grad accumulator pass, compiled unconditionally so - // `set_grad_accumulate` works without a session rebuild. - needed.insert(ShaderGroup::GradAccum); - entries_for_group - .entry(ShaderGroup::GradAccum) - .or_default() - .insert(ShaderEntry::GradAccum); - } - - let mut map: HashMap = HashMap::new(); - - // One place that turns a generated module into pipelines for the - // entries of a group. Every modifier axis - base, small tile, coop, - // weight format - differs only in which module it hands over and - // which `Variant` it files the result under. - let compile_variant = - |sm: crate::codegen::ShaderModule, - group: ShaderGroup, - key: &dyn Fn(ShaderEntry) -> Variant, - target: &mut HashMap| { - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - if let Some(entries) = entries_for_group.get(&group) { - for entry in entries { - let layout = shader_data_layout(entry); - let variant = key(entry.clone()); - let pipeline = create_profiled_pipeline( - gpu, - variant.label(), - &layout, - shader.at(entry.entry_point()), - ); - target.insert(variant, pipeline); + fn prepare( + &mut self, + gpu: &Gpu, + dispatch: &Dispatch, + coop_config: Option<&crate::codegen::CoopConfig>, + ) -> Result<(), String> { + use crate::codegen::ShaderGroup; + let key = Self::key(dispatch); + if self.map.contains_key(&key) { + return Ok(()); + } + let knobs = self.knobs; + let matmul_knobs = crate::codegen::MatmulKnobs { + k_stage: knobs.matmul_k_stage, + interleave_columns: knobs.matmul_interleave_columns, + integer_dot: gpu.capabilities().shader_integer_dot_product, + }; + let group = dispatch.shader.shader_group(); + let mut layout = shader_data_layout(&dispatch.shader); + let mut entry_point = dispatch.shader.entry_point(); + let cooperative = + || *coop_config.expect("cooperative dispatch needs a qualified device configuration"); + let module = match key { + Variant::Scalar(_) => crate::codegen::generate_module(group, matmul_knobs), + Variant::SmallTile(_) => crate::codegen::generate_module_small(group, matmul_knobs), + Variant::Weight(_, format) => { + crate::codegen::generate_module_weighted(group, format, matmul_knobs) + } + Variant::WeightSmall(..) => { + tuning::tile_module(dispatch, crate::tune::MatmulTile::Tile32, matmul_knobs) + } + Variant::ScalarMatmul(_, _, shape) => tuning::tile_module( + dispatch, + crate::tune::MatmulTile::Scalar(shape), + matmul_knobs, + ), + Variant::SpecializedConv(..) => { + let tile = crate::tune::MatmulTile::selected(dispatch, None) + .expect("specialized convolution"); + tuning::tile_module(dispatch, tile, matmul_knobs) + } + Variant::Gemv(_, _, shape) + | Variant::GemvIntDot(_, _, shape) + | Variant::GemvRmsNorm(_, _, shape) + | Variant::GemvRmsNormIntDot(_, _, shape) => { + if dispatch.gemv_rmsnorm.is_some() { + layout = ::layout(); + } + tuning::tile_module(dispatch, crate::tune::MatmulTile::Gemv(shape), matmul_knobs) + } + Variant::Coop(_) | Variant::CoopCompensated(_) => { + let mut config = cooperative(); + config.compensated = dispatch.use_coop_compensated(); + use crate::codegen::Conv2dCoopDirection; + match dispatch.shader { + ShaderEntry::Conv2dGemmCoopGen(kh, kw, stride) => { + crate::codegen::generate_conv2d_coop_module( + kh, + kw, + stride, + Conv2dCoopDirection::Forward, + &config, + ) } + ShaderEntry::Conv2dGradInputGemmCoopGen(kh, kw, stride) => { + crate::codegen::generate_conv2d_coop_module( + kh, + kw, + stride, + Conv2dCoopDirection::GradInput, + &config, + ) + } + _ => crate::codegen::generate_module_coop(group, &config), } - }; - let mut matmul_knobs = crate::codegen::MatmulKnobs { - k_stage: plan.knobs.matmul_k_stage, - interleave_columns: plan.knobs.matmul_interleave_columns, - ..Default::default() - }; - matmul_knobs.integer_dot = gpu.capabilities().shader_integer_dot_product; - let compile_group = - |group: ShaderGroup, - key: &dyn Fn(ShaderEntry) -> Variant, - target: &mut HashMap| { - compile_variant( - crate::codegen::generate_module(group, matmul_knobs), - group, - key, - target, - ); - }; - - for &group in &needed_small { - compile_variant( - crate::codegen::generate_module_small(group, matmul_knobs), + } + Variant::Epilogue(..) => crate::codegen::generate_matmul_with_epilogue( group, - &Variant::SmallTile, - &mut map, - ); - } - - for &group in &needed { - if matches!( + dispatch.matmul_epilogue.as_ref(), + crate::codegen::MatMulOptions { + format: dispatch.weight_format, + tile: epilogue_tile(dispatch), + knobs: matmul_knobs, + }, + ), + Variant::CoopEpilogue(..) => crate::codegen::generate_coop_matmul_with_dag_epilogue( group, - ShaderGroup::MultiHeadAttn - | ShaderGroup::FlashAttention - | ShaderGroup::FlashAttentionCoop - | ShaderGroup::FlashGradQ - | ShaderGroup::FlashGradQCoop - | ShaderGroup::FlashGradKV - | ShaderGroup::FlashGradKVCoop - | ShaderGroup::CachedBlockAttention - | ShaderGroup::CachedBlockAttentionSplit - | ShaderGroup::CachedBlockAttentionCombine - ) { - // Compiled below per (entry, head_dim), not once per group. - continue; - } else { - compile_group(group, &Variant::Scalar, &mut map); + &cooperative(), + dispatch + .matmul_epilogue + .as_ref() + .expect("cooperative epilogue"), + ), + Variant::CoopPrologue(..) => { + let prologue = dispatch + .matmul_prologue + .as_ref() + .expect("cooperative prologue"); + let (add, variant) = + crate::codegen::coop_shape(group).expect("cooperative matrix group"); + layout = matmul_with_prologue_layout(prologue.factors.len()); + crate::codegen::gen_matmul_coop_with_prologue( + add, + variant, + &cooperative(), + prologue, + ) } - } - - for (entry, hd) in attention_entries { - let group = entry.shader_group(); - let sm = match group { + Variant::Attention(_, hd) => match group { ShaderGroup::FlashAttention => { crate::codegen::generate_flash_attention_module(hd, knobs.flash_ept_cap) } @@ -1568,514 +1457,131 @@ impl Pipelines { crate::codegen::generate_cached_attention_module(group, Some(hd)) } _ => unreachable!("non-parameterized attention group {group:?}"), - }; - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = shader_data_layout(&entry); - let key = Variant::Attention(entry.clone(), hd); - let pipeline = - create_profiled_pipeline(gpu, key.label(), &layout, shader.at(entry.entry_point())); - map.insert(key, pipeline); - } - // Conv2d coop dispatches were rewritten by `select_variants` to - // generated per-(kernel, stride) entries, compiled individually - // below. Every other cooperative group compiles from its own - // module, keyed by the same entry as its scalar form. - let mut conv2d_gen_entries: Vec = Vec::new(); - for &group in &needed_coop { - let is_conv_gen = matches!( - group, - ShaderGroup::Conv2dGemmCoop | ShaderGroup::Conv2dGradInputGemmCoop - ); - match (coop_config, is_conv_gen) { - (Some(config), false) => compile_variant( - crate::codegen::generate_module_coop(group, config), - group, - &Variant::Coop, - &mut map, - ), - (Some(_), true) => conv2d_gen_entries - .extend(entries_for_group.get(&group).into_iter().flatten().cloned()), - // A deserialized plan can carry `use_coop` from a machine - // whose capabilities this one lacks. The scalar kernel - // computes the same thing. - (None, _) => compile_group(group, &Variant::Coop, &mut map), - } - } - if let Some(config) = coop_config { - for &group in &needed_coop_compensated { - if crate::codegen::coop_shape(group).is_none() { - continue; - } - let mut compensated = *config; - compensated.compensated = true; - compile_variant( - crate::codegen::generate_module_coop(group, &compensated), - group, - &Variant::CoopCompensated, - &mut map, - ); + }, + Variant::Pointwise(_) => { + let dag = dispatch.pointwise().expect("pointwise implementation"); + layout = pointwise_data_layout(dag.n_inputs); + entry_point = crate::schedule::POINTWISE_ENTRY; + crate::schedule::lower(&crate::schedule::KernelTemplate::Pointwise { + dag: dag.clone(), + grid: crate::schedule::GridShape::default(), + }) } - } - // Compile generated conv2d coop kernels individually. - if let Some(config) = coop_config { - use crate::codegen::Conv2dCoopDirection; - for entry in &conv2d_gen_entries { - let (kh, kw, stride, direction) = match *entry { - ShaderEntry::Conv2dGemmCoopGen(kh, kw, s) => { - (kh, kw, s, Conv2dCoopDirection::Forward) - } - ShaderEntry::Conv2dGradInputGemmCoopGen(kh, kw, s) => { - (kh, kw, s, Conv2dCoopDirection::GradInput) + Variant::Reduction(_) => { + let kernel = dispatch.reduction().expect("reduction implementation"); + layout = reduction_data_layout(kernel); + entry_point = crate::schedule::REDUCTION_ENTRY; + crate::schedule::lower(&kernel.to_template()) + } + Variant::Horizontal(_, count, kind) => { + let coop = match kind { + HorizMatMulKind::Scalar => None, + HorizMatMulKind::Coop | HorizMatMulKind::CoopCompensated => { + let mut config = cooperative(); + config.compensated = kind == HorizMatMulKind::CoopCompensated; + Some(config) } - _ => unreachable!(), - }; - let sm = - crate::codegen::generate_conv2d_coop_module(kh, kw, stride, direction, config); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = shader_data_layout(entry); - let key = Variant::Coop(entry.clone()); - let pipeline = create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(entry.entry_point()), - ); - map.insert(key, pipeline); - } - } - - // Compile the RmsNorm-fused GEMV when any dispatch asks for it. The - // module is derived from the plain GEMV, so it needs no ShaderGroup - // of its own; it is a variant of `MatMulGemv`, resolved by - // `get_pipeline` the same way a weight format is. Int-dot GEMVs - // have their own fused form, compiled below. - let mut fused_keys = std::collections::HashSet::new(); - for fused in plan - .dispatches - .iter() - .filter(|d| d.gemv_rmsnorm.is_some() && !d.gemv_int_dot) - { - let shape = fused - .gemv_shape - .unwrap_or_else(|| crate::codegen::GemvShape::initial(ShaderGroup::MatMulGemv)); - let key = Variant::GemvRmsNorm(fused.shader.clone(), fused.weight_format, shape); - if !fused_keys.insert(key.clone()) { - continue; - } - let sm = crate::codegen::generate_module_gemv_rmsnorm(shape, fused.weight_format); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = ::layout(); - let pipeline = create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(fused.shader.entry_point()), - ); - map.insert(key, pipeline); - } - - // Compile the int-dot GEMV for any dispatch that asked for it. - // Unlike a shape, this is not something measurement can turn on, so - // it is always present when the plan says so. The packed-dot - // intrinsic follows the device's `shader_integer_dot_product` - // capability; the scalar expansion is the exact fallback. - let knobs = matmul_knobs; - for dispatch in &plan.dispatches { - if !dispatch.gemv_int_dot { - continue; - } - let fused = dispatch.gemv_rmsnorm.is_some(); - let shape = dispatch.gemv_shape.unwrap_or_else(|| { - crate::codegen::GemvShape::initial(dispatch.shader.shader_group()) - }); - let key = if fused { - Variant::GemvRmsNormIntDot(dispatch.shader.clone(), dispatch.weight_format, shape) - } else { - Variant::GemvIntDot(dispatch.shader.clone(), dispatch.weight_format, shape) - }; - if let std::collections::hash_map::Entry::Vacant(slot) = map.entry(key.clone()) { - let sm = crate::codegen::generate_module_gemv_int_dot( - dispatch.shader.shader_group(), - dispatch.weight_format, - shape, - knobs.integer_dot, - fused, - ); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = if fused { - ::layout() - } else { - shader_data_layout(&dispatch.shader) - }; - slot.insert(create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(dispatch.shader.entry_point()), - )); - } - } - - // Compile any GEMV shape a plan already carries. Tuning inserts its - // own pipelines as it measures, so this is for a plan that arrives - // with a shape on it — deserialized, or rebuilt after a swap. Without - // it the dispatch would quietly fall back to the group's initial - // shape, which is correct but silently discards the measurement. - for dispatch in &plan.dispatches { - // Its shaped pipeline was compiled above and ordinary GEMV is an - // invalid fallback because it changes the activation arithmetic. - if dispatch.gemv_int_dot { - continue; - } - let Some(shape) = dispatch.gemv_shape else { - continue; - }; - // The RmsNorm-fused form has its own module and its own binding - // layout, and is not shaped. Tuning never sets a shape on one; - // this keeps a hand-built plan from producing a pipeline whose - // bindings do not match the dispatch. - let Some(group) = crate::tune::gemv_group(&dispatch.shader) - .filter(|_| dispatch.gemv_rmsnorm.is_none()) - else { - continue; - }; - let key = Variant::Gemv(dispatch.shader.clone(), dispatch.weight_format, shape); - if let std::collections::hash_map::Entry::Vacant(slot) = map.entry(key.clone()) { - let sm = crate::codegen::generate_module_gemv(group, dispatch.weight_format, shape); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = shader_data_layout(&dispatch.shader); - slot.insert(create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(dispatch.shader.entry_point()), - )); - } - } - - // Compile weight-format-specific pipelines (f16, Q4, Q8). - for (&format, groups) in &needed_weighted { - for &group in groups { - compile_variant( - crate::codegen::generate_module_weighted(group, format, matmul_knobs), - group, - &|entry| Variant::Weight(entry, format), - &mut map, - ); - } - } - for dispatch in &plan.dispatches { - if dispatch.use_small_tiles - && dispatch.weight_format.uses_reduced_storage() - && epilogue_pipeline_key(dispatch).is_none() - { - let key = Variant::WeightSmall(dispatch.shader.clone(), dispatch.weight_format); - if let std::collections::hash_map::Entry::Vacant(slot) = map.entry(key.clone()) { - let module = tuning::tile_module( - dispatch, - crate::tune::MatmulTile::Tile32, - matmul_knobs, - ); - let shader = create_gen_shader(gpu, module, wgsl_dump_dir); - slot.insert(create_profiled_pipeline( - gpu, - key.label(), - &shader_data_layout(&dispatch.shader), - shader.at(dispatch.shader.entry_point()), - )); - } - } - } - - // Compile epilogue-fused pipelines for dispatches with non-empty epilogue. - // Prefer the new MatMulEpilogue (PointwiseDAG); fall back to legacy - // Vec for cached plans that predate the DAG migration. - for dispatch in &plan.dispatches { - let Some(epilogue_key) = epilogue_pipeline_key(dispatch) else { - continue; - }; - let key = Variant::Epilogue(dispatch.shader.clone(), epilogue_key.clone()); - let coop_key = Variant::CoopEpilogue(dispatch.shader.clone(), epilogue_key); - if let std::collections::hash_map::Entry::Vacant(slot) = map.entry(key) { - let group = dispatch.shader.shader_group(); - let epilogue = match dispatch.matmul_epilogue { - Some(ref epi) => crate::codegen::EpilogueSource::Dag(epi), - None => crate::codegen::EpilogueSource::Ops(&dispatch.epilogue), - }; - let sm = crate::codegen::generate_matmul_with_epilogue( - group, - epilogue, - crate::codegen::MatMulOptions { - format: dispatch.weight_format, - tile: epilogue_tile(dispatch), - knobs: matmul_knobs, - }, - ); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = shader_data_layout(&dispatch.shader); - let pipeline = create_profiled_pipeline( - gpu, - slot.key().label(), - &layout, - shader.at(dispatch.shader.entry_point()), - ); - slot.insert(pipeline); - } - - if dispatch.use_coop && !map.contains_key(&coop_key) { - let config = coop_config - .expect("dispatch selected cooperative epilogue without a coop config"); - let epi = dispatch - .matmul_epilogue - .as_ref() - .expect("cooperative epilogues require the PointwiseDAG representation"); - let sm = crate::codegen::generate_coop_matmul_with_dag_epilogue( - dispatch.shader.shader_group(), - config, - epi, - ); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = shader_data_layout(&dispatch.shader); - let pipeline = create_profiled_pipeline( - gpu, - coop_key.label(), - &layout, - shader.at(dispatch.shader.entry_point()), - ); - map.insert(coop_key, pipeline); - } - } - - // Compile prologue-fused coop-matmul pipelines. One pipeline per - // (shader, prologue kind sequence); buffer IDs are bound at dispatch - // time via the existing matmul data layout extended with prologue - // factor buffers. - if let Some(coop_cfg) = coop_config { - for dispatch in &plan.dispatches { - let Some(ref prologue) = dispatch.matmul_prologue else { - continue; }; - if !dispatch.use_coop { - continue; - } - let kinds = prologue.factors.iter().map(|f| f.1.clone()).collect(); - let key = Variant::CoopPrologue(dispatch.shader.clone(), kinds); - if map.contains_key(&key) { - continue; - } - let group = dispatch.shader.shader_group(); - let Some((fused_add, variant)) = crate::codegen::coop_shape(group) else { - continue; - }; - let sm = crate::codegen::gen_matmul_coop_with_prologue( - fused_add, variant, coop_cfg, prologue, - ); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = matmul_with_prologue_layout(prologue.factors.len()); - let pipeline = create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(dispatch.shader.entry_point()), - ); - map.insert(key, pipeline); + layout = horizontal_matmul_layout(count); + entry_point = "main"; + crate::codegen::generate_horizontal_matmul(group, count, coop.as_ref()) } - } - - // Compile schedule-template pointwise pipelines. Each unique DAG - // (keyed by its content hash) gets one pipeline; the dispatch's - // existing `shader` field is used only to pick the data layout - // (UnaryData for n=1 inputs, BinaryData for n=2), which already - // matches the generated WGSL's binding names. - for dispatch in &plan.dispatches { - let dag = match dispatch.pointwise { - Some(ref d) => d, - None => continue, - }; - let key = Variant::Pointwise(dag.hash_key()); - if map.contains_key(&key) { - continue; - } - let template = crate::schedule::KernelTemplate::Pointwise { - dag: dag.clone(), - grid: crate::schedule::GridShape::default(), - }; - let sm = crate::schedule::lower(&template); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = pointwise_data_layout(dag.n_inputs); - let pipeline = create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(crate::schedule::POINTWISE_ENTRY), - ); - map.insert(key, pipeline); - } - - // Compile schedule-template reduction pipelines. - for dispatch in &plan.dispatches { - let kernel = match dispatch.reduction { - Some(ref k) => k, - None => continue, - }; - let key = Variant::Reduction(kernel.hash_key()); - if map.contains_key(&key) { - continue; - } - let sm = crate::schedule::lower(&kernel.to_template()); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = reduction_data_layout(kernel); - let pipeline = create_profiled_pipeline( - gpu, - key.label(), - &layout, - shader.at(crate::schedule::REDUCTION_ENTRY), - ); - map.insert(key, pipeline); - } - - for dispatch in &plan.dispatches { - let count = dispatch.horizontal_batch; - if count < 2 { - continue; - } - let kind = horiz_kind(dispatch); - let key = Variant::Horizontal(dispatch.shader.clone(), count, kind); - if map.contains_key(&key) { - continue; - } - let coop = match kind { - HorizMatMulKind::Scalar => None, - HorizMatMulKind::Coop | HorizMatMulKind::CoopCompensated => coop_config.map(|c| { - let mut cfg = *c; - cfg.compensated = matches!(kind, HorizMatMulKind::CoopCompensated); - cfg - }), - }; - let sm = crate::codegen::generate_horizontal_matmul( - dispatch.shader.shader_group(), - count, - coop.as_ref(), - ); - let shader = create_gen_shader(gpu, sm, wgsl_dump_dir); - let layout = horizontal_matmul_layout(count); - let pipeline = create_profiled_pipeline(gpu, key.label(), &layout, shader.at("main")); - map.insert(key, pipeline); - } - - let mut pipelines = Self { - map, - matmul_knobs, - dump_dir: wgsl_dump_dir.map(str::to_string), }; - for dispatch in &plan.dispatches { - if dispatch.conv_k_tile.is_some() { - let tile = crate::tune::MatmulTile::selected(dispatch, None) - .expect("scalar convolution specialization"); - pipelines - .ensure_tune_tile(gpu, dispatch, tile) - .expect("selected convolution pipeline"); + let shader = create_gen_shader(gpu, module, self.dump_dir.as_deref())?; + let pipeline = create_profiled_pipeline(gpu, key.label(), &layout, shader.at(entry_point)); + self.map.insert(key, pipeline); + Ok(()) + } + + fn attention_head_dim(dispatch: &Dispatch) -> Option { + use crate::codegen::ShaderGroup; + match dispatch.shader.shader_group() { + ShaderGroup::CachedBlockAttention + | ShaderGroup::CachedBlockAttentionSplit + | ShaderGroup::CachedBlockAttentionCombine => { + CachedBlockAttentionParams::from_words(&dispatch.params).map(|p| p.head_dim) } + ShaderGroup::MultiHeadAttn + | ShaderGroup::FlashAttention + | ShaderGroup::FlashAttentionCoop + | ShaderGroup::FlashGradQ + | ShaderGroup::FlashGradQCoop + | ShaderGroup::FlashGradKV + | ShaderGroup::FlashGradKVCoop => dispatch.params.get(3).copied(), + _ => None, } - pipelines } - /// Pipelines this dispatch can run, most specific first; `get` takes the - /// first one that was compiled. - /// - /// Every list ends with `Scalar`, which is always compiled - except for - /// epilogue fusion, where the plan has already deleted the standalone - /// ops, so running the unfused matmul would silently drop them. That - /// list stays a single entry and a miss is a panic. - fn candidates(dispatch: &Dispatch) -> Vec { - let entry = &dispatch.shader; - if let Some(k_tile) = dispatch.conv_k_tile { - return vec![Variant::SpecializedConv( - entry.clone(), - dispatch.params.clone(), - k_tile, - )]; + /// Resolve one implementation. Geometry and arithmetic must not depend on + /// which unrelated pipelines happen to have been compiled. + fn key(dispatch: &Dispatch) -> Variant { + let entry = dispatch.shader.clone(); + if let Some(shape) = dispatch.scalar_matmul() { + return Variant::ScalarMatmul(entry, dispatch.weight_format, shape); + } + if let Some(k_tile) = dispatch.conv_k_tile() { + return Variant::SpecializedConv(entry, dispatch.params.clone(), k_tile); } if dispatch.horizontal_batch >= 2 { - return vec![Variant::Horizontal( - dispatch.shader.clone(), - dispatch.horizontal_batch, - horiz_kind(dispatch), - )]; + return Variant::Horizontal(entry, dispatch.horizontal_batch, horiz_kind(dispatch)); } if let Some(epilogue) = epilogue_pipeline_key(dispatch) { - return vec![if dispatch.use_coop { - Variant::CoopEpilogue(entry.clone(), epilogue) + return if dispatch.use_coop() { + Variant::CoopEpilogue(entry, epilogue) } else { - Variant::Epilogue(entry.clone(), epilogue) - }]; - } - // The int-dot GEMV computes something no other variant here does, - // so it is the whole list and a miss is a panic — as for epilogue - // fusion. Falling through to the ordinary GEMV would put the - // activation back in f32 and change the session's arithmetic on the - // strength of a missing pipeline. - if dispatch.gemv_int_dot { + Variant::Epilogue(entry, epilogue) + }; + } + if dispatch.gemv_int_dot() || dispatch.gemv_rmsnorm.is_some() { let shape = dispatch - .gemv_shape + .gemv_shape() .unwrap_or_else(|| crate::codegen::GemvShape::initial(entry.shader_group())); - let format = dispatch.weight_format; - return vec![if dispatch.gemv_rmsnorm.is_some() { - Variant::GemvRmsNormIntDot(entry.clone(), format, shape) + return if dispatch.gemv_int_dot() { + if dispatch.gemv_rmsnorm.is_some() { + Variant::GemvRmsNormIntDot(entry, dispatch.weight_format, shape) + } else { + Variant::GemvIntDot(entry, dispatch.weight_format, shape) + } } else { - Variant::GemvIntDot(entry.clone(), format, shape) - }]; + Variant::GemvRmsNorm(entry, dispatch.weight_format, shape) + }; } - let mut out = Vec::new(); - if let Some(ref kernel) = dispatch.reduction { - out.push(Variant::Reduction(kernel.hash_key())); + if let Some(kernel) = dispatch.reduction() { + return Variant::Reduction(kernel.hash_key()); } - if let Some(ref dag) = dispatch.pointwise { - out.push(Variant::Pointwise(dag.hash_key())); + if let Some(dag) = dispatch.pointwise() { + return Variant::Pointwise(dag.hash_key()); } - if dispatch.params.len() >= 4 { - out.push(Variant::Attention(entry.clone(), dispatch.params[3])); + if let Some(dim) = Self::attention_head_dim(dispatch) { + return Variant::Attention(entry, dim); } - // A measured shape outranks the group's initial one, and the RmsNorm - // fusion outranks both: folding the norm in removes a whole dispatch, - // which no reduction choice can make up for. Shapes are only measured - // for the unfused forms, so the two never compete for the same - // dispatch — the ordering just makes that explicit. - if dispatch.gemv_rmsnorm.is_some() { - let shape = dispatch - .gemv_shape - .unwrap_or_else(|| crate::codegen::GemvShape::initial(entry.shader_group())); - return vec![Variant::GemvRmsNorm( - entry.clone(), - dispatch.weight_format, - shape, - )]; - } - if let Some(shape) = dispatch.gemv_shape { - out.push(Variant::Gemv(entry.clone(), dispatch.weight_format, shape)); + if let Some(shape) = dispatch.gemv_shape() { + return Variant::Gemv(entry, dispatch.weight_format, shape); } if dispatch.weight_format.uses_reduced_storage() { - if dispatch.use_small_tiles { - out.push(Variant::WeightSmall(entry.clone(), dispatch.weight_format)); - } - out.push(Variant::Weight(entry.clone(), dispatch.weight_format)); + return if dispatch.use_small_tiles() { + Variant::WeightSmall(entry, dispatch.weight_format) + } else { + Variant::Weight(entry, dispatch.weight_format) + }; } - if dispatch.use_coop { + if dispatch.use_coop() { if let Some(ref prologue) = dispatch.matmul_prologue { - let kinds = prologue.factors.iter().map(|f| f.1.clone()).collect(); - out.push(Variant::CoopPrologue(entry.clone(), kinds)); + return Variant::CoopPrologue( + entry, + prologue.factors.iter().map(|f| f.1.clone()).collect(), + ); } - if dispatch.use_coop_compensated { - out.push(Variant::CoopCompensated(entry.clone())); + return if dispatch.use_coop_compensated() { + Variant::CoopCompensated(entry) } else { - out.push(Variant::Coop(entry.clone())); - } + Variant::Coop(entry) + }; } - if dispatch.use_small_tiles { - out.push(Variant::SmallTile(entry.clone())); + if dispatch.use_small_tiles() { + return Variant::SmallTile(entry); } - out.push(Variant::Scalar(entry.clone())); - out + Variant::Scalar(entry) } /// The unmodified pipeline for an entry, for the fixed passes - @@ -2086,19 +1592,22 @@ impl Pipelines { &self.map[&Variant::Scalar(entry)] } - fn get(&self, dispatch: &Dispatch) -> &blade_graphics::ComputePipeline { - let candidates = Self::candidates(dispatch); - candidates + fn select(&mut self, dispatches: &[Dispatch]) { + self.selected = dispatches .iter() - .find_map(|variant| self.map.get(variant)) - .unwrap_or_else(|| panic!("no pipeline was compiled for any of {candidates:?}")) + .map(|dispatch| { + let key = Self::key(dispatch); + assert!( + self.map.contains_key(&key), + "no pipeline was compiled for {dispatch:?}" + ); + key + }) + .collect(); } - fn profile_key(&self, dispatch: &Dispatch) -> String { - Self::candidates(dispatch) - .into_iter() - .find(|variant| self.map.contains_key(variant)) - .map_or_else(|| format!("{:?}:scalar", dispatch.shader), |v| v.label()) + fn get(&self, dispatch_index: usize) -> &blade_graphics::ComputePipeline { + &self.map[&self.selected[dispatch_index]] } fn all_pipelines(&self) -> Vec<(&str, &blade_graphics::ComputePipeline)> { @@ -2357,6 +1866,7 @@ pub fn shader_data_layout(entry: &ShaderEntry) -> blade_graphics::ShaderDataLayo ShaderEntry::FusedMatMulAdd | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd + | ShaderEntry::MatMulGemvBTAdd | ShaderEntry::MatMulGemvAdd => FusedMatMulAddData::layout(), ShaderEntry::Relu | ShaderEntry::Sigmoid @@ -2522,6 +2032,54 @@ fn compute_groups(dispatches: &[Dispatch]) -> Vec> { groups } +/// Record the compiled graph, leaving the last chunk open for appended work. +#[allow(clippy::too_many_arguments)] +fn record_groups( + gpu: &Gpu, + encoder: &mut blade_graphics::CommandEncoder, + sync_point: &mut Option, + plan: &ExecutionPlan, + groups: &[std::ops::Range], + pipelines: &Pipelines, + buffers: &[blade_graphics::Buffer], + chunks: usize, +) { + let total = groups.len(); + let per_chunk = total.div_ceil(chunks.min(total).max(1)); + let chunk_count = if total == 0 { + 0 + } else { + total.div_ceil(per_chunk) + }; + let mut start = 0; + let mut chunk_index = 0; + while start < total { + let end = (start + per_chunk).min(total); + { + let label = format!("step {}/{}", chunk_index + 1, chunk_count); + let mut pass = encoder.compute(&label); + for (gi, group) in groups.iter().enumerate().take(end).skip(start) { + if gi > start { + pass.barrier(); + } + for i in group.clone() { + let dispatch = &plan.dispatches[i]; + let pipeline = pipelines.get(i); + let mut pc = pass.with(pipeline); + Session::bind_dispatch(buffers, dispatch, &mut pc); + pc.dispatch(dispatch.workgroups); + } + } + } + start = end; + chunk_index += 1; + if start < total { + *sync_point = Some(gpu.submit(encoder)); + encoder.start(); + } + } +} + // ---- Session ---- /// A compiled, ready-to-execute GPU session. @@ -2555,7 +2113,7 @@ pub(crate) fn select_variants( // iOS and future 8×8 f32 advertisers need the same veto. let apple_f32_coop = !config.use_f16_input && config.tile_size == 8; for dispatch in &mut plan.dispatches { - if dispatch.conv_k_tile.is_some() { + if dispatch.conv_k_tile().is_some() || dispatch.scalar_matmul().is_some() { continue; } // Autodiff marks derivative work as requiring f32 operands. A @@ -2571,9 +2129,9 @@ pub(crate) fn select_variants( // Cooperative epilogues are supported for the compiler's // current unary PointwiseDAG chains. They stage matrix // accumulators through workgroup memory, then apply scalar - // WGSL with bounds checks. Legacy-only plans and DAGs needing - // extra storage bindings stay on scalar geometry. - if !dispatch.epilogue.is_empty() || dispatch.matmul_epilogue.is_some() { + // WGSL with bounds checks. DAGs needing extra storage bindings + // stay on scalar geometry. + if dispatch.matmul_epilogue.is_some() { match dispatch.matmul_epilogue { Some(ref epilogue) if epilogue.inputs.is_empty() => {} _ => continue, @@ -2701,11 +2259,7 @@ pub(crate) fn select_variants( }; let _ = k; if coop_wgs >= min_wgs && !dispatch.weight_format.uses_reduced_storage() && vec4_ok { - // Retain pre-promotion geometry for diagnostics and future - // complete cooperative candidates. Survives reordering. - dispatch.scalar_fallback = Some((dispatch.shader.clone(), dispatch.workgroups)); - dispatch.use_coop = true; - dispatch.use_coop_compensated = false; + dispatch.kernel = crate::compile::Kernel::Cooperative; // Route conv2d coop dispatches to generated specialized kernels if is_conv_bwd { let kh = dispatch.params[5]; @@ -2759,8 +2313,9 @@ pub(crate) fn select_variants( { use crate::codegen::ShaderGroup; for dispatch in plan.dispatches.iter_mut() { - if dispatch.use_coop - || dispatch.use_small_tiles + if dispatch.use_coop() + || dispatch.use_small_tiles() + || dispatch.scalar_matmul().is_some() || dispatch.weight_format.uses_reduced_storage() { continue; @@ -2775,7 +2330,7 @@ pub(crate) fn select_variants( | ShaderGroup::MatMulBT ); if has_small && wgs_64 < 16 { - dispatch.use_small_tiles = true; + dispatch.kernel = crate::compile::Kernel::SmallTile; let (m, n) = match group { ShaderGroup::MatMul | ShaderGroup::MatMulAdd => { (dispatch.params[0], dispatch.params[2]) @@ -2827,8 +2382,8 @@ mod block_matmul_variant_tests { select_variants(&mut grouped, Some(&coop), false, false); let serial = &serial.dispatches[0]; let grouped = &grouped.dispatches[0]; - assert!(!serial.use_coop && !grouped.use_coop); - assert_eq!(serial.use_small_tiles, grouped.use_small_tiles); + assert!(!serial.use_coop() && !grouped.use_coop()); + assert_eq!(serial.use_small_tiles(), grouped.use_small_tiles()); assert_eq!(serial.workgroups[..2], grouped.workgroups[..2]); assert_eq!(grouped.workgroups[2], 8); } @@ -3708,8 +3263,8 @@ impl Session { d.label.clone() }, d.origin, - d.use_coop, - d.matmul_epilogue.is_some() || !d.epilogue.is_empty(), + d.use_coop(), + d.matmul_epilogue.is_some(), d.input_buffers.iter().map(|b| b.0).collect::>(), d.output_buffer.0, d.workgroups, @@ -4267,6 +3822,8 @@ impl Session { /// upper bound, since a short plan can produce fewer chunks. Profile the /// end-to-end workload on the target device: extra submissions have CPU /// and driver overhead and can make either workload slower. + /// [`crate::train::build_measured`] can measure this choice on initialized inputs + /// when optimizing this graph's latency without competing queue users. /// /// Correctness across the resulting submission boundaries does not need /// extra synchronization. Blade ends every command buffer with a @@ -4396,11 +3953,7 @@ impl Session { /// fused RmsNorm, fused epilogue or prologue, attention width, or a /// generated pointwise/reduction kernel. pub fn dispatch_pipeline_keys(&self) -> Vec { - self.plan - .dispatches - .iter() - .map(|dispatch| self.pipelines.profile_key(dispatch)) - .collect() + self.pipelines.selected.iter().map(Variant::label).collect() } /// Shared handle to the underlying Blade GPU context. @@ -4585,13 +4138,13 @@ mod variant_tests { } } - fn matmul(f: impl FnOnce(&mut Dispatch)) -> Vec { + fn matmul(f: impl FnOnce(&mut Dispatch)) -> Variant { let mut d = Dispatch { shader: ShaderEntry::MatMul, ..Default::default() }; f(&mut d); - Pipelines::candidates(&d) + Pipelines::key(&d) } /// The order here is what makes a modifier win over a less specific @@ -4599,21 +4152,18 @@ mod variant_tests { /// dispatch runs, so it is pinned rather than left to review. #[test] fn modifiers_outrank_the_scalar_base() { - assert_eq!(matmul(|_| {}), vec![Variant::Scalar(ShaderEntry::MatMul)]); + assert_eq!(matmul(|_| {}), Variant::Scalar(ShaderEntry::MatMul)); assert_eq!( matmul(|d| { - d.use_coop = true; - d.use_small_tiles = true; + d.kernel = crate::compile::Kernel::SmallTile; d.weight_format = crate::compile::WeightFormat::F16; }), - vec![ - Variant::WeightSmall(ShaderEntry::MatMul, crate::compile::WeightFormat::F16), - Variant::Weight(ShaderEntry::MatMul, crate::compile::WeightFormat::F16), - Variant::Coop(ShaderEntry::MatMul), - Variant::SmallTile(ShaderEntry::MatMul), - Variant::Scalar(ShaderEntry::MatMul), - ], + Variant::WeightSmall(ShaderEntry::MatMul, crate::compile::WeightFormat::F16), + ); + assert_eq!( + matmul(|d| d.kernel = crate::compile::Kernel::Cooperative), + Variant::Coop(ShaderEntry::MatMul), ); } @@ -4625,22 +4175,19 @@ mod variant_tests { let candidates = matmul(|d| { d.matmul_epilogue = Some(relu_epilogue()); }); - assert_eq!(candidates.len(), 1); - assert!(matches!(candidates[0], Variant::Epilogue(..))); + assert!(matches!(candidates, Variant::Epilogue(..))); let coop = matmul(|d| { d.matmul_epilogue = Some(relu_epilogue()); - d.use_coop = true; + d.kernel = crate::compile::Kernel::Cooperative; }); - assert_eq!(coop.len(), 1); - assert!(matches!(coop[0], Variant::CoopEpilogue(..))); + assert!(matches!(coop, Variant::CoopEpilogue(..))); let q4 = matmul(|d| { d.matmul_epilogue = Some(relu_epilogue()); d.weight_format = crate::compile::WeightFormat::Q4; }); - assert_eq!(q4.len(), 1); - assert!(matches!(q4[0], Variant::Epilogue(..))); + assert!(matches!(q4, Variant::Epilogue(..))); assert_ne!( epilogue_pipeline_key(&{ let mut d = Dispatch { @@ -4671,7 +4218,7 @@ mod variant_tests { ..Default::default() }; let small = Dispatch { - use_small_tiles: true, + kernel: crate::compile::Kernel::SmallTile, ..large.clone() }; assert_ne!( @@ -4690,36 +4237,23 @@ mod variant_tests { fn epilogue_dispatch_does_not_request_unreachable_variants() { let epilogue_small = matmul(|d| { d.matmul_epilogue = Some(relu_epilogue()); - d.use_small_tiles = true; + d.kernel = crate::compile::Kernel::SmallTile; d.weight_format = crate::compile::WeightFormat::Q4; }); - assert_eq!( - epilogue_small.len(), - 1, - "an epilogue dispatch offers exactly one pipeline, got {epilogue_small:?}" - ); - assert!(matches!(epilogue_small[0], Variant::Epilogue(..))); + assert!(matches!(epilogue_small, Variant::Epilogue(..))); - // The same modifiers without an epilogue do need those variants. + // Removing the epilogue still preserves both tile and storage format. let plain_small = matmul(|d| { - d.use_small_tiles = true; + d.kernel = crate::compile::Kernel::SmallTile; d.weight_format = crate::compile::WeightFormat::Q4; }); - assert!( - plain_small.contains(&Variant::SmallTile(ShaderEntry::MatMul)), - "a plain small-tile dispatch still needs the small-tile pipeline" - ); - assert!( - plain_small.contains(&Variant::Weight( - ShaderEntry::MatMul, - crate::compile::WeightFormat::Q4 - )), - "a plain weighted dispatch still needs the weighted pipeline" + assert_eq!( + plain_small, + Variant::WeightSmall(ShaderEntry::MatMul, crate::compile::WeightFormat::Q4), ); } - /// The compile-side collection has to mirror `pipeline_variants`. That - /// side is what actually builds modules, so the check goes through + /// Preparation and execution use the same key. Check this through /// `Pipelines::new`: a variant absent from the map is a kernel that was /// never compiled. #[test] @@ -4737,7 +4271,7 @@ mod variant_tests { crate::compile::compile(&g) }; select_variants(&mut demoted, None, false, false); - assert!(demoted.dispatches[0].use_small_tiles); + assert!(demoted.dispatches[0].use_small_tiles()); let pipelines = Pipelines::new(&gpu, &demoted, None, None); assert!( pipelines @@ -4799,7 +4333,7 @@ mod variant_tests { ); select_variants(&mut plan, None, false, false); let d = &plan.dispatches[0]; - assert!(d.use_small_tiles, "64×64 matmul should demote to 32×32"); + assert!(d.use_small_tiles(), "64×64 matmul should demote to 32×32"); assert_eq!( d.workgroups, [2, 2, 1], @@ -4817,30 +4351,21 @@ mod variant_tests { let scalar = matmul(|d| d.horizontal_batch = 3); let coop = matmul(|d| { d.horizontal_batch = 3; - d.use_coop = true; + d.kernel = crate::compile::Kernel::Cooperative; }); let compensated = matmul(|d| { d.horizontal_batch = 3; - d.use_coop = true; - d.use_coop_compensated = true; + d.kernel = crate::compile::Kernel::CooperativeCompensated; }); assert_eq!( scalar, - vec![Variant::Horizontal( - ShaderEntry::MatMul, - 3, - HorizMatMulKind::Scalar - )] + Variant::Horizontal(ShaderEntry::MatMul, 3, HorizMatMulKind::Scalar) ); assert_ne!(scalar, coop); assert_ne!(coop, compensated); assert_eq!( compensated, - vec![Variant::Horizontal( - ShaderEntry::MatMul, - 3, - HorizMatMulKind::CoopCompensated - )] + Variant::Horizontal(ShaderEntry::MatMul, 3, HorizMatMulKind::CoopCompensated) ); } } @@ -5577,6 +5102,18 @@ impl Session { for (derived_buf, sources, transform) in derived { let sources = sources.as_slice(); match transform { + crate::graph::ParamTransform::VerticalConcat => { + let ty = &self.plan.param_types[&derived_buf]; + let bytes = match ty.dtype { + crate::graph::DType::F32 => bytemuck::cast_slice(data).to_vec(), + crate::graph::DType::F16 => data + .iter() + .flat_map(|&v| half::f16::from_f32(v).to_le_bytes()) + .collect(), + _ => unreachable!("row concatenation requires dense weights"), + }; + self.copy_parameter_rows(derived_buf, name, &bytes, sources); + } crate::graph::ParamTransform::HorizontalConcat => { let total_cols: usize = sources.iter().map(|s| s.1).sum(); let derived_fmt = self @@ -5744,8 +5281,8 @@ impl Session { /// the values. The byte count is validated against the logical tensor /// size before anything is copied to the device. /// - /// HorizontalConcat derived parameters (the SwiGLU `gate+up` fusion) - /// are restaged in packed space. There is no K-quant encoder here, and + /// Derived gate/up concatenations are restaged in their storage format. + /// There is no K-quant encoder here, and /// Q4/Q6_K packed blobs are not a byte-append of their sources. pub fn set_parameter_packed(&mut self, name: &str, data: &[u8]) { self.wait(); @@ -5775,19 +5312,79 @@ impl Session { .derived_params .iter() .filter(|entry| { - matches!(entry.2, crate::graph::ParamTransform::HorizontalConcat) - && entry.1.iter().any(|s| s.0 == name) + matches!( + entry.2, + crate::graph::ParamTransform::HorizontalConcat + | crate::graph::ParamTransform::VerticalConcat + ) && entry.1.iter().any(|s| s.0 == name) }) - .map(|entry| (entry.0, entry.1.clone())) + .cloned() .collect(); - for (derived_buf, sources) in derived { - self.restage_packed_concat(derived_buf, name, data, &sources); + for (derived_buf, sources, transform) in derived { + match transform { + crate::graph::ParamTransform::HorizontalConcat => { + self.restage_packed_concat(derived_buf, name, data, &sources); + } + crate::graph::ParamTransform::VerticalConcat => { + self.copy_parameter_rows(derived_buf, name, data, &sources); + } + _ => unreachable!(), + } } return; } panic!("unknown parameter: {name}"); } + fn copy_parameter_rows( + &self, + derived_buf: BufferRef, + name: &str, + data: &[u8], + sources: &[(String, usize)], + ) { + let ty = &self.plan.param_types[&derived_buf]; + assert!(matches!( + ty.dtype, + crate::graph::DType::F32 | crate::graph::DType::F16 + )); + assert_eq!(ty.shape.len(), 2); + assert_eq!(ty.shape[0], sources.iter().map(|s| s.1).sum::()); + let row_bytes = ty.shape[1] + * if ty.dtype == crate::graph::DType::F16 { + 2 + } else { + 4 + }; + let host_visible = self.logical_host_visible(derived_buf); + let mut offset = 0usize; + for &(ref source, rows) in sources { + let bytes = rows * row_bytes; + if source == name { + assert_eq!(data.len(), bytes); + if host_visible || (offset.is_multiple_of(4) && bytes.is_multiple_of(4)) { + self.write_raw_buffer_at( + self.buffers[derived_buf.0 as usize].at(offset as u64), + data, + host_visible, + ); + } else { + // Odd f16 row ranges need a word-aligned transfer. Preserve + // the current device image, including checkpoint/shared + // updates, instead of keeping a stale host-side copy. + let capacity = self.plan.buffers[derived_buf.0 as usize]; + assert!(capacity.is_multiple_of(4)); + let mut image = vec![0.0f32; capacity / 4]; + self.read_buffer(derived_buf, &mut image); + let image: &mut [u8] = bytemuck::cast_slice_mut(&mut image); + image[offset..offset + bytes].copy_from_slice(data); + self.upload_buffer(derived_buf, image); + } + } + offset += bytes; + } + } + fn restage_packed_concat( &mut self, derived_buf: BufferRef, @@ -6441,6 +6038,18 @@ impl Session { ); } + /// Read complete buffers as F32 views in one staged transfer, preserving + /// request order. Includes any declared padding, just like [`Self::read_buffer`]. + /// Call after waiting for the producing step. Useful for full-output checks + /// without separately probing mapped reads for every allocation. + pub fn read_buffers(&self, buffers: &[BufferRef]) -> Vec> { + let requests: Vec<_> = buffers + .iter() + .map(|b| (self.buffers[b.0 as usize], self.plan.buffers[b.0 as usize])) + .collect(); + self.read_f32_buffers(&requests, "buffer_readback") + } + /// Read back a graph output by index. /// /// Index 0 is the primary output (logits/loss). Higher indices are @@ -6450,6 +6059,27 @@ impl Session { self.read_buffer(buf_ref, out); } + /// Wait for pending work and read a graph output. + /// + /// Queues a staged download before waiting on the CPU. Mapped reads and + /// the initial readback probe still wait before accessing the buffer. + pub fn wait_read_output(&mut self, index: usize, out: &mut [f32]) { + let buf_ref = self.plan.output_buffers[index]; + let buffer = self.buffers[buf_ref.0 as usize]; + let staged = !self.logical_host_visible(buf_ref) + || self + .readback + .borrow() + .staged + .get(&(buffer.data() as usize, std::mem::size_of_val(out))) + == Some(&true); + if !staged { + self.wait(); + } + self.read_buffer(buf_ref, out); + self.wait(); + } + /// Number of graph outputs. pub fn num_outputs(&self) -> usize { self.plan.output_buffers.len() @@ -6990,7 +6620,7 @@ impl Session { match encoded { ProfilePass::Timed(i) => { let dispatch = &self.plan.dispatches[i]; - let pipeline = self.pipelines.get(dispatch); + let pipeline = self.pipelines.get(i); let mut pass = self.encoder.compute(&dispatch.label); let mut pc = pass.with(pipeline); Self::bind_dispatch(&self.buffers, dispatch, &mut pc); @@ -7007,7 +6637,7 @@ impl Session { } for i in span { let dispatch = &self.plan.dispatches[i]; - let pipeline = self.pipelines.get(dispatch); + let pipeline = self.pipelines.get(i); let mut pc = pass.with(pipeline); Self::bind_dispatch(&self.buffers, dispatch, &mut pc); pc.dispatch(dispatch.workgroups); @@ -7028,41 +6658,16 @@ impl Session { // typically) can slot its own work between them. The chunk // boundary needs no explicit barrier: blade closes each command // buffer with a conservative global one. - let total = self.groups.len(); - let per_chunk = total.div_ceil(self.submission_chunks.min(total).max(1)); - let mut start = 0; - let chunk_count = if total == 0 { - 0 - } else { - total.div_ceil(per_chunk) - }; - let mut chunk_index = 0; - while start < total { - let end = (start + per_chunk).min(total); - { - let label = format!("step {}/{}", chunk_index + 1, chunk_count); - let mut pass = self.encoder.compute(&label); - for gi in start..end { - if gi > start { - pass.barrier(); - } - let group = self.groups[gi].clone(); - for i in group { - let dispatch = &self.plan.dispatches[i]; - let pipeline = self.pipelines.get(dispatch); - let mut pc = pass.with(pipeline); - Self::bind_dispatch(&self.buffers, dispatch, &mut pc); - pc.dispatch(dispatch.workgroups); - } - } - } - start = end; - chunk_index += 1; - if start < total { - self.sync_point = Some(self.gpu.submit(&mut self.encoder)); - self.encoder.start(); - } - } + record_groups( + &self.gpu, + &mut self.encoder, + &mut self.sync_point, + &self.plan, + &self.groups, + &self.pipelines, + &self.buffers, + self.submission_chunks, + ); } // Temporal grad accumulation: add this step's (overwritten) grads @@ -7416,6 +7021,11 @@ impl Session { // A GEMV with its RmsNorm folded in takes the norm's weight vector // as an extra binding and carries eps in the params' spare slot. if let Some(ref rn) = dispatch.gemv_rmsnorm { + let (n, k) = if dispatch.shader == ShaderEntry::MatMulGemvBT { + (dispatch.params[1], dispatch.params[2]) + } else { + (dispatch.params[2], dispatch.params[1]) + }; pc.bind( 0, &MatMulRmsNormData { @@ -7425,8 +7035,8 @@ impl Session { matrix_c: buf(dispatch.output_buffer), params: MatMulRmsNormParams { m: dispatch.params[0], - n: dispatch.params[2], - k: dispatch.params[1], + n, + k, eps_bits: rn.eps_bits, }, }, @@ -7435,7 +7045,7 @@ impl Session { } // Schedule-template reduction dispatches have priority and route // by kernel arity (n_per_elem, n_per_row, n_per_col). - if let Some(ref k) = dispatch.reduction { + if let Some(k) = dispatch.reduction() { let params = ReductionParams { outer: dispatch.params[0], inner: dispatch.params[1], @@ -7504,7 +7114,7 @@ impl Session { // Schedule-template pointwise dispatches route by DAG arity, not // by the dummy `shader` entry — arity may be 3 after fusion, which // no ShaderEntry variant represents. - if let Some(ref dag) = dispatch.pointwise { + if let Some(dag) = dispatch.pointwise() { let params = UnaryParams { len: dispatch.params[0], _pad0: 0, @@ -7552,7 +7162,7 @@ impl Session { // Prologue-fused coop matmul: 2-factor prologue (RmsNorm rsqrt + w_norm). // Only applies when use_coop is set AND a prologue is attached. if let Some(ref prologue) = dispatch.matmul_prologue { - if dispatch.use_coop && prologue.factors.len() == 2 { + if dispatch.use_coop() && prologue.factors.len() == 2 { let (m, n, k) = match dispatch.shader { ShaderEntry::MatMul | ShaderEntry::FusedMatMulAdd => { (dispatch.params[0], dispatch.params[2], dispatch.params[1]) @@ -7641,7 +7251,9 @@ impl Session { }, ); } - ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd => { + ShaderEntry::FusedMatMulATAdd + | ShaderEntry::FusedMatMulBTAdd + | ShaderEntry::MatMulGemvBTAdd => { // params layout: [m, n, k, 0] (same as AT/BT, no swizzle) pc.bind( 0, @@ -8536,16 +8148,8 @@ impl Session { kv_pos_buf: buf(dispatch.input_buffers[3]), valid_len_buf: buf(dispatch.input_buffers[4]), dst: buf(dispatch.output_buffer), - params: CachedBlockAttentionParams { - window_size: dispatch.params[0], - num_heads: dispatch.params[1], - num_kv_heads: dispatch.params[2], - head_dim: dispatch.params[3], - block_len: dispatch.params[4], - max_seq: dispatch.params[5], - splits: dispatch.params[6], - chunk: dispatch.params[7], - }, + params: CachedBlockAttentionParams::from_words(&dispatch.params) + .expect("cached attention parameter layout"), }, ); } @@ -8555,16 +8159,8 @@ impl Session { &CachedBlockAttentionCombineData { partials: buf(dispatch.input_buffers[0]), dst: buf(dispatch.output_buffer), - params: CachedBlockAttentionParams { - window_size: dispatch.params[0], - num_heads: dispatch.params[1], - num_kv_heads: dispatch.params[2], - head_dim: dispatch.params[3], - block_len: dispatch.params[4], - max_seq: dispatch.params[5], - splits: dispatch.params[6], - chunk: dispatch.params[7], - }, + params: CachedBlockAttentionParams::from_words(&dispatch.params) + .expect("cached attention parameter layout"), }, ); } diff --git a/src/runtime/search_state.rs b/src/runtime/search_state.rs new file mode 100644 index 00000000..bd0cf3ab --- /dev/null +++ b/src/runtime/search_state.rs @@ -0,0 +1,134 @@ +//! Reset explicit persistent state between complete-plan search trials. +use super::{PhysicalBuffer, Session, safe_device_memory_remaining}; +use crate::compile::{BufferRef, ExecutionPlan}; +use blade_graphics as bg; +use std::{collections::HashSet, sync::Arc}; + +pub(crate) fn persistent_writes(plan: &ExecutionPlan) -> Vec { + let writes: HashSet<_> = plan + .dispatches + .iter() + .flat_map(|d| std::iter::once(d.output_buffer).chain(d.extra_outputs.iter().copied())) + .collect(); + let mut buffers: Vec<_> = plan + .input_buffers + .iter() + .chain(&plan.param_buffers) + .map(|entry| entry.1) + .chain(plan.constant_buffers.iter().map(|entry| entry.0)) + .filter(|b| writes.contains(b)) + .collect(); + buffers.sort_unstable_by_key(|b| b.0); + buffers.dedup(); + buffers +} + +pub(crate) struct SearchState { + images: Vec, +} + +struct SavedBuffer { + live: Arc, + original: PhysicalBuffer, + bytes: usize, +} + +impl SearchState { + pub fn capture(session: &mut Session, max_bytes: usize) -> Result { + // These passes have state outside the compiled plan. Do not silently + // measure a different workload or advance their host-side counters. + if session.pending_lr.is_some() + || session.pending_adam.is_some() + || session.grad_accum_scale.is_some() + { + return Err( + "configure runtime optimizers and accumulation after program search".into(), + ); + } + session.wait(); + let mut physical: Vec<_> = persistent_writes(&session.plan) + .iter() + .map(|b| session.alias.map[b.0 as usize]) + .collect(); + physical.sort_unstable(); + physical.dedup(); + let mut bytes = 0usize; + for &p in &physical { + if Arc::strong_count(&session.physical_buffers[p]) != 1 { + return Err("program search requires private writable state; do not share mutable parameters".into()); + } + bytes = bytes + .checked_add(session.alias.sizes[p].max(4)) + .filter(|&n| n <= max_bytes) + .ok_or("program search state snapshot exceeds its byte budget")?; + } + let memory = session.gpu.memory_stats(); + if memory.budget != 0 + && bytes as u64 > safe_device_memory_remaining(memory.usage, memory.budget) + { + return Err("program search state snapshot exceeds the device memory budget".into()); + } + let state = Self { + images: physical + .into_iter() + .map(|p| { + let size = session.alias.sizes[p].max(4); + let saved = PhysicalBuffer { + gpu: session.gpu.clone(), + handle: session.gpu.create_buffer(bg::BufferDesc { + name: "program_search_state", + size: size as u64, + memory: bg::Memory::DeviceTransient, + }), + }; + SavedBuffer { + live: session.physical_buffers[p].clone(), + original: saved, + bytes: size, + } + }) + .collect(), + }; + state.copy(session, false)?; + Ok(state) + } + + pub fn restore(&self, session: &mut Session) -> Result<(), String> { + self.copy(session, true) + } + + pub fn bytes(&self) -> usize { + self.images.iter().map(|image| image.bytes).sum() + } + + fn copy(&self, session: &mut Session, restore: bool) -> Result<(), String> { + if self.images.is_empty() { + return Ok(()); + } + session.wait(); + session.encoder.start(); + { + let mut pass = session.encoder.transfer("program_search_state"); + for image in &self.images { + let (live, saved) = (image.live.handle, image.original.handle); + let (src, dst) = if restore { + (saved, live) + } else { + (live, saved) + }; + pass.copy_buffer_to_buffer(src.at(0), dst.at(0), image.bytes as u64); + } + } + let sync = session.gpu.submit(&mut session.encoder); + session.sync_point = Some(sync.clone()); + if !session + .gpu + .wait_for(&sync, !0) + .map_err(|_| "state copy GPU wait failed")? + { + return Err("state copy GPU wait did not complete".into()); + } + session.sync_point = None; + Ok(()) + } +} diff --git a/src/runtime/tuning.rs b/src/runtime/tuning.rs index 9766ffc0..8703d55d 100644 --- a/src/runtime/tuning.rs +++ b/src/runtime/tuning.rs @@ -11,6 +11,19 @@ use std::{ time::{Duration, Instant}, }; +/// Qualified private-scratch comparisons within one graph search, on one device +/// and under one numerical/timing policy. Never persisted or shared globally. +#[derive(Default)] +pub(crate) struct KernelMemo( + HashMap<(crate::compile::TuningKnobs, TuneClass, Vec), KernelProgress>, +); + +#[derive(Clone, Copy)] +struct KernelProgress { + selected: MatmulTile, + next_candidate: usize, +} + struct PhaseTimer<'a> { start: Instant, elapsed: &'a mut Option, @@ -38,41 +51,13 @@ impl Pipelines { dispatch: &Dispatch, tile: MatmulTile, ) -> Result<(), String> { - let key = tile_variant(dispatch, tile); - if self.map.contains_key(&key) { - return Ok(()); - } - let selected_entry = tile.shader(&dispatch.shader); - let mut knobs = self.matmul_knobs; - knobs.integer_dot = gpu.capabilities().shader_integer_dot_product; - let module = tile_module(dispatch, tile, knobs); - if let Some(dir) = self.dump_dir.as_deref() { - module.dump(dir); - } - let shader = gpu - .try_create_shader(bg::ShaderDesc { - source: &module.source, - naga_module: Some(module.module), - }) - .map_err(|error| error.to_string())?; - let layout = if dispatch.gemv_rmsnorm.is_some() { - use blade_graphics::ShaderData; - super::MatMulRmsNormData::layout() - } else { - super::shader_data_layout(&selected_entry) - }; - let pipeline = gpu.create_compute_pipeline(bg::ComputePipelineDesc { - name: &key.label(), - data_layouts: &[&layout], - compute: shader.at(selected_entry.entry_point()), - }); - self.map.insert(key, pipeline); - Ok(()) + let mut selected = dispatch.clone(); + tile.configure(&mut selected); + self.prepare(gpu, &selected, tile.coop_config().as_ref()) } fn discard_unused_convolutions(&mut self, gpu: &Gpu, plan: &crate::compile::ExecutionPlan) { - let used: std::collections::HashSet<_> = - plan.dispatches.iter().flat_map(Self::candidates).collect(); + let used: std::collections::HashSet<_> = plan.dispatches.iter().map(Self::key).collect(); self.map.retain(|key, pipeline| { let convolution = key.entry().is_some_and(|entry| { matches!( @@ -101,13 +86,32 @@ pub(super) fn tile_module( knobs: crate::codegen::MatmulKnobs, ) -> crate::codegen::ShaderModule { let entry = &dispatch.shader; + if let MatmulTile::Scalar(shape) = tile { + return crate::codegen::generate_matmul_with_epilogue( + entry.shader_group(), + None, + crate::codegen::MatMulOptions { + format: dispatch.weight_format, + tile: match shape.tile_size { + 32 => crate::codegen::MatMulTile::Small, + 64 => crate::codegen::MatMulTile::Large, + _ => unreachable!("unsupported scalar tile"), + }, + knobs: crate::codegen::MatmulKnobs { + k_stage: shape.k_stage, + interleave_columns: shape.interleave_columns, + ..knobs + }, + }, + ); + } if let MatmulTile::Gemv(shape) = tile { // The int-dot kernel is a different computation, not a different // route to the same one, so a shape candidate has to stay inside it. // Generating the ordinary GEMV here would quietly swap the // activation back to f32 and change what the plan computes. let group = crate::tune::gemv_group(entry).expect("GEMV candidate on a GEMV entry"); - if dispatch.gemv_int_dot { + if dispatch.gemv_int_dot() { return crate::codegen::generate_module_gemv_int_dot( group, dispatch.weight_format, @@ -117,7 +121,11 @@ pub(super) fn tile_module( ); } if dispatch.gemv_rmsnorm.is_some() { - return crate::codegen::generate_module_gemv_rmsnorm(shape, dispatch.weight_format); + return crate::codegen::generate_module_gemv_rmsnorm( + group, + shape, + dispatch.weight_format, + ); } return crate::codegen::generate_module_gemv(group, dispatch.weight_format, shape); } @@ -125,7 +133,7 @@ pub(super) fn tile_module( if dispatch.weight_format.uses_reduced_storage() { return crate::codegen::generate_matmul_with_epilogue( selected_entry.shader_group(), - crate::codegen::EpilogueSource::Ops(&[]), + None, crate::codegen::MatMulOptions { format: dispatch.weight_format, tile: match tile { @@ -157,52 +165,17 @@ pub(super) fn tile_module( entry.shader_group(), &tile.coop_config().expect("cooperative candidate"), ), - MatmulTile::SpecializedConv { .. } | MatmulTile::Gemv(_) => unreachable!(), + MatmulTile::SpecializedConv { .. } | MatmulTile::Gemv(_) | MatmulTile::Scalar(_) => { + unreachable!() + } } } } -fn tile_variant(dispatch: &Dispatch, tile: MatmulTile) -> Variant { - let entry = &dispatch.shader; - if let MatmulTile::Gemv(shape) = tile { - if dispatch.gemv_int_dot { - if dispatch.gemv_rmsnorm.is_some() { - return Variant::GemvRmsNormIntDot(entry.clone(), dispatch.weight_format, shape); - } - return Variant::GemvIntDot(entry.clone(), dispatch.weight_format, shape); - } - if dispatch.gemv_rmsnorm.is_some() { - return Variant::GemvRmsNorm(entry.clone(), dispatch.weight_format, shape); - } - return Variant::Gemv(entry.clone(), dispatch.weight_format, shape); - } - if let MatmulTile::SpecializedConv { k_tile, .. } = tile { - return Variant::SpecializedConv(tile.shader(entry), dispatch.params.clone(), k_tile); - } - if dispatch.weight_format.uses_reduced_storage() { - return match tile { - MatmulTile::Tile32 => Variant::WeightSmall(entry.clone(), dispatch.weight_format), - MatmulTile::Tile64 => Variant::Weight(entry.clone(), dispatch.weight_format), - _ => unreachable!("reduced storage only changes scalar tile geometry"), - }; - } - if matches!( - entry, - ShaderEntry::Conv2dGemm - | ShaderEntry::Conv2dGemmSmall - | ShaderEntry::Conv2dGradInputGemm - | ShaderEntry::Conv2dGradInputGemmSmall - | ShaderEntry::Conv2dGradWeightGemm - | ShaderEntry::Conv2dGradWeightGemmSmall - ) { - return Variant::Scalar(tile.shader(entry)); - } - match tile { - MatmulTile::Tile32 => Variant::SmallTile(entry.clone()), - MatmulTile::Tile64 => Variant::Scalar(entry.clone()), - MatmulTile::CooperativeF32 { .. } => Variant::Coop(entry.clone()), - MatmulTile::SpecializedConv { .. } | MatmulTile::Gemv(_) => unreachable!(), - } +pub(super) fn tile_variant(dispatch: &Dispatch, tile: MatmulTile) -> Variant { + let mut selected = dispatch.clone(); + tile.configure(&mut selected); + Pipelines::key(&selected) } struct SearchClass { @@ -310,11 +283,33 @@ fn collect_classes( .map(|b| plan.buffers[b.0 as usize]) .collect(); let initial = MatmulTile::selected(dispatch, coop_config).expect("checked class geometry"); + let equivalent = if key.conv2d.is_none() && !key.weight_format.is_quantized() { + match initial { + MatmulTile::Tile32 | MatmulTile::Tile64 => { + Some(MatmulTile::Scalar(crate::codegen::ScalarMatmulShape { + tile_size: if initial == MatmulTile::Tile32 { + 32 + } else { + 64 + }, + k_stage: plan.knobs.matmul_k_stage, + interleave_columns: plan.knobs.matmul_interleave_columns, + })) + } + _ => None, + } + } else { + None + }; if !initial.fits(&key) { excluded += 1; continue; } - let challengers = key.challengers(initial, coop_config); + let challengers = key + .challengers(initial, coop_config) + .into_iter() + .filter(|candidate| Some(*candidate) != equivalent) + .collect(); let next_index = classes.len(); let class_index = *indices.entry((key.clone(), initial)).or_insert(next_index); if class_index == next_index { @@ -401,11 +396,14 @@ impl Session { swap.left .apply(&mut other.plan.dispatches[swap.index], &swap.class); } + self.pipelines.select(&self.plan.dispatches); + other.pipelines.select(&other.plan.dispatches); Ok(swaps.len()) } /// Bounded kernel search with default options; logs skips and returns - /// per-comparison evidence. Use [`Self::tune_with`] for budgets and full reporting. + /// matrix/convolution evidence for compatibility. Use [`Self::tune_with`] + /// for budgets and the complete report, including attention choices. /// /// Unlike the former family-wide tuner, this never calls `step()` and /// never reads or writes live tensor, optimizer, accumulator, or KV state. @@ -417,7 +415,7 @@ impl Session { } /// Search scalar tiles and advertised, smoke-tested native-f32 cooperative - /// matmul, scalar convolution shapes, and GEMV shapes for eligible classes. + /// matmul, scalar convolution/GEMV shapes, and cached-attention split counts. /// Occupancy/large-shape thresholds only /// choose the starting implementation; they do not remove challengers. /// @@ -436,14 +434,22 @@ impl Session { /// Other prologues/epilogues, horizontal packs, f16-input cooperative, /// cooperative reduced-storage and overlapping-binding dispatches are excluded. /// Winners live in this session, not the plan cache. - /// Only selected dispatch geometry and pipeline resources change. No graph - /// execution occurs, including when an optimizer or external buffer is bound. + /// No live graph execution occurs, including with optimizers or external buffers. /// Cooperative padding must fit each binding's declared size; the live - /// allocation/alias plan is never resized. Sequential challenger + /// matrix bindings are never resized. Structural choices are made by + /// `build_measured` before allocation. Sequential challenger /// comparisons per class reuse the latest fully qualified winner as the /// incumbent. A soft deadline may be exceeded by one in-flight operation; /// an incomplete comparison always retains its incumbent. pub fn tune_with(&mut self, options: TuneOptions) -> Result { + self.tune_with_memo(options, None) + } + + pub(crate) fn tune_with_memo( + &mut self, + options: TuneOptions, + mut memo: Option<&mut KernelMemo>, + ) -> Result { options.validate()?; let start = Instant::now(); let (mut classes, mut excluded_dispatches) = @@ -471,8 +477,33 @@ impl Session { break; } report.visited_classes += 1; + let memo_key = ( + self.plan.knobs, + class.key.clone(), + std::iter::once(class.initial) + .chain(class.challengers.iter().copied()) + .collect(), + ); let mut incumbent = class.initial; - for &candidate in &class.challengers { + let mut next_candidate = 0; + if let Some(&progress) = memo.as_ref().and_then(|m| m.0.get(&memo_key)) { + let selected = progress.selected; + if progress.next_candidate <= class.challengers.len() + && (selected == class.initial || class.challengers.contains(&selected)) + && self + .pipelines + .ensure_tune_tile(&gpu, &self.plan.dispatches[class.members[0]], selected) + .is_ok() + { + for &index in &class.members { + selected.apply(&mut self.plan.dispatches[index], &class.key); + } + report.reused_classes.push((class.key.clone(), selected)); + incumbent = selected; + next_candidate = progress.next_candidate; + } + } + for (index, &candidate) in class.challengers.iter().enumerate().skip(next_candidate) { if start.elapsed() >= options.max_time { report.time_budget_exhausted = true; break; @@ -507,6 +538,28 @@ impl Session { if let Some(ref failure) = outcome.failure { log::warn!("tune: {failure}"); } + if outcome.qualified + && matches!( + outcome.decision, + TuneDecision::FasterCandidate | TuneDecision::KeepBaseline + ) + { + // Only skip a contiguous, fully measured prefix. A failed + // or interrupted challenger is retried; later qualified + // winners may still seed that retry. + if index == next_candidate { + next_candidate += 1; + } + if let Some(ref mut memo) = memo { + memo.0.insert( + memo_key.clone(), + KernelProgress { + selected: incumbent, + next_candidate, + }, + ); + } + } report.outcomes.push(outcome); } } @@ -521,6 +574,7 @@ impl Session { self.wait(); self.pipelines.discard_unused_convolutions(&gpu, &self.plan); } + self.pipelines.select(&self.plan.dispatches); report.scratch = Some(staging.stats); report.elapsed = start.elapsed(); report.time_budget_exhausted |= report.elapsed >= options.max_time; @@ -740,9 +794,16 @@ impl Session { return; } let mut scratch = Scratch::new( - &class.key, + &sizes + .iter() + .enumerate() + .map(|(i, _)| { + i > output_index + || class.key.device_local[if i == output_index { 3 } else { i }] + }) + .collect::>(), &sizes, - output_index, + (output_index, class.key.output_elements()), bytes, staging, prep, @@ -924,7 +985,7 @@ fn split_dispatches( .buffer_sizes() .ok_or(TuneError("invalid split-K extents"))?; let mut dispatch = dispatch.clone(); - dispatch.conv_k_tile = None; // Split-K has its own unspecialized K=16 shader. + dispatch.kernel = crate::compile::Kernel::Default; // Split-K has its own K=16 shader. dispatch.input_buffers = vec![BufferRef(0), BufferRef(1)]; dispatch.output_buffer = BufferRef(2); plan.dispatches.push(dispatch); @@ -1109,14 +1170,15 @@ struct Scratch<'gpu, 'trial> { impl<'gpu, 'trial> Scratch<'gpu, 'trial> { fn new( - class: &TuneClass, + device_local: &[bool], sizes: &[usize], - output_index: usize, + output: (usize, usize), bytes: usize, staging: &'trial mut Staging<'gpu>, preparation: &mut TunePreparationTimes, cleanup: &'trial mut Option, ) -> Self { + let (output_index, output_elements) = output; let gpu = staging.gpu; assert!( staging.buffer.is_none() @@ -1138,15 +1200,10 @@ impl<'gpu, 'trial> Scratch<'gpu, 'trial> { .iter() .enumerate() .map(|(i, &size)| { - let device_local = if i > output_index { - true - } else { - class.device_local[if i == output_index { 3 } else { i }] - }; gpu.create_buffer(bg::BufferDesc { name: "tune_scratch", size: size as u64, - memory: if device_local { + memory: if device_local[i] { bg::Memory::DeviceTransient } else { bg::Memory::Shared @@ -1172,7 +1229,7 @@ impl<'gpu, 'trial> Scratch<'gpu, 'trial> { staging, staging_reused, output_index, - output_elements: class.output_elements(), + output_elements, encoder, cleanup, } @@ -1420,7 +1477,10 @@ fn reference_dot(class: &TuneClass, inputs: &[Vec], row: usize, col: usize) // contiguous rows of B rather than forward `[K, N]` columns. let b = if matches!( class.shader, - ShaderEntry::MatMulBT | ShaderEntry::FusedMatMulBTAdd | ShaderEntry::MatMulGemvBT + ShaderEntry::MatMulBT + | ShaderEntry::FusedMatMulBTAdd + | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd ) { col * k + inner } else { @@ -1518,22 +1578,112 @@ fn qualify_output(class: &TuneClass, inputs: &[Vec], output: &[f32], scale: mod tests { use super::*; + #[test] + #[ignore = "GPU qualification of resumable private kernel searches"] + fn kernel_memo_resumes_only_qualified_comparisons() { + let gpu = std::sync::Arc::new( + crate::init_gpu_context_with(crate::GpuOptions::from_env()).unwrap(), + ); + let mut graph = crate::Graph::new(); + let x = graph.input("x", &[33, 17]); + let w = graph.parameter("w", &[17, 65]); + let y = graph.matmul(x, w); + graph.set_outputs(vec![y]); + let plan = crate::compile::compile(&graph); + let create = || { + Session::with_context_opts( + plan.clone(), + gpu.clone(), + crate::SessionOptions { + coop: crate::CoopPolicy::Disabled, + ..Default::default() + }, + ) + }; + let options = TuneOptions { + max_time: Duration::from_secs(30), + sample_pairs: 4, + dispatches_per_sample: 1, + ..Default::default() + }; + let mut memo = KernelMemo::default(); + let first = create() + .tune_with_memo(options.clone(), Some(&mut memo)) + .unwrap(); + assert!(first.outcomes.len() > 1); + assert!(first.outcomes.iter().all(|o| o.qualified)); + assert!(!first.time_budget_exhausted); + assert_eq!(memo.0.len(), 1); + // Keep an actual qualified prefix, without timing-dependent sleeps or + // an assumed GPU speed to interrupt the first search at this point. + *memo.0.values_mut().next().unwrap() = KernelProgress { + selected: first.outcomes[0].selected, + next_candidate: 1, + }; + let mut session = create(); + session.set_input("x", &[0.25; 33 * 17]); + session.set_parameter("w", &[0.125; 17 * 65]); + let skipped = session + .tune_with_memo( + TuneOptions { + max_time: Duration::ZERO, + ..options.clone() + }, + Some(&mut memo), + ) + .unwrap(); + assert!(skipped.outcomes.is_empty() && skipped.reused_classes.is_empty()); + assert_eq!(memo.0.values().next().unwrap().next_candidate, 1); + let resumed = session + .tune_with_memo(options.clone(), Some(&mut memo)) + .unwrap(); + assert_eq!(resumed.reused_classes.len(), 1); + assert_eq!(resumed.outcomes.len(), first.outcomes.len() - 1); + assert!(resumed.outcomes.iter().all(|o| o.qualified)); + assert_eq!( + resumed + .outcomes + .iter() + .map(|o| o.candidate) + .collect::>(), + first.outcomes[1..] + .iter() + .map(|o| o.candidate) + .collect::>() + ); + session.step(); + session.wait(); + assert!( + session + .read_output(33 * 65) + .iter() + .all(|&x| x == 17.0 / 32.0) + ); + assert_eq!(session.read_params(&["w"])[0], [0.125; 17 * 65]); + let complete = create().tune_with_memo(options, Some(&mut memo)).unwrap(); + assert_eq!(complete.reused_classes.len(), 1); + assert!(complete.outcomes.is_empty()); + } + #[test] fn int_dot_candidates_never_fall_back_or_change_arithmetic() { let shape = crate::codegen::GemvShape { threads: 64, reduction: crate::codegen::GemvReduction::Subgroup, + bt_rows: 1, }; let dispatch = Dispatch { shader: ShaderEntry::MatMulGemv, params: vec![1, 256, 16, 0], workgroups: [4, 1, 1], weight_format: crate::compile::WeightFormat::Q40, - gemv_int_dot: true, + kernel: crate::compile::Kernel::Gemv { + shape: crate::codegen::GemvShape::initial(crate::codegen::ShaderGroup::MatMulGemv), + integer_dot: true, + }, ..Default::default() }; - let candidates = Pipelines::candidates(&dispatch); - assert!(matches!(candidates.as_slice(), [Variant::GemvIntDot(..)])); + assert!(matches!(Pipelines::key(&dispatch), Variant::GemvIntDot(..))); assert!(matches!( tile_variant(&dispatch, MatmulTile::Gemv(shape)), Variant::GemvIntDot(_, format, selected) if format == crate::compile::WeightFormat::Q40 && selected == shape @@ -1549,17 +1699,13 @@ mod tests { ); let ordinary = Dispatch { - gemv_int_dot: false, + kernel: crate::compile::Kernel::Default, ..dispatch }; - let candidates = Pipelines::candidates(&ordinary); - assert!(candidates.len() > 1); - assert!(matches!(candidates.last(), Some(Variant::Scalar(_)))); - assert!( - !candidates - .iter() - .any(|v| matches!(v, Variant::GemvIntDot(..))) - ); + assert!(matches!( + Pipelines::key(&ordinary), + Variant::Weight(_, crate::compile::WeightFormat::Q40) + )); } #[test] @@ -1725,7 +1871,12 @@ mod tests { let mut right = left.clone(); for class in &classes { for &index in &class.members { - MatmulTile::Tile32.apply(&mut right.dispatches[index], &class.key); + MatmulTile::Scalar(crate::codegen::ScalarMatmulShape { + tile_size: 32, + k_stage: 8, + interleave_columns: true, + }) + .apply(&mut right.dispatches[index], &class.key); } } let right_classes = collect_classes(&right, &alias, None).0; @@ -1852,10 +2003,12 @@ mod tests { tile_size: 32, k_tile: 16, } - } else if class.initial == MatmulTile::Tile32 { - MatmulTile::Tile64 } else { - MatmulTile::Tile32 + MatmulTile::Scalar(crate::codegen::ScalarMatmulShape { + tile_size: 64, + k_stage: 8, + interleave_columns: true, + }) }; b.pipelines .ensure_tune_tile(&gpu, &b.plan.dispatches[class.members[0]], alternative) @@ -1869,6 +2022,7 @@ mod tests { alternative.apply(&mut b.plan.dispatches[index], &class.key); } b.pipelines.discard_unused_convolutions(&gpu, &b.plan); + b.pipelines.select(&b.plan.dispatches); if convolution { assert!(b.pipelines.map.len() < pipeline_count); } @@ -1933,7 +2087,7 @@ mod tests { let mut selected = left.dispatches[0].clone(); swaps[0].right.apply(&mut selected, &swaps[0].class); assert_eq!(selected, right.dispatches[0]); - assert!(selected.scalar_fallback.is_some()); + assert!(selected.use_coop()); let reduced = crate::codegen::CoopConfig { use_f16_input: true, ..config @@ -2016,9 +2170,9 @@ mod tests { let mut prep = TunePreparationTimes::default(); let mut cleanup = None; let mut scratch = Scratch::new( - &class, + &vec![device_local; sizes.len()], &sizes, - sizes.len() - 1, + (sizes.len() - 1, class.output_elements()), bytes, &mut staging, &mut prep, @@ -2088,9 +2242,9 @@ mod tests { let mut cleanup = None; let result = (|| { let mut scratch = Scratch::new( - &class, + &vec![true; sizes.len()], &sizes, - sizes.len() - 1, + (sizes.len() - 1, class.output_elements()), scratch_bytes(&sizes).unwrap(), staging, &mut prep, @@ -2153,7 +2307,7 @@ mod tests { let mut plan = crate::compile::compile_with(&graph, &crate::CompileOptions::default()); super::super::select_variants(&mut plan, None, false, false); let dispatch = &plan.dispatches[0]; - assert!(dispatch.use_small_tiles); + assert!(dispatch.use_small_tiles()); assert_eq!(dispatch.workgroups, [3, 2, 1]); assert!(TuneClass::from_dispatch(dispatch, None).is_some()); } @@ -2284,6 +2438,7 @@ mod tests { #[test] fn tuning_tiles_generate_valid_exact_binding_layouts() { + use crate::compile::WeightFormat; for entry in [ ShaderEntry::MatMul, ShaderEntry::FusedMatMulAdd, @@ -2295,7 +2450,7 @@ mod tests { ShaderEntry::Conv2dGradInputGemm, ShaderEntry::Conv2dGradWeightGemm, ] { - for tile in [ + let mut tiles = vec![ MatmulTile::Tile32, MatmulTile::Tile64, MatmulTile::CooperativeF32 { tile_size: 8 }, @@ -2308,7 +2463,30 @@ mod tests { tile_size: 64, k_tile: 32, }, - ] { + ]; + for tile_size in [32, 64] { + for k_stage in [8, 16, 32] { + for interleave_columns in [false, true] { + tiles.push(MatmulTile::Scalar(crate::codegen::ScalarMatmulShape { + tile_size, + k_stage, + interleave_columns, + })); + } + } + } + for (tile, format) in tiles.into_iter().flat_map(|tile| { + [WeightFormat::F32, WeightFormat::F16] + .into_iter() + .filter(move |format| { + *format == WeightFormat::F32 + || matches!( + tile, + MatmulTile::Scalar(_) | MatmulTile::Tile32 | MatmulTile::Tile64 + ) + }) + .map(move |format| (tile, format)) + }) { let convolution = matches!( entry, ShaderEntry::Conv2dGemm @@ -2319,11 +2497,16 @@ mod tests { if specialized && !convolution { continue; } - if convolution && tile.coop_config().is_some() { + if convolution + && (tile.coop_config().is_some() + || matches!(tile, MatmulTile::Scalar(_)) + || format != WeightFormat::F32) + { continue; } let dispatch = Dispatch { shader: entry.clone(), + weight_format: format, params: vec![2, 3, 7, 9, 5, 3, 2, 2, 0, 3, 5, 1], ..Default::default() }; diff --git a/src/shaders/matmul_coop.wgsl b/src/shaders/matmul_coop.wgsl index 0095b44d..ad1fd2c7 100644 --- a/src/shaders/matmul_coop.wgsl +++ b/src/shaders/matmul_coop.wgsl @@ -30,7 +30,7 @@ $RESULT_SHARED_DECL $PROLOGUE_CACHE_DECL @compute @workgroup_size(64) -fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) { +fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3, @builtin(subgroup_id) sg: u32) { let tile_row = wgid.x * $OUTPUT_TILE_U; let tile_col = wgid.y * $OUTPUT_TILE_U; let m = params.m; diff --git a/src/shaders/matmul_gemv_bt.wgsl b/src/shaders/matmul_gemv_bt.wgsl index 355855d8..626d8608 100644 --- a/src/shaders/matmul_gemv_bt.wgsl +++ b/src/shaders/matmul_gemv_bt.wgsl @@ -37,8 +37,8 @@ const LANES: u32 = 32u; var reduce_buf: array; @compute @workgroup_size(LANES) -fn main(@builtin(workgroup_id) wgid: vec3, @builtin(local_invocation_id) lid: vec3) { - let col = wgid.x; +fn main(@builtin(workgroup_id) wgid: vec3, @builtin(num_workgroups) grid: vec3, @builtin(local_invocation_id) lid: vec3) { + let col = wgid.x + grid.x * wgid.y; let lane = lid.x; if col >= params.n { return; } let k_v4 = params.k / 4u; diff --git a/src/train.rs b/src/train.rs index 3febec1b..9076cf20 100644 --- a/src/train.rs +++ b/src/train.rs @@ -7,6 +7,9 @@ use crate::{ }; use std::{path::Path, sync::Arc}; +mod search; +pub use search::{BuildSearchOptions, BuildSearchReport, BuildSearchTrial, build_measured}; + /// Optimizer selection. #[derive(Clone, Debug)] pub enum Optimizer { @@ -349,21 +352,51 @@ pub fn build(forward_graph: &Graph, cfg: SessionConfig<'_>) -> (Session, optimiz } } + let (final_graph, report) = + prepare_graph(forward_graph, mode, cfg.optimize, skip_full_optimize); + let plan = { + let _span = tracing::info_span!("compile").entered(); + compile::compile_owned_with_caps(final_graph, &options, coop_caps) + }; + log::info!( + "execution plan: {} buffers, {} dispatches", + plan.buffers.len(), + plan.dispatches.len() + ); + + if let Some(path) = cache_path { + if let Err(e) = cache::save_build_plan(&plan, forward_graph, build_hash, path) { + log::warn!("failed to save cache: {}", e); + } else { + log::info!("saved execution plan cache to {}", path.display()); + } + } + + let session = make_session(plan, gpu, cfg.runtime.clone(), cfg.tune); + (session, report) +} + +fn prepare_graph( + forward_graph: &Graph, + mode: Mode, + optimize: optimize::OptimizeConfig, + skip_full_optimize: bool, +) -> (Graph, optimize::OptimizeReport) { let (optimized_forward, forward_report) = { let _span = tracing::info_span!("optimize_forward").entered(); - optimize::optimize_with_config(forward_graph, cfg.optimize) + optimize::optimize_with_config(forward_graph, optimize) }; log::info!( "optimized forward: {} nodes", optimized_forward.nodes().len() ); - let (final_graph, report) = match mode { + match mode { Mode::Inference => { let mut g = optimized_forward; let mut fusions = Vec::new(); optimize::apply_group_norm_silu_fusions(&mut g, &mut fusions); - optimize::apply_winograd_conv_fusions(&mut g, &mut fusions, &cfg.optimize); + optimize::apply_winograd_conv_fusions(&mut g, &mut fusions, &optimize); for (name, count) in fusions.iter().fold( std::collections::BTreeMap::<&str, usize>::new(), |mut acc, entry| { @@ -394,31 +427,10 @@ pub fn build(forward_graph: &Graph, cfg: SessionConfig<'_>) -> (Session, optimiz (full, forward_report) } else { let _span = tracing::info_span!("optimize_full").entered(); - optimize::optimize_owned_with_config(full, cfg.optimize) + optimize::optimize_owned_with_config(full, optimize) } } - }; - - let plan = { - let _span = tracing::info_span!("compile").entered(); - compile::compile_owned_with_caps(final_graph, &options, coop_caps) - }; - log::info!( - "execution plan: {} buffers, {} dispatches", - plan.buffers.len(), - plan.dispatches.len() - ); - - if let Some(path) = cache_path { - if let Err(e) = cache::save_build_plan(&plan, forward_graph, build_hash, path) { - log::warn!("failed to save cache: {}", e); - } else { - log::info!("saved execution plan cache to {}", path.display()); - } } - - let session = make_session(plan, gpu, cfg.runtime.clone(), cfg.tune); - (session, report) } fn make_session( diff --git a/src/train/search.rs b/src/train/search.rs new file mode 100644 index 00000000..67096ba6 --- /dev/null +++ b/src/train/search.rs @@ -0,0 +1,382 @@ +//! Measured construction: graph alternatives survive until their kernels are tuned. +use super::{SessionConfig, prepare_graph}; +use crate::{ + Graph, Session, TuneOptions, + compile::{self, ExecutionPlan}, + optimize, runtime, +}; +use serde::Serialize; + +mod measure; +pub use measure::BuildSearchTrial; +use std::time::{Duration, Instant}; + +#[derive(Clone, Serialize)] +pub struct BuildSearchOptions { + /// Bound on logical forms, including ordinary extraction. + pub max_graphs: usize, + /// Per-program kernel tuning and paired whole-step decision policy. + pub tuning: TuneOptions, + /// Whole-program warmup pairs, independent of private kernel warmup. + pub warmup_runs: u32, + /// Soft total deadline, including construction, initialization and validation. + /// In-flight driver work and caller validation cannot be preempted. + pub max_time: Duration, + pub max_programs: usize, + /// Sum of declared logical bytes and persistent-state snapshots for both + /// incumbent and challenger. This is not a driver-heap bound: padding, + /// pipelines, staging and kernel-probe scratch are additional. + pub max_plan_bytes: usize, +} + +impl Default for BuildSearchOptions { + fn default() -> Self { + Self { + max_graphs: 4, + tuning: TuneOptions::default(), + warmup_runs: 2, + max_time: Duration::from_secs(30), + max_programs: 64, + max_plan_bytes: 512 << 20, + } + } +} + +#[derive(Default, Serialize)] +pub struct BuildSearchReport { + pub options: BuildSearchOptions, + /// Extracted forms, indexed by the `graph=N` trial descriptions. + pub graphs: Vec, + pub extraction_truncated: bool, + pub skipped_regions: Vec, + pub preparation_time: Duration, + pub selected: usize, + pub trials: Vec, + pub truncated: bool, + pub elapsed: Duration, +} + +/// Build from representative inputs, measuring logical and physical alternatives. +/// +/// Unlike [`super::build`], this executes private candidate sessions. +/// `initialize` writes representative inputs and weights once per candidate. Its +/// optional idle incumbent can donate identically represented immutable parameters +/// via [`Session::share_parameter_from`], but must not be modified. Inputs and +/// writable state must remain private. Configure runtime optimizers, accumulation +/// and external bindings after construction, not inside either callback. +/// +/// The runner executes one step before each read-only `qualify` call. Check every +/// observable output, gradient and persistent update against your numerical +/// contract. Qualification runs before/after kernel tuning and after measurement; +/// a failing challenger is discarded, a failing incumbent aborts construction. +/// Written inputs, parameters and constants are reset before each step, outside +/// timing. The returned session retains its initialized persistent state; outputs +/// hold its last qualified step. Timing includes fresh recording/submission/wait, +/// not readback. These search samples are not held-out benchmark results. +/// +/// The ordinary optimizer supplies the first candidate. Egglog retains alternatives +/// from the original forward graph (one bounded region, repeated where verified). +/// They are not greedily optimized again. Each training form is differentiated +/// separately, so parameter transformations and gradients stay consistent. +/// Complete plans explore dispatch fusion and cached-attention splits before +/// allocation, plus submission chunk counts. Each plan is kernel-tuned before +/// comparison. This bounded search does not promise a global optimum. +/// +/// `options.tuning` replaces `cfg.tune`. Build-plan caching is not yet supported: +/// calibration data and measured policy are not part of the ordinary cache key. +pub fn build_measured( + forward_graph: &Graph, + mut cfg: SessionConfig<'_>, + options: BuildSearchOptions, + initialize: impl FnMut(&mut Session, Option<&mut Session>) -> Result<(), String>, + qualify: impl FnMut(&Session) -> Result<(), String>, +) -> Result<(Session, BuildSearchReport), String> { + let start = Instant::now(); + let _span = tracing::info_span!("build_measured").entered(); + if cfg.cache.is_some() { + return Err("measured construction does not use the ordinary build-plan cache".into()); + } + if options.max_graphs == 0 || options.max_programs == 0 || options.max_time.is_zero() { + return Err("measured construction needs positive graph, program and time bounds".into()); + } + options.tuning.validate().map_err(|e| e.to_string())?; + if cfg.runtime.debug { + cfg.options.fuse_dispatches = false; + } + let gpu = cfg.gpu.take().unwrap_or_else(runtime::default_gpu_context); + let caps = cfg + .runtime + .coop + .filter_caps(runtime::auto_tune(&gpu, 0).coop_caps); + let (ordinary, _) = prepare_graph( + forward_graph, + cfg.mode, + cfg.optimize, + cfg.skip_full_optimize, + ); + let mut graphs = vec![optimize::search::Candidate { + graph: ordinary, + expression: "ordinary extraction".into(), + }]; + let mut extraction_truncated = false; + let mut skipped_regions = Vec::new(); + if cfg.optimize.mode != optimize::OptimizeMode::Off + && options.max_graphs > 1 + && start.elapsed() < options.max_time + { + let source = forward_graph.toposort(); + let limit = options.max_graphs - 1; + let spaces = if source.nodes().len() + <= cfg + .optimize + .saturation_cutoff + .min(optimize::SATURATION_CUTOFF) + { + vec![optimize::search::candidates(&source, cfg.optimize, limit)] + } else { + // Reuse outlining, not model names or a second pattern matcher. + let regions = crate::outline::detect_repeated_regions(&source); + extraction_truncated |= regions.len() > 1; + if regions.is_empty() { + skipped_regions.push("large graph has no bounded repeated region".into()); + } + regions + .into_iter() + .take(1) + .map(|region| { + optimize::search::repeated_candidates(&source, region, cfg.optimize, limit) + }) + .collect() + }; + for space in spaces { + match space { + Ok(space) => { + extraction_truncated |= space.truncated; + let unfused = optimize::OptimizeConfig { + mode: optimize::OptimizeMode::Off, + ..cfg.optimize + }; + for form in space.candidates { + if start.elapsed() >= options.max_time { + extraction_truncated = true; + break; + } + graphs.push(optimize::search::Candidate { + graph: prepare_graph( + &form.graph, + cfg.mode, + unfused, + cfg.skip_full_optimize, + ) + .0, + expression: form.expression, + }); + } + } + Err(error) => skipped_regions.push(error), + } + } + } + let expressions = graphs.iter().map(|form| form.expression.clone()).collect(); + // Compile graph forms only once. Physical alternatives remain lazy, and are + // interleaved across forms rather than spending the budget on the first form. + let mut seeds: Vec = Vec::new(); + for (index, form) in graphs.into_iter().enumerate() { + if start.elapsed() >= options.max_time { + extraction_truncated = true; + break; + } + let graph = std::sync::Arc::new(form.graph); + for &fusion in if cfg.options.fuse_dispatches { + &[true, false][..] + } else { + &[false][..] + } { + let options = compile::CompileOptions { + fuse_dispatches: fusion, + ..cfg.options.clone() + }; + let plan = compile::compile_with_caps(&graph, &options, caps); + if !seeds.iter().any(|seed| seed.plan == plan) { + seeds.push(Seed { + description: format!("graph={index}, dispatch_fusion={fusion}"), + plan, + graph: graph.clone(), + options, + }); + } + } + } + let preparation_time = start.elapsed(); + let programs = implementations(seeds, caps); + measure::select( + programs, + gpu, + cfg.runtime, + BuildSearchReport { + options, + graphs: expressions, + extraction_truncated, + skipped_regions, + preparation_time, + ..Default::default() + }, + start, + initialize, + qualify, + ) +} + +struct Seed { + graph: std::sync::Arc, + options: compile::CompileOptions, + plan: ExecutionPlan, + description: String, +} + +fn implementations( + seeds: Vec, + caps: crate::codegen::CoopCaps, +) -> impl Iterator { + let attention = seeds.iter().any(|p| { + p.graph + .nodes() + .iter() + .any(|n| matches!(n.op, crate::graph::Op::CachedBlockAttention { .. })) + }); + let splits: &[u32] = if attention { + &[0, 1, 2, 4, 8, 16] + } else { + &[0] + }; + let chunks = [1, 2, 4, 8, 16, 32, 64]; + // Interleave settings across logical forms, starting near ordinary lowering. + // Lower only the next candidate; do not allocate a Cartesian product of plans. + let mut choices = (0..splits.len() + chunks.len() - 1).flat_map(move |rank| { + (0..splits.len()).filter_map(move |a| { + let c = rank.checked_sub(a)?; + chunks.get(c).map(|&chunks| (splits[a], chunks)) + }) + }); + let mut current = (0, 1); + let mut seed_index = seeds.len(); + std::iter::from_fn(move || { + loop { + if seed_index == seeds.len() { + current = choices.next()?; + seed_index = 0; + } + let seed = seeds.get(seed_index)?; + seed_index += 1; + let (splits, chunks) = current; + let plan = if splits == 0 { + seed.plan.clone() + } else { + let options = compile::CompileOptions { + cached_attention_splits: Some(splits), + ..seed.options.clone() + }; + let plan = compile::compile_with_caps(&seed.graph, &options, caps); + if plan == seed.plan { + continue; + } + plan + }; + return Some(measure::Program { + description: format!( + "{}, attention_splits={splits}, submission_chunks={chunks}", + seed.description + ), + plan, + submission_chunks: chunks, + }); + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn measured_build_keeps_graph_choices_and_training_state() { + use crate::{Mode, TuneOptions}; + let gpu = std::sync::Arc::new( + crate::init_gpu_context_with(crate::GpuOptions::from_env()).unwrap(), + ); + let mut graph = Graph::new(); + let x = graph.input("x", &[3, 33]); + let w = graph.parameter("w", &[33, 5]); + let b = graph.input("b", &[3, 5]); + let y = graph.matmul(x, w); + let y = graph.add(y, b); + let loss = graph.mean_all(y); + graph.set_outputs(vec![loss]); + for mode in [Mode::Inference, Mode::Training] { + let (mut session, report) = build_measured( + &graph, + SessionConfig { + mode, + gpu: Some(gpu.clone()), + runtime: crate::SessionOptions { + coop: crate::CoopPolicy::Disabled, + ..Default::default() + }, + ..Default::default() + }, + BuildSearchOptions { + max_graphs: 4, + tuning: TuneOptions { + max_time: Duration::from_secs(1), + sample_pairs: 4, + warmup_runs: 1, + dispatches_per_sample: 1, + ..Default::default() + }, + warmup_runs: 1, + max_time: Duration::from_secs(60), + max_programs: 24, + max_plan_bytes: 4 << 20, + }, + |s, _| { + s.set_input("x", &[0.25; 99]); + s.set_input("b", &[0.0625; 15]); + s.set_parameter("w", &[0.125; 165]); + Ok(()) + }, + |s| { + if (s.read_output(1)[0] - 1.09375).abs() > 2e-6 { + return Err("forward result changed".into()); + } + assert_eq!(s.read_params(&["w"])[0], [0.125; 165]); + if mode == Mode::Training { + let mut grad = [0.0; 165]; + s.read_param_grad("w", &mut grad); + if grad.iter().any(|g| (g - 0.05).abs() > 2e-6) { + return Err("parameter gradient changed".into()); + } + } + Ok(()) + }, + ) + .unwrap(); + assert!(report.graphs.len() >= 3); + assert!(!report.extraction_truncated); + assert!(report.trials.len() <= 24); + assert!(report.skipped_regions.is_empty()); + assert!(report.trials.len() > report.graphs.len()); + assert!(report.trials.iter().all(|t| t.outcome.qualified)); + assert_eq!(session.read_params(&["w"])[0], [0.125; 165]); + if mode == Mode::Training { + session.set_learning_rate(0.1); + session.step(); + session.wait(); + assert!( + session.read_params(&["w"])[0] + .iter() + .all(|w| (w - 0.12).abs() < 2e-6) + ); + } + } + } +} diff --git a/src/train/search/measure.rs b/src/train/search/measure.rs new file mode 100644 index 00000000..c223099d --- /dev/null +++ b/src/train/search/measure.rs @@ -0,0 +1,550 @@ +use super::BuildSearchReport; +use crate::{ + Session, + compile::ExecutionPlan, + runtime::{ + SessionOptions, + search_state::{SearchState, persistent_writes}, + }, + tune::{TuneDecision, TuneOutcome, TuneReport, decide, measure_pairs}, +}; +use serde::Serialize; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +pub(super) struct Program { + pub description: String, + pub plan: ExecutionPlan, + pub submission_chunks: usize, +} + +#[derive(Serialize)] +pub struct BuildSearchTrial { + pub description: String, + /// CPU generation of this plan, when the input iterator lowers lazily. + pub lowering_time: Duration, + /// Full session construction, including allocations, not shader-only compile time. + pub construction_time: Duration, + pub initialization_time: Duration, + /// State capture and restoration, excluded from execution samples. + pub state_copy_time: Duration, + pub snapshot_bytes: usize, + pub qualification_time: Duration, + pub kernel_tuning: Option, + pub outcome: TuneOutcome<(), usize>, +} + +fn plan_bytes(plan: &ExecutionPlan) -> Result { + plan.buffers + .iter() + .chain( + persistent_writes(plan) + .iter() + .map(|b| &plan.buffers[b.0 as usize]), + ) + .try_fold(0usize, |sum, bytes| sum.checked_add((*bytes).max(4))) + .ok_or_else(|| "declared program bytes overflow".into()) +} + +fn restore( + session: &mut Session, + state: &SearchState, + elapsed: &mut Duration, +) -> Result<(), String> { + let start = Instant::now(); + let result = state.restore(session); + *elapsed += start.elapsed(); + result +} + +fn run(session: &mut Session, state: &SearchState, elapsed: &mut Duration) -> Result { + restore(session, state, elapsed)?; + let start = Instant::now(); + session.step(); + session.wait(); + Ok(start.elapsed().as_secs_f64() * 1000.0) +} + +fn validate( + session: &mut Session, + state: &SearchState, + trial: &mut BuildSearchTrial, + check: &mut impl FnMut(&Session) -> Result<(), String>, +) -> Result<(), String> { + restore(session, state, &mut trial.state_copy_time)?; + let start = Instant::now(); + session.step(); + session.wait(); + let result = check(session); + trial.qualification_time += start.elapsed(); + restore(session, state, &mut trial.state_copy_time)?; + result +} + +pub(super) fn select( + programs: impl IntoIterator, + gpu: Arc, + runtime: SessionOptions, + mut report: BuildSearchReport, + start: Instant, + mut initialize: impl FnMut(&mut Session, Option<&mut Session>) -> Result<(), String>, + mut qualify: impl FnMut(&Session) -> Result<(), String>, +) -> Result<(Session, BuildSearchReport), String> { + let options = report.options.clone(); + options + .tuning + .validate() + .map_err(|error| error.to_string())?; + if options.max_programs == 0 || options.max_time.is_zero() { + return Err("program search needs a positive program and time budget".into()); + } + let mut incumbent: Option<(Session, SearchState)> = None; + let mut incumbent_bytes = 0usize; + let mut kernels = crate::runtime::KernelMemo::default(); + let mut programs = programs.into_iter(); + for index in 0..options.max_programs { + if start.elapsed() >= options.max_time { + report.truncated = true; + break; + } + let lowering = Instant::now(); + let Some(program) = programs.next() else { + break; + }; + let lowering_time = lowering.elapsed(); + if start.elapsed() >= options.max_time { + report.truncated = true; + break; + } + let trial_start = Instant::now(); + let bytes = plan_bytes(&program.plan)?; + let mut trial = BuildSearchTrial { + description: program.description, + lowering_time, + construction_time: Duration::ZERO, + initialization_time: Duration::ZERO, + state_copy_time: Duration::ZERO, + snapshot_bytes: 0, + qualification_time: Duration::ZERO, + kernel_tuning: None, + outcome: TuneOutcome::new((), program.plan.dispatches.len(), report.selected, index), + }; + trial.outcome.phase_times = None; + if bytes + .checked_add(incumbent_bytes) + .is_none_or(|sum| sum > options.max_plan_bytes) + { + trial.outcome.decision = TuneDecision::ScratchLimit; + } else { + let build = Instant::now(); + let mut candidate = + Session::with_context_opts(program.plan, gpu.clone(), runtime.clone()); + candidate.set_submission_chunks(program.submission_chunks); + trial.construction_time = build.elapsed(); + let result = (|| { + let init = Instant::now(); + let initialized = + initialize(&mut candidate, incumbent.as_mut().map(|entry| &mut entry.0)); + trial.initialization_time = init.elapsed(); + initialized?; + let copy_start = Instant::now(); + let state = SearchState::capture(&mut candidate, options.max_plan_bytes)?; + trial.state_copy_time += copy_start.elapsed(); + trial.snapshot_bytes = state.bytes(); + validate(&mut candidate, &state, &mut trial, &mut qualify)?; + let mut policy = options.tuning.clone(); + policy.max_time = policy + .max_time + .min(options.max_time.saturating_sub(start.elapsed())); + trial.kernel_tuning = Some( + candidate + .tune_with_memo(policy, Some(&mut kernels)) + .map_err(|error| error.to_string())?, + ); + validate(&mut candidate, &state, &mut trial, &mut qualify)?; + trial.outcome.qualified = true; + if let Some((ref mut baseline, ref baseline_state)) = incumbent { + for _ in 0..options.warmup_runs { + if start.elapsed() >= options.max_time { + break; + } + run(baseline, baseline_state, &mut trial.state_copy_time)?; + run(&mut candidate, &state, &mut trial.state_copy_time)?; + } + let mut failure = None; + (trial.outcome.baseline_ms, trial.outcome.candidate_ms) = + measure_pairs(options.tuning.sample_pairs, |alternative| { + if start.elapsed() >= options.max_time { + return None; + } + let result = if alternative { + run(&mut candidate, &state, &mut trial.state_copy_time) + } else { + run(baseline, baseline_state, &mut trial.state_copy_time) + }; + match result { + Ok(ms) => Some(ms), + Err(error) => { + failure = Some(error); + None + } + } + }); + restore(baseline, baseline_state, &mut trial.state_copy_time)?; + if let Some(error) = failure { + return Err(error); + } + validate(&mut candidate, &state, &mut trial, &mut qualify)?; + decide(&mut trial.outcome, &options.tuning); + } else { + trial.outcome.selected = index; + } + Ok::<_, String>(state) + })(); + // A failing challenger is discarded. A failing incumbent invalidates + // the search: it must not survive as the supposedly safe fallback. + if let Some((ref mut baseline, ref state)) = incumbent { + validate(baseline, state, &mut trial, &mut qualify) + .map_err(|error| format!("incumbent failed repeated qualification: {error}"))?; + } + match result { + Ok(state) if incumbent.is_none() || trial.outcome.selected == index => { + incumbent_bytes = plan_bytes(candidate.plan())?; + incumbent = Some((candidate, state)); + report.selected = index; + } + Ok(_) => {} + Err(error) => { + if incumbent.is_none() { + return Err(format!("initial program failed qualification: {error}")); + } + if let Some((ref mut baseline, ref state)) = incumbent { + restore(baseline, state, &mut trial.state_copy_time)?; + } + log::warn!("program {index} qualification failed: {error}"); + trial.outcome.qualified = false; + trial.outcome.selected = report.selected; + trial.outcome.decision = TuneDecision::InvalidOutput; + trial.outcome.failure = Some(error); + } + } + } + trial.outcome.elapsed = trial_start.elapsed(); + log::info!( + "program {index}: {:?}, selected={}, {:?}/{:?} ms", + trial.outcome.decision, + trial.outcome.selected, + trial.outcome.baseline_median_ms, + trial.outcome.candidate_median_ms + ); + report.trials.push(trial); + } + // Do not lower one more program just to discover whether a bounded search + // is truncated. Lazy producers can allocate and compile in `next()`. + report.truncated |= start.elapsed() >= options.max_time + || (report.trials.len() == options.max_programs && programs.size_hint().1 != Some(0)); + incumbent + .map(|(mut session, state)| { + state.restore(&mut session)?; + report.elapsed = start.elapsed(); + Ok((session, report)) + }) + .ok_or_else(|| "no qualified program within the search bounds".into()) + .and_then(|result| result) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::TuneOptions; + use crate::train::BuildSearchOptions; + + fn select( + programs: impl IntoIterator, + gpu: Arc, + runtime: SessionOptions, + options: BuildSearchOptions, + initialize: impl FnMut(&mut Session, Option<&mut Session>) -> Result<(), String>, + qualify: impl FnMut(&Session) -> Result<(), String>, + ) -> Result<(Session, BuildSearchReport), String> { + super::select( + programs, + gpu, + runtime, + BuildSearchReport { + options, + ..Default::default() + }, + Instant::now(), + initialize, + qualify, + ) + } + #[test] + fn charges_persistent_state_before_allocation() { + use crate::compile::{BufferRef, Dispatch}; + let mut graph = crate::Graph::new(); + let input = graph.input("x", &[2]); + let output = graph.neg(input); + graph.set_outputs(vec![output]); + let mut plan = crate::compile::compile(&graph); + let bytes = super::plan_bytes(&plan).unwrap(); + plan.param_grad_pairs.push((BufferRef(0), BufferRef(1))); + assert_eq!(super::plan_bytes(&plan).unwrap(), bytes); + plan.param_grad_pairs.clear(); + plan.dispatches.push(Dispatch { + output_buffer: plan.input_buffers[0].1, + ..Default::default() + }); + assert_eq!(super::plan_bytes(&plan).unwrap(), bytes + 8); + } + + #[test] + fn program_search_preserves_update_and_failure_contracts() { + use crate::compile::{Dispatch, ShaderEntry}; + let gpu = Arc::new(crate::init_gpu_context_with(crate::GpuOptions::from_env()).unwrap()); + let mut graph = crate::Graph::new(); + let x = graph.input("x", &[2]); + let w = graph.parameter("w", &[2]); + let y = graph.mul(x, w); + let loss = graph.sum_all(y); + graph.set_outputs(vec![loss]); + let mut plan = crate::compile::compile(&crate::autodiff::differentiate(&graph)); + let (w, grad) = plan.param_grad_pairs[0]; + plan.dispatches.push(Dispatch { + shader: ShaderEntry::SgdUpdate, + input_buffers: vec![w, grad], + output_buffer: w, + workgroups: [1, 1, 1], + params: vec![2, 0.25f32.to_bits(), 0, 0], + ..Default::default() + }); + for invalidate in [false, true] { + let initialized = std::cell::Cell::new(0); + let programs = ["baseline", "candidate"].map(|name| { + let mut plan = plan.clone(); + plan.dispatches[0].label = name.into(); + Program { + description: name.into(), + plan, + submission_chunks: 1, + } + }); + let result = select( + programs, + gpu.clone(), + SessionOptions::default(), + BuildSearchOptions { + tuning: TuneOptions { + max_time: Duration::ZERO, + sample_pairs: 4, + ..Default::default() + }, + warmup_runs: 2, + max_time: Duration::from_secs(30), + max_programs: 2, + max_plan_bytes: 1 << 20, + ..Default::default() + }, + |s, _| { + initialized.set(initialized.get() + 1); + s.set_input("x", &[2.0, 4.0]); + s.set_parameter("w", &[3.0, 5.0]); + Ok(()) + }, + |s| { + if invalidate + && initialized.get() == 2 + && s.plan().dispatches.iter().any(|d| d.label == "baseline") + { + return Err("injected incumbent qualification failure".into()); + } + assert_eq!(s.read_loss(), 26.0); + let mut gradient = [0.0; 2]; + s.read_param_grad("w", &mut gradient); + assert_eq!(gradient, [2.0, 4.0]); + assert_eq!(s.read_params(&["w"])[0], [2.5, 4.0]); + Ok(()) + }, + ); + if invalidate { + assert!( + result + .err() + .unwrap() + .contains("incumbent failed repeated qualification") + ); + } else { + let (mut selected, report) = result.unwrap(); + assert!(report.trials.iter().all(|t| t.outcome.qualified)); + assert_eq!(selected.read_params(&["w"])[0], [3.0, 5.0]); + selected.step(); + selected.wait(); + assert_eq!(selected.read_params(&["w"])[0], [2.5, 4.0]); + } + } + } + + #[test] + fn cached_attention_search_restores_state() { + let gpu = Arc::new(crate::init_gpu_context_with(crate::GpuOptions::from_env()).unwrap()); + for (capacity, block, position, window, valid, dim) in + [(7, 3, 0, 3, 2, 4), (96, 1, 95, 37, 1, 4)] + { + let mut graph = crate::Graph::new(); + let q = graph.input("q", &[block, 2 * dim]); + let k = graph.parameter("k", &[capacity, dim]); + let v = graph.parameter("v", &[capacity, dim]); + let pos = graph.input_u32("position", &[1]); + let len = graph.input_u32("valid", &[1]); + // At position zero a missing reset doubles K and halves V again. + let new_k = graph.scale(k, 2.0); + let new_v = graph.scale(v, 0.5); + let k = graph.cache_write_prefix(new_k, k, pos, len); + let v = graph.cache_write_prefix(new_v, v, pos, len); + let output = graph.cached_block_attention(q, k, v, pos, len, 2, 1, dim as u32, window); + graph.set_outputs(vec![output]); + // Kernel numerics (unequal scores, ragged windows, wide heads) are + // covered by gemma_inference_ops. Here uniform attention isolates + // the search runner's mutation/reset and rejection contract. + let inputs = vec![0.0; block * 2 * dim]; + let original: Vec> = [0.2, 0.7] + .iter() + .map(|factor| { + (0..capacity * dim) + .map(|i| (i as f32 * factor).cos()) + .collect() + }) + .collect(); + let mut changed = original.clone(); + for (i, &factor) in [2.0, 0.5].iter().enumerate() { + for row in 0..valid { + for col in 0..dim { + changed[i][(position + row) * dim + col] = + original[i][row * dim + col] * factor; + } + } + } + let mut expected = Vec::new(); + for row in 0..valid { + let end = position + row + 1; + let begin = if window == 0 { + 0 + } else { + end.saturating_sub(window as usize) + }; + for col in 0..2 * dim { + expected.push( + (begin..end) + .map(|token| f64::from(changed[1][token * dim + col % dim])) + .sum::() + / (end - begin) as f64, + ); + } + } + let programs = [1, 1, 2, 4, 8, 16, 1].into_iter().map(|splits| { + let plan = crate::compile::compile_with( + &graph, + &crate::compile::CompileOptions { + cached_attention_splits: Some(splits), + ..Default::default() + }, + ); + Program { + description: splits.to_string(), + plan, + submission_chunks: splits as usize, + } + }); + let mut index = 0; + let (mut selected, report) = select( + programs, + gpu.clone(), + SessionOptions { + coop: crate::CoopPolicy::Disabled, + ..Default::default() + }, + BuildSearchOptions { + tuning: TuneOptions { + sample_pairs: 4, + max_time: Duration::from_secs(1), + ..Default::default() + }, + warmup_runs: 2, + max_time: Duration::from_secs(60), + max_programs: 7, + max_plan_bytes: 1 << 20, + ..Default::default() + }, + |s, incumbent| { + if let Some(ref baseline) = incumbent { + assert_eq!(baseline.read_params(&["k", "v"]), original); + } + s.set_input("q", &inputs); + s.set_input_u32("position", &[position as u32]); + s.set_input_u32("valid", &[valid as u32]); + s.set_parameter("k", &original[0]); + s.set_parameter("v", &original[1]); + if index == 1 { + s.set_input("q", &vec![f32::NAN; inputs.len()]); + } else if index == 6 { + s.share_parameter_from(incumbent.unwrap(), "k").unwrap(); + } + index += 1; + Ok(()) + }, + |s| { + let actual = s.read_output(valid * 2 * dim); + assert_eq!(s.read_params(&["k", "v"]), changed); + if actual + .iter() + .zip(&expected) + .all(|(&a, &b)| a.is_finite() && (f64::from(a) - b).abs() < 2e-5) + { + Ok(()) + } else { + Err("independent complete-output attention reference mismatch".into()) + } + }, + ) + .unwrap(); + assert_eq!(report.trials.len(), 7); + assert!(!report.truncated); + assert!(report.trials[0].outcome.qualified); + assert!(!report.trials[1].outcome.qualified); + assert!( + report.trials[2..6] + .iter() + .all(|t| t.outcome.qualified && t.snapshot_bytes == capacity * dim * 4 * 2) + ); + assert!(!report.trials[6].outcome.qualified); + assert!( + report.trials[6] + .outcome + .failure + .as_ref() + .unwrap() + .contains("private writable") + ); + assert_eq!(selected.read_params(&["k", "v"]), original); + selected.step(); + selected.wait(); + assert_eq!(selected.read_params(&["k", "v"]), changed); + assert!( + SearchState::capture(&mut selected, 0) + .err() + .unwrap() + .contains("byte budget") + ); + selected.set_learning_rate(0.1); + assert!( + SearchState::capture(&mut selected, 1 << 20) + .err() + .unwrap() + .contains("optimizers") + ); + } + } +} diff --git a/src/tune.rs b/src/tune.rs index 55bd99e8..f2b12d9a 100644 --- a/src/tune.rs +++ b/src/tune.rs @@ -2,8 +2,8 @@ //! //! The search space is deliberately small: scalar and native-f32 cooperative //! tiles for unpacked dense matmuls, shape-specialized scalar convolutions, -//! and workgroup width and cross-lane reduction for the K-split GEMV family, -//! with no precision or binding-layout changes. +//! workgroup geometry and cross-lane reduction for the K-split GEMV family, and +//! cached-attention split/combine sequences. Precision remains unchanged. //! Measurements use synthetic, private scratch, not a live training step. //! Explicit split-K probes measure complete sequences without installing them. //! Scalar tile and GEMV shape changes leave packed-weight decoding intact, so @@ -172,6 +172,7 @@ pub(crate) fn gemv_group(entry: &ShaderEntry) -> Option Some(ShaderGroup::MatMulGemv), ShaderEntry::MatMulGemvAdd => Some(ShaderGroup::MatMulGemvAdd), ShaderEntry::MatMulGemvBT => Some(ShaderGroup::MatMulGemvBT), + ShaderEntry::MatMulGemvBTAdd => Some(ShaderGroup::MatMulGemvBTAdd), _ => None, } } @@ -181,6 +182,7 @@ pub(crate) fn gemv_group(entry: &ShaderEntry) -> Option) -> Option { - let small = dispatch.use_small_tiles + use crate::compile::Kernel; + let small = dispatch.use_small_tiles() || matches!( dispatch.shader, ShaderEntry::Conv2dGemmSmall @@ -220,20 +222,24 @@ impl MatmulTile { | ShaderEntry::Conv2dGradWeightGemmSmall ); if let Some(group) = gemv_group(&dispatch.shader) { - Some(Self::Gemv(dispatch.gemv_shape.unwrap_or_else(|| { - crate::codegen::GemvShape::initial(group) - }))) - } else if let Some(k_tile) = dispatch.conv_k_tile { - Some(Self::SpecializedConv { - tile_size: if small { 32 } else { 64 }, - k_tile, - }) - } else if dispatch.use_coop { - Self::native_cooperative(config) - } else if small { - Some(Self::Tile32) + match dispatch.kernel { + Kernel::Default => Some(Self::Gemv(crate::codegen::GemvShape::initial(group))), + Kernel::Gemv { shape, .. } => Some(Self::Gemv(shape)), + _ => None, + } } else { - Some(Self::Tile64) + match dispatch.kernel { + Kernel::Default | Kernel::SmallTile => { + Some(if small { Self::Tile32 } else { Self::Tile64 }) + } + Kernel::ScalarMatmul(shape) => Some(Self::Scalar(shape)), + Kernel::SpecializedConv { k_tile } => Some(Self::SpecializedConv { + tile_size: if small { 32 } else { 64 }, + k_tile, + }), + Kernel::Cooperative => Self::native_cooperative(config), + _ => None, + } } } @@ -249,25 +255,38 @@ impl MatmulTile { } pub(crate) fn apply(self, dispatch: &mut Dispatch, class: &TuneClass) { + self.configure(dispatch); + dispatch.workgroups = self.workgroups(class); + } + + pub(crate) fn configure(self, dispatch: &mut Dispatch) { if let Self::Gemv(shape) = self { - // Width and reduction live inside the workgroup, so the entry, - // the workgroup count and every binding stay exactly as the - // compiler emitted them. - dispatch.gemv_shape = Some(shape); + dispatch.kernel = crate::compile::Kernel::Gemv { + shape, + integer_dot: dispatch.gemv_int_dot(), + }; return; } - dispatch.shader = self.shader(&class.shader); - dispatch.use_small_tiles = class.conv2d.is_none() && self == Self::Tile32; - dispatch.use_coop = matches!(self, Self::CooperativeF32 { .. }); - dispatch.use_coop_compensated = false; - dispatch.conv_k_tile = match self { - Self::SpecializedConv { k_tile, .. } => Some(k_tile), - _ => None, + dispatch.shader = self.shader(&dispatch.shader); + dispatch.kernel = match self { + Self::Tile32 + if !matches!( + dispatch.shader, + ShaderEntry::Conv2dGemmSmall + | ShaderEntry::Conv2dGradInputGemmSmall + | ShaderEntry::Conv2dGradWeightGemmSmall + ) => + { + crate::compile::Kernel::SmallTile + } + Self::Tile32 | Self::Tile64 => crate::compile::Kernel::Default, + Self::Scalar(shape) => crate::compile::Kernel::ScalarMatmul(shape), + Self::CooperativeF32 { .. } => crate::compile::Kernel::Cooperative, + Self::SpecializedConv { k_tile, .. } => { + crate::compile::Kernel::SpecializedConv { k_tile } + } + Self::Gemv(_) => unreachable!(), }; - dispatch.scalar_fallback = dispatch - .use_coop - .then(|| (dispatch.shader.clone(), Self::Tile64.workgroups(class))); - dispatch.workgroups = self.workgroups(class); } pub(crate) fn shader(self, entry: &ShaderEntry) -> ShaderEntry { @@ -302,12 +321,10 @@ impl MatmulTile { fn workgroups(self, class: &TuneClass) -> [u32; 3] { let tile = match self { - // One workgroup per output vec4 (or per row, transposed), set by - // N alone. Changing the threads per workgroup does not change how - // many there are. - Self::Gemv(_) => return class.gemv_workgroups(), + Self::Gemv(shape) => return class.gemv_workgroups(shape), Self::Tile32 => 32, Self::Tile64 => 64, + Self::Scalar(shape) => shape.tile_size, Self::SpecializedConv { tile_size, .. } => tile_size, Self::CooperativeF32 { tile_size } => { let tile = 2 * tile_size; @@ -325,8 +342,20 @@ impl MatmulTile { } pub(crate) fn buffer_sizes(self, class: &TuneClass) -> Option> { + if let Self::Scalar(shape) = self { + if class.conv2d.is_some() + || gemv_group(&class.shader).is_some() + || class.weight_format.is_quantized() + || !matches!(shape.tile_size, 32 | 64) + || !matches!(shape.k_stage, 8 | 16 | 32) + { + return None; + } + } if let Self::Gemv(shape) = self { - if gemv_group(&class.shader).is_none() || !matches!(shape.threads, 32 | 64 | 128 | 256) + if gemv_group(&class.shader).is_none() + || !matches!(shape.threads, 32 | 64 | 128 | 256) + || !matches!(shape.bt_rows, 1 | 2 | 4) { return None; } @@ -398,6 +427,7 @@ impl TuneClass { | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd | ShaderEntry::MatMulGemvAdd + | ShaderEntry::MatMulGemvBTAdd ); if !matches!( dispatch.shader, @@ -410,23 +440,21 @@ impl TuneClass { | ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvAdd | ShaderEntry::MatMulGemvBT + | ShaderEntry::MatMulGemvBTAdd | ShaderEntry::Conv2dGemm | ShaderEntry::Conv2dGemmSmall | ShaderEntry::Conv2dGradInputGemm | ShaderEntry::Conv2dGradInputGemmSmall | ShaderEntry::Conv2dGradWeightGemm | ShaderEntry::Conv2dGradWeightGemmSmall - ) || dispatch.use_coop_compensated - || (dispatch.use_coop && dispatch.use_small_tiles) - || (dispatch.use_coop && dispatch.weight_format.uses_reduced_storage()) + ) || dispatch.use_coop_compensated() + || (dispatch.use_coop() && dispatch.weight_format.uses_reduced_storage()) || dispatch.horizontal_batch >= 2 || dispatch.matmul_prologue.is_some() || dispatch.matmul_epilogue.is_some() - || !dispatch.epilogue.is_empty() - || !dispatch.epilogue_buffers.is_empty() || !dispatch.extra_outputs.is_empty() - || dispatch.pointwise.is_some() - || dispatch.reduction.is_some() + || dispatch.pointwise().is_some() + || dispatch.reduction().is_some() || dispatch.input_buffers.len() != if addend { 3 } else { 2 } { return None; @@ -443,16 +471,15 @@ impl TuneClass { | ShaderEntry::Conv2dGradInputGemm | ShaderEntry::Conv2dGradWeightGemm ) { - if dispatch.use_coop - || dispatch.use_small_tiles - || dispatch.scalar_fallback.is_some() + if dispatch.use_coop() + || dispatch.use_small_tiles() || dispatch.weight_format.uses_reduced_storage() { return None; } Some(TuneConv2d::from_params(&dispatch.params)?) } else { - if dispatch.conv_k_tile.is_some() + if dispatch.conv_k_tile().is_some() || dispatch.workgroups[2] != 1 || dispatch.params.len() != 4 || dispatch.params[3] != 0 @@ -510,7 +537,7 @@ impl TuneClass { conv2d, requires_full_precision: dispatch.requires_full_precision, weight_format: dispatch.weight_format, - gemv_int_dot: dispatch.gemv_int_dot, + gemv_int_dot: dispatch.gemv_int_dot(), gemv_rmsnorm: dispatch.gemv_rmsnorm.is_some(), gemv_rmsnorm_eps_bits: dispatch .gemv_rmsnorm @@ -539,6 +566,7 @@ impl TuneClass { | ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd | ShaderEntry::MatMulGemvAdd + | ShaderEntry::MatMulGemvBTAdd ) } @@ -596,18 +624,20 @@ impl TuneClass { Some(sizes) } - /// Workgroups a K-split GEMV dispatches: one per output vec4, or one per - /// output row for the transposed form. Independent of the thread count. - pub(crate) fn gemv_workgroups(&self) -> [u32; 3] { - if self.shader == ShaderEntry::MatMulGemvBT { - [self.n, 1, 1] + /// Workgroups per output vec4, or per group of transposed B rows. + pub(crate) fn gemv_workgroups(&self, shape: crate::codegen::GemvShape) -> [u32; 3] { + if matches!( + self.shader, + ShaderEntry::MatMulGemvBT | ShaderEntry::MatMulGemvBTAdd + ) { + crate::compile::row_gemv_workgroups(self.n.div_ceil(shape.bt_rows)) } else { [self.n / 4, 1, 1] } } - /// A small deterministic tournament: scalar alternative first, then - /// native f32 where it fits. Capability/precision/padding are legality; + /// A bounded tournament over scalar layouts and native f32 where it fits. + /// Capability/precision/padding are legality; /// occupancy and the static large-shape veto do not exclude candidates. pub(crate) fn challengers( &self, @@ -615,16 +645,25 @@ impl TuneClass { config: Option<&CoopConfig>, ) -> Vec { if let MatmulTile::Gemv(shape) = initial { - // Both axes, widest first: wide workgroups hide DRAM latency at - // M=1, and the subgroup reduction removes barriers in proportion - // to the wave width. Which trade wins is exactly what a device - // disagrees with another device about, so measure the cross - // product rather than guessing a rule. use crate::codegen::{GemvReduction, GemvShape}; let mut out = Vec::new(); - for reduction in [GemvReduction::Subgroup, GemvReduction::Tree] { - for threads in [256, 128, 64, 32] { - out.push(MatmulTile::Gemv(GemvShape { threads, reduction })); + let rows: &[u32] = if matches!( + self.shader, + ShaderEntry::MatMulGemvBT | ShaderEntry::MatMulGemvBTAdd + ) { + &GemvShape::BT_ROWS + } else { + &[1] + }; + for &bt_rows in rows { + for reduction in [GemvReduction::Subgroup, GemvReduction::Tree] { + for threads in [256, 128, 64, 32] { + out.push(MatmulTile::Gemv(GemvShape { + threads, + reduction, + bt_rows, + })); + } } } return out @@ -670,7 +709,7 @@ impl TuneClass { .filter(|&tile| tile != initial && tile.fits(self)) .collect(); } - [ + let mut candidates: Vec<_> = [ Some(MatmulTile::Tile64), Some(MatmulTile::Tile32), MatmulTile::native_cooperative(config) @@ -678,8 +717,25 @@ impl TuneClass { ] .into_iter() .flatten() - .filter(|&tile| tile != initial && tile.fits(self)) - .collect() + .collect(); + if !self.weight_format.is_quantized() { + candidates.retain(|tile| matches!(tile, MatmulTile::CooperativeF32 { .. })); + for interleave_columns in [false, true] { + for k_stage in [32, 16, 8] { + for tile_size in [64, 32] { + candidates.push(MatmulTile::Scalar(crate::codegen::ScalarMatmulShape { + tile_size, + k_stage, + interleave_columns, + })); + } + } + } + } + candidates + .into_iter() + .filter(|&tile| tile != initial && tile.fits(self)) + .collect() } } @@ -739,7 +795,7 @@ pub struct TuneOptions { #[serde(default)] pub scope: TuneScope, pub max_classes: usize, - /// GPU scratch including the upload/readback buffer, not pipelines. + /// GPU comparison scratch including staging, not pipelines. pub max_scratch_bytes: usize, /// Defaults to Download. Does not alter scratch binding placement, /// validation or kernel candidates. @@ -758,7 +814,7 @@ pub struct TuneOptions { /// Complete, alternating baseline/candidate pairs required for a decision. pub sample_pairs: usize, /// Separate, barrier-delimited dispatches in each timed submission; - /// complete sequences for explicit split-K measurements. + /// complete sequences for split-K and cached-attention measurements. pub dispatches_per_sample: u32, /// Required fractional improvement, in addition to a noise margin. pub min_improvement: f64, @@ -892,14 +948,15 @@ pub struct TuneQualificationTimes { /// Evidence for one candidate comparison within an exact class. Times are /// batched scratch wall times per dispatch (per complete sequence when -/// `candidate_split_k` is present), not GPU timestamps or whole-step latency. +/// `candidate_split_k` is present), not GPU +/// timestamps or whole-step latency. #[derive(Clone, Debug, Serialize, Deserialize)] -pub struct TuneOutcome { - pub class: TuneClass, +pub struct TuneOutcome { + pub class: Class, pub dispatches: usize, - pub initial: MatmulTile, - pub candidate: MatmulTile, - pub selected: MatmulTile, + pub initial: Choice, + pub candidate: Choice, + pub selected: Choice, /// Experimental two-pass dW challenger; baseline is unsplit with the same /// tile. Times are per complete sequence. `decision` reports the result; /// no live plan is changed and `selected` still describes only the tile. @@ -928,13 +985,8 @@ pub struct TuneOutcome { pub noise_margin_ms: Option, } -impl TuneOutcome { - pub(crate) fn new( - class: TuneClass, - dispatches: usize, - initial: MatmulTile, - candidate: MatmulTile, - ) -> Self { +impl TuneOutcome { + pub(crate) fn new(class: Class, dispatches: usize, initial: Choice, candidate: Choice) -> Self { Self { class, dispatches, @@ -986,6 +1038,9 @@ pub struct TuneReport { /// Live search visits a bounded candidate set per class; explicit split-K /// probes accept up to four counts against the same unsplit control. pub visited_classes: usize, + /// Qualified selections reused inside one calibrated build, never globally. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub reused_classes: Vec<(TuneClass, MatmulTile)>, pub excluded_dispatches: usize, pub class_limit_reached: bool, pub time_budget_exhausted: bool, @@ -1040,7 +1095,10 @@ pub(crate) fn measure_pairs( (baseline, candidate) } -pub(crate) fn decide(outcome: &mut TuneOutcome, options: &TuneOptions) { +pub(crate) fn decide( + outcome: &mut TuneOutcome, + options: &TuneOptions, +) { if outcome.baseline_ms.len() != options.sample_pairs || outcome.candidate_ms.len() != options.sample_pairs { @@ -1107,7 +1165,10 @@ mod tests { ..Default::default() }; let quantized = Dispatch { - gemv_int_dot: true, + kernel: crate::compile::Kernel::Gemv { + shape: crate::codegen::GemvShape::initial(crate::codegen::ShaderGroup::MatMulGemv), + integer_dot: true, + }, ..ordinary.clone() }; @@ -1179,7 +1240,11 @@ mod tests { seen.push(initial); for threads in GemvShape::WIDTHS { for reduction in [GemvReduction::Tree, GemvReduction::Subgroup] { - let shape = GemvShape { threads, reduction }; + let shape = GemvShape { + threads, + reduction, + bt_rows: 1, + }; assert!( seen.contains(&MatmulTile::Gemv(shape)), "{dtype:?}: missing {shape:?} from {seen:?}" @@ -1245,6 +1310,7 @@ mod tests { MatmulTile::Gemv(GemvShape { threads: 64, reduction: GemvReduction::Subgroup, + bt_rows: 1, }) .fits(&class), "64-wide subgroup fused Q4_0 GEMV must fit the class scratch" @@ -1299,14 +1365,41 @@ mod tests { #[test] fn complete_geometry_handles_edges_in_both_directions() { let mut d = dispatch(); - let class = TuneClass::from_dispatch(&d, None).unwrap(); + let mut class = TuneClass::from_dispatch(&d, None).unwrap(); assert_eq!((class.m, class.n, class.k), (33, 65, 17)); + class.binding_bytes = class.buffer_sizes().unwrap(); + for candidate in class.challengers(MatmulTile::Tile64, None) { + candidate.apply(&mut d, &class); + assert_eq!(MatmulTile::selected(&d, None), Some(candidate)); + assert_eq!(d.use_small_tiles(), d.workgroups == [3, 2, 1]); + let roundtrip: Dispatch = + serde_json::from_str(&serde_json::to_string(&d).unwrap()).unwrap(); + assert_eq!(roundtrip, d); + } MatmulTile::Tile32.apply(&mut d, &class); + assert!(d.scalar_matmul().is_none()); assert_eq!(d.workgroups, [3, 2, 1]); - assert!(d.use_small_tiles); + assert!(d.use_small_tiles()); MatmulTile::Tile64.apply(&mut d, &class); assert_eq!(d.workgroups, [2, 1, 1]); - assert!(!d.use_small_tiles); + assert!(!d.use_small_tiles()); + d.shader = ShaderEntry::MatMulGemvBT; + d.params = vec![1, 262_145, 4, 0]; + d.workgroups = crate::compile::row_gemv_workgroups(262_145); + let class = TuneClass::from_dispatch(&d, None).unwrap(); + for bt_rows in crate::codegen::GemvShape::BT_ROWS { + let shape = crate::codegen::GemvShape { + threads: 64, + reduction: crate::codegen::GemvReduction::Subgroup, + bt_rows, + }; + MatmulTile::Gemv(shape).apply(&mut d, &class); + assert_eq!( + d.workgroups, + crate::compile::row_gemv_workgroups(262_145u32.div_ceil(bt_rows)) + ); + assert_eq!(TuneClass::from_dispatch(&d, None), Some(class.clone())); + } } #[test] @@ -1356,9 +1449,9 @@ mod tests { fn unsupported_modifiers_never_enter_the_search() { let base = dispatch(); let mut variants = vec![base.clone(); 10]; - variants[0].use_coop = true; + variants[0].kernel = crate::compile::Kernel::Cooperative; variants[1].horizontal_batch = 2; - variants[2].use_coop_compensated = true; + variants[2].kernel = crate::compile::Kernel::CooperativeCompensated; variants[3].extra_outputs.push(crate::compile::BufferRef(3)); variants[4].workgroups[2] = 2; variants[5].shader = ShaderEntry::MatMulGemv; @@ -1377,7 +1470,7 @@ mod tests { fn conv_dispatch(shader: ShaderEntry) -> Dispatch { let mut d = dispatch(); d.shader = shader; - d.use_small_tiles = false; + d.kernel = crate::compile::Kernel::Default; d.params = vec![2, 3, 7, 9, 5, 3, 2, 2, 0, 3, 5, 1]; d.workgroups = if matches!( d.shader, @@ -1464,7 +1557,7 @@ mod tests { [1, 1, 1] } ); - assert!(!d.use_small_tiles && !d.use_coop && d.scalar_fallback.is_none()); + assert!(!d.use_small_tiles() && !d.use_coop()); assert_eq!(MatmulTile::selected(&d, None), Some(MatmulTile::Tile32)); let small = TuneClass::from_dispatch(&d, None).unwrap(); let mut expected = class.clone(); @@ -1493,9 +1586,9 @@ mod tests { ] { let base = conv_dispatch(shader); let mut variants = vec![base; 17]; - variants[0].use_small_tiles = true; - variants[1].use_coop = true; - variants[2].scalar_fallback = Some((ShaderEntry::MatMul, [1; 3])); + variants[0].kernel = crate::compile::Kernel::SmallTile; + variants[1].kernel = crate::compile::Kernel::Cooperative; + variants[2].kernel = crate::compile::Kernel::CooperativeCompensated; variants[3].params.pop(); variants[4].params[7] = 0; variants[5].params[9] += 1; @@ -1507,7 +1600,7 @@ mod tests { variants[11].params[1] = u32::MAX; variants[12].weight_format = crate::compile::WeightFormat::F16; variants[13].shader = ShaderEntry::Conv2dGradInputGemmCoopGen(3, 2, 2); - variants[14].conv_k_tile = Some(7); + variants[14].kernel = crate::compile::Kernel::SpecializedConv { k_tile: 7 }; variants[15].input_buffers.pop(); variants[16].params[5] = 100; for d in variants { @@ -1557,7 +1650,13 @@ mod tests { scope.includes(&dense), matches!(scope, TuneScope::Dense | TuneScope::All) ); - assert_eq!(scope.includes(&conv), scope != TuneScope::Dense); + assert_eq!( + scope.includes(&conv), + matches!( + scope, + TuneScope::ConvDerivatives | TuneScope::Convolution | TuneScope::All + ) + ); let options = TuneOptions { scope, ..Default::default() @@ -1875,7 +1974,7 @@ mod tests { } #[test] - fn native_geometry_flags_and_scalar_fallback_move_together() { + fn native_geometry_and_kernel_move_together() { let config = native_config(8); let native = MatmulTile::CooperativeF32 { tile_size: 8 }; let class = class(32, 64, 17); @@ -1883,13 +1982,12 @@ mod tests { d.params = vec![class.m, class.k, class.n, 0]; native.apply(&mut d, &class); assert_eq!(d.workgroups, [2, 4, 1]); - assert!(d.use_coop && !d.use_coop_compensated && !d.use_small_tiles); - assert_eq!(d.scalar_fallback, Some((ShaderEntry::MatMul, [1, 1, 1]))); + assert!(d.use_coop() && !d.use_coop_compensated() && !d.use_small_tiles()); assert!(TuneClass::from_dispatch(&d, Some(&config)).is_some()); assert!(TuneClass::from_dispatch(&d, None).is_none()); MatmulTile::Tile32.apply(&mut d, &class); assert_eq!(d.workgroups, [2, 1, 1]); - assert!(!d.use_coop && d.use_small_tiles && d.scalar_fallback.is_none()); + assert!(!d.use_coop() && d.use_small_tiles()); assert!(TuneClass::from_dispatch(&d, Some(&config)).is_some()); } diff --git a/tests/block_matmul.rs b/tests/block_matmul.rs index 26fb8781..adf097d0 100644 --- a/tests/block_matmul.rs +++ b/tests/block_matmul.rs @@ -218,7 +218,7 @@ fn cpu_all_three_chain_rules_match_finite_differences() { }) .collect(); assert_eq!(block_gradients.len(), 2); - assert!(block_gradients.iter().all(|d| !d.use_coop)); + assert!(block_gradients.iter().all(|d| !d.use_coop())); } } } @@ -245,12 +245,12 @@ fn cpu_production_shapes_use_one_dispatch_per_product() { | ShaderEntry::BlockMatMulBT )); assert_eq!(d.params, [rows as u32, n as u32, k as u32, 8]); - let tile = if d.use_small_tiles { 32 } else { 64 }; + let tile = if d.use_small_tiles() { 32 } else { 64 }; assert_eq!( d.workgroups, [(n as u32).div_ceil(tile), (rows as u32).div_ceil(tile), 8] ); - assert!(!d.use_coop); + assert!(!d.use_coop()); } } } diff --git a/tests/conv_derivatives.rs b/tests/conv_derivatives.rs index d7d34db1..a99ab98d 100644 --- a/tests/conv_derivatives.rs +++ b/tests/conv_derivatives.rs @@ -229,7 +229,7 @@ fn run_split( .plan() .dispatches .iter() - .any(|d| d.use_coop + .any(|d| d.use_coop() && matches!(d.shader, ShaderEntry::Conv2dGradInputGemmCoopGen(..))), "generated dX kernel must actually execute" ); @@ -870,7 +870,7 @@ fn generated_conv_indexing_matches_full_oracle_at_reciprocal_boundaries() { .plan() .dispatches .iter() - .any(|d| d.use_coop && matches!(d.shader, ShaderEntry::Conv2dGemmCoopGen(..))) + .any(|d| d.use_coop() && matches!(d.shader, ShaderEntry::Conv2dGemmCoopGen(..))) ); } } diff --git a/tests/device_local.rs b/tests/device_local.rs index 8541f095..efe0472e 100644 --- a/tests/device_local.rs +++ b/tests/device_local.rs @@ -18,7 +18,7 @@ fn readback_preserves_bits_across_sizes_and_updates() { graph.set_outputs(vec![output]); let mut session = meganeura::build(&graph, meganeura::SessionConfig::inference_from_env()).0; for seed in [0u32, 17] { - let values: Vec<_> = (0..len) + let mut values: Vec<_> = (0..len) .map(|i| { f32::from_bits(match i % 7 { 0 => 0x7fc0_0123, @@ -28,12 +28,13 @@ fn readback_preserves_bits_across_sizes_and_updates() { }) }) .collect(); - session.set_input("x", &values); - session.step(); - session.wait(); for count in [0, 17, len, 1024, len] { + values[3] = f32::from_bits(values[3].to_bits().wrapping_add(1)); + values[len - 1] = values[3]; + session.set_input("x", &values); + session.step(); let mut actual = vec![0.0; count]; - session.read_output_by_index(0, &mut actual); + session.wait_read_output(0, &mut actual); assert!( actual .iter() diff --git a/tests/flash_grad_kv_short.rs b/tests/flash_grad_kv_short.rs index 23b7a135..3d4c6562 100644 --- a/tests/flash_grad_kv_short.rs +++ b/tests/flash_grad_kv_short.rs @@ -1,4 +1,4 @@ -//! Short-sequence cooperative dK/dV parity on GPUs with 16x16 f16 tiles. +//! Short-sequence cooperative forward/backward parity on GPUs with f16 tiles. //! //! SmolVLA uses Q=50 with both KV=16 (cross attention) and KV=50 (self //! attention). These shapes must not inherit the scalar flash kernel's much @@ -12,10 +12,14 @@ fn run( gpu: Arc, q_seq: usize, kv_seq: usize, + window: u32, cooperative: bool, -) -> (Vec, Vec, bool) { +) -> ([Vec; 3], bool) { unsafe { - std::env::set_var("MEGANEURA_FLASH_FWD_COOP", "0"); + std::env::set_var( + "MEGANEURA_FLASH_FWD_COOP", + if cooperative { "1" } else { "0" }, + ); std::env::set_var( "MEGANEURA_FLASH_BWD_COOP", if cooperative { "1" } else { "0" }, @@ -27,8 +31,11 @@ fn run( let q = graph.parameter("q", &[q_seq, num_heads as usize * head_dim as usize]); let k = graph.parameter("k", &[kv_seq, num_kv_heads as usize * head_dim as usize]); let v = graph.parameter("v", &[kv_seq, num_kv_heads as usize * head_dim as usize]); - let attention = - graph.multi_head_attn(q, k, v, num_heads, num_kv_heads, head_dim, q_seq != kv_seq); + let attention = if window == 0 { + graph.multi_head_attn(q, k, v, num_heads, num_kv_heads, head_dim, q_seq != kv_seq) + } else { + graph.sliding_window_attention(q, k, v, num_heads, num_kv_heads, head_dim, window) + }; let loss = graph.mean_all(attention); graph.set_outputs(vec![loss]); @@ -61,11 +68,17 @@ fn run( session.step(); session.wait(); - let mut dk = vec![0.0; k_data.len()]; - let mut dv = vec![0.0; v_data.len()]; - session.read_param_grad("k", &mut dk); - session.read_param_grad("v", &mut dv); - (dk, dv, uses_coop) + let gradients = [ + ("q", q_data.len()), + ("k", k_data.len()), + ("v", v_data.len()), + ] + .map(|(name, len)| { + let mut gradient = vec![0.0; len]; + session.read_param_grad(name, &mut gradient); + gradient + }); + (gradients, uses_coop) } fn assert_close(label: &str, scalar: &[f32], cooperative: &[f32]) { @@ -81,24 +94,29 @@ fn assert_close(label: &str, scalar: &[f32], cooperative: &[f32]) { .max(1e-6); assert!( max_abs / scale < 0.02, - "{label}: cooperative dK/dV differs from scalar by {:.3}% (max abs {max_abs:.3e})", + "{label}: cooperative gradient differs from scalar by {:.3}% (max abs {max_abs:.3e})", max_abs / scale * 100.0, ); } #[test] -fn short_cross_and_self_attention_grad_kv_match_scalar() { +fn short_cross_self_and_window_attention_gradients_match_scalar() { let gpu = Arc::new( meganeura::init_gpu_context_with(meganeura::GpuOptions::from_env()).expect("GPU context"), ); let has_coop = gpu.capabilities().cooperative_matrix.f16_tile == 16; - for (label, q_seq, kv_seq) in [("cross", 50, 16), ("self", 50, 50)] { - let (scalar_k, scalar_v, scalar_used_coop) = run(gpu.clone(), q_seq, kv_seq, false); - let (coop_k, coop_v, coop_used_coop) = run(gpu.clone(), q_seq, kv_seq, true); + for (label, q_seq, kv_seq, window) in [ + ("cross", 50, 16, 0), + ("self", 50, 50, 0), + ("window", 50, 50, 17), + ] { + let (scalar, scalar_used_coop) = run(gpu.clone(), q_seq, kv_seq, window, false); + let (cooperative, coop_used_coop) = run(gpu.clone(), q_seq, kv_seq, window, true); assert!(!scalar_used_coop); assert_eq!(coop_used_coop, has_coop); - assert_close(&format!("{label} dK"), &scalar_k, &coop_k); - assert_close(&format!("{label} dV"), &scalar_v, &coop_v); + for (i, name) in ["dQ", "dK", "dV"].into_iter().enumerate() { + assert_close(&format!("{label} {name}"), &scalar[i], &cooperative[i]); + } } } diff --git a/tests/gemma_inference_ops.rs b/tests/gemma_inference_ops.rs index 3240c5ee..3a9842cb 100644 --- a/tests/gemma_inference_ops.rs +++ b/tests/gemma_inference_ops.rs @@ -227,17 +227,18 @@ fn chunked_relative_attention_matches_blocked_reference() { #[test] fn cached_block_writes_only_valid_rows_and_selects_last() { - for (max_seq, window, dim) in [ - (6, 0, 4), - (96, 0, 4), - (96, 37, 4), - (96, 0, 64), - (96, 37, 80), - (6, 0, 512), + for (max_seq, window, dim, splits) in [ + (6, 0, 4, 1), + (96, 0, 4, 2), + (96, 37, 4, 4), + (96, 0, 64, 8), + (96, 37, 80, 16), + (6, 0, 512, 4), ] { let block = 3; let mut graph = Graph::new(); let q = graph.input("q", &[block, 2 * dim]); + let q = graph.scale(q, 0.5); let new_k = graph.input("new_k", &[block, dim]); let new_v = graph.input("new_v", &[block, dim]); let k_cache = graph.parameter("k", &[max_seq, dim]); @@ -252,8 +253,10 @@ fn cached_block_writes_only_valid_rows_and_selects_last() { let output = graph.prefix_last(attended, valid); graph.set_outputs(vec![output]); - let mut session = - meganeura::build(&graph, meganeura::SessionConfig::inference_from_env()).0; + let mut config = meganeura::SessionConfig::inference_from_env(); + config.tune = false; + config.options.cached_attention_splits = Some(splits); + let mut session = meganeura::build(&graph, config).0; session.set_input("q", &vec![0.0; block * 2 * dim]); let new_k: Vec<_> = (0..block * dim) .map(|i| { @@ -295,7 +298,7 @@ fn cached_block_writes_only_valid_rows_and_selects_last() { let queries: Vec<_> = (0..block * 2 * dim) .map(|i| (i as f32 * 0.7).sin()) .collect(); - session.set_input("q", &queries); + session.set_input("q", &queries.iter().map(|x| x * 2.0).collect::>()); for (i, (k, v)) in initial_k.iter_mut().zip(&mut initial_v).enumerate() { *k = (i as f32 * 0.3).sin() * 4.0; *v = (i as f32 * 0.4).cos(); diff --git a/tests/gemv_parity.rs b/tests/gemv_parity.rs index 55243b13..16051c2b 100644 --- a/tests/gemv_parity.rs +++ b/tests/gemv_parity.rs @@ -152,7 +152,13 @@ fn q40_rmsnorm_folds_into_gemv() { let w = g.parameter_q40("w", &[K, N]); let y = g.matmul(h, w); g.set_outputs(vec![y]); - let mut session = meganeura::build(&g, meganeura::SessionConfig::inference_from_env()).0; + let mut config = meganeura::SessionConfig::inference_from_env(); + config.options.gemv_shape = Some(meganeura::GemvShape { + threads: 64, + reduction: meganeura::GemvReduction::Subgroup, + bt_rows: 1, + }); + let mut session = meganeura::build(&g, config).0; let fused = session .plan() .dispatches @@ -194,6 +200,165 @@ fn gemv_non_multiple_of_256() { test_shape(128, 1, 9); } +#[test] +fn dense_transposed_glu_matches_reference_after_restage() { + for (m, k, n) in [(1, 12, 7), (3, 5, 7)] { + for f16 in [false, true] { + for gelu in [false, true] { + let mut g = Graph::new(); + let x = g.input("x", &[m, k]); + let projected = if m == 1 { + let nw = g.parameter("norm", &[k]); + g.rms_norm(x, nw, 1e-5) + } else { + x + }; + let mut project = |name| { + let w = if f16 { + g.parameter_f16(name, &[n, k]) + } else { + g.parameter(name, &[n, k]) + }; + g.matmul_bt(projected, w) + }; + let gate = project("gate"); + let up = project("up"); + let out = if gelu { + g.geglu(gate, up) + } else { + g.swiglu(gate, up) + }; + g.set_outputs(vec![out]); + let mut config = meganeura::SessionConfig::inference_from_env(); + config.tune = false; + let mut session = meganeura::build(&g, config).0; + assert_eq!(session.plan().derived_params.len(), 1); + assert!(matches!( + session.plan().derived_params[0].2, + meganeura::graph::ParamTransform::VerticalConcat + )); + let x: Vec<_> = (0..m * k).map(|i| (i as f32 * 0.7).sin()).collect(); + let mut gate: Vec<_> = (0..n * k).map(|i| (i as f32 * 0.3).cos() * 0.25).collect(); + let mut up: Vec<_> = (0..n * k).map(|i| (i as f32 * 0.2).sin() * 0.25).collect(); + session.set_input("x", &x); + let nw: Vec<_> = (0..k).map(|i| 0.75 + (i % 5) as f32 * 0.125).collect(); + if m == 1 { + session.set_parameter("norm", &nw); + assert_eq!( + session + .plan() + .dispatches + .iter() + .filter(|d| d.gemv_rmsnorm.is_some()) + .count(), + 1 + ); + } + let inv_rms = if m == 1 { + (x.iter().map(|v| v * v).sum::() / k as f32 + 1e-5) + .sqrt() + .recip() + } else { + 1.0 + }; + for pass in 0..3 { + if pass == 1 { + gate.iter_mut().for_each(|v| *v *= -0.5); + } + if pass == 2 { + up.iter_mut().for_each(|v| *v *= 0.75); + } + // Reverse source order; mix f32 and exact f16 uploads. + if f16 && pass == 2 { + let bytes: Vec<_> = up + .iter() + .flat_map(|&v| half::f16::from_f32(v).to_le_bytes()) + .collect(); + session.set_parameter_packed("up", &bytes); + } else { + session.set_parameter("up", &up); + } + session.set_parameter("gate", &gate); + session.step(); + session.wait(); + let rounded = |v| { + if f16 { + half::f16::from_f32(v).to_f32() + } else { + v + } + }; + let mut expected = Vec::new(); + for row in 0..m { + for col in 0..n { + let a: f32 = (0..k) + .map(|i| { + x[row * k + i] + * inv_rms + * if m == 1 { nw[i] } else { 1.0 } + * rounded(gate[col * k + i]) + }) + .sum(); + let b: f32 = (0..k) + .map(|i| { + x[row * k + i] + * inv_rms + * if m == 1 { nw[i] } else { 1.0 } + * rounded(up[col * k + i]) + }) + .sum(); + let activated = if gelu { + 0.5 * a * (1.0 + (0.7978846 * (a + 0.044715 * a * a * a)).tanh()) + } else { + a / (1.0 + (-a).exp()) + }; + expected.push(activated * b); + } + } + assert_close_named( + "transposed GLU", + &session.read_output(m * n), + &expected, + 1e-4, + 1e-5, + ); + if m == 1 && !gelu && pass == 0 { + let before = session.read_output(m * n); + let report = session + .tune_with(meganeura::tune::TuneOptions { + scope: meganeura::tune::TuneScope::Dense, + max_time: std::time::Duration::from_secs(10), + sample_pairs: 4, + dispatches_per_sample: 1, + ..Default::default() + }) + .unwrap(); + assert_eq!(report.eligible_classes, 1); + assert!(!report.time_budget_exhausted); + assert!(!report.outcomes.is_empty()); + assert!( + report + .outcomes + .iter() + .all(|o| o.qualified && o.class.gemv_rmsnorm) + ); + assert_eq!(session.read_output(m * n), before); + session.step(); + session.wait(); + assert_close_named( + "tuned transposed GLU", + &session.read_output(m * n), + &expected, + 1e-4, + 1e-5, + ); + } + } + } + } + } +} + // ---- FusedMatMulAdd (GEMV + residual) ---- /// CPU reference: 1×K × K×N + D[1,N]. @@ -302,7 +467,13 @@ fn test_gemv_bt_shape(k: usize, n: usize, seed: u32) { let c = g.matmul_bt(a_n, b_n); g.set_outputs(vec![c]); - let mut session = meganeura::build(&g, meganeura::SessionConfig::inference_from_env()).0; + let mut config = meganeura::SessionConfig::inference_from_env(); + config.options.gemv_shape = Some(meganeura::GemvShape { + threads: 32, + reduction: meganeura::GemvReduction::Tree, + bt_rows: 4, + }); + let mut session = meganeura::build(&g, config).0; let plan = session.plan(); let gemv_bt_count = plan @@ -316,6 +487,11 @@ fn test_gemv_bt_shape(k: usize, n: usize, seed: u32) { "expected one MatMulGemvBT dispatch for k%4==0, got {}", gemv_bt_count, ); + assert!( + plan.dispatches + .iter() + .all(|d| d.workgroups.iter().all(|&n| n <= 65_535)) + ); } else { assert_eq!( gemv_bt_count, 0, @@ -343,6 +519,7 @@ fn test_gemv_bt_shape(k: usize, n: usize, seed: u32) { fn gemv_bt_smollm2_lm_head() { // SmolLM2-135M LM head (weight-tied): 1×576 × 49152×576^T → 1×49152. test_gemv_bt_shape(576, 49152, 200); + test_gemv_bt_shape(4, 262_145, 204); } #[test] @@ -369,11 +546,10 @@ fn gemv_non_multiple_k() { test_shape(513, 256, 12); } -/// Every GEMV kernel must be correct at every shape the search may install. -/// Plain, fused-add and transposed-B run together so eight sessions cover the -/// 24 combinations while failures retain the shape and kernel name. +/// Cover widths, reduction paths and row tails; live tuning qualifies each candidate. +/// Plain and transposed-B products, with and without addends, share each session. #[test] -fn every_gemv_shape_computes_the_same_product() { +fn gemv_shapes_cover_widths_reductions_and_row_tails() { use meganeura::compile::{CompileOptions, ShaderEntry}; use meganeura::train::{Mode, SessionConfig}; use meganeura::{GemvReduction, GemvShape}; @@ -393,6 +569,7 @@ fn every_gemv_shape_computes_the_same_product() { // and not multiples of every width, so the loop tails are exercised. const K: usize = 320; const N: usize = 36; + const BT_N: usize = N - 1; let a = data(K, 1); let b = data(K * N, 2); let addend = data(N, 3); @@ -407,67 +584,92 @@ fn every_gemv_shape_computes_the_same_product() { let want = cpu_gemv(&a, &b, K, N); let want_add = cpu_gemv_add(&a, &b, &addend, K, N); - let want_bt = cpu_gemv_bt(&a, &b_t, K, N); - - for threads in GemvShape::WIDTHS { - for reduction in [GemvReduction::Tree, GemvReduction::Subgroup] { - let shape = GemvShape { threads, reduction }; - let mut g = Graph::new(); - let x = g.input("x", &[1, K]); - let x_add = g.input("x_add", &[1, K]); - let x_bt = g.input("x_bt", &[1, K]); - let w = g.input("w", &[K, N]); - let w_add = g.input("w_add", &[K, N]); - let w_t = g.input("w_t", &[N, K]); - let d = g.input("d", &[1, N]); - let plain = g.matmul(x, w); - let product = g.matmul(x_add, w_add); - let add = g.add(product, d); - let bt = g.matmul_bt(x_bt, w_t); - g.set_outputs(vec![plain, add, bt]); - - let config = SessionConfig { - mode: Mode::Inference, - options: CompileOptions { - gemv_shape: Some(shape), - ..CompileOptions::from_env() - }, - ..SessionConfig::from_env() - }; - let mut s = meganeura::build(&g, config).0; - let shaders: Vec<_> = s.plan().dispatches.iter().map(|d| &d.shader).collect(); - for expected in [ - ShaderEntry::MatMulGemv, - ShaderEntry::MatMulGemvAdd, - ShaderEntry::MatMulGemvBT, - ] { - assert_eq!( - shaders - .iter() - .filter(|&&shader| *shader == expected) - .count(), - 1, - "{shape:?}: missing {expected:?}; got {shaders:?}" - ); - } - for name in ["x", "x_add", "x_bt"] { - s.set_input(name, &a); - } - s.set_input("w", &b); - s.set_input("w_add", &b); - s.set_input("w_t", &b_t); - s.set_input("d", &addend); - s.step(); - s.wait(); - for (index, (label, expected)) in - [("plain", &want), ("add", &want_add), ("bt", &want_bt)] - .into_iter() - .enumerate() - { - let mut got = vec![0.0; N]; - s.read_output_by_index(index, &mut got); - assert_close_named(&format!("{shape:?} {label}"), &got, expected, 2e-4, 2e-4); - } + let want_bt = cpu_gemv_bt(&a, &b_t, K, BT_N); + let want_bt_add: Vec<_> = want_bt.iter().zip(&addend).map(|(a, b)| a + b).collect(); + + use GemvReduction::{Subgroup, Tree}; + for (threads, reduction, bt_rows) in [ + (32, Tree, 1), + (32, Subgroup, 4), + (64, Tree, 2), + (64, Subgroup, 1), + (128, Tree, 4), + (128, Subgroup, 2), + (256, Tree, 1), + (256, Subgroup, 4), + ] { + let shape = GemvShape { + threads, + reduction, + bt_rows, + }; + let mut g = Graph::new(); + let x = g.input("x", &[1, K]); + let x_add = g.input("x_add", &[1, K]); + let x_bt = g.input("x_bt", &[1, K]); + let x_bt_add = g.input("x_bt_add", &[1, K]); + let w = g.input("w", &[K, N]); + let w_add = g.input("w_add", &[K, N]); + let w_t = g.input("w_t", &[BT_N, K]); + let w_t_add = g.input("w_t_add", &[BT_N, K]); + let d = g.input("d", &[1, N]); + let d_bt = g.input("d_bt", &[1, BT_N]); + let plain = g.matmul(x, w); + let product = g.matmul(x_add, w_add); + let add = g.add(product, d); + let bt = g.matmul_bt(x_bt, w_t); + let bt_product = g.matmul_bt(x_bt_add, w_t_add); + let bt_add = g.add(bt_product, d_bt); + g.set_outputs(vec![plain, add, bt, bt_add]); + + let config = SessionConfig { + mode: Mode::Inference, + options: CompileOptions { + gemv_shape: Some(shape), + ..CompileOptions::from_env() + }, + ..SessionConfig::from_env() + }; + let mut s = meganeura::build(&g, config).0; + let shaders: Vec<_> = s.plan().dispatches.iter().map(|d| &d.shader).collect(); + for expected in [ + ShaderEntry::MatMulGemv, + ShaderEntry::MatMulGemvAdd, + ShaderEntry::MatMulGemvBT, + ShaderEntry::MatMulGemvBTAdd, + ] { + assert_eq!( + shaders + .iter() + .filter(|&&shader| *shader == expected) + .count(), + 1, + "{shape:?}: missing {expected:?}; got {shaders:?}" + ); + } + for name in ["x", "x_add", "x_bt", "x_bt_add"] { + s.set_input(name, &a); + } + s.set_input("w", &b); + s.set_input("w_add", &b); + s.set_input("w_t", &b_t[..BT_N * K]); + s.set_input("w_t_add", &b_t[..BT_N * K]); + s.set_input("d", &addend); + s.set_input("d_bt", &addend[..BT_N]); + s.step(); + s.wait(); + for (index, (label, expected)) in [ + ("plain", &want), + ("add", &want_add), + ("bt", &want_bt), + ("bt_add", &want_bt_add), + ] + .into_iter() + .enumerate() + { + let mut got = vec![0.0; expected.len()]; + s.read_output_by_index(index, &mut got); + assert_close_named(&format!("{shape:?} {label}"), &got, expected, 2e-4, 2e-4); } } } diff --git a/tests/gpu_smoke.rs b/tests/gpu_smoke.rs index 8b41a901..ed6891f7 100644 --- a/tests/gpu_smoke.rs +++ b/tests/gpu_smoke.rs @@ -1687,6 +1687,13 @@ fn batched_parameter_read_matches_uploaded_values() { let values = session.read_params(&["b", "a"]); assert_eq!(values, [b_values.to_vec(), a_values.to_vec()]); + let a = session.param_buffer("a").unwrap(); + let b = session.param_buffer("b").unwrap(); + assert_eq!( + session.read_buffers(&[b, a, b]), + [b_values.to_vec(), a_values.to_vec(), b_values.to_vec()] + ); + assert!(session.read_buffers(&[]).is_empty()); } #[test] @@ -4055,7 +4062,7 @@ fn q6k_preserves_subnormal_block_scales() { } } -/// Default greedy packing concatenates SwiGLU gate/up into one matmul. +/// Default extraction concatenates SwiGLU gate/up into one matmul. /// `set_parameter_packed` has to restage that derived buffer; uploading /// only the named sources leaves the fused weight uninitialized. #[test] @@ -4175,7 +4182,7 @@ fn gguf_repacked_matmuls_match_reference() { } /// Packed blocks run along the parameter's first dimension, which differs -/// from K for transposed B. Cover the plain, greedy add-fused, and epilogue +/// from K for transposed B. Cover the plain, add-fused, and epilogue /// routes without emitting a GPU shader for any of them. #[test] fn block_quantized_matmul_bt_variants_are_refused() { diff --git a/tests/int_dot_gemv.rs b/tests/int_dot_gemv.rs index 29650382..a881009e 100644 --- a/tests/int_dot_gemv.rs +++ b/tests/int_dot_gemv.rs @@ -503,7 +503,7 @@ fn run( ShaderEntry::MatMulGemv } ); - assert_eq!(dispatch.gemv_int_dot, quantized_activations); + assert_eq!(dispatch.gemv_int_dot(), quantized_activations); session.set_input("x", a); if let Some(d) = addend { session.set_input("d", d); @@ -551,7 +551,7 @@ fn run_q40_rmsnorm( .iter() .find(|dispatch| dispatch.shader == ShaderEntry::MatMulGemv) .expect("RmsNorm output should feed a GEMV"); - assert!(dispatch.gemv_int_dot); + assert!(dispatch.gemv_int_dot()); assert_eq!(dispatch.gemv_rmsnorm.is_some(), !expose_normalized); session.set_input("x", x); diff --git a/tests/rmsnorm_matmul_prologue.rs b/tests/rmsnorm_matmul_prologue.rs index 464397fd..18bb9ac1 100644 --- a/tests/rmsnorm_matmul_prologue.rs +++ b/tests/rmsnorm_matmul_prologue.rs @@ -78,7 +78,7 @@ fn coop_rmsnorm_matmul_prologue_matches_scalar_after_updates() { .plan() .dispatches .iter() - .any(|dispatch| dispatch.use_coop && dispatch.matmul_prologue.is_some()); + .any(|dispatch| dispatch.use_coop() && dispatch.matmul_prologue.is_some()); let projection = (0..INNER * COLS) .map(|index| ((index * 17 % 101) as f32 - 50.0) * 0.001) diff --git a/tests/schedule_pointwise.rs b/tests/schedule_pointwise.rs index 45b7f554..cce05143 100644 --- a/tests/schedule_pointwise.rs +++ b/tests/schedule_pointwise.rs @@ -207,7 +207,7 @@ fn softplus_compiles_to_one_pointwise_dispatch() { assert_eq!(plan.dispatches.len(), 1); let dispatch = &plan.dispatches[0]; assert_eq!(dispatch.input_buffers.len(), 1); - assert!(dispatch.pointwise.is_some()); + assert!(dispatch.pointwise().is_some()); } #[test] @@ -290,7 +290,7 @@ fn fusion_reduces_dispatch_count() { 1, "expected pointwise chain to collapse to one dispatch" ); - assert!(fused_plan.dispatches[0].pointwise.is_some()); + assert!(fused_plan.dispatches[0].pointwise().is_some()); } /// Parity of a 3-op chain, once the fusion pass has run. @@ -358,8 +358,7 @@ fn ternary_fusion_add_of_mul() { "expected mul+add to collapse into a single arity-3 dispatch" ); let dag = fused.dispatches[0] - .pointwise - .as_ref() + .pointwise() .expect("fused dispatch should carry a DAG"); assert_eq!(dag.n_inputs, 3); @@ -453,8 +452,8 @@ fn softmax_schedule_emits_two_dispatches() { ); assert_eq!(baseline.dispatches.len(), 1); assert_eq!(schedule.dispatches.len(), 2); - assert!(schedule.dispatches[0].reduction.is_some()); - assert!(schedule.dispatches[1].reduction.is_some()); + assert!(schedule.dispatches[0].reduction().is_some()); + assert!(schedule.dispatches[1].reduction().is_some()); } // ---- Reduction archetype: RmsNorm ---- diff --git a/tests/schedule_reduction.rs b/tests/schedule_reduction.rs index 131e7e8c..a20cbe34 100644 --- a/tests/schedule_reduction.rs +++ b/tests/schedule_reduction.rs @@ -761,10 +761,10 @@ fn two_gather_reduction_actually_fuses() { let reductions: Vec<_> = plan .dispatches .iter() - .filter(|d| d.reduction.is_some()) + .filter(|d| d.reduction().is_some()) .collect(); assert_eq!(reductions.len(), 1, "expected exactly one fused reduction"); - let k = reductions[0].reduction.as_ref().unwrap(); + let k = reductions[0].reduction().unwrap(); assert_eq!( k.n_per_elem, 2, "mul producer should fold to 2 per-elem streams" @@ -780,7 +780,7 @@ fn two_gather_reduction_actually_fuses() { let embeds = plan .dispatches .iter() - .filter(|d| d.shader == ShaderEntry::Embedding && d.reduction.is_none()) + .filter(|d| d.shader == ShaderEntry::Embedding && d.reduction().is_none()) .count(); assert_eq!(embeds, 0, "embedding dispatches should be folded away"); } @@ -813,17 +813,17 @@ fn shared_gather_and_offset_fold_into_each_reduction() { let reductions: Vec<_> = plan .dispatches .iter() - .filter(|dispatch| dispatch.reduction.is_some()) + .filter(|dispatch| dispatch.reduction().is_some()) .collect(); assert_eq!(reductions.len(), 2); for reduction in reductions { - let kernel = reduction.reduction.as_ref().unwrap(); + let kernel = reduction.reduction().unwrap(); assert_eq!(kernel.n_per_elem, 3); assert_eq!(kernel.gather_elem, vec![true, false, false]); assert_eq!(reduction.input_buffers.len(), 4); } assert!(!plan.dispatches.iter().any(|dispatch| { - dispatch.shader == ShaderEntry::Embedding && dispatch.reduction.is_none() + dispatch.shader == ShaderEntry::Embedding && dispatch.reduction().is_none() })); assert!(!plan.dispatches.iter().any(|dispatch| { dispatch.shader == ShaderEntry::Add && dispatch.params[0] == (m * n) as u32 diff --git a/tests/submission_chunks.rs b/tests/submission_chunks.rs index e9c6220c..60543982 100644 --- a/tests/submission_chunks.rs +++ b/tests/submission_chunks.rs @@ -74,12 +74,19 @@ fn run(chunks: usize, layers: usize, rows: usize, dim: usize, steps: usize) -> V // buffer in the ring is still in flight, which only shows up once the // encoder has wrapped around at least once. let mut out = vec![0.0f32; rows * dim]; - for _ in 0..steps { + let mut reference = None; + for step in 0..steps { session.step(); session.wait(); session.read_output_by_index(0, &mut out); + if let Some(ref expected) = reference { + assert_eq!(&out, expected, "schedule changed repeated graph output"); + } + if chunks == 1 && step == 0 { + reference = Some(out.clone()); + } } - out + reference.unwrap_or(out) } #[test] diff --git a/tests/tune.rs b/tests/tune.rs index a48909e0..c1000c9f 100644 --- a/tests/tune.rs +++ b/tests/tune.rs @@ -225,11 +225,16 @@ fn reduced_storage_tiles_preserve_outputs() { }) .unwrap(); assert_eq!(report.eligible_classes, 1); - assert_eq!(report.outcomes.len(), 1); - let outcome = &report.outcomes[0]; - assert!(outcome.qualified, "{outcome:?}"); - assert_eq!(outcome.class.weight_format, format); - assert_eq!(outcome.class.shader, shader); + assert!(!report.time_budget_exhausted); + assert_eq!( + report.outcomes.len(), + if format.is_quantized() { 1 } else { 11 } + ); + for outcome in &report.outcomes { + assert!(outcome.qualified, "{outcome:?}"); + assert_eq!(outcome.class.weight_format, format); + assert_eq!(outcome.class.shader, shader); + } assert_eq!( before, session.read_output(m * n), @@ -388,22 +393,24 @@ fn qualify_scalar_entries(staging: TuneStaging, staging_reuse: TuneStagingReuse) .unwrap(); assert_eq!( report.outcomes.len(), - 1, + 11, "{shader:?} {m}x{n}x{k}: {report:?}" ); - let outcome = &report.outcomes[0]; + assert!(!report.time_budget_exhausted); assert_eq!(report.scratch.unwrap().retained_staging_bytes, 0); assert_eq!( report.scratch.unwrap().staging_allocations, report.scratch.unwrap().staging_releases ); assert!(report.scratch.unwrap().peak_bytes <= report.options.max_scratch_bytes); - assert_eq!(outcome.class.shader, shader); - assert!(outcome.qualified, "{outcome:?}"); - assert!(matches!( - outcome.decision, - TuneDecision::FasterCandidate | TuneDecision::KeepBaseline - )); + for outcome in &report.outcomes { + assert_eq!(outcome.class.shader, shader); + assert!(outcome.qualified, "{outcome:?}"); + assert!(matches!( + outcome.decision, + TuneDecision::FasterCandidate | TuneDecision::KeepBaseline + )); + } } } } @@ -501,14 +508,23 @@ fn tuning_releases_retained_staging_after_a_later_scratch_skip() { ..Default::default() }) .unwrap(); - assert_eq!(report.outcomes.len(), 2); - assert!(report.outcomes[0].qualified); - assert_eq!(report.outcomes[1].decision, TuneDecision::ScratchLimit); - assert_eq!(report.outcomes[1].scratch, None); + assert!(!report.time_budget_exhausted); + let qualified = report.outcomes.iter().filter(|o| o.qualified).count(); + assert!(qualified > 0); + assert!( + report + .outcomes + .iter() + .any(|o| o.decision == TuneDecision::ScratchLimit) + ); + for outcome in report.outcomes.iter().filter(|o| !o.qualified) { + assert_eq!(outcome.decision, TuneDecision::ScratchLimit); + assert_eq!(outcome.scratch, None); + } let stats = report.scratch.unwrap(); assert_eq!(stats.staging_allocations, 1); assert_eq!(stats.staging_releases, 1); - assert_eq!(stats.staging_reuses, 0); + assert_eq!(stats.staging_reuses, qualified - 1); assert_eq!(stats.retained_staging_bytes, 0); assert_eq!(stats.peak_bytes, 3 * 32 * 4096 * 4 + 32 * 32 * 4); assert!(report.final_cleanup.is_some()); @@ -562,7 +578,7 @@ fn tune_native_cooperative_f32() { }) .unwrap(); assert_eq!(report.visited_classes, 1, "{report:?}"); - assert_eq!(report.outcomes.len(), 2, "{report:?}"); + assert_eq!(report.outcomes.len(), 12, "{report:?}"); let native = report .outcomes .iter()