Skip to content

Commit bfcb14a

Browse files
xuyanwen2012claude
andcommitted
[ET-VK] Consolidate quantized coopmat shaders into two templates
The four quantized-linear coopmat shaders share one double-buffered skeleton since the previous change, so fold them into one template per MMA family, mirroring how coopmat_mm.glsl hosts its three variants: linear_qw_coopmat.glsl -> linear_q4gsw_coopmat_* (WEIGHT_NBITS=4) linear_q8csw_coopmat_* (WEIGHT_NBITS=8) linear_dq8ca_qw_coopmat.glsl -> linear_dq8ca_q4gsw_coopmat_* (4) linear_dq8ca_q8csw_coopmat_* (8) Generated variant names are unchanged, so dispatch code needs no changes. The WEIGHT_NBITS conditionals cover only the weight unpack and the quant bookkeeping: per-channel INT8 is expressed as the single-group special case of the grouped INT4 path (num_groups = 1), which collapses the nested groups-by-chunks loop, the wsum/wsc group-parity ping-pong, and the group epilog to exactly the previous per-channel control flow. Regression on ERD9975: all coopmat correctness cases pass and the M=1024 Llama-shape numbers match the pre-refactor run (4w 2.2x, 8da4w 4.0-4.4x, 8w 3.3x, 8da8w 7.0-8.0x vs tiled). Authored with Claude. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent b891057 commit bfcb14a

8 files changed

Lines changed: 306 additions & 971 deletions

backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q8csw_coopmat.glsl

Lines changed: 0 additions & 435 deletions
This file was deleted.

backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q8csw_coopmat.yaml

Lines changed: 0 additions & 30 deletions
This file was deleted.

backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl

Lines changed: 143 additions & 54 deletions
Large diffs are not rendered by default.

backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat.yaml renamed to backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,24 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
# coopmat<int8> x coopmat<int8> -> coopmat<int32> variant of
8-
# linear_dq8ca_q4gsw_tiled. Requires VK_COMPONENT_TYPE_SINT8_KHR cooperative
9-
# matrix property to be enumerated on the device (e.g. Radeon 780M / Mesa RADV
10-
# exposes int8 16x16x16 Subgroup).
7+
# coopmat<int8> x coopmat<int8> -> coopmat<int32> variants of the
8+
# dynamically-quantized-activation linear tiled shaders. One template, two
9+
# weight formats:
10+
# WEIGHT_NBITS=4 -> linear_dq8ca_q4gsw_coopmat (INT4 group-symmetric)
11+
# WEIGHT_NBITS=8 -> linear_dq8ca_q8csw_coopmat (INT8 per-channel symmetric)
12+
# Requires the VK_COMPONENT_TYPE_SINT8_KHR cooperative matrix property to be
13+
# enumerated on the device.
1114
# Loop structure follows the double-buffered reference (gemm_double_buf) at
12-
# a 128x64 tile with K-step 16, 4 subgroups x 64 threads. The reference's
15+
# a 128x64 tile with K-step 32, 4 subgroups x 64 threads. The reference's
1316
# subgroup-32 layout is NOT used — the Xclipse PAL compiler crashes in
1417
# vkCreateComputePipelines when int8 WMMA is compiled at forced subgroup
15-
# size 32 (fp16 WMMA at 32 is fine; see linear_q4gsw_coopmat).
18+
# size 32 (fp16 WMMA at 32 is fine; see linear_qw_coopmat).
1619

17-
linear_dq8ca_q4gsw_coopmat:
20+
linear_dq8ca_qw_coopmat:
1821
parameter_names_with_default_values:
1922
PRECISION: highp
2023
HAS_BIAS: false
24+
WEIGHT_NBITS: 4
2125
WEIGHT_STORAGE: texture2d
2226
MMA_M: 16
2327
MMA_N: 16
@@ -30,6 +34,14 @@ linear_dq8ca_q4gsw_coopmat:
3034
SUBGROUP_SIZE: 64
3135
shader_variants:
3236
- NAME: linear_dq8ca_q4gsw_coopmat_buffer_texture2d_half
37+
WEIGHT_NBITS: 4
3338
WEIGHT_STORAGE: texture2d
3439
- NAME: linear_dq8ca_q4gsw_coopmat_buffer_buffer_half
40+
WEIGHT_NBITS: 4
41+
WEIGHT_STORAGE: buffer
42+
- NAME: linear_dq8ca_q8csw_coopmat_buffer_texture2d_half
43+
WEIGHT_NBITS: 8
44+
WEIGHT_STORAGE: texture2d
45+
- NAME: linear_dq8ca_q8csw_coopmat_buffer_buffer_half
46+
WEIGHT_NBITS: 8
3547
WEIGHT_STORAGE: buffer

0 commit comments

Comments
 (0)