From e449b1ca08d868abcc4ab73070dac10fbb819f5a Mon Sep 17 00:00:00 2001
From: shuuul <6903405+shuuul@users.noreply.github.com>
Date: Wed, 4 Nov 2026 20:43:00 +0800
Subject: [PATCH 1/5] cuda: stage synchronous downloads through a thread's
 pinned buffer

On GB10 (driver 580.178.04, ATS) a synchronous cuMemcpyDtoH_v2 straight into
pageable host memory can return success while some pages keep their old
contents and the bytes land elsewhere; zspark's rank-round test met it 10
times in 60 runs, corrupting its allocator's metadata. Pin a bounce buffer
per thread (at most 4 MiB a piece) and copy on from it with the CPU.

The cost on an idle GB10 (2,000 copies, direct against staged): 4 B 5.90 vs
6.51 us, 32 KiB 5.33 vs 6.91, a 303,104 B logits row 10.59 vs 20.97, 4 MiB
84.74 vs 212.36, 64 MiB 1,354.59 vs 4,095.36.

The smoke test now downloads 6 MiB through two pieces.
---
 zig/src/cuda/download_stage.zig  | 39 ++++++++++++++++++++++++++++++++
 zig/src/cuda/memory.zig          | 12 +++++++++-
 zig/tests/cuda/runtime_tests.zig | 10 ++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 zig/src/cuda/download_stage.zig

diff --git a/zig/src/cuda/download_stage.zig b/zig/src/cuda/download_stage.zig
new file mode 100644
index 0000000000..77a8f9efd8
--- /dev/null
+++ b/zig/src/cuda/download_stage.zig
@@ -0,0 +1,39 @@
+//! What every synchronous download lands in first: this thread's page-locked memory, at most `most`
+//! bytes, copied on into the caller's by the CPU. On GB10 a `cuMemcpyDtoH` straight into pageable
+//! memory (the host allocator's, freshly mapped) can miss pages with no CUDA error: some keep what
+//! they held, the bytes land on other pages of the process. Calls to `DeviceBuffer.download` stage
+//! through here; `downloadAsync` still needs pinned memory, as documented.
+const std = @import("std");
+const abi = @import("abi.zig");
+const Driver = @import("driver.zig").Driver;
+const Error = @import("driver.zig").Error;
+
+/// The largest piece one copy takes; larger downloads go through in pieces.
+pub const most: usize = 4 << 20;
+
+const Held = struct { d: *const Driver, bytes: []align(16) u8 };
+threadlocal var held: ?Held = null;
+
+/// This thread's page-locked buffer of at least `len` bytes, reallocated when a larger piece is asked
+/// for. The driver allocates it directly, so it is not counted in `memory.usage`: it is one per
+/// thread, bounded by `most`, and kept for that thread's next download.
+pub fn take(d: *const Driver, len: usize) Error![]align(16) u8 {
+    if (held) |h| {
+        if (h.d == d and h.bytes.len >= len) return h.bytes[0..len];
+        drop();
+    }
+    const size = @min(most, std.math.ceilPowerOfTwoAssert(usize, @max(len, 64 << 10)));
+    var p: ?*anyopaque = null;
+    try d.check(d.api.cuMemHostAlloc(&p, size, abi.host_alloc_portable), "cuMemHostAlloc");
+    const base: [*]align(16) u8 = @ptrCast(@alignCast(p.?));
+    held = .{ .d = d, .bytes = base[0..size] };
+    return held.?.bytes[0..len];
+}
+
+/// Frees this thread's buffer, if one is held.
+pub fn drop() void {
+    if (held) |h| {
+        _ = h.d.api.cuMemFreeHost(h.bytes.ptr);
+        held = null;
+    }
+}
diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig
index 5c0022dfce..6e77c8dd88 100644
--- a/zig/src/cuda/memory.zig
+++ b/zig/src/cuda/memory.zig
@@ -190,10 +190,20 @@ pub const DeviceBuffer = struct {
         try @import("copy_source.zig").synchronousCopy(self.d, dst, bytes, std.heap.page_allocator);
     }
 
+    /// Returns once the bytes have landed; the copy stages through this thread's page-locked buffer, a
+    /// piece of at most `download_stage.most` bytes at a time (see that module for why).
     pub fn download(self: DeviceBuffer, offset: usize, out: []u8) Error!void {
         const src = try self.span(offset, out.len);
         if (out.len == 0) return;
-        try self.d.check(self.d.api.cuMemcpyDtoH_v2(out.ptr, src, out.len), "cuMemcpyDtoH");
+        const stage = @import("download_stage.zig");
+        var done: usize = 0;
+        while (done < out.len) {
+            const n = @min(out.len - done, stage.most);
+            const bounce = try stage.take(self.d, n);
+            try self.d.check(self.d.api.cuMemcpyDtoH_v2(bounce.ptr, src + done, n), "cuMemcpyDtoH");
+            @memcpy(out[done..][0..n], bounce[0..n]);
+            done += n;
+        }
     }
 
     /// Owns a pinned source copy until completion. Captured uploads must use copy_source.Stage or Buffer.
diff --git a/zig/tests/cuda/runtime_tests.zig b/zig/tests/cuda/runtime_tests.zig
index 53a843c812..e3189c27dd 100644
--- a/zig/tests/cuda/runtime_tests.zig
+++ b/zig/tests/cuda/runtime_tests.zig
@@ -25,6 +25,16 @@ pub fn smoke(gpu: Gpu) !void {
     try check.sameBytes("H2D, D2D, D2H round trip (1 MiB)", back, pattern);
     check.pass("blocking copies: 1 MiB host -> device -> device -> host equal", .{});
 
+    // Larger than one staged piece: the copy goes through in pieces and still lands whole.
+    const big_n: usize = 6 << 20;
+    var big = try cuda.DeviceBuffer.alloc(d, big_n);
+    defer big.free();
+    try big.fill8(0x3c, null);
+    const big_back = try check.download(gpu, big);
+    defer gpu.gpa.free(big_back);
+    try expect(std.mem.allEqual(u8, big_back, 0x3c), "6 MiB download", .{});
+    check.pass("blocking copies: a 6 MiB download through staged pieces equal", .{});
+
     var stream = try cuda.Stream.init(d, true);
     defer stream.deinit();
     var pinned = try cuda.HostBuffer.alloc(d, n);

From c4cefd8949db9bd94e49d5e16e81a19d72dd05d4 Mon Sep 17 00:00:00 2001
From: shuuul <6903405+shuuul@users.noreply.github.com>
Date: Wed, 4 Nov 2026 22:03:58 +0800
Subject: [PATCH 2/5] cuda: a fake driver pins the download staging path

The fake writes only pinned destinations; the test asserts the copy's
destination is the thread's pinned buffer and the pageable caller gets the
bytes, so a regression to the direct copy fails without needing the fault.
---
 zig/src/cuda/download_stage_test.zig | 46 ++++++++++++++++++++++++++++
 zig/src/cuda/root.zig                |  1 +
 2 files changed, 47 insertions(+)
 create mode 100644 zig/src/cuda/download_stage_test.zig

diff --git a/zig/src/cuda/download_stage_test.zig b/zig/src/cuda/download_stage_test.zig
new file mode 100644
index 0000000000..188abb6795
--- /dev/null
+++ b/zig/src/cuda/download_stage_test.zig
@@ -0,0 +1,46 @@
+//! The synchronous download's staging path, with a driver that only writes pinned destinations.
+const std = @import("std");
+const abi = @import("abi.zig");
+const Driver = @import("driver.zig").Driver;
+const memory = @import("memory.zig");
+const stage = @import("download_stage.zig");
+
+var pinned: [64 << 10]u8 align(16) = undefined;
+var destination: usize = 0;
+
+fn pinnedRange(ptr: usize) bool {
+    const base = @intFromPtr(&pinned[0]);
+    return ptr >= base and ptr < base + pinned.len;
+}
+
+fn hostAlloc(out: ?*?*anyopaque, len: usize, _: c_uint) callconv(.c) abi.Result {
+    if (len > pinned.len) return 1;
+    out.?.* = &pinned[0];
+    return 0;
+}
+
+fn hostFree(_: ?*anyopaque) callconv(.c) abi.Result {
+    return 0;
+}
+
+/// Writes only a pinned destination, as the GB10 fault leaves pageable memory untouched.
+fn copyToHost(dst: ?*anyopaque, _: abi.DevicePtr, len: usize) callconv(.c) abi.Result {
+    destination = @intFromPtr(dst.?);
+    if (pinnedRange(destination)) @memset(@as([*]u8, @ptrCast(dst.?))[0..len], 0x5a);
+    return 0;
+}
+
+test "a synchronous download lands in pinned memory, not the caller's pageable slice" {
+    var d: Driver = undefined;
+    d.api.cuMemHostAlloc = hostAlloc;
+    d.api.cuMemFreeHost = hostFree;
+    d.api.cuMemcpyDtoH_v2 = copyToHost;
+    stage.drop();
+    defer stage.drop();
+    var out: [64]u8 = @splat(0xaa);
+    const src: memory.DeviceBuffer = .{ .d = &d, .ptr = 0x1000, .len = out.len, .borrowed = true };
+    try src.download(0, &out);
+    try std.testing.expect(std.mem.allEqual(u8, &out, 0x5a));
+    try std.testing.expect(pinnedRange(destination));
+    try std.testing.expect(destination != @intFromPtr(&out));
+}
diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig
index b5faf7880e..a4b56c2947 100644
--- a/zig/src/cuda/root.zig
+++ b/zig/src/cuda/root.zig
@@ -89,6 +89,7 @@ test {
     _ = train_q4;
     _ = train_attn;
     _ = @import("memory.zig");
+    _ = @import("download_stage_test.zig");
     _ = memory_receipt;
     _ = @import("grouped.zig");
     _ = @import("qlinear.zig");

From d283418c57877496f04474481b5c35411ec27817 Mon Sep 17 00:00:00 2001
From: shuuul <6903405+shuuul@users.noreply.github.com>
Date: Thu, 5 Nov 2026 12:13:09 +0800
Subject: [PATCH 3/5] cuda: stage downloads only on integrated devices, and
 count the bounce

Discrete cards keep their direct DtoH copies and their speed. The bounce's
pinned bytes now count in memory.usage and are released with it. Comments
are one line each; tests cover both the integrated and discrete paths.
---
 zig/src/cuda/download_stage.zig      | 31 +++++++++++++-----
 zig/src/cuda/download_stage_test.zig | 48 +++++++++++++++++++++++++---
 zig/src/cuda/memory.zig              | 12 +++++--
 3 files changed, 77 insertions(+), 14 deletions(-)

diff --git a/zig/src/cuda/download_stage.zig b/zig/src/cuda/download_stage.zig
index 77a8f9efd8..532fbebd5a 100644
--- a/zig/src/cuda/download_stage.zig
+++ b/zig/src/cuda/download_stage.zig
@@ -1,22 +1,30 @@
-//! What every synchronous download lands in first: this thread's page-locked memory, at most `most`
-//! bytes, copied on into the caller's by the CPU. On GB10 a `cuMemcpyDtoH` straight into pageable
-//! memory (the host allocator's, freshly mapped) can miss pages with no CUDA error: some keep what
-//! they held, the bytes land on other pages of the process. Calls to `DeviceBuffer.download` stage
-//! through here; `downloadAsync` still needs pinned memory, as documented.
+//! A thread's page-locked bounce for synchronous downloads on integrated GPUs.
 const std = @import("std");
 const abi = @import("abi.zig");
 const Driver = @import("driver.zig").Driver;
 const Error = @import("driver.zig").Error;
+const memory = @import("memory.zig");
 
 /// The largest piece one copy takes; larger downloads go through in pieces.
 pub const most: usize = 4 << 20;
 
 const Held = struct { d: *const Driver, bytes: []align(16) u8 };
+const Device = struct { d: *const Driver, integrated: bool };
 threadlocal var held: ?Held = null;
+threadlocal var device: ?Device = null;
 
-/// This thread's page-locked buffer of at least `len` bytes, reallocated when a larger piece is asked
-/// for. The driver allocates it directly, so it is not counted in `memory.usage`: it is one per
-/// thread, bounded by `most`, and kept for that thread's next download.
+/// Whether this thread's current device shares memory with the host; read once per driver.
+pub fn integrated(d: *const Driver) Error!bool {
+    if (device) |one| if (one.d == d) return one.integrated;
+    var dev: abi.Device = 0;
+    try d.check(d.api.cuCtxGetDevice(&dev), "cuCtxGetDevice");
+    var v: c_int = 0;
+    try d.check(d.api.cuDeviceGetAttribute(&v, .integrated, dev), "cuDeviceGetAttribute");
+    device = .{ .d = d, .integrated = v != 0 };
+    return v != 0;
+}
+
+/// This thread's page-locked buffer of at least `len` bytes; it counts in `memory.usage`.
 pub fn take(d: *const Driver, len: usize) Error![]align(16) u8 {
     if (held) |h| {
         if (h.d == d and h.bytes.len >= len) return h.bytes[0..len];
@@ -27,6 +35,7 @@ pub fn take(d: *const Driver, len: usize) Error![]align(16) u8 {
     try d.check(d.api.cuMemHostAlloc(&p, size, abi.host_alloc_portable), "cuMemHostAlloc");
     const base: [*]align(16) u8 = @ptrCast(@alignCast(p.?));
     held = .{ .d = d, .bytes = base[0..size] };
+    memory.retainHost(size);
     return held.?.bytes[0..len];
 }
 
@@ -34,6 +43,12 @@ pub fn take(d: *const Driver, len: usize) Error![]align(16) u8 {
 pub fn drop() void {
     if (held) |h| {
         _ = h.d.api.cuMemFreeHost(h.bytes.ptr);
+        memory.releaseRetainedHost(h.bytes.len);
         held = null;
     }
 }
+
+/// Forgets the cached device reading (tests use distinct drivers, but a fake may reuse one).
+pub fn forget() void {
+    device = null;
+}
diff --git a/zig/src/cuda/download_stage_test.zig b/zig/src/cuda/download_stage_test.zig
index 188abb6795..a0c493e9e0 100644
--- a/zig/src/cuda/download_stage_test.zig
+++ b/zig/src/cuda/download_stage_test.zig
@@ -1,4 +1,4 @@
-//! The synchronous download's staging path, with a driver that only writes pinned destinations.
+//! The download staging path, with a fake driver that writes only pinned destinations.
 const std = @import("std");
 const abi = @import("abi.zig");
 const Driver = @import("driver.zig").Driver;
@@ -7,6 +7,8 @@ const stage = @import("download_stage.zig");
 
 var pinned: [64 << 10]u8 align(16) = undefined;
 var destination: usize = 0;
+var allocations: usize = 0;
+var integrated_flag: c_int = 1;
 
 fn pinnedRange(ptr: usize) bool {
     const base = @intFromPtr(&pinned[0]);
@@ -15,6 +17,7 @@ fn pinnedRange(ptr: usize) bool {
 
 fn hostAlloc(out: ?*?*anyopaque, len: usize, _: c_uint) callconv(.c) abi.Result {
     if (len > pinned.len) return 1;
+    allocations += 1;
     out.?.* = &pinned[0];
     return 0;
 }
@@ -23,24 +26,61 @@ fn hostFree(_: ?*anyopaque) callconv(.c) abi.Result {
     return 0;
 }
 
-/// Writes only a pinned destination, as the GB10 fault leaves pageable memory untouched.
+fn contextDevice(out: *abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
+
+fn deviceAttribute(out: *c_int, attribute: abi.DeviceAttribute, _: abi.Device) callconv(.c) abi.Result {
+    out.* = if (attribute == .integrated) integrated_flag else 0;
+    return 0;
+}
+
 fn copyToHost(dst: ?*anyopaque, _: abi.DevicePtr, len: usize) callconv(.c) abi.Result {
     destination = @intFromPtr(dst.?);
-    if (pinnedRange(destination)) @memset(@as([*]u8, @ptrCast(dst.?))[0..len], 0x5a);
+    if (pinnedRange(destination) or integrated_flag == 0) @memset(@as([*]u8, @ptrCast(dst.?))[0..len], 0x5a);
     return 0;
 }
 
-test "a synchronous download lands in pinned memory, not the caller's pageable slice" {
+fn driver() Driver {
     var d: Driver = undefined;
     d.api.cuMemHostAlloc = hostAlloc;
     d.api.cuMemFreeHost = hostFree;
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuMemcpyDtoH_v2 = copyToHost;
+    return d;
+}
+
+test "an integrated GPU's download lands in pinned memory, not the caller's pageable slice" {
+    integrated_flag = 1;
+    allocations = 0;
+    stage.forget();
     stage.drop();
     defer stage.drop();
+    const before = memory.usage(false).host;
+    var d = driver();
     var out: [64]u8 = @splat(0xaa);
     const src: memory.DeviceBuffer = .{ .d = &d, .ptr = 0x1000, .len = out.len, .borrowed = true };
     try src.download(0, &out);
     try std.testing.expect(std.mem.allEqual(u8, &out, 0x5a));
     try std.testing.expect(pinnedRange(destination));
     try std.testing.expect(destination != @intFromPtr(&out));
+    try std.testing.expectEqual(@as(usize, 1), allocations);
+    try std.testing.expect(memory.usage(false).host >= before + (64 << 10));
+}
+
+test "a discrete GPU's download copies straight into the caller's slice" {
+    integrated_flag = 0;
+    allocations = 0;
+    stage.forget();
+    stage.drop();
+    defer stage.drop();
+    var d = driver();
+    var out: [64]u8 = @splat(0xaa);
+    const src: memory.DeviceBuffer = .{ .d = &d, .ptr = 0x1000, .len = out.len, .borrowed = true };
+    try src.download(0, &out);
+    try std.testing.expect(std.mem.allEqual(u8, &out, 0x5a));
+    try std.testing.expectEqual(@intFromPtr(&out), destination);
+    try std.testing.expectEqual(@as(usize, 0), allocations);
 }
diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig
index 6e77c8dd88..2afea175b5 100644
--- a/zig/src/cuda/memory.zig
+++ b/zig/src/cuda/memory.zig
@@ -190,12 +190,15 @@ pub const DeviceBuffer = struct {
         try @import("copy_source.zig").synchronousCopy(self.d, dst, bytes, std.heap.page_allocator);
     }
 
-    /// Returns once the bytes have landed; the copy stages through this thread's page-locked buffer, a
-    /// piece of at most `download_stage.most` bytes at a time (see that module for why).
+    /// Returns once the bytes have landed; an integrated GPU stages through this thread's pinned bounce.
     pub fn download(self: DeviceBuffer, offset: usize, out: []u8) Error!void {
         const src = try self.span(offset, out.len);
         if (out.len == 0) return;
         const stage = @import("download_stage.zig");
+        if (!try stage.integrated(self.d)) {
+            try self.d.check(self.d.api.cuMemcpyDtoH_v2(out.ptr, src, out.len), "cuMemcpyDtoH");
+            return;
+        }
         var done: usize = 0;
         while (done < out.len) {
             const n = @min(out.len - done, stage.most);
@@ -554,6 +557,11 @@ pub fn releaseRetainedHost(bytes: usize) void {
     freed(&host_bytes, bytes);
 }
 
+/// Counts pinned bytes the driver allocated directly (a download's bounce), so `usage` and the limit see them.
+pub fn retainHost(bytes: usize) void {
+    held(&host_bytes, bytes);
+}
+
 test "Failure boundaries: a buffer keeps its allocation hook when the installed hook changes" {
     const Owner = struct {
         freed: usize = 0,

From 712b58f10d75340669efc60620b7ad0c90362b83 Mon Sep 17 00:00:00 2001
From: shuuul <6903405+shuuul@users.noreply.github.com>
Date: Thu, 5 Nov 2026 15:01:54 +0800
Subject: [PATCH 4/5] cuda: host fakes answer the integrated-device probe as
 discrete

download now asks cuCtxGetDevice and cuDeviceGetAttribute whether the
device is integrated. The fake driver tables that only set the entries
their test used left those two undefined, so the probe jumped into
garbage. Each now answers device 0 and attribute 0, which keeps the
direct copy path every host test expects.

zig build test passes.
---
 zig/src/families/deepseek_v41/load_mock.zig     | 10 ++++++++++
 zig/src/families/flashnext/cuda_pocket_test.zig | 10 ++++++++++
 zig/src/families/h3/cuda_runtime_fake.zig       | 10 ++++++++++
 zig/src/families/qwen3_5/cuda_pocket_test.zig   | 10 ++++++++++
 zig/tests/cuda/contract_q35_attn_host.zig       | 10 ++++++++++
 zig/tests/cuda/contract_rows_host.zig           | 10 ++++++++++
 zig/tests/cuda/contract_verify_host.zig         | 10 ++++++++++
 zig/tests/cuda/projection_f32_host_tests.zig    | 12 ++++++++++++
 zig/tests/cuda/quant_ref_host_tests.zig         | 12 ++++++++++++
 zig/tests/cuda/w4a8_exception_gate_host.zig     | 10 ++++++++++
 zig/tests/qwen4/cuda_draw_launch_test.zig       | 10 ++++++++++
 11 files changed, 114 insertions(+)

diff --git a/zig/src/families/deepseek_v41/load_mock.zig b/zig/src/families/deepseek_v41/load_mock.zig
index 40721ee84b..52dbb3a0ac 100644
--- a/zig/src/families/deepseek_v41/load_mock.zig
+++ b/zig/src/families/deepseek_v41/load_mock.zig
@@ -20,6 +20,14 @@ pub var fail_record = false;
 pub var live: [1024]bool = @splat(false);
 pub var memory: [1024][1024]u8 = @splat(@splat(0));
 pub var pinned: ?[]align(16) u8 = null;
+fn contextDevice(out: *abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: abi.DeviceAttribute, _: abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
 pub fn init() cuda.Driver {
     shared_calls = 0;
     shared_bytes = 0;
@@ -37,6 +45,8 @@ pub fn init() cuda.Driver {
     fail_record = false;
     @memset(&live, false);
     var d: cuda.Driver = undefined;
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuModuleLoadData = moduleLoad;
     d.api.cuModuleUnload = moduleUnload;
     d.api.cuModuleGetFunction = moduleFunction;
diff --git a/zig/src/families/flashnext/cuda_pocket_test.zig b/zig/src/families/flashnext/cuda_pocket_test.zig
index ceaff03bac..301381ceee 100644
--- a/zig/src/families/flashnext/cuda_pocket_test.zig
+++ b/zig/src/families/flashnext/cuda_pocket_test.zig
@@ -55,8 +55,18 @@ fn read(dst: ?*anyopaque, src: u64, n: usize) callconv(.c) abi.Result {
 fn sync(_: abi.Stream) callconv(.c) abi.Result {
     return if (fail_sync) abi.error_not_ready else 0;
 }
+fn contextDevice(out: *abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: abi.DeviceAttribute, _: abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
 fn driver() cuda.Driver {
     var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuMemAlloc_v2 = alloc;
     d.api.cuMemFree_v2 = free;
     d.api.cuMemHostAlloc = hostAlloc;
diff --git a/zig/src/families/h3/cuda_runtime_fake.zig b/zig/src/families/h3/cuda_runtime_fake.zig
index 24bdfe8a5b..3e008e78d1 100644
--- a/zig/src/families/h3/cuda_runtime_fake.zig
+++ b/zig/src/families/h3/cuda_runtime_fake.zig
@@ -26,6 +26,14 @@ fn sync() callconv(.c) c_int {
 fn streamSync(_: abi.Stream) callconv(.c) c_int {
     return sync();
 }
+fn contextDevice(out: *abi.Device) callconv(.c) c_int {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: abi.DeviceAttribute, _: abi.Device) callconv(.c) c_int {
+    out.* = 0;
+    return 0;
+}
 fn alloc(out: *u64, n: usize) callconv(.c) c_int {
     if (alloc_fail) return abi.error_not_ready;
     at = (at + 255) / 256 * 256;
@@ -103,6 +111,8 @@ pub fn model() !*h3.Model {
     m.gpa = std.testing.allocator;
     m.driver.lib = try std.DynLib.open(if (@import("builtin").os.tag == .macos) "/usr/lib/libSystem.B.dylib" else "libc.so.6");
     m.driver.api.cuCtxSetCurrent = current;
+    m.driver.api.cuCtxGetDevice = contextDevice;
+    m.driver.api.cuDeviceGetAttribute = deviceAttribute;
     m.driver.api.cuCtxSynchronize = sync;
     m.driver.api.cuStreamSynchronize = streamSync;
     m.driver.api.cuMemAlloc_v2 = alloc;
diff --git a/zig/src/families/qwen3_5/cuda_pocket_test.zig b/zig/src/families/qwen3_5/cuda_pocket_test.zig
index 7ebc99eba0..e22980fd2a 100644
--- a/zig/src/families/qwen3_5/cuda_pocket_test.zig
+++ b/zig/src/families/qwen3_5/cuda_pocket_test.zig
@@ -55,8 +55,18 @@ fn read(dst: ?*anyopaque, src: u64, n: usize) callconv(.c) abi.Result {
 fn sync(_: abi.Stream) callconv(.c) abi.Result {
     return if (fail_sync) abi.error_not_ready else 0;
 }
+fn contextDevice(out: *abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: abi.DeviceAttribute, _: abi.Device) callconv(.c) abi.Result {
+    out.* = 0;
+    return 0;
+}
 fn driver() cuda.Driver {
     var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuMemAlloc_v2 = alloc;
     d.api.cuMemFree_v2 = free;
     d.api.cuMemHostAlloc = hostAlloc;
diff --git a/zig/tests/cuda/contract_q35_attn_host.zig b/zig/tests/cuda/contract_q35_attn_host.zig
index 79dea869a0..75286faf32 100644
--- a/zig/tests/cuda/contract_q35_attn_host.zig
+++ b/zig/tests/cuda/contract_q35_attn_host.zig
@@ -68,8 +68,18 @@ const Trace = struct {
         return @as(*const V, @ptrCast(@alignCast(a.?[i].?))).*;
     }
 };
+fn contextDevice(out: *cuda.abi.Device) callconv(.c) cuda.abi.Result {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result {
+    out.* = 0;
+    return 0;
+}
 fn driver() cuda.Driver {
     var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuModuleGetFunction = Trace.function;
     d.api.cuFuncSetAttribute = Trace.attribute;
     d.api.cuLaunchKernel = Trace.launch;
diff --git a/zig/tests/cuda/contract_rows_host.zig b/zig/tests/cuda/contract_rows_host.zig
index 32921a7aaa..c97b9559d9 100644
--- a/zig/tests/cuda/contract_rows_host.zig
+++ b/zig/tests/cuda/contract_rows_host.zig
@@ -25,7 +25,17 @@ fn copy(dst: ?*anyopaque, src: u64, n: usize) callconv(.c) cuda.abi.Result {
     @memcpy(out[0..n], bytes[0..n]);
     return 0;
 }
+fn contextDevice(out: *cuda.abi.Device) callconv(.c) cuda.abi.Result {
+    out.* = 0;
+    return 0;
+}
+fn deviceAttribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result {
+    out.* = 0;
+    return 0;
+}
 fn output(d: *cuda.Driver, bytes: []u8) !rig.Rig {
+    d.api.cuCtxGetDevice = contextDevice;
+    d.api.cuDeviceGetAttribute = deviceAttribute;
     d.api.cuStreamSynchronize = sync;
     d.api.cuMemcpyDtoH_v2 = copy;
     var r: rig.Rig = .{ .gpu = .{ .d = d, .ctx = undefined, .gpa = std.testing.allocator, .io = undefined }, .s = .{ .d = d, .handle = null } };
diff --git a/zig/tests/cuda/contract_verify_host.zig b/zig/tests/cuda/contract_verify_host.zig
index c24dcc805a..daa4ed508e 100644
--- a/zig/tests/cuda/contract_verify_host.zig
+++ b/zig/tests/cuda/contract_verify_host.zig
@@ -165,10 +165,20 @@ const Baseline = struct {
         @memcpy(out[0..len], p[0..len]);
         return 0;
     }
+    fn contextDevice(out: *cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return 0;
+    }
+    fn deviceAttribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return 0;
+    }
 };
 
 test "verify contract: the actual replacement producer really runs 32 separate original one-row calls" {
     var d: cuda.Driver = undefined;
+    d.api.cuCtxGetDevice = Baseline.contextDevice;
+    d.api.cuDeviceGetAttribute = Baseline.deviceAttribute;
     d.api.cuLaunchKernel = Baseline.launch;
     d.api.cuStreamSynchronize = Baseline.sync;
     d.api.cuMemcpyDtoH_v2 = Baseline.copy;
diff --git a/zig/tests/cuda/projection_f32_host_tests.zig b/zig/tests/cuda/projection_f32_host_tests.zig
index 402ec4efff..a99d22dead 100644
--- a/zig/tests/cuda/projection_f32_host_tests.zig
+++ b/zig/tests/cuda/projection_f32_host_tests.zig
@@ -43,11 +43,23 @@ const Mock = struct {
         @memcpy(out[0..n], in[0..n]);
         return cuda.abi.success;
     }
+
+    fn contextDevice(out: *cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return cuda.abi.success;
+    }
+
+    fn deviceAttribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return cuda.abi.success;
+    }
 };
 
 test "every F32 and BF16 fixture dispatch rejects a no-op after a valid stale prefix" {
     const a = std.testing.allocator;
     var driver: cuda.Driver = undefined;
+    driver.api.cuCtxGetDevice = Mock.contextDevice;
+    driver.api.cuDeviceGetAttribute = Mock.deviceAttribute;
     driver.api.cuMemsetD8Async = Mock.fill;
     driver.api.cuLaunchKernel = Mock.plain;
     driver.api.cuLaunchKernelEx = Mock.extended;
diff --git a/zig/tests/cuda/quant_ref_host_tests.zig b/zig/tests/cuda/quant_ref_host_tests.zig
index eb72d34e94..aaecf1afbd 100644
--- a/zig/tests/cuda/quant_ref_host_tests.zig
+++ b/zig/tests/cuda/quant_ref_host_tests.zig
@@ -19,6 +19,16 @@ const Readback = struct {
         copies += 1;
         return cuda.abi.success;
     }
+
+    fn contextDevice(out: *cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return cuda.abi.success;
+    }
+
+    fn deviceAttribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result {
+        out.* = 0;
+        return cuda.abi.success;
+    }
 };
 
 const Closed = struct { failed: bool, last: Rig.Last };
@@ -26,6 +36,8 @@ const Closed = struct { failed: bool, last: Rig.Last };
 /// One quiet close against zero within one: whether it failed its gate, and the element it reported.
 fn close(comptime T: type, got: []const T) !Closed {
     var d: cuda.Driver = undefined;
+    d.api.cuCtxGetDevice = Readback.contextDevice;
+    d.api.cuDeviceGetAttribute = Readback.deviceAttribute;
     d.api.cuStreamSynchronize = Readback.sync;
     d.api.cuMemcpyDtoH_v2 = Readback.copy;
     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
diff --git a/zig/tests/cuda/w4a8_exception_gate_host.zig b/zig/tests/cuda/w4a8_exception_gate_host.zig
index 97a49c25fc..78d7d5b458 100644
--- a/zig/tests/cuda/w4a8_exception_gate_host.zig
+++ b/zig/tests/cuda/w4a8_exception_gate_host.zig
@@ -41,6 +41,14 @@ const Fake = struct {
     fn sync(_: abi.Stream) callconv(.c) abi.Result {
         return 0;
     }
+    fn contextDevice(out: *abi.Device) callconv(.c) abi.Result {
+        out.* = 0;
+        return 0;
+    }
+    fn deviceAttribute(out: *c_int, _: abi.DeviceAttribute, _: abi.Device) callconv(.c) abi.Result {
+        out.* = 0;
+        return 0;
+    }
     fn read(comptime T: type, args: ?[*]?*anyopaque, i: usize) T {
         return @as(*const T, @ptrCast(@alignCast(args.?[i].?))).*;
     }
@@ -59,6 +67,8 @@ const Fake = struct {
     }
     fn driver() cuda.Driver {
         var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+        d.api.cuCtxGetDevice = contextDevice;
+        d.api.cuDeviceGetAttribute = deviceAttribute;
         d.api.cuMemAlloc_v2 = alloc;
         d.api.cuMemFree_v2 = free;
         d.api.cuMemcpyHtoD_v2 = upload;
diff --git a/zig/tests/qwen4/cuda_draw_launch_test.zig b/zig/tests/qwen4/cuda_draw_launch_test.zig
index 2a51dc534a..76595c6f41 100644
--- a/zig/tests/qwen4/cuda_draw_launch_test.zig
+++ b/zig/tests/qwen4/cuda_draw_launch_test.zig
@@ -71,9 +71,19 @@ const Trace = struct {
     fn synchronize(_: C.abi.Stream) callconv(.c) c_int {
         return 0;
     }
+    fn contextDevice(out: *C.abi.Device) callconv(.c) c_int {
+        out.* = 0;
+        return 0;
+    }
+    fn deviceAttribute(out: *c_int, _: C.abi.DeviceAttribute, _: C.abi.Device) callconv(.c) c_int {
+        out.* = 0;
+        return 0;
+    }
 };
 fn driver() C.Driver {
     var d: C.Driver = .{ .api = undefined, .lib = undefined };
+    d.api.cuCtxGetDevice = Trace.contextDevice;
+    d.api.cuDeviceGetAttribute = Trace.deviceAttribute;
     d.api.cuModuleLoadData = Trace.load;
     d.api.cuModuleUnload = Trace.unload;
     d.api.cuModuleGetFunction = Trace.function;

From 5e0abe04581732bf1ddcd685e0a0ace70aaa86d9 Mon Sep 17 00:00:00 2001
From: shuuul <6903405+shuuul@users.noreply.github.com>
Date: Thu, 5 Nov 2026 18:46:33 +0800
Subject: [PATCH 5/5] cuda: forget the staged download's cached device after
 its tests

The integrated test cached its reading by the fake driver's address;
a later test whose fake lands there would reuse it and call
cuMemHostAlloc on a table that does not set it. Both tests clear it.
---
 zig/src/cuda/download_stage_test.zig | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/zig/src/cuda/download_stage_test.zig b/zig/src/cuda/download_stage_test.zig
index a0c493e9e0..0cdb01f464 100644
--- a/zig/src/cuda/download_stage_test.zig
+++ b/zig/src/cuda/download_stage_test.zig
@@ -57,6 +57,7 @@ test "an integrated GPU's download lands in pinned memory, not the caller's page
     allocations = 0;
     stage.forget();
     stage.drop();
+    defer stage.forget();
     defer stage.drop();
     const before = memory.usage(false).host;
     var d = driver();
@@ -75,6 +76,7 @@ test "a discrete GPU's download copies straight into the caller's slice" {
     allocations = 0;
     stage.forget();
     stage.drop();
+    defer stage.forget();
     defer stage.drop();
     var d = driver();
     var out: [64]u8 = @splat(0xaa);