日本語ニュース
最新
AI・テクノロジー

実践的 GPU エンジニアリング:Metal シェーダー、プロファイリング、エッジデプロイ

北米北西部の都市の灯りを示す軌道からのイメージ画像であり、AI コンピューティングシステムではありません。
北米北西部の都市の灯りを示す軌道からのイメージ画像であり、AI コンピューティングシステムではありません。

実践的な GPU エンジニアリングは、クロスプラットフォームのシェーダープログラミング、ハードウェアカウンタ診断、オンデバイスのエッジデプロイにまたがります。以下では、この 3 つのワークフローすべてを深く掘り下げます。

1. カスタム Apple Metal(MSL)タイル行列乗算シェーダー

Apple Silicon(M1–M4 / A シリーズ)では、コンピュートカーネルは Metal Shading Language(MSL)で記述され、Threadgroup Memory(CUDA Shared Memory に相当する Apple の機能)を活用して、遅いグローバルメモリの読み取りを回避します。

以下は、16 × 16 の threadgroup タイルを使用した MSL のタイル行列乗算(C = A × B)シェーダーです。

#include <metal_stdlib>
using namespace metal;

constant uint TILE_SIZE = 16;

struct MatrixDims {
    uint M; // Rows of A and C
    uint N; // Cols of B and C
    uint K; // Cols of A and Rows of B
};

