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();