論文の概要: Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU

• arxiv url: http://arxiv.org/abs/2604.15464v1
• Date: Thu, 16 Apr 2026 18:30:13 GMT
• ステータス: 翻訳完了
• システム内更新日: 2026-04-20 22:00:19.622613
• Title: Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU
• Title（参考訳）: タグ付きページ注意:TPUのための高性能でフレキシブルなLLM推論カーネル
• Authors: Jevin Jiang, Ying Chen, Blake A. Hechtman, Fenghui Zhang, Yarong Mu,
• Abstract要約: 我々はGoogleのTPUのための高性能で柔軟なアテンションカーネルであるRPA(Ragged Paged Attention)を提示する。 RPAはデコードで最大86%のメモリ帯域利用率(MBU)、プリフィルで73%のモデルFLOPs利用率(MFU)を達成する。
• 参考スコア（独自算出の注目度）: 3.9009842917937534
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Large Language Model (LLM) deployment is increasingly shifting to cost-efficient accelerators like Google's Tensor Processing Units (TPUs), prioritizing both performance and total cost of ownership (TCO). However, existing LLM inference kernels and serving systems remain largely GPU-centric, and there is no well-established approach for efficiently mapping LLM workloads onto TPU architectures--particularly under the dynamic and ragged execution patterns common in modern serving. In this paper, we present Ragged Paged Attention (RPA), a high-performance and flexible attention kernel for TPUs, implemented using Pallas and Mosaic. RPA addresses these challenges through three key techniques: (1) fine-grained tiling to enable efficient dynamic slicing over ragged memory, (2) a custom software pipeline that fuses KV cache updates with attention computation, and (3) a distribution-aware compilation strategy that generates specialized kernels for decode, prefill, and mixed workloads. Evaluated on Llama 3 8B on TPU7x, RPA achieves up to 86% memory bandwidth utilization (MBU) in decode and 73% model FLOPs utilization (MFU) in prefill. Integrated as the primary TPU backend in vLLM and SGLang, RPA provides a production-grade foundation for efficient TPU inference and offers practical insights into kernel design.
• Abstract（参考訳）: 大規模言語モデル(LLM)デプロイメントは、GoogleのTensor Processing Units(TPU)のようなコスト効率の高いアクセラレータにシフトし、パフォーマンスと総所有コストの両方を優先している。 しかし、既存のLLM推論カーネルとサービスシステムはほとんどGPU中心であり、特に現代のサービスで一般的な動的およびラグな実行パターンの下で、LLMワークロードをTPUアーキテクチャに効率的にマッピングするための確立されたアプローチは存在しない。 本稿では,Pallas と Mosaic を用いて実装した,TPU の高性能かつ柔軟な注目カーネルである Ragged Paged Attention (RPA) を提案する。 RPAは、(1)ラグメモリ上で効率的な動的スライシングを実現するためのきめ細かいタイリング、(2)KVキャッシュの更新と注意計算を融合するカスタムソフトウェアパイプライン、(3)デコード、プリフィル、混合ワークロードのための特別なカーネルを生成する分散対応コンパイル戦略である。 TPU7x上のLlama 3 8Bで評価されたRPAは、デコードで最大86%のメモリ帯域幅利用(MBU)、プリフィルで73%のモデルFLOPs利用(MFU)を達成する。 vLLMとSGLangの主要なTPUバックエンドとして統合されたRPAは、効率的なTPU推論のためのプロダクショングレードの基盤を提供し、カーネル設計に関する実践的な洞察を提供する。

関連論文リスト

• RollArt: Scaling Agentic RL Training via Disaggregated Infrastructure [49.88201789074532]
  エージェント強化学習(RL)は、大規模言語モデル(LLM)が自律的な意思決定と長期計画を行うことを可能にする。 分散インフラストラクチャ上でマルチタスクエージェントRLのスループットを最大化する分散システムであるRollArcを提案する。
  論文  参考訳（メタデータ） (2025-12-27T11:14:23Z)
• Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
  分析フレームワークとして'3つの税'(バルク同期、カーネル間データローカリティ、カーネルローンチオーバーヘッド)を紹介した。 我々は、分散GPU実行におけるキー非効率に対処するために、厳密なBSPモデルを超えて移動することを提案する。 BSPベースのアプローチによるエンドツーエンドのレイテンシの10-20%の高速化を観察する。
  論文  参考訳（メタデータ） (2025-11-04T01:15:44Z)
• xLLM Technical Report [57.13120905321185]
  我々は,知的かつ効率的なLarge Language Model (LLM)推論フレームワークであるxLLMを紹介する。 xLLMは、新しい分離されたサービスエンジンアーキテクチャを構築する。 xLLM-Engineは、システムとアルゴリズムの設計を最適化し、コンピュータリソースを完全に飽和させる。
  論文  参考訳（メタデータ） (2025-10-16T13:53:47Z)
• eIQ Neutron: Redefining Edge-AI Inference with Integrated NPU and Compiler Innovations [4.776283807742058]
  eIQ中性子効率NPUは商用フラッグシップMPUに統合される。 我々のソリューションは、標準AIベンチマークにおけるTOPSとメモリリソースの同等で平均1.8倍(4倍ピーク)のスピードアップを達成する。
  論文  参考訳（メタデータ） (2025-09-17T19:45:51Z)
• Serving Large Language Models on Huawei CloudMatrix384 [28.88558053380112]
  従来のAIクラスタは、計算強度、メモリ帯域幅、チップ間通信、レイテンシの制限に直面している。 本稿では,Huawei CloudMatrixを紹介する。Huawei CloudMatrixは,プロダクショングレードのCloudMatrix384スーパーノードで実現された次世代AIアーキテクチャである。 384 Ascend 910 NPUと192 Kunpeng CPUを超広帯域統一バス(UB)ネットワークを介して相互接続し、直接通信とリソースの動的プールを可能にする。
  論文  参考訳（メタデータ） (2025-06-15T03:41:34Z)
• EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference [49.94169109038806]
  本稿では,既存の並列処理方式を超越したMoE用パイプラインスケジューラであるEPS-MoEを紹介する。 その結果,既存の並列推論手法と比較して,プリフィルスループットは52.4%向上した。
  論文  参考訳（メタデータ） (2024-10-16T05:17:49Z)
• Flex-TPU: A Flexible TPU with Runtime Reconfigurable Dataflow Architecture [0.0]
  この作業はFlex-TPUと呼ばれる再構成可能なデータフローTPUの開発で構成されており、実行時に層ごとのデータフローを動的に変更することができる。 その結果,Flex-TPUの設計は従来のTPUに比べて2.75倍の大幅な性能向上を実現している。
  論文  参考訳（メタデータ） (2024-07-11T17:33:38Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  この研究は、効率的でポータブルなDeep LearningとHigh Performance Computingカーネルを開発するためのフレームワークを導入している。 1)プロセッシングプリミティブ(TPP)を用いた計算コアの表現と,2)高レベルな宣言的手法でTPPのまわりの論理ループの表現の2つのステップでカーネルの開発を分解する。 我々は、スタンドアロンカーネルと、さまざまなCPUプラットフォームにおける最先端実装よりも優れたエンドツーエンドワークロードを使用して、このアプローチの有効性を実証する。
  論文  参考訳（メタデータ） (2023-04-25T05:04:44Z)
• Heterogeneous Integration of In-Memory Analog Computing Architectures with Tensor Processing Units [0.0]
  本稿では,IMACユニットとエッジTPUを統合してモバイルCNNの性能を向上させる,新しい,異種,混合信号,混合精度アーキテクチャを提案する。 本稿では,TPU-IMACアーキテクチャ上にモデルをデプロイする際の潜在的な精度低下を軽減するために,混合精度トレーニング手法を取り入れた統合学習アルゴリズムを提案する。
  論文  参考訳（メタデータ） (2023-04-18T19:44:56Z)
• TPU-MLIR: A Compiler For TPU Using MLIR [2.6519283973116963]
  TPU-MLIRは、TPU(Processing Unit)と呼ばれるカスタムASICに、トレーニング済みニューラルネットワーク(NN)モデルをデプロイする NNモデルはTOP方言に変換され、チップの構成に応じて異なるTPUのためにTPU方言に低下する。 MLIRパスパイプラインを使用してTPU上で最適化を行い、マシンコードを生成する方法を示す。
  論文  参考訳（メタデータ） (2022-10-23T10:45:54Z)

関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。

指定された論文の情報です。
本サイトの運営者は本サイト（すべての情報・翻訳含む）の品質を保証せず、本サイト（すべての情報・翻訳含む）を使用して発生したあらゆる結果について一切の責任を負いません。