Skip to content

fix(cuda): allocate Qwen3.5 scratch lazily - #1879

Open
robertosantosx2 wants to merge 15 commits into
ashhart:mainfrom
robertosantosx2:fix/qwen35-cuda-lazy-scratch
Open

robertosantosx2 wants to merge 15 commits into
ashhart:mainfrom
robertosantosx2:fix/qwen35-cuda-lazy-scratch

Conversation

@robertosantosx2

@robertosantosx2 robertosantosx2 commented Nov 1, 2026 •

Copy link
Copy Markdown

What this changes

Separates Qwen3.5 CUDA forward scratch memory from the model weights arena and allocates it lazily, according to the number of rows required by each forward pass. The INT8 prompt-projection workspace is also allocated on demand and grows geometrically.

This avoids reserving the full forward scratch workspace during model loading, reducing the initial CUDA memory requirement on GPUs with limited VRAM.

Related to #1827.

Receipt

  • Environment: TensorFold commit 14f6500d1; Ubuntu 26.04.1; NVIDIA GeForce RTX 3050 Laptop GPU (4 GB VRAM, compute capability 8.6); Qwen3.5-4B MLX 4-bit checkpoint, revision 32f3e8ecf65426fc3306969496342d504bfa13f3.
  • Exactness: CUDA step and window outputs matched byte-for-byte at 64, 256 and 512 rows. This validates consistency between the tested execution paths; it is not a comparison against draft: false.
  • Decode speed: Approximately 56.5 tokens/s in the tested configuration. A before-and-after release comparison has not yet been performed.
  • Prompt speed: Prompt execution completed at 512 rows. A cold, long-prompt ttft_s comparison using tools/bench_cache.py --require-cold has not been performed.
  • Tests run: 43 passed, 0 failed, 0 skipped across the Qwen3.5 prompt INT8, prompt layout, attention path and prompt MoE host suites.
  • CUDA build: 114/114 build steps succeeded for sm_86 with ReleaseFast.
  • GPU execution: Qwen3.5-4B completed CUDA execution on the RTX 3050 Laptop GPU. Step and window paths were tested up to 512 rows; prompt execution completed at 512 rows. Progressive INT8 workspace growth completed through batches of up to 128 rows.
  • Not run: Metal (Apple Silicon), Vulkan and HIP; this change targets the Qwen3.5 CUDA implementation. A direct benchmark against TensorFold 1.1.2 and the affected later releases remains outstanding.

Checklist

  • No serial inference path was introduced.
  • Host tests cover lazy allocation, capacity growth, allocation failure and recovery.
  • CUDA execution was tested on real NVIDIA hardware.
  • Drafted output compared against draft: false.
  • Resumed prompts compared against fresh prompts.
  • Concurrent execution compared against solo execution.
  • Before-and-after decode and cold prompt benchmarks against the last release.
  • Metal and other non-CUDA backends tested.

The outstanding checks are explicitly marked as not yet validated.

@robertosantosx2

Copy link
Copy Markdown
Author

Ash, I’ve created this patch so that TensorFold continues to work on the smallest RTX cards on the market (my RTX 3050 Laptop -4GB-). This way, we make it possible for the vast majority of people to run local AI.

Thanks,
Best regards

@robertosantosx2

Copy link
Copy Markdown
Author

What this changes

Separates Qwen3.5 CUDA forward scratch memory from the model weights arena and allocates it lazily, according to the number of rows required by each forward pass. The INT8 prompt-projection workspace is also allocated on demand and grows geometrically.

This avoids reserving the full forward scratch workspace during model loading, reducing the initial CUDA memory requirement on GPUs with limited VRAM.

Related to #1827.

