From ba30ce77d233daa4d7202322b849616475bcca56 Mon Sep 17 00:00:00 2001 From: xycld Date: Wed, 4 Nov 2026 01:52:39 -0500 Subject: [PATCH 1/2] moss-infer-vulkan: the engine's request checks, --help/--version, --vram-budget Found by the moss-infer-hip review, which mirrors this program: - --help prints the usage to stdout and exits 0; --version prints the version. - pipeline::validate()'s bounds and messages: a whitespace-only prompt, at most 4.2 megapixels and 4096 pixels on a side, at most 10000 steps (was 1000), the seed, cfg, shift and max-seq messages; --device below -1 is refused. - The PNG metadata names the bundle by util::path_key(model).filename(), as the engine does (a trailing slash or "--model ." gave an empty name or "."), and records noise_layout official. - --vram-budget GiB, with vram_option_bytes' bounds: caps the device memory of the VAE and the DiT (weights, the DiT's activations, the decode); the DiT's blocks past it stream. A budget that cannot hold the decode or the activations is refused before the text encoder runs. RTX 4090 Laptop, 512x512 20 steps: --vram-budget 4 keeps 10 of 42 blocks resident and gives the same latents and PNG bytes as no budget; 0.5 GiB is refused before any work; 2 GiB at 1024x1024 streams every block and completes. --- docs/vulkan.md | 11 +++-- vulkan/generate.cpp | 34 +++++++++++-- vulkan/generate.h | 3 ++ vulkan/tools/moss_infer_vulkan.cpp | 79 ++++++++++++++++++++++++------ 4 files changed, 104 insertions(+), 23 deletions(-) diff --git a/docs/vulkan.md b/docs/vulkan.md index 49429f2..a0c1fd4 100644 --- a/docs/vulkan.md +++ b/docs/vulkan.md @@ -8,7 +8,7 @@ slower than CUDA; on a supported NVIDIA GPU, use `moss-infer`. It runs txt2img with the official Base and SFT bf16 files, the official euler sampler and schedule and classifier-free guidance with the uncond skip layers. The tier is `reference`. Everything else is refused with a message that names the option: other samplers and schedulers, LoRAs, img2img and inpainting, hires fix, control adapters, references, -the lossy accelerators, quantized or GGUF files, and the CUDA memory options. The tokenizer, the seed's initial noise +the lossy accelerators, quantized or GGUF files, and the CUDA memory options other than `--vram-budget`. The tokenizer, the seed's initial noise and the schedule are the engine's own, so a seed gives the same inputs as on CUDA. ```sh @@ -18,11 +18,12 @@ moss-infer-vulkan --list-devices The flags are those of `moss-infer generate` that the backend runs, with the same defaults (`--negative`, `--size`, `--width`, `--height`, `--steps`, `--cfg`, `--shift`, `--seed`, `--skip-layers`, `--max-seq`, `--latent-out`, -`--no-metadata`, `--quiet`; `--tier` accepts `reference`). `--device N` picks a device (default: the first usable one, a +`--no-metadata`, `--quiet`, `--vram-budget`; `--tier` accepts `reference`), checked with the engine's limits and +messages (at most 4.2 megapixels and 4096 pixels on a side, at most 10000 steps). `--device N` picks a device (default: the first usable one, a discrete GPU first) and `--kernels coop|coop-f16|plain` the kernels (default: the first in that order the device has); both are printed at the start of every run. Progress goes to stderr one line per event: `text encoder`, `denoise / `, `vae`, then a `done:` summary. The exit status is 0 on success, 2 for a usage error or an unsupported -request, 1 for any other failure. +request, 1 for any other failure. `--help` prints the usage to stdout, `--version` the version. ## Kernels @@ -46,7 +47,9 @@ of its 42 blocks on the device as fit next to the run's activations, from the dr before them runs. So a GPU with less than the DiT's 9.5 GiB still runs it, at the cost of host memory and bus traffic. The VAE decodes the whole image at once: about 1.6 GiB at 1024x1024, 3.1 GiB at 1456x1456 and 6 GiB at 2048x2048. A decode the device cannot hold is refused before the text encoder runs. The text encoder and the DiT are on the device -one after the other, never together. +one after the other, never together. `--vram-budget GiB` caps the device memory of the VAE and the DiT (their weights, +the DiT's activations and the decode), as `moss-infer generate`'s does: the DiT's blocks past it stream, and a run whose +activations or decode it cannot hold is refused before the text encoder runs. Long operations are split into submits of bounded work, so no submit runs long enough to trip the Windows display driver's watchdog on a slow GPU. diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp index 6ef986c..8cf0660 100644 --- a/vulkan/generate.cpp +++ b/vulkan/generate.cpp @@ -43,6 +43,19 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r vae.emplace(device, r.model_dir / "vae" / "diffusion_pytorch_model.safetensors", path == BlockPath::kCoop ? VaePath::kCoop : VaePath::kPlain); } + // What the DiT's activations and blocks and the VAE's decode may take: the device's free memory, within the budget + // less the VAE's weights (held through the run). + const std::uint64_t vae_weights = vae ? vae->weight_bytes() : 0; + const auto available = [&] { + const std::uint64_t free = device.free_device_memory(); + if (r.vram_budget == 0) return free; + return std::min(free, r.vram_budget > vae_weights ? r.vram_budget - vae_weights : 0); + }; + // ", within --vram-budget N MiB" when the budget, not the device, sets available(). + const auto budget_note = [&]() -> std::string { + if (r.vram_budget == 0 || available() >= device.free_device_memory()) return ""; + return " within --vram-budget " + std::to_string(r.vram_budget >> 20) + " MiB"; + }; if (r.decode) { const Capabilities& caps = device.info().caps; const std::uint64_t need = vae->workspace_bytes(lh, lw), largest = vae->largest_buffer_bytes(lh, lw); @@ -52,10 +65,24 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r std::to_string(largest >> 20) + " MiB buffer; " + device.info().name + " takes at most " + std::to_string(limit >> 20) + " MiB"); } - if (need > device.free_device_memory()) { + if (need > available()) { throw std::runtime_error("the VAE decode at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs " + std::to_string(need >> 20) + " MiB of device memory; " + device.info().name + " has " + - std::to_string(device.free_device_memory() >> 20) + " MiB free"); + std::to_string(available() >> 20) + " MiB free" + budget_note()); + } + } + + // A budget that cannot hold the DiT's activations (one text token per prompt token) is refused before any work. + constexpr std::uint64_t kReserve = std::uint64_t{512} << 20; // pipelines, command buffers, the driver + if (r.vram_budget != 0) { + std::uint64_t tokens = std::uint64_t{lh} * lw + r.cond_ids.size(); + if (cfg) tokens += std::uint64_t{lh} * lw + r.uncond_ids.size(); + const std::uint64_t need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; + if (need > available()) { + throw std::runtime_error("the DiT at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs " + + std::to_string(need >> 20) + " MiB of device memory for its activations; " + + std::to_string(available() >> 20) + " MiB are free" + + budget_note()); } } @@ -79,8 +106,7 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r // Blocks past what the device can hold with this run's activations stream from host memory. std::uint64_t tokens = 0; for (const auto& e : embeds) tokens += std::uint64_t{lh} * lw + e.size() / kDiTTextDim; - constexpr std::uint64_t kReserve = std::uint64_t{512} << 20; // pipelines, command buffers, the driver - const std::uint64_t free = device.free_device_memory(), need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; + const std::uint64_t free = available(), need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; DitOptions options; options.device_budget = free > need ? free - need : 1; Dit dit(device, kernels, r.model_dir / "transformer", options); diff --git a/vulkan/generate.h b/vulkan/generate.h index d770166..ab93597 100644 --- a/vulkan/generate.h +++ b/vulkan/generate.h @@ -28,6 +28,9 @@ struct GenerateRequest { float cfg = 3.0f; SkipSet uncond_skip; bool decode = true; // false: stop at the final latents + // Device memory the VAE and the DiT may use, in bytes: their weights, the DiT's activations and the VAE's decode; the + // DiT's blocks past it stream from host memory. 0: what the device has free. A decode it cannot hold is refused first. + std::uint64_t vram_budget = 0; // Called with the final latents before the decode. std::function& latents)> on_latents; // stage: "text encoder", "denoise", "vae"; done of total within it. diff --git a/vulkan/tools/moss_infer_vulkan.cpp b/vulkan/tools/moss_infer_vulkan.cpp index 6a1dc74..5e3072f 100644 --- a/vulkan/tools/moss_infer_vulkan.cpp +++ b/vulkan/tools/moss_infer_vulkan.cpp @@ -1,6 +1,6 @@ // moss-infer-vulkan: text to image on the Vulkan backend, a standalone program for GPUs the CUDA engine does not support. // It needs neither CUDA nor an NVIDIA driver, only a Vulkan driver. Its flags are those of `moss-infer generate` that the -// backend runs; the tokenizer, the seed's initial noise, the schedule and the PNG writer are the engine's own sources, +// backend runs, checked as the engine checks them; the tokenizer, the seed's initial noise, the schedule and the PNG writer are the engine's own sources, // so a seed gives the same inputs as on CUDA. Exit status: 0 done, 1 failed, 2 a usage error or an unsupported request. #include #include @@ -27,33 +27,46 @@ #include "tokenizer/prompt_tokens.h" #include "tokenizer/qwen_bpe.h" #include "util/path.h" +#include "util/regular_file.h" namespace { namespace fs = std::filesystem; using namespace mi; +// pipeline.h's request bounds (kMaxSide, kMaxPixels, kMaxSteps, kMaxSeed, kMaxVramBytes), which validate() applies; +// pipeline.h itself needs the CUDA headers. +constexpr int kMaxSide = 4096; +constexpr std::int64_t kMaxPixels = 4'200'000; +constexpr int kMaxSteps = 10000; +constexpr std::int64_t kMaxSeed = std::int64_t{1} << 53; +constexpr double kMaxVramBytes = static_cast(std::uint64_t{1024} << 30); + constexpr const char* kUsage = "usage: moss-infer-vulkan --model --prompt [-o ] [--latent-out ]\n" " [--negative ] [--size x] [--width N] [--height N] [--steps N] [--cfg X]\n" " [--shift X] [--seed N|-1] [--skip-layers \"22\"] [--max-seq N] [--tier reference]\n" - " [--device N] [--kernels coop|coop-f16|plain] [--no-metadata] [--quiet]\n" - " moss-infer-vulkan --list-devices\n" + " [--device N] [--kernels coop|coop-f16|plain] [--vram-budget GiB] [--no-metadata]\n" + " [--quiet]\n" + " moss-infer-vulkan --list-devices | --version | --help\n" "Text to image with MOSS-Image Base or SFT (the bundle's bf16 files) on a Vulkan GPU, with the official euler sampler\n" "and schedule: the fallback for GPUs the CUDA engine (moss-infer) does not support. The flags mean what they mean for\n" "`moss-infer generate`; the defaults are the same (1024x1024, 50 steps, cfg 3, shift 5, seed 42, skip layers 22).\n" "At least one of -o / --latent-out is required; without -o the VAE decode is skipped. --device picks a Vulkan device\n" "(default: the first usable one, a discrete GPU first); --kernels picks the cooperative-matrix kernels on bf16 or f16\n" - "fragments or the plain ones (default: the first of those the device has). The device and kernels are printed at the\n" - "start. Progress goes to stderr one line per event: `text encoder`, `denoise / `, `vae`, then\n" + "fragments or the plain ones (default: the first of those the device has). --vram-budget caps the device memory of\n" + "the VAE and the DiT: their weights, the DiT's activations and the VAE's decode; the DiT's blocks past it stream from\n" + "host memory every step (default: what the device has free). A run the budget or the device cannot hold is refused\n" + "before the text encoder runs. The device and kernels are printed at the start. Progress goes to stderr one line per event: `text encoder`, `denoise / `, `vae`, then\n" "a summary.\n"; struct Args { fs::path model, out_png, out_latent; std::string prompt, negative; - bool have_prompt = false, metadata = true, quiet = false, list = false; + bool have_prompt = false, metadata = true, quiet = false, list = false, help = false, version = false; int width = 1024, height = 1024, steps = 50, max_seq = 4096, device = -1; double cfg = 3.0, shift = 5.0; + std::uint64_t vram_budget = 0; std::int64_t seed = 42; std::vector skip = {22}; std::string kernels; @@ -110,6 +123,15 @@ std::vector layers(const std::string& v) { return out; } +// pipeline::vram_option_bytes: GiB to bytes, with its bounds and messages. +std::uint64_t vram_bytes(const std::string& flag, double gib) { + if (!std::isfinite(gib) || gib <= 0) usage(flag + " must be a finite positive number."); + const double bytes = gib * static_cast(std::uint64_t{1} << 30); + if (!(bytes <= kMaxVramBytes)) usage(flag + " must be at most 1024 GiB."); + if (bytes < static_cast(std::uint64_t{1} << 20)) usage(flag + " must be at least 1 MiB."); + return static_cast(bytes); +} + Args parse(const std::vector& argv) { Args a; for (std::size_t i = 1; i < argv.size(); ++i) { @@ -126,7 +148,14 @@ Args parse(const std::vector& argv) { a.list = true; continue; } - if (f == "-h" || f == "--help") usage(""); + if (f == "-h" || f == "--help") { + a.help = true; + return a; + } + if (f == "--version") { + a.version = true; + return a; + } if (i + 1 >= argv.size()) usage(f + " needs a value"); const std::string& v = argv[++i]; if (f == "--model") a.model = util::path_from_utf8(v); @@ -150,6 +179,7 @@ Args parse(const std::vector& argv) { else if (f == "--max-seq") a.max_seq = int32(f, v); else if (f == "--device") a.device = int32(f, v); else if (f == "--kernels") a.kernels = v; + else if (f == "--vram-budget") a.vram_budget = vram_bytes(f, number(f, v)); else if (f == "--tier") { // The backend runs the bf16 files exactly as the reference tier does; the lossy tiers need the CUDA kernels. if (v != "reference") usage("moss-infer-vulkan runs --tier reference only"); @@ -159,14 +189,23 @@ Args parse(const std::vector& argv) { if (a.list) return a; if (a.model.empty() || !a.have_prompt) usage("--model and --prompt are required"); if (a.out_png.empty() && a.out_latent.empty()) usage("at least one of -o / --latent-out is required"); - if (a.width < 16 || a.height < 16 || a.width % 16 != 0 || a.height % 16 != 0 || a.width > 4096 || a.height > 4096) { - usage("width and height must be multiples of 16 in [16, 4096]"); + if (a.device < -1) usage("--device takes a device index, or -1 for the first usable device"); + // pipeline::validate()'s checks and messages, for the settings this program takes. + if (tok::python_strip(a.prompt).empty()) usage("Enter a prompt."); + if (a.height <= 0 || a.width <= 0 || a.height % 16 != 0 || a.width % 16 != 0) { + usage("Height and width must be positive multiples of 16."); + } + if (a.height > kMaxSide || a.width > kMaxSide) usage("Height and width must be at most " + std::to_string(kMaxSide) + "."); + if (std::int64_t{a.height} * a.width > kMaxPixels) { + usage("Image size must be at most 4.2 megapixels (e.g. 2048x2048); the model is meant for 0.25 to 4 megapixels."); + } + if (a.steps < 1) usage("Steps and image count must be positive."); + if (a.steps > kMaxSteps) usage("Steps must be at most " + std::to_string(kMaxSteps) + "."); + if (a.cfg < 0 || a.shift <= 0) usage("CFG must be finite and nonnegative; shift must be finite and positive."); + if (a.seed < -1 || a.seed >= kMaxSeed) usage("Seed must be -1 (random) or an integer from 0 to 2^53-1 (browser-safe range)."); + if (a.max_seq < 1 || a.max_seq > static_cast(tok::kMossMaxSequenceLength)) { + usage("Maximum text length must be between 1 and 4096."); } - if (a.steps < 1 || a.steps > 1000) usage("--steps must be in [1, 1000]"); - if (a.cfg < 0) usage("--cfg must be >= 0"); - if (a.shift <= 0) usage("--shift must be > 0"); - if (a.seed < -1 || a.seed >= (std::int64_t{1} << 53)) usage("--seed must be -1 (random) or in [0, 2^53)"); - if (a.max_seq < 1 || a.max_seq > static_cast(tok::kMossMaxSequenceLength)) usage("--max-seq must be in [1, 4096]"); if (!a.kernels.empty() && a.kernels != "coop" && a.kernels != "coop-f16" && a.kernels != "plain") { usage("--kernels takes coop, coop-f16 or plain"); } @@ -238,6 +277,7 @@ int run(const Args& a) { g.timesteps = schedule.timesteps; g.cfg = static_cast(a.cfg); for (const int l : a.skip) g.uncond_skip.set(static_cast(l)); + g.vram_budget = a.vram_budget; g.decode = !a.out_png.empty(); if (!a.out_latent.empty()) { g.on_latents = [&](const std::vector& latents) { @@ -271,7 +311,7 @@ int run(const Args& a) { m.engine_version = std::to_string(MI_VERSION_MAJOR) + "." + std::to_string(MI_VERSION_MINOR) + "." + std::to_string(MI_VERSION_PATCH); m.tier = "reference"; - m.models.push_back({"bundle", util::path_utf8(a.model.filename()), ""}); + m.models.push_back({"bundle", util::path_utf8(util::path_key(a.model).filename()), ""}); m.prompt = a.prompt; m.negative_prompt = a.negative; m.sampler = "euler"; @@ -282,6 +322,7 @@ int run(const Args& a) { m.seed = seed; m.width = a.width; m.height = a.height; + m.extra.emplace_back("noise_layout", pipeline::noise_layout_name(pipeline::NoiseLayout::kOfficial)); m.extra.emplace_back("backend", std::string("vulkan ") + vk::block_path_name(path)); text = pipeline::generation_text_chunks(m); } @@ -305,6 +346,14 @@ int run_main(const std::vector& argv) { std::fputs(kUsage, stderr); return 2; } + if (a.help) { + std::fputs(kUsage, stdout); + return 0; + } + if (a.version) { + std::printf("moss-infer-vulkan %d.%d.%d\n", MI_VERSION_MAJOR, MI_VERSION_MINOR, MI_VERSION_PATCH); + return 0; + } try { return a.list ? list_devices() : run(a); } catch (const std::invalid_argument& e) { From 11d68fb68c722d47ef5bfbfa8095d5da9353dbbf Mon Sep 17 00:00:00 2001 From: xycld Date: Wed, 4 Nov 2026 01:57:50 -0500 Subject: [PATCH 2/2] moss-infer-vulkan: --vram-budget counts the DiT's weights with every block streamed Review of the previous commit: the early refusal counted the DiT's activations only, so a budget just above them passed, then the DiT's global tensors, text refiner and two block slots (about 1 GiB) went past it. Dit::minimum_weight_bytes() gives that footprint from the shape lists the constructor checks (now shared); the refusal before the text encoder and a second one before the DiT loads (should the encoder's lengths differ from the token counts) both include it, instead of falling back to every block streamed past the budget. RTX 4090 Laptop: 1.3 GiB at 512x512 and 2 GiB at 1024x1024 (both ran before, over the budget) are refused before the text encoder; 3 GiB at 1024x1024 runs with 3 blocks resident; 4 GiB at 512x512 gives the same latents as no budget. --- docs/vulkan.md | 5 +++-- vulkan/dit.cpp | 38 ++++++++++++++++++++++++++++---------- vulkan/dit.h | 2 ++ vulkan/generate.cpp | 19 +++++++++++-------- 4 files changed, 44 insertions(+), 20 deletions(-) diff --git a/docs/vulkan.md b/docs/vulkan.md index a0c1fd4..ca754a8 100644 --- a/docs/vulkan.md +++ b/docs/vulkan.md @@ -48,8 +48,9 @@ before them runs. So a GPU with less than the DiT's 9.5 GiB still runs it, at th The VAE decodes the whole image at once: about 1.6 GiB at 1024x1024, 3.1 GiB at 1456x1456 and 6 GiB at 2048x2048. A decode the device cannot hold is refused before the text encoder runs. The text encoder and the DiT are on the device one after the other, never together. `--vram-budget GiB` caps the device memory of the VAE and the DiT (their weights, -the DiT's activations and the decode), as `moss-infer generate`'s does: the DiT's blocks past it stream, and a run whose -activations or decode it cannot hold is refused before the text encoder runs. +the DiT's activations and the decode), as `moss-infer generate`'s does: the DiT's blocks past it stream, and a run it +cannot hold with every block streamed, or whose decode it cannot hold, is refused before the text encoder runs. With +`-o` the VAE's weights (about 0.1 GiB) count against it. Long operations are split into submits of bounded work, so no submit runs long enough to trip the Windows display driver's watchdog on a slow GPU. diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index 5932a62..555a7bd 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -16,6 +16,26 @@ namespace { constexpr std::uint32_t kDim = kDiTDim, kFfn = kDiTFfnDim, kHead = kDiTHeadDim, kQkv = 3 * kDiTDim; constexpr VkDeviceSize kBf16 = 2; +// The bf16 tensors the constructor loads, in its order: the global ones, one text refiner block, one DiT block. +using Shapes = std::vector>; +constexpr std::int64_t kD = kDim, kF = kFfn, kH = kHead; +const Shapes kGlobalShapes = {{kD, kDiTLatentChannels}, {kD}, {kD, kDiTFreqDim}, {kD}, {kD, kD}, {kD}, {6 * kD, kD}, + {6 * kD}, {2 * kD, kD}, {2 * kD}, {kD}, {kDiTLatentChannels, kD}, {kD, kDiTTextDim}, {kD}}; +const Shapes kTextBlockShapes = {{kD}, {kD}, {kH}, {kH}, {3 * kD, kD}, {3 * kD}, {kD, kD}, {kD}, + {kF, kD}, {kF}, {kF, kD}, {kF}, {kD, kF}, {kD}}; +const Shapes kBlockShapes = {{kD}, {kD}, {kH}, {kH}, {kD, kD}, {kD, kD}, {kD, kD}, {kD, kD}, {2 * kF, kD}, {kD, kF}}; + +// Device bytes of tensors placed at 256-byte aligned offsets, as the constructor places them. +std::uint64_t placed_bytes(const Shapes& shapes) { + std::uint64_t total = 0; + for (const auto& shape : shapes) { + std::uint64_t n = kBf16; + for (const std::int64_t e : shape) n *= static_cast(e); + total += (n + 255) / 256 * 256; + } + return total; +} + // The RoPE band frequencies 1 / 10000^(i / 16), as the reference builds them in fp32 on the CPU (kernels/basic/rope.h). constexpr std::uint32_t kBandBits[16] = {0x3f800000, 0x3f0ff59a, 0x3ea1e89b, 0x3e361887, 0x3dcccccd, 0x3d6655c2, 0x3d0186e3, 0x3c91ad39, 0x3c23d70a, 0x3bb8449c, 0x3b4f3e38, 0x3ae91528, @@ -104,15 +124,12 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat for (const VkDeviceSize o : offsets) addresses.push_back(owned_.back()->address() + o); return addresses; }; - const std::int64_t d = kDim, f = kFfn; - auto a = load({"video_embedder.weight", "video_embedder.bias", "time_embedding.0.weight", "time_embedding.0.bias", "time_embedding.2.weight", "time_embedding.2.bias", "shared_block_modulation.weight", "shared_block_modulation.bias", "video_head_modulation.weight", "video_head_modulation.bias", "final_norm_video.weight", "final_linear_video.weight", "text_embedder.proj.weight", "text_embedder.proj.bias"}, - {{d, kDiTLatentChannels}, {d}, {d, kDiTFreqDim}, {d}, {d, d}, {d}, {6 * d, d}, {6 * d}, {2 * d, d}, {2 * d}, - {d}, {kDiTLatentChannels, d}, {d, kDiTTextDim}, {d}}); + kGlobalShapes); g_ = {a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], a[9], a[10], a[11], a[12], a[13]}; for (int b = 0; b < 2; ++b) { @@ -120,14 +137,12 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat auto t = load({p + "norm1.weight", p + "norm2.weight", p + "q_norm.weight", p + "k_norm.weight", p + "qkv.weight", p + "qkv.bias", p + "proj.weight", p + "proj.bias", p + "ffn_up.weight", p + "ffn_up.bias", p + "ffn_gate.weight", p + "ffn_gate.bias", p + "ffn_down.weight", p + "ffn_down.bias"}, - {{d}, {d}, {kHead}, {kHead}, {3 * d, d}, {3 * d}, {d, d}, {d}, {f, d}, {f}, {f, d}, {f}, {d, f}, {d}}); + kTextBlockShapes); text_blocks_.push_back({t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7], t[8], t[9], t[10], t[11], t[12], t[13]}); } // Blocks: the first `resident` in device memory, the rest in host memory, each copied into one device slot before it // runs. q, k and v sit back to back (256-aligned sizes of 18 MiB each), so they form the [9216, 3072] QKV weight. - const std::vector> block_shapes = {{d}, {d}, {kHead}, {kHead}, {d, d}, {d, d}, {d, d}, {d, d}, - {2 * f, d}, {d, f}}; const auto block_names = [](int l) { const std::string p = "layers." + std::to_string(l) + "."; return std::vector{p + "norm1.weight", p + "norm2.weight", p + "q_norm.weight", p + "k_norm.weight", @@ -138,8 +153,7 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat if (o[5] != o[4] + kDim * kDim * kBf16 || o[6] != o[5] + kDim * kDim * kBf16) fail("q, k and v are not contiguous"); return BlockWeights{base + o[0], base + o[1], base + o[2], base + o[3], base + o[4], base + o[7], base + o[8], base + o[9]}; }; - VkDeviceSize block_bytes = 0; - for (const std::string& n : block_names(0)) block_bytes += (index.at(n).bytes + 255) / 256 * 256; + const VkDeviceSize block_bytes = placed_bytes(kBlockShapes); int resident = kDiTLayers; if (options.device_budget > 0) { const std::uint64_t free = options.device_budget > weight_bytes_ ? options.device_budget - weight_bytes_ : 0; @@ -149,7 +163,7 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat } for (int l = 0; l < kDiTLayers; ++l) { const bool on_device = l < resident; - const std::vector o = place(block_names(l), block_shapes, on_device ? Buffer::Kind::kDevice : Buffer::Kind::kStaging); + const std::vector o = place(block_names(l), kBlockShapes, on_device ? Buffer::Kind::kDevice : Buffer::Kind::kStaging); if (on_device) { blocks_.push_back(block_weights(owned_.back()->address(), o)); host_blocks_.push_back(nullptr); @@ -247,6 +261,10 @@ void Dit::run_block_op(int op, const BlockWeights& w, const BlockBuffers& b) { } } +std::uint64_t Dit::minimum_weight_bytes() { + return placed_bytes(kGlobalShapes) + 2 * placed_bytes(kTextBlockShapes) + 2 * placed_bytes(kBlockShapes); +} + std::uint64_t Dit::activation_bytes(std::uint64_t tokens, std::uint64_t hw) { // forward(): x, y and attn [tokens, 3072]; qkv [tokens, 9216]; ffn [tokens, 8192]; rope [tokens, 128]; the video tokens // and the head's projection [hw, 128]; the modulation and time scratch. encode_text() reuses them for a prompt no longer diff --git a/vulkan/dit.h b/vulkan/dit.h index 46c64b8..8008704 100644 --- a/vulkan/dit.h +++ b/vulkan/dit.h @@ -49,6 +49,8 @@ class Dit { std::uint64_t weight_bytes() const { return weight_bytes_; } // device memory, the streaming slot included std::uint64_t host_weight_bytes() const { return host_weight_bytes_; } int resident_blocks() const { return resident_blocks_; } + // Device memory of the weights with every block streamed: the global tensors, the text refiner and the two block slots. + static std::uint64_t minimum_weight_bytes(); // Device memory of the activations of one forward() over `tokens` joint tokens and an h x w latent grid. static std::uint64_t activation_bytes(std::uint64_t tokens, std::uint64_t hw); diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp index 8cf0660..3f06173 100644 --- a/vulkan/generate.cpp +++ b/vulkan/generate.cpp @@ -72,18 +72,20 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r } } - // A budget that cannot hold the DiT's activations (one text token per prompt token) is refused before any work. + // The least device memory the DiT runs in: its activations (one text token per prompt token, as the text encoder + // gives), the weights with every block streamed, and the reserve. A budget that cannot hold it is refused before any + // work, and again (should the lengths differ) before the DiT loads, rather than exceeded. constexpr std::uint64_t kReserve = std::uint64_t{512} << 20; // pipelines, command buffers, the driver + const auto refuse_dit = [&](std::uint64_t least) { + throw std::runtime_error("the DiT at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs at least " + + std::to_string(least >> 20) + " MiB of device memory (its activations, every block streamed); " + + std::to_string(available() >> 20) + " MiB are free" + budget_note()); + }; if (r.vram_budget != 0) { std::uint64_t tokens = std::uint64_t{lh} * lw + r.cond_ids.size(); if (cfg) tokens += std::uint64_t{lh} * lw + r.uncond_ids.size(); - const std::uint64_t need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; - if (need > available()) { - throw std::runtime_error("the DiT at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs " + - std::to_string(need >> 20) + " MiB of device memory for its activations; " + - std::to_string(available() >> 20) + " MiB are free" + - budget_note()); - } + const std::uint64_t least = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve + Dit::minimum_weight_bytes(); + if (least > available()) refuse_dit(least); } // Text encoder: streamed layer by layer, freed before the DiT loads. @@ -107,6 +109,7 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r std::uint64_t tokens = 0; for (const auto& e : embeds) tokens += std::uint64_t{lh} * lw + e.size() / kDiTTextDim; const std::uint64_t free = available(), need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; + if (r.vram_budget != 0 && need + Dit::minimum_weight_bytes() > free) refuse_dit(need + Dit::minimum_weight_bytes()); DitOptions options; options.device_budget = free > need ? free - need : 1; Dit dit(device, kernels, r.model_dir / "transformer", options);