[shuuul]

Relates to #2102.

What

DeviceBuffer.download now stages through a per-thread pinned bounce only on an integrated device; discrete cards keep their direct DtoH copy and their speed. The engine's CU_DEVICE_ATTRIBUTE_INTEGRATED reading gates it (cached per thread and driver). A new cuda/download_stage.zig owns the bounce, which counts in memory.usage and is released with the thread's buffer. downloadAsync is unchanged and still requires pinned memory, as its doc says.

The smoke test downloads 6 MiB through two pieces. A fake-driver test covers both paths: integrated stages into pinned memory and the pageable caller gets the bytes; discrete calls {{SECRET_kr7wjfjl}} with the caller's slice and allocates nothing.

Why

On GB10 (driver 580.178.04, ATS) a synchronous DtoH copy straight into pageable host memory can return success while some pages keep their old contents and the bytes land elsewhere. The downstream zspark fork hit it 10 times in 60 rank-round runs, corrupting an allocator's metadata; pinning the compare buffers gave 0/60, staging gave 0/100. Standalone loops and 200 unmodified tf-cuda-test smoke runs here never miss, so the issue has no small public reproducer; #2102 records what we have.

This is the defensive fix, not a claim that TensorFold's pointers are wrong.

Cost (idle GB10, 2,000 copies for the small sizes, direct against staged)

bytes	direct	staged
4	5.90 us	6.51 us
2,048	4.71 us	5.08 us
32,768	5.33 us	6.91 us
303,104	10.59 us	20.97 us
4 MiB	84.74 us	212.36 us
64 MiB	1,354.59 us	4,095.36 us

These are the integrated-device numbers; a discrete card takes the direct column unchanged. The 303,104-byte row is a logits row; the large rows are checks/test paths, not decoding.

Verified

• zig build test-runtime-host: 282/282 tests pass, including both fake-driver paths.
• zig build -Dnvcc=/usr/local/cuda/bin/nvcc compile-cuda-tests, then tf-cuda-test smoke on a GB10: all pass, including the 6 MiB case.
• python3 tools/zig/lean_check.py: 425 problems, the same as main; the new comments are one line each.
• 200 consecutive smoke runs pass.

Review resolution

• Integrated devices only; discrete cards unchanged.
• The bounce is counted in memory.usage and released by drop.
• No conflict with CUDA on integrated GPUs (GB10): keep prompt states and weights in place instead of copying them through pinned host memory #2043: qwen CUDA: an integrated GPU keeps prompt states in device memory, without pinning a copy each save (#2043) #2113 keeps states on the device and never downloads.
• Comments are one line each (lean stays at 425).

For the A/B on a GB10: replies should be byte-equal, decode unchanged, and the logits-row copy is the 303,104-byte row above.

[ashhart]

Thanks @shuuul, and sorry for the wait. A careful write-up of a fault that's hard to pin. Your three questions: stage only where the device is integrated: the engine already reads CU_DEVICE_ATTRIBUTE_INTEGRATED, so no new probe is needed, and discrete cards keep their direct copies and their speed. Count the bounce in memory.usage (or take it from copy_source's pinned spares), since it's pinned memory the budget doesn't see and it stays for the thread's life. And there's no conflict with #2043: #2113 keeps states on the device and never downloads. One small thing besides: on our main the new comments in download_stage.zig and memory.zig run over several lines, so lean reads 432; one line each, please. With those, we'll A/B it on a GB10: replies byte-equal, decode level, and the logits-row copy time.

[shuuul]

All four are in: staging is gated on CU_DEVICE_ATTRIBUTE_INTEGRATED (cached per thread/driver), the bounce counts in memory.usage and is released by drop, the comments are one line each (lean 425), and the fake-driver tests now cover both the integrated and discrete paths. zig build test-runtime-host is 282/282; tf-cuda-test smoke passes on a GB10 including the 6 MiB case. Marked ready for review.

[ashhart]

Thanks @shuuul, all four asks read as done. A note so this doesn't confuse anyone: we'd opened #2200 for the same issue without noticing this PR. I've closed it, because yours is the better fix and came first. This one's in review now; I'll come back to you with the result.

[ashhart]

@shuuul, one thing CI found on b76aa6373: zig and zig-windows crash in zig build test, outside test-runtime-host. download now calls stage.integrated, and cuCtxGetDevice at download_stage.zig:20 faults in host tests whose fake driver tables never set it. For example:

• tests/qwen4/cuda_draw_launch_test
• projection_f32_host_tests
• w4a8_exception_gate_host
• families/qwen3_5/cuda_pocket_test
• quant_ref_host_tests

Those fakes need cuCtxGetDevice and cuDeviceGetAttribute set, and they should answer discrete so each test keeps its direct copy. Please run the whole zig build test before pushing. The rest of the review will follow on the new head.

[shuuul]

Thanks. I'll set cuCtxGetDevice and cuDeviceGetAttribute in the fake driver tables named there, answering discrete so each test keeps its direct copy, and run the whole zig build test before pushing the new head. I'll address the rest of the review on that head.

[shuuul]

Pushed 0d5cd20. Every fake driver table that reaches download now sets cuCtxGetDevice and cuDeviceGetAttribute and answers device 0 / attribute 0, so the probe takes the discrete path and the direct copy: the five you named, plus the DeepSeek load mock, the H3 runtime fake, Flash Next's pocket test, contract_rows_host, contract_verify_host and contract_q35_attn_host (the last one also crashed in zig build test here). zig build test exits 0 on this head. The rest of the review can follow on it.

[ashhart]

Thanks @shuuul. That's every fake table, and answering discrete keeps each test on its old direct copy. One small thing: integrated() caches its reading per thread by the driver's address, and the integrated test ends with drop() but not forget(). A later test whose fake driver lands at the same stack address would reuse "integrated" and call cuMemHostAlloc on a table that doesn't set it. Could you add defer stage.forget(); to both download tests? The GB10 run is queued on our side.

[shuuul]

Done in 2d87813: both download tests now defer stage.forget() after defer stage.drop(), so the cached reading does not outlive the fake driver's stack address. zig build test-runtime-host passes.

[ashhart]

Thanks @shuuul, that's the one line. It's queued for after the 1.2.3 cut, with our GB10 run.

[ashhart]

Our GB10 run, main against this PR, Qwen3.8-27B: verdict GREEN on the head before your one-line forget() fix. 32K prompts and decode after them are within 0.1%, the 1- and 4-stream sweeps within 0.3%, and there are no token differences across runs in the 20 prompt rows I checked. A level result is what we hoped for. Thanks @shuuul: it's queued for after the 1.2.3 cut.

[ashhart]

Thank you @shuuul. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.