fix(vae): order Sol attention and FFN writes to reduce VRAM

This commit is contained in:
leejet committed 2026-11-10 23:56:53 +08:00
1 parent ddc367ab5d
commit 0678d04ae7
2 files changed
+2

No files matched your search

@@ -172,6 +172,7 @@ namespace MiniMaxH3VAE {
heads, nullptr, false, runner_ctx->flash_attn_enabled); heads, nullptr, false, runner_ctx->flash_attn_enabled);
auto indices = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows)); auto indices = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows));
output = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices); output = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices);
runner_ctx->expand_graph(output);
} }
// Later windows can only read this block and its immediate successor. // Later windows can only read this block and its immediate successor.
previous = current; previous = current;
+1
View File
@@ -430,6 +430,7 @@ namespace MiniMaxH3VAE {
auto input = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false); auto input = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false);
auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false); auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false);
output = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices); output = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices);
ctx->expand_graph(output);
} }
return output; return output;
} }