mirror of
https://github.com/gufo-org/gufo.git
synced 2026-11-07 09:35:59 -10:00
perf(embedding): run EmbeddingGemma 2 attention on matrix cores (#770)
This commit is contained in:
Shawn Wilsher authored and GitHub committed
1 parent
68bfe92101
commit
194db882a8
10 files changed
+708
-125
No files matched your search
@@ -25,7 +25,7 @@ All model documentation lives under [docs/models](docs/models/README.md):
|
||||
| [Qwen3.8 Flash-Next](docs/models/qwen3.8-flash-next/README.md) | Q4, images, AR, MTP | Unsloth [Q4_K_XL](https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF/tree/38bb39ee97821de2c9009abb7e93950eec396e66/UD-Q4_K_XL) · [MTP Q8_0](https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF/blob/38bb39ee97821de2c9009abb7e93950eec396e66/MTP/mtp-Qwen3.8-Flash-Next-shared-Q8_0.gguf) | **1,801.01 tok/s pp** (AR); up to **116.93 tok/s tg** single user and **245.81 aggregated tok/s** on 8 concurrent requests with MTP · [Benchmarks](docs/models/qwen3.8-flash-next/BENCHMARKS.md) | [Quality](docs/models/qwen3.8-flash-next/QUALITY.md) |
|
||||
| [DeepSeek V4 Flash](docs/models/deepseek-v4-flash/README.md) | AR, DSpark | [antirez Flash 0731 IQ2XXS](https://huggingface.co/antirez/deepseek-v4-gguf/blob/1cd7b564460821938add0475a60b942c409295e0/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf) · [DSpark](https://huggingface.co/antirez/deepseek-v4-gguf/blob/e7f04037032990db0346398d249baf9fb9df1ccc/DeepSeek-V4-Flash-DSpark-support-0731.gguf) | **484.62 tok/s pp**; up to **26.62 tok/s tg** single user and **54.74 aggregated tok/s** on 8 concurrent requests with DSpark · [Benchmarks](docs/models/deepseek-v4-flash/BENCHMARKS.md) | [Quality](docs/models/deepseek-v4-flash/QUALITY.md) |
|
||||
| [Qwen3-ASR 1.7B](docs/models/qwen3-asr/README.md) | Speech recognition | [BF16](https://huggingface.co/Qwen/Qwen3-ASR-1.7B/tree/7278e1e70fe206f11671096ffdd38061171dd6e5) | **15.27× realtime** · [Benchmarks](docs/models/qwen3-asr/BENCHMARKS.md) | [Quality](docs/models/qwen3-asr/QUALITY.md) |
|
||||
| [EmbeddingGemma 2](docs/models/embeddinggemma-2/README.md) | Text and image embeddings | [BF16](https://huggingface.co/google/embeddinggemma-2/tree/914f7f89142e33e77833254d9c9b90c3cef7303b) | Up to **23.9k tokens/s** over HTTP; **20.5 ms** for a 404-token request · [Benchmarks](docs/models/embeddinggemma-2/BENCHMARKS.md) | [Quality](docs/models/embeddinggemma-2/QUALITY.md) |
|
||||
| [EmbeddingGemma 2](docs/models/embeddinggemma-2/README.md) | Text and image embeddings | [BF16](https://huggingface.co/google/embeddinggemma-2/tree/914f7f89142e33e77833254d9c9b90c3cef7303b) | Up to **40.2k tokens/s** over HTTP; **12.9 ms** for a 404-token request · [Benchmarks](docs/models/embeddinggemma-2/BENCHMARKS.md) | [Quality](docs/models/embeddinggemma-2/QUALITY.md) |
|
||||
| [Qwen3-TTS 1.7B](docs/models/qwen3-tts/README.md) | Speech synthesis and voice cloning | BF16 [CustomVoice](https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) / [VoiceDesign](https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign) / [Base](https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base) | Up to **2.54× realtime**; **201 ms** to first audio (CustomVoice) · [Benchmarks](docs/models/qwen3-tts/BENCHMARKS.md) | [Quality](docs/models/qwen3-tts/QUALITY.md) |
|
||||
| [Qwen-Image-2.1](docs/models/qwen-image-2.1/README.md) | BF16 image generation and editing | [Complete pipeline](https://huggingface.co/Qwen/Qwen-Image-2.1/tree/b3179ad355be050328e483a9dfdd9e60cd62adfa) | In progress · [Benchmarks](docs/models/qwen-image-2.1/BENCHMARKS.md) | [Quality](docs/models/qwen-image-2.1/QUALITY.md) |
|
||||
| [MiniMax H3](docs/models/minimax-h3/README.md) | BF16 text to video/audio | [FL2VA pipeline](https://huggingface.co/MiniMaxAI/MiniMax-H3/tree/42ed227ee7df40d41602854ae760620d6eb651fe/FL2VA) | In progress · [Benchmarks](docs/models/minimax-h3/BENCHMARKS.md) | [Quality](docs/models/minimax-h3/QUALITY.md) |
|
||||
|
||||
@@ -4,6 +4,8 @@ Measured on 2026-10-23 on Linux with a Radeon 8060S GPU of architecture gfx1151,
|
||||
|
||||
This build's first pass stopped at the 16 × 1 × 16 row with one HTTP 503 that the server never logged. That is a race in the HTTP server's connection limit and is unrelated to this change, so the pass left one sample for the 4 × 8 × 128 and 16 × 8 × 128 rows. Those two rows come from a second pod on the same node after a reboot, with three passes per build. The [record](artifacts/gfx1151-2026-10-23-gemm-ab.json) holds both builds' rows, and the [earlier #700 record](artifacts/gfx1151-2026-10-23-nix.json) holds #700's own measurement.
|
||||
|
||||
The numbers for the matrix-core attention kernel are in the [Matrix-core attention](#matrix-core-attention) section. The table below is main with the hipBLASLt GEMMs and the vector-unit attention kernel.
|
||||
|
||||
Model load takes **2.2 s**, against 0.53 s for the baseline build. The baseline's time is mostly the 540 MB upload, and the other 1.7 s is the warm-up that times each hipBLASLt algorithm for every GEMM shape and row bucket.
|
||||
|
||||
## HTTP
|
||||
@@ -35,6 +37,27 @@ Against #700, this build is 1.5 to 2.3 times faster at the median on every row o
|
||||
| `ple_512` | 0.12 ms, 4.49 TFLOPS | 0.03 ms, 20.98 TFLOPS |
|
||||
| `final` | 0.17 ms, 4.72 TFLOPS | 0.05 ms, 17.41 TFLOPS |
|
||||
|
||||
## Matrix-core attention
|
||||
|
||||
Measured on 2026-10-24 on the node gpuworker01 with a Radeon 8060S GPU of architecture gfx1151, using BF16 safetensors and the production `gufo` from one Nix container image that holds both builds. Main is commit 479162ed, which has the hipBLASLt GEMMs from #751 and image inputs from #759. This PR is commit ec21a41f, which differs from the PR head only in documentation. The pod was the only GPU pod on the node. The builds ran in the order main, this PR, this PR, main, with 60 requests per row in each pass, and the table shows the mean of the two passes' medians. Speedup is main's median divided by this PR's median. The text rows use `http_bench.py` with the labels of the HTTP table above. The image rows send one 320 × 180 PNG followed by a newline in the typed `content` format, which is 269 tokens, at concurrency 1 and 4 after one warm-up request.
|
||||
|
||||
| Concurrency × inputs × words | Tokens/request | Main median | This PR median | Speedup | This PR tokens/s |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: |
|
||||
| 1 × 1 × 16 | 40 | 6.74 ms | 6.64 ms | 1.02 | 5,914 |
|
||||
| 1 × 1 × 128 | 404 | 19.95 ms | 12.94 ms | 1.54 | 30,391 |
|
||||
| 1 × 1 × 512 | 1,940 | 144.56 ms | 61.50 ms | 2.35 | 31,477 |
|
||||
| 1 × 1 × 2100 | 8,072 | 1,650.55 ms | 341.70 ms | 4.83 | 23,589 |
|
||||
| 1 × 8 × 128 | 3,495 | 148.01 ms | 94.27 ms | 1.57 | 36,954 |
|
||||
| 1 × 32 × 128 | 15,670 | 701.38 ms | 397.00 ms | 1.77 | 39,366 |
|
||||
| 4 × 1 × 128 | 404 | 77.88 ms | 48.45 ms | 1.61 | 32,919 |
|
||||
| 16 × 1 × 128 | 404 | 308.69 ms | 192.42 ms | 1.60 | 33,264 |
|
||||
| 4 × 8 × 128 | 3,495 | 574.09 ms | 346.99 ms | 1.65 | 40,215 |
|
||||
| 16 × 8 × 128 | 3,495 | 2,297.35 ms | 1,388.06 ms | 1.66 | 40,139 |
|
||||
| 1 × 1 image | 269 | 585.69 ms | 581.39 ms | 1.01 | 463 |
|
||||
| 4 × 1 image | 269 | 2,254.74 ms | 2,234.28 ms | 1.01 | 481 |
|
||||
|
||||
Image requests change by about 1 percent because the vision tower, which this PR does not change, takes most of their time. Server load took 1.98 s for this PR and 2.20 s for main. The [record](artifacts/gfx1151-2026-10-24-attention-ab.json) holds each pass's median, the p95 values, and the quality and test results.
|
||||
|
||||
## Reference comparison
|
||||
|
||||
The llama.cpp comparison is TODO. It requires a llama.cpp build with EmbeddingGemma 2 support and the matching GGUF. Measure both over HTTP with the same texts and report the ratio per row.
|
||||
|
||||
@@ -2,12 +2,13 @@
|
||||
|
||||
| Experiment | Decision / evidence |
|
||||
| --- | --- |
|
||||
| Float32 residual stream, BF16 matrix inputs | Retained. Its worst probe has cosine 0.999981 to the float32 reference, while Hugging Face BF16 has 0.999867. |
|
||||
| Float32 residual stream, BF16 matrix inputs | Retained. Its worst probe has cosine 0.999985 to the float32 reference, while Hugging Face BF16 has 0.999867. |
|
||||
| Fused QKV and gate/up projections | Retained. Each pair of projections that share an activation runs as one GEMM. |
|
||||
| BF16 safetensors only | Retained. The text backbone is 540 MB and 268 MB of that is the token table, which is a row lookup, so quantizing saves little memory or time. GGUF quants wait for published quality measurements in [discussion 478](https://github.com/gufo-org/gufo/discussions/478). |
|
||||
| Lazy vision tower | Retained. Text startup uploads only the text backbone. The first image request loads the 16-layer vision tower and projector; later requests reuse them. Audio tensors remain unindexed and are never uploaded. |
|
||||
| hipBLASLt instead of hipBLAS | Retained. hipBLAS has no matrix-core kernel for this BF16-in, float32-out transposed GEMM and runs at 4.5 to 5 TFLOPS in `gemm_bench_hip_test`, while hipBLASLt reaches 17 to 27 TFLOPS at 1,024 rows. No fixed heuristic rank wins across shapes, because the best rank moves between 4 and 8, so Gufo times each shape and row bucket once at load. Against #700 over HTTP, a 404-token request drops from 41.4 ms to 20.5 ms, which is 2.0 times faster. Concurrent single-input requests of 404 tokens are 2.0 to 2.1 times faster, and batches of 8 and 32 inputs are 2.2 to 2.3 times faster. A request of about 8,000 tokens is 1.3 times faster, and load time grows from 0.53 s to 2.2 s. |
|
||||
| Row-chunked GEMM and consumer | Retained. Chunks of 1,024 rows bound the float32 scratch buffers. Unchunked, the gate/up buffer for 8,192 tokens would be 8,192 x 4,096 x 4 bytes = 134 MB and the QKV buffer about 100 MB. |
|
||||
| Matrix-core attention | Retained. WMMA flash attention, with four waves per 16-query tile splitting the head width, replaced a VALU kernel. Each layer transposes the values once, so a lane reads 16 consecutive keys of one column. Against main with the hipBLASLt GEMMs, a 404-token request drops from 19.95 ms to 12.94 ms, which is 1.54 times faster. A request of 1,940 tokens is 2.35 times faster and one of 8,072 tokens is 4.83 times faster, while a 40-token request is unchanged at 1.02 times. Batches of 8 and 32 inputs are 1.57 and 1.77 times faster, and concurrent requests of 404 tokens are 1.60 to 1.61 times faster. Image requests change by about 1 percent, because the vision tower takes most of their time. |
|
||||
| Precomputed rotary tables | Retained. Double-precision cos/sin tables per position replace per-element `powf` and `sincosf`. |
|
||||
| Device turns per batch | Retained. One lock held for a whole request made a one-word request wait 58.4 s behind 32 inputs of 8,000 tokens. With arrival-order turns per device batch it waits 1.01 s, and the long request still takes 59.5 s. A plain mutex is not enough, because the releasing thread can take it straight back. |
|
||||
| Unmapping the checkpoint after upload | Retained. Keeping the mapping held 656 MB of file-backed RSS that nothing reads after the upload. RSS after load falls from 930 MB to 286 MB. |
|
||||
|
||||
@@ -1,16 +1,16 @@
|
||||
# EmbeddingGemma 2 quality
|
||||
|
||||
**On ten probe texts the GPU embeddings match sentence-transformers float32 output at cosine 0.999981 or better, and every token ID matches Hugging Face.** Qualified October 23, 2026, on gfx1151, with the original BF16 checkpoint at revision `914f7f89142e33e77833254d9c9b90c3cef7303b`.
|
||||
**On ten probe texts the GPU embeddings match sentence-transformers float32 output at cosine 0.999985 or better, and every token ID matches Hugging Face.** Qualified October 23, 2026, on gfx1151, with the original BF16 checkpoint at revision `914f7f89142e33e77833254d9c9b90c3cef7303b`. The cosines below are from the matrix-core attention build, measured on October 24, 2026.
|
||||
|
||||
| Check | Result |
|
||||
| --- | --- |
|
||||
| Token IDs versus Hugging Face `tokenizers` | 95/95 fixture cases identical, covering whitespace runs, byte fallback and special tokens |
|
||||
| GPU embedding versus sentence-transformers float32 | Cosine 0.999981 worst of 10 probes, 0.999999 best |
|
||||
| GPU embedding versus sentence-transformers float32 | Cosine 0.999985 worst of 10 probes, 0.999998 best |
|
||||
| Same probes, Hugging Face BF16 versus float32 | Cosine 0.999867 worst, which is the checkpoint's own precision noise |
|
||||
| Matryoshka 128 / 256 / 512 after renormalization | Worst cosine 0.999982 against the same truncation of the reference |
|
||||
| Matryoshka 128 / 256 / 512 after renormalization | Worst cosine 0.999987 against the same truncation of the reference |
|
||||
| GPU versus Gufo float32 CPU reference | Cosine 0.999992 or better on four sequences |
|
||||
| Packing sequences in one batch versus alone | Cosine 0.999992 or better. The tuned GEMM can choose a different algorithm for a different batch size |
|
||||
| Windowed attention beyond the 512-token radius | A 2,349-token probe matches the reference at cosine 0.999994 |
|
||||
| Packing sequences in one batch versus alone | Cosine 0.999987 or better. The tuned GEMM can choose a different algorithm for a different batch size |
|
||||
| Windowed attention beyond the 512-token radius | A 2,349-token probe matches the reference at cosine 0.999990 |
|
||||
|
||||
The probes are English, Spanish, and Chinese queries and passages, code, emoji, whitespace runs, literal special-token text, and one long document. They measure agreement with the reference implementation, not retrieval quality on a benchmark, and Gufo reports no MTEB score. Activations stay float32 between BF16 matrix products, so Gufo is closer to the float32 reference than the model's BF16 execution is.
|
||||
|
||||
@@ -31,7 +31,7 @@ two images in one input, and Frigate's trailing newline.
|
||||
| Check | Result |
|
||||
| --- | --- |
|
||||
| Resized patch RGB bytes versus official processor | SHA256 identical for every image, including the entire patch order |
|
||||
| GPU / HTTP embeddings versus official float32 model | Worst cosine **0.999961**, best **0.999996**; all exceed the retained 0.9999 bound |
|
||||
| GPU / HTTP embeddings versus official float32 model | Worst cosine **0.999957**, best **0.999995** with matrix-core attention, against 0.999961 and 0.999996 before it; all exceed the retained 0.9999 bound |
|
||||
| Mixed text/image batches and concurrent requests | Correct order, finite unit vectors and reference agreement |
|
||||
| Image Matryoshka 128 / 256 / 512, base64 | All exceed 0.9999 cosine against the same reference prefix |
|
||||
| Expanded context, image count, invalid destinations and malformed images | Rejected with 400; a subsequent text request succeeds |
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
{
|
||||
"date": "2026-10-24",
|
||||
"hardware": {
|
||||
"name": "AMD Radeon 8060S Graphics",
|
||||
"gfx": "gfx1151",
|
||||
"node": "gpuworker01 (Strix Halo)"
|
||||
},
|
||||
"model": "google/embeddinggemma-2 @ 914f7f89142e33e77833254d9c9b90c3cef7303b (BF16 safetensors)",
|
||||
"build": "Nix container image holding production gufo builds of main and this change, with both builds in one image. The PR head d2a9b84a differs from the measured attention commit ec21a41f only in documentation. The image and processor tests that ab8 could not run ran in a second image built the same way.",
|
||||
"image": {
|
||||
"embg2-ab8": "gufo-eg2@sha256:d4b75a8e9e615f3018f604101d8745fa349c29a00d9c93cc9a74cd4521923a80",
|
||||
"embg2-ab9": "gufo-eg2@sha256:d14fc2601ca9aecce8920161051107c4af5f263ae8308ebc758c285a111ce664"
|
||||
},
|
||||
"commits": {
|
||||
"base": "479162ed (main, with the hipBLASLt GEMMs from #751 and image inputs from #759)",
|
||||
"attn": "ec21a41f"
|
||||
},
|
||||
"log_sha256": {
|
||||
"embg2-ab8": "cc25bc529d2069a786430b013ad67f3867ad7509404b9b190e56aee1dba4f814",
|
||||
"embg2-ab9": "c3b4c57f4b5cf9887f103d5014584cf8ae7c10dd05da2dd5dd4b5b6712b21b4b"
|
||||
},
|
||||
"pod": "embg2-ab8 on gpuworker01, the only GPU pod on the node, for the tests and HTTP rows, and embg2-ab9 on gpuworker01 for the image processor and image golden tests",
|
||||
"method": "Each build ran its EmbeddingGemma 2 test binaries, then its HTTP rows twice of 60 requests each in the order base, attn, attn, base, and each value is the mean of the two runs' medians or p95s. The text rows come from tools/models/embedding_gemma2/http_bench.py. The two image rows are one 320 x 180 PNG followed by a newline in the typed content format, 269 tokens, at concurrency 1 and 4, after one warm-up request. Speedup is the base median divided by the attn median. The first server of each build also ran tests/functional/embeddings.py and tests/functional/embedding_images.py, and all passed. In ab8, embedding_gemma2_image_processor_test and embedding_gemma2_image_golden_hip_test failed for both builds because the pod script did not place their fixture at the compiled-in path, and the other seven test binaries passed. The ab9 pod ran those two tests for both builds, and all four runs passed. load_ms lists the two servers' load_completed elapsed_ms in ab8.",
|
||||
"load_ms": {
|
||||
"base": [
|
||||
2195,
|
||||
2197
|
||||
],
|
||||
"base_mean": 2196,
|
||||
"attn": [
|
||||
1992,
|
||||
1969
|
||||
],
|
||||
"attn_mean": 1981
|
||||
},
|
||||
"quality": {
|
||||
"base": {
|
||||
"golden_worst_cosine": 0.99998118,
|
||||
"golden_worst_truncated_cosine": 0.9999817,
|
||||
"long_document_cosine": 0.99999353,
|
||||
"cpu_reference_worst_cosine": 0.99999461,
|
||||
"batched_vs_alone_worst_cosine": 0.99999247,
|
||||
"image_worst_cosine": 0.9999613635260548,
|
||||
"image_best_cosine": 0.9999955757483984,
|
||||
"attention_hip_test_worst_error": 0.0153937
|
||||
},
|
||||
"attn": {
|
||||
"golden_worst_cosine": 0.99998484,
|
||||
"golden_worst_truncated_cosine": 0.99998681,
|
||||
"long_document_cosine": 0.99998986,
|
||||
"cpu_reference_worst_cosine": 0.99999234,
|
||||
"batched_vs_alone_worst_cosine": 0.99998681,
|
||||
"image_worst_cosine": 0.9999569516568002,
|
||||
"image_best_cosine": 0.9999950984605085,
|
||||
"attention_hip_test_worst_error": 0.0153937
|
||||
}
|
||||
},
|
||||
"http": [
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 1,
|
||||
"words": 16,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 40,
|
||||
"base": {
|
||||
"median_ms": 6.74,
|
||||
"p95_ms": 7.48,
|
||||
"tokens_per_s": 5768,
|
||||
"runs_median_ms": [
|
||||
6.72,
|
||||
6.76
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 6.64,
|
||||
"p95_ms": 6.98,
|
||||
"tokens_per_s": 5914,
|
||||
"runs_median_ms": [
|
||||
6.59,
|
||||
6.69
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.02
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 1,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 404,
|
||||
"base": {
|
||||
"median_ms": 19.95,
|
||||
"p95_ms": 21.13,
|
||||
"tokens_per_s": 20028,
|
||||
"runs_median_ms": [
|
||||
19.65,
|
||||
20.24
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 12.94,
|
||||
"p95_ms": 14.24,
|
||||
"tokens_per_s": 30391,
|
||||
"runs_median_ms": [
|
||||
13.03,
|
||||
12.85
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.54
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 1,
|
||||
"words": 512,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 1940,
|
||||
"base": {
|
||||
"median_ms": 144.56,
|
||||
"p95_ms": 147.28,
|
||||
"tokens_per_s": 13419,
|
||||
"runs_median_ms": [
|
||||
143.76,
|
||||
145.37
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 61.5,
|
||||
"p95_ms": 62.98,
|
||||
"tokens_per_s": 31477,
|
||||
"runs_median_ms": [
|
||||
61.66,
|
||||
61.35
|
||||
]
|
||||
},
|
||||
"median_speedup": 2.35
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 1,
|
||||
"words": 2100,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 8072,
|
||||
"base": {
|
||||
"median_ms": 1650.55,
|
||||
"p95_ms": 1657.01,
|
||||
"tokens_per_s": 4890,
|
||||
"runs_median_ms": [
|
||||
1641.25,
|
||||
1659.84
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 341.7,
|
||||
"p95_ms": 343.91,
|
||||
"tokens_per_s": 23589,
|
||||
"runs_median_ms": [
|
||||
341.54,
|
||||
341.85
|
||||
]
|
||||
},
|
||||
"median_speedup": 4.83
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 8,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 3495,
|
||||
"base": {
|
||||
"median_ms": 148.01,
|
||||
"p95_ms": 149.36,
|
||||
"tokens_per_s": 23537,
|
||||
"runs_median_ms": [
|
||||
147.9,
|
||||
148.13
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 94.27,
|
||||
"p95_ms": 94.94,
|
||||
"tokens_per_s": 36954,
|
||||
"runs_median_ms": [
|
||||
94.6,
|
||||
93.94
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.57
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": 32,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 15670,
|
||||
"base": {
|
||||
"median_ms": 701.38,
|
||||
"p95_ms": 704.08,
|
||||
"tokens_per_s": 22313,
|
||||
"runs_median_ms": [
|
||||
697.12,
|
||||
705.63
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 397.0,
|
||||
"p95_ms": 399.31,
|
||||
"tokens_per_s": 39366,
|
||||
"runs_median_ms": [
|
||||
399.26,
|
||||
394.73
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.77
|
||||
},
|
||||
{
|
||||
"concurrency": 4,
|
||||
"inputs": 1,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 404,
|
||||
"base": {
|
||||
"median_ms": 77.88,
|
||||
"p95_ms": 80.17,
|
||||
"tokens_per_s": 20554,
|
||||
"runs_median_ms": [
|
||||
77.73,
|
||||
78.02
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 48.45,
|
||||
"p95_ms": 49.19,
|
||||
"tokens_per_s": 32919,
|
||||
"runs_median_ms": [
|
||||
48.83,
|
||||
48.06
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.61
|
||||
},
|
||||
{
|
||||
"concurrency": 16,
|
||||
"inputs": 1,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 404,
|
||||
"base": {
|
||||
"median_ms": 308.69,
|
||||
"p95_ms": 312.88,
|
||||
"tokens_per_s": 20779,
|
||||
"runs_median_ms": [
|
||||
307.38,
|
||||
309.99
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 192.42,
|
||||
"p95_ms": 195.11,
|
||||
"tokens_per_s": 33264,
|
||||
"runs_median_ms": [
|
||||
193.58,
|
||||
191.26
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.6
|
||||
},
|
||||
{
|
||||
"concurrency": 4,
|
||||
"inputs": 8,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 3495,
|
||||
"base": {
|
||||
"median_ms": 574.09,
|
||||
"p95_ms": 575.45,
|
||||
"tokens_per_s": 24309,
|
||||
"runs_median_ms": [
|
||||
574.32,
|
||||
573.85
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 346.99,
|
||||
"p95_ms": 349.09,
|
||||
"tokens_per_s": 40215,
|
||||
"runs_median_ms": [
|
||||
348.51,
|
||||
345.47
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.65
|
||||
},
|
||||
{
|
||||
"concurrency": 16,
|
||||
"inputs": 8,
|
||||
"words": 128,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 3495,
|
||||
"base": {
|
||||
"median_ms": 2297.35,
|
||||
"p95_ms": 2304.22,
|
||||
"tokens_per_s": 24297,
|
||||
"runs_median_ms": [
|
||||
2298.56,
|
||||
2296.13
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 1388.06,
|
||||
"p95_ms": 1393.39,
|
||||
"tokens_per_s": 40139,
|
||||
"runs_median_ms": [
|
||||
1394.54,
|
||||
1381.58
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.66
|
||||
},
|
||||
{
|
||||
"concurrency": 1,
|
||||
"inputs": "image",
|
||||
"words": 0,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 269,
|
||||
"base": {
|
||||
"median_ms": 585.69,
|
||||
"p95_ms": 587.87,
|
||||
"tokens_per_s": 460,
|
||||
"runs_median_ms": [
|
||||
586.79,
|
||||
584.6
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 581.39,
|
||||
"p95_ms": 583.54,
|
||||
"tokens_per_s": 463,
|
||||
"runs_median_ms": [
|
||||
581.67,
|
||||
581.12
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.01
|
||||
},
|
||||
{
|
||||
"concurrency": 4,
|
||||
"inputs": "image",
|
||||
"words": 0,
|
||||
"requests_per_run": 60,
|
||||
"tokens_per_request": 269,
|
||||
"base": {
|
||||
"median_ms": 2254.74,
|
||||
"p95_ms": 2258.36,
|
||||
"tokens_per_s": 477,
|
||||
"runs_median_ms": [
|
||||
2259.24,
|
||||
2250.24
|
||||
]
|
||||
},
|
||||
"attn": {
|
||||
"median_ms": 2234.28,
|
||||
"p95_ms": 2239.15,
|
||||
"tokens_per_s": 481,
|
||||
"runs_median_ms": [
|
||||
2235.7,
|
||||
2232.86
|
||||
]
|
||||
},
|
||||
"median_speedup": 1.01
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -167,8 +167,8 @@ struct EncoderRuntime::Impl {
|
||||
// Scratch, sized for kMaxSequenceTokens.
|
||||
DeviceBuffer token_ids, positions, sequence_begin, sequence_end, offsets;
|
||||
DeviceBuffer state, embeddings_bf16, normed_bf16, qkv_f32, query_bf16,
|
||||
key_bf16, value_bf16, context_bf16, temp_f32, gate_up_f32, act_bf16,
|
||||
ple_f32, gate_f32, projected_f32, pooled_f32;
|
||||
key_bf16, value_bf16, value_t_bf16, context_bf16, temp_f32, gate_up_f32,
|
||||
act_bf16, ple_f32, gate_f32, projected_f32, pooled_f32;
|
||||
// Target of the device probe that runs after a failed batch.
|
||||
DeviceBuffer probe;
|
||||
std::atomic<bool> lost{false};
|
||||
@@ -431,7 +431,10 @@ void EncoderRuntime::Impl::Allocate() {
|
||||
qkv_f32 = DeviceBuffer(std::min(n, kChunkTokens) * qkv_width * sizeof(float));
|
||||
query_bf16 = DeviceBuffer(n * query_width * 2U);
|
||||
key_bf16 = DeviceBuffer(n * key_width * 2U);
|
||||
// The transposed copy the attention kernel reads has a stride that is a
|
||||
// multiple of 32 tokens. The transpose zero-fills the tail of each column.
|
||||
value_bf16 = DeviceBuffer(n * key_width * 2U);
|
||||
value_t_bf16 = DeviceBuffer(ValueStride(n) * key_width * 2U);
|
||||
context_bf16 = DeviceBuffer(n * query_width * 2U);
|
||||
temp_f32 = DeviceBuffer(n * hidden * sizeof(float));
|
||||
gate_up_f32 = DeviceBuffer(std::min(n, kChunkTokens) * 2U *
|
||||
@@ -532,7 +535,9 @@ void EncoderRuntime::Impl::RunBatch(
|
||||
rope_table.cos.get<float>(),
|
||||
rope_table.sin.get<float>(), epsilon, stream);
|
||||
}
|
||||
if (!LaunchAttention(query_bf16.get(), key_bf16.get(), value_bf16.get(),
|
||||
LaunchTransposeValues(value_bf16.get(), value_t_bf16.get(), tokens,
|
||||
static_cast<std::uint32_t>(key_width), stream);
|
||||
if (!LaunchAttention(query_bf16.get(), key_bf16.get(), value_t_bf16.get(),
|
||||
context_bf16.get(),
|
||||
sequence_begin.get<std::uint32_t>(),
|
||||
sequence_end.get<std::uint32_t>(), tokens, layout,
|
||||
|
||||
@@ -182,120 +182,240 @@ __global__ void QkvNormRopeKernel(
|
||||
}
|
||||
|
||||
using Vec16 = __attribute__((__vector_size__(16))) unsigned int;
|
||||
using Fragment = __attribute__((ext_vector_type(16))) unsigned short;
|
||||
using Accumulator = __attribute__((ext_vector_type(8))) float;
|
||||
|
||||
// One wave per (token, head). A lane owns HD / 32 elements held as 8-element
|
||||
// chunks, so a row is read with coalesced 16-byte loads. Keys are consumed
|
||||
// four at a time so their wave reductions overlap.
|
||||
// Loads 16 contiguous BF16 values. `source` must be 16-byte aligned.
|
||||
__device__ __forceinline__ Fragment LoadFragment(const std::uint16_t* source) {
|
||||
union {
|
||||
Fragment fragment;
|
||||
Vec16 words[2];
|
||||
} value;
|
||||
value.words[0] = *reinterpret_cast<const Vec16*>(source);
|
||||
value.words[1] = *reinterpret_cast<const Vec16*>(source + 8);
|
||||
return value.fragment;
|
||||
}
|
||||
|
||||
// Returns the value held by the lane 16 positions away, in the other
|
||||
// half-wave.
|
||||
__device__ __forceinline__ float SwapHalves(float value) {
|
||||
return __uint_as_float(__builtin_amdgcn_permlanex16(
|
||||
__float_as_uint(value), __float_as_uint(value), 0x76543210U, 0xfedcba98U,
|
||||
false, false));
|
||||
}
|
||||
|
||||
constexpr float kLog2e = 1.4426950408889634F;
|
||||
|
||||
// Transposes values from [token][column] to [column][token] through LDS so
|
||||
// both the reads and the writes are coalesced. Tokens at or beyond `tokens`
|
||||
// up to `ValueStride(tokens)` are written as zero.
|
||||
__global__ void TransposeValuesKernel(const std::uint16_t* value,
|
||||
std::uint16_t* value_t,
|
||||
std::uint32_t tokens,
|
||||
std::uint32_t columns,
|
||||
std::uint32_t value_stride) {
|
||||
__shared__ std::uint16_t tile[32][33];
|
||||
const std::uint32_t token0 = blockIdx.x * 32U;
|
||||
const std::uint32_t column0 = blockIdx.y * 32U;
|
||||
for (std::uint32_t row = threadIdx.y; row < 32U; row += blockDim.y) {
|
||||
const std::uint32_t token = token0 + row;
|
||||
tile[row][threadIdx.x] =
|
||||
token < tokens ? value[static_cast<std::size_t>(token) * columns +
|
||||
column0 + threadIdx.x]
|
||||
: static_cast<std::uint16_t>(0);
|
||||
}
|
||||
__syncthreads();
|
||||
for (std::uint32_t row = threadIdx.y; row < 32U; row += blockDim.y) {
|
||||
value_t[static_cast<std::size_t>(column0 + row) * value_stride + token0 +
|
||||
threadIdx.x] = tile[threadIdx.x][row];
|
||||
}
|
||||
}
|
||||
|
||||
// Bidirectional (windowed) attention with matrix cores.
|
||||
//
|
||||
// A block of four waves owns 16 consecutive query tokens of one head and
|
||||
// walks the keys in tiles of 32. Scores are held transposed, as keys by
|
||||
// queries, so a lane owns one query and the softmax needs no transposes. The
|
||||
// four waves split the head width. Each wave computes the partial scores over
|
||||
// its quarter of the dot product, and the partials are summed through LDS.
|
||||
// Each wave then accumulates the output for its own quarter of the columns,
|
||||
// which keeps the accumulators at 32 to 64 registers per lane.
|
||||
//
|
||||
// Values arrive transposed ([kv head][column][token]) so a lane reads sixteen
|
||||
// consecutive keys of one column. Scores use the attention scale of 1. The
|
||||
// softmax is the online base-two form, with FP32 statistics and accumulators.
|
||||
template<int HD>
|
||||
__global__ __launch_bounds__(128) void AttentionKernel(
|
||||
__global__ __launch_bounds__(128) void WmmaAttentionKernel(
|
||||
const std::uint16_t* query, const std::uint16_t* key,
|
||||
const std::uint16_t* value, std::uint16_t* context,
|
||||
const std::uint16_t* value_t, std::uint16_t* context,
|
||||
const std::uint32_t* sequence_begin, const std::uint32_t* sequence_end,
|
||||
std::uint32_t num_heads, std::uint32_t num_key_value_heads,
|
||||
std::uint32_t window) {
|
||||
constexpr int kChunks = HD / 256;
|
||||
constexpr int kPerLane = HD / kWaveSize;
|
||||
constexpr int kKeysPerStep = 4;
|
||||
std::uint32_t tokens, std::uint32_t value_stride, std::uint32_t num_heads,
|
||||
std::uint32_t num_key_value_heads, std::uint32_t window) {
|
||||
constexpr int kQuarter = HD / 4;
|
||||
constexpr int kChunks = kQuarter / 16;
|
||||
constexpr int kPartialStride = 20; // 16 floats per lane plus padding
|
||||
|
||||
const std::uint32_t lane = threadIdx.x;
|
||||
const std::uint32_t head = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (head >= num_heads) {
|
||||
return;
|
||||
}
|
||||
const std::size_t token = blockIdx.x;
|
||||
__shared__ __align__(16) float partial[2][4][32][kPartialStride];
|
||||
|
||||
const int lane = threadIdx.x & 31;
|
||||
const int wave = threadIdx.x >> 5;
|
||||
const int column = lane & 15;
|
||||
const std::uint32_t head = blockIdx.y;
|
||||
const std::uint32_t kv_head = head / (num_heads / num_key_value_heads);
|
||||
const std::int64_t begin = sequence_begin[token];
|
||||
const std::int64_t end = sequence_end[token];
|
||||
std::int64_t low = begin;
|
||||
std::int64_t high = end - 1;
|
||||
if (window != 0U) {
|
||||
low = max(begin, static_cast<std::int64_t>(token) - window);
|
||||
high = min(high, static_cast<std::int64_t>(token) + window);
|
||||
}
|
||||
const std::uint32_t first_query = blockIdx.x * 16U;
|
||||
const std::uint32_t query_token = first_query + column;
|
||||
const bool query_valid = query_token < tokens;
|
||||
|
||||
float q[kPerLane];
|
||||
const Vec16* query_row =
|
||||
reinterpret_cast<const Vec16*>(query + (token * num_heads + head) * HD);
|
||||
for (int j = 0; j < kChunks; ++j) {
|
||||
const Vec16 packed = query_row[lane + kWaveSize * j];
|
||||
const std::uint16_t* bits = reinterpret_cast<const std::uint16_t*>(&packed);
|
||||
for (int e = 0; e < 8; ++e) {
|
||||
q[j * 8 + e] = Bf16ToFloat(bits[e]);
|
||||
// Key range of this lane's query. Bounds only grow with the token index, so
|
||||
// the tile's range runs from its first query's low to its last query's high.
|
||||
const auto low_of = [&](std::uint32_t token) {
|
||||
const std::int64_t begin = sequence_begin[token];
|
||||
return window == 0U ? begin
|
||||
: max(begin, static_cast<std::int64_t>(token) - window);
|
||||
};
|
||||
const auto high_of = [&](std::uint32_t token) {
|
||||
const std::int64_t last =
|
||||
static_cast<std::int64_t>(sequence_end[token]) - 1;
|
||||
return window == 0U ? last
|
||||
: min(last, static_cast<std::int64_t>(token) + window);
|
||||
};
|
||||
const std::int64_t low = query_valid ? low_of(query_token) : 1;
|
||||
const std::int64_t high = query_valid ? high_of(query_token) : 0;
|
||||
const std::int64_t tile_low = low_of(first_query);
|
||||
const std::int64_t tile_high = high_of(min(first_query + 15U, tokens - 1U));
|
||||
|
||||
Fragment q[kChunks];
|
||||
for (int c = 0; c < kChunks; ++c) {
|
||||
if (query_valid) {
|
||||
q[c] = LoadFragment(
|
||||
query +
|
||||
(static_cast<std::size_t>(query_token) * num_heads + head) * HD +
|
||||
wave * kQuarter + c * 16);
|
||||
} else {
|
||||
q[c] = Fragment{};
|
||||
}
|
||||
}
|
||||
|
||||
float accumulator[kPerLane];
|
||||
for (int e = 0; e < kPerLane; ++e) {
|
||||
accumulator[e] = 0.0F;
|
||||
Accumulator result[kChunks];
|
||||
for (int t = 0; t < kChunks; ++t) {
|
||||
result[t] = Accumulator{};
|
||||
}
|
||||
float maximum = -INFINITY;
|
||||
float denominator = 0.0F;
|
||||
float sum = 0.0F;
|
||||
|
||||
for (std::int64_t first = low; first <= high; first += kKeysPerStep) {
|
||||
float score[kKeysPerStep];
|
||||
for (int u = 0; u < kKeysPerStep; ++u) {
|
||||
float partial = 0.0F;
|
||||
if (first + u <= high) {
|
||||
const Vec16* key_row = reinterpret_cast<const Vec16*>(
|
||||
key + ((first + u) * num_key_value_heads + kv_head) * HD);
|
||||
for (int j = 0; j < kChunks; ++j) {
|
||||
const Vec16 packed = key_row[lane + kWaveSize * j];
|
||||
const std::uint16_t* bits =
|
||||
reinterpret_cast<const std::uint16_t*>(&packed);
|
||||
for (int e = 0; e < 8; ++e) {
|
||||
partial += q[j * 8 + e] * Bf16ToFloat(bits[e]);
|
||||
}
|
||||
int parity = 0;
|
||||
for (std::int64_t base = tile_low & ~std::int64_t{31}; base <= tile_high;
|
||||
base += 32) {
|
||||
Accumulator scores[2] = {Accumulator{}, Accumulator{}};
|
||||
for (int c = 0; c < kChunks; ++c) {
|
||||
for (int h = 0; h < 2; ++h) {
|
||||
std::int64_t row = base + h * 16 + column;
|
||||
row = min(row, static_cast<std::int64_t>(tokens) - 1);
|
||||
const Fragment k = LoadFragment(
|
||||
key +
|
||||
(static_cast<std::size_t>(row) * num_key_value_heads + kv_head) *
|
||||
HD +
|
||||
wave * kQuarter + c * 16);
|
||||
scores[h] =
|
||||
__builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(k, q[c], scores[h]);
|
||||
}
|
||||
}
|
||||
|
||||
// Sum the four waves' partial scores in a fixed order, so every wave sees
|
||||
// identical values and keeps identical softmax statistics. The two buffers
|
||||
// alternate, so one barrier per tile separates the reads of one tile from
|
||||
// the writes of the next.
|
||||
float* mine = partial[parity][wave][lane];
|
||||
for (int h = 0; h < 2; ++h) {
|
||||
for (int i = 0; i < 8; i += 4) {
|
||||
*reinterpret_cast<float4*>(mine + h * 8 + i) = make_float4(
|
||||
scores[h][i], scores[h][i + 1], scores[h][i + 2], scores[h][i + 3]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
float total[16];
|
||||
for (int v = 0; v < 16; v += 4) {
|
||||
float4 sum4 =
|
||||
*reinterpret_cast<const float4*>(&partial[parity][0][lane][v]);
|
||||
for (int w = 1; w < 4; ++w) {
|
||||
const float4 other =
|
||||
*reinterpret_cast<const float4*>(&partial[parity][w][lane][v]);
|
||||
sum4.x += other.x;
|
||||
sum4.y += other.y;
|
||||
sum4.z += other.z;
|
||||
sum4.w += other.w;
|
||||
}
|
||||
total[v] = sum4.x;
|
||||
total[v + 1] = sum4.y;
|
||||
total[v + 2] = sum4.z;
|
||||
total[v + 3] = sum4.w;
|
||||
}
|
||||
parity ^= 1;
|
||||
|
||||
// Mask to this query's key range and move to the base-two domain. Register
|
||||
// i of half h holds the key at offset 2 * i + (lane / 16) of that half.
|
||||
float next = maximum;
|
||||
for (int h = 0; h < 2; ++h) {
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
const std::int64_t position = base + h * 16 + 2 * i + (lane >> 4);
|
||||
float score = total[h * 8 + i] * kLog2e;
|
||||
if (position < low || position > high) {
|
||||
score = -INFINITY;
|
||||
}
|
||||
}
|
||||
score[u] = partial;
|
||||
}
|
||||
for (int offset = kWaveSize / 2; offset > 0; offset >>= 1) {
|
||||
for (int u = 0; u < kKeysPerStep; ++u) {
|
||||
score[u] += __shfl_xor(score[u], offset, kWaveSize);
|
||||
total[h * 8 + i] = score;
|
||||
next = fmaxf(next, score);
|
||||
}
|
||||
}
|
||||
float step_maximum = maximum;
|
||||
for (int u = 0; u < kKeysPerStep; ++u) {
|
||||
if (first + u > high) {
|
||||
score[u] = -INFINITY;
|
||||
next = fmaxf(next, SwapHalves(next));
|
||||
// Until a query has seen a key, its maximum is -INFINITY. Shift by zero
|
||||
// then, so the exponentials below stay finite.
|
||||
const float shifted = next == -INFINITY ? 0.0F : next;
|
||||
const float correction = __builtin_amdgcn_exp2f(maximum - shifted);
|
||||
|
||||
float subtotal = 0.0F;
|
||||
Fragment probabilities[2];
|
||||
for (int h = 0; h < 2; ++h) {
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
const float p = __builtin_amdgcn_exp2f(total[h * 8 + i] - shifted);
|
||||
// The denominator sums the rounded weights the value product uses.
|
||||
const unsigned own = FloatToBf16(p);
|
||||
subtotal += Bf16ToFloat(static_cast<std::uint16_t>(own));
|
||||
const unsigned peer = __builtin_amdgcn_permlanex16(
|
||||
own, own, 0x76543210U, 0xfedcba98U, false, false);
|
||||
probabilities[h][2 * i] = lane < 16 ? own : peer;
|
||||
probabilities[h][2 * i + 1] = lane < 16 ? peer : own;
|
||||
}
|
||||
step_maximum = fmaxf(step_maximum, score[u]);
|
||||
}
|
||||
const float correction = expf(maximum - step_maximum);
|
||||
maximum = step_maximum;
|
||||
denominator *= correction;
|
||||
for (int e = 0; e < kPerLane; ++e) {
|
||||
accumulator[e] *= correction;
|
||||
subtotal += SwapHalves(subtotal);
|
||||
sum = sum * correction + subtotal;
|
||||
maximum = next;
|
||||
for (int t = 0; t < kChunks; ++t) {
|
||||
result[t] *= correction;
|
||||
}
|
||||
for (int u = 0; u < kKeysPerStep; ++u) {
|
||||
if (first + u > high) {
|
||||
break;
|
||||
}
|
||||
const float weight = expf(score[u] - maximum);
|
||||
denominator += weight;
|
||||
const Vec16* value_row = reinterpret_cast<const Vec16*>(
|
||||
value + ((first + u) * num_key_value_heads + kv_head) * HD);
|
||||
for (int j = 0; j < kChunks; ++j) {
|
||||
const Vec16 packed = value_row[lane + kWaveSize * j];
|
||||
const std::uint16_t* bits =
|
||||
reinterpret_cast<const std::uint16_t*>(&packed);
|
||||
for (int e = 0; e < 8; ++e) {
|
||||
accumulator[j * 8 + e] += weight * Bf16ToFloat(bits[e]);
|
||||
}
|
||||
|
||||
for (int t = 0; t < kChunks; ++t) {
|
||||
const std::size_t column_row = static_cast<std::size_t>(kv_head) * HD +
|
||||
wave * kQuarter + t * 16 + column;
|
||||
for (int h = 0; h < 2; ++h) {
|
||||
const Fragment v =
|
||||
LoadFragment(value_t + column_row * value_stride + base + h * 16);
|
||||
result[t] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(
|
||||
v, probabilities[h], result[t]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const float inverse = 1.0F / denominator;
|
||||
Vec16* out_row =
|
||||
reinterpret_cast<Vec16*>(context + (token * num_heads + head) * HD);
|
||||
for (int j = 0; j < kChunks; ++j) {
|
||||
Vec16 packed;
|
||||
std::uint16_t* bits = reinterpret_cast<std::uint16_t*>(&packed);
|
||||
for (int e = 0; e < 8; ++e) {
|
||||
bits[e] = FloatToBf16(accumulator[j * 8 + e] * inverse);
|
||||
if (query_valid) {
|
||||
const float inverse = 1.0F / sum;
|
||||
std::uint16_t* out =
|
||||
context +
|
||||
(static_cast<std::size_t>(query_token) * num_heads + head) * HD;
|
||||
for (int t = 0; t < kChunks; ++t) {
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
out[wave * kQuarter + t * 16 + 2 * i + (lane >> 4)] =
|
||||
FloatToBf16(result[t][i] * inverse);
|
||||
}
|
||||
}
|
||||
out_row[lane + kWaveSize * j] = packed;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -435,8 +555,20 @@ void LaunchQkvNormRope(const float* qkv, const void* query_norm_bf16,
|
||||
layout.num_key_value_heads, layout.head_dim, rope_cos, rope_sin, epsilon);
|
||||
}
|
||||
|
||||
void LaunchTransposeValues(const void* value_bf16, void* value_t_bf16,
|
||||
std::size_t tokens, std::uint32_t columns,
|
||||
hipStream_t stream) {
|
||||
const dim3 grid(static_cast<unsigned int>((tokens + 31U) / 32U),
|
||||
columns / 32U);
|
||||
TransposeValuesKernel<<<grid, dim3(32, 8), 0, stream>>>(
|
||||
static_cast<const std::uint16_t*>(value_bf16),
|
||||
static_cast<std::uint16_t*>(value_t_bf16),
|
||||
static_cast<std::uint32_t>(tokens), columns,
|
||||
static_cast<std::uint32_t>(ValueStride(tokens)));
|
||||
}
|
||||
|
||||
bool LaunchAttention(const void* query_bf16, const void* key_bf16,
|
||||
const void* value_bf16, void* context_bf16,
|
||||
const void* value_t_bf16, void* context_bf16,
|
||||
const std::uint32_t* sequence_begin,
|
||||
const std::uint32_t* sequence_end, std::size_t tokens,
|
||||
const QkvLayout& layout, std::uint32_t window,
|
||||
@@ -445,22 +577,22 @@ bool LaunchAttention(const void* query_bf16, const void* key_bf16,
|
||||
layout.num_heads % layout.num_key_value_heads != 0 || tokens == 0) {
|
||||
return false;
|
||||
}
|
||||
constexpr unsigned int kHeadsPerBlock = 4;
|
||||
const dim3 grid(static_cast<unsigned int>(tokens),
|
||||
(layout.num_heads + kHeadsPerBlock - 1) / kHeadsPerBlock);
|
||||
const dim3 block(kWaveSize, kHeadsPerBlock);
|
||||
const dim3 grid(static_cast<unsigned int>((tokens + 15U) / 16U),
|
||||
layout.num_heads);
|
||||
const auto* query = static_cast<const std::uint16_t*>(query_bf16);
|
||||
const auto* key = static_cast<const std::uint16_t*>(key_bf16);
|
||||
const auto* value = static_cast<const std::uint16_t*>(value_bf16);
|
||||
const auto* value_t = static_cast<const std::uint16_t*>(value_t_bf16);
|
||||
auto* context = static_cast<std::uint16_t*>(context_bf16);
|
||||
const auto token_count = static_cast<std::uint32_t>(tokens);
|
||||
const auto value_stride = static_cast<std::uint32_t>(ValueStride(tokens));
|
||||
if (layout.head_dim == 256) {
|
||||
AttentionKernel<256><<<grid, block, 0, stream>>>(
|
||||
query, key, value, context, sequence_begin, sequence_end,
|
||||
layout.num_heads, layout.num_key_value_heads, window);
|
||||
WmmaAttentionKernel<256><<<grid, 128, 0, stream>>>(
|
||||
query, key, value_t, context, sequence_begin, sequence_end, token_count,
|
||||
value_stride, layout.num_heads, layout.num_key_value_heads, window);
|
||||
} else if (layout.head_dim == 512) {
|
||||
AttentionKernel<512><<<grid, block, 0, stream>>>(
|
||||
query, key, value, context, sequence_begin, sequence_end,
|
||||
layout.num_heads, layout.num_key_value_heads, window);
|
||||
WmmaAttentionKernel<512><<<grid, 128, 0, stream>>>(
|
||||
query, key, value_t, context, sequence_begin, sequence_end, token_count,
|
||||
value_stride, layout.num_heads, layout.num_key_value_heads, window);
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -56,12 +56,29 @@ void LaunchQkvNormRope(const float* qkv, const void* query_norm_bf16,
|
||||
const float* rope_sin, float epsilon,
|
||||
hipStream_t stream);
|
||||
|
||||
/// Tokens per column of the transposed value buffer: `tokens` rounded up to a
|
||||
/// multiple of 32.
|
||||
inline constexpr std::size_t ValueStride(std::size_t tokens) {
|
||||
return (tokens + 31U) / 32U * 32U;
|
||||
}
|
||||
|
||||
/// Transposes values from [token][column] to [column][token] with
|
||||
/// `ValueStride(tokens)` tokens per column, zero-filling the tokens past
|
||||
/// `tokens`. `columns` is a multiple of 32. This is the layout
|
||||
/// `LaunchAttention` reads.
|
||||
void LaunchTransposeValues(const void* value_bf16, void* value_t_bf16,
|
||||
std::size_t tokens, std::uint32_t columns,
|
||||
hipStream_t stream);
|
||||
|
||||
/// Bidirectional attention with scale 1. A token attends to the keys of its
|
||||
/// own sequence within `window` positions on either side; `window` equal to
|
||||
/// zero means the whole sequence. Query, key and value are
|
||||
/// [token][head][column]. Returns false for an unsupported geometry.
|
||||
/// zero means the whole sequence. Query and key are [token][head][column];
|
||||
/// the value is transposed, [kv head][column][token], with
|
||||
/// `ValueStride(tokens)` entries per column, as `LaunchTransposeValues`
|
||||
/// writes it. Value reads run up to that stride, so the padding must be
|
||||
/// finite. Returns false for an unsupported geometry.
|
||||
[[nodiscard]] bool LaunchAttention(const void* query_bf16, const void* key_bf16,
|
||||
const void* value_bf16, void* context_bf16,
|
||||
const void* value_t_bf16, void* context_bf16,
|
||||
const std::uint32_t* sequence_begin,
|
||||
const std::uint32_t* sequence_end,
|
||||
std::size_t tokens, const QkvLayout& layout,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
// Checks the attention kernel against a direct float32 evaluation for both
|
||||
// head widths, packed sequences, and windowed and whole-sequence masks.
|
||||
// Checks the value transpose and the attention kernel against a direct
|
||||
// double-precision evaluation for both head widths, packed sequences, and
|
||||
// windowed and whole-sequence masks.
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
|
||||
@@ -53,7 +54,8 @@ T* ToDevice(const std::vector<T>& host) {
|
||||
}
|
||||
|
||||
void RunCase(std::uint32_t head_dim, std::uint32_t kv_heads,
|
||||
std::uint32_t window, const std::vector<std::uint32_t>& lengths) {
|
||||
std::uint32_t window, const std::vector<std::uint32_t>& lengths,
|
||||
float query_scale = 1.0F) {
|
||||
constexpr std::uint32_t kHeads = 4;
|
||||
std::size_t tokens = 0;
|
||||
for (const auto length : lengths) {
|
||||
@@ -61,14 +63,14 @@ void RunCase(std::uint32_t head_dim, std::uint32_t kv_heads,
|
||||
}
|
||||
std::mt19937 random(head_dim + window);
|
||||
std::normal_distribution<float> normal(0.0F, 1.0F);
|
||||
const auto fill = [&](std::size_t heads) {
|
||||
const auto fill = [&](std::size_t heads, float scale = 1.0F) {
|
||||
std::vector<std::uint16_t> values(tokens * heads * head_dim);
|
||||
for (auto& value : values) {
|
||||
value = ToBf16(normal(random));
|
||||
value = ToBf16(normal(random) * scale);
|
||||
}
|
||||
return values;
|
||||
};
|
||||
const auto query = fill(kHeads);
|
||||
const auto query = fill(kHeads, query_scale);
|
||||
const auto key = fill(kv_heads);
|
||||
const auto value = fill(kv_heads);
|
||||
|
||||
@@ -85,7 +87,32 @@ void RunCase(std::uint32_t head_dim, std::uint32_t kv_heads,
|
||||
|
||||
auto* d_query = ToDevice(query);
|
||||
auto* d_key = ToDevice(key);
|
||||
auto* d_value = ToDevice(value);
|
||||
// The kernel reads values transposed, [kv head][column][token], padded to
|
||||
// a multiple of 32 tokens with zeros. The device transpose must match the
|
||||
// host one, including the zero tail. The buffer starts as NaN to show that
|
||||
// the tail is written.
|
||||
const std::size_t stride = hip::ValueStride(tokens);
|
||||
const std::size_t columns = std::size_t{kv_heads} * head_dim;
|
||||
std::vector<std::uint16_t> expected_value_t(columns * stride);
|
||||
for (std::size_t t = 0; t < tokens; ++t) {
|
||||
for (std::size_t c = 0; c < columns; ++c) {
|
||||
expected_value_t[c * stride + t] = value[t * columns + c];
|
||||
}
|
||||
}
|
||||
auto* d_untransposed = ToDevice(value);
|
||||
void* d_value = nullptr;
|
||||
Check(hipMalloc(&d_value, expected_value_t.size() * 2U) == hipSuccess,
|
||||
"hipMalloc value_t");
|
||||
Check(hipMemset(d_value, 0xff, expected_value_t.size() * 2U) == hipSuccess,
|
||||
"memset value_t");
|
||||
hip::LaunchTransposeValues(d_untransposed, d_value, tokens,
|
||||
static_cast<std::uint32_t>(columns), nullptr);
|
||||
std::vector<std::uint16_t> value_t(expected_value_t.size());
|
||||
Check(hipMemcpy(value_t.data(), d_value, value_t.size() * 2U,
|
||||
hipMemcpyDeviceToHost) == hipSuccess,
|
||||
"download value_t");
|
||||
Check(value_t == expected_value_t,
|
||||
"device transpose differs from the host transpose");
|
||||
auto* d_begin = ToDevice(begin);
|
||||
auto* d_end = ToDevice(end);
|
||||
std::uint16_t* d_context = nullptr;
|
||||
@@ -144,10 +171,12 @@ void RunCase(std::uint32_t head_dim, std::uint32_t kv_heads,
|
||||
std::cout << "head " << head_dim << " window " << window << " tokens "
|
||||
<< tokens << ": worst error " << worst << '\n';
|
||||
// Outputs are rounded to BF16, whose relative step is 2^-8.
|
||||
Check(worst < 0.02, "attention differs from the float32 evaluation: " +
|
||||
std::to_string(worst));
|
||||
Check(worst < 0.02,
|
||||
"attention differs from the double-precision evaluation: " +
|
||||
std::to_string(worst));
|
||||
(void)hipFree(d_query);
|
||||
(void)hipFree(d_key);
|
||||
(void)hipFree(d_untransposed);
|
||||
(void)hipFree(d_value);
|
||||
(void)hipFree(d_begin);
|
||||
(void)hipFree(d_end);
|
||||
@@ -169,6 +198,9 @@ int main() {
|
||||
RunCase(256, 2, 100, {700});
|
||||
RunCase(256, 2, 0, {31, 32, 33, 1});
|
||||
RunCase(512, 1, 0, {32, 16, 15});
|
||||
RunCase(256, 2, 0, {32, 32});
|
||||
RunCase(512, 1, 5, {64});
|
||||
RunCase(256, 2, 0, {200}, 0.05F);
|
||||
Check(!hip::LaunchAttention(nullptr, nullptr, nullptr, nullptr, nullptr,
|
||||
nullptr, 1, {8, 2, 128}, 0, nullptr),
|
||||
"an unsupported head dimension is rejected");
|
||||
|
||||
@@ -26,6 +26,7 @@ WORKLOADS = [
|
||||
(1, 1, 16),
|
||||
(1, 1, 128),
|
||||
(1, 1, 512),
|
||||
(1, 1, 2100),
|
||||
(1, 8, 128),
|
||||
(1, 32, 128),
|
||||
(4, 1, 128),
|
||||
|
||||
Reference in new issue
Block a user