kernel void matmul_tiled(
    device const float* A        [[buffer(0)]],
    device const float* B        [[buffer(1)]],
    device float* C              [[buffer(2)]],
    constant MatrixDims& dims    [[buffer(3)]],
    uint2 gid                    [[thread_position_in_grid]],
    uint2 tid                    [[thread_position_in_threadgroup]]
) {
    // Fast on-chip Threadgroup (SRAM) tiles
    threadgroup float tileA[TILE_SIZE][TILE_SIZE];
    threadgroup float tileB[TILE_SIZE][TILE_SIZE];

    float acc = 0.0f;
    uint numTiles = (dims.K + TILE_SIZE - 1) / TILE_SIZE;

    for (uint t = 0; t < numTiles; ++t) {
        uint tiledCol = t * TILE_SIZE + tid.x;
        uint tiledRow = t * TILE_SIZE + tid.y;

        // Collaboratively load 16x16 tiles into Threadgroup memory
        tileA[tid.y][tid.x] = (gid.y < dims.M && tiledCol < dims.K)
                              ? A[gid.y * dims.K + tiledCol] : 0.0f;
        tileB[tid.y][tid.x] = (tiledRow < dims.K && gid.x < dims.N)
                              ? B[tiledRow * dims.N + gid.x] : 0.0f;

        // Synchronize all threads in the threadgroup before reading tiles
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Compute partial dot product entirely inside registers/SRAM
        for (uint k = 0; k < TILE_SIZE; ++k) {
            acc += tileA[tid.y][k] * tileB[k][tid.x];
        }

        // Wait before overwriting tiles in the next loop iteration
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (gid.y < dims.M && gid.x < dims.N) {
        C[gid.y * dims.N + gid.x] = acc;
    }
}
  • 重要なハードウェアマッピング:[[thread_position_in_grid]] は CUDA の blockIdx * blockDim + threadIdx に直接対応し、threadgroup_barrier(mem_flags::mem_threadgroup) は CUDA の __syncthreads() に対応します。本番環境の FP16 ワークロード向けに、Apple Silicon は simdgroup_float8x8 ハードウェア組み込み関数も公開しており、Apple の行列乗算ユニットに直接アクセスできます。

2. ncu Roofline 最適化のステップバイステップ手順

(4096, 4096)FP16 テンソルに対する行単位 LayerNorm の素朴な CUDA または Triton カーネルを書いたところ、PyTorch のネイティブ実装よりも明らかに遅くなったとします。

ステップ 1:ハードウェアカウンタのキャプチャ

NVIDIA Nsight Compute(ncu)を実行し、セットアップコードでのプロファイリングオーバーヘッドを避けるため、カスタムカーネル名のみを対象にします:

ncu --kernel-name regex:layernorm_naive --metrics \
sm__throughput.avg.pct_of_peak_sustained_elapsed,\
dram__throughput.avg.pct_of_peak_sustained_elapsed,\
l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\
smsp__warp_issue_stalled_long_scoreboard_per_warp_active.pct \
python run_layernorm.py

ステップ 2:Roofline 上でのボトルネックの診断

出力メトリクスを確認します:

  • 計算(SM)スループット:ピークの 11.4%。
  • メモリ(DRAM/HBM)スループット:ピークの 84.2%。
  • Warp 停滞理由(long_scoreboard):68.5%(warp がグローバル HBM メモリ読み取りの返却待ちで停滞)。
  • 診断:このカーネルは Roofline チャートの傾斜したメモリ帯域幅ルーフの真下に位置しています。さらに、l1tex__t_sectors を確認すると、メモリ読み取りがコアレースされていないことがわかります。同じ warp 内のスレッドが連続した行アドレスを読むのではなく列方向にストライドしており、フェッチされた 128 バイトの HBM キャッシュラインの 75% を無駄にしています。

ステップ 3:的を絞ったハードウェア修正の適用

1. メモリコアレッシングの修正:隣接するスレッド(threadIdx.x)が連続した内側の次元(n_cols)に沿って隣接する要素を読み取るようにし、32 の FP16 読み取りを 1 回の 64 バイトのコアレーストランザクションにマージします。

2. ベクトル化されたロード(float4 / tl.load ブロック):16 ビットではなく 1 命令あたり 128 ビット(8 つの FP16 値)をロードし、メモリ命令の総数を 8 分の 1 に削減します。

3. Warp シャッフルリダクション:遅い Shared Memory リダクションをレジスタ間の warp シャッフル(CUDA の __shfl_xor_sync または Triton の tl.sum)に置き換え、同期バリアなしで行の平均と分散を計算します。

ステップ 4:Roofline での検証

ncu を再実行して確認します:

  • グローバルロードセクタ:4 分の 1 に低下(キャッシュライン利用率 100%)。
  • ロングスコアボードストール:68.5% から < 15% に低下。
  • 実時間レイテンシ:3.8 倍改善し、実効帯域幅が H100 の物理 HBM 上限の 85% を超えます。

3. エッジデプロイパイプライン:モバイル/ノート PC での LLM の量子化と実行

8B パラメータの LLM をノート PC やスマートフォン(8 GB–16 GB の共有 RAM)にデプロイするには、重みを 4 ビット整数に圧縮し、llama.cpp や Apple MLX などのゼロオーバーヘッド C++ ランタイムで実行する必要があります。

[HF Safetensors (FP16, 16 GB)] 
       │  1. Convert & Apply Calibration Data (Imatrix)
       ▼
[GGUF Quantized Binary (Q4_K_M, 4.9 GB)]
       │  2. Memory-Map (mmap) into Unified RAM
       ▼
[Metal / Vulkan Compute Shaders (On-the-Fly Dequant + Matrix Mul)]

ステップ 1:変換とキャリブレーション(重要度行列)

単純な 4 ビット丸めは、推論やコーディングの精度を低下させます。llama.cpp を使用して、まず Hugging Face の重みを非圧縮の GGUF ファイルに変換し、次に小さなキャリブレーションテキストサンプルで重要度行列(imatrix)を計算して、どの重みチャネルが活性化に最も大きな影響を与えるかを測定します:

python convert_hf_to_gguf.py ./Llama-3.1-8B --outfile llama-8b-f16.gguf
./llama-imatrix -m llama-8b-f16.gguf -f calibration.txt -o imatrix.dat

ステップ 2:混合ビット K 量子化(Q4_K_M)

Q4_K_M(4 ビット K-quant Medium)でモデルを量子化します。これは、敏感な attention の value/output レイヤーを 6 ビット(Q6_K)に保ちながら、密なフィードフォワードレイヤーを 256 重みのスーパーブロックで 4 ビット(Q4_K)に圧縮します:

./llama-quantize --imatrix imatrix.dat llama-8b-f16.gguf llama-8b-q4_k_m.gguf Q4_K_M

これにより、モデルは 16.0 GB から約 4.9 GB に縮小し、パープレキシティの劣化は 1% 未満です。

ステップ 3:デバイス上でのゼロコピー統合メモリ実行

iPhone、Android デバイス、または MacBook で実行する場合:

  • mmap による即時ロード:ランタイムは 4.9 GB の .gguf ファイルをフラッシュストレージから統合メモリに直接メモリマップ(mmap)し、RAM の重複を排除して 1 秒未満でモデルをロードします。
  • 融合逆量子化シェーダー:トークン生成中に、カスタム Metal(ggml-metal.metal)または Vulkan(mul_mat_q4_K.comp)コンピュートシェーダーが、パックされた 4 ビット重みブロックを GPU レジスタに直接ロードし、ビットシフトでアンパックして、FP16 の活性化とその場で乗算します。これにより、最新のノート PC GPU で 40–70+ トークン/秒、フラッグシップ携帯電話で 15–25 トークン/秒を達成します。

その他のニュース: AI・テクノロジー