Skip to content

ROCm KQuantGemmKWmmaQ4K/Q6K: redesign around a cooperative shared tile, matching llama.cpp's mul_mat_q mechanism #3034

Description

@joral

Row: KERNEL-QUANT-CIQ-GEMM-ROCM

Row

KERNEL-QUANT-CIQ-GEMM-ROCM — already ACTIVE. Follow-on to #3032/#3033
(wider-block A/B, measured and rejected, geomean -4.3%). This issue is the
sharper hypothesis that measurement pointed at.

Confirmed mechanism (read, not inferred)

KQuantGemmKWmmaQ4K/Q6K give each WARP an independent 16x16 output tile:
its own shared-memory slice, its own load_matrix_sync calls for both the
weight and activation fragments, zero sharing with any other warp in the
same block — even when two warps share the same M-tile (it) and are
reading the identical 16 activation rows.

llama.cpp's mul_mat_q does not do this. At this row's shapes
(ggml/src/ggml-cuda/mmq-config-rdna4.cuh:127,
GGML_TYPE_Q4_K, 256 threads, I=128, J=128), ONE BLOCK owns a 128x128
output region:

  • ggml_cuda_mmq_load_tiles_q4_K (ggml/src/ggml-cuda/mmq-load-tiles.cuh:703-741)
    cooperatively loads and dequantizes the I=128 weight rows ONCE, striped
    across every warp's threads (i0 += nrows*nwarps) — no warp repeats work
    another warp already did.
  • mmq_get_nbytes_shared (ggml/src/ggml-cuda/mmq.cuh:1379-1383) sizes the
    activation-tile shared allocation by J alone, not J*nwarps — one
    shared copy of the activation tile serves the whole block too.
  • ggml_cuda_mmq_write_back_mma (ggml/src/ggml-cuda/mmq.cuh:476-500)
    splits the I rows across warps (rows_per_warp = I/nwarps); each warp
    loops the full J range doing 16x16 MMA against that one shared load.

Our kernel's per-warp-independent design is the actual gap the rejected
wider-block experiment could not reach, because widening only added more
independent warps to the same per-warp-load pattern — it never introduced
sharing.

The work

Redesign KQuantGemmKWmmaQ4K/KQuantGemmKWmmaQ6K so one BLOCK cooperatively
loads and dequantizes a bigger shared I x J tile (weight rows x activation
rows), with warps splitting the I dimension for compute and looping the
full J range per warp — matching the mechanism above, not necessarily
llama.cpp's exact bucket sizes or its multi-bucket (64x16 .. 128x128)
dispatch table on the first pass. A spec proposing a concrete, scoped tile
size (fitting this card's 64 KiB LDS budget, reusing the existing
per-superblock staging/epilogue shape where possible) lands before any
kernel code, per this row's own recorded PR-shape preference.

Named risk: llama.cpp's own multi-bucket dispatch exists because a single
big I/J tile wastes work on shapes smaller than it (our M is prompt length,
essentially never aligned to a large tile; our benchmarked N shapes happen
to already be multiples of 128, so N is likely fine with one bucket, M is
the harder axis). The M/N tail-fill mechanism this row already has absorbs
the remainder either way; whether one bucket size is enough or a second,
smaller bucket is needed is a measured decision, not assumed up front.

Confirmed capped at RDNA4 (gfx1200/gfx1201) by the same two existing gates
this row's kernels already use — #if defined(__gfx1200__) || defined(__gfx1201__) at compile time, Gfx12QuantWmmaHostOk at runtime —
so this redesign has no reach into CDNA, RDNA3, or RDNA3.5 as long as it
stays inside those same guards, which is also the only place rocwmma.hpp
is includable at all.

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

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions