diff --git a/docs/api/attention.md b/docs/api/attention.md index 1476afbe..bb5e4eab 100644 --- a/docs/api/attention.md +++ b/docs/api/attention.md @@ -45,6 +45,12 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] +::: tileops.ops.attention.gqa.GroupedQueryAttentionDenseFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ::: tileops.ops.attention.gqa.GroupedQueryAttentionPrefillFwdOp options: show_root_heading: true @@ -122,3 +128,11 @@ what runs when you call `op(...)`. show_root_heading: true heading_level: 3 members: ["__init__", "forward"] + +## Attention indexing + +::: tileops.ops.fp8_lightning_indexer.FP8LightningIndexerFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/convolution.md b/docs/api/convolution.md new file mode 100644 index 00000000..ed8851c6 --- /dev/null +++ b/docs/api/convolution.md @@ -0,0 +1,26 @@ +# Convolution Operators + +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## Forward convolution + +::: tileops.ops.convolution.Conv1dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.convolution.Conv2dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.convolution.Conv3dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/dropout.md b/docs/api/dropout.md new file mode 100644 index 00000000..adffdf6e --- /dev/null +++ b/docs/api/dropout.md @@ -0,0 +1,14 @@ +# Dropout Operators + +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## Dropout + +::: tileops.ops.dropout.DropoutFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/index.md b/docs/api/index.md index 0d4fbab1..4447410e 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -14,22 +14,26 @@ d = op(a, b) # the specialized kernel is built on first ``` The pages are ordered along the stack a caller works through: the dense matmul first, -then attention and its linear replacements, then the elementwise and reduction -primitives, then FFT. +then attention, its position encodings and its linear replacements, then the +elementwise and reduction primitives, then the convolutional family and FFT. | Page | What it covers | | --- | --- | | [GEMM](linear-algebra.md) | dense matmul — plain, batched, and the fp8 variants | | [Attention](attention.md) | forward and backward attention, including the paged and decode kernels | +| [RoPE](rope.md) | rotary position embedding — NeoX and interleaved layouts, Llama 3.1, YaRN, LongRoPE | | [Linear Attention](linear-attention.md) | the linear-attention family: DeltaNet, GLA, KDA, and their kin | | [Mamba](mamba.md) | the SSD scan, its decode step, and the chunked forms | | [MHC](mhc.md) | multi-head compression | | [Normalization](normalization.md) | RMSNorm, LayerNorm, GroupNorm, BatchNorm and the fused variants | | [Elementwise](elementwise.md) | unary and binary maps, activations, and the in-place forms | | [Reduction](reduction.md) | sums, extrema, arg-reductions, cumulative scans, softmax | +| [Convolution](convolution.md) | forward convolution over 1D, 2D and 3D inputs | +| [Pooling](pool.md) | average, max and adaptive pooling, with and without indices | +| [Dropout](dropout.md) | dropout with deterministic replay | | [FFT](fft.md) | the discrete transform | -| [Quantization](quantization.md) | fp8 quantization — interface not yet stable | -| [Top-k](topk.md) | top-k selection — interface not yet stable | +| [Quantization](quantization.md) | fp8 quantization | +| [Top-k](topk.md) | top-k selection | | [Trace](trace.md) | the in-kernel timeline tracer, a tool rather than an op | Two things this reference does not carry: diff --git a/docs/api/pool.md b/docs/api/pool.md new file mode 100644 index 00000000..7c92e898 --- /dev/null +++ b/docs/api/pool.md @@ -0,0 +1,84 @@ +# Pooling Operators + +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## Average pooling + +::: tileops.ops.pool.AvgPool1dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.AvgPool2dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.AvgPool3dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +## Max pooling + +::: tileops.ops.pool.MaxPool1dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.MaxPool1dIndicesFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.MaxPool2dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.MaxPool2dIndicesFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.MaxPool3dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.MaxPool3dIndicesFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +## Adaptive pooling + +::: tileops.ops.pool.AdaptiveAvgPool2dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.AdaptiveMaxPool2dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.pool.AdaptiveMaxPool2dIndicesFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/quantization.md b/docs/api/quantization.md index 2a9ea7c6..b64e823f 100644 --- a/docs/api/quantization.md +++ b/docs/api/quantization.md @@ -1,6 +1,14 @@ -# Quantization +# Quantization Operators -!!! note "Coming soon" - The `tileops.ops.fp8_quant` module is under active development in - [`tile-ai/TileOPs`](https://github.com/tile-ai/TileOPs) and will be - documented here once the public interface stabilizes. +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## FP8 quantization + +::: tileops.ops.fp8_quant.FP8QuantFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/rope.md b/docs/api/rope.md new file mode 100644 index 00000000..3ec188e3 --- /dev/null +++ b/docs/api/rope.md @@ -0,0 +1,48 @@ +# RoPE Operators + +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## NeoX layout + +::: tileops.ops.rope.RopeNeoxFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.rope.RopeNeoxPositionIdsFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +## Interleaved layout + +::: tileops.ops.rope.RopeNonNeoxFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +## Scaled frequencies + +::: tileops.ops.rope.RopeLlama31FwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.rope.RopeYarnFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + +::: tileops.ops.rope.RopeLongRopeFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/topk.md b/docs/api/topk.md index 6da26acf..8f403dc9 100644 --- a/docs/api/topk.md +++ b/docs/api/topk.md @@ -1,6 +1,14 @@ -# Top-k +# Top-k Operators -!!! note "Coming soon" - The `tileops.ops.topk_selector` module is under active development in - [`tile-ai/TileOPs`](https://github.com/tile-ai/TileOPs) and will be - documented here once the public interface stabilizes. +Every op on this page is used the same way: construct it once, then call it. The +constructor takes what the kernel is compiled with; the call takes the tensors. +Both are documented under each op — `__init__` and `forward`, where `forward` is +what runs when you call `op(...)`. + +## Top-k selection + +::: tileops.ops.topk_selector.TopkSelectorFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/mkdocs.yml b/mkdocs.yml index 04c96de9..f3f6373a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -37,19 +37,23 @@ nav: - Using torch.compile: torch-compile.md - Adding a Backend: backends.md # Ordered along the stack a reader is working through: the dense matmul, then - # attention and its linear replacements, then the elementwise and reduction - # primitives, then FFT. The two pages still awaiting a stable interface sit - # after those, and Trace last — it is a tool, not an op family. + # attention, its position encodings and its linear replacements, then the + # elementwise and reduction primitives, then the convolutional family, then + # FFT, quantization and top-k, and Trace last — it is a tool, not an op family. - API Reference: - api/index.md - GEMM: api/linear-algebra.md - Attention: api/attention.md + - RoPE: api/rope.md - Linear Attention: api/linear-attention.md - Mamba: api/mamba.md - MHC: api/mhc.md - Normalization: api/normalization.md - Elementwise: api/elementwise.md - Reduction: api/reduction.md + - Convolution: api/convolution.md + - Pooling: api/pool.md + - Dropout: api/dropout.md - FFT: api/fft.md - Quantization: api/quantization.md - Top-k: api/topk.md @@ -156,6 +160,10 @@ plugins: GEMM: GEMM Quantization: Quantization Attention: Attention + RoPE: RoPE + Convolution: Convolution + Pooling: Pooling + Dropout: Dropout Linear Attention: Linear Attention Mamba: Mamba FFT: FFT