論文の概要: Compiler-First State Space Duality and Portable $O(1)$ Autoregressive Caching for Inference

• arxiv url: http://arxiv.org/abs/2603.09555v1
• Date: Tue, 10 Mar 2026 12:03:00 GMT
• ステータス: 翻訳完了
• システム内更新日: 2026-03-11 15:25:24.267643
• Title: Compiler-First State Space Duality and Portable $O(1)$ Autoregressive Caching for Inference
• Title（参考訳）: コンパイラ-第一状態空間の双対性とポータブルな$O(1)$自動回帰キャッシング
• Authors: Cosmo Santoni,
• Abstract要約: Mamba-2のステートスペース双対性アルゴリズムは、XLAの融合とタイリングが実際に最適化したものにきれいにマッピングする。 我々は,XLA の標準プリミティブとして,完全な推論パス (プリフィル,キャッシュされた自己回帰復号化) を実装した。 この実装は、単一のJAXソースからCPU、NVIDIA GPU、Google Cloud TPUで無修正で実行される。
• 参考スコア（独自算出の注目度）: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: State-space model releases are typically coupled to fused CUDA and Triton kernels, inheriting a hard dependency on NVIDIA hardware. We show that Mamba-2's state space duality algorithm -- diagonal state structure, chunkable recurrence, and einsum-dominated compute with static control flow -- maps cleanly onto what XLA's fusion and tiling passes actually optimise, making custom kernels optional rather than required. We implement the full inference path (prefill, cached autoregressive decoding) as shaped standard primitives under XLA, without hand-written kernels, and realise the architecture's theoretical $O(1)$ state management as a compiled on-device cache requiring no host synchronisation during generation. The implementation runs unmodified on CPU, NVIDIA GPU, and Google Cloud TPU from a single JAX source. On TPU v6e across five model scales (130M--2.7B parameters), XLA-generated code reaches approximately 140 TFLOPS on single-stream prefill ($15%$ MFU) and up to $64%$ bandwidth utilisation on decode. Greedy decoding matches the PyTorch/CUDA reference token-for-token across 64 steps, with hidden-state agreement within float32 rounding tolerance. The pattern transfers to any SSM recurrence satisfying the same structural conditions, on any platform with a mature XLA backend. The implementation is publicly available at https://github.com/CosmoNaught/mamba2-jax and merged into the Bonsai JAX model library.
• Abstract（参考訳）: 状態空間モデルリリースは通常、融合したCUDAとTritonカーネルに結合され、NVIDIAハードウェアへのハード依存を継承する。 Mamba-2のステートスペース双対性アルゴリズム – 対角状態構造,チャンク可能な再帰性,および静的制御フローを備えたeinsumが支配する計算 – は,XLAの融合とタイリングパスが実際に最適化したものをクリーンにマップし,必要に応じてカスタムカーネルを任意にする。 我々は,XLAでは手書きのカーネルを使わずに,整形された標準プリミティブとして完全な推論パス(プリフィル,キャッシュされた自己回帰デコーディング)を実装し,生成時にホスト同期を必要としないコンパイルオンデバイスキャッシュとして,アーキテクチャの$O(1)$状態管理を実現する。 この実装は、単一のJAXソースからCPU、NVIDIA GPU、Google Cloud TPUで無修正で実行される。 5つのモデルスケール(130M--2.7Bパラメータ)にわたるTPU v6eでは、XLA生成コードがシングルストリームプリフィル(15%$MFU)で約140TFLOPS、デコードで最大64%の帯域幅利用を実現している。 Greedyデコーディングは、64ステップにわたってPyTorch/CUDA参照トークン-for-tokenと一致し、float32ラウンドリングトレランス内に隠された状態合意が記述される。 このパターンは、成熟したXLAバックエンドを持つ任意のプラットフォーム上で、同じ構造条件を満たす任意のSSM再発に転送される。 実装はhttps://github.com/CosmoNaught/mamba2-jaxで公開されており、Bonsai JAXモデルライブラリにマージされている。

関連論文リスト

• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling [20.849774181776414]
  FlashAttention-4はcuDNN 9.13で最大1.3$times$、BF16でB200 GPUで2.7$times$Tritonで最大2.7$times$を達成している。 従来のC++テンプレートベースのアプローチと比較して20～30$times$高速なコンパイルタイムを実現しています。
  論文  参考訳（メタデータ） (2026-03-05T18:24:49Z)
