vulkan: count GPU fallbacks per site, add GEIST_VK_STRICT (#474)

.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

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

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,

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;

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

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;
         }

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.

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