Skip to content

CUDA: TQ1_0 mat-vec is load-bound by the per-32-element chunk contract #139

Description

@bri-prism

Current behavior

Parent: #138.

A local patch adds a TQ1_0 dequantize kernel, a vec_dot for the quantized
mat-vec path, CUDA type traits and the capability entry. It is correct, and it
fixes prefill, but decode ends up slower than not having the kernel at all.

The cause is a layout mismatch with the mat-vec calling convention, not a tuning
problem. The quantized mat-vec path calls
vec_dot(vbq, bq8_1, kbx, iqs) once per 32-element chunk, and the callee is
expected to handle elements 32*iqs through 32*iqs+31. That contract assumes a
byte-aligned format where a 32-element chunk maps to a contiguous, exclusive slice
of the packed bytes.

TQ1_0 does not have that property. It packs five trits per byte in base 3, so the
digits of one byte are strided across the block: chunk c for c in 0 to 4 reads
digit c of bytes qs[0..31]. Chunks 0 through 4 therefore all read the same 32
packed bytes and extract one different trit from each. Five chunks, five full
re-reads of the same bytes, so the mat-vec path issues about 5x the global loads
the CPU reference needs for the same weights. On top of that, base 3 cannot be
extracted with a shift and a mask; each element costs a multiply and a shift plus
a three-way select.

The CPU reference wins precisely because it does not have this constraint.
dequantize_row_tq1_0 in ggml/src/ggml-quants.c walks the 256-element block once
and amortizes each packed byte across all five of its trits. The format was
designed for CPU SIMD working a block at a time, and the GPU mat-vec contract cuts
across it.

Measured impact

Same internal measurement as the umbrella issue: large mixture-of-experts model,
one NVIDIA data-center GPU, llama-bench with r=3. Tokens per second, without
and with the local patch:

no CUDA kernel (CPU fallback) local patch Q1_0 reference
pp512 36.75 207.06 (+/- 21.07) 2722.59
tg128 21.72 11.85 (+/- 0.01) 189.23

Prefill improves 5.6x, entirely because the ops stop being scheduled to the CPU.
Decode regresses 45 percent and lands below the CPU fallback. Both paths remain
more than an order of magnitude off the low-bit types that have proper kernels.
Correctness was verified against the CPU reference through test-backend-ops -o MUL_MAT
on all TQ1_0 shapes plus a model smoke run, so this is purely a performance shape
problem.

Upstream state

No CUDA implementation exists to compare against, and ggml-org#11183 for
TQ2_0 has been open and stalled since January 2025.

The one existing GPU solution to this exact layout problem is the Vulkan work in
ggml-org#27765 (open, approved). Its mat-vec shader inverts the loop:
instead of indexing by element and deriving the packed byte, each thread loads one
packed byte and expands all of its digits in place, accumulating each contribution
against the matching activation. Every packed byte is loaded exactly once. It also
factors the base-3 decode into shared helpers and packs the powers of three into a
single word so they extract with a shift and a mask instead of living in a constant
array that may not stay in registers.

Proposed approach

Do not try to make TQ1_0 fit the per-32-element chunk contract. Two options, in
order of preference:

  1. A dedicated mat-vec kernel for the type that consumes a whole 256-element block
    per invocation and iterates over packed bytes rather than elements, following
    the structure in vulkan: add TQ1_0 support (mm, mat-vec, mat-vec-id, dequant, get_rows) ggml-org/llama.cpp#27765. This means opting the type out of the shared chunked
    dispatch rather than registering a vec_dot into it.
  2. Failing that, keep the type off the mat-vec path entirely and let low-batch
    matmuls take the dequantize-and-cuBLAS route, which at least does not regress
    below the CPU. This is a stopgap, not a fix, and should be tied to a measurement
    showing it beats the CPU fallback.

One structural warning for whoever picks this up. There is no single choke point at
which the mat-vec path can be gated per type. GGML_OP_MUL_MAT and
GGML_OP_MUL_MAT_ID share one type switch in the capability check, so enabling the
type for plain matmul silently enables it for every MoE expert too. The two then
gate on different predicates: plain matmul consults ggml_cuda_should_use_mmvq,
while the expert path consults get_mmvq_mmid_max_batch, which on recent
architectures returns the maximum batch size without consulting the type at all. A
third gate lives in the fusion predicate. Guarding is the wrong shape of fix here;
implementing a kernel that suits the layout is the right one.

Acceptance criteria

  • test-backend-ops -b CUDA0 -o MUL_MAT passes for tq1_0, with the type
    uncommented in all_types and base_types.
  • test-backend-ops -b CUDA0 -o MUL_MAT_ID passes for tq1_0 at ne11 = 1 with
    more than one block per row and fewer experts used than available.
  • llama-bench tg128 on the same model and GPU: TQ1_0 within 25 percent of Q2_0.
    Below the CPU fallback number is a fail. Above the CPU fallback but still far
    from Q2_0 should be reported with a profile attributing the remaining gap, not
    merged as done.
  • Global load traffic for the mat-vec kernel no worse than one read per packed
    byte per output tile, confirmed by profiler counters rather than by inspection.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions