[ET-VK] Activation transpose preprocess shaders + dispatch helper (pytorch#20055)

[ET-VK][q4gsw] Add et_vk.q4gsw_linear.default with W_4X8 GEMM + adaptive nc-coop GEMV
Pull Request resolved: pytorch#19996

Adds `et_vk.q4gsw_linear.default` backed by production compute shaders and a single top-level dispatch in `Q4gswLinear.cpp`. The op uses two or three execute nodes per dtype path that share a 6-binding descriptor set layout: a GEMM `DynamicDispatchNode` that self-gates to `{0,0,0}` at M==1, and an adaptive nc-coop GEMV sibling `DynamicDispatchNode` that self-gates to `{0,0,0}` at M!=1. The fp16 path adds a transpose preprocess that also self-gates at M==1. Decode-time `virtual_resize` correctly routes M==1 through the nc-coop GEMV after a prefill M>1 set up the graph. Existing serialized models that reference `et_vk.linear_q4gsw.default` are routed to the same forwarder so the cutover requires no graph-level changes; the old `QuantizedLinear.cpp::linear_q4gsw` body is deleted.

Production shaders under `runtime/graph/ops/glsl/`:
- `q4gsw_linear_gemm__w_4x8_nc` — fp32 GEMM, 4M x 8N per-thread tile, 8x8 LWG, reads the raw `[M, K]` activation directly. The per-k4 weight tile loads as a single `ivec4` covering 2 consecutive 4x4 N-blocks (one 16B aligned LSU read replaces two 8B reads). Aliases the same buffer the prepack writes.
- `q4gsw_linear_gemm__tin__w_4x8_nc` — fp16 GEMM, transposed-input path, 8x4 per-thread tile, 1x128 LWG, reads the pre-transposed vec4 activation produced by the transpose preprocess. Nibble dequant stays in int16 end-to-end before the fp16 FMA, saving one fp32 register and lifting AOC fiber occupancy from 37% to 50% on Adreno 750 (validated at M=256 K=4096 N=4096: 6135.92 us -> 5637.20 us, -8.1% wall-clock).
- `q4gsw_linear_gemv_coop__w_4x8_nc_buffer_{g1w64,g4w16,g8w8}` — cooperative-reduction GEMV reading the same nc-buffer weight payload as the GEMM shaders, with three (NUM_GROUPS, WORKERS_PER_GROUP) decompositions chosen at dispatch time by the production picker based on output N. All keep total threads/WG = 64; each WG produces NUM_GROUPS * 8 outputs, and the K-loop strides by WORKERS_PER_GROUP.

Prepack: `pack_q4_linear_weight__w_4x8_nc_buffer` produces an int SOA nibble buffer; one `ivec4` covers a 4K x 8N block. The fp32 GEMM shader views the buffer as `ivec4`; the fp16 tin GEMM shader reads it as `ivec2`; the nc-coop GEMV shader rebinds the same bytes as scalar int arrays. Scales use a single `prepack_q4_scales` (dtype-matched vec4) for GEMM; GEMV rebinds the same bytes as a gvec2 via `vec_size=2`. The shared nc-buffer payload means prefill and decode both consume one prepack — no dual-format weight memory cost.

Dispatch structure:
- fp32 (`add_q4gsw_linear_w_4x8_node`): two nodes. The GEMM dispatch binds `q4gsw_linear_gemm__w_4x8_nc` unconditionally; its gated global WG gates the dispatch to `{0,0,0}` at M==1. The nc-coop GEMV sibling owns the decode.
- fp16 (`add_q4gsw_linear_tin_w_4x8_node`): three nodes. The transpose preprocess self-gates `{0,0,0}` when M==1 (no-op); the GEMM dispatch binds `q4gsw_linear_gemm__tin__w_4x8_nc` and its gated global WG returns `{0,0,0}` at M==1; the nc-coop GEMV sibling handles M==1.
- nc-coop GEMV picker (`add_q4gsw_linear_nc_coop_gemv_node`): shape-adaptive dispatch that self-gates `{0,0,0}` when M!=1. Heuristic on output N:
  * N <= 1024  -> `q4gsw_linear_gemv_coop__w_4x8_nc_buffer_g1w64`  LWG (1, 1, 64)
  * N <= 4096  -> `q4gsw_linear_gemv_coop__w_4x8_nc_buffer_g4w16`  LWG (1, 4, 16)
  * N >  4096  -> `q4gsw_linear_gemv_coop__w_4x8_nc_buffer_g8w8`   LWG (1, 8, 8)

The adaptive picker replaces sg-GEMV at decode. Cross-device sweep on Galaxy S24 (Adreno 750), S25 (Adreno 830), and Pixel 9 Pro XL (Mali-G715) shows shape-adaptive nc-coop beats sg-GEMV on every LLM-decode shape by 3-56%, with the largest wins at small-N (47-56% at K=2048 N=512 across both Adreno devices). The adaptive picker also unblocks Mali, where sg-GEMV cannot dispatch (subgroupSize=16 != required 64).

The op is registered under both `et_vk.q4gsw_linear.default` and `et_vk.linear_q4gsw.default` for backward compatibility. Also extends `validate_against_reference` in `test/custom_ops/utils.cpp` to handle fp16 tensors.

Demonstration shaders under `test/custom_ops/glsl/` (not linked into the production shader library):
- `q4gsw_linear_gemv__w_4x8_nc{,_nosg}` — sg / nosg GEMV variants from earlier iterations. The sg variant pins subgroupSize to 64 via `VK_EXT_subgroup_size_control` and uses `subgroupBroadcast`; the nosg variant has every thread load its own activation. Both are dispatched by the `test_fpa_q4gsw_linear` benchmark binary's forced-shader selectors `GEMV_W_4X8` and `GEMV_W_4X8_NOSG` (selectors 1 and 2). Retained as a demonstration of the single-thread-per-output GEMV layout the adaptive nc-coop now supersedes at decode.

Buffer-padding fix (latent OOB read on N % 8 != 0): The fp32 GEMM's 16B ivec4 weight load spans two consecutive (k4, n4) ivec2 tiles along N. Pad the SOA weight buffer's row stride to next-even N4 (no-op for the canonical N % 8 == 0 shapes); the prepack shader's existing `(n < N)` branch fills the OOB tiles with the bias-zero nibble pattern (0x88888888u). Adds GEMM accuracy coverage for N=12 and N=20.
ghstack-source-id: 390109017
@exported-using-ghexport

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

Author: pytorchbot

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

@@ -0,0 +1,179 @@
+/*
+ * 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.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define STORAGE ${STORAGE}
+
+layout(std430) buffer;
+
+// Output: W_4X8 block-packed nibble weight, written as 4K x 8N blocks. Each invocation
+// produces one full 4K x 8N block at logical position (k4, n8) — equivalent to
+// 2 consecutive ivec2 tiles at the SAME k4: n4 = 2*n8 (lower) and n4 = 2*n8+1
+// (upper). The 4 ints making up the block are:
+//     [0] = packed_x for (k4, n4 = 2*n8)        (rows N0, N1 of n4_a)
+//     [1] = packed_y for (k4, n4 = 2*n8)        (rows N2, N3 of n4_a)
+//     [2] = packed_x for (k4, n4 = 2*n8+1)      (rows N0, N1 of n4_b)
+//     [3] = packed_y for (k4, n4 = 2*n8+1)      (rows N2, N3 of n4_b)
+//
+// Buffer (nc) form: stored as a flat ivec4 buffer; one block per ivec4 at
+// index `k4 * N8 + n8`, where N8 = N4_padded/2 and N4_padded is the next-even
+// N4. This is byte-identical to writing 4 consecutive ints at scalar index
+// `4*(k4*N8 + n8)` (the legacy 2-tile layout).
+//
+// Buffer (kc dense) form: stored as a flat ivec4 buffer; one block per ivec4
+// at index `n8 * K4 + k4`. Adjacent ivec4s along K cover adjacent k4 (kc-
+// contiguous); adjacent n8 blocks are stride K4 apart.
+//
+// Texture2D (kc dense) form: stored as ivec4 texels. Each texel covers one
+// block; image position is (k4, n8). Adjacent texels along x are adjacent k4
+// (kc-contiguous), supplying the lane-stride reduction pattern of the coop
+// GEMV.
+//
+// Texture2D (nc) form: stored as ivec4 texels. Each texel covers one block;
+// image position is (n8, k4). Adjacent texels along x are adjacent n8 (nc-
+// contiguous). Lets nc-walking consumers route weight reads through the
+// texture cache.
+//
+// Interleaved (dp4a-style) byte-pair layout (same for all forms):
+//   Each byte of .x holds one (N_even, N_odd) nibble pair at a fixed K.
+//   .x byte b (b in {0,1,2,3}) = (N0, K=b) | (N1, K=b) << 4
+//   .y byte b                  = (N2, K=b) | (N3, K=b) << 4
+// The low nibble of each byte is the even-N row and the high nibble is the
+// odd-N row.
+${layout_declare_tensor(B, "w", "t_packed_weight", "int", STORAGE, is_scalar_array=False)}
+// Input: raw [N, K/2] uint8 data read as uint32.
+// Each uint32 holds 8 nibbles = 8 K-values for one N-row.
+// Indexed as t_int4_weight[n * K8 + k8] where K8 = ceil(K/8).
+${layout_declare_tensor(B, "r", "t_int4_weight", "uint", "buffer")}
+
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ivec2 orig_sizes; // {K, N}
+  // Unused — kept so both prepack call sites (buffer and texture2d) can share
+  // an identical push-constant layout. The block-row stride is implicit in
+  // ceil(N/8) on both paths: for buffer this matches N4_padded/2 (where
+  // N4_padded = (N4+1)&~1); for texture2d this is the image's N8 dimension.
+  int n4_pitch;
+};
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+// Returns the packed_x / packed_y uint pair for the (k4, n4) tile.
+void compute_tile_packed(
+    out uint packed_x,
+    out uint packed_y,
+    const int k4,
+    const int n4,
+    const int K8,
+    const int N) {
+  packed_x = 0u;
+  packed_y = 0u;
+
+  for (int ni = 0; ni < 4; ++ni) {
+    const int n = n4 * 4 + ni;
+
+    // k4 * 4 gives the starting K index. We need 4 consecutive K values.
+    // The source has 8 K-nibbles per uint32 at t_int4_weight[n * K8 + k8].
+    // k4 * 4 / 8 = k4 / 2 gives the uint32 index along K.
+    const int k_start = k4 * 4;
+    const int k8_idx = k_start / 8;
+    const uint src_word =
+        (n < N) ? t_int4_weight[n * K8 + k8_idx] : 0x88888888u;
+
+    // Within this uint32, extract 4 nibbles starting at position (k_start % 8)
+    const int nibble_offset = (k_start % 8);
+
+    // Interleaved: byte b holds one nibble from even-N row (low) and one
+    // nibble from odd-N row (high). The low/high selection is by ni parity.
+    // Even-ni rows go into packed_x's or packed_y's low nibble of each byte,
+    // odd-ni rows go into the high nibble.
+    for (int ki = 0; ki < 4; ++ki) {
+      const uint nibble = (src_word >> (4 * (nibble_offset + ki))) & 0xFu;
+      const int bit_offset = 8 * ki + (ni & 1) * 4;
+      if (ni < 2) {
+        packed_x |= nibble << bit_offset;
+      } else {
+        packed_y |= nibble << bit_offset;
+      }
+    }
+  }
+}
+
+void main() {
+  // One invocation = one full 4K x 8N block at logical (k4, n8).
+  const int k4 = int(gl_GlobalInvocationID.x);
+  const int n8 = int(gl_GlobalInvocationID.y);
+
+  const int K = orig_sizes.x;
+  const int N = orig_sizes.y;
+  const int K8 = (K + 7) / 8;
+  const int K4 = K / 4;
+  const int N4 = (N + 3) / 4;
+  // N8 = ceil(N4/2). Both buffer (where N4_padded = (N4+1)&~1, so N4_padded/2
+  // = (N4+1)/2) and texture2d paths use the same dispatch shape.
+  const int N8 = (N4 + 1) / 2;
+
+  if (k4 >= K4 || n8 >= N8) {
+    return;
+  }
+
+  const int n4_a = 2 * n8;
+  const int n4_b = n4_a + 1;
+
+  uint packed_x_a = 0u;
+  uint packed_y_a = 0u;
+  uint packed_x_b = 0u;
+  uint packed_y_b = 0u;
+
+  // Lower tile (n4_a) — always materialized. compute_tile_packed handles
+  // n >= N rows with the 0x88888888u (bias-zero) fallback per row.
+  compute_tile_packed(packed_x_a, packed_y_a, k4, n4_a, K8, N);
+
+  // Upper tile (n4_b). When n4_b >= N4 the entire tile is OOB along N — use
+  // the bias-zero pattern directly so the GEMV/GEMM consumers can safely
+  // read whole blocks even when N4 is odd.
+  if (n4_b < N4) {
+    compute_tile_packed(packed_x_b, packed_y_b, k4, n4_b, K8, N);
+  } else {
+    packed_x_b = 0x88888888u;
+    packed_y_b = 0x88888888u;
+  }
+
+  const ivec4 texel = ivec4(
+      int(packed_x_a),
+      int(packed_y_a),
+      int(packed_x_b),
+      int(packed_y_b));
+
+$if STORAGE == "texture2d":
+  $if WEIGHT_KC == 1:
+    // Texture2D (kc dense). Image position = (k4, n8); adjacent texels along
+    // x cover adjacent k4 (kc-contiguous).
+    imageStore(t_packed_weight, ivec2(k4, n8), texel);
+  $else:
+    // Texture2D (nc). Image position = (n8, k4); adjacent texels along x
+    // cover adjacent n8 (nc-contiguous). Same byte-pair payload as nc-buffer
+    // but stored as an ivec4 image2D so consumers route weight reads through
+    // the texture cache while keeping the nc walking pattern.
+    imageStore(t_packed_weight, ivec2(n8, k4), texel);
+$elif WEIGHT_KC == 1:
+  // Buffer (kc dense) form. One ivec4 per block at index `n8 * K4 + k4`.
+  // Adjacent ivec4s along K cover adjacent k4 (kc-contiguous); adjacent n8
+  // blocks are stride K4 apart. Mirrors the kc Tex2D layout so consumers can
+  // A/B-test SSBO ivec4 reads vs texelFetch on the same byte-pair payload.
+  t_packed_weight[n8 * K4 + k4] = texel;
+$else:
+  // Buffer (nc) form. One ivec4 per block at index `k4 * N8 + n8` — byte-
+  // identical to the legacy 2-ivec2-tile / 4-scalar-int layout because
+  // N4_padded = (N4 + 1) & ~1 is even, so 2 * (k4 * N4_padded + 2*n8)
+  // = 4 * (k4 * N8 + n8).
+  t_packed_weight[k4 * N8 + n8] = texel;
+}

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

@@ -0,0 +1,22 @@
+# 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.
+
+pack_q4_linear_weight__w_4x8:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+    WEIGHT_KC: 0
+  shader_variants:
+    # 4K x 8N weight blocks, interleaved nibble layout (dp4a-style byte pairs,
+    # low nibble = even-N row, high nibble = odd-N row).
+    #
+    #   nc_buffer — N dim contiguous in memory; one ivec4 per (k4, n8) at
+    #               flat index `k4 * N8 + n8`. Adjacent n8 blocks are
+    #               contiguous in memory; adjacent k4 are stride N8.
+    #               Required by the GEMM ivec4 weight load.
+    - NAME: pack_q4_linear_weight__w_4x8_nc_buffer
+      STORAGE: buffer
+      WEIGHT_KC: 0