Examples¶

This page collects five real specs from the manifest. Each spec is followed by a table that maps each form used in it to the section that describes it. The specs are taken from src/tileops/manifest/, with comments and some workload rows removed.

1. GemmFwdOp: type family¶

# src/tileops/manifest/spec/gemm.yaml
GemmFwdOp:
  ref_api: "torch.matmul"
  family: gemm
  status: implemented

  signature:
    types:
      Mat:
        params: {t: Bool, R: Dim, C: Dim}
        match: t
        cases:
        - {when: false, is: "[R, C]"}
        - {when: true, is: "[C, R]"}
    forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"}
    params:
      trans_a: {type: bool, default: false}
      trans_b: {type: bool, default: true}
    inputs:
      a: {dtype: T, shape: "Mat[trans_a, M, K]"}
      b: {dtype: T, shape: "Mat[trans_b, K, N]"}
    outputs:
      d: {dtype: T, shape: "[M, N]"}

  workloads:
  - {M: 1024, N: 1024, K: 1024, trans_a: false, trans_b: false, dtype_cases: [{T: float16}, {T: bfloat16}], label: square-1k}
  - {M: 128, N: 2112, K: 7168, trans_a: false, trans_b: true, dtype_cases: [{T: bfloat16}], label: ds-v3-decode-qkv-a}
  # remaining rows omitted

  roofline:
    flops: "2 * M * N * K"

Table 1 The forms used in GemmFwdOp

No.	Form in the spec	Described in
1	`types.Mat`: `trans_a` and `trans_b` decide the axis order of `a` and `b`	Extensions § 2
2	`forall` declares three axis lengths and one `DType` index	Spec fields § 3
3	`trans_a` and `trans_b` are both construction parameters and discriminants of `Mat`	Spec fields § 4
4	workload rows give the values of parameters and indices; the first row expands into two cases, with ids `square-1k-float16` and `square-1k-bfloat16`	Spec fields § 8
5	`roofline` gives only `flops`; `bytes` is derived from the signature and equals `(MK + KN + M*N)` times the bytes per element	Spec fields § 9

2. MoEPrePermuteFwdOp: ADT, let and generator¶

The type of the layout parameter is the ADT MGroupedLayout defined in spec/types.yaml; its definition is in Extensions § 3.

# src/tileops/manifest/spec/moe.yaml
MoEPrePermuteFwdOp:
  family: moe
  status: implemented
  signature:
    types:
      LayoutMetadata:
        params: {layout: MGroupedLayout, E: Dim, P: Dim}
        match: layout
        cases:
        - {when: {contiguous: {metadata_kind: physical_psum}}, is: "[E]"}
        - {when: {contiguous: {metadata_kind: per_row}}, is: "[P]"}
    forall: {T: Dim, H: Dim, K: Dim, D: "DType[float16 | bfloat16]"}
    params:
      layout: {type: MGroupedLayout}
      num_local_experts: {type: int}
    inputs:
      hidden_states: {dtype: D, shape: "[T, H]", contiguous: true}
      local_expert_ids: {dtype: int32, shape: "[T, K]", contiguous: true, values: "topk_ids(T, K, num_local_experts)", requires: ["in_range(0, num_local_experts)"]}
    let:
      P: "moe.capacity(layout, T * K, num_local_experts)"
    outputs:
      expert_input: {dtype: D, shape: "[P, H]", contiguous: true}
      layout_metadata: {dtype: int32, shape: "LayoutMetadata[layout, num_local_experts, P]", contiguous: true}
      inverse_indices: {dtype: int32, shape: "[T * K]", contiguous: true}
    shape_rules:
    - "layout.kind == 'contiguous'"
    - "num_local_experts > 0"
    - "T > 0"
    - "K > 0"
  workloads:
  - {T: 32, H: 128, K: 2, num_local_experts: 4, layout: {contiguous: {packing: tight, metadata_kind: physical_psum, alignment: 1}}, dtype_cases: [{D: float16}, {D: bfloat16}], label: "decode-t32"}
  # remaining rows omitted
  roofline:
    flops: "0"

Table 2 The forms used in MoEPrePermuteFwdOp

No.	Form in the spec	Described in
1	`layout.kind == 'contiguous'` is a domain restriction that excludes the masked layout, so `LayoutMetadata` needs no case for masked	Extensions § 2, § 3
2	`LayoutMetadata` pattern-matches on the constructor and the field `metadata_kind`	Extensions § 2
3	the `int` parameter `num_local_experts` appears directly in shapes	Spec fields § 4
4	`P` in `let` is computed by the primitive `moe.capacity` and used in output shapes	Extensions § 4
5	the values of `local_expert_ids` are generated by the generator `topk_ids`, and its contents are constrained by `requires`	Extensions § 7
6	each tensor declares `contiguous: true`	Extensions § 5
7	`layout` in the workload row is written as an ADT value	Extensions § 3

3. MeanPoolingFwdOp: optional inputs and metadata tensors¶

