Skip to content

cuda: keep exact Q8 MMA loads within tensor bounds - #978

Open
JordiPosthumus wants to merge 1 commit into
antirez:mainfrom
JordiPosthumus:codex/fix-q8-mma-bounded-loads
Open

cuda: keep exact Q8 MMA loads within tensor bounds#978
JordiPosthumus wants to merge 1 commit into
antirez:mainfrom
JordiPosthumus:codex/fix-q8-mma-bounded-loads

Conversation

@JordiPosthumus

@JordiPosthumus JordiPosthumus commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

Summary

Read each Q8 code word as two 16-bit loads, rather than reading two aligned
32-bit words and combining them with a funnel shift. Add a model-free CUDA
regression for exactly-sized weight allocations.

Why

ldu32_unaligned() can read beyond a raw tensor's allocation. For an aligned
four-byte address, its second load reads an entirely unused next word. For a
halfword-aligned address, that load can include two bytes beyond the tensor.
Padding in a larger model allocation can hide this; an exact tensor allocation
exposes it under Compute Sanitizer and can produce an illegal-access error.

The Q8 MMA caller already requires halfword-aligned weights, and every Q8 code
offset is even. Two 16-bit loads therefore reconstruct exactly the requested
four bytes without touching neighboring storage. The MMA instructions,
activation quantization, scales and floating-point reduction are unchanged.
There are no new options, allocations, cache formats or sampling changes.

Validation

  • Current candidate: 5c21574c248408d4d8752b0fc9a94cd72813f9c3, one commit
    on f4d03f6c. Runtime and fixtures are unchanged by the rebase; Makefile
    cleanup was reconciled with upstream.
  • On 2026-09-05 the exact head passed clean native GB10 builds of ds4-server,
    ds4_test and tests/test_cuda_q8_mma_bounds with CUDA_ARCH=sm_121,
    ./ds4_test --server, the focused test and memcheck (zero errors).
    No engine-source adapter was required on this base.
  • Exact-head M3 make -B -j2 all ds4_test, ./ds4_test --server and
    make -B -j2 cpu passed (CPU compile/link only). SDK 27 emits 27 existing
    Metal deprecation warnings. The red-to-green and model/timing results below
    are earlier unchanged-runtime measurements against b0a147a7, not newly
    repeated whole-model comparisons for this rebase.
  • NVIDIA GB10, 128 GB unified memory / CUDA 13, native sm_121a build. Public-API tests use eight
    prefill tokens, widths 32/96/4096, ranks 1/65/128, and odd/even tensor ends.
  • Original upstream kernel: memcheck reports 64 out-of-bounds reads on the
    first fixture. Patched kernel: all three fixtures pass with exact expected
    outputs and zero memory errors.
  • Memcheck uses --report-api-errors no only to suppress the engine's handled
    cudaHostRegister-unsupported fallback notices. Invalid memory accesses
    remain checked and cause failure.
  • The affected raw kernel microbenchmark (1024 rows, width 4096, 2048 tokens)
    remained approximately 1.92–1.98 ms versus 2.02 ms before; no observed speed
    regression. These are bounded measurements, not a claimed end-to-end gain.
  • A final production integration containing this fix and the separate
    prefill optimization matched the old binary on 8,403,200 float logits:
    32768 -> 36864 context, then 64 teacher-forced steps, Vision-Exp
    IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8, 262144 allocation, 2048 chunks, two
    resident sessions. This is combined integration evidence, not a bounds-only
    whole-model A/B. The standalone red-to-green memory test is above.
  • Default Metal build, make cpu, and ./ds4_test --server passed on M3 Ultra;
    no Mac model was loaded. The full model-backed aggregate suite was not run.
make tests/test_cuda_q8_mma_bounds CUDA_ARCH=sm_121
./tests/test_cuda_q8_mma_bounds
compute-sanitizer --tool memcheck --error-exitcode 99 --report-api-errors no \
  ./tests/test_cuda_q8_mma_bounds

make cuda-regression CUDA_ARCH=sm_121 passed after supplying the existing
tests/cuda_long_context_smoke link target's missing ds4_image.o dependency.
The earlier b0a147a7 native build also needed the unrelated non-Apple
g_tp_block_ctx platform guards, now fixed in the refreshed upstream base.
The old smoke-target link prerequisite is not included in this PR; the new
focused target links its own dependencies.

This was discovered while testing the separate GB10 Q8 prefill optimization #979.
The bounds fix does not require that optimization or any model/launcher change.

@JordiPosthumus
JordiPosthumus force-pushed the codex/fix-q8-mma-bounded-loads branch from e6f4700 to 5c21574 Compare September 5, 2026 20:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant