Fix GLM attention masking and reductions across GPU backends

Use the model-defined 2051-token dense boundary, handle padded selections, and synchronize reused reduction storage in Metal, CUDA and ROCm, including vision normalization.

Enable bounded ROCm selected-attention GEMM for GLM 5.3 without RoPE. Add independent attention references, boundary tests and finite, repeatable vision checks. Focused Strix continued prefill improves from about 21 to 68 tokens per second; stored official scores are unchanged.

.gitignore

@@ -37,6 +37,9 @@
 /tests/test_session_state_gpu
 /tests/test_tp_commands
 /tests/test_glm53_kda
+/tests/test_glm53_vision_engine
+/tests/test_glm_attention
+/tests/test_glm_attention_rocm
 /tests/test_linux_memory
 /tests/test_rocm_memory
 /tests/test_deepseek4_vision_image

Makefile

@@ -368,7 +368,7 @@ tests/test_glm53_kda.o: tests/test_glm53_kda.c ds4_gpu.h
 	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_kda.c
 
 tests/test_glm53_vision_engine.o: tests/test_glm53_vision_engine.c ds4.h ds4_image.h
-	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_vision_engine.c
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) -I. -c -o $@ tests/test_glm53_vision_engine.c
 
 tests/test_glm53_vision_engine: tests/test_glm53_vision_engine.o $(CORE_OBJS)
 ifeq ($(UNAME_S),Darwin)
@@ -414,6 +414,16 @@ $(GLM53_KDA_ROCM_TEST): tests/test_glm53_kda_rocm.o ds4_rocm.o ds4_image.o $(ROC
 test-glm53-kda-rocm: $(GLM53_KDA_ROCM_TEST)
 	./$(GLM53_KDA_ROCM_TEST)
 
+tests/test_glm_attention_rocm.o: tests/test_glm_attention.c ds4.h ds4_gpu.h ds4_linux_memory.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $<
+
+tests/test_glm_attention_rocm: tests/test_glm_attention_rocm.o ds4_rocm.o ds4_image.o $(ROCM_MMQ_OBJS)
+	$(HIPCC) $(ROCM_CFLAGS) -o $@ $^ $(ROCM_LDLIBS)
+
+.PHONY: test-glm-attention-rocm
+test-glm-attention-rocm: tests/test_glm_attention_rocm
+	./tests/test_glm_attention_rocm
+
 tests/test_linux_memory: tests/test_linux_memory.c ds4_linux_memory.h
 	$(CC) $(CFLAGS) -I. -o $@ $<
 
@@ -427,6 +437,21 @@ test-linux-memory: tests/test_linux_memory
 test-rocm-memory: tests/test_rocm_memory
 	./tests/test_rocm_memory
 
+tests/test_glm_attention.o: tests/test_glm_attention.c ds4.h ds4_gpu.h ds4_linux_memory.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) -I. -c -o $@ $<
+
+ifeq ($(UNAME_S),Darwin)
+tests/test_glm_attention: tests/test_glm_attention.o ds4_metal.o ds4_image.o
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+else
+tests/test_glm_attention: tests/test_glm_attention.o ds4_cuda.o ds4_image.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+endif
+
+.PHONY: test-glm-attention
+test-glm-attention: tests/test_glm_attention
+	./tests/test_glm_attention
+
 ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_glm53_vision_gpu.cuh ds4_deepseek4_vision_gpu.cuh ds4_image.h ds4_iq2_tables_cuda.inc cuda/mmq/ds4_mmq.h
 	$(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu
 
@@ -682,6 +707,7 @@ mxfp4-dot-test: tests/test_mxfp4_dot.c
 
 clean:
 	rm -f tests/test_linux_memory tests/test_rocm_memory
+	rm -f tests/test_glm_attention tests/test_glm_attention_rocm
 	rm -f tests/test_session_state tests/test_session_state_gpu tests/test_tp_commands
 	rm -f tests/test_metal_tp_spec
 	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_metal_moe_prefill tests/test_metal_dense_mpp tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_deepseek4_vision_image tests/test_prompt_prefix tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o

docs/SERVER.md

@@ -64,7 +64,7 @@ four times. Idle slots can be cached before reuse; active requests are not evict
 | --- | --- |
 | Metal, resident Flash | Native shared-expert/QKV batching where supported |
 | Metal, GLM 5.2 | Ordered fallback |
-| Metal, GLM 5.3 | Native batching below token 4096; ordered fallback afterward |
+| Metal, GLM 5.3 | Native batching through 2051 visible tokens; ordered fallback afterward |
 | CUDA, supported multi-GPU Flash TP layout | Native grouped decode and mixed prefill/decode |
 | Single-GPU CUDA, including Spark | Ordered fallback |
 

ds4.c

@@ -42057,10 +42057,10 @@ static bool glm53_graph_use_indexed_prefill(
 
 static uint32_t glm_graph_dense_compact_attention_limit(
         const ds4_glm_gpu_graph *g) {
-    /* GLM-5.3 attends densely through its normal 4K work window. Above it,
-     * the pool-4 selector supplies the sparse history. */
-    if (g && g->glm53 && !g->full_kv_cache) return g->ctx_cap;
-    return glm_graph_indexer_top_k_limit();
+    /* The incomplete GLM 5.3 pool is always included. Dense attention is
+     * equivalent through top_k + pool_size - 1, not the work-window size. */
+    return g && g->glm53 ? glm53_graph_indexer_selected_limit() :
+                          glm_graph_indexer_top_k_limit();
 }
 
 static uint32_t glm_graph_limit_indexed_prefill_chunk(
@@ -43270,10 +43270,11 @@ static bool glm_graph_alloc_slice(
     if (g->ctx_size > g->ctx_cap) {
         fprintf(stderr,
                 "ds4: GLM session ctx=%u (model max=%u); "
-                "full-attention prefill/work cap=%u; compact indexed decode is used beyond the cap\n",
+                "prefill/work cap=%u; dense attention limit=%u\n",
                 g->ctx_size,
                 glm_graph_model_context_limit(),
-                g->ctx_cap);
+                g->ctx_cap,
+                glm_graph_dense_compact_attention_limit(g));
     }
 
     const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float);
@@ -50342,6 +50343,11 @@ static bool glm_graph_forward_indexed_tokens(
                 uint32_t slice = n_tokens - t0;
                 if (slice > attn_slice_cap) slice = attn_slice_cap;
                 const bool slice_causal = slice_pos < dense_limit;
+#ifdef DS4_ROCM_BUILD
+                /* The selected-row GEMM gathers a bounded KV tile. Larger
+                 * sparse batches otherwise fall back to scalar attention. */
+                if (!slice_causal && slice > 256u) slice = 256u;
+#endif
                 if (slice_causal && slice > dense_limit - slice_pos) {
                     slice = dense_limit - slice_pos;
                 }
@@ -50422,7 +50428,10 @@ static bool glm_graph_forward_indexed_tokens(
                                 DS4_ROPE_YARN_BETA_FAST,
                                 DS4_ROPE_YARN_BETA_SLOW);
                     } else {
-                        rc = ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
+                        /* GLM 5.3 pads the incomplete pool with invalid IDs. */
+                        rc = (g->glm53 ?
+                                ds4_gpu_glm_attention_indexed_batch_lora_tensor :
+                                ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor)(
                                 attn_lora_view,
                                 q_view,
                                 qk_low_view,

ds4_cuda.cu

@@ -28137,6 +28137,7 @@ __global__ static void glm_dense_attn_causal_softmax_f16_kernel(
         __syncthreads();
     }
     const float max_score = reduce[0];
+    __syncthreads();
 
     float local_sum = 0.0f;
     for (uint32_t col = threadIdx.x; col < visible; col += blockDim.x) {
@@ -28491,7 +28492,7 @@ extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
     return cuda_ok(cudaGetLastError(), "glm attn lora causal launch");
 }
 
-extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
+extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_tensor(
         ds4_gpu_tensor       *lora_out,
         const ds4_gpu_tensor *q,
         const ds4_gpu_tensor *qk_low,
@@ -28559,6 +28560,25 @@ extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
     return cuda_ok(cudaGetLastError(), "glm attn lora selected launch");
 }
 
+extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
+        ds4_gpu_tensor *lora_out,
+        const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *qk_low,
+        const ds4_gpu_tensor *kv_lora_cache,
+        const ds4_gpu_tensor *k_rope_cache,
+        const ds4_gpu_tensor *selected,
+        uint32_t n_tokens, uint32_t n_selected, uint32_t cache_cap,
+        bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim,
+        uint32_t qk_nope, uint32_t qk_rope, uint32_t n_ctx_orig,
+        float freq_base, float freq_scale, float ext_factor,
+        float attn_factor, float beta_fast, float beta_slow) {
+    return ds4_gpu_glm_attention_indexed_batch_lora_tensor(
+        lora_out, q, qk_low, kv_lora_cache, k_rope_cache, selected,
+        n_tokens, n_selected, cache_cap, cache_f16, n_head, kv_lora_dim,
+        qk_nope, qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
+        attn_factor, beta_fast, beta_slow);
+}
+
 extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
         ds4_gpu_tensor       *heads,
         const ds4_gpu_tensor *q,

ds4_glm53_vision_gpu.cuh

@@ -261,6 +261,7 @@ __global__ static void glm53_vision_layernorm_gelu_kernel(
         __syncthreads();
     }
     const float mean = partial[0] / (float)width;
+    __syncthreads();
     float var = 0.0f;
     for (uint32_t d = tid; d < width; d += blockDim.x) {
         const float centered = xr[d] - mean;

metal/dsv4_misc.metal

@@ -874,6 +874,7 @@ kernel void kernel_glm_store_indexer_k(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
     const float mean = scratch[0] / (float)head_dim;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float ss = 0.0f;
     for (uint i = tid; i < head_dim; i += nth) {
@@ -1500,6 +1501,7 @@ kernel void kernel_glm_attention_full(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
     const float max_score = red[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float local_sum = 0.0f;
     for (uint s = tid; s < visible; s += nth) {
@@ -3082,6 +3084,7 @@ static void kernel_glm_attention_indexed_decode_split_group8_reduce_impl(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
     const float max_m = red[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float local_denom = 0.0f;
     if (tid < n_blocks) {
@@ -3266,6 +3269,7 @@ kernel void kernel_glm_attention_indexed_decode(
             threadgroup_barrier(mem_flags::mem_threadgroup);
         }
         const float max_score = red[0];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
 
         float local_sum = 0.0f;
         for (uint s = tid; s < args.n_selected; s += nth) {
@@ -3364,6 +3368,7 @@ kernel void kernel_glm_attention_indexed_decode(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
     const float max_score = red[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float local_sum = 0.0f;
     for (uint s = tid; s < args.n_selected; s += nth) {
@@ -3527,6 +3532,7 @@ kernel void kernel_glm_attention_indexed_batch(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
     const float max_score = red[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float local_sum = 0.0f;
     for (uint s = tid; s < args.n_selected; s += nth) {
@@ -3713,6 +3719,7 @@ kernel void kernel_glm_attention_indexed_batch_group2(
     }
     const float max_score0 = red0[0];
     const float max_score1 = red1[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     float local_sum0 = 0.0f;
     float local_sum1 = 0.0f;

metal/glm53_vision.metal

@@ -231,6 +231,7 @@ kernel void kernel_glm53_vision_layernorm_gelu(
     }
     threadgroup_barrier(mem_flags::mem_threadgroup);
     const float mean = partial[0];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
     float var = 0.0f;
     for (uint d = tid; d < args.width; d += 256u) {
         const float centered = xr[d] - mean;

rocm/ds4_rocm_glm.cuh

@@ -1110,6 +1110,7 @@ __global__ static void glm_store_indexer_k_kernel(
         __syncthreads();
     }
     const float mean = partial[0] / (float)head_dim;
+    __syncthreads();
     float ss = 0.0f;
     for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
         const float d = src[i] - mean;
@@ -1401,6 +1402,7 @@ __global__ static void glm_attention_full_kernel(
         __syncthreads();
     }
     const float max_score = red[0];
+    __syncthreads();
     float local_sum = 0.0f;
     for (uint32_t row = threadIdx.x; row < visible; row += blockDim.x) {
         const float w = expf(scores[row] - max_score);
@@ -1758,6 +1760,7 @@ __global__ static void glm_attention_indexed_decode_split_partial_kernel(
         __syncthreads();
     }
     const float max_score = red[0];
+    __syncthreads();
     float *out = partial_lora + ((uint64_t)block * n_head + head) * kv_lora_dim;
     float *ms = partial_ms + ((uint64_t)block * n_head + head) * 2u;
     if (!isfinite(max_score)) {
@@ -2000,6 +2003,7 @@ __global__ static void glm_attention_indexed_decode_split_reduce_kernel(
         __syncthreads();
     }
     const float max_m = red[0];
+    __syncthreads();
     if (!isfinite(max_m)) {
         for (uint32_t d = threadIdx.x; d < value_dim; d += blockDim.x) {
             heads[(uint64_t)head * value_dim + d] = 0.0f;
@@ -3324,6 +3328,7 @@ __global__ static void glm_causal_gemm_softmax_f16_kernel(
         __syncthreads();
     }
     const float max_score = reduce[0];
+    __syncthreads();
     float sum = 0.0f;
     for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) {
         const float v = s < visible ? expf(row[s] - max_score) : 0.0f;
@@ -3441,8 +3446,10 @@ __global__ static void glm_selected_gemm_rope_to_f16_kernel(
 __global__ static void glm_selected_gemm_softmax_f16_kernel(
         __half *probs,
         float *scores,
+        const int32_t *selected,
         uint32_t n_tokens,
         uint32_t n_selected,
+        uint32_t cache_cap,
         float scale) {
     const uint32_t token = blockIdx.x;
     if (token >= n_tokens) return;
@@ -3450,7 +3457,9 @@ __global__ static void glm_selected_gemm_softmax_f16_kernel(
     __shared__ float reduce[256];
     float vmax = -INFINITY;
     for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) {
-        const float v = row[s] * scale;
+        const int32_t cache_row = selected[(uint64_t)token * n_selected + s];
+        const float v = cache_row >= 0 && (uint32_t)cache_row < cache_cap ?
+            row[s] * scale : -INFINITY;
         row[s] = v;
         vmax = fmaxf(vmax, v);
     }
@@ -3464,9 +3473,10 @@ __global__ static void glm_selected_gemm_softmax_f16_kernel(
         __syncthreads();
     }
     const float max_score = reduce[0];
+    __syncthreads();
     float sum = 0.0f;
     for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) {
-        const float v = expf(row[s] - max_score);
+        const float v = isfinite(max_score) ? expf(row[s] - max_score) : 0.0f;
         row[s] = v;
         sum += v;
     }
@@ -3539,8 +3549,10 @@ __global__ static void glm_selected_gemm_gather_heads_f16_kernel(
 __global__ static void glm_selected_gemm_softmax_heads_f16_kernel(
         __half *probs,
         float *scores,
+        const int32_t *selected,
         uint32_t n_tokens,
         uint32_t n_selected,
+        uint32_t cache_cap,
         uint32_t head_count,
         uint32_t head_capacity,
         float scale) {
@@ -3553,7 +3565,9 @@ __global__ static void glm_selected_gemm_softmax_heads_f16_kernel(
     __shared__ float reduce[256];
     float vmax = -INFINITY;
     for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) {
-        const float v = row[s] * scale;
+        const int32_t cache_row = selected[(uint64_t)token * n_selected + s];
+        const float v = cache_row >= 0 && (uint32_t)cache_row < cache_cap ?
+            row[s] * scale : -INFINITY;
         row[s] = v;
         vmax = fmaxf(vmax, v);
     }
@@ -3567,9 +3581,10 @@ __global__ static void glm_selected_gemm_softmax_heads_f16_kernel(
         __syncthreads();
     }
     const float max_score = reduce[0];
+    __syncthreads();
     float sum = 0.0f;
     for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) {
-        const float v = expf(row[s] - max_score);
+        const float v = isfinite(max_score) ? expf(row[s] - max_score) : 0.0f;
         row[s] = v;
         sum += v;
     }
@@ -4004,11 +4019,11 @@ static int glm_attention_indexed_lora_selected_gemm(
         float beta_fast,
         float beta_slow) {
     if (!g_cublas_ready || !lora_out || !q || !qk_low ||
-        !kv_lora_cache || !k_rope_cache || !selected ||
+        !kv_lora_cache || (qk_rope != 0u && !k_rope_cache) || !selected ||
         n_tokens < 64u || n_tokens > 256u ||
         n_selected == 0u || n_selected > cache_cap ||
         n_head == 0u || kv_lora_dim == 0u ||
-        qk_rope == 0u || (qk_rope & 1u) != 0u ||
+        (qk_rope & 1u) != 0u ||
         n_tokens > INT_MAX || n_selected > INT_MAX ||
         kv_lora_dim > INT_MAX || qk_rope > INT_MAX) {
         return 0;
@@ -4105,32 +4120,34 @@ static int glm_attention_indexed_lora_selected_gemm(
     }
     glm_selected_gemm_profile_end(
         &profile, GLM_SELECTED_PROFILE_KV_GATHER);
-    const uint64_t rope_pairs = selected_rows * (qk_rope >> 1u);
-    const uint32_t rope_blocks =
-        (uint32_t)min((rope_pairs + 255u) / 256u, 4096ull);
-    glm_selected_gemm_profile_begin(&profile);
-    glm_selected_gemm_rope_to_f16_kernel<<<rope_blocks, 256>>>(
-        rope_h,
-        (const char *)k_rope_cache->ptr,
-        (const int32_t *)selected->ptr,
-        n_tokens,
-        n_selected,
-        qk_rope,
-        cache_cap,
-        cache_f16,
-        n_ctx_orig,
-        freq_base,
-        freq_scale,
-        ext_factor,
-        attn_factor,
-        beta_fast,
-        beta_slow);
-    if (!cuda_ok(cudaGetLastError(),
-                 "glm selected attention rope gather launch")) {
-        return 0;
+    if (qk_rope != 0u) {
+        const uint64_t rope_pairs = selected_rows * (qk_rope >> 1u);
+        const uint32_t rope_blocks =
+            (uint32_t)min((rope_pairs + 255u) / 256u, 4096ull);
+        glm_selected_gemm_profile_begin(&profile);
+        glm_selected_gemm_rope_to_f16_kernel<<<rope_blocks, 256>>>(
+            rope_h,
+            (const char *)k_rope_cache->ptr,
+            (const int32_t *)selected->ptr,
+            n_tokens,
+            n_selected,
+            qk_rope,
+            cache_cap,
+            cache_f16,
+            n_ctx_orig,
+            freq_base,
+            freq_scale,
+            ext_factor,
+            attn_factor,
+            beta_fast,
+            beta_slow);
+        if (!cuda_ok(cudaGetLastError(),
+                     "glm selected attention rope gather launch")) {
+            return 0;
+        }
+        glm_selected_gemm_profile_end(
+            &profile, GLM_SELECTED_PROFILE_ROPE_GATHER);
     }
-    glm_selected_gemm_profile_end(
-        &profile, GLM_SELECTED_PROFILE_ROPE_GATHER);
 
     const float one = 1.0f;
     const float zero = 0.0f;
@@ -4213,48 +4230,53 @@ static int glm_attention_indexed_lora_selected_gemm(
         }
         glm_selected_gemm_profile_end(
             &profile, GLM_SELECTED_PROFILE_LORA_SCORE);
-        glm_selected_gemm_profile_begin(&profile);
-        st = cublasGemmStridedBatchedEx(
-            g_cublas,
-            CUBLAS_OP_T,
-            CUBLAS_OP_N,
-            (int)n_selected,
-            (int)tile_heads,
-            (int)qk_rope,
-            &one,
-            rope_h,
-            CUDA_R_16F,
-            (int)qk_rope,
-            rope_stride,
-            qrope_h,
-            CUDA_R_16F,
-            (int)qk_rope,
-            qrope_stride,
-            &one,
-            scores,
-            CUDA_R_32F,
-            (int)n_selected,
-            score_stride,
-            (int)n_tokens,
-            CUBLAS_COMPUTE_32F,
-            CUBLAS_GEMM_DEFAULT);
-        if (!cublas_ok(st, "glm selected attention rope score gemm")) {
-            return 0;
+        if (qk_rope != 0u) {
+            glm_selected_gemm_profile_begin(&profile);
+            st = cublasGemmStridedBatchedEx(
+                g_cublas,
+                CUBLAS_OP_T,
+                CUBLAS_OP_N,
+                (int)n_selected,
+                (int)tile_heads,
+                (int)qk_rope,
+                &one,
+                rope_h,
+                CUDA_R_16F,
+                (int)qk_rope,
+                rope_stride,
+                qrope_h,
+                CUDA_R_16F,
+                (int)qk_rope,
+                qrope_stride,
+                &one,
+                scores,
+                CUDA_R_32F,
+                (int)n_selected,
+                score_stride,
+                (int)n_tokens,
+                CUBLAS_COMPUTE_32F,
+                CUBLAS_GEMM_DEFAULT);
+            if (!cublas_ok(st, "glm selected attention rope score gemm")) {
+                return 0;
+            }
+            glm_selected_gemm_profile_end(
+                &profile, GLM_SELECTED_PROFILE_ROPE_SCORE);
         }
-        glm_selected_gemm_profile_end(
-            &profile, GLM_SELECTED_PROFILE_ROPE_SCORE);
         glm_selected_gemm_profile_begin(&profile);
         if (head_tile == 1u) {
             glm_selected_gemm_softmax_f16_kernel<<<n_tokens, 256>>>(
-                probs, scores, n_tokens, n_selected, scale);
+                probs, scores, (const int32_t *)selected->ptr,
+                n_tokens, n_selected, cache_cap, scale);
         } else {
             const dim3 softmax_grid(n_tokens, tile_heads, 1u);
             glm_selected_gemm_softmax_heads_f16_kernel<<<
                 softmax_grid, 256>>>(
                     probs,
                     scores,
+                    (const int32_t *)selected->ptr,
                     n_tokens,
                     n_selected,
+                    cache_cap,
                     tile_heads,
                     head_tile,
                     scale);

tests/compare_glm53_vision_embeddings.py

@@ -26,7 +26,7 @@ def main():
     dot = ref_sq = candidate_sq = sum_abs = sum_sq = max_abs = 0.0
     nonfinite = 0
     for expected, actual in zip(reference, candidate):
-        if not math.isfinite(actual):
+        if not math.isfinite(expected) or not math.isfinite(actual):
             nonfinite += 1
             continue
         delta = abs(expected - actual)
@@ -38,12 +38,17 @@ def main():
         candidate_sq += actual * actual
 
     count = len(reference)
+    if nonfinite or ref_sq == 0 or candidate_sq == 0:
+        print(f"invalid vision embeddings: nonfinite={nonfinite} "
+              f"reference_norm_sq={ref_sq} candidate_norm_sq={candidate_sq}",
+              file=sys.stderr)
+        return 1
     cosine = dot / math.sqrt(ref_sq * candidate_sq)
     mean_abs = sum_abs / count
     rms = math.sqrt(sum_sq / count)
     print(f"values={count} nonfinite={nonfinite} max_abs={max_abs:.9g} "
           f"mean_abs={mean_abs:.9g} rms={rms:.9g} cosine={cosine:.12f}")
-    if nonfinite or max_abs > 0.06 or mean_abs > 0.001 or cosine < 0.995:
+    if max_abs > 0.06 or mean_abs > 0.001 or cosine < 0.995:
         print("vision embedding comparison failed", file=sys.stderr)
         return 1
     return 0

tests/test_compare_glm53_vision_embeddings.py

@@ -0,0 +1,36 @@
+"""The quality gate must reject corrupt references as well as outputs."""
+
+import array
+import pathlib
+import subprocess
+import sys
+import tempfile
+import unittest
+
+
+class VisionComparisonTest(unittest.TestCase):
+    def compare(self, reference, candidate):
+        with tempfile.TemporaryDirectory() as directory:
+            paths = [pathlib.Path(directory) / name for name in ('ref', 'out')]
+            for path, values in zip(paths, (reference, candidate)):
+                path.write_bytes(array.array('f', values).tobytes())
+            return subprocess.run(
+                [sys.executable,
+                 str(pathlib.Path(__file__).with_name('compare_glm53_vision_embeddings.py')),
+                 *map(str, paths)], capture_output=True, text=True).returncode
+
+    def test_equal(self):
+        self.assertEqual(self.compare([1, 2, 3], [1, 2, 3]), 0)
+
+    def test_invalid(self):
+        for bad in ([float('nan'), 2, 3], [float('inf'), 2, 3], [0, 0, 0], []):
+            with self.subTest(values=bad):
+                self.assertNotEqual(self.compare(bad, [1, 2, 3]), 0)
+                self.assertNotEqual(self.compare([1, 2, 3], bad), 0)
+
+    def test_different(self):
+        self.assertNotEqual(self.compare([1, 2, 3], [-1, -2, -3]), 0)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_glm53_vision_engine.c

@@ -2,6 +2,7 @@
 #include "ds4_image.h"
 
 #include <errno.h>
+#include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -55,6 +56,21 @@ int main(int argc, char **argv) {
             ds4_engine_close(engine);
             return 1;
         }
+        const size_t count = (size_t)next.token_count * 4096u;
+        bool valid = next.data != NULL && count != 0;
+        for (size_t j = 0; valid && j < count; j++)
+            valid = isfinite(next.data[j]);
+        if (valid && embedding.data) {
+            valid = next.token_count == embedding.token_count &&
+                memcmp(next.data, embedding.data, count * sizeof(float)) == 0;
+        }
+        if (!valid) {
+            fprintf(stderr, "vision embedding is nonfinite, empty, or not repeatable\n");
+            ds4_vision_embedding_free(&next);
+            ds4_vision_embedding_free(&embedding);
+            ds4_engine_close(engine);
+            return 1;
+        }
         if (repeats > 1u) {
             fprintf(stderr, "vision encode %u/%u: %.3f s\n",
                     i + 1u, repeats, wall_seconds() - start);

tests/test_glm_attention.c

@@ -0,0 +1,357 @@
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <time.h>
+
+#include "ds4.h"
+#include "ds4_gpu.h"
+#include "ds4_linux_memory.h"
+
+bool ds4_log_is_tty(FILE *fp) { (void)fp; return false; }
+
+static void check(int ok, const char *what) {
+    if (!ok) { fprintf(stderr, "%s failed\n", what); exit(1); }
+}
+
+static double seconds(void) {
+    struct timespec ts;
+    clock_gettime(CLOCK_MONOTONIC, &ts);
+    return ts.tv_sec + ts.tv_nsec * 1e-9;
+}
+
+static bool benchmark_has_headroom(void) {
+    uint64_t available;
+    return ds4_linux_nonmovable_memory(&available) && available >= (32ull << 30);
+}
+
+static ds4_gpu_tensor *upload(const void *data, size_t bytes) {
+    ds4_gpu_tensor *t = ds4_gpu_tensor_alloc(bytes);
+    check(t && ds4_gpu_tensor_write(t, 0, data, bytes), "tensor upload");
+    return t;
+}
+
+static ds4_gpu_tensor *upload_cache(const float *data, size_t count, bool f16) {
+    if (!f16) return upload(data, count * sizeof(float));
+    _Float16 *half = malloc(count * sizeof(*half));
+    check(half != NULL, "half cache allocation");
+    for (size_t i = 0; i < count; i++) half[i] = data[i];
+    ds4_gpu_tensor *t = upload(half, count * sizeof(*half));
+    free(half);
+    return t;
+}
+
+static void attention_reference(float *out, const float *low, const float *kv,
+        const int32_t *ids, uint32_t tokens, uint32_t heads, uint32_t dim,
+        uint32_t pos, uint32_t cap, uint32_t count, uint32_t nope, bool selected) {
+    double *weights = calloc(count, sizeof(*weights));
+    check(weights != NULL, "reference allocation");
+    for (uint32_t t = 0; t < tokens; t++) {
+        uint32_t visible = selected ? count : (pos + t + 1);
+        if (visible > count) visible = count;
+        for (uint32_t h = 0; h < heads; h++) {
+            const size_t base = ((size_t)t * heads + h) * dim;
+            double sum = 0;
+            for (uint32_t i = 0; i < visible; i++) {
+                const uint32_t row = selected ? (uint32_t)ids[t * count + i] : i;
+                double dot = 0;
+                weights[i] = 0;
+                if (row >= cap) continue;
+                for (uint32_t d = 0; d < dim; d++)
+                    dot += (double)low[base + d] * kv[(size_t)row * dim + d];
+                sum += weights[i] = exp(dot / sqrt(nope));
+            }
+            for (uint32_t d = 0; d < dim; d++) {
+                double value = 0;
+                for (uint32_t i = 0; i < visible; i++) {
+                    const uint32_t row = selected ? (uint32_t)ids[t * count + i] : i;
+                    if (row < cap) value += weights[i] * kv[(size_t)row * dim + d];
+                }
+                out[base + d] = sum > 0 ? value / sum : 0;
+            }
+        }
+    }
+    free(weights);
+}
+
+static void attention_case(uint32_t tokens, uint32_t heads, uint32_t dim,
+                           uint32_t pos, uint32_t cap, bool selected,
+                           bool zero_rope, bool bench, bool f16) {
+    const uint32_t nope = zero_rope ? 256 : 32, rope = zero_rope ? 0 : 8;
+    const uint32_t count = selected ? (cap > 2048 ? 2051 : 8) : cap;
+    const bool known_answer = selected && !f16;
+    const size_t out_n = (size_t)tokens * heads * dim;
+    const size_t q_n = (size_t)tokens * heads * (nope + rope);
+    const size_t kv_n = (size_t)cap * dim;
+    const size_t rope_n = (size_t)cap * rope;
+    float *q = calloc(q_n, sizeof(float));
+    float *low = calloc(out_n, sizeof(float));
+    float *kv = malloc(kv_n * sizeof(float));
+    float *kr = calloc(rope_n ? rope_n : 1, sizeof(float));
+    int32_t *ids = malloc((size_t)tokens * count * sizeof(int32_t));
+    float *ref = malloc(out_n * sizeof(float));
+    float *actual = malloc(out_n * sizeof(float));
+    float *repeat = malloc(out_n * sizeof(float));
+    check(q && low && kv && kr && ids && ref && actual && repeat,
+          "host allocation");
+    for (size_t i = 0; i < kv_n; i++)
+        kv[i] = known_answer ? 1.0f : ((int)(i % 31) - 15) * 0.03125f;
+    if (!known_answer) {
+        for (size_t i = 0; i < out_n; i++)
+            low[i] = ((int)(i % 23) - 11) * 0.015625f;
+        for (size_t i = 0; i < q_n; i++)
+            q[i] = ((int)(i % 17) - 8) * 0.03125f;
+        for (size_t i = 0; i < rope_n; i++)
+            kr[i] = ((int)(i % 13) - 6) * 0.0625f;
+    }
+    for (uint32_t t = 0; t < tokens; t++) {
+        for (uint32_t i = 0; i < count; i++) {
+            /* Include valid, padded, out-of-range and entirely empty rows. */
+            ids[(size_t)t * count + i] = t % 3 == 0 ? -1 :
+                (t % 3 == 1 && i % 2 == 0 ? (i == 0 ? -1 : (int32_t)cap) :
+                 (int32_t)(i % cap));
+        }
+    }
+    ds4_gpu_tensor *qg = upload(q, q_n * sizeof(float));
+    ds4_gpu_tensor *lg = upload(low, out_n * sizeof(float));
+    ds4_gpu_tensor *kg = upload_cache(kv, kv_n, f16);
+    ds4_gpu_tensor *rg = upload_cache(kr, rope_n ? rope_n : 1, f16);
+    ds4_gpu_tensor *ig = selected ? upload(ids, (size_t)tokens * count * sizeof(int32_t)) : NULL;
+    ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_n * sizeof(float));
+    check(out != NULL, "output allocation");
+    float *oracle = NULL;
+    if (zero_rope && !bench) {
+        oracle = malloc(out_n * sizeof(float));
+        check(oracle != NULL, "attention oracle allocation");
+        attention_reference(oracle, low, kv, ids, tokens, heads, dim,
+                            pos, cap, count, nope, selected);
+    }
+    const char *flag = selected ? "DS4_ROCM_GLM_SELECTED_ATTN_GEMM" :
+                                  "DS4_ROCM_GLM_CAUSAL_ATTN_GEMM";
+    double elapsed = 0;
+    const int runs = bench ? 6 : 3;
+    for (int run = 0; run < runs; run++) {
+        if (!bench && run == 0) setenv(flag, "0", 1);
+        else unsetenv(flag);
+        check(ds4_gpu_synchronize(), "attention pre-sync");
+        const double start = seconds();
+        int ok = selected ? ds4_gpu_glm_attention_indexed_batch_lora_tensor(
+            out, qg, lg, kg, rg, ig, tokens, count, cap, f16,
+            heads, dim, nope, rope, 4096, 10000, 1, 0, 1, 32, 1) :
+            ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
+                out, qg, lg, kg, rg, tokens, pos, count, cap, f16,
+                heads, dim, nope, rope, 4096, 10000, 1, 0, 1, 32, 1);
+        check(ok && ds4_gpu_synchronize(), "attention launch");
+        if (run > 0) elapsed += seconds() - start;
+        float *result = run == 0 ? ref : (run == 1 ? actual : repeat);
+        check(ds4_gpu_tensor_read(out, 0, result, out_n * sizeof(float)),
+              "attention read");
+        for (size_t i = 0; i < out_n; i++) {
+            const float expected = known_answer ?
+                (i / (heads * dim) % 3 == 0 ? 0.0f : 1.0f) :
+                (oracle ? oracle[i] : ref[i]);
+            if (!isfinite(result[i]) || fabsf(result[i] - expected) > 0.002f) {
+                fprintf(stderr, "attention selected=%d tokens=%u heads=%u pos=%u "
+                        "run=%d index=%zu: %.9g != %.9g\n", selected, tokens,
+                        heads, pos, run, i, result[i], expected);
+                exit(1);
+            }
+        }
+        if (run >= 2)
+            check(memcmp(actual, repeat, out_n * sizeof(float)) == 0,
+                  "attention repeat determinism");
+    }
+    unsetenv(flag);
+#ifndef DS4_ROCM_BUILD
+    if (!bench && !selected && zero_rope && f16) {
+        for (int run = 0; run < 3; run++) {
+            check(ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
+                out, lg, kg, pos, tokens, cap, cap, f16, heads, dim, nope),
+                "dense compact GEMM attention");
+            float *result = run == 0 ? actual : repeat;
+            check(ds4_gpu_tensor_read(out, 0, result, out_n * sizeof(float)),
+                  "dense compact attention read");
+            for (size_t i = 0; i < out_n; i++)
+                check(isfinite(result[i]) && fabsf(result[i] - oracle[i]) < 0.002f,
+                      "dense compact attention reference");
+            if (run > 0)
+                check(memcmp(actual, repeat, out_n * sizeof(float)) == 0,
+                      "dense compact attention repeat determinism");
+        }
+    }
+#endif
+    printf("attention selected=%d tokens=%u heads=%u dim=%u pos=%u cap=%u rope=%u f16=%d: "
+           "PASS, %.3f ms\n", selected, tokens, heads, dim, pos, cap, rope,
+           f16, elapsed * 1000 / (runs - 1));
+    ds4_gpu_tensor_free(out);
+    ds4_gpu_tensor_free(ig);
+    ds4_gpu_tensor_free(rg);
+    ds4_gpu_tensor_free(kg);
+    ds4_gpu_tensor_free(lg);
+    ds4_gpu_tensor_free(qg);
+    free(oracle);
+    free(repeat); free(actual); free(ref); free(ids);
+    free(kr); free(kv); free(low); free(q);
+}
+
+static void reduction_cases(void) {
+    enum { D = 512, N = 128, T = 3, C = 521, H = 8, V = 4, R = 64, MODEL_BYTES = 65536 };
+    unsigned char *model = mmap(NULL, MODEL_BYTES, PROT_READ | PROT_WRITE,
+                                MAP_PRIVATE | MAP_ANON, -1, 0);
+    check(model != MAP_FAILED, "test model allocation");
+    float *weight = (float *)model, *bias = weight + N;
+    float raw[T * N], result[T * N];
+    for (int i = 0; i < N; i++) { weight[i] = 1; bias[i] = 0.25f; }
+    for (int i = 0; i < T * N; i++) raw[i] = 128 + (i % 29) * 0.125f;
+    /* Q8 value projections: every weight is exactly 1. */
+    for (int i = 0; i < H * V * (D / 32); i++) {
+        unsigned char *block = model + 4096 + i * 34;
+        _Float16 scale = 1;
+        memcpy(block, &scale, sizeof(scale));
+        memset(block + 2, 1, 32);
+    }
+    check(ds4_gpu_set_model_map(model, MODEL_BYTES), "test model registration");
+    ds4_gpu_tensor *x = upload(raw, sizeof(raw));
+    ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(sizeof(raw));
+    for (int run = 0; run < 20; run++) {
+        check(ds4_gpu_glm_store_indexer_k_tensor(out, x, model, MODEL_BYTES,
+              0, N * sizeof(float), 0, T, T, N, R, 4096,
+              1e-5f, 10000, 1, 0, 1, 32, 1, false), "indexer normalization");
+        check(ds4_gpu_tensor_read(out, 0, result, sizeof(result)), "indexer read");
+        for (int t = 0; t < T; t++) {
+            double mean = 0, variance = 0;
+            for (int i = 0; i < N; i++) mean += raw[t * N + i];
+            mean /= N;
+            for (int i = 0; i < N; i++) {
+                double d = raw[t * N + i] - mean;
+                variance += d * d;
+            }
+            for (int i = 0; i < N; i++) {
+                double expected = (raw[t * N + i] - mean) /
+                    sqrt(variance / N + 1e-5) + 0.25;
+                if (i < R) {
+                    const int pair = i ^ 1;
+                    const double other = (raw[t * N + pair] - mean) /
+                        sqrt(variance / N + 1e-5) + 0.25;
+                    const double theta = t * pow(10000, -(double)(i & ~1) / R);
+                    expected = expected * cos(theta) +
+                        (i & 1 ? other : -other) * sin(theta);
+                }
+                check(isfinite(result[t * N + i]) &&
+                      fabs(result[t * N + i] - expected) < 1e-4,
+                      "indexer normalization reference");
+            }
+        }
+    }
+    ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(x);
+
+    float q[T * R], keys[C * R], values[C * V], full[T * V];
+    for (int i = 0; i < T * R; i++) q[i] = (i % 7 - 3) * 0.125f;
+    for (int i = 0; i < C * R; i++) keys[i] = (i % 11 - 5) * 0.125f;
+    for (int i = 0; i < C * V; i++) values[i] = (i % 19 - 9) * 0.125f;
+    ds4_gpu_tensor *qg = upload(q, sizeof(q));
+    ds4_gpu_tensor *kg = upload(keys, sizeof(keys));
+    ds4_gpu_tensor *vg = upload(values, sizeof(values));
+    out = ds4_gpu_tensor_alloc(sizeof(full));
+    for (int run = 0; run < 20; run++) {
+        check(ds4_gpu_glm_attention_full_tensor(out, qg, kg, vg,
+              C - T, T, C, C, 1, R, V, false), "full attention");
+        check(ds4_gpu_tensor_read(out, 0, full, sizeof(full)), "full attention read");
+        for (int t = 0; t < T; t++) {
+            double sum = 0, acc[V] = {0};
+            for (int r = 0; r <= C - T + t; r++) {
+                double score = 0;
+                for (int d = 0; d < R; d++) score += q[t * R + d] * keys[r * R + d];
+                double p = exp(score / sqrt(R));
+                sum += p;
+                for (int d = 0; d < V; d++) acc[d] += p * values[r * V + d];
+            }
+            for (int d = 0; d < V; d++)
+                check(isfinite(full[t * V + d]) &&
+                      fabs(full[t * V + d] - acc[d] / sum) < 1e-5,
+                      "full attention reference");
+        }
+    }
+    ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(vg);
+    ds4_gpu_tensor_free(kg); ds4_gpu_tensor_free(qg);
+
+#if defined(__APPLE__) || defined(DS4_ROCM_BUILD)
+    float *kv = malloc(C * D * sizeof(float));
+    float low[H * D] = {0}, sq[H * (32 + R)] = {0}, rope[C * R] = {0};
+    int32_t selected[C];
+    check(kv != NULL, "split cache allocation");
+    for (int i = 0; i < C * D; i++) kv[i] = 1;
+    kg = upload_cache(kv, C * D, true);
+    ds4_gpu_tensor *rg = upload_cache(rope, C * R, true);
+    qg = upload(sq, sizeof(sq));
+    ds4_gpu_tensor *lg = upload(low, sizeof(low));
+    ds4_gpu_tensor *ig = ds4_gpu_tensor_alloc(sizeof(selected));
+    ds4_gpu_tensor *pl = ds4_gpu_tensor_alloc(4 * H * D * sizeof(float));
+    ds4_gpu_tensor *pm = ds4_gpu_tensor_alloc(4 * H * 2 * sizeof(float));
+    out = ds4_gpu_tensor_alloc(H * V * sizeof(float));
+    float split[H * V];
+    for (int mode = 0; mode < 3; mode++) {
+        for (int i = 0; i < C; i++) selected[i] = mode == 2 ? -1 :
+            (mode == 1 && i % 2 == 0 ? (i % 4 == 0 ? -1 : C) : i);
+        check(ds4_gpu_tensor_write(ig, 0, selected, sizeof(selected)), "split ids");
+        for (int run = 0; run < 20; run++) {
+            check(ds4_gpu_glm_attention_indexed_decode_split_group8_tensor(
+                  out, pl, pm, qg, lg, kg, rg, model, MODEL_BYTES, 4096,
+                  ig, C, mode == 0, C, true, H, D, 32, R, V, 4096,
+                  192, 4, 10000, 1, 0, 1, 32, 1), "split attention");
+            check(ds4_gpu_tensor_read(out, 0, split, sizeof(split)), "split read");
+            for (int i = 0; i < H * V; i++)
+                check(isfinite(split[i]) && fabsf(split[i] - (mode == 2 ? 0 : D)) < 0.01f,
+                      "split attention reference");
+        }
+    }
+    ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(pm); ds4_gpu_tensor_free(pl);
+    ds4_gpu_tensor_free(ig); ds4_gpu_tensor_free(lg); ds4_gpu_tensor_free(qg);
+    ds4_gpu_tensor_free(rg); ds4_gpu_tensor_free(kg); free(kv);
+#endif
+    ds4_gpu_cleanup();
+    munmap(model, MODEL_BYTES);
+    puts("normalization and attention reductions: PASS");
+}
+
+int main(int argc, char **argv) {
+    const bool bench = argc == 2 && strcmp(argv[1], "--bench") == 0;
+    check(argc == 1 || bench, "arguments");
+    if (bench && !benchmark_has_headroom()) {
+        fprintf(stderr, "Unload the model before benchmarking; 32 GiB available memory required.\n");
+        return 1;
+    }
+    check(ds4_gpu_init(), "GPU initialization");
+    if (bench) {
+        attention_case(256, 64, 512, 0, 256, false, false, true, false);
+        attention_case(512, 64, 512, 0, 512, false, false, true, false);
+        attention_case(1024, 64, 512, 0, 1024, false, false, true, false);
+        attention_case(256, 64, 512, 3840, 4096, false, false, true, false);
+        attention_case(256, 64, 512, 0, 256, false, true, true, true);
+        attention_case(1024, 64, 512, 0, 1024, false, true, true, true);
+        attention_case(2048, 64, 512, 0, 2048, false, true, true, true);
+        attention_case(2048, 64, 512, 2048, 4096, false, true, true, true);
+        attention_case(256, 64, 512, 3840, 4096, false, true, true, true);
+        ds4_gpu_cleanup();
+    } else {
+#ifdef DS4_ROCM_BUILD
+        for (int f16 = 0; f16 < 2; f16++) {
+            attention_case(64, 17, 32, 0, 32, true, false, false, f16);
+            attention_case(67, 1, 32, 0, 32, true, false, false, f16);
+            attention_case(128, 17, 32, 0, 32, true, false, false, f16);
+            attention_case(256, 17, 32, 0, 32, true, false, false, f16);
+            attention_case(256, 17, 32, 0, 256, false, false, false, f16);
+            attention_case(257, 17, 32, 127, 384, false, false, false, f16);
+            attention_case(256, 1, 32, 128, 384, false, true, false, f16);
+        }
+#endif
+        attention_case(67, 17, 512, 0, 32, true, true, false, false);
+        attention_case(67, 17, 512, 0, 32, true, true, false, true);
+        attention_case(64, 17, 512, 0, 2052, true, true, false, true);
+        attention_case(257, 17, 512, 127, 384, false, true, false, true);
+        reduction_cases();
+    }
+    return 0;
+}

tests/test_session_state.c

@@ -116,6 +116,25 @@ static void test_payload_tokens(void) {
 }
 
 #ifndef DS4_NO_GPU
+static void test_glm_attention_budget(void) {
+    const ds4_shape saved_shape = g_ds4_shape;
+    g_ds4_shape = DS4_SHAPE_GLM53;
+    ds4_glm_gpu_graph *g = calloc(1, sizeof(*g));
+    assert(g);
+    g->glm53 = true;
+    g->compact_cache_cap = 16384;
+    const uint32_t capacities[] = {1024, 4096, 8192, 16384};
+    for (size_t i = 0; i < sizeof(capacities) / sizeof(*capacities); i++) {
+        g->ctx_cap = capacities[i];
+        assert(glm_graph_dense_compact_attention_limit(g) == 2051);
+        g->full_kv_cache = true;
+        assert(glm_graph_dense_compact_attention_limit(g) == 2051);
+        g->full_kv_cache = false;
+    }
+    free(g);
+    g_ds4_shape = saved_shape;
+}
+
 enum { POOL_DIM = 128, POOL_ROWS = 12, MODEL_BYTES = 16384 };
 static float pool_raw[POOL_ROWS * POOL_DIM], pool_gates[POOL_ROWS * POOL_DIM];
 
@@ -213,6 +232,7 @@ int main(void) {
     test_session_memory();
     test_payload_tokens();
 #ifndef DS4_NO_GPU
+    test_glm_attention_budget();
     test_glm_spec_rollback();
 #endif
     puts("session state tests: ok");