高速バックエンド

TileLang の可搬性

laya/tl_kernels.py の 5 つのカーネルは、compile_cpu を通じて TileLang の CPU（c）ターゲットで実行できます。これは明示的なスカラー fp32 特殊化であり、Agent.accelerate の CPU 高速化バックエンドではありません。TileLang とローカルの C++ コンパイラが必要です。ランタイムサービスの依存関係は導入しません。

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs) は、5 つの GPU ファクトリと同じ形状・演算オプションを受け取ります。cpu=True、dtype="float32"、ターゲット c を選択し、ベクトル化を無効にします。すべての浮動小数点の入力と出力は連続した CPU float32 テンソルでなければなりません。アテンション長は int32 のままです。呼び出す前に、16 ビットの活性化を .cpu().float().contiguous() で変換します。LayerNorm は依然として残差テンソルをその場で更新し、RoPE は依然として Q/K 列をその場で更新します。コンパイル済みカーネルを保持して、その動的次元を再利用します。

CPU 特殊化は、fragment と shared の割り当てをローカルバッファに置き換え、直列の T.grid ループを使い、アテンションの GPU swizzle 注釈を省きます。GEMM は TileLang のスカラー CPU 実装を使い、リダクションはローカルバッファを使います。元のタイル化アルゴリズム、パディング述語、アテンションマスク、オンラインソフトマックスは GPU 実装と共有されたままです。CPU の中間結果はアテンション確率を含めて fp32 のままです。GPU の中間結果は元の dtype を保ちます。CPU の高速化やモデル全体の CPU バックエンドは主張しません。

TileLang 0.1.14 での再調査

Linux、Python 3.12.13、torch 2.11.0+cu130、TileLang 0.1.14、RTX 4070 Ti SUPER で観測しました。以前の可搬性の懸念は、変更していない GPU 特殊化のコンパイルに関するもので、CPU への lower の可能性に関するものではありません。ベースコミット fa9a2a7 では、このドキュメントページは checkout に存在しませんでした。

テストは factory.get_tir(...) を target="c" と GPU の FAST パス設定でコンパイルします。以下は正確な診断テキストです（ソース位置とスタックトレースは省略）。すべてのカーネルで bf16 と fp16 の両方を調査しました。

最初の 4 つの fragment/リダクション失敗は tvm.error.InternalError です。bf16 のコード生成失敗も InternalError です。FP16 の RoPE は RuntimeError: Compilation Failed! を投げ、続いてコンパイラの起動とソースが出力されます。上記の診断は stderr に出力されます。生成されたベクトルキャストも、float ベクトルを half ベクトルに戻す変換に失敗します。

dtype のサポートを fragment のサポートから切り離すため、テストは各カーネルを cpu=True（ローカルバッファと直列ループ）でもう一度コンパイルし、bf16 を保持してベクトル化を無効にします。すると 5 つすべてが正確に次で失敗します：

Cannot convert type bfloat16 to C type

したがって、fragment は GEMM・GEGLU・アテンションの最初の障壁であり、fragment リダクションは LayerNorm の最初の障壁であり、bf16 は独立して 5 つすべての障壁です。RoPE には fragment の割り当てもリダクションもありません。GEMM と GEGLU には明示的な T.reduce_* 演算がありません。アテンションのリダクションは、当初はそのレイアウト失敗によって隠されています。fp32 のみの独立した fragment 調査は、GEMM も bf16 も使わずに、reduce_sum と reduce_max の両方でフィルエラーとリダクションエラーを再現します。scope だけを置き換えても、GEMM バッファでこの意味エラーが発生しました（他に影響を受けたバッファは Ci、x、s）：

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

直列 fp32 特殊化はこれらの障壁を取り除きます。診断アサートはバージョン 0.1.14 に固定され、別のバージョンではスキップされます。そこではエラーを再調査すべきです。数値 CPU テストは他のバージョンでも実行され続けます。

数値チェック

python -m pytest tests/test_fast_cpu.py -q -s を実行します。上記の環境では 58 件成功です。CPU のみのチェックは CUDA を必要としません。CUDA がないとスキップされるのは GPU 比較だけです。カバレッジには、不均一な GEMM の M/N/K タイル、すべての GEMM エピローグ、GEGLU、すべての残差/バイアス組み合わせ、大きな残差値、RoPE の位置ラップアラウンドと未変更の V 列、静的/動的アテンション形状、スライディングウィンドウ、部分タイル、不均一な長さ、空の系列、有限のパディング出力が含まれます。

シードは 1234 です。GPU 比較は同一の入力値を使い、bf16 または fp16 に丸めてから CPU 実行のために fp32 に昇格します。これらは有界なフィクスチャに対する絶対許容誤差であり、任意の大きさやモデル深さに対する保証ではありません。アテンション比較は tests/test_fast.py と同様に、有効な query 行を使います。

CPU/参照の許容誤差は 2e-5（RoPE は 2e-6）です。残差ストリームの更新は、残差なしの場合も含めて完全に等しいです。

GPU 保持の証拠

cpu キーワードの既定値は false です。既存の bf16/fp16 既定値、GPU 割り当て scope、並列ループ、swizzle、FAST オプションは変更されていません。既定の fp32 GPU 呼び出しは依然として ValueError を投げます。

前後比較の実行では、それぞれ git show fa9a2a7:laya/tl_kernels.py で取り出した元のモジュールと、変更後のモジュールを使いました。ベースライン実行では、元のモジュールを importlib 経由で laya.tl_kernels として読み込みました。他の Laya コードと Python 環境は同じままです。

• python -m pytest tests/test_fast.py -q。full-forward テストは以下に示すキャッシュ済みの英語チェックポイントを読み込むよう設定します：変更前 13 件成功、変更後 13 件成功。当初チェックポイントなしでは 11 成功 / 2 スキップでした。どちらの完全実行も、チェックポイント既存の温度クランプ警告（choice:11+=0.10058280825614929 -> 0.5）を出力しました。
• 既存カーネルテストのシード付きキャプチャ：残差更新を含めて 24 個の出力テンソルがビット単位で同一、前後の最大差 0。
• 上記 5 つの調査形状について、両 dtype で生成した CUDA ソース：10/10 がバイト単位で同一。これには元の fast スイートにない RoPE も含まれます。
• python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... と同等の fp16 コマンド：dtype あたり 60 状態・288 質問。前後のすべての JSON レコード（fp32、stock、fast の確率）は完全に等しく、前後の最大確率差は 0。

MODEL はキャッシュ済みの convaiinnovations/laya 英語スナップショット 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851 でした。実行には HF_HUB_OFFLINE=1 を使いました。多言語や型付き意思決定のチェックポイント比較はここでは主張しません。

リポジトリのチェック

すべてのコマンドは /home/ckl/projects/S/laya/.venv/bin/python を使いました。ruff と zensical はその仮想環境の bin ディレクトリから取得しました。

新しいスイートは、パッケージングテストの既存の除外項目に、オプションの TileLang extra と C++ コンパイラを必要とすることが登録されています。API 契約テストは、ベース CI 環境に TileLang をインポートせずに、追加されたキーワード専用 CPU セレクタ、変更していない GPU dtype 既定値、compile_cpu シグネチャを固定します。

AOTInductor

DecisionModel はエクスポートし、.pt2 パッケージにコンパイルし、torch._inductor.aoti_load_package で読み込めます。パッケージは意思決定ロジットと行動ロジットの両方を返します。トークン化、パディング、温度較正、回答フォーマットは呼び出し側の仕事のままで、Agent バックエンドを追加したり、既定の実行パスを変えたりしません。

クローズ済みの PR #472 は以前の dtype 障壁を記録しています。行動ヘッドのプールされた state と信頼度特徴は、モデルの重みが明示的に bf16 であっても fp32 で計算されます。autocast がない場合、最初の行動 linear は fp32 入力と bf16 重みを受け取っていました。フォワードは現在、連結した入力をヘッドの重み dtype に、autocast の外でのみキャストします。Softmax、エントロピー、返される意思決定ロジットは fp32 の計算を保ちます。既存の fp32 パラメータの eager 推論は fp16/bf16 AMP を含めて数値を保ちます。重み/autocast の混合 dtype も autocast の元の変換を保ちます。

bf16 に変換した別の評価用コピーを、autocast の外でエクスポートします。token 軸を 16 * Dim("tokens16", ...) として宣言し、アテンションの整列ガードを満たします。行数とマーカー数は独立に動的にできます。autocast の下でキャプチャしたエクスポートを、そのコンテキストの外でパッケージできると仮定しないでください。このレシピでは明示的な重み dtype を使います。

オフラインで再現

アサートベースのチェックは、既定でネットワークアクセスなしの、ランダムに初期化された小さな ModernBERT を使います。実モデルでの測定にはローカルのチェックポイントディレクトリを渡します。CUDA、AOTInductor API、または C++ コンパイラがない場合は明示的な SKIP になります。サポートされたインストールでのコンパイルとパリティの失敗は、チェックを失敗させます。

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

出力ディレクトリには decision.pt2、results.json、inputs.pt、eager.pt が入ります。JSON には完全な nvidia-smi スナップショット、エクスポート/パッケージ/ロード時間、パッケージのバイト数、レイテンシ、両ヘッドの最大絶対ロジット/確率デルタが含まれます。バイナリアーティファクトとコンパイラキャッシュは意図的にリポジトリの外に置かれます。

RTX 4070 Ti SUPER での測定

2026-10-03、Linux x86-64、Python 3.12.13、torch 2.11.0+cu130、transformers 5.17.0、NVIDIA ドライバ 615.71.09、16,376 MiB VRAM。チェックポイント：convaiinnovations/laya、リビジョン 1c5edc17a7acd8701df6fc341c0d179f1c62c982 の multilingual サブディレクトリ。ベースラインは main fa9a2a7 です。完全なマシンスナップショットと丸めていない測定は aoti_multilingual_rtx4070.json にあります。

再現された失敗は、エクスポート成功後のパッケージングで発生します：mat1 and mat2 must have the same dtype, but got Float and BFloat16。各実行は個別の初期状態が空の Inductor キャッシュを使いました。AMP コンパイルはパッケージングの前に実行されたので、パッケージの時間は新規インタプリタのコールドスタート測定ではありません。新規プロセスのチェックは、パッケージと保存済み入力のみを読み込み（チェックポイントなし）、torch.compile とパッケージングのエントリポイントをブロックしました。両ヘッドの回答を再現しました。PyTorch はそれでも小さな CPU AVX 能力プローブを初期状態が空のキャッシュにコンパイルしました。モデルグラフも CUDA カーネルもコンパイルされていません。したがって「ロード時にコンパイラ活動はない」という一般化した主張は、このランタイムでは不正確です。

同じ保存済み入力には、3 バッチにまたがる 7 つの意思決定行が含まれ、実際にトークン化された請求/返金テキスト、choice/score/noul の質問、パディング、不均一な有効マーカー数を持ちます。各レイテンシは、10 回のウォームアップ後に 30 フォワードを 5 グループ行った中央値で、同期したウォールクロック計測を使います。torch.compile(dynamic=True) は既定モードを使います。これらのレイテンシ数値には、明示的な CUDA グラフ、トークン化、エクスポート、コンパイルは含まれません。

ここでの AMP は fp32 パラメータと bf16 autocast を意味します。明示的 bf16 は、重みと残差ストリームが bf16 で autocast なしであることを意味します。パッケージに最も近い実行比較にはこれらの列を使ってください。明示的 bf16 eager/compiled と AOTI は、この修正以前は利用できませんでした。

両方の AOTI 比較で、意思決定と行動の argmax はともに 7/7 行で一致します。確率は生の softmax 出力で、較正されていません。行動分布はこれらの入力で飽和しているので、その確率デルタがゼロでも、AMP に対する基礎のロジットが同一であることを意味しません。明示的 bf16 compiled フォワードも明示的 bf16 eager と異なります（意思決定ロジットの最大デルタ 0.5、確率デルタ 0.00659859）。低精度コンパイルはビット厳密ではありません。

GPU はデスクトップアプリケーションや他の Python ジョブと共有されていました。CPU コンパイルも共有され、クロックは固定されていませんでした。これらは観測された時間であり、分離された高速化の主張ではありません。実行を挟む nvidia-smi スナップショット：

残る制約

• アーティファクトはこのチェックポイント、精度、PyTorch/ランタイムスタック、GPU ターゲットに固有です。他のハードウェアや PyTorch バージョンへの可搬性は検証していません。
• このチェックは、行 1–8、16 の倍数のトークン 32–512、マーカースロット 2–8 をエクスポートします。これらの範囲の各点ではなく、エクスポート例とは異なる形状を含む 3 つの形状を試します。1 つの有効なオプションはパディング済みマーカースロットを使えます。真の 1 スロットテンソルは別の単一オプション分岐を通り、別のエクスポートが必要です。任意のトークン長と本番のバケティング/ディスパッチ層は、この変更の範囲外です。
• 7 行はパッケージングの退行を検証するもので、広範なチェックポイント精度や較正済み信頼度のパリティではありません。エクスポートコピーを bf16 にキャストすることは、autocast 下で fp32 重みを保つこととは異なります。既存の eager パス自体は変更されていません。
• スクリプトはアーティファクトを作成するために、CUDA 対応の PyTorch ビルドとローカルのコンパイラツールチェーンを必要とします。.pt2 はモデルと CUDA カーネルを埋め込みます。ホストされたサービスは関与しません。