Skip to content

vulkan: staged device-local uploads for cards without Resizable BAR - #1871

Merged
JustVugg merged 8 commits into
devfrom
feat/vk-staged-uploads
Oct 3, 2026
Merged

JustVugg merged 8 commits into
devfrom
feat/vk-staged-uploads

Conversation

@JustVugg

@JustVugg JustVugg commented Oct 3, 2026

Copy link
Copy Markdown
Owner

On a discrete GPU without Resizable BAR, Vulkan put nothing on the card. The host-visible device-local window there is about 256 MB, and the backend placed resident weights in that memory type. An RTX 3070 on its launch VBIOS showed it (reported on #1784):

  • only 246 of 8192 MB VRAM is host-visible;
  • 0 matrices resident;
  • the expert tier failed 3004;
  • the chain "a matrix did not reach the device".

NVIDIA's driver refuses allocations past the window; RADV instead moves them to system RAM silently, so they are read over PCIe. Every Turing card and every Ampere card on its launch VBIOS is in that situation.

This adds staged device-local uploads: resident data goes to DEVICE_LOCAL memory through a host staging buffer and vkCmdCopyBuffer. The approach (a device-local arena, the staging copy, a submit lock, the first-touch fill of fresh blocks) comes from #1338 by @crichalchemist, rewritten for the current backend, expert tier and chain.

When it stages

  • COLI_VK_STAGED=1 forces staging and 0 forces the mapped path.
  • Unset, staging is used when the host-visible device-local heap is under a quarter of the largest device-local heap, or when no such type exists. That catches a 256 MB window on an 8 GB card, and never a ReBAR card, an integrated GPU or Lavapipe.
  • COLI_VK_HOST_VISIBLE_CAP_MB=N caps the heap size the rule sees, to emulate a card without ReBAR.

When the rule says mapped, the queues, memory types and command buffers are exactly as before.

What moves (only when staging)

  • The dense pool: device-local blocks, with rows and scales streamed through two 16 MiB staging slots.
  • The expert tier's pool: an expert's three matrices are committed in one submit by the tier's uploader thread (coli_vk_tensor_commit). On mapped memory the commit does nothing.
  • The second device's pool (COLI_VK_DEV2).
  • The MLA KV mirror and the q-prep norm weights: each host write becomes a copy recorded at the head of the next command buffer that reads it.
  • The dense chain: its host-written buffers use host staging memory instead of the window.

Readbacks and device-only scratches are unchanged. Each fresh device-local block is zero-filled before first use, #1338's Polaris finding; it does not reproduce on our three devices, but it costs one pass per 256 MB block, so it stays. With staging on, the exit line reports the resident data left in host memory, and the tests check it is 0.

Synchronization. Each device has one uploader with its own command pool, two command buffers and two fences, and does one upload at a time under a mutex. The queue it uses is, in order of preference:

  1. a transfer-only family;
  2. a spare queue of a family already in use;
  3. the main queue, in which case every submit, the chain's included, takes the same lock.

An upload waits on its fence before returning, so the tier hands an expert to the engine thread only after its copy has landed.

When an upload fails

A failure fails that upload, not the run:

  • a failed dense upload leaves that one matrix on the CPU;
  • a failed expert commit frees the slot and the expert stays on the CPU. The tier's budget is unchanged and the expert can be promoted again;
  • a failed fence wait means the device is lost, as for every other wait in the backend. The tier stops at its next step, and the chain rebuilds its state on the CPU at its next frame.

COLI_VK_STAGED_FAULT=<point>[:n] injects a failure at each point: stage, pwstage, block, kvbuf, record, submit, wait and commit.

Tests

  • Same bits: the backend harness's digest of every device result is identical for mapped, staged, staged without the fill, and the 246 MB cap, which picks staging on its own. This holds on Lavapipe, the Intel Iris Xe (Dozen) and the Radeon 780M; a patched copy of the backend before this change gives the same digest.

  • COLI_VK_STAGED=1 on Lavapipe:

    • shader;
    • qwen 48, qwen-sanitize 15, qwen-chain 55;
    • inkling-olmoe-chain 43;
    • glm-chain 53 (including COLI_VK_DEV2), glm-chain-sanitize 12, glm 34;
    • mimo-chain 53.

    The new staged and staged-window families pass. test_vk_tier and test_vk_chain pass staged on all three devices, with 0 MiB left in host memory.

  • staged-faults and staged-faults-sanitize (new): 32 runs over the dense path, the tier (awaited, and with its uploader thread running freely), the chain and the KV mirror. Each must have no crash and no sanitizer report, must show the fault fired, and must give the CPU's tokens. Results: 32/32 on Lavapipe, 32/32 under ASan+UBSan, 32/32 on the 780M.

  • New CI legs: staged, the shader family staged, qwen-staged, qwen-chain-staged, staged-faults and staged-faults-sanitize.

  • Default unchanged: on our devices the rule picks the mapped path, and its bits are the previous build's.

Measured on the Radeon 780M (Qwen3.6-35B-A3B, the docs' method)

Staging costs nothing where it is not needed:

mapped staged
Decode, tier 7.94 tok/s 7.96 tok/s
Decode, tier and chain 9.91 tok/s 9.96 tok/s
512-token prefill, tier 12.26 / 12.22 s 12.29 / 12.15 s
512-token prefill, tier and chain 9.53 s 9.43 s

Not measured

  • A real discrete card without ReBAR: we have none. @kreuzzelg offered to run it on the RTX 3070.
  • A transfer-only queue: none of our devices exposes one.
  • The second device under fault injection: it uses the same uploader code, with its own lock, but the fault family does not exercise it.

Without Resizable BAR a discrete card exposes HOST_VISIBLE|DEVICE_LOCAL as a
window of about 256 MB: NVIDIA refuses the backend's resident weights past it
(0 matrices resident, the tier failing every upload), RADV puts them in system
RAM. Staged uploads put resident data in a DEVICE_LOCAL type the host does not
map and copy it there from a host staging buffer with vkCmdCopyBuffer.

The rule: staged when the host-visible device-local heap is under a quarter of
the largest device-local heap (or there is none); COLI_VK_STAGED=1 forces it, 0
forbids it; COLI_VK_HOST_VISIBLE_CAP_MB=N treats the host-visible heap as at
most N MiB, so a device with Resizable BAR or unified memory takes the decision
a card without it takes. A device that stages nothing keeps the mapped path,
its queues and its bits.

What moves: the weight pools (dense tensors, the expert tier, COLI_VK_DEV2's),
streamed through two 16 MiB staging slots; the tier's tensors are filled in a
host image and committed (coli_vk_tensor_commit) by the uploader thread; the MLA
KV mirror and the q-prep norm weights go through pending copies recorded at the
head of the absorb and q-prep submits; the chain's host-written buffers (frame
staging, per-step rows) move to host staging memory instead of the window. A
fresh device-local block is zero-filled before its first tensor (#1338's
Polaris finding). The uploader takes a transfer-only queue, else a spare queue,
else the main queue under a lock that every backend and chain submit takes.
The approach follows #1338 by crichalchemist.
tests/vulkan_engines.sh staged: the VK_TEST harness's results bit for bit the
same mapped, staged twice and under an emulated 246 MB window that must choose
staging on its own (a digest of every result the device returns), then the
tier and the chain's ops staged, each run ending with no resident data in host
memory. <family>-staged runs a family with COLI_VK_STAGED=1; qwen-staged first
runs qwen36's tier and chain under the emulated window with the decision left
to the backend. CI: the staged family and the shader family staged in the
Vulkan job, qwen-staged and qwen-chain-staged in the engines matrix.
With staged uploads the backend reports at shutdown where the resident data
ended up; tests/vulkan_engines.sh staged checks that line after the tier's
test as after the chain's.
What a discrete card without Resizable BAR did to the mapped path, the rule
that turns staged uploads on, what moves where, how the uploader is
synchronized, the zero fill of a fresh block, the lines it prints, what was
tested on Lavapipe, the Iris Xe and the 780M, and mapped against staged on the
780M (the same: decode 7.94 against 7.96 tok/s with the tier, 9.91 against
9.96 with the chain). A discrete card without Resizable BAR was not measured.
ENVIRONMENT.md: COLI_VK_STAGED, COLI_VK_HOST_VISIBLE_CAP_MB, COLI_VK_TEST_NOFILL.
The section checks each format's numbers, and its twelve tiny steps could end
before the uploader landed any expert: with staged uploads a commit is a submit
and a fence wait (8.5 ms for the first, 0.6 ms after, on an Iris Xe through
Dozen), and 2 of 4 staged runs there failed one of the first formats on
'nothing ran on the device'. COLI_VK_TIER_SYNC=1 for that section only; warm,
adapt and partial keep the uploader running free.
A command buffer that would not record or a submit refused failed the uploader
for good and took the device down with it. Now it fails that upload only: the
uploader waits for what it had sent, the caller frees the tensors and the next
upload starts clean, so the matrix stays on the CPU, and an expert whose commit
failed stays on the CPU with the tier's budget intact (vk_tier tells a failed
copy from a refusal, which still shrinks the plan). A fence wait that fails
still loses the device, as every other wait does, now from any thread, and the
rest follows: the tier stops at its next step (coli_vk_xb_ready), the chain at
its next frame (vkc_begin), rebuilding its state on the CPU.

COLI_VK_STAGED_FAULT=<point>[:n] makes the n-th staged upload fail at the
staging buffer, the KV mirror's, a device-local block or buffer, a command
buffer's recording, a submit, a fence wait or a tier commit after its first
matrix. The exit line counts the submits.
tests/vulkan_engines.sh staged-faults: qwen36's matrices one by one, its tier
awaited and with the uploader thread free after a warm start, its chain, and
colibri's KV mirror, each through COLI_VK_STAGED_FAULT at every point it
reaches. Every run: no crash, the fault fired, the CPU's tokens; one matrix
fewer on the device where one failed, the tier's budget kept where a commit
failed, the state rebuilt where the device was lost mid-run.
staged-faults-sanitize: the same under ASan and UBSan. Both in the engines
matrix. docs/vulkan.md says what each failure does.
@JustVugg
JustVugg merged commit 0c4a751 into dev Oct 3, 2026
56 checks passed
JustVugg added a commit that referenced this pull request Oct 4, 2026
…eepSeek chain

Conflict: the usage lines and the dispatcher of tests/vulkan_engines.sh,
where both sides added families; the chain families come before the
staged entries, and the usage string names them all. The chain's one
submit already goes through coli_vk_queue_submit, the backend's lock.
@kreuzzelg

kreuzzelg commented Oct 4, 2026 •

Copy link
Copy Markdown
Contributor

@JustVugg Ran it on the RTX 3070 with its launch VBIOS (94.04.3A.00.9D, resource1_resize 0x1c0 = 256 MB BAR maximum), dev b86761bf, make qwen36 VK=1, Qwen3.6-35B-A3B int4, the usual 200-token prompt, the 3070 alone, Ollama stopped. Nothing set but COLI_VULKAN=1.

The rule picks staging on its own, and the card now holds everything:

[VK] memory: staged uploads, resident data in device-local memory (type 1, 8192 MiB heap) copied from host staging memory (type 3) on a transfer queue (246 of 8192 MiB of device-local memory is host-visible)
[VK] ready: NVIDIA GeForce RTX 3070, compute qfam 0, memtype 5, fused gate+up, absorb attention, tiled GEMM from S=2 and S*O>=4096, staged uploads
[VK] tier qwen36: on, NVIDIA GeForce RTX 3070, budget 4.43 GiB = 2995 experts of 1.5 MiB (fmt 2, down fmt 2), device memory, own queue, up to 16 promotions per token, balanced against the CPU
[VK] qwen36 chain: 40 layers on the device (10 attention), 4.4 MiB of parameters
[VK] qwen36: 0 matmuls on the GPU (361 matrices resident, 1863.8 MiB; placed int8 291, int4 0, f32 0)
[VK] tier qwen36 run: device 66271 of 79360 routed experts (83.5%) | ... | resident 2994 (budget 2994, 4.42 GiB of 4.43 GiB, 18 blocks, frag 0.62) | uploads 3174 (4.65 GiB, 2994 warm), evictions 180, skipped 0 queue + 0 rate, failed 0 | device 2075.5 ms, CPU share 16557.2 ms, waited 1681.2 ms (19% of device time hidden) | balance: device share 0.05, 439 rows handed to the CPU
[VK] memory at exit: weights 1863.8 MiB, expert tier 4535.2 MiB (peaks), KV mirror 0.0 MiB in device-local memory type 1 (not host-visible); the dense chain's state in type 1 (device-local); 6662.0 MiB staged in 19796 copies (3630 submits), 25 blocks zero-filled; resident data in host memory: 0.0 MiB

So this card has a transfer-only family and the uploader took it: the queue configuration you listed as not measured. nvidia-smi showed 7254 MiB in use on the 3070 during the run. Five runs (one 20-token smoke, four of 200 tokens), failed 0 in every one, and the text matches the CUDA path's.

3070, 200 tokens TTFT ms/token tok/s (TTFT included) decode tok/s routes on the device matrices resident
staged (unset), cold heat 6.52 s 76.5 7.43 13.1 78.7 % 361
staged (unset), warm, three runs 6.42–6.77 s 65.6–67.3 7.90–8.09 14.9–15.2 83.0–83.7 % 361
COLI_VK_STAGED=0, same build (control) 18.39 s 143.3 4.25 7.0 0 % 0
before #1871 (dev de4451e, 20 tokens, reported on #1784) 17.46 s 229.9 0.92 4.3 0 % 0
CUDA tier on the same card (Q36_DN_GPU=1, COLI_PLACE=auto) 1.33 s 32.1 25.6 31.1 83 % 553

The control run with COLI_VK_STAGED=0 reproduces the old state (the ReBAR warning, 0 matrices, tier failed, chain "a matrix did not reach the device"), so the rule and the mapped path are both doing what the PR says.

Two observations, for what they are worth:

  • Where the decode time goes. At 66 ms per token the run is in the routed-expert step: the chain reports ~5000 ms waiting for the device over 200 forwards (25 ms/token), the tier ~2050 ms of device time with ~1650 ms waited and balance: device share 0.05, although 83 % of the routes ran on the device. On the same card the CUDA tier decodes at 32 ms per token with the same hit rate, and your 780M does 9.9 tok/s with tier and chain, so the 3070 is running well under what it has. Not a staging matter, I think (the data is in VRAM and nothing is left in host memory), but the first Vulkan number from a discrete NVIDIA card; happy to run whatever would help pin it down (COLI_VK_TIER_BALANCE=0, COLI_VK_TIER_SYNC=1, chain off, tier queue shared).
  • TTFT 6.4–6.8 s against 1.33 s on CUDA. Each start copies 6.7–7.0 GB through the two 16 MiB staging slots (the tier's warm start of 2995 experts plus the 1.9 GB trunk, ~3700 submits); that is about 5 s of it. The 25 zero-filled blocks did not show (the smoke, which also filled them, started 7.45 s; the later runs 6.4–6.8 s).

Logs and the runner are in my workspace if you want any of the raw files.

Thanks to Claude Code for the help along the way.

@crichalchemist

crichalchemist commented Oct 4, 2026 •

Copy link
Copy Markdown
Contributor

Figures for this from a discrete card without Resizable BAR, since the PR's "Not measured" list has none. Staging turns on by itself here and decodes 2.9 times as fast as the mapped path.

Setup. dev at b86761b, make -C c qwen36 VK=1 (native, 0 warnings). Radeon Pro 580X (Polaris, 8 GB, 256 of 7936 MiB host-visible), Mesa 26.0.8 RADV, Ubuntu 26.04, i7-7700K (4C/8T, AVX2), 64 GB RAM. Qwen3.6-35B-A3B int4 gs64 on a spinning SATA disk, fully in the page cache. One 22-token prompt, greedy, 100 tokens. One warm-up per arm (discarded), then three runs, arm order rotated between rounds. The tier's usage history is one frozen file, copied fresh before every tier run.

# tier arms; mapped adds COLI_VK_STAGED=0, staged leaves it unset
SNAP=$M TOK=$M/tokenizer.json N_NEW=100 COLI_NO_OMP_TUNE=1 NOSTREAM=1 COLI_USAGE=usage.run \
VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/radeon_icd.json COLI_VULKAN=1 COLI_VK_SHADERS=shaders \
COLI_VK_TIER_GB=4 ./qwen36 64 4 prompt.txt
# CPU arm: the first line only. Control: COLI_VK_TIER_GB=0.15 COLI_VK_DENSE=0 COLI_VK_CHAIN=0

With the 4 GiB budget the device holds 2416 experts (3.98 GiB) and a 1863.8 MiB trunk, about 5.8 GiB against the 256 MiB window.

tok/s, median (three runs) whole process, median
CPU only 5.57 (5.50, 5.57, 5.79) 26.8 s
tier, mapped (COLI_VK_STAGED=0) 2.78 (2.76, 2.78, 2.83) 46.3 s
tier, staged (default here) 8.00 (7.99, 8.00, 8.68) 22.3 s
control: 0.15 GiB tier, trunk and chain on the CPU, mapped 5.44 (5.09, 5.44, 5.46) 27.3 s
the same control, staged 5.40 (5.14, 5.40, 5.41) 27.4 s
  • The decision. Unset, the log reads staged uploads, resident data in device-local memory (type 0, 7936 MiB heap) copied from host staging memory (type 2) on a queue of its own (256 of 7936 MiB of device-local memory is host-visible).
  • Mapped against staged. The generated text is identical in all eight tier runs (warm-ups included). With COLI_VK_TIER_SYNC=1 COLI_VK_TIER_BALANCE=0 the dumped last logits of the two modes are byte-identical, and both served 24378 of 38720 routed experts on the device.
  • Where everything fits in the window (the control, 91 experts, 153.6 MiB), the two modes are equal within run-to-run spread. The control's text varies between runs in both modes.
  • Against the CPU the comparison is only indicative: the device arms' text differs from the CPU-default text, as the "Arithmetic" section of docs/vulkan.md describes for f32 activations on the device.
  • Peak RSS: 8.10 GB CPU, 8.19 GB mapped, 8.22 GB staged. TTFT 2.7 to 3.7 s in the CPU and 4 GiB tier runs, 3.4 to 4.5 s in the control.

One card, one model, one prompt. The sweep results from the same machine are on #1867.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants