Repository navigation
3bit EXL3 エキスパートのプレフィルを融合カーネルに通す(#22 の A1) - #26
Open
GarupanOjisan wants to merge 2 commits into
Open
GarupanOjisan wants to merge 2 commits into
GarupanOjisan wants to merge 2 commits into
Conversation
融合プレフィル経路(rotate、gate/up、down の 3 カーネル)を trellis の幅でテンプレート化し、 3bit(uniform_k2 6、24 語のタイル)用の tf_exl3_prompt_gu3 と tf_exl3_prompt_down3 を足した。 これまで 3bit の層は grouped 経路の 6 カーネルに回り、入力の Hadamard 回転を 1 行につき 8 回繰り返していた。 Refs #22 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 26, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
#22 の A1 を実装した。3bit EXL3 のルーティッドエキスパートの層も、プレフィルで融合カーネル(rotate、gate/up、down の 3 つ)を通るようにした。
#22 の候補のうち、この PR で扱うのは A1 だけである。
PR #23(upstream v1.1.5 の取り込み)の上に積んでいる。v1.1.5 が
zig/src/families/glm/cuda_exl3.zigとzig/exl3_prompt_host.zigを変えているためである。#23 を先にマージしてほしい。なお、#22 の B1(GLM の KV キャッシュのページング)は v1.1.5 に入っているので、#23 で片付く。
Refs #22
変えたこと
これまで融合経路に入れたのは 4bit の層だけだった。条件が
uniform_k2 == 8だったためである。3bit の層(
uniform_k2 = 6)は grouped 経路に回っていた。grouped 経路は 6 つのカーネルを別々に起動し、入力の Hadamard 回転を 1 行につき 8 回(選ばれたエキスパートの数)繰り返す。融合カーネルで 4bit に固定されていたのは、タイルの語数(32 語)と、タイルを 16 バイトずつ 8 回で読む部分だけだった。
そこで、
zig/kernels/cuda/exl3_prompt.cuhのdecode_tile_p、tmain、tgateupを trellis の幅K2でテンプレート化した。3bit のタイルは 24 語なので、16 バイトずつ 6 回で読む。共有メモリのタイルの枠は 32 語のままにしている。
タイルの並び(
(kタイル × 列タイル数 + 列タイル) × 語数)は grouped カーネルと同じである。zig/kernels/cuda/exl3_experts.cu:tf_exl3_prompt_gu3とtf_exl3_prompt_down3(K2 = 6)を足した。4bit のtf_exl3_prompt_guとtf_exl3_prompt_downは、名前と引数を変えずに同じテンプレートから作っている。zig/src/cuda/exl3_prompt.zig:eligibleがuniform_k2の 6 も受け付ける。runは 3bit なら新しい 2 つのカーネルを起動する。dims % 256 == 0、width % 128 == 0)なので、ほかの条件は変わらない。出力について
#22 には「計算は同じなので出力は変わらない見込み」と書いたが、これは正しくなかった。
融合経路と grouped 経路では、K 方向の足し合わせの順番が違う(grouped 経路は split-K の部分和を後で足す)。
そのため 3bit の層の出力は fp32 の丸め誤差の範囲で変わり、greedy の出力トークンが今の配信と一致するとは限らない。
精度は、下のテストで float64 の参照に対する許容誤差の内側にあることを確認した。
検証
spark1(DGX Spark、sm_121)で、
nvcr.io/nvidia/pytorch:26.07-py3の中で Zig 0.17.0 を使った。zig fmt --check build.zig zig:成功zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc install native test:成功zig build test-kernel-contract:成功tf-cuda-test glm-exl3(GPU):成功。このテストを 4bit と 3bit の両方で回すようにした。3bit の融合プレフィルは 1、40、200、1030 行で、float64 の参照との差が許容誤差の 0.144 倍以内だった(4bit は 0.186 倍以内)。1024 行の区切りの前後の行が、その行だけで計算したときとビット単位で一致することも確認した。20bc843d)と変更後のexl3_experts.cuを、ビルドと同じフラグ(-O3 --expt-relaxed-constexpr)で sm_121 向けにコンパイルして比べた(nvcc 13.3)。tf_exl3_prompt_guは 128 レジスタ、tf_exl3_prompt_downは 127 レジスタで、共有メモリも含めて変更前と同じだった。新しいgu3とdown3はどちらも 128 レジスタで、共有メモリは 4bit 版と同じ大きさである。zig/exl3_prompt_host.zig:3bit のプレフィルがgu3とdown3を起動し、grouped カーネルを起動しないこと、デコードは grouped のままであることを確かめるテストを足した。uniform_k2 = 6を対象外としていた既存のテストは、対象に変えた。まだ確認していないこと
zig/src/core/contract/registry.zig):新しい 2 つのカーネルのレジスタ数と共有メモリの記録は足していない。記録には sm_80、86、89、120、121 の fatbin でtf-cuda-test kernel-contract-limitsを回す必要がある。登録していないカーネルは検査の対象にならないので、test-kernel-contractは通る。なお、nvcc 13.3 では 4bit の
guが sm_121 で 128 レジスタになり、記録(127)と合わない。これはこの PR の前から同じで、この PR では記録を直していない。🤖 Generated with Claude Code