Receipt

  • Environment: TensorFold commit 14f6500d1; Ubuntu 26.04.1; NVIDIA GeForce RTX 3050 Laptop GPU (4 GB VRAM, compute capability 8.6); Qwen3.5-4B MLX 4-bit checkpoint, revision 32f3e8ecf65426fc3306969496342d504bfa13f3.
  • Exactness: CUDA step and window outputs matched byte-for-byte at 64, 256 and 512 rows. This validates consistency between the tested execution paths; it is not a comparison against draft: false.
  • Decode speed: Approximately 56.5 tokens/s in the tested configuration. A before-and-after release comparison has not yet been performed.
  • Prompt speed: Prompt execution completed at 512 rows. A cold, long-prompt ttft_s comparison using tools/bench_cache.py --require-cold has not been performed.
  • Tests run: 43 passed, 0 failed, 0 skipped across the Qwen3.5 prompt INT8, prompt layout, attention path and prompt MoE host suites.
  • CUDA build: 114/114 build steps succeeded for sm_86 with ReleaseFast.
  • GPU execution: Qwen3.5-4B completed CUDA execution on the RTX 3050 Laptop GPU. Step and window paths were tested up to 512 rows; prompt execution completed at 512 rows. Progressive INT8 workspace growth completed through batches of up to 128 rows.
  • Not run: Metal (Apple Silicon), Vulkan and HIP; this change targets the Qwen3.5 CUDA implementation. A direct benchmark against TensorFold 1.1.2 and the affected later releases remains outstanding.

Checklist

  • No serial inference path was introduced.
  • Host tests cover lazy allocation, capacity growth, allocation failure and recovery.
  • CUDA execution was tested on real NVIDIA hardware.
  • Drafted output compared against draft: false.
  • Resumed prompts compared against fresh prompts.
  • Concurrent execution compared against solo execution.
  • Before-and-after decode and cold prompt benchmarks against the last release.
  • Metal and other non-CUDA backends tested.

The outstanding checks are explicitly marked as not yet validated.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2, this is the fix our numbers on #1827 pointed at. Main is frozen for the 1.2.1 release tonight, so we'll review and test it right after: a code review, and an A/B on a GB10 with Qwen3.8-27B and 35B at long contexts. That checks that the lazy allocation costs nothing on large GPUs (no stall at the first long prompt, the same speed later, the same replies drafted, resumed and concurrent). We'll post the numbers here.

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2, this is a clean way to stop the 4,096-row reservation on small cards. Every forward grows
the scratch first, and the graphs are torn down before the old buffer goes. A few things before we can land it:

  1. Admission. The server sizes streams and the prompt cache from the free memory it reads after load. With the
    scratch gone from load, that reading now includes the 4,096-row scratch, so on a large GPU the first long prompt
    allocates memory admission has already given out. Could you keep the full size counted, by reserving
    scratch(max_rows, .measure) in the load-time accounting, while allocating lazily? Or allocate eagerly when that
    fits, and lazily only when it doesn't. On GB10s, mid-life allocations are what slowed long prompts in Flash Next on a GB10: long-prompt prefill sometimes drops to about 0.73x and stays there until restart #1868.
  2. Failure. A failed grow now fails the whole lane round, every stream in it, after the old buffer is freed.
    Could a grow that can't fit refuse that request by name, and keep the old buffer?
  3. Graph re-capture. Each growth re-captures every decode graph, and the first growth now happens on a real
    request rather than at warm-up.
  4. Receipts. Could you add drafted-equals-"draft": false, resumed-equals-fresh and concurrent-equals-solo on
    the 4B? We'll run our own A/B on GB10s (27B and 35B at 16K-64K: speed and peak memory, once warm).

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

A second, independent read of this, @robertosantosx2, to make my three questions concrete. The lazy split itself is sound, and the Qwen3.5 host suites pass on it (test-qwen35-prompt-int8-host, -prompt-layout-host, -attention-path-host, -prompt-moe-host and -verify-host: 50/50). Every forward grows the scratch before use (cuda_engine.zig:484-486, cuda_prompt_int8.zig:68-69), and the graphs are torn down before the old buffer is freed.

