From 549653b1ba1ca056cf7e83093e76350001934804 Mon Sep 17 00:00:00 2001
From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com>
Date: Tue, 13 Oct 2026 14:47:01 -0400
Subject: [PATCH] qwen27 CUDA: weight tensors packed into 128 MiB blocks, not
 one cuMemAlloc each

The driver rounds each allocation up to whole 2 MiB pages, which costs ~1 GiB over the 27B's ~900 large tensors.
Both loaders (the 27B's and DFlash2's) now take tensors first fit out of 128 MiB blocks. The bytes and math are
unchanged; only the addresses move.
---
 zig/src/families/qwen27/cuda/draft_load.zig |  6 ++--
 zig/src/families/qwen27/cuda/weights.zig    | 36 +++++++++++++++++++--
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/zig/src/families/qwen27/cuda/draft_load.zig b/zig/src/families/qwen27/cuda/draft_load.zig
index 8f9b0a7404..7fd577b43f 100644
--- a/zig/src/families/qwen27/cuda/draft_load.zig
+++ b/zig/src/families/qwen27/cuda/draft_load.zig
@@ -114,11 +114,10 @@ const Loader = struct {
     scales: std.ArrayList(u16) = .empty,
     biases: std.ArrayList(u16) = .empty,
     staging: ?cuda.DeviceBuffer = null,
+    pool: @import("weights.zig").Pool = .{},
 
     fn alloc(L: *Loader, len: usize) !u64 {
-        const b = try cuda.DeviceBuffer.alloc(L.ops.k.d, len);
-        try L.w.buffers.append(L.gpa, b);
-        return b.ptr;
+        return L.pool.take(L.gpa, L.ops.k.d, &L.w.buffers, len);
     }
 
     fn raw(L: *Loader, t: Tensor) !u64 {
@@ -206,6 +205,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8,
         L.scales.deinit(gpa);
         L.biases.deinit(gpa);
         if (L.staging) |*b| b.free();
+        L.pool.deinit(gpa);
     }
     var ck = try core.Checkpoint.openModel(gpa, io, dir);
     defer ck.close();
diff --git a/zig/src/families/qwen27/cuda/weights.zig b/zig/src/families/qwen27/cuda/weights.zig
index 5c796bd43a..8f502405be 100644
--- a/zig/src/families/qwen27/cuda/weights.zig
+++ b/zig/src/families/qwen27/cuda/weights.zig
@@ -51,6 +51,36 @@ pub const Weights = struct {
     }
 };
 
+/// Weight tensors packed first fit into 128 MiB blocks: a cuMemAlloc each rounds a tensor up to whole 2 MiB pages.
+pub const Pool = struct {
+    pub const block: usize = 128 << 20;
+    const Gap = struct { at: u64, left: usize };
+    gaps: std.ArrayList(Gap) = .empty,
+
+    /// `len` device bytes, 256-byte aligned; the blocks go on `buffers`, which own them.
+    pub fn take(p: *Pool, gpa: std.mem.Allocator, d: *const cuda.Driver, buffers: *std.ArrayList(cuda.DeviceBuffer), len: usize) !u64 {
+        const n = std.mem.alignForward(usize, len, 256);
+        if (n == 0) return 0;
+        for (p.gaps.items) |*g| if (g.left >= n) {
+            const at = g.at;
+            g.at += n;
+            g.left -= n;
+            return at;
+        };
+        var b = try cuda.DeviceBuffer.alloc(d, @max(n, block));
+        {
+            errdefer b.free();
+            try buffers.append(gpa, b);
+        }
+        if (n < block) try p.gaps.append(gpa, .{ .at = b.ptr + n, .left = block - n });
+        return b.ptr;
+    }
+
+    pub fn deinit(p: *Pool, gpa: std.mem.Allocator) void {
+        p.gaps.deinit(gpa);
+    }
+};
+
 const Loader = struct {
     gpa: std.mem.Allocator,
     ops: kern.Ops,
@@ -58,11 +88,10 @@ const Loader = struct {
     src: *source.Source,
     scratch: cuda.DeviceBuffer,
     host: std.ArrayList(u8) = .empty,
+    pool: Pool = .{},
 
     fn alloc(L: *Loader, len: usize) !u64 {
-        const b = try cuda.DeviceBuffer.alloc(L.ops.k.d, len);
-        try L.w.buffers.append(L.gpa, b);
-        return b.ptr;
+        return L.pool.take(L.gpa, L.ops.k.d, &L.w.buffers, len);
     }
 
     /// Device scratch of at least `len` bytes; growing waits for the stream so pending packs keep their inputs.
@@ -204,6 +233,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8,
         ops.s.synchronize() catch {};
         L.scratch.free();
         L.host.deinit(gpa);
+        L.pool.deinit(gpa);
     }
     // the vision tower and the MTP head stay on disk, as weights._Tensors skips them
     var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, prefix);