diff --git a/CHANGELOG.md b/CHANGELOG.md index b4fe108d1f..c30f33e153 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -113,6 +113,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - GLM-5.3-Flash on CUDA: an EXL3 checkpoint whose `quantization_config` sits inside `text_config` loads as EXL3. Admission already read it there, but the config parser read only the top level, so such a pack was admitted and then asked for NVFP4 dense weights at load (#2327). - Docs: `docs/hermes-agent.md` shows how to use TensorFold from Hermes Agent through the community [hermes-tensorfold](https://github.com/atakann/hermes-tensorfold) plugin (a model provider and a dashboard tab), checked on a GB10 with Qwen3.8-Flash-Next (#2330, part of #2109). Thanks to @atakann. +- Qwen3.8-27B and the other dense Qwen3.5 models on CUDA: streams run on KV pool pages of 256 positions by default, the engine's own stream too, so `--parallel` is now a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. `TF_QWEN_KV_PAGES=0` keeps a window each, as before. Qwen3.6-35B-A3B and the other MoE checkpoints stay on windows by default, with pages opt-in (`TF_QWEN_KV_PAGES=256`). Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011). Thanks to @jschmied. +- Qwen3.5 and Qwen3.6 on CUDA: KV pool pages stay opt-in (`TF_QWEN_KV_PAGES=256`), dense and MoE alike; unset keeps a window each, as before, and `TF_QWEN_KV_PAGES_ALL` unset still leaves the engine's own stream a window. With pages set, `--parallel` is a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011). Thanks to @jschmied. +- Qwen3.5 and Qwen3.6 on CUDA: KV pool pages stay opt-in (`TF_QWEN_KV_PAGES=256`), dense and MoE alike; unset keeps a window each, as before, and `TF_QWEN_KV_PAGES_ALL` unset still leaves the engine's own stream a window. With pages set, `--parallel` is a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011, part of #1943). Thanks to @jschmied. - Metal: a request's finished stats list its prompt pass widths (`tensorfold.prefill_widths`, one entry per completed target pass) on every Metal text family except Flash Next, whose engine reports none. Telemetry only (#2293). diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index a0406a1e2e..52506315cb 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -41,9 +41,17 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B const attention = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen attention path:"} })); b.step("test-qwen35-attention-path-host", "Contiguous and paged CUDA attention selection and graph identity without devices").dependOn(&attention.step); step.dependOn(&attention.step); + const pages_filter = &.{"Qwen pages default:"}; + const pages_default = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = pages_filter })); + const pages_about = "TF_QWEN_KV_PAGES and _ALL unset as on main, --parallel a cap only with pages"; + b.step("test-qwen35-pages-default-host", pages_about).dependOn(&pages_default.step); + step.dependOn(&pages_default.step); const prompt_layout = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen prompt layout:"} })); b.step("test-qwen35-prompt-layout-host", "Actual whole and paged CUDA prompt ABI and stride calculations without devices").dependOn(&prompt_layout.step); step.dependOn(&prompt_layout.step); + const reuse_pages = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen paged reuse:"} })); + b.step("test-qwen35-reuse-pages-host", "Paged prompt-state saves and restores into host or device copies without devices").dependOn(&reuse_pages.step); + step.dependOn(&reuse_pages.step); const pending = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen pending keeps:"} })); pending.setEnvironmentVariable("TF_QWEN_KEEP_LATE", "1"); // deferred keeps are opt-in; the step tests them b.step("test-qwen35-pending-host", "Kept paths of every length inside the pending buffer and the log without devices").dependOn(&pending.step); diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 2719004aff..081e42b3f3 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -62,7 +62,9 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric int G, float scale, const int* __restrict__ PT, long long PS) { constexpr bool causal = true; constexpr int skip = 0; - constexpr bool HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int CAP = 0, PTOK = 1; const float *ks = nullptr, *vs = nullptr; const int HO = H; @@ -80,7 +82,9 @@ pattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict_ constexpr const int* PT = nullptr; // contiguous keys constexpr long long PS = 0; constexpr int PTOK = 1; - constexpr bool HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int CAP = 0; const float *ks = nullptr, *vs = nullptr; const int HO = H; @@ -115,7 +119,9 @@ full_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ o constexpr const int* PT = nullptr; // contiguous keys constexpr long long PS = 0; constexpr int PTOK = 1; - constexpr bool HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool P256 = false, HPAGED = false, F8 = false, HM = false; // bf16 keys, row-major + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int CAP = 0; const float *ks = nullptr, *vs = nullptr; const int H = 3 * HO, HK = H; @@ -132,7 +138,9 @@ __global__ void __launch_bounds__(32 * WARPS, 1) hpattn_whole_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, int G, float scale, int CAP) { - constexpr bool HPAGED = false, causal = true, F8 = false, HM = true; + constexpr bool P256 = false, HPAGED = false, causal = true, F8 = false, HM = true; + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int skip = 0; constexpr const int* PT = nullptr; constexpr long long PS = 0; @@ -149,7 +157,9 @@ __global__ void __launch_bounds__(32 * WARPS, 1) hpattn8_whole_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restrict__ kc, const uint8_t* __restrict__ vc, const float* __restrict__ ks, const float* __restrict__ vs, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, int G, float scale, int CAP) { - constexpr bool HPAGED = false, causal = true, F8 = true, HM = true; + constexpr bool P256 = false, HPAGED = false, causal = true, F8 = true, HM = true; + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int skip = 0; constexpr const int* PT = nullptr; constexpr long long PS = 0; @@ -165,7 +175,9 @@ __global__ void __launch_bounds__(32 * WARPS, 1) hpattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PTOK) { - constexpr bool HPAGED = true, causal = true, F8 = false, HM = true; + constexpr bool P256 = false, HPAGED = true, causal = true, F8 = false, HM = true; + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int skip = 0; const int HO = H; const float *ks = nullptr, *vs = nullptr; @@ -173,13 +185,46 @@ 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, IDS = false; + constexpr int PN = 0; + constexpr int skip = 0, PTOK = 256; + static_assert(PTOK % BN == 0, "a tile's keys, filled or not, in one page"); + const int HO = H; + const float *ks = nullptr, *vs = nullptr; + auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; +#include "prefill_attention_rows.cuh" +} + +// hpattn_p256_kernel with the chunk's PN page ids in shared memory past the queries (PN * 4 more dynamic bytes). +template +__global__ void __launch_bounds__(32 * WARPS, 1) +hpattn_p256_ids_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, + int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PN) { + constexpr bool P256 = true, HPAGED = true, causal = true, F8 = false, HM = true, IDS = true; + constexpr int skip = 0, PTOK = 256; + static_assert(PTOK % BN == 0, "a tile's keys, filled or not, in one page"); + const int HO = H; + const float *ks = nullptr, *vs = nullptr; + auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; +#include "prefill_attention_rows.cuh" +} + // hpattn_kernel over e4m3 rows with an fp32 scale a (key head, row) (qwen35_kv8.cu's layout, PS in bytes). template __global__ void __launch_bounds__(32 * WARPS, 1) hpattn8_kernel(const __nv_bfloat16* __restrict__ q, const uint8_t* __restrict__ kc, const uint8_t* __restrict__ vc, const float* __restrict__ ks, const float* __restrict__ vs, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, int G, float scale, int CAP, const int* __restrict__ PT, long long PS, int PTOK) { - constexpr bool HPAGED = true, causal = true, F8 = true, HM = true; + constexpr bool P256 = false, HPAGED = true, causal = true, F8 = true, HM = true; + constexpr bool IDS = false; // no page ids in shared memory + constexpr int PN = 0; constexpr int skip = 0; const int HO = H; auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; @@ -337,6 +382,12 @@ template __global__ void tf_prefill_attention::vision_kernel<128, 8>(const __nv_ template __global__ void tf_prefill_attention::hpattn_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float, int, const int* __restrict__, long long, int); +template __global__ void tf_prefill_attention::hpattn_p256_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, + int, float, int, const int* __restrict__, long long, int); +template __global__ void tf_prefill_attention::hpattn_p256_ids_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, + int, float, int, const int* __restrict__, long long, int); template __global__ void tf_prefill_attention::hpattn8_kernel<256, 8, 2, 2>(const __nv_bfloat16* __restrict__, const uint8_t* __restrict__, const uint8_t* __restrict__, const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float, int, const int* __restrict__, long long, int); diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index 53f2e0f0ab..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; @@ -26,6 +27,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 +38,8 @@ if (i < HALF * C && key < limit) { if constexpr (HPAGED) { const int row = key + (key < p0 ? 0 : skip); - const int64_t r = row_of(row), pg = page_of(row); + const bool near = HM && PT && row >= lo && row - lo < PTOK; + const int64_t r = near ? (int64_t)hk * PTOK + (row - lo) : row_of(row), pg = near ? first : page_of(row); pre[j] = __ldg(reinterpret_cast(src + pg + r * D + (i % C) * 8)); pre_s[j] = __ldg(reinterpret_cast(reinterpret_cast(sc) + pg) + r); } else { @@ -69,17 +74,33 @@ if constexpr (F8) { fetch(item); put(item); + } else if constexpr (P256) { + // pages of 256 positions: a tile's keys, filled or not, sit in one allocated page, so key0 needs no clamp + const __nv_bfloat16* src = reinterpret_cast(item % 4 < 2 ? kc : vc); + const int key0 = (item / 4) * BN + (item % 2) * HALF, p = key0 / PTOK; + uint4* dst = slots + (item % NS) * HALF * CS; + const __nv_bfloat16* base = src + (int64_t)(IDS ? ids[p] : PT[p]) * PS + ((int64_t)hk * PTOK + key0 % PTOK) * D; + for (int i = threadIdx.x; i < HALF * NB; i += THREADS) { + const int row = i / NB, chunk = i % NB; + cp16z(dst + sw<8 * CS>(i / NB, chunk), base + (i / NB) * D + chunk * 8, key0 + row < limit && (NB == C || chunk < C)); + } } else { const __nv_bfloat16* src = reinterpret_cast(item % 4 < 2 ? kc : vc); const int key0 = (item / 4) * BN + (item % 2) * HALF; uint4* dst = slots + (item % NS) * HALF * CS; + // HM with PT: the item's first page looked up once; a row past it (PTOK % HALF != 0) finds its own + const int lo = HM && PT ? min(key0 + skip, limit - 1 + skip) / PTOK * PTOK : 0; + const int64_t first = HM && PT ? (int64_t)PT[lo / PTOK] * PS : 0; for (int i = threadIdx.x; i < HALF * NB; i += THREADS) { const int key = key0 + i / NB, chunk = i % NB; const bool in = key < limit && (NB == C || chunk < C); // past the filled keys or the head dim: zeros const int row = in ? key + (key < p0 ? 0 : skip) : 0; if constexpr (HM) { - if constexpr (HPAGED) cp16z(dst + sw<8 * CS>(i / NB, chunk), src + page_of(row) + row_of(row) * D + chunk * 8, in); - else cp16z(dst + sw<8 * CS>(i / NB, chunk), src + row_of(row) * D + chunk * 8, in); + if constexpr (HPAGED) { + const bool near = PT && row >= lo && row - lo < PTOK; + const int64_t at = near ? first + ((int64_t)hk * PTOK + (row - lo)) * D : page_of(row) + row_of(row) * D; + cp16z(dst + sw<8 * CS>(i / NB, chunk), src + at + chunk * 8, in); + } else cp16z(dst + sw<8 * CS>(i / NB, chunk), src + row_of(row) * D + chunk * 8, in); } else { // a page's rows are contiguous: PT names each 512-row page, PS elements apart (null PT: one run) const int64_t at = PT ? (int64_t)PT[row / 512] * PS + (int64_t)(row % 512) * HK * D : (int64_t)row * HK * D; @@ -108,6 +129,10 @@ } commit(); } + if constexpr (IDS) { // every page id the chunk reads, once, so staging never waits on the table + for (int i = threadIdx.x; i < PN; i += THREADS) ids[i] = PT[i]; + __syncthreads(); + } #pragma unroll for (int i = 0; i < NS; ++i) { if (i < items) stage(i); diff --git a/zig/kernels/cuda/qwen35_attn_tiles.cuh b/zig/kernels/cuda/qwen35_attn_tiles.cuh index 2b679c44a5..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; @@ -45,8 +48,18 @@ struct Fp8Tiles { const int r = row(t0 + (tid >> 4) + THREADS / 16 * i); kr[i] = vr[i] = make_uint4(0u, 0u, 0u, 0u); if (r >= 0) { - kr[i] = *reinterpret_cast(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)); @@ -85,6 +98,9 @@ struct Bf16Tiles { size_t lblk, hoff, vgap; // a page's layer block (elements), the head's first row in it, keys to values int ptok; bool paged; + // the page a thread read last, kept across tiles: its first row is looked up once a page, not once a tile + const __nv_bfloat16* near = nullptr; + int lo = 0; __device__ Bf16Tiles(TileStore store, const Seg& sg, int layer, int kvh, int h) : s(store) { K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * HD + (size_t)h * sg.cap * HD; V = K + (size_t)kvh * sg.cap * HD; @@ -106,8 +122,18 @@ struct Bf16Tiles { for (int i = 0; i < TILE / THREADS; i++) { const int at = threadIdx.x + THREADS * i, u = at / CH, chunk = at % CH; const int r = row(t0 + u), pos = r < 0 ? 0 : r; // a key past the chunk reads row 0 and lands zeros - qmm_frag::cp16z(dst + swz(u, chunk), kp(pos) + chunk * 8, r >= 0); - qmm_frag::cp16z(dst + TILE + swz(u, chunk), vp(pos) + chunk * 8, r >= 0); + const __nv_bfloat16* k = K + (size_t)pos * HD; + const __nv_bfloat16* v = V + (size_t)pos * HD; + if (paged) { + if (near == nullptr || pos < lo || pos - lo >= ptok) { + lo = pos / ptok * ptok; + near = arena + (size_t)page.ids[lo / ptok] * page.stride + lblk + hoff * HD; + } + k = near + (size_t)(pos - lo) * HD; + v = k + vgap; + } + qmm_frag::cp16z(dst + swz(u, chunk), k + chunk * 8, r >= 0); + qmm_frag::cp16z(dst + TILE + swz(u, chunk), v + chunk * 8, r >= 0); } qmm_frag::commit(); } diff --git a/zig/src/core/contract/qwen_whole_records.zig b/zig/src/core/contract/qwen_whole_records.zig index 6de67b542a..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 }, @@ -47,11 +47,25 @@ 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 = 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 = 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 }, - .{ .regs = 231, .static = 0, .bound = 256 }, - .{ .regs = 231, .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/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..60e6527cab 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"); @@ -72,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 7d22ac08c7..8fff7322d7 100644 --- a/zig/src/families/qwen3_5/cuda_admit_test.zig +++ b/zig/src/families/qwen3_5/cuda_admit_test.zig @@ -56,3 +56,59 @@ 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)); +} + +/// 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 bff16dde85..99e57488f7 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -536,6 +536,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 47ca013c28..9a433a0ecc 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). +/// 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() !usize { - const v = @import("core").profile.get(.TF_QWEN_KV_PAGES) orelse return 0; - return parsePaging(v); + return pagingOf(@import("core").profile.get(.TF_QWEN_KV_PAGES)); +} + +pub fn pagingOf(v: ?[]const u8) !usize { + return parsePaging(v orelse return unset_pages); } 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=1: every stream on pages, the engine's own too: no window is reserved for it (unset: off). 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 false, "1"); } /// Pages a stream maps at most: a full window and one more forward. @@ -225,3 +234,18 @@ 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 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" { + 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/families/qwen3_5/cuda_prompt.zig b/zig/src/families/qwen3_5/cuda_prompt.zig index 36a9bc986b..d6e4fdd0bb 100644 --- a/zig/src/families/qwen3_5/cuda_prompt.zig +++ b/zig/src/families/qwen3_5/cuda_prompt.zig @@ -8,28 +8,34 @@ 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, 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"; 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 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), .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; } @@ -57,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); @@ -76,8 +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 }; - 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); + 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 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); } @@ -134,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/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); + } +} diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig index a0fb783fd7..12e26137a4 100644 --- a/zig/tests/cuda/contract_qwen_pages.zig +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -18,6 +18,8 @@ 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 + .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; @@ -196,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 { @@ -285,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)); @@ -413,6 +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 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; @@ -490,6 +503,22 @@ 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 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 var bad_r = flat; @@ -499,8 +528,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 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)\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}); } }