Repository navigation
[ET-VK] Cooperative-matrix kernels for 4-bit linear on Intel Xe2 and NVIDIA - #24000
Draft
xuyanwen2012 wants to merge 9 commits into
Draft
xuyanwen2012 wants to merge 9 commits into
xuyanwen2012 wants to merge 9 commits into
Conversation
… in vk_api - Pipelines with a required subgroup size are created with VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT_EXT when the device enabled computeFullSubgroups and the local size in X (spec constant 0) is a multiple of the subgroup size. Cooperative-matrix shaders below SPIR-V 1.6 need full subgroups (VUID-RuntimeSpirv-OpTypeCooperativeMatrixKHR-10770); the multiple-of condition keeps the flag valid (VUID 02757). The owned device passes what create_logical_device enabled; a device the embedder created keeps the previous behaviour (off). - Query vulkanMemoryModel when VK_KHR_vulkan_memory_model is advertised and enable it on the logical device, so shaders that declare the VulkanMemoryModel capability can be created. Other devices are unchanged. - Keep the enumerated cooperative-matrix property list instead of reducing it to a bool, query VkPhysicalDeviceDriverProperties (Vulkan 1.2 or VK_KHR_driver_properties), and expose vendor_id(), driver_id(), cooperative_matrix_properties() and max_compute_shared_memory_size() on the Adapter. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm
Adds a selection layer for cooperative-matrix kernels of the 4-bit
quantized linear ops, keyed by architecture family instead of device name.
- CoopmatFamily.{h,cpp} (pure C++, host-testable): CoopmatDeviceKey (vendor
ID, VkDriverId, subgroup size range, cooperative-matrix shapes, shared
memory) and coopmat_family(), which maps a key to Xe2 (Intel, Mesa driver,
fp16 8x16x16, int8 8x16x32, subgroup 16) or NVIDIA (proprietary driver,
fp16 16x16x16 with fp16 and fp32 results, int8 16x16x32, subgroup 32).
Each family has a short list of kernel rows with tile geometry, storage
combinations and an optional shape rule; select_wmma_kernel() returns the
first row that fits, or nothing.
- WmmaLinear.{h,cpp}: the ComputeGraph side (device key from the Adapter,
storage mapping, predicted weight storage, workgroup grid per family tile).
A device that matches no family, or a shape no row of its family accepts,
gets no choice, and the ops keep their existing kernels. Nothing calls this
yet.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm
At the top of q4gsw_linear (et_vk.q4gsw_linear / et_vk.linear_q4gsw), a
device of a cooperative-matrix family builds the op with its own node when a
family row applies to the build-time shape (fp16, no bias, batch 1, tile
aligned, matching storage). The node packs weights with
prepack_quantized_linear_weight and falls back to linear_q4gsw_{tiled,coop}
for shapes no row covers after a resize (decode, unaligned M). Every other
device and shape keeps the existing path unchanged.
Shaders (gemm/linear/qw/q4gsw/):
- linear_q4gsw_wmma: double-buffered fp16 cooperative-matrix GEMM with
in-shader int4 dequantization. Variants: Xe2 t128x128k16g82s16m8flib
(texture3d) and t128x128k16g44s16m8fli (buffer); NVIDIA
t256x128k16g42s32bt and t128x128k16g22s32 (texture3d).
- linear_q4gsw_wmma_bx: texel-wise weight staging with fp32 accumulation,
NVIDIA t128x128k32g42s32f32c (texture3d) for K > 8192, where fp16
accumulation loses accuracy.
Family rows: Xe2 uses the two Xe2 tiles for every aligned shape; NVIDIA uses
bx for K > 8192, bt for K <= 8192 when M is a multiple of 256 (and N a
multiple of 2048 at M = 256), and the 128x128 tile otherwise, only on the
Llama 3.2 1B / 3B and 3.1 8B prefill projections at group size 128 and
M <= 2048, which are the shapes validated on the family's devices.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm
At the top of linear_dq8ca_q4gsw, a device of a cooperative-matrix family
builds the whole op (activation quantize/pack and the int8
cooperative-matrix linear) with its own nodes when a family row applies to
the build-time shape. The chosen row also fixes the int8 activation layout:
4h4w (quantize_and_pack_4h4w_with_group_sums, as today) or row-major
(utils::kPackedInt8_4W, new packer). Shapes no row covers after a resize
use linear_dq8ca_q4gsw_{tiled,coop}; a row-major op keeps its row-major
kernel for unaligned M, and decode reads the fp activations through the
_coop kernel. Devices without int8 dot product, other devices and other
shapes keep the existing path unchanged.
Shaders (gemm/linear/qw/dq8ca_q4gsw/):
- linear_dq8ca_q4gsw_wmma: int8 x int8 -> int32 per quantization group, fp32
total, zero-point correction applied once; texel-wise weight staging and
an optional blocked workgroup order. Xe2 variants t128x128k64g84s16m8i and
...m8n13 (texture3d and buffer).
- linear_dq8ca_q4gsw_wmma_rowmajor: same arithmetic on row-major int8
activations staged as 16-byte vectors, whole-texel weight staging. NVIDIA
variant t128x128k64g24s32mk32 (texture3d).
- quantize_and_pack_4w_with_group_sums: row-major counterpart of
quantize_and_pack_4h4w_with_group_sums (same quantization and group sums).
linear_int8_input_block.glslh gets a SKIP_BLOCK_WRITE_HELPERS guard so the
packer can reuse its quantize helpers with a scalar output buffer.
Family rows: Xe2 uses the n13 order where K * N >= 24e6 and dispatch order
otherwise; NVIDIA uses the row-major kernel on the validated Llama
projections (group size 128, M <= 2048).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm
- test/custom_ops/coopmat_family_test.cpp: host-only gtest of the family key and the row selection (no GPU). Covers the Xe2 and NVIDIA keys, the same silicon under other drivers (Intel Windows, NVK), RADV, a key without the int8 shape or subgroup size, and the llvmpipe signature (vendor 0x10005, MESA_LLVMPIPE, 8x8x8 shapes only, subgroup 8), which must match no family and select nothing for any shape. Checks every Llama projection on both families, the NVIDIA edge shapes, the shapes that keep the existing kernels (unaligned M, decode, bias, fp32, mixed or buffer-weight storage, batch), the int8 layout rules, the shared-memory limit, the kernel-name to tile map and, when the sources are available, that every row has its shader variants. Buck target in targets.bzl; standalone CMake project in test/custom_ops/CMakeLists.txt (gtest only, no Vulkan backend). - etvk_bench q4gsw_linear:wmma_correctness: both ops, texture3d and buffer, group size 128, shapes aligned to every family tile that select each Xe2 and NVIDIA row, zero-mean data with exact dynamic quantization and varying zero points, and a threaded reference without the size limit of bench_reference. On a device without a family the same cases check the existing kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm
A family-built linear_q4gsw op fell back to linear_q4gsw_tiled for any M that is not a multiple of the row's tile, so on a family device an unaligned prompt ran slower than upstream's current 4w kernels (Llama 3.2 1B on an Arc B580, 511 tokens: -16 %). The texture-IO kernels (linear_q4gsw_wmma, linear_q4gsw_wmma_bx) now take a partial last M tile: rows of the tile at or past M read activation row M - 1 (each output row depends only on its own activation row) and are not stored. Selection treats an M of at least kWmmaPartialTileM (128) as M rounded up to a multiple of 128, so an unaligned M gets the row of the next aligned M. Smaller M, buffer IO and linear_dq8ca_q4gsw keep the alignment requirement and the existing fallback. Arc B580, Llama 3.2 1B 4w prefill tok/s (upstream / before / after, one run each): 511 tokens 3024 / 2542 / 10429, 300 tokens 2655 / 2174 / 8824, 129 tokens 1500 / 1817 / 5375, 512 tokens 3066 / 10449 / 10667. Tests: coopmat_family_test PartialMTile; wmma_correctness adds linear_q4gsw at M = 511, 300, 129 and 127. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
linear_dq8ca_q4gsw_wmma_rowmajor keeps serving an unaligned M after a resize (the grid rounds M up), and read the int8 activations, their scales and their zero points of whole 128-row tiles, past the M rows of those tensors. Rows at or past M now read row M - 1 and are not stored, so the kernel stays inside its inputs and output without padding the temporary activation buffer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…2 8da4w threshold 20e6 NVIDIA: the rows no longer require one of the Llama 3.2 1B / 3B and Llama 3.1 8B projections, group size 128 and M <= 2048; the tile alignment, K > 8192 for fp32 accumulation and the M = 256 rule remain. On an RTX 4070 Ti SUPER and a Jetson Orin, 30 projection shapes of Qwen2.5 0.5B / 1.5B, Gemma 2 2B, Qwen3 4B / 8B and Phi-3 mini at M = 512 and 2048, group sizes 64 and 256 and M = 3072 / 4096 pass against the CPU reference and run 2.2x to 13.5x faster than the existing kernels. Xe2: the blocked 8da4w workgroup order now starts at K * N >= 20e6 (was 24e6), between the largest shape where it is slower (16.8e6) and the smallest where it is faster (21.2e6) on both the Arc B580 and the Arc Pro B70. No Llama projection changes row. wmma_correctness gains Qwen2.5 0.5B, Gemma 2 2B and group size 64 cases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
cmake-format reflows plain comment lines; fence the commands with ~~~ as backends/vulkan/CMakeLists.txt does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Today the cooperative-matrix path of
q4gswanddq8ca_q4gswlinear only runs on devices with subgroup size 64whose name contains "Radeon" or "Xclipse". On Intel Xe2 (subgroup 16) and NVIDIA (subgroup 32) no
cooperative-matrix linear kernel runs, and prefill uses the tiled kernels.
This PR adds a cooperative-matrix path for 4-bit weight linear (
q4gsw, 4w) and 8-bit-activation 4-bit-weightlinear (
dq8ca_q4gsw, 8da4w) on two architecture families:A device matches a family by vendor ID, driver ID, the cooperative-matrix shapes it reports and the subgroup
sizes it supports, never by device name. A device that matches no family keeps exactly the kernels it runs
today.
What changes
vk_api: setREQUIRE_FULL_SUBGROUPSon cooperative-matrix pipelines below SPIR-V 1.6 when the local sizein X is a multiple of the required subgroup size (VUID-RuntimeSpirv-OpTypeCooperativeMatrixKHR-10770), enable
vulkanMemoryModelthroughVK_KHR_vulkan_memory_modelwhen the device advertises it, and expose driver ID(
VkPhysicalDeviceDriverProperties), the cooperative-matrix property list and the maximum compute sharedmemory.
impl/gemm/linear/qw/CoopmatFamily.{h,cpp}: the family key and one table of kernel rows per family. Each rownames a shader variant, its tile geometry and an optional shape rule. Selection checks op, layout, subgroup
size, shared-memory limit and tile alignment; a shape that no row fits falls back to the existing kernel.
q4gsw/Q4gswLinearWmma.*,dq8ca_q4gsw/Dq8caQ4gswLinearWmma.*,WmmaLinear.*: dispatch, hooked intoQ4gswLinear.cppandDq8caQ4gswLinear.cpp(7 lines each).glsl/gemm/linear/qw/: four cooperative-matrix linear templates(
linear_q4gsw_wmma,linear_q4gsw_wmma_bx,linear_dq8ca_q4gsw_wmma,linear_dq8ca_q4gsw_wmma_rowmajor)with shared
*_body.glslhfiles, 10 variants in total, plusquantize_and_pack_4w_with_group_sums, arow-major counterpart of the existing 4H4W activation packer.
vulkan_coopmat_family_test(host-only gtest of the key and the selection, including that an unmatchedkey, e.g. llvmpipe, gets no family) and an
etvk_benchsetq4gsw_linear:wmma_correctness.t128x128k16g82s16m8flib(texture3d),t128x128k16g44s16m8fli(buffer)t128x128k64g84s16m8i; blocked workgroup order...n13when K*N >= 20e6t256x128k16g42s32btfor K <= 8192;t128x128k16g22s32where the 256-row tile does not fit;bx_t128x128k32g42s32f32c(fp32 accumulation) for K > 8192rowmajor_t128x128k64g24s32mk32The kernel sets were chosen from per-device tuned kernels by measuring each candidate on every member of its
family: the accepted set is within 3 % geomean of each member's own best set, with no projection more than
8 % worse. The K > 8192 fp32-accumulation rule on NVIDIA is a correctness rule: the fp16-accumulate kernel exceeds
tolerance at K = 14336 (Llama 3.1 8B down projection).
Results
End-to-end prefill, Llama 3.2 1B, 512-token prompt, tokens/s (median of 3), this PR vs
main:Unaligned prompts (e.g. 511, 300, 129 tokens) take the family kernels through a partial last M tile; on the
4070 Ti, 1B 4w at 511 tokens goes from 6636 to 21292 tok/s. Outside the Llama projections, 30 projection shapes
of Qwen2.5 0.5B / 1.5B, Gemma 2 2B, Qwen3 4B / 8B and Phi-3 mini (M = 512 and 2048, group sizes 64, 128 and 256)
pass against the CPU reference on the 4070 Ti and the Orin and run 2.2x to 13.5x faster than the current kernels
in
etvk_bench. End-to-end numbers were measured before the last commit (selection on any aligned NVIDIA shape,Xe2 threshold 20e6), which leaves every Llama projection on the same kernel.
Testing
On this branch, rebased on
main43a7bc20bc:wmma_correctnessq4gsw_linear:correctnesscoopmat_correctnessmainmainThe 26 skips are the existing fp16 skips of the correctness set and are the same on
main.vulkan_coopmat_family_test: 11/11.Known limitations and open questions
matches any NVIDIA GPU on the proprietary driver that reports these shapes; it was measured on Ada and Ampere
only (Turing, Hopper and Blackwell would match unmeasured). Happy to narrow it if you prefer.
it is slower (16.8e6) and the smallest where it is faster (21.2e6) on both Xe2 cards.
the case on
mainas well and is not changed here.q4gsw_linear:wmma_correctnessbe added to the GPU job in.github/workflows/vulkan.yml? Itdoes not exercise the new kernels on llvmpipe (no family there).
etvk_benchon the 4070 Ti (driver 615.71) aborts at exit in about 1 of 10 runs, after all cases have passed:ComputePipelineCache::purgein~Adapterfrees a pipeline while the driver's background compiler thread isstill running.
mainshows the same rate (3 of 40 vs 4 of 40 for this branch).Not in this PR: fused attention, decode, fp16 linear.
🤖 Generated with Claude Code