[ET-VK] Address PR review: drop 8w/8da8w, add int8 coopmat gate, fix batch guard

Three changes addressing bot review on the quantized coopmat dispatch PR:

1. Drop 8w and 8da8w from the PR scope (reviewer request: keep this PR
   focused on int4 quantization). Removes linear_dq8ca_q8csw_tiled shader,
   the two 8w/8da8w YAML variants from the consolidated templates, all
   dispatch logic in QuantizedLinear.cpp, the op registration stack
   (custom_ops_lib.py / op_registry.py / patterns/quantized_linear.py),
   test_dq8ca_q8csw_linear.cpp, and the "8da8w" llm_config qmode entry.

2. Add int8 cooperative-matrix capability gate (Copilot #6): the existing
   can_use_q4gsw_coopmat() gate checked only supports_cooperative_matrix()
   which reports fp16 support, not whether VK_COMPONENT_TYPE_SINT8_KHR is
   enumerated. Adds supports_int8_cooperative_matrix() to Adapter.h (backed
   by a one-time vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR query in
   Device.cpp) and gates the dq8ca_q4gsw coopmat dispatch on it.

3. Reject batched outputs before coopmat dispatch (Codex P1 #8): adds a
   dim_of(output) > 2 check to can_use_q4gsw_coopmat(), matching the guard
   in is_coopmat_eligible(). The coopmat shaders use gl_WorkGroupID.xy only;
   a rank-3 output with aligned last two dims would silently miscompute.

Also updates the YAML/GLSL template comments to reflect single-variant scope,
fixes linear_qw_coopmat.yaml tile geometry to the tsweep optimum (128x64,
2x2 subgroup grid, sg32; ~+25% prefill on M5 EVT1 over the prior 128x128),
and updates quantized_linear_global/local_wg_size to use per-shader wg_size
from CoopmatTileDims rather than the hardcoded kCoopmatInvocations (which was
wrong after the tile change reduced q4gsw WG size from 256 to 128).

Author: xuyanwen2012

backends/vulkan/_passes/tag_memory_meta_pass.py

@@ -416,7 +416,10 @@ def constrain_op_arg_repset(self, arg_i: int, op_repsets: utils.OpRepSets) -> No
         over the upstream source's existing layout, falling back to the source only when
         downstream tracing does not fully constrain the repset.
         """
-        if self.force_fp16:
+        # An explicit whole-graph buffer override (e.g. ET_VK_FORCE_BUFFER) must win
+        # over the force_fp16 prefer-texture stopgap, otherwise buffer storage can
+        # never be selected when force_fp16 is enabled.
+        if self.force_fp16 and self.default_storage != VkStorageType.BUFFER:
             op_repsets.try_constrain_with_arg_repset(arg_i, utils.ANY_TEXTURE)
 
         # First, trace downstream users to discover what layout they prefer.

backends/vulkan/custom_ops_lib.py

@@ -307,41 +307,6 @@ def linear_dq8ca_q4gsw(
 linear_dq8ca_q4gsw_op = getattr(getattr(torch.ops, namespace), name)
 
 
-#######################
-## linear_dq8ca_q8csw ##
-#######################
-
-
-def linear_dq8ca_q8csw(
-    x: torch.Tensor,
-    input_scale: torch.Tensor,
-    input_zero_point: torch.Tensor,
-    weights: torch.Tensor,
-    weight_sums: torch.Tensor,
-    weight_scales: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
-):
-    # Per-channel symmetric INT8 weight: dequant = weight.to(fp) * scales (per output channel)
-    weights_dq = weights.to(x.dtype) * weight_scales.unsqueeze(-1)
-    return torch.nn.functional.linear(x, weights_dq, bias)
-
-
-name = "linear_dq8ca_q8csw"
-lib.define(
-    f"""
-            {name}(
-                Tensor input,
-                Tensor input_scales,
-                Tensor input_zp,
-                Tensor weights,
-                Tensor weight_sums,
-                Tensor weight_scales,
-                Tensor? bias = None) -> Tensor
-            """
-)
-lib.impl(name, linear_dq8ca_q8csw, "CompositeExplicitAutograd")
-linear_dq8ca_q8csw_op = getattr(getattr(torch.ops, namespace), name)
-
 #################
 ## qaqw_linear ##
 #################

backends/vulkan/op_registry.py

@@ -481,23 +481,6 @@ def register_linear_dq8ca_q4gsw():
     )
 
 
-@update_features(exir_ops.edge.et_vk.linear_dq8ca_q8csw.default)
-def register_linear_dq8ca_q8csw():
-    return OpFeatures(
-        inputs_storage=[
-            utils.CONTIGUOUS_ANY,  # input
-            utils.WIDTH_PACKED_TEXTURE,  # input_scale
-            utils.WIDTH_PACKED_TEXTURE,  # input_zero_point
-            utils.NO_STORAGE,  # weight (prepacked)
-            utils.NO_STORAGE,  # weight_sums (prepacked)
-            utils.NO_STORAGE,  # weight_scales (prepacked)
-            utils.NO_STORAGE,  # bias (prepacked)
-        ],
-        inputs_dtypes=utils.FP_T,
-        supports_prepacking=True,
-    )
-
-
 # =============================================================================
 # QuantizeDequantize.cpp
 # =============================================================================

backends/vulkan/patterns/quantized_linear.py

@@ -227,14 +227,11 @@ def is_weight_pergroup_quantized(self) -> bool:
     def is_weight_perchannel_quantized(self) -> bool:
         weight_shape = self.weight_node.meta["val"].shape
         scales_shape = self.weight_scales_node.meta["val"].shape
-        # Standard PT2E per-channel: scales is 1D [N].
-        if len(scales_shape) == 1:
-            return scales_shape[0] == weight_shape[-2]
-        # torchao source-transform with PerAxis(0) produces 2D [N, 1] (a
-        # single "group" covering the whole row). Treat that as per-channel.
-        if len(scales_shape) == 2 and scales_shape[-1] == 1:
-            return scales_shape[-2] == weight_shape[-2]
-        return False
+        if len(scales_shape) != 1:
+            return False
+
+        # scales should have same size as weight's output channels dim
+        return scales_shape[0] == weight_shape[-2]
 
     def is_input_static_per_tensor_quantized(self) -> bool:
         if self.dequantize_input_node is None:
@@ -492,85 +489,6 @@ def make_linear_dq8ca_q4gsw_op(
     match.output_node.replace_all_uses_with(qlinear_node)
 
 
-def make_linear_dq8ca_q8csw_op(
-    ep: ExportedProgram,
-    graph_module: torch.fx.GraphModule,
-    match: QuantizedLinearMatch,
-    weight_tensor: torch.Tensor,
-    weight_scales_tensor: torch.Tensor,
-):
-    # Per-channel symmetric INT8 weight: no group_size, no nibble packing.
-    # Align width to 4 so GPU shader reads don't go OOB.
-    utils.align_width_and_update_state_dict(
-        ep,
-        match.weight_node,
-        weight_tensor,
-        align_to=1,
-        force_update=True,
-    )
-
-    # torchao source-transform produces 2D [N, 1] scales; squeeze to 1D [N]
-    # so the runtime sees the same shape as the standard PT2E per-channel
-    # path.
-    if weight_scales_tensor.dim() == 2 and weight_scales_tensor.shape[-1] == 1:
-        weight_scales_tensor = weight_scales_tensor.squeeze(-1).contiguous()
-
-    utils.align_width_and_update_state_dict(
-        ep,
-        match.weight_scales_node,
-        weight_scales_tensor,
-        align_to=1,
-        force_update=True,
-    )
-
-    if match.bias_node is not None:
-        bias_tensor = get_param_tensor(ep, match.bias_node)
-        if bias_tensor is not None:
-            utils.align_width_and_update_state_dict(ep, match.bias_node, bias_tensor)
-
-    # Pre-compute per-output-channel weight sums for input zero-point
-    # correction during integer accumulation.
-    first_graph_node = list(graph_module.graph.nodes)[0]
-    with graph_module.graph.inserting_before(first_graph_node):
-        weight_tensor_name = utils.get_tensor_name(ep, match.weight_node)
-        sum_per_output_channel = weight_tensor.sum(dim=1).to(torch.int32).contiguous()
-        # Pad OC to multiple of 4 to keep shader loads in-bounds
-        oc = sum_per_output_channel.shape[0]
-        if oc % 4 != 0:
-            num_padding = 4 - (oc % 4)
-            sum_per_output_channel = F.pad(
-                sum_per_output_channel, (0, num_padding)
-            ).contiguous()
-
-        sums_name = weight_tensor_name + "_sums"
-        sums_name = sums_name.replace(".", "_")
-        weight_sums_node = create_constant_placeholder(
-            exp_program=ep,
-            graph=graph_module.graph,
-            kind=InputKind.PARAMETER,
-            name=sums_name,
-            data=sum_per_output_channel,
-        )
-
-    with graph_module.graph.inserting_before(match.output_node):
-        qlinear_node = graph_module.graph.create_node(
-            "call_function",
-            exir_ops.edge.et_vk.linear_dq8ca_q8csw.default,
-            args=(
-                match.pattern_input_node,
-                match.input_scales_node,
-                match.input_zeros_node,
-                match.weight_node,
-                weight_sums_node,
-                match.weight_scales_node,
-                match.bias_node,
-            ),
-        )
-
-    qlinear_node.meta["val"] = match.output_node.meta["val"]
-    match.output_node.replace_all_uses_with(qlinear_node)
-
-
 def make_linear_q8ta_q8csw_custom_op(
     ep: ExportedProgram,
     graph_module: torch.fx.GraphModule,
@@ -752,13 +670,6 @@ def replace_quantized_linear_patterns(
         make_linear_dq8ca_q4gsw_op(
             ep, graph_module, match, weight_tensor, weight_scales_tensor
         )
-    elif (
-        match.is_input_dynamic_perchannel_quantized()
-        and match.is_weight_perchannel_quantized()
-    ):
-        make_linear_dq8ca_q8csw_op(
-            ep, graph_module, match, weight_tensor, weight_scales_tensor
-        )
     elif (
         match.is_input_static_per_tensor_quantized()
         and match.is_weight_perchannel_quantized()