cuda window_attention: head dim 512 over a linear cache, for global attention

docs/kernel-contract.md

@@ -111,8 +111,11 @@ Mamba scan of a verify window (its host twin's bytes), qmm_group's cluster tile,
 bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a
 cluster, fused, or by `reduce_kernel<true>`), the checkpoint math's FP8 128-row prompt tiles (both slice meetings,
 bit-equal to the 64-row tiles) and NVFP4 experts: the decode-sized forms and the prompt tiles (up to down's NVFP4
-rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), and the
-reduce-scatter's `tf_rank_sum`.
+rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), the
+reduce-scatter's `tf_rank_sum`, and `window_attention`'s rows at head dims 256 and 512 (`window_rows.cu`: decode chunks
+and their merge, prompt tiles with P as bf16 or in two halves, a row's last key, windows over a ring or none over a
+linear cache), each output within float64's range over the spread its fp32 steps and tensor-core sums allow
+(`core/attention_range.zig`; seven controls, each one planted fault).
 
 ## CUDA records per nvcc release
 

zig/build/cuda.zig

@@ -321,6 +321,14 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     }) }));
     step.dependOn(&qsa_contract.step);
     flash_contract.dependOn(&qsa_contract.step);
+    const wrows_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{
+        .root_source_file = b.path("zig/tests/cuda/contract_window_rows.zig"),
+        .target = host,
+        .optimize = .debug,
+        .link_libc = true,
+        .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } },
+    }) }));
+    step.dependOn(&wrows_contract.step);
     const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)");
     const projection_step = b.step("test-cuda-projection-fixture", "Projection fixture dispatch controls with CPU-only CUDA calls");
     for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| {

zig/build/cuda_kernels.zig

@@ -82,6 +82,7 @@ pub const kernels = [_]Kernel{
     .{ .name = "tree_attention", .flags = &.{"-O3"} }, // ours: committed keys and paths (cuda/tree_attention.zig)
     .{ .name = "block_attention", .flags = &.{"-O3"} }, // ours: a drafter's block over its context (cuda/block_attention.zig)
     .{ .name = "window_attention", .flags = &.{"-O3"} }, // ours: sliding-window attention (cuda/window_attention.zig)
+    .{ .name = "window_rows", .flags = &.{"-O3"} }, // ours: window attention at head dims 256 and 512 (cuda/window_attention.zig)
     .{ .name = "moe_glue", .flags = &.{"-O3"} }, // ours: sigmoid top-k and the weighted combine (cuda/moe_glue.zig)
     .{ .name = "dense_rows", .flags = &.{"-O3"} }, // ours: bf16 rows times weight rows into fp32 (cuda/dense_rows.zig)
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts

zig/kernels/cuda/contract_controls.cu

@@ -85,3 +85,51 @@ extern "C" __global__ void tf_ctl_heavy(const float* X, float* Y) {
     for (int j = 0; j < 160; ++j) s += v[j] * v[159 - j];
     Y[t] = s;
 }
+
+// Gate 3: window_rows.cu's kernels at head dims 256 and 512, each with one planted fault (window_rows.cuh's Fault).
+#include "window_rows.cuh"
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_part_edge(const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                 const __nv_bfloat16* vc, const int* pos, const int* slot,
+                                                                                 float* po, float* pm, float* pl, int h, int hk, int win,
+                                                                                 long long ring, float scale) {
+    tf_wrows::part<256, tf_wrows::EDGE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_rescale(const float* po, const float* pm, const float* pl,
+                                                                                     __nv_bfloat16* out, int h, int hk, int chunks,
+                                                                                     const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                     const __nv_bfloat16* vc, const int* pos, const int* slot,
+                                                                                     int win, long long ring, float scale) {
+    tf_wrows::merge<256, tf_wrows::RESCALE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_trunc(const float* po, const float* pm, const float* pl,
+                                                                                   __nv_bfloat16* out, int h, int hk, int chunks,
+                                                                                   const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                   const __nv_bfloat16* vc, const int* pos, const int* slot,
+                                                                                   int win, long long ring, float scale) {
+    tf_wrows::merge<256, tf_wrows::TRUNC>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT, 1) tf_ctl_wrows_prompt_edge(const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                      const __nv_bfloat16* vc, __nv_bfloat16* out,
+                                                                                      const int* hi, int p0, int W, int H, int HK, int G,
+                                                                                      float scale, int win, int ring) {
+    tf_wrows::prompt::kernel<256, true, tf_wrows::EDGE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT, 1) tf_ctl_wrows_prompt_last(const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                      const __nv_bfloat16* vc, __nv_bfloat16* out,
+                                                                                      const int* hi, int p0, int W, int H, int HK, int G,
+                                                                                      float scale, int win, int ring) {
+    tf_wrows::prompt::kernel<256, true, tf_wrows::LAST>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}
+
+extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_rescale_512(const float* po, const float* pm, const float* pl,
+                                                                                         __nv_bfloat16* out, int h, int hk, int chunks,
+                                                                                         const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                                         const __nv_bfloat16* vc, const int* pos, const int* slot,
+                                                                                         int win, long long ring, float scale) {
+    tf_wrows::merge<512, tf_wrows::RESCALE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale);
+}

zig/kernels/cuda/window_rows.cu

@@ -0,0 +1,72 @@
+// window_attention's rows at head dims 256 and 512 (window_rows.cuh): decode partials and merge, prompt tiles.
+
+#include "window_rows.cuh"
+#include "window_rows512.cuh"
+
+using namespace tf_wrows;
+
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_part_256(const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                   const __nv_bfloat16* vc, const int* pos,
+                                                                   const int* slot, float* po, float* pm, float* pl,
+                                                                   int h, int hk, int win, long long ring, float scale) {
+    part<256, NONE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_merge_256(const float* po, const float* pm, const float* pl,
+                                                                    __nv_bfloat16* out, int h, int hk, int chunks,
+                                                                    const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                    const __nv_bfloat16* vc, const int* pos,
+                                                                    const int* slot, int win, long long ring,
+                                                                    float scale) {
+    merge<256, NONE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256(const __nv_bfloat16* q,
+                                                                        const __nv_bfloat16* kc,
+                                                                        const __nv_bfloat16* vc, __nv_bfloat16* out,
+                                                                        const int* hi, int p0, int W, int H, int HK,
+                                                                        int G, float scale, int win, int ring) {
+    prompt::kernel<256, false, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}
+
+extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256_p32(const __nv_bfloat16* q,
+                                                                            const __nv_bfloat16* kc,
+                                                                            const __nv_bfloat16* vc,
+                                                                            __nv_bfloat16* out, const int* hi, int p0,
+                                                                            int W, int H, int HK, int G, float scale,
+                                                                            int win, int ring) {
+    prompt::kernel<256, true, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}
+
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_part_512(const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                   const __nv_bfloat16* vc, const int* pos,
+                                                                   const int* slot, float* po, float* pm, float* pl,
+                                                                   int h, int hk, int win, long long ring, float scale) {
+    part_staged<512, NONE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_merge_512(const float* po, const float* pm, const float* pl,
+                                                                    __nv_bfloat16* out, int h, int hk, int chunks,
+                                                                    const __nv_bfloat16* q, const __nv_bfloat16* kc,
+                                                                    const __nv_bfloat16* vc, const int* pos,
+                                                                    const int* slot, int win, long long ring,
+                                                                    float scale) {
+    merge<512, NONE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale);
+}
+
+extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_512(const __nv_bfloat16* q,
+                                                                        const __nv_bfloat16* kc,
+                                                                        const __nv_bfloat16* vc, __nv_bfloat16* out,
+                                                                        const int* hi, int p0, int W, int H, int HK,
+                                                                        int G, float scale, int win, int ring) {
+    prompt::kernel<512, false, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}
+
+extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_512_p32(const __nv_bfloat16* q,
+                                                                            const __nv_bfloat16* kc,
+                                                                            const __nv_bfloat16* vc,
+                                                                            __nv_bfloat16* out, const int* hi, int p0,
+                                                                            int W, int H, int HK, int G, float scale,
+                                                                            int win, int ring) {
+    prompt::kernel<512, true, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+}