cuda: keep exact Q8 MMA loads within tensor bounds

.gitignore

@@ -23,6 +23,7 @@
 /tests/test_cuda_mixed_batch
 /tests/test_cuda_session_batch
 /tests/cuda_long_context_smoke
+/tests/test_cuda_q8_mma_bounds
 /tests/test_layer_pack
 /tests/test_metal_session_batch
 /tests/test_metal_tp_spec

Makefile

@@ -267,8 +267,9 @@ cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_eval_case
 	$(CC) $(CFLAGS) -o ds4-eval ds4_eval_cpu.o ds4_eval_cases.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
 	$(CC) $(CFLAGS) -o ds4-agent ds4_agent_cpu.o ds4_help.o ds4_prompt_prefix.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
 
-cuda-regression: tests/cuda_long_context_smoke
+cuda-regression: tests/cuda_long_context_smoke tests/test_cuda_q8_mma_bounds
 	./tests/cuda_long_context_smoke
+	./tests/test_cuda_q8_mma_bounds
 
 tests/test_mxfp4_cuda: tests/test_mxfp4_cuda.cu $(MMQ_OBJS)
 	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -o $@ $^ $(CUDA_LDLIBS)
@@ -524,6 +525,12 @@ ds4_rocm_unavailable.o: ds4_rocm_unavailable.cu
 tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o $(MMQ_OBJS)
 	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
 
+tests/test_cuda_q8_mma_bounds.o: tests/test_cuda_q8_mma_bounds.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_cuda_q8_mma_bounds: tests/test_cuda_q8_mma_bounds.o ds4_cuda.o ds4_image.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
 tests/test_layer_pack.o: tests/test_layer_pack.c ds4_layer_pack.h
 	$(CC) $(CFLAGS) -I. -c -o $@ $<
 
@@ -710,4 +717,4 @@ clean:
 	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
+	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 tests/test_cuda_q8_mma_bounds

ds4_cuda.cu

@@ -6042,11 +6042,10 @@ __global__ static void matmul_q8_0_preq_batch_tok2_exact_kernel(
  * kernels across shapes, including blocks < T and ragged out_dim/n_tok.
  * Rollback: DS4_CUDA_NO_Q8_MMA=1. */
 __device__ __forceinline__ static uint32_t ldu32_unaligned(const uint8_t *p) {
-    const uintptr_t addr = (uintptr_t)p;
-    const uint32_t *base = (const uint32_t *)(addr & ~(uintptr_t)3);
-    const uint32_t lo = base[0];
-    const uint32_t hi = base[1];
-    return __funnelshift_r(lo, hi, (uint32_t)(addr & 3u) * 8u);
+    /* Q8_0 codes are halfword-aligned. Read only the four requested bytes,
+     * including at the end of an exactly-sized raw tensor allocation. */
+    const uint16_t *q = (const uint16_t *)p;
+    return (uint32_t)q[0] | ((uint32_t)q[1] << 16);
 }
 
 __device__ __forceinline__ static void mma_m16n8k32_s8(

tests/test_cuda_q8_mma_bounds.c

@@ -0,0 +1,79 @@
+/* Q8 MMA must not read beyond an exactly-sized raw weight allocation. */
+#include "ds4_gpu.h"
+#include <cuda_runtime.h>
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define CHECK(c) do { if (!(c)) { \
+    fprintf(stderr, "Q8 bounds failure at line %d\n", __LINE__); goto cleanup; \
+} } while (0)
+
+static void ones(unsigned char *w, size_t blocks) {
+    for (size_t b = 0; b < blocks; b++) {
+        const uint16_t scale = 0x3c00; /* FP16 1 */
+        memcpy(w + b * 34, &scale, 2);
+        memset(w + b * 34 + 2, 1, 32);
+    }
+}
+
+static int check_shape(uint64_t width, uint64_t rank, uint32_t groups) {
+    const uint32_t tokens = 8;
+    const uint64_t low_dim = rank * groups, out_dim = 3;
+    const size_t a_bytes = low_dim * (width / 32) * 34;
+    const size_t b_bytes = out_dim * ((low_dim + 31) / 32) * 34;
+    const size_t nx = tokens * groups * width;
+    const size_t nl = tokens * low_dim, ny = tokens * out_dim;
+    unsigned char *model = malloc(a_bytes + b_bytes + 4);
+    float *x = malloc(nx * sizeof(float)), *values = malloc((nl + ny) * sizeof(float));
+    ds4_gpu_tensor *heads = NULL, *low = NULL, *out = NULL;
+    int rc = 1, initialized = 0;
+    CHECK(model && x && values);
+    ones(model, (a_bytes + b_bytes) / 34);
+    memset(model + a_bytes + b_bytes, 0, 4);
+    for (size_t i = 0; i < nx; i++) x[i] = 127.0f;
+    CHECK(ds4_gpu_init()); initialized = 1;
+    CHECK(ds4_gpu_set_model_map(model, a_bytes + b_bytes + 4));
+    /* Cache only the exact tensor spans, not the host buffer's end padding. */
+    CHECK(ds4_gpu_cache_model_range(model, a_bytes + b_bytes + 4,
+                                   0, a_bytes, "Q8 bounds A"));
+    CHECK(ds4_gpu_cache_model_range(model, a_bytes + b_bytes + 4,
+                                   a_bytes, b_bytes, "Q8 bounds B"));
+    heads = ds4_gpu_tensor_alloc(nx * sizeof(float));
+    low = ds4_gpu_tensor_alloc(nl * sizeof(float));
+    out = ds4_gpu_tensor_alloc(ny * sizeof(float));
+    CHECK(heads && low && out);
+    CHECK(ds4_gpu_tensor_write(heads, 0, x, nx * sizeof(float)));
+    CHECK(ds4_gpu_attention_output_q8_batch_tensor(out, low, NULL, NULL,
+        model, a_bytes + b_bytes + 4, 0, a_bytes, width, rank, groups, out_dim,
+        heads, tokens));
+    CHECK(ds4_gpu_synchronize());
+    CHECK(ds4_gpu_tensor_read(low, 0, values, nl * sizeof(float)));
+    CHECK(ds4_gpu_tensor_read(out, 0, values + nl, ny * sizeof(float)));
+    for (size_t i = 0; i < nl; i++) CHECK(values[i] == (float)(width * 127));
+    for (size_t i = 0; i < ny; i++) CHECK(values[nl + i] == (float)(low_dim * width * 127));
+    rc = 0;
+cleanup:
+    if (rc) fprintf(stderr, "width=%llu rank=%llu groups=%u\n",
+                    (unsigned long long)width, (unsigned long long)rank, groups);
+    ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(low); ds4_gpu_tensor_free(heads);
+    if (initialized) ds4_gpu_cleanup();
+    free(values); free(x); free(model);
+    return rc;
+}
+
+int main(void) {
+    struct cudaDeviceProp prop;
+    if (cudaGetDeviceProperties(&prop, 0) != cudaSuccess || prop.major < 8) {
+        fprintf(stderr, "SKIP: exact Q8 MMA requires an sm_80+ CUDA device\n");
+        return 0;
+    }
+    setenv("DS4_CUDA_NO_CUBLAS_ATTENTION_OUTPUT_A", "1", 1);
+    unsetenv("DS4_CUDA_COPY_MODEL");
+    if (check_shape(32, 1, 1) || check_shape(96, 65, 3) ||
+        check_shape(4096, 128, 4)) return 1;
+    fprintf(stderr, "PASS: Q8 MMA exact-size weights, including odd halfword tails\n");
+    return 0;
+}