[Refactor][Ops] read dtype at forward instead of taking it in op constructors

.claude/skills/implement-op/SKILL.md

@@ -92,7 +92,7 @@ Record `observations` — design knowledge discovered during migration:
 - Edge cases
 - Abstraction opportunities
 
-Do NOT modify manifest or design docs. Observations are returned to orchestrator and surfaced in PR for human review.
+Do not change the manifest or a design doc to match code that does not conform — the spec is the reference, not a record of the implementation. Observations are returned to the orchestrator and surfaced in the PR.
 
 ### 6. COMMIT
 

benchmarks/ops/attention/bench_deepseek_dsa_decode.py

@@ -81,7 +81,7 @@ def test_dsa_decode_bench(batch: int, heads: int, seq_len_q: int, seq_len_kv: in
 
     op = DeepSeekSparseAttentionDecodeWithKVCacheFwdOp(
         batch, heads, seq_len_q, seq_len_kv, dim, dim_tail, topk, stride_kv, heads_kv,
-        q_start_index_s, sm_scale=sm_scale, dtype=dtype, tune=tune)
+        q_start_index_s, sm_scale=sm_scale, tune=tune)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/attention/bench_deepseek_mla_decode.py

@@ -72,7 +72,7 @@ def test_mla_decode_bench(batch: int, heads: int, heads_kv: int, seq_len_kv: int
     inputs = test.gen_inputs()
 
     op = MultiHeadLatentAttentionDecodeWithKVCacheFwdOp(
-        batch, heads, heads_kv, seq_len_kv, dim, dim_pe, dtype, tune=tune)
+        batch, heads, heads_kv, seq_len_kv, dim, dim_pe, tune=tune)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/attention/bench_gqa.py

@@ -239,8 +239,6 @@ def _tileops_gqa_variant(op: GroupedQueryAttentionFwdOp) -> str:
         return "ws_noncausal"
     if isinstance(kernel, GQAFwdWgmmaPipelinedKernel):
         return "wgmma_pipelined"
-    if isinstance(kernel, GQAFwdKernel):
-        return "legacy"
     return kernel.__class__.__name__
 
 
