refactor: unify ggml operation helpers around runner context (#2314)

Co-authored-by: Maxime Dréan <maximedrean@proton.me>

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()
@@ -111,6 +126,12 @@ if(SD_CUDA)
 endif()
 
 if(SD_METAL)
+    if(NOT SD_USE_UPSTREAM_GGML)
+        if(CMAKE_VERSION VERSION_LESS 3.16)
+            message(FATAL_ERROR "SD_METAL with SD_USE_UPSTREAM_GGML=OFF requires CMake 3.16 or later for OBJC and OBJCXX support")
+        endif()
+        enable_language(OBJC OBJCXX)
+    endif()
     message("-- Use Metal as backend stable-diffusion")
     set(GGML_METAL ON)
 endif()
@@ -364,6 +385,11 @@ if (SD_CUDA)
     endif()
 endif()
 set(SD_BACKEND_REGISTRATIONS)
+if (SD_METAL AND NOT SD_USE_UPSTREAM_GGML)
+    add_subdirectory(src/core/metal)
+    target_link_libraries(${SD_LIB} PRIVATE "-framework Foundation" "-framework Metal")
+    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_metal_i8)
+endif()
 if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
     if(SD_HIPBLAS)
         include(cmake/hip.cmake)

CONTRIBUTING.md

@@ -29,6 +29,18 @@ PRs should include:
 - Linked issue/PR context when applicable.
 - Screenshots or sample outputs for UI/visual behavior changes.
 
+## Documentation
+
+Write model usage guides in `docs/` for users. Focus on supported weights,
+setup, runnable commands, parameter meanings, output behavior, and practical
+limits such as memory requirements. Include technical details only when they
+help users choose settings or troubleshoot a problem.
+
+Keep tensor layouts, internal class names, kernel implementation details,
+algorithm derivations, and development history out of usage guides. Keep task-specific
+profiling logs, numerical comparison procedures, and test artifacts local, and
+summarize relevant verification results in the PR description.
+
 ## README UI Listings
 
 The [UIs list in the README](README.md#uis) only includes open-source projects that use `stable-diffusion.cpp` as a backend. Do not submit PRs to add closed-source UIs.

README.md

@@ -88,6 +88,9 @@ API and command-line option may change frequently.***
     - [LTX-2.3/LTX-2.5](./docs/ltx2.md)
     - [HunyuanVideo 1.5](./docs/hunyuan_video.md)
     - [LingBot-Video](./docs/lingbot_video.md)
+  - Video Upscale Models
+    - [RelayVSR](./docs/relay_vsr.md)
+    - [SeedVR2](./docs/seedvr2.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
   - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
   - Control Net support with SD 1.5
@@ -96,7 +99,6 @@ API and command-line option may change frequently.***
   - Latent Consistency Models support (LCM/LCM-LoRA)
   - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
   - Upscale images generated with [ESRGAN](./docs/esrgan.md)
-  - Upscale videos with [RelayVSR](./docs/relay_vsr.md) (experimental, 2x and 4x)
 - Supported backends
   - CPU (AVX, AVX2 and AVX512 support for x86 architectures)
   - CUDA
@@ -199,6 +201,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
 

cmake/backend_extensions.h.in

@@ -3,6 +3,8 @@
 
 #include "core/cuda/registration.h"
 
+void sd_ggml_register_metal_i8(ggml_backend_dev_t device);
+
 void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
 void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device);
 void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device);

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
@@ -130,6 +130,9 @@ Services are discovered through `ggml_backend_reg_get_proc_address`:
 | `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. |
 | `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. |
 | `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. |
+| `ggml_ext_metal_device_v1` | Borrow the Metal device for application pipeline compilation. |
+| `ggml_ext_metal_buffer_v1` | Resolve a backend buffer range to its Metal buffer and offset, including private storage and split allocations. |
+| `ggml_ext_metal_submit_v1` | Encode application commands on a backend-owned command buffer, ordered with native Metal work and tracked by backend synchronization. |
 
 The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain
 or destroy it, or change its stream, pointer mode, math mode, or workspace. This
@@ -156,11 +159,34 @@ 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.
 
+`src/core/metal/int8.mm` registers Metal INT8 linear and ConvRot kernels. Their
+Metal source is embedded at build time and compiled for the selected device
+during registration. The public `ggml-ext-metal.h` services expose borrowed
+objects without private backend headers. The submit callback only encodes work;
+GGML owns submission, cancellation checks, command-buffer lifetime, and
+asynchronous error tracking. Rotation scratch and any overlapping output use
+the scheduler workspace. No full floating-point weight copy is allocated.
+When runtime compilation supports Metal 4 TensorOps, INT8 multiplication uses
+`matmul2d` with INT32 cooperative results and fused scale/bias restoration.
+The SIMD-group implementation remains the fallback for older systems and
+devices where the TensorOps pipeline is unavailable.
+TensorOps multiplies at most 1024 pairs per partial result and combines those
+results with wrapping integer addition. This preserves exact INT32 semantics
+even when TensorOps uses floating-point arithmetic internally.
+
 `ggml_ext_modulate(ctx, x, shift, scale)` emits `sd.modulate.v1` for F32 inputs
 when the runner context uses CPU, CUDA/HIP/MUSA, or Vulkan.
 Its explicit inputs are `x`, `scale`, and optional `shift`; it computes
@@ -243,13 +269,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
@@ -258,8 +284,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
@@ -472,8 +498,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.
 
@@ -491,7 +517,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

@@ -93,6 +93,7 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 ## Backend support
 
 - 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.
 - 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.

docs/minimax_h3.md

@@ -169,3 +169,18 @@ Notes for this mode:
   block's working memory must also fit within the VRAM budget.
 - The default video flow shift is 12. The audio stream is mapped internally to
   its shift of 3, so the regular samplers can operate on the packed AV latent.
+
+## HyperFlow H3 Flash
+
+[HyperFlow H3 Flash](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash)
+uses the shared components and input options described in this guide. Replace
+`--diffusion-model` with the matching HyperFlow FL2VA or Ref2VA checkpoint and
+use these sampling options:
+
+```sh
+--sampling-method euler --cfg-scale 1 --sigmas hyperflow-h3-flash-6
+```
+
+Use `hyperflow-h3-flash-4` or `hyperflow-h3-flash-8` for 4 or 8 steps. These
+presets also work with `--hires-sigmas`; their length determines the step count.
+They already include video flow shift 12, so keep the default flow shift.

docs/relay_vsr.md

@@ -1,24 +1,24 @@
-# RelayVSR video super-resolution
+# RelayVSR image and video super-resolution
 
-RelayVSR uses a Wan model to sample sparse reference latents, followed by the
-Flash conditional decoder to reconstruct every input frame. The model code
-follows [RelayVSR](https://github.com/kopperx/RelayVSR); the loader maps
-[ComfyUI repackaged weights](https://huggingface.co/Comfy-Org/RelayVSR) to the
-original module names.
+[RelayVSR](https://github.com/kopperx/RelayVSR) upscales images and videos by
+2x or 4x with `-M vid_upscale`.
 
-Use `-M vid_upscale` with an ordered directory of RGB images. Filenames are
-sorted lexicographically. Supply the source frame rate with `--fps`. Output
-keeps the input frame count. Use a video output supported by the CLI, such as
-`.avi` or `.webm`, or a numbered PNG output such as `output/%04d.png`.
+## Weights
 
-The scale is an explicit, required argument. It is not read from checkpoint
-metadata. Use matching Wan and Flash checkpoints:
+[ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/RelayVSR) load
+directly. Select the matching pair for the requested scale:
 
-| Scale | `--diffusion-model` | `--vae` |
+| `--upscale-factor` | `--diffusion-model` | `--vae` |
 | --- | --- | --- |
 | 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` |
 | 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` |
 
+No separate text encoder or ordinary Wan VAE is needed.
+
+## Usage
+
+For a video, extract numbered frames into a directory, then run:
+
 ```sh
 sd-cli -M vid_upscale \
   --diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \
@@ -27,44 +27,24 @@ sd-cli -M vid_upscale \
   --vsr-window-size 16 -s 42 -o output.avi
 ```
 
-`--vsr-window-size` includes both reference endpoints. The default 16 samples
-references at frame indices 0, 15, 30, and so on, always including the last
-frame. Shared endpoints update Flash history in the following window but are
-emitted only once. Single-frame input is also supported.
+Frames are read in filename order. Set `--fps` to the source video's frame rate.
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o output/%04d.png`.
 
-Wan performs one flow-matching Euler step at timestep 1000 with guidance 1,
-using the prompt embedding stored in the checkpoint. No text encoder or
-ordinary Wan VAE is needed. Generic generation sampling settings do not
-configure this fixed sampling path.
+For a single image, replace `--input-frames input_frames` with `-i input.png`
+and use `-o output.png` (or JPEG/WebP). No FPS setting is needed for an image.
+`-i` / `--init-img` and `--input-frames` are mutually exclusive.
 
-The Flash scale cannot be verified from its tensor shapes: the x2 and x4 files
-have identical shapes. The caller must select the correct pair. Wan's input
-projection is checked against the requested scale. Inputs are replicate-padded
-internally and outputs are cropped to exactly the input size times the scale.
+## Settings and limits
 
-INT8 convrot uses the existing [quantized-weight execution path](int8_convrot.md).
-Backend and parameter placement use `--backend` and `--params-backend`, with
-Wan assigned to `diffusion` and Flash to `vae`. CUDA builds with CUTLASS use
-fused window attention on compute capability 8.0 or newer. It loads spatial
-neighborhoods directly, with FP16 Q/K/V and FP32 accumulation, and preserves
-Wan's boundary-shifted windows and Flash's boundary-truncated windows.
-Other backends use bounded spatial tiles with ordinary attention operators.
-Set `GGML_CUDA_DISABLE_WINDOW_ATTENTION=1` to use that fallback on CUDA for
-comparison. Floating-point results can differ between these implementations.
-
-Library callers initialize `sd_video_upscaler_ctx_params_t` with
-`sd_video_upscaler_ctx_params_init()`, set the model paths, explicit scale and
-backend options, and create a `video_upscaler_ctx_t` with
-`new_video_upscaler_ctx()`. Model architecture is detected from
-weight names and shapes; the scale remains caller-supplied.
-
-Initialize `sd_video_upscale_params_t` with `sd_video_upscale_params_init()`,
-then set the RGB frames, frame count, FPS, seed and RelayVSR reference window.
-Call `upscale_video()` and free the result with
-`free_sd_video_generation_result()`. Defaults are 24 FPS, seed 42 and a
-16-frame reference window. No sampler or scheduler selection is required.
-
-Each call starts fresh
-temporal state. Model caches are bounded by the temporal windows; the public
-API currently holds the complete input and output frame arrays in host memory.
-Audio is not copied by this path.
+- `--upscale-factor` is required. Use 2 or 4 with the corresponding weights.
+  Output width and height are exactly the input dimensions times the scale.
+- `--vsr-window-size` sets the reference-frame window, including both endpoints.
+  The default 16 selects reference frames 0, 15, 30, and so on, and always
+  includes the final frame. Every input frame appears once in the output.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+  [backend selection](backend.md) to choose a device or weight placement.
+- Sampler, scheduler, and step settings do not configure RelayVSR.
+  `--vsr-temporal-chunk-size` is a SeedVR2 option.
+- Long videos require enough system memory to hold the input and output frames.
+  Audio is not copied to the output.

docs/seedvr2.md

@@ -0,0 +1,69 @@
+# SeedVR2 image and video super-resolution
+
+SeedVR2 upscales images and videos with `-M vid_upscale`. Support is experimental.
+
+## Weights
+
+These [ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/SeedVR2)
+load directly, without conversion or a separate text encoder:
+
+| Model | `--diffusion-model` | `--vae` |
+| --- | --- | --- |
+| 3B | `seedvr2_3b_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+| 7B Sharp | `seedvr2_7b_sharp_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+
+Use the matching SeedVR2 VAE for either model. The 7B model requires more memory.
+
+## Upscale an image
+
+```sh
+sd-cli -M vid_upscale \
+  --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+  --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+  -i input.png --upscale-factor 2 -s 42 -o restored.png
+```
+
+`-i` / `--init-img` accepts a single image. Outputs can be PNG, JPEG, or WebP.
+No FPS or temporal chunk setting is needed for an image.
+
+## Upscale a video
+
+Extract the video into numbered frames, then pass the frame directory:
+
+```sh
+mkdir frames
+ffmpeg -i input.mp4 -fps_mode passthrough frames/%06d.png
+sd-cli -M vid_upscale \
+  --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+  --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+  --input-frames frames --upscale-factor 2 --fps 24 \
+  --vsr-temporal-chunk-size 1 -s 42 -o restored.avi
+```
+
+Set `--fps` to the source video's frame rate; 24 is only the example above.
+Frames are read in filename order. Use either `-i` or `--input-frames`, not both.
+
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o restored/%06d.png`.
+
+## Settings and limits
+
+- `--upscale-factor` is required and must be a positive integer. Output width
+  and height are exactly the input dimensions multiplied by this value.
+- `--vsr-temporal-chunk-size` controls how many video frames are processed
+  together. Valid values are 1, 5, 9, 13, and so on; the default is 5. A value of
+  1 processes each frame independently and uses less memory. Larger values use
+  neighboring frames but require more memory. Chunk boundaries may be visible.
+- Frame count and order are preserved, including a shorter final chunk. Audio
+  is not copied to the output.
+- If you run out of GPU memory, try `--vsr-temporal-chunk-size 1`. A five-frame
+  chunk at 1728x960 output resolution exceeded 24 GB of GPU memory in testing.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+  [backend selection](backend.md) to choose a device or weight placement.
+- SeedVR2 does not use sampler, scheduler, or step settings.
+  `--vsr-window-size` is a RelayVSR option and does not affect SeedVR2.
+
+3B has been tested on CUDA and CPU; 7B has been tested on CUDA with small inputs.
+Other backends have not been tested with this integration.
+
+Model reference: [ByteDance-Seed/SeedVR](https://github.com/ByteDance-Seed/SeedVR).

examples/cli/main.cpp

@@ -35,12 +35,13 @@ std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
 
 struct SDCliParams {
     std::string input_frames;
-    int video_upscale_factor = 0;
-    int vsr_window_size      = 16;
-    SDMode mode              = IMG_GEN;
-    std::string output_path  = "output.png";
-    int output_begin_idx     = -1;
-    int compression_quality  = 90;
+    int video_upscale_factor    = 0;
+    int vsr_window_size         = 16;
+    int vsr_temporal_chunk_size = 5;
+    SDMode mode                 = IMG_GEN;
+    std::string output_path     = "output.png";
+    int output_begin_idx        = -1;
+    int compression_quality     = 90;
     std::string image_path;
     std::string metadata_format = "text";
 
@@ -66,7 +67,7 @@ struct SDCliParams {
         ArgOptions options;
 
         options.string_options = {
-            {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 0, &input_frames},
+            {"", "--input-frames", "directory of ordered RGB frames for vid_upscale (use -i for a single image)", 0, &input_frames},
             {"-o",
              "--output",
              "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
@@ -95,8 +96,9 @@ struct SDCliParams {
         };
 
         options.int_options = {
-            {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor},
+            {"", "--upscale-factor", "required positive video upscale factor (RelayVSR: 2 or 4; SeedVR2: integer)", &video_upscale_factor},
             {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size},
+            {"", "--vsr-temporal-chunk-size", "SeedVR2 input frames per independent chunk: 1 or 4n+1 (default: 5)", &vsr_temporal_chunk_size},
             {"",
              "--preview-interval",
              "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
@@ -227,8 +229,8 @@ struct SDCliParams {
     }
 
     bool validate() {
-        if (mode == VID_UPSCALE && (input_frames.empty() || video_upscale_factor <= 0 || vsr_window_size < 2)) {
-            LOG_ERROR("vid_upscale requires --input-frames, positive --upscale-factor, and --vsr-window-size >= 2");
+        if (mode == VID_UPSCALE && video_upscale_factor <= 0) {
+            LOG_ERROR("vid_upscale requires a positive --upscale-factor");
             return false;
         }
         if (mode != METADATA) {
@@ -540,7 +542,9 @@ bool save_results(const SDCliParams& cli_params,
         return sucessful_reults != 0;
     }
 
-    if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) {
+    const bool video_upscale_output = cli_params.mode == VID_UPSCALE &&
+                                      (gen_params.init_image_path.empty() || ext_lower == ".avi" || ext_lower == ".webm");
+    if ((cli_params.mode == VID_GEN && num_results > 1) || video_upscale_output) {
         if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
         std::string params = gen_params.embed_image_metadata
@@ -676,12 +680,23 @@ int main(int argc, const char* argv[]) {
     }
 
     if (cli_params.mode == VID_UPSCALE) {
+        if (cli_params.input_frames.empty() == gen_params.init_image_path.empty()) {
+            LOG_ERROR("vid_upscale requires exactly one input: -i/--init-img or --input-frames");
+            return 1;
+        }
+        ctx_params.prepare_backend_assignments();
         if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) {
             LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps");
             return 1;
         }
         std::vector<SDImageOwner> owned_frames;
-        if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
+        if (!gen_params.init_image_path.empty()) {
+            owned_frames.emplace_back();
+            if (!load_sd_image_from_file(owned_frames.back().put(), gen_params.init_image_path.c_str(), 0, 0, 3)) {
+                LOG_ERROR("load image from '%s' failed", gen_params.init_image_path.c_str());
+                return 1;
+            }
+        } else if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
             LOG_ERROR("vid_upscale input contains no readable frames");
             return 1;
         }
@@ -692,12 +707,22 @@ int main(int argc, const char* argv[]) {
             frames.push_back(frame.get());
         sd_video_upscaler_ctx_params_t params;
         sd_video_upscaler_ctx_params_init(&params);
-        params.diffusion_model_path = ctx_params.diffusion_model_path.c_str();
-        params.vae_path             = ctx_params.vae_path.c_str();
-        params.upscale_factor       = cli_params.video_upscale_factor;
-        params.n_threads            = ctx_params.n_threads;
-        params.backend              = ctx_params.effective_backend.c_str();
-        params.params_backend       = ctx_params.effective_params_backend.c_str();
+        params.diffusion_model_path  = ctx_params.diffusion_model_path.c_str();
+        params.vae_path              = ctx_params.vae_path.c_str();
+        params.upscale_factor        = cli_params.video_upscale_factor;
+        params.n_threads             = ctx_params.n_threads;
+        params.backend               = ctx_params.effective_backend.c_str();
+        params.params_backend        = ctx_params.effective_params_backend.c_str();
+        params.flash_attn_mode       = ctx_params.flash_attn_mode.c_str();
+        std::string conv_direct_mode = ctx_params.conv_direct_mode;
+        if (conv_direct_mode.empty()) {
+            conv_direct_mode = "auto";
+            if (ctx_params.diffusion_conv_direct)
+                conv_direct_mode += ",diffusion=on";
+            if (ctx_params.vae_conv_direct)
+                conv_direct_mode += ",vae=on";
+        }
+        params.conv_direct_mode = conv_direct_mode.c_str();
         std::unique_ptr<video_upscaler_ctx_t, decltype(&free_video_upscaler_ctx)> upscaler(
             new_video_upscaler_ctx(&params),
             free_video_upscaler_ctx);
@@ -706,11 +731,12 @@ int main(int argc, const char* argv[]) {
         sd_video_generation_result_t* output = nullptr;
         sd_video_upscale_params_t upscale_params;
         sd_video_upscale_params_init(&upscale_params);
-        upscale_params.frames      = frames.data();
-        upscale_params.frame_count = int(frames.size());
-        upscale_params.fps         = gen_params.fps;
-        upscale_params.window_size = cli_params.vsr_window_size;
-        upscale_params.seed        = gen_params.seed;
+        upscale_params.frames              = frames.data();
+        upscale_params.frame_count         = int(frames.size());
+        upscale_params.fps                 = gen_params.fps;
+        upscale_params.window_size         = cli_params.vsr_window_size;
+        upscale_params.temporal_chunk_size = cli_params.vsr_temporal_chunk_size;
+        upscale_params.seed                = gen_params.seed;
         if (!upscale_video(upscaler.get(), &upscale_params, &output))
             return 1;
         std::string video_info = output->generation_info;

examples/common/common.cpp

@@ -1136,7 +1136,7 @@ ArgOptions SDGenerationParams::get_options() {
          &extra_ad_args},
         {"-i",
          "--init-img",
-         "path to the init image",
+         "path to the init image, or a single input image for vid_upscale",
          0,
          &init_image_path},
         {"",
@@ -1540,6 +1540,12 @@ 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 == "hyperflow-h3-flash-4") {
+            sigmas_str = "1.0,0.962575853,0.755781889,0.312497854,0.0";
+        } else if (sigmas_str == "hyperflow-h3-flash-6") {
+            sigmas_str = "1.0,0.982655942,0.914642572,0.674755692,0.385223746,0.159055337,0.0";
+        } else if (sigmas_str == "hyperflow-h3-flash-8") {
+            sigmas_str = "1.0,0.993909776,0.9842875,0.966063678,0.923076928,0.834942341,0.696852028,0.468753338,0.0";
         }
         if (!sigmas_str.empty() && sigmas_str.front() == '[') {
             sigmas_str.erase(0, 1);
@@ -1787,11 +1793,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.",
+         "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.",
          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.",
+         "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.",
          on_hires_sigmas_arg},
         {"",
          "--skip-layers",

examples/server/README.md

@@ -76,6 +76,16 @@ Recipes declaring `video_frames` also switch the form to Video Generation;
 other recipes switch to Image Generation. Generation is enabled after the
 selected recipe is loaded. Changing a selection does not reload the model
 until you click Load.
+
+The generation status bar shows the current stage and elapsed time even when
+Preview Mode is `none`. During sampling it shows completed steps for the current
+pass; prompt encoding and VAE decoding show an activity indicator when their
+total work is unknown. Sampling reaching 100% does not mean the job is finished:
+decoding, ADetailer, and output encoding may still follow.
+The canvas keeps the submitted output aspect ratio, and stage changes keep the
+same layout. Editing the next request's resolution does not resize the current
+preview or result.
+
 The server can also start with the existing model arguments shown below;
 existing API clients then work without selecting a recipe. The UI automatically
 selects **Command-line model (all features)** and uses the startup generation

examples/server/api.md

@@ -775,6 +775,7 @@ Field types:
 | `queue_position` | `integer` |
 | `result` | `object \| null` |
 | `error` | `object \| null` |
+| `progress` | `object` |
 | `preview` | `object \| null` |
 
 `created`, `started`, and `completed` are integer Unix timestamps in
@@ -782,6 +783,57 @@ Field types:
 the elapsed generation time in seconds, including output encoding and excluding
 queue time.
 
+`progress` is available independently of the preview setting, for both image
+and video jobs. Continue polling until `status` is terminal: a sampling pass
+can finish before decoding, ADetailer, or output encoding finishes.
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `progress.stage` | `string` | Current execution stage; see below |
+| `progress.detail` | `string` | Stage detail, or an empty string |
+| `progress.pass` | `integer` | Sampling pass number within the job; zero before sampling |
+| `progress.completed` | `integer \| null` | Completed units in the current stage/pass |
+| `progress.total` | `integer \| null` | Actual total units, or null when unknown |
+| `progress.unit` | `string \| null` | `steps` for sampling, `images` for image output encoding; null when uncounted |
+| `progress.stage_elapsed_ms` | `integer` | Elapsed time in the current stage/detail/pass |
+| `progress.elapsed_ms` | `integer` | Job execution time, excluding queue time; zero before execution |
+
+Stages include `queued`, `waiting` (waiting for the execution lock), `preparing`,
+`conditioning`, `sampling`, `vae_encode`, `vae_decode`, `audio_decode`,
+`upscaling`, `adetailer`, and `encoding`. Stages may repeat or be skipped.
+Details include `positive_prompt` / `negative_prompt` / `image_condition` for
+conditioning, `high_noise` / `low_noise` for two-stage sampling, and `image` /
+`video` for output encoding. During ADetailer, `detail` identifies the inner
+stage, such as `conditioning` or `sampling`. Clients should tolerate unknown
+stage, detail, and unit values.
+
+Sampling counts completed logical steps using the actual schedule, including
+strength adjustments. Each new sampling pass increments `pass` and starts at
+zero, including batch items, highres passes, high/low-noise passes, and ADetailer
+passes. The ratio is **current-pass progress**, not overall job completion.
+The progress snapshot is published atomically. Preview metadata still describes
+the preview frame, which may be older than the current progress.
+
+Unknown totals are represented by null counts, not an estimated percentage.
+For example, prompt encoding can report:
+
+```json
+"progress": {
+  "stage": "conditioning",
+  "detail": "positive_prompt",
+  "pass": 0,
+  "completed": null,
+  "total": null,
+  "unit": null,
+  "stage_elapsed_ms": 3200,
+  "elapsed_ms": 4100
+}
+```
+
+Elapsed times freeze when a job completes, fails, or is cancelled; the last
+stage is retained for inspection. Use `status` to determine the final outcome.
+Elapsed time alone does not indicate that computation is advancing.
+
 `preview` sub-fields:
 
 | Field | Type | Notes |

examples/server/async_jobs.cpp

@@ -92,6 +92,63 @@ static void set_preview_callback(AsyncGenerationJob& job) {
     sd_set_preview_callback(preview_callback, mode, interval, true, false, &job);
 }
 
+static void update_job_progress(AsyncGenerationJob& job, const char* stage, const char* detail = "", int sample_pass = -1, int64_t completed = 0, int64_t total = 0, const char* unit = "") {
+    std::lock_guard<std::mutex> lock(job.progress_mutex);
+    auto& progress = job.progress;
+    if (sample_pass < 0) {
+        sample_pass = progress.sample_pass;
+    }
+    if (progress.stage != stage || progress.detail != detail || progress.sample_pass != sample_pass) {
+        progress.stage_started = std::chrono::steady_clock::now();
+    }
+    progress.stage       = stage;
+    progress.detail      = detail;
+    progress.sample_pass = sample_pass;
+    progress.completed   = completed;
+    progress.total       = total;
+    progress.unit        = unit;
+}
+
+static void finish_job_progress(AsyncGenerationJob& job) {
+    std::lock_guard<std::mutex> lock(job.progress_mutex);
+    job.progress.finished = std::chrono::steady_clock::now();
+}
+
+class JobCallbacksGuard {
+    sd_ctx_t* ctx_;
+    AsyncGenerationJob& job_;
+    bool detailing_ = false;
+
+    static void progress_callback(const sd_generation_progress_t* progress, void* data) {
+        auto& guard = *static_cast<JobCallbacksGuard*>(data);
+        update_job_progress(guard.job_, guard.detailing_ ? "adetailer" : progress->stage,
+                            guard.detailing_ ? progress->stage : progress->detail,
+                            progress->sample_pass, progress->completed, progress->total, progress->unit);
+    }
+
+public:
+    JobCallbacksGuard(sd_ctx_t* ctx, AsyncGenerationJob& job)
+        : ctx_(ctx), job_(job) {
+        clear_preview_callback();
+        set_preview_callback(job);
+        sd_set_generation_progress_callback(ctx_, progress_callback, this);
+    }
+
+    ~JobCallbacksGuard() {
+        sd_set_generation_progress_callback(ctx_, nullptr, nullptr);
+        clear_preview_callback();
+    }
+
+    void begin_detailing() {
+        clear_preview_callback();
+        detailing_ = true;
+        update_job_progress(job_, "adetailer");
+    }
+
+    JobCallbacksGuard(const JobCallbacksGuard&)            = delete;
+    JobCallbacksGuard& operator=(const JobCallbacksGuard&) = delete;
+};
+
 const char* async_job_kind_name(AsyncJobKind kind) {
     switch (kind) {
         case AsyncJobKind::ImgGen:
@@ -186,6 +243,7 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) {
     job.result_fps         = 0;
     job.error_code         = "cancelled";
     job.error_message      = "job cancelled by client";
+    finish_job_progress(job);
     return true;
 }
 
@@ -200,6 +258,24 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo
     result["completed"]        = job.completed_at == 0 ? json(nullptr) : json(job.completed_at);
     result["queue_position"]   = 0;
 
+    {
+        std::lock_guard<std::mutex> lock(job.progress_mutex);
+        const auto& progress = job.progress;
+        const auto end       = progress.finished == std::chrono::steady_clock::time_point{}
+                                   ? std::chrono::steady_clock::now()
+                                   : progress.finished;
+        result["progress"]   = {
+            {"stage", progress.stage},
+            {"detail", progress.detail},
+            {"pass", progress.sample_pass},
+            {"completed", progress.total > 0 ? json(progress.completed) : json(nullptr)},
+            {"total", progress.total > 0 ? json(progress.total) : json(nullptr)},
+            {"unit", progress.unit.empty() ? json(nullptr) : json(progress.unit)},
+            {"stage_elapsed_ms", std::chrono::duration_cast<std::chrono::milliseconds>(end - progress.stage_started).count()},
+            {"elapsed_ms", job.started_at == 0 ? 0 : std::max<int64_t>(0, (job.completed_at == 0 ? unix_timestamp_now_ms() : job.completed_at) - job.started_at)},
+        };
+    }
+
     if (job.status == AsyncJobStatus::Queued) {
         size_t position = 1;
         for (const auto& queued_id : manager.queue) {
@@ -278,17 +354,19 @@ bool execute_img_gen_job(ServerRuntime& runtime,
             return false;
         }
         JobCancellationMonitor cancellation(runtime.sd_ctx, job);
-        set_preview_callback(job);
+        JobCallbacksGuard callbacks(runtime.sd_ctx, job);
         sd_image_generation_result_t* generated = nullptr;
         generate_image(runtime.sd_ctx, &params, &generated);
         results.adopt(generated, runtime.ctx_params->lora_model_dir);
-        clear_preview_callback();
         if (job.cancel_requested.load()) {
             return false;
         }
-        if (!results.empty() && !apply_server_adetailer(runtime, job.img_gen, params, results)) {
-            error_message = "ADetailer failed";
-            return false;
+        if (!results.empty() && job.img_gen.adetailer_enabled && !runtime.default_gen_params->ad_model_path.empty()) {
+            callbacks.begin_detailing();
+            if (!apply_server_adetailer(runtime, job.img_gen, params, results)) {
+                error_message = "ADetailer failed";
+                return false;
+            }
         }
     }
 
@@ -305,6 +383,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         encoded_format = EncodedImageFormat::WEBP;
     }
 
+    update_job_progress(job, "encoding", "image", -1, 0, num_results, "images");
     for (int i = 0; i < num_results; ++i) {
         if (results[i].data == nullptr) {
             continue;
@@ -325,6 +404,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         }
         output_images.push_back(base64_encode(image_bytes));
         output_info.push_back(json::parse(results.generation_info(i)));
+        update_job_progress(job, "encoding", "image", -1, i + 1, num_results, "images");
     }
 
     if (output_images.empty()) {
@@ -356,7 +436,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
             return false;
         }
         JobCancellationMonitor cancellation(runtime.sd_ctx, job);
-        set_preview_callback(job);
+        JobCallbacksGuard callbacks(runtime.sd_ctx, job);
         sd_video_generation_result_t* generated = nullptr;
         if (generate_video(runtime.sd_ctx, &params, &generated)) {
             output_fps       = generated->fps;
@@ -367,7 +447,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
             if (job.vid_gen.gen_params.embed_image_metadata)
                 str_params = results.generation_info(0);
         }
-        clear_preview_callback();
     }
 
     num_results = results.count();
@@ -377,6 +456,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
         return false;
     }
 
+    update_job_progress(job, "encoding", "video");
     std::vector<uint8_t> video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format,
                                                                              results.data(),
                                                                              num_results,
@@ -425,6 +505,7 @@ void async_job_worker(ServerRuntime& runtime) {
             job             = it->second;
             job->status     = AsyncJobStatus::Generating;
             job->started_at = unix_timestamp_now_ms();
+            update_job_progress(*job, "waiting");
         }
 
         std::vector<std::string> output_images;
@@ -459,6 +540,7 @@ void async_job_worker(ServerRuntime& runtime) {
             }
 
             job->completed_at = unix_timestamp_now_ms();
+            finish_job_progress(*job);
             if (job->cancel_requested.load()) {
                 job->status        = AsyncJobStatus::Cancelled;
                 job->error_code    = "cancelled";

examples/server/async_jobs.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <atomic>
+#include <chrono>
 #include <condition_variable>
 #include <cstdint>
 #include <deque>
@@ -29,6 +30,17 @@ enum class AsyncJobStatus {
 const char* async_job_kind_name(AsyncJobKind kind);
 const char* async_job_status_name(AsyncJobStatus status);
 
+struct AsyncJobProgress {
+    std::string stage = "queued";
+    std::string detail;
+    std::string unit;
+    int sample_pass                                     = 0;
+    int64_t completed                                   = 0;
+    int64_t total                                       = 0;
+    std::chrono::steady_clock::time_point stage_started = std::chrono::steady_clock::now();
+    std::chrono::steady_clock::time_point finished{};
+};
+
 struct AsyncGenerationJob {
     std::string id;
     AsyncJobKind kind     = AsyncJobKind::ImgGen;
@@ -47,6 +59,8 @@ struct AsyncGenerationJob {
     int result_fps         = 0;
     std::string error_code;
     std::string error_message;
+    mutable std::mutex progress_mutex;
+    AsyncJobProgress progress;
     mutable std::mutex preview_mutex;
     std::string preview_b64;
     int preview_pass          = 0;

examples/server/recipes/defaults.md

@@ -42,6 +42,8 @@ For community edit adapters and other optional extensions, these are the base-mo
 | `flux3-klein` | 1 | 16 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/README.md)  |
 | `flux3-klein-base` | 4 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/src/flux3/pipeline.py)  |
 | `hidream-o1-dev` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev)  |
+| `hyperflow-h3-flash-fl2va` | 1 | 6 | 1344 × 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Two-time FL2VA checkpoint; native text and first/last-frame generation. Fixed 6-step video sigma grid with shift 12, 124 frames at 24 FPS. Reuses MiniMax-H3 components. |
+| `hyperflow-h3-flash-ref2va` | 1 | 6 | 1344 × 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Separate two-time Ref2VA checkpoint; image references. Same schedule and components as FL2VA; video/audio references remain unverified. |
 | `hunyuan-video-15-i2v` | 6 | 20 | 1280 × 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base I2V defaults; CFG/step-distilled and 480p checkpoints use different settings. |
 | `hunyuan-video-15-t2v` | 6 | 20 | 1280 × 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base T2V defaults; CFG-distilled / 480p checkpoints need their published settings. |
 | `ideogram4` | 7 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/ideogram-oss/ideogram4/main/docs/inference.md) Constant-guidance pipeline defaults; this recipe does not reproduce the CLI’s per-step guidance schedule. |

examples/server/recipes/hyperflow-h3-flash-fl2va.json

@@ -0,0 +1,75 @@
+{
+  "schema_version": 1,
+  "id": "hyperflow-h3-flash-fl2va",
+  "name": "HyperFlow H3 Flash FL2VA",
+  "size_multiple": 32,
+  "description": "Generate video and audio from text, optionally with start/end images. Use the HyperFlow H3 Flash FL2VA transformer. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.",
+  "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true,
+      "label": "HyperFlow H3 Flash FL2VA Transformer"
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder"
+    },
+    {
+      "id": "vae",
+      "required": false,
+      "label": "MiniMax H3 Video VAE"
+    },
+    {
+      "id": "tae",
+      "required": false
+    },
+    {
+      "id": "audio_vae",
+      "required": false,
+      "label": "MiniMax H3 Audio VAE"
+    }
+  ],
+  "required_any": [
+    [
+      "vae",
+      "tae"
+    ]
+  ],
+  "video_constraints": {
+    "frame_step": 17,
+    "frame_offset": 5,
+    "frame_rounding": "ceil",
+    "min_frames": 5,
+    "fixed_fps": 24
+  },
+  "defaults": {
+    "video_frames": 124,
+    "fps": 24,
+    "cfg_scale": 1,
+    "width": 1344,
+    "height": 768,
+    "sampling_method": "euler",
+    "steps": 6,
+    "flow_shift": 12,
+    "custom_sigmas": [
+      1.0,
+      0.982655942,
+      0.914642572,
+      0.674755692,
+      0.385223746,
+      0.159055337,
+      0.0
+    ]
+  },
+  "features": [
+    "end_image",
+    "flow_shift",
+    "init_image"
+  ],
+  "disabled_features": [
+    "img2img",
+    "mask_image"
+  ]
+}

examples/server/recipes/hyperflow-h3-flash-ref2va.json

@@ -0,0 +1,74 @@
+{
+  "schema_version": 1,
+  "id": "hyperflow-h3-flash-ref2va",
+  "name": "HyperFlow H3 Flash Ref2VA",
+  "size_multiple": 32,
+  "description": "Generate video and audio with reference images. Use the HyperFlow H3 Flash Ref2VA transformer; video/audio references are not validated for these weights. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.",
+  "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true,
+      "label": "HyperFlow H3 Flash REF2VA Transformer"
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder"
+    },
+    {
+      "id": "vae",
+      "required": false,
+      "label": "MiniMax H3 Video VAE"
+    },
+    {
+      "id": "tae",
+      "required": false
+    },
+    {
+      "id": "audio_vae",
+      "required": false,
+      "label": "MiniMax H3 Audio VAE"
+    }
+  ],
+  "required_any": [
+    [
+      "vae",
+      "tae"
+    ]
+  ],
+  "video_constraints": {
+    "frame_step": 17,
+    "frame_offset": 5,
+    "frame_rounding": "ceil",
+    "min_frames": 5,
+    "fixed_fps": 24
+  },
+  "defaults": {
+    "video_frames": 124,
+    "fps": 24,
+    "cfg_scale": 1,
+    "width": 1344,
+    "height": 768,
+    "sampling_method": "euler",
+    "steps": 6,
+    "flow_shift": 12,
+    "custom_sigmas": [
+      1.0,
+      0.982655942,
+      0.914642572,
+      0.674755692,
+      0.385223746,
+      0.159055337,
+      0.0
+    ]
+  },
+  "features": [
+    "flow_shift",
+    "ref_images"
+  ],
+  "disabled_features": [
+    "img2img",
+    "mask_image"
+  ]
+}

include/stable-diffusion.h

@@ -500,6 +500,15 @@ struct ggml_tensor;
 
 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
 typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
+typedef struct {
+    const char* stage;
+    const char* detail;
+    const char* unit;
+    int sample_pass;
+    int64_t completed;
+    int64_t total;
+} sd_generation_progress_t;
+typedef void (*sd_generation_progress_cb_t)(const sd_generation_progress_t* progress, void* data);
 typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
 
 typedef struct {
@@ -510,6 +519,11 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void*
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
 SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
+// Context-local generation events, independent of preview and the legacy progress callback.
+// Strings and the event are valid only during the synchronous callback. Callbacks must not throw.
+// total == 0 means unknown; sample_pass counts from 1 since callback registration.
+// Register/clear only while the context is idle. Passing nullptr clears the callback.
+SD_API void sd_set_generation_progress_callback(sd_ctx_t* ctx, sd_generation_progress_cb_t cb, void* data);
 // In each sampling pass, a positive interval previews every Nth denoiser step, while a
 // negative interval previews only completed logical step -interval. Zero previews the final
 // completed step of the first sampling pass (base-resolution or high-noise).
@@ -625,6 +639,9 @@ typedef struct {
     int n_threads;
     const char* backend;
     const char* params_backend;
+    // auto/on/off or diffusion/vae assignments; null selects auto.
+    const char* flash_attn_mode;
+    const char* conv_direct_mode;
 } sd_video_upscaler_ctx_params_t;
 
 typedef struct {
@@ -634,6 +651,8 @@ typedef struct {
     int64_t seed;
     // RelayVSR reference window, including both endpoints; defaults to 16.
     int window_size;
+    // SeedVR2 independent temporal chunks in input frames (1 or 4n+1); defaults to 5.
+    int temporal_chunk_size;
 } sd_video_upscale_params_t;
 
 SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params);

scripts/format-code.ps1

@@ -10,6 +10,8 @@ $patterns = @(
     "src/core/*.hpp"
     "src/core/cpu/*.cpp"
     "src/core/cpu/*.h"
+    "src/core/metal/*.mm"
+    "src/core/metal/*.metal"
     "src/core/cuda/*.cu"
     "src/core/cuda/*.cuh"
     "src/core/cuda/*.h"

scripts/format-code.sh

@@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \
          src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
          src/core/*.cpp src/core/*.h src/core/*.hpp \
          src/core/cpu/*.cpp src/core/cpu/*.h \
+         src/core/metal/*.mm src/core/metal/*.metal \
          src/core/cuda/*.cu src/core/cuda/*.cuh src/core/cuda/*.h \
          src/core/cuda/sageattention/*.cuh \
          src/core/probe/*.cpp src/core/probe/*.h \

src/conditioning/conditioner.hpp

@@ -3448,8 +3448,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/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,41 @@ 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);
+                              float scale         = 1.f,
+                              bool allow_conv2d   = false);
 
 // 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);
+std::vector<ggml_tensor*> ggml_ext_split_qkv(GGMLRunnerContext* runner_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);
+std::vector<ggml_tensor*> ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv);
 
 // 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 +190,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

@@ -13,22 +13,6 @@
 
 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());

src/core/ggml_runner.h

@@ -23,6 +23,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 +48,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,17 +130,6 @@ 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_;
@@ -142,11 +137,15 @@ private:
     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:

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/metal/CMakeLists.txt

@@ -0,0 +1,8 @@
+set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS int8.metal int8_mpp.metal)
+file(READ "${CMAKE_CURRENT_SOURCE_DIR}/int8.metal" SD_METAL_I8_SOURCE)
+file(READ "${CMAKE_CURRENT_SOURCE_DIR}/int8_mpp.metal" SD_METAL_I8_MPP_SOURCE)
+configure_file(int8_source.h.in "${CMAKE_CURRENT_BINARY_DIR}/int8_source.h" @ONLY)
+target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/int8.mm")
+target_compile_options(${SD_LIB} PRIVATE "$<$<COMPILE_LANGUAGE:OBJCXX>:-fno-objc-arc>")
+target_include_directories(${SD_LIB} PRIVATE "${CMAKE_CURRENT_BINARY_DIR}")
+set_target_properties(${SD_LIB} PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED ON)

src/core/metal/int8.metal

@@ -0,0 +1,143 @@
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+using namespace metal;
+
+struct QuantParams {
+    uint k;
+    uint rows;
+    uint group;
+    uint unused;
+    ulong scale_offset;
+};
+
+kernel void sd_i8_rotate(device const float* input [[buffer(0)]],
+                         device float* output [[buffer(1)]],
+                         constant QuantParams& p [[buffer(2)]],
+                         uint block [[threadgroup_position_in_grid]],
+                         uint tid [[thread_index_in_threadgroup]]) {
+    threadgroup float values[256];
+    const ulong base = ulong(block) * p.group;
+    values[tid]      = input[base + tid] * rsqrt(float(p.group));
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (uint stride = 1; stride < p.group; stride *= 4) {
+        uint index = tid / (4 * stride) * (4 * stride) + tid % stride;
+        float a = values[index], b = values[index + stride];
+        float c = values[index + 2 * stride], d = values[index + 3 * stride];
+        float value;
+        switch ((tid / stride) % 4) {
+            case 0:
+                value = a + b + c - d;
+                break;
+            case 1:
+                value = a + b - c + d;
+                break;
+            case 2:
+                value = a - b + c + d;
+                break;
+            default:
+                value = -a + b + c + d;
+                break;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        values[tid] = value;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+    output[base + tid] = values[tid];
+}
+
+kernel void sd_i8_quantize(device const float* input [[buffer(0)]],
+                           device char* output [[buffer(1)]],
+                           constant QuantParams& p [[buffer(2)]],
+                           uint row [[threadgroup_position_in_grid]],
+                           uint tid [[thread_index_in_threadgroup]],
+                           uint lane [[thread_index_in_simdgroup]],
+                           uint simd [[simdgroup_index_in_threadgroup]]) {
+    threadgroup float maxima[8];
+    const ulong base = ulong(row) * p.k;
+    float amax       = 0;
+    for (uint i = tid; i < p.k; i += 256) {
+        amax = max(amax, abs(input[base + i]));
+    }
+    amax = simd_max(amax);
+    if (lane == 0)
+        maxima[simd] = amax;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    amax          = simd_max(lane < 8 ? maxima[lane] : 0.f);
+    float scale   = amax / 127.f;
+    float inverse = scale == 0 ? 0 : 1.f / scale;
+    for (uint i = tid; i < p.k; i += 256) {
+        output[base + i] = char(clamp(rint(input[base + i] * inverse), -127.f, 127.f));
+    }
+    if (tid == 0)
+        ((device float*)(output + p.scale_offset))[row] = scale;
+}
+
+struct MatParams {
+    uint k;
+    uint n;
+    uint rows;
+    uint flags;
+};
+
+kernel void sd_i8_multiply(device const char* weight [[buffer(0)]],
+                           device const char* input [[buffer(1)]],
+                           device const float* xs [[buffer(2)]],
+                           device const float* ws [[buffer(3)]],
+                           device const float* bias [[buffer(4)]],
+                           device float* output [[buffer(5)]],
+                           constant MatParams& p [[buffer(6)]],
+                           uint2 tile [[threadgroup_position_in_grid]],
+                           uint tid [[thread_index_in_threadgroup]],
+                           uint simd [[simdgroup_index_in_threadgroup]]) {
+    threadgroup half a[32 * 32];
+    threadgroup half b[32 * 32];
+    threadgroup float partial[32 * 32];
+    uint sums[8]    = {};
+    const uint row0 = tile.y * 32, col0 = tile.x * 32;
+    const uint sr = (simd / 2) * 16, sc = (simd % 2) * 16;
+
+    // Half represents every INT8 value exactly. Each 256-term dot product fits
+    // in FP32's exact integer range; combine chunks as integers to avoid drift.
+    for (uint chunk = 0; chunk < p.k; chunk += 256) {
+        simdgroup_float8x8 acc[4];
+        for (uint i = 0; i < 4; ++i)
+            acc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
+        for (uint kk = chunk; kk < min(chunk + 256, p.k); kk += 32) {
+            for (uint index = tid; index < 1024; index += 128) {
+                uint r = index / 32, k = index % 32;
+                a[index] = row0 + r < p.rows && kk + k < p.k ? half(input[ulong(row0 + r) * p.k + kk + k]) : half(0);
+                b[index] = col0 + r < p.n && kk + k < p.k ? half(weight[ulong(col0 + r) * p.k + kk + k]) : half(0);
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            for (uint k = 0; k < 32; k += 8) {
+                simdgroup_half8x8 ma[2], mb[2];
+                for (uint i = 0; i < 2; ++i) {
+                    simdgroup_load(ma[i], a + (sr + i * 8) * 32 + k, 32);
+                    simdgroup_load(mb[i], b + (sc + i * 8) * 32 + k, 32, ulong2(0), true);
+                }
+                for (uint i = 0; i < 4; ++i) {
+                    simdgroup_multiply_accumulate(acc[i], ma[i / 2], mb[i % 2], acc[i]);
+                }
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+        for (uint i = 0; i < 4; ++i) {
+            simdgroup_store(acc[i], partial + (sr + (i / 2) * 8) * 32 + sc + (i % 2) * 8, 32);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (uint i = 0; i < 8; ++i)
+            sums[i] += uint(int(partial[tid + 128 * i]));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+    for (uint i = 0; i < 8; ++i) {
+        uint index = tid + 128 * i, row = row0 + index / 32, col = col0 + index % 32;
+        if (row < p.rows && col < p.n) {
+            float value = float(as_type<int>(sums[i])) * xs[row];
+            if (p.flags & 1)
+                value *= ws[col];
+            if (p.flags & 2)
+                value += bias[col];
+            output[ulong(row) * p.n + col] = value;
+        }
+    }
+}

src/core/metal/int8.mm

@@ -0,0 +1,298 @@
+#include "core/ggml_extensions.h"
+#include "ggml-ext-metal.h"
+#include "int8_source.h"
+
+#import <Foundation/Foundation.h>
+#import <Metal/Metal.h>
+
+#include <climits>
+#include <memory>
+#include <vector>
+
+namespace sd_metal {
+
+    struct Kernels {
+        id<MTLComputePipelineState> rotate       = nil;
+        id<MTLComputePipelineState> quantize     = nil;
+        id<MTLComputePipelineState> multiply     = nil;
+        id<MTLComputePipelineState> multiply_mpp = nil;
+        ggml_ext_metal_buffer_fn buffer          = nullptr;
+        ggml_ext_metal_submit_fn submit          = nullptr;
+
+        ~Kernels() {
+            [rotate release];
+            [quantize release];
+            [multiply release];
+            [multiply_mpp release];
+        }
+    };
+
+    static id<MTLComputePipelineState> tensor_multiply(id<MTLDevice> device) {
+        if (@available(macOS 26.0, iOS 26.0, *)) {
+            MTLCompileOptions* options = [[MTLCompileOptions alloc] init];
+            // Keep building with older SDKs; the optional shader is compiled at runtime.
+            options.languageVersion = static_cast<MTLLanguageVersion>(4 << 16);
+            options.fastMathEnabled = NO;
+            NSError* error          = nil;
+            id<MTLLibrary> library  = [device newLibraryWithSource:[NSString stringWithUTF8String:sd_metal_i8_mpp_source] options:options error:&error];
+            [options release];
+            id<MTLFunction> function             = [library newFunctionWithName:@"sd_i8_multiply_mpp"];
+            id<MTLComputePipelineState> pipeline = function ? [device newComputePipelineStateWithFunction:function error:&error] : nil;
+            [function release];
+            [library release];
+            if (pipeline && pipeline.threadExecutionWidth == 32 && pipeline.maxTotalThreadsPerThreadgroup >= 128) {
+                return pipeline;
+            }
+            [pipeline release];
+            NSLog(@"sd.cpp Metal INT8 TensorOps unavailable; using SIMD-group fallback: %@", error);
+        }
+        return nil;
+    }
+
+    struct QuantParams {
+        uint32_t k, rows, group, unused;
+        uint64_t scale_offset;
+    };
+
+    struct MatParams {
+        uint32_t k, n, rows, flags;
+    };
+
+    struct Execution {
+        const ggml_ext_execution_context* context;
+        ggml_tensor* node;
+        Kernels* kernels;
+        SDI8Operation op = {};
+        int group        = 0;
+        bool linear      = false;
+    };
+
+    static bool dimensions(const ggml_tensor* node) {
+        return node->ne[0] <= INT_MAX && ggml_nrows(node) <= INT_MAX;
+    }
+
+    static bool supports_i8(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDI8Operation op;
+        return sd_ggml_i8_describe(node, &op) && dimensions(node) && dimensions(op.weight) &&
+               uint64_t(op.weight->ne[0]) * ggml_nrows(node) <= UINT32_MAX;
+    }
+
+    static bool supports_convrot(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        int group;
+        return sd_ggml_convrot_describe(node, &group) && dimensions(node->src[0]) &&
+               uint64_t(ggml_nelements(node->src[0])) <= UINT32_MAX;
+    }
+
+    static size_t aligned(size_t bytes) {
+        return (bytes + 255) / 256 * 256;
+    }
+
+    static size_t packed_size(int64_t k, int64_t rows) {
+        return aligned(size_t(k) * ((rows + 3) / 4 * 4) + rows * sizeof(float));
+    }
+
+    static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        auto* node = nodes[0];
+        SDI8Operation op;
+        const bool linear  = sd_ggml_i8_describe(node, &op);
+        const auto* input  = linear ? op.input : node->src[0];
+        const int64_t rows = linear ? ggml_nrows(node) : ggml_nrows(input);
+        int group;
+        if (linear) {
+            group = op.group_size;
+        } else if (!sd_ggml_convrot_describe(node, &group)) {
+            return GGML_STATUS_FAILED;
+        }
+        *bytes = aligned(sd_ggml_modulate_workspace_size(node));
+        if (input->type == GGML_TYPE_F32) {
+            if (group > 1)
+                *bytes += aligned(ggml_nbytes(input));
+            if (linear)
+                *bytes += packed_size(input->ne[0], rows);
+        }
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static bool bind_tensor(Kernels* kernels, const ggml_tensor* tensor, ggml_ext_metal_buffer& result) {
+        auto buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+        if (!buffer || !tensor->data)
+            return false;
+        const size_t offset = uintptr_t(tensor->data) - uintptr_t(ggml_backend_buffer_get_base(buffer));
+        return kernels->buffer(buffer, offset, ggml_nbytes(tensor), &result);
+    }
+
+    static void set_buffer(id<MTLComputeCommandEncoder> encoder, const ggml_ext_metal_buffer& buffer, int index) {
+        [encoder setBuffer:(id<MTLBuffer>)buffer.buffer offset:buffer.offset atIndex:index];
+    }
+
+    static ggml_status encode(void* command_buffer, void* user_data) {
+        auto& e                      = *static_cast<Execution*>(user_data);
+        auto* kernels                = e.kernels;
+        const auto* input            = e.linear ? e.op.input : e.node->src[0];
+        const uint32_t k             = uint32_t(input->ne[0]);
+        const uint32_t rows          = uint32_t(e.linear ? ggml_nrows(e.node) : ggml_nrows(input));
+        id<MTLCommandBuffer> command = (id<MTLCommandBuffer>)command_buffer;
+        ggml_ext_metal_buffer x, dst, output, rotated, packed, xs, w, ws, bias;
+        if (!bind_tensor(kernels, input, x) || !bind_tensor(kernels, e.node, dst))
+            return GGML_STATUS_FAILED;
+        output        = dst;
+        size_t offset = 0;
+        auto scratch  = [&](size_t bytes, ggml_ext_metal_buffer& binding) {
+            bool ok = kernels->buffer(e.context->workspace, offset, bytes, &binding);
+            offset += aligned(bytes);
+            return ok;
+        };
+        const size_t output_scratch = sd_ggml_modulate_workspace_size(e.node);
+        if (output_scratch && !scratch(output_scratch, output))
+            return GGML_STATUS_FAILED;
+
+        if (input->type == GGML_TYPE_F32) {
+            rotated = x;
+            if (e.group > 1 && !scratch(ggml_nbytes(input), rotated))
+                return GGML_STATUS_FAILED;
+            packed = output;
+            if (e.linear && !scratch(packed_size(k, rows), packed))
+                return GGML_STATUS_FAILED;
+            QuantParams p                   = {k, rows, uint32_t(e.group), 0, uint64_t(k) * ((uint64_t(rows) + 3) / 4 * 4)};
+            id<MTLBlitCommandEncoder> clear = [command blitCommandEncoder];
+            [clear fillBuffer:(id<MTLBuffer>)packed.buffer range:NSMakeRange(packed.offset, e.linear ? packed_size(k, rows) : ggml_nbytes(e.node)) value:0];
+            [clear endEncoding];
+            if (e.group > 1) {
+                id<MTLComputeCommandEncoder> encoder = [command computeCommandEncoder];
+                [encoder setComputePipelineState:kernels->rotate];
+                set_buffer(encoder, x, 0);
+                set_buffer(encoder, rotated, 1);
+                [encoder setBytes:&p length:sizeof(p) atIndex:2];
+                [encoder dispatchThreadgroups:MTLSizeMake(uint64_t(k) * rows / e.group, 1, 1)
+                        threadsPerThreadgroup:MTLSizeMake(e.group, 1, 1)];
+                [encoder endEncoding];
+            }
+            id<MTLComputeCommandEncoder> encoder = [command computeCommandEncoder];
+            [encoder setComputePipelineState:kernels->quantize];
+            set_buffer(encoder, rotated, 0);
+            set_buffer(encoder, packed, 1);
+            [encoder setBytes:&p length:sizeof(p) atIndex:2];
+            [encoder dispatchThreadgroups:MTLSizeMake(rows, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+            [encoder endEncoding];
+            x  = packed;
+            xs = packed;
+            xs.offset += p.scale_offset;
+        } else if (e.op.packed) {
+            xs = x;
+            xs.offset += size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
+        } else if (!bind_tensor(kernels, e.op.input_scale, xs)) {
+            return GGML_STATUS_FAILED;
+        }
+
+        if (e.linear) {
+            if (!bind_tensor(kernels, e.op.weight, w))
+                return GGML_STATUS_FAILED;
+            ws = bias = xs;
+            if (e.op.weight_scale && !bind_tensor(kernels, e.op.weight_scale, ws))
+                return GGML_STATUS_FAILED;
+            if (e.op.bias && !bind_tensor(kernels, e.op.bias, bias))
+                return GGML_STATUS_FAILED;
+            MatParams p                          = {k, uint32_t(e.node->ne[0]), rows, uint32_t((e.op.weight_scale ? 1 : 0) | (e.op.bias ? 2 : 0))};
+            id<MTLComputeCommandEncoder> encoder = [command computeCommandEncoder];
+            [encoder setComputePipelineState:kernels->multiply_mpp ? kernels->multiply_mpp : kernels->multiply];
+            set_buffer(encoder, w, 0);
+            set_buffer(encoder, x, 1);
+            set_buffer(encoder, xs, 2);
+            set_buffer(encoder, ws, 3);
+            set_buffer(encoder, bias, 4);
+            set_buffer(encoder, output, 5);
+            [encoder setBytes:&p length:sizeof(p) atIndex:6];
+            [encoder dispatchThreadgroups:MTLSizeMake((p.n + 31) / 32, (rows + 31) / 32, 1)
+                    threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
+            [encoder endEncoding];
+        }
+        if (output_scratch) {
+            id<MTLBlitCommandEncoder> copy = [command blitCommandEncoder];
+            [copy copyFromBuffer:(id<MTLBuffer>)output.buffer
+                     sourceOffset:output.offset
+                         toBuffer:(id<MTLBuffer>)dst.buffer
+                destinationOffset:dst.offset
+                             size:ggml_nbytes(e.node)];
+            [copy endEncoding];
+        }
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static ggml_status execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void* user_data) {
+        Execution e = {context, nodes[0], static_cast<Kernels*>(user_data)};
+        e.linear    = sd_ggml_i8_describe(e.node, &e.op);
+        if (e.linear) {
+            e.group = e.op.group_size;
+        } else if (!sd_ggml_convrot_describe(e.node, &e.group)) {
+            return GGML_STATUS_FAILED;
+        }
+        return e.kernels->submit(context->backend, encode, &e);
+    }
+
+}
+
+void sd_ggml_register_metal_i8(ggml_backend_dev_t device) {
+    using namespace sd_metal;
+    auto reg        = ggml_backend_dev_backend_reg(device);
+    auto get_device = reinterpret_cast<ggml_ext_metal_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_device_v1"));
+    auto get_buffer = reinterpret_cast<ggml_ext_metal_buffer_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_buffer_v1"));
+    auto submit     = reinterpret_cast<ggml_ext_metal_submit_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_submit_v1"));
+    if (!get_device || !get_buffer || !submit)
+        return;
+    @autoreleasepool {
+        id<MTLDevice> metal = (id<MTLDevice>)get_device(device);
+        if (!metal || ![metal supportsFamily:MTLGPUFamilyApple7])
+            return;
+        auto kernels               = std::make_unique<Kernels>();
+        kernels->buffer            = get_buffer;
+        kernels->submit            = submit;
+        MTLCompileOptions* options = [[MTLCompileOptions alloc] init];
+        options.fastMathEnabled    = NO;
+        NSError* error             = nil;
+        id<MTLLibrary> library     = [metal newLibraryWithSource:[NSString stringWithUTF8String:sd_metal_i8_source] options:options error:&error];
+        [options release];
+        if (!library) {
+            NSLog(@"sd.cpp Metal INT8 compilation failed: %@", error);
+            return;
+        }
+        auto pipeline = [&](NSString* name) {
+            id<MTLFunction> function           = [library newFunctionWithName:name];
+            id<MTLComputePipelineState> result = function ? [metal newComputePipelineStateWithFunction:function error:&error] : nil;
+            [function release];
+            return result;
+        };
+        kernels->rotate   = pipeline(@"sd_i8_rotate");
+        kernels->quantize = pipeline(@"sd_i8_quantize");
+        kernels->multiply = pipeline(@"sd_i8_multiply");
+        [library release];
+        if (!kernels->rotate || !kernels->quantize || !kernels->multiply ||
+            kernels->quantize.threadExecutionWidth != 32 || kernels->multiply.threadExecutionWidth != 32 ||
+            kernels->rotate.maxTotalThreadsPerThreadgroup < 256 ||
+            kernels->quantize.maxTotalThreadsPerThreadgroup < 256 || kernels->multiply.maxTotalThreadsPerThreadgroup < 128) {
+            NSLog(@"sd.cpp Metal INT8 pipelines unavailable: %@", error);
+            return;
+        }
+        kernels->multiply_mpp = tensor_multiply(metal);
+        if (kernels->multiply_mpp) {
+            NSLog(@"sd.cpp Metal INT8: using TensorOps INT8 x INT8 -> INT32");
+        }
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.linear_i8.metal.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_i8_op();
+        desc.priority             = 100;
+        desc.supports             = supports_i8;
+        desc.prepare              = prepare;
+        desc.destroy              = [](void*, void*) {};
+        desc.execute              = execute;
+        desc.user_data            = kernels.get();
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+        desc.name      = "sd.quantize_i8_convrot.metal.v1";
+        desc.custom_op = sd_ggml_convrot_op();
+        desc.supports  = supports_convrot;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+        static std::vector<std::unique_ptr<Kernels>> registrations;
+        registrations.push_back(std::move(kernels));
+    }
+}

src/core/metal/int8_mpp.metal

@@ -0,0 +1,65 @@
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+#include <metal_stdlib>
+#include <metal_tensor>
+using namespace metal;
+using namespace mpp::tensor_ops;
+
+struct MatParams {
+    uint k;
+    uint n;
+    uint rows;
+    uint flags;
+};
+
+// TensorOps requires explicit signed INT8 and non-const tensor element types.
+kernel void sd_i8_multiply_mpp(device int8_t* weight [[buffer(0)]],
+                               device int8_t* input [[buffer(1)]],
+                               device const float* xs [[buffer(2)]],
+                               device const float* ws [[buffer(3)]],
+                               device const float* bias [[buffer(4)]],
+                               device float* output [[buffer(5)]],
+                               constant MatParams& p [[buffer(6)]],
+                               uint2 tile [[threadgroup_position_in_grid]]) {
+    auto a                    = tensor<device int8_t, dextents<int, 2>, tensor_inline>(input, dextents<int, 2>(p.k, p.rows));
+    auto b                    = tensor<device int8_t, dextents<int, 2>, tensor_inline>(weight, dextents<int, 2>(p.k, p.n));
+    constexpr int chunk_size  = 1024;
+    constexpr auto descriptor = matmul2d_descriptor(32, 32, chunk_size, false, true, false);
+    matmul2d<descriptor, execution_simdgroups<4>> multiply;
+    auto left  = a.slice(0, tile.y * 32);
+    auto right = b.slice(0, tile.x * 32);
+    auto acc   = multiply.get_destination_cooperative_tensor<decltype(left), decltype(right), int>();
+    auto sums  = multiply.get_destination_cooperative_tensor<decltype(left), decltype(right), int>();
+#pragma unroll
+    for (uint i = 0; i < sums.get_capacity(); ++i) {
+        if (sums.is_valid_element(i))
+            sums[i] = 0;
+    }
+    // TensorOps can round long integer dot products on older GPUs. Bound each
+    // partial sum to FP32's exact integer range, then combine with wrapping
+    // integer addition outside TensorOps.
+    for (uint k = 0; k < p.k; k += chunk_size) {
+        left  = a.slice(k, tile.y * 32);
+        right = b.slice(k, tile.x * 32);
+        multiply.run(left, right, acc);
+#pragma unroll
+        for (uint i = 0; i < acc.get_capacity(); ++i) {
+            if (acc.is_valid_element(i))
+                sums[i] = as_type<int>(as_type<uint>(sums[i]) + as_type<uint>(acc[i]));
+        }
+    }
+#pragma unroll
+    for (uint i = 0; i < acc.get_capacity(); ++i) {
+        if (acc.is_valid_element(i)) {
+            auto position = acc.get_multidimensional_index(i);
+            uint col = tile.x * 32 + position[0], row = tile.y * 32 + position[1];
+            if (col < p.n && row < p.rows) {
+                float value = float(sums[i]) * xs[row];
+                if (p.flags & 1)
+                    value *= ws[col];
+                if (p.flags & 2)
+                    value += bias[col];
+                output[ulong(row) * p.n + col] = value;
+            }
+        }
+    }
+}

src/core/metal/int8_source.h.in

@@ -0,0 +1,7 @@
+#ifndef __SD_CORE_METAL_INT8_SOURCE_H__
+#define __SD_CORE_METAL_INT8_SOURCE_H__
+
+static const char sd_metal_i8_source[] = R"SD_METAL(@SD_METAL_I8_SOURCE@)SD_METAL";
+static const char sd_metal_i8_mpp_source[] = R"SD_METAL(@SD_METAL_I8_MPP_SOURCE@)SD_METAL";
+
+#endif

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/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_split_qkv(ctx, qkv);
             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;
         }
 
@@ -1437,6 +1437,10 @@ namespace Flux {
             return "flux";
         }
 
+        int get_patch_size() const override {
+            return config.patch_size;
+        }
+
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             flux.get_param_tensors(tensors, prefix);
         }

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_split_qkv(ctx, qkv);
             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_split_qkv(ctx, qkv);
             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_split_qkv(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
             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_split_qkv(ctx, img_qkv->forward(ctx, img));
+            auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt));
 
             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

@@ -24,6 +24,19 @@ namespace MiniMaxH3 {
     constexpr float FRAME_RESCALE                 = 5.f / 3.f;
     constexpr float VISUAL_COND_TIMESTEP          = 0.999f;
 
+    struct LowRankShape {
+        int64_t q   = 0;
+        int64_t k   = 0;
+        int64_t v   = 0;
+        int64_t out = 0;
+        int64_t fc1 = 0;
+        int64_t fc2 = 0;
+
+        bool splits_qkv() const {
+            return q > 0 || k > 0 || v > 0;
+        }
+    };
+
     struct Config {
         int64_t hidden_size              = 5376;
         int64_t num_layers               = 50;
@@ -45,6 +58,8 @@ namespace MiniMaxH3 {
         float norm_eps                   = 1e-5f;
         float qk_norm_eps                = 1e-5f;
         float final_norm_eps             = 1e-5f;
+        bool two_time                    = false;
+        std::vector<LowRankShape> block_ranks;
 
         bool uses_adaln_curves() const {
             return adaln_curve_grid > 0;
@@ -86,12 +101,36 @@ namespace MiniMaxH3 {
             if (const auto* weight = find("blocks.0.attn.q_norm.weight")) {
                 config.attention_head_dim = weight->logical_ne0();
             }
+            auto rank_of = [&](const std::string& name) -> int64_t {
+                const auto* down = find(name + ".down.weight");
+                return down == nullptr ? 0 : down->ne[1];
+            };
+            config.block_ranks.resize(static_cast<size_t>(config.num_layers));
+            for (int64_t i = 0; i < config.num_layers; ++i) {
+                std::string block = "blocks." + std::to_string(i) + ".";
+                auto& ranks       = config.block_ranks[static_cast<size_t>(i)];
+                ranks.q           = rank_of(block + "attn.q_proj");
+                ranks.k           = rank_of(block + "attn.k_proj");
+                ranks.v           = rank_of(block + "attn.v_proj");
+                ranks.out         = rank_of(block + "attn.out_proj");
+                ranks.fc1         = rank_of(block + "mlp.fc1");
+                ranks.fc2         = rank_of(block + "mlp.fc2");
+            }
             if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) {
                 config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim);
+            } else {
+                const auto* q_weight = find("blocks.0.attn.q_proj.up.weight");
+                q_weight             = q_weight != nullptr ? q_weight : find("blocks.0.attn.q_proj.weight");
+                if (q_weight != nullptr) {
+                    config.num_attention_heads = q_weight->ne[1] / config.attention_head_dim;
+                }
             }
-            if (const auto* weight = find("blocks.0.mlp.fc1.weight")) {
-                config.ffn_hidden_size = weight->ne[1] / 2;
+            const auto* fc1 = find("blocks.0.mlp.fc1.weight");
+            fc1             = fc1 != nullptr ? fc1 : find("blocks.0.mlp.fc1.up.weight");
+            if (fc1 != nullptr) {
+                config.ffn_hidden_size = fc1->ne[1] / 2;
             }
+            config.two_time = find("endpoint_time_embedder.proj_in.weight") != nullptr;
             if (const auto* weight = find("condition_proj.weight")) {
                 config.text_dim = weight->logical_ne0();
             }
@@ -112,13 +151,14 @@ namespace MiniMaxH3 {
             }
 
             LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64
-                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64,
+                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d",
                         config.num_layers,
                         config.hidden_size,
                         config.num_attention_heads,
                         config.attention_head_dim,
                         config.ffn_hidden_size,
-                        config.adaln_curve_grid);
+                        config.adaln_curve_grid,
+                        config.two_time);
             return config;
         }
     };
@@ -141,31 +181,65 @@ namespace MiniMaxH3 {
         }
     };
 
+    struct LowRankLinear : public UnaryBlock {
+        LowRankLinear(int64_t in_features, int64_t rank, int64_t out_features, bool force_prec_f32 = false, float scale = 1.f) {
+            blocks["down"] = std::make_shared<Linear>(in_features, rank, false, false, force_prec_f32, scale);
+            blocks["up"]   = std::make_shared<Linear>(rank, out_features, false, false, force_prec_f32, scale);
+        }
+
+        std::shared_ptr<Linear> part(const char* name) {
+            return std::dynamic_pointer_cast<Linear>(blocks[name]);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto down = std::dynamic_pointer_cast<Linear>(blocks["down"]);
+            auto up   = std::dynamic_pointer_cast<Linear>(blocks["up"]);
+            return up->forward(ctx, down->forward(ctx, x));
+        }
+    };
+
+    static std::shared_ptr<UnaryBlock> make_linear(int64_t in_features, int64_t out_features, int64_t rank, bool force_prec_f32 = false, float scale = 1.f) {
+        if (rank > 0) {
+            return std::make_shared<LowRankLinear>(in_features, rank, out_features, force_prec_f32, scale);
+        }
+        return std::make_shared<Linear>(in_features, out_features, false, false, force_prec_f32, scale);
+    }
+
     struct MLP : public UnaryBlock {
         int64_t ffn_hidden_size;
 
         MLP(int64_t hidden_size,
-            int64_t ffn_hidden_size)
+            int64_t ffn_hidden_size,
+            int64_t fc1_rank = 0,
+            int64_t fc2_rank = 0)
             : ffn_hidden_size(ffn_hidden_size) {
-            blocks["fc1"] = std::make_shared<Linear>(hidden_size, ffn_hidden_size * 2, false, false, true, 1.f / 128.f);
-            blocks["fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f);
+            blocks["fc1"] = make_linear(hidden_size, ffn_hidden_size * 2, fc1_rank, true, 1.f / 128.f);
+            blocks["fc2"] = make_linear(ffn_hidden_size, hidden_size, fc2_rank, true, 1.f / 128.f);
         }
 
         void init_params(ggml_context* ctx,
                          const String2TensorStorage& tensors = {},
                          const std::string prefix            = "") override {
             // INT8 accumulates in INT32 and returns FP32, so it needs no FP16 overflow scaling.
-            for (const auto* name : {"fc1", "fc2"}) {
-                auto linear = std::dynamic_pointer_cast<Linear>(blocks[name]);
+            auto set_scale = [&](const std::shared_ptr<GGMLBlock>& block, const std::string& name) {
+                auto linear = std::dynamic_pointer_cast<Linear>(block);
                 linear->set_scale(get_type(prefix + name + ".weight", tensors, GGML_TYPE_F32) == GGML_TYPE_I8
                                       ? 1.f
                                       : 1.f / 128.f);
+            };
+            for (const std::string name : {"fc1", "fc2"}) {
+                if (auto low_rank = std::dynamic_pointer_cast<LowRankLinear>(blocks[name])) {
+                    set_scale(low_rank->part("down"), name + ".down");
+                    set_scale(low_rank->part("up"), name + ".up");
+                } else {
+                    set_scale(blocks[name], name);
+                }
             }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            auto fc1           = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
-            auto fc2           = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+            auto fc1           = std::dynamic_pointer_cast<UnaryBlock>(blocks["fc1"]);
+            auto fc2           = std::dynamic_pointer_cast<UnaryBlock>(blocks["fc2"]);
             auto forward_chunk = [&](ggml_tensor* input) {
                 auto projected = fc1->forward(ctx, input);
                 auto activated = ggml_swiglu(ctx->ggml_ctx, projected);
@@ -189,13 +263,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);
         }
     };
 
@@ -204,46 +278,43 @@ 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 {
         int64_t heads;
         int64_t head_dim;
         float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB;
+        bool split_projections;
 
         Attention(int64_t hidden_size,
                   int64_t heads,
                   int64_t head_dim,
-                  float eps)
-            : heads(heads), head_dim(head_dim) {
-            int64_t inner      = heads * head_dim;
-            blocks["qkv_proj"] = std::make_shared<Linear>(hidden_size, inner * 3, false);
+                  float eps,
+                  const LowRankShape& ranks = {})
+            : heads(heads), head_dim(head_dim), split_projections(ranks.splits_qkv()) {
+            int64_t inner = heads * head_dim;
+            if (split_projections) {
+                blocks["q_proj"] = make_linear(hidden_size, inner, ranks.q);
+                blocks["k_proj"] = make_linear(hidden_size, inner, ranks.k);
+                blocks["v_proj"] = make_linear(hidden_size, inner, ranks.v);
+            } else {
+                blocks["qkv_proj"] = std::make_shared<Linear>(hidden_size, inner * 3, false);
+            }
             blocks["q_norm"]   = std::make_shared<RMSNorm>(head_dim, eps);
             blocks["k_norm"]   = std::make_shared<RMSNorm>(head_dim, eps);
-            blocks["out_proj"] = std::make_shared<Linear>(inner, hidden_size, false);
+            blocks["out_proj"] = make_linear(inner, hidden_size, ranks.out);
+        }
+
+        ggml_tensor* project_packed_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            if (!split_projections) {
+                return std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"])->forward(ctx, x);
+            }
+            auto q = std::dynamic_pointer_cast<UnaryBlock>(blocks["q_proj"])->forward(ctx, x);
+            auto k = std::dynamic_pointer_cast<UnaryBlock>(blocks["k_proj"])->forward(ctx, x);
+            auto v = std::dynamic_pointer_cast<UnaryBlock>(blocks["v_proj"])->forward(ctx, x);
+            return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, q, k, 0), v, 0);
         }
 
         std::vector<ggml_tensor*> split_qkv(GGMLRunnerContext* ctx, ggml_tensor* projected) {
@@ -260,7 +331,6 @@ namespace MiniMaxH3 {
         std::vector<ggml_tensor*> project_qkv(GGMLRunnerContext* ctx,
                                               const std::vector<ggml_tensor*>& inputs,
                                               const std::function<ggml_tensor*(ggml_tensor*, size_t)>& prepare_input = {}) {
-            auto projection  = std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"]);
             int64_t sequence = 0;
             for (auto input : inputs) {
                 sequence += input->ne[1];
@@ -279,7 +349,7 @@ namespace MiniMaxH3 {
                 ggml_tensor* projected = nullptr;
                 for (size_t i = 0; i < inputs.size(); ++i) {
                     auto input = prepare_input ? prepare_input(inputs[i], i) : inputs[i];
-                    auto part  = projection->forward(ctx, input);
+                    auto part  = project_packed_qkv(ctx, input);
                     projected  = projected == nullptr ? part : ggml_concat(ctx->ggml_ctx, projected, part, 1);
                 }
                 return split_qkv(ctx, projected);
@@ -290,7 +360,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);
                     }
@@ -298,7 +368,7 @@ namespace MiniMaxH3 {
                     if (prepare_input) {
                         part = prepare_input(part, i);
                     }
-                    auto qkv = split_qkv(ctx, projection->forward(ctx, part));
+                    auto qkv = split_qkv(ctx, project_packed_qkv(ctx, part));
                     for (int j = 0; j < 3; ++j) {
                         chunks[j].push_back(ggml_cont(ctx->ggml_ctx, qkv[j]));
                     }
@@ -306,13 +376,13 @@ 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;
         }
 
         ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            return std::dynamic_pointer_cast<Linear>(blocks["out_proj"])->forward(ctx, x);
+            return std::dynamic_pointer_cast<UnaryBlock>(blocks["out_proj"])->forward(ctx, x);
         }
 
         ggml_tensor* attend(GGMLRunnerContext* ctx,
@@ -329,13 +399,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);
         }
@@ -471,34 +541,38 @@ 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 {
         Config config;
 
-        explicit TransformerBlock(const Config& config)
+        TransformerBlock(const Config& config, const LowRankShape& ranks)
             : config(config) {
             blocks["norm1"]      = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
             blocks["norm2"]      = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
             blocks["attn"]       = std::make_shared<Attention>(config.hidden_size,
                                                                config.num_attention_heads,
                                                                config.attention_head_dim,
-                                                               config.qk_norm_eps);
+                                                               config.qk_norm_eps,
+                                                               ranks);
             blocks["mlp"]        = std::make_shared<MLP>(config.hidden_size,
-                                                         config.ffn_hidden_size);
+                                                         config.ffn_hidden_size,
+                                                         ranks.fc1,
+                                                         ranks.fc2);
             blocks["adaln_proj"] = std::make_shared<AdaLayerNormModulation>(config.time_embed_dim,
                                                                             config.hidden_size,
                                                                             6,
@@ -522,9 +596,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]);
@@ -535,7 +609,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;
@@ -578,8 +652,8 @@ namespace MiniMaxH3 {
             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]);
@@ -602,10 +676,16 @@ namespace MiniMaxH3 {
                 blocks["time_embedder"] = std::make_shared<TimeEmbedder>(config.timestep_input_dim,
                                                                          config.time_embed_hidden_size,
                                                                          config.time_embed_dim);
+                if (config.two_time) {
+                    blocks["endpoint_time_embedder"] = std::make_shared<TimeEmbedder>(config.timestep_input_dim,
+                                                                                      config.time_embed_hidden_size,
+                                                                                      config.time_embed_dim);
+                }
             }
             blocks["token_refiner"] = std::make_shared<TokenRefiner>(config);
             for (int64_t i = 0; i < config.num_layers; ++i) {
-                blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(config);
+                blocks["blocks." + std::to_string(i)] =
+                    std::make_shared<TransformerBlock>(config, config.block_ranks.empty() ? LowRankShape{} : config.block_ranks[static_cast<size_t>(i)]);
             }
             blocks["final_layer"] = std::make_shared<FinalLayer>(config);
         }
@@ -615,6 +695,9 @@ namespace MiniMaxH3 {
                          const std::string prefix            = "") override {
             GGMLBlock::init_params(ctx, tensors, prefix);
             params["rope.inv_freq"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.rope_inv_freq_len);
+            if (config.two_time) {
+                params["hyperflow_gate"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+            }
             if (config.uses_adaln_curves()) {
                 params["adaln_t_table"] = ggml_new_tensor_2d(ctx,
                                                              GGML_TYPE_F32,
@@ -639,11 +722,20 @@ namespace MiniMaxH3 {
 
         ggml_tensor* time_embedding(GGMLRunnerContext* ctx,
                                     ggml_tensor* timestep_features,
+                                    ggml_tensor* endpoint_features,
                                     ggml_tensor* curve_indices,
                                     ggml_tensor* curve_upper_indices,
                                     ggml_tensor* curve_fractions) {
             if (!config.uses_adaln_curves()) {
-                return std::dynamic_pointer_cast<TimeEmbedder>(blocks["time_embedder"])->forward(ctx, timestep_features);
+                auto current = std::dynamic_pointer_cast<TimeEmbedder>(blocks["time_embedder"])->forward(ctx, timestep_features);
+                if (!config.two_time) {
+                    return current;
+                }
+                GGML_ASSERT(endpoint_features != nullptr);
+                auto endpoint = std::dynamic_pointer_cast<TimeEmbedder>(blocks["endpoint_time_embedder"])
+                                    ->forward(ctx, endpoint_features);
+                auto delta    = ggml_mul(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, endpoint, current), params["hyperflow_gate"]);
+                return ggml_add(ctx->ggml_ctx, current, delta);
             }
             auto lower = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_indices);
             auto upper = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_upper_indices);
@@ -662,7 +754,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],
@@ -689,6 +781,7 @@ namespace MiniMaxH3 {
                                                       const std::vector<ggml_tensor*>& condition_audios,
                                                       ggml_tensor* position_ids,
                                                       ggml_tensor* timestep_features,
+                                                      ggml_tensor* endpoint_features,
                                                       ggml_tensor* curve_indices,
                                                       ggml_tensor* curve_upper_indices,
                                                       ggml_tensor* curve_fractions,
@@ -703,7 +796,7 @@ namespace MiniMaxH3 {
             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,
@@ -714,7 +807,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,
@@ -727,7 +820,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],
@@ -760,20 +853,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,
@@ -784,6 +877,7 @@ namespace MiniMaxH3 {
 
             auto t_emb = time_embedding(ctx,
                                         timestep_features,
+                                        endpoint_features,
                                         curve_indices,
                                         curve_upper_indices,
                                         curve_fractions);
@@ -798,7 +892,7 @@ 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 video_out   = DiT::unpatchify_3d(ctx,
                                                   output.first,
                                                   video->ne[2] / config.patch_t,
                                                   video->ne[1] / config.patch_h,
@@ -813,9 +907,9 @@ namespace MiniMaxH3 {
                                                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);
+            audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3));
+            video_out        = ggml_ext_scale(ctx, video_out, -1.f);
+            audio_out        = ggml_ext_scale(ctx, audio_out, -1.f);
             return {video_out, audio_out};
         }
     };
@@ -826,7 +920,9 @@ namespace MiniMaxH3 {
         std::vector<SequenceSegment> sequence_segments;
         TokenModulationSpan video_segment{};
         TokenModulationSpan audio_segment{};
+        // Distinct (t, r) rows of the modulation table; r == t for single-time checkpoints.
         std::vector<float> timesteps;
+        std::vector<float> endpoints;
     };
 
     static float video_span(int64_t frame) {
@@ -846,13 +942,15 @@ namespace MiniMaxH3 {
         return result;
     }
 
-    static int find_or_add_timestep(std::vector<float>* values, float value) {
-        auto it = std::find(values->begin(), values->end(), value);
-        if (it != values->end()) {
-            return static_cast<int>(it - values->begin());
+    static int find_or_add_timestep(PackedSequenceLayout* layout, float value, float endpoint) {
+        for (size_t i = 0; i < layout->timesteps.size(); ++i) {
+            if (layout->timesteps[i] == value && layout->endpoints[i] == endpoint) {
+                return static_cast<int>(i);
+            }
         }
-        values->push_back(value);
-        return static_cast<int>(values->size() - 1);
+        layout->timesteps.push_back(value);
+        layout->endpoints.push_back(endpoint);
+        return static_cast<int>(layout->timesteps.size() - 1);
     }
 
     static PackedSequenceLayout build_layout(int64_t text_len,
@@ -866,7 +964,9 @@ namespace MiniMaxH3 {
                                              const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
                                              const sd::Tensor<int32_t>& text_tags,
                                              float video_t,
-                                             float audio_timestep) {
+                                             float audio_timestep,
+                                             float video_r,
+                                             float audio_r) {
         PackedSequenceLayout layout;
         float sqrt_area    = std::sqrt(static_cast<float>(latent_h * latent_w));
         auto h_axis        = spatial_axis(latent_h, sqrt_area);
@@ -884,12 +984,13 @@ namespace MiniMaxH3 {
         }
         layout.sequence_segments.push_back({0, text_len, SequenceKind::TEXT});
 
-        int video_time_row           = find_or_add_timestep(&layout.timesteps, video_t);
-        int audio_time_row           = find_or_add_timestep(&layout.timesteps, audio_timestep);
-        int condition_time_row       = find_or_add_timestep(&layout.timesteps,
-                                                            std::max(video_t, VISUAL_COND_TIMESTEP));
-        int audio_condition_time_row = find_or_add_timestep(&layout.timesteps,
-                                                            std::max(audio_timestep, 1.f));
+        // Conditions stay at their own noise level, so their endpoint equals their timestep.
+        float condition_t            = std::max(video_t, VISUAL_COND_TIMESTEP);
+        float audio_condition_t      = std::max(audio_timestep, 1.f);
+        int video_time_row           = find_or_add_timestep(&layout, video_t, video_r);
+        int audio_time_row           = find_or_add_timestep(&layout, audio_timestep, audio_r);
+        int condition_time_row       = find_or_add_timestep(&layout, condition_t, condition_t);
+        int audio_condition_time_row = find_or_add_timestep(&layout, audio_condition_t, audio_condition_t);
 
         int64_t run_start = 0;
         int current_tag   = text_tags.empty() ? 1 : text_tags[0];
@@ -1056,6 +1157,7 @@ namespace MiniMaxH3 {
         MiniMaxH3Transformer3DModel model;
         sd::Tensor<float> position_input_cache;
         sd::Tensor<float> timestep_feature_input_cache;
+        sd::Tensor<float> endpoint_feature_input_cache;
         sd::Tensor<int32_t> curve_index_input_cache;
         sd::Tensor<int32_t> curve_upper_index_input_cache;
         sd::Tensor<float> curve_fraction_input_cache;
@@ -1099,15 +1201,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,
@@ -1128,7 +1229,8 @@ namespace MiniMaxH3 {
                                  const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
                                  int audio_length,
                                  float video_shift,
-                                 float audio_shift) {
+                                 float audio_shift,
+                                 float video_sigma_next) {
             GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5);
             GGML_ASSERT(audio_length > 0);
             GGML_ASSERT(!context_tensor.empty());
@@ -1166,21 +1268,35 @@ namespace MiniMaxH3 {
             float audio_scale = video_shift / audio_shift;
             float t_v         = 1.f - sigma_v;
             float t_a         = 1.f - sigma_a;
+            float r_v         = t_v;
+            float r_a         = t_a;
+            if (config.two_time) {
+                if (std::isfinite(video_sigma_next)) {
+                    float next = std::clamp(video_sigma_next, 0.f, 1.f);
+                    r_v        = 1.f - next;
+                    r_a        = 1.f - time_shift_sigma(next, video_shift, audio_shift);
+                } else {
+                    LOG_WARN("two-time checkpoint without a step endpoint; using r = t");
+                }
+            }
             // 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);
+            auto audio_ctx = get_context();
+            auto audio     = 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)},
@@ -1188,6 +1304,7 @@ namespace MiniMaxH3 {
             auto positions = make_input(position_input_cache);
 
             ggml_tensor* timestep_features   = nullptr;
+            ggml_tensor* endpoint_features   = nullptr;
             ggml_tensor* curve_indices       = nullptr;
             ggml_tensor* curve_upper_indices = nullptr;
             ggml_tensor* curve_fractions     = nullptr;
@@ -1224,6 +1341,16 @@ namespace MiniMaxH3 {
                                        true,
                                        1.f));
                 timestep_features = make_input(timestep_feature_input_cache);
+                if (config.two_time) {
+                    endpoint_feature_input_cache = sd::Tensor<float>(
+                        {config.timestep_input_dim, static_cast<int64_t>(layout.endpoints.size())},
+                        timestep_embedding(layout.endpoints,
+                                           static_cast<int>(config.timestep_input_dim),
+                                           10000,
+                                           true,
+                                           1.f));
+                    endpoint_features = make_input(endpoint_feature_input_cache);
+                }
             }
 
             auto graph      = new_graph_custom(H3_GRAPH_SIZE);
@@ -1236,6 +1363,7 @@ namespace MiniMaxH3 {
                                             audio_condition_inputs,
                                             positions,
                                             timestep_features,
+                                            endpoint_features,
                                             curve_indices,
                                             curve_upper_indices,
                                             curve_fractions,
@@ -1245,11 +1373,11 @@ namespace MiniMaxH3 {
                                             layout.audio_segment);
             // 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,
+                                     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   = merge_av_latents(compute_ctx, output.first, output.second);
+            auto merged   = merge_av_latents(&runner_ctx, output.first, output.second);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }
@@ -1279,7 +1407,8 @@ namespace MiniMaxH3 {
                                    reference_blocks,
                                    extra->audio_length,
                                    extra->video_sigma_shift,
-                                   extra->audio_sigma_shift);
+                                   extra->audio_sigma_shift,
+                                   extra->video_sigma_next);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
                                                                        n_threads,

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_split_qkv(ctx, qkv);
         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/model.hpp

@@ -145,6 +145,7 @@ struct MiniMaxH3DiffusionExtra {
     int audio_length                                              = 0;
     float video_sigma_shift                                       = 12.f;
     float audio_sigma_shift                                       = 3.f;
+    float video_sigma_next                                        = NAN;
 };
 
 struct MiniT2IDiffusionExtra {
@@ -237,6 +238,10 @@ public:
     virtual sd::Tensor<float> compute(int n_threads,
                                       const DiffusionParams& diffusion_params) = 0;
 
+    virtual int get_patch_size() const {
+        return 1;
+    }
+
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
         get_param_tensors(tensors, prefix);
     }

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_split_qkv(ctx, qkv);
             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_split_qkv(ctx, qkv_x_proj->forward(ctx, x));
             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_split_qkv(ctx, qkv_y_proj->forward(ctx, y));
             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);
         }
     };
 

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

@@ -0,0 +1,325 @@
+#ifndef __SD_MODEL_DIFFUSION_SEEDVR2_HPP__
+#define __SD_MODEL_DIFFUSION_SEEDVR2_HPP__
+
+#include "model/common/vsr.hpp"
+
+namespace SeedVR2 {
+
+    struct SeedVR2Config {
+        int dim        = 0;
+        int layers     = 0;
+        int mm_layers  = 0;
+        int hidden     = 0;
+        int text_dim   = 0;
+        int text_len   = 0;
+        int rope_freqs = 0;
+        bool swiglu    = false;
+
+        static SeedVR2Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            SeedVR2Config config;
+            auto at = [&](const std::string& key) -> const TensorStorage& {
+                auto found = weights.find(prefix + "." + key);
+                if (found == weights.end())
+                    throw std::runtime_error("Missing SeedVR2 tensor: " + prefix + "." + key);
+                return found->second;
+            };
+            config.dim        = int(at("vid_in.proj.weight").ne[1]);
+            config.hidden     = int(at("blocks.0.mlp.vid.proj_in.weight").ne[1]);
+            config.text_dim   = int(at("positive_conditioning").ne[0]);
+            config.text_len   = int(at("positive_conditioning").ne[1]);
+            config.rope_freqs = int(at("blocks.0.attn.rope.rope.freqs").ne[0]);
+            config.swiglu     = weights.count(prefix + ".blocks.0.mlp.vid.proj_in_gate.weight") != 0;
+            while (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.rope.rope.freqs")) {
+                if (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.proj_qkv.vid.weight"))
+                    ++config.mm_layers;
+                ++config.layers;
+            }
+            if (at("vid_in.proj.weight").logical_ne0() != 132 || at("vid_out.proj.weight").ne[1] != 64 ||
+                config.text_dim != 5120 ||
+                !(config.swiglu ? config.dim == 2560 && config.layers == 32 && config.mm_layers == 10 && config.rope_freqs == 21
+                                : config.dim == 3072 && config.layers == 36 && config.mm_layers == 36 && config.rope_freqs == 10))
+                throw std::runtime_error("Unsupported SeedVR2 architecture");
+            LOG_VERBOSE("seedvr2: dim = %d, layers = %d, mm_layers = %d, hidden = %d, rope_freqs = %d", config.dim, config.layers, config.mm_layers, config.hidden, config.rope_freqs);
+            return config;
+        }
+    };
+
+    struct WindowPlan {
+        sd::Tensor<int32_t> order;
+        sd::Tensor<int32_t> reverse;
+        sd::Tensor<float> positions;
+        sd::Tensor<float> text_positions;
+        std::vector<int> lengths;
+
+        WindowPlan(int width, int height, int frames, int text_len, bool shifted, bool multimodal)
+            : order({width * height * frames}), reverse({width * height * frames}), positions({1, 3, width * height * frames}), text_positions({1, 3, text_len}) {
+            const double scale = std::sqrt(720.0 * 1280.0 / 256.0 / (width * height));
+            const int wh       = int(std::ceil(std::nearbyint(height * scale) / 3));
+            const int ww       = int(std::ceil(std::nearbyint(width * scale) / 3));
+            const int wt       = (std::min(frames, 30) + 3) / 4;
+            const double st    = shifted && wt < frames ? 0.5 : 0;
+            const double sh    = shifted && wh < height ? 0.5 : 0;
+            const double sw    = shifted && ww < width ? 0.5 : 0;
+            const int nt       = st ? int(std::ceil((frames - st) / wt)) + 1 : (frames + wt - 1) / wt;
+            const int nh       = sh ? int(std::ceil((height - sh) / wh)) + 1 : (height + wh - 1) / wh;
+            const int nw       = sw ? int(std::ceil((width - sw) / ww)) + 1 : (width + ww - 1) / ww;
+            int index          = 0;
+            for (int iw = 0; iw < nw; ++iw)
+                for (int ih = 0; ih < nh; ++ih)
+                    for (int it = 0; it < nt; ++it) {
+                        int x0 = std::max(0, int((iw - sw) * ww)), x1 = std::min(width, int((iw - sw + 1) * ww));
+                        int y0 = std::max(0, int((ih - sh) * wh)), y1 = std::min(height, int((ih - sh + 1) * wh));
+                        int t0 = std::max(0, int((it - st) * wt)), t1 = std::min(frames, int((it - st + 1) * wt));
+                        if (x1 <= x0 || y1 <= y0 || t1 <= t0)
+                            continue;
+                        lengths.push_back((x1 - x0) * (y1 - y0) * (t1 - t0));
+                        for (int t = t0; t < t1; ++t)
+                            for (int y = y0; y < y1; ++y)
+                                for (int x = x0; x < x1; ++x, ++index) {
+                                    int original             = (t * height + y) * width + x;
+                                    order.data()[index]      = original;
+                                    reverse.data()[original] = index;
+                                    const int pos[]          = {t - t0, y - y0, x - x0};
+                                    const int size[]         = {t1 - t0, y1 - y0, x1 - x0};
+                                    for (int a = 0; a < 3; ++a)
+                                        positions.data()[3 * index + a] = multimodal ? float(pos[a] + (a == 0 ? text_len : 0))
+                                                                                     : (size[a] == 1 ? -1.f : -1.f + 2.f * pos[a] / (size[a] - 1));
+                                }
+                    }
+            if (index != width * height * frames)
+                throw std::runtime_error("SeedVR2 window partition is incomplete");
+            for (int i = 0; i < text_len; ++i)
+                for (int a = 0; a < 3; ++a)
+                    text_positions.data()[3 * i + a] = multimodal ? float(i) : 0.f;
+        }
+    };
+
+    class AdaSingle : public GGMLBlock {
+        int dim_;
+        bool output_;
+
+    protected:
+        void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+            for (const auto& layer : output_ ? std::vector<std::string>{"out"} : std::vector<std::string>{"attn", "mlp"}) {
+                for (const auto& mode : output_ ? std::vector<std::string>{"shift", "scale"} : std::vector<std::string>{"shift", "scale", "gate"})
+                    params[layer + "_" + mode] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim_);
+            }
+        }
+
+    public:
+        AdaSingle(int dim, bool output = false)
+            : dim_(dim), output_(output) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb, const std::string& layer, bool gate = false) {
+            auto c = ctx->ggml_ctx;
+            // The reference reshapes embedding to [D, 2, 3], not [6, D].
+            const int offset = layer == "mlp" ? 3 : 0;
+            auto mod         = [&](int i, const char* name) {
+                auto value = ggml_cont(c, ggml_view_2d(c, emb, 1, dim_, 6 * emb->nb[0], (offset + i) * emb->nb[0]));
+                return ggml_add(c, ggml_reshape_1d(c, value, dim_), params[layer + "_" + name]);
+            };
+            if (gate)
+                return ggml_mul(c, x, mod(2, "gate"));
+            // AdaSingle uses scale directly; ggml_ext_modulate adds an identity scale.
+            return ggml_add(c, ggml_mul(c, x, mod(1, "scale")), mod(0, "shift"));
+        }
+    };
+
+    class SRMLP : public UnaryBlock {
+        bool swiglu_;
+
+    public:
+        explicit SRMLP(const SeedVR2Config& config)
+            : swiglu_(config.swiglu) {
+            blocks["proj_in"]  = std::make_shared<Linear>(config.dim, config.hidden, !swiglu_);
+            blocks["proj_out"] = std::make_shared<Linear>(config.hidden, config.dim, !swiglu_);
+            if (swiglu_)
+                blocks["proj_in_gate"] = std::make_shared<Linear>(config.dim, config.hidden, false);
+        }
+        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, up);
+            return std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, hidden);
+        }
+    };
+
+    class NaMMSRTransformerBlock : public GGMLBlock {
+        SeedVR2Config config_;
+        bool shared_;
+        bool last_;
+
+    protected:
+        void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+            params["attn.rope.rope.freqs"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config_.rope_freqs);
+        }
+
+    public:
+        NaMMSRTransformerBlock(const SeedVR2Config& config, int index)
+            : config_(config), shared_(index >= config.mm_layers), last_(config.swiglu && index == config.layers - 1) {
+            for (const auto& branch : shared_ ? std::vector<std::string>{"all"} : std::vector<std::string>{"vid", "txt"}) {
+                blocks["attn.proj_qkv." + branch] = std::make_shared<Linear>(config.dim, 3 * config.dim, false);
+                blocks["attn.proj_out." + branch] = std::make_shared<Linear>(config.dim, config.dim);
+                blocks["attn.norm_q." + branch]   = std::make_shared<RMSNorm>(128, 1e-5f);
+                blocks["attn.norm_k." + branch]   = std::make_shared<RMSNorm>(128, 1e-5f);
+                blocks["ada." + branch]           = std::make_shared<AdaSingle>(config.dim);
+                blocks["mlp." + branch]           = std::make_shared<SRMLP>(config);
+            }
+        }
+
+        ggml_tensor* rotation(GGMLRunnerContext* ctx, ggml_tensor* positions) {
+            auto c      = ctx->ggml_ctx;
+            const int n = config_.rope_freqs;
+            auto p      = ggml_repeat(c, positions, ggml_new_tensor_3d(c, GGML_TYPE_F32, n, 3, positions->ne[2]));
+            auto angle  = ggml_mul(c, p, params["attn.rope.rope.freqs"]);
+            angle       = ggml_reshape_2d(c, angle, 3 * n, positions->ne[2]);
+            angle       = ggml_pad(c, angle, 64 - 3 * n, 0, 0, 0);
+            // Keep tokens out of ne[3]: CUDA concat otherwise launches once per token.
+            auto cs     = ggml_reshape_3d(c, ggml_cos(c, angle), 1, 1, 64 * positions->ne[2]);
+            auto sn     = ggml_reshape_3d(c, ggml_sin(c, angle), 1, 1, 64 * positions->ne[2]);
+            auto matrix = ggml_concat(c, ggml_concat(c, cs, ggml_scale(c, sn, -1.f), 0), ggml_concat(c, sn, cs, 0), 1);
+            return ggml_reshape_4d(c, matrix, 2, 2, 64, positions->ne[2]);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vid, ggml_tensor* txt, ggml_tensor* emb, const WindowPlan& plan, VSR::GraphInputs& inputs) {
+            auto c        = ctx->ggml_ctx;
+            const int dim = config_.dim, heads = dim / 128;
+            auto branch = [&](bool text) -> std::string { return shared_ ? "all" : text ? "txt"
+                                                                                        : "vid"; };
+            auto linear = [&](const std::string& name, bool text, ggml_tensor* x) {
+                return std::dynamic_pointer_cast<Linear>(blocks[name + "." + branch(text)])->forward(ctx, x);
+            };
+            auto ada = [&](ggml_tensor* x, bool text, const std::string& layer, bool gate = false) {
+                if (text && last_)
+                    return x;
+                return std::dynamic_pointer_cast<AdaSingle>(blocks["ada." + branch(text)])->forward(ctx, x, emb, layer, gate);
+            };
+            auto vqkv = linear("attn.proj_qkv", false, ada(ggml_rms_norm(c, vid, 1e-5f), false, "attn"));
+            auto tqkv = linear("attn.proj_qkv", true, ada(ggml_rms_norm(c, txt, 1e-5f), true, "attn"));
+            vqkv      = ggml_get_rows(c, vqkv, inputs.input(ctx, plan.order));
+            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(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);
+                    h        = ggml_reshape_2d(c, h, dim, x->ne[1]);
+                    parts[i] = VSR::apply_rope(ctx, h, pe, heads);
+                }
+                return parts;
+            };
+            auto vp = qkv(vqkv, false, vpe), tp = qkv(tqkv, true, tpe);
+            std::vector<ggml_tensor*> video_windows;
+            ggml_tensor* tout = nullptr;
+            int start         = 0;
+            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(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(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));
+            txt       = ggml_add(c, txt, ada(linear("attn.proj_out", true, tout), true, "attn", true));
+            auto mlp  = [&](ggml_tensor* x, bool text) {
+                auto h = ada(ggml_rms_norm(c, x, 1e-5f), text, "mlp");
+                h      = std::dynamic_pointer_cast<SRMLP>(blocks["mlp." + branch(text)])->forward(ctx, h);
+                return ggml_add(c, x, ada(h, text, "mlp", true));
+            };
+            vid = mlp(vid, false);
+            if (!last_)
+                txt = mlp(txt, true);
+            return {vid, txt};
+        }
+    };
+
+    class NaDiT : public GGMLBlock {
+        SeedVR2Config config_;
+
+    protected:
+        void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+            params["positive_conditioning"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config_.text_dim, config_.text_len);
+        }
+
+    public:
+        explicit NaDiT(const SeedVR2Config& config)
+            : config_(config) {
+            blocks["vid_in.proj"]     = std::make_shared<Linear>(132, config.dim);
+            blocks["txt_in"]          = std::make_shared<Linear>(config.text_dim, config.dim);
+            blocks["emb_in.proj_in"]  = std::make_shared<Linear>(256, config.dim);
+            blocks["emb_in.proj_hid"] = std::make_shared<Linear>(config.dim, config.dim);
+            blocks["emb_in.proj_out"] = std::make_shared<Linear>(config.dim, 6 * config.dim);
+            for (int i = 0; i < config.layers; ++i)
+                blocks["blocks." + std::to_string(i)] = std::make_shared<NaMMSRTransformerBlock>(config, i);
+            blocks["vid_out.proj"] = std::make_shared<Linear>(config.dim, 64);
+            if (config.swiglu) {
+                blocks["vid_out_norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+                blocks["vid_out_ada"]  = std::make_shared<AdaSingle>(config.dim, true);
+            }
+        }
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* timestep, const std::vector<WindowPlan>& plans, VSR::GraphInputs& inputs) {
+            auto linear = [&](const std::string& name, ggml_tensor* x) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x); };
+            auto vid    = linear("vid_in.proj", input);
+            auto txt    = linear("txt_in", params["positive_conditioning"]);
+            auto emb    = linear("emb_in.proj_in", timestep);
+            emb         = linear("emb_in.proj_hid", ggml_silu(ctx->ggml_ctx, emb));
+            emb         = linear("emb_in.proj_out", ggml_silu(ctx->ggml_ctx, emb));
+            for (int i = 0; i < config_.layers; ++i) {
+                auto name          = "blocks." + std::to_string(i);
+                std::tie(vid, txt) = std::dynamic_pointer_cast<NaMMSRTransformerBlock>(blocks[name])->forward(ctx, vid, txt, emb, plans[i % 2], inputs);
+                sd::ggml_graph_cut::mark_graph_cut(vid, "seedvr2." + name, "vid");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "seedvr2." + name, "txt");
+            }
+            if (config_.swiglu) {
+                vid = std::dynamic_pointer_cast<RMSNorm>(blocks["vid_out_norm"])->forward(ctx, vid);
+                vid = std::dynamic_pointer_cast<AdaSingle>(blocks["vid_out_ada"])->forward(ctx, vid, emb, "out");
+            }
+            return linear("vid_out.proj", vid);
+        }
+    };
+
+    struct SeedVR2Runner : GGMLRunner {
+        SeedVR2Config config;
+        NaDiT model;
+        sd::Tensor<float> timestep;
+
+        SeedVR2Runner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr<DeviceResidencyManager> manager)
+            : GGMLRunner(backend, manager), config(SeedVR2Config::detect_from_weights(weights, "diffusion")), model(config), timestep({256}) {
+            model.init(params_ctx, weights, "diffusion");
+            for (int i = 0; i < 128; ++i) {
+                const double angle       = 1000.0 * std::pow(10000.0, -double(i) / 128);
+                timestep.data()[i]       = float(std::sin(angle));
+                timestep.data()[i + 128] = float(std::cos(angle));
+            }
+        }
+        std::string get_desc() override { return "SeedVR2 DiT"; }
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) { model.get_param_tensors(tensors, "diffusion"); }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input, int width, int height, int frames) {
+            std::vector<WindowPlan> plans;
+            for (bool shifted : {false, true})
+                plans.emplace_back(width, height, frames, config.text_len, shifted, config.swiglu);
+            VSR::GraphInputs inputs;
+            auto graph = [&]() {
+                inputs.begin_graph();
+                auto gf  = new_graph_custom(MAX_GRAPH_SIZE);
+                auto ctx = get_context(gf);
+                auto out = model.forward(&ctx, make_input(input), make_input(timestep), plans, inputs);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            auto output = GGMLRunner::compute(graph, n_threads);
+            return output ? *output : sd::Tensor<float>();
+        }
+    };
+}
+
+#endif  // __SD_MODEL_DIFFUSION_SEEDVR2_HPP__

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/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_split_qkv(ctx, qkv);
             }
 
             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;