Group 1 - Hardware-Aware Transformer Optimisation: Integrating Programmable Attention, Triton Kernel Fusion, and Multi-Objective NAS

[aahaidar01]

Authors:

Ali Haidar, Dorijan Donaj Magasic, Yash Agarwal, Mahmoud El Etreby

Summary

• FlexAttention integration pass — module-level transform replacing SDPA with torch.nn.attention.flex_attention, supporting causal, sliding-window, ALiBi, and document masking patterns with block-sparse acceleration. Up to 1.72x training speedup for sliding-window attention at seq=4096.

• Fused Add+RMSNorm Triton kernel — custom forward/backward Triton kernel fusing the residual-add → RMSNorm pattern in transformer decoder layers. 2.98x faster than unfused PyTorch, 1.42x faster than Liger-Kernel, with 60% peak memory reduction per fusion site. Includes both FX graph-level and module-level MASE passes.

• Automated multi-objective search pipeline — fills MASE's LatencyRunner stub with GPU timing, adds a ModuleSearchSpaceQuantizationFusion search space covering bit-width × fusion strategy, and wires everything into Optuna NSGA-II search. Produces Pareto frontiers over accuracy/perplexity, latency, and average bitwidth across BERT, TinyLlama, and Mistral.

Key results

Optimization	Best speedup	Condition
FlexAttention SWA (inference)	1.46x	seq=4096, Llama
FlexAttention SWA (training throughput)	1.73x	seq=4096, 25K vs 14.5K tok/s
FlexAttention document masking	2.25x	seq=8192 vs SDPA mask
Fused RMSNorm kernel vs PyTorch	2.98x	L40S, BF16
Fused RMSNorm kernel vs Liger-Kernel	1.42x	L40S, BF16
Fused RMSNorm memory reduction	60%	per fusion site (forward)
Mistral sliding-window (search pipeline)	7%	consistent from seq≥512

Files changed

New passes

• src/chop/passes/module/transforms/attention/flex_attention_transform.py — FlexAttention pass
• src/chop/passes/module/transforms/attention/score_mods.py — score/mask modification library
• src/chop/passes/graph/transforms/fused_rmsnorm/ — Triton kernel + FX graph pass
• src/chop/passes/module/transforms/fused_ops/rmsnorm_residual_fusion.py — module-level fusion pass

Search infrastructure

• src/chop/actions/search/strategies/runners/hardware/latency.py — GPU latency runner
• src/chop/actions/search/search_space/quantization/module_fusion.py — fusion search space
• src/chop/pipelines/optimization.py — pass-chain wrapper
• configs/search/quantization_fusion_{bert,llama,mistral}.toml — search configs

Experiments & benchmarks

• experiments/flex_attention/ — 12 experiments with JSON results and figures
• scripts/ — search runners, benchmarks, kernel profiling
• test/ — FlexAttention tests (40 tests), LatencyRunner tests

Documentations

• Flex Attention: https://github.com/aahaidar01/mase/blob/main/FLEXATTENTION_REPORT.md
• FUSED_RMS_NORM_RESIDUL_Triton: https://github.com/aahaidar01/mase/blob/main/FUSED_KERNEL_PROJECT_SPEC.md
• Automated Search Pipeline: https://github.com/aahaidar01/mase/blob/main/Automated_Search_Pipeline.md

Models tested

• BERT-base (SST-2 classification) — FlexAttention + quantization search
• TinyLlama-1.1B (WikiText-2 LM) — full pipeline with causal FlexAttention
• Mistral-7B (WikiText-2 LM) — sliding-window FlexAttention in float16

Hardware

NVIDIA L40S (48GB), PyTorch 2.6.0, Triton 3.3.1