diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ad33d0d8..f9a2517a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -895,6 +895,8 @@ jobs: test_backend_vulkan_deltanet_unit \ test_fused_probe_agreement_unit \ test_backend_vulkan_dispatch_fail_unit \ + test_backend_vulkan_fallback_unit \ + test_backend_vulkan_narrowing_unit \ test_backend_vulkan_vram_budget_unit \ test_backend_vulkan_weight_index_unit \ test_gguf_release_range_unit; do @@ -960,6 +962,8 @@ jobs: test_backend_vulkan_deltanet_unit \ test_fused_probe_agreement_unit \ test_backend_vulkan_dispatch_fail_unit \ + test_backend_vulkan_fallback_unit \ + test_backend_vulkan_narrowing_unit \ test_backend_vulkan_vram_budget_unit \ test_backend_vulkan_weight_index_unit \ test_gguf_release_range_unit; do diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c6b784a..eb78df76 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,13 @@ minor release. ## [Unreleased] ### Added +- **Vulkan fallback accounting and strict mode** (#474). Every place work + leaves the GPU (a declined fused op, a host view or copy of mapped memory, a + host-path weight) is counted per site and printed under `GEIST_VK_VERBOSE`; + the first host-path linear names how many weights and bytes run there. + `GEIST_VK_STRICT=1` makes each of them an error that names the site. A + failed argmax dispatch now returns the dispatch error instead of sending the + arch to a host scan of logits the GPU never wrote. - **Vulkan device-memory budget** (#466). Device-local allocations (weight copies, KV cache, x ring) are checked against the heap before they are made; a model that does not fit fails with the MiB the allocation needs, the MiB in @@ -711,6 +718,18 @@ minor release. vtable. ### Fixed +- **Vulkan: a failed submit no longer hands out stale results; sizes past + 32 bits fail instead of wrapping** (#474). After a dropped batch, + `buffer_map` returns nullptr with `GEIST_E_BACKEND` as the backend error + (and the host-path linear, a host buffer copy and a download report the + failure); every transformer caller that maps a buffer now turns nullptr into + an error instead of dereferencing it. The push constants, dispatch sizes and + offsets in `ops.c` go through a checked narrowing (`vk_ckd_u32`), so a value + the shaders cannot index makes the op return `GEIST_E_INVALID_ARG`. The x + ring is checked against `max_m` x `n_in` at weight resolve (a weight that + could not be staged now fails the load rather than running as a silent host + linear) and is created there, with the argmax word, instead of on the first + decode. - **A session's KV cache ignored its own `max_seq_len`** (#577). It was sized from the model's cap, so on a model loaded with a 32768-token cap a 64-token session still held 32768 rows of KV (64 MiB in FP32 for a diff --git a/docs/BACKENDS.md b/docs/BACKENDS.md index dacd53c5..2ab1e733 100644 --- a/docs/BACKENDS.md +++ b/docs/BACKENDS.md @@ -71,6 +71,23 @@ back to per-row matvecs for prefill (#471). Details and the side-by-side profile: `benchmark/results/TERNARY.md`. Phase-by-phase lab log: [`../benchmark/results/VULKAN.md`](../benchmark/results/VULKAN.md). +Work that leaves the GPU is counted per site: a fused op the shaders decline +(the arch then runs it on the host), a host loop over mapped memory, a host +buffer copy, and weights whose dtype or row length has no GPU kernel (Q3_K, a +large F16/BF16 matrix, a row that is not a whole number of blocks), which run +on a host row-dequant path. `GEIST_VK_VERBOSE=1` prints the counters at +destroy, and the first host-path linear prints how many weights and MiB took +that path. `GEIST_VK_STRICT=1` turns each of these into an error naming the +site (a host-path weight is refused at load), so a coverage gap fails loudly +instead of showing up only as a slowdown. + +A batch whose submit fails (device lost, out of memory) is reported at the +next host access: `buffer_map` returns nullptr with `GEIST_E_BACKEND` as the +backend error, and argmax, downloads and host views return `GEIST_E_BACKEND`. +Sizes and offsets that do not fit the shaders' 32-bit indices make the op +return `GEIST_E_INVALID_ARG`, and a weight whose `n_in` would not fit the +192 MB activation ring at the 512-row batch limit fails the load. + ## Resident memory per backend What stays in RAM once a model is loaded: the weights in one or more layouts, diff --git a/src/archs/transformer/arch.c b/src/archs/transformer/arch.c index 6a57d348..c3630e7d 100644 --- a/src/archs/transformer/arch.c +++ b/src/archs/transformer/arch.c @@ -284,13 +284,17 @@ static const float *op_peek_logits(size_t *n_logits, void *session) { *n_logits = 0; return nullptr; } + float *p = (float *) st->backend->desc->vtbl->buffer_map(sess->scratch_logits); + if (p == nullptr) { + *n_logits = 0; /* no mapping, or a failed batch: the backend error says which */ + return nullptr; + } *n_logits = (size_t) st->vocab_size; - float *p = (float *) st->backend->desc->vtbl->buffer_map(sess->scratch_logits); /* The greedy argmax path skips the Gemma final-logit softcap (monotonic, * so the argmax is invariant and it saves ~262k tanhf/token on decode). * peek_logits exposes the VALUES to scoring/perplexity consumers, which * need the model-conformant softcapped logits — apply it lazily, once. */ - if (p != nullptr && st->config.logit_softcap > 0.0f && !sess->logits_softcapped) { + if (st->config.logit_softcap > 0.0f && !sess->logits_softcapped) { const float c = st->config.logit_softcap; for (size_t i = 0; i < (size_t) st->vocab_size; i++) { p[i] = tanhf(p[i] / c) * c; diff --git a/src/archs/transformer/arch_ops.c b/src/archs/transformer/arch_ops.c index f3cdb6b8..93389655 100644 --- a/src/archs/transformer/arch_ops.c +++ b/src/archs/transformer/arch_ops.c @@ -377,7 +377,10 @@ static enum geist_status prefill_text_batch_inner(struct transformer_arch_sessio sess->kv_len += chunk; if (sess->kv_kivi_enabled) { sess->kivi_residual_count += chunk; - transformer_kivi_drain_full(sess); + s = transformer_kivi_drain_full(sess); + if (s != GEIST_OK) { + return s; + } } /* 5. Mean pooling is the one finish that spans chunks: every @@ -482,6 +485,9 @@ enum geist_status transformer_verify_forward(struct transformer_arch_session *se /* 1. Embed all k tokens into scratch_h_a [k, HIDDEN]. */ { float *h_dst = (float *) v->buffer_map(sess->scratch_h_a); + if (h_dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t t = 0; t < k; t++) { enum geist_status s = dequant_one_row(be, &st->embed_table, (size_t) ids[t], h_dst + t * st->d_model); @@ -642,7 +648,10 @@ enum geist_status transformer_kv_truncate(struct transformer_arch_session *sess, } /* Truncate may settle residual into commit-safe territory * (verify_forward burst → accept → truncate at kv_len_old + a). */ - transformer_kivi_drain_full(sess); + const enum geist_status ds = transformer_kivi_drain_full(sess); + if (ds != GEIST_OK) { + return ds; + } } if (!keep_pending) { sess->logits_valid = false; @@ -713,6 +722,10 @@ enum geist_status transformer_prefill_audio_batch(struct transformer_arch_sessio { const size_t bytes = chunk * st->d_model * sizeof(float); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_h_a); + if (dst == nullptr) { + rc = GEIST_E_BACKEND; /* the backend said why */ + goto cleanup; + } memcpy(dst, (const uint8_t *) (soft_tokens + off * st->d_model), bytes); v->buffer_unmap(sess->scratch_h_a); } @@ -739,7 +752,10 @@ enum geist_status transformer_prefill_audio_batch(struct transformer_arch_sessio sess->kv_len += chunk; if (sess->kv_kivi_enabled) { sess->kivi_residual_count += chunk; - transformer_kivi_drain_full(sess); + rc = transformer_kivi_drain_full(sess); + if (rc != GEIST_OK) { + goto cleanup; + } } if (off + chunk == n) { @@ -855,6 +871,10 @@ enum geist_status apply_awq_to_state(struct transformer_arch_state *st, const ch goto cleanup; } float *g = (float *) v->buffer_map(L->attn_norm.buffer); + if (g == nullptr) { + rc = GEIST_E_BACKEND; /* the backend said why */ + goto cleanup; + } for (size_t j = 0; j < st->d_model; j++) g[j] /= s[j]; v->buffer_unmap(L->attn_norm.buffer); @@ -875,6 +895,10 @@ enum geist_status apply_awq_to_state(struct transformer_arch_state *st, const ch goto cleanup; } float *g = (float *) v->buffer_map(L->ffn_norm.buffer); + if (g == nullptr) { + rc = GEIST_E_BACKEND; /* the backend said why */ + goto cleanup; + } for (size_t j = 0; j < st->d_model; j++) g[j] /= s[j]; v->buffer_unmap(L->ffn_norm.buffer); diff --git a/src/archs/transformer/arch_state.c b/src/archs/transformer/arch_state.c index b95efc4e..5c57e460 100644 --- a/src/archs/transformer/arch_state.c +++ b/src/archs/transformer/arch_state.c @@ -689,6 +689,9 @@ allocate_runtime_session(struct transformer_arch_session *sess) { } { float *p = (float *) be->desc->vtbl->buffer_map(sess->scratch_ones_headdim_max); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t i = 0; i < head_dim_max; i++) { p[i] = 1.0f; } diff --git a/src/archs/transformer/forward.h b/src/archs/transformer/forward.h index 5de9b02d..5ca63506 100644 --- a/src/archs/transformer/forward.h +++ b/src/archs/transformer/forward.h @@ -74,7 +74,7 @@ transformer_check_kv_room(struct transformer_arch_session *sess, size_t n_new) { /* Drain the per-session KIVI residual ring across all non-shared layers, * if KIVI mode is enabled and residual_count >= R. No-op otherwise. */ -void transformer_kivi_drain_full(struct transformer_arch_session *sess); +[[nodiscard]] enum geist_status transformer_kivi_drain_full(struct transformer_arch_session *sess); /* Copy the residual rows of a just-pinned prefix (kivi_pin_tail), for * transformer_kivi_pin_restore to write back when a reset returns to it. diff --git a/src/archs/transformer/forward/head.c b/src/archs/transformer/forward/head.c index ee4b76a9..8b5b2568 100644 --- a/src/archs/transformer/forward/head.c +++ b/src/archs/transformer/forward/head.c @@ -87,6 +87,9 @@ static struct transformer_forward_profile g_head_profile = { } else { const uint8_t *src = (const uint8_t *) v->buffer_map(sess->scratch_h_b); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_h_a); + if (src == nullptr || dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(dst, src + row_idx * bytes, bytes); v->buffer_unmap(sess->scratch_h_b); v->buffer_unmap(sess->scratch_h_a); @@ -170,8 +173,11 @@ static struct transformer_forward_profile g_head_profile = { * ~262 144 × tanhf calls per token (~5% of decode on Gemma 4). */ const bool sampler_needs_softcap = sess->temperature > 0.0f; if (st->config.logit_softcap > 0.0f && sampler_needs_softcap) { - t0 = profile ? transformer_profile_now_ns() : 0; - float *p = (float *) v->buffer_map(sess->scratch_logits); + t0 = profile ? transformer_profile_now_ns() : 0; + float *p = (float *) v->buffer_map(sess->scratch_logits); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } const float c = st->config.logit_softcap; for (size_t i = 0; i < (size_t) st->vocab_size; i++) { p[i] = tanhf(p[i] / c) * c; @@ -261,6 +267,9 @@ finalize_logits_batch(struct transformer_arch_session *sess, size_t k, geist_tok const size_t bytes = k * st->d_model * sizeof(float); const uint8_t *src = (const uint8_t *) v->buffer_map(sess->scratch_h_b); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_h_a); + if (src == nullptr || dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(dst, src, bytes); v->buffer_unmap(sess->scratch_h_b); v->buffer_unmap(sess->scratch_h_a); diff --git a/src/archs/transformer/forward/kv_store.c b/src/archs/transformer/forward/kv_store.c index 22aa1871..4741bb8c 100644 --- a/src/archs/transformer/forward/kv_store.c +++ b/src/archs/transformer/forward/kv_store.c @@ -90,9 +90,15 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c const float *k_src = (const float *) v->buffer_map(sess->scratch_k); const float *v_src = (const float *) v->buffer_map(sess->scratch_v); + if (k_src == nullptr || v_src == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } if (ctx->kv_kivi_enabled) { - float *k_res = (float *) v->buffer_map(ctx->k_residual_buf); - float *v_res = (float *) v->buffer_map(ctx->v_residual_buf); + float *k_res = (float *) v->buffer_map(ctx->k_residual_buf); + float *v_res = (float *) v->buffer_map(ctx->v_residual_buf); + if (k_res == nullptr || v_res == nullptr) { + return GEIST_E_BACKEND; + } const size_t row_elems = kv_out; for (size_t t = 0; t < seq; t++) { const size_t res_idx = (q_position + t) - sess->kivi_drained_count; @@ -105,10 +111,13 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c /* Packed 4-bit: 2 values/byte into the half-size int8 slots. Same * per-token per-head scale + optional rotation as INT8. denom 7 → * scale = amax/7, values in [-7,7]. */ - uint8_t *k_dst = (uint8_t *) v->buffer_map(ctx->k_cache_q8_buf); - uint8_t *v_dst = (uint8_t *) v->buffer_map(ctx->v_cache_q8_buf); - float *k_sca = (float *) v->buffer_map(ctx->k_cache_scale_buf); - float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); + uint8_t *k_dst = (uint8_t *) v->buffer_map(ctx->k_cache_q8_buf); + uint8_t *v_dst = (uint8_t *) v->buffer_map(ctx->v_cache_q8_buf); + float *k_sca = (float *) v->buffer_map(ctx->k_cache_scale_buf); + float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); + if (k_dst == nullptr || v_dst == nullptr || k_sca == nullptr || v_sca == nullptr) { + return GEIST_E_BACKEND; + } const size_t row_elems = kv_out; const size_t scales_per_row = st->n_kv_heads; const bool rot = sess->kv_rot_enabled && fwht_supported(hd) && hd <= 512; @@ -145,10 +154,13 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c v->buffer_unmap(ctx->k_cache_scale_buf); v->buffer_unmap(ctx->v_cache_scale_buf); } else if (ctx->kv_int8_enabled) { - int8_t *k_dst = (int8_t *) v->buffer_map(ctx->k_cache_q8_buf); - int8_t *v_dst = (int8_t *) v->buffer_map(ctx->v_cache_q8_buf); - float *k_sca = (float *) v->buffer_map(ctx->k_cache_scale_buf); - float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); + int8_t *k_dst = (int8_t *) v->buffer_map(ctx->k_cache_q8_buf); + int8_t *v_dst = (int8_t *) v->buffer_map(ctx->v_cache_q8_buf); + float *k_sca = (float *) v->buffer_map(ctx->k_cache_scale_buf); + float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); + if (k_dst == nullptr || v_dst == nullptr || k_sca == nullptr || v_sca == nullptr) { + return GEIST_E_BACKEND; + } const size_t row_elems = kv_out; const size_t scales_per_row = st->n_kv_heads; /* Issue #61: rotate each K/V head row before quantizing. Q is @@ -202,8 +214,11 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c v->buffer_unmap(ctx->k_cache_scale_buf); v->buffer_unmap(ctx->v_cache_scale_buf); } else { - uint8_t *k_dst = (uint8_t *) v->buffer_map(ctx->k_cache_buf); - uint8_t *v_dst = (uint8_t *) v->buffer_map(ctx->v_cache_buf); + uint8_t *k_dst = (uint8_t *) v->buffer_map(ctx->k_cache_buf); + uint8_t *v_dst = (uint8_t *) v->buffer_map(ctx->v_cache_buf); + if (k_dst == nullptr || v_dst == nullptr) { + return GEIST_E_BACKEND; + } const size_t row_bytes = kv_out * sizeof(float); const size_t span_bytes = seq * row_bytes; memcpy(k_dst + q_position * row_bytes, (const uint8_t *) k_src, span_bytes); @@ -219,15 +234,19 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c /* fwht_orthonormal on each of the first n_rows rows (hd floats) of `b`: * Q before a backend's attention kernel and its output after, where the * cache holds rotated rows (GEIST_KV_ROT). H is its own inverse. */ -static void kv_rotate_rows(const struct geist_backend_vtbl *v, - size_t n_rows, - size_t hd, - struct geist_buffer *b) { +[[nodiscard]] static enum geist_status kv_rotate_rows(const struct geist_backend_vtbl *v, + size_t n_rows, + size_t hd, + struct geist_buffer *b) { float *p = (float *) v->buffer_map(b); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t r = 0; r < n_rows; r++) { fwht_orthonormal(hd, p + r * hd); } v->buffer_unmap(b); + return GEIST_OK; } enum geist_status transformer_kv_store_attention(struct transformer_layer_forward_ctx *ctx, @@ -253,7 +272,12 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar const float *krp = (const float *) v->buffer_map(ctx->k_residual_buf); const float *vrp = (const float *) v->buffer_map(ctx->v_residual_buf); float *outp = (float *) v->buffer_map(sess->scratch_attn); - float *scores = + if (qp == nullptr || kqp == nullptr || vqp == nullptr || kscp == nullptr || + kzep == nullptr || vscp == nullptr || vzep == nullptr || krp == nullptr || + vrp == nullptr || outp == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } + float *scores = (float *) frame_arena_alloc(&sess->scratch_arena, kv_len_now * sizeof(float), 16); if (scores == nullptr) { geist_backend_set_error(be, @@ -299,7 +323,10 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar const bool rot = sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; const size_t n_rows = ctx->seq * st->n_q_heads; if (rot) { - kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_q); + const enum geist_status rs = kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_q); + if (rs != GEIST_OK) { + return rs; + } } const int64_t n_kv = (int64_t) kv_len_now; const int64_t n_kh = st->n_kv_heads; @@ -321,7 +348,7 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar return s; } if (rot) { - kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_attn); + return kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_attn); } } else if (sess->kv_int4_packed_enabled) { float *qp = (float *) v->buffer_map(sess->scratch_q); @@ -330,8 +357,12 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar const float *k_scalep = (const float *) v->buffer_map(ctx->k_cache_scale_buf); const float *v_scalep = (const float *) v->buffer_map(ctx->v_cache_scale_buf); float *outp = (float *) v->buffer_map(sess->scratch_attn); - const bool rot = sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; - const size_t n_rows = ctx->seq * st->n_q_heads; + if (qp == nullptr || k_q4p == nullptr || v_q4p == nullptr || k_scalep == nullptr || + v_scalep == nullptr || outp == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } + const bool rot = sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; + const size_t n_rows = ctx->seq * st->n_q_heads; if (rot) { for (size_t r = 0; r < n_rows; r++) { fwht_orthonormal(ctx->hd, qp + r * ctx->hd); @@ -367,7 +398,10 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar const bool rot = sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; const size_t n_rows = ctx->seq * st->n_q_heads; if (rot) { - kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_q); + const enum geist_status rs = kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_q); + if (rs != GEIST_OK) { + return rs; + } } const int64_t n_kv = (int64_t) kv_len_now; const int64_t n_kh = st->n_kv_heads; @@ -388,7 +422,7 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar return s; } if (rot) { - kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_attn); + return kv_rotate_rows(v, n_rows, ctx->hd, sess->scratch_attn); } } else if (ctx->kv_int8_enabled) { float *qp = (float *) v->buffer_map(sess->scratch_q); @@ -397,6 +431,10 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar const float *k_scalep = (const float *) v->buffer_map(ctx->k_cache_scale_buf); const float *v_scalep = (const float *) v->buffer_map(ctx->v_cache_scale_buf); float *outp = (float *) v->buffer_map(sess->scratch_attn); + if (qp == nullptr || k_q8p == nullptr || v_q8p == nullptr || k_scalep == nullptr || + v_scalep == nullptr || outp == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } /* Issue #61: rotate Q by the same H used on K/V so QK scores are * unchanged; the kernel then quantizes rotated Q, and we rotate the * (V-rotated) output back below. H is its own inverse. */ diff --git a/src/archs/transformer/forward/layer.c b/src/archs/transformer/forward/layer.c index 96e22c39..ccb42092 100644 --- a/src/archs/transformer/forward/layer.c +++ b/src/archs/transformer/forward/layer.c @@ -558,7 +558,10 @@ enum geist_status transformer_compute_per_layer_input(struct transformer_arch_se } if (!on_device) { float *dst = (float *) v->buffer_map(sess->scratch_ple_lookup); - s = dequant_one_row(be, &st->ple_table, (size_t) token_id, dst); + if (dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } + s = dequant_one_row(be, &st->ple_table, (size_t) token_id, dst); if (s != GEIST_OK) { v->buffer_unmap(sess->scratch_ple_lookup); return s; @@ -600,6 +603,9 @@ enum geist_status transformer_compute_per_layer_input(struct transformer_arch_se } } else { float *p = (float *) v->buffer_map(per_layer_input_buf); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t i = 0; i < (size_t) st->ple_out; i++) { p[i] *= st->config.ple_model_proj_scale; } @@ -632,6 +638,9 @@ enum geist_status transformer_compute_per_layer_input(struct transformer_arch_se } } else { float *p = (float *) v->buffer_map(per_layer_input_buf); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t i = 0; i < (size_t) st->ple_out; i++) { p[i] *= st->config.ple_input_scale; } @@ -806,6 +815,9 @@ compute_per_layer_inputs_batch(struct transformer_arch_session *sess, } } else { float *p = (float *) v->buffer_map(out_buf); + if (p == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t i = 0; i < n * PLE_OUT; i++) { p[i] *= st->config.ple_model_proj_scale; } diff --git a/src/archs/transformer/forward/step.c b/src/archs/transformer/forward/step.c index d9bee7c1..7a28d435 100644 --- a/src/archs/transformer/forward/step.c +++ b/src/archs/transformer/forward/step.c @@ -38,12 +38,12 @@ #include #include -void transformer_kivi_drain_full(struct transformer_arch_session *sess) { +enum geist_status transformer_kivi_drain_full(struct transformer_arch_session *sess) { struct transformer_arch_state *st = sess->model; if (!sess->kv_kivi_enabled) - return; + return GEIST_OK; if (sess->kivi_residual_count < KIVI_K_GROUP_SIZE) - return; + return GEIST_OK; struct geist_backend *be = st->backend; const struct geist_backend_vtbl *v = be->desc->vtbl; const size_t R = KIVI_K_GROUP_SIZE; @@ -61,6 +61,11 @@ void transformer_kivi_drain_full(struct transformer_arch_session *sess) { float *k_zeros = (float *) v->buffer_map(sess->k_kivi_zeros[li]); float *v_scales = (float *) v->buffer_map(sess->v_kivi_scales[li]); float *v_zeros = (float *) v->buffer_map(sess->v_kivi_zeros[li]); + if (k_res == nullptr || v_res == nullptr || k_q4 == nullptr || v_q4 == nullptr || + k_scales == nullptr || k_zeros == nullptr || v_scales == nullptr || + v_zeros == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } kivi_drain_one_layer(sess->kivi_drained_count, sess->kivi_residual_count, R, @@ -86,6 +91,7 @@ void transformer_kivi_drain_full(struct transformer_arch_session *sess) { sess->kivi_drained_count += R; sess->kivi_residual_count -= R; } + return GEIST_OK; } /* A pinned prefix that is no whole number of groups keeps its last @@ -94,7 +100,8 @@ void transformer_kivi_drain_full(struct transformer_arch_session *sess) { * moves later tokens over their rows, but a reset to the prefix reads them * from those rows again. So pinning copies them out, and reset writes them * back: the prefix is then the ring and groups it was at pin time. */ -static void kivi_pin_copy(struct transformer_arch_session *sess, bool save) { +[[nodiscard]] static enum geist_status kivi_pin_copy(struct transformer_arch_session *sess, + bool save) { const struct transformer_arch_state *st = sess->model; const struct geist_backend_vtbl *v = st->backend->desc->vtbl; float *tail = sess->kivi_pin_tail; @@ -106,11 +113,15 @@ static void kivi_pin_copy(struct transformer_arch_session *sess, bool save) { struct geist_buffer *ring[2] = {sess->k_residual[li], sess->v_residual[li]}; for (size_t i = 0; i < 2; i++) { float *rows = (float *) v->buffer_map(ring[i]); + if (rows == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(save ? tail : rows, save ? rows : tail, len * sizeof *tail); v->buffer_unmap(ring[i]); tail += len; } } + return GEIST_OK; } enum geist_status transformer_kivi_pin_save(struct transformer_arch_session *sess) { @@ -140,15 +151,16 @@ enum geist_status transformer_kivi_pin_save(struct transformer_arch_session *ses return GEIST_E_OOM; } sess->kivi_pin_rows = rows; - kivi_pin_copy(sess, true); - return GEIST_OK; + return kivi_pin_copy(sess, true); } void transformer_kivi_pin_restore(struct transformer_arch_session *sess) { /* Reset has put the counters back to the prefix: exactly the rows the * pin copied are residual again. Anything else is not that prefix. */ if (sess->kivi_pin_rows != 0 && sess->kivi_pin_rows == sess->kivi_residual_count) { - kivi_pin_copy(sess, false); + /* Reset has no status: an unmappable ring leaves the rows as they + * are, and the next drain or attention map reports the backend. */ + (void) kivi_pin_copy(sess, false); } } @@ -194,6 +206,9 @@ void transformer_kivi_pin_restore(struct transformer_arch_session *sess) { } else { const uint8_t *src = (const uint8_t *) v->buffer_map(initial_h_buf); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_h_a); + if (src == nullptr || dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(dst, src, seq * row_bytes_h); v->buffer_unmap(initial_h_buf); v->buffer_unmap(sess->scratch_h_a); @@ -234,6 +249,9 @@ void transformer_kivi_pin_restore(struct transformer_arch_session *sess) { } else { const uint8_t *src = (const uint8_t *) v->buffer_map(per_layer_input_buf); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_ple_lookup); + if (src == nullptr || dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } for (size_t t = 0; t < seq; t++) { memcpy(dst + t * row_bytes_ple, src + t * row_bytes_per_tok_ple + li * row_bytes_ple, @@ -274,6 +292,9 @@ void transformer_kivi_pin_restore(struct transformer_arch_session *sess) { } else { const uint8_t *src = (const uint8_t *) v->buffer_map(h_in); uint8_t *dst = (uint8_t *) v->buffer_map(out_h_buf); + if (src == nullptr || dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(dst, src, seq * row_bytes_h); v->buffer_unmap(h_in); v->buffer_unmap(out_h_buf); @@ -316,8 +337,11 @@ void transformer_kivi_pin_restore(struct transformer_arch_session *sess) { : GEIST_OK; } - float *dst = (float *) v->buffer_map(out_h_buf); - enum geist_status s = dequant_one_row(be, &st->embed_table, (size_t) token_id, dst); + float *dst = (float *) v->buffer_map(out_h_buf); + if (dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } + enum geist_status s = dequant_one_row(be, &st->embed_table, (size_t) token_id, dst); if (s != GEIST_OK) { v->buffer_unmap(out_h_buf); return s; @@ -393,7 +417,10 @@ transformer_run_one_step(struct transformer_arch_session *sess, sess->kv_len = q_position + 1; if (sess->kv_kivi_enabled) { sess->kivi_residual_count += 1; - transformer_kivi_drain_full(sess); + s = transformer_kivi_drain_full(sess); + if (s != GEIST_OK) { + return s; + } } sess->next_token_pending = best_id; sess->logits_valid = true; @@ -443,6 +470,9 @@ enum geist_status transformer_advance_audio_token(struct transformer_arch_sessio { const size_t bytes = (size_t) st->d_model * sizeof(float); uint8_t *dst = (uint8_t *) v->buffer_map(sess->scratch_h_a); + if (dst == nullptr) { + return GEIST_E_BACKEND; /* the backend said why */ + } memcpy(dst, h_in_host, bytes); v->buffer_unmap(sess->scratch_h_a); } diff --git a/src/backends/vulkan/lifecycle.c b/src/backends/vulkan/lifecycle.c index 86921ded..0559973a 100644 --- a/src/backends/vulkan/lifecycle.c +++ b/src/backends/vulkan/lifecycle.c @@ -446,6 +446,8 @@ static size_t vk_parse_bytes(const char *v) { st->backend = be; st->profile_enabled = getenv("GEIST_VK_PROFILE") != nullptr; st->pq2_f32_acc = getenv("GEIST_VK_PQ2_F32_ACC") != nullptr; + const char *strict = getenv("GEIST_VK_STRICT"); + st->strict = strict != nullptr && strcmp(strict, "0") != 0; st->vram_budget = vk_parse_bytes(getenv("GEIST_VK_VRAM_BUDGET")); /* Default on since #501's rollout validated cleanly (two models, * two GPUs, several misaligned chunk sizes); GEIST_VK_ATTN_CM=0 is the @@ -510,6 +512,14 @@ void vk_destroy(struct geist_backend *be) { (unsigned long long) st->stat_cpu_falls, (unsigned long long) st->stat_barriers, (unsigned long long) st->stat_barriers_elided); + for (size_t i = 0; i < VK_FB_COUNT; i++) { + if (st->fallbacks[i] != 0) { + fprintf(stderr, + "geist vulkan fallback: %s %llu\n", + vk_fallback_name((enum vk_fb) i), + (unsigned long long) st->fallbacks[i]); + } + } fprintf(stderr, "geist vulkan dset cache: %llu hits, %llu misses; submit+wait %.1f ms\n", (unsigned long long) st->stat_dset_hits, diff --git a/src/backends/vulkan/ops.c b/src/backends/vulkan/ops.c index 168dec5b..40e3fc34 100644 --- a/src/backends/vulkan/ops.c +++ b/src/backends/vulkan/ops.c @@ -61,41 +61,54 @@ struct vk_qinfo { size_t n_out, size_t blocks_per_row) { struct vk_state *st = be->state; + size_t x_bytes, y_bytes; + uint32_t n_in32, n_out32, bpr32, m32; + if (ckd_mul(&x_bytes, m, n_in) || ckd_mul(&x_bytes, x_bytes, sizeof(float)) || + ckd_mul(&y_bytes, m, n_out) || ckd_mul(&y_bytes, y_bytes, sizeof(float)) || + vk_ckd_u32(n_in, &n_in32) || vk_ckd_u32(n_out, &n_out32) || + vk_ckd_u32(blocks_per_row, &bpr32) || vk_ckd_u32(m, &m32)) { + return vk_too_wide(be, "linear"); + } vk_seq_flush(st); /* host x/y round-trip — must not interleave with a batch */ /* x_stage: the GPU reads it hot (GEMM B tiles), the host only writes — * SCRATCH role makes it BAR-eligible. y_stage stays in system RAM * (the host reads results back; CPU reads from BAR are uncached). */ - enum geist_status s = - vk_stage_reserve_role(be, &st->x_stage, m * n_in * sizeof(float), GEIST_BUFFER_SCRATCH); + enum geist_status s = vk_stage_reserve_role(be, &st->x_stage, x_bytes, GEIST_BUFFER_SCRATCH); if (s == GEIST_OK) { - s = vk_stage_reserve(be, &st->y_stage, m * n_out * sizeof(float)); + s = vk_stage_reserve(be, &st->y_stage, y_bytes); } if (s != GEIST_OK) { return s; } - memcpy(st->x_stage->mapped, x, m * n_in * sizeof(float)); + memcpy(st->x_stage->mapped, x, x_bytes); const VkDescriptorBufferInfo binfo[3] = { {.buffer = st->x_stage->buf, .range = VK_WHOLE_SIZE}, {.buffer = wbuf->buf, .range = VK_WHOLE_SIZE}, {.buffer = st->y_stage->buf, .range = VK_WHOLE_SIZE}, }; - const struct vk_push push = {.n_in = (uint32_t) n_in, - .n_out = (uint32_t) n_out, - .blocks_per_row = (uint32_t) blocks_per_row, - .rows = (uint32_t) m, - .x_stride = (uint32_t) n_in, - .y_stride = (uint32_t) n_out}; + const struct vk_push push = {.n_in = n_in32, + .n_out = n_out32, + .blocks_per_row = bpr32, + .rows = m32, + .x_stride = n_in32, + .y_stride = n_out32}; enum vk_pipe eff = pipe; - uint32_t gx = vk_linear_gx(pipe, (uint32_t) n_out); - uint32_t gy = vk_linear_gy(pipe, (uint32_t) m); - vk_linear_cm_route(st, &eff, (uint32_t) m, (uint32_t) n_out, &gx, &gy); + uint32_t gx = vk_linear_gx(pipe, n_out32); + uint32_t gy = vk_linear_gy(pipe, m32); + vk_linear_cm_route(st, &eff, m32, n_out32, &gx, &gy); s = vk_seq_dispatch(be, eff, binfo, &push, sizeof(push), gx, gy, 1); if (s != GEIST_OK) { return s; } vk_seq_flush(st); - memcpy(y, st->y_stage->mapped, m * n_out * sizeof(float)); + if (st->seq_failed) { + /* y_stage holds no result of this dispatch. The flag stays set: the + * resolved kernels have no status, so the next readback reports it. */ + geist_backend_set_error(be, GEIST_E_BACKEND, "vulkan: a submitted batch failed"); + return GEIST_E_BACKEND; + } + memcpy(y, st->y_stage->mapped, y_bytes); return GEIST_OK; } @@ -208,6 +221,15 @@ static void vk_w_cpu_mN(size_t m, struct vk_state *st = be->state; const size_t n_in = (size_t) w->n_in; const size_t n_out = (size_t) w->n_out; + st->fallbacks[VK_FB_HOST_LINEAR]++; + if (!st->host_weights_noted) { + st->host_weights_noted = true; + fprintf(stderr, + "geist vulkan: %zu weight(s), %zu MiB, run on the host row-dequant path " + "(no GPU kernel for their dtype or row length)\n", + st->host_weights, + st->host_weight_bytes >> 20); + } /* Row scratch lives in the backend state (grown on demand, freed at * destroy): the resolved kernels are allocation-free in steady state, and * a failed grow zeroes y and says why instead of leaving it unwritten. */ @@ -245,6 +267,29 @@ vk_w_cpu_m1(const float *x, const struct geist_weight *w, struct geist_backend * vk_w_cpu_mN(1, x, w, be, y); } +/* Install the host row-dequant kernels on `w`, or refuse them under + * GEIST_VK_STRICT=1. Counted so the first host linear can say how much of + * the model left the GPU (#474 item 9). */ +[[nodiscard]] static enum geist_status vk_resolve_host(struct geist_backend *be, + struct geist_weight *w) { + struct vk_state *st = be->state; + if (st->strict) { + geist_backend_set_error(be, + GEIST_E_BACKEND, + "vulkan: GEIST_VK_STRICT=1 and weight dtype %u (%dx%d) has no " + "GPU kernel", + (unsigned) w->dtype, + (int) w->n_out, + (int) w->n_in); + return GEIST_E_BACKEND; + } + st->host_weights++; + st->host_weight_bytes += w->raw_nbytes; + w->linear_m1 = vk_w_cpu_m1; + w->linear_mN = vk_w_cpu_mN; + return GEIST_OK; +} + /* ---- resolve_weight: upload GPU-supported dtypes to VRAM, register, */ /* install kernels; CPU fallback for the rest. */ @@ -395,6 +440,34 @@ vk_repack_weight(const struct geist_weight *w, size_t bytes, bool *failed) { } } +/* linear_t stages up to VK_MAX_M rows of x into the device x ring. Check that + * against this weight's n_in at load, so a batch that cannot fit fails the load + * instead of turning into a per-call UNSUPPORTED and a silent host linear, and + * create the ring and the argmax word here, off the hot path (#474 item 6). */ +[[nodiscard]] static enum geist_status vk_ring_reserve(struct geist_backend *be, size_t n_in) { + struct vk_state *st = be->state; + size_t need; + if (ckd_mul(&need, (size_t) VK_MAX_M, n_in) || ckd_mul(&need, need, sizeof(float)) || + need > VK_XRING_CAP) { + geist_backend_set_error(be, + GEIST_E_INVALID_ARG, + "vulkan: n_in=%zu at max_m=%d does not fit the %u MB x ring", + n_in, + (int) VK_MAX_M, + (unsigned) (VK_XRING_CAP >> 20)); + return GEIST_E_INVALID_ARG; + } + enum geist_status s = GEIST_OK; + if (st->xring == nullptr) { + s = vk_buffer_create( + be, VK_XRING_CAP, GEIST_BUFFER_SCRATCH, GEIST_MEMORY_DEVICE, &st->xring); + } + if (s == GEIST_OK && st->argmax_out == nullptr) { + s = vk_buffer_create(be, 16, GEIST_BUFFER_STAGING, GEIST_MEMORY_AUTO, &st->argmax_out); + } + return s; +} + [[nodiscard]] static enum geist_status vk_resolve_weight(struct geist_backend *be, struct geist_weight *w) { struct vk_state *st = be->state; @@ -423,9 +496,7 @@ vk_repack_weight(const struct geist_weight *w, size_t bytes, bool *failed) { } [[fallthrough]]; case GEIST_DTYPE_Q3_K: - w->linear_m1 = vk_w_cpu_m1; - w->linear_mN = vk_w_cpu_mN; - return GEIST_OK; + return vk_resolve_host(be, w); default: geist_backend_set_error(be, GEIST_E_UNSUPPORTED, @@ -441,12 +512,14 @@ vk_repack_weight(const struct geist_weight *w, size_t bytes, bool *failed) { * by block. Every dtype but the two native k-quants has a CPU dequant * row (vk_dequant_row) and keeps working through it. */ if (w->dtype != GEIST_DTYPE_Q4_K && w->dtype != GEIST_DTYPE_Q6_K) { - w->linear_m1 = vk_w_cpu_m1; - w->linear_mN = vk_w_cpu_mN; - return GEIST_OK; + return vk_resolve_host(be, w); } return GEIST_E_UNSUPPORTED; } + const enum geist_status rs = vk_ring_reserve(be, (size_t) w->n_in); + if (rs != GEIST_OK) { + return rs; + } /* Upload to VRAM and register. An existing entry for the same host * pointer is REPLACED, not reused: the same address can carry new bytes * after a model reload (or a freed+remalloc'd test blob) — the latest @@ -462,7 +535,9 @@ vk_repack_weight(const struct geist_weight *w, size_t bytes, bool *failed) { if (nw == nullptr) { return GEIST_E_OOM; } - memcpy(nw, st->weights, st->n_weights * sizeof(*nw)); + if (st->n_weights > 0) { /* weights is nullptr before the first grow */ + memcpy(nw, st->weights, st->n_weights * sizeof(*nw)); + } geist_backend_free(be, st->weights); st->weights = nw; st->cap_weights = cap; @@ -517,8 +592,9 @@ vk_repack_weight(const struct geist_weight *w, size_t bytes, bool *failed) { return GEIST_OK; } -static uint32_t vk_groups(size_t n) { - return (uint32_t) ((n + 255) / 256); +/* 256-lane workgroups covering n (already narrowed through vk_ckd_u32). */ +static uint32_t vk_groups(uint32_t n) { + return n / 256u + (n % 256u != 0u ? 1u : 0u); } /* Dispatch geometry of the linear pipes. matvec q4k/q6k: 8 rows per @@ -636,10 +712,10 @@ void vk_linear_cm_route(struct vk_state *st, !vk_t_geom(y, &ry, &cy, &sy)) { return GEIST_E_UNSUPPORTED; } - uint32_t cols = 0; + size_t cols = 0; if (sa != ca || sb != cb || sy != cy) { /* mixed contiguous/strided operands: unify on the strided cols */ - cols = (uint32_t) (sa != ca ? ca : (sb != cb ? cb : cy)); + cols = sa != ca ? ca : (sb != cb ? cb : cy); if ((sa != ca && ca != cols) || (sb != cb && cb != cols) || (sy != cy && cy != cols) || n % cols != 0) { return GEIST_E_UNSUPPORTED; @@ -660,17 +736,14 @@ void vk_linear_cm_route(struct vk_state *st, !vk_tensor_gpu(y, &bi[2], &off[2])) { return GEIST_E_UNSUPPORTED; } - const uint32_t push[8] = {(uint32_t) n, - off[0], - off[1], - off[2], - cols, - (uint32_t) sa, - (uint32_t) sb, - (uint32_t) sy}; - const struct vk_access acc[3] = { + uint32_t push[8] = {0, off[0], off[1], off[2]}; + if (vk_ckd_u32(n, &push[0]) || vk_ckd_u32(cols, &push[4]) || vk_ckd_u32(sa, &push[5]) || + vk_ckd_u32(sb, &push[6]) || vk_ckd_u32(sy, &push[7])) { + return vk_too_wide(be, "elementwise"); + } + const struct vk_access acc[3] = { vk_acc_tensor(a, false), vk_acc_tensor(b, false), vk_acc_tensor(y, true)}; - return vk_seq_dispatch_acc(be, pipe, bi, acc, push, sizeof(push), vk_groups(n), 1, 1); + return vk_seq_dispatch_acc(be, pipe, bi, acc, push, sizeof(push), vk_groups(push[0]), 1, 1); } /* GPU-first attempt for the unary elementwise family (gelu_tanh, silu, @@ -687,9 +760,12 @@ void vk_linear_cm_route(struct vk_state *st, !vk_tensor_gpu(y, &bi[1], &off[1])) { return GEIST_E_UNSUPPORTED; } - const uint32_t push[4] = {(uint32_t) n, off[0], off[1], 0}; - const struct vk_access acc[2] = {vk_acc_tensor(x, false), vk_acc_tensor(y, true)}; - return vk_seq_dispatch_acc(be, pipe, bi, acc, push, sizeof(push), vk_groups(n), 1, 1); + uint32_t push[4] = {0, off[0], off[1], 0}; + if (vk_ckd_u32(n, &push[0])) { + return vk_too_wide(be, "elementwise"); + } + const struct vk_access acc[2] = {vk_acc_tensor(x, false), vk_acc_tensor(y, true)}; + return vk_seq_dispatch_acc(be, pipe, bi, acc, push, sizeof(push), vk_groups(push[0]), 1, 1); } /* The 3-buffer elementwise ops (vk_try_ew3 above, vk_ew3_cpu below). */ @@ -921,17 +997,15 @@ vk_silu(struct geist_backend *be, const struct geist_tensor *x, struct geist_ten !vk_tensor_gpu(gate, &bi[2], &og)) { return GEIST_E_UNSUPPORTED; } - const size_t n = rows * heads * head_dim; - const uint32_t push[9] = {(uint32_t) rows, - (uint32_t) heads, - (uint32_t) head_dim, - oj, - oq, - og, - (uint32_t) joint->stride[0], - (uint32_t) q->stride[0], - (uint32_t) gate->stride[0]}; - const struct vk_access acc[3] = { + uint32_t n, push[9] = {0, 0, 0, oj, oq, og}; + if (vk_ckd_u32(rows * heads * head_dim, &n) || vk_ckd_u32(rows, &push[0]) || + vk_ckd_u32(heads, &push[1]) || vk_ckd_u32(head_dim, &push[2]) || + vk_ckd_u32((size_t) joint->stride[0], &push[6]) || + vk_ckd_u32((size_t) q->stride[0], &push[7]) || + vk_ckd_u32((size_t) gate->stride[0], &push[8])) { + return vk_too_wide(be, "attn_qgate_split"); + } + const struct vk_access acc[3] = { vk_acc_tensor(joint, false), vk_acc_tensor(q, true), vk_acc_tensor(gate, true)}; return vk_seq_dispatch_acc( be, VK_PIPE_QGATE_SPLIT, bi, acc, push, sizeof(push), vk_groups(n), 1, 1); @@ -950,14 +1024,17 @@ vk_silu(struct geist_backend *be, const struct geist_tensor *x, struct geist_ten if (feat != 0 && n % feat == 0 && vk_t_n(w) == feat && vk_t_n(y) == n && vk_tensor_gpu(x, &bi[0], &off[0]) && vk_tensor_gpu(w, &bi[1], &off[1]) && vk_tensor_gpu(y, &bi[2], &off[2])) { - const struct { + struct { uint32_t rows, feat, x, w, y; float eps; - } push = {(uint32_t) (n / feat), (uint32_t) feat, off[0], off[1], off[2], eps}; + } push = {0, 0, off[0], off[1], off[2], eps}; + if (vk_ckd_u32(n / feat, &push.rows) || vk_ckd_u32(feat, &push.feat)) { + return vk_too_wide(be, "rmsnorm"); + } const struct vk_access acc[3] = { vk_acc_tensor(x, false), vk_acc_tensor(w, false), vk_acc_tensor(y, true)}; return vk_seq_dispatch_acc( - be, VK_PIPE_RMSNORM, bi, acc, &push, sizeof(push), (uint32_t) (n / feat), 1, 1); + be, VK_PIPE_RMSNORM, bi, acc, &push, sizeof(push), push.rows, 1, 1); } } size_t nx = 0, nw = 0, ny = 0; @@ -1003,21 +1080,18 @@ vk_silu(struct geist_backend *be, const struct geist_tensor *x, struct geist_ten vk_t_n(sin) != 0 && rot != 0 && rot % 2 == 0 && rot <= (size_t) x->shape[2] && vk_tensor_gpu(x, &bi[0], &off[0]) && vk_tensor_gpu(cos, &bi[1], &off[1]) && vk_tensor_gpu(sin, &bi[2], &off[2])) { - const size_t seq = (size_t) x->shape[0]; - const size_t heads = (size_t) x->shape[1]; - const size_t hd = (size_t) x->shape[2]; - const size_t pairs = seq * heads * (rot / 2); - const uint32_t push[7] = {(uint32_t) pairs, - (uint32_t) heads, - (uint32_t) hd, - off[0], - off[1], - off[2], - (uint32_t) rot}; - const struct vk_access acc[3] = { + const size_t seq = (size_t) x->shape[0]; + const size_t heads = (size_t) x->shape[1]; + const size_t hd = (size_t) x->shape[2]; + uint32_t push[7] = {0, 0, 0, off[0], off[1], off[2]}; + if (vk_ckd_u32(seq * heads * (rot / 2), &push[0]) || vk_ckd_u32(heads, &push[1]) || + vk_ckd_u32(hd, &push[2]) || vk_ckd_u32(rot, &push[6])) { + return vk_too_wide(be, "rope_apply"); + } + const struct vk_access acc[3] = { vk_acc_tensor(x, true), vk_acc_tensor(cos, false), vk_acc_tensor(sin, false)}; return vk_seq_dispatch_acc( - be, VK_PIPE_ROPE, bi, acc, push, sizeof(push), vk_groups(pairs), 1, 1); + be, VK_PIPE_ROPE, bi, acc, push, sizeof(push), vk_groups(push[0]), 1, 1); } } size_t nx = 0, nc = 0, ns = 0; @@ -1054,12 +1128,14 @@ vk_silu(struct geist_backend *be, const struct geist_tensor *x, struct geist_ten return GEIST_E_UNSUPPORTED; } const size_t rows = (size_t) x->shape[0], cols = (size_t) x->shape[1], - stride = (size_t) x->stride[0]; - bi[1] = bi[0]; - const uint32_t push[4] = {(uint32_t) cols, off[0], (uint32_t) stride, 0}; - const struct vk_access acc[2] = {vk_acc_tensor(x, true), vk_acc_tensor(x, true)}; - return vk_seq_dispatch_acc( - be, VK_PIPE_ACT_QUANT, bi, acc, push, sizeof(push), (uint32_t) rows, 1, 1); + stride = (size_t) x->stride[0]; + bi[1] = bi[0]; + uint32_t push[4] = {0, off[0], 0, 0}, rows32; + if (vk_ckd_u32(cols, &push[0]) || vk_ckd_u32(stride, &push[2]) || vk_ckd_u32(rows, &rows32)) { + return vk_too_wide(be, "bitnet_act_quant"); + } + const struct vk_access acc[2] = {vk_acc_tensor(x, true), vk_acc_tensor(x, true)}; + return vk_seq_dispatch_acc(be, VK_PIPE_ACT_QUANT, bi, acc, push, sizeof(push), rows32, 1, 1); } /* Largest Hadamard block the shader stages in shared memory. */ @@ -1096,31 +1172,30 @@ enum { VK_HADAMARD_MAX_BLOCK = 1024 }; off[1] = 0; } const float scale = 1.0f / sqrtf((float) block); - uint32_t push[11] = {(uint32_t) width, - (uint32_t) block, - (uint32_t) (width / block), - (uint32_t) args->perm_hd, - (uint32_t) args->perm_nk, - (uint32_t) rep, + uint32_t push[11] = {0, + 0, + 0, + 0, + 0, + 0, (args->inverse ? 1u : 0u) | (has_signs ? 2u : 0u), off[0], off[1], off[2], 0}; + uint32_t groups; + if (vk_ckd_u32(width, &push[0]) || vk_ckd_u32(block, &push[1]) || + vk_ckd_u32(width / block, &push[2]) || vk_ckd_u32(args->perm_hd, &push[3]) || + vk_ckd_u32(args->perm_nk, &push[4]) || vk_ckd_u32(rep, &push[5]) || + vk_ckd_u32(rows * (width / block), &groups)) { + return vk_too_wide(be, "hadamard_rotate"); + } memcpy(&push[10], &scale, sizeof scale); const struct vk_access acc[3] = {vk_acc_tensor(args->x, false), has_signs ? vk_acc_tensor(args->signs, false) : vk_acc_tensor(args->x, false), vk_acc_tensor(args->y, true)}; - return vk_seq_dispatch_acc(be, - VK_PIPE_HADAMARD, - bi, - acc, - push, - sizeof(push), - (uint32_t) (rows * (width / block)), - 1, - 1); + return vk_seq_dispatch_acc(be, VK_PIPE_HADAMARD, bi, acc, push, sizeof(push), groups, 1, 1); } size_t nx = 0, ns = 0, ny = 0; const float *xp = vk_tensor_host(args->x, &nx); @@ -1173,11 +1248,13 @@ enum { VK_HADAMARD_MAX_BLOCK = 1024 }; q->shape[0] == 1 && k->dtype == GEIST_DTYPE_F16 && (size_t) k->shape[0] > 192 && q->shape[2] <= 512 && vk_t_n(q) != 0 && vk_t_n16(k) != 0 && stt->pipes[VK_PIPE_ATTN_PART_F16] != VK_NULL_HANDLE) { - const uint32_t qh = (uint32_t) q->shape[1]; - const uint32_t hd = (uint32_t) q->shape[2]; - const uint32_t n_kv = (uint32_t) k->shape[0]; - const uint32_t kvh = (uint32_t) k->shape[1]; - const uint32_t n_chunks = (n_kv + 127u) / 128u; + uint32_t qh, hd, n_kv, kvh, qpos, win; + if (vk_ckd_u32((size_t) q->shape[1], &qh) || vk_ckd_u32((size_t) q->shape[2], &hd) || + vk_ckd_u32((size_t) k->shape[0], &n_kv) || vk_ckd_u32((size_t) k->shape[1], &kvh) || + vk_ckd_u32(q_offset, &qpos) || vk_ckd_u32(sliding_window, &win)) { + return vk_too_wide(be, "attention"); + } + const uint32_t n_chunks = n_kv / 128u + (n_kv % 128u != 0u ? 1u : 0u); const size_t part_bytes = (size_t) qh * n_chunks * (hd + 2u) * 4u; VkDescriptorBufferInfo bq, bk, bv, bo; uint32_t qo, ko, vo, oo; @@ -1196,18 +1273,11 @@ enum { VK_HADAMARD_MAX_BLOCK = 1024 }; ((qo | ko | vo) & 3u) != 0u /* 4-wide K/V/Q streams */) { goto attn_generic; } - const uint32_t po = (uint32_t) (stt->xring_used / 4u); - const uint32_t push1[11] = {n_kv, - qh, - kvh, - hd, - (uint32_t) q_offset, - (uint32_t) sliding_window, - qo, - ko, - vo, - po, - n_chunks}; + uint32_t po; /* < VK_XRING_CAP / 4; checked all the same */ + if (vk_ckd_u32(stt->xring_used / 4u, &po)) { + return vk_too_wide(be, "attention"); + } + const uint32_t push1[11] = {n_kv, qh, kvh, hd, qpos, win, qo, ko, vo, po, n_chunks}; VkDescriptorBufferInfo bi1[4] = { bq, bk, bv, {.buffer = stt->xring->buf, .range = VK_WHOLE_SIZE}}; const struct vk_access acc1[4] = {vk_acc_tensor(q, false), @@ -1241,23 +1311,16 @@ attn_generic:; vk_tensor_gpu(out, &bi[3], &off[3]) && /* f16 kernel streams q/k/v as 4-wide vectors */ (!kv16 || ((off[0] | off[1] | off[2]) & 3u) == 0u)) { - const uint32_t n_q = (uint32_t) q->shape[0]; - const uint32_t qh = (uint32_t) q->shape[1]; - const uint32_t hd = (uint32_t) q->shape[2]; - const uint32_t n_kv = (uint32_t) k->shape[0]; - const uint32_t kvh = (uint32_t) k->shape[1]; - const uint32_t push[11] = {n_q, - n_kv, - qh, - kvh, - hd, - (uint32_t) q_offset, - (uint32_t) sliding_window, - off[0], - off[1], - off[2], - off[3]}; - const struct vk_access acc[4] = { + uint32_t n_q, qh, hd, n_kv, kvh, qpos, win; + if (vk_ckd_u32((size_t) q->shape[0], &n_q) || vk_ckd_u32((size_t) q->shape[1], &qh) || + vk_ckd_u32((size_t) q->shape[2], &hd) || vk_ckd_u32((size_t) k->shape[0], &n_kv) || + vk_ckd_u32((size_t) k->shape[1], &kvh) || vk_ckd_u32(q_offset, &qpos) || + vk_ckd_u32(sliding_window, &win)) { + return vk_too_wide(be, "attention"); + } + const uint32_t push[11] = { + n_q, n_kv, qh, kvh, hd, qpos, win, off[0], off[1], off[2], off[3]}; + const struct vk_access acc[4] = { vk_acc_tensor(q, false), kv16 ? vk_acc_tensor16(k, false) : vk_acc_tensor(k, false), kv16 ? vk_acc_tensor16(v, false) : vk_acc_tensor(v, false), @@ -1278,7 +1341,7 @@ attn_generic:; acc, push, sizeof(push), - (n_q + 15u) / 16u, + n_q / 16u + (n_q % 16u != 0u ? 1u : 0u), qh, 1); } @@ -1333,23 +1396,19 @@ attn_generic:; if (feat != 0 && n % feat == 0 && vk_t_n(w) == feat && vk_t_n(res) == n && vk_t_n(y) == n && vk_tensor_gpu(x, &bi[0], &off[0]) && vk_tensor_gpu(w, &bi[1], &off[1]) && vk_tensor_gpu(res, &bi[2], &off[2]) && vk_tensor_gpu(y, &bi[3], &off[3])) { - const struct { + struct { uint32_t rows, feat, x, w, r, y; float eps; - } push = {(uint32_t) (n / feat), (uint32_t) feat, off[0], off[1], off[2], off[3], eps}; + } push = {0, 0, off[0], off[1], off[2], off[3], eps}; + if (vk_ckd_u32(n / feat, &push.rows) || vk_ckd_u32(feat, &push.feat)) { + return vk_too_wide(be, "rmsnorm_add"); + } const struct vk_access acc[4] = {vk_acc_tensor(x, false), vk_acc_tensor(w, false), vk_acc_tensor(res, false), vk_acc_tensor(y, true)}; - return vk_seq_dispatch_acc(be, - VK_PIPE_RMSNORM_ADD, - bi, - acc, - &push, - sizeof(push), - (uint32_t) (n / feat), - 1, - 1); + return vk_seq_dispatch_acc( + be, VK_PIPE_RMSNORM_ADD, bi, acc, &push, sizeof(push), push.rows, 1, 1); } } /* CPU fallback: y = res + rmsnorm(x) * w */ @@ -1389,13 +1448,16 @@ attn_generic:; uint32_t off[2]; if (n != 0 && n == vk_t_n(y) && vk_tensor_gpu(x, &bi[0], &off[0]) && vk_tensor_gpu(y, &bi[1], &off[1])) { - const struct { + struct { uint32_t n, x, y; float scale; - } push = {(uint32_t) n, off[0], off[1], scale}; + } push = {0, off[0], off[1], scale}; + if (vk_ckd_u32(n, &push.n)) { + return vk_too_wide(be, "scale_f32"); + } const struct vk_access acc[2] = {vk_acc_tensor(x, false), vk_acc_tensor(y, true)}; return vk_seq_dispatch_acc( - be, VK_PIPE_SCALE, bi, acc, &push, sizeof(push), vk_groups(n), 1, 1); + be, VK_PIPE_SCALE, bi, acc, &push, sizeof(push), vk_groups(push.n), 1, 1); } } size_t nx = 0, ny = 0; @@ -1427,12 +1489,15 @@ vk_argmax_f32(struct geist_backend *be, const struct geist_tensor *logits, int32 } } bi[1] = (VkDescriptorBufferInfo) {.buffer = st->argmax_out->buf, .range = VK_WHOLE_SIZE}; - const uint32_t push[3] = {(uint32_t) n, off[0], 0}; - const struct vk_access acc[2] = {vk_acc_tensor(logits, false), vk_acc_all(true)}; + uint32_t push[3] = {0, off[0], 0}; + if (vk_ckd_u32(n, &push[0])) { + return vk_too_wide(be, "argmax"); + } + const struct vk_access acc[2] = {vk_acc_tensor(logits, false), vk_acc_all(true)}; enum geist_status s = vk_seq_dispatch_acc(be, VK_PIPE_ARGMAX, bi, acc, push, sizeof(push), 1, 1, 1); if (s != GEIST_OK) { - return GEIST_E_UNSUPPORTED; + return s; /* a failed dispatch is an error, not a host scan */ } vk_seq_flush(st); /* the one intended sync point per decoded token */ if (vk_seq_take_failure(st) != GEIST_OK) { @@ -1465,15 +1530,20 @@ vk_argmax_f32(struct geist_backend *be, const struct geist_tensor *logits, int32 !vk_tensor_gpu(t_y, &bi[2], &yo)) { return GEIST_E_UNSUPPORTED; } - const uint32_t n_in = (uint32_t) w->n_in; - const uint32_t n_out = (uint32_t) w->n_out; - uint32_t x_stride; + uint32_t n_in, n_out, m32, x_stride, y_stride; + if (vk_ckd_u32((size_t) w->n_in, &n_in) || vk_ckd_u32((size_t) w->n_out, &n_out) || + vk_ckd_u32(m, &m32) || + vk_ckd_u32(t_y->ndim >= 2 ? (size_t) t_y->stride[t_y->ndim - 2] : n_out, &y_stride)) { + return vk_too_wide(be, "linear_t"); + } if (t_x->buffer->device_mem) { /* BAR-resident activations: bind in place, no staging copy. */ if (!vk_tensor_gpu(t_x, &bi[0], &xo)) { return GEIST_E_UNSUPPORTED; } - x_stride = t_x->ndim >= 2 ? (uint32_t) t_x->stride[t_x->ndim - 2] : n_in; + if (vk_ckd_u32(t_x->ndim >= 2 ? (size_t) t_x->stride[t_x->ndim - 2] : n_in, &x_stride)) { + return vk_too_wide(be, "linear_t"); + } } else { if (!vk_xring_stage(be, t_x, m, n_in, &xo)) { return GEIST_E_UNSUPPORTED; @@ -1482,24 +1552,29 @@ vk_argmax_f32(struct geist_backend *be, const struct geist_tensor *logits, int32 x_stride = n_in; /* ring copy is contiguous */ } bi[1] = (VkDescriptorBufferInfo) {.buffer = wbuf->buf, .range = VK_WHOLE_SIZE}; - const uint32_t y_stride = t_y->ndim >= 2 ? (uint32_t) t_y->stride[t_y->ndim - 2] : n_out; - const struct vk_push push = {.n_in = n_in, - .n_out = n_out, - .blocks_per_row = n_in / qi.block_elems, - .rows = (uint32_t) m, - .x_offset = xo, - .y_offset = yo, - .x_stride = x_stride, - .y_stride = y_stride}; + const struct vk_push push = {.n_in = n_in, + .n_out = n_out, + .blocks_per_row = n_in / qi.block_elems, + .rows = m32, + .x_offset = xo, + .y_offset = yo, + .x_stride = x_stride, + .y_stride = y_stride}; if (m > 1 && !st->gemm_sg32) { /* The register-tiled GEMMs assume 32-lane subgroups (see * vk_linear_run). Elsewhere run the size-agnostic matvec once per - * batch row: correct on any device. */ - for (size_t r = 0; r < m; ++r) { + * batch row: correct on any device. The last row's offsets bound + * every row's. */ + uint32_t last; + if (vk_ckd_u32((size_t) xo + (size_t) (m32 - 1u) * x_stride, &last) || + vk_ckd_u32((size_t) yo + (size_t) (m32 - 1u) * y_stride, &last)) { + return vk_too_wide(be, "linear_t"); + } + for (uint32_t r = 0; r < m32; ++r) { struct vk_push row = push; row.rows = 1; - row.x_offset = push.x_offset + (uint32_t) r * push.x_stride; - row.y_offset = push.y_offset + (uint32_t) r * push.y_stride; + row.x_offset = push.x_offset + r * push.x_stride; + row.y_offset = push.y_offset + r * push.y_stride; const struct vk_access racc[3] = { vk_acc((uint64_t) row.x_offset * 4u, (uint64_t) n_in * 4u, false), vk_acc_all(false), @@ -1514,11 +1589,11 @@ vk_argmax_f32(struct geist_backend *be, const struct geist_tensor *logits, int32 } enum vk_pipe lpipe = m == 1 ? mv : mm; uint32_t gx = vk_linear_gx(lpipe, n_out); - uint32_t gy = vk_linear_gy(lpipe, (uint32_t) m); + uint32_t gy = vk_linear_gy(lpipe, m32); /* Tensor-core path for conforming GEMMs (shaders assume w_offset == 0, * which holds for all registry uploads). */ if (m > 1) { - vk_linear_cm_route(st, &lpipe, (uint32_t) m, n_out, &gx, &gy); + vk_linear_cm_route(st, &lpipe, m32, n_out, &gx, &gy); } const struct vk_access acc[3] = { t_x->buffer->device_mem ? vk_acc_tensor(t_x, false) @@ -1613,9 +1688,13 @@ vk_embedding_lookup_scaled(struct geist_backend *be, ? (const uint8_t *) embed_table->buffer->host_alias + embed_table->offset : nullptr; struct geist_buffer *wbuf = host != nullptr ? vk_weight_lookup(st, host) : nullptr; - if (wbuf == nullptr && host != nullptr && embed_table->buffer->bytes > embed_table->offset) { + struct vk_qinfo eqi; + if (wbuf == nullptr && host != nullptr && embed_table->buffer->bytes > embed_table->offset && + vk_qinfo_for((enum geist_dtype) embed_table->dtype, &eqi)) { /* An untied table (separate output.weight) is never resolved by the - * arch layer: register it now, so the repacked dtypes can be read. */ + * arch layer: register it now, so the repacked dtypes can be read. + * Only dtypes with a GPU copy: a host-path table is not a weight + * that left the GPU, and strict mode must not refuse it here. */ struct geist_weight ew = {.raw = host, .raw_nbytes = embed_table->buffer->bytes - embed_table->offset, .n_in = (int32_t) d, @@ -1640,9 +1719,11 @@ vk_embedding_lookup_scaled(struct geist_backend *be, (embed_table->buffer->base_off + embed_table->offset) % 4 != 0) { return GEIST_E_UNSUPPORTED; } - bi[0] = (VkDescriptorBufferInfo) {.buffer = embed_table->buffer->buf, - .range = VK_WHOLE_SIZE}; - w_elem_off = (uint32_t) (embed_table->buffer->base_off + embed_table->offset); + bi[0] = (VkDescriptorBufferInfo) {.buffer = embed_table->buffer->buf, + .range = VK_WHOLE_SIZE}; + if (vk_ckd_u32(embed_table->buffer->base_off + embed_table->offset, &w_elem_off)) { + return vk_too_wide(be, "embedding_lookup_scaled"); + } } uint32_t yo; if (vk_t_n(out) != (size_t) d || !vk_tensor_gpu(out, &bi[1], &yo)) { @@ -1650,25 +1731,23 @@ vk_embedding_lookup_scaled(struct geist_backend *be, } /* Blocks per row, from the same table the linears use (0 for the * non-block dtypes: F32/F16/BF16). */ - struct vk_qinfo qi; - const uint32_t bpr = dtype_code >= 3 && vk_qinfo_for(embed_table->dtype, &qi) - ? (uint32_t) d / qi.block_elems - : 0; - const struct { + struct { uint32_t n_in, token, dtype, bpr, w_byte, y; float scale; uint32_t n_rows; - } push = {(uint32_t) d, - (uint32_t) token_id, - dtype_code, - bpr, - w_elem_off, - yo, - scale, - (uint32_t) vocab}; + } push = {0, 0, dtype_code, 0, w_elem_off, yo, scale, 0}; + /* 0 <= token_id < vocab was checked above. */ + if (vk_ckd_u32((size_t) d, &push.n_in) || vk_ckd_u32((size_t) token_id, &push.token) || + vk_ckd_u32((size_t) vocab, &push.n_rows)) { + return vk_too_wide(be, "embedding_lookup_scaled"); + } + struct vk_qinfo qi; + if (dtype_code >= 3 && vk_qinfo_for(embed_table->dtype, &qi)) { + push.bpr = push.n_in / qi.block_elems; + } const struct vk_access acc[2] = {vk_acc_all(false), vk_acc_tensor(out, true)}; return vk_seq_dispatch_acc( - be, VK_PIPE_EMBED, bi, acc, &push, sizeof(push), vk_groups((size_t) d), 1, 1); + be, VK_PIPE_EMBED, bi, acc, &push, sizeof(push), vk_groups(push.n_in), 1, 1); } /* Geometry the decode gate/up kernels run, shared by the probe and both @@ -1706,10 +1785,13 @@ static bool vk_ffn_gate_up_geometry_ok(bool with_norm, up_w->shape[0])) { return GEIST_E_UNSUPPORTED; } - const uint32_t n_out = (uint32_t) gate_w->shape[0]; - const uint32_t n_in = (uint32_t) gate_w->shape[1]; - struct geist_buffer *gbuf = vk_weight_of(st, gate_w); - struct geist_buffer *ubuf = vk_weight_of(st, up_w); + uint32_t n_out, n_in; /* in range per vk_ffn_gate_up_geometry_ok; checked all the same */ + if (vk_ckd_u32((size_t) gate_w->shape[0], &n_out) || + vk_ckd_u32((size_t) gate_w->shape[1], &n_in)) { + return vk_too_wide(be, "ffn_gate_up"); + } + struct geist_buffer *gbuf = vk_weight_of(st, gate_w); + struct geist_buffer *ubuf = vk_weight_of(st, up_w); VkDescriptorBufferInfo bi[4]; uint32_t xo, yo; if (gbuf == nullptr || ubuf == nullptr || vk_t_n(y) < n_out || !vk_tensor_gpu(y, &bi[3], &yo)) { @@ -1765,8 +1847,11 @@ static bool vk_ffn_gate_up_geometry_ok(bool with_norm, up_w->shape[0])) { return GEIST_E_UNSUPPORTED; } - const uint32_t n_out = (uint32_t) gate_w->shape[0]; - const uint32_t n_in = (uint32_t) gate_w->shape[1]; + uint32_t n_out, n_in; /* in range per vk_ffn_gate_up_geometry_ok; checked all the same */ + if (vk_ckd_u32((size_t) gate_w->shape[0], &n_out) || + vk_ckd_u32((size_t) gate_w->shape[1], &n_in)) { + return vk_too_wide(be, "ffn_norm_gate_up"); + } if (vk_t_n(norm_w) != n_in) { return GEIST_E_UNSUPPORTED; } @@ -1818,10 +1903,13 @@ static bool vk_ffn_gate_up_geometry_ok(bool with_norm, gate_w->ndim != 2 || proj_w->ndim != 2) { return GEIST_E_UNSUPPORTED; } - const uint32_t d_in = (uint32_t) gate_w->shape[1]; - const uint32_t hpl = (uint32_t) gate_w->shape[0]; - const uint32_t feat = (uint32_t) proj_w->shape[0]; - if ((uint32_t) x->shape[1] != d_in || (uint32_t) proj_w->shape[1] != hpl || + uint32_t d_in, hpl, feat; + if (vk_ckd_u32((size_t) gate_w->shape[1], &d_in) || + vk_ckd_u32((size_t) gate_w->shape[0], &hpl) || + vk_ckd_u32((size_t) proj_w->shape[0], &feat)) { + return vk_too_wide(be, "ple_block"); + } + if (x->shape[1] != gate_w->shape[1] || proj_w->shape[1] != gate_w->shape[0] || vk_t_n(norm_w) != feat || vk_t_n(res) < feat || vk_t_n(y) < feat || vk_t_n(gate_scratch) < hpl || vk_t_n(proj_scratch) < feat) { return GEIST_E_UNSUPPORTED; @@ -1903,9 +1991,11 @@ static constexpr uint32_t VK_QKV_PREP_MAX_HEAD_DIM = 512; (!kv16 && (k_cache->dtype != GEIST_DTYPE_F32 || v_cache->dtype != GEIST_DTYPE_F32)))) { return GEIST_E_UNSUPPORTED; } - const uint32_t seq = (uint32_t) q->shape[0]; - const uint32_t qh = (uint32_t) q->shape[1]; - const uint32_t hd = (uint32_t) q->shape[2]; + uint32_t seq, qh, hd, qpos; + if (vk_ckd_u32((size_t) q->shape[0], &seq) || vk_ckd_u32((size_t) q->shape[1], &qh) || + vk_ckd_u32((size_t) q->shape[2], &hd) || vk_ckd_u32(q_position, &qpos)) { + return vk_too_wide(be, "attn_qkv_prep"); + } if (hd == 0u || (hd % 2u) != 0u || hd > VK_QKV_PREP_MAX_HEAD_DIM || vk_t_n(q) == 0) { return GEIST_E_UNSUPPORTED; } @@ -1932,7 +2022,9 @@ static constexpr uint32_t VK_QKV_PREP_MAX_HEAD_DIM = 512; b_vw.buffer != bi[0].buffer || b_kw.buffer != bi[1].buffer) { return GEIST_E_UNSUPPORTED; } - kvh = (uint32_t) k->shape[1]; + if (vk_ckd_u32((size_t) k->shape[1], &kvh)) { + return vk_too_wide(be, "attn_qkv_prep"); + } } else { bi[4] = bi[0]; /* unused bindings — anything valid */ bi[5] = bi[0]; @@ -1945,7 +2037,7 @@ static constexpr uint32_t VK_QKV_PREP_MAX_HEAD_DIM = 512; qh, kvh, hd, - (uint32_t) q_position, + qpos, has_kv ? 1u : 0u, qo, ko, @@ -2011,12 +2103,15 @@ static constexpr uint32_t VK_QKV_PREP_MAX_HEAD_DIM = 512; geist_backend_set_error(be, GEIST_E_UNSUPPORTED, "vulkan kv_append_f16: bad inputs"); return GEIST_E_UNSUPPORTED; } - const uint32_t push[5] = {(uint32_t) n, - kso, - vso, - kdo + (uint32_t) (q_position * kv_row), - vdo + (uint32_t) (q_position * kv_row)}; - const struct vk_access acc[4] = { + /* The cache row this append starts at, in f16 elements past the views. */ + size_t pos_elems; + uint32_t push[5] = {0, kso, vso, 0, 0}; + if (ckd_mul(&pos_elems, q_position, kv_row) || vk_ckd_u32(n, &push[0]) || + vk_ckd_u32((size_t) kdo + pos_elems, &push[3]) || + vk_ckd_u32((size_t) vdo + pos_elems, &push[4])) { + return vk_too_wide(be, "kv_append_f16"); + } + const struct vk_access acc[4] = { vk_acc_tensor(k_src, false), vk_acc_tensor(v_src, false), vk_acc(k_cache->buffer->base_off + k_cache->offset + q_position * kv_row * 2, @@ -2026,7 +2121,7 @@ static constexpr uint32_t VK_QKV_PREP_MAX_HEAD_DIM = 512; n * 2, true)}; return vk_seq_dispatch_acc( - be, VK_PIPE_KV_APPEND_F16, bi, acc, push, sizeof(push), vk_groups(n), 1, 1); + be, VK_PIPE_KV_APPEND_F16, bi, acc, push, sizeof(push), vk_groups(push[0]), 1, 1); } /* Gated-DeltaNet mixer on the device: causal conv + silu, then the delta-rule @@ -2084,21 +2179,29 @@ vk_deltanet_geometry_ok(size_t seq, size_t n_kh, size_t n_vh, size_t dk, size_t !vk_tensor_gpu(a->conv_state, &bcs, &ocs) || !vk_tensor_gpu(a->delta_state, &bs, &os)) { return GEIST_E_UNSUPPORTED; } + /* In range per vk_deltanet_geometry_ok; checked all the same. */ + uint32_t seq32, n_kh32, n_vh32, dk32, dv32, keyd32, convd32, K32; + if (vk_ckd_u32(seq, &seq32) || vk_ckd_u32(n_kh, &n_kh32) || vk_ckd_u32(n_vh, &n_vh32) || + vk_ckd_u32(dk, &dk32) || vk_ckd_u32(dv, &dv32) || vk_ckd_u32(keyd, &keyd32) || + vk_ckd_u32(convd, &convd32) || vk_ckd_u32(K, &K32)) { + return vk_too_wide(be, "deltanet_mix"); + } { - const VkDescriptorBufferInfo bi[3] = {bqkv, bw, bcs}; - const uint32_t push[6] = {(uint32_t) seq, (uint32_t) convd, (uint32_t) K, oqkv, ow, ocs}; - const struct vk_access acc[3] = {vk_acc_tensor(a->qkv, true), - vk_acc_tensor(a->conv_w, false), - vk_acc_tensor(a->conv_state, true)}; - const enum geist_status s = vk_seq_dispatch_acc(be, - VK_PIPE_DN_CONV, - bi, - acc, - push, - sizeof(push), - (uint32_t) ((convd + 127u) / 128u), - 1, - 1); + const VkDescriptorBufferInfo bi[3] = {bqkv, bw, bcs}; + const uint32_t push[6] = {seq32, convd32, K32, oqkv, ow, ocs}; + const struct vk_access acc[3] = {vk_acc_tensor(a->qkv, true), + vk_acc_tensor(a->conv_w, false), + vk_acc_tensor(a->conv_state, true)}; + const enum geist_status s = + vk_seq_dispatch_acc(be, + VK_PIPE_DN_CONV, + bi, + acc, + push, + sizeof(push), + convd32 / 128u + (convd32 % 128u != 0u ? 1u : 0u), + 1, + 1); if (s != GEIST_OK) { return s; } @@ -2108,13 +2211,13 @@ vk_deltanet_geometry_ok(size_t seq, size_t n_kh, size_t n_vh, size_t dk, size_t uint32_t seq, n_kh, n_vh, dk, dv, keyd, convd; uint32_t qkv_off, z_off, beta_off, alpha_off, a_off, dtb_off, nw_off, s_off; float eps, qscale; - } push = {(uint32_t) seq, - (uint32_t) n_kh, - (uint32_t) n_vh, - (uint32_t) dk, - (uint32_t) dv, - (uint32_t) keyd, - (uint32_t) convd, + } push = {seq32, + n_kh32, + n_vh32, + dk32, + dv32, + keyd32, + convd32, oqkv, oz, ob, @@ -2133,8 +2236,7 @@ vk_deltanet_geometry_ok(size_t seq, size_t n_kh, size_t n_vh, size_t dk, size_t vk_acc_tensor(a->dt_bias, false), vk_acc_tensor(a->norm_w, false), vk_acc_tensor(a->delta_state, true)}; - return vk_seq_dispatch_acc( - be, VK_PIPE_DN_DELTA, bi, acc, &push, sizeof(push), (uint32_t) n_vh, 1, 1); + return vk_seq_dispatch_acc(be, VK_PIPE_DN_DELTA, bi, acc, &push, sizeof(push), n_vh32, 1, 1); } /* ====================================================================== */ @@ -2238,6 +2340,71 @@ static bool vk_fused_supported(struct geist_backend *be, const struct geist_fusi } } +/* ---- Fused-op entry points: a decline is a fallback (#474 item 4). */ +/* The arch takes the host path on GEIST_E_UNSUPPORTED; vk_fallback counts */ +/* it per site and, under GEIST_VK_STRICT=1, turns it into an error. */ + +static enum geist_status +vk_declined(struct geist_backend *be, enum geist_status s, enum vk_fb site) { + return s == GEIST_E_UNSUPPORTED ? vk_fallback(be->state, site) : s; +} + +[[nodiscard]] static enum geist_status vk_fb_linear_t(struct geist_backend *be, + const struct geist_tensor *t_x, + const struct geist_weight *w, + const struct geist_tensor *t_w, + size_t m, + struct geist_tensor *t_y) { + return vk_declined(be, vk_linear_t(be, t_x, w, t_w, m, t_y), VK_FB_LINEAR_T); +} + +[[nodiscard]] static enum geist_status vk_fb_linear_t_pair(struct geist_backend *be, + const struct geist_tensor *t_x, + const struct geist_weight *w0, + const struct geist_tensor *t_w0, + const struct geist_weight *w1, + const struct geist_tensor *t_w1, + size_t m, + struct geist_tensor *t_y0, + struct geist_tensor *t_y1) { + return vk_declined( + be, vk_linear_t_pair(be, t_x, w0, t_w0, w1, t_w1, m, t_y0, t_y1), VK_FB_LINEAR_T); +} + +[[nodiscard]] static enum geist_status +vk_fb_embedding_lookup_scaled(struct geist_backend *be, + const struct geist_tensor *embed_table, + geist_token_t token_id, + float scale, + struct geist_tensor *out) { + return vk_declined( + be, vk_embedding_lookup_scaled(be, embed_table, token_id, scale, out), VK_FB_EMBED); +} + +[[nodiscard]] static enum geist_status +vk_fb_argmax_f32(struct geist_backend *be, const struct geist_tensor *logits, int32_t *out_index) { + return vk_declined(be, vk_argmax_f32(be, logits, out_index), VK_FB_ARGMAX); +} + +[[nodiscard]] static enum geist_status vk_fb_kv_append_f16(struct geist_backend *be, + const struct geist_tensor *k_src, + const struct geist_tensor *v_src, + size_t q_position, + struct geist_tensor *k_cache, + struct geist_tensor *v_cache) { + return vk_declined( + be, vk_kv_append_f16(be, k_src, v_src, q_position, k_cache, v_cache), VK_FB_KV_APPEND); +} + +[[nodiscard]] static enum geist_status vk_fb_attn_qgate_split(struct geist_backend *be, + const struct geist_tensor *joint, + size_t heads, + size_t head_dim, + struct geist_tensor *q, + struct geist_tensor *gate) { + return vk_declined(be, vk_attn_qgate_split(be, joint, heads, head_dim, q, gate), VK_FB_QGATE); +} + static const struct geist_backend_primitives vk_prims = { .rmsnorm = vk_rmsnorm, .add = vk_add, @@ -2256,18 +2423,18 @@ static const struct geist_backend_fused vk_fused = { .gelu_tanh_mul = vk_gelu_tanh_mul, .gelu_tanh_mul_scaled = vk_gelu_tanh_mul_scaled, /* Batched-submit paths: one flush per token (argmax). */ - .linear_t = vk_linear_t, - .linear_t_pair = vk_linear_t_pair, + .linear_t = vk_fb_linear_t, + .linear_t_pair = vk_fb_linear_t_pair, .rmsnorm_add = vk_rmsnorm_add, - .embedding_lookup_scaled = vk_embedding_lookup_scaled, - .argmax_f32 = vk_argmax_f32, + .embedding_lookup_scaled = vk_fb_embedding_lookup_scaled, + .argmax_f32 = vk_fb_argmax_f32, .ffn_gate_up = vk_ffn_gate_up, .ffn_norm_gate_up = vk_ffn_norm_gate_up, .ple_block = vk_ple_block, .attn_qkv_prep = vk_attn_qkv_prep, - .kv_append_f16 = vk_kv_append_f16, + .kv_append_f16 = vk_fb_kv_append_f16, .deltanet_mix = vk_deltanet_mix, - .attn_qgate_split = vk_attn_qgate_split, + .attn_qgate_split = vk_fb_attn_qgate_split, .sigmoid_mul = vk_sigmoid_mul, .silu_mul = vk_silu_mul, .bitnet_act_quant = vk_bitnet_act_quant, @@ -2283,7 +2450,7 @@ const struct geist_backend_descriptor geist_backend_vulkan = { .batched_submit = true, .weights_need_backend_arena = true, .weights_device_copy = true, - .max_m = 512, + .max_m = VK_MAX_M, /* the DeltaNet mixer is sequential over tokens: its cost does not * grow with the chunk, so GEIST_M_MAX above 64 is not capped for * qwen35 hybrids. The default chunk stays 64: 128 makes the diff --git a/src/backends/vulkan/resources.c b/src/backends/vulkan/resources.c index dbd5486c..b1ebc957 100644 --- a/src/backends/vulkan/resources.c +++ b/src/backends/vulkan/resources.c @@ -210,7 +210,9 @@ static void vk_vram_exhausted(struct geist_backend *be, struct geist_buffer **nb = geist_backend_alloc(be, cap * sizeof(*nb), alignof(struct geist_buffer *)); if (nb != nullptr) { - memcpy(nb, st->hostbufs, st->n_hostbufs * sizeof(*nb)); + if (st->n_hostbufs > 0) { /* hostbufs is nullptr before the first grow */ + memcpy(nb, st->hostbufs, st->n_hostbufs * sizeof(*nb)); + } geist_backend_free(be, st->hostbufs); st->hostbufs = nb; st->cap_hostbufs = cap; @@ -454,10 +456,17 @@ void *vk_buffer_map(struct geist_buffer *buf) { return nullptr; } vk_seq_flush(buf->owner); /* host is about to read/write — drain the batch */ - if (buf->host_alias != nullptr) { - return buf->host_alias; + void *p = buf->host_alias != nullptr ? buf->host_alias : buf->mapped; + if (p == nullptr) { + return nullptr; /* device-local — caller falls back (download reports a failure) */ + } + /* A dropped batch leaves the mapping holding whatever was there before: + * hand out no pointer rather than stale results. buffer_map has no status, + * so the failure is the backend error and the caller's nullptr check. */ + if (vk_seq_take_failure(buf->owner) != GEIST_OK) { + return nullptr; } - return buf->mapped; /* nullptr for device-local — caller must fall back */ + return p; } void vk_buffer_unmap(struct geist_buffer *buf) { @@ -603,17 +612,18 @@ struct vk_access vk_acc_tensor(const struct geist_tensor *t, bool write) { } /* GPU view of a tensor: VkBuffer + f32 element offset. False when the - * tensor's buffer has no VkBuffer behind it (e.g. GGUF-mmap aliases). */ + * tensor's buffer has no VkBuffer behind it (e.g. GGUF-mmap aliases), or + * when the offset does not fit the shaders' uint32 element index. */ bool vk_tensor_gpu(const struct geist_tensor *t, VkDescriptorBufferInfo *out, uint32_t *elem_off) { if (t == nullptr || t->buffer == nullptr || t->buffer->buf == VK_NULL_HANDLE) { return false; } - const size_t byte_off = t->buffer->base_off + t->offset; - if (byte_off % 4 != 0) { + size_t byte_off; + if (ckd_add(&byte_off, t->buffer->base_off, t->offset) || byte_off % 4 != 0 || + vk_ckd_u32(byte_off / 4, elem_off)) { return false; } - *out = (VkDescriptorBufferInfo) {.buffer = t->buffer->buf, .range = VK_WHOLE_SIZE}; - *elem_off = (uint32_t) (byte_off / 4); + *out = (VkDescriptorBufferInfo) {.buffer = t->buffer->buf, .range = VK_WHOLE_SIZE}; return true; } @@ -624,12 +634,12 @@ bool vk_tensor_gpu_f16(const struct geist_tensor *t, if (t == nullptr || t->buffer == nullptr || t->buffer->buf == VK_NULL_HANDLE) { return false; } - const size_t byte_off = t->buffer->base_off + t->offset; - if (byte_off % 2 != 0) { + size_t byte_off; + if (ckd_add(&byte_off, t->buffer->base_off, t->offset) || byte_off % 2 != 0 || + vk_ckd_u32(byte_off / 2, elem_off)) { return false; } - *out = (VkDescriptorBufferInfo) {.buffer = t->buffer->buf, .range = VK_WHOLE_SIZE}; - *elem_off = (uint32_t) (byte_off / 2); + *out = (VkDescriptorBufferInfo) {.buffer = t->buffer->buf, .range = VK_WHOLE_SIZE}; return true; } @@ -672,6 +682,9 @@ void *vk_tensor_host(const struct geist_tensor *t, size_t *out_n) { return nullptr; /* device-local — CPU ops can't touch it */ } t->buffer->owner->stat_cpu_falls++; + if (vk_fallback(t->buffer->owner, VK_FB_HOST_VIEW) != GEIST_E_UNSUPPORTED) { + return nullptr; /* GEIST_VK_STRICT: the error names the site */ + } vk_seq_flush(t->buffer->owner); /* host access — drain pending GPU work */ if (vk_seq_take_failure(t->buffer->owner) != GEIST_OK) { return nullptr; @@ -738,13 +751,20 @@ bool vk_t_geom(const struct geist_tensor *t, size_t *rows, size_t *cols, size_t return GEIST_OK; } /* Host fallback. */ - vk_seq_flush(st); uint8_t *d = dst->host_alias != nullptr ? dst->host_alias : dst->mapped; const uint8_t *sp = src->host_alias != nullptr ? src->host_alias : src->mapped; if (d == nullptr || sp == nullptr || dst_offset + n_bytes > dst->bytes || src_offset + n_bytes > src->bytes) { return GEIST_E_UNSUPPORTED; } + const enum geist_status fs = vk_fallback(st, VK_FB_HOST_COPY); + if (fs != GEIST_E_UNSUPPORTED) { + return fs; + } + vk_seq_flush(st); + if (vk_seq_take_failure(st) != GEIST_OK) { + return GEIST_E_BACKEND; /* the source holds no results of the dropped batch */ + } memcpy(d + dst_offset, sp + src_offset, n_bytes); return GEIST_OK; } @@ -762,19 +782,23 @@ bool vk_t_geom(const struct geist_tensor *t, size_t *rows, size_t *cols, size_t if (!vk_tensor_gpu(t_x, &src_bi, &src_elem)) { return false; } - const size_t bytes = m * n_in * sizeof(float); + size_t bytes; + if (ckd_mul(&bytes, m, n_in) || ckd_mul(&bytes, bytes, sizeof(float))) { + return false; + } if (st->xring == nullptr && vk_buffer_create(be, VK_XRING_CAP, GEIST_BUFFER_SCRATCH, GEIST_MEMORY_DEVICE, &st->xring) != GEIST_OK) { return false; } if (bytes > st->xring->bytes) { - return false; + return false; /* resolve_weight refuses an n_in whose max_m batch would not fit */ } if (st->xring_used + bytes > st->xring->bytes) { vk_seq_flush(st); /* drains the batch and resets the ring */ } - if (vk_seq_open_cmd(st) != GEIST_OK) { + uint32_t elem_off; /* < VK_XRING_CAP / 4; checked all the same */ + if (vk_ckd_u32(st->xring_used / sizeof(float), &elem_off) || vk_seq_open_cmd(st) != GEIST_OK) { return false; } { @@ -807,7 +831,7 @@ bool vk_t_geom(const struct geist_tensor *t, size_t *rows, size_t *cols, size_t st->seq_dispatches++; st->seq_in_cmd++; vk_prof_stamp(st, VK_PIPE_COUNT); - *out_elem_off = (uint32_t) (st->xring_used / sizeof(float)); + *out_elem_off = elem_off; st->xring_used = (st->xring_used + bytes + 63) & ~(size_t) 63; return true; } diff --git a/src/backends/vulkan/sequence.c b/src/backends/vulkan/sequence.c index 321599c1..cf51dfe8 100644 --- a/src/backends/vulkan/sequence.c +++ b/src/backends/vulkan/sequence.c @@ -73,6 +73,40 @@ void vk_seq_flush(struct vk_state *st) { return GEIST_E_BACKEND; } +const char *vk_fallback_name(enum vk_fb site) { + static const char *const names[VK_FB_COUNT] = { + [VK_FB_HOST_VIEW] = "host_view", + [VK_FB_HOST_LINEAR] = "host_linear", + [VK_FB_HOST_COPY] = "host_copy", + [VK_FB_LINEAR_T] = "linear_t", + [VK_FB_ARGMAX] = "argmax", + [VK_FB_EMBED] = "embedding_lookup", + [VK_FB_KV_APPEND] = "kv_append_f16", + [VK_FB_QGATE] = "attn_qgate_split", + }; + return (unsigned) site < VK_FB_COUNT ? names[site] : "?"; +} + +[[nodiscard]] enum geist_status vk_fallback(struct vk_state *st, enum vk_fb site) { + st->fallbacks[site]++; + if (!st->strict) { + return GEIST_E_UNSUPPORTED; + } + geist_backend_set_error(st->backend, + GEIST_E_BACKEND, + "vulkan: GEIST_VK_STRICT=1 and %s would leave the GPU", + vk_fallback_name(site)); + return GEIST_E_BACKEND; +} + +[[nodiscard]] enum geist_status vk_too_wide(struct geist_backend *be, const char *op) { + geist_backend_set_error(be, + GEIST_E_INVALID_ARG, + "vulkan %s: a size or offset exceeds the shader's 32-bit range", + op); + return GEIST_E_INVALID_ARG; +} + [[nodiscard]] enum geist_status vk_seq_open_cmd(struct vk_state *st) { if (st->seq_open) { return GEIST_OK; diff --git a/src/backends/vulkan/vk_internal.h b/src/backends/vulkan/vk_internal.h index 54035b4b..7bbca21b 100644 --- a/src/backends/vulkan/vk_internal.h +++ b/src/backends/vulkan/vk_internal.h @@ -17,6 +17,7 @@ #include #include +#include "checked.h" /* ckd_* size arithmetic (AGENT.md §3) */ #include "gemma4_kernels.h" /* shared reference rope/attention kernels */ #include "heap.h" #include "quant.h" /* CPU dequant helpers for the non-GPU dtype fallback */ @@ -207,6 +208,7 @@ struct vk_dirty { enum { VK_XRING_CAP = 192u << 20, /* a full prefill chunk stages ~124 MB */ + VK_MAX_M = 512, /* caps.max_m; resolve_weight checks it x n_in fits the ring */ VK_DIRTY_CAP = 96, VK_DSET_CACHE = 4096, VK_SEQ_CMDBUFS = 64, /* rolling submission ring */ @@ -228,6 +230,19 @@ enum { VK_PUSH_RANGE = 128, /* one push range covers every shader block */ }; +/* Where work can leave the GPU (vk_fallback). */ +enum vk_fb { + VK_FB_HOST_VIEW, /* a host loop over a mapped tensor (vk_tensor_host) */ + VK_FB_HOST_LINEAR, /* a linear on the host row-dequant path (vk_w_cpu_mN) */ + VK_FB_HOST_COPY, /* vk_buffer_copy through mapped memory */ + VK_FB_LINEAR_T, /* linear_t / linear_t_pair declined: arch host linear */ + VK_FB_ARGMAX, /* argmax declined: arch scans on the host */ + VK_FB_EMBED, /* embedding lookup declined: arch gathers on the host */ + VK_FB_KV_APPEND, /* kv_append_f16 declined */ + VK_FB_QGATE, /* attn_qgate_split declined */ + VK_FB_COUNT +}; + struct vk_state { struct geist_backend *backend; void *lib; /* dlopen handle, may be nullptr after create */ @@ -287,6 +302,18 @@ struct vk_state { uint64_t stat_dispatches; uint64_t stat_cpu_falls; + /* Work that left the GPU (#474 item 4), counted per site by + * vk_fallback: a coverage gap otherwise shows up only as a slowdown. + * GEIST_VK_STRICT=1 (strict) turns every such fallback into an error, + * and refuses host-path weights at resolve. Weights resolved onto the + * host row-dequant path are summed at resolve and reported once, when + * the first of them runs. */ + uint64_t fallbacks[VK_FB_COUNT]; + bool strict; + bool host_weights_noted; + size_t host_weights; + size_t host_weight_bytes; + /* GEIST_VK_PROFILE=1: GPU timestamps per dispatch, attributed by * pipeline (copies land in the extra slot). Execution is serialized by * the per-dispatch barriers, so consecutive deltas are exact. */ @@ -529,6 +556,28 @@ bool vk_t_geom(const struct geist_tensor *t, size_t *rows, size_t *cols, size_t void vk_seq_flush(struct vk_state *st); [[nodiscard]] enum geist_status vk_seq_take_failure(struct vk_state *st); +/* Record that the work at `site` leaves the GPU. Returns the status the + * caller hands on: GEIST_E_UNSUPPORTED (take the documented fallback), or + * under GEIST_VK_STRICT=1 GEIST_E_BACKEND with an error naming the site. */ +[[nodiscard]] enum geist_status vk_fallback(struct vk_state *st, enum vk_fb site); +const char *vk_fallback_name(enum vk_fb site); + +/* Checked size -> uint32_t narrowing for push constants, dispatch sizes and + * element offsets: the shaders index in uint32, so a value that does not fit + * fails the op instead of wrapping (AGENT.md §3/§5, #474). checked.h's + * convention: true when `v` does not fit; *out is written only when it does. */ +[[nodiscard]] static inline bool vk_ckd_u32(size_t v, uint32_t *out) { + if (v > UINT32_MAX) { + return true; + } + *out = (uint32_t) v; + return false; +} + +/* What an op returns when vk_ckd_u32 refused one of its values: + * GEIST_E_INVALID_ARG, with an error naming the op. */ +[[nodiscard]] enum geist_status vk_too_wide(struct geist_backend *be, const char *op); + [[nodiscard]] enum geist_status vk_seq_open_cmd(struct vk_state *st); void vk_prof_stamp(struct vk_state *st, uint32_t slot); diff --git a/tests/test_backend_vulkan_fallback_unit.c b/tests/test_backend_vulkan_fallback_unit.c new file mode 100644 index 00000000..ecd67e98 --- /dev/null +++ b/tests/test_backend_vulkan_fallback_unit.c @@ -0,0 +1,442 @@ +/* + * test_backend_vulkan_fallback_unit — work that leaves the GPU is counted + * per site, and GEIST_VK_STRICT=1 turns it into an error (#474 items 4, 9). + * + * The Vulkan backend declines a fused op (GEIST_E_UNSUPPORTED) when its + * shader does not apply, and the arch then runs the op on the host; weights + * without a GPU kernel run on a host row-dequant path. Both used to show up + * only as a slowdown. Checked here: + * + * - a declined fused op (argmax, linear_t, linear_t_pair, embedding, + * kv_append_f16, attn_qgate_split) counts one fallback at its site and + * still returns GEIST_E_UNSUPPORTED; under strict mode it returns + * GEIST_E_BACKEND with an error naming the site; + * - a failed argmax dispatch is GEIST_E_BACKEND, not a decline (it used to + * send the arch to a host scan of logits the GPU never wrote); + * - a host view of a mapped tensor and a host buffer copy are counted, and + * refused under strict mode; + * - a weight resolved onto the host path is summed at resolve, its linear + * counts at run time, and strict mode refuses it at resolve; + * - decoding the in-memory llama and Qwen3.5-hybrid fixtures leaves no + * fused op and no weight on the host, and runs with strict mode on. + * + * SKIPs when the Vulkan backend is not built or has no device. + */ +#include "test_helpers.h" +#include "model_fixtures.h" + +#include +#include +#include + +#include +#include +#include + +#if defined(GEIST_BACKEND_VULKAN) && GEIST_BACKEND_VULKAN +#include "src/backends/vulkan/vk_internal.h" +#include "src/quant/quant.h" + +static VkResult VKAPI_CALL fail_begin(VkCommandBuffer cmd, const VkCommandBufferBeginInfo *info) { + (void) cmd; + (void) info; + return VK_ERROR_DEVICE_LOST; +} + +enum { FEAT = 64 }; + +/* A [rows, FEAT] F32 view of a fresh buffer of `role` (nullptr buffer on + * failure). */ +static struct geist_tensor +tensor_2d(struct geist_backend *be, size_t rows, enum geist_buffer_role role) { + struct geist_tensor t = {.dtype = GEIST_DTYPE_F32, + .layout = GEIST_LAYOUT_DENSE, + .ndim = 2, + .shape = {(int64_t) rows, FEAT}, + .stride = {FEAT, 1}}; + if (be->desc->vtbl->buffer_create(be, rows * FEAT * sizeof(float), role, 0, &t.buffer) != + GEIST_OK) { + t.buffer = nullptr; + } + return t; +} + +static uint64_t total(const struct vk_state *st) { + uint64_t n = 0; + for (size_t i = 0; i < VK_FB_COUNT; i++) { + n += st->fallbacks[i]; + } + return n; +} + +/* `s` came from a declined op at `site`: one fallback counted there, and + * the status is the one strict mode asks for. */ +static int expect_declined(struct vk_state *st, + enum geist_status s, + enum vk_fb site, + uint64_t before, + const char *op) { + char msg[160]; + snprintf(msg, + sizeof msg, + "%s%s: declined -> %s and one %s fallback (got %d, count %llu)", + op, + st->strict ? " (strict)" : "", + st->strict ? "GEIST_E_BACKEND" : "GEIST_E_UNSUPPORTED", + vk_fallback_name(site), + (int) s, + (unsigned long long) (st->fallbacks[site] - before)); + int fails = geist_expect(s == (st->strict ? GEIST_E_BACKEND : GEIST_E_UNSUPPORTED) && + st->fallbacks[site] == before + 1, + msg); + if (st->strict) { + const char *err = geist_backend_errmsg(st->backend); + snprintf(msg, sizeof msg, "%s (strict): error names the site (%s)", op, err); + fails += geist_expect(err != nullptr && strstr(err, vk_fallback_name(site)) != nullptr && + strstr(err, "GEIST_VK_STRICT") != nullptr, + msg); + } + return fails; +} + +/* Every fused op the backend can decline, called so that it does. */ +static int check_declines(struct geist_backend *be) { + struct vk_state *st = be->state; + const struct geist_backend_fused *f = geist_backend_fused_tbl(be); + int fails = 0; + + /* argmax over an empty tensor: nothing to bind. */ + struct geist_tensor empty = {.dtype = GEIST_DTYPE_F32, + .layout = GEIST_LAYOUT_DENSE, + .ndim = 1, + .shape = {0}, + .stride = {1}}; + int32_t idx = -1; + uint64_t b = st->fallbacks[VK_FB_ARGMAX]; + fails += expect_declined(st, f->argmax_f32(be, &empty, &idx), VK_FB_ARGMAX, b, "argmax_f32"); + + /* linear_t / linear_t_pair with a weight never resolved to VRAM. */ + static float wraw[FEAT * FEAT]; + const struct geist_weight w = {.raw = wraw, + .raw_nbytes = sizeof wraw, + .n_in = FEAT, + .n_out = FEAT, + .dtype = GEIST_DTYPE_F32}; + struct geist_tensor x = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + struct geist_tensor y0 = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + struct geist_tensor y1 = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + if (x.buffer != nullptr && y0.buffer != nullptr && y1.buffer != nullptr) { + b = st->fallbacks[VK_FB_LINEAR_T]; + fails += expect_declined( + st, f->linear_t(be, &x, &w, nullptr, 1, &y0), VK_FB_LINEAR_T, b, "linear_t"); + b = st->fallbacks[VK_FB_LINEAR_T]; + fails += expect_declined(st, + f->linear_t_pair(be, &x, &w, nullptr, &w, nullptr, 1, &y0, &y1), + VK_FB_LINEAR_T, + b, + "linear_t_pair"); + + /* Embedding: a dtype the gather shader has no code for. */ + struct geist_tensor table = x; + table.dtype = GEIST_DTYPE_Q3_K; + b = st->fallbacks[VK_FB_EMBED]; + fails += expect_declined(st, + f->embedding_lookup_scaled(be, &table, 0, 1.0f, &y0), + VK_FB_EMBED, + b, + "embedding_lookup_scaled"); + + /* qgate split: the joint row is not 2 * heads * head_dim wide. */ + b = st->fallbacks[VK_FB_QGATE]; + fails += expect_declined( + st, f->attn_qgate_split(be, &x, 4, 16, &y0, &y1), VK_FB_QGATE, b, "qgate_split"); + + /* kv_append: caches with no VkBuffer behind them (host aliases). */ + struct geist_tensor kv = x; + kv.ndim = 3; + kv.shape[0] = 1; + kv.shape[1] = 4; + kv.shape[2] = FEAT / 4; + kv.stride[0] = FEAT; + kv.stride[1] = FEAT / 4; + kv.stride[2] = 1; + static uint16_t host_cache[FEAT]; + struct geist_buffer alias = { + .owner = st, .host_alias = host_cache, .bytes = sizeof host_cache}; + struct geist_tensor cache = y0; + cache.buffer = &alias; + b = st->fallbacks[VK_FB_KV_APPEND]; + fails += expect_declined(st, + f->kv_append_f16(be, &kv, &kv, 0, &cache, &cache), + VK_FB_KV_APPEND, + b, + "kv_append_f16"); + } else { + fails += geist_expect(false, "scratch buffers"); + } + struct geist_buffer *bufs[] = {x.buffer, y0.buffer, y1.buffer}; + for (size_t i = 0; i < 3; i++) { + if (bufs[i] != nullptr) { + be->desc->vtbl->buffer_destroy(be, bufs[i]); + } + } + return fails; +} + +/* Host views and host buffer copies (vk_tensor_host, vk_buffer_copy). */ +static int check_host_access(struct geist_backend *be) { + struct vk_state *st = be->state; + int fails = 0; + struct geist_tensor a = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + struct geist_tensor c = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + if (a.buffer == nullptr || c.buffer == nullptr || !a.buffer->host_visible) { + fails += geist_expect(a.buffer != nullptr && c.buffer != nullptr, "host-access buffers"); + } else { + uint64_t b = st->fallbacks[VK_FB_HOST_VIEW]; + size_t n = 0; + void *p = vk_tensor_host(&a, &n); + fails += geist_expect(p != nullptr && n == FEAT && st->fallbacks[VK_FB_HOST_VIEW] == b + 1, + "a host view is counted"); + /* Two bookkeeping aliases (no VkBuffer): the copy goes through host + * memory. */ + uint8_t src[32], dst[32] = {0}; + struct geist_buffer hs = {.owner = st, .host_alias = src, .bytes = sizeof src}; + struct geist_buffer hd = {.owner = st, .host_alias = dst, .bytes = sizeof dst}; + for (size_t i = 0; i < sizeof src; i++) { + src[i] = (uint8_t) i; + } + b = st->fallbacks[VK_FB_HOST_COPY]; + fails += geist_expect(vk_buffer_copy(&hd, 0, &hs, 0, 16) == GEIST_OK && + memcmp(dst, src, 16) == 0 && + st->fallbacks[VK_FB_HOST_COPY] == b + 1, + "a host buffer copy is counted"); + st->strict = true; + memset(dst, 0, sizeof dst); + fails += geist_expect(vk_buffer_copy(&hd, 0, &hs, 0, 16) == GEIST_E_BACKEND && + dst[1] == 0 && st->fallbacks[VK_FB_HOST_COPY] == b + 2, + "strict: a host buffer copy is refused"); + st->strict = false; + st->strict = true; + b = st->fallbacks[VK_FB_HOST_VIEW]; + fails += geist_expect(vk_tensor_host(&a, &n) == nullptr && + st->fallbacks[VK_FB_HOST_VIEW] == b + 1, + "strict: a host view is refused"); + st->strict = false; + } + if (a.buffer != nullptr) { + be->desc->vtbl->buffer_destroy(be, a.buffer); + } + if (c.buffer != nullptr) { + be->desc->vtbl->buffer_destroy(be, c.buffer); + } + return fails; +} + +/* A failed argmax dispatch is an error, not a decline. */ +static int check_argmax_dispatch_failure(struct geist_backend *be) { + struct vk_state *st = be->state; + const struct geist_backend_fused *f = geist_backend_fused_tbl(be); + struct geist_tensor logits = tensor_2d(be, 1, GEIST_BUFFER_SCRATCH); + int fails = 0; + if (logits.buffer == nullptr) { + return geist_expect(false, "logits buffer"); + } + float *p = be->desc->vtbl->buffer_map(logits.buffer); + for (size_t i = 0; i < FEAT; i++) { + p[i] = (float) (i == 37); + } + be->desc->vtbl->buffer_unmap(logits.buffer); + int32_t idx = -1; + fails += geist_expect(f->argmax_f32(be, &logits, &idx) == GEIST_OK && idx == 37, + "argmax on a working device"); + + const uint64_t before = total(st); + const PFN_vkBeginCommandBuffer real = st->fn.BeginCommandBuffer; + vk_seq_flush(st); + st->fn.BeginCommandBuffer = fail_begin; + const enum geist_status s = f->argmax_f32(be, &logits, &idx); + st->fn.BeginCommandBuffer = real; + (void) vk_seq_take_failure(st); + char msg[128]; + snprintf(msg, sizeof msg, "failed argmax dispatch -> GEIST_E_BACKEND (got %d)", (int) s); + fails += geist_expect(s == GEIST_E_BACKEND, msg); + fails += geist_expect(total(st) == before, "a failed dispatch is not counted as a fallback"); + fails += geist_expect(f->argmax_f32(be, &logits, &idx) == GEIST_OK && idx == 37, + "argmax after recovery"); + be->desc->vtbl->buffer_destroy(be, logits.buffer); + return fails; +} + +/* A Q3_K weight (no GPU kernel) resolves onto the host path. */ +static int check_host_weight(struct geist_backend *be) { + struct vk_state *st = be->state; + const struct geist_backend_vtbl *vt = be->desc->vtbl; + int fails = 0; + enum { N_OUT = 4 }; + static uint8_t raw[N_OUT * Q3_K_BLOCK_BYTES]; /* zero blocks: d = 0 */ + struct geist_weight w = {.raw = raw, + .raw_nbytes = sizeof raw, + .n_in = (int32_t) Q3_K_BLOCK_ELEMS, + .n_out = N_OUT, + .dtype = GEIST_DTYPE_Q3_K}; + + st->strict = true; + const enum geist_status ss = vt->resolve_weight(be, &w); + st->strict = false; + fails += geist_expect(ss == GEIST_E_BACKEND && w.linear_m1 == nullptr && st->host_weights == 0, + "strict: a host-path weight is refused at resolve"); + const char *err = geist_backend_errmsg(be); + fails += geist_expect(err != nullptr && strstr(err, "GEIST_VK_STRICT") != nullptr, + "strict: the resolve error names strict mode"); + + fails += geist_expect(vt->resolve_weight(be, &w) == GEIST_OK && w.linear_m1 != nullptr && + w.linear_mN != nullptr, + "a Q3_K weight resolves onto the host path"); + fails += geist_expect(st->host_weights == 1 && st->host_weight_bytes == sizeof raw && + !st->host_weights_noted, + "the host weight is summed at resolve"); + float x[Q3_K_BLOCK_ELEMS], y[2 * N_OUT]; + for (size_t i = 0; i < Q3_K_BLOCK_ELEMS; i++) { + x[i] = 1.0f; + } + const uint64_t b = st->fallbacks[VK_FB_HOST_LINEAR]; + w.linear_m1(x, &w, be, y); + fails += geist_expect(st->fallbacks[VK_FB_HOST_LINEAR] == b + 1 && st->host_weights_noted, + "the host linear is counted and the summary noted"); + fails += geist_expect(y[0] == 0.0f && y[N_OUT - 1] == 0.0f, "zero blocks give zero output"); + float x2[2 * Q3_K_BLOCK_ELEMS]; + memcpy(x2, x, sizeof x); + memcpy(x2 + Q3_K_BLOCK_ELEMS, x, sizeof x); + w.linear_mN(2, x2, &w, be, y); + fails += geist_expect(st->fallbacks[VK_FB_HOST_LINEAR] == b + 2, "linear_mN is counted too"); + return fails; +} + +/* Decode a fixture model; returns failures. `strict` sets the env var the + * backend reads at create. */ +static int check_model(const char *name, const struct tf_buf *g, size_t vocab, bool strict) { + if (strict) { + setenv("GEIST_VK_STRICT", "1", 1); + } else { + unsetenv("GEIST_VK_STRICT"); + } + struct geist_backend *be = nullptr; + if (geist_backend_create("vulkan", nullptr, nullptr, &be) != GEIST_OK) { + unsetenv("GEIST_VK_STRICT"); + return geist_expect(false, "vulkan backend for a model run"); + } + struct vk_state *st = be->state; + int fails = 0; + char msg[192]; + snprintf(msg, sizeof msg, "%s%s: strict flag read at create", name, strict ? " strict" : ""); + fails += geist_expect(st->strict == strict, msg); + + struct geist_model *m = nullptr; + if (geist_model_load_from_memory(g->b, g->n, be, &m) != GEIST_OK) { + snprintf(msg, sizeof msg, "%s: model load: %s", name, geist_last_create_error()); + fails += geist_expect(false, msg); + } else { + struct geist_session_opts o = {.kv_mode = GEIST_KV_F16, .top_p = 1.0f, .max_seq_len = 64}; + struct geist_session *s = nullptr; + static const geist_token_t prompt[] = {1, 5, 9, 13, 17, 21, 25, 29}; + geist_token_t tok = -1; + bool ok = geist_session_create(m, be, &o, &s) == GEIST_OK && + geist_session_prefill_tokens(s, sizeof prompt / sizeof prompt[0], prompt) == + GEIST_OK; + for (int i = 0; ok && i < 6; i++) { + ok = geist_session_decode_step(s, &tok) == GEIST_OK && tok >= 0 && (size_t) tok < vocab; + } + snprintf(msg, + sizeof msg, + "%s%s: prefill and decode (%s)", + name, + strict ? " strict" : "", + geist_backend_errmsg(be)); + fails += geist_expect(ok, msg); + for (size_t i = 0; i < VK_FB_COUNT; i++) { + snprintf(msg, + sizeof msg, + "%s%s: no %s fallback (%llu)", + name, + strict ? " strict" : "", + vk_fallback_name((enum vk_fb) i), + (unsigned long long) st->fallbacks[i]); + fails += geist_expect(st->fallbacks[i] == 0, msg); + } + snprintf(msg, sizeof msg, "%s: no weight on the host path", name); + fails += geist_expect(st->host_weights == 0 && st->fallbacks[VK_FB_HOST_LINEAR] == 0, msg); + if (s != nullptr) { + geist_session_destroy(s); + } + geist_model_destroy(m); + } + geist_backend_destroy(be); + unsetenv("GEIST_VK_STRICT"); + return fails; +} + +int main(void) { + unsetenv("GEIST_VK_STRICT"); + struct geist_backend *be = nullptr; + enum geist_status s = geist_backend_create("vulkan", nullptr, nullptr, &be); + if (s != GEIST_OK) { + fprintf(stderr, "SKIP: vulkan backend unavailable (%d)\n", (int) s); + return GEIST_TEST_SKIP; + } + struct vk_state *st = be->state; + int fails = geist_expect(!st->strict && total(st) == 0, "a fresh backend"); + fails += check_declines(be); + st->strict = true; + fails += check_declines(be); + st->strict = false; + fails += check_host_access(be); + fails += check_argmax_dispatch_failure(be); + fails += check_host_weight(be); + geist_backend_destroy(be); + + struct tf_vocab v = tf_make_vocab("\xc4\xa0", false); + struct tf_buf llama = mf_llama_gguf(&(struct mf_llama) {.layers = 2, + .d_model = 128, + .heads = 4, + .kv_heads = 2, + .ffn = 256, + .vocab = 512, + .context = 256, + .seed = 14}); + struct tf_buf qwen = mf_qwen35_gguf(&(struct mf_qwen35) {.layers = 4, + .interval = 4, + .d_model = 64, + .heads = 4, + .kv_heads = 2, + .head_dim = 16, + .rope_dims = 8, + .ffn = 128, + .dn_k_heads = 2, + .dn_v_heads = 4, + .dn_head_k = 16, + .dn_head_v = 16, + .dn_conv = 4, + .seed = 1, + .tok = &v}); + for (int strict = 0; strict < 2; strict++) { + fails += check_model("llama", &llama, 512, strict != 0); + fails += check_model("qwen35", &qwen, v.n_tok, strict != 0); + } + free(llama.b); + free(qwen.b); + tf_free_vocab(&v); + + if (fails > 0) { + fprintf(stderr, "%d check(s) failed\n", fails); + return GEIST_TEST_FAIL; + } + printf("vulkan fallback accounting: pass\n"); + return GEIST_TEST_PASS; +} + +#else +int main(void) { + fprintf(stderr, "SKIP: vulkan backend not built\n"); + return GEIST_TEST_SKIP; +} +#endif diff --git a/tests/test_backend_vulkan_narrowing_unit.c b/tests/test_backend_vulkan_narrowing_unit.c new file mode 100644 index 00000000..403c689e --- /dev/null +++ b/tests/test_backend_vulkan_narrowing_unit.c @@ -0,0 +1,233 @@ +/* + * test_backend_vulkan_narrowing_unit — a failed submit reaches buffer_map, + * and sizes the shaders cannot index are refused, not truncated (#474, + * items 1 and 6). + * + * 1. buffer_map after a failed submit: the batch was dropped, so the mapping + * still holds the bytes from before it. buffer_map has no status, so it + * must return nullptr with GEIST_E_BACKEND as the backend error instead + * of those stale bytes; a download reports the same failure. + * 2. vk_ckd_u32 refuses values past UINT32_MAX and leaves its output alone. + * 3. An op handed a position past UINT32_MAX (attention's q_offset and + * sliding_window: plain scalars, no large allocation needed) returns + * GEIST_E_INVALID_ARG instead of dispatching with the value wrapped. + * 4. resolve_weight refuses a weight whose max_m batch would not fit the + * x ring (it used to become a per-call UNSUPPORTED and a silent host + * linear), and creates the ring at load for one that fits. + * + * SKIPs (exit 77) when the Vulkan backend is not built or has no device. + */ +#include "test_helpers.h" + +#include +#include +#include + +#include +#include + +#if defined(GEIST_BACKEND_VULKAN) && GEIST_BACKEND_VULKAN +#include "src/backends/vulkan/vk_internal.h" + +static VkResult VKAPI_CALL fail_submit(VkQueue q, + uint32_t n, + const VkSubmitInfo *submits, + VkFence fence) { + (void) q; + (void) n; + (void) submits; + (void) fence; + return VK_ERROR_DEVICE_LOST; +} + +enum { N = 256 }; + +/* A host-visible F32 buffer of N elements viewed with the given shape. */ +static struct geist_tensor +host_tensor(struct geist_backend *be, size_t ndim, const int64_t shape[static ndim]) { + struct geist_tensor t = {.dtype = GEIST_DTYPE_F32, .layout = GEIST_LAYOUT_DENSE}; + t.ndim = (int) ndim; + int64_t stride = 1; + for (size_t d = ndim; d-- > 0;) { + t.shape[d] = shape[d]; + t.stride[d] = stride; + stride *= shape[d]; + } + if (be->desc->vtbl->buffer_create(be, N * sizeof(float), GEIST_BUFFER_SCRATCH, 0, &t.buffer) != + GEIST_OK) { + t.buffer = nullptr; + } + return t; +} + +static bool fill(struct geist_backend *be, struct geist_tensor *t, float v) { + float *p = be->desc->vtbl->buffer_map(t->buffer); + if (p == nullptr) { + return false; + } + for (size_t i = 0; i < N; i++) { + p[i] = v; + } + be->desc->vtbl->buffer_unmap(t->buffer); + return true; +} + +/* Part 1: a dropped batch is reported by buffer_map and by download. */ +static int check_failed_submit(struct geist_backend *be) { + const struct geist_backend_primitives *p = be->desc->prims; + const struct geist_backend_vtbl *vt = be->desc->vtbl; + struct vk_state *st = be->state; + const int64_t sh[1] = {N}; + struct geist_tensor x = host_tensor(be, 1, sh), z = host_tensor(be, 1, sh), + y = host_tensor(be, 1, sh); + int fails = geist_expect(x.buffer != nullptr && z.buffer != nullptr && y.buffer != nullptr, + "buffers"); + if (fails == 0) { + fails += geist_expect(fill(be, &x, 0.5f) && fill(be, &z, 2.0f) && fill(be, &y, 7.0f), + "fill on a working device"); + /* Sanity: a working device delivers the add through buffer_map. */ + fails += geist_expect(p->add(be, &x, &z, &y) == GEIST_OK, "add on a working device"); + const float *yp = vt->buffer_map(y.buffer); + fails += geist_expect(yp != nullptr && yp[0] == 2.5f && yp[N - 1] == 2.5f, + "buffer_map returns the add's result"); + vt->buffer_unmap(y.buffer); + fails += geist_expect(fill(be, &y, 7.0f), "reset y"); + + const PFN_vkQueueSubmit real = st->fn.QueueSubmit; + st->fn.QueueSubmit = fail_submit; + fails += geist_expect(p->add(be, &x, &z, &y) == GEIST_OK, "add records into the batch"); + be->err_code = GEIST_OK; + const void *stale = vt->buffer_map(y.buffer); /* flushes: the submit fails */ + fails += geist_expect(stale == nullptr, + "buffer_map after a failed submit returns nullptr, not stale bytes"); + fails += geist_expect(be->err_code == GEIST_E_BACKEND, + "buffer_map after a failed submit sets GEIST_E_BACKEND"); + + /* Same through download, which has a status. */ + fails += geist_expect(p->add(be, &x, &z, &y) == GEIST_OK, "second add records"); + float host[N]; + fails += geist_expect(vt->buffer_download(sizeof host, (uint8_t *) host, y.buffer) == + GEIST_E_BACKEND, + "download after a failed submit returns GEIST_E_BACKEND"); + st->fn.QueueSubmit = real; + + /* Reported once: the device works again and the mapping is live. */ + fails += geist_expect(p->add(be, &x, &z, &y) == GEIST_OK, "add after recovery"); + yp = vt->buffer_map(y.buffer); + fails += geist_expect(yp != nullptr && yp[0] == 2.5f, "buffer_map after recovery"); + vt->buffer_unmap(y.buffer); + } + struct geist_buffer *bufs[] = {x.buffer, z.buffer, y.buffer}; + for (size_t i = 0; i < 3; i++) { + if (bufs[i] != nullptr) { + vt->buffer_destroy(be, bufs[i]); + } + } + return fails; +} + +/* Part 2: the helper itself. */ +static int check_ckd_u32(void) { + uint32_t out = 42; + int fails = 0; + fails += geist_expect(!vk_ckd_u32(0, &out) && out == 0, "0 fits"); + fails += geist_expect(!vk_ckd_u32(UINT32_MAX, &out) && out == UINT32_MAX, "UINT32_MAX fits"); + out = 42; + fails += geist_expect(vk_ckd_u32((size_t) UINT32_MAX + 1u, &out) && out == 42, + "UINT32_MAX + 1 is refused and leaves *out alone"); + fails += geist_expect(vk_ckd_u32(SIZE_MAX, &out) && out == 42, "SIZE_MAX is refused"); + return fails; +} + +/* Part 3: an op refuses a position the shader would see wrapped. */ +static int check_attention_refuses(struct geist_backend *be) { + const struct geist_backend_primitives *p = be->desc->prims; + const struct geist_backend_vtbl *vt = be->desc->vtbl; + /* q [1, 2, 64], k/v [4, 1, 64]: 128 / 256 floats, inside one buffer each. */ + const int64_t qs[3] = {1, 2, 64}, ks[3] = {4, 1, 64}; + struct geist_tensor q = host_tensor(be, 3, qs), k = host_tensor(be, 3, ks), + v = host_tensor(be, 3, ks), o = host_tensor(be, 3, qs); + int fails = geist_expect(q.buffer != nullptr && k.buffer != nullptr && v.buffer != nullptr && + o.buffer != nullptr, + "attention buffers"); + if (fails == 0) { + fails += geist_expect(fill(be, &q, 0.1f) && fill(be, &k, 0.2f) && fill(be, &v, 0.3f) && + fill(be, &o, 0.0f), + "attention fill"); + fails += geist_expect(p->attention(be, &q, &k, &v, 3, 0, &o) == GEIST_OK, + "attention at q_offset 3"); + const size_t wide = (size_t) UINT32_MAX + 4u; /* wraps to 3 */ + enum geist_status s = p->attention(be, &q, &k, &v, wide, 0, &o); + fails += geist_expect(s == GEIST_E_INVALID_ARG && be->err_code == GEIST_E_INVALID_ARG, + "attention refuses a q_offset past UINT32_MAX"); + s = p->attention(be, &q, &k, &v, 3, wide, &o); + fails += geist_expect(s == GEIST_E_INVALID_ARG, + "attention refuses a sliding_window past UINT32_MAX"); + fails += geist_expect(vt->buffer_map(o.buffer) != nullptr, "no failure left pending"); + vt->buffer_unmap(o.buffer); + } + struct geist_buffer *bufs[] = {q.buffer, k.buffer, v.buffer, o.buffer}; + for (size_t i = 0; i < 4; i++) { + if (bufs[i] != nullptr) { + vt->buffer_destroy(be, bufs[i]); + } + } + return fails; +} + +/* Part 4: the x ring is sized against max_m x n_in at resolve. */ +static int check_ring_at_resolve(struct geist_backend *be) { + const struct geist_backend_vtbl *vt = be->desc->vtbl; + struct vk_state *st = be->state; + /* One F32 row: 512 x n_in x 4 B passes the 192 MB ring at 98304. */ + const size_t big_in = (size_t) VK_XRING_CAP / ((size_t) VK_MAX_M * sizeof(float)) + 256u; + float *raw = xmalloc(big_in * sizeof(float)); + memset(raw, 0, big_in * sizeof(float)); + struct geist_weight w = {.raw = raw, + .raw_nbytes = big_in * sizeof(float), + .n_in = (int32_t) big_in, + .n_out = 1, + .dtype = GEIST_DTYPE_F32}; + int fails = 0; + fails += geist_expect(vt->resolve_weight(be, &w) == GEIST_E_INVALID_ARG && + w.linear_m1 == nullptr, + "resolve refuses an n_in whose max_m batch overflows the x ring"); + struct geist_weight ok = {.raw = raw, + .raw_nbytes = 256 * sizeof(float), + .n_in = 256, + .n_out = 1, + .dtype = GEIST_DTYPE_F32}; + fails += geist_expect(vt->resolve_weight(be, &ok) == GEIST_OK && ok.linear_m1 != nullptr, + "resolve accepts a weight that fits"); + fails += geist_expect(st->xring != nullptr && st->argmax_out != nullptr, + "the x ring and the argmax word exist after resolve"); + free(raw); + return fails; +} + +int main(void) { + struct geist_backend *be = nullptr; + enum geist_status s = geist_backend_create("vulkan", nullptr, nullptr, &be); + if (s != GEIST_OK) { + fprintf(stderr, "SKIP: vulkan backend unavailable (%d)\n", (int) s); + return GEIST_TEST_SKIP; + } + int fails = check_failed_submit(be); + fails += check_ckd_u32(); + fails += check_attention_refuses(be); + fails += check_ring_at_resolve(be); + geist_backend_destroy(be); + if (fails > 0) { + fprintf(stderr, "%d check(s) failed\n", fails); + return GEIST_TEST_FAIL; + } + printf("vulkan narrowing and sticky flush: pass\n"); + return GEIST_TEST_PASS; +} + +#else +int main(void) { + fprintf(stderr, "SKIP: vulkan backend not built\n"); + return GEIST_TEST_SKIP; +} +#endif