Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e5eb5edb58 | ||
|
|
4442815c02 | ||
|
|
891760faba | ||
|
|
1bc7c581d8 | ||
|
|
3972da9cc9 | ||
|
|
42e52045dc | ||
|
|
1b8abd1c23 | ||
|
|
c3c913ba79 | ||
|
|
c440f0b925 | ||
|
|
02cff16681 | ||
|
|
5a0c899424 |
@@ -12,6 +12,77 @@
|
||||
|
||||
LLM inference in C/C++
|
||||
|
||||
## ⚡ This fork — Fable's MoE-offload prefill optimizations
|
||||
|
||||
Two **opt-in** optimizations for large MoE models whose experts are offloaded to system RAM
|
||||
(`--n-cpu-moe`), found and implemented by Fable. Both are **off by default**, toggled via
|
||||
environment variables, and produce **token-identical** output to mainline.
|
||||
|
||||
| Env var | What it does |
|
||||
| --- | --- |
|
||||
| `GGML_CUDA_REGISTER_HOST=1` | Page-locks (pins) the mmap'd CPU expert weights so host->device copies go straight over DMA instead of through the driver's hidden bounce buffer (~6-7 -> ~20 GB/s). Works on CUDA and Vulkan (also honored as `GGML_VK_REGISTER_HOST`). Note: it is a presence check, so `=0` still enables it. |
|
||||
| `GGML_SCHED_PREFETCH_EXPERTS=1` | Prefetches each layer's experts on a second stream, so the weight uploads overlap compute instead of stalling the GPU. **CUDA only** - on the Vulkan backend the second backend instance shares one device queue, giving no overlap, so this regresses (see Vulkan note below). Leave it off on Vulkan. |
|
||||
|
||||
### Benchmark
|
||||
|
||||
Measured on an **RTX 3060 12GB** with **Qwen3.6-35B-A3B** (`--n-cpu-moe 26`), prompt-processing at 2048 (`MODEL` = path to your `.gguf`):
|
||||
|
||||
```bash
|
||||
# baseline (patches off):
|
||||
./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
|
||||
|
||||
# patched (both optimizations on):
|
||||
GGML_CUDA_REGISTER_HOST=1 GGML_SCHED_PREFETCH_EXPERTS=1 \
|
||||
./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
|
||||
```
|
||||
|
||||
Result: **~1143 → ~1880 t/s** prefill (**+64%**) — same GPU, same settings, token-identical.
|
||||
|
||||
Branches: [`fable5/host-register`](https://github.com/thecodacus/llama.cpp/tree/fable5/host-register) (pinning only) · [`fable5/prefetch-experts`](https://github.com/thecodacus/llama.cpp/tree/fable5/prefetch-experts) (both — this branch).
|
||||
|
||||
### Vulkan (older AMD, e.g. RX 580 / Polaris)
|
||||
|
||||
On the Vulkan backend the CUDA-oriented flags above behave differently, and this fork adds a
|
||||
Polaris-specific flash-attention fix. Findings on an **RX 580 8GB** (Polaris / GCN, PCIe 3.0 x16,
|
||||
no fp16, no matrix cores) with **Qwen3.5-35B-A3B Q4_K_M**, `-b 2048 -ub 2048`:
|
||||
|
||||
- **Flash-attention `mask_opt` is now enabled for GCN large head sizes (this fork's own change).**
|
||||
Upstream disables it on GCN; it is a **lossless** win in high-context prefill - it skips
|
||||
fully-masked causal blocks and the per-block mask add on fully-visible ones, which is real work on
|
||||
a card whose attention is compute-bound (no matrix cores). Auto-on, no flag. On Qwen3.5-35B
|
||||
(head_dim 256): pp2048 **+8% @ 16k, +12% @ 32k**, growing with depth; perplexity bit-identical.
|
||||
- **For a long-running server, load with `--no-mmap`, not pinning.** `GGML_CUDA_REGISTER_HOST=1`
|
||||
(pinning) gives ~+17% in an isolated `llama-bench` run, but in a server the RADV host-pointer
|
||||
import fails and the fallback pre-stage buffer allocation fails for large / co-resident models, so
|
||||
it silently reverts to slow staging (and can trip warnings/OOM). `--no-mmap` (weights in RAM) is
|
||||
both faster and clean there. Pinning is still fine for one-off `llama-bench` numbers.
|
||||
- **`GGML_SCHED_PREFETCH_EXPERTS=1` regresses - do not use it** on Vulkan (its second backend shares
|
||||
one device queue, so uploads never overlap compute).
|
||||
- **`-b 2048 -ub 2048` is the biggest prefill lever** (the default `-ub 512` roughly halves pp).
|
||||
- **Tune `--n-cpu-moe` to context length.** Keep some expert layers resident in spare VRAM for short
|
||||
prompts (e.g. `ncmoe 28` on the 35B, ~+5% over all-host); at long context the KV cache needs that
|
||||
VRAM, so raise it (`ncmoe 40`, all experts on host). Keep flash attention on (`-fa 1`).
|
||||
|
||||
Prefill throughput (isolated `llama-bench`, pinned unless noted):
|
||||
|
||||
| Config | pp2048 (t/s) |
|
||||
| --- | ---: |
|
||||
| baseline, unpinned, `ncmoe 40` | ~252 |
|
||||
| pinned, `ncmoe 40` | ~294 |
|
||||
| pinned, `ncmoe 28` | ~308 |
|
||||
| server default (`--no-mmap`, `ncmoe 40`) | ~285 |
|
||||
| + `mask_opt`, @ 32k context | **+12%** |
|
||||
|
||||
**Recommended RX 580 / Polaris serving command** (per model):
|
||||
|
||||
```bash
|
||||
llama-server -hf <repo>:<quant> --no-mmap -ngl 99 --n-cpu-moe 40 -b 2048 -ub 2048 -fa 1
|
||||
```
|
||||
|
||||
Lower `--n-cpu-moe` (e.g. 28) if the model plus your context budget leave spare VRAM; keep it high
|
||||
for long-context / agentic use. At long context the bottleneck is attention compute (GPU-bound), so
|
||||
`mask_opt` (above) is where the remaining prefill wins come from, not the MoE-transfer path.
|
||||
|
||||
## Recent API changes
|
||||
|
||||
- [Changelog for `libllama` API](https://github.com/ggml-org/llama.cpp/issues/9289)
|
||||
|
||||
@@ -761,6 +761,10 @@ static bool ggml_is_view_op(enum ggml_op op) {
|
||||
#define GGML_SCHED_MAX_COPIES 4
|
||||
#endif
|
||||
|
||||
#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS
|
||||
#define GGML_SCHED_MAX_PREFETCH_SLOTS 8
|
||||
#endif
|
||||
|
||||
struct ggml_backend_sched_split {
|
||||
int backend_id;
|
||||
int i_start;
|
||||
@@ -818,6 +822,19 @@ struct ggml_backend_sched {
|
||||
|
||||
bool op_offload;
|
||||
|
||||
// full-tensor prefetch of offloaded MUL_MAT_ID weights (GGML_SCHED_PREFETCH_EXPERTS)
|
||||
// with a large batch virtually every expert is used, so the routing ids are not worth
|
||||
// waiting for; uploads run through a second backend instance on the same device so
|
||||
// they overlap compute, alternating between two staging slots
|
||||
bool prefetch_experts;
|
||||
ggml_backend_t prefetch_backend;
|
||||
int prefetch_n_slots;
|
||||
ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS];
|
||||
ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS];
|
||||
ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS];
|
||||
bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
|
||||
int prefetch_cur;
|
||||
|
||||
int debug;
|
||||
|
||||
// used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC]
|
||||
@@ -1538,6 +1555,94 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
|
||||
return true;
|
||||
}
|
||||
|
||||
static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) {
|
||||
sched->prefetch_experts = false;
|
||||
if (sched->prefetch_backend) {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
ggml_backend_synchronize(sched->prefetch_backend);
|
||||
}
|
||||
for (int i = 0; i < sched->prefetch_n_slots; i++) {
|
||||
ggml_backend_buffer_free(sched->prefetch_slots[i]);
|
||||
sched->prefetch_slots[i] = NULL;
|
||||
sched->prefetch_used[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// slots are sized once for the largest offloaded expert tensor in the current graph so
|
||||
// that they never need to grow mid-eval
|
||||
static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) {
|
||||
size_t max_size = 0;
|
||||
for (int split_id = 0; split_id < sched->n_splits; split_id++) {
|
||||
struct ggml_backend_sched_split * split = &sched->splits[split_id];
|
||||
if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) {
|
||||
continue;
|
||||
}
|
||||
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
|
||||
const ggml_tensor * input = split->inputs[input_id];
|
||||
if (input->buffer &&
|
||||
ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
|
||||
ggml_backend_buffer_is_host(input->buffer)) {
|
||||
max_size = std::max(max_size, ggml_nbytes(input));
|
||||
}
|
||||
}
|
||||
}
|
||||
return max_size;
|
||||
}
|
||||
|
||||
static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) {
|
||||
if (sched->prefetch_backend == NULL) {
|
||||
ggml_backend_dev_t dev = split_backend->device;
|
||||
ggml_backend_dev_props props;
|
||||
ggml_backend_dev_get_props(dev, &props);
|
||||
if (!props.caps.async || !props.caps.events) {
|
||||
sched->prefetch_experts = false;
|
||||
return false;
|
||||
}
|
||||
sched->prefetch_backend = ggml_backend_dev_init(dev, NULL);
|
||||
if (sched->prefetch_backend == NULL) {
|
||||
sched->prefetch_experts = false;
|
||||
return false;
|
||||
}
|
||||
for (int i = 0; i < sched->prefetch_n_slots; i++) {
|
||||
sched->prefetch_ready[i] = ggml_backend_event_new(dev);
|
||||
sched->prefetch_free[i] = ggml_backend_event_new(dev);
|
||||
if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) {
|
||||
sched->prefetch_experts = false;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size = std::max(size, ggml_backend_sched_prefetch_max_size(sched));
|
||||
|
||||
ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend);
|
||||
for (int i = 0; i < sched->prefetch_n_slots; i++) {
|
||||
if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) {
|
||||
// allocate before freeing so a failure leaves the old slot intact
|
||||
ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size);
|
||||
if (new_buf == NULL) {
|
||||
// overlap needs at least 2 slots, otherwise run with what fits
|
||||
if (i >= 2 && sched->prefetch_slots[0] != NULL &&
|
||||
ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) {
|
||||
sched->prefetch_n_slots = i;
|
||||
sched->prefetch_cur = 0;
|
||||
return true;
|
||||
}
|
||||
ggml_backend_sched_prefetch_disable(sched, split_backend);
|
||||
return false;
|
||||
}
|
||||
if (sched->prefetch_slots[i] != NULL) {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
ggml_backend_synchronize(sched->prefetch_backend);
|
||||
ggml_backend_buffer_free(sched->prefetch_slots[i]);
|
||||
}
|
||||
sched->prefetch_slots[i] = new_buf;
|
||||
sched->prefetch_used[i] = false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
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;
|
||||
@@ -1550,6 +1655,10 @@ 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];
|
||||
int split_prefetch_slot = -1;
|
||||
ggml_tensor * prefetch_input_cpy = NULL;
|
||||
ggml_backend_buffer_t prefetch_saved_buffer = NULL;
|
||||
void * prefetch_saved_data = NULL;
|
||||
|
||||
// copy the input tensors to the split backend
|
||||
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
|
||||
@@ -1566,6 +1675,41 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
|
||||
}
|
||||
ggml_backend_tensor_copy(input, input_cpy);
|
||||
} else {
|
||||
// with a large batch virtually every expert is used, so instead of waiting
|
||||
// for the routing ids, upload the full tensor through the prefetch backend
|
||||
// and let the copy overlap compute of the previous split
|
||||
if (sched->prefetch_experts && !sched->callback_eval && split_prefetch_slot == -1 && split->graph.n_nodes > 0) {
|
||||
ggml_tensor * node = split->graph.nodes[0];
|
||||
if (ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
|
||||
ggml_backend_buffer_is_host(input->buffer) &&
|
||||
node->op == GGML_OP_MUL_MAT_ID && node->src[0] == input_cpy) {
|
||||
const ggml_tensor * ids = node->src[2];
|
||||
const int64_t n_expert = input->ne[2];
|
||||
if (ids->ne[0]*ids->ne[1] >= 2*n_expert &&
|
||||
ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) {
|
||||
const int slot = sched->prefetch_cur;
|
||||
sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots;
|
||||
// wait for the previous user of this slot to finish computing
|
||||
if (sched->prefetch_used[slot]) {
|
||||
ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]);
|
||||
}
|
||||
// point the staging copy at the slot only for the duration of
|
||||
// this split, so a fallback to the regular path on a later
|
||||
// eval can never see a dangling slot pointer
|
||||
prefetch_input_cpy = input_cpy;
|
||||
prefetch_saved_buffer = input_cpy->buffer;
|
||||
prefetch_saved_data = input_cpy->data;
|
||||
input_cpy->buffer = sched->prefetch_slots[slot];
|
||||
input_cpy->data = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]);
|
||||
ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input));
|
||||
ggml_backend_event_record(sched->prefetch_ready[slot], sched->prefetch_backend);
|
||||
ggml_backend_event_wait(split_backend, sched->prefetch_ready[slot]);
|
||||
split_prefetch_slot = slot;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wait for the split backend to finish using the input before overwriting it
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
|
||||
@@ -1676,6 +1820,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
|
||||
|
||||
if (!sched->callback_eval) {
|
||||
enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
|
||||
if (split_prefetch_slot != -1) {
|
||||
// the kernels have captured the slot address at launch, safe to restore
|
||||
ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend);
|
||||
sched->prefetch_used[split_prefetch_slot] = true;
|
||||
prefetch_input_cpy->buffer = prefetch_saved_buffer;
|
||||
prefetch_input_cpy->data = prefetch_saved_data;
|
||||
}
|
||||
if (ec != GGML_STATUS_SUCCESS) {
|
||||
return ec;
|
||||
}
|
||||
@@ -1788,6 +1939,15 @@ ggml_backend_sched_t ggml_backend_sched_new(
|
||||
sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
|
||||
sched->op_offload = op_offload;
|
||||
|
||||
// GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it
|
||||
// directly; more slots let uploads run further ahead of compute at the cost of one
|
||||
// max-sized expert tensor of device memory per slot
|
||||
const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS");
|
||||
const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0;
|
||||
sched->prefetch_experts = op_offload && prefetch_n_slots > 0;
|
||||
// default of 3 covers the gate/up/down expert tensors of one MoE layer
|
||||
sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS);
|
||||
|
||||
ggml_backend_sched_reset(sched);
|
||||
|
||||
return sched;
|
||||
@@ -1802,6 +1962,16 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
|
||||
ggml_backend_event_free(sched->events[b][c]);
|
||||
}
|
||||
}
|
||||
if (sched->prefetch_backend) {
|
||||
ggml_backend_synchronize(sched->prefetch_backend);
|
||||
// the slot count may have been reduced after a failed allocation, free everything
|
||||
for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) {
|
||||
ggml_backend_event_free(sched->prefetch_ready[i]);
|
||||
ggml_backend_event_free(sched->prefetch_free[i]);
|
||||
ggml_backend_buffer_free(sched->prefetch_slots[i]);
|
||||
}
|
||||
ggml_backend_free(sched->prefetch_backend);
|
||||
}
|
||||
ggml_gallocr_free(sched->galloc);
|
||||
ggml_free(sched->ctx);
|
||||
ggml_hash_set_free(&sched->hash_set);
|
||||
@@ -1906,6 +2076,9 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
|
||||
for (int i = 0; i < sched->n_backends; i++) {
|
||||
ggml_backend_synchronize(sched->backends[i]);
|
||||
}
|
||||
if (sched->prefetch_backend) {
|
||||
ggml_backend_synchronize(sched->prefetch_backend);
|
||||
}
|
||||
if (!sched->is_alloc) {
|
||||
// if the graph is not already allocated, always use copy 0 after a synchronization
|
||||
// this ensures that during generation the same copy is used every time,
|
||||
|
||||
@@ -3283,6 +3283,7 @@ static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, const std
|
||||
import_info.setPNext(&mem_flags_info);
|
||||
buf->device_memory = device->device.allocateMemory({ size, memory_type_idx, &import_info });
|
||||
} catch (const vk::SystemError& e) {
|
||||
GGML_LOG_WARN("ggml_vulkan: host pointer memory import failed (%s)\n", e.what());
|
||||
}
|
||||
} else {
|
||||
for (auto it = req_flags_list.begin(); it != req_flags_list.end(); it++) {
|
||||
@@ -8128,25 +8129,32 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
|
||||
ggml_vk_host_get(dst->device, src, buf, buf_offset);
|
||||
|
||||
if (buf != nullptr) {
|
||||
// Memory is pinned, use as staging buffer
|
||||
std::vector<vk::BufferCopy> slices(1);
|
||||
if (width == spitch && width == dpitch) {
|
||||
// Only do single write if stride is equal
|
||||
slices[0].srcOffset = buf_offset;
|
||||
slices[0].dstOffset = offset;
|
||||
slices[0].size = width * height;
|
||||
} else {
|
||||
slices.resize(height);
|
||||
for (size_t i = 0; i < height; i++) {
|
||||
slices[i].srcOffset = buf_offset + i * spitch;
|
||||
slices[i].dstOffset = offset + i * dpitch;
|
||||
slices[i].size = width;
|
||||
// extent of the read in pinned source memory; guard against tensors that
|
||||
// straddle a pinned-chunk boundary (they fall back to staging below)
|
||||
size_t src_extent = (width == spitch) ? (size_t) width * height
|
||||
: (height > 0 ? (height - 1) * spitch + width : 0);
|
||||
if (buf_offset + src_extent <= buf->size) {
|
||||
// Memory is pinned, use as staging buffer
|
||||
std::vector<vk::BufferCopy> slices(1);
|
||||
if (width == spitch && width == dpitch) {
|
||||
// Only do single write if stride is equal
|
||||
slices[0].srcOffset = buf_offset;
|
||||
slices[0].dstOffset = offset;
|
||||
slices[0].size = width * height;
|
||||
} else {
|
||||
slices.resize(height);
|
||||
for (size_t i = 0; i < height; i++) {
|
||||
slices[i].srcOffset = buf_offset + i * spitch;
|
||||
slices[i].dstOffset = offset + i * dpitch;
|
||||
slices[i].size = width;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ggml_vk_sync_buffers(nullptr, subctx);
|
||||
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
|
||||
return true;
|
||||
ggml_vk_sync_buffers(nullptr, subctx);
|
||||
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
|
||||
return true;
|
||||
}
|
||||
// straddles a chunk boundary: fall through to staging
|
||||
}
|
||||
VK_LOG_DEBUG("STAGING");
|
||||
|
||||
@@ -10584,8 +10592,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
|
||||
}
|
||||
|
||||
// Only use mask opt when the mask is fairly large. This hasn't been tuned extensively.
|
||||
// GCN with large head size (>= 256) benefits in high-context prefill (skipping the
|
||||
// per-block mask add on fully-visible blocks dominates), so enable it there too.
|
||||
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
|
||||
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
|
||||
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK >= 256 || HSV >= 256);
|
||||
vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
|
||||
mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type);
|
||||
|
||||
@@ -12512,9 +12522,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
|
||||
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f });
|
||||
}
|
||||
|
||||
// Fast path for concat along dim 0 where one source is stored "transposed"
|
||||
// (nb[1] == type_size, dim1 innermost) and the other source + dst are
|
||||
// contiguous along dim 0. The generic concat shader reads the transposed
|
||||
// source with a catastrophic uncoalesced stride; here we instead copy the
|
||||
// contiguous source with copy.comp and transpose the other source into the
|
||||
// matching dst sub-region with the tiled copy_transpose shader (shared-memory
|
||||
// transpose, coalesced read+write). Mirrors the precedent in
|
||||
// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants.
|
||||
static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx,
|
||||
const ggml_tensor * ctg, const ggml_tensor * trp,
|
||||
ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) {
|
||||
const uint32_t ts = ggml_type_size(dst->type);
|
||||
|
||||
vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32
|
||||
: ctx->device->pipeline_cpy_f16_f16;
|
||||
vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32
|
||||
: ctx->device->pipeline_cpy_transpose_16;
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1);
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1);
|
||||
|
||||
vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true);
|
||||
vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true);
|
||||
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
|
||||
|
||||
const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts;
|
||||
const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts;
|
||||
const uint32_t d_misalign = get_misalign_bytes(ctx, dst) / ts;
|
||||
|
||||
// Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :]
|
||||
if (ctg->ne[0] > 0) {
|
||||
const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg);
|
||||
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg);
|
||||
pc.ne10 = (uint32_t) ctg->ne[0]; // only ctg's columns in dst
|
||||
pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg);
|
||||
init_pushconst_fastdiv(pc);
|
||||
std::array<uint32_t, 3> el = ne_ctg > 262144 ? std::array<uint32_t,3>{512, 512, CEIL_DIV(ne_ctg, 262144)}
|
||||
: ne_ctg > 512 ? std::array<uint32_t,3>{512, CEIL_DIV(ne_ctg, 512), 1}
|
||||
: std::array<uint32_t,3>{ne_ctg, 1, 1};
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el);
|
||||
}
|
||||
|
||||
// Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :]
|
||||
if (trp->ne[0] > 0) {
|
||||
vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp));
|
||||
pc.ne10 = (uint32_t) trp->ne[0]; // dst bound = trp columns (NOT dst->ne[0])
|
||||
pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp);
|
||||
init_pushconst_fastdiv(pc);
|
||||
std::array<uint32_t, 3> el = {
|
||||
(uint32_t) CEIL_DIV(trp->ne[0], 32),
|
||||
(uint32_t) CEIL_DIV(trp->ne[1], 32),
|
||||
(uint32_t) (trp->ne[2] * trp->ne[3]),
|
||||
};
|
||||
el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
|
||||
el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
|
||||
el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el);
|
||||
}
|
||||
|
||||
ggml_vk_sync_buffers(ctx, subctx);
|
||||
}
|
||||
|
||||
static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
|
||||
int * op_params = (int *)dst->op_params;
|
||||
|
||||
// Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size,
|
||||
// dim1 innermost) and the other source + dst contiguous along dim 0. The generic
|
||||
// shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for
|
||||
// a coalesced copy); here we instead use a tiled shared-memory transpose.
|
||||
auto src_transposed_2d = [&](const ggml_tensor * s) {
|
||||
const uint32_t ts = ggml_type_size(s->type);
|
||||
return s->nb[1] == ts // dim1 innermost
|
||||
&& s->nb[0] == s->ne[1] * ts // consistent 2D transpose
|
||||
&& s->ne[2] == 1 && s->ne[3] == 1;
|
||||
};
|
||||
const uint32_t dst_ts = ggml_type_size(dst->type);
|
||||
const bool dim0_ok = (op_params[0] == 0) && (dst->nb[0] == dst_ts);
|
||||
const bool types_ok = (dst_ts == 4 || dst_ts == 2)
|
||||
&& (src0->type == src1->type && src0->type == dst->type)
|
||||
&& (ggml_blck_size(dst->type) == 1);
|
||||
const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1])
|
||||
&& (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2])
|
||||
&& (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3])
|
||||
&& (src0->ne[0] + src1->ne[0] == dst->ne[0]);
|
||||
// doffset is 16-bit for unary shaders; guard against overflow
|
||||
const bool off_fits = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu;
|
||||
|
||||
const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits
|
||||
&& src_transposed_2d(src1) && ggml_is_contiguous(src0);
|
||||
const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits
|
||||
&& src_transposed_2d(src0) && ggml_is_contiguous(src1);
|
||||
|
||||
if (s1_trans || s0_trans) {
|
||||
const ggml_tensor * ctg = s1_trans ? src0 : src1; // contiguous source
|
||||
const ggml_tensor * trp = s1_trans ? src1 : src0; // transposed source
|
||||
// dst offsets (in elements) where each source region begins
|
||||
const uint32_t off_ctg = s1_trans ? 0u : (uint32_t) src1->ne[0];
|
||||
const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0] : 0u;
|
||||
ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp);
|
||||
return;
|
||||
}
|
||||
|
||||
const uint32_t src0_type_size = ggml_type_size(src0->type);
|
||||
const uint32_t src1_type_size = ggml_type_size(src1->type);
|
||||
const uint32_t dst_type_size = ggml_type_size(dst->type);
|
||||
@@ -16578,7 +16687,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
|
||||
std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr);
|
||||
std::fill(ctx->query_node_idx.begin(), ctx->query_node_idx.end(), 0);
|
||||
|
||||
GGML_ASSERT(ctx->compute_ctx.expired());
|
||||
// compute_ctx may hold async host uploads recorded before the graph; append the timestamp after them
|
||||
compute_ctx = ggml_vk_get_compute_ctx(ctx);
|
||||
ctx->query_idx = 0;
|
||||
compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
|
||||
@@ -18243,11 +18352,112 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
|
||||
return devices[device];
|
||||
}
|
||||
|
||||
// Import an mmap-backed host region as a Vulkan pinned buffer via
|
||||
// VK_EXT_external_memory_host so H2D uploads DMA straight from system RAM
|
||||
// instead of bouncing through the staging buffer + host memcpy. Mirrors the
|
||||
// GGML_CUDA_REGISTER_HOST path; populates device->pinned_memory, which the
|
||||
// existing pinned fast path in ggml_vk_buffer_write_2d_async looks up.
|
||||
//
|
||||
// A single Vulkan buffer cannot cover a whole multi-GB mmap (it is capped at
|
||||
// device->max_buffer_size), so the region is imported in page-aligned chunks.
|
||||
// ggml_vk_host_get resolves a tensor pointer to the chunk that contains it,
|
||||
// and ggml_vk_buffer_write_2d_async falls back to staging for any tensor that
|
||||
// straddles a chunk boundary, so correctness is preserved.
|
||||
static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
|
||||
if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
|
||||
return false;
|
||||
}
|
||||
if (size == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
bool success = false;
|
||||
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
|
||||
vk_device& device = vk_instance.devices[i];
|
||||
if (!device || !device->external_memory_host || device->max_buffer_size == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const size_t align = device->min_imported_host_pointer_alignment;
|
||||
size_t chunk = device->max_buffer_size & ~(align - 1);
|
||||
if (chunk == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
uint8_t * p = static_cast<uint8_t *>(buffer);
|
||||
size_t remaining = size;
|
||||
bool dev_success = false;
|
||||
bool import_ok = true; // flips to false once VK_EXT_external_memory_host import fails
|
||||
while (remaining > 0) {
|
||||
size_t cur = std::min(remaining, chunk);
|
||||
vk_buffer buf;
|
||||
if (import_ok) {
|
||||
buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
|
||||
}
|
||||
if (!buf || !buf->buffer) {
|
||||
// Fallback for drivers that can't import file-backed mmap pages (e.g. RADV):
|
||||
// make a one-time copy of the region into a host-visible Vulkan buffer. The GPU
|
||||
// then DMAs straight from it every eval at full PCIe bandwidth, instead of paying
|
||||
// a slow single-threaded pageable memcpy into the staging buffer each time.
|
||||
import_ok = false;
|
||||
buf = ggml_vk_create_buffer_check(device, cur,
|
||||
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
|
||||
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
|
||||
if (buf && buf->buffer && buf->ptr) {
|
||||
memcpy(buf->ptr, p, cur);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
{
|
||||
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
|
||||
device->pinned_memory.emplace_back(p, cur, buf);
|
||||
}
|
||||
dev_success = true;
|
||||
p += cur;
|
||||
remaining -= cur;
|
||||
}
|
||||
if (dev_success) {
|
||||
success = true;
|
||||
}
|
||||
}
|
||||
return success;
|
||||
}
|
||||
|
||||
static void ggml_backend_vk_unregister_host_buffer(void * buffer) {
|
||||
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
|
||||
vk_device& device = vk_instance.devices[i];
|
||||
if (!device) {
|
||||
continue;
|
||||
}
|
||||
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
|
||||
for (auto it = device->pinned_memory.begin(); it != device->pinned_memory.end(); ++it) {
|
||||
if (std::get<0>(*it) == buffer) {
|
||||
vk_buffer buf = std::get<2>(*it);
|
||||
device->pinned_memory.erase(it);
|
||||
ggml_vk_destroy_buffer(buf);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
|
||||
GGML_UNUSED(reg);
|
||||
if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
|
||||
return (void *) ggml_backend_vk_register_host_buffer;
|
||||
}
|
||||
if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) {
|
||||
return (void *) ggml_backend_vk_unregister_host_buffer;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
|
||||
/* .get_name = */ ggml_backend_vk_reg_get_name,
|
||||
/* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
|
||||
/* .get_device = */ ggml_backend_vk_reg_get_device,
|
||||
/* .get_proc_address = */ NULL,
|
||||
/* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
|
||||
};
|
||||
|
||||
ggml_backend_reg_t ggml_backend_vk_reg() {
|
||||
|
||||
+35
-1
@@ -618,13 +618,47 @@ struct llama_mmap::impl {
|
||||
};
|
||||
|
||||
llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
|
||||
llama_mmap::~llama_mmap() = default;
|
||||
|
||||
llama_mmap::~llama_mmap() {
|
||||
// unpin before the pages are unmapped by the impl destructor
|
||||
if (host_reg_addr && host_unreg_fn) {
|
||||
host_unreg_fn(host_reg_addr);
|
||||
}
|
||||
}
|
||||
|
||||
size_t llama_mmap::size() const { return pimpl->size; }
|
||||
void * llama_mmap::addr() const { return pimpl->addr; }
|
||||
|
||||
void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
|
||||
|
||||
size_t llama_mmap::register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *)) {
|
||||
#ifdef _POSIX_MAPPED_FILES
|
||||
if (host_reg_addr || !reg_fn || !unreg_fn || last <= first) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// expand outward to the page boundaries retained by unmap_fragment
|
||||
const size_t page_size = sysconf(_SC_PAGESIZE);
|
||||
first = first & ~(page_size - 1);
|
||||
last = (last + page_size - 1) & ~(page_size - 1);
|
||||
|
||||
void * reg_addr = (uint8_t *) pimpl->addr + first;
|
||||
if (!reg_fn(reg_addr, last - first)) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
host_reg_addr = reg_addr;
|
||||
host_unreg_fn = unreg_fn;
|
||||
return last - first;
|
||||
#else
|
||||
GGML_UNUSED(first);
|
||||
GGML_UNUSED(last);
|
||||
GGML_UNUSED(reg_fn);
|
||||
GGML_UNUSED(unreg_fn);
|
||||
return 0;
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
|
||||
const bool llama_mmap::SUPPORTED = true;
|
||||
#else
|
||||
|
||||
@@ -50,11 +50,19 @@ struct llama_mmap {
|
||||
|
||||
void unmap_fragment(size_t first, size_t last);
|
||||
|
||||
// pin the pages backing [first, last) with a backend allocator for faster H2D copies,
|
||||
// unpinned in the destructor before the pages are unmapped
|
||||
// returns the number of bytes registered, 0 on failure
|
||||
size_t register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *));
|
||||
|
||||
static const bool SUPPORTED;
|
||||
|
||||
private:
|
||||
struct impl;
|
||||
std::unique_ptr<impl> pimpl;
|
||||
|
||||
void * host_reg_addr = nullptr;
|
||||
void (*host_unreg_fn)(void *) = nullptr;
|
||||
};
|
||||
|
||||
struct llama_mlock {
|
||||
|
||||
@@ -1677,6 +1677,15 @@ bool llama_model_loader::load_all_data(
|
||||
if (size_done >= size_data) {
|
||||
// unmap offloaded tensors and metadata
|
||||
if (use_mmap) {
|
||||
// pin the pages backing the weights kept in system memory for faster H2D copies
|
||||
bool (*reg_fn)(void *, size_t) = nullptr;
|
||||
void (*unreg_fn)(void *) = nullptr;
|
||||
for (size_t i = 0; i < ggml_backend_dev_count() && !reg_fn; i++) {
|
||||
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_dev_get(i));
|
||||
reg_fn = (bool (*)(void *, size_t)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_register_host_buffer");
|
||||
unreg_fn = (void (*)(void *)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_unregister_host_buffer");
|
||||
}
|
||||
|
||||
for (uint32_t idx = 0; idx < mappings.size(); idx++) {
|
||||
const auto & mmap_used = mmaps_used.at(idx);
|
||||
auto & mapping = mappings.at(idx);
|
||||
@@ -1684,6 +1693,13 @@ bool llama_model_loader::load_all_data(
|
||||
if (mmap_used.second != 0) {
|
||||
mapping->unmap_fragment(mmap_used.second, mapping->size());
|
||||
}
|
||||
if (mmap_used.second > mmap_used.first) {
|
||||
size_t n_registered = mapping->register_host(mmap_used.first, mmap_used.second, reg_fn, unreg_fn);
|
||||
if (n_registered > 0) {
|
||||
LLAMA_LOG_INFO("%s: pinned %.2f MiB of mapped model memory for faster H2D transfers\n",
|
||||
__func__, n_registered / 1024.0 / 1024.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (progress_callback) {
|
||||
|
||||
Reference in New Issue
Block a user