slightly faster inference

inference.py

@@ -161,6 +161,7 @@ def forward_pass_quantized_linear(
         "num_input_groups",
         "num_input_groups_next_power_of_2",
         "compute_in_fp32",
+        "has_bias",
     ],
 )
 @triton.jit
@@ -180,6 +181,7 @@ def _aqlm_gemv_simple(
     num_input_groups: tl.constexpr,
     num_input_groups_next_power_of_2: tl.constexpr,
     compute_in_fp32: tl.constexpr,
+    has_bias: tl.constexpr,
     UNUSED: tl.constexpr,
 ):
     # variables ending with "_i" mean "for i-th output unit"
@@ -188,11 +190,11 @@ def _aqlm_gemv_simple(
     # Stage 1: load input data
     input_vec = tl.load(
         input_vec_ptr
-        + tl.arange(0, num_input_groups_next_power_of_2)[:, None, None] * in_group_size
-        + tl.arange(0, in_group_size)[None, None, :],
-        mask=tl.arange(0, num_input_groups_next_power_of_2)[:, None, None] < num_input_groups,
+        + tl.arange(0, num_input_groups_next_power_of_2)[:, None, None, None] * in_group_size
+        + tl.arange(0, in_group_size)[None, None, None, :],
+        mask=tl.arange(0, num_input_groups_next_power_of_2)[:, None, None, None] < num_input_groups,
     )
-    # [in_features//in_group_size, 1, group_size]
+    # [in_features//in_group_size, 1, 1, group_size]
     # Note: we could simply load input_vec then reshape
     #     input_vec = tl.load(input_vec_ptr + tl.arange(0, in_features))  # [in_features]
     #     input_vec = tl.view(input_vec, [num_input_groups, 1, in_group_size])
@@ -237,19 +239,17 @@ def _aqlm_gemv_simple(
         weights_i = weights_i.to(tl.float32)
         input_vec = input_vec.to(tl.float32)
     # ^-- [in_features // in_group_size, num_codebooks, out_group_size, in_group_size]
-    weights_i = tl.sum(weights_i, axis=1)  # sum codebooks as per additive quantization
-    # ^-- [in_features // in_group_size, out_group_size, in_group_size]
 
     if out_group_size == 1:
         scale = tl.load(scales_ptr + pid).to(weights_i.dtype)  # scalar
         output_i = tl.sum(weights_i * input_vec) * scale
-        if bias_ptr:
+        if has_bias:
             output_i += tl.load(bias_ptr + pid).to(weights_i.dtype)
         tl.store(output_vec_ptr + pid, output_i.to(input_vec.dtype))
     else:
         output_i = tl.sum(tl.sum(weights_i, axis=2) * input_vec, axis=0)  # [out_group_size]
         output_i *= tl.load(scales_ptr + pid).to(weights_i.dtype)
-        if bias_ptr:
+        if has_bias:
             output_i += tl.load(bias_ptr + pid).to(weights_i.dtype)
         tl.store(output_vec_ptr + pid * out_group_size + tl.arange(0, out_group_size), output_i.to(input_vec.dtype))
 
@@ -296,6 +296,7 @@ def aqlm_gemv_simple(
         num_input_groups,
         next_power_of_2(num_input_groups),
         compute_in_fp32,
+        bias is not None,
     )
 
     return output_vec
@@ -339,6 +340,7 @@ def aqlm_gemm_stupid(
             num_input_groups,
             next_power_of_2(num_input_groups),
             compute_in_fp32,
+            bias is not None,
         )
 
     return output