# src/tileops/manifest/spec/pool.yaml
MeanPoolingFwdOp:
  family: pool
  status: implemented

  signature:
    types:
      PoolOut:
        params: {p: Bool, B: Dim, NC: Dim, S: Dim, c: Dim, H: Dim, D: Dim}
        match: p
        cases:
        - {when: true, is: "[B, NC, H, D]"}
        - {when: false, is: "[B, ceil_div(S, c), H, D]"}
    forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[float16 | bfloat16 | float32]", seq_lens: "Seq[Int]"}
    params:
      chunk_size: {type: int}
      accum_dtype: {type: torch.dtype}
    inputs:
      x: {dtype: T, shape: "[B, S, H, D]"}
      offsets: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]}
      indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"}
    outputs:
      output: {dtype: T, shape: "PoolOut[present(offsets), B, NC, S, chunk_size, H, D]"}
    shape_rules:
    - "chunk_size > 0 and chunk_size % 32 == 0"
    - "D <= 128 or D % 128 == 0"
    - "category(accum_dtype) == 'float'"

  workloads:
  - {B: 1, S: 8192, H: 64, D: 128, chunk_size: 64, accum_dtype: float32, dtype_cases: [{T: float16}, {T: bfloat16}], label: uniform-8k}
  - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [offsets], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even}
  # remaining rows omitted

  roofline:
    flops: "B * S * H * D + B * (NC if present(offsets) else ceil_div(S, chunk_size)) * H * D"

Table 3 The forms used in MeanPoolingFwdOp

No.	Form in the spec	Described in
1	`indices` is declared `optional: "present(offsets)"`, so it is passed together with `offsets` or omitted together with it	Extensions § 1
2	the type family `PoolOut` uses `present(offsets)` as its discriminant	Extensions § 2
3	the first row does not pass `offsets` and does not use `seq_lens`, so the row does not give it; the second row passes `offsets` through `some: [offsets]` and gives `seq_lens`, written here as the value primitive call `"repeat(2048, 4)"`	Extensions § 1, Spec fields § 8
4	`NS` and `NC` are solved by unification of the generator results and are not given in the row	Extensions § 7
5	`category(accum_dtype) == 'float'` requires the accumulation dtype to be a floating-point type	Extensions § 9
6	the roofline formula distinguishes the two cases with `present(offsets)`	Spec fields § 9

4. ReluFwdOp: arbitrary rank and effect¶

# src/tileops/manifest/spec/elementwise_unary_activation.yaml
ReluFwdOp:
  ref_api: "torch.nn.functional.relu"
  family: elementwise
  status: implemented

  signature:
    forall: {S: Shape, T: "DType[float16 | bfloat16 | float32]"}
    params:
      inplace: {type: bool, default: false, kw_only: true}
    inputs:
      input: {dtype: T, shape: "[*S]", mutated: inplace}
    outputs:
      output: {dtype: T, shape: "[*S]", alias: input}

  workloads:
  - {S: [2048, 4096], dtype_cases: [{T: float16}, {T: bfloat16}], label: "hidden-state-prefill"}
  - {S: [1, 4096], dtype_cases: [{T: bfloat16}], label: "hidden-state-decode"}

  roofline:
    flops: "prod(S)"

Table 4 The forms used in ReluFwdOp

No.	Form in the spec	Described in
1	the kind of `S` is `Shape`, and the input and output use the same shape term `[*S]`	Spec fields § 3, § 5
2	`mutated: inplace` and `alias: input` mean that when `inplace` is true the op writes the input and returns it	Extensions § 6
3	`S` in the workload rows is written as a list of integers	Spec fields § 8

5. AlibiFwdOp: no tensor inputs¶

# src/tileops/manifest/spec/elementwise_generative.yaml
AlibiFwdOp:
  family: elementwise
  status: implemented

  signature:
    params:
      seq_len: {type: int, kw_only: true}
      num_heads: {type: int, kw_only: true}
      out_dtype: {type: "float16 | bfloat16 | float32", default: float32, kw_only: true}
      device: {type: "torch.device | str | None", default: null, kw_only: true}
    outputs:
      output: {dtype: out_dtype, shape: "[num_heads, seq_len, seq_len]"}
    shape_rules:
    - "seq_len > 0"
    - "num_heads > 0"

  workloads:
  - {seq_len: 2048, num_heads: 32, out_dtype: float16, label: "mpt-7b-2k"}
  - {seq_len: 2048, num_heads: 32, out_dtype: bfloat16, label: "mpt-7b-2k"}
  # remaining rows omitted

  roofline:
    flops: "3 * num_heads * seq_len * seq_len"

Table 5 The forms used in AlibiFwdOp

No.	Form in the spec	Described in
1	every index comes from a construction parameter, so the spec has no `forall`	Spec fields § 4
2	the output's `dtype` uses the dtype parameter `out_dtype` directly	Spec fields § 7
3	the op has no call-time tensor inputs, so it declares a `device` parameter	Extensions § 5, Calls and validation § 4
4	the two rows have the same `label`, and their case ids are distinguished by the value of the dtype parameter: `mpt-7b-2k-float16` and `mpt-7b-2k-bfloat16`	Spec fields § 8