From 73969529012464c5744defa1dc87299430b02912 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 16 Oct 2026 11:30:41 -0700 Subject: [PATCH 01/11] [ET-VK] Full subgroups, Vulkan memory model and family-key properties 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 Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm --- backends/vulkan/runtime/vk_api/Adapter.cpp | 29 +++++++++- backends/vulkan/runtime/vk_api/Adapter.h | 28 ++++++++++ backends/vulkan/runtime/vk_api/Device.cpp | 59 ++++++++++++++++++--- backends/vulkan/runtime/vk_api/Device.h | 8 +++ backends/vulkan/runtime/vk_api/Pipeline.cpp | 37 +++++++++++-- backends/vulkan/runtime/vk_api/Pipeline.h | 9 +++- 6 files changed, 155 insertions(+), 15 deletions(-) diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index eea61a8fe3f..8d4cbba72ea 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -134,6 +135,13 @@ VkDevice create_logical_device( #endif /* VK_EXT_subgroup_size_control */ }; +#ifdef VK_KHR_vulkan_memory_model + if (physical_device.vulkan_memory_model_features.vulkanMemoryModel) { + requested_device_extensions.push_back( + VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME); + } +#endif /* VK_KHR_vulkan_memory_model */ + std::vector enabled_device_extensions; find_requested_device_extensions( physical_device.handle, @@ -170,6 +178,20 @@ VkDevice create_logical_device( void* extension_list_top = nullptr; +#ifdef VK_KHR_vulkan_memory_model + VkPhysicalDeviceVulkanMemoryModelFeaturesKHR vulkan_memory_model_features{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_MEMORY_MODEL_FEATURES_KHR}; + for (const char* extension : enabled_device_extensions) { + if (std::strcmp(extension, VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME) == + 0) { + vulkan_memory_model_features.vulkanMemoryModel = VK_TRUE; + vulkan_memory_model_features.pNext = extension_list_top; + extension_list_top = &vulkan_memory_model_features; + break; + } + } +#endif /* VK_KHR_vulkan_memory_model */ + #ifdef VK_KHR_16bit_storage VkPhysicalDevice16BitStorageFeatures shader_16bit_storage{ physical_device.shader_16bit_storage}; @@ -300,7 +322,12 @@ Adapter::Adapter( shader_layout_cache_(device_.handle), shader_cache_(device_.handle), pipeline_layout_cache_(device_.handle), - compute_pipeline_cache_(device_.handle, cache_data_path), + compute_pipeline_cache_( + device_.handle, + cache_data_path, + // Full subgroups: what create_logical_device enabled + supports_subgroup_size_control() && + supports_compute_full_subgroups()), sampler_cache_(device_.handle), vma_(instance_, physical_device_.handle, device_.handle), linear_tiling_3d_enabled_{test_linear_tiling_3d_image_support( diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index 0c2784ae473..933b16ac70c 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -146,6 +146,16 @@ class Adapter final { return physical_device_.device_type; } + // VkPhysicalDeviceProperties::vendorID. + inline uint32_t vendor_id() const { + return physical_device_.properties.vendorID; + } + + // VkDriverId of the driver, or 0 when the device does not report it. + inline uint32_t driver_id() const { + return physical_device_.driver_id; + } + // Driver-reported physical device type. Use device_type() above for // vendor-specific quirks; this for capability-tier decisions. inline bool is_integrated_gpu() const { @@ -326,6 +336,20 @@ class Adapter final { #endif /* VK_KHR_cooperative_matrix */ } +#ifdef VK_KHR_cooperative_matrix + // Every cooperative-matrix configuration the device enumerates (empty when + // cooperative matrices are unsupported). + inline const std::vector& + cooperative_matrix_properties() const { +#ifdef ETVK_FORCE_NO_EXTENSIONS + static const std::vector kNone; + return kNone; +#else + return physical_device_.cooperative_matrix_properties; +#endif + } +#endif /* VK_KHR_cooperative_matrix */ + inline bool supports_int16_shader_types() { #ifdef ETVK_FORCE_NO_EXTENSIONS return false; @@ -491,6 +515,10 @@ class Adapter final { return physical_device_.properties.limits.maxComputeWorkGroupInvocations; } + inline uint32_t max_compute_shared_memory_size() const { + return physical_device_.properties.limits.maxComputeSharedMemorySize; + } + inline uint32_t recommended_lwg_nthreads() const { return 64u; } diff --git a/backends/vulkan/runtime/vk_api/Device.cpp b/backends/vulkan/runtime/vk_api/Device.cpp index 56902f05cd8..169d9b58c70 100644 --- a/backends/vulkan/runtime/vk_api/Device.cpp +++ b/backends/vulkan/runtime/vk_api/Device.cpp @@ -47,6 +47,10 @@ PhysicalDevice::PhysicalDevice( handle(physical_device_handle), properties{}, memory_properties{}, +#ifdef VK_KHR_vulkan_memory_model + vulkan_memory_model_features{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_MEMORY_MODEL_FEATURES_KHR}, +#endif /* VK_KHR_vulkan_memory_model */ #ifdef VK_KHR_16bit_storage shader_16bit_storage{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES, @@ -74,6 +78,7 @@ PhysicalDevice::PhysicalDevice( cooperative_matrix_features{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}, supports_int8_coopmat{false}, + cooperative_matrix_properties{}, #endif /* VK_KHR_cooperative_matrix */ #ifdef VK_NV_cooperative_matrix2 cooperative_matrix2_features{ @@ -107,7 +112,8 @@ PhysicalDevice::PhysicalDevice( supports_subgroup_size_control(false), supports_compute_full_subgroups(false), device_name{}, - device_type{DeviceType::UNKNOWN} { + device_type{DeviceType::UNKNOWN}, + driver_id(0) { // Extract physical device properties vkGetPhysicalDeviceProperties(handle, &properties); @@ -262,6 +268,19 @@ void PhysicalDevice::query_extensions_vk_1_1() { // Create linked list to query availability of extensions void* extension_list_top = nullptr; +#ifdef VK_KHR_vulkan_memory_model + // Runtime requests Vulkan 1.1, so use the advertised extension even on 1.2+. + std::vector memory_model_extensions; + find_requested_device_extensions( + handle, + memory_model_extensions, + {VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME}); + if (!memory_model_extensions.empty()) { + vulkan_memory_model_features.pNext = extension_list_top; + extension_list_top = &vulkan_memory_model_features; + } +#endif /* VK_KHR_vulkan_memory_model */ + #ifdef VK_KHR_16bit_storage shader_16bit_storage.pNext = extension_list_top; extension_list_top = &shader_16bit_storage; @@ -340,13 +359,15 @@ void PhysicalDevice::query_extensions_vk_1_1() { p.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; p.pNext = nullptr; } - vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR_fn( - handle, &count, props.data()); - for (const auto& p : props) { - if (p.AType == VK_COMPONENT_TYPE_SINT8_KHR) { - supports_int8_coopmat = true; - break; + if (vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR_fn( + handle, &count, props.data()) == VK_SUCCESS) { + props.resize(count); + for (const auto& p : props) { + if (p.AType == VK_COMPONENT_TYPE_SINT8_KHR) { + supports_int8_coopmat = true; + } } + cooperative_matrix_properties = std::move(props); } } } @@ -374,10 +395,34 @@ void PhysicalDevice::query_extensions_vk_1_1() { properties_list_top = &subgroup_size_control_properties; #endif /* VK_EXT_subgroup_size_control */ + // VkPhysicalDeviceDriverProperties is core in Vulkan 1.2; on a 1.1 device it + // needs VK_KHR_driver_properties. + VkPhysicalDeviceDriverProperties driver_properties{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DRIVER_PROPERTIES}; + bool query_driver_properties = properties.apiVersion >= VK_API_VERSION_1_2; +#ifdef VK_KHR_driver_properties + if (!query_driver_properties) { + std::vector driver_properties_extensions; + find_requested_device_extensions( + handle, + driver_properties_extensions, + {VK_KHR_DRIVER_PROPERTIES_EXTENSION_NAME}); + query_driver_properties = !driver_properties_extensions.empty(); + } +#endif /* VK_KHR_driver_properties */ + if (query_driver_properties) { + driver_properties.pNext = properties_list_top; + properties_list_top = &driver_properties; + } + properties2.pNext = properties_list_top; vkGetPhysicalDeviceProperties2(handle, &properties2); + if (query_driver_properties) { + driver_id = static_cast(driver_properties.driverID); + } + subgroup_size = subgroup_properties.subgroupSize; supported_subgroup_ops = subgroup_properties.supportedOperations; supported_subgroup_stages = subgroup_properties.supportedStages; diff --git a/backends/vulkan/runtime/vk_api/Device.h b/backends/vulkan/runtime/vk_api/Device.h index 587d3965c3e..f530c913065 100644 --- a/backends/vulkan/runtime/vk_api/Device.h +++ b/backends/vulkan/runtime/vk_api/Device.h @@ -36,6 +36,9 @@ struct PhysicalDevice final { VkPhysicalDeviceMemoryProperties memory_properties; // Additional features available from extensions +#ifdef VK_KHR_vulkan_memory_model + VkPhysicalDeviceVulkanMemoryModelFeaturesKHR vulkan_memory_model_features; +#endif /* VK_KHR_vulkan_memory_model */ #ifdef VK_KHR_16bit_storage VkPhysicalDevice16BitStorageFeatures shader_16bit_storage; #endif /* VK_KHR_16bit_storage */ @@ -57,6 +60,8 @@ struct PhysicalDevice final { // True when VK_COMPONENT_TYPE_SINT8_KHR appears in the enumerated coopmat // property list — required for coopmat shaders (e.g. dq8ca_q4gsw). bool supports_int8_coopmat; + // Every entry of vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR. + std::vector cooperative_matrix_properties; #endif /* VK_KHR_cooperative_matrix */ #ifdef VK_NV_cooperative_matrix2 @@ -103,6 +108,9 @@ struct PhysicalDevice final { // Device identity std::string device_name; DeviceType device_type; + // VkDriverId from VkPhysicalDeviceDriverProperties (Vulkan 1.2 or + // VK_KHR_driver_properties); 0 when the device does not report it. + uint32_t driver_id; explicit PhysicalDevice(VkInstance instance, VkPhysicalDevice); diff --git a/backends/vulkan/runtime/vk_api/Pipeline.cpp b/backends/vulkan/runtime/vk_api/Pipeline.cpp index 2b42e25c92e..bf469a49c15 100644 --- a/backends/vulkan/runtime/vk_api/Pipeline.cpp +++ b/backends/vulkan/runtime/vk_api/Pipeline.cpp @@ -266,6 +266,28 @@ void swap(PipelineLayout& lhs, PipelineLayout& rhs) noexcept { rhs.handle_ = tmp_handle; } +namespace { + +// Cooperative-matrix shaders below SPIR-V 1.6 need full subgroups +// (VUID-RuntimeSpirv-OpTypeCooperativeMatrixKHR-10770). The flag is only valid +// when the local size in X, specialization constant 0, is a multiple of the +// required subgroup size (VUID-VkPipelineShaderStageCreateInfo-pNext-02757). +VkPipelineShaderStageCreateFlags full_subgroups_flag( + const ComputePipeline::Descriptor& descriptor, + const bool full_subgroups) { +#ifdef VK_EXT_subgroup_size_control + const SpecVarList& spec_vars = descriptor.specialization_constants; + const uint32_t subgroup_size = descriptor.required_subgroup_size; + if (full_subgroups && subgroup_size > 0u && spec_vars.size() > 0u && + spec_vars.at(0).value.as_uint32 % subgroup_size == 0u) { + return VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT_EXT; + } +#endif /* VK_EXT_subgroup_size_control */ + return 0u; +} + +} // namespace + // // ComputePipeline // @@ -276,7 +298,8 @@ ComputePipeline::ComputePipeline(VkDevice device, VkPipeline handle) ComputePipeline::ComputePipeline( VkDevice device, const ComputePipeline::Descriptor& descriptor, - VkPipelineCache pipeline_cache) + VkPipelineCache pipeline_cache, + const bool full_subgroups) : device_(device), handle_{VK_NULL_HANDLE} { map_entries_ = descriptor.specialization_constants.generate_map_entries(); @@ -302,7 +325,7 @@ ComputePipeline::ComputePipeline( const VkPipelineShaderStageCreateInfo shader_stage_create_info{ VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, // sType shader_stage_pnext, // pNext - 0u, // flags + full_subgroups_flag(descriptor, full_subgroups), // flags VK_SHADER_STAGE_COMPUTE_BIT, // stage descriptor.shader_module, // module "main", // pName @@ -417,9 +440,11 @@ void PipelineLayoutCache::purge() { ComputePipelineCache::ComputePipelineCache( VkDevice device, - const std::string& cache_data_path) + const std::string& cache_data_path, + const bool full_subgroups) : cache_mutex_{}, device_(device), + full_subgroups_(full_subgroups), pipeline_cache_{VK_NULL_HANDLE}, cache_{}, cache_data_path_(cache_data_path) { @@ -443,6 +468,7 @@ ComputePipelineCache::ComputePipelineCache( ComputePipelineCache&& other) noexcept : cache_mutex_{}, device_(other.device_), + full_subgroups_(other.full_subgroups_), pipeline_cache_(other.pipeline_cache_), cache_(std::move(other.cache_)) { std::lock_guard lock(other.cache_mutex_); @@ -537,7 +563,7 @@ void ComputePipelineCache::create_pipelines( shader_stage_create_infos.push_back(VkPipelineShaderStageCreateInfo{ VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, // sType shader_stage_pnext, // pNext - 0u, // flags + full_subgroups_flag(key, full_subgroups_), // flags VK_SHADER_STAGE_COMPUTE_BIT, // stage key.shader_module, // module "main", // pName @@ -579,7 +605,8 @@ VkPipeline ComputePipelineCache::retrieve( it = cache_ .insert( {key, - ComputePipelineCache::Value(device_, key, pipeline_cache_)}) + ComputePipelineCache::Value( + device_, key, pipeline_cache_, full_subgroups_)}) .first; } return it->second.handle(); diff --git a/backends/vulkan/runtime/vk_api/Pipeline.h b/backends/vulkan/runtime/vk_api/Pipeline.h index 5e286ad77d7..fa37387a28c 100644 --- a/backends/vulkan/runtime/vk_api/Pipeline.h +++ b/backends/vulkan/runtime/vk_api/Pipeline.h @@ -172,7 +172,8 @@ class ComputePipeline final { explicit ComputePipeline( VkDevice device, const Descriptor& descriptor, - VkPipelineCache pipeline_cache); + VkPipelineCache pipeline_cache, + const bool full_subgroups = false); ComputePipeline(const ComputePipeline&) = delete; ComputePipeline& operator=(const ComputePipeline&) = delete; @@ -244,7 +245,8 @@ class ComputePipelineCache final { public: explicit ComputePipelineCache( VkDevice device, - const std::string& cache_data_path); + const std::string& cache_data_path, + const bool full_subgroups = false); ComputePipelineCache(const ComputePipelineCache&) = delete; ComputePipelineCache& operator=(const ComputePipelineCache&) = delete; @@ -306,6 +308,9 @@ class ComputePipelineCache final { std::mutex cache_mutex_; VkDevice device_; + // computeFullSubgroups of the device: pipelines with a required subgroup + // size are created with REQUIRE_FULL_SUBGROUPS. + bool full_subgroups_; VkPipelineCache pipeline_cache_; std::unordered_map cache_; const std::string cache_data_path_; From 0582c1cfb86aebba42b80310a205a22d6da7a006 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 16 Oct 2026 11:30:55 -0700 Subject: [PATCH 02/11] [ET-VK] Cooperative-matrix family key and kernel rows for 4-bit linear 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 Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm --- .../ops/impl/gemm/linear/qw/CoopmatFamily.cpp | 278 ++++++++++++++++++ .../ops/impl/gemm/linear/qw/CoopmatFamily.h | 175 +++++++++++ .../ops/impl/gemm/linear/qw/WmmaLinear.cpp | 118 ++++++++ .../ops/impl/gemm/linear/qw/WmmaLinear.h | 39 +++ 4 files changed, 610 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp new file mode 100644 index 00000000000..ed1e58399e2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp @@ -0,0 +1,278 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +namespace vkcompute { + +namespace { + +bool has_shape( + const CoopmatDeviceKey& device, + const uint32_t m, + const uint32_t n, + const uint32_t k, + const CoopmatComponentType ab, + const CoopmatComponentType result) { + for (const CoopmatShape& s : device.coopmat_shapes) { + if (s.m == m && s.n == n && s.k == k && s.ab == ab && s.result == result) { + return true; + } + } + return false; +} + +bool subgroup_size_available( + const CoopmatDeviceKey& device, + const uint32_t size) { + return device.subgroup_size == size || + (device.subgroup_size_control && size >= device.min_subgroup_size && + size <= device.max_subgroup_size); +} + +// +// Intel Xe2 rows. MMA 8x16x16 / 8x16x32 (the only fp16 / int8 cooperative +// matrix shapes Xe2 exposes), subgroup 16. +// + +// Weight-heavy projections run their workgroups in blocks of N tiles. +bool xe2_weight_heavy(const WmmaShape& s) { + return s.K * s.N >= 24000000; +} + +const std::vector& xe2_rows() { + static const std::vector rows = { + // 4w, texture3d: banded drain, 8 x 2 subgroups. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g82s16m8flib", + {128, 128, 16, 8, 2, 16, 8, false}, + kWmmaTex3dTex2d, + nullptr}, + // 4w, buffer IO: 4 x 4 subgroups. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g44s16m8fli", + {128, 128, 16, 4, 4, 16, 8, false}, + kWmmaBufTex2d, + nullptr}, + // 8da4w: texel-wise B staging and a full A thread map; blocked + // workgroup order on weight-heavy projections, dispatch order elsewhere. + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13", + {128, 128, 64, 8, 4, 16, 8, false}, + kWmmaTex3dTex2d | kWmmaBufTex2d, + xe2_weight_heavy}, + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i", + {128, 128, 64, 8, 4, 16, 8, false}, + kWmmaTex3dTex2d | kWmmaBufTex2d, + nullptr}, + }; + return rows; +} + +// +// NVIDIA rows. MMA 16x16x16 (fp16) / 16x16x32 (int8), subgroup 32, +// texture3d IO only. +// + +// Only the prefill projections of Llama 3.2 1B / 3B and Llama 3.1 8B at group +// size 128 and M <= 2048 were validated on the family's members; every other +// shape keeps the existing kernels. +bool nv_measured(const WmmaShape& s) { + const int64_t K = s.K; + const int64_t N = s.N; + const bool projection = (K == 2048 && (N == 512 || N == 2048 || N == 8192)) || + (K == 3072 && (N == 1024 || N == 3072 || N == 8192)) || + (K == 4096 && (N == 1024 || N == 4096 || N == 14336)) || + (K == 8192 && (N == 2048 || N == 3072)) || (K == 14336 && N == 4096); + return projection && s.group_size == 128 && s.M <= 2048; +} +// fp16 accumulation loses accuracy beyond K = 8192: fp32 accumulation there. +bool nv_large_k(const WmmaShape& s) { + return nv_measured(s) && s.K > 8192; +} +bool nv_256(const WmmaShape& s) { + return nv_measured(s) && s.K <= 8192 && (s.M != 256 || s.N % 2048 == 0); +} +bool nv_default(const WmmaShape& s) { + return nv_measured(s) && s.K <= 8192; +} + +const std::vector& nvidia_rows() { + static const std::vector rows = { + // 4w, K > 8192: fp32 accumulation, texel-wise B staging, drain in the + // A tile. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_bx_t128x128k32g42s32f32c", + {128, 128, 32, 4, 2, 32, 16, true}, + kWmmaTex3dTex2d, + nv_large_k}, + // 4w, K <= 8192: N-major B in shared memory. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t256x128k16g42s32bt", + {256, 128, 16, 4, 2, 32, 16, false}, + kWmmaTex3dTex2d, + nv_256}, + // 4w, the remaining validated shapes (M not a multiple of 256, and + // M = 256 at N % 2048 != 0). + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g22s32", + {128, 128, 16, 2, 2, 32, 16, false}, + kWmmaTex3dTex2d, + nv_default}, + // 8da4w: row-major int8 activations, 2 x 4 subgroups. + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_rowmajor_t128x128k64g24s32mk32", + {128, 128, 64, 2, 4, 32, 16, true}, + kWmmaTex3dTex2d, + nv_measured, + /*rowmajor_a=*/true}, + }; + return rows; +} + +} // namespace + +CoopmatFamily coopmat_family(const CoopmatDeviceKey& device) { + if (!device.coopmat || !device.int8_coopmat) { + return CoopmatFamily::kNone; + } + using T = CoopmatComponentType; + if (device.vendor_id == kVendorIdIntel && + device.driver_id == kDriverIdIntelMesa && + has_shape(device, 8, 16, 16, T::kFp16, T::kFp16) && + has_shape(device, 8, 16, 32, T::kSint8, T::kSint32) && + subgroup_size_available(device, 16)) { + return CoopmatFamily::kXe2; + } + if (device.vendor_id == kVendorIdNvidia && + device.driver_id == kDriverIdNvidiaProprietary && + has_shape(device, 16, 16, 16, T::kFp16, T::kFp16) && + has_shape(device, 16, 16, 16, T::kFp16, T::kFp32) && + has_shape(device, 16, 16, 32, T::kSint8, T::kSint32) && + subgroup_size_available(device, 32)) { + return CoopmatFamily::kNvidia; + } + return CoopmatFamily::kNone; +} + +const char* coopmat_family_name(const CoopmatFamily family) { + switch (family) { + case CoopmatFamily::kXe2: + return "xe2"; + case CoopmatFamily::kNvidia: + return "nvidia"; + case CoopmatFamily::kNone: + break; + } + return "none"; +} + +const std::vector& coopmat_family_rows(const CoopmatFamily family) { + static const std::vector kNoRows; + switch (family) { + case CoopmatFamily::kXe2: + return xe2_rows(); + case CoopmatFamily::kNvidia: + return nvidia_rows(); + case CoopmatFamily::kNone: + break; + } + return kNoRows; +} + +bool wmma_row_fits( + const CoopmatDeviceKey& device, + const WmmaShape& shape, + const WmmaRow& row) { + const WmmaTile& tile = row.tile; + if (shape.op != row.op) { + return false; + } + // The kernels are built without bias and read no leading dim. + if (shape.has_bias || shape.gemv || shape.batch != 1 || !shape.half) { + return false; + } + if (!device.coopmat) { + return false; + } + if (shape.op == WmmaOp::kDq8caQ4gsw) { + if (!device.int8_coopmat) { + return false; + } + const WmmaInt8Layout want = + row.rowmajor_a ? WmmaInt8Layout::kRowMajor : WmmaInt8Layout::k4H4W; + if (shape.int8_layout != WmmaInt8Layout::kAny && + shape.int8_layout != want) { + return false; + } + } + // The pipeline requires the kernel's subgroup size. + if (!subgroup_size_available(device, tile.subgroup_size)) { + return false; + } + // One IO_STORAGE parameter covers input and output. + uint8_t combo = 0; + if (shape.output == WmmaStorage::kBuffer && + shape.input == WmmaStorage::kBuffer && + shape.weight == WmmaStorage::kTexture2D) { + combo = kWmmaBufTex2d; + } else if ( + shape.output == WmmaStorage::kTexture3D && + shape.input == WmmaStorage::kTexture3D && shape.io_width_packed && + shape.weight == WmmaStorage::kTexture2D) { + combo = kWmmaTex3dTex2d; + // The texture epilogue stages SG_GRID_Y * MMA_M rows x WG_TILE_N fp16 on + // top of the A/B tiles unless it reuses the A tile. A tile over the + // shared-memory limit must not reach pipeline creation. + if (!tile.csh_in_ash) { + const uint64_t csh_bytes = + uint64_t(tile.sg_grid_y) * tile.mma_m * tile.n * 2u; + if (csh_bytes >= device.max_shared_bytes) { + return false; + } + } + } + if ((combo & row.storages) == 0) { + return false; + } + if (!shape.ignore_alignment && + (shape.M % tile.m != 0 || shape.N % tile.n != 0 || + shape.K % tile.k != 0 || shape.group_size % tile.k != 0)) { + return false; + } + return row.shape_ok == nullptr || row.shape_ok(shape); +} + +std::optional select_wmma_kernel( + const CoopmatDeviceKey& device, + const WmmaShape& shape) { + for (const WmmaRow& row : coopmat_family_rows(coopmat_family(device))) { + if (wmma_row_fits(device, shape, row)) { + return WmmaChoice{row.kernel_base, row.tile, row.rowmajor_a}; + } + } + return std::nullopt; +} + +std::optional wmma_tile_for_kernel(const std::string& kernel_name) { + for (const CoopmatFamily family : + {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(family)) { + const std::string base(row.kernel_base); + if (kernel_name.size() > base.size() && + kernel_name.compare(0, base.size(), base) == 0 && + kernel_name[base.size()] == '_') { + return row.tile; + } + } + } + return std::nullopt; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h new file mode 100644 index 00000000000..830166e02d0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h @@ -0,0 +1,175 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Cooperative-matrix architecture families for the 4-bit quantized linear ops +// (et_vk.linear_q4gsw / et_vk.q4gsw_linear and et_vk.linear_dq8ca_q4gsw). +// +// A device belongs to a family by its vendor ID, its driver ID and a property +// signature (cooperative-matrix shapes, available subgroup sizes), never by +// its name. Each family has a small set of kernel rows; a device that matches +// no family, or a shape no row of its family accepts, keeps the existing +// kernels unchanged. +// +// Pure C++ (no Vulkan or ComputeGraph types) so it can be unit-tested on a +// host without a GPU (test/custom_ops/coopmat_family_test.cpp). The Vulkan +// side builds a CoopmatDeviceKey from the adapter (WmmaLinear.h). + +#include +#include +#include +#include + +namespace vkcompute { + +// Component types of a cooperative-matrix configuration (VkComponentTypeKHR, +// reduced to the types the family keys test). +enum class CoopmatComponentType { kOther, kFp16, kFp32, kSint8, kSint32 }; + +// One subgroup-scope VkCooperativeMatrixPropertiesKHR entry with A == B and +// C == Result types. +struct CoopmatShape { + uint32_t m; + uint32_t n; + uint32_t k; + CoopmatComponentType ab; // A and B + CoopmatComponentType result; // C and Result +}; + +// VkPhysicalDeviceProperties::vendorID and VkDriverId values the family keys +// use. +constexpr uint32_t kVendorIdIntel = 0x8086; +constexpr uint32_t kVendorIdNvidia = 0x10de; +// VK_DRIVER_ID_NVIDIA_PROPRIETARY +constexpr uint32_t kDriverIdNvidiaProprietary = 4; +// VK_DRIVER_ID_INTEL_OPEN_SOURCE_MESA +constexpr uint32_t kDriverIdIntelMesa = 6; + +struct CoopmatDeviceKey { + uint32_t vendor_id = 0; + uint32_t driver_id = 0; // VkDriverId, 0 when not reported + uint32_t subgroup_size = 0; // default subgroup size + uint32_t min_subgroup_size = 0; + uint32_t max_subgroup_size = 0; + bool subgroup_size_control = false; // a required subgroup size can be set + bool coopmat = false; // VK_KHR_cooperative_matrix is usable + bool int8_coopmat = false; // an int8 cooperative-matrix shape is enumerated + uint32_t max_shared_bytes = 0; // maxComputeSharedMemorySize + std::vector coopmat_shapes; // subgroup scope only +}; + +enum class CoopmatFamily { + kNone, + // Intel Xe2 on the Mesa driver: fp16 8x16x16 (fp16 result), int8 8x16x32, + // subgroup size 16 available. + kXe2, + // NVIDIA proprietary driver: fp16 16x16x16 (fp16 and fp32 result), int8 + // 16x16x32, subgroup size 32. The key does not separate Ampere and Ada from + // other NVIDIA generations with the same shapes. + kNvidia, +}; + +// The only keying point. kNone keeps the existing kernels. +CoopmatFamily coopmat_family(const CoopmatDeviceKey& device); + +const char* coopmat_family_name(CoopmatFamily family); + +// +// Kernel selection +// + +enum class WmmaOp { + kQ4gsw, // 4w: fp16 activations, 4-bit group-symmetric weights + kDq8caQ4gsw, // 8da4w: dynamic int8 activations, 4-bit weights +}; + +enum class WmmaStorage { kBuffer, kTexture2D, kTexture3D, kOther }; + +// Layout of the int8 activations of kDq8caQ4gsw. The 4h4w kernels read +// utils::kPackedInt8_4H4W blocks; the row-major kernels read +// utils::kPackedInt8_4W. +enum class WmmaInt8Layout { kAny, k4H4W, kRowMajor }; + +struct WmmaShape { + WmmaOp op = WmmaOp::kQ4gsw; + int64_t M = 0; + int64_t N = 0; + int64_t K = 0; + int64_t group_size = 0; + int64_t batch = 1; // product of the leading dims + bool gemv = false; + bool has_bias = false; + bool half = false; + WmmaStorage input = WmmaStorage::kOther; + WmmaStorage output = WmmaStorage::kOther; + WmmaStorage weight = WmmaStorage::kOther; + bool io_width_packed = false; + WmmaInt8Layout int8_layout = WmmaInt8Layout::kAny; // kDq8caQ4gsw only + // Skip the M/N/K tile-alignment check (row-major 8da4w kernels, whose + // activation layout is fixed at build time; the grid rounds M up). + bool ignore_alignment = false; +}; + +struct WmmaTile { + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t sg_grid_x; + uint32_t sg_grid_y; + uint32_t subgroup_size; + uint32_t mma_m; + bool csh_in_ash; // texture3d drain staged in the A tile (no extra Csh) + uint32_t wg_size() const { + return sg_grid_x * sg_grid_y * subgroup_size; + } +}; + +// Storage combinations (IO_STORAGE x WEIGHT_STORAGE) a kernel is built for; +// the yaml contains ___half for each set bit. +enum WmmaStorageSet : uint8_t { + kWmmaTex3dTex2d = 1u << 0, // texture3d input/output, texture2d weight + kWmmaBufTex2d = 1u << 1, // buffer input/output, texture2d weight +}; + +struct WmmaRow { + WmmaOp op; + const char* kernel_base; // shader name without ___ + WmmaTile tile; + uint8_t storages; // WmmaStorageSet bits + // Extra shape requirement (beyond tile alignment), or null. Rows are tried + // in order, so a later row can serve the shapes an earlier one rejects. + bool (*shape_ok)(const WmmaShape&); + bool rowmajor_a = false; // kDq8caQ4gsw: reads row-major int8 activations +}; + +struct WmmaChoice { + std::string kernel_base; + WmmaTile tile; + bool rowmajor_a = false; +}; + +// The rows of a family, in selection order (empty for kNone). +const std::vector& coopmat_family_rows(CoopmatFamily family); + +// The device, storage, alignment and shape constraints of `row`. +bool wmma_row_fits( + const CoopmatDeviceKey& device, + const WmmaShape& shape, + const WmmaRow& row); + +// The first row of the device's family that fits `shape`, or std::nullopt +// (no family, or no row applies): the op keeps its existing kernels. +std::optional select_wmma_kernel( + const CoopmatDeviceKey& device, + const WmmaShape& shape); + +// Tile of a kernel named _ for any family row. +std::optional wmma_tile_for_kernel(const std::string& kernel_name); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp new file mode 100644 index 00000000000..02d94544a0e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp @@ -0,0 +1,118 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +namespace vkcompute { + +namespace { + +#ifdef VK_KHR_cooperative_matrix +CoopmatComponentType to_component_type(const VkComponentTypeKHR t) { + switch (t) { + case VK_COMPONENT_TYPE_FLOAT16_KHR: + return CoopmatComponentType::kFp16; + case VK_COMPONENT_TYPE_FLOAT32_KHR: + return CoopmatComponentType::kFp32; + case VK_COMPONENT_TYPE_SINT8_KHR: + return CoopmatComponentType::kSint8; + case VK_COMPONENT_TYPE_SINT32_KHR: + return CoopmatComponentType::kSint32; + default: + return CoopmatComponentType::kOther; + } +} +#endif /* VK_KHR_cooperative_matrix */ + +} // namespace + +CoopmatDeviceKey coopmat_device_key(ComputeGraph* graph) { + vkapi::Adapter* adapter = graph->context()->adapter_ptr(); + CoopmatDeviceKey key; + key.vendor_id = adapter->vendor_id(); + key.driver_id = adapter->driver_id(); + key.subgroup_size = adapter->subgroup_size(); + key.min_subgroup_size = adapter->min_subgroup_size(); + key.max_subgroup_size = adapter->max_subgroup_size(); + key.subgroup_size_control = adapter->supports_subgroup_size_control(); + key.coopmat = adapter->supports_cooperative_matrix(); + key.int8_coopmat = adapter->supports_int8_cooperative_matrix(); + key.max_shared_bytes = adapter->max_compute_shared_memory_size(); +#ifdef VK_KHR_cooperative_matrix + for (const VkCooperativeMatrixPropertiesKHR& p : + adapter->cooperative_matrix_properties()) { + if (p.scope != VK_SCOPE_SUBGROUP_KHR || p.AType != p.BType || + p.CType != p.ResultType) { + continue; + } + key.coopmat_shapes.push_back( + {p.MSize, + p.NSize, + p.KSize, + to_component_type(p.AType), + to_component_type(p.ResultType)}); + } +#endif /* VK_KHR_cooperative_matrix */ + return key; +} + +WmmaStorage to_wmma_storage(const utils::StorageType storage) { + switch (storage) { + case utils::kBuffer: + return WmmaStorage::kBuffer; + case utils::kTexture2D: + return WmmaStorage::kTexture2D; + case utils::kTexture3D: + return WmmaStorage::kTexture3D; + default: + return WmmaStorage::kOther; + } +} + +utils::StorageType predicted_q4_weight_storage( + ComputeGraph& graph, + const ValueRef weight_data) { + const std::vector sizes = graph.sizes_of(weight_data); + const int64_t K = utils::val_at(-1, sizes) * 2; + const int64_t N = utils::val_at(-2, sizes); + // Mirrors the 4-bit packed extents in prepack_quantized_linear_weight(). + const int64_t height = utils::div_up(N, int64_t(8)); + const int64_t width = utils::div_up(K, int64_t(4)) * 4; + const uint32_t max_extent = + graph.context()->adapter_ptr()->max_texture2d_dim(); + return (width > int64_t(max_extent) * 4 || height > int64_t(max_extent)) + ? utils::kBuffer + : utils::kTexture2D; +} + +GlobalWorkGrid wmma_linear_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const std::optional tile = wmma_tile_for_kernel(shader.kernel_name); + if (!tile.has_value()) { + return quantized_linear_gwg(graph, shader, args, resize_args); + } + const std::vector out_sizes = graph->sizes_of(args.at(0).refs.at(0)); + const uint32_t N = + utils::safe_downcast(utils::val_at(-1, out_sizes)); + const uint32_t M = + utils::safe_downcast(utils::val_at(-2, out_sizes)); + // The x extent is multiplied by the workgroup size because the framework + // divides it by the local size. + const uint32_t wg_size = tile->wg_size(); + return GlobalWorkGrid( + {utils::div_up(N, tile->n) * wg_size, utils::div_up(M, tile->m), 1u}, + kTiledWorkGrid, + LocalWorkGroup(wg_size, 1u, 1u)); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h new file mode 100644 index 00000000000..ea55cb17cc3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h @@ -0,0 +1,39 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// ComputeGraph side of the cooperative-matrix family kernels +// (CoopmatFamily.h), shared by the q4gsw and dq8ca_q4gsw node builders. + +#include +#include + +namespace vkcompute { + +// The family key inputs of the graph's device. +CoopmatDeviceKey coopmat_device_key(ComputeGraph* graph); + +WmmaStorage to_wmma_storage(utils::StorageType storage); + +// The storage prepack_quantized_linear_weight() picks for a 4-bit weight: +// texture2d unless the packed extents exceed the device limit. +utils::StorageType predicted_q4_weight_storage( + ComputeGraph& graph, + const ValueRef weight_data); + +// One workgroup per WG_TILE_M x WG_TILE_N output tile for a family kernel; +// quantized_linear_gwg() for any other kernel name. args.at(0).refs.at(0) is +// the output. +GlobalWorkGrid wmma_linear_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args); + +} // namespace vkcompute From 3b8c77fb25350fe5084e788ed7ef6a40961f5179 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 16 Oct 2026 11:30:55 -0700 Subject: [PATCH 03/11] [ET-VK] q4gsw linear: cooperative-matrix family kernels 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 Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm --- .../linear/qw/q4gsw/linear_q4gsw_wmma.glsl | 97 +++ .../linear/qw/q4gsw/linear_q4gsw_wmma.yaml | 60 ++ .../qw/q4gsw/linear_q4gsw_wmma_body.glslh | 577 ++++++++++++++++++ .../linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl | 83 +++ .../linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml | 36 ++ .../qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh | 416 +++++++++++++ .../impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp | 7 + .../gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp | 176 ++++++ .../gemm/linear/qw/q4gsw/Q4gswLinearWmma.h | 25 + 9 files changed, 1477 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl new file mode 100644 index 00000000000..84c1cb73fd6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl @@ -0,0 +1,97 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * fp16 q4gsw (4-bit group-symmetric weight) cooperative-matrix linear with + * per-variant tile geometry, used by the architecture-family path + * (impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp). + * + * This file is the template header: bindings, spec constants and the tile + * geometry. The kernel body is linear_q4gsw_wmma_body.glslh. + * + * Loop structure: double-buffered, prefetch-peeled. The prologue prefetches + * chunk 0; each iteration does barrier -> prefetch(next) -> MMA(cur) -> + * store(next). + * + * Feature defines (yaml parameters, all default off): + * B_COLMAJOR B staged N-major in shared memory, ColumnMajor B coopMatLoad + * CSH_BAND texture3d drain one SG_GRID_Y band at a time + * FRAG_LAYOUT fragment-contiguous shared memory, no padding + * IMG_A storage-image loads for A (texture3d IO only) + * MMA_M = 8 for devices whose fp16 cooperative matrix is 8x16x16 only + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: M % WG_TILE_M == 0, N % WG_TILE_N == 0, + * K % WG_TILE_K == 0, group_size % WG_TILE_K == 0, no bias, batch 1. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +$if IMG_A and IO_STORAGE == "texture3d": + ${layout_declare_image(B, "r", "t_input", "half")} +$else: + ${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +#define ACC_T float16_t + +$if B_COLMAJOR: + #define B_COLMAJOR + +$if CSH_BAND: + #define CSH_BAND + +$if FRAG_LAYOUT: + #define FRAG_LAYOUT + +$if IMG_A and IO_STORAGE == "texture3d": + #define IMG_A + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_q4gsw_wmma_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml new file mode 100644 index 00000000000..75ca4ee943c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the architecture-family q4gsw cooperative-matrix linear. Each +# one is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, which +# also carries its tile geometry; the two must stay in sync. +# +# Name: txkgs +# followed by m8 (MMA_M = 8), f (FRAG_LAYOUT), i (IMG_A), b (CSH_BAND) or +# bt (B_COLMAJOR). + +linear_q4gsw_wmma: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + FRAG_LAYOUT: false + IMG_A: false + CSH_BAND: false + B_COLMAJOR: false + shader_variants: + # Intel Xe2: subgroup 16, MMA 8x16x16 (the only fp16 shape Xe2 exposes), + # fragment-contiguous shared memory and image loads for A. + - NAME: linear_q4gsw_wmma_t128x128k16g82s16m8flib_texture3d_texture2d_half + IO_STORAGE: texture3d + MMA_M: 8 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 16 + FRAG_LAYOUT: true + IMG_A: true + CSH_BAND: true + - NAME: linear_q4gsw_wmma_t128x128k16g44s16m8fli_buffer_texture2d_half + MMA_M: 8 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 16 + FRAG_LAYOUT: true + IMG_A: true + # NVIDIA: subgroup 32, MMA 16x16x16 with fp16 accumulation (K <= 8192). + - NAME: linear_q4gsw_wmma_t256x128k16g42s32bt_texture3d_texture2d_half + IO_STORAGE: texture3d + WG_TILE_M: 256 + SG_GRID_X: 4 + SG_GRID_Y: 2 + B_COLMAJOR: true + - NAME: linear_q4gsw_wmma_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh new file mode 100644 index 00000000000..668f1b6da04 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh @@ -0,0 +1,577 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_q4gsw_wmma. Included by its template header +// (linear_q4gsw_wmma.glsl), which defines the bindings, the spec constants, +// ACC_T, the optional IO_TEXTURE / FRAG_LAYOUT / IMG_A / B_COLMAJOR / CSH_BAND +// macros and the tile constants MMA_{M,N,K}, WG_TILE_{M,N,K}, SG_GRID_{X,Y}, +// SUBGROUP_SIZE. Must not contain template syntax: .glslh files are not +// expanded per variant. + +#ifndef LINEAR_Q4GSW_WMMA_BODY_GLSLH +#define LINEAR_Q4GSW_WMMA_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +// fp16: 8 elements per uvec4 (128-bit) +const uint FP16_PER_VEC4 = 8; +#ifdef FRAG_LAYOUT +// Fragment-contiguous shared memory: A as [k16 block][m][16 fp16] and B as +// [n16 block][k][16 fp16], so every MMA_M x 16 A fragment and 16 x MMA_N B +// fragment is one contiguous run with a 32-byte row stride. The row-padded +// layout below scatters a fragment across strided 32-byte rows, which some +// compilers lower to many narrow shared-memory messages; where shared memory +// has no bank conflicts the padding buys nothing. Requires +// MMA_K == MMA_N == 16. +const uint FRAG_ROW_VEC4 = 16u / FP16_PER_VEC4; +const uint A_STRIDE_VEC4 = FRAG_ROW_VEC4; +const uint B_STRIDE_VEC4 = FRAG_ROW_VEC4; +const uint ASH_SLICE = WG_TILE_M * (WG_TILE_K / FP16_PER_VEC4); +const uint BSH_SLICE = WG_TILE_K * (WG_TILE_N / FP16_PER_VEC4); +// (row, uvec4 column) -> slice-relative index. +#define A_SH_IDX(m, c) ((((c) / FRAG_ROW_VEC4) * WG_TILE_M + (m)) * FRAG_ROW_VEC4 + ((c) % FRAG_ROW_VEC4)) +#define B_SH_IDX(k, c) ((((c) / FRAG_ROW_VEC4) * WG_TILE_K + (k)) * FRAG_ROW_VEC4 + ((c) % FRAG_ROW_VEC4)) +#else +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +// One ping-pong slice of each shared-memory buffer (in uvec4 units). +const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; +const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; +// No outer parentheses: every use is `base + X_SH_IDX(...)` or a subscript, +// so `base + row * STRIDE + col` keeps its natural association. +#define A_SH_IDX(m, c) (m) * A_STRIDE_VEC4 + (c) +#define B_SH_IDX(k, c) (k) * B_STRIDE_VEC4 + (c) +#endif + +// Double-buffered shared memory. +shared uvec4 Ash[2 * ASH_SLICE]; +#ifdef B_COLMAJOR +// B_COLMAJOR: on devices whose MMA B operand is one N column of MMA_K +// consecutive K values per lane, coopMatLoad from the row-major Bsh compiles +// to per-element gathers. Stage B N-major (BshT[n][k], K contiguous, padded +// like Ash) and load it ColumnMajor instead. +// Units: uvec2 = 4 fp16. +#if defined(FRAG_LAYOUT) +#error "B_COLMAJOR is not combined with FRAG_LAYOUT" +#endif +const uint BT_STRIDE_H = WG_TILE_K + 8u; +const uint BT_STRIDE_V2 = BT_STRIDE_H / 4u; +const uint BT_SLICE = WG_TILE_N * BT_STRIDE_V2; +shared uvec2 BshT[2 * BT_SLICE]; +#else +shared uvec4 Bsh[2 * BSH_SLICE]; +#endif +// Shared-memory stores of one staged uvec4. +#define ASH_STORE(idx, v) Ash[idx] = v +#define BSH_STORE(idx, v) Bsh[idx] = v + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. Not aliased +// with Ash/Bsh: GLSL has no shared unions and coopMatStore needs a +// float16_t-typed array. +#if defined(CSH_BAND) +// CSH_BAND: drain one SG_GRID_Y band at a time, so Csh is MMA_M rows x +// WG_TILE_N. +const uint CSH_ROWS = MMA_M; +#else +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +#endif +shared float16_t Csh[CSH_ROWS * WG_TILE_N]; +#endif + +// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. +const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; +const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; +const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; +const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; +const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; +const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; +#ifdef B_COLMAJOR +// One item = half of one packed-weight texel (4 K x 8 N nibbles): 4 K x 4 N, +// half 0 = N 0..3 (even nibbles), half 1 = N 4..7 (odd nibbles). K varies +// fastest across lanes: each (n8, half) group of BT_KQ lanes writes one +// contiguous run per column, and the groups land on disjoint shared-memory +// banks. +const uint BT_KQ = WG_TILE_K / 4u; +const uint BT_N8 = WG_TILE_N / 8u; +const uint BT_ITEMS = BT_KQ * BT_N8 * 2u; +const uint BT_PASSES = (BT_ITEMS + WG_SIZE - 1u) / WG_SIZE; +#endif + +// Accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per subgroup). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights +// (one Bsh uvec4). col_lo/col_hi select the K row within the block. +#define WFETCH(coord) texelFetch(t_packed_weight, coord, 0) + +uvec4 dequant_block( + const ivec4 wb, + const uint col_lo, + const uint col_hi, + const f16vec4 s0, + const f16vec4 s1) { + const f16vec4 v0 = f16vec4(((wb >> int(4u * col_lo)) & 0xF) - 8) * s0; + const f16vec4 v1 = f16vec4(((wb >> int(4u * col_hi)) & 0xF) - 8) * s1; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +} + +#ifdef B_COLMAJOR +// Store one item: 4 N columns (x..w of wb, nibble parity hf) x 4 K rows +// (nibbles 2*kk + hf), each column's 4 K values as one uvec2. +void store_bt_item( + const uint base, + const ivec4 wb, + const uint n_local, + const uint k_local, + const uint hf, + const f16vec4 sc) { + [[unroll]] for (uint j = 0; j < 4u; ++j) { + const int w = wb[j]; + const f16vec4 v = f16vec4( + ((w >> int(4u * (0u + hf))) & 0xF) - 8, + ((w >> int(4u * (2u + hf))) & 0xF) - 8, + ((w >> int(4u * (4u + hf))) & 0xF) - 8, + ((w >> int(4u * (6u + hf))) & 0xF) - 8) * sc[j]; + BshT[base + (n_local + j) * BT_STRIDE_V2 + k_local / 4u] = + uvec2(packFloat2x16(v.xy), packFloat2x16(v.zw)); + } +} +#define BT_ACTIVE(p) (BT_ITEMS % WG_SIZE == 0u || (p) * WG_SIZE + gl_LocalInvocationID.x < BT_ITEMS) +#endif + +// Fetch 8 consecutive fp16 activations of `row` starting at half-vec4 index +// `k_hv4`, packed into one uvec4. Both spellings address the same bytes: a +// width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), +// so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in +// two notations. +uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { +#ifdef IO_TEXTURE + // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 + // (fp32); narrowing immediately lets the compiler fold the conversion into + // a 16-bit image load. Lossless -- the source is rgba16f. +#ifdef IMG_A + // Storage-image load instead of a sampled fetch. + const f16vec4 v0 = f16vec4(imageLoad(t_input, ivec3(k_hv4, row, 0))); + const f16vec4 v1 = f16vec4(imageLoad(t_input, ivec3(k_hv4 + 1u, row, 0))); +#else + const f16vec4 v0 = f16vec4(texelFetch(t_input, ivec3(k_hv4, row, 0), 0)); + const f16vec4 v1 = f16vec4(texelFetch(t_input, ivec3(k_hv4 + 1u, row, 0), 0)); +#endif + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#else + const f16vec4 v0 = t_input[row * K4 + k_hv4]; + const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#endif +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint K4 = (K + 3u) / 4u; + const uint N4 = (uint(output_sizes.x) + 3u) / 4u; + + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; + const uint col_lo = 2u * (b_row_offset & 3u); + const uint col_hi = col_lo + 1u; + + const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; + uint cached_group = 0xFFFFFFFFu; + f16vec4 sc0; + f16vec4 sc1; + + uvec4 temp_A[A_PASSES]; +#ifdef B_COLMAJOR + ivec4 temp_B[BT_PASSES]; + f16vec4 sc_bt[BT_PASSES]; + uint bt_kq[BT_PASSES], bt_n8[BT_PASSES], bt_hf[BT_PASSES]; + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + const uint item = p * WG_SIZE + gl_LocalInvocationID.x; + bt_kq[p] = item % BT_KQ; + bt_hf[p] = (item / BT_KQ) & 1u; + bt_n8[p] = (item / BT_KQ) >> 1u; + } +#else + ivec4 temp_B[B_PASSES]; +#endif + + // ========================================================= + // PROLOGUE: prefetch chunk 0 into temp registers, then store to + // slice 0. No barrier here -- the main loop's first iteration barriers + // before reading slice 0. + // ========================================================= + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + cached_group = 0u; +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + const uint n8g = (tile_n_start >> 3u) + bt_n8[p]; + sc_bt[p] = t_weight_scales[n8g * 2u + bt_hf[p]]; + temp_B[p] = WFETCH(ivec2(bt_kq[p], n8g)); + } + } +#else + sc0 = t_weight_scales[sc_n4]; + sc1 = t_weight_scales[sc_n4 + 1u]; + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = p * B_ROWS_PER_PASS + b_row_offset; + ivec4 wblock; + wblock = WFETCH(ivec2(k_row >> 2u, n8_blk)); + temp_B[p] = wblock; + } +#endif + } + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), temp_A[p]); + } +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + store_bt_item(0u, temp_B[p], bt_n8[p] * 8u + 4u * bt_hf[p], bt_kq[p] * 4u, bt_hf[p], sc_bt[p]); + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(B_SH_IDX(p * B_ROWS_PER_PASS + b_row_offset, b_col), + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1)); + } +#endif + } + + // ========================================================= + // MAIN LOOP: one barrier per iteration, unconditional body, loop + // peeled (bound excludes the last chunk). Iteration `chunk` does: + // 1. barrier — slice (chunk%2) fully written + // 2. prefetch — chunk+1 from global into temp (in flight during math) + // 3. MMA math — on slice (chunk%2) + // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) + // ========================================================= + uint chunk; + for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; +#ifdef B_COLMAJOR + const uint cur_base_B = (chunk % 2u) * BT_SLICE; + const uint nxt_base_B = ((chunk + 1u) % 2u) * BT_SLICE; +#else + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; +#endif + const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; +#if !defined(B_COLMAJOR) + const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; +#endif + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- prefetch chunk+1 -> temp --- + { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } +#ifdef B_COLMAJOR + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + const bool new_group = group_nxt != cached_group; + cached_group = group_nxt; + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + const uint n8g = (tile_n_start >> 3u) + bt_n8[p]; + temp_B[p] = WFETCH(ivec2(chunkK_nxt / 4u + bt_kq[p], n8g)); + if (new_group) { + sc_bt[p] = t_weight_scales[group_nxt * N4 + n8g * 2u + bt_hf[p]]; + } + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = chunkK_nxt + p * B_ROWS_PER_PASS + b_row_offset; + temp_B[p] = WFETCH(ivec2(k_row >> 2u, n8_blk)); + } + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + if (group_nxt != cached_group) { + cached_group = group_nxt; + sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; + sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; + } +#endif + } + + // --- MMA math on the cur slice --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef B_COLMAJOR + const uint n_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, BshT, + cur_base_B + n_b * BT_STRIDE_V2 + k_start / 4u, + BT_STRIDE_V2, + gl_CooperativeMatrixLayoutColumnMajor); +#else + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#endif // B_COLMAJOR + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + + // --- store temp (chunk+1) -> nxt slice, dequantizing B --- + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(nxt_base_A + A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), + temp_A[p]); + } +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + store_bt_item(nxt_base_B, temp_B[p], bt_n8[p] * 8u + 4u * bt_hf[p], bt_kq[p] * 4u, bt_hf[p], sc_bt[p]); + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(nxt_base_B + B_SH_IDX(p * B_ROWS_PER_PASS + b_row_offset, b_col), + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1)); + } +#endif + } + } + + // --- epilogue: barrier, then MMA on the last chunk (loop peeled) --- + { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; +#ifdef B_COLMAJOR + const uint cur_base_B = (chunk % 2u) * BT_SLICE; +#else + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; +#endif + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef B_COLMAJOR + const uint n_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, BshT, + cur_base_B + n_b * BT_STRIDE_V2 + k_start / 4u, + BT_STRIDE_V2, + gl_CooperativeMatrixLayoutColumnMajor); +#else + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#endif // B_COLMAJOR + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + } + + + // --- Store result tile --- +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh at once, so the SG_GRID_Y bands it holds are disjoint global row + // ranges; the whole workgroup then imageStores them. The band-to-global-row + // map below reproduces the buffer path's gi exactly: + // buffer: gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i) + // = tile_m_start + warpInTile.y * SG_TILE_M + i * MMA_M + // and lr / MMA_M is the writing subgroup's warpInTile.y. + // + // PORTABILITY NOTE: [[unroll]] is only a hint and glslc does not honor it + // here -- the barrier() in the loop body keeps the loop rolled, so + // result[i][j] is dynamically indexed. Coopmat arrays are opaque per-lane + // storage and dynamic indexing is a construct shader compilers have broken + // before, so check this first if the texture variants miscompile on a new + // device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; +#if defined(CSH_BAND) + // Band b = the subgroups with warpInTile.y == b; per (i, b): barrier, that + // band's subgroups store accumulator row-block i into Csh, barrier, the + // whole workgroup imageStores the MMA_M rows. Same global row map as below. + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + for (uint b = 0; b < SG_GRID_Y; ++b) { + memoryBarrierShared(); + barrier(); + if (warpInTile.y == b) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopMatStore( + coopmat(result[i][j]), Csh, + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + } + memoryBarrierShared(); + barrier(); + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = tile_m_start + b * SG_TILE_M + i * MMA_M + lr; + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); + } + } + } +#else + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh against the previous iteration's readers. Inert on i == 0 but + // must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopMatStore( + coopmat(result[i][j]), Csh, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); + } + } +#endif // CSH_BAND +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_Q4GSW_WMMA_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl new file mode 100644 index 00000000000..ecc8fde5bd3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl @@ -0,0 +1,83 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * fp16 q4gsw cooperative-matrix linear with texel-wise weight staging and an + * optional fp32 accumulator, used by the architecture-family path + * (impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp) for long-K shapes where + * fp16 accumulation loses accuracy. Template header only; the body is + * linear_q4gsw_wmma_bx_body.glslh. + * + * Feature defines (yaml parameters, all default off): + * ACC_FP32 fp32 accumulator + * CSH_IN_ASH texture3d drain staged in the A shared-memory tile, which is + * dead after the last MMA + * + * Preconditions as for linear_q4gsw_wmma.glsl. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +$if ACC_FP32: + #define ACC_FP32 + +#ifdef ACC_FP32 +#define ACC_T float +#else +#define ACC_T float16_t +#endif + +$if CSH_IN_ASH: + #define CSH_IN_ASH + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_q4gsw_wmma_bx_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml new file mode 100644 index 00000000000..7951b48ff16 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml @@ -0,0 +1,36 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the texel-wise-staging q4gsw cooperative-matrix linear. Each one +# is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, which also +# carries its tile geometry; the two must stay in sync. Name as in +# linear_q4gsw_wmma.yaml, plus f32 (ACC_FP32) and c (CSH_IN_ASH). + +linear_q4gsw_wmma_bx: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + ACC_FP32: false + CSH_IN_ASH: false + shader_variants: + # NVIDIA, K > 8192: fp32 accumulation keeps the long-K error in bounds. + - NAME: linear_q4gsw_wmma_bx_t128x128k32g42s32f32c_texture3d_texture2d_half + IO_STORAGE: texture3d + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + ACC_FP32: true + CSH_IN_ASH: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh new file mode 100644 index 00000000000..10f258ab782 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh @@ -0,0 +1,416 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_q4gsw_wmma_bx: linear_q4gsw_wmma_body.glslh with texel-wise +// B staging (one packed-weight texel feeds every B pass of a thread). +// Included by its template header (linear_q4gsw_wmma_bx.glsl), which defines +// the bindings, the spec constants, ACC_T, the optional IO_TEXTURE / +// CSH_IN_ASH macros and the tile constants MMA_{M,N,K}, WG_TILE_{M,N,K}, +// SG_GRID_{X,Y}, SUBGROUP_SIZE. Must not contain template syntax. + +#ifndef LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH +#define LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +// fp16: 8 elements per uvec4 (128-bit) +const uint FP16_PER_VEC4 = 8; +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +// One ping-pong slice of each shared-memory buffer (in uvec4 units). +const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; +const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; +// No outer parentheses: every use is `base + X_SH_IDX(...)` or a subscript, +// so `base + row * STRIDE + col` keeps its natural association. +#define A_SH_IDX(m, c) (m) * A_STRIDE_VEC4 + (c) +#define B_SH_IDX(k, c) (k) * B_STRIDE_VEC4 + (c) + +// Double-buffered shared memory. +shared uvec4 Ash[2 * ASH_SLICE]; +shared uvec4 Bsh[2 * BSH_SLICE]; +// Shared-memory stores of one staged uvec4. +#define ASH_STORE(idx, v) Ash[idx] = v +#define BSH_STORE(idx, v) Bsh[idx] = v + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. Not aliased +// with Ash/Bsh: GLSL has no shared unions and coopMatStore needs a +// float16_t-typed array. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +// CSH_IN_ASH stages the drain band in Ash instead: Ash is dead after the +// last MMA and every drain iteration starts with a barrier. +#ifndef CSH_IN_ASH +shared float16_t Csh[CSH_ROWS * WG_TILE_N]; +#endif +#endif + +// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. +const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; +const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; +const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; +const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; +const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; +const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; +// Texel-wise B staging: one packed-weight texel (4 K rows x 8 N) feeds all +// B_PASSES rows of a thread; B_RP threads share a texel when B_PASSES < 4. +// Row of pass p within the chunk, and its first nibble index in the texel. +const uint B_RP = 4u / B_PASSES; +#define B_TROW(p) ((b_row_offset / B_RP) * 4u + (b_row_offset % B_RP) * B_PASSES + (p)) +#define B_TCOL(p) (2u * ((b_row_offset % B_RP) * B_PASSES + (p))) + +// Accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per subgroup). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights +// (one Bsh uvec4). col_lo/col_hi select the K row within the block. +#define WFETCH(coord) texelFetch(t_packed_weight, coord, 0) + +uvec4 dequant_block( + const ivec4 wb, + const uint col_lo, + const uint col_hi, + const f16vec4 s0, + const f16vec4 s1) { + const f16vec4 v0 = f16vec4(((wb >> int(4u * col_lo)) & 0xF) - 8) * s0; + const f16vec4 v1 = f16vec4(((wb >> int(4u * col_hi)) & 0xF) - 8) * s1; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +} + + +// Fetch 8 consecutive fp16 activations of `row` starting at half-vec4 index +// `k_hv4`, packed into one uvec4. Both spellings address the same bytes: a +// width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), +// so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in +// two notations. +uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { +#ifdef IO_TEXTURE + // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 + // (fp32); narrowing immediately lets the compiler fold the conversion into + // a 16-bit image load. Lossless -- the source is rgba16f. + const f16vec4 v0 = f16vec4(texelFetch(t_input, ivec3(k_hv4, row, 0), 0)); + const f16vec4 v1 = f16vec4(texelFetch(t_input, ivec3(k_hv4 + 1u, row, 0), 0)); + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#else + const f16vec4 v0 = t_input[row * K4 + k_hv4]; + const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#endif +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint K4 = (K + 3u) / 4u; + const uint N4 = (uint(output_sizes.x) + 3u) / 4u; + + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; + const uint col_lo = 2u * (b_row_offset & 3u); + const uint col_hi = col_lo + 1u; + + const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; + uint cached_group = 0xFFFFFFFFu; + f16vec4 sc0; + f16vec4 sc1; + + uvec4 temp_A[A_PASSES]; + ivec4 temp_B[B_PASSES]; + + // ========================================================= + // PROLOGUE: prefetch chunk 0 into temp registers, then store to + // slice 0. No barrier here -- the main loop's first iteration barriers + // before reading slice 0. + // ========================================================= + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + cached_group = 0u; + sc0 = t_weight_scales[sc_n4]; + sc1 = t_weight_scales[sc_n4 + 1u]; + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + if (p > 0u) { + temp_B[p] = temp_B[0]; + continue; + } + const uint k_row = B_TROW(0u); + ivec4 wblock; + wblock = WFETCH(ivec2(k_row >> 2u, n8_blk)); + temp_B[p] = wblock; + } + } + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), temp_A[p]); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(B_SH_IDX(B_TROW(p), b_col), + dequant_block(temp_B[p], B_TCOL(p), B_TCOL(p) + 1u, sc0, sc1)); + } + } + + // ========================================================= + // MAIN LOOP: one barrier per iteration, unconditional body, loop + // peeled (bound excludes the last chunk). Iteration `chunk` does: + // 1. barrier — slice (chunk%2) fully written + // 2. prefetch — chunk+1 from global into temp (in flight during math) + // 3. MMA math — on slice (chunk%2) + // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) + // ========================================================= + uint chunk; + for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; + const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- prefetch chunk+1 -> temp --- + { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + if (p > 0u) { + temp_B[p] = temp_B[0]; + continue; + } + const uint k_row = chunkK_nxt + B_TROW(0u); + temp_B[p] = WFETCH(ivec2(k_row >> 2u, n8_blk)); + } + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + if (group_nxt != cached_group) { + cached_group = group_nxt; + sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; + sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; + } + } + + // --- MMA math on the cur slice --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + + // --- store temp (chunk+1) -> nxt slice, dequantizing B --- + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(nxt_base_A + A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), + temp_A[p]); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(nxt_base_B + B_SH_IDX(B_TROW(p), b_col), + dequant_block(temp_B[p], B_TCOL(p), B_TCOL(p) + 1u, sc0, sc1)); + } + } + } + + // --- epilogue: barrier, then MMA on the last chunk (loop peeled) --- + { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + } + + + // --- Store result tile --- +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh at once, so the SG_GRID_Y bands it holds are disjoint global row + // ranges; the whole workgroup then imageStores them. The band-to-global-row + // map below reproduces the buffer path's gi exactly: + // buffer: gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i) + // = tile_m_start + warpInTile.y * SG_TILE_M + i * MMA_M + // and lr / MMA_M is the writing subgroup's warpInTile.y. + // + // PORTABILITY NOTE: [[unroll]] is only a hint and glslc does not honor it + // here -- the barrier() in the loop body keeps the loop rolled, so + // result[i][j] is dynamically indexed. Coopmat arrays are opaque per-lane + // storage and dynamic indexing is a construct shader compilers have broken + // before, so check this first if the texture variants miscompile on a new + // device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh against the previous iteration's readers. Inert on i == 0 but + // must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef CSH_IN_ASH + coopMatStore( + coopmat(result[i][j]), Ash, + (warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j)) / FP16_PER_VEC4, + WG_TILE_N / FP16_PER_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#else + coopMatStore( + coopmat(result[i][j]), Csh, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); +#endif + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; +#ifdef CSH_IN_ASH + // base is a multiple of 4 halves: the texel is one half of a uvec4. + const uvec4 q = Ash[base / FP16_PER_VEC4]; + const bool hi = (base % FP16_PER_VEC4) != 0u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4(unpackHalf2x16(hi ? q.z : q.x), unpackHalf2x16(hi ? q.w : q.y))); +#else + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); +#endif + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp index 8e7d16092c7..aebc0acf5b2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace vkcompute { @@ -634,6 +635,12 @@ void add_q4gsw_linear_tin_w_4x8_node( } void q4gsw_linear(ComputeGraph& graph, const std::vector& args) { + // Devices of a cooperative-matrix family (CoopmatFamily.h) take the + // family's kernels where one of its rows applies. + if (try_add_q4gsw_linear_wmma_node(graph, args)) { + return; + } + int32_t idx = 0; const ValueRef fp_input = args.at(idx++); const ValueRef weight_data = args.at(idx++); diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp new file mode 100644 index 00000000000..228bda2f6ac --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp @@ -0,0 +1,176 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vkcompute { + +namespace { + +// Dispatch node args: +// {{output}, {fp_input, packed_weight, packed_weight_scales, packed_bias}} +// resize_args: {group_size, weight_data, bias_data} +WmmaShape q4gsw_runtime_shape( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_weight = args.at(1).refs.at(1); + + const std::vector out_sizes = graph->sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph->sizes_of(fp_input)); + s.group_size = graph->extract_scalar(resize_args.at(0)); + for (int64_t d = 0; d < graph->dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(graph, fp_input); + s.has_bias = !graph->val_is_none(resize_args.at(2)); + s.half = graph->dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph->storage_type_of(fp_input)); + s.output = to_wmma_storage(graph->storage_type_of(output)); + s.weight = to_wmma_storage(graph->storage_type_of(packed_weight)); + s.io_width_packed = graph->packed_dim_of(output) == WHCN::kWidthDim && + graph->packed_dim_of(fp_input) == WHCN::kWidthDim; + return s; +} + +// args: {input, weight, weight_scales, group_size, bias, output} +WmmaShape q4gsw_build_time_shape( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef output = args.at(5); + const std::vector out_sizes = graph.sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph.sizes_of(fp_input)); + s.group_size = graph.extract_scalar(args.at(3)); + for (int64_t d = 0; d < graph.dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(&graph, fp_input); + s.has_bias = !graph.val_is_none(args.at(4)); + s.half = graph.dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph.storage_type_of(fp_input)); + s.output = to_wmma_storage(graph.storage_type_of(output)); + s.weight = to_wmma_storage(predicted_q4_weight_storage(graph, args.at(1))); + s.io_width_packed = graph.packed_dim_of(output) == WHCN::kWidthDim && + graph.packed_dim_of(fp_input) == WHCN::kWidthDim; + return s; +} + +vkapi::ShaderInfo pick_q4gsw_linear_wmma_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_weight = args.at(1).refs.at(1); + + const std::optional choice = select_wmma_kernel( + coopmat_device_key(graph), q4gsw_runtime_shape(graph, args, resize_args)); + + std::string kernel_name; + if (choice.has_value()) { + kernel_name = choice->kernel_base; + } else { + // Shapes no family row covers after a resize (decode, unaligned M): the + // kernels that read this packed weight layout. + kernel_name = + is_gemv(graph, fp_input) ? "linear_q4gsw_coop" : "linear_q4gsw_tiled"; + } + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + return VK_KERNEL_FROM_STR(kernel_name); +} + +} // namespace + +bool try_add_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef weight_data = args.at(1); + const ValueRef weight_scales_data = args.at(2); + const ValueRef group_size = args.at(3); + const ValueRef bias_data = args.at(4); + const ValueRef output = args.at(5); + + // The family kernels are fp16 only. + if (graph.dtype_of(fp_input) != vkapi::kHalf) { + return false; + } + // Take the op over only if a row applies to its build-time shape; otherwise + // it stays entirely on the existing path, including for later resizes. + if (!select_wmma_kernel( + coopmat_device_key(&graph), q4gsw_build_time_shape(graph, args)) + .has_value()) { + return false; + } + + const int64_t group_size_val = graph.extract_scalar(group_size); + const QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + const ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + + TmpTensor dummy_bias( + &graph, {}, graph.dtype_of(output), utils::kBuffer, utils::kWidthPacked); + ValueRef packed_bias = dummy_bias.vref; + uint32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + apply_bias = 1; + } + + const int32_t K4_per_group = + utils::div_up(static_cast(group_size_val), int32_t(4)); + const int32_t num_groups = graph.size_at(-1, fp_input) / + static_cast(group_size_val); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_q4gsw_linear_wmma_shader, + wmma_linear_gwg, + quantized_linear_lwg, + {{output, vkapi::kWrite}, + {{fp_input, packed_weight, packed_weight_scales, packed_bias}, + vkapi::kRead}}, + {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, + {}, + // The spec constants of the linear_q4gsw_{tiled,coop} kernels. The + // output width N is a spec constant for the coopMatStore address math. + {apply_bias, + K4_per_group, + num_groups, + graph.size_at(-1, output)}, + {group_size, weight_data, bias_data}, + resize_linear_qw_node)); + return true; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h new file mode 100644 index 00000000000..3d452c6bbed --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h @@ -0,0 +1,25 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +// Cooperative-matrix family path of et_vk.linear_q4gsw / et_vk.q4gsw_linear +// (CoopmatFamily.h). If the device belongs to a family and one of its rows +// applies to the op's build-time shape, adds the op's nodes and returns true; +// otherwise adds nothing and returns false, and the caller builds its usual +// nodes. `args` is the op's argument list: +// {input, weight, weight_scales, group_size, bias, output}. +bool try_add_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute From 68fcab55cbeae845b7155624b9cee914286b7874 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 16 Oct 2026 11:32:35 -0700 Subject: [PATCH 04/11] [ET-VK] dq8ca_q4gsw linear: cooperative-matrix family kernels 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 Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm --- .../dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl | 98 +++ .../dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml | 51 ++ .../linear_dq8ca_q4gsw_wmma_body.glslh | 559 ++++++++++++++++++ .../linear_dq8ca_q4gsw_wmma_rowmajor.glsl | 101 ++++ .../linear_dq8ca_q4gsw_wmma_rowmajor.yaml | 30 + ...inear_dq8ca_q4gsw_wmma_rowmajor_body.glslh | 491 +++++++++++++++ .../quantize_and_pack_4w_with_group_sums.glsl | 157 +++++ .../quantize_and_pack_4w_with_group_sums.yaml | 32 + .../tile_utils/linear_int8_input_block.glslh | 7 + .../graph/ops/impl/QuantizeDequantize.h | 8 + .../qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp | 7 + .../qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp | 303 ++++++++++ .../qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h | 26 + 13 files changed, 1870 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl new file mode 100644 index 00000000000..96c1915766e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl @@ -0,0 +1,98 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * 8da4w (dynamically quantized int8 activations, 4-bit group-symmetric + * weights) int8 cooperative-matrix linear for subgroup-16 devices with an + * 8x16x32 int8 MMA, used by the architecture-family path + * (impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp). Reads the 4h4w + * int8 activations of quantize_and_pack_4h4w_with_group_sums. Texel-wise B + * staging, multi-block A staging and an optional workgroup order + * (ORDER_NB_MAX). Template header only; the body is + * linear_dq8ca_q4gsw_wmma_body.glslh. + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: M % WG_TILE_M == 0, N % WG_TILE_N == 0, + * K % WG_TILE_K == 0, group_size % WG_TILE_K == 0, no bias, batch 1. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +// When the A staging thread map exactly covers the workgroup +// (A_ACTIVE_THREADS == WG_SIZE) the `a_active` guard is always true, but the +// driver compiler may not fold it because gl_LocalInvocationID.x's bound comes +// from a spec constant. Set A_MAP_FULL only for tiles where the equality holds. +$if A_MAP_FULL: + #define A_ALWAYS_ACTIVE + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings, in the argument order of the dispatch node: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +$if A_MULTI_BLOCK: + #define A_MULTI_BLOCK +const uint A_BLOCKS = ${A_BLOCKS}; + +// Workgroup order: 0 = dispatch order; > 0 = blocks of at most ORDER_NB_MAX +// N tiles, all M tiles of a block before the next block. +#define ORDER_NB_MAX ${ORDER_NB_MAX} + +#include "linear_dq8ca_q4gsw_wmma_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml new file mode 100644 index 00000000000..0143c0bfd32 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the architecture-family 8da4w int8 cooperative-matrix linear +# (4h4w int8 activations). Each one is referenced by a row in +# impl/gemm/linear/qw/CoopmatFamily.cpp, which also carries its tile geometry; +# the two must stay in sync. +# +# Name: txkgs +# m8 (MMA_M = 8), then i (dispatch order) or n (blocked order). + +linear_dq8ca_q4gsw_wmma: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + A_MAP_FULL: false + A_MULTI_BLOCK: false + A_BLOCKS: 1 + MMA_M: 8 + MMA_N: 16 + MMA_K: 32 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 16 + ORDER_NB_MAX: 0 + shader_variants: + # Intel Xe2: subgroup 16, int8 MMA 8x16x32, dispatch order. + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i_texture3d_texture2d_half + IO_STORAGE: texture3d + A_MAP_FULL: true + A_MULTI_BLOCK: true + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i_buffer_texture2d_half + A_MAP_FULL: true + A_MULTI_BLOCK: true + # Same tile, blocked workgroup order (weight-heavy projections). + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13_texture3d_texture2d_half + IO_STORAGE: texture3d + A_MAP_FULL: true + A_MULTI_BLOCK: true + ORDER_NB_MAX: 13 + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13_buffer_texture2d_half + A_MAP_FULL: true + A_MULTI_BLOCK: true + ORDER_NB_MAX: 13 diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh new file mode 100644 index 00000000000..9d3902b96c3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh @@ -0,0 +1,559 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_dq8ca_q4gsw_wmma; included by its template header +// (linear_dq8ca_q4gsw_wmma.glsl), which defines bindings, spec constants and +// tile constants. Must not contain template syntax. +// +// Per quantization group the int8 x int8 products accumulate in int32 +// coopmats; each group adds float(acc) * weight_scale into an fp32 total. The +// activation zero point and scale are applied once at the end: +// out = input_scale * (total - input_zp * SUM_g weight_scale_g * weight_sum_g) + +#ifndef LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH +#define LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// B rows are not padded (stride 4 words), and each thread writes contiguous +// shared-memory words (see the staging index setup below). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. +shared float wcorr_sh[WG_TILE_N]; + + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. float16_t-typed because +// coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +// All four components of a packed-weight texel at once: `shift` is 0 or 4 +// (nibble parity). Same arithmetic as widen_nibbles, per component. +uvec4 widen_nibbles4(const uvec4 w, const uint shift) { + const uvec4 nib = (w >> shift) & uvec4(0x0F0F0F0Fu); + const uvec4 p = nib ^ uvec4(0x08080808u); + const uvec4 sgn = p & uvec4(0x08080808u); + return p | (sgn * 0x1Eu); +} + +void main() { + // The tile this workgroup computes is a remap of its launch index (the grid + // is (N tiles, M tiles), launched x-fastest). ORDER_NB_MAX == 0: dispatch + // order. ORDER_NB_MAX > 0: blocks of at most ORDER_NB_MAX N tiles, evened + // out; all M tiles of a block run before the next block (N-fastest inside a + // block; the last block may be narrower). A bijection of the tiles. + const uint ord_nx = gl_NumWorkGroups.x; + const uint ord_lin = gl_WorkGroupID.y * ord_nx + gl_WorkGroupID.x; +#if ORDER_NB_MAX > 0 + const uint ord_ny = gl_NumWorkGroups.y; + const uint ord_nblk = (ord_nx + uint(ORDER_NB_MAX) - 1u) / uint(ORDER_NB_MAX); + const uint ord_nb = (ord_nx + ord_nblk - 1u) / ord_nblk; + const uint ord_full = ord_nb * ord_ny; + const uint ord_b = ord_lin / ord_full; + const uint ord_n0 = ord_b * ord_nb; + const uint ord_w = min(ord_nb, ord_nx - ord_n0); + const uint ord_r = ord_lin - ord_b * ord_full; + const uvec2 tileID = uvec2(ord_n0 + ord_r % ord_w, ord_r / ord_w); +#else + const uvec2 tileID = uvec2(ord_lin % ord_nx, ord_lin / ord_nx); +#endif + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; +#ifdef A_ALWAYS_ACTIVE + // A_ACTIVE_THREADS == WG_SIZE for this variant's tile, so every thread stages + // A and the guard is unconditionally true. + const bool a_active = true; +#else + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; +#endif + + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor shared uint (4 K-contiguous sign-extended int8) --- + // Texel-wise: a slot is (texel, nibble parity) and yields 4 shared words + // (the texel's 4 components = 4 output columns x 4 K). + const uint N8_PER_TILE = WG_TILE_N >> 3u; + const uint B_TOTAL_SLOTS = 2u * K_BLOCKS_PER_CHUNK * N8_PER_TILE; + // Round up; B_SLOT_ACTIVE masks the slots past the last one. + const uint B_SLOTS_PER_THREAD = (B_TOTAL_SLOTS + WG_SIZE - 1u) / WG_SIZE; +#define B_SLOT_ACTIVE(si) (B_TOTAL_SLOTS % WG_SIZE == 0u || gl_LocalInvocationID.x + (si) * WG_SIZE < B_TOTAL_SLOTS) + + // Loop-invariant staging index math, hoisted: every term below depends only + // on gl_LocalInvocationID.x, si, tile_n_start and tile_m_start, not on + // `chunk` or `group_i`. Only the k4 texel row varies with the chunk, by a + // constant stride. + uint b_lds_off[B_SLOTS_PER_THREAD]; // shared store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + // Coalesced writes: start from the contiguous shared-memory index this + // thread writes and derive which global element it needs, rather than + // deriving the shared address from the natural fetch grouping (which gives + // B_STRIDE_U32-strided writes). + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint s = gl_LocalInvocationID.x + si * WG_SIZE; + const uint par = s & 1u; + const uint tex = s >> 1u; + const uint k4_in_chunk = tex % K_BLOCKS_PER_CHUNK; + const uint n8_in_tile = tex / K_BLOCKS_PER_CHUNK; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + // Shared word of component 0; component c is c * B_STRIDE_U32 further. + b_lds_off[si] = slab_idx * B_SLAB_U32 + + (n8_in_tile * 8u + par * 4u) * B_STRIDE_U32 + k4_in_slab; + b_comp[si] = 0u; // unused: all four components are staged + b_par[si] = par * 4u; // nibble shift + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } + // A staging: same argument. base_row/slab/k_uint are all invariant. +#ifdef A_MULTI_BLOCK + // A_BLOCKS int8 blocks per thread (blocks tid + i * WG_SIZE). + uint a_lds_off[A_BLOCKS]; // slice-relative shared offset of the block's row 0 + uint a_glb[A_BLOCKS]; // global int8 block index at chunk 0 + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + const uint blk = gl_LocalInvocationID.x + ai * WG_SIZE; + const uint mb = blk / K_BLOCKS_PER_CHUNK; + const uint kb = blk % K_BLOCKS_PER_CHUNK; + a_lds_off[ai] = (kb / (MMA_K >> 2u)) * A_SLAB_U32 + + (mb * 4u) * A_STRIDE_U32 + + (kb % (MMA_K >> 2u)); + a_glb[ai] = ((tile_m_start >> 2u) + mb) * nblocks_x_A + kb; + } + ivec4 temp_A[A_BLOCKS]; +#else + const uint a_lds_off0 = + (a_k_block / (MMA_K >> 2u)) * A_SLAB_U32 + + (a_m_block * 4u) * A_STRIDE_U32 + + (a_k_block % (MMA_K >> 2u)); + const uint a_glb_row = ((tile_m_start >> 2u) + a_m_block) * nblocks_x_A; + + // Prefetch temp registers. + ivec4 temp_A; +#endif + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across all groups. The loop is + // prologue-only and replaces per-group wsum work inside the loop body. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are needed only after the group loop, so they are loaded there + // (not kept live across it). + + // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier + // here -- the main loop's first iteration barriers before reading slice 0). + if (a_active) { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + temp_A[ai] = t_packed_int8_input[a_glb[ai]]; + } +#else + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); + } + { + // store chunk 0 -> slice 0 +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[a_lds_off[ai] + m4i * A_STRIDE_U32] = uint(temp_A[ai][m4i]); + } + } +#else + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); + } +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uvec4 bw = widen_nibbles4(uvec4(temp_B[si]), b_par[si]); + [[unroll]] for (uint c = 0; c < 4u; ++c) { + Bsh_int8[b_lds_off[si] + c * B_STRIDE_U32] = bw[c]; + } + } + } + + // ========================================================= + // MAIN LOOP: nested groups x chunks (kept nested; flattening it with a + // conditional coopmat epilog has crashed a driver compiler at large + // spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + temp_A[ai] = t_packed_int8_input[a_glb[ai] + (chunkK_nxt >> 2u)]; + } +#else + temp_A = t_packed_int8_input[a_glb_row + (chunkK_nxt >> 2u) + a_k_block]; +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + a_lds_off[ai] + m4i * A_STRIDE_U32] = uint(temp_A[ai][m4i]); + } + } +#else + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); + } +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uvec4 bw = widen_nibbles4(uvec4(temp_B[si]), b_par[si]); + [[unroll]] for (uint c = 0; c < 4u; ++c) { + Bsh_int8[nxt_b + b_lds_off[si] + c * B_STRIDE_U32] = bw[c]; + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: --------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] is dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is a construct shader + // compilers have broken before -- check this first if the texture variants + // miscompile on a new device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl new file mode 100644 index 00000000000..35980ece92d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl @@ -0,0 +1,101 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * 8da4w int8 cooperative-matrix linear that reads row-major int8 activations + * (utils::kPackedInt8_4W, written by quantize_and_pack_4w_with_group_sums), + * used by the architecture-family path + * (impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp). The 4h4w layout + * is not row-major, so no coopMatLoad can address it; this kernel needs the + * row-major packer. Template header only; the body is + * linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh, which: + * - stages A as raw 16-byte vectors (uvec4 global load -> uvec4 shared + * store) instead of coopMatLoad(global) -> coopMatStore(shared); needs + * WG_TILE_M * WG_TILE_K / 16 to be a multiple of the workgroup size; + * - lets one weight texel fetch feed both nibble parities of a component + * (output columns c and c + 4); + * - stages the texture3d drain band in the A shared-memory tile, which is + * dead after the last MMA. + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, no bias, batch 1. M is rounded up by the grid. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// A is bound as a scalar int8_t array so that a coopMatLoad from it has a +// matching component type. +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings, in the argument order of the dispatch node: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// Row-major packed activations (row stride = K int8). +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +// uvec4 view of binding 2 (t_packed_int8_input) for 16-byte A staging. +// Spelled layout(std430, set...) on purpose: gen_vulkan_spv.py collects the +// descriptor list from every line matching ^layout\(set, so the usual spelling +// would add a descriptor. +layout(std430, set = 0, binding = 2) buffer restrict readonly t_packed_int8_input_v4Buffer { + uvec4 t_packed_int8_input_v4[]; +}; +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml new file mode 100644 index 00000000000..6351a1f65f0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml @@ -0,0 +1,30 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the row-major-activation 8da4w int8 cooperative-matrix linear. +# Each one is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, +# which also carries its tile geometry; the two must stay in sync. +# +# Name: txkgs +# mk32 (MMA_K = 32). + +linear_dq8ca_q4gsw_wmma_rowmajor: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 32 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + shader_variants: + # NVIDIA: subgroup 32, int8 MMA 16x16x32. + - NAME: linear_dq8ca_q4gsw_wmma_rowmajor_t128x128k64g24s32mk32_texture3d_texture2d_half diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh new file mode 100644 index 00000000000..f51484e34d1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh @@ -0,0 +1,491 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_dq8ca_q4gsw_wmma_rowmajor; included by its template header +// (linear_dq8ca_q4gsw_wmma_rowmajor.glsl), which defines bindings, spec +// constants and tile constants. Must not contain template syntax. +// +// Same arithmetic as linear_dq8ca_q4gsw_wmma_body.glslh (int32 accumulation +// per quantization group, fp32 total, zero-point correction applied once at +// the end); differs in the A layout (row-major int8, staged as raw 16-byte +// vectors) and in B staging (one thread per whole weight texel). +// +// A is read for whole WG_TILE_M row tiles, so the int8 activation buffer must +// cover M rounded up to WG_TILE_M rows. + +#ifndef LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH +#define LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// B rows are not padded (stride 4 words). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// MMA operand staging in two ping-pong slices. +const uint SLICES = 2u; +shared uvec4 Ash_int8[SLICES * ASH_SLICE_U32 / 4u]; +const uint ASH_UNIT = 4u; // u32 per Ash_int8 element +shared uint Bsh_int8[SLICES * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. +shared float wcorr_sh[WG_TILE_N]; + + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major, staged in Ash_int8 (dead after the last +// MMA; every drain iteration starts with a barrier). +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +// Compile-time check: the drain band (in u32) must fit in Ash_int8. +const int kCshFitsInAsh[(CSH_ROWS * WG_TILE_N / 2u <= SLICES * ASH_SLICE_U32) ? 1 : -1] = int[](0); +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in int8 elements (the binding is row-major int8), derived + // from nblocks_x_A so it matches the packer's `m_row * K4 + k4` addressing + // exactly (K % 4 == 0 makes it equal to K). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A tile counts (MMA_M x MMA_K tiles per chunk) --- + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor shared uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + // A slot is a whole texel (k4, n8) producing eight shared uints. Threads + // beyond the texel count of a chunk have no slot (b_on). + const uint B_TEXELS = B_TOTAL_SLOTS / 8u; + const uint B_SLOTS_PER_THREAD = (B_TEXELS + WG_SIZE - 1u) / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; + + // Loop-invariant B staging index math, hoisted out of the chunk loop. + uint b_lds_off[B_SLOTS_PER_THREAD]; // shared store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + bool b_on[B_SLOTS_PER_THREAD]; // this thread owns a texel in this slot + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint q = gl_LocalInvocationID.x + si * WG_SIZE; + const uint K4_PER_SLAB = MMA_K >> 2u; + const uint k4_in_slab = q % K4_PER_SLAB; + // Lanes run over k4 within a slab fastest, then slab, then the texel + // column block. + const uint slab_idx = (q / K4_PER_SLAB) % NUM_K_SLABS; + const uint n8_in_tile = (q / (K4_PER_SLAB * NUM_K_SLABS)) % N8_PER_TILE; + const uint k4_in_chunk = slab_idx * K4_PER_SLAB + k4_in_slab; + const uint n_col = n8_in_tile * 8u; // component c is column n_col + c, parity 1 is + 4 + b_lds_off[si] = slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab; + b_comp[si] = 0u; + b_on[si] = q < B_TEXELS; + b_par[si] = 0u; + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } + + // A staging thread map: 16-byte part fastest, then row, then K slab, so 8 + // lanes write 128 contiguous shared bytes and a subgroup reads whole + // 32-byte global sectors. + const uint A_PARTS = MMA_K / 16u; + const uint A_VECS = WG_TILE_M * WG_TILE_K / 16u; + const uint A_VECS_PER_THREAD = A_VECS / WG_SIZE; + const int kAVecsDivide[(A_VECS % WG_SIZE == 0u && A_VECS >= WG_SIZE) ? 1 : -1] = int[](0); + uint a_lds[A_VECS_PER_THREAD]; // Ash_int8 u32 offset within a slice + uint a_gv4[A_VECS_PER_THREAD]; // uvec4 index of this thread's chunk-0 data + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + const uint q = gl_LocalInvocationID.x + v * WG_SIZE; + const uint part = q % A_PARTS; + const uint row = (q / A_PARTS) % WG_TILE_M; + const uint slab = q / (A_PARTS * WG_TILE_M); + a_lds[v] = slab * A_SLAB_U32 + row * A_STRIDE_U32 + part * 4u; + a_gv4[v] = ((tile_m_start + row) * a_row_stride_i8 + slab * MMA_K + part * 16u) / 16u; + } + uvec4 temp_Av[A_VECS_PER_THREAD]; + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across all groups. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are needed only after the group loop, so they are loaded there. + + // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier + // here -- the main loop's first iteration barriers before reading slice 0). + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + temp_Av[v] = t_packed_int8_input_v4[a_gv4[v]]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); + } + } + { + // store chunk 0 -> slice 0 + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + Ash_int8[a_lds[v] / 4u] = temp_Av[v]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + Bsh_int8[b_lds_off[si] + 0u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 0u); + Bsh_int8[b_lds_off[si] + 4u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 1u); + Bsh_int8[b_lds_off[si] + 1u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 0u); + Bsh_int8[b_lds_off[si] + 5u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 1u); + Bsh_int8[b_lds_off[si] + 2u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 0u); + Bsh_int8[b_lds_off[si] + 6u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 1u); + Bsh_int8[b_lds_off[si] + 3u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 0u); + Bsh_int8[b_lds_off[si] + 7u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 1u); + } + } + } + + // ========================================================= + // MAIN LOOP: nested groups x chunks (kept nested; flattening it with a + // conditional coopmat epilog has crashed a driver compiler at large + // spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + temp_Av[v] = t_packed_int8_input_v4[a_gv4[v] + chunkK_nxt / 16u]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; + if (b_on[si]) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + (slab_a_base_u32 + row_a * A_STRIDE_U32) / ASH_UNIT, + A_STRIDE_U32 / ASH_UNIT, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + Ash_int8[(nxt_a + a_lds[v]) / 4u] = temp_Av[v]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + Bsh_int8[nxt_b + b_lds_off[si] + 0u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 4u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 1u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 5u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 2u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 6u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 3u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 7u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 1u); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into the drain band at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] is dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is a construct shader + // compilers have broken before -- check this first if the texture variants + // miscompile on a new device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards the drain band against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Ash_int8, + (warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j)) / (2u * ASH_UNIT), + WG_TILE_N / (2u * ASH_UNIT), + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4(unpackHalf2x16((base % 8u) != 0u ? Ash_int8[base / 8u].z : Ash_int8[base / 8u].x), + unpackHalf2x16((base % 8u) != 0u ? Ash_int8[base / 8u].w : Ash_int8[base / 8u].y))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl new file mode 100644 index 00000000000..25a48549cb3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl @@ -0,0 +1,157 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Row-major (utils::kPackedInt8_4W) counterpart of +// quantize_and_pack_4h4w_with_group_sums.glsl. +// +// Identical quantization and per-group sum reduction; only the output layout +// differs. The 4h4w version writes an ivec4 per (m4, k4) block whose component +// selects one of 4 rows -- as a scalar array that is +// index = m4*(4*K4) + k4*4 + r, which is not affine in the row index and so +// cannot be addressed by a cooperative-matrix load. This version writes each +// of the 4 quantized rows to its own row-major slot, +// t_packed_int8_input[m * K4 + k4], giving plain row-major int8 with 4 +// K-contiguous values per int32 and a row stride of K4. +// +// Consumed only by linear_dq8ca_q4gsw_wmma_rowmajor. Buffer output only. + +#version 450 core + +${define_required_extensions(INPUT_STORAGE, DTYPE)} +${define_required_extensions("texture3d", "int8")} + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#extension GL_EXT_integer_dot_product : require + +#define NUM_GROUPS_PER_WG ${NUM_GROUPS_PER_WG} +#define NUM_WORKERS_PER_GROUP ${NUM_WORKERS_PER_GROUP} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "w", "t_int8_input_sums", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8" if ZP_DTYPE_MODE == "zpint8" else DTYPE, "texture3d")} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +shared ivec4 shared_sums[NUM_GROUPS_PER_WG][NUM_WORKERS_PER_GROUP]; + +#define TILE_M4 1 +#define TILE_K4 1 + +#define TILE_M 4 + +// The shared helper's write_block() assumes an ivec4-typed output resource; +// this shader's output is a scalar int array, so opt out and write our own. +#define SKIP_BLOCK_WRITE_HELPERS +#include "gemm/tile_utils/linear_int8_input_block.glslh" + +// Scatters the 4 rows of a quantized block into row-major slots. Unlike the +// 4h4w layout, kPackedInt8_4W has NO row padding (outer_dim_align == 1), so +// the tail rows of a workgroup whose m4 block runs past M must be dropped -- +// hence the explicit bound check that write_block() does not need. +void write_block_rowmajor( + const Int8InputBlock block, + const int k4, + const int m, + const int K4, + const int M) { + for (int row = 0; row < 4; ++row) { + const int m_row = m + row; + if (m_row < M) { + t_packed_int8_input[m_row * K4 + k4] = block.data[row]; + } + } +} +#include "gemm/tile_utils/linear_int8_input_scales_zps_load.glslh" +#include "gemm/tile_utils/linear_fp_input_tile_load.glslh" + +void main() { + const int group_idx = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + const int worker_id = int(gl_LocalInvocationID.z); + const int group_offset = int(gl_LocalInvocationID.x); + + const int K = input_sizes.x; + const int M = input_sizes.y; + + // K4 and M4 represent the number of blocks in each dimension. + const int K4 = div_up_4(K); + const int M4 = div_up_4(M); + + const int num_groups = K4 / K4_per_group; + + if (group_idx >= num_groups || m4 >= M4) { + return; + } + + const int start_k4 = group_idx * K4_per_group + worker_id; + const int end_k4 = (group_idx + 1) * K4_per_group; + + Int8InputScales input_scales; + Int8InputZeroPoints input_zps; + load_int8_input_scales_and_zps(input_scales, input_zps, m4); + + // row of the input tensor to start loading from + const int m = mul_4(m4); + + FPInputTile in_tile; + Int8InputBlock packed_block; + + ivec4 local_sum = ivec4(0, 0, 0, 0); + const int packed_ones = 0x01010101; + + for (int k4 = start_k4; k4 < end_k4; k4 += NUM_WORKERS_PER_GROUP) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + quantize_and_pack(packed_block, in_tile, input_scales, input_zps); + + // Sum the quantized values in the block + [[unroll]] for (int m = 0; m < TILE_M; m++) { + local_sum[m] += dotPacked4x8AccSatEXT( + packed_block.data[m], packed_ones, local_sum[m]); + } + write_block_rowmajor(packed_block, k4, m, K4, M); + } + + shared_sums[group_offset][worker_id] = local_sum; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_GROUP / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_sums[group_offset][worker_id] = + shared_sums[group_offset][worker_id] + + shared_sums[group_offset][worker_id + i]; + } + memoryBarrierShared(); + barrier(); + } + + if (worker_id == 0) { + t_int8_input_sums[group_idx * M4 + m4] = shared_sums[group_offset][0]; + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml new file mode 100644 index 00000000000..dfa6f7313a5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml @@ -0,0 +1,32 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Row-major (utils::kPackedInt8_4W) counterpart of +# quantize_and_pack_4h4w_with_group_sums: same quantization, per-group sum +# reduction and workgroup shapes (_o2w32 / _o4w16); only the output layout +# differs. Consumed by linear_dq8ca_q4gsw_wmma_rowmajor, whose cooperative +# matrix A staging needs row-major int8 activations. Buffer output only. + +quantize_and_pack_4w_with_group_sums: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + NUM_GROUPS_PER_WG: 2 + NUM_WORKERS_PER_GROUP: 32 + ZP_DTYPE_MODE: zpint8 + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + ZP_DTYPE_MODE: + - VALUE: zpint8 + - VALUE: zpinherit + shader_variants: + - NAME: quantize_and_pack_4w_with_group_sums_o2w32_buffer_texture3d + - NAME: quantize_and_pack_4w_with_group_sums_o4w16_buffer_texture3d + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh index 8f19418cd19..2e0f12f95bb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh @@ -69,6 +69,12 @@ void quantize_and_pack( } } +// The 4h4w block writer assumes t_packed_int8_input is an ivec4-typed +// (block-packed) resource and would not compile against a scalar int array. +// A shader that reuses the quantize helpers above but writes another layout +// (quantize_and_pack_4w_with_group_sums.glsl) defines SKIP_BLOCK_WRITE_HELPERS +// before including this file. +#ifndef SKIP_BLOCK_WRITE_HELPERS void write_block( const Int8InputBlock block, const int block_x, @@ -80,5 +86,6 @@ void write_block( imageStore(t_packed_int8_input, ivec3(block_x, block_y, 0), block.data); #endif // OUTPUT_BUFFER } +#endif // SKIP_BLOCK_WRITE_HELPERS #endif // LINEAR_INT8_INPUT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h index 96e9cc7c1d3..2adf41469da 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h @@ -34,6 +34,14 @@ void add_quantize_and_pack_4h4w_node( const ValueRef packed_int_input, const ValueRef group_size); +// Global work grid of the *_with_group_sums packers (_o2w32 / _o4w16 +// variants); empty for GEMV inputs. +GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args); + void add_quantize_and_pack_4h4w_with_group_sums_node( ComputeGraph& graph, const QuantizationConfig& input_quant_config, diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp index dd13ec2ebb5..a1cfbe7c0ef 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace vkcompute { @@ -153,6 +154,12 @@ void add_linear_dqa_qw_node( void linear_dq8ca_q4gsw( ComputeGraph& graph, const std::vector& args) { + // Devices of a cooperative-matrix family (CoopmatFamily.h) take the + // family's kernels where one of its rows applies. + if (try_add_dq8ca_q4gsw_linear_wmma_node(graph, args)) { + return; + } + int32_t idx = 0; const ValueRef fp_input = args.at(idx++); const ValueRef input_scale = args.at(idx++); diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp new file mode 100644 index 00000000000..6800c6ceb2a --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp @@ -0,0 +1,303 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vkcompute { + +namespace { + +// Dispatch node args: +// {{output}, {fp_input, packed_int_input, int_input_sums, input_scale, +// input_zp, packed_weight, weight_sums, weight_scales, bias}} +// resize_args: {group_size, weight_data, bias_data} +WmmaShape dq8ca_runtime_shape( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_int_input = args.at(1).refs.at(1); + const ValueRef packed_weight = args.at(1).refs.at(5); + + const std::vector out_sizes = graph->sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kDq8caQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph->sizes_of(fp_input)); + s.group_size = graph->extract_scalar(resize_args.at(0)); + for (int64_t d = 0; d < graph->dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(graph, fp_input); + s.has_bias = !graph->val_is_none(resize_args.at(2)); + s.half = graph->dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph->storage_type_of(fp_input)); + s.output = to_wmma_storage(graph->storage_type_of(output)); + s.weight = to_wmma_storage(graph->storage_type_of(packed_weight)); + s.io_width_packed = graph->packed_dim_of(output) == WHCN::kWidthDim && + graph->packed_dim_of(fp_input) == WHCN::kWidthDim; + s.int8_layout = graph->estimate_memory_layout_of(packed_int_input) == + utils::kPackedInt8_4W + ? WmmaInt8Layout::kRowMajor + : WmmaInt8Layout::k4H4W; + return s; +} + +// args: {input, input_scale, input_zp, weight, weight_sums, weight_scales, +// group_size, bias, output} +WmmaShape dq8ca_build_time_shape( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef output = args.at(8); + const std::vector out_sizes = graph.sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kDq8caQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph.sizes_of(fp_input)); + s.group_size = graph.extract_scalar(args.at(6)); + for (int64_t d = 0; d < graph.dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(&graph, fp_input); + s.has_bias = !graph.val_is_none(args.at(7)); + s.half = graph.dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph.storage_type_of(fp_input)); + s.output = to_wmma_storage(graph.storage_type_of(output)); + s.weight = to_wmma_storage(predicted_q4_weight_storage(graph, args.at(3))); + s.io_width_packed = graph.packed_dim_of(output) == WHCN::kWidthDim && + graph.packed_dim_of(fp_input) == WHCN::kWidthDim; + s.int8_layout = WmmaInt8Layout::kAny; + return s; +} + +vkapi::ShaderInfo pick_dq8ca_q4gsw_linear_wmma_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef input_zp = args.at(1).refs.at(4); + const ValueRef packed_weight = args.at(1).refs.at(5); + + const CoopmatDeviceKey device = coopmat_device_key(graph); + WmmaShape shape = dq8ca_runtime_shape(graph, args, resize_args); + std::optional choice = select_wmma_kernel(device, shape); + if (!choice.has_value() && shape.int8_layout == WmmaInt8Layout::kRowMajor && + !shape.gemv) { + // The row-major activation layout was chosen at build time and, of the + // M > 1 kernels, only the row-major family kernels read it, so keep one + // for unaligned M: the grid rounds M up. Decode (M == 1) takes the + // linear_dq8ca_q4gsw_coop kernel below, which reads the fp activations + // (the quantize/pack node dispatches nothing at M == 1). + shape.ignore_alignment = true; + choice = select_wmma_kernel(device, shape); + VK_CHECK_COND( + choice.has_value(), "no row-major dq8ca_q4gsw kernel for this shape"); + } + + std::string kernel_name; + if (choice.has_value()) { + kernel_name = choice->kernel_base; + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + } else { + // The kernels for the 4h4w layout (as in pick_linear_dqa_qw_shader). + kernel_name = "linear_dq8ca_q4gsw"; + kernel_name += is_gemv(graph, fp_input) ? "_coop" : "_tiled"; + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zp)); + } + return VK_KERNEL_FROM_STR(kernel_name); +} + +// Row-major (utils::kPackedInt8_4W) counterpart of +// add_quantize_and_pack_4h4w_with_group_sums_node: same quantization, group +// sums and workgroup shape. +void add_quantize_and_pack_4w_with_group_sums_node( + ComputeGraph& graph, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size) { + const int32_t group_size_val = graph.extract_scalar(group_size); + const int32_t blocks_per_group = utils::div_up(group_size_val, int32_t(4)); + + std::string shader_name = "quantize_and_pack_4w_with_group_sums"; + shader_name += group_size_val >= 128 ? "_o2w32" : "_o4w16"; + add_storage_type_suffix(shader_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(shader_name, graph.storage_type_of(fp_input)); + add_dtype_suffix(shader_name, graph.dtype_of(fp_input)); + add_zp_dtype_mode_suffix(shader_name, graph.dtype_of(packed_input_zps)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(shader_name), + pick_quantize_and_pack_4h4w_with_group_sums_gwg, + pick_required_lwg, + {{{packed_int_input, int_input_sums}, vkapi::kWrite}, + {{fp_input, packed_input_scales, packed_input_zps}, vkapi::kRead}}, + {graph.sizes_ubo(fp_input)}, + {}, + {blocks_per_group}, + {group_size})); +} + +} // namespace + +bool try_add_dq8ca_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef input_scale = args.at(1); + const ValueRef input_zp = args.at(2); + const ValueRef weight_data = args.at(3); + const ValueRef weight_sums_data = args.at(4); + const ValueRef weight_scales_data = args.at(5); + const ValueRef group_size = args.at(6); + const ValueRef bias_data = args.at(7); + const ValueRef output = args.at(8); + + if (graph.dtype_of(fp_input) != vkapi::kHalf) { + return false; + } + // The decode and fallback kernels of this path use int8 dot products. + if (!graph.can_use_int8_dot_product()) { + return false; + } + const CoopmatDeviceKey device = coopmat_device_key(&graph); + const WmmaShape shape = dq8ca_build_time_shape(graph, args); + // Take the op over only if a row applies to its build-time shape; otherwise + // it stays entirely on the existing path. The kernel chosen here also fixes + // the int8 activation layout. + const std::optional build_choice = + select_wmma_kernel(device, shape); + if (!build_choice.has_value()) { + return false; + } + const bool rowmajor = build_choice->rowmajor_a; + + const int64_t group_size_val = graph.extract_scalar(group_size); + const QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true); + const QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + const ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + const ValueRef packed_weight_sums = prepack_standard( + graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); + + TmpTensor dummy_bias( + &graph, {}, graph.dtype_of(output), utils::kBuffer, utils::kWidthPacked); + ValueRef packed_bias = dummy_bias.vref; + uint32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + apply_bias = 1; + } + + ValueRef packed_input_scale = input_scale; + ValueRef packed_input_zp = input_zp; + if (graph.val_is_tref(input_scale)) { + VK_CHECK_COND(graph.val_is_tref(input_zp)); + packed_input_scale = prepack_standard( + graph, input_scale, utils::kTexture3D, utils::kWidthPacked); + packed_input_zp = prepack_standard( + graph, input_zp, utils::kTexture3D, utils::kWidthPacked); + } + + TmpTensor packed_int_input( + &graph, + graph.sizes_of(fp_input), + vkapi::kInt8x4, + utils::kBuffer, + rowmajor ? utils::kPackedInt8_4W : utils::kPackedInt8_4H4W); + + // Per-group int8 input sums, indexed as ivec4[group_idx * M4 + m4]. + const int64_t num_groups = graph.size_at(-2, weight_scales_data); + const int64_t M4 = utils::div_up(shape.M, int64_t(4)); + TmpTensor int_input_sums( + &graph, + {num_groups * M4 * 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + if (rowmajor) { + add_quantize_and_pack_4w_with_group_sums_node( + graph, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } else { + add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } + + const int32_t K4_per_group = + utils::div_up(static_cast(group_size_val), int32_t(4)); + const int32_t num_groups_arg = graph.size_at(-1, fp_input) / + static_cast(group_size_val); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_dq8ca_q4gsw_linear_wmma_shader, + wmma_linear_gwg, + quantized_linear_lwg, + {{output, vkapi::kWrite}, + {{fp_input, + packed_int_input.vref, + int_input_sums.vref, + packed_input_scale, + packed_input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, + {}, + // The spec constants of add_linear_dqa_qw_node. + {apply_bias, + K4_per_group, + num_groups_arg, + graph.size_at(-1, output)}, + {group_size, weight_data, bias_data}, + resize_linear_qw_node)); + return true; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h new file mode 100644 index 00000000000..ce94047574e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h @@ -0,0 +1,26 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +// Cooperative-matrix family path of et_vk.linear_dq8ca_q4gsw +// (CoopmatFamily.h). If the device belongs to a family and one of its rows +// applies to the op's build-time shape, adds the whole op (activation +// quantize/pack and the int8 cooperative-matrix linear) and returns true; +// otherwise adds nothing and returns false, and the caller builds its usual +// nodes. `args`: {input, input_scale, input_zp, weight, weight_sums, +// weight_scales, group_size, bias, output}. +bool try_add_dq8ca_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute From c20abb6d541be0e79155086bcf71d8d7304e2934 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 16 Oct 2026 11:35:03 -0700 Subject: [PATCH 05/11] [ET-VK] Tests for the cooperative-matrix family kernels - 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 Claude-Session: https://claude.ai/code/session_016AnsG4u2n3shZugcMcyUFm --- .../vulkan/test/custom_ops/CMakeLists.txt | 48 ++ .../test/custom_ops/coopmat_family_test.cpp | 440 ++++++++++++++++++ backends/vulkan/test/custom_ops/targets.bzl | 10 + .../q4gsw_linear/test_sets/wmma_sets.cpp | 175 +++++++ 4 files changed, 673 insertions(+) create mode 100644 backends/vulkan/test/custom_ops/CMakeLists.txt create mode 100644 backends/vulkan/test/custom_ops/coopmat_family_test.cpp create mode 100644 backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp diff --git a/backends/vulkan/test/custom_ops/CMakeLists.txt b/backends/vulkan/test/custom_ops/CMakeLists.txt new file mode 100644 index 00000000000..21fa561fb8d --- /dev/null +++ b/backends/vulkan/test/custom_ops/CMakeLists.txt @@ -0,0 +1,48 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Host-only tests that need neither a GPU nor the Vulkan backend library. +# +# cmake backends/vulkan/test/custom_ops -Bcmake-out/backends/vulkan/test/custom_ops +# cmake --build cmake-out/backends/vulkan/test/custom_ops +# ctest --test-dir cmake-out/backends/vulkan/test/custom_ops + +cmake_minimum_required(VERSION 3.19) +project(vulkan_custom_ops_tests CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) +endif() + +if(NOT TARGET GTest::gtest_main) + set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + add_subdirectory( + ${EXECUTORCH_ROOT}/third-party/googletest + ${CMAKE_CURRENT_BINARY_DIR}/googletest EXCLUDE_FROM_ALL + ) +endif() + +enable_testing() + +set(VULKAN_OPS_PATH ${EXECUTORCH_ROOT}/backends/vulkan/runtime/graph/ops) + +add_executable( + vulkan_coopmat_family_test + ${CMAKE_CURRENT_SOURCE_DIR}/coopmat_family_test.cpp + ${VULKAN_OPS_PATH}/impl/gemm/linear/qw/CoopmatFamily.cpp +) +target_include_directories( + vulkan_coopmat_family_test PRIVATE ${EXECUTORCH_ROOT}/.. +) +target_compile_definitions( + vulkan_coopmat_family_test + PRIVATE COOPMAT_FAMILY_TEST_GLSL_DIR="${VULKAN_OPS_PATH}/glsl" +) +target_link_libraries(vulkan_coopmat_family_test PRIVATE GTest::gtest_main) +add_test(vulkan_coopmat_family_test vulkan_coopmat_family_test) diff --git a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp new file mode 100644 index 00000000000..5a982bc1cff --- /dev/null +++ b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp @@ -0,0 +1,440 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Host-only tests of the cooperative-matrix family key and the family kernel +// selection (CoopmatFamily.h). No GPU needed. + +#include + +#include + +#include +#include +#include +#include + +namespace vkcompute { +namespace { + +using T = CoopmatComponentType; + +// VkDriverId values used below. +constexpr uint32_t kDriverMesaRadv = 3; +constexpr uint32_t kDriverIntelWindows = 5; +constexpr uint32_t kDriverMesaLlvmpipe = 13; +constexpr uint32_t kDriverMesaNvk = 24; + +CoopmatDeviceKey +base_key(const uint32_t sg, const uint32_t sg_min, const uint32_t sg_max) { + CoopmatDeviceKey d; + d.subgroup_size = sg; + d.min_subgroup_size = sg_min; + d.max_subgroup_size = sg_max; + d.subgroup_size_control = true; + d.coopmat = true; + d.int8_coopmat = true; + d.max_shared_bytes = 65536; + return d; +} + +// Keys as Xe2 (Arc B580 / Arc Pro B70) and NVIDIA (Ampere / Ada) devices +// report them. +CoopmatDeviceKey xe2_key(const uint32_t driver = kDriverIdIntelMesa) { + CoopmatDeviceKey d = base_key(32, 8, 32); + d.vendor_id = kVendorIdIntel; + d.driver_id = driver; + d.coopmat_shapes = { + {8, 16, 16, T::kFp16, T::kFp16}, + {8, 16, 16, T::kFp16, T::kFp32}, + {8, 16, 32, T::kSint8, T::kSint32}}; + return d; +} + +CoopmatDeviceKey nvidia_key( + const uint32_t driver = kDriverIdNvidiaProprietary) { + CoopmatDeviceKey d = base_key(32, 32, 32); + d.vendor_id = kVendorIdNvidia; + d.driver_id = driver; + d.coopmat_shapes = { + {16, 16, 16, T::kFp16, T::kFp16}, + {16, 8, 16, T::kFp16, T::kFp16}, + {16, 16, 16, T::kFp16, T::kFp32}, + {16, 16, 32, T::kSint8, T::kSint32}}; + return d; +} + +// Mesa llvmpipe (lavapipe): 8x8x8 shapes only, fixed subgroup 8. +CoopmatDeviceKey llvmpipe_key() { + CoopmatDeviceKey d = base_key(8, 8, 8); + d.vendor_id = 0x10005; + d.driver_id = kDriverMesaLlvmpipe; + d.max_shared_bytes = 32768; + d.coopmat_shapes = { + {8, 8, 8, T::kFp16, T::kFp16}, + {8, 8, 8, T::kFp16, T::kFp32}, + {8, 8, 8, T::kOther, T::kOther}, // u8 x u8 -> u32 + {8, 8, 8, T::kSint8, T::kSint32}}; + return d; +} + +// AMD RDNA3 under RADV: no family yet. +CoopmatDeviceKey radv_key() { + CoopmatDeviceKey d = base_key(64, 32, 64); + d.vendor_id = 0x1002; + d.driver_id = kDriverMesaRadv; + d.coopmat_shapes = { + {16, 16, 16, T::kFp16, T::kFp16}, + {16, 16, 16, T::kFp16, T::kFp32}, + {16, 16, 16, T::kSint8, T::kSint32}}; + return d; +} + +WmmaShape prefill( + const WmmaOp op, + const int64_t M, + const int64_t K, + const int64_t N, + const WmmaStorage io) { + WmmaShape s; + s.op = op; + s.M = M; + s.K = K; + s.N = N; + s.group_size = 128; + s.half = true; + s.input = io; + s.output = io; + s.weight = WmmaStorage::kTexture2D; + s.io_width_packed = true; + return s; +} + +std::string selected(const CoopmatDeviceKey& d, const WmmaShape& s) { + const std::optional c = select_wmma_kernel(d, s); + return c.has_value() ? c->kernel_base : ""; +} + +// (K, N) of every 4-bit linear in Llama 3.2 1B / 3B and Llama 3.1 8B. +const int64_t kShapes[][2] = { + {2048, 2048}, + {2048, 512}, + {2048, 8192}, + {8192, 2048}, + {3072, 3072}, + {3072, 1024}, + {3072, 8192}, + {8192, 3072}, + {4096, 4096}, + {4096, 1024}, + {4096, 14336}, + {14336, 4096}, +}; + +const std::string kQ = "linear_q4gsw_wmma_"; +const std::string kD = "linear_dq8ca_q4gsw_wmma_"; + +using Rule = std::function; + +// The intended kernel per family, restated independently of the tables. +std::string xe2_rule(const WmmaShape& s) { + if (s.op == WmmaOp::kQ4gsw) { + return kQ + + (s.output == WmmaStorage::kTexture3D ? "t128x128k16g82s16m8flib" + : "t128x128k16g44s16m8fli"); + } + return kD + + (s.K * s.N >= 24000000 ? "t128x128k64g84s16m8n13" + : "t128x128k64g84s16m8i"); +} + +std::string nvidia_rule(const WmmaShape& s) { + if (s.output != WmmaStorage::kTexture3D) { + return ""; + } + if (s.op == WmmaOp::kQ4gsw) { + return kQ + + (s.K > 8192 ? "bx_t128x128k32g42s32f32c" : "t256x128k16g42s32bt"); + } + return kD + "rowmajor_t128x128k64g24s32mk32"; +} + +std::string none_rule(const WmmaShape&) { + return ""; +} + +void check_llama_shapes( + const char* label, + const CoopmatDeviceKey& d, + const Rule& rule) { + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + for (const auto& kn : kShapes) { + for (const WmmaStorage io : + {WmmaStorage::kTexture3D, WmmaStorage::kBuffer}) { + const WmmaShape s = prefill(op, 2048, kn[0], kn[1], io); + EXPECT_EQ(selected(d, s), rule(s)) + << label << " op=" << static_cast(op) << " K=" << kn[0] + << " N=" << kn[1] + << " io=" << (io == WmmaStorage::kBuffer ? "buffer" : "texture3d"); + } + } + } +} + +TEST(CoopmatFamilyTest, FamilyKey) { + EXPECT_EQ(coopmat_family(xe2_key()), CoopmatFamily::kXe2); + EXPECT_EQ(coopmat_family(nvidia_key()), CoopmatFamily::kNvidia); + + // Same silicon under another driver: no family. + EXPECT_EQ(coopmat_family(xe2_key(kDriverIntelWindows)), CoopmatFamily::kNone); + EXPECT_EQ(coopmat_family(nvidia_key(kDriverMesaNvk)), CoopmatFamily::kNone); + + // The int8 shape is part of the key. + CoopmatDeviceKey no_int8 = nvidia_key(); + no_int8.int8_coopmat = false; + EXPECT_EQ(coopmat_family(no_int8), CoopmatFamily::kNone); + CoopmatDeviceKey no_int8_shape = xe2_key(); + no_int8_shape.coopmat_shapes.pop_back(); + EXPECT_EQ(coopmat_family(no_int8_shape), CoopmatFamily::kNone); + + // Xe2 needs subgroup size 16 to be available. + CoopmatDeviceKey xe2_sg8 = xe2_key(); + xe2_sg8.subgroup_size = 8; + xe2_sg8.min_subgroup_size = 8; + xe2_sg8.max_subgroup_size = 8; + EXPECT_EQ(coopmat_family(xe2_sg8), CoopmatFamily::kNone); + CoopmatDeviceKey xe2_no_control = xe2_key(); + xe2_no_control.subgroup_size_control = false; + EXPECT_EQ(coopmat_family(xe2_no_control), CoopmatFamily::kNone); + + // No family for vendors without validated kernels. + EXPECT_EQ(coopmat_family(radv_key()), CoopmatFamily::kNone); + + // No cooperative matrix at all. + CoopmatDeviceKey no_coopmat = xe2_key(); + no_coopmat.coopmat = false; + EXPECT_EQ(coopmat_family(no_coopmat), CoopmatFamily::kNone); +} + +TEST(CoopmatFamilyTest, LlvmpipeHasNoFamily) { + const CoopmatDeviceKey d = llvmpipe_key(); + EXPECT_EQ(coopmat_family(d), CoopmatFamily::kNone); + EXPECT_TRUE(coopmat_family_rows(coopmat_family(d)).empty()); + check_llama_shapes("llvmpipe", d, none_rule); + // Shapes aligned to every family tile. + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + for (const WmmaStorage io : + {WmmaStorage::kTexture3D, WmmaStorage::kBuffer}) { + EXPECT_FALSE( + select_wmma_kernel(d, prefill(op, 256, 256, 256, io)).has_value()); + } + } +} + +TEST(CoopmatFamilyTest, UnmatchedKeysSelectNothing) { + check_llama_shapes("radv", radv_key(), none_rule); + check_llama_shapes("xe2-windows", xe2_key(kDriverIntelWindows), none_rule); + check_llama_shapes("nvk", nvidia_key(kDriverMesaNvk), none_rule); +} + +TEST(CoopmatFamilyTest, Xe2Rows) { + check_llama_shapes("xe2", xe2_key(), xe2_rule); + // The weight-heavy rule: K * N >= 24e6. + const CoopmatDeviceKey d = xe2_key(); + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 11776, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8n13"); // 24.1e6 + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 11648, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8i"); // 23.9e6 + // Any aligned shape, not only the Llama projections. + EXPECT_EQ( + selected(d, prefill(WmmaOp::kQ4gsw, 128, 256, 384, WmmaStorage::kBuffer)), + kQ + "t128x128k16g44s16m8fli"); +} + +TEST(CoopmatFamilyTest, NvidiaRows) { + check_llama_shapes("nvidia", nvidia_key(), nvidia_rule); + const CoopmatDeviceKey d = nvidia_key(); + struct Edge { + int64_t M, K, N; + const char* want; // nullptr: existing kernels + }; + const Edge edges[] = { + // M not a multiple of 256: the 128-row tiles. + {384, 2048, 2048, "t128x128k16g22s32"}, + // M = 256 at N % 2048 != 0: not the 256-row tile. + {256, 3072, 1024, "t128x128k16g22s32"}, + {256, 2048, 8192, "t256x128k16g42s32bt"}, + // K > 8192 with M not a multiple of 256: fp32 accumulation. + {384, 14336, 4096, "bx_t128x128k32g42s32f32c"}, + // Shapes outside the validated projections. + {2048, 2048, 128256, nullptr}, + {2048, 2048, 4096, nullptr}, + // M beyond the validated prefill length. + {4096, 2048, 2048, nullptr}, + }; + for (const Edge& e : edges) { + const WmmaShape s = + prefill(WmmaOp::kQ4gsw, e.M, e.K, e.N, WmmaStorage::kTexture3D); + EXPECT_EQ(selected(d, s), e.want == nullptr ? "" : kQ + e.want) + << "M=" << e.M << " K=" << e.K << " N=" << e.N; + } + // Group size other than 128. + WmmaShape g64 = + prefill(WmmaOp::kQ4gsw, 2048, 2048, 2048, WmmaStorage::kTexture3D); + g64.group_size = 64; + EXPECT_EQ(selected(d, g64), ""); +} + +TEST(CoopmatFamilyTest, ShapesThatKeepTheExistingKernels) { + for (const CoopmatDeviceKey& d : {xe2_key(), nvidia_key()}) { + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + const WmmaShape s = + prefill(op, 2048, 4096, 4096, WmmaStorage::kTexture3D); + ASSERT_FALSE(selected(d, s).empty()); + + WmmaShape unaligned = s; + unaligned.M = 1304; + WmmaShape gemv = s; + gemv.M = 1; + gemv.gemv = true; + WmmaShape bias = s; + bias.has_bias = true; + WmmaShape fp32 = s; + fp32.half = false; + WmmaShape mixed = s; + mixed.input = WmmaStorage::kBuffer; + WmmaShape batched = s; + batched.batch = 2; + WmmaShape buffer_weight = s; + buffer_weight.weight = WmmaStorage::kBuffer; + WmmaShape channels_packed = s; + channels_packed.io_width_packed = false; + for (const WmmaShape* x : + {&unaligned, + &gemv, + &bias, + &fp32, + &mixed, + &batched, + &buffer_weight, + &channels_packed}) { + EXPECT_FALSE(select_wmma_kernel(d, *x).has_value()) + << coopmat_family_name(coopmat_family(d)) + << " op=" << static_cast(op); + } + } + } +} + +TEST(CoopmatFamilyTest, Int8Layout) { + const CoopmatDeviceKey d = nvidia_key(); + WmmaShape s = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D); + const std::optional aligned = select_wmma_kernel(d, s); + ASSERT_TRUE(aligned.has_value()); + EXPECT_TRUE(aligned->rowmajor_a); + + // A row-major op keeps its kernel for unaligned M when alignment is waived. + s.int8_layout = WmmaInt8Layout::kRowMajor; + s.M = 1304; + s.ignore_alignment = true; + const std::optional forced = select_wmma_kernel(d, s); + ASSERT_TRUE(forced.has_value()); + EXPECT_EQ(forced->kernel_base, aligned->kernel_base); + + // 4h4w activations never get a row-major kernel, and vice versa. + WmmaShape h = s; + h.M = 2048; + h.ignore_alignment = false; + h.int8_layout = WmmaInt8Layout::k4H4W; + EXPECT_FALSE(select_wmma_kernel(d, h).has_value()); + WmmaShape x = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D); + x.int8_layout = WmmaInt8Layout::kRowMajor; + EXPECT_FALSE(select_wmma_kernel(xe2_key(), x).has_value()); +} + +TEST(CoopmatFamilyTest, SharedMemoryLimit) { + // The texture3d drain of the Xe2 4w tile needs 2 * 8 * 128 * 2 = 4 KiB of + // shared memory on top of the A/B tiles. + CoopmatDeviceKey d = xe2_key(); + d.max_shared_bytes = 4096; + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D)), + ""); + EXPECT_EQ( + selected( + d, prefill(WmmaOp::kQ4gsw, 2048, 4096, 4096, WmmaStorage::kBuffer)), + kQ + "t128x128k16g44s16m8fli"); +} + +TEST(CoopmatFamilyTest, TileForKernel) { + for (const CoopmatFamily f : {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(f)) { + const std::optional tile = wmma_tile_for_kernel( + std::string(row.kernel_base) + "_texture3d_texture2d_half"); + ASSERT_TRUE(tile.has_value()) << row.kernel_base; + EXPECT_EQ(tile->m, row.tile.m) << row.kernel_base; + EXPECT_EQ(tile->n, row.tile.n) << row.kernel_base; + EXPECT_EQ(tile->wg_size(), row.tile.wg_size()) << row.kernel_base; + } + } + // The existing kernels are not family kernels. + EXPECT_FALSE( + wmma_tile_for_kernel("linear_q4gsw_coopmat_buffer_texture2d_half") + .has_value()); + EXPECT_FALSE( + wmma_tile_for_kernel("linear_q4gsw_tiled_texture3d_texture2d_half") + .has_value()); + EXPECT_FALSE(wmma_tile_for_kernel("linear_q4gsw_wmma_t128x128k16g82s16m8fli") + .has_value()); +} + +#ifdef COOPMAT_FAMILY_TEST_GLSL_DIR +std::string read_file(const std::string& path) { + std::ifstream in(path); + std::stringstream ss; + ss << in.rdbuf(); + return ss.str(); +} + +// Every kernel name a row can produce must be a variant in its yaml. +TEST(CoopmatFamilyTest, EveryRowHasShaderVariants) { + const std::string dir = + std::string(COOPMAT_FAMILY_TEST_GLSL_DIR) + "/gemm/linear/qw/"; + const std::string yamls = read_file(dir + "q4gsw/linear_q4gsw_wmma.yaml") + + read_file(dir + "q4gsw/linear_q4gsw_wmma_bx.yaml") + + read_file(dir + "dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml") + + read_file(dir + "dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml"); + if (yamls.empty()) { + GTEST_SKIP() << "shader sources not found under " << dir; + } + const std::pair combos[] = { + {kWmmaTex3dTex2d, "_texture3d_texture2d_half"}, + {kWmmaBufTex2d, "_buffer_texture2d_half"}}; + for (const CoopmatFamily f : {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(f)) { + for (const auto& c : combos) { + if (row.storages & c.first) { + const std::string name = std::string(row.kernel_base) + c.second; + EXPECT_NE(yamls.find("NAME: " + name + "\n"), std::string::npos) + << "no shader variant " << name; + } + } + } + } +} +#endif // COOPMAT_FAMILY_TEST_GLSL_DIR + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 2355e6a4f2a..77ff22a20c1 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -41,3 +41,13 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/vulkan:vulkan_graph_runtime", ], ) + + runtime.cxx_test( + name = "coopmat_family_test", + srcs = ["coopmat_family_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) diff --git a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp new file mode 100644 index 00000000000..42997884f02 --- /dev/null +++ b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp @@ -0,0 +1,175 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Correctness cases for the cooperative-matrix family kernels of linear_q4gsw +// and linear_dq8ca_q4gsw (impl/gemm/linear/qw/CoopmatFamily.cpp). The shapes +// are aligned to every family tile (M a multiple of 256, N of 128, K of 64, +// group size 128) and include the shapes that select each row: the Xe2 rows +// (any aligned shape; K * N >= 24e6 for the blocked 8da4w order) and the +// NVIDIA rows (the validated Llama projections, K > 8192 for fp32 +// accumulation). On a device without a family the same cases exercise the +// existing kernels. + +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace vulkan { +namespace prototyping { +namespace q4gsw_linear { + +namespace { + +std::vector to_float(const ValueSpec& s) { + if (s.dtype == vkapi::kFloat) { + return s.get_float_data(); + } + const auto& h = s.get_half_data(); + std::vector o(h.size()); + for (size_t i = 0; i < h.size(); ++i) { + o[i] = half_to_float(h[i]); + } + return o; +} + +// Zero-mean data with exactly representable activations: multiples of 1/16 +// in [-1, 1) (not periodic in K) and the random int4 weights and scales of +// make_case. Partial sums stay small, so the fp16-accumulating kernels (family +// and existing) stay within the tolerance; outputs are of order 1 to 5, so a +// structural error (a missing or misplaced tile, a wrong scale group) does +// not. +// For dq8ca the per-row activation scale is 1/16 so the dynamic int8 +// quantization is exact; the zero points vary per row to exercise the +// zero-point correction (quantized values stay within int8). +void set_family_data(TestCase& t, const CoopmatLinearConfig& cfg) { + auto& hin = t.inputs()[0].get_half_data(); + for (size_t i = 0; i < hin.size(); ++i) { + hin[i] = float_to_half(float(int((i * 7 + i / 61) % 32) - 16) / 16.0f); + } + if (is_dq8ca(cfg.op_name)) { + auto& hs = t.inputs()[1].get_half_data(); + std::fill(hs.begin(), hs.end(), float_to_half(0.0625f)); + auto& zp = t.inputs()[2].get_int8_data(); + for (size_t i = 0; i < zp.size(); ++i) { + zp[i] = int8_t(int(i % 17) - 8); + } + } + // fp16 accumulation error grows with K: at K = 14336 it reaches about 0.2 + // for the fp16 family kernels and 0.7 for the existing fp16 q4gsw kernel + // (llvmpipe), against outputs of order 10. + t.set_abs_tolerance(cfg.K > 8192 ? 1.0f : 0.25f); + t.set_rel_tolerance(0.02f); +} + +// Reference for both ops with 4-bit weights and no bias: +// q4gsw: out[m][n] = SUM_k a[m][k] * w[n][k] * s[k / group][n] +// dq8ca: a is first quantized per row (scale, zero point) and dequantized. +// Weights are dequantized once and the output columns are split across +// threads, so the larger shapes stay fast. +void family_reference(TestCase& tc) { + const bool dq = tc.operator_name().find("dq8ca") != std::string::npos; + const ValueSpec& in = tc.inputs()[0]; + ValueSpec& out = tc.outputs()[0]; + const auto in_sizes = in.get_tensor_sizes(); + const int64_t M = in_sizes[0]; + const int64_t K = in_sizes[1]; + const int64_t N = out.get_tensor_sizes()[1]; + const ValueSpec& w = tc.inputs()[dq ? 3 : 1]; + const std::vector scales = to_float(tc.inputs()[dq ? 5 : 2]); + const int64_t group = tc.inputs()[dq ? 6 : 3].get_int_value(); + const std::vector& w4 = w.get_uint8_data(); // [N, K/2] nibbles + + std::vector a = to_float(in); + std::vector row_scale(M, 1.0f); + if (dq) { + const std::vector in_scale = to_float(tc.inputs()[1]); + const std::vector& in_zp = tc.inputs()[2].get_int8_data(); + for (int64_t m = 0; m < M; ++m) { + const float s = in_scale[m]; + const int zp = int(in_zp[m]); + for (int64_t k = 0; k < K; ++k) { + float q = std::round(a[m * K + k] / s) + float(zp); + q = std::min(std::max(q, -128.0f), 127.0f); + a[m * K + k] = q - float(zp); + } + row_scale[m] = s; + } + } + + auto& ref = out.get_ref_float_data(); + ref.assign(M * N, 0.0f); + const int64_t num_threads = std::max( + 1, std::min(N, std::thread::hardware_concurrency())); + auto work = [&](const int64_t t) { + std::vector wn(K); + for (int64_t n = t; n < N; n += num_threads) { + for (int64_t k = 0; k < K; ++k) { + const uint8_t byte = w4[n * (K / 2) + k / 2]; + const int nib = (k & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); + wn[k] = float(nib - 8) * scales[(k / group) * N + n]; + } + for (int64_t m = 0; m < M; ++m) { + const float* am = a.data() + m * K; + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) { + acc += am[k] * wn[k]; + } + ref[m * N + n] = acc * row_scale[m]; + } + } + }; + std::vector threads; + for (int64_t t = 1; t < num_threads; ++t) { + threads.emplace_back(work, t); + } + work(0); + for (auto& th : threads) { + th.join(); + } +} + +} // namespace + +REGISTER_TEST_CASE_SET("q4gsw_linear", "wmma_correctness") { + // {M, K, N, group_size} + const std::vector shapes = { + {256, 256, 384, 128, ""}, // small, multi-workgroup + {2048, 2048, 512, 128, ""}, // long M + {256, 3072, 1024, 128, ""}, // M = 256 with N % 2048 != 0 + {256, 2048, 8192, 128, ""}, // M = 256 with N % 2048 == 0 + {256, 4096, 8192, 128, ""}, // K * N >= 24e6 + {256, 14336, 4096, 128, ""}, // K > 8192 + }; + std::vector cases; + for (const std::string op : {"linear_q4gsw", "linear_dq8ca_q4gsw"}) { + for (const auto& shape : shapes) { + const CoopmatLinearConfig cfg{ + shape.M, shape.K, shape.N, shape.group_size, op}; + for (const auto st : {utils::kTexture3D, utils::kBuffer}) { + TestCase t = make_case(cfg, st); + set_family_data(t, cfg); + cases.push_back(std::move(t)); + } + } + } + return { + std::move(cases), + family_reference, + flop_calc, + /*warmup_runs=*/1, + /*benchmark_runs=*/3}; +} + +} // namespace q4gsw_linear +} // namespace prototyping +} // namespace vulkan +} // namespace executorch From d537e4796ff336dc550b504ced4a38cf1cf1718b Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Sat, 17 Oct 2026 10:53:54 -0700 Subject: [PATCH 06/11] [ET-VK] q4gsw family kernels: partial last M tile on texture IO 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 --- .../qw/q4gsw/linear_q4gsw_wmma_body.glslh | 17 +++- .../qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh | 12 ++- .../ops/impl/gemm/linear/qw/CoopmatFamily.cpp | 14 ++- .../ops/impl/gemm/linear/qw/CoopmatFamily.h | 7 ++ .../gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp | 5 +- .../test/custom_ops/coopmat_family_test.cpp | 97 ++++++++++++++++++- .../q4gsw_linear/test_sets/wmma_sets.cpp | 32 ++++-- 7 files changed, 166 insertions(+), 18 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh index 668f1b6da04..2879f226cd0 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh @@ -162,8 +162,11 @@ void store_bt_item( // width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), // so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in // two notations. -uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { +uvec4 load_a_vec4(const uint row_in, const uint k_hv4, const uint K4) { #ifdef IO_TEXTURE + // Rows of a partial last M tile read row M - 1; their results are not + // stored (each output row depends only on its own A row). + const uint row = min(row_in, uint(input_sizes.y) - 1u); // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 // (fp32); narrowing immediately lets the compiler fold the conversion into // a 16-bit image load. Lossless -- the source is rgba16f. @@ -179,6 +182,8 @@ uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { packFloat2x16(v0.xy), packFloat2x16(v0.zw), packFloat2x16(v1.xy), packFloat2x16(v1.zw)); #else + // Buffer IO requires M to be a multiple of WG_TILE_M (CoopmatFamily.cpp). + const uint row = row_in; const f16vec4 v0 = t_input[row * K4 + k_hv4]; const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; return uvec4( @@ -507,6 +512,11 @@ void main() { const uint lc4 = t % CSH_TEXELS_PER_ROW; const uint m = tile_m_start + b * SG_TILE_M + i * MMA_M + lr; const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } imageStore( t_output, ivec3(tile_n_start / 4u + lc4, m, 0), @@ -544,6 +554,11 @@ void main() { const uint m = tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } imageStore( t_output, ivec3(tile_n_start / 4u + lc4, m, 0), diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh index 10f258ab782..adaa9790beb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh @@ -99,8 +99,11 @@ uvec4 dequant_block( // width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), // so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in // two notations. -uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { +uvec4 load_a_vec4(const uint row_in, const uint k_hv4, const uint K4) { #ifdef IO_TEXTURE + // Rows of a partial last M tile read row M - 1; their results are not + // stored (each output row depends only on its own A row). + const uint row = min(row_in, uint(input_sizes.y) - 1u); // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 // (fp32); narrowing immediately lets the compiler fold the conversion into // a 16-bit image load. Lossless -- the source is rgba16f. @@ -110,6 +113,8 @@ uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { packFloat2x16(v0.xy), packFloat2x16(v0.zw), packFloat2x16(v1.xy), packFloat2x16(v1.zw)); #else + // Buffer IO requires M to be a multiple of WG_TILE_M (CoopmatFamily.cpp). + const uint row = row_in; const f16vec4 v0 = t_input[row * K4 + k_hv4]; const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; return uvec4( @@ -374,6 +379,11 @@ void main() { const uint m = tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } #ifdef CSH_IN_ASH // base is a multiple of 4 halves: the texel is one half of a uvec4. const uvec4 q = Ash[base / FP16_PER_VEC4]; diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp index ed1e58399e2..1976b225049 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp @@ -241,12 +241,18 @@ bool wmma_row_fits( if ((combo & row.storages) == 0) { return false; } - if (!shape.ignore_alignment && - (shape.M % tile.m != 0 || shape.N % tile.n != 0 || - shape.K % tile.k != 0 || shape.group_size % tile.k != 0)) { + WmmaShape sel = shape; + if (shape.op == WmmaOp::kQ4gsw && combo == kWmmaTex3dTex2d && + shape.M >= kWmmaPartialTileM) { + sel.M = (shape.M + kWmmaPartialTileM - 1) / kWmmaPartialTileM * + kWmmaPartialTileM; + } + if (!sel.ignore_alignment && + (sel.M % tile.m != 0 || sel.N % tile.n != 0 || sel.K % tile.k != 0 || + sel.group_size % tile.k != 0)) { return false; } - return row.shape_ok == nullptr || row.shape_ok(shape); + return row.shape_ok == nullptr || row.shape_ok(sel); } std::optional select_wmma_kernel( diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h index 830166e02d0..6882f1bd205 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h @@ -116,6 +116,13 @@ struct WmmaShape { bool ignore_alignment = false; }; +// The kQ4gsw texture-IO kernels accept a partial last M tile. Their rows take +// an M of at least this many rows that is not a multiple of their tile, and are +// selected as for M rounded up to a multiple of it (the smallest WG_TILE_M of +// the family rows), so an unaligned M gets the kernel of the next aligned M. +// Smaller M, buffer IO and kDq8caQ4gsw keep the alignment requirement. +constexpr int64_t kWmmaPartialTileM = 128; + struct WmmaTile { uint32_t m; uint32_t n; diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp index 228bda2f6ac..65ddaa062bf 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp @@ -94,8 +94,9 @@ vkapi::ShaderInfo pick_q4gsw_linear_wmma_shader( if (choice.has_value()) { kernel_name = choice->kernel_base; } else { - // Shapes no family row covers after a resize (decode, unaligned M): the - // kernels that read this packed weight layout. + // Shapes no family row covers after a resize (decode, M below + // kWmmaPartialTileM, buffer IO at unaligned M): the kernels that read + // this packed weight layout. kernel_name = is_gemv(graph, fp_input) ? "linear_q4gsw_coop" : "linear_q4gsw_tiled"; } diff --git a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp index 5a982bc1cff..31053e57937 100644 --- a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp +++ b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp @@ -300,8 +300,9 @@ TEST(CoopmatFamilyTest, ShapesThatKeepTheExistingKernels) { prefill(op, 2048, 4096, 4096, WmmaStorage::kTexture3D); ASSERT_FALSE(selected(d, s).empty()); + // Below kWmmaPartialTileM no row takes an unaligned M. WmmaShape unaligned = s; - unaligned.M = 1304; + unaligned.M = 120; WmmaShape gemv = s; gemv.M = 1; gemv.gemv = true; @@ -334,6 +335,100 @@ TEST(CoopmatFamilyTest, ShapesThatKeepTheExistingKernels) { } } +TEST(CoopmatFamilyTest, PartialMTile) { + struct Case { + CoopmatDeviceKey device; + WmmaOp op; + int64_t M, K, N; + WmmaStorage io; + std::string want; // empty: existing kernels + }; + const Case cases[] = { + // 4w texture IO: selected as for M rounded up to 128 rows. + {xe2_key(), + WmmaOp::kQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g82s16m8flib"}, + {xe2_key(), + WmmaOp::kQ4gsw, + 129, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g82s16m8flib"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t256x128k16g42s32bt"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 300, + 3072, + 1024, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g22s32"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 129, + 14336, + 4096, + WmmaStorage::kTexture3D, + kQ + "bx_t128x128k32g42s32f32c"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 2047, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t256x128k16g42s32bt"}, + // Rounded up past the validated prefill length. + {nvidia_key(), + WmmaOp::kQ4gsw, + 2049, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + // Below kWmmaPartialTileM. + {xe2_key(), WmmaOp::kQ4gsw, 127, 2048, 2048, WmmaStorage::kTexture3D, ""}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 127, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + // Buffer IO and 8da4w keep the alignment requirement. + {xe2_key(), WmmaOp::kQ4gsw, 511, 2048, 2048, WmmaStorage::kBuffer, ""}, + {xe2_key(), + WmmaOp::kDq8caQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + {nvidia_key(), + WmmaOp::kDq8caQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + }; + for (const Case& c : cases) { + EXPECT_EQ(selected(c.device, prefill(c.op, c.M, c.K, c.N, c.io)), c.want) + << coopmat_family_name(coopmat_family(c.device)) + << " op=" << static_cast(c.op) << " M=" << c.M << " K=" << c.K + << " N=" << c.N << " io=" << static_cast(c.io); + } +} + TEST(CoopmatFamilyTest, Int8Layout) { const CoopmatDeviceKey d = nvidia_key(); WmmaShape s = diff --git a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp index 42997884f02..ba82e6a681f 100644 --- a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp @@ -10,8 +10,10 @@ // group size 128) and include the shapes that select each row: the Xe2 rows // (any aligned shape; K * N >= 24e6 for the blocked 8da4w order) and the // NVIDIA rows (the validated Llama projections, K > 8192 for fp32 -// accumulation). On a device without a family the same cases exercise the -// existing kernels. +// accumulation). linear_q4gsw also runs at M that is not a multiple of the +// tile: a partial last M tile on texture IO (kWmmaPartialTileM), the existing +// kernels below it and on buffer IO. On a device without a family the same +// cases exercise the existing kernels. #include @@ -149,18 +151,30 @@ REGISTER_TEST_CASE_SET("q4gsw_linear", "wmma_correctness") { {256, 4096, 8192, 128, ""}, // K * N >= 24e6 {256, 14336, 4096, 128, ""}, // K > 8192 }; + // linear_q4gsw at M not a multiple of the tile (Llama projections, so the + // NVIDIA rows apply). + const std::vector unaligned_shapes = { + {511, 2048, 2048, 128, ""}, // 256-row tile (NVIDIA), 255-row last tile + {300, 3072, 1024, 128, ""}, // 128-row tile, 44-row last tile + {129, 14336, 4096, 128, ""}, // K > 8192, 1-row last tile + {127, 2048, 512, 128, ""}, // below kWmmaPartialTileM + }; std::vector cases; + auto add = [&](const CoopmatLinearConfig& cfg) { + for (const auto st : {utils::kTexture3D, utils::kBuffer}) { + TestCase t = make_case(cfg, st); + set_family_data(t, cfg); + cases.push_back(std::move(t)); + } + }; for (const std::string op : {"linear_q4gsw", "linear_dq8ca_q4gsw"}) { for (const auto& shape : shapes) { - const CoopmatLinearConfig cfg{ - shape.M, shape.K, shape.N, shape.group_size, op}; - for (const auto st : {utils::kTexture3D, utils::kBuffer}) { - TestCase t = make_case(cfg, st); - set_family_data(t, cfg); - cases.push_back(std::move(t)); - } + add({shape.M, shape.K, shape.N, shape.group_size, op}); } } + for (const auto& shape : unaligned_shapes) { + add({shape.M, shape.K, shape.N, shape.group_size, "linear_q4gsw"}); + } return { std::move(cases), family_reference, From 2c79704c53049b3f173dfaa1b7f7ce41a3d0b65b Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Sat, 17 Oct 2026 10:53:54 -0700 Subject: [PATCH 07/11] [ET-VK] Row-major dq8ca family kernel: no reads past M rows 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 --- .../linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh index f51484e34d1..61658d0ede5 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh @@ -15,8 +15,9 @@ // the end); differs in the A layout (row-major int8, staged as raw 16-byte // vectors) and in B staging (one thread per whole weight texel). // -// A is read for whole WG_TILE_M row tiles, so the int8 activation buffer must -// cover M rounded up to WG_TILE_M rows. +// The last M tile may be partial: its rows at or past M read row M - 1 of the +// activations and their params (each output row depends only on its own +// activation row) and are not stored, so no buffer is read past M rows. #ifndef LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH #define LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH @@ -126,6 +127,7 @@ void main() { const uint tile_m_start = WG_TILE_M * tileID.y; const uint tile_n_start = WG_TILE_N * tileID.x; + const uint M_rows = uint(input_sizes.y); [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { @@ -191,7 +193,8 @@ void main() { const uint row = (q / A_PARTS) % WG_TILE_M; const uint slab = q / (A_PARTS * WG_TILE_M); a_lds[v] = slab * A_SLAB_U32 + row * A_STRIDE_U32 + part * 4u; - a_gv4[v] = ((tile_m_start + row) * a_row_stride_i8 + slab * MMA_K + part * 16u) / 16u; + const uint m_row = min(tile_m_start + row, M_rows - 1u); + a_gv4[v] = (m_row * a_row_stride_i8 + slab * MMA_K + part * 16u) / 16u; } uvec4 temp_Av[A_VECS_PER_THREAD]; ivec4 temp_B[B_SLOTS_PER_THREAD]; @@ -201,7 +204,8 @@ void main() { // PROLOGUE // ========================================================= if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const uint m4 = min( + (tile_m_start >> 2u) + gl_LocalInvocationID.x, (M_rows + 3u) / 4u - 1u); const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); const uint base = gl_LocalInvocationID.x * 4u; @@ -462,6 +466,11 @@ void main() { const uint m = tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } imageStore( t_output, ivec3(tile_n_start / 4u + lc4, m, 0), From 3345345a8aae5fc695850b0773a24b31708619c8 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Sun, 18 Oct 2026 17:03:09 -0700 Subject: [PATCH 08/11] [ET-VK] Family kernel selection: NVIDIA rows on any aligned shape, Xe2 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 --- .../ops/impl/gemm/linear/qw/CoopmatFamily.cpp | 29 +++++-------- .../ops/impl/gemm/linear/qw/CoopmatFamily.h | 5 ++- .../test/custom_ops/coopmat_family_test.cpp | 42 +++++++++++-------- .../q4gsw_linear/test_sets/wmma_sets.cpp | 22 +++++----- 4 files changed, 49 insertions(+), 49 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp index 1976b225049..92a7aa331fb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp @@ -40,9 +40,12 @@ bool subgroup_size_available( // matrix shapes Xe2 exposes), subgroup 16. // -// Weight-heavy projections run their workgroups in blocks of N tiles. +// Weight-heavy projections run their workgroups in blocks of N tiles. The +// threshold lies between the largest shape where the blocked order is slower +// (K * N = 16.8e6) and the smallest where it is faster (21.2e6) on both the +// Arc B580 and the Arc Pro B70. bool xe2_weight_heavy(const WmmaShape& s) { - return s.K * s.N >= 24000000; + return s.K * s.N >= 20000000; } const std::vector& xe2_rows() { @@ -80,27 +83,15 @@ const std::vector& xe2_rows() { // texture3d IO only. // -// Only the prefill projections of Llama 3.2 1B / 3B and Llama 3.1 8B at group -// size 128 and M <= 2048 were validated on the family's members; every other -// shape keeps the existing kernels. -bool nv_measured(const WmmaShape& s) { - const int64_t K = s.K; - const int64_t N = s.N; - const bool projection = (K == 2048 && (N == 512 || N == 2048 || N == 8192)) || - (K == 3072 && (N == 1024 || N == 3072 || N == 8192)) || - (K == 4096 && (N == 1024 || N == 4096 || N == 14336)) || - (K == 8192 && (N == 2048 || N == 3072)) || (K == 14336 && N == 4096); - return projection && s.group_size == 128 && s.M <= 2048; -} // fp16 accumulation loses accuracy beyond K = 8192: fp32 accumulation there. bool nv_large_k(const WmmaShape& s) { - return nv_measured(s) && s.K > 8192; + return s.K > 8192; } bool nv_256(const WmmaShape& s) { - return nv_measured(s) && s.K <= 8192 && (s.M != 256 || s.N % 2048 == 0); + return s.K <= 8192 && (s.M != 256 || s.N % 2048 == 0); } bool nv_default(const WmmaShape& s) { - return nv_measured(s) && s.K <= 8192; + return s.K <= 8192; } const std::vector& nvidia_rows() { @@ -118,7 +109,7 @@ const std::vector& nvidia_rows() { {256, 128, 16, 4, 2, 32, 16, false}, kWmmaTex3dTex2d, nv_256}, - // 4w, the remaining validated shapes (M not a multiple of 256, and + // 4w, K <= 8192, the remaining shapes (M not a multiple of 256, and // M = 256 at N % 2048 != 0). {WmmaOp::kQ4gsw, "linear_q4gsw_wmma_t128x128k16g22s32", @@ -130,7 +121,7 @@ const std::vector& nvidia_rows() { "linear_dq8ca_q4gsw_wmma_rowmajor_t128x128k64g24s32mk32", {128, 128, 64, 2, 4, 32, 16, true}, kWmmaTex3dTex2d, - nv_measured, + nullptr, /*rowmajor_a=*/true}, }; return rows; diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h index 6882f1bd205..826dcac58a3 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h @@ -70,8 +70,9 @@ enum class CoopmatFamily { // subgroup size 16 available. kXe2, // NVIDIA proprietary driver: fp16 16x16x16 (fp16 and fp32 result), int8 - // 16x16x32, subgroup size 32. The key does not separate Ampere and Ada from - // other NVIDIA generations with the same shapes. + // 16x16x32, subgroup size 32. Measured on Ampere (Jetson Orin) and Ada; the + // key does not separate them from other NVIDIA generations with the same + // shapes. kNvidia, }; diff --git a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp index 31053e57937..a3f05b3924b 100644 --- a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp +++ b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp @@ -146,7 +146,7 @@ std::string xe2_rule(const WmmaShape& s) { : "t128x128k16g44s16m8fli"); } return kD + - (s.K * s.N >= 24000000 ? "t128x128k64g84s16m8n13" + (s.K * s.N >= 20000000 ? "t128x128k64g84s16m8n13" : "t128x128k64g84s16m8i"); } @@ -241,18 +241,18 @@ TEST(CoopmatFamilyTest, UnmatchedKeysSelectNothing) { TEST(CoopmatFamilyTest, Xe2Rows) { check_llama_shapes("xe2", xe2_key(), xe2_rule); - // The weight-heavy rule: K * N >= 24e6. + // The weight-heavy rule: K * N >= 20e6. const CoopmatDeviceKey d = xe2_key(); EXPECT_EQ( selected( d, - prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 11776, WmmaStorage::kBuffer)), - kD + "t128x128k64g84s16m8n13"); // 24.1e6 + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 9856, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8n13"); // 20.2e6 EXPECT_EQ( selected( d, - prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 11648, WmmaStorage::kBuffer)), - kD + "t128x128k64g84s16m8i"); // 23.9e6 + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 9728, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8i"); // 19.9e6 // Any aligned shape, not only the Llama projections. EXPECT_EQ( selected(d, prefill(WmmaOp::kQ4gsw, 128, 256, 384, WmmaStorage::kBuffer)), @@ -274,11 +274,13 @@ TEST(CoopmatFamilyTest, NvidiaRows) { {256, 2048, 8192, "t256x128k16g42s32bt"}, // K > 8192 with M not a multiple of 256: fp32 accumulation. {384, 14336, 4096, "bx_t128x128k32g42s32f32c"}, - // Shapes outside the validated projections. - {2048, 2048, 128256, nullptr}, - {2048, 2048, 4096, nullptr}, - // M beyond the validated prefill length. - {4096, 2048, 2048, nullptr}, + // Any aligned shape, not only the Llama projections (Qwen2.5 0.5B, + // Gemma 2 2B), and M beyond 2048. + {512, 896, 4864, "t256x128k16g42s32bt"}, + {512, 9216, 2304, "bx_t128x128k32g42s32f32c"}, + {4096, 2048, 2048, "t256x128k16g42s32bt"}, + // N not a multiple of the 128-column tile. + {2048, 2048, 4160, nullptr}, }; for (const Edge& e : edges) { const WmmaShape s = @@ -286,11 +288,16 @@ TEST(CoopmatFamilyTest, NvidiaRows) { EXPECT_EQ(selected(d, s), e.want == nullptr ? "" : kQ + e.want) << "M=" << e.M << " K=" << e.K << " N=" << e.N; } - // Group size other than 128. - WmmaShape g64 = - prefill(WmmaOp::kQ4gsw, 2048, 2048, 2048, WmmaStorage::kTexture3D); - g64.group_size = 64; - EXPECT_EQ(selected(d, g64), ""); + // Group sizes other than 128, down to the K tile of the row. + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + WmmaShape g64 = prefill(op, 2048, 2048, 2048, WmmaStorage::kTexture3D); + g64.group_size = 64; + EXPECT_EQ(selected(d, g64), nvidia_rule(g64)); + } + WmmaShape g32 = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 2048, 2048, WmmaStorage::kTexture3D); + g32.group_size = 32; + EXPECT_EQ(selected(d, g32), ""); } TEST(CoopmatFamilyTest, ShapesThatKeepTheExistingKernels) { @@ -387,14 +394,13 @@ TEST(CoopmatFamilyTest, PartialMTile) { 2048, WmmaStorage::kTexture3D, kQ + "t256x128k16g42s32bt"}, - // Rounded up past the validated prefill length. {nvidia_key(), WmmaOp::kQ4gsw, 2049, 2048, 2048, WmmaStorage::kTexture3D, - ""}, + kQ + "t128x128k16g22s32"}, // Below kWmmaPartialTileM. {xe2_key(), WmmaOp::kQ4gsw, 127, 2048, 2048, WmmaStorage::kTexture3D, ""}, {nvidia_key(), diff --git a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp index ba82e6a681f..b24f111fbf2 100644 --- a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp @@ -7,13 +7,13 @@ // Correctness cases for the cooperative-matrix family kernels of linear_q4gsw // and linear_dq8ca_q4gsw (impl/gemm/linear/qw/CoopmatFamily.cpp). The shapes // are aligned to every family tile (M a multiple of 256, N of 128, K of 64, -// group size 128) and include the shapes that select each row: the Xe2 rows -// (any aligned shape; K * N >= 24e6 for the blocked 8da4w order) and the -// NVIDIA rows (the validated Llama projections, K > 8192 for fp32 -// accumulation). linear_q4gsw also runs at M that is not a multiple of the -// tile: a partial last M tile on texture IO (kWmmaPartialTileM), the existing -// kernels below it and on buffer IO. On a device without a family the same -// cases exercise the existing kernels. +// group size 64 or 128) and include the shapes that select each row (K * N +// >= 20e6 for the blocked Xe2 8da4w order, K > 8192 for the NVIDIA fp32 +// accumulation): Llama projections and some of Qwen2.5 0.5B and Gemma 2 2B. +// linear_q4gsw also runs at M that is not a multiple of the tile: a partial +// last M tile on texture IO (kWmmaPartialTileM), the existing kernels below it +// and on buffer IO. On a device without a family the same cases exercise the +// existing kernels. #include @@ -148,11 +148,13 @@ REGISTER_TEST_CASE_SET("q4gsw_linear", "wmma_correctness") { {2048, 2048, 512, 128, ""}, // long M {256, 3072, 1024, 128, ""}, // M = 256 with N % 2048 != 0 {256, 2048, 8192, 128, ""}, // M = 256 with N % 2048 == 0 - {256, 4096, 8192, 128, ""}, // K * N >= 24e6 + {256, 4096, 8192, 128, ""}, // K * N >= 20e6 {256, 14336, 4096, 128, ""}, // K > 8192 + {256, 896, 4864, 128, ""}, // Qwen2.5 0.5B MLP + {256, 9216, 2304, 128, ""}, // Gemma 2 2B down projection, K > 8192 + {256, 2048, 2048, 64, ""}, // group size 64 }; - // linear_q4gsw at M not a multiple of the tile (Llama projections, so the - // NVIDIA rows apply). + // linear_q4gsw at M not a multiple of the tile. const std::vector unaligned_shapes = { {511, 2048, 2048, 128, ""}, // 256-row tile (NVIDIA), 255-row last tile {300, 3072, 1024, 128, ""}, // 128-row tile, 44-row last tile From 1d8fc04442c7593999c55dafb77f3094276421a4 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Mon, 19 Oct 2026 16:52:42 -0700 Subject: [PATCH 09/11] [ET-VK] Fence the build commands in the custom_ops CMakeLists comment cmake-format reflows plain comment lines; fence the commands with ~~~ as backends/vulkan/CMakeLists.txt does. Co-Authored-By: Claude Opus 5.5 --- backends/vulkan/test/custom_ops/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backends/vulkan/test/custom_ops/CMakeLists.txt b/backends/vulkan/test/custom_ops/CMakeLists.txt index 21fa561fb8d..fdf2b485010 100644 --- a/backends/vulkan/test/custom_ops/CMakeLists.txt +++ b/backends/vulkan/test/custom_ops/CMakeLists.txt @@ -6,9 +6,11 @@ # Host-only tests that need neither a GPU nor the Vulkan backend library. # +# ~~~ # cmake backends/vulkan/test/custom_ops -Bcmake-out/backends/vulkan/test/custom_ops # cmake --build cmake-out/backends/vulkan/test/custom_ops # ctest --test-dir cmake-out/backends/vulkan/test/custom_ops +# ~~~ cmake_minimum_required(VERSION 3.19) project(vulkan_custom_ops_tests CXX) From dd34a6664cdaa035e01c9b4f0054ac5486b76252 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Thu, 22 Oct 2026 14:32:49 -0700 Subject: [PATCH 10/11] [ET-VK] Enable vulkanMemoryModelDeviceScope with the memory model Existing shaders use Device memory scope. With vulkanMemoryModel enabled and vulkanMemoryModelDeviceScope left disabled this violates VUID-RuntimeSpirv-vulkanMemoryModel-06265. Enable it whenever the device reports it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJnuDtQrcF8ZL9aHrH3VnY --- backends/vulkan/runtime/vk_api/Adapter.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index 8d4cbba72ea..501d80a7130 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -185,6 +185,11 @@ VkDevice create_logical_device( if (std::strcmp(extension, VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME) == 0) { vulkan_memory_model_features.vulkanMemoryModel = VK_TRUE; + // Existing shaders use Device memory scope + // (VUID-RuntimeSpirv-vulkanMemoryModel-06265). + vulkan_memory_model_features.vulkanMemoryModelDeviceScope = + physical_device.vulkan_memory_model_features + .vulkanMemoryModelDeviceScope; vulkan_memory_model_features.pNext = extension_list_top; extension_list_top = &vulkan_memory_model_features; break; From d5edd60a47351fb9678de499d6756a524350083d Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Thu, 22 Oct 2026 14:48:54 -0700 Subject: [PATCH 11/11] [ET-VK] Bind the Xe2 q4gsw texture3d input in the general layout The Xe2 texture3d 4w kernel reads A through a storage image (IMG_A), but the node bound the input read-only, so it sat in SHADER_READ_ONLY_OPTIMAL while the storage image descriptor uses GENERAL (VUID-vkCmdDispatch-imageLayout-00344). Mark such rows with image_a and bind the input read-write for them; the fallback kernels sample the input, which the general layout also allows. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJnuDtQrcF8ZL9aHrH3VnY --- .../ops/impl/gemm/linear/qw/CoopmatFamily.cpp | 6 +++-- .../ops/impl/gemm/linear/qw/CoopmatFamily.h | 4 ++++ .../gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp | 23 ++++++++++++------- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp index 92a7aa331fb..06a3efea361 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp @@ -55,7 +55,9 @@ const std::vector& xe2_rows() { "linear_q4gsw_wmma_t128x128k16g82s16m8flib", {128, 128, 16, 8, 2, 16, 8, false}, kWmmaTex3dTex2d, - nullptr}, + nullptr, + /*rowmajor_a=*/false, + /*image_a=*/true}, // 4w, buffer IO: 4 x 4 subgroups. {WmmaOp::kQ4gsw, "linear_q4gsw_wmma_t128x128k16g44s16m8fli", @@ -251,7 +253,7 @@ std::optional select_wmma_kernel( const WmmaShape& shape) { for (const WmmaRow& row : coopmat_family_rows(coopmat_family(device))) { if (wmma_row_fits(device, shape, row)) { - return WmmaChoice{row.kernel_base, row.tile, row.rowmajor_a}; + return WmmaChoice{row.kernel_base, row.tile, row.rowmajor_a, row.image_a}; } } return std::nullopt; diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h index 826dcac58a3..ed810af0149 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h @@ -154,12 +154,16 @@ struct WmmaRow { // in order, so a later row can serve the shapes an earlier one rejects. bool (*shape_ok)(const WmmaShape&); bool rowmajor_a = false; // kDq8caQ4gsw: reads row-major int8 activations + // kQ4gsw: the texture3d variant reads A through a storage image, so the + // input must be bound in VK_IMAGE_LAYOUT_GENERAL. + bool image_a = false; }; struct WmmaChoice { std::string kernel_base; WmmaTile tile; bool rowmajor_a = false; + bool image_a = false; }; // The rows of a family, in selection order (empty for kNone). diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp index 65ddaa062bf..c6ac5b46424 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp @@ -21,7 +21,7 @@ namespace vkcompute { namespace { // Dispatch node args: -// {{output}, {fp_input, packed_weight, packed_weight_scales, packed_bias}} +// {{output}, {fp_input}, {packed_weight, packed_weight_scales, packed_bias}} // resize_args: {group_size, weight_data, bias_data} WmmaShape q4gsw_runtime_shape( ComputeGraph* graph, @@ -29,7 +29,7 @@ WmmaShape q4gsw_runtime_shape( const std::vector& resize_args) { const ValueRef output = args.at(0).refs.at(0); const ValueRef fp_input = args.at(1).refs.at(0); - const ValueRef packed_weight = args.at(1).refs.at(1); + const ValueRef packed_weight = args.at(2).refs.at(0); const std::vector out_sizes = graph->sizes_of(output); WmmaShape s; @@ -85,7 +85,7 @@ vkapi::ShaderInfo pick_q4gsw_linear_wmma_shader( const std::vector& resize_args) { const ValueRef output = args.at(0).refs.at(0); const ValueRef fp_input = args.at(1).refs.at(0); - const ValueRef packed_weight = args.at(1).refs.at(1); + const ValueRef packed_weight = args.at(2).refs.at(0); const std::optional choice = select_wmma_kernel( coopmat_device_key(graph), q4gsw_runtime_shape(graph, args, resize_args)); @@ -124,11 +124,18 @@ bool try_add_q4gsw_linear_wmma_node( } // Take the op over only if a row applies to its build-time shape; otherwise // it stays entirely on the existing path, including for later resizes. - if (!select_wmma_kernel( - coopmat_device_key(&graph), q4gsw_build_time_shape(graph, args)) - .has_value()) { + const std::optional build_choice = select_wmma_kernel( + coopmat_device_key(&graph), q4gsw_build_time_shape(graph, args)); + if (!build_choice.has_value()) { return false; } + // A storage-image read of the input needs VK_IMAGE_LAYOUT_GENERAL + // (VUID-vkCmdDispatch-imageLayout-00344), which read-write access selects. + // The fallback kernels sample the input, which GENERAL also allows. + const vkapi::MemoryAccessFlags input_access = build_choice->image_a && + graph.storage_type_of(fp_input) == utils::kTexture3D + ? vkapi::kReadWrite + : vkapi::kRead; const int64_t group_size_val = graph.extract_scalar(group_size); const QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); @@ -159,8 +166,8 @@ bool try_add_q4gsw_linear_wmma_node( wmma_linear_gwg, quantized_linear_lwg, {{output, vkapi::kWrite}, - {{fp_input, packed_weight, packed_weight_scales, packed_bias}, - vkapi::kRead}}, + {fp_input, input_access}, + {{packed_weight, packed_weight_scales, packed_bias}, vkapi::kRead}}, {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, {}, // The spec constants of the linear_q4gsw_{tiled,coop} kernels. The