The three things to fix:

  1. Admission doesn't count the scratch. The server admits streams and the prompt cache from after.room(model) and served.stream_bytes (native/cuda.zig:379, :395). Loaded (qwen3_5/cuda_native.zig:60-67) carries no scratch bytes, and the arena now holds weights only. So the room measured after load still includes the 4,096-row scratch, and the first long prompt allocates memory that admission has already handed out. The smallest fix is a scratch_bytes from e.scratch(max_rows, .measure) in Loaded, subtracted before admitting. Or allocate the full scratch eagerly when it fits and lazily only when it doesn't, which also fixes 3 on large cards.
  2. A failed grow loses the old buffer. cuda_engine.zig:405-415 frees scratch_buf and sets it null before DeviceBuffer.alloc. If that allocation fails, the round fails with no buffer left. cuda_prompt_int8.zig:38-48 has the same order. Please allocate the new buffer first, keep the old one if that fails, and refuse the request with a named error (say error.ScratchGrowFailed). One detail: scratch(capacity, .measure) assigns 0 to the carved pointers, so save and restore them, or have the measure path return offsets only.
  3. Graph recapture on a real request. Each growth deinits every decode graph (cuda_engine.zig:406-409, cuda_prompt_int8.zig:34-37). Since the first growth now happens on a request, the first long prompt pays for every capture. The eager-if-fits path above, or growing and capturing before the first request, avoids it.

The 4B receipts are still open too: tools/bench_concurrent.py --alone --serial and tools/bench_cache.py --require-cold on the 4 GB card would cover drafted, resumed and concurrent requests. Our GB10 A/B on the 27B is almost done, and we'll post it here.

@ashhart

ashhart commented Nov 2, 2026 •

Copy link
Copy Markdown
Owner

Our GB10 run on Qwen3.8-27B (main against this branch, six servers, base, new, base, new, new, base) is in. Replies are identical in all 20 rows, and steady-state speed hasn't moved:

main this branch
16K prefill 1333.1 tok/s 1334.2 tok/s
64K prefill 1130.5 tok/s 1130.4 tok/s
8 streams, chat 92.0 tok/s 92.3 tok/s

But it does show the cost from point 3 above. Both arms' very first request is slow (decode after the first 2K prompt: 12.38 tok/s on main, 11.17 on this branch, against 24.5 later). On this branch, though, the first requests at longer prompts are slow too, where main's aren't:

decode in the first phase main, run 1 this branch, run 2 every later run, both
after a 16K prompt 13.57 tok/s 9.08 tok/s 13.55 to 13.59
after a 64K prompt (10K offset) 21.52 tok/s 10.98 tok/s 22.2 to 22.3
8 streams, code at T=1 87.1 tok/s 79.8 tok/s 87.6

That's the scratch growing on a real request and every decode graph being captured again. So allocating eagerly when the scratch fits (and lazily only when it doesn't), or growing and capturing before the first request, would remove it on large cards. With the two other fixes above, and your 4B receipts, this is ready to land.

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2, that was quick. We read the three commits:

  • A failed grow: fixed. The replacement is allocated first, the old buffer, capacity and pointers are kept, and ScratchGrowFailed is named. The INT8 workspace does the same.
  • Recapture on a real request: fixed on large cards. The full scratch and the INT8 workspace are allocated in warm-up, before any capture. We've queued the GB10 run again to confirm the first-request stall is gone.
  • Admission: fixed where the scratch fits, through that eager allocation. But the new scratch_bytes term only takes effect under a memory limit, because Pool.room ignores held bytes otherwise. Please subtract it from the room at native/cuda.zig:382 and :384, so a 4 GB card's lazy case is covered too.

One more, from our lean check: the branch adds 9 problems over main (433 against 424). They're mostly new multi-line comments, plus cuda_engine.zig at 687 lines against our 600-line limit. python3 tools/zig/lean_check.py lists them. Splitting the scratch code into its own file would cover the length.

With those two and your 4B receipts, it's ready to land.

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2. The new head doesn't compile in our CUDA build (zig build native on a GB10, Zig 0.17.0):

zig/src/families/qwen3_5/cuda_engine.zig:464:33: error: expected type '*memory.DeviceBuffer', found '*const memory.DeviceBuffer'
        if (previous) |*old| old.free();

previous is a const optional, so the capture is a const pointer, and DeviceBuffer.free takes *DeviceBuffer. Declaring it with var previous = ... (or copying it into a var inside the capture) fixes it. zig build check-native-cuda-host compile-cuda-tests -Dtarget=x86_64-linux-gnu catches this without a GPU. We'll rerun the GB10 A/B once it's pushed, together with the two points above.

@robertosantosx2
robertosantosx2 force-pushed the fix/qwen35-cuda-lazy-scratch branch from 4424405 to b250705 Compare November 2, 2026 19:30
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2: the compile fix is in. Two changes left:

  1. Admission still ignores the lazy scratch when no memory limit is set. Loaded.scratch_bytes is added into the model's bytes, but Pool.room ignores held bytes without a limit. So on the small cards this targets, streams and the prompt cache get sized into scratch that isn't allocated yet, and the first long prompt can fail with ScratchGrowFailed. Subtract scratch_bytes from the room itself and count the growth peak (old plus new, up to 1.5x the full scratch), or grow straight to max_rows.
  2. Lean is 434 against main's 424. cuda_engine.zig is at 692 lines and native/cuda.zig at 602. Please bring both to 600 or fewer, for example by moving the lazy-scratch code into its own file.

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2. The 20:54 commit's geometric INT8 growth and its tests look right, and everything we cleared before still holds: the compile fix, the failed-grow rollback, eager allocation in warm-up when the scratch fits, and identical tokens. Two asks from the last review are still open:

  1. Admission ignores the lazy scratch when no memory limit is set. native/cuda.zig:375-378 adds scratch_bytes into model, but Pool.room uses held only under a limit (native/cuda_memory.zig:118-121: with no limit it returns free − reserve). Without TENSORFOLD_CUDA_MEMORY_LIMIT_GB, which is the small-card case this PR is for, streams and the prompt cache are still sized into scratch that isn't allocated yet. So the first long prompt can still fail with ScratchGrowFailed. Geometric growth raises the peak, since a grow holds old and new buffers together. Please either subtract the unallocated scratch at its growth peak (old plus new, capped at the full scratch) from the room at :382/:384, or grow straight to max_rows on the first grow. And please add a test with no limit set.
  2. Lean is 434 against main's 424. families/qwen3_5/cuda_engine.zig is at 692 lines and native/cuda.zig at 602 (the limit is 600). Moving the lazy-scratch code into its own file would bring both back under.

@robertosantosx2

Copy link
Copy Markdown
Author

Independent RTX 3050 4 GiB Validation (Ubuntu, November 3, 2026)

The complete experimental report is documented in LEONES in both languages:

Controlled baseline vs. PR comparison

The experiment was conducted on the same NVIDIA RTX 3050 Laptop GPU (4,096 MiB VRAM, SM 86), using the same mlx-community/Qwen3.5-4B-MLX-4bit checkpoint and Hugging Face revision (32f3e8ecf65426fc3306969496342d504bfa13f3), CUDA backend, context 512, parallel 1, max-inflight 1, and TENSORFOLD_MEMORY_RESERVE_GIB=0.5.

Separate Ubuntu worktrees and ReleaseFast binaries were used. The stable TensorFold installation and the usual llama.cpp server on port 8080 were not modified.

Metric | Baseline | PR #1879 -- | -- | -- Commit | 089a18d | b880e52 Model loading | cuMemAlloc: CUDA_ERROR_OUT_OF_MEMORY | Loaded successfully in 10.4 s HTTP server | Failed to start | Started successfully Successful inference requests | 0 | 3 (HTTP 200) Decode speed | Not measurable | 54.32–56.05 tokens/s TTFT | Not measurable | 187.2–194.1 ms Maximum observed global GPU memory | Not captured during OOM | 3,227 MiB used; 542 MiB free

The three PR inference requests generated 128, 192, and 133 tokens from prompts of 51, 56, and 59 tokens, respectively. Prefill times were 188.1, 184.1, and 183.5 ms.

Baseline monitoring limitation

The baseline GPU monitoring file contains 142 lines (header + 141 samples), covering 00:54:04.339–00:56:24.418. It recorded 15 MiB used / 3,754 MiB free while idle.

The monitor did not capture the failed CUDA allocation itself. Therefore, no baseline OOM-time VRAM peak or allocation profile is available. The idle measurements must not be interpreted as the memory footprint at the time of failure.

Build and test status

The following dedicated host tests and CUDA build targets completed without visible errors:

  • test-qwen35-scratch-host

  • test-cuda-memory-admission-host

  • check-native-cuda-host

  • compile-cuda-tests

The native CUDA binary was built for Linux x86-64, SM 86, ReleaseFast.

The full test-cuda-host suite was not green and is not claimed as passing.

Conclusion

This experiment confirms a functional model-loading and memory-admission advantage on the tested RTX 3050 4 GiB configuration: the baseline fails with CUDA OOM, whereas PR #1879 successfully loads the same model and completes three real inference requests under equivalent conditions.

This is not a decode-speedup comparison, since the baseline never reached inference. It is also not a direct measurement of the exact number of scratch bytes saved, because the baseline OOM-time memory peak was not captured.

The LEONES reports contain the full scratch geometry and formulas, INT8 workspace calculations, memory-admission and buffer-growth analysis, commands, commits, GPU monitoring, inference metrics, evidence filenames, and experimental limitations.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2: both asks are answered. The pending lazy scratch is now subtracted from the room with or without a memory limit (with an unlimited-pool test), and the scratch moved to qwen3_5/cuda_scratch.zig with lean at 424. Our GB10 run (main against this head on the 27B) is queued, and it lands when green. Small points, not blocking: growth can still overshoot the scratch rows, and there's no start-up room wiring test.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @robertosantosx2. Our GB10 A/B is in: main against this head (b880e5270) on Qwen3.8-27B 4-bit, one GB10, warm then base-new-new-base.

  • Exact: 0 of 20 prompt rows differ across runs.
  • Prompt speed: level, 1414.6 → 1416.5 tok/s at 2K, 1394.4 → 1395.4 at 8K, 1260.0 → 1260.4 at 32K.
  • Decode after a prompt: level, 24.7 → 24.7 tok/s at 2K, 14.2 → 14.2 at 8K, 12.9 → 12.9 at 32K. Every arm's runs overlap.

So the lazy scratch costs nothing on a large card. It's in our landing queue now, with a changelog line crediting you.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

One thing before it can land, @robertosantosx2: main moved under this since the run. c7ee94f4ca (an EXL3 checkpoint's embedding read from host memory when the card has no room) moved the arena allocation into qwen3_5/cuda_host_embed.zig. Its arena() and its room check both still add the whole e.scratch(.measure).

Rebasing hits conflicts in cuda_engine.zig, cuda_native.zig, cli/qwen35_learn.zig and tests/qwen35_cuda_main.zig. The part that needs your judgement is cuda_host_embed.zig:

  • arena() should allocate only the weights.
  • mapped() should count the lazy scratch's peak the same way your admission change does.

Then an EXL3 27B on a small card still picks host-memory embeddings correctly. Could you rebase onto current main? We'll re-run the GB10 cell on the new head and land it when it's green.

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.

2 participants