qwen CUDA: a faster opt-in paged path (TF_QWEN_KV_PAGES=256), --parallel a cap with pages set

CHANGELOG.md

@@ -113,6 +113,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 - GLM-5.3-Flash on CUDA: an EXL3 checkpoint whose `quantization_config` sits inside `text_config` loads as EXL3. Admission already read it there, but the config parser read only the top level, so such a pack was admitted and then asked for NVFP4 dense weights at load (#2327).
 
 - Docs: `docs/hermes-agent.md` shows how to use TensorFold from Hermes Agent through the community [hermes-tensorfold](https://github.com/atakann/hermes-tensorfold) plugin (a model provider and a dashboard tab), checked on a GB10 with Qwen3.8-Flash-Next (#2330, part of #2109). Thanks to @atakann.
+- Qwen3.8-27B and the other dense Qwen3.5 models on CUDA: streams run on KV pool pages of 256 positions by default, the engine's own stream too, so `--parallel` is now a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. `TF_QWEN_KV_PAGES=0` keeps a window each, as before. Qwen3.6-35B-A3B and the other MoE checkpoints stay on windows by default, with pages opt-in (`TF_QWEN_KV_PAGES=256`). 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.
+- 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.
+- 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, part of #1943). Thanks to @jschmied.
 
 - Metal: a request's finished stats list its prompt pass widths (`tensorfold.prefill_widths`, one entry per completed
   target pass) on every Metal text family except Flash Next, whose engine reports none. Telemetry only (#2293).

zig/build/cuda_family_host.zig

@@ -41,9 +41,17 @@ 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);
+    const reuse_pages = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen paged reuse:"} }));
+    b.step("test-qwen35-reuse-pages-host", "Paged prompt-state saves and restores into host or device copies without devices").dependOn(&reuse_pages.step);
+    step.dependOn(&reuse_pages.step);
     const pending = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen pending keeps:"} }));
     pending.setEnvironmentVariable("TF_QWEN_KEEP_LATE", "1"); // deferred keeps are opt-in; the step tests them
     b.step("test-qwen35-pending-host", "Kept paths of every length inside the pending buffer and the log without devices").dependOn(&pending.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);