メモリ制約GPUで最大3.76倍高速化：BLR圧縮モデルの推論を可能にするTritonカーネル

Accelerating Block Low-Rank Foundation Model Inference on MemoryConstrained GPUs

大規模基盤モデルの推論は、メモリ容量と計算コストの制約に直面しています。Block Low-Rank (BLR)圧縮は、MonarchやBLASTなどの構造化分解を用いてモデルを圧縮し、精度を保ちながら計算量とメモリ使用量を削減します。しかし、実際のGPUでは、特に複数トークン推論時にメモリバウンドとなり、理論上の高速化が達成されないことがあります。本論文では、roofline解析によりこの問題を特定し、部分融合、演算順序の変更、メモリレイアウト最適化を施したカスタムTritonカーネルを提案します。Jetson Orin NanoやA40などのメモリ制約GPU上で、PyTorchのdenseベースラインと比較して最大3.76倍の高速化と3倍のモデル圧縮を達成し、Llama-7B/1B、GPT2-S、DiT-XL/2、ViT-Bなどのモデルで有効性を示しました。