From b388b8cc3ca696009993bfb236bb7fdae5795195 Mon Sep 17 00:00:00 2001 From: Christoph Sturm Date: Sun, 26 Jul 2026 20:24:35 +0200 Subject: [PATCH] Fix Laguna KV ring commits for oversized prefills Laguna prefill attention stages all KV rows before committing them to its 512-row sliding-window ring. When a chunk exceeds that capacity, the existing kernel dispatches every token and races multiple writers onto each ring row, so older rows can survive and corrupt subsequent decoding. Commit only the newest cache-capacity rows, cover a 1024-token prefill with a bit-exact Metal regression test, and bump the checkpoint payload ABI because earlier Laguna checkpoints may persist raced ring contents. Tested on an Apple M4 Max with the Metal backend and Poolside Laguna S 2.1 Q4_K_M: make clean && make; ./ds4_test --metal-kernels; and ./ds4_test --server. The regression test reported 30,912 key and 30,464 value mismatches on the old kernel and zero mismatches with this fix. --- README.md | 3 +- ds4_kvstore.c | 2 +- ds4_server.c | 2 +- metal/laguna.metal | 17 +++-- tests/ds4_test.c | 171 +++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 187 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index caedc85bb..03711d877 100644 --- a/README.md +++ b/README.md @@ -1447,7 +1447,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 diff --git a/ds4_kvstore.c b/ds4_kvstore.c index 6b663b51b..1e728d539 100644 --- a/ds4_kvstore.c +++ b/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 diff --git a/ds4_server.c b/ds4_server.c index d8c11a48c..ddacc022c 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -17049,7 +17049,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)); diff --git a/metal/laguna.metal b/metal/laguna.metal index 2190eae89..b75a3d854 100644 --- a/metal/laguna.metal +++ b/metal/laguna.metal @@ -626,14 +626,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 { diff --git a/tests/ds4_test.c b/tests/ds4_test.c index 31f5d5e68..9780ea4ae 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -4286,6 +4286,176 @@ static void test_cuda_laguna_moe_decode_prefill(void) { #if defined(__APPLE__) +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; @@ -5810,6 +5980,7 @@ static void test_metal_kernel_group(void) { test_metal_persistent_zero_attention_mask_exact(); test_metal_zero_prefix_prefill_mask_cache_exact(); test_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();