---
title: 'NineToothed: DSL for ML Kernels'
url: https://www.emergentmind.com/topics/ninetoothed
type: topic
---

# NineToothed: DSL for ML Kernels

Searching arXiv for the NineToothed paper and closely related context papers mentioned in the provided data.
NineToothed is a high-level domain-specific language for machine learning kernels that provides serial semantics for GPU-style workloads and automatically lowers those programs to parallel Triton kernels. Implemented as a Python-embedded DSL and integrated with PyTorch, it is centered on tensor-oriented metaprogramming (TOM) and an arrange-and-apply programming paradigm. The stated objective is to let developers express tiled tensor computations without explicit parallel reasoning, pointer arithmetic, or manual `load`/`store` logic, while preserving performance comparable to hand-written Triton kernels [2507.11978].

## 1. Definition and architectural position

NineToothed is presented as a domain-specific language that extends Triton upward rather than replacing it. In the formulation given for the system, developers write serial-looking code over symbolic and hierarchical tensors, and NineToothed then analyzes the symbolic tensor transformations, generates Triton kernel code, and launches that code on GPUs with an appropriate grid configuration [2507.11978].

The motivation is grounded in the observation that current deep learning DSLs such as Triton still require expertise in parallel programming and expose low-level details, including explicit `program_id` handling, pointer arithmetic, memory address calculation, masking, and direct `load`/`store` orchestration. NineToothed addresses this by replacing explicit parallel indexing with symbolic tensor layouts and tilings that the compiler can understand and lower automatically [2507.11978].

Within the surrounding compiler landscape, NineToothed is described as distinct from systems such as TVM, Halide, XLA, and MLIR-based compiler stacks. Those systems typically separate algorithm from schedule through loop nests, scheduling primitives, or graph rewrites. By contrast, NineToothed expresses schedule-like information through symbolic tensor metaprogramming together with a serial tile-level application function. The data also situates NineToothed relative to Graphene, noting that it borrows the idea of hierarchical tensors but uses that idea as a high-level DSL abstraction for developers rather than as a low-level IR [2507.11978].

## 2. Programming model and serial semantics

The central design principle is serial semantics for parallel workloads. The contrast drawn in the source material is between an intuitive serial model, in which an outer function iterates over objects and applies an inner function, and a typical CUDA- or Triton-like parallel model, in which each parallel instance must determine which data slice it owns. NineToothed preserves the former perspective: the programmer writes only the serial per-tile computation, while the logic for which tiles exist, how they are laid out, and which program processes which tile is encoded separately through symbolic tensor meta-operations [2507.11978].

This decoupling is formalized as the arrange-and-apply paradigm. In the arrangement phase, input and output tensors are transformed through compile-time meta-operations into aligned tile structures. In the application phase, a serial function operates per tile or per aligned tile group, without manual index calculations. The required semantic condition is that the arrangement produce matching outermost dimensions across all parameters; when that condition holds, NineToothed can launch one parallel program per outer tile and automatically map tiles to programs [2507.11978].

The programming workflow described in the paper consists of three programmer-supplied elements: symbolic parameter tensors such as `Tensor(2)`, an `arrangement` function that tiles, reshapes, and aligns tensors, and an `application` function containing the serial per-tile computation. These are passed to `ninetoothed.make(arrangement, application, tensors)`, which yields a kernel callable from PyTorch. NineToothed then infers the grid size from the shape of the outermost arrangement level, generates Triton code containing `program_id`, `arange`, `load`, and `store`, and performs the tile-to-program and source-to-target mappings needed to access actual tensor data [2507.11978].

A notable consequence is that the user does not write `program_id`, `arange`, `load`, `store`, or pointer arithmetic directly. This suggests a deliberate abstraction boundary: NineToothed operates at the program or block level, while Triton’s compiler remains responsible for lower-level parallel details such as threads and warps.

## 3. Tensor-oriented metaprogramming and hierarchical tensors

TOM is the compile-time framework used to encode parallel structure as symbolic tensor transformations rather than explicit index arithmetic. The source describes tensors in NineToothed as potentially hierarchical, meaning that a tensor’s `dtype` may itself be a tensor. In addition, tensor attributes such as `shape` and `strides` may be symbolic expressions rather than concrete integers. These symbolic attributes are built on Python’s `ast.AST`, allowing tensor arithmetic to construct expression trees suitable for code generation [2507.11978].

The TOM meta-operations named in the paper are `tile`, `expand`, `squeeze`, `permute`, `flatten`, and `ravel`. These operate on tensor shapes and layouts only; no runtime data is moved by them. Their purpose is to express tiling, partitioning, broadcasting, and alignment as first-class symbolic transformations [2507.11978].

| Operation | Description |
|---|---|
| `tile` | Forms a hierarchical tensor (tiling). |
| `expand` | Expands singleton dimensions (like PyTorch). |
| `squeeze` | Removes singleton dimensions. |
| `permute` | Permutes dimensions. |
| `flatten` | Flattens selected dimensions. |
| `ravel` | Flattens all hierarchical levels into one. |

