NineToothed: DSL for ML Kernels

• NineToothed is a domain-specific language that simplifies machine learning kernel development by using serial semantics and symbolic tensor metaprogramming.
• It employs an arrange-and-apply paradigm where tile arrangement is separated from serial per-tile computation, eliminating explicit parallel indexing and pointer arithmetic.
• Performance benchmarks show near-identical speeds to hand-written Triton kernels while significantly reducing code complexity and maintenance overhead.

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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,)) (Huang et al., 16 Jul 2025).

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:

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 (Huang et al., 16 Jul 2025).

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

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).

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 (Huang et al., 16 Jul 2025).