Implementation of two complementary papers for GPU kernel scheduling optimization:
┌─────────────────────────────────────────────────────────┐
│ 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 │
└─────────────────────────────────────────────────────────┘1 from twill . kernels import flash_attention_forward_simplified
2 from twill . twill_solver import twill_solve
3
4 graph = flash_attention_forward_simplified ( )
5 result = twill_solve ( graph , max_I = 5 , verbose = True )
6 # → I=2, schedule: S@0, P@2, O@3 (proves FA3 optimal)
1 from twill . kernels import flash_attention_forward_hopper
2 from twill . gaus_solver import gaus_solve_twill_graph
3
4 graph = flash_attention_forward_hopper ( )
5 result = gaus_solve_twill_graph ( graph , target_II = 4 , D = 20 , verbose = True )
6 # → II=4, feasible schedule found via gradient descent
1 from twill . gaus_solver import GauSSolver , generate_random_dag
2
3 graph = generate_random_dag ( num_nodes = 1000 , edge_density = 0.01 , num_back_edges = 50 )
4 solver = GauSSolver ( graph , D = 500 , lr = 0.01 )
5 result = solver . solve_modulo ( II = 10 , R_cap = 50.0 , max_iters = 2000 )
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)L_total = L_primary + Σ_i (λ_i · V_i + ρ/2 · ||V_i||²)
λ_i ← λ_i + ρ · V_i (dual update)1 python test_twill.py # Twill: 6/6 pass (~5s)
2 python test_gaus.py # GauS: 7/7 pass (~30s)
1 @article{soi2024twill,
2 title={Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs},
3 author={Soi, Rupanshu and others},
4 journal={arXiv preprint arXiv:2512.18134},
5 year={2024}
6 }
7
8 @article{cai2026gaus,
9 title={GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization},
10 author={Cai, Yaohui and others},
11 journal={arXiv preprint arXiv:2602.20427},
12 year={2026}
13 }
