xycld · Nov 4, 2026 · Nov 4, 2026 · Nov 4, 2026
@@ -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 <step>/<steps>
 <seconds>`, `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,10 @@ 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 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.

@@ -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<std::vector<std::int64_t>>;
+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<std::uint64_t>(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,30 +124,25 @@ 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) {
     const std::string p = "text_embedder.blocks." + std::to_string(b) + ".";
     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<std::vector<std::int64_t>> 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<std::string>{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<VkDeviceSize> o = place(block_names(l), block_shapes, on_device ? Buffer::Kind::kDevice : Buffer::Kind::kStaging);
+    const std::vector<VkDeviceSize> 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

@@ -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);
 

@@ -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,13 +65,29 @@ 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());
     }
   }
 
+  // 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 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.
   std::vector<std::vector<std::uint16_t>> embeds;
   {
@@ -79,8 +108,8 @@ 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;
+    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);

@@ -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<void(const std::vector<float>& latents)> on_latents;
   // stage: "text encoder", "denoise", "vae"; done of total within it.