• Scaling State-Space Models on Multiple GPUs with Tensor Parallelism [0.24148976266903474]
  選択状態空間モデル(SSM)は、大規模言語モデルにとって急速に魅力的なバックボーンとなっている。 しかし、デプロイメントでは、その推論性能は単一のGPUのメモリ容量、帯域幅、レイテンシ制限によって制限されることが多い。 本稿では,3つの実践的技術的課題に対処する,選択的SSM推論のための通信効率のよいTP設計法を提案する。
  論文  参考訳（メタデータ） (2026-02-24T17:47:54Z)
• Memory-Efficient Acceleration of Block Low-Rank Foundation Models on Resource Constrained GPUs [11.45717904490388]
  トランスフォーマーベースの基盤モデルの最近の進歩は、多くのタスクのデフォルト選択となった。 その急速に成長するサイズは、単一のGPUに完全なモデルを適合させることがますます難しくなり、計算コストが禁じられる。 ブロック低ランク(BLR)圧縮技術は、重み行列のコンパクト表現を学習することでこの問題に対処する。
  論文  参考訳（メタデータ） (2025-12-24T00:41:13Z)
• dInfer: An Efficient Inference Framework for Diffusion Language Models [54.80918957287927]
  拡散に基づく大規模言語モデル (dLLM) は自己回帰(AR) LLM に代わる有望な代替品として登場した。 本稿では、dLLM推論のための効率的かつ効率的なフレームワークであるdInferについて述べる。
  論文  参考訳（メタデータ） (2025-10-09T16:19:42Z)
• dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
  拡散大言語モデル(dLLM)は並列トークン予測と低推論遅延を提供する。 既存のオープンソースモデルは、パフォーマンスを確保するためにトークン長のデコードステップをほとんど必要としています。 高速サンプリングのためにdLLMs固有の並列性を解き放つシンプルで効果的な方法であるdParallelを導入する。
  論文  参考訳（メタデータ） (2025-09-30T16:32:52Z)
• Two-stream Beats One-stream: Asymmetric Siamese Network for Efficient Visual Tracking [54.124445709376154]
  効率的な追跡を行うために,textbfAsymTrack という新しい非対称なシームズトラッカーを提案する。 このアーキテクチャに基づいて、検索機能に重要な手がかりを注入する効率的なテンプレート変調機構を考案する。 実験によると、AsymTrackは異なるプラットフォーム間で、より優れたスピード精度のトレードオフを提供する。
  論文  参考訳（メタデータ） (2025-03-01T14:44:54Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  大規模言語モデル(LLM)は様々なドメインで広く使われ、数百万の日次要求を処理する。 大規模言語モデル(LLM)は様々なドメインで広く使われ、数百万の日次要求を処理する。
  論文  参考訳（メタデータ） (2024-07-22T14:37:58Z)
• Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference [23.633481089469836]
  LLM(Large Language Models)の自動回帰デコーディングは、ハードウェアの性能に大きなオーバーヘッドをもたらす。 トレーニング可能なパラメータを0.0002$%しか必要とせず,A100-40GBのGPUをたった16時間で効率的にトレーニングできる並列プロンプトデコーディングを提案する。 我々のアプローチでは、最大2.49$times$ スピードアップを示し、最小のメモリオーバーヘッドは0.0004$%である。
  論文  参考訳（メタデータ） (2024-05-28T22:19:30Z)
• FlexLLM: Token-Level Co-Serving of LLM Inference and Finetuning with SLO Guarantees [19.58773369944074]
  タスク適応には、大規模な言語モデル(LLM)の微調整が不可欠だが、今日のサービススタックは、推論と個別のGPUクラスタ上での微調整を分離している。 本稿では,LLM推論とPEFTに基づく共通GPUのファインタニングをトークンレベルで融合した最初のシステムであるFlexLLMを紹介する。 実行時に、新しいトークンレベルの微調整機構がハイブリッドトークンスケジューラと組み合わせて、共用イテレーション毎に推論とトレーニングトークンを動的にインターリーブする。
  論文  参考訳（メタデータ） (2024-02-29T01:33:08Z)
• QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models [64.34635279436054]
  Mixture-of-Experts (MoE)アーキテクチャは、スパースルーティングによる大規模言語モデル(LLM)の高推論コストに対する一般的な解決策を提供する。 本稿では,QMoEと呼ばれる新しい圧縮実行フレームワークの形で,このメモリ問題に対する解決策を提案する。
  論文  参考訳（メタデータ） (2023-10-25T17:24:53Z)

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

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