qwen CUDA: streams on pages by default (256 positions, the engine's own too), with the paged prompt and decode at the window path's speed

CHANGELOG.md

@@ -4,6 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- Qwen3.5 and Qwen3.6 on CUDA: KV pool pages stay opt-in (`TF_QWEN_KV_PAGES=256`), dense and MoE alike; unset keeps a window each, as before, and `TF_QWEN_KV_PAGES_ALL` unset still leaves the engine's own stream a window. With pages set, `--parallel` is a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011). Thanks to @jschmied.
 - Dory on CUDA takes its RoPE tables from core's exact helper, so they match transformers' words (they moved by up to 5 ulp) (#2035, part of #1852).
 
 ## 1.2.2

zig/build/cuda_family_host.zig

@@ -34,6 +34,11 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B
     const attention = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen attention path:"} }));
     b.step("test-qwen35-attention-path-host", "Contiguous and paged CUDA attention selection and graph identity without devices").dependOn(&attention.step);
     step.dependOn(&attention.step);
+    const pages_filter = &.{"Qwen pages default:"};
+    const pages_default = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = pages_filter }));
+    const pages_about = "TF_QWEN_KV_PAGES and _ALL unset as on main, --parallel a cap only with pages";
+    b.step("test-qwen35-pages-default-host", pages_about).dependOn(&pages_default.step);
+    step.dependOn(&pages_default.step);
     const prompt_layout = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen prompt layout:"} }));
     b.step("test-qwen35-prompt-layout-host", "Actual whole and paged CUDA prompt ABI and stride calculations without devices").dependOn(&prompt_layout.step);
     step.dependOn(&prompt_layout.step);

zig/kernels/cuda/prefill_attention.cu

@@ -62,7 +62,9 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric
              int G, float scale, const int* __restrict__ PT, long long PS) {
     constexpr bool causal = true;
     constexpr int skip = 0;
-    constexpr bool HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int CAP = 0, PTOK = 1;
     const float *ks = nullptr, *vs = nullptr;
     const int HO = H;
@@ -80,7 +82,9 @@ pattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict_
     constexpr const int* PT = nullptr;  // contiguous keys
     constexpr long long PS = 0;
     constexpr int PTOK = 1;
-    constexpr bool HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int CAP = 0;
     const float *ks = nullptr, *vs = nullptr;
     const int HO = H;
@@ -115,7 +119,9 @@ full_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ o
     constexpr const int* PT = nullptr;  // contiguous keys
     constexpr long long PS = 0;
     constexpr int PTOK = 1;
-    constexpr bool HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false;  // bf16 keys, row-major
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int CAP = 0;
     const float *ks = nullptr, *vs = nullptr;
     const int H = 3 * HO, HK = H;
@@ -132,7 +138,9 @@ __global__ void __launch_bounds__(32 * WARPS, 1)
 hpattn_whole_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
               const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK,
               int G, float scale, int CAP) {
-    constexpr bool HPAGED = false, causal = true, F8 = false, HM = true;
+    constexpr bool P256 = false, HPAGED = false, causal = true, F8 = false, HM = true;
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int skip = 0;
     constexpr const int* PT = nullptr;
     constexpr long long PS = 0;
@@ -149,7 +157,9 @@ __global__ void __launch_bounds__(32 * WARPS, 1)
 hpattn8_whole_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restrict__ kc, const uint8_t* __restrict__ vc,
                const float* __restrict__ ks, const float* __restrict__ vs, __nv_bfloat16* __restrict__ out, int p0,
                int W, int H, int HK, int G, float scale, int CAP) {
-    constexpr bool HPAGED = false, causal = true, F8 = true, HM = true;
+    constexpr bool P256 = false, HPAGED = false, causal = true, F8 = true, HM = true;
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int skip = 0;
     constexpr const int* PT = nullptr;
     constexpr long long PS = 0;
@@ -165,21 +175,56 @@ __global__ void __launch_bounds__(32 * WARPS, 1)
 hpattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
               const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK,
               int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PTOK) {
-    constexpr bool HPAGED = true, causal = true, F8 = false, HM = true;
+    constexpr bool P256 = false, HPAGED = true, causal = true, F8 = false, HM = true;
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int skip = 0;
     const int HO = H;
     const float *ks = nullptr, *vs = nullptr;
     auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); };
 #include "prefill_attention_rows.cuh"
 }
 