@@ -322,7 +320,7 @@ def test_gqa_bwd_bench(
     test = GroupedQueryAttentionBwdWorkload(batch, heads, heads_kv, seq_len, dim, causal, dtype)
     inputs = test.gen_inputs()
 
-    op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, dtype, tune=tune)
+    op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, tune=tune)
     bm = ManifestBenchmark(_GQA_BWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -457,7 +455,7 @@ def test_gqa_prefill_varlen_fwd_bench(
     inputs = test.gen_inputs()
 
     op = GroupedQueryAttentionPrefillVarlenFwdOp(
-        batch, heads, heads_kv, dim, test.max_seqlen_q, test.max_seqlen_kv, causal, dtype, tune=tune
+        batch, heads, heads_kv, dim, test.max_seqlen_q, test.max_seqlen_kv, causal, tune=tune
     )
     bm = GQAPrefillVarlenFwdBenchmark(test)
     result = bm.profile(op, *inputs)

benchmarks/ops/attention/bench_gqa_decode.py

@@ -145,7 +145,6 @@ def test_gqa_decode_bench(batch: int, heads: int, heads_kv: int, seq_len_kv: int
         heads_kv,
         seq_len_kv,
         dim,
-        dtype,
         sm_scale=sm_scale,
         softcap=softcap,
         tune=tune,

benchmarks/ops/attention/bench_gqa_decode_paged.py

@@ -150,7 +150,7 @@ def test_gqa_decode_paged_bench(batch: int, heads: int, heads_kv: int, seqlen_kv
     q, k, v, real_seqlen_kv, block_table = inputs
 
     op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp(
-        batch, heads, heads_kv, seqlen_kv, dim, page_size, dtype,
+        batch, heads, heads_kv, seqlen_kv, dim, page_size,
         sm_scale=sm_scale, softcap=softcap, tune=tune)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     result = bm.profile(op, *inputs)

benchmarks/ops/attention/bench_gqa_sliding_window.py

@@ -133,7 +133,6 @@ def test_gqa_sliding_window_fwd_bench(
         is_causal=is_causal,
         window_size_left=wl,
         window_size_right=wr,
-        dtype=dtype,
         tune=tune,
     )
     bm = ManifestBenchmark(_OP_NAME, op, test)

benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py

@@ -167,7 +167,6 @@ def test_gqa_sliding_window_varlen_fwd_bench(
         is_causal=is_causal,
         window_size_left=wl,
         window_size_right=wr,
-        dtype=dtype,
         tune=tune,
     )
     op.total_q = sum(seqlens_q)

benchmarks/ops/attention/bench_mha.py

@@ -139,7 +139,7 @@ def test_mha_bwd_bench(batch: int, seq_len: int, heads: int, dim: int, causal: b
     test = MhaBwdWorkload(batch, heads, seq_len, dim, causal, dtype)
     inputs = test.gen_inputs()
 
-    op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, dtype, tune=tune)
+    op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, tune=tune)
     bm = ManifestBenchmark(_MHA_BWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/attention/bench_mha_decode.py

@@ -76,7 +76,7 @@ def test_mha_decode_bench(b: int, h: int, s_q: int, s_kv: int, d: int, dtype: to
     test = MhaDecodeTestBaseline(b, h, s_q, s_kv, d, dtype)
     inputs = test.gen_inputs()
 
-    op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, dtype, tune=tune)
+    op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, tune=tune)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/attention/bench_mha_decode_paged.py

@@ -133,7 +133,7 @@ def test_mha_decode_paged_bench(batch: int, heads: int, seqlen_q: int, seqlen_kv
     q, k, v, real_seqlen_kv, block_table = inputs
 
     op = MultiHeadAttentionDecodePagedWithKVCacheFwdOp(
-        batch, heads, seqlen_q, seqlen_kv, dim, page_size, is_causal, dtype, tune=tune)
+        batch, heads, seqlen_q, seqlen_kv, dim, page_size, is_causal, tune=tune)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/bench_ada_layer_norm.py

@@ -29,7 +29,7 @@ def test_ada_layer_norm_bench(m: int, n: int, dtype: torch.dtype) -> None:
     test = AdaLayerNormWorkload(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = AdaLayerNormFwdOp(dtype=dtype)
+    op = AdaLayerNormFwdOp()
     bm = ManifestBenchmark(_ADA_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -48,7 +48,7 @@ def test_ada_layer_norm_zero_bench(m: int, n: int, dtype: torch.dtype) -> None:
     test = AdaLayerNormZeroWorkload(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = AdaLayerNormZeroFwdOp(dtype=dtype)
+    op = AdaLayerNormZeroFwdOp()
     bm = ManifestBenchmark(_ADA_ZERO_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/bench_argreduce.py

@@ -39,7 +39,7 @@ def test_argmax_bench(shape: tuple, dtype: torch.dtype, extra: dict) -> None:
     workload = ArgmaxWorkload(shape, dtype)
     inputs = workload.gen_inputs()
 
-    op = ArgmaxFwdOp(dtype=dtype, **extra)
+    op = ArgmaxFwdOp(**extra)
     bm = ManifestBenchmark(_ARGMAX_OP, op, workload)
     # FIXME(staged-rollout): ArgreduceKernel skips large-N manifest workloads
     #
@@ -72,7 +72,7 @@ def test_argmin_bench(shape: tuple, dtype: torch.dtype, extra: dict) -> None:
     workload = ArgminWorkload(shape, dtype)
     inputs = workload.gen_inputs()
 
-    op = ArgminFwdOp(dtype=dtype, **extra)
+    op = ArgminFwdOp(**extra)
     bm = ManifestBenchmark(_ARGMIN_OP, op, workload)
     # FIXME(staged-rollout): ArgreduceKernel skips large-N manifest workloads
     #

benchmarks/ops/bench_engram.py

@@ -82,7 +82,7 @@ def test_engram_gate_conv_fwd_bench(M, seq_len, d, dtype):
     test = EngramGateConvFwdWorkload(M, seq_len, d, dtype)
     inputs = test.gen_inputs()
 
-    op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=_TUNE)
+    op = EngramGateConvFwdOp(M, seq_len, d, tune=_TUNE)
     bm = ManifestBenchmark(_ENGRAM_GATE_CONV_FWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -156,7 +156,7 @@ def test_engram_gate_conv_bwd_bench(M, seq_len, d, dtype):
     test = EngramGateConvBwdTestBaseline(M, seq_len, d, dtype)
     inputs = test.gen_inputs()
 
-    op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=_TUNE)
+    op = EngramGateConvBwdOp(M, seq_len, d, tune=_TUNE)
     bm = ManifestBenchmark(_ENGRAM_GATE_CONV_BWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -242,7 +242,7 @@ def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, di
     inputs = test.gen_inputs()
 
     op = EngramDecodeOp(
-        batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune=_TUNE,
+        batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, tune=_TUNE,
     )
     bm = ManifestBenchmark(_ENGRAM_DECODE_OP, op, test)
     result = bm.profile(op, *inputs)

benchmarks/ops/bench_logical_reduce.py

@@ -32,7 +32,7 @@ def test_any_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = AnyFwdOp(dtype=dtype, **op_params)
+    op = AnyFwdOp(**op_params)
     bm = ManifestBenchmark(_ANY_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -66,7 +66,7 @@ def test_all_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = AllFwdOp(dtype=dtype, **op_params)
+    op = AllFwdOp(**op_params)
     bm = ManifestBenchmark(_ALL_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -100,7 +100,7 @@ def test_count_nonzero_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = CountNonzeroFwdOp(dtype=dtype, **op_params)
+    op = CountNonzeroFwdOp(**op_params)
     bm = ManifestBenchmark(_COUNT_NONZERO_OP, op, test)
     try:
         result = bm.profile(op, *inputs)

benchmarks/ops/bench_mean_pooling.py

@@ -111,7 +111,6 @@ def test_mean_pooling_bench(batch_size: int, seq_len: int, heads: int, dim: int,
         "chunks_per_batch": chunks_per_batch,
         "seq_num": seq_num,
         "use_offsets": use_offsets,
-        "dtype": dtype,
         "accum_dtype": accum_dtype,
         "tune": tune,
     }

benchmarks/ops/bench_moe_grouped_gemm_nopad.py

@@ -49,7 +49,7 @@ def test_moe_grouped_gemm_nopad_bench(
     workload = MoeGroupedGemmNopadWorkload(numel, num_experts, n, k, dtype)
     a, b, true_sizes, true_offsets = workload.gen_inputs()
 
-    op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k, dtype=dtype)
+    op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k)
     bm = ManifestBenchmark(_OP_NAME, op, workload)
 
     # Warmup: trigger JIT compilation before timed profiling.

benchmarks/ops/bench_moe_permute_nopad.py

@@ -66,7 +66,7 @@ def test_moe_permute_nopad_bench(
     hidden_states, topk_ids = workload.gen_inputs()
 
     # TileOPs
-    op = MoePermuteNopadFwdOp(num_experts=num_experts, dtype=dtype)
+    op = MoePermuteNopadFwdOp(num_experts=num_experts)
     bm = ManifestBenchmark(_OP_NAME, op, workload)
     op(hidden_states, topk_ids)  # warmup / JIT compile
     torch.cuda.synchronize()

benchmarks/ops/bench_moe_shared_fused_moe.py

@@ -133,7 +133,6 @@ def test_shared_fused_moe_bench(
         renormalize=renormalize,
         with_correction_bias=with_correction_bias,
         routed_scaling_factor=routed_scaling_factor,
-        dtype=dtype,
         shared_ffn_size=shared_ffn_size,
     )
     op(hidden, gating, w_gate_up, w_down, correction_bias,

benchmarks/ops/bench_moe_unpermute.py

@@ -63,7 +63,7 @@ def test_moe_unpermute_bench(total_tokens: int, top_k: int, hidden_size: int) ->
     mm2_pad, fwd_idx, topk_weights = test.gen_inputs()
 
     # TileOPs
-    op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size, dtype)
+    op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size)
     bm = ManifestBenchmark(_OP_NAME, op, test)
     op(mm2_pad, fwd_idx, topk_weights)  # warmup / JIT compile
     torch.cuda.synchronize()

benchmarks/ops/bench_norm.py

@@ -48,7 +48,7 @@ def test_rms_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool) -> None:
     test = RMSNormTestBaseline(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = RMSNormFwdOp(normalized_shape=(n,), dtype=dtype, tune=tune)
+    op = RMSNormFwdOp(normalized_shape=(n,), tune=tune)
     bm = ManifestBenchmark(_RMS_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -85,7 +85,7 @@ def test_fused_add_rms_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool
     test = FusedAddRMSNormWorkload(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = FusedAddRMSNormFwdOp(dtype=dtype, tune=tune)
+    op = FusedAddRMSNormFwdOp(tune=tune)
     bm = ManifestBenchmark(_FUSED_RMS_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -121,7 +121,7 @@ def test_layer_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool) -> Non
     test = LayerNormWorkload(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype, tune=tune)
+    op = LayerNormFwdOp(normalized_shape=(n,), tune=tune)
     bm = ManifestBenchmark(_LN_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
@@ -154,7 +154,7 @@ def test_fused_add_layer_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bo
     test = FusedAddLayerNormWorkload(m, n, dtype)
     inputs = test.gen_inputs()
 
-    op = FusedAddLayerNormFwdOp(dtype=dtype, tune=tune)
+    op = FusedAddLayerNormFwdOp(tune=tune)
     bm = ManifestBenchmark(_FUSED_LN_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")

benchmarks/ops/bench_reduce.py

@@ -55,7 +55,7 @@ def test_sum_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)  # baseline below reduces dim=-1
-    op = SumFwdOp(dtype=dtype, **op_params)
+    op = SumFwdOp(**op_params)
     bm = ManifestBenchmark(_SUM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -89,7 +89,7 @@ def test_mean_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)  # baseline below mirrors the op's dim
-    op = MeanFwdOp(dtype=dtype, **op_params)
+    op = MeanFwdOp(**op_params)
     bm = ManifestBenchmark(_MEAN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -123,7 +123,7 @@ def test_amax_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = AmaxFwdOp(dtype=dtype, **op_params)
+    op = AmaxFwdOp(**op_params)
     bm = ManifestBenchmark(_AMAX_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -157,7 +157,7 @@ def test_amin_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = AminFwdOp(dtype=dtype, **op_params)
+    op = AminFwdOp(**op_params)
     bm = ManifestBenchmark(_AMIN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -185,7 +185,7 @@ def test_prod_bench(shape: tuple, dtype: torch.dtype) -> None:
     test = ProdWorkload(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = ProdFwdOp(dtype=dtype)
+    op = ProdFwdOp()
     bm = ManifestBenchmark(_PROD_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -216,7 +216,7 @@ def test_std_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = StdFwdOp(dtype=dtype, correction=1, **op_params)
+    op = StdFwdOp(correction=1, **op_params)
     bm = ManifestBenchmark(_STD_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -250,7 +250,7 @@ def test_var_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = VarFwdOp(dtype=dtype, correction=1, **op_params)
+    op = VarFwdOp(correction=1, **op_params)
     bm = ManifestBenchmark(_VAR_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -284,7 +284,7 @@ def test_var_mean_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = VarMeanFwdOp(dtype=dtype, correction=1, **op_params)
+    op = VarMeanFwdOp(correction=1, **op_params)
     bm = ManifestBenchmark(_VAR_MEAN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)

benchmarks/ops/bench_softmax.py

@@ -87,7 +87,7 @@ def test_logsumexp_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = LogSumExpFwdOp(dtype=dtype, tune=True, **op_params)
+    op = LogSumExpFwdOp(tune=True, **op_params)
     bm = ManifestBenchmark(_LOGSUMEXP_OP, op, test)
     try:
         result = bm.profile(op, *inputs)

benchmarks/ops/bench_vector_norm.py

@@ -32,7 +32,7 @@ def test_l1_norm_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = L1NormFwdOp(dtype=dtype, **op_params)
+    op = L1NormFwdOp(**op_params)
     bm = ManifestBenchmark(_L1_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -68,7 +68,7 @@ def test_l2_norm_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = L2NormFwdOp(dtype=dtype, **op_params)
+    op = L2NormFwdOp(**op_params)
     bm = ManifestBenchmark(_L2_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -104,7 +104,7 @@ def test_inf_norm_bench(
     inputs = test.gen_inputs()
 
     op_params.setdefault("dim", -1)
-    op = InfNormFwdOp(dtype=dtype, **op_params)
+    op = InfNormFwdOp(**op_params)
     bm = ManifestBenchmark(_INF_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)

docs/design/manifest.md

@@ -14,7 +14,7 @@ One or more YAML files per family (single file by default; large families may sh
 
 ```mermaid
 flowchart LR
-    H["Human reviewer"] -->|writes / approves| M["tileops/manifest/"]
+    R["Authoritative reference"] -->|specified from| M["tileops/manifest/"]
     M -->|reads spec from| A["Agent (codegen)"]
     A -->|produces| C["Op code, tests, benchmarks"]
     M -->|validates against| V["Validator (CI)"]