Text Generation
GGUF
English
multilingual
ternary
speculative-decoding
mtp
vision
multimodal
custom-quant
qwen35
code
reasoning
agentic
conversational
Instructions to use talxcc/Tals-coder-flash-02 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- llama.cpp
How to use talxcc/Tals-coder-flash-02 with llama.cpp:
Install (macOS, Linux)
curl -LsSf https://llama.app/install.sh | sh # Start a local OpenAI-compatible server with a web UI: llama serve -hf talxcc/Tals-coder-flash-02 # Run inference directly in the terminal: llama cli -hf talxcc/Tals-coder-flash-02
Install from WinGet (Windows)
winget install llama.cpp # Start a local OpenAI-compatible server with a web UI: llama serve -hf talxcc/Tals-coder-flash-02 # Run inference directly in the terminal: llama cli -hf talxcc/Tals-coder-flash-02
Use pre-built binary
# Download pre-built binary from: # https://github.com/ggerganov/llama.cpp/releases # Start a local OpenAI-compatible server with a web UI: ./llama-server -hf talxcc/Tals-coder-flash-02 # Run inference directly in the terminal: ./llama-cli -hf talxcc/Tals-coder-flash-02
Build from source code
git clone https://github.com/ggerganov/llama.cpp.git cd llama.cpp cmake -B build cmake --build build -j --target llama-server llama-cli # Start a local OpenAI-compatible server with a web UI: ./build/bin/llama-server -hf talxcc/Tals-coder-flash-02 # Run inference directly in the terminal: ./build/bin/llama-cli -hf talxcc/Tals-coder-flash-02
Use Docker
docker model run hf.co/talxcc/Tals-coder-flash-02
- LM Studio
- Jan
- vLLM
How to use talxcc/Tals-coder-flash-02 with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "talxcc/Tals-coder-flash-02" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "talxcc/Tals-coder-flash-02", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker
docker model run hf.co/talxcc/Tals-coder-flash-02
- Ollama
How to use talxcc/Tals-coder-flash-02 with Ollama:
ollama run hf.co/talxcc/Tals-coder-flash-02
- Unsloth Desktop
- Pi
How to use talxcc/Tals-coder-flash-02 with Pi:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf talxcc/Tals-coder-flash-02
Configure the model in Pi
# Install Pi: npm install -g @earendil-works/pi-coding-agent # Add to ~/.pi/agent/models.json: { "providers": { "llama-cpp": { "baseUrl": "http://localhost:8080/v1", "api": "openai-completions", "apiKey": "none", "models": [ { "id": "talxcc/Tals-coder-flash-02" } ] } } }Run Pi
# Start Pi in your project directory: pi
- Docker Model Runner
How to use talxcc/Tals-coder-flash-02 with Docker Model Runner:
docker model run hf.co/talxcc/Tals-coder-flash-02
- Lemonade
How to use talxcc/Tals-coder-flash-02 with Lemonade:
Pull the model
# Download Lemonade from https://lemonade-server.ai/ lemonade pull talxcc/Tals-coder-flash-02
Run and chat with the model
lemonade run user.Tals-coder-flash-02-{{QUANT_TAG}}List all available models
lemonade list
- Hermes Agent
How to use talxcc/Tals-coder-flash-02 with Hermes Agent:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf talxcc/Tals-coder-flash-02
Configure Hermes
# Install Hermes: curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash hermes setup # Point Hermes at the local server: hermes config set model.provider custom hermes config set model.base_url http://127.0.0.1:8080/v1 hermes config set model.default talxcc/Tals-coder-flash-02
Run Hermes
hermes
- Atomic Chat
- OpenClaw
How to use talxcc/Tals-coder-flash-02 with OpenClaw:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf talxcc/Tals-coder-flash-02
Configure OpenClaw
# Install OpenClaw: npm install -g openclaw@latest # Register the local server and set it as the default model: openclaw onboard --non-interactive --mode local \ --auth-choice custom-api-key \ --custom-base-url http://127.0.0.1:8080/v1 \ --custom-model-id "talxcc/Tals-coder-flash-02" \ --custom-provider-id llama-cpp \ --custom-compatibility openai \ --custom-text-input \ --accept-risk \ --skip-health
Run OpenClaw
openclaw agent --local --agent main --message "Hello from Hugging Face"
Download bonsai2-pascal.patch from talxcc/Tals-coder-flash-02: direct link, hf CLI and curl.
- Browser
- Download file 259 kB
-
https://huggingface.co/talxcc/Tals-coder-flash-02/resolve/main/bonsai2-pascal.patch
- Command line
-
hf download hf://talxcc/Tals-coder-flash-02/bonsai2-pascal.patch
-
curl -L -o bonsai2-pascal.patch https://huggingface.co/talxcc/Tals-coder-flash-02/resolve/main/bonsai2-pascal.patch
259 kB
| diff --git a/common/speculative.cpp b/common/speculative.cpp | |
| index 05d9ff8..b0f10de 100644 | |
| --- a/common/speculative.cpp | |
| +++ b/common/speculative.cpp | |
| struct common_speculative_impl_draft_dflash : public common_speculative_impl { | |
| } | |
| }; | |
| +// ---- diagnostic MTP cycle timer: MTP_PROF=N prints every N cycles to stderr ---- | |
| +struct mtp_prof_state { | |
| + int every = -1; | |
| + int64_t t_draft_end = 0, t_process_end = 0; | |
| + double outside = 0, process = 0, post = 0, dec = 0, smp = 0, v_prep = 0, v_cpu = 0, v_wait = 0; | |
| + long cycles = 0, steps = 0, drafts = 0, width = 0, prefills = 0; | |
| +}; | |
| +static mtp_prof_state & mtp_prof() { | |
| + static mtp_prof_state st; | |
| + if (st.every < 0) { | |
| + const char * e = getenv("MTP_PROF"); | |
| + st.every = e ? atoi(e) : 0; | |
| + } | |
| + return st; | |
| +} | |
| +static void mtp_prof_report() { | |
| + auto & st = mtp_prof(); | |
| + if (st.cycles == 0) { | |
| + return; | |
| + } | |
| + const double c = (double) st.cycles; | |
| + fprintf(stderr, "MTPPROF cycles=%ld ms/cycle: verify %.2f catch-up %.2f server %.2f draft-decode %.2f " | |
| + "draft-sample %.2f | total %.2f | steps/cycle %.2f drafts/cycle %.2f width %.2f | per step: decode %.2f sample %.2f\n", | |
| + st.cycles, st.outside / c, st.process / c, st.post / c, st.dec / c, st.smp / c, | |
| + (st.outside + st.process + st.post + st.dec + st.smp) / c, | |
| + st.steps / c, st.drafts / c, st.width / c, st.steps ? st.dec / st.steps : 0.0, st.steps ? st.smp / st.steps : 0.0); | |
| + fprintf(stderr, "MTPPROF verify split ms/cycle: server-prep %.2f decode-call(cpu) %.2f gpu-wait %.2f (prefills skipped %ld)\n", | |
| + st.v_prep / c, st.v_cpu / c, st.v_wait / c, st.prefills); | |
| + fflush(stderr); | |
| + const int every = st.every; | |
| + const int64_t t_end = st.t_draft_end, t_pend = st.t_process_end; | |
| + st = mtp_prof_state(); | |
| + st.every = every; | |
| + st.t_draft_end = t_end; | |
| + st.t_process_end = t_pend; | |
| +} | |
| + | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| std::vector<int> i_last; | |
| std::vector<std::vector<float>> chain_h; | |
| + // fused catch-up (MTP_FUSE_CATCHUP, default on): process() records the verify batch instead of decoding it, | |
| + // accept() trims the record to the accepted prefix, and the next draft() decodes [prefix | draft-1] in one call | |
| + struct deferred_entry { | |
| + llama_token tok; | |
| + llama_pos pos; | |
| + std::vector<float> h; | |
| + }; | |
| + bool fuse_catchup = true; | |
| + int dbg_left = -1; // MTP_DEBUG=N: trace the first N driver calls to stderr | |
| + bool dbg() { | |
| + if (dbg_left < 0) { | |
| + const char * e = getenv("MTP_DEBUG"); | |
| + dbg_left = e ? atoi(e) : 0; | |
| + } | |
| + return dbg_left > 0; | |
| + } | |
| + std::vector<std::vector<deferred_entry>> deferred; // [n_seq] | |
| + std::vector<int32_t> deferred_new; // entries the latest process() appended, per seq | |
| + | |
| + // decode every recorded entry as a plain catch-up batch (no outputs) and forget them | |
| + bool flush_deferred() { | |
| + auto * ctx_dft = this->params.ctx_dft; | |
| + const size_t row_bytes = (size_t) n_embd * sizeof(float); | |
| + common_batch_clear(batch); | |
| + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { | |
| + for (const auto & e : deferred[seq_id]) { | |
| + common_batch_add(batch, e.tok, e.pos, { seq_id }, 0); | |
| + std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, e.h.data(), row_bytes); | |
| + } | |
| + deferred[seq_id].clear(); | |
| + deferred_new[seq_id] = 0; | |
| + } | |
| + if (batch.n_tokens == 0) { | |
| + return true; | |
| + } | |
| + const int32_t rc = llama_decode(ctx_dft, batch); | |
| + common_batch_clear(batch); | |
| + if (rc != 0) { | |
| + SPC_ERR("llama_decode(ctx_dft) deferred catch-up failed rc=%d\n", (int) rc); | |
| + return false; | |
| + } | |
| + return true; | |
| + } | |
| + | |
| common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) | |
| : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) | |
| , params(params.draft) | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| verify_h.assign(n_seq, {}); | |
| verify_h_rows.assign(n_seq, 0); | |
| + | |
| + { | |
| + const char * e = getenv("MTP_FUSE_CATCHUP"); | |
| + fuse_catchup = (!e || atoi(e) != 0) && !is_mem_shared && !chain_heads; | |
| + } | |
| + deferred.assign(n_seq, {}); | |
| + deferred_new.assign(n_seq, 0); | |
| } | |
| ~common_speculative_impl_draft_mtp() override { | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| if (batch_in.n_tokens <= 0) { | |
| return true; | |
| } | |
| + auto & prof = mtp_prof(); | |
| + if (prof.every > 0) { | |
| + llama_synchronize(this->params.ctx_tgt); // the verify decode is async: charge its GPU time to 'outside' | |
| + } | |
| + const int64_t prof_t0 = prof.every > 0 ? ggml_time_us() : 0; | |
| + if (prof.every > 0 && prof.t_draft_end > 0 && batch_in.n_tokens > this->params.n_max + 1) { | |
| + prof.prefills++; // a prompt prefill (and the client turnaround before it), not a verify: keep it out of the cycle | |
| + } else if (prof.every > 0 && prof.t_draft_end > 0) { | |
| + prof.outside += (prof_t0 - prof.t_draft_end) / 1000.0; | |
| + prof.width += batch_in.n_tokens; | |
| + int64_t te = 0, tr = 0; | |
| + llama_ext_last_decode_times(this->params.ctx_tgt, &te, &tr); | |
| + static const bool trace = getenv("MTP_PROF_TRACE") != nullptr; | |
| + if (trace) { | |
| + fprintf(stderr, "MTPTRACE w=%d prep=%.3f cpu=%.3f wait=%.3f (te-tde %lld us, tr-te %lld us)\n", (int) batch_in.n_tokens, | |
| + (te - prof.t_draft_end) / 1000.0, (tr - te) / 1000.0, (prof_t0 - tr) / 1000.0, | |
| + (long long) (te - prof.t_draft_end), (long long) (tr - te)); | |
| + } | |
| + if (te > prof.t_draft_end && tr >= te && prof_t0 >= tr) { | |
| + prof.v_prep += (te - prof.t_draft_end) / 1000.0; | |
| + prof.v_cpu += (tr - te) / 1000.0; | |
| + prof.v_wait += (prof_t0 - tr) / 1000.0; | |
| + } | |
| + } | |
| + struct prof_guard { | |
| + mtp_prof_state & st; int64_t t0; llama_context * dft; | |
| + ~prof_guard() { | |
| + if (st.every > 0) { | |
| + llama_synchronize(dft); // the catch-up decode too | |
| + const int64_t t1 = ggml_time_us(); | |
| + st.process += (t1 - t0) / 1000.0; | |
| + st.t_process_end = t1; | |
| + } | |
| + } | |
| + } prof_g { prof, prof_t0, this->params.ctx_dft }; | |
| // TODO: how to make it work with vision tokens? | |
| if (batch_in.token == nullptr || batch_in.embd != nullptr) { | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| const size_t row_bytes = (size_t) n_embd * sizeof(float); | |
| + // fused catch-up: a verify-sized batch (<= n_max + 1 rows per seq) is recorded, not decoded | |
| + bool defer = fuse_catchup; | |
| + for (llama_seq_id seq_id = 0; defer && seq_id < (llama_seq_id) n_seq; ++seq_id) { | |
| + if (i_batch_beg[seq_id] >= 0 && i_batch_end[seq_id] - i_batch_beg[seq_id] + 1 > params.n_max + 1) { | |
| + defer = false; | |
| + } | |
| + } | |
| + if (fuse_catchup) { | |
| + // recorded entries at or past this batch's first position are superseded (replay); entries that end right | |
| + // before it are kept (deferred) or flushed (prefill); anything else is a gap and is dropped | |
| + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { | |
| + auto & d = deferred[seq_id]; | |
| + if (i_batch_beg[seq_id] < 0 || d.empty()) { | |
| + continue; | |
| + } | |
| + const llama_pos p0 = batch_in.pos[i_batch_beg[seq_id]]; | |
| + while (!d.empty() && d.back().pos >= p0) { | |
| + d.pop_back(); | |
| + } | |
| + if (!d.empty() && d.back().pos != p0 - 1) { | |
| + SPC_WRN("seq %d: dropping %zu recorded catch-up rows (gap before pos %d)\n", (int) seq_id, d.size(), (int) p0); | |
| + d.clear(); | |
| + } | |
| + } | |
| + if (!defer && !flush_deferred()) { | |
| + return false; | |
| + } | |
| + } | |
| + if (dbg()) { | |
| + dbg_left--; | |
| + fprintf(stderr, "MTPDBG process n=%d pos0=%d defer=%d rec=%zu tok:", (int) n_tokens, (int) batch_in.pos[0], (int) defer, | |
| + deferred.empty() ? (size_t) 0 : deferred[0].size()); | |
| + for (int k = 0; k < n_tokens && k < 8; ++k) fprintf(stderr, " %d", batch_in.token[k]); | |
| + fprintf(stderr, " h0=%.4f\n", pending_h[0][0]); | |
| + } | |
| + if (defer) { | |
| + const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt); | |
| + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { | |
| + deferred_new[seq_id] = 0; | |
| + if (i_batch_beg[seq_id] < 0) { | |
| + continue; | |
| + } | |
| + for (int k = i_batch_beg[seq_id]; k <= i_batch_end[seq_id]; ++k) { | |
| + deferred_entry e; | |
| + e.tok = batch_in.token[k]; | |
| + e.pos = batch_in.pos[k]; | |
| + if (k == i_batch_beg[seq_id]) { | |
| + // a sequence start has no h_{p-1}: zeros, as in training (patch_fork_mtp_pos0.py) | |
| + static const int pos0_mode = getenv("MTP_POS0_H") ? atoi(getenv("MTP_POS0_H")) : 2; | |
| + if (e.pos == 0 && pos0_mode == 0) { | |
| + e.h.assign(n_embd, 0.0f); | |
| + } else if (e.pos == 0 && pos0_mode == 2) { | |
| + e.h.assign(h_tgt + (size_t) k * n_embd, h_tgt + (size_t) (k + 1) * n_embd); | |
| + } else { | |
| + e.h = pending_h[seq_id]; | |
| + } | |
| + } else { | |
| + e.h.assign(h_tgt + (size_t) (k - 1) * n_embd, h_tgt + (size_t) k * n_embd); | |
| + } | |
| + deferred[seq_id].push_back(std::move(e)); | |
| + deferred_new[seq_id]++; | |
| + } | |
| + } | |
| + // a sequence that stops drafting keeps recording: flush before the record grows large | |
| + size_t n_rec = 0; | |
| + for (const auto & d : deferred) { | |
| + n_rec += d.size(); | |
| + } | |
| + if (n_rec > 64 && !flush_deferred()) { | |
| + return false; | |
| + } | |
| + } | |
| + | |
| // if kv is shared with target (e.g Gemma4), then we can skip this catch-up decode | |
| - if (!is_mem_shared) { | |
| + if (!is_mem_shared && !defer) { | |
| common_batch_clear(batch); | |
| for (int k = 0; k < n_tokens; ++k) { | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| continue; | |
| } | |
| - set_h(i_batch_beg[seq_id], pending_h[seq_id].data()); | |
| + // a sequence start has no h_{p-1}: zeros, as in training -- not the previous request's last h | |
| + static const int pos0_mode = getenv("MTP_POS0_H") ? atoi(getenv("MTP_POS0_H")) : 2; | |
| + if (batch_in.pos[i_batch_beg[seq_id]] == 0 && pos0_mode == 0) { | |
| + std::memset(batch.embd + (size_t) i_batch_beg[seq_id] * n_embd, 0, row_bytes); | |
| + } else if (batch_in.pos[i_batch_beg[seq_id]] == 0 && pos0_mode == 2) { | |
| + const float * h_tgt0 = llama_get_embeddings_nextn(ctx_tgt); | |
| + set_h(i_batch_beg[seq_id], h_tgt0 + (size_t) i_batch_beg[seq_id] * n_embd); | |
| + } else { | |
| + set_h(i_batch_beg[seq_id], pending_h[seq_id].data()); | |
| + } | |
| } | |
| auto * mem_dft = llama_get_memory(ctx_dft); | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| void draft(common_speculative_draft_params_vec & dparams) override { | |
| auto & ctx_dft = params.ctx_dft; | |
| + { | |
| + auto & prof = mtp_prof(); | |
| + if (prof.every > 0 && prof.t_process_end > 0) { | |
| + prof.post += (ggml_time_us() - prof.t_process_end) / 1000.0; | |
| + prof.t_process_end = 0; | |
| + } | |
| + } | |
| + { | |
| + // diagnostic: MTP_FUSE_SPLIT=1 decodes the recorded catch-up rows as their own batch before draft 1 | |
| + static const bool split = getenv("MTP_FUSE_SPLIT") != nullptr && atoi(getenv("MTP_FUSE_SPLIT")) != 0; | |
| + if (split && fuse_catchup) { | |
| + flush_deferred(); | |
| + } | |
| + } | |
| common_batch_clear(batch); | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| drafting[seq_id] = true; | |
| common_sampler_reset(smpls[seq_id].get()); | |
| + if (fuse_catchup) { | |
| + auto & d = deferred[seq_id]; | |
| + if (!d.empty() && d.back().pos != dp.n_past - 1) { | |
| + SPC_WRN("seq %d: dropping %zu recorded catch-up rows (last pos %d, n_past %d)\n", | |
| + (int) seq_id, d.size(), (int) d.back().pos, (int) dp.n_past); | |
| + d.clear(); | |
| + } | |
| + for (const auto & e : d) { | |
| + common_batch_add(batch, e.tok, e.pos, { seq_id }, false); | |
| + std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, e.h.data(), row_bytes); | |
| + } | |
| + d.clear(); | |
| + deferred_new[seq_id] = 0; | |
| + } | |
| + | |
| common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); | |
| std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[seq_id].data(), row_bytes); | |
| i_last[seq_id] = batch.n_tokens - 1; | |
| + if (dbg()) { | |
| + dbg_left--; | |
| + fprintf(stderr, "MTPDBG draft n_past=%d id_last=%d batch=%d:", (int) dp.n_past, dp.id_last, batch.n_tokens); | |
| + for (int k = 0; k < batch.n_tokens; ++k) fprintf(stderr, " [%d@%d h0=%.4f]", batch.token[k], batch.pos[k], batch.embd[(size_t) k * n_embd]); | |
| + fprintf(stderr, "\n"); | |
| + } | |
| if (chain_heads) { | |
| chain_h[seq_id].assign(pending_h[seq_id].begin(), pending_h[seq_id].end()); | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| llama_set_nextn_layer_offset(ctx_dft, i); | |
| } | |
| + auto & prof = mtp_prof(); | |
| + const int64_t prof_td = prof.every > 0 ? ggml_time_us() : 0; | |
| int ret = llama_decode(ctx_dft, batch); | |
| if (ret != 0) { | |
| SPC_ERR("llama_decode[%d] returned %d\n", i, ret); | |
| break; | |
| } | |
| + // the decode is asynchronous until the sampler reads the logits: sync here so the split is honest | |
| + if (prof.every > 0) { | |
| + llama_synchronize(ctx_dft); | |
| + const int64_t t1 = ggml_time_us(); | |
| + prof.dec += (t1 - prof_td) / 1000.0; | |
| + prof.steps++; | |
| + } | |
| + const int64_t prof_ts = prof.every > 0 ? ggml_time_us() : 0; | |
| // rebuild the batch for the next step: the growing-KV paths re-add only the | |
| // new token (the KV already holds the prefix), while chained heads re-add the | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| // add drafted token for each sequence | |
| const llama_token id = cur_p->data[0].id; | |
| - // only collect very high-confidence draft tokens | |
| - if (cur_p->data[0].p < params.p_min) { | |
| + // only collect very high-confidence draft tokens; MTP_PMIN_TAIL = a stricter threshold from the 4th draft on | |
| + static const float pmin_tail = getenv("MTP_PMIN_TAIL") ? (float) atof(getenv("MTP_PMIN_TAIL")) : 0.0f; | |
| + const float pmin_i = (i >= 3 && pmin_tail > 0.0f) ? pmin_tail : params.p_min; | |
| + if (cur_p->data[0].p < pmin_i) { | |
| drafting[seq_id] = false; | |
| n_drafting--; | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| i_last[seq_id] = batch.n_tokens - 1; | |
| } | |
| + if (prof.every > 0) { | |
| + prof.smp += (ggml_time_us() - prof_ts) / 1000.0; | |
| + } | |
| + | |
| if (batch.n_tokens == 0) { | |
| break; | |
| } | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| dp.result->clear(); | |
| } | |
| } | |
| + | |
| + if (dbg() && n_seq > 0 && dparams[0].result) { | |
| + fprintf(stderr, "MTPDBG drafted:"); | |
| + for (auto t : *dparams[0].result) fprintf(stderr, " %d", t); | |
| + fprintf(stderr, "\n"); | |
| + } | |
| + auto & prof = mtp_prof(); | |
| + if (prof.every > 0) { | |
| + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { | |
| + if (dparams[seq_id].drafting && dparams[seq_id].result) { | |
| + prof.drafts += (long) dparams[seq_id].result->size(); | |
| + } | |
| + } | |
| + prof.cycles++; | |
| + prof.t_draft_end = ggml_time_us(); | |
| + if (prof.cycles >= prof.every) { | |
| + mtp_prof_report(); | |
| + } | |
| + } | |
| } | |
| void accept(llama_seq_id seq_id, uint16_t n_accepted, bool /*is_other*/) override { | |
| struct common_speculative_impl_draft_mtp : public common_speculative_impl { | |
| const int32_t i_h = std::min<int32_t>(n_accepted, n_rows - 1); | |
| const size_t row_bytes = (size_t) n_embd * sizeof(float); | |
| std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); | |
| + | |
| + if (dbg()) { | |
| + dbg_left--; | |
| + fprintf(stderr, "MTPDBG accept n_acc=%d n_rows=%d new=%d rec=%zu pend_h0=%.4f\n", (int) n_accepted, (int) n_rows, | |
| + (int) deferred_new[seq_id], deferred[seq_id].size(), pending_h[seq_id][0]); | |
| + } | |
| + if (fuse_catchup) { | |
| + // keep the sampled token + n_accepted drafts of the latest verify batch | |
| + auto & d = deferred[seq_id]; | |
| + const int32_t drop = deferred_new[seq_id] - (int32_t) (n_accepted + 1); | |
| + if (drop > 0 && (size_t) drop <= d.size()) { | |
| + d.resize(d.size() - (size_t) drop); | |
| + } | |
| + deferred_new[seq_id] = 0; | |
| + } | |
| } | |
| }; | |
| diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp | |
| index 3d6310f..021e31a 100644 | |
| --- a/ggml/src/ggml-backend.cpp | |
| +++ b/ggml/src/ggml-backend.cpp | |
| static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { | |
| return true; | |
| } | |
| +// ---- diagnostic: GGML_SCHED_PROF=N (patch_fork_schedprof.py) ---- | |
| +struct ggml_sched_prof_row { long n = 0, inputs = 0; double copy = 0, compute = 0; }; | |
| +static ggml_sched_prof_row g_sched_prof[2][2]; // [big][cpu] | |
| +static long g_sched_prof_calls[2] = {0, 0}; | |
| +static long g_sched_prof_n = 0; | |
| +static int ggml_sched_prof_every() { | |
| + static const int e = getenv("GGML_SCHED_PROF") ? atoi(getenv("GGML_SCHED_PROF")) : 0; | |
| + return e; | |
| +} | |
| +static void ggml_sched_prof_dump() { | |
| + for (int b = 0; b < 2; ++b) { | |
| + if (g_sched_prof_calls[b] == 0) { | |
| + continue; | |
| + } | |
| + for (int c = 0; c < 2; ++c) { | |
| + const auto & r = g_sched_prof[b][c]; | |
| + if (r.n == 0) { | |
| + continue; | |
| + } | |
| + fprintf(stderr, "SCHEDPROF %-5s %-3s splits/compute %.2f inputs/split %.2f | per compute: copy %.3f ms compute-call %.3f ms\n", | |
| + b ? "big" : "small", c ? "CPU" : "GPU", (double) r.n / g_sched_prof_calls[b], (double) r.inputs / r.n, | |
| + r.copy / g_sched_prof_calls[b] / 1000.0, r.compute / g_sched_prof_calls[b] / 1000.0); | |
| + } | |
| + } | |
| + fflush(stderr); | |
| + for (auto & rr : g_sched_prof) { for (auto & r : rr) { r = ggml_sched_prof_row(); } } | |
| + g_sched_prof_calls[0] = g_sched_prof_calls[1] = 0; | |
| +} | |
| + | |
| static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { | |
| GGML_ASSERT(sched); | |
| struct ggml_backend_sched_split * splits = sched->splits; | |
| + const bool sprof = ggml_sched_prof_every() > 0; | |
| + const int sprof_big = sched->graph.n_nodes > 500 ? 1 : 0; | |
| + if (sprof) { | |
| + g_sched_prof_calls[sprof_big]++; | |
| + } | |
| ggml_tensor * prev_ids_tensor = nullptr; | |
| std::vector<int32_t> ids; | |
| static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s | |
| struct ggml_backend_sched_split * split = &splits[split_id]; | |
| int split_backend_id = split->backend_id; | |
| ggml_backend_t split_backend = sched->backends[split_backend_id]; | |
| + const int64_t sp_t0 = sprof ? ggml_time_us() : 0; | |
| // ensure the previous split's async work has completed before we start | |
| // this split, the allocator may have reused buffer regions across splits | |
| static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s | |
| } | |
| // copy the input tensors to the split backend | |
| + // batched host inputs (patch_fork_batch_inputs.py): one sync before the first host copy, async copies, one sync after | |
| + static const bool batch_inputs_on = getenv("GGML_SCHED_BATCH_INPUTS") == nullptr || atoi(getenv("GGML_SCHED_BATCH_INPUTS")) != 0; | |
| + const bool bi = batch_inputs_on && sched->events[split_backend_id][sched->cur_copy] == NULL; | |
| + bool bi_synced = false; | |
| + bool bi_pending = false; | |
| for (int input_id = 0; input_id < split->n_inputs; input_id++) { | |
| ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); | |
| struct ggml_tensor * input = split->inputs[input_id]; | |
| struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); | |
| - if (input->flags & GGML_TENSOR_FLAG_INPUT) { | |
| + if (bi && input->buffer != NULL && ggml_backend_buffer_is_host(input->buffer) && | |
| + ggml_backend_buffer_get_usage(input->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { | |
| + ggml_backend_synchronize(input_backend); // a host split's compute is synchronous; kept for other host backends | |
| + if (!bi_synced) { | |
| + ggml_backend_synchronize(split_backend); | |
| + bi_synced = true; | |
| + } | |
| + ggml_backend_tensor_set_async(split_backend, input_cpy, input->data, 0, ggml_nbytes(input)); | |
| + bi_pending = true; | |
| + } else if (input->flags & GGML_TENSOR_FLAG_INPUT) { | |
| // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done | |
| if (sched->events[split_backend_id][sched->cur_copy] != NULL) { | |
| ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); | |
| static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s | |
| } | |
| } | |
| + if (bi_pending) { | |
| + ggml_backend_synchronize(split_backend); // every batched copy has landed before the graph is enqueued | |
| + } | |
| + const int64_t sp_t1 = sprof ? ggml_time_us() : 0; | |
| if (!sched->callback_eval) { | |
| enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); | |
| if (ec != GGML_STATUS_SUCCESS) { | |
| return ec; | |
| } | |
| + if (sprof) { | |
| + const int cpu = strncmp(ggml_backend_name(split_backend), "CPU", 3) == 0 ? 1 : 0; | |
| + auto & r = g_sched_prof[sprof_big][cpu]; | |
| + r.n++; | |
| + r.inputs += split->n_inputs; | |
| + r.copy += (double) (sp_t1 - sp_t0); | |
| + r.compute += (double) (ggml_time_us() - sp_t1); | |
| + } | |
| } else { | |
| // similar to ggml_backend_compare_graph_backend | |
| for (int j0 = 0; j0 < split->graph.n_nodes; j0++) { | |
| static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s | |
| prev_backend_id = split_backend_id; | |
| } | |
| + if (sprof && ++g_sched_prof_n % ggml_sched_prof_every() == 0) { | |
| + ggml_sched_prof_dump(); | |
| + } | |
| + | |
| return GGML_STATUS_SUCCESS; | |
| } | |
| diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh | |
| index ef929d3..1d347af 100644 | |
| --- a/ggml/src/ggml-cuda/common.cuh | |
| +++ b/ggml/src/ggml-cuda/common.cuh | |
| struct ggml_cuda_graph { | |
| std::vector<cudaGraphNode_t> nodes; | |
| bool disable_due_to_gpu_arch = false; | |
| bool warmup_complete = false; | |
| + int n_replays = 0; // replays since the last capture (patch_fork_graph_fastrecap.py) | |
| uint64_t uid = 0; | |
| int64_t last_used_time = 0; | |
| struct node_properties { | |
| diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu | |
| index f597088..2424e91 100644 | |
| --- a/ggml/src/ggml-cuda/concat.cu | |
| +++ b/ggml/src/ggml-cuda/concat.cu | |
| static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml | |
| dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); | |
| if constexpr (sizeof(T) == sizeof(uint32_t)) { | |
| - const bool transpose_dim0 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && | |
| + // DGX Spark, and sm_6x: the SSM conv input at a 4-token verify step (dst [7, 10240] from a transposed view) | |
| + // takes 19.7 us per layer in the generic kernel on a GTX 1080 Ti vs a few us tiled | |
| + const int cc_t = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; | |
| + const bool transpose_dim0 = (cc_t == GGML_CUDA_CC_DGX_SPARK || cc_t < GGML_CUDA_CC_VOLTA) && | |
| dim == 0 && src0->ne[2] == 1 && src0->ne[3] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1 && | |
| dst->ne[2] == 1 && dst->ne[3] == 1 && src0->ne[0] <= 8 && | |
| src0->nb[0] == sizeof(uint32_t) && src0->nb[1] == (uint64_t) src0->ne[0]*sizeof(uint32_t) && | |
| diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh | |
| index d1164b8..6d8fd3c 100644 | |
| --- a/ggml/src/ggml-cuda/fattn-tile.cuh | |
| +++ b/ggml/src/ggml-cuda/fattn-tile.cuh | |
| static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 2, 64, 64) | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 32, 256, 2, 64, 64) | |
| + // GQA 6 (patch_fork_fa_gqa6.py): one block per KV head instead of three | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 6, 192, 4, 64, 64) | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 12, 192, 3, 64, 64) | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 24, 384, 2, 64, 64) | |
| + | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 16, 256, 2, 64, 64) | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 2, 64, 2, 64, 64) | |
| static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 2, 32, 128) | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 32, 256, 2, 32, 64) | |
| + // GQA 6 (patch_fork_fa_gqa6.py): one block per KV head instead of three | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 6, 192, 4, 32, 64) | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 12, 192, 3, 32, 64) | |
| + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 24, 384, 2, 32, 64) | |
| + | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 16, 256, 2, 32, 64) | |
| GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 2, 64, 2, 32, 64) | |
| static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm | |
| constexpr size_t nbytes_shared = 0; | |
| + if constexpr (ncols2 == 6) { // patch_fork_fa_gqa6.py: 6 does not divide the power-of-two block widths | |
| + if (Q->ne[1] > 2) { | |
| + constexpr int cols_per_block = 24; | |
| + const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| + const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); | |
| + fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>; | |
| + launch_fattn<DV, cols_per_block/ncols2, ncols2> | |
| + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); | |
| + return; | |
| + } | |
| + if (Q->ne[1] > 1) { | |
| + constexpr int cols_per_block = 12; | |
| + const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| + const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); | |
| + fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>; | |
| + launch_fattn<DV, cols_per_block/ncols2, ncols2> | |
| + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); | |
| + return; | |
| + } | |
| + { | |
| + constexpr int cols_per_block = 6; | |
| + const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| + const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); | |
| + fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>; | |
| + launch_fattn<DV, cols_per_block/ncols2, ncols2> | |
| + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); | |
| + return; | |
| + } | |
| + } | |
| + | |
| + | |
| #ifdef GGML_USE_HIP | |
| if constexpr (DKQ <= 128) { | |
| if (Q->ne[1] > 32/ncols2) { | |
| static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm | |
| #endif // GGML_USE_HIP | |
| #ifndef GGML_USE_HIP | |
| - if constexpr (DKQ <= 256) | |
| + if constexpr (DKQ <= 256 && ncols2 != 6) | |
| #endif // GGML_USE_HIP | |
| { | |
| if (Q->ne[1] > 16/ncols2) { | |
| static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm | |
| } | |
| } | |
| - if constexpr (ncols2 <= 16) { | |
| + if constexpr (ncols2 <= 16 && ncols2 != 6) { | |
| if (Q->ne[1] > 8/ncols2) { | |
| constexpr int cols_per_block = 16; | |
| const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm | |
| } | |
| } | |
| - if constexpr (ncols2 <= 8) { | |
| + if constexpr (ncols2 <= 8 && ncols2 != 6) { | |
| if (Q->ne[1] > 4/ncols2) { | |
| constexpr int cols_per_block = 8; | |
| const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm | |
| } | |
| } | |
| - if constexpr (ncols2 <= 4) { | |
| + if constexpr (ncols2 <= 4 && ncols2 != 6) { | |
| if (Q->ne[1] > 2/ncols2) { | |
| constexpr int cols_per_block = 4; | |
| const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; | |
| static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggm | |
| return; | |
| } | |
| + // patch_fork_fa_gqa6.py: GQA 6 packs all 6 heads (one KV pass) instead of falling through to 2 (three passes). | |
| + // Q->ne[1] <= 4 keeps it to decode/verify batches; wider batches have no ncols = 30/48 config. | |
| + if constexpr (DKQ == 256 && DV == 256) { | |
| + static const bool gqa6 = getenv("GGML_CUDA_FA_GQA6") == nullptr || atoi(getenv("GGML_CUDA_FA_GQA6")) != 0; | |
| + if (gqa6 && use_gqa_opt && gqa_ratio % 6 == 0 && Q->ne[1] <= 4) { | |
| + launch_fattn_tile_switch_ncols1<DKQ, DV, 6, use_logit_softcap>(ctx, dst); | |
| + return; | |
| + } | |
| + } | |
| + | |
| if (use_gqa_opt && gqa_ratio % 2 == 0) { | |
| launch_fattn_tile_switch_ncols1<DKQ, DV, 2, use_logit_softcap>(ctx, dst); | |
| return; | |
| diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu | |
| index 467a84a..9da345d 100644 | |
| --- a/ggml/src/ggml-cuda/fwht.cu | |
| +++ b/ggml/src/ggml-cuda/fwht.cu | |
| #include "common.cuh" | |
| #include "fwht.cuh" | |
| +#include "unary.cuh" | |
| #include <cstdlib> | |
| +#include <unordered_map> | |
| + | |
| +// ---- q8_1 side output (see ggml_cuda_fwht_q8_find) ---- | |
| +struct fwht_q8_entry { | |
| + const void * q8; | |
| + int64_t nelem; | |
| +}; | |
| +static std::unordered_map<const ggml_tensor *, fwht_q8_entry> g_fwht_q8; | |
| + | |
| +#define FWHT_Q8_NSLOT 8 | |
| +#define FWHT_Q8_SLOT_BYTES (2u << 20) | |
| +static void * g_fwht_q8_ring [GGML_CUDA_MAX_DEVICES][FWHT_Q8_NSLOT] = {}; | |
| +static const ggml_tensor * g_fwht_q8_owner[GGML_CUDA_MAX_DEVICES][FWHT_Q8_NSLOT] = {}; | |
| +static int g_fwht_q8_next [GGML_CUDA_MAX_DEVICES] = {}; | |
| + | |
| +void ggml_cuda_fwht_q8_clear() { | |
| + g_fwht_q8.clear(); | |
| +} | |
| + | |
| +const void * ggml_cuda_fwht_q8_find(const ggml_tensor * t, int64_t nelem) { | |
| + if (g_fwht_q8.empty()) { | |
| + return nullptr; | |
| + } | |
| + const ggml_tensor * b = t; | |
| + while (b->view_src) { | |
| + b = b->view_src; | |
| + } | |
| + const auto it = g_fwht_q8.find(b); | |
| + if (it == g_fwht_q8.end() || it->second.nelem != nelem || t->data != b->data) { | |
| + return nullptr; | |
| + } | |
| + return it->second.q8; | |
| +} | |
| + | |
| +static bool fwht_q8_enabled() { | |
| + static const bool on = getenv("GGML_CUDA_FWHT_Q8") == nullptr || atoi(getenv("GGML_CUDA_FWHT_Q8")) != 0; | |
| + return on; | |
| +} | |
| + | |
| +// next ring slot for `owner` (its q8_1 copy lives there until FWHT_Q8_NSLOT more transforms have run); the previous | |
| +// owner of the slot is unregistered so a late reader falls back to quantizing itself | |
| +static void * fwht_q8_slot(const ggml_tensor * owner, size_t bytes, cudaStream_t stream) { | |
| + if (bytes > FWHT_Q8_SLOT_BYTES) { | |
| + return nullptr; | |
| + } | |
| + const int dev = ggml_cuda_get_device(); | |
| + int & k = g_fwht_q8_next[dev]; | |
| + void *& p = g_fwht_q8_ring[dev][k]; | |
| + if (p == nullptr) { | |
| + cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone; | |
| + CUDA_CHECK(cudaStreamIsCapturing(stream, &st)); | |
| + if (st != cudaStreamCaptureStatusNone) { | |
| + return nullptr; // no allocation inside a graph capture: this transform just skips the side output | |
| + } | |
| + CUDA_CHECK(cudaMalloc(&p, FWHT_Q8_SLOT_BYTES)); | |
| + } | |
| + if (g_fwht_q8_owner[dev][k] != nullptr) { | |
| + g_fwht_q8.erase(g_fwht_q8_owner[dev][k]); | |
| + } | |
| + g_fwht_q8_owner[dev][k] = owner; | |
| + void * res = p; | |
| + k = (k + 1) % FWHT_Q8_NSLOT; | |
| + return res; | |
| +} | |
| template <typename T> | |
| __device__ __forceinline__ float fwht_load(const T value) { | |
| __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, | |
| template <int N, int NT, typename T, bool has_signs> | |
| __launch_bounds__(NT, 1) | |
| __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale, | |
| - const float * signs, const int n_blk) { | |
| + const float * signs, const int n_blk, block_q8_1 * q8) { | |
| constexpr int warp_size = ggml_cuda_get_physical_warp_size(); | |
| constexpr int NE = N / NT; | |
| static_assert(NE >= 1 && N % NT == 0 && NT % warp_size == 0, "bad FWHT block shape"); | |
| __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows | |
| for (int i = 0; i < NE; ++i) { | |
| dst[i * NT + tid] = reg[i]; | |
| } | |
| + | |
| + if (q8 != nullptr) { | |
| + // element i*NT + tid: for a fixed i one warp holds 32 consecutive, 32-aligned elements = one q8_1 block; | |
| + // quantized exactly like quantize_q8_1 (amax / 127, roundf, ds = {d, sum}) | |
| + static_assert(NT % QK8_1 == 0 && N % QK8_1 == 0, "FWHT q8_1 side output needs 32-wide warps over the row"); | |
| + block_q8_1 * yb = q8 + r * (N / QK8_1); | |
| +#pragma unroll | |
| + for (int i = 0; i < NE; ++i) { | |
| + const float xi = reg[i]; | |
| + const float amax = warp_reduce_max<QK8_1>(fabsf(xi)); | |
| + const float sum = warp_reduce_sum<QK8_1>(xi); | |
| + const float d = amax / 127.0f; | |
| + const int e = i * NT + tid; | |
| + block_q8_1 & b = yb[e / QK8_1]; | |
| + b.qs[e % QK8_1] = amax == 0.0f ? 0 : (int8_t) roundf(xi / d); | |
| + if (e % QK8_1 == 0) { | |
| + b.ds = make_half2(d, sum); | |
| + } | |
| + } | |
| + } | |
| } | |
| template <typename T> | |
| static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float * dst_d, | |
| const int n, const int64_t rows, const float scale, | |
| - const float * signs, const int n_blk) { | |
| + const float * signs, const int n_blk, block_q8_1 * q8) { | |
| const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; | |
| const int rows_per_block = 4; | |
| const int64_t num_blocks = (rows + rows_per_block - 1) / rows_per_block; | |
| static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float | |
| const dim3 g((unsigned) rows, 1, 1), b(FWHT_BLOCK_THREADS, 1, 1); \ | |
| const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(g, b, 0, stream); \ | |
| if (signs) { \ | |
| - ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true>, lp, src_d, dst_d, rows, scale, signs, n_blk); \ | |
| + ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true>, lp, src_d, dst_d, rows, scale, signs, n_blk, q8); \ | |
| } else { \ | |
| - ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, false>, lp, src_d, dst_d, rows, scale, nullptr, 1); \ | |
| + ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, false>, lp, src_d, dst_d, rows, scale, nullptr, 1, q8); \ | |
| } \ | |
| return true; \ | |
| } | |
| static bool fwht_dispatch(ggml_backend_cuda_context & ctx, const ggml_tensor * s | |
| float * dst_d = (float *) dst->data; | |
| const float scale = 1 / sqrtf(n); | |
| + // q8_1 side output for the sm_6x PQ2_0 GEMV (block kernel widths only; legacy kernels do not write it) | |
| + block_q8_1 * q8 = nullptr; | |
| + static const bool legacy_q8 = getenv("GGML_CUDA_FWHT_LEGACY") != nullptr; | |
| + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; | |
| + if (fwht_q8_enabled() && !legacy_q8 && GGML_CUDA_CC_IS_NVIDIA(cc) && cc < GGML_CUDA_CC_VOLTA && | |
| + n >= 512 && n <= 8192 && ggml_is_contiguous(dst)) { | |
| + q8 = (block_q8_1 *) fwht_q8_slot(dst, (size_t) (ggml_nelements(dst) / QK8_1) * sizeof(block_q8_1), ctx.stream()); | |
| + } | |
| + | |
| + bool ok; | |
| if (src->type == GGML_TYPE_F32) { | |
| - return fwht_launch<float>(ctx, (const float *) src->data, dst_d, n, rows, scale, signs, n_blk); | |
| + ok = fwht_launch<float>(ctx, (const float *) src->data, dst_d, n, rows, scale, signs, n_blk, q8); | |
| + } else { | |
| + ok = fwht_launch<half>(ctx, (const half *) src->data, dst_d, n, rows, scale, signs, n_blk, q8); | |
| + } | |
| + if (ok && q8 != nullptr) { | |
| + g_fwht_q8[dst] = { q8, ggml_nelements(dst) }; | |
| + } | |
| + return ok; | |
| +} | |
| + | |
| + | |
| +// ---- fused residual add + RMS norm + weight + signs + FWHT_1024 (+q8_1) ---- | |
| +// a / b may alias add_out / norm_out / fwht_out at the same element (buffer reuse): no __restrict__ on those | |
| +static __global__ void __launch_bounds__(1024, 1) add_rmsnorm_fwht1024( | |
| + const float * a, const float * b, float * add_out, | |
| + const float * __restrict__ w, const float eps, float * norm_out, const float * __restrict__ signs, | |
| + float * fwht_out, block_q8_1 * __restrict__ q8, const int ncols, const float * __restrict__ row_scale) { | |
| + constexpr int N = 1024; | |
| + const int chunk = blockIdx.x; | |
| + const int row = blockIdx.y; | |
| + const int tid = threadIdx.x; | |
| + const int nchk = gridDim.x; | |
| + const float * ar = a + (size_t) row * ncols; | |
| + const float * br = b ? b + (size_t) row * ncols : nullptr; | |
| + | |
| + // 1) the row's RMS scale, exactly as rms_norm_f32<1024> computes it on the (a + b) tensor -- or precomputed by | |
| + // rmsnorm_row_scale when an output buffer overlaps an input (then this block only touches its own chunk) | |
| + float scale; | |
| + if (row_scale != nullptr) { | |
| + scale = row_scale[row]; | |
| + } else { | |
| + float tmp = 0.0f; | |
| + for (int col = tid; col < ncols; col += N) { | |
| + const float xi = br ? ar[col] + br[col] : ar[col]; | |
| + tmp += xi * xi; | |
| + } | |
| + __shared__ float s_sum[32]; | |
| + tmp = block_reduce<block_reduce_method::SUM, 1024>(tmp, s_sum); | |
| + const float mean = tmp / ncols; | |
| + scale = rsqrtf(mean + eps); | |
| + } | |
| + | |
| + // 2) this chunk: residual, norm output, Hadamard input -- same expressions/order as ADD, rms_norm_f32<1024, true>, | |
| + // and fwht_cuda_block (x * 1/sqrt(N) * sign) | |
| + const int col = chunk * N + tid; | |
| + const float r = br ? ar[col] + br[col] : ar[col]; | |
| + if (add_out) { | |
| + add_out[(size_t) row * ncols + col] = r; | |
| + } | |
| + const float y = scale * r * w[col]; | |
| + norm_out[(size_t) row * ncols + col] = y; | |
| + float v = y * (1.0f / 32.0f); | |
| + v *= signs[col]; | |
| + | |
| + // 3) the 1024-point butterfly: within the warp (h < 32), then across warps through smem (h = 32 .. 512) | |
| + const int lane = tid % 32; | |
| +#pragma unroll | |
| + for (int h = 1; h < 32; h *= 2) { | |
| + const float v2 = __shfl_xor_sync(0xFFFFFFFF, v, h, 32); | |
| + v = (lane & h) == 0 ? v + v2 : v2 - v; | |
| + } | |
| + __shared__ float s[N]; | |
| +#pragma unroll | |
| + for (int h = 32; h < N; h *= 2) { | |
| + s[tid] = v; | |
| + __syncthreads(); | |
| + const float v2 = s[tid ^ h]; | |
| + v = (tid & h) == 0 ? v + v2 : v2 - v; | |
| + __syncthreads(); | |
| + } | |
| + | |
| + // 4) outputs: FWHT row (row * nchk + chunk) of the [1024, rows*nchk] tensor, and its q8_1 blocks | |
| + const size_t frow = (size_t) row * nchk + chunk; | |
| + fwht_out[frow * N + tid] = v; | |
| + if (q8 != nullptr) { | |
| + const float amax = warp_reduce_max<QK8_1>(fabsf(v)); | |
| + const float sum = warp_reduce_sum<QK8_1>(v); | |
| + const float d = amax / 127.0f; | |
| + block_q8_1 & bq = q8[frow * (N / QK8_1) + tid / QK8_1]; | |
| + bq.qs[lane] = amax == 0.0f ? 0 : (int8_t) roundf(v / d); | |
| + if (lane == 0) { | |
| + bq.ds = make_half2(d, sum); | |
| + } | |
| + } | |
| +} | |
| + | |
| +static __global__ void __launch_bounds__(1024, 1) rmsnorm_row_scale( | |
| + const float * __restrict__ a, const float * __restrict__ b, float * __restrict__ scale_out, const int ncols, | |
| + const float eps) { | |
| + const int row = blockIdx.x; | |
| + const float * ar = a + (size_t) row * ncols; | |
| + const float * br = b ? b + (size_t) row * ncols : nullptr; | |
| + float tmp = 0.0f; | |
| + for (int col = threadIdx.x; col < ncols; col += 1024) { | |
| + const float xi = br ? ar[col] + br[col] : ar[col]; | |
| + tmp += xi * xi; | |
| + } | |
| + __shared__ float s_sum[32]; | |
| + tmp = block_reduce<block_reduce_method::SUM, 1024>(tmp, s_sum); | |
| + if (threadIdx.x == 0) { | |
| + scale_out[row] = rsqrtf(tmp / ncols + eps); | |
| + } | |
| +} | |
| + | |
| +void ggml_cuda_op_add_rmsnorm_fwht(ggml_backend_cuda_context & ctx, const float * a, const float * b, float * add_out, | |
| + const float * w, float eps, float * norm_out, const float * signs, | |
| + int64_t ncols, int64_t nrows, ggml_tensor * fwht_dst) { | |
| + cudaStream_t stream = ctx.stream(); | |
| + block_q8_1 * q8 = nullptr; | |
| + if (fwht_q8_enabled()) { | |
| + q8 = (block_q8_1 *) fwht_q8_slot(fwht_dst, (size_t) (ggml_nelements(fwht_dst) / QK8_1) * sizeof(block_q8_1), stream); | |
| + } | |
| + // outputs overlapping inputs (the graph allocator reuses buffers: in-place ADD, freed inputs): a block of the single | |
| + // kernel would overwrite another chunk of a row other blocks still read for the norm -> precompute the row scale | |
| + const size_t rb = (size_t) nrows * ncols * sizeof(float); | |
| + auto ovl = [rb](const void * x, const void * y) { | |
| + if (!x || !y) { | |
| + return false; | |
| + } | |
| + const char * p = (const char *) x; | |
| + const char * q = (const char *) y; | |
| + return p < q + rb && q < p + rb; | |
| + }; | |
| + const void * outs[3] = { add_out, norm_out, fwht_dst->data }; | |
| + bool overlap = false; | |
| + for (const void * o : outs) { | |
| + overlap = overlap || ovl(o, a) || ovl(o, b); | |
| + } | |
| + ggml_cuda_pool_alloc<float> scale_buf(ctx.pool()); | |
| + const float * row_scale = nullptr; | |
| + if (overlap) { | |
| + scale_buf.alloc(nrows); | |
| + rmsnorm_row_scale<<<(unsigned) nrows, 1024, 0, stream>>>(a, b, scale_buf.get(), (int) ncols, eps); | |
| + row_scale = scale_buf.get(); | |
| + } | |
| + add_rmsnorm_fwht1024<<<dim3((unsigned) (ncols / 1024), (unsigned) nrows), 1024, 0, stream>>>( | |
| + a, b, add_out, w, eps, norm_out, signs, (float *) fwht_dst->data, q8, (int) ncols, row_scale); | |
| + if (q8 != nullptr) { | |
| + g_fwht_q8[fwht_dst] = { q8, ggml_nelements(fwht_dst) }; | |
| + } | |
| +} | |
| + | |
| +// ---- fused SWIGLU + signs + FWHT_1024 (sm_6x): fwht_cuda_block<1024, 256, float, true> with a swiglu load ---- | |
| +static __global__ void __launch_bounds__(FWHT_BLOCK_THREADS, 1) swiglu_fwht1024( | |
| + const float * gate, const float * up, float * dst, const float scale, const float * signs, const int n_blk, | |
| + block_q8_1 * q8) { | |
| + constexpr int N = 1024; | |
| + constexpr int NT = FWHT_BLOCK_THREADS; | |
| + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); | |
| + constexpr int NE = N / NT; | |
| + | |
| + __shared__ float s[N]; | |
| + | |
| + const int64_t r = blockIdx.x; | |
| + const int tid = threadIdx.x; | |
| + const int lane = tid % warp_size; | |
| + | |
| + gate += r * N; | |
| + up += r * N; | |
| + dst += r * N; | |
| + const float * signs_row = signs + (r % n_blk) * N; | |
| + | |
| + float reg[NE]; | |
| +#pragma unroll | |
| + for (int i = 0; i < NE; ++i) { | |
| + const float v = ggml_cuda_op_silu_single(gate[i * NT + tid]) * up[i * NT + tid]; // the GLU kernel's value | |
| + reg[i] = v * scale; | |
| + reg[i] *= signs_row[i * NT + tid]; | |
| + } | |
| + | |
| +#pragma unroll | |
| + for (int h = 1; h < warp_size; h *= 2) { | |
| +#pragma unroll | |
| + for (int j = 0; j < NE; j++) { | |
| + const float val = reg[j]; | |
| + const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size); | |
| + reg[j] = (lane & h) == 0 ? val + val2 : val2 - val; | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int h = warp_size; h < NT; h *= 2) { | |
| +#pragma unroll | |
| + for (int j = 0; j < NE; j++) { | |
| + s[j * NT + tid] = reg[j]; | |
| + } | |
| + __syncthreads(); | |
| +#pragma unroll | |
| + for (int j = 0; j < NE; j++) { | |
| + const float val = reg[j]; | |
| + const float val2 = s[j * NT + (tid ^ h)]; | |
| + reg[j] = (tid & h) == 0 ? val + val2 : val2 - val; | |
| + } | |
| + __syncthreads(); | |
| + } | |
| +#pragma unroll | |
| + for (int h = NT; h < N; h *= 2) { | |
| + const int step = h / NT; | |
| +#pragma unroll | |
| + for (int j = 0; j < NE; j += 2 * step) { | |
| +#pragma unroll | |
| + for (int k = 0; k < step; k++) { | |
| + const float x = reg[j + k]; | |
| + const float y = reg[j + k + step]; | |
| + reg[j + k] = x + y; | |
| + reg[j + k + step] = x - y; | |
| + } | |
| + } | |
| + } | |
| + | |
| +#pragma unroll | |
| + for (int i = 0; i < NE; ++i) { | |
| + dst[i * NT + tid] = reg[i]; | |
| + } | |
| + | |
| + if (q8 != nullptr) { | |
| + block_q8_1 * yb = q8 + r * (N / QK8_1); | |
| +#pragma unroll | |
| + for (int i = 0; i < NE; ++i) { | |
| + const float xi = reg[i]; | |
| + const float amax = warp_reduce_max<QK8_1>(fabsf(xi)); | |
| + const float sum = warp_reduce_sum<QK8_1>(xi); | |
| + const float d = amax / 127.0f; | |
| + const int e = i * NT + tid; | |
| + block_q8_1 & b = yb[e / QK8_1]; | |
| + b.qs[e % QK8_1] = amax == 0.0f ? 0 : (int8_t) roundf(xi / d); | |
| + if (e % QK8_1 == 0) { | |
| + b.ds = make_half2(d, sum); | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| +void ggml_cuda_op_swiglu_fwht(ggml_backend_cuda_context & ctx, const float * gate, const float * up, const float * signs, | |
| + int64_t signs_n, ggml_tensor * fwht_dst) { | |
| + cudaStream_t stream = ctx.stream(); | |
| + const int n = 1024; | |
| + const int64_t rows = ggml_nelements(fwht_dst) / n; | |
| + block_q8_1 * q8 = nullptr; | |
| + if (fwht_q8_enabled()) { | |
| + q8 = (block_q8_1 *) fwht_q8_slot(fwht_dst, (size_t) (ggml_nelements(fwht_dst) / QK8_1) * sizeof(block_q8_1), stream); | |
| + } | |
| + const float scale = 1 / sqrtf(n); | |
| + swiglu_fwht1024<<<(unsigned) rows, FWHT_BLOCK_THREADS, 0, stream>>>(gate, up, (float *) fwht_dst->data, scale, signs, | |
| + (int) (signs_n / n), q8); | |
| + CUDA_CHECK(cudaGetLastError()); | |
| + if (q8 != nullptr) { | |
| + g_fwht_q8[fwht_dst] = { q8, ggml_nelements(fwht_dst) }; | |
| + } | |
| +} | |
| + | |
| +// ---- fused GDN output chain + FWHT_1024 (sm_6x) ---- | |
| +// One block of 1024 threads per token: thread tid holds element tid of each 1024 chunk c (NB chunks = the whole row), | |
| +// i.e. permuted head 8c + (tid >> 7), channel d = tid & 127. Each head spans 4 warps with warp k on channels 32k.., | |
| +// the thread/column layout of rms_norm_f32<256> on a 128-wide row: per-warp butterfly sums, combined | |
| +// (w0 + w2) + (w1 + w3). Every load of the row happens before the first barrier, so the output may alias o or z | |
| +// row-for-row. FWHT stages h < 32 by shuffle, h >= 32 through shared memory: the same pairwise sums as | |
| +// fwht_cuda_block; a warp holds 32 consecutive elements of a chunk = one q8_1 block, quantized as there. | |
| +template <int NB> | |
| +static __global__ void __launch_bounds__(1024, 1) gdn_out_fwht_tok( | |
| + const float * o, const int64_t o_st, const float * z, const int64_t z_st, const float * wn, const float eps, | |
| + const int ncols, const int nk, const int rep, const float * signs, float * dst, block_q8_1 * q8, | |
| + const float scale) { | |
| + constexpr int N = 1024; | |
| + constexpr int warp_size = WARP_SIZE; | |
| + constexpr int NW = N / warp_size; | |
| + | |
| + __shared__ float s[NB][N]; | |
| + __shared__ float sv[NB][NW]; | |
| + | |
| + const int64_t t = blockIdx.x; | |
| + const int tid = threadIdx.x; | |
| + const int lane = tid % warp_size; | |
| + const int warp = tid / warp_size; | |
| + const int d = tid & 127; | |
| + const int hsub = tid >> 7; | |
| + | |
| + float xv[NB]; | |
| + float zv[NB]; | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; ++c) { | |
| + const int P = 8 * c + hsub; // permuted (grouped) head | |
| + const int hh = P / rep + nk * (P % rep); // the same head in the tiled (GDN output) order | |
| + xv[c] = o[t * o_st + hh * 128 + d]; | |
| + zv[c] = z[t * z_st + hh * 128 + d]; | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; ++c) { | |
| + float tmp = 0.0f; | |
| + tmp += xv[c] * xv[c]; | |
| + tmp = warp_reduce_sum(tmp); | |
| + if (lane == 0) { | |
| + sv[c][warp] = tmp; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + | |
| + float reg[NB]; | |
| + const int w0 = hsub * 4; | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; ++c) { | |
| + const float tot = (sv[c][w0] + sv[c][w0 + 2]) + (sv[c][w0 + 1] + sv[c][w0 + 3]); | |
| + const float mean = tot / ncols; | |
| + const float sc = rsqrtf(mean + eps); | |
| + const float nrm = sc * xv[c] * wn[d]; | |
| + const float v = ggml_cuda_op_silu_single(zv[c]) * nrm; | |
| + reg[c] = v * scale; | |
| + reg[c] *= signs[c * N + tid]; | |
| + } | |
| + | |
| +#pragma unroll | |
| + for (int h = 1; h < warp_size; h *= 2) { | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; c++) { | |
| + const float val = reg[c]; | |
| + const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size); | |
| + reg[c] = (lane & h) == 0 ? val + val2 : val2 - val; | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int h = warp_size; h < N; h *= 2) { | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; c++) { | |
| + s[c][tid] = reg[c]; | |
| + } | |
| + __syncthreads(); | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; c++) { | |
| + const float val = reg[c]; | |
| + const float val2 = s[c][tid ^ h]; | |
| + reg[c] = (tid & h) == 0 ? val + val2 : val2 - val; | |
| + } | |
| + __syncthreads(); | |
| + } | |
| + | |
| + float * drow = dst + t * NB * N; | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; ++c) { | |
| + drow[c * N + tid] = reg[c]; | |
| + } | |
| + | |
| + if (q8 != nullptr) { | |
| +#pragma unroll | |
| + for (int c = 0; c < NB; ++c) { | |
| + block_q8_1 * yb = q8 + (t * NB + c) * (N / QK8_1); | |
| + const float xi = reg[c]; | |
| + const float amax = warp_reduce_max<QK8_1>(fabsf(xi)); | |
| + const float sum = warp_reduce_sum<QK8_1>(xi); | |
| + const float dd = amax / 127.0f; | |
| + block_q8_1 & bq = yb[tid / QK8_1]; | |
| + bq.qs[tid % QK8_1] = amax == 0.0f ? 0 : (int8_t) roundf(xi / dd); | |
| + if (tid % QK8_1 == 0) { | |
| + bq.ds = make_half2(dd, sum); | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| +void ggml_cuda_op_gdn_out_fwht(ggml_backend_cuda_context & ctx, const float * o, int64_t o_st, const float * z, int64_t z_st, | |
| + const float * wn, float eps, int nk, int rep, const float * signs, ggml_tensor * fwht_dst, | |
| + int64_t n_tokens) { | |
| + cudaStream_t stream = ctx.stream(); | |
| + const int nb = nk * rep * 128 / 1024; | |
| + block_q8_1 * q8 = nullptr; | |
| + if (fwht_q8_enabled()) { | |
| + q8 = (block_q8_1 *) fwht_q8_slot(fwht_dst, (size_t) (ggml_nelements(fwht_dst) / QK8_1) * sizeof(block_q8_1), stream); | |
| + } | |
| + const int n = 1024; | |
| + const float scale = 1 / sqrtf(n); | |
| + float * dst = (float *) fwht_dst->data; | |
| + switch (nb) { | |
| + case 6: gdn_out_fwht_tok<6><<<(unsigned) n_tokens, 1024, 0, stream>>>(o, o_st, z, z_st, wn, eps, 128, nk, rep, signs, dst, q8, scale); break; | |
| + case 4: gdn_out_fwht_tok<4><<<(unsigned) n_tokens, 1024, 0, stream>>>(o, o_st, z, z_st, wn, eps, 128, nk, rep, signs, dst, q8, scale); break; | |
| + case 8: gdn_out_fwht_tok<8><<<(unsigned) n_tokens, 1024, 0, stream>>>(o, o_st, z, z_st, wn, eps, 128, nk, rep, signs, dst, q8, scale); break; | |
| + default: GGML_ABORT("gdn_out_fwht: unsupported row width"); | |
| + } | |
| + CUDA_CHECK(cudaGetLastError()); | |
| + if (q8 != nullptr) { | |
| + g_fwht_q8[fwht_dst] = { q8, ggml_nelements(fwht_dst) }; | |
| } | |
| - return fwht_launch<half>(ctx, (const half *) src->data, dst_d, n, rows, scale, signs, n_blk); | |
| } | |
| bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { | |
| diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh | |
| index 62b2f28..0f1d8fe 100644 | |
| --- a/ggml/src/ggml-cuda/fwht.cuh | |
| +++ b/ggml/src/ggml-cuda/fwht.cuh | |
| // Returns whether the Fast Walsh-Hadamard transform could be used. | |
| bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst); | |
| + | |
| +// q8_1 side output of the block FWHT (sm_6x): cleared per graph evaluation; find() returns the q8_1 copy of `t`'s base | |
| +// tensor when it was produced by a FWHT in this evaluation and covers exactly nelem elements from its start, else nullptr | |
| +void ggml_cuda_fwht_q8_clear(); | |
| + | |
| +// fused [a + b] -> rms_norm(eps) * w -> * signs -> FWHT_1024 per 1024-chunk; b / add_out may be null (no residual add). | |
| +// Writes add_out (if b), norm_out, fwht_dst (+ its q8_1 side output). rows of ncols (multiple of 1024), contiguous. | |
| +void ggml_cuda_op_add_rmsnorm_fwht(ggml_backend_cuda_context & ctx, const float * a, const float * b, float * add_out, | |
| + const float * w, float eps, float * norm_out, const float * signs, | |
| + int64_t ncols, int64_t nrows, ggml_tensor * fwht_dst); | |
| +const void * ggml_cuda_fwht_q8_find(const ggml_tensor * t, int64_t nelem); | |
| + | |
| +// fused SWIGLU + signs + FWHT_1024: silu(gate) * up -> * signs -> FWHT per 1024 chunk into fwht_dst (+ q8_1 side | |
| +// output). gate / up / fwht_dst contiguous F32 of the same element count (a multiple of 1024); signs repeat every | |
| +// signs_n elements. fwht_dst may alias gate or up element-for-element. | |
| +void ggml_cuda_op_swiglu_fwht(ggml_backend_cuda_context & ctx, const float * gate, const float * up, const float * signs, | |
| + int64_t signs_n, ggml_tensor * fwht_dst); | |
| + | |
| +// fused GDN output chain: rms_norm(o per 128-channel head, eps) * wn -> silu(z) * . -> head permutation tiled | |
| +// [128, nk, rep] -> grouped [128, rep, nk] -> * signs -> FWHT_1024 into fwht_dst (+ its q8_1 side output). | |
| +// o / z: [128, nk*rep, n_tokens] with token strides o_st / z_st (floats); nk*rep*128 a multiple of 1024. | |
| +void ggml_cuda_op_gdn_out_fwht(ggml_backend_cuda_context & ctx, const float * o, int64_t o_st, const float * z, int64_t z_st, | |
| + const float * wn, float eps, int nk, int rep, const float * signs, ggml_tensor * fwht_dst, | |
| + int64_t n_tokens); | |
| bool ggml_cuda_op_fwht_signed(ggml_backend_cuda_context & ctx, const ggml_tensor * src, | |
| const ggml_tensor * signs, ggml_tensor * dst); | |
| diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu | |
| index 5cf6968..a568667 100644 | |
| --- a/ggml/src/ggml-cuda/gated_delta_net.cu | |
| +++ b/ggml/src/ggml-cuda/gated_delta_net.cu | |
| #include "gated_delta_net.cuh" | |
| #include "ggml-cuda/common.cuh" | |
| +#include <unordered_map> | |
| + | |
| +static std::unordered_map<const ggml_tensor *, ggml_cuda_gdn_state_rows> g_gdn_state_rows; | |
| + | |
| +void ggml_cuda_gdn_state_rows_clear() { | |
| + g_gdn_state_rows.clear(); | |
| +} | |
| + | |
| +void ggml_cuda_gdn_state_rows_set(const ggml_tensor * gdn, ggml_cuda_gdn_state_rows sr) { | |
| + g_gdn_state_rows[gdn] = sr; | |
| +} | |
| + | |
| static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_t n) { | |
| for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n; | |
| i += (int64_t) blockDim.x*gridDim.x) { | |
| static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_ | |
| // RAW: beta and g arrive pre-activation (ggml_gated_delta_net_set_raw_gates); the kernel applies | |
| // sigmoid(beta) and raw_a[h] * softplus(g + raw_dt_bias[h]) with the unary kernels' formulas. | |
| // G_PRECOMPUTED: g already holds exp(g) (GB10 long-prompt path); only used with RAW == false. | |
| -template <int S_v, bool KDA, bool keep_rs_t, bool RAW, bool G_PRECOMPUTED> | |
| -__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2) | |
| +template <int S_v, bool KDA, bool keep_rs_t, bool RAW, bool G_PRECOMPUTED, int CPW_T = 0, int NW_T = 4> | |
| +__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * NW_T, 2) | |
| gated_delta_net_cuda(const float * q, | |
| const float * k, | |
| const float * v, | |
| gated_delta_net_cuda(const float * q, | |
| const uint3 rq3_magic, | |
| float scale, | |
| int64_t state_slot_stride, | |
| - int K) { | |
| + int K, | |
| + const int32_t * state_rows, | |
| + int64_t state_row_stride) { | |
| const uint32_t h_idx = blockIdx.x; | |
| const uint32_t sequence = blockIdx.y; | |
| // Each warp owns one or more columns, using warp-level primitives to reduce across rows. | |
| const int lane = threadIdx.x; | |
| -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK | |
| - constexpr int cols_per_warp = S_v == 128 && !KDA ? 4 : 1; | |
| +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA) // sm_6x too: 4 columns per warp share q/k loads + overlap the reductions | |
| + constexpr int cols_per_warp = CPW_T > 0 ? CPW_T : (S_v == 128 && !KDA ? 4 : 1); | |
| #else | |
| - constexpr int cols_per_warp = 1; | |
| + constexpr int cols_per_warp = CPW_T > 0 ? CPW_T : 1; | |
| #endif | |
| const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp; | |
| gated_delta_net_cuda(const float * q, | |
| // input state holds s0 only: [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v. | |
| // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before. | |
| - const int64_t state_in_offset = sequence * H * S_v * S_v + h_idx * S_v * S_v; | |
| + // folded gather: seq's state row in the cache, else the gathered [S_v, S_v, H, n_seqs] tensor | |
| + const int64_t state_in_offset = (state_rows ? (int64_t) state_rows[sequence] * state_row_stride | |
| + : (int64_t) sequence * H * S_v * S_v) + h_idx * S_v * S_v; | |
| const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v; | |
| state += state_out_offset; | |
| curr_state += state_in_offset; | |
| static void launch_gated_delta_net( | |
| int64_t sv1, int64_t sv2, int64_t sv3, | |
| int64_t sb1, int64_t sb2, int64_t sb3, | |
| int64_t neqk1, int64_t rq3, | |
| - float scale, int64_t state_slot_stride, int K, cudaStream_t stream) { | |
| + float scale, int64_t state_slot_stride, int K, const int32_t * state_rows, int64_t state_row_stride, | |
| + cudaStream_t stream) { | |
| //TODO: Add chunked kernel for even faster pre-fill | |
| const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; | |
| const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; | |
| - const int num_warps = 4; | |
| - const int cols_per_warp = cc == GGML_CUDA_CC_DGX_SPARK && S_v == 128 && !KDA ? 4 : 1; | |
| + int num_warps = 4; | |
| + int cols_per_warp = (cc == GGML_CUDA_CC_DGX_SPARK || (GGML_CUDA_CC_IS_NVIDIA(cc) && cc < GGML_CUDA_CC_VOLTA)) && S_v == 128 && !KDA ? 4 : 1; | |
| + // layout knobs (patch_fork_gdn_layout.py): bit-exact variants of the S_v = 128 non-KDA kernel | |
| + static const int env_cpw = getenv("GGML_CUDA_GDN_CPW") ? atoi(getenv("GGML_CUDA_GDN_CPW")) : 0; | |
| + static const int env_nw = getenv("GGML_CUDA_GDN_NW") ? atoi(getenv("GGML_CUDA_GDN_NW")) : 0; | |
| + const bool tuned = S_v == 128 && !KDA && (env_cpw > 0 || env_nw > 0); | |
| + if (tuned) { | |
| + if (env_cpw > 0) { cols_per_warp = env_cpw; } | |
| + if (env_nw > 0) { num_warps = env_nw; } | |
| + } | |
| dim3 grid_dims(H, n_seqs, (S_v + num_warps * cols_per_warp - 1) / (num_warps * cols_per_warp)); | |
| dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1); | |
| static void launch_gated_delta_net( | |
| ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, | |
| q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, | |
| n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, | |
| - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); | |
| + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, state_rows, state_row_stride); | |
| break; | |
| case 32: | |
| ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, | |
| q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, | |
| n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, | |
| - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); | |
| + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, state_rows, state_row_stride); | |
| break; | |
| case 64: { | |
| ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, | |
| q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, | |
| n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, | |
| - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); | |
| + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, state_rows, state_row_stride); | |
| break; | |
| } | |
| case 128: { | |
| + if (tuned) { | |
| +#define GDN_LAUNCH_T(CPW, NW) \ | |
| + ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, RAW, G_PRECOMPUTED, CPW, NW>, launch_params, \ | |
| + q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, \ | |
| + n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \ | |
| + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, state_rows, state_row_stride) | |
| + if (cols_per_warp == 1 && num_warps == 4) { GDN_LAUNCH_T(1, 4); } | |
| + else if (cols_per_warp == 2 && num_warps == 4) { GDN_LAUNCH_T(2, 4); } | |
| + else if (cols_per_warp == 4 && num_warps == 4) { GDN_LAUNCH_T(4, 4); } | |
| + else if (cols_per_warp == 8 && num_warps == 4) { GDN_LAUNCH_T(8, 4); } | |
| + else if (cols_per_warp == 1 && num_warps == 8) { GDN_LAUNCH_T(1, 8); } | |
| + else if (cols_per_warp == 2 && num_warps == 8) { GDN_LAUNCH_T(2, 8); } | |
| + else if (cols_per_warp == 4 && num_warps == 2) { GDN_LAUNCH_T(4, 2); } | |
| + else if (cols_per_warp == 2 && num_warps == 2) { GDN_LAUNCH_T(2, 2); } | |
| + else if (cols_per_warp == 1 && num_warps == 2) { GDN_LAUNCH_T(1, 2); } | |
| + else { GGML_ABORT("GGML_CUDA_GDN_CPW/NW: unsupported combination"); } | |
| +#undef GDN_LAUNCH_T | |
| + break; | |
| + } | |
| ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, RAW, G_PRECOMPUTED>, launch_params, | |
| q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H, | |
| n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, | |
| - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); | |
| + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, state_rows, state_row_stride); | |
| break; | |
| } | |
| default: | |
| static void ggml_cuda_op_gated_delta_net_impl( | |
| const float * s_d = (const float *) src_state->data; | |
| float * dst_d = (float *) dst->data; | |
| + // folded state gather (ggml_cuda_gdn_fold_gather): read the cache rows directly, src[5] is never written | |
| + const int32_t * state_rows_d = nullptr; | |
| + int64_t state_row_stride = 0; | |
| + { | |
| + const auto it = g_gdn_state_rows.find(dst); | |
| + if (it != g_gdn_state_rows.end()) { | |
| + s_d = it->second.base; | |
| + state_rows_d = it->second.rows; | |
| + state_row_stride = it->second.row_stride; | |
| + } | |
| + } | |
| + | |
| GGML_ASSERT(ggml_is_contiguous_rows(src_q)); | |
| GGML_ASSERT(ggml_is_contiguous_rows(src_k)); | |
| GGML_ASSERT(ggml_is_contiguous_rows(src_v)); | |
| static void ggml_cuda_op_gated_delta_net_impl( | |
| #define GDN_LAUNCH(KDA_, KEEP_, RAW_, PRE_) \ | |
| launch_gated_delta_net<KDA_, KEEP_, RAW_, PRE_>(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, \ | |
| S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \ | |
| - sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream) | |
| + sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, state_rows_d, state_row_stride, stream) | |
| if (kda) { | |
| if (keep_rs) { GDN_LAUNCH(true, true, false, false); } else { GDN_LAUNCH(true, false, false, false); } | |
| diff --git a/ggml/src/ggml-cuda/gated_delta_net.cuh b/ggml/src/ggml-cuda/gated_delta_net.cuh | |
| index f9bf437..33f2bf6 100644 | |
| --- a/ggml/src/ggml-cuda/gated_delta_net.cuh | |
| +++ b/ggml/src/ggml-cuda/gated_delta_net.cuh | |
| struct ggml_cuda_gated_delta_net_fused_cache { | |
| int64_t slot_stride; // between rollback slots (0 when K==1) | |
| }; | |
| +// folded state gather: the GDN op reads seq s's initial state at base + rows[s] * row_stride (elements) instead of | |
| +// src[5] (see ggml_cuda_gdn_fold_gather in ggml-cuda.cu). Registered per GDN node for one graph evaluation. | |
| +struct ggml_cuda_gdn_state_rows { | |
| + const float * base; | |
| + const int32_t * rows; | |
| + int64_t row_stride; | |
| +}; | |
| +void ggml_cuda_gdn_state_rows_clear(); | |
| +void ggml_cuda_gdn_state_rows_set(const ggml_tensor * gdn, ggml_cuda_gdn_state_rows sr); | |
| + | |
| void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst); | |
| // same op, but writes the snapshot(s) into the cache instead of dst (see ggml_cuda_try_gdn_cache_fusion) | |
| diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu | |
| index 6b4453b..ec8881a 100644 | |
| --- a/ggml/src/ggml-cuda/getrows.cu | |
| +++ b/ggml/src/ggml-cuda/getrows.cu | |
| static void get_rows_cuda_kq( | |
| s10, s11, s12/*, s13*/); | |
| } | |
| +// one thread per index for 1-element rows (ne00 == 1, ne11 == ne12 == 1): the block-per-index launcher is | |
| +// pathological there (the pruned draft-vocabulary scatter gathers ~248K single floats per draft token) | |
| +template<typename src0_t, typename dst_t> | |
| +static __global__ void k_get_rows_scalar(const src0_t * __restrict__ src0, const int32_t * __restrict__ src1, | |
| + dst_t * __restrict__ dst, const int64_t ne10, const size_t nb01, const size_t s1, const size_t s10) { | |
| + const int64_t i10 = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; | |
| + if (i10 >= ne10) { | |
| + return; | |
| + } | |
| + const int i01 = src1[i10*s10]; | |
| + dst[i10*s1] = ggml_cuda_cast<dst_t>(*(const src0_t *) ((const char *) src0 + (size_t) i01*nb01)); | |
| +} | |
| + | |
| template<typename src0_t, typename dst_t> | |
| static void get_rows_cuda_float( | |
| const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d, | |
| static void get_rows_cuda_float( | |
| GGML_ASSERT(ne12 > 0); | |
| GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12); | |
| + | |
| + if (ne00 == 1 && ne11 == 1 && ne12 == 1) { | |
| + const int64_t nblk = (ne10 + 255) / 256; | |
| + k_get_rows_scalar<src0_t, dst_t><<<(unsigned int) nblk, 256, 0, stream>>>(src0_d, src1_d, dst_d, ne10, nb01, s1, s10); | |
| + return; | |
| + } | |
| + | |
| const uint3 ne12_fdv = init_fastdiv_values(ne12); | |
| if constexpr (std::is_same<src0_t, dst_t>::value) { | |
| diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu | |
| index 7536a5d..8d04cbd 100644 | |
| --- a/ggml/src/ggml-cuda/ggml-cuda.cu | |
| +++ b/ggml/src/ggml-cuda/ggml-cuda.cu | |
| #include "ggml-cuda/mmq.cuh" | |
| #include "ggml-cuda/mmvf.cuh" | |
| #include "ggml-cuda/mmvq.cuh" | |
| +#include "ggml-cuda/mmvq-rowlane.cuh" | |
| #include "ggml-cuda/norm.cuh" | |
| #include "ggml-cuda/opt-step-adamw.cuh" | |
| #include "ggml-cuda/opt-step-sgd.cuh" | |
| #include <cstdio> | |
| #include <cstdlib> | |
| #include <string> | |
| +#include <unordered_map> | |
| +#include <unordered_set> | |
| #include <vector> | |
| static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); | |
| static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { | |
| if (cc <= GGML_CUDA_CC_PASCAL) { | |
| return false; | |
| } | |
| + // the fused path runs the generic mmvq kernel; where the PQ2_0 rowlane kernel applies (sm_61) it is faster | |
| + // unfused: a 1-token Bonsai 2 decode step 38.9 -> 34.8 ms on a GTX 1080 Ti (80 mm+add, 40 gate/up/swiglu per step) | |
| + if (tensor->op == GGML_OP_MUL_MAT && ggml_cuda_rowlane_applicable_k(src0->type, cc, dst->ne[1], src0->ne[0])) { | |
| + return false; | |
| + } | |
| //we only support fusion for ncols_dst = 1 | |
| if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { | |
| return false; | |
| static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { | |
| } | |
| static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { | |
| - return cgraph->nodes[0]; | |
| + // The first node's address alone is not a graph identity: llama rebuilds a graph into the same memory whenever the | |
| + // ubatch width changes, and speculative decoding alternates verify widths (1..n_max+1) on every cycle -- one cache | |
| + // entry then sees a "property change" on nearly every call, resets its warmup and never replays (measured on a GTX | |
| + // 1080 Ti: ~4-5 ms of eager kernel launches per verify step). Mix in the graph's shape signature so every width gets | |
| + // its own entry; a rebuilt graph of a width seen before then matches that entry and replays. The key is only | |
| + // hashed, never dereferenced. | |
| + uint64_t h = 0xcbf29ce484222325ull ^ (uint64_t) cgraph->n_nodes; | |
| + const ggml_tensor * a = cgraph->nodes[0]; | |
| + const ggml_tensor * b = cgraph->nodes[cgraph->n_nodes - 1]; | |
| + for (int d = 0; d < GGML_MAX_DIMS; ++d) { | |
| + h = (h ^ (uint64_t) a->ne[d]) * 0x100000001b3ull; | |
| + h = (h ^ (uint64_t) b->ne[d]) * 0x100000001b3ull; | |
| + } | |
| + return (const void *) ((uintptr_t) cgraph->nodes[0] ^ (uintptr_t) (h << 4)); | |
| } | |
| static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { | |
| static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, | |
| return false; | |
| } | |
| +// GET_ROWS nodes folded into a GDN kernel for the graph being evaluated (set by the evaluation loop) | |
| +static const std::unordered_set<const ggml_tensor *> * g_cuda_folded_gathers = nullptr; | |
| + | |
| +// conv-state gathers folded into the fused conv step (patch_fork_conv_fold.py), keyed by the CONCAT that reads them | |
| +struct ggml_cuda_conv_fold { | |
| + const float * base; // the gather's source rows (the recurrent cache view) | |
| + const int32_t * rows; // its I32 row index (device) | |
| + int64_t stride; // row stride, floats | |
| + const ggml_tensor * gather; // the skipped GET_ROWS (run it if the fused step does not match after all) | |
| +}; | |
| +static std::unordered_map<const ggml_tensor *, ggml_cuda_conv_fold> g_cuda_conv_folds; | |
| + | |
| +// GDN conv step: CONCAT(conv_state, qkv^T) -> K x CPY(snapshot window -> cache) -> SSM_CONV -> SILU -> L2_NORM(q,k) | |
| +// as one kernel (ssm-conv.cu), sm_6x. Views, empty nodes and folded gathers may sit between; any other compute | |
| +// node aborts the match. The CONCAT output is not written, so its only readers must be the snapshot views and the | |
| +// SSM_CONV. Returns the nodes to skip. GGML_CUDA_CONV_STEP=0 disables; GGML_CUDA_CONV_STEP_DEBUG=1 logs. | |
| +static int ggml_cuda_try_fuse_conv_step(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, bool dry = false, | |
| + const ggml_cuda_conv_fold * cand = nullptr) { | |
| + static const bool on = getenv("GGML_CUDA_CONV_STEP") == nullptr || std::atoi(getenv("GGML_CUDA_CONV_STEP")) != 0; | |
| + static const bool dbg = getenv("GGML_CUDA_CONV_STEP_DEBUG") != nullptr && std::atoi(getenv("GGML_CUDA_CONV_STEP_DEBUG")) != 0; | |
| + ggml_tensor * cat = cgraph->nodes[i]; | |
| + if (!on || cat->op != GGML_OP_CONCAT) { | |
| + return 0; | |
| + } | |
| + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; | |
| + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc >= GGML_CUDA_CC_VOLTA) { | |
| + return 0; | |
| + } | |
| + auto reject = [&](int why) { | |
| + if (dbg && !dry) { | |
| + fprintf(stderr, "CONVSTEP %s: reject %d\n", cat->name, why); | |
| + } | |
| + return 0; | |
| + }; | |
| + const ggml_cuda_conv_fold * fold = cand; | |
| + if (fold == nullptr && !g_cuda_conv_folds.empty()) { | |
| + const auto it = g_cuda_conv_folds.find(cat); | |
| + if (it != g_cuda_conv_folds.end()) { | |
| + fold = &it->second; | |
| + } | |
| + } | |
| + const ggml_tensor * st = cat->src[0]; // [3, C, 1] conv state (gathered) | |
| + const ggml_tensor * qt = cat->src[1]; // [n, C, 1] qkv, transposed view | |
| + if (cat->type != GGML_TYPE_F32 || ggml_get_op_params_i32(cat, 0) != 0 || !st || !qt || | |
| + st->type != GGML_TYPE_F32 || qt->type != GGML_TYPE_F32 || !ggml_is_contiguous(cat) || !ggml_is_contiguous(st) || | |
| + (cat->flags & GGML_TENSOR_FLAG_OUTPUT)) { | |
| + return reject(1); | |
| + } | |
| + const int64_t C = st->ne[1]; | |
| + const int64_t n = qt->ne[0]; | |
| + if (st->ne[0] != 3 || qt->ne[1] != C || st->ne[2] != 1 || qt->ne[2] != 1 || st->ne[3] != 1 || qt->ne[3] != 1 || | |
| + n < 1 || n > 8 || C % 128 != 0 || cat->ne[0] != 3 + n || cat->ne[1] != C || | |
| + qt->nb[0] % sizeof(float) != 0 || qt->nb[1] % sizeof(float) != 0) { | |
| + return reject(2); | |
| + } | |
| + ggml_cuda_conv_step_snaps snaps = {}; | |
| + int i_conv = -1, i_silu = -1, i_l2 = -1; | |
| + for (int j = i + 1; j < cgraph->n_nodes && i_l2 < 0; ++j) { | |
| + ggml_tensor * t = cgraph->nodes[j]; | |
| + if (ggml_cuda_is_view_or_noop(t) || (t->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { | |
| + continue; | |
| + } | |
| + if (g_cuda_folded_gathers && g_cuda_folded_gathers->count(t)) { | |
| + continue; | |
| + } | |
| + if (t->op == GGML_OP_CPY && i_conv < 0) { | |
| + const ggml_tensor * sv = t->src[0]; | |
| + const ggml_tensor * dv = t->src[1]; | |
| + if (sv->op != GGML_OP_VIEW || sv->view_src != cat || sv->type != GGML_TYPE_F32 || dv->type != GGML_TYPE_F32 || | |
| + t->type != GGML_TYPE_F32 || sv->ne[0] != 3 || sv->ne[1] != C || sv->ne[2] != 1 || sv->ne[3] != 1 || | |
| + sv->nb[0] != sizeof(float) || sv->nb[1] != cat->nb[1] || sv->view_offs % sizeof(float) != 0 || | |
| + !ggml_is_contiguous(dv) || ggml_nelements(dv) != 3 * C || snaps.n >= 8 || (t->flags & GGML_TENSOR_FLAG_OUTPUT) || | |
| + t->data != dv->data) { | |
| + return reject(3); | |
| + } | |
| + const int64_t sidx = (int64_t) (sv->view_offs / sizeof(float)); | |
| + if (sidx < 0 || sidx > n) { | |
| + return reject(4); | |
| + } | |
| + snaps.dst[snaps.n] = (float *) dv->data; | |
| + snaps.sidx[snaps.n] = (int) sidx; | |
| + snaps.n++; | |
| + continue; | |
| + } | |
| + if (t->op == GGML_OP_SSM_CONV && i_conv < 0) { | |
| + if (t->src[0] != cat) { | |
| + return reject(5); | |
| + } | |
| + i_conv = j; | |
| + continue; | |
| + } | |
| + if (t->op == GGML_OP_UNARY && i_conv >= 0 && i_silu < 0) { | |
| + if (ggml_get_unary_op(t) != GGML_UNARY_OP_SILU || t->src[0] != cgraph->nodes[i_conv]) { | |
| + return reject(6); | |
| + } | |
| + i_silu = j; | |
| + continue; | |
| + } | |
| + if (t->op == GGML_OP_L2_NORM && i_silu >= 0) { | |
| + i_l2 = j; | |
| + continue; | |
| + } | |
| + if (dbg && !dry) { | |
| + fprintf(stderr, "CONVSTEP %s: compute node in between: %s (%s)\n", cat->name, t->name, ggml_op_desc(t)); | |
| + } | |
| + return reject(7); | |
| + } | |
| + if (i_l2 < 0 || snaps.n == 0) { | |
| + return reject(8); | |
| + } | |
| + const ggml_tensor * conv = cgraph->nodes[i_conv]; | |
| + const ggml_tensor * silu = cgraph->nodes[i_silu]; | |
| + const ggml_tensor * l2n = cgraph->nodes[i_l2]; | |
| + const ggml_tensor * wt = conv->src[1]; // [4, C] | |
| + const ggml_tensor * lv = l2n->src[0]; // [128, H, n] view of the silu output | |
| + if (wt->type != GGML_TYPE_F32 || wt->ne[0] != 4 || wt->ne[1] != C || wt->nb[0] != sizeof(float) || | |
| + conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32 || l2n->type != GGML_TYPE_F32 || | |
| + silu->ne[0] != C || silu->ne[1] != n || silu->ne[2] != 1 || silu->nb[0] != sizeof(float) || | |
| + silu->nb[1] % sizeof(float) != 0 || (silu->flags & GGML_TENSOR_FLAG_OUTPUT) || (conv->flags & GGML_TENSOR_FLAG_OUTPUT)) { | |
| + return reject(9); | |
| + } | |
| + if (lv->op != GGML_OP_VIEW || lv->view_src != silu || lv->view_offs != 0 || lv->ne[0] != 128 || lv->ne[2] != n || | |
| + lv->ne[3] != 1 || lv->ne[1] * 128 > C || lv->nb[0] != sizeof(float) || lv->nb[1] != 128 * sizeof(float) || | |
| + lv->nb[2] != silu->nb[1] || !ggml_is_contiguous(l2n) || !ggml_are_same_shape(l2n, lv)) { | |
| + return reject(10); | |
| + } | |
| + // readers: the CONCAT output only through the snapshot views + the SSM_CONV; the raw conv output only by the SILU | |
| + if (ggml_node_get_use_count(cgraph, i) != snaps.n + 1 || ggml_node_get_use_count(cgraph, i_conv) != 1) { | |
| + if (dbg) { | |
| + fprintf(stderr, "CONVSTEP %s: uses cat %d (snaps %d) conv %d\n", cat->name, ggml_node_get_use_count(cgraph, i), | |
| + snaps.n, ggml_node_get_use_count(cgraph, i_conv)); | |
| + } | |
| + return reject(11); | |
| + } | |
| + // inputs must not overlap outputs (blocks read and write different channels) | |
| + auto span = [](const void * p, size_t bytes) { return std::make_pair((const char *) p, (const char *) p + bytes); }; | |
| + auto ovl = [](std::pair<const char *, const char *> a, std::pair<const char *, const char *> b) { | |
| + return a.first < b.second && b.first < a.second; | |
| + }; | |
| + // folded gather: the state is read from the cache row itself (not the compute buffer): no overlap with y / l2 / | |
| + // qkv; a snapshot slot equal to that row is read before written by the same thread | |
| + const auto s_st = fold ? span(nullptr, 0) : span(st->data, ggml_nbytes(st)); | |
| + const auto s_qt = span(qt->data, (size_t) ((n - 1) * qt->nb[0] + (C - 1) * qt->nb[1] + sizeof(float))); | |
| + const auto s_y = span(silu->data, ggml_nbytes(silu)); | |
| + const auto s_l2 = span(l2n->data, ggml_nbytes(l2n)); | |
| + bool bad = ovl(s_st, s_y); | |
| + // the qkv input may alias the silu output only exactly (same element -> same thread, read before write) | |
| + if (ovl(s_qt, s_y) && !(qt->data == silu->data && qt->nb[0] == silu->nb[1] && qt->nb[1] == sizeof(float))) { | |
| + bad = true; | |
| + } | |
| + for (int s = 0; s < snaps.n && !bad; ++s) { | |
| + const auto s_sn = span(snaps.dst[s], 3 * C * sizeof(float)); | |
| + bad = ovl(s_sn, s_st) || ovl(s_sn, s_qt) || ovl(s_sn, s_y) || ovl(s_sn, s_l2); | |
| + } | |
| + if (bad) { | |
| + if (dbg && !dry) { | |
| + fprintf(stderr, "CONVSTEP %s: overlap st=%p qt=%p y=%p l2=%p\n", cat->name, st->data, qt->data, silu->data, l2n->data); | |
| + } | |
| + return reject(12); | |
| + } | |
| + if (dry) { | |
| + return 1; // would fuse | |
| + } | |
| + // the allocator often gives the L2_NORM output the (by then dead) gathered conv state's buffer: blocks would then | |
| + // overwrite other blocks' inputs, so the L2_NORM stays a separate launch (fused up to the SILU) | |
| + const bool with_l2 = !ovl(s_l2, s_st) && !ovl(s_l2, s_qt) && !ovl(s_l2, s_y); | |
| + float eps; | |
| + memcpy(&eps, l2n->op_params, sizeof(float)); | |
| + ggml_cuda_conv_step_f32(fold ? fold->base : (const float *) st->data, fold ? fold->rows : nullptr, fold ? fold->stride : 0, | |
| + (const float *) qt->data, (int64_t) (qt->nb[0] / sizeof(float)), | |
| + (int64_t) (qt->nb[1] / sizeof(float)), (const float *) wt->data, (int64_t) (wt->nb[1] / sizeof(float)), | |
| + (float *) silu->data, (int64_t) (silu->nb[1] / sizeof(float)), (float *) l2n->data, | |
| + with_l2 ? (int) lv->ne[1] : 0, eps, snaps, C, n, cuda_ctx->stream()); | |
| + const int skip = with_l2 ? i_l2 - i : i_silu - i; | |
| + if (dbg) { | |
| + static int logged = 0; | |
| + if (logged++ < 4) { | |
| + fprintf(stderr, "CONVSTEP %s: fused n=%lld C=%lld snaps=%d l2 %s gather %s skip=%d\n", cat->name, (long long) n, | |
| + (long long) C, snaps.n, with_l2 ? "fused" : "separate", fold ? "folded" : "separate", skip); | |
| + } | |
| + } | |
| + return skip; | |
| +} | |
| + | |
| +// GDN output chain: RMS_NORM(128/head) -> MUL(w) -> SWIGLU(z, .) -> RESHAPE.. -> PERMUTE -> CONT -> RESHAPE -> MUL(signs) | |
| +// -> RESHAPE -> MUL_MAT(H_1024) as one kernel (fwht.cu), sm_6x. Every node up to the MUL_MAT must belong to the chain. | |
| +// Returns the nodes to skip. GGML_CUDA_GDN_OUT_FWHT=0 disables; GGML_CUDA_GDN_OUT_DEBUG=1 logs. | |
| +static int ggml_cuda_try_fuse_gdn_out_fwht(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { | |
| + static const bool on = getenv("GGML_CUDA_GDN_OUT_FWHT") == nullptr || std::atoi(getenv("GGML_CUDA_GDN_OUT_FWHT")) != 0; | |
| + static const bool dbg = getenv("GGML_CUDA_GDN_OUT_DEBUG") != nullptr && std::atoi(getenv("GGML_CUDA_GDN_OUT_DEBUG")) != 0; | |
| + ggml_tensor * norm = cgraph->nodes[i]; | |
| + if (!on || norm->op != GGML_OP_RMS_NORM) { | |
| + return 0; | |
| + } | |
| + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; | |
| + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc >= GGML_CUDA_CC_VOLTA) { | |
| + return 0; | |
| + } | |
| + auto reject = [&](int why) { | |
| + if (dbg) { | |
| + fprintf(stderr, "GDNOUT %s: reject %d\n", norm->name, why); | |
| + } | |
| + return 0; | |
| + }; | |
| + const ggml_tensor * o = norm->src[0]; // [128, H, n] view of the GDN output | |
| + if (norm->type != GGML_TYPE_F32 || o->type != GGML_TYPE_F32 || o->ne[0] != 128 || o->ne[3] != 1 || | |
| + o->nb[0] != sizeof(float) || o->nb[1] != 128 * sizeof(float) || o->nb[2] % sizeof(float) != 0 || | |
| + o->nb[2] < (size_t) o->ne[1] * 128 * sizeof(float) || !ggml_is_contiguous(norm) || (o->ne[1] * 128) % 1024 != 0) { | |
| + return 0; // not this chain (the other RMS_NORMs of the model land here too): no log | |
| + } | |
| + const int64_t H = o->ne[1]; | |
| + const int64_t n = o->ne[2]; | |
| + int idx[16]; | |
| + ggml_op ops[16]; | |
| + int cnt = 0; | |
| + idx[cnt] = i; | |
| + ops[cnt++] = GGML_OP_RMS_NORM; | |
| + ggml_tensor * mulw = nullptr, * glu = nullptr, * perm = nullptr, * cont = nullptr, * muls = nullptr, * mm = nullptr; | |
| + int i_mm = -1; | |
| + auto has_src = [](const ggml_tensor * t, const ggml_tensor * s) { | |
| + for (int k = 0; k < GGML_MAX_SRC; ++k) { | |
| + if (t->src[k] == s) { | |
| + return true; | |
| + } | |
| + } | |
| + return false; | |
| + }; | |
| + for (int j = i + 1; j < cgraph->n_nodes && !mm; ++j) { | |
| + ggml_tensor * t = cgraph->nodes[j]; | |
| + if (cnt >= 16) { | |
| + return reject(2); | |
| + } | |
| + if (!mulw) { | |
| + if (t->op != GGML_OP_MUL || !has_src(t, norm)) { | |
| + return 0; // a plain RMS_NORM + MUL elsewhere: no log | |
| + } | |
| + mulw = t; | |
| + } else if (!glu) { | |
| + if (t->op != GGML_OP_GLU || ggml_get_glu_op(t) != GGML_GLU_OP_SWIGLU || t->src[1] != mulw || | |
| + ggml_get_op_params_i32(t, 1) != 0) { | |
| + return 0; | |
| + } | |
| + glu = t; | |
| + } else if (!cont) { | |
| + if (t->op == GGML_OP_RESHAPE && !perm) { | |
| + // a reshape of the swiglu output on the way to the permute | |
| + } else if (t->op == GGML_OP_PERMUTE && !perm) { | |
| + perm = t; | |
| + } else if (t->op == GGML_OP_CONT && perm && t->src[0] == perm) { | |
| + cont = t; | |
| + } else { | |
| + return reject(3); | |
| + } | |
| + } else if (!muls) { | |
| + if (t->op == GGML_OP_RESHAPE) { | |
| + } else if (t->op == GGML_OP_MUL) { | |
| + muls = t; | |
| + } else { | |
| + return reject(4); | |
| + } | |
| + } else { | |
| + if (t->op == GGML_OP_RESHAPE) { | |
| + } else if (t->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(t, 1) == GGML_HINT_SRC0_IS_HADAMARD) { | |
| + mm = t; | |
| + i_mm = j; | |
| + } else { | |
| + return reject(5); | |
| + } | |
| + } | |
| + idx[cnt] = j; | |
| + ops[cnt++] = t->op; | |
| + } | |
| + if (!mm) { | |
| + return reject(6); | |
| + } | |
| + auto f32c = [](const ggml_tensor * t) { return t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); }; | |
| + const ggml_tensor * wn = mulw->src[0] == norm ? mulw->src[1] : mulw->src[0]; | |
| + const ggml_tensor * z = glu->src[0]; | |
| + const ggml_tensor * p4 = perm->src[0]; // [128, nk, rep, n] | |
| + if ((H * 128 / 1024) != 4 && (H * 128 / 1024) != 6 && (H * 128 / 1024) != 8) { | |
| + return reject(12); | |
| + } | |
| + if (!f32c(wn) || wn->ne[0] != 128 || ggml_nrows(wn) != 1 || !f32c(mulw) || !ggml_are_same_shape(mulw, norm) || | |
| + !f32c(glu) || !ggml_are_same_shape(glu, norm) || z->type != GGML_TYPE_F32 || !ggml_are_same_shape(z, norm) || | |
| + z->nb[0] != sizeof(float) || z->nb[1] != 128 * sizeof(float) || z->nb[2] % sizeof(float) != 0) { | |
| + return reject(7); | |
| + } | |
| + const int64_t nk = p4->ne[1]; | |
| + const int64_t rep = p4->ne[2]; | |
| + if (p4->op != GGML_OP_RESHAPE || p4->view_src != glu || p4->ne[0] != 128 || nk * rep != H || p4->ne[3] != n || | |
| + perm->view_src != glu || perm->ne[0] != 128 || perm->ne[1] != rep || perm->ne[2] != nk || perm->ne[3] != n || | |
| + perm->nb[1] != p4->nb[2] || perm->nb[2] != p4->nb[1] || perm->nb[3] != p4->nb[3] || | |
| + !f32c(cont) || !ggml_are_same_shape(cont, perm)) { | |
| + return reject(8); | |
| + } | |
| + const ggml_tensor * sg = nullptr; | |
| + for (int k = 0; k < 2; ++k) { | |
| + const ggml_tensor * a = muls->src[k]; | |
| + const ggml_tensor * c = muls->src[1 - k]; | |
| + if (a && c && a->op == GGML_OP_RESHAPE && a->view_src == cont && a->ne[0] == H * 128) { | |
| + sg = c; | |
| + } | |
| + } | |
| + const ggml_tensor * mr = mm->src[1]; | |
| + if (!sg || !f32c(sg) || sg->ne[0] != H * 128 || ggml_nrows(sg) != 1 || !f32c(muls) || muls->ne[0] != H * 128 || | |
| + mr->op != GGML_OP_RESHAPE || mr->view_src != muls || mr->ne[0] != 1024 || mm->src[0]->ne[0] != 1024 || | |
| + mm->src[0]->ne[1] != 1024 || !f32c(mm) || mm->ne[0] != 1024 || ggml_nelements(mm) != H * 128 * n) { | |
| + return reject(9); | |
| + } | |
| + const int outs[1] = { i_mm }; | |
| + if (!ggml_can_fuse_subgraph_ext(cgraph, idx, cnt, ops, outs, 1)) { | |
| + return reject(10); | |
| + } | |
| + // one block per token reads its whole row before writing: the output may alias o or z only row-for-row | |
| + const char * d0 = (const char *) mm->data; | |
| + const char * d1 = d0 + ggml_nbytes(mm); | |
| + auto ovl = [&](const void * p, size_t bytes) { | |
| + const char * a = (const char *) p; | |
| + return a < d1 && d0 < a + bytes; | |
| + }; | |
| + const size_t row_b = (size_t) H * 128 * sizeof(float); | |
| + auto same_rows = [&](const ggml_tensor * x) { return x->data == mm->data && x->nb[2] == row_b; }; | |
| + if ((ovl(o->data, (size_t) ((n - 1) * o->nb[2]) + row_b) && !same_rows(o)) || | |
| + (ovl(z->data, (size_t) ((n - 1) * z->nb[2]) + row_b) && !same_rows(z))) { | |
| + if (dbg) { | |
| + fprintf(stderr, "GDNOUT %s: output overlaps an input (o=%p z=%p dst=%p)\n", norm->name, o->data, z->data, mm->data); | |
| + } | |
| + return reject(11); | |
| + } | |
| + float eps; | |
| + memcpy(&eps, norm->op_params, sizeof(float)); | |
| + ggml_cuda_op_gdn_out_fwht(*cuda_ctx, (const float *) o->data, (int64_t) (o->nb[2] / sizeof(float)), | |
| + (const float *) z->data, (int64_t) (z->nb[2] / sizeof(float)), (const float *) wn->data, eps, | |
| + (int) nk, (int) rep, (const float *) sg->data, mm, n); | |
| + if (dbg) { | |
| + static int logged = 0; | |
| + if (logged++ < 4) { | |
| + fprintf(stderr, "GDNOUT %s: fused n=%lld H=%lld nk=%lld rep=%lld skip=%d\n", norm->name, (long long) n, | |
| + (long long) H, (long long) nk, (long long) rep, i_mm - i); | |
| + } | |
| + } | |
| + return i_mm - i; | |
| +} | |
| + | |
| +// Two BF16 MUL_MATs of the same shape on the same F32 activations (the GDN alpha / beta projections), only views in | |
| +// between, n <= 8 columns: one launch (mmvf.cu, bit-identical to mul_mat_vec_f). sm_6x. GGML_CUDA_BF16_DUAL=0 | |
| +// disables; GGML_CUDA_BF16_DUAL_RPB picks rows per block (2, 4 default, 8). | |
| +static int ggml_cuda_try_fuse_bf16_dual(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { | |
| + static const bool on = getenv("GGML_CUDA_BF16_DUAL") == nullptr || std::atoi(getenv("GGML_CUDA_BF16_DUAL")) != 0; | |
| + static const int rpb = getenv("GGML_CUDA_BF16_DUAL_RPB") ? std::atoi(getenv("GGML_CUDA_BF16_DUAL_RPB")) : 2; | |
| + ggml_tensor * a = cgraph->nodes[i]; | |
| + if (!on || a->op != GGML_OP_MUL_MAT || a->src[0]->type != GGML_TYPE_BF16) { | |
| + return 0; | |
| + } | |
| + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; | |
| + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc >= GGML_CUDA_CC_VOLTA) { | |
| + return 0; | |
| + } | |
| + int j = i + 1; | |
| + while (j < cgraph->n_nodes && ggml_cuda_is_view_or_noop(cgraph->nodes[j])) { | |
| + ++j; | |
| + } | |
| + if (j >= cgraph->n_nodes) { | |
| + return 0; | |
| + } | |
| + ggml_tensor * b = cgraph->nodes[j]; | |
| + const ggml_tensor * xa = a->src[0]; | |
| + const ggml_tensor * xb = b->src[0]; | |
| + const ggml_tensor * y = a->src[1]; | |
| + if (b->op != GGML_OP_MUL_MAT || b->src[1] != y || xb->type != GGML_TYPE_BF16 || !ggml_are_same_shape(xa, xb) || | |
| + xa->nb[1] != xb->nb[1] || !ggml_is_contiguous(xa) || !ggml_is_contiguous(xb) || | |
| + y->type != GGML_TYPE_F32 || !ggml_is_contiguous(y) || y->ne[2] != 1 || y->ne[3] != 1 || xa->ne[2] != 1 || xa->ne[3] != 1 || | |
| + y->ne[1] < 1 || y->ne[1] > 8 || xa->ne[0] != y->ne[0] || xa->ne[0] % 2 != 0 || (2 * xa->ne[1]) % rpb != 0 || | |
| + a->type != GGML_TYPE_F32 || b->type != GGML_TYPE_F32 || !ggml_is_contiguous(a) || !ggml_is_contiguous(b) || | |
| + (a->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || (b->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { | |
| + return 0; | |
| + } | |
| + // outputs must not overlap the activations (they are tiny; the allocator keeps them apart while y is live) | |
| + const char * y0 = (const char *) y->data; | |
| + const char * y1 = y0 + ggml_nbytes(y); | |
| + auto ovl_y = [&](const ggml_tensor * t) { | |
| + const char * p0 = (const char *) t->data; | |
| + return p0 < y1 && y0 < p0 + ggml_nbytes(t); | |
| + }; | |
| + if (ovl_y(a) || ovl_y(b)) { | |
| + return 0; | |
| + } | |
| + ggml_cuda_mul_mat_vec_bf16_dual(xa->data, xb->data, (int64_t) (xa->nb[1] / ggml_type_size(GGML_TYPE_BF16)), xa->ne[1], | |
| + xa->ne[0], (const float *) y->data, (int64_t) (y->nb[1] / sizeof(float)), y->ne[1], | |
| + (float *) a->data, (float *) b->data, (int64_t) (a->nb[1] / sizeof(float)), rpb, | |
| + cuda_ctx->stream()); | |
| + return j - i; | |
| +} | |
| + | |
| +// use count of any graph tensor (0 if it is not in the graph's hash set) | |
| +static int ggml_cuda_tensor_use_count(const ggml_cgraph * cgraph, const ggml_tensor * t) { | |
| + const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, t); | |
| + if (pos == GGML_HASHSET_FULL || !ggml_bitset_get(cgraph->visited_hash_set.used, pos)) { | |
| + return 0; | |
| + } | |
| + return cgraph->use_counts[pos]; | |
| +} | |
| + | |
| +// pre-pass: fold the conv-state GET_ROWS of every CONCAT the conv-step fusion will take (dry run) into it | |
| +static void ggml_cuda_conv_fold_gather(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, | |
| + std::unordered_set<const ggml_tensor *> & skip) { | |
| + g_cuda_conv_folds.clear(); | |
| + static const bool on = (getenv("GGML_CUDA_CONV_FOLD") == nullptr || std::atoi(getenv("GGML_CUDA_CONV_FOLD")) != 0) && | |
| + !(getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")) != 0); | |
| + if (!on) { | |
| + return; | |
| + } | |
| + for (int i = 0; i < cgraph->n_nodes; ++i) { | |
| + ggml_tensor * cat = cgraph->nodes[i]; | |
| + if (cat->op != GGML_OP_CONCAT || (cat->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { | |
| + continue; | |
| + } | |
| + const ggml_tensor * st = cat->src[0]; | |
| + if (!st || st->op != GGML_OP_RESHAPE || st->view_src == nullptr) { | |
| + continue; | |
| + } | |
| + const ggml_tensor * g = st->view_src; | |
| + if (g->op != GGML_OP_GET_ROWS || g->type != GGML_TYPE_F32 || g->src[0]->type != GGML_TYPE_F32 || | |
| + g->src[1]->type != GGML_TYPE_I32 || g->ne[1] != 1 || g->ne[2] != 1 || g->ne[3] != 1 || g->src[1]->ne[0] != 1 || | |
| + !ggml_is_contiguous(g) || g->src[0]->nb[0] != sizeof(float) || g->src[0]->ne[0] != g->ne[0] || | |
| + g->src[0]->nb[1] % sizeof(float) != 0 || (g->flags & GGML_TENSOR_FLAG_OUTPUT) || skip.count(g) || | |
| + ggml_cuda_tensor_use_count(cgraph, g) != 1 || ggml_cuda_tensor_use_count(cgraph, st) != 1) { | |
| + continue; | |
| + } | |
| + const ggml_cuda_conv_fold cand = { (const float *) g->src[0]->data, (const int32_t *) g->src[1]->data, | |
| + (int64_t) (g->src[0]->nb[1] / sizeof(float)), g }; | |
| + if (ggml_cuda_try_fuse_conv_step(cuda_ctx, cgraph, i, /*dry =*/ true, &cand) > 0) { | |
| + g_cuda_conv_folds[cat] = cand; | |
| + skip.insert(g); | |
| + } | |
| + } | |
| +} | |
| + | |
| // try and fuse nodes and return the number of nodes to skip | |
| +static int ggml_cuda_try_fuse_impl(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i); | |
| +// diagnostic: GGML_CUDA_FUSE_LOG=N logs the first N fusions (ops fused, first node name, its ne[1]) | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { | |
| + const int r = ggml_cuda_try_fuse_impl(cuda_ctx, cgraph, i); | |
| + static const int log_n = getenv("GGML_CUDA_FUSE_LOG") ? std::atoi(getenv("GGML_CUDA_FUSE_LOG")) : 0; | |
| + static int logged = 0; | |
| + if (r > 0 && logged < log_n) { | |
| + logged++; | |
| + std::string ops; | |
| + for (int k = 0; k <= r && i + k < cgraph->n_nodes; ++k) { | |
| + ops += ggml_op_desc(cgraph->nodes[i + k]); | |
| + ops += "("; | |
| + ops += ggml_type_name(cgraph->nodes[i + k]->src[0] ? cgraph->nodes[i + k]->src[0]->type : GGML_TYPE_F32); | |
| + ops += ") "; | |
| + } | |
| + GGML_LOG_INFO("FUSE n_nodes=%d i=%d skip=%d ne1=%lld %s| %s\n", cgraph->n_nodes, i, r, | |
| + (long long) cgraph->nodes[i]->ne[1], ops.c_str(), cgraph->nodes[i]->name); | |
| + } | |
| + return r; | |
| +} | |
| + | |
| +// [ADD] -> RMS_NORM -> MUL(w) -> MUL(signs) -> RESHAPE -> MUL_MAT(H_1024) as one kernel (fwht.cu), sm_6x. | |
| +// Returns the nodes to skip (0 = no match). Only views/no-ops may sit between the matched nodes. | |
| +static int ggml_cuda_try_fuse_norm_fwht(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { | |
| + static const bool on = getenv("GGML_CUDA_NORM_FWHT") == nullptr || std::atoi(getenv("GGML_CUDA_NORM_FWHT")) != 0; | |
| + static const bool nf_dbg = getenv("GGML_CUDA_NF_DEBUG") != nullptr; // read once: this runs for every ADD / RMS_NORM | |
| + if (!on) { | |
| + return 0; | |
| + } | |
| + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; | |
| + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc >= GGML_CUDA_CC_VOLTA) { | |
| + return 0; | |
| + } | |
| + ggml_tensor * n0 = cgraph->nodes[i]; | |
| + const bool has_add = n0->op == GGML_OP_ADD; | |
| + if (!has_add && n0->op != GGML_OP_RMS_NORM) { | |
| + return 0; | |
| + } | |
| + auto folded = [](const ggml_tensor * t) { return g_cuda_folded_gathers != nullptr && g_cuda_folded_gathers->count(t) > 0; }; | |
| + auto next = [&](int j) { | |
| + for (++j; j < cgraph->n_nodes; ++j) { | |
| + if (!ggml_cuda_is_view_or_noop(cgraph->nodes[j]) && !folded(cgraph->nodes[j])) { | |
| + return j; | |
| + } | |
| + } | |
| + return -1; | |
| + }; | |
| + const int i_norm = has_add ? next(i) : i; | |
| + const int i_mulw = i_norm >= 0 ? next(i_norm) : -1; | |
| + const int i_muls = i_mulw >= 0 ? next(i_mulw) : -1; | |
| + const int i_mm = i_muls >= 0 ? next(i_muls) : -1; | |
| + if (i_mm < 0) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 1\n", n0->name); } | |
| + return 0; | |
| + } | |
| + ggml_tensor * norm = cgraph->nodes[i_norm]; | |
| + ggml_tensor * mulw = cgraph->nodes[i_mulw]; | |
| + ggml_tensor * muls = cgraph->nodes[i_muls]; | |
| + ggml_tensor * mm = cgraph->nodes[i_mm]; | |
| + if (norm->op != GGML_OP_RMS_NORM || mulw->op != GGML_OP_MUL || muls->op != GGML_OP_MUL || mm->op != GGML_OP_MUL_MAT || | |
| + ggml_get_op_params_i32(mm, 1) != GGML_HINT_SRC0_IS_HADAMARD) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 2\n", n0->name); } | |
| + return 0; | |
| + } | |
| + if (has_add && norm->src[0] != n0) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 3\n", n0->name); } | |
| + return 0; | |
| + } | |
| + const ggml_tensor * x = has_add ? n0 : norm->src[0]; | |
| + const ggml_tensor * wt = mulw->src[0] == norm ? mulw->src[1] : (mulw->src[1] == norm ? mulw->src[0] : nullptr); | |
| + const ggml_tensor * sg = muls->src[0] == mulw ? muls->src[1] : (muls->src[1] == mulw ? muls->src[0] : nullptr); | |
| + const ggml_tensor * rs = mm->src[1]; | |
| + if (!wt || !sg || rs->op != GGML_OP_RESHAPE || rs->src[0] != muls) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 4\n", n0->name); } | |
| + return 0; | |
| + } | |
| + int i_rs = -1; | |
| + for (int j = i_muls + 1; j < i_mm; ++j) { | |
| + if (cgraph->nodes[j] == rs) { | |
| + i_rs = j; | |
| + } | |
| + } | |
| + if (i_rs < 0) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 5\n", n0->name); } | |
| + return 0; | |
| + } | |
| + const bool dbg = nf_dbg; | |
| + for (int j = i + 1; j < i_mm; ++j) { | |
| + if (j != i_norm && j != i_mulw && j != i_muls && j != i_rs && !ggml_cuda_is_view_or_noop(cgraph->nodes[j]) && | |
| + !folded(cgraph->nodes[j])) { | |
| + if (dbg) { fprintf(stderr, "NFDBG %s: compute node in between: %s (%s)\n", n0->name, cgraph->nodes[j]->name, ggml_op_desc(cgraph->nodes[j])); } | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 6\n", n0->name); } | |
| + return 0; // another compute node in between would be skipped | |
| + } | |
| + } | |
| + const int64_t ncols = x->ne[0]; | |
| + const int64_t nrows = ggml_nrows(x); | |
| + auto f32c = [](const ggml_tensor * t) { return t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); }; | |
| + if (ncols % 1024 != 0 || mm->ne[0] != 1024 || mm->src[0]->ne[0] != 1024 || mm->type != GGML_TYPE_F32 || | |
| + !f32c(x) || !f32c(norm) || !f32c(mulw) || !f32c(muls) || !f32c(mm) || !f32c(wt) || !f32c(sg) || | |
| + wt->ne[0] != ncols || ggml_nrows(wt) != 1 || sg->ne[0] != ncols || ggml_nrows(sg) != 1 || | |
| + !ggml_are_same_shape(x, mulw) || !ggml_are_same_shape(x, muls) || ggml_nelements(mm) != ggml_nelements(x) || | |
| + nrows > 65535) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 7\n", n0->name); } | |
| + return 0; | |
| + } | |
| + const ggml_tensor * a = x; | |
| + const ggml_tensor * b = nullptr; | |
| + if (has_add) { | |
| + a = n0->src[0]; | |
| + b = n0->src[1]; | |
| + if (!f32c(a) || !f32c(b) || !ggml_are_same_shape(a, n0) || !ggml_are_same_shape(b, n0)) { | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 8\n", n0->name); } | |
| + return 0; | |
| + } | |
| + } | |
| + int idxs[6]; | |
| + ggml_op ops[6]; | |
| + int n = 0; | |
| + if (has_add) { | |
| + idxs[n] = i; ops[n++] = GGML_OP_ADD; | |
| + } | |
| + idxs[n] = i_norm; ops[n++] = GGML_OP_RMS_NORM; | |
| + idxs[n] = i_mulw; ops[n++] = GGML_OP_MUL; | |
| + idxs[n] = i_muls; ops[n++] = GGML_OP_MUL; | |
| + idxs[n] = i_rs; ops[n++] = GGML_OP_RESHAPE; | |
| + idxs[n] = i_mm; ops[n++] = GGML_OP_MUL_MAT; | |
| + int outs[3]; | |
| + int n_out = 0; | |
| + if (has_add) { | |
| + outs[n_out++] = i; | |
| + } | |
| + outs[n_out++] = i_mulw; | |
| + outs[n_out++] = i_mm; | |
| + if (!ggml_can_fuse_subgraph_ext(cgraph, idxs, n, ops, outs, n_out)) { | |
| + if (dbg) { fprintf(stderr, "NFDBG %s: subgraph check failed (norm %s uses %d, muls uses %d, rs uses %d)\n", n0->name, norm->name, | |
| + ggml_node_get_use_count(cgraph, i_norm), ggml_node_get_use_count(cgraph, i_muls), ggml_node_get_use_count(cgraph, i_rs)); } | |
| + if (nf_dbg && has_add && strncmp(n0->name, "l_out", 5) == 0) { fprintf(stderr, "NFDBG %s: reject 9\n", n0->name); } | |
| + return 0; | |
| + } | |
| + float eps; | |
| + memcpy(&eps, norm->op_params, sizeof(float)); | |
| + ggml_cuda_op_add_rmsnorm_fwht(*cuda_ctx, (const float *) a->data, b ? (const float *) b->data : nullptr, | |
| + has_add ? (float *) n0->data : nullptr, (const float *) wt->data, eps, | |
| + (float *) mulw->data, (const float *) sg->data, ncols, nrows, mm); | |
| + return i_mm - i; | |
| +} | |
| + | |
| +static int ggml_cuda_try_fuse_impl(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { | |
| static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); | |
| if (disable_fusion) { | |
| return 0; | |
| } | |
| + { | |
| + const int skip = ggml_cuda_try_fuse_norm_fwht(cuda_ctx, cgraph, i); | |
| + if (skip > 0) { | |
| + return skip; | |
| + } | |
| + } | |
| + { | |
| + const int skip = ggml_cuda_try_fuse_conv_step(cuda_ctx, cgraph, i); | |
| + if (skip > 0) { | |
| + return skip; | |
| + } | |
| + if (!g_cuda_conv_folds.empty() && cgraph->nodes[i]->op == GGML_OP_CONCAT) { | |
| + // the gather was folded for a fused step that did not match after all: produce its output now | |
| + const auto it = g_cuda_conv_folds.find(cgraph->nodes[i]); | |
| + if (it != g_cuda_conv_folds.end()) { | |
| + ggml_cuda_op_get_rows(*cuda_ctx, (ggml_tensor *) it->second.gather); | |
| + g_cuda_conv_folds.erase(it); | |
| + } | |
| + } | |
| + } | |
| + { | |
| + const int skip = ggml_cuda_try_fuse_gdn_out_fwht(cuda_ctx, cgraph, i); | |
| + if (skip > 0) { | |
| + return skip; | |
| + } | |
| + } | |
| + { | |
| + const int skip = ggml_cuda_try_fuse_bf16_dual(cuda_ctx, cgraph, i); | |
| + if (skip > 0) { | |
| + return skip; | |
| + } | |
| + } | |
| + // diagnostic: GGML_CUDA_FUSE_OFF bitmask disables single fusion families (1 fwht, 2 gdn-cache, 4 multi-add/mul, | |
| + // 8 rope/set_rows, 16 rms_norm+mul, 32 ssm_conv+silu, 64 unary+mul) | |
| + static const int fuse_off = getenv("GGML_CUDA_FUSE_OFF") ? std::atoi(getenv("GGML_CUDA_FUSE_OFF")) : 0; | |
| static const bool dual_rms_q8_enabled = [] { | |
| const char * env = getenv("GGML_CUDA_GB10_DUAL_RMS_Q8"); | |
| return !env || std::atoi(env) != 0; | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; | |
| const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }; | |
| const int out_nodes[] = { i, i + 2 }; | |
| - if (cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && | |
| + if ((cc == GGML_CUDA_CC_DGX_SPARK || cc < GGML_CUDA_CC_VOLTA) && rms_norm->op == GGML_OP_RMS_NORM && | |
| mul->op == GGML_OP_MUL && (mul->src[0] == rms_norm || mul->src[1] == rms_norm) && | |
| rms_norm->src[0] == node && node->src[0] && node->src[1] && | |
| node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| } | |
| // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache | |
| - if (node->op == GGML_OP_GATED_DELTA_NET) { | |
| + if (!(fuse_off & 2) && node->op == GGML_OP_GATED_DELTA_NET) { | |
| ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; | |
| const int nodes_to_skip = ggml_cuda_try_gdn_cache_fusion(cgraph, i, fused_state_cpy); | |
| if (nodes_to_skip > 0) { | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| } | |
| } | |
| + // SWIGLU + sign flip + reshape + FWHT_1024 (sm_6x): the swiglu is computed in the transform's load stage | |
| + // (fwht.cu swiglu_fwht1024). GGML_CUDA_SWIGLU_FWHT=0 disables. | |
| + { | |
| + static const bool sf_on = getenv("GGML_CUDA_SWIGLU_FWHT") == nullptr || std::atoi(getenv("GGML_CUDA_SWIGLU_FWHT")) != 0; | |
| + const int sf_cc = ggml_cuda_info().devices[cuda_ctx->device].cc; | |
| + if (sf_on && GGML_CUDA_CC_IS_NVIDIA(sf_cc) && sf_cc < GGML_CUDA_CC_VOLTA && !(fuse_off & 1) && | |
| + cgraph->nodes[i]->op == GGML_OP_GLU && i + 3 < cgraph->n_nodes && | |
| + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_GLU, GGML_OP_MUL, GGML_OP_RESHAPE, GGML_OP_MUL_MAT }, { i + 3 })) { | |
| + const ggml_tensor * glu = cgraph->nodes[i]; | |
| + const ggml_tensor * mul = cgraph->nodes[i + 1]; | |
| + const ggml_tensor * reshape = cgraph->nodes[i + 2]; | |
| + ggml_tensor * mm = cgraph->nodes[i + 3]; | |
| + const ggml_tensor * g = glu->src[0]; | |
| + const ggml_tensor * u = glu->src[1]; | |
| + const ggml_tensor * sg = mul->src[0] == glu ? mul->src[1] : mul->src[0]; | |
| + auto f32c = [](const ggml_tensor * t) { return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); }; | |
| + const bool ok = ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU && ggml_get_op_params_i32(glu, 1) == 0 && | |
| + f32c(g) && f32c(u) && f32c(glu) && f32c(mul) && f32c(sg) && f32c(mm) && | |
| + ggml_are_same_shape(g, glu) && ggml_are_same_shape(u, glu) && ggml_are_same_shape(mul, glu) && | |
| + (mul->src[0] == glu || mul->src[1] == glu) && | |
| + ggml_get_op_params_i32(mm, 1) == GGML_HINT_SRC0_IS_HADAMARD && mm->src[1] == reshape && reshape->src[0] == mul && | |
| + mm->src[0]->ne[0] == 1024 && mm->src[0]->ne[1] == 1024 && mm->ne[0] == 1024 && | |
| + ggml_nelements(mm) == ggml_nelements(glu) && glu->ne[0] % 1024 == 0 && | |
| + sg->ne[0] == glu->ne[0] && ggml_nrows(sg) == 1; | |
| + if (ok) { | |
| + // the output may alias gate or up only element-for-element (each block reads its chunk before writing it) | |
| + const char * d0 = (const char *) mm->data; | |
| + const char * d1 = d0 + ggml_nbytes(mm); | |
| + auto bad = [&](const ggml_tensor * t) { | |
| + const char * a = (const char *) t->data; | |
| + return a < d1 && d0 < a + ggml_nbytes(t) && a != d0; | |
| + }; | |
| + if (!bad(g) && !bad(u)) { | |
| + ggml_cuda_op_swiglu_fwht(*cuda_ctx, (const float *) g->data, (const float *) u->data, (const float *) sg->data, | |
| + sg->ne[0], mm); | |
| + return 3; | |
| + } | |
| + } | |
| + } | |
| + } | |
| + | |
| // Hadamard sign flip + reshape + FWHT-hint matmul: multiply the sign | |
| // vector during the transform's load instead of a separate full pass | |
| - if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_RESHAPE, GGML_OP_MUL_MAT }, { i + 2 })) { | |
| + if (!(fuse_off & 1) && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_RESHAPE, GGML_OP_MUL_MAT }, { i + 2 })) { | |
| const ggml_tensor * mul = cgraph->nodes[i]; | |
| const ggml_tensor * reshape = cgraph->nodes[i + 1]; | |
| ggml_tensor * mm = cgraph->nodes[i + 2]; | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| } | |
| //RoPE + view + set-rows | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { | |
| + if (!(fuse_off & 8) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { | |
| ggml_tensor * rope = cgraph->nodes[i]; | |
| ggml_tensor * set_rows = cgraph->nodes[i + 2]; | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| } | |
| // multi-(add or mul) | |
| - if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { | |
| + if (!(fuse_off & 4) && (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL)) { | |
| int n_fuse = 0; | |
| ggml_op ops[8]; | |
| std::fill(ops, ops + 8, node->op); | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| return fused_node_count - 1; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { | |
| + if (!(fuse_off & 8) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { | |
| ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]); | |
| return 4; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {})) { | |
| + if (!(fuse_off & 8) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {})) { | |
| ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], nullptr); | |
| return 2; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { | |
| + if (!(fuse_off & 16) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { | |
| ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); | |
| return 2; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { | |
| + if (!(fuse_off & 16) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { | |
| ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); | |
| return 1; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { | |
| + if (!(fuse_off & 32) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { | |
| ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); | |
| return 2; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { | |
| + if (!(fuse_off & 32) && ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { | |
| ggml_cuda_op_ssm_conv(*cuda_ctx, node, /*bias_add_node=*/ nullptr, cgraph->nodes[i + 1]); | |
| return 1; | |
| } | |
| - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || | |
| + if (!(fuse_off & 64) && (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || | |
| ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || | |
| - ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { | |
| + ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS }))) { | |
| ggml_cuda_op_unary_mul(*cuda_ctx, node, cgraph->nodes[i + 1]); | |
| return 1; | |
| } | |
| static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph | |
| return 0; | |
| } | |
| +// ---- diagnostic per-op GPU timer: GGML_CUDA_OP_PROF=1 (with GGML_CUDA_DISABLE_GRAPHS=1); table printed at exit ---- | |
| +struct ggml_cuda_op_prof_state { | |
| + std::vector<cudaEvent_t> events; | |
| + std::vector<int> idx; | |
| + std::map<std::string, std::pair<double, long>> acc; | |
| + std::map<int, long> graphs_by_nodes; | |
| + long graphs = 0; | |
| + double sum_ms = 0.0; | |
| + double wall_ms = 0.0; | |
| + long graphs_c[2] = {0, 0}; | |
| + double sum_c[2] = {0.0, 0.0}; | |
| + double wall_c[2] = {0.0, 0.0}; | |
| +}; | |
| +static ggml_cuda_op_prof_state & ggml_cuda_op_prof() { | |
| + static ggml_cuda_op_prof_state st; | |
| + return st; | |
| +} | |
| +static bool ggml_cuda_op_prof_on() { | |
| + static const bool on = getenv("GGML_CUDA_OP_PROF") != nullptr && std::atoi(getenv("GGML_CUDA_OP_PROF")) != 0; | |
| + return on; | |
| +} | |
| +static void ggml_cuda_op_prof_dump() { | |
| + auto & st = ggml_cuda_op_prof(); | |
| + if (st.graphs == 0) { | |
| + return; | |
| + } | |
| + for (int c = 0; c < 2; ++c) { | |
| + if (st.graphs_c[c] > 0) { | |
| + fprintf(stderr, "OPPROF class %s: graphs=%ld gpu-sum %.3f ms/graph wall %.3f ms/graph\n", c ? "T(>1000 nodes)" : "H(small)", | |
| + st.graphs_c[c], st.sum_c[c] / st.graphs_c[c], st.wall_c[c] / st.graphs_c[c]); | |
| + } | |
| + } | |
| + std::vector<std::pair<std::string, std::pair<double, long>>> rows(st.acc.begin(), st.acc.end()); | |
| + std::sort(rows.begin(), rows.end(), [](const auto & a, const auto & b) { return a.second.first > b.second.first; }); | |
| + const double g = (double) st.graphs; | |
| + fprintf(stderr, "OPPROF graphs=%ld gpu-sum %.3f ms/graph wall %.3f ms/graph (n_nodes:", st.graphs, st.sum_ms / g, st.wall_ms / g); | |
| + for (const auto & [n, c] : st.graphs_by_nodes) { | |
| + fprintf(stderr, " %d x%ld", n, c); | |
| + } | |
| + fprintf(stderr, ")\n"); | |
| + fprintf(stderr, "OPPROF %9s %6s %8s %7s %s\n", "ms/graph", "share", "calls/g", "us/call", "label"); | |
| + for (const auto & r : rows) { | |
| + // per-class labels ("T " trunk-sized graphs, "H " small graphs) average over that class's graphs | |
| + const int cls = r.first.rfind("T ", 0) == 0 ? 1 : (r.first.rfind("H ", 0) == 0 ? 0 : -1); | |
| + const double gc = cls >= 0 && st.graphs_c[cls] > 0 ? (double) st.graphs_c[cls] : g; | |
| + const double ms = r.second.first / gc; | |
| + fprintf(stderr, "OPPROF %9.3f %5.1f%% %8.1f %7.1f %s\n", ms, 100.0 * r.second.first / st.sum_ms, | |
| + (double) r.second.second / gc, 1000.0 * r.second.first / (double) r.second.second, r.first.c_str()); | |
| + } | |
| +} | |
| +static std::string ggml_cuda_op_prof_label(const ggml_tensor * n) { | |
| + std::string l = ggml_op_desc(n); | |
| + if ((n->op == GGML_OP_MUL_MAT || n->op == GGML_OP_MUL_MAT_ID) && n->src[0] && n->src[1]) { | |
| + char buf[128]; | |
| + snprintf(buf, sizeof(buf), "(%s %lldx%lld n=%lld)", ggml_type_name(n->src[0]->type), | |
| + (long long) n->src[0]->ne[1], (long long) n->src[0]->ne[0], (long long) n->src[1]->ne[1]); | |
| + l += buf; | |
| + } else if (n->op == GGML_OP_GET_ROWS && n->src[0]) { | |
| + l += "("; | |
| + l += ggml_type_name(n->src[0]->type); | |
| + l += ")"; | |
| + } | |
| + return l; | |
| +} | |
| +static void ggml_cuda_op_prof_mark(cudaStream_t stream, const int i, const int k) { | |
| + auto & st = ggml_cuda_op_prof(); | |
| + if ((int) st.events.size() <= k) { | |
| + cudaEvent_t e; | |
| + CUDA_CHECK(cudaEventCreate(&e)); | |
| + st.events.push_back(e); | |
| + st.idx.push_back(i); | |
| + } | |
| + st.idx[k] = i; | |
| + CUDA_CHECK(cudaEventRecord(st.events[k], stream)); | |
| +} | |
| +static void ggml_cuda_op_prof_finish(cudaStream_t stream, const ggml_cgraph * cgraph, const int nmarks) { | |
| + if (nmarks == 0) { | |
| + return; | |
| + } | |
| + auto & st = ggml_cuda_op_prof(); | |
| + static bool registered = false; | |
| + if (!registered) { | |
| + registered = true; | |
| + atexit(ggml_cuda_op_prof_dump); | |
| + } | |
| + ggml_cuda_op_prof_mark(stream, cgraph->n_nodes, nmarks); | |
| + CUDA_CHECK(cudaEventSynchronize(st.events[nmarks])); | |
| + for (int k = 0; k < nmarks; ++k) { | |
| + float ms = 0.0f; | |
| + CUDA_CHECK(cudaEventElapsedTime(&ms, st.events[k], st.events[k + 1])); | |
| + std::string label; | |
| + for (int j = st.idx[k]; j < st.idx[k + 1]; ++j) { | |
| + const ggml_tensor * n = cgraph->nodes[j]; | |
| + if (ggml_cuda_is_view_or_noop(n) || (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { | |
| + continue; | |
| + } | |
| + if (!label.empty()) { | |
| + label += " + "; | |
| + } | |
| + label += ggml_cuda_op_prof_label(n); | |
| + } | |
| + const int cls = cgraph->n_nodes > 1000 ? 1 : 0; | |
| + auto & a = st.acc[(cls ? "T " : "H ") + label]; | |
| + a.first += ms; | |
| + a.second += 1; | |
| + st.sum_ms += ms; | |
| + st.sum_c[cls] += ms; | |
| + } | |
| + float wall = 0.0f; | |
| + CUDA_CHECK(cudaEventElapsedTime(&wall, st.events[0], st.events[nmarks])); | |
| + st.wall_ms += wall; | |
| + st.graphs++; | |
| + st.graphs_by_nodes[cgraph->n_nodes]++; | |
| + { | |
| + const int cls = cgraph->n_nodes > 1000 ? 1 : 0; | |
| + st.graphs_c[cls]++; | |
| + st.wall_c[cls] += wall; | |
| + } | |
| + static const int every = getenv("GGML_CUDA_OP_PROF") ? std::atoi(getenv("GGML_CUDA_OP_PROF")) : 0; | |
| + if (every > 1 && st.graphs >= every) { | |
| + ggml_cuda_op_prof_dump(); | |
| + st.acc.clear(); | |
| + st.graphs_by_nodes.clear(); | |
| + st.graphs = 0; st.sum_ms = 0.0; st.wall_ms = 0.0; | |
| + st.graphs_c[0] = st.graphs_c[1] = 0; st.sum_c[0] = st.sum_c[1] = 0.0; st.wall_c[0] = st.wall_c[1] = 0.0; | |
| + } | |
| +} | |
| + | |
| +// Fold the recurrent-state gather into the GDN kernel: a GET_ROWS of F32 rows (src0 contiguous rows, src1 I32) whose | |
| +// only non-view consumer is a GATED_DELTA_NET reading it as src[5] (through views) is skipped, and that GDN reads the | |
| +// rows from src0 directly. Returns the GET_ROWS nodes to skip. GGML_CUDA_GDN_GATHER=1 disables. | |
| +static std::unordered_set<const ggml_tensor *> ggml_cuda_gdn_fold_gather(const ggml_cgraph * cgraph) { | |
| + std::unordered_set<const ggml_tensor *> skip; | |
| + ggml_cuda_gdn_state_rows_clear(); | |
| + static const bool disabled = getenv("GGML_CUDA_GDN_GATHER") != nullptr && std::atoi(getenv("GGML_CUDA_GDN_GATHER")) != 0; | |
| + if (disabled) { | |
| + return skip; | |
| + } | |
| + // walk from each GDN's state input down its view chain to the gather; every link (and the gather itself) must have | |
| + // exactly one reader, so the GDN is the gather's only non-view reader (the graph's use_counts) | |
| + auto uses = [&](const ggml_tensor * t) { | |
| + const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, t); | |
| + if (pos == GGML_HASHSET_FULL || !ggml_bitset_get(cgraph->visited_hash_set.used, pos)) { | |
| + return -1; | |
| + } | |
| + return (int) cgraph->use_counts[pos]; | |
| + }; | |
| + for (int i = 0; i < cgraph->n_nodes; ++i) { | |
| + const ggml_tensor * gdn = cgraph->nodes[i]; | |
| + if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->src[5] == nullptr || gdn->src[6] != nullptr) { | |
| + continue; | |
| + } | |
| + const ggml_tensor * st = gdn->src[5]; | |
| + const ggml_tensor * g = st; | |
| + bool chain_ok = true; | |
| + while (g != nullptr && g->op != GGML_OP_GET_ROWS) { | |
| + if (!ggml_cuda_is_view_or_noop(g) || g->view_src == nullptr || uses(g) != 1) { | |
| + chain_ok = false; | |
| + break; | |
| + } | |
| + g = g->src[0]; | |
| + } | |
| + if (!chain_ok || g == nullptr || skip.count(g)) { | |
| + continue; | |
| + } | |
| + if (!(g->type == GGML_TYPE_F32 && g->src[0]->type == GGML_TYPE_F32 && g->src[1]->type == GGML_TYPE_I32 && | |
| + g->src[0]->nb[0] == sizeof(float) && ggml_is_contiguous(g) && g->ne[2] == 1 && g->ne[3] == 1 && | |
| + g->src[1]->ne[1] == 1 && (g->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && !ggml_is_empty(g)) || uses(g) != 1) { | |
| + continue; | |
| + } | |
| + // the GDN must see the gather as its [S_v, S_v, H, n_seqs] state: one gathered row per seq | |
| + if (!ggml_is_contiguous(st) || ggml_nelements(st) != ggml_nelements(g) || | |
| + g->ne[0] != st->ne[0] * st->ne[1] * st->ne[2] || g->ne[1] != st->ne[3] || | |
| + g->src[1]->ne[0] != g->ne[1]) { | |
| + continue; | |
| + } | |
| + ggml_cuda_gdn_state_rows sr; | |
| + sr.base = (const float *) g->src[0]->data; | |
| + sr.rows = (const int32_t *) g->src[1]->data; | |
| + sr.row_stride = (int64_t) (g->src[0]->nb[1] / sizeof(float)); | |
| + ggml_cuda_gdn_state_rows_set(gdn, sr); | |
| + skip.insert(g); | |
| + } | |
| + return skip; | |
| +} | |
| + | |
| +// diagnostic: GGML_CUDA_PREPASS_PROF=N -- host time of the per-evaluation setup before the first node, split in parts | |
| +struct ggml_cuda_prepass_prof { | |
| + int every = -1; | |
| + long n = 0; | |
| + double acc[4] = {0, 0, 0, 0}; | |
| + bool on() { | |
| + if (every < 0) { | |
| + every = getenv("GGML_CUDA_PREPASS_PROF") ? atoi(getenv("GGML_CUDA_PREPASS_PROF")) : 0; | |
| + } | |
| + return every > 0; | |
| + } | |
| + void add(int k, int64_t us) { acc[k] += us; } | |
| + void done(int n_nodes) { | |
| + if (++n >= every) { | |
| + fprintf(stderr, "PREPASS n=%ld (last graph %d nodes) us/eval: consumer-map %.1f | gdn-fold %.1f | conv-fold %.1f | total-to-first-node %.1f\n", | |
| + n, n_nodes, acc[0] / n, acc[1] / n, acc[2] / n, acc[3] / n); | |
| + fflush(stderr); | |
| + n = 0; | |
| + acc[0] = acc[1] = acc[2] = acc[3] = 0; | |
| + } | |
| + } | |
| +}; | |
| +static ggml_cuda_prepass_prof g_prepass_prof; | |
| + | |
| static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { | |
| + const int64_t pp_t0 = g_prepass_prof.on() ? ggml_time_us() : 0; | |
| bool graph_evaluated_or_captured = false; | |
| // flag used to determine whether it is an integrated_gpu | |
| static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud | |
| std::vector<std::unique_ptr<ggml_cuda_pool_alloc<char>>> gb10_pool_allocations; | |
| std::map<const ggml_tensor *, std::array<int, GGML_TYPE_COUNT>> gb10_shared_q8_consumer_counts; | |
| - for (int j = 0; j < cgraph->n_nodes; ++j) { | |
| - const ggml_tensor * candidate = cgraph->nodes[j]; | |
| - if (candidate->op == GGML_OP_MUL_MAT && candidate->src[0] && candidate->src[1]) { | |
| - ++gb10_shared_q8_consumer_counts[candidate->src[1]][candidate->src[0]->type]; | |
| + if (ggml_cuda_info().devices[cuda_ctx->device].cc == GGML_CUDA_CC_DGX_SPARK) { // only the GB10 paths read it | |
| + for (int j = 0; j < cgraph->n_nodes; ++j) { | |
| + const ggml_tensor * candidate = cgraph->nodes[j]; | |
| + if (candidate->op == GGML_OP_MUL_MAT && candidate->src[0] && candidate->src[1]) { | |
| + ++gb10_shared_q8_consumer_counts[candidate->src[1]][candidate->src[0]->type]; | |
| + } | |
| } | |
| } | |
| + if (pp_t0) { | |
| + g_prepass_prof.add(0, ggml_time_us() - pp_t0); | |
| + } | |
| const auto gb10_shared_q8_consumer_count = [&](const ggml_tensor * src1, ggml_type type) { | |
| const auto it = gb10_shared_q8_consumer_counts.find(src1); | |
| return it == gb10_shared_q8_consumer_counts.end() ? 0 : it->second[type]; | |
| static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud | |
| stream_ctx.concurrent_events.clear(); | |
| } | |
| + const bool op_prof = ggml_cuda_op_prof_on(); | |
| + int op_prof_marks = 0; | |
| + { | |
| + // diagnostic: GGML_CUDA_DUMP_NODES=a:b prints nodes [a, b) of the first graph with >= 1000 nodes, once | |
| + static bool dumped = false; | |
| + static const char * e = getenv("GGML_CUDA_DUMP_NODES"); | |
| + if (e && !dumped && cgraph->n_nodes >= 1000) { | |
| + dumped = true; | |
| + int a = 0, b = 0; | |
| + sscanf(e, "%d:%d", &a, &b); | |
| + for (int j = a; j < b && j < cgraph->n_nodes; ++j) { | |
| + const ggml_tensor * n = cgraph->nodes[j]; | |
| + fprintf(stderr, "NODE %5d %-16s %-28s [%lld,%lld,%lld] %s <-", j, ggml_op_desc(n), n->name, | |
| + (long long) n->ne[0], (long long) n->ne[1], (long long) n->ne[2], | |
| + ggml_cuda_is_view_or_noop(n) ? "(view)" : ""); | |
| + for (int k = 0; k < GGML_MAX_SRC && n->src[k]; ++k) { | |
| + fprintf(stderr, " %s", n->src[k]->name); | |
| + } | |
| + fprintf(stderr, "\n"); | |
| + } | |
| + } | |
| + } | |
| + const int64_t pp_t1 = pp_t0 ? ggml_time_us() : 0; | |
| + std::unordered_set<const ggml_tensor *> gdn_gather_skip = ggml_cuda_gdn_fold_gather(cgraph); | |
| + g_cuda_folded_gathers = &gdn_gather_skip; | |
| + const int64_t pp_t2 = pp_t0 ? ggml_time_us() : 0; | |
| + ggml_cuda_conv_fold_gather(cuda_ctx, cgraph, gdn_gather_skip); | |
| + if (pp_t0) { | |
| + const int64_t pp_t3 = ggml_time_us(); | |
| + g_prepass_prof.add(1, pp_t2 - pp_t1); | |
| + g_prepass_prof.add(2, pp_t3 - pp_t2); | |
| + g_prepass_prof.add(3, pp_t3 - pp_t0); | |
| + g_prepass_prof.done(cgraph->n_nodes); | |
| + } | |
| + ggml_cuda_fwht_q8_clear(); | |
| for (int i = 0; i < cgraph->n_nodes; i++) { | |
| ggml_tensor * node = cgraph->nodes[i]; | |
| if (is_concurrent_event_active) { | |
| static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud | |
| continue; | |
| } | |
| + if (!gdn_gather_skip.empty() && gdn_gather_skip.count(node)) { | |
| + continue; // folded into the GDN kernel's state read | |
| + } | |
| + | |
| + if (op_prof && !use_cuda_graph) { | |
| + ggml_cuda_op_prof_mark(cuda_ctx->stream(), i, op_prof_marks++); | |
| + } | |
| + | |
| // The normalized pre-attention residual is consumed only by a | |
| // group of low-bit projections. Preserve residual + one scale per | |
| // row and let their shared Q8 quantizer apply the norm weight. | |
| static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud | |
| try_launch_concurrent_event(node); | |
| } | |
| } | |
| + if (op_prof && !use_cuda_graph) { | |
| + ggml_cuda_op_prof_finish(cuda_ctx->stream(), cgraph, op_prof_marks); | |
| + } | |
| } | |
| #ifdef USE_CUDA_GRAPH | |
| static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co | |
| ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); | |
| if (graph->graph == nullptr) { | |
| - if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) { | |
| + // sm_6x uses CUDA graphs too (GGML_CUDA_GRAPHS_PASCAL=0 = eager). The 09-18 "no gain / -1%" reading on a GTX 1080 Ti | |
| + // was desktop-compositor noise; on a quiet screen, with per-width graph keys, graphs win 1-2.5% end to end with MTP | |
| + // (patch_fork_graphs_pascal_on.py) | |
| + static const bool graphs_pascal = getenv("GGML_CUDA_GRAPHS_PASCAL") == nullptr || atoi(getenv("GGML_CUDA_GRAPHS_PASCAL")) != 0; | |
| + if (ggml_cuda_info().devices[cuda_ctx->device].cc < (graphs_pascal ? 600 : GGML_CUDA_CC_VOLTA)) { | |
| if (!graph->disable_due_to_gpu_arch) { | |
| GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); | |
| } | |
| static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co | |
| } | |
| #endif // USE_CUDA_GRAPH | |
| +// ---- diagnostic: GGML_CUDA_LAUNCH_PROF=N prints every N graph computes the mean host ms of the validity work (key + | |
| +// compatibility + property check) and of the evaluate step, per graph class (big/small) and path (eager/capture/replay) | |
| +struct ggml_cuda_launch_prof { | |
| + int every = -1; | |
| + long calls = 0; | |
| + long n[2][3] = {}; | |
| + double chk[2][3] = {}, ev[2][3] = {}; | |
| + bool on() { | |
| + if (every < 0) { every = getenv("GGML_CUDA_LAUNCH_PROF") ? atoi(getenv("GGML_CUDA_LAUNCH_PROF")) : 0; } | |
| + return every > 0; | |
| + } | |
| + void add(int cls, int kind, double c, double e) { | |
| + n[cls][kind]++; chk[cls][kind] += c; ev[cls][kind] += e; | |
| + if (++calls >= every) { | |
| + static const char * cn[2] = { "small", "big" }; | |
| + static const char * kn[3] = { "eager", "capture", "replay" }; | |
| + for (int a = 0; a < 2; ++a) { | |
| + for (int b = 0; b < 3; ++b) { | |
| + if (n[a][b]) { | |
| + fprintf(stderr, "LAUNCHPROF %-5s %-7s n=%5ld check %.3f ms eval %.3f ms\n", cn[a], kn[b], n[a][b], | |
| + chk[a][b] / n[a][b] / 1000.0, ev[a][b] / n[a][b] / 1000.0); | |
| + } | |
| + } | |
| + } | |
| + fflush(stderr); | |
| + calls = 0; | |
| + for (int a = 0; a < 2; ++a) { for (int b = 0; b < 3; ++b) { n[a][b] = 0; chk[a][b] = ev[a][b] = 0; } } | |
| + } | |
| + } | |
| +}; | |
| +static ggml_cuda_launch_prof g_launch_prof; | |
| + | |
| static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { | |
| ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; | |
| + const int64_t lp_t0 = g_launch_prof.on() ? ggml_time_us() : 0; | |
| ggml_cuda_set_device(cuda_ctx->device); | |
| static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, | |
| // else: properties changed or first call - execute directly (use_cuda_graph stays false) | |
| } else { | |
| // Post-warmup: normal CUDA graph operation | |
| - if (properties_changed) { | |
| + static const int fastrecap = getenv("GGML_CUDA_GRAPH_FASTRECAP") ? atoi(getenv("GGML_CUDA_GRAPH_FASTRECAP")) : 0; // default off: neutral in A/B (09-19) | |
| + if (properties_changed && fastrecap > 0 && graph->n_replays >= fastrecap) { | |
| + // a key that was stable: re-capture on this call instead of one eager pass first | |
| + use_cuda_graph = true; | |
| + cuda_graph_update_required = true; | |
| + graph->n_replays = 0; | |
| + } else if (properties_changed) { | |
| // Properties changed - reset warmup, execute directly until stable again | |
| graph->warmup_complete = false; | |
| + graph->n_replays = 0; | |
| GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); | |
| } else { | |
| use_cuda_graph = true; | |
| cuda_graph_update_required = graph->instance == nullptr; | |
| + graph->n_replays++; | |
| } | |
| } | |
| } | |
| static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, | |
| CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); | |
| } | |
| + const int64_t lp_t1 = lp_t0 ? ggml_time_us() : 0; | |
| ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); | |
| + if (lp_t0) { | |
| + const int kind = !use_cuda_graph ? 0 : (cuda_graph_update_required ? 1 : 2); | |
| + g_launch_prof.add(cgraph->n_nodes > 500 ? 1 : 0, kind, (double) (lp_t1 - lp_t0), (double) (ggml_time_us() - lp_t1)); | |
| + } | |
| return GGML_STATUS_SUCCESS; | |
| } | |
| diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu | |
| index d7dbc8b..67f623a 100644 | |
| --- a/ggml/src/ggml-cuda/mmvf.cu | |
| +++ b/ggml/src/ggml-cuda/mmvf.cu | |
| bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 | |
| return false; | |
| } | |
| } | |
| + | |
| +// ---- two BF16 GEMVs on the same activations, RPB rows per 256-thread block (see ggml_cuda_mul_mat_vec_bf16_dual) ---- | |
| +template <int NCOLS, int RPB, int ITERS> | |
| +static __global__ void __launch_bounds__(256) mmv_bf16_dual( | |
| + const nv_bfloat16 * xa, const nv_bfloat16 * xb, const int64_t stride_row, const int nrows, | |
| + const float * y, const int64_t stride_col_y2, float * da, float * db, const int64_t stride_col_dst, const int ncols2) { | |
| + constexpr int block_size = 256; | |
| + constexpr int warp_size = WARP_SIZE; | |
| + constexpr int nwarps = block_size / warp_size; | |
| + const int tid = threadIdx.x; | |
| + const int lane = tid % warp_size; | |
| + const int warp = tid / warp_size; | |
| + | |
| + const nv_bfloat162 * xr[RPB]; | |
| +#pragma unroll | |
| + for (int r = 0; r < RPB; ++r) { | |
| + const int row = blockIdx.x * RPB + r; // rows [0, nrows) of xa, then [nrows, 2 nrows) of xb | |
| + xr[r] = (const nv_bfloat162 *) (row < nrows ? xa + (int64_t) row * stride_row : xb + (int64_t) (row - nrows) * stride_row); | |
| + } | |
| + const float2 * y2 = (const float2 *) y; | |
| + | |
| + float sumf[RPB][NCOLS]; | |
| +#pragma unroll | |
| + for (int r = 0; r < RPB; ++r) { | |
| +#pragma unroll | |
| + for (int j = 0; j < NCOLS; ++j) { | |
| + sumf[r][j] = 0.0f; | |
| + } | |
| + } | |
| + | |
| + const int n_iter = ITERS > 0 ? ITERS : (ncols2 - tid + block_size - 1) / block_size; | |
| +#pragma unroll | |
| + for (int it = 0; it < (ITERS > 0 ? ITERS : n_iter); ++it) { | |
| + const int col2 = tid + it * block_size; | |
| + nv_bfloat162 tmpx[RPB]; | |
| +#pragma unroll | |
| + for (int r = 0; r < RPB; ++r) { | |
| + tmpx[r] = xr[r][col2]; | |
| + } | |
| +#pragma unroll | |
| + for (int j = 0; j < NCOLS; ++j) { | |
| + const float2 tmpy = y2[j * stride_col_y2 + col2]; | |
| +#pragma unroll | |
| + for (int r = 0; r < RPB; ++r) { | |
| + ggml_cuda_mad(sumf[r][j], tmpx[r].x, tmpy.x); | |
| + ggml_cuda_mad(sumf[r][j], tmpx[r].y, tmpy.y); | |
| + } | |
| + } | |
| + } | |
| + | |
| + __shared__ float buf[RPB][NCOLS][nwarps]; | |
| +#pragma unroll | |
| + for (int r = 0; r < RPB; ++r) { | |
| +#pragma unroll | |
| + for (int j = 0; j < NCOLS; ++j) { | |
| + const float v = warp_reduce_sum<warp_size>(sumf[r][j]); | |
| + if (lane == 0) { | |
| + buf[r][j][warp] = v; | |
| + } | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + // second stage of mul_mat_vec_f: one warp per (r, j), lanes >= nwarps read zero | |
| + for (int q = warp; q < RPB * NCOLS; q += nwarps) { | |
| + const int r = q / NCOLS; | |
| + const int j = q % NCOLS; | |
| + float v = lane < nwarps ? buf[r][j][lane] : 0.0f; | |
| + v = warp_reduce_sum<warp_size>(v); | |
| + if (lane == 0) { | |
| + const int row = blockIdx.x * RPB + r; | |
| + float * d = row < nrows ? da + row : db + (row - nrows); | |
| + d[j * stride_col_dst] = v; | |
| + } | |
| + } | |
| +} | |
| + | |
| +template <int RPB> | |
| +static void mmv_bf16_dual_launch(const nv_bfloat16 * xa, const nv_bfloat16 * xb, int64_t stride_row, int64_t nrows, int64_t ncols, | |
| + const float * y, int64_t stride_col_y, int64_t ncols_y, float * da, float * db, | |
| + int64_t stride_col_dst, cudaStream_t stream) { | |
| + const dim3 grid((unsigned) (2 * nrows / RPB)); | |
| + const int ncols2 = (int) (ncols / 2); | |
| + const int64_t sy2 = stride_col_y / 2; | |
| + switch (ncols_y) { | |
| +#define BF16_DUAL_CASE(NC) case NC: \ | |
| + if (ncols2 == 2560) { mmv_bf16_dual<NC, RPB, 10><<<grid, 256, 0, stream>>>(xa, xb, stride_row, (int) nrows, y, sy2, da, db, stride_col_dst, ncols2); } \ | |
| + else { mmv_bf16_dual<NC, RPB, 0><<<grid, 256, 0, stream>>>(xa, xb, stride_row, (int) nrows, y, sy2, da, db, stride_col_dst, ncols2); } \ | |
| + break; | |
| + BF16_DUAL_CASE(1) | |
| + BF16_DUAL_CASE(2) | |
| + BF16_DUAL_CASE(3) | |
| + BF16_DUAL_CASE(4) | |
| + BF16_DUAL_CASE(5) | |
| + BF16_DUAL_CASE(6) | |
| + BF16_DUAL_CASE(7) | |
| + BF16_DUAL_CASE(8) | |
| +#undef BF16_DUAL_CASE | |
| + default: GGML_ABORT("mmv_bf16_dual: ncols_y"); | |
| + } | |
| + CUDA_CHECK(cudaGetLastError()); | |
| +} | |
| + | |
| +void ggml_cuda_mul_mat_vec_bf16_dual(const void * xa, const void * xb, int64_t stride_row, int64_t nrows, int64_t ncols, | |
| + const float * y, int64_t stride_col_y, int64_t ncols_y, float * da, float * db, | |
| + int64_t stride_col_dst, int rpb, cudaStream_t stream) { | |
| + GGML_ASSERT(ncols % 2 == 0 && stride_row % 2 == 0 && stride_col_y % 2 == 0 && ncols_y >= 1 && ncols_y <= 8); | |
| + const nv_bfloat16 * a = (const nv_bfloat16 *) xa; | |
| + const nv_bfloat16 * b = (const nv_bfloat16 *) xb; | |
| + switch (rpb) { | |
| + case 1: mmv_bf16_dual_launch<1>(a, b, stride_row, nrows, ncols, y, stride_col_y, ncols_y, da, db, stride_col_dst, stream); break; | |
| + case 2: GGML_ASSERT((2 * nrows) % 2 == 0); mmv_bf16_dual_launch<2>(a, b, stride_row, nrows, ncols, y, stride_col_y, ncols_y, da, db, stride_col_dst, stream); break; | |
| + case 8: GGML_ASSERT((2 * nrows) % 8 == 0); mmv_bf16_dual_launch<8>(a, b, stride_row, nrows, ncols, y, stride_col_y, ncols_y, da, db, stride_col_dst, stream); break; | |
| + default: GGML_ASSERT((2 * nrows) % 4 == 0); mmv_bf16_dual_launch<4>(a, b, stride_row, nrows, ncols, y, stride_col_y, ncols_y, da, db, stride_col_dst, stream); break; | |
| + } | |
| +} | |
| diff --git a/ggml/src/ggml-cuda/mmvf.cuh b/ggml/src/ggml-cuda/mmvf.cuh | |
| index a50f7c0..3467cd5 100644 | |
| --- a/ggml/src/ggml-cuda/mmvf.cuh | |
| +++ b/ggml/src/ggml-cuda/mmvf.cuh | |
| void ggml_cuda_op_mul_mat_vec_f( | |
| const int64_t src1_padded_row_size, cudaStream_t stream); | |
| bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11); | |
| + | |
| +// two BF16 weights xa, xb ([ncols, nrows] each, row stride stride_row elements) times the same F32 y ([ncols, ncols_y], | |
| +// column stride stride_col_y floats) into da, db ([nrows, ncols_y], column stride stride_col_dst): one launch, results | |
| +// bit-identical to two mul_mat_vec_f calls. ncols_y <= 8, ncols even. | |
| +void ggml_cuda_mul_mat_vec_bf16_dual(const void * xa, const void * xb, int64_t stride_row, int64_t nrows, int64_t ncols, | |
| + const float * y, int64_t stride_col_y, int64_t ncols_y, float * da, float * db, | |
| + int64_t stride_col_dst, int rpb, cudaStream_t stream); | |
| diff --git a/ggml/src/ggml-cuda/mmvq-rowlane.cu b/ggml/src/ggml-cuda/mmvq-rowlane.cu | |
| new file mode 100644 | |
| index 0000000..df4c164 | |
| --- /dev/null | |
| +++ b/ggml/src/ggml-cuda/mmvq-rowlane.cu | |
| +#include <map> | |
| +#include <vector> | |
| +// mmvq-rowlane.cu -- smem-tiled GEMV for PQ2_0 on pre-Volta NVIDIA (sm_6x). ("rowlane" is the file's historical name.) | |
| +// | |
| +// WHY (measured 2026-09-18 on a GTX 1080 Ti, Bonsai 2 27B PQ2_0, streambench.cu in the kernel-lane): | |
| +// * read-only streaming ceiling on the card ~335 GB/s (6003 MHz memory); the generic mul_mat_vec_q LOAD pattern alone | |
| +// reaches 321 GB/s, but the full generic inner loop reaches only ~210: the loss is the q8_1 activation side -- | |
| +// PQ2_0 reads 4 bytes of int8 activations per byte of weight, every row re-reads them through L1 with a 36-byte | |
| +// lane stride, and ncols_dst > 1 re-unpacks the weights per column (a 4-wide verify costs 2.15x a single token). | |
| +// * a lane-per-row layout (v1) is catastrophic for DRAM (32-62 GB/s): 32 distinct rows per warp instruction. | |
| +// * this design: each thread owns one whole 34-byte block of ONE row (8 aligned 32-bit loads + two 16-bit tails, | |
| +// every load inside the block, 325-335 GB/s as a pure load pattern); the q8_1 activations of a K tile are staged | |
| +// ONCE per CUDA block in shared memory (16-byte aligned, 36 words per K-block -> conflict-free LDS.128) and shared | |
| +// by R=8 rows; weights are unpacked once per row-block and dotted against every column; the next tile's weights | |
| +// are prefetched before the current tile is computed. Microbench: 300-305 GB/s at K=5120 vs 210 generic. | |
| +// | |
| +// Scope: MUL_MAT only (no ids), no fusion, ne02 == ne03 == 1, ncols_dst <= 8, type PQ2_0. Disabled with | |
| +// GGML_CUDA_ROWLANE=0 (A/B on one binary). Not bit-exact with the generic kernel: per-chunk products are summed | |
| +// in a different order (gate by KL, not by bytes). | |
| + | |
| +#include "common.cuh" | |
| +#include "mmvq-rowlane.cuh" | |
| + | |
| +#include <cstdlib> | |
| + | |
| +#define TILED_MAX_NCOLS 8 | |
| +#define TILED_ROWS_MAX 32 // rows per CUDA block sharing one staged activation tile (R by ncols: 8 / 16 / 32) | |
| +#define TILED_TB_MAX 40 // K-blocks per tile (threads per row); smem = ncols * TB * 144 B | |
| +#define TILED_WORDS_PER_KB 36 // 4 chunks x 8 qs words + 4 ds floats (144 B, 16-B aligned) | |
| + | |
| +namespace { | |
| + | |
| +static __device__ __forceinline__ void pq2_unpack16(const uint32_t q16, int & qx, int & qy) { | |
| + const int qe = __byte_perm(0x020100FF, 0x020100FF, q16 >> 0); | |
| + const int qo = __byte_perm(0x020100FF, 0x020100FF, q16 >> 2); | |
| + qx = __byte_perm(qe, qo, 0x5140); // elements 0..3 | |
| + qy = __byte_perm(qe, qo, 0x7362); // elements 4..7 | |
| +} | |
| + | |
| +struct pq2_regs { | |
| + uint32_t w[8]; | |
| + uint32_t t0; // bytes 30..31 of the block | |
| + uint32_t t1; // bytes 32..33 of the block | |
| +}; | |
| + | |
| +static __device__ __forceinline__ void pq2_load(const uint8_t * __restrict__ blk, pq2_regs & r) { | |
| + const uint32_t * base = (const uint32_t *) ((size_t) blk & ~(size_t) 3); | |
| +#pragma unroll | |
| + for (int i = 0; i < 8; ++i) { | |
| + r.w[i] = __ldg(base + i); | |
| + } | |
| + r.t0 = __ldg((const unsigned short *) (blk + 30)); | |
| + r.t1 = __ldg((const unsigned short *) (blk + 32)); | |
| +} | |
| + | |
| +static __device__ __forceinline__ void pq2_assemble(const pq2_regs & r, const uint32_t misal, float & d, uint32_t (&q)[8]) { | |
| + if (misal == 0) { | |
| + d = __half2float(__ushort_as_half((unsigned short) (r.w[0] & 0xFFFFu))); | |
| +#pragma unroll | |
| + for (int j = 0; j < 7; ++j) { | |
| + q[j] = __funnelshift_r(r.w[j], r.w[j + 1], 16); | |
| + } | |
| + q[7] = (r.w[7] >> 16) | (r.t1 << 16); | |
| + } else { | |
| + d = __half2float(__ushort_as_half((unsigned short) (r.w[0] >> 16))); | |
| +#pragma unroll | |
| + for (int j = 0; j < 7; ++j) { | |
| + q[j] = r.w[j + 1]; | |
| + } | |
| + q[7] = r.t0 | (r.t1 << 16); | |
| + } | |
| +} | |
| + | |
| +// grid.x = ceil(nrows/R); block = (TB, R/RT); dynamic smem = ncols * TB * 144 B. Thread (t, rr) owns K-block t of the | |
| +// RT consecutive rows rr*RT..rr*RT+RT-1 of the block's R rows, so one staged activation read serves RT rows (the wide | |
| +// verify widths are shared-memory-bandwidth bound with RT=1) and 2x the weight bytes are in flight per thread. | |
| +// Launch bounds pin registers: 320x3 -> 68 regs (ncols<=2), 192x2 -> 170, 256x2 -> 128. | |
| +template <int ncols, int R, int RT> | |
| +__launch_bounds__((ncols <= 2 ? 320 : 256), (ncols <= 2 ? 3 : 2)) | |
| +static __global__ void mul_mat_vec_pq2_tiled( | |
| + const uint8_t * __restrict__ vx, const block_q8_1 * __restrict__ vy, float * __restrict__ dst, | |
| + const int nblocks, const int nrows, const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst) { | |
| + extern __shared__ uint32_t ytile[]; // [ncols][TB][36]; reused as the reduction buffer after the last tile | |
| + static_assert(R % RT == 0, "rows per block must be a multiple of rows per thread"); | |
| + constexpr int RB = R / RT; // thread rows per block | |
| + | |
| + const int TB = blockDim.x; | |
| + const int t = threadIdx.x; | |
| + const int rr = threadIdx.y; | |
| + const int tid = rr * TB + t; | |
| + const int nthr = TB * RB; | |
| + | |
| + const uint8_t * xrow[RT]; | |
| + bool row_ok[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = blockIdx.x * R + rr * RT + i; | |
| + row_ok[i] = row < nrows; | |
| + xrow[i] = vx + (size_t) (row_ok[i] ? row : nrows - 1) * stride_row_x_bytes; | |
| + } | |
| + | |
| + float acc[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] = 0.0f; | |
| + } | |
| + } | |
| + | |
| + pq2_regs cur[RT]; | |
| + if (t < nblocks) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_load(xrow[i] + (size_t) t * 34, cur[i]); | |
| + } | |
| + } | |
| + | |
| + for (int tile0 = 0; tile0 < nblocks; tile0 += TB) { | |
| + // stage the activation tile: ncols x TB K-blocks x (32 qs words + 4 ds floats) | |
| + for (int k = tid; k < ncols * TB * 4; k += nthr) { | |
| + const int c = k / (TB * 4); | |
| + const int rem = k - c * (TB * 4); | |
| + const int kb = rem >> 2; | |
| + const int j = rem & 3; | |
| + uint32_t * d = ytile + (c * TB + kb) * TILED_WORDS_PER_KB; | |
| + const int kbg = tile0 + kb; | |
| + if (kbg < nblocks) { | |
| + const block_q8_1 * src = vy + c * stride_col_y + kbg * 4 + j; | |
| + const int * qs = (const int *) src->qs; | |
| +#pragma unroll | |
| + for (int w = 0; w < 8; ++w) { | |
| + d[j * 8 + w] = (uint32_t) __ldg(qs + w); | |
| + } | |
| + const uint32_t ds = __ldg((const unsigned int *) &src->ds); | |
| + d[32 + j] = __float_as_uint(__half2float(__ushort_as_half((unsigned short) (ds & 0xFFFFu)))); | |
| + } else { | |
| +#pragma unroll | |
| + for (int w = 0; w < 8; ++w) { | |
| + d[j * 8 + w] = 0u; | |
| + } | |
| + d[32 + j] = 0u; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + | |
| + // prefetch this thread's blocks of the NEXT tile before computing the current one | |
| + const int kb_cur = tile0 + t; | |
| + const int kb_nxt = kb_cur + TB; | |
| + pq2_regs nxt[RT]; | |
| + if (kb_nxt < nblocks) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_load(xrow[i] + (size_t) kb_nxt * 34, nxt[i]); | |
| + } | |
| + } | |
| + | |
| + if (kb_cur < nblocks) { | |
| + const uint32_t misal = (uint32_t) ((size_t) (xrow[0] + (size_t) kb_cur * 34) & 3); | |
| + float d[RT]; | |
| + uint32_t q[RT][8]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_assemble(cur[i], misal, d[i], q[i]); | |
| + } | |
| + | |
| + float accd[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + accd[i][c] = 0.0f; | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int j = 0; j < 4; ++j) { | |
| + int qx0[RT], qy0[RT], qx1[RT], qy1[RT], qx2[RT], qy2[RT], qx3[RT], qy3[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_unpack16(q[i][2*j] & 0xFFFFu, qx0[i], qy0[i]); | |
| + pq2_unpack16(q[i][2*j] >> 16, qx1[i], qy1[i]); | |
| + pq2_unpack16(q[i][2*j+1] & 0xFFFFu, qx2[i], qy2[i]); | |
| + pq2_unpack16(q[i][2*j+1] >> 16, qx3[i], qy3[i]); | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + const uint32_t * yt = ytile + (c * TB + t) * TILED_WORDS_PER_KB; | |
| + const uint4 u0 = *(const uint4 *) (yt + j * 8); | |
| + const uint4 u1 = *(const uint4 *) (yt + j * 8 + 4); | |
| + const float d8 = __uint_as_float(yt[32 + j]); | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + int s = 0; | |
| + s = ggml_cuda_dp4a((int) u0.x, qx0[i], s); | |
| + s = ggml_cuda_dp4a((int) u0.y, qy0[i], s); | |
| + s = ggml_cuda_dp4a((int) u0.z, qx1[i], s); | |
| + s = ggml_cuda_dp4a((int) u0.w, qy1[i], s); | |
| + s = ggml_cuda_dp4a((int) u1.x, qx2[i], s); | |
| + s = ggml_cuda_dp4a((int) u1.y, qy2[i], s); | |
| + s = ggml_cuda_dp4a((int) u1.z, qx3[i], s); | |
| + s = ggml_cuda_dp4a((int) u1.w, qy3[i], s); | |
| + accd[i][c] += d8 * (float) s; | |
| + } | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] += d[i] * accd[i][c]; | |
| + } | |
| + } | |
| + } | |
| + __syncthreads(); // the next staging overwrites the tile | |
| + if (kb_nxt < nblocks) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + cur[i] = nxt[i]; | |
| + } | |
| + } | |
| + } | |
| + | |
| + // reduce over the TB threads of each row (the tile is dead after the loop's final barrier; R*4 <= 144 B per K-block) | |
| + float * red = (float *) ytile; // [R][TB][ncols] | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + red[((rr * RT + i) * TB + t) * ncols + c] = acc[i][c]; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + if constexpr (ncols >= 4) { | |
| + // two-level reduction: every thread folds the upper half onto the lower half, log2(TB) steps (+11% at n=4; | |
| + // at n=1..3 the serial sum by one thread per row is faster because the extra barriers cost more than it saves) | |
| + for (int half = (TB + 1) / 2, n = TB; n > 1; n = half, half = (half + 1) / 2) { | |
| + if (t < half && t + half < n) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + red[((rr * RT + i) * TB + t) * ncols + c] += red[((rr * RT + i) * TB + t + half) * ncols + c]; | |
| + } | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + } | |
| + if (t == 0) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + if (!row_ok[i]) { | |
| + continue; | |
| + } | |
| + const int row = blockIdx.x * R + rr * RT + i; | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + dst[c * stride_col_dst + row] = red[((rr * RT + i) * TB) * ncols + c]; | |
| + } | |
| + } | |
| + } | |
| + return; | |
| + } | |
| + if (t == 0) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + if (!row_ok[i]) { | |
| + continue; | |
| + } | |
| + const int row = blockIdx.x * R + rr * RT + i; | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + float sum = 0.0f; | |
| + for (int k = 0; k < TB; ++k) { | |
| + sum += red[((rr * RT + i) * TB + k) * ncols + c]; | |
| + } | |
| + dst[c * stride_col_dst + row] = sum; | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| + | |
| +// Persistent variant (ncols <= 3) whenever the whole activation vector fits shared memory: the y tile is staged ONCE | |
| +// per resident block; the block loops over row groups (rg, rg+grid, ...) and each thread walks K-blocks t, t+TB, ... | |
| +// of its RT rows, with the next (row group, K-block) item's weights prefetched into registers before the current | |
| +// one is computed. Removes the per-block staging + launch/drain cost that capped the short K=5120 rows at ~255 GB/s. | |
| +// Measured on the 1080 Ti (vbench.cu): RT=1 at ncols=1 330-370 GB/s (tiled 235-300); RT=2 with 20 threads per row | |
| +// at ncols=2 305-333 (tiled ~190) and ncols=3 270-296 (tiled ~210); at ncols>=4 the tiled kernel stays ahead. | |
| +// The reduction is double-buffered so one barrier per row group suffices. Per-row summation order equals the tiled | |
| +// kernel's when one K-block per thread and differs only by fp reordering beyond that (KL-gated). | |
| +template <int ncols, int RT, int MAXT, int MINB> | |
| +__launch_bounds__(MAXT, MINB) | |
| +static __global__ void mul_mat_vec_pq2_persist( | |
| + const uint8_t * __restrict__ vx, const block_q8_1 * __restrict__ vy, float * __restrict__ dst, | |
| + const int nblocks, const int nrows, const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst) { | |
| + constexpr int RB = 8; // thread rows per block | |
| + constexpr int R = RT * RB; // rows per block | |
| + extern __shared__ uint32_t ytile[]; // [ncols][nblocks][36] then red[2][R][TB][ncols] | |
| + const int TB = blockDim.x; | |
| + const int t = threadIdx.x; | |
| + const int rr = threadIdx.y; | |
| + const int tid = rr * TB + t; | |
| + const uint32_t vx_mis = (uint32_t) ((uintptr_t) vx & 3u); | |
| + const uint32_t half_stride = (uint32_t) stride_row_x_bytes / 2u; // row stride is a whole number of 34-byte blocks | |
| + const int nthr = TB * RB; | |
| + float * red = (float *) (ytile + ncols * nblocks * TILED_WORDS_PER_KB); | |
| + const int nrg = (nrows + R - 1) / R; | |
| + | |
| + for (int k = tid; k < ncols * nblocks * 4; k += nthr) { | |
| + const int c = k / (nblocks * 4); | |
| + const int rem = k - c * (nblocks * 4); | |
| + const int kb = rem >> 2; | |
| + const int j = rem & 3; | |
| + uint32_t * d = ytile + (c * nblocks + kb) * TILED_WORDS_PER_KB; | |
| + const block_q8_1 * src = vy + c * stride_col_y + kb * 4 + j; | |
| + const int * qs = (const int *) src->qs; | |
| +#pragma unroll | |
| + for (int w = 0; w < 8; ++w) { | |
| + d[j * 8 + w] = (uint32_t) __ldg(qs + w); | |
| + } | |
| + const uint32_t ds = __ldg((const unsigned int *) &src->ds); | |
| + d[32 + j] = __float_as_uint(__half2float(__ushort_as_half((unsigned short) (ds & 0xFFFFu)))); | |
| + } | |
| + | |
| + int rg = blockIdx.x; | |
| + int kb = t; | |
| + pq2_regs cur[RT]; | |
| + auto load_item = [&](const int g, const int b, pq2_regs (&dstr)[RT]) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = g * R + rr * RT + i; | |
| + const uint8_t * xrow = vx + (size_t) (row < nrows ? row : nrows - 1) * stride_row_x_bytes; | |
| + pq2_load(xrow + (size_t) b * 34, dstr[i]); | |
| + } | |
| + }; | |
| + if (rg < nrg && kb < nblocks) { | |
| + load_item(rg, kb, cur); | |
| + } | |
| + __syncthreads(); | |
| + | |
| + int it = 0; | |
| + for (; rg < nrg; rg += gridDim.x, ++it) { | |
| + float acc[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] = 0.0f; | |
| + } | |
| + } | |
| + for (kb = t; kb < nblocks; kb += TB) { | |
| + int nkb = kb + TB; | |
| + int nrg_ = rg; | |
| + if (nkb >= nblocks) { | |
| + nkb = t; | |
| + nrg_ = rg + gridDim.x; | |
| + } | |
| + const bool has_next = nrg_ < nrg; | |
| + pq2_regs nxt[RT]; | |
| + if (has_next) { | |
| + load_item(nrg_, nkb, nxt); | |
| + } | |
| + // alignment class per ROW: rows are only 2-byte aligned when nblocks is odd (K=128 test shapes) | |
| + float d[RT]; | |
| + uint32_t q[RT][8]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row_i = rg * R + rr * RT + i; | |
| + // block address = vx + row*stride + kb*34 with 34-byte blocks: only bit 1 varies, so the misalignment is | |
| + // (vx & 3) + 2*((row*stride/2 + kb*17) & 1), computed in 32 bits (the 64-bit form cost registers at ncols 3) | |
| + const uint32_t rowc = (uint32_t) (row_i < nrows ? row_i : nrows - 1); | |
| + const uint32_t misal = (vx_mis + (((rowc * half_stride) + (uint32_t) kb) & 1u) * 2u) & 3u; | |
| + pq2_assemble(cur[i], misal, d[i], q[i]); | |
| + } | |
| + float accd[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + accd[i][c] = 0.0f; | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int j = 0; j < 4; ++j) { | |
| + int qx0[RT], qy0[RT], qx1[RT], qy1[RT], qx2[RT], qy2[RT], qx3[RT], qy3[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_unpack16(q[i][2*j] & 0xFFFFu, qx0[i], qy0[i]); | |
| + pq2_unpack16(q[i][2*j] >> 16, qx1[i], qy1[i]); | |
| + pq2_unpack16(q[i][2*j+1] & 0xFFFFu, qx2[i], qy2[i]); | |
| + pq2_unpack16(q[i][2*j+1] >> 16, qx3[i], qy3[i]); | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + const uint32_t * yt = ytile + (c * nblocks + kb) * TILED_WORDS_PER_KB; | |
| + const uint4 u0 = *(const uint4 *) (yt + j * 8); | |
| + const uint4 u1 = *(const uint4 *) (yt + j * 8 + 4); | |
| + const float d8 = __uint_as_float(yt[32 + j]); | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + int sacc = 0; | |
| + sacc = ggml_cuda_dp4a((int) u0.x, qx0[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.y, qy0[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.z, qx1[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.w, qy1[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.x, qx2[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.y, qy2[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.z, qx3[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.w, qy3[i], sacc); | |
| + accd[i][c] += d8 * (float) sacc; | |
| + } | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] += d[i] * accd[i][c]; | |
| + } | |
| + } | |
| + if (has_next) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + cur[i] = nxt[i]; | |
| + } | |
| + } | |
| + } | |
| + float * rb = red + (size_t) (it & 1) * R * TB * ncols; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + rb[((rr * RT + i) * TB + t) * ncols + c] = acc[i][c]; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + if (t == 0) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = rg * R + rr * RT + i; | |
| + if (row >= nrows) { | |
| + continue; | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + float sum = 0.0f; | |
| + for (int k = 0; k < TB; ++k) { | |
| + sum += rb[((rr * RT + i) * TB + k) * ncols + c]; | |
| + } | |
| + dst[c * stride_col_dst + row] = sum; | |
| + } | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| + | |
| +// v15: shuffle-reduced persistent kernel with K-chunked activation staging | |
| +// (tools/kernel-lane/bonsai/patch_fork_shfl.py + patch_fork_shfl2.py). | |
| +// Block = RB thread-rows x 8 lanes; thread-row rr owns RT rows of the block's R = RB*RT rows; the 8 lanes of a team walk | |
| +// the K-blocks of the current chunk lane, lane+8, ... (nblocks % 8 == 0 and kc_len % 8 == 0); partial sums are reduced | |
| +// with xor-shuffles inside the team, so there is no reduction buffer. The activations of one K-chunk (all columns) are | |
| +// staged once per block: one chunk (the whole row) unless ncols * nblocks * 144 B exceeds the budget. Chunks after | |
| +// the first accumulate into dst, read-add-write by the thread that wrote the value (same block, team and rows). | |
| +// Measured on a GTX 1080 Ti (vbench_s.cu, model shapes) against v12: n=1 +0-17%, n=2 +12-26%, n=3 +16-30%, | |
| +// n=4 (vs the tiled kernel) +15-45%; in the model a 2/3/4-token step 42.1/48.4/51.3 -> 37.3/41.8/44.5 ms. | |
| +template <int ncols, int RT, int RB, int MINB, bool CHUNKED, bool SPLIT> | |
| +__launch_bounds__(RB * 8, MINB) | |
| +static __global__ void mul_mat_vec_pq2_shfl( | |
| + const uint8_t * __restrict__ vx, const block_q8_1 * __restrict__ vy, float * __restrict__ dst, | |
| + const int nblocks, const int nrows, const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, | |
| + const int kc_len, unsigned long long * stamp, float * __restrict__ part, int * cnt) { | |
| + constexpr int R = RT * RB; | |
| + __shared__ int s_last; | |
| + extern __shared__ uint32_t ytile[]; // [ncols][kc_len][36] | |
| + if (stamp != nullptr && threadIdx.x == 0) { | |
| + unsigned long long t0; | |
| + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); | |
| + atomicMin(stamp, t0); | |
| + } | |
| + const int tid = threadIdx.x; | |
| + const int nthr = RB * 8; | |
| + const int lane8 = tid & 7; | |
| + const int rr = tid >> 3; | |
| + const int nrg = (nrows + R - 1) / R; | |
| + | |
| + const uint32_t vx_mis = (uint32_t) ((uintptr_t) vx & 3u); | |
| + const uint32_t half_stride = (uint32_t) stride_row_x_bytes / 2u; | |
| + | |
| + const int nch = CHUNKED ? (nblocks + kc_len - 1) / kc_len : 1; | |
| + // SPLIT: this block owns one chunk (stages its tile once) and every (gridDim.x / nch)-th row group | |
| + const int ch_first = SPLIT ? (int) (blockIdx.x % nch) : 0; | |
| + const int ch_last = SPLIT ? ch_first + 1 : nch; | |
| + const int rg_first = SPLIT ? (int) (blockIdx.x / nch) : (int) blockIdx.x; | |
| + const int rg_step = SPLIT ? (int) (gridDim.x / nch) : (int) gridDim.x; | |
| + for (int ch = ch_first; ch < ch_last; ++ch) { | |
| + const int kc0 = CHUNKED ? ch * kc_len : 0; | |
| + const int kcn = CHUNKED ? min(kc_len, nblocks - kc0) : nblocks; | |
| + if (!SPLIT && kc0 > 0) { | |
| + __syncthreads(); // every team is done reading the previous chunk's tile | |
| + } | |
| + for (int k = tid; k < ncols * kcn * 4; k += nthr) { | |
| + const int c = k / (kcn * 4); | |
| + const int rem = k - c * (kcn * 4); | |
| + const int kb = rem >> 2; | |
| + const int j = rem & 3; | |
| + uint32_t * d = ytile + (c * kcn + kb) * TILED_WORDS_PER_KB; | |
| + const block_q8_1 * src = vy + c * stride_col_y + (kc0 + kb) * 4 + j; | |
| + const int * qs = (const int *) src->qs; | |
| +#pragma unroll | |
| + for (int w = 0; w < 8; ++w) { | |
| + d[j * 8 + w] = (uint32_t) __ldg(qs + w); | |
| + } | |
| + const uint32_t ds = __ldg((const unsigned int *) &src->ds); | |
| + d[32 + j] = __float_as_uint(__half2float(__ushort_as_half((unsigned short) (ds & 0xFFFFu)))); | |
| + } | |
| + __syncthreads(); | |
| + | |
| + for (int rg = rg_first; rg < nrg; rg += rg_step) { | |
| + const uint8_t * xrow[RT]; | |
| + uint32_t rowc[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = rg * R + rr * RT + i; | |
| + rowc[i] = (uint32_t) (row < nrows ? row : nrows - 1); | |
| + xrow[i] = vx + (size_t) rowc[i] * stride_row_x_bytes; | |
| + } | |
| + float acc[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] = 0.0f; | |
| + } | |
| + } | |
| + pq2_regs cur[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_load(xrow[i] + (size_t) (kc0 + lane8) * 34, cur[i]); | |
| + } | |
| + for (int kb = lane8; kb < kcn; kb += 8) { | |
| + const bool hn = kb + 8 < kcn; | |
| + pq2_regs nxt[RT]; | |
| + if (hn) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_load(xrow[i] + (size_t) (kc0 + kb + 8) * 34, nxt[i]); | |
| + } | |
| + } | |
| + float d[RT]; | |
| + uint32_t q[RT][8]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const uint32_t misal = (vx_mis + (((rowc[i] * half_stride) + (uint32_t) (kc0 + kb)) & 1u) * 2u) & 3u; | |
| + pq2_assemble(cur[i], misal, d[i], q[i]); | |
| + } | |
| + float accd[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + accd[i][c] = 0.0f; | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int j = 0; j < 4; ++j) { | |
| + int qx0[RT], qy0[RT], qx1[RT], qy1[RT], qx2[RT], qy2[RT], qx3[RT], qy3[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + pq2_unpack16(q[i][2*j] & 0xFFFFu, qx0[i], qy0[i]); | |
| + pq2_unpack16(q[i][2*j] >> 16, qx1[i], qy1[i]); | |
| + pq2_unpack16(q[i][2*j+1] & 0xFFFFu, qx2[i], qy2[i]); | |
| + pq2_unpack16(q[i][2*j+1] >> 16, qx3[i], qy3[i]); | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + const uint32_t * yt = ytile + (c * kcn + kb) * TILED_WORDS_PER_KB; | |
| + const uint4 u0 = *(const uint4 *) (yt + j * 8); | |
| + const uint4 u1 = *(const uint4 *) (yt + j * 8 + 4); | |
| + const float d8 = __uint_as_float(yt[32 + j]); | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + int sacc = 0; | |
| + sacc = ggml_cuda_dp4a((int) u0.x, qx0[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.y, qy0[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.z, qx1[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u0.w, qy1[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.x, qx2[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.y, qy2[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.z, qx3[i], sacc); | |
| + sacc = ggml_cuda_dp4a((int) u1.w, qy3[i], sacc); | |
| + accd[i][c] += d8 * (float) sacc; | |
| + } | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] += d[i] * accd[i][c]; | |
| + } | |
| + } | |
| + if (hn) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + cur[i] = nxt[i]; | |
| + } | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + float v = acc[i][c]; | |
| + v += __shfl_xor_sync(0xffffffffu, v, 4); | |
| + v += __shfl_xor_sync(0xffffffffu, v, 2); | |
| + v += __shfl_xor_sync(0xffffffffu, v, 1); | |
| + acc[i][c] = v; | |
| + } | |
| + } | |
| + if (lane8 == 0) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = rg * R + rr * RT + i; | |
| + if (row < nrows) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + if (SPLIT) { | |
| + part[((size_t) ch * ncols + c) * nrows + row] = acc[i][c]; | |
| + } else { | |
| + float * o = dst + c * stride_col_dst + row; | |
| + if (CHUNKED && kc0 > 0) { | |
| + *o += acc[i][c]; | |
| + } else { | |
| + *o = acc[i][c]; | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + if (SPLIT && cnt != nullptr) { | |
| + __threadfence(); | |
| + __syncthreads(); // this block's partials of rg are written and fenced | |
| + if (tid == 0) { | |
| + const int prev = atomicAdd(cnt + rg, 1); | |
| + s_last = prev == nch - 1; | |
| + if (s_last) { | |
| + cnt[rg] = 0; // every chunk has arrived: reset for the next launch | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + if (s_last) { | |
| + __threadfence(); | |
| + for (int k = tid; k < R * ncols; k += nthr) { | |
| + const int c = k / R; | |
| + const int row = rg * R + (k - c * R); | |
| + if (row < nrows) { | |
| + float v = __ldcg(part + (size_t) c * nrows + row); | |
| + for (int ch2 = 1; ch2 < nch; ++ch2) { | |
| + v += __ldcg(part + ((size_t) ch2 * ncols + c) * nrows + row); | |
| + } | |
| + dst[c * stride_col_dst + row] = v; | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + if (stamp != nullptr) { | |
| + __syncthreads(); | |
| + if (threadIdx.x == 0) { | |
| + unsigned long long t1; | |
| + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t1)); | |
| + atomicMax(stamp + 1, t1); | |
| + } | |
| + } | |
| +} | |
| + | |
| +// per-call globaltimer slot for GGML_CUDA_RL_PROF (null = off) | |
| +static unsigned long long * g_rl_stamp = nullptr; | |
| + | |
| +// ---- Q4_0 row-lane shuffle GEMV (sm_6x), see patch_fork_q4_rowlane.py ---- | |
| +#define Q4R_WPU 44 // tile words per 128-weight unit and column: 32 qs + 4 d8 + 4 (-8*s8) + 4 pad (bank spread) | |
| + | |
| +struct q4_regs { | |
| + uint32_t w[18]; | |
| +}; | |
| + | |
| +static __device__ __forceinline__ void q4_load(const uint8_t * __restrict__ unit, q4_regs & r) { | |
| + const uint32_t * p = (const uint32_t *) unit; | |
| +#pragma unroll | |
| + for (int i = 0; i < 18; ++i) { | |
| + r.w[i] = __ldg(p + i); | |
| + } | |
| +} | |
| + | |
| +// qs word j (bytes 4j..4j+3 of the block's 16 qs bytes) of block b of a 72-byte unit held in r | |
| +static __device__ __forceinline__ uint32_t q4_qs(const q4_regs & r, const int b, const int j) { | |
| + switch (b) { | |
| + case 0: return __funnelshift_r(r.w[j], r.w[j + 1], 16); // qs at bytes 2..17 | |
| + case 1: return r.w[5 + j]; // qs at bytes 20..35 | |
| + case 2: return __funnelshift_r(r.w[9 + j], r.w[10 + j], 16); // qs at bytes 38..53 | |
| + default: return r.w[14 + j]; // qs at bytes 56..71 | |
| + } | |
| +} | |
| + | |
| +static __device__ __forceinline__ void q4_scales(const q4_regs & r, float (&d)[4]) { | |
| + d[0] = __half2float(__ushort_as_half((unsigned short) (r.w[0] & 0xFFFFu))); | |
| + d[1] = __half2float(__ushort_as_half((unsigned short) (r.w[4] >> 16))); | |
| + d[2] = __half2float(__ushort_as_half((unsigned short) (r.w[9] & 0xFFFFu))); | |
| + d[3] = __half2float(__ushort_as_half((unsigned short) (r.w[13] >> 16))); | |
| +} | |
| + | |
| +template <int ncols, int RT, int RB, int MINB, bool CHUNKED, bool SPLIT> | |
| +__launch_bounds__(RB * 8, MINB) | |
| +static __global__ void mul_mat_vec_q4_0_shfl( | |
| + const uint8_t * __restrict__ vx, const block_q8_1 * __restrict__ vy, float * __restrict__ dst, | |
| + const int nunits, const int nrows, const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, | |
| + const int kc_len, float * __restrict__ part, int * cnt) { | |
| + constexpr int R = RT * RB; | |
| + extern __shared__ uint32_t ytile[]; // [ncols][kc_len][Q4R_WPU] | |
| + __shared__ int s_last; | |
| + const int tid = threadIdx.x; | |
| + const int nthr = RB * 8; | |
| + const int lane8 = tid & 7; | |
| + const int rr = tid >> 3; | |
| + const int nrg = (nrows + R - 1) / R; | |
| + | |
| + const int nch = CHUNKED ? (nunits + kc_len - 1) / kc_len : 1; | |
| + const int ch_first = SPLIT ? (int) (blockIdx.x % nch) : 0; | |
| + const int ch_last = SPLIT ? ch_first + 1 : nch; | |
| + const int rg_first = SPLIT ? (int) (blockIdx.x / nch) : (int) blockIdx.x; | |
| + const int rg_step = SPLIT ? (int) (gridDim.x / nch) : (int) gridDim.x; | |
| + for (int ch = ch_first; ch < ch_last; ++ch) { | |
| + const int kc0 = CHUNKED ? ch * kc_len : 0; | |
| + const int kcn = CHUNKED ? min(kc_len, nunits - kc0) : nunits; | |
| + if (!SPLIT && kc0 > 0) { | |
| + __syncthreads(); // every team is done reading the previous chunk's tile | |
| + } | |
| + // stage: one q8_1 block (unit kb, block b, column c) per iteration | |
| + for (int k = tid; k < ncols * kcn * 4; k += nthr) { | |
| + const int c = k / (kcn * 4); | |
| + const int rem = k - c * (kcn * 4); | |
| + const int kb = rem >> 2; | |
| + const int b = rem & 3; | |
| + uint32_t * d = ytile + (c * kcn + kb) * Q4R_WPU; | |
| + const block_q8_1 * src = vy + c * stride_col_y + (kc0 + kb) * 4 + b; | |
| + const int * qs = (const int *) src->qs; | |
| +#pragma unroll | |
| + for (int w = 0; w < 8; ++w) { | |
| + d[b * 8 + w] = (uint32_t) __ldg(qs + w); | |
| + } | |
| + const uint32_t ds = __ldg((const unsigned int *) &src->ds); | |
| + const float d8 = __half2float(__ushort_as_half((unsigned short) (ds & 0xFFFFu))); | |
| + const float s8 = __half2float(__ushort_as_half((unsigned short) (ds >> 16))); | |
| + d[32 + b] = __float_as_uint(d8); | |
| + d[36 + b] = __float_as_uint(-8.0f * s8); | |
| + } | |
| + __syncthreads(); | |
| + | |
| + for (int rg = rg_first; rg < nrg; rg += rg_step) { | |
| + const uint8_t * xrow[RT]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = rg * R + rr * RT + i; | |
| + xrow[i] = vx + (size_t) (row < nrows ? row : nrows - 1) * stride_row_x_bytes; | |
| + } | |
| + float acc[RT][ncols]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + acc[i][c] = 0.0f; | |
| + } | |
| + } | |
| + q4_regs cur[RT]; | |
| + if (lane8 < kcn) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + q4_load(xrow[i] + (size_t) (kc0 + lane8) * 72, cur[i]); | |
| + } | |
| + } | |
| + for (int kb = lane8; kb < kcn; kb += 8) { | |
| + const bool hn = kb + 8 < kcn; | |
| + q4_regs nxt[RT]; | |
| + if (hn) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + q4_load(xrow[i] + (size_t) (kc0 + kb + 8) * 72, nxt[i]); | |
| + } | |
| + } | |
| + float d4[RT][4]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + q4_scales(cur[i], d4[i]); | |
| + } | |
| +#pragma unroll | |
| + for (int b = 0; b < 4; ++b) { | |
| + int lo[RT][4], hi[RT][4]; | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int j = 0; j < 4; ++j) { | |
| + const uint32_t q = q4_qs(cur[i], b, j); | |
| + lo[i][j] = (int) (q & 0x0F0F0F0Fu); | |
| + hi[i][j] = (int) ((q >> 4) & 0x0F0F0F0Fu); | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + const uint32_t * yt = ytile + (c * kcn + kb) * Q4R_WPU; | |
| + const uint4 ul = *(const uint4 *) (yt + b * 8); | |
| + const uint4 uh = *(const uint4 *) (yt + b * 8 + 4); | |
| + const float d8 = __uint_as_float(yt[32 + b]); | |
| + const float m8 = __uint_as_float(yt[36 + b]); | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + int s = 0; | |
| + s = ggml_cuda_dp4a(lo[i][0], (int) ul.x, s); | |
| + s = ggml_cuda_dp4a(lo[i][1], (int) ul.y, s); | |
| + s = ggml_cuda_dp4a(lo[i][2], (int) ul.z, s); | |
| + s = ggml_cuda_dp4a(lo[i][3], (int) ul.w, s); | |
| + s = ggml_cuda_dp4a(hi[i][0], (int) uh.x, s); | |
| + s = ggml_cuda_dp4a(hi[i][1], (int) uh.y, s); | |
| + s = ggml_cuda_dp4a(hi[i][2], (int) uh.z, s); | |
| + s = ggml_cuda_dp4a(hi[i][3], (int) uh.w, s); | |
| + acc[i][c] += d4[i][b] * fmaf((float) s, d8, m8); | |
| + } | |
| + } | |
| + } | |
| + if (hn) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + cur[i] = nxt[i]; | |
| + } | |
| + } | |
| + } | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + float v = acc[i][c]; | |
| + v += __shfl_xor_sync(0xffffffffu, v, 4); | |
| + v += __shfl_xor_sync(0xffffffffu, v, 2); | |
| + v += __shfl_xor_sync(0xffffffffu, v, 1); | |
| + acc[i][c] = v; | |
| + } | |
| + } | |
| + if (lane8 == 0) { | |
| +#pragma unroll | |
| + for (int i = 0; i < RT; ++i) { | |
| + const int row = rg * R + rr * RT + i; | |
| + if (row < nrows) { | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + if (SPLIT) { | |
| + part[((size_t) ch * ncols + c) * nrows + row] = acc[i][c]; | |
| + } else { | |
| + float * o = dst + c * stride_col_dst + row; | |
| + if (CHUNKED && kc0 > 0) { | |
| + *o += acc[i][c]; | |
| + } else { | |
| + *o = acc[i][c]; | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + if (SPLIT && cnt != nullptr) { | |
| + __threadfence(); | |
| + __syncthreads(); | |
| + if (tid == 0) { | |
| + const int prev = atomicAdd(cnt + rg, 1); | |
| + s_last = prev == nch - 1; | |
| + if (s_last) { | |
| + cnt[rg] = 0; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + if (s_last) { | |
| + __threadfence(); | |
| + for (int k = tid; k < R * ncols; k += nthr) { | |
| + const int c = k / R; | |
| + const int row = rg * R + (k - c * R); | |
| + if (row < nrows) { | |
| + float v = __ldcg(part + (size_t) c * nrows + row); | |
| + for (int ch2 = 1; ch2 < nch; ++ch2) { | |
| + v += __ldcg(part + ((size_t) ch2 * ncols + c) * nrows + row); | |
| + } | |
| + dst[c * stride_col_dst + row] = v; | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| +// SPLIT partials -> dst in chunk order: v = p0; v += p1; ... (the sequential kernel's *o = ..., *o += ... order) | |
| +template <int ncols> | |
| +static __global__ void pq2_split_reduce(const float * __restrict__ part, float * __restrict__ dst, const int nrows, | |
| + const int nch, const int stride_col_dst) { | |
| + const int row = blockIdx.x * blockDim.x + threadIdx.x; | |
| + if (row >= nrows) { | |
| + return; | |
| + } | |
| +#pragma unroll | |
| + for (int c = 0; c < ncols; ++c) { | |
| + float v = part[(size_t) c * nrows + row]; | |
| + for (int ch = 1; ch < nch; ++ch) { | |
| + v += part[((size_t) ch * ncols + c) * nrows + row]; | |
| + } | |
| + dst[(size_t) c * stride_col_dst + row] = v; | |
| + } | |
| +} | |
| + | |
| +// per-device row-group counters for the SPLIT last-arrival reduction (zeroed once; every launch leaves them zero) | |
| +static int * rl_split_counters(int n, cudaStream_t stream) { | |
| + static int * buf[GGML_CUDA_MAX_DEVICES] = {}; | |
| + static int cap[GGML_CUDA_MAX_DEVICES] = {}; | |
| + const int dev = ggml_cuda_get_device(); | |
| + if (cap[dev] < n) { | |
| + cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone; | |
| + CUDA_CHECK(cudaStreamIsCapturing(stream, &st)); | |
| + if (st != cudaStreamCaptureStatusNone) { | |
| + return nullptr; | |
| + } | |
| + if (buf[dev] != nullptr) { | |
| + CUDA_CHECK(cudaStreamSynchronize(stream)); | |
| + CUDA_CHECK(cudaFree(buf[dev])); | |
| + } | |
| + const int want = n > 65536 ? n : 65536; | |
| + CUDA_CHECK(cudaMalloc(&buf[dev], (size_t) want * sizeof(int))); | |
| + CUDA_CHECK(cudaMemsetAsync(buf[dev], 0, (size_t) want * sizeof(int), stream)); | |
| + cap[dev] = want; | |
| + } | |
| + return buf[dev]; | |
| +} | |
| + | |
| +// per-device scratch for the SPLIT partials (grown outside graph capture; null -> use the sequential kernel) | |
| +static float * rl_split_scratch(size_t nfloats, cudaStream_t stream) { | |
| + static float * buf[GGML_CUDA_MAX_DEVICES] = {}; | |
| + static size_t cap[GGML_CUDA_MAX_DEVICES] = {}; | |
| + const int dev = ggml_cuda_get_device(); | |
| + if (cap[dev] < nfloats) { | |
| + cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone; | |
| + CUDA_CHECK(cudaStreamIsCapturing(stream, &st)); | |
| + if (st != cudaStreamCaptureStatusNone) { | |
| + return nullptr; | |
| + } | |
| + if (buf[dev] != nullptr) { | |
| + CUDA_CHECK(cudaStreamSynchronize(stream)); | |
| + CUDA_CHECK(cudaFree(buf[dev])); | |
| + } | |
| + const size_t want = nfloats > ((size_t) 1 << 20) ? nfloats : ((size_t) 1 << 20); | |
| + CUDA_CHECK(cudaMalloc(&buf[dev], want * sizeof(float))); | |
| + cap[dev] = want; | |
| + } | |
| + return buf[dev]; | |
| +} | |
| + | |
| +template <int ncols, int RT, int MINB, bool CHUNKED, int RB = 16, bool SPLIT = false> | |
| +static void launch_shfl_impl(const uint8_t * x, const block_q8_1 * y, float * dst, const int nblocks, const int nrows, | |
| + const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, const int nsm, const int kc_len, | |
| + cudaStream_t stream) { | |
| + constexpr int R = RT * RB; | |
| + const size_t smem = (size_t) ncols * kc_len * TILED_WORDS_PER_KB * sizeof(uint32_t); | |
| + // resident blocks per SM for this instance at this tile size (occupancy API, cached; a handful of sizes per model) | |
| + static size_t cache_smem[16]; | |
| + static int cache_nb[16]; | |
| + static int cache_n = 0; | |
| + int per_sm = 0; | |
| + for (int i = 0; i < cache_n; ++i) { | |
| + if (cache_smem[i] == smem) { | |
| + per_sm = cache_nb[i]; | |
| + break; | |
| + } | |
| + } | |
| + if (per_sm == 0) { | |
| + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, mul_mat_vec_pq2_shfl<ncols, RT, RB, MINB, CHUNKED, SPLIT>, RB * 8, smem)); | |
| + if (per_sm < 1) { | |
| + per_sm = 1; | |
| + } | |
| + if (cache_n < 16) { | |
| + cache_smem[cache_n] = smem; | |
| + cache_nb[cache_n] = per_sm; | |
| + cache_n++; | |
| + } | |
| + } | |
| + const int nrg = (nrows + R - 1) / R; | |
| + if (SPLIT) { | |
| + const int nch = (nblocks + kc_len - 1) / kc_len; | |
| + static const int splitk_mode = getenv("GGML_CUDA_RL_SPLITK") ? atoi(getenv("GGML_CUDA_RL_SPLITK")) : 1; | |
| + float * part = rl_split_scratch((size_t) nch * ncols * nrows, stream); | |
| + int * cnt = splitk_mode == 2 ? nullptr : rl_split_counters(nrg, stream); | |
| + if (part != nullptr && (cnt != nullptr || splitk_mode == 2)) { | |
| + int per_ch = (nsm * per_sm) / nch; | |
| + per_ch = per_ch < 1 ? 1 : (per_ch > nrg ? nrg : per_ch); | |
| + mul_mat_vec_pq2_shfl<ncols, RT, RB, MINB, true, true><<<nch * per_ch, RB * 8, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, kc_len, g_rl_stamp, part, cnt); | |
| + if (cnt == nullptr) { | |
| + pq2_split_reduce<ncols><<<(nrows + 255) / 256, 256, 0, stream>>>(part, dst, nrows, nch, stride_col_dst); | |
| + } | |
| + return; | |
| + } | |
| + } | |
| + int grid = nsm * per_sm; | |
| + if (grid > nrg) { | |
| + grid = nrg; | |
| + } | |
| + mul_mat_vec_pq2_shfl<ncols, RT, RB, MINB, CHUNKED, false><<<grid, RB * 8, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, kc_len, g_rl_stamp, nullptr, nullptr); | |
| +} | |
| + | |
| +template <int ncols, int RT, int MINB, bool CHUNKED, int RB, bool SPLIT> | |
| +static void launch_q4_impl(const uint8_t * x, const block_q8_1 * y, float * dst, const int nunits, const int nrows, | |
| + const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, const int nsm, const int kc_len, | |
| + cudaStream_t stream) { | |
| + constexpr int R = RT * RB; | |
| + const size_t smem = (size_t) ncols * kc_len * Q4R_WPU * sizeof(uint32_t); | |
| + static size_t cache_smem[16]; | |
| + static int cache_nb[16]; | |
| + static int cache_n = 0; | |
| + int per_sm = 0; | |
| + for (int i = 0; i < cache_n; ++i) { | |
| + if (cache_smem[i] == smem) { | |
| + per_sm = cache_nb[i]; | |
| + break; | |
| + } | |
| + } | |
| + if (per_sm == 0) { | |
| + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per_sm, mul_mat_vec_q4_0_shfl<ncols, RT, RB, MINB, CHUNKED, SPLIT>, RB * 8, smem)); | |
| + if (per_sm < 1) { | |
| + per_sm = 1; | |
| + } | |
| + if (cache_n < 16) { | |
| + cache_smem[cache_n] = smem; | |
| + cache_nb[cache_n] = per_sm; | |
| + cache_n++; | |
| + } | |
| + } | |
| + const int nrg = (nrows + R - 1) / R; | |
| + if (SPLIT) { | |
| + const int nch = (nunits + kc_len - 1) / kc_len; | |
| + float * part = rl_split_scratch((size_t) nch * ncols * nrows, stream); | |
| + int * cnt = rl_split_counters(nrg, stream); | |
| + if (part != nullptr && cnt != nullptr) { | |
| + int per_ch = (nsm * per_sm) / nch; | |
| + per_ch = per_ch < 1 ? 1 : (per_ch > nrg ? nrg : per_ch); | |
| + mul_mat_vec_q4_0_shfl<ncols, RT, RB, MINB, true, true><<<nch * per_ch, RB * 8, smem, stream>>>(x, y, dst, nunits, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, kc_len, part, cnt); | |
| + return; | |
| + } | |
| + } | |
| + int grid = nsm * per_sm; | |
| + if (grid > nrg) { | |
| + grid = nrg; | |
| + } | |
| + mul_mat_vec_q4_0_shfl<ncols, RT, RB, MINB, CHUNKED, false><<<grid, RB * 8, smem, stream>>>(x, y, dst, nunits, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, kc_len, nullptr, nullptr); | |
| +} | |
| + | |
| +template <int ncols, int RT, int MINB, int RB = 16> | |
| +static void launch_q4(const uint8_t * x, const block_q8_1 * y, float * dst, const int nunits, const int nrows, | |
| + const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, const int nsm, const int kc_len, | |
| + cudaStream_t stream) { | |
| + if (kc_len >= nunits) { | |
| + launch_q4_impl<ncols, RT, MINB, false, RB, false>(x, y, dst, nunits, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, nsm, nunits, stream); | |
| + } else { | |
| + launch_q4_impl<ncols, RT, MINB, true, RB, true>(x, y, dst, nunits, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, nsm, kc_len, stream); | |
| + } | |
| +} | |
| + | |
| +template <int ncols, int RT, int MINB, int RB = 16> | |
| +static void launch_shfl(const uint8_t * x, const block_q8_1 * y, float * dst, const int nblocks, const int nrows, | |
| + const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, const int nsm, const int kc_len, | |
| + cudaStream_t stream) { | |
| + static const bool splitk = getenv("GGML_CUDA_RL_SPLITK") == nullptr || atoi(getenv("GGML_CUDA_RL_SPLITK")) != 0; | |
| + if (kc_len >= nblocks) { | |
| + launch_shfl_impl<ncols, RT, MINB, false, RB>(x, y, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, nsm, nblocks, stream); | |
| + } else if (splitk) { | |
| + launch_shfl_impl<ncols, RT, MINB, true, RB, true>(x, y, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, nsm, kc_len, stream); | |
| + } else { | |
| + launch_shfl_impl<ncols, RT, MINB, true, RB>(x, y, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst, nsm, kc_len, stream); | |
| + } | |
| +} | |
| + | |
| +static int rowlane_env_int(const char * name, const int dflt) { | |
| + const char * e = getenv(name); | |
| + return e ? atoi(e) : dflt; | |
| +} | |
| + | |
| +// widest ncols the shuffle kernel takes (0 = off): GGML_CUDA_ROWLANE_SHFL=0 disables, GGML_CUDA_ROWLANE_SHFL_MAXN caps | |
| +static int rowlane_env_shfl() { | |
| + static int maxn = -1; | |
| + if (maxn < 0) { | |
| + maxn = rowlane_env_int("GGML_CUDA_ROWLANE_SHFL", 1) ? rowlane_env_int("GGML_CUDA_ROWLANE_SHFL_MAXN", TILED_MAX_NCOLS) : 0; | |
| + } | |
| + return maxn; | |
| +} | |
| + | |
| +// activation tile budget per block in bytes (GGML_CUDA_ROWLANE_SMEM_KB, default 32: 3 blocks/SM on 96 KB; median step | |
| +// on a 1080 Ti, widths 3/4/5: 48 KB 40.45/44.97/53.24 ms, 32 KB 40.17/44.08/52.45, 24 KB 40.31/45.62/56.17) | |
| +static int rowlane_env_smem_budget() { | |
| + static int budget = -1; | |
| + if (budget < 0) { | |
| + int kb = rowlane_env_int("GGML_CUDA_ROWLANE_SMEM_KB", 32); | |
| + kb = kb < 10 ? 10 : (kb > 48 ? 48 : kb); | |
| + budget = kb * 1024; | |
| + } | |
| + return budget; | |
| +} | |
| + | |
| +// rows per block R and rows per thread RT by column count: ncols<=2 -> 8 rows, 1 per thread (320 threads at TB 40); | |
| +// 3-4 -> 16 rows, 2 per thread (192 threads at TB 24); 5-8 -> 32 rows, 2 per thread (256 threads at TB 16). | |
| +template <int ncols> static constexpr int tiled_rows() { return ncols <= 2 ? 8 : 32; } | |
| +template <int ncols> static constexpr int tiled_rows_thr() { return ncols <= 2 ? 1 : 2; } | |
| + | |
| +template <int ncols> | |
| +static void launch_tiled(const uint8_t * vx, const block_q8_1 * vy, float * dst, | |
| + const int nblocks, const int nrows, const int stride_row_x_bytes, const int stride_col_y, const int stride_col_dst, | |
| + const int TB, cudaStream_t stream) { | |
| + constexpr int R = tiled_rows<ncols>(); | |
| + constexpr int RT = tiled_rows_thr<ncols>(); | |
| + const dim3 grid((nrows + R - 1) / R, 1, 1); | |
| + const dim3 block(TB, R / RT, 1); | |
| + const size_t smem = (size_t) ncols * TB * TILED_WORDS_PER_KB * sizeof(uint32_t); | |
| + mul_mat_vec_pq2_tiled<ncols, R, RT><<<grid, block, smem, stream>>>(vx, vy, dst, nblocks, nrows, stride_row_x_bytes, stride_col_y, stride_col_dst); | |
| +} | |
| + | |
| +static int rowlane_env_persist() { | |
| + static int mode = -1; | |
| + if (mode < 0) { | |
| + const char * s = getenv("GGML_CUDA_ROWLANE_PERSIST"); | |
| + mode = s ? atoi(s) : 1; | |
| + } | |
| + return mode; | |
| +} | |
| + | |
| +static int rowlane_env_mode() { | |
| + static int mode = -1; | |
| + if (mode < 0) { | |
| + const char * s = getenv("GGML_CUDA_ROWLANE"); | |
| + mode = s ? atoi(s) : 1; | |
| + } | |
| + return mode; | |
| +} | |
| + | |
| +} // namespace | |
| + | |
| +// minimum column count for the Q4_0 row-lane kernel (0 = off, the default): microbench wins only from 5 columns | |
| +static int rowlane_env_q4_min() { | |
| + static int v = -1; | |
| + if (v < 0) { | |
| + const char * s = getenv("GGML_CUDA_ROWLANE_Q4"); | |
| + v = s ? atoi(s) : 0; // opt-in: lost end-to-end on E4B MTP (see patch_fork_q4_off.py) | |
| + } | |
| + return v; | |
| +} | |
| + | |
| +bool ggml_cuda_rowlane_applicable(const ggml_type type, const int cc, const int64_t ncols_dst) { | |
| + if (rowlane_env_mode() == 0) { | |
| + return false; | |
| + } | |
| + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc >= GGML_CUDA_CC_VOLTA) { | |
| + return false; | |
| + } | |
| + if (type != GGML_TYPE_PQ2_0 && | |
| + !(type == GGML_TYPE_Q4_0 && rowlane_env_q4_min() > 0 && ncols_dst >= rowlane_env_q4_min())) { | |
| + return false; | |
| + } | |
| + return ncols_dst >= 1 && ncols_dst <= TILED_MAX_NCOLS; | |
| +} | |
| + | |
| +bool ggml_cuda_rowlane_applicable_k(const ggml_type type, const int cc, const int64_t ncols_dst, const int64_t k) { | |
| + return ggml_cuda_rowlane_applicable(type, cc, ncols_dst) && k % 128 == 0; | |
| +} | |
| + | |
| +static void rowlane_impl( | |
| + const void * vx, const ggml_type type, const void * vy_q8_1, float * dst, | |
| + const int64_t ncols_x, const int64_t nrows_x, const int64_t ncols_dst, | |
| + const int64_t stride_row_x, const int64_t stride_col_y, const int64_t stride_col_dst, | |
| + const int nsm, cudaStream_t stream); | |
| + | |
| +// ---- diagnostic: GGML_CUDA_RL_PROF=N -- GPU time of the rowlane kernel alone, per shape, printed every N calls | |
| +static long g_rl_quant_miss = 0; | |
| +void ggml_cuda_rowlane_note_quantize() { | |
| + g_rl_quant_miss++; | |
| +} | |
| +void ggml_cuda_mul_mat_vec_q_rowlane( | |
| + const void * vx, const ggml_type type, const void * vy_q8_1, float * dst, | |
| + const int64_t ncols_x, const int64_t nrows_x, const int64_t ncols_dst, | |
| + const int64_t stride_row_x, const int64_t stride_col_y, const int64_t stride_col_dst, | |
| + const int nsm, cudaStream_t stream) { | |
| + static const int every = getenv("GGML_CUDA_RL_PROF") ? atoi(getenv("GGML_CUDA_RL_PROF")) : 0; | |
| + if (every <= 0) { | |
| + rowlane_impl(vx, type, vy_q8_1, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, nsm, stream); | |
| + return; | |
| + } | |
| + struct rec { long long key; cudaEvent_t a, b; int slot; }; | |
| + static std::vector<rec> pend; | |
| + static std::vector<cudaEvent_t> pool; | |
| + static std::map<long long, std::pair<double, long>> acc; | |
| + auto get_ev = [&]() { | |
| + if (pool.empty()) { | |
| + cudaEvent_t e; | |
| + CUDA_CHECK(cudaEventCreate(&e)); | |
| + return e; | |
| + } | |
| + cudaEvent_t e = pool.back(); | |
| + pool.pop_back(); | |
| + return e; | |
| + }; | |
| + static unsigned long long * stamps = nullptr; | |
| + static std::map<long long, double> span_acc; | |
| + if (stamps == nullptr) { | |
| + CUDA_CHECK(cudaMalloc(&stamps, (size_t) every * 2 * sizeof(unsigned long long))); | |
| + } | |
| + rec r; | |
| + r.key = (long long) nrows_x * 1000000LL + (long long) (ncols_x / 128) * 100LL + (long long) ncols_dst; | |
| + r.a = get_ev(); | |
| + r.b = get_ev(); | |
| + r.slot = (int) pend.size(); | |
| + { | |
| + const unsigned long long init[2] = { ~0ull, 0ull }; | |
| + CUDA_CHECK(cudaMemcpyAsync(stamps + 2 * r.slot, init, sizeof(init), cudaMemcpyHostToDevice, stream)); | |
| + } | |
| + g_rl_stamp = stamps + 2 * r.slot; | |
| + CUDA_CHECK(cudaEventRecord(r.a, stream)); | |
| + rowlane_impl(vx, type, vy_q8_1, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, nsm, stream); | |
| + CUDA_CHECK(cudaEventRecord(r.b, stream)); | |
| + g_rl_stamp = nullptr; | |
| + pend.push_back(r); | |
| + if ((int) pend.size() >= every) { | |
| + CUDA_CHECK(cudaEventSynchronize(pend.back().b)); | |
| + std::vector<unsigned long long> hs((size_t) every * 2); | |
| + CUDA_CHECK(cudaMemcpy(hs.data(), stamps, hs.size() * sizeof(unsigned long long), cudaMemcpyDeviceToHost)); | |
| + for (auto & q : pend) { | |
| + float ms = 0.0f; | |
| + CUDA_CHECK(cudaEventElapsedTime(&ms, q.a, q.b)); | |
| + auto & v = acc[q.key]; | |
| + v.first += ms; | |
| + v.second++; | |
| + const unsigned long long t0 = hs[2 * q.slot], t1 = hs[2 * q.slot + 1]; | |
| + span_acc[q.key] += (t1 > t0 && t0 != ~0ull) ? (double) (t1 - t0) * 1e-6 : 0.0; // ns -> ms | |
| + pool.push_back(q.a); | |
| + pool.push_back(q.b); | |
| + } | |
| + pend.clear(); | |
| + fprintf(stderr, "RLPROF %d calls (quantize misses so far %ld)\n", every, g_rl_quant_miss); | |
| + for (const auto & [k, v] : acc) { | |
| + const long long rows = k / 1000000LL, kb = (k / 100LL) % 10000LL, nc = k % 100LL; | |
| + const double us = 1000.0 * v.first / v.second; | |
| + const double span_us = 1000.0 * span_acc[k] / v.second; | |
| + fprintf(stderr, "RLPROF rows %6lld K %6lld n %lld : %8.1f us (%5.0f GB/s) span %8.1f us (%5.0f GB/s) x%ld\n", rows, kb * 128, nc, us, | |
| + (double) rows * kb * 34.0 / (us * 1e3), span_us, span_us > 0 ? (double) rows * kb * 34.0 / (span_us * 1e3) : 0.0, v.second); | |
| + } | |
| + acc.clear(); | |
| + span_acc.clear(); | |
| + } | |
| +} | |
| + | |
| +static void rowlane_impl( | |
| + const void * vx, const ggml_type type, const void * vy_q8_1, float * dst, | |
| + const int64_t ncols_x, const int64_t nrows_x, const int64_t ncols_dst, | |
| + const int64_t stride_row_x, const int64_t stride_col_y, const int64_t stride_col_dst, | |
| + const int nsm, cudaStream_t stream) { | |
| + if (type == GGML_TYPE_Q4_0) { | |
| + GGML_ASSERT(ncols_x % 128 == 0); | |
| + const int nunits = (int) (ncols_x / 128); | |
| + const int nrows = (int) nrows_x; | |
| + const int srb = (int) (stride_row_x * 18); // stride in q4_0 blocks -> bytes | |
| + const int budget = rowlane_env_smem_budget(); | |
| + const int per_u = (int) ncols_dst * Q4R_WPU * (int) sizeof(uint32_t); | |
| + int nch = (nunits * per_u + budget - 1) / budget; | |
| + int kc = nunits; | |
| + for (;;) { | |
| + kc = (nunits + nch - 1) / nch; | |
| + if (kc * per_u <= budget || kc == 1) { | |
| + break; | |
| + } | |
| + ++nch; | |
| + } | |
| + const uint8_t * x = (const uint8_t *) vx; | |
| + const block_q8_1 * y = (const block_q8_1 *) vy_q8_1; | |
| + const int scy = (int) stride_col_y; | |
| + const int scd = (int) stride_col_dst; | |
| + switch (ncols_dst) { | |
| + case 1: launch_q4<1, 1, 4>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 2: launch_q4<2, 1, 4>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 3: launch_q4<3, 2, 4>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 4: launch_q4<4, 2, 4>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 5: launch_q4<5, 2, 3>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 6: launch_q4<6, 2, 3>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + case 7: launch_q4<7, 2, 3>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + default: launch_q4<8, 2, 3>(x, y, dst, nunits, nrows, srb, scy, scd, nsm, kc, stream); break; | |
| + } | |
| + return; | |
| + } | |
| + GGML_ASSERT(type == GGML_TYPE_PQ2_0); | |
| + GGML_ASSERT(ncols_x % 128 == 0); | |
| + const int nblocks = (int) (ncols_x / 128); | |
| + const int nrows = (int) nrows_x; | |
| + const int stride_row_x_bytes = (int) (stride_row_x * 34); | |
| + | |
| + // K-blocks per tile: as large as the smem budget allows, sized so the tiles divide K evenly | |
| + // (K=5120 -> 40 x1, 6144 -> 24 x2, 10240 -> 40 x2, 12288 -> 32 x3, 17408 -> 34 x4). | |
| + // smem = ncols * TB * 144 B; keep >= 2-3 blocks per SM (48 KB): 40 -> 5.8 KB/col, 24 -> 3.5, 16 -> 2.3 | |
| + const int tb_max = ncols_dst <= 2 ? TILED_TB_MAX : 16; | |
| + const int ntiles = (nblocks + tb_max - 1) / tb_max; | |
| + int TB = (nblocks + ntiles - 1) / ntiles; | |
| + TB = (TB + 3) / 4 * 4; // keep the block width a multiple of 4 for the staging loop | |
| + if (TB > tb_max) { | |
| + TB = tb_max; | |
| + } | |
| + | |
| + const uint8_t * x = (const uint8_t *) vx; | |
| + const block_q8_1 * y = (const block_q8_1 *) vy_q8_1; | |
| + if (ncols_dst <= rowlane_env_shfl() && (nblocks & 7) == 0) { | |
| + // K-chunk length: the fewest chunks whose tile fits the budget, balanced, rounded up to a multiple of 8 blocks | |
| + const int budget = rowlane_env_smem_budget(); | |
| + const int per_kb = (int) ncols_dst * TILED_WORDS_PER_KB * (int) sizeof(uint32_t); | |
| + int nch = (nblocks * per_kb + budget - 1) / budget; | |
| + int kc = nblocks; | |
| + for (;;) { | |
| + kc = (nblocks + nch - 1) / nch; | |
| + kc = (kc + 7) & ~7; | |
| + if (kc * per_kb <= budget || kc == 8) { | |
| + break; | |
| + } | |
| + ++nch; | |
| + } | |
| + const int scy = (int) stride_col_y; | |
| + const int scd = (int) stride_col_dst; | |
| + static const int chunk_ok = rowlane_env_int("GGML_CUDA_ROWLANE_SHFL_CHUNK", 1); // 0: shapes needing chunks fall back | |
| + if (kc < nblocks && !chunk_ok) { | |
| + goto no_shfl; | |
| + } | |
| + // in-model layout A/B (patch_fork_rl_n4var.py): GGML_CUDA_RL_N<w>[_LONG]=v, 0 = production | |
| + static const int var_n[5][2] = { | |
| + {0, 0}, {0, 0}, | |
| + { rowlane_env_int("GGML_CUDA_RL_N2", 0), rowlane_env_int("GGML_CUDA_RL_N2_LONG", 0) }, | |
| + { rowlane_env_int("GGML_CUDA_RL_N3", 0), rowlane_env_int("GGML_CUDA_RL_N3_LONG", 0) }, | |
| + { rowlane_env_int("GGML_CUDA_RL_N4", 0), rowlane_env_int("GGML_CUDA_RL_N4_LONG", 0) } }; | |
| + const int var = ncols_dst >= 2 && ncols_dst <= 4 ? var_n[ncols_dst][nblocks > 48 ? 1 : 0] : 0; | |
| +#define RL_VAR_LAUNCH(NC, RT0) \ | |
| + switch (var) { \ | |
| + case 1: launch_shfl<NC, 2, 2, 32>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + case 2: launch_shfl<NC, 2, 3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + case 3: launch_shfl<NC, 1, 4>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + case 4: launch_shfl<NC, 2, 2>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + case 5: launch_shfl<NC, 1, 2, 32>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + default: launch_shfl<NC, RT0, 4>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; \ | |
| + } | |
| + switch (ncols_dst) { | |
| + case 1: launch_shfl<1, 1, 4>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; | |
| + case 2: RL_VAR_LAUNCH(2, 1) break; | |
| + case 3: RL_VAR_LAUNCH(3, 2) break; | |
| + case 4: RL_VAR_LAUNCH(4, 2) break; | |
| + case 5: | |
| + if (nblocks <= 40) { | |
| + launch_shfl<5, 2, 2, 32>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); | |
| + } else { | |
| + launch_shfl<5, 2, 3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); | |
| + } | |
| + break; | |
| + case 6: launch_shfl<6, 2, 3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; | |
| + case 7: launch_shfl<7, 2, 3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; | |
| + default: launch_shfl<8, 2, 3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, scy, scd, nsm, kc, stream); break; | |
| + } | |
| + return; | |
| + } | |
| +no_shfl: | |
| + if (ncols_dst <= 3 && rowlane_env_persist()) { | |
| + // ncols 1: 1 row/thread, <= 48 threads per row, 3 resident blocks (long rows at ncols 2 prefer 2 with RT=1); | |
| + // ncols 2-3 on rows that fit: 2 rows/thread, <= 20 threads per row, 3 resident blocks (160-thread blocks) | |
| + const int RB = 8; | |
| + const bool wide = ncols_dst >= 2 && nblocks <= 48; | |
| + const int RT = wide ? 2 : 1; | |
| + const int tb_max = wide ? 20 : 48; | |
| + const int nt = (nblocks + tb_max - 1) / tb_max; | |
| + const int TB = (nblocks + nt - 1) / nt; | |
| + const int R = RT * RB; | |
| + const size_t smem = (size_t) ncols_dst * nblocks * TILED_WORDS_PER_KB * sizeof(uint32_t) + (size_t) 2 * R * TB * ncols_dst * sizeof(float); | |
| + if (smem <= 48 * 1024 && (ncols_dst <= 2 || wide)) { | |
| + const int nrg = (nrows + R - 1) / R; | |
| + const int minb = (RT == 2 || ncols_dst == 1 || nblocks <= 48) ? 3 : 2; | |
| + int grid = nsm * minb; | |
| + if (grid > nrg) { | |
| + grid = nrg; | |
| + } | |
| + const dim3 block(TB, RB, 1); | |
| + if (RT == 2) { | |
| + if (ncols_dst == 2) { | |
| + mul_mat_vec_pq2_persist<2, 2, 160, 3><<<grid, block, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst); | |
| + } else { | |
| + mul_mat_vec_pq2_persist<3, 2, 160, 3><<<grid, block, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst); | |
| + } | |
| + } else if (ncols_dst == 1) { | |
| + mul_mat_vec_pq2_persist<1, 1, 384, 3><<<grid, block, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst); | |
| + } else if (minb == 3) { | |
| + mul_mat_vec_pq2_persist<2, 1, 384, 3><<<grid, block, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst); | |
| + } else { | |
| + mul_mat_vec_pq2_persist<2, 1, 384, 2><<<grid, block, smem, stream>>>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst); | |
| + } | |
| + return; | |
| + } | |
| + } | |
| + switch (ncols_dst) { | |
| + case 1: launch_tiled<1>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 2: launch_tiled<2>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 3: launch_tiled<3>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 4: launch_tiled<4>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 5: launch_tiled<5>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 6: launch_tiled<6>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + case 7: launch_tiled<7>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + default: launch_tiled<8>(x, y, dst, nblocks, nrows, stride_row_x_bytes, (int) stride_col_y, (int) stride_col_dst, TB, stream); break; | |
| + } | |
| +} | |
| diff --git a/ggml/src/ggml-cuda/mmvq-rowlane.cuh b/ggml/src/ggml-cuda/mmvq-rowlane.cuh | |
| new file mode 100644 | |
| index 0000000..425785e | |
| --- /dev/null | |
| +++ b/ggml/src/ggml-cuda/mmvq-rowlane.cuh | |
| +#pragma once | |
| + | |
| +#include "common.cuh" | |
| + | |
| +// Row-per-lane GEMV for PQ2_0 / PTQ1_0 on pre-Volta NVIDIA. See mmvq-rowlane.cu. | |
| +bool ggml_cuda_rowlane_applicable(ggml_type type, int cc, int64_t ncols_dst); | |
| +// as above, and the row length K (Q4_0 needs K % 128 == 0) | |
| +bool ggml_cuda_rowlane_applicable_k(ggml_type type, int cc, int64_t ncols_dst, int64_t k); | |
| + | |
| +void ggml_cuda_mul_mat_vec_q_rowlane( | |
| + const void * vx, ggml_type type, const void * vy_q8_1, float * dst, | |
| + int64_t ncols_x, int64_t nrows_x, int64_t ncols_dst, | |
| + int64_t stride_row_x, int64_t stride_col_y, int64_t stride_col_dst, | |
| + int nsm, cudaStream_t stream); | |
| + | |
| +// diagnostic counter (GGML_CUDA_RL_PROF): a rowlane-eligible GEMV had to quantize its activations itself | |
| +void ggml_cuda_rowlane_note_quantize(); | |
| diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu | |
| index bf51b61..7c16489 100644 | |
| --- a/ggml/src/ggml-cuda/mmvq.cu | |
| +++ b/ggml/src/ggml-cuda/mmvq.cu | |
| #include "mmvq.cuh" | |
| +#include "mmvq-rowlane.cuh" | |
| +#include "fwht.cuh" | |
| #include "quantize.cuh" | |
| #include "unary.cuh" | |
| #include "vecdotq.cuh" | |
| bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { | |
| return false; | |
| } | |
| #if !defined(GGML_USE_HIP) | |
| - if (type == GGML_TYPE_PTQ1_0 && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) { | |
| + if (type == GGML_TYPE_PTQ1_0 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // PASCAL PATCH: allow the mmvq (GEMV) path below Turing | |
| return ne11 <= 7; | |
| } | |
| #endif | |
| void ggml_cuda_mul_mat_vec_q( | |
| } | |
| const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); | |
| + | |
| + // the Hadamard transform that produced src1 already wrote its q8_1 copy (fwht.cu): skip the quantize | |
| + if (!ids && !fusion && ne02 == 1 && ne03 == 1 && ne12 == 1 && ne13 == 1 && ne10_padded == ne10 && | |
| + src1->nb[1] == (size_t) ne10 * sizeof(float) && | |
| + ggml_cuda_rowlane_applicable_k(src0->type, ggml_cuda_info().devices[ctx.device].cc, ne1, ne00)) { | |
| + const void * preq = ggml_cuda_fwht_q8_find(src1, ne10 * ne11); | |
| + if (preq != nullptr) { | |
| + ggml_cuda_mul_mat_vec_q_rowlane(src0->data, src0->type, preq, dst_d, | |
| + ne00, ne01, ne1, src0->nb[1] / ts_src0, ne10_padded / QK8_1, dst->nb[1] / ts_dst, | |
| + ggml_cuda_info().devices[ctx.device].nsm, stream); | |
| + return; | |
| + } | |
| + } | |
| + | |
| + if (src0->type == GGML_TYPE_PQ2_0 && ggml_cuda_rowlane_applicable_k(src0->type, ggml_cuda_info().devices[ctx.device].cc, ne1, ne00)) { | |
| + ggml_cuda_rowlane_note_quantize(); | |
| + } | |
| ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); | |
| { | |
| const int64_t s11 = src1->nb[1] / ts_src1; | |
| void ggml_cuda_mul_mat_vec_q( | |
| const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; | |
| + // PASCAL PATCH: row-per-lane GEMV for the Prism ternary formats (see mmvq-rowlane.cu) | |
| + if (!ids && !fusion && ne02 == 1 && ne03 == 1 && ne12 == 1 && ne13 == 1 && | |
| + ggml_cuda_rowlane_applicable_k(src0->type, ggml_cuda_info().devices[ctx.device].cc, ne1, ne00)) { | |
| + ggml_cuda_mul_mat_vec_q_rowlane(src0->data, src0->type, src1_q8_1.get(), dst_d, | |
| + ne00, ne01, ne1, s01, s11, s1, ggml_cuda_info().devices[ctx.device].nsm, stream); | |
| + return; | |
| + } | |
| + | |
| mul_mat_vec_q_switch_type( | |
| src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, | |
| ne01, ncols_dst, s01, stride_col_y, stride_col_dst, | |
| diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu | |
| index 1463169..20bda71 100644 | |
| --- a/ggml/src/ggml-cuda/ssm-conv.cu | |
| +++ b/ggml/src/ggml-cuda/ssm-conv.cu | |
| void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, g | |
| out->nb[2], nc, nr, n_t, n_s, stream); | |
| } | |
| } | |
| + | |
| + | |
| +// ---- fused GDN conv step (sm_6x), one sequence, conv kernel DC = 4, NT tokens. Block = one 128-channel head, | |
| +// thread = one channel. Per channel c: | |
| +// x[i] = i < 3 ? cs[3c + i] : qkv[(i-3)*qkv_st + c*qkv_sc] (the CONCAT, kept in registers) | |
| +// snap_s = x[sidx_s .. sidx_s+2] -> snaps.dst[s][3c ..] (the rollback-snapshot CPYs) | |
| +// y[t] = silu(sum_j x[t+j] * w[c][j]) (SSM_CONV + SILU, as ssm_conv_f32) | |
| +// l2[t][h] = y[t] * rsqrt(max(sum_d y^2, eps^2)) for heads h < n_l2_heads (L2_NORM, as l2_norm_f32<32>) | |
| +template <int NT> | |
| +static __global__ void __launch_bounds__(128) conv_step_f32( | |
| + const float * cs, const int32_t * cs_rows, const int64_t cs_row_stride, const float * __restrict__ qkv, | |
| + const int64_t qkv_st, const int64_t qkv_sc, const float * __restrict__ w, const int64_t w_sc, float * __restrict__ y, | |
| + const int64_t y_st, float * __restrict__ l2, const int n_l2_heads, const float eps, const ggml_cuda_conv_step_snaps snaps) { | |
| + constexpr int DC = 4; | |
| + if (cs_rows != nullptr) { | |
| + cs += (int64_t) cs_rows[0] * cs_row_stride; // folded gather: the state row itself (may be a snapshot slot) | |
| + } | |
| + constexpr int NX = DC - 1 + NT; | |
| + const int h = blockIdx.x; | |
| + const int d = threadIdx.x; | |
| + const int64_t c = (int64_t) h * 128 + d; | |
| + | |
| + float x[NX]; | |
| +#pragma unroll | |
| + for (int j = 0; j < DC - 1; ++j) { | |
| + x[j] = cs[c * (DC - 1) + j]; | |
| + } | |
| +#pragma unroll | |
| + for (int t = 0; t < NT; ++t) { | |
| + x[DC - 1 + t] = qkv[t * qkv_st + c * qkv_sc]; | |
| + } | |
| + float wr[DC]; | |
| +#pragma unroll | |
| + for (int j = 0; j < DC; ++j) { | |
| + wr[j] = w[c * w_sc + j]; | |
| + } | |
| + | |
| + for (int s = 0; s < snaps.n; ++s) { | |
| + const int sidx = snaps.sidx[s]; | |
| + float * dst = snaps.dst[s] + c * (DC - 1); | |
| +#pragma unroll | |
| + for (int i = 0; i < NX; ++i) { | |
| + if (i >= sidx && i < sidx + DC - 1) { | |
| + dst[i - sidx] = x[i]; | |
| + } | |
| + } | |
| + } | |
| + | |
| + float yv[NT]; | |
| +#pragma unroll | |
| + for (int t = 0; t < NT; ++t) { | |
| + float sumf = 0.0f; | |
| +#pragma unroll | |
| + for (int j = 0; j < DC; ++j) { | |
| + sumf += x[t + j] * wr[j]; | |
| + } | |
| + const float b = 0.0f; // ssm_conv_f32 without a bias adds b = 0.0f | |
| + sumf += b; | |
| + yv[t] = ggml_cuda_op_silu_single(sumf); | |
| + y[t * y_st + c] = yv[t]; | |
| + } | |
| + | |
| + if (h >= n_l2_heads) { | |
| + return; // block-uniform | |
| + } | |
| + __shared__ float sy[NT][128]; | |
| + __shared__ float sscale[NT]; | |
| +#pragma unroll | |
| + for (int t = 0; t < NT; ++t) { | |
| + sy[t][d] = yv[t]; | |
| + } | |
| + __syncthreads(); | |
| + if (d < WARP_SIZE) { | |
| +#pragma unroll | |
| + for (int t = 0; t < NT; ++t) { | |
| + float tmp = 0.0f; | |
| +#pragma unroll | |
| + for (int col = d; col < 128; col += WARP_SIZE) { | |
| + const float xi = sy[t][col]; | |
| + tmp += xi * xi; | |
| + } | |
| + tmp = warp_reduce_sum(tmp); | |
| + if (d == 0) { | |
| + sscale[t] = rsqrtf(fmaxf(tmp, eps * eps)); | |
| + } | |
| + } | |
| + } | |
| + __syncthreads(); | |
| +#pragma unroll | |
| + for (int t = 0; t < NT; ++t) { | |
| + l2[((int64_t) t * n_l2_heads + h) * 128 + d] = sscale[t] * yv[t]; | |
| + } | |
| +} | |
| + | |
| +void ggml_cuda_conv_step_f32(const float * cs, const int32_t * cs_rows, int64_t cs_row_stride, const float * qkv, int64_t qkv_st, | |
| + int64_t qkv_sc, const float * w, int64_t w_sc, float * y, int64_t y_st, float * l2, int n_l2_heads, | |
| + float eps, const ggml_cuda_conv_step_snaps & snaps, int64_t C, int64_t n, cudaStream_t stream) { | |
| + GGML_ASSERT(C % 128 == 0 && n >= 1 && n <= 8 && snaps.n <= 8); | |
| + const dim3 grid((unsigned) (C / 128)); | |
| + switch (n) { | |
| + case 1: conv_step_f32<1><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 2: conv_step_f32<2><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 3: conv_step_f32<3><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 4: conv_step_f32<4><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 5: conv_step_f32<5><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 6: conv_step_f32<6><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 7: conv_step_f32<7><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + case 8: conv_step_f32<8><<<grid, 128, 0, stream>>>(cs, cs_rows, cs_row_stride, qkv, qkv_st, qkv_sc, w, w_sc, y, y_st, l2, n_l2_heads, eps, snaps); break; | |
| + } | |
| + CUDA_CHECK(cudaGetLastError()); | |
| +} | |
| diff --git a/ggml/src/ggml-cuda/ssm-conv.cuh b/ggml/src/ggml-cuda/ssm-conv.cuh | |
| index 8514ca8..e6cbbd8 100644 | |
| --- a/ggml/src/ggml-cuda/ssm-conv.cuh | |
| +++ b/ggml/src/ggml-cuda/ssm-conv.cuh | |
| #include "common.cuh" | |
| +// fused GDN conv step (CONCAT + snapshot CPYs + SSM_CONV + SILU + L2_NORM), one sequence, conv kernel 4 | |
| +struct ggml_cuda_conv_step_snaps { | |
| + float * dst[8]; | |
| + int sidx[8]; | |
| + int n; | |
| +}; | |
| +// cs_rows != null: cs is a row base, the state is row cs_rows[0] (row stride cs_row_stride floats) -- the folded gather | |
| +void ggml_cuda_conv_step_f32(const float * cs, const int32_t * cs_rows, int64_t cs_row_stride, const float * qkv, int64_t qkv_st, | |
| + int64_t qkv_sc, const float * w, int64_t w_sc, float * y, int64_t y_st, float * l2, int n_l2_heads, | |
| + float eps, const ggml_cuda_conv_step_snaps & snaps, int64_t C, int64_t n, cudaStream_t stream); | |
| + | |
| void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * bias_add_node = nullptr, ggml_tensor * silu_dst = nullptr); | |
| diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu | |
| index 9681cd2..a01aeb3 100644 | |
| --- a/ggml/src/ggml-cuda/top-k.cu | |
| +++ b/ggml/src/ggml-cuda/top-k.cu | |
| #include "argsort.cuh" | |
| #include "top-k.cuh" | |
| +#include <climits> | |
| +#include <cstdlib> | |
| +#include <utility> | |
| + | |
| #ifdef GGML_CUDA_USE_CUB | |
| # include <cub/cub.cuh> | |
| # if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) | |
| static int next_power_of_2(int x) { | |
| #endif // CUB_TOP_K_AVAILABLE | |
| + | |
| +// ---- fast top-k for small k on long rows (see tools/kernel-lane/bonsai/patch_fork_topk.py) ---- | |
| +#define TOPK_FAST_THREADS 256 | |
| +#define TOPK_FAST_PT 16 // elements per thread per block: a block covers 4096 elements | |
| +#define TOPK_FAST_KMAX 32 | |
| + | |
| +// block-wide argmax of (v, i) pairs, larger v wins, equal v -> smaller i wins; result broadcast to every thread | |
| +static __device__ __forceinline__ void topk_block_argmax(float & v, int & i, float * sv, int * si) { | |
| +#pragma unroll | |
| + for (int o = 16; o > 0; o >>= 1) { | |
| + const float v2 = __shfl_xor_sync(0xffffffffu, v, o); | |
| + const int i2 = __shfl_xor_sync(0xffffffffu, i, o); | |
| + if (v2 > v || (v2 == v && i2 < i)) { | |
| + v = v2; | |
| + i = i2; | |
| + } | |
| + } | |
| + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; | |
| + __syncthreads(); // sv/si reuse across rounds | |
| + if (lane == 0) { | |
| + sv[warp] = v; | |
| + si[warp] = i; | |
| + } | |
| + __syncthreads(); | |
| + if (warp == 0) { | |
| + const int nw = blockDim.x / 32; | |
| + v = lane < nw ? sv[lane] : -INFINITY; | |
| + i = lane < nw ? si[lane] : INT_MAX; | |
| +#pragma unroll | |
| + for (int o = 16; o > 0; o >>= 1) { | |
| + const float v2 = __shfl_xor_sync(0xffffffffu, v, o); | |
| + const int i2 = __shfl_xor_sync(0xffffffffu, i, o); | |
| + if (v2 > v || (v2 == v && i2 < i)) { | |
| + v = v2; | |
| + i = i2; | |
| + } | |
| + } | |
| + if (lane == 0) { | |
| + sv[0] = v; | |
| + si[0] = i; | |
| + } | |
| + } | |
| + __syncthreads(); | |
| + v = sv[0]; | |
| + i = si[0]; | |
| +} | |
| + | |
| +// src: [n rows of ncols]; candidates out: per (row, block) k (value, index) pairs, best first. | |
| +// in_idx != nullptr: src holds candidate values whose original indices are in_idx (merge pass). | |
| +static __global__ void topk_fast_select(const float * __restrict__ src, const int * __restrict__ in_idx, | |
| + float * __restrict__ out_v, int * __restrict__ out_i, | |
| + const int ncols, const int k) { | |
| + __shared__ float sv[TOPK_FAST_THREADS / 32]; | |
| + __shared__ int si[TOPK_FAST_THREADS / 32]; | |
| + const int row = blockIdx.y; | |
| + const int chunk = TOPK_FAST_THREADS * TOPK_FAST_PT; | |
| + const int beg = blockIdx.x * chunk; | |
| + const float * x = src + (size_t) row * ncols; | |
| + const int * xi = in_idx ? in_idx + (size_t) row * ncols : nullptr; | |
| + | |
| + float v[TOPK_FAST_PT]; | |
| + int id[TOPK_FAST_PT]; | |
| +#pragma unroll | |
| + for (int e = 0; e < TOPK_FAST_PT; ++e) { | |
| + const int j = beg + e * TOPK_FAST_THREADS + threadIdx.x; | |
| + const bool ok = j < ncols; | |
| + const float xv = ok ? x[j] : -INFINITY; | |
| + // a NaN would never be selected (all comparisons false): load it as -inf with its real index, so the output | |
| + // indices are always real columns (bit test: -use_fast_math may fold x != x) | |
| + v[e] = (__float_as_uint(xv) & 0x7fffffffu) > 0x7f800000u ? -INFINITY : xv; | |
| + id[e] = ok ? (xi ? xi[j] : j) : INT_MAX; | |
| + } | |
| + float * ov = out_v + ((size_t) row * gridDim.x + blockIdx.x) * k; | |
| + int * oi = out_i + ((size_t) row * gridDim.x + blockIdx.x) * k; | |
| + for (int r = 0; r < k; ++r) { | |
| + float bv = -INFINITY; | |
| + int bi = INT_MAX; | |
| +#pragma unroll | |
| + for (int e = 0; e < TOPK_FAST_PT; ++e) { | |
| + if (v[e] > bv || (v[e] == bv && id[e] < bi)) { | |
| + bv = v[e]; | |
| + bi = id[e]; | |
| + } | |
| + } | |
| + topk_block_argmax(bv, bi, sv, si); | |
| + if (threadIdx.x == 0) { | |
| + ov[r] = bv; | |
| + oi[r] = bi; | |
| + } | |
| +#pragma unroll | |
| + for (int e = 0; e < TOPK_FAST_PT; ++e) { | |
| + if (id[e] == bi) { | |
| + v[e] = -INFINITY; | |
| + id[e] = INT_MAX; | |
| + } | |
| + } | |
| + } | |
| +} | |
| + | |
| +static bool top_k_fast(ggml_cuda_pool & pool, const float * src, int * dst, const int64_t ncols, const int64_t nrows, | |
| + const int k, cudaStream_t stream) { | |
| + static const bool on = getenv("GGML_CUDA_TOPK_FAST") == nullptr || atoi(getenv("GGML_CUDA_TOPK_FAST")) != 0; | |
| + const int chunk = TOPK_FAST_THREADS * TOPK_FAST_PT; | |
| + if (!on || k < 1 || k > TOPK_FAST_KMAX || ncols < 2048 || ncols > (int64_t) chunk * 1024 || nrows > 65535) { | |
| + return false; | |
| + } | |
| + int nb = (int) ((ncols + chunk - 1) / chunk); | |
| + ggml_cuda_pool_alloc<float> cv(pool, (size_t) nrows * nb * k); | |
| + ggml_cuda_pool_alloc<int> ci(pool, (size_t) nrows * nb * k); | |
| + topk_fast_select<<<dim3(nb, (unsigned) nrows), TOPK_FAST_THREADS, 0, stream>>>(src, nullptr, cv.get(), ci.get(), (int) ncols, k); | |
| + // merge passes until one block per row remains; the last one writes the indices | |
| + int m = nb * k; | |
| + ggml_cuda_pool_alloc<float> cv2(pool, (size_t) nrows * k * ((m + chunk - 1) / chunk)); | |
| + ggml_cuda_pool_alloc<int> ci2(pool, (size_t) nrows * k * ((m + chunk - 1) / chunk)); | |
| + float * in_v = cv.get(); | |
| + int * in_i = ci.get(); | |
| + float * o_v = cv2.get(); | |
| + int * o_i = ci2.get(); | |
| + while (true) { | |
| + const int nb2 = (m + chunk - 1) / chunk; | |
| + topk_fast_select<<<dim3(nb2, (unsigned) nrows), TOPK_FAST_THREADS, 0, stream>>>(in_v, in_i, o_v, o_i, m, k); | |
| + if (nb2 == 1) { | |
| + CUDA_CHECK(cudaMemcpyAsync(dst, o_i, (size_t) nrows * k * sizeof(int), cudaMemcpyDeviceToDevice, stream)); | |
| + return true; | |
| + } | |
| + m = nb2 * k; | |
| + std::swap(in_v, o_v); | |
| + std::swap(in_i, o_i); | |
| + } | |
| +} | |
| + | |
| void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { | |
| const ggml_tensor * src0 = dst->src[0]; | |
| const float * src0_d = (const float *) src0->data; | |
| void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { | |
| const int64_t nrows = ggml_nrows(src0); | |
| const int64_t k = dst->ne[0]; | |
| ggml_cuda_pool & pool = ctx.pool(); | |
| + | |
| + if (top_k_fast(pool, src0_d, dst_d, ncols, nrows, (int) k, stream)) { | |
| + return; | |
| + } | |
| #ifdef CUB_TOP_K_AVAILABLE | |
| // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented | |
| // https://github.com/NVIDIA/cccl/issues/6391 | |
| diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp | |
| index 8e0b369..49918da 100644 | |
| --- a/src/llama-arch.cpp | |
| +++ b/src/llama-arch.cpp | |
| static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = { | |
| { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, | |
| { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, | |
| { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, | |
| + { LLM_TENSOR_NEXTN_DRAFT_VOCAB, "blk.%d.nextn.draft_vocab" }, | |
| { LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" }, | |
| { LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" }, | |
| { LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" }, | |
| static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = { | |
| {LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, | |
| {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, | |
| {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, | |
| + {LLM_TENSOR_NEXTN_DRAFT_VOCAB, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, | |
| // Nemotron 3 Super | |
| // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU | |
| {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, | |
| diff --git a/src/llama-arch.h b/src/llama-arch.h | |
| index 2b7172e..fb17b4e 100644 | |
| --- a/src/llama-arch.h | |
| +++ b/src/llama-arch.h | |
| enum llm_tensor { | |
| LLM_TENSOR_NEXTN_HNORM, | |
| LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, | |
| LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, | |
| + LLM_TENSOR_NEXTN_DRAFT_VOCAB, | |
| LLM_TENSOR_MASKED_EMBD_CENTROIDS, | |
| LLM_TENSOR_MASKED_EMBD_ORDERING, | |
| LLM_TENSOR_FC, | |
| diff --git a/src/llama-context.cpp b/src/llama-context.cpp | |
| index 2c8fbd9..d840144 100644 | |
| --- a/src/llama-context.cpp | |
| +++ b/src/llama-context.cpp | |
| #include "llama-context.h" | |
| +#include <map> | |
| + | |
| #include "ggml.h" | |
| #include "llama-arch.h" | |
| #include "llama-graph.h" | |
| llama_context::llama_context( | |
| } | |
| cparams.n_rs_seq = params.n_rs_seq; | |
| + // diagnostic: LLAMA_N_RS_SEQ=n forces the rollback window, so llama-bench can time the speculative-verify graph shape | |
| + if (const char * e = getenv("LLAMA_N_RS_SEQ")) { | |
| + cparams.n_rs_seq = (uint32_t) atoi(e); | |
| + } | |
| if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { | |
| LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", | |
| __func__, cparams.n_rs_seq); | |
| void llama_context::sched_reserve() { | |
| synchronize(); | |
| + gcache_clear(); | |
| + | |
| const int64_t t_start_us = ggml_time_us(); | |
| const uint32_t n_seqs = cparams.n_seq_max; | |
| void llama_context::sched_reserve() { | |
| gf_res_prev.reset(new llm_graph_result(max_nodes)); | |
| gf_res_reserve.reset(new llm_graph_result(max_nodes)); | |
| + gcache_max_nodes = max_nodes; | |
| sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); | |
| bool llama_context::memory_update(bool optimize) { | |
| return false; | |
| } | |
| + gcache_select(0); // the memory update computes on the scheduler: keep it on the main one | |
| + | |
| { | |
| const auto mctx = memory->init_update(this, optimize); | |
| switch (mctx->get_status()) { | |
| bool llama_context::set_adapter_cvec( | |
| return res; | |
| } | |
| +// ---- diagnostic: LLAMA_UBATCH_PROF=N (see tools/kernel-lane/bonsai/patch_fork_ubprof.py) ---- | |
| +struct ubprof_row { long calls = 0, reused = 0; double build = 0, inputs = 0, compute = 0, gbuild = 0, galloc = 0; }; | |
| +static std::map<std::pair<const void *, int>, ubprof_row> g_ubprof; | |
| +static long g_ubprof_n = 0; | |
| +static int ubprof_every() { | |
| + static const int n = getenv("LLAMA_UBATCH_PROF") ? atoi(getenv("LLAMA_UBATCH_PROF")) : 0; | |
| + return n; | |
| +} | |
| +static void ubprof_dump() { | |
| + for (const auto & [k, r] : g_ubprof) { | |
| + const double nb = r.calls > r.reused ? (double) (r.calls - r.reused) : 1.0; | |
| + fprintf(stderr, "UBPROF ctx=%p w=%3d calls=%5ld reused=%5ld | build+alloc %.3f ms (per rebuild %.3f = graph %.3f + sched alloc %.3f) | inputs %.3f | compute(cpu) %.3f ms\n", | |
| + k.first, k.second, r.calls, r.reused, r.build / r.calls, r.calls > r.reused ? r.build / (r.calls - r.reused) : 0.0, | |
| + r.gbuild / nb, r.galloc / nb, r.inputs / r.calls, r.compute / r.calls); | |
| + } | |
| + fflush(stderr); | |
| + g_ubprof.clear(); | |
| +} | |
| + | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret) { | |
| + const bool ubprof = ubprof_every() > 0; | |
| + const int64_t ubprof_t0 = ubprof ? ggml_time_us() : 0; | |
| + bool ubprof_reused = false; | |
| + int64_t ubprof_t1 = 0, ubprof_t2 = 0; | |
| + double ubprof_gbuild_ms = 0.0, ubprof_galloc_ms = 0.0; | |
| if (mctx && !mctx->apply()) { | |
| LLAMA_LOG_ERROR("%s: failed to apply memory context\n", __func__); | |
| ret = GGML_STATUS_FAILED; | |
| return nullptr; | |
| } | |
| + { | |
| + static const uint32_t gcache_w = getenv("LLAMA_GRAPH_CACHE") ? (uint32_t) atoi(getenv("LLAMA_GRAPH_CACHE")) : 8; // 0 = off | |
| + if (gcache_w > 0 && gcache_max_nodes > 0) { | |
| + gcache_select(gtype != LLM_GRAPH_TYPE_ENCODER && ubatch.n_tokens <= gcache_w ? ubatch.n_tokens : 0); | |
| + } | |
| + } | |
| + | |
| auto * res = gf_res_prev.get(); | |
| auto * gf = res->get_gf(); | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll | |
| } | |
| n_reused++; | |
| + ubprof_reused = true; | |
| } else { | |
| res->reset(); | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll | |
| //const auto t_start_us = ggml_time_us(); | |
| + const int64_t ubprof_tb0 = ubprof ? ggml_time_us() : 0; | |
| gf = model.build_graph(gparams); | |
| + const int64_t ubprof_tb1 = ubprof ? ggml_time_us() : 0; | |
| //LLAMA_LOG_INFO("graph build time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll | |
| ret = GGML_STATUS_ALLOC_FAILED; | |
| return nullptr; | |
| } | |
| + if (ubprof) { | |
| + ubprof_gbuild_ms = (ubprof_tb1 - ubprof_tb0) / 1000.0; | |
| + ubprof_galloc_ms = (ggml_time_us() - ubprof_tb1) / 1000.0; | |
| + } | |
| } | |
| + if (ubprof) { | |
| + ubprof_t1 = ggml_time_us(); | |
| + } | |
| // set the input data for the input tensors | |
| { | |
| //const auto t_start_us = ggml_time_us(); | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll | |
| //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); | |
| } | |
| + if (ubprof) { | |
| + ubprof_t2 = ggml_time_us(); | |
| + } | |
| const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); | |
| llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll | |
| ret = GGML_STATUS_SUCCESS; | |
| + if (ubprof) { | |
| + const int64_t t3 = ggml_time_us(); | |
| + auto & r = g_ubprof[{ (const void *) this, (int) ubatch.n_tokens }]; | |
| + r.calls++; | |
| + r.reused += ubprof_reused ? 1 : 0; | |
| + r.build += (ubprof_t1 - ubprof_t0) / 1000.0; | |
| + r.inputs += (ubprof_t2 - ubprof_t1) / 1000.0; | |
| + r.compute += (t3 - ubprof_t2) / 1000.0; | |
| + r.gbuild += ubprof_gbuild_ms; | |
| + r.galloc += ubprof_galloc_ms; | |
| + if (++g_ubprof_n % ubprof_every() == 0) { | |
| + ubprof_dump(); | |
| + } | |
| + } | |
| + | |
| return res; | |
| } | |
| static void ubatch_prepare_reserve( | |
| } | |
| } | |
| +void llama_context::gcache_select(uint32_t key) { | |
| + if (key == gcache_cur) { | |
| + return; | |
| + } | |
| + // park the active scheduler + graph result under its key, then bring in (or create) the one for `key` | |
| + gcache[gcache_cur] = { std::move(sched), std::move(gf_res_prev) }; | |
| + auto it = gcache.find(key); | |
| + if (it != gcache.end()) { | |
| + sched = std::move(it->second.first); | |
| + gf_res_prev = std::move(it->second.second); | |
| + gcache.erase(it); | |
| + } else { | |
| + GGML_ASSERT(key != 0 && "the main scheduler is always active or parked"); | |
| + sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), gcache_max_nodes, false, cparams.op_offload)); | |
| + gf_res_prev.reset(new llm_graph_result(gcache_max_nodes)); | |
| + } | |
| + ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data); | |
| + gcache_cur = key; | |
| +} | |
| + | |
| +void llama_context::gcache_clear() { | |
| + if (gcache.empty() && gcache_cur == 0) { | |
| + return; | |
| + } | |
| + synchronize(); | |
| + gcache_select(0); | |
| + gcache.clear(); | |
| +} | |
| + | |
| ggml_cgraph * llama_context::graph_reserve( | |
| uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) { | |
| LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs); | |
| ggml_cgraph * llama_context::graph_reserve( | |
| LLAMA_LOG_DEBUG("%s: making n_tokens a multiple of n_seqs - n_tokens = %u, n_seqs = %u, n_outputs = %u\n", __func__, n_tokens, n_seqs, n_outputs); | |
| } | |
| + gcache_select(0); // worst-case reserves size the main scheduler only (patch_fork_graph_cache.py) | |
| + | |
| ggml_backend_sched_reset(sched.get()); | |
| // when the scheduler is reset, we cannot reuse the old graph, so we reset the previous graph result to prevent that | |
| int32_t llama_encode( | |
| return ret; | |
| } | |
| +static std::map<const llama_context *, std::pair<int64_t, int64_t>> g_decode_times; | |
| + | |
| +void llama_ext_last_decode_times(const struct llama_context * ctx, int64_t * t_entry, int64_t * t_return) { | |
| + const auto it = g_decode_times.find(ctx); | |
| + *t_entry = it == g_decode_times.end() ? 0 : it->second.first; | |
| + *t_return = it == g_decode_times.end() ? 0 : it->second.second; | |
| +} | |
| + | |
| int32_t llama_decode( | |
| llama_context * ctx, | |
| llama_batch batch) { | |
| + const int64_t t_entry = ggml_time_us(); | |
| const int ret = ctx->decode(batch); | |
| + g_decode_times[ctx] = { t_entry, ggml_time_us() }; | |
| if (ret != 0 && ret != 1) { | |
| LLAMA_LOG_ERROR("%s: failed to decode, ret = %d\n", __func__, ret); | |
| } | |
| diff --git a/src/llama-context.h b/src/llama-context.h | |
| index e2eb74d..503da37 100644 | |
| --- a/src/llama-context.h | |
| +++ b/src/llama-context.h | |
| private: | |
| llm_graph_result_ptr gf_res_prev; | |
| llm_graph_result_ptr gf_res_reserve; | |
| + // per-width graph cache (patch_fork_graph_cache.py, LLAMA_GRAPH_CACHE=W): parked scheduler + graph result per ubatch | |
| + // width <= W; key 0 = the main scheduler. Declared after the backends, so it is destroyed before them. | |
| + std::map<uint32_t, std::pair<ggml_backend_sched_ptr, llm_graph_result_ptr>> gcache; | |
| + uint32_t gcache_cur = 0; | |
| + size_t gcache_max_nodes = 0; | |
| + void gcache_select(uint32_t key); | |
| + void gcache_clear(); | |
| + | |
| // one-time Hadamard transform-coverage check on the first built graph | |
| bool hadamard_verified = false; | |
| diff --git a/src/llama-ext.h b/src/llama-ext.h | |
| index f1d9800..79f2a7e 100644 | |
| --- a/src/llama-ext.h | |
| +++ b/src/llama-ext.h | |
| #include <cstdint> | |
| #include <map> | |
| +// diagnostic: wall-clock (ggml_time_us) of the last llama_decode call on ctx -- entry and return | |
| +LLAMA_API void llama_ext_last_decode_times(const struct llama_context * ctx, int64_t * t_entry, int64_t * t_return); | |
| + | |
| // Reserve a new compute graph. It is valid until the next call to llama_graph_reserve. | |
| LLAMA_API struct ggml_cgraph * llama_graph_reserve( | |
| struct llama_context * ctx, | |
| diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp | |
| index f7c14f5..5f4858e 100644 | |
| --- a/src/llama-graph.cpp | |
| +++ b/src/llama-graph.cpp | |
| #include "llama-graph.h" | |
| +#include <typeinfo> | |
| + | |
| #include "llama-impl.h" | |
| #include "llama-model.h" | |
| #include "llama-batch.h" | |
| void llm_graph_result::set_outputs(const llm_graph_params & params) { | |
| } | |
| bool llm_graph_result::can_reuse(const llm_graph_params & params) { | |
| + static const bool trace = getenv("LLAMA_REUSE_TRACE") != nullptr; | |
| if (!this->params.allow_reuse(params)) { | |
| if (debug > 1) { | |
| LLAMA_LOG_DEBUG("%s: cannot reuse graph due to incompatible graph parameters\n", __func__); | |
| } | |
| + if (trace) { | |
| + fprintf(stderr, "REUSE-NO params: n_tokens %d -> %d, n_outputs %d -> %d, samplers %zu -> %zu, nextn_off %d -> %d\n", | |
| + (int) this->params.ubatch.n_tokens, (int) params.ubatch.n_tokens, (int) this->params.n_outputs, (int) params.n_outputs, | |
| + this->params.samplers.size(), params.samplers.size(), | |
| + (int) this->params.cparams.nextn_layer_offset, (int) params.cparams.nextn_layer_offset); | |
| + } | |
| return false; | |
| } | |
| bool llm_graph_result::can_reuse(const llm_graph_params & params) { | |
| for (auto & input : inputs) { | |
| const bool cur = input->can_reuse(params); | |
| + if (trace && !cur) { | |
| + fprintf(stderr, "REUSE-NO input %s (n_tokens %d)\n", typeid(*input).name(), (int) params.ubatch.n_tokens); | |
| + } | |
| if (debug > 1) { | |
| LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", "placeholder", cur); | |
| diff --git a/src/llama-model.cpp b/src/llama-model.cpp | |
| index f30522c..c7d7509 100644 | |
| --- a/src/llama-model.cpp | |
| +++ b/src/llama-model.cpp | |
| void llama_model_base::load_hparams(llama_model_loader & ml) { | |
| "ffn_gate_exps", "ffn_up_exps", "ffn_down_exps", "ffn_gate_up_exps", | |
| "ffn_gate_shexp", "ffn_up_shexp", "ffn_down_shexp", | |
| "ssm_out", | |
| + "nextn.shared_head_head", // MTP draft head (full or pruned vocabulary): built through build_lora_mm | |
| }; | |
| if (name == "output.weight") { | |
| return true; // the output head is built through build_lora_mm in every arch | |
| diff --git a/src/llama-model.h b/src/llama-model.h | |
| index ff1bc9b..195f219 100644 | |
| --- a/src/llama-model.h | |
| +++ b/src/llama-model.h | |
| struct llama_layer_nextn { | |
| struct ggml_tensor * shared_head_head_s = nullptr; | |
| struct ggml_tensor * shared_head_head_in_s = nullptr; | |
| struct ggml_tensor * shared_head_norm = nullptr; | |
| + struct ggml_tensor * draft_vocab = nullptr; // I32 [n_vocab]: row in a pruned shared_head_head, or n_rows = not drafted | |
| }; | |
| struct llama_layer_switch_lora { | |
| diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp | |
| index 8e0944b..44dafa6 100644 | |
| --- a/src/models/qwen35.cpp | |
| +++ b/src/models/qwen35.cpp | |
| void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { | |
| layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags); | |
| layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags); | |
| layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); | |
| - layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); | |
| + // optional pruned draft vocabulary: when draft_vocab is present, shared_head_head holds only the rows of the | |
| + // n_small most frequent tokens (its real width is read from the file) and the draft graph scatters back | |
| + int64_t n_head_rows = n_vocab; | |
| + { | |
| + const std::string vname = tn(LLM_TENSOR_NEXTN_DRAFT_VOCAB, "weight", il); | |
| + const std::string hname = tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il); | |
| + if (ml.get_tensor_meta(vname.c_str()) != nullptr) { | |
| + if (const ggml_tensor * hm = ml.get_tensor_meta(hname.c_str())) { | |
| + n_head_rows = hm->ne[1]; | |
| + } | |
| + } | |
| + } | |
| + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_head_rows }, mtp_flags|TENSOR_NOT_REQUIRED); | |
| + layer.nextn.draft_vocab = create_tensor(tn(LLM_TENSOR_NEXTN_DRAFT_VOCAB, "weight", il), { n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED); | |
| layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED); | |
| }; | |
| llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr | |
| cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); | |
| cb(cur, "h_nextn", -1); | |
| - res->t_h_nextn = cur; | |
| + res->t_h_nextn = cur; // every row: unmasked contexts read one row per token | |
| cur = ggml_get_rows(ctx0, cur, inp_out_ids); | |
| cb(cur, "mtp_shared_head_norm", -1); | |
| + if (cparams.embeddings_nextn_masked) { | |
| + // masked contexts (the MTP draft context) copy n_outputs rows in output order (llama_context::decode): | |
| + // hand them the selected rows, or a batch whose output is not its first row -- catch-up rows followed | |
| + // by a draft row -- reads the hidden state of the batch's first row | |
| + res->t_h_nextn = cur; | |
| + } | |
| ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; | |
| ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; | |
| GGML_ASSERT(head_w && "QWEN35 MTP: missing LM head (nextn.shared_head_head or model.output)"); | |
| cur = build_lora_mm(head_w, cur, head_s); | |
| + if (layer.nextn.draft_vocab && layer.nextn.shared_head_head) { | |
| + // pruned draft vocabulary: cur is [n_small, n_out]; scatter to [n_vocab, n_out] with -1e30 outside the set | |
| + const int64_t n_out = cur->ne[1]; | |
| + ggml_tensor * pad = ggml_view_2d(ctx0, cur, 1, n_out, cur->nb[1], 0); | |
| + pad = ggml_scale_bias(ctx0, ggml_cont(ctx0, pad), 0.0f, -1e30f); // [1, n_out] | |
| + cur = ggml_concat(ctx0, cur, pad, 0); // [n_small + 1, n_out] | |
| + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_out, n_small + 1] | |
| + cur = ggml_get_rows(ctx0, cur, layer.nextn.draft_vocab); // [n_out, n_vocab] | |
| + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_vocab, n_out] | |
| + } | |
| cb(cur, "result_output", -1); | |
| res->t_logits = cur; | |
| diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp | |
| index b28bc29..8f6427d 100644 | |
| --- a/tests/test-backend-ops.cpp | |
| +++ b/tests/test-backend-ops.cpp | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() { | |
| } | |
| } | |
| + // Bonsai 2 (PQ2_0) decode/verify shapes on sm_6x: K a multiple of 1024 (8 K-blocks) reaches the shuffle kernel, | |
| + // K=17408 at n>=3 (and 6144 at n=8) its K-chunked activation staging, M=50 the row tail, M=4000 the persistent | |
| + // row-group loop (more row groups than resident blocks). | |
| + for (int n = 1; n <= 8; ++n) { | |
| + for (int k : {1024, 5120, 6144, 17408}) { | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 50, n, k, {1, 1}, {1, 1})); | |
| + } | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 4000, n, 5120, {1, 1}, {1, 1})); | |
| + } | |
| + | |
| + // Q4_0 row-lane kernel (sm_6x): K=256 (2 units, idle lanes), 2560, 10240 (chunked/split-K at n>=3), 15360 (n>=2); | |
| + // M=50 the row tail, M=3000 the persistent row-group loop | |
| + for (int n = 1; n <= 8; ++n) { | |
| + for (int k : {256, 2560, 10240, 15360}) { | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 50, n, k, {1, 1}, {1, 1})); | |
| + } | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 3000, n, 2560, {1, 1}, {1, 1})); | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 3000, n, 15360, {1, 1}, {1, 1})); | |
| + } | |
| + | |
| test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); | |
| test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); | |
| test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() { | |
| test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2048, 2, 1, 3}, k)); | |
| test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2049, 2, 1, 3}, k)); | |
| } | |
| + // the MTP draft sampler's shape: top_k(10) over the draft-vocab-scattered logits (Bonsai 2 / Qwen3.8 vocab 248320) | |
| + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {248320, 1, 1, 1}, 10)); | |
| + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {248320, 4, 1, 1}, 10)); | |
| + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {248320, 1, 1, 1}, 10, true)); | |
| // exhaustive top_k tests | |
| //for (int i = 1; i < 9999; ++i) { | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() { | |
| } | |
| } | |
| + // GQA 6 at head size 256 (Bonsai 2 27B: 24 q heads / 4 KV heads) -- the ncols2 = 6 tile path | |
| + // (patch_fork_fa_gqa6.py). nb 1..4 is plain decode and every speculative verify width; nb 8 takes the | |
| + // fallback. kv covers multiples and non-multiples of the 256-token FA stride. | |
| + for (int kv : { 512, 1024, 4096, 8192, 18176, 18000 }) { | |
| + for (int nb : { 1, 2, 3, 4, 8 }) { | |
| + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); | |
| + } | |
| + } | |
| + | |
| // mixed quant and Q1_0 test cases | |
| test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); | |
| test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() { | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() { | |
| std::vector<std::unique_ptr<test_case>> test_cases; | |
| + // Bonsai 2 native MTP head (Q4_0): eh_proj 5120x10240, q 12288x5120, k/v 1024x5120, o 5120x6144, ffn 17408x5120 / 5120x17408 | |
| + for (int n : {1, 2, 3, 4}) { | |
| + for (auto mk : std::vector<std::pair<int, int>>{{5120, 10240}, {12288, 5120}, {1024, 5120}, {5120, 6144}, | |
| + {17408, 5120}, {5120, 17408}}) { | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, mk.first, n, mk.second, {1, 1}, {1, 1})); | |
| + } | |
| + } | |
| + | |
| + // Gemma 4 E4B / 12B (Q4_0) decode/verify matmuls on sm_6x: rows x K | |
| + for (int n : {1, 2, 4, 5}) { | |
| + for (auto mk : std::vector<std::pair<int, int>>{{10240, 2560}, {2560, 10240}, {2048, 2560}, {2560, 2048}, {512, 2560}, | |
| + {256, 2560}, {262144, 2560}, {15360, 3840}, {3840, 15360}, {4096, 3840}, | |
| + {3840, 4096}, {2048, 3840}}) { | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, mk.first, n, mk.second, {1, 1}, {1, 1})); | |
| + } | |
| + } | |
| + | |
| + // Bonsai 2 27B (PQ2_0) decode/verify matmuls: rows x K of every ternary weight, 1-8 tokens of one sequence | |
| + for (int n : {1, 2, 3, 4, 5, 8}) { | |
| + for (auto mk : std::vector<std::pair<int, int>>{{17408, 5120}, {5120, 17408}, {10240, 5120}, {6144, 5120}, | |
| + {5120, 6144}, {12288, 5120}, {1024, 5120}, {248320, 5120}}) { | |
| + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, mk.first, n, mk.second, {1, 1}, {1, 1})); | |
| + } | |
| + } | |
| + | |
| + | |
| // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands | |
| // note: same bytes either way, so a backend that indexes them differently shows it here | |
| for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { | |
| static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() { | |
| } | |
| } | |
| + // Bonsai 2 27B attention shape (GQA 6, head size 256) for tuning the ncols2 = 6 tile config | |
| + for (int kv : { 8192, 18176, }) { | |
| + for (int nb : { 1, 4, }) { | |
| + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); | |
| + } | |
| + } | |
| + | |
| for (int col : {8192, 16384, 32768, 65536, 131072, 262144, 524288}) { | |
| for (int rows : {1, 4, 16}){ | |
| test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {col, rows, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); | |
| diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp | |
| index a9edbd7..091dd10 100644 | |
| --- a/tools/server/server-context.cpp | |
| +++ b/tools/server/server-context.cpp | |
| static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch | |
| // server_context_impl (private implementation) | |
| // | |
| + | |
| +// ---- diagnostic: SRV_PHASE_PROF=N prints every N draft cycles the mean ms of the server phases between drafting and decode | |
| +struct srv_phase_prof { | |
| + int every = -1; | |
| + long n = 0; | |
| + double acc[8] = {0}; | |
| + int64_t t_last = 0; | |
| + bool on() { | |
| + if (every < 0) { every = getenv("SRV_PHASE_PROF") ? atoi(getenv("SRV_PHASE_PROF")) : 0; } | |
| + return every > 0; | |
| + } | |
| + void mark_start() { t_last = ggml_time_us(); } | |
| + void mark(int k) { const int64_t t = ggml_time_us(); acc[k] += (t - t_last) / 1000.0; t_last = t; } | |
| + void cycle() { | |
| + if (++n >= every) { | |
| + fprintf(stderr, "SRVPHASE n=%ld ms: ckpt+dft_seq_rm %.3f | last_sampled %.3f | prompt_batching %.3f | render %.3f | lora+embd+view %.3f | decode-pre %.3f\n", | |
| + n, acc[0] / n, acc[1] / n, acc[2] / n, acc[3] / n, acc[4] / n, acc[5] / n); | |
| + fflush(stderr); | |
| + n = 0; for (auto & a : acc) a = 0; | |
| + } | |
| + } | |
| +}; | |
| +static srv_phase_prof g_srv_phase; | |
| + | |
| struct server_context_impl { | |
| friend struct server_context; | |
| private: | |
| try { | |
| scoped_timer t(t_pre_decode, n_pre_decode); | |
| pre_decode(); | |
| + if (g_srv_phase.t_last > 0 && g_srv_phase.every > 0) { | |
| + g_srv_phase.mark(2); | |
| + } | |
| batch.render(); | |
| + if (g_srv_phase.t_last > 0 && g_srv_phase.every > 0) { | |
| + g_srv_phase.mark(3); | |
| + } | |
| } catch (const std::exception & e) { | |
| SRV_ERR("pre_decode() failed: %s\n", e.what()); | |
| abort_all_slots("pre_decode() failed: " + std::string(e.what())); | |
| private: | |
| // TODO @ngxson : maybe handle n_batch == 1 here instead of inside decode() | |
| batch_view = batch.get_view(off, n_tokens); | |
| + if (g_srv_phase.t_last > 0 && g_srv_phase.every > 0) { | |
| + g_srv_phase.mark(4); | |
| + } | |
| bool ok = decode(n_batch, off, batch_view); | |
| #ifdef DEBUG_TIMINGS | |
| llama_synchronize(ctx_tgt); | |
| private: | |
| common_speculative_draft(spec.get()); | |
| }); | |
| } | |
| + const bool srv_prof = g_srv_phase.on() && !drafting.empty(); | |
| + if (srv_prof) { | |
| + g_srv_phase.mark_start(); | |
| + } | |
| // make checkpoints if needed | |
| iterate(drafting, [&](server_slot & slot) { | |
| private: | |
| } | |
| }); | |
| + if (srv_prof) { | |
| + g_srv_phase.mark(0); | |
| + } | |
| // update the batch with the sampled/drafted tokens | |
| iterate(generating, [&](server_slot & slot) { | |
| slot.handle_last_sampled_token(batch); | |
| }); | |
| + if (srv_prof) { | |
| + g_srv_phase.mark(1); | |
| + } | |
| // process in chunks of params.n_batch | |
| int32_t n_batch = llama_n_batch(ctx_tgt); | |
| private: | |
| // note: the sync is done here too, so that the wait is also covered by the yield | |
| int ret = 0; | |
| queue_tasks.yield_to_queue([&]() { | |
| + if (g_srv_phase.t_last > 0 && g_srv_phase.every > 0) { | |
| + g_srv_phase.mark(5); | |
| + g_srv_phase.t_last = 0; | |
| + g_srv_phase.cycle(); | |
| + } | |
| ret = llama_decode(ctx_tgt, batch_view); | |
| if (ret == 0 && has_output) { | |
| llama_synchronize(ctx_tgt); | |
| diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp | |
| index 78169e9..188c65b 100644 | |
| --- a/tools/server/server-queue.cpp | |
| +++ b/tools/server/server-queue.cpp | |
| void server_queue::worker_stop() { | |
| worker.thread.join(); | |
| } | |
| +// diagnostic: SRV_YIELD_PROF=N prints every N yields: mean ms spent entering (lock + notify), in the work, and exiting | |
| +// (waiting for the worker thread to hand the queue back) | |
| +static void srv_yield_prof(double enter_ms, double work_ms, double exit_ms) { | |
| + static const int every = getenv("SRV_YIELD_PROF") ? atoi(getenv("SRV_YIELD_PROF")) : 0; | |
| + if (every <= 0) { | |
| + return; | |
| + } | |
| + static double se = 0, sw = 0, sx = 0, mx = 0; | |
| + static long n = 0; | |
| + se += enter_ms; sw += work_ms; sx += exit_ms; mx = std::max(mx, exit_ms); n++; | |
| + if (n >= every) { | |
| + fprintf(stderr, "YIELDPROF n=%ld enter %.3f ms work %.3f ms exit %.3f ms (max %.3f)\n", n, se / n, sw / n, sx / n, mx); | |
| + fflush(stderr); | |
| + se = sw = sx = mx = 0; n = 0; | |
| + } | |
| +} | |
| + | |
| void server_queue::yield_to_queue(std::function<void()> && work) { | |
| GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running"); | |
| + const int64_t yp_t0 = ggml_time_us(); | |
| QUE_DBG("%s", "yielding to queue\n"); | |
| void server_queue::yield_to_queue(std::function<void()> && work) { | |
| } | |
| worker.cv.notify_one(); | |
| + const int64_t yp_t1 = ggml_time_us(); | |
| // run the work on the current thread, so that all ggml compute stays on the same thread | |
| std::exception_ptr exception; | |
| try { | |
| void server_queue::yield_to_queue(std::function<void()> && work) { | |
| } catch (...) { | |
| exception = std::current_exception(); | |
| } | |
| + const int64_t yp_t2 = ggml_time_us(); | |
| { | |
| std::unique_lock<std::mutex> lock(mutex_tasks); | |
| void server_queue::yield_to_queue(std::function<void()> && work) { | |
| } | |
| QUE_DBG("%s", "done yielding to queue\n"); | |
| + { | |
| + const int64_t yp_t3 = ggml_time_us(); | |
| + srv_yield_prof((yp_t1 - yp_t0) / 1000.0, (yp_t2 - yp_t1) / 1000.0, (yp_t3 - yp_t2) / 1000.0); | |
| + } | |
| // note: rethrow only after the declined tasks are back in the queue, so they are not lost | |
| if (exception) { | |