[ET-VK][Ops] torchao.dequantize_affine vulkan impl and shader and cleanup

# Changes
* Implement `torchao.dequantize_affine` operator in Vulkan backend with comprehensive texture and buffer storage support
* Add block-wise dequantization mode in `dequantize_texture.glsl` and `dequantize_buffer.glsl` shaders for configurable tensor block dequantization
* Extend dequantization infrastructure in `Dequantize.cpp` to handle affine transformations with configurable block sizes and quantization parameters
* Support integer-to-floating-point conversion with precise reconstruction of original values

BE: Improved the documentation in the shader logic which is more detailed and clear


# Motivation
The existing Vulkan quantization infrastructure lacked support for the `torchao.dequantize_affine` operator, which is essential for completing the quantization-dequantization cycle in dynamic quantization workflows. The `dequantize_affine` operator provides flexible block-wise dequantization that reconstructs floating-point values from quantized integer blocks, enabling:

* **Block-wise Dequantization**: Reconstructs floating-point values from configurable tensor blocks using separate scale and zero-point parameters, enabling precise recovery of original data distributions
* **Affine Transformation**: Uses the formula `value = (qvalue - zero_point) * scale` for accurate integer-to-floating-point mapping
* **TorchAO Integration**: Seamless compatibility with TorchAO quantization workflows and completes the quantization-dequantization round-trip

# Operator Description
The `dequantize_affine` operator converts n-bit integer tensor values back to floating-point representations using pre-computed quantization parameters (scale and zero_point) applied to configurable tensor blocks. Block-wise dequantization divides tensors into blocks and applies separate dequantization parameters to each block, allowing fine-grained reconstruction of the original floating-point precision.

The dequantization formula is: `value = (qvalue - zero_point) * scale`

**Storage Requirements**: Scale and zero_point tensors must use buffer storage with width-packed layout. Input/output tensors support both buffer and texture storage with standard axis mapping. Input tensors must be integer types (kByte, kChar, kInt).

# Block-wise Dequantization Implementation
Block-wise dequantization enables fine-grained reconstruction by dividing tensors into blocks and applying separate dequantization parameters to each block. The implementation uses the same key data structures computed in `Dequantize.cpp`:

* **`block_size_vec`**: WHCN-ordered block dimensions converted from PyTorch NCHW layout (e.g., [3,3,2,1] for 3×3×2×1 blocks)
* **`tensor_size_whcn`**: Input tensor dimensions converted to WHCN layout using `utils::make_whcn_ivec4()`
* **`num_blocks_vec`**: Number of blocks per dimension calculated as `tensor_size_whcn / block_size_vec`
* **`block_stride_vec`**: Pre-computed linear strides for block grid indexing `{1, #W, #W*#H, #W*#H*#C}` to enable efficient block ID calculation

The block coordinate calculation uses: `bcoord = tidx / blockSize` where `tidx` is the tensor coordinate in WHCN layout, then the linear block ID is computed as: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`

# Shader Algorithm Overview

## Texture Storage Implementation (`dequantize_texture.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on texture dimensions
- **Local WG Size**: Default with special handling for batch dimension dequantization (Z dimension set to 1 for proper workgroup dispatching when `global_workgroup_size[2] > 1`)

**Block-wise Mode Algorithm**:
The shader processes 3D texture positions where each position represents a texel containing 4 width-packed integer components. For each texel at position `pos`, it calculates a base tensor index `base_tidx = ivec4(pos.x * 4, pos.y, pos.z, 0)` to account for width-packing.

For each of the 4 components in the texel, it computes the actual tensor coordinate: `tidx = ivec4(base_tidx.x + i, base_tidx.y, (foldedZ % C_total), (foldedZ / C_total))` where `foldedZ = pos.z` handles batch-channel folding in 4D tensors and `C_total = numBlocks.z * blockSize.z` represents the total channel dimension.

The block coordinate is calculated using integer division: `bcoord = tidx / blockSize`, then the linear block ID uses pre-computed strides: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

Each integer component is dequantized using its corresponding block's parameters: `value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id])` where `dequantize_val()` applies the formula `(qvalue - zero_point) * scale`. The reconstructed floating-point values are written to the output texel with proper type handling for double precision outputs.

## Buffer Storage Implementation (`dequantize_buffer.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on buffer element count
- **Local WG Size**: Default sizing without special constraints

**Block-wise Mode Algorithm**:
The shader processes linear buffer indices using `gl_GlobalInvocationID.x` as the output buffer index. It converts this to tensor coordinates using `bufi_to_tidx(out_bufi, t_out_strides, out_dim_order)` which handles the buffer-to-tensor index mapping with proper stride calculations.

For each element, it computes the block coordinate directly: `bcoord = out_tidx / blockSize` where `out_tidx` is the 4D tensor coordinate in WHCN layout. The linear block ID calculation uses the same pre-computed stride approach: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

The quantized integer value is loaded using the corresponding input buffer index: `qvalue = t_in[in_bufi]` where `in_bufi = tidx_to_bufi(out_tidx, t_in_strides)`. Dequantization applies the block-specific parameters: `value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id])` to reconstruct the original floating-point value.

**Future Improvements**: Dynamic workgroup sizing based on block dimensions

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

[ghstack-poisoned]

Author: morelos

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

@@ -53,6 +53,17 @@ $if MODE == "per_channel":
     int quant_min;
     int quant_max;
   };
+$if MODE == "block_wise":
+  ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")}
+  ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")}
+
+  layout(push_constant) uniform restrict Block {
+    ivec4 blockSize;     // bW, bH, bC, bN
+    ivec4 numBlocks;     // tW/bW, tH/bH, tC/bC, tN/bN
+    ivec4 blockStride;   // pre-computed linear strides for the block grid
+    int quant_min;
+    int quant_max;
+  };
 
 ${layout_declare_ubo(B, "int", "out_numel")}
 ${layout_declare_ubo(B, "ivec4", "t_in_sizes")}
@@ -71,68 +82,60 @@ const lowp ivec4 out_dim_order = unhash_dim_order(out_layout);
 const lowp ivec4 in_dim_order = unhash_dim_order(in_layout);
 
 /*
- * DEQUANTIZATION SHADER (BUFFER STORAGE)
- *
- * This shader converts n-bit integer tensor values back to floating-point representations
- * using pre-computed quantization parameters (scale and zero_point). The dequantization
- * reconstructs the original floating-point values from their discrete integer representations
- * with minimal precision loss.
- *
- * ALGORITHM:
- * 1. Load quantized integer value from buffer
- * 2. Apply dequantization formula: value = (qvalue - zero_point) * scale
- * 3. Store reconstructed floating-point value to output buffer
- *
- * WORKGROUP CONFIGURATION:
- * - Per-Tensor Mode:
- *   - Global WG Size: {num_elements, 1, 1} (one thread per tensor element)
- *   - Local WG Size: Default (typically {64, 1, 1} or based on global WG size)
- * - Per-Token Mode:
- *   - Global WG Size: {num_elements, 1, 1} (one thread per tensor element)
- *   - Local WG Size: Default (typically {64, 1, 1} or based on global WG size)
- *
- * SUPPORTED CONFIGURATIONS:
- * - Buffer Storage: Uses linear buffer indexing with stride-based tensor access
- * - Per-Tensor: Supports any tensor layout through stride calculations and dimension ordering
- * - Per-Token: Supports only width packed tensors (packed_dim = 0) and standard axis mapping
- * - Scale/zero_point tensors: Must use buffer storage with width packing (packed_dim = 0)
- *
- * DEQUANTIZATION FORMULA VISUALIZATION:
- * For integer range [quant_min, quant_max] mapped back to [min_val, max_val]:
- *
- * Integer Domain:           Floating Point Domain:
- * quant_min ──────────────► min_val
- *    │                         │
- *    │    scale = (max_val - min_val) / (quant_max - quant_min)
- *    │    zero_point = quant_min - round(min_val / scale)
- *    │                         │
- * quant_max ──────────────► max_val
- *
- * Dequantization Process:
- * Input: -103 (int8)
- * Step 1: qvalue - zero_point = -103 - (-128) = 25
- * Step 2: result * scale = 25 * 0.1 = 2.5
- * Output: 2.5 (float)
- *
- * PER-TENSOR DEQUANTIZATION:
- * - Single scale and zero_point values for entire tensor
- * - All elements use same dequantization parameters
- * - Parameters passed as push constants for efficiency
- * - Formula: value = (qvalue - zero_point) * scale
- *
- * PER-TOKEN DEQUANTIZATION:
- * - Separate scale and zero_point for each token
- * - Token = all elements except last dimension (e.g., for [B,S,H]: B*S tokens of H elements)
- * - Parameters stored in buffer arrays indexed by token_id
- * - Each thread calculates its token_id from tensor coordinates
- * - Formula: value = (qvalue - zero_point[token_id]) * scale[token_id]
- *
- * Token ID calculation for element at tensor index (w, z, y, x):
- * - 4D tensor: token_id = w * (sizes.z * sizes.y) + z * sizes.y + y
- * - 3D tensor: token_id = z * sizes.y + y
- * - 2D tensor: token_id = y
- * - 1D tensor: token_id = 0
- */
+  Dequantization Shader (Buffer Storage)
+    This shader converts n-bit integer tensor values back to floating-point representations
+    using pre-computed quantization parameters (scale and zero_point). The dequantization
+    reconstructs the original floating-point values from their discrete integer representations
+    with minimal precision loss.
+
+  Important Considerations:
+    (+) All input tensors are assumed to be WIDTH_PACKED (i.e., contiguous in the last dimension)
+    (+) The axis map layout is assumed to be a standard layout for scales and zero_points
+    (++) The scale and zero_point tensors must be implemented as buffers
+
+  Workgroup Configuration:
+  - dequantize_per_tensor
+      This mode reverses the uniform quantization applied across the entire tensor by using the
+      single scale and zero_point values to convert quantized integer values back to their original
+      floating-point representation.
+
+    (*) global_wg_size: default
+    (*) local_wg_size: default
+
+  - dequantize_per_token
+      This mode reverses the quantization applied individually to each token (or element) in the
+      input by using separate scale and zero_point values for each token. For a tensor of shape
+      [B, S, H], it applies the inverse transformation token-wise across the B*S tokens, converting
+      quantized values back to their original floating-point representation for each group of H
+      elements independently.
+
+    (*) global_wg_size: default
+    (*) local_wg_size: default
+
+  - dequantize_per_channel
+      This mode reverses the quantization applied separately to each channel of the input tensor
+      by using distinct scale and zero_point values for each channel. For a tensor of shape
+      [B, C, H, W] with axis = 1, it applies the inverse transformation channel-wise across the C
+      channels, converting quantized values back to their original floating-point representation
+      independently for each channel.
+
+    (*) global_wg_size: default
+    (*) local_wg_size: default
+
+  - dequantize_block_wise
+      This mode reverses the block-wise quantization applied to groups of elements by using separate
+      scale and zero_point values for each block. Equivalent to dequantize_affine, it applies the
+      inverse affine transformation per block to convert quantized values back to their original
+      floating-point representation. For example, if the tensor shape is [6, 9, 4] and
+      blockSize = [3, 3, 2], the tensor is divided into 12 blocks, each containing 18 elements,
+      and dequantization is performed independently on each block.
+
+    (*) global_wg_size: default
+    (*) local_wg_size: default
+
+  Dequantization Formula:
+    value = (qvalue - zero_point) * scale
+*/
 
 #ifdef per_tensor
 
@@ -187,7 +190,7 @@ void dequantize_per_token() {
   t_out[out_bufi] = value;
 }
 
-#else // per_channel
+#elif defined(per_channel)
 
 void dequantize_per_channel() {
   const int out_bufi = int(gl_GlobalInvocationID.x);
@@ -226,6 +229,29 @@ void dequantize_per_channel() {
   t_out[out_bufi] = value;
 }
 
+#else // block_wise
+
+void dequantize_block_wise() {
+  const int out_bufi = int(gl_GlobalInvocationID.x);
+
+  if (out_bufi >= out_numel) {
+    return;
+  }
+
+  const ivec4 out_tidx = bufi_to_tidx(out_bufi, t_out_strides, out_dim_order);
+  const int in_bufi = tidx_to_bufi(out_tidx, t_in_strides);
+
+  IN_T qvalue = t_in[in_bufi];
+
+  const ivec4 bcoord = out_tidx / blockSize;
+
+  const int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w;
+
+  const OUT_T value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id]);
+
+  t_out[out_bufi] = value;
+}
+
 #endif
 
 void main() {

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

@@ -19,3 +19,5 @@ dequantize_buffer:
       MODE: per_token
     - NAME: dequantize_per_channel_buffer
       MODE: per_channel
+    - NAME: dequantize_block_wise_buffer
+      MODE: block_wise

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

@@ -56,6 +56,17 @@ $if MODE == "per_channel":
     int quant_min;
     int quant_max;
   };
+$if MODE == "block_wise":
+  ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")}
+  ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")}
+
+  layout(push_constant) uniform restrict Block {
+    ivec4 blockSize;     // bW, bH, bC, bN
+    ivec4 numBlocks;     // tW/bW, tH/bH, tC/bC, tN/bN
+    ivec4 blockStride;   // pre-computed linear strides for the block grid
+    int quant_min;
+    int quant_max;
+  };
 
 ${layout_declare_ubo(B, "ivec3", "t_in_limits")}
 ${layout_declare_ubo(B, "ivec3", "t_out_limits")}
@@ -201,7 +212,7 @@ void dequantize_per_token() {
   write_texel(t_out, pos, outtex);
 }
 
-#else // per_channel
+#elif defined(per_channel)
 
 void dequantize_per_channel() {
   const ivec3 pos = ivec3(gl_GlobalInvocationID);
@@ -292,6 +303,39 @@ void dequantize_per_channel() {
   write_texel(t_out, pos, outtex);
 }
 
+#else // block_wise
+
+void dequantize_block_wise() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  if (any(greaterThanEqual(pos, t_in_limits)))
+    return;
+
+  IVEC4_T intex = load_texel(t_in, pos);
+  FVEC4_T outtex;
+
+  ivec4 base_tidx = ivec4(pos.x * 4, pos.y, pos.z, 0);
+  int foldedZ = pos.z;
+
+  int C_total = numBlocks.z * blockSize.z;
+
+  [[unroll]] for (int i = 0; i < 4; ++i) {
+    ivec4 tidx = ivec4(base_tidx.x + i, base_tidx.y, (foldedZ % C_total), (foldedZ / C_total));
+
+    ivec4 bcoord = tidx / blockSize;
+    int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w;
+
+    IN_T qvalue = IN_T(intex[i]);
+    OUT_T value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id]);
+    $if OUT_DTYPE == "double":
+      outtex[i] = float(value);
+    $else:
+      outtex[i] = value;
+  }
+
+  write_texel(t_out, pos, outtex);
+}
+
 #endif
 
 void main() {

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

@@ -19,3 +19,5 @@ dequantize_texture:
       MODE: per_token
     - NAME: dequantize_per_channel_texture3d
       MODE: per_channel
+    - NAME: dequantize_block_wise_texture3d
+      MODE: block_wise