From 11bbcefae8b0e8a339b78884d74eca44f05a6b1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 2 Nov 2026 17:41:44 +0100 Subject: [PATCH 01/16] cuda prefill_attention: a paged item's page looked up once, not every 16-byte chunk hpattn_kernel and hpattn8_kernel staged each chunk through PT[row / PTOK], a runtime divide and a dependent table load per chunk. An item's 32 keys sit in one page whenever PTOK is a multiple of 32, so its first page is read once and a row's offset is row - lo; a row past that page (PTOK 48) still looks up its own. Addresses only: the same bytes land in the same slots. --- zig/kernels/cuda/prefill_attention_rows.cuh | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index 53f2e0f0ab..03fcd45c40 100644 --- a/zig/kernels/cuda/prefill_attention_rows.cuh +++ b/zig/kernels/cuda/prefill_attention_rows.cuh @@ -26,6 +26,9 @@ const uint8_t* src = reinterpret_cast(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 +37,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(src + pg + r * D + (i % C) * 8)); pre_s[j] = __ldg(reinterpret_cast(reinterpret_cast(sc) + pg) + r); } else { @@ -73,13 +77,19 @@ const __nv_bfloat16* src = reinterpret_cast(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; From e3620ba990158ac75c5b793d75c89f3dd915beb7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:57:03 +0100 Subject: [PATCH 02/16] cuda prefill_attention: hpattn_p256_kernel for pages of 256 positions, with the page size a constant TF_QWEN_KV_PAGES's default page holds 256 positions. hpattn_p256_kernel is hpattn_kernel with PTOK a compile-time 256: an item's keys sit in one page (static_assert), so it reads the table once and finds its base once, with no runtime divide or modulo. cuda_prompt.zig takes it for bf16 pages of 256; every other page size keeps hpattn_kernel. Addresses only: the same bytes land in the same slots. The kernel contract checks it against contiguous rows twice, with a swapped page table as its control, and registers its usage. --- zig/kernels/cuda/prefill_attention.cu | 32 +++++++++++++++----- zig/kernels/cuda/prefill_attention_rows.cuh | 11 +++++++ zig/src/core/contract/qwen_whole_records.zig | 17 ++++++++--- zig/src/families/qwen3_5/cuda_prompt.zig | 9 ++++-- zig/tests/cuda/contract_qwen_pages.zig | 21 +++++++++++-- 5 files changed, 73 insertions(+), 17 deletions(-) diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 2719004aff..53346ee6bc 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -62,7 +62,7 @@ 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 int CAP = 0, PTOK = 1; const float *ks = nullptr, *vs = nullptr; const int HO = H; @@ -80,7 +80,7 @@ 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 int CAP = 0; const float *ks = nullptr, *vs = nullptr; const int HO = H; @@ -115,7 +115,7 @@ 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 int CAP = 0; const float *ks = nullptr, *vs = nullptr; const int H = 3 * HO, HK = H; @@ -132,7 +132,7 @@ __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 int skip = 0; constexpr const int* PT = nullptr; constexpr long long PS = 0; @@ -149,7 +149,7 @@ __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 int skip = 0; constexpr const int* PT = nullptr; constexpr long long PS = 0; @@ -165,7 +165,7 @@ __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 int skip = 0; const int HO = H; const float *ks = nullptr, *vs = nullptr; @@ -173,13 +173,28 @@ hpattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restri #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 +__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; + constexpr int skip = 0, PTOK = 256; + static_assert(PTOK % HALF == 0, "an item's keys 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 __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 int skip = 0; const int HO = H; auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; @@ -337,6 +352,9 @@ 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::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); diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index 03fcd45c40..ef144d8e26 100644 --- a/zig/kernels/cuda/prefill_attention_rows.cuh +++ b/zig/kernels/cuda/prefill_attention_rows.cuh @@ -73,6 +73,17 @@ if constexpr (F8) { fetch(item); put(item); + } else if constexpr (P256) { + // pages of 256 positions: an item's HALF keys sit in one page, so one table read and one base an item + const __nv_bfloat16* src = reinterpret_cast(item % 4 < 2 ? kc : vc); + const int key0 = (item / 4) * BN + (item % 2) * HALF, k0 = min(key0, limit - 1); + uint4* dst = slots + (item % NS) * HALF * CS; + const __nv_bfloat16* base = src + (int64_t)PT[k0 / PTOK] * PS + ((int64_t)hk * PTOK + k0 % PTOK) * D; + 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 + cp16z(dst + sw<8 * CS>(i / NB, chunk), base + (int64_t)(in ? key - k0 : 0) * D + chunk * 8, in); + } } else { const __nv_bfloat16* src = reinterpret_cast(item % 4 < 2 ? kc : vc); const int key0 = (item / 4) * BN + (item % 2) * HALF; diff --git a/zig/src/core/contract/qwen_whole_records.zig b/zig/src/core/contract/qwen_whole_records.zig index 6de67b542a..945894f766 100644 --- a/zig/src/core/contract/qwen_whole_records.zig +++ b/zig/src/core/contract/qwen_whole_records.zig @@ -47,11 +47,18 @@ pub const records = [_]reg.CudaRecord{ .{ .regs = 255, .static = 0, .bound = 256 }, }, .later_pending = true }, .{ .launch = .{ .name = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 98304, .opt_in = true, .min_sm = 80 }, .usage = .{ - .{ .regs = 253, .static = 0, .bound = 256 }, - .{ .regs = 253, .static = 0, .bound = 256 }, - .{ .regs = 253, .static = 0, .bound = 256 }, - .{ .regs = 231, .static = 0, .bound = 256 }, - .{ .regs = 231, .static = 0, .bound = 256 }, + .{ .regs = 251, .static = 0, .bound = 256 }, + .{ .regs = 251, .static = 0, .bound = 256 }, + .{ .regs = 251, .static = 0, .bound = 256 }, + .{ .regs = 233, .static = 0, .bound = 256 }, + .{ .regs = 233, .static = 0, .bound = 256 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 98304, .opt_in = true, .min_sm = 80 }, .usage = .{ + .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 255, .static = 0, .bound = 256 }, }, .later_pending = true }, .{ .launch = .{ .name = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 98304, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, diff --git a/zig/src/families/qwen3_5/cuda_prompt.zig b/zig/src/families/qwen3_5/cuda_prompt.zig index 36a9bc986b..1554932022 100644 --- a/zig/src/families/qwen3_5/cuda_prompt.zig +++ b/zig/src/families/qwen3_5/cuda_prompt.zig @@ -12,24 +12,26 @@ pub const smem: u32 = 98304; const hpc = 2; const block_rows = 16 * (8 / hpc); -pub const Fns = struct { bf16: cuda.Function, fp8: cuda.Function, bf16_whole: cuda.Function, fp8_whole: cuda.Function, chain: cuda.Function }; +pub const Fns = struct { bf16: cuda.Function, fp8: cuda.Function, bf16_whole: cuda.Function, fp8_whole: cuda.Function, bf16_p256: cuda.Function, chain: cuda.Function }; pub const symbols = struct { pub const bf16 = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi"; pub const fp8 = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi"; pub const bf16_whole = "_ZN20tf_prefill_attention19hpattn_whole_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifi"; pub const fp8_whole = "_ZN20tf_prefill_attention20hpattn8_whole_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifi"; + pub const bf16_p256 = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi"; pub const chain = "_ZN14tf_gdn_prefill12chain_kernelI13__nv_bfloat16S1_Li64ELi32EEEvPKT_S4_PKS1_PKfPKT0_S8_PfPS1_iii"; }; /// gdn_prefill.cu's value rows a block (two threads a row). const chain_rows = 64; pub fn resolve(attention: cuda.Module, gdn: cuda.Module) !Fns { - const f: Fns = .{ .bf16 = try attention.function(symbols.bf16), .fp8 = try attention.function(symbols.fp8), .bf16_whole = try attention.function(symbols.bf16_whole), .fp8_whole = try attention.function(symbols.fp8_whole), .chain = try gdn.function(symbols.chain) }; + const f: Fns = .{ .bf16 = try attention.function(symbols.bf16), .fp8 = try attention.function(symbols.fp8), .bf16_whole = try attention.function(symbols.bf16_whole), .fp8_whole = try attention.function(symbols.fp8_whole), .bf16_p256 = try attention.function(symbols.bf16_p256), .chain = try gdn.function(symbols.chain) }; try f.bf16.allowDynamicShared(smem); try f.fp8.allowDynamicShared(smem); try f.bf16_whole.allowDynamicShared(smem); try f.fp8_whole.allowDynamicShared(smem); + try f.bf16_p256.allowDynamicShared(smem); return f; } @@ -77,7 +79,8 @@ pub fn runFrom(e: *E.Engine, layer: usize, from: usize, end: usize) !void { const paged = sg.pt != 0; var a = arguments(e, sg, layer, rows, from * g.q_heads * head_dim * 2); const cfg: cuda.Config = .{ .grid = .{ .x = @intCast((rows + block_rows - 1) / block_rows), .y = @intCast(g.kv_heads * (group / hpc)) }, .block = .{ .x = 256 }, .shared = smem }; - const f = if (paged) (if (e.kv8) e.prompt_attn.fp8 else e.prompt_attn.bf16) else (if (e.kv8) e.prompt_attn.fp8_whole else e.prompt_attn.bf16_whole); + // pages of 256 positions (the default) take the kernel whose page size is a constant + const f = if (paged) (if (e.kv8) e.prompt_attn.fp8 else if (sg.ptok == 256) e.prompt_attn.bf16_p256 else e.prompt_attn.bf16) else (if (e.kv8) e.prompt_attn.fp8_whole else e.prompt_attn.bf16_whole); try cuda.launch.launch(f, cfg, e.stream, &a); } diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig index a0fb783fd7..a1a03f4e41 100644 --- a/zig/tests/cuda/contract_qwen_pages.zig +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -18,6 +18,7 @@ const prompt_shared = 98304; // cuda_prompt.zig's smem const prompt_symbols = .{ .bf16 = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .fp8 = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", + .bf16_p256 = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", // pages of 256 }; const e4m3 = core.quant.fp8.e4m3; @@ -413,6 +414,10 @@ pub fn run(gpu: check.Gpu) !void { try whole.rows.allowDynamicShared(attn_tiles); } try whole.prompt.allowDynamicShared(prompt_shared); + var p256 = tc; // cuda_prompt.zig's kernel for pages of 256 positions (bf16) + p256.prompt = try prompt.function(prompt_symbols.bf16_p256); + try p256.prompt.allowDynamicShared(prompt_shared); + var red_p256 = format != .bf16; var red_tc = false; var red_tail = false; var red_table = false; @@ -490,6 +495,18 @@ pub fn run(gpu: check.Gpu) !void { red_tc = red_tc or moved; try check.expect(moved, "qwen35 pages {t}: a swapped page table still gave the tensor-core or prompt tiles contiguous bytes (control missed) at lens {any}", .{ format, c.lens }); } + if (format == .bf16 and c.ptok == 256) { + for (0..2) |rep| { + var repeated = std.heap.ArenaAllocator.init(gpu.gpa); + defer repeated.deinit(); + const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, repeated.allocator(), 20 + rep, p256); + try check.sameBytes("qwen35 pages: the 256-position prompt rows, paged against contiguous", pg.prompt, flat_tc.prompt); + } + const swapped = try a.dupe(u32, f.ids[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9, p256); + red_p256 = red_p256 or !std.mem.eql(u8, bad.prompt, flat_tc.prompt); + } const o = try a.dupe(u8, flat.out); std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 var bad_r = flat; @@ -499,8 +516,8 @@ pub fn run(gpu: check.Gpu) !void { try check.expect(caught, "qwen35 pages {t}: a planted attention row passed float64 (control missed)", .{format}); std.debug.print("PASS qwen35 pages {t}: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep; tensor-core and prompt rows), repeats, float64 bound, guards\n", .{ format, c.lens, c.rows, c.ptok }); } - const red = red_table and red_bound and red_tc and red_tail; + const red = red_table and red_bound and red_tc and red_tail and red_p256; try check.expect(red, "qwen35 pages {t}: a control never fired", .{format}); - std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row, uncounted tail keys)\n", .{format}); + std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row, uncounted tail keys, swapped table on 256-position pages)\n", .{format}); } } From eaf6db2aa92d1572107befcfc9b6b0fd5c30d4c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:46:13 +0100 Subject: [PATCH 03/16] qwen35 attention tiles: a decode fetch looks its page up once per page, not once per key Bf16Tiles::fetch and Fp8Tiles::fetch found each key's page with a runtime pos / ptok and pos % ptok. A thread's keys run in order, so it keeps the page it is in (its first position and base) and looks a new one up only when a key leaves it. Addresses only: the same bytes are copied. On a GB10 this brings paged decode on Qwen3.8-27B to the window path's speed (attn_partial was 0.327 against 0.266 ms a call before). --- zig/kernels/cuda/qwen35_attn_tiles.cuh | 34 +++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/qwen35_attn_tiles.cuh b/zig/kernels/cuda/qwen35_attn_tiles.cuh index 2b679c44a5..523f48981d 100644 --- a/zig/kernels/cuda/qwen35_attn_tiles.cuh +++ b/zig/kernels/cuda/qwen35_attn_tiles.cuh @@ -40,13 +40,26 @@ struct Fp8Tiles { template __device__ void fetch(int t0, Row row) { const int tid = threadIdx.x, part = tid & 15; + // a thread's rows mostly share a page: its head's first row there is looked up once, not per load + const unsigned char* near = nullptr; + int lo = 0; #pragma unroll for (int i = 0; i < PER; i++) { 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(kp(r) + part * 16); - vr[i] = *reinterpret_cast(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(k + part * 16); + vr[i] = *reinterpret_cast(v + part * 16); } } const int r = row(t0 + (tid & 31)); @@ -102,12 +115,25 @@ struct Bf16Tiles { template __device__ void fetch(int t0, Row row) { uint4* dst = s.kt + (cur ^ 1) * 2 * TILE; + // a thread's rows mostly share a page: its head's first row there is looked up once, not per 16-byte chunk + const __nv_bfloat16* near = nullptr; + int lo = 0; #pragma unroll 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(); } From a57254f50527c2d09510932b5dbbb1241466cc57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 2 Nov 2026 17:41:44 +0100 Subject: [PATCH 04/16] qwen CUDA: streams on pages by default, the engine's own too TF_QWEN_KV_PAGES unset now pages at 256 positions (a divisor of the 4096-row prompt chunk, as TF_QWEN_KV_SHARE needs) and TF_QWEN_KV_PAGES_ALL unset puts the engine's stream on pages too. 0 keeps a window each (PAGES) or the engine's window (ALL). Beside --slide, whose learner reads whole windows, an unset TF_QWEN_KV_PAGES keeps windows; set explicitly it is still refused. Both are still read through the profile registry, which supplies no value. Its reporting follows the new default: a count entry names what its reader takes when unset (256 for TF_QWEN_KV_PAGES), and TF_QWEN_KV_PAGES_ALL takes the `off` syntax, on unless exactly 0, so the startup line reports 0 and not 256 or 1. --- docs/profiles.md | 13 ++++++---- zig/src/core/profile.zig | 11 ++++++--- zig/src/core/profile_test.zig | 12 ++++++++- zig/src/families/qwen3_5/cuda_native.zig | 10 ++++---- zig/src/families/qwen3_5/cuda_pages.zig | 31 ++++++++++++++++++------ zig/src/server/profile_line.zig | 6 ++--- 6 files changed, 58 insertions(+), 25 deletions(-) diff --git a/docs/profiles.md b/docs/profiles.md index 05ed1b3bf2..2087aad653 100644 --- a/docs/profiles.md +++ b/docs/profiles.md @@ -49,8 +49,9 @@ Initial coverage is: `identical` records the existing switch's token-preserving design and linked evidence; it is not a claim that the switch is faster for every shape, device, or workload. The `*_ALL` checks stay `debug`; Flash Next paging remains -`experimental`. No class turns a switch on. GLM's existing automatic paging default is preserved; other default-on -paging outside this registry is unchanged. The two GLM host-memory switches carry [#1795](https://github.com/ashhart/TensorFold/pull/1795)'s +`experimental`. No class turns a switch on. GLM's existing automatic paging default is preserved; Qwen3.5's family +reader pages at 256 positions with the engine's stream on pages when both switches are unset (windows beside +`--slide`), and `0` keeps windows; other default-on paging outside this registry is unchanged. The two GLM host-memory switches carry [#1795](https://github.com/ashhart/TensorFold/pull/1795)'s identical tokens at TP=2 on two GB10s, and `TF_GLM_PROMPT_TAIL` carries [#2226](https://github.com/ashhart/TensorFold/pull/2226)'s. `TF_GLM_KV_HOST_GIB` stays outside the registry: it sizes those pages rather than choosing a policy, and its reader keeps its own default. @@ -64,13 +65,15 @@ still carries each resolved value, so ranks that disagree still stop before weig The serving process prints its profile and applicable registered overrides, for example: ```text -[tensorfold] profile: stable; registered overrides: TF_QWEN_KV_PAGES=256 (identical, env); TF_QWEN_KV_SHARE=16 (identical, env) +[tensorfold] profile: stable; registered overrides: TF_QWEN_KV_PAGES=128 (identical, env); TF_QWEN_KV_SHARE=16 (identical, env) ``` When no applicable registered setting differs from its default, it prints `registered overrides: none`. This is explicitly the registry's coverage, not an assertion that the whole environment is at defaults. Settings for another -model type/backend and arbitrary environment variables are omitted. Numeric zero for count switches and disabled exact-one -flags do not count as changes. Dependency checks still determine whether a configured switch can be used. +model type/backend and arbitrary environment variables are omitted. A count equal to what the family takes when it is +unset (zero, or 256 for `TF_QWEN_KV_PAGES`) and a disabled exact-one flag do not count as changes; `TF_QWEN_KV_PAGES_ALL` +is on unless it is exactly `0`, so only `0` is reported. Dependency checks still determine whether a configured switch +can be used. For GLM CUDA, `TF_KV_PAGES` uses the registry's `auto` syntax: unset or exactly `auto` leaves the family's automatic choice, while every other explicit value is a policy override, including `off` and `0`. An explicit policy may select the diff --git a/zig/src/core/profile.zig b/zig/src/core/profile.zig index fb71e61df6..becefa7c02 100644 --- a/zig/src/core/profile.zig +++ b/zig/src/core/profile.zig @@ -15,7 +15,7 @@ pub const Switch = enum { TF_GLM_KV_HOST_SHARED, TF_GLM_PROMPT_TAIL, }; -pub const Syntax = enum { count, one, auto }; +pub const Syntax = enum { count, one, auto, off }; pub const Entry = struct { id: Switch, /// When present, these exact values alone differ from the family default; otherwise use syntax. @@ -26,6 +26,8 @@ pub const Entry = struct { effect: []const u8, evidence: []const u8, syntax: Syntax = .count, + /// The count the family's reader takes when the switch is unset (`count` syntax only). + unset: usize = 0, stable: ?[:0]const u8 = null, pub fn name(e: Entry) [:0]const u8 { @@ -45,8 +47,9 @@ pub const Entry = struct { } return switch (e.syntax) { .one => std.mem.eql(u8, raw, "1"), - .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != 0, + .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != e.unset, .auto => !std.mem.eql(u8, raw, "auto"), + .off => std.mem.eql(u8, raw, "0"), }; } }; @@ -55,8 +58,8 @@ const qwen: []const []const u8 = &.{ "qwen3_5", "qwen3_5_moe" }; const flash: []const []const u8 = &.{"qwen4_exp"}; pub const entries = [_]Entry{ - .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .effect = "KV pool page size; memory and workload tradeoffs", .evidence = "#1357" }, - .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .one, .effect = "Put the engine's own stream on pages too", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .unset = 256, .effect = "KV pool page size (unset: 256; 0: a window each); memory and workload tradeoffs", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .off, .effect = "The engine's own stream on pages too (unset); 0 keeps it a window", .evidence = "#1357" }, .{ .id = .TF_QWEN_KV_SHARE, .families = qwen, .class = .identical, .effect = "Keep prompt pages and bounded recurrent snapshots", .evidence = "#1357" }, .{ .id = .TF_KV_SHARE, .families = &.{ "gemma4", "gemma4_unified", "kolibri1", "spark2_5", "glm5_next" }, .class = .identical, .effect = "Keep prompt states over the core KV pool", .evidence = "#1657" }, .{ .id = .TF_FLASHNEXT_KV_PAGES, .families = flash, .class = .experimental, .effect = "KV pool page size; fast-profile qualification pending", .evidence = "#1660" }, diff --git a/zig/src/core/profile_test.zig b/zig/src/core/profile_test.zig index 3b67ff71df..b6818ca81a 100644 --- a/zig/src/core/profile_test.zig +++ b/zig/src/core/profile_test.zig @@ -73,11 +73,21 @@ test "stable profile: reporting follows the original count and exact-one flag se const count = p.entry(.TF_QWEN_KV_SHARE); for ([_][]const u8{ "0", "00", "+0" }) |v| try t.expect(!count.differs(v)); for ([_][]const u8{ "1", "256", "", "-1", "not-a-count" }) |v| try t.expect(count.differs(v)); - const flag = p.entry(.TF_QWEN_KV_PAGES_ALL); + const flag = p.entry(.TF_FLASHNEXT_KV_PAGES_ALL); try t.expect(flag.differs("1")); for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |v| try t.expect(!flag.differs(v)); } +test "stable profile: Qwen's page switches report a change from the family's paged default" { + const pages = p.entry(.TF_QWEN_KV_PAGES); + for ([_][]const u8{ "256", "0256", "+256" }) |v| try t.expect(!pages.differs(v)); + for ([_][]const u8{ "0", "128", "", "not-a-count" }) |v| try t.expect(pages.differs(v)); + const all = p.entry(.TF_QWEN_KV_PAGES_ALL); + try t.expectEqual(p.Syntax.off, all.syntax); + try t.expect(all.differs("0")); + for ([_][]const u8{ "1", "", "00", "false" }) |v| try t.expect(!all.differs(v)); +} + test "stable profile: GLM paging registers automatic policy without supplying or coercing its value" { const e = p.entry(.TF_KV_PAGES); try t.expectEqual(p.Syntax.auto, e.syntax); diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index bff16dde85..06b8e6c81b 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -83,7 +83,7 @@ pub const Loaded = struct { own_prompt_cache: bool = false, /// TF_QWEN_KV_SHARE or TF_QWEN_KV_PAGES_ALL: every stream runs on the pool's pages; the engine's window sits idle. all_on_pages: bool = false, - /// TF_QWEN_KV_PAGES_ALL=1: the engine's own stream takes pages too, a window of them set aside by the plan. + /// TF_QWEN_KV_PAGES_ALL (unset or 1; 0 keeps a window): the engine's own stream takes pages, set aside by the plan. engine_paged: bool = false, /// An integrated GPU: kept states in device memory on the host budget (Target.device), the startup line says so. device_snaps: bool = false, @@ -136,7 +136,7 @@ pub fn ledgerPool(p: *anyopaque) ?*@import("core").kv_pool.Pool { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; - const page_tokens = try pages_mod.paging(); + const page_tokens = try pages_mod.paging(o.slide); const engine_paged = page_tokens > 0 and pages_mod.pagedOnly(); const page_test = try cache_mod.pageTest(); _ = try cache_mod.cacheRows(o.context, page_test); @@ -221,11 +221,11 @@ pub const openDrafter = drafter_open.openDrafter; pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { if (ctx == null) return switch (err) { error.NoQwenVision => "--vision: this checkpoint has no image tower (no vision_config): serve it without --vision", - error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset it", + error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset TF_QWEN_KV_PAGES", error.QwenShareNeedsPages => "TF_QWEN_KV_SHARE keeps prompts as pool pages: set TF_QWEN_KV_PAGES to a divisor of the 4096-row prompt chunk", error.BadQwenKvShare => "TF_QWEN_KV_SHARE takes the most recurrent-state snapshots the prompt cache keeps (0 or unset: off)", error.QwenYarnNoLearner => "--slide: the learner's attention ropes without YaRN; serve without --slide, or without --rope-yarn and a YaRN config", - error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0 or unset: a window each)", + error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0: a window each; unset: 256)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, }; @@ -266,7 +266,7 @@ pub const Lanes = struct { /// TF_QWEN_KV_SHARE's snapshots (0: off) and the prompt cache over the pool's pages `attach` made. share_snaps: usize = 0, prefix: ?*prefix_mod.Prefix = null, - engine_paged: bool = false, // TF_QWEN_KV_PAGES_ALL=1: no window reserved, the engine's own stream a pool lane + engine_paged: bool = false, // TF_QWEN_KV_PAGES_ALL unset or 1: no window reserved, the engine's own stream a pool lane /// The DFlash2 drafter, when one is attached: its rings go into the pool lanes' runs. draft: ?*Draft = null, /// The MTP drafter (qwen3_5_moe), when attached: its ring and state go into the pool lanes' runs. diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 47ca013c28..9a673824a3 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -11,10 +11,16 @@ const Cache = caches.Cache; /// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). pub const window_rows: u32 = engine.max_rows; -/// TF_QWEN_KV_PAGES=n: streams beside the first take n-position pages from one pool (unset or 0: a window each). -pub fn paging() !usize { - const v = @import("core").profile.get(.TF_QWEN_KV_PAGES) orelse return 0; - return parsePaging(v); +/// Positions a page holds with TF_QWEN_KV_PAGES unset: divides the 4096-row prompt chunk, as TF_QWEN_KV_SHARE needs. +pub const default_page_tokens: usize = 256; + +/// TF_QWEN_KV_PAGES=n: streams take n-position pages from one pool (0: a window each; unset: 256, windows by --slide). +pub fn paging(slide: bool) !usize { + return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES), slide); +} + +pub fn pagingOf(v: ?[]const u8, slide: bool) !usize { + return parsePaging(v orelse return if (slide) 0 else default_page_tokens); } pub fn parsePaging(v: []const u8) !usize { @@ -23,10 +29,13 @@ pub fn parsePaging(v: []const u8) !usize { return n; } -/// TF_QWEN_KV_PAGES_ALL=1: every stream on pages, the engine's own too: no window is reserved for it. +/// TF_QWEN_KV_PAGES_ALL: every stream on pages, the engine's own too; 0 keeps a window for the engine's stream. pub fn pagedOnly() bool { - const v = @import("core").profile.get(.TF_QWEN_KV_PAGES_ALL) orelse return false; - return std.mem.eql(u8, v, "1"); + return pagedOnlyOf(@import("core").profile.get(.TF_QWEN_KV_PAGES_ALL)); +} + +pub fn pagedOnlyOf(v: ?[]const u8) bool { + return !std.mem.eql(u8, v orelse return true, "0"); } /// Pages a stream maps at most: a full window and one more forward. @@ -225,3 +234,11 @@ test "KV page dimensions: pool rows fit the signed kernel argument and zero disa try std.testing.expectError(error.BadQwenKvPages, parsePaging("2147483648")); try std.testing.expectError(error.BadQwenKvPages, parsePaging("18446744073709551615")); } + +test "Qwen pages by default: 256 positions and the engine's stream too, windows beside --slide or at 0" { + try std.testing.expectEqual(default_page_tokens, try pagingOf(null, false)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, true)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", false)); + try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", true)); + try std.testing.expect(pagedOnlyOf(null) and pagedOnlyOf("1") and !pagedOnlyOf("0")); +} diff --git a/zig/src/server/profile_line.zig b/zig/src/server/profile_line.zig index a33b2eb8d8..d0ffca2e8b 100644 --- a/zig/src/server/profile_line.zig +++ b/zig/src/server/profile_line.zig @@ -53,13 +53,13 @@ test "profile line: only applicable allowlisted changes are shown, without claim var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); - const env: TestEnv = .{ .qwen = "256", .all = "1", .share = "8" }; + const env: TestEnv = .{ .qwen = "128", .all = "0", .share = "8" }; const line = try format(a, .stable, "qwen3_5", "cuda", env); - try std.testing.expectEqualStrings("profile: stable; registered overrides: TF_QWEN_KV_PAGES=256 (identical, env); TF_QWEN_KV_PAGES_ALL=1 (debug, env); TF_QWEN_KV_SHARE=8 (identical, env)", line); + try std.testing.expectEqualStrings("profile: stable; registered overrides: TF_QWEN_KV_PAGES=128 (identical, env); TF_QWEN_KV_PAGES_ALL=0 (debug, env); TF_QWEN_KV_SHARE=8 (identical, env)", line); for ([_][2][]const u8{ .{ "unknown", "cuda" }, .{ "qwen3_5", "metal" }, .{ "qwen3_5", "hip" } }) |scope| { try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, scope[0], scope[1], env)); } - try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, "qwen3_5", "cuda", TestEnv{ .qwen = "00", .all = "true", .share = "+0" })); + try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, "qwen3_5", "cuda", TestEnv{ .qwen = "256", .all = "1", .share = "+0" })); } test "profile line: GLM reports explicit page policies including off without treating auto as an override" { From e7c001231bf309035a9dada42500f0717b8b4d9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:30:41 +0100 Subject: [PATCH 05/16] qwen35 attention tiles: a decode thread keeps its last page across tiles, so its first row is looked up once a page, not once a tile --- zig/kernels/cuda/qwen35_attn_tiles.cuh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/zig/kernels/cuda/qwen35_attn_tiles.cuh b/zig/kernels/cuda/qwen35_attn_tiles.cuh index 523f48981d..917cc79124 100644 --- a/zig/kernels/cuda/qwen35_attn_tiles.cuh +++ b/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(sg.kv) + (size_t)layer * 2 * kvh_ * sg.cap * (HD + 4); K = slab + (size_t)h * sg.cap * HD; @@ -40,9 +43,6 @@ struct Fp8Tiles { template __device__ void fetch(int t0, Row row) { const int tid = threadIdx.x, part = tid & 15; - // a thread's rows mostly share a page: its head's first row there is looked up once, not per load - const unsigned char* near = nullptr; - int lo = 0; #pragma unroll for (int i = 0; i < PER; i++) { const int r = row(t0 + (tid >> 4) + THREADS / 16 * i); @@ -98,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(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * HD + (size_t)h * sg.cap * HD; V = K + (size_t)kvh * sg.cap * HD; @@ -115,9 +118,6 @@ struct Bf16Tiles { template __device__ void fetch(int t0, Row row) { uint4* dst = s.kt + (cur ^ 1) * 2 * TILE; - // a thread's rows mostly share a page: its head's first row there is looked up once, not per 16-byte chunk - const __nv_bfloat16* near = nullptr; - int lo = 0; #pragma unroll for (int i = 0; i < TILE / THREADS; i++) { const int at = threadIdx.x + THREADS * i, u = at / CH, chunk = at % CH; From 7249217517967d1e76227fb3c076b1824b03445f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:30:41 +0100 Subject: [PATCH 06/16] cuda prefill_attention: hpattn_p256_ids_kernel stages a chunk's page ids in shared memory (up to 768 pages), hpattn_p256_kernel past that --- zig/kernels/cuda/prefill_attention.cu | 37 ++++++++++++- zig/kernels/cuda/prefill_attention_rows.cuh | 16 +++--- zig/src/core/contract/qwen_whole_records.zig | 7 +++ zig/src/families/qwen3_5/cuda_prompt.zig | 55 +++++++++++++++++--- zig/tests/cuda/contract_qwen_pages.zig | 24 ++++++--- 5 files changed, 119 insertions(+), 20 deletions(-) diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 53346ee6bc..081e42b3f3 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -63,6 +63,8 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric constexpr bool causal = true; constexpr int skip = 0; 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; @@ -81,6 +83,8 @@ pattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict_ constexpr long long PS = 0; constexpr int PTOK = 1; 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; @@ -116,6 +120,8 @@ full_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ o constexpr long long PS = 0; constexpr int PTOK = 1; 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; @@ -133,6 +139,8 @@ hpattn_whole_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __ 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 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; @@ -150,6 +158,8 @@ hpattn8_whole_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restr 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 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; @@ -166,6 +176,8 @@ hpattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restri 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 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; @@ -179,9 +191,25 @@ __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; + 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 % HALF == 0, "an item's keys in one page"); + 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 +__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); }; @@ -195,6 +223,8 @@ hpattn8_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restrict__ 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 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); }; @@ -355,6 +385,9 @@ template __global__ void tf_prefill_attention::hpattn_kernel<256, 8, 2, 2>(const 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); diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index ef144d8e26..df8e28861c 100644 --- a/zig/kernels/cuda/prefill_attention_rows.cuh +++ b/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(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; @@ -74,15 +75,14 @@ fetch(item); put(item); } else if constexpr (P256) { - // pages of 256 positions: an item's HALF keys sit in one page, so one table read and one base an item + // 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(item % 4 < 2 ? kc : vc); - const int key0 = (item / 4) * BN + (item % 2) * HALF, k0 = min(key0, limit - 1); + 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)PT[k0 / PTOK] * PS + ((int64_t)hk * PTOK + k0 % PTOK) * D; + 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 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 - cp16z(dst + sw<8 * CS>(i / NB, chunk), base + (int64_t)(in ? key - k0 : 0) * D + chunk * 8, in); + 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(item % 4 < 2 ? kc : vc); @@ -129,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); diff --git a/zig/src/core/contract/qwen_whole_records.zig b/zig/src/core/contract/qwen_whole_records.zig index 945894f766..054140f1dc 100644 --- a/zig/src/core/contract/qwen_whole_records.zig +++ b/zig/src/core/contract/qwen_whole_records.zig @@ -56,9 +56,16 @@ pub const records = [_]reg.CudaRecord{ .{ .launch = .{ .name = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 98304, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 253, .static = 0, .bound = 256 }, + .{ .regs = 253, .static = 0, .bound = 256 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "_ZN20tf_prefill_attention22hpattn_p256_ids_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention22hpattn_p256_ids_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 101376, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, + .{ .regs = 253, .static = 0, .bound = 256 }, + .{ .regs = 253, .static = 0, .bound = 256 }, }, .later_pending = true }, .{ .launch = .{ .name = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", .symbol = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", .image = "prefill_attention", .threads = 256, .dynamic = 98304, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, diff --git a/zig/src/families/qwen3_5/cuda_prompt.zig b/zig/src/families/qwen3_5/cuda_prompt.zig index 1554932022..e39f24cbf8 100644 --- a/zig/src/families/qwen3_5/cuda_prompt.zig +++ b/zig/src/families/qwen3_5/cuda_prompt.zig @@ -8,11 +8,13 @@ const head_dim = geometry.head_dim; /// Two 32-key slots of 256 dims beside eight warps' staged queries (hpattn_kernel<256, 8, 2, 2>). pub const smem: u32 = 98304; +/// hpattn_p256_ids_kernel's page ids in shared memory past smem: 3 KiB, so the block stays within 99 KiB. +pub const page_ids: u32 = 768; /// Query heads a block (the kernels' HPC) and rows a block (16 a warp, four warp rows). const hpc = 2; const block_rows = 16 * (8 / hpc); -pub const Fns = struct { bf16: cuda.Function, fp8: cuda.Function, bf16_whole: cuda.Function, fp8_whole: cuda.Function, bf16_p256: cuda.Function, chain: cuda.Function }; +pub const Fns = struct { bf16: cuda.Function, fp8: cuda.Function, bf16_whole: cuda.Function, fp8_whole: cuda.Function, bf16_p256: cuda.Function, bf16_p256_ids: cuda.Function, chain: cuda.Function }; pub const symbols = struct { pub const bf16 = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi"; @@ -20,18 +22,20 @@ pub const symbols = struct { pub const bf16_whole = "_ZN20tf_prefill_attention19hpattn_whole_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifi"; pub const fp8_whole = "_ZN20tf_prefill_attention20hpattn8_whole_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifi"; pub const bf16_p256 = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi"; + pub const bf16_p256_ids = "_ZN20tf_prefill_attention22hpattn_p256_ids_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi"; pub const chain = "_ZN14tf_gdn_prefill12chain_kernelI13__nv_bfloat16S1_Li64ELi32EEEvPKT_S4_PKS1_PKfPKT0_S8_PfPS1_iii"; }; /// gdn_prefill.cu's value rows a block (two threads a row). const chain_rows = 64; pub fn resolve(attention: cuda.Module, gdn: cuda.Module) !Fns { - const f: Fns = .{ .bf16 = try attention.function(symbols.bf16), .fp8 = try attention.function(symbols.fp8), .bf16_whole = try attention.function(symbols.bf16_whole), .fp8_whole = try attention.function(symbols.fp8_whole), .bf16_p256 = try attention.function(symbols.bf16_p256), .chain = try gdn.function(symbols.chain) }; + const f: Fns = .{ .bf16 = try attention.function(symbols.bf16), .fp8 = try attention.function(symbols.fp8), .bf16_whole = try attention.function(symbols.bf16_whole), .fp8_whole = try attention.function(symbols.fp8_whole), .bf16_p256 = try attention.function(symbols.bf16_p256), .bf16_p256_ids = try attention.function(symbols.bf16_p256_ids), .chain = try gdn.function(symbols.chain) }; try f.bf16.allowDynamicShared(smem); try f.fp8.allowDynamicShared(smem); try f.bf16_whole.allowDynamicShared(smem); try f.fp8_whole.allowDynamicShared(smem); try f.bf16_p256.allowDynamicShared(smem); + try f.bf16_p256_ids.allowDynamicShared(smem + 4 * page_ids); return f; } @@ -59,11 +63,22 @@ fn arguments(e: *const E.Engine, sg: Seg, layer: usize, rows: usize, at: u64) cu if (paged) { a.add(sg.pt); a.add(@as(i64, @intCast(sg.ps))); - a.add(@as(c_int, sg.ptok)); + a.add(@as(c_int, if (p256(e, sg)) @intCast(pageIds(sg, rows)) else sg.ptok)); } return a; } +/// Whether a chunk takes hpattn_p256_kernel: bf16 pages of 256 positions. +fn p256(e: *const E.Engine, sg: Seg) bool { + return sg.pt != 0 and !e.kv8 and sg.ptok == 256; +} + +/// The page ids hpattn_p256_ids_kernel keeps in shared memory: every page the chunk reads; 0 past page_ids. +pub fn pageIds(sg: Seg, rows: usize) u32 { + const pages = (@as(usize, @intCast(sg.len)) + rows + 255) / 256; + return if (pages <= page_ids) @intCast(pages) else 0; +} + /// Attention layer `layer` into e.y: a chunk's queries (e.qn) over cache positions 0 .. len + row (rope wrote its own). pub fn run(e: *E.Engine, layer: usize, rows: usize) !void { return runFrom(e, layer, 0, rows); @@ -78,9 +93,10 @@ pub fn runFrom(e: *E.Engine, layer: usize, from: usize, end: usize) !void { sg.len += @intCast(from); // row `from` is the first query: its keys end at len + from const paged = sg.pt != 0; var a = arguments(e, sg, layer, rows, from * g.q_heads * head_dim * 2); - const cfg: cuda.Config = .{ .grid = .{ .x = @intCast((rows + block_rows - 1) / block_rows), .y = @intCast(g.kv_heads * (group / hpc)) }, .block = .{ .x = 256 }, .shared = smem }; - // pages of 256 positions (the default) take the kernel whose page size is a constant - const f = if (paged) (if (e.kv8) e.prompt_attn.fp8 else if (sg.ptok == 256) e.prompt_attn.bf16_p256 else e.prompt_attn.bf16) else (if (e.kv8) e.prompt_attn.fp8_whole else e.prompt_attn.bf16_whole); + const ids = if (p256(e, sg)) pageIds(sg, rows) else 0; + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast((rows + block_rows - 1) / block_rows), .y = @intCast(g.kv_heads * (group / hpc)) }, .block = .{ .x = 256 }, .shared = smem + 4 * ids }; + // 256-position pages (the default): the kernel with the page size a constant, page ids in shared memory if they fit + const f = if (paged) (if (e.kv8) e.prompt_attn.fp8 else if (sg.ptok == 256) (if (ids > 0) e.prompt_attn.bf16_p256_ids else e.prompt_attn.bf16_p256) else e.prompt_attn.bf16) else (if (e.kv8) e.prompt_attn.fp8_whole else e.prompt_attn.bf16_whole); try cuda.launch.launch(f, cfg, e.stream, &a); } @@ -137,3 +153,30 @@ test "Qwen prompt layout: whole ABI omits page metadata and keeps original layer }; }; } + +test "Qwen prompt layout: 256-position pages pass their page-id count while every id fits shared memory, else 0" { + const t = std.testing; + var e: E.Engine = undefined; + e.g.q_heads = 24; + e.g.kv_heads = 4; + e.qn = 0x10000; + e.y = 0x20000; + var sg = std.mem.zeroes(Seg); + sg.kv = 0x400000; + sg.cap = 1 << 20; + sg.pt = 0xdead0000; + sg.ps = 0x100000; + sg.ptok = 256; + // filled positions (prefix + chunk) against the ids it passes: 768 pages of 256 is the shared table's ceiling + for ([_][3]usize{ .{ 0, 1, 1 }, .{ 255, 1, 1 }, .{ 256, 1, 2 }, .{ 28672, 4096, 128 }, .{ 192512, 4096, 768 }, .{ 192513, 4096, 0 }, .{ 196608, 1, 0 } }) |c| { + sg.len = @intCast(c[0]); + for ([_]bool{ false, true }) |fp8| { + e.kv8 = fp8; + const a = arguments(&e, sg, 1, c[1], 0); + const np: usize = if (fp8) 6 else 4; + // fp8 pages take hpattn8_kernel, which takes the page size + try t.expectEqual(@as(u64, if (fp8) 256 else c[2]), a.integer(np + 9).?); + } + try t.expectEqual(@as(u32, @intCast(c[2])), pageIds(sg, c[1])); + } +} diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig index a1a03f4e41..dda8979610 100644 --- a/zig/tests/cuda/contract_qwen_pages.zig +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -19,6 +19,7 @@ const prompt_symbols = .{ .bf16 = "_ZN20tf_prefill_attention13hpattn_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", .fp8 = "_ZN20tf_prefill_attention14hpattn8_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16PKhS5_PKfS7_PS1_iiiiifiPKixi", .bf16_p256 = "_ZN20tf_prefill_attention18hpattn_p256_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", // pages of 256 + .bf16_p256_ids = "_ZN20tf_prefill_attention22hpattn_p256_ids_kernelILi256ELi8ELi2ELi2EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifiPKixi", // their ids in shared memory }; const e4m3 = core.quant.fp8.e4m3; @@ -197,8 +198,9 @@ pub fn tailFrom(rows: usize) usize { } const Fns = struct { rope: cuda.Function, partial: cuda.Function, merge: cuda.Function, compact: cuda.Function }; -/// The engine's attention kernels: short streams' rows in `partial`, long streams' in `rows`, a chunk's in `prompt`. -const Tc = struct { partial: cuda.Function, rows: cuda.Function, prompt: cuda.Function, shared: u32, whole: bool = false }; +/// Attention kernels: short streams in `partial`, long ones in `rows`, a chunk in `prompt` (`ids`: by page-id count). +const Tc = struct { partial: cuda.Function, rows: cuda.Function, prompt: cuda.Function, shared: u32, whole: bool = false, ids: bool = false }; +const page_ids = 768; // cuda_prompt.zig's ceiling on hpattn_p256_ids_kernel's ids /// rope, attention (`fns.partial`, or `tc`'s kernels), keep on layer 1 of both streams: contiguous, or `tables`' pages. fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]const u32, given: ?[2][]const u32, a: std.mem.Allocator, seed: u64, tc: ?Tc) !Result { @@ -286,13 +288,14 @@ fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]co inline for (.{ len, f.c.rows - from, q_heads, kv_heads, group }) |v| pa.add(@as(c_int, @intCast(v))); pa.add(@as(f32, 1.0 / 16.0)); pa.add(@as(c_int, @intCast(f.cap))); + const ids: u32 = if (t.ids) @intCast((f.c.lens[st] + f.c.rows + 255) / 256) else 0; if (!t.whole) { pa.add(segs[st].pt); pa.add(@as(i64, @intCast(segs[st].ps))); - pa.add(@as(c_int, if (segs[st].pt != 0) @intCast(f.c.ptok) else 1)); + pa.add(@as(c_int, if (t.ids) @intCast(ids) else if (segs[st].pt != 0) @intCast(f.c.ptok) else 1)); } const grid: cuda.launch.Dim3 = .{ .x = @intCast((f.c.rows - from + 63) / 64), .y = kv_heads * (group / 2) }; - try cuda.launch.launch(t.prompt, .{ .grid = grid, .block = .{ .x = 256 }, .shared = prompt_shared }, s, &pa); + try cuda.launch.launch(t.prompt, .{ .grid = grid, .block = .{ .x = 256 }, .shared = prompt_shared + 4 * ids }, s, &pa); }; var r: Result = undefined; r.q = try a.dupe(u8, try rig.bytes(qn, f.q.len * 2)); @@ -414,10 +417,15 @@ pub fn run(gpu: check.Gpu) !void { try whole.rows.allowDynamicShared(attn_tiles); } try whole.prompt.allowDynamicShared(prompt_shared); - var p256 = tc; // cuda_prompt.zig's kernel for pages of 256 positions (bf16) + var p256 = tc; // cuda_prompt.zig's kernels for pages of 256 positions (bf16): page ids from the table p256.prompt = try prompt.function(prompt_symbols.bf16_p256); try p256.prompt.allowDynamicShared(prompt_shared); + var p256_ids = p256; // and every page id in shared memory + p256_ids.prompt = try prompt.function(prompt_symbols.bf16_p256_ids); + p256_ids.ids = true; + try p256_ids.prompt.allowDynamicShared(prompt_shared + 4 * page_ids); var red_p256 = format != .bf16; + var red_ids = format != .bf16; var red_tc = false; var red_tail = false; var red_table = false; @@ -501,11 +509,15 @@ pub fn run(gpu: check.Gpu) !void { defer repeated.deinit(); const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, repeated.allocator(), 20 + rep, p256); try check.sameBytes("qwen35 pages: the 256-position prompt rows, paged against contiguous", pg.prompt, flat_tc.prompt); + const shared = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, repeated.allocator(), 24 + rep, p256_ids); + try check.sameBytes("qwen35 pages: the 256-position prompt rows, page ids in shared memory", shared.prompt, flat_tc.prompt); } const swapped = try a.dupe(u32, f.ids[1]); std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9, p256); red_p256 = red_p256 or !std.mem.eql(u8, bad.prompt, flat_tc.prompt); + const bad_ids = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9, p256_ids); + red_ids = red_ids or !std.mem.eql(u8, bad_ids.prompt, flat_tc.prompt); } const o = try a.dupe(u8, flat.out); std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 @@ -516,7 +528,7 @@ pub fn run(gpu: check.Gpu) !void { try check.expect(caught, "qwen35 pages {t}: a planted attention row passed float64 (control missed)", .{format}); std.debug.print("PASS qwen35 pages {t}: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep; tensor-core and prompt rows), repeats, float64 bound, guards\n", .{ format, c.lens, c.rows, c.ptok }); } - const red = red_table and red_bound and red_tc and red_tail and red_p256; + const red = red_table and red_bound and red_tc and red_tail and red_p256 and red_ids; try check.expect(red, "qwen35 pages {t}: a control never fired", .{format}); std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row, uncounted tail keys, swapped table on 256-position pages)\n", .{format}); } From 91f5e80b2f0b93ec682262978e1b20b3d37c2e2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:30:41 +0100 Subject: [PATCH 07/16] qwen CUDA: a paged prompt-cache save and restore copy one 2D block a page, not one run a page and piece --- zig/build/cuda_family_host.zig | 3 + zig/src/cuda/abi.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_reuse.zig | 36 ++++++ .../qwen3_5/cuda_reuse_pages_test.zig | 120 ++++++++++++++++++ 5 files changed, 161 insertions(+) create mode 100644 zig/src/families/qwen3_5/cuda_reuse_pages_test.zig diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index a0406a1e2e..665e44d78b 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -44,6 +44,9 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B 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); diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index c464348d9d..7c7d98a38e 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -155,6 +155,7 @@ pub const ExecUpdateResultInfo = extern struct { pub const event_record_external: c_uint = 0x1; /// CU_MEMORYTYPE_DEVICE +pub const memory_type_host: c_uint = 1; pub const memory_type_device: c_uint = 2; pub const jit_info_log_buffer: c_int = 3; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index c2b1abe75c..e4092f8eec 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -47,6 +47,7 @@ test { _ = @import("cuda_prefix.zig"); _ = @import("cuda_prefix_fault_test.zig"); _ = @import("cuda_pend_test.zig"); + _ = @import("cuda_reuse_pages_test.zig"); } pub const slide = @import("cuda_slide.zig"); diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index 3dc29d0874..2c30be44d7 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -135,6 +135,34 @@ const Runs = struct { } }; +/// Whether a cache's keys and values go a page a copy: bf16 pages, whose pieces lie back to back a page's rows apart. +fn wholePages(cache: *const Cache, l: Layout) bool { + return cache.paged != null and l.format == .bf16 and l.at > 0; +} + +/// Every bf16 piece's rows from `at`'s pages, one 2D copy a page (`tokens` rows apart there, `at` apart in `copy`). +fn pageCopies(e: *Engine, cache: *const Cache, l: Layout, copy: *const Copy, off: usize, way: enum { save, restore }) !void { + const pg = cache.paged.?; + const row = head_dim * 2; + const page_bytes = (Layout{ .g = l.g, .format = l.format, .cap = pg.tokens, .at = pg.tokens }).kvBytes(); + const pieces = l.count() - 1; + var p: usize = 0; + while (p * pg.tokens < l.at) : (p += 1) { + const rows = @min(pg.tokens, l.at - p * pg.tokens); + const dev = pg.base + pg.ids[p] * page_bytes; + const o = off + p * pg.tokens * row; + // the kept side: device memory on an integrated GPU (copy.dev), else the pinned host copy + const kept: cuda.abi.DevicePtr = if (copy.dev) |b| b.ptr + o else 0; + const host: ?*u8 = if (copy.dev == null) ©.host.bytes.ptr[o] else null; + const kind: c_int = @intCast(if (copy.dev == null) cuda.abi.memory_type_host else cuda.abi.memory_type_device); + const m: cuda.abi.Memcpy2D = switch (way) { + .save => .{ .src = dev, .src_pitch = pg.tokens * row, .dst_type = kind, .dst_host = host, .dst = kept, .dst_pitch = l.at * row, .width = rows * row, .height = pieces }, + .restore => .{ .src_type = kind, .src_host = host, .src = kept, .src_pitch = l.at * row, .dst = dev, .dst_pitch = pg.tokens * row, .width = rows * row, .height = pieces }, + }; + try e.d.check(e.d.api.cuMemcpy2DAsync_v2(&m, e.stream.handle), "cuMemcpy2DAsync"); + } +} + /// Whether a request's images match a kept state's: the same identities at the same rows, for every image before `at`. pub fn sameImages(kept: []const Seen, request: []const Seen, at: usize) bool { var n: usize = 0; @@ -198,6 +226,10 @@ pub fn saveCache(e: *Engine, cache: *Cache, at: usize, seen: []const Seen, devic var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); + if (i == 1 and wholePages(cache, l)) { + try pageCopies(e, cache, l, copy, off, .save); + break; + } var it = Runs.of(cache, l, i, p); while (it.next()) |r| { const s, const o = .{ e.stream.handle, off + r.at }; @@ -248,6 +280,10 @@ pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); + if (i == 1 and wholePages(cache, l)) { + try pageCopies(e, cache, l, copy, off, .restore); + break; + } var it = Runs.of(cache, l, i, p); while (it.next()) |r| { const s, const o = .{ e.stream.handle, off + r.at }; diff --git a/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig b/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig new file mode 100644 index 0000000000..6bc3a14607 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig @@ -0,0 +1,120 @@ +//! A paged bf16 cache's prompt-state save and restore with a mock driver: the kept side is the copy's own memory. +const std = @import("std"); +const cuda = @import("cuda"); +const reuse = @import("cuda_reuse.zig"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); +const t = std.testing; + +const Mock = struct { + const dev_at: u64 = 0x4000000; + var dev_len: usize = 0; + var host: []u8 = &.{}; + var pinned: usize = 0; + var flat_host: usize = 0; + var flat_dev: usize = 0; + var bad: usize = 0; + var pages: usize = 0; + + fn alloc(p: *cuda.abi.DevicePtr, len: usize) callconv(.c) cuda.abi.Result { + dev_len = len; + p.* = dev_at; + return cuda.abi.success; + } + fn free(_: cuda.abi.DevicePtr) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn hostAlloc(p: *?*anyopaque, len: usize, _: c_uint) callconv(.c) cuda.abi.Result { + pinned += 1; + host = std.heap.page_allocator.alloc(u8, len) catch return cuda.abi.error_not_ready; + p.* = host.ptr; + return cuda.abi.success; + } + fn freeHost(_: ?*anyopaque) callconv(.c) cuda.abi.Result { + std.heap.page_allocator.free(host); + host = &.{}; + return cuda.abi.success; + } + fn dtod(_: cuda.abi.DevicePtr, _: cuda.abi.DevicePtr, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + flat_dev += 1; + return cuda.abi.success; + } + fn htod(_: cuda.abi.DevicePtr, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + flat_host += 1; + return cuda.abi.success; + } + fn dtoh(_: ?*anyopaque, _: cuda.abi.DevicePtr, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + flat_host += 1; + return cuda.abi.success; + } + /// A page copy's kept side must lie in the copy's own buffer, as its memory type says. + fn copy2d(m: *const cuda.abi.Memcpy2D, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + pages += 1; + const save = m.src_type == cuda.abi.memory_type_device and m.src_host == null and m.src < dev_at; + const kind = if (save) m.dst_type else m.src_type; + const hp: usize = if (save) (if (m.dst_host) |h| @intFromPtr(h) else 0) else (if (m.src_host) |h| @intFromPtr(h) else 0); + const at: usize = if (kind == cuda.abi.memory_type_device) @intCast((if (save) m.dst else m.src) -% dev_at) else hp -% @intFromPtr(host.ptr); + const room = if (kind == cuda.abi.memory_type_device) dev_len else host.len; + if (at + (m.height - 1) * (if (save) m.dst_pitch else m.src_pitch) + m.width > room) bad += 1; + return cuda.abi.success; + } + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } +}; + +test "Qwen paged reuse: a bf16 save and restore copy a page each into the copy's own memory: device memory too, pinning nothing" { + var api: cuda.abi.Api = undefined; + api.cuMemAlloc_v2 = Mock.alloc; + api.cuMemFree_v2 = Mock.free; + api.cuMemHostAlloc = Mock.hostAlloc; + api.cuMemFreeHost = Mock.freeHost; + api.cuMemcpyDtoDAsync_v2 = Mock.dtod; + api.cuMemcpyHtoDAsync_v2 = Mock.htod; + api.cuMemcpyDtoHAsync_v2 = Mock.dtoh; + api.cuMemcpy2DAsync_v2 = Mock.copy2d; + api.cuStreamSynchronize = Mock.sync; + const d: cuda.Driver = .{ .lib = undefined, .api = api }; + var e: engine.Engine = undefined; + e.kit = null; + e.gpa = t.allocator; + e.d = &d; + e.g = undefined; + e.g.hidden = 5120; + e.g.k_heads = 16; + e.g.v_heads = 48; + e.g.q_heads = 32; + e.g.kv_heads = 4; + e.g.layers = 8; + @memset(&e.g.linear, false); + for (0..8) |i| e.g.linear[i] = i % 4 != 3; + e.stream = .{ .d = &d, .handle = null }; + if (comptime @hasField(engine.Engine, "kv8")) e.kv8 = false; + var ids = [_]u32{ 3, 0, 2, 1 }; + var cache: engine.Cache = undefined; + cache.cap = 64; + cache.len = 40; // two full pages of 16 and one of 8 + cache.last = null; + cache.pend = 0; + cache.delta = 3; + cache.paged = .{ .base = 0x100000, .table = .{ .d = &d, .ptr = 0x90000, .len = 0 }, .ids = &ids, .tokens = 16, .own_ids = false }; + cache.kv = .{ .d = &d, .ptr = 0x10000, .len = 0 }; + cache.state = .{ .d = &d, .ptr = 0x20000, .len = 0 }; + for ([_]bool{ true, false }) |device| { + Mock.pinned = 0; + Mock.flat_host = 0; + Mock.flat_dev = 0; + Mock.bad = 0; + Mock.pages = 0; + const copy = try reuse.saveCache(&e, &cache, 40, &.{}, device); + try t.expectEqual(@as(usize, 3), Mock.pages); + try reuse.restoreCache(&e, &cache, copy, &.{}); + try t.expectEqual(@as(usize, 6), Mock.pages); + try t.expectEqual(@as(usize, 0), Mock.bad); + // the state piece goes flat: on the device for a device copy, else through the pinned host copy + try t.expectEqual(@as(usize, if (device) 2 else 0), Mock.flat_dev); + try t.expectEqual(@as(usize, if (device) 0 else 2), Mock.flat_host); + try t.expectEqual(@as(usize, @intFromBool(!device)), Mock.pinned); + reuse.dropCopy(&e, copy); + } +} From 3337929f13c071f295266b01941376e5ec56f548 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 4 Nov 2026 08:50:30 +0100 Subject: [PATCH 08/16] qwen CUDA: a MoE checkpoint stays on windows by default, pages opt-in; a dense one pages at 256 positions TF_QWEN_KV_PAGES unset now pages a dense Qwen3.5 (the 27B) and keeps a MoE (the 35B-A3B) on windows: the 35B's warm turn measured slower on pages, the 27B's no different. The reader takes the MoE bit from the loaded geometry, so a malformed value is still refused before the load. The profile registry reports a count against what this family takes unset (unset_in), and TF_QWEN_KV_PAGES_ALL takes 1 or 0, refusing anything else by name. --- docs/profiles.md | 4 +-- zig/src/core/profile.zig | 17 +++++++---- zig/src/core/profile_test.zig | 34 ++++++++++----------- zig/src/families/qwen3_5/cuda_native.zig | 17 ++++++----- zig/src/families/qwen3_5/cuda_pages.zig | 38 ++++++++++++++---------- zig/src/server/profile_line.zig | 2 +- 6 files changed, 64 insertions(+), 48 deletions(-) diff --git a/docs/profiles.md b/docs/profiles.md index 2087aad653..e7f9ccb190 100644 --- a/docs/profiles.md +++ b/docs/profiles.md @@ -71,8 +71,8 @@ The serving process prints its profile and applicable registered overrides, for When no applicable registered setting differs from its default, it prints `registered overrides: none`. This is explicitly the registry's coverage, not an assertion that the whole environment is at defaults. Settings for another model type/backend and arbitrary environment variables are omitted. A count equal to what the family takes when it is -unset (zero, or 256 for `TF_QWEN_KV_PAGES`) and a disabled exact-one flag do not count as changes; `TF_QWEN_KV_PAGES_ALL` -is on unless it is exactly `0`, so only `0` is reported. Dependency checks still determine whether a configured switch +unset (zero, or 256 for `TF_QWEN_KV_PAGES` on a dense Qwen; zero on a MoE) and a disabled exact-one flag do not count as +changes; `TF_QWEN_KV_PAGES_ALL` takes `1` (as unset) or `0`, so only `0` is reported. Dependency checks still determine whether a configured switch can be used. For GLM CUDA, `TF_KV_PAGES` uses the registry's `auto` syntax: unset or exactly `auto` leaves the family's automatic choice, diff --git a/zig/src/core/profile.zig b/zig/src/core/profile.zig index becefa7c02..5cd6a54f1e 100644 --- a/zig/src/core/profile.zig +++ b/zig/src/core/profile.zig @@ -26,8 +26,9 @@ pub const Entry = struct { effect: []const u8, evidence: []const u8, syntax: Syntax = .count, - /// The count the family's reader takes when the switch is unset (`count` syntax only). + /// The count the reader takes when the switch is unset in `unset_in`'s families (`count` syntax; 0 in the others). unset: usize = 0, + unset_in: []const []const u8 = &.{}, stable: ?[:0]const u8 = null, pub fn name(e: Entry) [:0]const u8 { @@ -40,26 +41,32 @@ pub const Entry = struct { return false; } - pub fn differs(e: Entry, raw: []const u8) bool { + pub fn differs(e: Entry, family: []const u8, raw: []const u8) bool { if (e.report_values) |values| { for (values) |v| if (std.mem.eql(u8, raw, v)) return true; return false; } return switch (e.syntax) { .one => std.mem.eql(u8, raw, "1"), - .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != e.unset, + .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != e.unsetIn(family), .auto => !std.mem.eql(u8, raw, "auto"), .off => std.mem.eql(u8, raw, "0"), }; } + + /// The count `family`'s reader takes with the switch unset. + pub fn unsetIn(e: Entry, family: []const u8) usize { + for (e.unset_in) |f| if (std.mem.eql(u8, f, family)) return e.unset; + return 0; + } }; const qwen: []const []const u8 = &.{ "qwen3_5", "qwen3_5_moe" }; const flash: []const []const u8 = &.{"qwen4_exp"}; pub const entries = [_]Entry{ - .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .unset = 256, .effect = "KV pool page size (unset: 256; 0: a window each); memory and workload tradeoffs", .evidence = "#1357" }, - .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .off, .effect = "The engine's own stream on pages too (unset); 0 keeps it a window", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .unset = 256, .unset_in = &.{"qwen3_5"}, .effect = "KV pool page size (unset: 256 dense, windows MoE; 0: a window each); memory and workload tradeoffs", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .off, .effect = "With pages, the engine's own stream on them too (unset or 1); 0 keeps it a window", .evidence = "#1357" }, .{ .id = .TF_QWEN_KV_SHARE, .families = qwen, .class = .identical, .effect = "Keep prompt pages and bounded recurrent snapshots", .evidence = "#1357" }, .{ .id = .TF_KV_SHARE, .families = &.{ "gemma4", "gemma4_unified", "kolibri1", "spark2_5", "glm5_next" }, .class = .identical, .effect = "Keep prompt states over the core KV pool", .evidence = "#1657" }, .{ .id = .TF_FLASHNEXT_KV_PAGES, .families = flash, .class = .experimental, .effect = "KV pool page size; fast-profile qualification pending", .evidence = "#1660" }, diff --git a/zig/src/core/profile_test.zig b/zig/src/core/profile_test.zig index b6818ca81a..c4cd5065b9 100644 --- a/zig/src/core/profile_test.zig +++ b/zig/src/core/profile_test.zig @@ -39,18 +39,18 @@ test "stable profile: shared paging variable has disjoint GLM and Nemotron repor if (e.applies("glm5_next", "cuda")) { glm += 1; try t.expectEqual(p.Syntax.auto, e.syntax); - try t.expect(e.differs("on")); - try t.expect(e.differs("off")); - try t.expect(!e.differs("auto")); + try t.expect(e.differs("glm5_next", "on")); + try t.expect(e.differs("glm5_next", "off")); + try t.expect(!e.differs("glm5_next", "auto")); } if (e.applies("nemotron_h", "cuda")) { nemotron += 1; try t.expectEqual(@as(usize, 2), e.report_values.?.len); try t.expectEqualStrings("0", e.report_values.?[0]); try t.expectEqualStrings("all", e.report_values.?[1]); - for ([_][]const u8{ "0", "all" }) |raw| try t.expect(e.differs(raw)); + for ([_][]const u8{ "0", "all" }) |raw| try t.expect(e.differs("nemotron_h", raw)); for ([_][]const u8{ "", "on", "off", "1", "auto", "00", "+0", "ALL", "false", " all ", "0\n" }) |raw| - try t.expect(!e.differs(raw)); + try t.expect(!e.differs("nemotron_h", raw)); } if (e.applies("nemotron_h_omni", "cuda")) omni += 1; try t.expectEqual(p.Class.identical, e.class); @@ -71,21 +71,21 @@ test "stable profile: explicit off, empty, malformed and ordinary values are ret test "stable profile: reporting follows the original count and exact-one flag semantics" { const count = p.entry(.TF_QWEN_KV_SHARE); - for ([_][]const u8{ "0", "00", "+0" }) |v| try t.expect(!count.differs(v)); - for ([_][]const u8{ "1", "256", "", "-1", "not-a-count" }) |v| try t.expect(count.differs(v)); + for ([_][]const u8{ "0", "00", "+0" }) |v| try t.expect(!count.differs("qwen3_5", v)); + for ([_][]const u8{ "1", "256", "", "-1", "not-a-count" }) |v| try t.expect(count.differs("qwen3_5", v)); const flag = p.entry(.TF_FLASHNEXT_KV_PAGES_ALL); - try t.expect(flag.differs("1")); - for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |v| try t.expect(!flag.differs(v)); + try t.expect(flag.differs("qwen3_5", "1")); + for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |v| try t.expect(!flag.differs("qwen3_5", v)); } test "stable profile: Qwen's page switches report a change from the family's paged default" { const pages = p.entry(.TF_QWEN_KV_PAGES); - for ([_][]const u8{ "256", "0256", "+256" }) |v| try t.expect(!pages.differs(v)); - for ([_][]const u8{ "0", "128", "", "not-a-count" }) |v| try t.expect(pages.differs(v)); + for ([_][]const u8{ "256", "0256", "+256" }) |v| try t.expect(!pages.differs("qwen3_5", v)); + for ([_][]const u8{ "0", "128", "", "not-a-count" }) |v| try t.expect(pages.differs("qwen3_5", v)); const all = p.entry(.TF_QWEN_KV_PAGES_ALL); try t.expectEqual(p.Syntax.off, all.syntax); - try t.expect(all.differs("0")); - for ([_][]const u8{ "1", "", "00", "false" }) |v| try t.expect(!all.differs(v)); + try t.expect(all.differs("qwen3_5", "0")); + for ([_][]const u8{ "1", "", "00", "false" }) |v| try t.expect(!all.differs("qwen3_5", v)); } test "stable profile: GLM paging registers automatic policy without supplying or coercing its value" { @@ -95,9 +95,9 @@ test "stable profile: GLM paging registers automatic policy without supplying or try t.expect(e.applies("glm5_next", "cuda")); try t.expect(!e.applies("qwen3_5", "cuda")); try t.expect(!e.applies("glm5_next", "metal")); - try t.expect(!e.differs("auto")); + try t.expect(!e.differs("qwen3_5", "auto")); for ([_][]const u8{ "off", "0", "on", "1", "all", "", "AUTO", "false", " 0 " }) |raw| - try t.expect(e.differs(raw)); + try t.expect(e.differs("qwen3_5", raw)); } test "stable profile: the GLM host KV tier registers as an exact-one switch a family, off unless set" { @@ -110,8 +110,8 @@ test "stable profile: the GLM host KV tier registers as an exact-one switch a fa try t.expect(!e.applies("glm5_next", "mlx")); try t.expectEqualStrings("#1795", e.evidence); try t.expect(p.resolve(.stable, id, null) == null); - try t.expect(e.differs("1")); - for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |raw| try t.expect(!e.differs(raw)); + try t.expect(e.differs("glm5_next", "1")); + for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |raw| try t.expect(!e.differs("glm5_next", raw)); } } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 06b8e6c81b..c2c3a3086b 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -83,7 +83,7 @@ pub const Loaded = struct { own_prompt_cache: bool = false, /// TF_QWEN_KV_SHARE or TF_QWEN_KV_PAGES_ALL: every stream runs on the pool's pages; the engine's window sits idle. all_on_pages: bool = false, - /// TF_QWEN_KV_PAGES_ALL (unset or 1; 0 keeps a window): the engine's own stream takes pages, set aside by the plan. + /// TF_QWEN_KV_PAGES_ALL (unset or 1; 0 keeps a window): with pages, the engine's own stream takes them too. engine_paged: bool = false, /// An integrated GPU: kept states in device memory on the host budget (Target.device), the startup line says so. device_snaps: bool = false, @@ -136,16 +136,18 @@ pub fn ledgerPool(p: *anyopaque) ?*@import("core").kv_pool.Pool { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; - const page_tokens = try pages_mod.paging(o.slide); - const engine_paged = page_tokens > 0 and pages_mod.pagedOnly(); + const pages_all = try pages_mod.pagedOnly(); + _ = try pages_mod.paging(o.slide, false); // a malformed TF_QWEN_KV_PAGES is refused before the load const page_test = try cache_mod.pageTest(); _ = try cache_mod.cacheRows(o.context, page_test); - if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; const share_snaps = try prefix_mod.sharing(); - // a kept state sits on a cold pass's chunk cut, so later chunks keep a cold pass's bits: pages divide the chunk - if (share_snaps > 0 and (page_tokens == 0 or E.max_rows % page_tokens != 0)) return error.QwenShareNeedsPages; const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8, fitOf(o)); errdefer e.deinit(); + const page_tokens = try pages_mod.paging(o.slide, e.g.moe()); // a MoE pages only when asked + const engine_paged = page_tokens > 0 and pages_all; + if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; + // a kept state sits on a cold pass's chunk cut, so later chunks keep a cold pass's bits: pages divide the chunk + if (share_snaps > 0 and (page_tokens == 0 or E.max_rows % page_tokens != 0)) return error.QwenShareNeedsPages; if (o.rope_yarn) |f| { // --rope-yarn over the downloaded config: its table, and the window it allows var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -225,7 +227,8 @@ pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { error.QwenShareNeedsPages => "TF_QWEN_KV_SHARE keeps prompts as pool pages: set TF_QWEN_KV_PAGES to a divisor of the 4096-row prompt chunk", error.BadQwenKvShare => "TF_QWEN_KV_SHARE takes the most recurrent-state snapshots the prompt cache keeps (0 or unset: off)", error.QwenYarnNoLearner => "--slide: the learner's attention ropes without YaRN; serve without --slide, or without --rope-yarn and a YaRN config", - error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0: a window each; unset: 256)", + error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0: a window each; unset: 256 for a dense model, windows for a MoE)", + error.BadQwenKvPagesAll => "TF_QWEN_KV_PAGES_ALL takes 1 (the engine's own stream on pages too, as unset) or 0 (a window for it)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, }; diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 9a673824a3..4da5363b2d 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -11,16 +11,16 @@ const Cache = caches.Cache; /// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). pub const window_rows: u32 = engine.max_rows; -/// Positions a page holds with TF_QWEN_KV_PAGES unset: divides the 4096-row prompt chunk, as TF_QWEN_KV_SHARE needs. +/// A dense model's page positions with TF_QWEN_KV_PAGES unset: divides the 4096-row chunk, as TF_QWEN_KV_SHARE needs. pub const default_page_tokens: usize = 256; -/// TF_QWEN_KV_PAGES=n: streams take n-position pages from one pool (0: a window each; unset: 256, windows by --slide). -pub fn paging(slide: bool) !usize { - return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES), slide); +/// TF_QWEN_KV_PAGES=n: n-position pages from one pool (0: a window each; unset: 256 dense, windows for MoE or --slide). +pub fn paging(slide: bool, moe: bool) !usize { + return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES), slide, moe); } -pub fn pagingOf(v: ?[]const u8, slide: bool) !usize { - return parsePaging(v orelse return if (slide) 0 else default_page_tokens); +pub fn pagingOf(v: ?[]const u8, slide: bool, moe: bool) !usize { + return parsePaging(v orelse return if (slide or moe) 0 else default_page_tokens); } pub fn parsePaging(v: []const u8) !usize { @@ -29,13 +29,16 @@ pub fn parsePaging(v: []const u8) !usize { return n; } -/// TF_QWEN_KV_PAGES_ALL: every stream on pages, the engine's own too; 0 keeps a window for the engine's stream. -pub fn pagedOnly() bool { +/// TF_QWEN_KV_PAGES_ALL: every stream on pages, the engine's own too (unset or 1); 0 keeps the engine's a window. +pub fn pagedOnly() !bool { return pagedOnlyOf(@import("core").profile.get(.TF_QWEN_KV_PAGES_ALL)); } -pub fn pagedOnlyOf(v: ?[]const u8) bool { - return !std.mem.eql(u8, v orelse return true, "0"); +pub fn pagedOnlyOf(v: ?[]const u8) !bool { + const s = v orelse return true; + if (std.mem.eql(u8, s, "1")) return true; + if (std.mem.eql(u8, s, "0")) return false; + return error.BadQwenKvPagesAll; } /// Pages a stream maps at most: a full window and one more forward. @@ -235,10 +238,13 @@ test "KV page dimensions: pool rows fit the signed kernel argument and zero disa try std.testing.expectError(error.BadQwenKvPages, parsePaging("18446744073709551615")); } -test "Qwen pages by default: 256 positions and the engine's stream too, windows beside --slide or at 0" { - try std.testing.expectEqual(default_page_tokens, try pagingOf(null, false)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, true)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", false)); - try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", true)); - try std.testing.expect(pagedOnlyOf(null) and pagedOnlyOf("1") and !pagedOnlyOf("0")); +test "Qwen pages by default: a dense model at 256 positions, the engine's stream too; windows for a MoE, --slide or 0" { + try std.testing.expectEqual(default_page_tokens, try pagingOf(null, false, false)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, true, false)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, false, true)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", false, false)); + try std.testing.expectEqual(@as(usize, 256), try pagingOf("256", false, true)); + try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", true, false)); + try std.testing.expect(try pagedOnlyOf(null) and try pagedOnlyOf("1") and !try pagedOnlyOf("0")); + for ([_][]const u8{ "2", "", "true", "01" }) |v| try std.testing.expectError(error.BadQwenKvPagesAll, pagedOnlyOf(v)); } diff --git a/zig/src/server/profile_line.zig b/zig/src/server/profile_line.zig index d0ffca2e8b..b181564361 100644 --- a/zig/src/server/profile_line.zig +++ b/zig/src/server/profile_line.zig @@ -20,7 +20,7 @@ pub fn format(a: std.mem.Allocator, profile: p.Profile, family: []const u8, back for (p.entries) |e| { if (!e.applies(family, backend)) continue; const raw = environment.get(e.name()) orelse continue; - if (!e.differs(raw)) continue; + if (!e.differs(family, raw)) continue; try out.writer.print("{s} {s}=", .{ if (found) ";" else "", e.name() }); try shown(&out.writer, raw); try out.writer.print(" ({s}, env)", .{@tagName(e.class)}); From c86df9ddcad628a42da7e8abbb6e897ab43cae24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 4 Nov 2026 08:50:30 +0100 Subject: [PATCH 09/16] changelog: Qwen pages by default on the dense models, --parallel a cap (#2011) --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4dcf39c832..6777b9b5b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- 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. + - 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. Telemetry only (#2293). From 35c028daf4eca546f747817ac44eeeeec5e5380d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:32:50 +0100 Subject: [PATCH 10/16] qwen CUDA: pages opt-in on a dense model too, TF_QWEN_KV_PAGES_ALL unset off again, both as on main TF_QWEN_KV_PAGES unset keeps a window each on the 27B and the 35B-A3B alike: the 27B measured about 0.3 % slower on 8k and 32k prompts and the warm turn on pages. The reader still takes the MoE bit from the loaded geometry, so a per-model default is one constant (unset_pages) once #2043 closes the gap. With pages unset there is no pool, so --parallel reserves a window a stream as on main; set, it is a cap and a request can wait for pages. TF_QWEN_KV_PAGES_ALL is on only at exactly 1 again, so setting only TF_QWEN_KV_PAGES gives what it gave on main. The profile registry, its startup line and docs/profiles.md are main's again. --- docs/profiles.md | 13 +++---- zig/src/core/profile.zig | 20 +++-------- zig/src/core/profile_test.zig | 38 ++++++++------------- zig/src/families/qwen3_5/cuda_native.zig | 17 ++++++---- zig/src/families/qwen3_5/cuda_pages.zig | 43 +++++++++++++----------- zig/src/server/profile_line.zig | 8 ++--- 6 files changed, 61 insertions(+), 78 deletions(-) diff --git a/docs/profiles.md b/docs/profiles.md index e7f9ccb190..05ed1b3bf2 100644 --- a/docs/profiles.md +++ b/docs/profiles.md @@ -49,9 +49,8 @@ Initial coverage is: `identical` records the existing switch's token-preserving design and linked evidence; it is not a claim that the switch is faster for every shape, device, or workload. The `*_ALL` checks stay `debug`; Flash Next paging remains -`experimental`. No class turns a switch on. GLM's existing automatic paging default is preserved; Qwen3.5's family -reader pages at 256 positions with the engine's stream on pages when both switches are unset (windows beside -`--slide`), and `0` keeps windows; other default-on paging outside this registry is unchanged. The two GLM host-memory switches carry [#1795](https://github.com/ashhart/TensorFold/pull/1795)'s +`experimental`. No class turns a switch on. GLM's existing automatic paging default is preserved; other default-on +paging outside this registry is unchanged. The two GLM host-memory switches carry [#1795](https://github.com/ashhart/TensorFold/pull/1795)'s identical tokens at TP=2 on two GB10s, and `TF_GLM_PROMPT_TAIL` carries [#2226](https://github.com/ashhart/TensorFold/pull/2226)'s. `TF_GLM_KV_HOST_GIB` stays outside the registry: it sizes those pages rather than choosing a policy, and its reader keeps its own default. @@ -65,15 +64,13 @@ still carries each resolved value, so ranks that disagree still stop before weig The serving process prints its profile and applicable registered overrides, for example: ```text -[tensorfold] profile: stable; registered overrides: TF_QWEN_KV_PAGES=128 (identical, env); TF_QWEN_KV_SHARE=16 (identical, env) +[tensorfold] profile: stable; registered overrides: TF_QWEN_KV_PAGES=256 (identical, env); TF_QWEN_KV_SHARE=16 (identical, env) ``` When no applicable registered setting differs from its default, it prints `registered overrides: none`. This is explicitly the registry's coverage, not an assertion that the whole environment is at defaults. Settings for another -model type/backend and arbitrary environment variables are omitted. A count equal to what the family takes when it is -unset (zero, or 256 for `TF_QWEN_KV_PAGES` on a dense Qwen; zero on a MoE) and a disabled exact-one flag do not count as -changes; `TF_QWEN_KV_PAGES_ALL` takes `1` (as unset) or `0`, so only `0` is reported. Dependency checks still determine whether a configured switch -can be used. +model type/backend and arbitrary environment variables are omitted. Numeric zero for count switches and disabled exact-one +flags do not count as changes. Dependency checks still determine whether a configured switch can be used. For GLM CUDA, `TF_KV_PAGES` uses the registry's `auto` syntax: unset or exactly `auto` leaves the family's automatic choice, while every other explicit value is a policy override, including `off` and `0`. An explicit policy may select the diff --git a/zig/src/core/profile.zig b/zig/src/core/profile.zig index 5cd6a54f1e..fb71e61df6 100644 --- a/zig/src/core/profile.zig +++ b/zig/src/core/profile.zig @@ -15,7 +15,7 @@ pub const Switch = enum { TF_GLM_KV_HOST_SHARED, TF_GLM_PROMPT_TAIL, }; -pub const Syntax = enum { count, one, auto, off }; +pub const Syntax = enum { count, one, auto }; pub const Entry = struct { id: Switch, /// When present, these exact values alone differ from the family default; otherwise use syntax. @@ -26,9 +26,6 @@ pub const Entry = struct { effect: []const u8, evidence: []const u8, syntax: Syntax = .count, - /// The count the reader takes when the switch is unset in `unset_in`'s families (`count` syntax; 0 in the others). - unset: usize = 0, - unset_in: []const []const u8 = &.{}, stable: ?[:0]const u8 = null, pub fn name(e: Entry) [:0]const u8 { @@ -41,32 +38,25 @@ pub const Entry = struct { return false; } - pub fn differs(e: Entry, family: []const u8, raw: []const u8) bool { + pub fn differs(e: Entry, raw: []const u8) bool { if (e.report_values) |values| { for (values) |v| if (std.mem.eql(u8, raw, v)) return true; return false; } return switch (e.syntax) { .one => std.mem.eql(u8, raw, "1"), - .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != e.unsetIn(family), + .count => (std.fmt.parseInt(usize, raw, 10) catch return true) != 0, .auto => !std.mem.eql(u8, raw, "auto"), - .off => std.mem.eql(u8, raw, "0"), }; } - - /// The count `family`'s reader takes with the switch unset. - pub fn unsetIn(e: Entry, family: []const u8) usize { - for (e.unset_in) |f| if (std.mem.eql(u8, f, family)) return e.unset; - return 0; - } }; const qwen: []const []const u8 = &.{ "qwen3_5", "qwen3_5_moe" }; const flash: []const []const u8 = &.{"qwen4_exp"}; pub const entries = [_]Entry{ - .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .unset = 256, .unset_in = &.{"qwen3_5"}, .effect = "KV pool page size (unset: 256 dense, windows MoE; 0: a window each); memory and workload tradeoffs", .evidence = "#1357" }, - .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .off, .effect = "With pages, the engine's own stream on them too (unset or 1); 0 keeps it a window", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES, .families = qwen, .class = .identical, .effect = "KV pool page size; memory and workload tradeoffs", .evidence = "#1357" }, + .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = qwen, .class = .debug, .syntax = .one, .effect = "Put the engine's own stream on pages too", .evidence = "#1357" }, .{ .id = .TF_QWEN_KV_SHARE, .families = qwen, .class = .identical, .effect = "Keep prompt pages and bounded recurrent snapshots", .evidence = "#1357" }, .{ .id = .TF_KV_SHARE, .families = &.{ "gemma4", "gemma4_unified", "kolibri1", "spark2_5", "glm5_next" }, .class = .identical, .effect = "Keep prompt states over the core KV pool", .evidence = "#1657" }, .{ .id = .TF_FLASHNEXT_KV_PAGES, .families = flash, .class = .experimental, .effect = "KV pool page size; fast-profile qualification pending", .evidence = "#1660" }, diff --git a/zig/src/core/profile_test.zig b/zig/src/core/profile_test.zig index c4cd5065b9..3b67ff71df 100644 --- a/zig/src/core/profile_test.zig +++ b/zig/src/core/profile_test.zig @@ -39,18 +39,18 @@ test "stable profile: shared paging variable has disjoint GLM and Nemotron repor if (e.applies("glm5_next", "cuda")) { glm += 1; try t.expectEqual(p.Syntax.auto, e.syntax); - try t.expect(e.differs("glm5_next", "on")); - try t.expect(e.differs("glm5_next", "off")); - try t.expect(!e.differs("glm5_next", "auto")); + try t.expect(e.differs("on")); + try t.expect(e.differs("off")); + try t.expect(!e.differs("auto")); } if (e.applies("nemotron_h", "cuda")) { nemotron += 1; try t.expectEqual(@as(usize, 2), e.report_values.?.len); try t.expectEqualStrings("0", e.report_values.?[0]); try t.expectEqualStrings("all", e.report_values.?[1]); - for ([_][]const u8{ "0", "all" }) |raw| try t.expect(e.differs("nemotron_h", raw)); + for ([_][]const u8{ "0", "all" }) |raw| try t.expect(e.differs(raw)); for ([_][]const u8{ "", "on", "off", "1", "auto", "00", "+0", "ALL", "false", " all ", "0\n" }) |raw| - try t.expect(!e.differs("nemotron_h", raw)); + try t.expect(!e.differs(raw)); } if (e.applies("nemotron_h_omni", "cuda")) omni += 1; try t.expectEqual(p.Class.identical, e.class); @@ -71,21 +71,11 @@ test "stable profile: explicit off, empty, malformed and ordinary values are ret test "stable profile: reporting follows the original count and exact-one flag semantics" { const count = p.entry(.TF_QWEN_KV_SHARE); - for ([_][]const u8{ "0", "00", "+0" }) |v| try t.expect(!count.differs("qwen3_5", v)); - for ([_][]const u8{ "1", "256", "", "-1", "not-a-count" }) |v| try t.expect(count.differs("qwen3_5", v)); - const flag = p.entry(.TF_FLASHNEXT_KV_PAGES_ALL); - try t.expect(flag.differs("qwen3_5", "1")); - for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |v| try t.expect(!flag.differs("qwen3_5", v)); -} - -test "stable profile: Qwen's page switches report a change from the family's paged default" { - const pages = p.entry(.TF_QWEN_KV_PAGES); - for ([_][]const u8{ "256", "0256", "+256" }) |v| try t.expect(!pages.differs("qwen3_5", v)); - for ([_][]const u8{ "0", "128", "", "not-a-count" }) |v| try t.expect(pages.differs("qwen3_5", v)); - const all = p.entry(.TF_QWEN_KV_PAGES_ALL); - try t.expectEqual(p.Syntax.off, all.syntax); - try t.expect(all.differs("qwen3_5", "0")); - for ([_][]const u8{ "1", "", "00", "false" }) |v| try t.expect(!all.differs("qwen3_5", v)); + for ([_][]const u8{ "0", "00", "+0" }) |v| try t.expect(!count.differs(v)); + for ([_][]const u8{ "1", "256", "", "-1", "not-a-count" }) |v| try t.expect(count.differs(v)); + const flag = p.entry(.TF_QWEN_KV_PAGES_ALL); + try t.expect(flag.differs("1")); + for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |v| try t.expect(!flag.differs(v)); } test "stable profile: GLM paging registers automatic policy without supplying or coercing its value" { @@ -95,9 +85,9 @@ test "stable profile: GLM paging registers automatic policy without supplying or try t.expect(e.applies("glm5_next", "cuda")); try t.expect(!e.applies("qwen3_5", "cuda")); try t.expect(!e.applies("glm5_next", "metal")); - try t.expect(!e.differs("qwen3_5", "auto")); + try t.expect(!e.differs("auto")); for ([_][]const u8{ "off", "0", "on", "1", "all", "", "AUTO", "false", " 0 " }) |raw| - try t.expect(e.differs("qwen3_5", raw)); + try t.expect(e.differs(raw)); } test "stable profile: the GLM host KV tier registers as an exact-one switch a family, off unless set" { @@ -110,8 +100,8 @@ test "stable profile: the GLM host KV tier registers as an exact-one switch a fa try t.expect(!e.applies("glm5_next", "mlx")); try t.expectEqualStrings("#1795", e.evidence); try t.expect(p.resolve(.stable, id, null) == null); - try t.expect(e.differs("glm5_next", "1")); - for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |raw| try t.expect(!e.differs("glm5_next", raw)); + try t.expect(e.differs("1")); + for ([_][]const u8{ "0", "", "01", "true", " 1 " }) |raw| try t.expect(!e.differs(raw)); } } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index c2c3a3086b..c5f02751fa 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -83,7 +83,7 @@ pub const Loaded = struct { own_prompt_cache: bool = false, /// TF_QWEN_KV_SHARE or TF_QWEN_KV_PAGES_ALL: every stream runs on the pool's pages; the engine's window sits idle. all_on_pages: bool = false, - /// TF_QWEN_KV_PAGES_ALL (unset or 1; 0 keeps a window): with pages, the engine's own stream takes them too. + /// TF_QWEN_KV_PAGES_ALL=1: the engine's own stream takes pages too, a window of them set aside by the plan. engine_paged: bool = false, /// An integrated GPU: kept states in device memory on the host budget (Target.device), the startup line says so. device_snaps: bool = false, @@ -136,14 +136,14 @@ pub fn ledgerPool(p: *anyopaque) ?*@import("core").kv_pool.Pool { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; - const pages_all = try pages_mod.pagedOnly(); + const pages_all = pages_mod.pagedOnly(); _ = try pages_mod.paging(o.slide, false); // a malformed TF_QWEN_KV_PAGES is refused before the load const page_test = try cache_mod.pageTest(); _ = try cache_mod.cacheRows(o.context, page_test); const share_snaps = try prefix_mod.sharing(); const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8, fitOf(o)); errdefer e.deinit(); - const page_tokens = try pages_mod.paging(o.slide, e.g.moe()); // a MoE pages only when asked + const page_tokens = try pages_mod.paging(o.slide, e.g.moe()); // pages only when asked, dense or MoE const engine_paged = page_tokens > 0 and pages_all; if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; // a kept state sits on a cold pass's chunk cut, so later chunks keep a cold pass's bits: pages divide the chunk @@ -223,12 +223,11 @@ pub const openDrafter = drafter_open.openDrafter; pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { if (ctx == null) return switch (err) { error.NoQwenVision => "--vision: this checkpoint has no image tower (no vision_config): serve it without --vision", - error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset TF_QWEN_KV_PAGES", + error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset it", error.QwenShareNeedsPages => "TF_QWEN_KV_SHARE keeps prompts as pool pages: set TF_QWEN_KV_PAGES to a divisor of the 4096-row prompt chunk", error.BadQwenKvShare => "TF_QWEN_KV_SHARE takes the most recurrent-state snapshots the prompt cache keeps (0 or unset: off)", error.QwenYarnNoLearner => "--slide: the learner's attention ropes without YaRN; serve without --slide, or without --rope-yarn and a YaRN config", - error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0: a window each; unset: 256 for a dense model, windows for a MoE)", - error.BadQwenKvPagesAll => "TF_QWEN_KV_PAGES_ALL takes 1 (the engine's own stream on pages too, as unset) or 0 (a window for it)", + error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0 or unset: a window each)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, }; @@ -269,7 +268,7 @@ pub const Lanes = struct { /// TF_QWEN_KV_SHARE's snapshots (0: off) and the prompt cache over the pool's pages `attach` made. share_snaps: usize = 0, prefix: ?*prefix_mod.Prefix = null, - engine_paged: bool = false, // TF_QWEN_KV_PAGES_ALL unset or 1: no window reserved, the engine's own stream a pool lane + engine_paged: bool = false, // TF_QWEN_KV_PAGES_ALL=1: no window reserved, the engine's own stream a pool lane /// The DFlash2 drafter, when one is attached: its rings go into the pool lanes' runs. draft: ?*Draft = null, /// The MTP drafter (qwen3_5_moe), when attached: its ring and state go into the pool lanes' runs. @@ -539,6 +538,10 @@ test "Qwen share faults: --parallel N admits N streams at once with every stream try @import("cuda_admit_test.zig").admits(Lanes); } +test "Qwen pages default: --parallel a cap with TF_QWEN_KV_PAGES set, a request waits for pages; unset, a window each" { + try @import("cuda_admit_test.zig").caps(Lanes); +} + test "Qwen stepped prefill: cancel between host iterations retains only the last completed cut" { try @import("cuda_prefill_test.zig").cancelled(Lanes, Lanes.backend); } diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 4da5363b2d..41df033e45 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -11,16 +11,16 @@ const Cache = caches.Cache; /// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). pub const window_rows: u32 = engine.max_rows; -/// A dense model's page positions with TF_QWEN_KV_PAGES unset: divides the 4096-row chunk, as TF_QWEN_KV_SHARE needs. -pub const default_page_tokens: usize = 256; +/// TF_QWEN_KV_PAGES unset: windows for a dense model and a MoE alike, as on main, until #2043 closes the page cost. +pub const unset_pages: struct { dense: usize, moe: usize } = .{ .dense = 0, .moe = 0 }; -/// TF_QWEN_KV_PAGES=n: n-position pages from one pool (0: a window each; unset: 256 dense, windows for MoE or --slide). +/// TF_QWEN_KV_PAGES=n: n-position pages from one pool, --parallel a cap (0 or unset: a window each, as on main). pub fn paging(slide: bool, moe: bool) !usize { return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES), slide, moe); } pub fn pagingOf(v: ?[]const u8, slide: bool, moe: bool) !usize { - return parsePaging(v orelse return if (slide or moe) 0 else default_page_tokens); + return parsePaging(v orelse return if (slide) 0 else if (moe) unset_pages.moe else unset_pages.dense); } pub fn parsePaging(v: []const u8) !usize { @@ -29,16 +29,13 @@ pub fn parsePaging(v: []const u8) !usize { return n; } -/// TF_QWEN_KV_PAGES_ALL: every stream on pages, the engine's own too (unset or 1); 0 keeps the engine's a window. -pub fn pagedOnly() !bool { +/// TF_QWEN_KV_PAGES_ALL=1: every stream on pages, the engine's own too: no window is reserved for it (unset: off). +pub fn pagedOnly() bool { return pagedOnlyOf(@import("core").profile.get(.TF_QWEN_KV_PAGES_ALL)); } -pub fn pagedOnlyOf(v: ?[]const u8) !bool { - const s = v orelse return true; - if (std.mem.eql(u8, s, "1")) return true; - if (std.mem.eql(u8, s, "0")) return false; - return error.BadQwenKvPagesAll; +pub fn pagedOnlyOf(v: ?[]const u8) bool { + return std.mem.eql(u8, v orelse return false, "1"); } /// Pages a stream maps at most: a full window and one more forward. @@ -238,13 +235,19 @@ test "KV page dimensions: pool rows fit the signed kernel argument and zero disa try std.testing.expectError(error.BadQwenKvPages, parsePaging("18446744073709551615")); } -test "Qwen pages by default: a dense model at 256 positions, the engine's stream too; windows for a MoE, --slide or 0" { - try std.testing.expectEqual(default_page_tokens, try pagingOf(null, false, false)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, true, false)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, false, true)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", false, false)); - try std.testing.expectEqual(@as(usize, 256), try pagingOf("256", false, true)); - try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", true, false)); - try std.testing.expect(try pagedOnlyOf(null) and try pagedOnlyOf("1") and !try pagedOnlyOf("0")); - for ([_][]const u8{ "2", "", "true", "01" }) |v| try std.testing.expectError(error.BadQwenKvPagesAll, pagedOnlyOf(v)); +test "Qwen pages default: unset or 0 keeps windows on a dense model and a MoE, a set size pages either" { + for ([_]bool{ false, true }) |moe| for ([_]bool{ false, true }) |slide| { + try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, slide, moe)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", slide, moe)); + try std.testing.expectEqual(@as(usize, 256), try pagingOf("256", slide, moe)); + try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", slide, moe)); + try std.testing.expectEqual(@as(usize, 1), try pagingOf("1", slide, moe)); + try std.testing.expectError(error.BadQwenKvPages, pagingOf("", slide, moe)); + }; +} + +test "Qwen pages default: TF_QWEN_KV_PAGES_ALL unset keeps the engine's stream a window, as on main; 1 pages it" { + try std.testing.expect(!pagedOnlyOf(null)); + try std.testing.expect(pagedOnlyOf("1")); + for ([_][]const u8{ "0", "", "true", "01", " 1 " }) |v| try std.testing.expect(!pagedOnlyOf(v)); } diff --git a/zig/src/server/profile_line.zig b/zig/src/server/profile_line.zig index b181564361..a33b2eb8d8 100644 --- a/zig/src/server/profile_line.zig +++ b/zig/src/server/profile_line.zig @@ -20,7 +20,7 @@ pub fn format(a: std.mem.Allocator, profile: p.Profile, family: []const u8, back for (p.entries) |e| { if (!e.applies(family, backend)) continue; const raw = environment.get(e.name()) orelse continue; - if (!e.differs(family, raw)) continue; + if (!e.differs(raw)) continue; try out.writer.print("{s} {s}=", .{ if (found) ";" else "", e.name() }); try shown(&out.writer, raw); try out.writer.print(" ({s}, env)", .{@tagName(e.class)}); @@ -53,13 +53,13 @@ test "profile line: only applicable allowlisted changes are shown, without claim var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); - const env: TestEnv = .{ .qwen = "128", .all = "0", .share = "8" }; + const env: TestEnv = .{ .qwen = "256", .all = "1", .share = "8" }; const line = try format(a, .stable, "qwen3_5", "cuda", env); - try std.testing.expectEqualStrings("profile: stable; registered overrides: TF_QWEN_KV_PAGES=128 (identical, env); TF_QWEN_KV_PAGES_ALL=0 (debug, env); TF_QWEN_KV_SHARE=8 (identical, env)", line); + try std.testing.expectEqualStrings("profile: stable; registered overrides: TF_QWEN_KV_PAGES=256 (identical, env); TF_QWEN_KV_PAGES_ALL=1 (debug, env); TF_QWEN_KV_SHARE=8 (identical, env)", line); for ([_][2][]const u8{ .{ "unknown", "cuda" }, .{ "qwen3_5", "metal" }, .{ "qwen3_5", "hip" } }) |scope| { try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, scope[0], scope[1], env)); } - try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, "qwen3_5", "cuda", TestEnv{ .qwen = "256", .all = "1", .share = "+0" })); + try std.testing.expectEqualStrings("profile: stable; registered overrides: none", try format(a, .stable, "qwen3_5", "cuda", TestEnv{ .qwen = "00", .all = "true", .share = "+0" })); } test "profile line: GLM reports explicit page policies including off without treating auto as an override" { From 54d186941461b72003d48cf45051c6beeada0403 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:32:50 +0100 Subject: [PATCH 11/16] qwen CUDA host tests: pages default (unset, 0, 1, a size; dense and MoE), _ALL unset off, --parallel a cap only with pages test-qwen35-pages-default-host runs them, part of test-cuda-host. A planted wrong admit count fails the cap test. --- zig/build/cuda_family_host.zig | 5 +++ zig/src/families/qwen3_5/cuda_admit_test.zig | 39 ++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index 665e44d78b..52506315cb 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -41,6 +41,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); diff --git a/zig/src/families/qwen3_5/cuda_admit_test.zig b/zig/src/families/qwen3_5/cuda_admit_test.zig index 7d22ac08c7..f549637366 100644 --- a/zig/src/families/qwen3_5/cuda_admit_test.zig +++ b/zig/src/families/qwen3_5/cuda_admit_test.zig @@ -56,3 +56,42 @@ pub fn admits(comptime Lanes: type) !void { try t.expectEqual(@as(usize, parallel), try admitted(Lanes, parallel, 0)); // the checks } } + +/// --parallel `parallel` with pages for `paged` streams: those admit, the rest wait for pages (--parallel a cap). +fn capped(comptime Lanes: type, parallel: u32, paged: u32) !usize { + var d = D.driver(); + var e = engine(&d); + var l: Lanes = .{ + .gpa = t.allocator, + .io = undefined, + .e = &e, + .capacity = capacity, + .page_tokens = page, + .engine_paged = true, + .target = undefined, + }; + const lay = pages.layout(&e, capacity, page, 0, 0); + try native.attach(&l, parallel - 1, paged * (lay.pages(capacity + pages.window_rows) + lay.runPages())); + var open: std.ArrayList(*engine_mod.Cache) = .empty; + defer { + for (open.items) |x| _ = l.pages.?.close(x); + open.deinit(t.allocator); + l.pages.?.deinit(); + } + while (open.items.len <= parallel) { + if (!native.fits(&l, capacity - pages.window_rows, 16)) break; + try open.append(t.allocator, try l.take(capacity - pages.window_rows, 16)); + } + return open.items.len; +} + +/// TF_QWEN_KV_PAGES set: --parallel caps the streams and a request waits for pages. Unset: no pool, nothing waits. +pub fn caps(comptime Lanes: type) !void { + try t.expectEqual(@as(usize, 2), try capped(Lanes, 4, 2)); + try t.expectEqual(@as(usize, 4), try capped(Lanes, 4, 4)); + var d = D.driver(); + var e = engine(&d); + var l: Lanes = .{ .gpa = t.allocator, .io = undefined, .e = &e, .capacity = capacity, .target = undefined }; + try t.expect(native.pagedLayout(&l, 0) == null); // no pool: the plan reserves a window for each of --parallel's + try t.expect(native.fits(&l, capacity, 16)); +} From 16f3241cbde260c62c58b176eff909f66b88cbde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:32:50 +0100 Subject: [PATCH 12/16] changelog: Qwen pages opt-in, --parallel a cap with TF_QWEN_KV_PAGES set (#2011) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6777b9b5b9..eba49f3d08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- 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. - 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. Telemetry only (#2293). From f149c84baa7d8ae21ebbf1e4c9d5fcab1d702144 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 6 Nov 2026 09:24:47 +0100 Subject: [PATCH 13/16] qwen CUDA: the two pages refusals stop before the weights load again, unset_pages a plain 0 With no per-model default left, TF_QWEN_KV_PAGES no longer depends on the model or --slide, so paging() is main's again and open() reads it, and refuses QwenPagesNoLearner and QwenShareNeedsPages, before Engine.init. --- zig/src/families/qwen3_5/cuda.zig | 4 +++ zig/src/families/qwen3_5/cuda_admit_test.zig | 17 ++++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 12 ++++----- zig/src/families/qwen3_5/cuda_pages.zig | 28 +++++++++----------- zig/src/families/qwen3_5/cuda_prompt.zig | 2 +- 5 files changed, 40 insertions(+), 23 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index e4092f8eec..60e6527cab 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -73,3 +73,7 @@ test { _ = train; _ = @import("cuda_read_cost_test.zig"); } + +test "Qwen pages default: the learner's and the prefix share's refusals stop before the weights load" { + try @import("cuda_admit_test.zig").refusals(); +} diff --git a/zig/src/families/qwen3_5/cuda_admit_test.zig b/zig/src/families/qwen3_5/cuda_admit_test.zig index f549637366..8fff7322d7 100644 --- a/zig/src/families/qwen3_5/cuda_admit_test.zig +++ b/zig/src/families/qwen3_5/cuda_admit_test.zig @@ -95,3 +95,20 @@ pub fn caps(comptime Lanes: type) !void { try t.expect(native.pagedLayout(&l, 0) == null); // no pool: the plan reserves a window for each of --parallel's try t.expect(native.fits(&l, capacity, 16)); } + +/// The refusal `open` gives with these switches: no model at the path, so reaching the load names another error. +fn refused(slide: bool, kv_pages: ?[:0]const u8, share: ?[:0]const u8) !void { + defer _ = env.unsetenv("TF_QWEN_KV_PAGES"); + defer _ = env.unsetenv("TF_QWEN_KV_SHARE"); + if (kv_pages) |v| _ = env.setenv("TF_QWEN_KV_PAGES", v, 1); + if (share) |v| _ = env.setenv("TF_QWEN_KV_SHARE", v, 1); + const o: native.Options = .{ .context = capacity, .drafts = false, .slide = slide }; + _ = try native.open(t.allocator, t.io, undefined, "/nonexistent", null, o); +} + +/// A set switch's refusal needs no model: it stops before the weights load, as the other switch checks do. +pub fn refusals() !void { + try t.expectError(error.QwenPagesNoLearner, refused(true, "256", null)); + try t.expectError(error.QwenShareNeedsPages, refused(false, null, "16")); + try t.expectError(error.QwenShareNeedsPages, refused(false, "100", "16")); +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index c5f02751fa..99e57488f7 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -136,18 +136,16 @@ pub fn ledgerPool(p: *anyopaque) ?*@import("core").kv_pool.Pool { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; - const pages_all = pages_mod.pagedOnly(); - _ = try pages_mod.paging(o.slide, false); // a malformed TF_QWEN_KV_PAGES is refused before the load + const page_tokens = try pages_mod.paging(); + const engine_paged = page_tokens > 0 and pages_mod.pagedOnly(); const page_test = try cache_mod.pageTest(); _ = try cache_mod.cacheRows(o.context, page_test); - const share_snaps = try prefix_mod.sharing(); - const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8, fitOf(o)); - errdefer e.deinit(); - const page_tokens = try pages_mod.paging(o.slide, e.g.moe()); // pages only when asked, dense or MoE - const engine_paged = page_tokens > 0 and pages_all; if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; + const share_snaps = try prefix_mod.sharing(); // a kept state sits on a cold pass's chunk cut, so later chunks keep a cold pass's bits: pages divide the chunk if (share_snaps > 0 and (page_tokens == 0 or E.max_rows % page_tokens != 0)) return error.QwenShareNeedsPages; + const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8, fitOf(o)); + errdefer e.deinit(); if (o.rope_yarn) |f| { // --rope-yarn over the downloaded config: its table, and the window it allows var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 41df033e45..9a433a0ecc 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -11,16 +11,16 @@ const Cache = caches.Cache; /// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). pub const window_rows: u32 = engine.max_rows; -/// TF_QWEN_KV_PAGES unset: windows for a dense model and a MoE alike, as on main, until #2043 closes the page cost. -pub const unset_pages: struct { dense: usize, moe: usize } = .{ .dense = 0, .moe = 0 }; +/// TF_QWEN_KV_PAGES unset: a window each, dense or MoE, as on main, until #2043 closes the page cost. +pub const unset_pages: usize = 0; /// TF_QWEN_KV_PAGES=n: n-position pages from one pool, --parallel a cap (0 or unset: a window each, as on main). -pub fn paging(slide: bool, moe: bool) !usize { - return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES), slide, moe); +pub fn paging() !usize { + return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES)); } -pub fn pagingOf(v: ?[]const u8, slide: bool, moe: bool) !usize { - return parsePaging(v orelse return if (slide) 0 else if (moe) unset_pages.moe else unset_pages.dense); +pub fn pagingOf(v: ?[]const u8) !usize { + return parsePaging(v orelse return unset_pages); } pub fn parsePaging(v: []const u8) !usize { @@ -235,15 +235,13 @@ test "KV page dimensions: pool rows fit the signed kernel argument and zero disa try std.testing.expectError(error.BadQwenKvPages, parsePaging("18446744073709551615")); } -test "Qwen pages default: unset or 0 keeps windows on a dense model and a MoE, a set size pages either" { - for ([_]bool{ false, true }) |moe| for ([_]bool{ false, true }) |slide| { - try std.testing.expectEqual(@as(usize, 0), try pagingOf(null, slide, moe)); - try std.testing.expectEqual(@as(usize, 0), try pagingOf("0", slide, moe)); - try std.testing.expectEqual(@as(usize, 256), try pagingOf("256", slide, moe)); - try std.testing.expectEqual(@as(usize, 128), try pagingOf("128", slide, moe)); - try std.testing.expectEqual(@as(usize, 1), try pagingOf("1", slide, moe)); - try std.testing.expectError(error.BadQwenKvPages, pagingOf("", slide, moe)); - }; +test "Qwen pages default: unset or 0 keeps windows, a set size pages" { + try std.testing.expectEqual(@as(usize, 0), try pagingOf(null)); + try std.testing.expectEqual(@as(usize, 0), try pagingOf("0")); + try std.testing.expectEqual(@as(usize, 256), try pagingOf("256")); + try std.testing.expectEqual(@as(usize, 128), try pagingOf("128")); + try std.testing.expectEqual(@as(usize, 1), try pagingOf("1")); + try std.testing.expectError(error.BadQwenKvPages, pagingOf("")); } test "Qwen pages default: TF_QWEN_KV_PAGES_ALL unset keeps the engine's stream a window, as on main; 1 pages it" { diff --git a/zig/src/families/qwen3_5/cuda_prompt.zig b/zig/src/families/qwen3_5/cuda_prompt.zig index e39f24cbf8..d6e4fdd0bb 100644 --- a/zig/src/families/qwen3_5/cuda_prompt.zig +++ b/zig/src/families/qwen3_5/cuda_prompt.zig @@ -95,7 +95,7 @@ pub fn runFrom(e: *E.Engine, layer: usize, from: usize, end: usize) !void { var a = arguments(e, sg, layer, rows, from * g.q_heads * head_dim * 2); const ids = if (p256(e, sg)) pageIds(sg, rows) else 0; const cfg: cuda.Config = .{ .grid = .{ .x = @intCast((rows + block_rows - 1) / block_rows), .y = @intCast(g.kv_heads * (group / hpc)) }, .block = .{ .x = 256 }, .shared = smem + 4 * ids }; - // 256-position pages (the default): the kernel with the page size a constant, page ids in shared memory if they fit + // 256-position pages: the kernel with the page size a constant, page ids in shared memory if they fit const f = if (paged) (if (e.kv8) e.prompt_attn.fp8 else if (sg.ptok == 256) (if (ids > 0) e.prompt_attn.bf16_p256_ids else e.prompt_attn.bf16_p256) else e.prompt_attn.bf16) else (if (e.kv8) e.prompt_attn.fp8_whole else e.prompt_attn.bf16_whole); try cuda.launch.launch(f, cfg, e.stream, &a); } From 448b6d05d9a7d5aada2a5dc7d7bb5a643606e52c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 6 Nov 2026 09:24:47 +0100 Subject: [PATCH 14/16] qwen pages contract: the PASS line names the ids kernel's control --- zig/tests/cuda/contract_qwen_pages.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig index dda8979610..12e26137a4 100644 --- a/zig/tests/cuda/contract_qwen_pages.zig +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -530,6 +530,6 @@ pub fn run(gpu: check.Gpu) !void { } const red = red_table and red_bound and red_tc and red_tail and red_p256 and red_ids; try check.expect(red, "qwen35 pages {t}: a control never fired", .{format}); - std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row, uncounted tail keys, swapped table on 256-position pages)\n", .{format}); + std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row, uncounted tail keys, swapped table on 256-position pages, swapped table with page ids in shared memory)\n", .{format}); } } From 215ae3b01b8994539d43624494f50b9c46642085 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 6 Nov 2026 09:24:47 +0100 Subject: [PATCH 15/16] changelog: Qwen pages cites (#2011, part of #1943) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eba49f3d08..34d245d68f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +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. +- 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. Telemetry only (#2293). From 19caa06c942717e2d3a5bfd73b1b867118446b4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:26:27 +0100 Subject: [PATCH 16/16] qwen contract records: tf_q35_attn_rows8_mt2 at 243 registers on sm_121, measured with the fp8 tiles' one page lookup --- zig/src/core/contract/qwen_whole_records.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/core/contract/qwen_whole_records.zig b/zig/src/core/contract/qwen_whole_records.zig index 054140f1dc..48559b3936 100644 --- a/zig/src/core/contract/qwen_whole_records.zig +++ b/zig/src/core/contract/qwen_whole_records.zig @@ -2,7 +2,7 @@ const std = @import("std"); const reg = @import("registry.zig"); pub const records = [_]reg.CudaRecord{ - .{ .launch = .{ .name = "tf_q35_attn_rows8_mt2", .symbol = "tf_q35_attn_rows8_mt2", .image = "qwen35_kv8", .threads = 128, .min_sm = 120 }, .usage = .{ .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 242, .static = 36992, .bound = 128 }, .{ .regs = 242, .static = 36992, .bound = 128 } }, .later_pending = true }, + .{ .launch = .{ .name = "tf_q35_attn_rows8_mt2", .symbol = "tf_q35_attn_rows8_mt2", .image = "qwen35_kv8", .threads = 128, .min_sm = 120 }, .usage = .{ .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 242, .static = 36992, .bound = 128 }, .{ .regs = 243, .static = 36992, .bound = 128 } }, .later_pending = true }, .{ .launch = .{ .name = "tf_q35_attn_rows8_mt2_whole", .symbol = "tf_q35_attn_rows8_mt2_whole", .image = "qwen35_kv8", .threads = 128, .min_sm = 120 }, .usage = .{ .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 0, .static = 0 }, .{ .regs = 239, .static = 36992, .bound = 128 }, .{ .regs = 239, .static = 36992, .bound = 128 } }, .later_pending = true }, .{ .launch = .{ .name = "tf_q35_attn_partial_whole", .symbol = "tf_q35_attn_partial_whole", .image = "qwen35", .threads = 128, .dynamic = 65536, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 167, .static = 11648, .bound = 128 },