Metadata-Version: 2.4
Name: piper-kernels
Version: 0.7.0
Summary: Reusable PyTorch inference operators and Triton kernels.
Keywords: inference,pytorch,quantization,triton
License-Expression: Apache-2.0
License-File: LICENSE
License-File: NOTICE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: torch>=2.14
Requires-Dist: torchao>=0.17 ; extra == 'convrot'
Requires-Dist: numpy>=2 ; extra == 'nvfp4'
Requires-Dist: torchao>=0.18 ; extra == 'nvfp4'
Requires-Dist: triton-windows>=3.8,<3.9 ; sys_platform == 'win32' and extra == 'triton'
Requires-Python: >=3.13
Project-URL: Homepage, https://github.com/Boffee/piper-kernels
Project-URL: Repository, https://github.com/Boffee/piper-kernels.git
Project-URL: Issues, https://github.com/Boffee/piper-kernels/issues
Project-URL: Changelog, https://github.com/Boffee/piper-kernels/blob/main/CHANGELOG.md
Provides-Extra: convrot
Provides-Extra: nvfp4
Provides-Extra: triton
Description-Content-Type: text/markdown

# Piper Kernels

Reusable PyTorch inference operators and optimized kernels for the Piper ecosystem and
other consumers.

Piper Kernels requires Python 3.13 or newer and PyTorch 2.14 or newer.

The package owns operator semantics, portable PyTorch references, tensor subclasses,
and optimized backends. It deliberately does not know about model repositories,
checkpoint metadata, pipeline frameworks, or device-offloading policy.

Fused ConvRot preparation, GGUF conversion, and NVFP4 weight updates use FP32 arithmetic
instead of reproducing the extra FP16/BF16 rounding of an eager PyTorch composition.
Their portable references use FP32 arithmetic. Lower-precision tensor-core operands and
compact workspaces remain where they reduce storage or execution cost.
Sparse attention's fused coarse residual keeps the fine output and gated coarse contribution
in FP32 until the final BF16 store.

## Validation contract

This is a library-wide API and development contract. It applies to inference operators,
dispatch, compiler rewrites, fake/meta implementations, and weight wrappers, including
construction, reconstruction, views, and device moves. New implementations must preserve it.

- Validation may inspect host metadata: shapes, dtypes, devices, layouts/strides, gradient
  flags, and Python configuration values. Keep checks for supported storage and operations.
- Numerical tensor contents are caller preconditions on every device. For example, callers
  must supply a finite positive ConvRot INT8 static input scale. This requirement does not
  promise runtime rejection of zero, negative, NaN, or infinite supplied scales.
- These paths must not inspect tensor contents solely for input validation. Do not introduce
  host readbacks (`.item()`, `bool(tensor)`, `.cpu()`), synchronization, tensor scans/reductions,
  device assertions, validation kernel launches, or temporary device allocations for that
  purpose. Validation must work without tensor contents during tracing and fake/meta execution
  and must not introduce barriers to CUDA graph capture.
- Exporters, checkpoint loaders, and other callers own any required numerical validation at
  ingestion. Any dedicated tensor-content validation API must be explicitly invoked outside
  inference, compilation, and weight wrapping; it must not run implicitly in those paths.

GPU value readbacks synchronize execution, and additional validation kernels and allocations
consume inference time and memory. Guards needed by the numerical algorithm remain required:
for example, deriving a usable dynamic scale for an all-zero input is valid-input handling.
Documented quantization/conversion work outside the paths above may also check the values it
uses to construct a quantized representation. Neither permits adding content-validation work
to inference. Callers can rely on this boundary when composing and capturing kernels.

## Operators

| Package | Role |
|---|---|
| `piper_kernels` | Public dense and sparse Piper Attention plus SageAttention2++ operators |
| `piper_kernels.attention` | Attention dispatch, portable references, and optimized backends |
| `piper_kernels.weights` | Quantized weight formats, conversion, updates, and sharding |
| `piper_kernels.linear` | Linear operators and optimized backends |
| `piper_kernels.linear.convrot` | ConvRot linear operators and compiler integrations |

## Triton setup

Install the optimized backends with `piper-kernels[triton]`, or include ConvRot's tensor
format with `piper-kernels[convrot,triton]`. On Linux, first install a CUDA or ROCm PyTorch
distribution with its matching Triton; Piper does not pin a competing Linux Triton version.
The extra selects Triton 3.8 via
[`triton-windows`](https://github.com/triton-lang/triton-windows) on 64-bit Windows.

Optimized Windows execution requires Windows 10 or 11, a supported NVIDIA GPU with a
current driver, and the Visual C++ Redistributable for Visual Studio 2015-2022. The
Windows wheel bundles its CUDA toolchain and TinyCC, so a separate CUDA toolkit or Visual
Studio install is not required for Piper's Triton kernels. The base package remains
portable and does not require either Triton distribution.

## Shared weight formats

Import weight types independently of the operator that executes them:

```python
from piper_kernels.weights.convrot.int8 import ConvRotInt8Tensor
from piper_kernels.weights.convrot.nvfp4 import ConvRotNVFP4Tensor
from piper_kernels.weights.nvfp4 import PiperNVFP4Tensor
```

The weight packages own packed storage, scales, quantization/dequantization,
GGUF conversion, in-place updates, and sharding. Importing, quantizing, loading,
or updating a weight does not load Piper's linear operators. Tensor dispatch
loads linear execution when `torch.nn.functional.linear` or a matrix product
uses the weight. Operator kernels and graph optimizations remain under `linear`;
reusable rotation and packing primitives live under `_triton`.

The former tensor exports under `linear` have been removed. Callers must update
imports together with their Piper Kernels dependency. Packed checkpoint data and
scale layouts are unchanged, including direct wrapping of mmap storage. Pickled
Python tensor subclasses now use the new class module paths; checkpoints saved
with the old paths must be re-exported. No legacy import aliases are provided.

## ConvRot INT8

INT8-only fusion packages use the explicit `convrot_int8` prefix:

| Previous package/helper prefix | Current prefix |
| --- | --- |
| `convrot_swiglu_ffn` | `convrot_int8_swiglu_ffn` |
| `convrot_sparse_piper` | `convrot_int8_sparse_piper` |
| `convrot_sage_qk` | `convrot_int8_sage_qk` |

Update imports under `piper_kernels.fusions`, compile-option helper names, and any direct
`torch.ops.piper_kernels` calls using these prefixes. The old names are not retained as aliases.
Shared weight rotation lives under `weights.convrot`, with reusable accelerator
primitives under `_triton`. Linear operators and compile options live under `linear.convrot`.

Quantize a dense weight, or wrap existing checkpoint storage without dequantizing it,
then use the resulting tensor as a normal linear weight:

```python
import torch

from piper_kernels.weights.convrot.int8 import ConvRotInt8Tensor
from piper_kernels.linear.convrot import convrot_int8_compile_options, convrot_int8_linear

weight = ConvRotInt8Tensor.from_hp(dense_weight, group_size=256)
checkpoint_weight = ConvRotInt8Tensor.from_quantized(
    qdata,
    scale,
    group_size=256,
    logical_dtype=torch.bfloat16,
)
output = torch.nn.functional.linear(activation, weight, bias)

# Let Inductor optimize repeated inputs and absorb supported input activations.
compiled_block = torch.compile(block, options=convrot_int8_compile_options())

# Optionally fuse a raw [up | gate] SwiGLU input with ConvRot preparation.
mlp_output = convrot_int8_linear(up_gate, weight, bias, activation_fn="swiglu")

# GELU with tanh approximation uses the same activation/preparation boundary.
mlp_output = convrot_int8_linear(activation, weight, bias, activation_fn="gelu_tanh")

# The explicit API also supports an ordinary linear.
output = convrot_int8_linear(activation, weight, bias)

# In-place low-rank update with the standard Tensor.addmm_ contract.
weight.addmm_(lora_b, lora_a, alpha=lora_strength)

# Reproducible stochastic terminal-code selection for a quantized LoRA merge.
weight.addmm_(lora_b, lora_a, alpha=lora_strength, rounding_seed=seed)

# In-place full-rank logical update for a materialized adapter delta.
weight.add_(dense_update, alpha=adapter_strength, rounding_seed=seed)
```

Use `from_quantized(..., logical_dtype=...)` to construct a weight from checkpoint storage.
Pass `act_per_tensor_scale=input_scale` to either constructor to use a calibrated static input
scale. It must be a finite positive FP32 scalar tensor on the weight device, calibrated after
any input activation and ConvRot rotation. The scalar moves and serializes with the weight;
conversion and dequantization of the weight do not depend on its value. Omitting it selects
dynamic per-row input scaling. Static preparation skips the row maximum reduction and preserves
the existing prepared-input and matrix-multiply contracts.

For a weight with shape `[out_features, in_features]`, ordinary and GELU-tanh inputs have
shape `[..., in_features]`; the SwiGLU input has shape `[..., 2 * in_features]`. The output
always has shape `[..., out_features]`.
`convrot_int8_linear(...)` applies an ordinary linear when `activation_fn` is omitted, matching
`torch.nn.functional.linear`. `activation_fn="gelu_tanh"` applies tanh-approximate GELU, while
`activation_fn="swiglu"` computes `up * silu(gate)` from `[up | gate]`. Portable paths use
PyTorch operations; optimized NVIDIA preparation uses shared Triton activation primitives and
native approximate tanh, so GELU preparation may differ from the portable path by one INT8 code
rather than being bitwise identical. NVIDIA preparation uses up to three equal
power-of-two chunks of at most 16,384 columns, fusing rows through 49,152 columns across every
supported ConvRot group size, logical dtype, row count, and NVIDIA target. This selection is
measured on exact SM120 and optimistic on other targets. Larger rows materialize the activation
and retain the same semantics. Both
`F.linear` with a ConvRot INT8 weight and the explicit INT8 entry point are inference-only and
reject autograd inputs.

For compiled inference, `convrot_int8_compile_options()` installs deterministic post-AOT Inductor
rewrites. An exclusive tanh-approximate GELU feeding a ConvRot linear becomes an activated
input-preparation node followed by a prepared linear. This avoids the materialized activated
input and lets its source die before the linear output is allocated.
Separately, two or more ordinary ConvRot linears fed by the same graph value become one explicit
input preparation followed by independent prepared GEMMs at the original operation positions.
Static inputs share preparation only when they use the same scale graph value; different
static scales and dynamic scaling remain separate.
GELU input fusion, SwiGLU FFNs, and sparse-attention region fusions support static, dynamic,
and mixed input scaling. Each projection retains its own input scale. Compatible gate/value
inputs share FFN preparation and a paired GEMM; distinct scales reuse the bounded preparation
workspace for separate projections. Sparse attention shares preparation only across compatible
Q/K/V and coarse-gate inputs, centers V using its own prepared input, and applies the output
projection's scale inside each attention chunk.
Static ConvRot INT8 weights bypass PyTorch's AOTAutograd disk cache because its wrapper
cache key does not distinguish shared from independent input-scale tensors. Compilation,
Inductor caching, and reuse of the compiled graph remain supported.
Prepared tensors are ordinary graph values—there is no hidden runtime cache—and unmatched,
eager, and training paths remain unchanged. Existing post-grad compiler passes in the supplied
options mapping are preserved. Pass the result through `torch.compile(options=...)`; PyTorch
treats `mode` and `options` as mutually exclusive, so do not also supply `mode`.

The ConvRot INT8, NVFP4, and ConvRot NVFP4 FFN compiler integrations support FP16 and BF16
activations. Their `*_swiglu_ffn_compile_options()` helpers match separate gate, value, and down
projections. Their `*_gelu_ffn_compile_options()` helpers fold an exclusive
`down(gelu(up(input), approximate="tanh"))` region. Each helper installs FFN fusion before
ordinary linear rewriting and preserves every projection's static or dynamic input scale.
Feature-width-aware GELU row chunks bound reusable temporary storage near 1 GiB for INT8 and
512 MiB for NVFP4 while amortizing long-sequence launches. Fused activation preparation retains
FP32 arithmetic and quantizes directly into the down-projection input; outputs retain the input
dtype.

ConvRot NVFP4 FFN fusion includes SwiGLU or tanh-GELU, rotation, and NVFP4 preparation for the
down projection. It reads the FFN's private projection workspace directly. Dynamic
preparation reuses rotated FP32 values in a scratch buffer capped at 32 MiB per
chunk and recomputes them in small tiles above that limit; static preparation
uses those tiles directly. The scratch limit was selected from RTX 5090
measurements: reuse helped smaller working sets, while larger buffers added
enough memory traffic to favor recomputation. See the
[FFN benchmark](benchmarks/README.md#nvfp4-ffn) for measurement commands.
Source dynamic scaling covers the full input, while down dynamic scaling remains
per chunk. The shared NVFP4 scale reduction handles arrays of up to 8,192 elements
in one GPU launch.

The cross-operator ConvRot-to-sparse-Piper optimization is enabled explicitly by importing
`convrot_int8_sparse_piper_compile_options` from
`piper_kernels.fusions.convrot_int8_sparse_piper`. It installs the fusion pass before the ordinary
ConvRot pass. On exact SM120, it recognizes a compatible H3-style region containing three
ConvRot Q/K/V projections with optional FP16/BF16/FP32 bias, D64/D128 RMSNorm and split-half
RoPE for Q/K, followed by `sparse_piper_attention`. The rewrite shares input preparation and
emits quantized Q/K/V plus routing summaries directly, avoiding the three materialized projection
outputs. Arbitrary logical sequence lengths are written directly into internally K64-padded
attention storage; only the final projection tile is masked, and the result retains the exact logical
length. It fails closed for unsupported shapes, layouts, or parameters; the ordinary ConvRot and
sparse-attention APIs remain independent.

Because no projected activation is externally observable in the fused region, projection,
RMSNorm, and RoPE stay in FP32 until the final INT8 Q/K/V encoding. This removes otherwise
redundant FP32-to-BF16-to-FP32 round trips without materializing FP32 activation tensors.

The internal `piper_kernels.fusions.projected_qk` layer owns projection-independent RMSNorm and
RoPE. The existing Sage Q/K quantization layer owns signed-Hadamard grouped Q/K encoding shared
with dense Piper, while `piper_kernels.attention.kernels.sparse_piper` owns only sparse Piper's
tile-scaled V encoding. `piper_kernels.fusions.convrot_int8_sage_qk` adapts ConvRot projection tiles
to those boundaries and owns ConvRot validation; the explicit sparse fusion adds routing summaries,
storage, and graph rewriting. Another projection backend can therefore compose the same pieces
without depending on ConvRot internals or adding a backend protocol to attention.

`addmm_` computes `weight = beta * weight + alpha * (mat1 @ mat2)`, while `add_`
accepts an exact-shape dense logical update and computes `weight = weight + alpha * update`.
Both operations requantize the result and preserve the ConvRot tensor and quantized storage
identities, allowing offload integrations to keep their existing buffers. Repeated updates are
lossy, so reload a pristine base weight before changing or removing a previously merged adapter.
Passing an unsigned 64-bit `rounding_seed` stochastically selects one of the two adjacent
INT8 codes with probability proportional to distance, without changing the deterministic
row scales or consuming PyTorch's process-global random-number generator. Omitting the seed
retains nearest-integer rounding. This is an inference operation and does not support
autograd. Torch and Triton each replay for a fixed seed, device, and backend; their random
samples are not promised to match each other or different Triton versions byte-for-byte.
The standard `add_(update, alpha=...)` signature is compatible with `torch.compile`; its
`rounding_seed` extension is intended for eager merge code because Dynamo enforces the built-in
`Tensor.add_` keyword schema while tracing.

The operator selects its Triton implementation on supported NVIDIA CUDA or Linux AMD HIP devices and otherwise
uses the portable PyTorch reference. Install the tensor format and optimized backend with
`piper-kernels[convrot,triton]`. The base package does not require TorchAO or Triton, and
attention-only consumers do not inherit the TorchAO dependency.

### ROCm INT8 support

The AMD implementation lives in `linear/convrot/int8/_amd/`, alongside `_nvidia/`.
Both implement the same preparation/projection interface; reusable INT8 arithmetic lives in
`int8/_kernels/`. Public operators and compiler rewrites do not select launch schedules.

ROCm coverage includes ordinary, GELU-tanh, and SwiGLU input preparation; INT8 linear and
prepared/paired projections; caller-owned output buffers; dense and low-rank weight updates;
and base `torch.compile` preparation sharing. FP16, BF16, and FP32 are supported.
The shared chunked INT8 SwiGLU and GELU FFNs also run on ROCm, including indexed gated updates
and automatic fusion of compatible FP16/BF16 graphs via their compile-option helpers.
The RX 9070 XT (`gfx1201`) has on-device validation. `gfx942`, `gfx1100`, `gfx1151`,
and `gfx1200` have compiler coverage only, not hardware correctness or performance validation.
Unknown AMD architectures retain the portable reference for linear execution.

Rotation, quantization, dequantization, and weight-update arithmetic are shared.
Standalone preparation and updates do not require a tuned INT8 GEMM target: they
use conservative shared Triton launchers when the installed driver can handle the
device, and PyTorch otherwise. Wide rows use the PyTorch path to bound kernel resources.
These generic paths require the device's underlying operations and dtypes, not a GPU-model
allowlist; tuned preparation and GEMM policies remain accelerator-specific.

AMD fused preparation supports widths through 16,384; larger widths use separate rotation
and quantization. RDNA4 uses BF16 ragged-row chunks and its own measured GEMM schedule.
FP32 reduction ordering and fused activations can differ from the reference at INT8 rounding
boundaries. GGUF-to-INT8 conversion uses a shared fused decoding/rotation kernel, using
fused rows through 8,192 columns on ROCm and bounded tiled conversion for wider rows,
including on AMD targets without a tuned GEMM policy. This integration does
not enable ROCm dequantized-input means, specialized attention fusions, attention kernels,
or NVFP4 kernels.

The repository's default `uv` development sources still select CUDA; use a separate
ROCm environment rather than `uv sync` in that environment.

## ConvRot INT8 Conv3D

`ConvRotInt8Tensor` also supports causal 3×3×3 convolution weights through the same
`from_hp()`, `from_quantized()`, and `dequantize()` API. It carries packed INT8
weights, FP32 weight scales, and an optional FP32 `act_per_tensor_scale` tensor.
`piper_kernels.conv3d.convrot.int8.ConvRotInt8Conv3d` consumes that weight with a
fixed activation scale. The optimized backend targets SM120, with a portable
reference elsewhere. Loading contiguous checkpoint tensors preserves mmap
storage. H3 encoder compile options fuse framewise GroupNorm, SiLU, padding,
and residuals around explicitly installed quantized convolutions.

See [ConvRot INT8 Conv3D](src/piper_kernels/conv3d/convrot/int8/README.md) for
checkpoint conversion, loading, supported shapes, and engine integration.

## NVFP4 construction

Piper's ordinary and ConvRot NVFP4 wrappers can quantize a floating-point weight without
exposing TorchAO storage construction to the caller:

```python
from piper_kernels.weights.convrot.nvfp4 import ConvRotNVFP4Tensor
from piper_kernels.weights.nvfp4 import PiperNVFP4Tensor
from torchao.prototype.mx_formats.nvfp4_tensor import QuantizeTensorToNVFP4Kwargs

activation_quantization = QuantizeTensorToNVFP4Kwargs(
    block_size=16,
    is_swizzled_scales=True,
    use_triton_kernel=False,
    use_dynamic_per_tensor_scale=True,
)

weight = PiperNVFP4Tensor.from_hp(
    dense_weight,
    compute_per_tensor_scale=True,
    is_swizzled_scales=True,
    act_quant_kwargs=activation_quantization,
)
rotated_weight = ConvRotNVFP4Tensor.from_hp(
    dense_weight,
    group_size=64,
    compute_per_tensor_scale=True,
    is_swizzled_scales=True,
    act_quant_kwargs=activation_quantization,
)
```

For ConvRot, the global NVFP4 scale is derived after rotation. This keeps rotation and
quantization in one package-owned operation and prevents callers from accidentally scaling the
logical basis instead of the stored basis. `SUPPORTED_GROUP_SIZES` is exported from
`piper_kernels.weights.convrot` for format-policy validation.

`ConvRotInt8Tensor`, `PiperNVFP4Tensor`, and `ConvRotNVFP4Tensor` support same-shape
`view` and `view_as`, preserving the concrete wrapper, quantization metadata, and shared
storage. Matrix transposes (`t`, `transpose`, `mT`, and `permute`) also share storage and
preserve the represented weight, including ConvRot's rotation axis and NVFP4's packing order.
`as_strided` supports the existing layout or its matrix transpose, with unchanged storage offset.
Other shape/layout changes raise `NotImplementedError`; transposed weights cannot be made
contiguous or updated in place.

`F.linear` on DTensors constructed with `DTensor.from_local(..., run_check=False)` uses
Piper's local quantized linear implementation through the transpose and `mm`/`addmm` path.
`addmm` supports the linear case: `alpha=1`, `beta=1`, and a bias vector with one value per
output feature. Other coefficients and bias shapes raise `NotImplementedError`.
Eager and fullgraph compiled execution support replicated weights, output-feature weight
shards (`Shard(0)`), and input-feature weight shards (`Shard(1)`) with matching activation
placements. ConvRot feature shards must align with rotation groups. Input-feature sharding
produces partial outputs that must be summed.
The numerical reference is the corresponding local quantized computation on each rank.
Redistributing quantized weights is unsupported. Transposed weights support dense activation
matrix products; using one as the weight of another `linear` is unsupported.

To partition an **already quantized full weight**, use
`piper_kernels.weights.sharding.shard_quantized_weight(weight, dim=..., start=..., length=...)`.
It copies packed data and repacks NVFP4 scales without requantizing, preserving the wrapper,
nibble order, global scales, rotation, and activation-quantization configuration. Each shard
owns its tensor storage. Create shards on CPU during loading, then move them to CUDA for
execution, or create them directly on CUDA.

Row partitions (`dim=0`) accept any nonempty contiguous interval, including cuts inside
NVFP4's 128-row scale tiles. Input-channel partitions (`dim=1`) must align to NVFP4's 16-value
blocks and ConvRot's rotation groups. NVFP4 accepts ordinary or swizzled scales, flat or
canonical 2-D, and returns canonical 2-D scales with fresh padding. Empty partitions,
transposed/noncontiguous storage, nonstandard NVFP4 blocks, and per-expert scales are rejected.
Ordinary slice/narrow views remain unsupported. Execution requirements still apply; for
example, ConvRot NVFP4 linear requires swizzled scales.

For standard DTensor plans, install the prepared DTensor parameter before calling
`parallelize_module`. Each rank must load the same quantized full weight. This example
assumes an initialized 1-D mesh and an evenly partitioned weight:

```python
from torch.distributed.tensor import DTensor, Shard
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, parallelize_module

from piper_kernels.weights.sharding import shard_quantized_weight

dim = 0  # 0: output rows / ColwiseParallel; 1: input channels / RowwiseParallel
parts = mesh.size()
assert full_weight.shape[dim] % parts == 0
length = full_weight.shape[dim] // parts
local_weight = shard_quantized_weight(
    full_weight, dim=dim, start=mesh.get_local_rank() * length, length=length
).to(mesh.device_type)
distributed_weight = DTensor.from_local(
    local_weight,
    mesh,
    [Shard(dim)],
    run_check=False,
    shape=full_weight.shape,
    stride=full_weight.stride(),
)
out_features, in_features = full_weight.shape
linear = torch.nn.Linear(in_features, out_features, bias=False, device="meta")
linear.weight = torch.nn.Parameter(distributed_weight, requires_grad=False)
plan = ColwiseParallel() if dim == 0 else RowwiseParallel()
linear = parallelize_module(linear, mesh, plan)
```

Install bias as a DTensor too: `Shard(0)` for output-row sharding, `Replicate()` for
input-channel sharding. By default, `ColwiseParallel` takes replicated inputs and returns
local output shards; `RowwiseParallel` takes input-channel shards and sums partial outputs.
For sequential execution, use local `F.linear` calls and concatenate row-shard outputs or
sum column-shard outputs, adding bias once. Column partitions preserve weight values but
can change dynamic activation scales and accumulation order, so comparisons with an
unsharded quantized linear require numerical tolerances.

With the corresponding `*_swiglu_ffn_compile_options()`, `*_gelu_ffn_compile_options()`,
or `*_sparse_piper_compile_options()`, batched DTensor projections retain the existing FFN,
QKV-preparation, and attention/output-projection fusions. Shared compiler normalization
removes redundant row flatten/restore pairs around semantic linears, including symbolic
leading dimensions. Feature order, quantization metadata, and externally consumed values
are preserved; the existing layout and intermediate-consumer restrictions still apply.

For sharded attention, unwrap QKV projection outputs before the local attention region and
wrap its head-flattened result for the output projection. Each rank computes its own complete
heads or FFN intermediate-feature shard. Sum partial outputs after the output/down projection,
keeping collectives outside the local fused region. The numerical reference is the corresponding
local **fused** quantized computation; fused and unfused quantization boundaries can differ.

Supported eager and compiled NVFP4 linears share the same prepared projection backend,
including ConvRot and the affine projections used by fused SwiGLU FFNs and sparse attention.
Global scales and bias are applied in FP32 before the final FP16/BF16 output conversion.
No-bias and matching-dtype-bias projections fuse this epilogue into GEMM. Mixed bias retains
its dtype until FP32 addition and uses a reusable FP32 workspace bounded by 32 MiB or one
128-row scale block, whichever is larger. This avoids a full-size FP32 FFN intermediate,
although small mixed-bias projections can be slower. FP32 outputs reuse their output buffer
for bias addition. Autocast converts eligible operands at the public linear boundary.

PyTorch 2.14's native two-level NVFP4 GEMM keeps each call's scaling tensor independent
through the [upstream concurrency fix](https://github.com/pytorch/pytorch/commit/7add580915ff1a547f3c8bfd25afaca4207ac832).
Affine projections retain their fused GEMM epilogue across concurrent threads and CUDA
streams. Sparse-attention fusions can overlap gate and output projections on separate streams.

`piper_kernels.linear.nvfp4.reference` provides independent PyTorch-only activation preparation,
prepared projections, and ordinary/ConvRot projections. It does not invoke the optimized
preparation or GEMM operators. Its portable quantization can choose neighboring FP8 block scales
or FP4 values at rounding boundaries compared with the optimized quantizer; prepared-projection
tests separately check affine precision using identical packed operands.

Both NVFP4 wrappers support in-place adapter merges:

```python
weight.add_(dense_update, alpha=adapter_strength, rounding_seed=seed)
rotated_weight.addmm_(lora_b, lora_a, alpha=lora_strength, rounding_seed=seed)
```

On NVIDIA compute capability 10.0 or newer with Triton installed, these updates fuse
dequantization, optional ConvRot rotation, merging, and NVFP4 packing in tiles. `addmm_`
accumulates the matrix product in FP32 without allocating a dense product. Stochastic E2M1
rounding runs in registers during packing and requires no per-element temporary tensors.
Two-level weight scaling uses a read-only tile-amax pass, a small reduction, and a second pass
that recomputes and packs each tile into the existing buffers. One-level scaling needs only the
packing pass. Other devices use the portable PyTorch implementation.

Updates preserve the wrapper, packed data, block scales, global-scale buffer, activation
calibration, and packed-pair order. `add_` requires an exact-shape dense update; both operations
require inputs matching the weight's logical dtype and device and do not support autograd.
An unsigned 64-bit `rounding_seed` enables reproducible stochastic rounding without consuming
the global RNG or changing scale selection. Omit it for nearest rounding. Kernel reduction
orders and random samples can differ from the PyTorch backend. Standard signatures support
`torch.compile`; the `add_` seed extension is intended for eager updates. Repeated requantization
is lossy, so restore pristine weights before replacing or removing an adapter.

### Packed GGUF weights

Both ConvRot formats accept a two-dimensional tensor of packed GGUF bytes. Pass the GGML
quantization type explicitly, or omit it when the tensor has a `quant_type` attribute:

```python
int8_weight = ConvRotInt8Tensor.from_gguf(
    packed_gguf,
    quant_type=ggml_quant_type,
    group_size=64,
)
nvfp4_weight = ConvRotNVFP4Tensor.from_gguf(
    packed_gguf,
    quant_type=ggml_quant_type,
    group_size=64,
    compute_per_tensor_scale=True,
    is_swizzled_scales=True,
    act_quant_kwargs=activation_quantization,
)

# Streaming runtimes can refill the same device allocations.
int8_weight.copy_from_gguf_(next_packed_gguf, quant_type=ggml_quant_type)
nvfp4_weight.copy_from_gguf_(
    next_packed_gguf,
    quant_type=ggml_quant_type,
    compute_per_tensor_scale=True,
)
```

CUDA/ROCm INT8 conversion and exact-SM120 NVFP4 conversion decode GGUF values in registers and feed
the existing ConvRot quantization epilogues, so no dense weight is allocated. Computing an NVFP4
global scale requires one row-amax pass and one packing pass. Unsupported devices fail closed
instead of allocating a dense fallback: INT8 requires a compatible Triton accelerator, while
NVFP4 requires exact SM120. INT8 uses one shared fused converter across NVIDIA and AMD;
NVIDIA retains its existing fused schedule through 49,152 columns. The wide-row fallback
recomputes decoded/rotated tiles after reducing their maxima, retaining at most 1 MiB of
temporary maxima or one row's maxima, whichever is larger. Existing output buffers can be
refilled without replacing their storage.
Piper Kernels does not parse GGUF files and does not require a GGUF parser at runtime; the caller
owns file loading, tensor-name mapping, and the packed bytes plus quant-type metadata.

## Piper Attention

Piper Attention is the package's key-scaled integer-PV attention algorithm:

```python
from piper_kernels import piper_attention

output = piper_attention(query, key, value, is_causal=False)
```

It follows FlashAttention's fused online-softmax structure and SageAttention's K
smoothing plus INT8 QK quantization. Before quantization, it applies the same fixed
signed, normalized Hadamard transform across each Q head and centered K head. This
orthogonal change of basis preserves their exact dot products while smoothing outliers
for the subsequent integer quantizers. Its distinct PV path quantizes each V key row
with one signed-INT8 scale, folds those scales into nonnegative probabilities, and uses
`UINT8 x INT8 -> INT32` tensor-core products. The probability multiplier remains FP32
so every finite FP16 input scale is representable without a conversion in the hot loop.
The online-softmax state, denominator, and PV numerator also remain FP32. The numerator
stays in UINT8 probability-code units during the recurrence, and the common factor of
255 is removed once in the output epilogue.

For centered V, Piper Attention uses the exact identity

```text
softmax(QK) @ V = softmax(QK) @ (V - mean_sequence(V)) + mean_sequence(V)
```

For non-causal attention, it stores only the compact FP32
`[batch, head, feature]` mean, subtracts it while quantizing V, and restores it in the
attention epilogue. This improves signed-INT8 precision when V has a large feature bias
and preserves constant V exactly. Causal attention leaves V uncentered so per-row INT8
rounding cannot make an earlier output depend on future V rows. Both paths preserve the
original K/V sequence order.

Native mixed-sign MMA is selected on the supported NVIDIA backend through the packaged
stock-Triton extension. The integer-PV benchmark retains the exact affine identity
`u @ v = (u - 128) @ v + 128 * sum(v)` as a signed-INT8 correctness control; unsupported
production targets use the portable quantized reference instead. The public optimized
dispatch supports NVIDIA SM8x and consumer Blackwell SM12x, whose Triton lowering uses
the MMAv2 instruction rewritten by the packaged extension. Exact SM120 uses packed
four-code probability conversion for D64 and non-causal D128, while causal D128 retains
the faster stock conversion. SM89 and exact SM120 have measured schedules; other
supported targets use the generic schedule. Production plan selection depends on target, head
dimension, and causal mode, not sequence length. Hopper lowers the operation through
unsupported WGMMA and therefore uses the slow portable quantized reference. Native ROCm
mixed-sign lowering remains future work.

Piper Attention is an independently developed Sage-derived design. The per-key
quantizer, centering identity, and online-softmax lineage are not claimed as novel in
isolation; the name identifies this package's selected combination and fused recurrence.

### Grouped-query attention

`piper_attention` and `SparsePiperAttention` accept `Hq = groups * Hkv` query
heads with matching K/V head counts, including multi-query attention (`Hkv = 1`).
Head `h` reads K/V head `h // groups`; the output retains `Hq` heads. This is
inferred from tensor shapes, without a separate enable flag. Dense attention uses
`[B, H, S, D]`; sparse attention uses `[B, S, H, D]` and still requires matching
Q/K/V sequence lengths. Existing D64/D128, dtype, device, and layout restrictions apply.

K/V means, quantization, and storage are computed once per KV head. Sparse keep
ratios and routing decisions remain per query head, including mean/minmax routing,
dense suffixes, ragged tails, and valid-front padded blocks. No K/V repetition is
needed. Dense causal attention retains its existing equal-sequence-length contract.

This core support does not extend the projection-fusion graph patterns, the
separate coarse-residual API, or SageAttention2++ to GQA.

## Sparse Piper Attention

Sparse Piper is a separate non-causal SM120 operator for pre-tiled H3-style self-attention:

```python
from piper_kernels import SparsePiperAttention

attention = SparsePiperAttention(
    (0.2, 0.4, 0.6),
    routing="mean",
)
output = attention(
    query,
    key,
    value,
    sparse_key_blocks=1036,
    sparse_query_blocks=1024,  # optional leading routed-query K64 blocks
    block_lengths=block_lengths,  # optional valid-front padded K64 storage
)
```

Inputs use `[batch, sequence, heads, head_dim]` FP16, BF16, or FP32 layout with head dimensions 64 or 128.
The output preserves the input dtype; internal quantization and attention arithmetic are unchanged.
Without `block_lengths`, every row participates in attention and the sequence length may be
arbitrary. The operator pads only its internal quantized storage to K64. Supplying one contiguous
device INT32 length in `[1, 64]` per physical K64 block instead selects valid-front padded
storage. The output retains that physical
layout so the caller can apply its existing gather; padded query rows are unspecified.
`sparse_key_blocks` is a runtime count of complete routeable physical K64 prefix tiles, so any
compact partial final tile belongs to the dense suffix. Routing defaults to FP32 min/max pooling;
passing `routing="mean"` instead scores FP32 Q64/K64 mean summaries. Both policies select the
same per-head block budget over the sparse prefix, after which every query attends to every
remaining K/V row in the same softmax. By default every query block uses that policy. Supplying
`sparse_query_blocks` makes only that many leading K64 query blocks routed; later query blocks
attend every K/V block densely. This supports packed video-first layouts followed by dense
non-video queries using one runtime scalar rather than a per-block mask. Engine owns only the
semantic per-layer ratio profile.
Each opaque attention call derives its temporary physical keep counts, packed offsets, and exact
route storage from that immutable model configuration and the current prefix length. Dynamic
compiled graphs accept changed prefix lengths and their resulting route capacities without compiling
another graph or SM120 attention kernel. Routes remain call-local because both policies depend on
the current Q/K values. Compatible ConvRot INT8, NVFP4, and ConvRot NVFP4 compiler rewrites preserve
the selected policy while producing its summaries directly from fused projections.

When every head's physical budget includes every sparse key block, fine routing skips its
scores and top-k selection. Standalone attention also skips routing summaries in this case.
This includes ratios that round to a full physical budget. Coarse-attention scores still run
because they contribute to the coarse output.

Sparse Piper also exposes a routing-selectable Q/K/V-derived coarse-attention residual:

```python
from piper_kernels import sparse_piper_coarse_residual

coarse_output = sparse_piper_coarse_residual(
    query,
    key,
    value,
    coarse_gate,
    routing="mean",  # or "minmax"
    coarse_key_blocks=total_key_blocks,
    coarse_scale=coarse_scale,
    block_lengths=block_lengths,
)
output = fine_output + coarse_output
```

The selected policy derives mean- or extrema-based Q/K block scores, mean-pools V blocks, applies
dense coarse attention, expands each result over its physical K64 query block, multiplies the
caller-provided gate directly without an implicit activation, and returns the independent residual
for the caller to compose. The optional `coarse_key_blocks` prefix may include a partial compact
tail and defaults to every available block.
`block_lengths` is optional for compact storage and selects valid-front internally padded storage
when supplied. `coarse_attention_residual` remains available for learned or already-materialized
block scores.
These composable implementations define the operations and training behavior; compatible compiled
ConvRot INT8, NVFP4, and ConvRot NVFP4 graphs fuse the shared route scores, wider coarse attention,
and gated residual, including valid-front padded storage. The fused residual combines both terms
in FP32 and rounds once on output, avoiding intermediate activation rounding.
When a compatible ConvRot INT8 projection or statically scaled NVFP4/ConvRot NVFP4 projection
immediately consumes the quantized attention result, the bounded output rewrite supports
`block_lengths` and the coarse residual together with `sparse_query_blocks`.
It passes the coarse result and coarse gate into
each ranged attention launch and projects that chunk directly, so the full attention output is
not materialized. ConvRot INT8 retains this bounded path with either static or dynamic per-row
input scaling. Independent Q/coarse-gate input scales are prepared within each query window,
avoiding two full-sequence prepared inputs. When the floating-point source is a fresh,
exclusive intermediate with the same shape and dtype as the projected output, its consumed
rows become output storage. Compiler ownership checks exclude caller inputs, aliases, and
escaping values; other cases allocate a separate output. These bounds describe live tensors;
allocator-reserved memory can additionally depend on cache and library initialization history.
NVFP4 and ConvRot NVFP4 also fuse dynamically scaled output projections:
they materialize attention, compute one global activation scale (after rotation for ConvRot),
and pack/project successive chunks. When the output width does not exceed the attention width,
the final contiguous output reuses the attention allocation; narrower outputs retain that larger
backing storage. Wider outputs use a separate allocation. Dynamic scaling defaults to 32,768-row
query windows, while static NVFP4 output keeps 8,192-row windows and attention/projection overlap.
These full fusion paths support FP16, BF16, and FP32 activations, preserving
the dtype through attention, coarse-gate buffers, and the final output projection. Quantized
Q/K/V storage and FP32 accumulation are unchanged; internal operators default to BF16 when
`output_dtype` is omitted.
Q and K RMSNorm may independently use `weight=None`; the fused kernels omit the affine
weight load and multiply. Standalone fused Q/K projection operators require `head_dim=64`
or `head_dim=128` for weightless norms; compiled graphs infer it from the attention shape.
Affine norms continue to infer head width from their weight when `head_dim` is omitted.
ConvRot INT8 Q/K/V projection biases are added in FP32 inside the existing fused kernels,
including the global V mean and coarse block means used by centered attention.

The SM120 path supports both head widths, pairs two logical K64 tiles in one physical K128
recurrence, and uses one centered-V INT8 scale per logical tile. It normally reads packed UINT16
routes. Full-keep D64 calls use `skip_dense_routing` to visit all blocks without a route list;
D128 retains the list. The online numerator and pre-rounding denominator remain FP32.
The RDNA4 native path supports D64 and D128, including ConvRot INT8 projection and output
fusions. Both widths retain packed route lists and the four-wave Q64 schedule.
Unsupported devices use a slow portable implementation of the same quantized Sparse Piper
arithmetic. A separate exact-BF16 sparse reference serves as its
quality oracle; it is not the public fallback.

## SageAttention2++

The package provides an independently written, pure-Triton backend for the
canonical [SageAttention2++](https://github.com/thu-ml/SageAttention) 8+8 algorithm:

```python
from piper_kernels import sage_attention_2pp

output = sage_attention_2pp(query, key, value, is_causal=False)
```

Inputs use `[batch, heads, sequence, head_dim]` layout and may be FP16 or BF16. The
optimized backend requires NVIDIA FP8 tensor cores with FP16 accumulation (SM89 or
newer); measured schedules currently cover consumer SM89 and SM120 GPUs, while other
SM12x targets retain grouped Q/K quantization with generic scheduling. It supports head
dimensions 64 and 128, equal query/KV head counts, arbitrary positive sequence lengths,
rectangular non-causal attention, strided sequence dimensions, and `torch.compile`. It
is inference-only and does not support autograd. Its production execution plans are also
sequence-length invariant.

This is SageAttention2++, not a Piper Attention-specific algorithm: K is smoothed, the
same fixed signed, normalized Hadamard transform is applied to Q and centered K, Q/K are
quantized to INT8 with the canonical architecture-specific granularity, V and the
online-softmax probabilities are quantized to E4M3, each 64-key P x V tile accumulates
in FP16, and tile results are buffered in FP32. All optimized device code is Triton; the
package contains no CUDA extension. Unsupported devices use the slow portable quantized
reference.

Install either optimized attention backend with `piper-kernels[triton]`. The official CUDA
SageAttention package is a revision-pinned, optional benchmark dependency only; it is
not imported by production code. See [benchmarks/README.md](benchmarks/README.md) for
the reproducible provider comparison.

## Dependency direction

Applications such as Piper consume this package. Integrations such as torch-offload may
optionally recognize its tensor types, but `piper-kernels` does not depend on either
project.

## Development

```shell
uv sync --dev
uv run pytest
uv run ruff check .
uv run ruff format --check .
uv run pyright
uv build
```

GPU tests use the `gpu` pytest marker. The pre-commit test hook hides CUDA so commits run
the portable suite; run `uv run pytest` directly to exercise installed GPU backends.

## Releases

Releases follow the compatibility and release policy in [VERSIONING.md](VERSIONING.md).
Distribution artifacts are built from version tags and published to PyPI by GitHub Actions
using Trusted Publishing; maintainers do not upload releases from local environments.
