Fix Laguna KV ring commits for oversized prefills

README.md

@@ -1405,7 +1405,8 @@ The fixed header is little-endian:
 8   u32    cached token count
 12  u32    hit count
 16  u32    context size the snapshot was written for
-20  u8[4]  reserved
+20  u8     graph-payload ABI, currently 3
+21  u8[3]  reserved
 24  u64    creation Unix time
 32  u64    last-used Unix time
 40  u64    DS4 session payload byte count

ds4_kvstore.c

@@ -29,7 +29,7 @@
 /* Header byte 20 carries the graph-payload ABI.  It is separate from the outer
  * file version because the KVC envelope can remain stable while the serialized
  * ds4_session internals become unsafe to restore across runtime changes. */
-#define KV_CACHE_PAYLOAD_ABI 2u
+#define KV_CACHE_PAYLOAD_ABI 3u
 #define KV_CACHE_DEFAULT_MIN_TOKENS 512
 #define KV_CACHE_DEFAULT_COLD_MAX_TOKENS 30000
 /* Tokenizers may merge text across the prompt boundary. Trimming a small tail

ds4_server.c

@@ -17044,7 +17044,7 @@ static void test_kv_cache_lookup_rejects_stale_payload_abi(void) {
     if (fp) {
         uint8_t h[KV_CACHE_FIXED_HEADER];
         kv_fill_header(h, 2, KV_REASON_COLD, 0, 512, 0, 32768, 100, 100, 0);
-        h[20] = 0; /* pre-ABI-guard files used this byte as reserved zero. */
+        h[20] = 2; /* Prior Laguna checkpoints may contain raced SWA rows. */
         uint8_t text_len[4];
         le_put32(text_len, (uint32_t)strlen(text));
         TEST_ASSERT(fwrite(h, 1, sizeof(h), fp) == sizeof(h));

metal/laguna.metal

@@ -605,14 +605,21 @@ kernel void kernel_laguna_commit_kv_f16(
         device half *value_cache,
         uint gid [[thread_position_in_grid]]) {
     const uint width = args.n_head_kv * args.head_dim;
-    const uint values = args.n_tokens * width;
+    // Only the newest cache_cap rows survive in the ring. Dispatching older
+    // rows as well races multiple threads onto the same destination when a
+    // prefill chunk is wider than the sliding window.
+    const uint commit_tokens = min(args.n_tokens, args.cache_cap);
+    const uint token0 = args.n_tokens - commit_tokens;
+    const uint values = commit_tokens * width;
     if (gid >= values) return;
-    const uint token = gid / width;
-    const uint col = gid - token * width;
+    const uint local_token = gid / width;
+    const uint token = token0 + local_token;
+    const uint col = gid - local_token * width;
     const uint cache_row = (args.pos0 + token) % args.cache_cap;
     const uint64_t dst = (uint64_t)cache_row * width + col;
-    key_cache[dst] = staged_key[gid];
-    value_cache[dst] = staged_value[gid];
+    const uint64_t src = (uint64_t)token * width + col;
+    key_cache[dst] = staged_key[src];
+    value_cache[dst] = staged_value[src];
 }
 
 struct ds4_metal_args_laguna_attention {

tests/ds4_test.c

@@ -3750,6 +3750,176 @@ static void test_metal_laguna_gqa3_decode_numeric(void) {
     ds4_gpu_tensor_free(heads);
 }
 
+static void test_metal_laguna_prefill_ring_commit_exact(void) {
+    const uint32_t n_tokens = 1024;
+    const uint32_t cache_cap = 512;
+    const uint32_t n_head = 3;
+    const uint32_t n_head_kv = 1;
+    const uint32_t head_dim = 128;
+    const uint32_t cache_width = n_head_kv * head_dim;
+    const uint64_t q_values = (uint64_t)n_tokens * n_head * head_dim;
+    const uint64_t kv_values = (uint64_t)n_tokens * cache_width;
+    const uint64_t cache_values = (uint64_t)cache_cap * cache_width;
+
+    ds4_gpu_tensor *heads =
+        ds4_gpu_tensor_alloc(q_values * sizeof(float));
+    ds4_gpu_tensor *key_cache =
+        ds4_gpu_tensor_alloc(cache_values * sizeof(uint16_t));
+    ds4_gpu_tensor *value_cache =
+        ds4_gpu_tensor_alloc(cache_values * sizeof(uint16_t));
+    ds4_gpu_tensor *staged_key =
+        ds4_gpu_tensor_alloc(kv_values * sizeof(uint16_t));
+    ds4_gpu_tensor *staged_value =
+        ds4_gpu_tensor_alloc(kv_values * sizeof(uint16_t));
+    ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_values * sizeof(float));
+    ds4_gpu_tensor *k = ds4_gpu_tensor_alloc(kv_values * sizeof(float));
+    ds4_gpu_tensor *v = ds4_gpu_tensor_alloc(kv_values * sizeof(float));
+    ds4_gpu_tensor *gate =
+        ds4_gpu_tensor_alloc((uint64_t)n_tokens * n_head * sizeof(float));
+    float *q_host = calloc((size_t)q_values, sizeof(float));
+    float *k_host = malloc((size_t)kv_values * sizeof(float));
+    float *v_host = malloc((size_t)kv_values * sizeof(float));
+    float *gate_host =
+        calloc((size_t)n_tokens * n_head, sizeof(float));
+    uint16_t *key_actual =
+        malloc((size_t)cache_values * sizeof(uint16_t));
+    uint16_t *value_actual =
+        malloc((size_t)cache_values * sizeof(uint16_t));
+    uint16_t *key_staged =
+        malloc((size_t)kv_values * sizeof(uint16_t));
+    uint16_t *value_staged =
+        malloc((size_t)kv_values * sizeof(uint16_t));
+
+    TEST_ASSERT(heads != NULL);
+    TEST_ASSERT(key_cache != NULL);
+    TEST_ASSERT(value_cache != NULL);
+    TEST_ASSERT(staged_key != NULL);
+    TEST_ASSERT(staged_value != NULL);
+    TEST_ASSERT(q != NULL);
+    TEST_ASSERT(k != NULL);
+    TEST_ASSERT(v != NULL);
+    TEST_ASSERT(gate != NULL);
+    TEST_ASSERT(q_host != NULL);
+    TEST_ASSERT(k_host != NULL);
+    TEST_ASSERT(v_host != NULL);
+    TEST_ASSERT(gate_host != NULL);
+    TEST_ASSERT(key_actual != NULL);
+    TEST_ASSERT(value_actual != NULL);
+    TEST_ASSERT(key_staged != NULL);
+    TEST_ASSERT(value_staged != NULL);
+
+    size_t key_mismatches = 0;
+    size_t value_mismatches = 0;
+    if (heads && key_cache && value_cache && staged_key && staged_value &&
+        q && k && v && gate && q_host && k_host && v_host && gate_host &&
+        key_actual && value_actual && key_staged && value_staged) {
+        for (uint32_t token = 0; token < n_tokens; token++) {
+            for (uint32_t col = 0; col < cache_width; col++) {
+                const uint64_t index =
+                    (uint64_t)token * cache_width + col;
+                k_host[index] =
+                    (float)(token + 1u) / 64.0f +
+                    (float)col / 8192.0f;
+                v_host[index] =
+                    -(float)(token + 1u) / 32.0f -
+                    (float)col / 4096.0f;
+            }
+        }
+
+        TEST_ASSERT(ds4_gpu_tensor_write(
+                        q, 0, q_host, q_values * sizeof(float)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_write(
+                        k, 0, k_host, kv_values * sizeof(float)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_write(
+                        v, 0, v_host, kv_values * sizeof(float)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_write(
+                        gate,
+                        0,
+                        gate_host,
+                        (uint64_t)n_tokens * n_head * sizeof(float)) != 0);
+        TEST_ASSERT(ds4_gpu_laguna_attention_prefill_tensor(
+                        heads,
+                        key_cache,
+                        value_cache,
+                        staged_key,
+                        staged_value,
+                        q,
+                        k,
+                        v,
+                        gate,
+                        0,
+                        n_tokens,
+                        cache_cap,
+                        n_head,
+                        n_head_kv,
+                        head_dim,
+                        1.0f / sqrtf((float)head_dim),
+                        0) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_read(
+                        key_cache,
+                        0,
+                        key_actual,
+                        cache_values * sizeof(uint16_t)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_read(
+                        value_cache,
+                        0,
+                        value_actual,
+                        cache_values * sizeof(uint16_t)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_read(
+                        staged_key,
+                        0,
+                        key_staged,
+                        kv_values * sizeof(uint16_t)) != 0);
+        TEST_ASSERT(ds4_gpu_tensor_read(
+                        staged_value,
+                        0,
+                        value_staged,
+                        kv_values * sizeof(uint16_t)) != 0);
+
+        for (uint32_t row = 0; row < cache_cap; row++) {
+            const uint32_t latest_token = row + cache_cap;
+            for (uint32_t col = 0; col < cache_width; col++) {
+                const uint64_t cache_index =
+                    (uint64_t)row * cache_width + col;
+                const uint64_t source_index =
+                    (uint64_t)latest_token * cache_width + col;
+                if (key_actual[cache_index] != key_staged[source_index]) {
+                    key_mismatches++;
+                }
+                if (value_actual[cache_index] != value_staged[source_index]) {
+                    value_mismatches++;
+                }
+            }
+        }
+    }
+
+    fprintf(stderr,
+            "ds4-test: Laguna oversized prefill ring commit exact "
+            "key_mismatches=%zu value_mismatches=%zu\n",
+            key_mismatches,
+            value_mismatches);
+    TEST_ASSERT(key_mismatches == 0);
+    TEST_ASSERT(value_mismatches == 0);
+
+    free(value_staged);
+    free(key_staged);
+    free(value_actual);
+    free(key_actual);
+    free(gate_host);
+    free(v_host);
+    free(k_host);
+    free(q_host);
+    ds4_gpu_tensor_free(gate);
+    ds4_gpu_tensor_free(v);
+    ds4_gpu_tensor_free(k);
+    ds4_gpu_tensor_free(q);
+    ds4_gpu_tensor_free(staged_value);
+    ds4_gpu_tensor_free(staged_key);
+    ds4_gpu_tensor_free(value_cache);
+    ds4_gpu_tensor_free(key_cache);
+    ds4_gpu_tensor_free(heads);
+}
+
 static void test_metal_laguna_qk_norm_rope_pair_exact(void) {
     typedef struct {
         uint32_t n_tokens;
@@ -5274,6 +5444,7 @@ static void test_metal_kernel_group(void) {
     test_metal_persistent_zero_attention_mask_exact();
     test_metal_zero_prefix_prefill_mask_cache_exact();
     test_metal_laguna_gqa3_decode_numeric();
+    test_metal_laguna_prefill_ring_commit_exact();
     test_metal_laguna_qk_norm_rope_pair_exact();
     test_metal_hc_split_weighted_sum_norm_batch_exact();
     test_metal_output_hc_weights4_exact();