Update

[ghstack-poisoned]

Author: vkuzo

.github/workflows/torchao_experimental_test.yml

@@ -36,7 +36,7 @@ jobs:
           # Install executorch first because it installs its own version
           # of torch and torchao, which we do not want to use
           pip install executorch
-          pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu" --force-reinstall
+          pip install torch==2.7.0.dev20250311 --index-url "https://download.pytorch.org/whl/nightly/cpu" --force-reinstall
           pip install numpy
           pip install pytest
           pip install parameterized

benchmarks/mx_formats/cast_bench.py

@@ -6,7 +6,7 @@
 from torch._inductor.utils import do_bench_using_profiling
 
 from torchao.prototype.mx_formats.custom_cast import (
-    to_mxfp8_dim1,
+    triton_to_mxfp8_dim1,
 )
 from torchao.prototype.mx_formats.mx_tensor import to_mx
 
@@ -172,12 +172,12 @@ def run(
         bps = (bytes_r + bytes_w) / (time_us / 1e6)
 
     elif mode == "dim1_mx_triton":
-        y_d1, s_d1 = to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE)
+        y_d1, s_d1 = triton_to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE)
 
         for _ in range(2):
-            __ = to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE)
+            __ = triton_to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE)
         time_us = benchmark_cuda_function_in_microseconds(
-            lambda x, b: to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE),
+            lambda x, b: triton_to_mxfp8_dim1(x, inner_block_size=BLOCK_SIZE),
             x,
             BLOCK_SIZE,
         )

test/prototype/mx_formats/test_custom_cast.py

@@ -26,12 +26,11 @@
     get_bits,
     pack_uint4,
     pack_uint6,
-    # TODO(before land): better name?
-    to_mxfp8_dim1,
-    to_mxfp8_dim1_reference,
     triton_f4_to_bf16,
     triton_f6_e2m3_to_bf16,
     triton_f6_e3m2_to_bf16,
+    triton_to_mxfp8_dim1,
+    triton_to_mxfp8_dim1_reference,
     unpack_uint4,
 )
 from torchao.prototype.mx_formats.fp_format_spec import (
@@ -460,9 +459,11 @@ def test_fp6_e3m2_pack_unpack():
 )
 @pytest.mark.parametrize("M", (256, 2048))
 @pytest.mark.parametrize("K", (256, 2048))
+# @pytest.mark.parametrize("M", (256,))
+# @pytest.mark.parametrize("K", (256,))
 def test_triton_mxfp8_dim1(M, K):
     x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
-    x_mx_ref, x_s_ref = to_mxfp8_dim1_reference(x, block_size=32)
-    x_mx_t, x_s_t = to_mxfp8_dim1(x, inner_block_size=32)
+    x_mx_ref, x_s_ref = triton_to_mxfp8_dim1_reference(x, block_size=32)
+    x_mx_t, x_s_t = triton_to_mxfp8_dim1(x, inner_block_size=32)
     torch.testing.assert_close(x_mx_t, x_mx_ref, rtol=0, atol=0)
     torch.testing.assert_close(x_s_t, x_s_ref, rtol=0, atol=0)

torchao/experimental/kernels/cpu/aarch64/embedding/embedding.h

@@ -10,8 +10,10 @@
 
 #include <arm_neon.h>
 #include <torchao/experimental/kernels/cpu/aarch64/bitpacking/bitpack.h>
+#include <torchao/experimental/kernels/cpu/aarch64/linear/pack_weights.h>
 #include <torchao/experimental/kernels/cpu/aarch64/macro.h>
 #include <cassert>
+#include <vector>
 
 namespace torchao::kernels::cpu::aarch64::embedding {
 
@@ -177,6 +179,7 @@ inline void embedding_(
     if (weight_zeros != nullptr) {
       zero = weight_zeros[group_idx];
     }
+
     internal::vec_dequantize_and_store_16_values(out + i, qvals0, scale, zero);
     internal::vec_dequantize_and_store_16_values(
         out + i + 16, qvals1, scale, zero);
@@ -322,6 +325,61 @@ inline void pack_embedding_weight_qvals(
       qvals + index * embedding_dim);
 }
 
+// Embedding op that shares weights with unembedding linear op
+template <int weight_nbit, int nr, int kr, int sr>
+inline void shared_embedding(
+    // Output
+    float* out,
+    // Inputs
+    const void* packed_weights,
+    int n,
+    int k,
+    int group_size,
+    bool has_weight_zeros,
+    bool has_bias,
+    int index) {
+  assert(k % group_size == 0);
+  assert(group_size % 16 == 0);
+
+  int groups_per_k = k / group_size;
+  std::vector<int8_t> weight_qvals(k * nr);
+  std::vector<float> weight_scales(groups_per_k * nr);
+  std::vector<int8_t> weight_zeros(groups_per_k * nr);
+  std::vector<float> bias(nr);
+
+  // Set n_idx to multiple of nr that is at most index
+  // j is index of "index" in nr group
+  int n_idx = index / nr;
+  n_idx = n_idx * nr;
+  int j = index - n_idx;
+
+  torchao::kernels::cpu::aarch64::linear::packing::
+      unpack_weights_at_n_idx<weight_nbit, nr, kr, sr>(
+          weight_qvals.data(),
+          weight_scales.data(),
+          has_weight_zeros ? weight_zeros.data() : nullptr,
+          has_bias ? bias.data() : nullptr,
+          n_idx,
+          n,
+          k,
+          group_size,
+          has_weight_zeros,
+          has_bias,
+          packed_weights);
+
+  // Dequantize and store to output (size k)
+  int8x16_t qvals;
+  for (int i = 0; i < k; i += 16) {
+    qvals = vld1q_s8(weight_qvals.data() + j * k + i);
+    float scale = weight_scales[j * groups_per_k + i / group_size];
+    float zero = 0.0;
+    if (has_weight_zeros) {
+      zero = weight_zeros[j * groups_per_k + i / group_size];
+    }
+    internal::vec_dequantize_and_store_16_values(out + i, qvals, scale, zero);
+  }
+}
+
 } // namespace torchao::kernels::cpu::aarch64::embedding
 
 #endif // defined(__aarch64__) || defined(__ARM_NEON)

torchao/experimental/kernels/cpu/aarch64/linear/channelwise_8bit_activation_groupwise_lowbit_weight_1x1x32_f32_neondot-impl.h

@@ -10,6 +10,7 @@
 
 #include <torchao/experimental/kernels/cpu/aarch64/bitpacking/bitpack.h>
 #include <torchao/experimental/kernels/cpu/aarch64/linear/channelwise_8bit_activation_prepare_activation_data_1xk_f32-impl.h>
+#include <torchao/experimental/kernels/cpu/aarch64/linear/pack_weights.h>
 #include <torchao/experimental/kernels/cpu/aarch64/quantization/quantize.h>
 #include <torchao/experimental/kernels/cpu/aarch64/reduction/reduction.h>
 #include <cassert>
@@ -149,8 +150,8 @@ void kernel_impl(
           int32_t activation_qvals_sum = *((int32_t*)activation_ptr);
           activation_ptr += sizeof(int32_t);
 
-          int8_t weight_zero = *((int8_t*)weight_data_byte_ptr);
-          weight_data_byte_ptr += sizeof(int8_t);
+          int32_t weight_zero = *((int32_t*)weight_data_byte_ptr);
+          weight_data_byte_ptr += sizeof(int32_t);
 
           res += (weight_scale * activation_scale) *
               (qval_dot - (activation_zero * weight_qvals_sum) -
@@ -190,31 +191,14 @@ size_t inline weight_data_size_impl(
     int weight_nbit,
     bool has_weight_zeros,
     bool has_bias) {
-  assert(k % group_size == 0);
-  assert(k % 32 == 0);
-  int groups_per_col = k / group_size;
-  int col_size = 0;
-
-  // qvals
-  // (k * weight_bit) bits -> ((k / 8) * weight_bit) bytes
-  col_size += (k / 8) * weight_nbit;
-
-  // scales
-  col_size += sizeof(float) * groups_per_col;
-
-  // qvals_sum
-  col_size += sizeof(int32_t) * groups_per_col;
-
-  // zeros
-  if (has_weight_zeros) {
-    col_size += sizeof(int8_t) * groups_per_col;
-  }
-
-  if (has_bias) {
-    col_size += sizeof(float);
-  }
-
-  return col_size * n;
+  return torchao::kernels::cpu::aarch64::linear::packing::packed_weights_size(
+      n,
+      k,
+      group_size,
+      weight_nbit,
+      has_weight_zeros,
+      has_bias,
+      /*nr*/ 1);
 }
 
 template <int weight_nbit>
@@ -227,56 +211,19 @@ void prepare_weight_data_impl(
     int group_size,
     const int8_t* weight_qvals,
     const float* weight_scales,
+    // Ignored if has_weight_zeros = false
     const int8_t* weight_zeros,
     const float* bias) {
-  assert(k % group_size == 0);
-  assert(group_size % 32 == 0);
-
-  bool has_weight_zeros = (weight_zeros != nullptr);
-  bool has_bias = (bias != nullptr);
-
-  auto weight_data_byte_ptr = (char*)weight_data;
-  constexpr int bytes_per_32_weight_values = 4 * weight_nbit;
-
-  int8x16_t wq0, wq1;
-
-  const int8_t* qvals_ptr = weight_qvals;
-  const float* scales_ptr = weight_scales;
-  const int8_t* zeros_ptr = weight_zeros;
-  const float* bias_ptr = bias;
-
-  for (int n_idx = 0; n_idx < n; n_idx++) {
-    for (int k_idx = 0; k_idx < k; k_idx += group_size) {
-      int32_t group_qvals_sum = 0;
-      for (int i = 0; i < group_size; i += 32) {
-        wq0 = vld1q_s8(qvals_ptr);
-        wq1 = vld1q_s8(qvals_ptr + 16);
-        qvals_ptr += 32;
-
-        group_qvals_sum += vaddlvq_s8(wq0) + vaddlvq_s8(wq1);
-
-        torchao::bitpacking::vec_pack_32_lowbit_values<weight_nbit>(
-            /*packed=*/(uint8_t*)weight_data_byte_ptr,
-            /*unpacked0=*/wq0,
-            /*unpacked1=*/wq1);
-        weight_data_byte_ptr += bytes_per_32_weight_values;
-      }
-      *((float*)weight_data_byte_ptr) = *scales_ptr++;
-      weight_data_byte_ptr += sizeof(float);
-
-      *((int32_t*)weight_data_byte_ptr) = group_qvals_sum;
-      weight_data_byte_ptr += sizeof(int32_t);
-
-      if (has_weight_zeros) {
-        *((int8_t*)weight_data_byte_ptr) = *zeros_ptr++;
-        weight_data_byte_ptr += sizeof(int8_t);
-      }
-    }
-    if (has_bias) {
-      *((float*)weight_data_byte_ptr) = *bias_ptr++;
-      weight_data_byte_ptr += sizeof(float);
-    }
-  }
+  torchao::kernels::cpu::aarch64::linear::packing::
+      pack_weights<weight_nbit, /*nr*/ 1, /*kr*/ 32, /*sr*/ 2>(
+          weight_data,
+          n,
+          k,
+          group_size,
+          weight_qvals,
+          weight_scales,
+          weight_zeros,
+          bias);
 }
 
 } // namespace