MiniMax-H3: flash attention in the video VAE, faster tile stitching, … · unslothai/stable-diffusion.cpp@dcb4f1c · GitHub
Skip to content

Commit dcb4f1c

Browse files
danielhanchenoobabooga
danielhanchen
and
oobabooga
authored
MiniMax-H3: flash attention in the video VAE, faster tile stitching, ggml-cuda speed patch (#17)
* MiniMax-H3: flash attention in the video VAE under --diffusion-fa; carry ggml-cuda speed patch The H3 video VAE decoder is a 36-layer ViT run per 16x16 latent tile. With only --diffusion-fa (what Studio passes) it used mul_mat + scale + softmax attention, materialising a 32 x L x L f32 score matrix per tile and layer. --diffusion-fa now also covers it; SD_H3_VAE_FLASH_ATTN=0 restores the old path. The text encoder is unchanged (that is what --fa would also switch, which moves the conditioning). scripts/unsloth/ggml-patches/ carries two ggml-cuda commits on top of the pinned submodule (BF16 cuBLAS for large-batch quantized matmul behind GGML_CUDA_QUANT_CUBLAS_MIN_BATCH, a cached device integrated flag, a vectorized row copy); the prebuilt workflow applies them after the submodule checkout. * tiling: walk tile planes contiguously when splitting and merging process_tiles_2d copied and blended each pixel across all planes in the inner loop, striding a whole output plane (960x544 floats for H3) per element. A video tile has frames x channels planes, so every element was a cache and TLB miss: on MiniMax-H3 960x544x124 (105 tiles) this was about 10 s of single-threaded host time per decode. Loop planes outermost and precompute the per-column and per-row blend weights; each element still computes old + new * wy * wx in the same order, so the output is bitwise identical. * Tighten H3 VAE and tiling comments * Keep the H3 VAE on its old attention outside CUDA and ROCm --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
1 parent 6321a69 commit dcb4f1c

5 files changed

Lines changed: 330 additions & 26 deletions

File tree

‎.github/workflows/unsloth-sd-prebuilt.yml‎

Lines changed: 6 additions & 0 deletions

‎docs/minimax_h3.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,12 @@ are detected from their weights.
4646
Omitting `--audio-vae` still runs the joint diffusion model but produces video without a
4747
decoded audio track.
4848

49+
On CUDA and ROCm, `--diffusion-fa` also enables flash attention in the video VAE decoder (a ViT),
50+
which removes most of its attention cost without touching the text encoder. Set
51+
`SD_H3_VAE_FLASH_ATTN=0` to decode with the previous mul_mat + softmax attention. Other backends
52+
keep that attention by default (on Vulkan the flash attention decode was slower);
53+
`SD_H3_VAE_FLASH_ATTN=1` turns it on there.
54+
4955
## First/last-frame conditioning
5056

5157
Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA:

‎scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
From 22b560daa5e90bdb72c492f0b55667468ad3fd6b Mon Sep 17 00:00:00 2001
2+
From: Daniel Han <danielhanchen@gmail.com>
3+
Date: Wed, 30 Sep 2026 22:28:21 -0700
4+
Subject: [PATCH 1/2] ggml-cuda: BF16 cuBLAS for large-batch quantized matmul,
5+
cache device integrated flag
6+
7+
GGML_CUDA_QUANT_CUBLAS_MIN_BATCH=<rows> dequantizes a quantized weight to BF16
8+
and runs a cuBLAS GEMM (FP32 accumulate, FP32 output) instead of MMQ when src1
9+
has at least that many rows; 0 or unset keeps MMQ everywhere.
10+
GGML_CUDA_QUANT_CUBLAS_TYPE=f32 dequantizes to FP32 instead (accuracy oracle).
11+
At diffusion-transformer sizes (MiniMax-H3 960x544x124: 19108 tokens per call)
12+
MMQ reaches 130-240 TFLOP/s on a B200 while the BF16 GEMM reaches ~1300, and
13+
the result is closer to an FP32 reference because activations stay BF16
14+
instead of being quantized to q8_1.
15+
16+
ggml_backend_cuda_device_get_memory / get_type called cudaGetDeviceProperties
17+
(~1-2 ms) on every call; stable-diffusion.cpp polls them per graph, 2798 times
18+
in one tiled H3 VAE decode. The flag is now read once at registration.
19+
---
20+
src/ggml-cuda/ggml-cuda.cu | 55 ++++++++++++++++++++++++++++++++------
21+
1 file changed, 47 insertions(+), 8 deletions(-)
22+
23+
diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu
24+
index e15fdc5e..517d0894 100644
25+
--- a/src/ggml-cuda/ggml-cuda.cu
26+
+++ b/src/ggml-cuda/ggml-cuda.cu
27+
@@ -2427,6 +2427,25 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
28+
return use_mul_mat_vec_q;
29+
}
30+
31+
+// Large-batch matmul with a quantized weight: dequantize the weight to BF16 once and run a cuBLAS
32+
+// GEMM (FP32 accumulate, FP32 output) instead of MMQ. MMQ quantizes the activations to q8_1 and runs
33+
+// int8 MMA, which is the right call for LLM prompt sizes, but at diffusion-transformer sizes (10^4
34+
+// tokens per call) a tensor-core BF16 GEMM is faster on GPUs whose BF16 rate is not far below their
35+
+// int8 rate, and it is more accurate (activations stay BF16 instead of 8-bit).
36+
+// GGML_CUDA_QUANT_CUBLAS_MIN_BATCH=<rows>: use the BF16 path when src1 has at least this many rows;
37+
+// 0 disables it (stock MMQ everywhere). Unset = the per-architecture default below.
38+
+static int64_t ggml_cuda_quant_cublas_min_batch(int cc) {
39+
+ static const int64_t env = [] {
40+
+ const char * e = getenv("GGML_CUDA_QUANT_CUBLAS_MIN_BATCH");
41+
+ return e != nullptr ? (int64_t) atoll(e) : (int64_t) -1;
42+
+ }();
43+
+ if (env >= 0) {
44+
+ return env;
45+
+ }
46+
+ GGML_UNUSED(cc);
47+
+ return 0;
48+
+}
49+
+
50+
static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
51+
GGML_TENSOR_BINARY_OP_LOCALS
52+
53+
@@ -2492,6 +2511,23 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
54+
ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
55+
return;
56+
}
57+
+ if (ggml_is_quantized(src0->type) && ne11 > MMVQ_MAX_BATCH_SIZE && GGML_CUDA_CC_IS_NVIDIA(cc) &&
58+
+ bf16_mma_hardware_available(cc) && ggml_get_to_bf16_cuda(src0->type) != nullptr) {
59+
+ const int64_t min_batch = ggml_cuda_quant_cublas_min_batch(cc);
60+
+ if (min_batch > 0 && ne11*ne12*ne13 >= min_batch) {
61+
+ // GGML_CUDA_QUANT_CUBLAS_TYPE=f32 dequantizes to FP32 instead: slow, for accuracy oracles only.
62+
+ static const bool f32_oracle = [] {
63+
+ const char * e = getenv("GGML_CUDA_QUANT_CUBLAS_TYPE");
64+
+ return e != nullptr && (strcmp(e, "f32") == 0 || strcmp(e, "fp32") == 0);
65+
+ }();
66+
+ if (f32_oracle) {
67+
+ ggml_cuda_mul_mat_cublas_impl<GGML_TYPE_F32>(ctx, src0, src1, dst);
68+
+ } else {
69+
+ ggml_cuda_mul_mat_cublas_impl<GGML_TYPE_BF16>(ctx, src0, src1, dst);
70+
+ }
71+
+ return;
72+
+ }
73+
+ }
74+
if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
75+
ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
76+
return;
77+
@@ -5200,6 +5236,10 @@ struct ggml_backend_cuda_device_context {
78+
std::string description;
79+
std::string pci_bus_id;
80+
int op_offload_min_batch_size;
81+
+ // cudaDeviceProp::integrated, read once at registration. cudaGetDeviceProperties costs ~1-2 ms per call, and
82+
+ // get_memory / get_type are polled per graph by callers that budget VRAM (stable-diffusion.cpp asks thousands
83+
+ // of times per tiled VAE decode, ~5 s of host time per H3 video on a B200 host).
84+
+ bool integrated = false;
85+
};
86+
87+
static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) {
88+
@@ -5303,12 +5343,9 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
89+
// ref: https://github.com/ggml-org/llama.cpp/pull/17368
90+
#if defined(__linux__)
91+
// Check if this is a UMA (Unified Memory Architecture) system
92+
- cudaDeviceProp prop;
93+
- CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
94+
-
95+
// Check if UMA is explicitly enabled via environment variable
96+
bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr;
97+
- bool is_uma = prop.integrated > 0 || uma_env;
98+
+ bool is_uma = ctx->integrated || uma_env;
99+
100+
if (is_uma) {
101+
// For UMA systems (like DGX Spark), use system memory info
102+
@@ -5332,10 +5369,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
103+
static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) {
104+
ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *) dev->context;
105+
106+
- cudaDeviceProp prop;
107+
- CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
108+
-
109+
- return prop.integrated
110+
+ return ctx->integrated
111+
? GGML_BACKEND_DEVICE_TYPE_IGPU
112+
: GGML_BACKEND_DEVICE_TYPE_GPU;
113+
}
114+
@@ -6095,6 +6129,11 @@ ggml_backend_reg_t ggml_backend_cuda_reg() {
115+
dev_ctx->device = i;
116+
dev_ctx->name = GGML_CUDA_NAME + std::to_string(i);
117+
dev_ctx->description = ggml_cuda_device_description(i);
118+
+ {
119+
+ cudaDeviceProp prop;
120+
+ CUDA_CHECK(cudaGetDeviceProperties(&prop, physical_id));
121+
+ dev_ctx->integrated = prop.integrated > 0;
122+
+ }
123+
124+
char pci_bus_id[32] = {};
125+
CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), physical_id));
126+
--
127+
2.43.0
128+
129+
130+
From 0bf02a4e324ef837685c63ebaeddf2b4a4c0f317 Mon Sep 17 00:00:00 2001
131+
From: Daniel Han <danielhanchen@gmail.com>
132+
Date: Wed, 30 Sep 2026 22:54:01 -0700
133+
Subject: [PATCH 2/2] ggml-cuda: vectorized row copy for strided F32->F32/F16
134+
copies with contiguous rows
135+
136+
cpy_scalar recomputes the 4-D source and destination index (64-bit div/mod) for
137+
every element, so the permuted ggml_cont copies of a DiT's Q/K/V (rows of
138+
head_dim floats) run at a few percent of memory bandwidth: a 411 MB
139+
[128,42,19108] permute copy takes 1.67 ms on a B200, 0.30 ms as one float4 per
140+
thread with the index math done once per row. Pure copy / same per-element
141+
conversion, so the output is bit-identical; GGML_CUDA_CPY_ROWS=0 restores
142+
cpy_scalar.
143+
---
144+
src/ggml-cuda/cpy.cu | 83 ++++++++++++++++++++++++++++++++++++++++++++
145+
1 file changed, 83 insertions(+)
146+
147+
diff --git a/src/ggml-cuda/cpy.cu b/src/ggml-cuda/cpy.cu
148+
index 4c6fb4a2..c26a081c 100644
149+
--- a/src/ggml-cuda/cpy.cu
150+
+++ b/src/ggml-cuda/cpy.cu
151+
@@ -1,4 +1,5 @@
152+
#include "cpy.cuh"
153+
+#include <type_traits>
154+
#include "dequantize.cuh"
155+
#include "cpy-utils.cuh"
156+
#include "fp8.cuh"
157+
@@ -203,6 +204,84 @@ cudaStream_t stream) {
158+
ggml_cuda_kernel_launch(cpy_scalar_contiguous<src_t, dst_t>, launch_params, cx, cdst, ne);
159+
}
160+
161+
+// Strided copy whose rows (dim 0) are contiguous on both sides and of equal length: one float4 per thread, the
162+
+// 4-D index math done once per row instead of per element. cpy_scalar spends most of its time in the 64-bit
163+
+// div/mod of that per-element math; the permuted ggml_cont copies of a DiT's Q/K/V (rows of head_dim floats) run
164+
+// at a few percent of memory bandwidth through it. A pure copy (or the same per-element conversion), so the
165+
+// result is bit-identical to cpy_scalar. GGML_CUDA_CPY_ROWS=0 disables it.
166+
+template <typename dst_t>
167+
+static __global__ void cpy_rows_vec4(const char * cx, char * cdst, const int64_t nrows, const int64_t ne00_4,
168+
+ const int64_t ne01, const int64_t ne02, const int64_t nb01, const int64_t nb02, const int64_t nb03,
169+
+ const int64_t ne11, const int64_t ne12, const int64_t nb11, const int64_t nb12, const int64_t nb13) {
170+
+ ggml_cuda_pdl_lc();
171+
+ const int64_t row = (int64_t) blockIdx.x*blockDim.y + threadIdx.y;
172+
+ if (row >= nrows) {
173+
+ return;
174+
+ }
175+
+ const int64_t i01 = row % ne01;
176+
+ const int64_t i02 = (row / ne01) % ne02;
177+
+ const int64_t i03 = row / (ne01*ne02);
178+
+ const int64_t i11 = row % ne11;
179+
+ const int64_t i12 = (row / ne11) % ne12;
180+
+ const int64_t i13 = row / (ne11*ne12);
181+
+ const float4 * src = (const float4 *) (cx + i01*nb01 + i02*nb02 + i03*nb03);
182+
+ char * dst_row = cdst + i11*nb11 + i12*nb12 + i13*nb13;
183+
+ ggml_cuda_pdl_sync();
184+
+ for (int64_t c = threadIdx.x; c < ne00_4; c += blockDim.x) {
185+
+ const float4 v = src[c];
186+
+ if constexpr (std::is_same_v<dst_t, float>) {
187+
+ ((float4 *) dst_row)[c] = v;
188+
+ } else {
189+
+ dst_t * d = (dst_t *) dst_row + 4*c;
190+
+ cpy_1_scalar<float, dst_t>((const char *) &v.x, (char *) (d + 0));
191+
+ cpy_1_scalar<float, dst_t>((const char *) &v.y, (char *) (d + 1));
192+
+ cpy_1_scalar<float, dst_t>((const char *) &v.z, (char *) (d + 2));
193+
+ cpy_1_scalar<float, dst_t>((const char *) &v.w, (char *) (d + 3));
194+
+ }
195+
+ }
196+
+}
197+
+
198+
+static bool ggml_cuda_cpy_rows_enabled() {
199+
+ static const bool enabled = [] {
200+
+ const char * e = getenv("GGML_CUDA_CPY_ROWS");
201+
+ return e == nullptr || atoi(e) != 0;
202+
+ }();
203+
+ return enabled;
204+
+}
205+
+
206+
+// True (and launched) when src0 -> src1 is a row copy cpy_rows_vec4 can do; false leaves the copy to the caller.
207+
+template <typename dst_t>
208+
+static bool ggml_cpy_rows_vec4_cuda(const ggml_tensor * src0, ggml_tensor * src1, cudaStream_t stream) {
209+
+ if (!ggml_cuda_cpy_rows_enabled() || src0->type != GGML_TYPE_F32) {
210+
+ return false;
211+
+ }
212+
+ const int64_t ne00 = src0->ne[0];
213+
+ if (src0->nb[0] != (int64_t) sizeof(float) || src1->nb[0] != (int64_t) sizeof(dst_t) || src1->ne[0] != ne00 || ne00 % 4 != 0) {
214+
+ return false;
215+
+ }
216+
+ const size_t dst_align = std::is_same_v<dst_t, float> ? 16 : 2*sizeof(dst_t);
217+
+ if (((uintptr_t) src0->data | src0->nb[1] | src0->nb[2] | src0->nb[3]) % 16 != 0 ||
218+
+ ((uintptr_t) src1->data | src1->nb[1] | src1->nb[2] | src1->nb[3]) % dst_align != 0) {
219+
+ return false;
220+
+ }
221+
+ const int64_t nrows = ggml_nrows(src0);
222+
+ if (nrows != ggml_nrows(src1) || src0->ne[1] == 0 || src0->ne[2] == 0 || src1->ne[1] == 0 || src1->ne[2] == 0) {
223+
+ return false;
224+
+ }
225+
+ const int64_t ne00_4 = ne00 / 4;
226+
+ const int threads_x = ne00_4 >= 32 ? 32 : (int) ne00_4;
227+
+ const int threads_y = 256 / threads_x;
228+
+ const int64_t blocks = (nrows + threads_y - 1) / threads_y;
229+
+ if (blocks > INT_MAX) {
230+
+ return false;
231+
+ }
232+
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(dim3((unsigned) blocks), dim3(threads_x, threads_y), 0, stream);
233+
+ ggml_cuda_kernel_launch(cpy_rows_vec4<dst_t>, launch_params, (const char *) src0->data, (char *) src1->data, nrows, ne00_4,
234+
+ src0->ne[1], src0->ne[2], (int64_t) src0->nb[1], (int64_t) src0->nb[2], (int64_t) src0->nb[3],
235+
+ src1->ne[1], src1->ne[2], (int64_t) src1->nb[1], (int64_t) src1->nb[2], (int64_t) src1->nb[3]);
236+
+ return true;
237+
+}
238+
+
239+
template<typename src_t, typename dst_t, bool transposed = false>
240+
static void ggml_cpy_scalar_cuda(
241+
const char * cx, char * cdst, const int64_t ne,
242+
@@ -509,6 +588,10 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
243+
ggml_cpy_scalar_cuda<ggml_fp8_e5m2_cuda, nv_bfloat16>
244+
(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
245+
}
246+
+ } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && ggml_cpy_rows_vec4_cuda<float>(src0, src1, main_stream)) {
247+
+ // row copy done
248+
+ } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && !contiguous_srcs && ggml_cpy_rows_vec4_cuda<half>(src0, src1, main_stream)) {
249+
+ // row copy with conversion done
250+
} else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) {
251+
if (can_be_transposed) {
252+
ggml_cpy_scalar_cuda<float, float, true>
253+
--
254+
2.43.0
255+

‎src/pipeline/diffusion_engine.cpp‎

Lines changed: 14 additions & 0 deletions

0 commit comments

Comments
 (0)