diff --git a/README.md b/README.md index 2a1b29ce5..e8ac8fec0 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ This repo does **not** contain the full llama.cpp tree (too large to mirror here ``` patches/ — apply these on top of the upstream fork modified-files/ — the exact modified files (drop-in replacements) +tools/ — patch/revert scripts and the rocprofv3 trace analyser tests/ — KIVI2 correctness tests BUILD.md — how to build for gfx90a in Docker ``` @@ -71,14 +72,251 @@ The recommended platform-agnostic alternative is the **Triton** implementation i --- +### 4. Chunked SSD Mamba-2 prefill on CDNA → [`patches/04-ssd-mamba2-prefill-cdna.patch`](patches/04-ssd-mamba2-prefill-cdna.patch) + +**+19% prompt processing** on hybrid Mamba-2 models. Measured on +NVIDIA-Nemotron-3-Super-120B-A12B (`i1-Q4_K_M`, 80 GiB) across 2× MI210. + +Upstream [PR #22675](https://github.com/ggml-org/llama.cpp/pull/22675) (merged +2026-07-28) replaces the sequential SSM scan with a chunked **State-Space +Duality** formulation: per chunk, the intra-chunk output and chunk-final state +become batched GEMMs (FP16 in, FP32 accumulate), leaving only a short scan over +`n_tok / 256` chunk boundaries. It is gated off for HIP: + +```c +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // kernels + dispatch +const bool use_ssd = ... && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING; +``` + +The PR author states the change "does not affect ... HIP" — a **scoping +decision, not a technical limitation**. Nobody had tried it on AMD. + +**Why CDNA2 wants this.** `rocprofv3` on a 4096-token prefill put +`ssm_scan_f32_group` at **22.8% of all GPU time** (1759 ms of 7703 ms, 160 calls +at 11 ms each) — the single largest kernel, scalar FP32, with zero matrix-core +use. It is not occupancy-limited (2048 blocks over 104 CUs), so the cost is the +work itself. The SSD path converts that into FP16 GEMMs, which on gfx90a is the +181 TFLOPS `v_mfma_f32_16x16x16f16` path rather than the 22.6 TFLOPS vector +path. Every cuBLAS symbol involved already has a hipBLAS alias in +`ggml-cuda/vendors/hip.h`, exercised today by `ggml_cuda_mul_mat_batched_cublas` +— so **the kernels compiled for HIP with no changes to the GEMM calls at all**. + +The patch admits HIP to both preprocessor guards and adds CDNA to the runtime +capability test, leaving NVIDIA's condition byte-for-byte unchanged: + +```c +&& ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) + || GGML_CUDA_CC_IS_CDNA(cc)) +``` + +Scoped to CDNA, **not** blanket AMD: RDNA uses WMMA with different tile shapes +and is unvalidated here. + +`USE_CUB` is deliberately left HIP-excluded; the file's shared-memory sequential +scan fallback is used instead, because hipCUB collides with ggml's own `__trap` +macro (the reason upstream PR #26388 stalled). **Watch out:** line 1 of +`ssm-scan.cu` *contains* the same guard text as the two that need editing, so +patch tooling must anchor to end-of-line or it will silently enable hipCUB too. + +| | before | after | +|---|---:|---:| +| pp4096 | 1366 t/s | **1624 t/s** (+18.9%) | +| pp16384 | 1775 t/s | **2114 t/s** (+19.1%) | +| total GPU kernel time | 7703 ms | 6538 ms | +| SSM scan share | 22.8% (1759 ms) | 1.0% (67 ms) | + +Total SSM cost including the new GEMM and mask kernels is ~540 ms, down from +1759 ms — a **3.2× reduction** on that component. `rocprofv3` after the change +confirms `ssm_ssd_pre_matmul_kernel<256, __half>` dispatching alongside new +`Cijk_*_MI16x16x16x1_*` kernels, i.e. the FP16 matrix cores really are in use. + +**Correctness:** `test-backend-ops -o SSM_SCAN` 7/7 on both MI210s, including +the multi-chunk shapes (`n_seq_tokens` = 256, 512, 300) that exercise this path +past its 128-token threshold. Additionally verified by reading generated tokens +at `temperature 0` — see "Verifying correctness" below. + +**Tunables** (compile-time `#define`s in `ssm-scan.cu`, no CLI flag): +`SSM_SSD_MIN_TOKENS` (128) and `SSM_SSD_CHUNK_SIZE` (256). The chunk size was +tuned against cuBLAS on NVIDIA; retuning it for rocBLAS on gfx90a is an obvious +next experiment, at one rebuild per value. + +**Files (2):** `ggml/src/ggml-cuda/ssm-scan.cu`, `tests/test-backend-ops.cpp`. + +--- + +### 5. Disable stream-k for K-quants on CDNA → [`patches/05-mmq-cdna-no-streamk.patch`](patches/05-mmq-cdna-no-streamk.patch) + +**+3.7% prompt processing** on MoE models. + +Every real `CASE` entry in `mmq-config-cdna.cuh` sets `stream_k = true` (8th +positional argument); only the unreachable `GGML_TYPE_COUNT` sentinel is false. +A single `mmq-config-cdna.cuh` covers all CDNA generations, so gfx90a inherits +tuning done elsewhere. + +Upstream [PR #26199](https://github.com/ggml-org/llama.cpp/pull/26199), which +retuned the RDNA configs, reports: *"I have also found that stream_k true helps +a lot for Dense models and hurts MoE models."* This workload is MoE, and after +change set 4 the MMQ kernels are **53.3%** of prefill GPU time. + +Stream-k splits the K dimension across more workgroups than there are output +tiles, then reconciles partial sums in `mul_mat_q_stream_k_fixup`. It exists to +fill a GPU that has too few output tiles to saturate it. In an MoE prefill each +expert already produces many tiles, so the fixup pass and the extra +global-memory traffic for partial accumulators are pure overhead. + +Scoped to the K-quants (Q2_K–Q6_K, 35 entries), which is what an `i1-Q4_K_M` +model dispatches; leaving other types alone keeps the result attributable. + +| | before | after | +|---|---:|---:| +| pp4096 | 1624 t/s | **1684 t/s** (+3.7%) | +| pp16384 | 2114 t/s | **2192 t/s** (+3.7%) | + +**Stated plainly:** the stream-k/MoE evidence is from RDNA3.5/RDNA4, *not* +CDNA2, and no CDNA stream-k benchmark exists upstream. This was an experiment +that happened to pay off, not the transfer of a known result. An earlier +experiment on the same theory — extending CDNA3's rocBLAS carve-out to CDNA2 — +came back **6.5% slower** and was discarded. + +**Correctness:** `test-backend-ops -o MUL_MAT_ID` 790/790 on both MI210s. + +**Files (1):** `ggml/src/ggml-cuda/mmq-config-cdna.cuh`. + +--- + +## Combined result (change sets 4 + 5) + +`llama-bench`, 2× MI210, Nemotron-3-Super-120B-A12B `i1-Q4_K_M`, +`-b 4096 -ub 2048 -fa 1 -sm layer -ctk q8_0 -ctv q8_0 -t 24 -r 2`: + +| build | pp4096 (t/s) | pp16384 (t/s) | vs base | +|---|---:|---:|---:| +| upstream `67b9b0e` | 1366 | 1775 | — | +| + SSD on CDNA | 1624.07 ± 0.81 | 2114.21 ± 0.73 | +19.1% | +| + stream-k off | 1683.57 ± 0.40 | **2192.26 ± 0.84** | **+23.5%** | + +For reference, vLLM with its AITER fast paths reaches 4,070 t/s at 16k on the +same hardware with an AWQ-INT4 Nemotron. These changes narrow the gap from +2.29× to 1.86×; they do not close it. + +--- + +## Verifying correctness + +Both change sets were gated on **reading generated tokens**, not only on the +benchmark number. This is not ceremony: a fast kernel emitting garbage already +cost this project one published benchmark. + +The specific risk in change set 4 is that the SSD path chains batched GEMMs with +`beta=1` accumulation for inter-chunk state propagation and materialises a +causal decay mask in a helper kernel. A wrong transpose flag, stride, or +alpha/beta does not crash — it propagates a subtly wrong SSM state and yields +fluent, confident, **wrong** text. + +Procedure: `llama-server` at `temperature 0` with a 240-token prompt (past the +128-token SSD threshold), asking the model to explain why prefill is +compute-bound and decode is bandwidth-bound. The answer had to use the figures +supplied in the prompt correctly — something a corrupted SSM state would not do. +Multi-request runs were done on a single card because of the pre-existing fault +below; single-request verification was done on both cards. + +--- + +## Known pre-existing issue: multi-GPU fault on sequential requests + +**Not caused by change sets 4 or 5**, but you will hit it, so it is recorded here. + +Running `llama-server` across both MI210s (`-sm layer`), the *second* sequential +request faults: + +``` +Memory access fault by GPU node-1 (Agent handle: 0x...) on address 0x... Reason: Unknown. +``` + +The first request returns correct output; the next one launches +(`slot launch_slot_: id 3 | task 263`) and the GPU faults. The process survives +and keeps answering `/health` with `ok`, so it looks alive while being unable to +serve — which makes it easy to misdiagnose. + +Isolated by bisecting configuration rather than assuming: + +| build | GPUs | result | +|---|---|---| +| upstream `67b9b0e`, unpatched | 2 | **faults on request 2** | +| + SSD | 2 | faults on request 2 | +| + SSD + stream-k off | 2 | faults on request 2 | +| + SSD | 1 (`ROCR_VISIBLE_DEVICES=0`, `-ngl 62`) | 3 sequential requests clean | + +Since the **unpatched baseline faults identically**, this is pre-existing on +this tree and unrelated to either change set. It is specific to the multi-GPU +split; the same build on a single card handles repeated requests without +incident. + +`llama-bench` does not surface it across many prefills, which points at state +reuse between requests rather than any prefill kernel — and means the throughput +numbers above are unaffected. + +Not root-caused. A good starting point is bisecting llama.cpp between `67b9b0e` +and current master with a two-request `llama-server` script on two cards. + +--- + +## Tested and rejected (change sets 4–5) + +Recorded so they are not re-attempted. + +| change | result | +|---|---| +| Extend CDNA3's rocBLAS carve-out to CDNA2 in `mmq.cu` (`ggml_cuda_should_use_mmq` true for Q4_K/Q5_K at any `ne11`) | **6.5% slower** (1277 vs 1366 t/s at pp4096). rocBLAS genuinely beats MMQ at `ne11=2048` on gfx90a. | +| `ROCBLAS_USE_HIPBLASLT=1` | no-op | +| `-sm row` | unsupported on ROCm | +| Crossing the attention MFMA gate by batching (`fattn.cu` needs `Q->ne[1] * gqa_ratio > 16`; Nemotron's ratio is exactly 16 at batch 1) | no gain; flash attention is only 0.9% of prefill | +| `-ub 4096` | slower than 2048 | +| `-ub 1024` | better at pp4096, worse at pp16384 | +| W4A8 weights | dead end on CDNA2 twice over: `mfma.i32.16x16x32.i8` (K=32) is MI300-only and fails to select on gfx90a, and CDNA2 gives INT8 and BF16 the *same* 181 TOPS peak, so there is no throughput to win | +| Making `ssm_scan_f32_group` wavefront-aware (it indexes lanes with the hardcoded `WARP_SIZE` 32 while gfx90a's wavefront is 64, so `warp_reduce_sum` takes a slow path once per token) | abandoned: `c_factor` serves double duty as warps-per-block *and* state-elements-per-lane, so changing it requires reworking the grid too or `state[]` reads out of bounds. Change set 4 makes it moot for prefill — the scalar kernel now only handles sequences under 128 tokens. | + +Confirmed already optimal at base, so not worth revisiting: MMQ already uses +int8 MFMA (`mfma_i32_16x16x16i8` via `AMD_MFMA_AVAILABLE`), CUDA graphs are +active (188 graph reuses per request), `GGML_HIP_MMQ_MFMA=ON`. + +--- + +## What is left + +MMQ is now 53.3% of prefill and is the obvious next target: + +1. **Retune `SSM_SSD_CHUNK_SIZE`** (256) for rocBLAS's kernel-selection sweet + spots on gfx90a. +2. **Extend the stream-k experiment** beyond K-quants, and sweep the other MMQ + config fields (`nthreads`, `occupancy`, `I`/`J` tile shape) for CDNA2 — one + shared CDNA config file is unlikely to suit all three CDNA generations. +3. **`mm_ids_helper`** was 4.5% at base; re-measure its share now. +4. Root-cause the multi-GPU fault above. + +--- + ## Base commit -All patches are generated against the `TheTom/llama-cpp-turboquant` fork at: +Change sets **1-3** are generated against the `TheTom/llama-cpp-turboquant` +fork at: ``` c26cbdffcf6fc9b7430cd6b117757e9a3f70b7ea Merge pull request #225 from TheTom/fix-ui-assets-partial-dist ``` +Change sets **4-5** are generated against **upstream `ggml-org/llama.cpp`** at: + +``` +67b9b0e7f6ce45d929a4411907d3c48ec719e81c llama-arch: fix DeepSeek4 APE tensor op (#25945) +``` + +These are different bases. Change sets 4-5 were developed and measured on the +upstream tree, **not** on the TurboQuant fork, and have not been tested there — +`ssm-scan.cu` in particular changed substantially upstream in the interim, so +expect `patches/04-*` to need rebasing before it applies to the fork. The two +groups touch disjoint files, so there is no conflict between them in principle. + ## How to apply See [`BUILD.md`](BUILD.md) for the full Docker build procedure with ccache. Short version: @@ -93,6 +331,37 @@ git apply 03-turboquant-wave64-fixes.patch # build for gfx90a (see BUILD.md) ``` +Change sets 4-5 target upstream llama.cpp instead (see "Base commit" above): + +```bash +git clone https://github.com/ggml-org/llama.cpp.git +cd llama.cpp +git checkout 67b9b0e7f6ce45d929a4411907d3c48ec719e81c +git apply patches/04-ssd-mamba2-prefill-cdna.patch +git apply patches/05-mmq-cdna-no-streamk.patch +cmake -B build -DGGML_HIP=ON -DAMDGPU_TARGETS=gfx90a -DGGML_HIP_MMQ_MFMA=ON \ + -DCMAKE_BUILD_TYPE=Release +cmake --build build --target llama-bench llama-server test-backend-ops -j +``` + +`patches/04-*` bundles the upstream SSD kernels together with the CDNA +enablement, so it applies to a bare `67b9b0e` checkout with no cherry-pick +first — verified with `git apply --check`, and the resulting files are +byte-identical to the `modified-files/` copies. If you would rather keep the +upstream work as its own commit, `git cherry-pick b62b350` instead and then +apply only the guard changes via `tools/patch_ssm_ssd_cdna.py`. + +Verify before trusting the build: + +```bash +./build/bin/test-backend-ops -o SSM_SCAN # expect 7/7 per device +./build/bin/test-backend-ops -o MUL_MAT_ID # expect 790/790 per device +``` + +`tools/patch_ssm_ssd_cdna.py` and `tools/patch_mmq_cdna_no_streamk.py` apply and +revert the same changes against a clean tree, with `--check` / `--revert`, and +refuse to apply if their anchors have moved upstream. + The [`modified-files/`](modified-files/) directory contains the final state of every changed file if you prefer drop-in replacement over `git apply`. --- diff --git a/modified-files/ggml/src/ggml-cuda/mmq-config-cdna.cuh b/modified-files/ggml/src/ggml-cuda/mmq-config-cdna.cuh new file mode 100644 index 000000000..9b713b227 --- /dev/null +++ b/modified-files/ggml/src/ggml-cuda/mmq-config-cdna.cuh @@ -0,0 +1,177 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/modified-files/ggml/src/ggml-cuda/ssm-scan.cu b/modified-files/ggml/src/ggml-cuda/ssm-scan.cu new file mode 100644 index 000000000..45f172c38 --- /dev/null +++ b/modified-files/ggml/src/ggml-cuda/ssm-scan.cu @@ -0,0 +1,853 @@ +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#define USE_CUB +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 + +#ifdef USE_CUB +#include +using namespace cub; +#endif // USE_CUB + +#include "ssm-scan.cuh" + + +// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. +// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). +#define SSM_SSD_MIN_TOKENS 128 + +// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items. +#define SSM_SSD_DT_BLOCK 256 +#define SSM_SSD_DT_MAX_ITEMS 32 + +// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity. +#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) + +// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. +#define SSM_SSD_CHUNK_SIZE 256 + +// We would like to keep pragma unroll for cases where L_template is not 0, +// so we suppress the clang transformation warning. +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wpass-failed" +#endif // __clang__ +template +__global__ void __launch_bounds__(splitD, 1) + ssm_scan_f32(const float * src0_ptr, const float * src1_ptr, const float * src2_ptr, + const float * src3_ptr, const float * src4_ptr, const float * src5_ptr, + const int32_t * src6_ptr, float * dst_ptr, + const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3, + const int src2_nb1, const int src2_nb2, const int src3_nb1, + const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, + const int64_t s_off, const int64_t d_inner, const int64_t L_param) +{ + const float * GGML_CUDA_RESTRICT src0 = src0_ptr; + const float * GGML_CUDA_RESTRICT src1 = src1_ptr; + const float * GGML_CUDA_RESTRICT src2 = src2_ptr; + const float * GGML_CUDA_RESTRICT src3 = src3_ptr; + const float * GGML_CUDA_RESTRICT src4 = src4_ptr; + const float * GGML_CUDA_RESTRICT src5 = src5_ptr; + const int32_t * GGML_CUDA_RESTRICT src6 = src6_ptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + const size_t L = L_template == 0 ? L_param : L_template; + ggml_cuda_pdl_sync(); + const float *s0_block = (const float *)((const char *)src0 + src6[blockIdx.x] * src0_nb3 + blockIdx.y * splitD * src0_nb2); + const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb3) + blockIdx.y * splitD * sizeof(float)); + const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); + const float *A_block = (const float *)((const char *)src3 + blockIdx.y * splitD * src3_nb1); + const float *B_block = (const float *)((const char *)src4 + (blockIdx.x * src4_nb3)); + const float *C_block = (const float *)((const char *)src5 + (blockIdx.x * src5_nb3)); + float *y_block = (float *)((char *)dst + (blockIdx.x * d_inner * L * sizeof(float)) + blockIdx.y * splitD * sizeof(float)); + float *s_block = (float *)((char *)dst + s_off + blockIdx.x * src0_nb3 + blockIdx.y * splitD * src0_nb2); + + const int stride_x = src1_nb2 / sizeof(float); + const int stride_dt = src2_nb1 / sizeof(float); + const int stride_B = src4_nb2 / sizeof(float); + const int stride_C = src5_nb2 / sizeof(float); + const int stride_y = d_inner; + + float regA[N]; + float regs0[N]; + + __shared__ float smemB[N]; + __shared__ float smemC[N]; + +#ifdef USE_CUB + using BlockLoad = cub::BlockLoad; + using BlockStore = cub::BlockStore; + + union CubTempStorage { + typename BlockLoad::TempStorage load_temp; + typename BlockStore::TempStorage store_temp; + }; + __shared__ CubTempStorage cub_temp_storage; + + BlockLoad(cub_temp_storage.load_temp).Load(A_block, regA); + __syncthreads(); + BlockLoad(cub_temp_storage.load_temp).Load(s0_block, regs0); +#else + const int stride_s0 = src0_nb2 / sizeof(float); + const int stride_A = src3_nb1 / sizeof(float); +#pragma unroll + for (size_t n = 0; n < N; ++n) + { + regA[n] = A_block[threadIdx.x * stride_A + n]; + regs0[n] = s0_block[threadIdx.x * stride_s0 + n]; + } +#endif + +#pragma unroll + for (size_t i = 0; i < L; i++) + { + if (threadIdx.x < N) + { + smemB[threadIdx.x] = B_block[i * stride_B + threadIdx.x]; + smemC[threadIdx.x] = C_block[i * stride_C + threadIdx.x]; + } + __syncthreads(); + + float dt_soft_plus = dt_block[i * stride_dt + threadIdx.x]; + if (dt_soft_plus <= 20.0f) + { + dt_soft_plus = log1pf(expf(dt_soft_plus)); + } + float x_dt = x_block[i * stride_x + threadIdx.x] * dt_soft_plus; + + float sumf = 0.0f; +#pragma unroll + for (size_t n = 0; n < N; n++) + { + float state = regs0[n] * expf(dt_soft_plus * regA[n]) + smemB[n] * x_dt; + sumf += state * smemC[n]; + regs0[n] = state; + } + y_block[i * stride_y + threadIdx.x] = sumf; + __syncthreads(); + } + +#ifdef USE_CUB + BlockStore(cub_temp_storage.store_temp).Store(s_block, regs0); +#else + const int stride_s = stride_s0; +#pragma unroll + for (size_t n = 0; n < N; ++n) + { + s_block[threadIdx.x * stride_s + n] = regs0[n]; + } +#endif +} +#ifdef __clang__ +#pragma clang diagnostic pop +#endif // __clang__ + +// assumes as many threads as d_state +template +__global__ void __launch_bounds__(d_state, 1) + ssm_scan_f32_group( + const float * src0_ptr, const float * src1_ptr, const float * src2_ptr, + const float * src3_ptr, const float * src4_ptr, const float * src5_ptr, + const int32_t * src6_ptr, float * dst_ptr, + const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3, + const int src2_nb1, const int src2_nb2, const int src3_nb1, + const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, + const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) { + const float * GGML_CUDA_RESTRICT src0 = src0_ptr; + const float * GGML_CUDA_RESTRICT src1 = src1_ptr; + const float * GGML_CUDA_RESTRICT src2 = src2_ptr; + const float * GGML_CUDA_RESTRICT src3 = src3_ptr; + const float * GGML_CUDA_RESTRICT src4 = src4_ptr; + const float * GGML_CUDA_RESTRICT src5 = src5_ptr; + const int32_t * GGML_CUDA_RESTRICT src6 = src6_ptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int warp_idx = blockIdx.x * c_factor + warp; + + const int head_idx = warp_idx / d_head; + const int head_off = (warp_idx % d_head) * sizeof(float); + const int seq_idx = blockIdx.y; + + const int group_off = (head_idx / (n_head / n_group)) * d_state * sizeof(float); + + ggml_cuda_pdl_sync(); + // TODO: refactor strides to be in elements/floats instead of bytes to be cleaner and consistent with the rest of the codebase + const float * s0_warp = (const float *) ((const char *) src0 + src6[seq_idx] * src0_nb3 + head_idx * src0_nb2 + head_off * d_state); + const float * x_warp = (const float *) ((const char *) src1 + (seq_idx * src1_nb3) + (warp_idx * sizeof(float))); + const float * dt_warp = (const float *) ((const char *) src2 + (seq_idx * src2_nb2) + head_idx * sizeof(float)); + const float * A_warp = (const float *) ((const char *) src3 + head_idx * src3_nb1); + const float * B_warp = (const float *) ((const char *) src4 + (seq_idx * src4_nb3) + (group_off)); + const float * C_warp = (const float *) ((const char *) src5 + (seq_idx * src5_nb3) + (group_off)); + float * y_warp = dst + (seq_idx * n_tok * n_head * d_head) + warp_idx; + float * s_warp = (float *) ((char *) dst + s_off + seq_idx * src0_nb3 + head_idx * src0_nb2 + head_off * d_state); + + // strides across n_seq_tokens + const int stride_x = src1_nb2 / sizeof(float); + const int stride_dt = src2_nb1 / sizeof(float); + const int stride_B = src4_nb2 / sizeof(float); + const int stride_C = src5_nb2 / sizeof(float); + const int stride_y = n_head * d_head; + + float state[c_factor]; + float state_sum = 0.0f; + +#pragma unroll + for (int j = 0; j < c_factor; j++) { + state[j] = s0_warp[WARP_SIZE * j + lane]; + } + + for (int64_t i = 0; i < n_tok; i++) { + // NOTE: dt_soft_plus, dA and x_dt have the same value for a warp here. + // Recalculation is intentional; sharing via shuffles/smem proved slower due to sync overhead. + const float dt_soft_plus = (dt_warp[i * stride_dt] <= 20.0f ? log1pf(expf(dt_warp[i * stride_dt])) : dt_warp[i * stride_dt]); + + state_sum = 0.0f; + const float dA = expf(dt_soft_plus * A_warp[0]); + const float x_dt = x_warp[i * stride_x] * dt_soft_plus; +#pragma unroll + for (int j = 0; j < c_factor; j++) { + const float B_val = B_warp[i * stride_B + WARP_SIZE * j + lane]; + const float C_val = C_warp[i * stride_C + WARP_SIZE * j + lane]; + state[j] = (state[j] * dA) + (B_val * x_dt); + state_sum += state[j] * C_val; + } + + // parallel accumulation for output + state_sum = warp_reduce_sum(state_sum); + + if (lane == 0) { + y_warp[i * stride_y] = state_sum; + } + } + + // write back the state +#pragma unroll + for (int j = 0; j < c_factor; j++) { + s_warp[WARP_SIZE * j + lane] = state[j]; + } +} + +static void ssm_scan_f32_cuda(const float * src0, const float * src1, const float * src2, const float * src3, + const float * src4, const float * src5, const int32_t * src6, float * dst, + const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3, const int src2_nb1, + const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, + const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim, + const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq, + cudaStream_t stream) { + // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition! + if (src3_nb1 == sizeof(float)) { + // Mamba-2 + if (d_state == 128) { + constexpr int threads = 128; + constexpr int num_warps = threads/WARP_SIZE; + + const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream); + ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1, + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok); + } else if (d_state == 256) { // Falcon-H1 + constexpr int threads = 256; + constexpr int num_warps = threads/WARP_SIZE; + + const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream); + ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1, + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok); + } else { + GGML_ABORT("doesn't support d_state!=(128 or 256)."); + } + } else { + // Mamba-1 + constexpr int threads = 128; + GGML_ASSERT(n_head % threads == 0); + GGML_ASSERT(head_dim == 1); + GGML_ASSERT(n_group == 1); + const dim3 blocks(n_seq, (n_head + threads - 1) / threads, 1); + if (d_state == 16) { + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream); + switch (n_tok) + { + case 1: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 2: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 3: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 4: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 5: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 6: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 7: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + case 8: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + default: + ggml_cuda_kernel_launch(ssm_scan_f32, launch_params, + src0, src1, src2, src3, src4, src5, src6, dst, + src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, + src3_nb1, src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, n_tok); + break; + } + } else { + GGML_ABORT("doesn't support d_state!=16."); + } + } +} + +// SSD_CDNA: HIP admitted here; the runtime `use_ssd` test below still +// restricts this to CDNA. MUSA remains excluded (untested). +#if !defined(GGML_USE_MUSA) +// ============================================================================ +// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) +// +// Instead of a sequential scan, SSD reformulates the output as: +// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init +// where L is a causal decay mask derived from A and dt. +// +// This converts the O(T*N) sequential scan into parallel matmuls. +// ============================================================================ +// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan. +// Grid: (n_head, n_seqs) +template +__global__ void ssm_ssd_prepare_dt_kernel( + const float * __restrict__ dt_raw, + float * __restrict__ dt_sp_out, + float * __restrict__ cs_out, + const int n_head, const int n_tok, + const int dt_stride_tok, // elements between tokens in dt + const int dt_stride_seq) { // elements between sequences in dt + + const int h = blockIdx.x; + const int s = blockIdx.y; + + const float * dt_seq = dt_raw + s * dt_stride_seq; + + float * dt_sp_seq = dt_sp_out + s * n_tok * n_head; + float * cs_seq = cs_out + s * n_tok * n_head; + + const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x). + // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads + // (stride n_head between threads vs. items_per_thread*n_head in blocked layout). + float local_vals[MAX_ITEMS]; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + float val = dt_seq[h + t * dt_stride_tok]; + float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; + local_vals[i] = sp; + dt_sp_seq[t * n_head + h] = sp; + } else { + local_vals[i] = 0.0f; + } + } + + // Phase 2+3: per-step inclusive scan to build cs[] in token order. + // With interleaved distribution the per-thread total scan would not give token-order + // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total. +#ifdef USE_CUB + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage scan_temp; + __shared__ float step_total; + + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + float inclusive; + BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive); + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + cs_seq[t * n_head + h] = running + inclusive; + } + if (threadIdx.x == BLOCK_SIZE - 1) { + step_total = inclusive; + } + __syncthreads(); + running += step_total; + } +#else + // Fallback: sequential prefix scan in shared memory, one slab at a time. + __shared__ float sdata[BLOCK_SIZE]; + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + sdata[threadIdx.x] = local_vals[i]; + __syncthreads(); + if (threadIdx.x == 0) { + for (int j = 1; j < BLOCK_SIZE; j++) { + sdata[j] += sdata[j - 1]; + } + } + __syncthreads(); + if (t < n_tok) { + cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; + } + running += sdata[BLOCK_SIZE - 1]; + __syncthreads(); + } +#endif +} + +// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. +// T_matmul controls precision for X_dt, B_weighted (float or half). +// C_scaled is always float (pairs with float s_cur in step 3c). +// Computation is always FP32; only the final store converts to T_matmul. +// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch). +// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_pre_matmul_kernel( + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt) + const float * __restrict__ A, // {1, n_head} + const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs} + const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs} + const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs} + T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest + T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end + float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) + const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} + half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs} + const int chunk_len, const int head_dim, const int n_head, const int n_group, + const int d_state, const int A_stride, + const int x_stride_tok, const int x_stride_seq, + const int B_stride_tok, const int B_stride_seq, + const int C_stride_tok, const int C_stride_seq, + const int chunk_offset, + const int n_tok_total) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int g = h / (n_head / n_group); + + const float A_h = A[h * A_stride]; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + + // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes. + const int n_xdt = chunk_len * head_dim; + if (idx < n_xdt) { + const int d = idx % head_dim; + const int t = idx / head_dim; + + const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim]; + const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h]; + + X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); + } + + // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len) + // and the same cs_t load, so merging halves the cs[] global memory traffic. + const int n_bw = d_state * chunk_len; + if (idx < n_bw) { + const int n = idx % d_state; + const int t = idx / d_state; + + const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; + + const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; + B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t))); + + const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; + C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t); + } + + // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. + const int n_M = chunk_len * chunk_len; + if (idx < n_M) { + const int t_out = idx % chunk_len; + const int t_in = idx / chunk_len; + + half val; + if (t_in <= t_out) { + const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base; + const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base; + const float decay = __expf(A_h * (cs_out - cs_in)); + const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group + + (int64_t)g * chunk_len * chunk_len; + const float cb_val = CB_g[t_out + t_in * chunk_len]; + val = __float2half(decay * cb_val); + } else { + val = __float2half(0.0f); + } + + M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val; + } +} + +// Scale running state in-place: s_cur *= decay_total(chunk). +// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. +// Eliminates the s_old buffer and D2D memcpy vs the old approach of: +// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old +// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_scale_state_kernel( + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ A, // {1, n_head} + const int d_state, const int head_dim, const int n_head, + const int chunk_offset, const int chunk_len, + const int n_tok_total, const int A_stride) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + const int state_per_head = d_state * head_dim; + if (idx >= state_per_head) return; + + const float A_h = A[h * A_stride]; + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + const float decay_total = __expf(A_h * cs_last); + + const int off = s * state_per_head * n_head + h * state_per_head + idx; + s_cur[off] *= decay_total; +} + +// Copy initial state from src0[ids[s]] into s_cur for each sequence. +// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) +template +__global__ void ssm_ssd_init_state_kernel( + const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs} + const int32_t * __restrict__ ids, // {n_seqs} + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const int state_size, // d_state * head_dim * n_head + const int64_t s0_stride_seq) { // elements between state rows + const int s = blockIdx.y; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + if (idx >= state_size) return; + + const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq; + s_cur[s * state_size + idx] = s_src[idx]; +} + +// SSD (State Space Duality) dispatch for Mamba-2 prefill. +// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt. +// All strides are in elements (floats), not bytes. +static void ssm_scan_ssd_f32_cuda( + ggml_backend_cuda_context & ctx, + const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, + const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, + const int64_t s0_stride_seq, // state (src0) stride between seqs + const int x_stride_tok, const int x_stride_seq, // x (src1) strides + const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides + const int A_stride, // A (src3) stride between heads + const int B_stride_tok, const int B_stride_seq, // B (src4) strides + const int C_stride_tok, const int C_stride_seq, // C (src5) strides + const int64_t s_off, const int64_t d_state, const int64_t head_dim, + const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) { + + cudaStream_t stream = ctx.stream(); + const int64_t d_inner = head_dim * n_head; + + const int64_t chunk_size = SSM_SSD_CHUNK_SIZE; + const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size; + + const int64_t state_per_head = d_state * head_dim; + + using matmul_t = half; + static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F; + + ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq); + ggml_cuda_pool_alloc X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq); + ggml_cuda_pool_alloc B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + ggml_cuda_pool_alloc C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + float * dt_sp = dt_sp_buf.get(); + float * cs = cs_buf.get(); + float * CB = CB_buf.get(); + matmul_t * X_dt = X_dt_buf.get(); + matmul_t * B_weighted = B_w_buf.get(); + float * C_scaled = C_s_buf.get(); + float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst + + // Step 1: softplus(dt) and parallel prefix sum over full sequence + { + dim3 grid(n_head, n_seq); + ssm_ssd_prepare_dt_kernel<<>>( + src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 2: initialize running state from src0[ids[s]] + { + constexpr int BLOCK = 256; + const int64_t state_size = d_state * head_dim * n_head; + dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq); + ssm_ssd_init_state_kernel<<>>( + src0_d, src6_d, s_cur, state_size, s0_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 3: chunked SSD loop + // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state + cublasHandle_t handle = ctx.cublas_handle(); + CUBLAS_CHECK(cublasSetStream(handle, stream)); + const float alpha_one = 1.0f; + const float beta_zero = 0.0f; + const float beta_one = 1.0f; + const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B + const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B + + // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size) + const int64_t n_M_max = chunk_size * chunk_size; + ggml_cuda_pool_alloc M_buf(ctx.pool(), n_M_max * n_head * n_seq); + half * M_mat = M_buf.get(); + + for (int64_t k = 0; k < n_chunks; k++) { + const int64_t chunk_offset = k * chunk_size; + const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); + + // 3a: CB = C^T @ B per group + for (int64_t s = 0; s < n_seq; s++) { + const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok; + const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok; + float * CB_s = CB + s * chunk_len * chunk_len * n_group; + + if (n_group == 1) { + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, + &beta_zero, CB_s, (int)chunk_len)); + } else { + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, + C_s, CUDA_R_32F, lda_C_src, d_state, + B_s, CUDA_R_32F, ldb_B_src, d_state, + &beta_zero, + CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len), + n_group, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix + const int64_t n_M = chunk_len * chunk_len; + { + constexpr int BLOCK = 256; + const int64_t n_xdt = chunk_len * head_dim; + const int64_t n_bw = d_state * chunk_len; + int64_t max_work = n_xdt; + if (n_bw > max_work) max_work = n_bw; + if (n_M > max_work) max_work = n_M; + dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_pre_matmul_kernel<<>>( + cs, dt_sp, src3_d, src1_d, src4_d, src5_d, + X_dt, B_weighted, C_scaled, + CB, M_mat, + chunk_len, head_dim, n_head, n_group, d_state, A_stride, + x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, + chunk_offset, n_tok); + CUDA_CHECK(cudaGetLastError()); + } + + // 3c: dst = S_cur^T @ C_scaled (state contribution) + { + const int64_t stride_S = state_per_head; + const int64_t stride_Cs = d_state * chunk_len; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + head_dim, chunk_len, d_state, + &alpha_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs, + &beta_zero, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result) + // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T. + { + const int64_t stride_M = n_M; + const int64_t stride_X_h = (int64_t)chunk_len * head_dim; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + head_dim, chunk_len, chunk_len, + &alpha_one, + X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h, + M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M, + &beta_one, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) + { + // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it + constexpr int BLOCK = 256; + dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_scale_state_kernel<<>>( + s_cur, cs, src3_d, + d_state, head_dim, n_head, + chunk_offset, chunk_len, n_tok, A_stride); + CUDA_CHECK(cudaGetLastError()); + + // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled) + const int64_t stride_Bw = d_state * chunk_len; + const int64_t stride_X = chunk_len * head_dim; + const int64_t stride_S = state_per_head; + + for (int64_t s = 0; s < n_seq; s++) { + // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd} + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + d_state, head_dim, chunk_len, + &alpha_one, + B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw, + X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X, + &beta_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + } +} +#endif // SSD_CDNA: !defined(GGML_USE_MUSA) + +void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const struct ggml_tensor * src0 = dst->src[0]; // s + const struct ggml_tensor * src1 = dst->src[1]; // x + const struct ggml_tensor * src2 = dst->src[2]; // dt + const struct ggml_tensor * src3 = dst->src[3]; // A + const struct ggml_tensor * src4 = dst->src[4]; // B + const struct ggml_tensor * src5 = dst->src[5]; // C + const struct ggml_tensor * src6 = dst->src[6]; // ids + + const int64_t nc = src0->ne[0]; // d_state + const int64_t nr = src0->ne[1]; // head_dim or 1 + const int64_t nh = src1->ne[1]; // n_head + const int64_t ng = src4->ne[1]; // n_group + const int64_t n_t = src1->ne[2]; // number of tokens per sequence + const int64_t n_s = src1->ne[3]; // number of sequences in the batch + + const int64_t s_off = ggml_nelements(src1) * sizeof(float); + + GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(src3->nb[0] == sizeof(float)); + GGML_ASSERT(src4->nb[0] == sizeof(float)); + GGML_ASSERT(src5->nb[0] == sizeof(float)); + GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); + + const float * src0_d = (const float *) src0->data; + const float * src1_d = (const float *) src1->data; + const float * src2_d = (const float *) src2->data; + const float * src3_d = (const float *) src3->data; + const float * src4_d = (const float *) src4->data; + const float * src5_d = (const float *) src5->data; + const int32_t * src6_d = (const int32_t *) src6->data; + float * dst_d = (float *) dst->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src6->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + // Byte strides are narrowed to int for both scan and SSD paths. + GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); + +// SSD_CDNA: HIP admitted here; the runtime `use_ssd` test below still +// restricts this to CDNA. MUSA remains excluded (untested). +#if !defined(GGML_USE_MUSA) + // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. + // Requires NVIDIA Turing+ otherwise fallback to scan. + const bool is_mamba2 = (src3->nb[1] == sizeof(float)); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && n_t <= SSM_SSD_MAX_TOKENS + && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) + // SSD_CDNA: CDNA has the FP16 matrix cores this path wants + // (v_mfma_f32_16x16x16f16) and full hipBLAS aliases for the + // batched GEMMs. Deliberately NOT all of AMD: RDNA's WMMA + // path is unvalidated here. + || GGML_CUDA_CC_IS_CDNA(cc)) + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment + + if (use_ssd) { + // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence, + // so src0 must be fully contiguous across all inner dimensions. + // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1]. + GGML_ASSERT(src0->nb[1] == nc * sizeof(float)); + GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float)); + + ssm_scan_ssd_f32_cuda(ctx, + src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + (int64_t)(src0->nb[3] / sizeof(float)), + (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)), + (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)), + (int)(src3->nb[1] / sizeof(float)), + (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)), + (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)), + s_off, nc, nr, nh, ng, n_t, n_s); + return; + } +#endif + ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], + src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], + s_off, nc, nr, nh, ng, n_t, n_s, stream); +} diff --git a/patches/04-ssd-mamba2-prefill-cdna.patch b/patches/04-ssd-mamba2-prefill-cdna.patch new file mode 100644 index 000000000..d54ac5e22 --- /dev/null +++ b/patches/04-ssd-mamba2-prefill-cdna.patch @@ -0,0 +1,587 @@ +diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu +index 3022249c7..45f172c38 100644 +--- a/ggml/src/ggml-cuda/ssm-scan.cu ++++ b/ggml/src/ggml-cuda/ssm-scan.cu +@@ -9,6 +9,21 @@ using namespace cub; + + #include "ssm-scan.cuh" + ++ ++// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. ++// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). ++#define SSM_SSD_MIN_TOKENS 128 ++ ++// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items. ++#define SSM_SSD_DT_BLOCK 256 ++#define SSM_SSD_DT_MAX_ITEMS 32 ++ ++// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity. ++#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) ++ ++// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. ++#define SSM_SSD_CHUNK_SIZE 256 ++ + // We would like to keep pragma unroll for cases where L_template is not 0, + // so we suppress the clang transformation warning. + #ifdef __clang__ +@@ -316,6 +331,431 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa + } + } + ++// SSD_CDNA: HIP admitted here; the runtime `use_ssd` test below still ++// restricts this to CDNA. MUSA remains excluded (untested). ++#if !defined(GGML_USE_MUSA) ++// ============================================================================ ++// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) ++// ++// Instead of a sequential scan, SSD reformulates the output as: ++// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init ++// where L is a causal decay mask derived from A and dt. ++// ++// This converts the O(T*N) sequential scan into parallel matmuls. ++// ============================================================================ ++// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan. ++// Grid: (n_head, n_seqs) ++template ++__global__ void ssm_ssd_prepare_dt_kernel( ++ const float * __restrict__ dt_raw, ++ float * __restrict__ dt_sp_out, ++ float * __restrict__ cs_out, ++ const int n_head, const int n_tok, ++ const int dt_stride_tok, // elements between tokens in dt ++ const int dt_stride_seq) { // elements between sequences in dt ++ ++ const int h = blockIdx.x; ++ const int s = blockIdx.y; ++ ++ const float * dt_seq = dt_raw + s * dt_stride_seq; ++ ++ float * dt_sp_seq = dt_sp_out + s * n_tok * n_head; ++ float * cs_seq = cs_out + s * n_tok * n_head; ++ ++ const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; ++ ++ // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x). ++ // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads ++ // (stride n_head between threads vs. items_per_thread*n_head in blocked layout). ++ float local_vals[MAX_ITEMS]; ++ for (int i = 0; i < items_per_thread; i++) { ++ const int t = i * BLOCK_SIZE + threadIdx.x; ++ if (t < n_tok) { ++ float val = dt_seq[h + t * dt_stride_tok]; ++ float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; ++ local_vals[i] = sp; ++ dt_sp_seq[t * n_head + h] = sp; ++ } else { ++ local_vals[i] = 0.0f; ++ } ++ } ++ ++ // Phase 2+3: per-step inclusive scan to build cs[] in token order. ++ // With interleaved distribution the per-thread total scan would not give token-order ++ // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total. ++#ifdef USE_CUB ++ using BlockScan = cub::BlockScan; ++ __shared__ typename BlockScan::TempStorage scan_temp; ++ __shared__ float step_total; ++ ++ float running = 0.0f; ++ for (int i = 0; i < items_per_thread; i++) { ++ float inclusive; ++ BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive); ++ const int t = i * BLOCK_SIZE + threadIdx.x; ++ if (t < n_tok) { ++ cs_seq[t * n_head + h] = running + inclusive; ++ } ++ if (threadIdx.x == BLOCK_SIZE - 1) { ++ step_total = inclusive; ++ } ++ __syncthreads(); ++ running += step_total; ++ } ++#else ++ // Fallback: sequential prefix scan in shared memory, one slab at a time. ++ __shared__ float sdata[BLOCK_SIZE]; ++ float running = 0.0f; ++ for (int i = 0; i < items_per_thread; i++) { ++ const int t = i * BLOCK_SIZE + threadIdx.x; ++ sdata[threadIdx.x] = local_vals[i]; ++ __syncthreads(); ++ if (threadIdx.x == 0) { ++ for (int j = 1; j < BLOCK_SIZE; j++) { ++ sdata[j] += sdata[j - 1]; ++ } ++ } ++ __syncthreads(); ++ if (t < n_tok) { ++ cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; ++ } ++ running += sdata[BLOCK_SIZE - 1]; ++ __syncthreads(); ++ } ++#endif ++} ++ ++// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. ++// T_matmul controls precision for X_dt, B_weighted (float or half). ++// C_scaled is always float (pairs with float s_cur in step 3c). ++// Computation is always FP32; only the final store converts to T_matmul. ++// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch). ++// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs) ++template ++__global__ void ssm_ssd_pre_matmul_kernel( ++ const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums ++ const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt) ++ const float * __restrict__ A, // {1, n_head} ++ const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs} ++ const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs} ++ const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs} ++ T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest ++ T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end ++ float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) ++ const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} ++ half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs} ++ const int chunk_len, const int head_dim, const int n_head, const int n_group, ++ const int d_state, const int A_stride, ++ const int x_stride_tok, const int x_stride_seq, ++ const int B_stride_tok, const int B_stride_seq, ++ const int C_stride_tok, const int C_stride_seq, ++ const int chunk_offset, ++ const int n_tok_total) { ++ ++ const int h = blockIdx.y; ++ const int s = blockIdx.z; ++ const int g = h / (n_head / n_group); ++ ++ const float A_h = A[h * A_stride]; ++ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; ++ ++ const int cs_seq_off = s * n_tok_total * n_head; ++ const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; ++ const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; ++ ++ // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes. ++ const int n_xdt = chunk_len * head_dim; ++ if (idx < n_xdt) { ++ const int d = idx % head_dim; ++ const int t = idx / head_dim; ++ ++ const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim]; ++ const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h]; ++ ++ X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); ++ } ++ ++ // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len) ++ // and the same cs_t load, so merging halves the cs[] global memory traffic. ++ const int n_bw = d_state * chunk_len; ++ if (idx < n_bw) { ++ const int n = idx % d_state; ++ const int t = idx / d_state; ++ ++ const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; ++ ++ const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; ++ B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t))); ++ ++ const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; ++ C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t); ++ } ++ ++ // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. ++ const int n_M = chunk_len * chunk_len; ++ if (idx < n_M) { ++ const int t_out = idx % chunk_len; ++ const int t_in = idx / chunk_len; ++ ++ half val; ++ if (t_in <= t_out) { ++ const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base; ++ const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base; ++ const float decay = __expf(A_h * (cs_out - cs_in)); ++ const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group ++ + (int64_t)g * chunk_len * chunk_len; ++ const float cb_val = CB_g[t_out + t_in * chunk_len]; ++ val = __float2half(decay * cb_val); ++ } else { ++ val = __float2half(0.0f); ++ } ++ ++ M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val; ++ } ++} ++ ++// Scale running state in-place: s_cur *= decay_total(chunk). ++// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. ++// Eliminates the s_old buffer and D2D memcpy vs the old approach of: ++// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old ++// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) ++template ++__global__ void ssm_ssd_scale_state_kernel( ++ float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} ++ const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums ++ const float * __restrict__ A, // {1, n_head} ++ const int d_state, const int head_dim, const int n_head, ++ const int chunk_offset, const int chunk_len, ++ const int n_tok_total, const int A_stride) { ++ ++ const int h = blockIdx.y; ++ const int s = blockIdx.z; ++ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; ++ const int state_per_head = d_state * head_dim; ++ if (idx >= state_per_head) return; ++ ++ const float A_h = A[h * A_stride]; ++ const int cs_seq_off = s * n_tok_total * n_head; ++ const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; ++ const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; ++ const float decay_total = __expf(A_h * cs_last); ++ ++ const int off = s * state_per_head * n_head + h * state_per_head + idx; ++ s_cur[off] *= decay_total; ++} ++ ++// Copy initial state from src0[ids[s]] into s_cur for each sequence. ++// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) ++template ++__global__ void ssm_ssd_init_state_kernel( ++ const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs} ++ const int32_t * __restrict__ ids, // {n_seqs} ++ float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} ++ const int state_size, // d_state * head_dim * n_head ++ const int64_t s0_stride_seq) { // elements between state rows ++ const int s = blockIdx.y; ++ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; ++ if (idx >= state_size) return; ++ ++ const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq; ++ s_cur[s * state_size + idx] = s_src[idx]; ++} ++ ++// SSD (State Space Duality) dispatch for Mamba-2 prefill. ++// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt. ++// All strides are in elements (floats), not bytes. ++static void ssm_scan_ssd_f32_cuda( ++ ggml_backend_cuda_context & ctx, ++ const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, ++ const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, ++ const int64_t s0_stride_seq, // state (src0) stride between seqs ++ const int x_stride_tok, const int x_stride_seq, // x (src1) strides ++ const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides ++ const int A_stride, // A (src3) stride between heads ++ const int B_stride_tok, const int B_stride_seq, // B (src4) strides ++ const int C_stride_tok, const int C_stride_seq, // C (src5) strides ++ const int64_t s_off, const int64_t d_state, const int64_t head_dim, ++ const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) { ++ ++ cudaStream_t stream = ctx.stream(); ++ const int64_t d_inner = head_dim * n_head; ++ ++ const int64_t chunk_size = SSM_SSD_CHUNK_SIZE; ++ const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size; ++ ++ const int64_t state_per_head = d_state * head_dim; ++ ++ using matmul_t = half; ++ static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F; ++ ++ ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); ++ ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); ++ ggml_cuda_pool_alloc CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq); ++ ggml_cuda_pool_alloc X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq); ++ ggml_cuda_pool_alloc B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); ++ ggml_cuda_pool_alloc C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); ++ float * dt_sp = dt_sp_buf.get(); ++ float * cs = cs_buf.get(); ++ float * CB = CB_buf.get(); ++ matmul_t * X_dt = X_dt_buf.get(); ++ matmul_t * B_weighted = B_w_buf.get(); ++ float * C_scaled = C_s_buf.get(); ++ float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst ++ ++ // Step 1: softplus(dt) and parallel prefix sum over full sequence ++ { ++ dim3 grid(n_head, n_seq); ++ ssm_ssd_prepare_dt_kernel<<>>( ++ src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); ++ CUDA_CHECK(cudaGetLastError()); ++ } ++ ++ // Step 2: initialize running state from src0[ids[s]] ++ { ++ constexpr int BLOCK = 256; ++ const int64_t state_size = d_state * head_dim * n_head; ++ dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq); ++ ssm_ssd_init_state_kernel<<>>( ++ src0_d, src6_d, s_cur, state_size, s0_stride_seq); ++ CUDA_CHECK(cudaGetLastError()); ++ } ++ ++ // Step 3: chunked SSD loop ++ // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state ++ cublasHandle_t handle = ctx.cublas_handle(); ++ CUBLAS_CHECK(cublasSetStream(handle, stream)); ++ const float alpha_one = 1.0f; ++ const float beta_zero = 0.0f; ++ const float beta_one = 1.0f; ++ const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B ++ const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B ++ ++ // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size) ++ const int64_t n_M_max = chunk_size * chunk_size; ++ ggml_cuda_pool_alloc M_buf(ctx.pool(), n_M_max * n_head * n_seq); ++ half * M_mat = M_buf.get(); ++ ++ for (int64_t k = 0; k < n_chunks; k++) { ++ const int64_t chunk_offset = k * chunk_size; ++ const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); ++ ++ // 3a: CB = C^T @ B per group ++ for (int64_t s = 0; s < n_seq; s++) { ++ const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok; ++ const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok; ++ float * CB_s = CB + s * chunk_len * chunk_len * n_group; ++ ++ if (n_group == 1) { ++ CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, ++ chunk_len, chunk_len, d_state, ++ &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, ++ &beta_zero, CB_s, (int)chunk_len)); ++ } else { ++ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, ++ chunk_len, chunk_len, d_state, ++ &alpha_one, ++ C_s, CUDA_R_32F, lda_C_src, d_state, ++ B_s, CUDA_R_32F, ldb_B_src, d_state, ++ &beta_zero, ++ CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len), ++ n_group, ++ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); ++ } ++ } ++ ++ // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix ++ const int64_t n_M = chunk_len * chunk_len; ++ { ++ constexpr int BLOCK = 256; ++ const int64_t n_xdt = chunk_len * head_dim; ++ const int64_t n_bw = d_state * chunk_len; ++ int64_t max_work = n_xdt; ++ if (n_bw > max_work) max_work = n_bw; ++ if (n_M > max_work) max_work = n_M; ++ dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); ++ ssm_ssd_pre_matmul_kernel<<>>( ++ cs, dt_sp, src3_d, src1_d, src4_d, src5_d, ++ X_dt, B_weighted, C_scaled, ++ CB, M_mat, ++ chunk_len, head_dim, n_head, n_group, d_state, A_stride, ++ x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, ++ chunk_offset, n_tok); ++ CUDA_CHECK(cudaGetLastError()); ++ } ++ ++ // 3c: dst = S_cur^T @ C_scaled (state contribution) ++ { ++ const int64_t stride_S = state_per_head; ++ const int64_t stride_Cs = d_state * chunk_len; ++ ++ for (int64_t s = 0; s < n_seq; s++) { ++ float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; ++ ++ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, ++ head_dim, chunk_len, d_state, ++ &alpha_one, ++ s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, ++ C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs, ++ &beta_zero, ++ dst_chunk, CUDA_R_32F, d_inner, head_dim, ++ n_head, ++ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); ++ } ++ } ++ ++ // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result) ++ // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T. ++ { ++ const int64_t stride_M = n_M; ++ const int64_t stride_X_h = (int64_t)chunk_len * head_dim; ++ ++ for (int64_t s = 0; s < n_seq; s++) { ++ float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; ++ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, ++ head_dim, chunk_len, chunk_len, ++ &alpha_one, ++ X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h, ++ M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M, ++ &beta_one, ++ dst_chunk, CUDA_R_32F, d_inner, head_dim, ++ n_head, ++ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); ++ } ++ } ++ ++ // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) ++ { ++ // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it ++ constexpr int BLOCK = 256; ++ dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq); ++ ssm_ssd_scale_state_kernel<<>>( ++ s_cur, cs, src3_d, ++ d_state, head_dim, n_head, ++ chunk_offset, chunk_len, n_tok, A_stride); ++ CUDA_CHECK(cudaGetLastError()); ++ ++ // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled) ++ const int64_t stride_Bw = d_state * chunk_len; ++ const int64_t stride_X = chunk_len * head_dim; ++ const int64_t stride_S = state_per_head; ++ ++ for (int64_t s = 0; s < n_seq; s++) { ++ // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd} ++ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, ++ d_state, head_dim, chunk_len, ++ &alpha_one, ++ B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw, ++ X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X, ++ &beta_one, ++ s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, ++ n_head, ++ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); ++ } ++ } ++ } ++} ++#endif // SSD_CDNA: !defined(GGML_USE_MUSA) ++ + void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const struct ggml_tensor * src0 = dst->src[0]; // s + const struct ggml_tensor * src1 = dst->src[1]; // x +@@ -357,6 +797,55 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_ASSERT(src6->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + ++ // Byte strides are narrowed to int for both scan and SSD paths. ++ GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX); ++ GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX); ++ GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX); ++ GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX); ++ GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX); ++ GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX); ++ GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX); ++ GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX); ++ GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX); ++ GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); ++ GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); ++ ++// SSD_CDNA: HIP admitted here; the runtime `use_ssd` test below still ++// restricts this to CDNA. MUSA remains excluded (untested). ++#if !defined(GGML_USE_MUSA) ++ // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. ++ // Requires NVIDIA Turing+ otherwise fallback to scan. ++ const bool is_mamba2 = (src3->nb[1] == sizeof(float)); ++ const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; ++ const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS ++ && n_t <= SSM_SSD_MAX_TOKENS ++ && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) ++ // SSD_CDNA: CDNA has the FP16 matrix cores this path wants ++ // (v_mfma_f32_16x16x16f16) and full hipBLAS aliases for the ++ // batched GEMMs. Deliberately NOT all of AMD: RDNA's WMMA ++ // path is unvalidated here. ++ || GGML_CUDA_CC_IS_CDNA(cc)) ++ && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment ++ ++ if (use_ssd) { ++ // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence, ++ // so src0 must be fully contiguous across all inner dimensions. ++ // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1]. ++ GGML_ASSERT(src0->nb[1] == nc * sizeof(float)); ++ GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float)); ++ ++ ssm_scan_ssd_f32_cuda(ctx, ++ src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, ++ (int64_t)(src0->nb[3] / sizeof(float)), ++ (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)), ++ (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)), ++ (int)(src3->nb[1] / sizeof(float)), ++ (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)), ++ (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)), ++ s_off, nc, nr, nh, ng, n_t, n_s); ++ return; ++ } ++#endif + ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], + src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index e7cd6d0cb..9d6ebd963 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -4000,7 +4000,7 @@ struct test_ssm_scan : public test_case { + + test_ssm_scan(ggml_type type = GGML_TYPE_F32, + int64_t d_state = 32, +- int64_t head_dim = 1, // non-zero for Mamba-2 ++ int64_t head_dim = 1, // 1 = Mamba-1; > 1 = Mamba-2 (scalar A per head) + int64_t n_head = 32, + int64_t n_group = 1, + int64_t n_seq_tokens = 32, +@@ -4008,6 +4008,11 @@ struct test_ssm_scan : public test_case { + bool xbc_overlap = false) + : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {} + ++ double max_nmse_err() override { ++ // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32. ++ return (head_dim > 1) ? 2e-7 : 1e-7; ++ } ++ + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs); + ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs); +@@ -4034,14 +4039,14 @@ struct test_ssm_scan : public test_case { + return out; + } + +- // similar to test_mul_mat_id ++ + void initialize_tensors(ggml_context * ctx) override { + std::random_device rd; + std::default_random_engine rng(rd()); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + if (ggml_is_view_op(t->op)) { continue; } +- // ids ++ // ids: permutation of [0..n_seqs) + for (int64_t r = 0; r < ggml_nrows(t); r++) { + std::vector data(t->ne[0]); + for (int i = 0; i < t->ne[0]; i++) { +@@ -4050,6 +4055,11 @@ struct test_ssm_scan : public test_case { + std::shuffle(data.begin(), data.end(), rng); + ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); + } ++ } else if (ggml_is_view_op(t->op)) { ++ continue; ++ } else if (t->ne[1] == n_head && t->ne[2] == 1) { ++ // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors) ++ init_tensor_uniform(t, -1.0f, -0.5f); + } else { + init_tensor_uniform(t); + } +@@ -8770,6 +8780,9 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap ++ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path ++ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks) ++ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs) + + test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); + test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); +@@ -9974,6 +9987,8 @@ static std::vector> make_test_cases_perf() { + test_cases.emplace_back(new test_ssm_conv_bias_silu(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1}, true)); // generate + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate ++ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B prefill ++ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 1, 1)); // Nemotron-9B generate + + // acc + test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); diff --git a/patches/05-mmq-cdna-no-streamk.patch b/patches/05-mmq-cdna-no-streamk.patch new file mode 100644 index 000000000..a0dc89965 --- /dev/null +++ b/patches/05-mmq-cdna-no-streamk.patch @@ -0,0 +1,89 @@ +diff --git a/ggml/src/ggml-cuda/mmq-config-cdna.cuh b/ggml/src/ggml-cuda/mmq-config-cdna.cuh +index 46ec6aa9d..9b713b227 100644 +--- a/ggml/src/ggml-cuda/mmq-config-cdna.cuh ++++ b/ggml/src/ggml-cuda/mmq-config-cdna.cuh +@@ -49,45 +49,45 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf + + // --------------------------------------------------------------------------------------------- + +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false); +- +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false); +- +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); +- +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); +- CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false); ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK ++ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // MI210_NO_STREAMK + + // --------------------------------------------------------------------------------------------- + diff --git a/tools/analyze_prefill_trace.py b/tools/analyze_prefill_trace.py new file mode 100644 index 000000000..8b92ab1e8 --- /dev/null +++ b/tools/analyze_prefill_trace.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""Break down GPU kernel time from a rocprofv3 kernel trace. + +Answers the only question that matters for optimisation work: which kernels +actually consume the wall clock. Guessing at this is how you spend a week +speeding up something that was 3% of the runtime. +""" +import sqlite3 +import sys + +db = sys.argv[1] if len(sys.argv) > 1 else "/hosttmp/pf/pf_results.db" +con = sqlite3.connect(db) + +tables = [r[0] for r in con.execute( + "SELECT name FROM sqlite_master WHERE type='table'").fetchall()] +sym = next((t for t in tables if "kernel_symbol" in t), None) +disp = next((t for t in tables if "kernel_dispatch" in t), None) +if not sym or not disp: + print("no kernel tables in", db) + print("tables:", tables[:10]) + raise SystemExit(1) + +rows = con.execute(f""" + SELECT s.display_name, COUNT(*), SUM(d.end - d.start) / 1e6 + FROM {disp} d JOIN {sym} s ON d.kernel_id = s.id + GROUP BY s.display_name + ORDER BY 3 DESC +""").fetchall() + +total = sum(r[2] for r in rows) +print(f"total GPU kernel time: {total:.0f} ms across {sum(r[1] for r in rows)} dispatches") +print() +print(f"{'%':>6} {'ms':>9} {'calls':>8} kernel") +print("-" * 78) +for name, n, ms in rows[:18]: + print(f"{ms/total*100:>5.1f}% {ms:>9.0f} {n:>8} {name[:56]}") + +# Group by what the kernel actually is, so the fix target is obvious. +print() +print("by category:") +cats = { + "SSM / Mamba scan": ("ssm_scan", "ssm_conv", "selective"), + "quantized GEMM (MMQ)": ("mul_mat_q", "mmq", "vec_dot"), + "flash attention": ("flash_attn", "fattn"), + "dequant / convert": ("dequantize", "cpy", "convert", "quantize"), + "norms / activations": ("rms_norm", "norm", "silu", "glu", "soft_max"), + "MoE routing": ("mul_mat_id", "argsort", "top_k", "moe"), +} +seen = set() +for label, keys in cats.items(): + ms = sum(r[2] for r in rows if any(k in r[0].lower() for k in keys)) + n = sum(r[1] for r in rows if any(k in r[0].lower() for k in keys)) + seen |= {r[0] for r in rows if any(k in r[0].lower() for k in keys)} + if ms: + print(f" {label:<24} {ms/total*100:>5.1f}% {ms:>8.0f} ms {n:>7} calls") +other = sum(r[2] for r in rows if r[0] not in seen) +if other: + print(f" {'other':<24} {other/total*100:>5.1f}% {other:>8.0f} ms") diff --git a/tools/patch_mmq_cdna_no_streamk.py b/tools/patch_mmq_cdna_no_streamk.py new file mode 100644 index 000000000..b69195c72 --- /dev/null +++ b/tools/patch_mmq_cdna_no_streamk.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Disable stream-k decomposition for K-quants in llama.cpp's CDNA MMQ config. + +WHY. Every real CASE entry in ggml-cuda/mmq-config-cdna.cuh sets stream_k=true +(the 8th positional argument); only the unreachable GGML_TYPE_COUNT sentinel is +false. There is a single mmq-config-cdna.cuh covering all CDNA generations, so +gfx90a inherits whatever was tuned elsewhere. + +Upstream PR #26199 (merged 2026-07-29), which retuned the RDNA configs, reports +directly: "I have also found that stream_k true helps a lot for Dense models and +hurts MoE models." Our workload is MoE -- Nemotron-3-Super-120B-A12B -- and MMQ +is 45.2% of prefill GPU time in the rocprofv3 profile (of 7703 ms total). That +makes this the largest single lever identified, ahead of the 22.8% SSM scan. + +Stream-k splits the K dimension across more workgroups than there are output +tiles, then fixes up partial sums in a second pass (mul_mat_q_stream_k_fixup). +It exists to fill the GPU when there are too few output tiles to saturate it. In +an MoE prefill each expert already produces many tiles, so the machine is +saturated without it and the fixup pass plus the extra global-memory traffic for +partial accumulators is pure overhead. + +IMPORTANT -- THIS IS AN EXPERIMENT, NOT A KNOWN WIN. The stream-k/MoE evidence is +from RDNA3.5/RDNA4, NOT CDNA2. No CDNA stream-k benchmark exists upstream. It may +well be slower on gfx90a, exactly as the earlier CDNA2 rocBLAS carve-out patch +turned out 6.5% slower than the assumption behind it. A/B against the same +binary, and read generated tokens both ways. + +Scoped to the K-quants (Q2_K..Q6_K) because that is what an i1-Q4_K_M model +actually dispatches; leaving the other types alone keeps the experiment narrow +and the result attributable. + + python patch_mmq_cdna_no_streamk.py [--check] [--revert] +""" +import re +import sys + +TARGET = "ggml/src/ggml-cuda/mmq-config-cdna.cuh" +TYPES = ("Q2_K", "Q3_K", "Q4_K", "Q5_K", "Q6_K") +MARKER = "MI210_NO_STREAMK" + +# CASE(type, nthreads, occupancy, I, J, sram_layout, K_vram, stream_k, fallback) +# Rewrite only the 8th argument, and only on lines for the listed types. Anchor +# on the two trailing args so a config whose shape changes upstream fails to +# match instead of silently corrupting a different field. +LINE_RE = re.compile( + r"^(\s*CASE\(GGML_TYPE_(?:" + "|".join(TYPES) + r"),[^;]*?,\s*)true(\s*,\s*(?:true|false)\s*\);)$" +) + + +def main() -> int: + check = "--check" in sys.argv + revert = "--revert" in sys.argv + + with open(TARGET) as f: + lines = f.readlines() + patched = any(MARKER in ln for ln in lines) + + if check: + n = sum(1 for ln in lines if MARKER in ln) + print(f"{'PATCHED' if patched else 'not patched'} {TARGET} ({n} entries)") + return 0 if patched else 1 + + if revert: + if not patched: + print("not patched; nothing to revert") + return 0 + out = [] + for ln in lines: + if MARKER in ln: + ln = ln.split(" // " + MARKER)[0] + "\n" + ln = re.sub(r",\s*false(\s*,\s*(?:true|false)\s*\);)$", r", true\1", ln) + out.append(ln) + with open(TARGET, "w") as f: + f.writelines(out) + print("reverted") + return 0 + + if patched: + print("already patched") + return 0 + + out, n = [], 0 + for ln in lines: + m = LINE_RE.match(ln.rstrip("\n")) + if m: + out.append(f"{m.group(1)}false{m.group(2)} // {MARKER}\n") + n += 1 + else: + out.append(ln) + + if n == 0: + print("ERROR: no CASE lines matched -- the config layout changed upstream; " + "re-derive rather than forcing.", file=sys.stderr) + return 1 + + with open(TARGET, "w") as f: + f.writelines(out) + print(f"patched: stream_k disabled on {n} CDNA K-quant MMQ configs") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/patch_ssm_ssd_cdna.py b/tools/patch_ssm_ssd_cdna.py new file mode 100644 index 000000000..ba027947b --- /dev/null +++ b/tools/patch_ssm_ssd_cdna.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Enable llama.cpp's chunked-SSD Mamba-2 prefill path on CDNA (gfx90a / MI210). + +BACKGROUND. Upstream commit b62b350 (PR #22675, merged 2026-07-28) replaces the +sequential SSM scan with a chunked State-Space-Duality formulation: per-chunk +intra-chunk output and chunk-final-state become batched GEMMs, leaving only a +short scan over n_tok/256 chunk boundaries. The heavy GEMMs run FP16-in / +FP32-accumulate via cublasGemmStridedBatchedEx. + +It is gated off for HIP by `#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)` +plus a runtime `GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING`. The PR +author states the change "does not affect ... HIP" -- a scoping decision, not a +finding that it cannot work on AMD. + +WHY IT SHOULD WORK ON CDNA2. Profiling pp4096 of Nemotron-3-Super-120B (Q4_K_M) +on 2x MI210 puts ssm_scan_f32_group at 22.8% of GPU time (1759 ms / 7703 ms), and +that kernel is scalar FP32 with zero matrix-core use. The SSD path converts that +work into FP16 GEMMs, which is gfx90a's 181 TFLOPS v_mfma_f32_16x16x16f16 path +rather than its 22.6 TFLOPS vector path. Every cuBLAS symbol the kernel uses +(cublasGemmStridedBatchedEx, cublasSgemm, cublasSetStream, CUBLAS_OP_*, +CUDA_R_16F/32F, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT) already has a hipBLAS +alias in ggml-cuda/vendors/hip.h, and ggml exercises those same aliases today in +ggml_cuda_mul_mat_batched_cublas. CUB is not required: USE_CUB stays undefined on +HIP and the file ships a working shared-memory sequential-scan fallback. + +SCOPE. Gated to CDNA specifically, not blanket AMD. RDNA's matrix cores behave +differently (WMMA, different tile shapes) and are entirely unvalidated here; +opening this up for all AMD would be shipping an untested path to other people's +hardware. + +CORRECTNESS IS THE WHOLE RISK. The SSD path chains batched GEMMs with beta=1 +accumulation for inter-chunk state propagation and materializes a causal decay +mask in a helper kernel. A wrong transpose flag, stride, or alpha/beta does not +crash -- it propagates a subtly wrong SSM state and yields fluent, confident, +WRONG text. Gate on `test-backend-ops -o SSM_SCAN` (which carries the +Nemotron-9B-shaped multi-chunk cases added by the same commit) BEFORE trusting +any throughput number, then read real generated tokens. + + python patch_ssm_ssd_cdna.py [--check] [--revert] +""" +import sys + +TARGET = "ggml/src/ggml-cuda/ssm-scan.cu" + +# The two `#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)` guards that +# wrap (a) the SSD kernel definitions and (b) the dispatch site. MUSA stays +# excluded -- it is untested here and not ours to enable. +# +# The trailing newline is load-bearing. Line 1 of the file is +# #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +# which *contains* this text as a substring and controls USE_CUB. Matching it +# would enable hipCUB, whose headers collide with ggml's own __trap macro (the +# reason upstream PR #26388 stalled). Anchoring to end-of-line excludes it. +GUARD_ORIG = "#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)\n" +GUARD_NEW = ( + "// SSD_CDNA: HIP admitted here; the runtime `use_ssd` test below still\n" + "// restricts this to CDNA. MUSA remains excluded (untested).\n" + "#if !defined(GGML_USE_MUSA)\n" +) + +# Matching #endif trailer comments, so the file still reads correctly. Same +# substring hazard as above -- the USE_CUB #endif shares this prefix. +ENDIF_ORIG = "#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)\n" +ENDIF_NEW = "#endif // SSD_CDNA: !defined(GGML_USE_MUSA)\n" + +# The runtime capability test. Keep NVIDIA's condition exactly as upstream has +# it and add CDNA alongside, so this cannot change behaviour on NVIDIA. +COND_ORIG = """ && GGML_CUDA_CC_IS_NVIDIA(cc) + && cc >= GGML_CUDA_CC_TURING + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment""" + +COND_NEW = """ && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) + // SSD_CDNA: CDNA has the FP16 matrix cores this path wants + // (v_mfma_f32_16x16x16f16) and full hipBLAS aliases for the + // batched GEMMs. Deliberately NOT all of AMD: RDNA's WMMA + // path is unvalidated here. + || GGML_CUDA_CC_IS_CDNA(cc)) + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment""" + +EDITS = [ + ("HIP guards", GUARD_ORIG, GUARD_NEW, 2), + # Only the kernel block's #endif carries a trailer comment; the dispatch + # block closes with a bare `#endif`, which stays balanced and needs no edit. + ("endif trailers", ENDIF_ORIG, ENDIF_NEW, 1), + ("use_ssd capability test", COND_ORIG, COND_NEW, 1), +] +MARKER = "SSD_CDNA" + + +def main() -> int: + check = "--check" in sys.argv + revert = "--revert" in sys.argv + + with open(TARGET) as f: + src = f.read() + patched = MARKER in src + + if check: + print(f"{'PATCHED' if patched else 'not patched'} {TARGET}") + return 0 if patched else 1 + + if revert: + if not patched: + print("not patched; nothing to revert") + return 0 + for _, orig, new, _ in EDITS: + src = src.replace(new, orig) + with open(TARGET, "w") as f: + f.write(src) + print("reverted") + return 0 + + if patched: + print("already patched") + return 0 + + # Verify every anchor's occurrence count up front. A half-applied set of + # preprocessor guards produces an unbalanced #if/#endif and a wall of + # confusing compiler errors far from the real cause. + for name, orig, _, want in EDITS: + n = src.count(orig) + if n != want: + print(f"ERROR: anchor '{name}' matched {n} times, expected {want}. " + "Upstream moved; re-derive rather than forcing.", file=sys.stderr) + return 1 + for name, orig, new, _ in EDITS: + src = src.replace(orig, new) + with open(TARGET, "w") as f: + f.write(src) + print("patched: chunked-SSD Mamba-2 prefill now enabled on CDNA") + print("NEXT: test-backend-ops -o SSM_SCAN must pass before any benchmark.") + return 0 + + +if __name__ == "__main__": + sys.exit(main())