+// hpattn_kernel over pages of 256 positions (TF_QWEN_KV_PAGES's default): an item's keys sit in one page.
+template <int D, int WARPS, int HPC, int NS>
+__global__ void __launch_bounds__(32 * WARPS, 1)
+hpattn_p256_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
+                   const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK,
+                   int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int) {
+    constexpr bool P256 = true, HPAGED = true, causal = true, F8 = false, HM = true, IDS = false;
+    constexpr int PN = 0;
+    constexpr int skip = 0, PTOK = 256;
+    static_assert(PTOK % BN == 0, "a tile's keys, filled or not, in one page");
+    const int HO = H;
+    const float *ks = nullptr, *vs = nullptr;
+    auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); };
+#include "prefill_attention_rows.cuh"
+}
+
+// hpattn_p256_kernel with the chunk's PN page ids in shared memory past the queries (PN * 4 more dynamic bytes).
+template <int D, int WARPS, int HPC, int NS>
+__global__ void __launch_bounds__(32 * WARPS, 1)
+hpattn_p256_ids_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
+                       const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK,
+                       int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PN) {
+    constexpr bool P256 = true, HPAGED = true, causal = true, F8 = false, HM = true, IDS = true;
+    constexpr int skip = 0, PTOK = 256;
+    static_assert(PTOK % BN == 0, "a tile's keys, filled or not, in one page");
+    const int HO = H;
+    const float *ks = nullptr, *vs = nullptr;
+    auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); };
+#include "prefill_attention_rows.cuh"
+}
+
 // hpattn_kernel over e4m3 rows with an fp32 scale a (key head, row) (qwen35_kv8.cu's layout, PS in bytes).
 template <int D, int WARPS, int HPC, int NS>
 __global__ void __launch_bounds__(32 * WARPS, 1)
 hpattn8_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restrict__ kc, const uint8_t* __restrict__ vc,
                const float* __restrict__ ks, const float* __restrict__ vs, __nv_bfloat16* __restrict__ out, int p0,
                int W, int H, int HK, int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PTOK) {
-    constexpr bool HPAGED = true, causal = true, F8 = true, HM = true;
+    constexpr bool P256 = false, HPAGED = true, causal = true, F8 = true, HM = true;
+    constexpr bool IDS = false;  // no page ids in shared memory
+    constexpr int PN = 0;
     constexpr int skip = 0;
     const int HO = H;
     auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); };
@@ -337,6 +382,12 @@ template __global__ void tf_prefill_attention::vision_kernel<128, 8>(const __nv_
 template __global__ void tf_prefill_attention::hpattn_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__,
     const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int,
     int, float, int, const int* __restrict__, long long, int);
+template __global__ void tf_prefill_attention::hpattn_p256_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__,
+    const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int,
+    int, float, int, const int* __restrict__, long long, int);
+template __global__ void tf_prefill_attention::hpattn_p256_ids_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__,
+    const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int,
+    int, float, int, const int* __restrict__, long long, int);
 template __global__ void tf_prefill_attention::hpattn8_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__,
     const uint8_t* __restrict__, const uint8_t* __restrict__, const float* __restrict__, const float* __restrict__,
     __nv_bfloat16* __restrict__, int, int, int, int, int, float, int, const int* __restrict__, long long, int);

zig/kernels/cuda/prefill_attention_rows.cuh

@@ -5,6 +5,7 @@
     extern __shared__ uint4 smem[];
     uint4* slots = smem;
     uint4* qs = smem + NS * HALF * CS;
+    int* ids = reinterpret_cast<int*>(qs + WARPS * 16 * CS);  // IDS: the chunk's PN page ids, past the queries
     const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g = lane / 4, qd = lane % 4, mt = lane / 8;
     const int groups = G / HPC, hk = blockIdx.y / groups, hg = blockIdx.y % groups;
     const int row0 = first_row(), last = min(row0 + 16 * RB, W) - 1;
@@ -26,6 +27,9 @@
             const uint8_t* src = reinterpret_cast<const uint8_t*>(item % 4 < 2 ? kc : vc);
             const float* sc = item % 4 < 2 ? ks : vs;
             const int key0 = (item / 4) * BN + (item % 2) * HALF;
+            // HM with PT: the item's first page looked up once, as stage does
+            const int lo = HM && PT ? min(key0 + skip, limit - 1 + skip) / PTOK * PTOK : 0;
+            const int64_t first = HM && PT ? (int64_t)PT[lo / PTOK] * PS : 0;
 #pragma unroll
             for (int j = 0; j < PER; ++j) {
                 const int i = threadIdx.x + j * THREADS, key = key0 + i / C;
@@ -34,7 +38,8 @@
                 if (i < HALF * C && key < limit) {
                     if constexpr (HPAGED) {
                         const int row = key + (key < p0 ? 0 : skip);
-                        const int64_t r = row_of(row), pg = page_of(row);
+                        const bool near = HM && PT && row >= lo && row - lo < PTOK;
+                        const int64_t r = near ? (int64_t)hk * PTOK + (row - lo) : row_of(row), pg = near ? first : page_of(row);
                         pre[j] = __ldg(reinterpret_cast<const uint2*>(src + pg + r * D + (i % C) * 8));
                         pre_s[j] = __ldg(reinterpret_cast<const float*>(reinterpret_cast<const uint8_t*>(sc) + pg) + r);
                     } else {
@@ -69,17 +74,33 @@
         if constexpr (F8) {
             fetch(item);
             put(item);
+        } else if constexpr (P256) {
+            // pages of 256 positions: a tile's keys, filled or not, sit in one allocated page, so key0 needs no clamp
+            const __nv_bfloat16* src = reinterpret_cast<const __nv_bfloat16*>(item % 4 < 2 ? kc : vc);
+            const int key0 = (item / 4) * BN + (item % 2) * HALF, p = key0 / PTOK;
+            uint4* dst = slots + (item % NS) * HALF * CS;
+            const __nv_bfloat16* base = src + (int64_t)(IDS ? ids[p] : PT[p]) * PS + ((int64_t)hk * PTOK + key0 % PTOK) * D;
+            for (int i = threadIdx.x; i < HALF * NB; i += THREADS) {
+                const int row = i / NB, chunk = i % NB;
+                cp16z(dst + sw<8 * CS>(i / NB, chunk), base + (i / NB) * D + chunk * 8, key0 + row < limit && (NB == C || chunk < C));
+            }
         } else {
             const __nv_bfloat16* src = reinterpret_cast<const __nv_bfloat16*>(item % 4 < 2 ? kc : vc);
             const int key0 = (item / 4) * BN + (item % 2) * HALF;
             uint4* dst = slots + (item % NS) * HALF * CS;
+            // HM with PT: the item's first page looked up once; a row past it (PTOK % HALF != 0) finds its own
+            const int lo = HM && PT ? min(key0 + skip, limit - 1 + skip) / PTOK * PTOK : 0;
+            const int64_t first = HM && PT ? (int64_t)PT[lo / PTOK] * PS : 0;
             for (int i = threadIdx.x; i < HALF * NB; i += THREADS) {
                 const int key = key0 + i / NB, chunk = i % NB;
                 const bool in = key < limit && (NB == C || chunk < C);  // past the filled keys or the head dim: zeros
                 const int row = in ? key + (key < p0 ? 0 : skip) : 0;
                 if constexpr (HM) {
-                    if constexpr (HPAGED) cp16z(dst + sw<8 * CS>(i / NB, chunk), src + page_of(row) + row_of(row) * D + chunk * 8, in);
-                    else cp16z(dst + sw<8 * CS>(i / NB, chunk), src + row_of(row) * D + chunk * 8, in);
+                    if constexpr (HPAGED) {
+                        const bool near = PT && row >= lo && row - lo < PTOK;
+                        const int64_t at = near ? first + ((int64_t)hk * PTOK + (row - lo)) * D : page_of(row) + row_of(row) * D;
+                        cp16z(dst + sw<8 * CS>(i / NB, chunk), src + at + chunk * 8, in);
+                    } else cp16z(dst + sw<8 * CS>(i / NB, chunk), src + row_of(row) * D + chunk * 8, in);
                 } else {
                     // a page's rows are contiguous: PT names each 512-row page, PS elements apart (null PT: one run)
                     const int64_t at = PT ? (int64_t)PT[row / 512] * PS + (int64_t)(row % 512) * HK * D : (int64_t)row * HK * D;
@@ -108,6 +129,10 @@
         }
         commit();
     }
+    if constexpr (IDS) {  // every page id the chunk reads, once, so staging never waits on the table
+        for (int i = threadIdx.x; i < PN; i += THREADS) ids[i] = PT[i];
+        __syncthreads();
+    }
 #pragma unroll
     for (int i = 0; i < NS; ++i) {
         if (i < items) stage(i);

zig/kernels/cuda/qwen35_attn_tiles.cuh

@@ -15,6 +15,9 @@ struct Fp8Tiles {
     size_t lblk, hoff, vgap, sgap;  // a page's layer block (bytes), the head's first row in it, keys to values, codes to scales
     int ptok, kvh;
     bool paged;
+    // the page a thread read last, kept across tiles: its first row is looked up once a page, not once a tile
+    const unsigned char* near = nullptr;
+    int lo = 0;
     __device__ Fp8Tiles(TileStore store, const Seg& sg, int layer, int kvh_, int h) : s(store) {
         const unsigned char* slab = reinterpret_cast<const unsigned char*>(sg.kv) + (size_t)layer * 2 * kvh_ * sg.cap * (HD + 4);
         K = slab + (size_t)h * sg.cap * HD;
@@ -45,8 +48,18 @@ struct Fp8Tiles {
             const int r = row(t0 + (tid >> 4) + THREADS / 16 * i);
             kr[i] = vr[i] = make_uint4(0u, 0u, 0u, 0u);
             if (r >= 0) {
-                kr[i] = *reinterpret_cast<const uint4*>(kp(r) + part * 16);
-                vr[i] = *reinterpret_cast<const uint4*>(vp(r) + part * 16);
+                const unsigned char* k = K + (size_t)r * HD;
+                const unsigned char* v = V + (size_t)r * HD;
+                if (paged) {
+                    if (near == nullptr || r < lo || r - lo >= ptok) {
+                        lo = r / ptok * ptok;
+                        near = arena + (size_t)page.ids[lo / ptok] * page.stride + lblk + hoff * HD;
+                    }
+                    k = near + (size_t)(r - lo) * HD;
+                    v = k + vgap;
+                }
+                kr[i] = *reinterpret_cast<const uint4*>(k + part * 16);
+                vr[i] = *reinterpret_cast<const uint4*>(v + part * 16);
             }
         }
         const int r = row(t0 + (tid & 31));
@@ -85,6 +98,9 @@ struct Bf16Tiles {
     size_t lblk, hoff, vgap;  // a page's layer block (elements), the head's first row in it, keys to values
     int ptok;
     bool paged;
+    // the page a thread read last, kept across tiles: its first row is looked up once a page, not once a tile
+    const __nv_bfloat16* near = nullptr;
+    int lo = 0;
     __device__ Bf16Tiles(TileStore store, const Seg& sg, int layer, int kvh, int h) : s(store) {
         K = reinterpret_cast<const __nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * HD + (size_t)h * sg.cap * HD;
         V = K + (size_t)kvh * sg.cap * HD;
@@ -106,8 +122,18 @@ struct Bf16Tiles {
         for (int i = 0; i < TILE / THREADS; i++) {
             const int at = threadIdx.x + THREADS * i, u = at / CH, chunk = at % CH;
             const int r = row(t0 + u), pos = r < 0 ? 0 : r;  // a key past the chunk reads row 0 and lands zeros
-            qmm_frag::cp16z(dst + swz(u, chunk), kp(pos) + chunk * 8, r >= 0);
-            qmm_frag::cp16z(dst + TILE + swz(u, chunk), vp(pos) + chunk * 8, r >= 0);
+            const __nv_bfloat16* k = K + (size_t)pos * HD;
+            const __nv_bfloat16* v = V + (size_t)pos * HD;
+            if (paged) {
+                if (near == nullptr || pos < lo || pos - lo >= ptok) {
+                    lo = pos / ptok * ptok;
+                    near = arena + (size_t)page.ids[lo / ptok] * page.stride + lblk + hoff * HD;
+                }
+                k = near + (size_t)(pos - lo) * HD;
+                v = k + vgap;
+            }
+            qmm_frag::cp16z(dst + swz(u, chunk), k + chunk * 8, r >= 0);
+            qmm_frag::cp16z(dst + TILE + swz(u, chunk), v + chunk * 8, r >= 0);
         }
         qmm_frag::commit();
     }