HaoyiZhu GitHub avatar

Piecewise-Taylor-Attention

HaoyiZhu

Piecewise-Taylor Attention (PWT) is a training-free and trainable sparse attention mechanism that adds a second-order Taylor correction to PISA's block-wise attention approximation.

Stars

25

7-day growth

No data

Forks

0

Open issues

0

License

MIT

Last updated

2026-07-29

AI repository intelligence
FR-AI / ANALYSIS

Why it is worth attention

It achieves higher PSNR than PISA in a training-free setting on a 14B video diffusion model while maintaining nearly identical speed, and it ships a full backward pass for fine-tuning at high sparsity.

Who it is for

  • Researchers in efficient attention mechanisms
  • Engineers working on diffusion transformers for video generation
  • Practitioners deploying large-scale video models with memory constraints
  • Developers interested in trainable sparse attention kernels

Use cases

  • Accelerating attention in large video diffusion transformers like Wan2.1-T2V-14B
  • Fine-tuning pre-trained models with high sparsity using a differentiable sparse attention layer
  • Replacing dense attention in training loops where variable sequence lengths are common

Strengths

  • Demonstrates 21.03 dB PSNR on Wan2.1-T2V-14B at 87.5% sparsity, outperforming PISA's 20.83 dB
  • Adds negligible overhead over PISA: only a precomputed broadcast add, no extra tensor-core work
  • Includes a full backward pass that flows gradients through all block statistics for trainable fine-tuning
  • Kernels never recompile across variable sequence lengths, making them suitable for video-DiT training loops

Considerations

  • Requires torch >= 2.7, triton >= 3.4, and a CUDA GPU, limiting platform compatibility
  • The second-order correction is gated to zero on blocks where the Taylor expansion is invalid, falling back to PISA's behavior
  • End-to-end speedup (1.54x) is slightly lower than PISA (1.57x) on the tested hardware and configuration

README quick start

Installation

Description

Piecewise-Taylor Attention

Related repositories

Similar projects matched by category, topics, and programming language.

lopopolo
Featured
lopopolo GitHub avatar

harness-engineering

Harness Engineering is a methodology for improving coding agent outputs by carefully crafting the environment around them—providing curated context, tools, and executable constraints that encode an organization’s nonfunctional requirements and cumulative lessons.

AI & Machine LearningAI Agents
2,390
slvDev
Featured
slvDev GitHub avatar

esp32-ai

A 28.9 million parameter language model runs on an $8 ESP32-S3 microcontroller entirely on-device, generating simple stories at about 9.5 tokens per second.

AI & Machine LearningLarge Language Models
1,960
littledivy
Featured
littledivy GitHub avatar

mimic

mimic captures traffic from any iOS or web app and automatically generates a Python client library that lets you call the app's API like a regular library.

AI & Machine Learning
1,482