Update

[ghstack-poisoned]

Author: ssjia

backends/vulkan/patterns/quantized_linear.py

@@ -350,13 +350,10 @@ def make_linear_q4gsw_op(
         ep, match.weight_scales_node, "4 bit linear scales"
     ):
         weight_scales_tensor = weight_scales_tensor.transpose(0, 1).contiguous()
-        # Align to multiple of 8 to ensure that data loads from the weight scales
-        # tensor do not go out of bounds. Each thread computes 8 output channels.
         utils.align_width_and_update_state_dict(
             ep,
             match.weight_scales_node,
             weight_scales_tensor,
-            align_to=8,
             force_update=True,
         )
 
@@ -414,7 +411,6 @@ def make_linear_dq8ca_q4gsw_op(
             ep,
             match.weight_scales_node,
             weight_scales_tensor,
-            align_to=1,
             force_update=True,
         )
 

backends/vulkan/patterns/weight_packing_utils.py

@@ -13,19 +13,15 @@
 def pack_4bit_weight_tensor(
     weight_tensor: torch.Tensor,
     *,
-    even_value_high_nibble: bool = False,
     inner_dim_padding: Optional[int] = 8,
 ) -> torch.Tensor:
     """Pack signed 4-bit values stored in int8 into uint8 byte pairs.
 
     The input tensor stores one quantized value per byte in the range [-8, 7].
     The returned tensor stores two 4-bit values per byte along the innermost dim.
 
-    By default this preserves the legacy linear q4gsw packing convention:
+    This preserves the legacy linear q4gsw packing convention:
       packed_byte = (odd_val + 8) << 4 | (even_val + 8)
-
-    Set `even_value_high_nibble=True` for the embedding q4gsw convention:
-      packed_byte = (even_val + 8) << 4 | (odd_val + 8)
     """
     min_val, max_val = weight_tensor.min().item(), weight_tensor.max().item()
     assert (
@@ -52,6 +48,4 @@ def pack_4bit_weight_tensor(
     shifted_weight_tensor = weight_tensor.to(dtype=torch.uint8) + 8
     even_values = shifted_weight_tensor[:, ::2]
     odd_values = shifted_weight_tensor[:, 1::2]
-    if even_value_high_nibble:
-        return even_values << 4 | odd_values
     return odd_values << 4 | even_values

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

@@ -24,22 +24,13 @@ namespace vkcompute {
 
 namespace {
 
-/*
- * Currently, only a few shader variants are generated for "compatible" tensor
- * dtype / scalar dtype pairs. In particular float/half tensor + float scalar,
- * and int32 tensor + int32 scalar. This function coerces the scalar dtype so
- * that a "higher precision" tensor dtype (i.e. float, half) can be used with a
- * "lower precision" scalar dtype (i.e. int32) with the current generated shader
- * variants. "Lower precision" tensor dtype (i.e. int32) + "higher precision"
- * scalar dtype (i.e. float/half) are currently not supported because the shader
- * currently casts the scalar to the tensor dtype before computation. Though
- * these combinations are not expected to be needed, if they are required in the
- * future, update the shaders with more robust dtype handling.
- */
 vkapi::ScalarType resolve_scalar_extract_dtype(
     ComputeGraph& graph,
     const ValueRef scalar,
     const vkapi::ScalarType tensor_dtype) {
+  // For float tensors, ensure that the scalar argument is extracted as a float
+  // to avoid having to generate additional shader variants for float/half
+  // tensor + int scalar.
   if (tensor_dtype == vkapi::kFloat || tensor_dtype == vkapi::kHalf) {
     return vkapi::kFloat;
   }

backends/vulkan/test/op_tests/cases.py

@@ -2247,6 +2247,6 @@ def get_eq_scalar_inputs():
         "utils::kWidthPacked",
         "utils::kChannelsPacked",
     ]
-    test_suite.dtypes = ["at::kInt"]
+    test_suite.dtypes = ["at::kInt", "at::kFloat"]
     test_suite.data_gen = "make_seq_tensor"
     return test_suite