Spec fields¶
Most specs use only the following fields, and this page describes how to write them:
forall;- the
shapeanddtypeof tensors; - construction parameters;
- simple
shape_rules; - workload rows;
- inline roofline formulas.
How to write optional inputs, shapes that vary with parameters, side effects, metadata tensors and similar cases is in Extensions.
1. Structure of a spec¶
The SiluAndMul spec uses only the fields above:
SiluAndMulFwdOp:
family: elementwise
status: implemented
signature:
forall: {M: Dim, N: Dim, T: "DType[float16 | bfloat16 | float32]"}
inputs:
x: {dtype: T, shape: "[M, 2 * N]"}
outputs:
output: {dtype: T, shape: "[M, N]"}
workloads:
- {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill}
- {M: 1, N: 14336, dtype_cases: [{T: bfloat16}], label: llama-8b-ffn-decode}
roofline:
flops: "6 * M * N"
The following lists every top-level field a spec can contain. Each comment gives the section that describes the field:
<Op>:
family: <module name> # 2
status: implemented | spec-only # 2
ref_api: <fully qualified name> # 2, optional
signature:
forall: {<index>: <kind>} # 3
params: {<p>: {type, default, kw_only}} # 4
inputs: {<t>: {dtype, shape, optional, mutated, values, requires, ...}} # 5
outputs: {<t>: {dtype, shape, nullable, buffer, alias, ...}} # 5
types: {<Family>: {params, match, cases}} # Extensions 2
let: {<name>: "<expression>"} # Extensions 4
shape_rules: ["<refinement>"] # 6
dtype_combos: [{<DType index>: <dtype>}] # Extensions 8
workloads: [{<params and indices>, some, dtype_cases, label}] # 8
roofline: {flops, bytes} | {func} # 9
composition: {kind: composite, stages} # Extensions 10
2. Top-level fields¶
family, status, signature, workloads and roofline are required, and workloads contains at least one row.
familyis the name of the public module the op belongs to. The op can be imported by class name fromtileops.<family>, and the__all__of that module agrees with the manifest. The name of the YAML file that holds the spec is also determined byfamily, see Reading and writing the manifest § 6.statustakes one of two values:implementedmeans an implementation that conforms to the spec exists;spec-onlymeans no conforming implementation exists yet; the code may be absent or only partly done. Checks that depend on code are skipped only forspec-onlyops.statusdecides only which code-dependent checks run, and does not affect the methods generated from the signature: any class that has a spec gets all the generated methods, includingeval_roofline().ref_apiis optional and records the fully qualified name of the API the op follows semantically, for exampletorch.matmul. The validator checks its format, and when the corresponding module can be imported, it also checks that the name exists.
3. forall and kind¶
forall declares every free type index in the signature and its kind, for example forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"}.
Table 1 Kinds available in forall
| No. | kind | Values | How it is determined at call time | How it is written in a workload row |
|---|---|---|---|---|
| 1 | Dim |
a non-negative integer, an axis length | solved by unification of the inputs | an integer |
| 2 | Shape |
a tuple of Dim |
solved by unification of the inputs | a list of integers |
| 3 | DType[a \| b] |
one of the listed dtypes | solved by unification of the inputs | dtype_cases |
| 4 | Seq[Int] |
a list of integers, such as q_lens |
exists only at instantiation, as an argument to a generator (see Extensions § 7) | a list of integers, or a call to a value primitive that returns a list, such as "repeat(512, 64)" |
- A
Dimcan be used where anIntis required, and aShapecan be used where aSeq[Int]is required. - Every axis in a shape is an integer expression. An axis whose kind is not
Dim, and every element of a sequence expanded with*p, must be non-negative; the generated checks confirm this at construction or at call time. - Values in YAML are converted to Python values by the declared kind or
type:Shapeandtuple[...]become tuples,Seq[Int]andlist[...]stay lists, dtype names becometorch.dtype, and ADT values become the corresponding Python objects. - Within one spec, the names of indices,
letentries and tensors are all distinct.outis a reserved name and cannot be used as the name of a tensor, parameter or index.
4. params¶
signature.params corresponds one to one with the parameter list of __init__. An ordinary construction parameter declares type, and optionally default and kw_only (keyword-only). A tensor passed at construction declares dtype and shape instead, see Extensions § 5.
- The spec is the authority. For an
implementedop, the validator comparesparamswith__init__item by item, and requires the same set of parameters, order,defaultandkw_only. - The code may have only two kinds of extra parameter: the execution-policy parameter
target, which every op has, written as*, target=Noneafter the manifest parameters; and the implementation objects injected by the caller that the class attributeinjected_parameterslists. - The call-time inputs and output buffers in the signature form, in order, the leading part of the
forwardparameter list.forwardmay append code-defined execution parameters after them; those parameters are not part of the signature.
Construction parameters can appear directly in types without extra annotation:
- an
intparameter can be written in a shape, for examplenum_local_expertsin MoE; - a
list[int]ortuple[int, ...]parameter can be expanded in a shape as*p, for examplenormalized_shapein RMSNorm; - a dtype parameter (whose
typeis a union of dtype names) can be used directly as a tensor'sdtype, for exampleout_dtypein Alibi; - how to write an
int | Noneparameter is in Extensions § 1.
The full mapping from type to kind is in the design document Manifest Table 3. Parameters of types such as float or an unconstrained str do not take part in type inference, but they can appear in shape_rules and roofline formulas; the validator checks their type, default, and how they are used in those expressions.
5. inputs and outputs¶
Each tensor is written {dtype: ..., shape: "..."}, which corresponds to the type Tensor[T, s].
- A shape is written
"[" axis, ... "]", or as an application of a type family. Each axis can be an expression,*Sor*primitive(...);"[]"denotes a zero-dimensional tensor. - Tensors with the same shape use the same shape term; for example, the input and output of an elementwise op are both written
[*S]. - Each spec has exactly one signature, so every call returns outputs with the same names and the same count. If the number of outputs of an op varies with some switch, or the same parameter can be either a scalar or a tensor, the op is split into several specs.
- An op selects its implementation from its parameters and from which tensors are passed. The contents of a tensor are only input to the computation and do not affect implementation selection.
6. shape_rules¶
Each entry in shape_rules is a refinement, a predicate that constrains the values of indices, checked after unification.
- Shapes,
let, type families, refinements and inline roofline formulas use the same closed expression language, with the same operator precedence as Python. The parts of the expression language are in the design document Manifest Table 10. - A refinement can depend only on quantities available at run time. A refinement that can be evaluated at construction is checked at construction; the others are checked on every call.
- Constraints on the contents of a metadata tensor are written in the tensor's
requires, not inshape_rules, see Extensions § 7. - A condition narrows kinds within the branch it selects: in the branch where
present(v)is true,Maybe[X]narrows toX; in the branch wherex == 'a'orx in ('a', 'b')holds, the values ofxnarrow to those literals. If a comparison uses a literal that is not a member of the corresponding enum or dtype set, the validator rejects the refinement. - Whether a refinement can be satisfied is the responsibility of the spec's author.
- Forms such as
x.shape == (...),x is Noneandisinstanceare rejected; the corresponding rewrites are in Calls and validation § 7.
7. dtype¶
A tensor's dtype is a dtype expression, which takes one of four forms:
Table 2 Forms of a dtype expression
| No. | Form | Meaning | Example |
|---|---|---|---|
| 1 | a DType index in forall |
takes a value in the declared set, solved by unification of the inputs | a: {dtype: T} |
| 2 | a dtype parameter | takes the value of the construction parameter | output: {dtype: out_dtype} |
| 3 | a constant | a fixed dtype | cu_seqlens_q: {dtype: int32} |
| 4 | a dtype primitive | computed from other dtypes | promote_int_to_float(T), coalesce_dtype(out_dtype, D) |
If the spec has no dtype_combos, each DType index takes values in its own set independently. How to write a spec in which several DType indices allow only specific combinations is in Extensions § 8.
8. Workload rows¶
Each workload row determines one call. The benchmarks, nightly and the manifest tests all generate their calls from workload rows.
workloads:
- {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill}
Table 3 Keys of a workload row
| No. | Key | Value |
|---|---|---|
| 1 | a construction parameter name | the parameter's value; parameters without a default must be given |
| 2 | some |
the optional tensors passed in this call, see Extensions § 1 |
| 3 | a Dim, Shape or Seq[Int] index in forall |
the indices used on the current branch that are not solved by a generator |
| 4 | dtype_cases |
a non-empty list, each item of which is one assignment of the DType indices used on the current branch, for example [{T: float16}, {T: bfloat16}]; written only when such indices exist; dtype parameters are given as construction parameters |
| 5 | label |
the name of the row |
- An index is used on a branch when it appears in one of the following places on that branch: a shape, a dtype, a refinement, a generator argument,
requires, or an inline roofline formula. A workload row must give every index used on the current branch, and only those indices;letentries and indices solved by a generator are not given in the row. - For this decision, each expression is first simplified for the current branch. If the condition of a refinement is always true on that branch, the indices in it do not count as used.
- When a
letis used, the indices in its expression also count as used. - A roofline in
funcform makes no index used. - A discriminant that selects a type family branch, decides whether a tensor exists, or decides whether an output is
Nonealways counts as used. -
Each row expands into several cases according to
dtype_cases. The case id joins the following three parts in order with-: -
label; - each dtype value in
dtype_cases, in the declaration order offorall; - the values of dtype parameters, in the declaration order of
params.
For example, the id of the first case in the example above is llama-8b-ffn-prefill-float16.
- Nightly history is keyed by case id, so changing a label breaks the history of that row.
- A label is non-empty, at most 24 characters long, and contains only characters in [A-Za-z0-9._-]. A label describes only the scenario the row models, that is, the model and its use or the synthetic purpose, plus the qualifiers needed to tell similar rows apart, for example llama-8b-ffn-prefill. The op name and dtype already appear in the case id and are not repeated. Rows that share a label may differ only in dtype.
- Case ids within one spec are all distinct.
- At instantiation, the workload row determines the shapes, dtypes, parameter values, which optional tensors are passed, and the values of metadata tensors. The rest is generated by fixed rules: the device is determined as described in Calls and validation § 4, strides are contiguous, tensors do not alias each other, and ordinary data is random.
- The validator re-infers the call from the instantiated inputs and requires the result to agree with the workload row.
- Workload rows are not responsible for unit-test coverage; the shapes needed to cover every branch of the kernel are chosen separately by each op's tests.
9. roofline¶
roofline gives the FLOPs and byte count of one call. It is written either as inline formulas or as a func:
roofline:
flops: "2 * M * N * K" # inline formula; when bytes is omitted it is derived from the signature
# or
roofline:
func: "tileops.perf.formulas.gqa_fwd_roofline"
- Inline formulas are written in the expression language and can refer to the signature's indices, construction parameters and
letentries, as well aspresent(t),bytes(t)(the byte count of tensort) and built-in primitives. When the cost depends on whether a tensor is passed, the formula can distinguish the cases withpresent(...). - When
bytesis omitted, the byte count is derived by reading or writing each tensor in full once: - an input that is not written counts as one read, an output counts as one write, a
mutatedinput counts as one read and one write, and awrite_onlyinput counts as one write only; - each tensor name in the signature counts as separate storage, even if the caller passes the same tensor to two parameters; only aliases declared in the signature (
buffer,alias) count as the same storage; - the byte count of each tensor is
prod(shape) * bits(dtype) / 8, and packed dtypes are counted by their carrier. - The byte count is the algorithm's minimum memory traffic: intermediate results are not counted, and an input that the algorithm reads only in part counts the distinct elements actually read. When the two differ by less than 1% on every workload row, the whole tensor may be counted.
- If the derived byte count does not match this minimum traffic, the spec writes
bytesexplicitly, with a corresponding test. flopsis the minimum amount of arithmetic the algorithm needs for the call, not the number of instructions the hardware executes:- values already computed count as reused;
- recurrent ops such as linear attention and state-space scans count by the chunked algorithm at the chunk size when the signature has a chunk size, and by per-token recurrence otherwise;
- elementwise operations, attention and MoE are counted by the conventions in the design document Roofline § 1.3;
- a path that is an identity map on the call's dtype counts as 0.
funcis used when the formula needs Python logic. It points to a module-level functionf(call) -> tuple[int, int]intileops.perf.formulas. The function reads only its argumentcalland does not read the op instance.callis one checked call and provides the following interface:
Table 4 The argument call of a roofline func
| No. | Interface | Contents |
|---|---|---|
| 1 | call.indices |
the parameters, the indices and dtype indices solved for this call, and the let entries used, that is, the names an inline formula can refer to |
| 2 | call.present(t) |
whether tensor t is passed, held or returned; call.present("out") says whether the caller passed out |
| 3 | call.tensors[t] |
the (shape, dtype name) of tensor t |
| 4 | call.bytes(t) |
the byte count of tensor t |
| 5 | call.metadata_values(t) |
the contents of metadata tensor t; calling it on meta tensors raises an error, because meta tensors have no values |
| 6 | call.stages |
the calls completed by each sub-op of a composite op in this call, indexed by stage name |
- Every spec generates an
eval_roofline()method, which computes FLOPs and byte count from the op's most recently completed call. The benchmarks obtain the numbers through this method and write them into the results; the roofline tools read the benchmark results and do not call the op directly.
The full roofline rules are in the design document Roofline.