feat: add MiniMax-H3-Image support (#2324)

CMakeLists.txt

@@ -9,13 +9,28 @@ if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
 endif()
 
 if (MSVC)
+    cmake_host_system_information(RESULT SD_LOGICAL_CORES QUERY NUMBER_OF_LOGICAL_CORES)
+    if (SD_LOGICAL_CORES GREATER 8)
+        math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 2")
+    elseif (SD_LOGICAL_CORES GREATER 2)
+        math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 1")
+    else()
+        set(SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES}")
+    endif()
+    if (NOT SD_DEFAULT_MSVC_MP_COUNT GREATER 0)
+        set(SD_DEFAULT_MSVC_MP_COUNT 1)
+    endif()
+    set(SD_MSVC_MP_COUNT "${SD_DEFAULT_MSVC_MP_COUNT}" CACHE STRING "Maximum parallel MSVC compilations per project")
+    if (NOT SD_MSVC_MP_COUNT MATCHES "^[1-9][0-9]*$" OR SD_MSVC_MP_COUNT GREATER 65536)
+        message(FATAL_ERROR "SD_MSVC_MP_COUNT must be an integer between 1 and 65536")
+    endif()
     add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
     add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
     # /MP is MSVC-only: icx rejects it outright once offloading is enabled.
     add_compile_options(
-        $<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP>
+        $<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP${SD_MSVC_MP_COUNT}>
         $<$<COMPILE_LANGUAGE:C>:/utf-8>
-        $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP>
+        $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP${SD_MSVC_MP_COUNT}>
         $<$<COMPILE_LANGUAGE:CXX>:/utf-8>
     )
 endif()

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/11/11** 🚀 stable-diffusion.cpp adds **Day-0 support for MiniMax-H3-Image**
 * **2026/10/16** 🚀 stable-diffusion.cpp adds **Day-0 support for MOSS-Image-Base / MOSS-Image-SFT**
 * **2026/10/16** 🚀 stable-diffusion.cpp adds **Day-0 support for FLUX.3 Image Klein / Klein Base**
 * **2026/09/20** 🚀 stable-diffusion.cpp adds **Day-0 support for Qwen-Image-2.1**
@@ -60,6 +61,7 @@ API and command-line option may change frequently.***
     - [Z-Image L2P](./docs/z_image_l2p.md)
     - [MiniT2I](./docs/minit2i.md)
     - [P-H3-Image](./docs/p_h3_image.md)
+    - [MiniMax-H3-Image](./docs/minimax_h3_image.md)
     - [SenseNova U1.5](./docs/sensenova_u1.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
@@ -201,6 +203,9 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
 - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel)
 - [KoboldCpp](https://github.com/LostRuins/koboldcpp)
 - [Kokon](https://github.com/GabrieleLozupone/kokon)
+- [NodeTool](https://github.com/nodetool-ai/nodetool)
+- [Quo](https://github.com/quo-ai/quo)
+- [Prisme](https://github.com/d3Ita/prisme)
 
 ## Contributors
 

docs/backend.md

@@ -160,10 +160,12 @@ the runner's graph-cut capacity checks.
 Runtime capacity checks also leave 512 MiB of currently free device memory for
 backend scratch buffers and pipelines, including with explicit backend assignments.
 They cap free-memory reports by the device's total memory minus tracked
-resident allocations. Vulkan reports exceeding total memory are rejected because
-its heap-budget subtraction can underflow. Other backends use the cap instead of
-treating such reports as zero free memory. Failed checks log the reported free and
-total memory alongside tracked weight and runtime allocations.
+resident allocations. Vulkan reports exceeding total memory are rejected on
+discrete GPUs because its heap-budget subtraction can underflow. On integrated
+GPUs, a wrapped report triggers reclamation before the tracked estimate can be
+accepted, subject to the byte checks and configured budgets. Other backends use
+the cap instead of treating such reports as zero free memory. Failed checks log
+the reported free and total memory alongside tracked weight and runtime allocations.
 With `--mmap`, device-backed mappings count toward these budgets at their full
 mapped-file size, once per device buffer even when multiple parameter blocks
 share it. Mappings retained in the loader cache continue to count.
@@ -181,8 +183,18 @@ GPU cache space follows the same component priority. Before a lower-priority
 component can become permanently resident, the planner leaves room for the full
 weights and estimated compute space of higher-priority offloaded components.
 If offloaded diffusion already needs the entire main GPU budget, TE and VAE also
-use offloaded parameters. Their GPU copies can then be released after their
-phases, leaving more room to reuse diffusion weights across sampling steps.
+use offloaded parameters. GPU copies remain cached across sampling steps while
+a component's execution is active and are released under memory pressure.
+Ending a component's execution releases its staged GPU copies, disk-backed
+weight buffers, temporary compute buffers and execution caches. Parameters
+assigned permanently to a backend remain resident. Hires model upscaling shares
+the runtime budget and can
+reclaim idle weight copies from other components.
+Reclamation releases idle compute buffers first, then cached weight allocations
+until the current work fits. Weights used by active computation or pending
+transfers stay protected.
+For CPU computation with disk-backed parameters, cached weights are released
+when execution ends so they do not accumulate in RAM across components.
 CPU parameter residency allows GPU weight caching; it does not force every
 weight to be copied again at every step.
 

docs/build.md

@@ -21,6 +21,25 @@ git pull origin master
 git submodule update --init --recursive
 ```
 
+## MSVC compilation parallelism
+
+MSVC builds limit parallel C/C++ compilation within each project by default.
+For `n` logical processors, the limit is `n` when `n <= 2`, `n - 1` when
+`3 <= n <= 8`, and `n - 2` when `n > 8`. The default is calculated when the
+build directory is first configured with this setting. Override it with
+`-DSD_MSVC_MP_COUNT=4`, or pass `-USD_MSVC_MP_COUNT` to recalculate it.
+
+For Visual Studio builds, keep project-level parallelism at one to avoid
+multiplying the per-project limit:
+
+```shell
+cmake --build build --config Release --parallel 1
+```
+
+This setting controls MSVC's `/MP` option, not CUDA compiler parallelism or a
+hard CPU usage limit. Other build tools, such as Ninja, schedule compilation
+jobs separately; use `cmake --build build --parallel N` to limit those jobs.
+
 ## Selecting a GGML source tree
 
 By default, sd.cpp builds the patched GGML submodule in `ggml/`. To build with

docs/ggml_extensions.md

@@ -110,7 +110,7 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions.
 ## Backend services
 
 Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by
-`ggml_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the
+`ggml_ext_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the
 explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization
 node during graph construction. Logical output dimensions are stored in the
 operator attributes, so backend copies of the packed activation do not need to
@@ -159,7 +159,15 @@ rebuild a workspace whose actual allocation is smaller than the estimate.
 ## sd.cpp integration
 
 `src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas,
-validation, graph construction helpers, and backend registration entry point.
+validation, and backend registration entry point. Graph construction helpers live
+in `src/core/ggml_extend.cpp` and `.h`; they take `GGMLRunnerContext*` so backend
+selection, compatibility workarounds, and shared ConvRot nodes use the same
+context as the calling model. Per-call precision, scaling, attention, and padding
+options remain explicit. Model blocks retain weight loading and adapter
+orchestration, while the shared helpers own FP8 fallback, embedding lookup,
+full/partial RoPE, grouped convolution, and backend-dependent repeat handling.
+LoRA/LoKr construction, n-mode products, and Kronecker products stay in
+`src/model/adapter/lora_ops.cpp` and `.h`, with the same runner-context interface.
 CPU kernels and fusion rules live in `src/core/cpu/`, with registration collected
 in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
 respective backend directories.
@@ -207,11 +215,13 @@ HIP, MUSA, and Vulkan. Other backends can use the scheduler's CPU fallback.
 The former `ggml_rope_apply` API, builtin opcode, and CPU implementation have
 been removed from ggml.
 
-`src/core/cuda/rope.cu` supplies standalone CUDA RoPE
-and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
+`src/core/cuda/rope.cu` supplies standalone CUDA RoPE,
+`RMSNorm -> RoPE` fusion for non-affine normalization, and
+`RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
 its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
 position-encoding strides are supported by the standalone kernel; RMSNorm fusion
-requires contiguous position encoding and weights. Copy fusion uses workspace
+requires contiguous position encoding and, when present, weights. The non-affine
+fusion uses 64-channel F32 inputs and split-half RoPE. Copy fusion uses workspace
 when the destination overlaps an input. These registrations do not require
 CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
 
@@ -261,13 +271,13 @@ implementations. The CPU kernels use the
 configured thread pool and scheduler workspace; the old CPU INT8 GEMM and ConvRot
 implementations have been removed from ggml.
 Model code keeps calling `ggml_ext_linear_i8_tensorwise`,
-`ggml_mul_mat_i8_tensorwise`, and `ggml_quantize_i8_convrot`. The latter two
+`ggml_ext_mul_mat_i8_tensorwise`, and `ggml_ext_quantize_i8_convrot`. The latter two
 constructors are declared and implemented in sd.cpp and create custom nodes.
 Their former ggml declarations, implementations, ConvRot opcode, and native
 backend paths have been removed. CPU, CUDA/HIP/MUSA, and Vulkan register kernels
 and fusion rules through the extension interfaces.
 
-For already quantized activations, `ggml_mul_mat_i8_tensorwise` accepts ordinary
+For already quantized activations, `ggml_ext_mul_mat_i8_tensorwise` accepts ordinary
 I8 weight and activation tensors. Its optional final `input_scale` parameter is a
 separate F32 tensor with one value per activation row. Weight scales have one
 value per output channel; weight scales and bias can be null. Existing calls
@@ -276,8 +286,8 @@ order. All scales are explicit graph dependencies; there is no quantized-tensor
 wrapper. Shapes use ggml's low-to-high dimension order. Call
 `sd_ggml_register_extensions` before scheduling these custom nodes.
 
-`ggml_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary
-workspace) or the output of `ggml_quantize_i8_convrot`. ConvRot keeps the existing
+`ggml_ext_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary
+workspace) or the output of `ggml_ext_quantize_i8_convrot`. ConvRot keeps the existing
 packed I8 activation rows, padded to a multiple of four rows, followed by F32 row
 scales. Weight scales remain separate. CPU ConvRot groups are 1, 4, 16, 64, or 256.
 Unsupported layouts and malformed scale shapes are rejected. Separate CPU quantize
@@ -490,8 +500,8 @@ weight scales; unsupported custom nodes fall back to CPU.
 
 ## SageAttention
 
-The existing `ggml_sage_attn` entry point is declared and implemented in sd.cpp's
-common extension files. It constructs `sd.sage_attn.v1` with explicit Q/K/V inputs
+The `ggml_ext_sage_attn` entry point is declared and implemented in sd.cpp's
+graph construction helpers. It constructs `sd.sage_attn.v1` with explicit Q/K/V inputs
 and scale/mode attributes. ggml has no Sage-specific opcode, allocation rule, or
 backend dispatch branch.
 
@@ -509,7 +519,7 @@ layers; a custom Sage node itself does not silently change algorithms.
 
 ## Sol-Attn
 
-The existing `ggml_sol_attn` entry point now lives in the common sd.cpp extension
+The `ggml_ext_sol_attn` entry point lives in the common sd.cpp graph construction
 files and constructs `sd.sol_attn.v1`. It accepts contiguous F32 Q/K/V tensors
 with shape `[128, tokens, heads, batch]`, a positive finite scale, and finite tau.
 The output remains F32 `[128, heads, tokens, batch]`. No current sd.cpp model or

docs/int8_convrot.md

@@ -94,7 +94,7 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 
 - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
 - Metal supports INT8 tensorwise and convrot on Apple silicon (Apple GPU family 7 or newer). It keeps weights in INT8 and rotates and quantizes activations on the GPU. On macOS/iOS 26 or newer, it probes Metal TensorOps `matmul2d` with INT8 inputs and INT32 cooperative results. It combines 1024-product partial sums using wrapping integer addition to preserve exact accumulation on older GPUs, then applies scales and bias in registers. Older systems or unavailable TensorOps pipelines use SIMD-group matrix multiplication: INT8 tiles are converted to FP16, and FP32 partial sums cover at most 256 products before being combined as integers, preserving the I32 accumulation contract. Convrot groups 1, 4, 16, 64, and 256 are supported. These kernels require the patched GGML Metal extension services and are enabled by `SD_METAL=ON`.
-- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
+- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled. For linear layers, it restores scales and adds any aligned bias before writing the result; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
 - AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication.
 - MUSA shares the CUDA-style H256 activation kernels. INT8 GEMM falls back to CPU because the MUSA backend does not advertise INT8 BLAS support.
 - Vulkan provides H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU.

docs/minimax_h3.md

@@ -99,6 +99,24 @@ The C API exposes the same inputs through `ref_images`, `ref_videos`, and
 `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
 frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
 
+## Sol-H3-VAE decoding
+
+Pass a [Sol-H3-VAE checkpoint](https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae)
+with `--vae` to use the distilled video decoder. The ComfyUI FP16 and INT8
+ConvRot checkpoints include the MiniMax-H3 encoder for image and video inputs.
+Supply `--audio-vae` as usual to decode the soundtrack.
+
+Sol decodes the whole clip by default. Use `--vae-tiling` for spatial tiling or
+`--temporal-tiling` for temporal tiling. Either can reduce memory use but may
+affect output quality because the decoder has less context at tile boundaries.
+
+Temporal tiles default to 10 latent frames with 5 latent frames of overlap.
+Set `temporal_tile_frames` and `temporal_tile_overlap` in
+`--extra-tiling-args` to override them. Both must be multiples of 5 latent
+frames, and overlap must be smaller than the tile. The last tile may be shorter.
+Circular decoding requires spatial tiling on each circular axis, with a tile
+size smaller than the input on that axis.
+
 ## Latent hires second pass (two-stage sampling)
 
 MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers:

docs/minimax_h3_image.md

@@ -0,0 +1,65 @@
+# MiniMax-H3-Image
+
+MiniMax-H3-Image supports text-to-image generation and reference-image editing
+with the same checkpoint. It reuses the MiniMax-H3 text encoder and uses a
+dedicated image VAE.
+
+## Model files
+
+Pass the three MiniMax-H3-Image components separately:
+
+- `--diffusion-model`: MiniMax-H3-Image diffusion transformer
+- `--vae`: MiniMax-H3-Image VAE
+- `--llm`: the MiniMax-H3 Qwen3-VL-32B text encoder checkpoint
+
+Use the MiniMax-H3 text encoder, including its Qwen3-VL vision tower and
+DeepStack mergers. An existing MiniMax-H3 text encoder can be reused. If the
+vision tower is stored separately, pass it with `--llm_vision`.
+
+Use a single-file diffusion checkpoint with the six-step PDD output heads
+included, such as `minimax_h3_image_pruned_bf16.safetensors` or its INT8 ConvRot
+version. Use the MiniMax-H3-Image VAE; the MiniMax-H3 video VAE and P-H3-Image
+VAE are incompatible. No audio VAE or P-H3 text connector is needed.
+
+### Download weights
+
+- Download minimax_h3_image_pruned
+    - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/diffusion_models
+- Download qwen3vl_32b_minimax_h3
+    - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/text_encoders
+    - gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/vae
+
+## Text-to-image
+
+```sh
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\minimax_h3_image_pruned_bf16.safetensors --vae ..\models\vae\minimax_h3_image_vae_fp16.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "A silver tabby kitten holding a sign reading 'sd.cpp', natural daylight" --cfg-scale 1.0 --steps 6 --sampling-method euler --sigmas minimax-h3-image -v -W 1024 -H 1024 --offload-to-cpu --rng cpu --diffusion-fa --vae-tiling
+```
+
+Use Euler, CFG 1, and `--sigmas minimax-h3-image` for the model's six-step
+schedule. The default discrete scheduler also produces this schedule with
+`--steps 6`. Prompt expansion is not included.
+
+## Reference-image editing
+
+Add `--ref-image` (or `-r`) for reference-image editing:
+
+```sh
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\minimax_h3_image_pruned_bf16.safetensors --vae ..\models\vae\minimax_h3_image_vae_fp16.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "Edit <Picture 1>. Put a red scarf on the cat, preserving its face and the background." -r ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 --steps 6 --sampling-method euler --sigmas minimax-h3-image -v -W 1024 -H 1024 --offload-to-cpu --rng cpu --diffusion-fa --vae-tiling
+```
+
+`--ref-image` can be repeated. Reference images are numbered in argument order
+as `<Picture 1>`, `<Picture 2>`, and so on, with up to 16 references supported
+by the model. Prompts are passed through without an automatic editing prefix.
+
+Reference images keep their aspect ratio. The text encoder sees a 1024-pixel
+short edge, while VAE inputs use a short edge between 1024 and 2048 pixels.
+Small references are enlarged and dimensions are rounded to multiples of 32.
+
+## Shape and runtime notes
+
+- Width and height are aligned upward to a multiple of 32.
+- Start at 1024 × 1024. Larger canvases and more references increase memory use.
+- Use quantized weights, `--offload-to-cpu`, and `--vae-tiling` to reduce GPU
+  memory requirements.

docs/performance.md

@@ -92,10 +92,12 @@ dimensions down to the minimum size. Each spatial retry must reduce the effectiv
 tile size. These runtime adjustments do not change the caller's parameters.
 Execution failures are not retried, and encoding has no automatic OOM retry.
 
-`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses
-spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal
-windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit
-spatial sizes are honored.
+`--temporal-tiling` remains independent of spatial tiling. The original MiniMax-H3
+video VAE uses spatial tiling (256x256 pixels and 25% overlap by default) and its own
+temporal windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap` and
+explicit spatial sizes are honored. [Sol-H3-VAE](minimax_h3.md#sol-h3-vae-decoding)
+decodes the whole clip by default and supports opt-in spatial and block-aligned
+temporal tiling. Tiling may affect output quality.
 
 ## Offload weights to the CPU to save VRAM without reducing generation speed.
 

examples/common/common.cpp

@@ -1540,6 +1540,8 @@ ArgOptions SDGenerationParams::get_options() {
             sigmas_str = "1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0";
         } else if (sigmas_str == "p-h3-image") {
             sigmas_str = "1.0,0.984753,0.967905,0.951688,0.925372,0.881923,0.810083,0.528718,0.0";
+        } else if (sigmas_str == "minimax-h3-image") {
+            sigmas_str = "1.0,0.9675236940383911,0.8875485062599182,0.75,0.5327739119529724,0.23200887441635132,0.0";
         } else if (sigmas_str == "hyperflow-h3-flash-4") {
             sigmas_str = "1.0,0.962575853,0.755781889,0.312497854,0.0";
         } else if (sigmas_str == "hyperflow-h3-flash-6") {
@@ -1793,11 +1795,11 @@ ArgOptions SDGenerationParams::get_options() {
          on_scheduler_arg},
         {"",
          "--sigmas",
-         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
+         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, minimax-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
          on_sigmas_arg},
         {"",
          "--hires-sigmas",
-         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
+         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, minimax-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
          on_hires_sigmas_arg},
         {"",
          "--skip-layers",

examples/server/README.md

@@ -287,6 +287,16 @@ If you want to use a different host or port, pass:
 --listen-ip <ip> --listen-port <port>
 ```
 
+The native async job queue can be tuned with:
+
+```bash
+--max-pending-jobs <count> --completed-ttl-seconds <seconds>
+```
+
+`--max-pending-jobs` defaults to 64 and limits how many async jobs may be
+queued or generating. `--completed-ttl-seconds` defaults to 600 and controls
+how long completed jobs remain available through the async API.
+
 # Frontend
 
 ## Build with Frontend

examples/server/main.cpp

@@ -113,6 +113,10 @@ int main(int argc, const char** argv) {
     std::vector<UpscalerEntry> upscaler_cache;
     std::mutex upscaler_mutex;
     AsyncJobManager async_job_manager;
+    async_job_manager.max_pending_jobs =
+        static_cast<size_t>(svr_params.max_pending_jobs);
+    async_job_manager.completed_ttl_seconds =
+        static_cast<int64_t>(svr_params.completed_ttl_seconds);
     ServerRuntime runtime = {
         models.context.get(),
         &sd_ctx_mutex,

examples/server/recipes/defaults.md

@@ -70,6 +70,7 @@ For community edit adapters and other optional extensions, these are the base-mo
 | `ming-image` | 1 | 12 | 1024 × 1024 | [Source](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design)  |
 | `minimax-h3-fl2va` | 1 | — | 1344 × 768 | [Source](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3) Guidance-distilled; official 768p canvas. No fixed step count is prescribed here; the upstream API counts sigma grid points rather than sd.cpp denoising steps. |
 | `minimax-h3-ref2va` | 1 | — | 1344 × 768 | [Source](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3) Guidance-distilled; official 768p canvas. No fixed step count is prescribed here; the upstream API counts sigma grid points rather than sd.cpp denoising steps. |
+| `minimax-h3-image` | 1 | 6 | 1024 × 1024 | [Runtime guide](../../../docs/minimax_h3_image.md) Dedicated image VAE and fixed PDD sigma schedule. |
 | `minit2i` | 2.5 | 20 | 512 × 512 | [Source](https://huggingface.co/MiniT2I/MiniT2I/raw/main/pipeline.py) B/16 example and native model size; L/16 uses CFG 6. |
 | `moss-image-base` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base)  |
 | `moss-image-sft` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)  |

examples/server/recipes/minimax-h3-image.json

@@ -0,0 +1,48 @@
+{
+  "schema_version": 1,
+  "id": "minimax-h3-image",
+  "name": "MiniMax H3 Image",
+  "size_multiple": 32,
+  "description": "Text-to-image and reference-image editing with the six-step PDD checkpoint. Requires the MiniMax-H3-Image 2D VAE and MiniMax-H3 Qwen3-VL-32B text encoder with its vision tower.",
+  "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-H3-Image",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder"
+    },
+    {
+      "id": "llm_vision",
+      "required": false,
+      "label": "Vision tower (required when not included in the text encoder)"
+    },
+    {
+      "id": "vae",
+      "required": true,
+      "label": "MiniMax-H3-Image 2D VAE"
+    }
+  ],
+  "defaults": {
+    "sampling_method": "euler",
+    "cfg_scale": 1,
+    "steps": 6,
+    "width": 1024,
+    "height": 1024,
+    "custom_sigmas": [
+      1.0,
+      0.9675236940383911,
+      0.8875485062599182,
+      0.75,
+      0.5327739119529724,
+      0.23200887441635132,
+      0.0
+    ]
+  },
+  "features": [
+    "ref_images"
+  ]
+}

examples/server/runtime.cpp

@@ -351,6 +351,12 @@ ArgOptions SDSvrParams::get_options() {
 
     options.int_options = {
         {"", "--listen-port", "server listen port (default: 1234)", &listen_port},
+        {"", "--max-pending-jobs",
+         "maximum number of pending async jobs (default: 64)",
+         &max_pending_jobs},
+        {"", "--completed-ttl-seconds",
+         "retention time for completed async jobs in seconds (default: 600)",
+         &completed_ttl_seconds},
     };
 
     options.bool_options = {
@@ -388,6 +394,16 @@ bool SDSvrParams::validate() {
         return false;
     }
 
+    if (max_pending_jobs <= 0) {
+        LOG_ERROR("error: max_pending_jobs must be greater than 0");
+        return false;
+    }
+
+    if (completed_ttl_seconds < 0) {
+        LOG_ERROR("error: completed_ttl_seconds must be non-negative");
+        return false;
+    }
+
     if (!serve_html_path.empty() && !fs::exists(fs::u8path(serve_html_path))) {
         LOG_ERROR("error: serve_html_path file does not exist: %s", serve_html_path.c_str());
         return false;
@@ -408,6 +424,8 @@ std::string SDSvrParams::to_string() const {
         << "  log_level: " << log_level_name(log_level) << ",\n"
         << "  listen_ip: " << listen_ip << ",\n"
         << "  listen_port: \"" << listen_port << "\",\n"
+        << "  max_pending_jobs: " << max_pending_jobs << ",\n"
+        << "  completed_ttl_seconds: " << completed_ttl_seconds << ",\n"
         << "  serve_html_path: \"" << serve_html_path << "\",\n"
         << "}";
     return oss.str();

examples/server/runtime.h

@@ -20,8 +20,10 @@ struct AsyncJobManager;
 struct ServerModels;
 
 struct SDSvrParams {
-    std::string listen_ip = "127.0.0.1";
-    int listen_port       = 1234;
+    std::string listen_ip     = "127.0.0.1";
+    int listen_port           = 1234;
+    int max_pending_jobs      = 64;
+    int completed_ttl_seconds = 600;
     std::string serve_html_path;
     std::vector<std::string> model_dirs;
     bool normal_exit         = false;

src/conditioning/conditioner.hpp

@@ -2589,7 +2589,9 @@ struct LLMEmbedder : public Conditioner {
                 } else if (conditioner_params.ref_images != nullptr) {
                     for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
                         auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
-                        prompt += "<Picture " + std::to_string(i + 1) + ">: ";
+                        prompt += version == VERSION_MINIMAX_H3_IMAGE
+                                      ? "Picture " + std::to_string(i + 1) + ": "
+                                      : "<Picture " + std::to_string(i + 1) + ">: ";
                         if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
                                                                           resized,
                                                                           false),
@@ -3448,8 +3450,8 @@ struct LTXAVTextProjection : public GGMLBlock {
 
         auto video_projection = std::dynamic_pointer_cast<Linear>(blocks["video_aggregate_embed"]);
         auto audio_projection = std::dynamic_pointer_cast<Linear>(blocks["audio_aggregate_embed"]);
-        auto video_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
-        auto audio_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
+        auto video_in         = ggml_ext_scale(ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
+        auto audio_in         = ggml_ext_scale(ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
         auto video            = video_projection->forward(ctx, video_in);
         auto audio            = audio_projection->forward(ctx, audio_in);
         return ggml_concat(ctx->ggml_ctx, video, audio, 0);

src/convert.cpp

@@ -463,11 +463,15 @@ static bool load_model_into_memory(ModelLoader& model_loader,
         const ggml_type export_type = get_export_tensor_type(model_loader, storage, output_type, type_rules);
         const bool quant_src        = ggml_is_quantized(storage.type);
         const bool quant_dst        = ggml_is_quantized(export_type);
+        const bool half_src         = storage.type == GGML_TYPE_F16 || storage.type == GGML_TYPE_BF16;
         if (quant_src && quant_dst) {
             load_type = GGML_TYPE_F16;
-        } else if (quant_src || quant_dst) {
-            load_type = quant_src ? export_type : storage.type;
-        } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32) {
+        } else if (quant_src) {
+            load_type = export_type;
+        } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32 ||
+                   (half_src && export_type != storage.type)) {
+            // LoRA deltas can be smaller than half an F16/BF16 ulp of the weight. Merge in F32
+            // when the export converts the tensor again, so the merged weight is rounded once.
             load_type = GGML_TYPE_F32;
         }
         ggml_tensor* tensor = ggml_new_tensor(ctx, load_type, storage.n_dims, storage.ne);

src/core/cuda/i8_gemm.cu

@@ -176,9 +176,22 @@ namespace sd_i8_cutlass {
             dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)];
     }
 
+    template <class Fragment>
+    struct AddBias {
+        CUTLASS_DEVICE Fragment operator()(const Fragment& value, const Fragment& bias) const {
+            Fragment result;
+            // Prevent contraction with the preceding scale multiplication.
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < Fragment::kElements; ++i)
+                result[i] = __fadd_rn(value[i], bias[i]);
+            return result;
+        }
+    };
+
     template <class Config, bool StreamK, bool PackedW4 = false>
     struct DenseTypes {
         static constexpr bool kSplit = false;
+        static constexpr bool kBias  = !PackedW4;
         using Weight                 = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
         using Tile                   = typename Config::Tile;
         using Warp                   = typename Config::Warp;
@@ -190,7 +203,9 @@ namespace sd_i8_cutlass {
         using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
         using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
         using Store                  = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
-        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Add                    = cutlass::epilogue::threadblock::VisitorCompute<AddBias, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Biased                 = cutlass::epilogue::threadblock::Sm80EVT<Add, ScaleW, Row>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, std::conditional_t<kBias, Biased, ScaleW>>;
         using Kernel                 = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
             int8_t,
             cutlass::layout::RowMajor,
@@ -219,6 +234,7 @@ namespace sd_i8_cutlass {
     template <class Config, bool StreamK, bool Split = false>
     struct SwigluTypes {
         static constexpr bool kSplit = Split;
+        static constexpr bool kBias  = false;
         using Weight                 = int8_t;
         using Tile                   = typename Config::Tile;
         using Warp                   = typename Config::Warp;
@@ -259,24 +275,30 @@ namespace sd_i8_cutlass {
         bool packed_w4;
         bool split;
         size_t kernel_bytes = 0, scratch_offset = 0;
-        int occupancy                                                                                                                                    = 0;
-        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr;
+        int occupancy                                                                                                                                                  = 0;
+        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*, const float*) = nullptr;
     };
     template <class Types>
-    static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
-        typename Types::ScaleW::Arguments scales     = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
-        typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+    static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, const float* bias) {
+        typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+        typename Types::Epilogue::Arguments epilogue;
+        if constexpr (Types::kBias) {
+            // A negative zero default preserves signed-zero outputs.
+            epilogue = {{scales, {bias, -0.f, {}}, {}}, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+        } else {
+            epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+        }
         return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast<const typename Types::Weight*>(w), nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
     }
     template <class Types>
-    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales) {
+    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales, const float* bias) {
         using Kernel = typename Types::Kernel;
         if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
-             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales)) %
+             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales) | reinterpret_cast<uintptr_t>(bias)) %
             16)
             return GGML_STATUS_FAILED;
         auto* paired = static_cast<float*>(scratch);
-        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
+        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out, bias);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return GGML_STATUS_FAILED;
         if (!p.occupancy) {
@@ -312,7 +334,7 @@ namespace sd_i8_cutlass {
         using Kernel = typename Types::Kernel;
         if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
             return false;
-        auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
+        auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return false;
         size_t bytes = 0;
@@ -394,11 +416,13 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64
         return nullptr;
     return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
 }
-ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
-    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr);
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const float* bias) {
+    if (bias && (plan->swiglu || plan->packed_w4))
+        return GGML_STATUS_FAILED;
+    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr, bias);
 }
 ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) {
-    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales);
+    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales, nullptr);
 }
 void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
     delete plan;

src/core/cuda/i8_gemm.h

@@ -4,7 +4,7 @@
 #include "ggml-ext.h"
 struct SDI8CudaGemmPlan;
 SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false);
-ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const float* bias = nullptr);
 // n is the width of each branch; gate/up weights and scales remain separate.
 SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace);
 ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream);

src/core/cuda/int8.cu

@@ -447,7 +447,7 @@ namespace SD_CUDA_NAMESPACE {
         }
 #ifdef SD_CUDA_I8_CUTLASS
         static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) {
-            return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 &&
+            return op.weight_scale && op.weight->ne[0] <= INT_MAX / 16384 && aligned(op.bias, 16) &&
                    aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
                    (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16)));
         }
@@ -531,7 +531,7 @@ namespace SD_CUDA_NAMESPACE {
 #ifdef SD_CUDA_I8_CUTLASS
                 size_t gemm_bytes = 0;
                 // Retain cuBLAS when I32 accumulation can overflow.
-                if (cutlass_compatible(op, dst) &&
+                if (cutlass_compatible(op, dst) && (!op.bias || p->kind == MAT) &&
                     (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
                     if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
                         p->gemm        = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
@@ -819,7 +819,7 @@ namespace SD_CUDA_NAMESPACE {
 #ifdef SD_CUDA_I8_CUTLASS
                 if (p.gemm) {
                     auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
-                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream);
+                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream, bias);
                     if (status != GGML_STATUS_SUCCESS)
                         return status;
                 } else

src/core/cuda/rope.cu

@@ -12,7 +12,7 @@ namespace SD_CUDA_NAMESPACE {
             v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
         return v;
     }
-    template <int D, bool interleaved, typename T>
+    template <int D, bool interleaved, typename T, bool affine = true>
     static __global__ void rms_norm_mul_rope_apply_fused(
         const float* x,
         const float* weight,
@@ -51,7 +51,7 @@ namespace SD_CUDA_NAMESPACE {
 #pragma unroll
         for (int j = 0; j < D / 32; ++j) {
             const int col     = j * 32 + lane;
-            const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
+            const float value = affine ? __fmul_rn(__fmul_rn(scale, v[j]), weight[col]) : __fmul_rn(scale, v[j]);
             float rotated     = value;
             if constexpr (interleaved) {
                 const float other = cuda_shuffle_xor(value, 1, 32);
@@ -65,7 +65,7 @@ namespace SD_CUDA_NAMESPACE {
             } else {
                 if (col < 2 * rot_half) {
                     const int partner   = col < rot_half ? col + rot_half : col - rot_half;
-                    const float other   = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
+                    const float other   = affine ? __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]) : __fmul_rn(scale, x_row[partner]);
                     const float* m      = pe + 4 * (token * rot_half + col % rot_half);
                     const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
                     rotated             = col < rot_half
@@ -178,6 +178,20 @@ namespace SD_CUDA_NAMESPACE {
                (ggml_nrows(x) + 3) / 4 <= INT_MAX;
     }
 
+    static bool rope_non_affine_fusion_supports(const ggml_tensor* norm, const ggml_tensor* rope) {
+        if (!norm || norm->op != GGML_OP_RMS_NORM || !rope_is_node(rope) || rope->src[0] != norm ||
+            !rope_supports(nullptr, rope, nullptr) || rope_interleaved(rope))
+            return false;
+        const auto* x  = norm->src[0];
+        const auto* pe = rope->src[1];
+        // The kernel reads contiguous channels and position matrices, with row strides in floats.
+        return x && x->type == GGML_TYPE_F32 && x->ne[0] == 64 &&
+               x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 &&
+               x->nb[2] % sizeof(float) == 0 && x->nb[3] % sizeof(float) == 0 &&
+               ggml_are_same_shape(x, norm) && ggml_is_contiguous(pe) &&
+               ggml_nrows(x) > 0 && (ggml_nrows(x) + 3) / 4 <= INT_MAX;
+    }
+
     template <typename T, bool Interleaved>
     static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
         const auto* x      = nodes[0]->src[0];
@@ -220,6 +234,9 @@ namespace SD_CUDA_NAMESPACE {
         if (disabled)
             return 0;
         auto* g = const_cast<ggml_cgraph*>(graph);
+        if (first + 1 < ggml_graph_n_nodes(g) &&
+            rope_non_affine_fusion_supports(ggml_graph_node(g, first), ggml_graph_node(g, first + 1)))
+            return 2;
         if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
                                                       ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
             return 0;
@@ -253,6 +270,16 @@ namespace SD_CUDA_NAMESPACE {
         for (int i = 2; i + 1 < total; ++i) {
             const auto* mul  = ggml_graph_node(graph, i);
             const auto* rope = ggml_graph_node(graph, i + 1);
+            if (rope_non_affine_fusion_supports(rope->src[0], rope)) {
+                for (int j = i; j >= 0 && j >= i - 64; --j) {
+                    if (ggml_graph_node(graph, j) == rope->src[0]) {
+                        if (j < i)
+                            ggml_ext_graph_delay(graph, j, 1, i + 1);
+                        break;
+                    }
+                }
+                continue;
+            }
             if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
                 !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
                 continue;
@@ -296,13 +323,20 @@ namespace SD_CUDA_NAMESPACE {
         if (!cuda_context(execution->backend, &context))
             return GGML_STATUS_FAILED;
         auto stream      = static_cast<CUDAStream>(context.stream);
-        const auto* rope = nodes[count == 1 ? 0 : 2];
+        const auto* rope = nodes[count <= 2 ? count - 1 : 2];
         const auto* x    = count == 1 ? rope->src[0] : nodes[0]->src[0];
         auto* dst        = nodes[count - 1];
-        const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
+        const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 2 && overlaps(dst, nodes[1]->src[1]));
         void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
         CUDAError error;
-        if (count > 1) {
+        if (count == 2) {
+            const auto* pe     = rope->src[1];
+            const int64_t rows = ggml_nrows(x);
+            rms_norm_mul_rope_apply_fused<64, false, float, false><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>(
+                static_cast<const float*>(x->data), nullptr, static_cast<const float*>(pe->data), static_cast<float*>(output),
+                x->ne[1], x->ne[2], rows, x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), 1.f);
+            error = cuda_last_error();
+        } else if (count > 2) {
             error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
         } else {
             RopeApplyStrides strides;

src/core/execution_engine.cpp

@@ -0,0 +1,992 @@
+#include "core/execution_engine.h"
+
+#include <algorithm>
+#include <chrono>
+#include <exception>
+#include <thread>
+#include <unordered_set>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/layer_split_partition.h"
+#include "core/segment_graph_bindings.h"
+#include "core/segment_weight_pipeline.h"
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+
+std::shared_ptr<sd::ExecutionEngine> DeviceResidencyManager::execution_engine() {
+    if (execution_engine_ == nullptr) {
+        execution_engine_ = std::make_shared<sd::ExecutionEngine>(this);
+    }
+    return execution_engine_;
+}
+
+namespace sd {
+    class GraphExecution {
+        ExecutionEngine& engine_;
+        ExecutionSession& session_;
+        const ComputeTask& task_;
+        std::shared_ptr<DeviceResidencyManager> manager_;
+        std::unordered_set<const ggml_tensor*> params_tensor_set_;
+        ggml_status last_compute_status_ = GGML_STATUS_FAILED;
+        Tensor<float> output_;
+
+        using GraphCutPlan = ggml_graph_cut::Plan;
+        ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+        std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* graph);
+        ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
+        // Weightless ops follow the most recent weight device. Views must stay
+        // unpinned so the scheduler preserves cross-device copies for consumers.
+        void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* graph, ggml_cgraph* original_graph = nullptr);
+        Tensor<float> read_graph_tensor(ggml_tensor* tensor, const char* label);
+        void copy_data_to_backend_tensor(ggml_cgraph* graph);
+        const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* graph);
+        bool assign_graph_cut_layer_split_backends(ggml_cgraph* graph);
+        ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
+        std::vector<DeviceMemoryRequest> memory_requests(const std::vector<BackendBufferSize>& sizes, size_t pending_cache_bytes) const;
+        bool fits(const std::vector<DeviceMemoryRequest>& requests, const std::vector<ggml_tensor*>& params) const;
+        bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
+        bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
+        bool is_multi_device() const { return !session_.extra_backends_.empty(); }
+        size_t retained_runtime_buffer_bytes(ggml_backend_t backend) const { return session_.resident_bytes(backend); }
+        void sync_runtime_residency() { engine_.sync_residency(session_); }
+
+    public:
+        GraphExecution(ExecutionEngine& engine, ExecutionSession& session, const ComputeTask& task)
+            : engine_(engine), session_(session), task_(task), manager_(session.residency_manager_.lock()) {
+            params_tensor_set_.insert(task.weights.begin(), task.weights.end());
+            if (manager_ != nullptr) {
+                for (int i = 0; i < ggml_graph_cut::leaf_count(task.graph); ++i) {
+                    auto* parameter = manager_->resolve_param_tensor(ggml_graph_cut::leaf_tensor(task.graph, i));
+                    if (parameter != nullptr) {
+                        params_tensor_set_.insert(parameter);
+                    }
+                }
+                for (auto* tensor : params_tensor_set_) {
+                    auto* parameter = manager_->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
+                    if (parameter != nullptr) {
+                        session_.params_.insert(parameter);
+                    }
+                }
+            }
+        }
+
+        ComputeResult run() {
+            if (!execute_graph(task_.graph, task_.n_threads, task_.no_return, task_.read_outputs)) {
+                return {last_compute_status_, {}};
+            }
+            return {GGML_STATUS_SUCCESS, std::move(output_)};
+        }
+    };
+
+    static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
+        // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
+        if (sd_backend_is(backend, "Vulkan")) {
+            return true;
+        }
+        auto buft        = ggml_backend_get_default_buffer_type(backend);
+        size_t alignment = ggml_backend_buft_get_alignment(buft);
+        size_t max_size  = ggml_backend_buft_get_max_size(buft);
+        max_size -= max_size % alignment;
+        if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
+            return false;
+        }
+        // Cross-backend copies must also fit the destination's allocation limit.
+        for (auto src : node->src) {
+            if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static size_t add_bytes(size_t a, size_t b) {
+        return b > SIZE_MAX - a ? SIZE_MAX : a + b;
+    }
+
+    ggml_tensor* GraphExecution::canonical_param_tensor(ggml_tensor* tensor) {
+        for (auto* current = tensor; current != nullptr; current = current->view_src) {
+            if (params_tensor_set_.count(current) != 0)
+                return current;
+        }
+        return nullptr;
+    }
+
+    std::vector<ggml_tensor*> GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) {
+        std::vector<ggml_tensor*> used_params;
+        if (gf == nullptr || params_tensor_set_.empty()) {
+            return used_params;
+        }
+
+        std::unordered_set<const ggml_tensor*> seen_params;
+        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+        seen_params.reserve(static_cast<size_t>(n_leafs));
+        for (int i = 0; i < n_leafs; ++i) {
+            ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
+            ggml_tensor* param_leaf = canonical_param_tensor(leaf);
+            if (param_leaf != nullptr &&
+                seen_params.insert(param_leaf).second) {
+                used_params.push_back(param_leaf);
+            }
+        }
+        return used_params;
+    }
+
+    ggml_backend_t GraphExecution::backend_for_weight(const ggml_tensor* tensor) const {
+        if (tensor == nullptr || tensor->buffer == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+            ggml_backend_buffer_is_host(tensor->buffer)) {
+            return nullptr;
+        }
+        ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+        if (dev == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_get_device(session_.backend_) == dev) {
+            return session_.backend_;
+        }
+        for (ggml_backend_t backend : session_.extra_backends_) {
+            if (ggml_backend_get_device(backend) == dev) {
+                return backend;
+            }
+        }
+        return nullptr;
+    }
+
+    void GraphExecution::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
+        if (sched == nullptr || gf == nullptr) {
+            return;
+        }
+        ggml_backend_t current = session_.backend_;
+        const int n_nodes      = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; i++) {
+            ggml_tensor* node    = ggml_graph_node(gf, i);
+            auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+            if (node_assignment != session_.node_assignments_.end()) {
+                current = node_assignment->second;
+            }
+            for (int s = 0; s < GGML_MAX_SRC; s++) {
+                ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+                if (weight_backend != nullptr) {
+                    if (node_assignment == session_.node_assignments_.end()) {
+                        current = weight_backend;
+                    }
+                }
+            }
+            if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+                node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+                continue;
+            }
+            if (!backend_can_allocate_op(current, node)) {
+                auto cpu = session_.workspace_.cpu_backend();
+                if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
+                    ggml_backend_sched_set_tensor_backend(sched, node, cpu);
+                }
+            } else if (ggml_backend_supports_op(current, node)) {
+                ggml_backend_sched_set_tensor_backend(sched, node, current);
+            }
+        }
+    }
+
+    Tensor<float> GraphExecution::read_graph_tensor(ggml_tensor* tensor, const char* label) {
+        if (tensor == nullptr) {
+            LOG_ERROR("%s %s tensor is null", task_.description.c_str(), label);
+            return {};
+        }
+        if (tensor->type != GGML_TYPE_F32) {
+            LOG_ERROR("%s %s tensor type mismatch: got %s",
+                      task_.description.c_str(),
+                      label,
+                      ggml_type_name(tensor->type));
+            return {};
+        }
+        ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+        if (buf == nullptr) {
+            LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                      task_.description.c_str(),
+                      label,
+                      tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                      ggml_op_name(tensor->op),
+                      tensor->buffer,
+                      tensor->view_src,
+                      tensor->view_src ? tensor->view_src->buffer : nullptr,
+                      tensor->data);
+            return {};
+        }
+
+        return sd::make_sd_tensor_from_ggml<float>(tensor);
+    }
+
+    void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) {
+        GGML_ASSERT(gf != nullptr);
+        std::unordered_set<const ggml_tensor*> graph_tensor_set;
+        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
+        for (int i = 0; i < n_leafs; ++i) {
+            graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+        }
+        for (int i = 0; i < n_nodes; ++i) {
+            graph_tensor_set.insert(ggml_graph_node(gf, i));
+        }
+
+        for (auto& kv : task_.inputs) {
+            auto tensor = kv.first;
+            auto data   = kv.second;
+            if (tensor == nullptr || data == nullptr) {
+                continue;
+            }
+            const char* name = ggml_get_name(tensor);
+            if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+                continue;
+            }
+            if (tensor->buffer == nullptr) {
+                LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
+                         task_.description.c_str(),
+                         name != nullptr ? name : "",
+                         (long long)tensor->ne[0],
+                         (long long)tensor->ne[1],
+                         (long long)tensor->ne[2],
+                         (long long)tensor->ne[3],
+                         ggml_type_name(tensor->type));
+                continue;
+            }
+
+            ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+            if (buf == nullptr) {
+                LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
+                         task_.description.c_str(),
+                         tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                         tensor ? ggml_op_name(tensor->op) : "<null>",
+                         data,
+                         tensor ? tensor->view_src : nullptr,
+                         (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
+                continue;
+            }
+
+            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+        }
+    }
+
+    const GraphExecution::GraphCutPlan& GraphExecution::resolve_graph_cut_plan(ggml_cgraph* gf) {
+        GGML_ASSERT(gf != nullptr);
+        return sd::ggml_graph_cut::resolve_plan(session_.backend_,
+                                                gf,
+                                                &session_.plan_cache_,
+                                                params_tensor_set_,
+                                                task_.description.c_str());
+    }
+
+    bool GraphExecution::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+        session_.node_assignments_.clear();
+        if (!session_.layer_split_enabled_) {
+            return true;
+        }
+        if (!is_multi_device()) {
+            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", task_.description.c_str());
+            return false;
+        }
+
+        const auto& plan = resolve_graph_cut_plan(gf);
+        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+            auto manager = manager_;
+            if (manager == nullptr) {
+                LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str());
+                return false;
+            }
+            std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+            if (!graph_params.empty() &&
+                !manager->assign_compute_backend(graph_params, session_.backend_)) {
+                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                          task_.description.c_str(),
+                          sd::layer_split_backend_device_display_name(session_.backend_).c_str());
+                return false;
+            }
+            for (ggml_tensor* param : graph_params) {
+                if (param != nullptr) {
+                    session_.weight_assignments_[param] = session_.backend_;
+                }
+            }
+            const int n_nodes = ggml_graph_n_nodes(gf);
+            for (int i = 0; i < n_nodes; i++) {
+                ggml_tensor* node = ggml_graph_node(gf, i);
+                if (node != nullptr) {
+                    session_.node_assignments_[node] = session_.backend_;
+                }
+            }
+            if (!session_.primary_notice_logged_) {
+                LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                         task_.description.c_str(),
+                         sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
+                         graph_params.size());
+                session_.primary_notice_logged_ = true;
+            } else {
+                LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                            task_.description.c_str(),
+                            sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
+                            graph_params.size());
+            }
+            return true;
+        }
+
+        std::vector<ggml_backend_t> split_backends;
+        split_backends.reserve(session_.extra_backends_.size() + 1);
+        split_backends.push_back(session_.backend_);
+        for (ggml_backend_t backend : session_.extra_backends_) {
+            if (backend != nullptr) {
+                split_backends.push_back(backend);
+            }
+        }
+
+        auto manager = manager_;
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str());
+            return false;
+        }
+
+        sd::GraphCutLayerSplitAssignment assignment;
+        auto canonicalize_param = [this](ggml_tensor* tensor) {
+            return canonical_param_tensor(tensor);
+        };
+        if (!sd::partition_graph_cut_layer_split(task_.description.c_str(),
+                                                 gf,
+                                                 plan,
+                                                 split_backends,
+                                                 session_.backend_limits_,
+                                                 session_.max_backend_bytes_,
+                                                 session_.weight_assignments_,
+                                                 canonicalize_param,
+                                                 &assignment)) {
+            return false;
+        }
+
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+                LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                          task_.description.c_str(),
+                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+                return false;
+            }
+        }
+
+        session_.node_assignments_ = std::move(assignment.node_assignments);
+        sd::log_graph_cut_layer_split_assignment(task_.description.c_str(), split_backends, assignment);
+
+        return true;
+    }
+
+    ComputeWorkspace::Measurement GraphExecution::measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs) {
+        auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
+            if (external_outputs != nullptr && external_outputs->count(tensor)) {
+                return session_.backend_;
+            }
+            if (!params_tensor_set_.count(tensor)) {
+                return nullptr;
+            }
+            auto placement = session_.weight_assignments_.find(tensor);
+            return placement == session_.weight_assignments_.end() ? session_.backend_ : placement->second;
+        };
+        auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
+            pin_multi_device_nodes(scheduler, copy, graph);
+        };
+        return session_.workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
+    }
+
+    std::vector<DeviceMemoryRequest> GraphExecution::memory_requests(
+        const std::vector<BackendBufferSize>& sizes,
+        size_t pending_cache_bytes) const {
+        std::vector<DeviceMemoryRequest> requests;
+        for (const auto& size : sizes) {
+            const size_t retained    = retained_runtime_buffer_bytes(size.backend);
+            const size_t reusable    = session_.workspace_.bytes(size.backend);
+            const size_t cache_bytes = size.backend == session_.backend_ ? pending_cache_bytes : 0;
+            const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
+            size_t limit             = session_.max_backend_bytes_;
+            if (is_multi_device()) {
+                size_t index = 0;
+                if (size.backend != session_.backend_) {
+                    auto position = std::find(session_.extra_backends_.begin(), session_.extra_backends_.end(), size.backend);
+                    index         = static_cast<size_t>(position - session_.extra_backends_.begin()) + 1;
+                }
+                if (index < session_.backend_limits_.size()) {
+                    limit = session_.backend_limits_[index];
+                }
+            }
+            requests.push_back({size.backend, session_.owner_id(), pending,
+                                retained, limit});
+        }
+        return requests;
+    }
+
+    bool GraphExecution::fits(const std::vector<DeviceMemoryRequest>& requests,
+                              const std::vector<ggml_tensor*>& params) const {
+        auto manager = manager_;
+        if (manager == nullptr) {
+            return params.empty();
+        }
+        for (const auto& request : requests) {
+            if (!manager->fits_compute_backend_capacity(request, params)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool GraphExecution::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
+        if (sd_backend_is_cpu(session_.backend_)) {
+            sd_backend_cpu_set_n_threads(session_.backend_, n_threads);
+        }
+        if (session_.workspace_.cpu_backend() != nullptr) {
+            sd_backend_cpu_set_n_threads(session_.workspace_.cpu_backend(), n_threads);
+        }
+        auto scheduler = session_.workspace_.scheduler();
+        GGML_ASSERT(scheduler != nullptr);
+        ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
+#ifndef SD_USE_UPSTREAM_GGML
+        struct BeforeComputeCallback {
+            ggml_backend_t backend;
+            const std::function<void()>& callback;
+            bool invoked = false;
+        } before_compute{session_.backend_, on_before_compute};
+        ggml_backend_sched_set_compute_callback(
+            scheduler, [](ggml_backend_t backend, void* data) {
+                auto& hook = *static_cast<BeforeComputeCallback*>(data);
+                if (backend == hook.backend && !hook.invoked) {
+                    hook.invoked = true;
+                    hook.callback();
+                }
+            },
+            &before_compute);
+        const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+        ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr);
+#else
+        on_before_compute();
+        const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+#endif
+        session_.workspace_.synchronize();
+        if (status != GGML_STATUS_SUCCESS) {
+            last_compute_status_ = status;
+            LOG_ERROR("%s compute failed: %s", task_.description.c_str(), ggml_status_to_string(status));
+            return false;
+        }
+        const std::string description = task_.description;
+        if (!task_.debug_tensors.empty()) {
+            std::unordered_set<const ggml_tensor*> graph_tensors;
+            const int leaf_count = ggml_graph_cut::leaf_count(graph);
+            const int node_count = ggml_graph_n_nodes(graph);
+            graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
+            for (int index = 0; index < leaf_count; ++index) {
+                graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
+            }
+            for (int index = 0; index < node_count; ++index) {
+                graph_tensors.insert(ggml_graph_node(graph, index));
+            }
+
+            for (const auto& entry : task_.debug_tensors) {
+                ggml_tensor* tensor = entry.first;
+                if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
+                    continue;
+                }
+                ggml_backend_buffer_t buffer =
+                    tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+                if (buffer == nullptr) {
+                    LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
+                             description.c_str(),
+                             entry.second.c_str());
+                    continue;
+                }
+                if (tensor->type != GGML_TYPE_F32) {
+                    LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
+                             description.c_str(),
+                             entry.second.c_str(),
+                             ggml_type_name(tensor->type));
+                    continue;
+                }
+                auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
+                print_sd_tensor(debug_tensor, false, entry.second.c_str());
+            }
+        }
+
+        return true;
+    }
+
+    bool GraphExecution::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
+        if (!assign_graph_cut_layer_split_backends(graph)) {
+            return false;
+        }
+        const auto params       = collect_used_param_tensors(graph);
+        const auto& cached_plan = resolve_graph_cut_plan(graph);
+        auto manager            = manager_;
+        const bool can_segment  = !is_multi_device() && !sd_backend_is_cpu(session_.backend_) &&
+                                  manager != nullptr && manager->segmented_compute_enabled() &&
+                                  cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1;
+        // An empty workspace request is a lower bound on the full graph's memory demand.
+        bool segmented = can_segment && !fits(memory_requests({{session_.backend_, 0}}, 0), params);
+        ComputeWorkspace::Measurement full_measurement;
+        if (!segmented) {
+            full_measurement = measure(graph);
+            if (full_measurement.buffers.empty()) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return false;
+            }
+            if (can_segment) {
+                // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
+                constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
+                auto requests                      = memory_requests(full_measurement.buffers, session_.cache_.pending_bytes(graph));
+                for (auto& request : requests) {
+                    request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+                }
+                segmented = !fits(requests, params);
+            }
+        }
+        ggml_graph_cut::Plan monolithic_plan;
+        if (!segmented) {
+            monolithic_plan.segments.emplace_back();
+            auto& segment               = monolithic_plan.segments.back();
+            segment.group_name          = "graph";
+            segment.compute_buffer_size = cached_plan.compute_buffer_size;
+            segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
+            segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
+            for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+                segment.internal_node_indices.push_back(i);
+            }
+            for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
+                auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
+                ggml_graph_cut::Segment::InputRef input;
+                input.leaf_index = i;
+                input.type       = canonical_param_tensor(tensor) != nullptr
+                                       ? ggml_graph_cut::Segment::INPUT_PARAM
+                                       : ggml_graph_cut::Segment::INPUT_EXTERNAL;
+                segment.input_refs.push_back(input);
+            }
+        }
+        const auto& plan            = segmented ? cached_plan : monolithic_plan;
+        const bool segments_changed = plan.segments.size() != session_.logged_segment_count_;
+        if (segments_changed && (segmented || session_.logged_segment_count_ > 1)) {
+            LOG_VERBOSE("%s using %zu segment%s", task_.description.c_str(),
+                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+        }
+        SegmentGraphBindings bindings(session_.cut_cache_, plan, graph);
+        SegmentWeightPipeline weights(manager, session_.backend_, session_.owner_id(),
+                                      graph, plan, params_tensor_set_,
+                                      segmented && manager != nullptr && manager->prefetch_enabled());
+
+        std::map<ggml_backend_t, size_t> peak_compute_bytes;
+        auto track_compute_buffer = [&](ggml_backend_t backend) {
+            if (backend != nullptr) {
+                auto& peak = peak_compute_bytes[backend];
+                peak       = std::max(peak, session_.workspace_.bytes(backend));
+            }
+        };
+        for (size_t index = 0; index < plan.segments.size(); ++index) {
+            const auto& segment = plan.segments[index];
+            const bool last     = index + 1 == plan.segments.size();
+            auto fail_segment   = [&](const char* phase) {
+                LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", task_.description.c_str(),
+                          index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
+                return false;
+            };
+            session_.cut_cache_.prune(segment.live_cut_names);
+            bindings.reset(segment);
+            if (!bindings.bind_cached_inputs(segment, task_.description.c_str())) {
+                return fail_segment("input binding");
+            }
+            ggml_context* segment_context = nullptr;
+            auto segment_graph            = segmented
+                                                ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
+                                                : graph;
+            struct SegmentCleanup {
+                GraphExecution& runner;
+                SegmentWeightPipeline& weights;
+                SegmentGraphBindings& bindings;
+                ggml_context* context;
+                ~SegmentCleanup() {
+                    if (!runner.session_.workspace_.segment_end()) {
+                        runner.last_compute_status_ = GGML_STATUS_FAILED;
+                    }
+                    bindings.restore();
+                    weights.segment_end();
+                    ggml_free(context);
+                    runner.sync_runtime_residency();
+                }
+            } segment_cleanup{*this, weights, bindings, segment_context};
+
+            std::unordered_set<const ggml_tensor*> cache_outputs;
+            if (segmented) {
+                for (int output_index : segment.output_node_indices) {
+                    auto tensor = ggml_graph_node(graph, output_index);
+                    if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) &&
+                        GraphCutTensorCache::can_bind_output(session_.backend_, tensor)) {
+                        cache_outputs.insert(tensor);
+                    }
+                }
+            }
+            auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
+            if (!session_.workspace_.prepare(measurement)) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace preparation");
+            }
+            const size_t cut_bytes = last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment);
+            size_t new_cache_bytes = add_bytes(cut_bytes, session_.cache_.pending_bytes(segment_graph));
+            auto ensure_capacity   = [&]() {
+                sync_runtime_residency();
+                auto requests = memory_requests(measurement.buffers, new_cache_bytes);
+                if (fits(requests, weights.params(index))) {
+                    return true;
+                }
+                if (session_.workspace_.release_excess(measurement)) {
+                    sync_runtime_residency();
+                    requests = memory_requests(measurement.buffers, new_cache_bytes);
+                }
+                const bool ready = weights.ensure_segment_capacity(index, requests);
+                if (!ready && manager != nullptr) {
+                    last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                }
+                return ready;
+            };
+            if (!weights.segment_start(index, ensure_capacity)) {
+                return fail_segment("weight preparation");
+            }
+            // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
+            if (!session_.workspace_.prepare(measurement)) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace preparation");
+            }
+            if (!ensure_capacity()) {
+                return fail_segment("workspace capacity check");
+            }
+            if (!session_.cut_cache_.bind_outputs(cache_outputs)) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("output cache allocation");
+            }
+            new_cache_bytes = add_bytes(last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph));
+            if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
+                    pin_multi_device_nodes(scheduler, current);
+                })) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace allocation");
+            }
+            for (const auto& size : measurement.buffers) {
+                track_compute_buffer(size.backend);
+            }
+            if (session_.workspace_.scheduler() != nullptr) {
+                track_compute_buffer(session_.workspace_.cpu_backend());
+            }
+            if (!ensure_capacity()) {
+                return fail_segment("allocated capacity check");
+            }
+            copy_data_to_backend_tensor(segment_graph);
+            auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
+            if (!prefetch_requests.empty()) {
+                weights.enqueue_next(index, prefetch_requests.front());
+            }
+            LOG_DEBUG("%s executing segment %zu/%zu: %s", task_.description.c_str(),
+                      index + 1, plan.segments.size(), segment.group_name.c_str());
+            if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
+                return fail_segment("execution");
+            }
+            auto cache_status = session_.cache_.capture(segment_graph);
+            if (cache_status == GGML_STATUS_SUCCESS) {
+                cache_status = session_.cut_cache_.capture(graph, segment, task_.description.c_str());
+            }
+            if (cache_status != GGML_STATUS_SUCCESS) {
+                last_compute_status_ = cache_status;
+                return fail_segment("output caching");
+            }
+            sync_runtime_residency();
+            if (last) {
+                if (read_outputs && !read_outputs()) {
+                    return fail_segment("output finalization");
+                }
+                if (!no_return) {
+                    output_ = read_graph_tensor(task_.output, "output");
+                    if (output_.empty()) {
+                        return fail_segment("output readback");
+                    }
+                }
+            }
+            if (!session_.workspace_.segment_end()) {
+                last_compute_status_ = GGML_STATUS_FAILED;
+                return fail_segment("workspace synchronization");
+            }
+            // Final outputs and their callbacks may still be views of consumed cuts.
+            session_.cut_cache_.prune(segment.future_cut_names);
+        }
+        if (segments_changed || peak_compute_bytes != session_.logged_compute_bytes_) {
+            for (const auto& entry : peak_compute_bytes) {
+                LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
+                            task_.description.c_str(), entry.second / (1024.0 * 1024.0),
+                            sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
+                            plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+            }
+            session_.logged_compute_bytes_ = std::move(peak_compute_bytes);
+            session_.logged_segment_count_ = plan.segments.size();
+        }
+        return true;
+    }
+
+    ExecutionSession::ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager)
+        : backend_(backend), residency_manager_(manager), cache_(backend), cut_cache_(backend), workspace_(backend) {
+        GGML_ASSERT(backend_ != nullptr);
+    }
+
+    size_t ExecutionSession::resident_bytes(ggml_backend_t backend) const {
+        size_t bytes = workspace_.bytes(backend);
+        if (backend == backend_) {
+            const auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
+            bytes          = add(bytes, cache_.resident_bytes(ggml_backend_get_device(backend)));
+            bytes          = add(bytes, cut_cache_.resident_bytes(ggml_backend_get_device(backend)));
+        }
+        return bytes;
+    }
+
+    void ExecutionSession::set_layer_split_enabled(bool enabled) {
+        layer_split_enabled_ = enabled;
+        if (!enabled) {
+            weight_assignments_.clear();
+            node_assignments_.clear();
+            primary_notice_logged_ = false;
+        }
+    }
+
+    void ExecutionSession::set_backend_limits(const std::vector<size_t>& limits) {
+        backend_limits_ = limits;
+        weight_assignments_.clear();
+        node_assignments_.clear();
+        primary_notice_logged_ = false;
+    }
+
+    void ExecutionSession::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+        extra_backends_.clear();
+        for (auto backend : backends) {
+            if (backend != nullptr && backend != backend_ &&
+                std::find(extra_backends_.begin(), extra_backends_.end(), backend) == extra_backends_.end()) {
+                extra_backends_.push_back(backend);
+            }
+        }
+        workspace_.set_extra_backends(extra_backends_);
+        weight_assignments_.clear();
+        node_assignments_.clear();
+        primary_notice_logged_ = false;
+    }
+
+    ExecutionEngine::~ExecutionEngine() {
+        for (auto& entry : sessions_) {
+            end(*entry.second);
+        }
+    }
+
+    ExecutionSession& ExecutionEngine::create_session(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager) {
+        GGML_ASSERT(manager.get() == manager_);
+        auto session = std::make_unique<ExecutionSession>(backend, manager);
+        auto& result = *session;
+        sessions_.emplace(session->owner_id(), std::move(session));
+        return result;
+    }
+
+    void ExecutionEngine::release_session(ExecutionSession& session) {
+        end(session);
+        sessions_.erase(session.owner_id());
+    }
+
+    bool ExecutionEngine::start(ExecutionSession& session) {
+        if (session.started_) {
+            return true;
+        }
+        sd_backend_set_buffer_cache_enabled(session.backend_, true);
+        for (auto backend : session.extra_backends_) {
+            sd_backend_set_buffer_cache_enabled(backend, true);
+        }
+        session.workspace_.set_extra_backends(session.extra_backends_);
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->register_runtime_owner(session.owner_id());
+        }
+        session.started_ = true;
+        return true;
+    }
+
+    void ExecutionEngine::end(ExecutionSession& session) {
+        GGML_ASSERT(!session.executing_);
+        if (!session.started_) {
+            return;
+        }
+        session.workspace_.release();
+        session.cache_.clear();
+        session.cut_cache_.clear();
+        session.logged_compute_bytes_.clear();
+        session.logged_segment_count_ = 0;
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->clear_prefetched_params(session.owner_id());
+            std::vector<ggml_tensor*> params(session.params_.begin(), session.params_.end());
+            manager->evict_compute_backend_params(params);
+            manager->remove_runtime_owner(session.owner_id());
+        }
+        session.params_.clear();
+        sd_backend_clear_buffer_cache(session.backend_);
+        for (auto backend : session.extra_backends_) {
+            sd_backend_clear_buffer_cache(backend);
+        }
+        session.started_ = false;
+    }
+
+    void ExecutionEngine::sync_residency(ExecutionSession& session) {
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->update_runtime_residency(session.owner_id(), session.backend_, session.resident_bytes(session.backend_));
+            for (auto backend : session.extra_backends_) {
+                manager->update_runtime_residency(session.owner_id(), backend, session.resident_bytes(backend));
+            }
+        }
+    }
+
+    bool ExecutionEngine::ensure_capacity(
+        DeviceResidencyManager& manager,
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& required_params,
+        const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+        const std::vector<ggml_tensor*>& protected_params) {
+        GGML_ASSERT(&manager == manager_);
+        auto capacity      = [&]() { return manager.inspect_compute_backend_capacity(request, required_params); };
+        const auto initial = capacity();
+        if (!initial.valid) {
+            return false;
+        }
+        if (initial.fits()) {
+            return true;
+        }
+        for (auto& entry : sessions_) {
+            auto& session = *entry.second;
+            if (session.owner_id() != request.owner_id && session.workspace_.release()) {
+                sync_residency(session);
+                if (capacity().fits(true)) {
+                    return true;
+                }
+            }
+        }
+        auto protected_weights = protected_params;
+        protected_weights.insert(protected_weights.end(), required_params.begin(), required_params.end());
+        auto candidates = manager.evictable_compute_backend_blocks(request.compute_backend, protected_weights);
+        std::unordered_map<const ggml_tensor*, size_t> preferred_positions;
+        for (size_t index = 0; index < preferred_eviction_order.size(); ++index) {
+            for (auto* tensor : preferred_eviction_order[index]) {
+                auto* weight = manager.resolve_param_tensor(tensor);
+                preferred_positions.emplace(weight != nullptr ? weight : tensor, index);
+            }
+        }
+        auto preferred_position = [&](const WeightResidencyBlock& block) {
+            size_t position = 0;
+            // A shared allocation follows the last of its members in the eviction order.
+            for (auto* tensor : block.tensors) {
+                const auto found = preferred_positions.find(tensor);
+                if (found == preferred_positions.end()) {
+                    return SIZE_MAX;
+                }
+                position = std::max(position, found->second);
+            }
+            return position;
+        };
+        std::stable_sort(candidates.begin(), candidates.end(), [&](const auto& lhs, const auto& rhs) {
+            if (lhs.requires_demotion != rhs.requires_demotion) {
+                return !lhs.requires_demotion;
+            }
+            if (!lhs.requires_demotion) {
+                const auto left = preferred_position(lhs), right = preferred_position(rhs);
+                if (left != right) {
+                    return left < right;
+                }
+                if (left != SIZE_MAX) {
+                    return false;
+                }
+            }
+            return lhs.last_use_epoch < rhs.last_use_epoch;
+        });
+        for (const auto& block : candidates) {
+            if (manager.evict_compute_backend_block(block) && capacity().fits(true)) {
+                return true;
+            }
+        }
+        auto available = capacity();
+        if (available.fits(true)) {
+            return true;
+        }
+        if (sd_backend_is(request.compute_backend, "Vulkan") &&
+            available.required_budget_bytes <= available.available_budget_bytes) {
+            ggml_backend_synchronize(request.compute_backend);
+            sd_backend_clear_buffer_cache(request.compute_backend);
+            const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
+            do {
+                available = capacity();
+                if (available.fits(true)) {
+                    return true;
+                }
+                if (available.required_budget_bytes > available.available_budget_bytes) {
+                    break;
+                }
+                std::this_thread::sleep_for(std::chrono::milliseconds(10));
+            } while (std::chrono::steady_clock::now() < deadline);
+        }
+        available               = manager.inspect_compute_backend_capacity(request, required_params, true);
+        const auto device_bytes = available.available_device_bytes == SIZE_MAX
+                                      ? std::string("unknown")
+                                      : sd_format("%.2f MB", available.available_device_bytes / (1024.0 * 1024.0));
+        const auto budget_bytes = available.available_budget_bytes == SIZE_MAX
+                                      ? std::string("unlimited")
+                                      : sd_format("%.2f MB", available.available_budget_bytes / (1024.0 * 1024.0));
+        LOG_WARN("execution engine cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
+                 ggml_backend_name(request.compute_backend),
+                 available.required_device_bytes / (1024.0 * 1024.0),
+                 available.required_budget_bytes / (1024.0 * 1024.0), device_bytes.c_str(), budget_bytes.c_str());
+        return false;
+    }
+
+    void ExecutionEngine::clear_cache(ExecutionSession& session) {
+        session.cache_.clear();
+        sync_residency(session);
+    }
+
+    void ExecutionEngine::discard_graph(ExecutionSession& session) {
+        GGML_ASSERT(!session.executing_);
+        session.cache_.graph_end(false);
+        session.cut_cache_.clear();
+        session.workspace_.release();
+        sync_residency(session);
+    }
+
+    ComputeResult ExecutionEngine::execute(ExecutionSession& session, const ComputeTask& task) {
+        if (session.executing_ || task.graph == nullptr || !start(session)) {
+            return {};
+        }
+        session.executing_ = true;
+        ComputeResult result;
+        try {
+            GraphExecution execution(*this, session, task);
+            result = execution.run();
+        } catch (const std::bad_alloc&) {
+            result.status = GGML_STATUS_ALLOC_FAILED;
+            LOG_ERROR("%s graph allocation failed", task.description.c_str());
+        } catch (const std::exception& error) {
+            LOG_ERROR("%s graph execution failed on %s: %s", task.description.c_str(),
+                      ggml_backend_name(session.backend_), error.what());
+        } catch (...) {
+            LOG_ERROR("%s graph execution failed on %s: unknown exception", task.description.c_str(),
+                      ggml_backend_name(session.backend_));
+        }
+        if (!session.workspace_.segment_end()) {
+            result.status = GGML_STATUS_FAILED;
+        }
+        const bool success = result.status == GGML_STATUS_SUCCESS;
+        session.cache_.graph_end(success);
+        session.cut_cache_.clear();
+        if (!success) {
+            result.output = {};
+            session.workspace_.release();
+        }
+        session.executing_ = false;
+        sync_residency(session);
+        return result;
+    }
+}

src/core/execution_engine.h

@@ -0,0 +1,107 @@
+#ifndef __SD_CORE_EXECUTION_ENGINE_H__
+#define __SD_CORE_EXECUTION_ENGINE_H__
+
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "core/compute_workspace.h"
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/tensor.hpp"
+#include "device_residency_manager.h"
+
+namespace sd {
+    struct ComputeTask {
+        std::string description;
+        ggml_cgraph* graph  = nullptr;
+        ggml_tensor* output = nullptr;
+        std::vector<ggml_tensor*> weights;
+        std::map<ggml_tensor*, const void*> inputs;
+        std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+        int n_threads  = 1;
+        bool no_return = false;
+        std::function<bool()> read_outputs;
+    };
+
+    struct ComputeResult {
+        ggml_status status = GGML_STATUS_FAILED;
+        Tensor<float> output;
+    };
+
+    class ExecutionSession {
+        friend class ExecutionEngine;
+        friend class GraphExecution;
+
+        ggml_backend_t backend_;
+        std::weak_ptr<DeviceResidencyManager> residency_manager_;
+        RunnerCache cache_;
+        GraphCutTensorCache cut_cache_;
+        ComputeWorkspace workspace_;
+        std::unordered_set<ggml_tensor*> params_;
+        std::vector<ggml_backend_t> extra_backends_;
+        size_t max_backend_bytes_ = 0;
+        bool layer_split_enabled_ = false;
+        std::vector<size_t> backend_limits_;
+        ggml_graph_cut::PlanCache plan_cache_;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> weight_assignments_;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments_;
+        bool primary_notice_logged_ = false;
+        std::map<ggml_backend_t, size_t> logged_compute_bytes_;
+        size_t logged_segment_count_ = 0;
+        bool started_                = false;
+        bool executing_              = false;
+
+        size_t resident_bytes(ggml_backend_t backend) const;
+        uintptr_t owner_id() const { return reinterpret_cast<uintptr_t>(this); }
+
+    public:
+        ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager);
+        ExecutionSession(const ExecutionSession&)            = delete;
+        ExecutionSession& operator=(const ExecutionSession&) = delete;
+
+        RunnerCache& cache() { return cache_; }
+        bool started() const { return started_; }
+        size_t compute_bytes() const { return workspace_.bytes(backend_); }
+        void clear_plan_cache() { plan_cache_.graph_cut_plans.clear(); }
+        void set_max_backend_bytes(size_t bytes) { max_backend_bytes_ = bytes; }
+        void set_layer_split_enabled(bool enabled);
+        void set_backend_limits(const std::vector<size_t>& limits);
+        void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+    };
+
+    class ExecutionEngine {
+        friend class GraphExecution;
+
+        std::map<uintptr_t, std::unique_ptr<ExecutionSession>> sessions_;
+        DeviceResidencyManager* manager_;
+
+        void sync_residency(ExecutionSession& session);
+
+    public:
+        explicit ExecutionEngine(DeviceResidencyManager* manager = nullptr)
+            : manager_(manager) {}
+        ~ExecutionEngine();
+        ExecutionSession& create_session(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager);
+        void release_session(ExecutionSession& session);
+        bool start(ExecutionSession& session);
+        void end(ExecutionSession& session);
+        void clear_cache(ExecutionSession& session);
+        void discard_graph(ExecutionSession& session);
+        bool ensure_capacity(DeviceResidencyManager& manager,
+                             const DeviceMemoryRequest& request,
+                             const std::vector<ggml_tensor*>& required_params,
+                             const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+                             const std::vector<ggml_tensor*>& protected_params);
+        // Graphs, input data and output callbacks remain borrowed until execute returns.
+        ComputeResult execute(ExecutionSession& session, const ComputeTask& task);
+    };
+}
+
+#endif  // __SD_CORE_EXECUTION_ENGINE_H__

src/core/ggml_extend.h

@@ -10,33 +10,24 @@
 
 struct GGMLRunnerContext;
 
+constexpr uint32_t GGML_PRECISION_F32_VULKAN = 1u << 0;
+constexpr uint32_t GGML_PRECISION_F32_ROCM   = 1u << 1;
+
 #define EPS 1e-05f
 
 static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
 
-// n-mode tensor-matrix product
-// example: 2-mode product
-// A: [ne03, k, ne01, ne00]
-// B: k rows, m columns => [k, m]
-// result is [ne03, m, ne01, ne00]
-ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0);
-
-// Kronecker product
-// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10]
-ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b);
-
-ggml_tensor* ggml_ext_cont(ggml_context* ctx,
-                           ggml_tensor* x);
+ggml_tensor* ggml_ext_cont(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
 
 // torch like permute
-ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+ggml_tensor* ggml_ext_torch_permute(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     int axis0,
                                     int axis1,
                                     int axis2,
                                     int axis3);
 
-ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+ggml_tensor* ggml_ext_slice(GGMLRunnerContext* runner_ctx,
                             ggml_tensor* x,
                             int dim,
                             int64_t start,
@@ -44,48 +35,38 @@ ggml_tensor* ggml_ext_slice(ggml_context* ctx,
                             bool cont = true);
 
 // example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W])
-std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+std::vector<ggml_tensor*> ggml_ext_chunk(GGMLRunnerContext* runner_ctx,
                                          ggml_tensor* x,
                                          int num,
                                          int64_t dim,
                                          bool cont = true);
 
-ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true);
+ggml_tensor* ggml_ext_silu_act(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool gate_first = true);
 
 // Without gate, hidden contains [hidden, gate] packed along ne[0].
-ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr);
+ggml_tensor* ggml_ext_situ(GGMLRunnerContext* runner_ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr);
 
-ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
-                                    ggml_tensor* a);
-
-ggml_tensor* ggml_ext_scale(ggml_context* ctx,
-                            ggml_tensor* x,
-                            float factor,
-                            bool inplace = false);
+ggml_tensor* ggml_ext_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* x, float factor, bool inplace = false);
 
 // scale and optional shift must already be broadcastable to x in ggml layout.
 // Computes (x + x * scale) + shift; F32 uses sd.modulate.v1 when supported by ctx->backend.
-ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx,
+ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* runner_ctx,
                                ggml_tensor* x,
                                ggml_tensor* shift,
                                ggml_tensor* scale);
 
-ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
-                           ggml_tensor* x,
-                           bool inplace = false);
+ggml_tensor* ggml_ext_gelu(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace = false);
 
-ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
-                                 ggml_tensor* x,
-                                 bool inplace = false);
+ggml_tensor* ggml_ext_gelu_quick(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace = false);
 
-ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+ggml_tensor* ggml_ext_linear(GGMLRunnerContext* runner_ctx,
                              ggml_tensor* x,
                              ggml_tensor* w,
                              ggml_tensor* b,
                              bool force_prec_f32 = false,
                              float scale         = 1.f);
 
-ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+ggml_tensor* ggml_ext_linear_i8_tensorwise(GGMLRunnerContext* runner_ctx,
                                            ggml_tensor* x,
                                            ggml_tensor* w,
                                            ggml_tensor* weight_scale,
@@ -93,8 +74,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
                                            int convrot_group_size,
                                            float scale = 1.f);
 
-ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
-                              ggml_backend_t backend,
+ggml_tensor* ggml_ext_pad_ext(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               int lp0,
                               int rp0,
@@ -107,7 +87,7 @@ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
                               bool circular_x = false,
                               bool circular_y = false);
 
-ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+ggml_tensor* ggml_ext_pad(GGMLRunnerContext* runner_ctx,
                           ggml_tensor* x,
                           int p0,
                           int p1,
@@ -118,7 +98,7 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
 
 // ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N].
 // force_prec_f32 keeps both input patches and weights in F32.
-ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
+ggml_tensor* ggml_ext_conv_1d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -132,7 +112,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
 // x: [N, IC, IH, IW]
 // b: [OC,]
 // result: [N, OC, OH, OW]
-ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+ggml_tensor* ggml_ext_conv_2d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -151,8 +131,7 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
 // x: [N, IC, IH, IW]
 // b: [OC,]
 // result: [N*OC, OD, OH, OW]
-ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
-                              ggml_backend_t backend,
+ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -168,65 +147,33 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d2              = 1,
                               bool force_prec_f32 = false,
                               bool direct         = false,
-                              float scale         = 1.f);
-
-// w: [OC，IC, KD, 1 * 1]
-// x: [N, IC, ID, IH*IW]
-// b: [OC,]
-// result: [N, OC, OD, OH*OW]
-ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
-                                    ggml_tensor* x,
-                                    ggml_tensor* w,
-                                    ggml_tensor* b,
-                                    int s2 = 1,
-                                    int p2 = 1,
-                                    int d2 = 1);
-
-// qkv: [N, L, 3*C]
-// return: ([N, L, C], [N, L, C], [N, L, C])
-std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
-                                    ggml_tensor* qkv);
-
-// qkv: [N, 3*C, H, W]
-// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W])
-std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
-                                          ggml_tensor* qkv);
+                              float scale         = 1.f,
+                              bool allow_conv2d   = false);
 
 // Constant and cast helpers require the built-in tensors initialized by GGMLRunner.
-ggml_tensor* ggml_ext_full(ggml_context* ctx,
+ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx,
                            float value,
                            int64_t ne0,
                            int64_t ne1,
                            int64_t ne2,
                            int64_t ne3);
 
-ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
-                            int64_t ne0,
-                            int64_t ne1,
-                            int64_t ne2,
-                            int64_t ne3);
+ggml_tensor* ggml_ext_zeros(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3);
 
-ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
-                                 ggml_tensor* x);
+ggml_tensor* ggml_ext_zeros_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
 
-ggml_tensor* ggml_ext_ones(ggml_context* ctx,
-                           int64_t ne0,
-                           int64_t ne1,
-                           int64_t ne2,
-                           int64_t ne3);
+ggml_tensor* ggml_ext_ones(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3);
 
-ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
-                                ggml_tensor* x);
+ggml_tensor* ggml_ext_ones_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
 
-ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a);
+ggml_tensor* ggml_ext_cast_f32(GGMLRunnerContext* runner_ctx, ggml_tensor* a);
 
 // q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head]
 // k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head]
 // v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head]
 // mask: [N, L_q, L_k]
 // return: [N, L_q, C]
-ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
-                                    ggml_backend_t backend,
+ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* q,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
@@ -235,31 +182,138 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     bool skip_reshape     = false,
                                     bool flash_attn       = false,
                                     float kv_scale        = 1.0f,
-                                    bool sage_attn        = false,
                                     bool* used_flash_attn = nullptr);
 
-ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+ggml_tensor* ggml_ext_layer_norm(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
                                  float eps = EPS);
 
-ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+ggml_tensor* ggml_ext_group_norm(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
                                  int num_groups = 32,
                                  float eps      = 1e-6f);
 
-ggml_tensor* ggml_ext_timestep_embedding(
-    ggml_context* ctx,
-    ggml_tensor* timesteps,
-    int dim,
-    int max_period    = 10000,
-    float time_factor = 1.0f);
+ggml_tensor* ggml_ext_timestep_embedding(GGMLRunnerContext* runner_ctx,
+                                         ggml_tensor* timesteps,
+                                         int dim,
+                                         int max_period    = 10000,
+                                         float time_factor = 1.0f);
 
-ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
-                                 std::vector<ggml_tensor*>& tensors,
-                                 int dim);
+ggml_tensor* ggml_ext_vec_concat(GGMLRunnerContext* runner_ctx, std::vector<ggml_tensor*>& tensors, int dim);
+
+#ifndef SD_USE_UPSTREAM_GGML
+ggml_tensor* ggml_ext_mul_mat_i8_tensorwise(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* weight,
+                                            ggml_tensor* input,
+                                            ggml_tensor* weight_scale,
+                                            ggml_tensor* bias,
+                                            int group_size,
+                                            ggml_tensor* input_scale = nullptr);
+
+ggml_tensor* ggml_ext_quantize_i8_convrot(GGMLRunnerContext* runner_ctx, ggml_tensor* input, int group_size);
+
+// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
+// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
+ggml_tensor* ggml_ext_mul_mat_w4_convrot(GGMLRunnerContext* runner_ctx,
+                                         ggml_tensor* weight,
+                                         ggml_tensor* input,
+                                         ggml_tensor* weight_scale,
+                                         ggml_tensor* s_channel,
+                                         ggml_tensor* s_rel,
+                                         ggml_tensor* bias,
+                                         SDW4Kind kind,
+                                         int group_size);
+
+// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout.
+// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference.
+ggml_tensor* ggml_ext_sage_attn(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* q,
+                                ggml_tensor* k,
+                                ggml_tensor* v,
+                                float scale,
+                                ggml_sage_attn_mode mode);
+
+// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
+ggml_tensor* ggml_ext_sol_attn(GGMLRunnerContext* runner_ctx,
+                               ggml_tensor* q,
+                               ggml_tensor* k,
+                               ggml_tensor* v,
+                               float scale,
+                               float tau);
+
+// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes].
+// Returns Q's layout, or nullptr when the backend has no compatible kernel.
+ggml_tensor* ggml_ext_window_attn(GGMLRunnerContext* runner_ctx,
+                                  ggml_tensor* q,
+                                  ggml_tensor* k,
+                                  ggml_tensor* v,
+                                  const SDWindowAttentionParams& params);
+#endif
+
+// x: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
+// Returns [head_dim, tokens, heads*batch]; channels beyond rot_dim are unchanged.
+ggml_tensor* ggml_ext_rope_apply(GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* pe,
+                                 bool interleaved = true);
+
+ggml_tensor* ggml_ext_repeat_4d(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* x,
+                                int64_t ne0,
+                                int64_t ne1,
+                                int64_t ne2,
+                                int64_t ne3);
+
+ggml_tensor* ggml_ext_repeat(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shape);
+
+ggml_tensor* ggml_ext_embedding(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* weight,
+                                ggml_tensor* input_ids,
+                                ggml_tensor* weight_scale = nullptr,
+                                int group_size            = 0);
+
+ggml_tensor* ggml_ext_prepare_linear_weight(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* w,
+                                            ggml_tensor* x,
+                                            bool force_prec_f32);
+
+ggml_tensor* ggml_ext_prepare_i8_input(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
+
+bool ggml_ext_force_prec_f32(GGMLRunnerContext* runner_ctx,
+                             bool requested,
+                             uint32_t flags);
+
+#ifndef SD_USE_UPSTREAM_GGML
+ggml_tensor* ggml_ext_linear_w4_convrot(GGMLRunnerContext* runner_ctx,
+                                        ggml_tensor* x,
+                                        ggml_tensor* w,
+                                        ggml_tensor* ws,
+                                        ggml_tensor* sc,
+                                        ggml_tensor* sr,
+                                        ggml_tensor* bias,
+                                        SDW4Kind kind,
+                                        int group,
+                                        float scale = 1.f);
+#endif
+
+ggml_tensor* ggml_ext_conv_2d_grouped(GGMLRunnerContext* runner_ctx,
+                                      ggml_tensor* x,
+                                      ggml_tensor* w,
+                                      ggml_tensor* b,
+                                      int groups,
+                                      int s0          = 1,
+                                      int s1          = 1,
+                                      int p0          = 0,
+                                      int p1          = 0,
+                                      int d0          = 1,
+                                      int d1          = 1,
+                                      bool direct     = false,
+                                      bool circular_x = false,
+                                      bool circular_y = false,
+                                      float scale     = 1.f);
 
 #endif  // __SD_CORE_GGML_EXTEND_H__

src/core/ggml_extend_backend.cpp

@@ -7,6 +7,7 @@
 #include <mutex>
 #include <sstream>
 #include <stdexcept>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -447,7 +448,19 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
         return false;
     }
     auto dev = ggml_backend_get_device(backend);
-    return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
+    if (dev == nullptr) {
+        return false;
+    }
+    // ggml_backend_dev_type() can query the driver on every call (CUDA reads all device
+    // properties), and this check runs inside per-tensor residency loops.
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, bool> is_cpu_by_device;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto it = is_cpu_by_device.find(dev);
+    if (it == is_cpu_by_device.end()) {
+        it = is_cpu_by_device.emplace(dev, ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU).first;
+    }
+    return it->second;
 }
 
 static int cuda_compiled_arch(ggml_backend_t backend) {

src/core/ggml_extensions.cpp

@@ -131,20 +131,12 @@ bool sd_ggml_rope_describe(const ggml_tensor* node) {
            std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0;
 }
 
-struct I8Attributes {
-    int64_t shape[3];
-    int32_t group_size;
-    int32_t flags;
-};
-enum I8Flags { I8_PACKED       = 1,
-               I8_INPUT_SCALE  = 2,
-               I8_WEIGHT_SCALE = 4,
-               I8_BIAS         = 8 };
-
-static bool valid_group(int group) {
+bool sd_ggml_valid_convrot_group(int group) {
     return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
 }
+
 static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX)));
+
 static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
     if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4)
         return false;
@@ -155,15 +147,17 @@ static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
     *height = padded + scales;
     return true;
 }
+
 static bool scale_matches(const ggml_tensor* scale, int64_t count) {
     return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count);
 }
+
 static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
     const auto* w = op.weight;
     const auto* x = op.input;
     if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 ||
         !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] ||
-        (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size)))
+        (op.group_size != 0 && (!sd_ggml_valid_convrot_group(op.group_size) || w->ne[0] % op.group_size)))
         return false;
     int64_t rows = 1;
     for (int i = 0; i < 3; ++i) {
@@ -177,7 +171,7 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
         rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height))
         return false;
     if (op.packed) {
-        if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale ||
+        if (x->type != GGML_TYPE_I8 || !sd_ggml_valid_convrot_group(op.group_size) || op.input_scale ||
             x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1)
             return false;
     } else {
@@ -191,14 +185,15 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
     }
     return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]);
 }
+
 static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-    I8Attributes attr;
+    SDI8Attributes attr;
     std::memcpy(&attr, attributes, sizeof(attr));
     if (count < 2 || (attr.flags & ~15))
         return false;
-    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0};
+    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & SD_I8_PACKED) != 0};
     int index        = 2;
-    for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) {
+    for (auto entry : {std::make_pair(SD_I8_INPUT_SCALE, &op.input_scale), std::make_pair(SD_I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(SD_I8_BIAS, &op.bias)}) {
         if (attr.flags & entry.first) {
             if (index >= count)
                 return false;
@@ -212,10 +207,11 @@ static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attribut
     std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape));
     return true;
 }
+
 static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
     int32_t group;
     std::memcpy(&group, attributes, sizeof(group));
-    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
+    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !sd_ggml_valid_convrot_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
         return false;
     const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]);
     int64_t height;
@@ -227,38 +223,42 @@ static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* att
     output->ne[2] = output->ne[3] = 1;
     return true;
 }
+
 ggml_ext_op_t sd_ggml_i8_op() {
-    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(SDI8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
     return op;
 }
+
 ggml_ext_op_t sd_ggml_convrot_op() {
     static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
     return op;
 }
+
 bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
     if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
         return false;
     *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
     if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
-        I8Attributes attr;
+        SDI8Attributes attr;
         std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
         if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape)))
             return false;
         op->group_size = attr.group_size;
-        op->packed     = (attr.flags & I8_PACKED) != 0;
+        op->packed     = (attr.flags & SD_I8_PACKED) != 0;
         int index      = 2;
-        if (attr.flags & I8_INPUT_SCALE)
+        if (attr.flags & SD_I8_INPUT_SCALE)
             op->input_scale = node->src[index++];
-        if (attr.flags & I8_WEIGHT_SCALE)
+        if (attr.flags & SD_I8_WEIGHT_SCALE)
             op->weight_scale = node->src[index++];
-        if (attr.flags & I8_BIAS)
+        if (attr.flags & SD_I8_BIAS)
             op->bias = node->src[index++];
-        if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias))
+        if (((attr.flags & SD_I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & SD_I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & SD_I8_BIAS) && !op->bias))
             return false;
     } else
         return false;
     return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1];
 }
+
 bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
     if (!node || !node->src[0] || !ggml_is_contiguous(node))
         return false;
@@ -269,48 +269,11 @@ bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
     ggml_ext_tensor_desc output;
     return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne));
 }
-ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group, ggml_tensor* xs) {
-    if (!w || !x)
-        return nullptr;
-    const auto* logical = x;
-    if (x->type == GGML_TYPE_I8 && !xs) {
-        int packed_group;
-        if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
-            return nullptr;
-        logical = x->src[0];
-    }
-    I8Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
-    ggml_tensor* inputs[5] = {w, x};
-    int count              = 2;
-    if (xs) {
-        attr.flags |= I8_INPUT_SCALE;
-        inputs[count++] = xs;
-    }
-    if (ws) {
-        attr.flags |= I8_WEIGHT_SCALE;
-        inputs[count++] = ws;
-    }
-    if (bias) {
-        attr.flags |= I8_BIAS;
-        inputs[count++] = bias;
-    }
-    return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
-}
-ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
-    int32_t attr = group;
-    return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
-}
-
-struct W4Attributes {
-    int64_t shape[3];
-    SDW4Kind kind;
-    int32_t has_bias;
-};
 
 static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (!attributes || count < 3)
         return false;
-    W4Attributes attr;
+    SDW4Attributes attr;
     std::memcpy(&attr, attributes, sizeof(attr));
     const bool codebook = attr.kind == SDW4Kind::W4A8;
     if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) ||
@@ -348,7 +311,7 @@ static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attribut
 
 ggml_ext_op_t sd_ggml_w4_op() {
     static const auto op = [] {
-        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer};
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(SDW4Attributes), w4_infer};
         auto result                 = ggml_ext_register_op(&desc);
         GGML_ASSERT(result);
         return result;
@@ -359,7 +322,7 @@ ggml_ext_op_t sd_ggml_w4_op() {
 bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
     if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node))
         return false;
-    W4Attributes attr;
+    SDW4Attributes attr;
     std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
     const bool codebook = attr.kind == SDW4Kind::W4A8;
     const int count     = (codebook ? 5 : 3) + (attr.has_bias != 0);
@@ -372,30 +335,6 @@ bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
     return true;
 }
 
-ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) {
-    if (!w || !x || group < 0 || (group && !valid_group(group)))
-        return nullptr;
-    const int quant_group = group ? group : 1;
-    if (x->type == GGML_TYPE_F32)
-        x = ggml_quantize_i8_convrot(ctx, x, quant_group);
-    int packed_group;
-    if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group)
-        return nullptr;
-    const auto* logical    = x->src[0];
-    W4Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0};
-    ggml_tensor* inputs[6] = {w, x, ws};
-    int count              = 3;
-    if (kind == SDW4Kind::W4A8) {
-        inputs[count++] = sc;
-        inputs[count++] = sr;
-    } else if (sc || sr) {
-        return nullptr;
-    }
-    if (bias)
-        inputs[count++] = bias;
-    return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &attr);
-}
-
 static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
@@ -443,12 +382,6 @@ bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams*
     return true;
 }
 
-ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode) {
-    ggml_tensor* inputs[]              = {q, k, v};
-    const SDSageAttentionParams params = {scale, mode};
-    return ggml_ext_op_create(ctx, sd_ggml_sage_attn_op(), inputs, 3, &params);
-}
-
 static bool sol_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
@@ -495,13 +428,10 @@ bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* pa
     return true;
 }
 
-ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau) {
-    ggml_tensor* inputs[]             = {q, k, v};
-    const SDSolAttentionParams params = {scale, tau};
-    return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, &params);
-}
-
-static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs,
+                               int count,
+                               const void* attributes,
+                               ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
     SDWindowAttentionParams p;
@@ -531,7 +461,7 @@ static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void*
 
 ggml_ext_op_t sd_ggml_window_attn_op() {
     static const auto op = [] {
-        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), window_attn_infer};
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), sd_ggml_window_attn_infer};
         auto result                 = ggml_ext_register_op(&desc);
         GGML_ASSERT(result);
         return result;
@@ -544,22 +474,12 @@ bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionPara
         return false;
     const auto* attributes = ggml_ext_op_attributes(node);
     ggml_ext_tensor_desc out;
-    if (!window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne)))
+    if (!sd_ggml_window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne)))
         return false;
     std::memcpy(params, attributes, sizeof(*params));
     return true;
 }
 
-ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params) {
-    ggml_tensor* inputs[] = {q, k, v};
-    ggml_ext_tensor_desc out;
-    if (!backend || !window_attn_infer(inputs, 3, &params, &out))
-        return nullptr;
-    sd_ggml_register_extensions(backend);
-    auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, &params);
-    return ggml_backend_supports_op(backend, result) ? result : nullptr;
-}
-
 void sd_ggml_register_extensions(ggml_backend_t backend) {
     if (!backend) {
         return;
@@ -583,13 +503,6 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
     sd_cpu::register_extensions(device);
 }
 
-ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) {
-    sd_ggml_register_extensions(backend);
-    ggml_tensor* inputs[] = {x, pe};
-    int32_t attributes    = interleaved ? 1 : 0;
-    return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
-}
-
 #else
 void sd_ggml_register_extensions(ggml_backend_t) {}
 #endif

src/core/ggml_extensions.h

@@ -9,18 +9,24 @@ void sd_ggml_register_extensions(ggml_backend_t backend);
 #include "ggml-ext.h"
 
 ggml_ext_op_t sd_ggml_modulate_op();
+
 bool sd_ggml_modulate_describe(const ggml_tensor* node);
+
 size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node);
 
 ggml_ext_op_t sd_ggml_rope_op();
+
 bool sd_ggml_rope_describe(const ggml_tensor* node);
-// Apply precomputed rotary matrices, with shapes in ne[] order:
-// a: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
-// pe stores [[cos, -sin], [sin, cos]] for each pair and token, shared across heads/batches.
-// Returns contiguous F32 [head_dim, tokens, heads*batch]. Input strides are preserved.
-// Interleaved pairs are (2*i, 2*i+1); otherwise pairs are (i, i+rot_dim/2).
-// Channels from rot_dim to head_dim are copied unchanged.
-ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
+
+struct SDI8Attributes {
+    int64_t shape[3];
+    int32_t group_size;
+    int32_t flags;
+};
+enum SDI8Flags { SD_I8_PACKED       = 1,
+                 SD_I8_INPUT_SCALE  = 2,
+                 SD_I8_WEIGHT_SCALE = 4,
+                 SD_I8_BIAS         = 8 };
 
 struct SDI8Operation {
     const ggml_tensor* weight;
@@ -32,16 +38,25 @@ struct SDI8Operation {
     bool packed;
 };
 
+bool sd_ggml_valid_convrot_group(int group);
+
 ggml_ext_op_t sd_ggml_i8_op();
+
 ggml_ext_op_t sd_ggml_convrot_op();
+
 bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
+
 bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
-ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr);
-ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
 
 enum class SDW4Kind : int32_t { W4A4 = 1,
                                 W4A8 = 2 };
 
+struct SDW4Attributes {
+    int64_t shape[3];
+    SDW4Kind kind;
+    int32_t has_bias;
+};
+
 struct SDW4Operation {
     const ggml_tensor* weight;
     const ggml_tensor* input;
@@ -53,10 +68,8 @@ struct SDW4Operation {
 };
 
 ggml_ext_op_t sd_ggml_w4_op();
+
 bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation);
-// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
-// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
-ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size);
 
 enum ggml_sage_attn_mode {
     GGML_SAGE_ATTN_AUTO = 0,
@@ -70,10 +83,8 @@ struct SDSageAttentionParams {
 };
 
 ggml_ext_op_t sd_ggml_sage_attn_op();
+
 bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params);
-// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout.
-// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference.
-ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode);
 
 struct SDSolAttentionParams {
     float scale;
@@ -86,16 +97,19 @@ struct SDWindowAttentionParams {
     int32_t preserve_boundary;
 };
 
+bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs,
+                               int count,
+                               const void* attributes,
+                               ggml_ext_tensor_desc* out);
+
 ggml_ext_op_t sd_ggml_window_attn_op();
+
 bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params);
-// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes].
-// Returns Q's layout, or nullptr when the backend has no compatible kernel.
-ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params);
 
 ggml_ext_op_t sd_ggml_sol_attn_op();
+
 bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params);
-// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
-ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau);
+
 #endif
 
 #endif

src/core/ggml_runner.cpp

@@ -1,34 +1,14 @@
 #include <algorithm>
 #include <exception>
-#include <map>
 #include <utility>
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
-#include "core/layer_split_partition.h"
-#include "core/segment_graph_bindings.h"
-#include "core/segment_weight_pipeline.h"
 
 using namespace sd;
 
-ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
-                                    ggml_tensor* q,
-                                    ggml_tensor* k,
-                                    ggml_tensor* v,
-                                    int64_t n_head,
-                                    ggml_tensor* mask,
-                                    bool skip_reshape,
-                                    bool flash_attn,
-                                    float kv_scale,
-                                    bool* used_flash_attn) {
-    if (ctx->attn_scale > 0.f) {
-        kv_scale = ctx->attn_scale;
-    }
-    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
-}
-
 void GGMLRunner::alloc_params_ctx() {
     ggml_init_params params;
     params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
@@ -117,16 +97,6 @@ std::vector<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf
     return used_params;
 }
 
-void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors) {
-    if (tensors.empty()) {
-        return;
-    }
-    auto manager = residency_manager.lock();
-    if (manager != nullptr) {
-        manager->evict_compute_backend_params(tensors);
-    }
-}
-
 void GGMLRunner::prepare_build_in_tensor_before() {
     one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
     ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one");
@@ -164,7 +134,7 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
             ggml_build_forward_expand(gf, entry.first);
         }
     }
-    for (const auto& entry : cache_.outputs()) {
+    for (const auto& entry : execution_session_.cache().outputs()) {
         if (entry.second != nullptr) {
             ggml_build_forward_expand(gf, entry.second);
         }
@@ -188,365 +158,22 @@ bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph,
     return true;
 }
 
-ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
-    if (tensor == nullptr || tensor->buffer == nullptr) {
-        return nullptr;
-    }
-    if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
-        ggml_backend_buffer_is_host(tensor->buffer)) {
-        return nullptr;
-    }
-    ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
-    if (dev == nullptr) {
-        return nullptr;
-    }
-    if (ggml_backend_get_device(runtime_backend) == dev) {
-        return runtime_backend;
-    }
-    for (ggml_backend_t backend : extra_runtime_backends) {
-        if (ggml_backend_get_device(backend) == dev) {
-            return backend;
-        }
-    }
-    return nullptr;
-}
-
-static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
-    // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
-    if (sd_backend_is(backend, "Vulkan")) {
-        return true;
-    }
-    auto buft        = ggml_backend_get_default_buffer_type(backend);
-    size_t alignment = ggml_backend_buft_get_alignment(buft);
-    size_t max_size  = ggml_backend_buft_get_max_size(buft);
-    max_size -= max_size % alignment;
-    if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
-        return false;
-    }
-    // Cross-backend copies must also fit the destination's allocation limit.
-    for (auto src : node->src) {
-        if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
-            return false;
-        }
-    }
-    return true;
-}
-
-void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
-    if (sched == nullptr || gf == nullptr) {
-        return;
-    }
-    ggml_backend_t current = runtime_backend;
-    const int n_nodes      = ggml_graph_n_nodes(gf);
-    for (int i = 0; i < n_nodes; i++) {
-        ggml_tensor* node    = ggml_graph_node(gf, i);
-        auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
-        if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
-            current = node_assignment->second;
-        }
-        for (int s = 0; s < GGML_MAX_SRC; s++) {
-            ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
-            if (weight_backend != nullptr) {
-                if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
-                    current = weight_backend;
-                }
-            }
-        }
-        if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
-            node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
-            continue;
-        }
-        if (!backend_can_allocate_op(current, node)) {
-            auto cpu = workspace_.cpu_backend();
-            if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
-                ggml_backend_sched_set_tensor_backend(sched, node, cpu);
-            }
-        } else if (ggml_backend_supports_op(current, node)) {
-            ggml_backend_sched_set_tensor_backend(sched, node, current);
-        }
-    }
-}
-
-size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const {
-    backend      = backend == nullptr ? runtime_backend : backend;
-    size_t bytes = workspace_.bytes(backend);
-    if (backend == runtime_backend) {
-        const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend));
-        bytes                    = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes;
-        const size_t cut_bytes   = cut_cache_.resident_bytes(ggml_backend_get_device(backend));
-        bytes                    = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes;
-    }
-    return bytes;
-}
-
-void GGMLRunner::sync_runtime_residency() {
-    if (auto manager = residency_manager.lock()) {
-        manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
-                                          runtime_backend, retained_runtime_buffer_bytes());
-        for (auto backend : extra_runtime_backends) {
-            manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
-                                              backend, retained_runtime_buffer_bytes(backend));
-        }
-    }
-}
-
-std::optional<sd::Tensor<float>> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) {
-    if (tensor == nullptr) {
-        LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
-        return std::nullopt;
-    }
-    if (tensor->type != GGML_TYPE_F32) {
-        LOG_ERROR("%s %s tensor type mismatch: got %s",
-                  get_desc().c_str(),
-                  label,
-                  ggml_type_name(tensor->type));
-        return std::nullopt;
-    }
-    ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
-    if (buf == nullptr) {
-        LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
-                  get_desc().c_str(),
-                  label,
-                  tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
-                  ggml_op_name(tensor->op),
-                  tensor->buffer,
-                  tensor->view_src,
-                  tensor->view_src ? tensor->view_src->buffer : nullptr,
-                  tensor->data);
-        return std::nullopt;
-    }
-
-    return sd::make_sd_tensor_from_ggml<float>(tensor);
-}
-
-void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
-    GGML_ASSERT(gf != nullptr);
-    std::unordered_set<const ggml_tensor*> graph_tensor_set;
-    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
-    const int n_nodes = ggml_graph_n_nodes(gf);
-    graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
-    for (int i = 0; i < n_leafs; ++i) {
-        graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
-    }
-    for (int i = 0; i < n_nodes; ++i) {
-        graph_tensor_set.insert(ggml_graph_node(gf, i));
-    }
-
-    for (auto& kv : backend_tensor_data_map) {
-        auto tensor = kv.first;
-        auto data   = kv.second;
-        if (tensor == nullptr || data == nullptr) {
-            continue;
-        }
-        const char* name = ggml_get_name(tensor);
-        if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
-            continue;
-        }
-        if (tensor->buffer == nullptr) {
-            LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
-                     get_desc().c_str(),
-                     name != nullptr ? name : "",
-                     (long long)tensor->ne[0],
-                     (long long)tensor->ne[1],
-                     (long long)tensor->ne[2],
-                     (long long)tensor->ne[3],
-                     ggml_type_name(tensor->type));
-            continue;
-        }
-
-        ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
-        if (buf == nullptr) {
-            LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
-                     get_desc().c_str(),
-                     tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
-                     tensor ? ggml_op_name(tensor->op) : "<null>",
-                     data,
-                     tensor ? tensor->view_src : nullptr,
-                     (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
-            continue;
-        }
-
-        ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
-    }
-
-    if (clear_after_copy) {
-        backend_tensor_data_map.clear();
-    }
-}
-
-const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) {
-    GGML_ASSERT(gf != nullptr);
-    return sd::ggml_graph_cut::resolve_plan(runtime_backend,
-                                            gf,
-                                            &graph_cut_plan_cache_,
-                                            params_tensor_set_,
-                                            get_desc().c_str());
-}
-
-const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) {
-    return resolve_graph_cut_plan(gf);
-}
-
-bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
-    graph_cut_layer_split_node_assignments_.clear();
-    if (!graph_cut_layer_split_enabled) {
-        return true;
-    }
-    if (!is_multi_device()) {
-        LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
-        return false;
-    }
-
-    const auto& plan = resolve_graph_cut_layer_split_plan(gf);
-    if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
-        auto manager = residency_manager.lock();
-        if (manager == nullptr) {
-            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
-            return false;
-        }
-        std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
-        if (!graph_params.empty() &&
-            !manager->assign_compute_backend(graph_params, runtime_backend)) {
-            LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
-                      get_desc().c_str(),
-                      sd::layer_split_backend_device_display_name(runtime_backend).c_str());
-            return false;
-        }
-        for (ggml_tensor* param : graph_params) {
-            if (param != nullptr) {
-                graph_cut_layer_split_assignments_[param] = runtime_backend;
-            }
-        }
-        const int n_nodes = ggml_graph_n_nodes(gf);
-        for (int i = 0; i < n_nodes; i++) {
-            ggml_tensor* node = ggml_graph_node(gf, i);
-            if (node != nullptr) {
-                graph_cut_layer_split_node_assignments_[node] = runtime_backend;
-            }
-        }
-        if (!graph_cut_layer_split_primary_notice_logged_) {
-            LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
-                     get_desc().c_str(),
-                     sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
-                     graph_params.size());
-            graph_cut_layer_split_primary_notice_logged_ = true;
-        } else {
-            LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
-                        get_desc().c_str(),
-                        sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
-                        graph_params.size());
-        }
-        return true;
-    }
-
-    std::vector<ggml_backend_t> split_backends;
-    split_backends.reserve(extra_runtime_backends.size() + 1);
-    split_backends.push_back(runtime_backend);
-    for (ggml_backend_t backend : extra_runtime_backends) {
-        if (backend != nullptr) {
-            split_backends.push_back(backend);
-        }
-    }
-
-    auto manager = residency_manager.lock();
-    if (manager == nullptr) {
-        LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
-        return false;
-    }
-
-    sd::GraphCutLayerSplitAssignment assignment;
-    auto canonicalize_param = [this](ggml_tensor* tensor) {
-        return canonical_param_tensor(tensor);
-    };
-    if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
-                                             gf,
-                                             plan,
-                                             split_backends,
-                                             graph_cut_layer_split_backend_vram_limits_,
-                                             max_graph_vram_bytes,
-                                             graph_cut_layer_split_assignments_,
-                                             canonicalize_param,
-                                             &assignment)) {
-        return false;
-    }
-
-    for (size_t i = 0; i < split_backends.size(); i++) {
-        if (assignment.tensors_by_backend[i].empty()) {
-            continue;
-        }
-        if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
-            LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
-                      get_desc().c_str(),
-                      sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
-            return false;
-        }
-    }
-
-    graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
-    sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
-
-    return true;
-}
-
 bool GGMLRunner::runner_start() {
-    if (runner_started_) {
-        return true;
-    }
-    sd_backend_set_buffer_cache_enabled(runtime_backend, true);
-    for (auto backend : extra_runtime_backends) {
-        sd_backend_set_buffer_cache_enabled(backend, true);
-    }
-    cache_.clear();
-    workspace_.set_extra_backends(extra_runtime_backends);
-    if (auto manager = residency_manager.lock()) {
-        manager->set_workspace_reclaimer(reinterpret_cast<uintptr_t>(this), [this]() {
-            if (!workspace_.release()) {
-                return false;
-            }
-            sync_runtime_residency();
-            return true;
-        });
-    }
-    runner_started_ = true;
-    return true;
+    return execution_engine_->start(execution_session_);
 }
 
 void GGMLRunner::runner_end() {
     GGML_ASSERT(!graph_active_);
-    if (!runner_started_) {
-        return;
-    }
-    workspace_.release();
+    execution_engine_->end(execution_session_);
     free_compute_ctx();
     compute_ctx_buffer_.reset();
-    cache_.clear();
-    cut_cache_.clear();
-    logged_compute_bytes_.clear();
-    logged_segment_count_ = 0;
-    if (auto manager = residency_manager.lock()) {
-        manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
-        std::vector<ggml_tensor*> tensors;
-        for (auto tensor : params_tensor_set_) {
-            auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
-            if (parameter != nullptr)
-                tensors.push_back(parameter);
-        }
-        manager->evict_compute_backend_params(tensors);
-        manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
-    }
-    sd_backend_clear_buffer_cache(runtime_backend);
-    for (auto backend : extra_runtime_backends) {
-        sd_backend_clear_buffer_cache(backend);
-    }
-    runner_started_ = false;
 }
 
 GGMLRunner::GGMLRunner(ggml_backend_t backend,
                        std::shared_ptr<DeviceResidencyManager> manager)
-    : runtime_backend(backend),
-      cache_(backend),
-      cut_cache_(backend),
-      workspace_(backend),
+    : execution_engine_(manager != nullptr ? manager->execution_engine() : std::make_shared<sd::ExecutionEngine>()),
+      execution_session_(execution_engine_->create_session(backend, manager)),
+      runtime_backend(backend),
       residency_manager(manager) {
     GGML_ASSERT(runtime_backend != nullptr);
     alloc_params_ctx();
@@ -554,6 +181,7 @@ GGMLRunner::GGMLRunner(ggml_backend_t backend,
 
 GGMLRunner::~GGMLRunner() {
     runner_end();
+    execution_engine_->release_session(execution_session_);
     free_compute_ctx();
     free_params_ctx();
 }
@@ -591,8 +219,7 @@ void GGMLRunner::reset_compute_ctx() {
 }
 
 void GGMLRunner::free_cache_ctx_and_buffer() {
-    cache_.clear();
-    sync_runtime_residency();
+    execution_engine_->clear_cache(execution_session_);
 }
 
 void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
@@ -626,7 +253,7 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph*
     if (tensor != nullptr) {
         ggml_set_output(tensor);
     }
-    cache_.stage(name, tensor);
+    execution_session_.cache().stage(name, tensor);
     if (graph != nullptr && tensor != nullptr) {
         // Schedule the cache output here so its source can be reused before graph end.
         ggml_build_forward_expand(graph, tensor);
@@ -662,439 +289,54 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
         GGMLRunner& runner;
         const bool& success;
         ~GraphEndGuard() {
-            if (!runner.workspace_.segment_end()) {
-                runner.last_compute_status_ = GGML_STATUS_FAILED;
+            if (!success) {
+                runner.execution_engine_->discard_graph(runner.execution_session_);
             }
-            runner.cache_.graph_end(false);
-            runner.cut_cache_.clear();
             runner.free_compute_ctx();
             runner.graph_active_ = false;
-            if (!success) {
-                runner.workspace_.release();
-            }
-            runner.sync_runtime_residency();
         }
     } graph_guard{*this, success};
 
-    ggml_cgraph* graph = nullptr;
-    if (!prepare_compute_graph(get_graph, &graph)) {
-        return std::nullopt;
-    }
-    params_tensor_set_dirty_ = true;
-    rebuild_params_tensor_set();
-    if (auto manager = residency_manager.lock()) {
-        for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) {
-            auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i));
-            if (parameter != nullptr)
-                params_tensor_set_.insert(parameter);
-        }
-    }
-    std::optional<sd::Tensor<float>> output;
     try {
-        output = execute_graph(graph, n_threads, no_return, read_outputs);
+        ggml_cgraph* graph = nullptr;
+        if (!prepare_compute_graph(get_graph, &graph)) {
+            return std::nullopt;
+        }
+        params_tensor_set_dirty_ = true;
+        rebuild_params_tensor_set();
+        sd::ComputeTask task;
+        task.description     = get_desc();
+        task.graph           = graph;
+        task.output          = no_return ? nullptr : ggml_get_tensor(compute_ctx, final_result_name.c_str());
+        task.weights         = collect_used_param_tensors(graph);
+        task.inputs          = backend_tensor_data_map;
+        task.debug_tensors   = debug_tensors;
+        task.n_threads       = n_threads;
+        task.no_return       = no_return;
+        task.read_outputs    = read_outputs;
+        auto result          = execution_engine_->execute(execution_session_, task);
+        last_compute_status_ = result.status;
+        success              = result.status == GGML_STATUS_SUCCESS;
+        if (success) {
+            return std::move(result.output);
+        }
     } catch (const std::bad_alloc&) {
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         LOG_ERROR("%s graph allocation failed", get_desc().c_str());
-        return std::nullopt;
     } catch (const std::exception& error) {
-        last_compute_status_ = GGML_STATUS_FAILED;
-        LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
-                  ggml_backend_name(runtime_backend), error.what());
-        return std::nullopt;
+        LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what());
     }
-    success = output.has_value();
-    if (success) {
-        cache_.graph_end(true);
-        last_compute_status_ = GGML_STATUS_SUCCESS;
-    }
-    return output;
+    return std::nullopt;
 }
 
 void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) {
-    graph_cut_layer_split_enabled = enabled;
-    if (!enabled) {
-        graph_cut_layer_split_assignments_.clear();
-        graph_cut_layer_split_node_assignments_.clear();
-        graph_cut_layer_split_primary_notice_logged_ = false;
-    }
+    execution_session_.set_layer_split_enabled(enabled);
 }
 
 void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
-    graph_cut_layer_split_backend_vram_limits_ = limits;
-    graph_cut_layer_split_assignments_.clear();
-    graph_cut_layer_split_node_assignments_.clear();
-    graph_cut_layer_split_primary_notice_logged_ = false;
+    execution_session_.set_backend_limits(limits);
 }
 
 void GGMLRunner::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
-    extra_runtime_backends.clear();
-    for (ggml_backend_t backend : backends) {
-        if (backend == nullptr || backend == runtime_backend) {
-            continue;
-        }
-        if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
-            extra_runtime_backends.end()) {
-            extra_runtime_backends.push_back(backend);
-        }
-    }
-    workspace_.set_extra_backends(extra_runtime_backends);
-    graph_cut_layer_split_assignments_.clear();
-    graph_cut_layer_split_node_assignments_.clear();
-    graph_cut_layer_split_primary_notice_logged_ = false;
-}
-
-static size_t add_bytes(size_t a, size_t b) {
-    return b > SIZE_MAX - a ? SIZE_MAX : a + b;
-}
-
-ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs) {
-    auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
-        if (external_outputs != nullptr && external_outputs->count(tensor)) {
-            return runtime_backend;
-        }
-        if (!params_tensor_set_.count(tensor)) {
-            return nullptr;
-        }
-        auto placement = graph_cut_layer_split_assignments_.find(tensor);
-        return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second;
-    };
-    auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
-        pin_multi_device_nodes(scheduler, copy, graph);
-    };
-    return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
-}
-
-std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
-    const std::vector<BackendBufferSize>& sizes,
-    size_t pending_cache_bytes) const {
-    std::vector<DeviceMemoryRequest> requests;
-    for (const auto& size : sizes) {
-        const size_t retained    = retained_runtime_buffer_bytes(size.backend);
-        const size_t reusable    = workspace_.bytes(size.backend);
-        const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0;
-        const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
-        size_t limit             = max_graph_vram_bytes;
-        if (is_multi_device()) {
-            size_t index = 0;
-            if (size.backend != runtime_backend) {
-                auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend);
-                index         = static_cast<size_t>(position - extra_runtime_backends.begin()) + 1;
-            }
-            if (index < graph_cut_layer_split_backend_vram_limits_.size()) {
-                limit = graph_cut_layer_split_backend_vram_limits_[index];
-            }
-        }
-        requests.push_back({size.backend, reinterpret_cast<uintptr_t>(this), pending,
-                            retained, limit});
-    }
-    return requests;
-}
-
-bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
-                      const std::vector<ggml_tensor*>& params) const {
-    auto manager = residency_manager.lock();
-    if (manager == nullptr) {
-        return params.empty();
-    }
-    for (const auto& request : requests) {
-        if (!manager->fits_compute_backend_capacity(request, params)) {
-            return false;
-        }
-    }
-    return true;
-}
-
-bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
-    if (sd_backend_is_cpu(runtime_backend)) {
-        sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
-    }
-    if (workspace_.cpu_backend() != nullptr) {
-        sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
-    }
-    auto scheduler = workspace_.scheduler();
-    GGML_ASSERT(scheduler != nullptr);
-    ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
-#ifndef SD_USE_UPSTREAM_GGML
-    struct BeforeComputeCallback {
-        ggml_backend_t backend;
-        const std::function<void()>& callback;
-        bool invoked = false;
-    } before_compute{runtime_backend, on_before_compute};
-    ggml_backend_sched_set_compute_callback(
-        scheduler, [](ggml_backend_t backend, void* data) {
-            auto& hook = *static_cast<BeforeComputeCallback*>(data);
-            if (backend == hook.backend && !hook.invoked) {
-                hook.invoked = true;
-                hook.callback();
-            }
-        },
-        &before_compute);
-    const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
-    ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr);
-#else
-    on_before_compute();
-    const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
-#endif
-    workspace_.synchronize();
-    if (status != GGML_STATUS_SUCCESS) {
-        last_compute_status_ = status;
-        LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
-        return false;
-    }
-    const std::string description = get_desc();
-    if (!debug_tensors.empty()) {
-        std::unordered_set<const ggml_tensor*> graph_tensors;
-        const int leaf_count = ggml_graph_cut::leaf_count(graph);
-        const int node_count = ggml_graph_n_nodes(graph);
-        graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
-        for (int index = 0; index < leaf_count; ++index) {
-            graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
-        }
-        for (int index = 0; index < node_count; ++index) {
-            graph_tensors.insert(ggml_graph_node(graph, index));
-        }
-
-        for (const auto& entry : debug_tensors) {
-            ggml_tensor* tensor = entry.first;
-            if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
-                continue;
-            }
-            ggml_backend_buffer_t buffer =
-                tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
-            if (buffer == nullptr) {
-                LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
-                         description.c_str(),
-                         entry.second.c_str());
-                continue;
-            }
-            if (tensor->type != GGML_TYPE_F32) {
-                LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
-                         description.c_str(),
-                         entry.second.c_str(),
-                         ggml_type_name(tensor->type));
-                continue;
-            }
-            auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
-            print_sd_tensor(debug_tensor, false, entry.second.c_str());
-        }
-    }
-
-    return true;
-}
-
-std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
-    if (!assign_graph_cut_layer_split_backends(graph)) {
-        return std::nullopt;
-    }
-    const auto params       = collect_used_param_tensors(graph);
-    const auto& cached_plan = resolve_graph_cut_plan(graph);
-    auto manager            = residency_manager.lock();
-    const bool can_segment  = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
-                              manager != nullptr && manager->segmented_compute_enabled() &&
-                              cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1;
-    // An empty workspace request is a lower bound on the full graph's memory demand.
-    bool segmented = can_segment && !fits(memory_requests({{runtime_backend, 0}}, 0), params);
-    ComputeWorkspace::Measurement full_measurement;
-    if (!segmented) {
-        full_measurement = measure(graph);
-        if (full_measurement.buffers.empty()) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return std::nullopt;
-        }
-        if (can_segment) {
-            // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
-            constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
-            auto requests                      = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
-            for (auto& request : requests) {
-                request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
-            }
-            segmented = !fits(requests, params);
-        }
-    }
-    ggml_graph_cut::Plan monolithic_plan;
-    if (!segmented) {
-        monolithic_plan.segments.emplace_back();
-        auto& segment               = monolithic_plan.segments.back();
-        segment.group_name          = "graph";
-        segment.compute_buffer_size = cached_plan.compute_buffer_size;
-        segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
-        segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
-        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
-            segment.internal_node_indices.push_back(i);
-        }
-        for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
-            auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
-            ggml_graph_cut::Segment::InputRef input;
-            input.leaf_index = i;
-            input.type       = canonical_param_tensor(tensor) != nullptr
-                                   ? ggml_graph_cut::Segment::INPUT_PARAM
-                                   : ggml_graph_cut::Segment::INPUT_EXTERNAL;
-            segment.input_refs.push_back(input);
-        }
-    }
-    const auto& plan            = segmented ? cached_plan : monolithic_plan;
-    const bool segments_changed = plan.segments.size() != logged_segment_count_;
-    if (segments_changed && (segmented || logged_segment_count_ > 1)) {
-        LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(),
-                    plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
-    }
-    SegmentGraphBindings bindings(cut_cache_, plan, graph);
-    SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast<uintptr_t>(this),
-                                  graph, plan, params_tensor_set_,
-                                  segmented && manager != nullptr && manager->prefetch_enabled());
-
-    std::map<ggml_backend_t, size_t> peak_compute_bytes;
-    auto track_compute_buffer = [&](ggml_backend_t backend) {
-        if (backend != nullptr) {
-            auto& peak = peak_compute_bytes[backend];
-            peak       = std::max(peak, workspace_.bytes(backend));
-        }
-    };
-    std::optional<Tensor<float>> output = Tensor<float>();
-    for (size_t index = 0; index < plan.segments.size(); ++index) {
-        const auto& segment = plan.segments[index];
-        const bool last     = index + 1 == plan.segments.size();
-        auto fail_segment   = [&](const char* phase) {
-            LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(),
-                      index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
-            return std::nullopt;
-        };
-        cut_cache_.prune(segment.live_cut_names);
-        bindings.reset(segment);
-        if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) {
-            return fail_segment("input binding");
-        }
-        ggml_context* segment_context = nullptr;
-        auto segment_graph            = segmented
-                                            ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
-                                            : graph;
-        struct SegmentCleanup {
-            GGMLRunner& runner;
-            SegmentWeightPipeline& weights;
-            SegmentGraphBindings& bindings;
-            ggml_context* context;
-            ~SegmentCleanup() {
-                if (!runner.workspace_.segment_end()) {
-                    runner.last_compute_status_ = GGML_STATUS_FAILED;
-                }
-                bindings.restore();
-                weights.segment_end();
-                ggml_free(context);
-                runner.sync_runtime_residency();
-            }
-        } segment_cleanup{*this, weights, bindings, segment_context};
-
-        std::unordered_set<const ggml_tensor*> cache_outputs;
-        if (segmented) {
-            for (int output_index : segment.output_node_indices) {
-                auto tensor = ggml_graph_node(graph, output_index);
-                if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) &&
-                    GraphCutTensorCache::can_bind_output(runtime_backend, tensor)) {
-                    cache_outputs.insert(tensor);
-                }
-            }
-        }
-        auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
-        if (!workspace_.prepare(measurement)) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("workspace preparation");
-        }
-        const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
-        size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph));
-        auto ensure_capacity   = [&]() {
-            sync_runtime_residency();
-            auto requests = memory_requests(measurement.buffers, new_cache_bytes);
-            if (fits(requests, weights.params(index))) {
-                return true;
-            }
-            if (workspace_.release_excess(measurement)) {
-                sync_runtime_residency();
-                requests = memory_requests(measurement.buffers, new_cache_bytes);
-            }
-            const bool ready = weights.ensure_segment_capacity(index, requests);
-            if (!ready && manager != nullptr) {
-                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            }
-            return ready;
-        };
-        if (!weights.segment_start(index, ensure_capacity)) {
-            return fail_segment("weight preparation");
-        }
-        // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
-        if (!workspace_.prepare(measurement)) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("workspace preparation");
-        }
-        if (!ensure_capacity()) {
-            return fail_segment("workspace capacity check");
-        }
-        if (!cut_cache_.bind_outputs(cache_outputs)) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("output cache allocation");
-        }
-        new_cache_bytes = add_bytes(last ? 0 : cut_cache_.estimate_output_bytes(graph, segment), cache_.pending_bytes(segment_graph));
-        if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
-                pin_multi_device_nodes(scheduler, current);
-            })) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("workspace allocation");
-        }
-        for (const auto& size : measurement.buffers) {
-            track_compute_buffer(size.backend);
-        }
-        if (workspace_.scheduler() != nullptr) {
-            track_compute_buffer(workspace_.cpu_backend());
-        }
-        if (!ensure_capacity()) {
-            return fail_segment("allocated capacity check");
-        }
-        copy_data_to_backend_tensor(segment_graph, false);
-        auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
-        if (!prefetch_requests.empty()) {
-            weights.enqueue_next(index, prefetch_requests.front());
-        }
-        LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
-                  index + 1, plan.segments.size(), segment.group_name.c_str());
-        if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
-            return fail_segment("execution");
-        }
-        auto cache_status = cache_.capture(segment_graph);
-        if (cache_status == GGML_STATUS_SUCCESS) {
-            cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
-        }
-        if (cache_status != GGML_STATUS_SUCCESS) {
-            last_compute_status_ = cache_status;
-            return fail_segment("output caching");
-        }
-        sync_runtime_residency();
-        if (last) {
-            if (read_outputs && !read_outputs()) {
-                return fail_segment("output finalization");
-            }
-            if (!no_return) {
-                auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
-                output      = read_graph_tensor(result, "output");
-                if (!output.has_value()) {
-                    return fail_segment("output readback");
-                }
-            }
-        }
-        if (!workspace_.segment_end()) {
-            last_compute_status_ = GGML_STATUS_FAILED;
-            return fail_segment("workspace synchronization");
-        }
-        // Final outputs and their callbacks may still be views of consumed cuts.
-        cut_cache_.prune(segment.future_cut_names);
-    }
-    if (segments_changed || peak_compute_bytes != logged_compute_bytes_) {
-        for (const auto& entry : peak_compute_bytes) {
-            LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
-                        get_desc().c_str(), entry.second / (1024.0 * 1024.0),
-                        sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
-                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
-        }
-        logged_compute_bytes_ = std::move(peak_compute_bytes);
-        logged_segment_count_ = plan.segments.size();
-    }
-    return output;
+    execution_session_.set_runtime_backends(backends);
 }

src/core/ggml_runner.h

@@ -12,9 +12,7 @@
 #include <utility>
 #include <vector>
 
-#include "core/compute_workspace.h"
-#include "core/ggml_graph_cut.h"
-#include "core/runner_cache.h"
+#include "core/execution_engine.h"
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
 #include "device_residency_manager.h"
@@ -23,6 +21,8 @@
 #define MAX_PARAMS_TENSOR_NUM 32768
 #define MAX_GRAPH_SIZE 327680
 
+struct GGMLRunnerContext;
+
 struct WeightAdapter {
     struct ForwardParams {
         enum class op_type_t {
@@ -46,22 +46,26 @@ struct WeightAdapter {
             float scale     = 1.f;
         } conv2d;
     };
-    virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0;
-    virtual ggml_tensor* forward_with_lora(ggml_context* ctx,
-                                           ggml_backend_t backend,
+
+    virtual ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                                      ggml_tensor* weight,
+                                      const std::string& weight_name) = 0;
+
+    virtual ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx,
                                            ggml_tensor* x,
                                            ggml_tensor* w,
                                            ggml_tensor* b,
                                            const std::string& prefix,
-                                           ForwardParams forward_params)                                                              = 0;
-    virtual ggml_tensor* add_lora_to_output(ggml_context* ctx,
-                                            ggml_backend_t backend,
+                                           ForwardParams forward_params) = 0;
+
+    virtual ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx,
                                             ggml_tensor* x,
                                             ggml_tensor* w,
                                             ggml_tensor* output,
                                             const std::string& prefix,
-                                            ForwardParams forward_params)                                                             = 0;
-    virtual size_t get_extra_graph_size()                                                                                             = 0;
+                                            ForwardParams forward_params) = 0;
+
+    virtual size_t get_extra_graph_size() = 0;
 };
 
 struct GGMLRunnerContext {
@@ -124,52 +128,22 @@ struct GGMLRunnerContext {
     }
 };
 
-ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
-                                    ggml_tensor* q,
-                                    ggml_tensor* k,
-                                    ggml_tensor* v,
-                                    int64_t n_head,
-                                    ggml_tensor* mask     = nullptr,
-                                    bool skip_reshape     = false,
-                                    bool flash_attn       = false,
-                                    float kv_scale        = 1.f,
-                                    bool* used_flash_attn = nullptr);
-
 struct GGMLRunner {
 private:
-    std::map<ggml_backend_t, size_t> logged_compute_bytes_;
-    size_t logged_segment_count_     = 0;
+    std::shared_ptr<sd::ExecutionEngine> execution_engine_;
+    sd::ExecutionSession& execution_session_;
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
-    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
-    std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
-                                                     size_t pending_cache_bytes) const;
-    bool fits(const std::vector<DeviceMemoryRequest>& requests,
-              const std::vector<ggml_tensor*>& params) const;
-    bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
-    std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
-
 protected:
     typedef std::function<ggml_cgraph*()> get_graph_cb_t;
-    using GraphCutPlan = sd::ggml_graph_cut::Plan;
 
     ggml_backend_t runtime_backend = nullptr;
 
     ggml_context* params_ctx = nullptr;
 
-    sd::RunnerCache cache_;
-    sd::GraphCutTensorCache cut_cache_;
-    sd::ComputeWorkspace workspace_;
     ggml_context* compute_ctx = nullptr;
     std::unique_ptr<std::max_align_t[]> compute_ctx_buffer_;
-    bool runner_started_ = false;
-    bool graph_active_   = false;
-
-    size_t max_graph_vram_bytes        = 0;
-    bool graph_cut_layer_split_enabled = false;
-    std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
-
-    std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
+    bool graph_active_ = false;
 
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<DeviceResidencyManager> residency_manager;
@@ -194,11 +168,7 @@ protected:
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
 
-    sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
     std::unordered_set<const ggml_tensor*> params_tensor_set_;
-    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
-    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
-    bool graph_cut_layer_split_primary_notice_logged_ = false;
 
     template <typename T>
     static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -240,8 +210,6 @@ protected:
 
     std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
 
-    void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors);
-
     void prepare_build_in_tensor_before();
 
     void prepare_build_in_tensor_after(ggml_cgraph* gf);
@@ -253,39 +221,18 @@ protected:
     bool prepare_compute_graph(get_graph_cb_t get_graph,
                                ggml_cgraph** gf_out);
 
-    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
-
-    // Weightless ops have no scheduler anchor, so pin them to the most recent
-    // weight device. Views must stay unpinned or cross-device copies can be
-    // skipped for their consumers.
-    void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr);
-
-    bool is_multi_device() const {
-        return !extra_runtime_backends.empty();
-    }
-
     size_t reusable_compute_buffer_bytes() const {
-        return workspace_.bytes(runtime_backend);
+        return execution_session_.compute_bytes();
     }
 
-    size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const;
-
-    void sync_runtime_residency();
-
-    std::optional<sd::Tensor<float>> read_graph_tensor(ggml_tensor* tensor, const char* label);
-
-    void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
-
-    const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* gf);
-
-    const GraphCutPlan& resolve_graph_cut_layer_split_plan(ggml_cgraph* gf);
-
-    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
+    bool has_cache_tensors() const {
+        return !execution_session_.cache().empty();
+    }
 
 public:
     bool runner_start();
 
-    bool runner_started() const { return runner_started_; }
+    bool runner_started() const { return execution_session_.started(); }
 
     void runner_end();
 
@@ -336,7 +283,7 @@ public:
     void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
 
     ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
-        return cache_.get(name);
+        return execution_session_.cache().get(name);
     }
 
     std::optional<sd::Tensor<float>> compute(get_graph_cb_t get_graph,
@@ -354,7 +301,7 @@ public:
     void set_sage_attention_enabled(bool enabled) {
         if (sage_attn_enabled != enabled) {
             free_cache_ctx_and_buffer();
-            graph_cut_plan_cache_.graph_cut_plans.clear();
+            execution_session_.clear_plan_cache();
             sage_attn_enabled = enabled;
         }
     }
@@ -382,7 +329,7 @@ public:
     }
 
     void set_max_graph_vram_bytes(size_t max_vram_bytes) {
-        max_graph_vram_bytes = max_vram_bytes;
+        execution_session_.set_max_backend_bytes(max_vram_bytes);
     }
 
     void set_graph_cut_layer_split_enabled(bool enabled);

src/core/layer_split_partition.cpp

@@ -73,10 +73,13 @@ namespace sd {
         return backend_index == 0 ? primary_backend_vram_limit : 0;
     }
 
+    // Capacity of a device that reports no free memory and has no VRAM limit.
+    static constexpr int64_t graph_cut_layer_split_unknown_capacity = std::numeric_limits<int64_t>::max() / 4;
+
     static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
                                                                          const std::vector<size_t>& backend_vram_limits,
                                                                          size_t primary_backend_vram_limit) {
-        std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
+        std::vector<int64_t> capacities(backends.size(), graph_cut_layer_split_unknown_capacity);
         constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
         for (size_t i = 0; i < backends.size(); i++) {
             ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
@@ -97,6 +100,57 @@ namespace sd {
         return capacities;
     }
 
+    // Whether filling the devices in order, each up to `fraction` of its capacity, places every
+    // segment. Mirrors the placement loop in partition_graph_cut_layer_split().
+    static bool graph_cut_layer_split_fits(const std::vector<int64_t>& segment_bytes,
+                                           const std::vector<int64_t>& capacities,
+                                           double fraction) {
+        size_t backend = 0;
+        int64_t used   = 0;
+        for (int64_t bytes : segment_bytes) {
+            while (backend + 1 < capacities.size() && bytes > 0 &&
+                   used + bytes > (int64_t)((double)capacities[backend] * fraction)) {
+                backend++;
+                used = 0;
+            }
+            if (bytes > 0 && used + bytes > (int64_t)((double)capacities[backend] * fraction)) {
+                return false;
+            }
+            used += bytes;
+        }
+        return true;
+    }
+
+    // Fill limits that size the contiguous ranges in proportion to the device capacities: the same
+    // fraction of every capacity, as small as possible while all segments still fit. Filling the
+    // first device up to its capacity would leave only the fixed headroom for compute, which is too
+    // small for large video graphs. A fraction of 1 is plain capacity-based placement, so every
+    // layout that fits the capacities is still found.
+    static std::vector<int64_t> graph_cut_layer_split_fill_limits(const std::vector<int64_t>& segment_bytes,
+                                                                  const std::vector<int64_t>& capacities) {
+        const bool capacities_known = std::all_of(capacities.begin(), capacities.end(), [](int64_t capacity) {
+            return capacity < graph_cut_layer_split_unknown_capacity;
+        });
+        if (!capacities_known || !graph_cut_layer_split_fits(segment_bytes, capacities, 1.0)) {
+            return capacities;
+        }
+        double low  = 0.0;
+        double high = 1.0;
+        for (int i = 0; i < 32; i++) {
+            const double mid = (low + high) / 2;
+            if (graph_cut_layer_split_fits(segment_bytes, capacities, mid)) {
+                high = mid;
+            } else {
+                low = mid;
+            }
+        }
+        std::vector<int64_t> limits(capacities.size());
+        for (size_t i = 0; i < capacities.size(); i++) {
+            limits[i] = (int64_t)((double)capacities[i] * high);
+        }
+        return limits;
+    }
+
     bool partition_graph_cut_layer_split(const char* desc,
                                          ggml_cgraph* gf,
                                          const sd::ggml_graph_cut::Plan& plan,
@@ -151,6 +205,11 @@ namespace sd {
             return param_assignments.count(param) != 0;
         });
 
+        // Move to the next device at the proportional fill limit; the capacity stays the hard limit.
+        const std::vector<int64_t> backend_limits = reuse_assignments
+                                                        ? backend_capacities
+                                                        : graph_cut_layer_split_fill_limits(segment_param_bytes, backend_capacities);
+
         std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
         size_t current_backend = 0;
         int64_t current_used   = 0;
@@ -158,7 +217,7 @@ namespace sd {
             int64_t bytes = segment_param_bytes[seg_idx];
             while (!reuse_assignments && current_backend + 1 < split_backends.size() &&
                    bytes > 0 &&
-                   current_used + bytes > backend_capacities[current_backend]) {
+                   current_used + bytes > backend_limits[current_backend]) {
                 current_backend++;
                 current_used = 0;
             }

src/core/probe/flash_attn.cpp

@@ -9,6 +9,7 @@
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
 #include "core/probe/backend_probe.h"
 #include "core/util.h"
 
@@ -41,8 +42,11 @@ static void build_flash_attention_probe(BackendProbeGraph& graph, ggml_backend_t
     auto k              = ggml_dup_tensor(graph.ctx, q);
     auto v              = ggml_dup_tensor(graph.ctx, q);
     graph.inputs        = {q, k, v};
+    GGMLRunnerContext runner_ctx;
+    runner_ctx.ggml_ctx = graph.ctx;
+    runner_ctx.backend  = backend;
     bool used_flash     = false;
-    auto output         = ggml_ext_attention_ext(graph.ctx, backend, q, k, v, heads, nullptr, false, flash, 1.f, false, &used_flash);
+    auto output         = ggml_ext_attention_ext(&runner_ctx, q, k, v, heads, nullptr, false, flash, 1.f, &used_flash);
     graph.output        = flash && !used_flash ? nullptr : output;
 }
 

src/device_residency_manager.h

@@ -3,12 +3,17 @@
 
 #include <cstdint>
 #include <functional>
+#include <memory>
 #include <vector>
 
 #include "ggml-backend.h"
 
 struct ggml_tensor;
 
+namespace sd {
+    class ExecutionEngine;
+}
+
 enum class WeightPrefetchResult {
     Scheduled,
     AlreadyResident,
@@ -36,13 +41,45 @@ struct DeviceMemoryRequest {
     }
 };
 
+struct DeviceMemoryCapacity {
+    size_t required_device_bytes  = 0;
+    size_t required_budget_bytes  = 0;
+    size_t available_device_bytes = SIZE_MAX;
+    size_t available_budget_bytes = SIZE_MAX;
+    bool requires_reclamation     = false;
+    bool valid                    = true;
+    bool wrapped_device_report    = false;
+
+    bool fits(bool after_reclamation = false) const {
+        return valid && !requires_reclamation &&
+               (after_reclamation || !wrapped_device_report) &&
+               required_device_bytes <= available_device_bytes &&
+               required_budget_bytes <= available_budget_bytes;
+    }
+};
+
+struct WeightResidencyBlock {
+    ggml_backend_buffer_t buffer = nullptr;
+    ggml_backend_t backend       = nullptr;
+    std::vector<ggml_tensor*> tensors;
+    size_t bytes            = 0;
+    uint64_t last_use_epoch = 0;
+    bool requires_demotion  = false;
+};
+
 struct DeviceResidencyManager {
+private:
+    std::shared_ptr<sd::ExecutionEngine> execution_engine_;
+
+public:
     virtual ~DeviceResidencyManager() = default;
 
-    virtual bool segmented_compute_enabled() const                                          = 0;
-    virtual bool prefetch_enabled() const                                                   = 0;
-    virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) = 0;
-    virtual void remove_runtime_owner(uintptr_t owner_id)                                   = 0;
+    std::shared_ptr<sd::ExecutionEngine> execution_engine();
+
+    virtual bool segmented_compute_enabled() const          = 0;
+    virtual bool prefetch_enabled() const                   = 0;
+    virtual void register_runtime_owner(uintptr_t owner_id) = 0;
+    virtual void remove_runtime_owner(uintptr_t owner_id)   = 0;
     // Capacity requests select their backend's weights; protection spans all backends.
     virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
                                                const std::vector<ggml_tensor*>& required_params) const = 0;
@@ -62,6 +99,18 @@ struct DeviceResidencyManager {
         const std::vector<ggml_tensor*>& required_params,
         const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
         const std::vector<ggml_tensor*>& protected_params) = 0;
+    virtual DeviceMemoryCapacity inspect_compute_backend_capacity(
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& params,
+        bool log_details = false) const {
+        DeviceMemoryCapacity result;
+        result.valid = false;
+        return result;
+    }
+    virtual std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
+        ggml_backend_t backend,
+        const std::vector<ggml_tensor*>& protected_params) const { return {}; }
+    virtual bool evict_compute_backend_block(const WeightResidencyBlock& block) { return false; }
     virtual WeightPrefetchResult prefetch_params(
         uintptr_t owner_id,
         const std::vector<ggml_tensor*>& tensors) = 0;

src/model.h

@@ -71,6 +71,7 @@ enum SDVersion {
     VERSION_FIBO2,
     VERSION_CHROMA,
     VERSION_P_H3_IMAGE,
+    VERSION_MINIMAX_H3_IMAGE,
     VERSION_COUNT,
 };
 
@@ -167,7 +168,11 @@ static inline bool sd_version_is_ltxav(SDVersion version) {
 }
 
 static inline bool sd_version_is_minimax_h3(SDVersion version) {
-    return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE;
+    return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE || version == VERSION_MINIMAX_H3_IMAGE;
+}
+
+static inline bool sd_version_is_minimax_h3_image(SDVersion version) {
+    return version == VERSION_P_H3_IMAGE || version == VERSION_MINIMAX_H3_IMAGE;
 }
 
 static inline bool sd_version_is_wan(SDVersion version) {

src/model/adapter/lora.hpp

@@ -218,7 +218,8 @@ struct LoraModel : public GGMLRunner {
         lora_tensors = std::move(new_lora_tensors);
     }
 
-    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -241,21 +242,21 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lora_up_name);
             if (iter != lora_tensors.end()) {
-                lora_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_up = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (swapped_lora_up_tensors.count(lora_up_name)) {
-                    auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1);
+                    auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1);
                     lora_up     = ggml_concat(ctx, halves[1], halves[0], 1);
                 }
             }
 
             iter = lora_tensors.find(lora_mid_name);
             if (iter != lora_tensors.end()) {
-                lora_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(lora_down_name);
             if (iter != lora_tensors.end()) {
-                lora_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             if (lora_up == nullptr || lora_down == nullptr) {
@@ -288,8 +289,8 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            auto curr_updown = ggml_ext_merge_lora(ctx, lora_down, lora_up, lora_mid);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            auto curr_updown = ggml_ext_merge_lora(runner_ctx, lora_down, lora_up, lora_mid);
+            curr_updown      = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -302,7 +303,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -319,7 +321,7 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(diff_name);
             if (iter != lora_tensors.end()) {
-                curr_updown = ggml_ext_cast_f32(ctx, backend, iter->second);
+                curr_updown = ggml_ext_cast_f32(runner_ctx, iter->second);
             } else {
                 break;
             }
@@ -329,7 +331,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             scale_value *= multiplier;
 
-            curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -342,7 +344,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -370,17 +373,17 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(hada_1_down_name);
             if (iter != lora_tensors.end()) {
-                hada_1_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_up_name);
             if (iter != lora_tensors.end()) {
-                hada_1_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_up = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (hada_1_up != nullptr) {
                     hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
                 }
@@ -388,17 +391,17 @@ struct LoraModel : public GGMLRunner {
 
             iter = lora_tensors.find(hada_2_down_name);
             if (iter != lora_tensors.end()) {
-                hada_2_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_up_name);
             if (iter != lora_tensors.end()) {
-                hada_2_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_up = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (hada_2_up != nullptr) {
                     hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
                 }
@@ -435,10 +438,10 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
-            ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
+            ggml_tensor* updown_1 = ggml_ext_merge_lora(runner_ctx, hada_1_down, hada_1_up, hada_1_mid);
+            ggml_tensor* updown_2 = ggml_ext_merge_lora(runner_ctx, hada_2_down, hada_2_up, hada_2_mid);
             auto curr_updown      = ggml_mul_inplace(ctx, updown_1, updown_2);
-            curr_updown           = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown           = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
             if (updown == nullptr) {
                 updown = curr_updown;
             } else {
@@ -449,7 +452,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -476,24 +480,24 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lokr_w1_name);
             if (iter != lora_tensors.end()) {
-                lokr_w1 = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lokr_w1 = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(lokr_w2_name);
             if (iter != lora_tensors.end()) {
-                lokr_w2 = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lokr_w2 = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             int64_t rank = 1;
             if (lokr_w1 == nullptr) {
                 iter = lora_tensors.find(lokr_w1_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_a = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w1_a = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w1_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_b = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w1_b = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 if (lokr_w1_a == nullptr || lokr_w1_b == nullptr) {
@@ -502,18 +506,18 @@ struct LoraModel : public GGMLRunner {
 
                 rank = lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1];
 
-                lokr_w1 = ggml_ext_merge_lora(ctx, lokr_w1_b, lokr_w1_a);
+                lokr_w1 = ggml_ext_merge_lora(runner_ctx, lokr_w1_b, lokr_w1_a);
             }
 
             if (lokr_w2 == nullptr) {
                 iter = lora_tensors.find(lokr_w2_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_a = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w2_a = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w2_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_b = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w2_b = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 if (lokr_w2_a == nullptr || lokr_w2_b == nullptr) {
@@ -522,7 +526,7 @@ struct LoraModel : public GGMLRunner {
 
                 rank = lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1];
 
-                lokr_w2 = ggml_ext_merge_lora(ctx, lokr_w2_b, lokr_w2_a);
+                lokr_w2 = ggml_ext_merge_lora(runner_ctx, lokr_w2_b, lokr_w2_a);
             }
 
             if (!lokr_w1_a) {
@@ -553,8 +557,8 @@ struct LoraModel : public GGMLRunner {
 
             scale_value *= multiplier;
 
-            auto curr_updown = ggml_ext_kronecker(ctx, lokr_w1, lokr_w2);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            auto curr_updown = ggml_ext_kronecker(runner_ctx, lokr_w1, lokr_w2);
+            curr_updown      = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -566,29 +570,33 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_backend_t backend, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
+    ggml_tensor* get_weight_diff(const std::string& model_tensor_name,
+                                 GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* model_tensor,
+                                 bool with_lora_and_lokr = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // lora
         ggml_tensor* diff = nullptr;
         if (with_lora_and_lokr) {
-            diff = get_lora_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_lora_weight_diff(model_tensor_name, runner_ctx);
         }
         // diff
         if (diff == nullptr) {
-            diff = get_raw_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_raw_weight_diff(model_tensor_name, runner_ctx);
         }
         // loha
         if (diff == nullptr) {
-            diff = get_loha_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_loha_weight_diff(model_tensor_name, runner_ctx);
         }
         // lokr
         if (diff == nullptr && with_lora_and_lokr) {
-            diff = get_lokr_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_lokr_weight_diff(model_tensor_name, runner_ctx);
         }
         if (diff != nullptr) {
             if (ggml_nelements(diff) < ggml_nelements(model_tensor)) {
                 if (ggml_n_dims(diff) == 2 && ggml_n_dims(model_tensor) == 2 && diff->ne[0] == model_tensor->ne[0]) {
                     LOG_WARN("pad for %s", model_tensor_name.c_str());
-                    auto pad_tensor = ggml_ext_zeros(ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1);
+                    auto pad_tensor = ggml_ext_zeros(runner_ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1);
                     diff            = ggml_concat(ctx, diff, pad_tensor, 1);
                 }
             }
@@ -619,12 +627,12 @@ struct LoraModel : public GGMLRunner {
         return diff;
     }
 
-    ggml_tensor* get_out_diff(ggml_context* ctx,
-                              ggml_backend_t backend,
+    ggml_tensor* get_out_diff(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* model_weight,
                               WeightAdapter::ForwardParams forward_params,
                               const std::string& model_tensor_name) {
+        auto ctx              = runner_ctx->ggml_ctx;
         ggml_tensor* out_diff = nullptr;
         int index             = 0;
 
@@ -712,7 +720,7 @@ struct LoraModel : public GGMLRunner {
                 }
                 scale_value *= multiplier;
 
-                auto curr_out_diff = ggml_ext_lokr_forward(ctx, backend, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
+                auto curr_out_diff = ggml_ext_lokr_forward(runner_ctx, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
                 if (out_diff == nullptr) {
                     out_diff = curr_out_diff;
                 } else {
@@ -761,7 +769,7 @@ struct LoraModel : public GGMLRunner {
             if (iter != lora_tensors.end()) {
                 lora_up = iter->second;
                 if (swapped_lora_up_tensors.count(lora_up_name)) {
-                    auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1);
+                    auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1);
                     lora_up     = ggml_concat(ctx, halves[1], halves[0], 1);
                 }
                 if (is_conv2d && lora_up->type != GGML_TYPE_F16) {
@@ -850,13 +858,13 @@ struct LoraModel : public GGMLRunner {
 
             ggml_tensor* lx;
             if (!is_conv2d) {
-                lx = ggml_ext_linear(ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                lx = ggml_ext_linear(runner_ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
                 if (lora_mid) {
-                    lx = ggml_ext_linear(ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                    lx = ggml_ext_linear(runner_ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
                 }
-                lx = ggml_ext_linear(ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                lx = ggml_ext_linear(runner_ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
             } else {  // OP_CONV2D
-                lx = ggml_ext_conv_2d(ctx,
+                lx = ggml_ext_conv_2d(runner_ctx,
                                       x,
                                       lora_down,
                                       nullptr,
@@ -871,7 +879,7 @@ struct LoraModel : public GGMLRunner {
                                       forward_params.conv2d.circular_y,
                                       forward_params.conv2d.scale);
                 if (lora_mid) {
-                    lx = ggml_ext_conv_2d(ctx,
+                    lx = ggml_ext_conv_2d(runner_ctx,
                                           lx,
                                           lora_mid,
                                           nullptr,
@@ -886,7 +894,7 @@ struct LoraModel : public GGMLRunner {
                                           forward_params.conv2d.circular_y,
                                           forward_params.conv2d.scale);
                 }
-                lx = ggml_ext_conv_2d(ctx,
+                lx = ggml_ext_conv_2d(runner_ctx,
                                       lx,
                                       lora_up,
                                       nullptr,
@@ -902,7 +910,7 @@ struct LoraModel : public GGMLRunner {
                                       forward_params.conv2d.scale);
             }
 
-            auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
+            auto curr_out_diff = ggml_ext_scale(runner_ctx, lx, scale_value, true);
             if (out_diff == nullptr) {
                 out_diff = curr_out_diff;
             } else {
@@ -951,6 +959,7 @@ struct LoraModel : public GGMLRunner {
         size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
         ggml_cgraph* gf        = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
 
+        auto runner_ctx = get_context();
         preprocess_lora_tensors(model_tensor_names);
 
         original_tensor_to_final_tensor.clear();
@@ -961,7 +970,7 @@ struct LoraModel : public GGMLRunner {
             ggml_tensor* model_tensor     = it.second;
 
             // lora
-            ggml_tensor* diff = get_weight_diff(model_tensor_name, runtime_backend, compute_ctx, model_tensor);
+            ggml_tensor* diff = get_weight_diff(model_tensor_name, &runner_ctx, model_tensor);
             if (diff == nullptr) {
                 continue;
             }
@@ -974,7 +983,7 @@ struct LoraModel : public GGMLRunner {
 
             ggml_tensor* final_tensor;
             if (model_tensor->type != GGML_TYPE_F32 && model_tensor->type != GGML_TYPE_F16) {
-                final_tensor = ggml_ext_cast_f32(compute_ctx, runtime_backend, model_tensor);
+                final_tensor = ggml_ext_cast_f32(&runner_ctx, model_tensor);
                 final_tensor = ggml_add_inplace(compute_ctx, final_tensor, diff);
                 final_tensor = ggml_cpy(compute_ctx, final_tensor, model_tensor);
             } else {
@@ -1061,41 +1070,47 @@ public:
         : lora_models(lora_models) {
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
+    ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                              ggml_tensor* weight,
+                              const std::string& weight_name,
+                              bool with_lora_and_lokr) {
+        auto ctx = runner_ctx->ggml_ctx;
         for (auto& lora_model : lora_models) {
-            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, backend, ctx, weight, with_lora_and_lokr);
+            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, runner_ctx, weight, with_lora_and_lokr);
             if (diff == nullptr) {
                 continue;
             }
 
             if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
-                weight = ggml_ext_cast_f32(ctx, backend, weight);
+                weight = ggml_ext_cast_f32(runner_ctx, weight);
             }
             weight = ggml_add(ctx, weight, diff);
         }
         return weight;
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) override {
-        return patch_weight(ctx, backend, weight, weight_name, true);
+    ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                              ggml_tensor* weight,
+                              const std::string& weight_name) override {
+        return patch_weight(runner_ctx, weight, weight_name, true);
     }
 
-    ggml_tensor* forward_with_lora(ggml_context* ctx,
-                                   ggml_backend_t backend,
+    ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx,
                                    ggml_tensor* x,
                                    ggml_tensor* w,
                                    ggml_tensor* b,
                                    const std::string& prefix,
                                    WeightAdapter::ForwardParams forward_params) override {
-        w = patch_weight(ctx, backend, w, prefix + "weight", false);
+        auto ctx = runner_ctx->ggml_ctx;
+        w        = patch_weight(runner_ctx, w, prefix + "weight", false);
         if (b) {
-            b = patch_weight(ctx, backend, b, prefix + "bias", false);
+            b = patch_weight(runner_ctx, b, prefix + "bias", false);
         }
         ggml_tensor* out;
         if (forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR) {
-            out = ggml_ext_linear(ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+            out = ggml_ext_linear(runner_ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale);
         } else {  // OP_CONV2D
-            out = ggml_ext_conv_2d(ctx,
+            out = ggml_ext_conv_2d(runner_ctx,
                                    x,
                                    w,
                                    b,
@@ -1111,7 +1126,7 @@ public:
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight");
             if (out_diff == nullptr) {
                 continue;
             }
@@ -1120,13 +1135,13 @@ public:
         return out;
     }
 
-    ggml_tensor* add_lora_to_output(ggml_context* ctx,
-                                    ggml_backend_t backend,
+    ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     ggml_tensor* w,
                                     ggml_tensor* output,
                                     const std::string& prefix,
                                     WeightAdapter::ForwardParams forward_params) override {
+        auto ctx = runner_ctx->ggml_ctx;
         // Packed INT4 stores half the logical input width. Adapter shape checks
         // and reshapes need the logical weight dimensions, without decoding it.
         ggml_tensor logical_weight = *w;
@@ -1138,10 +1153,10 @@ public:
             w = &logical_weight;
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
+            ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", runner_ctx, w, false);
             if (weight_diff != nullptr) {
                 GGML_ASSERT(forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR);
-                ggml_tensor* out_diff = ggml_ext_linear(ctx,
+                ggml_tensor* out_diff = ggml_ext_linear(runner_ctx,
                                                         x,
                                                         weight_diff,
                                                         nullptr,
@@ -1150,7 +1165,7 @@ public:
                 output                = ggml_add_inplace(ctx, output, out_diff);
             }
 
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight");
             if (out_diff != nullptr) {
                 output = ggml_add_inplace(ctx, output, out_diff);
             }

src/model/adapter/lora_ops.cpp

@@ -4,11 +4,46 @@
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
 
-ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+ggml_tensor* ggml_ext_mul_n_mode(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
+    auto ctx = runner_ctx->ggml_ctx;
+    // reshape A
+    // swap 0th and nth axis
+    a           = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0));
+    int64_t ne1 = a->ne[1];
+    int64_t ne2 = a->ne[2];
+    int64_t ne3 = a->ne[3];
+    // make 2D
+    a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1)));
+
+    ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b)));
+
+    // reshape output (same shape as a after permutation except first dim)
+    result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3);
+    // swap back 0th and nth axis
+    result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0);
+    return result;
+}
+
+ggml_tensor* ggml_ext_kronecker(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b) {
+    auto ctx = runner_ctx->ggml_ctx;
+    return ggml_mul(ctx,
+                    ggml_interpolate(ctx,
+                                     a,
+                                     a->ne[0] * b->ne[0],
+                                     a->ne[1] * b->ne[1],
+                                     a->ne[2] * b->ne[2],
+                                     a->ne[3] * b->ne[3],
+                                     GGML_SCALE_MODE_NEAREST),
+                    b);
+}
+
+ggml_tensor* ggml_ext_merge_lora(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* lora_down,
                                  ggml_tensor* lora_up,
                                  ggml_tensor* lora_mid) {
+    auto ctx = runner_ctx->ggml_ctx;
     ggml_tensor* updown;
     // flat lora tensors to multiply it
     // ggml_n_dims drops the output dimension when a linear layer has one output.
@@ -32,25 +67,25 @@ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
         // lora_down has shape (Rank, In,  1,    1)
         // lora_up   has shape (Rank, Out, 1,    1)
         // conv layer shape is (3,    3,   Out,  In)
-        updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2);
+        updown = ggml_ext_mul_n_mode(runner_ctx, ggml_ext_mul_n_mode(runner_ctx, lora_mid, lora_down, 3), lora_up, 2);
         updown = ggml_cont(ctx, updown);
     }
     return updown;
 }
 
-ggml_tensor* ggml_ext_lokr_forward(
-    ggml_context* ctx,
-    ggml_backend_t backend,
-    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
-    ggml_tensor* w1,   // Outer C (Full rank)
-    ggml_tensor* w1a,  // Outer A (Low rank part 1)
-    ggml_tensor* w1b,  // Outer B (Low rank part 2)
-    ggml_tensor* w2,   // Inner BA (Full rank)
-    ggml_tensor* w2a,  // Inner A (Low rank part 1)
-    ggml_tensor* w2b,  // Inner B (Low rank part 2)
-    bool is_conv,
-    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
-    float scale) {
+ggml_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx,
+                                   ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+                                   ggml_tensor* w1,   // Outer C (Full rank)
+                                   ggml_tensor* w1a,  // Outer A (Low rank part 1)
+                                   ggml_tensor* w1b,  // Outer B (Low rank part 2)
+                                   ggml_tensor* w2,   // Inner BA (Full rank)
+                                   ggml_tensor* w2a,  // Inner A (Low rank part 1)
+                                   ggml_tensor* w2b,  // Inner B (Low rank part 2)
+                                   bool is_conv,
+                                   WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+                                   float scale) {
+    auto ctx     = runner_ctx->ggml_ctx;
+    auto backend = runner_ctx->backend;
     GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
     GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
 
@@ -73,7 +108,6 @@ ggml_tensor* ggml_ext_lokr_forward(
 
         if (sd_backend_is(backend, "Vulkan")) {
             if (batch > 1) {
-                // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend
                 int max_batch    = 65535;
                 int max_batch_uq = max_batch / uq;
                 merge_batch_uq   = 1;
@@ -121,14 +155,14 @@ ggml_tensor* ggml_ext_lokr_forward(
 
         ggml_tensor* hc  = ggml_transpose(ctx, hc_t);
         ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
-        return ggml_ext_scale(ctx, out, scale);
+        return ggml_ext_scale(runner_ctx, out, scale);
     } else {
         int batch = (int)h->ne[3];
         // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
         ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch);
 
         if (w2 != nullptr) {
-            hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr,
+            hb = ggml_ext_conv_2d(runner_ctx, h_split, w2, nullptr,
                                   conv_params.s0,
                                   conv_params.s1,
                                   conv_params.p0,
@@ -154,7 +188,7 @@ ggml_tensor* ggml_ext_lokr_forward(
                 GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]);
                 a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]);
             }
-            ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr,
+            ggml_tensor* ha = ggml_ext_conv_2d(runner_ctx, h_split, a, nullptr,
                                                conv_params.s0,
                                                conv_params.s1,
                                                conv_params.p0,
@@ -167,8 +201,7 @@ ggml_tensor* ggml_ext_lokr_forward(
                                                conv_params.scale);
 
             // not supporting lora_mid here
-            hb = ggml_ext_conv_2d(ctx,
-                                  ha,
+            hb = ggml_ext_conv_2d(runner_ctx, ha,
                                   b,
                                   nullptr,
                                   1,
@@ -204,6 +237,6 @@ ggml_tensor* ggml_ext_lokr_forward(
         ggml_tensor* hc = ggml_transpose(ctx, hc_t);
         // ungroup
         ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
-        return ggml_ext_scale(ctx, out, scale);
+        return ggml_ext_scale(runner_ctx, out, scale);
     }
 }

src/model/adapter/lora_ops.h

@@ -3,23 +3,32 @@
 
 #include "core/ggml_runner.h"
 
-ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+// n-mode tensor-matrix product
+// example: 2-mode product
+// A: [ne03, k, ne01, ne00]
+// B: k rows, m columns => [k, m]
+// result is [ne03, m, ne01, ne00]
+ggml_tensor* ggml_ext_mul_n_mode(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0);
+
+// Kronecker product
+// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10]
+ggml_tensor* ggml_ext_kronecker(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b);
+
+ggml_tensor* ggml_ext_merge_lora(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* lora_down,
                                  ggml_tensor* lora_up,
                                  ggml_tensor* lora_mid = nullptr);
 
-ggml_tensor* ggml_ext_lokr_forward(
-    ggml_context* ctx,
-    ggml_backend_t backend,
-    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
-    ggml_tensor* w1,   // Outer C (Full rank)
-    ggml_tensor* w1a,  // Outer A (Low rank part 1)
-    ggml_tensor* w1b,  // Outer B (Low rank part 2)
-    ggml_tensor* w2,   // Inner BA (Full rank)
-    ggml_tensor* w2a,  // Inner A (Low rank part 1)
-    ggml_tensor* w2b,  // Inner B (Low rank part 2)
-    bool is_conv,
-    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
-    float scale);
+ggml_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx,
+                                   ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+                                   ggml_tensor* w1,   // Outer C (Full rank)
+                                   ggml_tensor* w1a,  // Outer A (Low rank part 1)
+                                   ggml_tensor* w1b,  // Outer B (Low rank part 2)
+                                   ggml_tensor* w2,   // Inner BA (Full rank)
+                                   ggml_tensor* w2a,  // Inner A (Low rank part 1)
+                                   ggml_tensor* w2b,  // Inner B (Low rank part 2)
+                                   bool is_conv,
+                                   WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+                                   float scale);
 
 #endif  // __SD_MODEL_ADAPTER_LORA_OPS_H__

src/model/adapter/pmid.hpp

@@ -36,7 +36,7 @@ public:
         x = layer_norm->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x),  fc1_b);
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         x = fc2->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x),  fc2_b);
         if (use_residue)
@@ -132,8 +132,8 @@ public:
         k             = reshape_tensor(ctx->ggml_ctx, k, heads);
         v             = reshape_tensor(ctx->ggml_ctx, v, heads);
         scale         = 1.f / sqrt(sqrt((float)dim_head));
-        k             = ggml_ext_scale(ctx->ggml_ctx, k, scale, true);
-        q             = ggml_ext_scale(ctx->ggml_ctx, q, scale, true);
+        k             = ggml_ext_scale(ctx, k, scale, true);
+        q             = ggml_ext_scale(ctx, q, scale, true);
         // auto weight = ggml_mul_mat(ctx, q, k);
         auto weight = ggml_mul_mat(ctx->ggml_ctx, k, q);  // NOTE order of mul is opposite to pytorch
 

src/model/audio/wav2vec2.hpp

@@ -67,7 +67,7 @@ namespace Wav2Vec2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
             x         = conv->forward(ctx, x);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -85,7 +85,7 @@ namespace Wav2Vec2 {
             x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
             x = layer_norm->forward(ctx, x);
             x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -103,7 +103,7 @@ namespace Wav2Vec2 {
             x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
             x = layer_norm->forward(ctx, x);
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -198,7 +198,7 @@ namespace Wav2Vec2 {
             auto w = weight(ctx);
             auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr;
             x      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-            x      = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
+            x      = ggml_ext_conv_1d(ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
             // Apply GELU out of place before cropping to keep graph buffer reuse safe.
             x = ggml_gelu_erf(ctx->ggml_ctx, x);
             x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0);
@@ -216,7 +216,7 @@ namespace Wav2Vec2 {
             auto intermediate_dense = std::dynamic_pointer_cast<Linear>(blocks["intermediate_dense"]);
             auto output_dense       = std::dynamic_pointer_cast<Linear>(blocks["output_dense"]);
             x                       = intermediate_dense->forward(ctx, x);
-            x                       = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x                       = ggml_ext_gelu(ctx, x, true);
             x                       = output_dense->forward(ctx, x);
             return x;
         }

src/model/common/block.hpp

@@ -32,7 +32,7 @@ public:
         if (vae_downsample) {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
-            x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
             x = conv->forward(ctx, x);
         } else {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);
@@ -220,7 +220,7 @@ public:
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
         x = proj->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         return x;
     }
 };
@@ -250,7 +250,7 @@ public:
         auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
 
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         x = fc2->forward(ctx, x);
         return x;
     }
@@ -293,9 +293,7 @@ public:
 
         auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
         auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
-        if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-            net_2->set_force_prec_f32(true);
-        }
+        net_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
         x = net_0->forward(ctx, x);  // [ne3, ne2, ne1, inner_dim]
         x = net_2->forward(ctx, x);  // [ne3, ne2, ne1, dim_out]
@@ -597,7 +595,7 @@ public:
         x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, ln_1->forward(ctx, x)));
 
         auto h = mlp_c_fc->forward(ctx, ln_2->forward(ctx, x));
-        h      = ggml_ext_gelu_quick(ctx->ggml_ctx, h);
+        h      = ggml_ext_gelu_quick(ctx, h);
         h      = mlp_c_proj->forward(ctx, h);
         return ggml_add(ctx->ggml_ctx, x, h);
     }
@@ -614,7 +612,7 @@ protected:
     ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
         // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,]
         // so learned_with_images is same as learned
-        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]);
         return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 

src/model/common/ggml_block.hpp

@@ -147,6 +147,7 @@ protected:
     bool bias;
     bool force_f32;
     bool force_prec_f32;
+    uint32_t precision_flags    = 0;
     bool has_weight_scale       = false;
     int w4_convrot_kind         = 0;
     bool int8_convrot           = false;
@@ -215,6 +216,10 @@ public:
         scale = scale_;
     }
 
+    void set_precision_flags(uint32_t flags) {
+        precision_flags = flags;
+    }
+
     void set_force_prec_f32(bool force_prec_f32_) {
         force_prec_f32 = force_prec_f32_;
     }
@@ -225,72 +230,34 @@ public:
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        const bool force_prec_f32 = ggml_ext_force_prec_f32(ctx, this->force_prec_f32, precision_flags);
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
-#ifndef SD_USE_UPSTREAM_GGML
-        if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
-            bool supports_fp8_matmul = false;
-            if (ctx->backend != nullptr) {
-                ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x);
-                if (force_prec_f32) {
-                    ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32);
-                }
-                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul) && !sd_backend_is_cpu(ctx->backend);
-            }
-            if (!supports_fp8_matmul) {
-                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
-            }
-        }
-#endif
-        ggml_tensor* b = nullptr;
+        w                         = ggml_ext_prepare_linear_weight(ctx, w, x, force_prec_f32);
+        ggml_tensor* b            = nullptr;
         if (bias) {
             b = params["bias"];
         }
         ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
         ggml_tensor* out         = nullptr;
         if (w->type == GGML_TYPE_I8) {
-            if (x->type != GGML_TYPE_F32) {
-                x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x);
-            }
-            if (!ggml_is_contiguous(x)) {
-                x = ggml_cont(ctx->ggml_ctx, x);
-            }
+            x                       = ggml_ext_prepare_i8_input(ctx, x);
             ggml_tensor* lora_input = x;
             if (ctx->weight_adapter && b != nullptr) {
-                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
-#ifndef SD_USE_UPSTREAM_GGML
-            if (int8_convrot && scale == 1.f) {
-                const auto cache_key = std::make_pair(x, int8_convrot_group_size);
-                auto cached          = ctx->int8_convrot_cache.find(cache_key);
-                if (cached == ctx->int8_convrot_cache.end()) {
-                    x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size);
-                    ctx->int8_convrot_cache.emplace(cache_key, x);
-                } else {
-                    x = cached->second;
-                }
-            }
-#endif
+
 #ifndef SD_USE_UPSTREAM_GGML
             if (w4_convrot_kind) {
-                if (scale != 1.f)
-                    x = ggml_ext_scale(ctx->ggml_ctx, x, scale);
-                out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x,
-                                              w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
-                                              w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
-                                              w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
-                                              scale == 1.f ? b : nullptr, static_cast<SDW4Kind>(w4_convrot_kind),
-                                              int8_convrot ? int8_convrot_group_size : 0);
-                GGML_ASSERT(out);
-                if (scale != 1.f) {
-                    out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale);
-                    if (b)
-                        out = ggml_add_inplace(ctx->ggml_ctx, out, b);
-                }
+                out = ggml_ext_linear_w4_convrot(ctx, x, w,
+                                                 w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
+                                                 w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
+                                                 w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
+                                                 b, static_cast<SDW4Kind>(w4_convrot_kind), int8_convrot ? int8_convrot_group_size : 0, scale);
             } else
 #endif
-                out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                out = ggml_ext_linear_i8_tensorwise(ctx,
                                                     x,
                                                     w,
                                                     weight_scale,
@@ -302,33 +269,21 @@ public:
                 forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
                 forward_params.linear.force_prec_f32 = force_prec_f32;
                 forward_params.linear.scale          = scale;
-                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
-                                                                                               ctx->backend,
-                                                                                               lora_input,
-                                                                                               w,
-                                                                                               out,
-                                                                                               prefix,
-                                                                                               forward_params);
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx, lora_input, w, out, prefix, forward_params);
             }
             return out;
         }
         if (has_weight_scale) {
-            out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale);
+            out = ggml_ext_linear(ctx, x, w, nullptr, force_prec_f32, scale);
             out = ggml_mul(ctx->ggml_ctx, out, weight_scale);
             if (ctx->weight_adapter) {
                 WeightAdapter::ForwardParams forward_params;
                 forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
                 forward_params.linear.force_prec_f32 = force_prec_f32;
                 forward_params.linear.scale          = scale;
-                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
-                                                                                               ctx->backend,
-                                                                                               x,
-                                                                                               w,
-                                                                                               out,
-                                                                                               prefix,
-                                                                                               forward_params);
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx, x, w, out, prefix, forward_params);
                 if (b != nullptr) {
-                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
             }
             if (b != nullptr) {
@@ -341,9 +296,9 @@ public:
             forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
             forward_params.linear.force_prec_f32 = force_prec_f32;
             forward_params.linear.scale          = scale;
-            out                                  = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx, x, w, linear_bias, prefix, forward_params);
         } else {
-            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
+            out = ggml_ext_linear(ctx, x, w, linear_bias, force_prec_f32, scale);
         }
         return out;
     }
@@ -376,42 +331,6 @@ protected:
     bool is_int8_tensorwise     = false;
     int int8_convrot_group_size = 0;
 
-    static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) {
-        const auto* embedding     = static_cast<const Embedding*>(userdata);
-        const int group_size      = embedding->int8_convrot_group_size;
-        const auto* weight        = dst->src[0];
-        const auto* input_ids     = static_cast<const int32_t*>(dst->src[1]->data);
-        const auto* scales        = static_cast<const float*>(dst->src[2]->data);
-        const bool scalar_scale   = ggml_nelements(dst->src[2]) == 1;
-        const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast<float>(group_size)) : 1.f;
-        for (int64_t row = ith; row < dst->ne[1]; row += nth) {
-            const int32_t token = input_ids[row];
-            GGML_ASSERT(token >= 0 && token < weight->ne[1]);
-            const auto* src   = reinterpret_cast<const int8_t*>(static_cast<const char*>(weight->data) + token * weight->nb[1]);
-            auto* out         = reinterpret_cast<float*>(static_cast<char*>(dst->data) + row * dst->nb[1]);
-            const float scale = scales[scalar_scale ? 0 : token] * normalization;
-            for (int64_t i = 0; i < weight->ne[0]; ++i) {
-                out[i] = static_cast<float>(src[i]) * scale;
-            }
-            // The regular Hadamard rotation is symmetric and its own inverse.
-            for (int stride = 1; stride < group_size; stride *= 4) {
-                for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) {
-                    for (int j = 0; j < stride; ++j) {
-                        float* values      = out + base + j;
-                        const float a      = values[0];
-                        const float b      = values[stride];
-                        const float c      = values[2 * stride];
-                        const float d      = values[3 * stride];
-                        values[0]          = a + b + c - d;
-                        values[stride]     = a + b - c + d;
-                        values[2 * stride] = a - b + c + d;
-                        values[3 * stride] = -a + b + c + d;
-                    }
-                }
-            }
-        }
-    }
-
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
         auto weight_storage     = tensor_storage_map.find(prefix + "weight");
         is_int8_tensorwise      = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
@@ -451,30 +370,9 @@ public:
           num_embeddings(num_embeddings) {
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids) override {
-        // input_ids: [N, n_token]
-        auto weight = params["weight"];
-
-        // There are issues with ggml batch inference, so we are expanding it here first.
-        // TODO: fix ggml batch inference
-        int64_t n = input_ids->ne[1];
-        input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]);
-
-        ggml_tensor* embedding;
-        if (is_int8_tensorwise) {
-            GGML_ASSERT(input_ids->type == GGML_TYPE_I32);
-            ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]};
-            embedding           = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1,
-                                                 args, 3, get_rows_i8, GGML_N_TASKS_MAX, this);
-        } else {
-            input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
-            embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
-        }
-        embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
-
-        // [N, n_token, embedding_dim]
-        return embedding;
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input_ids) override {
+        return ggml_ext_embedding(ctx, params["weight"], input_ids,
+                                  is_int8_tensorwise ? params["weight_scale"] : nullptr, int8_convrot_group_size);
     }
 };
 
@@ -528,7 +426,7 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         GGML_ASSERT(x->ne[1] == in_channels);
-        return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr,
+        return ggml_ext_conv_1d(ctx, x, params["weight"], bias ? params["bias"] : nullptr,
                                 stride, padding, dilation, groups, force_prec_f32);
     }
 };
@@ -598,9 +496,9 @@ public:
             forward_params.conv2d.circular_x = ctx->circular_x_enabled;
             forward_params.conv2d.circular_y = ctx->circular_y_enabled;
             forward_params.conv2d.scale      = scale;
-            return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            return ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params);
         }
-        return ggml_ext_conv_2d(ctx->ggml_ctx,
+        return ggml_ext_conv_2d(ctx,
                                 x,
                                 w,
                                 b,
@@ -687,9 +585,9 @@ public:
                 forward_params.conv2d.circular_x = ctx->circular_x_enabled;
                 forward_params.conv2d.circular_y = ctx->circular_y_enabled;
                 forward_params.conv2d.scale      = scale;
-                return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+                return ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params);
             }
-            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+            return ggml_ext_conv_2d(ctx, x, w, b,
                                     stride.second, stride.first,
                                     padding.second, padding.first,
                                     dilation.second, dilation.first,
@@ -699,24 +597,15 @@ public:
                                     scale);
         }
 
-        if (groups == in_channels && groups == out_channels) {
-            ggml_tensor* res;
-            if (ctx->conv2d_direct_enabled) {
-                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
-                                             stride.second, stride.first,
-                                             padding.second, padding.first,
-                                             dilation.second, dilation.first);
-            } else {
-                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
-                                      stride.second, stride.first,
-                                      padding.second, padding.first,
-                                      dilation.second, dilation.first);
-            }
-            if (b) {
-                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
-                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
-            }
-            return res;
+        if (!ctx->weight_adapter || (groups == in_channels && groups == out_channels)) {
+            return ggml_ext_conv_2d_grouped(ctx, x, w, b, groups,
+                                            stride.second, stride.first,
+                                            padding.second, padding.first,
+                                            dilation.second, dilation.first,
+                                            ctx->conv2d_direct_enabled,
+                                            ctx->circular_x_enabled,
+                                            ctx->circular_y_enabled,
+                                            scale);
         }
 
         int64_t ic_g = in_channels / groups;
@@ -743,33 +632,22 @@ public:
                 b_i             = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
             }
 
-            if (ctx->weight_adapter) {
-                WeightAdapter::ForwardParams forward_params;
-                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
-                forward_params.conv2d.s0         = stride.second;
-                forward_params.conv2d.s1         = stride.first;
-                forward_params.conv2d.p0         = padding.second;
-                forward_params.conv2d.p1         = padding.first;
-                forward_params.conv2d.d0         = dilation.second;
-                forward_params.conv2d.d1         = dilation.first;
-                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
-                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
-                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
-                forward_params.conv2d.scale      = scale;
-                out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
-            } else {
-                out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
-                                                 stride.second, stride.first,
-                                                 padding.second, padding.first,
-                                                 dilation.second, dilation.first,
-                                                 ctx->conv2d_direct_enabled,
-                                                 ctx->circular_x_enabled,
-                                                 ctx->circular_y_enabled,
-                                                 scale);
-            }
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+            forward_params.conv2d.s0         = stride.second;
+            forward_params.conv2d.s1         = stride.first;
+            forward_params.conv2d.p0         = padding.second;
+            forward_params.conv2d.p1         = padding.first;
+            forward_params.conv2d.d0         = dilation.second;
+            forward_params.conv2d.d1         = dilation.first;
+            forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+            forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+            forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+            forward_params.conv2d.scale      = scale;
+            out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx, x_i, w_i, b_i, prefix, forward_params);
         }
 
-        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+        ggml_tensor* out = ggml_ext_vec_concat(ctx, out_slices, 2);
 
         return out;
     }
@@ -823,7 +701,7 @@ public:
         ggml_tensor* w = params["weight"];
         ggml_tensor* b = nullptr;
         if (ctx->weight_adapter) {
-            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             if (w->type != GGML_TYPE_F16) {
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16);
             }
@@ -831,10 +709,10 @@ public:
         if (bias) {
             b = params["bias"];
             if (ctx->weight_adapter) {
-                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
         }
-        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+        return ggml_ext_conv_3d(ctx, x, w, b, in_channels,
                                 std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                 std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
                                 std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
@@ -879,16 +757,16 @@ public:
         if (elementwise_affine) {
             w = params["weight"];
             if (ctx->weight_adapter) {
-                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             }
             if (bias) {
                 b = params["bias"];
                 if (ctx->weight_adapter) {
-                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
             }
         }
-        return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps);
+        return ggml_ext_layer_norm(ctx, x, w, b, eps);
     }
 };
 
@@ -927,11 +805,11 @@ public:
             w = params["weight"];
             b = params["bias"];
             if (ctx->weight_adapter) {
-                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
-                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
         }
-        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps);
+        return ggml_ext_group_norm(ctx, x, w, b, num_groups, eps);
     }
 };
 
@@ -971,7 +849,7 @@ public:
         }
         ggml_tensor* w = params["weight"];
         if (ctx->weight_adapter) {
-            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
         }
         x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
         // The affine multiply may run on a different backend after normalization falls back.
@@ -1032,7 +910,7 @@ public:
         if (proj_in) {
             auto in_proj = std::dynamic_pointer_cast<Linear>(blocks[in_proj_name]);
             auto qkv     = in_proj->forward(ctx, x);
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             q            = qkv_vec[0];
             k            = qkv_vec[1];
             v            = qkv_vec[2];

src/model/common/rope.hpp

@@ -1071,58 +1071,7 @@ namespace Rope {
                                               ggml_tensor* x,
                                               ggml_tensor* pe,
                                               bool rope_interleaved = true) {
-        auto ctx = runner_ctx->ggml_ctx;
-#ifndef SD_USE_UPSTREAM_GGML
-        auto backend = runner_ctx->backend;
-        if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
-            x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 &&
-            pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) {
-            auto out = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved);
-            if (out && ggml_backend_supports_op(backend, out)) {
-                return out;
-            }
-        }
-#endif
-        // x: [N, L, n_head, d_head]
-        // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
-        int64_t d_head = x->ne[0];
-        int64_t n_head = x->ne[1];
-        int64_t L      = x->ne[2];
-        int64_t N      = x->ne[3];
-        x              = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, n_head, L, d_head]
-        if (rope_interleaved) {
-            x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
-            x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]
-        } else {
-            auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0);
-            auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]);
-
-            auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1));
-            auto cos  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0);
-            auto sin  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]);
-
-            auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin));
-            auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos));
-            return ggml_concat(ctx, part_0, part_1, 0);
-        }
-
-        int64_t offset = x->nb[2] * x->ne[2];
-        auto x_0       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0);  // [N * n_head, L, d_head/2]
-        auto x_1       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1);  // [N * n_head, L, d_head/2]
-        x_0            = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        x_1            = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        auto temp_x    = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]);
-        x_0            = ggml_repeat(ctx, x_0, temp_x);  // [N * n_head, L, d_head/2, 2]
-        x_1            = ggml_repeat(ctx, x_1, temp_x);  // [N * n_head, L, d_head/2, 2]
-
-        pe        = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2));  // [2, L, d_head/2, 2]
-        offset    = pe->nb[2] * pe->ne[2];
-        auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0);  // [L, d_head/2, 2]
-        auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1);  // [L, d_head/2, 2]
-
-        auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1));  // [N * n_head, L, d_head/2, 2]
-        x_out      = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);                         // [N*n_head, L, d_head]
-        return x_out;
+        return ggml_ext_rope_apply(runner_ctx, x, pe, rope_interleaved);
     }
 
     __STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,

src/model/common/vsr.hpp

@@ -177,7 +177,7 @@ namespace VSR {
                     params.window_height[i] = windows[i].first;
                     params.window_width[i]  = windows[i].second;
                 }
-                if (auto out = sd_ggml_window_attn(c, ctx->backend, q, k, v, params))
+                if (auto out = ggml_ext_window_attn(ctx, q, k, v, params))
                     return out;
             }
 #endif
@@ -192,7 +192,7 @@ namespace VSR {
                 auto tk   = ggml_get_rows(c, k, ki);
                 auto tv   = ggml_get_rows(c, v, ki);
                 auto mask = inputs.input(ctx, tile.mask);
-                auto out  = ggml_ext_attention_ext(c, ctx->backend, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled);
+                auto out  = ggml_ext_attention_ext(ctx, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled);
                 out       = ggml_reshape_3d(c, out, out->ne[0], tile.width, tile.height);
                 row       = row ? ggml_concat(c, row, out, 1) : out;
                 row_width += tile.width;

src/model/diffusion/anima.hpp

@@ -123,7 +123,7 @@ namespace Anima {
                 emb = ggml_add(ctx->ggml_ctx, emb, temb);
             }
 
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 3, 0);
+            auto emb_chunks = ggml_ext_chunk(ctx, emb, 3, 0);
             auto shift      = emb_chunks[0];
             auto scale      = emb_chunks[1];
             auto gate       = emb_chunks[2];
@@ -164,7 +164,7 @@ namespace Anima {
                 emb          = ggml_add(ctx->ggml_ctx, emb, temb_2c);
             }
 
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto emb_chunks = ggml_ext_chunk(ctx, emb, 2, 0);
             auto shift      = emb_chunks[0];
             auto scale      = emb_chunks[1];
 
@@ -275,7 +275,7 @@ namespace Anima {
             auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
 
             x = layer1->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
             x = layer2->forward(ctx, x);
             return x;
         }
@@ -293,7 +293,7 @@ namespace Anima {
             auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
 
             x = layer0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
             x = layer2->forward(ctx, x);
             return x;
         }
@@ -500,14 +500,14 @@ namespace Anima {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
 
-            auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
+            auto padding_mask = ggml_ext_zeros(ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
             x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
 
             x               = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
             int64_t img_len = x->ne[1];
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
+                    auto padding_mask = ggml_ext_zeros(ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
                     ref               = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2);  // [N, C + 1, H, W]
                     ref               = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
                     x                 = ggml_concat(ctx->ggml_ctx, x, ref, 1);
@@ -515,7 +515,7 @@ namespace Anima {
             }
             x = x_embedder->forward(ctx, x);
 
-            auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
+            auto timestep_proj     = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(config.hidden_size));
             auto temb              = t_embedder->forward(ctx, timestep_proj);
             auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
 
@@ -530,14 +530,14 @@ namespace Anima {
                     adapted_context = ggml_mul(ctx->ggml_ctx, adapted_context, w);
                 }
                 if (adapted_context->ne[1] < 512) {
-                    auto pad_ctx    = ggml_ext_zeros(ctx->ggml_ctx,
+                    auto pad_ctx    = ggml_ext_zeros(ctx,
                                                      adapted_context->ne[0],
                                                      512 - adapted_context->ne[1],
                                                      adapted_context->ne[2],
                                                      1);
                     adapted_context = ggml_concat(ctx->ggml_ctx, adapted_context, pad_ctx, 1);
                 } else if (adapted_context->ne[1] > 512) {
-                    adapted_context = ggml_ext_slice(ctx->ggml_ctx, adapted_context, 1, 0, 512);
+                    adapted_context = ggml_ext_slice(ctx, adapted_context, 1, 0, 512);
                 }
                 encoder_hidden_states = adapted_context;
             }
@@ -552,11 +552,11 @@ namespace Anima {
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
                 sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
             }
-            x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
+            x = ggml_ext_slice(ctx, x, 1, 0, img_len);
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
 
-            x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false);  // [N, C, H, W]
+            x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, false);  // [N, C, H, W]
 
             return x;
         }

src/model/diffusion/animatediff.hpp

@@ -121,13 +121,13 @@ namespace AnimateDiff {
             auto residual = x;
             auto h        = norm->forward(ctx, x);
 
-            h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
+            h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
             h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H);
             h = proj_in->forward(ctx, h);
             h = tb0->forward(ctx, h);
             h = proj_out->forward(ctx, h);
             h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H);
-            h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
+            h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
 
             return ggml_add(ctx->ggml_ctx, h, residual);
         }

src/model/diffusion/boogu.hpp

@@ -161,7 +161,7 @@ namespace Boogu {
             auto caption_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["caption_embedder.0"]);
             auto caption_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["caption_embedder.1"]);
 
-            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
             auto time_embed    = timestep_embedder->forward(ctx, timestep_proj);
             auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states));
             return {time_embed, caption_embed};
@@ -179,7 +179,7 @@ namespace Boogu {
             auto norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
 
             emb       = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));
-            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0);
+            auto mods = ggml_ext_chunk(ctx, emb, 4, 0);
 
             auto scale_msa = mods[0];
             auto gate_msa  = mods[1];
@@ -204,9 +204,7 @@ namespace Boogu {
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
             auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                linear_2->set_force_prec_f32(true);
-            }
+            linear_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto h1 = linear_1->forward(ctx, x);
             auto h2 = linear_3->forward(ctx, x);
@@ -264,9 +262,7 @@ namespace Boogu {
             auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t N  = hidden_states->ne[2];
             int64_t Lq = hidden_states->ne[1];
@@ -388,9 +384,7 @@ namespace Boogu {
             auto instruct_out  = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
             auto img_out       = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t N          = img_hidden_states->ne[2];
             int64_t L_img      = img_hidden_states->ne[1];
@@ -417,8 +411,8 @@ namespace Boogu {
             k      = norm_k->forward(ctx, k);
 
             auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
-            auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct);
-            auto img_attn      = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
+            auto instruct_attn = ggml_ext_slice(ctx, hidden_states, 1, 0, L_instruct);
+            auto img_attn      = ggml_ext_slice(ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
 
             instruct_attn = instruct_out->forward(ctx, instruct_attn);
             img_attn      = img_out->forward(ctx, img_attn);
@@ -504,8 +498,8 @@ namespace Boogu {
             auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec);
 
             auto joint_attn_out    = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb);
-            auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct);
-            auto img_attn_out      = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
+            auto instruct_attn_out = ggml_ext_slice(ctx, joint_attn_out, 1, 0, L_instruct);
+            auto img_attn_out      = ggml_ext_slice(ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
 
             auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb);
 
@@ -634,7 +628,7 @@ namespace Boogu {
             auto time_caption_embed = std::dynamic_pointer_cast<LuminaCombinedTimestepCaptionEmbedding>(blocks["time_caption_embed"]);
             auto norm_out           = std::dynamic_pointer_cast<LuminaLayerNormContinuous>(blocks["norm_out"]);
 
-            auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps);
+            auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, timesteps), timesteps);
             auto embeds   = time_caption_embed->forward(ctx, timestep, context);
             auto temb     = embeds.first;
             auto txt      = embeds.second;
@@ -647,8 +641,8 @@ namespace Boogu {
             int64_t txt_len = txt->ne[1];
 
             GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len);
-            auto txt_pe   = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len);
-            auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
+            auto txt_pe   = ggml_ext_slice(ctx, pe, 3, 0, txt_len);
+            auto noise_pe = ggml_ext_slice(ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
 
             for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
@@ -664,7 +658,7 @@ namespace Boogu {
 
             ggml_tensor* combined_img = img;
             if (ref_img != nullptr) {
-                auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len);
+                auto ref_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + ref_len);
                 for (int i = 0; i < config.num_refiner_layers; i++) {
                     auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["ref_image_refiner." + std::to_string(i)]);
                     ref_img    = block->forward(ctx, ref_img, ref_pe, temb);
@@ -673,7 +667,7 @@ namespace Boogu {
                 combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1);
             }
 
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
             for (int i = 0; i < config.num_double_stream_layers; i++) {
                 auto block   = std::dynamic_pointer_cast<BooguImageDoubleStreamBlock>(blocks["double_stream_layers." + std::to_string(i)]);
                 auto result  = block->forward(ctx, combined_img, txt, pe, img_pe, temb);
@@ -691,9 +685,9 @@ namespace Boogu {
             }
 
             hidden_states = norm_out->forward(ctx, hidden_states, temb);
-            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
-            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
-            hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f);
+            hidden_states = ggml_ext_slice(ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
+            hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
+            hidden_states = ggml_ext_scale(ctx, hidden_states, -1.f);
             return hidden_states;
         }
     };

src/model/diffusion/control.hpp

@@ -230,7 +230,7 @@ public:
         auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["control_add_embedding.linear_2"]);
 
         // [N*num_control_types, CONTROL_TYPE_EMBED_DIM]
-        auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM);
+        auto e = ggml_ext_timestep_embedding(ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM);
         e      = ggml_reshape_2d(ctx->ggml_ctx, e, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, batch);
         e      = linear_1->forward(ctx, e);
         e      = ggml_silu_inplace(ctx->ggml_ctx, e);
@@ -238,8 +238,9 @@ public:
     }
 
     // [N, C, H, W] => [N, C]
-    static ggml_tensor* spatial_mean(ggml_context* ctx, ggml_tensor* x) {
-        auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+    static ggml_tensor* spatial_mean(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+        auto ctx  = runner_ctx->ggml_ctx;
+        auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(runner_ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
         return ggml_reshape_2d(ctx, ggml_mean(ctx, flat), x->ne[2], x->ne[3]);
     }
 
@@ -255,11 +256,11 @@ public:
         auto task_embedding = params["task_embedding"];  // [num_control_types, C]
         auto task           = ggml_view_2d(c, task_embedding, channels, 1, task_embedding->nb[1], control_type_index * task_embedding->nb[1]);
 
-        auto hint_feat = ggml_add(c, spatial_mean(c, hint), task);  // [N, C]
+        auto hint_feat = ggml_add(c, spatial_mean(ctx, hint), task);  // [N, C]
         if (hint_feat->ne[1] != batch) {
             hint_feat = ggml_repeat(c, hint_feat, ggml_new_tensor_2d(c, GGML_TYPE_F32, channels, batch));
         }
-        auto sample_feat = spatial_mean(c, sample);  // [N, C]
+        auto sample_feat = spatial_mean(ctx, sample);  // [N, C]
 
         // as in the reference (nn.MultiheadAttention without batch_first) the attention is over
         // the batch, and the 2 features (hint, sample) are its batch: [2, N, C]
@@ -273,7 +274,7 @@ public:
         }
 
         auto projection = std::dynamic_pointer_cast<Linear>(blocks["spatial_ch_projs"]);
-        auto alpha      = ggml_reshape_2d(c, ggml_ext_slice(c, x, 2, 0, 1), channels, batch);
+        auto alpha      = ggml_reshape_2d(c, ggml_ext_slice(ctx, x, 2, 0, 1), channels, batch);
         alpha           = projection->forward(ctx, alpha);
         alpha           = ggml_reshape_4d(c, alpha, 1, 1, channels, batch);
 
@@ -346,7 +347,7 @@ public:
 
         auto middle_block_out = std::dynamic_pointer_cast<Conv2d>(blocks["middle_block_out.0"]);
 
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);

src/model/diffusion/dit.hpp

@@ -5,11 +5,12 @@
 #include "core/ggml_runner.h"
 
 namespace DiT {
-    inline ggml_tensor* patchify(ggml_context* ctx,
+    inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  int pw,
                                  int ph,
                                  bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, C, H, W]
         // return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         int64_t N = x->ne[3];
@@ -28,19 +29,20 @@ namespace DiT {
             x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, h*w, C, ph*pw]
             x = ggml_reshape_3d(ctx, x, pw * ph * C, w * h, N);    // [N, h*w, C*ph*pw]
         } else {
-            x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, h*w, C, ph*pw]
-            x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N);              // [N, h*w, ph*pw*C]
+            x = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [N, h*w, C, ph*pw]
+            x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N);                     // [N, h*w, ph*pw*C]
         }
         return x;
     }
 
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
+    inline ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                    ggml_tensor* x,
                                    int64_t h,
                                    int64_t w,
                                    int ph,
                                    int pw,
                                    bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         // return: [N, C, H, W]
         int64_t N = x->ne[2];
@@ -74,7 +76,7 @@ namespace DiT {
 
         int pad_h = (ph - H % ph) % ph;
         int pad_w = (pw - W % pw) % pw;
-        x         = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+        x         = ggml_ext_pad(ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
         return x;
     }
 
@@ -84,11 +86,11 @@ namespace DiT {
                                          int pw,
                                          bool patch_last = true) {
         x = pad_to_patch_size(ctx, x, ph, pw);
-        x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last);
+        x = patchify(ctx, x, ph, pw, patch_last);
         return x;
     }
 
-    inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx,
+    inline ggml_tensor* unpatchify_and_crop(GGMLRunnerContext* runner_ctx,
                                             ggml_tensor* x,
                                             int64_t H,
                                             int64_t W,
@@ -99,19 +101,20 @@ namespace DiT {
         int pad_w = (pw - W % pw) % pw;
         int64_t h = ((H + pad_h) / ph);
         int64_t w = ((W + pad_w) / pw);
-        x         = unpatchify(ctx, x, h, w, ph, pw, patch_last);  // [N, C, H + pad_h, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 1, 0, H);               // [N, C, H, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 0, 0, W);               // [N, C, H, W]
+        x         = unpatchify(runner_ctx, x, h, w, ph, pw, patch_last);  // [N, C, H + pad_h, W + pad_w]
+        x         = ggml_ext_slice(runner_ctx, x, 1, 0, H);               // [N, C, H, W + pad_w]
+        x         = ggml_ext_slice(runner_ctx, x, 0, 0, W);               // [N, C, H, W]
         return x;
     }
 
-    inline ggml_tensor* patchify_3d(ggml_context* ctx,
+    inline ggml_tensor* patchify_3d(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     int pt,
                                     int ph,
                                     int pw,
                                     int64_t N       = 1,
                                     bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N*C, T, H, W]
         // return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
         int64_t C     = x->ne[3] / N;
@@ -125,24 +128,24 @@ namespace DiT {
         GGML_ASSERT(C * N == x->ne[3]);
         GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
 
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);  // [N*C*t_len, pt, h_len*ph, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, h_len*ph, pt, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);  // [N*C*t_len*h_len, ph, pt, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt, ph, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt*ph, w_len, pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, w_len, pt*ph, pw]
-        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);            // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);            // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);            // [N, C, t_len*h_len*w_len, pt*ph*pw]
         if (patch_last) {
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N, t_len*h_len*w_len, C, pt*ph*pw]
         } else {
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [N, t_len*h_len*w_len, pt*ph*pw, C]
         }
         // [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
         x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
         return x;
     }
 
-    inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
+    inline ggml_tensor* unpatchify_3d(GGMLRunnerContext* runner_ctx,
                                       ggml_tensor* x,
                                       int64_t t_len,
                                       int64_t h_len,
@@ -151,6 +154,7 @@ namespace DiT {
                                       int ph,
                                       int pw,
                                       bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
         int64_t N = x->ne[2];
@@ -159,20 +163,20 @@ namespace DiT {
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
         if (patch_last) {
-            x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, C, pt*ph*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
         } else {
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
         }
 
-        x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
         return x;
     }
 }  // namespace DiT

src/model/diffusion/ernie_image.hpp

@@ -85,11 +85,12 @@ namespace ErnieImage {
         }
     };
 
-    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx,
                                                               ggml_tensor* timesteps,
                                                               int dim,
                                                               int max_period = 10000) {
-        auto emb       = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f);
+        auto ctx       = runner_ctx->ggml_ctx;
+        auto emb       = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, max_period, 1.0f);
         int64_t half   = dim / 2;
         auto cos_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0);
         auto sin_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]);
@@ -97,7 +98,8 @@ namespace ErnieImage {
         return sin_first;
     }
 
-    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) {
+    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, S, heads, head_dim]
         // pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2].
         int64_t head_dim = x->ne[0];
@@ -110,8 +112,8 @@ namespace ErnieImage {
         GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2);
 
         x           = ggml_cont(ctx, x);
-        auto x_rot  = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false);
-        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr;
+        auto x_rot  = ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim, false);
+        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(runner_ctx, x, 0, rot_dim, head_dim, false) : nullptr;
 
         int64_t half = rot_dim / 2;
         auto x1      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0);
@@ -120,8 +122,8 @@ namespace ErnieImage {
         x2           = ggml_cont(ctx, x2);
         auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0);
 
-        auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false);
-        auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false);
+        auto cos_emb = ggml_ext_slice(runner_ctx, pe, 3, 0, 1, false);
+        auto sin_emb = ggml_ext_slice(runner_ctx, pe, 3, 1, 2, false);
 
         auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb));
         if (x_pass != nullptr) {
@@ -175,8 +177,8 @@ namespace ErnieImage {
             q = norm_q->forward(ctx, q);
             k = norm_k->forward(ctx, k);
 
-            q = apply_rotary_emb(ctx->ggml_ctx, q, pe);
-            k = apply_rotary_emb(ctx->ggml_ctx, k, pe);
+            q = apply_rotary_emb(ctx, q, pe);
+            k = apply_rotary_emb(ctx, k, pe);
 
             q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, heads, S, head_dim]
             q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
@@ -269,7 +271,7 @@ namespace ErnieImage {
             auto norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
 
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0);
+            auto mods  = ggml_ext_chunk(ctx, linear->forward(ctx, conditioning), 2, 0);
             auto scale = mods[0];
             auto shift = mods[1];
 
@@ -336,9 +338,9 @@ namespace ErnieImage {
             auto final_norm      = std::dynamic_pointer_cast<ErnieImageAdaLNContinuous>(blocks["final_norm"]);
             auto final_linear    = std::dynamic_pointer_cast<Linear>(blocks["final_linear"]);
 
-            auto img = x_embedder_proj->forward(ctx, x);                                                  // [N, hidden_size, Hp, Wp]
-            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);       // [N, hidden_size, image_tokens]
-            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));  // [N, image_tokens, hidden_size]
+            auto img = x_embedder_proj->forward(ctx, x);                                             // [N, hidden_size, Hp, Wp]
+            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);  // [N, hidden_size, image_tokens]
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));       // [N, image_tokens, hidden_size]
 
             auto txt       = context;
             auto text_proj = std::dynamic_pointer_cast<Linear>(blocks["text_proj"]);
@@ -348,13 +350,13 @@ namespace ErnieImage {
 
             auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);  // [N, image_tokens + text_tokens, hidden_size]
 
-            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
+            auto sample = timestep_embedding_sin_cos(ctx, timestep, static_cast<int>(config.hidden_size));
             auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
 
             auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
             sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states");
             // sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params");
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
+            auto chunks = ggml_ext_chunk(ctx, mod_params, 6, 0);
             std::vector<ggml_tensor*> temb;
             temb.reserve(6);
             for (auto chunk : chunks) {
@@ -368,10 +370,10 @@ namespace ErnieImage {
             }
 
             hidden_states = final_norm->forward(ctx, hidden_states, c);
-            hidden_states = final_linear->forward(ctx, hidden_states);                  // [N, image_tokens, p*p*out_channels]
-            auto patches  = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
+            hidden_states = final_linear->forward(ctx, hidden_states);        // [N, image_tokens, p*p*out_channels]
+            auto patches  = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
 
-            auto out = DiT::unpatchify(ctx->ggml_ctx,
+            auto out = DiT::unpatchify(ctx,
                                        patches,
                                        Hp,
                                        Wp,

src/model/diffusion/fibo2.hpp

@@ -108,14 +108,14 @@ namespace Fibo2 {
             };
             ggml_tensor *q, *k, *v;
             if (fused_qkv) {
-                auto parts = ggml_ext_chunk(c, project("qkv", x), 3, 0);
+                auto parts = ggml_ext_chunk(ctx, project("qkv", x), 3, 0);
                 q          = parts[0];
                 k          = parts[1];
                 v          = parts[2];
             } else {
                 q = project("to_q", x);
                 if (fused_kv) {
-                    auto parts = ggml_ext_chunk(c, project("to_kv", kv), 2, 0);
+                    auto parts = ggml_ext_chunk(ctx, project("to_kv", kv), 2, 0);
                     k          = parts[0];
                     v          = parts[1];
                 } else {
@@ -158,9 +158,7 @@ namespace Fibo2 {
                 return ZImage::FeedForward::forward(ctx, x);
             }
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                w2->set_force_prec_f32(true);
-            }
+            w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
             x = std::dynamic_pointer_cast<Linear>(blocks["w13"])->forward(ctx, x);
             return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, x));
         }
@@ -206,7 +204,7 @@ namespace Fibo2 {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* text, int gist_count) {
-            auto x = ggml_ext_slice(ctx->ggml_ctx, params["latent_queries"], 1, 0, gist_count);
+            auto x = ggml_ext_slice(ctx, params["latent_queries"], 1, 0, gist_count);
             for (int i = 0; i < config.perceiver_depth; ++i) {
                 x = std::dynamic_pointer_cast<PerceiverLayer>(blocks["layers." + std::to_string(i)])->forward(ctx, x, text);
             }
@@ -244,8 +242,8 @@ namespace Fibo2 {
                 return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h);
             };
             auto modulation = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
-            auto mn         = ggml_ext_chunk(c, modulation->forward(ctx, noisy), 4, 0);
-            auto mc         = ggml_ext_chunk(c, modulation->forward(ctx, clean), 4, 0);
+            auto mn         = ggml_ext_chunk(ctx, modulation->forward(ctx, noisy), 4, 0);
+            auto mc         = ggml_ext_chunk(ctx, modulation->forward(ctx, clean), 4, 0);
             auto select     = [&](int index, bool gate) {
                 auto n      = gate ? ggml_tanh(c, mn[index]) : mn[index];
                 auto k      = gate ? ggml_tanh(c, mc[index]) : mc[index];
@@ -311,12 +309,12 @@ namespace Fibo2 {
             auto embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
             auto noisy    = embedder->forward(ctx, t);
             auto clean    = embedder->forward(ctx, ggml_scale(c, t, 0.f));
-            auto image_pe = ggml_ext_slice(c, pe, 3, gist_count, pe->ne[3]);
+            auto image_pe = ggml_ext_slice(ctx, pe, 3, gist_count, pe->ne[3]);
             for (int i = 0; i < config.refiner_depth; ++i) {
                 x = std::dynamic_pointer_cast<TransformerBlock>(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, x, image_pe, noisy, clean, 0, image_count);
                 sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.noise_refiner." + std::to_string(i), "x");
             }
-            auto bundle      = [&](int i) { return ggml_ext_slice(c, text, 0, i * config.text_dim, (i + 1) * config.text_dim); };
+            auto bundle      = [&](int i) { return ggml_ext_slice(ctx, text, 0, i * config.text_dim, (i + 1) * config.text_dim); };
             auto gist        = std::dynamic_pointer_cast<Perceiver>(blocks["perceiver"])->forward(ctx, bundle(0), gist_count);
             x                = ggml_concat(c, gist, x, 1);
             int bundle_index = 1;
@@ -326,7 +324,7 @@ namespace Fibo2 {
                 x            = block->forward(ctx, x, pe, noisy, clean, gist_count, image_count, context);
                 sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.layers." + std::to_string(i), "x");
             }
-            x = ggml_ext_slice(c, x, 1, gist_count, gist_count + image_count);
+            x = ggml_ext_slice(ctx, x, 1, gist_count, gist_count + image_count);
             x = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"])->forward(ctx, x, noisy);
             x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3));
             return ggml_reshape_4d(c, x, width, height, config.channels, 1);

src/model/diffusion/flux.hpp

@@ -335,9 +335,9 @@ namespace Flux {
 
             x = mlp_0->forward(ctx, x);
             if (use_mlp_silu_act) {
-                x = ggml_ext_silu_act(ctx->ggml_ctx, x);
+                x = ggml_ext_silu_act(ctx, x);
             } else {
-                x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+                x = ggml_ext_gelu(ctx, x, true);
             }
             x = mlp_2->forward(ctx, x);
             return x;
@@ -686,11 +686,11 @@ namespace Flux {
 
             auto mlp = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, mlp_hidden_dim * mlp_mult_factor, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * 3 * qkv_mlp->nb[0]);
             if (use_yak_mlp) {
-                mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp, false);
+                mlp = ggml_ext_silu_act(ctx, mlp, false);
             } else if (use_mlp_silu_act) {
-                mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp);
+                mlp = ggml_ext_silu_act(ctx, mlp);
             } else {
-                mlp = ggml_ext_gelu(ctx->ggml_ctx, mlp, true);
+                mlp = ggml_ext_gelu(ctx, mlp, true);
             }
             auto attn_mlp = ggml_concat(ctx->ggml_ctx, attn, mlp, 0);  // [N, n_token, hidden_size + mlp_hidden_dim]
             auto output   = linear2->forward(ctx, attn_mlp);           // [N, n_token, hidden_size]
@@ -745,7 +745,7 @@ namespace Flux {
                 auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
                 auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 2 * hidden_size]
-                auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
+                auto m_vec = ggml_ext_chunk(ctx, m, 2, 0);
                 shift      = m_vec[0];
                 scale      = m_vec[1];
             }
@@ -832,16 +832,16 @@ namespace Flux {
             int64_t hidden_size_x = x->ne[0];
 
             auto mlp_params = param_generator->forward(ctx, s);
-            auto fc_params  = ggml_ext_chunk(ctx->ggml_ctx, mlp_params, 3, 0);
+            auto fc_params  = ggml_ext_chunk(ctx, mlp_params, 3, 0);
             auto fc1_gate   = ggml_reshape_3d(ctx->ggml_ctx, fc_params[0], hidden_size_x * mlp_ratio, hidden_size_x, batch_size);
             auto fc1_value  = ggml_reshape_3d(ctx->ggml_ctx, fc_params[1], hidden_size_x * mlp_ratio, hidden_size_x, batch_size);
             auto fc2        = ggml_reshape_3d(ctx->ggml_ctx, fc_params[2], hidden_size_x, mlp_ratio * hidden_size_x, batch_size);
 
-            fc1_gate  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_gate, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
+            fc1_gate  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_gate, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
             fc1_gate  = ggml_l2_norm(ctx->ggml_ctx, fc1_gate, 1e-12f);
-            fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_value, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
+            fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_value, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
             fc1_value = ggml_l2_norm(ctx->ggml_ctx, fc1_value, 1e-12f);
-            fc2       = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc2, 1, 0, 2, 3));  // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio]
+            fc2       = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc2, 1, 0, 2, 3));  // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio]
             fc2       = ggml_l2_norm(ctx->ggml_ctx, fc2, 1e-12f);
 
             auto res_x = x;
@@ -892,9 +892,9 @@ namespace Flux {
             auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [N, H, W, C]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, H, W, C]
             x = norm->forward(ctx, x);
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));  // [N, C, H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));  // [N, C, H, W]
             x = conv->forward(ctx, x);
 
             return x;
@@ -1028,13 +1028,13 @@ namespace Flux {
             if (config.is_chroma) {
                 int64_t mod_index_length = 344;
                 auto approx              = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
-                auto distill_timestep    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f);
-                auto distill_guidance    = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 16, 10000, 1000.f);
+                auto distill_timestep    = ggml_ext_timestep_embedding(ctx, timesteps, 16, 10000, 1000.f);
+                auto distill_guidance    = ggml_ext_timestep_embedding(ctx, guidance, 16, 10000, 1000.f);
 
                 // auto mod_index_arange  = ggml_arange(ctx, 0, (float)mod_index_length, 1);
                 // ggml_arange tot working on a lot of backends, precomputing it on CPU instead
                 GGML_ASSERT(mod_index_arange != nullptr);
-                auto modulation_index = ggml_ext_timestep_embedding(ctx->ggml_ctx, mod_index_arange, 32, 10000, 1000.f);  // [1, 344, 32]
+                auto modulation_index = ggml_ext_timestep_embedding(ctx, mod_index_arange, 32, 10000, 1000.f);  // [1, 344, 32]
 
                 // Batch broadcast (will it ever be useful)
                 modulation_index = ggml_repeat(ctx->ggml_ctx, modulation_index, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, modulation_index->ne[0], modulation_index->ne[1], img->ne[2]));  // [N, 344, 32]
@@ -1057,12 +1057,12 @@ namespace Flux {
                 vec                  = dual_time_embed->forward(ctx, timestep_sem, timestep_tex);
             } else {
                 auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
-                vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
+                vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f));
                 if (config.guidance_embed) {
                     GGML_ASSERT(guidance != nullptr);
                     auto guidance_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["guidance_in"]);
                     // bf16 and fp16 result is different
-                    auto g_in = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1000.f);
+                    auto g_in = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1000.f);
                     vec       = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in));
                 }
 
@@ -1223,7 +1223,7 @@ namespace Flux {
 
             img = img_in_patch->forward(ctx, img);                                                       // [N, hidden_size, H/patch_size, W/patch_size]
             img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], img->ne[3]);  // [N, hidden_size, H/patch_size*W/patch_size]
-            img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));      // [N, H/patch_size*W/patch_size, hidden_size]
+            img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));                // [N, H/patch_size*W/patch_size, hidden_size]
 
             auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers,
                                     pulid_id, pulid_id_weight);  // [N, n_img_token, hidden_size]
@@ -1232,14 +1232,14 @@ namespace Flux {
             auto nerf_image_embedder   = std::dynamic_pointer_cast<NerfEmbedder>(blocks["nerf_image_embedder"]);
             auto nerf_final_layer_conv = std::dynamic_pointer_cast<NerfFinalLayerConv>(blocks["nerf_final_layer_conv"]);
 
-            auto nerf_pixels    = DiT::patchify(ctx->ggml_ctx, orig_img, patch_size, patch_size);  // [N, num_patches, C * patch_size * patch_size]
+            auto nerf_pixels    = DiT::patchify(ctx, orig_img, patch_size, patch_size);  // [N, num_patches, C * patch_size * patch_size]
             int64_t num_patches = nerf_pixels->ne[1];
             nerf_pixels         = ggml_reshape_3d(ctx->ggml_ctx,
                                                   nerf_pixels,
                                                   nerf_pixels->ne[0] / C,
                                                   C,
-                                                  nerf_pixels->ne[1] * nerf_pixels->ne[2]);                                  // [N*num_patches, C, patch_size*patch_size]
-            nerf_pixels         = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, nerf_pixels, 1, 0, 2, 3));  // [N*num_patches, patch_size*patch_size, C]
+                                                  nerf_pixels->ne[1] * nerf_pixels->ne[2]);                        // [N*num_patches, C, patch_size*patch_size]
+            nerf_pixels         = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, nerf_pixels, 1, 0, 2, 3));  // [N*num_patches, patch_size*patch_size, C]
 
             auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]);  // [N*num_patches, hidden_size]
             auto img_dct     = nerf_image_embedder->forward(ctx, nerf_pixels, dct);                       // [N*num_patches, patch_size*patch_size, nerf_hidden_size]
@@ -1250,9 +1250,9 @@ namespace Flux {
                 img_dct = block->forward(ctx, img_dct, nerf_hidden);
             }
 
-            img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img_dct, 1, 0, 2, 3));                                 // [N*num_patches, nerf_hidden_size, patch_size*patch_size]
+            img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img_dct, 1, 0, 2, 3));                                           // [N*num_patches, nerf_hidden_size, patch_size*patch_size]
             img_dct = ggml_reshape_3d(ctx->ggml_ctx, img_dct, img_dct->ne[0] * img_dct->ne[1], num_patches, img_dct->ne[2] / num_patches);  // [N, num_patches, nerf_hidden_size*patch_size*patch_size]
-            img_dct = DiT::unpatchify(ctx->ggml_ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size);  // [N, nerf_hidden_size, H, W]
+            img_dct = DiT::unpatchify(ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size);            // [N, nerf_hidden_size, H, W]
 
             out = nerf_final_layer_conv->forward(ctx, img_dct);  // [N, C, H, W]
 
@@ -1331,7 +1331,7 @@ namespace Flux {
                 out = ggml_cont(ctx->ggml_ctx, out);
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size);  // [N, C, H, W]
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size);  // [N, C, H, W]
             return out;
         }
 

src/model/diffusion/flux3.hpp

@@ -95,7 +95,7 @@ namespace Flux3 {
             auto k           = norm->key_norm(ctx, split(linear(ctx, "k_proj", x_mod)));
             auto v           = split(linear(ctx, "v_proj", x_mod));
             auto attn        = Rope::attention(ctx, q, k, v, pe, nullptr);
-            auto mlp         = ggml_ext_silu_act(gc, linear(ctx, "mlp_in", x_mod));
+            auto mlp         = ggml_ext_silu_act(ctx, linear(ctx, "mlp_in", x_mod));
             auto update      = ggml_add(gc, linear(ctx, "attn_out", attn), linear(ctx, "mlp_out", mlp));
             ggml_tensor* out = nullptr;
             offset           = 0;
@@ -140,16 +140,16 @@ namespace Flux3 {
             auto gc             = ctx->ggml_ctx;
             const int64_t width = image->ne[0], height = image->ne[1];
             auto time_in   = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
-            auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, timestep, 256, 10000, 1000.f));
-            auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, zero, 256, 10000, 1000.f));
+            auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1000.f));
+            auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, zero, 256, 10000, 1000.f));
             if (config.flow_map) {
                 auto ref_time = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["ref_time_in"]);
-                auto ref_vec  = ref_time->forward(ctx, ggml_ext_timestep_embedding(gc, target, 256, 10000, 1000.f));
+                auto ref_vec  = ref_time->forward(ctx, ggml_ext_timestep_embedding(ctx, target, 256, 10000, 1000.f));
                 image_vec     = ggml_add(gc, image_vec, ref_vec);
                 clean_vec     = ggml_add(gc, clean_vec, ref_vec);
             }
             std::vector<std::string> modes{"image"};
-            std::vector<ggml_tensor*> inputs{DiT::patchify(gc, image, 1, 1)};
+            std::vector<ggml_tensor*> inputs{DiT::patchify(ctx, image, 1, 1)};
             if (reference) {
                 modes.push_back("image_cond");
                 inputs.push_back(reference);
@@ -184,7 +184,7 @@ namespace Flux3 {
                 sd::ggml_graph_cut::mark_graph_cut(joint, "flux3.single_blocks." + std::to_string(i), "joint");
             }
             auto out = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer.image"])->forward(ctx, token_slice(gc, joint, 0, lengths[0]), image_vec);
-            return DiT::unpatchify(gc, out, height, width, 1, 1);
+            return DiT::unpatchify(ctx, out, height, width, 1, 1);
         }
     };
 
@@ -215,9 +215,10 @@ namespace Flux3 {
             GGML_ASSERT(params.x && params.context && params.timesteps);
             GGML_ASSERT(!config.flow_map || extra->target_timestep);
             auto get_graph = [&]() {
-                auto gf   = new_graph_custom(FLUX_GRAPH_SIZE);
-                auto x    = make_input(*params.x);
-                auto text = make_input(*params.context);
+                auto gf         = new_graph_custom(FLUX_GRAPH_SIZE);
+                auto runner_ctx = get_context();
+                auto x          = make_input(*params.x);
+                auto text       = make_input(*params.context);
                 GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1);
                 std::vector<ggml_tensor*> refs;
                 ggml_tensor* references = nullptr;
@@ -225,13 +226,13 @@ namespace Flux3 {
                     for (const auto& ref : *params.ref_latents) {
                         auto r = make_input(ref);
                         refs.push_back(r);
-                        r          = DiT::patchify(compute_ctx, r, 1, 1);
+                        r          = DiT::patchify(&runner_ctx, r, 1, 1);
                         references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
                     }
                 }
                 sd::PECacheKey pe_key("flux3.pe");
                 pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled);
-                auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto pe  = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     Rope::Embedding embedding;
                     auto append_image = [&](int64_t w, int64_t h, float index) {
                         embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
@@ -254,9 +255,8 @@ namespace Flux3 {
                     const int64_t tokens = static_cast<int64_t>(values.size() / 256);
                     return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
                 }));
-                auto runner_ctx = get_context();
-                auto out        = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
-                                                config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
+                auto out = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
+                                         config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
                 ggml_build_forward_expand(gf, out);
                 return gf;
             };

src/model/diffusion/hidream_o1.hpp

@@ -197,7 +197,7 @@ namespace HiDreamO1 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto emb   = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f);
+            auto emb   = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1000.0f);
             emb        = mlp_0->forward(ctx, emb);
             emb        = ggml_silu_inplace(ctx->ggml_ctx, emb);
             emb        = mlp_2->forward(ctx, emb);
@@ -292,7 +292,6 @@ namespace HiDreamO1 {
 
         ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             return LLM::LLMRunner::encode_image_common(this,
-                                                       compute_ctx,
                                                        runner_ctx,
                                                        image,
                                                        config.llm.vision,
@@ -398,7 +397,7 @@ namespace HiDreamO1 {
             auto t_emb          = t_embedder1->forward(&runner_ctx, timestep);
             int64_t txt_seq_len = txt->ne[1];
             if (txt_seq_len > 1) {
-                auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
+                auto prefix = ggml_ext_slice(&runner_ctx, txt, 1, 0, txt_seq_len - 1);
                 txt         = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
             } else {
                 txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1);
@@ -425,8 +424,8 @@ namespace HiDreamO1 {
                 }
                 x_pred_start = first_vinput;
             }
-            auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
-            x_pred      = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
+            auto x_pred = ggml_ext_slice(&runner_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
+            x_pred      = DiT::unpatchify_and_crop(&runner_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
 
             float sigma = 1.0f - timestep_tensor.values()[0];
             sigma       = std::max(1e-6f, sigma);

src/model/diffusion/hunyuan.hpp

@@ -50,7 +50,7 @@ namespace Hunyuan {
 
             // self attn
             auto qkv     = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt));
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             auto q       = qkv_vec[0];
             auto k       = qkv_vec[1];
             auto v       = qkv_vec[2];
@@ -60,7 +60,7 @@ namespace Hunyuan {
 
             // adaLN_modulation
             auto emb  = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb));
-            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto mods = ggml_ext_chunk(ctx, emb, 2, 0);
 
             txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0]));
 
@@ -122,11 +122,11 @@ namespace Hunyuan {
             auto c_embedder               = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["c_embedder"]);
             auto individual_token_refiner = std::dynamic_pointer_cast<IndividualTokenRefiner>(blocks["individual_token_refiner"]);
 
-            auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+            auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f));
 
-            auto h                  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3));
+            auto h                  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, txt, 1, 0, 2, 3));
             auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]);
-            pooled_projections      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3));
+            pooled_projections      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, pooled_projections, 1, 0, 2, 3));
             auto c_emb              = c_embedder->forward(ctx, pooled_projections);
 
             t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb);
@@ -151,9 +151,9 @@ namespace Hunyuan {
             auto fc3       = std::dynamic_pointer_cast<Linear>(blocks["fc3"]);
 
             x = fc1->forward(ctx, layernorm->forward(ctx, x));
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             x = fc2->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             return fc3->forward(ctx, x);
         }
     };
@@ -364,11 +364,12 @@ namespace Hunyuan {
             return x;
         }
 
-        ggml_tensor* unpatchify(struct ggml_context* ctx,
+        ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                 ggml_tensor* x,
                                 int64_t t_len,
                                 int64_t h_len,
                                 int64_t w_len) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
             // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
             int64_t N  = x->ne[3];
@@ -379,15 +380,15 @@ namespace Hunyuan {
 
             GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
-            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
             return x;
         }
 
@@ -431,10 +432,10 @@ namespace Hunyuan {
 
             img      = img_in->forward(ctx, img);                     // [N*C, t_len*h_len*w_len, hidden_size]
             txt      = txt_in->forward(ctx, txt, timestep, nullptr);  // [N, n_txt_token, hidden_size]
-            auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+            auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f));
             if (config.use_meanflow && timestep_r != nullptr) {
                 auto time_r_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_r_in"]);
-                auto vec_r     = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f));
+                auto vec_r     = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep_r, 256, 10000, 1000.f));
                 vec            = ggml_add(ctx->ggml_ctx, vec, vec_r);
                 if (!config.use_meanflow_sum) {
                     vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f);
@@ -446,7 +447,7 @@ namespace Hunyuan {
             }
             if (config.guidance_embed && guidance != nullptr) {
                 auto guidance_in  = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["guidance_in"]);
-                auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f);
+                auto guidance_emb = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1.f);
                 vec               = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb));
             }
 
@@ -529,12 +530,12 @@ namespace Hunyuan {
 
             auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N);
 
-            out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
+            out = unpatchify(ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
             // slice
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            out = ggml_ext_slice(ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
 
             return out;
         }

src/model/diffusion/ideogram4.hpp

@@ -73,12 +73,13 @@ namespace Ideogram4 {
         }
     };
 
-    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx,
                                                               ggml_tensor* timesteps,
                                                               int dim) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(dim % 2 == 0);
-        auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
-        auto chunks    = ggml_ext_chunk(ctx, embedding, 2, 0);
+        auto embedding = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
+        auto chunks    = ggml_ext_chunk(runner_ctx, embedding, 2, 0);
         return ggml_concat(ctx, chunks[1], chunks[0], 0);
     }
 
@@ -108,7 +109,10 @@ namespace Ideogram4 {
         return ggml_ext_modulate(ctx, x, nullptr, scale);
     }
 
-    __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
+    __STATIC_INLINE__ ggml_tensor* patchify(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* x,
+                                            const Ideogram4Config& config) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, 128, H, W] with channel order [ae, ph, pw].
         // return: [N, H*W, 128] with token channel order [ph, pw, ae].
         const int64_t W = x->ne[0];
@@ -126,11 +130,12 @@ namespace Ideogram4 {
         return x;
     }
 
-    __STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                               ggml_tensor* x,
                                               int64_t H,
                                               int64_t W,
                                               const Ideogram4Config& config) {
+        auto ctx        = runner_ctx->ggml_ctx;
         const int64_t C = x->ne[0];
         const int64_t N = x->ne[2];
 
@@ -217,7 +222,7 @@ namespace Ideogram4 {
             auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
 
             auto qkv     = qkv_proj->forward(ctx, x);
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             auto q       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
             auto k       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
             auto v       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
@@ -278,7 +283,7 @@ namespace Ideogram4 {
             auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
 
             auto mod       = adaln_modulation->forward(ctx, adaln_input);
-            auto mods      = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0);
+            auto mods      = ggml_ext_chunk(ctx, mod, 4, 0);
             auto scale_msa = mods[0];
             auto gate_msa  = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1]));
             auto scale_mlp = mods[2];
@@ -315,7 +320,7 @@ namespace Ideogram4 {
             auto mlp_in  = std::dynamic_pointer_cast<Linear>(blocks["mlp_in"]);
             auto mlp_out = std::dynamic_pointer_cast<Linear>(blocks["mlp_out"]);
 
-            x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast<int>(dim));
+            x = timestep_embedding_sin_cos(ctx, x, static_cast<int>(dim));
             x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x));
             x = mlp_out->forward(ctx, x);
             return x;
@@ -383,7 +388,7 @@ namespace Ideogram4 {
             auto embed_image_indicator = std::dynamic_pointer_cast<Embedding>(blocks["embed_image_indicator"]);
             auto final_layer           = std::dynamic_pointer_cast<Ideogram4FinalLayer>(blocks["final_layer"]);
 
-            auto img = patchify(ctx->ggml_ctx, x, config);
+            auto img = patchify(ctx, x, config);
             img      = input_proj->forward(ctx, img);
 
             ggml_tensor* h      = img;
@@ -413,11 +418,11 @@ namespace Ideogram4 {
 
             h = final_layer->forward(ctx, h, adaln_input);
             if (context_len > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]);
+                h = ggml_ext_slice(ctx, h, 1, context_len, h->ne[1]);
             }
 
-            h = unpatchify(ctx->ggml_ctx, h, H, W, config);
-            h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f);
+            h = unpatchify(ctx, h, H, W, config);
+            h = ggml_ext_scale(ctx, h, -1.f);
             return h;
         }
     };

src/model/diffusion/iris.hpp

@@ -92,7 +92,7 @@ namespace Iris {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) {
             auto c   = ctx->ggml_ctx;
-            auto qkv = split_qkv(c, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
+            auto qkv = ggml_ext_chunk(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x), 3, 0);
             auto q   = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]);
             auto k   = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]);
             auto v   = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]);
@@ -247,24 +247,24 @@ namespace Iris {
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* y, ggml_tensor* img_mod, ggml_tensor* txt_mod, ggml_tensor* pe) {
             auto c     = ctx->ggml_ctx;
             int64_t nt = y->ne[1];
-            auto mx    = ggml_ext_chunk(c, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0);
+            auto mx    = ggml_ext_chunk(ctx, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0);
             if (!dual) {
                 auto tokens = ggml_concat(c, y, x, 1);
                 auto h      = Pid::apply_adaln(c, norm(ctx, "norm1", tokens), mx[0], mx[1]);
                 auto qkv    = project(ctx, h, "");
                 auto out    = Rope::attention(ctx, qkv[0], qkv[1], qkv[2], pe, nullptr);
                 tokens      = finish(ctx, tokens, h, out, mx, "");
-                return {ggml_ext_slice(c, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(c, tokens, 1, 0, nt)};
+                return {ggml_ext_slice(ctx, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(ctx, tokens, 1, 0, nt)};
             }
-            auto my  = ggml_ext_chunk(c, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0);
+            auto my  = ggml_ext_chunk(ctx, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0);
             auto hx  = Pid::apply_adaln(c, norm(ctx, "norm_x1", x), mx[0], mx[1]);
             auto hy  = Pid::apply_adaln(c, norm(ctx, "norm_y1", y), my[0], my[1]);
             auto qx  = project(ctx, hx, "_x");
             auto qy  = project(ctx, hy, "_y");
             auto out = Rope::attention(ctx, ggml_concat(c, qy[0], qx[0], 2), ggml_concat(c, qy[1], qx[1], 2),
                                        ggml_concat(c, qy[2], qx[2], 2), pe, nullptr);
-            x        = finish(ctx, x, hx, ggml_ext_slice(c, out, 1, nt, out->ne[1]), mx, "_x");
-            y        = finish(ctx, y, hy, ggml_ext_slice(c, out, 1, 0, nt), my, "_y");
+            x        = finish(ctx, x, hx, ggml_ext_slice(ctx, out, 1, nt, out->ne[1]), mx, "_x");
+            y        = finish(ctx, y, hy, ggml_ext_slice(ctx, out, 1, 0, nt), my, "_y");
             return {x, y};
         }
     };
@@ -294,7 +294,7 @@ namespace Iris {
             auto batch   = cond->ne[2];
             auto mods    = std::dynamic_pointer_cast<Linear>(blocks["adaln"])->forward(ctx, cond);
             mods         = ggml_reshape_3d(c, mods, 4 * d, p2, patches * batch);
-            auto m       = ggml_ext_chunk(c, mods, 4, 0);
+            auto m       = ggml_ext_chunk(ctx, mods, 4, 0);
             auto h       = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"])->forward(ctx, x);
             h            = ggml_reshape_3d(c, h, d * p2, patches, batch);
             h            = std::dynamic_pointer_cast<Linear>(blocks["compress"])->forward(ctx, h);
@@ -337,7 +337,7 @@ namespace Iris {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* t, ggml_tensor* text, ggml_tensor* mask, ggml_tensor* joint_pe, ggml_tensor* pixel_pe, ggml_tensor* pixel_pos) {
             auto c    = ctx->ggml_ctx;
             int p     = static_cast<int>(config.patch_size);
-            auto x    = DiT::patchify(c, image, p, p, true);
+            auto x    = DiT::patchify(ctx, image, p, p, true);
             x         = std::dynamic_pointer_cast<Pid::PatchTokenEmbedder>(blocks["s_embedder"])->forward(ctx, x);
             auto te   = std::dynamic_pointer_cast<Pid::PixelDiTTimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
             te        = ggml_reshape_3d(c, te, config.hidden_size, 1, image->ne[3]);
@@ -360,7 +360,7 @@ namespace Iris {
             }
             x = std::dynamic_pointer_cast<Pid::FinalLayer>(blocks["final_layer"])->forward(ctx, x);
             x = ggml_reshape_3d(c, x, 3 * p * p, cond->ne[1], image->ne[3]);
-            return DiT::unpatchify(c, x, image->ne[1] / p, image->ne[0] / p, p, p, false);
+            return DiT::unpatchify(ctx, x, image->ne[1] / p, image->ne[0] / p, p, p, false);
         }
     };
 

src/model/diffusion/krea2.hpp

@@ -185,9 +185,9 @@ namespace Krea2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* scale = params["scale"];
             if (ctx->weight_adapter) {
-                scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
+                scale = ctx->weight_adapter->patch_weight(ctx, scale, prefix + "scale.weight");
             }
-            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
+            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx, scale->ne[0], 1, 1, 1));
             x     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
             x     = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
             return x;
@@ -273,9 +273,7 @@ namespace Krea2 {
             auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
             auto wo    = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                wo->set_force_prec_f32(true);
-            }
+            wo->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t L = x->ne[1];
             int64_t N = x->ne[2];
@@ -316,11 +314,11 @@ namespace Krea2 {
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             auto lin = params["lin"];
             if (ctx->weight_adapter) {
-                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+                lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight");
             }
             lin      = ggml_repeat(ctx->ggml_ctx, lin, vec);
             auto out = ggml_add(ctx->ggml_ctx, vec, lin);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
+            return ggml_ext_chunk(ctx, out, 6, 0);
         }
     };
 
@@ -342,10 +340,10 @@ namespace Krea2 {
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             auto lin = params["lin"];
             if (ctx->weight_adapter) {
-                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+                lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight");
             }
             auto out = ggml_add(ctx->ggml_ctx, lin, vec);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
     };
 
@@ -541,7 +539,7 @@ namespace Krea2 {
             auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
             x             = linear_0->forward(ctx, x);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = ggml_ext_gelu(ctx, x, false);
             x             = linear_2->forward(ctx, x);
             return x;
         }
@@ -555,7 +553,7 @@ namespace Krea2 {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = ggml_ext_gelu(ctx, x, false);
             x             = linear_1->forward(ctx, x);
             return x;
         }
@@ -575,7 +573,7 @@ namespace Krea2 {
             auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["3"]);
             x             = norm->forward(ctx, x);
             x             = linear_1->forward(ctx, x);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x             = ggml_ext_gelu(ctx, x, true);
             x             = linear_3->forward(ctx, x);
             return x;
         }
@@ -656,7 +654,7 @@ namespace Krea2 {
             int64_t ref_len = img->ne[1] - img_len;
             img             = first->forward(ctx, img);
 
-            auto t    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+            auto t    = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
             t         = tmlp->forward(ctx, t);
             t         = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
             auto tvec = tproj->forward(ctx, t);
@@ -665,7 +663,7 @@ namespace Krea2 {
             if (ref_latents.size() > 0 && zero_timestep_refs) {
                 // "index_timestep_zero" mode: use timestep = 0 for ref latents
                 auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
-                auto t_0        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+                auto t_0        = ggml_ext_timestep_embedding(ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
                 t_0             = tmlp->forward(ctx, t_0);
                 t_0             = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
                 tvec_0          = tproj->forward(ctx, t_0);
@@ -683,9 +681,9 @@ namespace Krea2 {
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
             }
 
-            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
+            hidden_states = ggml_ext_slice(ctx, hidden_states, 1, txt_len, txt_len + img_len);
             hidden_states = last->forward(ctx, hidden_states, t);
-            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
+            hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
             return hidden_states;
         }
     };

src/model/diffusion/lens.hpp

@@ -87,7 +87,7 @@ namespace Lens {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
             auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
-            auto timesteps_proj    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f);
+            auto timesteps_proj    = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f);
             return timestep_embedder->forward(ctx, timesteps_proj);
         }
     };
@@ -151,8 +151,8 @@ namespace Lens {
             int64_t N       = img->ne[2];
             int64_t inner   = dim_head * num_heads;
 
-            auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img));
-            auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt));
+            auto img_qkv_vec = ggml_ext_chunk(ctx, img_qkv->forward(ctx, img), 3, 0);
+            auto txt_qkv_vec = ggml_ext_chunk(ctx, txt_qkv->forward(ctx, txt), 3, 0);
 
             auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
             auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
@@ -232,9 +232,9 @@ namespace Lens {
             auto temb = ggml_silu(ctx->ggml_ctx, t_emb);
 
             auto img_mod_params = img_mod_1->forward(ctx, temb);
-            auto img_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
+            auto img_mod_vec    = ggml_ext_chunk(ctx, img_mod_params, 6, 0);
             auto txt_mod_params = txt_mod_1->forward(ctx, temb);
-            auto txt_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
+            auto txt_mod_vec    = ggml_ext_chunk(ctx, txt_mod_params, 6, 0);
 
             auto img_normed    = img_norm1->forward(ctx, img);
             auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
@@ -268,7 +268,7 @@ namespace Lens {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
-            auto mods   = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
+            auto mods   = ggml_ext_chunk(ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
             auto scale  = mods[0];
             auto shift  = mods[1];
             x           = ggml_norm(ctx->ggml_ctx, x, eps);
@@ -319,10 +319,10 @@ namespace Lens {
             auto t_emb = time_text_embed->forward(ctx, timestep);
 
             auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N);
-            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));
             img      = img_in->forward(ctx, img);
 
-            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0);
+            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx, context, config.selected_layer_count, 0);
             ggml_tensor* txt                     = nullptr;
             for (int i = 0; i < config.selected_layer_count; ++i) {
                 auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
@@ -346,7 +346,7 @@ namespace Lens {
             img = norm_out->forward(ctx, img, t_emb);
             img = proj_out->forward(ctx, img);
 
-            auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
+            auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));
             out      = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N);
             return out;
         }

src/model/diffusion/lingbot_video.hpp

@@ -365,7 +365,7 @@ namespace LingBotVideo {
             }
             selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
 
-            ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
+            ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(ctx, selected_mask));
             ggml_tensor* masked_scores       = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
             return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
         }
@@ -476,7 +476,7 @@ namespace LingBotVideo {
 
             auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
             auto mod   = ggml_add(ctx->ggml_ctx, temb6, table);
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
+            auto mods  = ggml_ext_chunk(ctx, mod, 6, 0);
 
             auto shift_msa = mods[0];
             auto scale_msa = mods[1];
@@ -556,13 +556,13 @@ namespace LingBotVideo {
             int64_t w_len = W / config.patch_w;
             int64_t n_img = t_len * h_len * w_len;
 
-            auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
+            auto img = DiT::patchify_3d(ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
             img      = patch_embedder->forward(ctx, img);
 
             auto txt           = text_embedder->forward(ctx, context);
             auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
 
-            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx,
                                                              timestep,
                                                              static_cast<int>(config.freq_dim),
                                                              10000,
@@ -582,16 +582,16 @@ namespace LingBotVideo {
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
             }
 
-            auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
+            auto final_mods = ggml_ext_chunk(ctx,
                                              norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
                                              2,
                                              0);
             hidden_states   = norm_out->forward(ctx, hidden_states);
             hidden_states   = Flux::modulate(ctx, hidden_states, final_mods[0], final_mods[1], true);
             hidden_states   = proj_out->forward(ctx, hidden_states);
-            hidden_states   = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
+            hidden_states   = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img);
 
-            auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
+            auto out = DiT::unpatchify_3d(ctx,
                                           hidden_states,
                                           t_len,
                                           h_len,

src/model/diffusion/llada_image.hpp

@@ -217,8 +217,8 @@ namespace LLaDAImage {
 
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
 
             for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
@@ -246,7 +246,7 @@ namespace LLaDAImage {
 
             txt_img = final_layer->forward(ctx, txt_img, t_emb);
 
-            return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
+            return ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
         }
 
         ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
@@ -308,8 +308,8 @@ namespace LLaDAImage {
 
             GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
 
-            auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
+            auto cap_pe = ggml_ext_slice(ctx, pe, 3, 0, cap_len * 2);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
 
             auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
 
@@ -322,7 +322,7 @@ namespace LLaDAImage {
                 img        = block->forward(ctx, img, img_pe, nullptr, img_adaln);
             }
             if (sig != nullptr) {
-                auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
+                auto sig_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
                 for (int i = 0; i < config.num_refiner_layers; i++) {
                     auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
                     sig        = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
@@ -350,7 +350,7 @@ namespace LLaDAImage {
             // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
             // the latent grid that unpatchify reconstructs.
             int64_t target_start = cap_len * 2 + img_len;
-            return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
+            return ggml_ext_slice(ctx, seq, 1, target_start, target_start + n_img_token);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -372,10 +372,10 @@ namespace LLaDAImage {
 
             auto out = forward_core(ctx, img, timestep, context, pe);
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false);
 
             // The reference pipeline negates the model output before the scheduler step.
-            return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
+            return ggml_ext_scale(ctx, out, -1.f);
         }
     };
 
@@ -488,8 +488,8 @@ namespace LLaDAImage {
             auto src    = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
 
             auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
-            out      = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
-            out      = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
+            out      = DiT::unpatchify_and_crop(&runner_ctx, out, H, W, config.patch_size, config.patch_size, false);
+            out      = ggml_ext_scale(&runner_ctx, out, -1.f);
 
             ggml_build_forward_expand(gf, out);
             return gf;

src/model/diffusion/ltxv.hpp

@@ -581,7 +581,7 @@ namespace LTXV {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
 
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, frequency_embedding_size);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, timestep, frequency_embedding_size);
             t_emb      = linear_1->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             t_emb      = linear_2->forward(ctx, t_emb);
@@ -631,7 +631,7 @@ namespace LTXV {
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
 
             caption = linear_1->forward(ctx, caption);
-            caption = ggml_ext_gelu(ctx->ggml_ctx, caption, true);
+            caption = ggml_ext_gelu(ctx, caption, true);
             caption = linear_2->forward(ctx, caption);
             return caption;
         }
@@ -650,7 +650,7 @@ namespace LTXV {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* caption) {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
             caption       = ggml_rms_norm(ctx->ggml_ctx, caption, 1e-6f);
-            caption       = ggml_ext_scale(ctx->ggml_ctx, caption, std::sqrt(static_cast<float>(hidden_size) / static_cast<float>(in_features)));
+            caption       = ggml_ext_scale(ctx, caption, std::sqrt(static_cast<float>(hidden_size) / static_cast<float>(in_features)));
             return linear_1->forward(ctx, caption);
         }
     };
@@ -724,7 +724,7 @@ namespace LTXV {
                 auto to_gate_logits = std::dynamic_pointer_cast<Linear>(blocks["to_gate_logits"]);
                 auto gate_logits    = to_gate_logits->forward(ctx, x);
                 auto gates          = ggml_sigmoid(ctx->ggml_ctx, gate_logits);
-                gates               = ggml_ext_scale(ctx->ggml_ctx, gates, 2.0f, true);
+                gates               = ggml_ext_scale(ctx, gates, 2.0f, true);
                 gates               = ggml_reshape_4d(ctx->ggml_ctx, gates, 1, heads, gate_logits->ne[1], gate_logits->ne[2]);
 
                 auto out4 = ggml_reshape_4d(ctx->ggml_ctx, out, dim_head, heads, out->ne[1], out->ne[2]);
@@ -783,7 +783,7 @@ namespace LTXV {
             s             = ggml_repeat(ctx->ggml_ctx, s, e);
             t             = ggml_repeat(ctx->ggml_ctx, t, e);
             auto out      = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, static_cast<int>(coeff), 1);
+            return ggml_ext_chunk(ctx, out, static_cast<int>(coeff), 1);
         }
 
         std::vector<ggml_tensor*> get_prompt_scale_shift_values(GGMLRunnerContext* ctx,
@@ -797,7 +797,7 @@ namespace LTXV {
             s        = ggml_repeat(ctx->ggml_ctx, s, e);
             t        = ggml_repeat(ctx->ggml_ctx, t, e);
             auto out = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -945,7 +945,7 @@ namespace LTXV {
             for (int64_t i = 1; i < duplications; i++) {
                 regs_full = ggml_concat(ctx->ggml_ctx, regs_full, regs, 1);
             }
-            regs_full = ggml_ext_slice(ctx->ggml_ctx, regs_full, 1, seq_len, seq_len + total_to_keep);
+            regs_full = ggml_ext_slice(ctx, regs_full, 1, seq_len, seq_len + total_to_keep);
             return ggml_concat(ctx->ggml_ctx, hidden_states, regs_full, 1);
         }
 
@@ -1182,7 +1182,7 @@ namespace LTXV {
             t           = ggml_repeat(ctx->ggml_ctx, t, e);
             s           = ggml_repeat(ctx->ggml_ctx, s, e);
             auto out    = ggml_add(ctx->ggml_ctx, s, t);
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, out, static_cast<int>(coeff), 1);
+            auto chunks = ggml_ext_chunk(ctx, out, static_cast<int>(coeff), 1);
             return std::vector<ggml_tensor*>(chunks.begin() + start, chunks.begin() + start + count);
         }
 
@@ -1287,27 +1287,27 @@ namespace LTXV {
                 auto ax_norm3 = rms_norm(ctx->ggml_ctx, ax);
 
                 if (run_a2v) {
-                    auto a2v_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
-                    auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
+                    auto a2v_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
+                    auto a2v_video_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto a2v_audio       = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto a2v_video       = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
                     auto vx_scaled       = LTXV::modulate(ctx, vx_norm3, a2v_video[1], a2v_video[0]);
                     auto ax_scaled       = LTXV::modulate(ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
                     auto a2v_out         = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
-                    auto a2v_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
+                    auto a2v_gate_table  = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
                     auto a2v_gate        = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
                     vx                   = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, a2v_out, a2v_gate));
                 }
 
                 if (run_v2a) {
-                    auto v2a_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
-                    auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
+                    auto v2a_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
+                    auto v2a_video_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto v2a_audio       = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto v2a_video       = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
                     auto ax_scaled       = LTXV::modulate(ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
                     auto vx_scaled       = LTXV::modulate(ctx, vx_norm3, v2a_video[3], v2a_video[2]);
                     auto v2a_out         = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
-                    auto v2a_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
+                    auto v2a_gate_table  = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
                     auto v2a_gate        = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
                     ax                   = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, v2a_out, v2a_gate));
                 }
@@ -1427,7 +1427,7 @@ namespace LTXV {
 
         ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) {
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / n, n);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             return x;
         }
 
@@ -1436,14 +1436,14 @@ namespace LTXV {
                                       int64_t width,
                                       int64_t height,
                                       int64_t frames) {
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, x->ne[1] * x->ne[2]);
             return x;
         }
 
         ggml_tensor* patchify_audio(GGMLRunnerContext* ctx, ggml_tensor* ax) {
             // ax: [b, c, t, f]
-            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));  // [b, t, c, f]
+            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3));            // [b, t, c, f]
             ax = ggml_reshape_3d(ctx->ggml_ctx, ax, ax->ne[0] * ax->ne[1], ax->ne[2], ax->ne[3]);  // [b, t, c*f]
             return ax;
         }
@@ -1462,7 +1462,7 @@ namespace LTXV {
                 return nullptr;
             }
             ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]);  // [b, t, c, f]
-            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));                                      // [b, c, t, f]
+            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3));                                                // [b, c, t, f]
             return ax;
         }
 
@@ -1483,10 +1483,10 @@ namespace LTXV {
                 context->ne[1] < 1024;
 
             if (is_fully_processed_context) {
-                auto v_context         = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
+                auto v_context         = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim);
                 ggml_tensor* a_context = nullptr;
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
                 return {v_context, a_context};
             }
@@ -1494,14 +1494,14 @@ namespace LTXV {
             ggml_tensor* v_context = context;
             ggml_tensor* a_context = process_audio_context ? context : nullptr;
             if (is_unprocessed_dual_context) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
+                v_context = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
             } else if (context->ne[0] == config.caption_channels * 2) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels);
+                v_context = ggml_ext_slice(ctx, context, 0, 0, config.caption_channels);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.caption_channels, config.caption_channels * 2);
                 }
             }
 
@@ -1604,7 +1604,7 @@ namespace LTXV {
             auto t    = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, embedded_timestep, dim, 1, embedded_timestep->ne[1]), temp);
             auto s    = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, table, dim, 2, 1), temp);
             auto out  = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
@@ -1653,13 +1653,13 @@ namespace LTXV {
                 a_context = ggml_cont(ctx->ggml_ctx, a_context);
             }
 
-            auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier);
+            auto v_timestep_scaled = ggml_ext_scale(ctx, timestep, config.timestep_scale_multiplier);
             auto v_pair            = adaln_single->forward(ctx, v_timestep_scaled);
             auto v_timestep_mod    = v_pair.first;
             auto v_embedded_time   = v_pair.second;
 
             ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep;
-            auto a_timestep_scaled                = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier);
+            auto a_timestep_scaled                = ggml_ext_scale(ctx, effective_audio_timestep, config.timestep_scale_multiplier);
             auto a_pair                           = audio_adaln_single->forward(ctx, a_timestep_scaled);
             auto a_timestep_mod                   = a_pair.first;
             auto a_embedded_time                  = a_pair.second;
@@ -1680,13 +1680,13 @@ namespace LTXV {
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first;
             auto av_ca_a2v_gate_noise_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_a2v_gate_adaln_single"])
-                    ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_video_timestep, av_ca_factor))
+                    ->forward(ctx, ggml_ext_scale(ctx, av_ca_video_timestep, av_ca_factor))
                     .first;
             auto av_ca_audio_scale_shift_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_audio_scale_shift_adaln_single"])->forward(ctx, av_ca_audio_timestep).first;
             auto av_ca_v2a_gate_noise_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_v2a_gate_adaln_single"])
-                    ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_audio_timestep, av_ca_factor))
+                    ->forward(ctx, ggml_ext_scale(ctx, av_ca_audio_timestep, av_ca_factor))
                     .first;
 
             sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx");
@@ -1797,9 +1797,8 @@ namespace LTXV {
             return {vx, ax};
         }
 
-        ggml_tensor* merge_av_latents(ggml_context* ctx,
-                                      ggml_tensor* vx,
-                                      ggml_tensor* ax) const {
+        ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* vx, ggml_tensor* ax) const {
+            auto ctx = runner_ctx->ggml_ctx;
             if (ax == nullptr || ggml_nelements(ax) == 0 || ax->ne[1] == 0) {
                 return vx;
             }
@@ -1815,7 +1814,7 @@ namespace LTXV {
             ax = ggml_cont(ctx, ax);
             ax = ggml_reshape_4d(ctx, ax, audio_values, 1, 1, 1);
             if (pad_values > 0) {
-                ax = ggml_ext_pad(ctx, ax, static_cast<int>(pad_values), 0, 0, 0);
+                ax = ggml_ext_pad(runner_ctx, ax, static_cast<int>(pad_values), 0, 0, 0);
             }
             int64_t extra_channels = padded_len / divisor;
             ax                     = ggml_reshape_4d(ctx, ax, width, height, frames, extra_channels);
@@ -2008,7 +2007,7 @@ namespace LTXV {
                                             audio_cross_pe,
                                             video_connector_pe,
                                             audio_connector_pe);
-            auto out        = merge_av_latents(compute_ctx, out_pair.first, out_pair.second);
+            auto out        = merge_av_latents(&runner_ctx, out_pair.first, out_pair.second);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/diffusion/ming_image.hpp

@@ -50,8 +50,8 @@ namespace MingImage {
             txt                 = ggml_concat(gctx, txt, direct, 1);
             auto t              = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, timestep);
             const int64_t n_txt = txt->ne[1], n_img = img->ne[1];
-            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, n_txt);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, n_txt, n_txt + n_img);
             for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
                 txt = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe);
                 sd::ggml_graph_cut::mark_graph_cut(txt, "ming_image.context_refiner." + std::to_string(i), "txt");
@@ -65,9 +65,9 @@ namespace MingImage {
                 combined = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t);
                 sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined");
             }
-            img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img);
+            img = ggml_ext_slice(ctx, combined, 1, n_txt, n_txt + n_img);
             img = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"])->forward(ctx, img, t);
-            img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false);
+            img = DiT::unpatchify_and_crop(ctx, img, height, width, config.patch_size, config.patch_size, false);
             return ggml_scale(gctx, img, -1.f);
         }
     };

src/model/diffusion/minimax_h3.hpp

@@ -2,8 +2,10 @@
 #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__
 
 #include <algorithm>
+#include <array>
 #include <cinttypes>
 #include <cmath>
+#include <cstddef>
 #include <functional>
 #include <set>
 #include <string>
@@ -24,6 +26,10 @@ namespace MiniMaxH3 {
     constexpr float FRAME_RESCALE                 = 5.f / 3.f;
     constexpr float VISUAL_COND_TIMESTEP          = 0.999f;
 
+    inline constexpr std::array<float, 7> pdd_sigmas = {
+        1.f, 0.9675236940383911f, 0.8875485062599182f, 0.75f,
+        0.5327739119529724f, 0.23200887441635132f, 0.f};
+
     struct LowRankShape {
         int64_t q   = 0;
         int64_t k   = 0;
@@ -52,6 +58,7 @@ namespace MiniMaxH3 {
         int64_t time_embed_dim           = 2688;
         int64_t rope_inv_freq_len        = 16;
         int64_t adaln_curve_grid         = 0;
+        int64_t pdd_num_inference_steps  = 1;
         int patch_t                      = 1;
         int patch_h                      = 2;
         int patch_w                      = 2;
@@ -95,6 +102,21 @@ namespace MiniMaxH3 {
             }
             if (const auto* weight = find("audio_patch_proj.weight")) {
                 config.audio_latent_channels = weight->logical_ne0();
+            } else {
+                config.audio_latent_channels = 0;
+            }
+            if (const auto* weight = find("final_layer.video_out.weight")) {
+                const int64_t patch_dim = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w;
+                if (patch_dim <= 0 || weight->ne[1] % patch_dim != 0) {
+                    throw std::runtime_error("invalid MiniMax-H3 output head shape");
+                }
+                config.pdd_num_inference_steps = weight->ne[1] / patch_dim;
+                if (config.pdd_num_inference_steps != 1 && (config.audio_latent_channels != 0 || config.pdd_num_inference_steps != 6)) {
+                    throw std::runtime_error("unsupported MiniMax-H3 output head bank");
+                }
+            }
+            if (config.audio_latent_channels == 0 && config.pdd_num_inference_steps != 6) {
+                throw std::runtime_error("MiniMax-H3-Image requires a checkpoint with its six-step PDD heads included");
             }
             config.num_layers               = count_blocks(tensors, prefix + ".blocks.");
             config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks.");
@@ -151,14 +173,16 @@ namespace MiniMaxH3 {
             }
 
             LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64
-                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d",
+                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d, audio_latent_channels=%" PRId64 ", pdd_num_inference_steps=%" PRId64,
                         config.num_layers,
                         config.hidden_size,
                         config.num_attention_heads,
                         config.attention_head_dim,
                         config.ffn_hidden_size,
                         config.adaln_curve_grid,
-                        config.two_time);
+                        config.two_time,
+                        config.audio_latent_channels,
+                        config.pdd_num_inference_steps);
             return config;
         }
     };
@@ -263,13 +287,13 @@ namespace MiniMaxH3 {
 
             std::vector<ggml_tensor*> chunks;
             for (int64_t start = 0; start < x->ne[1]; start += chunk_size) {
-                auto input = ggml_ext_slice(ctx->ggml_ctx, x, 1, start, std::min(start + chunk_size, x->ne[1]), false);
+                auto input = ggml_ext_slice(ctx, x, 1, start, std::min(start + chunk_size, x->ne[1]), false);
                 if (!ggml_is_contiguous(input)) {
                     input = ggml_cont(ctx->ggml_ctx, input);
                 }
                 chunks.push_back(forward_chunk(input));
             }
-            return ggml_ext_vec_concat(ctx->ggml_ctx, chunks, 1);
+            return ggml_ext_vec_concat(ctx, chunks, 1);
         }
     };
 
@@ -278,29 +302,8 @@ namespace MiniMaxH3 {
         return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
     }
 
-    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
-                                           ggml_tensor* x,
-                                           ggml_tensor* pe) {
-        auto ctx        = runner_ctx->ggml_ctx;
-        int64_t rot_dim = pe->ne[2] * 2;
-        GGML_ASSERT(rot_dim <= x->ne[0]);
-#ifndef SD_USE_UPSTREAM_GGML
-        if (runner_ctx->backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32) {
-            auto out = sd_ggml_rope_apply(ctx, runner_ctx->backend, x, pe, false);
-            if (out && ggml_backend_supports_op(runner_ctx->backend, out)) {
-                return out;
-            }
-        }
-#endif
-        auto rotated = Rope::apply_rope(runner_ctx,
-                                        ggml_ext_slice(ctx, x, 0, 0, rot_dim),
-                                        pe,
-                                        false);
-        if (rot_dim == x->ne[0]) {
-            return rotated;
-        }
-        auto tail = attention_layout(ctx, ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0]));
-        return ggml_concat(ctx, rotated, tail, 0);
+    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) {
+        return ggml_ext_rope_apply(runner_ctx, x, pe, false);
     }
 
     struct Attention : public GGMLBlock {
@@ -381,7 +384,7 @@ namespace MiniMaxH3 {
             for (size_t i = 0; i < inputs.size(); ++i) {
                 auto input = inputs[i];
                 for (int64_t start = 0; start < input->ne[1]; start += chunk_size) {
-                    auto part = ggml_ext_slice(ctx->ggml_ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false);
+                    auto part = ggml_ext_slice(ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false);
                     if (!ggml_is_contiguous(part)) {
                         part = ggml_cont(ctx->ggml_ctx, part);
                     }
@@ -397,7 +400,7 @@ namespace MiniMaxH3 {
             }
             std::vector<ggml_tensor*> result;
             for (auto& parts : chunks) {
-                result.push_back(ggml_ext_vec_concat(ctx->ggml_ctx, parts, 2));
+                result.push_back(ggml_ext_vec_concat(ctx, parts, 2));
             }
             return result;
         }
@@ -420,13 +423,13 @@ namespace MiniMaxH3 {
                 for (int64_t start = 0; start < heads; start += head_chunk) {
                     std::vector<ggml_tensor*> part;
                     for (auto tensor : qkv) {
-                        part.push_back(ggml_ext_slice(ctx->ggml_ctx, tensor, 1, start, std::min(start + head_chunk, heads), false));
+                        part.push_back(ggml_ext_slice(ctx, tensor, 1, start, std::min(start + head_chunk, heads), false));
                     }
                     outputs.push_back(attend_heads(ctx, part, pe));
                     // Keep each head group's temporaries out of the concatenation tree's live set.
                     ctx->expand_graph(outputs.back());
                 }
-                return ggml_ext_vec_concat(ctx->ggml_ctx, outputs, 0);
+                return ggml_ext_vec_concat(ctx, outputs, 0);
             }
             return attend_heads(ctx, qkv, pe);
         }
@@ -562,19 +565,20 @@ namespace MiniMaxH3 {
         int32_t source_index = -1;
     };
 
-    static std::vector<ggml_tensor*> modulation_row(ggml_context* ctx,
+    static std::vector<ggml_tensor*> modulation_row(GGMLRunnerContext* runner_ctx,
                                                     ggml_tensor* projection,
                                                     int64_t hidden_size,
                                                     int expand,
                                                     int modalities,
                                                     int64_t row) {
+        auto ctx              = runner_ctx->ggml_ctx;
         int64_t timestep_rows = projection->ne[1];
         auto reshaped         = ggml_reshape_2d(ctx,
                                                 projection,
                                                 hidden_size * expand,
                                                 timestep_rows * modalities);
-        auto selected         = ggml_ext_slice(ctx, reshaped, 1, row, row + 1);
-        return ggml_ext_chunk(ctx, selected, expand, 0);
+        auto selected         = ggml_ext_slice(runner_ctx, reshaped, 1, row, row + 1);
+        return ggml_ext_chunk(runner_ctx, selected, expand, 0);
     }
 
     struct TransformerBlock : public GGMLBlock {
@@ -616,9 +620,9 @@ namespace MiniMaxH3 {
             std::vector<std::vector<ggml_tensor*>> rows;
             std::vector<ggml_tensor*> inputs;
             for (const auto& segment : segments) {
-                rows.push_back(modulation_row(ctx->ggml_ctx, mods, config.hidden_size,
+                rows.push_back(modulation_row(ctx, mods, config.hidden_size,
                                               6, 3, segment.modulation_row));
-                inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false));
+                inputs.push_back(ggml_ext_slice(ctx, x, 1, segment.start, segment.end, false));
             }
             auto modulate = [&](ggml_tensor* value, const std::vector<ggml_tensor*>& row, int shift, int scale) {
                 return ggml_ext_modulate(ctx, value, row[shift], row[scale]);
@@ -629,7 +633,7 @@ namespace MiniMaxH3 {
             auto attended = attn->attend(ctx, qkv, pe);
             std::vector<ggml_tensor*> projections;
             for (const auto& segment : segments) {
-                auto part = ggml_ext_slice(ctx->ggml_ctx, attended, 1, segment.start, segment.end, false);
+                auto part = ggml_ext_slice(ctx, attended, 1, segment.start, segment.end, false);
                 projections.push_back(attn->project_output(ctx, part));
             }
             ggml_tensor* output = nullptr;
@@ -644,6 +648,27 @@ namespace MiniMaxH3 {
         }
     };
 
+    static std::array<float, 6> pdd_head_weights(float sigma, float sigma_next) {
+        auto nearest = [](float value) {
+            size_t index = 0;
+            for (size_t i = 1; i < pdd_sigmas.size(); ++i) {
+                if (std::abs(pdd_sigmas[i] - value) < std::abs(pdd_sigmas[index] - value)) {
+                    index = i;
+                }
+            }
+            return index;
+        };
+        const size_t start = std::min(nearest(sigma), pdd_sigmas.size() - 2);
+        const size_t stop  = std::max(nearest(sigma_next), start + 1);
+        std::array<float, 6> weights{};
+        // Comfy checkpoints keep head zero as the base and store offsets in the remaining heads.
+        weights[0] = 1.f;
+        for (size_t i = std::max<size_t>(start, 1); i < stop; ++i) {
+            weights[i] = (pdd_sigmas[i] - pdd_sigmas[i + 1]) / (pdd_sigmas[start] - pdd_sigmas[stop]);
+        }
+        return weights;
+    }
+
     struct FinalLayer : public GGMLBlock {
         Config config;
 
@@ -657,29 +682,50 @@ namespace MiniMaxH3 {
                                                                             1,
                                                                             !config.uses_adaln_curves(),
                                                                             config.uses_adaln_curves());
-            blocks["video_out"]  = std::make_shared<Linear>(config.hidden_size, video_dim, true, true);
-            blocks["audio_out"]  = std::make_shared<Linear>(config.hidden_size, config.audio_latent_channels, true, true);
+            blocks["video_out"] =
+                std::make_shared<Linear>(config.hidden_size, video_dim * config.pdd_num_inference_steps, true, true);
+            if (config.audio_latent_channels > 0) {
+                blocks["audio_out"] = std::make_shared<Linear>(config.hidden_size, config.audio_latent_channels, true, true);
+            }
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
                                                       ggml_tensor* x,
                                                       ggml_tensor* t_emb,
                                                       const TokenModulationSpan& video,
-                                                      const TokenModulationSpan& audio) {
+                                                      const TokenModulationSpan& audio,
+                                                      float sigma,
+                                                      float sigma_next) {
             auto norm      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
             auto adaln     = std::dynamic_pointer_cast<AdaLayerNormModulation>(blocks["adaln_proj"]);
             auto video_out = std::dynamic_pointer_cast<Linear>(blocks["video_out"]);
-            auto audio_out = std::dynamic_pointer_cast<Linear>(blocks["audio_out"]);
             auto mods      = adaln->forward(ctx, t_emb);
             auto apply     = [&](const TokenModulationSpan& segment) {
-                auto row   = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row);
-                auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end));
+                auto row   = modulation_row(ctx, mods, config.hidden_size, 2, 1, segment.modulation_row);
+                auto value = norm->forward(ctx, ggml_ext_slice(ctx, x, 1, segment.start, segment.end));
                 return ggml_add(ctx->ggml_ctx,
                                 ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])),
                                 row[0]);
             };
-            return {video_out->forward(ctx, apply(video)),
-                    audio_out->forward(ctx, apply(audio))};
+            auto prediction = video_out->forward(ctx, apply(video));
+            if (config.pdd_num_inference_steps > 1) {
+                GGML_ASSERT(std::isfinite(sigma_next));
+                auto weights      = pdd_head_weights(sigma, sigma_next);
+                const int64_t dim = prediction->ne[0] / config.pdd_num_inference_steps;
+                auto blended      = ggml_ext_scale(ctx, ggml_ext_slice(ctx, prediction, 0, 0, dim), weights[0]);
+                for (int64_t i = 1; i < config.pdd_num_inference_steps; ++i) {
+                    if (weights[i] != 0.f) {
+                        auto offset = ggml_ext_slice(ctx, prediction, 0, i * dim, (i + 1) * dim);
+                        blended     = ggml_add(ctx->ggml_ctx, blended, ggml_ext_scale(ctx, offset, weights[i]));
+                    }
+                }
+                prediction = blended;
+            }
+            ggml_tensor* audio_prediction = nullptr;
+            if (config.audio_latent_channels > 0) {
+                audio_prediction = std::dynamic_pointer_cast<Linear>(blocks["audio_out"])->forward(ctx, apply(audio));
+            }
+            return {prediction, audio_prediction};
         }
     };
 
@@ -690,8 +736,10 @@ namespace MiniMaxH3 {
             : config(config) {
             int64_t video_dim          = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w;
             blocks["video_patch_proj"] = std::make_shared<Linear>(video_dim, config.hidden_size, true, true);
-            blocks["audio_patch_proj"] = std::make_shared<Linear>(config.audio_latent_channels, config.hidden_size, true, true);
-            blocks["condition_proj"]   = std::make_shared<Linear>(config.text_dim, config.hidden_size, true);
+            if (config.audio_latent_channels > 0) {
+                blocks["audio_patch_proj"] = std::make_shared<Linear>(config.audio_latent_channels, config.hidden_size, true, true);
+            }
+            blocks["condition_proj"] = std::make_shared<Linear>(config.text_dim, config.hidden_size, true);
             if (!config.uses_adaln_curves()) {
                 blocks["time_embedder"] = std::make_shared<TimeEmbedder>(config.timestep_input_dim,
                                                                          config.time_embed_hidden_size,
@@ -774,7 +822,7 @@ namespace MiniMaxH3 {
                                                   1);
             ggml_tensor* angles = nullptr;
             for (int axis = 0; axis < 3; ++axis) {
-                auto pos          = ggml_ext_slice(ctx->ggml_ctx, position_ids, 0, axis, axis + 1);
+                auto pos          = ggml_ext_slice(ctx, position_ids, 0, axis, axis + 1);
                 auto expanded_inv = ggml_repeat_4d(ctx->ggml_ctx,
                                                    inv,
                                                    inv->ne[0],
@@ -808,15 +856,16 @@ namespace MiniMaxH3 {
                                                       const std::vector<TokenModulationSpan>& segments,
                                                       const std::vector<SequenceSegment>& sequence_segments,
                                                       const TokenModulationSpan& video_segment,
-                                                      const TokenModulationSpan& audio_segment) {
+                                                      const TokenModulationSpan& audio_segment,
+                                                      float sigma,
+                                                      float sigma_next) {
             auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
-            auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);
 
             std::vector<std::pair<int64_t, int64_t>> condition_video_ranges;
             ggml_tensor* video_rows = nullptr;
             int64_t video_offset    = 0;
             for (auto condition : condition_videos) {
-                auto rows = DiT::patchify_3d(ctx->ggml_ctx,
+                auto rows = DiT::patchify_3d(ctx,
                                              condition,
                                              config.patch_t,
                                              config.patch_h,
@@ -827,7 +876,7 @@ namespace MiniMaxH3 {
                 video_offset += rows->ne[1];
                 video_rows = video_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, video_rows, rows, 1);
             }
-            auto target_video_rows                         = DiT::patchify_3d(ctx->ggml_ctx,
+            auto target_video_rows                         = DiT::patchify_3d(ctx,
                                                                               video,
                                                                               config.patch_t,
                                                                               config.patch_h,
@@ -840,7 +889,7 @@ namespace MiniMaxH3 {
             auto video_embeds                              = video_proj->forward(ctx, video_rows);
 
             auto pack_audio_rows = [&](ggml_tensor* value) {
-                value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, value, 2, 0, 1, 3));
+                value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, value, 2, 0, 1, 3));
                 return ggml_reshape_3d(ctx->ggml_ctx,
                                        value,
                                        value->ne[0],
@@ -856,11 +905,15 @@ namespace MiniMaxH3 {
                 audio_offset += rows->ne[1];
                 audio_rows = audio_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, audio_rows, rows, 1);
             }
-            audio                                          = pack_audio_rows(audio);
-            std::pair<int64_t, int64_t> target_audio_range = {audio_offset, audio_offset + audio->ne[1]};
-            audio_rows                                     = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1);
-            auto audio_embeds                              = audio_proj->forward(ctx, audio_rows);
-            context                                        = refine_context(ctx, context);
+            std::pair<int64_t, int64_t> target_audio_range{};
+            ggml_tensor* audio_embeds = nullptr;
+            if (audio != nullptr) {
+                audio              = pack_audio_rows(audio);
+                target_audio_range = {audio_offset, audio_offset + audio->ne[1]};
+                audio_rows         = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1);
+                audio_embeds       = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"])->forward(ctx, audio_rows);
+            }
+            context = refine_context(ctx, context);
 
             ggml_tensor* h = nullptr;
             auto append    = [&](ggml_tensor* value) {
@@ -873,20 +926,20 @@ namespace MiniMaxH3 {
                     GGML_ASSERT(sequence.source_index >= 0 &&
                                 sequence.source_index < static_cast<int32_t>(condition_video_ranges.size()));
                     auto range = condition_video_ranges[static_cast<size_t>(sequence.source_index)];
-                    append(ggml_ext_slice(ctx->ggml_ctx, video_embeds, 1, range.first, range.second));
+                    append(ggml_ext_slice(ctx, video_embeds, 1, range.first, range.second));
                 } else if (sequence.kind == SequenceKind::CONDITION_AUDIO) {
                     GGML_ASSERT(sequence.source_index >= 0 &&
                                 sequence.source_index < static_cast<int32_t>(condition_audio_ranges.size()));
                     auto range = condition_audio_ranges[static_cast<size_t>(sequence.source_index)];
-                    append(ggml_ext_slice(ctx->ggml_ctx, audio_embeds, 1, range.first, range.second));
+                    append(ggml_ext_slice(ctx, audio_embeds, 1, range.first, range.second));
                 } else if (sequence.kind == SequenceKind::TARGET_AUDIO) {
-                    append(ggml_ext_slice(ctx->ggml_ctx,
+                    append(ggml_ext_slice(ctx,
                                           audio_embeds,
                                           1,
                                           target_audio_range.first,
                                           target_audio_range.second));
                 } else {
-                    append(ggml_ext_slice(ctx->ggml_ctx,
+                    append(ggml_ext_slice(ctx,
                                           video_embeds,
                                           1,
                                           target_video_range.first,
@@ -911,8 +964,8 @@ namespace MiniMaxH3 {
             }
 
             auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
-            auto output      = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment);
-            auto video_out   = DiT::unpatchify_3d(ctx->ggml_ctx,
+            auto output      = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment, sigma, sigma_next);
+            auto video_out   = DiT::unpatchify_3d(ctx,
                                                   output.first,
                                                   video->ne[2] / config.patch_t,
                                                   video->ne[1] / config.patch_h,
@@ -921,15 +974,18 @@ namespace MiniMaxH3 {
                                                   config.patch_h,
                                                   config.patch_w,
                                                   true);
-            auto audio_out   = ggml_reshape_4d(ctx->ggml_ctx,
-                                               output.second,
-                                               config.audio_latent_channels,
-                                               audio->ne[1] / 2,
-                                               2,
-                                               audio->ne[2]);
-            audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
-            video_out        = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
-            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
+            video_out        = ggml_ext_scale(ctx, video_out, -1.f);
+            if (audio == nullptr) {
+                return {video_out, nullptr};
+            }
+            auto audio_out = ggml_reshape_4d(ctx->ggml_ctx,
+                                             output.second,
+                                             config.audio_latent_channels,
+                                             audio->ne[1] / 2,
+                                             2,
+                                             audio->ne[2]);
+            audio_out      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3));
+            audio_out      = ggml_ext_scale(ctx, audio_out, -1.f);
             return {video_out, audio_out};
         }
     };
@@ -1015,7 +1071,7 @@ namespace MiniMaxH3 {
         int64_t run_start = 0;
         int current_tag   = text_tags.empty() ? 1 : text_tags[0];
         for (int64_t i = 1; i <= text_len; ++i) {
-            int tag = i < text_len && !text_tags.empty() ? text_tags[i] : -1;
+            int tag = i < text_len ? (text_tags.empty() ? 1 : text_tags[i]) : -1;
             if (i == text_len || tag != current_tag) {
                 layout.segments.push_back({run_start,
                                            i,
@@ -1141,16 +1197,18 @@ namespace MiniMaxH3 {
             }
         }
 
-        int64_t audio_start = row;
-        append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back());
-        layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row};
-        layout.sequence_segments.push_back({audio_start,
-                                            row + audio_t * 2,
-                                            SequenceKind::TARGET_AUDIO});
-        layout.segments.push_back({audio_start,
-                                   row + audio_t * 2,
-                                   audio_time_row * 3 + 2});
-        row += audio_t * 2;
+        if (audio_t > 0) {
+            int64_t audio_start = row;
+            append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back());
+            layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row};
+            layout.sequence_segments.push_back({audio_start,
+                                                row + audio_t * 2,
+                                                SequenceKind::TARGET_AUDIO});
+            layout.segments.push_back({audio_start,
+                                       row + audio_t * 2,
+                                       audio_time_row * 3 + 2});
+            row += audio_t * 2;
+        }
 
         int64_t video_start = row;
         for (int64_t t = 0; t < latent_t; ++t) {
@@ -1221,15 +1279,14 @@ namespace MiniMaxH3 {
             model.get_param_tensors(tensors, prefix);
         }
 
-        ggml_tensor* merge_av_latents(ggml_context* ctx,
-                                      ggml_tensor* video,
-                                      ggml_tensor* audio) const {
+        ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* video, ggml_tensor* audio) const {
+            auto ctx        = runner_ctx->ggml_ctx;
             int64_t divisor = video->ne[0] * video->ne[1] * video->ne[2];
             int64_t values  = ggml_nelements(audio);
             int64_t padding = (divisor - values % divisor) % divisor;
             audio           = ggml_reshape_4d(ctx, ggml_cont(ctx, audio), values, 1, 1, 1);
             if (padding > 0) {
-                audio = ggml_ext_pad(ctx, audio, static_cast<int>(padding), 0, 0, 0);
+                audio = ggml_ext_pad(runner_ctx, audio, static_cast<int>(padding), 0, 0, 0);
             }
             audio = ggml_reshape_4d(ctx,
                                     audio,
@@ -1253,26 +1310,29 @@ namespace MiniMaxH3 {
                                  float audio_shift,
                                  float video_sigma_next) {
             GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5);
-            GGML_ASSERT(audio_length > 0);
+            GGML_ASSERT(audio_length > 0 || config.audio_latent_channels == 0);
             GGML_ASSERT(!context_tensor.empty());
             const int64_t spatial      = packed.shape()[0] * packed.shape()[1] * packed.shape()[2];
             const int64_t video_values = spatial * config.video_latent_channels;
             const int64_t audio_values = int64_t(audio_length) * 2 * config.audio_latent_channels;
             GGML_ASSERT(packed.numel() >= video_values + audio_values);
-            auto packed_input  = make_input(packed);
-            auto video         = ggml_view_4d(compute_ctx, packed_input,
-                                              packed.shape()[0], packed.shape()[1], packed.shape()[2],
-                                              config.video_latent_channels,
-                                              packed.shape()[0] * sizeof(float),
-                                              packed.shape()[0] * packed.shape()[1] * sizeof(float),
-                                              spatial * sizeof(float), 0);
-            auto audio_carrier = ggml_view_4d(compute_ctx, packed_input,
-                                              audio_length, 2, config.audio_latent_channels, 1,
-                                              audio_length * sizeof(float),
-                                              int64_t(audio_length) * 2 * sizeof(float),
-                                              audio_values * sizeof(float),
-                                              video_values * sizeof(float));
-            auto context       = make_input(context_tensor);
+            auto packed_input          = make_input(packed);
+            auto video                 = ggml_view_4d(compute_ctx, packed_input,
+                                                      packed.shape()[0], packed.shape()[1], packed.shape()[2],
+                                                      config.video_latent_channels,
+                                                      packed.shape()[0] * sizeof(float),
+                                                      packed.shape()[0] * packed.shape()[1] * sizeof(float),
+                                                      spatial * sizeof(float), 0);
+            ggml_tensor* audio_carrier = nullptr;
+            if (config.audio_latent_channels > 0) {
+                audio_carrier = ggml_view_4d(compute_ctx, packed_input,
+                                             audio_length, 2, config.audio_latent_channels, 1,
+                                             audio_length * sizeof(float),
+                                             int64_t(audio_length) * 2 * sizeof(float),
+                                             audio_values * sizeof(float),
+                                             video_values * sizeof(float));
+            }
+            auto context = make_input(context_tensor);
             std::vector<ggml_tensor*> condition_inputs;
             condition_inputs.reserve(condition_videos.size());
             for (const auto& condition : condition_videos) {
@@ -1302,21 +1362,22 @@ namespace MiniMaxH3 {
             }
             // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
             // latent follows one sigma schedule. Restore x_a for the H3 network.
-            auto audio  = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
-            auto layout = build_layout(context_tensor.shape()[1],
-                                       packed.shape()[2],
-                                       packed.shape()[1],
-                                       packed.shape()[0],
-                                       audio_length,
-                                       condition_videos,
-                                       condition_audios,
-                                       keyframe_indices,
-                                       reference_blocks,
-                                       text_tags,
-                                       t_v,
-                                       t_a,
-                                       r_v,
-                                       r_a);
+            auto audio_ctx = get_context();
+            auto audio     = audio_carrier == nullptr ? nullptr : ggml_ext_scale(&audio_ctx, audio_carrier, sigma_a / sigma_v);
+            auto layout    = build_layout(context_tensor.shape()[1],
+                                          packed.shape()[2],
+                                          packed.shape()[1],
+                                          packed.shape()[0],
+                                          audio_length,
+                                          condition_videos,
+                                          condition_audios,
+                                          keyframe_indices,
+                                          reference_blocks,
+                                          text_tags,
+                                          t_v,
+                                          t_a,
+                                          r_v,
+                                          r_a);
 
             position_input_cache = sd::Tensor<float>(
                 {3, static_cast<int64_t>(layout.positions.size() / 3)},
@@ -1390,14 +1451,18 @@ namespace MiniMaxH3 {
                                             layout.segments,
                                             layout.sequence_segments,
                                             layout.video_segment,
-                                            layout.audio_segment);
+                                            layout.audio_segment,
+                                            sigma_v,
+                                            video_sigma_next);
             // Convert the model's audio velocity to d(c_a) / d(sigma_v).
-            output.second = ggml_add(compute_ctx,
-                                     ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
-                                     ggml_ext_scale(compute_ctx,
-                                                    output.second,
-                                                    1.f + (audio_scale - 1.f) * sigma_a));
-            auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
+            if (output.second != nullptr) {
+                output.second = ggml_add(compute_ctx,
+                                         ggml_ext_scale(&runner_ctx, audio, 1.f - audio_scale),
+                                         ggml_ext_scale(&runner_ctx,
+                                                        output.second,
+                                                        1.f + (audio_scale - 1.f) * sigma_a));
+            }
+            auto merged = output.second == nullptr ? output.first : merge_av_latents(&runner_ctx, output.first, output.second);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }

src/model/diffusion/minit2i.hpp

@@ -257,14 +257,18 @@ namespace MiniT2I {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1.0f);
             t_emb      = mlp_0->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             return mlp_2->forward(ctx, t_emb);
         }
     };
 
-    inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
+    inline std::vector<ggml_tensor*> split_qkv(GGMLRunnerContext* runner_ctx,
+                                               ggml_tensor* qkv,
+                                               int64_t num_heads,
+                                               int64_t head_dim) {
+        auto ctx  = runner_ctx->ggml_ctx;
         int64_t N = qkv->ne[2];
         int64_t L = qkv->ne[1];
         auto q    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
@@ -301,7 +305,7 @@ namespace MiniT2I {
             auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
             auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
 
-            auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto qkv = split_qkv(ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
             auto q   = q_norm->forward(ctx, qkv[0]);
             auto k   = k_norm->forward(ctx, qkv[1]);
             auto v   = qkv[2];
@@ -353,16 +357,16 @@ namespace MiniT2I {
             int64_t li = img->ne[1];
             int64_t lt = txt->ne[1];
 
-            auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
-            auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto img_qkv = split_qkv(ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
+            auto txt_qkv = split_qkv(ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
 
             auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
             auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
             auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
 
             auto out     = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
-            auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
-            auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
+            auto out_txt = ggml_ext_slice(ctx, out, 1, 0, lt);
+            auto out_img = ggml_ext_slice(ctx, out, 1, lt, lt + li);
 
             img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
             txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
@@ -416,7 +420,7 @@ namespace MiniT2I {
             mask            = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
             mask            = ggml_repeat(ctx->ggml_ctx, mask, context);
             auto keep       = ggml_mul(ctx->ggml_ctx, context, mask);
-            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
+            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, mask), mask);
             auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
             return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
         }
@@ -468,8 +472,8 @@ namespace MiniT2I {
             }
             auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
             auto out      = final_layer->forward(ctx, combined);
-            auto img_out  = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
-            return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
+            auto img_out  = ggml_ext_slice(ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
+            return DiT::unpatchify(ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
         }
     };
 

src/model/diffusion/mmdit.hpp

@@ -272,7 +272,7 @@ public:
         auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
         auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
 
-        auto t_freq = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size);  // [N, frequency_embedding_size]
+        auto t_freq = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size);  // [N, frequency_embedding_size]
 
         auto t_emb = mlp_0->forward(ctx, t_freq);
         t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
@@ -334,7 +334,7 @@ public:
         auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
         auto qkv         = qkv_proj->forward(ctx, x);
-        auto qkv_vec     = split_qkv(ctx->ggml_ctx, qkv);
+        auto qkv_vec     = ggml_ext_chunk(ctx, qkv, 3, 0);
         int64_t head_dim = qkv_vec[0]->ne[0] / num_heads;
         auto q           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);  // [N, n_token, n_head, d_head]
         auto k           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]);  // [N, n_token, n_head, d_head]
@@ -439,7 +439,7 @@ public:
 
         int n_mods = 9;
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0);
 
         auto shift_msa  = m_vec[0];  // [N, hidden_size]
         auto scale_msa  = m_vec[1];  // [N, hidden_size]
@@ -476,7 +476,7 @@ public:
             n_mods = 2;
         }
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0);
 
         auto shift_msa = m_vec[0];  // [N, hidden_size]
         auto scale_msa = m_vec[1];  // [N, hidden_size]
@@ -750,7 +750,7 @@ public:
         auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 2 * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, 2, 0);
         auto shift = m_vec[0];  // [N, hidden_size]
         auto scale = m_vec[1];  // [N, hidden_size]
 
@@ -926,7 +926,7 @@ public:
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
-        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false);  // [N, C, H, W]
+        x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false);  // [N, C, H, W]
 
         return x;
     }

src/model/diffusion/moss_image.hpp

@@ -164,9 +164,9 @@ namespace MOSSImage {
             auto c            = ctx->ggml_ctx;
             auto linear       = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
             int64_t image_len = image->ne[0] * image->ne[1];
-            auto x            = linear("video_embedder", DiT::patchify(c, image, 1, 1));
+            auto x            = linear("video_embedder", DiT::patchify(ctx, image, 1, 1));
             text              = linear("text_embedder.proj", text);
-            auto text_pe      = ggml_ext_slice(c, pe, 3, image_len, image_len + text->ne[1]);
+            auto text_pe      = ggml_ext_slice(ctx, pe, 3, image_len, image_len + text->ne[1]);
             for (int i = 0; i < 2; ++i) {
                 text = std::dynamic_pointer_cast<TextBlock>(blocks["text_embedder.blocks." + std::to_string(i)])->forward(ctx, text, text_pe);
             }
@@ -175,7 +175,7 @@ namespace MOSSImage {
             auto head = linear("video_head_modulation", time);
             std::array<ggml_tensor*, 6> parts;
             for (int i = 0; i < 6; ++i) {
-                parts[i] = ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim);
+                parts[i] = ggml_ext_slice(ctx, mod, 0, i * config.dim, (i + 1) * config.dim);
                 if (ctx->graph != nullptr) {
                     ggml_build_forward_expand(ctx->graph, parts[i]);
                 }
@@ -187,13 +187,13 @@ namespace MOSSImage {
                     sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x");
                 }
             }
-            x          = ggml_ext_slice(c, x, 1, 0, image_len);
+            x          = ggml_ext_slice(ctx, x, 1, 0, image_len);
             x          = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
-            auto shift = ggml_ext_slice(c, head, 0, 0, config.dim);
-            auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim);
+            auto shift = ggml_ext_slice(ctx, head, 0, 0, config.dim);
+            auto scale = ggml_ext_slice(ctx, head, 0, config.dim, 2 * config.dim);
             x          = ggml_ext_modulate(ctx, x, shift, scale);
             x          = linear("final_linear_video", x);
-            return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1);
+            return DiT::unpatchify(ctx, x, image->ne[1], image->ne[0], 1, 1);
         }
     };
 

src/model/diffusion/pid.hpp

@@ -185,14 +185,15 @@ namespace Pid {
         return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
     }
 
-    inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) {
-        auto left  = ggml_ext_slice(ctx, x, 0, 0, 1);
-        auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+    inline ggml_tensor* replicate_pad_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+        auto ctx   = runner_ctx->ggml_ctx;
+        auto left  = ggml_ext_slice(runner_ctx, x, 0, 0, 1);
+        auto right = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
         x          = ggml_concat(ctx, left, x, 0);
         x          = ggml_concat(ctx, x, right, 0);
 
-        auto top    = ggml_ext_slice(ctx, x, 1, 0, 1);
-        auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+        auto top    = ggml_ext_slice(runner_ctx, x, 1, 0, 1);
+        auto bottom = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]);
         x           = ggml_concat(ctx, top, x, 1);
         return ggml_concat(ctx, x, bottom, 1);
     }
@@ -235,7 +236,7 @@ namespace Pid {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10);
             t_emb      = mlp_0->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             return mlp_2->forward(ctx, t_emb);
@@ -292,7 +293,7 @@ namespace Pid {
             auto proj     = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
             auto qkv         = qkv_proj->forward(ctx, x);
-            auto qkv_vec     = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec     = ggml_ext_chunk(ctx, qkv, 3, 0);
             int64_t L        = x->ne[1];
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
@@ -342,14 +343,14 @@ namespace Pid {
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
 
-            auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x));
+            auto qkv_x = ggml_ext_chunk(ctx, qkv_x_proj->forward(ctx, x), 3, 0);
             auto qx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
             auto kx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
             auto vx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
             qx         = q_norm_x->forward(ctx, qx);
             kx         = k_norm_x->forward(ctx, kx);
 
-            auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y));
+            auto qkv_y = ggml_ext_chunk(ctx, qkv_y_proj->forward(ctx, y), 3, 0);
             auto qy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
             auto ky    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
             auto vy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
@@ -362,8 +363,8 @@ namespace Pid {
             auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3);
             auto out       = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true);
 
-            auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny);
-            auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx);
+            auto out_y = ggml_ext_slice(ctx, out, 1, 0, Ny);
+            auto out_x = ggml_ext_slice(ctx, out, 1, Ny, Ny + Nx);
             return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)};
         }
     };
@@ -400,8 +401,8 @@ namespace Pid {
             auto ada_img = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_img.0"]);
             auto ada_txt = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_txt.0"]);
 
-            auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0);
-            auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0);
+            auto mx = ggml_ext_chunk(ctx, ada_img->forward(ctx, c), 6, 0);
+            auto my = ggml_ext_chunk(ctx, ada_txt->forward(ctx, c), 6, 0);
 
             auto x_norm   = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]);
             auto y_norm   = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]);
@@ -438,13 +439,13 @@ namespace Pid {
             int64_t L  = (W / patch_size) * (H / patch_size);
             int64_t P2 = patch_size * patch_size;
 
-            auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3));
+            auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, inputs, 2, 0, 1, 3));
             x      = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B);
             x      = proj->forward(ctx, x);
             x      = ggml_add(ctx->ggml_ctx, x, pos_full);
             x      = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B);
-            x      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-            x      = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
+            x      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));
+            x      = DiT::patchify(ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
             x      = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B);
             return x;
         }
@@ -500,7 +501,7 @@ namespace Pid {
 
             auto ada_params = ada->forward(ctx, s_cond);
             ada_params      = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL);
-            auto mod        = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0);
+            auto mod        = ggml_ext_chunk(ctx, ada_params, 6, 0);
 
             auto x_norm    = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]);
             auto x_flat    = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL);
@@ -564,12 +565,12 @@ namespace Pid {
             auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
             auto h     = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
             if (replicate_padding) {
-                h = replicate_pad_2d(ctx->ggml_ctx, h);
+                h = replicate_pad_2d(ctx, h);
             }
             h = conv1->forward(ctx, h);
             h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
             if (replicate_padding) {
-                h = replicate_pad_2d(ctx->ggml_ctx, h);
+                h = replicate_pad_2d(ctx, h);
             }
             h = conv2->forward(ctx, h);
             return ggml_add(ctx->ggml_ctx, x, h);
@@ -628,13 +629,13 @@ namespace Pid {
             if (unpatchify_factor > 1) {
                 int64_t latent_h = lq_latent->ne[1];
                 int64_t latent_w = lq_latent->ne[0];
-                lq_latent        = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
+                lq_latent        = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, lq_latent, 2, 0, 1, 3));
                 lq_latent        = ggml_reshape_3d(ctx->ggml_ctx,
                                                    lq_latent,
                                                    lq_latent->ne[0],
                                                    lq_latent->ne[1] * lq_latent->ne[2],
                                                    lq_latent->ne[3]);
-                lq_latent        = DiT::unpatchify(ctx->ggml_ctx,
+                lq_latent        = DiT::unpatchify(ctx,
                                                    lq_latent,
                                                    latent_h,
                                                    latent_w,
@@ -658,12 +659,12 @@ namespace Pid {
             }
 
             if (config.lq_replicate_padding) {
-                lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent);
+                lq_latent = replicate_pad_2d(ctx, lq_latent);
             }
             auto feat = conv0->forward(ctx, lq_latent);
             feat      = ggml_silu_inplace(ctx->ggml_ctx, feat);
             if (config.lq_replicate_padding) {
-                feat = replicate_pad_2d(ctx->ggml_ctx, feat);
+                feat = replicate_pad_2d(ctx, feat);
             }
             feat = conv2->forward(ctx, feat);
             for (int i = 0; i < config.lq_num_res_blocks; ++i) {
@@ -674,7 +675,7 @@ namespace Pid {
             int64_t B   = feat->ne[3];
             int64_t C   = feat->ne[2];
             int64_t L   = target_pH * target_pW;
-            auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3));
+            auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, feat, 2, 0, 1, 3));
             tokens      = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B);
 
             int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
@@ -754,15 +755,15 @@ namespace Pid {
             int64_t L      = Hs * Ws;
             int64_t P2     = config.patch_size * config.patch_size;
 
-            auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
+            auto x_patches = DiT::patchify(ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
             auto t_emb     = t_embedder->forward(ctx, timesteps);
             auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
 
             GGML_ASSERT(context != nullptr);
             int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
-            auto y       = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt);
+            auto y       = ggml_ext_slice(ctx, context, 1, 0, Ltxt);
             auto y_emb   = y_embedder->forward(ctx, y);
-            auto y_pos   = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt);
+            auto y_pos   = ggml_ext_slice(ctx, params["y_pos_embedding"], 1, 0, Ltxt);
             y_emb        = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
 
             std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
@@ -809,15 +810,15 @@ namespace Pid {
 
             pixels   = final_layer->forward(ctx, pixels);
             pixels   = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B);
-            auto out = DiT::unpatchify(ctx->ggml_ctx,
+            auto out = DiT::unpatchify(ctx,
                                        pixels,
                                        Hs,
                                        Ws,
                                        static_cast<int>(config.patch_size),
                                        static_cast<int>(config.patch_size),
                                        false);
-            out      = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig);
-            out      = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig);
+            out      = ggml_ext_slice(ctx, out, 1, 0, H_orig);
+            out      = ggml_ext_slice(ctx, out, 0, 0, W_orig);
             return out;
         }
     };

src/model/diffusion/pixart.hpp

@@ -140,11 +140,11 @@ namespace PixArt {
             std::vector<ggml_tensor*> qkv;
             if (self_attention) {
                 auto projected = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
-                qkv            = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0);
+                qkv            = ggml_ext_chunk(ctx, projected, 3, 0);
             } else {
                 auto q     = std::dynamic_pointer_cast<Linear>(blocks["q_linear"])->forward(ctx, x);
                 auto kv    = std::dynamic_pointer_cast<Linear>(blocks["kv_linear"])->forward(ctx, context);
-                auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0);
+                auto parts = ggml_ext_chunk(ctx, kv, 2, 0);
                 qkv        = {q, parts[0], parts[1]};
             }
             auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled);
@@ -186,7 +186,7 @@ namespace PixArt {
             }
             table   = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1);
             auto m  = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table);
-            auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0);
+            auto mv = ggml_ext_chunk(ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, m), dim * 6, N), 6, 0);
 
             auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn"]);
             auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["cross_attn"]);
@@ -204,7 +204,7 @@ namespace PixArt {
             x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
             h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
             h = proj->forward(ctx, h);
-            h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
+            h = ggml_ext_gelu(ctx, h, true);
             h = fc2->forward(ctx, h);
             return gate(h, mv[5]);
         }
@@ -256,18 +256,18 @@ namespace PixArt {
             int64_t hp = H / p;
 
             auto h = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"])->forward(ctx, x);  // [N, hidden, hp, wp]
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));       // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
+            h      = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));                 // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
             h      = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N);              // [N, hp*wp, hidden]
             h      = ggml_add(ctx->ggml_ctx, h, pos_embed);
 
-            auto t   = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000);
+            auto t   = ggml_ext_timestep_embedding(ctx, timesteps, ADALN_EMBED_DIM, 10000);
             auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["t_embedder"])->forward(ctx, t);
 
             auto mod = std::dynamic_pointer_cast<Linear>(blocks["t_block.1"])
                            ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));  // [N, 6 * hidden]
 
             auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context);
-            ctx_emb      = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
+            ctx_emb      = ggml_ext_gelu(ctx, ctx_emb, true);
             ctx_emb      = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb);
 
             for (int i = 0; i < config.num_layers; ++i) {
@@ -284,13 +284,13 @@ namespace PixArt {
             auto ss    = ggml_add(ctx->ggml_ctx,
                                   ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1),
                                   ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N));  // [2, hidden, N]
-            auto parts = ggml_ext_chunk(ctx->ggml_ctx,
-                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
+            auto parts = ggml_ext_chunk(ctx,
+                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, ss), config.hidden_size * 2, N),
                                         2, 0);
             h          = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
             h          = modulate(ctx, h, parts[0], parts[1]);
             h          = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h);  // [N, hp*wp, p*p*out_ch]
-            h          = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
+            h          = DiT::unpatchify(ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
             return h;  // [N, out_channels, H, W]
         }
     };
@@ -369,7 +369,7 @@ namespace PixArt {
             auto runner_ctx  = get_context();
             ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);
             // learn_sigma: keep the noise prediction half of the output channels
-            out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels);
+            out = ggml_ext_slice(&runner_ctx, out, 2, 0, config.in_channels);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/diffusion/qwen_image.hpp

@@ -113,7 +113,7 @@ namespace Qwen {
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
 
-            auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
+            auto timesteps_proj = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1.f);
             auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
             if (use_additional_t_cond) {
                 GGML_ASSERT(addition_t_cond != nullptr);
@@ -186,9 +186,7 @@ namespace Qwen {
             auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto norm_added_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_q"]);
             auto norm_added_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_k"]);
@@ -289,14 +287,17 @@ namespace Qwen {
                                                                                eps));
         }
 
-        std::vector<ggml_tensor*> get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) {
+        std::vector<ggml_tensor*> get_mod_params_vec(GGMLRunnerContext* runner_ctx,
+                                                     ggml_tensor* mod_params,
+                                                     ggml_tensor* index = nullptr) {
+            auto ctx = runner_ctx->ggml_ctx;
             // index: [N, n_img_token]
             // mod_params: [N, hidden_size * 12]
             if (index == nullptr) {
-                return ggml_ext_chunk(ctx, mod_params, 6, 0);
+                return ggml_ext_chunk(runner_ctx, mod_params, 6, 0);
             }
             mod_params          = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params));
-            auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0);
+            auto mod_params_vec = ggml_ext_chunk(runner_ctx, mod_params, 12, 0);
             index               = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]);                                      // [N, n_img_token, 1]
             index               = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]);  // [N, n_img_token, hidden_size]
             std::vector<ggml_tensor*> mod_results;
@@ -339,15 +340,15 @@ namespace Qwen {
 
             auto img_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             img_mod_params         = img_mod_1->forward(ctx, img_mod_params);
-            auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index);
+            auto img_mod_param_vec = get_mod_params_vec(ctx, img_mod_params, modulate_index);
 
             if (zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
+                t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0];
             }
 
             auto txt_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             txt_mod_params         = txt_mod_1->forward(ctx, txt_mod_params);
-            auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
+            auto txt_mod_param_vec = get_mod_params_vec(ctx, txt_mod_params);
 
             auto img_normed    = img_norm1->forward(ctx, img);
             auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
@@ -402,7 +403,7 @@ namespace Qwen {
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
 
             auto emb   = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto mods  = ggml_ext_chunk(ctx, emb, 2, 0);
             auto scale = mods[0];
             auto shift = mods[1];
 
@@ -461,7 +462,7 @@ namespace Qwen {
 
             auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
             if (config.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
+                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx, timestep), addition_t_cond);
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
             auto img = img_in->forward(ctx, x);
@@ -482,7 +483,7 @@ namespace Qwen {
             }
 
             if (config.zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
+                t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0];
             }
 
             img = norm_out->forward(ctx, img, t_emb);
@@ -519,12 +520,12 @@ namespace Qwen {
             auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
                 input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
                 if (!has_time_axis) {
-                    return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
+                    return DiT::patchify(ctx, input, config.patch_size, config.patch_size);
                 }
                 if (input->ne[3] == 1) {
                     input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
                 }
-                return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+                return DiT::patchify_3d(ctx, input, 1, config.patch_size, config.patch_size, N);
             };
 
             auto img           = patchify_input(x);
@@ -550,11 +551,11 @@ namespace Qwen {
                 int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
                 int h_len = static_cast<int>((H + pad_h) / config.patch_size);
                 int w_len = static_cast<int>((W + pad_w) / config.patch_size);
-                out       = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
-                out       = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-                out       = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+                out       = DiT::unpatchify_3d(ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
+                out       = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+                out       = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
             } else {
-                out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+                out = DiT::unpatchify_and_crop(ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
             }
 
             return out;

src/model/diffusion/qwen_image_2_1.hpp

@@ -146,7 +146,7 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto weight = params["weight"];
             if (ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
             }
             weight = ggml_scale_bias(ctx->ggml_ctx, weight, 1.f, 1.f);
             return ggml_mul(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, x, eps), weight);
@@ -164,7 +164,7 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             x = std::dynamic_pointer_cast<QwenImage21ZeroCenterRMSNorm>(blocks["text_norm"])->forward(ctx, x);
             x = std::dynamic_pointer_cast<Linear>(blocks["in_layer"])->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             return std::dynamic_pointer_cast<Linear>(blocks["out_layer"])->forward(ctx, x);
         }
     };
@@ -196,7 +196,7 @@ namespace Qwen {
                 // prefix before its full-sequence K/V can accumulate across layers.
                 ctx->expand_graph(q);
                 auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
-                    auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
+                    auto part = ggml_ext_slice(ctx, tensor, axis, 0, cache.prefix_length);
                     // Pack the contiguous data into wider rows so quantization blocks
                     // can exceed head_dim without padding or changing element order.
                     part      = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
@@ -236,17 +236,15 @@ namespace Qwen {
             } else {
                 for (size_t i = 0; i < segments.size(); ++i) {
                     const auto& segment = segments[i];
-                    auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
-                    auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
-                    auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
+                    auto sq             = ggml_ext_slice(ctx, q, 1, segment.start, segment.end);
+                    auto sk             = ggml_ext_slice(ctx, k, 1, 0, segment.end);
+                    auto sv             = ggml_ext_slice(ctx, v, 2, 0, segment.end);
                     auto out            = attend(sq, sk, sv, masks[i]);
                     result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
                 }
             }
             auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out->set_force_prec_f32(true);
-            }
+            to_out->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
             return to_out->forward(ctx, result);
         }
     };
@@ -268,7 +266,7 @@ namespace Qwen {
 
         static ggml_tensor* modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
             auto ctx   = runner_ctx->ggml_ctx;
-            auto rows  = ggml_ext_chunk(ctx, params, 2, 1);
+            auto rows  = ggml_ext_chunk(runner_ctx, params, 2, 1);
             auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
                 row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
                 runner_ctx->expand_graph(row);
@@ -277,8 +275,8 @@ namespace Qwen {
             if (prefix_length == 0) {
                 return apply(x, rows[0]);
             }
-            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
-            auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
+            auto target = apply(ggml_ext_slice(runner_ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
+            auto prefix = apply(ggml_ext_slice(runner_ctx, x, 1, 0, prefix_length), rows[1]);
             return ggml_concat(ctx, prefix, target, 1);
         }
 
@@ -323,27 +321,27 @@ namespace Qwen {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
-            auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
+            auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx, timestep), 0);
             // Runtime flow timesteps already use the [0, 1000] scale.
-            time               = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
+            time               = ggml_ext_timestep_embedding(ctx, time, 256, 10000, 1.f);
             time               = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["time_text_embed.timestep_embedder"])->forward(ctx, time);
             time               = ggml_silu(ctx->ggml_ctx, time);
             auto modulation    = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
-            auto mod           = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
+            auto mod           = ggml_ext_chunk(ctx, modulation, 4, 0);
             auto img_in        = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
             ggml_tensor* joint = nullptr;
             if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
-                joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
+                joint = img_in->forward(ctx, DiT::patchify(ctx, x, 1, 1));
             } else {
                 auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
                 for (const auto& segment : layout.segments) {
                     ggml_tensor* h;
                     if (segment.image_index < 0) {
-                        h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
+                        h = ggml_ext_slice(ctx, text, 1, segment.context_start,
                                            segment.context_start + segment.end - segment.start);
                     } else {
                         auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
-                        h          = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
+                        h          = img_in->forward(ctx, DiT::patchify(ctx, image, 1, 1));
                     }
                     joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
                 }
@@ -359,13 +357,13 @@ namespace Qwen {
                 sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
             }
             if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
-                joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
+                joint = ggml_ext_slice(ctx, joint, 1, layout.prefix_length, joint->ne[1]);
             }
-            auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
+            auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx, time, 2, 1)[0]);
             joint      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
             joint      = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
             joint      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, joint);
-            return DiT::unpatchify_and_crop(ctx->ggml_ctx, joint, x->ne[1], x->ne[0], 1, 1);
+            return DiT::unpatchify_and_crop(ctx, joint, x->ne[1], x->ne[0], 1, 1);
         }
     };
 
@@ -578,7 +576,7 @@ namespace Qwen {
                 result = run(cache);
             }
             if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
-                (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
+                (cache.mode != QwenImage21PrefixCache::Mode::NONE || has_cache_tensors())) {
                 // The failed graph has ended before persistent inputs are released.
                 free_cache_ctx_and_buffer();
                 prefix_cache_disabled = true;

src/model/diffusion/relay_vsr.hpp

@@ -124,7 +124,7 @@ namespace RelayVSR {
                 k   = VSR::apply_rope(ctx, k, key_rope, heads_);
                 out = local->forward(ctx, q, k, v, heads_);
             } else {
-                out = ggml_ext_attention_ext(c, ctx->backend, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled);
+                out = ggml_ext_attention_ext(ctx, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled);
             }
             return to_out_0->forward(ctx, out);
         }
@@ -157,7 +157,7 @@ namespace RelayVSR {
             auto ffn_net_2      = std::dynamic_pointer_cast<Linear>(blocks["ffn.net.2"]);
             auto c              = ctx->ggml_ctx;
             auto modulation     = ggml_add(c, timestep_proj, params["scale_shift_table"]);
-            auto pieces         = ggml_ext_chunk(c, modulation, 6, 1);
+            auto pieces         = ggml_ext_chunk(ctx, modulation, 6, 1);
             auto attn1          = std::dynamic_pointer_cast<WanAttention>(blocks["attn1"]);
             auto attn2          = std::dynamic_pointer_cast<WanAttention>(blocks["attn2"]);
             auto y              = attn1->forward(ctx, ggml_ext_modulate(ctx, norm1->forward(ctx, x), pieces[0], pieces[1]), nullptr, qpe, kpe, &local, name + ".attn1");
@@ -165,7 +165,7 @@ namespace RelayVSR {
             x                   = ggml_add(c, x, attn2->forward(ctx, norm2->forward(ctx, x), text, nullptr, nullptr, nullptr, name + ".attn2"));
             y                   = ggml_ext_modulate(ctx, norm3->forward(ctx, x), pieces[3], pieces[4]);
             y                   = ffn_net_0_proj->forward(ctx, y);
-            y                   = ffn_net_2->forward(ctx, ggml_ext_gelu(c, y));
+            y                   = ffn_net_2->forward(ctx, ggml_ext_gelu(ctx, y));
             return ggml_add(c, x, ggml_mul(c, y, pieces[5]));
         }
     };
@@ -213,7 +213,7 @@ namespace RelayVSR {
             auto e0   = condition_embedder_time_proj->forward(ctx, ggml_silu(c, e));
             e0        = ggml_reshape_3d(c, e0, config_.dim, 6, 1);
             auto text = condition_embedder_text_embedder_linear_1->forward(ctx, params["prompt_embeds"]);
-            text      = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(c, text));
+            text      = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(ctx, text));
             for (int i = 0; i < config_.num_layers; ++i) {
                 auto name = "blocks." + std::to_string(i);
                 x         = std::dynamic_pointer_cast<WanTransformerBlock>(blocks[name])->forward(ctx, x, text, e0, qpe, kpe, local, name);
@@ -221,7 +221,7 @@ namespace RelayVSR {
             }
             e               = ggml_reshape_3d(c, e, config_.dim, 1, 1);
             auto modulation = ggml_add(c, params["scale_shift_table"], e);
-            auto pieces     = ggml_ext_chunk(c, modulation, 2, 1);
+            auto pieces     = ggml_ext_chunk(ctx, modulation, 2, 1);
             x               = norm_out->forward(ctx, x);
             x               = ggml_ext_modulate(ctx, x, pieces[0], pieces[1]);
             x               = proj_out->forward(ctx, x);

src/model/diffusion/seedvr2.hpp

@@ -139,7 +139,7 @@ namespace SeedVR2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto up     = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, x);
             auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast<Linear>(blocks["proj_in_gate"])->forward(ctx, x), up)
-                                  : ggml_ext_gelu(ctx->ggml_ctx, up);
+                                  : ggml_ext_gelu(ctx, up);
             return std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, hidden);
         }
     };
@@ -200,7 +200,7 @@ namespace SeedVR2 {
             auto vpe  = rotation(ctx, inputs.input(ctx, plan.positions));
             auto tpe  = rotation(ctx, inputs.input(ctx, plan.text_positions));
             auto qkv  = [&](ggml_tensor* x, bool text, ggml_tensor* pe) {
-                auto parts = ggml_ext_chunk(c, x, 3, 0);
+                auto parts = ggml_ext_chunk(ctx, x, 3, 0);
                 for (int i = 0; i < 2; ++i) {
                     auto h   = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]);
                     h        = std::dynamic_pointer_cast<RMSNorm>(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h);
@@ -216,15 +216,15 @@ namespace SeedVR2 {
             for (int length : plan.lengths) {
                 std::vector<ggml_tensor*> joined;
                 for (int i = 0; i < 3; ++i)
-                    joined.push_back(ggml_concat(c, ggml_ext_slice(c, vp[i], 1, start, start + length), tp[i], 1));
-                auto out = ggml_ext_attention_ext(c, ctx->backend, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled);
-                auto vo  = ggml_ext_slice(c, out, 1, 0, length);
-                auto to  = ggml_ext_slice(c, out, 1, length, length + txt->ne[1]);
+                    joined.push_back(ggml_concat(c, ggml_ext_slice(ctx, vp[i], 1, start, start + length), tp[i], 1));
+                auto out = ggml_ext_attention_ext(ctx, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled);
+                auto vo  = ggml_ext_slice(ctx, out, 1, 0, length);
+                auto to  = ggml_ext_slice(ctx, out, 1, length, length + txt->ne[1]);
                 video_windows.push_back(vo);
                 tout = tout ? ggml_add(c, tout, to) : to;
                 start += length;
             }
-            auto vout = ggml_ext_vec_concat(c, video_windows, 1);
+            auto vout = ggml_ext_vec_concat(ctx, video_windows, 1);
             vout      = ggml_get_rows(c, vout, inputs.input(ctx, plan.reverse));
             tout      = ggml_scale(c, tout, 1.f / plan.lengths.size());
             vid       = ggml_add(c, vid, ada(linear("attn.proj_out", false, vout), false, "attn", true));

src/model/diffusion/sefi_image.hpp

@@ -77,8 +77,8 @@ namespace SefiImage {
             auto semantic_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["semantic_embedder"]);
             auto texture_embedder  = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["texture_embedder"]);
 
-            auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
-            auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto sem_proj = ggml_ext_timestep_embedding(ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto tex_proj = ggml_ext_timestep_embedding(ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
             auto sem_emb  = semantic_embedder->forward(ctx, sem_proj);
             auto tex_emb  = texture_embedder->forward(ctx, tex_proj);
             return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0);

src/model/diffusion/sensenova_u1.h

@@ -143,7 +143,7 @@ namespace SenseNovaU1 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto x     = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+            auto x     = ggml_ext_timestep_embedding(ctx,
                                                      timesteps,
                                                      static_cast<int>(frequency_embedding_size),
                                                      10000,
@@ -166,8 +166,8 @@ namespace SenseNovaU1 {
         // head axis used by the RoPE kernel: [hidden, 1, tokens, batch].
         x                  = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
         const int64_t half = x->ne[0] / 2;
-        auto x_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half);
-        auto y_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]);
+        auto x_part        = ggml_ext_slice(ctx, x, 0, 0, half);
+        auto y_part        = ggml_ext_slice(ctx, x, 0, half, x->ne[0]);
         x_part             = ggml_rope_ext(ctx->ggml_ctx,
                                            x_part,
                                            position_x,
@@ -255,10 +255,10 @@ namespace SenseNovaU1 {
         const int64_t h = x->ne[1];
         const int64_t w = x->ne[0];
         GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
-        x = ggml_ext_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        x = ggml_ext_cont(ctx,
+                          ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
         x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);
-        return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
+        return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true);
     }
 
     struct PixelDecoder : public GGMLBlock {
@@ -352,13 +352,13 @@ namespace SenseNovaU1 {
             const int64_t spatial_dim  = config.head_dim - temporal_dim;
             const int64_t axis_dim     = spatial_dim / 2;
 
-            auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim);
-            auto spatial  = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim);
+            auto temporal = ggml_ext_slice(ctx, x, 0, 0, temporal_dim);
+            auto spatial  = ggml_ext_slice(ctx, x, 0, temporal_dim, config.head_dim);
             temporal      = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_name])->forward(ctx, temporal);
             spatial       = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_hw_name])->forward(ctx, spatial);
 
-            auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim);
-            auto width  = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim);
+            auto height = ggml_ext_slice(ctx, spatial, 0, 0, axis_dim);
+            auto width  = ggml_ext_slice(ctx, spatial, 0, axis_dim, spatial_dim);
             temporal    = apply_axis_rope(ctx,
                                           temporal,
                                           position_t,
@@ -448,14 +448,13 @@ namespace SenseNovaU1 {
             }
 
             q = ggml_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));
+                          ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3));
             q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
             k = ggml_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));
+                          ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3));
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
 
-            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                              ctx->backend,
+            auto out = ggml_ext_attention_ext(ctx,
                                               q,
                                               k,
                                               v,

src/model/diffusion/unet.hpp

@@ -247,7 +247,7 @@ public:
 
         auto num_frames = ggml_arange(ctx->ggml_ctx, 0.f, static_cast<float>(timesteps), 1.f);
         // since b is 1, no need to do repeat
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period);  // [N, in_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period);  // [N, in_channels]
 
         auto emb = time_pos_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);
@@ -572,7 +572,7 @@ public:
         auto out_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out.0"]);
         auto out_2 = std::dynamic_pointer_cast<Conv2d>(blocks["out.2"]);
 
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);
@@ -685,7 +685,7 @@ public:
         sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
         if (controls.size() > 0) {
             // not in place, it segfaults the Vulkan backend with SDXL ControlNets (seen on RADV)
-            auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, false);
+            auto cs = ggml_ext_scale(ctx, controls[controls.size() - 1], control_strength, false);
             h       = ggml_add(ctx->ggml_ctx, h, cs);  // middle control
         }
         int control_offset = static_cast<int>(controls.size() - 2);
@@ -699,7 +699,7 @@ public:
 
                 if (controls.size() > 0) {
                     // not in place, see the middle control
-                    auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, false);
+                    auto cs = ggml_ext_scale(ctx, controls[control_offset], control_strength, false);
                     h_skip  = ggml_add(ctx->ggml_ctx, h_skip, cs);  // control net condition
                     control_offset--;
                 }

src/model/diffusion/wan.hpp

@@ -362,7 +362,7 @@ namespace WAN {
 
             auto modulation = params["modulation"];
             e               = ggml_add(ctx->ggml_ctx, e, modulation);  // [N, 6, dim] or [N, T, 6, dim]
-            auto es         = ggml_ext_chunk(ctx->ggml_ctx, e, 6, 1);  // ([N, 1, dim], ...) or [N, T, 1, dim]
+            auto es         = ggml_ext_chunk(ctx, e, 6, 1);            // ([N, 1, dim], ...) or [N, T, 1, dim]
 
             auto norm1      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
             auto self_attn  = std::dynamic_pointer_cast<WanSelfAttention>(blocks["self_attn"]);
@@ -391,7 +391,7 @@ namespace WAN {
             y = modulate_add(ctx->ggml_ctx, y, es[3]);
 
             y = ffn_0->forward(ctx, y);
-            y = ggml_ext_gelu(ctx->ggml_ctx, y, true);
+            y = ggml_ext_gelu(ctx, y, true);
             y = ffn_2->forward(ctx, y);
 
             x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[5]));
@@ -484,7 +484,7 @@ namespace WAN {
             e               = ggml_repeat_4d(ctx->ggml_ctx, e, e->ne[0], 2, e->ne[2], e->ne[3]);   // [N, 2, dim] or [N, T, 2, dim]
 
             e       = ggml_add(ctx->ggml_ctx, e, modulation);  // [N, 2, dim] or [N, T, 2, dim]
-            auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 2, 1);  // ([N, 1, dim], ...) or ([N, T, 1, dim], ...)
+            auto es = ggml_ext_chunk(ctx, e, 2, 1);            // ([N, 1, dim], ...) or ([N, T, 1, dim], ...)
 
             auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
             auto head = std::dynamic_pointer_cast<Linear>(blocks["head"]);
@@ -525,8 +525,8 @@ namespace WAN {
             if (flf_pos_embed_token_number > 0) {
                 auto emb_pos = params["emb_pos"];
 
-                auto a = ggml_ext_slice(ctx->ggml_ctx, image_embeds, 1, 0, emb_pos->ne[1]);
-                auto b = ggml_ext_slice(ctx->ggml_ctx, emb_pos, 1, 0, image_embeds->ne[1]);
+                auto a = ggml_ext_slice(ctx, image_embeds, 1, 0, emb_pos->ne[1]);
+                auto b = ggml_ext_slice(ctx, emb_pos, 1, 0, image_embeds->ne[1]);
 
                 image_embeds = ggml_add(ctx->ggml_ctx, a, b);
             }
@@ -538,7 +538,7 @@ namespace WAN {
 
             auto x = proj_0->forward(ctx, image_embeds);
             x      = proj_1->forward(ctx, x);
-            x      = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x      = ggml_ext_gelu(ctx, x, true);
             x      = proj_3->forward(ctx, x);
             x      = proj_4->forward(ctx, x);
 
@@ -639,15 +639,16 @@ namespace WAN {
             int pad_t = (std::get<0>(config.patch_size) - T % std::get<0>(config.patch_size)) % std::get<0>(config.patch_size);
             int pad_h = (std::get<1>(config.patch_size) - H % std::get<1>(config.patch_size)) % std::get<1>(config.patch_size);
             int pad_w = (std::get<2>(config.patch_size) - W % std::get<2>(config.patch_size)) % std::get<2>(config.patch_size);
-            ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            ggml_ext_pad(ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
             return x;
         }
 
-        ggml_tensor* unpatchify(ggml_context* ctx,
+        ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                 ggml_tensor* x,
                                 int64_t t_len,
                                 int64_t h_len,
                                 int64_t w_len) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
             // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
             int64_t N  = x->ne[3];
@@ -658,15 +659,15 @@ namespace WAN {
 
             GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
-            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
             return x;
         }
 
@@ -707,7 +708,7 @@ namespace WAN {
             // patch_embedding
             x = patch_embedding->forward(ctx, x);                                                    // [N*dim, t_len, h_len, w_len]
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));                      // [N, t_len*h_len*w_len, dim]
 
             ggml_tensor* audio_local  = nullptr;
             ggml_tensor* audio_global = nullptr;
@@ -725,26 +726,26 @@ namespace WAN {
 
                 // video tokens get cond_mask[0], reference tokens cond_mask[1]
                 auto cond_mask = params["trainable_cond_mask.weight"];
-                auto cm0       = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
+                auto cm0       = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
                 x              = ggml_add(ctx->ggml_ctx, x, cm0);
 
                 if (reference_latent != nullptr) {
                     t_ref_len = reference_latent->ne[2];
                     auto ref  = patch_embedding->forward(ctx, reference_latent);
                     ref       = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N);
-                    ref       = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3));  // [N, t_ref*h_len*w_len, dim]
-                    auto cm1  = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
+                    ref       = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, ref, 1, 0, 2, 3));  // [N, t_ref*h_len*w_len, dim]
+                    auto cm1  = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
                     ref       = ggml_add(ctx->ggml_ctx, ref, cm1);
                     x         = ggml_concat(ctx->ggml_ctx, x, ref, 1);
 
                     // Reference tokens use timestep 0.
                     GGML_ASSERT(timestep->ne[0] == T);
-                    timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0);
+                    timestep = ggml_ext_pad(ctx, timestep, (int)t_ref_len, 0, 0, 0);
                 }
             }
 
             // time_embedding
-            auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim);
+            auto e = ggml_ext_timestep_embedding(ctx, timestep, config.freq_dim);
             e      = time_embedding_0->forward(ctx, e);
             e      = ggml_silu_inplace(ctx->ggml_ctx, e);
             e      = time_embedding_2->forward(ctx, e);  // [N, dim] or [N, T, dim]
@@ -755,7 +756,7 @@ namespace WAN {
             e0      = ggml_reshape_4d(ctx->ggml_ctx, e0, e0->ne[0] / 6, 6, e0->ne[1], e0->ne[2]);  //  [N, 6, dim] or [N, T, 6, dim]
 
             context = text_embedding_0->forward(ctx, context);
-            context = ggml_ext_gelu(ctx->ggml_ctx, context);
+            context = ggml_ext_gelu(ctx, context);
             context = text_embedding_2->forward(ctx, context);  // [N, context_txt_len, dim]
 
             int64_t context_img_len = 0;
@@ -775,7 +776,7 @@ namespace WAN {
 
                 c = vace_patch_embedding->forward(ctx, vace_context);                                    // [N*dim, t_len, h_len, w_len]
                 c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
-                c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
+                c = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, c, 1, 0, 2, 3));                      // [N, t_len*h_len*w_len, dim]
             }
             sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x");
             // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e");
@@ -806,7 +807,7 @@ namespace WAN {
                     auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len);
                     auto c_skip = result.first;
                     c           = result.second;
-                    c_skip      = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
+                    c_skip      = ggml_ext_scale(ctx, c_skip, vace_strength);
                     x           = ggml_add(ctx->ggml_ctx, x, c_skip);
                 }
 
@@ -870,16 +871,16 @@ namespace WAN {
 
             if (reference_latent != nullptr) {
                 // Exclude reference tokens from the generated video.
-                out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len);
+                out = ggml_ext_slice(ctx, out, 1, 0, t_len * h_len * w_len);
             }
 
-            out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
+            out = unpatchify(ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
             // slice
 
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            out = ggml_ext_slice(ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
 
             return out;
         }

src/model/diffusion/wan_audio.hpp

@@ -27,7 +27,7 @@ namespace WAN {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             // Replicate the first sample for causal left padding.
             if (kernel_size_ > 1) {
-                auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1);
+                auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
                 for (int i = 0; i < kernel_size_ - 1; i++) {
                     x = ggml_concat(ctx->ggml_ctx, first, x, 0);
                 }
@@ -57,7 +57,7 @@ namespace WAN {
             x = std::dynamic_pointer_cast<LayerNorm>(blocks[norm_key])->forward(ctx, x);
             x = ggml_silu(ctx->ggml_ctx, x);
             if (to_conv_layout) {
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+                x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
             }
             return x;
         }
@@ -87,11 +87,11 @@ namespace WAN {
             auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
             std::vector<ggml_tensor*> tokens;
             // Each token group is normalized independently over channels.
-            for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) {
+            for (auto& group : ggml_ext_chunk(ctx, local, num_token_, 1)) {
                 ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3);
                 s              = norm1->forward(ctx, s);
                 s              = ggml_silu(ctx->ggml_ctx, s);
-                s              = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
+                s              = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
                 s              = conv_norm_silu(ctx, s, "conv2", "norm2", true);
                 s              = conv_norm_silu(ctx, s, "conv3", "norm3", false);
                 tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1]));
@@ -99,7 +99,7 @@ namespace WAN {
             auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1);
             padding      = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]);
             tokens.push_back(padding);
-            ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1);
+            ggml_tensor* local_out = ggml_ext_vec_concat(ctx, tokens, 1);
 
             if (!need_global_) {
                 return {local_out, nullptr};
@@ -141,10 +141,10 @@ namespace WAN {
             auto x       = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_));
             x            = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights));
             // Move the layer axis to ggml dimension 0 for reduction.
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
             x = ggml_sum_rows(ctx->ggml_ctx, x);
             x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             return std::dynamic_pointer_cast<WanMotionEncoder>(blocks["encoder"])->forward(ctx, x);
         }
     };
@@ -187,12 +187,12 @@ namespace WAN {
 
             auto temb  = ggml_silu(ctx->ggml_ctx, audio_global);
             temb       = adain_linear->forward(ctx, temb);
-            auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_);
-            auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2);
+            auto shift = ggml_ext_slice(ctx, temb, 0, 0, dim_);
+            auto scale = ggml_ext_slice(ctx, temb, 0, dim_, dim_ * 2);
             shift      = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T);
             scale      = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T);
 
-            auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len);
+            auto x_vid = ggml_ext_slice(ctx, x, 1, 0, seq_len);
             auto h     = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T);
             h          = adain_norm->forward(ctx, h);
             h          = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale));
@@ -203,7 +203,7 @@ namespace WAN {
 
             auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res);
             if (seq_len < n_token) {
-                auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token);
+                auto x_tail = ggml_ext_slice(ctx, x, 1, seq_len, n_token);
                 return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1);
             }
             return x_head;

src/model/diffusion/z_image.hpp

@@ -179,12 +179,10 @@ namespace ZImage {
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
                 auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
 
-                if (sd_backend_is(ctx->backend, "ROCm")) {
-                    out_proj->set_force_prec_f32(true);
-                    q_proj->set_force_prec_f32(true);
-                    k_proj->set_force_prec_f32(true);
-                    v_proj->set_force_prec_f32(true);
-                }
+                out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                q_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                k_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                v_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
                 auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N);
                 auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
@@ -201,10 +199,8 @@ namespace ZImage {
 
             auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
-            if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_force_prec_f32(true);
-                qkv_proj->set_force_prec_f32(true);
-            }
+            out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            qkv_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
             auto qkv = qkv_proj->forward(ctx, x);                                                                            // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
             qkv      = ggml_reshape_4d(ctx->ggml_ctx, qkv, head_dim, num_heads + num_kv_heads * 2, qkv->ne[1], qkv->ne[2]);  // [N, n_token, num_heads + num_kv_heads*2, head_dim]
@@ -281,9 +277,7 @@ namespace ZImage {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                w2->set_force_prec_f32(true);
-            }
+            w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto x1 = w1->forward(ctx, x);
             auto x3 = w3->forward(ctx, x);
@@ -353,7 +347,7 @@ namespace ZImage {
                 auto adaLN_modulation_0 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
 
                 auto m         = adaLN_modulation_0->forward(ctx, adaln_input);  // [N, 4 * hidden_size]
-                auto mods      = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
+                auto mods      = ggml_ext_chunk(ctx, m, 4, 0);
                 auto scale_msa = mods[0];
                 auto gate_msa  = mods[1];
                 auto scale_mlp = mods[2];
@@ -457,12 +451,10 @@ namespace ZImage {
             int64_t n_token = c->ne[1];
             int64_t N       = c->ne[2];
 
-            if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_force_prec_f32(true);
-                to_q->set_force_prec_f32(true);
-                to_k->set_force_prec_f32(true);
-                to_v->set_force_prec_f32(true);
-            }
+            out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_q->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_k->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_v->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
             auto q = ggml_reshape_4d(ctx->ggml_ctx, to_q->forward(ctx, c), head_dim, num_heads, n_token, N);
             auto k = ggml_reshape_4d(ctx->ggml_ctx, to_k->forward(ctx, c), head_dim, num_kv_heads, n_token, N);
@@ -497,7 +489,7 @@ namespace ZImage {
             }
 
             auto m         = adaLN_modulation_0->forward(ctx, t_emb);
-            auto mods      = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
+            auto mods      = ggml_ext_chunk(ctx, m, 4, 0);
             auto scale_msa = mods[0];
             auto gate_msa  = mods[1];
             auto scale_mlp = mods[2];
@@ -713,7 +705,7 @@ namespace ZImage {
 
                 auto out = block->forward(ctx, c, img_pe, img, t_emb);
                 c        = out.second;
-                ref_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+                ref_hints.push_back(ggml_ext_scale(ctx, out.first, strength));
                 sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control");
             }
 
@@ -757,7 +749,7 @@ namespace ZImage {
 
                 auto out     = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb);
                 ctrl_unified = out.second;
-                main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+                main_hints.push_back(ggml_ext_scale(ctx, out.first, strength));
                 sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control");
             }
 
@@ -807,8 +799,8 @@ namespace ZImage {
 
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
 
             std::vector<ggml_tensor*> ref_hints;
             ggml_tensor* control_ctx = nullptr;
@@ -858,7 +850,7 @@ namespace ZImage {
 
             txt_img = final_layer->forward(ctx, txt_img, t_emb);  // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
 
-            img = ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);  // [N, n_img_token, ph*pw*C]
+            img = ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);  // [N, n_img_token, ph*pw*C]
 
             return img;
         }
@@ -906,10 +898,10 @@ namespace ZImage {
 
             auto out = forward_core(ctx, img, timestep, context, pe, control, control_strength);
 
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token);                              // [N, n_img_token, ph*pw*C]
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);  // [N, C, H, W]
+            out = ggml_ext_slice(ctx, out, 1, 0, n_img_token);                              // [N, n_img_token, ph*pw*C]
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false);  // [N, C, H, W]
 
-            out = ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
+            out = ggml_ext_scale(ctx, out, -1.f);
 
             return out;
         }

src/model/diffusion/z_image_l2p.hpp

@@ -134,7 +134,7 @@ namespace ZImageL2P {
             const int64_t n_img = w * h;
             const int64_t n_txt = context->ne[1];
 
-            auto img = DiT::patchify(gctx, x, config.patch_size, config.patch_size, false);
+            auto img = DiT::patchify(ctx, x, config.patch_size, config.patch_size, false);
             img      = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"])->forward(ctx, img);
             auto txt = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"])->forward(ctx, context);
             txt      = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"])->forward(ctx, txt);
@@ -155,8 +155,8 @@ namespace ZImageL2P {
             }
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
             for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
                 txt = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe);
                 sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.context_refiner." + std::to_string(i), "txt");
@@ -175,11 +175,11 @@ namespace ZImageL2P {
             // The local decoder consumes the raw last-layer hidden states: there is no final
             // norm or adaLN modulation in front of it, unlike Z-Image's final_layer.
             const int64_t img_start = n_txt + n_txt_pad;
-            auto feat               = ggml_ext_slice(gctx, txt_img, 1, img_start, img_start + n_img);                           // [N, h*w, hidden_size]
+            auto feat               = ggml_ext_slice(ctx, txt_img, 1, img_start, img_start + n_img);                            // [N, h*w, hidden_size]
             feat                    = ggml_reshape_4d(gctx, feat, config.hidden_size, w, h, N);                                 // [N, h, w, hidden_size]
             feat                    = ggml_cont(gctx, ggml_permute(gctx, feat, 2, 0, 1, 3));                                    // [N, hidden_size, h, w]
             auto out                = std::dynamic_pointer_cast<LocalDecoder>(blocks["local_decoder"])->forward(ctx, x, feat);  // [N, C, H, W]
-            return ggml_ext_scale(gctx, out, -1.f);
+            return ggml_ext_scale(ctx, out, -1.f);
         }
     };
 

src/model/integrations/minimax_h3.cpp

@@ -41,6 +41,15 @@ const std::vector<ModelDefinition>& minimax_h3_definitions() {
             definition.build                            = build_minimax_h3;
             result.push_back(std::move(definition));
         }
+        {
+            ModelDefinition definition                  = result.front();
+            definition.version                          = VERSION_MINIMAX_H3_IMAGE;
+            definition.name                             = "MiniMax-H3-Image";
+            definition.traits.supports_video_generation = false;
+            definition.sampling.default_flow_shift      = 1.f;
+            definition.ref_image_preset                 = "qwen";
+            result.push_back(std::move(definition));
+        }
         {
             ModelDefinition definition                  = result.front();
             definition.version                          = VERSION_P_H3_IMAGE;

src/model/te/clip.hpp

@@ -35,9 +35,9 @@ public:
 
         x = fc1->forward(ctx, x);
         if (use_gelu) {
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
         } else {
-            x = ggml_ext_gelu_quick(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu_quick(ctx, x, true);
         }
         x = fc2->forward(ctx, x);
         return x;
@@ -228,10 +228,10 @@ public:
         // concat(patch_embedding, class_embedding) + position_embedding
         ggml_tensor* patch_embedding;
         int64_t N       = pixel_values->ne[3];
-        patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size);  // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
-        patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N);                          // [N, embed_dim, num_patches]
-        patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3));                  // [N, num_patches, embed_dim]
-        patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N);                       // [N, num_patches, embed_dim, 1]
+        patch_embedding = ggml_ext_conv_2d(ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size);  // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
+        patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N);                // [N, embed_dim, num_patches]
+        patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3));        // [N, num_patches, embed_dim]
+        patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N);             // [N, num_patches, embed_dim, 1]
 
         ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
         class_embedding              = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding);      // [N, embed_dim]
@@ -325,7 +325,7 @@ public:
             auto text_projection = params["text_projection"];
             ggml_tensor* pooled  = ggml_view_1d(ctx->ggml_ctx, x, hidden_size, x->nb[1] * max_token_idx);
             if (text_projection != nullptr) {
-                pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr);
+                pooled = ggml_ext_linear(ctx, pooled, text_projection, nullptr);
             } else {
                 LOG_VERBOSE("identity projection");
             }
@@ -426,7 +426,7 @@ public:
         if (transpose_weight) {
             w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
         }
-        return ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr);
+        return ggml_ext_linear(ctx, x, w, nullptr);
     }
 };
 

src/model/te/llada_image_te.hpp

@@ -78,15 +78,15 @@ namespace LLaDAImageTE {
             auto b = params["in_proj_bias"];
 
             auto slice_w = [&](int64_t index) {
-                return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
+                return ggml_ext_slice(ctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
             };
             auto slice_b = [&](int64_t index) {
-                return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
+                return ggml_ext_slice(ctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
             };
 
-            auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0));
-            auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1));
-            auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2));
+            auto q = ggml_ext_linear(ctx, query, slice_w(0), slice_b(0));
+            auto k = ggml_ext_linear(ctx, context, slice_w(1), slice_b(1));
+            auto v = ggml_ext_linear(ctx, context, slice_w(2), slice_b(2));
 
             auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, num_queries, hidden_size]
             return out_proj->forward(ctx, x);
@@ -297,9 +297,9 @@ namespace LLaDAImageTE {
 
             int64_t hidden_size = num_heads * head_dim;
             auto qkv            = qkv_proj->forward(ctx, x);
-            auto q              = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size);
-            auto k              = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2);
-            auto v              = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3);
+            auto q              = ggml_ext_slice(ctx, qkv, 0, 0, hidden_size);
+            auto k              = ggml_ext_slice(ctx, qkv, 0, hidden_size, hidden_size * 2);
+            auto v              = ggml_ext_slice(ctx, qkv, 0, hidden_size * 2, hidden_size * 3);
 
             auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
             return out_proj->forward(ctx, out);

src/model/te/llm.hpp

@@ -485,7 +485,7 @@ namespace LLM {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* w = params["weight"];
             if (ctx->weight_adapter) {
-                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             }
             x           = ggml_rms_norm(ctx->ggml_ctx, x, eps);
             auto scaled = ggml_mul(ctx->ggml_ctx, x, w);
@@ -941,7 +941,7 @@ namespace LLM {
             // An embed can sit flush against the previous one or at the very start/end of the
             // sequence, leaving no text tokens to splice around it.
             if (txt_token_end > txt_token_start) {
-                auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+                auto txt_embed = ggml_ext_slice(ctx, raw_x, 1, txt_token_start, txt_token_end);
                 if (input_embed == nullptr) {
                     input_embed = txt_embed;
                 } else {
@@ -960,7 +960,7 @@ namespace LLM {
         txt_token_end   = raw_x->ne[1];
 
         if (txt_token_end > txt_token_start) {
-            auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+            auto final_txt_embed = ggml_ext_slice(ctx, raw_x, 1, txt_token_start, txt_token_end);
             input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
         }
         GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]);
@@ -982,16 +982,16 @@ namespace LLM {
             GGML_ASSERT(index >= token_start);
             GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]);
             if (index > token_start) {
-                auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index);
+                auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, index);
                 output          = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
             }
-            auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]);
+            auto visual_embed = ggml_ext_slice(ctx, raw_x, 1, index, index + image_embed->ne[1]);
             visual_embed      = ggml_add(ctx->ggml_ctx, visual_embed, image_embed);
             output            = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1);
             token_start       = index + image_embed->ne[1];
         }
         if (token_start < raw_x->ne[1]) {
-            auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]);
+            auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, raw_x->ne[1]);
             output          = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
         }
         GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]);
@@ -1020,7 +1020,7 @@ namespace LLM {
                 auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
                 auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
                 x               = linear_fc1->forward(ctx, x);
-                x               = ggml_ext_gelu(ctx->ggml_ctx, x);
+                x               = ggml_ext_gelu(ctx, x);
                 x               = linear_fc2->forward(ctx, x);
             } else {
                 auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
@@ -1105,11 +1105,11 @@ namespace LLM {
                 auto proj_0 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.0"]);
                 auto proj_1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.1"]);
 
-                auto x0 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto x0 = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x0      = ggml_reshape_4d(ctx->ggml_ctx, x0, x0->ne[0], x0->ne[1], in_channels, x0->ne[3] / in_channels);
                 x0      = proj_0->forward(ctx, x0);
 
-                auto x1 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, 2);
+                auto x1 = ggml_ext_slice(ctx, x, 2, 1, 2);
                 x1      = ggml_reshape_4d(ctx->ggml_ctx, x1, x1->ne[0], x1->ne[1], in_channels, x1->ne[3] / in_channels);
                 x1      = proj_1->forward(ctx, x1);
 
@@ -1173,7 +1173,7 @@ namespace LLM {
             x = ln_q->forward(ctx, x);
             x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size);
             x = mlp_0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             x = mlp_2->forward(ctx, x);
             return x;
         }
@@ -1252,7 +1252,7 @@ namespace LLM {
             } else {
                 auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
                 auto qkv      = qkv_proj->forward(ctx, x);
-                qkv_vec       = split_qkv(ctx->ggml_ctx, qkv);
+                qkv_vec       = ggml_ext_chunk(ctx, qkv, 3, 0);
             }
 
             auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);  // [N, n_token, n_head, d_head]
@@ -1487,12 +1487,12 @@ namespace LLM {
             for (int i = 0; i < kernel; ++i) {
                 auto w    = ggml_cont(c, ggml_view_2d(c, params["conv1d.weight"], 1, conv_dim, kernel * sizeof(float), i * sizeof(float)));
                 w         = ggml_reshape_1d(c, w, conv_dim);
-                auto part = ggml_mul(c, ggml_ext_slice(c, padded, 1, i, i + n), w);
+                auto part = ggml_mul(c, ggml_ext_slice(ctx, padded, 1, i, i + n), w);
                 mixed     = mixed ? ggml_add(c, mixed, part) : part;
             }
             mixed      = ggml_silu(c, mixed);
             auto heads = [&](int start, int count) {
-                auto h = ggml_ext_slice(c, mixed, 0, start, start + head_dim * count);
+                auto h = ggml_ext_slice(ctx, mixed, 0, start, start + head_dim * count);
                 return ggml_reshape_4d(c, h, head_dim, count, n, 1);
             };
             auto repeat_keys = [&](ggml_tensor* h) {
@@ -1514,7 +1514,7 @@ namespace LLM {
             auto a     = gguf_layout ? params["A"] : ggml_neg(c, ggml_exp(c, params["A_log"]));
             decay      = ggml_mul(c, decay, a);
             decay      = ggml_reshape_4d(c, decay, 1, value_heads, n, 1);
-            auto state = ggml_ext_zeros(c, head_dim, head_dim, value_heads, 1);
+            auto state = ggml_ext_zeros(ctx, head_dim, head_dim, value_heads, 1);
             auto out   = ggml_gated_delta_net(c, q, k, v, decay, beta, state, 1);
             out        = ggml_view_3d(c, out, head_dim, value_heads, n, head_dim * sizeof(float), head_dim * value_heads * sizeof(float), 0);
             out        = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])->forward(ctx, out);
@@ -1594,13 +1594,14 @@ namespace LLM {
         // ggml applies theta_i / freq_factors[i], so a huge factor collapses the angle to zero and
         // leaves that pair unrotated. This reproduces transformers' "proportional" RoPE, whose
         // inv_freq is zero-padded past `rope_pairs`, without reordering the head.
-        ggml_tensor* rope_freq_factors(ggml_context* ctx) const {
+        ggml_tensor* rope_freq_factors(GGMLRunnerContext* runner_ctx) const {
+            auto ctx  = runner_ctx->ggml_ctx;
             int pairs = head_dim / 2;
             if (rope_pairs >= pairs) {
                 return nullptr;
             }
-            auto rotated   = ggml_ext_ones(ctx, rope_pairs, 1, 1, 1);
-            auto unrotated = ggml_ext_full(ctx, 1e30f, pairs - rope_pairs, 1, 1, 1);
+            auto rotated   = ggml_ext_ones(runner_ctx, rope_pairs, 1, 1, 1);
+            auto unrotated = ggml_ext_full(runner_ctx, 1e30f, pairs - rope_pairs, 1, 1, 1);
             return ggml_concat(ctx, rotated, unrotated, 0);
         }
 
@@ -1624,9 +1625,9 @@ namespace LLM {
                 auto qkv      = qkv_proj->forward(ctx, x);  // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
                 int64_t q_len = num_heads * head_dim;
                 int64_t k_len = num_kv_heads * head_dim;
-                q             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len);
-                k             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len);
-                v             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
+                q             = ggml_ext_slice(ctx, qkv, 0, 0, q_len);
+                k             = ggml_ext_slice(ctx, qkv, 0, q_len, q_len + k_len);
+                v             = ggml_ext_slice(ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
             } else {
                 auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
@@ -1640,9 +1641,9 @@ namespace LLM {
             ggml_tensor* output_gate = nullptr;
             if (arch == LLMArch::QWEN3_5) {
                 q           = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim * 2, num_heads, n_token, N);
-                output_gate = ggml_ext_slice(ctx->ggml_ctx, q, 0, head_dim, head_dim * 2);
+                output_gate = ggml_ext_slice(ctx, q, 0, head_dim, head_dim * 2);
                 output_gate = ggml_reshape_3d(ctx->ggml_ctx, output_gate, head_dim * num_heads, n_token, N);
-                q           = ggml_ext_slice(ctx->ggml_ctx, q, 0, 0, head_dim);
+                q           = ggml_ext_slice(ctx, q, 0, 0, head_dim);
             }
             q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N);     // [N, n_token, num_heads, head_dim]
             k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N);  // [N, n_token, num_kv_heads, head_dim]
@@ -1666,12 +1667,12 @@ namespace LLM {
                 auto rotate = [&](ggml_tensor* input) {
                     auto gctx    = ctx->ggml_ctx;
                     int64_t half = rope_cos->ne[0];
-                    auto first   = ggml_ext_slice(gctx, input, 0, 0, half);
-                    auto second  = ggml_ext_slice(gctx, input, 0, half, 2 * half);
+                    auto first   = ggml_ext_slice(ctx, input, 0, 0, half);
+                    auto second  = ggml_ext_slice(ctx, input, 0, half, 2 * half);
                     auto left    = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin));
                     auto right   = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin));
                     auto rotated = ggml_concat(gctx, left, right, 0);
-                    return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0);
+                    return ggml_concat(gctx, rotated, ggml_ext_slice(ctx, input, 0, 2 * half, input->ne[0]), 0);
                 };
                 q = rotate(q);
                 k = rotate(k);
@@ -1746,7 +1747,7 @@ namespace LLM {
                                                  1.f);
             } else if (arch == LLMArch::GEMMA4_12B) {
                 float rope_theta  = (rope_index == 1 ? 10000.0f : 1000000.0f);
-                auto freq_factors = rope_freq_factors(ctx->ggml_ctx);
+                auto freq_factors = rope_freq_factors(ctx);
                 q                 = ggml_rope_ext(ctx->ggml_ctx,
                                                   q,
                                                   input_pos,
@@ -1844,18 +1845,18 @@ namespace LLM {
 
             if (unscaled_attention) {
                 // Gemma 4 attends with scaling=1.0; undo the helper's own 1/sqrt(head_dim).
-                q = ggml_ext_scale(ctx->ggml_ctx, q, std::sqrt(static_cast<float>(head_dim)));
+                q = ggml_ext_scale(ctx, q, std::sqrt(static_cast<float>(head_dim)));
             }
 
-            q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, num_heads, n_token, head_dim]
-            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);      // [N*num_heads, n_token, head_dim]
+            q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3));        // [N, num_heads, n_token, head_dim]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);  // [N*num_heads, n_token, head_dim]
 
-            k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, num_kv_heads, n_token, head_dim]
-            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);      // [N*num_kv_heads, n_token, head_dim]
+            k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3));        // [N, num_kv_heads, n_token, head_dim]
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);  // [N*num_kv_heads, n_token, head_dim]
 
             if (arch == LLMArch::GPT_OSS_20B) {
                 GGML_ASSERT(N == 1);
-                auto v_attn = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, kv_heads, head_dim, tokens]
+                auto v_attn = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, kv_heads, head_dim, tokens]
                 v_attn      = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N);
 
                 auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q);
@@ -1870,7 +1871,7 @@ namespace LLM {
                 auto kqv = ggml_mul_mat(ctx->ggml_ctx, v_attn, kq);
                 kqv      = ggml_reshape_4d(ctx->ggml_ctx, kqv, head_dim, n_token, num_heads, N);
                 kqv      = ggml_permute(ctx->ggml_ctx, kqv, 0, 2, 1, 3);
-                x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
+                x        = ggml_ext_cont(ctx, kqv);
                 x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
             } else {
                 x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
@@ -2059,7 +2060,7 @@ namespace LLM {
             std::vector<ggml_tensor*> intermediate_outputs;
 
             if (config.normalize_input) {
-                x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
+                x = ggml_ext_scale(ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
             }
             if (return_all_hidden_states) {
                 intermediate_outputs.push_back(x);
@@ -2207,9 +2208,10 @@ namespace LLM {
         std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
         std::array<std::vector<float>, 4> pos_embed_weight_data_;
 
-        static ggml_tensor* process_image_common(ggml_context* ctx,
+        static ggml_tensor* process_image_common(GGMLRunnerContext* runner_ctx,
                                                  ggml_tensor* image,
                                                  const LLMVisionConfig& vision_params) {
+            auto ctx = runner_ctx->ggml_ctx;
             // image: [C, H, W]
             // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
             int64_t C  = image->ne[2];
@@ -2222,16 +2224,16 @@ namespace LLM {
             int64_t pw = vision_params.patch_size;
 
             image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C);                               // [C*H, (W/mw/pw), mw, pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [mw, C*H, (W/mw/pw), pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1));                   // [mw, C*H, (W/mw/pw), pw]
             image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw);                               // [mw, C, H, (W/mw/pw)*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [H, mw, C, (W/mw/pw)*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1));                   // [H, mw, C, (W/mw/pw)*pw]
             image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph));         // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
             image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph));         // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw]
             image = ggml_concat(ctx, image, image, 0);                                                       // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
             image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph));    // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
             image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph));  // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
             return image;
         }
@@ -2301,7 +2303,6 @@ namespace LLM {
         }
 
         static ggml_tensor* encode_image_common(GGMLRunner* runner,
-                                                ggml_context* compute_ctx,
                                                 GGMLRunnerContext* runner_ctx,
                                                 ggml_tensor* image,
                                                 const LLMVisionConfig& vision_params,
@@ -2313,13 +2314,14 @@ namespace LLM {
                                                 std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
                                                 std::array<std::vector<float>, 4>& pos_embed_weight_data,
                                                 std::vector<ggml_tensor*>* output_tensors = nullptr) {
+            auto compute_ctx = runner_ctx->ggml_ctx;
             GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
             GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
 
             int grid_h = static_cast<int>(image->ne[1]) / vision_params.patch_size;
             int grid_w = static_cast<int>(image->ne[0]) / vision_params.patch_size;
 
-            auto pixel_values = process_image_common(compute_ctx, image, vision_params);
+            auto pixel_values = process_image_common(runner_ctx, image, vision_params);
             int head_dim      = static_cast<int>(vision_params.hidden_size / vision_params.num_heads);
 
             if (vision_params.arch == LLMVisionArch::QWEN3_VL) {
@@ -2678,8 +2680,8 @@ namespace LLM {
             return grid_t * grid_h * grid_w;
         }
 
-        ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
-            return process_image_common(ctx, image, config.vision);
+        ggml_tensor* process_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
+            return process_image_common(runner_ctx, image, config.vision);
         }
 
         ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx,
@@ -2698,7 +2700,6 @@ namespace LLM {
 
         ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             return encode_image_common(this,
-                                       compute_ctx,
                                        runner_ctx,
                                        image,
                                        config.vision,
@@ -2714,7 +2715,6 @@ namespace LLM {
         std::vector<ggml_tensor*> encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             std::vector<ggml_tensor*> outputs;
             encode_image_common(this,
-                                compute_ctx,
                                 runner_ctx,
                                 image,
                                 config.vision,

src/model/te/ming_image_te.hpp

@@ -73,16 +73,16 @@ namespace MingImageTE {
             auto connector            = std::dynamic_pointer_cast<ConnectorModel>(blocks["connector"]);
             auto x                    = backbone->embed(ctx, ids);
             const int64_t query_start = prompt_length + 1;
-            auto before               = ggml_ext_slice(gctx, x, 1, 0, query_start);
-            auto after                = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]);
+            auto before               = ggml_ext_slice(ctx, x, 1, 0, query_start);
+            auto after                = ggml_ext_slice(ctx, x, 1, query_start + config.num_queries, x->ne[1]);
             x                         = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1);
             // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1.
             x            = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin);
-            auto direct  = ggml_ext_slice(gctx, x, 1, 0, prompt_length);
+            auto direct  = ggml_ext_slice(ctx, x, 1, 0, prompt_length);
             direct       = std::dynamic_pointer_cast<RMSNorm>(blocks["proj_directvlm.0"])->forward(ctx, direct);
             direct       = std::dynamic_pointer_cast<Linear>(blocks["proj_directvlm.1"])->forward(ctx, direct);
-            auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3);
-            queries      = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries);
+            auto queries = ggml_ext_slice(ctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3);
+            queries      = ggml_ext_slice(ctx, queries, 1, query_start, query_start + config.num_queries);
             queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, queries);
             queries      = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {});
             queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, queries);

src/model/te/ph3_image_connector.hpp

@@ -28,11 +28,11 @@ struct PH3ImageTextConnector : public UnaryBlock {
         x         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, x, params["input_mean"]), params["input_std"]);
         x         = proj->forward(ctx, x);
         x         = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, params["output_std"]), params["output_mean"]);
-        auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, 1));
+        auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx, x, 1, 0, 1));
         if (x->ne[1] == 1) {
             return sink;
         }
-        return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]), 1);
+        return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx, x, 1, 1, x->ne[1]), 1);
     }
 };
 

src/model/te/t5.hpp

@@ -249,7 +249,7 @@ public:
             }
         }
 
-        k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
+        k = ggml_ext_scale(ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
 
         x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
 

src/model/upscaler/esrgan.hpp

@@ -127,7 +127,7 @@ public:
         x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
         auto x5    = conv5->forward(ctx, x_cat);
 
-        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
+        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, x5, 0.2f), x);
         return x5;
     }
 };
@@ -152,7 +152,7 @@ public:
         out      = rdb2->forward(ctx, out);
         out      = rdb3->forward(ctx, out);
 
-        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
+        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, out, 0.2f), x);
         return out;
     }
 };

src/model/upscaler/flash_vsr.hpp

@@ -146,7 +146,7 @@ namespace RelayVSR {
             auto y                         = attention->forward(ctx, attention_norm->forward(ctx, x), qpe, kpe, local, name, condition_slot, start_slot, end_slot, history);
             x                              = ggml_add(c, x, y);
             y                              = feed_forward_gate_up_proj->forward(ctx, feed_forward_norm->forward(ctx, x));
-            auto parts                     = ggml_ext_chunk(c, y, 2, 0);
+            auto parts                     = ggml_ext_chunk(ctx, y, 2, 0);
             y                              = ggml_mul(c, ggml_silu(c, parts[0]), parts[1]);
             return ggml_add(c, x, feed_forward_down->forward(ctx, y));
         }

src/model/upscaler/h3_latent_upscaler.hpp

@@ -195,8 +195,8 @@ namespace H3VUpsampler {
             ggml_tensor* emb_silu = ggml_silu(ctx->ggml_ctx, emb);
             ggml_tensor* emb_out  = emb_layers->forward(ctx, emb_silu);
             ggml_tensor* emb_1d   = ggml_reshape_1d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, emb_out), 2 * channels);
-            ggml_tensor* scale    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels);
-            ggml_tensor* shift    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels);
+            ggml_tensor* scale    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels);
+            ggml_tensor* shift    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels);
             ggml_tensor* gain     = ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f);
 
             h = out_norm->forward(ctx, h);
@@ -236,7 +236,7 @@ namespace H3VUpsampler {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* w = params["weight"];
             if (ctx->weight_adapter) {
-                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             }
             ggml_tensor* w32 = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32);
             ggml_tensor* xa  = x;
@@ -250,7 +250,7 @@ namespace H3VUpsampler {
             // ggml_pad only pads the end; concatenate zeros for symmetric padding.
             ggml_tensor* xp = xa;
             if (pad > 0) {
-                ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, xa, 2, 0, 1), 0.0f);
+                ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx, xa, 2, 0, 1), 0.0f);
                 ggml_tensor* lead = zero;
                 for (int i = 1; i < pad; ++i) {
                     lead = ggml_concat(ctx->ggml_ctx, lead, zero, 2);
@@ -261,15 +261,15 @@ namespace H3VUpsampler {
 
             ggml_tensor* acc = nullptr;
             for (int j = 0; j < kernel; ++j) {
-                ggml_tensor* xs   = ggml_ext_slice(ctx->ggml_ctx, xp, 2, j, j + T);
-                ggml_tensor* wj   = ggml_ext_slice(ctx->ggml_ctx, w32, 2, j, j + 1);  // [1, 1, 1, C]
+                ggml_tensor* xs   = ggml_ext_slice(ctx, xp, 2, j, j + T);
+                ggml_tensor* wj   = ggml_ext_slice(ctx, w32, 2, j, j + 1);  // [1, 1, 1, C]
                 ggml_tensor* term = ggml_mul(ctx->ggml_ctx, xs, wj);
                 acc               = acc == nullptr ? term : ggml_add(ctx->ggml_ctx, acc, term);
             }
             if (has_bias) {
                 ggml_tensor* b = params["bias"];
                 if (ctx->weight_adapter) {
-                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
                 acc = ggml_add(ctx->ggml_ctx, acc, ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, 1, channels));
             }

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -158,15 +158,15 @@ namespace LTXVUpsampler {
             const int64_t W = x->ne[0];
             const int64_t H = x->ne[1];
             const int64_t T = x->ne[2];
-            x               = ggml_ext_cont(ctx->ggml_ctx, x);
+            x               = ggml_ext_cont(ctx, x);
             x               = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1);
             x               = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps);
 
             ggml_tensor* weight = params["weight"];
             ggml_tensor* bias   = params["bias"];
             if (ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
-                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias");
             }
             weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1);
             bias   = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1);
@@ -218,9 +218,9 @@ namespace LTXVUpsampler {
             int64_t w = x->ne[0];
             GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
             // x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*p1*p2]
-            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*p1*p2]
-            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));              // [b*f, h, w, c*p1*p2]
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);  // [b*f, h*w, c*p1*p2]
+            return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true);
         }
     };
 
@@ -241,9 +241,9 @@ namespace LTXVUpsampler {
             const int64_t C = x->ne[3] / upscale_factor;
 
             // x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
-            x = ggml_ext_cont(ctx->ggml_ctx, x);
+            x = ggml_ext_cont(ctx, x);
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
             return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
         }
     };
@@ -324,11 +324,11 @@ namespace LTXVUpsampler {
             auto blur_down     = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
 
             // rearrange(x, "b c f h w -> (b f) c h w")
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
             x = conv->forward(ctx, x);
             x = pixel_shuffle->forward(ctx, x);
             x = blur_down->forward(ctx, x);
-            return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            return ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
         }
     };
 
@@ -403,18 +403,18 @@ namespace LTXVUpsampler {
             } else if (config.temporal_upsample) {
                 auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
                 auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
-                x                  = upsample_conv->forward(ctx, x);                    // [b, c*2, f, h, w]
-                x                  = pixel_shuffle->forward(ctx, x);                    // [b, c, f*2, h, w]
-                x                  = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);  // x[:, :, 1:, :, :]
+                x                  = upsample_conv->forward(ctx, x);          // [b, c*2, f, h, w]
+                x                  = pixel_shuffle->forward(ctx, x);          // [b, c, f*2, h, w]
+                x                  = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);  // x[:, :, 1:, :, :]
             } else {
                 auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
                 auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
 
                 // rearrange(x, "b c f h w -> (b f) c h w"),
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
-                x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
-                x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
+                x = upsample_conv->forward(ctx, x);                                  // [b*f, c*4, h, w]
+                x = pixel_shuffle->forward(ctx, x);                                  // [b*f, c, h*2, w*2]
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
             }
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
 

src/model/upscaler/mosr.hpp

@@ -118,9 +118,9 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         // [N, C, H, W] -> [N, H, W, C] so the norm reduces over channels.
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-        x = ggml_ext_layer_norm(ctx->ggml_ctx, x, params["weight"], params["bias"], 1e-6f);
-        return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+        x = ggml_ext_layer_norm(ctx, x, params["weight"], params["bias"], 1e-6f);
+        return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
     }
 };
 
@@ -141,9 +141,8 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         const int pad = kernel_size / 2;
-        x             = ggml_conv_2d_dw_direct(ctx->ggml_ctx, params["weight"], ggml_ext_cont(ctx->ggml_ctx, x), 1, 1, pad, pad, 1, 1);
-        auto b        = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, channels, 1);
-        return ggml_add_inplace(ctx->ggml_ctx, x, b);
+        return ggml_ext_conv_2d_grouped(ctx, ggml_ext_cont(ctx, x), params["weight"], params["bias"], channels,
+                                        1, 1, pad, pad, 1, 1, true);
     }
 };
 
@@ -179,10 +178,10 @@ public:
         auto i = channel_view(hidden, hidden - conv_ch);
         auto c = conv->forward(ctx, channel_view(2 * hidden - conv_ch, conv_ch));
         if (hidden > conv_ch) {
-            c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, i), c, 2);
+            c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx, i), c, 2);
         }
 
-        h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, g)));
+        h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx, g)));
         h = mosr_mish(ctx->ggml_ctx, fc2->forward(ctx, h));
         return ggml_add(ctx->ggml_ctx, h, ggml_scale_bias(ctx->ggml_ctx, x, 1.f, -0.5f));
     }

src/model/vae/auto_encoder_kl.hpp

@@ -216,7 +216,7 @@ protected:
     }
 
     ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
-        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]);
         return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 
@@ -601,11 +601,11 @@ public:
             int64_t H = h * p;
             int64_t W = w * p;
 
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N);                           // [N, C, p*p, h*w]
-            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 1, 0, 2, 3));  // [N, C, h*w, p*p]
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N);                           // [N*C*h, w, p, p]
-            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));  // [N*C*h, p, w, p]
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N);                                   // [N, C, h*p, w*p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N);                 // [N, C, p*p, h*w]
+            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 1, 0, 2, 3));  // [N, C, h*w, p*p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N);                 // [N*C*h, w, p, p]
+            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 0, 2, 1, 3));  // [N*C*h, p, w, p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N);                         // [N, C, h*p, w*p]
         }
 
         if (use_quant) {
@@ -632,7 +632,7 @@ public:
             // sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z");
         }
         if (sd_version_uses_flux2_vae(version)) {
-            z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
+            z = ggml_ext_chunk(ctx, z, 2, 2)[0];
 
             // [N, C, H, W] -> [N, C*p*p, H/p, W/p]
             int64_t p = 2;

src/model/vae/hunyuan_vae.hpp

@@ -96,20 +96,20 @@ namespace Hunyuan {
                     auto previous = carry->take();
                     GGML_ASSERT(previous->ne[2] <= pad_t);
                     for (int64_t frame = 0; frame < previous->ne[2]; frame++) {
-                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1));
+                        temporal_frames.push_back(ggml_ext_slice(ctx, previous, 2, frame, frame + 1));
                     }
                     for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) {
-                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1));
+                        temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, 0, 1));
                     }
                 } else {
-                    auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                    auto first = ggml_ext_slice(ctx, x, 2, 0, 1);
                     for (int frame = 0; frame < pad_t; frame++) {
                         temporal_frames.push_back(first);
                     }
                 }
             }
             for (int64_t frame = 0; frame < x->ne[2]; frame++) {
-                temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1));
+                temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, frame, frame + 1));
             }
 
             if (pad_t > 0 && carry != nullptr && carry->output != nullptr) {
@@ -127,13 +127,13 @@ namespace Hunyuan {
             }
             auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) {
                 if (left > 0) {
-                    auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1);
+                    auto first = ggml_ext_slice(ctx, input, dim, 0, 1);
                     for (int i = 0; i < left; i++) {
                         input = ggml_concat(ctx->ggml_ctx, first, input, dim);
                     }
                 }
                 if (right > 0) {
-                    auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
+                    auto last = ggml_ext_slice(ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
                     for (int i = 0; i < right; i++) {
                         input = ggml_concat(ctx->ggml_ctx, input, last, dim);
                     }
@@ -184,19 +184,19 @@ namespace Hunyuan {
             auto k = k_proj->forward(ctx, x);  // [b*c, t, h, w]
             auto v = v_proj->forward(ctx, x);  // [b*c, t, h, w]
 
-            q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b);                                  // [b, c, t*h*w]
-            q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3));  // [b, t*h*w, c]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b);              // [b, c, t*h*w]
+            q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 1, 0, 2, 3));  // [b, t*h*w, c]
 
-            k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b);                                  // [b, c, t*h*w]
-            k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3));  // [b, t*h*w, c]
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b);              // [b, c, t*h*w]
+            k = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, k, 1, 0, 2, 3));  // [b, t*h*w, c]
 
-            v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);                                  // [b, c, t*h*w]
-            v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
+            v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);              // [b, c, t*h*w]
+            v = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
 
             x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [b, t*h*w, c]
 
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
-            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);                         // [b*c, t, h, w]
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);               // [b*c, t, h, w]
 
             x = proj_out->forward(ctx, x);
 
@@ -301,13 +301,13 @@ namespace Hunyuan {
             int64_t W      = x->ne[0];
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor);                          // [factor, C, T*H, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [C, factor, T*H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));            // [C, factor, T*H, W]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C);  // [C*factor_t*factor_s, factor_s, T*H, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [C*factor_t*factor_s, T*H, W, factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));            // [C*factor_t*factor_s, T*H, W, factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C);  // [C*factor_t, factor_s, T*H, W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C*factor_t, T*H, factor_s, W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C*factor_t, T*H, factor_s, W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C);  // [C, factor_t, T, H*factor_s*W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C, T, factor_t, H*factor_s*W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C, T, factor_t, H*factor_s*W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C);  // [C, T*factor_t, H*factor_s, W*factor_s]
 
             return x;
@@ -328,11 +328,11 @@ namespace Hunyuan {
 
             ggml_tensor* shortcut = nullptr;
             if (factor_t == 2 && !continuation) {
-                auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);                   // [B*factor*OC, 1, H, W]
-                h_first      = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B);             // [B*2*OC, 1, H*factor_s, W*factor_s]
-                h_first      = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels);  // [B*OC, 1, H*factor_s, W*factor_s]
+                auto h_first = ggml_ext_slice(ctx, h, 2, 0, 1);                   // [B*factor*OC, 1, H, W]
+                h_first      = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B);   // [B*2*OC, 1, H*factor_s, W*factor_s]
+                h_first      = ggml_ext_slice(ctx, h_first, 3, 0, out_channels);  // [B*OC, 1, H*factor_s, W*factor_s]
 
-                auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto x_first = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x_first      = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1);
                 x_first      = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B);
 
@@ -340,12 +340,12 @@ namespace Hunyuan {
                     return ggml_add(ctx->ggml_ctx, h_first, x_first);
                 }
 
-                auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);       // [B*factor*OC, T - 1, H, W]
+                auto h_next = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]);                 // [B*factor*OC, T - 1, H, W]
                 h_next      = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B);  // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s]
 
                 h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2);  // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
 
-                auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+                auto x_next = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
                 x_next      = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1);
                 x_next      = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B);
 
@@ -374,13 +374,13 @@ namespace Hunyuan {
         int64_t factor = factor_t * factor_s * factor_s;
 
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
         return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor);
     }
 
@@ -398,9 +398,9 @@ namespace Hunyuan {
         int64_t spatial = W * H * T;
         int64_t groups  = x->ne[3] / group_size;
         x               = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups);
-        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
         x               = ggml_sum_rows(ctx->ggml_ctx, x);
-        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
         x               = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups);
         return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast<float>(group_size));
     }
@@ -435,19 +435,19 @@ namespace Hunyuan {
             ggml_tensor* h_first = nullptr;
             ggml_tensor* x_first = nullptr;
             if (factor_t == 2) {
-                h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);
+                h_first = ggml_ext_slice(ctx, h, 2, 0, 1);
                 h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s);
                 h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3);
 
-                x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                x_first = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s);
                 x_first = mean_channel_groups(ctx, x_first, group_size / 2);
 
                 if (x->ne[2] == 1) {
                     return ggml_add(ctx->ggml_ctx, h_first, x_first);
                 }
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);
-                x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+                h = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]);
+                x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
             }
 
             GGML_ASSERT(h->ne[2] % factor_t == 0);
@@ -648,7 +648,7 @@ namespace Hunyuan {
             x             = ggml_silu_inplace(ctx->ggml_ctx, x);
             x             = conv_out->forward(ctx, x);
             x             = ggml_add(ctx->ggml_ctx, x, shortcut);
-            return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels);
+            return ggml_ext_slice(ctx, x, 3, 0, z_channels);
         }
     };
 
@@ -707,7 +707,7 @@ namespace Hunyuan {
             const int64_t frames = h->ne[2];
             for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) {
                 const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames);
-                auto chunk        = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end);
+                auto chunk        = ggml_ext_slice(ctx, h, 2, start, end);
 
                 std::vector<ggml_tensor*> carry_output;
                 TemporalConvCarry carry{

src/model/vae/ltx_audio_vae.hpp

@@ -220,7 +220,7 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
         if (left_pad > 0) {
-            x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
+            x = ggml_ext_pad_ext(runner_ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
         }
 
         auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@@ -228,8 +228,8 @@ namespace LTXV {
         GGML_ASSERT(frames->ne[1] == stft_channels);
         GGML_ASSERT(frames->ne[2] == channels * batch);
 
-        auto real      = ggml_ext_slice(ctx, frames, 1, 0, n_freqs);
-        auto imag      = ggml_ext_slice(ctx, frames, 1, n_freqs, stft_channels);
+        auto real      = ggml_ext_slice(runner_ctx, frames, 1, 0, n_freqs);
+        auto imag      = ggml_ext_slice(runner_ctx, frames, 1, n_freqs, stft_channels);
         auto magnitude = ggml_sqrt(ctx,
                                    ggml_add(ctx,
                                             ggml_sqr(ctx, real),
@@ -265,42 +265,19 @@ namespace LTXV {
         return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type;
     }
 
-    static ggml_tensor* repeat_with_vulkan_f32_workaround(ggml_backend_t backend,
-                                                          ggml_context* ctx,
-                                                          ggml_tensor* x,
-                                                          int64_t ne0,
-                                                          int64_t ne1,
-                                                          int64_t ne2,
-                                                          int64_t ne3) {
-        if (x->type != GGML_TYPE_F32 &&
-            (x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
-            sd_backend_is(backend, "vulkan")) {
-            auto x_f32    = ggml_cast(ctx, x, GGML_TYPE_F32);
-            auto repeated = ggml_repeat_4d(ctx,
-                                           x_f32,
-                                           ne0,
-                                           ne1,
-                                           ne2,
-                                           ne3);
-            return ggml_cast(ctx, repeated, x->type);
-        }
-        return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3);
-    }
-
     static ggml_tensor* repeat_1d_value(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t count) {
-        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(x->ne[0] == 1);
-        return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, x, count, x->ne[1], x->ne[2], x->ne[3]);
+        return ggml_ext_repeat_4d(runner_ctx, x, count, x->ne[1], x->ne[2], x->ne[3]);
     }
 
     static ggml_tensor* replicate_pad_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t left, int64_t right) {
         auto ctx = runner_ctx->ggml_ctx;
         if (left > 0) {
-            auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
+            auto first = ggml_ext_slice(runner_ctx, x, 0, 0, 1);
             x          = ggml_concat(ctx, repeat_1d_value(runner_ctx, first, left), x, 0);
         }
         if (right > 0) {
-            auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            auto last = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
             x         = ggml_concat(ctx, x, repeat_1d_value(runner_ctx, last, right), 0);
         }
         return x;
@@ -314,7 +291,7 @@ namespace LTXV {
         } else if (ggml_n_dims(base) == 1) {
             base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
         }
-        return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, base, base->ne[0], 1, channels, 1);
+        return ggml_ext_repeat_4d(runner_ctx, base, base->ne[0], 1, channels, 1);
     }
 
     static ggml_tensor* depthwise_conv1d(GGMLRunnerContext* runner_ctx,
@@ -329,7 +306,8 @@ namespace LTXV {
         return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1);
     }
 
-    static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) {
+    static ggml_tensor* reverse_1d_filter(GGMLRunnerContext* runner_ctx, ggml_tensor* filter) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(ctx != nullptr);
         GGML_ASSERT(filter != nullptr);
         GGML_ASSERT(filter->ne[1] == 1);
@@ -338,16 +316,17 @@ namespace LTXV {
 
         ggml_tensor* reversed = nullptr;
         for (int64_t k = filter->ne[0] - 1; k >= 0; --k) {
-            auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1);
+            auto slice = ggml_ext_slice(runner_ctx, filter, 0, k, k + 1);
             reversed   = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0);
         }
         return reversed;
     }
 
-    static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx,
+    static ggml_tensor* depthwise_conv_transpose1d(GGMLRunnerContext* runner_ctx,
                                                    ggml_tensor* x,
                                                    ggml_tensor* filter,
                                                    int stride) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1);
         GGML_ASSERT(filter->ne[1] == 1);
         GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1);
@@ -359,7 +338,7 @@ namespace LTXV {
 
         auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels);
         if (stride > 1) {
-            auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f);
+            auto zero_unit = ggml_ext_scale(runner_ctx, x_flat, 0.0f);
             auto zero_tail = zero_unit;
             for (int i = 1; i < stride - 1; ++i) {
                 zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0);
@@ -368,16 +347,16 @@ namespace LTXV {
         }
         x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels);
 
-        auto reversed_filter = reverse_1d_filter(ctx, filter);
+        auto reversed_filter = reverse_1d_filter(runner_ctx, filter);
         auto out             = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast<int>(kernel_size - 1), 1);
         if (out->ne[0] > out_time) {
-            out = ggml_ext_slice(ctx, out, 0, 0, out_time);
+            out = ggml_ext_slice(runner_ctx, out, 0, 0, out_time);
         }
         GGML_ASSERT(out->ne[0] == out_time);
         GGML_ASSERT(out->ne[1] == 1);
         GGML_ASSERT(out->ne[2] == channels);
 
-        out = ggml_ext_scale(ctx, out, static_cast<float>(stride));
+        out = ggml_ext_scale(runner_ctx, out, static_cast<float>(stride));
         return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1);
     }
 
@@ -409,21 +388,22 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, channels * batch, 1);
         x      = replicate_pad_1d(runner_ctx, x, pad, pad);
-        x      = depthwise_conv_transpose1d(ctx, x, filter, ratio);
-        x      = ggml_ext_slice(ctx, x, 0, pad_left, x->ne[0] - pad_right);
+        x      = depthwise_conv_transpose1d(runner_ctx, x, filter, ratio);
+        x      = ggml_ext_slice(runner_ctx, x, 0, pad_left, x->ne[0] - pad_right);
         return ggml_reshape_3d(ctx, x, x->ne[0], channels, batch);
     }
 
-    static ggml_tensor* crop_waveform_samples(ggml_context* ctx,
+    static ggml_tensor* crop_waveform_samples(GGMLRunnerContext* runner_ctx,
                                               ggml_tensor* waveform,
                                               int64_t target_samples) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(ctx != nullptr);
         GGML_ASSERT(waveform != nullptr);
         if (waveform->ne[0] == target_samples) {
             return waveform;
         }
         GGML_ASSERT(waveform->ne[0] > target_samples);
-        return ggml_ext_slice(ctx, waveform, 0, 0, target_samples);
+        return ggml_ext_slice(runner_ctx, waveform, 0, 0, target_samples);
     }
 
     struct PixelNorm2D : public UnaryBlock {
@@ -433,9 +413,9 @@ namespace LTXV {
             : eps(eps) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
             h      = ggml_rms_norm(ctx->ggml_ctx, h, eps);
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));
+            h      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 0, 3));
             return h;
         }
     };
@@ -456,8 +436,7 @@ namespace LTXV {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
             int pad_h = kernel_size.first - 1;
             int pad_w = kernel_size.second - 1;
-            x         = ggml_ext_pad_ext(ctx->ggml_ctx,
-                                         ctx->backend,
+            x         = ggml_ext_pad_ext(ctx,
                                          x,
                                          pad_w / 2,
                                          pad_w - pad_w / 2,
@@ -481,7 +460,7 @@ namespace LTXV {
             auto conv = std::dynamic_pointer_cast<HeightCausalConv2D>(blocks["conv"]);
             x         = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST);
             x         = conv->forward(ctx, x);
-            return ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]);
+            return ggml_ext_slice(ctx, x, 1, 1, x->ne[1]);
         }
     };
 
@@ -606,7 +585,7 @@ namespace LTXV {
             GGML_ASSERT(dilation == 1);
             x = ggml_conv_transpose_1d(ctx->ggml_ctx, params["weight"], x, stride, 0, dilation);
             if (padding > 0) {
-                x = ggml_ext_slice(ctx->ggml_ctx, x, 0, padding, x->ne[0] - padding);
+                x = ggml_ext_slice(ctx, x, 0, padding, x->ne[0] - padding);
             }
             if (bias) {
                 auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, params["bias"]->ne[0], 1, 1);
@@ -639,7 +618,7 @@ namespace LTXV {
             beta             = ggml_reshape_4d(ctx->ggml_ctx, beta, 1, beta->ne[0], 1, 1);
             auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
             oscillation      = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
-            auto eps_tensor  = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), eps);
+            auto eps_tensor  = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), eps);
             oscillation      = ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, beta, eps_tensor));
             return ggml_add(ctx->ggml_ctx, x, oscillation);
         }
@@ -675,8 +654,8 @@ namespace LTXV {
             int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2;
 
             x = replicate_pad_1d(ctx, x, up_pad, up_pad);
-            x = depthwise_conv_transpose1d(ctx->ggml_ctx, x, up_filter, up_ratio);
-            x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right);
+            x = depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio);
+            x = ggml_ext_slice(ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right);
 
             x = act->forward(ctx, x);
 
@@ -778,7 +757,7 @@ namespace LTXV {
 
         ggml_tensor* prepare_input(GGMLRunnerContext* ctx, ggml_tensor* mel) {
             mel       = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, mel, 1, 0, 2, 3));
-            auto mels = ggml_ext_chunk(ctx->ggml_ctx, mel, 2, 2);
+            auto mels = ggml_ext_chunk(ctx, mel, 2, 2);
             mel       = ggml_concat(ctx->ggml_ctx, mels[0], mels[1], 1);
             // mel = ggml_reshape_4d(ctx->ggml_ctx,
             //                       mel,
@@ -810,7 +789,7 @@ namespace LTXV {
                     auto block_out = resblock->forward(ctx, x);
                     sum            = sum == nullptr ? block_out : ggml_add(ctx->ggml_ctx, sum, block_out);
                 }
-                x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.0f / static_cast<float>(resblock_kernel_sizes.size()));
+                x = ggml_ext_scale(ctx, sum, 1.0f / static_cast<float>(resblock_kernel_sizes.size()));
             }
 
             x = act_post->forward(ctx, x);
@@ -876,8 +855,8 @@ namespace LTXV {
                                          int target_freq) {
             int64_t time = std::min<int64_t>(decoded->ne[1], target_time);
             int64_t freq = std::min<int64_t>(decoded->ne[0], target_freq);
-            decoded      = ggml_ext_slice(ctx->ggml_ctx, decoded, 0, 0, freq);
-            decoded      = ggml_ext_slice(ctx->ggml_ctx, decoded, 1, 0, time);
+            decoded      = ggml_ext_slice(ctx, decoded, 0, 0, freq);
+            decoded      = ggml_ext_slice(ctx, decoded, 1, 0, time);
             return decoded;
         }
 
@@ -995,7 +974,7 @@ namespace LTXV {
                                        ggml_add(ctx->ggml_ctx, residual, skip),
                                        -1.0f,
                                        1.0f);
-                waveform  = crop_waveform_samples(ctx->ggml_ctx, waveform, out_time);
+                waveform  = crop_waveform_samples(ctx, waveform, out_time);
             }
 
             return waveform;

src/model/vae/ltx_vae.hpp

@@ -32,28 +32,26 @@ namespace LTXVAE {
         return ggml_reshape_4d(ctx, x, 1, 1, 1, ggml_nelements(x));
     }
 
-    static inline std::pair<ggml_tensor*, ggml_tensor*> get_shift_scale(ggml_context* ctx,
-                                                                        ggml_tensor* table,
-                                                                        ggml_tensor* timestep,
-                                                                        int64_t channels,
-                                                                        int parts) {
+    static inline std::pair<ggml_tensor*, ggml_tensor*> get_shift_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* table, ggml_tensor* timestep, int64_t channels, int parts) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(timestep != nullptr);
         GGML_ASSERT(ggml_nelements(timestep) == channels * parts);
 
         auto timestep_view = ggml_reshape_2d(ctx, timestep, channels, parts);
         auto values        = ggml_add(ctx, table, timestep_view);
-        auto chunks        = ggml_ext_chunk(ctx, values, parts, 1, false);
+        auto chunks        = ggml_ext_chunk(runner_ctx, values, parts, 1, false);
         auto shift         = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[0]));
         auto scale         = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[1]));
         return {shift, scale};
     }
 
-    static inline ggml_tensor* depth_to_space_3d(ggml_context* ctx,
+    static inline ggml_tensor* depth_to_space_3d(GGMLRunnerContext* runner_ctx,
                                                  ggml_tensor* x,
                                                  int64_t c,
                                                  int factor_t,
                                                  int factor_s,
                                                  bool drop_first_temporal_frame) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [B*c*p1*p2*p3, T, H, W], B == 1, p2 == p3 == factor_s, p1 == factor_t
         // return: [B*c, T*p1, H*p2, W*p2]
         // Match: rearrange(x, "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)")
@@ -61,26 +59,24 @@ namespace LTXVAE {
         const int64_t H = x->ne[1];
         const int64_t W = x->ne[0];
 
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));        // [T, C, H, W]
-        x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T);  // [T*c*p1*p2, p3, H, W]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));        // [T*c*p1*p2, H, W, p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T);  // [T*c*p1, p2, H, W*p3]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));        // [T*c*p1, H, p2, W*p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T);  // [T, c, p1, H*p2*W*p3]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));        // [c, T, p1, H*p2*W*p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c);  // [T, c, T*p1, H*p2*W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2));  // [T, C, H, W]
+        x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T);          // [T*c*p1*p2, p3, H, W]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [T*c*p1*p2, H, W, p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T);          // [T*c*p1, p2, H, W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [T*c*p1, H, p2, W*p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T);          // [T, c, p1, H*p2*W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2));  // [c, T, p1, H*p2*W*p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c);          // [T, c, T*p1, H*p2*W*p3]
 
         if (drop_first_temporal_frame && factor_t > 1 && x->ne[2] > 0) {
-            x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
+            x = ggml_ext_slice(runner_ctx, x, 2, 1, x->ne[2]);
         }
 
         return x;
     }
 
-    static inline ggml_tensor* patchify(ggml_context* ctx,
-                                        ggml_tensor* x,
-                                        int patch_size) {
-        return WAN::WanVAE::patchify(ctx, x, patch_size, 1);
+    static inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int patch_size) {
+        return WAN::WanVAE::patchify(runner_ctx, x, patch_size, 1);
     }
 
     class CausalConv3d : public GGMLBlock {
@@ -113,20 +109,20 @@ namespace LTXVAE {
             auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
 
             if (causal) {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < time_kernel_size - 1; i++) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
                 }
                 x = ggml_concat(ctx->ggml_ctx, first_frame_pad, x, 2);
             } else {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < (time_kernel_size - 1) / 2; i++) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
                 }
 
-                auto last_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]);
+                auto last_frame     = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]);
                 auto last_frame_pad = last_frame;
                 for (int i = 1; i < (time_kernel_size - 1) / 2; i++) {
                     last_frame_pad = ggml_concat(ctx->ggml_ctx, last_frame_pad, last_frame, 2);
@@ -165,7 +161,7 @@ namespace LTXVAE {
                 GGML_ASSERT(start_idx >= 0);
                 GGML_ASSERT(end_idx > 0);
 
-                auto slice         = ggml_ext_slice(ctx->ggml_ctx, x, 2, start_idx, end_idx);
+                auto slice         = ggml_ext_slice(ctx, x, 2, start_idx, end_idx);
                 feat_map[feat_idx] = ggml_cont(ctx->ggml_ctx, slice);
             }
             feat_idx++;
@@ -175,7 +171,7 @@ namespace LTXVAE {
                 if (prev != nullptr) {
                     left_pad = prev;
                 } else {
-                    auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                    auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1);
                     left_pad         = first_frame;
                     for (int i = 1; i < pad; i++) {
                         left_pad = ggml_concat(ctx->ggml_ctx, left_pad, first_frame, 2);
@@ -185,7 +181,7 @@ namespace LTXVAE {
             }
 
             if (!causal && pad > 0) {
-                auto last_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]);
+                auto last_frame = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]);
                 auto right_pad  = last_frame;
                 for (int i = 1; i < pad; i++) {
                     right_pad = ggml_concat(ctx->ggml_ctx, right_pad, last_frame, 2);
@@ -204,9 +200,9 @@ namespace LTXVAE {
             : eps(eps) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2));
+            auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 3, 0, 1, 2));
             h      = ggml_rms_norm(ctx->ggml_ctx, h, eps);
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
+            h      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 3, 0));
             return h;
         }
     };
@@ -267,7 +263,7 @@ namespace LTXVAE {
                 auto values = ggml_add(ctx->ggml_ctx,
                                        params["scale_shift_table"],
                                        ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4));
-                auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false);
+                auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false);
                 shift1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0]));
                 scale1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1]));
                 shift2      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2]));
@@ -313,7 +309,7 @@ namespace LTXVAE {
                 auto values = ggml_add(ctx->ggml_ctx,
                                        params["scale_shift_table"],
                                        ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4));
-                auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false);
+                auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false);
                 shift1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0]));
                 scale1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1]));
                 shift2      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2]));
@@ -430,7 +426,7 @@ namespace LTXVAE {
 
             ggml_tensor* x_in = nullptr;
             if (residual) {
-                x_in       = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1);
+                x_in       = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1);
                 int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor;
                 auto res   = x_in;
                 for (int i = 1; i < repeat; i++) {
@@ -440,7 +436,7 @@ namespace LTXVAE {
             }
 
             x = conv->forward(ctx, x, causal);
-            x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1);
+            x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1);
             if (residual) {
                 x = ggml_add(ctx->ggml_ctx, x, x_in);
             }
@@ -460,7 +456,7 @@ namespace LTXVAE {
 
             ggml_tensor* x_in = nullptr;
             if (residual) {
-                x_in       = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first);
+                x_in       = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first);
                 int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor;
                 auto res   = x_in;
                 for (int i = 1; i < repeat; i++) {
@@ -470,7 +466,7 @@ namespace LTXVAE {
             }
 
             x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
-            x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
+            x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
             if (residual) {
                 x = ggml_add(ctx->ggml_ctx, x, x_in);
             }
@@ -515,7 +511,7 @@ namespace LTXVAE {
             auto conv_downsample = std::dynamic_pointer_cast<WAN::AvgDown3D>(blocks["conv_downsample"]);
 
             if (factor_t > 1 && x->ne[2] > 0) {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < factor_t - 1; ++i) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
@@ -860,7 +856,7 @@ namespace LTXVAE {
             x = ggml_silu_inplace(ctx->ggml_ctx, x);
             x = conv_out->forward(ctx, x, true);
 
-            auto last_channel = ggml_ext_slice(ctx->ggml_ctx, x, 3, x->ne[3] - 1, x->ne[3]);
+            auto last_channel = ggml_ext_slice(ctx, x, 3, x->ne[3] - 1, x->ne[3]);
             auto repeat_shape = ggml_new_tensor_4d(ctx->ggml_ctx, last_channel->type, last_channel->ne[0], last_channel->ne[1], last_channel->ne[2], latent_channels - 1);
             auto repeated     = ggml_repeat(ctx->ggml_ctx, last_channel, repeat_shape);
             return ggml_concat(ctx->ggml_ctx, x, repeated, 3);
@@ -959,7 +955,7 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning) {
-                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]);
                 scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
@@ -981,7 +977,7 @@ namespace LTXVAE {
             if (timestep_conditioning) {
                 auto last_time_embedder = std::dynamic_pointer_cast<PixArtAlphaCombinedTimestepSizeEmbeddings>(blocks["last_time_embedder"]);
                 auto timestep_embed     = last_time_embedder->forward(ctx, scaled_timestep);
-                auto [shift, scale]     = get_shift_scale(ctx->ggml_ctx,
+                auto [shift, scale]     = get_shift_scale(ctx,
                                                           params["last_scale_shift_table"],
                                                           timestep_embed,
                                                           hidden_channels,
@@ -1010,7 +1006,7 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning && timestep != nullptr) {
-                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]);
                 scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
@@ -1038,7 +1034,7 @@ namespace LTXVAE {
             if (timestep_conditioning) {
                 auto last_time_embedder = std::dynamic_pointer_cast<PixArtAlphaCombinedTimestepSizeEmbeddings>(blocks["last_time_embedder"]);
                 auto timestep_embed     = last_time_embedder->forward(ctx, scaled_timestep);
-                auto [shift, scale]     = get_shift_scale(ctx->ggml_ctx,
+                auto [shift, scale]     = get_shift_scale(ctx,
                                                           params["last_scale_shift_table"],
                                                           timestep_embed,
                                                           hidden_channels,
@@ -1092,7 +1088,7 @@ namespace LTXVAE {
             auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
             auto latents   = processor->un_normalize(ctx, z);
             auto out       = decoder->forward(ctx, latents, timestep);
-            out            = WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+            out            = WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
             return out;
         }
 
@@ -1112,7 +1108,7 @@ namespace LTXVAE {
             const int64_t T = z->ne[2];
             if (T <= 1) {
                 auto out = decoder->forward(ctx, latents, timestep);
-                return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+                return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
             }
 
             // feat_map holds ggml_tensor* nodes (contiguous copies at each conv layer).
@@ -1145,20 +1141,20 @@ namespace LTXVAE {
 
                 int chunk_overlap = overlap;  // modified by forward_tiled_frame temporal inflation
 
-                auto z_chunk = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, end_i);
+                auto z_chunk = ggml_ext_slice(ctx, latents, 2, i, end_i);
 
                 auto out_chunk = decoder->forward_tiled_frame(ctx, z_chunk, timestep,
                                                               feat_map, feat_idx, i, chunk_overlap);
 
                 // discard overlap frames if it's not the final chunk
                 if (overlap > 0 && end_i < (int)T) {
-                    out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+                    out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
                 }
 
                 out = (out == nullptr) ? out_chunk : ggml_concat(ctx->ggml_ctx, out, out_chunk, 2);
             }
 
-            return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+            return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
         }
 
         ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
@@ -1177,9 +1173,9 @@ namespace LTXVAE {
             auto out_chunk    = decoder->forward_tiled_frame(ctx, latents, timestep,
                                                              feat_map, feat_idx, chunk_idx, chunk_overlap);
             if (chunk_overlap > 0) {
-                out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+                out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
             }
-            return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out_chunk, patch_size, 1);
+            return WAN::WanVAE::unpatchify(ctx, out_chunk, patch_size, 1);
         }
 
         ggml_tensor* encode(GGMLRunnerContext* ctx,
@@ -1188,9 +1184,9 @@ namespace LTXVAE {
             auto encoder   = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
             auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
 
-            x         = patchify(ctx->ggml_ctx, x, patch_size);
+            x         = patchify(ctx, x, patch_size);
             auto out  = encoder->forward(ctx, x);
-            auto mean = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3, false)[0];
+            auto mean = ggml_ext_chunk(ctx, out, 2, 3, false)[0];
             mean      = ggml_cont(ctx->ggml_ctx, mean);
             return processor->normalize(ctx, mean);
         }

src/model/vae/mage_vae.hpp

@@ -32,23 +32,24 @@ namespace MageVAE {
             ggml_tensor* weight = affine ? params["weight"] : nullptr;
             ggml_tensor* bias   = affine ? params["bias"] : nullptr;
             if (affine && ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
-                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias");
             }
             // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-            x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
-            return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+            x = ggml_ext_layer_norm(ctx, x, weight, bias, 1e-6f);
+            return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
         }
     };
 
-    inline ggml_tensor* modulate_2d(ggml_context* ctx,
+    inline ggml_tensor* modulate_2d(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     ggml_tensor* shift,
                                     ggml_tensor* scale) {
-        shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
-        scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
-        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
+        auto ctx = runner_ctx->ggml_ctx;
+        shift    = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
+        scale    = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
+        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(runner_ctx, 1, 1, 1, 1))), shift);
     }
 
     inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
@@ -70,7 +71,7 @@ namespace MageVAE {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
             auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto x        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
+            auto x        = ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f);
             x             = linear_0->forward(ctx, x);
             x             = ggml_silu_inplace(ctx->ggml_ctx, x);
             return linear_2->forward(ctx, x);
@@ -139,9 +140,9 @@ namespace MageVAE {
             auto ada   = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
             auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
+            auto chunks = ggml_ext_chunk(ctx, params, 6, 0);
             auto x      = norm1->forward(ctx, input);
-            x           = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
+            x           = modulate_2d(ctx, x, chunks[0], chunks[1]);
             x           = conv1->forward(ctx, x);
             x           = conv2->forward(ctx, x);
             x           = ggml_gelu(ctx->ggml_ctx, x);
@@ -151,7 +152,7 @@ namespace MageVAE {
             x           = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
 
             auto h      = norm2->forward(ctx, x);
-            h           = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
+            h           = modulate_2d(ctx, h, chunks[3], chunks[4]);
             h           = conv4->forward(ctx, h);
             h           = ggml_gelu(ctx->ggml_ctx, h);
             h           = conv5->forward(ctx, h);
@@ -179,17 +180,18 @@ namespace MageVAE {
         }
     };
 
-    inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
+    inline ggml_tensor* replicate_pad_right_bottom(GGMLRunnerContext* runner_ctx,
                                                    ggml_tensor* x,
                                                    int pad_w,
                                                    int pad_h) {
+        auto ctx = runner_ctx->ggml_ctx;
         if (pad_w > 0) {
-            auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            auto edge = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
             edge      = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
             x         = ggml_concat(ctx, x, edge, 0);
         }
         if (pad_h > 0) {
-            auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+            auto edge = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]);
             edge      = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
             x         = ggml_concat(ctx, x, edge, 1);
         }
@@ -209,25 +211,27 @@ namespace MageVAE {
             blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
         }
 
-        ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
-            x = DiT::patchify(ctx, x, patch_size, patch_size);
-            x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
+        ggml_tensor* to_patches(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+            auto ctx = runner_ctx->ggml_ctx;
+            x        = DiT::patchify(runner_ctx, x, patch_size, patch_size);
+            x        = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
             // [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
-            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
             return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
         }
 
-        ggml_tensor* from_patches(ggml_context* ctx,
+        ggml_tensor* from_patches(GGMLRunnerContext* runner_ctx,
                                   ggml_tensor* x,
                                   int64_t patch_count,
                                   int64_t batch_size,
                                   int64_t h_patches,
                                   int64_t w_patches) {
-            x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
+            auto ctx = runner_ctx->ggml_ctx;
+            x        = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
             // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
-            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
-            return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
+            return DiT::unpatchify(runner_ctx, x, h_patches, w_patches, patch_size, patch_size);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
@@ -247,19 +251,19 @@ namespace MageVAE {
             int64_t np     = wp * hp;
 
             auto h = norm->forward(ctx, input);
-            auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
-            auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
-            auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
-            q      = to_patches(ctx->ggml_ctx, q);
-            k      = to_patches(ctx->ggml_ctx, k);
-            v      = to_patches(ctx->ggml_ctx, v);
+            auto q = replicate_pad_right_bottom(ctx, q_proj->forward(ctx, h), pad_w, pad_h);
+            auto k = replicate_pad_right_bottom(ctx, k_proj->forward(ctx, h), pad_w, pad_h);
+            auto v = replicate_pad_right_bottom(ctx, v_proj->forward(ctx, h), pad_w, pad_h);
+            q      = to_patches(ctx, q);
+            k      = to_patches(ctx, k);
+            v      = to_patches(ctx, v);
             h      = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
-            h      = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
+            h      = from_patches(ctx, h, np, batch, hp, wp);
             if (pad_h > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
+                h = ggml_ext_slice(ctx, h, 1, 0, height);
             }
             if (pad_w > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
+                h = ggml_ext_slice(ctx, h, 0, 0, width);
             }
             return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
         }
@@ -311,11 +315,11 @@ namespace MageVAE {
                 cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
             }
             cond   = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
-            auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
+            auto z = ggml_ext_zeros(ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
             z      = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
             z      = ggml_concat(ctx->ggml_ctx, cond, z, 2);
             z      = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
-            auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
+            auto t = ggml_ext_zeros(ctx, image->ne[3], 1, 1, 1);
             t      = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
             auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
             for (int i = 0; i < 21; ++i) {
@@ -336,9 +340,9 @@ namespace MageVAE {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
             auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
+            auto chunks = ggml_ext_chunk(ctx, params, 3, 0);
             auto h      = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
-            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
+            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx, 1, 1, 1, 1))), chunks[0]);
             h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
             h           = ggml_silu(ctx->ggml_ctx, h);
             h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
@@ -373,8 +377,8 @@ namespace MageVAE {
             int64_t n      = cond->ne[3];
             int64_t length = w * h;
 
-            auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
-            auto t     = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
+            auto image = ggml_ext_zeros(ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
+            auto t     = ggml_ext_zeros(ctx, n, 1, 1, 1);
             t          = ggml_reshape_1d(ctx->ggml_ctx, t, n);
             auto c     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
 
@@ -385,15 +389,15 @@ namespace MageVAE {
                 s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
             }
             // [N, C, H, W] -> [N*H*W, C].
-            s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
+            s = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
             s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
 
             auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
             // Split 32*P channels as [32, P], then produce [N*L, P, 32].
             y          = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
-            y          = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
+            y          = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
             y          = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
-            auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            auto zeros = ggml_ext_zeros(ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
             dct        = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
             auto x     = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
             x          = ggml_concat(ctx->ggml_ctx, x, dct, 0);
@@ -409,9 +413,9 @@ namespace MageVAE {
             x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
             // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
             x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
-            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
+            return DiT::unpatchify(ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
         }
     };
 

src/model/vae/minimax_h3_audio_vae.hpp

@@ -28,7 +28,7 @@ namespace MiniMaxH3 {
             auto alpha       = params["alpha"];
             auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
             oscillation      = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
-            auto eps         = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f);
+            auto eps         = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), 1e-9f);
             return ggml_add(ctx->ggml_ctx,
                             x,
                             ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps)));
@@ -59,7 +59,7 @@ namespace MiniMaxH3 {
             auto h     = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x))));
             if (x->ne[0] != h->ne[0]) {
                 int64_t pad = (x->ne[0] - h->ne[0]) / 2;
-                x           = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad);
+                x           = ggml_ext_slice(ctx, x, 0, pad, x->ne[0] - pad);
             }
             return ggml_add(ctx->ggml_ctx, x, h);
         }
@@ -162,7 +162,7 @@ namespace MiniMaxH3 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto qkv_layer  = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
             auto proj       = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
-            auto qkv        = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0);
+            auto qkv        = ggml_ext_chunk(ctx, qkv_layer->forward(ctx, x), 3, 0);
             auto bias_shape = [&](ggml_tensor* bias) {
                 return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
             };
@@ -172,7 +172,7 @@ namespace MiniMaxH3 {
 
             int64_t sequence = x->ne[1];
             auto mask        = ggml_diag_mask_inf(ctx->ggml_ctx,
-                                                  ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1),
+                                                  ggml_ext_zeros(ctx, sequence, sequence, 1, 1),
                                                   0);
             auto attn_out    = ggml_ext_attention_ext(ctx,
                                                       q,
@@ -323,7 +323,7 @@ namespace MiniMaxH3 {
                     auto value = block->forward(ctx, x);
                     sum        = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value);
                 }
-                x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels);
+                x = ggml_ext_scale(ctx, sum, 1.f / num_kernels);
             }
             auto activation = std::dynamic_pointer_cast<LTXV::Activation1D>(blocks["activation_post"]);
             auto conv_post  = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_post"]);
@@ -412,7 +412,7 @@ namespace MiniMaxH3 {
                                                     streams);
             ggml_tensor* waveform = nullptr;
             for (int64_t stream = 0; stream < streams; ++stream) {
-                auto stream_latent   = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1);
+                auto stream_latent   = ggml_ext_slice(ctx, latent, 2, stream, stream + 1);
                 auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent));
                 waveform             = waveform == nullptr
                                            ? stream_waveform

src/model/vae/minimax_h3_sol_attention.hpp

@@ -0,0 +1,206 @@
+#ifndef __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
+#define __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
+
+#include <algorithm>
+#include <array>
+#include <map>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/tensor.hpp"
+
+namespace MiniMaxH3VAE {
+
+    constexpr int64_t sol_window_size    = 16;
+    constexpr int64_t sol_block_size     = 5;
+    constexpr int64_t sol_global_rows    = 5;
+    constexpr int64_t sol_spatial_shift  = sol_window_size / 2;
+    constexpr int64_t sol_temporal_shift = 2;
+
+    struct SolAttentionPlan {
+        template <typename Tensor>
+        struct Batch {
+            Tensor output_rows;
+            Tensor queries;
+            Tensor keys;
+            Tensor previous_queries;
+            Tensor previous_keys;
+        };
+
+        template <typename Tensor>
+        struct Block {
+            Tensor rows;
+            std::vector<Batch<Tensor>> batches;
+        };
+
+        std::vector<Block<sd::Tensor<int32_t>>> blocks;
+
+        SolAttentionPlan(int64_t width, int64_t height, int64_t frames, bool shifted) {
+            const int64_t frame_size = width * height;
+            const int64_t pixels     = frames * frame_size;
+            const int64_t block_rows = sol_block_size * frame_size;
+            const int64_t count      = (frames + sol_block_size - 1) / sol_block_size;
+            const int64_t shift      = shifted ? sol_spatial_shift : 0;
+            const int64_t time_shift = shifted ? sol_temporal_shift : 0;
+
+            auto pack = [](std::vector<int32_t> rows) {
+                const auto size = static_cast<int64_t>(rows.size());
+                return sd::Tensor<int32_t>({size}, std::move(rows));
+            };
+            for (int64_t block = 0; block < count; ++block) {
+                std::vector<int32_t> rows;
+                const int64_t start = block * block_rows;
+                const int64_t end   = std::min(pixels, start + block_rows);
+                for (int64_t i = start; i < end; ++i) {
+                    rows.push_back(static_cast<int32_t>(i));
+                }
+                std::vector<int32_t> globals;
+                std::vector<int32_t> queries;
+                for (int64_t i = 0; i < sol_global_rows; ++i) {
+                    globals.push_back(static_cast<int32_t>(pixels + block * sol_global_rows + i));
+                    queries.push_back(static_cast<int32_t>(rows.size()));
+                    rows.push_back(globals.back());
+                }
+                blocks.push_back({pack(std::move(rows)), {}});
+                blocks.back().batches.push_back({pack(std::move(globals)), pack(std::move(queries)), {}, {}, {}});
+            }
+
+            struct Window {
+                std::vector<int32_t> output_rows;
+                std::vector<int32_t> queries;
+                std::vector<int32_t> keys;
+                std::vector<int32_t> previous_queries;
+                std::vector<int32_t> previous_keys;
+            };
+            for (int64_t t = -time_shift; t < frames; t += sol_block_size) {
+                const int64_t first       = std::max<int64_t>(0, t);
+                const int64_t end         = std::min(frames, t + sol_block_size);
+                const int64_t first_block = first / sol_block_size;
+                const int64_t last_block  = (end - 1) / sol_block_size;
+                std::map<int64_t, std::vector<Window>> groups;
+                for (int64_t y = -shift; y < height; y += sol_window_size) {
+                    for (int64_t x = -shift; x < width; x += sol_window_size) {
+                        Window window;
+                        for (int64_t i = first; i < end; ++i) {
+                            for (int64_t j = std::max<int64_t>(0, y); j < std::min(height, y + sol_window_size); ++j) {
+                                for (int64_t k = std::max<int64_t>(0, x); k < std::min(width, x + sol_window_size); ++k) {
+                                    const int64_t row = (i * height + j) * width + k;
+                                    window.output_rows.push_back(static_cast<int32_t>(row));
+                                    const int64_t block = i / sol_block_size;
+                                    auto& queries       = block == last_block ? window.queries : window.previous_queries;
+                                    queries.push_back(static_cast<int32_t>(row - block * block_rows));
+                                }
+                            }
+                        }
+                        window.keys          = window.queries;
+                        window.previous_keys = window.previous_queries;
+                        for (int64_t block = first_block; block <= last_block; ++block) {
+                            const int64_t start = std::min(block_rows, pixels - block * block_rows);
+                            auto& keys          = block == last_block ? window.keys : window.previous_keys;
+                            for (int64_t i = 0; i < sol_global_rows; ++i) {
+                                keys.push_back(static_cast<int32_t>(start + i));
+                            }
+                        }
+                        groups[window.output_rows.size()].push_back(std::move(window));
+                    }
+                }
+                // Bound the workspace of backends that fall back to unfused attention.
+                constexpr size_t max_windows = 4;
+                for (const auto& group : groups) {
+                    const auto& windows = group.second;
+                    for (size_t first_window = 0; first_window < windows.size(); first_window += max_windows) {
+                        const auto n      = static_cast<int64_t>(std::min(max_windows, windows.size() - first_window));
+                        auto pack_windows = [&](auto member) -> sd::Tensor<int32_t> {
+                            const auto rows = static_cast<int64_t>((windows[first_window].*member).size());
+                            if (rows == 0) {
+                                return {};
+                            }
+                            sd::Tensor<int32_t> indices({rows, n});
+                            for (int64_t i = 0; i < n; ++i) {
+                                const auto& source = windows[first_window + i].*member;
+                                std::copy(source.begin(), source.end(), indices.data() + i * rows);
+                            }
+                            return indices;
+                        };
+                        blocks[last_block].batches.push_back({pack_windows(&Window::output_rows),
+                                                              pack_windows(&Window::queries),
+                                                              pack_windows(&Window::keys),
+                                                              pack_windows(&Window::previous_queries),
+                                                              pack_windows(&Window::previous_keys)});
+                    }
+                }
+            }
+        }
+
+        struct Inputs {
+            std::vector<Block<ggml_tensor*>> blocks;
+            ggml_tensor* row_ids = nullptr;
+
+            template <typename Project>
+            ggml_tensor* forward(GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* pe,
+                                 int64_t heads,
+                                 Project project) const {
+                auto ctx          = runner_ctx->ggml_ctx;
+                const int64_t dim = x->ne[0];
+                auto output       = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, x->ne[1]);
+                auto positions    = ggml_reshape_2d(ctx, pe, pe->ne[0] * pe->ne[1] * pe->ne[2], pe->ne[3]);
+                auto gather       = [&](ggml_tensor* rows, ggml_tensor* indices) {
+                    auto flat = ggml_reshape_1d(ctx, indices, ggml_nelements(indices));
+                    return ggml_reshape_3d(ctx, ggml_get_rows(ctx, rows, flat), dim, indices->ne[0], indices->ne[1]);
+                };
+                std::array<ggml_tensor*, 3> previous{};
+                for (const auto& block : blocks) {
+                    auto local_pe = ggml_reshape_4d(ctx, ggml_get_rows(ctx, positions, block.rows),
+                                                    pe->ne[0], pe->ne[1], pe->ne[2], block.rows->ne[0]);
+                    auto current  = project(ggml_get_rows(ctx, x, block.rows), local_pe);
+                    for (const auto& batch : block.batches) {
+                        // Gather first so a crossing window never copies either full block.
+                        auto window = [&](size_t i, ggml_tensor* now, ggml_tensor* before) {
+                            // Global queries read the whole native block in its existing order.
+                            auto rows = now == nullptr ? current[i] : gather(current[i], now);
+                            if (before != nullptr) {
+                                rows = ggml_concat(ctx, gather(previous[i], before), rows, 1);
+                            }
+                            return rows;
+                        };
+                        auto attended = ggml_ext_attention_ext(runner_ctx,
+                                                               window(0, batch.queries, batch.previous_queries),
+                                                               window(1, batch.keys, batch.previous_keys),
+                                                               window(2, batch.keys, batch.previous_keys),
+                                                               heads, nullptr, false, runner_ctx->flash_attn_enabled);
+                        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);
+                        runner_ctx->expand_graph(output);
+                    }
+                    // Later windows can only read this block and its immediate successor.
+                    previous = current;
+                }
+                return output;
+            }
+        };
+
+        template <typename BindInput>
+        Inputs bind(BindInput bind_input) const {
+            Inputs inputs;
+            for (const auto& block : blocks) {
+                Block<ggml_tensor*> bound{bind_input(block.rows), {}};
+                for (const auto& batch : block.batches) {
+                    auto bind_optional = [&](const auto& indices) {
+                        return indices.empty() ? nullptr : bind_input(indices);
+                    };
+                    bound.batches.push_back({bind_input(batch.output_rows), bind_input(batch.queries),
+                                             bind_optional(batch.keys), bind_optional(batch.previous_queries),
+                                             bind_optional(batch.previous_keys)});
+                }
+                inputs.blocks.push_back(std::move(bound));
+            }
+            return inputs;
+        }
+    };
+
+}  // namespace MiniMaxH3VAE
+
+#endif  // __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__