The `tile` operation is described as producing an outer level with shape `ceil(shape / tile_shape)` and an inner level with shape `tile_shape`, with default `strides == tile_shape` for non-overlapping tiles. Smaller strides allow sliding-window patterns, which the paper connects to convolution-like computations. `expand`, `squeeze`, `permute`, and `flatten` behave analogously to PyTorch layout transformations but remain purely symbolic and compile-time. `flatten` acts within a hierarchical level, whereas `ravel` collapses all hierarchical levels into one [2507.11978].

The deeper significance of TOM in this system is that it turns tilings and layouts into composable symbolic objects. The source explicitly states that this can enable reuse of arrangements, including the reuse of a matrix multiplication arrangement to implement convolution via implicit GEMM. A plausible implication is that NineToothed’s abstraction is not merely syntactic sugar over Triton, but a distinct representation strategy for tiled tensor programs.

## 4. Mapping, compilation, and code generation

The compilation pipeline begins in the Python-embedded front end, where the user writes `arrangement` and `application` functions using symbolic tensors. Through Python’s `ast` infrastructure, NineToothed captures symbolic expressions and meta-operations as expression trees. From these trees it deduces tiling and alignment relationships, then derives tile-to-program mappings and source-to-target mappings before emitting Triton code that is further compiled through Triton’s MLIR/LLVM backend [2507.11978].

Tile-to-program mapping is defined by the outermost arrangement shape shared by all arranged argument tensors. If that common shape is `outer_shape`, NineToothed launches $\prod_i \text{outer\_shape}[i]$ programs, and each program is assigned the tiles at one position in that outer shape. Inner hierarchical levels are then accessed by indexing within the program. This mechanism is the basis for automatic parallelization from a serial tile-oriented specification [2507.11978].

Source-to-target mapping handles the problem that arranged tensors may be far removed from the original tensor layout. According to the source, each tensor carries source dimensions, indicating which original tensor dimensions generated each current dimension, and target dimensions, indicating where broadcasted offsets are applied in the arranged tensor space. The code generator uses program IDs and `arange` to generate 1D offsets, source dimensions to compute stride-based offsets in the original layout, and target dimensions to broadcast offsets to multidimensional tile shapes; these are then combined into final offsets for `load` and `store` [2507.11978].

The paper emphasizes that NineToothed does not perform automatic schedule search in the style of TVM. Tiling strategy and vectorization are instead expressed by the programmer through the chosen meta-operations and tile sizes. The system is therefore described as constructive rather than dependence-analysis-driven: the programmer constructs the desired tiling pattern symbolically, and the compiler lowers that construction faithfully.

## 5. Kernel expression and representative examples

The simplest example given is vector addition. The arrangement tiles each one-dimensional tensor by `BLOCK_SIZE`, and the application is just `output = input + other`. The invocation allocates an output tensor with `torch.empty_like(input)` and calls the kernel with `BLOCK_SIZE=1024`. The explicit significance of the example is that there is no grid configuration and no explicit indexing; the parallel decomposition is entirely encoded in `tile((BLOCK_SIZE,))` [2507.11978].

Matrix multiplication provides the paper’s principal illustration of expressiveness. The arrangement tiles the output into `(BLOCK_SIZE_M, BLOCK_SIZE_N)`, tiles the left input into `(BLOCK_SIZE_M, BLOCK_SIZE_K)` and expands it to align with all output column tiles, and tiles the right input into `(BLOCK_SIZE_K, BLOCK_SIZE_N)` and expands it similarly. The application then performs a serial loop over the tile-level `k` dimension:
```python
accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
for k in range(input.shape[0]):
    accumulator += ntl.dot(input[k], other[k])
output = accumulator
```
The source explicitly contrasts this with Triton: the NineToothed GEMM kernel contains no explicit `pid_m`, `pid_n`, `pid_k`, no `offs_m`, `offs_n`, `offs_k`, and no direct pointer-offset or `load`/`store` calls. GEMM is expressed as arranging tiles and performing dot products over `k` [2507.11978].

For `conv2d`, the paper describes a transformation to implicit GEMM. The input has shape `(N, C, H, W)`, the filter `(K, C, R, S)`, and the output `(N, K, P, Q)`. The arrangement uses
```python
input.tile((1, *filter.shape[1:]), strides=(-1, -1, 1, 1))
```
to build sliding windows over height and width, then applies `squeeze`, `ravel`, and `flatten` operations to form a 2D tensor with shape `(N * P * Q, C * R * S)`. The filter is flattened and permuted to shape `(C * R * S, K)`, and the output is permuted and flattened to shape `(N * P * Q, K)`. The resulting computation is then passed through the existing `mm.arrangement` and `mm.application`, so no new application function is required [2507.11978].

The paper also states that the same style is applied to `bmm`, `rms_norm`, `rotary_position_embedding`, `scaled_dot_product_attention` using FlashAttention-2, `silu`, and `softmax`. The detailed listings are not all reproduced in the source material, but the characterization is consistent: tiling and alignment are handled through TOM, while arithmetic logic remains serial and tile-local [2507.11978].

## 6. Performance and programmability evidence

The evaluation setup reported in the paper uses an NVIDIA A100 80GB PCIe, Triton 3.0.0, and PyTorch 2.4.1. The microbenchmarks implement the same algorithms in NineToothed and Triton for ten kernels: `add`, `addmm`, `bmm`, `conv2d` using implicit GEMM, `mm`, `rms_norm`, `rotary_position_embedding`, `scaled_dot_product_attention` using FlashAttention-2, `silu`, and `softmax`. All tensors in these microbenchmarks are `float16` [2507.11978].

For single-kernel performance, the reported relative difference between NineToothed and Triton ranges from **–1.58% to +3.93%**, with an average of **+0.37%**. The conclusion drawn in the source is that the generated kernels are essentially as fast as hand-written Triton kernels. PyTorch is shown as a reference, but the source explicitly identifies the NineToothed-versus-Triton comparison as the meaningful one [2507.11978].

For end-to-end inference, the paper uses `deepseek-ai/DeepSeek-R1-Distill-Llama-8B` with batch size 2, input 32 tokens, and output lengths `{128, 512, 2048}`. PyTorch modules are replaced with custom kernels in both NineToothed and Triton for `Attention` including rope, `Linear`, `RMSNorm`, and `SiLU`. Throughput is measured in tokens/s as the average of 3 runs after warmup. The reported relative difference between NineToothed and Triton ranges from **–5.32% to +0.33%**, with an average of **–1.79%**, which the source characterizes as near-identical end-to-end performance, slightly slower on average but within a few percent [2507.11978].

The programmability results are framed through code metrics: LOC, LLOC, SLOC, cyclomatic complexity `G`, Halstead vocabulary `η`, length `N`, volume `V`, difficulty `D`, and Maintainability Index `MI`. NineToothed’s Halstead volume ranges from **0.25% to 56.33%** of the Triton version across the tested kernels; LOC, LLOC, and SLOC are generally lower, especially for `conv2d`, `addmm`, `bmm`, `mm`, and `sdpa`; cyclomatic complexity remains comparable and generally low; Halstead metrics are consistently smaller; and `MI` is higher for NineToothed in all tested kernels [2507.11978].

A specific example given is `conv2d`, where Triton has `LOC=129, LLOC=56, SLOC=110, V=1625.34, MI=39.11`, while NineToothed has `LOC=25, LLOC=18, SLOC=16, V=4.00, MI=68.13`. The stated interpretation is that TOM and arrange-and-apply shorten kernels substantially by abstracting away pointer arithmetic and `load`/`store` boilerplate [2507.11978].

## 7. Scope, limitations, and relation to adjacent systems

The paper’s stated scope is tiled tensor computations typical in deep learning, including GEMM, convolutions, attention, and normalization. It also identifies constraints. NineToothed is designed primarily for tiled tensor computations; very irregular kernels or those with dynamic control-flow-heavy structure may not fit as naturally into arrange-and-apply. Expert Triton programmers may still hand-tune low-level details, fuse operations, or exploit hardware-specific quirks beyond what NineToothed exposes. Moreover, NineToothed does not automatically search for optimal tilings; programmers still choose tile sizes such as `BLOCK_SIZE_M`, `BLOCK_SIZE_N`, and `BLOCK_SIZE_K` [2507.11978].

The system also inherits Triton’s backend support. The source states that it relies on Triton’s supported backends, currently GPUs such as NVIDIA, and does not mention CPU targets or accelerators beyond Triton’s ecosystem. Symbolic shapes are supported, but the actual runtime values must be available when the kernel executes, and compile-time constants such as `BLOCK_SIZE` are designated with `constexpr=True` [2507.11978].

Relative to Triton, NineToothed is positioned as a higher-level serial programming model that treats Triton’s program abstraction as an implementation detail while automatically emitting Triton code with the same parallel workflow. Relative to Graphene, it borrows hierarchical tensors and tiling ideas but reorients them toward a DSL for machine learning kernel developers rather than backend compiler specialists. Relative to Numba, it is specialized for deep learning kernels, PyTorch integration, and tensor-oriented tiling abstractions rather than general-purpose numerical JIT compilation with explicit CUDA programming concepts. Relative to TVM, Halide, XLA, and MLIR-based systems, its novelty is identified as serial semantics at tile level combined with tensor-oriented metaprogramming as the vehicle for schedule expression [2507.11978].

The implementation is reported as open source at `https://github.com/InfiniTensor/ninetoothed`. The paper’s overall conclusion is that serial semantics, TOM, and arrange-and-apply together provide a higher-level interface on top of Triton that reduces code volume and improves maintainability while preserving performance close to hand-written Triton kernels [2507.11978].

Source: https://www.emergentmind.com/topics/ninetoothed