[ET-VK][ez] Add support for buffer backed qparams in int4 linear + add checks for physical limits when allocating

## Context

Currently, the groupwise quantized int4 linear op implementation forces the scales and zero tensor to be a `Texture3D`. However, for i.e. transformer models that have a logit linear layer, the image extents required may exceed the maximum image extents available on the device.

## Changes

* Add support for the scales and zero tensor being a `Buffer` instead of a `Texture3D`
* Add checks when allocating buffers or images for tensors that the requested resource fits within the physical device limits

Differential Revision: [D72662176](https://our.internmc.facebook.com/intern/diff/D72662176/)

ghstack-source-id: 276858281
Pull Request resolved: #9974

Author: SS-JIA

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -260,6 +260,11 @@ vkapi::VulkanImage allocate_image(
       return vkapi::VulkanImage();
   }
 
+  utils::uvec3 max_extents = adapter_ptr->max_texture_extents();
+  VK_CHECK_COND(
+      image_extents[0] <= max_extents[0] &&
+      image_extents[1] <= max_extents[1] && image_extents[2] <= max_extents[2]);
+
   VkSampler sampler = adapter_ptr->sampler_cache().retrieve(sampler_props);
 
   return adapter_ptr->vma().create_image(
@@ -291,6 +296,8 @@ vkapi::VulkanBuffer allocate_buffer(
       return vkapi::VulkanBuffer();
   }
 
+  VK_CHECK_COND(numel <= context_ptr->adapter_ptr()->max_buffer_numel());
+
   return adapter_ptr->vma().create_storage_buffer(
       element_size(dtype) * numel, allocate_memory);
 }

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

@@ -21,7 +21,7 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, OUT_STORAGE, is_scalar_array=False)}
 ${layout_declare_tensor(B, "r", "t_mat1", DTYPE, IN_STORAGE, is_scalar_array=False)}
 ${layout_declare_tensor(B, "r", "t_qmat2", "uint8", WEIGHT_STORAGE, is_scalar_array=False)}
-${layout_declare_tensor(B, "r", "t_qparams", DTYPE, "texture3D")}
+${layout_declare_tensor(B, "r", "t_qparams", DTYPE, PARAMS_STORAGE, is_scalar_array=False)}
 
 layout(push_constant) uniform restrict Block {
   ivec4 out_sizes;
@@ -79,13 +79,23 @@ void main() {
 
   $if WEIGHT_STORAGE == "buffer":
     const int qmat2_stride = qmat2_sizes.x >> 2;
+  $if PARAMS_STORAGE == "buffer":
+    const int qparams_y_stride = out_sizes.x >> 2;
+    const int qparams_z_stride = qparams_y_stride * 2;
 
   for (int block_idx = 0; block_idx < num_blocks; ++block_idx) {
-    scales[0] = texelFetch(t_qparams, ivec3(out_col_texel_idx, 0, block_idx), 0);
-    zeros[0] = texelFetch(t_qparams, ivec3(out_col_texel_idx, 1, block_idx), 0);
-
-    scales[1] = texelFetch(t_qparams, ivec3(out_col_texel_idx + 1, 0, block_idx), 0);
-    zeros[1] = texelFetch(t_qparams, ivec3(out_col_texel_idx + 1, 1, block_idx), 0);
+    $if PARAMS_STORAGE == "buffer":
+      scales[0] = t_qparams[block_idx * qparams_z_stride + out_col_texel_idx];
+      zeros[0] = t_qparams[block_idx * qparams_z_stride + out_col_texel_idx + qparams_y_stride];
+
+      scales[1] = t_qparams[block_idx * qparams_z_stride + out_col_texel_idx + 1];
+      zeros[1] = t_qparams[block_idx * qparams_z_stride + out_col_texel_idx + 1 + qparams_y_stride];
+    $else:
+      scales[0] = texelFetch(t_qparams, ivec3(out_col_texel_idx, 0, block_idx), 0);
+      zeros[0] = texelFetch(t_qparams, ivec3(out_col_texel_idx, 1, block_idx), 0);
+
+      scales[1] = texelFetch(t_qparams, ivec3(out_col_texel_idx + 1, 0, block_idx), 0);
+      zeros[1] = texelFetch(t_qparams, ivec3(out_col_texel_idx + 1, 1, block_idx), 0);
 
     for (int g_idx = 0; g_idx < group_size; g_idx += 4) {
       const int k = block_idx * group_size + g_idx;

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

@@ -10,14 +10,13 @@ q_4w_linear:
     OUT_STORAGE: texture3d
     IN_STORAGE: texture3d
     WEIGHT_STORAGE: texture3d
+    PARAMS_STORAGE: texture3d
   shader_variants:
-    - NAME: q_4w_linear_texture3d_texture3d_texture3d_float
-    - NAME: q_4w_linear_texture3d_buffer_texture3d_float
-      IN_STORAGE: buffer
-    - NAME: q_4w_linear_buffer_buffer_texture3d_float
+    - NAME: q_4w_linear_texture3d_texture3d_texture3d_texture3d_float
+    - NAME: q_4w_linear_buffer_buffer_texture3d_texture3d_float
       OUT_STORAGE: buffer
       IN_STORAGE: buffer
-    - NAME: q_4w_linear_buffer_buffer_buffer_float
+    - NAME: q_4w_linear_buffer_buffer_texture3d_buffer_float
       OUT_STORAGE: buffer
       IN_STORAGE: buffer
-      WEIGHT_STORAGE: buffer
+      PARAMS_STORAGE: buffer

backends/vulkan/runtime/graph/ops/impl/QuantizedLinearGroupwiseInt4.cpp

@@ -132,13 +132,22 @@ void add_q_4w_linear_node(
   ValueRef mat2 =
       prepack_int4_linear_weight_transposed_interleaved(graph, mat2_data);
 
+  utils::StorageType qparams_storage_type = utils::kTexture3D;
+  utils::uvec3 max_extents =
+      graph.context()->adapter_ptr()->max_texture_extents();
+  if (graph.size_at<uint32_t>(-2, scales_and_zeros_data) > max_extents[0] * 4 ||
+      graph.size_at<uint32_t>(-3, scales_and_zeros_data) > max_extents[2]) {
+    qparams_storage_type = utils::kBuffer;
+  }
+
   ValueRef scales_and_zeros = prepack_standard_hw_transposed(
-      graph, scales_and_zeros_data, utils::kTexture3D, utils::kWidthPacked);
+      graph, scales_and_zeros_data, qparams_storage_type, utils::kWidthPacked);
 
   std::string kernel_name = "q_4w_linear";
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_storage_type_suffix(kernel_name, graph.storage_type_of(mat1));
   add_storage_type_suffix(kernel_name, graph.storage_type_of(mat2));
+  add_storage_type_suffix(kernel_name, qparams_storage_type);
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
 
   const uint32_t group_size_val = graph.extract_scalar<uint32_t>(group_size);

backends/vulkan/runtime/vk_api/Adapter.h

@@ -218,6 +218,10 @@ class Adapter final {
         physical_device_.properties.limits.maxImageDimension3D};
   }
 
+  inline uint32_t max_buffer_numel() const {
+    return physical_device_.properties.limits.maxStorageBufferRange;
+  }
+
   // Command Buffer Submission
 
   void