Build uploaded using `kernels`.

build/torch210-cxx11-cu126-aarch64-linux/{_megablocks_cuda_6e04dec.abi3.so → _megablocks_cuda_a45325d.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d43ea617155587acccc47750e126596b0438c63c7ada6f3607a2ed4603337f72
+oid sha256:77772a8c5e37277e5f14ff74222bfd44b64d78b0ecd36a4c4ce21008acc71d50
 size 15124328

build/torch210-cxx11-cu126-aarch64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_cuda_6e04dec
-ops = torch.ops._megablocks_cuda_6e04dec
+from . import _megablocks_cuda_a45325d
+ops = torch.ops._megablocks_cuda_a45325d
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_cuda_6e04dec::{op_name}"
+    return f"_megablocks_cuda_a45325d::{op_name}"

build/torch210-cxx11-cu126-aarch64-linux/megablocks/__init__.py

@@ -1,10 +1,10 @@
 import ctypes
+import importlib.util
 import sys
-
-import importlib
 from pathlib import Path
 from types import ModuleType
 
+
 def _import_from_path(file_path: Path) -> ModuleType:
     # We cannot use the module name as-is, after adding it to `sys.modules`,
     # it would also be used for other imports. So, we make a module name that

build/torch210-cxx11-cu126-aarch64-linux/ops/histogram_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -47,31 +47,31 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class HistogramBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testTorchHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, 128, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class HistogramBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testTorchHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, 128, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/ops/matmul_benchmark.py

@@ -17,7 +17,7 @@ import unittest
 from .. import stk
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -48,367 +48,367 @@ def log_benchmark(name, arguments, time, std, flops):
     print('=' * 60)
 
 
-class MatmulBenchmark(parameterized.TestCase):
-
-    def build_sparse_matrix(self, x, padded_bins, fhs, ne):
-        blocking = 128
-        padded_tokens, _ = x.size()
-        assert padded_tokens % blocking == 0
-        assert fhs % blocking == 0
-
-        # Offsets for the sparse matrix. All rows have the
-        # same number of nonzero blocks dictated by the
-        # dimensionality of a single expert.
-        block_rows = padded_tokens // blocking
-        blocks_per_row = fhs // blocking
-        offsets = torch.arange(
-            0,
-            block_rows * blocks_per_row + 1,
-            blocks_per_row,
-            dtype=torch.int32,
-            device=x.device,
-        )
-
-        # Indices for the sparse matrix. The indices for
-        # the intermediate matrix are dynamic depending
-        # on the mapping of tokens to experts.
-        column_indices = ops.topology(
-            padded_bins,
-            blocking,
-            block_rows,
-            blocks_per_row,
-        )
-        data = torch.empty(
-            column_indices.numel(),
-            blocking,
-            blocking,
-            dtype=torch.float16,
-            device=x.device,
-        )
-        shape = (padded_tokens, fhs * ne)
-        row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
-        return stk.Matrix(shape, data, row_indices, column_indices, offsets)
-
-    def build_input_matrix(self, sl, hs, ne):
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Assign tokens to experts uniformly.
-        top_expert = torch.arange(0, sl).cuda().int() % ne
-
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
-        return out, padded_bins
-
-    def build_weight_matrix(self, ne, hs, fhs):
-        return torch.randn((hs, ne * fhs)).cuda().half()
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(x, w, topo)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(topo, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradX::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        topo = topo.t()
-
-        def benchmark():
-            return stk.ops.dsd(topo, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(out, w, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        x = x.t()
-
-        def benchmark():
-            return stk.ops.dsd(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-
-        w = w.transpose(1, 2).contiguous()
-        w = w.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd:DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = w.transpose(1, 2).contiguous()
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradX:DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        out = out.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(out, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradW:DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = torch.transpose(w, 1, 2)
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        x = torch.transpose(x, 1, 2)
-
-        def benchmark():
-            return torch.bmm(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
+# class MatmulBenchmark(parameterized.TestCase):
+#
+#     def build_sparse_matrix(self, x, padded_bins, fhs, ne):
+#         blocking = 128
+#         padded_tokens, _ = x.size()
+#         assert padded_tokens % blocking == 0
+#         assert fhs % blocking == 0
+#
+#         # Offsets for the sparse matrix. All rows have the
+#         # same number of nonzero blocks dictated by the
+#         # dimensionality of a single expert.
+#         block_rows = padded_tokens // blocking
+#         blocks_per_row = fhs // blocking
+#         offsets = torch.arange(
+#             0,
+#             block_rows * blocks_per_row + 1,
+#             blocks_per_row,
+#             dtype=torch.int32,
+#             device=x.device,
+#         )
+#
+#         # Indices for the sparse matrix. The indices for
+#         # the intermediate matrix are dynamic depending
+#         # on the mapping of tokens to experts.
+#         column_indices = ops.topology(
+#             padded_bins,
+#             blocking,
+#             block_rows,
+#             blocks_per_row,
+#         )
+#         data = torch.empty(
+#             column_indices.numel(),
+#             blocking,
+#             blocking,
+#             dtype=torch.float16,
+#             device=x.device,
+#         )
+#         shape = (padded_tokens, fhs * ne)
+#         row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
+#         return stk.Matrix(shape, data, row_indices, column_indices, offsets)
+#
+#     def build_input_matrix(self, sl, hs, ne):
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Assign tokens to experts uniformly.
+#         top_expert = torch.arange(0, sl).cuda().int() % ne
+#
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
+#         return out, padded_bins
+#
+#     def build_weight_matrix(self, ne, hs, fhs):
+#         return torch.randn((hs, ne * fhs)).cuda().half()
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(x, w, topo)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradX::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         topo = topo.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(out, w, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         x = x.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#
+#         w = w.transpose(1, 2).contiguous()
+#         w = w.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd:DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = w.transpose(1, 2).contiguous()
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradX:DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         out = out.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradW:DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = torch.transpose(w, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         x = torch.transpose(x, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/ops/padded_scatter_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -16,50 +16,50 @@ _PADDED_SCATTER_BENCHMARK = (
 )
 
 
-class PaddedScatterTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
-    def testPaddedScatter(self, sl, hs, ne, top_k):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        # Sample weights for the scatter reduce.
-        weights = torch.rand((sl * top_k,)).cuda().half()
-
-        # Gather the data to prepare for backwards.
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-
-        def benchmark():
-            return ops.padded_scatter(
-                x,
-                indices,
-                bin_ids,
-                weights,
-                bins,
-                padded_bins,
-                top_k,
-            )
-
-        time, std = benchmark_util.benchmark_function(benchmark)
-        benchmark_util.log_benchmark(
-            'Padded Scatter',
-            {
-                'sequence_length': sl,
-                'hidden_size': hs,
-                'num_experts': ne,
-                'top_k': top_k,
-            },
-            time,
-            std,
-        )
+# class PaddedScatterTest(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
+#     def testPaddedScatter(self, sl, hs, ne, top_k):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         # Sample weights for the scatter reduce.
+#         weights = torch.rand((sl * top_k,)).cuda().half()
+#
+#         # Gather the data to prepare for backwards.
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+#
+#         def benchmark():
+#             return ops.padded_scatter(
+#                 x,
+#                 indices,
+#                 bin_ids,
+#                 weights,
+#                 bins,
+#                 padded_bins,
+#                 top_k,
+#             )
+#
+#         time, std = benchmark_util.benchmark_function(benchmark)
+#         benchmark_util.log_benchmark(
+#             'Padded Scatter',
+#             {
+#                 'sequence_length': sl,
+#                 'hidden_size': hs,
+#                 'num_experts': ne,
+#                 'top_k': top_k,
+#             },
+#             time,
+#             std,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/ops/permute_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -26,123 +26,123 @@ _PERMUTE_TESTS = (
 )
 
 
-class PermuteBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedGather(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.binned_gather(x, indices, bins, ec)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedScatter(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.binned_gather(x, indices, bins, ec)
-
-        def benchmark():
-            return ops.binned_scatter(x, indices, bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedGather(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedScatter(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        def benchmark():
-            return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testCopy(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        # ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        y = x.clone()
-
-        def benchmark():
-            return y.copy_(x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
+# class PermuteBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedGather(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.binned_gather(x, indices, bins, ec)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedScatter(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.binned_gather(x, indices, bins, ec)
+#
+#         def benchmark():
+#             return ops.binned_scatter(x, indices, bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedGather(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedScatter(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         def benchmark():
+#             return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testCopy(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         # ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         y = x.clone()
+#
+#         def benchmark():
+#             return y.copy_(x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/ops/sort_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -53,32 +53,32 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class SortBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_SORT_TESTS)
-    def testSort(self, n, dtype, max_val):
-        if max_val is None:
-            max_val = np.iinfo(numpy_dtype(dtype)).max
-        end_bit = int(np.ceil(np.log2(max_val)))
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_BASELINE_SORT_TESTS)
-    def testTorchSort(self, n):
-        x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
-        arguments = {
-            'n': n,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class SortBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_SORT_TESTS)
+#     def testSort(self, n, dtype, max_val):
+#         if max_val is None:
+#             max_val = np.iinfo(numpy_dtype(dtype)).max
+#         end_bit = int(np.ceil(np.log2(max_val)))
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_BASELINE_SORT_TESTS)
+#     def testTorchSort(self, n):
+#         x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
+#         arguments = {
+#             'n': n,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/stk/ops/eltwise_ops_test.py

@@ -1,7 +1,7 @@
 import unittest
 import itertools
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 from stk.ops.linear_ops_test import allclose, _dense_and_sparse
@@ -47,40 +47,40 @@ def _dense_and_sparse_like(x, std=0.1):
     return (dense.requires_grad_(True),
             sparse.requires_grad_(True))
 
-@parameterized.parameters(_ELTWISE_OP_TESTS)
-class EltwiseOpsTest(parameterized.TestCase):
-
-    def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
-
-        a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-        b_dense, b = _dense_and_sparse_like(a)
-
-        out = stk.ops.mul(a, b)
-        expected_out = torch.mul(a_dense, b_dense)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size(), out.size())
-        self.assertTrue(allclose(out, expected_out)) 
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = a_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad =  stk.ops.to_dense(b.grad)
-        expected_grad = b_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(_ELTWISE_OP_TESTS)
+# class EltwiseOpsTest(parameterized.TestCase):
+#
+#     def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
+#
+#         a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#         b_dense, b = _dense_and_sparse_like(a)
+#
+#         out = stk.ops.mul(a, b)
+#         expected_out = torch.mul(a_dense, b_dense)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size(), out.size())
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = a_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad =  stk.ops.to_dense(b.grad)
+#         expected_grad = b_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu126-aarch64-linux/stk/ops/linear_ops_test.py

@@ -2,7 +2,7 @@ import unittest
 import itertools
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 
@@ -96,121 +96,121 @@ def _mask(x, mask):
     return x * mask
 
 
-@parameterized.parameters(*_LINEAR_OP_TESTS)
-class LinearOpsTest(parameterized.TestCase):
-
-    def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = _mask(a_dense.grad, a.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = stk.ops.to_dense(b.grad)
-        expected_grad = _mask(b_dense.grad, b.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-        _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
-        expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(*_LINEAR_OP_TESTS)
+# class LinearOpsTest(parameterized.TestCase):
+#
+#     def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = _mask(a_dense.grad, a.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = stk.ops.to_dense(b.grad)
+#         expected_grad = _mask(b_dense.grad, b.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#         _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
+#         expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu126-aarch64-linux/stk/ops/matrix_ops_test.py

@@ -1,61 +1,61 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 import stk
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class MatrixOpsTest(parameterized.TestCase):
-
-    def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
-        mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
-        x = (torch.randn(rows, cols) * mask).type(torch.float16)
-
-        # Convert the matrix to sparse format.
-        sparse_x = stk.ops.to_sparse(x, blocking)
-
-        # Validate the matrix.
-        sparse_x.validate()
-
-        # Validate the shape.
-        self.assertEqual(sparse_x.dim(), 2)
-        self.assertEqual(sparse_x.size()[0], rows)
-        self.assertEqual(sparse_x.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(sparse_x.nnz, nnz)
-
-        # Convert back to dense format.
-        dense_x = stk.ops.to_dense(sparse_x)
-
-        # Validate the shape.
-        self.assertEqual(dense_x.dim(), 2)
-        self.assertEqual(dense_x.size()[0], rows)
-        self.assertEqual(dense_x.size()[1], cols)
-
-        # Validate the sparsity
-        self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
-
-        # Validate the output.
-        self.assertTrue(torch.all(torch.eq(x, dense_x)))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class MatrixOpsTest(parameterized.TestCase):
+#
+#     def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
+#         mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
+#         x = (torch.randn(rows, cols) * mask).type(torch.float16)
+#
+#         # Convert the matrix to sparse format.
+#         sparse_x = stk.ops.to_sparse(x, blocking)
+#
+#         # Validate the matrix.
+#         sparse_x.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(sparse_x.dim(), 2)
+#         self.assertEqual(sparse_x.size()[0], rows)
+#         self.assertEqual(sparse_x.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(sparse_x.nnz, nnz)
+#
+#         # Convert back to dense format.
+#         dense_x = stk.ops.to_dense(sparse_x)
+#
+#         # Validate the shape.
+#         self.assertEqual(dense_x.dim(), 2)
+#         self.assertEqual(dense_x.size()[0], rows)
+#         self.assertEqual(dense_x.size()[1], cols)
+#
+#         # Validate the sparsity
+#         self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
+#
+#         # Validate the output.
+#         self.assertTrue(torch.all(torch.eq(x, dense_x)))
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu126-aarch64-linux/stk/random/random_ops_test.py

@@ -1,72 +1,72 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 from . import random
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class RandomOpsTest(parameterized.TestCase):
-
-    def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
-        mask = random.dense_mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(
-            torch.count_nonzero(mask).item(),
-            nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask, 0),
-                torch.eq(mask, 1))))
-
-    def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
-        mask = random.mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the matrix.
-        mask.validate()
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(mask.nnz, nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask.data, 0),
-                torch.eq(mask.data, 1))))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class RandomOpsTest(parameterized.TestCase):
+#
+#     def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.dense_mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(
+#             torch.count_nonzero(mask).item(),
+#             nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask, 0),
+#                 torch.eq(mask, 1))))
+#
+#     def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the matrix.
+#         mask.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(mask.nnz, nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask.data, 0),
+#                 torch.eq(mask.data, 1))))
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/{_megablocks_cuda_6e04dec.abi3.so → _megablocks_cuda_a45325d.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12705f4547b6a55442c52e081a303d4407202cdc26522f7269c983b627946ab9
+oid sha256:4352757f386e2d948559c737927fa969b2cc9674ee6255487ba8f67fb3470199
 size 21088232

build/torch210-cxx11-cu128-aarch64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_cuda_6e04dec
-ops = torch.ops._megablocks_cuda_6e04dec
+from . import _megablocks_cuda_a45325d
+ops = torch.ops._megablocks_cuda_a45325d
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_cuda_6e04dec::{op_name}"
+    return f"_megablocks_cuda_a45325d::{op_name}"

build/torch210-cxx11-cu128-aarch64-linux/megablocks/__init__.py

@@ -1,10 +1,10 @@
 import ctypes
+import importlib.util
 import sys
-
-import importlib
 from pathlib import Path
 from types import ModuleType
 
+
 def _import_from_path(file_path: Path) -> ModuleType:
     # We cannot use the module name as-is, after adding it to `sys.modules`,
     # it would also be used for other imports. So, we make a module name that

build/torch210-cxx11-cu128-aarch64-linux/ops/histogram_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -47,31 +47,31 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class HistogramBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testTorchHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, 128, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class HistogramBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testTorchHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, 128, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/ops/matmul_benchmark.py

@@ -17,7 +17,7 @@ import unittest
 from .. import stk
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -48,367 +48,367 @@ def log_benchmark(name, arguments, time, std, flops):
     print('=' * 60)
 
 
-class MatmulBenchmark(parameterized.TestCase):
-
-    def build_sparse_matrix(self, x, padded_bins, fhs, ne):
-        blocking = 128
-        padded_tokens, _ = x.size()
-        assert padded_tokens % blocking == 0
-        assert fhs % blocking == 0
-
-        # Offsets for the sparse matrix. All rows have the
-        # same number of nonzero blocks dictated by the
-        # dimensionality of a single expert.
-        block_rows = padded_tokens // blocking
-        blocks_per_row = fhs // blocking
-        offsets = torch.arange(
-            0,
-            block_rows * blocks_per_row + 1,
-            blocks_per_row,
-            dtype=torch.int32,
-            device=x.device,
-        )
-
-        # Indices for the sparse matrix. The indices for
-        # the intermediate matrix are dynamic depending
-        # on the mapping of tokens to experts.
-        column_indices = ops.topology(
-            padded_bins,
-            blocking,
-            block_rows,
-            blocks_per_row,
-        )
-        data = torch.empty(
-            column_indices.numel(),
-            blocking,
-            blocking,
-            dtype=torch.float16,
-            device=x.device,
-        )
-        shape = (padded_tokens, fhs * ne)
-        row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
-        return stk.Matrix(shape, data, row_indices, column_indices, offsets)
-
-    def build_input_matrix(self, sl, hs, ne):
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Assign tokens to experts uniformly.
-        top_expert = torch.arange(0, sl).cuda().int() % ne
-
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
-        return out, padded_bins
-
-    def build_weight_matrix(self, ne, hs, fhs):
-        return torch.randn((hs, ne * fhs)).cuda().half()
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(x, w, topo)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(topo, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradX::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        topo = topo.t()
-
-        def benchmark():
-            return stk.ops.dsd(topo, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(out, w, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        x = x.t()
-
-        def benchmark():
-            return stk.ops.dsd(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-
-        w = w.transpose(1, 2).contiguous()
-        w = w.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd:DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = w.transpose(1, 2).contiguous()
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradX:DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        out = out.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(out, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradW:DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = torch.transpose(w, 1, 2)
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        x = torch.transpose(x, 1, 2)
-
-        def benchmark():
-            return torch.bmm(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
+# class MatmulBenchmark(parameterized.TestCase):
+#
+#     def build_sparse_matrix(self, x, padded_bins, fhs, ne):
+#         blocking = 128
+#         padded_tokens, _ = x.size()
+#         assert padded_tokens % blocking == 0
+#         assert fhs % blocking == 0
+#
+#         # Offsets for the sparse matrix. All rows have the
+#         # same number of nonzero blocks dictated by the
+#         # dimensionality of a single expert.
+#         block_rows = padded_tokens // blocking
+#         blocks_per_row = fhs // blocking
+#         offsets = torch.arange(
+#             0,
+#             block_rows * blocks_per_row + 1,
+#             blocks_per_row,
+#             dtype=torch.int32,
+#             device=x.device,
+#         )
+#
+#         # Indices for the sparse matrix. The indices for
+#         # the intermediate matrix are dynamic depending
+#         # on the mapping of tokens to experts.
+#         column_indices = ops.topology(
+#             padded_bins,
+#             blocking,
+#             block_rows,
+#             blocks_per_row,
+#         )
+#         data = torch.empty(
+#             column_indices.numel(),
+#             blocking,
+#             blocking,
+#             dtype=torch.float16,
+#             device=x.device,
+#         )
+#         shape = (padded_tokens, fhs * ne)
+#         row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
+#         return stk.Matrix(shape, data, row_indices, column_indices, offsets)
+#
+#     def build_input_matrix(self, sl, hs, ne):
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Assign tokens to experts uniformly.
+#         top_expert = torch.arange(0, sl).cuda().int() % ne
+#
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
+#         return out, padded_bins
+#
+#     def build_weight_matrix(self, ne, hs, fhs):
+#         return torch.randn((hs, ne * fhs)).cuda().half()
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(x, w, topo)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradX::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         topo = topo.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(out, w, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         x = x.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#
+#         w = w.transpose(1, 2).contiguous()
+#         w = w.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd:DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = w.transpose(1, 2).contiguous()
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradX:DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         out = out.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradW:DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = torch.transpose(w, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         x = torch.transpose(x, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/ops/padded_scatter_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -16,50 +16,50 @@ _PADDED_SCATTER_BENCHMARK = (
 )
 
 
-class PaddedScatterTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
-    def testPaddedScatter(self, sl, hs, ne, top_k):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        # Sample weights for the scatter reduce.
-        weights = torch.rand((sl * top_k,)).cuda().half()
-
-        # Gather the data to prepare for backwards.
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-
-        def benchmark():
-            return ops.padded_scatter(
-                x,
-                indices,
-                bin_ids,
-                weights,
-                bins,
-                padded_bins,
-                top_k,
-            )
-
-        time, std = benchmark_util.benchmark_function(benchmark)
-        benchmark_util.log_benchmark(
-            'Padded Scatter',
-            {
-                'sequence_length': sl,
-                'hidden_size': hs,
-                'num_experts': ne,
-                'top_k': top_k,
-            },
-            time,
-            std,
-        )
+# class PaddedScatterTest(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
+#     def testPaddedScatter(self, sl, hs, ne, top_k):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         # Sample weights for the scatter reduce.
+#         weights = torch.rand((sl * top_k,)).cuda().half()
+#
+#         # Gather the data to prepare for backwards.
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+#
+#         def benchmark():
+#             return ops.padded_scatter(
+#                 x,
+#                 indices,
+#                 bin_ids,
+#                 weights,
+#                 bins,
+#                 padded_bins,
+#                 top_k,
+#             )
+#
+#         time, std = benchmark_util.benchmark_function(benchmark)
+#         benchmark_util.log_benchmark(
+#             'Padded Scatter',
+#             {
+#                 'sequence_length': sl,
+#                 'hidden_size': hs,
+#                 'num_experts': ne,
+#                 'top_k': top_k,
+#             },
+#             time,
+#             std,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/ops/permute_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -26,123 +26,123 @@ _PERMUTE_TESTS = (
 )
 
 
-class PermuteBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedGather(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.binned_gather(x, indices, bins, ec)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedScatter(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.binned_gather(x, indices, bins, ec)
-
-        def benchmark():
-            return ops.binned_scatter(x, indices, bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedGather(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedScatter(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        def benchmark():
-            return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testCopy(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        # ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        y = x.clone()
-
-        def benchmark():
-            return y.copy_(x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
+# class PermuteBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedGather(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.binned_gather(x, indices, bins, ec)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedScatter(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.binned_gather(x, indices, bins, ec)
+#
+#         def benchmark():
+#             return ops.binned_scatter(x, indices, bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedGather(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedScatter(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         def benchmark():
+#             return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testCopy(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         # ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         y = x.clone()
+#
+#         def benchmark():
+#             return y.copy_(x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/ops/sort_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -53,32 +53,32 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class SortBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_SORT_TESTS)
-    def testSort(self, n, dtype, max_val):
-        if max_val is None:
-            max_val = np.iinfo(numpy_dtype(dtype)).max
-        end_bit = int(np.ceil(np.log2(max_val)))
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_BASELINE_SORT_TESTS)
-    def testTorchSort(self, n):
-        x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
-        arguments = {
-            'n': n,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class SortBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_SORT_TESTS)
+#     def testSort(self, n, dtype, max_val):
+#         if max_val is None:
+#             max_val = np.iinfo(numpy_dtype(dtype)).max
+#         end_bit = int(np.ceil(np.log2(max_val)))
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_BASELINE_SORT_TESTS)
+#     def testTorchSort(self, n):
+#         x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
+#         arguments = {
+#             'n': n,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/stk/ops/eltwise_ops_test.py

@@ -1,7 +1,7 @@
 import unittest
 import itertools
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 from stk.ops.linear_ops_test import allclose, _dense_and_sparse
@@ -47,40 +47,40 @@ def _dense_and_sparse_like(x, std=0.1):
     return (dense.requires_grad_(True),
             sparse.requires_grad_(True))
 
-@parameterized.parameters(_ELTWISE_OP_TESTS)
-class EltwiseOpsTest(parameterized.TestCase):
-
-    def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
-
-        a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-        b_dense, b = _dense_and_sparse_like(a)
-
-        out = stk.ops.mul(a, b)
-        expected_out = torch.mul(a_dense, b_dense)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size(), out.size())
-        self.assertTrue(allclose(out, expected_out)) 
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = a_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad =  stk.ops.to_dense(b.grad)
-        expected_grad = b_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(_ELTWISE_OP_TESTS)
+# class EltwiseOpsTest(parameterized.TestCase):
+#
+#     def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
+#
+#         a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#         b_dense, b = _dense_and_sparse_like(a)
+#
+#         out = stk.ops.mul(a, b)
+#         expected_out = torch.mul(a_dense, b_dense)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size(), out.size())
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = a_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad =  stk.ops.to_dense(b.grad)
+#         expected_grad = b_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu128-aarch64-linux/stk/ops/linear_ops_test.py

@@ -2,7 +2,7 @@ import unittest
 import itertools
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 
@@ -96,121 +96,121 @@ def _mask(x, mask):
     return x * mask
 
 
-@parameterized.parameters(*_LINEAR_OP_TESTS)
-class LinearOpsTest(parameterized.TestCase):
-
-    def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = _mask(a_dense.grad, a.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = stk.ops.to_dense(b.grad)
-        expected_grad = _mask(b_dense.grad, b.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-        _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
-        expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(*_LINEAR_OP_TESTS)
+# class LinearOpsTest(parameterized.TestCase):
+#
+#     def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = _mask(a_dense.grad, a.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = stk.ops.to_dense(b.grad)
+#         expected_grad = _mask(b_dense.grad, b.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#         _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
+#         expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu128-aarch64-linux/stk/ops/matrix_ops_test.py

@@ -1,61 +1,61 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 import stk
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class MatrixOpsTest(parameterized.TestCase):
-
-    def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
-        mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
-        x = (torch.randn(rows, cols) * mask).type(torch.float16)
-
-        # Convert the matrix to sparse format.
-        sparse_x = stk.ops.to_sparse(x, blocking)
-
-        # Validate the matrix.
-        sparse_x.validate()
-
-        # Validate the shape.
-        self.assertEqual(sparse_x.dim(), 2)
-        self.assertEqual(sparse_x.size()[0], rows)
-        self.assertEqual(sparse_x.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(sparse_x.nnz, nnz)
-
-        # Convert back to dense format.
-        dense_x = stk.ops.to_dense(sparse_x)
-
-        # Validate the shape.
-        self.assertEqual(dense_x.dim(), 2)
-        self.assertEqual(dense_x.size()[0], rows)
-        self.assertEqual(dense_x.size()[1], cols)
-
-        # Validate the sparsity
-        self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
-
-        # Validate the output.
-        self.assertTrue(torch.all(torch.eq(x, dense_x)))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class MatrixOpsTest(parameterized.TestCase):
+#
+#     def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
+#         mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
+#         x = (torch.randn(rows, cols) * mask).type(torch.float16)
+#
+#         # Convert the matrix to sparse format.
+#         sparse_x = stk.ops.to_sparse(x, blocking)
+#
+#         # Validate the matrix.
+#         sparse_x.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(sparse_x.dim(), 2)
+#         self.assertEqual(sparse_x.size()[0], rows)
+#         self.assertEqual(sparse_x.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(sparse_x.nnz, nnz)
+#
+#         # Convert back to dense format.
+#         dense_x = stk.ops.to_dense(sparse_x)
+#
+#         # Validate the shape.
+#         self.assertEqual(dense_x.dim(), 2)
+#         self.assertEqual(dense_x.size()[0], rows)
+#         self.assertEqual(dense_x.size()[1], cols)
+#
+#         # Validate the sparsity
+#         self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
+#
+#         # Validate the output.
+#         self.assertTrue(torch.all(torch.eq(x, dense_x)))
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu128-aarch64-linux/stk/random/random_ops_test.py

@@ -1,72 +1,72 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 from . import random
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class RandomOpsTest(parameterized.TestCase):
-
-    def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
-        mask = random.dense_mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(
-            torch.count_nonzero(mask).item(),
-            nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask, 0),
-                torch.eq(mask, 1))))
-
-    def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
-        mask = random.mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the matrix.
-        mask.validate()
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(mask.nnz, nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask.data, 0),
-                torch.eq(mask.data, 1))))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class RandomOpsTest(parameterized.TestCase):
+#
+#     def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.dense_mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(
+#             torch.count_nonzero(mask).item(),
+#             nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask, 0),
+#                 torch.eq(mask, 1))))
+#
+#     def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the matrix.
+#         mask.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(mask.nnz, nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask.data, 0),
+#                 torch.eq(mask.data, 1))))
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/{_megablocks_cuda_6e04dec.abi3.so → _megablocks_cuda_a45325d.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ca7f2de93adbb930ffecaea6953cb94c870333295d05eade3c9c17296aa766a0
+oid sha256:afae994771a9eb3f5ce01fbee1d46fc21b83dbcb3c556cd331cb3fbfde0ff604
 size 12073200

build/torch210-cxx11-cu130-aarch64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_cuda_6e04dec
-ops = torch.ops._megablocks_cuda_6e04dec
+from . import _megablocks_cuda_a45325d
+ops = torch.ops._megablocks_cuda_a45325d
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_cuda_6e04dec::{op_name}"
+    return f"_megablocks_cuda_a45325d::{op_name}"

build/torch210-cxx11-cu130-aarch64-linux/megablocks/__init__.py

@@ -1,10 +1,10 @@
 import ctypes
+import importlib.util
 import sys
-
-import importlib
 from pathlib import Path
 from types import ModuleType
 
+
 def _import_from_path(file_path: Path) -> ModuleType:
     # We cannot use the module name as-is, after adding it to `sys.modules`,
     # it would also be used for other imports. So, we make a module name that

build/torch210-cxx11-cu130-aarch64-linux/ops/histogram_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -47,31 +47,31 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class HistogramBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testTorchHistogram(self, n, dtype, max_val):
-        x = torch.randint(0, 128, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class HistogramBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.histogram(x, max_val),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_HISTOGRAM_TESTS)
+#     def testTorchHistogram(self, n, dtype, max_val):
+#         x = torch.randint(0, 128, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.histc(x, max_val, 0, max_val - 1),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/ops/matmul_benchmark.py

@@ -17,7 +17,7 @@ import unittest
 from .. import stk
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -48,367 +48,367 @@ def log_benchmark(name, arguments, time, std, flops):
     print('=' * 60)
 
 
-class MatmulBenchmark(parameterized.TestCase):
-
-    def build_sparse_matrix(self, x, padded_bins, fhs, ne):
-        blocking = 128
-        padded_tokens, _ = x.size()
-        assert padded_tokens % blocking == 0
-        assert fhs % blocking == 0
-
-        # Offsets for the sparse matrix. All rows have the
-        # same number of nonzero blocks dictated by the
-        # dimensionality of a single expert.
-        block_rows = padded_tokens // blocking
-        blocks_per_row = fhs // blocking
-        offsets = torch.arange(
-            0,
-            block_rows * blocks_per_row + 1,
-            blocks_per_row,
-            dtype=torch.int32,
-            device=x.device,
-        )
-
-        # Indices for the sparse matrix. The indices for
-        # the intermediate matrix are dynamic depending
-        # on the mapping of tokens to experts.
-        column_indices = ops.topology(
-            padded_bins,
-            blocking,
-            block_rows,
-            blocks_per_row,
-        )
-        data = torch.empty(
-            column_indices.numel(),
-            blocking,
-            blocking,
-            dtype=torch.float16,
-            device=x.device,
-        )
-        shape = (padded_tokens, fhs * ne)
-        row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
-        return stk.Matrix(shape, data, row_indices, column_indices, offsets)
-
-    def build_input_matrix(self, sl, hs, ne):
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Assign tokens to experts uniformly.
-        top_expert = torch.arange(0, sl).cuda().int() % ne
-
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
-        return out, padded_bins
-
-    def build_weight_matrix(self, ne, hs, fhs):
-        return torch.randn((hs, ne * fhs)).cuda().half()
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(x, w, topo)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(topo, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradX::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        topo = topo.t()
-
-        def benchmark():
-            return stk.ops.dsd(topo, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-
-        def benchmark():
-            return stk.ops.dsd(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DSD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        w = transpose_view(w)
-
-        def benchmark():
-            return stk.ops.sdd(out, w, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::SDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
-        x, padded_bins = self.build_input_matrix(sl, hs, ne)
-        w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
-        x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
-        out = stk.ops.dsd(x, w)
-        x = x.t()
-
-        def benchmark():
-            return stk.ops.dsd(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DSD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.nnz * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-
-        w = w.transpose(1, 2).contiguous()
-        w = w.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0::Fwd:DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = w.transpose(1, 2).contiguous()
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradX:DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, hs)).cuda().half()
-        w = torch.randn((ne, hs, fhs)).cuda().half()
-        out = torch.bmm(x, w)
-        out = out.transpose(1, 2)
-
-        def benchmark():
-            return torch.bmm(out, x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '0:GradW:DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * fhs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-
-        def benchmark():
-            return torch.bmm(x, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::Fwd::DDD::NN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        w = torch.transpose(w, 1, 2)
-
-        def benchmark():
-            return torch.bmm(out, w)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradX::DDD::NT',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
-
-    @parameterized.parameters(*_MATMUL_TESTS)
-    def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
-        assert (sl % ne) == 0
-        x = torch.randn((ne, sl // ne, fhs)).cuda().half()
-        w = torch.randn((ne, fhs, hs)).cuda().half()
-        out = torch.bmm(x, w)
-        x = torch.transpose(x, 1, 2)
-
-        def benchmark():
-            return torch.bmm(x, out)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'ffn_hidden_size': fhs,
-            'num_experts': ne,
-        }
-        log_benchmark(
-            '1::GradW::DDD::TN',
-            arguments,
-            mean_t,
-            std_t,
-            x.numel() * hs * 2,
-        )
+# class MatmulBenchmark(parameterized.TestCase):
+#
+#     def build_sparse_matrix(self, x, padded_bins, fhs, ne):
+#         blocking = 128
+#         padded_tokens, _ = x.size()
+#         assert padded_tokens % blocking == 0
+#         assert fhs % blocking == 0
+#
+#         # Offsets for the sparse matrix. All rows have the
+#         # same number of nonzero blocks dictated by the
+#         # dimensionality of a single expert.
+#         block_rows = padded_tokens // blocking
+#         blocks_per_row = fhs // blocking
+#         offsets = torch.arange(
+#             0,
+#             block_rows * blocks_per_row + 1,
+#             blocks_per_row,
+#             dtype=torch.int32,
+#             device=x.device,
+#         )
+#
+#         # Indices for the sparse matrix. The indices for
+#         # the intermediate matrix are dynamic depending
+#         # on the mapping of tokens to experts.
+#         column_indices = ops.topology(
+#             padded_bins,
+#             blocking,
+#             block_rows,
+#             blocks_per_row,
+#         )
+#         data = torch.empty(
+#             column_indices.numel(),
+#             blocking,
+#             blocking,
+#             dtype=torch.float16,
+#             device=x.device,
+#         )
+#         shape = (padded_tokens, fhs * ne)
+#         row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
+#         return stk.Matrix(shape, data, row_indices, column_indices, offsets)
+#
+#     def build_input_matrix(self, sl, hs, ne):
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Assign tokens to experts uniformly.
+#         top_expert = torch.arange(0, sl).cuda().int() % ne
+#
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, 1)
+#         return out, padded_bins
+#
+#     def build_weight_matrix(self, ne, hs, fhs):
+#         return torch.randn((hs, ne * fhs)).cuda().half()
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(x, w, topo)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradX::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         topo = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         topo = topo.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(topo, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DSD_NN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DSD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_SDD_NT(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         w = transpose_view(w)
+#
+#         def benchmark():
+#             return stk.ops.sdd(out, w, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::SDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DSD_TN(self, sl, hs, fhs, ne):
+#         x, padded_bins = self.build_input_matrix(sl, hs, ne)
+#         w = self.build_weight_matrix(ne, hs, fhs).t().contiguous()
+#         x = self.build_sparse_matrix(x, padded_bins, fhs, ne)
+#         out = stk.ops.dsd(x, w)
+#         x = x.t()
+#
+#         def benchmark():
+#             return stk.ops.dsd(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DSD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.nnz * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_Fwd_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#
+#         w = w.transpose(1, 2).contiguous()
+#         w = w.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0::Fwd:DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradX_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = w.transpose(1, 2).contiguous()
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradX:DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear0_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, hs)).cuda().half()
+#         w = torch.randn((ne, hs, fhs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         out = out.transpose(1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '0:GradW:DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * fhs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_Fwd_DDD_NN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#
+#         def benchmark():
+#             return torch.bmm(x, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::Fwd::DDD::NN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradX_DDD_NT(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         w = torch.transpose(w, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(out, w)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradX::DDD::NT',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
+#
+#     @parameterized.parameters(*_MATMUL_TESTS)
+#     def testFFN_Linear1_GradW_DDD_TN(self, sl, hs, fhs, ne):
+#         assert (sl % ne) == 0
+#         x = torch.randn((ne, sl // ne, fhs)).cuda().half()
+#         w = torch.randn((ne, fhs, hs)).cuda().half()
+#         out = torch.bmm(x, w)
+#         x = torch.transpose(x, 1, 2)
+#
+#         def benchmark():
+#             return torch.bmm(x, out)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'ffn_hidden_size': fhs,
+#             'num_experts': ne,
+#         }
+#         log_benchmark(
+#             '1::GradW::DDD::TN',
+#             arguments,
+#             mean_t,
+#             std_t,
+#             x.numel() * hs * 2,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/ops/padded_scatter_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -16,50 +16,50 @@ _PADDED_SCATTER_BENCHMARK = (
 )
 
 
-class PaddedScatterTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
-    def testPaddedScatter(self, sl, hs, ne, top_k):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        # Sample weights for the scatter reduce.
-        weights = torch.rand((sl * top_k,)).cuda().half()
-
-        # Gather the data to prepare for backwards.
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-
-        def benchmark():
-            return ops.padded_scatter(
-                x,
-                indices,
-                bin_ids,
-                weights,
-                bins,
-                padded_bins,
-                top_k,
-            )
-
-        time, std = benchmark_util.benchmark_function(benchmark)
-        benchmark_util.log_benchmark(
-            'Padded Scatter',
-            {
-                'sequence_length': sl,
-                'hidden_size': hs,
-                'num_experts': ne,
-                'top_k': top_k,
-            },
-            time,
-            std,
-        )
+# class PaddedScatterTest(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PADDED_SCATTER_BENCHMARK)
+#     def testPaddedScatter(self, sl, hs, ne, top_k):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         # Sample weights for the scatter reduce.
+#         weights = torch.rand((sl * top_k,)).cuda().half()
+#
+#         # Gather the data to prepare for backwards.
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+#
+#         def benchmark():
+#             return ops.padded_scatter(
+#                 x,
+#                 indices,
+#                 bin_ids,
+#                 weights,
+#                 bins,
+#                 padded_bins,
+#                 top_k,
+#             )
+#
+#         time, std = benchmark_util.benchmark_function(benchmark)
+#         benchmark_util.log_benchmark(
+#             'Padded Scatter',
+#             {
+#                 'sequence_length': sl,
+#                 'hidden_size': hs,
+#                 'num_experts': ne,
+#                 'top_k': top_k,
+#             },
+#             time,
+#             std,
+#         )
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/ops/permute_benchmark.py

@@ -4,7 +4,7 @@
 import unittest
 
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import benchmark_util, ops
 
@@ -26,123 +26,123 @@ _PERMUTE_TESTS = (
 )
 
 
-class PermuteBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedGather(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.binned_gather(x, indices, bins, ec)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testBinnedScatter(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(indices, ne)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.binned_gather(x, indices, bins, ec)
-
-        def benchmark():
-            return ops.binned_scatter(x, indices, bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedGather(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def benchmark():
-            return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testPaddedScatter(self, sl, hs, ne):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
-
-        def benchmark():
-            return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_PERMUTE_TESTS)
-    def testCopy(self, sl, hs, ne):
-        # NOTE: Capacity factor == 1.
-        # ec = sl // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-        y = x.clone()
-
-        def benchmark():
-            return y.copy_(x)
-
-        mean_t, std_t = benchmark_util.benchmark_function(benchmark)
-        arguments = {
-            'sequence_length': sl,
-            'hidden_size': hs,
-            'num_experts': ne,
-        }
-        benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
+# class PermuteBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedGather(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.binned_gather(x, indices, bins, ec)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testBinnedScatter(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(indices, ne)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.binned_gather(x, indices, bins, ec)
+#
+#         def benchmark():
+#             return ops.binned_scatter(x, indices, bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('BinnedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedGather(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#
+#         def benchmark():
+#             return ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedGather', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testPaddedScatter(self, sl, hs, ne):
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#
+#         # Randomly assign tokens to experts.
+#         top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+#         bin_ids, indices = ops.sort(top_expert)
+#         tokens_per_expert = ops.histogram(top_expert, ne)
+#         padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+#         padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+#         bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+#         x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins)
+#
+#         def benchmark():
+#             return ops.padded_scatter(x, indices, bin_ids, bins, padded_bins)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('PaddedScatter', arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_PERMUTE_TESTS)
+#     def testCopy(self, sl, hs, ne):
+#         # NOTE: Capacity factor == 1.
+#         # ec = sl // ne
+#
+#         # Create the data and indices.
+#         x = torch.randn((sl, hs)).cuda().half()
+#         y = x.clone()
+#
+#         def benchmark():
+#             return y.copy_(x)
+#
+#         mean_t, std_t = benchmark_util.benchmark_function(benchmark)
+#         arguments = {
+#             'sequence_length': sl,
+#             'hidden_size': hs,
+#             'num_experts': ne,
+#         }
+#         benchmark_util.log_benchmark('Copy', arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/ops/sort_benchmark.py

@@ -5,7 +5,7 @@ import unittest
 
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 from .. import ops
 
@@ -53,32 +53,32 @@ def log_benchmark(arguments, mean_t, std_t):
     print('=' * 60)
 
 
-class SortBenchmark(parameterized.TestCase):
-
-    @parameterized.parameters(*_SORT_TESTS)
-    def testSort(self, n, dtype, max_val):
-        if max_val is None:
-            max_val = np.iinfo(numpy_dtype(dtype)).max
-        end_bit = int(np.ceil(np.log2(max_val)))
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
-        arguments = {
-            'n': n,
-            'dtype': dtype,
-            'max_val': max_val,
-        }
-        log_benchmark(arguments, mean_t, std_t)
-
-    @parameterized.parameters(*_BASELINE_SORT_TESTS)
-    def testTorchSort(self, n):
-        x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
-
-        mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
-        arguments = {
-            'n': n,
-        }
-        log_benchmark(arguments, mean_t, std_t)
+# class SortBenchmark(parameterized.TestCase):
+#
+#     @parameterized.parameters(*_SORT_TESTS)
+#     def testSort(self, n, dtype, max_val):
+#         if max_val is None:
+#             max_val = np.iinfo(numpy_dtype(dtype)).max
+#         end_bit = int(np.ceil(np.log2(max_val)))
+#         x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: ops.sort(x, end_bit),)
+#         arguments = {
+#             'n': n,
+#             'dtype': dtype,
+#             'max_val': max_val,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
+#
+#     @parameterized.parameters(*_BASELINE_SORT_TESTS)
+#     def testTorchSort(self, n):
+#         x = torch.randint(0, 128, (n,)).cuda().to(torch.int32)
+#
+#         mean_t, std_t, max_t, min_t = benchmark_function(lambda: torch.sort(x))
+#         arguments = {
+#             'n': n,
+#         }
+#         log_benchmark(arguments, mean_t, std_t)
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/stk/ops/eltwise_ops_test.py

@@ -1,7 +1,7 @@
 import unittest
 import itertools
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 from stk.ops.linear_ops_test import allclose, _dense_and_sparse
@@ -47,40 +47,40 @@ def _dense_and_sparse_like(x, std=0.1):
     return (dense.requires_grad_(True),
             sparse.requires_grad_(True))
 
-@parameterized.parameters(_ELTWISE_OP_TESTS)
-class EltwiseOpsTest(parameterized.TestCase):
-
-    def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
-
-        a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-        b_dense, b = _dense_and_sparse_like(a)
-
-        out = stk.ops.mul(a, b)
-        expected_out = torch.mul(a_dense, b_dense)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size(), out.size())
-        self.assertTrue(allclose(out, expected_out)) 
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = a_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad =  stk.ops.to_dense(b.grad)
-        expected_grad = b_dense.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size(), grad.size())
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(_ELTWISE_OP_TESTS)
+# class EltwiseOpsTest(parameterized.TestCase):
+#
+#     def testEltwiseMul(self, m, n, sparsity, blocking, dtype):
+#
+#         a_dense, a = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#         b_dense, b = _dense_and_sparse_like(a)
+#
+#         out = stk.ops.mul(a, b)
+#         expected_out = torch.mul(a_dense, b_dense)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size(), out.size())
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = a_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad =  stk.ops.to_dense(b.grad)
+#         expected_grad = b_dense.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size(), grad.size())
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu130-aarch64-linux/stk/ops/linear_ops_test.py

@@ -2,7 +2,7 @@ import unittest
 import itertools
 import numpy as np
 import torch
-from absl.testing import parameterized
+# from absl.testing import parameterized
 
 import stk
 
@@ -96,121 +96,121 @@ def _mask(x, mask):
     return x * mask
 
 
-@parameterized.parameters(*_LINEAR_OP_TESTS)
-class LinearOpsTest(parameterized.TestCase):
-
-    def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = stk.ops.to_dense(a.grad)
-        expected_grad = _mask(a_dense.grad, a.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
-        expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        out.sum().backward()
-
-        # Validate the results.
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = stk.ops.to_dense(b.grad)
-        expected_grad = _mask(b_dense.grad, b.grad)
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-    def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
-        # Construct the operands.
-        a_shape = (k, m) if trans_a else (m, k)
-        a, acp = _dense_2x(*a_shape, dtype)
-        b_shape = (n, k) if trans_b else (k, n)
-        b, bcp = _dense_2x(*b_shape, dtype)
-        _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
-
-        # Execute the matmul.
-        out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
-        expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
-
-        # Compute the gradients w.r.t. the inputs.
-        expected_out.sum().backward()
-        stk.ops.sum(out).backward()
-
-        # Validate the results.
-        out = stk.ops.to_dense(out)
-        self.assertEqual(out.dim(), 2)
-        self.assertEqual(expected_out.size()[0], out.size()[0])
-        self.assertEqual(expected_out.size()[1], out.size()[1])
-        self.assertTrue(allclose(out, expected_out))
-
-        # LHS gradient.
-        grad = a.grad
-        expected_grad = acp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
-
-        # RHS gradient.
-        grad = b.grad
-        expected_grad = bcp.grad
-        self.assertEqual(grad.dim(), 2)
-        self.assertEqual(expected_grad.size()[0], grad.size()[0])
-        self.assertEqual(expected_grad.size()[1], grad.size()[1])
-        self.assertTrue(allclose(grad, expected_grad))
+# @parameterized.parameters(*_LINEAR_OP_TESTS)
+# class LinearOpsTest(parameterized.TestCase):
+#
+#     def testLinearOps_Dsd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a_dense, a = _dense_and_sparse(*a_shape, sparsity, blocking, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dsd, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, a_dense, bcp, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = stk.ops.to_dense(a.grad)
+#         expected_grad = _mask(a_dense.grad, a.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Dds(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b_dense, b = _dense_and_sparse(*b_shape, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _with_transpose(stk.ops.dds, a, b, trans_a, trans_b)
+#         expected_out = _with_transpose(torch.mm, acp, b_dense, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         out.sum().backward()
+#
+#         # Validate the results.
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = stk.ops.to_dense(b.grad)
+#         expected_grad = _mask(b_dense.grad, b.grad)
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#     def testLinearOps_Sdd(self, m, k, n, sparsity, trans_a, trans_b, blocking, dtype):
+#         # Construct the operands.
+#         a_shape = (k, m) if trans_a else (m, k)
+#         a, acp = _dense_2x(*a_shape, dtype)
+#         b_shape = (n, k) if trans_b else (k, n)
+#         b, bcp = _dense_2x(*b_shape, dtype)
+#         _, topo = _dense_and_sparse(m, n, sparsity, blocking, dtype)
+#
+#         # Execute the matmul.
+#         out = _sparse_out_with_transpose(stk.ops.sdd, a, b, topo, trans_a, trans_b)
+#         expected_out = _sparse_out_with_transpose(_mmm, acp, bcp, topo, trans_a, trans_b)
+#
+#         # Compute the gradients w.r.t. the inputs.
+#         expected_out.sum().backward()
+#         stk.ops.sum(out).backward()
+#
+#         # Validate the results.
+#         out = stk.ops.to_dense(out)
+#         self.assertEqual(out.dim(), 2)
+#         self.assertEqual(expected_out.size()[0], out.size()[0])
+#         self.assertEqual(expected_out.size()[1], out.size()[1])
+#         self.assertTrue(allclose(out, expected_out))
+#
+#         # LHS gradient.
+#         grad = a.grad
+#         expected_grad = acp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
+#
+#         # RHS gradient.
+#         grad = b.grad
+#         expected_grad = bcp.grad
+#         self.assertEqual(grad.dim(), 2)
+#         self.assertEqual(expected_grad.size()[0], grad.size()[0])
+#         self.assertEqual(expected_grad.size()[1], grad.size()[1])
+#         self.assertTrue(allclose(grad, expected_grad))
 
 if __name__ == '__main__':
     unittest.main()

build/torch210-cxx11-cu130-aarch64-linux/stk/ops/matrix_ops_test.py

@@ -1,61 +1,61 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 import stk
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class MatrixOpsTest(parameterized.TestCase):
-
-    def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
-        mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
-        x = (torch.randn(rows, cols) * mask).type(torch.float16)
-
-        # Convert the matrix to sparse format.
-        sparse_x = stk.ops.to_sparse(x, blocking)
-
-        # Validate the matrix.
-        sparse_x.validate()
-
-        # Validate the shape.
-        self.assertEqual(sparse_x.dim(), 2)
-        self.assertEqual(sparse_x.size()[0], rows)
-        self.assertEqual(sparse_x.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(sparse_x.nnz, nnz)
-
-        # Convert back to dense format.
-        dense_x = stk.ops.to_dense(sparse_x)
-
-        # Validate the shape.
-        self.assertEqual(dense_x.dim(), 2)
-        self.assertEqual(dense_x.size()[0], rows)
-        self.assertEqual(dense_x.size()[1], cols)
-
-        # Validate the sparsity
-        self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
-
-        # Validate the output.
-        self.assertTrue(torch.all(torch.eq(x, dense_x)))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class MatrixOpsTest(parameterized.TestCase):
+#
+#     def testMatrixOps_FormatConversion(self, rows, cols, sparsity, blocking):
+#         mask = stk.random.dense_mask(rows, cols, sparsity, blocking)
+#         x = (torch.randn(rows, cols) * mask).type(torch.float16)
+#
+#         # Convert the matrix to sparse format.
+#         sparse_x = stk.ops.to_sparse(x, blocking)
+#
+#         # Validate the matrix.
+#         sparse_x.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(sparse_x.dim(), 2)
+#         self.assertEqual(sparse_x.size()[0], rows)
+#         self.assertEqual(sparse_x.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(sparse_x.nnz, nnz)
+#
+#         # Convert back to dense format.
+#         dense_x = stk.ops.to_dense(sparse_x)
+#
+#         # Validate the shape.
+#         self.assertEqual(dense_x.dim(), 2)
+#         self.assertEqual(dense_x.size()[0], rows)
+#         self.assertEqual(dense_x.size()[1], cols)
+#
+#         # Validate the sparsity
+#         self.assertEqual(torch.count_nonzero(dense_x).item(), nnz)
+#
+#         # Validate the output.
+#         self.assertTrue(torch.all(torch.eq(x, dense_x)))
 
 
 if __name__ == '__main__':

build/torch210-cxx11-cu130-aarch64-linux/stk/random/random_ops_test.py

@@ -1,72 +1,72 @@
 import unittest
 
-from absl.testing import parameterized
+# from absl.testing import parameterized
 from . import random
 import torch
 
 
-@parameterized.parameters(
-    (8, 16, 0.0, 1),
-    (8, 16, 0.5, 1),
-    (8, 16, .95, 1),
-    (16, 8, 0.0, 1),
-    (16, 8, 0.5, 1),
-    (16, 8, .95, 1),
-    (8, 16, 0.0, 8),
-    (8, 16, 0.5, 8),
-    (8, 16, 1.0, 8),
-    (16, 8, 0.0, 8),
-    (16, 8, 0.5, 8),
-    (16, 8, 1.0, 8),
-    (128, 256, 0.5, 16),
-    (256, 128, 0.75, 32),
-    (512, 512, .875, 128))
-class RandomOpsTest(parameterized.TestCase):
-
-    def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
-        mask = random.dense_mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(
-            torch.count_nonzero(mask).item(),
-            nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask, 0),
-                torch.eq(mask, 1))))
-
-    def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
-        mask = random.mask(
-            rows, cols, sparsity, blocking)
-
-        # Validate the matrix.
-        mask.validate()
-
-        # Validate the shape.
-        self.assertEqual(mask.dim(), 2)
-        self.assertEqual(mask.size()[0], rows)
-        self.assertEqual(mask.size()[1], cols)
-
-        # Validate the sparsity.
-        numblocks = rows // blocking * cols // blocking
-        nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
-        self.assertEqual(mask.nnz, nnz)
-
-        # Check values are zero or one.
-        self.assertTrue(
-            torch.all(torch.logical_or(
-                torch.eq(mask.data, 0),
-                torch.eq(mask.data, 1))))
+# @parameterized.parameters(
+#     (8, 16, 0.0, 1),
+#     (8, 16, 0.5, 1),
+#     (8, 16, .95, 1),
+#     (16, 8, 0.0, 1),
+#     (16, 8, 0.5, 1),
+#     (16, 8, .95, 1),
+#     (8, 16, 0.0, 8),
+#     (8, 16, 0.5, 8),
+#     (8, 16, 1.0, 8),
+#     (16, 8, 0.0, 8),
+#     (16, 8, 0.5, 8),
+#     (16, 8, 1.0, 8),
+#     (128, 256, 0.5, 16),
+#     (256, 128, 0.75, 32),
+#     (512, 512, .875, 128))
+# class RandomOpsTest(parameterized.TestCase):
+#
+#     def testRandomOps_DenseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.dense_mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(
+#             torch.count_nonzero(mask).item(),
+#             nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask, 0),
+#                 torch.eq(mask, 1))))
+#
+#     def testRandomOps_SparseMask(self, rows, cols, sparsity, blocking):
+#         mask = random.mask(
+#             rows, cols, sparsity, blocking)
+#
+#         # Validate the matrix.
+#         mask.validate()
+#
+#         # Validate the shape.
+#         self.assertEqual(mask.dim(), 2)
+#         self.assertEqual(mask.size()[0], rows)
+#         self.assertEqual(mask.size()[1], cols)
+#
+#         # Validate the sparsity.
+#         numblocks = rows // blocking * cols // blocking
+#         nnz = round(numblocks * (1 - sparsity)) * blocking ** 2
+#         self.assertEqual(mask.nnz, nnz)
+#
+#         # Check values are zero or one.
+#         self.assertTrue(
+#             torch.all(torch.logical_or(
+#                 torch.eq(mask.data, 0),
+#                 torch.eq(mask.data, 1))))
 
 
 if __name__ == '__main__':

build/torch211-cxx11-cu126-aarch64-linux/__init__.py

@@ -0,0 +1,202 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from ._ops import ops
+
+from .grouped_gemm import backend as gg_backend
+from .grouped_gemm import ops as gg_ops
+
+
+from ._layers.arguments import Arguments
+from ._layers.dmoe import ParallelDroplessMLP, dMoE
+from ._layers.glu import SparseGLU
+from ._layers.mlp import MLP, SparseMLP
+from ._layers.moe import MoE, ParallelMLP, get_load_balancing_loss
+
+from . import layers
+
+# This section contains the direct kernel exports (not inlcuded in the original code)
+def exclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:
+    """
+    Compute exclusive cumulative sum along the specified dimension.
+
+    Args:
+        x: Input tensor
+        dim: Dimension along which to compute cumsum
+        out: Output tensor (modified in-place)
+
+    Returns:
+        The output tensor
+    """
+    result = ops.exclusive_cumsum(x, dim)
+    out.copy_(result)
+    return out
+
+
+def inclusive_cumsum(x: torch.Tensor, dim: int, out: torch.Tensor) -> torch.Tensor:
+    """
+    Compute inclusive cumulative sum along the specified dimension.
+
+    Args:
+        x: Input tensor
+        dim: Dimension along which to compute cumsum
+        out: Output tensor (modified in-place)
+
+    Returns:
+        The output tensor
+    """
+    result = ops.inclusive_cumsum(x, dim)
+    out.copy_(result)
+    return out
+
+
+def histogram(x: torch.Tensor, num_bins: int) -> torch.Tensor:
+    """
+    Compute histogram of input tensor values.
+
+    Args:
+        x: Input tensor
+        num_bins: Number of histogram bins
+
+    Returns:
+        Histogram tensor with counts for each bin
+    """
+    return ops.histogram(x, num_bins)
+
+
+def indices(
+    padded_bins: torch.Tensor,
+    block_size: int,
+    output_block_rows: int,
+    output_block_columns: int,
+) -> torch.Tensor:
+    """
+    Construct indices from padded bins for sparse operations.
+
+    Args:
+        padded_bins: Tensor containing bin boundaries
+        block_size: Size of each block
+        output_block_rows: Number of rows in output blocks
+        output_block_columns: Number of columns in output blocks
+
+    Returns:
+        Tensor containing constructed indices
+    """
+    return ops.indices(padded_bins, block_size, output_block_rows, output_block_columns)
+
+
+def replicate_forward(
+    x: torch.Tensor, bins: torch.Tensor, out: torch.Tensor
+) -> torch.Tensor:
+    """
+    Forward pass of replicate operation - replicate values according to bin sizes.
+
+    Args:
+        x: Input tensor with values to replicate
+        bins: Tensor containing bin sizes
+        out: Output tensor (modified in-place)
+
+    Returns:
+        The output tensor
+    """
+    return ops.replicate_forward(x, bins, out)
+
+
+def replicate_backward(
+    grad: torch.Tensor, bins: torch.Tensor, out: torch.Tensor
+) -> torch.Tensor:
+    """
+    Backward pass of replicate operation - reduce gradients back to bins.
+
+    Args:
+        grad: Gradient tensor to reduce
+        bins: Tensor containing bin sizes
+        out: Output tensor (modified in-place)
+
+    Returns:
+        The output tensor
+    """
+    return ops.replicate_backward(grad, bins, out)
+
+
+def sort(
+    x: torch.Tensor, end_bit: int, x_out: torch.Tensor, iota_out: torch.Tensor
+) -> torch.Tensor:
+    """
+    Radix sort with index tracking.
+
+    Args:
+        x: Input tensor to sort
+        end_bit: Number of bits to consider in sorting
+        x_out: Output tensor for sorted values
+        iota_out: Output tensor for sorted indices
+
+    Returns:
+        The sorted values tensor
+    """
+    return ops.sort(x, end_bit, x_out, iota_out)
+
+
+# Convenience functions for common use cases
+def cumsum(x: torch.Tensor, dim: int = -1, exclusive: bool = False) -> torch.Tensor:
+    """
+    Compute cumulative sum with automatic output allocation.
+
+    Args:
+        x: Input tensor
+        dim: Dimension along which to compute cumsum (default: last dimension)
+        exclusive: Whether to compute exclusive (True) or inclusive (False) cumsum
+
+    Returns:
+        New tensor containing the cumulative sum
+    """
+    out = torch.empty_like(x)
+    if exclusive:
+        return exclusive_cumsum(x, dim, out)
+    else:
+        return inclusive_cumsum(x, dim, out)
+
+
+def argsort(x: torch.Tensor, end_bit: int = 32) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Sort tensor and return both sorted values and indices.
+
+    Args:
+        x: Input tensor to sort
+        end_bit: Number of bits to consider in sorting
+
+    Returns:
+        Tuple of (sorted_values, sorted_indices)
+    """
+    x_out = torch.empty_like(x)
+    iota_out = torch.empty_like(x)
+    sort(x, end_bit, x_out, iota_out)
+    return x_out, iota_out
+
+
+# Export public API
+__all__ = [
+    "MyReplacementLayer",
+    # Direct kernel exports
+    "exclusive_cumsum",
+    "inclusive_cumsum",
+    "histogram",
+    "indices",
+    "replicate_forward",
+    "replicate_backward",
+    "sort",
+    "cumsum",
+    "argsort",
+    # Original exports
+    "Arguments",
+    "ParallelDroplessMLP",
+    "dMoE",
+    "SparseGLU",
+    "MLP",
+    "SparseMLP",
+    "MoE",
+    "ParallelMLP",
+    "get_load_balancing_loss",
+]

build/torch211-cxx11-cu126-aarch64-linux/_layers/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+# from megablocks.layers.dmoe import dMoE
+from .moe import MoE
+
+__all__ = [
+    'MoE',
+    # 'dMoE',
+]

build/torch211-cxx11-cu126-aarch64-linux/_layers/activation_fn.py

@@ -0,0 +1,33 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Callable, Union
+
+import torch
+from ..stk import Matrix
+
+
+def act_fn(
+    x: Matrix,
+    function: Callable,
+    return_grad_fn: bool = False,
+    **kwargs,
+) -> Union[tuple[Matrix, Any] | Matrix]:
+    assert isinstance(x, Matrix)
+    with torch.set_grad_enabled(torch.is_grad_enabled() or return_grad_fn):
+        if return_grad_fn:
+            x.data.requires_grad = True
+        out = function(x.data, **kwargs)
+        y = Matrix(
+            x.size(),
+            out,
+            x.row_indices,
+            x.column_indices,
+            x.offsets,
+            x.column_indices_t,
+            x.offsets_t,
+            x.block_offsets_t,
+        )
+        if return_grad_fn:
+            return y, out.backward
+        return y

build/torch211-cxx11-cu126-aarch64-linux/_layers/all_to_all.py

@@ -0,0 +1,54 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty((sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype)
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )

build/torch211-cxx11-cu126-aarch64-linux/_layers/arguments.py

@@ -0,0 +1,101 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from functools import partial
+from typing import Any, Callable, Optional, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+# import megablocks.grouped_gemm_util as grouped_gemm
+from .. import grouped_gemm_util as grouped_gemm
+
+# Type annotation for in-place Tensor initialization function.
+InitFn = Union[Callable[[torch.Tensor], None], partial[torch.Tensor]]
+
+_ALLOWED_BITWIDTHS = (-1, 4, 8)
+
+DEFAULT_ACTIVATION_FN = partial(F.gelu, approximate='tanh')
+
+
+@dataclasses.dataclass
+class Arguments:
+    # Model arguments.
+    hidden_size: int = 1024
+    ffn_hidden_size: int = 4096
+    num_layers: int = 1
+    bias: bool = True
+    return_bias: bool = True
+    activation_fn: Optional[Callable] = DEFAULT_ACTIVATION_FN
+
+    # MoE arguments.
+    moe_num_experts: int = 1
+    moe_top_k: int = 1
+    moe_capacity_factor: int = 1
+    moe_normalize_expert_weights: Optional[Union[int, float]] = None
+    moe_loss_weight: float = 0.1
+    moe_jitter_eps: Optional[float] = None
+    moe_lbl_in_fp32: bool = False
+
+    # Parallelism arguments.
+    moe_expert_model_parallelism: bool = False
+    expert_parallel_group: Optional[dist.ProcessGroup] = None
+    pipeline_model_parallel_size: int = 1
+    num_layers_per_virtual_pipeline_stage: Optional[int] = None
+
+    # Compute arguments.
+    memory_optimized_mlp: bool = False
+    mlp_type: str = 'mlp'
+    mlp_impl: str = 'sparse'
+
+    # Initialization arguments.
+    fp16: bool = True
+    bf16: bool = False
+    device: Union[int, torch.device] = dataclasses.field(default_factory=torch.cuda.current_device)
+    init_method: InitFn = partial(torch.nn.init.normal_, mean=0.0, std=0.02)
+    output_layer_init_method: InitFn = init_method
+
+    # Benchmarking arguments.
+    uniform_expert_assignment: bool = False
+
+    # shared expert arguments
+    shared_expert: bool = False  # enable using shared expert
+    fc_cls: Any = torch.nn.Linear  # class of the fully connected layer in shared expert (purpose: to allow using custom FC layer eg te.Linear (for FP8))
+    fc_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict,)  # kwargs for custom fc layers
+    remat_act_fn: bool = True  # enable act fn to be rematerialized instead of stored
+    shared_expert_hidden_size: Optional[
+        int] = None  # hidden size of the shared expert IF we want to set it to something different from hidden_size
+    shared_expert_weighted_sum: bool = False  # enable using weighted sum for shared expert output (wieghted by number of experts used)
+
+    # Router Z-loss arguments
+    moe_zloss_weight: float = 0  # 1e-3 is a reasonable value
+    moe_zloss_in_fp32: bool = False
+
+    def __post_init__(self):
+        # Sparse MLP is not supported with triton >=3.2.0
+        # TODO: Remove this once sparse is supported with triton >=3.2.0
+        if self.__getattribute__('mlp_impl') == 'sparse':
+            try:
+                import triton
+                if triton.__version__ >= '3.2.0':
+                    raise ValueError(
+                        'Sparse MLP is not supported with triton >=3.2.0. Please use mlp_impl="grouped" instead.',
+                    )
+            except ImportError:
+                raise ImportError('Triton is required for sparse MLP implementation')
+
+        if self.__getattribute__('mlp_impl') == 'grouped':
+            grouped_gemm.assert_grouped_gemm_is_available()
+
+        if self.shared_expert_hidden_size is None:
+            self.shared_expert_hidden_size = self.ffn_hidden_size
+
+
+def from_megatron(megatron_args: Any):
+    args = Arguments()
+    for field in dataclasses.fields(args):
+        if hasattr(megatron_args, field.name):
+            setattr(args, field.name, getattr(megatron_args, field.name))
+    return args

build/torch211-cxx11-cu126-aarch64-linux/_layers/common.py

@@ -0,0 +1,26 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from .arguments import Arguments
+
+
+def dtype(args: Arguments):
+    if args.fp16:
+        return torch.float16
+    elif args.bf16:
+        return torch.bfloat16
+    return None
+
+
+def cast_if_autocast_enabled(tensor):
+    if torch.is_autocast_enabled():
+        if tensor.device.type == 'cuda':
+            dtype = torch.get_autocast_gpu_dtype()
+        elif tensor.device.type == 'cpu':
+            dtype = torch.get_autocast_cpu_dtype()
+        else:
+            raise NotImplementedError()
+        return tensor.to(dtype=dtype)
+    return tensor

build/torch211-cxx11-cu126-aarch64-linux/_layers/dmlp_registry.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Union
+
+from . import glu, mlp
+from .arguments import Arguments
+
+MlpType = Union[mlp.SparseMLP, glu.SparseGLU]
+
+_REGISTRY = {
+    'mlp': {
+        'grouped': mlp.GroupedMLP,
+        'sparse': mlp.SparseMLP,
+    },
+    'glu': {
+        'grouped': glu.GroupedGLU,
+        'sparse': glu.SparseGLU,
+    },
+}
+
+
+def get(args: Arguments) -> MlpType:
+    """Returns an MLP for use in a dMoE instance.
+
+    Uses the provided arguments to instantiate the appropriate
+    MLP instance. This only contains MLPs for use in dMoEs
+    (ie. only for the dropless versions of MoEs).
+
+    Args:
+        args: propagated Arguments dataclass.
+
+    Returns:
+        An instantiated MLP constructed using the input args.
+    """
+    if args.mlp_type not in _REGISTRY:
+        raise ValueError(f'Unsupported mlp type: {args.mlp_type}')
+
+    if args.mlp_impl not in _REGISTRY[args.mlp_type]:
+        raise ValueError(f'{args.mlp_type} does not support {args.mlp_impl} backend.',)
+
+    return _REGISTRY[args.mlp_type][args.mlp_impl](args)

build/torch211-cxx11-cu126-aarch64-linux/_layers/dmoe.py

@@ -0,0 +1,337 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+# try:
+#     import stk.ops
+# except ImportError:
+#     import warnings
+#     warnings.warn(
+#         'Please add `stanford-stk` if megablocks/_layers/dmoe.py is needed.',
+#     )
+
+# import megablocks.ops as ops
+# # from megablocks.ops import ops
+# from megablocks.layers import common, dmlp_registry, moe, mpu
+# from megablocks.layers.arguments import Arguments
+
+from .. import stk
+from .. import ops
+from . import common, dmlp_registry, moe, mpu
+from .arguments import Arguments
+
+def promote_scalar(x):
+    return x.view(1) if not len(x.size()) else x
+
+
+class ParallelDroplessMLP(moe.ParallelMLP):
+
+    def __init__(self, args: Arguments):
+        super(ParallelDroplessMLP, self).__init__(args)
+        self.hidden_size = args.hidden_size
+        self.ffn_hidden_size = mpu.features_per_rank(args)
+        self.blocking = 128
+        self.mlp = dmlp_registry.get(args)
+
+        # Calculate the number of bits needed to represent the column indices
+        # in the intermediate sparse matrix.
+        max_column_index = ((self.ffn_hidden_size * self.num_experts) // self.blocking)
+        self.transpose_sort_end_bit = max(
+            int(np.ceil(np.log2(max_column_index))),
+            1,
+        )
+
+    def sparse_transpose(self, size, row_indices, column_indices, offsets):
+        block_columns = size[1] // self.blocking
+
+        # Sort row indices by column indices to get the transposed matrix's
+        # column indices.
+        #
+        # NOTE: Our sort operation uses the same width indices as the input values.
+        # To avoid overflow when we have large activation matrices we cast to
+        # 32-bit before sorting.
+        _, gather_indices = ops.sort(
+            column_indices.int(),
+            self.transpose_sort_end_bit,
+        )
+
+        # There are a constant number of blocks in every row of the sparse matrix.
+        # A blocks offset is:
+        #
+        # row_index * blocks_per_row + column_index % blocks_per_row
+        #
+        # Once we have the block offsets ordered for transposition we can divide
+        # by blocks_per_row to get the transposed column indices.
+        column_indices_t = row_indices.gather(0, gather_indices.long())
+        block_offsets_t = gather_indices.int()
+
+        zero = torch.zeros((1,), dtype=torch.int32, device=row_indices.device)
+        nnz_per_column = ops.histogram(column_indices, block_columns)
+        nnz_per_column = ops.inclusive_cumsum(nnz_per_column, 0)
+        if nnz_per_column.dim() == 0:
+            # This addresses an edge case when ffn_hidden_size is equal to self.blocking.
+            nnz_per_column = nnz_per_column.unsqueeze(0)
+        offsets_t = torch.cat([zero, nnz_per_column])
+        return column_indices_t, offsets_t, block_offsets_t
+
+    def topology(self, x, padded_bins):
+        padded_tokens, _ = x.size()
+        assert padded_tokens % self.blocking == 0
+        if self.ffn_hidden_size % self.blocking != 0:
+            raise ValueError(
+                f'The ffn_hidden_size {self.ffn_hidden_size} must be divisible by ' +
+                f'the block size {self.blocking}. Please update your configuration.',
+            )
+
+        # Offsets for the sparse matrix. All rows have the
+        # same number of nonzero blocks dictated by the
+        # dimensionality of a single expert.
+        block_rows = padded_tokens // self.blocking
+        blocks_per_row = self.ffn_hidden_size // self.blocking
+        offsets = torch.arange(
+            0,
+            block_rows * blocks_per_row + 1,
+            blocks_per_row,
+            dtype=torch.int32,
+            device=x.device,
+        )
+
+        # Indices for the sparse matrix. The indices for
+        # the intermediate matrix are dynamic depending
+        # on the mapping of tokens to experts.
+        column_indices = ops.topology(
+            padded_bins,
+            self.blocking,
+            block_rows,
+            blocks_per_row,
+        )
+
+        # TODO(tgale): This is unused. Remove the need for this in stk.
+        # For now, use meta init to save the device memory.
+        data = torch.empty(
+            column_indices.numel(),
+            self.blocking,
+            self.blocking,
+            dtype=common.dtype(self.args),
+            device='meta',
+        )
+        shape = (
+            padded_tokens,
+            self.ffn_hidden_size * mpu.experts_per_rank(self.args),
+        )
+        row_indices = stk.ops.row_indices(shape, data, offsets, column_indices)
+        column_indices_t, offsets_t, block_offsets_t = self.sparse_transpose(
+            shape,
+            row_indices,
+            column_indices,
+            offsets,
+        )
+        return stk.Matrix(
+            shape,
+            data,
+            row_indices,
+            column_indices,
+            offsets,
+            column_indices_t,
+            offsets_t,
+            block_offsets_t,
+        )
+
+    def indices_and_padded_bins(self, top_experts):
+        # Sort the expert ids to produce the scatter/gather
+        # indices for the permutation.
+        top_experts = top_experts.int()
+        bin_ids, indices = ops.sort(top_experts, self.sort_end_bit)
+
+        # Histogram the expert ids to identify the number of
+        # tokens routed to each expert.
+        tokens_per_expert = ops.histogram(top_experts, self.num_experts)
+
+        # Round the token counts up to the block size used in
+        # the matrix muliplications. Caculate the starting
+        # position of each bin.
+        padded_tokens_per_expert = ops.round_up(
+            tokens_per_expert,
+            self.blocking,
+        )
+        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+        padded_bins = promote_scalar(padded_bins)
+
+        # Calculate the bin bounds for the sorted tokens.
+        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+        bins = promote_scalar(bins)
+        return indices, bin_ids, bins, padded_bins, tokens_per_expert
+
+    def sparse_forward_once(self, x, expert_weights, top_experts):
+        # x: [sl, bs, hs]
+        # expert_weights: [sl * bs, top-k]
+        # top_experts: [sl * bs, top-k]
+        expert_weights = expert_weights.flatten()
+        top_experts = top_experts.flatten()
+        with torch.no_grad():
+            indices, bin_ids, bins, padded_bins, tokens_per_expert = (self.indices_and_padded_bins(top_experts))
+
+        # Route the tokens for MoE computation.
+        x = x.view(-1, x.shape[-1])
+        x = ops.padded_gather(
+            x,
+            indices,
+            bin_ids,
+            bins,
+            padded_bins,
+            self.top_k,
+        )
+
+        # Create the sparse matrix topology.
+        with torch.no_grad():
+            topo = self.topology(x, padded_bins)
+
+        # Perform the expert computation.
+        x = self.mlp(x, topo)
+
+        # Un-route the data for the MoE output.
+        x = ops.padded_scatter(
+            x,
+            indices,
+            bin_ids,
+            expert_weights,
+            bins,
+            padded_bins,
+            self.top_k,
+        )
+        return x, tokens_per_expert
+
+    # For use in the base-class parallel_forward_once.
+    def sparse_permute_and_compute(
+        self,
+        x,
+        tokens_per_expert,
+        indices,
+        bin_ids,
+        expert_weights,
+        bins,
+        expert_capactiy,  # unused
+        top_k,
+    ):
+
+        # Round the token counts up to the block size used in the matrix
+        # multiplication. Calculate the starting position of each bin.
+        padded_tokens_per_expert = ops.round_up(
+            tokens_per_expert,
+            self.blocking,
+        )
+        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+        padded_bins = promote_scalar(padded_bins)
+
+        # Route the tokens for MoE computation.
+        x = x.view(-1, x.shape[-1])
+        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+
+        # Create the sparse matrix topology.
+        with torch.no_grad():
+            topo = self.topology(x, padded_bins)
+
+        # Perform the expert computation.
+        x = self.mlp(x, topo)
+
+        # Un-route the data for the MoE output.
+        return ops.padded_scatter(
+            x,
+            indices,
+            bin_ids,
+            expert_weights,
+            bins,
+            padded_bins,
+            top_k,
+        )
+
+    def grouped_forward_once(self, x, expert_weights, top_experts):
+        # x: [sl, bs, hs]
+        # expert_weights: [sl * bs, top-k]
+        # top_experts: [sl * bs, top-k]
+        expert_weights = expert_weights.flatten()
+        top_experts = top_experts.flatten()
+        with torch.no_grad():
+            indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts))
+
+        out = self.grouped_permute_and_compute(
+            x,
+            tokens_per_expert,
+            indices,
+            bin_ids,
+            expert_weights,
+            bins,
+            -1,  # unused
+            self.args.moe_top_k,
+        )
+        return out, tokens_per_expert
+
+    def grouped_permute_and_compute(
+        self,
+        x,
+        tokens_per_expert,
+        indices,
+        bin_ids,
+        expert_weights,
+        bins,
+        expert_capactiy,  # unused
+        top_k,
+    ):
+
+        # Route the tokens for MoE computation.
+        x = x.view(-1, x.shape[-1])
+        x = ops.gather(x, indices, bin_ids, bins, top_k)
+
+        # Perform the expert computation.
+        x = self.mlp(x, tokens_per_expert)
+
+        # Un-route the data for the MoE output.
+        return ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k)
+
+    def forward_once(self, x, expert_weights, top_experts):
+        if self.args.mlp_impl == 'sparse':
+            return self.sparse_forward_once(x, expert_weights, top_experts)
+        else:
+            return self.grouped_forward_once(x, expert_weights, top_experts)
+
+    def permute_and_compute(
+        self,
+        x,
+        tokens_per_expert,
+        indices,
+        bin_ids,
+        expert_weights,
+        bins,
+        expert_capactiy,
+        top_k,
+    ):
+        if self.args.mlp_impl == 'sparse':
+            return self.sparse_permute_and_compute(
+                x,
+                tokens_per_expert,
+                indices,
+                bin_ids,
+                expert_weights,
+                bins,
+                expert_capactiy,
+                top_k,
+            )
+        else:
+            return self.grouped_permute_and_compute(
+                x,
+                tokens_per_expert,
+                indices,
+                bin_ids,
+                expert_weights,
+                bins,
+                expert_capactiy,
+                top_k,
+            )
+
+
+class dMoE(moe.MoE):
+
+    def _init_experts_mlp(self, args: Arguments):
+        return ParallelDroplessMLP(args)

build/torch211-cxx11-cu126-aarch64-linux/_layers/gelu.py

@@ -0,0 +1,52 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+# try:
+#     import stk
+# except ImportError:
+#     import warnings
+#     warnings.warn(
+#         'Please add `stanford-stk` if megablocks/_layers/gelu.py is needed.',
+#     )
+
+from .. import stk
+
+import torch
+import torch.nn.functional as F
+
+
+@torch.jit.script
+def _gelu_backward_inplace(g, x):
+    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+    ff = (0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out))
+    return g.mul_(ff)
+
+
+def gelu_backward_(grad: stk.Matrix, x: stk.Matrix):
+    # NOTE: The two sparse matrices must have the same topology.
+    if isinstance(grad, stk.Matrix) and isinstance(x, stk.Matrix):
+        return stk.Matrix(
+            x.size(),
+            _gelu_backward_inplace(grad.data, x.data),
+            x.row_indices,
+            x.column_indices,
+            x.offsets,
+            x.column_indices_t,
+            x.offsets_t,
+            x.block_offsets_t,
+        )
+    return _gelu_backward_inplace(grad, x)
+
+
+def gelu(x: stk.Matrix):
+    assert isinstance(x, stk.Matrix)
+    return stk.Matrix(
+        x.size(),
+        F.gelu(x.data, approximate='tanh'),
+        x.row_indices,
+        x.column_indices,
+        x.offsets,
+        x.column_indices_t,
+        x.offsets_t,
+        x.block_offsets_t,
+    )

build/torch211-cxx11-cu126-aarch64-linux/_layers/glu.py

@@ -0,0 +1,244 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+# import stk.ops
+# try:
+#     import stk.ops
+# except ImportError:
+#     import warnings
+#     warnings.warn(
+#         'Please add `stanford-stk` if megablocks/_layers/glu.py is needed.',
+#     )
+
+from .. import stk
+
+import torch
+
+# from megablocks import grouped_gemm_util as gg
+# from megablocks.layers import common, mpu
+# from megablocks.layers.activation_fn import act_fn
+# from megablocks.layers.arguments import Arguments
+# from megablocks.layers.mlp import (
+#     SharedMLP,
+#     SparseMLP,
+#     create_dmoe_expert_weights,
+#     resolve_dtensor,
+# )
+
+from .. import grouped_gemm_util as gg
+from . import common, mpu
+from .activation_fn import act_fn
+from .arguments import Arguments
+from .mlp import (
+    SharedMLP,
+    SparseMLP,
+    create_dmoe_expert_weights,
+    resolve_dtensor,
+)
+
+
+class SparseGLU(SparseMLP):
+
+    def __init__(self, args: Arguments):
+        super().__init__(args)
+        self.v1 = torch.nn.Parameter(
+            torch.empty(
+                self._num_rows_per_rank,
+                args.hidden_size,
+                device=args.device,
+                dtype=common.dtype(args),
+            ),
+        )
+        with torch.no_grad():
+            self.v1.copy_(
+                create_dmoe_expert_weights(
+                    args,
+                    args.moe_num_experts,
+                    args.ffn_hidden_size,
+                    args.hidden_size,
+                    args.init_method,
+                ),
+            )
+
+        mpu.set_expert_model_parallel_attributes(
+            self.v1,
+            self._should_set_parallelism_attribute,
+        )
+
+    def forward(self, x, topo):
+        if self.args.memory_optimized_mlp:
+            raise NotImplementedError(
+                'Memory optimized implementation not yet supported with GLU with sparse kernels.',
+            )
+
+        w1, v1, w2 = self.scale_grad(self.w1), self.scale_grad(self.v1,), self.scale_grad(self.w2)
+        w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2)
+
+        # Compute the GLU.
+        x1 = stk.ops.sdd(x, w1.t(), topo)
+        x2 = stk.ops.sdd(x, v1.t(), topo)
+
+        activation_fn_out = act_fn(x1, self.args.activation_fn)
+        x1 = stk.ops.mul(activation_fn_out, x2)
+
+        return stk.ops.dsd(x1, w2)
+
+
+class MemoryOptimizedGroupedGLU(torch.autograd.Function):
+    """GroupedMLP with manually scheduled memory reuse."""
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_fwd(device_type='cuda')
+    def forward(ctx, x, w1, v1, w2, batch_sizes, activation_fn):
+        # Cast inputs using ctx dtype from AMP
+        if ctx._fwd_used_autocast:
+            x = x.to(ctx._dtype)
+            w1 = w1.to(ctx._dtype)
+            v1 = v1.to(ctx._dtype)
+            w2 = w2.to(ctx._dtype)
+        # x: [m, k], w1: [n, k], v1: [n, k], w2: [n, k]
+        if (not x.is_contiguous() or not w1.is_contiguous() or not v1.is_contiguous() or not w2.is_contiguous()):
+            raise ValueError("Expected contiguous 'x', 'w1', 'v1' and 'w2'.")
+
+        # Layer 0: x @ w1.t().
+        assert gg.backend is not None
+        sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True)
+        v1_out = gg.backend.gmm(x, v1, batch_sizes, trans_b=True)
+
+        # GeLU.
+        activation_fn_out = activation_fn(sdd_out) * v1_out
+
+        # Layer 1: x @ w2.
+        dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes)
+
+        # NOTE: Save the input to the layer and the activation_fn input for
+        # gradient computation. We'll re-compute the activation_fn forward
+        # pass in the backward pass to avoid materializing another
+        # intermediate.
+        ctx.x_shape = x.shape
+        ctx.sdd_out_shape = sdd_out.shape
+        ctx.dtype = x.dtype
+        ctx.activation_fn = activation_fn
+        ctx.save_for_backward(w1, v1, w2, batch_sizes, x, sdd_out, v1_out)
+        return dsd_out
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_bwd(device_type='cuda')
+    def backward(ctx, ddsd_out):
+        if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]):
+            raise ValueError('Expected all MLP inputs to need grad.')
+
+        # Unpack saved tensors
+        # dtype = ctx.dtype
+        saved_tensors = ctx.saved_tensors
+        w1, v1, w2 = saved_tensors[:3]
+        batch_sizes = saved_tensors[3]
+        x = saved_tensors[4]
+        sdd_out, v1_out = saved_tensors[5:7]
+
+        # Rematerialize activation_fn output.
+        activation_fn = ctx.activation_fn
+        with torch.set_grad_enabled(True):
+            sdd_out.requires_grad = True
+            v1_out.requires_grad = True
+            activation_fn_out = activation_fn(sdd_out) * v1_out
+            activation_grad_fn = activation_fn_out.backward
+
+        # Compute dw2 with recomputed activation_fn output.
+        assert gg.backend is not None
+        dw2 = gg.backend.gmm(
+            activation_fn_out,
+            ddsd_out,
+            batch_sizes,
+            trans_a=True,
+        )
+
+        # Compute dactivation_fn_out.
+        #
+        # NOTE: We reuse the activation_fn_out allocation.
+        dactivation_fn_out = activation_fn_out
+        gg.backend.gmm(
+            ddsd_out,
+            w2,
+            batch_sizes,
+            trans_b=True,
+            c=dactivation_fn_out,
+        )
+
+        # Compute dsdd_out.
+        #
+        # NOTE: This reuses the dactivation_fn_out allocation.
+        assert activation_grad_fn is not None
+        activation_grad_fn(dactivation_fn_out)
+        dsdd_out = sdd_out.grad
+        dv1_out = v1_out.grad
+
+        # Compute dw1.
+        dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True)
+
+        # Compute dv1.
+        dv1 = gg.backend.gmm(dv1_out, x, batch_sizes, trans_a=True)
+
+        # Compute dx.
+        #
+        # NOTE: This reuses the ddsd_out allocation.
+        dx = ddsd_out
+        gg.backend.gmm(dsdd_out, w1, batch_sizes, c=dx)
+        dx += gg.backend.gmm(dv1_out, v1, batch_sizes)
+        return dx, dw1, dv1, dw2, None, None
+
+
+memory_optimized_grouped_glu = MemoryOptimizedGroupedGLU.apply
+
+
+class GroupedGLU(SparseGLU):
+
+    def forward(self, x, tokens_per_expert):
+        batch_sizes = tokens_per_expert.cpu().to(torch.long)
+        w1, v1, w2 = (
+            self.scale_grad(self.w1),
+            self.scale_grad(self.v1),
+            self.scale_grad(self.w2),
+        )
+        w1, v1, w2 = resolve_dtensor(w1), resolve_dtensor(v1,), resolve_dtensor(w2)
+
+        # Re-shape the weights for the grouped GEMMs.
+        ne = mpu.experts_per_rank(self.args)
+        w1 = w1.view(ne, -1, self.args.hidden_size)
+        v1 = v1.view(ne, -1, self.args.hidden_size)
+        w2 = w2.view(ne, -1, self.args.hidden_size)
+
+        if self.args.memory_optimized_mlp:
+            return memory_optimized_grouped_glu(
+                x,
+                w1,
+                v1,
+                w2,
+                batch_sizes,
+                self.args.activation_fn,
+            )
+
+        # Compute the MLP.
+        assert gg.ops is not None
+        x1 = gg.ops.gmm(x, w1, batch_sizes, trans_b=True)
+        x2 = gg.ops.gmm(x, v1, batch_sizes, trans_b=True)
+        x1 = self.args.activation_fn(x1) * x2
+        return gg.ops.gmm(x1, w2, batch_sizes)
+
+
+class SharedGLU(SharedMLP):
+    """GPU for shared expert.
+
+    Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTGLU class
+    """
+
+    def __init__(self, args: Arguments):
+        super().__init__(args)
+        self.gate_proj = args.fc_cls(
+            args.hidden_size,
+            self.args.shared_expert_hidden_size,
+            **self.fc_kwargs,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))

build/torch211-cxx11-cu126-aarch64-linux/_layers/memory_test.py

@@ -0,0 +1,103 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+
+import torch
+import torch.distributed as dist
+
+# from megablocks.layers import arguments, dmoe
+from . import arguments, dmoe
+
+_TESTS = ((8, 2048, 4096, 4096, 32, 4),)
+
+
+def get_tensors():
+    ptrs = set()
+    out = []
+    for obj in gc.get_objects():
+        if torch.is_tensor(obj):
+            if not obj.is_contiguous() or obj.data_ptr() in ptrs:
+                continue
+            out.append(obj)
+            ptrs.add(obj.data_ptr())
+    return out
+
+
+def test_memory(
+    group,
+    batch_size,
+    sequence_length,
+    hidden_size,
+    ffn_hidden_size,
+    num_experts,
+    top_k,
+):
+    args = arguments.Arguments(
+        hidden_size=hidden_size,
+        ffn_hidden_size=ffn_hidden_size,
+        moe_num_experts=num_experts,
+        moe_top_k=top_k,
+        moe_expert_model_parallelism=True,
+        expert_parallel_group=group,
+        fp16=False,
+        bf16=True,
+        device=torch.cuda.current_device(),
+    )
+    layer = dmoe.dMoE(args).cuda()
+
+    x = torch.randn((batch_size, sequence_length, hidden_size),
+                    device=torch.cuda.current_device(),
+                    dtype=torch.bfloat16).requires_grad_(True)
+    torch.cuda.empty_cache()
+
+    # Run forward + backward.
+    # with torch.autograd.detect_anomaly():
+    out, _ = layer(x)
+    out.mean().backward()
+
+    # Report peak memory.
+    mem = torch.cuda.max_memory_allocated()
+    print('Max Memory Allocated = {:0.0f}MiB'.format(mem / 1e6))
+    print('Max Memory Reserved = {:0.0f}MiB'.format(torch.cuda.max_memory_reserved() / 1e6,),)
+
+    # Calculate weight and gradient memory usage.
+    weight_memory = 2 * (
+        layer.router.layer.weight.numel() + layer.experts.mlp.w1.numel() + layer.experts.mlp.w2.numel()
+    )
+
+    def grad_numel(x):
+        if x.grad is not None:
+            return x.grad.numel()
+        return 0
+
+    grad_memory = 2 * (
+        grad_numel(layer.router.layer.weight) + grad_numel(layer.experts.mlp.w1) + grad_numel(layer.experts.mlp.w2)
+    )
+    weight_memory += grad_memory
+
+    print('Weight Memory Allocated = {:0.0f}MiB'.format(weight_memory / 1e6))
+    print('Activation Memory Allocated = {:0.0f}MiB'.format((mem - weight_memory) / 1e6,),)
+
+    # Manually calculate GPU memory usage from the garbage
+    # collector.
+    gc.collect()
+    total = 0
+    tensors = get_tensors()
+    tensors = sorted(tensors, key=lambda x: -x.numel())
+    for i, t in enumerate(tensors):
+        total += t.numel()
+        print(f'{i}: {t.shape}, {t.numel() * 2}')
+    del tensors
+
+    print('Total Bytes Found = {:0.0f}MiB'.format(total * 2 / 1e6))
+
+
+if __name__ == '__main__':
+    assert dist.is_available()
+    group = dist.init_process_group(backend='nccl')
+    local_rank = dist.get_rank(group)
+    torch.cuda.set_device(local_rank)
+
+    for args in _TESTS:
+        test_memory(group, *args)

build/torch211-cxx11-cu126-aarch64-linux/_layers/mlp.py

@@ -0,0 +1,587 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+# try:
+#     import stk
+#     import stk.backend.triton_kernels
+#     import stk.ops
+# except ImportError:
+#     import warnings
+#     warnings.warn(
+#         'Please add `stanford-stk` if megablocks/_layers/mlp.py is needed.',
+#     )
+
+from .. import stk
+
+import torch
+from packaging import version
+
+# from megablocks import grouped_gemm_util as gg
+# from megablocks.layers import common, gelu, mpu
+# from megablocks.layers.activation_fn import act_fn
+# from megablocks.layers.arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn
+
+from .. import grouped_gemm_util as gg
+from . import common, gelu, mpu
+from .activation_fn import act_fn
+from .arguments import DEFAULT_ACTIVATION_FN, Arguments, InitFn
+
+class ScaleGradient(torch.autograd.Function):
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_fwd(device_type='cuda')
+    def forward(ctx: Any, x: torch.Tensor, scale: float):
+        ctx.scale = scale
+        return x
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_bwd(device_type='cuda')
+    def backward(ctx: torch.Tensor, grad: torch.Tensor):
+        return grad * ctx.scale, None
+
+
+scale_gradient = ScaleGradient.apply
+
+
+def resolve_dtensor(weight: torch.Tensor):
+    if version.parse(torch.__version__) >= version.parse('2.0.0'):
+        from torch.distributed._tensor import DTensor
+        if isinstance(weight, DTensor):
+            return weight.to_local()
+    return weight
+
+
+def create_moe_expert_weights(
+    args: Arguments,
+    num_experts: int,
+    ffn_hidden_size: int,
+    hidden_size: int,
+    init_method: InitFn,
+):
+    # Create the entire weight matrix such that the sampled weights will
+    # not vary between data parallelism and expert model parallelism for
+    # the same random seed.
+    master_weights = torch.empty(
+        num_experts,
+        ffn_hidden_size,
+        hidden_size,
+        device=args.device,
+        dtype=common.dtype(args),
+    )
+    init_method(master_weights)
+
+    if not args.moe_expert_model_parallelism:
+        return master_weights
+
+    # Calculate the amount of sharding in each dimension.
+    expert_sharding_degree = mpu.expert_sharding_degree(args)
+    hidden_sharding_degree = mpu.hidden_sharding_degree(args)
+
+    # Calculate the experts per rank.
+    #
+    # NOTE: We assign ranks to be expert parallel before going
+    # tensor parallel.
+    rank = mpu.get_expert_parallel_rank(args)
+    expert_rank = rank % expert_sharding_degree
+    num_experts_per_rank = num_experts // expert_sharding_degree
+    start_expert = expert_rank * num_experts_per_rank
+    end_expert = (expert_rank + 1) * num_experts_per_rank
+
+    # Calculate the rows per rank.
+    row_rank = rank // expert_sharding_degree
+    num_rows_per_rank = ffn_hidden_size // hidden_sharding_degree
+    start_row = row_rank * num_rows_per_rank
+    end_row = (row_rank + 1) * num_rows_per_rank
+
+    # Slice the weight matrix to get the chunk for this rank.
+    with torch.no_grad():
+        weights = master_weights[start_expert:end_expert, start_row:end_row]
+    return weights
+
+
+class MLP(torch.nn.Module):
+
+    def __init__(self, args: Arguments):
+        super().__init__()
+        self.args = args
+        # expert_parallel_world_size = mpu.get_expert_parallel_world_size(args)
+        experts_per_rank = mpu.experts_per_rank(args)
+
+        self.w1 = torch.nn.Parameter(
+            torch.empty(
+                experts_per_rank,
+                args.hidden_size,
+                mpu.features_per_rank(args),
+                device=args.device,
+                dtype=common.dtype(args),
+            ),
+        )
+        self.w2 = torch.nn.Parameter(
+            torch.empty(
+                experts_per_rank,
+                mpu.features_per_rank(args),
+                args.hidden_size,
+                device=args.device,
+                dtype=common.dtype(args),
+            ),
+        )
+        mpu.set_expert_model_parallel_attributes(
+            self.w1,
+            args.moe_expert_model_parallelism,
+        )
+        mpu.set_expert_model_parallel_attributes(
+            self.w2,
+            args.moe_expert_model_parallelism,
+        )
+
+        # Initialize the parameters for the MLP.
+        #
+        # NOTE: It is important that we create the weight tensors prior
+        # to creating the master weights and slicing our the piece for
+        # this rank. If the master weights are created first the PyTorch
+        # caching allocator appears to use the same memory block for these
+        # and the slice which causes large increases in our peak memory
+        # usage.
+        with torch.no_grad():
+            w1 = create_moe_expert_weights(
+                args,
+                args.moe_num_experts,
+                args.ffn_hidden_size,
+                args.hidden_size,
+                args.init_method,
+            )
+            self.w1.copy_(w1.transpose(1, 2).contiguous())
+            self.w2.copy_(
+                create_moe_expert_weights(
+                    args,
+                    args.moe_num_experts,
+                    args.ffn_hidden_size,
+                    args.hidden_size,
+                    args.output_layer_init_method,
+                ),
+            )
+
+        self.gradient_scale = None
+        if self.args.moe_expert_model_parallelism:
+            self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,)
+
+    def scale_grad(self, w):
+        if self.gradient_scale is None:
+            return w
+        return scale_gradient(w, self.gradient_scale)
+
+    def forward(self, x):
+        w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2)
+        w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2)
+        x = torch.bmm(x, w1)
+        x = self.args.activation_fn(x)
+        return torch.bmm(x, w2)
+
+
+def create_dmoe_expert_weights(
+    args: Arguments,
+    num_experts: int,
+    rows: int,
+    columns: int,
+    init_method: InitFn,
+):
+    weights = create_moe_expert_weights(
+        args,
+        num_experts,
+        rows,
+        columns,
+        init_method,
+    )
+    return weights.view([-1, columns])
+
+
+class MemoryOptimizedMLP(torch.autograd.Function):
+    """Sparse MLP with manually scheduled memory reuse."""
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_fwd(device_type='cuda')
+    def forward(ctx, x, w1, w2, topo, activation_fn):
+        # Cast inputs using ctx dtype from AMP
+        if ctx._fwd_used_autocast:
+            x = x.to(ctx._dtype)
+            w1 = w1.to(ctx._dtype)
+            w2 = w2.to(ctx._dtype)
+        # x: [m, k], w1: [n, k], w2: [n, k]
+        if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()):
+            raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.")
+
+        topo_tensors = (
+            topo.row_indices,
+            topo.column_indices,
+            topo.offsets,
+            topo.column_indices_t,
+            topo.offsets_t,
+            topo.block_offsets_t,
+        )
+
+        # Layer 0: x @ w1.t().
+        sdd_out = stk.ops.sdd(x, w1.t(), topo)
+
+        # GeLU.
+        activation_fn_out = act_fn(sdd_out, activation_fn)
+
+        # Layer 1: x @ w2.
+        dsd_out = stk.ops.dsd(activation_fn_out, w2)
+
+        # NOTE: Save the input to the layer and the activation_fn input for
+        # gradient computation. We'll re-compute the activation_fn forward
+        # pass in the backward pass to avoid materializing another
+        # intermediate.
+        ctx.shape = topo.shape
+        ctx.x_shape = x.shape
+        ctx.sdd_out_shape = sdd_out.data.shape
+        ctx.dtype = x.dtype
+        ctx.activation_fn = activation_fn
+        ctx.save_for_backward(w1, w2, *topo_tensors, x, sdd_out.data)
+        return dsd_out
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_bwd(device_type='cuda')
+    def backward(ctx, ddsd_out):
+        if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]):
+            raise ValueError('Expected all MLP inputs to need grad.')
+
+        # unpack saved tensors
+        # dtype = ctx.dtype
+        saved_tensors = ctx.saved_tensors
+        w1, w2 = saved_tensors[:2]
+        topo_tensors = saved_tensors[2:8]
+        x = saved_tensors[8]
+        sdd_out_data = saved_tensors[9]
+
+        # rematerialize activation function output
+        activation_fn = ctx.activation_fn
+        sdd_out = stk.Matrix(ctx.shape, sdd_out_data, *topo_tensors)
+        activation_fn_out, activation_grad_fn = act_fn(
+            sdd_out,
+            activation_fn,
+            return_grad_fn=True,
+        )
+
+        # Compute dw2 with recomputed activation_fn output.
+        dw2 = stk.ops.dsd(activation_fn_out.t(), ddsd_out)
+
+        # Compute dactivation_fn_out.
+        #
+        # NOTE: We reuse the activation_fn_out allocation.
+        dactivation_fn_out = activation_fn_out
+        stk.backend.triton_kernels.sdd(
+            ddsd_out,
+            w2.t(),
+            dactivation_fn_out.shape,
+            dactivation_fn_out.data,
+            dactivation_fn_out.offsets,
+            dactivation_fn_out.row_indices,
+            dactivation_fn_out.column_indices,
+        )
+
+        # Compute dsdd_out.
+        #
+        # NOTE: This reuses the dactivation_fn_out allocation.
+        if activation_fn is DEFAULT_ACTIVATION_FN:
+            dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out)
+        else:
+            assert activation_grad_fn is not None
+            activation_grad_fn(dactivation_fn_out.data)
+            dsdd_out = stk.Matrix(ctx.shape, sdd_out.data.grad, *topo_tensors)
+
+        # Compute dw1.
+        dw1 = stk.ops.dsd(dsdd_out.t(), x)
+
+        # Compute dx.
+        #
+        # NOTE: This reuses the ddsd_out allocation.
+        stk.backend.triton_kernels.dsd(
+            dsdd_out.shape,
+            dsdd_out.data,
+            dsdd_out.offsets,
+            dsdd_out.row_indices,
+            dsdd_out.column_indices,
+            dsdd_out.offsets_t,
+            dsdd_out.column_indices_t,
+            dsdd_out.block_offsets_t,
+            False,
+            w1,
+            ddsd_out,
+        )
+        dx = ddsd_out
+        return dx, dw1, dw2, None, None
+
+
+memory_optimized_mlp = MemoryOptimizedMLP.apply
+
+
+class SparseMLP(torch.nn.Module):
+
+    def __init__(self, args: Arguments):
+        super().__init__()
+        self.args = args
+        self._num_rows_per_rank = mpu.experts_per_rank(args) * mpu.features_per_rank(args)
+
+        self.w1 = torch.nn.Parameter(
+            torch.empty(
+                self._num_rows_per_rank,
+                args.hidden_size,
+                device=args.device,
+                dtype=common.dtype(args),
+            ),
+        )
+        self.w2 = torch.nn.Parameter(
+            torch.empty(
+                self._num_rows_per_rank,
+                args.hidden_size,
+                device=args.device,
+                dtype=common.dtype(args),
+            ),
+        )
+
+        # Initialize the parameters for the MLP.
+        #
+        # NOTE: It is important that we create the weight tensors prior
+        # to creating the master weights and slicing our the piece for
+        # this rank. If the master weights are created first the PyTorch
+        # caching allocator appears to use the same memory block for these
+        # and the slice which causes large increases in our peak memory
+        # usage.
+        with torch.no_grad():
+            self.w1.copy_(
+                create_dmoe_expert_weights(
+                    args,
+                    args.moe_num_experts,
+                    args.ffn_hidden_size,
+                    args.hidden_size,
+                    args.init_method,
+                ),
+            )
+            self.w2.copy_(
+                create_dmoe_expert_weights(
+                    args,
+                    args.moe_num_experts,
+                    args.ffn_hidden_size,
+                    args.hidden_size,
+                    args.output_layer_init_method,
+                ),
+            )
+
+        self._should_set_parallelism_attribute = args.moe_expert_model_parallelism
+        mpu.set_expert_model_parallel_attributes(
+            self.w1,
+            self._should_set_parallelism_attribute,
+        )
+        mpu.set_expert_model_parallel_attributes(
+            self.w2,
+            self._should_set_parallelism_attribute,
+        )
+
+        self.gradient_scale = None
+        if self.args.moe_expert_model_parallelism:
+            self.gradient_scale = 1 / mpu.get_expert_parallel_world_size(self.args,)
+
+    def scale_grad(self, w):
+        if self.gradient_scale is None:
+            return w
+        return scale_gradient(w, self.gradient_scale)
+
+    def forward(self, x, topo):
+        w1, w2 = self.scale_grad(self.w1), self.scale_grad(self.w2)
+        w1, w2 = resolve_dtensor(w1), resolve_dtensor(w2)
+        if self.args.memory_optimized_mlp:
+            return memory_optimized_mlp(
+                x,
+                w1,
+                w2,
+                topo,
+                self.args.activation_fn,
+            )
+
+        # Compute the MLP.
+        x = stk.ops.sdd(x, w1.t(), topo)
+        activation_fn_out = act_fn(x, self.args.activation_fn)
+        return stk.ops.dsd(activation_fn_out, w2)
+
+
+class MemoryOptimizedGroupedMLP(torch.autograd.Function):
+    """GroupedMLP with manually scheduled memory reuse."""
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_fwd(device_type='cuda')
+    def forward(ctx, x, w1, w2, batch_sizes, activation_fn):
+        # Cast inputs using ctx dtype from AMP
+        if ctx._fwd_used_autocast:
+            x = x.to(ctx._dtype)
+            w1 = w1.to(ctx._dtype)
+            w2 = w2.to(ctx._dtype)
+        # x: [m, k], w1: [n, k], w2: [n, k]
+        if (not x.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous()):
+            raise ValueError("Expected contiguous 'x', 'w1' and 'w2'.")
+
+        # Layer 0: x @ w1.t().
+        assert gg.backend is not None
+        sdd_out = gg.backend.gmm(x, w1, batch_sizes, trans_b=True)
+
+        # activation_fn
+        activation_fn_out = activation_fn(sdd_out)
+
+        # Layer 1: x @ w2.
+        dsd_out = gg.backend.gmm(activation_fn_out, w2, batch_sizes)
+
+        # NOTE: Save the input to the layer and the activation_fn input for
+        # gradient computation. We'll re-compute the activation_fn forward
+        # pass in the backward pass to avoid materializing another
+        # intermediate.
+        ctx.x_shape = x.shape
+        ctx.sdd_out_shape = sdd_out.shape
+        ctx.dtype = x.dtype
+        ctx.activation_fn = activation_fn
+        ctx.save_for_backward(w1, w2, batch_sizes, x, sdd_out)
+        return dsd_out
+
+    @staticmethod
+    @torch.amp.autocast_mode.custom_bwd(device_type='cuda')
+    def backward(ctx: Any, ddsd_out: torch.Tensor):
+        if (not ctx.needs_input_grad[0] or not ctx.needs_input_grad[1] or not ctx.needs_input_grad[2]):
+            raise ValueError('Expected all MLP inputs to need grad.')
+
+        # Unpack saved tensors
+        # dtype = ctx.dtype
+        saved_tensors = ctx.saved_tensors
+        w1, w2 = saved_tensors[:2]
+        batch_sizes = saved_tensors[2]
+        x = saved_tensors[3]
+        sdd_out = saved_tensors[4]
+
+        # Rematerialize activation_fn output.
+        activation_fn = ctx.activation_fn
+        with torch.set_grad_enabled(True):
+            sdd_out.requires_grad = True
+            activation_fn_out = activation_fn(sdd_out)
+            activation_grad_fn = activation_fn_out.backward
+
+        # Compute dw2 with recomputed activation_fn output.
+        assert gg.backend is not None
+        dw2 = gg.backend.gmm(
+            activation_fn_out,
+            ddsd_out,
+            batch_sizes,
+            trans_a=True,
+        )
+
+        # Compute dactivation_fn_out.
+        #
+        # NOTE: We reuse the activation_fn_out allocation.
+        dactivation_fn_out = activation_fn_out
+        gg.backend.gmm(
+            ddsd_out,
+            w2,
+            batch_sizes,
+            trans_b=True,
+            c=dactivation_fn_out,
+        )
+
+        # Compute dsdd_out.
+        #
+        # NOTE: This reuses the dactivation_fn_out allocation.
+        if activation_fn is DEFAULT_ACTIVATION_FN:
+            dsdd_out = gelu.gelu_backward_(dactivation_fn_out, sdd_out)
+        else:
+            assert activation_grad_fn is not None
+            activation_grad_fn(dactivation_fn_out)
+            dsdd_out = sdd_out.grad
+
+        # Compute dw1.
+        dw1 = gg.backend.gmm(dsdd_out, x, batch_sizes, trans_a=True)
+
+        # Compute dx.
+        #
+        # NOTE: This reuses the ddsd_out allocation.
+        gg.backend.gmm(dsdd_out, w1, batch_sizes, c=ddsd_out)
+        dx = ddsd_out
+        return dx, dw1, dw2, None, None
+
+
+memory_optimized_grouped_mlp = MemoryOptimizedGroupedMLP.apply
+
+
+class GroupedMLP(SparseMLP):
+
+    def forward(self, x, tokens_per_expert):
+        batch_sizes = tokens_per_expert.cpu().to(torch.long)
+        w1, w2 = (self.scale_grad(self.w1), self.scale_grad(self.w2))
+
+        # Re-shape the weights for the grouped GEMMs.
+        ne = mpu.experts_per_rank(self.args)
+        w1 = resolve_dtensor(w1).view(ne, -1, self.args.hidden_size)
+        w2 = resolve_dtensor(w2).view(ne, -1, self.args.hidden_size)
+
+        if self.args.memory_optimized_mlp:
+            return memory_optimized_grouped_mlp(
+                x,
+                w1,
+                w2,
+                batch_sizes,
+                self.args.activation_fn,
+            )
+
+        # Compute the MLP.
+        assert gg.ops is not None
+        x = gg.ops.gmm(x, w1, batch_sizes, trans_b=True)
+        x = self.args.activation_fn(x)
+        return gg.ops.gmm(x, w2, batch_sizes)
+
+
+class SharedMLP(torch.nn.Module):
+    """MLP for shared expert.
+
+    Note: this is a copy -> pasta -> modify of the LLM-Foundry MPTMLP class
+    """
+
+    def __init__(self, args: Arguments):
+        super().__init__()
+        self.args = args
+        self.fc_kwargs: dict[str, Any] = {
+            'bias': args.bias,
+            'device': args.device,
+        }
+        self.fc_kwargs.update(args.fc_kwargs)
+
+        self.up_proj = args.fc_cls(
+            args.hidden_size,
+            args.shared_expert_hidden_size,
+            **self.fc_kwargs,
+        )
+        self.act = args.activation_fn
+        self.down_proj = args.fc_cls(
+            args.shared_expert_hidden_size,
+            args.hidden_size,
+            **self.fc_kwargs,
+        )
+        self.down_proj._is_residual = True  # a flag for llm-foundry init
+
+    def add_experts_sharedexpert(
+        self,
+        shared_expert_out: torch.Tensor,
+        expert_out: torch.Tensor,
+    ) -> torch.Tensor:
+        # Helper function to add expert output to shared expert output
+        # with optional weighted sum.
+        if self.args.shared_expert_weighted_sum:
+            # enable using weighted sum for shared expert output
+            # wieghted by number of experts used
+            t_experts = self.args.moe_top_k + 1
+            sh_mlp_out = shared_expert_out / t_experts
+            return sh_mlp_out.add(
+                expert_out,
+                alpha=(self.args.moe_top_k / t_experts),
+            )
+
+        return shared_expert_out + expert_out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act(self.up_proj(x)))

build/torch211-cxx11-cu126-aarch64-linux/_layers/moe.py

@@ -0,0 +1,507 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+# import megablocks.ops as ops
+# from megablocks.layers import common, mlp, mpu, router, sharedexpert_registry
+# from megablocks.layers.all_to_all import all_to_all
+# from megablocks.layers.arguments import Arguments
+
+from ..ops import (
+    sort,
+    histogram,
+    inclusive_cumsum,
+    exclusive_cumsum,
+    binned_gather,
+    binned_scatter,
+    gather,
+    scatter,
+    repeat,
+    replicate,
+)
+
+from . import common, mlp, mpu, router, sharedexpert_registry
+from .arguments import Arguments
+from .all_to_all import all_to_all
+
+_LOAD_BALANCING_LOSS = []
+
+
+def save_load_balancing_loss(loss):
+    global _LOAD_BALANCING_LOSS
+    _LOAD_BALANCING_LOSS.append(loss)
+
+
+def get_load_balancing_loss():
+    global _LOAD_BALANCING_LOSS
+    return _LOAD_BALANCING_LOSS
+
+
+def clear_load_balancing_loss():
+    global _LOAD_BALANCING_LOSS
+    _LOAD_BALANCING_LOSS.clear()
+
+
+def batched_load_balancing_loss(args: Arguments):
+    if args.moe_loss_weight == 0:
+        return 0.0
+
+    # tokens_per_expert[i].shape = (num_experts)
+    # expert_scores[i].shape = (tokens, num_experts)
+    tokens_per_expert, expert_scores = zip(*get_load_balancing_loss())
+    num_layers_per_pipeline_stage = (args.num_layers // args.pipeline_model_parallel_size)
+    if args.num_layers_per_virtual_pipeline_stage is not None:
+        num_layers_per_pipeline_stage = args.num_layers_per_virtual_pipeline_stage
+
+    if len(tokens_per_expert) != num_layers_per_pipeline_stage:
+        raise ValueError(
+            f'Expected {num_layers_per_pipeline_stage} token_per_experts '
+            f'but found {len(tokens_per_expert)}.\nnum_layers = '
+            f'{args.num_layers}\npipeline_model_parallel_size = '
+            f'{args.pipeline_model_parallel_size}\n'
+            'num_layers_per_virtual_pipeline_stage'
+            f' = {args.num_layers_per_virtual_pipeline_stage}',
+        )
+    if len(expert_scores) != num_layers_per_pipeline_stage:
+        raise ValueError(
+            f'Expected {num_layers_per_pipeline_stage} expert_scores '
+            f'but found {len(tokens_per_expert)}.\nnum_layers = '
+            f'{args.num_layers}\npipeline_model_parallel_size = '
+            f'{args.pipeline_model_parallel_size}\n'
+            'num_layers_per_virtual_pipeline_stage'
+            f' = {args.num_layers_per_virtual_pipeline_stage}',
+        )
+
+    # Verify the shape of the tokens_per_expert and expert_scores tensors.
+    assert all((x.ndim == 1 and x.numel() == args.moe_num_experts for x in tokens_per_expert))
+
+    tokens = expert_scores[0].shape[0]
+    assert all(((x.ndim == 2 and x.shape[1] == args.moe_num_experts and x.shape[0] == tokens) for x in expert_scores))
+
+    # Concatenate the contributions of each layer and convert to
+    # the correct types and formats for the dot product.
+    expert_scores = torch.cat(expert_scores, dim=1)
+    if args.moe_lbl_in_fp32:
+        expert_scores = expert_scores.float()
+    if tokens != 0:
+        expert_scores = expert_scores.mean(dim=0)
+    else:
+        expert_scores = expert_scores.sum(dim=0)
+    tokens_per_expert = torch.cat(tokens_per_expert).to(expert_scores.dtype)
+
+    expected_values = num_layers_per_pipeline_stage * args.moe_num_experts
+    assert tokens_per_expert.numel() == expected_values
+    assert expert_scores.numel() == expected_values
+
+    # Calculate the total scale across all factors.
+    #
+    # loss_weight * num_experts / (num_layers * tokens * top_k)
+    scale_numerator = (args.moe_num_experts * args.moe_loss_weight)
+    scale_denominator = (args.num_layers * tokens * args.moe_top_k)
+    scale = scale_numerator / scale_denominator
+    return scale * torch.dot(tokens_per_expert, expert_scores)
+
+
+# NOTE: This class defines MoE expert computation, including expert model parallel
+# communication. When using FSDP on top of MegaBlocks this is the module that should
+# be wrapped s.t. the weight all-gathers can be scheduled *before* the expert model
+# parallel all2all.
+class ParallelMLP(torch.nn.Module):
+
+    def __init__(self, args: Arguments):
+        super(ParallelMLP, self).__init__()
+        self.args = args
+
+        # Calculate the number of experts in total and the number of experts
+        # owned by this rank.
+        # world_size = mpu.get_expert_parallel_world_size(args)
+        self.num_experts = args.moe_num_experts
+        self.top_k = self.args.moe_top_k
+
+        # Calculate the number of bits needed to represent the expert indices
+        # so that we can pass it to radix sort.
+        self.sort_end_bit = max(int(np.ceil(np.log2(self.num_experts))), 1)
+
+        # Expert MLP.
+        self.mlp = mlp.MLP(args)
+
+        self.bias: Optional[torch.Tensor]
+        if self.args.bias:
+            # Note that the output bias is not parallelized with expert
+            # model parallelism.
+            self.bias = torch.nn.Parameter(
+                torch.empty(
+                    args.hidden_size,
+                    device=args.device,
+                    dtype=common.dtype(args),
+                ),
+            )
+            torch.nn.init.zeros_(self.bias)
+        else:
+            self.register_parameter('bias', None)
+
+        # Select the forward function for the operating mode.
+        self.forward_fn = (self.parallel_forward_once if args.moe_expert_model_parallelism else self.forward_once)
+
+    def expert_capacity(self, tokens: int) -> int:
+        world_size = mpu.get_expert_parallel_world_size(self.args)
+        tokens_per_expert = (self.top_k * tokens * world_size / self.num_experts)
+        return int(self.args.moe_capacity_factor * tokens_per_expert)
+
+    def load_balancing_loss(self, tokens_per_expert: torch.Tensor, expert_scores: torch.Tensor):
+        """Calculate the load balancing loss contribution."""
+        assert len(expert_scores.size()) == 2
+        tokens, num_experts = expert_scores.size()
+        assert num_experts == self.num_experts
+        assert len(tokens_per_expert.size()) == 1
+        num_experts, = tokens_per_expert.size()
+        assert num_experts == self.num_experts
+        scale = self.num_experts / (tokens * self.top_k)
+        return scale * torch.dot(
+            tokens_per_expert.to(expert_scores.dtype),
+            expert_scores.mean(dim=0),
+        )
+
+    def indices_and_bins(self,
+                         top_expert: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Sort the expert ids to produce the scatter/gather
+        # indices for the permutation.
+        #
+        # TODO(tgale): Is it worth doing this conversion to 32-bit
+        # prior? Could we place the `torch.max` operation to return
+        # 32-bit expert indices?
+        top_expert = top_expert.int()
+        # output = ops.sort(top_expert, self.sort_end_bit)
+        output = sort(top_expert, self.sort_end_bit)
+        assert output is not None
+        bin_ids, indices = output
+
+        # Histogram the expert ids to identify the number of
+        # tokens routed to each expert.
+        #
+        # TODO(tgale): Does the sorted data produce a more favorable
+        # data distribution for histogram? Or is the op parallelism
+        # worth more?
+        # tokens_per_expert = ops.histogram(top_expert, self.num_experts)
+        tokens_per_expert = histogram(top_expert, self.num_experts)
+
+        # Calculate the bin bounds for the sorted tokens.
+        # bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+        bins = inclusive_cumsum(tokens_per_expert, 0)
+        assert bins is not None
+        bins = bins.view(1) if not len(bins.size()) else bins
+
+        assert isinstance(indices, torch.Tensor)
+        assert isinstance(bin_ids, torch.Tensor)
+        assert isinstance(bins, torch.Tensor)
+        assert isinstance(tokens_per_expert, torch.Tensor)
+
+        return indices, bin_ids, bins, tokens_per_expert
+
+    def permute_and_compute(
+        self,
+        x: torch.Tensor,
+        tokens_per_expert: int,  # unused
+        indices: torch.Tensor,
+        bin_ids: torch.Tensor,  # unused
+        expert_weights: torch.Tensor,
+        bins: torch.Tensor,
+        expert_capacity: int,
+        top_k: int,
+    ):
+        # Route the tokens for MoE computation.
+        x = x.view(-1, x.shape[-1])
+        # output = ops.binned_gather(x, indices, bins, expert_capacity, top_k)
+        output = binned_gather(x, indices, bins, expert_capacity, top_k)
+        assert output is not None
+        x = output
+
+        # Perform the expert computation. Note that we don't
+        # use biases for these linear operations.
+        x = self.mlp(x)
+
+        # Un-route the data for the MoE output.
+        # return ops.binned_scatter(x, indices, expert_weights, bins, top_k)
+        return binned_scatter(x, indices, expert_weights, bins, top_k)
+    
+
+    def forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor):
+        # x: [sl, bs, hs]
+        # expert_weights: [sl * bs, top-k]
+        # top_experts: [sl * bs, top-k]
+        expert_weights = expert_weights.flatten()
+        top_experts = top_experts.flatten()
+        with torch.no_grad():
+            indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts))
+
+            # If expert_capacity is set to zero, set the number of tokens
+            # per expert to the maximum we need to avoid dropping tokens.
+            sl, bs, _ = x.size()
+            expert_capacity = self.expert_capacity(sl * bs)
+            if expert_capacity == 0:
+                expert_capacity = torch.max(tokens_per_expert).item()
+
+        x = self.permute_and_compute(
+            x,
+            tokens_per_expert,
+            indices,
+            bin_ids,
+            expert_weights,
+            bins,
+            expert_capacity,
+            self.top_k,
+        )
+        return x, tokens_per_expert
+
+    def parallel_forward_once(self, x: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor):
+        # NOTE: This function implements the same computation as forward_once
+        # but with expert model parallelism.
+        #
+        # 1. Permute the tokens locally so that they are grouped by their
+        # expert assignments. This allows us to transfer all of the tokens
+        # for a remote device in one communication primitive.
+        #
+        # 2. Permute the tokens across the expert parallel devices. After
+        # this is completed each device has all of the tokens assigned to
+        # its set of experts in its local HBM.
+        #
+        # 3. Permute the tokens locally so that they are grouped by their
+        # expert assignement. After the distributed permutation the tokens
+        # are grouped by which device they came from. We re-order them
+        # locally to allow for efficient computation.
+        #
+        # After this series of permutations we compute the linear layers
+        # and then repeat these three steps in reverse to produce the final
+        # output.
+        #
+        # Compute the mapping of local tokens to experts.
+        expert_weights = expert_weights.flatten()
+        top_experts = top_experts.flatten()
+        with torch.no_grad():
+            indices, bin_ids, bins, tokens_per_expert = (self.indices_and_bins(top_experts))
+
+            # If we're sharding the experts along the hidden dimension
+            # multiple devices own parts of the same sets of experts.
+            # Replicate the token counts so every device gets the counts.
+            # repeated_tokens_per_expert = ops.repeat(
+            repeated_tokens_per_expert = repeat(
+                tokens_per_expert,
+                (mpu.hidden_sharding_degree(self.args),),
+            )
+
+            # Pass token count information to the device on which the
+            # target expert resides.
+            parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert,)
+            tpe_handle = dist.all_to_all_single(
+                parallel_tokens_per_expert,
+                repeated_tokens_per_expert,
+                group=self.args.expert_parallel_group,
+                async_op=True,
+            )
+
+        # Permute locally and without any padding so that tokens for each
+        # parallel device are stored contiguously.
+        #
+        # This view updates the shape of the tensor from [sl, bs, hs] to
+        # [sl * bs, hs] prior to the permutation.
+        x = x.view(-1, x.shape[-1])
+        # output = ops.gather(x, indices, bin_ids, bins, self.top_k)
+        output = gather(x, indices, bin_ids, bins, self.top_k)
+        assert output is not None
+        x = output
+
+        # Compute the number of tokens that will be received from each
+        # device and permute the input data across the devices.
+        with torch.no_grad():
+            tpe_handle.wait()
+            experts_per_rank = mpu.experts_per_rank(self.args)
+
+            # Reshape to [world_size, num_experts_per_rank].
+            world_size = mpu.get_expert_parallel_world_size(self.args)
+            repeated_tokens_per_expert = (repeated_tokens_per_expert.view(world_size, experts_per_rank))
+            parallel_tokens_per_expert = (parallel_tokens_per_expert.view(world_size, experts_per_rank))
+
+            # TODO(tgale): It might be faster to do this on the GPU and
+            # then communicate the results back to the host.
+            send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1)
+            parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu()
+            recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1)
+
+            # Convert the send/recv counts to lists.
+            send_counts = send_counts.tolist()
+            recv_counts = recv_counts.tolist()
+            tokens_received = sum(recv_counts)
+
+        # If we're sharding the experts along the hidden dimension
+        # multiple devices own parts of the same sets of experts.
+        # Replicate the token counts so devices that share experts
+        # get all of the tokens assigned to them.
+        #
+        # TODO(tgale): Fuse this into the prior, local permutation.
+        # x = ops.repeat(x, (mpu.hidden_sharding_degree(self.args), 1))
+        x = repeat(x, (mpu.hidden_sharding_degree(self.args), 1))
+
+        # Start the cross-device permutation asynchronously so we can
+        # overlap communication with computation.
+        parallel_x, parallel_x_handle = all_to_all(
+            x,
+            recv_counts,
+            send_counts,
+            self.args.expert_parallel_group,
+            async_op=True,
+        )
+
+        with torch.no_grad():
+            # After we do the cross-device permutation we have the tokens on the
+            # correct device but not yet grouped by expert because we received
+            # tokens from each device as contiguous chunks. To group the tokens
+            # for expert computation we'll do one more local permutation. The
+            # rest of this torch.no_grad() scope sets up the indices and bins
+            # for this permutation.
+            # replicate_bins = ops.inclusive_cumsum(
+            replicate_bins = inclusive_cumsum(
+                parallel_tokens_per_expert.flatten(),
+                0,
+            )
+            replicate_bins = (replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins)
+
+            # Construct the expert indices for the permuted tokens.
+            parallel_top_expert = torch.remainder(
+                torch.arange(
+                    self.num_experts * mpu.hidden_sharding_degree(self.args),
+                    dtype=torch.int32,
+                    device=indices.device,
+                ),
+                mpu.experts_per_rank(self.args),
+            )
+            # parallel_top_expert = ops.replicate(
+            parallel_top_expert = replicate(
+                parallel_top_expert.unsqueeze(dim=0),
+                replicate_bins,
+                tokens_received,
+            ).flatten()
+
+            # TODO(tgale): The sort_end_bit here can be reduced.
+            # parallel_bin_ids, parallel_indices = ops.sort(
+            parallel_bin_ids, parallel_indices = sort(
+                parallel_top_expert,
+                self.sort_end_bit,
+            )
+
+            # Calculate the bins boundaries from the token counts.
+            parallel_tokens_per_expert = parallel_tokens_per_expert.sum(
+                dim=0,
+                dtype=torch.int,
+            )
+            # parallel_bins = ops.inclusive_cumsum(parallel_tokens_per_expert, 0)
+            parallel_bins = inclusive_cumsum(parallel_tokens_per_expert, 0)
+            parallel_bins = (parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins)
+
+            # If expert_capacity is set to zero, set the number of tokens
+            # per expert to the maximum we need to avoid dropping tokens.
+            tokens, _ = x.size()
+            expert_capacity = self.expert_capacity(tokens)
+            if expert_capacity == 0:
+                expert_capacity = torch.max(parallel_tokens_per_expert).item()
+
+        # Locally permute the tokens and perform the expert computation.
+        # Block to make sure that the cross-device permutation is complete.
+        if self.args.mlp_impl == 'grouped':
+            # GroupedMLP requires counts on CPU. We can use the tensor already
+            # moved to CPU for the prior all_to_all, which avoids an extra
+            # device synchronization.
+            parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum(
+                dim=0,
+                dtype=torch.int,
+            )
+        parallel_x_handle.wait()
+        parallel_x = self.permute_and_compute(
+            parallel_x,
+            parallel_tokens_per_expert,
+            parallel_indices,
+            parallel_bin_ids,
+            None,  # expert_weights
+            parallel_bins,
+            expert_capacity,
+            top_k=1,
+        )
+
+        # Un-permute the tokens across the devices.
+        x, _ = all_to_all(
+            parallel_x,
+            send_counts,
+            recv_counts,
+            self.args.expert_parallel_group,
+        )
+
+        # Reduce along the hidden sharding to get the final outputs.
+        #
+        # TODO(tgale): Fuse this into the following local permutation.
+        shape = (
+            mpu.hidden_sharding_degree(self.args),
+            -1,
+            self.args.hidden_size,
+        )
+        # x = ops.sum(x.view(shape), dim=0)
+        x = x.view(shape).sum(dim=0)
+
+        # Un-permute locally to setup for the next series of operations.
+        # x = ops.scatter(x, indices, bin_ids, expert_weights, bins, self.top_k)
+        x = scatter(x, indices, bin_ids, expert_weights, bins, self.top_k)
+        return x, tokens_per_expert.flatten()
+
+    def forward(self, x: torch.Tensor, scores: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor):
+        in_shape = x.size()
+
+        # Compute the experts.
+        x, tokens_per_expert = self.forward_fn(x, expert_weights, top_experts)
+        if self.training and self.args.moe_loss_weight > 0:
+            save_load_balancing_loss((tokens_per_expert, scores))
+        x = x.view(in_shape)
+        if self.bias is not None:
+            if self.args.return_bias:
+                return x, self.bias
+            return x + self.bias
+        return x
+
+
+class MoE(torch.nn.Module):
+
+    def __init__(self, args: Arguments):
+        super(MoE, self).__init__()
+
+        # Token router.
+        self.router = router.LearnedRouter(args)
+
+        # Expert computation helper.
+        self.experts = self._init_experts_mlp(args)
+
+        self.shared_expert = None
+        if args.shared_expert:
+            # SharedExpert computation helper.
+            self.shared_expert = sharedexpert_registry.get(args)
+
+    def _init_experts_mlp(self, args: Arguments):
+        return ParallelMLP(args)
+
+    def forward(self, x: torch.Tensor):
+        # NOTE: If we're going to cast the activations to lower precision
+        # do it before we permute the tokens to save bandwidth.
+        x = common.cast_if_autocast_enabled(x)
+
+        # Compute the expert scores and assignments.
+        scores, expert_weights, top_experts = self.router(x)
+
+        # Compute the experts.
+        out = self.experts(x, scores, expert_weights, top_experts)
+        if self.shared_expert is not None:
+            shared_expert_out = self.shared_expert(x)
+            out = self.shared_expert.add_experts_sharedexpert(
+                shared_expert_out,
+                out,
+            )
+        return out

build/torch211-cxx11-cu126-aarch64-linux/_layers/mpu.py

@@ -0,0 +1,94 @@
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+
+# from megablocks.layers.arguments import Arguments
+from .arguments import Arguments
+
+
+class MoeParam(torch.Tensor):
+
+    def __init__(self):
+        super().__init__(self)
+        self.expert_model_parallel: bool
+
+
+def is_moe_param(tensor: torch.Tensor) -> bool:
+    return hasattr(tensor, 'expert_model_parallel')
+
+
+def get_expert_parallel_world_size(args: Arguments) -> int:
+    return (dist.get_world_size(args.expert_parallel_group) if args.moe_expert_model_parallelism else 1)
+
+
+def get_expert_parallel_rank(args: Arguments) -> int:
+    return (dist.get_rank(args.expert_parallel_group) if args.moe_expert_model_parallelism else 0)
+
+
+def set_expert_model_parallel_attributes(
+    tensor: torch.Tensor,
+    is_parallel: bool,
+):
+    assert not hasattr(tensor, 'expert_model_parallel')
+    setattr(tensor, 'expert_model_parallel', is_parallel)
+
+
+def param_is_expert_model_parallel(param: MoeParam) -> bool:
+    return (hasattr(param, 'expert_model_parallel') and param.expert_model_parallel)
+
+
+def copy_expert_model_parallel_attributes(
+    destination_tensor: torch.Tensor,
+    source_tensor: torch.Tensor,
+):
+    if hasattr(source_tensor, 'expert_model_parallel'):
+        setattr(
+            destination_tensor,
+            'expert_model_parallel',
+            getattr(source_tensor, 'expert_model_parallel'),
+        )
+
+
+def synchronized_print(group: Optional[dist.ProcessGroup], *x: torch.Tensor):
+    world_size = dist.get_world_size(group)
+    rank = dist.get_rank(group)
+    for i in range(world_size):
+        dist.barrier(group)
+        if i == rank:
+            print(f'rank = {rank}', *x)
+
+
+# Helpers for expert/tensor sharding.
+def expert_sharding_degree(args: Arguments) -> int:
+    world_size = get_expert_parallel_world_size(args)
+    esd = min(world_size, args.moe_num_experts)
+
+    if (args.moe_num_experts % esd) != 0:
+        raise ValueError(f'Cannot shard {args.moe_num_experts} experts {esd} ways.',)
+    return esd
+
+
+def hidden_sharding_degree(args: Arguments) -> int:
+    world_size = get_expert_parallel_world_size(args)
+    esd = expert_sharding_degree(args)
+    hsd = world_size // esd
+
+    if (args.ffn_hidden_size % hsd) != 0:
+        raise ValueError(f'Cannot shard {args.ffn_hidden_size} features {hsd} ways.',)
+    if (esd * hsd) != world_size:
+        raise ValueError(
+            f"Invalid sharding. 'expert_sharding_degree' ({esd}) * hidden_sharding_degree ({hsd}) != world_size ({world_size}).",
+        )
+    return hsd
+
+
+def experts_per_rank(args: Arguments) -> int:
+    return args.moe_num_experts // expert_sharding_degree(args)
+
+
+def features_per_rank(args: Arguments) -> int:
+    return args.ffn_hidden_size // hidden_sharding_degree(args)