diff --git a/docs/api/attention.md b/docs/api/attention.md index ae04883e..cf1658fa 100644 --- a/docs/api/attention.md +++ b/docs/api/attention.md @@ -53,6 +53,12 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] +::: tileops.attention.MLAVarlenFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ## Native Sparse Attention ::: tileops.attention.NSACompressedVarlenFwdOp @@ -81,6 +87,12 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] +::: tileops.attention.DSAPagedFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ## Attention indexing ::: tileops.attention.FP8LightningIndexerFwdOp @@ -89,6 +101,14 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] +## Merging attention states + +::: tileops.attention.MergeAttentionStatesFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ## Paged KV Cache ::: tileops.attention.PagedKVCacheWriteFwdOp @@ -96,3 +116,9 @@ what runs when you call `op(...)`. show_root_heading: true heading_level: 3 members: ["__init__", "forward"] + +::: tileops.attention.PagedKVCacheGatherFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/api/index.md b/docs/api/index.md index 3e1ceedb..5cf3f254 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -17,7 +17,7 @@ The pages run in the order a reader meets the ops in a model: the pointwise transforms and the positional rotation, then the axis reductions and the normalizations built on them, the windowed kernels, the matmul and the quantization around it, then attention, the expert routing after it and sampling, -then the sequence-model kernels. FFT, mHC and Engram follow, and Trace, +then the sequence-model kernels. FFT, Transform, mHC and Engram follow, and Trace, a tool rather than an op, comes last. The Benchmarks pages use the same order. | Page | What it covers | @@ -28,14 +28,15 @@ a tool rather than an op, comes last. The Benchmarks pages use the same order. | [Normalization](normalization.md) | RMSNorm, LayerNorm, GroupNorm, BatchNorm and the fused variants | | [Pooling](pool.md) | average, max and adaptive pooling, with and without indices, plus the chunked sequence mean | | [Convolution](convolution.md) | forward convolution over 1D, 2D and 3D inputs | -| [GEMM](linear-algebra.md) | dense matmul — plain, batched, and the fp8 variants | +| [GEMM](linear-algebra.md) | dense matmul — plain, batched and grouped, with the fp8, W4A16 and W4A8 variants | | [Quantization & Dequantization](quantization.md) | INT8, FP8 and INT4 quantization, and INT8 dequantization | -| [Attention](attention.md) | forward and backward attention, including the paged and decode kernels | +| [Attention](attention.md) | forward and backward attention, including the paged and decode kernels, merging attention states, and writing and gathering the paged KV cache | | [MoE](moe.md) | the routed mixture-of-experts FFN and its separately callable stages | | [Top-k & Sampling](sampling.md) | top-k selection, logits masks (top-k, top-p, min-p) and sampling, including chain speculative sampling | | [Linear Attention](linear-attention.md) | DeltaNet, Gated DeltaNet, Kimi Delta Attention and Gated Linear Attention | -| [Mamba](mamba.md) | the SSD scan, its decode step, and the chunked forms | +| [Mamba](mamba.md) | the causal convolution and the SSD scan, their decode steps, and the chunked forms | | [FFT](fft.md) | the discrete transform | +| [Transform](transform.md) | the fast Walsh-Hadamard transform along the last axis | | [mHC](mhc.md) | Manifold-Constrained Hyper-Connections — the pre/post pair around a layer | | [Engram](engram.md) | the Engram GateConv pair and its decode step | | [Trace](trace.md) | the in-kernel timeline tracer, a tool rather than an op | diff --git a/docs/api/linear-algebra.md b/docs/api/linear-algebra.md index 0cbb66c7..f2e6d9d9 100644 --- a/docs/api/linear-algebra.md +++ b/docs/api/linear-algebra.md @@ -31,6 +31,14 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false +::: tileops.gemm.GemmW4A8FwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + merge_init_into_class: false + show_signature_annotations: false + ## Batched GEMM ::: tileops.gemm.BmmFwdOp diff --git a/docs/api/mamba.md b/docs/api/mamba.md index 4a839384..3bb566f1 100644 --- a/docs/api/mamba.md +++ b/docs/api/mamba.md @@ -13,6 +13,14 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] +## Causal convolution + +::: tileops.mamba.CausalConv1dFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ## SSD stages ::: tileops.mamba.SSDChunkCumsumFwdOp @@ -47,6 +55,12 @@ what runs when you call `op(...)`. ## Decode +::: tileops.mamba.CausalConv1dUpdateFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] + ::: tileops.mamba.SSDRecurrentFwdOp options: show_root_heading: true diff --git a/docs/api/transform.md b/docs/api/transform.md index e4fe7858..e5b5cf71 100644 --- a/docs/api/transform.md +++ b/docs/api/transform.md @@ -6,12 +6,11 @@ quantized; the quantization that follows one is on the ## Walsh-Hadamard -`HadamardTransformFwdOp` is specified and not yet implemented, so it has no -constructor to document here. The signature it will serve is in -`src/tileops/manifest/spec/transform.yaml`: a fast Walsh-Hadamard transform along -the last axis scaled by $1/\sqrt{n}$, where $n$ is `base_order` times a power of -two, in float16 or bfloat16 with the butterfly accumulated in float32. It is the -rotation step of rotation-based quantization, which spreads outliers across +The rotation step of rotation-based quantization, which spreads outliers across channels so a lower bit width stays accurate. -This page gains the op's `__init__` and `forward` when a kernel serves it. +::: tileops.transform.HadamardTransformFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/performance-guides/trace-timeline.zh.md b/docs/performance-guides/trace-timeline.zh.md index 775c307e..827edc38 100644 --- a/docs/performance-guides/trace-timeline.zh.md +++ b/docs/performance-guides/trace-timeline.zh.md @@ -6,7 +6,7 @@ 1. 在 kernel 代码中用标记(`trace.range`、`trace.group` 等)标注要记录的区间。 1. 标记**总是以占位的形式生成**。构建时,kernel 要么被 **lower**:标记变成真正调用 `clock64()` 记录时间的代码,并在输出末尾多出一个 `slots`;要么被 **strip**:标记变成空操作,生成的 CUDA 与未插桩的版本完全相同。 -1. 运行时,`trace.run` 执行 kernel,解码 `slots` 缓冲区,并写出一个自包含的 Plotly HTML 时间线。 +1. 运行时,`trace.run` 执行 kernel;开启追踪时,它解码 `slots` 缓冲区,并写出一个自包含的 Plotly HTML 时间线。 1. 由进程内的开关 `trace.enable()` 决定 lower 还是 strip,因此关闭时追踪**没有任何开销**,标记可以留在生产代码中。 时间戳来自 `clock64()`,即每个 SM 的周期计数器。 @@ -81,7 +81,7 @@ def forward(self, a, b): ``` 1. 按开关构建对应的 kernel:[`trace.enabled`](../api/trace.md#tileops.trace.api._Trace.enabled) 决定选用带追踪的版本还是 strip 后的版本,它也是标记 `(7)` 所说的缓存键的一部分。 -1. [`trace.run(compiled, inputs, stem=...)`](../api/trace.md#tileops.trace.api._Trace.run) 运行 kernel。带追踪时,它把末尾的 `slots` 拆出来解码,写出 `debug/.html`,每次调用都生成一个不会重名的新文件。它内部由 [`decode`](../api/trace.md#tileops.trace.api._Trace.decode) 和 [`dump`](../api/trace.md#tileops.trace.api._Trace.dump) 组成,两者也可以直接调用。 +1. [`trace.run(compiled, inputs, stem=...)`](../api/trace.md#tileops.trace.api._Trace.run) 运行 kernel。带追踪构建的 kernel,末尾的 `slots` 总会被拆出来;开启追踪时,它被解码并写进 `debug/.html`,每次调用都生成一个不会重名的新文件。它内部由 [`decode`](../api/trace.md#tileops.trace.api._Trace.decode) 和 [`dump`](../api/trace.md#tileops.trace.api._Trace.dump) 组成,两者也可以直接调用。 ## 开启追踪 { #enabling-tracing }