diff --git a/benchmarks/ops/attention/bench_deepseek_dsa_decode.py b/benchmarks/ops/attention/bench_deepseek_dsa_decode.py index a8244c98..f3e56b22 100644 --- a/benchmarks/ops/attention/bench_deepseek_dsa_decode.py +++ b/benchmarks/ops/attention/bench_deepseek_dsa_decode.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import DeepSeekSparseAttentionDecodeWithKVCacheFwdOp -from workloads.ops.deepseek_dsa_decode import DsaDecodeTest +from workloads.attention.deepseek_dsa_decode import DsaDecodeTest class _DsaDecodeTestBaseline(DsaDecodeTest): diff --git a/benchmarks/ops/attention/bench_deepseek_mla_decode.py b/benchmarks/ops/attention/bench_deepseek_mla_decode.py index 9bdc907b..9c28c8c5 100644 --- a/benchmarks/ops/attention/bench_deepseek_mla_decode.py +++ b/benchmarks/ops/attention/bench_deepseek_mla_decode.py @@ -7,7 +7,7 @@ from einops import einsum, rearrange from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MultiHeadLatentAttentionDecodeWithKVCacheFwdOp -from workloads.ops.deepseek_mla_decode import MlaDecodeTest +from workloads.attention.deepseek_mla_decode import MlaDecodeTest class _MlaDecodeTestBaseline(MlaDecodeTest): diff --git a/benchmarks/ops/attention/bench_deepseek_nsa.py b/benchmarks/ops/attention/bench_deepseek_nsa.py index bb5f692d..78a43173 100644 --- a/benchmarks/ops/attention/bench_deepseek_nsa.py +++ b/benchmarks/ops/attention/bench_deepseek_nsa.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import NSAFwdVarlenOp -from workloads.ops.deepseek_nsa_fwd import NsaFwdTest +from workloads.attention.deepseek_nsa import NsaFwdTest class _NsaFwdTestBaseline(NsaFwdTest): diff --git a/benchmarks/ops/attention/bench_deepseek_nsa_cmp.py b/benchmarks/ops/attention/bench_deepseek_nsa_cmp.py index 3a0beea0..ec782e4f 100644 --- a/benchmarks/ops/attention/bench_deepseek_nsa_cmp.py +++ b/benchmarks/ops/attention/bench_deepseek_nsa_cmp.py @@ -5,8 +5,8 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import NSACmpFwdVarlenOp +from workloads.attention.deepseek_nsa_cmp import NsaCmpFwdTest from workloads.nsa_utils import prepare_chunk_offsets -from workloads.ops.deepseek_nsa_cmp_fwd import NsaCmpFwdTest def _parallel_nsa_compression_fwd_pytorch(test, q, k_cmp, v_cmp, block_size, scale, offsets): diff --git a/benchmarks/ops/attention/bench_deepseek_nsa_topk.py b/benchmarks/ops/attention/bench_deepseek_nsa_topk.py index c859e08b..c46d4aba 100644 --- a/benchmarks/ops/attention/bench_deepseek_nsa_topk.py +++ b/benchmarks/ops/attention/bench_deepseek_nsa_topk.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import NSATopkVarlenOp -from workloads.ops.deepseek_nsa_topk import NsaTopkTest +from workloads.attention.deepseek_nsa_topk import NsaTopkTest def _nsa_topk_torch(test, q, k_cmp, lse, block_counts, block_size, scale, diff --git a/benchmarks/ops/attention/bench_gqa.py b/benchmarks/ops/attention/bench_gqa.py index c412e194..0f9e116c 100644 --- a/benchmarks/ops/attention/bench_gqa.py +++ b/benchmarks/ops/attention/bench_gqa.py @@ -6,7 +6,7 @@ from torch.nn import functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GroupedQueryAttentionBwdOp, GroupedQueryAttentionFwdOp -from workloads.ops.gqa import ( +from workloads.attention.gqa import ( GqaBwdTest, GqaFwdTest, ) diff --git a/benchmarks/ops/attention/bench_gqa_decode.py b/benchmarks/ops/attention/bench_gqa_decode.py index 532cc8fa..cb4aae5c 100644 --- a/benchmarks/ops/attention/bench_gqa_decode.py +++ b/benchmarks/ops/attention/bench_gqa_decode.py @@ -7,7 +7,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GroupedQueryAttentionDecodeWithKVCacheFwdOp -from workloads.ops.gqa_decode import GqaDecodeTest +from workloads.attention.gqa_decode import GqaDecodeTest class _GqaDecodeTestBaseline(GqaDecodeTest): diff --git a/benchmarks/ops/attention/bench_gqa_decode_paged.py b/benchmarks/ops/attention/bench_gqa_decode_paged.py index c672a3f2..6c376900 100644 --- a/benchmarks/ops/attention/bench_gqa_decode_paged.py +++ b/benchmarks/ops/attention/bench_gqa_decode_paged.py @@ -8,7 +8,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GroupedQueryAttentionDecodePagedWithKVCacheFwdOp -from workloads.ops.gqa_decode_paged import GqaDecodePagedTest +from workloads.attention.gqa_decode_paged import GqaDecodePagedTest class _GqaDecodePagedTestBaseline(GqaDecodePagedTest): diff --git a/benchmarks/ops/attention/bench_gqa_sliding_window.py b/benchmarks/ops/attention/bench_gqa_sliding_window.py index d77493c5..e2a5283a 100644 --- a/benchmarks/ops/attention/bench_gqa_sliding_window.py +++ b/benchmarks/ops/attention/bench_gqa_sliding_window.py @@ -7,7 +7,7 @@ from torch.nn import functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GqaSlidingWindowFwdOp -from workloads.ops.gqa_sliding_window_fwd import GqaSlidingWindowFwdTest +from workloads.attention.gqa_sliding_window import GqaSlidingWindowFwdTest class GqaSlidingWindowFwdBenchmark(BenchmarkBase): diff --git a/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py b/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py index 8582cf9d..2115ee2b 100644 --- a/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py +++ b/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py @@ -7,7 +7,7 @@ from torch.nn import functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GqaSlidingWindowVarlenFwdOp -from workloads.ops.gqa_sliding_window_varlen_fwd import GqaSlidingWindowVarlenFwdTest +from workloads.attention.gqa_sliding_window_varlen import GqaSlidingWindowVarlenFwdTest _GQA_SLIDING_WINDOW_VARLEN_FWD_BENCH_PARAMS = [ pytest.param(1, [3000], [3000], 32, 8, 128, True, -1, -1, torch.float16, False, id="single-seq-causal"), diff --git a/benchmarks/ops/attention/bench_mean_pooling.py b/benchmarks/ops/attention/bench_mean_pooling.py index 3a39b9d6..15e99f15 100644 --- a/benchmarks/ops/attention/bench_mean_pooling.py +++ b/benchmarks/ops/attention/bench_mean_pooling.py @@ -5,8 +5,8 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MeanPoolingForwardOp +from workloads.attention.mean_pooling import MeanPoolingTest from workloads.nsa_utils import prepare_chunk_indices -from workloads.ops.mean_pooling_ops import MeanPoolingTest class _MeanPoolingTestBaseline(MeanPoolingTest): diff --git a/benchmarks/ops/attention/bench_mha.py b/benchmarks/ops/attention/bench_mha.py index 6348d644..1b405c03 100644 --- a/benchmarks/ops/attention/bench_mha.py +++ b/benchmarks/ops/attention/bench_mha.py @@ -6,7 +6,7 @@ from torch.nn import functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MultiHeadAttentionBwdOp, MultiHeadAttentionFwdOp -from workloads.ops.mha import ( +from workloads.attention.mha import ( MhaBwdTest, MhaFwdTest, ) diff --git a/benchmarks/ops/attention/bench_mha_decode.py b/benchmarks/ops/attention/bench_mha_decode.py index e4432155..70bb86ea 100644 --- a/benchmarks/ops/attention/bench_mha_decode.py +++ b/benchmarks/ops/attention/bench_mha_decode.py @@ -7,7 +7,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MultiHeadAttentionDecodeWithKVCacheFwdOp -from workloads.ops.mha_decode import MhaDecodeTest +from workloads.attention.mha_decode import MhaDecodeTest class _MhaDecodeTestBaseline(MhaDecodeTest): diff --git a/benchmarks/ops/attention/bench_mha_decode_paged.py b/benchmarks/ops/attention/bench_mha_decode_paged.py index 31ace9f3..abef1233 100644 --- a/benchmarks/ops/attention/bench_mha_decode_paged.py +++ b/benchmarks/ops/attention/bench_mha_decode_paged.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MultiHeadAttentionDecodePagedWithKVCacheFwdOp -from workloads.ops.mha_decode_paged import MhaDecodePagedTest +from workloads.attention.mha_decode_paged import MhaDecodePagedTest class _MhaDecodePagedTestBaseline(MhaDecodePagedTest): diff --git a/benchmarks/ops/bench_activation.py b/benchmarks/ops/bench_activation.py index 369d8bbd..c27fe8ba 100644 --- a/benchmarks/ops/bench_activation.py +++ b/benchmarks/ops/bench_activation.py @@ -26,8 +26,8 @@ from tileops.kernels.elementwise import ( _make_unary_explicit, ) from tileops.ops.elementwise import ErfOp, GeluOp, MishOp, ReluOp +from workloads.activation import ReluTest from workloads.base import FixtureBase -from workloads.ops.activation import ReluTest # --------------------------------------------------------------------------- # LLM-realistic shapes (LLaMA-family defaults) diff --git a/benchmarks/ops/bench_ada_layer_norm.py b/benchmarks/ops/bench_ada_layer_norm.py index a35d5bed..9bf1cb3f 100644 --- a/benchmarks/ops/bench_ada_layer_norm.py +++ b/benchmarks/ops/bench_ada_layer_norm.py @@ -8,8 +8,8 @@ from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.ada_layer_norm import AdaLayerNormFwdOp from tileops.ops.norm.ada_layer_norm_zero import AdaLayerNormZeroFwdOp -from workloads.ops.ada_layer_norm import AdaLayerNormTest -from workloads.ops.ada_layer_norm_zero import AdaLayerNormZeroTest +from workloads.ada_layer_norm import AdaLayerNormTest +from workloads.ada_layer_norm_zero import AdaLayerNormZeroTest _ADA_OP_NAME = "AdaLayerNormFwdOp" _ADA_ZERO_OP_NAME = "AdaLayerNormZeroFwdOp" diff --git a/benchmarks/ops/bench_argreduce.py b/benchmarks/ops/bench_argreduce.py index f2f92c0a..a645f564 100644 --- a/benchmarks/ops/bench_argreduce.py +++ b/benchmarks/ops/bench_argreduce.py @@ -10,7 +10,7 @@ import torch from benchmarks.benchmark import BenchmarkReport, ManifestBenchmark, workloads_to_params from tileops.ops.reduction.argmax import ArgmaxFwdOp from tileops.ops.reduction.argmin import ArgminFwdOp -from workloads.ops.argreduce import ArgmaxTest, ArgminTest +from workloads.argreduce import ArgmaxTest, ArgminTest _ARGMAX_OP = "ArgmaxFwdOp" _ARGMIN_OP = "ArgminFwdOp" diff --git a/benchmarks/ops/bench_batch_norm.py b/benchmarks/ops/bench_batch_norm.py index 5c324d27..6c5af245 100644 --- a/benchmarks/ops/bench_batch_norm.py +++ b/benchmarks/ops/bench_batch_norm.py @@ -15,7 +15,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.batch_norm import BatchNormBwdOp, BatchNormFwdOp -from workloads.ops.batch_norm import BatchNormBwdTest, BatchNormFwdTest +from workloads.batch_norm import BatchNormBwdTest, BatchNormFwdTest _FWD_OP_NAME = "BatchNormFwdOp" _BWD_OP_NAME = "BatchNormBwdOp" diff --git a/benchmarks/ops/bench_binary_arith.py b/benchmarks/ops/bench_binary_arith.py index 64cbf8d3..a3fb868b 100644 --- a/benchmarks/ops/bench_binary_arith.py +++ b/benchmarks/ops/bench_binary_arith.py @@ -20,7 +20,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops.elementwise import AddOp, WhereOp from workloads.base import FixtureBase -from workloads.ops.binary_arith import AddSameShapeTest +from workloads.binary_arith import AddSameShapeTest # --------------------------------------------------------------------------- # LLM-realistic shapes (LLaMA-family defaults) diff --git a/benchmarks/ops/bench_deltanet.py b/benchmarks/ops/bench_deltanet.py index 55f9cd60..15c5dc59 100644 --- a/benchmarks/ops/bench_deltanet.py +++ b/benchmarks/ops/bench_deltanet.py @@ -21,7 +21,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import DeltaNetBwdOp, DeltaNetFwdOp, DeltaNetOp from workloads.base import FixtureBase -from workloads.ops.deltanet_fwd import DeltaNetFwdTest +from workloads.deltanet import DeltaNetFwdTest def _differentiable_fwd(q, k, v, beta, chunk_size): diff --git a/benchmarks/ops/bench_deltanet_recurrence.py b/benchmarks/ops/bench_deltanet_recurrence.py index f88dcab8..ae7ef57b 100644 --- a/benchmarks/ops/bench_deltanet_recurrence.py +++ b/benchmarks/ops/bench_deltanet_recurrence.py @@ -6,7 +6,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import DeltaNetDecodeOp from workloads.base import FixtureBase -from workloads.ops.deltanet_recurrence import DeltaNetDecodeTest +from workloads.deltanet import DeltaNetDecodeTest def deltanet_decode_torch( diff --git a/benchmarks/ops/bench_engram.py b/benchmarks/ops/bench_engram.py index 38acc0b5..38f066a7 100644 --- a/benchmarks/ops/bench_engram.py +++ b/benchmarks/ops/bench_engram.py @@ -7,9 +7,12 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops.engram import EngramGateConvBwdOp, EngramGateConvFwdOp from tileops.ops.engram_decode import EngramDecodeOp -from workloads.ops.engram_bwd import EngramGateConvBwdTest -from workloads.ops.engram_decode import EngramDecodeTest -from workloads.ops.engram_fwd import CONV_KERNEL_SIZE, EngramGateConvFwdTest +from workloads.engram import ( + CONV_KERNEL_SIZE, + EngramDecodeTest, + EngramGateConvBwdTest, + EngramGateConvFwdTest, +) def _rmsnorm(x, w, eps=1e-6): diff --git a/benchmarks/ops/bench_fft.py b/benchmarks/ops/bench_fft.py index 5eef7ca0..b6424ea0 100644 --- a/benchmarks/ops/bench_fft.py +++ b/benchmarks/ops/bench_fft.py @@ -7,7 +7,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import FFTC2COp from workloads.base import FixtureBase -from workloads.ops.fft import FFTTest +from workloads.fft import FFTTest class _FFTTestBaseline(FFTTest): diff --git a/benchmarks/ops/bench_fp8_lighting_indexer.py b/benchmarks/ops/bench_fp8_lighting_indexer.py index deb0567c..1cf4d1c4 100644 --- a/benchmarks/ops/bench_fp8_lighting_indexer.py +++ b/benchmarks/ops/bench_fp8_lighting_indexer.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import FP8LightingIndexerOp -from workloads.ops.fp8_lighting_indexer import FP8LightingIndexerTest +from workloads.fp8_lighting_indexer import FP8LightingIndexerTest class _FP8LightingIndexerTestBaseline(FP8LightingIndexerTest): diff --git a/benchmarks/ops/bench_fp8_quant.py b/benchmarks/ops/bench_fp8_quant.py index 58340aaf..d3f38214 100644 --- a/benchmarks/ops/bench_fp8_quant.py +++ b/benchmarks/ops/bench_fp8_quant.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import FP8QuantOp -from workloads.ops.fp8_quant import FP8QuantTest +from workloads.fp8_quant import FP8QuantTest class _FP8QuantTestBaseline(FP8QuantTest): diff --git a/benchmarks/ops/bench_fused_add_layer_norm.py b/benchmarks/ops/bench_fused_add_layer_norm.py index 1b1e419b..b49dbcce 100644 --- a/benchmarks/ops/bench_fused_add_layer_norm.py +++ b/benchmarks/ops/bench_fused_add_layer_norm.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.fused_add_layer_norm import FusedAddLayerNormFwdOp -from workloads.ops.fused_add_layer_norm import FusedAddLayerNormTest +from workloads.fused_add_layer_norm import FusedAddLayerNormTest _OP_NAME = "FusedAddLayerNormFwdOp" diff --git a/benchmarks/ops/bench_fused_add_rmsnorm.py b/benchmarks/ops/bench_fused_add_rmsnorm.py index bd6e58c9..3a11fb0c 100644 --- a/benchmarks/ops/bench_fused_add_rmsnorm.py +++ b/benchmarks/ops/bench_fused_add_rmsnorm.py @@ -6,7 +6,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.fused_add_rmsnorm import FusedAddRMSNormFwdOp -from workloads.ops.fused_add_rmsnorm import FusedAddRmsNormTest +from workloads.fused_add_rmsnorm import FusedAddRmsNormTest _OP_NAME = "FusedAddRMSNormFwdOp" diff --git a/benchmarks/ops/bench_gated_deltanet.py b/benchmarks/ops/bench_gated_deltanet.py index b7dafb4f..33ac5ce5 100644 --- a/benchmarks/ops/bench_gated_deltanet.py +++ b/benchmarks/ops/bench_gated_deltanet.py @@ -21,7 +21,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GatedDeltaNetBwdOp, GatedDeltaNetFwdOp, GatedDeltaNetOp from workloads.base import FixtureBase -from workloads.ops.gated_deltanet_fwd import GatedDeltaNetFwdTest +from workloads.gated_deltanet import GatedDeltaNetFwdTest def _differentiable_fwd(q, k, v, g_raw, beta, chunk_size): diff --git a/benchmarks/ops/bench_gated_deltanet_recurrence.py b/benchmarks/ops/bench_gated_deltanet_recurrence.py index bd2bb3a6..93033087 100644 --- a/benchmarks/ops/bench_gated_deltanet_recurrence.py +++ b/benchmarks/ops/bench_gated_deltanet_recurrence.py @@ -6,7 +6,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GatedDeltaNetDecodeOp from workloads.base import FixtureBase -from workloads.ops.gated_deltanet_recurrence import GatedDeltaNetDecodeTest +from workloads.gated_deltanet import GatedDeltaNetDecodeTest def gated_deltanet_decode_torch( diff --git a/benchmarks/ops/bench_gemm.py b/benchmarks/ops/bench_gemm.py index 248323cd..9a9f1271 100644 --- a/benchmarks/ops/bench_gemm.py +++ b/benchmarks/ops/bench_gemm.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GemmOp -from workloads.ops.gemm import GemmTest +from workloads.gemm import GemmTest class _GemmTestBaseline(GemmTest): diff --git a/benchmarks/ops/bench_gla_recurrence.py b/benchmarks/ops/bench_gla_recurrence.py index 385238f4..262e6932 100644 --- a/benchmarks/ops/bench_gla_recurrence.py +++ b/benchmarks/ops/bench_gla_recurrence.py @@ -13,7 +13,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GLADecodeOp from workloads.base import FixtureBase -from workloads.ops.gla_recurrence import GLADecodeTest +from workloads.gla import GLADecodeTest def gla_decode_torch( diff --git a/benchmarks/ops/bench_group_norm.py b/benchmarks/ops/bench_group_norm.py index 0eab3b22..13f0f754 100644 --- a/benchmarks/ops/bench_group_norm.py +++ b/benchmarks/ops/bench_group_norm.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.group_norm import GroupNormFwdOp -from workloads.ops.group_norm import GroupNormTest +from workloads.group_norm import GroupNormTest _OP_NAME = "GroupNormFwdOp" diff --git a/benchmarks/ops/bench_grouped_gemm.py b/benchmarks/ops/bench_grouped_gemm.py index 3bda799a..d99185dc 100644 --- a/benchmarks/ops/bench_grouped_gemm.py +++ b/benchmarks/ops/bench_grouped_gemm.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import GroupedGemmOp -from workloads.ops.grouped_gemm import ( +from workloads.grouped_gemm import ( GroupedGemmCompleteTest, GroupedGemmTest, ) diff --git a/benchmarks/ops/bench_instance_norm.py b/benchmarks/ops/bench_instance_norm.py index 9edbc35c..e0e00b9e 100644 --- a/benchmarks/ops/bench_instance_norm.py +++ b/benchmarks/ops/bench_instance_norm.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.instance_norm import InstanceNormFwdOp -from workloads.ops.instance_norm import InstanceNormTest +from workloads.instance_norm import InstanceNormTest _OP_NAME = "InstanceNormFwdOp" diff --git a/benchmarks/ops/bench_layer_norm.py b/benchmarks/ops/bench_layer_norm.py index 30eb5678..2f256d53 100644 --- a/benchmarks/ops/bench_layer_norm.py +++ b/benchmarks/ops/bench_layer_norm.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.layer_norm import LayerNormFwdOp -from workloads.ops.layer_norm import LayerNormTest +from workloads.layer_norm import LayerNormTest _OP_NAME = "LayerNormFwdOp" diff --git a/benchmarks/ops/bench_logical_reduce.py b/benchmarks/ops/bench_logical_reduce.py index 5192f385..cfd29fc4 100644 --- a/benchmarks/ops/bench_logical_reduce.py +++ b/benchmarks/ops/bench_logical_reduce.py @@ -11,7 +11,7 @@ from benchmarks.benchmark import BenchmarkReport, ManifestBenchmark, workloads_t from tileops.ops.reduction.all_op import AllFwdOp from tileops.ops.reduction.any_op import AnyFwdOp from tileops.ops.reduction.count_nonzero import CountNonzeroFwdOp -from workloads.ops.logical_reduce import AllTest, AnyTest, CountNonzeroTest +from workloads.logical_reduce import AllTest, AnyTest, CountNonzeroTest # =================================================================== # Op name constants diff --git a/benchmarks/ops/bench_mamba.py b/benchmarks/ops/bench_mamba.py index 9d0b0eb9..31a74034 100644 --- a/benchmarks/ops/bench_mamba.py +++ b/benchmarks/ops/bench_mamba.py @@ -9,11 +9,18 @@ from tileops.ops.ssd_chunk_scan import SsdChunkScanFwdOp from tileops.ops.ssd_chunk_state import SsdChunkStateFwdOp from tileops.ops.ssd_decode import SsdDecodeOp from tileops.ops.ssd_state_passing import SsdStatePassingFwdOp -from workloads.ops.da_cumsum_fwd import DaCumsumFwdFixture, DaCumsumFwdTest -from workloads.ops.ssd_chunk_scan_fwd import SsdChunkScanFwdFixture, SsdChunkScanFwdTest -from workloads.ops.ssd_chunk_state_fwd import SsdChunkStateFwdFixture, SsdChunkStateFwdTest -from workloads.ops.ssd_decode import SsdDecodeFixture, SsdDecodeTest -from workloads.ops.ssd_state_passing_fwd import SsdStatePassingFwdFixture, SsdStatePassingFwdTest +from workloads.mamba import ( + DaCumsumFwdFixture, + DaCumsumFwdTest, + SsdChunkScanFwdFixture, + SsdChunkScanFwdTest, + SsdChunkStateFwdFixture, + SsdChunkStateFwdTest, + SsdDecodeFixture, + SsdDecodeTest, + SsdStatePassingFwdFixture, + SsdStatePassingFwdTest, +) def da_cumsum_fwd_ref( diff --git a/benchmarks/ops/bench_mhc_post.py b/benchmarks/ops/bench_mhc_post.py index 0329bd76..ecc86b95 100644 --- a/benchmarks/ops/bench_mhc_post.py +++ b/benchmarks/ops/bench_mhc_post.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MHCPostOp -from workloads.ops.mhc_post import MHCPostTest +from workloads.mhc import MHCPostTest class _MHCPostTestBaseline(MHCPostTest): diff --git a/benchmarks/ops/bench_mhc_pre.py b/benchmarks/ops/bench_mhc_pre.py index f0e77eaa..2c5d9e87 100644 --- a/benchmarks/ops/bench_mhc_pre.py +++ b/benchmarks/ops/bench_mhc_pre.py @@ -6,7 +6,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import MHCPreOp -from workloads.ops.mhc_pre import MHCPreTest +from workloads.mhc import MHCPreTest class _MHCPreTestBaseline(MHCPreTest): diff --git a/benchmarks/ops/bench_moe_fused_topk.py b/benchmarks/ops/bench_moe_fused_topk.py index f96a35da..22f0f615 100644 --- a/benchmarks/ops/bench_moe_fused_topk.py +++ b/benchmarks/ops/bench_moe_fused_topk.py @@ -29,7 +29,7 @@ except ImportError: from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops.moe import FusedTopKOp from workloads.base import FixtureBase -from workloads.ops.moe_fused_topk import FusedTopKTest +from workloads.moe import FusedTopKTest def fused_topk_torch( diff --git a/benchmarks/ops/bench_moe_permute.py b/benchmarks/ops/bench_moe_permute.py index 6af35139..fc02b651 100644 --- a/benchmarks/ops/bench_moe_permute.py +++ b/benchmarks/ops/bench_moe_permute.py @@ -32,7 +32,7 @@ except ImportError: from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.moe import MoePermutePaddedFwdOp -from workloads.ops.moe_permute import MoePermuteTest +from workloads.moe import MoePermuteTest _OP_NAME = "MoePermutePaddedFwdOp" diff --git a/benchmarks/ops/bench_moe_permute_align.py b/benchmarks/ops/bench_moe_permute_align.py index 9b1526ad..e71dbd39 100644 --- a/benchmarks/ops/bench_moe_permute_align.py +++ b/benchmarks/ops/bench_moe_permute_align.py @@ -28,7 +28,7 @@ except ImportError: from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.moe import MoePermuteAlignFwdOp -from workloads.ops.moe_permute_align import MoePermuteAlignTest +from workloads.moe import MoePermuteAlignTest _OP_NAME = "MoePermuteAlignFwdOp" diff --git a/benchmarks/ops/bench_moe_unpermute.py b/benchmarks/ops/bench_moe_unpermute.py index f26349da..a70e925d 100644 --- a/benchmarks/ops/bench_moe_unpermute.py +++ b/benchmarks/ops/bench_moe_unpermute.py @@ -33,7 +33,7 @@ except ImportError: from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.moe import MoeUnpermuteFwdOp -from workloads.ops.moe_unpermute import MoeUnpermuteTest +from workloads.moe import MoeUnpermuteTest _OP_NAME = "MoeUnpermuteFwdOp" diff --git a/benchmarks/ops/bench_reduce.py b/benchmarks/ops/bench_reduce.py index be6a9bfa..c20791dc 100644 --- a/benchmarks/ops/bench_reduce.py +++ b/benchmarks/ops/bench_reduce.py @@ -18,7 +18,7 @@ from tileops.ops.reduction.reduce import ( VarFwdOp, VarMeanFwdOp, ) -from workloads.ops.reduce import ( +from workloads.reduce import ( AmaxTest, AminTest, MeanTest, diff --git a/benchmarks/ops/bench_rms_norm.py b/benchmarks/ops/bench_rms_norm.py index 4184e9d3..7eb38ebc 100644 --- a/benchmarks/ops/bench_rms_norm.py +++ b/benchmarks/ops/bench_rms_norm.py @@ -6,7 +6,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.manifest import eval_roofline, load_workloads from tileops.ops.norm.rms_norm import RMSNormFwdOp -from workloads.ops.rms_norm import RmsNormTest +from workloads.rms_norm import RmsNormTest class _RmsNormTestBaseline(RmsNormTest): diff --git a/benchmarks/ops/bench_softmax.py b/benchmarks/ops/bench_softmax.py index 5519fb24..6056a3cb 100644 --- a/benchmarks/ops/bench_softmax.py +++ b/benchmarks/ops/bench_softmax.py @@ -12,7 +12,7 @@ from benchmarks.benchmark import BenchmarkReport, ManifestBenchmark, workloads_t from tileops.ops.reduction.log_softmax import LogSoftmaxFwdOp from tileops.ops.reduction.logsumexp import LogSumExpFwdOp from tileops.ops.reduction.softmax import SoftmaxFwdOp -from workloads.ops.softmax import ( +from workloads.softmax import ( LogSoftmaxTest, LogSumExpTest, SoftmaxTest, diff --git a/benchmarks/ops/bench_topk_selector.py b/benchmarks/ops/bench_topk_selector.py index 12cb2c21..786613d1 100644 --- a/benchmarks/ops/bench_topk_selector.py +++ b/benchmarks/ops/bench_topk_selector.py @@ -5,7 +5,7 @@ import torch from benchmarks.benchmark import BenchmarkBase, BenchmarkReport from tileops.ops import TopkSelectorOp -from workloads.ops.topk_selector import TopkSelectorTest +from workloads.topk_selector import TopkSelectorTest class _TopkSelectorTestBaseline(TopkSelectorTest): diff --git a/benchmarks/ops/bench_vector_norm.py b/benchmarks/ops/bench_vector_norm.py index e3ea2a8a..87338cda 100644 --- a/benchmarks/ops/bench_vector_norm.py +++ b/benchmarks/ops/bench_vector_norm.py @@ -11,7 +11,7 @@ from benchmarks.benchmark import BenchmarkReport, ManifestBenchmark, workloads_t from tileops.ops.reduction.inf_norm import InfNormFwdOp from tileops.ops.reduction.l1_norm import L1NormFwdOp from tileops.ops.reduction.l2_norm import L2NormFwdOp -from workloads.ops.vector_norm import InfNormTest, L1NormTest, L2NormTest +from workloads.vector_norm import InfNormTest, L1NormTest, L2NormTest # =================================================================== # Op name constants diff --git a/docs/testing.md b/docs/testing.md index 18490660..e69ea617 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -15,7 +15,7 @@ Tests and benchmarks are separated by concern: `pytest tests/` validates correct ## Test/Benchmark Pattern ```python -# workloads/ops/mha.py +# workloads/attention/mha.py class MhaFwdTest(WorkloadBase): def __init__(self, batch, heads, seq_len, dim, causal, dtype): ... def gen_inputs(self): ... @@ -23,7 +23,7 @@ class MhaFwdTest(WorkloadBase): # tests/ops/test_mha.py from tileops.ops import MhaFwdOp -from workloads.ops.mha import MhaFwdTest +from workloads.attention.mha import MhaFwdTest class MhaFwdTestCase(MhaFwdTest, TestBase): @@ -43,7 +43,7 @@ def test_mha_fwd(batch, seq_len, heads, dim, causal, dtype, tune): # benchmarks/ops/bench_mha.py from tileops.ops import MhaFwdOp -from workloads.ops.mha import MhaFwdTest # import workload, NOT test +from workloads.attention.mha import MhaFwdTest # import workload, NOT test class MhaFwdBenchmark(BenchmarkBase): @@ -151,7 +151,7 @@ python scripts/test_node_delta.py --base origin/release # different base branc ### File checklist -1. **Workload class** in `workloads/ops/` — subclass `WorkloadBase`, implement `gen_inputs()`. +1. **Workload class** in `workloads/` — subclass `WorkloadBase`, implement `gen_inputs()`. 1. **Fixture class** — subclass `FixtureBase`, define `PARAMS` with `smoke`/`full` marks. 1. **Test class** in `tests/ops/test_.py` — inherit `(MyWorkload, TestBase)`, implement `ref_program()` locally. 1. **Test function** — `@YourFixture` decorated, call `test.check(op, *test.gen_inputs())`. @@ -177,7 +177,7 @@ TestBase (tests/test_base.py, inherits WorkloadBase) ### File checklist -1. **Workload class** in `workloads/ops/` — reuse the `WorkloadBase` subclass from the test. +1. **Workload class** in `workloads/` — reuse the `WorkloadBase` subclass from the test. 1. **Fixture class** — reuse the `FixtureBase` subclass from the test. 1. **Benchmark class** in `benchmarks/ops/bench_.py` — subclass `BenchmarkBase`, implement `calculate_flops()` and `calculate_memory()` (return `None` if not applicable). 1. **Benchmark function** — `@YourFixture` decorated, construct workload + benchmark, call `inputs = workload.gen_inputs()`, then `bm.profile(op, *inputs)` and `BenchmarkReport.record(op, locals(), result, tag="tileops")`. diff --git a/tests/ops/attention/test_deepseek_dsa_decode.py b/tests/ops/attention/test_deepseek_dsa_decode.py index c58f9269..33eaa1fa 100644 --- a/tests/ops/attention/test_deepseek_dsa_decode.py +++ b/tests/ops/attention/test_deepseek_dsa_decode.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import DeepSeekSparseAttentionDecodeWithKVCacheFwdOp -from workloads.ops.deepseek_dsa_decode import DsaDecodeTest as _DsaDecodeTestWorkload +from workloads.attention.deepseek_dsa_decode import DsaDecodeTest as _DsaDecodeTestWorkload class DsaDecodeTest(_DsaDecodeTestWorkload, TestBase): diff --git a/tests/ops/attention/test_deepseek_mla_decode.py b/tests/ops/attention/test_deepseek_mla_decode.py index 564e6c84..1b6fa9a7 100644 --- a/tests/ops/attention/test_deepseek_mla_decode.py +++ b/tests/ops/attention/test_deepseek_mla_decode.py @@ -6,7 +6,7 @@ from einops import einsum, rearrange from tests.test_base import FixtureBase, TestBase from tileops.ops import MultiHeadLatentAttentionDecodeWithKVCacheFwdOp -from workloads.ops.deepseek_mla_decode import MlaDecodeTest as _MlaDecodeTestWorkload +from workloads.attention.deepseek_mla_decode import MlaDecodeTest as _MlaDecodeTestWorkload class MlaDecodeTest(_MlaDecodeTestWorkload, TestBase): diff --git a/tests/ops/attention/test_deepseek_nsa.py b/tests/ops/attention/test_deepseek_nsa.py index edbc5855..6cf99ddf 100644 --- a/tests/ops/attention/test_deepseek_nsa.py +++ b/tests/ops/attention/test_deepseek_nsa.py @@ -6,7 +6,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import NSAFwdVarlenOp -from workloads.ops.deepseek_nsa_fwd import NsaFwdTest as _NsaFwdTestWorkload +from workloads.attention.deepseek_nsa import NsaFwdTest as _NsaFwdTestWorkload class NsaFwdTest(_NsaFwdTestWorkload, TestBase): diff --git a/tests/ops/attention/test_deepseek_nsa_cmp.py b/tests/ops/attention/test_deepseek_nsa_cmp.py index b8402025..660490c4 100644 --- a/tests/ops/attention/test_deepseek_nsa_cmp.py +++ b/tests/ops/attention/test_deepseek_nsa_cmp.py @@ -6,8 +6,8 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import NSACmpFwdVarlenOp +from workloads.attention.deepseek_nsa_cmp import NsaCmpFwdTest as _NsaCmpFwdTestWorkload from workloads.nsa_utils import prepare_chunk_offsets -from workloads.ops.deepseek_nsa_cmp_fwd import NsaCmpFwdTest as _NsaCmpFwdTestWorkload def _parallel_nsa_compression_fwd_pytorch(test, q, k_cmp, v_cmp, block_size, scale, offsets): diff --git a/tests/ops/attention/test_deepseek_nsa_topk.py b/tests/ops/attention/test_deepseek_nsa_topk.py index 7c2c049f..5cd7999c 100644 --- a/tests/ops/attention/test_deepseek_nsa_topk.py +++ b/tests/ops/attention/test_deepseek_nsa_topk.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import NSATopkVarlenOp -from workloads.ops.deepseek_nsa_topk import NsaTopkTest as _NsaTopkTestWorkload +from workloads.attention.deepseek_nsa_topk import NsaTopkTest as _NsaTopkTestWorkload def _nsa_topk_torch(test, q, k_cmp, lse, block_counts, block_size, scale, diff --git a/tests/ops/attention/test_gqa.py b/tests/ops/attention/test_gqa.py index bf24f1c1..364913bb 100644 --- a/tests/ops/attention/test_gqa.py +++ b/tests/ops/attention/test_gqa.py @@ -8,8 +8,8 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from tests.test_base import FixtureBase, TestBase from tileops.ops import GroupedQueryAttentionBwdOp, GroupedQueryAttentionFwdOp -from workloads.ops.gqa import GqaBwdTest as _GqaBwdTestWorkload -from workloads.ops.gqa import GqaFwdTest as _GqaFwdTestWorkload +from workloads.attention.gqa import GqaBwdTest as _GqaBwdTestWorkload +from workloads.attention.gqa import GqaFwdTest as _GqaFwdTestWorkload class GqaBwdTest(_GqaBwdTestWorkload, TestBase): diff --git a/tests/ops/attention/test_gqa_decode.py b/tests/ops/attention/test_gqa_decode.py index f808bffc..c5ae587b 100644 --- a/tests/ops/attention/test_gqa_decode.py +++ b/tests/ops/attention/test_gqa_decode.py @@ -6,7 +6,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from tests.test_base import FixtureBase, TestBase from tileops.ops import GroupedQueryAttentionDecodeWithKVCacheFwdOp -from workloads.ops.gqa_decode import GqaDecodeTest as _GqaDecodeTestWorkload +from workloads.attention.gqa_decode import GqaDecodeTest as _GqaDecodeTestWorkload class GqaDecodeTest(_GqaDecodeTestWorkload, TestBase): diff --git a/tests/ops/attention/test_gqa_decode_paged.py b/tests/ops/attention/test_gqa_decode_paged.py index 8e6a0655..d66c066b 100644 --- a/tests/ops/attention/test_gqa_decode_paged.py +++ b/tests/ops/attention/test_gqa_decode_paged.py @@ -10,7 +10,9 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from tests.test_base import FixtureBase, TestBase from tileops.ops import GroupedQueryAttentionDecodePagedWithKVCacheFwdOp -from workloads.ops.gqa_decode_paged import GqaDecodePagedTest as _GqaDecodePagedTestWorkload +from workloads.attention.gqa_decode_paged import ( + GqaDecodePagedTest as _GqaDecodePagedTestWorkload, +) class GqaDecodePagedTest(_GqaDecodePagedTestWorkload, TestBase): diff --git a/tests/ops/attention/test_gqa_sliding_window.py b/tests/ops/attention/test_gqa_sliding_window.py index 8625d4f3..f5c7cdd6 100644 --- a/tests/ops/attention/test_gqa_sliding_window.py +++ b/tests/ops/attention/test_gqa_sliding_window.py @@ -5,7 +5,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GqaSlidingWindowFwdOp -from workloads.ops.gqa_sliding_window_fwd import ( +from workloads.attention.gqa_sliding_window import ( GqaSlidingWindowFwdTest as _GqaSlidingWindowFwdTestWorkload, ) diff --git a/tests/ops/attention/test_gqa_sliding_window_varlen.py b/tests/ops/attention/test_gqa_sliding_window_varlen.py index d1efec84..24470016 100644 --- a/tests/ops/attention/test_gqa_sliding_window_varlen.py +++ b/tests/ops/attention/test_gqa_sliding_window_varlen.py @@ -6,7 +6,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GqaSlidingWindowVarlenFwdOp -from workloads.ops.gqa_sliding_window_varlen_fwd import ( +from workloads.attention.gqa_sliding_window_varlen import ( GqaSlidingWindowVarlenFwdTest as _GqaSlidingWindowVarlenFwdTestWorkload, ) diff --git a/tests/ops/attention/test_mean_pooling.py b/tests/ops/attention/test_mean_pooling.py index a167b28c..d4ba8ec7 100644 --- a/tests/ops/attention/test_mean_pooling.py +++ b/tests/ops/attention/test_mean_pooling.py @@ -5,8 +5,8 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import MeanPoolingForwardOp +from workloads.attention.mean_pooling import MeanPoolingTest as _MeanPoolingTestWorkload from workloads.nsa_utils import prepare_chunk_indices -from workloads.ops.mean_pooling_ops import MeanPoolingTest as _MeanPoolingTestWorkload class MeanPoolingTest(_MeanPoolingTestWorkload, TestBase): diff --git a/tests/ops/attention/test_mha.py b/tests/ops/attention/test_mha.py index 1190b7ef..6b897500 100644 --- a/tests/ops/attention/test_mha.py +++ b/tests/ops/attention/test_mha.py @@ -8,8 +8,8 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from tests.test_base import FixtureBase, TestBase from tileops.ops import MultiHeadAttentionBwdOp, MultiHeadAttentionFwdOp -from workloads.ops.mha import MhaBwdTest as _MhaBwdTestWorkload -from workloads.ops.mha import MhaFwdTest as _MhaFwdTestWorkload +from workloads.attention.mha import MhaBwdTest as _MhaBwdTestWorkload +from workloads.attention.mha import MhaFwdTest as _MhaFwdTestWorkload class MhaBwdTest(_MhaBwdTestWorkload, TestBase): diff --git a/tests/ops/attention/test_mha_decode.py b/tests/ops/attention/test_mha_decode.py index dc45b6c3..1263deb4 100644 --- a/tests/ops/attention/test_mha_decode.py +++ b/tests/ops/attention/test_mha_decode.py @@ -6,7 +6,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from tests.test_base import FixtureBase, TestBase from tileops.ops import MultiHeadAttentionDecodeWithKVCacheFwdOp -from workloads.ops.mha_decode import MhaDecodeTest as _MhaDecodeTestWorkload +from workloads.attention.mha_decode import MhaDecodeTest as _MhaDecodeTestWorkload class MhaDecodeTest(_MhaDecodeTestWorkload, TestBase): diff --git a/tests/ops/attention/test_mha_decode_paged.py b/tests/ops/attention/test_mha_decode_paged.py index 1168f007..10f1220f 100644 --- a/tests/ops/attention/test_mha_decode_paged.py +++ b/tests/ops/attention/test_mha_decode_paged.py @@ -9,7 +9,9 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops import MultiHeadAttentionDecodePagedWithKVCacheFwdOp -from workloads.ops.mha_decode_paged import MhaDecodePagedTest as _MhaDecodePagedTestWorkload +from workloads.attention.mha_decode_paged import ( + MhaDecodePagedTest as _MhaDecodePagedTestWorkload, +) class MhaDecodePagedTest(_MhaDecodePagedTestWorkload, TestBase): diff --git a/tests/ops/test_activation.py b/tests/ops/test_activation.py index 43843dc2..1bd48682 100644 --- a/tests/ops/test_activation.py +++ b/tests/ops/test_activation.py @@ -9,7 +9,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.elementwise import ReluOp -from workloads.ops.activation import ReluTest as _ReluTestWorkload +from workloads.activation import ReluTest as _ReluTestWorkload class ReluTest(_ReluTestWorkload, TestBase): diff --git a/tests/ops/test_ada_layer_norm.py b/tests/ops/test_ada_layer_norm.py index b1ad9e4e..b346f16e 100644 --- a/tests/ops/test_ada_layer_norm.py +++ b/tests/ops/test_ada_layer_norm.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.ada_layer_norm import AdaLayerNormFwdOp -from workloads.ops.ada_layer_norm import AdaLayerNormTest as _AdaLayerNormTestWorkload +from workloads.ada_layer_norm import AdaLayerNormTest as _AdaLayerNormTestWorkload class AdaLayerNormTest(_AdaLayerNormTestWorkload, TestBase): diff --git a/tests/ops/test_ada_layer_norm_zero.py b/tests/ops/test_ada_layer_norm_zero.py index 51728be6..95ba7216 100644 --- a/tests/ops/test_ada_layer_norm_zero.py +++ b/tests/ops/test_ada_layer_norm_zero.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.ada_layer_norm_zero import AdaLayerNormZeroFwdOp -from workloads.ops.ada_layer_norm_zero import AdaLayerNormZeroTest as _AdaLayerNormZeroTestWorkload +from workloads.ada_layer_norm_zero import AdaLayerNormZeroTest as _AdaLayerNormZeroTestWorkload class AdaLayerNormZeroTest(_AdaLayerNormZeroTestWorkload, TestBase): diff --git a/tests/ops/test_argreduce.py b/tests/ops/test_argreduce.py index d4808f5b..2845c59e 100644 --- a/tests/ops/test_argreduce.py +++ b/tests/ops/test_argreduce.py @@ -9,7 +9,7 @@ import pytest import torch from tests.test_base import FixtureBase, TestBase -from workloads.ops.argreduce import ArgmaxTest as _ArgmaxWorkload +from workloads.argreduce import ArgmaxTest as _ArgmaxWorkload # --------------------------------------------------------------------------- # Fixtures diff --git a/tests/ops/test_batch_norm.py b/tests/ops/test_batch_norm.py index 0ec5158a..e99b84c0 100644 --- a/tests/ops/test_batch_norm.py +++ b/tests/ops/test_batch_norm.py @@ -13,10 +13,10 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.batch_norm import BatchNormBwdOp, BatchNormFwdOp -from workloads.ops.batch_norm import ( +from workloads.batch_norm import ( BatchNormBwdTest as _BatchNormBwdTestWorkload, ) -from workloads.ops.batch_norm import ( +from workloads.batch_norm import ( BatchNormFwdTest as _BatchNormFwdTestWorkload, ) diff --git a/tests/ops/test_binary_arith.py b/tests/ops/test_binary_arith.py index 2fda078d..3df60b92 100644 --- a/tests/ops/test_binary_arith.py +++ b/tests/ops/test_binary_arith.py @@ -22,7 +22,7 @@ from tileops.ops.elementwise import ( SubOp, coalesce_broadcast_dims, ) -from workloads.ops.binary_arith import AddSameShapeTest as _AddSameShapeTestWorkload +from workloads.binary_arith import AddSameShapeTest as _AddSameShapeTestWorkload class AddSameShapeTest(_AddSameShapeTestWorkload, TestBase): diff --git a/tests/ops/test_deltanet_fwd.py b/tests/ops/test_deltanet_fwd.py index 1e755b99..66039c42 100644 --- a/tests/ops/test_deltanet_fwd.py +++ b/tests/ops/test_deltanet_fwd.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import DeltaNetFwdOp -from workloads.ops.deltanet_fwd import DeltaNetFwdTest as _DeltaNetFwdTestWorkload +from workloads.deltanet import DeltaNetFwdTest as _DeltaNetFwdTestWorkload def compute_w_u_torch(Aw, Au, k, v, beta, chunk_size): diff --git a/tests/ops/test_deltanet_recurrence.py b/tests/ops/test_deltanet_recurrence.py index ceaebbfd..8308ad1b 100644 --- a/tests/ops/test_deltanet_recurrence.py +++ b/tests/ops/test_deltanet_recurrence.py @@ -6,7 +6,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import DeltaNetDecodeOp -from workloads.ops.deltanet_recurrence import DeltaNetDecodeTest as _DeltaNetDecodeTestWorkload +from workloads.deltanet import DeltaNetDecodeTest as _DeltaNetDecodeTestWorkload def deltanet_decode_torch( diff --git a/tests/ops/test_engram.py b/tests/ops/test_engram.py index b813deec..0c19af81 100644 --- a/tests/ops/test_engram.py +++ b/tests/ops/test_engram.py @@ -5,10 +5,16 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.engram import EngramGateConvBwdOp, EngramGateConvFwdOp from tileops.ops.engram_decode import EngramDecodeOp -from workloads.ops.engram_bwd import EngramGateConvBwdTest as _EngramGateConvBwdTestWorkload -from workloads.ops.engram_decode import EngramDecodeTest as _EngramDecodeTestWorkload -from workloads.ops.engram_fwd import CONV_KERNEL_SIZE -from workloads.ops.engram_fwd import ( +from workloads.engram import ( + CONV_KERNEL_SIZE, +) +from workloads.engram import ( + EngramDecodeTest as _EngramDecodeTestWorkload, +) +from workloads.engram import ( + EngramGateConvBwdTest as _EngramGateConvBwdTestWorkload, +) +from workloads.engram import ( EngramGateConvFwdTest as _EngramGateConvFwdTestWorkload, ) diff --git a/tests/ops/test_fft.py b/tests/ops/test_fft.py index 38a0d2c5..5b995af3 100644 --- a/tests/ops/test_fft.py +++ b/tests/ops/test_fft.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import FFTC2COp -from workloads.ops.fft import FFTTest as _FFTTestWorkload +from workloads.fft import FFTTest as _FFTTestWorkload class FFTTest(_FFTTestWorkload, TestBase): diff --git a/tests/ops/test_fp8_lighting_indexer.py b/tests/ops/test_fp8_lighting_indexer.py index dac876d5..822c0862 100644 --- a/tests/ops/test_fp8_lighting_indexer.py +++ b/tests/ops/test_fp8_lighting_indexer.py @@ -5,7 +5,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import FP8LightingIndexerOp -from workloads.ops.fp8_lighting_indexer import ( +from workloads.fp8_lighting_indexer import ( FP8LightingIndexerTest as _FP8LightingIndexerTestWorkload, ) diff --git a/tests/ops/test_fp8_quant.py b/tests/ops/test_fp8_quant.py index aa482726..738d8fc0 100644 --- a/tests/ops/test_fp8_quant.py +++ b/tests/ops/test_fp8_quant.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops import FP8QuantOp -from workloads.ops.fp8_quant import FP8QuantTest as _FP8QuantTestWorkload +from workloads.fp8_quant import FP8QuantTest as _FP8QuantTestWorkload class FP8QuantTest(_FP8QuantTestWorkload, TestBase): diff --git a/tests/ops/test_fused_add_layer_norm.py b/tests/ops/test_fused_add_layer_norm.py index 84e992e4..8e8ce653 100644 --- a/tests/ops/test_fused_add_layer_norm.py +++ b/tests/ops/test_fused_add_layer_norm.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.fused_add_layer_norm import FusedAddLayerNormFwdOp -from workloads.ops.fused_add_layer_norm import ( +from workloads.fused_add_layer_norm import ( FusedAddLayerNormTest as _FusedAddLayerNormTestWorkload, ) diff --git a/tests/ops/test_fused_add_rmsnorm.py b/tests/ops/test_fused_add_rmsnorm.py index d9652528..d11699b3 100644 --- a/tests/ops/test_fused_add_rmsnorm.py +++ b/tests/ops/test_fused_add_rmsnorm.py @@ -3,7 +3,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.fused_add_rmsnorm import FusedAddRMSNormFwdOp -from workloads.ops.fused_add_rmsnorm import FusedAddRmsNormTest as _FusedAddRmsNormTestWorkload +from workloads.fused_add_rmsnorm import FusedAddRmsNormTest as _FusedAddRmsNormTestWorkload class FusedAddRmsNormTest(_FusedAddRmsNormTestWorkload, TestBase): diff --git a/tests/ops/test_gated_deltanet_fwd.py b/tests/ops/test_gated_deltanet_fwd.py index 0ddc3fa1..3e5b0e65 100644 --- a/tests/ops/test_gated_deltanet_fwd.py +++ b/tests/ops/test_gated_deltanet_fwd.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GatedDeltaNetFwdOp -from workloads.ops.gated_deltanet_fwd import ( +from workloads.gated_deltanet import ( GatedDeltaNetFwdTest as _GatedDeltaNetFwdTestWorkload, ) diff --git a/tests/ops/test_gated_deltanet_recurrence.py b/tests/ops/test_gated_deltanet_recurrence.py index e3f35157..237c3886 100644 --- a/tests/ops/test_gated_deltanet_recurrence.py +++ b/tests/ops/test_gated_deltanet_recurrence.py @@ -6,7 +6,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GatedDeltaNetDecodeOp -from workloads.ops.gated_deltanet_recurrence import ( +from workloads.gated_deltanet import ( GatedDeltaNetDecodeTest as _GatedDeltaNetDecodeTestWorkload, ) diff --git a/tests/ops/test_gemm.py b/tests/ops/test_gemm.py index cbf89eb9..733f843d 100644 --- a/tests/ops/test_gemm.py +++ b/tests/ops/test_gemm.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GemmOp -from workloads.ops.gemm import GemmTest as _GemmTestWorkload +from workloads.gemm import GemmTest as _GemmTestWorkload class GemmTest(_GemmTestWorkload, TestBase): diff --git a/tests/ops/test_gla_recurrence.py b/tests/ops/test_gla_recurrence.py index 7ced5434..6eaa064f 100644 --- a/tests/ops/test_gla_recurrence.py +++ b/tests/ops/test_gla_recurrence.py @@ -6,7 +6,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import GLADecodeOp -from workloads.ops.gla_recurrence import GLADecodeTest as _GLADecodeTestWorkload +from workloads.gla import GLADecodeTest as _GLADecodeTestWorkload def gla_decode_torch( diff --git a/tests/ops/test_group_norm.py b/tests/ops/test_group_norm.py index 921c8f76..f2050664 100644 --- a/tests/ops/test_group_norm.py +++ b/tests/ops/test_group_norm.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.group_norm import GroupNormFwdOp -from workloads.ops.group_norm import GroupNormTest as _GroupNormTestWorkload +from workloads.group_norm import GroupNormTest as _GroupNormTestWorkload class GroupNormTest(_GroupNormTestWorkload, TestBase): diff --git a/tests/ops/test_grouped_gemm.py b/tests/ops/test_grouped_gemm.py index 32e5bb95..f266686d 100644 --- a/tests/ops/test_grouped_gemm.py +++ b/tests/ops/test_grouped_gemm.py @@ -4,7 +4,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.grouped_gemm import GroupedGemmOp -from workloads.ops.grouped_gemm import ( +from workloads.grouped_gemm import ( GroupedGemmTest as _GroupedGemmTestWorkload, ) diff --git a/tests/ops/test_instance_norm.py b/tests/ops/test_instance_norm.py index 11b82b9b..70d748fb 100644 --- a/tests/ops/test_instance_norm.py +++ b/tests/ops/test_instance_norm.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.instance_norm import InstanceNormFwdOp -from workloads.ops.instance_norm import InstanceNormTest as _InstanceNormTestWorkload +from workloads.instance_norm import InstanceNormTest as _InstanceNormTestWorkload class InstanceNormTest(_InstanceNormTestWorkload, TestBase): diff --git a/tests/ops/test_layer_norm.py b/tests/ops/test_layer_norm.py index 17676023..d4b41cd4 100644 --- a/tests/ops/test_layer_norm.py +++ b/tests/ops/test_layer_norm.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.layer_norm import LayerNormFwdOp -from workloads.ops.layer_norm import LayerNormTest as _LayerNormTestWorkload +from workloads.layer_norm import LayerNormTest as _LayerNormTestWorkload class LayerNormTest(_LayerNormTestWorkload, TestBase): diff --git a/tests/ops/test_logical_reduce.py b/tests/ops/test_logical_reduce.py index 2a35275e..313b61d6 100644 --- a/tests/ops/test_logical_reduce.py +++ b/tests/ops/test_logical_reduce.py @@ -10,7 +10,7 @@ import pytest import torch from tests.test_base import FixtureBase, TestBase -from workloads.ops.logical_reduce import AnyTest as _AnyWorkload +from workloads.logical_reduce import AnyTest as _AnyWorkload # --------------------------------------------------------------------------- # Fixtures diff --git a/tests/ops/test_mamba.py b/tests/ops/test_mamba.py index 5b5e88c4..1fc14723 100644 --- a/tests/ops/test_mamba.py +++ b/tests/ops/test_mamba.py @@ -7,16 +7,18 @@ from tileops.ops.ssd_chunk_scan import SsdChunkScanFwdOp from tileops.ops.ssd_chunk_state import SsdChunkStateFwdOp from tileops.ops.ssd_decode import SsdDecodeOp from tileops.ops.ssd_state_passing import SsdStatePassingFwdOp -from workloads.ops.da_cumsum_fwd import DaCumsumFwdFixture -from workloads.ops.da_cumsum_fwd import DaCumsumFwdTest as _DaCumsumFwdTestWorkload -from workloads.ops.ssd_chunk_scan_fwd import SsdChunkScanFwdFixture -from workloads.ops.ssd_chunk_scan_fwd import SsdChunkScanFwdTest as _SsdChunkScanFwdTestWorkload -from workloads.ops.ssd_chunk_state_fwd import SsdChunkStateFwdFixture -from workloads.ops.ssd_chunk_state_fwd import SsdChunkStateFwdTest as _SsdChunkStateFwdTestWorkload -from workloads.ops.ssd_decode import SsdDecodeFixture -from workloads.ops.ssd_decode import SsdDecodeTest as _SsdDecodeTestWorkload -from workloads.ops.ssd_state_passing_fwd import SsdStatePassingFwdFixture -from workloads.ops.ssd_state_passing_fwd import ( +from workloads.mamba import ( + DaCumsumFwdFixture, + SsdChunkScanFwdFixture, + SsdChunkStateFwdFixture, + SsdDecodeFixture, + SsdStatePassingFwdFixture, +) +from workloads.mamba import DaCumsumFwdTest as _DaCumsumFwdTestWorkload +from workloads.mamba import SsdChunkScanFwdTest as _SsdChunkScanFwdTestWorkload +from workloads.mamba import SsdChunkStateFwdTest as _SsdChunkStateFwdTestWorkload +from workloads.mamba import SsdDecodeTest as _SsdDecodeTestWorkload +from workloads.mamba import ( SsdStatePassingFwdTest as _SsdStatePassingFwdTestWorkload, ) diff --git a/tests/ops/test_mhc_post.py b/tests/ops/test_mhc_post.py index 7f1e350e..b44044ce 100644 --- a/tests/ops/test_mhc_post.py +++ b/tests/ops/test_mhc_post.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops import MHCPostOp -from workloads.ops.mhc_post import MHCPostTest as _MHCPostTestWorkload +from workloads.mhc import MHCPostTest as _MHCPostTestWorkload class MHCPostTest(_MHCPostTestWorkload, TestBase): diff --git a/tests/ops/test_mhc_pre.py b/tests/ops/test_mhc_pre.py index 433d07cb..3a66e8a3 100644 --- a/tests/ops/test_mhc_pre.py +++ b/tests/ops/test_mhc_pre.py @@ -10,7 +10,7 @@ import torch.nn.functional as F from tests.test_base import FixtureBase, TestBase from tileops.ops import MHCPreOp -from workloads.ops.mhc_pre import MHCPreTest as _MHCPreTestWorkload +from workloads.mhc import MHCPreTest as _MHCPreTestWorkload class MHCPreTest(_MHCPreTestWorkload, TestBase): diff --git a/tests/ops/test_moe_fused_topk.py b/tests/ops/test_moe_fused_topk.py index 0f3f6cde..00cbda6a 100644 --- a/tests/ops/test_moe_fused_topk.py +++ b/tests/ops/test_moe_fused_topk.py @@ -16,7 +16,7 @@ import torch from tests.test_base import FixtureBase from tileops.ops.moe import FusedTopKOp -from workloads.ops.moe_fused_topk import FusedTopKTest +from workloads.moe import FusedTopKTest def fused_topk_torch( diff --git a/tests/ops/test_moe_permute.py b/tests/ops/test_moe_permute.py index 5877e23e..d5dffb46 100644 --- a/tests/ops/test_moe_permute.py +++ b/tests/ops/test_moe_permute.py @@ -16,7 +16,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.moe import MoePermutePaddedFwdOp -from workloads.ops.moe_permute import MoePermuteTest as _MoePermuteTestWorkload +from workloads.moe import MoePermuteTest as _MoePermuteTestWorkload def _ref_moe_permute( diff --git a/tests/ops/test_moe_permute_align.py b/tests/ops/test_moe_permute_align.py index 3ed86bd0..3341796d 100644 --- a/tests/ops/test_moe_permute_align.py +++ b/tests/ops/test_moe_permute_align.py @@ -15,7 +15,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.moe import MoePermuteAlignFwdOp -from workloads.ops.moe_permute_align import MoePermuteAlignTest as _MoePermuteAlignTestWorkload +from workloads.moe import MoePermuteAlignTest as _MoePermuteAlignTestWorkload def _ref_permute_align( diff --git a/tests/ops/test_moe_unpermute.py b/tests/ops/test_moe_unpermute.py index 6d8e167a..8aabaf61 100644 --- a/tests/ops/test_moe_unpermute.py +++ b/tests/ops/test_moe_unpermute.py @@ -18,7 +18,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.moe import MoeUnpermuteFwdOp -from workloads.ops.moe_unpermute import MoeUnpermuteTest as _MoeUnpermuteTestWorkload +from workloads.moe import MoeUnpermuteTest as _MoeUnpermuteTestWorkload def _ref_moe_unpermute( diff --git a/tests/ops/test_reduce.py b/tests/ops/test_reduce.py index d54b90a3..beff67f8 100644 --- a/tests/ops/test_reduce.py +++ b/tests/ops/test_reduce.py @@ -8,13 +8,13 @@ import pytest import torch from tests.test_base import FixtureBase, TestBase -from workloads.ops.reduce import ( +from workloads.reduce import ( ProdTest as _ProdTest, ) -from workloads.ops.reduce import ( +from workloads.reduce import ( StdTest as _StdTest, ) -from workloads.ops.reduce import ( +from workloads.reduce import ( SumTest as _SumTest, ) diff --git a/tests/ops/test_rms_norm.py b/tests/ops/test_rms_norm.py index 86140ae4..1b57bcde 100644 --- a/tests/ops/test_rms_norm.py +++ b/tests/ops/test_rms_norm.py @@ -3,7 +3,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops.norm.rms_norm import RMSNormFwdOp -from workloads.ops.rms_norm import RmsNormTest as _RmsNormTestWorkload +from workloads.rms_norm import RmsNormTest as _RmsNormTestWorkload class RmsNormTest(_RmsNormTestWorkload, TestBase): diff --git a/tests/ops/test_softmax.py b/tests/ops/test_softmax.py index efc03771..e817f172 100644 --- a/tests/ops/test_softmax.py +++ b/tests/ops/test_softmax.py @@ -21,9 +21,9 @@ from tests.test_base import FixtureBase, TestBase from tileops.ops.reduction.log_softmax import LogSoftmaxFwdOp from tileops.ops.reduction.logsumexp import LogSumExpFwdOp from tileops.ops.reduction.softmax import SoftmaxFwdOp -from workloads.ops.softmax import LogSoftmaxTest as _LogSoftmaxTestWorkload -from workloads.ops.softmax import LogSumExpTest as _LogSumExpTestWorkload -from workloads.ops.softmax import SoftmaxTest as _SoftmaxTestWorkload +from workloads.softmax import LogSoftmaxTest as _LogSoftmaxTestWorkload +from workloads.softmax import LogSumExpTest as _LogSumExpTestWorkload +from workloads.softmax import SoftmaxTest as _SoftmaxTestWorkload # --------------------------------------------------------------------------- # Tolerances (from docs/testing.md) diff --git a/tests/ops/test_topk_selector.py b/tests/ops/test_topk_selector.py index d16138e8..186e8cff 100644 --- a/tests/ops/test_topk_selector.py +++ b/tests/ops/test_topk_selector.py @@ -5,7 +5,7 @@ import torch from tests.test_base import FixtureBase, TestBase from tileops.ops import TopkSelectorOp from tileops.utils import str2dtype -from workloads.ops.topk_selector import TopkSelectorTest as _TopkSelectorTestWorkload +from workloads.topk_selector import TopkSelectorTest as _TopkSelectorTestWorkload class TopkSelectorTest(_TopkSelectorTestWorkload, TestBase): diff --git a/tests/ops/test_vector_norm.py b/tests/ops/test_vector_norm.py index 46148780..bfa096ca 100644 --- a/tests/ops/test_vector_norm.py +++ b/tests/ops/test_vector_norm.py @@ -9,7 +9,7 @@ import pytest import torch from tests.test_base import FixtureBase, TestBase, allclose_compare -from workloads.ops.vector_norm import L1NormTest as _L1NormWorkload +from workloads.vector_norm import L1NormTest as _L1NormWorkload # --------------------------------------------------------------------------- # Fixtures diff --git a/workloads/ops/activation.py b/workloads/activation.py similarity index 100% rename from workloads/ops/activation.py rename to workloads/activation.py diff --git a/workloads/ops/ada_layer_norm.py b/workloads/ada_layer_norm.py similarity index 100% rename from workloads/ops/ada_layer_norm.py rename to workloads/ada_layer_norm.py diff --git a/workloads/ops/ada_layer_norm_zero.py b/workloads/ada_layer_norm_zero.py similarity index 100% rename from workloads/ops/ada_layer_norm_zero.py rename to workloads/ada_layer_norm_zero.py diff --git a/workloads/ops/argreduce.py b/workloads/argreduce.py similarity index 100% rename from workloads/ops/argreduce.py rename to workloads/argreduce.py diff --git a/workloads/ops/__init__.py b/workloads/attention/__init__.py similarity index 100% rename from workloads/ops/__init__.py rename to workloads/attention/__init__.py diff --git a/workloads/ops/deepseek_dsa_decode.py b/workloads/attention/deepseek_dsa_decode.py similarity index 100% rename from workloads/ops/deepseek_dsa_decode.py rename to workloads/attention/deepseek_dsa_decode.py diff --git a/workloads/ops/deepseek_mla_decode.py b/workloads/attention/deepseek_mla_decode.py similarity index 100% rename from workloads/ops/deepseek_mla_decode.py rename to workloads/attention/deepseek_mla_decode.py diff --git a/workloads/ops/deepseek_nsa_fwd.py b/workloads/attention/deepseek_nsa.py similarity index 100% rename from workloads/ops/deepseek_nsa_fwd.py rename to workloads/attention/deepseek_nsa.py diff --git a/workloads/ops/deepseek_nsa_cmp_fwd.py b/workloads/attention/deepseek_nsa_cmp.py similarity index 100% rename from workloads/ops/deepseek_nsa_cmp_fwd.py rename to workloads/attention/deepseek_nsa_cmp.py diff --git a/workloads/ops/deepseek_nsa_topk.py b/workloads/attention/deepseek_nsa_topk.py similarity index 100% rename from workloads/ops/deepseek_nsa_topk.py rename to workloads/attention/deepseek_nsa_topk.py diff --git a/workloads/ops/gqa.py b/workloads/attention/gqa.py similarity index 100% rename from workloads/ops/gqa.py rename to workloads/attention/gqa.py diff --git a/workloads/ops/gqa_decode.py b/workloads/attention/gqa_decode.py similarity index 100% rename from workloads/ops/gqa_decode.py rename to workloads/attention/gqa_decode.py diff --git a/workloads/ops/gqa_decode_paged.py b/workloads/attention/gqa_decode_paged.py similarity index 100% rename from workloads/ops/gqa_decode_paged.py rename to workloads/attention/gqa_decode_paged.py diff --git a/workloads/ops/gqa_sliding_window_fwd.py b/workloads/attention/gqa_sliding_window.py similarity index 100% rename from workloads/ops/gqa_sliding_window_fwd.py rename to workloads/attention/gqa_sliding_window.py diff --git a/workloads/ops/gqa_sliding_window_varlen_fwd.py b/workloads/attention/gqa_sliding_window_varlen.py similarity index 100% rename from workloads/ops/gqa_sliding_window_varlen_fwd.py rename to workloads/attention/gqa_sliding_window_varlen.py diff --git a/workloads/ops/mean_pooling_ops.py b/workloads/attention/mean_pooling.py similarity index 100% rename from workloads/ops/mean_pooling_ops.py rename to workloads/attention/mean_pooling.py diff --git a/workloads/ops/mha.py b/workloads/attention/mha.py similarity index 100% rename from workloads/ops/mha.py rename to workloads/attention/mha.py diff --git a/workloads/ops/mha_decode.py b/workloads/attention/mha_decode.py similarity index 100% rename from workloads/ops/mha_decode.py rename to workloads/attention/mha_decode.py diff --git a/workloads/ops/mha_decode_paged.py b/workloads/attention/mha_decode_paged.py similarity index 100% rename from workloads/ops/mha_decode_paged.py rename to workloads/attention/mha_decode_paged.py diff --git a/workloads/ops/batch_norm.py b/workloads/batch_norm.py similarity index 100% rename from workloads/ops/batch_norm.py rename to workloads/batch_norm.py diff --git a/workloads/ops/binary_arith.py b/workloads/binary_arith.py similarity index 100% rename from workloads/ops/binary_arith.py rename to workloads/binary_arith.py diff --git a/workloads/ops/deltanet_fwd.py b/workloads/deltanet.py similarity index 53% rename from workloads/ops/deltanet_fwd.py rename to workloads/deltanet.py index f239e01c..b07556ea 100644 --- a/workloads/ops/deltanet_fwd.py +++ b/workloads/deltanet.py @@ -32,3 +32,29 @@ class DeltaNetFwdTest(WorkloadBase): v = torch.randn(B, H, S, DV, device="cuda", dtype=self.dtype) * 0.1 beta = torch.rand(B, H, S, device="cuda", dtype=self.dtype) * 0.5 return q, k, v, beta + + +class DeltaNetDecodeTest(WorkloadBase): + + def __init__( + self, + batch: int, + heads: int, + dim_k: int, + dim_v: int, + dtype: torch.dtype, + ) -> None: + self.batch = batch + self.heads = heads + self.dim_k = dim_k + self.dim_v = dim_v + self.dtype = dtype + + def gen_inputs(self) -> Tuple[torch.Tensor, ...]: + B, H, DK, DV = self.batch, self.heads, self.dim_k, self.dim_v + q = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 + k = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 + v = torch.randn(B, H, DV, device="cuda", dtype=self.dtype) * 0.1 + beta = torch.rand(B, H, device="cuda", dtype=self.dtype) * 0.5 + state = torch.randn(B, H, DK, DV, device="cuda", dtype=self.dtype) * 0.1 + return q, k, v, beta, state diff --git a/workloads/engram.py b/workloads/engram.py new file mode 100644 index 00000000..a7f453a3 --- /dev/null +++ b/workloads/engram.py @@ -0,0 +1,91 @@ +import torch + +from workloads.base import WorkloadBase + +CONV_KERNEL_SIZE = 4 + + +class EngramGateConvFwdTest(WorkloadBase): + def __init__(self, M, seq_len, d, dtype, eps=1e-6): + self.M = M + self.seq_len = seq_len + self.d = d + self.dtype = dtype + self.eps = eps + + def gen_inputs(self): + H = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") + k = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 + v = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 + rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") + rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") + conv_w = torch.randn(CONV_KERNEL_SIZE, self.d, dtype=self.dtype, device="cuda") * 0.02 + return H, k, v, rms_w_h, rms_w_v, conv_w + + +class EngramGateConvBwdTest(WorkloadBase): + def __init__(self, M, seq_len, d, dtype, eps=1e-6): + self.M = M + self.seq_len = seq_len + self.d = d + self.dtype = dtype + self.eps = eps + + def gen_inputs(self): + """Generate inputs including saved intermediates from a reference forward.""" + H = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") + k = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 + v = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 + rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") + rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") + conv_w = torch.randn(CONV_KERNEL_SIZE, self.d, dtype=self.dtype, device="cuda") * 0.02 + dY = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 + + # Compute saved intermediates via reference forward + def _rmsnorm(x, w): + x_f = x.float() + rrms = (x_f ** 2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() + return x_f * rrms * w.float(), rrms.squeeze(-1) + + h_norm, rrms_h = _rmsnorm(H, rms_w_h) + k_norm, rrms_k = _rmsnorm(k, rms_w_h) + dot = (h_norm * k_norm).sum(dim=-1, keepdim=True) + alpha = torch.sigmoid(dot / (self.d ** 0.5)) + v_hat = alpha * v.float() + _, rrms_v = _rmsnorm(v_hat.to(self.dtype), rms_w_v) + + vhat = v_hat.to(self.dtype) + alpha_squeezed = alpha.squeeze(-1).float() + rrms_h = rrms_h.float() + rrms_k = rrms_k.float() + rrms_v = rrms_v.float() + + return (dY, H, k, v, rms_w_h, rms_w_v, conv_w, + vhat, alpha_squeezed, rrms_h, rrms_k, rrms_v) + + +class EngramDecodeTest(WorkloadBase): + def __init__(self, batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, eps=1e-6): + self.batch = batch + self.d_mem = d_mem + self.d = d + self.max_conv_len = max_conv_len + self.conv_kernel_size = conv_kernel_size + self.dilation = dilation + self.dtype = dtype + self.eps = eps + + def gen_inputs(self): + e_t = torch.randn(self.batch, self.d_mem, dtype=self.dtype, device="cuda") * 0.1 + h_t = torch.randn(self.batch, self.d, dtype=self.dtype, device="cuda") + # Full conv_state (max_conv_len entries) + conv_state = torch.randn( + self.batch, self.max_conv_len, self.d, + dtype=self.dtype, device="cuda" + ) * 0.1 + W_K = torch.randn(self.d_mem, self.d, dtype=self.dtype, device="cuda") * 0.02 + W_V = torch.randn(self.d_mem, self.d, dtype=self.dtype, device="cuda") * 0.02 + rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") + rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") + conv_w = torch.randn(self.conv_kernel_size, self.d, dtype=self.dtype, device="cuda") * 0.02 + return e_t, h_t, conv_state, W_K, W_V, rms_w_h, rms_w_v, conv_w diff --git a/workloads/ops/fft.py b/workloads/fft.py similarity index 100% rename from workloads/ops/fft.py rename to workloads/fft.py diff --git a/workloads/ops/fp8_lighting_indexer.py b/workloads/fp8_lighting_indexer.py similarity index 100% rename from workloads/ops/fp8_lighting_indexer.py rename to workloads/fp8_lighting_indexer.py diff --git a/workloads/ops/fp8_quant.py b/workloads/fp8_quant.py similarity index 100% rename from workloads/ops/fp8_quant.py rename to workloads/fp8_quant.py diff --git a/workloads/ops/fused_add_layer_norm.py b/workloads/fused_add_layer_norm.py similarity index 100% rename from workloads/ops/fused_add_layer_norm.py rename to workloads/fused_add_layer_norm.py diff --git a/workloads/ops/fused_add_rmsnorm.py b/workloads/fused_add_rmsnorm.py similarity index 100% rename from workloads/ops/fused_add_rmsnorm.py rename to workloads/fused_add_rmsnorm.py diff --git a/workloads/ops/gated_deltanet_fwd.py b/workloads/gated_deltanet.py similarity index 53% rename from workloads/ops/gated_deltanet_fwd.py rename to workloads/gated_deltanet.py index 35f5d655..d208f273 100644 --- a/workloads/ops/gated_deltanet_fwd.py +++ b/workloads/gated_deltanet.py @@ -33,3 +33,30 @@ class GatedDeltaNetFwdTest(WorkloadBase): g = -torch.rand(B, H, S, device="cuda", dtype=self.dtype) beta = torch.rand(B, H, S, device="cuda", dtype=self.dtype) * 0.5 return q, k, v, g, beta + + +class GatedDeltaNetDecodeTest(WorkloadBase): + + def __init__( + self, + batch: int, + heads: int, + dim_k: int, + dim_v: int, + dtype: torch.dtype, + ) -> None: + self.batch = batch + self.heads = heads + self.dim_k = dim_k + self.dim_v = dim_v + self.dtype = dtype + + def gen_inputs(self) -> Tuple[torch.Tensor, ...]: + B, H, DK, DV = self.batch, self.heads, self.dim_k, self.dim_v + q = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 + k = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 + v = torch.randn(B, H, DV, device="cuda", dtype=self.dtype) * 0.1 + g = -torch.rand(B, H, device="cuda", dtype=self.dtype) + beta = torch.rand(B, H, device="cuda", dtype=self.dtype) * 0.5 + state = torch.randn(B, H, DK, DV, device="cuda", dtype=self.dtype) * 0.1 + return q, k, v, g, beta, state diff --git a/workloads/ops/gemm.py b/workloads/gemm.py similarity index 100% rename from workloads/ops/gemm.py rename to workloads/gemm.py diff --git a/workloads/ops/gla_recurrence.py b/workloads/gla.py similarity index 100% rename from workloads/ops/gla_recurrence.py rename to workloads/gla.py diff --git a/workloads/ops/group_norm.py b/workloads/group_norm.py similarity index 100% rename from workloads/ops/group_norm.py rename to workloads/group_norm.py diff --git a/workloads/ops/grouped_gemm.py b/workloads/grouped_gemm.py similarity index 100% rename from workloads/ops/grouped_gemm.py rename to workloads/grouped_gemm.py diff --git a/workloads/ops/instance_norm.py b/workloads/instance_norm.py similarity index 100% rename from workloads/ops/instance_norm.py rename to workloads/instance_norm.py diff --git a/workloads/ops/layer_norm.py b/workloads/layer_norm.py similarity index 100% rename from workloads/ops/layer_norm.py rename to workloads/layer_norm.py diff --git a/workloads/ops/logical_reduce.py b/workloads/logical_reduce.py similarity index 100% rename from workloads/ops/logical_reduce.py rename to workloads/logical_reduce.py diff --git a/workloads/mamba.py b/workloads/mamba.py new file mode 100644 index 00000000..fdd93bee --- /dev/null +++ b/workloads/mamba.py @@ -0,0 +1,237 @@ +import torch + +from workloads.base import FixtureBase, WorkloadBase + + +class DaCumsumFwdFixture(FixtureBase): + @classmethod + def get_params(cls): + import pytest + return [ + ("batch, num_chunks, chunk_len, n_heads, tune", [ + pytest.param(1, 2, 64, 4, False, marks=pytest.mark.smoke), + pytest.param(2, 4, 64, 8, False, marks=pytest.mark.full), + pytest.param(1, 2, 128, 4, False, marks=pytest.mark.full), + pytest.param(2, 4, 128, 16, False, marks=pytest.mark.full), + ]), + ] + +class DaCumsumFwdTest(WorkloadBase): + def __init__( + self, + batch: int, + num_chunks: int, + chunk_len: int, + n_heads: int, + ): + self.batch = batch + self.num_chunks = num_chunks + self.chunk_len = chunk_len + self.n_heads = n_heads + + def gen_inputs(self): + b, C, Q, h = self.batch, self.num_chunks, self.chunk_len, self.n_heads + seq_len = C * Q + # dt > 0 (softplus output in Mamba-2), A <= 0 (negative decay) + dt = torch.rand(b, seq_len, h, dtype=torch.float32, device="cuda") * 0.1 + 0.01 + A = -torch.rand(h, dtype=torch.float32, device="cuda") + return dt, A + + +class SsdChunkScanFwdFixture(FixtureBase): + @classmethod + def get_params(cls): + import pytest + return [ + ("batch, num_chunks, chunk_len, n_heads, d_head, d_state, dtype, tune", [ + pytest.param(1, 2, 64, 4, 64, 32, torch.float16, False, marks=pytest.mark.smoke), + pytest.param(2, 4, 64, 8, 64, 64, torch.float16, False, marks=pytest.mark.full), + pytest.param(1, 2, 128, 4, 128, 32, torch.bfloat16, False, marks=pytest.mark.full), + pytest.param(2, 2, 64, 4, 64, 32, torch.bfloat16, False, marks=pytest.mark.full), + ]), + ] + +class SsdChunkScanFwdTest(WorkloadBase): + def __init__( + self, + batch: int, + num_chunks: int, + chunk_len: int, + n_heads: int, + d_head: int, + d_state: int, + dtype: torch.dtype, + ): + self.batch = batch + self.num_chunks = num_chunks + self.chunk_len = chunk_len + self.n_heads = n_heads + self.d_head = d_head + self.d_state = d_state + self.dtype = dtype + + def gen_inputs(self): + b, c, L, h, p, n = ( + self.batch, self.num_chunks, self.chunk_len, + self.n_heads, self.d_head, self.d_state, + ) + x = torch.randn(b, c, L, h, p, dtype=self.dtype, device="cuda") * 0.1 + cb = torch.randn(b, c, h, L, L, dtype=self.dtype, device="cuda") * 0.1 + dA_cumsum = torch.zeros(b, h, c, L, dtype=torch.float32, device="cuda") + # fill with plausible negative cumsum values (decaying system) + dA_cumsum = -torch.rand(b, h, c, L, dtype=torch.float32, device="cuda").cumsum(-1) + C = torch.randn(b, c, L, h, n, dtype=self.dtype, device="cuda") * 0.1 + prev_states = torch.randn(b, c, h, n, p, dtype=self.dtype, device="cuda") * 0.1 + dt = torch.rand(b, c, L, h, dtype=self.dtype, device="cuda") * 0.1 + 0.01 + return x, cb, dA_cumsum, C, prev_states, dt + +class SsdChunkStateFwdFixture(FixtureBase): + @classmethod + def get_params(cls): + import pytest + return [ + ("batch, num_chunks, chunk_len, n_heads, d_head, d_state, n_groups, dtype, tune, has_seq_idx", [ + pytest.param( + 1, 2, 64, 4, 64, 32, 1, torch.float16, False, False, marks=pytest.mark.smoke, + ), + pytest.param( + 2, 4, 64, 8, 64, 64, 2, torch.float16, False, False, marks=pytest.mark.full, + ), + pytest.param( + 1, 2, 128, 4, 128, 32, 1, torch.bfloat16, False, False, marks=pytest.mark.full, + ), + pytest.param( + 2, 2, 64, 4, 64, 32, 2, torch.bfloat16, False, False, marks=pytest.mark.full, + ), + pytest.param( + 2, 4, 64, 8, 64, 64, 2, torch.float16, False, True, marks=pytest.mark.full, + ), + ]), + ] + +class SsdChunkStateFwdTest(WorkloadBase): + def __init__( + self, + batch: int, + num_chunks: int, + chunk_len: int, + n_heads: int, + d_head: int, + d_state: int, + n_groups: int, + dtype: torch.dtype, + has_seq_idx: bool = False, + ): + self.batch = batch + self.num_chunks = num_chunks + self.chunk_len = chunk_len + self.n_heads = n_heads + self.d_head = d_head + self.d_state = d_state + self.n_groups = n_groups + self.dtype = dtype + self.has_seq_idx = has_seq_idx + + def gen_inputs(self): + b, c, Q, h, p, n, g = ( + self.batch, self.num_chunks, self.chunk_len, + self.n_heads, self.d_head, self.d_state, self.n_groups, + ) + seq_len = c * Q + x = torch.randn(b, seq_len, h, p, dtype=self.dtype, device="cuda") * 0.1 + Bmat = torch.randn(b, seq_len, g, n, dtype=self.dtype, device="cuda") * 0.1 + # dA_cumsum: monotonically non-increasing (negative values, cumsum of negatives) + dA_cumsum = -torch.rand(b, h, c, Q, dtype=torch.float32, device="cuda").cumsum(-1) + dt = torch.rand(b, h, c, Q, dtype=torch.float32, device="cuda") * 0.1 + 0.01 + seq_idx = None + if self.has_seq_idx: + # simulate two packed sequences per batch row, split at midpoint + seq_idx = torch.zeros(b, seq_len, dtype=torch.int32, device="cuda") + seq_idx[:, seq_len // 2:] = 1 + return x, Bmat, dt, dA_cumsum, seq_idx + +class SsdDecodeFixture(FixtureBase): + @classmethod + def get_params(cls): + import pytest + return [ + ("batch, n_heads, d_head, d_state, n_groups, dtype, tune", [ + pytest.param( + 1, 4, 64, 16, 1, torch.float16, False, marks=pytest.mark.smoke, + ), + pytest.param( + 2, 8, 64, 32, 2, torch.float16, False, marks=pytest.mark.full, + ), + pytest.param( + 1, 4, 64, 16, 1, torch.bfloat16, False, marks=pytest.mark.full, + ), + pytest.param( + 2, 8, 128, 64, 4, torch.bfloat16, False, marks=pytest.mark.full, + ), + ]), + ] + +class SsdDecodeTest(WorkloadBase): + def __init__( + self, + batch: int, + n_heads: int, + d_head: int, + d_state: int, + n_groups: int, + dtype: torch.dtype, + ): + self.batch = batch + self.n_heads = n_heads + self.d_head = d_head + self.d_state = d_state + self.n_groups = n_groups + self.dtype = dtype + + def gen_inputs(self): + b, h, p, n, g = ( + self.batch, self.n_heads, self.d_head, self.d_state, self.n_groups, + ) + # A <= 0 (negative decay), dt > 0 (post-softplus) + A = -torch.rand(h, dtype=torch.float32, device="cuda") + dt = torch.rand(b, h, dtype=torch.float32, device="cuda") * 0.1 + 0.01 + x = torch.randn(b, h, p, dtype=self.dtype, device="cuda") * 0.1 + B_in = torch.randn(b, g, n, dtype=self.dtype, device="cuda") * 0.1 + C_in = torch.randn(b, g, n, dtype=self.dtype, device="cuda") * 0.1 + state = torch.randn(b, h, p, n, dtype=torch.float32, device="cuda") * 0.1 + return A, dt, x, B_in, C_in, state + +class SsdStatePassingFwdFixture(FixtureBase): + @classmethod + def get_params(cls): + import pytest + return [ + ("batch, num_chunks, n_heads, d_state, dtype, tune", [ + pytest.param(1, 2, 4, 32, torch.float16, False, marks=pytest.mark.smoke), + pytest.param(2, 4, 8, 64, torch.float16, False, marks=pytest.mark.full), + pytest.param(1, 2, 4, 32, torch.bfloat16, False, marks=pytest.mark.full), + pytest.param(2, 4, 8, 64, torch.bfloat16, False, marks=pytest.mark.full), + ]), + ] + +class SsdStatePassingFwdTest(WorkloadBase): + def __init__( + self, + batch: int, + num_chunks: int, + n_heads: int, + d_state: int, + dtype: torch.dtype, + ): + self.batch = batch + self.num_chunks = num_chunks + self.n_heads = n_heads + self.d_state = d_state + self.dtype = dtype + + def gen_inputs(self): + b, c, h, d = self.batch, self.num_chunks, self.n_heads, self.d_state + states = torch.randn(b, c, h, d, dtype=self.dtype, device="cuda") * 0.1 + dA_chunk_cumsum = -torch.rand(b, h, c, dtype=torch.float32, device="cuda").cumsum(-1) + initial_states = torch.randn(b, h, d, dtype=torch.float32, device="cuda") * 0.1 + return states, dA_chunk_cumsum, initial_states diff --git a/workloads/ops/mhc_pre.py b/workloads/mhc.py similarity index 62% rename from workloads/ops/mhc_pre.py rename to workloads/mhc.py index aab175d8..6f59e54e 100644 --- a/workloads/ops/mhc_pre.py +++ b/workloads/mhc.py @@ -1,4 +1,3 @@ - from typing import Tuple import torch @@ -31,3 +30,22 @@ class MHCPreTest(WorkloadBase): sinkhorn_repeat = 20 eps = 0.02 return phi, x, b, alpha_pre, alpha_post, alpha_res, sinkhorn_repeat, eps + + +class MHCPostTest(WorkloadBase): + + def __init__(self, batch: int, n_expand: int, c_x: int, dtype: torch.dtype): + self.batch = batch + self.n_expand = n_expand + self.c_x = c_x + self.dtype = dtype + + def gen_inputs(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch = self.batch + n_expand = self.n_expand + c_x = self.c_x + + x_layer_out = torch.randn([batch, c_x], device="cuda", dtype=self.dtype) + h_post = torch.randn([batch, n_expand], device="cuda", dtype=torch.float32) + x_res = torch.randn([batch, n_expand * c_x], device="cuda", dtype=self.dtype) + return x_layer_out, h_post, x_res diff --git a/workloads/moe.py b/workloads/moe.py new file mode 100644 index 00000000..18cda49a --- /dev/null +++ b/workloads/moe.py @@ -0,0 +1,79 @@ +from typing import Tuple + +import torch + +from workloads.base import WorkloadBase + + +class FusedTopKTest(WorkloadBase): + def __init__(self, num_tokens, num_experts, top_k, scoring_func, renormalize, dtype): + self.num_tokens = num_tokens + self.num_experts = num_experts + self.top_k = top_k + self.scoring_func = scoring_func + self.renormalize = renormalize + self.dtype = dtype + + def gen_inputs(self): + torch.manual_seed(42) + return torch.randn(self.num_tokens, self.num_experts, dtype=self.dtype, device="cuda") + + +class MoePermuteTest(WorkloadBase): + + def __init__(self, total_tokens, top_k, num_experts, hidden_size, dtype): + self.total_tokens = total_tokens + self.top_k = top_k + self.num_experts = num_experts + self.hidden_size = hidden_size + self.dtype = dtype + + def gen_inputs(self): + hidden_states = torch.randn( + self.total_tokens, self.hidden_size, dtype=self.dtype, device="cuda" + ) + topk_ids = torch.randint( + 0, self.num_experts, + (self.total_tokens, self.top_k), + dtype=torch.int32, device="cuda", + ) + return hidden_states, topk_ids + + +class MoePermuteAlignTest(WorkloadBase): + + def __init__(self, total_tokens: int, top_k: int, num_experts: int, block_size: int): + self.total_tokens = total_tokens + self.top_k = top_k + self.num_experts = num_experts + self.block_size = block_size + + def gen_inputs(self) -> Tuple[torch.Tensor]: + topk_ids = torch.randint( + 0, self.num_experts, + (self.total_tokens, self.top_k), + dtype=torch.int32, device="cuda", + ) + return (topk_ids,) + + +class MoeUnpermuteTest(WorkloadBase): + + def __init__(self, total_tokens, top_k, hidden_size, dtype): + self.total_tokens = total_tokens + self.top_k = top_k + self.hidden_size = hidden_size + self.dtype = dtype + # Use padded_batch_sum = T*K (no actual padding) for standalone tests. + self.padded_batch_sum = total_tokens * top_k + + def gen_inputs(self): + numel = self.total_tokens * self.top_k + mm2_pad = torch.randn(numel, self.hidden_size, dtype=self.dtype, device="cuda") + # fwd_idx: each flat_idx maps to a padded_slot in [0, padded_batch_sum) + # simulate a valid mapping: random shuffle of [0, numel) + fwd_idx = torch.randperm(numel, dtype=torch.int32, device="cuda") + topk_weights = torch.rand( + self.total_tokens, self.top_k, dtype=torch.float32, device="cuda" + ) + return mm2_pad, fwd_idx, topk_weights diff --git a/workloads/ops/da_cumsum_fwd.py b/workloads/ops/da_cumsum_fwd.py deleted file mode 100644 index 984da895..00000000 --- a/workloads/ops/da_cumsum_fwd.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch - -from workloads.base import FixtureBase, WorkloadBase - - -class DaCumsumFwdFixture(FixtureBase): - @classmethod - def get_params(cls): - import pytest - return [ - ("batch, num_chunks, chunk_len, n_heads, tune", [ - pytest.param(1, 2, 64, 4, False, marks=pytest.mark.smoke), - pytest.param(2, 4, 64, 8, False, marks=pytest.mark.full), - pytest.param(1, 2, 128, 4, False, marks=pytest.mark.full), - pytest.param(2, 4, 128, 16, False, marks=pytest.mark.full), - ]), - ] - -class DaCumsumFwdTest(WorkloadBase): - def __init__( - self, - batch: int, - num_chunks: int, - chunk_len: int, - n_heads: int, - ): - self.batch = batch - self.num_chunks = num_chunks - self.chunk_len = chunk_len - self.n_heads = n_heads - - def gen_inputs(self): - b, C, Q, h = self.batch, self.num_chunks, self.chunk_len, self.n_heads - seq_len = C * Q - # dt > 0 (softplus output in Mamba-2), A <= 0 (negative decay) - dt = torch.rand(b, seq_len, h, dtype=torch.float32, device="cuda") * 0.1 + 0.01 - A = -torch.rand(h, dtype=torch.float32, device="cuda") - return dt, A diff --git a/workloads/ops/deltanet_bwd.py b/workloads/ops/deltanet_bwd.py deleted file mode 100644 index 51d38c68..00000000 --- a/workloads/ops/deltanet_bwd.py +++ /dev/null @@ -1,3 +0,0 @@ -# This module previously contained oracle/reference functions (_differentiable_fwd, -# deltanet_autograd_bwd_torch). Those have been moved to their respective test and -# benchmark files as part of the trust-model separation (no shared oracle surface). diff --git a/workloads/ops/deltanet_recurrence.py b/workloads/ops/deltanet_recurrence.py deleted file mode 100644 index a59aeb2f..00000000 --- a/workloads/ops/deltanet_recurrence.py +++ /dev/null @@ -1,31 +0,0 @@ -from typing import Tuple - -import torch - -from workloads.base import WorkloadBase - - -class DeltaNetDecodeTest(WorkloadBase): - - def __init__( - self, - batch: int, - heads: int, - dim_k: int, - dim_v: int, - dtype: torch.dtype, - ) -> None: - self.batch = batch - self.heads = heads - self.dim_k = dim_k - self.dim_v = dim_v - self.dtype = dtype - - def gen_inputs(self) -> Tuple[torch.Tensor, ...]: - B, H, DK, DV = self.batch, self.heads, self.dim_k, self.dim_v - q = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 - k = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 - v = torch.randn(B, H, DV, device="cuda", dtype=self.dtype) * 0.1 - beta = torch.rand(B, H, device="cuda", dtype=self.dtype) * 0.5 - state = torch.randn(B, H, DK, DV, device="cuda", dtype=self.dtype) * 0.1 - return q, k, v, beta, state diff --git a/workloads/ops/engram_bwd.py b/workloads/ops/engram_bwd.py deleted file mode 100644 index 3152fb45..00000000 --- a/workloads/ops/engram_bwd.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - -CONV_KERNEL_SIZE = 4 - - -class EngramGateConvBwdTest(WorkloadBase): - def __init__(self, M, seq_len, d, dtype, eps=1e-6): - self.M = M - self.seq_len = seq_len - self.d = d - self.dtype = dtype - self.eps = eps - - def gen_inputs(self): - """Generate inputs including saved intermediates from a reference forward.""" - H = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") - k = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 - v = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 - rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") - rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") - conv_w = torch.randn(CONV_KERNEL_SIZE, self.d, dtype=self.dtype, device="cuda") * 0.02 - dY = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 - - # Compute saved intermediates via reference forward - def _rmsnorm(x, w): - x_f = x.float() - rrms = (x_f ** 2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() - return x_f * rrms * w.float(), rrms.squeeze(-1) - - h_norm, rrms_h = _rmsnorm(H, rms_w_h) - k_norm, rrms_k = _rmsnorm(k, rms_w_h) - dot = (h_norm * k_norm).sum(dim=-1, keepdim=True) - alpha = torch.sigmoid(dot / (self.d ** 0.5)) - v_hat = alpha * v.float() - _, rrms_v = _rmsnorm(v_hat.to(self.dtype), rms_w_v) - - vhat = v_hat.to(self.dtype) - alpha_squeezed = alpha.squeeze(-1).float() - rrms_h = rrms_h.float() - rrms_k = rrms_k.float() - rrms_v = rrms_v.float() - - return (dY, H, k, v, rms_w_h, rms_w_v, conv_w, - vhat, alpha_squeezed, rrms_h, rrms_k, rrms_v) diff --git a/workloads/ops/engram_decode.py b/workloads/ops/engram_decode.py deleted file mode 100644 index bc5963c3..00000000 --- a/workloads/ops/engram_decode.py +++ /dev/null @@ -1,30 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - - -class EngramDecodeTest(WorkloadBase): - def __init__(self, batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, eps=1e-6): - self.batch = batch - self.d_mem = d_mem - self.d = d - self.max_conv_len = max_conv_len - self.conv_kernel_size = conv_kernel_size - self.dilation = dilation - self.dtype = dtype - self.eps = eps - - def gen_inputs(self): - e_t = torch.randn(self.batch, self.d_mem, dtype=self.dtype, device="cuda") * 0.1 - h_t = torch.randn(self.batch, self.d, dtype=self.dtype, device="cuda") - # Full conv_state (max_conv_len entries) - conv_state = torch.randn( - self.batch, self.max_conv_len, self.d, - dtype=self.dtype, device="cuda" - ) * 0.1 - W_K = torch.randn(self.d_mem, self.d, dtype=self.dtype, device="cuda") * 0.02 - W_V = torch.randn(self.d_mem, self.d, dtype=self.dtype, device="cuda") * 0.02 - rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") - rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") - conv_w = torch.randn(self.conv_kernel_size, self.d, dtype=self.dtype, device="cuda") * 0.02 - return e_t, h_t, conv_state, W_K, W_V, rms_w_h, rms_w_v, conv_w diff --git a/workloads/ops/engram_fwd.py b/workloads/ops/engram_fwd.py deleted file mode 100644 index c9515bbd..00000000 --- a/workloads/ops/engram_fwd.py +++ /dev/null @@ -1,23 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - -CONV_KERNEL_SIZE = 4 - - -class EngramGateConvFwdTest(WorkloadBase): - def __init__(self, M, seq_len, d, dtype, eps=1e-6): - self.M = M - self.seq_len = seq_len - self.d = d - self.dtype = dtype - self.eps = eps - - def gen_inputs(self): - H = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") - k = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 - v = torch.randn(self.M, self.seq_len, self.d, dtype=self.dtype, device="cuda") * 0.1 - rms_w_h = torch.ones(self.d, dtype=self.dtype, device="cuda") - rms_w_v = torch.ones(self.d, dtype=self.dtype, device="cuda") - conv_w = torch.randn(CONV_KERNEL_SIZE, self.d, dtype=self.dtype, device="cuda") * 0.02 - return H, k, v, rms_w_h, rms_w_v, conv_w diff --git a/workloads/ops/gated_deltanet_bwd.py b/workloads/ops/gated_deltanet_bwd.py deleted file mode 100644 index 81e88af6..00000000 --- a/workloads/ops/gated_deltanet_bwd.py +++ /dev/null @@ -1,3 +0,0 @@ -# This module previously contained oracle/reference functions (_differentiable_fwd, -# gated_deltanet_autograd_bwd_torch). Those have been moved to their respective test -# and benchmark files as part of the trust-model separation (no shared oracle surface). diff --git a/workloads/ops/gated_deltanet_recurrence.py b/workloads/ops/gated_deltanet_recurrence.py deleted file mode 100644 index ab40e9d2..00000000 --- a/workloads/ops/gated_deltanet_recurrence.py +++ /dev/null @@ -1,32 +0,0 @@ -from typing import Tuple - -import torch - -from workloads.base import WorkloadBase - - -class GatedDeltaNetDecodeTest(WorkloadBase): - - def __init__( - self, - batch: int, - heads: int, - dim_k: int, - dim_v: int, - dtype: torch.dtype, - ) -> None: - self.batch = batch - self.heads = heads - self.dim_k = dim_k - self.dim_v = dim_v - self.dtype = dtype - - def gen_inputs(self) -> Tuple[torch.Tensor, ...]: - B, H, DK, DV = self.batch, self.heads, self.dim_k, self.dim_v - q = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 - k = torch.randn(B, H, DK, device="cuda", dtype=self.dtype) * 0.1 - v = torch.randn(B, H, DV, device="cuda", dtype=self.dtype) * 0.1 - g = -torch.rand(B, H, device="cuda", dtype=self.dtype) - beta = torch.rand(B, H, device="cuda", dtype=self.dtype) * 0.5 - state = torch.randn(B, H, DK, DV, device="cuda", dtype=self.dtype) * 0.1 - return q, k, v, g, beta, state diff --git a/workloads/ops/gla_bwd.py b/workloads/ops/gla_bwd.py deleted file mode 100644 index 113cbe8d..00000000 --- a/workloads/ops/gla_bwd.py +++ /dev/null @@ -1,3 +0,0 @@ -# This module previously contained oracle/reference functions (gla_autograd_bwd_torch, -# gla_fwd_chunked_torch). Those have been moved to their respective test and benchmark -# files as part of the trust-model separation (no shared oracle surface). diff --git a/workloads/ops/mhc_post.py b/workloads/ops/mhc_post.py deleted file mode 100644 index 19be0a7b..00000000 --- a/workloads/ops/mhc_post.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import Tuple - -import torch - -from workloads.base import WorkloadBase - - -class MHCPostTest(WorkloadBase): - - def __init__(self, batch: int, n_expand: int, c_x: int, dtype: torch.dtype): - self.batch = batch - self.n_expand = n_expand - self.c_x = c_x - self.dtype = dtype - - def gen_inputs(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - batch = self.batch - n_expand = self.n_expand - c_x = self.c_x - - x_layer_out = torch.randn([batch, c_x], device="cuda", dtype=self.dtype) - h_post = torch.randn([batch, n_expand], device="cuda", dtype=torch.float32) - x_res = torch.randn([batch, n_expand * c_x], device="cuda", dtype=self.dtype) - return x_layer_out, h_post, x_res diff --git a/workloads/ops/moe_fused_topk.py b/workloads/ops/moe_fused_topk.py deleted file mode 100644 index c6a43476..00000000 --- a/workloads/ops/moe_fused_topk.py +++ /dev/null @@ -1,17 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - - -class FusedTopKTest(WorkloadBase): - def __init__(self, num_tokens, num_experts, top_k, scoring_func, renormalize, dtype): - self.num_tokens = num_tokens - self.num_experts = num_experts - self.top_k = top_k - self.scoring_func = scoring_func - self.renormalize = renormalize - self.dtype = dtype - - def gen_inputs(self): - torch.manual_seed(42) - return torch.randn(self.num_tokens, self.num_experts, dtype=self.dtype, device="cuda") diff --git a/workloads/ops/moe_permute.py b/workloads/ops/moe_permute.py deleted file mode 100644 index a0d4e2bc..00000000 --- a/workloads/ops/moe_permute.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - - -class MoePermuteTest(WorkloadBase): - - def __init__(self, total_tokens, top_k, num_experts, hidden_size, dtype): - self.total_tokens = total_tokens - self.top_k = top_k - self.num_experts = num_experts - self.hidden_size = hidden_size - self.dtype = dtype - - def gen_inputs(self): - hidden_states = torch.randn( - self.total_tokens, self.hidden_size, dtype=self.dtype, device="cuda" - ) - topk_ids = torch.randint( - 0, self.num_experts, - (self.total_tokens, self.top_k), - dtype=torch.int32, device="cuda", - ) - return hidden_states, topk_ids diff --git a/workloads/ops/moe_permute_align.py b/workloads/ops/moe_permute_align.py deleted file mode 100644 index 1314cdd6..00000000 --- a/workloads/ops/moe_permute_align.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import Tuple - -import torch - -from workloads.base import WorkloadBase - - -class MoePermuteAlignTest(WorkloadBase): - - def __init__(self, total_tokens: int, top_k: int, num_experts: int, block_size: int): - self.total_tokens = total_tokens - self.top_k = top_k - self.num_experts = num_experts - self.block_size = block_size - - def gen_inputs(self) -> Tuple[torch.Tensor]: - topk_ids = torch.randint( - 0, self.num_experts, - (self.total_tokens, self.top_k), - dtype=torch.int32, device="cuda", - ) - return (topk_ids,) diff --git a/workloads/ops/moe_unpermute.py b/workloads/ops/moe_unpermute.py deleted file mode 100644 index b1a961fc..00000000 --- a/workloads/ops/moe_unpermute.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch - -from workloads.base import WorkloadBase - - -class MoeUnpermuteTest(WorkloadBase): - - def __init__(self, total_tokens, top_k, hidden_size, dtype): - self.total_tokens = total_tokens - self.top_k = top_k - self.hidden_size = hidden_size - self.dtype = dtype - # Use padded_batch_sum = T*K (no actual padding) for standalone tests. - self.padded_batch_sum = total_tokens * top_k - - def gen_inputs(self): - numel = self.total_tokens * self.top_k - mm2_pad = torch.randn(numel, self.hidden_size, dtype=self.dtype, device="cuda") - # fwd_idx: each flat_idx maps to a padded_slot in [0, padded_batch_sum) - # simulate a valid mapping: random shuffle of [0, numel) - fwd_idx = torch.randperm(numel, dtype=torch.int32, device="cuda") - topk_weights = torch.rand( - self.total_tokens, self.top_k, dtype=torch.float32, device="cuda" - ) - return mm2_pad, fwd_idx, topk_weights diff --git a/workloads/ops/ssd_chunk_scan_fwd.py b/workloads/ops/ssd_chunk_scan_fwd.py deleted file mode 100644 index 77c64b93..00000000 --- a/workloads/ops/ssd_chunk_scan_fwd.py +++ /dev/null @@ -1,51 +0,0 @@ -import torch - -from workloads.base import FixtureBase, WorkloadBase - - -class SsdChunkScanFwdFixture(FixtureBase): - @classmethod - def get_params(cls): - import pytest - return [ - ("batch, num_chunks, chunk_len, n_heads, d_head, d_state, dtype, tune", [ - pytest.param(1, 2, 64, 4, 64, 32, torch.float16, False, marks=pytest.mark.smoke), - pytest.param(2, 4, 64, 8, 64, 64, torch.float16, False, marks=pytest.mark.full), - pytest.param(1, 2, 128, 4, 128, 32, torch.bfloat16, False, marks=pytest.mark.full), - pytest.param(2, 2, 64, 4, 64, 32, torch.bfloat16, False, marks=pytest.mark.full), - ]), - ] - -class SsdChunkScanFwdTest(WorkloadBase): - def __init__( - self, - batch: int, - num_chunks: int, - chunk_len: int, - n_heads: int, - d_head: int, - d_state: int, - dtype: torch.dtype, - ): - self.batch = batch - self.num_chunks = num_chunks - self.chunk_len = chunk_len - self.n_heads = n_heads - self.d_head = d_head - self.d_state = d_state - self.dtype = dtype - - def gen_inputs(self): - b, c, L, h, p, n = ( - self.batch, self.num_chunks, self.chunk_len, - self.n_heads, self.d_head, self.d_state, - ) - x = torch.randn(b, c, L, h, p, dtype=self.dtype, device="cuda") * 0.1 - cb = torch.randn(b, c, h, L, L, dtype=self.dtype, device="cuda") * 0.1 - dA_cumsum = torch.zeros(b, h, c, L, dtype=torch.float32, device="cuda") - # fill with plausible negative cumsum values (decaying system) - dA_cumsum = -torch.rand(b, h, c, L, dtype=torch.float32, device="cuda").cumsum(-1) - C = torch.randn(b, c, L, h, n, dtype=self.dtype, device="cuda") * 0.1 - prev_states = torch.randn(b, c, h, n, p, dtype=self.dtype, device="cuda") * 0.1 - dt = torch.rand(b, c, L, h, dtype=self.dtype, device="cuda") * 0.1 + 0.01 - return x, cb, dA_cumsum, C, prev_states, dt diff --git a/workloads/ops/ssd_chunk_state_fwd.py b/workloads/ops/ssd_chunk_state_fwd.py deleted file mode 100644 index b5084190..00000000 --- a/workloads/ops/ssd_chunk_state_fwd.py +++ /dev/null @@ -1,69 +0,0 @@ -import torch - -from workloads.base import FixtureBase, WorkloadBase - - -class SsdChunkStateFwdFixture(FixtureBase): - @classmethod - def get_params(cls): - import pytest - return [ - ("batch, num_chunks, chunk_len, n_heads, d_head, d_state, n_groups, dtype, tune, has_seq_idx", [ - pytest.param( - 1, 2, 64, 4, 64, 32, 1, torch.float16, False, False, marks=pytest.mark.smoke, - ), - pytest.param( - 2, 4, 64, 8, 64, 64, 2, torch.float16, False, False, marks=pytest.mark.full, - ), - pytest.param( - 1, 2, 128, 4, 128, 32, 1, torch.bfloat16, False, False, marks=pytest.mark.full, - ), - pytest.param( - 2, 2, 64, 4, 64, 32, 2, torch.bfloat16, False, False, marks=pytest.mark.full, - ), - pytest.param( - 2, 4, 64, 8, 64, 64, 2, torch.float16, False, True, marks=pytest.mark.full, - ), - ]), - ] - -class SsdChunkStateFwdTest(WorkloadBase): - def __init__( - self, - batch: int, - num_chunks: int, - chunk_len: int, - n_heads: int, - d_head: int, - d_state: int, - n_groups: int, - dtype: torch.dtype, - has_seq_idx: bool = False, - ): - self.batch = batch - self.num_chunks = num_chunks - self.chunk_len = chunk_len - self.n_heads = n_heads - self.d_head = d_head - self.d_state = d_state - self.n_groups = n_groups - self.dtype = dtype - self.has_seq_idx = has_seq_idx - - def gen_inputs(self): - b, c, Q, h, p, n, g = ( - self.batch, self.num_chunks, self.chunk_len, - self.n_heads, self.d_head, self.d_state, self.n_groups, - ) - seq_len = c * Q - x = torch.randn(b, seq_len, h, p, dtype=self.dtype, device="cuda") * 0.1 - Bmat = torch.randn(b, seq_len, g, n, dtype=self.dtype, device="cuda") * 0.1 - # dA_cumsum: monotonically non-increasing (negative values, cumsum of negatives) - dA_cumsum = -torch.rand(b, h, c, Q, dtype=torch.float32, device="cuda").cumsum(-1) - dt = torch.rand(b, h, c, Q, dtype=torch.float32, device="cuda") * 0.1 + 0.01 - seq_idx = None - if self.has_seq_idx: - # simulate two packed sequences per batch row, split at midpoint - seq_idx = torch.zeros(b, seq_len, dtype=torch.int32, device="cuda") - seq_idx[:, seq_len // 2:] = 1 - return x, Bmat, dt, dA_cumsum, seq_idx diff --git a/workloads/ops/ssd_decode.py b/workloads/ops/ssd_decode.py deleted file mode 100644 index b8d1f2dc..00000000 --- a/workloads/ops/ssd_decode.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch - -from workloads.base import FixtureBase, WorkloadBase - - -class SsdDecodeFixture(FixtureBase): - @classmethod - def get_params(cls): - import pytest - return [ - ("batch, n_heads, d_head, d_state, n_groups, dtype, tune", [ - pytest.param( - 1, 4, 64, 16, 1, torch.float16, False, marks=pytest.mark.smoke, - ), - pytest.param( - 2, 8, 64, 32, 2, torch.float16, False, marks=pytest.mark.full, - ), - pytest.param( - 1, 4, 64, 16, 1, torch.bfloat16, False, marks=pytest.mark.full, - ), - pytest.param( - 2, 8, 128, 64, 4, torch.bfloat16, False, marks=pytest.mark.full, - ), - ]), - ] - -class SsdDecodeTest(WorkloadBase): - def __init__( - self, - batch: int, - n_heads: int, - d_head: int, - d_state: int, - n_groups: int, - dtype: torch.dtype, - ): - self.batch = batch - self.n_heads = n_heads - self.d_head = d_head - self.d_state = d_state - self.n_groups = n_groups - self.dtype = dtype - - def gen_inputs(self): - b, h, p, n, g = ( - self.batch, self.n_heads, self.d_head, self.d_state, self.n_groups, - ) - # A <= 0 (negative decay), dt > 0 (post-softplus) - A = -torch.rand(h, dtype=torch.float32, device="cuda") - dt = torch.rand(b, h, dtype=torch.float32, device="cuda") * 0.1 + 0.01 - x = torch.randn(b, h, p, dtype=self.dtype, device="cuda") * 0.1 - B_in = torch.randn(b, g, n, dtype=self.dtype, device="cuda") * 0.1 - C_in = torch.randn(b, g, n, dtype=self.dtype, device="cuda") * 0.1 - state = torch.randn(b, h, p, n, dtype=torch.float32, device="cuda") * 0.1 - return A, dt, x, B_in, C_in, state diff --git a/workloads/ops/ssd_state_passing_fwd.py b/workloads/ops/ssd_state_passing_fwd.py deleted file mode 100644 index 7858e643..00000000 --- a/workloads/ops/ssd_state_passing_fwd.py +++ /dev/null @@ -1,39 +0,0 @@ -import torch - -from workloads.base import FixtureBase, WorkloadBase - - -class SsdStatePassingFwdFixture(FixtureBase): - @classmethod - def get_params(cls): - import pytest - return [ - ("batch, num_chunks, n_heads, d_state, dtype, tune", [ - pytest.param(1, 2, 4, 32, torch.float16, False, marks=pytest.mark.smoke), - pytest.param(2, 4, 8, 64, torch.float16, False, marks=pytest.mark.full), - pytest.param(1, 2, 4, 32, torch.bfloat16, False, marks=pytest.mark.full), - pytest.param(2, 4, 8, 64, torch.bfloat16, False, marks=pytest.mark.full), - ]), - ] - -class SsdStatePassingFwdTest(WorkloadBase): - def __init__( - self, - batch: int, - num_chunks: int, - n_heads: int, - d_state: int, - dtype: torch.dtype, - ): - self.batch = batch - self.num_chunks = num_chunks - self.n_heads = n_heads - self.d_state = d_state - self.dtype = dtype - - def gen_inputs(self): - b, c, h, d = self.batch, self.num_chunks, self.n_heads, self.d_state - states = torch.randn(b, c, h, d, dtype=self.dtype, device="cuda") * 0.1 - dA_chunk_cumsum = -torch.rand(b, h, c, dtype=torch.float32, device="cuda").cumsum(-1) - initial_states = torch.randn(b, h, d, dtype=torch.float32, device="cuda") * 0.1 - return states, dA_chunk_cumsum, initial_states diff --git a/workloads/ops/reduce.py b/workloads/reduce.py similarity index 100% rename from workloads/ops/reduce.py rename to workloads/reduce.py diff --git a/workloads/ops/rms_norm.py b/workloads/rms_norm.py similarity index 100% rename from workloads/ops/rms_norm.py rename to workloads/rms_norm.py diff --git a/workloads/ops/softmax.py b/workloads/softmax.py similarity index 100% rename from workloads/ops/softmax.py rename to workloads/softmax.py diff --git a/workloads/ops/topk_selector.py b/workloads/topk_selector.py similarity index 100% rename from workloads/ops/topk_selector.py rename to workloads/topk_selector.py diff --git a/workloads/ops/vector_norm.py b/workloads/vector_norm.py similarity index 100% rename from workloads/ops/vector_norm.py rename to workloads/vector_norm.py