twill-swp-ws

Twill + GauS: Optimal GPU Kernel Scheduling

Implementation of two complementary papers for GPU kernel scheduling optimization:

1. Twill — "Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs" (Soi et al., NVIDIA Research, 2024). Exact ILP+SMT solver that finds provably optimal schedules.

2. GauS — "Differentiable Scheduling Optimization via Gaussian Reparameterization" (Cai et al., 2026). Scalable gradient-based solver using Gaussian distributions and Augmented Lagrangian Method.

When to Use Which

Architecture

┌─────────────────────────────────────────────────────────┐
│                    DependenceGraph                       │
│  G = (V, E), Machine Description, RRTs                  │
├────────────────────┬────────────────────────────────────┤
│    Twill Solver    │         GauS Solver                │
│                    │                                    │
│  Phase 1: ILP      │  Gaussian Reparameterization       │
│  (CBC/PuLP)        │  X_i ~ N(μ_i, σ_i²)               │
│  → Modulo Schedule │  → P_i^d = Φ((d+0.5-μ)/σ)        │
│                    │        - Φ((d-0.5-μ)/σ)            │
│  Phase 2: SMT      │                                    │
│  (Z3/QFLIA)       │  Augmented Lagrangian Method        │
│  → Joint SWP + WS │  → Adam optimizer on (μ, σ)        │
│                    │  → Dependency + Resource +          │
│  Cost Norm (§5.2)  │    Modulo + Recurrence losses      │
│  → Ratio-preserving│                                    │
│    cycle count     │  Legalization Heuristics            │
│    reduction       │  → Topological pass (regular)      │
│                    │  → Fixed-point iteration (modulo)   │
├────────────────────┴────────────────────────────────────┤
│               Code Generation                           │
│  Pseudocode · CUDA Skeleton · Pipelined Schedule        │
└─────────────────────────────────────────────────────────┘

Quick Start

pip install pulp z3-solver numpy matplotlib torch

Twill (exact solver)

from twill.kernels import flash_attention_forward_simplified
from twill.twill_solver import twill_solve

graph = flash_attention_forward_simplified()
result = twill_solve(graph, max_I=5, verbose=True)
# → I=2, schedule: S@0, P@2, O@3 (proves FA3 optimal)

GauS (differentiable solver)

from twill.kernels import flash_attention_forward_hopper
from twill.gaus_solver import gaus_solve_twill_graph

graph = flash_attention_forward_hopper()
result = gaus_solve_twill_graph(graph, target_II=4, D=20, verbose=True)
# → II=4, feasible schedule found via gradient descent

GauS on custom large graphs

from twill.gaus_solver import GauSSolver, generate_random_dag

graph = generate_random_dag(num_nodes=1000, edge_density=0.01, num_back_edges=50)
solver = GauSSolver(graph, D=500, lr=0.01)
result = solver.solve_modulo(II=10, R_cap=50.0, max_iters=2000)

Pre-built Kernel Descriptions

Module Structure

twill/
├── __init__.py              # Package exports (v0.2.0)
├── graph.py                 # DependenceGraph, Instruction, RRT, MachineDescription
├── cost_normalization.py    # §5.2: ILP-based cycle count normalization
├── modulo_scheduler.py      # Twill Phase 1: ILP modulo scheduling (CBC)
├── smt_joint.py             # Twill Phase 2: SMT joint SWP+WS (Z3)
├── twill_solver.py          # Twill Algorithm 1: Main search procedure
├── gaus_solver.py           # GauS: Differentiable scheduling (PyTorch)
├── codegen.py               # Code generation (pseudocode, CUDA skeleton)
├── visualization.py         # Schedule visualization (text + matplotlib)
└── kernels.py               # Pre-built kernel descriptions (FMHA, GEMM)

GauS: Technical Details

Gaussian Reparameterization (§3.1)

Each operator v_i is modeled as X_i ~ N(μ_i, σ_i²) with only 2|V| parameters (vs. D·|V| for categorical approaches). The probability of scheduling at step d:

P_i^d = Φ((d+0.5-μ_i)/σ_i) - Φ((d-0.5-μ_i)/σ_i)

Differentiable Loss Functions

• Dependency: Expected topological violations via CDF products
• Resource: LogSumExp smooth-max of per-step usage + ReLU violations
• Modulo Resource: Wrapped Gaussian probabilities into II-slot reservation table
• Recurrence: Expected back-edge constraint violations
• Memory: CDF-based active lifetime estimation
• Communication: Expected edge lengths (compactness)

Augmented Lagrangian Method (ALM)

L_total = L_primary + Σ_i (λ_i · V_i + ρ/2 · ||V_i||²)
λ_i ← λ_i + ρ · V_i    (dual update)

Hyperparameters (from paper): lr=0.01, ρ=1e-4, τ=1e-2, κ=1/6

Tests

python test_twill.py   # Twill: 6/6 pass (~5s)
python test_gaus.py    # GauS:  7/7 pass (~30s)

Solvers Used

Citations

@article{soi2024twill,
  title={Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs},
  author={Soi, Rupanshu and others},
  journal={arXiv preprint arXiv:2512.18134},
  year={2024}
}

@article{cai2026gaus,
  title={GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization},
  author={Cai, Yaohui and others},
  journal={arXiv preprint arXiv:2602.20427},
  year={2026}
}

Related Work

• FlashAttention-3 — Hopper FMHA schedule that Twill rediscovers
• FlashAttention-4 — Blackwell FMHA schedule that Twill rediscovers
• ThunderKittens — Warp-level kernel framework
• CUTLASS 3.x — NVIDIA GEMM templates with WS
• Tawa — Automatic WS compiler (downstream of Twill)
• Cypress — Task-based GPU programming model
• Nautilus — Auto-scheduling tensor compiler (fills Twill's tile-size gap)
• MPK — Cross-kernel software pipelining
• GS-Schedule — GauS predecessor (categorical approach)